app.py

import os

os.environ.setdefault("HF_HOME", "/tmp/hf_home")
os.environ.setdefault("HF_MODULES_CACHE", "/tmp/hf_modules")
os.environ.setdefault("MPLCONFIGDIR", "/tmp/matplotlib")
# LingBot-Video runtime knobs (read at import / first forward).
os.environ.setdefault("DIFFUSERS_ATTN_BACKEND", "_native_flash")

import spaces  # noqa: E402 — must be imported before torch

import base64
import io
import json
import random
import re
import time
import urllib.request
from pathlib import Path

import gradio as gr
import numpy as np
import torch
from PIL import Image
from diffusers import AutoencoderKLWan
from diffusers.utils import export_to_video
from huggingface_hub import snapshot_download
from transformers import AutoProcessor, Qwen3VLForConditionalGeneration

from lingbot_video import (
    FlowUniPCMultistepScheduler,
    LingBotVideoImageToVideoPipeline,
    LingBotVideoPipeline,
    LingBotVideoTransformer3DModel,
)
from lingbot_video.pipeline_lingbot_video import (
    DEFAULT_NEGATIVE_PROMPT,
    DEFAULT_NEGATIVE_PROMPT_IMAGE,
)
from lingbot_video.utils import num_frames_from_duration

from rewriter_prompts import (
    IMAGE_STEP1_EXPAND,
    IMAGE_STEP2_MAP,
    VIDEO_STEP1_EXPAND,
    VIDEO_STEP2_MAP,
)

torch.backends.cuda.matmul.allow_tf32 = True
torch.set_float32_matmul_precision("high")

MODEL_ID = "robbyant/lingbot-video-dense-1.3b"
FPS = 24
MAX_SEED = 2**31 - 1
MAX_GPU_SECONDS = 180
REWRITER_MODEL = os.environ.get("REWRITER_MODEL", "Qwen/Qwen3.6-27B:deepinfra")

# (height, width), multiples of 16 — official 480p buckets.
VIDEO_SIZES = {
    "832 × 480 (16:9)": (480, 832),
    "480 × 832 (9:16)": (832, 480),
    "640 × 480 (4:3)": (480, 640),
    "480 × 480 (1:1)": (480, 480),
}
IMAGE_SIZES = {
    "832 × 480 (16:9)": (480, 832),
    "480 × 832 (9:16)": (832, 480),
    "480 × 480 (1:1)": (480, 480),
    "1280 × 736 (16:9, 720p)": (736, 1280),
    "736 × 1280 (9:16, 720p)": (1280, 736),
    "1088 × 1088 (1:1, 1080p)": (1088, 1088),
}

EXAMPLES_DIR = Path(__file__).parent / "examples"


def _read_example(name: str) -> str:
    return (EXAMPLES_DIR / name).read_text(encoding="utf-8").strip()


print(f"[startup] downloading {MODEL_ID} ...", flush=True)
t0 = time.perf_counter()
model_dir = snapshot_download(
    MODEL_ID,
    ignore_patterns=["*.bak_lingbot_video_diffusers"],
)
print(f"[startup] snapshot ready in {time.perf_counter() - t0:.1f}s", flush=True)

t0 = time.perf_counter()
transformer = LingBotVideoTransformer3DModel.from_pretrained(
    model_dir, subfolder="transformer", torch_dtype=torch.bfloat16
)
text_encoder = Qwen3VLForConditionalGeneration.from_pretrained(
    model_dir, subfolder="text_encoder", dtype=torch.bfloat16, attn_implementation="sdpa"
)
processor = AutoProcessor.from_pretrained(model_dir, subfolder="processor")
vae = AutoencoderKLWan.from_pretrained(model_dir, subfolder="vae", torch_dtype=torch.float32)
print(f"[startup] components loaded in {time.perf_counter() - t0:.1f}s", flush=True)

