fix load

app.py

@@ -9,7 +9,7 @@ import time
 import logging
 from pathlib import Path
 from typing import Tuple, Optional, Dict, Any
-
+import gc
 import gradio as gr
 import numpy as np
 import soundfile as sf
@@ -208,156 +208,197 @@ def amp_autocast(device):
 # ---------------------------------------------------------
 # 冷启动：load+cache pipeline（缓存 CPU 上的 model）
 # ---------------------------------------------------------
-def load_pipeline_cpu() -> Tuple[object, object, int]:
-    # 延迟导入（避免启动阶段触发 CUDA 初始化）
+# 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
+
+#     register_omegaconf_resolvers()
+
+#     cache_key = f"{MMEDIT_REPO_ID}@{MMEDIT_REVISION}::{QWEN_REPO_ID}@{QWEN_REVISION}"
+#     if cache_key in _PIPELINE_CACHE:
+#         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}")
+
+#     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(f"patched pretrained_ckpt = {exp_cfg['model']['autoencoder'].get('pretrained_ckpt')}")
+#     logger.info(f"patched qwen model_path = {exp_cfg['model']['content_encoder']['text_encoder'].get('model_path')}")
+
+#     model: LoadPretrainedBase = hydra.utils.instantiate(exp_cfg["model"], _convert_="all")
+
+#     ckpt_path = repo_root / "model.safetensors"
+#     sd = load_file(str(ckpt_path))
+#     model.load_pretrained(sd)
+#     logger.info(f"Model loaded from safetensors: {ckpt_path}")
+#     # ZeroGPU：缓存 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.")
+#     return model, scheduler, target_sr
+def load_pipeline_cpu():
+    # 延迟导入
     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
-
-    register_omegaconf_resolvers()
+    
+    # 尝试导入项目模块
+    try:
+        from utils.config import register_omegaconf_resolvers
+        register_omegaconf_resolvers()
+    except: pass
 
     cache_key = f"{MMEDIT_REPO_ID}@{MMEDIT_REVISION}::{QWEN_REPO_ID}@{QWEN_REVISION}"
-    if cache_key in _PIPELINE_CACHE:
-        return _PIPELINE_CACHE[cache_key]
+    if cache_key in _PIPELINE_CACHE: 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}")
-
-    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(f"patched pretrained_ckpt = {exp_cfg['model']['autoencoder'].get('pretrained_ckpt')}")
-    logger.info(f"patched qwen model_path = {exp_cfg['model']['content_encoder']['text_encoder'].get('model_path')}")
+    
+    # 加载 Config
+    exp_cfg = OmegaConf.to_container(OmegaConf.load(repo_root / "config.yaml"), resolve=True)
+    
+    # 路径修复
+    vae_ckpt = exp_cfg["model"]["autoencoder"].get("pretrained_ckpt", "")
+    if vae_ckpt:
+        potential_paths = [repo_root / "vae" / Path(vae_ckpt).name, repo_root / Path(vae_ckpt).name]
+        for p in potential_paths:
+            if p.exists():
+                exp_cfg["model"]["autoencoder"]["pretrained_ckpt"] = str(p)
+                break
+                
+    exp_cfg["model"]["content_encoder"]["text_encoder"]["model_path"] = str(qwen_root)
 
-    model: LoadPretrainedBase = hydra.utils.instantiate(exp_cfg["model"], _convert_="all")
+    logger.info("Instantiating model...")
+    model = hydra.utils.instantiate(exp_cfg["model"], _convert_="all")
+    
+    # 加载权重并立即释放 state_dict 内存
+    ckpt_path = str(repo_root / "model.safetensors")
+    logger.info(f"Loading state_dict from {ckpt_path}...")
+    sd = load_file(ckpt_path)
 
-    ckpt_path = repo_root / "model.safetensors"
-    sd = load_file(str(ckpt_path))
+    logger.info(f"Model loaded from safetensors: {ckpt_path}")
     model.load_pretrained(sd)
+    del sd  # <--- 关键：立即删除 state_dict 释放 20GB+ 内存
+    gc.collect() # <--- 关键：强制回收
 
-    # ZeroGPU：缓��� CPU 版
-    model = model.to(torch.device("cpu")).eval()
+    # 确保在 CPU
+    model = model.to("cpu").eval()
+    
+    # Scheduler
+    import diffusers.schedulers as noise_schedulers
+    try:
+        scheduler = noise_schedulers.DDIMScheduler.from_pretrained(
+            exp_cfg["model"].get("noise_scheduler_name", "stabilityai/stable-diffusion-2-1"), 
+            subfolder="scheduler", token=HF_TOKEN
+        )
+    except:
+        scheduler = noise_schedulers.DDIMScheduler(num_train_timesteps=1000)
 
