model_loader.py

"""Load IntrinsicWeather diffusers checkpoints from Hugging Face Hub for ZeroGPU."""

from __future__ import annotations

import gc
import importlib.util
import inspect
import os
import sys
from pathlib import Path
from typing import Any, get_args, get_origin

import numpy as np
import torch
from diffusers import DiffusionPipeline
import diffusers.pipelines.pipeline_utils as pipeline_utils
from huggingface_hub import snapshot_download
from PIL import Image
from transformers import AutoImageProcessor, AutoModel

from model_catalog import AOVS, DINO_MODEL_ID, MAP_LABELS, ModelProfile, get_profile


def _patch_diffusers_custom_pipeline_type_check() -> None:
    """Work around diffusers KeyError when custom pipelines omit parsed annotations."""

    if getattr(pipeline_utils, "_intrinsic_weather_type_check_patch", False):
        return

    @classmethod
    def patched_get_signature_types(cls):
        signature_types = {}
        for name, param in inspect.signature(cls.__init__).parameters.items():
            if name == "self":
                continue
            annotation = param.annotation
            if annotation is inspect.Parameter.empty:
                signature_types[name] = (inspect.Signature.empty,)
                continue
            origin = get_origin(annotation)
            if inspect.isclass(annotation):
                signature_types[name] = (annotation,)
            elif origin is not None:
                args = get_args(annotation)
                signature_types[name] = args if args else (annotation,)
            else:
                signature_types[name] = (inspect.Signature.empty,)
        return signature_types

    original_from_pretrained = DiffusionPipeline.from_pretrained.__func__

    @classmethod
    def from_pretrained_patched(cls, pretrained_model_name_or_path, *args, **kwargs):
        original_get_signature_types = DiffusionPipeline._get_signature_types
        DiffusionPipeline._get_signature_types = patched_get_signature_types
        try:
            return original_from_pretrained(cls, pretrained_model_name_or_path, *args, **kwargs)
        finally:
            DiffusionPipeline._get_signature_types = original_get_signature_types

    DiffusionPipeline.from_pretrained = from_pretrained_patched
    pipeline_utils._intrinsic_weather_type_check_patch = True


_patch_diffusers_custom_pipeline_type_check()


def _local_models_root() -> Path | None:
    raw = os.environ.get("LOCAL_MODELS_ROOT", "").strip()
    if not raw:
        return None
    path = Path(raw).expanduser()
    return path if path.is_dir() else None


LOCAL_MODELS_ROOT = _local_models_root()


def _is_model_checkpoint_dir(path: Path) -> bool:
    return (path / "pipeline_intrinsic_weather.py").is_file()


DTYPE = torch.bfloat16


def _to_uint8(image: np.ndarray | Image.Image) -> np.ndarray:
    if isinstance(image, Image.Image):
        return np.asarray(image.convert("RGB"))
    arr = np.asarray(image)
    if arr.dtype != np.uint8 and float(np.max(arr)) <= 1.0 + 1e-3:
        arr = arr * 255.0
    return np.clip(arr, 0, 255).astype(np.uint8)


def _import_intrinsic_weather_pipeline(model_dir: Path):
    pipeline_file = model_dir / "pipeline_intrinsic_weather.py"
    if not pipeline_file.is_file():
        raise FileNotFoundError(f"Missing unified pipeline module: {pipeline_file}")

    model_dir_str = str(model_dir.resolve())
    if model_dir_str not in sys.path:
        sys.path.insert(0, model_dir_str)

    module_name = "pipeline_intrinsic_weather"
    spec = importlib.util.spec_from_file_location(module_name, pipeline_file)
    if spec is None or spec.loader is None:
        raise ImportError(f"Cannot import IntrinsicWeather pipeline from {pipeline_file}")
    module = importlib.util.module_from_spec(spec)
    sys.modules[module_name] = module
    spec.loader.exec_module(module)
    return module.IntrinsicWeatherPipeline


def _load_dino(model_dir: Path) -> tuple[AutoModel, AutoImageProcessor]:
    dino_dir = model_dir / "dinov2"
    local_only = dino_dir.is_dir()
    source = dino_dir if local_only else DINO_MODEL_ID
    load_kwargs: dict[str, Any] = {}
    if local_only:
        load_kwargs["local_files_only"] = True

    processor = AutoImageProcessor.from_pretrained(source, **load_kwargs)
    model = AutoModel.from_pretrained(source, **load_kwargs)
    model.eval()
    return model, processor


