"""Kandinsky 6.0 Lite (distill 5s) — text/image-to-audio-video demo. Text → video with synchronized audio, or image → video when a reference first frame is supplied. The pipeline (a multimodal DiT with Qwen2.5-VL + CLIP text encoders, a HunyuanVideo VAE for the frames and MMAudioVAE + vocoder for sound) is built at module level on ZeroGPU's emulated CUDA device; each request runs inside ``@spaces.GPU``. The Lite checkpoint is ~3.2B params (bf16), so every component stays resident on the GPU (no offload): DiT ~6.4 GB + Qwen2.5-VL ~16.6 GB + VAEs/vocoder ~3 GB ≈ 26 GB, well inside the 48 GB ``large`` slice. The 10-step distilled PiFlow schedule (guidance_scale=1.0) produces a 5-second 480x864 24 fps clip with 44.1 kHz audio in about a minute. The vendored ``diffusers/`` in this repo is the patched copy that ships ``Kandinsky6TI2VAPipeline`` (Kandinsky6 support is not yet in a released diffusers); it is pinned against one checkpoint revision because upstream re-keys the weights whenever module names change (last time 2026-10-03). Local debug: ``python app.py --smoke "a prompt"`` runs one request without the UI and prints stage timings. """ from __future__ import annotations import os # Before torch is imported: freed weights go back to the device instead of staying reserved # inside one large allocator segment. os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True") import argparse # noqa: E402 import importlib.util # noqa: E402 import subprocess # noqa: E402 import sys # noqa: E402 import tempfile # noqa: E402 import time # noqa: E402 from pathlib import Path # noqa: E402 # SageAttention 2.2.0 for sm_120 has no PyPI wheel; ours (torch 2.8.0+cu128, cp312) ships in # wheels/. The Space build step sees only requirements.txt, so it is installed here, before # anything imports it. _WHEELS = sorted((Path(__file__).resolve().parent / "wheels").glob("sageattention-*.whl")) if _WHEELS and importlib.util.find_spec("sageattention") is None: subprocess.check_call([sys.executable, "-m", "pip", "install", "--no-deps", "--quiet", str(_WHEELS[-1])]) importlib.invalidate_caches() import gradio as gr # noqa: E402 import spaces # noqa: E402 import torch # noqa: E402 from diffusers import Kandinsky6TI2VAPipeline, Kandinsky6Transformer3DModel # noqa: E402 from diffusers.utils import encode_video # noqa: E402 from PIL import Image # noqa: E402 import torch._dynamo.utils # noqa: E402 # ZeroGPU replaces `torch.device` with a wrapper class. If dynamo is imported after that, # its constant-type set holds the wrapper instead of the real class, and any compiled # function reading `tensor.device` fails with "Cannot construct `ConstantVariable` for value # of type ". torch._dynamo.utils.common_constant_types.add(torch._C.device) K6_MODEL = os.environ.get("K6_MODEL", "kandinskylab/Kandinsky-6.0-Lite-distill-5s-Diffusers") # Pinned to the checkpoint revision this vendored diffusers copy was checked against, key for # key: upstream re-keys the weights whenever the module names change (last time 2026-10-03). K6_MODEL_REVISION = os.environ.get("K6_MODEL_REVISION", "314ea02f07cdbc60aa932cb5a0fc9d088b801753") NUM_STEPS = int(os.environ.get("K6_NUM_STEPS", "10")) ATTENTION = os.environ.get("K6_ATTENTION", "native") # diffusers backend name; "native" = SDPA # Only the video self-attention (the dominant cost) gets Sage; text/audio/cross stay on ATTENTION. VIDEO_SELF_ATTENTION = os.environ.get("K6_VIDEO_SELF_ATTENTION", "sage") # ZeroGPU scales the requested duration by the GPU's duration_factor (1.5 on the RTX PRO 6000 # `large` slice). Measured on the live Space: T2AV/TI2AV/expand-prompt all run in ~72-77 s inside # the GPU call; 90 s is that + ~15% margin (cap 135 s). GPU_DURATION = int(os.environ.get("K6_GPU_DURATION", "90")) GPU_SIZE = os.environ.get("K6_ZEROGPU_SIZE", "large") HEIGHT, WIDTH, NUM_FRAMES, FPS = 480, 864, 121, 24 CUDA = torch.device("cuda") # One intra-op thread per core oversubscribes big hosts; 16 is the sweet spot on the Space host. torch.set_num_threads(min(16, os.cpu_count() or 16)) # The pipeline picks its device from its components; pin it to CUDA so the sampler, latents and # generator land on the request's device. Kandinsky6TI2VAPipeline._execution_device = property(lambda self: CUDA) # The checkpoint is all bf16; the default would upcast every modulation to fp32 (+9 GB, slower # steps). Only the timestep embedding needs fp32. Kandinsky6Transformer3DModel._keep_in_fp32_modules = ["time_embeddings"] TITLE = "Kandinsky 6.0 Lite (distill 5s) · text/image → video + audio" DESCRIPTION = """Generate a **5-second video with synchronized audio** from a text prompt (T2AV), or from a prompt plus a reference **first frame** (TI2AV). Kandinsky 6.0 Lite is a ~3.2B multimodal DiT that produces 480x864 24 fps video and 44.1 kHz audio jointly, in 10 distilled PiFlow steps (CFG 1.0). An optional *expand prompt* toggle rewrites short prompts into the detailed style the model was trained on, using the pipeline's own Qwen2.5-VL.""" CSS = """ #col-container { max-width: 1100px; margin: 0 auto; } .dark .gradio-container { color: var(--body-text-color); } """ def build_k6() -> Kandinsky6TI2VAPipeline: """Load the pipeline: everything resident on the GPU (Lite fits the 48 GB large slice).""" pipe = Kandinsky6TI2VAPipeline.from_pretrained( K6_MODEL, revision=K6_MODEL_REVISION, torch_dtype=torch.bfloat16 ) pipe.audio_vae.to(dtype=torch.float32) pipe.vocoder.to(dtype=torch.float32) for name in ("transformer", "vae", "audio_vae", "vocoder", "text_encoder_2", "text_encoder"): if getattr(pipe, name, None) is not None: getattr(pipe, name).to(CUDA) pipe.transformer.set_attention_backend(ATTENTION) if VIDEO_SELF_ATTENTION: for block in pipe.transformer.visual_transformer_blocks: block.video_dec_block.self_attention.set_attention_backend(VIDEO_SELF_ATTENTION) return pipe t0 = time.perf_counter() pipe = build_k6() print(f"[startup] models ready in {time.perf_counter() - t0:.1f}s", flush=True) def ensure_video_attention(timings: dict) -> None: """Fall back to ATTENTION if the Sage kernel cannot run on this GPU, instead of failing.""" if not VIDEO_SELF_ATTENTION or VIDEO_SELF_ATTENTION == ATTENTION: return try: from sageattention import sageattn q = torch.randn(1, 256, 4, 128, device=CUDA, dtype=torch.bfloat16) ref = torch.nn.functional.scaled_dot_product_attention(*(x.transpose(1, 2) for x in (q, q, q))).transpose(1, 2) err = ((sageattn(q, q, q, tensor_layout="NHD") - ref).norm() / ref.norm()).item() if not err < 0.1: raise RuntimeError(f"rel error {err:.3f}") timings["sage_ok"] = 1.0 except Exception as e: # noqa: BLE001 print(f"[request] {VIDEO_SELF_ATTENTION} unavailable ({type(e).__name__}: {e}); using {ATTENTION}", flush=True) for block in pipe.transformer.visual_transformer_blocks: block.video_dec_block.self_attention.set_attention_backend(ATTENTION) timings["sage_ok"] = 0.0 def size_for_image(image: Image.Image) -> tuple[int, int]: """Output (height, width) for a reference image: the area of 480x864 at the image's aspect ratio, multiples of 16, so portrait inputs give portrait videos at the same cost.""" w, h = image.size scale = (HEIGHT * WIDTH / (w * h)) ** 0.5 return max(16, round(h * scale / 16) * 16), max(16, round(w * scale / 16) * 16) def _sync_time() -> float: torch.cuda.synchronize() return time.perf_counter() @spaces.GPU(duration=GPU_DURATION, size=GPU_SIZE) def generate(progress: gr.Progress = gr.Progress(), prompt: str = "", image: Image.Image | None = None, expand_prompt: bool = False, seed: float = 42) -> tuple[str, str, str, str]: """Generate a 5-second video with synchronized audio from a prompt and an optional first frame. Args: prompt: what to generate — describe the scene and the audio you want. image: optional reference first frame (image-to-video); leave empty for text-to-video. expand_prompt: rewrite a short prompt into the detailed style the model was trained on. seed: RNG seed for reproducibility. Returns: (mp4 path, prompt actually used, stats html, timing report) """ # One progress bar driven by hand: track_tqdm would show every tqdm inside the call (pipeline # steps, mp4 encoding, ...) as separate bars. prompt = (prompt or "").strip() if not prompt: raise gr.Error("Enter a prompt.") timings: dict[str, float] = {} start = _sync_time() torch.cuda.reset_peak_memory_stats() ensure_video_attention(timings) image = image.convert("RGB") if image is not None else None height, width = size_for_image(image) if image is not None else (HEIGHT, WIDTH) expanded = prompt gen = torch.Generator(device="cuda").manual_seed(int(seed)) if expand_prompt: progress(0, desc="Expanding the prompt") expanded = pipe.expand_prompts(prompt, pipe.tokenizer, pipe.text_encoder, CUDA, image=image, generator=gen) def on_step(_pipe, i, _t, _kwargs): progress((i + 1, NUM_STEPS), desc="Decoding video" if i + 1 == NUM_STEPS else "Generating", unit="steps") return {} progress((0, NUM_STEPS), desc="Generating", unit="steps") t = _sync_time() with torch.no_grad(): result = pipe(prompt=expanded, image=image, height=height, width=width, num_frames=NUM_FRAMES, frame_rate=float(FPS), num_inference_steps=NUM_STEPS, guidance_scale=1.0, generator=gen, sample_audio=True, output_type="pt", callback_on_step_end=on_step) timings["k6_s"] = _sync_time() - t progress(1.0, desc="Writing MP4") t = time.perf_counter() out = Path(tempfile.mkdtemp(prefix="k6_lite_")) / "video.mp4" # frames[0]: (T, C, H, W) floats in [0, 1]; encode_video wants (T, H, W, C) uint8 and stereo # audio (2, samples). video = (result.frames[0].float().clamp(0, 1) * 255).round().to(torch.uint8).permute(0, 2, 3, 1).cpu() audio = result.audio[0][None].float().repeat(2, 1).cpu() if result.audio is not None else None encode_video(video, fps=FPS, output_path=str(out), audio=audio, audio_sample_rate=pipe.audio_sample_rate if audio is not None else None) timings["mp4_s"] = time.perf_counter() - t timings["total_s"] = time.perf_counter() - start timings["peak_allocated_gib"] = torch.cuda.max_memory_allocated() / 2**30 report = " · ".join(f"{k} {v:.1f}" for k, v in timings.items()) print(f"[request] {report}", flush=True) return str(out), expanded, _stats_html(height, width, timings), report def _stats_html(height: int, width: int, timings: dict[str, float]) -> str: """Small stats line rendered under the video.""" parts = [ f"{NUM_FRAMES} frames · {FPS} fps · {width}x{height}", f"{NUM_STEPS} steps · CFG 1.0 (distilled PiFlow)", ] if "total_s" in timings: parts.append(f"{timings['total_s']:.0f}s on {torch.cuda.get_device_name(0)}") return "  ·  ".join(parts) # --- UI ---------------------------------------------------------------------------------------------------- EXAMPLES_T2AV = [ ["cinematic shot: a giant stone samurai on a stormy cliff above a neon city opens glowing golden eyes and raises a katana. Blue lightning strikes the blade, creating a massive shockwave through the clouds. The camera rapidly pulls back from a low angle. Photorealistic, epic scale, dark blue and gold lighting, rain, sparks, volumetric lightning, blockbuster quality. Audio: heavy rain, deep thunder, metallic sword hum, rising brass and choir, electrical crackle, perfectly synchronized lightning impact, sub-bass shockwave. No dialogue, text, or logos."], ["cinematic shot: a red dragon wakes on a mound of gold coins in a sunlit cave, slowly unfolds its wings and lowers its head toward the glowing stream below. Coins slide and clink down the pile as the dragon shifts its weight. The camera slowly pushes in toward the dragon's eyes. Photorealistic, fantasy, warm golden light, drifting dust motes, lush cave flowers. Audio: deep rumbling breath, leathery wing flaps, clinking coins, trickling water, distant echoing roar. No dialogue, text, or logos."], ] EXAMPLES_TI2AV = [ ["cinematic shot: a red dragon wakes on a mound of gold coins in a sunlit cave, slowly unfolds its wings and lowers its head toward the glowing stream below. Coins slide and clink down the pile as the dragon shifts its weight. The camera slowly pushes in toward the dragon's eyes. Photorealistic, fantasy, warm golden light, drifting dust motes, lush cave flowers. Audio: deep rumbling breath, leathery wing flaps, clinking coins, trickling water, distant echoing roar. No dialogue, text, or logos.", "examples/i2va_input.jpg"], ] def build_demo() -> gr.Blocks: with gr.Blocks(title=TITLE) as demo: gr.Markdown(f"## {TITLE}\n{DESCRIPTION}") with gr.Column(elem_id="col-container"): with gr.Row(): prompt_box = gr.Textbox(label="Prompt", placeholder="Describe the scene and its audio…", lines=4, scale=4) run_btn = gr.Button("Generate", variant="primary", scale=1) image_box = gr.Image(label="First frame (optional — image-to-video)", type="pil") with gr.Accordion("Advanced settings", open=False): expand_box = gr.Checkbox(value=False, label="Expand prompt", info="Rewrite a short prompt into the detailed style the model " "was trained on (uses the pipeline's own Qwen2.5-VL; adds time)") seed_box = gr.Number(value=42, precision=0, label="Seed") video_out = gr.Video(label="Generated video + audio") used_out = gr.Textbox(label="Prompt used", lines=2, visible=True, interactive=False) stats_out = gr.HTML() timing_out = gr.Textbox(label="Timings", visible=False, interactive=False) run_btn.click(generate, inputs=[prompt_box, image_box, expand_box, seed_box], outputs=[video_out, used_out, stats_out, timing_out], api_name="generate") gr.Examples(examples=EXAMPLES_T2AV, inputs=[prompt_box], fn=generate, outputs=[video_out, used_out, stats_out, timing_out], cache_examples=True, cache_mode="lazy", label="Text → video + audio", examples_per_page=2) gr.Examples(examples=EXAMPLES_TI2AV, inputs=[prompt_box, image_box], fn=generate, outputs=[video_out, used_out, stats_out, timing_out], cache_examples=True, cache_mode="lazy", label="Image + text → video + audio", examples_per_page=2) return demo demo = build_demo() if __name__ == "__main__": cli = argparse.ArgumentParser() cli.add_argument("--smoke", metavar="PROMPT", help="run one request without the UI and exit") cli.add_argument("--image", help="reference first frame for --smoke (image-to-video)") args = cli.parse_args() if args.smoke: path, used, _, report = generate(args.smoke, Image.open(args.image) if args.image else None, False, 42) print(f"[smoke] video: {path}\n[smoke] prompt: {used}\n[smoke] {report}") else: demo.queue(default_concurrency_limit=1, max_size=8).launch( theme=gr.themes.Citrus(), css=CSS, mcp_server=True, ssr_mode=False, show_error=True)