ltx/maintainer_probe.py

from __future__ import annotations

import gc
import json
import time
from collections.abc import Callable
from typing import Any

import gradio as gr
import torch

from . import generation as generation_backend
from . import probe_artifacts
from . import runtime_utils

FULL_SFT_30S_STAGE1_PROFILE = "full_sft_30s_stage1_5step"
FULL_SFT_30S_STAGE1_STEPS = 5
FULL_SFT_30S_STAGE1_FRAMES = 721
FULL_SFT_30S_STAGE1_ZEROGPU_DURATION_SECONDS = 75


FULL_SFT_30S_STAGE2_CONV_PROFILE = "full_sft_30s_stage2_conv_release_before_decode_2step"
FULL_SFT_30S_STAGE2_DIFFUSION_PROFILE = "full_sft_30s_stage2_diffusion_decoder_2step"
FULL_SFT_30S_STAGE2_CONV_FRAMEWISE_PROFILE = "full_sft_30s_stage2_conv_framewise_2step"
FULL_SFT_30S_STAGE2_LATENT_PROFILE = "full_sft_30s_stage2_latent_2step"
FULL_SFT_30S_UPSAMPLE_PROFILE = "full_sft_30s_upsample_2step"
FULL_SFT_30S_STAGE2_CONV_ZEROGPU_DURATION_SECONDS = 100
FULL_SFT_30S_STAGE2_DIFFUSION_ZEROGPU_DURATION_SECONDS = 140
FULL_SFT_30S_STAGE2_CONV_FRAMEWISE_ZEROGPU_DURATION_SECONDS = 110
FULL_SFT_30S_STAGE2_LATENT_ZEROGPU_DURATION_SECONDS = 75
FULL_SFT_30S_UPSAMPLE_ZEROGPU_DURATION_SECONDS = 50
FULL_SFT_30S_FLF2V_STAGE1_PROFILE = "full_sft_30s_flf2v_stage1_2step"
FULL_SFT_30S_FLF2V_STAGE1_STEPS = 2
FULL_SFT_30S_FLF2V_STAGE1_ZEROGPU_DURATION_SECONDS = 50
FULL_SFT_30S_TAIL_STAGE1_STEPS = 2


