app.py

# ---------------------------------------------------------------------------
# Krea Realtime Video 14B — Hugging Face Space Demo
# ZeroGPU compatibility version for Diffusers ModularPipeline.
# ---------------------------------------------------------------------------

import os

# ---------------------------------------------------------------------------
# HF Spaces / cache configuration — must happen before HF imports
# ---------------------------------------------------------------------------

_ASF_HF_CACHE_ROOT = os.environ.get("ASF_HF_CACHE_DIR") or "/tmp/asf-hf-cache"

os.environ.setdefault("HF_HOME", _ASF_HF_CACHE_ROOT)
os.environ.setdefault("HF_HUB_CACHE", os.path.join(_ASF_HF_CACHE_ROOT, "hub"))
os.environ.setdefault("HUGGINGFACE_HUB_CACHE", os.path.join(_ASF_HF_CACHE_ROOT, "hub"))
os.environ.setdefault("TRANSFORMERS_CACHE", os.path.join(_ASF_HF_CACHE_ROOT, "transformers"))
os.environ.setdefault("DIFFUSERS_CACHE", os.path.join(_ASF_HF_CACHE_ROOT, "diffusers"))
os.environ.setdefault("HF_MODULES_CACHE", "/tmp/hf_modules")
os.environ.setdefault("MPLCONFIGDIR", "/tmp/matplotlib")

# ZeroGPU compatibility mode:
# - torch.compile is disabled below.
# - hub kernels / torchao optimized path is intentionally not used.
os.environ.setdefault("DIFFUSERS_ENABLE_HUB_KERNELS", "0")
os.environ.setdefault("USE_HUB_KERNELS", "NO")

os.makedirs(_ASF_HF_CACHE_ROOT, exist_ok=True)
os.makedirs(os.path.join(_ASF_HF_CACHE_ROOT, "hub"), exist_ok=True)
os.makedirs(os.path.join(_ASF_HF_CACHE_ROOT, "transformers"), exist_ok=True)
os.makedirs(os.path.join(_ASF_HF_CACHE_ROOT, "diffusers"), exist_ok=True)
os.makedirs(os.environ["HF_MODULES_CACHE"], exist_ok=True)
os.makedirs(os.environ["MPLCONFIGDIR"], exist_ok=True)

HF_TOKEN = os.environ.get("HF_TOKEN") or os.environ.get("HUGGING_FACE_HUB_TOKEN")

# ---------------------------------------------------------------------------
# Safe spaces import
# ---------------------------------------------------------------------------

try:
    import spaces

    HAS_SPACES = True
except Exception:
    HAS_SPACES = False

    class _DummySpaces:
        def GPU(self, *args, **kwargs):
            def decorator(fn):
                return fn

            return decorator

    spaces = _DummySpaces()


def _spaces_gpu(*args, **kwargs):
    """
    Wrapper around spaces.GPU.

    Some versions of the spaces package may not support size=...
    In that case, we fall back to the same decorator without size.
    """
    try:
        return spaces.GPU(*args, **kwargs)
    except TypeError:
        kwargs.pop("size", None)
        return spaces.GPU(*args, **kwargs)


# ---------------------------------------------------------------------------
# Imports + ZeroGPU torch.compile bypass
# ---------------------------------------------------------------------------

import sys
import re
import time
import threading
import traceback
import importlib.util
import importlib.metadata

import torch

_ORIG_TORCH_COMPILE = getattr(torch, "compile", None)
_ASF_COMPILE_BYPASS_ANNOUNCED = False


def _asf_zerogpu_compile_bypass(fn=None, *args, **kwargs):
    """
    ZeroGPU compatibility shim.

    Supports both call styles:
      torch.compile(fn, ...)
      @torch.compile(...)
      def fn(...): ...

    Default behavior is quiet because FlexAttention may call this repeatedly
    during generation.
    """
    global _ASF_COMPILE_BYPASS_ANNOUNCED

    if fn is None:
        def decorator(real_fn):
            return _asf_zerogpu_compile_bypass(real_fn, *args, **kwargs)

        return decorator

    if os.environ.get("ASF_VERBOSE_COMPILE_BYPASS", "0") == "1":
        name = getattr(fn, "__name__", repr(fn))
        module = getattr(fn, "__module__", "")
        print(
            f"[ASF] ZeroGPU compatibility: bypassing torch.compile for {module}.{name}",
            flush=True,
        )
    elif not _ASF_COMPILE_BYPASS_ANNOUNCED:
        print(
            "[ASF] ZeroGPU compatibility: torch.compile bypass is active.",
            flush=True,
        )
        _ASF_COMPILE_BYPASS_ANNOUNCED = True

    return fn


