Upload folder using huggingface_hub

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+examples/ti2v_frame.png filter=lfs diff=lfs merge=lfs -text

README.md

@@ -1,13 +1,29 @@
 ---
-title: Lingbot Video Dense
-emoji: 📉
-colorFrom: green
+title: LingBot-Video Dense 1.3B
+emoji: 🎬
+colorFrom: red
 colorTo: blue
 sdk: gradio
 sdk_version: 6.20.0
-python_version: '3.12'
 app_file: app.py
-pinned: false
+license: apache-2.0
+short_description: Dense 1.3B embodied video generation — T2V, I2V, T2I
+models:
+  - robbyant/lingbot-video-dense-1.3b
+python_version: "3.12"
+startup_duration_timeout: 1h
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# LingBot-Video Dense 1.3B — ZeroGPU demo
+
+Text-to-video, image-to-video and text-to-image with
+[LingBot-Video Dense 1.3B](https://huggingface.co/robbyant/lingbot-video-dense-1.3b),
+a lightweight dense video generation model for embodied intelligence (1.3B parameters).
+
+This demo runs the dense DiT with:
+
+- bf16 transformer + text encoder (Qwen3-VL), fp32 Wan VAE
+- 480p resolution buckets for video, up to 5 s at 24 fps (plus 720p/1080p text-to-image)
+- batched CFG with SDPA attention, TF32
+- automatic prompt enhancement: the official two-step rewriter recipe (expand → JSON map) runs
+  on the official rewriter base model (Qwen3.6-27B) via HF Inference Providers

app.py

@@ -0,0 +1,771 @@
+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 (CFG included), calibrated relative to 81f @ 480x832.
+    # Dense 1.3B is ~3x faster than MoE 30B-A3B per the paper.
+    r = _latent_tokens(height, width, num_frames) / 32760.0
+    return 2.3 * r * r + 0.45 * r + 0.15
+
+
+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 {}
+    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=True,
+        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=True,
+        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 = """
+<div id="lb-header">
+  <div class="lb-top">
+    <div class="lb-title">
+      LingBot-Video
+      <span class="lb-badge">Dense&nbsp;1.3B</span>
+    </div>
+    <nav class="lb-links">
+      <a href="https://huggingface.co/robbyant/lingbot-video-dense-1.3b" target="_blank" rel="noopener">Model</a>
+      <a href="https://github.com/Robbyant/lingbot-video" target="_blank" rel="noopener">GitHub</a>
+      <a href="https://technology.robbyant.com/lingbot-video" target="_blank" rel="noopener">Project</a>
+    </nav>
+  </div>
+  <p class="lb-sub">
+    Embodied-intelligence video generation from a lightweight dense model — text-to-video, image-to-video &amp; text-to-image.
+    Prompts are auto-expanded into structured JSON captions
+    (<a href="https://huggingface.co/Qwen/Qwen3.6-27B" target="_blank" rel="noopener">Qwen3.6-27B</a>);
+    paste a raw LingBot caption to skip that. 480p video, ZeroGPU.
+  </p>
+</div>
+"""
+
+CSS = """
+#col-container { max-width: 1100px; margin: 0 auto; }
+.dark .gradio-container { color: var(--body-text-color); }
+/* Zero out Gradio's wrapper padding around the header HTML block. */
+#lb-header-wrap, #lb-header-wrap .html-container {
+  padding: 0 !important;
+  border: none !important;
+  background: transparent !important;
+}
+/* Borderless header, flush with the app container. */
+#lb-header {
+  margin: 0;
+  padding: 2px 0 8px 0;
+}
+#lb-header .lb-top {
+  display: flex;
+  align-items: center;
+  justify-content: flex-start;
+  flex-wrap: wrap;
+  gap: 6px 18px;
+}
+#lb-header .lb-title {
+  font-size: 1.35rem;
+  font-weight: 700;
+  line-height: 1.2;
+  color: var(--body-text-color);
+  display: flex;
+  align-items: center;
+  gap: 10px;
+}
+#lb-header .lb-badge {
+  font-size: 0.66rem;
+  font-weight: 600;
+  letter-spacing: 0.02em;
+  text-transform: uppercase;
+  color: var(--body-text-color-subdued);
+  background: var(--background-fill-primary);
+  border: 1px solid var(--border-color-primary);
+  border-radius: 999px;
+  padding: 2px 9px;
+  white-space: nowrap;
+}
+#lb-header .lb-links {
+  display: flex;
+  gap: 14px;
+  font-size: 0.85rem;
+  padding-left: 4px;
+}
+#lb-header .lb-links a {
+  color: var(--body-text-color);
+  opacity: 0.85;
+  text-decoration: none;
+  border-bottom: 1px solid transparent;
+}
+#lb-header .lb-links a:hover {
+  color: var(--link-text-color);
+  border-bottom-color: currentColor;
+}
+#lb-header .lb-sub {
+  margin: 7px 0 0 0;
+  font-size: 0.82rem;
+  line-height: 1.45;
+  color: var(--body-text-color);
+  opacity: 0.78;
+}
+#lb-header .lb-sub a { color: var(--link-text-color); text-decoration: none; }
+#lb-header .lb-sub a:hover { text-decoration: underline; }
+/* Gradio's theme pads <a> with 0 8px, which shows up as fake space inside the
+   "(Qwen3.6-27B)" parentheses — strip it on all header links. */
+#lb-header a { padding: 0 !important; }
+
+/* Compact, tidy examples: no wide horizontal overflow, subtle scrollbar */
+.lb-examples .gr-samples-table td, .lb-examples table td { white-space: normal; }
+.lb-examples { scrollbar-width: thin; }
+.lb-examples ::-webkit-scrollbar { height: 8px; width: 8px; }
+.lb-examples ::-webkit-scrollbar-thumb {
+  background: var(--border-color-primary);
+  border-radius: 999px;
+}
+.lb-examples ::-webkit-scrollbar-track { background: transparent; }
+"""
+
+
+def _seed_row():
+    with gr.Row():
+        seed = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=42)
+        randomize = gr.Checkbox(label="Randomize seed", value=True)
+    return seed, randomize
+
+
+def _advanced(negative_default: str, steps_value: int):
+    with gr.Accordion("Advanced settings", open=False):
+        steps = gr.Slider(label="Inference steps", minimum=4, maximum=50, step=1, value=steps_value)
+        guidance = gr.Slider(label="Guidance scale", minimum=1.0, maximum=10.0, step=0.5, value=3.0)
+        shift = gr.Slider(label="Flow shift", minimum=1.0, maximum=9.0, step=0.5, value=3.0)
+        negative = gr.Textbox(label="Negative prompt", value=negative_default, lines=3)
+        seed, randomize = _seed_row()
+    return steps, guidance, shift, negative, seed, randomize
+
+
+# Example runners: examples carry only the natural-language prompt (and first
+# frame for I2V) — enhancement is on and every other control keeps its default,
+# so the examples table stays to one or two clean columns instead of ten.
+def _example_t2v(prompt):
+    return generate_t2v(prompt, "832 × 480 (16:9)", 2.0, 30, 3.0, 3.0,
+                        DEFAULT_NEGATIVE_PROMPT, 42, False, True)
+
+
+def _example_i2v(image, prompt):
+    return generate_i2v(image, prompt, "832 × 480 (16:9)", 2.0, 30, 3.0, 3.0,
+                        DEFAULT_NEGATIVE_PROMPT, 42, False, True)
+
+
+def _example_t2i(prompt):
+    return generate_t2i(prompt, "832 × 480 (16:9)", 30, 3.0, 3.0,
+                        DEFAULT_NEGATIVE_PROMPT_IMAGE, 42, False, True)
+
+
+# Pre-built example prompts derived from the official LingBot-Video example cases.
+# For t2v/t2i, we use a short natural-language prompt that the rewriter expands.
+# For i2v, we pass the first-frame image and a natural-language motion description.
+T2V_EXAMPLES = [
+    "A young woman with long, wavy brown hair stands in a bright modern apartment living room, "
+    "wearing an oversized cream knit cardigan over a white tank top with beige trousers. She "
+    "smiles at the camera, shifts her weight, and adjusts her collar, showcasing her outfit.",
+    "A young child with short brown hair plays outdoors on a sunny day, blowing shimmering soap "
+    "bubbles with a wand. The bubbles float upwards, catching the warm afternoon light against a "
+    "soft-focus green background.",
+    "A robotic arm on a workbench reaches forward, grasps a black game controller, lifts it, "
+    "moves it to the right, and lowers it into an open box. Another robotic arm and headphones "
+    "remain stationary on the desk, top-down view.",
+]
+T2V_LABELS = ["👗 Woman in apartment", "🫧 Child blowing bubbles", "🤖 Robot arm sorting"]
+
+I2V_EXAMPLES = [
+    [str(EXAMPLES_DIR / "ti2v_frame.png"),
+     "A fit young man and a sleek white humanoid robot run side by side along a cherry blossom "
+     "lined promenade toward the camera, which tracks backward. Energetic, futuristic, daylight."],
+]
+I2V_LABELS = ["🏃 Man and robot running"]
+
+T2I_EXAMPLES = [
+    "A clear glass bottle of water on a sunlit wooden table acts as a lens, refracting bright "
+    "sunlight into a warm glow, extreme close-up, photorealistic.",
+    "A humanoid robot chef flipping a pancake in a bright modern kitchen, dramatic side lighting, "
+    "steam rising from the pan, shallow depth of field, 85mm lens.",
+]
+T2I_LABELS = ["🔆 Bottle as a lens", "🤖 Robot chef"]
+
+
+with gr.Blocks(title="LingBot-Video Dense 1.3B", theme=gr.themes.Citrus(), css=CSS) as demo:
+    gr.HTML(HEADER_HTML, elem_id="lb-header-wrap")
+
+    with gr.Tab("Text → Video"):
+        with gr.Row():
+            with gr.Column():
+                t2v_prompt = gr.Textbox(
+                    label="Prompt",
+                    placeholder="Describe the scene in detail (plain text or a LingBot JSON caption)...",
+                    lines=5,
+                    max_lines=12,
+                )
+                with gr.Row():
+                    t2v_size = gr.Dropdown(
+                        label="Resolution", choices=list(VIDEO_SIZES), value="832 × 480 (16:9)"
+                    )
+                    t2v_dur = gr.Slider(
+                        label="Video duration (s)", minimum=1.0, maximum=5.0, step=0.5, value=2.0
+                    )
+                t2v_enhance = gr.Checkbox(
+                    label="Enhance prompt (official rewriter — required for plain prompts)", value=True
+                )
+                t2v_steps, t2v_guidance, t2v_shift, t2v_negative, t2v_seed, t2v_rand = _advanced(
+                    DEFAULT_NEGATIVE_PROMPT, 30
+                )
+                t2v_btn = gr.Button("Generate video", variant="primary")
+            with gr.Column():
+                t2v_out = gr.Video(label="Generated video", autoplay=True)
+                t2v_seed_out = gr.Number(label="Seed used", interactive=False)
+                with gr.Accordion("Structured caption used", open=False):
+                    t2v_caption_out = gr.Textbox(label="Caption", lines=4)
+
+        t2v_inputs = [
+            t2v_prompt, t2v_size, t2v_dur, t2v_steps, t2v_guidance, t2v_shift,
+            t2v_negative, t2v_seed, t2v_rand, t2v_enhance,
+        ]
+        t2v_outputs = [t2v_out, t2v_seed_out, t2v_caption_out]
+        t2v_btn.click(generate_t2v, inputs=t2v_inputs, outputs=t2v_outputs,
+                      concurrency_id="gpu", concurrency_limit=1)
+        with gr.Column(elem_classes="lb-examples"):
+            gr.Examples(
+                examples=T2V_EXAMPLES,
+                example_labels=T2V_LABELS,
+                fn=_example_t2v,
+                inputs=[t2v_prompt],
+                outputs=t2v_outputs,
+                cache_examples=True,
+                cache_mode="lazy",
+            )
+
+    with gr.Tab("Image → Video"):
+        with gr.Row():
+            with gr.Column():
+                i2v_image = gr.Image(label="First frame", type="pil")
+                i2v_prompt = gr.Textbox(
+                    label="Prompt",
+                    placeholder="Describe how the scene should evolve...",
+                    lines=4,
+                    max_lines=12,
+                )
+                with gr.Row():
+                    i2v_size = gr.Dropdown(
+                        label="Resolution", choices=list(VIDEO_SIZES), value="832 × 480 (16:9)"
+                    )
+                    i2v_dur = gr.Slider(
+                        label="Video duration (s)", minimum=1.0, maximum=5.0, step=0.5, value=2.0
+                    )
+                i2v_enhance = gr.Checkbox(
+                    label="Enhance prompt (official rewriter — required for plain prompts)", value=True
+                )
+                i2v_steps, i2v_guidance, i2v_shift, i2v_negative, i2v_seed, i2v_rand = _advanced(
+                    DEFAULT_NEGATIVE_PROMPT, 30
+                )
+                i2v_btn = gr.Button("Generate video", variant="primary")
+            with gr.Column():
+                i2v_out = gr.Video(label="Generated video", autoplay=True)
+                i2v_seed_out = gr.Number(label="Seed used", interactive=False)
+                with gr.Accordion("Structured caption used", open=False):
+                    i2v_caption_out = gr.Textbox(label="Caption", lines=4)
+
+        i2v_inputs = [
+            i2v_image, i2v_prompt, i2v_size, i2v_dur, i2v_steps, i2v_guidance, i2v_shift,
+            i2v_negative, i2v_seed, i2v_rand, i2v_enhance,
+        ]
+        i2v_outputs = [i2v_out, i2v_seed_out, i2v_caption_out]
+        i2v_btn.click(generate_i2v, inputs=i2v_inputs, outputs=i2v_outputs,
+                      concurrency_id="gpu", concurrency_limit=1)
+        with gr.Column(elem_classes="lb-examples"):
+            gr.Examples(
+                examples=I2V_EXAMPLES,
+                example_labels=I2V_LABELS,
+                fn=_example_i2v,
+                inputs=[i2v_image, i2v_prompt],
+                outputs=i2v_outputs,
+                cache_examples=True,
+                cache_mode="lazy",
+            )
+
+    with gr.Tab("Text → Image"):
+        with gr.Row():
+            with gr.Column():
+                t2i_prompt = gr.Textbox(
+                    label="Prompt",
+                    placeholder="Describe the image in detail...",
+                    lines=5,
+                    max_lines=12,
+                )
+                t2i_size = gr.Dropdown(
+                    label="Resolution", choices=list(IMAGE_SIZES), value="832 × 480 (16:9)"
+                )
+                t2i_enhance = gr.Checkbox(
+                    label="Enhance prompt (official rewriter — required for plain prompts)", value=True
+                )
+                t2i_steps, t2i_guidance, t2i_shift, t2i_negative, t2i_seed, t2i_rand = _advanced(
+                    DEFAULT_NEGATIVE_PROMPT_IMAGE, 30
+                )
+                t2i_btn = gr.Button("Generate image", variant="primary")
+            with gr.Column():
+                t2i_out = gr.Image(label="Generated image")
+                t2i_seed_out = gr.Number(label="Seed used", interactive=False)
+                with gr.Accordion("Structured caption used", open=False):
+                    t2i_caption_out = gr.Textbox(label="Caption", lines=4)
+
+        t2i_inputs = [
+            t2i_prompt, t2i_size, t2i_steps, t2i_guidance, t2i_shift,
+            t2i_negative, t2i_seed, t2i_rand, t2i_enhance,
+        ]
+        t2i_outputs = [t2i_out, t2i_seed_out, t2i_caption_out]
+        t2i_btn.click(generate_t2i, inputs=t2i_inputs, outputs=t2i_outputs,
+                      concurrency_id="gpu", concurrency_limit=1)
+        with gr.Column(elem_classes="lb-examples"):
+            gr.Examples(
+                examples=T2I_EXAMPLES,
+                example_labels=T2I_LABELS,
+                fn=_example_t2i,
+                inputs=[t2i_prompt],
+                outputs=t2i_outputs,
+                cache_examples=True,
+                cache_mode="lazy",
+            )
+
+demo.queue(max_size=30).launch(css=CSS, mcp_server=True)

examples/t2v_bubbles.json

@@ -0,0 +1,101 @@
+{
+  "caption": {
+    "comprehensive_description": {
+      "scene_content_description": "A young child with short brown hair is outdoors on a bright, sunny day, playing with bubbles. The child is wearing a colorful striped shirt with horizontal bands of red, yellow, green, and blue. They are holding a small bottle of bubble solution and a wand, dipping the wand into the bottle and then blowing to create a stream of shimmering, iridescent bubbles. The background is a soft-focus outdoor setting with greenery and a hint of a building, creating a warm and joyful atmosphere. The lighting is bright and natural, highlighting the child's focused expression and the delicate, translucent nature of the bubbles.",
+      "camera_movement_description": "The camera is essentially stationary throughout the video, maintaining a medium close-up shot of the child from an eye-level angle. There is a very slight handheld tremor, but no intentional panning, tilting, or zooming occurs."
+    },
+    "camera_info": {
+      "color": "Saturated",
+      "frame_size": "Medium Close Up",
+      "shot_type_angle": "Low angle",
+      "lens_size": "Long Lens",
+      "composition": "Center",
+      "lighting": "Hard light",
+      "lighting_type": "Daylight"
+    },
+    "world_knowledge": [],
+    "prominent_elements": [
+      {
+        "name": "young child",
+        "description": "A young child with short brown hair and a joyful expression, focused on playing with bubbles.",
+        "actions": [
+          {
+            "timestamp": "[0.0s - 1.5s]",
+            "action": "Dips the bubble wand into the solution bottle and lifts it out."
+          },
+          {
+            "timestamp": "[1.5s - 2.5s]",
+            "action": "Brings the wand to their lips and blows to create bubbles."
+          },
+          {
+            "timestamp": "[2.5s - 5.0s]",
+            "action": "Watches the bubbles float away, smiling slightly."
+          }
+        ],
+        "location": "Center of the frame",
+        "relative_size": "dominant",
+        "shape_and_color": "Human form; wearing a multi-colored striped shirt (red, yellow, green, blue).",
+        "texture": "Soft skin, fine hair, fabric texture of the shirt.",
+        "appearance_details": "Short brown hair, bright eyes, colorful horizontal stripes on the shirt.",
+        "relationship": "Holding the bubble wand and bottle, interacting with the bubbles.",
+        "orientation": "Facing forward and slightly to the right.",
+        "pose": "Standing or sitting upright, arms raised to hold the bubble wand.",
+        "expression": "Focused and happy.",
+        "clothing": "A short-sleeved shirt with horizontal stripes in red, yellow, green, and blue.",
+        "gender": "Male",
+        "skin_tone_and_texture": "Fair skin with a smooth texture."
+      },
+      {
+        "name": "bubble wand and bottle",
+        "description": "A small plastic bottle containing bubble solution and a wand with a circular loop.",
+        "actions": [
+          {
+            "timestamp": "[0.0s - 1.5s]",
+            "action": "The wand is dipped into the bottle and then lifted out."
+          },
+          {
+            "timestamp": "[1.5s - 2.5s]",
+            "action": "The wand is held at the child's mouth as bubbles are blown."
+          }
+        ],
+        "location": "Lower center of the frame, held by the child.",
+        "relative_size": "small",
+        "shape_and_color": "Cylindrical bottle with a blue cap; circular wand loop.",
+        "texture": "Smooth plastic.",
+        "appearance_details": "The bottle is partially filled with clear liquid; the wand has a thin handle.",
+        "relationship": "Held by the child's hands.",
+        "orientation": "Vertical bottle, horizontal wand loop.",
+        "pose": "",
+        "expression": "",
+        "clothing": "",
+        "gender": "",
+        "skin_tone_and_texture": ""
+      },
+      {
+        "name": "bubbles",
+        "description": "A cluster of small, iridescent, translucent bubbles floating in the air.",
+        "actions": [
+          {
+            "timestamp": "[1.5s - 5.0s]",
+            "action": "The bubbles are blown from the wand and float upwards and towards the right side of the frame."
+          }
+        ],
+        "location": "Upper center and right side of the frame.",
+        "relative_size": "medium",
+        "shape_and_color": "Spherical and iridescent with rainbow-like reflections.",
+        "texture": "Glossy and translucent.",
+        "appearance_details": "Shimmering surfaces that catch the sunlight.",
+        "relationship": "Created by the child's blowing action.",
+        "orientation": "Floating in various directions.",
+        "pose": "",
+        "expression": "",
+        "clothing": "",
+        "gender": "",
+        "skin_tone_and_texture": "",
+        "is_cluster": true,
+        "number_of_objects": "many"
+      }
+    ]
+  },
+  "duration": 5
+}

examples/t2v_robot.json

@@ -0,0 +1,155 @@
+{
+  "caption": {
+    "comprehensive_description": {
+      "scene_content_description": "The video presents a first-person perspective of a workspace, likely a desk, viewed from a fixed, slightly elevated angle. The environment is well-lit with neutral, even lighting, creating a clear and focused atmosphere. The desk surface is a light grey color. On the left side of the desk, there is a black game controller with a glowing blue light bar, resting on a black mousepad. In the center of the desk lies a pair of black over-ear headphones. To the right, there is an open, empty black box with a white interior. In the foreground, two robotic arms are visible. The left robotic arm, featuring a black and silver body with a two-fingered gripper, remains completely stationary throughout the video. The right robotic arm, similar in design, is the active subject. It begins by moving forward and to the left, positioning its gripper over the game controller. It then grasps the controller, lifts it off the desk, and transports it to the right, moving it over the open box. Finally, the right robotic arm lowers the controller into the box and releases it. The headphones and the left robotic arm remain undisturbed during this entire sequence.",
+      "camera_movement_description": ""
+    },
+    "camera_info": {
+      "color": "Cyan",
+      "frame_size": "Wide",
+      "shot_type_angle": "High angle",
+      "lens_size": "Ultra Wide / Fisheye",
+      "composition": "Balanced",
+      "lighting": "Soft light",
+      "lighting_type": "Artificial light"
+    },
+    "world_knowledge": [],
+    "prominent_elements": [
+      {
+        "name": "right robotic arm",
+        "description": "A mechanical arm with a black and silver body, equipped with a two-fingered gripper and visible wiring.",
+        "actions": [
+          {
+            "timestamp": "[0.0s - 1.5s]",
+            "action": "Moves forward and to the left towards the game controller."
+          },
+          {
+            "timestamp": "[1.5s - 2.5s]",
+            "action": "Grasps the game controller and lifts it upward."
+          },
+          {
+            "timestamp": "[2.5s - 4.0s]",
+            "action": "Moves to the right, carrying the game controller over the open box."
+          },
+          {
+            "timestamp": "[4.0s - 5.0s]",
+            "action": "Lowers the game controller into the box and releases it."
+          }
+        ],
+        "location": "Originates from the bottom right, moves to the center, then to the right.",
+        "relative_size": "large",
+        "shape_and_color": "Cylindrical and angular segments, black and silver.",
+        "texture": "Metallic and matte plastic.",
+        "appearance_details": "Visible joints, wiring, and a two-fingered gripper mechanism.",
+        "relationship": "Interacts directly with the game controller.",
+        "orientation": "Extends forward and slightly upward from the bottom right.",
+        "pose": "",
+        "expression": "",
+        "clothing": "",
+        "is_cluster": false,
+        "number_of_objects": ""
+      },
+      {
+        "name": "left robotic arm",
+        "description": "A mechanical arm with a black and silver body, equipped with a two-fingered gripper.",
+        "actions": [
+          {
+            "timestamp": "[0.0s - 5.0s]",
+            "action": "Remains stationary."
+          }
+        ],
+        "location": "Bottom left corner of the frame.",
+        "relative_size": "large",
+        "shape_and_color": "Cylindrical and angular segments, black and silver.",
+        "texture": "Metallic and matte plastic.",
+        "appearance_details": "Visible joints and a two-fingered gripper mechanism.",
+        "relationship": "Positioned opposite the right robotic arm, not interacting with other objects.",
+        "orientation": "Extends forward and slightly upward from the bottom left.",
+        "pose": "",
+        "expression": "",
+        "clothing": "",
+        "is_cluster": false,
+        "number_of_objects": ""
+      },
+      {
+        "name": "game controller",
+        "description": "A standard video game controller with joysticks, buttons, and a glowing light bar.",
+        "actions": [
+          {
+            "timestamp": "[0.0s - 1.5s]",
+            "action": "Rests stationary on the desk."
+          },
+          {
+            "timestamp": "[1.5s - 2.5s]",
+            "action": "Is grasped and lifted upward by the right robotic arm."
+          },
+          {
+            "timestamp": "[2.5s - 4.0s]",
+            "action": "Is moved to the right by the right robotic arm."
+          },
+          {
+            "timestamp": "[4.0s - 5.0s]",
+            "action": "Is lowered into the open box and released."
+          }
+        ],
+        "location": "Initially on the left side of the desk, moved to the right side inside the box.",
+        "relative_size": "medium",
+        "shape_and_color": "Contoured shape, black with a blue light bar.",
+        "texture": "Matte plastic.",
+        "appearance_details": "Two joysticks, directional pad, action buttons, and a glowing blue light bar in the center.",
+        "relationship": "Initially on the desk, then grasped and moved by the right robotic arm, finally placed in the open box.",
+        "orientation": "Horizontal on the desk, then tilted while being carried.",
+        "pose": "",
+        "expression": "",
+        "clothing": "",
+        "is_cluster": false,
+        "number_of_objects": ""
+      },
+      {
+        "name": "headphones",
+        "description": "A pair of over-ear headphones with a headband and ear cups.",
+        "actions": [
+          {
+            "timestamp": "[0.0s - 5.0s]",
+            "action": "Remains stationary on the desk."
+          }
+        ],
+        "location": "Center of the desk.",
+        "relative_size": "medium",
+        "shape_and_color": "Curved headband with circular ear cups, black.",
+        "texture": "Matte plastic and soft ear cushions.",
+        "appearance_details": "Visible ear cushions and a headband.",
+        "relationship": "Rests on the desk, untouched by the robotic arms.",
+        "orientation": "Lying flat on the desk.",
+        "pose": "",
+        "expression": "",
+        "clothing": "",
+        "is_cluster": false,
+        "number_of_objects": ""
+      },
+      {
+        "name": "open box",
+        "description": "A rectangular box with the lid open, revealing a white interior.",
+        "actions": [
+          {
+            "timestamp": "[0.0s - 5.0s]",
+            "action": "Remains stationary on the desk."
+          }
+        ],
+        "location": "Right side of the desk.",
+        "relative_size": "medium",
+        "shape_and_color": "Rectangular, black exterior with a white interior.",
+        "texture": "Smooth cardboard or plastic.",
+        "appearance_details": "Open lid, empty interior.",
+        "relationship": "Serves as the receptacle for the game controller.",
+        "orientation": "Horizontal on the desk.",
+        "pose": "",
+        "expression": "",
+        "clothing": "",
+        "is_cluster": false,
+        "number_of_objects": ""
+      }
+    ]
+  },
+  "duration": 5
+}

examples/t2v_woman.json

@@ -0,0 +1,77 @@
+{
+  "caption": {
+    "comprehensive_description": {
+      "scene_content_description": "A young woman with long, wavy brown hair is standing in a bright, modern apartment living room. She is wearing a stylish, oversized cream-colored knit cardigan over a white tank top, paired with high-waisted, wide-leg beige trousers. She holds a small, structured tan leather handbag in her left hand. The background features a neutral-toned interior with a beige sofa, a potted plant, and large windows that let in soft, natural light, creating a warm and inviting atmosphere. The woman is smiling and looking directly at the camera, showcasing her outfit with a confident and friendly demeanor.",
+      "camera_movement_description": "The camera is positioned at eye level and remains essentially stationary throughout the video, maintaining a medium shot that captures the subject from the waist up. There is a very shallow depth of field, keeping the woman in sharp focus while the background remains softly blurred."
+    },
+    "camera_info": {
+      "color": "Warm",
+      "frame_size": "Medium",
+      "shot_type_angle": "Eye level",
+      "lens_size": "Medium",
+      "composition": "Center",
+      "lighting": "Soft light",
+      "lighting_type": "Daylight"
+    },
+    "world_knowledge": [],
+    "prominent_elements": [
+      {
+        "name": "young woman",
+        "description": "A woman with long, wavy brown hair and a friendly expression, modeling a fashion outfit.",
+        "actions": [
+          {
+            "timestamp": "[0.0s - 0.5s]",
+            "action": "stands still, smiling at the camera"
+          },
+          {
+            "timestamp": "[0.5s - 2.0s]",
+            "action": "shifts her weight and turns her body slightly to the right"
+          },
+          {
+            "timestamp": "[2.0s - 3.5s]",
+            "action": "adjusts the collar of her cardigan with her right hand"
+          },
+          {
+            "timestamp": "[3.5s - 5.0s]",
+            "action": "returns to a neutral pose, smiling at the camera"
+          }
+        ],
+        "location": "center of the frame",
+        "relative_size": "dominant",
+        "shape_and_color": "slender build; wearing cream, white, and beige",
+        "texture": "soft knit cardigan, smooth fabric trousers",
+        "appearance_details": "long wavy brown hair, gold hoop earrings, tan leather handbag",
+        "relationship": "the main subject of the video, standing in front of a blurred apartment background",
+        "orientation": "upright, facing the camera",
+        "pose": "standing, shifting weight, and adjusting clothing",
+        "expression": "smiling and confident",
+        "clothing": "oversized cream knit cardigan, white tank top, high-waisted wide-leg beige trousers",
+        "gender": "female",
+        "skin_tone_and_texture": "fair skin with a smooth texture"
+      },
+      {
+        "name": "tan handbag",
+        "description": "A small, structured leather handbag with a top handle.",
+        "actions": [
+          {
+            "timestamp": "[0.0s - 5.0s]",
+            "action": "held steady in the woman's left hand"
+          }
+        ],
+        "location": "held in the woman's left hand, lower center of the frame",
+        "relative_size": "small",
+        "shape_and_color": "rectangular, tan or light brown",
+        "texture": "smooth leather",
+        "appearance_details": "structured shape with a top handle",
+        "relationship": "held by the woman as an accessory",
+        "orientation": "upright",
+        "pose": "",
+        "expression": "",
+        "clothing": "",
+        "gender": "",
+        "skin_tone_and_texture": ""
+      }
+    ]
+  },
+  "duration": 5
+}

examples/ti2v_robot.json

@@ -0,0 +1,131 @@
+{
+  "caption": {
+    "comprehensive_description": {
+      "scene_content_description": "A dynamic, high-fidelity sequence of a fit young man and a sleek white humanoid robot running side-by-side along a paved promenade lined with blooming cherry blossom trees. The man is dressed in black athletic wear, while the robot features a glossy white chassis with black mechanical joints. They are running towards the camera, which tracks backward to maintain their position in the frame. The background features a blurred cityscape and a bridge over a river, creating a sense of depth and urban vitality. The atmosphere is energetic and futuristic, blending organic nature with advanced technology.",
+      "camera_movement_description": "The camera executes a smooth, continuous backward tracking shot, moving parallel to the subjects' forward motion to keep them centered in the frame. There is a subtle, rhythmic vertical bobbing motion synchronized with the cadence of the runners' footsteps, adding a visceral sense of speed and physical exertion. The camera maintains a steady focus on the two runners, allowing the background elements to exhibit motion blur."
+    },
+    "camera_info": {
+      "color": "Natural",
+      "frame_size": "Wide",
+      "shot_type_angle": "Eye level",
+      "lens_size": "Telephoto",
+      "composition": "Symmetrical",
+      "lighting": "Bright sunlight",
+      "lighting_type": "Daylight"
+    },
+    "world_knowledge": [],
+    "prominent_elements": [
+      {
+        "name": "fit young man",
+        "description": "A muscular male with short dark hair, wearing a tight black t-shirt and black running shorts.",
+        "actions": [
+          {
+            "timestamp": "[0.0s - 5.0s]",
+            "action": "He runs steadily towards the camera, his arms pumping rhythmically at his sides and his legs cycling with powerful strides, maintaining a consistent pace."
+          }
+        ],
+        "location": "Left side of the frame (viewer's perspective)",
+        "relative_size": "large",
+        "shape_and_color": "Athletic build, dark clothing contrasting with the bright surroundings.",
+        "texture": "Realistic skin and fabric",
+        "appearance_details": "White running shoes, focused and determined facial expression.",
+        "relationship": "The primary human subject",
+        "orientation": "Facing forward",
+        "pose": "Running posture, arms bent at the elbows",
+        "expression": "Focused and determined",
+        "clothing": "Black t-shirt, black shorts, white sneakers",
+        "gender": "Male",
+        "skin_tone_and_texture": "Tanned, muscular"
+      },
+      {
+        "name": "humanoid robot",
+        "description": "A sleek, anthropomorphic machine with a glossy white exterior and exposed black mechanical joints at the shoulders, elbows, knees, and ankles.",
+        "actions": [
+          {
+            "timestamp": "[0.0s - 5.0s]",
+            "action": "It runs with mechanical precision, its limbs moving in a fluid, lifelike gait that perfectly matches the man's stride, its arms swinging in coordination with its legs."
+          }
+        ],
+        "location": "Right side of the frame (viewer's perspective)",
+        "relative_size": "large",
+        "shape_and_color": "Angular and futuristic, white with black accents.",
+        "texture": "Smooth, reflective metal and matte composite materials",
+        "appearance_details": "A smooth, featureless black visor for a face; the label '07' is visible on its chest plate.",
+        "relationship": "The companion subject running alongside the man",
+        "orientation": "Facing forward",
+        "pose": "Running posture, mirroring the human",
+        "expression": "",
+        "clothing": "",
+        "gender": "",
+        "skin_tone_and_texture": ""
+      },
+      {
+        "name": "cherry blossom trees",
+        "description": "A dense row of mature trees with dark trunks and canopies full of vibrant pink flowers.",
+        "actions": [
+          {
+            "timestamp": "[0.0s - 5.0s]",
+            "action": "The trees appear to rush past the camera due to the backward tracking motion, with the pink blossoms blurring slightly to emphasize speed."
+          }
+        ],
+        "location": "Background, lining both sides of the path",
+        "relative_size": "large",
+        "shape_and_color": "Organic and sprawling, with dark brown bark and bright pink blossoms.",
+        "texture": "Rough bark and delicate petals",
+        "appearance_details": "Hanging branches that frame the top of the shot.",
+        "relationship": "The environmental setting",
+        "orientation": "Receding into the distance",
+        "pose": "",
+        "expression": "",
+        "clothing": "",
+        "gender": "",
+        "skin_tone_and_texture": ""
+      },
+      {
+        "name": "paved promenade",
+        "description": "A wide, brick-paved walkway stretching from the foreground into the distance.",
+        "actions": [
+          {
+            "timestamp": "[0.0s - 5.0s]",
+            "action": "The ground rushes towards the bottom of the frame, providing a strong sense of forward momentum."
+          }
+        ],
+        "location": "Lower center of the frame",
+        "relative_size": "large",
+        "shape_and_color": "Rectangular and linear, composed of reddish-brown bricks.",
+        "texture": "Rough and uneven",
+        "appearance_details": "Distinct pattern of the paving stones.",
+        "relationship": "The ground surface for the runners",
+        "orientation": "Leading the eye towards the horizon",
+        "pose": "",
+        "expression": "",
+        "clothing": "",
+        "gender": "",
+        "skin_tone_and_texture": ""
+      },
+      {
+        "name": "urban background",
+        "description": "A blurred cityscape featuring tall buildings and a steel bridge spanning a body of water.",
+        "actions": [
+          {
+            "timestamp": "[0.0s - 5.0s]",
+            "action": "The background remains relatively static but shifts slightly due to the camera's lateral movement, remaining out of focus to keep attention on the runners."
+          }
+        ],
+        "location": "Far background, visible through the gaps in the trees and to the right",
+        "relative_size": "medium",
+        "shape_and_color": "Geometric and muted, with grey, blue, and beige tones.",
+        "texture": "Indistinct due to depth of field",
+        "appearance_details": "Silhouette of skyscrapers and the truss structure of the bridge.",
+        "relationship": "The distant context",
+        "orientation": "Static in the distance",
+        "pose": "",
+        "expression": "",
+        "clothing": "",
+        "gender": "",
+        "skin_tone_and_texture": ""
+      }
+    ]
+  },
+  "duration": 5
+}

lingbot_video/__init__.py

@@ -0,0 +1,36 @@
+from __future__ import annotations
+
+import importlib
+from typing import Any
+
+
+_EXPORTS = {
+    "FlowUniPCMultistepScheduler": (
+        "lingbot_video.scheduling_flow_unipc",
+        "FlowUniPCMultistepScheduler",
+    ),
+    "LingBotVideoImageToVideoPipeline": (
+        "lingbot_video.pipeline_lingbot_video_i2v",
+        "LingBotVideoImageToVideoPipeline",
+    ),
+    "LingBotVideoPipeline": (
+        "lingbot_video.pipeline_lingbot_video",
+        "LingBotVideoPipeline",
+    ),
+    "LingBotVideoTransformer3DModel": (
+        "lingbot_video.transformer_lingbot_video",
+        "LingBotVideoTransformer3DModel",
+    ),
+}
+
+__all__ = sorted(_EXPORTS)
+
+
+def __getattr__(name: str) -> Any:
+    try:
+        module_name, attr_name = _EXPORTS[name]
+    except KeyError as exc:
+        raise AttributeError(name) from exc
+    value = getattr(importlib.import_module(module_name), attr_name)
+    globals()[name] = value
+    return value

lingbot_video/default_negative_prompt.json

@@ -0,0 +1,57 @@
+{
+  "universal_negative": {
+    "visual_quality": [
+      "low quality",
+      "worst quality",
+      "blurry",
+      "pixelated",
+      "jpeg artifacts",
+      "low resolution",
+      "unstable color",
+      "color flicker",
+      "underexposed",
+      "overexposed",
+      "invisible subject",
+      "subject hidden in darkness"
+    ],
+    "artistic_style": [
+      "painting",
+      "illustration",
+      "drawing",
+      "cartoon",
+      "3d render",
+      "cgi",
+      "sketch",
+      "digital art"
+    ],
+    "composition_and_content": [
+      "text",
+      "watermark",
+      "signature",
+      "logo",
+      "subtitles",
+      "pillarboxed",
+      "side bars",
+      "portrait image in landscape frame"
+    ],
+    "temporal_and_motion_stability": [
+      "flickering",
+      "jittery",
+      "motion blur",
+      "temporal inconsistency",
+      "warping",
+      "morphing",
+      "incoherent motion",
+      "unnatural movement",
+      "static object with sudden jump",
+      "frame-to-frame inconsistency"
+    ],
+    "material_and_structure": [
+      "plastic-like glass",
+      "unrealistic texture",
+      "deformed bottle",
+      "liquid freezing improperly",
+      "distorted reflections"
+    ]
+  }
+}

