Update app.py

app.py

@@ -35,8 +35,8 @@ def stream_process_output(process, log_list):
                 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)
@@ -44,96 +44,130 @@ def stream_process_output(process, log_list):
         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)):
+def convert_safetensors_to_onnx_gradio(
+    safetensors_file_obj,
+    config_file_obj,
+    style_vectors_file_obj
+    ): # gr.Progress は削除
     """
-    アップロードされたSafetensorsモデルをONNXに変換し、結果をダウンロード可能にする。
+    アップロードされた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: No safetensors file uploaded. Please upload a .safetensors file.")
-        # エラーメッセージを表示し、ダウンロードは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 パスの確認
+    # --- 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.")
+        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}")
+            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 # 拡張子を除いた部分
 
-            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"   - Created model directory: {model_dir_in_temp.relative_to(assets_root)}")
+            yield "\n".join(current_log), None
 
-            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}")
+            # --- 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)
-            paths_config = {"dataset_root": str(dataset_root.resolve()), "assets_root": str(assets_root.resolve())}
+            # 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}")
+            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}")
+            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}'.")
+                 current_log.append(f"❌ Error: convert_onnx.py not found at '{convert_script}'. Check repository setup.")
                  yield "\n".join(current_log), None
                  return # tryブロックを抜ける
 
@@ -142,40 +176,35 @@ def convert_safetensors_to_onnx_gradio(safetensors_file_obj, progress=gr.Progres
                 python_executable,
                 str(convert_script.resolve()),
                 "--model",
-                str(temp_safetensors_path.resolve()) # 一時ディレクトリ内のモデルパス
+                str(temp_safetensors_path.resolve()) # 一時ディレクトリ内の .safetensors パス
             ]
-            current_log.append(f"\nRunning command: {' '.join(command)}")
+            current_log.append(f"\n   Running 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 = []
+            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でGradio UIを更新 (現在のログ + プロセスからのログ)
                 yield "\n".join(current_log + process_output_lines), None
-                time.sleep(0.2) # 更新頻度
+                time.sleep(0.3) # 更新頻度
 
-            # スレッド終了待ち
+            # スレッド終了待ちとプロセス終了待ち
             thread.join()
-            # プロセス終了待ち (タイムアウト設定)
             try:
-                 process.wait(timeout=900) # 15分タイムアウト
+                 process.wait(timeout=12000) # 20分タイムアウト (モデルサイズにより調整)
             except subprocess.TimeoutExpired:
                  current_log.extend(process_output_lines) # ここまでのログを追加
-                 current_log.append("\nError: Conversion process timed out after 15 minutes.")
-                 process.kill() # タイムアウトしたらプロセスを強制終了
+                 current_log.append("\n❌ Error: Conversion process timed out after 20 minutes.")
+                 process.kill()
                  yield "\n".join(current_log), None
                  return # tryブロックを抜ける
 
@@ -187,102 +216,119 @@ def convert_safetensors_to_onnx_gradio(safetensors_file_obj, progress=gr.Progres
                  processed_stderr = []
                  for line in final_stderr.strip().split('\n'):
                      processed_line = f"stderr: {line.strip()}"
-                     if "warning" not in line.lower():
+                     if "warning" not in line.lower() and line.strip(): # 空行と警告以外
                          processed_line = f"ERROR (stderr): {line.strip()}"
                      processed_stderr.append(processed_line)
-                 if processed_stderr:
-                      process_output_lines.append("--- stderr ---")
+                 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("--------------")
+                      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.")
+                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}")
+                    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) # ダウンロード用ファイルパスを設定
+                    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}'. Please check the logs.")
+                    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').")
+                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
             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"\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"\nAn unexpected error occurred: {e}")
+        current_log.append(f"\n❌ An 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
+        # 最後の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 `.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."
+        "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="1. Upload Safetensors Model",
+                label="Safetensors Model (.safetensors)",
                 file_types=[".safetensors"],
-                # file_count="single" (default)
             )
-            convert_button = gr.Button("2. Convert to ONNX", variant="primary")
+            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="3. Download ONNX Model",
+                label="ONNX Model (.onnx)",
                 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."
+                "**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=25, # 少し高さを増やす
+                lines=30, # 高さをさらに増やす
                 interactive=False,
                 autoscroll=True,
-                max_lines=1500 # ログが多くなる可能性を考慮
+                max_lines=2000 # ログが増える可能性
             )
 
     # ボタンクリック時のアクション設定
     convert_button.click(
         convert_safetensors_to_onnx_gradio,
-        inputs=[safetensors_upload],
+        inputs=[safetensors_upload, config_upload, style_vectors_upload],
         outputs=[output_log, onnx_download] # ログとダウンロードファイルの2つを出力
     )
 
@@ -295,6 +341,4 @@ if __name__ == "__main__":
     print(f"Output directory: {OUTPUT_DIR.resolve()}")
 
     # Gradioアプリを起動
-    # share=True にするとパブリックリンクが生成される（HF Spacesでは不要）
-    # queue() を使うと複数ユーザーのリクエストを処理しやすくなる
     demo.queue().launch()