import sys sys.stdout.reconfigure(line_buffering=True) 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 try: import spaces def gpu_decorator(func): return spaces.GPU(func, duration=120) except ImportError: def gpu_decorator(func): return func 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's full stack (text encoder + DiT + VAE) is kept on the GPU at a # time: the two variants need different Qwen2.5-Omni-7B layer-sampling depths, so # the encoder can't be shared as-is (QwenOmniThinkerExtractor bakes dit_depth into # which LLM layers it extracts at construction time). Switching variants evicts the # previous one and rebuilds from the local HF cache (no re-download). _active_name = None _active_entry = None _loading = True _load_error = None # Only one request runs at a time. The scheduler object is shared and stateful, # so two requests running together would corrupt each other's results. This lock # also covers _active_name/_active_entry below, so no separate lock is needed. _gpu_lock = threading.Lock() def _build_variant(name): """Download and construct one variant's full stack: text encoder, DiT backbone, and VAE. Mirrors the model-loading steps in unison/pipelines/infer.py's main().""" 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, DEVICE, 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), DEVICE, 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"], ) audio_vae.to(DEVICE).eval() # probe latent length for a MAX_AUDIO_DURATION-long clip (needed for generation mode). sample_rate = spec["sample_rate"] 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]) scheduler = FlowMatchEulerDiscreteScheduler() print(f"Variant ready: {name}") return (model, extractor, audio_vae, 0.5, sample_rate, gen_target_frames, scheduler) def load_default_variant(): """Background-thread target: build DEFAULT_VARIANT at startup and publish it as the active variant, so the Space doesn't block its HTTP server on the load.""" 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, building and switching to `name` first if a different variant is currently active. Caller must hold _gpu_lock — that's what makes the unlocked reads/writes of _active_name/_active_entry here safe.""" global _active_name, _active_entry if name == _active_name: return _active_entry print(f"Switching variant: {_active_name!r} -> {name!r}") entry = _build_variant(name) _active_name, _active_entry = name, entry gc.collect() torch.cuda.empty_cache() return 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 with _gpu_lock: 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)