ltx/ic_lora.py

"""Shared IC-LoRA runtime primitives.

Keep task semantics (prompting, preprocessing, task evidence and UI policy) in the
individual task module. This module only owns reusable adapter/reference mechanics
that are expected to be shared by future IC-LoRA tabs.
"""
from __future__ import annotations

import gc
import time
from dataclasses import dataclass
from pathlib import Path

import numpy as np
import PIL.Image
import torch
from diffusers.pipelines.ltx2.pipeline_ltx2_ic_lora import LTX2ReferenceCondition


@dataclass(frozen=True)
class PreparedRequestFiles:
    """Validated request-local files needed by an IC task GPU callback."""

    reference_cache: Path
    local_lora: Path
    video_path: Path
    diagnostics_path: Path


def adapter_state(pipe) -> dict:
    """Return the public adapter inventory without assuming one Diffusers version."""
    state = {}
    for name in ("get_active_adapters", "get_list_adapters"):
        fn = getattr(pipe, name, None)
        if callable(fn):
            try:
                state[name] = fn()
            except Exception as exc:
                state[name] = {"error": f"{type(exc).__name__}: {exc}"}
    return state


def adapter_names(state: dict) -> set[str]:
    """Return adapter names from Diffusers public adapter-state snapshots."""
    names: set[str] = set()

    def _collect(value) -> None:
        if isinstance(value, dict):
            for item in value.values():
                _collect(item)
        elif isinstance(value, (list, tuple, set)):
            for item in value:
                _collect(item)
        elif isinstance(value, str) and value.strip() and not value.strip().lower().startswith("error"):
            names.add(value.strip())

    _collect(state.get("get_active_adapters"))
    _collect(state.get("get_list_adapters"))
    return names


def adapter_state_has_any(state: dict) -> bool:
    """True when a Diffusers adapter-state snapshot contains a real adapter."""
    def _nonempty(value) -> bool:
        if isinstance(value, dict):
            return any(_nonempty(item) for item in value.values())
        if isinstance(value, (list, tuple, set)):
            return len(value) > 0
        if isinstance(value, str):
            return bool(value.strip()) and not value.strip().lower().startswith("error")
        return bool(value)

    return _nonempty(state.get("get_active_adapters")) or _nonempty(state.get("get_list_adapters"))


def load_adapter_weights(*, pipe, local_lora: Path, adapter_name: str) -> None:
    """Load one already-local IC adapter; never performs Hub I/O."""
    pipe.load_lora_weights(
        str(local_lora.parent),
        weight_name=local_lora.name,
        adapter_name=adapter_name,
    )


def activate_adapter(*, pipe, adapter_name: str, strength: float) -> None:
    """Select and enable one loaded request-scoped IC adapter."""
    pipe.set_adapters([adapter_name], adapter_weights=[float(strength)])
    pipe.enable_lora()

def release_adapter(*, pipe, adapter_name: str, state_pipe=None) -> tuple[float, dict]:
    """Disable/delete one request-scoped IC adapter and release cached CUDA memory."""
    started = time.monotonic()
    pipe.disable_lora()
    pipe.delete_adapters([adapter_name])
    gc.collect()
    if torch.cuda.is_available():
        torch.cuda.empty_cache()
    return time.monotonic() - started, adapter_state(state_pipe if state_pipe is not None else pipe)


def force_release_adapter(*, pipe, adapter_name: str) -> None:
    """Best-effort finalizer used only after an IC callback failed mid-lifecycle."""
    try:
        pipe.disable_lora()
        pipe.delete_adapters([adapter_name])
    finally:
        gc.collect()
        if torch.cuda.is_available():
            torch.cuda.empty_cache()


def reference_condition_from_cache(
    cache_path: Path,
    *,
    num_frames: int,
    height: int,
    width: int,
    strength: float,
) -> LTX2ReferenceCondition:
    """Restore an RGB frame cache and build the standard LTX2 reference condition."""
    with np.load(cache_path, allow_pickle=False) as cached:
        reference_array = np.asarray(cached["frames"], dtype=np.uint8)
    expected_shape = (int(num_frames), int(height), int(width), 3)
    if tuple(reference_array.shape) != expected_shape:
        raise RuntimeError(f"Prepared reference shape mismatch: {reference_array.shape} != {expected_shape}")
    reference_frames = [PIL.Image.fromarray(frame).convert("RGB") for frame in reference_array]
    return LTX2ReferenceCondition(frames=reference_frames, strength=float(strength))