t0 = time.perf_counter()
pipe = LingBotVideoPipeline(
    transformer=transformer,
    vae=vae,
    text_encoder=text_encoder,
    processor=processor,
    scheduler=FlowUniPCMultistepScheduler.from_pretrained(model_dir, subfolder="scheduler"),
).to("cuda")
pipe_i2v = LingBotVideoImageToVideoPipeline(
    transformer=transformer,
    vae=vae,
    text_encoder=text_encoder,
    processor=processor,
    scheduler=FlowUniPCMultistepScheduler.from_pretrained(model_dir, subfolder="scheduler"),
)
print(f"[startup] pipelines on cuda in {time.perf_counter() - t0:.1f}s", flush=True)


def _resolve_seed(seed: float, randomize: bool) -> int:
    if randomize:
        return random.randint(0, MAX_SEED)
    return int(seed)


# ---------------------------------------------------------------------------
# Prompt enhancement — the official LingBot rewriter recipe (expand -> JSON map)
# run on the official rewriter base model (Qwen3.6-27B) via HF Inference
# Providers. The DiT was trained on these structured JSON captions; plain
# prompts are far out of distribution and produce severe artifacts.
# ---------------------------------------------------------------------------


def _router_chat(text: str, max_tokens: int, image: Image.Image | None = None) -> str:
    token = os.environ.get("HF_TOKEN")
    if not token:
        raise RuntimeError("HF_TOKEN is not configured for prompt enhancement.")
    if image is not None:
        image = image.convert("RGB")
        image.thumbnail((768, 768))
        buf = io.BytesIO()
        image.save(buf, format="JPEG", quality=90)
        data_uri = "data:image/jpeg;base64," + base64.b64encode(buf.getvalue()).decode()
        content = [
            {"type": "image_url", "image_url": {"url": data_uri}},
            {"type": "text", "text": text},
        ]
    else:
        content = text
    body = {
        "model": REWRITER_MODEL,
        "messages": [{"role": "user", "content": content}],
        "max_tokens": max_tokens,
        "temperature": 0.0,
        "chat_template_kwargs": {"enable_thinking": False},
    }
    last_error = None
    for attempt in range(2):
        try:
            req = urllib.request.Request(
                "https://router.huggingface.co/v1/chat/completions",
                data=json.dumps(body).encode(),
                headers={"Authorization": f"Bearer {token}", "Content-Type": "application/json"},
            )
            with urllib.request.urlopen(req, timeout=180) as resp:
                data = json.load(resp)
            out = (data["choices"][0]["message"].get("content") or "").strip()
            if out:
                return out
            last_error = RuntimeError("empty rewriter response")
        except Exception as exc:  # noqa: BLE001
            last_error = exc
            if image is not None:
                # provider may not accept image input — retry text-only
                body["messages"][0]["content"] = text
                image = None
    raise RuntimeError(f"prompt enhancement failed: {last_error}")


def _extract_caption_json(raw: str) -> dict:
    match = re.search(r"```(?:json)?\s*(\{.*\})\s*```", raw, re.DOTALL)
    s = match.group(1) if match else raw
    start = s.find("{")
    if start < 0:
        raise ValueError("no JSON object in rewriter output")
    s = s[start:]
    try:
        return json.loads(s)
    except json.JSONDecodeError:
        from json_repair import repair_json

        obj = repair_json(s, return_objects=True)
        if not isinstance(obj, dict):
            raise ValueError("rewriter output is not a JSON object")
        return obj


def enhance_prompt(prompt: str, mode: str, duration_s: float | None, image: Image.Image | None = None) -> str:
    """Two-stage official rewrite: plain prompt -> detailed prose -> JSON caption string."""
    if mode == "t2i":
        step1 = IMAGE_STEP1_EXPAND + "\n\nUser image prompt:\n" + prompt
    else:
        step1 = VIDEO_STEP1_EXPAND + "\n\n" + prompt + f"\n\nVideo Duration: {duration_s:g} seconds"
    prose = _router_chat(step1, 1200, image=image)
    if mode == "t2i":
        step2 = IMAGE_STEP2_MAP + "\n\nDETAILED CAPTION:\n" + prose
    else:
        step2 = (
            VIDEO_STEP2_MAP
            + f"\n\nVideo Duration: {duration_s:g} seconds\n\nDETAILED CAPTION:\n"
            + prose
            + "\n\nOutput the JSON now."
        )
    caption = _extract_caption_json(_router_chat(step2, 6000, image=image))
    return json.dumps(caption, ensure_ascii=False, separators=(",", ":"))