if _ORIG_TORCH_COMPILE is not None and os.environ.get("ASF_ENABLE_TORCH_COMPILE", "0") != "1":
    torch.compile = _asf_zerogpu_compile_bypass

try:
    import torch._dynamo

    torch._dynamo.config.suppress_errors = True
except Exception:
    pass


def _installed_version(package_name):
    try:
        return importlib.metadata.version(package_name)
    except Exception:
        return None


def _version_tuple(version):
    """
    Minimal semver-ish parser.
    Handles strings like 0.12.0, 0.16.0, 0.16.0.dev...
    """
    if not version:
        return None
    parts = re.findall(r"\d+", version)
    if not parts:
        return None
    return tuple(int(p) for p in parts[:3])


PEFT_AVAILABLE = importlib.util.find_spec("peft") is not None
PEFT_VERSION = _installed_version("peft")

TORCHAO_VERSION = _installed_version("torchao")
TORCHAO_INSTALLED = TORCHAO_VERSION is not None
TORCHAO_VERSION_TUPLE = _version_tuple(TORCHAO_VERSION)

# Recent PEFT rejects torchao < 0.16.0 when torchao is installed.
# In this Space, torchao is not needed because we are not using Krea's optimized
# torch.compile / FP8 path on ZeroGPU.
TORCHAO_COMPATIBLE_FOR_PEFT = (
    not TORCHAO_INSTALLED
    or (
        TORCHAO_VERSION_TUPLE is not None
        and TORCHAO_VERSION_TUPLE >= (0, 16, 0)
    )
)

LORA_BACKEND_READY = PEFT_AVAILABLE and TORCHAO_COMPATIBLE_FOR_PEFT

if not PEFT_AVAILABLE:
    LORA_BACKEND_ERROR = "PEFT is not installed. Add `peft` to requirements.txt and rebuild."
elif not TORCHAO_COMPATIBLE_FOR_PEFT:
    LORA_BACKEND_ERROR = (
        f"Incompatible torchao version detected: {TORCHAO_VERSION}. "
        "Remove `torchao==0.12.0` from requirements.txt, or upgrade torchao to >=0.16.0. "
        "For this ZeroGPU compatibility Space, removing torchao is recommended."
    )
else:
    LORA_BACKEND_ERROR = ""


import gradio as gr

# ---------------------------------------------------------------------------
# Diffusers imports
# ---------------------------------------------------------------------------

_DIFFUSERS_OK = False
_DIFFUSERS_IMPORT_ERROR = None

try:
    from diffusers import ModularPipeline
    from diffusers.modular_pipelines import PipelineState
    from diffusers.utils import export_to_video

    _DIFFUSERS_OK = True
except Exception as e:
    _DIFFUSERS_IMPORT_ERROR = f"{type(e).__name__}: {e}"
    traceback.print_exc()


# ---------------------------------------------------------------------------
# Model / LoRA configuration
# ---------------------------------------------------------------------------

MODEL_ID = "krea/krea-realtime-video"

KNOWN_LORAS = {
    "Base model": None,
    "Origami": {
        "repo_id": "shauray/Origami_WanLora",
        "prefix": "diffusion_model",
        "weight_name": "origami_000000500.safetensors",
        "adapter_name": "origami",
        "trigger": "[origami]",
    },
}

_pipeline = None
_pipeline_error = None
_pipeline_lock = threading.Lock()

_loaded_loras = set()
_active_lora = None
_active_lora_label = "Base model"
_lora_lock = threading.Lock()


def _log(msg):
    print(f"[KreaRealtimeVideo] {msg}", flush=True)


def _runtime_report():
    return {
        "python": sys.version.replace("\n", " "),
        "torch": getattr(torch, "__version__", "unknown"),
        "cuda_available": bool(torch.cuda.is_available()),
        "cuda_device_count": int(torch.cuda.device_count()) if torch.cuda.is_available() else 0,
        "has_spaces": HAS_SPACES,
        "torch_compile_bypassed": torch.compile is _asf_zerogpu_compile_bypass,
        "peft_available": PEFT_AVAILABLE,
        "peft_version": PEFT_VERSION,
        "torchao_installed": TORCHAO_INSTALLED,
        "torchao_version": TORCHAO_VERSION,
        "lora_backend_ready": LORA_BACKEND_READY,
        "lora_backend_error": LORA_BACKEND_ERROR,
        "hf_home": os.environ.get("HF_HOME", ""),
        "hf_modules_cache": os.environ.get("HF_MODULES_CACHE", ""),
    }


