"""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()