garmentiq.matting.load_model

 1"""Loading a matting 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 matting model from a local checkpoint and prepares it for inference.
20
21    This mirrors `garmentiq.segmentation.load_model`: the model class is instantiated with
22    safely filtered configuration arguments, weights are read from a local `.pth` or
23    `.safetensors` file, the model is moved to the requested device and set to evaluation
24    mode, and common weight prefixes are stripped for compatibility.
25
26    Matting Anything checkpoints are not loaded here because they pair a decoder with a
27    separate SAM model; use `garmentiq.matting.model_definition.mam.load_mam` for those.
28
29    Args:
30        model_class (Type[nn.Module]): The uninstantiated model class, typically
31                                       `VitMatteForImageMatting`.
32        model_path (str): Local path to the checkpoint weights, ending in `.pth` or
33                          `.safetensors`.
34        model_args (dict, optional): Configuration arguments for initializing the model,
35                                     e.g. `{"config": load_vitmatte_config(...)}`.
36                                     Incompatible arguments are safely ignored.
37                                     Default is None.
38        device (Union[str, torch.device], optional): The device to load the model onto, e.g.
39                                                     `"cpu"`, `"cuda"`, or `"mps"`. Hardware
40                                                     acceleration is opt-in. Default is `"cpu"`.
41        **kwargs: Additional arbitrary keyword arguments.
42
43    Raises:
44        ValueError: If the requested `device` is invalid or unavailable on this machine.
45        Exception: If the weights cannot be loaded or the file format is unsupported.
46
47    Returns:
48        nn.Module: The loaded model in evaluation mode, placed on `device`.
49    """
50    model_args = model_args or {}
51
52    if not isinstance(model_args, dict):
53        if hasattr(model_args, "to_dict"):
54            model_args = model_args.to_dict()
55        elif hasattr(model_args, "__dict__"):
56            model_args = vars(model_args)
57        else:
58            raise TypeError(
59                "model_args must be a dictionary or a configuration object."
60            )
61
62    device = resolve_device(device)
63
64    sig = inspect.signature(model_class.__init__)
65    valid_params = set(sig.parameters.keys())
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    model = model_class(**filtered_args).to(device)
76
77    if model_path.endswith(".safetensors"):
78        state_dict = load_file(model_path, device=str(device))
79    else:
80        state_dict = torch.load(model_path, map_location=device, weights_only=True)
81
82    new_state_dict = {
83        k.removeprefix("module.").removeprefix("model."): v
84        for k, v in state_dict.items()
85    }
86
87    load_state_dict_checked(model, new_state_dict, model_path)
88    model.eval()
89
90    return model