Blame for app.py

"""SANA-Video 2.0 5B 720p — four-step text-to-video preview.

The DMD sampler, prompt construction, latent layout and VAE post-processing
mirror the authors' verified four-step inference command. The original
50-step T2V + TI2V release remains available from the linked base model.
"""

from __future__ import annotations

import os

os.environ.setdefault("DISABLE_XFORMERS", "1")
os.environ.setdefault("TOKENIZERS_PARALLELISM", "false")

# isort: off
import spaces  # noqa: E402  (must precede torch / CUDA-touching imports)
from preview_controls import (  # noqa: E402
    DEFAULT_MOTION_SCORE,
    DEFAULT_RL_LORA_SCALE,
    DEFAULT_VIDEO_DURATION,
    MAX_MOTION_SCORE,
    MAX_RL_LORA_SCALE,
    MIN_MOTION_SCORE,
    MIN_RL_LORA_SCALE,
    REFERENCE_RL_LORA_SCALE,
    RL_LORA_SCALE_STEP,
    VIDEO_PROFILES,
    get_video_profile,
    lora_adjustment_coefficient,
    motion_suffix,
    normalize_base_prompt,
    normalize_lora_scale,
    normalize_motion_score,
)

# isort: on

import gc  # noqa: E402
import random  # noqa: E402
import tempfile  # noqa: E402
import threading  # noqa: E402
import time  # noqa: E402

import gradio as gr  # noqa: E402
import imageio  # noqa: E402
import pyrallis  # noqa: E402
import torch  # noqa: E402
from accelerate import init_empty_weights  # noqa: E402
from huggingface_hub import hf_hub_download, snapshot_download  # noqa: E402

from diffusion import FastVideoDMD4Step  # noqa: E402
from diffusion.model.builder import (  # noqa: E402
    build_model,
    get_tokenizer_and_text_encoder,
    get_vae,
    vae_decode,
)
from diffusion.model.utils import get_weight_dtype  # noqa: E402
from diffusion.utils.config import SanaVideoConfig, model_video_init_config  # noqa: E402

# --------------------------------------------------------------------------------------
# Constants taken from the released config / the authors' verified commands
# --------------------------------------------------------------------------------------
MODEL_ID = "Efficient-Large-Model/SANA-Video_2.0_5B_720p_4step"
CKPT_FILE = "checkpoints/SANA_Video_2.0_5B_720p_4step.pth"
RL_LORA_FILE = "adapters/SANA_Video_2.0_5B_720p_RL500_LoRA_rank128.pt"
VAE_ID = "Efficient-Large-Model/LTX-2.3-Diffusers"
CONFIG_PATH = "configs/sana_video2/SanaVideo2_5B_720p.yaml"

# Hack2Cortex portrait test bucket: 1280 x 704
VIDEO_HEIGHT, VIDEO_WIDTH = 1280, 704
MODEL_CONSTRUCTION_IMAGE_SIZE = 480

# Authors' verified DMD preview settings. The schedule is fixed by training.
DEFAULT_STEPS = 4
DEFAULT_CFG = 1.0
GENERATOR_SIGMA_PROFILE = "sana_shift6_dpm"
MAX_SEED = 2**31 - 1
DEFAULT_SEED = 4

EXPECTED_LORA_FORMAT = "sana_video20_refl_lora_v1"
EXPECTED_LORA_GLOBAL_STEP = 500
EXPECTED_LORA_RANK = 128
EXPECTED_LORA_ALPHA = 256.0
EXPECTED_LORA_MODULES = 256
EXPECTED_LORA_TARGET_ELEMENTS = 4_194_304_000
EXPECTED_PUBLIC_LORA_BYTES = 587_418_315

torch.set_grad_enabled(False)


