Load the pretrained KumoAnomaly checkpoint

Co-authored-by: Aditi Gautam <adgautam@nvidia.com>

Author: aw471

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

@@ -0,0 +1,32 @@
+# 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.")
+            for old, new in (
+                ("layers.0.self_attn.in_proj_", "attn.qkv_lin."),
+                ("layers.0.self_attn.out_proj.", "attn.out_lin."),
+                ("layers.0.linear1.", "mlp.0."),
+                ("layers.0.linear2.", "mlp.3."),
+                ("layers.0.norm1.", "attn_residual_norm."),
+                ("layers.0.norm2.", "mlp_residual_norm."),
+            ):
+                key = key.replace(old, new)
+            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

@@ -19,6 +19,7 @@
 from sdm._warnings import warn_once
 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,9 @@ def __init__(
     ) -> None:
         super().__init__()
 
-        self.model = _KumoAnomaly(
-            device="meta" if pretrained else device,
-        )
+        # Build on the target device, since the checkpoint does not contain
+        # the non-persistent sinusoidal lookup of the step embedding:
+        self.model = _KumoAnomaly(device=device)
 
         if pretrained:
             self._load_from_pretrained(device=device)
@@ -81,7 +82,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