Spaces:
Running on Zero
Running on Zero
| from __future__ import annotations | |
| from typing import Any, Mapping | |
| import torch | |
| from torch import nn | |
| def model_state_from_checkpoint(ckpt: Any, *, use_ema: bool = False) -> Any: | |
| """Extract online or EMA model weights from supported checkpoint formats.""" | |
| if use_ema: | |
| if not isinstance(ckpt, Mapping): | |
| raise KeyError("EMA weights require a wrapped checkpoint.") | |
| ema_payload = ckpt.get("ema") | |
| if isinstance(ema_payload, Mapping): | |
| state = ema_payload.get("state_dict") | |
| if isinstance(state, Mapping): | |
| return state | |
| ema_state = ckpt.get("ema_model") | |
| if isinstance(ema_state, Mapping): | |
| return ema_state | |
| raise KeyError("Checkpoint does not contain EMA weights.") | |
| if ( | |
| isinstance(ckpt, Mapping) | |
| and "state_dict" in ckpt | |
| and not any(str(key).startswith("model.") for key in ckpt.keys()) | |
| ): | |
| return ckpt["state_dict"] | |
| if isinstance(ckpt, Mapping) and ckpt and all( | |
| str(key).startswith("model.") for key in ckpt.keys() | |
| ): | |
| return { | |
| str(key)[len("model.") :]: value | |
| for key, value in ckpt.items() | |
| if str(key).startswith("model.") | |
| } | |
| if isinstance(ckpt, Mapping) and any( | |
| key in ckpt for key in ("model", "net", "weights") | |
| ): | |
| return ckpt.get("model", ckpt.get("net", ckpt.get("weights", {}))) | |
| return ckpt | |
| def load_checkpoint(path: str, *, map_location: torch.device | str) -> Any: | |
| """Load a trusted GEM checkpoint across supported PyTorch versions.""" | |
| try: | |
| return torch.load(path, map_location=map_location, weights_only=False) | |
| except TypeError: | |
| return torch.load(path, map_location=map_location) | |
| def load_model_checkpoint( | |
| model: nn.Module, | |
| path: str, | |
| *, | |
| map_location: torch.device | str, | |
| use_ema: bool = False, | |
| ) -> Any: | |
| """Load a checkpoint and require an exact match with ``model``.""" | |
| checkpoint = load_checkpoint(path, map_location=map_location) | |
| state = model_state_from_checkpoint(checkpoint, use_ema=use_ema) | |
| if not isinstance(state, Mapping) or not state: | |
| raise TypeError(f"Checkpoint '{path}' does not contain a model state dictionary.") | |
| model.load_state_dict(state, strict=True) | |
| return checkpoint | |