app.py

import gc
import os
import subprocess
import tempfile
from dataclasses import dataclass
from fractions import Fraction

import av
import gradio as gr
import numpy as np
import spaces
import torch
from PIL import Image

from diffusers import LTX2InContextPipeline
from diffusers.pipelines.ltx2.pipeline_ltx2_condition import LTX2VideoCondition
from diffusers.pipelines.ltx2.pipeline_ltx2_ic_lora import LTX2ReferenceCondition
from diffusers.pipelines.ltx2.utils import (
    DEFAULT_NEGATIVE_PROMPT,
    DISTILLED_SIGMA_VALUES,
)

MODEL_ID = "diffusers/LTX-2.3-Distilled-Diffusers"
MODEL_REVISION = "432e0d3c2d1769aaa4d295f9243f7062bf6b47ee"
LORA_ID = "Zlikwid/LTX_2.3_Upscale_IC_Lora"
LORA_WEIGHT = "ltx2.3_upscale_ic-lora_06250.safetensors"
REFERENCE_DOWNSCALE_FACTOR = 2

DISTILLED_STEPS = 8
DISTILLED_GUIDANCE_SCALE = 1.0

MAX_INPUT_DURATION_SECONDS = 12.0
MAX_UPLOAD_BYTES = 200 * 1024 * 1024
MAX_FRAME_COUNT = 65
MAX_FRAME_PIXELS = 1024 * 576
MAX_WORK_UNITS = MAX_FRAME_PIXELS * MAX_FRAME_COUNT
FALLBACK_DECODE_FRAME_LIMIT = 900

INPUT_FACTOR_PRESET = "Source × Bicubic Factor"
BICUBIC_RESAMPLE = (
    Image.Resampling.BICUBIC if hasattr(Image, "Resampling") else Image.BICUBIC
)

RESOLUTION_PRESETS = {
    INPUT_FACTOR_PRESET: (None, None),
    "768×512  (3:2)": (768, 512),
    "512×768  (2:3)": (512, 768),
    "1024×576 (16:9)": (1024, 576),
    "576×1024 (9:16)": (576, 1024),
    "768×768  (1:1)": (768, 768),
    "512×512  (1:1)": (512, 512),
    "960×544  (16:9)": (960, 544),
    "544×960  (9:16)": (544, 960),
    "Custom": (None, None),
}

PROCESSING_PRESETS = {
    "Preview": {
        "resolution": "512×512  (1:1)",
        "max_frames": 17,
        "lora_strength": 0.45,
        "input_preservation_strength": 0.90,
    },
    "Balanced": {
        "resolution": "768×512  (3:2)",
        "max_frames": 41,
        "lora_strength": 0.55,
        "input_preservation_strength": 0.82,
    },
    "Quality": {
        "resolution": "768×768  (1:1)",
        "max_frames": 65,
        "lora_strength": 0.65,
        "input_preservation_strength": 0.75,
    },
}


print("Loading distilled LTX-2.3 pipeline...")
pipe = LTX2InContextPipeline.from_pretrained(
    MODEL_ID,
    revision=MODEL_REVISION,
    torch_dtype=torch.bfloat16,
)
print("Loading upscale IC-LoRA...")
pipe.load_lora_weights(
    LORA_ID,
    adapter_name="upscale",
    weight_name=LORA_WEIGHT,
)
pipe.vae.enable_tiling()
print("Pipeline ready!")

_offload_ready = False


@dataclass
class VideoMeta:
    path: str
    width: int
    height: int
    duration: float | None
    fps: float | None
    frame_count: int | None
    has_audio: bool
    file_size: int


def video_path_from_input(input_video) -> str:
    if isinstance(input_video, dict):
        path = input_video.get("video") or input_video.get("path")
    elif isinstance(input_video, (list, tuple)):
        path = input_video[0] if input_video else None
    else:
        path = input_video

    if not path:
        raise gr.Error("Please upload a video first.")
    return str(path)