def _lora_report():
    return {
        "active_lora": _active_lora_label,
        "active_adapter": _active_lora,
        "loaded_loras": sorted(list(_loaded_loras)),
        "available_loras": list(KNOWN_LORAS.keys()),
        "backend_ready": LORA_BACKEND_READY,
        "backend_error": LORA_BACKEND_ERROR,
    }


def _call_from_pretrained_compat(*args, **kwargs):
    """
    Compatibility wrapper because some diffusers/HF Hub combinations
    may use token= while older ones expect use_auth_token= or no token.
    """
    try:
        return ModularPipeline.from_pretrained(*args, **kwargs)
    except TypeError as e:
        if "token" in str(e):
            kwargs.pop("token", None)
            if HF_TOKEN:
                kwargs["use_auth_token"] = HF_TOKEN
            return ModularPipeline.from_pretrained(*args, **kwargs)
        raise


def _load_components_compat(pipe, **kwargs):
    """
    Compatibility wrapper around pipe.load_components().
    """
    try:
        return pipe.load_components(**kwargs)
    except TypeError as e:
        msg = str(e)
        if "token" in msg:
            kwargs.pop("token", None)
            if HF_TOKEN:
                kwargs["use_auth_token"] = HF_TOKEN
            return pipe.load_components(**kwargs)
        raise


def _load_pipeline():
    """
    Load the ModularPipeline once at app startup.

    For ZeroGPU, this gives the best UX:
    - model warms up when the app starts;
    - generation remains protected by @spaces.GPU;
    - LoRAs can be loaded manually before generation.
    """
    global _pipeline, _pipeline_error

    with _pipeline_lock:
        if _pipeline is not None:
            return _pipeline

        if not _DIFFUSERS_OK:
            _pipeline_error = _DIFFUSERS_IMPORT_ERROR or "Diffusers import failed"
            _log(f"Pipeline load skipped: {_pipeline_error}")
            return None

        try:
            _log(f"Runtime report: {_runtime_report()}")
            _log(f"Loading ModularPipeline from {MODEL_ID} ...")

            pipe = _call_from_pretrained_compat(
                MODEL_ID,
                trust_remote_code=True,
                token=HF_TOKEN,
            )

            _log("Skeleton loaded; attaching components ...")

            try:
                _load_components_compat(
                    pipe,
                    trust_remote_code=True,
                    device_map="cuda",
                    torch_dtype={
                        "default": torch.bfloat16,
                        "vae": torch.float16,
                    },
                    token=HF_TOKEN,
                )
            except RuntimeError as err:
                msg = str(err)
                cuda_load_failed = (
                    "Found no NVIDIA driver" in msg
                    or "No CUDA GPUs are available" in msg
                    or "libcudart" in msg
                    or "CUDA error" in msg
                    or "CUDA driver" in msg
                )

                if cuda_load_failed:
                    _log(
                        "device_map='cuda' failed during startup. "
                        "Retrying CPU-load + manual .to('cuda') ..."
                    )
                    _log(f"CUDA load error was: {msg}")

                    _load_components_compat(
                        pipe,
                        trust_remote_code=True,
                        torch_dtype={
                            "default": torch.bfloat16,
                            "vae": torch.float16,
                        },
                        token=HF_TOKEN,
                    )

                    pipe = pipe.to("cuda")
                else:
                    raise

            # Krea model-card optimization: fuse projections.
            # This is safe; it is not torch.compile.
            try:
                if hasattr(pipe, "transformer") and hasattr(pipe.transformer, "blocks"):
                    fused = 0
                    for block in pipe.transformer.blocks:
                        self_attn = getattr(block, "self_attn", None)
                        if self_attn is not None and hasattr(self_attn, "fuse_projections"):
                            self_attn.fuse_projections()
                            fused += 1
                    _log(f"Fused attention projections on {fused} blocks.")
            except Exception as e:
                _log(f"fuse_projections warning: {type(e).__name__}: {e}")

            _pipeline = pipe
            _pipeline_error = None
            _log("Pipeline ready.")
            return _pipeline

        except Exception as e:
            _pipeline_error = f"{type(e).__name__}: {e}"
            _log(f"Pipeline load FAILED: {_pipeline_error}")
            traceback.print_exc()
            return None