def _prepare_caption(prompt: str, mode: str, duration_s: float | None, enhance: bool,
                     image: Image.Image | None = None) -> str:
    prompt = (prompt or "").strip()
    if not prompt:
        raise gr.Error("Please enter a prompt.")
    if prompt.startswith("{") or not enhance:
        return prompt  # already a structured JSON caption, or enhancement disabled
    t0 = time.perf_counter()
    try:
        caption = enhance_prompt(prompt, mode, duration_s, image=image)
    except Exception as exc:  # noqa: BLE001
        print(f"[rewriter] failed: {exc}", flush=True)
        raise gr.Error(
            "Prompt enhancement failed (the model needs structured captions to work well). "
            "Please try again, or paste a LingBot JSON caption and disable enhancement."
        )
    print(f"[rewriter] {mode} enhanced in {time.perf_counter() - t0:.1f}s ({len(caption)} chars)", flush=True)
    return caption


def _latent_tokens(height: int, width: int, num_frames: int) -> int:
    latent_frames = (num_frames - 1) // 4 + 1
    return latent_frames * (height // 16) * (width // 16)


def _step_seconds(height: int, width: int, num_frames: int) -> float:
    # Per-step cost (sequential CFG included), calibrated on the live Space:
    # 30 steps @ 832x480 f=49 took 92.4s → 3.08 s/step at r=0.619.
    r = _latent_tokens(height, width, num_frames) / 32760.0
    return 5.5 * r * r + 1.0 * r + 0.2


def _estimate_video_seconds(size_label: str, duration_s: float, steps: int, sizes=VIDEO_SIZES) -> int:
    height, width = sizes[size_label]
    num_frames = num_frames_from_duration(duration_s, FPS)
    return int(20.0 + steps * _step_seconds(height, width, num_frames))


def _estimate_image_seconds(size_label: str, steps: int) -> int:
    height, width = IMAGE_SIZES[size_label]
    return int(10.0 + steps * _step_seconds(height, width, 1))


def _check_budget(estimate: int) -> None:
    if estimate > MAX_GPU_SECONDS:
        raise gr.Error(
            f"These settings need ~{estimate}s of GPU time (max {MAX_GPU_SECONDS}s). "
            "Reduce the video duration, steps, or resolution."
        )


def _export_video(frames: np.ndarray) -> str:
    path = f"/tmp/lingbot_{int(time.time() * 1000)}.mp4"
    export_to_video(frames, path, fps=FPS)
    return path


def _video_gpu_duration(caption, height, width, num_frames, steps, *args, **kwargs) -> int:
    return min(MAX_GPU_SECONDS, int(20.0 + int(steps) * _step_seconds(height, width, num_frames)) + 20)


@spaces.GPU(duration=_video_gpu_duration)
def _gpu_generate_video(caption, height, width, num_frames, steps, guidance_scale, shift,
                        negative_prompt, seed, image=None):
    target = pipe_i2v if image is not None else pipe
    kwargs = {"image": image} if image is not None else {}
    # Sequential CFG. Batched CFG (batch_cfg=True) triggers the packed attention
    # path which requires flash_attn_varlen_func_v3 (FA3), and the FA3 binaries
    # only target sm_80/sm_90a — ZeroGPU's RTX PRO 6000 is sm_120, so the packed
    # path cannot run here.
    t0 = time.perf_counter()
    output = target(
        prompt=caption,
        negative_prompt=negative_prompt.strip() or DEFAULT_NEGATIVE_PROMPT,
        height=height,
        width=width,
        num_frames=num_frames,
        num_inference_steps=int(steps),
        guidance_scale=float(guidance_scale),
        shift=float(shift),
        batch_cfg=False,
        generator=torch.Generator().manual_seed(seed),
        **kwargs,
    )
    torch.cuda.synchronize()
    print(
        f"[timing] {'i2v' if image is not None else 't2v'} {width}x{height} f={num_frames} "
        f"steps={int(steps)} tokens={_latent_tokens(height, width, num_frames)} "
        f"took {time.perf_counter() - t0:.1f}s",
        flush=True,
    )
    return _export_video(output.frames[0])


def _image_gpu_duration(caption, height, width, steps, *args, **kwargs) -> int:
    return min(MAX_GPU_SECONDS, int(10.0 + int(steps) * _step_seconds(height, width, 1)) + 15)


@spaces.GPU(duration=_image_gpu_duration)
def _gpu_generate_image(caption, height, width, steps, guidance_scale, shift, negative_prompt, seed):
    t0 = time.perf_counter()
    output = pipe(
        prompt=caption,
        negative_prompt=negative_prompt.strip() or DEFAULT_NEGATIVE_PROMPT_IMAGE,
        height=height,
        width=width,
        num_frames=1,
        num_inference_steps=int(steps),
        guidance_scale=float(guidance_scale),
        shift=float(shift),
        batch_cfg=False,
        generator=torch.Generator().manual_seed(seed),
    )
    torch.cuda.synchronize()
    print(f"[timing] t2i {width}x{height} steps={int(steps)} "
          f"took {time.perf_counter() - t0:.1f}s", flush=True)
    frame = output.frames[0][0]
    return Image.fromarray((np.clip(frame, 0, 1) * 255).astype(np.uint8))


def generate_t2v(
    prompt: str,
    size_label: str = "832 × 480 (16:9)",
    duration_s: float = 2.0,
    steps: int = 30,
    guidance_scale: float = 3.0,
    shift: float = 3.0,
    negative_prompt: str = "",
    seed: int = 42,
    randomize_seed: bool = True,
    enhance: bool = True,
    progress=gr.Progress(track_tqdm=True),
):
    """Generate a short video from a text prompt with LingBot-Video Dense 1.3B.

    Args:
        prompt: Scene description in natural language (it is auto-expanded into the
            structured caption the model expects), or a raw LingBot JSON caption.
        size_label: Resolution/aspect preset, e.g. "832 × 480 (16:9)", "480 × 832 (9:16)".
        duration_s: Video length in seconds (1.0-5.0) at 24 fps.
        steps: Number of denoising steps (more = more detail, slower).
        guidance_scale: Classifier-free guidance strength.
        shift: Flow-matching timestep shift.
        negative_prompt: What to avoid; empty uses the model default.
        seed: Random seed for reproducibility.
        randomize_seed: If true, ignore seed and use a random one.
        enhance: If true, expand a plain prompt into a structured caption before generation.

    Returns:
        The generated MP4 video, the seed used, and the structured caption fed to the model.
    """
    estimate = _estimate_video_seconds(size_label, duration_s, int(steps))
    _check_budget(estimate)
    height, width = VIDEO_SIZES[size_label]
    num_frames = num_frames_from_duration(duration_s, FPS)
    seed = _resolve_seed(seed, randomize_seed)
    caption = _prepare_caption(prompt, "t2v", duration_s, enhance)
    video = _gpu_generate_video(
        caption, height, width, num_frames, steps, guidance_scale, shift, negative_prompt, seed,
    )
    return video, seed, caption


def generate_i2v(
    image,
    prompt: str,
    size_label: str = "832 × 480 (16:9)",
    duration_s: float = 2.0,
    steps: int = 30,
    guidance_scale: float = 3.0,
    shift: float = 3.0,
    negative_prompt: str = "",
    seed: int = 42,
    randomize_seed: bool = True,
    enhance: bool = True,
    progress=gr.Progress(track_tqdm=True),
):
    """Animate a first-frame image into a short video with LingBot-Video Dense 1.3B.

    Args:
        image: The first frame to animate (filepath or PIL image).
        prompt: How the scene should evolve, in natural language (auto-expanded), or a raw LingBot JSON caption.
        size_label: Resolution/aspect preset, e.g. "832 × 480 (16:9)".
        duration_s: Video length in seconds (1.0-5.0) at 24 fps.
        steps: Number of denoising steps.
        guidance_scale: Classifier-free guidance strength.
        shift: Flow-matching timestep shift.
        negative_prompt: What to avoid; empty uses the model default.
        seed: Random seed for reproducibility.
        randomize_seed: If true, ignore seed and use a random one.
        enhance: If true, expand a plain prompt into a structured caption before generation.

    Returns:
        The generated MP4 video, the seed used, and the structured caption fed to the model.
    """
    if image is None:
        raise gr.Error("Please upload a first-frame image.")
    estimate = _estimate_video_seconds(size_label, duration_s, int(steps)) + 15
    _check_budget(estimate)
    height, width = VIDEO_SIZES[size_label]
    num_frames = num_frames_from_duration(duration_s, FPS)
    seed = _resolve_seed(seed, randomize_seed)
    caption = _prepare_caption(prompt, "ti2v", duration_s, enhance, image=image)
    video = _gpu_generate_video(
        caption, height, width, num_frames, steps, guidance_scale, shift, negative_prompt, seed,
        image=image,
    )
    return video, seed, caption


def generate_t2i(
    prompt: str,
    size_label: str = "832 × 480 (16:9)",
    steps: int = 30,
    guidance_scale: float = 3.0,
    shift: float = 3.0,
    negative_prompt: str = "",
    seed: int = 42,
    randomize_seed: bool = True,
    enhance: bool = True,
    progress=gr.Progress(track_tqdm=True),
):
    """Generate an image from a text prompt with LingBot-Video Dense 1.3B.

    Args:
        prompt: Image description in natural language (auto-expanded), or a raw LingBot JSON caption.
        size_label: Resolution/aspect preset, e.g. "832 × 480 (16:9)", "1088 × 1088 (1:1, 1080p)".
        steps: Number of denoising steps.
        guidance_scale: Classifier-free guidance strength.
        shift: Flow-matching timestep shift.
        negative_prompt: What to avoid; empty uses the model default.
        seed: Random seed for reproducibility.
        randomize_seed: If true, ignore seed and use a random one.
        enhance: If true, expand a plain prompt into a structured caption before generation.

    Returns:
        The generated image, the seed used, and the structured caption fed to the model.
    """
    height, width = IMAGE_SIZES[size_label]
    seed = _resolve_seed(seed, randomize_seed)
    caption = _prepare_caption(prompt, "t2i", None, enhance)
    image = _gpu_generate_image(
        caption, height, width, steps, guidance_scale, shift, negative_prompt, seed,
    )
    return image, seed, caption


HEADER_HTML = """
<div id="lb-header">
  <div class="lb-top">
    <div class="lb-title">
      LingBot-Video
      <span class="lb-badge">Dense&nbsp;1.3B</span>
    </div>
    <nav class="lb-links">
      <a href="https://huggingface.co/robbyant/lingbot-video-dense-1.3b" target="_blank" rel="noopener">Model</a>
      <a href="https://github.com/Robbyant/lingbot-video" target="_blank" rel="noopener">GitHub</a>
      <a href="https://technology.robbyant.com/lingbot-video" target="_blank" rel="noopener">Project</a>
    </nav>
  </div>
  <p class="lb-sub">
    Embodied-intelligence video generation from a lightweight dense model — text-to-video, image-to-video &amp; text-to-image.
    Prompts are auto-expanded into structured JSON captions
    (<a href="https://huggingface.co/Qwen/Qwen3.6-27B" target="_blank" rel="noopener">Qwen3.6-27B</a>);
    paste a raw LingBot caption to skip that. 480p video, ZeroGPU.
  </p>
</div>
"""

CSS = """
#col-container { max-width: 1100px; margin: 0 auto; }
.dark .gradio-container { color: var(--body-text-color); }
/* Zero out Gradio's wrapper padding around the header HTML block. */
#lb-header-wrap, #lb-header-wrap .html-container {
  padding: 0 !important;
  border: none !important;
  background: transparent !important;
}
/* Borderless header, flush with the app container. */
#lb-header {
  margin: 0;
  padding: 2px 0 8px 0;
}
#lb-header .lb-top {
  display: flex;
  align-items: center;
  justify-content: flex-start;
  flex-wrap: wrap;
  gap: 6px 18px;
}
#lb-header .lb-title {
  font-size: 1.35rem;
  font-weight: 700;
  line-height: 1.2;
  color: var(--body-text-color);
  display: flex;
  align-items: center;
  gap: 10px;
}
#lb-header .lb-badge {
  font-size: 0.66rem;
  font-weight: 600;
  letter-spacing: 0.02em;
  text-transform: uppercase;
  color: var(--body-text-color-subdued);
  background: var(--background-fill-primary);
  border: 1px solid var(--border-color-primary);
  border-radius: 999px;
  padding: 2px 9px;
  white-space: nowrap;
}
#lb-header .lb-links {
  display: flex;
  gap: 14px;
  font-size: 0.85rem;
  padding-left: 4px;
}
#lb-header .lb-links a {
  color: var(--body-text-color);
  opacity: 0.85;
  text-decoration: none;
  border-bottom: 1px solid transparent;
}
#lb-header .lb-links a:hover {
  color: var(--link-text-color);
  border-bottom-color: currentColor;
}
#lb-header .lb-sub {
  margin: 7px 0 0 0;
  font-size: 0.82rem;
  line-height: 1.45;
  color: var(--body-text-color);
  opacity: 0.78;
}
#lb-header .lb-sub a { color: var(--link-text-color); text-decoration: none; }
#lb-header .lb-sub a:hover { text-decoration: underline; }
/* Gradio's theme pads <a> with 0 8px, which shows up as fake space inside the
   "(Qwen3.6-27B)" parentheses — strip it on all header links. */
#lb-header a { padding: 0 !important; }

/* Compact, tidy examples: no wide horizontal overflow, subtle scrollbar */
.lb-examples .gr-samples-table td, .lb-examples table td { white-space: normal; }
.lb-examples { scrollbar-width: thin; }
.lb-examples ::-webkit-scrollbar { height: 8px; width: 8px; }
.lb-examples ::-webkit-scrollbar-thumb {
  background: var(--border-color-primary);
  border-radius: 999px;
}
.lb-examples ::-webkit-scrollbar-track { background: transparent; }
"""


def _seed_row():
    with gr.Row():
        seed = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=42)
        randomize = gr.Checkbox(label="Randomize seed", value=True)
    return seed, randomize