class RLLoRAScaleController:
    """Reconstruct the latest DMD model at a requested RL LoRA multiplier.

    The latest full checkpoint is the scale=1 reference. For every scale change
    we start from immutable CPU BF16 copies of the 256 target weights and apply
    ``(scale - 1) * (alpha / rank) * (B @ A)``. This keeps repeated requests
    deterministic and prevents cumulative BF16 drift.
    """

    def __init__(self, path: str) -> None:
        if os.path.getsize(path) != EXPECTED_PUBLIC_LORA_BYTES:
            raise RuntimeError("Public RL LoRA file size does not match the validated release artifact.")

        payload = torch.load(path, map_location="cpu", mmap=True, weights_only=True)
        if not isinstance(payload, dict):
            raise TypeError("RL LoRA checkpoint must be a mapping.")
        if payload.get("format") != EXPECTED_LORA_FORMAT or payload.get("global_step") != EXPECTED_LORA_GLOBAL_STEP:
            raise RuntimeError("RL LoRA checkpoint identity does not match the release contract.")

        config = payload.get("lora_config")
        state = payload.get("lora_state_dict")
        release_metadata = payload.get("release_metadata")
        if not isinstance(config, dict) or not isinstance(state, dict):
            raise TypeError("RL LoRA checkpoint is missing its config or tensor mapping.")
        if not isinstance(release_metadata, dict):
            raise TypeError("RL LoRA checkpoint is missing public release metadata.")
        if (
            release_metadata.get("intended_model") != MODEL_ID
            or float(release_metadata.get("full_model_reference_scale", -1)) != REFERENCE_RL_LORA_SCALE
            or float(release_metadata.get("space_default_scale", -1)) != DEFAULT_RL_LORA_SCALE
            or release_metadata.get("space_supported_scale_range") != [MIN_RL_LORA_SCALE, MAX_RL_LORA_SCALE]
        ):
            raise RuntimeError("RL LoRA release metadata does not target this preview model.")

        module_names = config.get("matched_module_names")
        rank = int(config.get("rank", -1))
        alpha = float(config.get("alpha", float("nan")))
        if (
            not isinstance(module_names, list)
            or len(module_names) != EXPECTED_LORA_MODULES
            or len(set(module_names)) != EXPECTED_LORA_MODULES
            or rank != EXPECTED_LORA_RANK
            or alpha != EXPECTED_LORA_ALPHA
        ):
            raise RuntimeError("RL LoRA rank, alpha, or module coverage has drifted.")

        expected_keys = {f"{name}.{suffix}" for name in module_names for suffix in ("lora_A", "lora_B")}
        if set(state) != expected_keys or any(not torch.is_tensor(value) for value in state.values()):
            raise RuntimeError("RL LoRA tensor coverage does not match the release contract.")

        self.module_names = tuple(module_names)
        self.rank = rank
        self.alpha = alpha
        self.intrinsic_scaling = alpha / rank
        self.state = state
        self.payload = payload  # Keep mmap-backed adapter tensors alive.
        self.reference: dict[str, torch.Tensor] | None = None
        self.parameters: dict[str, torch.nn.Parameter] = {}
        self.current_scale = REFERENCE_RL_LORA_SCALE

    def capture_reference(self, state_dict: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]:
        """Capture immutable scale=1 target weights before model assignment."""
        references: dict[str, torch.Tensor] = {}
        target_elements = 0
        for name in self.module_names:
            key = f"{name}.weight"
            if key not in state_dict:
                raise KeyError(f"Full checkpoint is missing RL LoRA target {key}.")
            weight = state_dict[key]
            lora_a = self.state[f"{name}.lora_A"]
            lora_b = self.state[f"{name}.lora_B"]
            if tuple(lora_a.shape) != (self.rank, weight.shape[1]) or tuple(lora_b.shape) != (
                weight.shape[0],
                self.rank,
            ):
                raise RuntimeError(f"RL LoRA shape mismatch for {name}.")
            references[key] = weight.detach().to(device="cpu", dtype=torch.bfloat16).clone()
            target_elements += weight.numel()
        if target_elements != EXPECTED_LORA_TARGET_ELEMENTS:
            raise RuntimeError("RL LoRA target parameter coverage does not match the release contract.")
        return references

    def register_model(
        self,
        model: torch.nn.Module,
        reference: dict[str, torch.Tensor],
    ) -> None:
        """Bind validated target parameters after the full model reaches CUDA."""
        named_parameters = dict(model.named_parameters())
        selected: dict[str, torch.nn.Parameter] = {}
        for name in self.module_names:
            key = f"{name}.weight"
            parameter = named_parameters.get(key)
            if parameter is None or tuple(parameter.shape) != tuple(reference[key].shape):
                raise RuntimeError(f"Loaded model does not expose RL LoRA target {key}.")
            if parameter.device.type != "cuda":
                raise RuntimeError(f"RL LoRA target is not on CUDA: {key}.")
            selected[key] = parameter
        self.reference = reference
        self.parameters = selected

    @torch.inference_mode()
    def apply_scale(self, requested_scale: int | float) -> float:
        """Reconstruct weights from the immutable scale=1 reference."""
        scale = normalize_lora_scale(requested_scale)
        if scale == self.current_scale:
            return 0.0
        if self.reference is None or len(self.parameters) != len(self.module_names):
            raise RuntimeError("RL LoRA controller is not bound to the model.")

        coefficient = lora_adjustment_coefficient(scale, self.intrinsic_scaling)
        started = time.perf_counter()
        print(
            f"[lora] Reconstructing scale={scale:.1f}; relative BA coefficient={coefficient:+.2f}",
            flush=True,
        )
        try:
            for index, name in enumerate(self.module_names, start=1):
                key = f"{name}.weight"
                reference = self.reference[key].to("cuda", dtype=torch.float32)
                lora_a = self.state[f"{name}.lora_A"].to("cuda", dtype=torch.float32)
                lora_b = self.state[f"{name}.lora_B"].to("cuda", dtype=torch.float32)
                reference.addmm_(lora_b, lora_a, alpha=coefficient)
                self.parameters[key].copy_(reference.to(dtype=self.parameters[key].dtype))
                del reference, lora_a, lora_b
                if index % 32 == 0 or index == len(self.module_names):
                    print(f"[lora] targets={index}/{len(self.module_names)}", flush=True)
            torch.cuda.synchronize()
        except Exception:
            self.current_scale = float("nan")
            raise

        self.current_scale = scale
        elapsed = time.perf_counter() - started
        torch.cuda.empty_cache()
        print(f"[lora] Scale {scale:.1f} ready in {elapsed:.1f}s", flush=True)
        return elapsed


