import os os.environ.setdefault("HF_HOME", "/tmp/hf_home") os.environ.setdefault("HF_MODULES_CACHE", "/tmp/hf_modules") os.environ.setdefault("MPLCONFIGDIR", "/tmp/matplotlib") # LingBot-Video runtime knobs (read at import / first forward). os.environ.setdefault("DIFFUSERS_ATTN_BACKEND", "_native_flash") import spaces # noqa: E402 — must be imported before torch import base64 import io import json import random import re import time import urllib.request from pathlib import Path import gradio as gr import numpy as np import torch from PIL import Image from diffusers import AutoencoderKLWan from diffusers.utils import export_to_video from huggingface_hub import snapshot_download from transformers import AutoProcessor, Qwen3VLForConditionalGeneration from lingbot_video import ( FlowUniPCMultistepScheduler, LingBotVideoImageToVideoPipeline, LingBotVideoPipeline, LingBotVideoTransformer3DModel, ) from lingbot_video.pipeline_lingbot_video import ( DEFAULT_NEGATIVE_PROMPT, DEFAULT_NEGATIVE_PROMPT_IMAGE, ) from lingbot_video.utils import num_frames_from_duration from rewriter_prompts import ( IMAGE_STEP1_EXPAND, IMAGE_STEP2_MAP, VIDEO_STEP1_EXPAND, VIDEO_STEP2_MAP, ) torch.backends.cuda.matmul.allow_tf32 = True torch.set_float32_matmul_precision("high") MODEL_ID = "robbyant/lingbot-video-dense-1.3b" FPS = 24 MAX_SEED = 2**31 - 1 MAX_GPU_SECONDS = 180 REWRITER_MODEL = os.environ.get("REWRITER_MODEL", "Qwen/Qwen3.6-27B:deepinfra") # (height, width), multiples of 16 — official 480p buckets. VIDEO_SIZES = { "832 × 480 (16:9)": (480, 832), "480 × 832 (9:16)": (832, 480), "640 × 480 (4:3)": (480, 640), "480 × 480 (1:1)": (480, 480), } IMAGE_SIZES = { "832 × 480 (16:9)": (480, 832), "480 × 832 (9:16)": (832, 480), "480 × 480 (1:1)": (480, 480), "1280 × 736 (16:9, 720p)": (736, 1280), "736 × 1280 (9:16, 720p)": (1280, 736), "1088 × 1088 (1:1, 1080p)": (1088, 1088), } EXAMPLES_DIR = Path(__file__).parent / "examples" def _read_example(name: str) -> str: return (EXAMPLES_DIR / name).read_text(encoding="utf-8").strip() print(f"[startup] downloading {MODEL_ID} ...", flush=True) t0 = time.perf_counter() model_dir = snapshot_download( MODEL_ID, ignore_patterns=["*.bak_lingbot_video_diffusers"], ) print(f"[startup] snapshot ready in {time.perf_counter() - t0:.1f}s", flush=True) t0 = time.perf_counter() transformer = LingBotVideoTransformer3DModel.from_pretrained( model_dir, subfolder="transformer", torch_dtype=torch.bfloat16 ) text_encoder = Qwen3VLForConditionalGeneration.from_pretrained( model_dir, subfolder="text_encoder", dtype=torch.bfloat16, attn_implementation="sdpa" ) processor = AutoProcessor.from_pretrained(model_dir, subfolder="processor") vae = AutoencoderKLWan.from_pretrained(model_dir, subfolder="vae", torch_dtype=torch.float32) print(f"[startup] components loaded in {time.perf_counter() - t0:.1f}s", flush=True) t0 = time.perf_counter() pipe = LingBotVideoPipeline( transformer=transformer, vae=vae, text_encoder=text_encoder, processor=processor, scheduler=FlowUniPCMultistepScheduler.from_pretrained(model_dir, subfolder="scheduler"), ).to("cuda") pipe_i2v = LingBotVideoImageToVideoPipeline( transformer=transformer, vae=vae, text_encoder=text_encoder, processor=processor, scheduler=FlowUniPCMultistepScheduler.from_pretrained(model_dir, subfolder="scheduler"), ) print(f"[startup] pipelines on cuda in {time.perf_counter() - t0:.1f}s", flush=True) def _resolve_seed(seed: float, randomize: bool) -> int: if randomize: return random.randint(0, MAX_SEED) return int(seed) # --------------------------------------------------------------------------- # Prompt enhancement — the official LingBot rewriter recipe (expand -> JSON map) # run on the official rewriter base model (Qwen3.6-27B) via HF Inference # Providers. The DiT was trained on these structured JSON captions; plain # prompts are far out of distribution and produce severe artifacts. # --------------------------------------------------------------------------- def _router_chat(text: str, max_tokens: int, image: Image.Image | None = None) -> str: token = os.environ.get("HF_TOKEN") if not token: raise RuntimeError("HF_TOKEN is not configured for prompt enhancement.") if image is not None: image = image.convert("RGB") image.thumbnail((768, 768)) buf = io.BytesIO() image.save(buf, format="JPEG", quality=90) data_uri = "data:image/jpeg;base64," + base64.b64encode(buf.getvalue()).decode() content = [ {"type": "image_url", "image_url": {"url": data_uri}}, {"type": "text", "text": text}, ] else: content = text body = { "model": REWRITER_MODEL, "messages": [{"role": "user", "content": content}], "max_tokens": max_tokens, "temperature": 0.0, "chat_template_kwargs": {"enable_thinking": False}, } last_error = None for attempt in range(2): try: req = urllib.request.Request( "https://router.huggingface.co/v1/chat/completions", data=json.dumps(body).encode(), headers={"Authorization": f"Bearer {token}", "Content-Type": "application/json"}, ) with urllib.request.urlopen(req, timeout=180) as resp: data = json.load(resp) out = (data["choices"][0]["message"].get("content") or "").strip() if out: return out last_error = RuntimeError("empty rewriter response") except Exception as exc: # noqa: BLE001 last_error = exc if image is not None: # provider may not accept image input — retry text-only body["messages"][0]["content"] = text image = None raise RuntimeError(f"prompt enhancement failed: {last_error}") def _extract_caption_json(raw: str) -> dict: match = re.search(r"```(?:json)?\s*(\{.*\})\s*```", raw, re.DOTALL) s = match.group(1) if match else raw start = s.find("{") if start < 0: raise ValueError("no JSON object in rewriter output") s = s[start:] try: return json.loads(s) except json.JSONDecodeError: from json_repair import repair_json obj = repair_json(s, return_objects=True) if not isinstance(obj, dict): raise ValueError("rewriter output is not a JSON object") return obj def enhance_prompt(prompt: str, mode: str, duration_s: float | None, image: Image.Image | None = None) -> str: """Two-stage official rewrite: plain prompt -> detailed prose -> JSON caption string.""" if mode == "t2i": step1 = IMAGE_STEP1_EXPAND + "\n\nUser image prompt:\n" + prompt else: step1 = VIDEO_STEP1_EXPAND + "\n\n" + prompt + f"\n\nVideo Duration: {duration_s:g} seconds" prose = _router_chat(step1, 1200, image=image) if mode == "t2i": step2 = IMAGE_STEP2_MAP + "\n\nDETAILED CAPTION:\n" + prose else: step2 = ( VIDEO_STEP2_MAP + f"\n\nVideo Duration: {duration_s:g} seconds\n\nDETAILED CAPTION:\n" + prose + "\n\nOutput the JSON now." ) caption = _extract_caption_json(_router_chat(step2, 6000, image=image)) return json.dumps(caption, ensure_ascii=False, separators=(",", ":")) def _prepare_caption(prompt: str, mode: str, duration_s: float | None, enhance: bool, image: Image.Image | None = None) -> str: prompt = (prompt or "").strip() if not prompt: raise gr.Error("Please enter a prompt.") if prompt.startswith("{") or not enhance: return prompt # already a structured JSON caption, or enhancement disabled t0 = time.perf_counter() try: caption = enhance_prompt(prompt, mode, duration_s, image=image) except Exception as exc: # noqa: BLE001 print(f"[rewriter] failed: {exc}", flush=True) raise gr.Error( "Prompt enhancement failed (the model needs structured captions to work well). " "Please try again, or paste a LingBot JSON caption and disable enhancement." ) print(f"[rewriter] {mode} enhanced in {time.perf_counter() - t0:.1f}s ({len(caption)} chars)", flush=True) return caption def _latent_tokens(height: int, width: int, num_frames: int) -> int: latent_frames = (num_frames - 1) // 4 + 1 return latent_frames * (height // 16) * (width // 16) def _step_seconds(height: int, width: int, num_frames: int) -> float: # Per-step cost (sequential CFG included), calibrated on the live Space: # 30 steps @ 832x480 f=49 took 92.4s → 3.08 s/step at r=0.619. r = _latent_tokens(height, width, num_frames) / 32760.0 return 5.5 * r * r + 1.0 * r + 0.2 def _estimate_video_seconds(size_label: str, duration_s: float, steps: int, sizes=VIDEO_SIZES) -> int: height, width = sizes[size_label] num_frames = num_frames_from_duration(duration_s, FPS) return int(20.0 + steps * _step_seconds(height, width, num_frames)) def _estimate_image_seconds(size_label: str, steps: int) -> int: height, width = IMAGE_SIZES[size_label] return int(10.0 + steps * _step_seconds(height, width, 1)) def _check_budget(estimate: int) -> None: if estimate > MAX_GPU_SECONDS: raise gr.Error( f"These settings need ~{estimate}s of GPU time (max {MAX_GPU_SECONDS}s). " "Reduce the video duration, steps, or resolution." ) def _export_video(frames: np.ndarray) -> str: path = f"/tmp/lingbot_{int(time.time() * 1000)}.mp4" export_to_video(frames, path, fps=FPS) return path def _video_gpu_duration(caption, height, width, num_frames, steps, *args, **kwargs) -> int: return min(MAX_GPU_SECONDS, int(20.0 + int(steps) * _step_seconds(height, width, num_frames)) + 20) @spaces.GPU(duration=_video_gpu_duration) def _gpu_generate_video(caption, height, width, num_frames, steps, guidance_scale, shift, negative_prompt, seed, image=None): target = pipe_i2v if image is not None else pipe kwargs = {"image": image} if image is not None else {} # Sequential CFG. Batched CFG (batch_cfg=True) triggers the packed attention # path which requires flash_attn_varlen_func_v3 (FA3), and the FA3 binaries # only target sm_80/sm_90a — ZeroGPU's RTX PRO 6000 is sm_120, so the packed # path cannot run here. t0 = time.perf_counter() output = target( prompt=caption, negative_prompt=negative_prompt.strip() or DEFAULT_NEGATIVE_PROMPT, height=height, width=width, num_frames=num_frames, num_inference_steps=int(steps), guidance_scale=float(guidance_scale), shift=float(shift), batch_cfg=False, generator=torch.Generator().manual_seed(seed), **kwargs, ) torch.cuda.synchronize() print( f"[timing] {'i2v' if image is not None else 't2v'} {width}x{height} f={num_frames} " f"steps={int(steps)} tokens={_latent_tokens(height, width, num_frames)} " f"took {time.perf_counter() - t0:.1f}s", flush=True, ) return _export_video(output.frames[0]) def _image_gpu_duration(caption, height, width, steps, *args, **kwargs) -> int: return min(MAX_GPU_SECONDS, int(10.0 + int(steps) * _step_seconds(height, width, 1)) + 15) @spaces.GPU(duration=_image_gpu_duration) def _gpu_generate_image(caption, height, width, steps, guidance_scale, shift, negative_prompt, seed): t0 = time.perf_counter() output = pipe( prompt=caption, negative_prompt=negative_prompt.strip() or DEFAULT_NEGATIVE_PROMPT_IMAGE, height=height, width=width, num_frames=1, num_inference_steps=int(steps), guidance_scale=float(guidance_scale), shift=float(shift), batch_cfg=False, generator=torch.Generator().manual_seed(seed), ) torch.cuda.synchronize() print(f"[timing] t2i {width}x{height} steps={int(steps)} " f"took {time.perf_counter() - t0:.1f}s", flush=True) frame = output.frames[0][0] return Image.fromarray((np.clip(frame, 0, 1) * 255).astype(np.uint8)) def generate_t2v( prompt: str, size_label: str = "832 × 480 (16:9)", duration_s: float = 2.0, steps: int = 30, guidance_scale: float = 3.0, shift: float = 3.0, negative_prompt: str = "", seed: int = 42, randomize_seed: bool = True, enhance: bool = True, progress=gr.Progress(track_tqdm=True), ): """Generate a short video from a text prompt with LingBot-Video Dense 1.3B. Args: prompt: Scene description in natural language (it is auto-expanded into the structured caption the model expects), or a raw LingBot JSON caption. size_label: Resolution/aspect preset, e.g. "832 × 480 (16:9)", "480 × 832 (9:16)". duration_s: Video length in seconds (1.0-5.0) at 24 fps. steps: Number of denoising steps (more = more detail, slower). guidance_scale: Classifier-free guidance strength. shift: Flow-matching timestep shift. negative_prompt: What to avoid; empty uses the model default. seed: Random seed for reproducibility. randomize_seed: If true, ignore seed and use a random one. enhance: If true, expand a plain prompt into a structured caption before generation. Returns: The generated MP4 video, the seed used, and the structured caption fed to the model. """ estimate = _estimate_video_seconds(size_label, duration_s, int(steps)) _check_budget(estimate) height, width = VIDEO_SIZES[size_label] num_frames = num_frames_from_duration(duration_s, FPS) seed = _resolve_seed(seed, randomize_seed) caption = _prepare_caption(prompt, "t2v", duration_s, enhance) video = _gpu_generate_video( caption, height, width, num_frames, steps, guidance_scale, shift, negative_prompt, seed, ) return video, seed, caption def generate_i2v( image, prompt: str, size_label: str = "832 × 480 (16:9)", duration_s: float = 2.0, steps: int = 30, guidance_scale: float = 3.0, shift: float = 3.0, negative_prompt: str = "", seed: int = 42, randomize_seed: bool = True, enhance: bool = True, progress=gr.Progress(track_tqdm=True), ): """Animate a first-frame image into a short video with LingBot-Video Dense 1.3B. Args: image: The first frame to animate (filepath or PIL image). prompt: How the scene should evolve, in natural language (auto-expanded), or a raw LingBot JSON caption. size_label: Resolution/aspect preset, e.g. "832 × 480 (16:9)". duration_s: Video length in seconds (1.0-5.0) at 24 fps. steps: Number of denoising steps. guidance_scale: Classifier-free guidance strength. shift: Flow-matching timestep shift. negative_prompt: What to avoid; empty uses the model default. seed: Random seed for reproducibility. randomize_seed: If true, ignore seed and use a random one. enhance: If true, expand a plain prompt into a structured caption before generation. Returns: The generated MP4 video, the seed used, and the structured caption fed to the model. """ if image is None: raise gr.Error("Please upload a first-frame image.") estimate = _estimate_video_seconds(size_label, duration_s, int(steps)) + 15 _check_budget(estimate) height, width = VIDEO_SIZES[size_label] num_frames = num_frames_from_duration(duration_s, FPS) seed = _resolve_seed(seed, randomize_seed) caption = _prepare_caption(prompt, "ti2v", duration_s, enhance, image=image) video = _gpu_generate_video( caption, height, width, num_frames, steps, guidance_scale, shift, negative_prompt, seed, image=image, ) return video, seed, caption def generate_t2i( prompt: str, size_label: str = "832 × 480 (16:9)", steps: int = 30, guidance_scale: float = 3.0, shift: float = 3.0, negative_prompt: str = "", seed: int = 42, randomize_seed: bool = True, enhance: bool = True, progress=gr.Progress(track_tqdm=True), ): """Generate an image from a text prompt with LingBot-Video Dense 1.3B. Args: prompt: Image description in natural language (auto-expanded), or a raw LingBot JSON caption. size_label: Resolution/aspect preset, e.g. "832 × 480 (16:9)", "1088 × 1088 (1:1, 1080p)". steps: Number of denoising steps. guidance_scale: Classifier-free guidance strength. shift: Flow-matching timestep shift. negative_prompt: What to avoid; empty uses the model default. seed: Random seed for reproducibility. randomize_seed: If true, ignore seed and use a random one. enhance: If true, expand a plain prompt into a structured caption before generation. Returns: The generated image, the seed used, and the structured caption fed to the model. """ height, width = IMAGE_SIZES[size_label] seed = _resolve_seed(seed, randomize_seed) caption = _prepare_caption(prompt, "t2i", None, enhance) image = _gpu_generate_image( caption, height, width, steps, guidance_scale, shift, negative_prompt, seed, ) return image, seed, caption HEADER_HTML = """
Embodied-intelligence video generation from a lightweight dense model — text-to-video, image-to-video & text-to-image. Prompts are auto-expanded into structured JSON captions (Qwen3.6-27B); paste a raw LingBot caption to skip that. 480p video, ZeroGPU.