ltx/zerogpu_duration.py

"""ZeroGPU reservation estimator, CPU preflight, and user-facing preview."""
from __future__ import annotations

from math import ceil, isfinite

import gradio as gr

from space_config import FRAME_RATE, STANDARD_MAX_SECONDS
from .app_config import (
    EFFECTIVE_AUTO_DURATION_MAX_SECONDS,
    IS_FULL_SFT_PROFILE,
)
from .conditioning import mode_from_condition_images
from .ui_controls import frames_from_seconds, parse_resolution_value

# Explicit policy coefficients. Runtime evidence records estimator revision and environment identity.
DURATION_ESTIMATOR_REVISION = "p34r0-v2"
FULL_SFT_DURATION_25F_FLOOR_SECONDS = 90
FULL_SFT_DURATION_LONG_FLOOR_SECONDS = 110
FULL_SFT_DURATION_SHORT_FIXED_SECONDS = 24.0
FULL_SFT_DURATION_SHORT_SECONDS_PER_VIDEO_SECOND = 2.5
FULL_SFT_DURATION_LONG_FIXED_SECONDS = 16.0
FULL_SFT_DURATION_LONG_SECONDS_PER_VIDEO_SECOND = 10.2
FULL_SFT_DURATION_LONG_GUARD_RATIO = 1.10
FULL_SFT_DURATION_LONG_GUARD_SECONDS = 5.0
# P29 long center already includes the measured Conv-VAE tail. P33 measured the
# 721f Diffusion Decoder tail on the same P28R4 two-step contract. After
# normalizing P33 to P28R4's cold Stage-2 adapter transfer, the decoder path adds
# ~26.1s over Conv VAE at 30s. Keep the historical 6s fixed term and scale the
# incremental temporal cost at 2/3s per realized video second.
FULL_SFT_DURATION_LONG_DIFFUSION_DECODER_FIXED_SECONDS = 6.0
FULL_SFT_DURATION_LONG_DIFFUSION_DECODER_SECONDS_PER_VIDEO_SECOND = 2.0 / 3.0
DURATION_LORA_ALLOWANCE_SECONDS = 8.0
DURATION_CONDITION_ALLOWANCE_SECONDS = {
    "T2V": 0.0, "I2V": 2.0, "FLF2V": 4.0, "KEYFRAME": 4.0, "KEYFRAME_FLF2V": 6.0,
}
DISTILLED_DURATION_SHORT_FIXED_SECONDS = 8.0
DISTILLED_DURATION_SHORT_SECONDS_PER_VIDEO_SECOND = 1.2
DISTILLED_DURATION_LONG_FIXED_SECONDS = 18.0
DISTILLED_DURATION_LONG_SECONDS_PER_VIDEO_SECOND = 1.8


class DurationEstimateError(ValueError):
    """A user/config input prevented a safe ZeroGPU reservation estimate."""


