ltx/generation.py

from __future__ import annotations

import gc
import json
import random
import time
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Callable

import gradio as gr
import torch
from diffusers import FlowMatchEulerDiscreteScheduler
from diffusers.pipelines.ltx2.pipeline_ltx2_condition import LTX2VideoCondition
from diffusers.pipelines.ltx2.utils import (
    DEFAULT_NEGATIVE_PROMPT,
    DISTILLED_SIGMA_VALUES,
    STAGE_2_DISTILLED_SIGMA_VALUES,
)
from diffusers.utils import encode_video

from ltx import probe_artifacts
from ltx.conditioning import keyframe_pixel_frame_index, middle_keyframe_latent_index, uses_condition_pipeline
from ltx.runtime_utils import execution_environment_identity, process_rss_kib, reset_gpu_peak_memory


@dataclass(frozen=True)
class GenerationRuntime:
    candidate_id: str
    diffusion_decode_pipe: Any
    auto_duration_enabled: bool
    auto_duration_max_seconds: float
    auto_duration_min_seconds: float
    full_sft_stage2_lora_strength: float
    experimental_max_seconds: float
    frame_rate: float
    full_sft_stage2_adapter_name: str
    is_full_sft_profile: bool
    pipe: Any
    pipe_condition: Any
    pipe_i2v: Any
    preload_state: dict
    runtime_profile: str
    standard_max_seconds: float
    upsample_pipe: Any
    worker_uuid: str


@dataclass(frozen=True)
class GenerationHooks:
    cleanup_request_loras: Callable
    close_request_logger: Callable
    disk_state: Callable
    duration: Callable
    frames_from_seconds: Callable
    gpu_state: Callable
    history_outputs: Callable
    load_conditioning_image: Callable
    load_request_loras: Callable
    lora_adapter_state: Callable
    mode_from_images: Callable
    open_request_logger: Callable
    resolution: Callable
    settings_snapshot: Callable
    sha256_file: Callable
    supports_experimental_long: Callable
    update_history: Callable
    create_request_paths: Callable


def _build_stage1_call_kwargs(
    *,
    prompt: str,
    generator: torch.Generator,
    mode: str,
    start_image_path,
    middle_image_path,
    end_image_path,
    num_frames: int | None,
    width: int,
    height: int,
    runtime: GenerationRuntime,
    hooks: GenerationHooks,
) -> dict[str, Any]:
    if runtime.is_full_sft_profile:
        call_kwargs = dict(
            prompt=prompt, negative_prompt=DEFAULT_NEGATIVE_PROMPT, frame_rate=runtime.frame_rate,
            guidance_scale=3.0, audio_guidance_scale=7.0, stg_scale=1.0, audio_stg_scale=1.0,
            modality_scale=3.0, audio_modality_scale=3.0, guidance_rescale=0.7, audio_guidance_rescale=0.7,
            spatio_temporal_guidance_blocks=[28], use_cross_timestep=True, generator=generator, return_dict=False,
        )
    else:
        call_kwargs = dict(
            prompt=prompt, negative_prompt=DEFAULT_NEGATIVE_PROMPT, frame_rate=runtime.frame_rate,
            guidance_scale=1.0, audio_guidance_scale=1.0, stg_scale=0.0, audio_stg_scale=0.0,
            modality_scale=1.0, audio_modality_scale=1.0, guidance_rescale=0.0, audio_guidance_rescale=0.0,
            spatio_temporal_guidance_blocks=None, use_cross_timestep=False, generator=generator, return_dict=False,
        )
    if mode == "I2V":
        call_kwargs["image"] = hooks.load_conditioning_image(str(start_image_path), width, height)
    elif uses_condition_pipeline(mode):
        first_image = hooks.load_conditioning_image(str(start_image_path), width, height)
        conditions = [LTX2VideoCondition(frames=first_image, index=0, strength=1.0)]
        if middle_image_path:
            if num_frames is None:
                raise gr.Error("Middle keyframe currently requires Manual Duration so its timeline midpoint is known before Stage 1.")
            middle_image = hooks.load_conditioning_image(str(middle_image_path), width, height)
            middle_index = middle_keyframe_latent_index(num_frames)
            conditions.append(LTX2VideoCondition(frames=middle_image, index=middle_index, strength=1.0))
        if end_image_path:
            last_image = hooks.load_conditioning_image(str(end_image_path), width, height)
            conditions.append(LTX2VideoCondition(frames=last_image, index=-1, strength=1.0))
        call_kwargs["conditions"] = conditions
    return call_kwargs