# --------------------------------------------------------------------------------------
# Load everything at module scope and move it to CUDA eagerly (ZeroGPU requirement)
# --------------------------------------------------------------------------------------
config: SanaVideoConfig = pyrallis.load(SanaVideoConfig, open(CONFIG_PATH))
# The distilled checkpoint was trained/evaluated with the 480 source tower and
# an explicit 736×1280 inference bucket. SanaVideo2 uses dynamic RoPE at runtime.
config.model.image_size = MODEL_CONSTRUCTION_IMAGE_SIZE

weight_dtype = get_weight_dtype(config.model.mixed_precision)  # bfloat16
vae_dtype = get_weight_dtype(config.vae.weight_dtype)  # bfloat16

# Inference needs the LTX-2.3 decoder even though the training recipe selects the
# memory-efficient causal encoder-only path.
config.vae.use_causal_encode = False
config.vae.vae_pretrained = snapshot_download(VAE_ID, allow_patterns=["vae/*"])

print("Loading LTX-2.3 VAE …", flush=True)
vae = get_vae(
    config.vae.vae_type,
    config.vae.vae_pretrained,
    device="cuda",
    dtype=vae_dtype,
    config=config.vae,
)

print("Loading Gemma-2-2B-it text encoder …", flush=True)
tokenizer, text_encoder = get_tokenizer_and_text_encoder(name=config.text_encoder.text_encoder_name, device="cuda")

print("Building SANA-Video 2.0 5B transformer …", flush=True)
latent_size = config.model.image_size // config.vae.vae_downsample_rate
# `init_empty_weights` puts only *parameters* on the meta device, so the 17.9 GB
# checkpoint can be assigned straight in without ever materialising fp32 weights.
# Crucially it leaves plain tensor attributes real (e.g. `WanRotaryPosEmbed.freqs`,
# which is not a registered buffer) — `torch.device("meta")` would not.
with init_empty_weights(include_buffers=False):
    model = build_model(
        config.model.model,
        use_fp32_attention=config.model.fp32_attention,
        **model_video_init_config(config, latent_size=latent_size),
    )
print(f"Model parameters: {sum(p.numel() for p in model.parameters()):,}", flush=True)

print("Downloading transformer weights (~17.9 GB) …", flush=True)
ckpt_path = os.environ.get("SANA_VIDEO2_PREVIEW_CKPT") or hf_hub_download(MODEL_ID, CKPT_FILE)
print("Downloading RL LoRA scale controller (~587 MB) …", flush=True)
rl_lora_path = os.environ.get("SANA_VIDEO2_PREVIEW_RL_LORA") or hf_hub_download(MODEL_ID, RL_LORA_FILE)
rl_lora_controller = RLLoRAScaleController(rl_lora_path)
print("Loading transformer weights …", flush=True)
_raw = torch.load(ckpt_path, map_location="cpu", mmap=True, weights_only=True)
_raw = _raw.get("state_dict_ema", _raw.get("state_dict", _raw))
state_dict = {
    k[len("model.") :] if k.startswith("model.") else k: (v.to(weight_dtype) if torch.is_floating_point(v) else v)
    for k, v in _raw.items()
}
del _raw
# `pos_embed` is an unused all-zeros buffer on this architecture (it uses Wan RoPE);
# the reference drops it from the checkpoint, we just materialise it at the right shape.
state_dict["pos_embed"] = torch.zeros(1, latent_size * latent_size, model.hidden_size, dtype=weight_dtype)
lora_reference = rl_lora_controller.capture_reference(state_dict)

missing, unexpected = model.load_state_dict(state_dict, strict=False, assign=True)
print(f"Missing keys: {missing}", flush=True)
print(f"Unexpected keys: {unexpected}", flush=True)
# Positional embeddings are generated dynamically from the requested video bucket.
unexpected_missing = sorted(set(missing) - {"pos_embed"})
assert not unexpected_missing, f"checkpoint does not cover the model: {unexpected_missing}"
assert not unexpected, f"checkpoint has unexpected model keys: {unexpected}"
del state_dict
gc.collect()

model = model.eval().to(weight_dtype).to("cuda")
model.requires_grad_(False)
rl_lora_controller.register_model(model, lora_reference)
generation_lock = threading.Lock()

reference_gib = sum(t.numel() * t.element_size() for t in lora_reference.values()) / 2**30
print(
    f"Model ready on CUDA; immutable scale=1 RL target reference: {reference_gib:.2f} GiB CPU.",
    flush=True,
)


def _encode_prompt(prompt: str):
    """Reproduce the reference's complex-human-instruction prompt encoding."""
    chi_prompt = "\n".join(config.text_encoder.chi_prompt)
    num_chi_prompt_tokens = len(tokenizer.encode(chi_prompt))
    max_length_all = num_chi_prompt_tokens + config.text_encoder.model_max_length - 2
    tok = tokenizer(
        [chi_prompt + prompt],
        max_length=max_length_all,
        padding="max_length",
        truncation=True,
        return_tensors="pt",
    ).to("cuda")
    select_index = [0] + list(range(-config.text_encoder.model_max_length + 1, 0))
    embs = text_encoder(tok.input_ids, tok.attention_mask)[0][:, None][:, :, select_index]
    return embs, tok.attention_mask[:, select_index]