# ---------------------------------------------------------------------------
# LoRA helpers
# ---------------------------------------------------------------------------

def _load_lora_if_needed(pipe, lora_label):
    """
    Load a known LoRA adapter once.

    This is intentionally not decorated with @spaces.GPU when called through
    the UI load button, so it does not reserve ZeroGPU generation time.
    """
    global _loaded_loras

    cfg = KNOWN_LORAS.get(lora_label)
    if not cfg:
        return None

    if not LORA_BACKEND_READY:
        raise RuntimeError(LORA_BACKEND_ERROR)

    adapter_name = cfg["adapter_name"]

    if adapter_name in _loaded_loras:
        return adapter_name

    transformer = getattr(pipe, "transformer", None)
    if transformer is None or not hasattr(transformer, "load_lora_adapter"):
        raise RuntimeError("This pipeline transformer does not expose load_lora_adapter().")

    _log(f"Loading LoRA adapter: {lora_label} ({adapter_name})")

    transformer.load_lora_adapter(
        cfg["repo_id"],
        prefix=cfg["prefix"],
        weight_name=cfg["weight_name"],
        adapter_name=adapter_name,
    )

    _loaded_loras.add(adapter_name)
    _log(f"LoRA loaded: {adapter_name}")

    return adapter_name


def _safe_disable_lora(transformer):
    """
    Disable PEFT LoRA if available.

    Some diffusers methods raise if PEFT is not installed or incompatible, so
    this is defensive.
    """
    if transformer is None:
        return

    if hasattr(transformer, "disable_lora"):
        try:
            transformer.disable_lora()
            return
        except Exception as e:
            _log(f"disable_lora warning: {type(e).__name__}: {e}")

    if hasattr(transformer, "set_adapters"):
        try:
            transformer.set_adapters([])
            return
        except Exception as e:
            _log(f"set_adapters([]) warning: {type(e).__name__}: {e}")


def _activate_lora_adapter(transformer, adapter_name):
    """
    Activate a LoRA adapter across several possible Diffusers/PEFT APIs.

    This intentionally avoids adapter_weights because this Space's PEFT
    PeftAdapterMixin.set_adapters() does not accept that keyword.
    """
    if hasattr(transformer, "enable_lora"):
        try:
            transformer.enable_lora()
        except Exception as e:
            _log(f"enable_lora warning: {type(e).__name__}: {e}")

    if hasattr(transformer, "set_adapters"):
        # Diffusers-style API often accepts a list.
        try:
            transformer.set_adapters([adapter_name])
            return
        except TypeError as e:
            _log(f"set_adapters([adapter_name]) TypeError: {e}")
        except Exception as e:
            _log(f"set_adapters([adapter_name]) warning: {type(e).__name__}: {e}")

        # PEFT-style API may accept a string.
        try:
            transformer.set_adapters(adapter_name)
            return
        except TypeError as e:
            _log(f"set_adapters(adapter_name) TypeError: {e}")
        except Exception as e:
            _log(f"set_adapters(adapter_name) warning: {type(e).__name__}: {e}")

    if hasattr(transformer, "set_adapter"):
        transformer.set_adapter(adapter_name)
        return

    raise RuntimeError("Transformer does not expose set_adapters() or set_adapter().")


def _set_lora(pipe, lora_label, allow_load=True):
    """
    Activate the selected LoRA, or disable LoRA for base model.

    If allow_load=False, this function will not download/load a missing adapter.
    This keeps generate() fast and avoids hidden loading inside @spaces.GPU.
    """
    global _active_lora, _active_lora_label

    transformer = getattr(pipe, "transformer", None)
    if transformer is None:
        raise RuntimeError("Pipeline has no transformer.")

    cfg = KNOWN_LORAS.get(lora_label)

    if not cfg:
        _safe_disable_lora(transformer)

        _active_lora = None
        _active_lora_label = "Base model"
        return ""

    adapter_name = cfg["adapter_name"]

    if adapter_name not in _loaded_loras:
        if not allow_load:
            raise RuntimeError(
                f"LoRA '{lora_label}' is selected but not loaded. "
                "Click 'Load style' before generating."
            )
        adapter_name = _load_lora_if_needed(pipe, lora_label)

    _activate_lora_adapter(transformer, adapter_name)

    _active_lora = adapter_name
    _active_lora_label = lora_label

    return cfg.get("trigger", "").strip()