def _advanced(negative_default: str, steps_value: int):
    with gr.Accordion("Advanced settings", open=False):
        steps = gr.Slider(label="Inference steps", minimum=4, maximum=50, step=1, value=steps_value)
        guidance = gr.Slider(label="Guidance scale", minimum=1.0, maximum=10.0, step=0.5, value=3.0)
        shift = gr.Slider(label="Flow shift", minimum=1.0, maximum=9.0, step=0.5, value=3.0)
        negative = gr.Textbox(label="Negative prompt", value=negative_default, lines=3)
        seed, randomize = _seed_row()
    return steps, guidance, shift, negative, seed, randomize


# Example runners: examples carry only the natural-language prompt (and first
# frame for I2V) — enhancement is on and every other control keeps its default,
# so the examples table stays to one or two clean columns instead of ten.
def _example_t2v(prompt):
    return generate_t2v(prompt, "832 × 480 (16:9)", 2.0, 30, 3.0, 3.0,
                        DEFAULT_NEGATIVE_PROMPT, 42, False, True)


def _example_i2v(image, prompt):
    return generate_i2v(image, prompt, "832 × 480 (16:9)", 2.0, 30, 3.0, 3.0,
                        DEFAULT_NEGATIVE_PROMPT, 42, False, True)


def _example_t2i(prompt):
    return generate_t2i(prompt, "832 × 480 (16:9)", 30, 3.0, 3.0,
                        DEFAULT_NEGATIVE_PROMPT_IMAGE, 42, False, True)


