ltxc/model_runtime.py

from __future__ import annotations

import gc
import inspect
import math
import os
import time
from dataclasses import dataclass
from typing import Any

from .config import STRICT_WEIGHT_AUDIT, WEIGHT_AUDIT_MODE
from .diagnostics import bnb_model_audit, count_adapter_state, raw_lora_audit
from .runtime_reporting import startup_event


@dataclass
class LoadedColorizer:
    """One process-wide NF4 + unfused-LoRA colorization pipeline."""

    pipe: Any
    transformer: Any
    text_encoder: Any
    tokenizer: Any
    gemma_audit: dict
    transformer_audit: dict
    adapter_audit: dict
    lora_audit: dict
    startup_load_seconds: float


_GLOBAL_COLORIZER: LoadedColorizer | None = None
_BNB_TUPLE_SHAPE_COMPAT_INSTALLED = False


def _install_diffusers_bnb_tuple_shape_compat() -> bool:
    """Preserve Diffusers' NF4 shape check when ZeroGPU exposes tuple shapes.

    The pinned Diffusers revision calls ``current_param.shape.numel()`` while
    restoring prequantized BnB 4-bit weights. ZeroGPU startup CUDA emulation can
    expose the same shape as a plain tuple. Replace only that check with an
    equivalent implementation using ``math.prod`` for tuple-like shapes.
    """
    global _BNB_TUPLE_SHAPE_COMPAT_INSTALLED
    if _BNB_TUPLE_SHAPE_COMPAT_INSTALLED:
        return False

    from diffusers.quantizers.bitsandbytes.bnb_quantizer import (
        BnB4BitDiffusersQuantizer,
    )

    original = BnB4BitDiffusersQuantizer.check_quantized_param_shape
    if getattr(original, "_ltxc_tuple_shape_compatible", False):
        _BNB_TUPLE_SHAPE_COMPAT_INSTALLED = True
        return False

    def check_quantized_param_shape(self, param_name, current_param, loaded_param):
        current_shape = current_param.shape
        loaded_shape = loaded_param.shape
        n = (
            current_shape.numel()
            if hasattr(current_shape, "numel")
            else math.prod(current_shape)
        )
        inferred_shape = (n,) if "bias" in param_name else ((n + 1) // 2, 1)
        if loaded_shape != inferred_shape:
            raise ValueError(
                "Expected the flattened shape of the current param "
                f"({param_name}) to be {loaded_shape} but is {inferred_shape}."
            )
        return True

    check_quantized_param_shape._ltxc_tuple_shape_compatible = True
    check_quantized_param_shape._ltxc_original = original
    BnB4BitDiffusersQuantizer.check_quantized_param_shape = check_quantized_param_shape
    _BNB_TUPLE_SHAPE_COMPAT_INSTALLED = True
    return True


def _dtype_kwargs(loader, torch_module) -> dict:
    parameter = (
        "dtype" if "dtype" in inspect.signature(loader).parameters else "torch_dtype"
    )
    return {parameter: torch_module.bfloat16}


def _attention_policy(text_encoder) -> dict:
    text_config = getattr(text_encoder.config, "text_config", None)
    return {
        "requested_attn_implementation": None,
        "model_config_resolved": getattr(
            text_encoder.config,
            "_attn_implementation",
            None,
        ),
        "text_config_resolved": (
            getattr(text_config, "_attn_implementation", None)
            if text_config is not None
            else None
        ),
        "compile_disabled": os.environ.get("TORCH_COMPILE_DISABLE") == "1",
        "dynamo_disabled": os.environ.get("TORCHDYNAMO_DISABLE") == "1",
    }


def _quantization_config_summary(model) -> dict | None:
    config = getattr(model, "config", None)
    quantization_config = getattr(config, "quantization_config", None)
    if hasattr(quantization_config, "to_dict"):
        quantization_config = quantization_config.to_dict()
    return quantization_config


def _minimal_model_summary(model) -> dict:
    """Record cheap load metadata without traversing modules or parameters."""
    return {
        "audit_mode": "minimal",
        "strict_weight_audit": False,
        "class": model.__class__.__name__,
        "is_loaded_in_4bit": bool(getattr(model, "is_loaded_in_4bit", False)),
        "device_map": getattr(model, "hf_device_map", None),
        "quantization_config": _quantization_config_summary(model),
        "module_scan_performed": False,
        "parameter_scan_performed": False,
    }


def _minimal_adapter_summary(pipe) -> dict:
    """Record adapter API state without walking LoRA modules or reading tensors."""
    active = None
    listed = None
    try:
        active = pipe.get_active_adapters()
    except Exception:
        pass
    try:
        listed = pipe.get_list_adapters()
    except Exception:
        pass
    return {
        "audit_mode": "minimal",
        "strict_weight_audit": False,
        "load_lora_weights_succeeded": True,
        "set_adapters_succeeded": True,
        "active_adapters": active,
        "listed_adapters": listed,
        "adapter_modules": None,
        "lora_A_tensors": None,
        "lora_B_tensors": None,
        "sampled_nonzero_tensors": None,
        "module_scan_performed": False,
        "value_sampling_performed": False,
    }


def _validate_loaded_models(
    gemma_audit: dict,
    transformer_audit: dict,
    adapter_audit: dict,
    *,
    require_adapter_values: bool,
) -> None:
    failures = []
    if not gemma_audit["is_loaded_in_4bit"] and gemma_audit["params4bit_count"] == 0:
        failures.append("Gemma did not report a BnB 4-bit load.")
    if transformer_audit["params4bit_count"] <= 0:
        failures.append("Transformer did not load prequantized BnB 4-bit weights.")
    if adapter_audit["adapter_modules"] <= 0:
        failures.append("No PEFT LoRA adapter modules were found in the transformer.")
    if require_adapter_values and adapter_audit["sampled_nonzero_tensors"] <= 0:
        failures.append("Sampled LoRA tensors were all zero or unavailable.")
    if failures:
        raise RuntimeError("; ".join(failures))


def _collect_cuda_garbage() -> None:
    gc.collect()
    try:
        import torch

        if torch.cuda.is_available():
            torch.cuda.empty_cache()
            torch.cuda.ipc_collect()
    except Exception:
        pass


def _load_colorizer_pipeline(assets: dict) -> LoadedColorizer:
    """Load once at module startup and explicitly retain the pipeline on CUDA."""
    import torch
    from diffusers import LTX2InContextPipeline, LTX2VideoTransformer3DModel
    from transformers import AutoTokenizer, Gemma3ForConditionalGeneration

    compat_installed = _install_diffusers_bnb_tuple_shape_compat()
    startup_event(
        "compat.diffusers_bnb_tuple_shape",
        "Installed tuple-compatible prequantized NF4 shape validation.",
        installed_now=compat_installed,
    )

    started = time.perf_counter()
    pipe = None
    transformer = None
    text_encoder = None
    tokenizer = None
    try:
        tokenizer = AutoTokenizer.from_pretrained(
            assets["gemma_dir"],
            local_files_only=True,
            use_fast=True,
        )

        gemma_kwargs = {
            "local_files_only": True,
            "low_cpu_mem_usage": True,
            "device_map": {"": 0},
            **_dtype_kwargs(
                Gemma3ForConditionalGeneration.from_pretrained,
                torch,
            ),
        }
        text_encoder = Gemma3ForConditionalGeneration.from_pretrained(
            assets["gemma_dir"],
            **gemma_kwargs,
        )
        gemma_audit = (
            bnb_model_audit(text_encoder)
            if STRICT_WEIGHT_AUDIT
            else _minimal_model_summary(text_encoder)
        )
        gemma_audit["audit_mode"] = WEIGHT_AUDIT_MODE
        gemma_audit["strict_weight_audit"] = STRICT_WEIGHT_AUDIT
        gemma_audit["attention_policy"] = _attention_policy(text_encoder)

        transformer_kwargs = {
            "local_files_only": True,
            "low_cpu_mem_usage": True,
            "device_map": {"": 0},
            **_dtype_kwargs(
                LTX2VideoTransformer3DModel.from_pretrained,
                torch,
            ),
        }
        transformer = LTX2VideoTransformer3DModel.from_pretrained(
            assets["transformer_dir"],
            subfolder="transformer",
            **transformer_kwargs,
        )
        transformer_audit = (
            bnb_model_audit(transformer)
            if STRICT_WEIGHT_AUDIT
            else _minimal_model_summary(transformer)
        )
        transformer_audit["audit_mode"] = WEIGHT_AUDIT_MODE
        transformer_audit["strict_weight_audit"] = STRICT_WEIGHT_AUDIT

        pipe_kwargs = {
            "transformer": transformer,
            "text_encoder": text_encoder,
            "tokenizer": tokenizer,
            "local_files_only": True,
            "low_cpu_mem_usage": True,
            "use_safetensors": True,
            "device_map": "cuda",
        }
        try:
            pipe = LTX2InContextPipeline.from_pretrained(
                assets["base_dir"],
                torch_dtype=torch.bfloat16,
                **pipe_kwargs,
            )
        except TypeError as exc:
            if "torch_dtype" not in str(exc):
                raise
            pipe = LTX2InContextPipeline.from_pretrained(
                assets["base_dir"],
                dtype=torch.bfloat16,
                **pipe_kwargs,
            )

        if hasattr(pipe.vae, "enable_tiling"):
            pipe.vae.enable_tiling()
        lora_audit = raw_lora_audit(
            assets["lora_path"],
            strict=STRICT_WEIGHT_AUDIT,
        )
        if not lora_audit.get("ok"):
            raise RuntimeError(
                "Unable to read Colorization LoRA metadata: "
                f"{lora_audit.get('error_type')}: {lora_audit.get('error_message')}"
            )
        pipe.load_lora_weights(
            assets["lora_path"],
            adapter_name="colorizer",
        )
        pipe.set_adapters("colorizer", 1.0)

        # Keep this explicit even though the prequantized components use CUDA device maps.
        # On ZeroGPU this is handled by startup CUDA emulation; on a dedicated GPU it
        # performs the normal persistent CUDA placement. No warmup or inference runs here.
        pipe = pipe.to("cuda")

        if STRICT_WEIGHT_AUDIT:
            adapter_audit = count_adapter_state(pipe, sample_values=False)
            adapter_audit["audit_mode"] = "strict"
            adapter_audit["strict_weight_audit"] = True
            _validate_loaded_models(
                gemma_audit,
                transformer_audit,
                adapter_audit,
                require_adapter_values=False,
            )
        else:
            adapter_audit = _minimal_adapter_summary(pipe)
        startup_load_seconds = time.perf_counter() - started

        return LoadedColorizer(
            pipe=pipe,
            transformer=transformer,
            text_encoder=text_encoder,
            tokenizer=tokenizer,
            gemma_audit=gemma_audit,
            transformer_audit=transformer_audit,
            adapter_audit=adapter_audit,
            lora_audit=lora_audit,
            startup_load_seconds=startup_load_seconds,
        )
    except Exception:
        if pipe is not None:
            try:
                pipe.maybe_free_model_hooks()
            except Exception:
                pass
        pipe = transformer = text_encoder = tokenizer = None
        _collect_cuda_garbage()
        raise


def initialize_global_colorizer(assets: dict) -> LoadedColorizer:
    """Initialize the process-wide CUDA-resident pipeline exactly once."""
    global _GLOBAL_COLORIZER
    if _GLOBAL_COLORIZER is not None:
        return _GLOBAL_COLORIZER

    startup_event(
        "model.startup_load.start",
        "Loading the fixed colorizer onto CUDA before the Gradio UI is built.",
    )
    loaded = _load_colorizer_pipeline(assets)
    _GLOBAL_COLORIZER = loaded
    startup_event(
        "model.startup_load.complete",
        "The fixed colorizer is resident on CUDA.",
        startup_load_seconds=round(loaded.startup_load_seconds, 3),
        weight_audit_mode=WEIGHT_AUDIT_MODE,
        strict_weight_audit=STRICT_WEIGHT_AUDIT,
        gemma_params4bit=loaded.gemma_audit.get("params4bit_count"),
        transformer_params4bit=loaded.transformer_audit.get("params4bit_count"),
        adapter_modules=loaded.adapter_audit.get("adapter_modules"),
        lora_metadata_loaded=bool(loaded.lora_audit.get("ok")),
        lora_sha256_performed=bool(loaded.lora_audit.get("sha256_performed")),
        explicit_pipeline_to_cuda=True,
        adapter_value_sampling_performed=False,
    )
    return loaded


def get_global_colorizer() -> LoadedColorizer:
    """Return the startup-loaded pipeline; never lazy-load inside a GPU callback."""
    if _GLOBAL_COLORIZER is None:
        raise RuntimeError(
            "Global colorizer is not initialized. Initialize it before building the UI."
        )
    return _GLOBAL_COLORIZER


def audit_global_colorizer_for_callback(loaded: LoadedColorizer) -> dict:
    """Run expensive LoRA value checks only when explicitly requested."""
    if not STRICT_WEIGHT_AUDIT:
        return _minimal_adapter_summary(loaded.pipe)

    adapter_audit = count_adapter_state(loaded.pipe, sample_values=True)
    adapter_audit["audit_mode"] = "strict"
    adapter_audit["strict_weight_audit"] = True
    _validate_loaded_models(
        loaded.gemma_audit,
        loaded.transformer_audit,
        adapter_audit,
        require_adapter_values=True,
    )
    return adapter_audit