def load_selected_lora(lora_style):
    """
    Manual LoRA loading button.

    Not decorated with @spaces.GPU on purpose:
    loading the adapter should happen before generation and not consume
    the generation reservation window.
    """
    pipe = _load_pipeline()
    if pipe is None:
        return (
            "Model is not loaded.",
            {
                "status": "error",
                "error": _pipeline_error or "Pipeline failed to load",
                "lora": _lora_report(),
            },
        )

    if lora_style == "Base model":
        with _lora_lock:
            _set_lora(pipe, "Base model", allow_load=False)

        return (
            "Base model active.",
            {
                "status": "ready",
                "message": "Base model active. No LoRA selected.",
                "lora": _lora_report(),
            },
        )

    try:
        with _lora_lock:
            trigger = _set_lora(
                pipe,
                lora_style,
                allow_load=True,
            )

        return (
            f"{lora_style} loaded. Trigger `{trigger}` will be added automatically.",
            {
                "status": "ready",
                "message": f"LoRA loaded and activated: {lora_style}",
                "trigger": trigger,
                "lora": _lora_report(),
            },
        )

    except Exception as e:
        traceback.print_exc()
        return (
            f"LoRA load failed: {type(e).__name__}: {e}",
            {
                "status": "error",
                "message": "LoRA load failed.",
                "error": f"{type(e).__name__}: {e}",
                "lora": _lora_report(),
            },
        )


def disable_lora():
    """
    Disable LoRA and return to base model.

    This does not necessarily remove the adapter from memory; it only disables it.
    Keeping the adapter cached makes switching back faster.
    """
    pipe = _load_pipeline()
    if pipe is None:
        return (
            "Model is not loaded.",
            {
                "status": "error",
                "error": _pipeline_error or "Pipeline failed to load",
                "lora": _lora_report(),
            },
        )

    try:
        with _lora_lock:
            _set_lora(pipe, "Base model", allow_load=False)

        return (
            "Base model active.",
            {
                "status": "ready",
                "message": "LoRA disabled. Base model active.",
                "lora": _lora_report(),
            },
        )

    except Exception as e:
        traceback.print_exc()
        return (
            f"Could not switch to base model cleanly: {type(e).__name__}: {e}",
            {
                "status": "error",
                "message": "Could not disable LoRA cleanly.",
                "error": f"{type(e).__name__}: {e}",
                "lora": _lora_report(),
            },
        )


# ---------------------------------------------------------------------------
# Eager app runtime warm-up
# ---------------------------------------------------------------------------

if os.environ.get("SKIP_MODEL_LOAD") != "1":
    _load_pipeline()


# ---------------------------------------------------------------------------
# Health / warm-up endpoints
# ---------------------------------------------------------------------------

def health():
    return {
        "status": "ready" if _pipeline is not None else "not_loaded",
        "model_ready": _pipeline is not None,
        "pipeline_ready": _pipeline is not None,
        "model_id": MODEL_ID,
        "runtime_mode": "zerogpu_compatibility_compile_bypass",
        "last_error": _pipeline_error or "",
        "runtime": _runtime_report(),
        "lora": _lora_report(),
    }


def warmup_model():
    """
    Manual refresh button.

    Usually the model is already loaded at app startup.
    """
    pipe = _load_pipeline()
    if pipe is None:
        return {
            "status": "error",
            "message": "Pipeline failed to load.",
            "error": _pipeline_error or "Pipeline failed to load",
            "runtime": _runtime_report(),
            "lora": _lora_report(),
        }

    return {
        "status": "ready",
        "message": "Model loaded and cached in this Space process.",
        "model_id": MODEL_ID,
        "runtime": _runtime_report(),
        "lora": _lora_report(),
    }


# ---------------------------------------------------------------------------
# Generation endpoint — real inference guarded by @spaces.GPU
# ---------------------------------------------------------------------------

def _gpu_duration(
    prompt,
    lora_style,
    num_blocks,
    num_inference_steps,
    seed,
    *args,
    **kwargs,
):
    try:
        blocks = int(num_blocks)
        steps = int(num_inference_steps)
    except Exception:
        blocks = 9
        steps = 6

    # Model and LoRA are expected to be loaded before generation.
    # Observed on this Space:
    #   9 blocks × 4 steps  < 75s
    #   9 blocks × 8 steps  < 80s
    return min(120, max(30, int(35 + blocks * steps * 1.2)))


