import sys sys.stdout.reconfigure(line_buffering=True) try: import spaces def gpu_decorator(func): return spaces.GPU(func, duration=120) except ImportError: def gpu_decorator(func): return func import gc import tempfile import threading import traceback import types import torch import yaml import gradio as gr from huggingface_hub import hf_hub_download from diffusers import FlowMatchEulerDiscreteScheduler from pyharp import ModelCard, build_endpoint from unison.models.mmaudio.features_utils import FeaturesUtils from unison.pipelines.infer import ( init_text_hidden_extractor, sync_omni_dim_with_text_encoder, _load_model, sample_latents, decode_and_save, decode_and_save_full, load_source_audio, load_ref_audio, make_edit_mask, downsample_mask, join_ref_target_text, transcribe_ref_audio, MAX_AUDIO_DURATION, ) DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu") OMNI_MODEL_ID = "Qwen/Qwen2.5-Omni-7B" UNISON_REPO = "jac22/UNISON" MMAUDIO_REPO = "hkchengrex/MMAudio" DEFAULT_VARIANT = "Balanced (44kHz)" VARIANTS = { "Balanced (44kHz)": { "model_config": "unison/config/D20S0_O_40ch.yaml", "ckpt_file": "unison_D20S0_O_40ch/model.safetensors", "vae_mode": "44k", "vae_ckpt_file": "ext_weights/v1-44.pth", "vocoder_ckpt_file": None, # 44k mode auto-pulls BigVGANv2 from HF Hub "sample_rate": 44100, }, "High detail (16kHz)": { "model_config": "unison/config/D24S0_O_20ch.yaml", "ckpt_file": "unison_D24S0_O_20ch/model.safetensors", "vae_mode": "16k", "vae_ckpt_file": "ext_weights/v1-16.pth", "vocoder_ckpt_file": "ext_weights/best_netG.pt", "sample_rate": 16000, }, } # Only one variant is kept on the GPU at a time; switching evicts the other. _active_name = None _active_entry = None # built on CPU at first; moved to GPU on first real use _loading = True _load_error = None _active_ready = False # has _active_entry been moved onto the GPU yet? def _build_variant(name): """Download and construct one variant's full stack on CPU: text encoder, DiT backbone, and VAE. Safe to call from the background thread — nothing here touches CUDA, since ZeroGPU only intercepts CUDA calls made from inside an @spaces.GPU-decorated call. Moving onto the GPU happens later, in _activate_variant().""" spec = VARIANTS[name] print(f"Building variant: {name} ...") with open(spec["model_config"]) as f: model_config = yaml.safe_load(f) dit_depth = model_config.get("mm_double_blocks_depth", 0) + model_config.get("mm_single_blocks_depth", 0) omni_last_layer_idx = model_config.get("omni_last_layer_idx", -1) extractor = init_text_hidden_extractor( "omni", OMNI_MODEL_ID, None, dit_depth, "cpu", omni_last_layer_idx=omni_last_layer_idx, ) sync_omni_dim_with_text_encoder(model_config, extractor) ckpt_path = hf_hub_download(repo_id=UNISON_REPO, filename=spec["ckpt_file"]) model = _load_model(types.SimpleNamespace(model_ckpt=ckpt_path), "cpu", model_config) vae_ckpt_path = hf_hub_download(repo_id=MMAUDIO_REPO, filename=spec["vae_ckpt_file"]) vocoder_ckpt_path = ( hf_hub_download(repo_id=MMAUDIO_REPO, filename=spec["vocoder_ckpt_file"]) if spec["vocoder_ckpt_file"] else None ) audio_vae = FeaturesUtils( tod_vae_ckpt=vae_ckpt_path, bigvgan_vocoder_ckpt=vocoder_ckpt_path, mode=spec["vae_mode"], ).eval() scheduler = FlowMatchEulerDiscreteScheduler() print(f"Variant built (CPU): {name}") # gen_target_frames is None until _activate_variant() probes it on the GPU. return (model, extractor, audio_vae, 0.5, spec["sample_rate"], None, scheduler) def _activate_variant(entry): """Move a CPU-built variant onto the GPU and probe its latent frame count. Must only be called from inside process_fn (i.e. inside @spaces.GPU) — this is where it's actually safe to touch CUDA.""" model, extractor, audio_vae, vae_scale_factor, sample_rate, _, scheduler = entry extractor.text_backbone = extractor.text_backbone.to(DEVICE) extractor.device = DEVICE model = model.to(device=DEVICE, dtype=torch.bfloat16).eval() audio_vae = audio_vae.to(DEVICE).eval() # probe latent length for a MAX_AUDIO_DURATION-long clip (needed for generation mode). dummy_len = int(MAX_AUDIO_DURATION * sample_rate) with torch.no_grad(): dummy_lat = audio_vae.wrapped_encode(torch.zeros(1, dummy_len, device=DEVICE)) gen_target_frames = int(dummy_lat.shape[-1]) print("Variant activated (GPU)") return (model, extractor, audio_vae, vae_scale_factor, sample_rate, gen_target_frames, scheduler) def load_default_variant(): """Background-thread target: build DEFAULT_VARIANT (CPU only) at startup and publish it as the active variant, so the Space doesn't block its HTTP server on the download. Actually moving it to the GPU happens on the first request.""" global _active_name, _active_entry, _loading, _load_error try: entry = _build_variant(DEFAULT_VARIANT) # No lock needed: _loading stays True until right after this write, and # process_fn won't touch _active_name/_active_entry while _loading is True. _active_name, _active_entry = DEFAULT_VARIANT, entry except Exception: _load_error = traceback.format_exc() print(f"Variant load error: {_load_error}") finally: _loading = False threading.Thread(target=load_default_variant, daemon=True).start() def get_variant(name: str): """Return the active variant's (model, extractor, vae, ...) tuple, ready to use on the GPU — building and/or activating it first if needed. Must be called from inside process_fn (i.e. inside @spaces.GPU) — that's what makes it safe to touch CUDA in _activate_variant().""" global _active_name, _active_entry, _active_ready if name != _active_name: print(f"Switching variant: {_active_name!r} -> {name!r}") _active_name, _active_entry, _active_ready = name, _build_variant(name), False gc.collect() torch.cuda.empty_cache() if not _active_ready: _active_entry = _activate_variant(_active_entry) _active_ready = True return _active_entry model_card = ModelCard( name="UNISON", description=( "Unified sound generation and editing: text-to-audio, text-to-speech, " "audio-scene editing, and zero-shot voice cloning from a single model." ), author="Zhaoqing Li, Haoning Xu, Jingran Su, Yaofang Liu, Zhefan Rao, Huimeng Wang, " "Jiajun Deng, Tianzi Wang, Zengrui Jin, Rui Liu, Haoxuan Che, Xunying Liu", tags=["text-to-audio", "text-to-speech", "audio-editing", "zero-shot-tts"], ) @gpu_decorator @torch.inference_mode() def process_fn( input_audio_path: str, mode: str, sound_type: str, voice: str, prompt: str, background: str, ref_text: str, model_variant: str, steps: int, guidance: float, duration: float, ) -> str: if _loading: raise gr.Error("Model is still loading, please wait a moment and try again.") if _active_entry is None: raise gr.Error(f"Model failed to load: {_load_error}") with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as f: out_path = f.name model, omni_extractor, audio_vae, vae_scale_factor, vae_sample_rate, gen_target_frames, scheduler = get_variant(model_variant) if mode == "Generate": # Text-to-audio/speech, no input audio involved. sound_type picks which of # the model's trained prompt templates to build (see README task table) — # voice/background only apply to the Speech templates, unused otherwise. if sound_type == "Sound Effect": tagged_prompt = f"[Audio] {prompt}" elif sound_type == "Speech": tagged_prompt = f'[Speech] A {voice.lower()} voice saying "{prompt}"' else: # "Speech + Background" tagged_prompt = f'[Speech] A {voice.lower()} voice saying "{prompt}" [Audio] {background}' latents = sample_latents( model, scheduler, omni_extractor, [tagged_prompt], num_inference_steps=steps, guidance_scale=guidance, device=DEVICE, target_frames=gen_target_frames, ) latents = latents * (1.0 / vae_scale_factor) decode_and_save(audio_vae, latents, [duration], [out_path], sample_rate=vae_sample_rate) elif mode == "Edit": # Source audio + instruction. Mask is all-ones: the whole clip is editable, # source_latents just gives the model something to condition on. # sound_type again picks [Audio] vs [Speech] as the edit's sub-tag; "Speech # + Background" isn't a real edit template, so it also falls back to [Audio]. if not input_audio_path: raise gr.Error("Edit mode requires an input audio track.") edit_target = "Speech" if sound_type == "Speech" else "Audio" tgt_wav_len = int(MAX_AUDIO_DURATION * vae_sample_rate) src_wav = load_source_audio(input_audio_path, target_sr=vae_sample_rate, target_length=tgt_wav_len, device=DEVICE) src_latent = audio_vae.wrapped_encode(src_wav) * vae_scale_factor mask_wav = make_edit_mask(src_wav.shape[-1], DEVICE) mask_lat = downsample_mask(mask_wav, src_latent.shape[-1]) latents = sample_latents( model, scheduler, omni_extractor, [f"[Edit] [{edit_target}] {prompt}"], source_latents=src_latent, masks=mask_lat, num_inference_steps=steps, guidance_scale=guidance, device=DEVICE, ) latents = latents * (1.0 / vae_scale_factor) edit_duration = src_wav.shape[-1] / vae_sample_rate decode_and_save(audio_vae, latents, [edit_duration], [out_path], sample_rate=vae_sample_rate) elif mode == "Clone Voice": # Reference clip + text to speak in that voice. if not input_audio_path: raise gr.Error("Clone Voice mode requires a reference audio track.") # 3.0s matches REF_DURATION, the reference-clip length the model was trained with. ref_wav = load_ref_audio(input_audio_path, target_sr=vae_sample_rate, max_ref_duration=3.0, device=DEVICE) ref_audio_duration = ref_wav.shape[-1] / vae_sample_rate if ref_audio_duration + 1.0 > duration: # 1.0s floor so some cloned speech always fits raise gr.Error(f"Reference clip ({ref_audio_duration:.1f}s) leaves under 1s for " f"the cloned speech at Duration={duration:.1f}s. Increase Duration.") # Get the ref transcript (typed one wins over auto-transcription), then # combine it with what the user wants said next into one prompt. # Exclude load_ref_audio's trailing silence pad (tail_pad_s=0.1) so Whisper # doesn't hallucinate tokens over silence. speech_samples = max(ref_wav.shape[-1] - int(0.1 * vae_sample_rate), 1) resolved_ref_text = ref_text.strip() or transcribe_ref_audio( ref_wav[..., :speech_samples], sr=vae_sample_rate, ) combined_text = join_ref_target_text(resolved_ref_text, prompt) full_prompt = f"[Speech with voice] {combined_text}" # Build one waveform [ref audio | silence] and encode it as a single clip — # the model generates the target portion conditioned on the ref portion. total_wav_len = int(duration * vae_sample_rate) ref_wav_1d = ref_wav.squeeze(0) if ref_wav.dim() == 2 else ref_wav ref_wav_len = min(ref_wav_1d.shape[-1], total_wav_len) source_wav = torch.zeros(1, total_wav_len, device=DEVICE) source_wav[:, :ref_wav_len] = ref_wav_1d[:ref_wav_len] source_latent = audio_vae.wrapped_encode(source_wav) * vae_scale_factor mask_wav = torch.zeros(1, 1, total_wav_len, device=DEVICE) mask_wav[:, :, :ref_wav_len] = 2.0 # 2 = reference region, 0 = target (see sample_latents docstring) mask_latent = torch.nn.functional.interpolate( mask_wav, size=source_latent.shape[-1] ).to(torch.bfloat16) latents = sample_latents( model, scheduler, omni_extractor, [full_prompt], source_latents=source_latent, masks=mask_latent, num_inference_steps=steps, guidance_scale=guidance, device=DEVICE, ) latents_full = latents * (1.0 / vae_scale_factor) # Decode the full ref+target latent, then crop out just the target — # the decoder needs the ref portion as context to decode cleanly. ref_samples = int(ref_audio_duration * vae_sample_rate) decode_and_save_full( audio_vae, latents_full, ref_samples, [duration - ref_audio_duration], [out_path], sample_rate=vae_sample_rate, ) else: raise gr.Error(f"Unknown mode: {mode}") return out_path with gr.Blocks() as demo: input_components = [ gr.Audio(type="filepath", label="Input Audio").harp_required(False) .set_info("Source track for Edit mode, or reference voice for Clone Voice mode. Unused in Generate mode."), gr.Dropdown(choices=["Generate", "Edit", "Clone Voice"], value="Generate", label="Mode"), gr.Dropdown(choices=["Sound Effect", "Speech", "Speech + Background"], value="Sound Effect", label="Sound Type", info="Generate/Edit only: what kind of content Prompt describes."), # Voice/Background are only meaningful for the Speech sound types. HARP has no # way to hide a control based on another control's value, so they're always # shown and just ignored (e.g. for Sound Effect) rather than hidden. gr.Dropdown(choices=["Female", "Male"], value="Female", label="Voice", info="Speech sound types only."), gr.Textbox(label="Prompt", info="Generate: describe the sound, or what's said. Edit: describe the change. Clone Voice: text to speak."), gr.Textbox(label="Background (optional)", info="\"Speech + Background\" sound type only: the background sound to mix in."), gr.Textbox(label="Reference Transcript (optional)", info="Only used in Clone Voice mode. Leave blank to auto-transcribe the input audio."), gr.Dropdown(choices=list(VARIANTS), value="Balanced (44kHz)", label="Model"), gr.Slider(minimum=10, maximum=100, step=5, value=50, label="Generation Steps"), gr.Slider(minimum=1.0, maximum=10.0, step=0.5, value=4.5, label="Prompt Strength"), gr.Slider(minimum=1.0, maximum=float(MAX_AUDIO_DURATION), step=0.5, value=10.0, label="Duration (s)"), ] output_components = [ gr.Audio(type="filepath", label="Output Audio").set_info("Generated or edited audio."), ] build_endpoint( model_card=model_card, input_components=input_components, output_components=output_components, process_fn=process_fn, ) demo.queue().launch(pwa=True)