def estimate_zerogpu_duration_seconds(
    start_image_path, middle_image_path, end_image_path, duration_seconds, resolution_key, selected_loras,
    use_diffusion_decoder, use_auto_duration,
) -> int:
    """Strict estimator shared by UI, CPU preflight, and the ZeroGPU scheduler callable."""
    try:
        width, height = parse_resolution_value(resolution_key)
        reservation_seconds = (
            EFFECTIVE_AUTO_DURATION_MAX_SECONDS if bool(use_auto_duration) else float(duration_seconds)
        )
    except Exception as exc:
        raise DurationEstimateError(f"invalid duration/resolution input: {exc}") from exc
    if not isfinite(reservation_seconds) or reservation_seconds <= 0:
        raise DurationEstimateError(f"duration must be a finite positive value, got {reservation_seconds!r}")
    frames = frames_from_seconds(reservation_seconds)
    realized_seconds = (frames - 1) / FRAME_RATE
    area_ratio = max(1.0, (width * height) / float(512 * 512))
    lora_count = len(selected_loras or [])
    mode = mode_from_condition_images(start_image_path, middle_image_path, end_image_path)
    if mode == "INVALID_CONDITION_WITHOUT_START":
        raise DurationEstimateError("Middle/End frame requires a Start frame")
    if middle_image_path and bool(use_auto_duration):
        raise DurationEstimateError("Middle keyframe requires Manual Duration so its midpoint is known before Stage 1")
    condition_overhead = DURATION_CONDITION_ALLOWANCE_SECONDS.get(mode)
    if condition_overhead is None:
        raise DurationEstimateError(f"unsupported conditioning mode for duration estimate: {mode}")
    lora_overhead = DURATION_LORA_ALLOWANCE_SECONDS * lora_count
    # The distilled profile keeps the P0R1 absolute decoder calibration. For
    # Full/SFT long requests, the P29 center already contains the Conv-VAE tail,
    # so only the measured Diffusion-vs-Conv increment belongs on top.
    distilled_decoder_overhead = (6.0 + realized_seconds) if bool(use_diffusion_decoder) else 0.0
    full_sft_decoder_increment = (
        FULL_SFT_DURATION_LONG_DIFFUSION_DECODER_FIXED_SECONDS
        + (FULL_SFT_DURATION_LONG_DIFFUSION_DECODER_SECONDS_PER_VIDEO_SECOND * realized_seconds)
        if bool(use_diffusion_decoder) else 0.0
    )
    if IS_FULL_SFT_PROFILE:
        short_estimate = (
            FULL_SFT_DURATION_SHORT_FIXED_SECONDS
            + (FULL_SFT_DURATION_SHORT_SECONDS_PER_VIDEO_SECOND * realized_seconds * area_ratio)
            + condition_overhead + lora_overhead + full_sft_decoder_increment
        )
        short_requested = int(ceil(short_estimate))
        if frames <= 25:
            return int(max(FULL_SFT_DURATION_25F_FLOOR_SECONDS, short_requested))
        long_center = (
            FULL_SFT_DURATION_LONG_FIXED_SECONDS
            + (FULL_SFT_DURATION_LONG_SECONDS_PER_VIDEO_SECOND * realized_seconds * area_ratio)
            + condition_overhead + lora_overhead + full_sft_decoder_increment
        )
        long_requested = int(ceil(
            (FULL_SFT_DURATION_LONG_GUARD_RATIO * long_center) + FULL_SFT_DURATION_LONG_GUARD_SECONDS
        ))
        return int(max(FULL_SFT_DURATION_LONG_FLOOR_SECONDS, short_requested, long_requested))
    if realized_seconds > STANDARD_MAX_SECONDS:
        estimate = (
            DISTILLED_DURATION_LONG_FIXED_SECONDS
            + (DISTILLED_DURATION_LONG_SECONDS_PER_VIDEO_SECOND * realized_seconds * area_ratio)
            + condition_overhead + lora_overhead + distilled_decoder_overhead
        )
    else:
        estimate = (
            DISTILLED_DURATION_SHORT_FIXED_SECONDS
            + (DISTILLED_DURATION_SHORT_SECONDS_PER_VIDEO_SECOND * realized_seconds * area_ratio)
            + condition_overhead + lora_overhead + distilled_decoder_overhead
        )
    return int(max(10, ceil(estimate)))


def generation_duration(prompt, start_image_path, middle_image_path, end_image_path, duration_seconds, experimental_long, resolution_key, seed, randomize_seed,
                        selected_loras, lora_strength, custom_loras, prepared_loras, use_diffusion_decoder, use_auto_duration, session_id, history, *args) -> int:  # noqa: ARG001
    """ZeroGPU scheduler callable. Strict by design; CPU preflight runs first in product UI."""
    del prompt, experimental_long, seed, randomize_seed, lora_strength, custom_loras, prepared_loras, session_id, history, args
    return estimate_zerogpu_duration_seconds(
        start_image_path, middle_image_path, end_image_path, duration_seconds, resolution_key, selected_loras,
        use_diffusion_decoder, use_auto_duration,
    )


def generation_duration_preflight(
    start_image_path, middle_image_path, end_image_path, duration_seconds, resolution_key, selected_loras,
    use_diffusion_decoder, use_auto_duration,
):
    try:
        requested = estimate_zerogpu_duration_seconds(
            start_image_path, middle_image_path, end_image_path, duration_seconds, resolution_key, selected_loras,
            use_diffusion_decoder, use_auto_duration,
        )
    except DurationEstimateError as exc:
        raise gr.Error(
            "ZeroGPU duration could not be estimated safely. GPU quota was not requested. "
            f"Check Duration/Resolution/input mode and try again. Details: {exc}"
        ) from exc
    return f"Preflight ready · ZeroGPU reservation **{requested}s**. GPU request starts next."