@_spaces_gpu(duration=_gpu_duration, size="xlarge")
def generate(
    prompt,
    lora_style,
    num_blocks,
    num_inference_steps,
    seed,
    progress=gr.Progress(track_tqdm=True),
):
    pipe = _load_pipeline()

    if pipe is None:
        err = _pipeline_error or "Pipeline not loaded (unknown failure)"
        raise RuntimeError(f"Generation unavailable: {err}")

    if not isinstance(prompt, str) or not prompt.strip():
        raise ValueError("Prompt must be a non-empty string.")

    num_blocks = int(num_blocks)
    num_inference_steps = int(num_inference_steps)
    seed = int(seed)

    if num_blocks < 1 or num_blocks > 12:
        raise ValueError("num_blocks must be between 1 and 12.")

    if num_inference_steps < 1 or num_inference_steps > 8:
        raise ValueError("num_inference_steps must be between 1 and 8.")

    device = "cuda"

    try:
        pipe = pipe.to(device)
    except Exception as e:
        _log(f"Pipeline .to('cuda') warning: {type(e).__name__}: {e}")

    with _lora_lock:
        trigger = _set_lora(
            pipe,
            lora_style,
            allow_load=False,
        )

    final_prompt = prompt.strip()
    if trigger and not final_prompt.startswith(trigger):
        final_prompt = f"{trigger} {final_prompt}"

    frames = []
    state = PipelineState()

    try:
        generator = torch.Generator(device=device).manual_seed(seed)
    except Exception as e:
        _log(f"CUDA generator failed, falling back to CPU generator: {type(e).__name__}: {e}")
        generator = torch.Generator(device="cpu").manual_seed(seed)

    try:
        progress(0, desc="Preparing generation")

        for block_idx in progress.tqdm(
            range(num_blocks),
            desc="Generating video blocks",
        ):
            _log(f"Block {block_idx + 1}/{num_blocks}")

            progress(
                block_idx / max(1, num_blocks),
                desc=f"Generating block {block_idx + 1}/{num_blocks}",
            )

            state = pipe(
                state,
                prompt=[final_prompt],
                num_inference_steps=num_inference_steps,
                num_blocks=num_blocks,
                block_idx=block_idx,
                generator=generator,
            )

            videos = state.values.get("videos")
            if not videos:
                raise RuntimeError("Pipeline state did not contain `videos` after inference.")

            frames.extend(videos[0])

    except Exception as e:
        _log(f"Inference failed at block {locals().get('block_idx', 'unknown')}: {e}")
        traceback.print_exc()
        raise RuntimeError(
            f"Inference error at block {locals().get('block_idx', 'unknown')}: {e}"
        )

    if not frames:
        raise RuntimeError("No frames were generated.")

    progress(0.95, desc="Exporting video")

    output_path = f"/tmp/krea_output_{int(time.time())}.mp4"
    export_to_video(frames, output_path, fps=24)

    progress(1.0, desc="Done")
    _log(f"Saved video to {output_path}")

    return output_path


# ---------------------------------------------------------------------------
# Gradio app
# ---------------------------------------------------------------------------

