Improve trainer UI and add GPT dependency

app.py

@@ -14,6 +14,7 @@ import os
 import shutil
 import subprocess
 import sys
+from datetime import datetime
 from pathlib import Path
 
 import gradio as gr
@@ -142,14 +143,65 @@ def latest_file(directory: Path, suffix: str):
     return str(files[-1]) if files else None
 
 
+def list_files(directory: Path, suffix: str):
+    return sorted(directory.glob(f"*{suffix}"), key=lambda item: item.stat().st_mtime, reverse=True)
+
+
+def format_size(size_bytes: int):
+    value = float(size_bytes)
+    units = ["B", "KB", "MB", "GB"]
+    for unit in units:
+        if value < 1024 or unit == units[-1]:
+            if unit == "B":
+                return f"{int(value)} {unit}"
+            return f"{value:.1f} {unit}"
+        value /= 1024
+    return f"{int(size_bytes)} B"
+
+
+def format_mtime(path: Path):
+    return datetime.fromtimestamp(path.stat().st_mtime).strftime("%Y-%m-%d %H:%M:%S")
+
+
+def directory_overview():
+    return "\n".join(
+        [
+            f"工作目录: {WORK_DIR}",
+            f"数据集目录: {DATASET_DIR}",
+            f"日志目录: {EXP_DIR}",
+            f"SoVITS 导出目录: {SOVITS_OUTPUT_DIR}",
+            f"GPT 导出目录: {GPT_OUTPUT_DIR}",
+        ]
+    )
+
+
+def artifact_lines(label: str, files):
+    if not files:
+        return [f"{label}: 暂无"]
+    lines = [f"{label}: 共 {len(files)} 个"]
+    for index, path in enumerate(files[:10], start=1):
+        lines.append(f"{index}. {path.name} | {format_size(path.stat().st_size)} | {format_mtime(path)}")
+    return lines
+
+
 def artifacts_summary():
-    sovits = latest_file(SOVITS_OUTPUT_DIR, ".pth")
-    gpt = latest_file(GPT_OUTPUT_DIR, ".ckpt")
+    sovits_files = list_files(SOVITS_OUTPUT_DIR, ".pth")
+    gpt_files = list_files(GPT_OUTPUT_DIR, ".ckpt")
+    sovits = str(sovits_files[0]) if sovits_files else None
+    gpt = str(gpt_files[0]) if gpt_files else None
     lines = [
-        f"SoVITS: {sovits or '暂无'}",
-        f"GPT: {gpt or '暂无'}",
+        "训练结果总览",
+        *artifact_lines("SoVITS", sovits_files),
+        *artifact_lines("GPT", gpt_files),
     ]
-    return "\n".join(lines), sovits, gpt
+    return (
+        "\n".join(lines),
+        directory_overview(),
+        sovits,
+        gpt,
+        [str(item) for item in sovits_files] or None,
+        [str(item) for item in gpt_files] or None,
+    )
 
 
 def dataset_prepared():
@@ -509,6 +561,10 @@ def create_ui():
             "# 🎤 GPT-SoVITS 训练器 — 达妮娅语音\n"
             "这个 Space 按当前 GPT-SoVITS 训练链路执行。先拿到 SoVITS 权重，再按需继续训练 GPT。"
         )
+        gr.Markdown(
+            "打开页面会自动加载当前输出。训练完成后，最新模型会直接出现在下载框里，"
+            "下面的“输出与目录”也会自动刷新，避免找不到文件。"
+        )
 
         with gr.Row():
             with gr.Column(scale=1):
@@ -531,37 +587,50 @@ def create_ui():
                 sovits_save_every = gr.Slider(1, 5, value=1, step=1, label="每隔多少轮导出")
                 sovits_lr = gr.Slider(1e-5, 5e-4, value=1e-4, step=1e-5, label="学习率")
                 sovits_btn = gr.Button("开始 SoVITS 训练", variant="primary", size="lg")
