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, ''): # WARNING以外のstderrはエラーとして強調表示しても良い 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, progress=gr.Progress(track_tqdm=True)): """ アップロードされたSafetensorsモデルをONNXに変換し、結果をダウンロード可能にする。 """ log = ["Starting ONNX conversion..."] # 初期状態ではダウンロードファイルは空 yield "\n".join(log), None if safetensors_file_obj is None: log.append("Error: No safetensors file uploaded. Please upload a .safetensors file.") # エラーメッセージを表示し、ダウンロードはNoneのまま 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) # 以前の出力ファイルを削除 (オプション、ディスクスペース節約のため) # for item in OUTPUT_DIR.glob("*.onnx"): # try: # item.unlink() # log.append(f"Cleaned up previous output: {item.name}") # except OSError as e: # log.append(f"Warning: Could not delete previous file {item.name}: {e}") onnx_output_path_str = None # 最終的なONNXファイルパス (文字列) current_log = log[:] # ログリストをコピー try: # 一時ディレクトリを作成して処理 with tempfile.TemporaryDirectory() as temp_dir: temp_dir_path = Path(temp_dir) # Gradio Fileコンポーネントは一時ファイルパスを .name で持つ original_filename = Path(safetensors_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 temp_safetensors_path_in_root = temp_dir_path / original_filename current_log.append(f"Processing uploaded file: {original_filename}") current_log.append(f"Copying to temporary location: {temp_safetensors_path_in_root}") yield "\n".join(current_log), None # アップロードされたファイルオブジェクトから一時ディレクトリにコピー shutil.copy(safetensors_file_obj.name, temp_safetensors_path_in_root) # --- SBV2が期待するディレクトリ構造を一時ディレクトリ内に作成 --- # モデル名をファイル名から取得 (拡張子なし) model_name = temp_safetensors_path_in_root.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) # safetensorsファイルを model_name ディレクトリに移動 temp_safetensors_path = model_dir_in_temp / original_filename shutil.move(temp_safetensors_path_in_root, temp_safetensors_path) # dataset_root も assets_root と同じにしておく (今回は使用しない) dataset_root = assets_root current_log.append(f"Using temporary assets_root: {assets_root}") current_log.append(f"Created temporary model directory: {model_dir_in_temp}") current_log.append(f"Using temporary model path: {temp_safetensors_path}") 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) paths_config = {"dataset_root": str(dataset_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"\nStarting ONNX conversion script for: {temp_safetensors_path.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}'.") 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()) # 一時ディレクトリ内のモデルパス ] current_log.append(f"\nRunning command: {' '.join(command)}") yield "\n".join(current_log), None process_env = os.environ.copy() # メモリリーク対策？ (あまり効果ないかも) # process_env["PYTORCH_CUDA_ALLOC_CONF"] = "max_split_size_mb:128" 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 = [] thread = threading.Thread(target=stream_process_output, args=(process, process_output_lines)) thread.start() # 進捗表示のためのループ while thread.is_alive(): # yieldでGradio UIを更新 (現在のログ + プロセスからのログ) yield "\n".join(current_log + process_output_lines), None time.sleep(0.2) # 更新頻度 # スレッド終了待ち thread.join() # プロセス終了待ち (タイムアウト設定) try: process.wait(timeout=900) # 15分タイムアウト except subprocess.TimeoutExpired: current_log.extend(process_output_lines) # ここまでのログを追加 current_log.append("\nError: Conversion process timed out after 15 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(): processed_line = f"ERROR (stderr): {line.strip()}" processed_stderr.append(processed_line) if processed_stderr: process_output_lines.append("--- stderr ---") process_output_lines.extend(processed_stderr) process_output_lines.append("--------------") # 全てのプロセスログをメインログに追加 current_log.extend(process_output_lines) 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 shutil.copy(expected_onnx_path_in_temp, final_onnx_path) current_log.append(f"Copied ONNX file for download to: {final_onnx_path}") onnx_output_path_str = str(final_onnx_path) # ダウンロード用ファイルパスを設定 else: current_log.append(f"Warning: Expected ONNX file not found at '{expected_onnx_path_in_temp}'. Please check the logs.") 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"\nError: 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"\nAn unexpected error occurred: {e}") import traceback current_log.append(traceback.format_exc()) # エラー発生時も最終ログとNoneを返す yield "\n".join(current_log), None finally: # ガベージコレクションを試みる (メモリ解放目的) gc.collect() # 最終的な状態をUIに反映させるための最後のyield (重要) # tryブロック内で既にyieldしている場合でも、finallyで再度yieldすることで # UIが最終状態（エラーメッセージや成功メッセージ＋ダウンロードリンク）に更新される print("Conversion function finished.") # サーバーログ用 # ここで再度yieldするとUIの更新が確実になることがある # yield "\n".join(current_log), onnx_output_path_str # --- Gradio Interface --- with gr.Blocks(theme=gr.themes.Soft()) as demo: gr.Markdown("# Style-Bert-VITS2 Safetensors to ONNX Converter") gr.Markdown( "Upload your `.safetensors` model file, convert it to ONNX format, and download the result. " "The environment setup (cloning repo, installing dependencies, running initialize.py) " "is handled automatically when this Space starts." ) with gr.Row(): with gr.Column(scale=1): safetensors_upload = gr.File( label="1. Upload Safetensors Model", file_types=[".safetensors"], # file_count="single" (default) ) convert_button = gr.Button("2. Convert to ONNX", variant="primary") gr.Markdown("---") onnx_download = gr.File( label="3. Download ONNX Model", interactive=False, # 出力専用 ) gr.Markdown( "**Note:** Conversion can take several minutes, especially on free hardware. " "Please be patient. The log on the right will update during the process." ) with gr.Column(scale=2): output_log = gr.Textbox( label="Conversion Log", lines=25, # 少し高さを増やす interactive=False, autoscroll=True, max_lines=1500 # ログが多くなる可能性を考慮 ) # ボタンクリック時のアクション設定 convert_button.click( convert_safetensors_to_onnx_gradio, inputs=[safetensors_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アプリを起動 # share=True にするとパブリックリンクが生成される（HF Spacesでは不要） # queue() を使うと複数ユーザーのリクエストを処理しやすくなる demo.queue().launch()