zerogpu

app.py

@@ -20,6 +20,9 @@ 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
 
@@ -45,17 +48,22 @@ 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"
 
-_PIPELINE_CACHE: Dict[str, Tuple[LoadPretrainedBase, object, int, torch.device]] = {}
+# ZeroGPU：缓存 CPU pipeline（不要缓存在 CUDA）
+_PIPELINE_CACHE: Dict[str, Tuple[LoadPretrainedBase, object, int]] = {}
+_MODEL_DIR_CACHE: Dict[str, Tuple[Path, Path]] = {}
 
 
 # ---------------------------------------------------------
-# 下载 repo
+# 下载 repo（只下载一次；huggingface_hub 自带缓存）
 # ---------------------------------------------------------
 def resolve_model_dirs() -> Tuple[Path, Path]:
     """
@@ -63,12 +71,17 @@ 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]
+
     logger.info(f"Downloading MMEdit repo: {MMEDIT_REPO_ID} (revision={MMEDIT_REVISION})")
     repo_root = snapshot_download(
         repo_id=MMEDIT_REPO_ID,
         revision=MMEDIT_REVISION,
         local_dir=None,
         local_dir_use_symlinks=False,
+        token=HF_TOKEN,  # 私有 repo 时也可用
     )
     repo_root = Path(repo_root).resolve()
 
@@ -78,9 +91,11 @@ def resolve_model_dirs() -> Tuple[Path, Path]:
         revision=QWEN_REVISION,
         local_dir=None,
         local_dir_use_symlinks=False,
+        token=HF_TOKEN,  # gated 模型必须
     )
     qwen_root = Path(qwen_root).resolve()
 
+    _MODEL_DIR_CACHE[cache_key] = (repo_root, qwen_root)
     return repo_root, qwen_root
 
 
@@ -155,21 +170,15 @@ def patch_paths_in_exp_config(exp_cfg: Dict[str, Any], repo_root: Path, qwen_roo
       - 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("\\", "/")
 
-        # 你这里最稳定的做法：找到 "vae/" 子串之后的后缀
-        # 例如：
-        #   ckpt/mmedit/vae/epoch=13-step=1000000.ckpt  -> vae/epoch=13-step=1000000.ckpt
         idx = vae_ckpt.find("vae/")
         if idx != -1:
             vae_rel = vae_ckpt[idx:]  # 从 vae/ 开始截断
         else:
-            # 兜底：如果有人直接写 epoch=xx.ckpt，那就放到 repo_root/vae/
-            # 或者写 vae/xxx.ckpt
             if vae_ckpt.endswith(".ckpt") and "/" not in vae_ckpt:
                 vae_rel = f"vae/{vae_ckpt}"
             else:
@@ -188,20 +197,35 @@ def patch_paths_in_exp_config(exp_cfg: Dict[str, Any], repo_root: Path, qwen_roo
             )
 
     # ---- 2) Qwen2-Audio model_path ----
-    # 你的 config 里写的是 ckpt/qwen2-audio-7B-instruct，但 Space 上我们直接用下载后的 qwen_root
     exp_cfg["model"]["content_encoder"]["text_encoder"]["model_path"] = str(qwen_root)
 
 
 # ---------------------------------------------------------
 # Scheduler（与你 exp_cfg.model.noise_scheduler_name 对齐）
+# 注意：有些 repo_id 不存在 scheduler 子目录会 404。
+# 这里给一个 fallback，避免直接炸。
 # ---------------------------------------------------------
 def build_scheduler(exp_cfg: Dict[str, Any]):
     name = exp_cfg["model"].get("noise_scheduler_name", "stabilityai/stable-diffusion-2-1")
-    scheduler = noise_schedulers.DDIMScheduler.from_pretrained(name, subfolder="scheduler")
-    return scheduler
+    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
+        return noise_schedulers.DDIMScheduler(
+            num_train_timesteps=1000,
+            beta_start=0.00085,
+            beta_end=0.012,
+            beta_schedule="scaled_linear",
+            clip_sample=False,
+            set_alpha_to_one=False,
+            steps_offset=1,
+        )
 
 
 def _amp_ctx(device: torch.device):
+    # ZeroGPU：只有在 device=cuda 且你明确开启 USE_AMP 才 autocast
     if not USE_AMP:
         return torch.autocast("cuda", enabled=False)
     if device.type != "cuda":
@@ -211,9 +235,10 @@ def _amp_ctx(device: torch.device):
 
 
 # ---------------------------------------------------------
-# 冷启动：load+cache pipeline
+# 冷启动：load+cache pipeline（缓存 CPU 上的 model）
+# ZeroGPU 启动阶段一般没有 CUDA，所以这里不要 model.to("cuda")
 # ---------------------------------------------------------
-def load_pipeline() -> Tuple[LoadPretrainedBase, object, int, torch.device]:
+def load_pipeline_cpu() -> Tuple[LoadPretrainedBase, object, int]:
     cache_key = f"{MMEDIT_REPO_ID}@{MMEDIT_REVISION}::{QWEN_REPO_ID}@{QWEN_REVISION}"
     if cache_key in _PIPELINE_CACHE:
         return _PIPELINE_CACHE[cache_key]
@@ -221,10 +246,9 @@ def load_pipeline() -> Tuple[LoadPretrainedBase, object, int, torch.device]:
     repo_root, qwen_root = resolve_model_dirs()
     assert_repo_layout(repo_root)
 
-    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
     logger.info(f"repo_root = {repo_root}")
-    logger.info(f"device    = {device}")
     logger.info(f"qwen_root = {qwen_root}")
+    logger.info(f"torch.cuda.is_available (startup) = {torch.cuda.is_available()}")
 
     exp_cfg = OmegaConf.load(repo_root / "config.yaml")
     exp_cfg = OmegaConf.to_container(exp_cfg, resolve=True)
@@ -233,25 +257,31 @@ def load_pipeline() -> Tuple[LoadPretrainedBase, object, int, torch.device]:
     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)
 
-    model = model.to(device).eval()
+    # 强制留在 CPU（ZeroGPU 关键）
+    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, device)
-    logger.info("Pipeline loaded and cached.")
-    return model, scheduler, target_sr, device
+    _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
+# 并且：函数内再把模型搬到 cuda，推完搬回 cpu
 # ---------------------------------------------------------
+@spaces.GPU
 @torch.no_grad()
 def run_edit(
     audio_file: str,
@@ -268,12 +298,25 @@ def run_edit(
     if not caption:
         return None, "Error: caption is empty."
 
-    model, scheduler, target_sr, device = load_pipeline()
+    # 1) 取 CPU 缓存
+    model_cpu, scheduler, target_sr = load_pipeline_cpu()
+
+    # 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."
+
+    device = torch.device("cuda")
+    logger.info(f"[GPU] torch.cuda.is_available={torch.cuda.is_available()}, device={device}")
+
+    # 3) 把模型搬到 GPU（临时）
+    model = model_cpu.to(device).eval()
 
+    # seed
     seed = int(seed)
     torch.manual_seed(seed)
     np.random.seed(seed)
 
+    # audio preprocess
     wav = load_and_process_audio(audio_file, target_sr=target_sr).to(device)
 
     batch = {
@@ -282,7 +325,7 @@ def run_edit(
         "task": ["audio_editing"],
     }
 
-    # 和你给的 infer.config 对齐
+    # 与 infer.config 对齐
     kwargs = {
         "num_steps": int(num_steps),
         "guidance_scale": float(guidance_scale),
@@ -301,6 +344,15 @@ def run_edit(
     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）
+    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)
+
     return str(out_path), f"OK | saved={out_path.name} | time={dt:.2f}s | sr={target_sr} | seed={seed}"
 
 
@@ -308,25 +360,24 @@ def run_edit(
 # UI
 # ---------------------------------------------------------
 def build_demo():
-    with gr.Blocks(title="MMEdit Space Simulator") as demo:
-        gr.Markdown("# MMEdit Space 模拟（audio + caption → edited audio）")
-        gr.Markdown(
-            "点下面的示例即可自动填充音频路径与编辑指令，然后点击 Run Editing。"
-        )
+    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)
 
-                # 一键填充示例：点一下就把 audio_in + caption 填好
+                # 注意：Spaces 不允许你 push 大的 wav 示例。
+                # 最稳的方式：你自己在 Space repo 放一个很小的 demo wav（几百 KB）。
                 gr.Examples(
                     label="example inputs",
                     examples=[
-                        ["example/Ym8O802VvJes.wav", "Mix in dog barking in the middle."],
+                        ["./Ym8O802VvJes.wav", "Mix in dog barking in the middle."],
                     ],
                     inputs=[audio_in, caption],
-                    cache_examples=False,  # 本地/Space 都更稳，不提前缓存
+                    cache_examples=False,
                 )
 
                 with gr.Row():
@@ -351,15 +402,20 @@ def build_demo():
 
         gr.Markdown(
             "## 注意事项\n"
-            "- 首次加载较慢\n"
-            "- Space 上有一些bug，某些情况会损失原始音频\n"
+            "1) ZeroGPU 首次点击会分配 GPU，可能稍慢。\n"
+            "2) 如果遇到错误，请重试（尤其是首次启动时）。\n"
+            "3) 原始音频保留可能有bug\n"
         )
-
     return demo
 
 
-
 if __name__ == "__main__":
     demo = build_demo()
-    port = int(os.environ.get("PORT", "7860"))  # Space 默认 7860
-    demo.launch(server_name="0.0.0.0", server_port=port, share=False)
+    port = int(os.environ.get("PORT", "7860"))
+    # ZeroGPU：强烈建议 queue；并禁用 SSR 更稳
+    demo.queue().launch(
+        server_name="0.0.0.0",
+        server_port=port,
+        share=False,
+        ssr_mode=False,
+    )