[Kumo-Anomaly 4/n] Load the pretrained KumoAnomaly checkpoint

sdm/models/kumo/timeseries/anomaly/ckpt.py

@@ -0,0 +1,24 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+from torch import Tensor
+
+
+def remap_ckpt(ckpt: dict[str, Tensor]) -> dict[str, Tensor]:  # noqa: D103
+    out = {}
+    for key, value in ckpt.items():
+        if key.startswith("diffmodel."):
+            key = "denoiser." + key.removeprefix("diffmodel.")
+            key = key.replace(".layers.0.", ".")
+            if key in (
+                "denoiser.input_projection.weight",
+                "denoiser.output_projection1.weight",
+                "denoiser.output_projection2.weight",
+            ):
+                value = value.squeeze(-1)  # `Conv1d` to `Linear`.
+        elif key.startswith(
+            ("embed_layer.", "embed_norm.", "time_embed_norm.")
+        ):
+            key = "side_embedding." + key
+        out[key] = value
+    return out

sdm/models/kumo/timeseries/anomaly/model.py

@@ -18,6 +18,7 @@
 from sdm._inference import inference_mode
 from sdm.models import Callback, Model
 from sdm.models._huggingface import download_checkpoint
+from sdm.models.kumo.timeseries.anomaly.ckpt import remap_ckpt
 from sdm.models.kumo.timeseries.anomaly.denoiser import DiffusionDenoiser
 from sdm.models.kumo.timeseries.anomaly.embedding import SideEmbedding
 from sdm.models.kumo.timeseries.anomaly.recipe import default_recipe
@@ -61,9 +62,7 @@ def __init__(
     ) -> None:
         super().__init__()
 
-        self.model = _KumoAnomaly(
-            device="meta" if pretrained else device,
-        )
+        self.model = _KumoAnomaly(device=device)
 
         if pretrained:
             self._load_from_pretrained(device=device)
@@ -81,7 +80,8 @@ def _load_from_pretrained(
             filename="final_model.pth",
             revision="226c5003b0582adc1cac433123f47324bfae39f2",
         )
-        torch.load(path, map_location=device, weights_only=True)  # TODO
+        ckpt = torch.load(path, map_location=device, weights_only=True)
+        self.model.load_state_dict(remap_ckpt(ckpt["model"]))
 
         return self