def _estimate_duration(
    prompt: str = "",
    video_duration: str = DEFAULT_VIDEO_DURATION,
    rl_lora_scale: float = DEFAULT_RL_LORA_SCALE,
    motion_score: int = DEFAULT_MOTION_SCORE,
    seed: int = DEFAULT_SEED,
    randomize_seed: bool = False,
    *args,
    **kwargs,
) -> int:
    # Five-second live validation completed in about 22 seconds end to end. The
    # 193-frame profile reserves extra decoder headroom while staying bounded.
    try:
        profile = get_video_profile(video_duration)
    except ValueError:
        return 120
    return 75 if profile.num_frames == 81 else 120


@spaces.GPU(duration=_estimate_duration)
@torch.inference_mode()
def generate(
    prompt: str,
    video_duration: str = DEFAULT_VIDEO_DURATION,
    rl_lora_scale: float = DEFAULT_RL_LORA_SCALE,
    motion_score: int = DEFAULT_MOTION_SCORE,
    seed: int = DEFAULT_SEED,
    randomize_seed: bool = False,
    progress=gr.Progress(track_tqdm=True),
):
    """Generate a four-step 720p T2V preview with SANA-Video 2.0 5B.

    Args:
        prompt: Text description of the video to generate.
        video_duration: Release-validated five- or eight-second temporal profile.
        rl_lora_scale: Post-hoc RL multiplier around the full scale=1 checkpoint.
        motion_score: Numeric prompt suffix; zero disables the suffix.
        seed: RNG seed.
        randomize_seed: Draw a fresh random seed for this run.

    Returns:
        A tuple of (path to the generated MP4, the seed that was used).
    """
    if not prompt or not prompt.strip():
        raise gr.Error("Please enter a prompt.")

    try:
        profile = get_video_profile(video_duration)
        score = normalize_motion_score(motion_score)
        scale = normalize_lora_scale(rl_lora_scale)
    except (TypeError, ValueError) as exc:
        raise gr.Error(str(exc)) from exc

    base_prompt = normalize_base_prompt(prompt)
    if not base_prompt:
        raise gr.Error("Please enter a scene description, not only a motion-score suffix.")
    conditioned_prompt = base_prompt + motion_suffix(score)
    seed = random.randint(0, MAX_SEED) if randomize_seed else int(seed)

    # Scale changes and inference share a lock: another request must not rewrite
    # target weights while this request is sampling.
    with generation_lock:
        t0 = time.perf_counter()
        progress(0.02, desc="Applying RL LoRA scale")
        scale_elapsed = rl_lora_controller.apply_scale(scale)
        generator = torch.Generator(device="cuda").manual_seed(seed)

        latent_h = VIDEO_HEIGHT // config.vae.vae_stride[1]
        latent_w = VIDEO_WIDTH // config.vae.vae_stride[2]
        latent_t = (profile.num_frames - 1) // config.vae.vae_stride[0] + 1

        progress(0.08, desc="Encoding prompt")
        stage_started = time.perf_counter()
        caption_embs, emb_masks = _encode_prompt(conditioned_prompt)
        torch.cuda.synchronize()
        prompt_elapsed = time.perf_counter() - stage_started

        hw = torch.tensor([[float(VIDEO_HEIGHT), float(VIDEO_WIDTH)]], dtype=torch.float, device="cuda")
        z = torch.randn(
            1,
            config.vae.vae_latent_dim,
            latent_t,
            latent_h,
            latent_w,
            device="cuda",
            dtype=weight_dtype,
            generator=generator,
        )
        model_kwargs = dict(data_info={"img_hw": hw}, mask=emb_masks)
        solver = FastVideoDMD4Step(
            model,
            caption_embs,
            cfg_scale=DEFAULT_CFG,
            generator_sigma_profile=GENERATOR_SIGMA_PROFILE,
            model_kwargs=model_kwargs,
        )
        progress(0.16, desc=f"Running {DEFAULT_STEPS} denoising steps")
        stage_started = time.perf_counter()
        samples = solver.sample(z, steps=DEFAULT_STEPS, generator=generator)
        torch.cuda.synchronize()
        dit_elapsed = time.perf_counter() - stage_started

        progress(0.58, desc=f"Decoding {profile.num_frames} video frames")
        stage_started = time.perf_counter()
        samples = vae_decode(config.vae.vae_type, vae, samples.to(vae_dtype))
        if isinstance(samples, list):
            samples = torch.stack(samples, dim=0)
        torch.cuda.synchronize()
        vae_elapsed = time.perf_counter() - stage_started

        progress(0.84, desc="Transferring decoded frames")
        stage_started = time.perf_counter()
        video = (
            torch.clamp(127.5 * samples + 127.5, 0, 255).permute(0, 2, 3, 4, 1).to("cpu", dtype=torch.uint8)[0].numpy()
        )
        transfer_elapsed = time.perf_counter() - stage_started
        del samples
        torch.cuda.empty_cache()

    progress(0.92, desc="Encoding MP4")
    stage_started = time.perf_counter()
    out_path = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False).name
    writer = imageio.get_writer(
        out_path,
        fps=profile.fps,
        codec="libx264",
        quality=8,
        output_params=["-preset", "veryfast"],
    )
    try:
        for frame in video:
            writer.append_data(frame)
    finally:
        writer.close()
    mp4_elapsed = time.perf_counter() - stage_started
    total_elapsed = time.perf_counter() - t0
    output_mib = os.path.getsize(out_path) / 2**20
    progress(1.0, desc="Video ready")

    print(
        f"[timing] t2v {profile.num_frames}f/{profile.fps}fps x {DEFAULT_STEPS} steps, "
        f"motion={score}, RL scale={scale:.1f}: scale={scale_elapsed:.2f}s, "
        f"prompt={prompt_elapsed:.2f}s, DiT={dit_elapsed:.2f}s, VAE={vae_elapsed:.2f}s, "
        f"GPU-to-CPU={transfer_elapsed:.2f}s, MP4={mp4_elapsed:.2f}s, "
        f"total={total_elapsed:.2f}s, output={output_mib:.1f}MiB",
        flush=True,
    )
    return out_path, seed