# Pre-built example prompts derived from the official LingBot-Video example cases.
# For t2v/t2i, we use a short natural-language prompt that the rewriter expands.
# For i2v, we pass the first-frame image and a natural-language motion description.
T2V_EXAMPLES = [
    "A young woman with long, wavy brown hair stands in a bright modern apartment living room, "
    "wearing an oversized cream knit cardigan over a white tank top with beige trousers. She "
    "smiles at the camera, shifts her weight, and adjusts her collar, showcasing her outfit.",
    "A young child with short brown hair plays outdoors on a sunny day, blowing shimmering soap "
    "bubbles with a wand. The bubbles float upwards, catching the warm afternoon light against a "
    "soft-focus green background.",
    "A robotic arm on a workbench reaches forward, grasps a black game controller, lifts it, "
    "moves it to the right, and lowers it into an open box. Another robotic arm and headphones "
    "remain stationary on the desk, top-down view.",
]
T2V_LABELS = ["👗 Woman in apartment", "🫧 Child blowing bubbles", "🤖 Robot arm sorting"]

I2V_EXAMPLES = [
    [str(EXAMPLES_DIR / "ti2v_frame.png"),
     "A fit young man and a sleek white humanoid robot run side by side along a cherry blossom "
     "lined promenade toward the camera, which tracks backward. Energetic, futuristic, daylight."],
]
I2V_LABELS = ["🏃 Man and robot running"]