def _run_stage1(
    *,
    active_pipeline,
    mode: str,
    stage1_call_kwargs: dict[str, Any],
    num_frames: int | None,
    use_auto_duration: bool,
    width: int,
    height: int,
    runtime: GenerationRuntime,
    progress,
    full_sft_metrics: dict[str, Any],
    request_logger,
):
    # Stage 1: half-resolution video/audio latent generation.
    progress(0.10, desc=f"Stage 1 · {mode} {'Full/SFT guided' if runtime.is_full_sft_profile else 'distilled'} half-resolution")
    stage1_duration_kwargs = {"num_frames": num_frames}
    if use_auto_duration:
        stage1_duration_kwargs.update(
            num_frames=None, min_seconds=runtime.auto_duration_min_seconds, max_seconds=runtime.auto_duration_max_seconds,
        )
    stage1_sampling_kwargs = ({"num_inference_steps": 30} if runtime.is_full_sft_profile else {"sigmas": DISTILLED_SIGMA_VALUES})
    stage1_started = time.monotonic()
    stage1_video_latents, stage1_audio_latents = active_pipeline(
        height=height // 2, width=width // 2, output_type="latent",
        **stage1_sampling_kwargs, **stage1_duration_kwargs, **stage1_call_kwargs,
    )
    full_sft_metrics["stage1_seconds"] = time.monotonic() - stage1_started
    request_logger.info("generate.stage1.complete seconds=%.3f", full_sft_metrics["stage1_seconds"])
    if use_auto_duration:
        latent_frames = int(stage1_video_latents.shape[2])
        temporal_ratio = int(getattr(active_pipeline, "vae_temporal_compression_ratio", 8))
        num_frames = ((latent_frames - 1) * temporal_ratio) + 1
        if num_frames < 1 or (num_frames - 1) % temporal_ratio != 0:
            raise gr.Error("Auto Duration returned an unexpected temporal grid.")
    return stage1_video_latents, stage1_audio_latents, num_frames


def _run_stage2(
    *,
    active_pipeline,
    stage1_call_kwargs: dict[str, Any],
    num_frames: int,
    upscaled_video_latents,
    stage1_audio_latents,
    loaded_loras: list[dict[str, Any]],
    lora_strength: float,
    mode: str,
    use_diffusion_decoder: bool,
    width: int,
    height: int,
    runtime: GenerationRuntime,
    hooks: GenerationHooks,
    progress,
    full_sft_metrics: dict[str, Any],
):
    # Stage 2: full-resolution refinement and synchronized audio.
    progress(0.60, desc="Stage 2 · video + synchronized audio")
    stage2_call_kwargs = dict(stage1_call_kwargs)
    if runtime.is_full_sft_profile:
        full_sft_metrics["gpu_before_stage2_adapter_transfer"] = hooks.gpu_state()
        transfer_started = time.monotonic()
        runtime.pipe.set_lora_device([runtime.full_sft_stage2_adapter_name], device="cuda:0")
        user_adapter_names = [item["adapter_name"] for item in loaded_loras]
        stage2_adapter_names = [*user_adapter_names, runtime.full_sft_stage2_adapter_name]
        stage2_adapter_weights = [float(lora_strength)] * len(user_adapter_names) + [
            runtime.full_sft_stage2_lora_strength
        ]
        runtime.pipe.set_adapters(stage2_adapter_names, adapter_weights=stage2_adapter_weights)
        runtime.pipe.enable_lora()
        full_sft_metrics["stage2_adapter_names"] = list(stage2_adapter_names)
        full_sft_metrics["stage2_adapter_weights"] = list(stage2_adapter_weights)
        full_sft_metrics["adapter_state_stage2"] = hooks.lora_adapter_state()
        full_sft_metrics["stage2_adapter_gpu_transfer_seconds"] = time.monotonic() - transfer_started
        full_sft_metrics["gpu_after_stage2_adapter_transfer"] = hooks.gpu_state()
        distilled_scheduler = FlowMatchEulerDiscreteScheduler.from_config(
            runtime.pipe.scheduler.config, use_dynamic_shifting=False, shift_terminal=None
        )
        for pipeline in (runtime.pipe, runtime.pipe_i2v, runtime.pipe_condition):
            if pipeline is not None:
                pipeline.scheduler = distilled_scheduler
        if runtime.diffusion_decode_pipe is not None:
            runtime.diffusion_decode_pipe.scheduler = distilled_scheduler
        stage2_call_kwargs.update(
            guidance_scale=1.0, audio_guidance_scale=1.0, stg_scale=0.0, audio_stg_scale=0.0,
            modality_scale=1.0, audio_modality_scale=1.0, guidance_rescale=0.0, audio_guidance_rescale=0.0,
            spatio_temporal_guidance_blocks=None, use_cross_timestep=False,
        )
    # Official distilled FLF2V two-stage flow applies image conditions in Stage 1;
    # Stage 2 refines the upscaled latent without re-appending keyframe conditions.
    if uses_condition_pipeline(mode):
        stage2_call_kwargs.pop("conditions", None)
        # The official two-stage condition example supplies the final target
        # size explicitly during Stage 2. Keep validated T2V/I2V semantics
        # unchanged and apply this only to the new FLF2V branch.
        stage2_call_kwargs["height"] = height
        stage2_call_kwargs["width"] = width
    stage2_output_type = "latent" if (bool(use_diffusion_decoder) or runtime.is_full_sft_profile) else "np"
    stage2_started = time.monotonic()
    video_or_latents, audio_or_latents = active_pipeline(
        num_frames=num_frames, sigmas=STAGE_2_DISTILLED_SIGMA_VALUES, latents=upscaled_video_latents,
        audio_latents=stage1_audio_latents, noise_scale=STAGE_2_DISTILLED_SIGMA_VALUES[0],
        output_type=stage2_output_type, **stage2_call_kwargs,
    )
    full_sft_metrics["stage2_seconds"] = time.monotonic() - stage2_started
    if runtime.is_full_sft_profile:
        offload_started = time.monotonic()
        runtime.pipe.disable_lora()
        runtime.pipe.set_lora_device([runtime.full_sft_stage2_adapter_name], device="cpu")
        gc.collect()
        if torch.cuda.is_available():
            torch.cuda.empty_cache()
        full_sft_metrics["stage2_adapter_cpu_release_seconds"] = time.monotonic() - offload_started
        full_sft_metrics["gpu_after_stage2_adapter_release"] = hooks.gpu_state()
        full_sft_metrics["adapter_state_after_stage2_internal_release"] = hooks.lora_adapter_state()
    return video_or_latents, audio_or_latents