def run_example(prompt: str):
    """Run an example row at the UI defaults.

    `gr.Examples` calls its `fn` with exactly the values in the row, but Gradio inserts the
    `gr.Progress` object at the *positional index* it occupies in the callee's signature.
    With short example rows that index is past the end of the list, which puts Progress in
    the wrong slot (and then IndexErrors when Gradio reads it back). Routing examples
    through this Progress-free wrapper — rather than reordering `generate`, which would
    also shift every API/MCP parameter name by one — keeps both paths correct.

    Args:
        prompt: Text description of the video to generate.
    Returns:
        Path to the generated MP4 and the seed that produced it.
    """
    return generate(prompt=prompt)


# --------------------------------------------------------------------------------------
# UI
# --------------------------------------------------------------------------------------
CSS = """
:root {
  --sana-green: #76b900;
  --sana-green-bright: #a7e132;
  --sana-ink: #070a08;
  --sana-muted: #6b756d;
}

.gradio-container {
  background:
    radial-gradient(circle at 50% -10%, rgba(118, 185, 0, 0.09), transparent 30rem),
    var(--background-fill-primary) !important;
}

#col-container {
  max-width: 1180px;
  margin: 0 auto;
  padding: 18px 12px 44px;
  gap: 20px;
}

#sana-hero {
  position: relative;
  overflow: hidden;
  padding: clamp(28px, 5vw, 54px);
  border: 1px solid rgba(167, 225, 50, 0.28);
  border-radius: 26px;
  color: #f7faf7 !important;
  color-scheme: dark;
  background:
    radial-gradient(circle at 88% 12%, rgba(118, 185, 0, 0.34), transparent 30%),
    linear-gradient(135deg, #071008 0%, #101713 58%, #172019 100%);
  box-shadow: 0 24px 70px rgba(0, 0, 0, 0.2);
}

#sana-hero::after {
  position: absolute;
  right: -70px;
  bottom: -110px;
  width: 300px;
  height: 300px;
  border: 1px solid rgba(167, 225, 50, 0.18);
  border-radius: 50%;
  content: "";
}

.sana-eyebrow {
  display: flex;
  gap: 9px;
  align-items: center;
  margin: 0 0 13px;
  color: var(--sana-green-bright) !important;
  font-size: 0.72rem;
  font-weight: 800;
  letter-spacing: 0.13em;
  text-transform: uppercase;
}

.sana-eyebrow::before {
  width: 8px;
  height: 8px;
  border-radius: 2px;
  background: var(--sana-green);
  box-shadow: 0 0 18px rgba(167, 225, 50, 0.75);
  content: "";
}

#sana-hero h1 {
  position: relative;
  z-index: 1;
  margin: 0;
  font-size: clamp(2.55rem, 6vw, 5rem);
  font-weight: 760;
  line-height: 0.98;
  letter-spacing: -0.055em;
  color: #f7faf7 !important;
}

#sana-hero h1 span { color: var(--sana-green-bright) !important; }

.sana-title-row {
  position: relative;
  z-index: 1;
  display: flex;
  flex-wrap: wrap;
  gap: 14px 20px;
  align-items: flex-end;
}

.sana-preview-badge {
  display: inline-flex;
  align-items: center;
  min-height: 42px;
  margin-bottom: 4px;
  padding: 8px 15px;
  border: 1px solid var(--sana-green-bright);
  border-radius: 999px;
  color: #071008 !important;
  background: var(--sana-green-bright);
  box-shadow: 0 0 24px rgba(167, 225, 50, 0.2);
  font-size: clamp(0.88rem, 1.5vw, 1.05rem);
  font-weight: 850;
  line-height: 1;
  letter-spacing: 0.035em;
  text-transform: uppercase;
  white-space: nowrap;
}

.sana-lede {
  position: relative;
  z-index: 1;
  max-width: 760px;
  margin: 18px 0 0;
  color: rgba(247, 250, 247, 0.78) !important;
  font-size: clamp(0.95rem, 1.6vw, 1.08rem);
  line-height: 1.65;
}

.sana-specs,
.sana-nav {
  position: relative;
  z-index: 1;
  display: flex;
  flex-wrap: wrap;
  gap: 8px;
}

.sana-specs { margin-top: 22px; }

.sana-specs span {
  padding: 6px 10px;
  border: 1px solid rgba(255, 255, 255, 0.15);
  border-radius: 999px;
  color: rgba(255, 255, 255, 0.82) !important;
  background: rgba(255, 255, 255, 0.06);
  font-size: 0.68rem;
  font-weight: 750;
  letter-spacing: 0.04em;
  text-transform: uppercase;
}

.sana-nav { margin-top: 25px; }

.sana-nav a {
  display: inline-flex;
  align-items: center;
  min-height: 38px;
  padding: 0 15px;
  border: 1px solid rgba(255, 255, 255, 0.2);
  border-radius: 10px;
  color: #f7faf7 !important;
  background: rgba(255, 255, 255, 0.055);
  text-decoration: none !important;
  font-size: 0.78rem;
  font-weight: 720;
  transition: border-color 150ms ease, background 150ms ease, transform 150ms ease;
}

.sana-nav a:first-child {
  border-color: var(--sana-green);
  color: #071008 !important;
  background: var(--sana-green);
}

.sana-nav a:hover,
.sana-nav a:focus-visible {
  border-color: var(--sana-green-bright);
  background: rgba(118, 185, 0, 0.2);
  transform: translateY(-1px);
}

.sana-nav a:first-child:hover,
.sana-nav a:first-child:focus-visible { background: var(--sana-green-bright); }

.sana-section-heading {
  display: flex;
  gap: 14px;
  align-items: end;
  justify-content: space-between;
  margin-top: 4px;
  padding: 0 4px;
}

.sana-section-heading h2 {
  margin: 0;
  font-size: 1.35rem;
  letter-spacing: -0.025em;
}

.sana-section-heading p {
  margin: 0;
  color: var(--body-text-color-subdued);
  font-size: 0.8rem;
}

.sana-create-heading { margin: 18px 4px -8px; }

#examples-heading {
  min-height: 0 !important;
  padding: 0 !important;
}

#examples-heading > .html-container {
  min-height: 0 !important;
  padding: 0 !important;
}

#generation-workspace { gap: 16px; align-items: stretch; }

.sana-card {
  padding: 19px !important;
  border: 1px solid rgba(118, 185, 0, 0.2) !important;
  border-radius: 20px !important;
  background: var(--block-background-fill) !important;
  box-shadow: 0 12px 34px rgba(0, 0, 0, 0.07);
}

.sana-panel-label {
  margin-bottom: 13px;
  color: var(--sana-green);
  font-size: 0.69rem;
  font-weight: 820;
  letter-spacing: 0.11em;
  text-transform: uppercase;
}

#prompt-input textarea {
  min-height: 132px;
  border-color: rgba(118, 185, 0, 0.24);
  font-size: 0.96rem;
  line-height: 1.55;
}

#generate-button {
  min-height: 48px;
  margin-top: 7px;
  border: 1px solid #679f00 !important;
  color: #071008 !important;
  background: var(--sana-green) !important;
  box-shadow: 0 10px 22px rgba(118, 185, 0, 0.22);
  font-weight: 820;
}

#generate-button:hover,
#generate-button:focus-visible {
  border-color: var(--sana-green-bright) !important;
  background: var(--sana-green-bright) !important;
}

#result-video { flex: 1; }

.sana-helper,
.sana-output-meta {
  margin: 10px 2px 0;
  color: var(--body-text-color-subdued);
  font-size: 0.73rem;
  line-height: 1.5;
}

.sana-output-meta {
  display: flex;
  flex-wrap: wrap;
  gap: 6px 14px;
  padding-top: 10px;
  border-top: 1px solid rgba(118, 185, 0, 0.14);
}

.sana-output-meta b { color: var(--sana-green); }

.sana-advanced {
  border-color: rgba(118, 185, 0, 0.2) !important;
  border-radius: 16px !important;
}

.sana-examples { margin-top: 8px; }

.sana-footer {
  padding: 4px;
  color: var(--body-text-color-subdued);
  text-align: center;
  font-size: 0.72rem;
}

#col-container :is(a, button, input, textarea):focus-visible {
  outline: 2px solid var(--sana-green-bright) !important;
  outline-offset: 2px;
}

.dark .gradio-container { color: var(--body-text-color); }
.dark .sana-card { box-shadow: 0 16px 42px rgba(0, 0, 0, 0.24); }

@media (max-width: 760px) {
  #col-container { padding-inline: 4px; }
  #sana-hero { padding: 26px 22px; border-radius: 20px; }
  .sana-title-row { gap: 12px; }
  .sana-preview-badge { min-height: 38px; margin-bottom: 1px; padding: 7px 13px; }
  .sana-section-heading { display: block; }
  .sana-section-heading p { margin-top: 5px; }
  .sana-create-heading { margin: 14px 0 -8px; }
  #generation-workspace { flex-direction: column; }
  .sana-card { padding: 15px !important; }
  .sana-nav a { flex: 1 1 auto; justify-content: center; }
}
"""

