die

app.py

@@ -20,17 +20,14 @@ import librosa
 from huggingface_hub import snapshot_download
 
 # -----------------------------
-# Logging 配置
+# Logging
 # -----------------------------
-logging.basicConfig(
-    level=logging.INFO, 
-    format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
-    datefmt="%H:%M:%S"
-)
+logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")
 logger = logging.getLogger("mmedit_space")
 
+
 # ---------------------------------------------------------
-# 配置信息
+# HF Repo IDs
 # ---------------------------------------------------------
 MMEDIT_REPO_ID = os.environ.get("MMEDIT_REPO_ID", "CocoBro/MMEdit")
 MMEDIT_REVISION = os.environ.get("MMEDIT_REVISION", None)
@@ -38,44 +35,58 @@ 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"))
 OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
 
 # ---------------------------------------------------------
-# 全局缓存
+# 缓存定义
 # ---------------------------------------------------------
-# 存储 (model_cpu, scheduler, target_sr)
-# 警告：此缓存中的 model 必须始终保持在 "cpu" 设备上！
+# 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
 # ---------------------------------------------------------
 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}...")
+    logger.info(f"Downloading MMEdit repo: {MMEDIT_REPO_ID} (revision={MMEDIT_REVISION})")
     repo_root = snapshot_download(
-        repo_id=MMEDIT_REPO_ID, revision=MMEDIT_REVISION, token=HF_TOKEN
+        repo_id=MMEDIT_REPO_ID,
+        revision=MMEDIT_REVISION,
+        local_dir=None,
+        local_dir_use_symlinks=False,
+        token=HF_TOKEN,
     )
-    
-    logger.info(f"Downloading Qwen repo: {QWEN_REPO_ID}...")
+    repo_root = Path(repo_root).resolve()
+
+    logger.info(f"Downloading Qwen repo: {QWEN_REPO_ID} (revision={QWEN_REVISION})")
     qwen_root = snapshot_download(
-        repo_id=QWEN_REPO_ID, revision=QWEN_REVISION, token=HF_TOKEN
+        repo_id=QWEN_REPO_ID,
+        revision=QWEN_REVISION,
+        local_dir=None,
+        local_dir_use_symlinks=False,
+        token=HF_TOKEN,  # gated 模型必须
     )
-    
-    res = (Path(repo_root).resolve(), Path(qwen_root).resolve())
-    _MODEL_DIR_CACHE[cache_key] = res
-    return res
+    qwen_root = Path(qwen_root).resolve()
+
+    _MODEL_DIR_CACHE[cache_key] = (repo_root, qwen_root)
+    return repo_root, qwen_root
 
-def load_and_process_audio(audio_path: str, target_sr: int) -> torch.Tensor:
+
+# ---------------------------------------------------------
+# 2. 音频加载（保留你的逻辑，增强鲁棒性）
+# ---------------------------------------------------------
+def load_and_process_audio(audio_path: str, target_sr: int):
     import torchaudio
     
     path = Path(audio_path)
@@ -85,34 +96,87 @@ def load_and_process_audio(audio_path: str, target_sr: int) -> torch.Tensor:
     waveform, orig_sr = torchaudio.load(str(path))  # (C, T)
 
     # Convert to mono
-    if waveform.ndim > 1:
-        waveform = waveform.mean(dim=0)
-    
-    # Resample logic (robust method)
-    if int(orig_sr) != int(target_sr):
-        wav_np = waveform.cpu().numpy()
-        
-        # Intermediate resampling to 16k if needed (for better stability)
-        if int(orig_sr) != 16000:
-            wav_np = librosa.resample(wav_np, orig_sr=int(orig_sr), target_sr=16000)
-            orig_sr_mid = 16000
+    if waveform.ndim == 2:
+        waveform = waveform.mean(dim=0)  # (T,)
+    elif waveform.ndim > 2:
+        waveform = waveform.reshape(-1)
+
+    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)
+            orig_sr_mid = sr_mid
         else:
             orig_sr_mid = int(orig_sr)
