kill gpu

app.py

@@ -7,24 +7,30 @@ import spaces
 import os
 import time
 import logging
+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
 from huggingface_hub import snapshot_download
 
-
 # -----------------------------
-# Logging
+# Logging 配置
 # -----------------------------
-logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")
+logging.basicConfig(
+    level=logging.INFO, 
+    format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
+    datefmt="%H:%M:%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)
@@ -32,63 +38,46 @@ 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)
 
-USE_AMP = os.environ.get("USE_AMP", "0") == "1"
-AMP_DTYPE = os.environ.get("AMP_DTYPE", "bf16")  # "bf16" or "fp16"
-
-# ZeroGPU：缓存 CPU pipeline（不要缓存 CUDA Tensor）
-# cache: key -> (model_cpu, scheduler, target_sr)
+# ---------------------------------------------------------
+# 全局缓存
+# ---------------------------------------------------------
+# 存储 (model_cpu, scheduler, target_sr)
+# 警告：此缓存中的 model 必须始终保持在 "cpu" 设备上！
 _PIPELINE_CACHE: Dict[str, Tuple[object, object, int]] = {}
-# cache: key -> (repo_root, qwen_root)
 _MODEL_DIR_CACHE: Dict[str, Tuple[Path, Path]] = {}
 
 
 # ---------------------------------------------------------
-# 下载 repo（只下载一次；huggingface_hub 自带缓存）
+# 辅助函数
 # ---------------------------------------------------------
 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_id=MMEDIT_REPO_ID, revision=MMEDIT_REVISION, 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 模型必须
+        repo_id=QWEN_REPO_ID, revision=QWEN_REVISION, token=HF_TOKEN
     )
-    qwen_root = Path(qwen_root).resolve()
-
-    _MODEL_DIR_CACHE[cache_key] = (repo_root, qwen_root)
-    return repo_root, qwen_root
-
+    
+    res = (Path(repo_root).resolve(), Path(qwen_root).resolve())
+    _MODEL_DIR_CACHE[cache_key] = res
+    return res
 
-# ---------------------------------------------------------
-# 你的音频加载（按你要求：orig -> 16k -> target_sr）
-# ---------------------------------------------------------
-def load_and_process_audio(audio_path: str, target_sr: int):
-    # 延迟导入（避免启动阶段触发 CUDA 初始化）
-    import torch
+def load_and_process_audio(audio_path: str, target_sr: int) -> torch.Tensor:
     import torchaudio
-    import librosa
-
+    
     path = Path(audio_path)
     if not path.exists():
         raise FileNotFoundError(f"Audio file not found: {audio_path}")
@@ -96,91 +85,34 @@ def load_and_process_audio(audio_path: str, target_sr: int):
     waveform, orig_sr = torchaudio.load(str(path))  # (C, T)
 
     # Convert to mono
-    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()
-
-        # 1) 先到 16k
-        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
+    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
         else:
             orig_sr_mid = int(orig_sr)
-
-        # 2) 再到 target_sr（如 24k）
+        
         if int(target_sr) != orig_sr_mid:
-            waveform_np = librosa.resample(waveform_np, orig_sr=orig_sr_mid, target_sr=int(target_sr))
-
-        waveform = torch.from_numpy(waveform_np)
-
+            wav_np = librosa.resample(wav_np, orig_sr=orig_sr_mid, target_sr=int(target_sr))
+            
+        waveform = torch.from_numpy(wav_np)
+        
     return waveform
 
-
-# ---------------------------------------------------------
-# 校验 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'}")
-
-
-# ---------------------------------------------------------
-# 适配 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():
-            raise FileNotFoundError(
-                f"VAE ckpt not found after patch:\n"
-                f"  original: {vae_ckpt}\n"
-                f"  patched : {vae_path}\n"
-                f"Repo root: {repo_root}\n"
-                f"Expected:  {repo_root/'vae'/'*.ckpt'}"
-            )
-
-    # ---- 2) Qwen2-Audio model_path ----
-    exp_cfg["model"]["content_encoder"]["text_encoder"]["model_path"] = str(qwen_root)
-
-
-# ---------------------------------------------------------
-# Scheduler（与你 exp_cfg.model.noise_scheduler_name 对齐）
-# 带 fallback：避免 404
-# ---------------------------------------------------------
-def build_scheduler(exp_cfg: Dict[str, Any]):
+def build_scheduler(exp_cfg):
     import diffusers.schedulers as noise_schedulers