# Five prompts selected from the official SANA-Video 2.0 project page.
# Source: Efficient-Large-Model/Sana-assets, Video2/assets/curated-20260812.
PROJECT_PAGE_EXAMPLES = [
    (
        "A close-up shot of an astronaut standing on the moon's surface, wearing a white spacesuit "
        "with an American flag patch on the arm. The helmet visor reflects the lunar landscape and "
        "stars in the night sky. The astronaut's face is partially visible through the clear visor, "
        "showing a serious expression. The background features the rugged, rocky terrain of the moon, "
        "with small craters and hills stretching into the distance under a starry sky. The scene "
        "captures the vastness and isolation of space exploration.",
        "examples/project-page-ti2v/00_ffc54a70a0fc48678532d50b2e9c5beb.png",
    ),
    (
        "A vibrant, iridescent dragon with large, expressive eyes and spiky scales is seen in a sunny "
        "outdoor setting. The dragon is positioned on a patch of dirt, surrounded by greenery and a "
        "bright rainbow arching in the background. The dragon's mouth opens and closes, revealing its "
        "small teeth, as it looks around curiously. The sunlight casts a warm glow on the dragon's "
        "scales, highlighting its intricate patterns. The camera remains steady, focusing on the "
        "dragon's face and upper body, capturing its movements and expressions in detail.",
        "examples/project-page-ti2v/01_37b21f2cc8f64f4a81f3b39c8ed8cb42.png",
    ),
    (
        "In a cozy, vintage room adorned with floral wallpaper, a cartoon rooster sits comfortably "
        "in a floral-patterned armchair, sipping from a bottle of beer. The rooster, with its vibrant "
        "red comb and wattle, displays a range of expressions—smiling, nodding, and opening its beak "
        "wide in a cheerful manner. The setting includes wooden furniture and another beer bottle on "
        "the table, adding to the relaxed atmosphere. The camera captures the rooster from a close-up "
        "angle, emphasizing its animated movements and lively demeanor.",
        "examples/project-page-ti2v/03_afa93d2c952d4cdeaf1356225b27064b.png",
    ),
    (
        "In a cluttered laboratory filled with various scientific equipment, a small gray alien with "
        "large black eyes and a curious gray cat with striking yellow eyes are engaged in an "
        "experiment. The alien, holding a magnifying glass, examines a test tube filled with a yellow "
        "liquid, while the cat watches attentively, its tail swishing slightly. The background is "
        "adorned with shelves stocked with jars, bottles, and other lab supplies, creating a whimsical "
        "and imaginative setting. The scene captures a moment of scientific exploration and curiosity "
        "between the two characters. A medium shot of a dynamic interaction.",
        "examples/project-page-ti2v/05_962ea333452948c4963c46727fafef3b.png",
    ),
    (
        "A sleek red Jaguar sports car is captured in motion along a winding coastal road during a "
        "breathtaking sunset. The car, with its headlights on, glides smoothly down the wet asphalt, "
        "leaving a trail of motion blur behind it. The driver, visible through the windshield, appears "
        "focused and relaxed. The backdrop features a stunning ocean view with waves crashing against "
        "rocky cliffs, and the sky is painted with hues of orange and pink as the sun sets. The camera "
        "follows the car from a low angle, emphasizing its speed and the dynamic nature of the scene.",
        "examples/project-page-ti2v/07_e4fcd27037fa430a9e9aba96594245aa.png",
    ),
]

