fix some bugs

app.py

@@ -1,6 +1,9 @@
 #!/usr/bin/env python3
 # -*- coding: utf-8 -*-
 
+# ZeroGPU 关键：必须最先导入
+import spaces
+
 import os
 import time
 import logging
@@ -10,34 +13,15 @@ from typing import Tuple, Optional, Dict, Any
 import gradio as gr
 import numpy as np
 import soundfile as sf
-import torch
-import torchaudio
-import librosa
-
-import hydra
-from omegaconf import OmegaConf
-from safetensors.torch import load_file
-import diffusers.schedulers as noise_schedulers
 from huggingface_hub import snapshot_download
 
-# ZeroGPU 关键：spaces
-import spaces
-
-from models.common import LoadPretrainedBase
-from utils.config import register_omegaconf_resolvers
-
 
 # -----------------------------
 # Logging
 # -----------------------------
-logging.basicConfig(
-    level=logging.INFO,
-    format="%(asctime)s - %(levelname)s - %(message)s"
-)
+logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")
 logger = logging.getLogger("mmedit_space")
 
-register_omegaconf_resolvers()
-
 
 # ---------------------------------------------------------
 # HF Repo IDs（按你的默认需求）
@@ -57,8 +41,10 @@ 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）
-_PIPELINE_CACHE: Dict[str, Tuple[LoadPretrainedBase, object, int]] = {}
+# ZeroGPU：缓存 CPU pipeline（不要缓存 CUDA Tensor）
+# cache: key -> (model_cpu, scheduler, target_sr)
+_PIPELINE_CACHE: Dict[str, Tuple[object, object, int]] = {}
+# cache: key -> (repo_root, qwen_root)
 _MODEL_DIR_CACHE: Dict[str, Tuple[Path, Path]] = {}
 
 
@@ -66,11 +52,6 @@ _MODEL_DIR_CACHE: Dict[str, Tuple[Path, Path]] = {}
 # 下载 repo（只下载一次；huggingface_hub 自带缓存）
 # ---------------------------------------------------------
 def resolve_model_dirs() -> Tuple[Path, Path]:
-    """
-    返回：
-      repo_root: 你的 MMEdit repo 的本地目录（包含 config.yaml / model.safetensors / vae/）
-      qwen_root: Qwen2-Audio repo 的本地目录
-    """
     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]
@@ -81,7 +62,7 @@ def resolve_model_dirs() -> Tuple[Path, Path]:
         revision=MMEDIT_REVISION,
         local_dir=None,
         local_dir_use_symlinks=False,
-        token=HF_TOKEN,  # 私有 repo 时也可用
+        token=HF_TOKEN,
     )
     repo_root = Path(repo_root).resolve()
 
@@ -102,7 +83,12 @@ def resolve_model_dirs() -> Tuple[Path, Path]:
 # ---------------------------------------------------------
 # 你的音频加载（按你要求：orig -> 16k -> target_sr）
 # ---------------------------------------------------------
-def load_and_process_audio(audio_path: str, target_sr: int) -> torch.Tensor:
+def load_and_process_audio(audio_path: str, target_sr: int):
+    # 延迟导入（避免启动阶段触发 CUDA 初始化）
+    import torch
+    import torchaudio
+    import librosa
+
     path = Path(audio_path)
     if not path.exists():
         raise FileNotFoundError(f"Audio file not found: {audio_path}")
@@ -121,22 +107,14 @@ def load_and_process_audio(audio_path: str, target_sr: int) -> torch.Tensor:
         # 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
-            )
+            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)
 
         # 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_np = librosa.resample(waveform_np, orig_sr=orig_sr_mid, target_sr=int(target_sr))
 
         waveform = torch.from_numpy(waveform_np)
 
@@ -147,11 +125,7 @@ def load_and_process_audio(audio_path: str, target_sr: int) -> torch.Tensor:
 # 校验 repo 结构
 # ---------------------------------------------------------
 def assert_repo_layout(repo_root: Path) -> None:
-    must = [
-        repo_root / "config.yaml",
-        repo_root / "model.safetensors",
-        repo_root / "vae",
-    ]
+    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}")
@@ -162,19 +136,13 @@ def assert_repo_layout(repo_root: Path) -> None:
 
 
 # ---------------------------------------------------------
-# 关键：适配你这个 config.yaml 的路径写法
+# 适配 config.yaml 的路径写法
 # ---------------------------------------------------------
 def patch_paths_in_exp_config(exp_cfg: Dict[str, Any], repo_root: Path, qwen_root: Path) -> None:
-    """
-    适配你 config.yaml：
-      - pretrained_ckpt: ckpt/mmedit/vae/epoch=xx.ckpt  -> repo_root/vae/epoch=xx.ckpt
-      - model_path:      ckpt/qwen2-audio-7B-instruct   -> qwen_root (snapshot_download 结果)
-    """
     # ---- 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/ 开始截断