def _run_conv_vae_decoder(
    *,
    video_or_latents,
    audio_or_latents,
    runtime: GenerationRuntime,
    hooks: GenerationHooks,
    progress,
    decoder_metrics: dict[str, Any],
    force_framewise: bool = False,
):
    """Finish a latent Stage-2 result with the upstream Conv VAE/audio decode contract.

    Full/SFT Stage 2 is intentionally returned as latents so its internal/user LoRA
    residency can end before the memory-heavy video decode begins. `output_type=latent`
    already denormalizes both video and audio latents in the pinned LTX2 pipeline.
    The normal pipeline decode defaults to decode_timestep=0.0, so no decode noise is
    injected here; timestep conditioning, when present, receives the equivalent zero.
    """
    decoder_metrics["conv_vae_manual_after_stage2_release"] = True
    decoder_metrics["conv_vae_force_framewise"] = bool(force_framewise)
    decoder_metrics["gpu_before_conv_vae_decode"] = hooks.gpu_state()
    prior_framewise = bool(getattr(runtime.pipe.vae, "use_framewise_decoding", False))
    if force_framewise:
        runtime.pipe.vae.use_framewise_decoding = True
    started = time.monotonic()
    try:
        progress(0.80, desc="Conv VAE · video decode")
        with torch.no_grad():
            video_latents = video_or_latents.to(runtime.pipe.vae.dtype)
            if bool(getattr(runtime.pipe.vae.config, "timestep_conditioning", False)):
                timestep = torch.zeros(
                    video_latents.shape[0], device=video_latents.device, dtype=video_latents.dtype
                )
            else:
                timestep = None
            video = runtime.pipe.vae.decode(video_latents, timestep, return_dict=False)[0]
            video = runtime.pipe.video_processor.postprocess_video(video, output_type="np")

        decoder_metrics["gpu_after_video_decode"] = hooks.gpu_state()
        progress(0.86, desc="Conv VAE · audio decode")
        with torch.no_grad():
            audio_latents = audio_or_latents.to(runtime.pipe.audio_vae.dtype)
            mel = runtime.pipe.audio_vae.decode(audio_latents, return_dict=False)[0]
            audio = runtime.pipe.vocoder(mel).detach().cpu()
        decoder_metrics["conv_vae_decode_seconds"] = time.monotonic() - started
        decoder_metrics["gpu_after_conv_vae_decode"] = hooks.gpu_state()
        return video, audio
    finally:
        if force_framewise:
            runtime.pipe.vae.use_framewise_decoding = prior_framewise


