"""DreamX-Creator 1.0 — native joint audio-video generation on ZeroGPU. Image + prompt -> a video whose soundtrack is denoised jointly with the frames by the same model (gated A2V / V2A cross-attention), so speech, foley and ambience stay in sync with the picture. The inference path is the authors' own release code (`videox_fun/` + `dreamx_inference.py`, copied verbatim from the reference Space `hugging-apps/gd-ml-dreamx-creator` / `AMAP-ML/DreamX-Creator`); this file only wires it into Gradio and stages the checkpoint download so the container never holds all 43 GB of fp32 weights on disk at once. Adapted for the AI Shorts Factory backend (Space `text_amon_API`): - `image` is optional: an empty value yields a neutral keyframe (Option A), so scenes without a first-frame image can still generate. - Returns ``(mp4, last-frame PNG, seed)`` instead of ``(mp4, seed)``: the last frame is what scene 2..N sends back as the first frame for continuity. - Checkpoint root can be redirected to persistent storage with ``DREAMX_CKPT_DIR`` (default: this repo's ``./checkpoints``, the proven upstream layout). Served through Gradio's `/call/generate` protocol (``api_name="generate"``). """ import os os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True") # videox_fun falls back to torch SDPA when flash-attn is absent; the numerics are # identical here because every sequence in the batch is full-length (no padding). os.environ.setdefault("VIDEOX_ATTENTION_TYPE", "FLASH_ATTENTION") import spaces # noqa: E402 (must precede torch) import math # noqa: E402 import random # noqa: E402 import shutil # noqa: E402 import subprocess # noqa: E402 import tempfile # noqa: E402 import time # noqa: E402 from pathlib import Path # noqa: E402 from types import SimpleNamespace # noqa: E402 import numpy as np # noqa: E402 import torch # noqa: E402 # Trusted upstream .pth checkpoints (UMT5 encoder, Wan2.2 VAE) are plain tensor # dicts saved before torch 2.6 flipped `weights_only` to True. _torch_load = torch.load def _torch_load_compat(*args, **kwargs): kwargs.setdefault("weights_only", False) return _torch_load(*args, **kwargs) torch.load = _torch_load_compat import gradio as gr # noqa: E402 from diffusers import FlowMatchEulerDiscreteScheduler # noqa: E402 from huggingface_hub import snapshot_download # noqa: E402 from omegaconf import OmegaConf # noqa: E402 from PIL import Image # noqa: E402 from transformers import AutoTokenizer # noqa: E402 from dreamx_inference import ( # noqa: E402 DEFAULT_NEGATIVE_PROMPT, DirectionalMultimodalCFGAdapter, filter_kwargs, generate_joint_audio_video, ) from videox_fun.models import AutoencoderKLWan3_8, WanT5EncoderModel # noqa: E402 from videox_fun.models.creator_dac_vae import CreatorDACVAE # noqa: E402 from videox_fun.models.creator_gating import WanCreatorGatingAVModel # noqa: E402 REPO_ID = "GD-ML/DreamX-Creator" APP_DIR = Path(__file__).parent.resolve() CKPT_DIR = Path(os.environ.get("DREAMX_CKPT_DIR", str(APP_DIR / "checkpoints"))) CONFIG_PATH = APP_DIR / "config" / "config.yaml" WEIGHT_DTYPE = torch.bfloat16 FPS = 24 CACHE_VERSION = 1 # Authors' defaults (audio_video_generation/inference.py + inference.sh). VIDEO_BRIDGE_GUIDANCE = 3.5 AUDIO_BRIDGE_GUIDANCE = 3.5 VIDEO_SHIFT = 5.0 AUDIO_SHIFT = 5.0 RESOLUTION_CHOICES = [ ("Fast — ~360p", 220), ("Balanced — ~480p", 440), ("Sharp — ~600p", 660), ] # Neutral keyframe (Option A) dimensions — portrait 9:16, shorts-first. KEYFRAME_WIDTH = int(os.environ.get("DREAMX_KEYFRAME_WIDTH", "720")) KEYFRAME_HEIGHT = int(os.environ.get("DREAMX_KEYFRAME_HEIGHT", "1280")) def _log(msg: str) -> None: print(f"[dreamx] {msg}", flush=True) def _disk() -> str: total, used, free = shutil.disk_usage("/") return f"disk used={used / 2**30:.1f}GB free={free / 2**30:.1f}GB" # --------------------------------------------------------------------------- # # Model loading (module scope, eagerly on "cuda" — ZeroGPU packs from here) # --------------------------------------------------------------------------- # _cfg = OmegaConf.load(CONFIG_PATH) _video_kwargs = OmegaConf.to_container(_cfg["video_transformer_additional_kwargs"], resolve=True) _audio_kwargs = OmegaConf.to_container(_cfg["audio_transformer_additional_kwargs"], resolve=True) _gating_kwargs = OmegaConf.to_container(_cfg["creator_gating_kwargs"], resolve=True) _video_vae_kwargs = OmegaConf.to_container(_cfg["video_vae_kwargs"], resolve=True) _text_encoder_kwargs = OmegaConf.to_container(_cfg["text_encoder_kwargs"], resolve=True) _scheduler_kwargs = OmegaConf.to_container(_cfg["scheduler_kwargs"], resolve=True) MAX_SEQUENCE_LENGTH = int(_text_encoder_kwargs.get("text_length", 512)) _log(f"checkpoints in {CKPT_DIR} ({_disk()})") _log(f"downloading joint AV generator ... ({_disk()})") snapshot_download(REPO_ID, local_dir=str(CKPT_DIR), allow_patterns=["creator/*", "creator/**/*"]) _log(f"loading joint AV generator ... ({_disk()})") transformer = WanCreatorGatingAVModel.from_pretrained( pretrained_model_path=str(CKPT_DIR / "creator"), video_pretrained_model_path=str(CKPT_DIR / "creator" / "video_model"), audio_pretrained_model_path=str(CKPT_DIR / "creator" / "audio_model"), video_subfolder=_video_kwargs.get("transformer_low_noise_model_subpath", None), audio_subfolder=_audio_kwargs.get("transformer_low_noise_model_subpath", None), video_kwargs=_video_kwargs, audio_kwargs=_audio_kwargs, low_cpu_mem_usage=True, torch_dtype=WEIGHT_DTYPE, use_temporal_rope=_gating_kwargs.get("use_temporal_rope", True), audio_fps=_gating_kwargs.get("audio_fps", 48000.0 / 960.0), vae_temporal_stride=_gating_kwargs.get("vae_temporal_stride", 4), a2v_cross_attn_layers=_gating_kwargs.get("a2v_cross_attn_layers", None), v2a_cross_attn_layers=_gating_kwargs.get("v2a_cross_attn_layers", None), use_gating=_gating_kwargs.get("use_gating", True), zero_init_cross_attn=_gating_kwargs.get("zero_init_cross_attn", False), zero_init_gating=_gating_kwargs.get("zero_init_gating", True), gate_init_value=_gating_kwargs.get("gate_init_value", 0.0), a2v_gate_alphas=_gating_kwargs.get("a2v_gate_alphas", None), v2a_gate_alphas=_gating_kwargs.get("v2a_gate_alphas", None), ) transformer.eval() # fp32 source shards are no longer needed once the bf16 model is in RAM. shutil.rmtree(CKPT_DIR / "creator", ignore_errors=True) _log(f"joint AV generator loaded ({_disk()})") _log("downloading VAEs + UMT5-xxl text encoder ...") snapshot_download( REPO_ID, local_dir=str(CKPT_DIR), allow_patterns=["audio_vae/*", "wan2.2_ti2v_5b/*", "wan2.2_ti2v_5b/**/*"], ) _log(f"loading VAEs + text encoder ... ({_disk()})") _wan_dir = CKPT_DIR / "wan2.2_ti2v_5b" _video_vae_path = _wan_dir / _video_vae_kwargs.get("vae_subpath", "Wan2.2_VAE.pth") video_vae = AutoencoderKLWan3_8.from_pretrained( str(_video_vae_path), additional_kwargs=_video_vae_kwargs ).eval() audio_vae = CreatorDACVAE.from_pretrained(str(CKPT_DIR / "audio_vae"), strict=False).eval() tokenizer = AutoTokenizer.from_pretrained( str(_wan_dir / _text_encoder_kwargs.get("tokenizer_subpath", "google/umt5-xxl")) ) _text_encoder_path = _wan_dir / _text_encoder_kwargs.get( "text_encoder_subpath", "models_t5_umt5-xxl-enc-bf16.pth" ) text_encoder = WanT5EncoderModel.from_pretrained( str(_text_encoder_path), additional_kwargs=_text_encoder_kwargs, low_cpu_mem_usage=True, torch_dtype=WEIGHT_DTYPE, ).eval() for _stale in (_video_vae_path, _text_encoder_path): try: os.remove(_stale) except OSError: pass shutil.rmtree(CKPT_DIR / "audio_vae", ignore_errors=True) transformer = DirectionalMultimodalCFGAdapter( transformer, video_scale=VIDEO_BRIDGE_GUIDANCE, audio_scale=AUDIO_BRIDGE_GUIDANCE, enable_a2v=True, enable_v2a=True, ).eval() transformer.to("cuda") text_encoder.to("cuda") video_vae.to("cuda") audio_vae.to("cuda") _log(f"all models resident on cuda ({_disk()})") # --------------------------------------------------------------------------- # # Helpers # --------------------------------------------------------------------------- # def _latent_frames(duration: float) -> int: num_frames = int(duration * FPS) num_frames = int((num_frames - 1) // 4 * 4) + 1 return (num_frames - 1) // 4 + 1 DURATION_CAP = 300 def _raw_gpu_seconds(seconds, num_inference_steps, resolution_tokens) -> float: """Predicted GPU seconds, fit to two measured ZeroGPU runs. Measured on RTX PRO 6000 (sm_120), bf16, 3-branch directional multimodal CFG: 2640 video tokens x 10 steps -> 10.77s denoise, 16.0s total 7560 video tokens x 30 steps -> 95.04s denoise, 104.7s total """ tokens = _latent_frames(float(seconds)) * int(resolution_tokens) # per denoising step: linear (FFN/proj) + quadratic (self-attention) terms per_step = 4.02e-4 * tokens + 2.256e-9 * tokens * tokens denoise = per_step * int(num_inference_steps) overhead = 4.5 + 7.0e-4 * tokens # T5 encode, first-frame VAE encode, decode, mux return denoise + overhead def _estimate_duration( image=None, prompt="", seconds=3.0, num_inference_steps=30, resolution_tokens=440, *args, **kwargs, ): """GPU seconds to reserve — calibrated against measured runs on ZeroGPU.""" raw = _raw_gpu_seconds(seconds, num_inference_steps, resolution_tokens) return int(min(DURATION_CAP, math.ceil(raw * 1.15) + 5)) def _write_mp4(frames: np.ndarray, audio: np.ndarray, sample_rate: int, fps: int) -> str: """Mux uint8 RGB frames + mono float audio into a single H.264/AAC mp4.""" height, width = frames.shape[1], frames.shape[2] out_path = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False).name wav_path = tempfile.NamedTemporaryFile(suffix=".wav", delete=False).name import soundfile as sf sf.write(wav_path, np.clip(audio, -1.0, 1.0), sample_rate) base = [ "ffmpeg", "-hide_banner", "-loglevel", "error", "-y", "-f", "rawvideo", "-pix_fmt", "rgb24", "-s", f"{width}x{height}", "-r", str(fps), "-i", "-", "-i", wav_path, ] for vcodec in ("libx264", "mpeg4"): cmd = base + [ "-c:v", vcodec, "-pix_fmt", "yuv420p", "-crf", "18", "-c:a", "aac", "-b:a", "192k", "-shortest", out_path, ] if vcodec == "mpeg4": cmd.remove("-crf") cmd.remove("18") proc = subprocess.run(cmd, input=frames.tobytes(), capture_output=True) if proc.returncode == 0 and os.path.getsize(out_path) > 0: os.remove(wav_path) return out_path _log(f"ffmpeg ({vcodec}) failed: {proc.stderr.decode()[-600:]}") raise gr.Error("ffmpeg failed to encode the generated video.") def _neutral_keyframe(width: int = KEYFRAME_WIDTH, height: int = KEYFRAME_HEIGHT) -> str: """Deterministic dark diagonal gradient — the "Option A" neutral first frame. By default a 9:16 portrait so generated clips keep a shorts-friendly aspect ratio when no real first-frame image is provided. """ y, x = np.mgrid[0:height, 0:width] norm = np.sqrt((x / max(width - 1, 1)) ** 2 + (y / max(height - 1, 1)) ** 2) tone = (0.06 + 0.16 * norm[..., None] * np.array([0.86, 0.95, 1.0])).clip(0, 1) frame = (tone * 255.0).astype(np.uint8) path = tempfile.NamedTemporaryFile(suffix=".png", delete=False).name Image.fromarray(frame).save(path) return path # --------------------------------------------------------------------------- # # Inference # --------------------------------------------------------------------------- # @spaces.GPU(duration=_estimate_duration) def generate( image: str, prompt: str, seconds: float = 3.0, num_inference_steps: int = 30, resolution_tokens: int = 440, guidance_scale: float = 5.0, negative_prompt: str = DEFAULT_NEGATIVE_PROMPT, seed: int = 113, randomize_seed: bool = False, progress=gr.Progress(track_tqdm=True), ): """Generate a video with a jointly-denoised soundtrack from a first frame and a prompt. Args: image: path to the image used as the first frame of the video. Empty means "no image": a neutral keyframe is generated in-app instead. prompt: description of the action AND the sound to generate; put spoken lines in quotes (e.g. `Man says, 'hello there.'`). seconds: length of the clip in seconds (24 fps). num_inference_steps: number of flow-matching denoising steps. resolution_tokens: spatial token budget; higher means higher resolution. guidance_scale: classifier-free guidance strength for text. negative_prompt: what to avoid in both the video and the audio. seed: RNG seed for reproducible sampling. randomize_seed: draw a fresh random seed instead of using `seed`. Returns: A tuple of (path to the generated mp4 with audio, path to the last-frame PNG for chaining into the next scene, the seed actually used). """ if not prompt or not prompt.strip(): raise gr.Error("Please provide a prompt describing the motion and the sound.") if not image: _log("empty first-frame image -> neutral keyframe") image = _neutral_keyframe() # ZeroGPU kills a task that outruns its reservation, and the reservation is # capped, so refuse the few extreme knob combinations that cannot fit. if _raw_gpu_seconds(seconds, num_inference_steps, resolution_tokens) * 1.15 + 5 > DURATION_CAP: raise gr.Error( f"That combination needs more than the {DURATION_CAP}s GPU budget. " "Lower the duration, the resolution, or the number of steps." ) if randomize_seed: seed = random.randint(0, 2**31 - 1) seed = int(seed) device = torch.device("cuda") video_scheduler_kwargs = dict(_scheduler_kwargs, shift=VIDEO_SHIFT) audio_scheduler_kwargs = dict(_scheduler_kwargs, shift=AUDIO_SHIFT) models = { "config": _cfg, "transformer": transformer, "video_vae": video_vae, "audio_vae": audio_vae, "tokenizer": tokenizer, "text_encoder": text_encoder, # fresh schedulers per request: they carry mutable step state "video_scheduler": FlowMatchEulerDiscreteScheduler( **filter_kwargs(FlowMatchEulerDiscreteScheduler, video_scheduler_kwargs) ), "audio_scheduler": FlowMatchEulerDiscreteScheduler( **filter_kwargs(FlowMatchEulerDiscreteScheduler, audio_scheduler_kwargs) ), "max_sequence_length": MAX_SEQUENCE_LENGTH, } args = SimpleNamespace( config_path=str(CONFIG_PATH), image=image, output="output.mp4", negative_prompt=negative_prompt or DEFAULT_NEGATIVE_PROMPT, duration=float(seconds), target_spatial_tokens=int(resolution_tokens), min_token_ratio=0.95, fps=FPS, num_inference_steps=int(num_inference_steps), guidance_scale=float(guidance_scale), cfg_mode="multimodal", video_bridge_guidance_scale=VIDEO_BRIDGE_GUIDANCE, audio_bridge_guidance_scale=AUDIO_BRIDGE_GUIDANCE, seed=seed, video_shift=VIDEO_SHIFT, audio_shift=AUDIO_SHIFT, flow_match_mu=None, sampler_name="Flow", weight_dtype="bfloat16", GPU_memory_mode="model_full_load", text_encoder_cpu_offload=False, video_vae_cpu_offload=False, audio_vae_cpu_offload=False, vae_cpu_offload=False, use_temporal_rope=True, audio_fps=48000.0 / 960.0, vae_temporal_stride=4, disable_a2v_cross_attn=False, disable_v2a_cross_attn=False, suppress_aux_writes=True, skip_output_decode=False, disable_progress=False, synchronize_noise=False, ulysses_degree=1, ring_degree=1, fsdp_dit=False, ) item = { "prompt": prompt.strip(), "video_prompt": prompt.strip(), "audio_prompt": prompt.strip(), "negative_prompt": args.negative_prompt, "audio_negative_prompt": args.negative_prompt, "duration": args.duration, "guidance_scale": args.guidance_scale, "num_inference_steps": args.num_inference_steps, "seed": seed, "first_frame_path": image, "name": "sample", } started = time.perf_counter() video_decoded, audio_decoded, num_frames = generate_joint_audio_video( args, models, device, WEIGHT_DTYPE, item ) gpu_seconds = time.perf_counter() - started frames = ( video_decoded[0].permute(1, 2, 3, 0).clamp(0, 1).numpy() * 255.0 ).astype(np.uint8) waveform = audio_decoded.detach().float().cpu() while waveform.ndim > 1: waveform = waveform[0] audio = waveform.numpy() out_path = _write_mp4(frames, audio, int(audio_vae.sample_rate), FPS) last_frame_path = tempfile.NamedTemporaryFile(suffix=".png", delete=False).name Image.fromarray(frames[-1]).save(last_frame_path) _log( f"done in {gpu_seconds:.1f}s | {num_frames} frames @ {frames.shape[2]}x{frames.shape[1]} " f"| steps={args.num_inference_steps} tokens={args.target_spatial_tokens} " f"| reserved={_estimate_duration(image, prompt, seconds, num_inference_steps, resolution_tokens)}s" ) return out_path, last_frame_path, seed # --------------------------------------------------------------------------- # # UI # --------------------------------------------------------------------------- # EXAMPLES = [ [ "examples/case1.jpg", "A man in a dark suit and white shirt is seated on a yellow couch, speaking about " "the language of Americans. He uses hand gestures to emphasize his points, and the " "background shows a cityscape with illuminated buildings, suggesting an urban setting, " "possibly a studio with a city view. Man says, 'The thing about Americans that I've " "thought about the language is that they speak, they say they speak English.'.", ], [ "examples/case4.jpg", "The video shows a wolf standing on its hind legs, howling with its mouth wide open, " "showing its teeth and tongue. The wolf's ears are perked up, and its eyes are focused " "on something in the distance. The background consists of trees with green and yellow " "leaves, indicating it is autumn. The wolf's howl is loud and resonant, filling the air " "with its powerful voice. The sound of a dog howls and howls can be heard.", ], [ "examples/case3.jpg", "The video captures a dramatic night scene with a series of lightning strikes " "illuminating the dark sky and revealing the city lights below. The clouds move across " "the sky, and the lightning strikes again, followed by a thunderclap. The sound of a " "thunderstorm and rain falling can be heard.", ], [ "examples/case2.jpg", "A humanoid robot is cooking in a modern kitchen. The robot, with a white body and blue " "eyes, is stirring food in a pan on the stove. The kitchen is equipped with dark cabinets " "and various utensils. The robot's movements are smooth and precise as it stirs the food, " "causing steam to rise. The ambient sound of cooking can be heard throughout the scene.", ], [ "examples/case5.jpg", "A person in a red plaid shirt is typing on a white keyboard placed on a wooden table. " "The scene is set in a room with a wooden floor and a part of a white blanket visible in " "the background. The person's hands are actively moving across the keyboard, indicating " "typing activity. The ambient sound is the distinct sound of keys being pressed.", ], [ "examples/case6.jpg", "A man is playing an acoustic guitar in a modern kitchen setting. He is focused on his " "playing, with his hands moving along the strings and fretboard. The background features " "a well-lit kitchen with wooden cabinets and hanging lights.", ], ] CSS = """ #col-container { max-width: 1180px; margin: 0 auto; } .dark .gradio-container { color: var(--body-text-color); } """ with gr.Blocks(title="DreamX-Creator 1.0") as demo: with gr.Column(elem_id="col-container"): gr.Markdown( "# 🎬 DreamX-Creator 1.0\n" "Turn **one image + one prompt** into a video whose **soundtrack is generated " "jointly with the frames** — speech, foley and ambience come out of the same " "denoiser as the picture, so they stay in sync.\n\n" "Describe the *sound* as well as the action. For speech, quote the line: " "`Man says, 'hello there.'`\n\n" "[Model](https://huggingface.co/GD-ML/DreamX-Creator) · " "[Code](https://github.com/AMAP-ML/DreamX-Creator)" ) with gr.Row(): with gr.Column(scale=1): image = gr.Image(label="First frame (optional — leave empty for a neutral keyframe)", type="filepath", height=320) prompt = gr.Textbox( label="Prompt", placeholder="Describe the motion and the sound you want to hear…", lines=4, ) run = gr.Button("Generate audio + video", variant="primary") with gr.Column(scale=1): video_out = gr.Video(label="Result (video + generated audio)", height=380) last_frame_out = gr.Image(label="Last frame (feeds the next scene)", height=380) with gr.Accordion("Advanced settings", open=False): with gr.Row(): seconds = gr.Slider( label="Duration (seconds)", minimum=2.0, maximum=5.0, step=0.5, value=3.0 ) num_inference_steps = gr.Slider( label="Denoising steps", minimum=10, maximum=50, step=1, value=30 ) with gr.Row(): resolution_tokens = gr.Dropdown( label="Resolution budget (spatial tokens)", choices=RESOLUTION_CHOICES, value=440, ) guidance_scale = gr.Slider( label="Guidance scale", minimum=1.0, maximum=10.0, step=0.1, value=5.0 ) negative_prompt = gr.Textbox( label="Negative prompt", value=DEFAULT_NEGATIVE_PROMPT, lines=2 ) with gr.Row(): seed = gr.Number(label="Seed", value=113, precision=0) randomize_seed = gr.Checkbox(label="Randomize seed", value=False) gr.Markdown( "Longer clips, more steps and a bigger resolution budget all cost GPU time " "roughly linearly (and attention grows quadratically with tokens × frames)." ) gr.Examples( examples=EXAMPLES, inputs=[image, prompt], outputs=[video_out, last_frame_out, seed], fn=generate, cache_examples=True, cache_mode="lazy", label=f"Official Verse-Bench cases from the DreamX-Creator repo (v{CACHE_VERSION})", ) run.click( fn=generate, inputs=[ image, prompt, seconds, num_inference_steps, resolution_tokens, guidance_scale, negative_prompt, seed, randomize_seed, ], outputs=[video_out, last_frame_out, seed], api_name="generate", ) if __name__ == "__main__": demo.launch(theme=gr.themes.Citrus(), css=CSS, mcp_server=True)