T2V_EXAMPLES = [[prompt] for prompt, _ in PROJECT_PAGE_EXAMPLES]

with gr.Blocks(title="SANA-Video 2.0 · 5B · 720p · 4-Step Preview") as demo:
    with gr.Column(elem_id="col-container"):
        gr.HTML(
            """
            <header id="sana-hero">
              <p class="sana-eyebrow">NVIDIA Research · Efficient AI</p>
              <div class="sana-title-row">
                <h1>SANA-Video <span>2.0</span></h1>
                <span class="sana-preview-badge">4-Step Preview</span>
              </div>
              <p class="sana-lede">
                A four-step research preview of the open 5B video model with hybrid
                linear/softmax attention and Attention Residuals. Generate a 720p
                video from text in only four steps, with adjustable motion conditioning
                and RL LoRA strength.
              </p>
              <div class="sana-specs" aria-label="Model capabilities">
                <span>5B open weights</span><span>704 × 1280</span>
                <span>4-step preview</span><span>5 / 8 seconds</span>
                <span>81 / 193 frames</span><span>RL scale 0.3–1.0</span><span>T2V</span>
              </div>
              <nav class="sana-nav" aria-label="SANA-Video 2.0 resources">
                <a href="https://nvlabs.github.io/Sana/Video2/" target="_blank" rel="noopener noreferrer">Project Page ↗</a>
                <a href="https://huggingface.co/Efficient-Large-Model/SANA-Video_2.0_5B_720p_4step" target="_blank" rel="noopener noreferrer">4-Step Model ↗</a>
                <a href="https://huggingface.co/Efficient-Large-Model/SANA-Video_2.0_5B_720p" target="_blank" rel="noopener noreferrer">50-Step Model ↗</a>
                <a href="https://nvlabs.github.io/Sana/docs/sana_video2/" target="_blank" rel="noopener noreferrer">Documentation ↗</a>
                <a href="https://github.com/NVlabs/Sana" target="_blank" rel="noopener noreferrer">GitHub ↗</a>
                <a href="https://arxiv.org/abs/2607.21553" target="_blank" rel="noopener noreferrer">Paper ↗</a>
              </nav>
            </header>
            <div class="sana-section-heading sana-create-heading">
              <h2>Create a video</h2>
              <p>Four-step text-to-video research preview.</p>
            </div>
            """
        )

        with gr.Row(elem_id="generation-workspace"):
            with gr.Column(scale=5, elem_classes=["sana-card", "sana-input-card"]):
                gr.HTML('<div class="sana-panel-label">01 · Prompt & conditioning</div>')
                prompt = gr.Textbox(
                    label="Prompt",
                    show_label=False,
                    placeholder="Describe the scene, subject, motion, lighting, and camera…",
                    lines=5,
                    container=False,
                    elem_id="prompt-input",
                )
                with gr.Row():
                    video_duration = gr.Radio(
                        choices=list(VIDEO_PROFILES),
                        value=DEFAULT_VIDEO_DURATION,
                        label="Video duration",
                        info="5 s: 81 frames at 16 FPS · 8 s: 193 frames at 24 FPS",
                    )
                    rl_lora_scale = gr.Slider(
                        label="RL LoRA scale",
                        info="0.7 is the preview default; 1.0 restores the full checkpoint.",
                        minimum=MIN_RL_LORA_SCALE,
                        maximum=MAX_RL_LORA_SCALE,
                        step=RL_LORA_SCALE_STEP,
                        value=DEFAULT_RL_LORA_SCALE,
                    )
                run_button = gr.Button(
                    "Generate 4-step preview",
                    variant="primary",
                    elem_id="generate-button",
                )
                gr.HTML(
                    """
                    <p class="sana-helper">
                      The distilled schedule stays fixed at 4 steps and CFG 1.
                      Motion conditioning and the RL LoRA multiplier are adjustable;
                      the default multiplier is 0.7.
                    </p>
                    """
                )
            with gr.Column(scale=7, elem_classes=["sana-card", "sana-output-card"]):
                gr.HTML('<div class="sana-panel-label">02 · Generated result</div>')
                result = gr.Video(
                    label="Result",
                    show_label=False,
                    autoplay=True,
                    height=470,
                    elem_id="result-video",
                )
                gr.HTML(
                    """
                    <div class="sana-output-meta" aria-label="Default output settings">
                      <span><b>Resolution</b> 704 × 1280</span>
                      <span><b>Length</b> 5 or 8 seconds</span>
                      <span><b>Playback</b> 16 or 24 FPS</span>
                      <span><b>Inference</b> 4 steps</span>
                      <span><b>Default RL scale</b> 0.7</span>
                      <span><b>Default motion</b> 20</span>
                    </div>
                    """
                )

        with gr.Accordion("Reproducibility", open=True, elem_classes=["sana-advanced"]):
            with gr.Row():
                motion_score = gr.Slider(
                    label="Motion score",
                    info="Appended as a prompt suffix; 0 disables it.",
                    minimum=MIN_MOTION_SCORE,
                    maximum=MAX_MOTION_SCORE,
                    step=1,
                    value=DEFAULT_MOTION_SCORE,
                )
                seed = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=DEFAULT_SEED)
                randomize_seed = gr.Checkbox(label="Randomize seed", value=False)

        gr.HTML(
            """
            <div class="sana-section-heading sana-examples">
              <h2>Start from an example</h2>
              <p>Select a card to load and run it with the UI defaults.</p>
            </div>
            """,
            elem_id="examples-heading",
            min_height=0,
            container=False,
            padding=False,
        )
        gr.Examples(
            examples=T2V_EXAMPLES,
            inputs=[prompt],
            outputs=[result, seed],
            fn=run_example,
            cache_examples=True,
            cache_mode="lazy",
        )

        gr.HTML(
            """
            <footer class="sana-footer">
              SANA-Video 2.0 · 4-step research preview · Apache-2.0 · Hugging Face ZeroGPU
            </footer>
            """
        )

    inputs = [
        prompt,
        video_duration,
        rl_lora_scale,
        motion_score,
        seed,
        randomize_seed,
    ]
    gr.on(
        triggers=[run_button.click, prompt.submit],
        fn=generate,
        inputs=inputs,
        outputs=[result, seed],
        api_name="generate",
    )

if __name__ == "__main__":
    demo.launch(theme=gr.themes.Citrus(), css=CSS, mcp_server=True)