garmentiq.segmentation.load_model

  1"""Loading a segmentation model onto a device."""
  2import inspect
  3import torch
  4import torch.nn as nn
  5from typing import Type, Union
  6from safetensors.torch import load_file
  7from garmentiq.utils.device import resolve_device
  8from garmentiq.utils.checkpoint import load_state_dict_checked
  9
 10
 11def load_model(
 12    model_class: Type[nn.Module],
 13    model_path: str,
 14    model_args: dict = None,
 15    device: Union[str, torch.device] = "cpu",
 16    **kwargs,
 17):
 18    """
 19    Loads a PyTorch model from a local checkpoint and prepares it for inference.
 20
 21    This function instantiates the provided model class using safely filtered configuration
 22    arguments, loads the weights from a local `.pth` or `.safetensors` file, moves the model
 23    to the requested device, and sets it to evaluation mode. It automatically
 24    strips common weight prefixes (e.g., "module.", "model.") to ensure compatibility.
 25
 26    Args:
 27        model_class (Type[nn.Module]): The uninstantiated PyTorch model class to be used.
 28        model_path (str): The local file path to the model checkpoint weights, typically
 29                          ending in `.pth` or `.safetensors`.
 30        model_args (dict, optional): A dictionary of configuration arguments for initializing
 31                                     the model. Incompatible arguments are safely ignored.
 32                                     Default is None.
 33        device (Union[str, torch.device], optional): The device to load the model onto, e.g.
 34                                                     `"cpu"`, `"cuda"`, `"cuda:0"`, or `"mps"`.
 35                                                     Hardware acceleration is opt-in; pass it
 36                                                     explicitly to use a GPU or Apple Silicon.
 37                                                     Default is `"cpu"`.
 38        **kwargs: Additional arbitrary keyword arguments.
 39
 40    Raises:
 41        ValueError: If the requested `device` is invalid or unavailable on this machine.
 42        Exception: If the model weights cannot be loaded from the specified local path or if
 43                   the file format is unsupported.
 44
 45    Returns:
 46        nn.Module: The loaded and prepared PyTorch model instance, placed on `device`.
 47    """
 48    model_args = model_args or {}
 49    
 50    # --- THE SMART CONVERTER ---
 51    # If the user passes a Config object instead of a dictionary, safely convert it.
 52    if not isinstance(model_args, dict):
 53        if hasattr(model_args, "to_dict"):
 54            model_args = model_args.to_dict()  # Hugging Face standard
 55        elif hasattr(model_args, "__dict__"):
 56            model_args = vars(model_args)      # Standard Python objects
 57        else:
 58            raise TypeError("model_args must be a dictionary or a configuration object.")
 59    # ---------------------------
 60
 61    device = resolve_device(device)
 62
 63    sig = inspect.signature(model_class.__init__)
 64    valid_params = set(sig.parameters.keys())
 65
 66    has_kwargs = any(
 67        p.kind == inspect.Parameter.VAR_KEYWORD for p in sig.parameters.values()
 68    )
 69
 70    if not has_kwargs:
 71        filtered_args = {k: v for k, v in model_args.items() if k in valid_params}
 72    else:
 73        filtered_args = model_args
 74
 75    # ... [Keep the rest of your loading logic exactly the same] ...
 76    model = model_class(**filtered_args).to(device)
 77
 78    if model_path.endswith(".safetensors"):
 79        state_dict = load_file(model_path, device=str(device))
 80    else:
 81        state_dict = torch.load(model_path, map_location=device, weights_only=True)
 82
 83    # SAM 3 is published as a video checkpoint that nests the image ("detector") model
 84    # alongside tracker weights. Keep only the detector half when that layout is seen,
 85    # otherwise the image model would silently load nothing.
 86    from garmentiq.segmentation.model_definition.sam.sam import SAM3_DETECTOR_PREFIX
 87
 88    if any(k.startswith(SAM3_DETECTOR_PREFIX) for k in state_dict):
 89        state_dict = {
 90            k[len(SAM3_DETECTOR_PREFIX) :]: v
 91            for k, v in state_dict.items()
 92            if k.startswith(SAM3_DETECTOR_PREFIX)
 93        }
 94
 95    new_state_dict = {
 96        k.removeprefix("module.").removeprefix("model."): v
 97        for k, v in state_dict.items()
 98    }
 99
100    load_state_dict_checked(model, new_state_dict, model_path)
101    model.eval()
102
103    return model