why

app.py

@@ -7,18 +7,18 @@ import spaces
 import os
 import time
 import logging
-import traceback
-import gc
+import traceback  # [新增] 用于打印报错堆栈
+import gc         # [新增] 用于显存清理
 from pathlib import Path
 from typing import Tuple, Optional, Dict, Any
 
 import gradio as gr
 import numpy as np
 import soundfile as sf
-import torch
-import librosa
+# [修改] 移除了顶部的 hydra/models 导入，防止启动时触发 CUDA
 from huggingface_hub import snapshot_download
 
+
 # -----------------------------
 # Logging
 # -----------------------------
@@ -35,7 +35,6 @@ MMEDIT_REVISION = os.environ.get("MMEDIT_REVISION", None)
 QWEN_REPO_ID = os.environ.get("QWEN_REPO_ID", "Qwen/Qwen2-Audio-7B-Instruct")
 QWEN_REVISION = os.environ.get("QWEN_REVISION", None)
 
-# 如果 Qwen gated：Space 里把 HF_TOKEN 设为 Secret
 HF_TOKEN = os.environ.get("HF_TOKEN", None)
 
 OUTPUT_DIR = Path(os.environ.get("OUTPUT_DIR", "./outputs"))
@@ -45,37 +44,32 @@ OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
 # 缓存定义
 # ---------------------------------------------------------
 # cache: key -> (model_cpu, scheduler, target_sr)
-# 注意：model_cpu 必须始终在 CPU 上
 _PIPELINE_CACHE: Dict[str, Tuple[object, object, int]] = {}
 # cache: key -> (repo_root, qwen_root)
 _MODEL_DIR_CACHE: Dict[str, Tuple[Path, Path]] = {}
 
 
 # ---------------------------------------------------------
-# 1. 下载 repo
+# 下载 Repo
 # ---------------------------------------------------------
 def resolve_model_dirs() -> Tuple[Path, Path]:
     cache_key = f"{MMEDIT_REPO_ID}@{MMEDIT_REVISION}::{QWEN_REPO_ID}@{QWEN_REVISION}"
     if cache_key in _MODEL_DIR_CACHE:
         return _MODEL_DIR_CACHE[cache_key]
 
-    logger.info(f"Downloading MMEdit repo: {MMEDIT_REPO_ID} (revision={MMEDIT_REVISION})")
+    logger.info(f"Downloading MMEdit repo: {MMEDIT_REPO_ID}")
     repo_root = snapshot_download(
         repo_id=MMEDIT_REPO_ID,
         revision=MMEDIT_REVISION,
-        local_dir=None,
-        local_dir_use_symlinks=False,
         token=HF_TOKEN,
     )
     repo_root = Path(repo_root).resolve()
 
-    logger.info(f"Downloading Qwen repo: {QWEN_REPO_ID} (revision={QWEN_REVISION})")
+    logger.info(f"Downloading Qwen repo: {QWEN_REPO_ID}")
     qwen_root = snapshot_download(
         repo_id=QWEN_REPO_ID,
         revision=QWEN_REVISION,
-        local_dir=None,
-        local_dir_use_symlinks=False,
-        token=HF_TOKEN,  # gated 模型必须
+        token=HF_TOKEN,
     )
     qwen_root = Path(qwen_root).resolve()
 
@@ -84,27 +78,26 @@ def resolve_model_dirs() -> Tuple[Path, Path]:
 
 
 # ---------------------------------------------------------
-# 2. 音频加载（保留你的逻辑，增强鲁棒性）
+# 音频处理
 # ---------------------------------------------------------
 def load_and_process_audio(audio_path: str, target_sr: int):
+    # 延迟导入，防止干扰
+    import torch
     import torchaudio
-    
+    import librosa
+
     path = Path(audio_path)
     if not path.exists():
         raise FileNotFoundError(f"Audio file not found: {audio_path}")
 
-    waveform, orig_sr = torchaudio.load(str(path))  # (C, T)
+    waveform, orig_sr = torchaudio.load(str(path))
 