T2I_EXAMPLES = [
    "A clear glass bottle of water on a sunlit wooden table acts as a lens, refracting bright "
    "sunlight into a warm glow, extreme close-up, photorealistic.",
    "A humanoid robot chef flipping a pancake in a bright modern kitchen, dramatic side lighting, "
    "steam rising from the pan, shallow depth of field, 85mm lens.",
]
T2I_LABELS = ["🔆 Bottle as a lens", "🤖 Robot chef"]


with gr.Blocks(title="LingBot-Video Dense 1.3B") as demo:
    gr.HTML(HEADER_HTML, elem_id="lb-header-wrap")

    with gr.Tab("Text → Video"):
        with gr.Row():
            with gr.Column():
                t2v_prompt = gr.Textbox(
                    label="Prompt",
                    placeholder="Describe the scene in detail (plain text or a LingBot JSON caption)...",
                    lines=5,
                    max_lines=12,
                )
                with gr.Row():
                    t2v_size = gr.Dropdown(
                        label="Resolution", choices=list(VIDEO_SIZES), value="832 × 480 (16:9)"
                    )
                    t2v_dur = gr.Slider(
                        label="Video duration (s)", minimum=1.0, maximum=5.0, step=0.5, value=2.0
                    )
                t2v_enhance = gr.Checkbox(
                    label="Enhance prompt (official rewriter — required for plain prompts)", value=True
                )
                t2v_steps, t2v_guidance, t2v_shift, t2v_negative, t2v_seed, t2v_rand = _advanced(
                    DEFAULT_NEGATIVE_PROMPT, 30
                )
                t2v_btn = gr.Button("Generate video", variant="primary")
            with gr.Column():
                t2v_out = gr.Video(label="Generated video", autoplay=True)
                t2v_seed_out = gr.Number(label="Seed used", interactive=False)
                with gr.Accordion("Structured caption used", open=False):
                    t2v_caption_out = gr.Textbox(label="Caption", lines=4)

        t2v_inputs = [
            t2v_prompt, t2v_size, t2v_dur, t2v_steps, t2v_guidance, t2v_shift,
            t2v_negative, t2v_seed, t2v_rand, t2v_enhance,
        ]
        t2v_outputs = [t2v_out, t2v_seed_out, t2v_caption_out]
        t2v_btn.click(generate_t2v, inputs=t2v_inputs, outputs=t2v_outputs,
                      concurrency_id="gpu", concurrency_limit=1)
        with gr.Column(elem_classes="lb-examples"):
            gr.Examples(
                examples=T2V_EXAMPLES,
                example_labels=T2V_LABELS,
                fn=_example_t2v,
                inputs=[t2v_prompt],
                outputs=t2v_outputs,
                cache_examples=True,
                cache_mode="lazy",
            )

    with gr.Tab("Image → Video"):
        with gr.Row():
            with gr.Column():
                i2v_image = gr.Image(label="First frame", type="pil")
                i2v_prompt = gr.Textbox(
                    label="Prompt",
                    placeholder="Describe how the scene should evolve...",
                    lines=4,
                    max_lines=12,
                )
                with gr.Row():
                    i2v_size = gr.Dropdown(
                        label="Resolution", choices=list(VIDEO_SIZES), value="832 × 480 (16:9)"
                    )
                    i2v_dur = gr.Slider(
                        label="Video duration (s)", minimum=1.0, maximum=5.0, step=0.5, value=2.0
                    )
                i2v_enhance = gr.Checkbox(
                    label="Enhance prompt (official rewriter — required for plain prompts)", value=True
                )
                i2v_steps, i2v_guidance, i2v_shift, i2v_negative, i2v_seed, i2v_rand = _advanced(
                    DEFAULT_NEGATIVE_PROMPT, 30
                )
                i2v_btn = gr.Button("Generate video", variant="primary")
            with gr.Column():
                i2v_out = gr.Video(label="Generated video", autoplay=True)
                i2v_seed_out = gr.Number(label="Seed used", interactive=False)
                with gr.Accordion("Structured caption used", open=False):
                    i2v_caption_out = gr.Textbox(label="Caption", lines=4)

        i2v_inputs = [
            i2v_image, i2v_prompt, i2v_size, i2v_dur, i2v_steps, i2v_guidance, i2v_shift,
            i2v_negative, i2v_seed, i2v_rand, i2v_enhance,
        ]
        i2v_outputs = [i2v_out, i2v_seed_out, i2v_caption_out]
        i2v_btn.click(generate_i2v, inputs=i2v_inputs, outputs=i2v_outputs,
                      concurrency_id="gpu", concurrency_limit=1)
        with gr.Column(elem_classes="lb-examples"):
            gr.Examples(
                examples=I2V_EXAMPLES,
                example_labels=I2V_LABELS,
                fn=_example_i2v,
                inputs=[i2v_image, i2v_prompt],
                outputs=i2v_outputs,
                cache_examples=True,
                cache_mode="lazy",
            )

    with gr.Tab("Text → Image"):
        with gr.Row():
            with gr.Column():
                t2i_prompt = gr.Textbox(
                    label="Prompt",
                    placeholder="Describe the image in detail...",
                    lines=5,
                    max_lines=12,
                )
                t2i_size = gr.Dropdown(
                    label="Resolution", choices=list(IMAGE_SIZES), value="832 × 480 (16:9)"
                )
                t2i_enhance = gr.Checkbox(
                    label="Enhance prompt (official rewriter — required for plain prompts)", value=True
                )
                t2i_steps, t2i_guidance, t2i_shift, t2i_negative, t2i_seed, t2i_rand = _advanced(
                    DEFAULT_NEGATIVE_PROMPT_IMAGE, 30
                )
                t2i_btn = gr.Button("Generate image", variant="primary")
            with gr.Column():
                t2i_out = gr.Image(label="Generated image")
                t2i_seed_out = gr.Number(label="Seed used", interactive=False)
                with gr.Accordion("Structured caption used", open=False):
                    t2i_caption_out = gr.Textbox(label="Caption", lines=4)

        t2i_inputs = [
            t2i_prompt, t2i_size, t2i_steps, t2i_guidance, t2i_shift,
            t2i_negative, t2i_seed, t2i_rand, t2i_enhance,
        ]
        t2i_outputs = [t2i_out, t2i_seed_out, t2i_caption_out]
        t2i_btn.click(generate_t2i, inputs=t2i_inputs, outputs=t2i_outputs,
                      concurrency_id="gpu", concurrency_limit=1)
        with gr.Column(elem_classes="lb-examples"):
            gr.Examples(
                examples=T2I_EXAMPLES,
                example_labels=T2I_LABELS,
                fn=_example_t2i,
                inputs=[t2i_prompt],
                outputs=t2i_outputs,
                cache_examples=True,
                cache_mode="lazy",
            )

demo.queue(max_size=30).launch(theme=gr.themes.Citrus(), css=CSS, mcp_server=True)