-
     name = exp_cfg["model"].get("noise_scheduler_name", "stabilityai/stable-diffusion-2-1")
     try:
-        scheduler = noise_schedulers.DDIMScheduler.from_pretrained(name, subfolder="scheduler", token=HF_TOKEN)
-        return scheduler
+        return noise_schedulers.DDIMScheduler.from_pretrained(name, subfolder="scheduler", token=HF_TOKEN)
     except Exception as e:
-        logger.warning(f"DDIMScheduler.from_pretrained failed for '{name}', fallback. err={e}")
+        logger.warning(f"Scheduler init failed: {e}, using fallback.")
         return noise_schedulers.DDIMScheduler(
             num_train_timesteps=1000,
             beta_start=0.00085,
@@ -191,73 +123,67 @@ def build_scheduler(exp_cfg: Dict[str, Any]):
             steps_offset=1,
         )
 
-
-def amp_autocast(device):
-    import torch
-
-    if not USE_AMP:
-        return torch.autocast("cuda", enabled=False)
-
-    if device.type != "cuda":
-        return torch.autocast("cpu", enabled=False)
-
-    dtype = torch.bfloat16 if AMP_DTYPE.lower() == "bf16" else torch.float16
-    return torch.autocast("cuda", dtype=dtype, enabled=True)
-
-
-# ---------------------------------------------------------
-# 冷启动：load+cache pipeline（缓存 CPU 上的 model）
-# ---------------------------------------------------------
 def load_pipeline_cpu() -> Tuple[object, object, int]:
-    # 延迟导入（避免启动阶段触发 CUDA 初始化）
-    import torch
+    """加载模型到 RAM（CPU），并建立全局缓存"""
     import hydra
     from omegaconf import OmegaConf
     from safetensors.torch import load_file
-
-    # 你的项目依赖也延迟导入
-    from models.common import LoadPretrainedBase
     from utils.config import register_omegaconf_resolvers
+    from models.common import LoadPretrainedBase
 
-    register_omegaconf_resolvers()
+    # 注册 omegaconf
+    try:
+        register_omegaconf_resolvers()
+    except Exception:
+        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]
 
     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)
+    
+    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)
 
-    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')}")
+    # Fix Qwen path
+    exp_cfg["model"]["content_encoder"]["text_encoder"]["model_path"] = str(qwen_root)
 
+    logger.info("Instantiating model architecture...")
     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)