def round_to_8k_plus_1(n: int) -> int:
    n = int(n)
    k = max(0, (n - 1) // 8)
    return max(9, k * 8 + 1)


def snap_to_32(n: int) -> int:
    return max(32, round(n / 32) * 32)


def format_duration(seconds: float | None) -> str:
    if seconds is None:
        return "unknown"
    return f"{seconds:.2f}s"


def format_fps(fps: float | None) -> str:
    if fps is None:
        return "unknown"
    return f"{fps:.2f} fps"


def update_stage_progress(progress, start, end, index, total, desc):
    if total <= 0:
        progress(end, desc=desc)
        return
    fraction = start + (end - start) * (index + 1) / total
    progress(fraction, desc=f"{desc} {index + 1}/{total}")


def cleanup_cuda():
    gc.collect()
    if torch.cuda.is_available():
        torch.cuda.empty_cache()


def ensure_cpu_offload():
    global _offload_ready
    if not _offload_ready:
        pipe.enable_model_cpu_offload()
        _offload_ready = True


def probe_video(input_video) -> VideoMeta:
    path = video_path_from_input(input_video)
    if not os.path.exists(path):
        raise gr.Error(f"Input video does not exist: {path}")

    file_size = os.path.getsize(path)
    try:
        with av.open(path) as container:
            if not container.streams.video:
                raise gr.Error("Uploaded file has no video stream.")

            stream = container.streams.video[0]
            fps = float(stream.average_rate) if stream.average_rate else None

            duration = None
            if stream.duration is not None and stream.time_base is not None:
                duration = float(stream.duration * stream.time_base)
            elif container.duration is not None:
                duration = float(container.duration / 1_000_000)

            frame_count = int(stream.frames) if stream.frames else None
            if frame_count is None and duration and fps:
                frame_count = max(1, int(round(duration * fps)))

            width = int(stream.codec_context.width or stream.width)
            height = int(stream.codec_context.height or stream.height)
            has_audio = bool(container.streams.audio)
    except gr.Error:
        raise
    except Exception as e:
        raise gr.Error(f"Failed to inspect video: {e}")

    return VideoMeta(
        path=path,
        width=width,
        height=height,
        duration=duration,
        fps=fps,
        frame_count=frame_count,
        has_audio=has_audio,
        file_size=file_size,
    )


def resolve_output_size(
    meta: VideoMeta,
    resolution_preset: str,
    output_width,
    output_height,
    zoom,
    bicubic_upscale_factor,
) -> tuple[int, int, int, int]:
    zoom = float(zoom)
    bicubic_upscale_factor = float(bicubic_upscale_factor)

    source_width = max(32, int(meta.width / zoom)) if zoom > 1.0 else meta.width
    source_height = max(32, int(meta.height / zoom)) if zoom > 1.0 else meta.height

    if resolution_preset == INPUT_FACTOR_PRESET:
        width = snap_to_32(int(source_width * bicubic_upscale_factor))
        height = snap_to_32(int(source_height * bicubic_upscale_factor))
    elif resolution_preset == "Custom":
        width = snap_to_32(int(output_width))
        height = snap_to_32(int(output_height))
    else:
        width, height = RESOLUTION_PRESETS[resolution_preset]

    return width, height, source_width, source_height


def validate_plan(meta: VideoMeta, output_width: int, output_height: int, num_frames: int):
    if meta.file_size > MAX_UPLOAD_BYTES:
        raise gr.Error(
            f"Input file is {meta.file_size / 1024 / 1024:.1f} MB. "
            f"Limit is {MAX_UPLOAD_BYTES / 1024 / 1024:.0f} MB."
        )

    if meta.duration and meta.duration > MAX_INPUT_DURATION_SECONDS:
        raise gr.Error(
            f"Input duration is {meta.duration:.1f}s. "
            f"Limit is {MAX_INPUT_DURATION_SECONDS:.0f}s for this ZeroGPU Space."
        )

    if meta.frame_count is not None and meta.frame_count < 9:
        raise gr.Error(
            f"Video has only {meta.frame_count} frames. Need at least 9 (8k+1)."
        )

    if num_frames > MAX_FRAME_COUNT:
        raise gr.Error(f"Frame count {num_frames} exceeds limit {MAX_FRAME_COUNT}.")

    pixels = output_width * output_height
    if pixels > MAX_FRAME_PIXELS:
        raise gr.Error(
            f"Output {output_width}x{output_height} is too large for zero-a10g. "
            f"Limit is about {MAX_FRAME_PIXELS:,} pixels per frame."
        )

    work_units = pixels * num_frames
    if work_units > MAX_WORK_UNITS:
        raise gr.Error(
            f"Requested {num_frames} frames at {output_width}x{output_height}. "
            "Reduce resolution or frame count."
        )


def planned_frame_count(meta: VideoMeta, max_frames) -> int:
    max_frames = min(MAX_FRAME_COUNT, int(max_frames))
    available_frames = meta.frame_count or max_frames
    return round_to_8k_plus_1(min(max_frames, available_frames))


def output_fps_for_duration(meta: VideoMeta, num_frames: int) -> float:
    if meta.duration and meta.duration > 0:
        return max(0.1, (num_frames - 1) / meta.duration)
    if meta.fps:
        return meta.fps
    return 24.0


def encoded_fps_for_duration(meta: VideoMeta, frame_count: int) -> float:
    if meta.duration and meta.duration > 0:
        return max(0.1, frame_count / meta.duration)
    if meta.fps:
        return meta.fps
    return 24.0


def ffmpeg_rate_arg(fps: float) -> str:
    rate = Fraction(float(fps)).limit_denominator(100000)
    return f"{rate.numerator}/{rate.denominator}"


def encoded_rate_arg_for_duration(
    meta: VideoMeta,
    frame_count: int,
    fallback_fps: float,
) -> str:
    if meta.duration and meta.duration > 0:
        duration = Fraction(float(meta.duration)).limit_denominator(10000)
        if duration > 0:
            rate = Fraction(frame_count, 1) / duration
            return f"{rate.numerator}/{rate.denominator}"
    return ffmpeg_rate_arg(fallback_fps)


def normalize_output_frame(frame) -> np.ndarray:
    arr = np.asarray(frame)

    if arr.ndim != 3:
        raise gr.Error(f"Unexpected output frame shape: {arr.shape}")
    if arr.shape[0] in (1, 3, 4) and arr.shape[-1] not in (1, 3, 4):
        arr = np.moveaxis(arr, 0, -1)
    if arr.shape[-1] == 4:
        arr = arr[..., :3]
    elif arr.shape[-1] == 1:
        arr = np.repeat(arr, 3, axis=-1)
    elif arr.shape[-1] != 3:
        raise gr.Error(f"Unexpected output frame channels: {arr.shape}")

    if arr.dtype != np.uint8:
        arr = np.nan_to_num(arr)
        if arr.max(initial=0.0) <= 1.0 and arr.min(initial=0.0) >= 0.0:
            arr = arr * 255.0
        arr = np.clip(arr, 0, 255).astype(np.uint8)

    return np.ascontiguousarray(arr)


def encode_video_frames(frames, fps: float, rate_arg: str, output_path: str):
    frame_count = len(frames)
    if frame_count <= 0:
        raise gr.Error("Pipeline returned no output frames.")

    first_frame = normalize_output_frame(frames[0])
    height, width = first_frame.shape[:2]
    raw_frames = [first_frame.tobytes()]
    for frame in frames[1:]:
        frame_arr = normalize_output_frame(frame)
        if frame_arr.shape[:2] != (height, width):
            raise gr.Error(
                f"Output frame size changed from {width}x{height} "
                f"to {frame_arr.shape[1]}x{frame_arr.shape[0]}."
            )
        raw_frames.append(frame_arr.tobytes())
    raw_video = b"".join(raw_frames)

    def command(use_demux_time_base: bool) -> list[str]:
        cmd = [
            "ffmpeg",
            "-y",
            "-loglevel",
            "error",
            "-f",
            "rawvideo",
            "-pix_fmt",
            "rgb24",
            "-s:v",
            f"{width}x{height}",
            "-framerate",
            rate_arg,
            "-i",
            "pipe:0",
            "-an",
            "-c:v",
            "libx264",
            "-preset",
            "medium",
            "-crf",
            "18",
            "-pix_fmt",
            "yuv420p",
        ]
        if use_demux_time_base:
            cmd.extend(["-enc_time_base", "demux"])
        cmd.extend([
            "-movflags",
            "+faststart",
            "-video_track_timescale",
            "90000",
            output_path,
        ])
        return cmd

    errors = []
    for use_demux_time_base in (True, False):
        result = subprocess.run(
            command(use_demux_time_base),
            input=raw_video,
            capture_output=True,
        )
        if result.returncode == 0:
            return
        stderr = result.stderr.decode("utf-8", errors="replace")
        errors.append(stderr[-1000:])

    raise gr.Error(
        "Failed to encode output video after ffmpeg retry: "
        + " | ".join(error for error in errors if error)
    )


def decode_sampled_frames(
    meta: VideoMeta,
    num_frames: int,
    progress=gr.Progress(),
) -> list[Image.Image]:
    progress(0.02, desc=f"Decoding {num_frames} sampled frames...")

    if meta.frame_count:
        indices = np.linspace(0, meta.frame_count - 1, num_frames).astype(int).tolist()
        wanted = set(indices)
        selected: dict[int, Image.Image] = {}

        with av.open(meta.path) as container:
            stream = container.streams.video[0]
            stream.thread_type = "AUTO"
            for frame_index, frame in enumerate(container.decode(stream)):
                if frame_index in wanted:
                    selected[frame_index] = frame.to_image().convert("RGB")
                    update_stage_progress(
                        progress,
                        0.02,
                        0.08,
                        len(selected) - 1,
                        len(wanted),
                        "Decoded sampled frames",
                    )
                    if len(selected) == len(wanted):
                        break

        frames = [selected[index] for index in indices if index in selected]
        if len(frames) == num_frames:
            return frames

    decoded = []
    with av.open(meta.path) as container:
        stream = container.streams.video[0]
        stream.thread_type = "AUTO"
        for frame in container.decode(stream):
            decoded.append(frame.to_image().convert("RGB"))
            if len(decoded) > FALLBACK_DECODE_FRAME_LIMIT:
                raise gr.Error("Too many frames to decode safely in this Space.")

    if len(decoded) < 9:
        raise gr.Error(f"Video has only {len(decoded)} frames. Need at least 9 (8k+1).")

    num_frames = round_to_8k_plus_1(min(num_frames, len(decoded)))
    indices = np.linspace(0, len(decoded) - 1, num_frames).astype(int)
    frames = [decoded[index] for index in indices]
    progress(0.08, desc=f"Decoded {num_frames} sampled frames.")
    return frames


def mux_source_audio(source_path: str, video_path: str, meta: VideoMeta, preserve_audio: bool):
    if not preserve_audio or not meta.has_audio:
        return video_path

    output_path = tempfile.mktemp(suffix=".mp4")
    cmd = [
        "ffmpeg",
        "-y",
        "-i",
        video_path,
        "-i",
        source_path,
        "-map",
        "0:v:0",
        "-map",
        "1:a:0",
        "-c:v",
        "copy",
        "-c:a",
        "aac",
        "-shortest",
        output_path,
    ]
    result = subprocess.run(cmd, capture_output=True, text=True)
    if result.returncode != 0:
        raise gr.Error(f"Failed to mux source audio: {result.stderr[-1000:]}")
    return output_path


def estimate_settings(
    input_video,
    resolution_preset,
    output_width,
    output_height,
    zoom,
    bicubic_upscale_factor,
    max_frames,
    input_preservation_strength,
    preserve_audio,
):
    if input_video is None:
        return "Upload a video to see the planned output."

    try:
        meta = probe_video(input_video)
        num_frames = planned_frame_count(meta, max_frames)
        width, height, source_width, source_height = resolve_output_size(
            meta,
            resolution_preset,
            output_width,
            output_height,
            zoom,
            bicubic_upscale_factor,
        )
        validate_plan(meta, width, height, num_frames)
        output_fps = output_fps_for_duration(meta, num_frames)
    except gr.Error as e:
        return f"Warning: {e}"
    except Exception as e:
        return f"Warning: failed to estimate settings: {e}"

    audio_status = "source audio" if preserve_audio and meta.has_audio else "silent/video only"
    return (
        "**Run summary**\n\n"
        f"- Input: {meta.width}x{meta.height}, {format_duration(meta.duration)}, "
        f"{format_fps(meta.fps)}, {meta.frame_count or 'unknown'} frames\n"
        f"- Source after zoom: {source_width}x{source_height}\n"
        f"- Output: {width}x{height}, {num_frames} frames, "
        f"{output_fps:.2f} fps, duration preserved when metadata is available\n"
        f"- Model: distilled LTX 2.3, {DISTILLED_STEPS} steps, CFG {DISTILLED_GUIDANCE_SCALE}\n"
        f"- Input preservation: {float(input_preservation_strength):.2f}\n"
        f"- Audio: {audio_status}"
    )


@spaces.GPU(duration=600)
def upscale_video(
    input_video,
    resolution_preset,
    output_width,
    output_height,
    zoom,
    bicubic_upscale_factor,
    lora_strength,
    input_preservation_strength,
    max_frames,
    seed,
    preserve_audio,
    progress=gr.Progress(track_tqdm=True),
):
    if input_video is None:
        raise gr.Error("Please upload a video first.")

    zoom = float(zoom)
    bicubic_upscale_factor = float(bicubic_upscale_factor)
    lora_strength = float(lora_strength)
    input_preservation_strength = float(input_preservation_strength)
    max_frames = int(max_frames)
    seed = int(seed)
    preserve_audio = bool(preserve_audio)

    progress(0.0, desc="Inspecting input video...")
    meta = probe_video(input_video)
    num_frames = planned_frame_count(meta, max_frames)
    output_width, output_height, _, _ = resolve_output_size(
        meta,
        resolution_preset,
        output_width,
        output_height,
        zoom,
        bicubic_upscale_factor,
    )
    validate_plan(meta, output_width, output_height, num_frames)
    output_fps = output_fps_for_duration(meta, num_frames)

    try:
        frames = decode_sampled_frames(meta, num_frames, progress)

        progress(0.09, desc="Preparing bicubic pre-upscale...")

        if zoom > 1.0:
            crop_w = max(32, int(frames[0].width / zoom))
            crop_h = max(32, int(frames[0].height / zoom))
            left = (frames[0].width - crop_w) // 2
            top = (frames[0].height - crop_h) // 2
            cropped_frames = []
            for index, frame in enumerate(frames):
                cropped_frames.append(
                    frame.crop((left, top, left + crop_w, top + crop_h))
                )
                update_stage_progress(
                    progress, 0.09, 0.12, index, len(frames), "Cropping frames"
                )
            frames = cropped_frames
        else:
            progress(0.12, desc="No crop applied.")

        resized_frames = []
        resize_desc = f"Bicubic pre-upscale to {output_width}x{output_height}"
        for index, frame in enumerate(frames):
            resized_frames.append(
                frame.resize((output_width, output_height), BICUBIC_RESAMPLE)
            )
            update_stage_progress(
                progress, 0.12, 0.20, index, len(frames), resize_desc
            )
        frames = resized_frames

        progress(0.21, desc="Building reference condition...")
        ref_cond = LTX2ReferenceCondition(frames=frames, strength=1.0)
        input_cond = LTX2VideoCondition(
            frames=frames,
            index=0,
            strength=input_preservation_strength,
        )

        progress(0.23, desc="Applying LoRA strength...")
        pipe.set_adapters("upscale", lora_strength)

        progress(0.24, desc="Preparing GPU offload...")
        ensure_cpu_offload()

        generator = torch.Generator("cuda").manual_seed(seed)
        denoise_start = 0.26
        denoise_end = 0.93

        def step_callback(pipeline, step, timestep, kwargs):
            progress(
                denoise_start
                + (denoise_end - denoise_start) * (step + 1) / DISTILLED_STEPS,
                desc=f"Denoising step {step + 1}/{DISTILLED_STEPS}",
            )
            return kwargs

        progress(denoise_start, desc=f"Denoising step 0/{DISTILLED_STEPS}")

        video, _audio = pipe(
            prompt="upscale",
            negative_prompt=DEFAULT_NEGATIVE_PROMPT,
            reference_conditions=[ref_cond],
            conditions=[input_cond],
            reference_downscale_factor=REFERENCE_DOWNSCALE_FACTOR,
            height=output_height,
            width=output_width,
            num_frames=num_frames,
            frame_rate=output_fps,
            num_inference_steps=DISTILLED_STEPS,
            sigmas=DISTILLED_SIGMA_VALUES,
            guidance_scale=DISTILLED_GUIDANCE_SCALE,
            stg_scale=1.0,
            modality_scale=3.0,
            guidance_rescale=0.7,
            spatio_temporal_guidance_blocks=[28],
            use_cross_timestep=True,
            generator=generator,
            output_type="np",
            return_dict=False,
            callback_on_step_end=step_callback,
        )

        progress(0.95, desc="Encoding output video...")
        fd, video_path = tempfile.mkstemp(suffix=".mp4")
        os.close(fd)
        output_frames = video[0]
        actual_frame_count = len(output_frames)
        encoding_fps = encoded_fps_for_duration(meta, actual_frame_count)
        encoding_rate_arg = encoded_rate_arg_for_duration(
            meta, actual_frame_count, encoding_fps
        )
        encode_video_frames(output_frames, encoding_fps, encoding_rate_arg, video_path)

        progress(0.97, desc="Preserving source audio..." if preserve_audio else "Finalizing...")
        output_path = mux_source_audio(meta.path, video_path, meta, preserve_audio)

    except torch.cuda.OutOfMemoryError:
        cleanup_cuda()
        raise gr.Error(
            "GPU out of memory. Reduce resolution or frame count. "
            "The ZeroGPU-safe limit is enforced, but this input still exceeded memory."
        )
    except gr.Error:
        cleanup_cuda()
        raise
    except Exception as e:
        cleanup_cuda()
        raise gr.Error(f"Pipeline error: {e}")

    cleanup_cuda()
    progress(
        1.0,
        desc=(
            f"Done: {output_width}x{output_height}, "
            f"{actual_frame_count} output frames, {encoding_fps:.2f} fps."
        ),
    )
    return output_path


with gr.Blocks(
    title="LTX 2.3 Distilled Upscale IC-LoRA",
    theme=gr.themes.Soft(primary_hue="blue", secondary_hue="purple"),
) as demo:
    gr.Markdown(
        """
        # LTX 2.3 Distilled Video Upscaler

        Upload a short video to bicubic pre-upscale and refine it with the
        Zlikwid LTX 2.3 Upscale IC-LoRA.
        """
    )

    with gr.Row():
        with gr.Column(scale=1):
            input_video = gr.Video(label="Input Video", sources=["upload"])
            run_summary = gr.Markdown("Upload a video to see the planned output.")

            with gr.Accordion("Settings", open=True):
                processing_preset = gr.Radio(
                    choices=list(PROCESSING_PRESETS.keys()),
                    value="Balanced",
                    label="Processing Preset",
                )
                resolution_preset = gr.Dropdown(
                    choices=list(RESOLUTION_PRESETS.keys()),
                    value="768×512  (3:2)",
                    label="Output Resolution",
                )
                with gr.Row():
                    output_width = gr.Slider(
                        64, 1024, value=768, step=32,
                        label="Width (px)", visible=False,
                    )
                    output_height = gr.Slider(
                        64, 1024, value=512, step=32,
                        label="Height (px)", visible=False,
                    )

                zoom = gr.Slider(
                    1.0, 4.0, value=1.0, step=0.1,
                    label="Zoom",
                    info="1.0 keeps the full frame; higher values crop the center.",
                )
                bicubic_upscale_factor = gr.Slider(
                    1.0, 4.0, value=2.0, step=0.25,
                    label="Bicubic Pre-Upscale Factor",
                    info="Used when Output Resolution is Source × Bicubic Factor.",
                    visible=False,
                )
                lora_strength = gr.Slider(
                    0.0, 1.0, value=0.55, step=0.05,
                    label="Refinement Strength (LoRA)",
                    info="Higher values add more generative detail; lower values stay closer to input.",
                )
                input_preservation_strength = gr.Slider(
                    0.0, 1.0, value=0.82, step=0.05,
                    label="Input Preservation Strength",
                    info="Higher values anchor the output to the bicubic input more strongly.",
                )
                max_frames = gr.Slider(
                    9, MAX_FRAME_COUNT, value=41, step=8,
                    label="Sampled Frames (8k+1)",
                )
                with gr.Row():
                    seed = gr.Number(
                        value=42, label="Seed", precision=0,
                    )
                    randomize_seed = gr.Checkbox(
                        label="Randomize", value=False,
                    )
                preserve_audio = gr.Checkbox(
                    label="Preserve Source Audio",
                    value=True,
                )

            generate_btn = gr.Button(
                "Upscale Video", variant="primary", size="lg",
            )

        with gr.Column(scale=1):
            output_video = gr.Video(label="Output Video")

    def controls_for_resolution(preset):
        if preset == "Custom":
            return (
                gr.update(visible=True),
                gr.update(visible=True),
                gr.update(visible=False),
            )
        if preset == INPUT_FACTOR_PRESET:
            return (
                gr.update(visible=False),
                gr.update(visible=False),
                gr.update(visible=True),
            )
        return (
            gr.update(visible=False),
            gr.update(visible=False),
            gr.update(visible=False),
        )

    def toggle_custom_resolution(preset):
        return controls_for_resolution(preset)

    resolution_preset.change(
        toggle_custom_resolution,
        inputs=[resolution_preset],
        outputs=[output_width, output_height, bicubic_upscale_factor],
        queue=False,
        show_progress="hidden",
    )

    def apply_processing_preset(
        preset,
        input_video_value,
        output_width_value,
        output_height_value,
        zoom_value,
        bicubic_upscale_factor_value,
        preserve_audio_value,
    ):
        values = PROCESSING_PRESETS[preset]
        resolution = values["resolution"]
        width_update, height_update, factor_update = controls_for_resolution(resolution)
        summary = estimate_settings(
            input_video_value,
            resolution,
            output_width_value,
            output_height_value,
            zoom_value,
            bicubic_upscale_factor_value,
            values["max_frames"],
            values["input_preservation_strength"],
            preserve_audio_value,
        )
        return (
            gr.update(value=resolution),
            width_update,
            height_update,
            factor_update,
            gr.update(value=values["max_frames"]),
            gr.update(value=values["lora_strength"]),
            gr.update(value=values["input_preservation_strength"]),
            summary,
        )

    processing_preset.change(
        apply_processing_preset,
        inputs=[
            processing_preset,
            input_video,
            output_width,
            output_height,
            zoom,
            bicubic_upscale_factor,
            preserve_audio,
        ],
        outputs=[
            resolution_preset,
            output_width,
            output_height,
            bicubic_upscale_factor,
            max_frames,
            lora_strength,
            input_preservation_strength,
            run_summary,
        ],
        queue=False,
        show_progress="hidden",
    )

    summary_inputs = [
        input_video,
        resolution_preset,
        output_width,
        output_height,
        zoom,
        bicubic_upscale_factor,
        max_frames,
        input_preservation_strength,
        preserve_audio,
    ]
    for component in [
        input_video,
        resolution_preset,
        zoom,
        bicubic_upscale_factor,
        max_frames,
        input_preservation_strength,
        preserve_audio,
        output_width,
        output_height,
    ]:
        component.change(
            estimate_settings,
            inputs=summary_inputs,
            outputs=[run_summary],
            queue=False,
            show_progress="hidden",
        )

    def maybe_randomize(seed_val, do_randomize):
        if do_randomize:
            return int(np.random.randint(0, 2**32))
        return seed_val

    generate_btn.click(
        maybe_randomize,
        inputs=[seed, randomize_seed],
        outputs=[seed],
        queue=False,
        show_progress="hidden",
    ).then(
        upscale_video,
        inputs=[
            input_video,
            resolution_preset,
            output_width,
            output_height,
            zoom,
            bicubic_upscale_factor,
            lora_strength,
            input_preservation_strength,
            max_frames,
            seed,
            preserve_audio,
        ],
        outputs=[output_video],
        concurrency_limit=1,
    )

demo.queue(default_concurrency_limit=1, max_size=4)

if __name__ == "__main__":
    demo.launch()