-    # Convert to mono
-    if waveform.ndim == 2:
-        waveform = waveform.mean(dim=0)  # (T,)
-    elif waveform.ndim > 2:
-        waveform = waveform.reshape(-1)
+    if waveform.ndim > 1:
+        waveform = waveform.mean(dim=0)
 
     if target_sr and int(target_sr) != int(orig_sr):
         waveform_np = waveform.cpu().numpy()
-
-        # 稳健的两步重采样逻辑
+        # 稳健的重采样逻辑
         sr_mid = 16000
         if int(orig_sr) != sr_mid:
             waveform_np = librosa.resample(waveform_np, orig_sr=int(orig_sr), target_sr=sr_mid)
@@ -121,52 +114,33 @@ def load_and_process_audio(audio_path: str, target_sr: int):
 
 
 # ---------------------------------------------------------
-# 3. 校验 repo 结构（保留你的逻辑）
-# ---------------------------------------------------------
-def assert_repo_layout(repo_root: Path) -> None:
-    must = [repo_root / "config.yaml", repo_root / "model.safetensors", repo_root / "vae"]
-    for p in must:
-        if not p.exists():
-            raise FileNotFoundError(f"Missing required path: {p}")
-
-    vae_files = list((repo_root / "vae").glob("*.ckpt"))
-    if len(vae_files) == 0:
-        raise FileNotFoundError(f"No .ckpt found under: {repo_root/'vae'}")
-
-
-# ---------------------------------------------------------
-# 4. 适配 config.yaml（保留你的逻辑）
+# 路径适配
 # ---------------------------------------------------------
 def patch_paths_in_exp_config(exp_cfg: Dict[str, Any], repo_root: Path, qwen_root: Path) -> None:
-    # ---- 1) VAE ckpt ----
     vae_ckpt = exp_cfg["model"]["autoencoder"].get("pretrained_ckpt", None)
     if vae_ckpt:
         vae_ckpt = str(vae_ckpt).replace("\\", "/")
-        idx = vae_ckpt.find("vae/")
-        if idx != -1:
-            vae_rel = vae_ckpt[idx:]  # 从 vae/ 开始截断
+        if "vae/" in vae_ckpt:
+            vae_rel = vae_ckpt[vae_ckpt.find("vae/"):]
+        elif vae_ckpt.endswith(".ckpt"):
+            vae_rel = f"vae/{vae_ckpt}" if "/" not in vae_ckpt else vae_ckpt
         else:
-            if vae_ckpt.endswith(".ckpt") and "/" not in vae_ckpt:
-                vae_rel = f"vae/{vae_ckpt}"
-            else:
-                vae_rel = vae_ckpt
+            vae_rel = vae_ckpt
 
         vae_path = (repo_root / vae_rel).resolve()
-        exp_cfg["model"]["autoencoder"]["pretrained_ckpt"] = str(vae_path)
-
+        # 鲁棒性检查：如果算出来的路径不存在，尝试在根目录找文件名
         if not vae_path.exists():
-            # Fallback check (鲁棒性增强)
-            if (repo_root / Path(vae_ckpt).name).exists():
-                 exp_cfg["model"]["autoencoder"]["pretrained_ckpt"] = str(repo_root / Path(vae_ckpt).name)
-            else:
-                logger.warning(f"VAE ckpt warning: {vae_path} not found. Model loading might fail.")
+             fallback = repo_root / Path(vae_ckpt).name
+             if fallback.exists():
+                 vae_path = fallback
+        
+        exp_cfg["model"]["autoencoder"]["pretrained_ckpt"] = str(vae_path)
 
-    # ---- 2) Qwen2-Audio model_path ----
     exp_cfg["model"]["content_encoder"]["text_encoder"]["model_path"] = str(qwen_root)
 
 
 # ---------------------------------------------------------
-# 5. Scheduler（保留你的逻辑）
+# Scheduler
 # ---------------------------------------------------------
 def build_scheduler(exp_cfg: Dict[str, Any]):
     import diffusers.schedulers as noise_schedulers