-
-    # ZeroGPU：缓存 CPU 版
-    model = model.to(torch.device("cpu")).eval()
-
+    
+    # 关键：确保初始状态在 CPU
+    model = model.to("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("✅ Model loaded and cached in CPU RAM.")
     return model, scheduler, target_sr
 
 
 # ---------------------------------------------------------
-# 推理：audio + caption -> edited audio
-# ZeroGPU：必须用 @spaces.GPU
+# ZeroGPU 推理函数
 # ---------------------------------------------------------
 @spaces.GPU
 def run_edit(
@@ -268,148 +194,147 @@ 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, "Error: caption is empty."
+    if not caption: return None, "Please enter an instruction caption."
 
-    # 2. 获取缓存模型
-    # 注意：此时 model_cpu 在 CPU 上
+    # 1. 获取 CPU 上的模型引用
     model_cpu, scheduler, target_sr = load_pipeline_cpu()
+    
+    # 2. 准备设备 - 强制使用 float16
+    device = torch.device("cuda")
+    dtype = torch.float16  # <--- 强制 FP16
+    
+    logger.info(f"🚀 [GPU Task Start] Device: {device}, Precision: {dtype}")
+
+    # 用于 finally 清理的变量
+    model_on_gpu = None 
+    wav_on_gpu = None
 
-    # 使用 try-finally 确保无论是否出错，最后都把模型搬回 CPU
-    # 使用 try-except 确保捕获所有推理错误，打印日志
     try:
-        # --- 检查 GPU ---
+        # --- GPU 环境检查 ---
         if not torch.cuda.is_available():
-            return None, "Error: ZeroGPU did not allocate CUDA."
+            raise RuntimeError("ZeroGPU assigned but CUDA unavailable.")
 
-        device = torch.device("cuda")
-        logger.info(f"[GPU] Assigned device: {device}")
-
-        # --- 关键修改：模型上 GPU ---
-        # model_cpu.to(device) 是原位操作！会修改全局缓存！
-        # 所以必须在 finally 里搬回去，或者在这里使用深拷贝（深拷贝太慢，建议搬回去）
+        # --- 3. 模型搬运 (CPU -> GPU) ---
+        # 显式清理，为大模型腾出完整空间
+        gc.collect()
+        torch.cuda.empty_cache()
         
-        model = model_cpu.to(device).eval()
-        logger.info("Moving model to GPU for inference...")
-        # --- 数据预处理 ---
-        seed = int(seed)
-        torch.manual_seed(seed)
-        np.random.seed(seed)
+        logger.info("Moving model to GPU...")
         
-        # 加载音频并转到 GPU
-        wav = load_and_process_audio(audio_file, target_sr=target_sr).to(device)
+        # ⚠️ 核心逻辑：这里虽然用了 to(device)，这会修改 model_cpu 的设备属性
+        # 所以我们在 finally 块中必须将其搬回 CPU，否则下次运行会因为设备失效而崩溃
+        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)
+        
         batch = {
             "audio_id": [Path(audio_file).stem],
-            "content": [{"audio": wav, "caption": caption}],
+            "content": [{"audio": wav_on_gpu, "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
         }
-        kwargs.update(batch)
 
-        # --- 推理 ---
+        # --- 5. 推理 ---
+        logger.info("Starting inference...")
         t0 = time.time()
+        
         with torch.no_grad():
-            with amp_autocast(device):
-                # 这里的报错现在能被捕获了
-                out = model.inference(scheduler=scheduler, **kwargs)
+            # 使用 float16
+            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
         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"Success | Time: {dt:.2f}s | Seed: {seed}"
 
     except Exception as e:
-        # 这里会打印完整的堆栈信息，让你看到真正的报错原因
-        logger.exception("Error during inference")
-        return None, f"Runtime Error: {str(e)}"
+        # 🔥 捕捉所有错误，防止 spaces 吞掉报错，打印完整堆栈
+        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:
-        # --- 关键修改：清理现场 ---
-        # 无论 try 里面是否成功，这里都会执行
-        # 必须把模型搬回 CPU，否则全局缓存 _PIPELINE_CACHE 将指向损坏的 CUDA 地址
-        if 'model_cpu' in locals() and model_cpu is not None:
-            logger.info("Moving model back to CPU to preserve cache integrity...")
-            model_cpu.to("cpu") 
+        # --- 7. 关键：现场恢复 ---
+        # 无论成功还是失败，必须把模型搬回 CPU，否则全局缓存 _PIPELINE_CACHE 将指向已释放的显存
+        logger.info("♻️ Cleaning up resources...")
+        try:
+            # 只要 model_cpu 还在，就强制搬回 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}")
+
+        # 删除局部引用
+        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 启动
 # ---------------------------------------------------------
 def build_demo():
-    with gr.Blocks(title="MMEdit (ZeroGPU)") as demo:
-        gr.Markdown("# MMEdit ZeroGPU（audio + caption → edited audio）")
-
+    with gr.Blocks(title="MMEdit ZeroGPU") as demo:
+        gr.Markdown("## MMEdit")
+        gr.Markdown("ZeroGPU environment detected. Resources are allocated dynamically.")
+        
         with gr.Row():
             with gr.Column():
                 audio_in = gr.Audio(label="Input Audio", type="filepath")
-                caption = gr.Textbox(label="Caption (Edit Instruction)", lines=3)
-
-                # 注意：Space 不建议推大 wav；你可以换成更小的 demo wav
-                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")
+                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")
 
                 run_btn = gr.Button("Run Editing", variant="primary")
 
             with gr.Column():
-                audio_out = gr.Audio(label="Edited Audio", type="filepath")
-                status = gr.Textbox(label="Status")
+                audio_out = gr.Audio(label="Result", type="filepath")
+                status = gr.Textbox(label="Status Logs")
 
         run_btn.click(
-            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"
+            run_edit,
+            inputs=[audio_in, caption, steps, cfg, rescale, seed],
+            outputs=[audio_out, status]
         )
-
     return demo
 
 
 if __name__ == "__main__":
     demo = build_demo()
-    port = int(os.environ.get("PORT", "7860"))
+    # 兼容性设置：去掉 ssr_mode，让 Gradio 自动处理
+    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,
-        ssr_mode=False,
-    )
+        share=False
+    )