lingbot_video/default_negative_prompt_image.json

@@ -0,0 +1,41 @@
+{
+    "universal_negative": {
+        "visual_quality": [
+            "low quality",
+            "worst quality",
+            "blurry",
+            "pixelated",
+            "jpeg artifacts",
+            "low resolution",
+            "underexposed",
+            "overexposed",
+            "invisible subject",
+            "subject hidden in darkness"
+        ],
+        "artistic_style": [
+            "painting",
+            "illustration",
+            "drawing",
+            "cartoon",
+            "3d render",
+            "cgi",
+            "sketch",
+            "digital art"
+        ],
+        "composition_and_content": [
+            "text",
+            "watermark",
+            "signature",
+            "logo",
+            "pillarboxed",
+            "side bars",
+            "portrait image in landscape frame"
+        ],
+        "material_and_structure": [
+            "plastic-like glass",
+            "unrealistic texture",
+            "deformed bottle",
+            "distorted reflections"
+        ]
+    }
+}

lingbot_video/fsdp_inference.py

@@ -0,0 +1,113 @@
+from __future__ import annotations
+
+import gc
+from dataclasses import dataclass
+from typing import Any
+
+import torch
+import torch.distributed as dist
+from torch.distributed.device_mesh import DeviceMesh, init_device_mesh
+
+try:
+    from torch.distributed._composable.fsdp import fully_shard
+except Exception as exc:  # pragma: no cover - depends on the installed torch build
+    fully_shard = None
+    FSDP_IMPORT_ERROR = exc
+else:
+    FSDP_IMPORT_ERROR = None
+
+
+@dataclass(frozen=True)
+class FSDPInferenceInfo:
+    enabled: bool
+    world_size: int
+    wrapped_blocks: int
+    ignored_params: int
+
+
+def init_fsdp_inference_mesh() -> DeviceMesh | None:
+    if not dist.is_available() or not dist.is_initialized():
+        return None
+    world_size = dist.get_world_size()
+    if world_size <= 1:
+        return None
+    return init_device_mesh("cuda", (world_size,), mesh_dim_names=("fsdp",))
+
+
+def _move_buffers_to_device(module: torch.nn.Module, device: torch.device) -> None:
+    for submodule in module.modules():
+        for name, buffer in tuple(submodule.named_buffers(recurse=False)):
+            if buffer is not None and buffer.device != device:
+                submodule._buffers[name] = buffer.to(device=device)
+
+
+def _move_parameters_to_device(
+    parameters: set[torch.nn.Parameter],
+    device: torch.device,
+) -> None:
+    for param in parameters:
+        if param.device != device:
+            param.data = param.data.to(device=device)
+        if param.grad is not None and param.grad.device != device:
+            param.grad.data = param.grad.data.to(device=device)
+
+
+def _current_cuda_device() -> torch.device | None:
+    if not torch.cuda.is_available():
+        return None
+    return torch.device("cuda", torch.cuda.current_device())
+
+
+def apply_fsdp_inference(
+    transformer: torch.nn.Module,
+    mesh: DeviceMesh | None,
+) -> FSDPInferenceInfo:
+    if mesh is None:
+        return FSDPInferenceInfo(enabled=False, world_size=1, wrapped_blocks=0, ignored_params=0)
+    if bool(getattr(transformer, "_lingbot_fsdp_inference_enabled", False)):
+        blocks = getattr(transformer, "blocks", ())
+        return FSDPInferenceInfo(
+            enabled=True,
+            world_size=int(mesh.size()),
+            wrapped_blocks=len(blocks),
+            ignored_params=int(getattr(transformer, "_lingbot_fsdp_inference_ignored_params", 0)),
+        )
+
+    if fully_shard is None:
+        raise RuntimeError("PyTorch composable FSDP is not importable.") from FSDP_IMPORT_ERROR
+
+    cuda_device = _current_cuda_device()
+    if cuda_device is not None:
+        _move_buffers_to_device(transformer, cuda_device)
+
+    dtype_counts: dict[torch.dtype, int] = {}
+    for param in transformer.parameters():
+        dtype_counts[param.dtype] = dtype_counts.get(param.dtype, 0) + param.numel()
+    if dtype_counts:
+        primary_dtype = max(dtype_counts.items(), key=lambda item: item[1])[0]
+        ignored_params = {param for param in transformer.parameters() if param.dtype != primary_dtype}
+    else:
+        ignored_params = set()
+    if cuda_device is not None and ignored_params:
+        _move_parameters_to_device(ignored_params, cuda_device)
+
+    blocks: Any = getattr(transformer, "blocks", ())
+    wrapped_blocks = 0
+    for block in blocks:
+        block_ignored_params = {
+            param for param in block.parameters() if param in ignored_params
+        }
+        fully_shard(block, mesh=mesh, ignored_params=block_ignored_params)
+        wrapped_blocks += 1
+    fully_shard(transformer, mesh=mesh, ignored_params=ignored_params)
+    gc.collect()
+    if cuda_device is not None:
+        torch.cuda.empty_cache()
+    transformer._lingbot_fsdp_inference_enabled = True
+    transformer._lingbot_fsdp_inference_ignored_params = len(ignored_params)
+    return FSDPInferenceInfo(
+        enabled=True,
+        world_size=int(mesh.size()),
+        wrapped_blocks=wrapped_blocks,
+        ignored_params=len(ignored_params),
+    )

lingbot_video/inference_backend.py

@@ -0,0 +1,114 @@
+from __future__ import annotations
+
+import importlib.util
+import json
+import sys
+from pathlib import Path
+from typing import Sequence, TextIO
+
+
+BACKEND_TO_ENGINE = {
+    "diffusers": "diffusers",
+    "sglang": "sglang-native",
+}
+ENGINE_CHOICES = frozenset(BACKEND_TO_ENGINE.values())
+
+
+def _option_present(argv: Sequence[str], option: str) -> bool:
+    return any(arg == option or arg.startswith(f"{option}=") for arg in argv)
+
+
+def default_backend_argv(argv: Sequence[str]) -> list[str]:
+    """Default the public CLI to direct diffusers when no backend is requested."""
+
+    normalized = list(argv)
+    if _option_present(normalized, "--backend") or _option_present(normalized, "--engine"):
+        return normalized
+    return ["--backend", "diffusers", *normalized]
+
+
+def _module_available(name: str) -> bool:
+    if name in sys.modules:
+        return True
+    try:
+        return importlib.util.find_spec(name) is not None
+    except (ImportError, ModuleNotFoundError, KeyError, ValueError):
+        return False
+
+
+def sglang_native_available() -> bool:
+    """Return whether the local SGLang-native adapter can be imported."""
+
+    if not _module_available("sglang"):
+        return False
+    try:
+        from lingbot_video.native_backend import LingBotVideoNativePipeline
+    except Exception:
+        return False
+    return LingBotVideoNativePipeline is not None
+
+
+def resolve_backend_engine(
+    *,
+    engine: str | None,
+    backend: str | None,
+    sglang_available: bool | None = None,
+    stderr: TextIO | None = None,
+) -> str:
+    """Resolve public ``--backend`` and internal ``--engine`` into a runner engine."""
+
+    explicit_engine = backend is None and engine is not None
+    if backend is not None:
+        if backend not in BACKEND_TO_ENGINE:
+            choices = ", ".join(sorted(BACKEND_TO_ENGINE))
+            raise ValueError(f"unsupported backend: {backend!r}; choices: {choices}")
+        resolved = BACKEND_TO_ENGINE[backend]
+    elif engine is not None:
+        if engine not in ENGINE_CHOICES:
+            choices = ", ".join(sorted(ENGINE_CHOICES))
+            raise ValueError(f"unsupported engine: {engine!r}; choices: {choices}")
+        resolved = engine
+    else:
+        resolved = "sglang-native"
+
+    if resolved != "sglang-native":
+        return resolved
+    if explicit_engine:
+        return resolved
+
+    available = sglang_available
+    if available is None:
+        available = sglang_native_available()
+    if available:
+        return resolved
+
+    stream = stderr if stderr is not None else sys.stderr
+    print(
+        "WARNING: SGLang backend requested but SGLang is not installed or its "
+        "native diffusion API is unavailable; falling back to diffusers. "
+        "Install requirements-sglang.txt to enable SGLang Diffusion.",
+        file=stream,
+        flush=True,
+    )
+    return "diffusers"
+
+
+def load_negative_prompt_json(path: str | Path) -> str:
+    """Load an auto-negative JSON file and serialize it as the prompt string."""
+
+    with Path(path).open(encoding="utf-8") as f:
+        payload = json.load(f)
+    return json.dumps(payload, ensure_ascii=False, separators=(",", ":"))
+
+
+def resolve_negative_prompt_arg(
+    negative_prompt: str | None,
+    negative_prompt_json: str | None,
+) -> str | None:
+    """Resolve manual and auto-negative prompt inputs for the runner."""
+
+    if negative_prompt and negative_prompt_json:
+        raise ValueError("Use either --negative_prompt or --negative_prompt_json, not both.")
+    if negative_prompt_json:
+        return load_negative_prompt_json(negative_prompt_json)
+    return negative_prompt

lingbot_video/model_paths.py

@@ -0,0 +1,28 @@
+from __future__ import annotations
+
+from pathlib import Path
+from typing import Any
+
+
+def model_component_dir(model_dir: str | Path, component: str) -> Path:
+    """Return a named component directory under a public model root."""
+
+    path = Path(model_dir) / component
+    if not path.is_dir():
+        raise FileNotFoundError(f"missing model component `{component}`: {path}")
+    return path
+
+
+def effective_refiner_model_dir(args: Any) -> Path | None:
+    """Resolve the model root to use for refiner loading.
+
+    Passing --run_refiner means the public runner should load `refiner/` from
+    the same model root by default. --refiner_model_dir remains an override for
+    nonstandard package layouts.
+    """
+
+    refiner_model_dir = getattr(args, "refiner_model_dir", None)
+    requested = bool(getattr(args, "run_refiner", False) or refiner_model_dir)
+    if not requested:
+        return None
+    return Path(refiner_model_dir or getattr(args, "model_dir"))

lingbot_video/moe_pack_kernels.py

@@ -0,0 +1,142 @@
+from __future__ import annotations
+
+import torch
+
+try:
+    import triton
+    import triton.language as tl
+except Exception:  # pragma: no cover - Triton is optional outside GPU deployments.
+    triton = None
+    tl = None
+
+
+if triton is not None and tl is not None:
+    @triton.jit
+    def _moe_route_count_slots_kernel(
+        flat_indices,
+        counts,
+        route_slots,
+        num_routes: tl.constexpr,
+        block_m: tl.constexpr,
+    ):
+        offsets = tl.program_id(0) * block_m + tl.arange(0, block_m)
+        mask = offsets < num_routes
+        experts = tl.load(flat_indices + offsets, mask=mask, other=0)
+        slots = tl.atomic_add(counts + experts, 1, sem="relaxed", mask=mask)
+        tl.store(route_slots + offsets, slots, mask=mask)
+
+    @triton.jit
+    def _moe_pack_tokens_kernel(
+        tokens,
+        flat_scores,
+        flat_indices,
+        offsets,
+        route_slots,
+        permuted_tokens,
+        sorted_positions,
+        sorted_scores,
+        hidden_size: tl.constexpr,
+        top_k: tl.constexpr,
+        block_h: tl.constexpr,
+    ):
+        route_idx = tl.program_id(0)
+        hidden_block = tl.program_id(1)
+        offsets_h = hidden_block * block_h + tl.arange(0, block_h)
+        hidden_mask = offsets_h < hidden_size
+
+        expert_idx = tl.load(flat_indices + route_idx)
+        slot = tl.load(route_slots + route_idx)
+        expert_offset = tl.load(offsets + expert_idx)
+        dest_idx = expert_offset + slot
+        token_idx = route_idx // top_k
+
+        values = tl.load(
+            tokens + token_idx * hidden_size + offsets_h,
+            mask=hidden_mask,
+            other=0.0,
+        )
+        tl.store(
+            permuted_tokens + dest_idx * hidden_size + offsets_h,
+            values,
+            mask=hidden_mask,
+        )
+
+        first_hidden_block = hidden_block == 0
+        tl.store(sorted_positions + dest_idx, route_idx, mask=first_hidden_block)
+        score = tl.load(flat_scores + route_idx, mask=first_hidden_block, other=0.0)
+        tl.store(sorted_scores + dest_idx, score, mask=first_hidden_block)
+else:
+    _moe_route_count_slots_kernel = None
+    _moe_pack_tokens_kernel = None
+
+
+def reorder_tokens_triton_pack(
+    tokens: torch.Tensor,
+    top_scores: torch.Tensor,
+    top_indices: torch.Tensor,
+    num_experts: int,
+    block_h: int = 64,
+    block_m: int = 256,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, int, int]:
+    if _moe_route_count_slots_kernel is None or _moe_pack_tokens_kernel is None:
+        raise RuntimeError("LINGBOT_MOE_REORDER_BACKEND=triton_pack requires Triton")
+    if tokens.device.type != "cuda":
+        raise RuntimeError("LINGBOT_MOE_REORDER_BACKEND=triton_pack requires CUDA tensors")
+    if tokens.ndim != 2:
+        raise ValueError(f"Expected 2D tokens, got {tokens.ndim}D")
+    if top_scores.shape != top_indices.shape:
+        raise ValueError("top_scores and top_indices must have the same shape")
+    if top_indices.ndim != 2:
+        raise ValueError(f"Expected 2D top_indices, got {top_indices.ndim}D")
+
+    num_tokens = tokens.shape[0]
+    hidden_size = tokens.shape[1]
+    top_k = top_indices.shape[1]
+    num_routes = top_indices.numel()
+    counts = torch.zeros(num_experts, dtype=torch.int32, device=tokens.device)
+
+    if num_routes == 0:
+        return (
+            tokens.new_empty((0, hidden_size)),
+            counts,
+            torch.empty(0, dtype=torch.int64, device=tokens.device),
+            top_scores.new_empty((0,)),
+            num_tokens,
+            top_k,
+        )
+
+    tokens = tokens.contiguous()
+    flat_scores = top_scores.contiguous().reshape(-1)
+    flat_indices = top_indices.contiguous().to(torch.int32).reshape(-1)
+    route_slots = torch.empty(num_routes, dtype=torch.int32, device=tokens.device)
+
+    _moe_route_count_slots_kernel[(triton.cdiv(num_routes, block_m),)](
+        flat_indices,
+        counts,
+        route_slots,
+        num_routes,
+        block_m,
+    )
+
+    counts_i64 = counts.to(torch.int64)
+    offsets = torch.empty(num_experts + 1, dtype=torch.int64, device=tokens.device)
+    offsets[0] = 0
+    offsets[1:] = torch.cumsum(counts_i64, dim=0)
+
+    permuted_tokens = torch.empty((num_routes, hidden_size), dtype=tokens.dtype, device=tokens.device)
+    sorted_positions = torch.empty(num_routes, dtype=torch.int64, device=tokens.device)
+    sorted_scores = torch.empty(num_routes, dtype=top_scores.dtype, device=tokens.device)
+    _moe_pack_tokens_kernel[(num_routes, triton.cdiv(hidden_size, block_h))](
+        tokens,
+        flat_scores,
+        flat_indices,
+        offsets,
+        route_slots,
+        permuted_tokens,
+        sorted_positions,
+        sorted_scores,
+        hidden_size,
+        top_k,
+        block_h,
+    )
+    return permuted_tokens, counts, sorted_positions, sorted_scores, num_tokens, top_k

lingbot_video/moe_restore_kernels.py

@@ -0,0 +1,92 @@
+from __future__ import annotations
+
+import torch
+
+try:
+    import triton
+    import triton.language as tl
+except Exception:  # pragma: no cover - Triton is optional outside GPU deployments.
+    triton = None
+    tl = None
+
+
+if triton is not None and tl is not None:
+    @triton.jit
+    def _moe_restore_weighted_sum_kernel(
+        expert_output,
+        sorted_scores,
+        route_for_position,
+        output,
+        hidden_size: tl.constexpr,
+        top_k: tl.constexpr,
+        block_h: tl.constexpr,
+    ):
+        token_idx = tl.program_id(0)
+        hidden_block = tl.program_id(1)
+        offsets_h = hidden_block * block_h + tl.arange(0, block_h)
+        hidden_mask = offsets_h < hidden_size
+        acc = tl.zeros((block_h,), dtype=tl.float32)
+
+        for route_slot in tl.range(0, top_k):
+            route_pos = token_idx * top_k + route_slot
+            route_idx = tl.load(route_for_position + route_pos)
+            active = route_idx >= 0
+            safe_route_idx = tl.maximum(route_idx, 0)
+            values = tl.load(
+                expert_output + safe_route_idx * hidden_size + offsets_h,
+                mask=active & hidden_mask,
+                other=0.0,
+            ).to(tl.float32)
+            score = tl.load(sorted_scores + safe_route_idx, mask=active, other=0.0).to(tl.float32)
+            acc += values * score
+
+        tl.store(output + token_idx * hidden_size + offsets_h, acc, mask=hidden_mask)
+else:
+    _moe_restore_weighted_sum_kernel = None
+
+
+def restore_tokens_triton(
+    expert_output: torch.Tensor,
+    sorted_positions: torch.Tensor,
+    sorted_scores: torch.Tensor,
+    num_tokens: int,
+    top_k: int,
+    block_h: int = 64,
+) -> torch.Tensor:
+    if _moe_restore_weighted_sum_kernel is None:
+        raise RuntimeError("LINGBOT_MOE_RESTORE_BACKEND=triton requires Triton")
+    if expert_output.ndim != 2:
+        raise ValueError(f"Expected 2D expert_output, got {expert_output.ndim}D")
+    if sorted_positions.numel() != sorted_scores.numel():
+        raise ValueError("sorted_positions and sorted_scores must have the same length")
+    if sorted_positions.numel() == 0:
+        return expert_output.new_zeros((num_tokens, expert_output.shape[-1]))
+
+    hidden_size = expert_output.shape[-1]
+    route_for_position = torch.full(
+        (num_tokens * top_k,),
+        -1,
+        dtype=torch.int32,
+        device=expert_output.device,
+    )
+    route_for_position[sorted_positions] = torch.arange(
+        sorted_positions.numel(),
+        dtype=torch.int32,
+        device=expert_output.device,
+    )
+    output = torch.empty(
+        (num_tokens, hidden_size),
+        dtype=expert_output.dtype,
+        device=expert_output.device,
+    )
+    grid = (num_tokens, triton.cdiv(hidden_size, block_h))
+    _moe_restore_weighted_sum_kernel[grid](
+        expert_output.contiguous(),
+        sorted_scores.contiguous(),
+        route_for_position,
+        output,
+        hidden_size,
+        top_k,
+        block_h,
+    )
+    return output

lingbot_video/native_backend.py

@@ -0,0 +1,347 @@
+from __future__ import annotations
+
+import os
+from contextlib import contextmanager
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any
+
+import torch
+from diffusers import DiffusionPipeline
+
+from lingbot_video.model_paths import model_component_dir
+from lingbot_video.pipeline_lingbot_video import (
+    DEFAULT_NEGATIVE_PROMPT,
+    LingBotVideoPipeline,
+    LingBotVideoPipelineOutput,
+)
+from lingbot_video.pipeline_lingbot_video_i2v import LingBotVideoImageToVideoPipeline
+from lingbot_video.transformer_lingbot_video import LingBotVideoTransformer3DModel
+
+try:
+    from sglang.multimodal_gen import registry as sglang_registry
+except Exception:  # pragma: no cover - SGLang is an optional deployment dep
+    sglang_registry = None
+
+try:
+    from sglang.multimodal_gen.runtime.server_args import (
+        Backend,
+        ServerArgs,
+        set_global_server_args,
+    )
+except Exception:  # pragma: no cover - SGLang is an optional deployment dep
+    Backend = None
+    ServerArgs = None
+    set_global_server_args = None
+
+
+@dataclass
+class LingBotVideoNativePipelineConfig:
+    """Minimal config object for the LingBotVideo native adapter."""
+
+    flow_shift: float = 3.0
+
+
+@dataclass
+class LingBotVideoNativeSamplingParams:
+    """Default sampling parameters for the native pipeline."""
+
+    prompt: str | None = None
+    negative_prompt: str = DEFAULT_NEGATIVE_PROMPT
+    height: int = 480
+    width: int = 832
+    num_frames: int = 81
+    fps: int = 24
+    num_inference_steps: int = 40
+    guidance_scale: float = 3.0
+    seed: int = 42
+    shift: float = 3.0
+
+
+@contextmanager
+def _patch_qwen3vl_from_pretrained():
+    try:
+        from transformers import Qwen3VLForConditionalGeneration
+    except Exception:
+        yield
+        return
+
+    original_from_pretrained = Qwen3VLForConditionalGeneration.from_pretrained
+    attn_implementation = os.environ.get("LINGBOT_QWEN_ATTN_IMPLEMENTATION", "flash_attention_3")
+
+    @classmethod
+    def patched_from_pretrained(cls, pretrained_model_name_or_path, *args, **kwargs):
+        if attn_implementation:
+            kwargs.setdefault("attn_implementation", attn_implementation)
+        if "torch_dtype" in kwargs and "dtype" not in kwargs:
+            kwargs["dtype"] = kwargs.pop("torch_dtype")
+        return original_from_pretrained(pretrained_model_name_or_path, *args, **kwargs)
+
+    Qwen3VLForConditionalGeneration.from_pretrained = patched_from_pretrained
+    try:
+        yield
+    finally:
+        Qwen3VLForConditionalGeneration.from_pretrained = original_from_pretrained
+
+
+def _module_dtype_name(module: Any) -> str | None:
+    if module is None:
+        return None
+    dtype = getattr(module, "dtype", None)
+    if isinstance(dtype, torch.dtype):
+        return str(dtype).replace("torch.", "")
+    if isinstance(module, torch.nn.Module):
+        try:
+            return str(next(module.parameters()).dtype).replace("torch.", "")
+        except StopIteration:
+            return None
+    return None
+
+
+def _default_device() -> torch.device:
+    if not torch.cuda.is_available():
+        return torch.device("cpu")
+    return torch.device("cuda", torch.cuda.current_device())
+
+
+def _pipeline_class_for_mode(mode: str) -> type[DiffusionPipeline]:
+    if mode == "ti2v":
+        return LingBotVideoImageToVideoPipeline
+    return LingBotVideoPipeline
+
+
+def _load_lingbot_diffusers_pipe(
+    model_path: str | Path,
+    dtype_map: dict[str, torch.dtype] | torch.dtype | None,
+    mode: str,
+    transformer_subfolder: str,
+) -> DiffusionPipeline:
+    model_root = Path(model_path)
+    model_component_dir(model_root, transformer_subfolder)
+    transformer_dtype = (
+        dtype_map.get("transformer", dtype_map.get("default", torch.float32))
+        if isinstance(dtype_map, dict)
+        else dtype_map
+    )
+    transformer = LingBotVideoTransformer3DModel.from_pretrained(
+        str(model_root),
+        subfolder=transformer_subfolder,
+        torch_dtype=transformer_dtype,
+    )
+    load_kwargs: dict[str, Any] = {
+        "transformer": transformer,
+        "trust_remote_code": True,
+    }
+    if dtype_map is not None:
+        load_kwargs["torch_dtype"] = dtype_map
+    with _patch_qwen3vl_from_pretrained():
+        pipe = _pipeline_class_for_mode(mode).from_pretrained(
+            str(model_root),
+            **load_kwargs,
+        )
+    return pipe.to(_default_device())
+
+
+class LingBotVideoNativeExecutionStage:
+    """Single-stage LingBotVideo execution."""
+
+    def __init__(self, diffusers_pipe: DiffusionPipeline):
+        self.diffusers_pipe = diffusers_pipe
+
+    def __call__(self, **kwargs: Any) -> LingBotVideoPipelineOutput:
+        return self.forward(**kwargs)
+
+    @torch.no_grad()
+    def forward(self, **kwargs: Any) -> LingBotVideoPipelineOutput:
+        output = self.diffusers_pipe(**kwargs)
+        frames = output.frames if hasattr(output, "frames") else output[0]
+        return LingBotVideoPipelineOutput(frames=frames)
+
+
+class LingBotVideoNativePipeline:
+    """SGLang-native LingBotVideo pipeline adapter."""
+
+    pipeline_name = "LingBotVideoNativePipeline"
+    pipeline_config_cls = LingBotVideoNativePipelineConfig
+    sampling_params_cls = LingBotVideoNativeSamplingParams
+    is_video_pipeline = True
+
+    def __init__(
+        self,
+        diffusers_pipe: DiffusionPipeline | str | Path,
+        *,
+        model_path: str | Path | None = None,
+        server_args: Any | None = None,
+        dtype_map: dict[str, torch.dtype] | torch.dtype | None = None,
+        mode: str = "t2v",
+        transformer_subfolder: str = "transformer",
+        **_: Any,
+    ):
+        if isinstance(diffusers_pipe, (str, Path)):
+            model_path = diffusers_pipe if model_path is None else model_path
+            diffusers_pipe = _load_lingbot_diffusers_pipe(
+                model_path,
+                dtype_map,
+                mode=mode,
+                transformer_subfolder=transformer_subfolder,
+            )
+        if model_path is None:
+            raise ValueError("model_path is required for LingBotVideoNativePipeline")
+        self.diffusers_pipe = diffusers_pipe
+        self.model_path = str(model_path)
+        self.server_args = server_args
+        self.execution_stage = LingBotVideoNativeExecutionStage(diffusers_pipe)
+        self.modules = {"lingbot_video_pipeline": diffusers_pipe}
+        self.memory_usages: dict[str, float] = {}
+
+    @property
+    def transformer(self) -> Any:
+        return getattr(self.diffusers_pipe, "transformer", None)
+
+    @property
+    def text_encoder(self) -> Any:
+        return getattr(self.diffusers_pipe, "text_encoder", None)
+
+    @property
+    def vae(self) -> Any:
+        return getattr(self.diffusers_pipe, "vae", None)
+
+    @property
+    def scheduler(self) -> Any:
+        return getattr(self.diffusers_pipe, "scheduler", None)
+
+    def encode_prompt(self, *args: Any, **kwargs: Any) -> Any:
+        return self.diffusers_pipe.encode_prompt(*args, **kwargs)
+
+    def encode_video_latent(self, *args: Any, **kwargs: Any) -> Any:
+        return self.diffusers_pipe.encode_video_latent(*args, **kwargs)
+
+    @classmethod
+    def from_diffusers_pipe(
+        cls,
+        diffusers_pipe: DiffusionPipeline,
+        *,
+        model_path: str | Path,
+        server_args: Any | None = None,
+    ) -> "LingBotVideoNativePipeline":
+        return cls(diffusers_pipe, model_path=model_path, server_args=server_args)
+
+    @classmethod
+    def from_pretrained(
+        cls,
+        model_path: str | Path,
+        *,
+        torch_dtype: dict[str, torch.dtype] | torch.dtype | None = None,
+        server_args: Any | None = None,
+        mode: str = "t2v",
+        transformer_subfolder: str = "transformer",
+        **kwargs: Any,
+    ) -> "LingBotVideoNativePipeline":
+        return cls(
+            model_path,
+            model_path=model_path,
+            server_args=server_args,
+            dtype_map=torch_dtype,
+            mode=mode,
+            transformer_subfolder=transformer_subfolder,
+            **kwargs,
+        )
+
+    @torch.no_grad()
+    def __call__(
+        self,
+        *,
+        prompt: str,
+        negative_prompt: str = DEFAULT_NEGATIVE_PROMPT,
+        height: int = 480,
+        width: int = 832,
+        num_frames: int = 81,
+        num_inference_steps: int = 40,
+        guidance_scale: float = 3.0,
+        shift: float = 3.0,
+        generator: torch.Generator | None = None,
+        image: Any | None = None,
+        output_type: str = "np",
+        **kwargs: Any,
+    ) -> LingBotVideoPipelineOutput:
+        call_kwargs = dict(
+            prompt=prompt,
+            negative_prompt=negative_prompt,
+            height=height,
+            width=width,
+            num_frames=num_frames,
+            num_inference_steps=num_inference_steps,
+            guidance_scale=guidance_scale,
+            shift=shift,
+            generator=generator,
+            output_type=output_type,
+        )
+        if image is not None:
+            call_kwargs["image"] = image
+        call_kwargs.update(kwargs)
+        return self.execution_stage(**call_kwargs)
+
+    def component_dtypes(self) -> dict[str, str | None]:
+        return {
+            "transformer": _module_dtype_name(self.transformer),
+            "text_encoder": _module_dtype_name(self.text_encoder),
+            "vae": _module_dtype_name(self.vae),
+        }
+
+
+def register_lingbot_native_pipeline() -> bool:
+    """Register the adapter in SGLang's native registry when SGLang is present."""
+
+    if sglang_registry is None:
+        return False
+
+    try:
+        sglang_registry._PIPELINE_REGISTRY[LingBotVideoNativePipeline.pipeline_name] = (  # type: ignore[attr-defined]
+            LingBotVideoNativePipeline
+        )
+        sglang_registry._PIPELINE_CONFIG_REGISTRY[LingBotVideoNativePipeline.pipeline_name] = (  # type: ignore[attr-defined]
+            LingBotVideoNativePipelineConfig,
+            LingBotVideoNativeSamplingParams,
+        )
+    except Exception:
+        return False
+    return True
+
+
+def load_lingbot_native_pipeline(
+    model_dir: Path,
+    dtype_map: dict[str, torch.dtype],
+    mode: str = "t2v",
+    transformer_subfolder: str = "transformer",
+) -> LingBotVideoNativePipeline:
+    """Load a LingBotVideo native adapter from a diffusers-format model dir."""
+
+    server_args = None
+    if Backend is not None and ServerArgs is not None and set_global_server_args is not None:
+        try:
+            server_args = ServerArgs(
+                model_path=str(model_dir),
+                backend=Backend.SGLANG,
+                trust_remote_code=True,
+                pipeline_class_name=LingBotVideoNativePipeline.pipeline_name,
+                pipeline_config=LingBotVideoNativePipelineConfig(),
+                dit_cpu_offload=False,
+                dit_layerwise_offload=False,
+                layerwise_offload_components=[],
+                text_encoder_cpu_offload=False,
+                image_encoder_cpu_offload=False,
+                vae_cpu_offload=False,
+            )
+            set_global_server_args(server_args)
+        except Exception:
+            server_args = None
+
+    register_lingbot_native_pipeline()
+
+    return LingBotVideoNativePipeline.from_pretrained(
+        model_path=model_dir,
+        torch_dtype=dtype_map,
+        server_args=server_args,
+        mode=mode,
+        transformer_subfolder=transformer_subfolder,
+    )

lingbot_video/pipeline_lingbot_video.py

