init: LingBot-Video MoE 30B-A3B ZeroGPU demo

.gitattributes

@@ -33,3 +33,5 @@ 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_1.png filter=lfs diff=lfs merge=lfs -text
+examples/ti2v_5.png filter=lfs diff=lfs merge=lfs -text

README.md

@@ -1,13 +1,32 @@
 ---
-title: Lingbot Video
-emoji: 🐠
-colorFrom: yellow
-colorTo: green
+title: LingBot-Video MoE 30B-A3B
+emoji: 🎬
+colorFrom: indigo
+colorTo: purple
 sdk: gradio
-sdk_version: 6.20.0
-python_version: '3.12'
+sdk_version: 6.10.0
 app_file: app.py
-pinned: false
+license: apache-2.0
+short_description: Embodied MoE video generation — T2V, I2V, T2I
+models:
+  - robbyant/lingbot-video-moe-30b-a3b
+startup_duration_timeout: 1h
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# LingBot-Video MoE 30B-A3B — ZeroGPU demo
+
+Text-to-video, image-to-video and text-to-image with
+[LingBot-Video MoE 30B-A3B](https://huggingface.co/robbyant/lingbot-video-moe-30b-a3b),
+the first open-source large-scale MoE video generation model dedicated to embodied intelligence
+(30B parameters, ~3B active per token).
+
+This demo runs the base MoE DiT (the refiner stage is skipped to fit ZeroGPU) with:
+
+- bf16 transformer + text encoder (Qwen3-VL), fp32 Wan VAE
+- 480p resolution buckets, up to 5 s at 24 fps
+- sequential CFG with SDPA flash attention
+- `torch._grouped_mm` MoE expert dispatch
+
+The model was trained on detailed structured captions — long, specific prompts (or JSON captions
+from the [LingBot rewriter](https://huggingface.co/robbyant/lingbot-video-rewriter-lora)) work
+much better than short casual ones.

app.py

@@ -0,0 +1,456 @@
+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("LINGBOT_MOE_EXPERT_BACKEND", "grouped_mm")
+os.environ.setdefault("LINGBOT_MOE_PAD_BACKEND", "vectorized")
+os.environ.setdefault("DIFFUSERS_ATTN_BACKEND", "_native_flash")
+
+import spaces  # noqa: E402 — must be imported before torch
+
+import random  # noqa: E402
+import time  # noqa: E402
+from pathlib import Path  # noqa: E402
+
+import gradio as gr  # noqa: E402
+import numpy as np  # noqa: E402
+import torch  # noqa: E402
+from PIL import Image  # noqa: E402
+from diffusers import AutoencoderKLWan  # noqa: E402
+from diffusers.utils import export_to_video  # noqa: E402
+from huggingface_hub import snapshot_download  # noqa: E402
+from transformers import AutoProcessor, Qwen3VLForConditionalGeneration  # noqa: E402
+
+from lingbot_video import (  # noqa: E402
+    FlowUniPCMultistepScheduler,
+    LingBotVideoImageToVideoPipeline,
+    LingBotVideoPipeline,
+    LingBotVideoTransformer3DModel,
+)
+from lingbot_video.pipeline_lingbot_video import (  # noqa: E402
+    DEFAULT_NEGATIVE_PROMPT,
+    DEFAULT_NEGATIVE_PROMPT_IMAGE,
+)
+from lingbot_video.utils import num_frames_from_duration  # noqa: E402
+
+MODEL_ID = "robbyant/lingbot-video-moe-30b-a3b"
+FPS = 24
+MAX_SEED = 2**31 - 1
+MAX_GPU_SECONDS = 300
+
+# (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} (refiner excluded) ...", flush=True)
+t0 = time.perf_counter()
+model_dir = snapshot_download(
+    MODEL_ID,
+    ignore_patterns=["refiner/*", "*.bak_lingbot_video_diffusers"],
+)
+print(f"[startup] snapshot ready in {time.perf_counter() - t0:.1f}s", flush=True)
+
+print(f"[startup] torch {torch.__version__}, grouped_mm available: {hasattr(torch, '_grouped_mm')}", 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)
+
+
+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:
+    # Empirical per-step cost model (CFG included): quadratic attention term +
+    # linear MoE/MLP term, normalized to the 81-frame 480x832 token count.
+    r = _latent_tokens(height, width, num_frames) / 32760.0
+    return 9.0 * r * r + 6.0 * r
+
+
+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(40.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(30.0 + steps * max(_step_seconds(height, width, 1), 0.35))
+
+
+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 _t2v_duration(prompt, size_label, duration_s, steps, *args, **kwargs) -> int:
+    return min(MAX_GPU_SECONDS, _estimate_video_seconds(size_label, duration_s, int(steps)) + 20)
+
+
+@spaces.GPU(duration=_t2v_duration, size="xlarge")
+def generate_t2v(
+    prompt,
+    size_label,
+    duration_s,
+    steps,
+    guidance_scale,
+    shift,
+    negative_prompt,
+    seed,
+    randomize_seed,
+    progress=gr.Progress(track_tqdm=True),
+):
+    if not prompt or not prompt.strip():
+        raise gr.Error("Please enter a prompt.")
+    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)
+    t0 = time.perf_counter()
+    output = pipe(
+        prompt=prompt.strip(),
+        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),
+        generator=torch.Generator().manual_seed(seed),
+    )
+    torch.cuda.synchronize()
+    elapsed = time.perf_counter() - t0
+    print(
+        f"[timing] t2v {width}x{height} f={num_frames} steps={int(steps)} "
+        f"tokens={_latent_tokens(height, width, num_frames)} took {elapsed:.1f}s "
+        f"(estimated {estimate}s)",
+        flush=True,
+    )
+    return _export_video(output.frames[0]), seed
+
+
+def _i2v_duration(image, prompt, size_label, duration_s, steps, *args, **kwargs) -> int:
+    return min(MAX_GPU_SECONDS, _estimate_video_seconds(size_label, duration_s, int(steps)) + 35)
+
+
+@spaces.GPU(duration=_i2v_duration, size="xlarge")
+def generate_i2v(
+    image,
+    prompt,
+    size_label,
+    duration_s,
+    steps,
+    guidance_scale,
+    shift,
+    negative_prompt,
+    seed,
+    randomize_seed,
+    progress=gr.Progress(track_tqdm=True),
+):
+    if image is None:
+        raise gr.Error("Please upload a first-frame image.")
+    if not prompt or not prompt.strip():
+        raise gr.Error("Please enter a prompt.")
+    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)
+    t0 = time.perf_counter()
+    output = pipe_i2v(
+        prompt=prompt.strip(),
+        image=image,
+        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),
+        generator=torch.Generator().manual_seed(seed),
+    )
+    torch.cuda.synchronize()
+    elapsed = time.perf_counter() - t0
+    print(
+        f"[timing] i2v {width}x{height} f={num_frames} steps={int(steps)} took {elapsed:.1f}s "
+        f"(estimated {estimate}s)",
+        flush=True,
+    )
+    return _export_video(output.frames[0]), seed
+
+
+def _t2i_duration(prompt, size_label, steps, *args, **kwargs) -> int:
+    return min(MAX_GPU_SECONDS, _estimate_image_seconds(size_label, int(steps)) + 20)
+
+
+@spaces.GPU(duration=_t2i_duration, size="xlarge")
+def generate_t2i(
+    prompt,
+    size_label,
+    steps,
+    guidance_scale,
+    shift,
+    negative_prompt,
+    seed,
+    randomize_seed,
+    progress=gr.Progress(track_tqdm=True),
+):
+    if not prompt or not prompt.strip():
+        raise gr.Error("Please enter a prompt.")
+    height, width = IMAGE_SIZES[size_label]
+    seed = _resolve_seed(seed, randomize_seed)
+    t0 = time.perf_counter()
+    output = pipe(
+        prompt=prompt.strip(),
+        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),
+        generator=torch.Generator().manual_seed(seed),
+    )
+    torch.cuda.synchronize()
+    elapsed = time.perf_counter() - t0
+    print(f"[timing] t2i {width}x{height} steps={int(steps)} took {elapsed:.1f}s", flush=True)
+    frame = output.frames[0][0]
+    image = Image.fromarray((np.clip(frame, 0, 1) * 255).astype(np.uint8))
+    return image, seed
+
+
+DESCRIPTION = """
+# LingBot-Video MoE 30B-A3B
+
+First open-source large-scale MoE video generation model dedicated to **embodied intelligence** — 30B parameters,
+3B active per token. Text-to-video, image-to-video and text-to-image with a single model.
+
+[Model card](https://huggingface.co/robbyant/lingbot-video-moe-30b-a3b) ·
+[GitHub](https://github.com/Robbyant/lingbot-video) ·
+[Project page](https://technology.robbyant.com/lingbot-video)
+
+> 💡 The model was trained on **detailed structured captions**: long, specific prompts (or the JSON captions produced
+> by the [LingBot rewriter](https://huggingface.co/robbyant/lingbot-video-rewriter-lora)) work much better than short ones.
+> This demo runs the base MoE model (no refiner) at 480p on ZeroGPU.
+"""
+
+
+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=40, 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
+
+
+with gr.Blocks(title="LingBot-Video MoE 30B-A3B") as demo:
+    gr.Markdown(DESCRIPTION)
+
+    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_steps, t2v_guidance, t2v_shift, t2v_negative, t2v_seed, t2v_rand = _advanced(
+                    DEFAULT_NEGATIVE_PROMPT, 24
+                )
+                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)
+
+        t2v_inputs = [
+            t2v_prompt, t2v_size, t2v_dur, t2v_steps, t2v_guidance, t2v_shift,
+            t2v_negative, t2v_seed, t2v_rand,
+        ]
+        t2v_btn.click(generate_t2v, inputs=t2v_inputs, outputs=[t2v_out, t2v_seed_out],
+                      concurrency_id="gpu", concurrency_limit=1)
+        gr.Examples(
+            examples=[
+                [_read_example("t2v_bubbles.txt"), "832 × 480 (16:9)", 2.0, 24, 3.0, 3.0,
+                 DEFAULT_NEGATIVE_PROMPT, 42, False],
+                [_read_example("t2v_robot_desk.txt"), "832 × 480 (16:9)", 2.0, 24, 3.0, 3.0,
+                 DEFAULT_NEGATIVE_PROMPT, 42, False],
+                ["A corgi puppy runs through shallow water on a beach at golden hour. Slow motion, "
+                 "splashes glistening in warm sunlight, shallow depth of field, cinematic color grading. "
+                 "The camera tracks alongside the puppy at a low angle close to the water surface.",
+                 "832 × 480 (16:9)", 2.0, 24, 3.0, 3.0, DEFAULT_NEGATIVE_PROMPT, 42, False],
+            ],
+            fn=generate_t2v,
+            inputs=t2v_inputs,
+            outputs=[t2v_out, t2v_seed_out],
+            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_steps, i2v_guidance, i2v_shift, i2v_negative, i2v_seed, i2v_rand = _advanced(
+                    DEFAULT_NEGATIVE_PROMPT, 24
+                )
+                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)
+
+        i2v_inputs = [
+            i2v_image, i2v_prompt, i2v_size, i2v_dur, i2v_steps, i2v_guidance, i2v_shift,
+            i2v_negative, i2v_seed, i2v_rand,
+        ]
+        i2v_btn.click(generate_i2v, inputs=i2v_inputs, outputs=[i2v_out, i2v_seed_out],
+                      concurrency_id="gpu", concurrency_limit=1)
+        gr.Examples(
+            examples=[
+                [str(EXAMPLES_DIR / "ti2v_1.png"), _read_example("ti2v_1.txt"),
+                 "832 × 480 (16:9)", 2.0, 24, 3.0, 3.0, DEFAULT_NEGATIVE_PROMPT, 42, False],
+                [str(EXAMPLES_DIR / "ti2v_5.png"), _read_example("ti2v_5.txt"),
+                 "832 × 480 (16:9)", 2.0, 24, 3.0, 3.0, DEFAULT_NEGATIVE_PROMPT, 42, False],
+            ],
+            fn=generate_i2v,
+            inputs=i2v_inputs,
+            outputs=[i2v_out, i2v_seed_out],
+            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="1280 × 736 (16:9, 720p)"
+                )
+                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)
+
+        t2i_inputs = [
+            t2i_prompt, t2i_size, t2i_steps, t2i_guidance, t2i_shift,
+            t2i_negative, t2i_seed, t2i_rand,
+        ]
+        t2i_btn.click(generate_t2i, inputs=t2i_inputs, outputs=[t2i_out, t2i_seed_out],
+                      concurrency_id="gpu", concurrency_limit=1)
+        gr.Examples(
+            examples=[
+                [_read_example("t2i_bottle.txt"), "1280 × 736 (16:9, 720p)", 30, 3.0, 3.0,
+                 DEFAULT_NEGATIVE_PROMPT_IMAGE, 42, False],
+                ["A humanoid robot chef flipping a pancake in a bright modern kitchen, "
+                 "photorealistic, dramatic side lighting, steam rising from the pan, "
+                 "shallow depth of field, 85mm lens.",
+                 "1280 × 736 (16:9, 720p)", 30, 3.0, 3.0, DEFAULT_NEGATIVE_PROMPT_IMAGE, 42, False],
+            ],
+            fn=generate_t2i,
+            inputs=t2i_inputs,
+            outputs=[t2i_out, t2i_seed_out],
+            cache_examples=True,
+            cache_mode="lazy",
+        )
+
+demo.queue(max_size=30).launch()

