ltx/ui_controls.py

"""Small reusable UI helpers for product controls, resolution, and duration."""
from __future__ import annotations

import gradio as gr
import PIL.Image

from space_config import EXPERIMENTAL_MAX_SECONDS, FRAME_RATE, RESOLUTIONS, STANDARD_MAX_SECONDS
from . import app_helpers
from .app_config import EFFECTIVE_DEFAULT_DURATION_SECONDS, MAX_FRAMES


def parse_resolution_value(value: str) -> tuple[int, int]:
    return app_helpers.parse_resolution_value(value, resolutions=RESOLUTIONS)


def frames_from_seconds(seconds: float) -> int:
    return app_helpers.frames_from_seconds(
        seconds,
        frame_rate=FRAME_RATE,
        experimental_max_seconds=EXPERIMENTAL_MAX_SECONDS,
        max_frames=MAX_FRAMES,
    )


def supports_experimental_long(resolution_key: str) -> bool:
    return app_helpers.supports_experimental_long(resolution_key, resolutions=RESOLUTIONS)


def duration_slider_update(experimental_long, duration_seconds, resolution_key):
    allow_long = bool(experimental_long) and supports_experimental_long(resolution_key)
    maximum = EXPERIMENTAL_MAX_SECONDS if allow_long else STANDARD_MAX_SECONDS
    try:
        value = min(maximum, max(1.0, float(duration_seconds)))
    except Exception:
        value = min(maximum, EFFECTIVE_DEFAULT_DURATION_SECONDS)
    return gr.Slider(
        minimum=1.0,
        maximum=maximum,
        step=0.5,
        value=value,
        label="Duration",
        info=(
            f"1–{STANDARD_MAX_SECONDS:g}s standard; {STANDARD_MAX_SECONDS:g}–{EXPERIMENTAL_MAX_SECONDS:g}s "
            f"experimental at 512×512 only. Rounded to the required 8k+1 frame grid at {FRAME_RATE:g} fps."
        ),
    )


def resolution_duration_controls(resolution_key, experimental_long, duration_seconds):
    long_capable = supports_experimental_long(resolution_key)
    enabled = bool(experimental_long) and long_capable
    checkbox = gr.Checkbox(
        value=enabled,
        interactive=long_capable,
        label=f"Experimental long duration (>{STANDARD_MAX_SECONDS:g}s)",
        info=(
            f"Allows {STANDARD_MAX_SECONDS:g}–{EXPERIMENTAL_MAX_SECONDS:g}s only at 512×512. "
            "The 30s / 721f endpoint is live-passed for distilled T2V and I2V."
        ),
    )
    return checkbox, duration_slider_update(enabled, duration_seconds, resolution_key)


def resolution(value: str) -> tuple[int, int]:
    try:
        return parse_resolution_value(value)
    except Exception as exc:
        raise gr.Error(
            f"Invalid resolution `{value}`. Use WIDTH × HEIGHT with both dimensions multiples of 64."
        ) from exc


def _resolution_dimension(value, label: str) -> int:
    if value is None:
        raise gr.Error(f"{label} is required.")
    try:
        numeric = float(value)
        dimension = int(numeric)
    except Exception as exc:
        raise gr.Error(f"{label} must be an integer multiple of 64.") from exc
    if numeric != dimension or dimension < 64 or dimension % 64 != 0:
        raise gr.Error(f"{label} must be an integer multiple of 64 (64, 128, 192, ...).")
    return dimension


def apply_custom_resolution(width, height):
    """Validate explicit dimensions and return the canonical Resolution dropdown value."""
    width = _resolution_dimension(width, "Width")
    height = _resolution_dimension(height, "Height")
    for label, dimensions in RESOLUTIONS.items():
        if dimensions == (width, height):
            return label, f"Using preset **{width}×{height}**."
    return f"{width} × {height}", f"Using custom **{width}×{height}** · unvalidated; large sizes may be slow or OOM."


def resolution_editor_update(value: str):
    """Synchronize explicit Width/Height fields from the active Resolution value."""
    width, height = resolution(value)
    preset = any(dimensions == (width, height) for dimensions in RESOLUTIONS.values())
    status = (
        f"Preset **{width}×{height}**."
        if preset else f"Custom **{width}×{height}** · unvalidated; large sizes may be slow or OOM."
    )
    return width, height, status