-                sovits_log = gr.Textbox(label="SoVITS 训练日志", lines=18, interactive=False, autoscroll=True)
-                sovits_file = gr.File(label="最新 SoVITS 权重", interactive=False)
 
                 gr.Markdown("### 5. GPT 训练（可选）")
                 gpt_epochs = gr.Slider(1, 10, value=1, step=1, label="训练轮数")
                 gpt_batch = gr.Slider(1, 4, value=1, step=1, label="批次大小")
                 gpt_save_every = gr.Slider(1, 5, value=1, step=1, label="每隔多少轮导出")
                 gpt_btn = gr.Button("开始 GPT 训练", variant="secondary")
-                gpt_log = gr.Textbox(label="GPT 训练日志", lines=14, interactive=False, autoscroll=True)
-                gpt_file = gr.File(label="最新 GPT 权重", interactive=False)
 
-                gr.Markdown("### 6. 当前输出")
-                refresh_btn = gr.Button("刷新最新权重")
-                refresh_text = gr.Textbox(label="输出摘要", lines=3, interactive=False)
-                refresh_sovits = gr.File(label="SoVITS 输出", interactive=False)
-                refresh_gpt = gr.File(label="GPT 输出", interactive=False)
+        gr.Markdown("### 6. SoVITS 实时日志与下载")
+        sovits_log = gr.Textbox(label="SoVITS 训练日志", lines=22, interactive=False, autoscroll=True)
+        sovits_file = gr.File(label="最新 SoVITS 权重", interactive=False)
+
+        gr.Markdown("### 7. GPT 实时日志与下载")
+        gpt_log = gr.Textbox(label="GPT 训练日志", lines=22, interactive=False, autoscroll=True)
+        gpt_file = gr.File(label="最新 GPT 权重", interactive=False)
+
+        gr.Markdown("### 8. 输出与目录")
+        refresh_btn = gr.Button("刷新输出与目录", variant="secondary")
+        refresh_text = gr.Textbox(label="模型列表与状态", lines=14, interactive=False, autoscroll=True)
+        output_dirs = gr.Textbox(label="工作目录与输出目录", lines=6, interactive=False)
+        with gr.Row():
+            refresh_sovits = gr.File(label="最新 SoVITS 输出", interactive=False)
+            refresh_gpt = gr.File(label="最新 GPT 输出", interactive=False)
+        with gr.Row():
+            all_sovits = gr.File(label="全部 SoVITS 文件", interactive=False, file_count="multiple")
+            all_gpt = gr.File(label="全部 GPT 文件", interactive=False, file_count="multiple")
 
         env_btn.click(check_environment, outputs=env_out)
         dataset_btn.click(download_dataset, outputs=dataset_out)
         prep_btn.click(prepare_data, outputs=prep_out)
-        sovits_btn.click(
+        refresh_outputs_targets = [refresh_text, output_dirs, refresh_sovits, refresh_gpt, all_sovits, all_gpt]
+        sovits_event = sovits_btn.click(
             start_training,
             inputs=[sovits_epochs, sovits_batch, sovits_save_every, sovits_lr],
             outputs=[sovits_log, sovits_file],
         )
-        gpt_btn.click(
+        gpt_event = gpt_btn.click(
             start_gpt_training,
             inputs=[gpt_epochs, gpt_batch, gpt_save_every],
             outputs=[gpt_log, gpt_file],
         )
-        refresh_btn.click(refresh_outputs, outputs=[refresh_text, refresh_sovits, refresh_gpt])
+        refresh_btn.click(refresh_outputs, outputs=refresh_outputs_targets)
+        sovits_event.then(refresh_outputs, outputs=refresh_outputs_targets)
+        gpt_event.then(refresh_outputs, outputs=refresh_outputs_targets)
+        demo.load(refresh_outputs, outputs=refresh_outputs_targets)
 
     return demo
 

requirements.txt

@@ -10,6 +10,7 @@ torchaudio>=2.0.0
 pytorch-lightning>=2.4
 torchmetrics<=1.5
 tensorboard
+matplotlib
 transformers>=4.43,<=4.50
 sentencepiece>=0.1.99
 accelerate>=0.20.0