def _run_full_sft_30s_tail_scout(
    *,
    profile_id: str,
    stop_after: str,
    current_gate: bool,
    requested_duration_seconds: int,
    session_id,
    progress,
    candidate_id: str,
    runtime: generation_backend.GenerationRuntime,
    hooks: generation_backend.GenerationHooks,
    cats_example: tuple[str, str],
    create_request_paths: Callable[[Any], runtime_utils.RequestPaths],
    open_request_logger: Callable,
    close_request_logger: Callable,
):
    """Measure the exact 721f tail path without repeating the 30-step Stage 1."""
    if not runtime.is_full_sft_profile:
        raise gr.Error("This Probe requires MODEL_RUNTIME_PROFILE=full_sft_nf4 and a Space restart.")
    if runtime.preload_state.get("status") != "ready" or runtime.pipe is None or runtime.pipe_i2v is None:
        raise gr.Error(f"Diffusers preload is not ready: {runtime.preload_state}")
    if runtime.upsample_pipe is None:
        raise gr.Error("Latent upsampler is unavailable on this worker.")
    if stop_after == "stage2_diffusion_decoder" and runtime.diffusion_decode_pipe is None:
        raise gr.Error("Diffusion decoder is unavailable on this worker.")

    start_image_path, probe_prompt = cats_example
    paths = create_request_paths(session_id)
    request_logger, request_handler = open_request_logger(paths.log, paths.request_id)
    stage1_video_latents = None
    stage1_audio_latents = None
    upscaled_video_latents = None
    video_or_latents = None
    audio_or_latents = None
    started = time.monotonic()
    status = "FAIL"
    full_sft_metrics: dict[str, Any] = {}
    decoder_metrics: dict[str, Any] = {}
    try:
        request_logger.info(
            "maintainer_probe.start profile=%s mode=I2V size=512x512 frames=721 stage1_steps=%s stop_after=%s seed=42 no_user_lora=true",
            profile_id,
            FULL_SFT_30S_TAIL_STAGE1_STEPS,
            stop_after,
        )
        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()
        runtime_utils.reset_gpu_peak_memory()
        gpu_before = hooks.gpu_state()
        execution_environment = runtime_utils.execution_environment_identity(runtime.preload_state.get("environment"))

        generator = torch.Generator("cuda").manual_seed(42)
        stage1_call_kwargs = generation_backend._build_stage1_call_kwargs(
            prompt=probe_prompt,
            generator=generator,
            mode="I2V",
            start_image_path=start_image_path,
            middle_image_path=None,
            end_image_path=None,
            num_frames=FULL_SFT_30S_STAGE1_FRAMES,
            width=512,
            height=512,
            runtime=runtime,
            hooks=hooks,
        )

        def _tail_stage1_progress_callback(_pipeline, step_index, _timestep, callback_kwargs):
            completed = int(step_index) + 1
            progress(
                0.08 + (0.24 * (completed / FULL_SFT_30S_TAIL_STAGE1_STEPS)),
                desc=f"Maintainer Probe · Stage 1 step {completed}/{FULL_SFT_30S_TAIL_STAGE1_STEPS}",
            )
            return callback_kwargs

        progress(0.08, desc="Maintainer Probe · two-step 721f latent preparation")
        stage1_started = time.monotonic()
        stage1_video_latents, stage1_audio_latents = runtime.pipe_i2v(
            height=256,
            width=256,
            output_type="latent",
            num_inference_steps=FULL_SFT_30S_TAIL_STAGE1_STEPS,
            num_frames=FULL_SFT_30S_STAGE1_FRAMES,
            callback_on_step_end=_tail_stage1_progress_callback,
            callback_on_step_end_tensor_inputs=[],
            **stage1_call_kwargs,
        )
        stage1_seconds = time.monotonic() - stage1_started
        gpu_after_stage1 = hooks.gpu_state()

        progress(0.34, desc="Maintainer Probe · exact 721f latent upscale ×2")
        upscale_started = time.monotonic()
        upscaled_video_latents = runtime.upsample_pipe(
            latents=stage1_video_latents, output_type="latent", return_dict=False
        )[0]
        upscale_seconds = time.monotonic() - upscale_started
        gpu_after_upscale = hooks.gpu_state()

        stage2_seconds = None
        gpu_after_stage2 = None
        output_contract = "upscaled_latent_only"
        if stop_after != "upsample":
            stage2_step_count = max(1, len(generation_backend.STAGE_2_DISTILLED_SIGMA_VALUES) - 1)

            def _tail_stage2_progress_callback(_pipeline, step_index, _timestep, callback_kwargs):
                completed = int(step_index) + 1
                progress(
                    min(0.76, 0.60 + (0.16 * (completed / stage2_step_count))),
                    desc=f"Maintainer Probe · Stage 2 step {completed}/{stage2_step_count}",
                )
                return callback_kwargs

            stage2_call_kwargs = dict(stage1_call_kwargs)
            stage2_call_kwargs.update(
                callback_on_step_end=_tail_stage2_progress_callback,
                callback_on_step_end_tensor_inputs=[],
            )
            progress(0.50, desc="Maintainer Probe · real 721f Stage 2")
            stage2_started = time.monotonic()
            video_or_latents, audio_or_latents = generation_backend._run_stage2(
                active_pipeline=runtime.pipe_i2v,
                stage1_call_kwargs=stage2_call_kwargs,
                num_frames=FULL_SFT_30S_STAGE1_FRAMES,
                upscaled_video_latents=upscaled_video_latents,
                stage1_audio_latents=stage1_audio_latents,
                loaded_loras=[],
                lora_strength=1.0,
                mode="I2V",
                use_diffusion_decoder=(stop_after in {"stage2_latent", "stage2_diffusion_decoder"}),
                width=512,
                height=512,
                runtime=runtime,
                hooks=hooks,
                progress=progress,
                full_sft_metrics=full_sft_metrics,
            )
            stage2_seconds = time.monotonic() - stage2_started
            gpu_after_stage2_latent = hooks.gpu_state()
            gpu_after_stage2 = gpu_after_stage2_latent
            output_contract = "stage2_latent_only"
            if stop_after in {"stage2_conv_vae", "stage2_conv_framewise", "stage2_diffusion_decoder"}:
                # Full/SFT Stage 2 returns latents and releases transformer-only
                # adapter residency before the final decoder. Drop all earlier latent
                # references as well so decode sees the real post-Stage2 headroom.
                stage1_video_latents = None
                stage1_audio_latents = None
                upscaled_video_latents = None
                gc.collect()
                if torch.cuda.is_available():
                    torch.cuda.empty_cache()
                if stop_after == "stage2_diffusion_decoder":
                    video_or_latents, audio_or_latents = generation_backend._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,
                    )
                    output_contract = "stage2_diffusion_decoder_decoded_np"
                else:
                    video_or_latents, audio_or_latents = generation_backend._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,
                        force_framewise=(stop_after == "stage2_conv_framewise"),
                    )
                    output_contract = (
                        "stage2_conv_vae_framewise_decoded_np"
                        if stop_after == "stage2_conv_framewise"
                        else "stage2_conv_vae_release_before_decode_decoded_np"
                    )
                gpu_after_stage2 = hooks.gpu_state()

        elapsed = time.monotonic() - started
        run_info = {
            "schema_version": "ltx25-maintainer-probe-v2",
            "candidate_id": candidate_id,
            "request_id": paths.request_id,
            "status": "PASS",
            "surface": "maintainer_probe",
            "profile_id": profile_id,
            "current_gate": bool(current_gate),
            "contract": {
                "runtime_profile": runtime.runtime_profile,
                "mode": "I2V",
                "example": "Cats · I2V",
                "width": 512,
                "height": 512,
                "frames": FULL_SFT_30S_STAGE1_FRAMES,
                "shape_video_seconds": 30.0,
                "stage1_inference_steps": FULL_SFT_30S_TAIL_STAGE1_STEPS,
                "stage2_sigmas": "STAGE_2_DISTILLED_SIGMA_VALUES",
                "seed": 42,
                "user_lora_count": 0,
                "stop_after": stop_after,
                "output_contract": output_contract,
                "mp4_encode": False,
                "zerogpu_requested_duration_seconds": int(requested_duration_seconds),
            },
            "timing": {
                "elapsed_seconds": elapsed,
                "stage1_2step_seconds": stage1_seconds,
                "latent_upscale_seconds": upscale_seconds,
                "stage2_path_seconds": stage2_seconds,
                "stage2_backend_seconds": full_sft_metrics.get("stage2_seconds"),
                "stage2_adapter_gpu_transfer_seconds": full_sft_metrics.get("stage2_adapter_gpu_transfer_seconds"),
                "stage2_adapter_cpu_release_seconds": full_sft_metrics.get("stage2_adapter_cpu_release_seconds"),
            },
            "gpu_before": gpu_before,
            "gpu_after_stage1": gpu_after_stage1,
            "gpu_after_upscale": gpu_after_upscale,
            "gpu_after_stage2": gpu_after_stage2,
            "full_sft_metrics": full_sft_metrics,
            "decoder_metrics": decoder_metrics,
            "ru_maxrss_kib": runtime_utils.process_rss_kib(),
            "input_image_sha256": hooks.sha256_file(start_image_path),
            "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"),
        }
        paths.run_info.write_text(json.dumps(run_info, indent=2, sort_keys=True), encoding="utf-8")
        paths.diagnostics.write_text(json.dumps(run_info, indent=2, sort_keys=True), encoding="utf-8")
        status = "PASS"
        request_logger.info(
            "maintainer_probe.success profile=%s elapsed_seconds=%.3f stage1_2step_seconds=%.3f upscale_seconds=%.3f stage2_path_seconds=%s",
            profile_id, elapsed, stage1_seconds, upscale_seconds,
            "n/a" if stage2_seconds is None else f"{stage2_seconds:.3f}",
        )
        try:
            request_handler.flush()
        except Exception:
            pass
        progress(0.94, desc="Maintainer Probe · packaging evidence")
        probe_path = probe_artifacts.build_request_probe(
            paths,
            candidate_id=candidate_id,
            status="PASS",
            sha256_file=hooks.sha256_file,
            metadata={"surface": "maintainer_probe", "profile_id": profile_id},
        )
        peak_source = gpu_after_stage2 or gpu_after_upscale or gpu_after_stage1
        peak_reserved = float((peak_source or {}).get("max_reserved_bytes") or 0) / (1024**3)
        stage2_text = "not run" if stage2_seconds is None else f"{stage2_seconds:.2f}s"
        progress(1.0, desc="Maintainer Probe · complete")
        return (
            f"**PASS · {profile_id}**  \n"
            f"Stage1 2-step: **{stage1_seconds:.2f}s** · upscale: **{upscale_seconds:.2f}s** · Stage2/decode path: **{stage2_text}** · peak reserved: **{peak_reserved:.2f} GiB**.  \n"
            "Return the Probe ZIP; this is a structural/timing witness, not a visual-quality run.",
            str(probe_path),
        )
    except Exception as exc:
        elapsed = time.monotonic() - started
        request_logger.exception("maintainer_probe.failure profile=%s error_type=%s error=%s", profile_id, type(exc).__name__, exc)
        failure = {
            "schema_version": "ltx25-maintainer-probe-v2",
            "candidate_id": candidate_id,
            "request_id": paths.request_id,
            "status": "FAIL",
            "surface": "maintainer_probe",
            "profile_id": profile_id,
            "current_gate": bool(current_gate),
            "elapsed_seconds": elapsed,
            "error_type": type(exc).__name__,
            "error": str(exc),
            "gpu_after_failure": hooks.gpu_state(),
            "full_sft_metrics": full_sft_metrics,
            "decoder_metrics": decoder_metrics,
        }
        try:
            paths.diagnostics.write_text(json.dumps(failure, indent=2, sort_keys=True), encoding="utf-8")
            request_handler.flush()
            probe_artifacts.build_request_probe(
                paths, candidate_id=candidate_id, status="FAIL", sha256_file=hooks.sha256_file,
                metadata={"surface": "maintainer_probe", "profile_id": profile_id},
            )
        except Exception as evidence_exc:
            request_logger.exception(
                "maintainer_probe.evidence.failure error_type=%s error=%s",
                type(evidence_exc).__name__, evidence_exc,
            )
        if isinstance(exc, gr.Error):
            raise
        raise gr.Error(f"Maintainer Probe failed: {type(exc).__name__}: {exc}") from exc
    finally:
        stage1_video_latents = None
        stage1_audio_latents = None
        upscaled_video_latents = None
        video_or_latents = None
        audio_or_latents = None
        try:
            generation_backend._restore_shared_runtime(
                runtime=runtime, hooks=hooks, request_logger=request_logger, loaded_loras=[],
                lora_metrics={}, lora_cleanup_done=True,
                use_diffusion_decoder=(stop_after == "stage2_diffusion_decoder"),
            )
        except Exception as cleanup_exc:
            request_logger.warning("maintainer_probe.cleanup warning=%s: %s", type(cleanup_exc).__name__, cleanup_exc)
        gc.collect()
        if torch.cuda.is_available():
            torch.cuda.empty_cache()
        request_logger.info("maintainer_probe.end profile=%s status=%s", profile_id, status)
        close_request_logger(request_logger, request_handler)