def resolution_input_preflight(value: str):
    """Reject malformed resolution before LoRA acquisition or GPU scheduling begins."""
    width, height = resolution(value)
    return f"Input preflight ready · **{width}×{height}** resolution is on the required 64-pixel grid."


def load_conditioning_image(path: str, width: int, height: int) -> PIL.Image.Image:
    img = PIL.Image.open(path).convert("RGB")
    scale = max(width / img.width, height / img.height)
    img = img.resize((round(img.width * scale), round(img.height * scale)), PIL.Image.LANCZOS)
    left, top = (img.width - width) // 2, (img.height - height) // 2
    return img.crop((left, top, left + width, top + height))


def mode_from_images(start_image_path, end_image_path) -> str:
    if end_image_path and not start_image_path:
        return "INVALID_END_ONLY"
    if start_image_path and end_image_path:
        return "FLF2V"
    if start_image_path:
        return "I2V"
    return "T2V"




def bind_custom_resolution_editor(
    *, resolution_component, width_input, height_input, apply_button, status_output,
    experimental_long, duration_seconds, api_visibility,
):
    """Bind the small custom-resolution editor without leaking product generation wiring here."""
    resolution_component.blur(
        resolution_editor_update, [resolution_component], [width_input, height_input, status_output],
        queue=False, show_progress="hidden", api_visibility=api_visibility,
    )
    apply_event = apply_button.click(
        apply_custom_resolution, [width_input, height_input], [resolution_component, status_output],
        queue=False, show_progress="hidden", api_visibility=api_visibility,
    )
    return apply_event.then(
        resolution_duration_controls, [resolution_component, experimental_long, duration_seconds],
        [experimental_long, duration_seconds], queue=False, show_progress="hidden", api_visibility=api_visibility,
    )


def sync_resolution_editor_after(events, *, resolution_component, width_input, height_input, status_output, api_visibility):
    """Keep helper dimensions aligned after settings/history restore events."""
    for event in events:
        event.then(
            resolution_editor_update, [resolution_component], [width_input, height_input, status_output],
            queue=False, show_progress="hidden", api_visibility=api_visibility,
        )


def build_custom_resolution_editor(default_resolution: str):
    """Render explicit Width/Height controls while preserving the canonical Resolution string."""
    width, height = parse_resolution_value(default_resolution)
    with gr.Accordion("Custom resolution", open=False):
        with gr.Row():
            width_input = gr.Number(value=width, minimum=64, step=64, precision=0, label="Width")
            height_input = gr.Number(value=height, minimum=64, step=64, precision=0, label="Height")
            apply_button = gr.Button("Use custom resolution", variant="secondary")
        status = gr.Markdown(f"Preset **{width}×{height}**.", elem_classes=["ltx-subtle"])
        gr.Markdown(
            "LTX-2.5 requires a 64-pixel spatial grid. No product-side upper cap is added here; "
            "unvalidated large dimensions can be slow or run out of VRAM.",
            elem_classes=["ltx-subtle"],
        )
    return width_input, height_input, apply_button, status


def build_seed_controls(*, seed_value: int, randomize_value: bool):
    """Render the standard Seed / Randomize seed row used by product generation tabs."""
    with gr.Row():
        seed = gr.Number(value=seed_value, precision=0, label="Seed")
        randomize_seed = gr.Checkbox(value=randomize_value, label="Randomize seed")
    return seed, randomize_seed


def build_result_panel(*, result_label: str, height: int = 480, probe_note: str | None = None):
    """Render the common result video, used-seed field, and Probe-artifact disclosure."""
    result = gr.Video(label=result_label, autoplay=True, height=height)
    used_seed = gr.Textbox(label="Used seed", interactive=False)
    with gr.Accordion("Run details / Probe artifacts", open=False):
        probe_files = gr.File(
            label="Probe artifacts · one ZIP per request",
            file_count="multiple",
            interactive=False,
        )
        if probe_note:
            gr.Markdown(probe_note, elem_classes=["ltx-subtle"])
    return result, used_seed, probe_files