@@ -202,17 +170,17 @@ def patch_paths_in_exp_config(exp_cfg: Dict[str, Any], repo_root: Path, qwen_roo
 
 # ---------------------------------------------------------
 # Scheduler（与你 exp_cfg.model.noise_scheduler_name 对齐）
-# 注意：有些 repo_id 不存在 scheduler 子目录会 404。
-# 这里给一个 fallback，避免直接炸。
+# 带 fallback：避免 404
 # ---------------------------------------------------------
 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:
         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 to default DDIM config. err={e}")
-        # fallback：不依赖远端 repo
+        logger.warning(f"DDIMScheduler.from_pretrained failed for '{name}', fallback. err={e}")
         return noise_schedulers.DDIMScheduler(
             num_train_timesteps=1000,
             beta_start=0.00085,
@@ -224,21 +192,35 @@ def build_scheduler(exp_cfg: Dict[str, Any]):
         )
 
 
-def _amp_ctx(device: torch.device):
-    # ZeroGPU：只有在 device=cuda 且你明确开启 USE_AMP 才 autocast
+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）
-# ZeroGPU 启动阶段一般没有 CUDA，所以这里不要 model.to("cuda")
 # ---------------------------------------------------------
-def load_pipeline_cpu() -> Tuple[LoadPretrainedBase, object, int]:
+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]
@@ -257,15 +239,13 @@ def load_pipeline_cpu() -> Tuple[LoadPretrainedBase, object, int]:
     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')}")
 
-    # instantiate model（在 CPU 上构建）
     model: LoadPretrainedBase = hydra.utils.instantiate(exp_cfg["model"], _convert_="all")
 
-    # load weights（你的 mmedit 权重）
     ckpt_path = repo_root / "model.safetensors"
     sd = load_file(str(ckpt_path))
     model.load_pretrained(sd)
 
-    # 强制留在 CPU（ZeroGPU 关键）
+    # ZeroGPU：缓存 CPU 版
     model = model.to(torch.device("cpu")).eval()
 
     scheduler = build_scheduler(exp_cfg)
@@ -279,10 +259,8 @@ def load_pipeline_cpu() -> Tuple[LoadPretrainedBase, object, int]:
 # ---------------------------------------------------------
 # 推理：audio + caption -> edited audio
 # ZeroGPU：必须用 @spaces.GPU
-# 并且：函数内再把模型搬到 cuda，推完搬回 cpu
 # ---------------------------------------------------------
 @spaces.GPU
-@torch.no_grad()
 def run_edit(
     audio_file: str,
     caption: str,
@@ -291,6 +269,8 @@ def run_edit(
     guidance_rescale: float,
     seed: int,
 ) -> Tuple[Optional[str], str]:
+    import torch
+
     if audio_file is None or not Path(audio_file).exists():
         return None, "Error: please upload an audio file."
 
@@ -303,7 +283,7 @@ def run_edit(
 
     # 2) ZeroGPU 进入 GPU 区域后，cuda 才会 available
     if not torch.cuda.is_available():
-        return None, "Error: ZeroGPU did not allocate CUDA. Please retry (queue) or check Space hardware."
+        return None, "Error: ZeroGPU did not allocate CUDA. Please retry or check Space hardware."
 
     device = torch.device("cuda")
     logger.info(f"[GPU] torch.cuda.is_available={torch.cuda.is_available()}, device={device}")
@@ -325,7 +305,6 @@ def run_edit(
         "task": ["audio_editing"],
     }
 
-    # 与 infer.config 对齐
     kwargs = {
         "num_steps": int(num_steps),
         "guidance_scale": float(guidance_scale),
@@ -336,20 +315,20 @@ def run_edit(
     kwargs.update(batch)
 
     t0 = time.time()
-    with _amp_ctx(device):
-        out = model.inference(scheduler=scheduler, **kwargs)
+    with torch.no_grad():
+        with amp_autocast(device):
+            out = model.inference(scheduler=scheduler, **kwargs)
     dt = time.time() - t0
 
     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)
 
-    # 4) 推完立刻把模型搬回 CPU（ZeroGPU 关键：避免缓存里残留 cuda tensor）
+    # 4) 推完立刻把模型搬回 CPU（避免缓存残留 cuda tensor）
     model_cpu = model.to("cpu")
     del model
     torch.cuda.empty_cache()
 
-    # 5) 更新缓存（仍然缓存 CPU 版本）
     cache_key = f"{MMEDIT_REPO_ID}@{MMEDIT_REVISION}::{QWEN_REPO_ID}@{QWEN_REVISION}"
     _PIPELINE_CACHE[cache_key] = (model_cpu, scheduler, target_sr)
 
@@ -363,14 +342,12 @@ def build_demo():
     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="Caption (Edit Instruction)", lines=3)
 
-                # 注意：Spaces 不允许你 push 大的 wav 示例。
-                # 最稳的方式：你自己在 Space repo 放一个很小的 demo wav（几百 KB）。
+                # 注意：Space 不建议推大 wav；你可以换成更小的 demo wav
                 gr.Examples(
                     label="example inputs",
                     examples=[
@@ -403,16 +380,15 @@ def build_demo():
         gr.Markdown(
             "## 注意事项\n"
             "1) ZeroGPU 首次点击会分配 GPU，可能稍慢。\n"
-            "2) 如果遇到错误，请重试（尤其是首次启动时）。\n"
-            "3) 原始音频保留可能有bug\n"
+            "2) 如果首次报 cuda 不可用，通常重试一次即可。\n"
         )
+
     return demo
 
 
 if __name__ == "__main__":
     demo = build_demo()
     port = int(os.environ.get("PORT", "7860"))
-    # ZeroGPU：强烈建议 queue；并禁用 SSR 更稳
     demo.queue().launch(
         server_name="0.0.0.0",
         server_port=port,

requirements.txt

@@ -1,4 +1,5 @@
 gradio==4.26.0
+spaces>=0.13.0
 # --- Core Framework (Pinned Versions) ---
 torch==2.5.1
 torchvision==0.20.1