"""OmniVAE T2AV — text to synchronized audio + video. Thin Gradio wrapper around the reference inference path shipped in https://github.com/OpenMOSS/OmniVAE (`generation/infer/t2av/t2av_pipeline.py`), running the released `t2av_recon_distill_avclip` joint checkpoint from https://huggingface.co/OpenMOSS-Team/OmniVAE. Defaults mirror the release smoke test documented in `generation/docs/inference.md` (`validate_checkpoints.sh --cfg 4`): dual CFG, 50 steps, all four guidance scales at 4.0, 121 frames @ 256x256 @ 24 fps with a 5.04 s waveform. """ import os import sys sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) os.environ.setdefault("HF_HUB_ENABLE_HF_TRANSFER", "1") os.environ.setdefault("TOKENIZERS_PARALLELISM", "false") import spaces # noqa: E402 (must precede torch / CUDA-touching imports) import gc # noqa: E402 import json # noqa: E402 import logging # noqa: E402 import random # noqa: E402 import tempfile # noqa: E402 import time # noqa: E402 from pathlib import Path # noqa: E402 import gradio as gr # noqa: E402 import torch # noqa: E402 from huggingface_hub import snapshot_download # noqa: E402 logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(name)s: %(message)s") logger = logging.getLogger("omnivae-demo") MODEL_REPO = "OpenMOSS-Team/OmniVAE" EXPERIMENT = "t2av_recon_distill_avclip" # Only the four subtrees the T2AV path touches (~18 GB of the ~64 GB release). ALLOW_PATTERNS = [ f"models/dit/t2av/{EXPERIMENT}/*", f"models/dit/t2av/{EXPERIMENT}/**", "models/text_encoder/Qwen3.5-0.8B-Base/*", "models/vae/audio_video/recon_distill_avclip/*", "models/vae/audio_only/recon_distill_avclip_ft_decoder/*", # `t2av_pipeline._release_root()` requires both `models/` and `eval/` to # exist under OMNIVAE_RELEASE_ROOT before it will resolve relative paths. "eval/data/t2av/versebench_minimal/*", ] logger.info("Downloading OmniVAE release assets ...") _t0 = time.time() RELEASE_ROOT = snapshot_download(MODEL_REPO, allow_patterns=ALLOW_PATTERNS, max_workers=8) logger.info("Release assets ready in %.1fs at %s", time.time() - _t0, RELEASE_ROOT) os.environ["OMNIVAE_RELEASE_ROOT"] = RELEASE_ROOT CHECKPOINT_DIR = os.path.join(RELEASE_ROOT, "models", "dit", "t2av", EXPERIMENT) import torchaudio # noqa: E402 try: # pragma: no cover import torchcodec # noqa: F401 except ImportError: # torchaudio >= 2.10 delegates `save` to torchcodec, which is not installed # (its wheels are pinned to a specific FFmpeg ABI). The pipeline only needs # a plain 48 kHz WAV write, so route it through soundfile instead. import numpy as _np # noqa: E402 import soundfile as _sf # noqa: E402 def _save_with_soundfile(uri, src, sample_rate, **_kwargs): data = src.detach().to("cpu", torch.float32).numpy() if data.ndim == 2: # torchaudio is (channels, time); soundfile wants (time, channels) data = data.T _sf.write(str(uri), _np.ascontiguousarray(data), int(sample_rate)) torchaudio.save = _save_with_soundfile logger.info("torchcodec unavailable; torchaudio.save patched to use soundfile") from t2av_pipeline import generate_one_av, load_joint_av_pipeline # noqa: E402 # --------------------------------------------------------------------------- # Module-scope load. Components are materialised on CPU (mmap / low_cpu_mem_usage # path) and then moved with `.to("cuda")` so ZeroGPU can intercept the # placement; the upstream `device="cuda"` path uses `device_map={"": "cuda:0"}` # plus `torch.cuda.set_device`, neither of which is ZeroGPU-compatible. # --------------------------------------------------------------------------- logger.info("Loading T2AV pipeline from %s ...", CHECKPOINT_DIR) _t0 = time.time() PIPE = load_joint_av_pipeline(CHECKPOINT_DIR, device="cpu") logger.info("Pipeline loaded on CPU in %.1fs", time.time() - _t0) # `load_univae_ckpt` memoises the raw 4.5 GB + 1.5 GB `state_dict.pt` parses via # an lru_cache; the VAEs are built, so drop them before the ZeroGPU pack step. from omnivae_generation.trainer.vae.univae import _load_univae_raw # noqa: E402 _load_univae_raw.cache_clear() gc.collect() for _name, _module in ( ("text_encoder", PIPE.text_encoder), ("video_vae", PIPE.video_vae), ("audio_vae", PIPE.audio_vae), ("joint_model", PIPE.joint_model), ): _module.to("cuda") _module.eval() logger.info("moved %s to cuda", _name) PIPE.device = torch.device("cuda") gc.collect() logger.info("Pipeline ready (checkpoint_step=%s)", PIPE.checkpoint_step) # Native shapes the released checkpoint was trained / validated at. NUM_FRAMES = 121 FPS = 24.0 HEIGHT = 256 WIDTH = 256 AUDIO_SECONDS = 5.0417 MAX_SEED = 2**31 - 1 DEFAULT_SEED = 20260508 # Measured on this Space's zero-a10g: 20 steps -> 40.9 s, 50 steps -> 97.6 s # wall inside `generate` => ~1.89 s/step plus ~3 s fixed (text encode, VAE # decode of 121 frames, WAV write, ffmpeg mux). Sized with a ~15-20% margin. STEP_SECONDS = 2.2 BASE_SECONDS = 5.0 def _duration(*args, **kwargs): """Reserve GPU time proportional to the requested step count. Tolerant of partial call shapes: `gr.Examples` invokes `generate` with only the prompt, and Gradio appends its `Progress` object. """ steps = kwargs.get("num_inference_steps") if steps is None and len(args) >= 3: steps = args[2] try: steps = int(steps) except (TypeError, ValueError): steps = 50 return int(min(300, BASE_SECONDS + STEP_SECONDS * steps)) @spaces.GPU(duration=_duration) def generate( prompt: str, negative_prompt: str = "", num_inference_steps: int = 50, guidance_scale: float = 4.0, seed: int = DEFAULT_SEED, randomize_seed: bool = False, progress=gr.Progress(track_tqdm=True), ): if not prompt or not prompt.strip(): raise gr.Error("Please enter a prompt describing the scene and its sound.") seed = random.randint(0, MAX_SEED) if randomize_seed else int(seed) guidance_scale = float(guidance_scale) out_dir = tempfile.mkdtemp(prefix="omnivae_t2av_") wall_t0 = time.perf_counter() record = generate_one_av( PIPE, prompt=prompt.strip(), negative_prompt=(negative_prompt or "").strip(), mode="joint_av", # Release smoke test (`--cfg 4`) selects BridgeDiT dual CFG (NFE=3) # and ties all four guidance scales to the same value. cfg_mode="dual", num_inference_steps=int(num_inference_steps), video_guidance_scale=guidance_scale, audio_guidance_scale=guidance_scale, cfg_normalization=False, video_text_guidance=guidance_scale, video_modality_guidance=guidance_scale, audio_text_guidance=guidance_scale, audio_modality_guidance=guidance_scale, num_frames=NUM_FRAMES, fps=FPS, height=HEIGHT, width=WIDTH, audio_duration_seconds=AUDIO_SECONDS, seed=seed, output_dir=out_dir, file_stem="omnivae_t2av", video_quality=8, # The joint model is trained on task-prefixed prompts; the released # validation config keeps the prefix on and the duration suffix off. wrap_task_prefix=True, task_prefix_kind="t2av", append_duration_suffix=False, ) wall = time.perf_counter() - wall_t0 video_path = record.get("av_path") or record.get("video_path") if not video_path or not Path(video_path).is_file(): raise gr.Error("Generation produced no video file. Check the Space logs.") if not record.get("av_path"): logger.warning("ffmpeg mux unavailable; returning silent video.") audio_path = record.get("audio_path") or None details = json.dumps( { "wrapped_prompt": record.get("wrapped_prompt"), "seed": record.get("seed"), "cfg_mode": record.get("cfg_mode"), "num_inference_steps": record.get("num_inference_steps"), "guidance_scale": guidance_scale, "frames": record.get("decoded_num_frames"), "resolution": f"{record.get('width')}x{record.get('height')}", "fps": record.get("fps"), "audio_sample_rate": record.get("sample_rate"), "denoise_seconds": round(float(record.get("elapsed_s", 0.0)), 2), "wall_seconds": round(wall, 2), }, indent=2, ) logger.info("generate() finished in %.1fs (denoise %.1fs)", wall, record.get("elapsed_s", 0.0)) return video_path, audio_path, seed, details EXAMPLES = [ # The two prompts the authors ship as their T2AV demo set # (generation/examples/prompts/t2av_valid.jsonl). ["A street musician plays acoustic guitar on a busy sidewalk while traffic hums in the background."], ["Ocean waves roll onto a sandy beach at sunset with soft wind and distant seabirds."], # From the released validation config's authored prompt list # (models/dit/t2av/t2av_recon_distill_avclip/resolved_config.json). ["An old man telling a story to a group of children sitting around him in a park."], ["A young woman laughing while chatting with a friend at a sunny outdoor cafe."], ["A bowl of steaming noodles on a wooden table in a cozy small restaurant."], ] CSS = """ #col-container { max-width: 900px; margin: 0 auto; } .dark .gradio-container { color: var(--body-text-color); } """ # Gradio 6 moved `theme` / `css` from the Blocks constructor to `launch()`. with gr.Blocks() as demo: with gr.Column(elem_id="col-container"): gr.Markdown( """ # OmniVAE — Text to Audio + Video Generate a short clip with **natively synchronized sound** from a single text prompt. [`OpenMOSS-Team/OmniVAE`](https://huggingface.co/OpenMOSS-Team/OmniVAE) · [code](https://github.com/OpenMOSS/OmniVAE) OmniVAE is a unified audio-video tokenizer; this Space runs the released `t2av_recon_distill_avclip` joint text-to-audio-video model built on top of it — two Z-Image diffusion-transformer branches (video + audio) coupled by bridge cross-attention, so the picture and the soundtrack are denoised together rather than dubbed afterwards. Output is the checkpoint's native resolution: **121 frames · 256x256 · 24 fps · 48 kHz audio (~5 s)**. A default 50-step run takes about 1.5 minutes. """ ) prompt = gr.Textbox( label="Prompt", placeholder="Describe the scene and what it sounds like…", lines=3, ) run_button = gr.Button("Generate", variant="primary") video_out = gr.Video(label="Audio + video", autoplay=True) audio_out = gr.Audio(label="Audio track (48 kHz)") with gr.Accordion("Advanced settings", open=False): negative_prompt = gr.Textbox( label="Negative prompt", value="", placeholder="Left empty in the reference configuration", lines=2, ) num_inference_steps = gr.Slider( label="Inference steps", minimum=10, maximum=100, step=1, value=50 ) guidance_scale = gr.Slider( label="Guidance scale (dual CFG — text & cross-modal)", minimum=1.0, maximum=10.0, step=0.5, value=4.0, ) with gr.Row(): seed = gr.Slider( label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=DEFAULT_SEED ) randomize_seed = gr.Checkbox(label="Randomize seed", value=False) details_out = gr.Code(label="Run details", language="json") gr.Examples( examples=EXAMPLES, inputs=[prompt], outputs=[video_out, audio_out, seed, details_out], fn=generate, cache_examples=True, cache_mode="lazy", ) gr.on( triggers=[run_button.click, prompt.submit], fn=generate, inputs=[prompt, negative_prompt, num_inference_steps, guidance_scale, seed, randomize_seed], outputs=[video_out, audio_out, seed, details_out], ) demo.queue(max_size=8).launch(theme=gr.themes.Citrus(), css=CSS, mcp_server=True)