#!/usr/bin/env python3 # -*- coding: utf-8 -*- # ZeroGPU 关键：必须最先导入 import spaces import traceback import os 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 from huggingface_hub import snapshot_download # ----------------------------- # Logging # ----------------------------- logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s") logger = logging.getLogger("mmedit_space") MMEDIT_REPO_ID = os.environ.get("MMEDIT_REPO_ID", "CocoBro/MMEdit") 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) 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[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})") 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} (revision={QWEN_REVISION})") qwen_root = snapshot_download( repo_id=QWEN_REPO_ID, 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 # --------------------------------------------------------- # 你的音频加载（按你要求：orig -> 16k -> target_sr） # --------------------------------------------------------- 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}") 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 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) 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) @spaces.GPU def run_edit(audio_file, caption, num_steps, guidance_scale, guidance_rescale, seed): import torch import hydra from omegaconf import OmegaConf from safetensors.torch import load_file import diffusers.schedulers as noise_schedulers logger.info("🚀 Starting ..") torch.backends.cuda.matmul.allow_tf32 = False torch.backends.cudnn.allow_tf32 = False try: from utils.config import register_omegaconf_resolvers register_omegaconf_resolvers() except: pass if not audio_file: return None, "Please upload audio." model = None try: # ========================================== logger.info("🚀 Starting ZeroGPU Task...") # 路径准备 repo_root, qwen_root = resolve_model_dirs() 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: p1 = repo_root / "vae" / Path(vae_ckpt).name p2 = repo_root / Path(vae_ckpt).name if p1.exists(): exp_cfg["model"]["autoencoder"]["pretrained_ckpt"] = str(p1) elif p2.exists(): exp_cfg["model"]["autoencoder"]["pretrained_ckpt"] = str(p2) exp_cfg["model"]["content_encoder"]["text_encoder"]["model_path"] = str(qwen_root) # logger.info("Instantiating model (Hydra)...") model = hydra.utils.instantiate(exp_cfg["model"], _convert_="all") # 加载权重 ckpt_path = str(repo_root / "model.safetensors") logger.info(f"Loading weights from {ckpt_path}...") sd = load_file(ckpt_path) model.load_pretrained(sd) del sd # 立即释放 gc.collect() # ========================================== # ========================================== device = torch.device("cuda") logger.info("Moving model to CUDA (FP16)...") # 这一步将模型送入显卡 def safe_move_model(m, dev): logger.info("🛡️ Moving model to GPU in FP32...") for name, child in m.named_children(): child.to(dev, dtype=torch.float32) logger.info(f"Moving {name} to GPU (fp32)...") m.to(dev, dtype=torch.float32) return m model = safe_move_model(model, device) model.eval() logger.info("Model is moved to CUDA.") # Scheduler try: scheduler = noise_schedulers.DDIMScheduler.from_pretrained( exp_cfg["model"].get("noise_scheduler_name", ""), subfolder="scheduler", token=HF_TOKEN ) except: scheduler = noise_schedulers.DDIMScheduler(num_train_timesteps=1000) # ========================================== # 3. 开始推理 # ========================================== target_sr = int(exp_cfg.get("sample_rate", 24000)) torch.manual_seed(int(seed)) np.random.seed(int(seed)) wav = load_and_process_audio(audio_file, target_sr).to(device, dtype=torch.float32) batch = { "audio_id": [Path(audio_file).stem], "content": [{"audio": wav, "caption": caption}], "task": ["audio_editing"], "num_steps": int(num_steps), "guidance_scale": float(guidance_scale), "guidance_rescale": float(guidance_rescale), "use_gt_duration": False, "mask_time_aligned_content": False } logger.info("Inference running...") t0 = time.time() with torch.no_grad(): out = model.inference(scheduler=scheduler, **batch) 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"Success | {time.time()-t0:.2f}s" except Exception as e: err = traceback.format_exc() logger.error(f"❌ ERROR:\n{err}") return None, f"Runtime Error: {e}" finally: # 强制清理，防止下一次任务显存不够 logger.info("Cleaning up...") if model is not None: del model torch.cuda.empty_cache() gc.collect() # ----------------------------- # UI # ----------------------------- def build_demo(): with gr.Blocks(title="MMEdit") as demo: gr.Markdown("# MMEdit ZeroGPU (Direct Load)") with gr.Row(): with gr.Column(): audio_in = gr.Audio(label="Input", type="filepath") caption = gr.Textbox(label="Instruction", lines=3) gr.Examples( label="Examples (Click to load)", # 格式：[ [音频路径1, 提示词1], [音频路径2, 提示词2], ... ] examples=[ # 示例 1 (原本的) ["./Ym8O802VvJes.wav", "Mix in dog barking around the middle."], ["./YDKM2KjNkX18.wav", "Incorporate Telephone bell ringing into the background."], ["./drop_audiocaps_1.wav", "Erase the rain falling sound from the background."], ["./reorder_audiocaps_1.wav", "Switch the positions of the woman's voice and whistling."] ], inputs=[audio_in, caption], # 对应上面列表的顺序：第一个是 Audio，第二个是 Textbox cache_examples=False, # ZeroGPU 环境建议设为 False，避免启动时耗时计算 ) with gr.Row(): num_steps = gr.Slider(10, 100, 50, step=1, label="Steps") guidance_scale = gr.Slider(1.0, 12.0, 5.0, step=0.5, label="Guidance") guidance_rescale = gr.Slider(0.0, 1.0, 0.5, step=0.05, label="Rescale") seed = gr.Number(42, label="Seed") run_btn = gr.Button("Run", variant="primary") with gr.Column(): out = gr.Audio(label="Output") status = gr.Textbox(label="Status") run_btn.click(run_edit, [audio_in, caption, num_steps, guidance_scale, guidance_rescale, seed], [out, status]) return demo if __name__ == "__main__": print("[BOOT] entering main()", flush=True) demo = build_demo() port = int(os.environ.get("PORT", "7860")) print(f"[BOOT] launching gradio on 0.0.0.0:{port}", flush=True) demo.queue().launch( server_name="0.0.0.0", server_port=port, share=False, ssr_mode=False, )