examples/t2i_bottle.txt

@@ -0,0 +1 @@
+{"comprehensive_description":"A clear glass bottle filled with water sits on a light-colored wooden table, acting as a lens to focus sunlight. The bottle is positioned slightly to the left of the center, with its base resting on the table. Bright, direct sunlight from the upper right creates a brilliant, starburst-like lens flare on the table's surface, casting a warm, golden glow. The background is a soft, out-of-focus blur of green foliage and bright light, suggesting an outdoor garden or patio setting. The overall atmosphere is warm, serene, and summery, emphasizing the interplay between light and water.","camera_info":{"color":"Warm","frame_size":"Extreme Close Up","shot_type_angle":"Low angle","lens_size":"Long Lens","composition":"Left heavy","lighting":"Hard light","lighting_type":"Daylight"},"world_knowledge":[],"prominent_elements":[{"name":"glass bottle","description":"A clear, cylindrical glass bottle with a narrow neck and a flat base, partially filled with water.","location":"center-left","relative_size":"medium","shape_and_color":"Cylindrical with a tapered neck; transparent and clear","texture":"smooth, glossy","appearance_details":"The bottle contains clear water that reflects the surrounding light. The glass surface shows subtle highlights and reflections from the sun.","relationship":"Sits on the wooden table and acts as the focal point for the sunlight, creating a lens flare.","orientation":"upright","pose":"","expression":"","clothing":"","gender":"","skin_tone_and_texture":""},{"name":"wooden table","description":"A flat surface made of light-colored wood with visible grain and horizontal planks.","location":"spanning the bottom half of the frame","relative_size":"large","shape_and_color":"Rectangular planks; light tan and beige","texture":"matte, grainy","appearance_details":"The surface is illuminated by a bright, star-shaped lens flare caused by the bottle focusing the sunlight.","relationship":"Provides the base upon which the bottle rests.","orientation":"horizontal","pose":"","expression":"","clothing":"","gender":"","skin_tone_and_texture":""},{"name":"lens flare","description":"A bright, star-shaped burst of light created by sunlight passing through the water in the bottle.","location":"center, on the table surface","relative_size":"medium","shape_and_color":"Star-shaped; bright white and golden yellow","texture":"ethereal, glowing","appearance_details":"The flare has multiple sharp rays extending outwards from the point where the bottle meets the table.","relationship":"Originates from the interaction of sunlight and the water inside the glass bottle.","orientation":"radiating","pose":"","expression":"","clothing":"","gender":"","skin_tone_and_texture":""},{"name":"blurred foliage","description":"A soft-focus background consisting of green leaves and bright light spots.","location":"top half of the frame","relative_size":"large","shape_and_color":"Amorphous shapes; various shades of green and bright white","texture":"soft, blurry","appearance_details":"The background is heavily out of focus, creating a bokeh effect that emphasizes the sharpness of the bottle in the foreground.","relationship":"Provides a natural, outdoor context for the scene.","orientation":"upright","pose":"","expression":"","clothing":"","gender":"","skin_tone_and_texture":"","is_cluster":true,"number_of_objects":"numerous"}]}