def _run_diffusion_decoder(
    *,
    video_or_latents,
    audio_or_latents,
    generator: torch.Generator,
    runtime: GenerationRuntime,
    hooks: GenerationHooks,
    progress,
    decoder_metrics: dict[str, Any],
):
    decoder_metrics["gpu_before_audio_decode"] = hooks.gpu_state()

    progress(0.78, desc="Diffusion decoder · audio decode")
    decoder_phase_started = time.monotonic()
    # The upstream LTX2 pipeline performs its normal VAE/vocoder decode inside
    # @torch.no_grad(). Stage-2 output_type="latent" returns before that block,
    # so this manual completion must restore the same inference-only contract.
    decoder_metrics["audio_decode_autograd"] = "disabled"
    with torch.no_grad():
        audio_decode_latents = audio_or_latents.to(runtime.pipe.audio_vae.dtype)
        mel = runtime.pipe.audio_vae.decode(audio_decode_latents, return_dict=False)[0]
        audio = runtime.pipe.vocoder(mel).detach().cpu()
    decoder_metrics["audio_decode_seconds"] = time.monotonic() - decoder_phase_started

    # The final waveform is CPU-resident. Drop audio GPU intermediates and cached
    # allocator blocks before bringing the optional video decoder onto the device.
    audio_or_latents = None
    audio_decode_latents = None
    mel = None
    gc.collect()
    if torch.cuda.is_available():
        torch.cuda.empty_cache()
    decoder_metrics["gpu_after_audio_release"] = hooks.gpu_state()

    progress(0.84, desc="Diffusion decoder · GPU transfer")
    decoder_metrics["gpu_before_transfer"] = hooks.gpu_state()
    decoder_phase_started = time.monotonic()
    runtime.diffusion_decode_pipe.diffusion_decoder.to("cuda")
    decoder_metrics["gpu_transfer_seconds"] = time.monotonic() - decoder_phase_started
    decoder_metrics["device_after_transfer"] = str(next(runtime.diffusion_decode_pipe.diffusion_decoder.parameters()).device)
    decoder_metrics["gpu_after_transfer"] = hooks.gpu_state()

    progress(0.87, desc="Diffusion decoder · NATTEN tiled decode")
    decoder_phase_started = time.monotonic()
    video = runtime.diffusion_decode_pipe(
        video_or_latents,
        generator=generator,
        output_type="np",
        denormalize=False,
        return_dict=False,
    )[0]
    decoder_metrics["decode_seconds"] = time.monotonic() - decoder_phase_started
    decoder_metrics["gpu_after_decode"] = hooks.gpu_state()

    # The decoder's GPU residency ends with the decode phase, not at request end.
    video_or_latents = None
    decoder_phase_started = time.monotonic()
    runtime.diffusion_decode_pipe.diffusion_decoder.to("cpu")
    gc.collect()
    if torch.cuda.is_available():
        torch.cuda.empty_cache()
    decoder_metrics["gpu_release_seconds"] = time.monotonic() - decoder_phase_started
    decoder_metrics["device_after_release"] = str(next(runtime.diffusion_decode_pipe.diffusion_decoder.parameters()).device)
    decoder_metrics["gpu_after_release"] = hooks.gpu_state()
    return video, audio


def _restore_shared_runtime(
    *,
    runtime: GenerationRuntime,
    hooks: GenerationHooks,
    request_logger,
    loaded_loras: list[dict[str, Any]],
    lora_metrics: dict[str, Any],
    lora_cleanup_done: bool,
    use_diffusion_decoder: bool,
) -> None:
    # Restore shared runtime state even when inference/encoding fails.
    if runtime.is_full_sft_profile and runtime.pipe is not None:
        try:
            runtime.pipe.disable_lora()
            runtime.pipe.set_lora_device([runtime.full_sft_stage2_adapter_name], device="cpu")
        except Exception as cleanup_exc:
            request_logger.warning("generate.full_sft.stage2_cleanup warning=%s: %s", type(cleanup_exc).__name__, cleanup_exc)
        try:
            sft_scheduler = FlowMatchEulerDiscreteScheduler.from_config(
                runtime.pipe.scheduler.config, use_dynamic_shifting=True, shift_terminal=0.1
            )
            for pipeline in (runtime.pipe, runtime.pipe_i2v, runtime.pipe_condition):
                if pipeline is not None:
                    pipeline.scheduler = sft_scheduler
            if runtime.diffusion_decode_pipe is not None:
                runtime.diffusion_decode_pipe.scheduler = sft_scheduler
        except Exception as scheduler_exc:
            request_logger.warning("generate.full_sft.scheduler_restore warning=%s: %s", type(scheduler_exc).__name__, scheduler_exc)
    if runtime.diffusion_decode_pipe is not None:
        try:
            runtime.diffusion_decode_pipe.diffusion_decoder.to("cpu")
            gc.collect()
            if torch.cuda.is_available():
                torch.cuda.empty_cache()
            if bool(use_diffusion_decoder):
                request_logger.info("generate.diffusion_decoder.cleaned cpu_resident=true")
        except Exception as decoder_cleanup_exc:
            request_logger.exception("generate.diffusion_decoder.cleanup.failure error_type=%s error=%s", type(decoder_cleanup_exc).__name__, decoder_cleanup_exc)
    if not lora_cleanup_done:
        try:
            hooks.cleanup_request_loras(loaded_loras, lora_metrics)
        except Exception as cleanup_exc:
            request_logger.exception("generate.lora.cleanup.failure error_type=%s error=%s", type(cleanup_exc).__name__, cleanup_exc)


