import gradio as gr import subprocess import os import sys import yaml from pathlib import Path import time import threading import tempfile import shutil import gc # ガベージコレクション用 # --- 定数 --- # Dockerfile内でクローンされるパスに合わせる SCRIPT_DIR = Path(__file__).parent SBV2_REPO_PATH = SCRIPT_DIR / "Style-Bert-VITS2" # ダウンロード用ファイルの一時置き場 (コンテナ内に作成) OUTPUT_DIR = SCRIPT_DIR / "outputs" # --- ヘルパー関数 --- def add_sbv2_to_path(): """Style-Bert-VITS2リポジトリのパスを sys.path に追加""" repo_path_str = str(SBV2_REPO_PATH.resolve()) if SBV2_REPO_PATH.exists() and repo_path_str not in sys.path: sys.path.insert(0, repo_path_str) print(f"Added {repo_path_str} to sys.path") elif not SBV2_REPO_PATH.exists(): print(f"Warning: Style-Bert-VITS2 repository not found at {SBV2_REPO_PATH}") def stream_process_output(process, log_list): """サブプロセスの標準出力/エラーをリアルタイムでリストに追加""" try: if process.stdout: for line in iter(process.stdout.readline, ''): log_list.append(line.strip()) # 余分な改行を削除 if process.stderr: for line in iter(process.stderr.readline, ''): processed_line = f"stderr: {line.strip()}" # 警告はそのまま、他はエラーとして強調 (任意) if "warning" not in line.lower(): processed_line = f"ERROR (stderr): {line.strip()}" log_list.append(processed_line) except Exception as e: log_list.append(f"Error reading process stream: {e}") # --- Gradio アプリのバックエンド関数 --- def convert_safetensors_to_onnx_gradio( safetensors_file_obj, config_file_obj, style_vectors_file_obj ): # gr.Progress は削除 """ アップロードされたSafetensors, config.json, style_vectors.npy を使って ONNXに変換し、結果をダウンロード可能にする。 """ log = ["Starting ONNX conversion..."] # 初期状態ではダウンロードファイルは空 yield "\n".join(log), None # --- ファイルアップロードの検証 --- if safetensors_file_obj is None: log.append("❌ Error: Safetensors file is missing. Please upload the .safetensors file.") yield "\n".join(log), None return if config_file_obj is None: log.append("❌ Error: config.json file is missing. Please upload the config.json file.") yield "\n".join(log), None return if style_vectors_file_obj is None: log.append("❌ Error: style_vectors.npy file is missing. Please upload the style_vectors.npy file.") yield "\n".join(log), None return # --- Style-Bert-VITS2 パスの確認 --- add_sbv2_to_path() if not SBV2_REPO_PATH.exists(): log.append(f"❌ Error: Style-Bert-VITS2 repository not found at {SBV2_REPO_PATH}. Check Space build logs.") yield "\n".join(log), None return # --- 出力ディレクトリ作成 --- OUTPUT_DIR.mkdir(parents=True, exist_ok=True) onnx_output_path_str = None # 最終的なONNXファイルパス (文字列) current_log = log[:] # ログリストをコピー try: # --- 一時ディレクトリを作成して処理 --- with tempfile.TemporaryDirectory() as temp_dir: temp_dir_path = Path(temp_dir) current_log.append(f"📁 Created temporary directory: {temp_dir_path}") yield "\n".join(current_log), None # UI更新 # --- SBV2が期待するディレクトリ構造を一時ディレクトリ内に作成 --- # モデル名を .safetensors ファイル名から取得 (拡張子なし) safetensors_filename = Path(safetensors_file_obj.name).name if not safetensors_filename.lower().endswith(".safetensors"): current_log.append(f"❌ Error: Invalid safetensors filename: {safetensors_filename}") yield "\n".join(current_log), None return model_name = Path(safetensors_filename).stem # 拡張子を除いた部分 # assets_root を一時ディレクトリ自体にする assets_root = temp_dir_path # assets_root の下に model_name のディレクトリを作成 model_dir_in_temp = assets_root / model_name model_dir_in_temp.mkdir(exist_ok=True) current_log.append(f" - Created model directory: {model_dir_in_temp.relative_to(assets_root)}") yield "\n".join(current_log), None # --- 3つのファイルを model_dir_in_temp にコピー --- files_to_copy = { "safetensors": safetensors_file_obj, "config.json": config_file_obj, "style_vectors.npy": style_vectors_file_obj, } copied_paths = {} for file_key, file_obj in files_to_copy.items(): original_filename = Path(file_obj.name).name # ファイル名の基本的な検証 (サニタイズはより厳密に行うことも可能) if "/" in original_filename or "\\" in original_filename or ".." in original_filename: current_log.append(f"❌ Error: Invalid characters found in filename: {original_filename}") yield "\n".join(current_log), None return # tryブロックを抜ける # 期待されるファイル名と一致しているか確認 (config と style_vectors) if file_key == "config.json" and original_filename.lower() != "config.json": current_log.append(f"⚠️ Warning: Uploaded JSON file name is '{original_filename}', expected 'config.json'. Using uploaded name.") if file_key == "style_vectors.npy" and original_filename.lower() != "style_vectors.npy": current_log.append(f"⚠️ Warning: Uploaded NPY file name is '{original_filename}', expected 'style_vectors.npy'. Using uploaded name.") destination_path = model_dir_in_temp / original_filename try: shutil.copy(file_obj.name, destination_path) current_log.append(f" - Copied '{original_filename}' to model directory.") # .safetensorsファイルのパスを保存しておく if file_key == "safetensors": copied_paths["safetensors"] = destination_path except Exception as e: current_log.append(f"❌ Error copying file '{original_filename}': {e}") yield "\n".join(current_log), None return # tryブロックを抜ける yield "\n".join(current_log), None # 各ファイルコピー後にUI更新 # safetensorsファイルがコピーされたか確認 temp_safetensors_path = copied_paths.get("safetensors") if not temp_safetensors_path: current_log.append("❌ Error: Failed to locate the copied safetensors file in the temporary directory.") yield "\n".join(current_log), None return current_log.append(f"✅ All required files copied to temporary model directory.") current_log.append(f" - Using temporary assets_root: {assets_root}") yield "\n".join(current_log), None # --- paths.yml を一時的に設定 --- config_path = SBV2_REPO_PATH / "configs" / "paths.yml" config_path.parent.mkdir(parents=True, exist_ok=True) # dataset_root は今回は使わないが設定はしておく (assets_rootと同じ場所) paths_config = {"dataset_root": str(assets_root.resolve()), "assets_root": str(assets_root.resolve())} with open(config_path, "w", encoding="utf-8") as f: yaml.dump(paths_config, f) current_log.append(f" - Saved temporary paths config to {config_path}") yield "\n".join(current_log), None # --- ONNX変換スクリプト実行 --- current_log.append(f"\n🚀 Starting ONNX conversion script for model '{model_name}'...") convert_script = SBV2_REPO_PATH / "convert_onnx.py" if not convert_script.exists(): current_log.append(f"❌ Error: convert_onnx.py not found at '{convert_script}'. Check repository setup.") yield "\n".join(current_log), None return # tryブロックを抜ける python_executable = sys.executable command = [ python_executable, str(convert_script.resolve()), "--model", str(temp_safetensors_path.resolve()) # 一時ディレクトリ内の .safetensors パス ] current_log.append(f"\n Running command: {' '.join(command)}") yield "\n".join(current_log), None process_env = os.environ.copy() process = subprocess.Popen(command, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True, encoding='utf-8', errors='replace', cwd=SBV2_REPO_PATH, # スクリプトの場所で実行 env=process_env) # ログ出力用リスト (スレッドと共有) process_output_lines = ["\n--- Conversion Script Output ---"] thread = threading.Thread(target=stream_process_output, args=(process, process_output_lines)) thread.start() # 進捗表示のためのループ while thread.is_alive(): yield "\n".join(current_log + process_output_lines), None time.sleep(0.3) # 更新頻度 # スレッド終了待ちとプロセス終了待ち thread.join() try: process.wait(timeout=12000) # 20分タイムアウト (モデルサイズにより調整) except subprocess.TimeoutExpired: current_log.extend(process_output_lines) # ここまでのログを追加 current_log.append("\n❌ Error: Conversion process timed out after 20 minutes.") process.kill() yield "\n".join(current_log), None return # tryブロックを抜ける # 最終的なプロセス出力を取得 final_stdout, final_stderr = process.communicate() if final_stdout: process_output_lines.extend(final_stdout.strip().split('\n')) if final_stderr: processed_stderr = [] for line in final_stderr.strip().split('\n'): processed_line = f"stderr: {line.strip()}" if "warning" not in line.lower() and line.strip(): # 空行と警告以外 processed_line = f"ERROR (stderr): {line.strip()}" processed_stderr.append(processed_line) if any(line.startswith("ERROR") for line in processed_stderr): process_output_lines.append("--- Errors/Warnings (stderr) ---") process_output_lines.extend(processed_stderr) process_output_lines.append("-----------------------------") elif processed_stderr: # 警告のみの場合 process_output_lines.append("--- Warnings (stderr) ---") process_output_lines.extend(processed_stderr) process_output_lines.append("------------------------") # 全てのプロセスログをメインログに追加 current_log.extend(process_output_lines) current_log.append("--- End Script Output ---") current_log.append("\n-------------------------------") # --- 結果の確認と出力ファイルのコピー --- if process.returncode == 0: current_log.append("✅ ONNX conversion command finished successfully.") # 期待されるONNXファイルパス (入力と同じディレクトリ内) expected_onnx_path_in_temp = temp_safetensors_path.with_suffix(".onnx") if expected_onnx_path_in_temp.exists(): current_log.append(f" - Found converted ONNX file: {expected_onnx_path_in_temp.name}") # 一時ディレクトリから永続的な出力ディレクトリにコピー final_onnx_path = OUTPUT_DIR / expected_onnx_path_in_temp.name try: shutil.copy(expected_onnx_path_in_temp, final_onnx_path) current_log.append(f" - Copied ONNX file for download to: {final_onnx_path.relative_to(SCRIPT_DIR)}") onnx_output_path_str = str(final_onnx_path) # ダウンロード用ファイルパスを設定 except Exception as e: current_log.append(f"❌ Error copying ONNX file to output directory: {e}") else: current_log.append(f"⚠️ Warning: Expected ONNX file not found at '{expected_onnx_path_in_temp.name}'. Check script output above.") else: current_log.append(f"❌ ONNX conversion command failed with return code {process.returncode}.") current_log.append(" Please check the logs above for errors (especially lines starting with 'ERROR').") # 一時ディレクトリが自動で削除される前に最終結果をyield yield "\n".join(current_log), onnx_output_path_str except FileNotFoundError as e: current_log.append(f"\n❌ Error: A required command or file was not found: {e.filename}. Check Dockerfile setup and PATH.") current_log.append(f"{e}") yield "\n".join(current_log), None except Exception as e: current_log.append(f"\n❌ An unexpected error occurred: {e}") import traceback current_log.append(traceback.format_exc()) yield "\n".join(current_log), None finally: # ガベージコレクション gc.collect() print("Conversion function finished.") # サーバーログ用 # 最後のyieldでUIを最終状態に更新 # yield "\n".join(current_log), onnx_output_path_str # tryブロック内で既に返している # --- Gradio Interface --- with gr.Blocks(theme=gr.themes.Soft()) as demo: gr.Markdown("# Style-Bert-VITS2 Safetensors to ONNX Converter") gr.Markdown( "Upload your model's `.safetensors`, `config.json`, and `style_vectors.npy` files. " "The application will convert the model to ONNX format, and you can download the resulting `.onnx` file." ) gr.Markdown( "_(Environment setup is handled automatically when this Space starts.)_" ) with gr.Row(): with gr.Column(scale=1): gr.Markdown("### 1. Upload Model Files") safetensors_upload = gr.File( label="Safetensors Model (.safetensors)", file_types=[".safetensors"], ) config_upload = gr.File( label="Config File (config.json)", file_types=[".json"], ) style_vectors_upload = gr.File( label="Style Vectors (style_vectors.npy)", file_types=[".npy"], ) convert_button = gr.Button("2. Convert to ONNX", variant="primary", elem_id="convert_button") gr.Markdown("---") gr.Markdown("### 3. Download Result") onnx_download = gr.File( label="ONNX Model (.onnx)", interactive=False, # 出力専用 ) gr.Markdown( "**Note:** Conversion can take **several minutes** (5-20+ min depending on model size and hardware). " "Please be patient. The log on the right shows the progress." ) with gr.Column(scale=2): output_log = gr.Textbox( label="Conversion Log", lines=30, # 高さをさらに増やす interactive=False, autoscroll=True, max_lines=2000 # ログが増える可能性 ) # ボタンクリック時のアクション設定 convert_button.click( convert_safetensors_to_onnx_gradio, inputs=[safetensors_upload, config_upload, style_vectors_upload], outputs=[output_log, onnx_download] # ログとダウンロードファイルの2つを出力 ) # --- アプリの起動 --- if __name__ == "__main__": # Style-Bert-VITS2 へのパスを追加 add_sbv2_to_path() # 出力ディレクトリ作成 (存在確認含む) OUTPUT_DIR.mkdir(parents=True, exist_ok=True) print(f"Output directory: {OUTPUT_DIR.resolve()}") # Gradioアプリを起動 demo.queue().launch()