examples/t2v_bubbles.txt

@@ -0,0 +1 @@
+{"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"}]}

examples/t2v_robot_desk.txt

@@ -0,0 +1 @@
+{"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":""}]}

examples/ti2v_1.txt

@@ -0,0 +1 @@
+{"comprehensive_description":{"scene_content_description":"A high-end, cinematic still-life advertisement featuring a thick-bottomed, transparent whiskey glass resting on a dark wooden table. The glass contains amber-colored liquid, and a single transparent ice cube is suspended in mid-air directly above the rim, poised to fall. The background is dark and out of focus, creating a shallow depth of field that isolates the subject. The atmosphere is sophisticated, moody, and luxurious, emphasizing the clarity of the glass and the rich color of the whiskey.","camera_movement_description":"The camera remains in a stable close-up perspective throughout the duration. There is no significant panning, tilting, or zooming. Instead, the camera exhibits a very subtle, smooth, slow-motion breathing effect, maintaining a steady focus on the glass and the falling ice."},"camera_info":{"color":"Warm","frame_size":"Close Up","shot_type_angle":"Eye level","lens_size":"Macro","composition":"Center","lighting":"Soft light","lighting_type":"Studio lighting"},"world_knowledge":[],"prominent_elements":[{"name":"transparent ice cube","description":"A single, clear, roughly cubic block of ice","actions":[{"timestamp":"[0.0s - 0.8s]","action":"falls from above the rim, rotating slightly as it descends"},{"timestamp":"[0.8s - 1.6s]","action":"strikes the surface of the liquid, creating a circular depression and initial splash"},{"timestamp":"[1.6s - 2.8s]","action":"partially submerges, generating splashes and bubbles"},{"timestamp":"[2.8s - 4.0s]","action":"bobs slightly up and down and rotates within the liquid"},{"timestamp":"[4.0s - 5.0s]","action":"floats gently in the liquid, slowly coming to a stop"}],"location":"initially in the upper center of the frame, then moves into the center of the glass","relative_size":"medium","shape_and_color":"geometric cube, transparent with internal refractions","texture":"crystalline and solid","appearance_details":"visible internal air bubbles, cracks, and sharp edges that catch the light","relationship":"the dynamic object interacting with the liquid","orientation":"slightly tilted","pose":"","expression":"","clothing":"","gender":"","skin_tone_and_texture":""},{"name":"whiskey glass","description":"A heavy, short tumbler with a thick, textured base","actions":[{"timestamp":"[0.0s - 5.0s]","action":""}],"location":"center of the frame, resting on the table","relative_size":"large","shape_and_color":"cylindrical glass, transparent","texture":"smooth and reflective","appearance_details":"thick, grid-patterned base and clear walls that refract the amber liquid inside","relationship":"the container for the liquid and ice","orientation":"upright","pose":"","expression":"","clothing":"","gender":"","skin_tone_and_texture":""},{"name":"amber liquid","description":"A rich, golden-brown fluid filling the lower half of the glass","actions":[{"timestamp":"[0.0s - 0.8s]","action":"remains relatively still with a flat surface"},{"timestamp":"[0.8s - 1.6s]","action":"surface is pushed down by the ice, forming a circular wave"},{"timestamp":"[1.6s - 2.8s]","action":"creates splashes and fine droplets as the ice enters"},{"timestamp":"[2.8s - 4.0s]","action":"exhibits ripples that bounce off the glass walls and cross-ripples"},{"timestamp":"[4.0s - 5.0s]","action":"surface gradually calms down, becoming smooth again"}],"location":"inside the glass","relative_size":"large","shape_and_color":"fluid volume, deep amber","texture":"viscous and glossy","appearance_details":"surface tension and the way light passes through it","relationship":"contained within the glass","orientation":"horizontal","pose":"","expression":"","clothing":"","gender":"","skin_tone_and_texture":""},{"name":"dark wooden table","description":"A polished, dark wood surface","actions":[{"timestamp":"[0.0s - 5.0s]","action":""}],"location":"foreground and bottom of the frame","relative_size":"large","shape_and_color":"flat surface, dark brown/black","texture":"smooth with a slight sheen","appearance_details":"wood grain and reflections of the glass and light","relationship":"the supporting surface for the glass","orientation":"horizontal","pose":"","expression":"","clothing":"","gender":"","skin_tone_and_texture":""}]}