@@ -176,7 +150,7 @@ def build_scheduler(exp_cfg: Dict[str, Any]):
         scheduler = noise_schedulers.DDIMScheduler.from_pretrained(name, subfolder="scheduler", token=HF_TOKEN)
         return scheduler
     except Exception as e:
-        logger.warning(f"DDIMScheduler.from_pretrained failed for '{name}', fallback. err={e}")
+        logger.warning(f"Scheduler fallback: {e}")
         return noise_schedulers.DDIMScheduler(
             num_train_timesteps=1000,
             beta_start=0.00085,
@@ -189,20 +163,22 @@ def build_scheduler(exp_cfg: Dict[str, Any]):
 
 
 # ---------------------------------------------------------
-# 6. 冷启动：Load Pipeline to CPU
+# [核心] 冷启动：Load to CPU
 # ---------------------------------------------------------
 def load_pipeline_cpu() -> Tuple[object, object, int]:
-    # 延迟导入
+    # [修改] 所有的库都在这里导入，防止全局导入触发 CUDA 初始化
     import torch
     import hydra
     from omegaconf import OmegaConf
     from safetensors.torch import load_file
-
-    from models.common import LoadPretrainedBase
-    from utils.config import register_omegaconf_resolvers
-
+    
+    # 你的项目依赖
     try:
+        from utils.config import register_omegaconf_resolvers
+        from models.common import LoadPretrainedBase
         register_omegaconf_resolvers()
+    except ImportError:
+        logger.warning("Could not import project utils/models. Ensure they are in the python path.")
     except Exception:
         pass
 
@@ -211,36 +187,35 @@ def load_pipeline_cpu() -> Tuple[object, object, int]:
         return _PIPELINE_CACHE[cache_key]
 
     repo_root, qwen_root = resolve_model_dirs()
-    assert_repo_layout(repo_root)
-
-    logger.info(f"repo_root = {repo_root}")
-    logger.info(f"qwen_root = {qwen_root}")
-
+    
+    logger.info(f"repo_root: {repo_root}")
+    
     exp_cfg = OmegaConf.load(repo_root / "config.yaml")
     exp_cfg = OmegaConf.to_container(exp_cfg, resolve=True)
 
     patch_paths_in_exp_config(exp_cfg, repo_root, qwen_root)
 
     logger.info("Instantiating model...")
-    model: LoadPretrainedBase = hydra.utils.instantiate(exp_cfg["model"], _convert_="all")
+    model = hydra.utils.instantiate(exp_cfg["model"], _convert_="all")
 
     ckpt_path = repo_root / "model.safetensors"
+    logger.info(f"Loading weights: {ckpt_path}")
     sd = load_file(str(ckpt_path))
     model.load_pretrained(sd)
 
-    # 关键：确保模型在 CPU 上，并且是 eval 模式
+    # [修改] 确保加载到 CPU
     model = model.to(torch.device("cpu")).eval()
 
     scheduler = build_scheduler(exp_cfg)
     target_sr = int(exp_cfg.get("sample_rate", 24000))
 
     _PIPELINE_CACHE[cache_key] = (model, scheduler, target_sr)
-    logger.info("CPU pipeline loaded and cached.")
+    logger.info("CPU pipeline cached.")
     return model, scheduler, target_sr
 
 
 # ---------------------------------------------------------
-# 7. ZeroGPU 推理核心（修复版）
+# [核心] 推理函数 (ZeroGPU 适配版)
 # ---------------------------------------------------------
 @spaces.GPU
 def run_edit(
@@ -253,57 +228,45 @@ def run_edit(
 ) -> Tuple[Optional[str], str]:
     import torch
 
-    # 1. 基础检查
-    if audio_file is None or not Path(audio_file).exists():
-        return None, "Error: please upload an audio file."
-    
-    caption = (caption or "").strip()
-    if not caption:
-        return None, "Error: caption is empty."
+    if not audio_file: return None, "Error: Upload audio first."
+    if not caption: return None, "Error: Input caption."
 
-    # 2. 获取缓存模型 (CPU)
+    # 1. 获取 CPU 模型
     model_cpu, scheduler, target_sr = load_pipeline_cpu()
 
-    # 强制使用 float16，兼容性最好
+    # 2. 准备设备 (强制 float16 以防 OOM 和兼容问题)
     device = torch.device("cuda")
     dtype = torch.float16
+    
+    logger.info(f"🚀 [GPU Start] Device: {device}, Dtype: {dtype}")
 
-    logger.info(f"🚀 [GPU Task Start] Device: {device}, Dtype: {dtype}")
-
-    # 用于 finally 清理
     model_on_gpu = None
-    wav_on_gpu = None
-
+    
     try:
-        # --- 检查环境 ---
         if not torch.cuda.is_available():
             raise RuntimeError("ZeroGPU assigned but CUDA not found!")
 
-        # --- 3. 模型搬运 (CPU -> GPU) ---
+        # --- 3. 搬运模型 (CPU -> GPU) ---
         gc.collect()
         torch.cuda.empty_cache()
-
-        logger.info("Moving model to GPU...")
         
-        # ⚠️ 关键点：这里 model_cpu.to(device) 是原位操作，
-        # 我们必须在 finally 里搬回去，才能保证全局缓存不坏。
-        # 同时做 dtype 转换以节省显存。
+        logger.info("Moving model to GPU...")
+        # [关键] 原位操作警告：model_cpu.to() 会改变 cpu 对象
+        # 我们必须在 finally 里搬回去！
         model_on_gpu = model_cpu.to(device, dtype=dtype)
 
-        # --- 4. 数据预处理 ---
-        seed = int(seed)
-        torch.manual_seed(seed)
-        np.random.seed(seed)
+        # --- 4. 数据准备 ---
+        torch.manual_seed(int(seed))
+        np.random.seed(int(seed))
+        
+        wav = load_and_process_audio(audio_file, target_sr=target_sr).to(device, dtype=dtype)
         
-        # 加载音频并转到 GPU
-        wav_on_gpu = load_and_process_audio(audio_file, target_sr=target_sr).to(device, dtype=dtype)
-
         batch = {
             "audio_id": [Path(audio_file).stem],
-            "content": [{"audio": wav_on_gpu, "caption": caption}],
+            "content": [{"audio": wav, "caption": caption}],
             "task": ["audio_editing"],
         }
-
+        
         kwargs = {
             "num_steps": int(num_steps),
             "guidance_scale": float(guidance_scale),
@@ -314,84 +277,72 @@ def run_edit(
         }
 
         # --- 5. 推理 ---
-        logger.info("Starting inference...")
+        logger.info("Inference start...")
         t0 = time.time()
-        
         with torch.no_grad():
-            # 使用 float16 autocast
             with torch.autocast("cuda", dtype=dtype):
                 out = model_on_gpu.inference(scheduler=scheduler, **kwargs)
-        
         dt = time.time() - t0
-        logger.info(f"✅ Inference finished in {dt:.2f}s")
+        logger.info(f"✅ Inference done: {dt:.2f}s")
 
-        # --- 6. 后处理 ---
+        # --- 6. 结果保存 ---
         out_audio = out[0, 0].detach().float().cpu().numpy()
         out_path = OUTPUT_DIR / f"{Path(audio_file).stem}_edited.wav"
         sf.write(str(out_path), out_audio, samplerate=target_sr)
 
-        return str(out_path), f"OK | time={dt:.2f}s | seed={seed}"
+        return str(out_path), f"OK | {dt:.2f}s | Seed: {seed}"
 
     except Exception as e:
-        # 🔥 打印完整堆栈，防止 404 掩盖真实错误
+        # [关键] 打印完整堆栈，不再报 404
         err_msg = traceback.format_exc()
-        logger.error(f"❌ CRITICAL ERROR:\n{err_msg}")
-        return None, f"Runtime Error: {str(e)}\n(See logs for details)"
+        logger.error(f"❌ ERROR:\n{err_msg}")
+        return None, f"Runtime Error: {str(e)}\nCheck Logs."
 
     finally:
-        # --- 7. 关键：现场恢复（必须执行）---
-        logger.info("♻️ Cleaning up resources...")
+        # --- 7. [关键] 现场恢复 ---
+        logger.info("♻️ Restoring CPU state...")
         try:
-            # 必须把模型搬回 CPU，否则全局缓存 _PIPELINE_CACHE 指向已释放的显存
+            # 必须搬回 CPU，否则缓存中的指针指向已释放的显存，下次必崩
             if 'model_cpu' in locals() and model_cpu is not None:
                 model_cpu.to("cpu")
                 logger.info("Model restored to CPU.")
         except Exception as e:
-            logger.error(f"Failed to restore model to CPU: {e}")
+            logger.error(f"Failed to restore model: {e}")
 
-        # 删除引用
+        # 清理显存
         if 'model_on_gpu' in locals(): del model_on_gpu
-        if 'wav_on_gpu' in locals(): del wav_on_gpu
-
-        # 强制清理显存
         torch.cuda.empty_cache()
         gc.collect()
 
 
 # ---------------------------------------------------------
-# UI (完全保留你的 Examples)
+# UI
 # ---------------------------------------------------------
 def build_demo():
     with gr.Blocks(title="MMEdit (ZeroGPU)") as demo:
-        gr.Markdown("# MMEdit ZeroGPU（audio + caption → edited audio）")
+        gr.Markdown("# MMEdit ZeroGPU")
 
         with gr.Row():
             with gr.Column():
                 audio_in = gr.Audio(label="Input Audio", type="filepath")
-                caption = gr.Textbox(label="Caption (Edit Instruction)", lines=3)
+                caption = gr.Textbox(label="Caption", lines=3)
 
-                # 恢复了你的 Examples
                 gr.Examples(
-                    label="example inputs",
-                    examples=[
-                        ["./Ym8O802VvJes.wav", "Mix in dog barking around the middle."],
-                    ],
+                    label="Examples",
+                    examples=[["./Ym8O802VvJes.wav", "Mix in dog barking around the middle."]],
                     inputs=[audio_in, caption],
-                    cache_examples=False,
                 )
 
                 with gr.Row():
-                    num_steps = gr.Slider(1, 100, value=50, step=1, label="num_steps")
-                    guidance_scale = gr.Slider(1.0, 12.0, value=5.0, step=0.5, label="guidance_scale")
-
-                with gr.Row():
-                    guidance_rescale = gr.Slider(0.0, 1.0, value=0.5, step=0.05, label="guidance_rescale")
-                    seed = gr.Number(value=42, precision=0, label="seed")
+                    num_steps = gr.Slider(1, 100, value=50, step=1, label="Steps")
+                    guidance_scale = gr.Slider(1.0, 12.0, value=5.0, step=0.5, label="Guidance")
+                    rescale = gr.Slider(0.0, 1.0, 0.5, step=0.05, label="Rescale")
+                    seed = gr.Number(42, label="Seed")
 
-                run_btn = gr.Button("Run Editing", variant="primary")
+                run_btn = gr.Button("Run", variant="primary")
 
             with gr.Column():
-                audio_out = gr.Audio(label="Edited Audio", type="filepath")
+                audio_out = gr.Audio(label="Output", type="filepath")
                 status = gr.Textbox(label="Status")
 
         run_btn.click(
@@ -400,12 +351,6 @@ def build_demo():
             outputs=[audio_out, status],
         )
 
-        gr.Markdown(
-            "## 注意事项\n"
-            "1) ZeroGPU 首次点击会分配 GPU，可能稍慢。\n"
-            "2) 如果首次报 cuda 不可用，通常重试一次即可。\n"
-        )
-
     return demo