@@ -0,0 +1,591 @@
+from __future__ import annotations
+
+from contextlib import nullcontext
+from dataclasses import dataclass
+from typing import Any, Dict, List, Optional, Tuple, Union
+
+import numpy as np
+import torch
+import torch.distributed as dist
+from torchvision.transforms.functional import normalize as normalize_image_tensor
+
+from diffusers import DiffusionPipeline
+from diffusers.utils import BaseOutput
+from diffusers.utils.torch_utils import randn_tensor
+
+from .utils import (
+    LOW_NOISE_TAIL_V1_DEFAULT_STEPS,
+    batch_cfg_prompt_inputs,
+    compute_refiner_sigmas,
+)
+from .scheduling_flow_unipc import FlowUniPCMultistepScheduler
+
+
+TOKEN_LENGTH = 37698
+HIDDEN_STATE_SKIP_LAYER = 0
+
+PROMPT_TEMPLATE = (
+    "<|im_start|>system\nGiven a user input that may include a text prompt alone, "
+    "a text prompt with an image reference, or a text prompt with a video reference "
+    "or a video reference alone, generate an \"Enhanced prompt\" that provides detailed "
+    "visual descriptions suitable for video generation. Evaluate the level of detail "
+    "in the user's input: if it is simple, enrich it by adding specifics about colors, "
+    "shapes, sizes, textures, lighting, motion dynamics, camera movement, temporal "
+    "progression, and spatial relationships to create vivid, concrete, and temporally "
+    "coherent scenes to create vivid and concrete scenes. Please generate only the "
+    "enhanced description for the prompt below and avoid including any additional "
+    "commentary or evaluations:<|im_end|>\n<|im_start|>user\n{}<|im_end|>\n"
+    "<|im_start|>assistant\n"
+)
+IMG_PROMPT_TEMPLATE = "<|vision_start|><|image_pad|><|vision_end|>"
+VIDEO_PROMPT_TEMPLATE = "<|vision_start|><|video_pad|><|vision_end|>"
+
+DEFAULT_NEGATIVE_PROMPT = (
+    '{"universal_negative": {"visual_quality": ["low quality", "worst quality", "blurry", "pixelated", "jpeg artifacts", "low resolution", "unstable color", "color flicker", "underexposed", "overexposed", "invisible subject", "subject hidden in darkness"], "artistic_style": ["painting", "illustration", "drawing", "cartoon", "3d render", "cgi", "sketch", "digital art"], "composition_and_content": ["text", "watermark", "signature", "logo", "subtitles", "pillarboxed", "side bars", "portrait image in landscape frame"], "temporal_and_motion_stability": ["flickering", "jittery", "motion blur", "temporal inconsistency", "warping", "morphing", "incoherent motion", "unnatural movement", "static object with sudden jump", "frame-to-frame inconsistency"], "material_and_structure": ["plastic-like glass", "unrealistic texture", "deformed bottle", "liquid freezing improperly", "distorted reflections"]}}'
+)
+
+# Still-image default (t2i): drops the whole temporal/motion block and the video-only
+# codec/temporal terms that cannot apply to a single frame.
+DEFAULT_NEGATIVE_PROMPT_IMAGE = (
+    '{"universal_negative": {"visual_quality": ["low quality", "worst quality", "blurry", "pixelated", "jpeg artifacts", "low resolution", "underexposed", "overexposed", "invisible subject", "subject hidden in darkness"], "artistic_style": ["painting", "illustration", "drawing", "cartoon", "3d render", "cgi", "sketch", "digital art"], "composition_and_content": ["text", "watermark", "signature", "logo", "pillarboxed", "side bars", "portrait image in landscape frame"], "material_and_structure": ["plastic-like glass", "unrealistic texture", "deformed bottle", "distorted reflections"]}}'
+)
+
+
+@dataclass
+class LingBotVideoPipelineOutput(BaseOutput):
+    frames: Union[List[np.ndarray], torch.Tensor]
+
+
+def _module_dtype(module: torch.nn.Module) -> torch.dtype:
+    try:
+        return next(module.parameters()).dtype
+    except StopIteration:
+        return torch.float32
+
+
+def _transformer_timestep(timestep: torch.Tensor, transformer_dtype: torch.dtype) -> torch.Tensor:
+    sigma = timestep.float() / 1000.0
+    if transformer_dtype in {torch.bfloat16, torch.float16}:
+        sigma = sigma.to(transformer_dtype)
+    return (sigma * 1000.0).float()
+
+
+def _transformer_autocast(device: torch.device, transformer_dtype: torch.dtype):
+    if device.type != "cuda" or transformer_dtype not in {torch.bfloat16, torch.float16}:
+        return nullcontext()
+    return torch.autocast(device_type="cuda", dtype=transformer_dtype)
+
+
+def _module_device(module: torch.nn.Module) -> torch.device:
+    try:
+        return next(module.parameters()).device
+    except StopIteration:
+        return torch.device("cpu")
+
+
+def _group_global_rank(group: Optional[Any], group_rank: int) -> int:
+    if group is None:
+        return group_rank
+    get_global_rank = getattr(dist, "get_global_rank", None)
+    if get_global_rank is None:
+        return group_rank
+    return int(get_global_rank(group, group_rank))
+
+
+class LingBotVideoPipeline(DiffusionPipeline):
+    """Minimal LingBotVideo t2v/t2i pipeline.
+
+    Standard CFG runs as two independent transformer forwards unless batched CFG
+    or CFG parallelism is explicitly requested.
+    """
+
+    model_cpu_offload_seq = "text_encoder->transformer->vae"
+
+    def __init__(self, transformer, vae, text_encoder, processor, scheduler):
+        super().__init__()
+        if (
+            scheduler is not None
+            and scheduler.__class__.__name__ != FlowUniPCMultistepScheduler.__name__
+        ):
+            raise TypeError(
+                "LingBotVideoPipeline requires vendored FlowUniPCMultistepScheduler; "
+                f"got {scheduler.__class__.__name__}."
+            )
+        self.register_modules(
+            transformer=transformer,
+            vae=vae,
+            text_encoder=text_encoder,
+            processor=processor,
+            scheduler=scheduler,
+        )
+        self.vae_scale_factor_temporal = 4
+        self.vae_scale_factor_spatial = 8
+        self.token_length = TOKEN_LENGTH
+        self.hidden_state_skip_layer = HIDDEN_STATE_SKIP_LAYER
+        self.prompt_template = PROMPT_TEMPLATE
+        self.img_prompt_template = IMG_PROMPT_TEMPLATE
+        self.video_prompt_template = VIDEO_PROMPT_TEMPLATE
+        self._crop_start: Optional[int] = None
+
+    @staticmethod
+    def check_inputs(height: int, width: int, num_frames: int) -> None:
+        if num_frames != 1 and (num_frames - 1) % 4 != 0:
+            raise ValueError(f"`num_frames` must be 1 or 4n+1, got {num_frames}.")
+        if height % 16 != 0 or width % 16 != 0:
+            raise ValueError(f"`height` and `width` must be multiples of 16, got {height}x{width}.")
+
+    @staticmethod
+    def _apply_inpainting(latents: torch.Tensor, cond_latent: torch.Tensor) -> torch.Tensor:
+        cond_t = cond_latent.shape[2]
+        latents[:, :, :cond_t] = cond_latent.float()
+        return latents
+
+    @staticmethod
+    def apply_text_to_template(text: str, template: str = PROMPT_TEMPLATE) -> str:
+        return template.format(text)
+
+    def _compute_crop_start(self) -> int:
+        if self._crop_start is None:
+            marker = "<|USER_INPUT_MARKER|>"
+            marked = self.prompt_template.format(marker)
+            marker_pos = marked.find(marker)
+            if marker_pos < 0:
+                self._crop_start = 0
+            else:
+                prefix = self.processor(
+                    text=marked[:marker_pos],
+                    images=None,
+                    videos=None,
+                    return_tensors="pt",
+                )
+                self._crop_start = int(prefix["input_ids"].shape[1])
+        return self._crop_start
+
+    def _build_prompt_inputs(
+        self,
+        prompt: Union[str, List[str]],
+        images: Optional[Any] = None,
+        videos: Optional[Any] = None,
+        video_metadata: Optional[Any] = None,
+        video_kwargs: Optional[Dict[str, Any]] = None,
+    ):
+        if isinstance(prompt, str):
+            prompts = [prompt]
+        else:
+            prompts = list(prompt)
+
+        visual_template = ""
+        if images is not None:
+            visual_template = self.img_prompt_template
+        elif videos is not None:
+            visual_template = self.video_prompt_template
+
+        texts = [
+            self.apply_text_to_template(visual_template + text, self.prompt_template)
+            for text in prompts
+        ]
+        kwargs = dict(video_kwargs or {})
+        return self.processor(
+            text=texts,
+            images=images,
+            videos=videos,
+            video_metadata=video_metadata,
+            do_resize=False,
+            truncation=True,
+            max_length=self.token_length,
+            padding="longest",
+            return_tensors="pt",
+            **kwargs,
+        )
+
+    @torch.no_grad()
+    def encode_prompt(
+        self,
+        prompt: Union[str, List[str]],
+        *,
+        images: Optional[Any] = None,
+        videos: Optional[Any] = None,
+        video_metadata: Optional[Any] = None,
+        video_kwargs: Optional[Dict[str, Any]] = None,
+        device: Optional[Union[str, torch.device]] = None,
+        return_inputs: bool = False,
+    ):
+        if self.text_encoder is None or self.processor is None:
+            raise ValueError("`text_encoder` and `processor` are required for encode_prompt().")
+
+        device = torch.device(device) if device is not None else self._execution_device
+        inputs = self._build_prompt_inputs(
+            prompt,
+            images=images,
+            videos=videos,
+            video_metadata=video_metadata,
+            video_kwargs=video_kwargs,
+        )
+        inputs = inputs.to(device)
+        outputs = self.text_encoder(
+            **inputs,
+            output_hidden_states=self.hidden_state_skip_layer is not None,
+        )
+        if self.hidden_state_skip_layer is not None:
+            prompt_embeds = outputs.hidden_states[-(self.hidden_state_skip_layer + 1)]
+        else:
+            prompt_embeds = outputs.last_hidden_state
+
+        prompt_mask = inputs["attention_mask"]
+        crop_start = self._compute_crop_start()
+        if crop_start > 0:
+            prompt_embeds = prompt_embeds[:, crop_start:]
+            prompt_mask = prompt_mask[:, crop_start:]
+
+        # Batch=1 can drop right padding before DiT inference.
+        if prompt_embeds.shape[0] == 1:
+            true_len = int(prompt_mask[0].sum().item())
+            prompt_embeds = prompt_embeds[:, :true_len]
+            prompt_mask = prompt_mask[:, :true_len]
+
+        if return_inputs:
+            return prompt_embeds, prompt_mask, inputs
+        return prompt_embeds, prompt_mask
+
+    def prepare_latents(
+        self,
+        num_frames: int,
+        height: int,
+        width: int,
+        generator: Optional[torch.Generator],
+        latents: Optional[torch.Tensor],
+        device: torch.device,
+    ) -> torch.Tensor:
+        latent_frames = (num_frames - 1) // self.vae_scale_factor_temporal + 1
+        latent_height = height // self.vae_scale_factor_spatial
+        latent_width = width // self.vae_scale_factor_spatial
+        shape = (
+            1,
+            self.transformer.config.in_channels,
+            latent_frames,
+            latent_height,
+            latent_width,
+        )
+        if latents is None:
+            return randn_tensor(shape, generator=generator, device=device, dtype=torch.float32)
+        if tuple(latents.shape) != shape:
+            raise ValueError(f"`latents` shape must be {shape}, got {tuple(latents.shape)}.")
+        return latents.to(device=device, dtype=torch.float32)
+
+    def _dit_latent_to_vae(self, latents: torch.Tensor) -> torch.Tensor:
+        mean = torch.tensor(self.vae.config.latents_mean, device=latents.device, dtype=torch.float32)
+        std_inv = 1.0 / torch.tensor(
+            self.vae.config.latents_std, device=latents.device, dtype=torch.float32
+        )
+        mean = mean.view(1, -1, 1, 1, 1)
+        std_inv = std_inv.view(1, -1, 1, 1, 1)
+        return latents.float() / std_inv + mean
+
+    def _vae_latent_to_dit(self, latents: torch.Tensor) -> torch.Tensor:
+        mean = torch.tensor(self.vae.config.latents_mean, device=latents.device, dtype=torch.float32)
+        std_inv = 1.0 / torch.tensor(
+            self.vae.config.latents_std, device=latents.device, dtype=torch.float32
+        )
+        mean = mean.view(1, -1, 1, 1, 1)
+        std_inv = std_inv.view(1, -1, 1, 1, 1)
+        return (latents.float() - mean) * std_inv
+
+    @torch.no_grad()
+    def encode_video_latent(
+        self,
+        video: torch.Tensor,
+        generator: Optional[torch.Generator] = None,
+    ) -> torch.Tensor:
+        if self.vae is None:
+            raise ValueError("`vae` is required to encode video latents.")
+        vae_device = _module_device(self.vae)
+        video = video.to(device=vae_device, dtype=torch.float32)
+        bsz, channels, frames, height, width = video.shape
+        flat_video = video.permute(0, 2, 1, 3, 4).reshape(bsz * frames, channels, height, width)
+        norm_flat_video = normalize_image_tensor(
+            flat_video,
+            [0.5, 0.5, 0.5],
+            [0.5, 0.5, 0.5],
+            inplace=False,
+        )
+        norm_video = (
+            norm_flat_video.reshape(bsz, frames, channels, height, width)
+            .permute(0, 2, 1, 3, 4)
+            .contiguous()
+        )
+        with torch.autocast(
+            "cuda",
+            dtype=torch.bfloat16,
+            enabled=vae_device.type == "cuda",
+        ):
+            encoded = self.vae.encode(norm_video)
+        if hasattr(encoded, "latent_dist"):
+            latents = encoded.latent_dist.sample(generator)
+        else:
+            latents = encoded[0] if isinstance(encoded, tuple) else encoded
+        return self._vae_latent_to_dit(latents).to(latents)
+
+    @torch.no_grad()
+    def _decode_latents(
+        self,
+        latents: torch.Tensor,
+    ) -> List[np.ndarray]:
+        vae_device = _module_device(self.vae)
+        vae_dtype = _module_dtype(self.vae)
+        vae_latents = self._dit_latent_to_vae(latents).to(device=vae_device, dtype=torch.float32)
+        if vae_latents.ndim == 5:
+            vae_latents = vae_latents.contiguous(memory_format=torch.channels_last_3d)
+        autocast_dtype = (
+            vae_dtype
+            if vae_device.type == "cuda" and vae_dtype in {torch.bfloat16, torch.float16}
+            else None
+        )
+        with torch.autocast(
+            "cuda",
+            dtype=autocast_dtype or torch.bfloat16,
+            enabled=autocast_dtype is not None,
+        ):
+            decoded = self.vae.decode(vae_latents)
+        frames = decoded[0] if isinstance(decoded, tuple) else decoded.sample
+        frames = frames.float().clamp_(-1, 1)
+        frames = (frames + 1.0) / 2.0
+        frames = frames.permute(0, 2, 3, 4, 1).cpu().numpy()
+        return [video for video in frames]
+
+    @torch.no_grad()
+    def __call__(
+        self,
+        prompt: str,
+        negative_prompt: str = DEFAULT_NEGATIVE_PROMPT,
+        height: int = 480,
+        width: int = 480,
+        num_frames: int = 81,
+        num_inference_steps: int = 40,
+        guidance_scale: float = 6.0,
+        shift: float = 3.0,
+        generator: Optional[torch.Generator] = None,
+        latents: Optional[torch.Tensor] = None,
+        cond_latent: Optional[torch.Tensor] = None,
+        prompt_embeds: Optional[torch.Tensor] = None,
+        prompt_mask: Optional[torch.Tensor] = None,
+        negative_prompt_embeds: Optional[torch.Tensor] = None,
+        negative_prompt_mask: Optional[torch.Tensor] = None,
+        output_type: str = "np",
+        cfg_parallel_group: Optional[Any] = None,
+        batch_cfg: bool = False,
+        null_cond_clone_zero: bool = False,
+        t_thresh: Optional[float] = None,
+        refiner_sigma_tail_steps: int = LOW_NOISE_TAIL_V1_DEFAULT_STEPS,
+        offload_vae_during_denoise: bool = False,
+        return_dict: bool = True,
+    ) -> Union[LingBotVideoPipelineOutput, Tuple[Union[List[np.ndarray], torch.Tensor]]]:
+        self.check_inputs(height, width, num_frames)
+        if self.transformer is None or self.scheduler is None:
+            raise ValueError("`transformer` and `scheduler` are required for generation.")
+
+        device = self._execution_device
+        do_cfg = guidance_scale > 1.0
+        requested_batch_cfg = bool(batch_cfg)
+        effective_batch_cfg = requested_batch_cfg
+        batch_cfg_fallback_reason = None
+        self._last_batch_cfg_requested = requested_batch_cfg
+        self._last_effective_batch_cfg = bool(effective_batch_cfg and do_cfg)
+        self._last_batch_cfg_fallback_reason = batch_cfg_fallback_reason
+
+        cfg_parallel = cfg_parallel_group is not None
+        if cfg_parallel and effective_batch_cfg:
+            raise ValueError("`cfg_parallel_group` and `batch_cfg` are mutually exclusive.")
+        cfg_parallel_rank = 0
+        cfg_parallel_world_size = 1
+        if cfg_parallel:
+            if not dist.is_available() or not dist.is_initialized():
+                raise ValueError("`cfg_parallel_group` requires an initialized process group.")
+            if not do_cfg:
+                raise ValueError("CFG parallel requires `guidance_scale > 1.0`.")
+            cfg_parallel_rank = dist.get_rank(cfg_parallel_group)
+            cfg_parallel_world_size = dist.get_world_size(cfg_parallel_group)
+            if cfg_parallel_world_size != 2:
+                raise ValueError(
+                    f"CFG parallel currently requires exactly 2 ranks, got {cfg_parallel_world_size}."
+                )
+
+        if prompt_embeds is not None:
+            if prompt_mask is None:
+                raise ValueError("`prompt_mask` is required when `prompt_embeds` is provided.")
+            prompt_embeds = prompt_embeds.to(device=device)
+            prompt_mask = prompt_mask.to(device=device)
+        if negative_prompt_embeds is not None:
+            if negative_prompt_mask is None:
+                raise ValueError(
+                    "`negative_prompt_mask` is required when `negative_prompt_embeds` is provided."
+                )
+            negative_prompt_embeds = negative_prompt_embeds.to(device=device)
+            negative_prompt_mask = negative_prompt_mask.to(device=device)
+
+        if cfg_parallel and cfg_parallel_rank == 1:
+            if negative_prompt_embeds is not None:
+                negative_embeds, negative_mask = negative_prompt_embeds, negative_prompt_mask
+            else:
+                negative_embeds, negative_mask = self.encode_prompt(negative_prompt, device=device)
+            prompt_embeds = prompt_mask = None
+        else:
+            if prompt_embeds is None:
+                prompt_embeds, prompt_mask = self.encode_prompt(prompt, device=device)
+            if do_cfg and not cfg_parallel:
+                if null_cond_clone_zero:
+                    negative_embeds = torch.zeros_like(prompt_embeds)
+                    negative_mask = prompt_mask.clone()
+                elif negative_prompt_embeds is not None:
+                    negative_embeds, negative_mask = negative_prompt_embeds, negative_prompt_mask
+                else:
+                    negative_embeds, negative_mask = self.encode_prompt(negative_prompt, device=device)
+
+        latents = self.prepare_latents(num_frames, height, width, generator, latents, device)
+        # Clean temporal-prefix condition (e.g. the ti2v refiner's first-frame
+        # latent): written into the latent before sampling and after every
+        # scheduler step, so the fixed frames stay clean while the rest denoise
+        # against them through attention.
+        if cond_latent is not None:
+            cond_latent = cond_latent.to(device=device, dtype=torch.float32)
+            latents = self._apply_inpainting(latents, cond_latent)
+        sigmas = compute_refiner_sigmas(
+            sigma_max=float(self.scheduler.sigma_max),
+            sigma_min=float(self.scheduler.sigma_min),
+            num_inference_steps=num_inference_steps,
+            shift=shift,
+            t_thresh=t_thresh,
+            tail_steps=refiner_sigma_tail_steps,
+        )
+        if sigmas is None:
+            self.scheduler.set_timesteps(num_inference_steps, device=device, shift=shift)
+        else:
+            self.scheduler.set_timesteps(
+                int(sigmas.shape[0]),
+                device=device,
+                sigmas=sigmas,
+                shift=1.0,
+            )
+        transformer_dtype = _module_dtype(self.transformer)
+        vae_restore_device: Optional[torch.device] = None
+        vae_offloaded = False
+        if offload_vae_during_denoise and output_type == "np" and self.vae is not None:
+            vae_device = _module_device(self.vae)
+            if vae_device.type == "cuda":
+                self.vae.to("cpu")
+                torch.cuda.empty_cache()
+                vae_restore_device = vae_device
+                vae_offloaded = True
+        cfg_latent_src = _group_global_rank(cfg_parallel_group, 0)
+        cfg_uncond_src = _group_global_rank(cfg_parallel_group, 1)
+        for i, timestep in enumerate(self.progress_bar(self.scheduler.timesteps)):
+            if cfg_parallel:
+                dist.broadcast(latents, src=cfg_latent_src, group=cfg_parallel_group)
+            timestep_batch = _transformer_timestep(timestep, transformer_dtype).expand(1).to(device)
+            latent_model_input = latents
+            if cfg_parallel:
+                if cfg_parallel_rank == 0:
+                    branch_embeds = prompt_embeds
+                    branch_mask = prompt_mask
+                    branch_name = "transformer.cond"
+                else:
+                    branch_embeds = negative_embeds
+                    branch_mask = negative_mask
+                    branch_name = "transformer.uncond"
+                branch_model_input = branch_embeds.to(transformer_dtype)
+                with _transformer_autocast(device, transformer_dtype):
+                    branch_noise_pred = self.transformer(
+                        latent_model_input,
+                        timestep_batch,
+                        branch_model_input,
+                        encoder_attention_mask=branch_mask,
+                        return_dict=False,
+                    )[0].float()
+                if cfg_parallel_rank == 0:
+                    noise_pred = branch_noise_pred
+                    noise_pred_uncond = torch.empty_like(noise_pred)
+                else:
+                    noise_pred_uncond = branch_noise_pred
+                dist.broadcast(noise_pred_uncond, src=cfg_uncond_src, group=cfg_parallel_group)
+                if cfg_parallel_rank != 0:
+                    continue
+                noise_pred = noise_pred_uncond + guidance_scale * (noise_pred - noise_pred_uncond)
+            else:
+                prompt_model_input = prompt_embeds.to(transformer_dtype)
+                if do_cfg and effective_batch_cfg:
+                    negative_model_input = negative_embeds.to(transformer_dtype)
+                    cfg_embeds, cfg_mask = batch_cfg_prompt_inputs(
+                        prompt_model_input,
+                        prompt_mask,
+                        negative_model_input,
+                        negative_mask,
+                        null_cond_clone_zero=False,
+                    )
+                    cfg_latents = torch.cat([latent_model_input, latent_model_input], dim=0)
+                    cfg_timesteps = torch.cat([timestep_batch, timestep_batch], dim=0)
+                    with _transformer_autocast(device, transformer_dtype):
+                        noise_batched = self.transformer(
+                            cfg_latents,
+                            cfg_timesteps,
+                            cfg_embeds,
+                            encoder_attention_mask=cfg_mask,
+                            return_dict=False,
+                        )[0].float()
+                    noise_pred, noise_pred_uncond = noise_batched.chunk(2, dim=0)
+                    noise_pred = noise_pred_uncond + guidance_scale * (
+                        noise_pred - noise_pred_uncond
+                    )
+                else:
+                    with _transformer_autocast(device, transformer_dtype):
+                        noise_pred = self.transformer(
+                            latent_model_input,
+                            timestep_batch,
+                            prompt_model_input,
+                            encoder_attention_mask=prompt_mask,
+                            return_dict=False,
+                        )[0].float()
+
+                if do_cfg and not effective_batch_cfg:
+                    negative_model_input = negative_embeds.to(transformer_dtype)
+                    with _transformer_autocast(device, transformer_dtype):
+                        noise_pred_uncond = self.transformer(
+                            latent_model_input,
+                            timestep_batch,
+                            negative_model_input,
+                            encoder_attention_mask=negative_mask,
+                            return_dict=False,
+                        )[0].float()
+                    noise_pred = noise_pred_uncond + guidance_scale * (noise_pred - noise_pred_uncond)
+
+            latents = self.scheduler.step(
+                noise_pred,
+                timestep,
+                latents,
+                return_dict=False,
+                generator=generator,
+            )[0]
+            if cond_latent is not None:
+                latents = self._apply_inpainting(latents, cond_latent)
+
+        if cfg_parallel:
+            dist.barrier(group=cfg_parallel_group)
+            if cfg_parallel_rank != 0:
+                frames = latents if output_type == "latent" else []
+                self.maybe_free_model_hooks()
+                if not return_dict:
+                    return (frames,)
+                return LingBotVideoPipelineOutput(frames=frames)
+
+        if output_type == "latent":
+            frames = latents
+        elif output_type == "np":
+            if vae_offloaded and vae_restore_device is not None:
+                self.vae.to(device=vae_restore_device)
+                torch.cuda.empty_cache()
+            frames = self._decode_latents(latents)
+        else:
+            raise ValueError(f"Unsupported output_type: {output_type}")
+
+        self.maybe_free_model_hooks()
+        if not return_dict:
+            return (frames,)
+        return LingBotVideoPipelineOutput(frames=frames)

lingbot_video/pipeline_lingbot_video_i2v.py

@@ -0,0 +1,352 @@
+from __future__ import annotations
+
+import math
+from typing import Any, Optional, Tuple, Union
+
+import numpy as np
+import torch
+import torch.distributed as dist
+import torch.nn.functional as F
+from PIL import Image
+from diffusers.utils.torch_utils import randn_tensor
+
+from .pipeline_lingbot_video import (
+    DEFAULT_NEGATIVE_PROMPT,
+    LingBotVideoPipeline,
+    LingBotVideoPipelineOutput,
+    _group_global_rank,
+    _module_device,
+    _module_dtype,
+    _transformer_autocast,
+    _transformer_timestep,
+)
+from .utils import batch_cfg_prompt_inputs
+
+
+IMAGE_MIN_TOKEN_NUM = 4
+IMAGE_MAX_TOKEN_NUM = 16384
+MAX_RATIO = 200
+SPATIAL_MERGE_SIZE = 2
+
+
+def _round_by_factor(number: float, factor: int) -> int:
+    return round(number / factor) * factor
+
+
+def _ceil_by_factor(number: float, factor: int) -> int:
+    return math.ceil(number / factor) * factor
+
+
+def _floor_by_factor(number: float, factor: int) -> int:
+    return math.floor(number / factor) * factor
+
+
+def smart_resize(
+    height: int,
+    width: int,
+    factor: int,
+    min_pixels: Optional[int] = None,
+    max_pixels: Optional[int] = None,
+) -> Tuple[int, int]:
+    max_pixels = max_pixels if max_pixels is not None else IMAGE_MAX_TOKEN_NUM * factor**2
+    min_pixels = min_pixels if min_pixels is not None else IMAGE_MIN_TOKEN_NUM * factor**2
+    if max_pixels < min_pixels:
+        raise ValueError("max_pixels must be greater than or equal to min_pixels.")
+    if max(height, width) / min(height, width) > MAX_RATIO:
+        raise ValueError(f"absolute aspect ratio must be smaller than {MAX_RATIO}.")
+
+    resized_height = max(factor, _round_by_factor(height, factor))
+    resized_width = max(factor, _round_by_factor(width, factor))
+    if resized_height * resized_width > max_pixels:
+        beta = math.sqrt((height * width) / max_pixels)
+        resized_height = _floor_by_factor(height / beta, factor)
+        resized_width = _floor_by_factor(width / beta, factor)
+    elif resized_height * resized_width < min_pixels:
+        beta = math.sqrt(min_pixels / (height * width))
+        resized_height = _ceil_by_factor(height * beta, factor)
+        resized_width = _ceil_by_factor(width * beta, factor)
+    return resized_height, resized_width
+
+
+def _pixel_tensor_to_pil(pixel: torch.Tensor) -> Image.Image:
+    """Match torchvision.transforms.ToPILImage for a float CHW image in [0, 1]."""
+    frame = pixel[0, :, 0].detach().cpu().clamp(0, 1)
+    array = frame.permute(1, 2, 0).mul(255).byte().numpy()
+    return Image.fromarray(array, mode="RGB")
+
+
+class LingBotVideoImageToVideoPipeline(LingBotVideoPipeline):
+    """LingBotVideo ti2v pipeline.
+
+    The condition frame is used twice: as visual input for Qwen3-VL and as a
+    clean latent that is written into the beginning of the diffusion latent
+    before sampling and after every scheduler step.
+    """
+
+    def preprocess_image(self, image: Image.Image, height: int, width: int) -> torch.Tensor:
+        if image is None:
+            raise ValueError("`image` is required when `image_tensor` is not provided.")
+        raw = torch.from_numpy(np.array(image.convert("RGB"))).permute(2, 0, 1).unsqueeze(0).contiguous()
+        old_h, old_w = raw.shape[-2:]
+        scale = max(height / old_h, width / old_w)
+        new_h = max(math.ceil(old_h * scale), height)
+        new_w = max(math.ceil(old_w * scale), width)
+        resized = F.interpolate(raw, size=(new_h, new_w), mode="bilinear", align_corners=False)
+        top = int(round((new_h - height) / 2.0))
+        left = int(round((new_w - width) / 2.0))
+        cropped = resized[:, :, top : top + height, left : left + width].float() / 255.0
+        return cropped.unsqueeze(2)
+
+    def _vision_patch_size(self) -> int:
+        for obj in (
+            getattr(getattr(self.text_encoder, "config", None), "vision_config", None),
+            getattr(getattr(self.processor, "image_processor", None), "config", None),
+            getattr(self.processor, "image_processor", None),
+        ):
+            patch = getattr(obj, "patch_size", None)
+            if patch is not None:
+                return int(patch)
+        return 16
+
+    def _vlm_image(self, pixel: torch.Tensor) -> Image.Image:
+        image = _pixel_tensor_to_pil(pixel)
+        patch_factor = self._vision_patch_size() * SPATIAL_MERGE_SIZE
+        width, height = image.size
+        resized_height, resized_width = smart_resize(height, width, factor=patch_factor)
+        return image.resize((resized_width, resized_height))
+
+    @torch.no_grad()
+    def encode_image_latent(
+        self,
+        pixel: torch.Tensor,
+        generator: Optional[torch.Generator] = None,
+    ) -> torch.Tensor:
+        if self.vae is None:
+            raise ValueError("`vae` is required to encode image latents.")
+        device = _module_device(self.vae)
+        pixel = pixel.to(device=device, dtype=torch.float32)
+        norm_pixel = (pixel - 0.5) / 0.5
+        with torch.autocast("cuda", dtype=torch.bfloat16, enabled=device.type == "cuda"):
+            latents = self.vae.encode(norm_pixel).latent_dist.sample(generator)
+
+        mean = torch.tensor(self.vae.config.latents_mean, device=latents.device, dtype=torch.float32)
+        std_inv = 1.0 / torch.tensor(
+            self.vae.config.latents_std, device=latents.device, dtype=torch.float32
+        )
+        mean = mean.view(1, -1, 1, 1, 1)
+        std_inv = std_inv.view(1, -1, 1, 1, 1)
+        return (latents.float() - mean) * std_inv
+
+    @torch.no_grad()
+    def __call__(
+        self,
+        prompt: str,
+        image: Optional[Image.Image] = None,
+        image_tensor: Optional[torch.Tensor] = None,
+        cond_latent: Optional[torch.Tensor] = None,
+        negative_prompt: str = DEFAULT_NEGATIVE_PROMPT,
+        height: int = 480,
+        width: int = 480,
+        num_frames: int = 81,
+        num_inference_steps: int = 40,
+        guidance_scale: float = 6.0,
+        shift: float = 3.0,
+        generator: Optional[torch.Generator] = None,
+        latents: Optional[torch.Tensor] = None,
+        prompt_embeds: Optional[torch.Tensor] = None,
+        prompt_mask: Optional[torch.Tensor] = None,
+        negative_prompt_embeds: Optional[torch.Tensor] = None,
+        negative_prompt_mask: Optional[torch.Tensor] = None,
+        output_type: str = "np",
+        cfg_parallel_group: Optional[Any] = None,
+        batch_cfg: bool = False,
+        null_cond_clone_zero: bool = False,
+        return_dict: bool = True,
+    ) -> Union[LingBotVideoPipelineOutput, Tuple[Union[list, torch.Tensor]]]:
+        self.check_inputs(height, width, num_frames)
+        if self.transformer is None or self.scheduler is None:
+            raise ValueError("`transformer` and `scheduler` are required for generation.")
+
+        device = self._execution_device
+        do_cfg = guidance_scale > 1.0
+        requested_batch_cfg = bool(batch_cfg)
+        effective_batch_cfg = requested_batch_cfg and do_cfg
+        self._last_batch_cfg_requested = requested_batch_cfg
+        self._last_effective_batch_cfg = effective_batch_cfg
+        self._last_batch_cfg_fallback_reason = None
+
+        cfg_parallel = cfg_parallel_group is not None
+        if cfg_parallel and effective_batch_cfg:
+            raise ValueError("`cfg_parallel_group` and `batch_cfg` are mutually exclusive.")
+        cfg_parallel_rank = 0
+        cfg_parallel_world_size = 1
+        if cfg_parallel:
+            if not dist.is_available() or not dist.is_initialized():
+                raise ValueError("`cfg_parallel_group` requires an initialized process group.")
+            if not do_cfg:
+                raise ValueError("CFG parallel requires `guidance_scale > 1.0`.")
+            cfg_parallel_rank = dist.get_rank(cfg_parallel_group)
+            cfg_parallel_world_size = dist.get_world_size(cfg_parallel_group)
+            if cfg_parallel_world_size != 2:
+                raise ValueError(
+                    f"CFG parallel currently requires exactly 2 ranks, got {cfg_parallel_world_size}."
+                )
+
+        pixel = image_tensor if image_tensor is not None else self.preprocess_image(image, height, width)
+        pixel = pixel.to(device=device, dtype=torch.float32)
+        vlm_image = self._vlm_image(pixel)
+
+        if prompt_embeds is not None:
+            if prompt_mask is None:
+                raise ValueError("`prompt_mask` is required when `prompt_embeds` is provided.")
+            prompt_embeds = prompt_embeds.to(device=device)
+            prompt_mask = prompt_mask.to(device=device)
+        if negative_prompt_embeds is not None:
+            if negative_prompt_mask is None:
+                raise ValueError(
+                    "`negative_prompt_mask` is required when `negative_prompt_embeds` is provided."
+                )
+            negative_prompt_embeds = negative_prompt_embeds.to(device=device)
+            negative_prompt_mask = negative_prompt_mask.to(device=device)
+
+        if cfg_parallel and cfg_parallel_rank == 1:
+            if negative_prompt_embeds is not None:
+                negative_embeds, negative_mask = negative_prompt_embeds, negative_prompt_mask
+            else:
+                negative_embeds, negative_mask = self.encode_prompt(
+                    negative_prompt, images=[vlm_image], device=device
+                )
+            prompt_embeds = prompt_mask = None
+        else:
+            if prompt_embeds is None:
+                prompt_embeds, prompt_mask = self.encode_prompt(prompt, images=[vlm_image], device=device)
+            if do_cfg and not cfg_parallel:
+                if null_cond_clone_zero:
+                    negative_embeds = torch.zeros_like(prompt_embeds)
+                    negative_mask = prompt_mask.clone()
+                elif negative_prompt_embeds is not None:
+                    negative_embeds, negative_mask = negative_prompt_embeds, negative_prompt_mask
+                else:
+                    negative_embeds, negative_mask = self.encode_prompt(
+                        negative_prompt, images=[vlm_image], device=device
+                    )
+
+        if cond_latent is None:
+            cond_latent = self.encode_image_latent(pixel, generator=generator)
+        cond_latent = cond_latent.to(device=device, dtype=torch.float32)
+
+        latents = self.prepare_latents(num_frames, height, width, generator, latents, device)
+        latents = self._apply_inpainting(latents, cond_latent)
+        self.scheduler.set_timesteps(num_inference_steps, device=device, shift=shift)
+        transformer_dtype = _module_dtype(self.transformer)
+        cfg_latent_src = _group_global_rank(cfg_parallel_group, 0)
+        cfg_uncond_src = _group_global_rank(cfg_parallel_group, 1)
+
+        for i, timestep in enumerate(self.progress_bar(self.scheduler.timesteps)):
+            if cfg_parallel:
+                dist.broadcast(latents, src=cfg_latent_src, group=cfg_parallel_group)
+            timestep_batch = _transformer_timestep(timestep, transformer_dtype).expand(1).to(device)
+            latent_model_input = latents
+            if cfg_parallel:
+                if cfg_parallel_rank == 0:
+                    branch_embeds = prompt_embeds
+                    branch_mask = prompt_mask
+                else:
+                    branch_embeds = negative_embeds
+                    branch_mask = negative_mask
+                branch_model_input = branch_embeds.to(transformer_dtype)
+                with _transformer_autocast(device, transformer_dtype):
+                    branch_noise_pred = self.transformer(
+                        latent_model_input,
+                        timestep_batch,
+                        branch_model_input,
+                        encoder_attention_mask=branch_mask,
+                        return_dict=False,
+                    )[0].float()
+                if cfg_parallel_rank == 0:
+                    noise_pred = branch_noise_pred
+                    noise_pred_uncond = torch.empty_like(noise_pred)
+                else:
+                    noise_pred_uncond = branch_noise_pred
+                dist.broadcast(noise_pred_uncond, src=cfg_uncond_src, group=cfg_parallel_group)
+                if cfg_parallel_rank != 0:
+                    continue
+                noise_pred = noise_pred_uncond + guidance_scale * (noise_pred - noise_pred_uncond)
+            else:
+                prompt_model_input = prompt_embeds.to(transformer_dtype)
+                if do_cfg and effective_batch_cfg:
+                    negative_model_input = negative_embeds.to(transformer_dtype)
+                    cfg_embeds, cfg_mask = batch_cfg_prompt_inputs(
+                        prompt_model_input,
+                        prompt_mask,
+                        negative_model_input,
+                        negative_mask,
+                        null_cond_clone_zero=null_cond_clone_zero,
+                    )
+                    cfg_latents = torch.cat([latent_model_input, latent_model_input], dim=0)
+                    cfg_timesteps = torch.cat([timestep_batch, timestep_batch], dim=0)
+                    with _transformer_autocast(device, transformer_dtype):
+                        noise_batched = self.transformer(
+                            cfg_latents,
+                            cfg_timesteps,
+                            cfg_embeds,
+                            encoder_attention_mask=cfg_mask,
+                            return_dict=False,
+                        )[0].float()
+                    noise_pred, noise_pred_uncond = noise_batched.chunk(2, dim=0)
+                    noise_pred = noise_pred_uncond + guidance_scale * (
+                        noise_pred - noise_pred_uncond
+                    )
+                else:
+                    with _transformer_autocast(device, transformer_dtype):
+                        noise_pred = self.transformer(
+                            latent_model_input,
+                            timestep_batch,
+                            prompt_model_input,
+                            encoder_attention_mask=prompt_mask,
+                            return_dict=False,
+                        )[0].float()
+
+                    if do_cfg:
+                        negative_model_input = negative_embeds.to(transformer_dtype)
+                        with _transformer_autocast(device, transformer_dtype):
+                            noise_pred_uncond = self.transformer(
+                                latent_model_input,
+                                timestep_batch,
+                                negative_model_input,
+                                encoder_attention_mask=negative_mask,
+                                return_dict=False,
+                            )[0].float()
+                        noise_pred = noise_pred_uncond + guidance_scale * (
+                            noise_pred - noise_pred_uncond
+                        )
+
+            latents = self.scheduler.step(
+                noise_pred,
+                timestep,
+                latents,
+                return_dict=False,
+                generator=generator,
+            )[0]
+            latents = self._apply_inpainting(latents, cond_latent)
+
+        if cfg_parallel:
+            dist.barrier(group=cfg_parallel_group)
+            if cfg_parallel_rank != 0:
+                frames = latents if output_type == "latent" else []
+                self.maybe_free_model_hooks()
+                if not return_dict:
+                    return (frames,)
+                return LingBotVideoPipelineOutput(frames=frames)
+
+        if output_type == "latent":
+            frames = latents
+        elif output_type == "np":
+            frames = self._decode_latents(latents)
+        else:
+            raise ValueError(f"Unsupported output_type: {output_type}")
+
+        self.maybe_free_model_hooks()
+        if not return_dict:
+            return (frames,)
+        return LingBotVideoPipelineOutput(frames=frames)

lingbot_video/runner.py