-    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.")
     return model, scheduler, target_sr
 
-
 # ---------------------------------------------------------
 # 推理：audio + caption -> edited audio
 # ZeroGPU：必须用 @spaces.GPU
 # ---------------------------------------------------------
 # ---------------------------------------------------------
 @spaces.GPU
-def run_edit(
-    audio_file: str,
-    caption: str,
-    num_steps: int,
-    guidance_scale: float,
-    guidance_rescale: float,
-    seed: int,
-) -> Tuple[Optional[str], str]:
+def run_edit(audio_file, caption, num_steps, guidance_scale, guidance_rescale, seed):
     import torch
-    import gc
-
-    if not audio_file: return None, "Error: Upload audio first."
-    if not caption: return None, "Error: Input caption."
-
-    # 1. 获取 CPU 模型
-    model_cpu, scheduler, target_sr = load_pipeline_cpu()
-
-    # 2. 准备设备 (强制 float16 以防 OOM 和兼容问题)
-    device = torch.device("cuda")
-    dtype = torch.float16
     
-    logger.info(f"🚀 [GPU Start] Device: {device}, Dtype: {dtype}")
-
+    if not audio_file: return None, "Please upload audio."
+    if not caption: return None, "Please input caption."
+    
+    # 局部变量初始化，防 finally 报错
+    model_cpu = None
     model_on_gpu = None
     
     try:
+        # --- 1. 将加载过程放入 try 块保护 ---
+        logger.info("Loading pipeline (CPU)...")
+        model_cpu, scheduler, target_sr = load_pipeline_cpu()
+    
+        # --- 2. 准备 GPU 环境 ---
+        device = torch.device("cuda")
+        dtype = torch.float16 
+        
         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() 会改变 cpu 对象
-        # 我们必须在 finally 里搬回去！
+        
+        # 原位操作，finally 必须移回
         model_on_gpu = model_cpu.to(device, dtype=dtype)
-
-        # --- 4. 数据准备 ---
+        
+        # --- 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)
+        wav = load_and_process_audio(audio_file, target_sr).to(device, dtype=dtype)
         
         batch = {
             "audio_id": [Path(audio_file).stem],
             "content": [{"audio": wav, "caption": caption}],
             "task": ["audio_editing"],
-        }
-        
-        kwargs = {
             "num_steps": int(num_steps),
             "guidance_scale": float(guidance_scale),
             "guidance_rescale": float(guidance_rescale),
             "use_gt_duration": False,
-            "mask_time_aligned_content": False,
-            **batch
+            "mask_time_aligned_content": False
         }
-
+        
         # --- 5. 推理 ---
-        logger.info("Inference start...")
+        logger.info("Running inference...")
         t0 = time.time()
-        with torch.no_grad():
-            with torch.autocast("cuda", dtype=dtype):
-                out = model_on_gpu.inference(scheduler=scheduler, **kwargs)
-        dt = time.time() - t0
-        logger.info(f"✅ Inference done: {dt:.2f}s")
-
-        # --- 6. 结果保存 ---
+        with torch.no_grad(), torch.autocast("cuda", dtype=dtype):
+            out = model_on_gpu.inference(scheduler=scheduler, **batch)
+            
+        # --- 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 | {dt:.2f}s | Seed: {seed}"
+        
+        return str(out_path), f"Success | {time.time()-t0:.2f}s"
 
     except Exception as e:
-        # [关键] 打印完整堆栈，不再报 404
+        # 🔥 现在你可以看到真正的报错了！
         err_msg = traceback.format_exc()
         logger.error(f"❌ ERROR:\n{err_msg}")
-        return None, f"Runtime Error: {str(e)}\nCheck Logs."
-
+        return None, f"Error: {str(e)}\n(Check Logs for Traceback)"
+        
     finally:
-        # --- 7. [关键] 现场恢复 ---
-        logger.info("♻️ Restoring CPU state...")
+        # --- 7. 还原现场 ---
+        logger.info("Restoring CPU state...")
         try:
-            # 必须搬回 CPU，否则缓存中的指针指向已释放的显存，下次必崩
-            if 'model_cpu' in locals() and model_cpu is not None:
+            if 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: {e}")
+            logger.error(f"Restore failed: {e}")
 
-        # 清理显存
-        if 'model_on_gpu' in locals(): del model_on_gpu
+        if model_on_gpu is not None: del model_on_gpu
         torch.cuda.empty_cache()
         gc.collect()
-
 # ---------------------------------------------------------
 # UI
 # ---------------------------------------------------------