def zerogpu_duration_panel(start_image_path, middle_image_path, end_image_path, duration_seconds, experimental_long, resolution_key, selected_loras, use_diffusion_decoder, use_auto_duration):  # noqa: ARG001
    try:
        width, height = parse_resolution_value(resolution_key)
    except Exception:
        return "⚠️ Resolution must be entered as WIDTH × HEIGHT with both dimensions multiples of 64."
    reservation_seconds = EFFECTIVE_AUTO_DURATION_MAX_SECONDS if bool(use_auto_duration) else float(duration_seconds)
    frames = frames_from_seconds(reservation_seconds)
    realized_seconds = (frames - 1) / FRAME_RATE
    try:
        requested = estimate_zerogpu_duration_seconds(
            start_image_path, middle_image_path, end_image_path, duration_seconds, resolution_key, selected_loras,
            use_diffusion_decoder, use_auto_duration,
        )
    except DurationEstimateError as exc:
        return (
            "⚠️ **ZeroGPU duration estimate unavailable.** "
            f"GPU generation will stop in CPU preflight instead of requesting a fallback duration. Details: `{exc}`"
        )
    mode = mode_from_condition_images(start_image_path, middle_image_path, end_image_path)
    lora_count = len(selected_loras or [])
    mode_note = (
        "⚠️ Middle/End frame requires a Start frame; generation will be rejected until a Start frame is supplied."
        if mode == "INVALID_CONDITION_WITHOUT_START" else f"Input mode: **{mode}**."
    )
    if requested <= 60:
        quota_note = "tight/ordinary reservation range for this app"
    elif requested <= 120:
        quota_note = "heavy reservation — quota impact is significant"
    elif requested <= 300:
        quota_note = "very heavy reservation — may consume a large share of a visitor's daily quota"
    else:
        quota_note = "aggressive reservation — visitor tier caps or remaining quota may reject it before execution"
    full_sft_probe_note = ""
    if IS_FULL_SFT_PROFILE and frames > 25:
        full_sft_probe_note = (
            "  \n**Full/SFT >25f reservation model:** P27/P28 Stage-1 scaling + P28R4 cold Conv-VAE tail define the base model; P33's 721f Diffusion Decoder witness now calibrates decoder-vs-Conv incremental cost. "
            "The 110s value is only the minimum floor; long requests scale above it. Runtime evidence is hardware-specific; unverified combinations remain available rather than being product-blocked."
        )
    long_note = ""
    if realized_seconds > STANDARD_MAX_SECONDS:
        long_note = (
            "  \n⚠️ **Experimental video length.** The 30s / 721f endpoint is live-passed for distilled T2V and I2V at 512×512; "
            "long durations remain opt-in and quota/VRAM cost rises with length."
        )
    flf_note = ""
    if mode == "FLF2V":
        flf_note = (
            "  \n⚠️ **FLF2V remains experimental.** 512×512/25f current-NF4 has a returned ZeroGPU 48GB PASS; "
            "longer durations and other FLF2V resolutions remain unvalidated."
        )
    keyframe_note = ""
    if middle_image_path:
        keyframe_note = (
            "  \n⚠️ **Middle keyframe is experimental.** It is fixed at the nearest latent-grid midpoint of Manual Duration; "
            "Auto Duration is intentionally rejected before GPU reservation while a Middle image is present."
        )
    vram_note = ""
    if width == 512 and height == 512 and mode == "T2V":
        predicted_gib = 26.5189 + (0.0254508 * frames)
        vram_note = (
            f"  \nRough 512×512 T2V peak-allocation extrapolation: **{predicted_gib:.1f} GiB** "
            "(measured linear trend through 241f; not a guarantee)."
        )
    duration_contract = (
        f"Auto Duration · reservation basis {EFFECTIVE_AUTO_DURATION_MAX_SECONDS:.1f}s cap"
        if bool(use_auto_duration) else f"Manual · {frames}f · {realized_seconds:.2f}s video"
    )
    auto_note = (
        "  \nAuto Duration is predicted only after GPU allocation; ZeroGPU therefore reserves against the configured max cap, "
        "not the eventual model-predicted length. The actual output is snapped to the upstream 8k+1 grid and may be shorter than the cap."
        if bool(use_auto_duration) else ""
    )
    return (
        f"**Estimated ZeroGPU request: {requested} seconds** — {quota_note}.  \n"
        f"{mode_note} Contract: **{width}×{height} · {duration_contract} · {lora_count} LoRA(s) · decoder {'Diffusion' if use_diffusion_decoder else 'Conv VAE'}**.  \n"
        "This is the requested GPU reservation/quota budget, not the video duration or a guaranteed wall-clock runtime."
        f"{auto_note}{full_sft_probe_note}{long_note}{flf_note}{keyframe_note}{vram_note}"
    )