with gr.Blocks(title="Krea Realtime Video 14B") as demo:
    gr.Markdown(
        "# Krea Realtime Video 14B\n"
        "Text-to-video with Diffusers ModularPipeline on ZeroGPU."
    )

    with gr.Row():
        with gr.Column(scale=4):
            prompt = gr.Textbox(
                label="Prompt",
                placeholder="Describe the video you want to generate...",
                lines=4,
            )

            with gr.Row():
                lora_style = gr.Dropdown(
                    choices=list(KNOWN_LORAS.keys()),
                    value="Base model",
                    label="Style",
                    scale=3,
                )

                load_lora_btn = gr.Button("Load style", variant="secondary", scale=1)

            style_status = gr.Textbox(
                label="Style status",
                value="Base model active.",
                interactive=False,
                lines=1,
            )

            with gr.Row():
                disable_lora_btn = gr.Button("Base model", variant="secondary")
                generate_btn = gr.Button("Generate video", variant="primary")

            with gr.Row():
                num_blocks = gr.Slider(
                    minimum=1,
                    maximum=12,
                    value=9,
                    step=1,
                    label="Length",
                )

                num_inference_steps = gr.Slider(
                    minimum=1,
                    maximum=8,
                    value=6,
                    step=1,
                    label="Quality",
                )

            seed = gr.Number(value=42, precision=0, label="Seed")

            gr.Markdown(
                "For Origami: select **Origami**, click **Load style**, then generate. "
                "The `[origami]` trigger is added automatically."
            )

        with gr.Column(scale=5):
            output_video = gr.Video(label="Generated video")

            with gr.Accordion("Advanced / status", open=False):
                model_status = gr.JSON(label="Model")
                lora_status = gr.JSON(label="LoRA")
                warmup_btn = gr.Button("Refresh status", variant="secondary")

    gr.Examples(
        examples=[
            [
                "An astronaut runs through a dense jungle, pushing through wet foliage while glowing insects swirl around and mist rises from the ground, the camera following as branches snap and leaves shake in his path, cinematic action scene",
                "Base model",
                9,
                6,
                42,
            ],
            [
                "A man sprints through a narrow neon alley at night while being chased, jumping over trash bins and splashing through puddles as loose papers swirl in the wind and signs flicker around him, intense handheld action shot",
                "Base model",
                9,
                6,
                123,
            ],
            [
                "A sports car drifts aggressively around tight corners on a snowy mountain road, kicking up clouds of snow and ice as headlights flash through the blizzard and the vehicle accelerates downhill, high-speed dynamic tracking shot",
                "Base model",
                9,
                6,
                99,
            ],
            [
                "A surfer rides a huge crashing wave, carving sharply across the water and spraying foam into the air as the wave curls behind him and sunlight glitters across the ocean, energetic action scene",
                "Base model",
                9,
                6,
                314,
            ],
            [
                "A small fighter spaceship races through an asteroid field, dodging spinning rocks and firing bright lasers while explosions flash in the distance, fast-paced sci-fi action sequence",
                "Base model",
                9,
                6,
                777,
            ],
            [
                "A skateboarder speeds through an urban plaza, jumping down stairs and landing tricks while pigeons scatter and people turn to watch, lively street action video",
                "Base model",
                9,
                6,
                888,
            ],
            [
                "A group of dancers performs an explosive street routine, spinning, jumping, and moving in sync while dust rises from the ground and bystanders cheer around them, high-energy choreography",
                "Base model",
                9,
                6,
                2026,
            ],
            [
                "An explorer runs across a collapsing rope bridge in the jungle, grabbing the ropes as wooden planks snap beneath his feet while mist rises from the canyon below and leaves whip in the wind, adventure action sequence",
                "Base model",
                9,
                6,
                515,
            ],
            [
                "a small origami fox runs through a paper forest, leaping over folded rocks and darting between swaying paper trees while loose paper leaves flutter around it, whimsical animated motion",
                "Origami",
                9,
                6,
                616,
            ],
            [
                "a giant origami dragon flies over a paper village, flapping its folded wings and sending paper rooftops trembling, then diving low before soaring upward again, dramatic handcrafted action",
                "Origami",
                9,
                6,
                717,
            ],
            [
                "an origami boat rides rough folded-paper waves during a storm, rocking violently as paper water splashes upward and the sail bends in the wind, dynamic handmade animation",
                "Origami",
                9,
                6,
                818,
            ],
        ],
        inputs=[
            prompt,
            lora_style,
            num_blocks,
            num_inference_steps,
            seed,
        ],
    )

    warmup_btn.click(
        warmup_model,
        inputs=None,
        outputs=model_status,
        api_name="warmup",
    )

    load_lora_btn.click(
        load_selected_lora,
        inputs=[lora_style],
        outputs=[style_status, lora_status],
        api_name="load_lora",
    )

    disable_lora_btn.click(
        disable_lora,
        inputs=None,
        outputs=[style_status, lora_status],
        api_name="disable_lora",
    )

    generate_btn.click(
        generate,
        inputs=[
            prompt,
            lora_style,
            num_blocks,
            num_inference_steps,
            seed,
        ],
        outputs=output_video,
        api_name="generate",
    )

    demo.load(
        lambda: health(),
        inputs=None,
        outputs=model_status,
        api_name="health",
    )


if __name__ == "__main__":
    demo.queue().launch(
        server_name="0.0.0.0",
        server_port=7860,
        show_error=True,
    )