def run_full_sft_30s_stage2_diffusion_scout(**kwargs):
    return _run_full_sft_30s_tail_scout(
        profile_id=FULL_SFT_30S_STAGE2_DIFFUSION_PROFILE, stop_after="stage2_diffusion_decoder", current_gate=False,
        requested_duration_seconds=FULL_SFT_30S_STAGE2_DIFFUSION_ZEROGPU_DURATION_SECONDS, **kwargs
    )


def run_full_sft_30s_stage2_conv_scout(**kwargs):
    return _run_full_sft_30s_tail_scout(
        profile_id=FULL_SFT_30S_STAGE2_CONV_PROFILE, stop_after="stage2_conv_vae", current_gate=False,
        requested_duration_seconds=FULL_SFT_30S_STAGE2_CONV_ZEROGPU_DURATION_SECONDS, **kwargs
    )


def run_full_sft_30s_stage2_conv_framewise_scout(**kwargs):
    return _run_full_sft_30s_tail_scout(
        profile_id=FULL_SFT_30S_STAGE2_CONV_FRAMEWISE_PROFILE, stop_after="stage2_conv_framewise", current_gate=False,
        requested_duration_seconds=FULL_SFT_30S_STAGE2_CONV_FRAMEWISE_ZEROGPU_DURATION_SECONDS, **kwargs
    )


