fix rescale

app.py

@@ -7,15 +7,12 @@ 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
-# [修改] 移除了顶部的 hydra/models 导入，防止启动时触发 CUDA
 from huggingface_hub import snapshot_download
 
 
@@ -27,7 +24,7 @@ logger = logging.getLogger("mmedit_space")
 
 
 # ---------------------------------------------------------
-# HF Repo IDs
+# HF Repo IDs（按你的默认需求）
 # ---------------------------------------------------------
 MMEDIT_REPO_ID = os.environ.get("MMEDIT_REPO_ID", "CocoBro/MMEdit")
 MMEDIT_REVISION = os.environ.get("MMEDIT_REVISION", None)
@@ -35,14 +32,16 @@ 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)
 _PIPELINE_CACHE: Dict[str, Tuple[object, object, int]] = {}
 # cache: key -> (repo_root, qwen_root)
@@ -50,26 +49,30 @@ _MODEL_DIR_CACHE: Dict[str, Tuple[Path, Path]] = {}
 
 
 # ---------------------------------------------------------
-# 下载 Repo
+# 下载 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}")
+    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_root = Path(repo_root).resolve()
 
-    logger.info(f"Downloading Qwen repo: {QWEN_REPO_ID}")
+    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,
+        local_dir=None,
+        local_dir_use_symlinks=False,
+        token=HF_TOKEN,  # gated 模型必须
     )
     qwen_root = Path(qwen_root).resolve()
 
@@ -78,10 +81,10 @@ def resolve_model_dirs() -> Tuple[Path, Path]:
 
 
 # ---------------------------------------------------------
-# 音频处理
+# 你的音频加载（按你要求：orig -> 16k -> target_sr）
 # ---------------------------------------------------------
 def load_and_process_audio(audio_path: str, target_sr: int):
-    # 延迟导入，防止干扰
+    # 延迟导入（避免启动阶段触发 CUDA 初始化）
     import torch
     import torchaudio
     import librosa
@@ -90,14 +93,18 @@ def load_and_process_audio(audio_path: str, target_sr: int):
     if not path.exists():
         raise FileNotFoundError(f"Audio file not found: {audio_path}")
 
-    waveform, orig_sr = torchaudio.load(str(path))
+    waveform, orig_sr = torchaudio.load(str(path))  # (C, T)
 
-    if waveform.ndim > 1:
-        waveform = waveform.mean(dim=0)
+    # 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)
@@ -105,6 +112,7 @@ def load_and_process_audio(audio_path: str, target_sr: int):
         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))
 
@@ -114,33 +122,55 @@ def load_and_process_audio(audio_path: str, target_sr: int):
 
 
 # ---------------------------------------------------------
-# 路径适配
+# 校验 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("\\", "/")
-        if "vae/" in vae_ckpt:
-            vae_rel = vae_ckpt[vae_ckpt.find("vae/"):]
-        elif vae_ckpt.endswith(".ckpt"):
-            vae_rel = f"vae/{vae_ckpt}" if "/" not in vae_ckpt else vae_ckpt
+        idx = vae_ckpt.find("vae/")
+        if idx != -1:
+            vae_rel = vae_ckpt[idx:]  # 从 vae/ 开始截断
         else:
-            vae_rel = vae_ckpt
+            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()
-        # 鲁棒性检查：如果算出来的路径不存在，尝试在根目录找文件名
-        if not vae_path.exists():
-             fallback = repo_root / Path(vae_ckpt).name
-             if fallback.exists():
-                 vae_path = fallback
-        
         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
+# Scheduler（与你 exp_cfg.model.noise_scheduler_name 对齐）
+# 带 fallback：避免 404
 # ---------------------------------------------------------
 def build_scheduler(exp_cfg: Dict[str, Any]):
     import diffusers.schedulers as noise_schedulers
@@ -150,7 +180,7 @@ def build_scheduler(exp_cfg: Dict[str, Any]):
         scheduler = noise_schedulers.DDIMScheduler.from_pretrained(name, subfolder="scheduler", token=HF_TOKEN)
         return scheduler
     except Exception as e:
-        logger.warning(f"Scheduler fallback: {e}")
+        logger.warning(f"DDIMScheduler.from_pretrained failed for '{name}', fallback. err={e}")
         return noise_schedulers.DDIMScheduler(
             num_train_timesteps=1000,
             beta_start=0.00085,
@@ -162,60 +192,73 @@ def build_scheduler(exp_cfg: Dict[str, Any]):
         )
 
 
+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 to CPU
+# 冷启动：load+cache pipeline（缓存 CPU 上的 model）
 # ---------------------------------------------------------
 def load_pipeline_cpu() -> Tuple[object, object, int]:
-    # [修改] 所有的库都在这里导入，防止全局导入触发 CUDA 初始化
+    # 延迟导入（避免启动阶段触发 CUDA 初始化）
     import torch
     import hydra
     from omegaconf import OmegaConf
     from safetensors.torch import load_file
-    
-    # 你的项目依赖
-    try:
-        from utils.config import register_omegaconf_resolvers
-        from models.common import LoadPretrainedBase
-        register_omegaconf_resolvers()
-    except ImportError:
-        logger.warning("Could not import project utils/models. Ensure they are in the python path.")
-    except Exception:
-        pass
+
+    # 你的项目依赖也延迟导入
+    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()
-    
-    logger.info(f"repo_root: {repo_root}")
-    
+    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')}")
 
-    logger.info("Instantiating model...")
-    model = hydra.utils.instantiate(exp_cfg["model"], _convert_="all")
+    model: LoadPretrainedBase = hydra.utils.instantiate(exp_cfg["model"], _convert_="all")
 
     ckpt_path = repo_root / "model.safetensors"
-    logger.info(f"Loading weights: {ckpt_path}")
     sd = load_file(str(ckpt_path))
     model.load_pretrained(sd)
 
-    # [修改] 确保加载到 CPU
+    # 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 cached.")
+    logger.info("CPU pipeline loaded and cached.")
     return model, scheduler, target_sr
 
 
 # ---------------------------------------------------------
-# [核心] 推理函数 (ZeroGPU 适配版)
+# 推理：audio + caption -> edited audio
+# ZeroGPU：必须用 @spaces.GPU
+# ---------------------------------------------------------
 # ---------------------------------------------------------
 @spaces.GPU
 def run_edit(
@@ -227,6 +270,7 @@ def run_edit(
     seed: int,
 ) -> Tuple[Optional[str], str]:
     import torch
+    import gc
 
     if not audio_file: return None, "Error: Upload audio first."
     if not caption: return None, "Error: Input caption."
@@ -314,43 +358,54 @@ def run_edit(
         torch.cuda.empty_cache()
         gc.collect()
 
-
 # ---------------------------------------------------------
 # UI
 # ---------------------------------------------------------
 def build_demo():
     with gr.Blocks(title="MMEdit (ZeroGPU)") as demo:
-        gr.Markdown("# MMEdit ZeroGPU")
+        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", lines=3)
+                caption = gr.Textbox(label="Caption (Edit Instruction)", lines=3)
 
+                # 注意：Space 不建议推大 wav；你可以换成更小的 demo wav
                 gr.Examples(
-                    label="Examples",
-                    examples=[["./Ym8O802VvJes.wav", "Mix in dog barking around the middle."]],
+                    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="Steps")
-                    guidance_scale = gr.Slider(1.0, 12.0, value=5.0, step=0.5, label="Guidance")
-                    rescale = gr.Slider(0.0, 1.0, 0.5, step=0.05, label="Rescale")
-                    seed = gr.Number(42, label="Seed")
+                    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", variant="primary")
+                run_btn = gr.Button("Run Editing", variant="primary")
 
             with gr.Column():
-                audio_out = gr.Audio(label="Output", type="filepath")
+                audio_out = gr.Audio(label="Edited Audio", type="filepath")
                 status = gr.Textbox(label="Status")
 
         run_btn.click(
             fn=run_edit,
-            inputs=[audio_in, caption, num_steps, guidance_scale, rescale, seed],
+            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
 
 
@@ -361,4 +416,5 @@ if __name__ == "__main__":
         server_name="0.0.0.0",
         server_port=port,
         share=False,
-    )
+        ssr_mode=False,
+    )