app.py

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)