-        
+
         if int(target_sr) != orig_sr_mid:
-            wav_np = librosa.resample(wav_np, orig_sr=orig_sr_mid, target_sr=int(target_sr))
-            
-        waveform = torch.from_numpy(wav_np)
-        
+            waveform_np = librosa.resample(waveform_np, orig_sr=orig_sr_mid, target_sr=int(target_sr))
+
+        waveform = torch.from_numpy(waveform_np)
+
     return waveform
 
-def build_scheduler(exp_cfg):
+
+# ---------------------------------------------------------
+# 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/ 开始截断
+        else:
+            if vae_ckpt.endswith(".ckpt") and "/" not in vae_ckpt:
+                vae_rel = f"vae/{vae_ckpt}"
+            else:
+                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.")
+
+    # ---- 2) Qwen2-Audio model_path ----
+    exp_cfg["model"]["content_encoder"]["text_encoder"]["model_path"] = str(qwen_root)
+
+
+# ---------------------------------------------------------
+# 5. Scheduler（保留你的逻辑）
+# ---------------------------------------------------------
+def build_scheduler(exp_cfg: Dict[str, Any]):
     import diffusers.schedulers as noise_schedulers
+
     name = exp_cfg["model"].get("noise_scheduler_name", "stabilityai/stable-diffusion-2-1")
     try:
-        return noise_schedulers.DDIMScheduler.from_pretrained(name, subfolder="scheduler", token=HF_TOKEN)
+        scheduler = noise_schedulers.DDIMScheduler.from_pretrained(name, subfolder="scheduler", token=HF_TOKEN)
+        return scheduler
     except Exception as e:
-        logger.warning(f"Scheduler init failed: {e}, using fallback.")
+        logger.warning(f"DDIMScheduler.from_pretrained failed for '{name}', fallback. err={e}")
         return noise_schedulers.DDIMScheduler(
             num_train_timesteps=1000,
             beta_start=0.00085,
@@ -123,15 +187,20 @@ def build_scheduler(exp_cfg):
             steps_offset=1,
         )
 
+
+# ---------------------------------------------------------
+# 6. 冷启动：Load Pipeline to CPU
+# ---------------------------------------------------------
 def load_pipeline_cpu() -> Tuple[object, object, int]:
-    """加载模型到 RAM（CPU），并建立全局缓存"""
+    # 延迟导入
+    import torch
     import hydra
     from omegaconf import OmegaConf
     from safetensors.torch import load_file
-    from utils.config import register_omegaconf_resolvers
+
     from models.common import LoadPretrainedBase
+    from utils.config import register_omegaconf_resolvers
 
-    # 注册 omegaconf
     try:
         register_omegaconf_resolvers()
     except Exception:
@@ -142,48 +211,36 @@ def load_pipeline_cpu() -> Tuple[object, object, int]:
         return _PIPELINE_CACHE[cache_key]
 
     repo_root, qwen_root = resolve_model_dirs()
-    
-    cfg_path = repo_root / "config.yaml"
-    exp_cfg = OmegaConf.to_container(OmegaConf.load(cfg_path), resolve=True)
-    
-    # --- Config Patching ---
-    # Fix VAE ckpt path
-    vae_ckpt = exp_cfg["model"]["autoencoder"].get("pretrained_ckpt", "")
-    if vae_ckpt:
-        # 简单暴力的路径修复：只要是 ckpt 就去 vae 目录下找
-        fname = Path(vae_ckpt).name
-        local_vae = repo_root / "vae" / fname
-        if local_vae.exists():
-            exp_cfg["model"]["autoencoder"]["pretrained_ckpt"] = str(local_vae)
-        else:
-            # 尝试直接在 repo_root 下找
-            if (repo_root / fname).exists():
-                exp_cfg["model"]["autoencoder"]["pretrained_ckpt"] = str(repo_root / fname)
+    assert_repo_layout(repo_root)
 