def run(
    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,
    progress,
    *,
    runtime: GenerationRuntime,
    hooks: GenerationHooks,
):
    # Validate the request and resolve the active product pipeline.
    if runtime.preload_state.get("status") != "ready" or runtime.pipe is None or runtime.upsample_pipe is None:
        raise gr.Error(f"Diffusers preload is not ready: {runtime.preload_state}")

    prompt = str(prompt or "").strip()
    if not prompt:
        raise gr.Error("Prompt is required.")
    mode = hooks.mode_from_images(start_image_path, middle_image_path, end_image_path)
    if mode == "INVALID_CONDITION_WITHOUT_START":
        raise gr.Error("Middle/End frame requires a Start frame. Remove the extra condition or provide a Start frame.")
    if middle_image_path and bool(use_auto_duration):
        raise gr.Error("Middle keyframe currently requires Manual Duration so its midpoint can be resolved before Stage 1.")
    use_auto_duration = bool(use_auto_duration)
    requested_video_seconds = float(duration_seconds)
    if use_auto_duration:
        if not runtime.auto_duration_enabled or getattr(runtime.pipe, "duration_head", None) is None:
            record = (runtime.preload_state.get("model_sources") or {}).get("duration_head") or {}
            raise gr.Error(
                f"Auto Duration is unavailable: {record.get('fallback_reason') or 'duration_head did not initialize'}"
            )
        num_frames = None
    else:
        if requested_video_seconds > runtime.standard_max_seconds and not bool(experimental_long):
            raise gr.Error("Enable Experimental long duration to generate more than 15 seconds.")
        if requested_video_seconds > runtime.experimental_max_seconds:
            raise gr.Error(f"Maximum exposed duration is {runtime.experimental_max_seconds:.0f} seconds.")
        if requested_video_seconds > runtime.standard_max_seconds and not hooks.supports_experimental_long(resolution_key):
            raise gr.Error("15–30 second experimental generation is currently restricted to 512×512.")
        num_frames = hooks.frames_from_seconds(requested_video_seconds)
        if num_frames % 8 != 1:
            raise gr.Error("Internal frame-grid error: expected 8k+1 frames.")
    width, height = hooks.resolution(resolution_key)
    zerogpu_requested_duration_seconds = hooks.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
    )
    if mode == "I2V" and runtime.pipe_i2v is None:
        raise gr.Error("I2V pipeline did not initialize.")
    if uses_condition_pipeline(mode) and runtime.pipe_condition is None:
        raise gr.Error("Timeline condition pipeline did not initialize.")
    if bool(use_diffusion_decoder) and runtime.diffusion_decode_pipe is None:
        record = (runtime.preload_state.get("model_sources") or {}).get("diffusion_decoder") or {}
        raise gr.Error(f"Diffusion decoder is unavailable: {record.get('fallback_reason') or 'startup preparation did not complete'}")

    paths = hooks.create_request_paths(session_id)
    seed = random.SystemRandom().randrange(0, 2**31 - 1) if bool(randomize_seed) else int(seed)
    if uses_condition_pipeline(mode):
        active_pipeline = runtime.pipe_condition
    elif mode == "I2V":
        active_pipeline = runtime.pipe_i2v
    else:
        active_pipeline = runtime.pipe
    generator = torch.Generator("cuda").manual_seed(seed)
    stage1_call_kwargs = _build_stage1_call_kwargs(
        prompt=prompt, generator=generator, mode=mode, start_image_path=start_image_path,
        middle_image_path=middle_image_path, end_image_path=end_image_path, num_frames=num_frames,
        width=width, height=height, runtime=runtime, hooks=hooks,
    )

    # Initialize request-scoped observability and adapter state.
    loaded_loras = []
    lora_metrics = {
        "requested_labels": [str(x) for x in (selected_loras or [])],
        "requested_count": len(selected_loras or []),
        "hub_download_inside_gpu_callback": False,
    }
    lora_cleanup_done = False
    decoder_metrics = {
        "requested": bool(use_diffusion_decoder),
        "available": runtime.diffusion_decode_pipe is not None,
        "startup_residency": "CPU RAM-ready / GPU-lazy" if runtime.diffusion_decode_pipe is not None else "unavailable",
    }
    full_sft_metrics = {
        "profile": runtime.runtime_profile,
        "active": bool(runtime.is_full_sft_profile),
        "stage2_adapter": runtime.full_sft_stage2_adapter_name if runtime.is_full_sft_profile else None,
        "stage2_adapter_startup_residency": "CPU RAM-ready / GPU-lazy" if runtime.is_full_sft_profile else None,
        "stage2_adapter_strength": runtime.full_sft_stage2_lora_strength if runtime.is_full_sft_profile else None,
        "user_lora_stage_policy": (
            "Stage 1 user LoRAs; Stage 2 same user LoRAs + internal distilled adapter"
            if runtime.is_full_sft_profile else None
        ),
    }
    request_logger, request_handler = hooks.open_request_logger(paths.log, paths.request_id)
    request_logger.info(
        "generate.start mode=%s size=%sx%s requested_video_seconds=%s auto_duration=%s profile=%s seed=%s",
        mode, width, height, requested_video_seconds, use_auto_duration, runtime.runtime_profile, seed,
    )
    reset_gpu_peak_memory()
    gpu_before = hooks.gpu_state()
    execution_environment = execution_environment_identity(runtime.preload_state.get("environment"))
    request_started = time.monotonic()
    try:
        progress(0.04, desc="Loading prepared LoRA" if selected_loras else "Preparing request")
        loaded_loras, lora_metrics = hooks.load_request_loras(
            selected_loras, custom_loras, prepared_loras, lora_strength, paths.request_id
        )
        request_logger.info("generate.lora.ready count=%s", len(loaded_loras))
        if runtime.is_full_sft_profile:
            full_sft_metrics["stage1_user_adapter_names"] = [item["adapter_name"] for item in loaded_loras]
            full_sft_metrics["adapter_state_stage1"] = hooks.lora_adapter_state()

        stage1_video_latents, stage1_audio_latents, num_frames = _run_stage1(
            active_pipeline=active_pipeline, mode=mode, stage1_call_kwargs=stage1_call_kwargs, num_frames=num_frames,
            use_auto_duration=use_auto_duration, width=width, height=height, runtime=runtime,
            progress=progress, full_sft_metrics=full_sft_metrics, request_logger=request_logger,
        )

        # Spatial latent upscale between the two diffusion stages.
        progress(0.50, desc="Latent upscale ×2")
        upscaled_video_latents = runtime.upsample_pipe(
            latents=stage1_video_latents, output_type="latent", return_dict=False
        )[0]

        video_or_latents, audio_or_latents = _run_stage2(
            active_pipeline=active_pipeline, stage1_call_kwargs=stage1_call_kwargs, num_frames=num_frames,
            upscaled_video_latents=upscaled_video_latents, stage1_audio_latents=stage1_audio_latents,
            loaded_loras=loaded_loras, lora_strength=lora_strength, mode=mode,
            use_diffusion_decoder=use_diffusion_decoder, width=width, height=height, runtime=runtime,
            hooks=hooks, progress=progress, full_sft_metrics=full_sft_metrics,
        )

        # Stage-2 transformer inference is complete. For Full/SFT, Stage 2 deliberately
        # returns latents so all transformer-only adapter residency can end before decode.
        # User LoRAs are likewise no longer needed by either Conv VAE or diffusion decoder.
        if loaded_loras:
            lora_metrics = hooks.cleanup_request_loras(loaded_loras, lora_metrics)
            lora_cleanup_done = True

        if runtime.is_full_sft_profile or bool(use_diffusion_decoder):
            # Drop Stage-1/upscale references and allocator cache before entering a decoder.
            stage1_video_latents = None
            stage1_audio_latents = None
            upscaled_video_latents = None
            gc.collect()
            if torch.cuda.is_available():
                torch.cuda.empty_cache()

        # Optional diffusion-decoder completion replaces the default Conv VAE video decode.
        if bool(use_diffusion_decoder):
            video, audio = _run_diffusion_decoder(
                video_or_latents=video_or_latents, audio_or_latents=audio_or_latents, generator=generator,
                runtime=runtime, hooks=hooks, progress=progress, decoder_metrics=decoder_metrics,
            )
        elif runtime.is_full_sft_profile:
            video, audio = _run_conv_vae_decoder(
                video_or_latents=video_or_latents, audio_or_latents=audio_or_latents,
                runtime=runtime, hooks=hooks, progress=progress, decoder_metrics=decoder_metrics,
            )
        else:
            video, audio = video_or_latents, audio_or_latents

        # Encode product output, then persist machine-readable evidence/history.
        progress(0.92, desc="Encoding MP4")
        encode_video(
            video[0],
            fps=int(runtime.frame_rate),
            output_path=str(paths.video),
            audio=audio[0].float().cpu(),
            audio_sample_rate=runtime.pipe.vocoder.config.output_sampling_rate,
        )
        elapsed = time.monotonic() - request_started
        run_settings = hooks.settings_snapshot(
            prompt, duration_seconds, experimental_long, resolution_key, seed, False,
            selected_loras, lora_strength, custom_loras, use_diffusion_decoder, use_auto_duration,
            start_image_path, middle_image_path, end_image_path,
        )
        run_settings["seed"]["randomize_requested"] = bool(randomize_seed)
        middle_latent_index = middle_keyframe_latent_index(num_frames) if middle_image_path else None
        middle_pixel_frame = keyframe_pixel_frame_index(middle_latent_index) if middle_latent_index is not None else None
        conditioning = {
            "start_image_name": Path(str(start_image_path)).name if start_image_path else None,
            "start_image_sha256": hooks.sha256_file(start_image_path) if start_image_path else None,
            "middle_image_name": Path(str(middle_image_path)).name if middle_image_path else None,
            "middle_image_sha256": hooks.sha256_file(middle_image_path) if middle_image_path else None,
            "middle_keyframe_latent_index": middle_latent_index,
            "middle_keyframe_pixel_frame": middle_pixel_frame,
            "middle_keyframe_seconds": (middle_pixel_frame / runtime.frame_rate) if middle_pixel_frame is not None else None,
            "end_image_name": Path(str(end_image_path)).name if end_image_path else None,
            "end_image_sha256": hooks.sha256_file(end_image_path) if end_image_path else None,
        }
        run_info = {
            "schema_version": "ltx25-run-v1",
            "candidate_id": runtime.candidate_id,
            "request_id": paths.request_id,
            "mode": mode,
            "width": width,
            "height": height,
            "frames": num_frames,
            "duration_mode": "auto" if use_auto_duration else "manual",
            "requested_duration_seconds": None if use_auto_duration else float(duration_seconds),
            "auto_duration_bounds_seconds": (
                {"min": runtime.auto_duration_min_seconds, "max": runtime.auto_duration_max_seconds}
                if use_auto_duration else None
            ),
            "fps": runtime.frame_rate,
            "realized_duration_seconds": (num_frames - 1) / runtime.frame_rate,
            "seed": seed,
            "elapsed_seconds": elapsed,
            "conditioning": conditioning,
            "settings": run_settings,
            "loras_effective": loaded_loras,
            "lora_metrics": lora_metrics,
            "runtime_model_sources": runtime.preload_state.get("model_sources"),
            "attention_backend": runtime.preload_state.get("attention_backend"),
            "diffusers": (runtime.preload_state.get("packages") or {}).get("diffusers"),
            "runtime_packages": runtime.preload_state.get("packages"),
            "runtime_environment": execution_environment,
            "preload_environment": runtime.preload_state.get("environment"),
            "decoder": "diffusion" if bool(use_diffusion_decoder) else "conv_vae",
            "decoder_metrics": decoder_metrics,
            "runtime_profile": runtime.runtime_profile,
            "full_sft_metrics": full_sft_metrics,
            "output_file": paths.video.name,
            "log_file": paths.log.name,
        }
        paths.run_info.write_text(json.dumps(run_info, indent=2, sort_keys=True), encoding="utf-8")

        long_duration_metrics = {"active": bool(((num_frames - 1) / runtime.frame_rate) > runtime.standard_max_seconds)}
        if long_duration_metrics["active"] and torch.cuda.is_available():
            long_duration_metrics["gpu_before_cache_trim"] = hooks.gpu_state()
            cache_trim_started = time.monotonic()
            gc.collect()
            torch.cuda.empty_cache()
            long_duration_metrics["cache_trim_seconds"] = time.monotonic() - cache_trim_started
            long_duration_metrics["gpu_after_cache_trim"] = hooks.gpu_state()

        diag = {
            "candidate_id": runtime.candidate_id,
            "status": "PASS",
            "worker_uuid": runtime.worker_uuid,
            "session_id": paths.session_id,
            "request_id": paths.request_id,
            "mode": mode,
            "width": width,
            "height": height,
            "frames": num_frames,
            "duration_mode": "auto" if use_auto_duration else "manual",
            "requested_duration_seconds": None if use_auto_duration else float(duration_seconds),
            "auto_duration_bounds_seconds": (
                {"min": runtime.auto_duration_min_seconds, "max": runtime.auto_duration_max_seconds}
                if use_auto_duration else None
            ),
            "realized_duration_seconds": (num_frames - 1) / runtime.frame_rate,
            "zerogpu_requested_duration_seconds": zerogpu_requested_duration_seconds,
            "fps": runtime.frame_rate,
            "seed": seed,
            "loras": loaded_loras,
            "lora_strength": float(lora_strength) if loaded_loras else None,
            "lora_metrics": lora_metrics,
            "decoder": "diffusion" if bool(use_diffusion_decoder) else "conv_vae",
            "decoder_metrics": decoder_metrics,
            "runtime_profile": runtime.runtime_profile,
            "full_sft_metrics": full_sft_metrics,
            "long_duration_metrics": long_duration_metrics,
            "elapsed_seconds": elapsed,
            "gpu_before": gpu_before,
            "gpu_after": hooks.gpu_state(),
            "ru_maxrss_kib": process_rss_kib(),
            "disk_after_run": hooks.disk_state(),
            "preload_elapsed_seconds": runtime.preload_state.get("elapsed_seconds"),
            "preload_phases": runtime.preload_state.get("preload_phases"),
            "runtime_model_sources": runtime.preload_state.get("model_sources"),
            "attention_backend": runtime.preload_state.get("attention_backend"),
            "runtime_environment": execution_environment,
            "preload_environment": runtime.preload_state.get("environment"),
            "runtime_packages": runtime.preload_state.get("packages"),
            "run_info_file": paths.run_info.name,
            "log_file": paths.log.name,
        }
        diag_json = json.dumps(diag, indent=2, sort_keys=True)
        paths.diagnostics.write_text(diag_json, encoding="utf-8")
        request_logger.info("generate.success elapsed_seconds=%.3f output=%s", elapsed, paths.video.name)
        try:
            request_handler.flush()
        except Exception:
            pass
        probe_path = probe_artifacts.build_request_probe(
            paths,
            candidate_id=runtime.candidate_id,
            status="PASS",
            sha256_file=hooks.sha256_file,
            metadata={"surface": "generation", "mode": mode},
        )

        history_record = {
            "request_id": paths.request_id,
            "request_root": str(paths.root),
            "video": str(paths.video),
            "probe": str(probe_path),
            "mode": mode,
            "width": width,
            "height": height,
            "seed": seed,
            "elapsed_seconds": elapsed,
            "settings": run_settings,
        }
        new_history = hooks.update_history(history, history_record)
        history_dropdown, history_summary = hooks.history_outputs(new_history)
        progress(1.0, desc="Done")
        return (
            str(paths.video), seed, str(seed), str(probe_path),
            new_history, history_dropdown, history_summary,
        )
    except Exception as exc:
        if loaded_loras and not lora_cleanup_done:
            try:
                lora_metrics = hooks.cleanup_request_loras(loaded_loras, lora_metrics)
                lora_cleanup_done = True
            except Exception as cleanup_exc:
                lora_metrics["cleanup_status"] = "FAIL"
                lora_metrics["cleanup_error"] = f"{type(cleanup_exc).__name__}: {cleanup_exc}"
        diag = {
            "candidate_id": runtime.candidate_id,
            "status": "FAIL",
            "worker_uuid": runtime.worker_uuid,
            "session_id": paths.session_id,
            "request_id": paths.request_id,
            "mode": mode,
            "duration_mode": "auto" if use_auto_duration else "manual",
            "auto_duration_bounds_seconds": (
                {"min": runtime.auto_duration_min_seconds, "max": runtime.auto_duration_max_seconds}
                if use_auto_duration else None
            ),
            "predicted_frames_if_available": num_frames,
            "decoder": "diffusion" if bool(use_diffusion_decoder) else "conv_vae",
            "decoder_metrics": decoder_metrics,
            "lora_metrics": lora_metrics,
            "runtime_profile": runtime.runtime_profile,
            "full_sft_metrics": full_sft_metrics,
            "log_file": paths.log.name,
            "error_type": type(exc).__name__,
            "error": str(exc),
            "elapsed_seconds": time.monotonic() - request_started,
            "gpu": hooks.gpu_state(),
        }
        try:
            paths.diagnostics.write_text(json.dumps(diag, indent=2, sort_keys=True), encoding="utf-8")
        except Exception:
            pass
        request_logger.exception("generate.failure error_type=%s error=%s", type(exc).__name__, exc)
        try:
            if request_handler is not None:
                request_handler.flush()
            probe_artifacts.build_request_probe(
                paths,
                candidate_id=runtime.candidate_id,
                status="FAIL",
                sha256_file=hooks.sha256_file,
                metadata={"surface": "generation", "mode": mode, "error_type": type(exc).__name__},
            )
        except Exception as probe_exc:
            request_logger.warning("generate.probe_bundle.failure error_type=%s error=%s", type(probe_exc).__name__, probe_exc)
        raise
    finally:
        _restore_shared_runtime(
            runtime=runtime, hooks=hooks, request_logger=request_logger, loaded_loras=loaded_loras,
            lora_metrics=lora_metrics, lora_cleanup_done=lora_cleanup_done,
            use_diffusion_decoder=bool(use_diffusion_decoder),
        )
        hooks.close_request_logger(request_logger, request_handler)