class PipelineManager:
    def __init__(self) -> None:
        self._pipe: DiffusionPipeline | None = None
        self._dino_model: AutoModel | None = None
        self._dino_processor: AutoImageProcessor | None = None
        self._loaded_variant: str | None = None
        self._model_dir: Path | None = None
        self._on_cuda: bool = False

    @property
    def loaded_variant(self) -> str | None:
        return self._loaded_variant

    @property
    def pipe(self) -> DiffusionPipeline | None:
        return self._pipe

    def _resolve_model_source(self, profile: ModelProfile) -> tuple[str, str, bool]:
        if LOCAL_MODELS_ROOT is not None:
            for candidate in (LOCAL_MODELS_ROOT / profile.variant, LOCAL_MODELS_ROOT):
                if _is_model_checkpoint_dir(candidate):
                    return str(candidate.resolve()), str(candidate), True

        repo_id = profile.hub_repo
        cached_repo = snapshot_download(repo_id, repo_type="model")
        return str(Path(cached_repo).resolve()), repo_id, False

    def unload(self) -> None:
        for attr in ("_pipe", "_dino_model", "_dino_processor"):
            obj = getattr(self, attr)
            if obj is not None:
                del obj
                setattr(self, attr, None)
        self._loaded_variant = None
        self._model_dir = None
        self._on_cuda = False
        gc.collect()
        if torch.cuda.is_available():
            torch.cuda.empty_cache()

    def move_to_cuda(self) -> None:
        if self._pipe is not None and not self._on_cuda:
            for name in self._pipe.components.keys():
                module = getattr(self._pipe, name, None)
                if module is not None and hasattr(module, "to"):
                    if name == "imaa":
                        module.to(device="cuda")
                    else:
                        module.to(device="cuda", dtype=DTYPE)
            self._on_cuda = True
        if self._dino_model is not None:
            self._dino_model = self._dino_model.to("cuda")

    def load(self, variant: str) -> tuple[str, ModelProfile]:
        profile = get_profile(variant)
        if self._loaded_variant == variant and self._pipe is not None:
            return f"Model already loaded: `{variant}`", profile

        self.unload()
        model_source, source_label, from_local = self._resolve_model_source(profile)
        model_dir = Path(model_source)

        IntrinsicWeatherPipeline = _import_intrinsic_weather_pipeline(model_dir)
        load_kwargs: dict[str, Any] = {
            "inverse_transformer_subfolder": "inverse-512",
            "forward_transformer_subfolder": "forward",
            "torch_dtype": DTYPE,
            "load_lora": True,
            "load_imaa": True,
        }
        if from_local:
            load_kwargs["local_files_only"] = True

        pipe = IntrinsicWeatherPipeline.from_pretrained(model_dir.as_posix(), **load_kwargs)
        pipe.set_progress_bar_config(disable=True)
        dino_model, dino_processor = _load_dino(model_dir)

        self._pipe = pipe
        self._dino_model = dino_model
        self._dino_processor = dino_processor
        self._loaded_variant = variant
        self._model_dir = model_dir
        self._on_cuda = False

        return f"Loaded `{variant}` from {source_label}", profile

    @torch.no_grad()
    def run_weather_edit(
        self,
        profile: ModelProfile,
        image: Image.Image,
        *,
        weather: str,
        num_inverse_steps: int,
        num_forward_steps: int,
        guidance_scale: float,
        image_guidance_scale: float,
        image_size: int,
        render_size: int,
        seed: int,
        return_maps: bool,
    ) -> tuple[np.ndarray, list[tuple[np.ndarray, str]] | None]:
        if self._pipe is None or self._dino_model is None or self._dino_processor is None:
            raise RuntimeError("Pipeline is not loaded.")

        self.move_to_cuda()
        device = torch.device("cuda")
        generator = torch.Generator(device=device).manual_seed(int(seed))

        result = self._pipe(
            image=image.convert("RGB"),
            weather=weather,
            dino_model=self._dino_model,
            dino_processor=self._dino_processor,
            num_inverse_steps=int(num_inverse_steps),
            num_forward_steps=int(num_forward_steps),
            guidance_scale=float(guidance_scale),
            image_guidance_scale=float(image_guidance_scale),
            image_size=int(image_size),
            render_size=int(render_size),
            generator=generator,
            return_maps=return_maps,
            output_type="pil",
        )

        rendered = _to_uint8(result.images[0])
        maps_gallery = None
        if return_maps and result.maps is not None:
            maps_gallery = [
                (_to_uint8(result.maps[aov]), MAP_LABELS[aov]) for aov in AOVS if aov in result.maps
            ]
        return rendered, maps_gallery

    @torch.no_grad()
    def run_decompose(
        self,
        profile: ModelProfile,
        image: Image.Image,
        *,
        num_inverse_steps: int,
        image_size: int,
    ) -> list[tuple[np.ndarray, str]]:
        if self._pipe is None or self._dino_model is None or self._dino_processor is None:
            raise RuntimeError("Pipeline is not loaded.")

        self.move_to_cuda()
        maps = self._pipe.decompose(
            image=image.convert("RGB"),
            dino_model=self._dino_model,
            dino_processor=self._dino_processor,
            num_inference_steps=int(num_inverse_steps),
            image_size=int(image_size),
            output_type="np",
        )
        return [(_to_uint8(maps[aov]), MAP_LABELS[aov]) for aov in AOVS if aov in maps]


PIPELINE_MANAGER = PipelineManager()