@@ -0,0 +1,1401 @@
+from __future__ import annotations
+
+import argparse
+import gc
+import json
+import logging
+import os
+import sys
+import warnings
+from contextlib import contextmanager
+from datetime import timedelta
+from pathlib import Path
+from typing import Any
+
+import numpy as np
+import torch
+import torch.distributed as dist
+from torch.distributed.device_mesh import DeviceMesh
+from PIL import Image
+
+REPO_ROOT = Path(__file__).resolve().parents[1]
+sys.path.insert(0, str(REPO_ROOT))
+
+
+def _env_flag(name: str, default: bool = False) -> bool:
+    raw = os.environ.get(name)
+    if raw is None:
+        return default
+    return raw.lower() in {"1", "true", "yes", "on"}
+
+
+def _configure_concise_import_logs() -> None:
+    if _env_flag("LINGBOT_VERBOSE_LOGS"):
+        return
+    os.environ.setdefault("HF_HUB_DISABLE_PROGRESS_BARS", "1")
+    os.environ.setdefault("TOKENIZERS_PARALLELISM", "false")
+    os.environ.setdefault("DIFFUSERS_VERBOSITY", "error")
+    os.environ.setdefault("TRANSFORMERS_VERBOSITY", "error")
+    warnings.filterwarnings(
+        "ignore",
+        message=r"Unable to import `torchao` Tensor objects.*",
+    )
+    warnings.filterwarnings(
+        "ignore",
+        message=r"`enable_parallelism` is an experimental feature.*",
+    )
+    warnings.filterwarnings(
+        "ignore",
+        message=r"barrier\(\): using the device under current context.*",
+    )
+    logging.getLogger(
+        "sglang.srt.layers.moe.moe_runner.triton_utils.fused_moe_triton_config"
+    ).setLevel(logging.ERROR)
+
+
+_configure_concise_import_logs()
+
+from lingbot_video.inference_backend import (  # noqa: E402
+    resolve_backend_engine,
+    resolve_negative_prompt_arg,
+)
+
+try:
+    import transformers as _transformers
+    import transformers.cache_utils as _transformers_cache_utils
+except Exception as exc:  # pragma: no cover - deployment dependency guard
+    _transformers = None
+    _transformers_cache_utils = None
+    _TRANSFORMERS_IMPORT_ERROR = exc
+else:
+    _TRANSFORMERS_IMPORT_ERROR = None
+
+try:
+    import diffusers.utils.import_utils as _diffusers_import_utils
+except Exception as exc:  # pragma: no cover - deployment dependency guard
+    _diffusers_import_utils = None
+    _DIFFUSERS_IMPORT_UTILS_ERROR = exc
+else:
+    _DIFFUSERS_IMPORT_UTILS_ERROR = None
+
+try:
+    from diffusers.utils import logging as _diffusers_logging
+except Exception:
+    _diffusers_logging = None
+
+try:
+    from transformers.utils import logging as _transformers_logging
+except Exception:
+    _transformers_logging = None
+
+try:
+    import huggingface_hub.utils as _hf_hub_utils
+except Exception:
+    _hf_hub_utils = None
+
+
+def _install_sglang_import_shims() -> None:
+    if _transformers is not None and _transformers_cache_utils is not None:
+        hybrid_cache = getattr(_transformers_cache_utils, "HybridCache", None)
+        if hybrid_cache is None:
+            hybrid_cache = getattr(_transformers_cache_utils, "DynamicCache", object)
+            _transformers_cache_utils.HybridCache = hybrid_cache
+        import_structure = getattr(_transformers, "_import_structure", {})
+        cache_exports = import_structure.setdefault("cache_utils", [])
+        if "HybridCache" not in cache_exports:
+            cache_exports.append("HybridCache")
+        _transformers.HybridCache = hybrid_cache
+    if _diffusers_import_utils is not None:
+        _diffusers_import_utils._peft_available = False
+
+
+_install_sglang_import_shims()
+
+try:
+    from lingbot_video.fsdp_inference import (
+        apply_fsdp_inference,
+        init_fsdp_inference_mesh,
+    )
+    from lingbot_video.model_paths import (
+        effective_refiner_model_dir,
+        model_component_dir,
+    )
+    from lingbot_video.pipeline_lingbot_video import (
+        DEFAULT_NEGATIVE_PROMPT,
+        DEFAULT_NEGATIVE_PROMPT_IMAGE,
+        LingBotVideoPipeline,
+    )
+    from lingbot_video.pipeline_lingbot_video_i2v import (
+        LingBotVideoImageToVideoPipeline,
+    )
+    from lingbot_video.utils import (
+        caption_from_sample,
+        load_first_frame_condition_tensor,
+        load_refiner_video_tensor,
+        num_frames_from_duration,
+        prepare_refiner_latent,
+    )
+except Exception as exc:  # pragma: no cover - reported when generation is attempted
+    DEFAULT_NEGATIVE_PROMPT = ""
+    DEFAULT_NEGATIVE_PROMPT_IMAGE = ""
+    LingBotVideoPipeline = None
+    LingBotVideoImageToVideoPipeline = None
+    apply_fsdp_inference = None
+    init_fsdp_inference_mesh = None
+    effective_refiner_model_dir = None
+    model_component_dir = None
+    caption_from_sample = None
+    load_first_frame_condition_tensor = None
+    load_refiner_video_tensor = None
+    num_frames_from_duration = None
+    prepare_refiner_latent = None
+    _LINGBOT_PIPELINE_IMPORT_ERROR = exc
+else:
+    _LINGBOT_PIPELINE_IMPORT_ERROR = None
+
+try:
+    from lingbot_video.transformer_lingbot_video import (
+        LingBotVideoTransformer3DModel,
+    )
+except Exception as exc:  # pragma: no cover - reported when generation is attempted
+    LingBotVideoTransformer3DModel = None
+    _TRANSFORMER_IMPORT_ERROR = exc
+else:
+    _TRANSFORMER_IMPORT_ERROR = None
+
+try:
+    from diffusers import ContextParallelConfig
+    from diffusers.hooks.context_parallel import apply_context_parallel
+    from diffusers.models._modeling_parallel import ParallelConfig
+    from diffusers.models.attention import AttentionModuleMixin
+    from diffusers.models.attention_processor import Attention, MochiAttention
+    from diffusers.utils import export_to_video
+except Exception as exc:  # pragma: no cover - reported by the selected engine
+    ContextParallelConfig = None
+    ParallelConfig = None
+    Attention = None
+    AttentionModuleMixin = None
+    MochiAttention = None
+    apply_context_parallel = None
+    export_to_video = None
+    _DIFFUSERS_IMPORT_ERROR = exc
+else:
+    _DIFFUSERS_IMPORT_ERROR = None
+
+
+def _cuda_sdp_state() -> dict[str, bool | None]:
+    state: dict[str, bool | None] = {}
+    cuda_backends = getattr(torch.backends, "cuda", None)
+    if cuda_backends is None:
+        return state
+    getters = {
+        "flash": "flash_sdp_enabled",
+        "mem_efficient": "mem_efficient_sdp_enabled",
+        "math": "math_sdp_enabled",
+        "cudnn": "cudnn_sdp_enabled",
+    }
+    for key, getter in getters.items():
+        fn = getattr(cuda_backends, getter, None)
+        if fn is None:
+            state[key] = None
+            continue
+        try:
+            state[key] = bool(fn())
+        except Exception:
+            state[key] = None
+    return state
+
+
+def _restore_cuda_sdp_state(state: dict[str, bool | None]) -> None:
+    cuda_backends = getattr(torch.backends, "cuda", None)
+    if cuda_backends is None:
+        return
+    setters = {
+        "flash": "enable_flash_sdp",
+        "mem_efficient": "enable_mem_efficient_sdp",
+        "math": "enable_math_sdp",
+        "cudnn": "enable_cudnn_sdp",
+    }
+    for key, setter in setters.items():
+        value = state.get(key)
+        if value is None:
+            continue
+        fn = getattr(cuda_backends, setter, None)
+        if fn is None:
+            continue
+        try:
+            fn(value)
+        except Exception:
+            pass
+
+
+_BASELINE_CUDA_SDP_STATE = _cuda_sdp_state()
+
+
+RESOLUTION_BUCKETS: dict[str, dict[str, tuple[int, int]]] = {
+    "192p": {
+        "1:1": (192, 192),
+        "9:16": (192, 320),
+        "16:9": (320, 192),
+        "3:4": (192, 256),
+        "4:3": (256, 192),
+    },
+    "480p": {
+        "1:1": (480, 480),
+        "9:16": (480, 832),
+        "16:9": (832, 480),
+        "3:4": (480, 640),
+        "4:3": (640, 480),
+    },
+    "720p": {
+        "1:1": (736, 736),
+        "9:16": (736, 1280),
+        "16:9": (1280, 736),
+        "3:4": (736, 960),
+        "4:3": (960, 736),
+    },
+    "1080p": {
+        "1:1": (1088, 1088),
+        "9:16": (1088, 1920),
+        "16:9": (1920, 1088),
+        "3:4": (1088, 1440),
+        "4:3": (1440, 1088),
+    },
+    "2k": {
+        "1:1": (1440, 1440),
+        "9:16": (1440, 2560),
+        "16:9": (2560, 1440),
+        "3:4": (1440, 1920),
+        "4:3": (1920, 1440),
+    },
+    "4k": {
+        "1:1": (2176, 2176),
+        "9:16": (2176, 3840),
+        "16:9": (3840, 2176),
+        "3:4": (2176, 2880),
+        "4:3": (2880, 2176),
+    },
+}
+
+
+def _distributed_env() -> tuple[int, int, int]:
+    return (
+        int(os.environ.get("RANK", "0")),
+        int(os.environ.get("LOCAL_RANK", "0")),
+        int(os.environ.get("WORLD_SIZE", "1")),
+    )
+
+
+def _default_device() -> torch.device:
+    if not torch.cuda.is_available():
+        return torch.device("cpu")
+    return torch.device("cuda", torch.cuda.current_device())
+
+
+def _init_parallel(
+    cfg_degree: int,
+    context_degree: int,
+    enable_fsdp_inference: bool,
+) -> tuple[int, int, int, Any | None, DeviceMesh | None, int, int]:
+    rank, local_rank, world_size = _distributed_env()
+    if cfg_degree < 1 or context_degree < 1:
+        raise ValueError("Parallel degrees must be >= 1.")
+    expected_world_size = cfg_degree * context_degree
+    if expected_world_size <= 1:
+        if enable_fsdp_inference and world_size > 1:
+            if not torch.cuda.is_available():
+                raise RuntimeError("FSDP inference requires CUDA devices.")
+            torch.cuda.set_device(local_rank)
+            if not dist.is_initialized():
+                dist.init_process_group("nccl", timeout=timedelta(minutes=30))
+            return rank, local_rank, world_size, None, None, 0, 0
+        if torch.cuda.is_available():
+            torch.cuda.set_device(0)
+        return rank, local_rank, world_size, None, None, 0, 0
+    if world_size != expected_world_size:
+        raise ValueError(
+            f"Parallel topology cfg={cfg_degree}, context={context_degree} requires "
+            f"WORLD_SIZE={expected_world_size}, got {world_size}."
+        )
+    if not torch.cuda.is_available():
+        raise RuntimeError("Distributed parallel inference requires CUDA devices.")
+    torch.cuda.set_device(local_rank)
+    if not dist.is_initialized():
+        dist.init_process_group("nccl", timeout=timedelta(minutes=30))
+
+    cfg_branch_rank = rank // context_degree
+    context_rank = rank % context_degree
+    cfg_parallel_group = None
+    if cfg_degree > 1:
+        lane_groups = [
+            dist.new_group(
+                ranks=[cfg_rank * context_degree + context_rank for cfg_rank in range(cfg_degree)]
+            )
+            for context_rank in range(context_degree)
+        ]
+        cfg_parallel_group = lane_groups[context_rank]
+
+    context_mesh = None
+    if context_degree > 1 and cfg_degree > 1:
+        mesh_ranks = [
+            [
+                cfg_rank * context_degree + context_rank
+                for context_rank in range(context_degree)
+            ]
+            for cfg_rank in range(cfg_degree)
+        ]
+        context_mesh = DeviceMesh(
+            "cuda",
+            mesh_ranks,
+            mesh_dim_names=("ring", "ulysses"),
+        )
+
+    return rank, local_rank, world_size, cfg_parallel_group, context_mesh, cfg_branch_rank, context_rank
+
+try:
+    from lingbot_video.native_backend import (
+        LingBotVideoNativePipeline,
+        register_lingbot_native_pipeline,
+    )
+except Exception as exc:  # pragma: no cover - only needed for --engine sglang-native
+    LingBotVideoNativePipeline = None
+    register_lingbot_native_pipeline = None
+    _NATIVE_BACKEND_IMPORT_ERROR = exc
+else:
+    _NATIVE_BACKEND_IMPORT_ERROR = None
+finally:
+    _restore_cuda_sdp_state(_BASELINE_CUDA_SDP_STATE)
+
+
+def _parse_dtype(name: str) -> torch.dtype:
+    normalized = name.lower()
+    if normalized in {"bf16", "bfloat16"}:
+        return torch.bfloat16
+    if normalized in {"fp16", "float16"}:
+        return torch.float16
+    if normalized in {"fp32", "float32"}:
+        return torch.float32
+    raise ValueError(f"unsupported dtype: {name}")
+
+
+def _dtype_name(dtype: torch.dtype) -> str:
+    return str(dtype).replace("torch.", "")
+
+
+def _make_dtype_map(args: argparse.Namespace) -> dict[str, torch.dtype]:
+    default_dtype = _parse_dtype(args.default_dtype)
+    return {
+        "default": default_dtype,
+        "transformer": _parse_dtype(args.transformer_dtype),
+        "text_encoder": _parse_dtype(args.text_encoder_dtype),
+        "vae": _parse_dtype(args.vae_dtype),
+    }
+
+
+def _make_default_image(height: int, width: int) -> Image.Image:
+    yy, xx = np.mgrid[0:height, 0:width]
+    red = (xx / max(width - 1, 1) * 255).astype(np.uint8)
+    green = (yy / max(height - 1, 1) * 255).astype(np.uint8)
+    blue = (((xx // 24 + yy // 24) % 2) * 180 + 40).astype(np.uint8)
+    return Image.fromarray(np.stack([red, green, blue], axis=-1), mode="RGB")
+
+
+def _load_prompt_sample(path: Path) -> dict[str, Any]:
+    data = json.loads(path.read_text(encoding="utf-8"))
+    if isinstance(data, list):
+        if not data:
+            raise ValueError(f"`--prompt_json` is empty: {path}")
+        data = data[0]
+    if not isinstance(data, dict):
+        raise ValueError(f"`--prompt_json` must contain a dict or a non-empty list of dicts: {path}")
+    return data
+
+
+def _caption_from_sample(sample: dict[str, Any]) -> str:
+    if caption_from_sample is not None:
+        return caption_from_sample(sample)
+    if "caption" in sample:
+        caption = sample["caption"]
+    else:
+        runtime_keys = {
+            "duration",
+            "fps",
+            "height",
+            "width",
+            "num_frames",
+            "resolution",
+            "ratio",
+        }
+        caption = {key: value for key, value in sample.items() if key not in runtime_keys}
+    if isinstance(caption, dict):
+        return str(caption.get("qwen_long_caption") or caption.get("comprehensive_description") or caption)
+    return str(caption)
+
+
+def _height_width_from_bucket(resolution: str, ratio: str) -> tuple[int, int]:
+    if resolution not in RESOLUTION_BUCKETS:
+        choices = ", ".join(sorted(RESOLUTION_BUCKETS))
+        raise ValueError(f"unsupported resolution {resolution!r}; choices: {choices}")
+    ratios = RESOLUTION_BUCKETS[resolution]
+    if ratio not in ratios:
+        choices = ", ".join(sorted(ratios))
+        raise ValueError(f"unsupported ratio {ratio!r} for {resolution}; choices: {choices}")
+    return ratios[ratio]
+
+
+def _module_dtype(module: Any) -> str | None:
+    if module is None:
+        return None
+    if hasattr(module, "dtype"):
+        dtype = getattr(module, "dtype")
+        if isinstance(dtype, torch.dtype):
+            return _dtype_name(dtype)
+    if isinstance(module, torch.nn.Module):
+        try:
+            return _dtype_name(next(module.parameters()).dtype)
+        except StopIteration:
+            return None
+    return None
+
+
+def _component_dtypes(pipe: Any) -> dict[str, str | None]:
+    return {
+        "transformer": _module_dtype(getattr(pipe, "transformer", None)),
+        "text_encoder": _module_dtype(getattr(pipe, "text_encoder", None)),
+        "vae": _module_dtype(getattr(pipe, "vae", None)),
+    }
+
+
+def _disable_external_progress_bars() -> None:
+    for logging_module in (_diffusers_logging, _transformers_logging):
+        fn = getattr(logging_module, "disable_progress_bar", None)
+        if callable(fn):
+            fn()
+    hf_fn = getattr(_hf_hub_utils, "disable_progress_bars", None)
+    if callable(hf_fn):
+        hf_fn()
+
+
+def _configure_pipeline_logs(pipe: Any) -> None:
+    inner_pipe = _inner_diffusers_pipe(pipe)
+    set_progress_bar_config = getattr(inner_pipe, "set_progress_bar_config", None)
+    if callable(set_progress_bar_config):
+        set_progress_bar_config(
+            disable=not _show_progress(),
+            desc="denoising",
+            dynamic_ncols=True,
+        )
+
+
+@contextmanager
+def _patch_qwen3vl_from_pretrained():
+    try:
+        from transformers import Qwen3VLForConditionalGeneration
+    except Exception:
+        yield
+        return
+
+    original_from_pretrained = Qwen3VLForConditionalGeneration.from_pretrained
+    attn_implementation = os.environ.get("LINGBOT_QWEN_ATTN_IMPLEMENTATION", "flash_attention_3")
+
+    @classmethod
+    def patched_from_pretrained(cls, pretrained_model_name_or_path, *args, **kwargs):
+        if attn_implementation:
+            kwargs.setdefault("attn_implementation", attn_implementation)
+        if "torch_dtype" in kwargs and "dtype" not in kwargs:
+            kwargs["dtype"] = kwargs.pop("torch_dtype")
+        return original_from_pretrained(pretrained_model_name_or_path, *args, **kwargs)
+
+    Qwen3VLForConditionalGeneration.from_pretrained = patched_from_pretrained
+    try:
+        yield
+    finally:
+        Qwen3VLForConditionalGeneration.from_pretrained = original_from_pretrained
+
+
+def _destroy_parallel_if_needed() -> None:
+    if dist.is_available() and dist.is_initialized():
+        dist.destroy_process_group()
+
+
+def _sync_parallel_if_needed() -> None:
+    if dist.is_available() and dist.is_initialized():
+        dist.barrier()
+
+
+def _current_rank() -> int:
+    if dist.is_available() and dist.is_initialized():
+        return dist.get_rank()
+    return 0
+
+
+def _is_main_process() -> bool:
+    return _current_rank() == 0
+
+
+def _show_progress() -> bool:
+    return _is_main_process() and not _env_flag("LINGBOT_QUIET_PROGRESS")
+
+
+def _log_progress(message: str) -> None:
+    if _show_progress():
+        print(message, flush=True)
+
+
+def _apply_fsdp_inference_if_requested(
+    pipe: Any,
+    enabled: bool,
+    mesh: DeviceMesh | None,
+) -> Any | None:
+    if not enabled:
+        return None
+    if apply_fsdp_inference is None:
+        raise RuntimeError("FSDP inference helpers are not importable.") from _LINGBOT_PIPELINE_IMPORT_ERROR
+    inner_pipe = _inner_diffusers_pipe(pipe)
+    transformer = getattr(inner_pipe, "transformer", None)
+    if not isinstance(transformer, torch.nn.Module):
+        raise ValueError(
+            "FSDP inference requires a pipeline with a torch.nn.Module transformer."
+        )
+    _log_progress("applying FSDP inference sharding")
+    info = apply_fsdp_inference(transformer, mesh)
+    _log_progress(f"applied FSDP inference sharding: {info}")
+    return info
+
+
+def _cache_prompt_conditions(
+    pipe: Any,
+    prompt: str,
+    negative_prompt: str,
+    *,
+    device: torch.device,
+    null_cond_clone_zero: bool,
+    images: list[Any] | None = None,
+) -> dict[str, torch.Tensor]:
+    encode_kwargs: dict[str, Any] = {"device": device}
+    if images is not None:
+        encode_kwargs["images"] = images
+    prompt_embeds, prompt_mask = pipe.encode_prompt(prompt, **encode_kwargs)
+    if null_cond_clone_zero:
+        negative_embeds = torch.zeros_like(prompt_embeds)
+        negative_mask = prompt_mask.clone()
+    else:
+        negative_embeds, negative_mask = pipe.encode_prompt(negative_prompt, **encode_kwargs)
+    return {
+        "prompt_embeds": prompt_embeds.detach().cpu(),
+        "prompt_mask": prompt_mask.detach().cpu(),
+        "negative_prompt_embeds": negative_embeds.detach().cpu(),
+        "negative_prompt_mask": negative_mask.detach().cpu(),
+    }
+
+
+def _inner_diffusers_pipe(pipe: Any) -> Any:
+    return getattr(pipe, "diffusers_pipe", pipe)
+
+
+def _cache_ti2v_prompt_conditions(
+    pipe: Any,
+    prompt: str,
+    negative_prompt: str,
+    image: Image.Image,
+    *,
+    height: int,
+    width: int,
+    device: torch.device,
+    null_cond_clone_zero: bool,
+) -> tuple[dict[str, torch.Tensor], torch.Tensor]:
+    condition_pipe = _inner_diffusers_pipe(pipe)
+    if not hasattr(condition_pipe, "preprocess_image") or not hasattr(condition_pipe, "_vlm_image"):
+        raise ValueError("TI2V condition reuse requires a LingBotVideo image-to-video pipeline.")
+    pixel = condition_pipe.preprocess_image(image, height, width).to(
+        device=device,
+        dtype=torch.float32,
+    )
+    vlm_image = condition_pipe._vlm_image(pixel)
+    return (
+        _cache_prompt_conditions(
+            condition_pipe,
+            prompt,
+            negative_prompt,
+            device=device,
+            null_cond_clone_zero=null_cond_clone_zero,
+            images=[vlm_image],
+        ),
+        pixel.detach().cpu(),
+    )
+
+
+def _condition_call_kwargs(
+    cache: dict[str, torch.Tensor] | None,
+    device: torch.device,
+) -> dict[str, torch.Tensor]:
+    if not cache:
+        return {}
+    return {key: value.to(device=device) for key, value in cache.items()}
+
+
+def _pipeline_class_for_mode(mode: str) -> Any:
+    if mode == "ti2v":
+        if LingBotVideoImageToVideoPipeline is None:
+            raise RuntimeError(
+                "LingBotVideoImageToVideoPipeline is not importable."
+            ) from _LINGBOT_PIPELINE_IMPORT_ERROR
+        return LingBotVideoImageToVideoPipeline
+    if LingBotVideoPipeline is None:
+        raise RuntimeError("LingBotVideoPipeline is not importable.") from _LINGBOT_PIPELINE_IMPORT_ERROR
+    return LingBotVideoPipeline
+
+
+def _load_transformer_component(
+    model_dir: Path,
+    transformer_subfolder: str,
+    dtype_map: dict[str, torch.dtype],
+) -> Any:
+    if LingBotVideoTransformer3DModel is None:
+        raise RuntimeError(
+            "LingBotVideoTransformer3DModel is not importable."
+        ) from _TRANSFORMER_IMPORT_ERROR
+    if model_component_dir is None:
+        raise RuntimeError("model path helpers are not importable.") from _LINGBOT_PIPELINE_IMPORT_ERROR
+    model_component_dir(model_dir, transformer_subfolder)
+    transformer_dtype = dtype_map.get(
+        "transformer",
+        dtype_map.get("default", torch.float32),
+    )
+    _disable_external_progress_bars()
+    _log_progress(
+        f"loading transformer subfolder={transformer_subfolder} "
+        f"dtype={_dtype_name(transformer_dtype)} model_dir={model_dir}"
+    )
+    transformer = LingBotVideoTransformer3DModel.from_pretrained(
+        str(model_dir),
+        subfolder=transformer_subfolder,
+        torch_dtype=transformer_dtype,
+    )
+    _log_progress(f"loaded transformer subfolder={transformer_subfolder}")
+    return transformer
+
+
+def _move_pipeline_aux_modules_to_device(pipe: Any, device: torch.device) -> Any:
+    inner_pipe = _inner_diffusers_pipe(pipe)
+    for name in ("text_encoder", "vae"):
+        module = getattr(inner_pipe, name, None)
+        if isinstance(module, torch.nn.Module):
+            _log_progress(f"moving {name} to {device}")
+            module.to(device)
+    return pipe
+
+
+def _load_diffusers_pipe(
+    model_dir: Path,
+    dtype_map: dict[str, torch.dtype],
+    mode: str,
+    transformer_subfolder: str,
+    defer_transformer_to_device: bool = False,
+) -> Any:
+    pipeline_class = _pipeline_class_for_mode(mode)
+    transformer = _load_transformer_component(model_dir, transformer_subfolder, dtype_map)
+
+    _disable_external_progress_bars()
+    _log_progress(f"loading pipeline mode={mode} model_dir={model_dir}")
+    with _patch_qwen3vl_from_pretrained():
+        pipe = pipeline_class.from_pretrained(
+            str(model_dir),
+            transformer=transformer,
+            trust_remote_code=True,
+            torch_dtype=dtype_map,
+        )
+    _log_progress(f"loaded pipeline mode={mode}")
+    device = _default_device()
+    if defer_transformer_to_device:
+        return _move_pipeline_aux_modules_to_device(pipe, device)
+    _log_progress(f"moving pipeline to {device}")
+    return pipe.to(device)
+
+
+def _load_sglang_native_pipe(
+    model_dir: Path,
+    dtype_map: dict[str, torch.dtype],
+    mode: str,
+    transformer_subfolder: str,
+    defer_transformer_to_device: bool = False,
+) -> Any:
+    if LingBotVideoNativePipeline is None or register_lingbot_native_pipeline is None:
+        raise RuntimeError(
+            "LingBotVideo SGLang native backend is not importable."
+        ) from _NATIVE_BACKEND_IMPORT_ERROR
+    _restore_cuda_sdp_state(_BASELINE_CUDA_SDP_STATE)
+    try:
+        register_lingbot_native_pipeline()
+        diffusers_pipe = _load_diffusers_pipe(
+            model_dir,
+            dtype_map,
+            mode=mode,
+            transformer_subfolder=transformer_subfolder,
+            defer_transformer_to_device=defer_transformer_to_device,
+        )
+        return LingBotVideoNativePipeline.from_diffusers_pipe(
+            diffusers_pipe,
+            model_path=model_dir,
+        )
+    finally:
+        _restore_cuda_sdp_state(_BASELINE_CUDA_SDP_STATE)
+
+
+def _load_pipe(
+    args: argparse.Namespace,
+    dtype_map: dict[str, torch.dtype],
+    *,
+    defer_transformer_to_device: bool = False,
+):
+    model_dir = Path(args.model_dir).resolve()
+    if args.engine == "diffusers":
+        return (
+            _load_diffusers_pipe(
+                model_dir,
+                dtype_map,
+                mode=args.mode,
+                transformer_subfolder=args.transformer_subfolder,
+                defer_transformer_to_device=defer_transformer_to_device,
+            ),
+            "diffusers-reference",
+        )
+    if args.engine == "sglang-native":
+        return (
+            _load_sglang_native_pipe(
+                model_dir,
+                dtype_map,
+                mode=args.mode,
+                transformer_subfolder=args.transformer_subfolder,
+                defer_transformer_to_device=defer_transformer_to_device,
+            ),
+            "sglang-native",
+        )
+    raise ValueError(f"unsupported engine: {args.engine}")
+
+
+def _refiner_model_available(
+    model_dir: str | None,
+    transformer_subfolder: str = "refiner",
+) -> bool:
+    if not model_dir:
+        return False
+    root = Path(model_dir)
+    if not root.is_dir():
+        return False
+    required = ("model_index.json", "scheduler", "text_encoder", "vae", transformer_subfolder)
+    return all((root / name).exists() for name in required)
+
+
+def _refiner_skip_reason(
+    requested: bool,
+    model_dir: str | None,
+    transformer_subfolder: str = "refiner",
+) -> str | None:
+    if not requested:
+        return None
+    if not model_dir:
+        return "missing_refiner_model_dir"
+    if not _refiner_model_available(model_dir, transformer_subfolder):
+        return f"missing_refiner_component:{transformer_subfolder}"
+    return None
+
+
+def _maybe_preload_refiner(
+    args: argparse.Namespace,
+    dtype_map: dict[str, torch.dtype],
+    device: torch.device,
+    context_parallel_rank: int,
+    context_parallel_mesh: DeviceMesh | None,
+    defer_transformer_to_device: bool,
+) -> dict[str, Any]:
+    requested = bool(args.run_refiner)
+    available = (
+        _refiner_model_available(args.refiner_model_dir, args.refiner_transformer_subfolder)
+        if requested
+        else False
+    )
+    state: dict[str, Any] = {
+        "refiner_requested": requested,
+        "refiner_available": available,
+        "refiner_skipped_reason": _refiner_skip_reason(
+            requested,
+            args.refiner_model_dir,
+            args.refiner_transformer_subfolder,
+        ),
+        "pipe": None,
+        "engine_name": None,
+        "component_dtypes": None,
+    }
+    if not requested or not available:
+        return state
+
+    refiner_args = argparse.Namespace(**vars(args))
+    refiner_args.model_dir = args.refiner_model_dir
+    refiner_args.transformer_subfolder = args.refiner_transformer_subfolder
+    refiner_args.vae_dtype = args.refiner_vae_dtype
+    # The refiner always samples through the text-only t2v pipeline; ti2v
+    # first-frame conditioning is injected as a clean frame-0 latent
+    # (`cond_latent`), not through the image-to-video pipeline.
+    refiner_args.mode = "t2v"
+    refiner_dtype_map = dict(dtype_map)
+    refiner_dtype_map["vae"] = _parse_dtype(args.refiner_vae_dtype)
+    refiner_pipe, refiner_engine_name = _load_pipe(
+        refiner_args,
+        refiner_dtype_map,
+        defer_transformer_to_device=defer_transformer_to_device,
+    )
+    refiner_component_dtypes = _component_dtypes(refiner_pipe)
+    expected_refiner_vae_dtype = _dtype_name(refiner_dtype_map["vae"])
+    if refiner_component_dtypes.get("vae") != expected_refiner_vae_dtype:
+        raise AssertionError(
+            f"Refiner VAE dtype mismatch: requested {expected_refiner_vae_dtype}, "
+            f"got {refiner_component_dtypes}"
+        )
+    if args.context_parallel_degree > 1:
+        _enable_context_parallel(
+            refiner_pipe.transformer,
+            args.context_parallel_degree,
+            args.context_parallel_ulysses_anything,
+            context_parallel_rank,
+            device,
+            context_parallel_mesh,
+            args.cfg_parallel_degree == 1,
+        )
+    state.update(
+        {
+            "pipe": refiner_pipe,
+            "engine_name": refiner_engine_name,
+            "component_dtypes": refiner_component_dtypes,
+        }
+    )
+    return state
+
+
+def _extract_frames(result: Any) -> np.ndarray:
+    frames = result.frames if hasattr(result, "frames") else result[0]
+    if isinstance(frames, torch.Tensor):
+        tensor = frames.detach().cpu().float()
+        if tensor.ndim == 5:
+            # Accept B,T,C,H,W or B,C,T,H,W.
+            if tensor.shape[2] in (1, 3, 4):
+                tensor = tensor[0].permute(0, 2, 3, 1)
+            else:
+                tensor = tensor[0].permute(1, 2, 3, 0)
+        elif tensor.ndim == 4:
+            if tensor.shape[1] in (1, 3, 4):
+                tensor = tensor.permute(0, 2, 3, 1)
+        return tensor.numpy()
+    if isinstance(frames, list):
+        return np.asarray(frames[0])
+    return np.asarray(frames)
+
+
+def _save_frames(
+    frames: np.ndarray,
+    mode: str,
+    output: Path,
+    fps: int,
+) -> None:
+    output.parent.mkdir(parents=True, exist_ok=True)
+
+    if mode == "t2i":
+        Image.fromarray((frames[0] * 255).clip(0, 255).astype(np.uint8)).save(output)
+    else:
+        if export_to_video is None:
+            raise RuntimeError("diffusers.utils.export_to_video is not importable.") from _DIFFUSERS_IMPORT_ERROR
+        export_to_video(frames, str(output), fps=fps)
+
+    print(f"saved {output} shape={tuple(frames.shape)}", flush=True)
+
+
+def _enable_context_parallel(
+    transformer: torch.nn.Module,
+    degree: int,
+    ulysses_anything: bool,
+    rank: int,
+    device: torch.device,
+    mesh: DeviceMesh | None,
+    use_native_mesh: bool,
+) -> None:
+    if degree <= 1:
+        return
+    if (
+        ContextParallelConfig is None
+        or ParallelConfig is None
+        or apply_context_parallel is None
+        or Attention is None
+        or MochiAttention is None
+        or AttentionModuleMixin is None
+    ):
+        raise RuntimeError("diffusers context-parallel helpers are unavailable.") from _DIFFUSERS_IMPORT_ERROR
+    cp_config = ContextParallelConfig(
+        ulysses_degree=degree,
+        ulysses_anything=ulysses_anything,
+    )
+    if use_native_mesh:
+        transformer.enable_parallelism(config=cp_config)
+        return
+
+    if mesh is None:
+        raise ValueError("A context-parallel mesh is required when context_parallel_degree > 1.")
+
+    config = ParallelConfig(context_parallel_config=cp_config)
+    config.setup(rank, degree, device, mesh=mesh)
+    if cp_config.ring_degree == 1 and int(mesh.mesh.numel()) != degree:
+        cp_config._flattened_mesh = cp_config._ulysses_mesh._flatten()
+    transformer._parallel_config = config
+
+    attention_classes = (Attention, MochiAttention, AttentionModuleMixin)
+    for module in transformer.modules():
+        if not isinstance(module, attention_classes):
+            continue
+        processor = module.processor
+        if processor is not None and hasattr(processor, "_parallel_config"):
+            processor._parallel_config = config
+
+    cp_plan = getattr(transformer, "_cp_plan", None)
+    if cp_plan is None:
+        raise ValueError("Transformer does not define a context-parallel plan.")
+    apply_context_parallel(transformer, cp_config, cp_plan)
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--model_dir", required=True, help="LingBot-Video model root directory.")
+    parser.add_argument(
+        "--backend",
+        choices=["diffusers", "sglang"],
+        default=None,
+        help=(
+            "public backend selector; `sglang` uses SGLang Diffusion when "
+            "available and falls back to diffusers otherwise"
+        ),
+    )
+    parser.add_argument(
+        "--engine",
+        choices=["sglang-native", "diffusers"],
+        default=None,
+        help=argparse.SUPPRESS,
+    )
+    parser.add_argument("--mode", choices=["t2i", "t2v", "ti2v"], required=True)
+    parser.add_argument("--prompt", default=None)
+    parser.add_argument("--prompt_json", default=None)
+    parser.add_argument("--negative_prompt", default=None,
+                        help="negative prompt; if unset, uses the mode's default "
+                             "(image default for t2i, video default for t2v/ti2v)")
+    parser.add_argument(
+        "--negative_prompt_json",
+        default=None,
+        help="path to the JSON file produced by rewriter/auto_negative.py",
+    )
+    parser.add_argument("--image", default=None)
+    parser.add_argument("--output", required=True)
+    parser.add_argument("--resolution", default=None)
+    parser.add_argument("--ratio", default=None)
+    parser.add_argument("--duration", type=float, default=None)
+    parser.add_argument("--height", type=int, default=192)
+    parser.add_argument("--width", type=int, default=320)
+    parser.add_argument("--num_frames", type=int, default=9)
+    parser.add_argument("--steps", type=int, default=40)
+    parser.add_argument("--guidance_scale", type=float, default=3.0)
+    parser.add_argument("--shift", type=float, default=3.0)
+    parser.add_argument("--seed", type=int, default=42)
+    parser.add_argument("--fps", type=int, default=24)
+    parser.add_argument("--default_dtype", default="bf16")
+    parser.add_argument("--transformer_dtype", default="bf16")
+    parser.add_argument("--transformer_subfolder", default="transformer")
+    parser.add_argument("--text_encoder_dtype", default="bf16")
+    parser.add_argument("--vae_dtype", default="fp32")
+    parser.add_argument("--diffusers_attn_backend", default=os.environ.get("DIFFUSERS_ATTN_BACKEND", ""))
+    parser.add_argument("--allow_tf32", action=argparse.BooleanOptionalAction, default=True)
+    parser.add_argument(
+        "--quiet_progress",
+        action="store_true",
+        help="Disable model-loading logs and denoising progress bars.",
+    )
+    parser.add_argument("--cfg_parallel_degree", type=int, default=1)
+    parser.add_argument("--context_parallel_degree", type=int, default=1)
+    parser.add_argument("--context_parallel_ulysses_anything", action="store_true")
+    parser.add_argument(
+        "--enable_fsdp_inference",
+        action="store_true",
+        help="Shard the base/refiner DiT transformers with PyTorch composable FSDP2.",
+    )
+    parser.add_argument("--batch_cfg", action="store_true")
+    parser.add_argument("--null_cond_clone_zero", action="store_true")
+    parser.add_argument("--reuse_condition_features", action="store_true")
+    parser.add_argument("--run_refiner", action="store_true")
+    parser.add_argument("--refiner_model_dir", default=None)
+    parser.add_argument("--refiner_transformer_subfolder", default="refiner")
+    parser.add_argument("--refiner_output", default=None)
+    parser.add_argument("--refiner_height", type=int, default=1088)
+    parser.add_argument("--refiner_width", type=int, default=1920)
+    parser.add_argument("--refiner_steps", type=int, default=8)
+    parser.add_argument("--refiner_guidance_scale", type=float, default=3.0)
+    parser.add_argument("--refiner_shift", type=float, default=3.0)
+    parser.add_argument("--refiner_t_thresh", type=float, default=0.85)
+    parser.add_argument("--refiner_sigma_tail_steps", type=int, default=2)
+    parser.add_argument("--refiner_fps", type=int, default=24)
+    parser.add_argument("--refiner_sample_fps", type=int, default=24)
+    parser.add_argument("--refiner_max_video_frames", type=int, default=None)
+    parser.add_argument("--refiner_vae_dtype", default="fp32")
+    parser.add_argument("--refiner_batch_cfg", action="store_true")
+    parser.add_argument("--refiner_no_null_cond_clone_zero", action="store_true")
+    parser.add_argument("--refiner_offload_vae_during_denoise", action="store_true")
+    args = parser.parse_args()
+    if args.quiet_progress:
+        os.environ["LINGBOT_QUIET_PROGRESS"] = "1"
+    args.engine = resolve_backend_engine(
+        engine=args.engine,
+        backend=args.backend,
+        stderr=sys.stderr,
+    )
+    args.negative_prompt = resolve_negative_prompt_arg(
+        args.negative_prompt,
+        args.negative_prompt_json,
+    )
+
+    prompt_sample = None
+    if args.prompt_json:
+        prompt_sample = _load_prompt_sample(Path(args.prompt_json))
+        args.prompt = _caption_from_sample(prompt_sample)
+        if args.duration is None and "duration" in prompt_sample:
+            args.duration = float(prompt_sample["duration"])
+    if args.prompt is None:
+        raise ValueError(
+            "Provide `--prompt_json <path_to_structured_prompt.json>` or an explicit `--prompt`."
+        )
+    if args.resolution or args.ratio:
+        if not args.resolution or not args.ratio:
+            raise ValueError("`--resolution` and `--ratio` must be provided together.")
+        args.height, args.width = _height_width_from_bucket(args.resolution, args.ratio)
+    if args.negative_prompt is None:
+        args.negative_prompt = (
+            DEFAULT_NEGATIVE_PROMPT_IMAGE if args.mode == "t2i" else DEFAULT_NEGATIVE_PROMPT
+        )
+    if args.mode != "t2i" and args.duration is not None:
+        if num_frames_from_duration is None:
+            raise RuntimeError("num_frames_from_duration is not importable.") from _LINGBOT_PIPELINE_IMPORT_ERROR
+        args.num_frames = num_frames_from_duration(args.duration, args.fps)
+    args.run_refiner = bool(args.run_refiner or args.refiner_model_dir)
+    if args.run_refiner:
+        if effective_refiner_model_dir is None:
+            raise RuntimeError("model path helpers are not importable.") from _LINGBOT_PIPELINE_IMPORT_ERROR
+        args.refiner_model_dir = str(effective_refiner_model_dir(args))
+
+    (
+        rank,
+        local_rank,
+        world_size,
+        cfg_parallel_group,
+        context_parallel_mesh,
+        cfg_branch_rank,
+        context_parallel_rank,
+    ) = _init_parallel(
+        args.cfg_parallel_degree,
+        args.context_parallel_degree,
+        args.enable_fsdp_inference,
+    )
+    if args.cfg_parallel_degree > 1 and args.batch_cfg:
+        raise ValueError("`--cfg_parallel_degree > 1` and `--batch_cfg` are mutually exclusive.")
+    if args.refiner_model_dir and args.cfg_parallel_degree > 1 and args.refiner_batch_cfg:
+        raise ValueError(
+            "`--cfg_parallel_degree > 1` and `--refiner_batch_cfg` are mutually exclusive."
+        )
+    if args.run_refiner and args.mode == "t2i":
+        raise ValueError("Refiner is only supported for video modes.")
+    if args.run_refiner and args.mode == "ti2v" and not args.image:
+        raise ValueError(
+            "The ti2v refiner conditions on the clean first frame; `--image` is required."
+        )
+    if args.diffusers_attn_backend:
+        os.environ["DIFFUSERS_ATTN_BACKEND"] = args.diffusers_attn_backend
+    if args.allow_tf32:
+        torch.backends.cuda.matmul.allow_tf32 = True
+        torch.set_float32_matmul_precision("high")
+    if args.mode == "t2i":
+        args.num_frames = 1
+
+    fsdp_mesh = None
+    if args.enable_fsdp_inference:
+        if init_fsdp_inference_mesh is None:
+            raise RuntimeError("FSDP inference helpers are not importable.") from _LINGBOT_PIPELINE_IMPORT_ERROR
+        fsdp_mesh = init_fsdp_inference_mesh()
+
+    dtype_map = _make_dtype_map(args)
+    defer_transformer_to_device = fsdp_mesh is not None
+    pipe, engine_name = _load_pipe(
+        args,
+        dtype_map,
+        defer_transformer_to_device=defer_transformer_to_device,
+    )
+    _configure_pipeline_logs(pipe)
+    component_dtypes = _component_dtypes(pipe)
+    expected_vae_dtype = _dtype_name(dtype_map["vae"])
+    if component_dtypes.get("vae") != expected_vae_dtype:
+        raise AssertionError(
+            f"VAE dtype mismatch: requested {expected_vae_dtype}, got {component_dtypes}"
+        )
+    device = _default_device()
+    if args.context_parallel_degree > 1:
+        _enable_context_parallel(
+            pipe.transformer,
+            args.context_parallel_degree,
+            args.context_parallel_ulysses_anything,
+            context_parallel_rank,
+            device,
+            context_parallel_mesh,
+            args.cfg_parallel_degree == 1,
+        )
+    base_fsdp_info = _apply_fsdp_inference_if_requested(pipe, args.enable_fsdp_inference, fsdp_mesh)
+
+    refiner_state = _maybe_preload_refiner(
+        args,
+        dtype_map,
+        device,
+        context_parallel_rank,
+        context_parallel_mesh,
+        defer_transformer_to_device,
+    )
+    if refiner_state["pipe"] is not None:
+        _configure_pipeline_logs(refiner_state["pipe"])
+    refiner_fsdp_info = None
+    if refiner_state["pipe"] is not None:
+        refiner_fsdp_info = _apply_fsdp_inference_if_requested(
+            refiner_state["pipe"],
+            args.enable_fsdp_inference,
+            fsdp_mesh,
+        )
+    if rank == 0 and refiner_state["refiner_requested"] and not refiner_state["refiner_available"]:
+        print(
+            "WARNING: refiner requested but unavailable; "
+            f"reason={refiner_state['refiner_skipped_reason']} "
+            f"refiner_model_dir={args.refiner_model_dir}",
+            flush=True,
+        )
+
+    generator = torch.Generator(device=device).manual_seed(args.seed)
+    condition_cache = None
+    input_image = None
+    ti2v_image_tensor_cpu = None
+    should_cache_conditions = (
+        args.reuse_condition_features
+        or bool(refiner_state["refiner_available"])
+        or args.batch_cfg
+        or args.null_cond_clone_zero
+    )
+    if args.mode == "ti2v":
+        input_image = (
+            Image.open(args.image).convert("RGB")
+            if args.image
+            else _make_default_image(args.height, args.width)
+        )
+        if should_cache_conditions:
+            condition_cache, ti2v_image_tensor_cpu = _cache_ti2v_prompt_conditions(
+                pipe,
+                args.prompt,
+                args.negative_prompt,
+                input_image,
+                height=args.height,
+                width=args.width,
+                device=device,
+                null_cond_clone_zero=args.null_cond_clone_zero,
+            )
+    elif should_cache_conditions:
+        condition_cache = _cache_prompt_conditions(
+            pipe,
+            args.prompt,
+            args.negative_prompt,
+            device=device,
+            null_cond_clone_zero=args.null_cond_clone_zero,
+        )
+    call_kwargs = dict(
+        prompt=args.prompt,
+        negative_prompt=args.negative_prompt,
+        height=args.height,
+        width=args.width,
+        num_frames=args.num_frames,
+        num_inference_steps=args.steps,
+        guidance_scale=args.guidance_scale,
+        shift=args.shift,
+        generator=generator,
+        output_type="np",
+        **_condition_call_kwargs(condition_cache, device),
+    )
+    call_kwargs["batch_cfg"] = args.batch_cfg
+    if args.mode != "ti2v":
+        call_kwargs["null_cond_clone_zero"] = args.null_cond_clone_zero
+    if args.mode == "ti2v":
+        call_kwargs["image"] = input_image
+        if ti2v_image_tensor_cpu is not None:
+            call_kwargs["image_tensor"] = ti2v_image_tensor_cpu.to(device=device)
+    if args.cfg_parallel_degree > 1:
+        call_kwargs["cfg_parallel_group"] = cfg_parallel_group
+
+    if rank == 0:
+        print(
+            "runtime "
+            f"engine={engine_name} model_dir={args.model_dir} mode={args.mode} device={device} "
+            f"rank={rank}/{world_size} local_rank={local_rank} "
+            f"cfg_branch_rank={cfg_branch_rank} context_parallel_rank={context_parallel_rank} "
+            f"cfg_parallel_degree={args.cfg_parallel_degree} "
+            f"context_parallel_degree={args.context_parallel_degree} "
+            f"height={args.height} width={args.width} frames={args.num_frames} steps={args.steps} "
+            f"guidance={args.guidance_scale} shift={args.shift} seed={args.seed} "
+            f"attn_backend={os.environ.get('DIFFUSERS_ATTN_BACKEND')} "
+            f"allow_tf32={torch.backends.cuda.matmul.allow_tf32} "
+            f"fsdp_inference={base_fsdp_info} "
+            f"component_dtypes={component_dtypes}",
+            flush=True,
+        )
+
+    with torch.no_grad():
+        result = pipe(**call_kwargs)
+    frames = _extract_frames(result) if rank == 0 else None
+    if rank == 0:
+        _save_frames(
+            frames,
+            args.mode,
+            Path(args.output),
+            args.fps,
+        )
+
+    if not refiner_state["refiner_requested"] or not refiner_state["refiner_available"]:
+        _sync_parallel_if_needed()
+        _destroy_parallel_if_needed()
+        return
+
+    _sync_parallel_if_needed()
+    del result, frames, call_kwargs
+    gc.collect()
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
+
+    if load_refiner_video_tensor is None or prepare_refiner_latent is None:
+        raise RuntimeError("Refiner helpers are not importable.") from _LINGBOT_PIPELINE_IMPORT_ERROR
+    refiner_pipe = refiner_state["pipe"]
+    if refiner_pipe is None:
+        raise RuntimeError("Refiner was marked available but was not preloaded.")
+    refiner_engine_name = str(refiner_state["engine_name"])
+    refiner_component_dtypes = refiner_state["component_dtypes"]
+
+    refiner_generator = torch.Generator(device=device).manual_seed(args.seed)
+    lowres_video, lowres_meta = load_refiner_video_tensor(
+        args.output,
+        args.refiner_height,
+        args.refiner_width,
+        sample_fps=args.refiner_sample_fps,
+        vae_tc=getattr(refiner_pipe, "vae_scale_factor_temporal", 4),
+        max_frames=args.refiner_max_video_frames,
+    )
+
+    # TI2V always refines with the clean input first frame as a fixed frame-0
+    # latent: injected into the initial latent here and re-clamped after every
+    # scheduler step via `cond_latent`.
+    first_frame_condition_enabled = args.mode == "ti2v"
+    refiner_cond_latent = None
+    with torch.no_grad():
+        x_up = refiner_pipe.encode_video_latent(lowres_video, generator=refiner_generator)
+        if first_frame_condition_enabled:
+            condition_pixels = load_first_frame_condition_tensor(
+                args.image,
+                args.refiner_height,
+                args.refiner_width,
+                geometry_height=args.height,
+                geometry_width=args.width,
+            )
+            clean_x0 = refiner_pipe.encode_video_latent(
+                condition_pixels,
+                generator=refiner_generator,
+            )
+            refiner_cond_latent = clean_x0[:, :, 0:1].contiguous()
+            x_up[:, :, 0:1] = refiner_cond_latent.to(dtype=x_up.dtype)
+            del condition_pixels, clean_x0
+        noise = torch.randn(
+            x_up.shape,
+            device=device,
+            dtype=x_up.dtype,
+            generator=refiner_generator,
+        )
+        initial_latent = prepare_refiner_latent(x_up, noise, args.refiner_t_thresh)
+    del lowres_video, x_up, noise
+    refiner_null_clone_zero = not args.refiner_no_null_cond_clone_zero
+    # The ti2v condition cache carries image tokens; the refiner conditions on
+    # text only, so it rebuilds t2v-style embeddings instead of reusing it.
+    refiner_condition_cache = None if args.mode == "ti2v" else condition_cache
+    should_cache_refiner_conditions = not (
+        args.cfg_parallel_degree > 1 and cfg_branch_rank == 1 and not refiner_null_clone_zero
+    )
+    if refiner_condition_cache is None and should_cache_refiner_conditions:
+        refiner_condition_cache = _cache_prompt_conditions(
+            refiner_pipe,
+            args.prompt,
+            args.negative_prompt,
+            device=device,
+            null_cond_clone_zero=refiner_null_clone_zero,
+        )
+    elif refiner_null_clone_zero:
+        refiner_condition_cache = dict(refiner_condition_cache)
+        refiner_condition_cache["negative_prompt_embeds"] = torch.zeros_like(
+            refiner_condition_cache["prompt_embeds"]
+        )
+        refiner_condition_cache["negative_prompt_mask"] = refiner_condition_cache["prompt_mask"].clone()
+
+    if rank == 0:
+        print(
+            "runtime "
+            f"engine={refiner_engine_name} model_dir={args.refiner_model_dir} mode=refiner "
+            f"height={args.refiner_height} width={args.refiner_width} "
+            f"frames={lowres_meta['sample_frame']} steps={args.refiner_steps} "
+            f"guidance={args.refiner_guidance_scale} shift={args.refiner_shift} "
+            f"t_thresh={args.refiner_t_thresh} tail_steps={args.refiner_sigma_tail_steps} "
+            f"seed={args.seed} "
+            f"batch_cfg={args.refiner_batch_cfg} "
+            f"first_frame_condition={first_frame_condition_enabled} "
+            f"fsdp_inference={refiner_fsdp_info} "
+            f"component_dtypes={refiner_component_dtypes}",
+            flush=True,
+        )
+
+    refiner_call_kwargs = dict(
+        prompt=args.prompt,
+        negative_prompt=args.negative_prompt,
+        height=args.refiner_height,
+        width=args.refiner_width,
+        num_frames=int(lowres_meta["sample_frame"]),
+        num_inference_steps=args.refiner_steps,
+        guidance_scale=args.refiner_guidance_scale,
+        shift=args.refiner_shift,
+        t_thresh=args.refiner_t_thresh,
+        refiner_sigma_tail_steps=args.refiner_sigma_tail_steps,
+        generator=refiner_generator,
+        latents=initial_latent,
+        output_type="np",
+        batch_cfg=args.refiner_batch_cfg,
+        null_cond_clone_zero=refiner_null_clone_zero,
+        offload_vae_during_denoise=args.refiner_offload_vae_during_denoise,
+        **_condition_call_kwargs(refiner_condition_cache, device),
+    )
+    if refiner_cond_latent is not None:
+        refiner_call_kwargs["cond_latent"] = refiner_cond_latent
+    if args.cfg_parallel_degree > 1:
+        refiner_call_kwargs["cfg_parallel_group"] = cfg_parallel_group
+
+    with torch.no_grad():
+        refiner_result = refiner_pipe(**refiner_call_kwargs)
+    if rank == 0:
+        refiner_frames = _extract_frames(refiner_result)
+        _save_frames(
+            refiner_frames,
+            args.mode,
+            Path(args.refiner_output or str(Path(args.output).with_name(Path(args.output).stem + "_refined.mp4"))),
+            args.refiner_fps,
+        )
+
+    _sync_parallel_if_needed()
+    _destroy_parallel_if_needed()
+
+
+if __name__ == "__main__":
+    main()

lingbot_video/scheduling_flow_unipc.py

@@ -0,0 +1,796 @@
+import math
+from typing import List, Optional, Tuple, Union
+
+import numpy as np
+import torch
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+from diffusers.schedulers.scheduling_utils import (KarrasDiffusionSchedulers,
+                                                   SchedulerMixin,
+                                                   SchedulerOutput)
+from diffusers.utils import deprecate, is_scipy_available
+
+if is_scipy_available():
+    import scipy.stats
+
+
+class FlowUniPCMultistepScheduler(SchedulerMixin, ConfigMixin):
+    """
+    `UniPCMultistepScheduler` is a training-free framework designed for the fast sampling of diffusion models.
+
+    This model inherits from [`SchedulerMixin`] and [`ConfigMixin`]. Check the superclass documentation for the generic
+    methods the library implements for all schedulers such as loading and saving.
+
+    Args:
+        num_train_timesteps (`int`, defaults to 1000):
+            The number of diffusion steps to train the model.
+        solver_order (`int`, default `2`):
+            The UniPC order which can be any positive integer. The effective order of accuracy is `solver_order + 1`
+            due to the UniC. It is recommended to use `solver_order=2` for guided sampling, and `solver_order=3` for
+            unconditional sampling.
+        prediction_type (`str`, defaults to "flow_prediction"):
+            Prediction type of the scheduler function; must be `flow_prediction` for this scheduler, which predicts
+            the flow of the diffusion process.
+        thresholding (`bool`, defaults to `False`):
+            Whether to use the "dynamic thresholding" method. This is unsuitable for latent-space diffusion models such
+            as Stable Diffusion.
+        dynamic_thresholding_ratio (`float`, defaults to 0.995):
+            The ratio for the dynamic thresholding method. Valid only when `thresholding=True`.
+        sample_max_value (`float`, defaults to 1.0):
+            The threshold value for dynamic thresholding. Valid only when `thresholding=True` and `predict_x0=True`.
+        predict_x0 (`bool`, defaults to `True`):
+            Whether to use the updating algorithm on the predicted x0.
+        solver_type (`str`, default `bh2`):
+            Solver type for UniPC. It is recommended to use `bh1` for unconditional sampling when steps < 10, and `bh2`
+            otherwise.
+        lower_order_final (`bool`, default `True`):
+            Whether to use lower-order solvers in the final steps. Only valid for < 15 inference steps. This can
+            stabilize the sampling of DPMSolver for steps < 15, especially for steps <= 10.
+        disable_corrector (`list`, default `[]`):
+            Decides which step to disable the corrector to mitigate the misalignment between `epsilon_theta(x_t, c)`
+            and `epsilon_theta(x_t^c, c)` which can influence convergence for a large guidance scale. Corrector is
+            usually disabled during the first few steps.
+        solver_p (`SchedulerMixin`, default `None`):
+            Any other scheduler that if specified, the algorithm becomes `solver_p + UniC`.
+        use_karras_sigmas (`bool`, *optional*, defaults to `False`):
+            Whether to use Karras sigmas for step sizes in the noise schedule during the sampling process. If `True`,
+            the sigmas are determined according to a sequence of noise levels {σi}.
+        use_exponential_sigmas (`bool`, *optional*, defaults to `False`):
+            Whether to use exponential sigmas for step sizes in the noise schedule during the sampling process.
+        timestep_spacing (`str`, defaults to `"linspace"`):
+            The way the timesteps should be scaled. Refer to Table 2 of the [Common Diffusion Noise Schedules and
+            Sample Steps are Flawed](https://huggingface.co/papers/2305.08891) for more information.
+        steps_offset (`int`, defaults to 0):
+            An offset added to the inference steps, as required by some model families.
+        final_sigmas_type (`str`, defaults to `"zero"`):
+            The final `sigma` value for the noise schedule during the sampling process. If `"sigma_min"`, the final
+            sigma is the same as the last sigma in the training schedule. If `zero`, the final sigma is set to 0.
+    """
+
+    _compatibles = [e.name for e in KarrasDiffusionSchedulers]
+    order = 1
+
+    @register_to_config
+    def __init__(
+            self,
+            num_train_timesteps: int = 1000,
+            solver_order: int = 2,
+            prediction_type: str = "flow_prediction",
+            shift: Optional[float] = 1.0,
+            use_dynamic_shifting=False,
+            thresholding: bool = False,
+            dynamic_thresholding_ratio: float = 0.995,
+            sample_max_value: float = 1.0,
+            predict_x0: bool = True,
+            solver_type: str = "bh2",
+            lower_order_final: bool = True,
+            disable_corrector: List[int] = [],
+            solver_p: SchedulerMixin = None,
+            timestep_spacing: str = "linspace",
+            steps_offset: int = 0,
+            final_sigmas_type: Optional[str] = "zero",  # "zero", "sigma_min"
+    ):
+
+        if solver_type not in ["bh1", "bh2"]:
+            if solver_type in ["midpoint", "heun", "logrho"]:
+                self.register_to_config(solver_type="bh2")
+            else:
+                raise NotImplementedError(
+                    f"{solver_type} is not implemented for {self.__class__}")
+
+        self.predict_x0 = predict_x0
+        # setable values
+        self.num_inference_steps = None
+        alphas = np.linspace(1, 1 / num_train_timesteps,
+                             num_train_timesteps)[::-1].copy()
+        sigmas = 1.0 - alphas
+        sigmas = torch.from_numpy(sigmas).to(dtype=torch.float32)
+
+        if not use_dynamic_shifting:
+            # when use_dynamic_shifting is True, we apply the timestep shifting on the fly based on the image resolution
+            sigmas = shift * sigmas / (1 +
+                                       (shift - 1) * sigmas)  # pyright: ignore
+
+        self.sigmas = sigmas
+        self.timesteps = sigmas * num_train_timesteps
+
+        self.model_outputs = [None] * solver_order
+        self.timestep_list = [None] * solver_order
+        self.lower_order_nums = 0
+        self.disable_corrector = disable_corrector
+        self.solver_p = solver_p
+        self.last_sample = None
+        self._step_index = None
+        self._begin_index = None
+
+        self.sigmas = self.sigmas.to(
+            "cpu")  # to avoid too much CPU/GPU communication
+        self.sigma_min = self.sigmas[-1].item()
+        self.sigma_max = self.sigmas[0].item()
+
+    @property
+    def step_index(self):
+        """
+        The index counter for current timestep. It will increase 1 after each scheduler step.
+        """
+        return self._step_index
+
+    @property
+    def begin_index(self):
+        """
+        The index for the first timestep. It should be set from pipeline with `set_begin_index` method.
+        """
+        return self._begin_index
+
+    # Copied from diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler.set_begin_index
+    def set_begin_index(self, begin_index: int = 0):
+        """
+        Sets the begin index for the scheduler. This function should be run from pipeline before the inference.
+
+        Args:
+            begin_index (`int`):
+                The begin index for the scheduler.
+        """
+        self._begin_index = begin_index
+
+    # Modified from diffusers.schedulers.scheduling_flow_match_euler_discrete.FlowMatchEulerDiscreteScheduler.set_timesteps
+    def set_timesteps(
+        self,
+        num_inference_steps: Union[int, None] = None,
+        device: Union[str, torch.device] = None,
+        sigmas: Optional[List[float]] = None,
+        mu: Optional[Union[float, None]] = None,
+        shift: Optional[Union[float, None]] = None,
+    ):
+        """
+        Sets the discrete timesteps used for the diffusion chain (to be run before inference).
+        Args:
+            num_inference_steps (`int`):
+                Total number of the spacing of the time steps.
+            device (`str` or `torch.device`, *optional*):
+                The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
+        """
+
+        if self.config.use_dynamic_shifting and mu is None:
+            raise ValueError(
+                " you have to pass a value for `mu` when `use_dynamic_shifting` is set to be `True`"
+            )
+
+        if sigmas is None:
+            sigmas = np.linspace(self.sigma_max, self.sigma_min,
+                                 num_inference_steps +
+                                 1).copy()[:-1]  # pyright: ignore
+
+        if self.config.use_dynamic_shifting:
+            sigmas = self.time_shift(mu, 1.0, sigmas)  # pyright: ignore
+        else:
+            if shift is None:
+                shift = self.config.shift
+            sigmas = shift * sigmas / (1 +
+                                       (shift - 1) * sigmas)  # pyright: ignore
+
+        if self.config.final_sigmas_type == "sigma_min":
+            sigma_last = ((1 - self.alphas_cumprod[0]) /
+                          self.alphas_cumprod[0])**0.5
+        elif self.config.final_sigmas_type == "zero":
+            sigma_last = 0
+        else:
+            raise ValueError(
+                f"`final_sigmas_type` must be one of 'zero', or 'sigma_min', but got {self.config.final_sigmas_type}"
+            )
+
+        timesteps = sigmas * self.config.num_train_timesteps
+        sigmas = np.concatenate([sigmas, [sigma_last]
+                                ]).astype(np.float32)  # pyright: ignore
+
+        self.sigmas = torch.from_numpy(sigmas)
+        self.timesteps = torch.from_numpy(timesteps).to(
+            device=device, dtype=torch.int64)
+
+        self.num_inference_steps = len(timesteps)
+
+        self.model_outputs = [
+            None,
+        ] * self.config.solver_order
+        self.lower_order_nums = 0
+        self.last_sample = None
+        if self.solver_p:
+            self.solver_p.set_timesteps(self.num_inference_steps, device=device)
+
+        # add an index counter for schedulers that allow duplicated timesteps
+        self._step_index = None
+        self._begin_index = None
+        self.sigmas = self.sigmas.to(
+            "cpu")  # to avoid too much CPU/GPU communication
+
+    # Copied from diffusers.schedulers.scheduling_ddpm.DDPMScheduler._threshold_sample
+    def _threshold_sample(self, sample: torch.Tensor) -> torch.Tensor:
+        """
+        "Dynamic thresholding: At each sampling step we set s to a certain percentile absolute pixel value in xt0 (the
+        prediction of x_0 at timestep t), and if s > 1, then we threshold xt0 to the range [-s, s] and then divide by
+        s. Dynamic thresholding pushes saturated pixels (those near -1 and 1) inwards, thereby actively preventing
+        pixels from saturation at each step. We find that dynamic thresholding results in significantly better
+        photorealism as well as better image-text alignment, especially when using very large guidance weights."
+
+        https://arxiv.org/abs/2205.11487
+        """
+        dtype = sample.dtype
+        batch_size, channels, *remaining_dims = sample.shape
+
+        if dtype not in (torch.float32, torch.float64):
+            sample = sample.float(
+            )  # upcast for quantile calculation, and clamp not implemented for cpu half
+
+        # Flatten sample for doing quantile calculation along each image
+        sample = sample.reshape(batch_size, channels * np.prod(remaining_dims))
+
+        abs_sample = sample.abs()  # "a certain percentile absolute pixel value"
+
+        s = torch.quantile(
+            abs_sample, self.config.dynamic_thresholding_ratio, dim=1)
+        s = torch.clamp(
+            s, min=1, max=self.config.sample_max_value
+        )  # When clamped to min=1, equivalent to standard clipping to [-1, 1]
+        s = s.unsqueeze(
+            1)  # (batch_size, 1) because clamp will broadcast along dim=0
+        sample = torch.clamp(
+            sample, -s, s
+        ) / s  # "we threshold xt0 to the range [-s, s] and then divide by s"
+
+        sample = sample.reshape(batch_size, channels, *remaining_dims)
+        sample = sample.to(dtype)
+
+        return sample
+
+    # Copied from diffusers.schedulers.scheduling_flow_match_euler_discrete.FlowMatchEulerDiscreteScheduler._sigma_to_t
+    def _sigma_to_t(self, sigma):
+        return sigma * self.config.num_train_timesteps
+
+    def _sigma_to_alpha_sigma_t(self, sigma):
+        return 1 - sigma, sigma
+
+    # Copied from diffusers.schedulers.scheduling_flow_match_euler_discrete.set_timesteps
+    def time_shift(self, mu: float, sigma: float, t: torch.Tensor):
+        return math.exp(mu) / (math.exp(mu) + (1 / t - 1)**sigma)
+
+    def convert_model_output(
+        self,
+        model_output: torch.Tensor,
+        *args,
+        sample: torch.Tensor = None,
+        **kwargs,
+    ) -> torch.Tensor:
+        r"""
+        Convert the model output to the corresponding type the UniPC algorithm needs.
+
+        Args:
+            model_output (`torch.Tensor`):
+                The direct output from the learned diffusion model.
+            timestep (`int`):
+                The current discrete timestep in the diffusion chain.
+            sample (`torch.Tensor`):
+                A current instance of a sample created by the diffusion process.
+
+        Returns:
+            `torch.Tensor`:
+                The converted model output.
+        """
+        timestep = args[0] if len(args) > 0 else kwargs.pop("timestep", None)
+        if sample is None:
+            if len(args) > 1:
+                sample = args[1]
+            else:
+                raise ValueError(
+                    "missing `sample` as a required keyward argument")
+        if timestep is not None:
+            deprecate(
+                "timesteps",
+                "1.0.0",
+                "Passing `timesteps` is deprecated and has no effect as model output conversion is now handled via an internal counter `self.step_index`",
+            )
+
+        sigma = self.sigmas[self.step_index]
+        alpha_t, sigma_t = self._sigma_to_alpha_sigma_t(sigma)
+
+        if self.predict_x0:
+            if self.config.prediction_type == "flow_prediction":
+                sigma_t = self.sigmas[self.step_index]
+                x0_pred = sample - sigma_t * model_output
+            else:
+                raise ValueError(
+                    f"prediction_type given as {self.config.prediction_type} must be one of `epsilon`, `sample`,"
+                    " `v_prediction` or `flow_prediction` for the UniPCMultistepScheduler."
+                )
+
+            if self.config.thresholding:
+                x0_pred = self._threshold_sample(x0_pred)
+
+            return x0_pred
+        else:
+            if self.config.prediction_type == "flow_prediction":
+                sigma_t = self.sigmas[self.step_index]
+                epsilon = sample - (1 - sigma_t) * model_output
+            else:
+                raise ValueError(
+                    f"prediction_type given as {self.config.prediction_type} must be one of `epsilon`, `sample`,"
+                    " `v_prediction` or `flow_prediction` for the UniPCMultistepScheduler."
+                )
+
+            if self.config.thresholding:
+                sigma_t = self.sigmas[self.step_index]
+                x0_pred = sample - sigma_t * model_output
+                x0_pred = self._threshold_sample(x0_pred)
+                epsilon = model_output + x0_pred
+
+            return epsilon
+
+    def multistep_uni_p_bh_update(
+        self,
+        model_output: torch.Tensor,
+        *args,
+        sample: torch.Tensor = None,
+        order: int = None,  # pyright: ignore
+        **kwargs,
+    ) -> torch.Tensor:
+        """
+        One step for the UniP (B(h) version). Alternatively, `self.solver_p` is used if is specified.
+
+        Args:
+            model_output (`torch.Tensor`):
+                The direct output from the learned diffusion model at the current timestep.
+            prev_timestep (`int`):
+                The previous discrete timestep in the diffusion chain.
+            sample (`torch.Tensor`):
+                A current instance of a sample created by the diffusion process.
+            order (`int`):
+                The order of UniP at this timestep (corresponds to the *p* in UniPC-p).
+
+        Returns:
+            `torch.Tensor`:
+                The sample tensor at the previous timestep.
+        """
+        prev_timestep = args[0] if len(args) > 0 else kwargs.pop(
+            "prev_timestep", None)
+        if sample is None:
+            if len(args) > 1:
+                sample = args[1]
+            else:
+                raise ValueError(
+                    " missing `sample` as a required keyward argument")
+        if order is None:
+            if len(args) > 2:
+                order = args[2]
+            else:
+                raise ValueError(
+                    " missing `order` as a required keyward argument")
+        if prev_timestep is not None:
+            deprecate(
+                "prev_timestep",
+                "1.0.0",
+                "Passing `prev_timestep` is deprecated and has no effect as model output conversion is now handled via an internal counter `self.step_index`",
+            )
+        model_output_list = self.model_outputs
+
+        s0 = self.timestep_list[-1]
+        m0 = model_output_list[-1]
+        x = sample
+
+        if self.solver_p:
+            x_t = self.solver_p.step(model_output, s0, x).prev_sample
+            return x_t
+
+        sigma_t, sigma_s0 = self.sigmas[self.step_index + 1], self.sigmas[
+            self.step_index]  # pyright: ignore
+        alpha_t, sigma_t = self._sigma_to_alpha_sigma_t(sigma_t)
+        alpha_s0, sigma_s0 = self._sigma_to_alpha_sigma_t(sigma_s0)
+
+        lambda_t = torch.log(alpha_t) - torch.log(sigma_t)
+        lambda_s0 = torch.log(alpha_s0) - torch.log(sigma_s0)
+
+        h = lambda_t - lambda_s0
+        device = sample.device
+
+        rks = []
+        D1s = []
+        for i in range(1, order):
+            si = self.step_index - i  # pyright: ignore
+            mi = model_output_list[-(i + 1)]
+            alpha_si, sigma_si = self._sigma_to_alpha_sigma_t(self.sigmas[si])
+            lambda_si = torch.log(alpha_si) - torch.log(sigma_si)
+            rk = (lambda_si - lambda_s0) / h
+            rks.append(rk)
+            D1s.append((mi - m0) / rk)  # pyright: ignore
+
+        rks.append(1.0)
+        rks = torch.tensor(rks, device=device)
+
+        R = []
+        b = []
+
+        hh = -h if self.predict_x0 else h
+        h_phi_1 = torch.expm1(hh)  # h\phi_1(h) = e^h - 1
+        h_phi_k = h_phi_1 / hh - 1
+
+        factorial_i = 1
+
+        if self.config.solver_type == "bh1":
+            B_h = hh
+        elif self.config.solver_type == "bh2":
+            B_h = torch.expm1(hh)
+        else:
+            raise NotImplementedError()
+
+        for i in range(1, order + 1):
+            R.append(torch.pow(rks, i - 1))
+            b.append(h_phi_k * factorial_i / B_h)
+            factorial_i *= i + 1
+            h_phi_k = h_phi_k / hh - 1 / factorial_i
+
+        R = torch.stack(R)
+        b = torch.tensor(b, device=device)
+
+        if len(D1s) > 0:
+            D1s = torch.stack(D1s, dim=1)  # (B, K)
+            # for order 2, we use a simplified version
+            if order == 2:
+                rhos_p = torch.tensor([0.5], dtype=x.dtype, device=device)
+            else:
+                rhos_p = torch.linalg.solve(R[:-1, :-1],
+                                            b[:-1]).to(device).to(x.dtype)
+        else:
+            D1s = None
+
+        if self.predict_x0:
+            x_t_ = sigma_t / sigma_s0 * x - alpha_t * h_phi_1 * m0
+            if D1s is not None:
+                pred_res = torch.einsum("k,bkc...->bc...", rhos_p,
+                                        D1s)  # pyright: ignore
+            else:
+                pred_res = 0
+            x_t = x_t_ - alpha_t * B_h * pred_res
+        else:
+            x_t_ = alpha_t / alpha_s0 * x - sigma_t * h_phi_1 * m0
+            if D1s is not None:
+                pred_res = torch.einsum("k,bkc...->bc...", rhos_p,
+                                        D1s)  # pyright: ignore
+            else:
+                pred_res = 0
+            x_t = x_t_ - sigma_t * B_h * pred_res
+
+        x_t = x_t.to(x.dtype)
+        return x_t
+
+    def multistep_uni_c_bh_update(
+        self,
+        this_model_output: torch.Tensor,
+        *args,
+        last_sample: torch.Tensor = None,
+        this_sample: torch.Tensor = None,
+        order: int = None,  # pyright: ignore
+        **kwargs,
+    ) -> torch.Tensor:
+        """
+        One step for the UniC (B(h) version).
+
+        Args:
+            this_model_output (`torch.Tensor`):
+                The model outputs at `x_t`.
+            this_timestep (`int`):
+                The current timestep `t`.
+            last_sample (`torch.Tensor`):
+                The generated sample before the last predictor `x_{t-1}`.
+            this_sample (`torch.Tensor`):
+                The generated sample after the last predictor `x_{t}`.
+            order (`int`):
+                The `p` of UniC-p at this step. The effective order of accuracy should be `order + 1`.
+
+        Returns:
+            `torch.Tensor`:
+                The corrected sample tensor at the current timestep.
+        """
+        this_timestep = args[0] if len(args) > 0 else kwargs.pop(
+            "this_timestep", None)
+        if last_sample is None:
+            if len(args) > 1:
+                last_sample = args[1]
+            else:
+                raise ValueError(
+                    " missing`last_sample` as a required keyward argument")
+        if this_sample is None:
+            if len(args) > 2:
+                this_sample = args[2]
+            else:
+                raise ValueError(
+                    " missing`this_sample` as a required keyward argument")
+        if order is None:
+            if len(args) > 3:
+                order = args[3]
+            else:
+                raise ValueError(
+                    " missing`order` as a required keyward argument")
+        if this_timestep is not None:
+            deprecate(
+                "this_timestep",
+                "1.0.0",
+                "Passing `this_timestep` is deprecated and has no effect as model output conversion is now handled via an internal counter `self.step_index`",
+            )
+
+        model_output_list = self.model_outputs
+
+        m0 = model_output_list[-1]
+        x = last_sample
+        x_t = this_sample
+        model_t = this_model_output
+
+        sigma_t, sigma_s0 = self.sigmas[self.step_index], self.sigmas[
+            self.step_index - 1]  # pyright: ignore
+        alpha_t, sigma_t = self._sigma_to_alpha_sigma_t(sigma_t)
+        alpha_s0, sigma_s0 = self._sigma_to_alpha_sigma_t(sigma_s0)
+
+        lambda_t = torch.log(alpha_t) - torch.log(sigma_t)
+        lambda_s0 = torch.log(alpha_s0) - torch.log(sigma_s0)
+
+        h = lambda_t - lambda_s0
+        device = this_sample.device
+
+        rks = []
+        D1s = []
+        for i in range(1, order):
+            si = self.step_index - (i + 1)  # pyright: ignore
+            mi = model_output_list[-(i + 1)]
+            alpha_si, sigma_si = self._sigma_to_alpha_sigma_t(self.sigmas[si])
+            lambda_si = torch.log(alpha_si) - torch.log(sigma_si)
+            rk = (lambda_si - lambda_s0) / h
+            rks.append(rk)
+            D1s.append((mi - m0) / rk)  # pyright: ignore
+
+        rks.append(1.0)
+        rks = torch.tensor(rks, device=device)
+
+        R = []
+        b = []
+
+        hh = -h if self.predict_x0 else h
+        h_phi_1 = torch.expm1(hh)  # h\phi_1(h) = e^h - 1
+        h_phi_k = h_phi_1 / hh - 1
+
+        factorial_i = 1
+
+        if self.config.solver_type == "bh1":
+            B_h = hh
+        elif self.config.solver_type == "bh2":
+            B_h = torch.expm1(hh)
+        else:
+            raise NotImplementedError()
+
+        for i in range(1, order + 1):
+            R.append(torch.pow(rks, i - 1))
+            b.append(h_phi_k * factorial_i / B_h)
+            factorial_i *= i + 1
+            h_phi_k = h_phi_k / hh - 1 / factorial_i
+
+        R = torch.stack(R)
+        b = torch.tensor(b, device=device)
+
+        if len(D1s) > 0:
+            D1s = torch.stack(D1s, dim=1)
+        else:
+            D1s = None
+
+        # for order 1, we use a simplified version
+        if order == 1:
+            rhos_c = torch.tensor([0.5], dtype=x.dtype, device=device)
+        else:
+            rhos_c = torch.linalg.solve(R, b).to(device).to(x.dtype)
+
+        if self.predict_x0:
+            x_t_ = sigma_t / sigma_s0 * x - alpha_t * h_phi_1 * m0
+            if D1s is not None:
+                corr_res = torch.einsum("k,bkc...->bc...", rhos_c[:-1], D1s)
+            else:
+                corr_res = 0
+            D1_t = model_t - m0
+            x_t = x_t_ - alpha_t * B_h * (corr_res + rhos_c[-1] * D1_t)
+        else:
+            x_t_ = alpha_t / alpha_s0 * x - sigma_t * h_phi_1 * m0
+            if D1s is not None:
+                corr_res = torch.einsum("k,bkc...->bc...", rhos_c[:-1], D1s)
+            else:
+                corr_res = 0
+            D1_t = model_t - m0
+            x_t = x_t_ - sigma_t * B_h * (corr_res + rhos_c[-1] * D1_t)
+        x_t = x_t.to(x.dtype)
+        return x_t
+
+    def index_for_timestep(self, timestep, schedule_timesteps=None):
+        if schedule_timesteps is None:
+            schedule_timesteps = self.timesteps
+
+        indices = (schedule_timesteps == timestep).nonzero()
+
+        # The sigma index that is taken for the **very** first `step`
+        # is always the second index (or the last index if there is only 1)
+        # This way we can ensure we don't accidentally skip a sigma in
+        # case we start in the middle of the denoising schedule (e.g. for image-to-image)
+        pos = 1 if len(indices) > 1 else 0
+
+        return indices[pos].item()
+
+    # Copied from diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler._init_step_index
+    def _init_step_index(self, timestep):
+        """
+        Initialize the step_index counter for the scheduler.
+        """
+
+        if self.begin_index is None:
+            if isinstance(timestep, torch.Tensor):
+                timestep = timestep.to(self.timesteps.device)
+            self._step_index = self.index_for_timestep(timestep)
+        else:
+            self._step_index = self._begin_index
+
+    def step(self,
+             model_output: torch.Tensor,
+             timestep: Union[int, torch.Tensor],
+             sample: torch.Tensor,
+             return_dict: bool = True,
+             generator=None) -> Union[SchedulerOutput, Tuple]:
+        """
+        Predict the sample from the previous timestep by reversing the SDE. This function propagates the sample with
+        the multistep UniPC.
+
+        Args:
+            model_output (`torch.Tensor`):
+                The direct output from learned diffusion model.
+            timestep (`int`):
+                The current discrete timestep in the diffusion chain.
+            sample (`torch.Tensor`):
+                A current instance of a sample created by the diffusion process.
+            return_dict (`bool`):
+                Whether or not to return a [`~schedulers.scheduling_utils.SchedulerOutput`] or `tuple`.
+
+        Returns:
+            [`~schedulers.scheduling_utils.SchedulerOutput`] or `tuple`:
+                If return_dict is `True`, [`~schedulers.scheduling_utils.SchedulerOutput`] is returned, otherwise a
+                tuple is returned where the first element is the sample tensor.
+
+        """
+        if self.num_inference_steps is None:
+            raise ValueError(
+                "Number of inference steps is 'None', you need to run 'set_timesteps' after creating the scheduler"
+            )
+
+        if self.step_index is None:
+            self._init_step_index(timestep)
+
+        use_corrector = (
+            self.step_index > 0 and
+            self.step_index - 1 not in self.disable_corrector and
+            self.last_sample is not None  # pyright: ignore
+        )
+
+        model_output_convert = self.convert_model_output(
+            model_output, sample=sample)
+        if use_corrector:
+            sample = self.multistep_uni_c_bh_update(
+                this_model_output=model_output_convert,
+                last_sample=self.last_sample,
+                this_sample=sample,
+                order=self.this_order,
+            )
+
+        for i in range(self.config.solver_order - 1):
+            self.model_outputs[i] = self.model_outputs[i + 1]
+            self.timestep_list[i] = self.timestep_list[i + 1]
+
+        self.model_outputs[-1] = model_output_convert
+        self.timestep_list[-1] = timestep  # pyright: ignore
+
+        if self.config.lower_order_final:
+            this_order = min(self.config.solver_order,
+                             len(self.timesteps) -
+                             self.step_index)  # pyright: ignore
+        else:
+            this_order = self.config.solver_order
+
+        self.this_order = min(this_order,
+                              self.lower_order_nums + 1)  # warmup for multistep
+        assert self.this_order > 0
+
+        self.last_sample = sample
+        prev_sample = self.multistep_uni_p_bh_update(
+            model_output=model_output,  # pass the original non-converted model output, in case solver-p is used
+            sample=sample,
+            order=self.this_order,
+        )
+
+        if self.lower_order_nums < self.config.solver_order:
+            self.lower_order_nums += 1
+
+        # upon completion increase step index by one
+        self._step_index += 1  # pyright: ignore
+
+        if not return_dict:
+            return (prev_sample,)
+
+        return SchedulerOutput(prev_sample=prev_sample)
+
+    def scale_model_input(self, sample: torch.Tensor, *args,
+                          **kwargs) -> torch.Tensor:
+        """
+        Ensures interchangeability with schedulers that need to scale the denoising model input depending on the
+        current timestep.
+
+        Args:
+            sample (`torch.Tensor`):
+                The input sample.
+
+        Returns:
+            `torch.Tensor`:
+                A scaled input sample.
+        """
+        return sample
+
+    # Copied from diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler.add_noise
+    def add_noise(
+        self,
+        original_samples: torch.Tensor,
+        noise: torch.Tensor,
+        timesteps: torch.IntTensor,
+    ) -> torch.Tensor:
+        # Make sure sigmas and timesteps have the same device and dtype as original_samples
+        sigmas = self.sigmas.to(
+            device=original_samples.device, dtype=original_samples.dtype)
+        if original_samples.device.type == "mps" and torch.is_floating_point(
+                timesteps):
+            # mps does not support float64
+            schedule_timesteps = self.timesteps.to(
+                original_samples.device, dtype=torch.float32)
+            timesteps = timesteps.to(
+                original_samples.device, dtype=torch.float32)
+        else:
+            schedule_timesteps = self.timesteps.to(original_samples.device)
+            timesteps = timesteps.to(original_samples.device)
+
+        # begin_index is None when the scheduler is used for training or pipeline does not implement set_begin_index
+        if self.begin_index is None:
+            step_indices = [
+                self.index_for_timestep(t, schedule_timesteps)
+                for t in timesteps
+            ]
+        elif self.step_index is not None:
+            # add_noise is called after first denoising step (for inpainting)
+            step_indices = [self.step_index] * timesteps.shape[0]
+        else:
+            # add noise is called before first denoising step to create initial latent(img2img)
+            step_indices = [self.begin_index] * timesteps.shape[0]
+
+        sigma = sigmas[step_indices].flatten()
+        while len(sigma.shape) < len(original_samples.shape):
+            sigma = sigma.unsqueeze(-1)
+
+        alpha_t, sigma_t = self._sigma_to_alpha_sigma_t(sigma)
+        noisy_samples = alpha_t * original_samples + sigma_t * noise
+        return noisy_samples
+
+    def __len__(self):
+        return self.config.num_train_timesteps

lingbot_video/sglang_moe_shim.py

@@ -0,0 +1,297 @@
+from __future__ import annotations
+
+import importlib.util
+import os
+import sys
+from contextlib import contextmanager
+from dataclasses import dataclass
+from types import ModuleType, SimpleNamespace
+from typing import NamedTuple, Optional
+
+import torch
+import torch.nn.functional as F
+
+
+@dataclass
+class LightSglangMoeRunnerConfig:
+    num_experts: Optional[int] = None
+    num_local_experts: Optional[int] = None
+    hidden_size: Optional[int] = None
+    intermediate_size_per_partition: Optional[int] = None
+    layer_id: Optional[int] = None
+    top_k: Optional[int] = None
+    num_fused_shared_experts: Optional[int] = None
+    params_dtype: Optional[torch.dtype] = None
+    routing_method_type: Optional[object] = None
+    activation: str = "silu"
+    is_gated: bool = True
+    apply_router_weight_on_input: bool = False
+    inplace: bool = False
+    no_combine: bool = False
+    routed_scaling_factor: Optional[float] = None
+    gemm1_alpha: Optional[float] = None
+    gemm1_clamp_limit: Optional[float] = None
+    swiglu_limit: Optional[float] = None
+    gate_up_interleaved: bool = False
+
+
+class LightSglangStandardTopKOutput(NamedTuple):
+    topk_weights: torch.Tensor
+    topk_ids: torch.Tensor
+    router_logits: torch.Tensor
+
+
+SGLANG_MOE_SERVER_ARGS = SimpleNamespace(
+    enable_deterministic_inference=False,
+    enable_fused_moe_sum_all_reduce=False,
+)
+FP8_E4M3_MAX = 448.0
+_SERVER_ARGS_READY = False
+
+
+def sglang_env_flag(name: str, default: bool = False) -> bool:
+    raw = os.environ.get(name)
+    if raw is None:
+        return default
+    return raw.lower() in {"1", "true", "yes", "on"}
+
+
+class SglangEnvFlag:
+    def __init__(self, name: str, default: bool = False):
+        self.name = name
+        self.default = default
+
+    def get(self) -> bool:
+        return sglang_env_flag(self.name, self.default)
+
+
+def ensure_module(name: str, package_path: str | None = None) -> ModuleType:
+    module = sys.modules.get(name)
+    if module is None:
+        module = ModuleType(name)
+        sys.modules[name] = module
+    if package_path is not None:
+        module.__path__ = [package_path]
+    return module
+
+
+def fp8_scale_from_amax(amax: torch.Tensor) -> torch.Tensor:
+    return torch.clamp(amax.float() / FP8_E4M3_MAX, min=1e-12)
+
+
+def quantize_to_fp8_e4m3fn(input: torch.Tensor, scale: torch.Tensor) -> torch.Tensor:
+    return torch.clamp(input.float() / scale, -FP8_E4M3_MAX, FP8_E4M3_MAX).to(
+        torch.float8_e4m3fn
+    )
+
+
+def raise_unsupported_sglang_quantization(*args, **kwargs):
+    raise RuntimeError("This LingBotVideo SGLang MoE shim only enables FP8 W8A8 quantization")
+
+
+def sglang_scaled_fp8_quant(
+    input: torch.Tensor,
+    scale: Optional[torch.Tensor] = None,
+    num_token_padding: Optional[int] = None,
+    use_per_token_if_dynamic: bool = False,
+):
+    if input.ndim != 2:
+        raise ValueError(f"Expected 2D input tensor, got {input.ndim}D")
+    rows = input.shape[0]
+    output_rows = max(int(num_token_padding or 0), rows)
+    output = torch.empty((output_rows, input.shape[1]), device=input.device, dtype=torch.float8_e4m3fn)
+    input_contiguous = input.contiguous()
+    if scale is None:
+        if use_per_token_if_dynamic:
+            scale = fp8_scale_from_amax(input_contiguous.float().abs().amax(dim=1, keepdim=True))
+            if output_rows > rows:
+                padded_scale = torch.ones((output_rows, 1), device=input.device, dtype=torch.float32)
+                padded_scale[:rows] = scale
+                scale = padded_scale
+        else:
+            scale = fp8_scale_from_amax(input_contiguous.float().abs().amax()).reshape(1)
+    quant_scale = scale[:rows] if scale.ndim == 2 else scale
+    output[:rows] = quantize_to_fp8_e4m3fn(input_contiguous, quant_scale)
+    if output_rows > rows:
+        output[rows:] = torch.zeros(
+            (output_rows - rows, input.shape[1]), device=input.device, dtype=torch.float8_e4m3fn
+        )
+    return output, scale.contiguous()
+
+
+def sglang_per_token_group_quant_fp8(
+    x: torch.Tensor,
+    group_size: int,
+    eps: float = 1e-10,
+    column_major_scales: bool = False,
+    scale_tma_aligned: bool = False,
+    scale_ue8m0: bool = False,
+    fuse_silu_and_mul: bool = False,
+    masked_m: Optional[torch.Tensor] = None,
+    enable_v2: Optional[bool] = None,
+):
+    if fuse_silu_and_mul:
+        half = x.shape[-1] // 2
+        x = (F.silu(x[..., :half]) * x[..., half:]).contiguous()
+    if x.shape[-1] % group_size != 0:
+        raise ValueError("The last dimension must be divisible by group_size")
+    x_view = x.contiguous().view(*x.shape[:-1], x.shape[-1] // group_size, group_size)
+    scales = fp8_scale_from_amax(x_view.float().abs().amax(dim=-1)).clamp_min(eps)
+    x_q = quantize_to_fp8_e4m3fn(x_view, scales.unsqueeze(-1)).view(x.shape)
+    if column_major_scales:
+        scales = scales.transpose(0, 1).contiguous()
+    return x_q.contiguous(), scales.contiguous()
+
+
+def sglang_silu_and_mul(x: torch.Tensor, out: torch.Tensor | None = None, *args, **kwargs):
+    half = x.shape[-1] // 2
+    result = F.silu(x[..., :half]) * x[..., half:]
+    if out is not None:
+        out.copy_(result)
+        return out
+    return result
+
+
+def sglang_gelu_and_mul(x: torch.Tensor, out: torch.Tensor | None = None, *args, **kwargs):
+    half = x.shape[-1] // 2
+    result = F.gelu(x[..., :half]) * x[..., half:]
+    if out is not None:
+        out.copy_(result)
+        return out
+    return result
+
+
+@contextmanager
+def null_sglang_config_override(config):
+    yield
+
+
+def install_sglang_moe_import_shims() -> None:
+    try:
+        spec = importlib.util.find_spec("sglang")
+    except Exception:
+        spec = None
+    if spec is None or not spec.submodule_search_locations:
+        return
+
+    package_dir = next(iter(spec.submodule_search_locations))
+    srt_dir = os.path.join(package_dir, "srt")
+    layers_dir = os.path.join(srt_dir, "layers")
+    moe_dir = os.path.join(layers_dir, "moe")
+    moe_runner_dir = os.path.join(moe_dir, "moe_runner")
+    triton_utils_dir = os.path.join(moe_runner_dir, "triton_utils")
+    jit_kernel_dir = os.path.join(package_dir, "jit_kernel")
+
+    ensure_module("sglang", package_dir)
+    ensure_module("sglang.srt", srt_dir)
+    ensure_module("sglang.srt.layers", layers_dir)
+    ensure_module("sglang.srt.layers.moe", moe_dir)
+    moe_runner_module = ensure_module("sglang.srt.layers.moe.moe_runner", moe_runner_dir)
+    moe_runner_module.MoeRunnerConfig = LightSglangMoeRunnerConfig
+    triton_utils_module = ensure_module(
+        "sglang.srt.layers.moe.moe_runner.triton_utils", triton_utils_dir
+    )
+    triton_utils_module.get_config = lambda: None
+    triton_utils_module.override_config = null_sglang_config_override
+    ensure_module("sglang.jit_kernel", jit_kernel_dir)
+    activation_module = ensure_module("sglang.jit_kernel.activation")
+    activation_module.silu_and_mul = sglang_silu_and_mul
+    activation_module.gelu_and_mul = sglang_gelu_and_mul
+
+    server_args_module = ensure_module("sglang.srt.server_args")
+    server_args_module.get_global_server_args = lambda: SGLANG_MOE_SERVER_ARGS
+    server_args_module.set_global_server_args_for_scheduler = lambda args: None
+
+    batch_module = ensure_module("sglang.srt.batch_invariant_ops")
+    batch_module.is_batch_invariant_mode_enabled = lambda: False
+
+    environ_module = ensure_module("sglang.srt.environ")
+    environ_module.envs = SimpleNamespace(
+        SGLANG_OPT_SWIGLU_CLAMP_FUSION=SglangEnvFlag("SGLANG_OPT_SWIGLU_CLAMP_FUSION"),
+        SGLANG_EXPERIMENTAL_LORA_OPTI=SglangEnvFlag("SGLANG_EXPERIMENTAL_LORA_OPTI"),
+    )
+
+    utils_module = ensure_module("sglang.srt.utils")
+    utils_module.cpu_has_amx_support = lambda: False
+    utils_module.get_bool_env_var = sglang_env_flag
+    utils_module.is_cpu = lambda: not torch.cuda.is_available()
+    utils_module.is_cuda = torch.cuda.is_available
+    utils_module.is_hip = lambda: False
+    utils_module.is_musa = lambda: False
+    utils_module.is_xpu = lambda: False
+    utils_module.use_intel_xpu_backend = lambda: False
+    utils_module.get_device_name = lambda: torch.cuda.get_device_name() if torch.cuda.is_available() else "cpu"
+    utils_module.is_sm90_supported = (
+        lambda: torch.cuda.is_available() and torch.cuda.get_device_capability()[0] >= 9
+    )
+
+    custom_op_module = ensure_module("sglang.srt.utils.custom_op")
+    custom_op_module.register_custom_op = lambda *args, **kwargs: (lambda fn: fn)
+
+    moe_utils_module = ensure_module("sglang.srt.layers.moe.utils")
+    moe_utils_module.get_moe_padding_size = lambda is_aiter_moe: 0
+
+    fp8_module = ensure_module("sglang.srt.layers.quantization.fp8_kernel")
+    fp8_module.per_token_group_quant_fp8 = sglang_per_token_group_quant_fp8
+    fp8_module.scaled_fp8_quant = sglang_scaled_fp8_quant
+    fp8_module.sglang_per_token_group_quant_fp8 = sglang_per_token_group_quant_fp8
+
+    int8_module = ensure_module("sglang.srt.layers.quantization.int8_kernel")
+    int8_module.per_token_group_quant_int8 = raise_unsupported_sglang_quantization
+    int8_module.per_token_quant_int8 = raise_unsupported_sglang_quantization
+    int8_module.sglang_per_token_group_quant_int8 = raise_unsupported_sglang_quantization
+
+
+try:
+    from sglang.srt.layers.moe.moe_runner.triton_utils.fused_moe import (
+        fused_experts as sglang_fused_experts,
+    )
+except Exception as direct_exc:  # pragma: no cover - optional deployment dependency
+    install_sglang_moe_import_shims()
+    try:
+        from sglang.srt.layers.moe.moe_runner.triton_utils.fused_moe import (
+            fused_experts as sglang_fused_experts,
+        )
+    except Exception as shim_exc:  # pragma: no cover - optional deployment dependency
+        sglang_fused_experts = None
+        SGLANG_MOE_IMPORT_ERROR = shim_exc
+    else:
+        SGLANG_MOE_IMPORT_ERROR = None
+else:
+    SGLANG_MOE_IMPORT_ERROR = None
+
+try:
+    from sglang.srt.server_args import (
+        get_global_server_args as sglang_get_global_server_args,
+        set_global_server_args_for_scheduler as sglang_set_global_server_args_for_scheduler,
+    )
+except Exception:  # pragma: no cover - optional deployment dependency
+    sglang_get_global_server_args = None
+    sglang_set_global_server_args_for_scheduler = None
+
+
+def _ensure_moe_server_args_attrs(args) -> None:
+    if not hasattr(args, "enable_deterministic_inference"):
+        args.enable_deterministic_inference = False
+    if not hasattr(args, "enable_fused_moe_sum_all_reduce"):
+        args.enable_fused_moe_sum_all_reduce = False
+
+
+def ensure_sglang_moe_ready() -> None:
+    global _SERVER_ARGS_READY
+    if sglang_fused_experts is None:
+        raise RuntimeError(
+            "LINGBOT_MOE_EXPERT_BACKEND=sglang_triton requires SGLang MoE runtime"
+        ) from SGLANG_MOE_IMPORT_ERROR
+    if _SERVER_ARGS_READY:
+        return
+    _ensure_moe_server_args_attrs(SGLANG_MOE_SERVER_ARGS)
+    if sglang_get_global_server_args is not None:
+        try:
+            server_args = sglang_get_global_server_args()
+        except Exception:
+            server_args = SGLANG_MOE_SERVER_ARGS
+            if sglang_set_global_server_args_for_scheduler is not None:
+                sglang_set_global_server_args_for_scheduler(server_args)
+        _ensure_moe_server_args_attrs(server_args)
+    _SERVER_ARGS_READY = True

lingbot_video/transformer_lingbot_video.py

@@ -0,0 +1,1312 @@
+import math
+import os
+from typing import Optional, Tuple
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.distributed as dist
+
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+from diffusers.models.attention_dispatch import dispatch_attention_fn
+from diffusers.models._modeling_parallel import ContextParallelInput, ContextParallelOutput
+from diffusers.models.embeddings import TimestepEmbedding, Timesteps
+from diffusers.models.modeling_outputs import Transformer2DModelOutput
+from diffusers.models.modeling_utils import ModelMixin
+
+try:
+    from flash_attn_interface import flash_attn_varlen_func as flash_attn_varlen_func_v3
+except Exception:  # pragma: no cover - optional CUDA kernel.
+    flash_attn_varlen_func_v3 = None
+
+try:
+    from .moe_pack_kernels import reorder_tokens_triton_pack
+    from .moe_restore_kernels import restore_tokens_triton
+    from .sglang_moe_shim import (
+        LightSglangMoeRunnerConfig,
+        LightSglangStandardTopKOutput,
+        ensure_sglang_moe_ready,
+        fp8_scale_from_amax,
+        quantize_to_fp8_e4m3fn,
+        sglang_fused_experts,
+    )
+except ImportError:  # pragma: no cover - allows direct file loading in diagnostics.
+    from moe_pack_kernels import reorder_tokens_triton_pack
+    from moe_restore_kernels import restore_tokens_triton
+    from sglang_moe_shim import (
+        LightSglangMoeRunnerConfig,
+        LightSglangStandardTopKOutput,
+        ensure_sglang_moe_ready,
+        fp8_scale_from_amax,
+        quantize_to_fp8_e4m3fn,
+        sglang_fused_experts,
+    )
+
+
+LINGBOT_VIDEO_FP32_MODULES = (
+    "time_embedder",
+    "time_modulation",
+    "scale_shift_table",
+    "norm",
+    "norm1",
+    "norm2",
+    "norm_q",
+    "norm_k",
+    "norm_post_attn",
+    "norm_post_ffn",
+    "norm_out",
+    "norm_out_modulation",
+    "router",
+)
+
+
+def should_keep_in_fp32(name: str) -> bool:
+    return any(module_name in name.split(".") for module_name in LINGBOT_VIDEO_FP32_MODULES)
+
+
+def _moe_expert_backend() -> str:
+    return os.environ.get("LINGBOT_MOE_EXPERT_BACKEND", "grouped_mm").lower().strip()
+
+
+def _moe_pad_backend() -> str:
+    return os.environ.get("LINGBOT_MOE_PAD_BACKEND", "loop").lower().strip()
+
+
+def _moe_reorder_backend() -> str:
+    return os.environ.get("LINGBOT_MOE_REORDER_BACKEND", "sort").lower().strip()
+
+
+def _moe_restore_backend() -> str:
+    return os.environ.get("LINGBOT_MOE_RESTORE_BACKEND", "scatter").lower().strip()
+
+
+def _all_to_all_split_cat(
+    local_input: torch.Tensor,
+    scatter_dim: int,
+    gather_dim: int,
+    group: dist.ProcessGroup,
+) -> torch.Tensor:
+    world_size = dist.get_world_size(group)
+    input_list = [
+        tensor.contiguous()
+        for tensor in torch.tensor_split(local_input, world_size, scatter_dim)
+    ]
+    output_list = [torch.empty_like(input_list[0]) for _ in range(world_size)]
+    dist.all_to_all(output_list, input_list, group=group)
+    return torch.cat(output_list, dim=gather_dim).contiguous()
+
+
+class LingBotVideoRMSNorm(nn.Module):
+    """RMSNorm with fp32 accumulation."""
+
+    def __init__(self, dim: int, eps: float = 1e-6):
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(dim))
+        self.variance_epsilon = eps
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        input_dtype = hidden_states.dtype
+        hidden_states = hidden_states.to(torch.float32)
+        variance = hidden_states.pow(2).mean(-1, keepdim=True)
+        hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
+        return (self.weight * hidden_states).to(input_dtype)
+
+
+def apply_rotary_emb(x: torch.Tensor, freqs_cis: torch.Tensor) -> torch.Tensor:
+    """Apply complex RoPE to `(B, S, H, D)` attention tensors."""
+    with torch.amp.autocast("cuda", enabled=False):
+        x_c = torch.view_as_complex(x.float().reshape(*x.shape[:-1], -1, 2))
+        out = torch.view_as_real(x_c * freqs_cis.unsqueeze(2)).flatten(3)
+        return out.type_as(x)
+
+
+class LingBotVideoRotaryEmbedding(nn.Module):
+    """Complex64 RoPE table indexed by position ids."""
+
+    def __init__(self, axes_dims: Tuple[int, ...], axes_lens: Tuple[int, ...], theta: float):
+        super().__init__()
+        self.axes_dims = tuple(axes_dims)
+        self.axes_lens = list(axes_lens)
+        self.theta = theta
+        self.freqs_cis = None
+
+    @staticmethod
+    def precompute_freqs_cis(dim: Tuple[int, ...], end: Tuple[int, ...], theta: float):
+        freqs_cis = []
+        for d, e in zip(dim, end):
+            freqs = 1.0 / (
+                theta ** (torch.arange(0, d, 2, dtype=torch.float64, device="cpu") / d)
+            )
+            timestep = torch.arange(e, device=freqs.device, dtype=torch.float64)
+            freqs = torch.outer(timestep, freqs).float()
+            freqs_cis.append(torch.polar(torch.ones_like(freqs), freqs).to(torch.complex64))
+        return freqs_cis
+
+    def forward(self, position_ids: torch.Tensor) -> torch.Tensor:
+        # position_ids: (S, 3) int → (S, head_dim/2) complex64
+        device = position_ids.device
+        max_vals = position_ids.max(dim=0).values.tolist()
+        needs_rebuild = self.freqs_cis is None or any(m >= l for m, l in zip(max_vals, self.axes_lens))
+        if needs_rebuild:
+            for i in range(len(self.axes_lens)):
+                if max_vals[i] >= self.axes_lens[i]:
+                    self.axes_lens[i] = int(max_vals[i] * 1.5) + 1
+            self.freqs_cis = self.precompute_freqs_cis(
+                self.axes_dims, tuple(self.axes_lens), theta=self.theta
+            )
+            self.freqs_cis = [freqs_cis.to(device) for freqs_cis in self.freqs_cis]
+        elif self.freqs_cis[0].device != device:
+            self.freqs_cis = [freqs_cis.to(device) for freqs_cis in self.freqs_cis]
+
+        return torch.cat([self.freqs_cis[i][position_ids[:, i]] for i in range(len(self.axes_dims))], dim=-1)
+
+
+def make_joint_position_ids(
+    text_len: int, grid_t: int, grid_h: int, grid_w: int, device: torch.device
+) -> torch.Tensor:
+    """3D positions in [video; text] order. Text t-axis is 1..text_len; video t-axis starts at text_len+1.
+
+    Matches patchify_and_embed: cap start (1,0,0); vision start (cap_len+1,0,0);
+    freqs ordered with x first and cap second (same order as cat_interleave).
+    """
+    tt = torch.arange(grid_t, device=device, dtype=torch.int32) + (text_len + 1)
+    hh = torch.arange(grid_h, device=device, dtype=torch.int32)
+    ww = torch.arange(grid_w, device=device, dtype=torch.int32)
+    grid = torch.stack(torch.meshgrid(tt, hh, ww, indexing="ij"), dim=-1).flatten(0, 2)
+    text_t = torch.arange(text_len, device=device, dtype=torch.int32) + 1
+    text_pos = torch.stack(
+        [text_t, torch.zeros_like(text_t), torch.zeros_like(text_t)], dim=-1
+    )
+    return torch.cat([grid, text_pos], dim=0)  # (Nx + L, 3)
+
+
+def _cat_interleave(
+    a: torch.Tensor,
+    len_a: list[int],
+    b: torch.Tensor,
+    len_b: list[int],
+) -> torch.Tensor:
+    a_split = torch.split(a, len_a, dim=1)
+    b_split = torch.split(b, len_b, dim=1)
+    blocks: list[torch.Tensor] = []
+    for x_part, text_part in zip(a_split, b_split):
+        blocks.extend([x_part, text_part])
+    return torch.cat(blocks, dim=1)
+
+
+class LingBotVideoTextEmbedder(nn.Module):
+    """Matches CondProjection: RMSNorm(text_dim, eps=1e-6 fixed) -> Linear-SiLU-Linear."""
+
+    def __init__(self, text_dim: int, hidden_size: int):
+        super().__init__()
+        self.norm = LingBotVideoRMSNorm(text_dim, eps=1e-6)
+        self.linear_1 = nn.Linear(text_dim, hidden_size, bias=True)
+        self.linear_2 = nn.Linear(hidden_size, hidden_size, bias=True)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = self.norm(x)
+        return self.linear_2(F.silu(self.linear_1(x)))
+
+
+class LingBotVideoAttention(nn.Module):
+    def __init__(self, hidden_size, num_heads, norm_eps, qkv_bias, out_bias):
+        super().__init__()
+        self.num_heads = num_heads
+        self.head_dim = hidden_size // num_heads
+        self.to_q = nn.Linear(hidden_size, hidden_size, bias=qkv_bias)
+        self.to_k = nn.Linear(hidden_size, hidden_size, bias=qkv_bias)
+        self.to_v = nn.Linear(hidden_size, hidden_size, bias=qkv_bias)
+        self.norm_q = LingBotVideoRMSNorm(self.head_dim, norm_eps)
+        self.norm_k = LingBotVideoRMSNorm(self.head_dim, norm_eps)
+        self.to_out = nn.Linear(hidden_size, hidden_size, bias=out_bias)
+
+    def forward(
+        self,
+        x,
+        rotary_emb,
+        attention_mask=None,
+        packed_indices: Optional[dict[str, torch.Tensor]] = None,
+        parallel_config=None,
+    ):
+        B, S, _ = x.shape
+        if os.environ.get("LINGBOT_FUSED_QKV_LINEAR") == "1":
+            weight = torch.cat(
+                (self.to_q.weight, self.to_k.weight, self.to_v.weight),
+                dim=0,
+            )
+            bias = None
+            if self.to_q.bias is not None:
+                bias = torch.cat(
+                    (self.to_q.bias, self.to_k.bias, self.to_v.bias),
+                    dim=0,
+                )
+            qkv = F.linear(x, weight, bias)
+            q, k, v = qkv.view(B, S, 3, self.num_heads, self.head_dim).unbind(2)
+        else:
+            q = self.to_q(x).unflatten(2, (self.num_heads, self.head_dim))
+            k = self.to_k(x).unflatten(2, (self.num_heads, self.head_dim))
+            v = self.to_v(x).unflatten(2, (self.num_heads, self.head_dim))
+        q = apply_rotary_emb(self.norm_q(q), rotary_emb)
+        k = apply_rotary_emb(self.norm_k(k), rotary_emb)
+        # dispatch_attention_fn expects (B, S, H, D) in and out (same as the diffusers Wan processor)
+        if packed_indices is None:
+            out = dispatch_attention_fn(
+                q,
+                k,
+                v,
+                attn_mask=attention_mask,
+                parallel_config=parallel_config,
+            )
+        else:
+            if flash_attn_varlen_func_v3 is None:
+                raise RuntimeError("flash_attn_interface.flash_attn_varlen_func is required.")
+            if parallel_config is None:
+                result = flash_attn_varlen_func_v3(
+                    q=q.reshape(-1, self.num_heads, self.head_dim),
+                    k=k.reshape(-1, self.num_heads, self.head_dim),
+                    v=v.reshape(-1, self.num_heads, self.head_dim),
+                    cu_seqlens_q=packed_indices["cu_seqlens_kv"],
+                    cu_seqlens_k=packed_indices["cu_seqlens_kv"],
+                    max_seqlen_q=packed_indices["max_seqlen_in_batch_kv"],
+                    max_seqlen_k=packed_indices["max_seqlen_in_batch_kv"],
+                    causal=False,
+                )
+                out = result[0] if isinstance(result, tuple) else result
+                out = out.reshape(B, S, self.num_heads, self.head_dim)
+            else:
+                group = parallel_config.context_parallel_config._ulysses_mesh.get_group()
+                world_size = dist.get_world_size(group)
+                local_heads = self.num_heads // world_size
+                q_global = _all_to_all_split_cat(
+                    q.reshape(B, S, self.num_heads * self.head_dim),
+                    scatter_dim=2,
+                    gather_dim=1,
+                    group=group,
+                ).view(B, S * world_size, local_heads, self.head_dim)
+                k_global = _all_to_all_split_cat(
+                    k.reshape(B, S, self.num_heads * self.head_dim),
+                    scatter_dim=2,
+                    gather_dim=1,
+                    group=group,
+                ).view(B, S * world_size, local_heads, self.head_dim)
+                v_global = _all_to_all_split_cat(
+                    v.reshape(B, S, self.num_heads * self.head_dim),
+                    scatter_dim=2,
+                    gather_dim=1,
+                    group=group,
+                ).view(B, S * world_size, local_heads, self.head_dim)
+                q_flat = q_global.reshape(-1, local_heads, self.head_dim)
+                k_flat = k_global.reshape(-1, local_heads, self.head_dim)
+                v_flat = v_global.reshape(-1, local_heads, self.head_dim)
+                result = flash_attn_varlen_func_v3(
+                    q=q_flat,
+                    k=k_flat,
+                    v=v_flat,
+                    cu_seqlens_q=packed_indices["cu_seqlens_kv"],
+                    cu_seqlens_k=packed_indices["cu_seqlens_kv"],
+                    max_seqlen_q=packed_indices["max_seqlen_in_batch_kv"],
+                    max_seqlen_k=packed_indices["max_seqlen_in_batch_kv"],
+                    causal=False,
+                )
+                out_global = result[0] if isinstance(result, tuple) else result
+                out_global = out_global.reshape(B, S * world_size, local_heads * self.head_dim)
+                out = _all_to_all_split_cat(
+                    out_global,
+                    scatter_dim=1,
+                    gather_dim=2,
+                    group=group,
+                ).view(B, S, self.num_heads, self.head_dim)
+        return self.to_out(out.flatten(2, 3).type_as(x))
+
+
+class LingBotVideoMLP(nn.Module):
+    def __init__(self, hidden_size, intermediate_size):
+        super().__init__()
+        self.gate_proj = nn.Linear(hidden_size, intermediate_size, bias=False)
+        self.up_proj = nn.Linear(hidden_size, intermediate_size, bias=False)
+        self.down_proj = nn.Linear(intermediate_size, hidden_size, bias=False)
+
+    def forward(self, x):
+        return self.down_proj(F.silu(self.gate_proj(x)) * self.up_proj(x))
+
+
+class LingBotVideoRouter(nn.Module):
+    """Matches the TokenChoiceTopKRouter inference path (no capacity/jitter/load stats).
+
+    The asymmetry must be preserved: selection uses the bias-added score, while gating
+    weights gather the bias-free score.
+    """
+
+    def __init__(self, hidden_size, num_experts, top_k, score_func, norm_topk_prob,
+                 n_group, topk_group, route_scale):
+        super().__init__()
+        self.num_experts = num_experts
+        self.top_k = top_k
+        self.score_func = score_func
+        self.norm_topk_prob = norm_topk_prob
+        self.n_group = n_group
+        self.topk_group = topk_group
+        self.route_scale = route_scale
+        self.weight = nn.Parameter(torch.empty(num_experts, hidden_size))
+        self.register_buffer("e_score_correction_bias", torch.zeros(num_experts), persistent=True)
+
+    def _group_limited_topk(self, scores_for_choice):
+        seq_len = scores_for_choice.shape[0]
+        experts_per_group = self.num_experts // self.n_group
+        grouped = scores_for_choice.view(seq_len, self.n_group, experts_per_group)
+        group_scores = grouped.topk(2, dim=-1)[0].sum(dim=-1)
+        group_idx = torch.topk(group_scores, k=self.topk_group, dim=-1, sorted=False)[1]
+        group_mask = torch.zeros_like(group_scores)
+        group_mask.scatter_(1, group_idx, 1)
+        score_mask = (
+            group_mask.unsqueeze(-1)
+            .expand(seq_len, self.n_group, experts_per_group)
+            .reshape(seq_len, -1)
+        )
+        masked = scores_for_choice.masked_fill(~score_mask.bool(), float("-inf"))
+        return torch.topk(masked, k=self.top_k, dim=-1, sorted=False)[1]
+
+    def forward(self, tokens: torch.Tensor):
+        with torch.amp.autocast(tokens.device.type, enabled=False):
+            logits = F.linear(tokens.float(), self.weight.float())
+        if self.score_func == "softmax":
+            scores = F.softmax(logits, dim=-1)
+        else:
+            scores = logits.sigmoid()
+        scores_for_choice = scores + self.e_score_correction_bias.unsqueeze(0)
+        if self.n_group is not None and self.n_group > 1:
+            top_indices = self._group_limited_topk(scores_for_choice)
+        else:
+            top_indices = torch.topk(scores_for_choice, k=self.top_k, dim=-1, sorted=False)[1]
+        top_scores = scores.gather(1, top_indices)
+        if self.top_k > 1 and self.norm_topk_prob:
+            top_scores = top_scores / (top_scores.sum(dim=-1, keepdim=True) + 1e-20)
+        top_scores = top_scores * self.route_scale
+        return top_indices, top_scores.to(tokens.dtype), logits, scores, scores_for_choice
+
+
+class LingBotVideoGroupedExperts(nn.Module):
+    """Weight layout matches GroupedExperts: w1 [E,I,H], w2 [E,H,I], w3 [E,I,H]. Eager per-expert compute."""
+
+    def __init__(self, num_experts, hidden_size, intermediate_size):
+        super().__init__()
+        self.num_experts = num_experts
+        self.w1 = nn.Parameter(torch.empty(num_experts, intermediate_size, hidden_size))
+        self.w2 = nn.Parameter(torch.empty(num_experts, hidden_size, intermediate_size))
+        self.w3 = nn.Parameter(torch.empty(num_experts, intermediate_size, hidden_size))
+
+
+def _round_up_to_multiple(value: int, multiple: int) -> int:
+    return ((value + multiple - 1) // multiple) * multiple
+
+
+class LingBotVideoSparseMoeBlock(nn.Module):
+    def __init__(self, hidden_size, intermediate_size, num_experts, top_k,
+                 moe_intermediate_size, score_func, norm_topk_prob, n_group, topk_group,
+                 routed_scaling_factor, n_shared_experts):
+        super().__init__()
+        self.hidden_size = hidden_size
+        self.num_experts = num_experts
+        self.router = LingBotVideoRouter(
+            hidden_size, num_experts, top_k, score_func, norm_topk_prob,
+            n_group, topk_group, routed_scaling_factor,
+        )
+        self.experts = LingBotVideoGroupedExperts(num_experts, hidden_size, moe_intermediate_size)
+        self._sglang_w13_cache: Optional[torch.Tensor] = None
+        self._sglang_w13_cache_key = None
+        self._sglang_fp8_cache = None
+        self._sglang_fp8_cache_key = None
+        self.shared_experts = None
+        if n_shared_experts is not None and n_shared_experts > 0:
+            self.shared_experts = LingBotVideoMLP(
+                hidden_size, moe_intermediate_size * n_shared_experts
+            )
+
+    @staticmethod
+    def _reorder_tokens(tokens: torch.Tensor, top_scores: torch.Tensor, top_indices: torch.Tensor, num_experts: int):
+        backend = _moe_reorder_backend()
+        if backend in {"triton_pack", "pack", "triton"}:
+            return reorder_tokens_triton_pack(tokens, top_scores, top_indices, num_experts)
+        if backend not in {"sort", "argsort", "default"}:
+            raise ValueError(
+                f"Unsupported LINGBOT_MOE_REORDER_BACKEND={backend!r}; "
+                "expected sort or triton_pack"
+            )
+        num_tokens = tokens.shape[0]
+        top_k = top_indices.shape[1]
+        flat_scores = top_scores.reshape(-1)
+        flat_indices = top_indices.reshape(-1)
+        active_positions = torch.where(flat_scores != 0)[0]
+        active_experts = flat_indices[active_positions]
+
+        counts = torch.zeros(num_experts, device=tokens.device, dtype=torch.int64)
+        counts.scatter_add_(0, active_experts, torch.ones_like(active_experts, dtype=torch.int64))
+
+        sort_order = torch.argsort(active_experts, stable=True)
+        sorted_positions = active_positions[sort_order]
+        sorted_scores = flat_scores[sorted_positions]
+        original_token_idx = sorted_positions // top_k
+        permuted_tokens = tokens[original_token_idx]
+        return permuted_tokens, counts, sorted_positions, sorted_scores, num_tokens, top_k
+
+    @staticmethod
+    def _pad_grouped_tokens_loop(tokens: torch.Tensor, counts: torch.Tensor, align: int = 8):
+        num_tokens = tokens.shape[0]
+        num_experts = int(counts.shape[0])
+        max_len = _round_up_to_multiple(num_tokens + num_experts * align, align)
+        counts_i64 = counts.to(torch.int64)
+        total_per_expert = torch.clamp_min(counts_i64, align)
+        aligned_counts = (
+            (total_per_expert + align - 1) // align * align
+        ).to(torch.int32)
+        write_offsets = torch.cumsum(aligned_counts, dim=0) - aligned_counts
+        start_indices = torch.cumsum(counts_i64, dim=0) - counts_i64
+
+        fill_value = num_tokens
+        permuted_indices = torch.full(
+            (max_len,), fill_value, dtype=torch.int64, device=tokens.device
+        )
+        for expert_idx in range(num_experts):
+            length = int(counts_i64[expert_idx].item())
+            if length == 0:
+                continue
+            write_start = int(write_offsets[expert_idx].item())
+            start = int(start_indices[expert_idx].item())
+            permuted_indices[write_start:write_start + length] = torch.arange(
+                start, start + length, device=tokens.device, dtype=torch.int64
+            )
+
+        tokens_with_pad = torch.vstack((tokens, tokens.new_zeros((tokens.shape[-1],))))
+        input_shape = tokens_with_pad.shape
+        return input_shape, tokens_with_pad[permuted_indices], permuted_indices, aligned_counts
+
+    @staticmethod
+    def _pad_grouped_tokens_vectorized(tokens: torch.Tensor, counts: torch.Tensor, align: int = 8):
+        num_tokens = tokens.shape[0]
+        num_experts = int(counts.shape[0])
+        max_len = _round_up_to_multiple(num_tokens + num_experts * align, align)
+        counts_i64 = counts.to(torch.int64)
+        total_per_expert = torch.clamp_min(counts_i64, align)
+        aligned_counts_i64 = (total_per_expert + align - 1) // align * align
+        write_offsets = torch.cumsum(aligned_counts_i64, dim=0) - aligned_counts_i64
+        end_offsets = torch.cumsum(aligned_counts_i64, dim=0)
+        start_indices = torch.cumsum(counts_i64, dim=0) - counts_i64
+
+        slots = torch.arange(max_len, dtype=torch.int64, device=tokens.device)
+        expert_idx = torch.bucketize(slots, end_offsets, right=True)
+        valid_expert = expert_idx < num_experts
+        safe_expert_idx = expert_idx.clamp(max=num_experts - 1)
+        local_idx = slots - write_offsets[safe_expert_idx]
+        source_idx = start_indices[safe_expert_idx] + local_idx
+        valid = valid_expert & (local_idx < counts_i64[safe_expert_idx])
+        fill = torch.full_like(source_idx, num_tokens)
+        permuted_indices = torch.where(valid, source_idx, fill)
+
+        tokens_with_pad = torch.vstack((tokens, tokens.new_zeros((tokens.shape[-1],))))
+        input_shape = tokens_with_pad.shape
+        return (
+            input_shape,
+            tokens_with_pad[permuted_indices],
+            permuted_indices,
+            aligned_counts_i64.to(torch.int32),
+        )
+
+    @staticmethod
+    def _pad_grouped_tokens(tokens: torch.Tensor, counts: torch.Tensor, align: int = 8):
+        backend = _moe_pad_backend()
+        if backend in {"loop", "default"}:
+            return LingBotVideoSparseMoeBlock._pad_grouped_tokens_loop(tokens, counts, align)
+        if backend in {"vectorized", "torch"}:
+            return LingBotVideoSparseMoeBlock._pad_grouped_tokens_vectorized(tokens, counts, align)
+        raise ValueError(
+            f"Unsupported LINGBOT_MOE_PAD_BACKEND={backend!r}; expected loop or vectorized"
+        )
+
+    @staticmethod
+    def _unpad_grouped_tokens(output: torch.Tensor, input_shape: torch.Size, permuted_indices: torch.Tensor):
+        unpermuted = output.new_empty(input_shape)
+        unpermuted[permuted_indices, :] = output
+        return unpermuted[:-1]
+
+    def _run_grouped_experts(self, tokens: torch.Tensor, counts: torch.Tensor) -> torch.Tensor:
+        if not hasattr(torch, "_grouped_mm"):
+            return self._run_experts_for_loop(tokens, counts)
+        input_shape, padded_tokens, permuted_indices, aligned_counts = self._pad_grouped_tokens(tokens, counts)
+        offsets = torch.cumsum(aligned_counts, dim=0, dtype=torch.int32)
+        h = F.silu(
+            torch._grouped_mm(
+                padded_tokens.bfloat16(),
+                self.experts.w1.bfloat16().transpose(-2, -1),
+                offs=offsets,
+            )
+        )
+        h = h * torch._grouped_mm(
+            padded_tokens.bfloat16(),
+            self.experts.w3.bfloat16().transpose(-2, -1),
+            offs=offsets,
+        )
+        out = torch._grouped_mm(
+            h,
+            self.experts.w2.bfloat16().transpose(-2, -1),
+            offs=offsets,
+        ).type_as(padded_tokens)
+        return self._unpad_grouped_tokens(out, input_shape, permuted_indices)
+
+    def _run_experts_for_loop(self, tokens: torch.Tensor, counts: torch.Tensor) -> torch.Tensor:
+        count_list = counts.tolist()
+        splits = torch.split(tokens, count_list, dim=0)
+        outputs = []
+        for expert_idx, expert_tokens in enumerate(splits):
+            if expert_tokens.numel() == 0:
+                continue
+            h = F.silu(expert_tokens @ self.experts.w1[expert_idx].transpose(-2, -1))
+            h = h * (expert_tokens @ self.experts.w3[expert_idx].transpose(-2, -1))
+            h = h @ self.experts.w2[expert_idx].transpose(-2, -1)
+            outputs.append(h)
+        if not outputs:
+            return tokens.new_zeros(tokens.shape)
+        return torch.cat(outputs, dim=0)
+
+    def _get_sglang_w13(self) -> torch.Tensor:
+        key = (
+            self.experts.w1.data_ptr(),
+            self.experts.w3.data_ptr(),
+            self.experts.w1.device,
+            self.experts.w3.device,
+            self.experts.w1.dtype,
+            self.experts.w3.dtype,
+        )
+        if self._sglang_w13_cache is None or self._sglang_w13_cache_key != key:
+            self._sglang_w13_cache = torch.cat(
+                (self.experts.w1.bfloat16(), self.experts.w3.bfloat16()), dim=1
+            ).contiguous()
+            self._sglang_w13_cache_key = key
+        return self._sglang_w13_cache
+
+    @staticmethod
+    def _quantize_fp8_weight_per_expert(weight: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
+        weight_float = weight.float()
+        scale = fp8_scale_from_amax(weight_float.abs().amax(dim=(1, 2)))
+        quantized = quantize_to_fp8_e4m3fn(weight_float, scale[:, None, None]).contiguous()
+        return quantized, scale.contiguous()
+
+    def _get_sglang_fp8_weights(self) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+        key = (
+            self.experts.w1.data_ptr(),
+            self.experts.w2.data_ptr(),
+            self.experts.w3.data_ptr(),
+            self.experts.w1.device,
+            self.experts.w2.device,
+            self.experts.w3.device,
+            self.experts.w1.dtype,
+            self.experts.w2.dtype,
+            self.experts.w3.dtype,
+        )
+        if self._sglang_fp8_cache is None or self._sglang_fp8_cache_key != key:
+            w13 = torch.cat((self.experts.w1.float(), self.experts.w3.float()), dim=1).contiguous()
+            w13_fp8, w13_scale = self._quantize_fp8_weight_per_expert(w13)
+            w2_fp8, w2_scale = self._quantize_fp8_weight_per_expert(self.experts.w2)
+            self._sglang_fp8_cache = (w13_fp8, w2_fp8, w13_scale, w2_scale)
+            self._sglang_fp8_cache_key = key
+        return self._sglang_fp8_cache
+
+    def _run_sglang_triton_experts(
+        self,
+        tokens: torch.Tensor,
+        top_scores: torch.Tensor,
+        top_indices: torch.Tensor,
+    ) -> torch.Tensor:
+        ensure_sglang_moe_ready()
+        topk_output = LightSglangStandardTopKOutput(
+            top_scores.float(),
+            top_indices.to(torch.int32),
+            torch.empty(0, device=tokens.device),
+        )
+        runner_config = LightSglangMoeRunnerConfig(
+            num_experts=self.num_experts,
+            num_local_experts=self.num_experts,
+            activation="silu",
+            is_gated=True,
+            inplace=False,
+        )
+        return sglang_fused_experts(
+            tokens.contiguous().bfloat16(),
+            self._get_sglang_w13(),
+            self.experts.w2.bfloat16().contiguous(),
+            topk_output,
+            runner_config,
+        ).type_as(tokens)
+
+    def _run_sglang_triton_fp8_experts(
+        self,
+        tokens: torch.Tensor,
+        top_scores: torch.Tensor,
+        top_indices: torch.Tensor,
+    ) -> torch.Tensor:
+        ensure_sglang_moe_ready()
+        topk_output = LightSglangStandardTopKOutput(
+            top_scores.float(),
+            top_indices.to(torch.int32),
+            torch.empty(0, device=tokens.device),
+        )
+        runner_config = LightSglangMoeRunnerConfig(
+            num_experts=self.num_experts,
+            num_local_experts=self.num_experts,
+            activation="silu",
+            is_gated=True,
+            inplace=False,
+        )
+        w13_fp8, w2_fp8, w13_scale, w2_scale = self._get_sglang_fp8_weights()
+        return sglang_fused_experts(
+            tokens.contiguous().bfloat16(),
+            w13_fp8,
+            w2_fp8,
+            topk_output,
+            runner_config,
+            use_fp8_w8a8=True,
+            w1_scale=w13_scale,
+            w2_scale=w2_scale,
+        ).type_as(tokens)
+
+    def _run_selected_experts(
+        self,
+        tokens: torch.Tensor,
+        top_scores: torch.Tensor,
+        top_indices: torch.Tensor,
+    ) -> torch.Tensor:
+        backend = _moe_expert_backend()
+        if backend in {"grouped_mm", "torch_grouped_mm", "default"}:
+            (
+                permuted_tokens,
+                counts,
+                sorted_positions,
+                sorted_scores,
+                num_tokens,
+                top_k,
+            ) = self._reorder_tokens(tokens, top_scores, top_indices, self.router.num_experts)
+            expert_output = self._run_grouped_experts(permuted_tokens, counts)
+            return self._restore_tokens(
+                expert_output,
+                sorted_positions,
+                sorted_scores,
+                num_tokens,
+                top_k,
+            )
+        if backend in {"sglang_triton", "triton", "sglang"}:
+            return self._run_sglang_triton_experts(tokens, top_scores, top_indices)
+        if backend in {"sglang_triton_fp8", "triton_fp8", "sglang_fp8"}:
+            return self._run_sglang_triton_fp8_experts(tokens, top_scores, top_indices)
+        raise ValueError(
+            f"Unsupported LINGBOT_MOE_EXPERT_BACKEND={backend!r}; "
+            "expected grouped_mm, sglang_triton, or sglang_triton_fp8"
+        )
+
+    @staticmethod
+    def _restore_tokens(
+        expert_output: torch.Tensor,
+        sorted_positions: torch.Tensor,
+        sorted_scores: torch.Tensor,
+        num_tokens: int,
+        top_k: int,
+    ) -> torch.Tensor:
+        backend = _moe_restore_backend()
+        if backend in {"triton", "triton_fused", "fused"}:
+            return LingBotVideoSparseMoeBlock._restore_tokens_triton(
+                expert_output,
+                sorted_positions,
+                sorted_scores,
+                num_tokens,
+                top_k,
+            )
+        if backend in {"index_add", "index_add_", "scatter_add"}:
+            return LingBotVideoSparseMoeBlock._restore_tokens_index_add(
+                expert_output,
+                sorted_positions,
+                sorted_scores,
+                num_tokens,
+                top_k,
+            )
+        if backend in {"weighted_scatter", "weighted", "fast_scatter"}:
+            return LingBotVideoSparseMoeBlock._restore_tokens_weighted_scatter(
+                expert_output,
+                sorted_positions,
+                sorted_scores,
+                num_tokens,
+                top_k,
+            )
+        if backend in {"chunked", "chunked_scatter", "scatter_chunked"}:
+            return LingBotVideoSparseMoeBlock._restore_tokens_chunked_scatter(
+                expert_output,
+                sorted_positions,
+                sorted_scores,
+                num_tokens,
+                top_k,
+            )
+        if backend not in {"scatter", "default"}:
+            raise ValueError(
+                f"Unsupported LINGBOT_MOE_RESTORE_BACKEND={backend!r}; "
+                "expected scatter, chunked_scatter, weighted_scatter, index_add, or triton"
+            )
+        dim = expert_output.shape[-1]
+        unsorted = torch.zeros(
+            (num_tokens * top_k, dim),
+            dtype=expert_output.dtype,
+            device=expert_output.device,
+        )
+        unsorted[sorted_positions] = expert_output
+        unsorted = unsorted.reshape(num_tokens, top_k, dim)
+
+        scores_unsorted = torch.zeros(
+            num_tokens * top_k,
+            dtype=sorted_scores.dtype,
+            device=sorted_scores.device,
+        )
+        scores_unsorted[sorted_positions] = sorted_scores
+        scores_unsorted = scores_unsorted.reshape(num_tokens, top_k, 1)
+        return (unsorted.float() * scores_unsorted).sum(dim=1).to(expert_output.dtype)
+
+    @staticmethod
+    def _restore_tokens_chunked_scatter(
+        expert_output: torch.Tensor,
+        sorted_positions: torch.Tensor,
+        sorted_scores: torch.Tensor,
+        num_tokens: int,
+        top_k: int,
+    ) -> torch.Tensor:
+        dim = expert_output.shape[-1]
+        chunk_size = int(os.environ.get("LINGBOT_MOE_RESTORE_CHUNK_SIZE", "128"))
+        if chunk_size <= 0:
+            raise ValueError("LINGBOT_MOE_RESTORE_CHUNK_SIZE must be positive")
+
+        scores_unsorted = torch.zeros(
+            num_tokens * top_k,
+            dtype=sorted_scores.dtype,
+            device=sorted_scores.device,
+        )
+        scores_unsorted[sorted_positions] = sorted_scores
+        scores_unsorted = scores_unsorted.reshape(num_tokens, top_k, 1)
+        output = expert_output.new_empty((num_tokens, dim))
+        for start in range(0, dim, chunk_size):
+            end = min(start + chunk_size, dim)
+            unsorted = torch.zeros(
+                (num_tokens * top_k, end - start),
+                dtype=expert_output.dtype,
+                device=expert_output.device,
+            )
+            unsorted[sorted_positions] = expert_output[:, start:end]
+            unsorted = unsorted.reshape(num_tokens, top_k, end - start)
+            output[:, start:end] = (unsorted.float() * scores_unsorted).sum(dim=1).to(
+                expert_output.dtype
+            )
+        return output
+
+    @staticmethod
+    def _restore_tokens_triton(
+        expert_output: torch.Tensor,
+        sorted_positions: torch.Tensor,
+        sorted_scores: torch.Tensor,
+        num_tokens: int,
+        top_k: int,
+    ) -> torch.Tensor:
+        return restore_tokens_triton(
+            expert_output,
+            sorted_positions,
+            sorted_scores,
+            num_tokens,
+            top_k,
+        )
+
+    @staticmethod
+    def _restore_tokens_weighted_scatter(
+        expert_output: torch.Tensor,
+        sorted_positions: torch.Tensor,
+        sorted_scores: torch.Tensor,
+        num_tokens: int,
+        top_k: int,
+    ) -> torch.Tensor:
+        weighted = (expert_output * sorted_scores[:, None].to(expert_output.dtype)).to(expert_output.dtype)
+        unsorted = torch.zeros(
+            (num_tokens * top_k, expert_output.shape[-1]),
+            dtype=expert_output.dtype,
+            device=expert_output.device,
+        )
+        unsorted[sorted_positions] = weighted
+        return unsorted.reshape(num_tokens, top_k, expert_output.shape[-1]).sum(dim=1)
+
+    @staticmethod
+    def _restore_tokens_index_add(
+        expert_output: torch.Tensor,
+        sorted_positions: torch.Tensor,
+        sorted_scores: torch.Tensor,
+        num_tokens: int,
+        top_k: int,
+    ) -> torch.Tensor:
+        token_indices = torch.div(sorted_positions, top_k, rounding_mode="floor")
+        weighted = expert_output.float() * sorted_scores[:, None].float()
+        out = torch.zeros(
+            (num_tokens, expert_output.shape[-1]),
+            dtype=torch.float32,
+            device=expert_output.device,
+        )
+        out.index_add_(0, token_indices, weighted)
+        return out.to(expert_output.dtype)
+
+    def forward(self, hidden_states: torch.Tensor, padding_mask: Optional[torch.Tensor] = None):
+        # hidden_states: (B, S, H); padding_mask: (B*S,) with 1=valid (only needed when B>1)
+        B = hidden_states.shape[0]
+        tokens = hidden_states.view(-1, self.hidden_size)
+        top_indices, top_scores, logits, scores, scores_for_choice = self.router(tokens)
+        del logits, scores, scores_for_choice
+        if padding_mask is not None:
+            pm = padding_mask.unsqueeze(-1).to(top_scores.dtype)
+            top_scores = top_scores * pm
+            top_scores = top_scores / (top_scores.sum(dim=-1, keepdim=True) + 1e-9)
+            top_scores = top_scores * self.router.route_scale
+
+        out = self._run_selected_experts(tokens, top_scores, top_indices)
+
+        out = out.view(B, -1, self.hidden_size)
+        if self.shared_experts is not None:
+            shared_output = self.shared_experts(hidden_states)
+            out = out + shared_output
+        return out
+
+
+class LingBotVideoBlock(nn.Module):
+    def __init__(
+        self,
+        hidden_size,
+        num_attention_heads,
+        intermediate_size,
+        norm_eps,
+        qkv_bias,
+        out_bias,
+        num_experts,
+        num_experts_per_tok,
+        moe_intermediate_size,
+        decoder_sparse_step,
+        mlp_only_layers,
+        n_shared_experts,
+        score_func,
+        norm_topk_prob,
+        n_group,
+        topk_group,
+        routed_scaling_factor,
+        layer_idx: int,
+    ):
+        super().__init__()
+        self.layer_idx = layer_idx
+        h = hidden_size
+        self.scale_shift_table = nn.Parameter(torch.zeros(1, 6 * h))
+        self.norm1 = LingBotVideoRMSNorm(h, norm_eps)
+        self.attn = LingBotVideoAttention(
+            h, num_attention_heads, norm_eps, qkv_bias, out_bias
+        )
+        self.norm_post_attn = LingBotVideoRMSNorm(h, norm_eps)
+        self.norm2 = LingBotVideoRMSNorm(h, norm_eps)
+        # Sparsity decision matches MoEBlock: mlp_only_layers + decoder_sparse_step + num_experts
+        if layer_idx not in mlp_only_layers and (
+            num_experts > 0 and (layer_idx + 1) % decoder_sparse_step == 0
+        ):
+            self.ffn = LingBotVideoSparseMoeBlock(
+                h, intermediate_size, num_experts, num_experts_per_tok,
+                moe_intermediate_size, score_func, norm_topk_prob,
+                n_group, topk_group, routed_scaling_factor,
+                n_shared_experts,
+            )
+        else:
+            self.ffn = LingBotVideoMLP(h, intermediate_size)
+        self.norm_post_ffn = LingBotVideoRMSNorm(h, norm_eps)
+
+    def forward(
+        self,
+        x,
+        temb6,
+        rotary_emb,
+        attention_mask=None,
+        moe_padding_mask=None,
+        packed_indices: Optional[dict[str, torch.Tensor]] = None,
+        parallel_config=None,
+    ):
+        expected_tokens = x.shape[0] * x.shape[1]
+        if temb6.ndim != 2 or temb6.shape[0] != expected_tokens:
+            raise ValueError(
+                "LingBotVideoBlock expects token-level temb6 with shape "
+                f"(B*S, 6D); got {tuple(temb6.shape)} for hidden states {tuple(x.shape)}."
+            )
+        # AdaLN mod: dense and MoE both keep scale_shift_table fp32 (master
+        # moe/models.py:80 dropped the accidental `.to(dtype=c.dtype)` cast).
+        mod = temb6.view(x.shape[0], x.shape[1], -1) + self.scale_shift_table.unsqueeze(0)
+        shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = mod.chunk(6, dim=-1)
+        gate_msa, gate_mlp = gate_msa.tanh(), gate_mlp.tanh()
+        scale_msa, scale_mlp = 1.0 + scale_msa, 1.0 + scale_mlp
+
+        # AdaLN modulation / norms run in fp32 (sensitive path); cast to the bulk
+        # compute dtype only at the bf16 Linear boundary. This replaces the old
+        # ambient autocast, which rounded Linear inputs to bf16 at the same point.
+        bulk_dtype = self.attn.to_q.weight.dtype
+        attn_in = (self.norm1(x) * scale_msa + shift_msa).to(bulk_dtype)
+        attn_out = self.attn(
+            attn_in,
+            rotary_emb,
+            attention_mask,
+            packed_indices=packed_indices,
+            parallel_config=parallel_config,
+        )
+        x = x + (gate_msa * self.norm_post_attn(attn_out)).to(x.dtype)
+
+        ffn_in = (self.norm2(x) * scale_mlp + shift_mlp).to(bulk_dtype)
+        if isinstance(self.ffn, LingBotVideoSparseMoeBlock):
+            ffn_out = self.ffn(ffn_in, padding_mask=moe_padding_mask)
+        else:
+            ffn_out = self.ffn(ffn_in)
+        ffn_normed = self.norm_post_ffn(ffn_out)
+        x = x + (gate_mlp * ffn_normed).to(x.dtype)
+        return x
+
+
+class LingBotVideoTransformer3DModel(ModelMixin, ConfigMixin):
+    _supports_gradient_checkpointing = False
+    _no_split_modules = ["LingBotVideoBlock"]
+    _keep_in_fp32_modules = list(LINGBOT_VIDEO_FP32_MODULES)
+
+    def to(self, *args, **kwargs):
+        device, dtype, non_blocking, _ = torch._C._nn._parse_to(*args, **kwargs)
+        if dtype is None or dtype == torch.float32:
+            return super().to(*args, **kwargs)
+
+        dtype_is_floating = torch.is_floating_point(torch.empty((), dtype=dtype))
+        if not dtype_is_floating:
+            return super().to(*args, **kwargs)
+
+        if device is not None:
+            super().to(device=device, non_blocking=non_blocking)
+
+        for name, param in self.named_parameters():
+            if not torch.is_floating_point(param):
+                continue
+            target_dtype = torch.float32 if should_keep_in_fp32(name) else dtype
+            param.data = param.data.to(dtype=target_dtype, non_blocking=non_blocking)
+            if param.grad is not None:
+                param.grad.data = param.grad.data.to(dtype=target_dtype, non_blocking=non_blocking)
+
+        for name, buffer in self.named_buffers():
+            if not torch.is_floating_point(buffer):
+                continue
+            target_dtype = torch.float32 if should_keep_in_fp32(name) else dtype
+            buffer.data = buffer.data.to(dtype=target_dtype, non_blocking=non_blocking)
+
+        return self
+
+    @register_to_config
+    def __init__(
+        self,
+        patch_size: Tuple[int, int, int] = (1, 2, 2),
+        in_channels: int = 16,
+        out_channels: int = 16,
+        hidden_size: int = 2048,
+        num_attention_heads: int = 16,
+        depth: int = 24,
+        intermediate_size: int = 6144,
+        text_dim: int = 2560,
+        freq_dim: int = 256,
+        norm_eps: float = 1e-6,
+        rope_theta: float = 256.0,
+        axes_dims: Tuple[int, int, int] = (32, 48, 48),
+        axes_lens: Tuple[int, int, int] = (8192, 1024, 1024),
+        qkv_bias: bool = False,
+        out_bias: bool = True,
+        patch_embed_bias: bool = True,
+        timestep_mlp_bias: bool = True,
+        num_experts: int = 0,
+        num_experts_per_tok: int = 8,
+        moe_intermediate_size: int = 512,
+        decoder_sparse_step: int = 1,
+        mlp_only_layers: Tuple[int, ...] = (),
+        n_shared_experts: Optional[int] = None,
+        score_func: str = "sigmoid",
+        norm_topk_prob: bool = True,
+        n_group: Optional[int] = None,
+        topk_group: Optional[int] = None,
+        routed_scaling_factor: float = 1.0,
+    ):
+        super().__init__()
+        head_dim = hidden_size // num_attention_heads
+        assert head_dim == sum(axes_dims), f"head_dim {head_dim} != sum(axes_dims) {sum(axes_dims)}"
+        mlp_only_layers = tuple(mlp_only_layers)
+
+        self.patch_embedder = nn.Linear(
+            in_channels * math.prod(patch_size), hidden_size, bias=patch_embed_bias
+        )
+        self.time_proj = Timesteps(freq_dim, flip_sin_to_cos=True, downscale_freq_shift=0)
+        self.time_embedder = TimestepEmbedding(
+            freq_dim, hidden_size, act_fn="silu", sample_proj_bias=timestep_mlp_bias
+        )
+        self.time_modulation = nn.Sequential(nn.SiLU(), nn.Linear(hidden_size, 6 * hidden_size))
+        self.text_embedder = LingBotVideoTextEmbedder(text_dim, hidden_size)
+        self.rope = LingBotVideoRotaryEmbedding(axes_dims, axes_lens, rope_theta)
+        self.blocks = nn.ModuleList(
+            [
+                LingBotVideoBlock(
+                    hidden_size=hidden_size,
+                    num_attention_heads=num_attention_heads,
+                    intermediate_size=intermediate_size,
+                    norm_eps=norm_eps,
+                    qkv_bias=qkv_bias,
+                    out_bias=out_bias,
+                    num_experts=num_experts,
+                    num_experts_per_tok=num_experts_per_tok,
+                    moe_intermediate_size=moe_intermediate_size,
+                    decoder_sparse_step=decoder_sparse_step,
+                    mlp_only_layers=mlp_only_layers,
+                    n_shared_experts=n_shared_experts,
+                    score_func=score_func,
+                    norm_topk_prob=norm_topk_prob,
+                    n_group=n_group,
+                    topk_group=topk_group,
+                    routed_scaling_factor=routed_scaling_factor,
+                    layer_idx=i,
+                )
+                for i in range(depth)
+            ]
+        )
+        self.norm_out = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=norm_eps)
+        self.norm_out_modulation = nn.Sequential(nn.SiLU(), nn.Linear(hidden_size, 2 * hidden_size))
+        self.proj_out = nn.Linear(hidden_size, math.prod(patch_size) * out_channels)
+        self.cp_joint = nn.Identity()
+        self.cp_rotary = nn.Identity()
+        self.cp_temb_input = nn.Identity()
+        self.cp_temb6 = nn.Identity()
+        self.cp_out = nn.Identity()
+        self._cp_plan = {
+            "cp_joint": {
+                "input": ContextParallelInput(split_dim=1, expected_dims=3),
+            },
+            "cp_rotary": {
+                "input": ContextParallelInput(split_dim=1, expected_dims=3),
+            },
+            "cp_temb_input": {
+                "input": ContextParallelInput(split_dim=1, expected_dims=3),
+            },
+            "cp_temb6": {
+                "input": ContextParallelInput(split_dim=1, expected_dims=3),
+            },
+            "cp_out": ContextParallelOutput(gather_dim=1, expected_dims=3),
+        }
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,             # (B, C, T, H, W)
+        timestep: torch.Tensor,                  # (B,) ∈ [0, 1000](= sigma*1000)
+        encoder_hidden_states: torch.Tensor,     # (B, L, text_dim)
+        encoder_attention_mask: Optional[torch.Tensor] = None,  # (B, L) 1=valid
+        return_dict: bool = True,
+    ):
+        B, C, T, H, W = hidden_states.shape
+        pF, pH, pW = self.config.patch_size
+        gt, gh, gw = T // pF, H // pH, W // pW
+        n_video = gt * gh * gw
+        L = encoder_hidden_states.shape[1]
+        device = hidden_states.device
+        if encoder_attention_mask is not None:
+            text_lens = encoder_attention_mask.sum(dim=-1).long()
+        else:
+            text_lens = torch.full((B,), L, dtype=torch.long, device=device)
+        text_lens_list = [int(v) for v in text_lens.detach().cpu().tolist()]
+        packed_batch = B > 1
+
+        # patchify: token order (f h w), feature order (pf ph pw c) -- matches patchify_and_embed
+        patch_tokens = hidden_states.reshape(B, C, gt, pF, gh, pH, gw, pW)
+        patch_tokens = patch_tokens.permute(0, 2, 4, 6, 3, 5, 7, 1).reshape(
+            B,
+            n_video,
+            pF * pH * pW * C,
+        )
+        if packed_batch:
+            packed_patch_tokens = patch_tokens.reshape(1, B * n_video, -1)
+            x = torch.cat(
+                [self.patch_embedder(patch_tokens[i : i + 1]) for i in range(B)],
+                dim=1,
+            )
+        else:
+            x = self.patch_embedder(patch_tokens)
+
+        if packed_batch:
+            text_parts = [
+                self.text_embedder(encoder_hidden_states[i : i + 1, : text_lens_list[i], :])
+                for i in range(B)
+            ]
+            text = torch.cat(text_parts, dim=1)
+            joint = _cat_interleave(
+                x,
+                [n_video] * B,
+                text,
+                text_lens_list,
+            )
+        else:
+            text = self.text_embedder(encoder_hidden_states)
+            joint = torch.cat([x, text], dim=1)  # [video; text]
+        joint_seq_len = joint.shape[1]
+
+        # Per-sample RoPE: video t-axis start = real text length of this sample + 1
+        rotary_parts = [
+            self.rope(make_joint_position_ids(text_lens_list[i], gt, gh, gw, device))
+            for i in range(B)
+        ]
+        if packed_batch:
+            rotary = torch.cat(rotary_parts, dim=0).unsqueeze(0)
+        else:
+            rotary = torch.stack(rotary_parts, dim=0)  # (B, S, head_dim/2) complex64
+
+        parallel_config = getattr(self, "_parallel_config", None)
+        use_packed_attention = parallel_config is not None
+
+        attention_mask = None
+        moe_padding_mask = None
+        packed_indices = None
+        has_padding = encoder_attention_mask is not None and bool((text_lens < L).any())
+        if packed_batch or use_packed_attention:
+            sample_seq_lens = [n_video + text_len for text_len in text_lens_list]
+            cu_seqlens = torch.zeros(B + 1, device=device, dtype=torch.int32)
+            cu_seqlens[1:] = torch.cumsum(
+                torch.tensor(sample_seq_lens, device=device, dtype=torch.int32),
+                dim=0,
+            )
+            packed_indices = {
+                "cu_seqlens_kv": cu_seqlens,
+                "max_seqlen_in_batch_kv": max(sample_seq_lens),
+            }
+            has_padding = False
+        if has_padding:
+            key_mask = torch.cat(
+                [torch.ones(B, n_video, dtype=torch.bool, device=device),
+                 encoder_attention_mask.bool()],
+                dim=1,
+            )
+            attention_mask = key_mask[:, None, None, :]      # (B,1,1,S) → SDPA broadcast
+            moe_padding_mask = key_mask.reshape(-1).float()  # (B*S,)
+        packed_cp = packed_indices is not None and parallel_config is not None
+        padding_size = 0
+        if packed_cp:
+            cp_config = parallel_config.context_parallel_config
+            cp_world_size = int(getattr(cp_config, "ulysses_degree", getattr(cp_config, "_world_size", 1)))
+            padding_size = (cp_world_size - (joint_seq_len % cp_world_size)) % cp_world_size
+            if padding_size:
+                joint = torch.cat(
+                    [
+                        joint,
+                        torch.zeros(
+                            joint.shape[0],
+                            padding_size,
+                            joint.shape[2],
+                            device=joint.device,
+                            dtype=joint.dtype,
+                        ),
+                    ],
+                    dim=1,
+                )
+                rotary = torch.cat(
+                    [
+                        rotary,
+                        torch.zeros(
+                            rotary.shape[0],
+                            padding_size,
+                            rotary.shape[2],
+                            device=rotary.device,
+                            dtype=rotary.dtype,
+                        ),
+                    ],
+                    dim=1,
+                )
+                if packed_indices is None:
+                    raise RuntimeError("packed_indices must be initialized for packed context parallel.")
+                packed_indices["cu_seqlens_kv"] = torch.cat(
+                    [
+                        packed_indices["cu_seqlens_kv"],
+                        packed_indices["cu_seqlens_kv"][-1:] + padding_size,
+                    ],
+                    dim=0,
+                )
+                packed_indices["max_seqlen_in_batch_kv"] = max(
+                    int(packed_indices["max_seqlen_in_batch_kv"]),
+                    int(padding_size),
+                )
+                joint_seq_len = joint.shape[1]
+
+        timestep_for_embed = timestep.float()
+        timestep_proj = self.time_proj(timestep_for_embed)
+        t_emb = self.time_embedder(timestep_proj)                            # (B, D)
+        if packed_batch:
+            temb_input = torch.cat(
+                [
+                    t_emb[i : i + 1].unsqueeze(1).expand(1, n_video + text_lens_list[i], -1)
+                    for i in range(B)
+                ],
+                dim=1,
+            )
+            if padding_size:
+                temb_input = torch.cat(
+                    [
+                        temb_input,
+                        torch.zeros(
+                            temb_input.shape[0],
+                            padding_size,
+                            temb_input.shape[2],
+                            device=temb_input.device,
+                            dtype=temb_input.dtype,
+                        ),
+                    ],
+                    dim=1,
+                )
+            temb6 = self.time_modulation(temb_input.reshape(joint_seq_len, -1))
+            temb6 = temb6.reshape(1, joint_seq_len, -1)
+        else:
+            temb_input = t_emb.unsqueeze(1).expand(B, joint_seq_len, -1)       # (B, S, D)
+            temb6 = self.time_modulation(temb_input.reshape(B * joint_seq_len, -1))
+            temb6 = temb6.reshape(B, joint_seq_len, -1)                        # (B, S, 6D)
+
+        joint = self.cp_joint(joint)
+        rotary = self.cp_rotary(rotary)
+        if packed_cp:
+            temb_input = self.cp_temb_input(temb_input)
+        temb6 = self.cp_temb6(temb6)
+        temb6 = temb6.reshape(temb6.shape[0] * temb6.shape[1], -1)
+
+        for block in self.blocks:
+            joint = block(
+                joint,
+                temb6,
+                rotary,
+                attention_mask,
+                moe_padding_mask,
+                packed_indices=packed_indices,
+                parallel_config=parallel_config,
+            )
+        if not packed_cp:
+            joint = self.cp_out(joint)
+
+        final_mod = self.norm_out_modulation(temb_input.reshape(joint.shape[0] * joint.shape[1], -1))
+        shift, scale = final_mod.reshape(joint.shape[0], joint.shape[1], -1).chunk(2, dim=-1)
+        final_hidden = self.norm_out(joint) * (1.0 + scale) + shift
+        projected = self.proj_out(final_hidden.to(self.proj_out.weight.dtype))
+        if packed_cp:
+            projected = self.cp_out(projected)
+            if padding_size:
+                projected = projected[:, :-padding_size, :]
+        if packed_batch:
+            split_lengths: list[int] = []
+            for text_len in text_lens_list:
+                split_lengths.extend([n_video, text_len])
+            parts = torch.split(projected, split_lengths, dim=1)
+            x = torch.cat(parts[::2], dim=1).reshape(B, n_video, -1)
+        else:
+            x = projected[:, :n_video]
+
+        # unpatchify (matches the rearrange in postprocess)
+        Cout = self.config.out_channels
+        x = x.reshape(B, gt, gh, gw, pF, pH, pW, Cout)
+        x = x.permute(0, 7, 1, 4, 2, 5, 3, 6).reshape(B, Cout, T, H, W)
+
+        if not return_dict:
+            return (x,)
+        return Transformer2DModelOutput(sample=x)

lingbot_video/utils.py

@@ -0,0 +1,287 @@
+from __future__ import annotations
+
+import json
+from pathlib import Path
+from typing import Any, Sequence
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+from PIL import Image
+
+LOW_NOISE_TAIL_V1_NAME = "low_noise_tail_v1"
+LOW_NOISE_TAIL_V1_DEFAULT_STEPS = 2
+
+
+def num_frames_from_duration(duration: float, fps: int) -> int:
+    frame_count = int(float(duration) * int(fps))
+    return ((frame_count - 1) // 4 + 1) * 4 + 1
+
+
+def validate_refiner_sigmas(
+    sigmas: Sequence[float] | np.ndarray,
+    t_thresh: float | None = None,
+) -> np.ndarray:
+    arr = np.asarray(list(sigmas), dtype=np.float64)
+    if arr.ndim != 1 or arr.size == 0:
+        raise ValueError("refiner sigma schedule must be a non-empty 1D list")
+    if not np.all(np.isfinite(arr)):
+        raise ValueError("refiner sigma schedule contains non-finite values")
+    if np.any(arr < 0.0) or np.any(arr > 1.0):
+        raise ValueError(f"refiner sigma schedule values must be in [0, 1], got {arr.tolist()}")
+    if arr.size > 1 and not np.all(np.diff(arr) < 0.0):
+        raise ValueError(f"refiner sigma schedule must be strictly descending, got {arr.tolist()}")
+    if t_thresh is not None and abs(float(arr[0]) - float(t_thresh)) > 1e-6:
+        raise ValueError(f"refiner sigma schedule must start at t_thresh={float(t_thresh)}, got {float(arr[0])}")
+    return arr
+
+
+def compute_refiner_sigmas(
+    *,
+    sigma_max: float,
+    sigma_min: float,
+    num_inference_steps: int,
+    shift: float,
+    t_thresh: float | None,
+    tail_steps: int = 0,
+) -> np.ndarray | None:
+    if t_thresh is None:
+        return None
+    t_value = float(t_thresh)
+    if not (0.0 < t_value <= 1.0):
+        raise ValueError(f"refiner t_thresh must lie in (0, 1], got {t_value}")
+    steps = int(num_inference_steps)
+    if steps < 1:
+        raise ValueError(f"num_inference_steps must be >= 1, got {steps}")
+    tail = int(tail_steps or 0)
+    if tail < 0:
+        raise ValueError(f"refiner_sigma_tail_steps must be >= 0, got {tail}")
+
+    base = np.linspace(float(sigma_max), float(sigma_min), steps + 1).copy()[:-1]
+    shift_value = float(shift)
+    shifted = shift_value * base / (1.0 + (shift_value - 1.0) * base)
+    eps = 1e-6
+    sigmas = shifted[shifted <= t_value + eps]
+    if sigmas.size == 0 or abs(float(sigmas[0]) - t_value) > eps:
+        sigmas = np.concatenate([[t_value], sigmas])
+    if tail > 0:
+        start = float(sigmas[-1])
+        stop = min(float(sigma_min), start)
+        extra = np.linspace(start, stop, tail + 2, dtype=np.float64)[1:-1]
+        sigmas = np.concatenate([sigmas, extra])
+    return validate_refiner_sigmas(sigmas, t_value).astype(np.float32)
+
+
+def prepare_refiner_latent(
+    x_up: torch.Tensor,
+    noise: torch.Tensor,
+    t_thresh: float | torch.Tensor,
+) -> torch.Tensor:
+    if not torch.is_tensor(t_thresh):
+        t_thresh = torch.tensor(float(t_thresh), device=x_up.device, dtype=x_up.dtype)
+    while t_thresh.ndim < x_up.ndim:
+        t_thresh = t_thresh.view(*t_thresh.shape, *([1] * (x_up.ndim - t_thresh.ndim)))
+    return (1.0 - t_thresh) * x_up + t_thresh * noise
+
+
+def _pad_prompt_embeds(
+    embeds: torch.Tensor,
+    mask: torch.Tensor,
+    target_length: int,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    if embeds.ndim != 3:
+        raise ValueError(f"prompt embeds must be rank-3, got {tuple(embeds.shape)}")
+    if mask.ndim != 2:
+        raise ValueError(f"prompt mask must be rank-2, got {tuple(mask.shape)}")
+    if embeds.shape[:2] != mask.shape:
+        raise ValueError(
+            f"prompt embeds/mask shape mismatch: {tuple(embeds.shape)} vs {tuple(mask.shape)}"
+        )
+    if embeds.shape[0] != 1:
+        raise ValueError(f"batched CFG helper expects batch=1 inputs, got {embeds.shape[0]}")
+    if embeds.shape[1] > target_length:
+        raise ValueError(f"cannot pad length {embeds.shape[1]} down to {target_length}")
+    pad_len = target_length - embeds.shape[1]
+    if pad_len == 0:
+        return embeds, mask
+    embed_pad = torch.zeros(
+        embeds.shape[0],
+        pad_len,
+        embeds.shape[2],
+        dtype=embeds.dtype,
+        device=embeds.device,
+    )
+    mask_pad = torch.zeros(mask.shape[0], pad_len, dtype=mask.dtype, device=mask.device)
+    return torch.cat([embeds, embed_pad], dim=1), torch.cat([mask, mask_pad], dim=1)
+
+
+def batch_cfg_prompt_inputs(
+    prompt_embeds: torch.Tensor,
+    prompt_mask: torch.Tensor,
+    negative_embeds: torch.Tensor,
+    negative_mask: torch.Tensor,
+    *,
+    null_cond_clone_zero: bool = False,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    if null_cond_clone_zero:
+        zero_negative = torch.zeros_like(prompt_embeds)
+        return (
+            torch.cat([prompt_embeds, zero_negative], dim=0),
+            torch.cat([prompt_mask, prompt_mask.clone()], dim=0),
+        )
+
+    target_length = max(int(prompt_embeds.shape[1]), int(negative_embeds.shape[1]))
+    prompt_padded, prompt_mask_padded = _pad_prompt_embeds(
+        prompt_embeds, prompt_mask, target_length
+    )
+    negative_padded, negative_mask_padded = _pad_prompt_embeds(
+        negative_embeds, negative_mask, target_length
+    )
+    return (
+        torch.cat([prompt_padded, negative_padded], dim=0),
+        torch.cat([prompt_mask_padded, negative_mask_padded], dim=0),
+    )
+
+
+def caption_from_sample(sample: dict[str, Any]) -> str:
+    if "caption" in sample:
+        caption = sample["caption"]
+    else:
+        runtime_keys = {
+            "duration",
+            "fps",
+            "height",
+            "width",
+            "num_frames",
+            "resolution",
+            "ratio",
+        }
+        caption = {key: value for key, value in sample.items() if key not in runtime_keys}
+    if isinstance(caption, (dict, list)):
+        return json.dumps(caption, ensure_ascii=False, separators=(",", ":"))
+    return str(caption)
+
+
+def compute_training_frame_budget(
+    num_source_frames: int,
+    source_fps: float,
+    sample_fps: int = 24,
+    vae_tc: int = 4,
+) -> tuple[int, float, int]:
+    if num_source_frames <= 0:
+        return 1, 0.0, 1
+    if source_fps > sample_fps:
+        raw_val = int(num_source_frames / source_fps * sample_fps)
+    else:
+        raw_val = int(num_source_frames)
+    sample_frame = ((raw_val - 1) // vae_tc) * vae_tc + 1
+    sample_frame = max(sample_frame, 1)
+    vae_fps = sample_frame / num_source_frames * float(source_fps)
+    t_vae = (sample_frame - 1) // vae_tc + 1
+    return int(sample_frame), float(vae_fps), int(t_vae)
+
+
+def compute_training_aligned_indices(
+    num_source_frames: int,
+    sample_frame: int,
+) -> np.ndarray:
+    if sample_frame <= 0:
+        return np.zeros(0, dtype=int)
+    if num_source_frames <= 0:
+        return np.zeros(sample_frame, dtype=int)
+    if num_source_frames >= sample_frame:
+        return np.linspace(0, num_source_frames - 1, sample_frame, dtype=int)
+    head = np.arange(num_source_frames, dtype=int)
+    pad = np.full(sample_frame - num_source_frames, num_source_frames - 1, dtype=int)
+    return np.concatenate([head, pad])
+
+
+def resize_video_tensor(video: torch.Tensor, height: int, width: int) -> torch.Tensor:
+    if video.ndim != 5:
+        raise ValueError(f"video tensor must have shape [B,C,T,H,W], got {tuple(video.shape)}")
+    bsz, channels, frames, _height, _width = video.shape
+    flat = video.permute(0, 2, 1, 3, 4).reshape(bsz * frames, channels, _height, _width)
+    resized = F.interpolate(flat, size=(height, width), mode="bicubic", align_corners=False)
+    resized = resized.clamp(0.0, 1.0)
+    return resized.reshape(bsz, frames, channels, height, width).permute(0, 2, 1, 3, 4).contiguous()
+
+
+def load_refiner_video_tensor(
+    path: str | Path,
+    height: int,
+    width: int,
+    *,
+    sample_fps: int = 24,
+    vae_tc: int = 4,
+    max_frames: int | None = None,
+) -> tuple[torch.Tensor, dict[str, Any]]:
+    from decord import VideoReader, cpu
+
+    vr = VideoReader(str(path), ctx=cpu(0))
+    total = len(vr)
+    if total <= 0:
+        raise ValueError(f"Video has no frames: {path}")
+    src_fps = float(vr.get_avg_fps())
+    sample_frame, vae_fps, t_vae = compute_training_frame_budget(
+        total,
+        src_fps,
+        sample_fps=sample_fps,
+        vae_tc=vae_tc,
+    )
+    sample_frame_uncapped = int(sample_frame)
+    truncated = False
+    if max_frames is not None and sample_frame > int(max_frames):
+        sample_frame = int(max_frames)
+        truncated = True
+        vae_fps = float(sample_frame) / max(total, 1) * src_fps
+        t_vae = (sample_frame - 1) // vae_tc + 1
+    indices = compute_training_aligned_indices(total, sample_frame)
+    frames = torch.from_numpy(vr.get_batch(indices).asnumpy()).permute(0, 3, 1, 2).float()
+    frames = frames / 255.0
+    video = frames.permute(1, 0, 2, 3).unsqueeze(0).contiguous()
+    video = resize_video_tensor(video, height=height, width=width)
+    meta = {
+        "src_fps": float(src_fps),
+        "sample_frame": int(sample_frame),
+        "sample_frame_uncapped": int(sample_frame_uncapped),
+        "max_frames": None if max_frames is None else int(max_frames),
+        "truncated_by_max_frames": bool(truncated),
+        "vae_fps": float(vae_fps),
+        "t_vae": int(t_vae),
+        "num_source_frames": int(total),
+        "align_to_training": True,
+    }
+    return video, meta
+
+
+def load_first_frame_condition_tensor(
+    path: str | Path,
+    target_height: int,
+    target_width: int,
+    geometry_height: int,
+    geometry_width: int,
+) -> torch.Tensor:
+    """Load the clean first frame center-cropped to the lowres video's aspect.
+
+    The refiner latent grid is encoded from the lowres video, so the injected
+    frame-0 condition must share that video's geometry, not the raw image's.
+    """
+    image = Image.open(path).convert("RGB")
+    image_width, image_height = image.size
+    geometry_aspect = float(geometry_width) / float(geometry_height)
+    image_aspect = float(image_width) / float(image_height)
+    if image_aspect > geometry_aspect:
+        crop_height = image_height
+        crop_width = max(1, int(round(crop_height * geometry_aspect)))
+        left = int(round((image_width - crop_width) / 2.0))
+        top = 0
+    else:
+        crop_width = image_width
+        crop_height = max(1, int(round(crop_width / geometry_aspect)))
+        left = 0
+        top = int(round((image_height - crop_height) / 2.0))
+    crop = image.crop((left, top, left + crop_width, top + crop_height))
+    crop = crop.resize((target_width, target_height), resample=Image.BICUBIC)
+    arr = np.asarray(crop, dtype=np.float32) / 255.0
+    frame = torch.from_numpy(arr).permute(2, 0, 1).unsqueeze(0).contiguous()
+    return frame.permute(1, 0, 2, 3).unsqueeze(0)

requirements.txt

@@ -0,0 +1,11 @@
+diffusers==0.39.0
+transformers==5.8.1
+accelerate>=1.10
+safetensors>=0.4.5
+torchvision
+imageio>=2.35.0
+imageio-ffmpeg>=0.5.1
+numpy>=1.26
+pillow>=10.4.0
+scipy>=1.11
+json_repair>=0.30

rewriter_prompts.py

@@ -0,0 +1,293 @@
+VIDEO_STEP1_EXPAND = (
+    'You write ONE short, natural, standalone video caption in English from a brief user video\n'
+    'prompt — as if briefly recounting what happens in a real clip of that idea.\n'
+    '\n'
+    "INPUT: a short user prompt (and, when provided, the video's first frame as a visual anchor).\n"
+    '\n'
+    'HARD LENGTH LIMIT: the output MUST be UNDER 1000 characters. Aim for roughly 400-800\n'
+    'characters. Concise is a top priority.\n'
+    '\n'
+    'STYLE — a flowing story, NOT a checklist:\n'
+    '- Connected, natural English prose. NO headings, NO bullets, NO field labels.\n'
+    '- The MAIN THREAD is what HAPPENS — lead with the action. Do NOT march subject-by-subject\n'
+    '  listing attributes; weave the few details you keep into the action naturally.\n'
+    '- One cohesive paragraph; never pad to add length.\n'
+    '\n'
+    'WHAT TO COVER (only these):\n'
+    '- The scene in a brief phrase.\n'
+    '- Each main subject by name/what-it-is plus ONE most defining visual trait — not a full\n'
+    '  appearance list.\n'
+    '- What happens over the clip, in correct chronological order as the backbone, with an EXPLICIT\n'
+    '  timestamp in seconds on each action, distributed within the given video duration — e.g.\n'
+    '  "at 0.0s", "from 1.2s to 2.7s", "around the 3.4s mark", "finally from 4.0s to 5.0s". Span the\n'
+    '  actions from 0s up to (but never beyond) the stated duration; the last one ends at or before it.\n'
+    '  Attach a timestamp ONLY to a real action or change; mention a static/unchanging element once\n'
+    '  with NO timestamp — never put a whole-clip span on something that merely exists.\n'
+    '- A one-word shot type only if it matters; otherwise skip the camera.\n'
+    "- Named entities only if clearly implied — don't invent identities.\n"
+    '\n'
+    'FAITHFUL EXPANSION:\n'
+    '- Stay consistent with the prompt; elaborate plausibly but never contradict it. If a first\n'
+    '  frame is given, ground the scene and subject in it.\n'
+    '- Keep actions in their natural, real-world direction — never reverse them. Commit to one\n'
+    "  coherent interpretation; don't hedge with alternatives.\n"
+    '- Keep subject identities and counts consistent.\n'
+    '\n'
+    'DOMAIN NOTES (apply whichever fits):\n'
+    '- Multiple events / sequence: lay out consecutive actions in STRICT chronological order as\n'
+    '  distinct, separated steps with approximate timing — never blur them into one vague action.\n'
+    '- Robot manipulation (VLA): subjects are robotic arm(s), gripper(s) and workspace objects,\n'
+    '  NOT people — no clothing/skin/gender/expression. When two arms are present, ALWAYS keep the\n'
+    '  LEFT and RIGHT arm distinct and state which arm/gripper does each motion; never swap, merge,\n'
+    '  or leave it ambiguous. Manipulation is STRICTLY NOT reversible.\n'
+    "- First-person (EGO): a head-mounted first-person view; camera motion IS the wearer's head\n"
+    '  movement. The agent is the camera-wearer, a PERSON shown through their own hands/arms and\n'
+    '  viewpoint; never a robot or external third-person subject. Stay strictly first-person.\n'
+    '\n'
+    'Output ONLY the caption text — no preamble, no headings, no explanation.\n'
+    '\n'
+    "## EXAMPLES — one representative example per domain. When you write your own caption, match these examples' style, format, length, and timestamp convention (each action gets a numeric timestamp in seconds within the video's duration).\n"
+    '\n'
+    '### Example 1 — general\n'
+    'USER PROMPT:\n'
+    '低角度广角镜头，镜头静止。中央大型条纹热气球缓慢上升并飘移（气球呈红橙黄绿蓝条纹，底部深色），背景中多个热气球在天空漂移。地面草地上散布着正在充气的热气球，一辆红色皮卡停放在旁。随后一辆蓝色拖车进入画面并横穿前景（覆盖蓝色防水布）。饱和色彩，硬光，日光，电影质感。\n'
+    '\n'
+    'DETAILED CAPTION:\n'
+    'At a vibrant hot air balloon festival on a sunny day, a wide shot shows a grassy field under a blue sky. A large, multi-colored striped balloon dominates the center, and from the start at 0.0 seconds until 5.3 seconds, it slowly rises and drifts slightly upward and to the right. Simultaneously during the 0.0s to 5.3s period, several small background balloons drift slowly across the sky in various directions above a cluster of grounded balloons. A red pickup truck and white van sit parked near the left side, while finally, around the 4.3s mark until 5.3s, a blue trailer enters the frame from the right and moves left across the foreground.\n'
+    '\n'
+    '### Example 2 — multi-event\n'
+    'USER PROMPT:\n'
+    'Subtitle: Outdoor High-Intensity Workout — All live-action, hyper-realistic, fixed shot with slight handheld unsteadiness.\n'
+    '\n'
+    'Style: Hard daylight, high-angle wide shot, center composition, sharp shadows, athletic tension.\n'
+    '\n'
+    '1. Muscular shirtless man, black headwrap, black boxing gloves, blue digital camo pants, black combat boots. Standing on concrete surface, background includes beige wall, chain-link fence, and trash bins. Stands in a boxing stance with hands up.\n'
+    '\n'
+    '2. Drops into a deep squat, then jumps upward explosively.\n'
+    '\n'
+    '3. Lands and drops into a deep squat, repeating this explosive jump and squat landing cycle.\n'
+    '\n'
+    '4. After the final landing and squat, begins to stand up.\n'
+    '\n'
+    'Overall movements coherent and natural, high contrast lighting, focused and athletic atmosphere.\n'
+    '\n'
+    'DETAILED CAPTION:\n'
+    'In an outdoor urban setting against a beige wall, a shirtless man wearing boxing gloves performs a workout on concrete. From 0.0s to 0.6s he stands in a boxing stance with hands up. Then from 0.6s to 1.2s he drops into a deep squat, followed by the period from 1.2s to 1.8s where he jumps upward explosively. From 1.8s to 2.4s he lands and drops into a deep squat, then from 2.4s to 3.0s jumps upward explosively again. Around the 3.0s to 3.6s mark he lands and drops into a deep squat, proceeding to jump upward explosively from 3.6s to 4.2s. From 4.2s to 4.8s he lands and drops into a deep squat once more, and finally from 4.8s to 5.0s he begins to stand up.\n'
+    '\n'
+    '### Example 3 — VLA\n'
+    'USER PROMPT:\n'
+    '环境：自动化超市补货工作站，俯视视角，明亮均匀人工光，清晰功能化氛围。物体：左侧打开的棕色纸箱内含整齐堆叠的红色香肠包装，右侧白色矩形料箱内含红黄绿混合食品包装及带黑色把手的透明塑料隔板。机器人：双臂系统。左臂黑色机身银色底座黑色夹爪（全程静止悬停），右臂白色机身黑色夹爪腕部蓝色指示灯（活动主体）。动作：右臂向下向左移动伸入白色料箱，抓取隔板黑色把手，向右滑动隔板，释放把手，向上向右收回复位。左臂悬停于纸箱上方保持静止。相机：固定高角度俯视视角，全程静止镜头，宽画幅，超广角镜头，柔和人工光，极致细节。\n'
+    '\n'
+    'DETAILED CAPTION:\n'
+    'A top-down wide shot shows an automated workspace with a cardboard box of sausage packages on the left and a white bin on the right. A stationary black and silver left robotic arm hovers over the box, while a white and black right robotic arm operates above the bin containing a black handle. From 0.0s to 2.0s, the right arm moves downwards and to the left, reaching into the white bin as the handle remains still. Then from 2.0s to 6.0s, the right arm grasps the black handle and moves to the right, pulling the handle and sliding the divider across the bin. Finally from 6.0s to 9.2s, the right arm releases the handle and moves upwards and to the right, returning to a resting position while the handle stays stationary. The box and bin remain fixed throughout.\n'
+    '\n'
+    '### Example 4 — EGO\n'
+    'USER PROMPT:\n'
+    'First-person POV, brightly lit grocery store produce section, soft artificial lighting. Two rectangular bins side-by-side in front; left filled with red/yellow apples, right piled with bright orange oranges. Background shows aisles and shoppers, one in red shirt with basket. Right hand enters from bottom, reaches into right orange pile, grasps one orange, lifts it upwards and slightly left. Camera moves forward and slightly down approaching bins, then remains stable with minor panning/tilting following hand movement.\n'
+    '\n'
+    'DETAILED CAPTION:\n'
+    "From a first-person perspective in a grocery store produce section, the camera approaches display bins containing apples on the left and bright oranges on the right. In the background aisles, a shopper wearing a red shirt walks from the left side of the frame towards the right between 0.0s and 2.0s, while another shopper in a striped shirt walks away from the camera down the aisle from 0.0s to 5.0s. The operator's right hand enters the frame from the bottom from 0.0s to 2.5s, then reaches into the pile of oranges and grasps one from 2.5s to 3.5s. Finally, the hand lifts the grasped orange upwards and to the left from 3.5s to 5.0s, selecting it from the cluster."
+)
+
+VIDEO_STEP2_MAP = (
+    'You are a structuring engine. You convert a DETAILED natural-language video caption (which concisely states the scene and the full motion/timing plan) into a STRUCTURED JSON caption. Your job is FAITHFUL structural extraction ONLY: do NOT drop, alter, reorder, or re-time anything the prose states (keep every subject, every action, and every timestamp exactly), and do NOT re-plan motion; BUT the prose is brief and intentionally omits fine visual attributes and camera settings — you MUST fill those omitted fields (texture, skin tone, precise colors, relative size, pose/orientation, clothing, and all camera_info) with plausible values that stay consistent with and never contradict the prose — just map what the prose already states into the schema.\n'
+    '\n'
+    '## OUTPUT FORMAT (JSON)\n'
+    '\n'
+    '```json\n'
+    '{\n'
+    '  "comprehensive_description": {\n'
+    '    "scene_content_description": "(String) A detailed description focusing on scene content, subject appearance, lighting, atmosphere, narrative, and interactions between elements. Text physically printed on visual objects should be mentioned alongside the object and detailed further in prominent_elements. For standalone OCR/text elements, briefly describe their role, relative scale, font style, color, and orientation; categorize them as \'static overlays\' (like watermarks) or \'integrated scene text\' (like subtitles/scrolling text); and describe their temporal behavior if they appear or disappear. Provide an exact transcription for prominent text (preserving spelling, punctuation, and capitalization), but summarize long or dense text blocks instead of transcribing them fully. **Do not describe camera movement here**. Maximum 800 words.",\n'
+    '    "camera_movement_description": "(String) A detailed description focusing on camera behavior. Include camera movement types (Pan, Tilt, Zoom, Dolly, Truck, Roll), shooting angles (high/low/eye-level), shot size changes, and stability. Maximum 100 words. If the camera is essentially stationary, set to \'\'."\n'
+    '  },\n'
+    '\n'
+    '  "prominent_elements": [\n'
+    '    {\n'
+    '      "name": "(String) Short label for the object (e.g., \'red sports car\', \'elderly man\')",\n'
+    '      "description": "(String) Detailed visual description of this specific element",\n'
+    '      "actions": [\n'
+    '        {\n'
+    '          "timestamp": "(String) Time range when this action occurs, e.g., \'[0.0s - 3.0s]\'",\n'
+    '          "action": "(String) Specific action description during this time period. **Direction must be described from the observer\'s perspective**. If the element has no action throughout, the entire actions array contains only one element with action set to \'\'."\n'
+    '        }\n'
+    '      ],\n'
+    '      "location": "(String) Precise position in the frame or main area of activity (from observer\'s perspective)",\n'
+    '      "relative_size": "(String) small / medium / large / dominant",\n'
+    '      "shape_and_color": "(String) Basic geometric shape and dominant colors",\n'
+    '      "texture": "(String) e.g., smooth, rough, metallic, furry, glossy, matte",\n'
+    '      "appearance_details": "(String) Specific details such as patterns, text physically printed on the object, wear marks, or distinctive markings",\n'
+    '      "relationship": "(String) This object\'s spatial or contextual relationship with other elements in the scene. If this object obscures text, specify exactly what it covers (e.g., \'blocking the letter O in COW). If it is obscured by floating text, describe what part of the object is covered.",\n'
+    '      "orientation": "(String) e.g., upright, tilted, horizontal, facing away, diagonal",\n'
+    '      \n'
+    '      "pose": "(String) Body posture and its changes (human/humanoid only, otherwise empty)",\n'
+    '      "expression": "(String) Facial expression and emotional changes (human/humanoid only, otherwise empty)",\n'
+    '      "clothing": "(String) Clothing description, including colors and styles (human/humanoid only, otherwise empty)",\n'
+    '      "gender": "(String) Apparent gender (human/humanoid only, otherwise empty)",\n'
+    '      "skin_tone_and_texture": "(String) Skin appearance (human/humanoid only, otherwise empty)",\n'
+    '      \n'
+    '      "is_cluster": "true (only for cluster objects, omit otherwise)",\n'
+    '      "number_of_objects": "(String) Exact number if countable, otherwise \'several\' (3-6), \'many\' (7-20), or \'numerous\' (20+)"\n'
+    '    }\n'
+    '  ],\n'
+    '\n'
+    '  "camera_info": {\n'
+    '    "color": "(String) Warm, Cool, Mixed, Saturated, Desaturated, Black and White, Red, Orange, Yellow, Green, Cyan, Blue, Magenta, or Pink",\n'
+    '    "frame_size": "(String) Extreme Wide, Wide, Medium Wide, Medium, Medium Close Up, Close Up, or Extreme Close Up",\n'
+    '    "shot_type_angle": "(String) High angle, Low angle, Dutch angle, Overhead, Aerial, or Eye level",\n'
+    '    "lens_size": "(String) Ultra Wide / Fisheye, Wide, Medium, Long Lens, or Telephoto",\n'
+    '    "composition": "(String) Center, Balanced, Symmetrical, Left heavy, Right heavy, or Short side",\n'
+    '    "lighting": "(String) Hard light, Soft light, High contrast, Low contrast, Side light, Top light, Underlight, Backlight, Edge light, or Silhouette",\n'
+    '    "lighting_type": "(String) Daylight, Sunny, Overcast, Moonlight, Artificial light, Practical light, Tungsten, Fluorescent, Firelight, or Mixed light"\n'
+    '  }\n'
+    '}\n'
+    '```\n'
+    '\n'
+    '## IMPORTANT RULES\n'
+    '\n'
+    '1.  **Think First:** Before generating the JSON, perform an internal "Let\'s think step by step" synthesis.\n'
+    '2.  **Output Only JSON:** Do not output any thinking process or Markdown text outside of the JSON code block.\n'
+    '3.  **Strict Fidelity (No Modification)**:\n'
+    '- Identity: Never alter the identity of subjects (e.g., a "real tiger" remains "real tiger").\n'
+    '- Actions: Verbs must be preserved exactly (e.g., "walking" cannot be changed to "running"). Sequence Integrity: You must preserve the complete action chain defined in the prompt. Do not summarize, skip, or merge distinct sequential steps into a single state. No Collective Summarization: Do not use summary phrases (e.g., "a sequence of...", "various movements") to cover multiple steps. If the prompt defines A -> B -> C, every step must be represented as a separate action segment. Any skipped step or summarized sequence is a CRITICAL ERROR.\n'
+    '- Quantity: If a number is specified, it must be exact. Quantifiers must be strictly preserved. "All" means all, "every" means every, "a single" means exactly 1. Do not paraphrase quantifiers into vague terms.\n'
+    '- Spatial Integrity: Relative positions (e.g., "A on the left of B") must be fixed as stated.\n'
+    '- Color & Text: Colors must be accurate; OCR text must be transcribed character-for-character (case-sensitive).\n'
+    '4.  **Creative Expansion Constraints:** When expanding a simple prompt, ensure added details (e.g., "sun-drenched oak windowsill") never contradict, replace, or occlude the user\'s explicit subject requirements.\n'
+    '5.  **Text Placement Logic:** \n'
+    '- Surface-Bound Text (Printed on objects): Must be transcribed exactly within the appearance_details field of its respective object in prominent_elements.\n'
+    '- Standalone Graphic Text (Floating/Poster text): Must NOT be an entry in prominent_elements. Describe its role, style, and exact transcription exclusively within the comprehensive_description.\n'
+    "- Dynamic Text: If the scene contains moving or updating text (e.g., a scrolling ticker or a subtitle), you must explicitly describe the text's movement path and content change frequency within the actions array of the respective prominent_elements or the comprehensive_description.\n"
+    '6.  **Occlusion & Relational Mapping:** Explicitly document overlaps. If text occludes an object, or an object occludes text, specify which character(s) or object parts are affected in both relationship and comprehensive_description.\n'
+    '7.  **Handle Missing Data:** If an attribute is not applicable to an element (e.g., pose for a mountain), set its value to "" (an empty string). Never use "N/A", "unknown", "not applicable", etc.\n'
+    '8.  **Counting:** For number_of_objects, provide an exact number if specified by the user. Otherwise, use: "several" (3-6), "many" (7-20), or "numerous" (20+).\n'
+    '9.  **Perspective:** Always describe positions (location) and directions (orientation) as "left" and "right" from the viewer\'s perspective.\n'
+    '10. **Visual Descriptive Realism:** Avoid abstract or emotional adjectives (e.g., "beautiful," "sad"). Instead, translate them into visually observable details (e.g., "soft golden-hour rim lighting," "desaturated cool blue tones with falling rain droplets").\n'
+    '11. **Clusters:** When describing a cluster (e.g., "a forest"), describe the collective appearance of the group rather than listing every individual tree.\n'
+    '12. **Attribute Consistency:** Ensure that the comprehensive_description and the individual prominent_elements are perfectly synchronized. Every element mentioned in the elements list must exist in the paragraph description and vice versa.\n'
+    '13. **Negative Constraints:** Any explicit prohibitions (e.g., "no text," "no background") must be strictly upheld. Never add elements that the user has explicitly requested to exclude.\n'
+    '14. **Subject Priority:** The primary subject specified by the user must be the anchor of the scene. Its description and prominence in the JSON must reflect its status as the most important element.\n'
+    '15. **Integrity:** Before generating JSON, you must cross-verify all keywords, actions, and states against the User Prompt; if any item is missing or any state intensity deviates from the original, you MUST regenerate the content to ensure 100% fidelity. \n'
+    '16. **Non-OCR Classification:** Do not force objects into "text" status. If a word appears in the prompt without context of "sign," "label," or "document," describe it as a physical attribute or brand marking on the object, not as OCR/Typography.\n'
+    "17. **Camera Fidelity:** The camera_info must be logically consistent with the visual scene. For example, if the scene is a vast landscape, the camera should reflect a 'Wide' or 'Ultra Wide' frame size and appropriate lens settings. Avoid contradictory pairings (e.g., 'Extreme Close Up' with a 'Wide' lens).\n"
+    '18. **Temporal Continuity:** Ensure temporal continuity for all subjects and environment settings. If an object is introduced, it must not disappear or reappear without a logical reason (e.g., leaving the frame or being occluded). Texture, color, and lighting must remain consistent across all action segments.\n'
+    "19. **Action-Camera Sync:** The camera movement must logically justify the subject's displacement. If the camera follows a subject (tracking shot), the subject's position in the frame should remain relatively stable. If the camera is static, the subject must show movement within the frame.\n"
+    "20. **Physical Logic:** All movements must respect physical laws unless the prompt specifies a surreal or fantasy context. Avoid 'gliding' motions; ensure footsteps or object interactions (e.g., picking up an item) align with the timestamped actions.\n"
+    '21. **Motion and Velocity:** Verbs must describe the process of an action, not just the result (e.g., use \'rising from the chair\' instead of just \'standing\'). Every action in the actions array must include a velocity modifier (e.g., "slowly walking," "abruptly turning," "smoothly panning," "rapidly accelerating"). Avoid generic verbs without speed or force descriptors.\n'
+    '22. **Language:** Output all JSON content in professional English, except for OCR text, which must be transcribed exactly as provided to maintain absolute fidelity.\n'
+    '23. **Audio Neglect:** The User Prompt may contain audio, music, or voiceover instructions. Strictly ignore all audio-related instructions. Do not attempt to rewrite, describe, or represent these as part of the JSON output. Focus exclusively on the visual, temporal, and narrative elements and OCR texts.\n'
+    '24. **Duration Adherence:** A target video duration is provided in the input. Every timestamp in the `actions` array MUST fall within `[0.0s, the given duration]`, and the action timeline should span up to (but never beyond) this duration. Distribute the action segments proportionally to the given length; the final segment must end at or before the given duration.\n'
+    '\n'
+    "## THIS STEP'S INPUT/OUTPUT\n"
+    '- Input: the detailed prose caption + the target video duration.\n'
+    '- The expansion/thinking is already done in the prose — preserve the stated core EXACTLY (scene, subjects, every action + timestamp, named entities); ONLY the fine visual attributes and camera fields that the brief prose omits may be completed — plausibly, consistently, never contradicting the prose.\n'
+    '- Output ONLY the JSON object (valid and parseable), following the schema and rules above.'
+)
+
+IMAGE_STEP1_EXPAND = (
+    'You write ONE short, natural, standalone IMAGE caption in English from a brief user image\n'
+    'prompt — as if briefly describing a real photo of that idea.\n'
+    '\n'
+    'This is a SINGLE STILL IMAGE: there is NO motion, NO time, NO timestamps, NO camera movement.\n'
+    'Describe only what the photo shows; never invent action, sequence, or a moving camera.\n'
+    '\n'
+    'HARD LENGTH LIMIT: the output MUST be UNDER 1000 characters. Aim for roughly 400-800\n'
+    'characters. Concise is a top priority.\n'
+    '\n'
+    'STYLE — a flowing description, NOT a checklist:\n'
+    '- Connected, natural English prose. NO headings, NO bullets, NO field labels.\n'
+    '- The MAIN THREAD is the main subject and the scene — lead with WHAT THE IMAGE SHOWS. Do NOT\n'
+    '  march subject-by-subject listing attributes; weave the few details you keep into the\n'
+    '  description naturally. One cohesive paragraph; never pad to add length.\n'
+    '\n'
+    'WHAT TO COVER (only these):\n'
+    '- The scene in a brief phrase (where it is / overall vibe).\n'
+    '- Each main subject by name/what-it-is plus ONE most defining visual trait — not a full\n'
+    '  appearance list. The key spatial arrangement only if it defines the image.\n'
+    '- A one-word shot type only if it matters (e.g. close-up, wide shot, aerial); otherwise skip the camera.\n'
+    "- Named entities (real people / places / landmarks / brands) only if clearly implied — don't invent identities.\n"
+    '\n'
+    'FAITHFUL EXPANSION:\n'
+    '- Stay consistent with the prompt; elaborate plausibly but never contradict it. Commit to one\n'
+    "  coherent interpretation; don't hedge with alternatives.\n"
+    '- Keep subject identities and counts consistent.\n'
+    '\n'
+    'Output ONLY the caption text — no preamble, no headings, no explanation.\n'
+    '\n'
+    '\n'
+    "## EXAMPLES — one representative example per type. When you write your own caption, match these examples' style, format, and length (one flowing paragraph, no timestamps, no camera movement — it is a single still image).\n"
+    '\n'
+    '### Example 1 — animal / wildlife\n'
+    'USER PROMPT:\n'
+    'A leopard tortoise dominates the center of the frame, captured from a high angle with its head and front legs extended toward the left. Its domed shell displays a striking mosaic of black and tan geometric patterns, contrasting with the rough, yellowish-tan scales of its skin and sharp claws. The creature is positioned on reddish-brown sandy soil scattered with dry twigs and pebbles, illuminated by warm daylight that creates a soft shadow beneath it against a backdrop of green grass on the right.\n'
+    '\n'
+    'DETAILED CAPTION:\n'
+    'A leopard tortoise is positioned centrally in a natural outdoor setting, resting on reddish-brown sandy soil scattered with small pebbles and dry twigs. Facing toward the left with its head and front legs extended, the animal displays a domed shell marked by intricate black and tan geometric patterns reminiscent of leopard spots. Its yellowish-tan skin contrasts with the earthy ground beneath its clawed feet, while a patch of green grass and low-lying vegetation rises in the background to the right. The distinct markings of the shell stand out against the textured ground and sparse greenery, capturing the tortoise within its earthy environment.\n'
+    '\n'
+    '### Example 2 — architecture / landmark\n'
+    'USER PROMPT:\n'
+    '极度广角镜头，高角度拍摄。明亮日光，硬光，暖色调，清晰阴影。宏伟的古典风格宫殿庭院，浅色石材。中央巨大的多层拱形入口，饰有金色和蓝色装饰带，深色凹陷内部可见小金门。宽阔石阶通向圆形分层平台。两侧对称亭阁，顶部为反射金色圆顶，蓝色瓷砖拱门，装饰攀爬绿藤和粉色花朵。露台布满茂盛绿植和小花。若干穿着传统长袍（白、灰、棕、红）的小人物散布庭院以示比例。背景岩石山坡，远处建筑，晴朗明亮天空。写实建筑摄影。\n'
+    '\n'
+    'DETAILED CAPTION:\n'
+    'A grand, sun-drenched courtyard of a classical-style palace complex is dominated by a massive central building featuring a large, multi-tiered arched entrance. Wide stone stairs lead up to the structure, flanked on either side by identical pavilions topped with brilliant, reflective golden domes and blue-tiled arches. The light-colored stone complex includes various terraces filled with lush greenery, while several small figures dressed in traditional robes are scattered throughout the courtyard, providing a sense of scale against the monumental architecture. In the background, a rocky hillside rises under a clear bright sky, dotted with additional buildings and vegetation surrounding the temple-like grounds.\n'
+    '\n'
+    '### Example 3 — landscape / nature\n'
+    'USER PROMPT:\n'
+    "Make me an image of a serene coastal scene featuring a classic white lighthouse attached to a keeper's house, with a vintage red pickup truck and a small boat on a trailer parked on a dirt path nearby, set against a lush green lawn, wooden fence, and bright blue sky over calm waters.\n"
+    '\n'
+    'DETAILED CAPTION:\n'
+    "A wide aerial view captures a serene coastal scene dominated by a classic white lighthouse and its attached keeper's house. The tall cylindrical tower features a black lantern room and rises from a white residence with a brown shingled roof and green shutters. To the left, a classic red pickup truck is parked on a dirt path with a small white and red boat resting on a trailer behind it. A lush green lawn slopes down toward a rocky shoreline where a weathered wooden fence runs along the edge of the deep blue water. In the background, a calm sea stretches to the horizon, meeting a distant tree-covered coastline under a bright blue sky filled with soft clouds.\n"
+    '\n'
+    '### Example 4 — person / portrait\n'
+    'USER PROMPT:\n'
+    'A stylish woman with long wavy blonde hair and bold red lipstick stands confidently on a city sidewalk, wearing a black long-sleeved dress with sheer patterned sleeves and a ruffled waistline. She carries a small black crossbody bag with a gold Saint Laurent logo slung over her shoulder, her left arm slightly extended and right arm by her side in a high-angle medium wide shot. The background features brick buildings and a blurred crosswalk with distant pedestrians under soft daylight, creating a sophisticated urban atmosphere with a left-heavy composition.\n'
+    '\n'
+    'DETAILED CAPTION:\n'
+    'A stylish woman with long blonde hair and bold red lipstick stands on a city sidewalk, gazing directly forward in a black long-sleeved dress with sheer patterned sleeves and a ruffled waistline. A small black crossbody bag featuring a gold Saint Laurent logo hangs at her side, complementing her sophisticated look. The setting is a classic urban street scene characterized by brick buildings and a black wall-mounted lamp visible behind her. In the distance, a blurred crosswalk shows several pedestrians, adding depth to the modern atmosphere surrounding the central figure.'
+)
+
+IMAGE_STEP2_MAP = (
+    'You convert a SHORT natural-language STILL-IMAGE caption (the DETAILED CAPTION) into ONE\n'
+    'structured JSON caption describing the image. Output ONLY the JSON object — no prose, no code fence.\n'
+    '\n'
+    '## TASK SEMANTICS (read carefully)\n'
+    'The detailed caption is SHORT and LOSSY: it states the scene, the main subjects, and each\n'
+    "subject's single most defining trait, but it deliberately OMITS fine attributes (exact colors,\n"
+    'texture, relative size, orientation, clothing, skin tone, pose, expression) and the camera\n'
+    'metadata. Your job:\n'
+    '- KEEP faithfully everything the prose states (scene, every named subject, named entities) — do\n'
+    '  not drop, rename, merge, or split subjects, and do not contradict the prose.\n'
+    '- REASONABLY COMPLETE the fields the prose omits (fine attributes + camera_info) with plausible,\n'
+    '  self-consistent values that do NOT contradict the prose. This is expected: the target JSON is a\n'
+    '  FULL caption, so every schema field must be filled even when the prose did not mention it.\n'
+    '- This is a STILL IMAGE: no actions, no timestamps, no camera movement anywhere.\n'
+    '\n'
+    '## OUTPUT SCHEMA (exact keys)\n'
+    'A single JSON object with EXACTLY these top-level keys:\n'
+    '- "comprehensive_description": string — one flowing prose paragraph describing the whole image\n'
+    '  (scene, subjects, their arrangement and look). No field labels inside it.\n'
+    '- "camera_info": object with EXACTLY these 7 string keys, each set to ONE value from its list:\n'
+    '  - "color": Warm | Cool | Mixed | Saturated | Desaturated | White | Green | Blue | Red | Cyan\n'
+    '  - "frame_size": Extreme Close Up | Close Up | Medium Close Up | Medium | Medium Wide | Wide | Extreme Wide\n'
+    '  - "shot_type_angle": Eye level | Low angle | High angle | Overhead | Aerial\n'
+    '  - "lens_size": Ultra Wide / Fisheye | Wide | Medium | Long Lens | Telephoto\n'
+    '  - "composition": Center | Balanced | Symmetrical | Left heavy | Right heavy\n'
+    '  - "lighting": Soft light | Hard light | Top light | Underlight | Backlight\n'
+    '  - "lighting_type": Daylight | Artificial light\n'
+    '- "world_knowledge": array of strings — named entities / real-world facts (people, places,\n'
+    '  landmarks, brands) stated or clearly implied by the prose; [] if none. NEVER invent identities.\n'
+    '- "prominent_elements": array of objects, one per notable subject/object. Each element has:\n'
+    '  - always: "name", "description", "location", "relative_size" (one of: dominant | large | medium | small),\n'
+    '    "shape_and_color", "texture", "appearance_details", "relationship", "orientation"\n'
+    '  - for a PERSON, additionally: "pose", "expression", "clothing", "gender" (male | female),\n'
+    '    "skin_tone_and_texture"\n'
+    '  - for a GROUP/crowd of like items, additionally: "is_cluster": true and "number_of_objects": "<count or range as string>"\n'
+    '\n'
+    '## RULES\n'
+    '- Output STRICT valid JSON only (double quotes, no trailing commas, no comments, no code fence).\n'
+    '- Fill EVERY field; never leave a required field empty. Values you complete must be consistent\n'
+    '  with the prose and with each other.\n'
+    '- Keep the number and identity of prominent_elements aligned with the subjects the prose names\n'
+    '  (you may add a clearly-implied background element, but do not invent unrelated subjects).\n'
+    '- Copy named entities verbatim into world_knowledge.\n'
+    '- No actions, no timestamps, no camera-movement fields — this is a still image.'
+)