-    # Fix Qwen path
-    exp_cfg["model"]["content_encoder"]["text_encoder"]["model_path"] = str(qwen_root)
+    logger.info(f"repo_root = {repo_root}")
+    logger.info(f"qwen_root = {qwen_root}")
+
+    exp_cfg = OmegaConf.load(repo_root / "config.yaml")
+    exp_cfg = OmegaConf.to_container(exp_cfg, resolve=True)
 
-    logger.info("Instantiating model architecture...")
+    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")
-    
+
     ckpt_path = repo_root / "model.safetensors"
-    logger.info(f"Loading weights from {ckpt_path.name}...")
     sd = load_file(str(ckpt_path))
     model.load_pretrained(sd)
-    
-    # 关键：确保初始状态在 CPU
-    model = model.to("cpu").eval()
-    
+
+    # 关键：确保模型在 CPU 上，并且是 eval 模式
+    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("✅ Model loaded and cached in CPU RAM.")
+    logger.info("CPU pipeline loaded and cached.")
     return model, scheduler, target_sr
 
 
 # ---------------------------------------------------------
-# ZeroGPU 推理函数
+# 7. ZeroGPU 推理核心（修复版）
 # ---------------------------------------------------------
 @spaces.GPU
 def run_edit(
@@ -194,52 +251,59 @@ def run_edit(
     guidance_rescale: float,
     seed: int,
 ) -> 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."
     
-    if not audio_file: return None, "Please upload an audio file."
     caption = (caption or "").strip()
-    if not caption: return None, "Please enter an instruction caption."
+    if not caption:
+        return None, "Error: caption is empty."
 
-    # 1. 获取 CPU 上的模型引用
+    # 2. 获取缓存模型 (CPU)
     model_cpu, scheduler, target_sr = load_pipeline_cpu()
-    
-    # 2. 准备设备 - 强制使用 float16
+
+    # 强制使用 float16，兼容性最好
     device = torch.device("cuda")
-    dtype = torch.float16  # <--- 强制 FP16
-    
-    logger.info(f"🚀 [GPU Task Start] Device: {device}, Precision: {dtype}")
+    dtype = torch.float16
 
-    # 用于 finally 清理的变量
-    model_on_gpu = None 
+    logger.info(f"🚀 [GPU Task Start] Device: {device}, Dtype: {dtype}")
+
+    # 用于 finally 清理
+    model_on_gpu = None
     wav_on_gpu = None
 
     try:
-        # --- GPU 环境检查 ---
+        # --- 检查环境 ---
         if not torch.cuda.is_available():
-            raise RuntimeError("ZeroGPU assigned but CUDA unavailable.")
+            raise RuntimeError("ZeroGPU assigned but CUDA not found!")
 
         # --- 3. 模型搬运 (CPU -> GPU) ---
-        # 显式清理，为大模型腾出完整空间
         gc.collect()
         torch.cuda.empty_cache()
-        
+
         logger.info("Moving model to GPU...")
         
-        # ⚠️ 核心逻辑：这里虽然用了 to(device)，这会修改 model_cpu 的设备属性
-        # 所以我们在 finally 块中必须将其搬回 CPU，否则下次运行会因为设备失效而崩溃
+        # ⚠️ 关键点：这里 model_cpu.to(device) 是原位操作，
+        # 我们必须在 finally 里搬回去，才能保证全局缓存不坏。
+        # 同时做 dtype 转换以节省显存。
         model_on_gpu = model_cpu.to(device, dtype=dtype)
 
-        # --- 4. 数据准备 ---
-        torch.manual_seed(int(seed))
-        np.random.seed(int(seed))
-        
-        wav_on_gpu = load_and_process_audio(audio_file, target_sr).to(device, dtype=dtype)
+        # --- 4. 数据预处理 ---
+        seed = int(seed)
+        torch.manual_seed(seed)
+        np.random.seed(seed)
         
+        # 加载音频并转到 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}],
             "task": ["audio_editing"],
         }