def run_full_sft_30s_stage2_latent_scout(**kwargs):
    return _run_full_sft_30s_tail_scout(
        profile_id=FULL_SFT_30S_STAGE2_LATENT_PROFILE, stop_after="stage2_latent", current_gate=False,
        requested_duration_seconds=FULL_SFT_30S_STAGE2_LATENT_ZEROGPU_DURATION_SECONDS, **kwargs
    )


def run_full_sft_30s_upsample_scout(**kwargs):
    return _run_full_sft_30s_tail_scout(
        profile_id=FULL_SFT_30S_UPSAMPLE_PROFILE, stop_after="upsample", current_gate=False,
        requested_duration_seconds=FULL_SFT_30S_UPSAMPLE_ZEROGPU_DURATION_SECONDS, **kwargs
    )



def run_full_sft_30s_flf2v_stage1_scout(
    *,
    session_id,
    progress,
    candidate_id: str,
    runtime: generation_backend.GenerationRuntime,
    hooks: generation_backend.GenerationHooks,
    flf2v_example: tuple[str, str, str],
    create_request_paths: Callable[[Any], runtime_utils.RequestPaths],
    open_request_logger: Callable,
    close_request_logger: Callable,
):
    """Exercise the exact 721f Start+End condition-pipeline Stage 1 with two denoising steps."""
    if not runtime.is_full_sft_profile:
        raise gr.Error("This Probe requires MODEL_RUNTIME_PROFILE=full_sft_nf4 and a Space restart.")
    if runtime.preload_state.get("status") != "ready" or runtime.pipe is None or runtime.pipe_condition is None:
        raise gr.Error(f"Diffusers condition-pipeline preload is not ready: {runtime.preload_state}")

    start_image_path, end_image_path, probe_prompt = flf2v_example
    paths = create_request_paths(session_id)
    request_logger, request_handler = open_request_logger(paths.log, paths.request_id)
    video_latents = None
    audio_latents = None
    started = time.monotonic()
    status = "FAIL"
    try:
        request_logger.info(
            "maintainer_probe.start profile=%s mode=FLF2V size=512x512 frames=%s steps=%s seed=42 no_user_lora=true",
            FULL_SFT_30S_FLF2V_STAGE1_PROFILE,
            FULL_SFT_30S_STAGE1_FRAMES,
            FULL_SFT_30S_FLF2V_STAGE1_STEPS,
        )
        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()
        runtime_utils.reset_gpu_peak_memory()
        gpu_before = hooks.gpu_state()
        execution_environment = runtime_utils.execution_environment_identity(runtime.preload_state.get("environment"))

        generator = torch.Generator("cuda").manual_seed(42)
        call_kwargs = generation_backend._build_stage1_call_kwargs(
            prompt=probe_prompt,
            generator=generator,
            mode="FLF2V",
            start_image_path=start_image_path,
            middle_image_path=None,
            end_image_path=end_image_path,
            num_frames=FULL_SFT_30S_STAGE1_FRAMES,
            width=512,
            height=512,
            runtime=runtime,
            hooks=hooks,
        )
        conditions = list(call_kwargs.get("conditions") or [])
        condition_indices = [int(item.index) for item in conditions]
        if condition_indices != [0, -1]:
            raise RuntimeError(f"Expected FLF2V condition indices [0, -1], observed {condition_indices!r}.")

        step_end_times: list[float] = []

        def _step_probe_callback(_pipeline, step_index, _timestep, callback_kwargs):
            step_end_times.append(time.monotonic())
            completed = int(step_index) + 1
            progress(
                0.12 + (0.72 * (completed / FULL_SFT_30S_FLF2V_STAGE1_STEPS)),
                desc=f"Maintainer Probe · FLF2V Stage 1 step {completed}/{FULL_SFT_30S_FLF2V_STAGE1_STEPS}",
            )
            return callback_kwargs

        progress(0.08, desc="Maintainer Probe · 30s FLF2V Start+End condition preparation")
        stage1_started = time.monotonic()
        video_latents, audio_latents = runtime.pipe_condition(
            height=256,
            width=256,
            output_type="latent",
            num_inference_steps=FULL_SFT_30S_FLF2V_STAGE1_STEPS,
            num_frames=FULL_SFT_30S_STAGE1_FRAMES,
            callback_on_step_end=_step_probe_callback,
            callback_on_step_end_tensor_inputs=[],
            **call_kwargs,
        )
        stage1_seconds = time.monotonic() - stage1_started
        gpu_after_stage1 = hooks.gpu_state()
        if len(step_end_times) != FULL_SFT_30S_FLF2V_STAGE1_STEPS:
            raise RuntimeError(
                f"Expected {FULL_SFT_30S_FLF2V_STAGE1_STEPS} step callbacks, observed {len(step_end_times)}."
            )

        temporal_ratio = int(getattr(runtime.pipe_condition, "vae_temporal_compression_ratio", 8) or 8)
        expected_last_latent_index = (FULL_SFT_30S_STAGE1_FRAMES - 1) // temporal_ratio
        run_info = {
            "schema_version": "ltx25-maintainer-probe-v2",
            "candidate_id": candidate_id,
            "request_id": paths.request_id,
            "status": "PASS",
            "surface": "maintainer_probe",
            "profile_id": FULL_SFT_30S_FLF2V_STAGE1_PROFILE,
            "current_gate": True,
            "contract": {
                "runtime_profile": runtime.runtime_profile,
                "mode": "FLF2V",
                "example": "Blue bird · first + last frame",
                "width": 512,
                "height": 512,
                "frames": FULL_SFT_30S_STAGE1_FRAMES,
                "shape_video_seconds": 30.0,
                "stage1_inference_steps": FULL_SFT_30S_FLF2V_STAGE1_STEPS,
                "seed": 42,
                "user_lora_count": 0,
                "scope": "stage1_only_latent_output",
                "condition_indices_requested": condition_indices,
                "expected_last_latent_index_after_negative_index_resolution": expected_last_latent_index,
                "zerogpu_requested_duration_seconds": FULL_SFT_30S_FLF2V_STAGE1_ZEROGPU_DURATION_SECONDS,
            },
            "timing": {"elapsed_seconds": time.monotonic() - started, "stage1_2step_seconds": stage1_seconds},
            "video_latent_shape": list(video_latents.shape),
            "audio_latent_shape": list(audio_latents.shape),
            "gpu_before": gpu_before,
            "gpu_after_stage1": gpu_after_stage1,
            "ru_maxrss_kib": runtime_utils.process_rss_kib(),
            "start_image_sha256": hooks.sha256_file(start_image_path),
            "end_image_sha256": hooks.sha256_file(end_image_path),
            "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"),
        }
        paths.run_info.write_text(json.dumps(run_info, indent=2, sort_keys=True), encoding="utf-8")
        paths.diagnostics.write_text(json.dumps(run_info, indent=2, sort_keys=True), encoding="utf-8")
        status = "PASS"
        request_logger.info(
            "maintainer_probe.success profile=%s elapsed_seconds=%.3f stage1_2step_seconds=%.3f last_latent_index=%s",
            FULL_SFT_30S_FLF2V_STAGE1_PROFILE,
            run_info["timing"]["elapsed_seconds"],
            stage1_seconds,
            expected_last_latent_index,
        )
        try:
            request_handler.flush()
        except Exception:
            pass
        progress(0.94, desc="Maintainer Probe · packaging FLF2V evidence")
        probe_path = probe_artifacts.build_request_probe(
            paths,
            candidate_id=candidate_id,
            status="PASS",
            sha256_file=hooks.sha256_file,
            metadata={"surface": "maintainer_probe", "profile_id": FULL_SFT_30S_FLF2V_STAGE1_PROFILE},
        )
        peak_reserved = float(gpu_after_stage1.get("max_reserved_bytes") or 0) / (1024**3)
        progress(1.0, desc="Maintainer Probe · FLF2V scout complete")
        return (
            f"**PASS · {FULL_SFT_30S_FLF2V_STAGE1_PROFILE}**  \n"
            f"721f Start+End / two-step Stage 1: **{stage1_seconds:.2f}s** · peak reserved: **{peak_reserved:.2f} GiB**.  \n"
            "This is a structural long-FLF2V witness only; it does not claim 30-step visual-quality closure.",
            str(probe_path),
        )
    except Exception as exc:
        elapsed = time.monotonic() - started
        request_logger.exception("maintainer_probe.failure error_type=%s error=%s", type(exc).__name__, exc)
        failure = {
            "schema_version": "ltx25-maintainer-probe-v2",
            "candidate_id": candidate_id,
            "request_id": paths.request_id,
            "status": "FAIL",
            "surface": "maintainer_probe",
            "profile_id": FULL_SFT_30S_FLF2V_STAGE1_PROFILE,
            "current_gate": True,
            "elapsed_seconds": elapsed,
            "error_type": type(exc).__name__,
            "error": str(exc),
            "gpu_after_failure": hooks.gpu_state(),
        }
        try:
            paths.diagnostics.write_text(json.dumps(failure, indent=2, sort_keys=True), encoding="utf-8")
            request_handler.flush()
            probe_artifacts.build_request_probe(
                paths,
                candidate_id=candidate_id,
                status="FAIL",
                sha256_file=hooks.sha256_file,
                metadata={"surface": "maintainer_probe", "profile_id": FULL_SFT_30S_FLF2V_STAGE1_PROFILE},
            )
        except Exception as evidence_exc:
            request_logger.exception(
                "maintainer_probe.evidence.failure error_type=%s error=%s",
                type(evidence_exc).__name__, evidence_exc,
            )
        if isinstance(exc, gr.Error):
            raise
        raise gr.Error(f"Maintainer Probe failed: {type(exc).__name__}: {exc}") from exc
    finally:
        video_latents = None
        audio_latents = 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(
                "maintainer_probe.cleanup warning=%s: %s", type(cleanup_exc).__name__, cleanup_exc
            )
        gc.collect()
        if torch.cuda.is_available():
            torch.cuda.empty_cache()
        request_logger.info("maintainer_probe.end profile=%s status=%s", FULL_SFT_30S_FLF2V_STAGE1_PROFILE, status)
        close_request_logger(request_logger, request_handler)