examples/ti2v_5.txt

@@ -0,0 +1 @@
+{"comprehensive_description":{"scene_content_description":"A sleek, vibrant red sports car is captured in a high-speed motion sequence. The car begins by emerging from a dark, shadowed cave entrance, its glossy surface reflecting the intense light of the desert environment. As it accelerates, it kicks up a significant amount of dust and gravel from the sandy ground. The car then bursts into the bright, sun-drenched exterior, where the background reveals a vast, arid landscape with distant mountains under a clear sky. The lighting is dramatic, with high-contrast shadows inside the cave and brilliant, almost blinding sunlight outside, creating a sense of power and speed.","camera_movement_description":"The camera performs a dynamic tracking shot, moving parallel to the car's side. It starts with a close-up side profile and gradually pulls back and pans to the right, following the car's trajectory. As the car exits the cave, the camera angle shifts slightly upward and pulls back further to capture the car's front and side as it speeds into the open desert. The movement is smooth and fluid, maintaining a consistent distance from the subject."},"camera_info":{"color":"Warm","frame_size":"Wide","shot_type_angle":"Low angle","lens_size":"Wide","composition":"Balanced","lighting":"Hard light","lighting_type":"Daylight"},"world_knowledge":[],"prominent_elements":[{"name":"red sports car","description":"A modern, high-performance coupe with a glossy red finish and black accents","actions":[{"timestamp":"[0.0s - 4.0s]","action":"moves right while emerging from a dark cave"},{"timestamp":"[4.0s - 6.0s]","action":"accelerates and exits the cave into bright sunlight"},{"timestamp":"[6.0s - 8.0s]","action":"speeds right across the desert landscape"},{"timestamp":"[8.0s - 10.0s]","action":"continues moving right with the camera tracking its side"}],"location":"center of the frame, moving from left to right","relative_size":"dominant","shape_and_color":"aerodynamic coupe shape, vibrant red with black wheels and trim","texture":"glossy and metallic","appearance_details":"features a low-slung profile, large black multi-spoke wheels, and a visible driver inside","relationship":"the main subject of the video, interacting with the cave and desert environment","orientation":"facing right","pose":"","expression":"","clothing":"","gender":"","skin_tone_and_texture":""},{"name":"cave entrance","description":"A dark, rocky opening in a cliff face","actions":[{"timestamp":"[0.0s - 4.0s]","action":""}],"location":"left side of the frame","relative_size":"large","shape_and_color":"irregular rocky shape, dark brown and black","texture":"rough and rocky","appearance_details":"deep shadows inside the opening, contrasting with the bright exterior","relationship":"the starting point for the car's movement","orientation":"upright","pose":"","expression":"","clothing":"","gender":"","skin_tone_and_texture":""},{"name":"desert landscape","description":"A vast, arid environment with sandy ground and distant mountains","actions":[{"timestamp":"[0.0s - 10.0s]","action":""}],"location":"background and foreground","relative_size":"large","shape_and_color":"sandy beige ground, light blue sky, and brown mountains","texture":"sandy and dusty","appearance_details":"bright sunlight creates a hazy, high-contrast atmosphere","relationship":"the setting for the car's high-speed run","orientation":"horizontal","pose":"","expression":"","clothing":"","gender":"","skin_tone_and_texture":""}]}

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/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/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/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 @@
+spaces>=0.41
+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