-        
+
         kwargs = {
             "num_steps": int(num_steps),
             "guidance_scale": float(guidance_scale),
@@ -254,87 +318,102 @@ def run_edit(
         t0 = time.time()
         
         with torch.no_grad():
-            # 使用 float16
+            # 使用 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")
 
-        # --- 6. 保存结果 ---
-        # 立即 detach 并转回 CPU
+        # --- 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"Success | Time: {dt:.2f}s | Seed: {seed}"
+
+        return str(out_path), f"OK | time={dt:.2f}s | seed={seed}"
 
     except Exception as e:
-        # 🔥 捕捉所有错误，防止 spaces 吞掉报错，打印完整堆栈
+        # 🔥 打印完整堆栈，防止 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)"
 
     finally:
-        # --- 7. 关键：现场恢复 ---
-        # 无论成功还是失败，必须把模型搬回 CPU，否则全局缓存 _PIPELINE_CACHE 将指向已释放的显存
+        # --- 7. 关键：现场恢复（必须执行）---
         logger.info("♻️ Cleaning up resources...")
         try:
-            # 只要 model_cpu 还在，就强制搬回 CPU
+            # 必须把模型搬回 CPU，否则全局缓存 _PIPELINE_CACHE 指向已释放的显存
             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}")
 
-        # 删除局部引用
+        # 删除引用
         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 启动
+# UI (完全保留你的 Examples)
 # ---------------------------------------------------------
 def build_demo():
-    with gr.Blocks(title="MMEdit ZeroGPU") as demo:
-        gr.Markdown("## MMEdit")
-        gr.Markdown("ZeroGPU environment detected. Resources are allocated dynamically.")
-        
+    with gr.Blocks(title="MMEdit (ZeroGPU)") as demo:
+        gr.Markdown("# MMEdit ZeroGPU（audio + caption → edited audio）")
+
         with gr.Row():
             with gr.Column():
                 audio_in = gr.Audio(label="Input Audio", type="filepath")
-                caption = gr.Textbox(label="Editing Instruction", placeholder="e.g., Add rain sound in the background")
-                
-                with gr.Accordion("Advanced Settings", open=False):
-                    steps = gr.Slider(10, 100, 50, step=1, label="Steps")
-                    cfg = gr.Slider(1.0, 15.0, 5.0, step=0.5, label="Guidance Scale")
-                    rescale = gr.Slider(0.0, 1.0, 0.5, step=0.05, label="Guidance Rescale")
-                    seed = gr.Number(42, label="Seed")
+                caption = gr.Textbox(label="Caption (Edit Instruction)", lines=3)
+
+                # 恢复了你的 Examples
+                gr.Examples(
+                    label="example inputs",
+                    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")
 
                 run_btn = gr.Button("Run Editing", variant="primary")
 
             with gr.Column():
-                audio_out = gr.Audio(label="Result", type="filepath")
-                status = gr.Textbox(label="Status Logs")
+                audio_out = gr.Audio(label="Edited Audio", type="filepath")
+                status = gr.Textbox(label="Status")
 
         run_btn.click(
-            run_edit,
-            inputs=[audio_in, caption, steps, cfg, rescale, seed],
-            outputs=[audio_out, status]
+            fn=run_edit,
+            inputs=[audio_in, caption, num_steps, guidance_scale, guidance_rescale, seed],
+            outputs=[audio_out, status],
+        )
+
+        gr.Markdown(
+            "## 注意事项\n"
+            "1) ZeroGPU 首次点击会分配 GPU，可能稍慢。\n"
+            "2) 如果首次报 cuda 不可用，通常重试一次即可。\n"
         )
+
     return demo
 
 
 if __name__ == "__main__":
     demo = build_demo()
-    # 兼容性设置：去掉 ssr_mode，让 Gradio 自动处理
-    port = int(os.environ.get("PORT", 7860))
+    port = int(os.environ.get("PORT", "7860"))
     demo.queue().launch(
-        server_name="0.0.0.0", 
+        server_name="0.0.0.0",
         server_port=port,
-        share=False
+        share=False,
     )