def run_full_sft_30s_stage1_scout(
    *,
    session_id,
    progress,
    candidate_id: str,
    runtime: generation_backend.GenerationRuntime,
    hooks: generation_backend.GenerationHooks,
    cats_example: tuple[str, str],
    create_request_paths: Callable[[Any], runtime_utils.RequestPaths],
    open_request_logger: Callable,
    close_request_logger: Callable,
):
    """Run the one current low-quota gate Probe; no Stage 2/decode/encode."""
    if not runtime.is_full_sft_profile:
        raise gr.Error("This Probe requires MODEL_RUNTIME_PROFILE=full_sft_nf4 and a Space restart.")
    if runtime.preload_state.get("status") != "ready" or runtime.pipe is None or runtime.pipe_i2v is None:
        raise gr.Error(f"Diffusers preload is not ready: {runtime.preload_state}")

    start_image_path, probe_prompt = cats_example
    paths = create_request_paths(session_id)
    request_logger, request_handler = open_request_logger(paths.log, paths.request_id)
    video_latents = None
    audio_latents = None
    step_end_times: list[float] = []
    started = time.monotonic()
    status = "FAIL"
    try:
        request_logger.info(
            "maintainer_probe.start profile=%s mode=I2V size=512x512 frames=%s steps=%s seed=42 no_user_lora=true",
            FULL_SFT_30S_STAGE1_PROFILE,
            FULL_SFT_30S_STAGE1_FRAMES,
            FULL_SFT_30S_STAGE1_STEPS,
        )
        # Match the ordinary Full/SFT Stage-1 baseline: no request adapters active,
        # internal Stage-2 distilled adapter retained CPU-resident and inactive.
        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()
        runtime_utils.reset_gpu_peak_memory()
        gpu_before = hooks.gpu_state()
        execution_environment = runtime_utils.execution_environment_identity(runtime.preload_state.get("environment"))

        generator = torch.Generator("cuda").manual_seed(42)
        call_kwargs = generation_backend._build_stage1_call_kwargs(
            prompt=probe_prompt,
            generator=generator,
            mode="I2V",
            start_image_path=start_image_path,
            middle_image_path=None,
            end_image_path=None,
            num_frames=FULL_SFT_30S_STAGE1_FRAMES,
            width=512,
            height=512,
            runtime=runtime,
            hooks=hooks,
        )

        def _step_probe_callback(_pipeline, step_index, _timestep, callback_kwargs):
            step_end_times.append(time.monotonic())
            progress(
                0.12 + (0.72 * ((int(step_index) + 1) / FULL_SFT_30S_STAGE1_STEPS)),
                desc=f"Maintainer Probe · Stage 1 step {int(step_index) + 1}/{FULL_SFT_30S_STAGE1_STEPS}",
            )
            return callback_kwargs

        progress(0.08, desc="Maintainer Probe · preparing exact 721f Stage-1 shape")
        stage1_started = time.monotonic()
        video_latents, audio_latents = runtime.pipe_i2v(
            height=256,
            width=256,
            output_type="latent",
            num_inference_steps=FULL_SFT_30S_STAGE1_STEPS,
            num_frames=FULL_SFT_30S_STAGE1_FRAMES,
            callback_on_step_end=_step_probe_callback,
            callback_on_step_end_tensor_inputs=[],
            **call_kwargs,
        )
        stage1_finished = time.monotonic()
        stage1_seconds = stage1_finished - stage1_started
        gpu_after_stage1 = hooks.gpu_state()

        if len(step_end_times) != FULL_SFT_30S_STAGE1_STEPS:
            raise RuntimeError(
                f"Expected {FULL_SFT_30S_STAGE1_STEPS} step callbacks, observed {len(step_end_times)}."
            )
        first_step_end_seconds = step_end_times[0] - stage1_started
        steady_intervals = [
            step_end_times[idx] - step_end_times[idx - 1] for idx in range(1, len(step_end_times))
        ]
        steady_step_seconds = sum(steady_intervals) / len(steady_intervals)
        post_step_seconds = stage1_finished - step_end_times[-1]
        projected_stage1_30 = first_step_end_seconds + (29.0 * steady_step_seconds) + post_step_seconds

        # Explicit P20R0 Full/SFT no-user-LoRA historical witnesses. Keep these
        # numbers visible here rather than introducing a generic estimator layer.
        stage2_49f_seconds = 3.820
        stage2_121f_seconds = 7.054
        stage2_slope_seconds_per_frame = (stage2_121f_seconds - stage2_49f_seconds) / (121 - 49)
        projected_stage2_721 = stage2_49f_seconds + (
            (FULL_SFT_30S_STAGE1_FRAMES - 49) * stage2_slope_seconds_per_frame
        )
        adapter_transfer_release_guard = 7.946
        cpu_encode_misc_guard = 3.0
        projected_e2e_center = (
            projected_stage1_30 + projected_stage2_721 + adapter_transfer_release_guard + cpu_encode_misc_guard
        )
        suggested_exact_reservation = int(projected_e2e_center * 1.20 + 0.999999)

        projection = {
            "purpose": "scout_projection_not_live_runtime_claim",
            "stage1_30step_formula": "first_step_end + 29 * mean(step2_to_step5_intervals) + post_last_step",
            "first_step_end_seconds": first_step_end_seconds,
            "steady_step_intervals_seconds": steady_intervals,
            "steady_step_mean_seconds": steady_step_seconds,
            "post_last_step_seconds": post_step_seconds,
            "projected_stage1_30step_seconds": projected_stage1_30,
            "historical_stage2_points": {"49f_seconds": stage2_49f_seconds, "121f_seconds": stage2_121f_seconds},
            "stage2_linear_slope_seconds_per_frame": stage2_slope_seconds_per_frame,
            "projected_stage2_721f_seconds": projected_stage2_721,
            "adapter_transfer_release_guard_seconds": adapter_transfer_release_guard,
            "cpu_encode_misc_guard_seconds": cpu_encode_misc_guard,
            "projected_e2e_center_seconds": projected_e2e_center,
            "suggested_exact_run_reservation_seconds_20pct_guard": suggested_exact_reservation,
        }
        run_info = {
            "schema_version": "ltx25-maintainer-probe-v1",
            "candidate_id": candidate_id,
            "request_id": paths.request_id,
            "status": "PASS",
            "surface": "maintainer_probe",
            "profile_id": FULL_SFT_30S_STAGE1_PROFILE,
            "current_gate": False,
            "contract": {
                "runtime_profile": runtime.runtime_profile,
                "mode": "I2V",
                "example": "Cats · I2V",
                "width": 512,
                "height": 512,
                "frames": FULL_SFT_30S_STAGE1_FRAMES,
                "shape_video_seconds": 30.0,
                "stage1_inference_steps": FULL_SFT_30S_STAGE1_STEPS,
                "seed": 42,
                "user_lora_count": 0,
                "scope": "stage1_only_latent_output",
                "zerogpu_requested_duration_seconds": FULL_SFT_30S_STAGE1_ZEROGPU_DURATION_SECONDS,
            },
            "stage1_seconds": stage1_seconds,
            "projection": projection,
            "gpu_before": gpu_before,
            "gpu_after_stage1": gpu_after_stage1,
            "ru_maxrss_kib": runtime_utils.process_rss_kib(),
            "input_image_sha256": hooks.sha256_file(start_image_path),
            "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"),
        }
        paths.run_info.write_text(json.dumps(run_info, indent=2, sort_keys=True), encoding="utf-8")
        paths.diagnostics.write_text(json.dumps(run_info, indent=2, sort_keys=True), encoding="utf-8")
        status = "PASS"
        request_logger.info(
            "maintainer_probe.success stage1_seconds=%.3f projected_stage1_30step_seconds=%.3f projected_e2e_center_seconds=%.3f suggested_reservation_seconds=%s",
            stage1_seconds,
            projected_stage1_30,
            projected_e2e_center,
            suggested_exact_reservation,
        )
        try:
            request_handler.flush()
        except Exception:
            pass
        probe_path = probe_artifacts.build_request_probe(
            paths,
            candidate_id=candidate_id,
            status="PASS",
            sha256_file=hooks.sha256_file,
            metadata={"surface": "maintainer_probe", "profile_id": FULL_SFT_30S_STAGE1_PROFILE},
        )
        peak_reserved = float(gpu_after_stage1.get("max_reserved_bytes") or 0) / (1024**3)
        return (
            f"**PASS · {FULL_SFT_30S_STAGE1_PROFILE}**  \n"
            f"5-step Stage 1: **{stage1_seconds:.2f}s** · peak reserved: **{peak_reserved:.2f} GiB**.  \n"
            f"Projected 30-step Stage 1: **{projected_stage1_30:.1f}s**. Provisional E2E center: **{projected_e2e_center:.1f}s**; "
            f"20% guard suggests **{suggested_exact_reservation}s** for the later exact run. Return the Probe ZIP before treating that projection as adopted.",
            str(probe_path),
        )
    except Exception as exc:
        elapsed = time.monotonic() - started
        request_logger.exception("maintainer_probe.failure error_type=%s error=%s", type(exc).__name__, exc)
        failure = {
            "schema_version": "ltx25-maintainer-probe-v1",
            "candidate_id": candidate_id,
            "request_id": paths.request_id,
            "status": "FAIL",
            "surface": "maintainer_probe",
            "profile_id": FULL_SFT_30S_STAGE1_PROFILE,
            "elapsed_seconds": elapsed,
            "error_type": type(exc).__name__,
            "error": str(exc),
            "gpu_after_failure": hooks.gpu_state(),
        }
        try:
            paths.diagnostics.write_text(json.dumps(failure, indent=2, sort_keys=True), encoding="utf-8")
            request_handler.flush()
            probe_artifacts.build_request_probe(
                paths,
                candidate_id=candidate_id,
                status="FAIL",
                sha256_file=hooks.sha256_file,
                metadata={"surface": "maintainer_probe", "profile_id": FULL_SFT_30S_STAGE1_PROFILE},
            )
        except Exception as evidence_exc:
            request_logger.exception(
                "maintainer_probe.evidence.failure error_type=%s error=%s",
                type(evidence_exc).__name__,
                evidence_exc,
            )
        if isinstance(exc, gr.Error):
            raise
        raise gr.Error(f"Maintainer Probe failed: {type(exc).__name__}: {exc}") from exc
    finally:
        video_latents = None
        audio_latents = 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(
                "maintainer_probe.cleanup warning=%s: %s", type(cleanup_exc).__name__, cleanup_exc
            )
        gc.collect()
        if torch.cuda.is_available():
            torch.cuda.empty_cache()
        request_logger.info("maintainer_probe.end status=%s", status)
        close_request_logger(request_logger, request_handler)