Update landmarkdiff/model_registry.py to v0.3.2

landmarkdiff/model_registry.py

@@ -139,7 +139,9 @@ class ModelRegistry:
             step = int(parts[-1])
 
         # Compute size
-        size_mb = sum(f.stat().st_size for f in ckpt_dir.rglob("*") if f.is_file()) / (1024 * 1024)
+        size_mb = sum(
+            f.stat().st_size for f in ckpt_dir.rglob("*") if f.is_file()
+        ) / (1024 * 1024)
 
         return ModelEntry(
             name=ckpt_dir.name,
@@ -193,15 +195,16 @@ class ModelRegistry:
         Returns:
             Best ModelEntry, or None if no models have the metric.
         """
-        candidates = [m for m in self._models.values() if metric in m.metrics]
+        candidates = [
+            m for m in self._models.values()
+            if metric in m.metrics
+        ]
         if not candidates:
             return None
 
-        return (
-            min(candidates, key=lambda m: m.metrics[metric])
-            if lower_is_better
-            else max(candidates, key=lambda m: m.metrics[metric])
-        )
+        return min(candidates, key=lambda m: m.metrics[metric])  \
+            if lower_is_better else \
+            max(candidates, key=lambda m: m.metrics[metric])
 
     def get_by_step(self, step: int) -> ModelEntry | None:
         """Get a model by its training step."""
@@ -246,30 +249,40 @@ class ModelRegistry:
         self,
         name: str,
         use_ema: bool = True,
+        torch_dtype: torch.dtype | None = None,
     ) -> Any:
         """Load a ControlNet model from checkpoint.
 
         Args:
             name: Checkpoint name.
             use_ema: If True, load EMA weights (preferred for inference).
+            torch_dtype: Weight dtype (e.g. torch.float16). Defaults to
+                float16 on CUDA, float32 on CPU.
 
         Returns:
             ControlNetModel instance.
         """
         from diffusers import ControlNetModel
 
+        if torch_dtype is None:
+            torch_dtype = torch.float16 if torch.cuda.is_available() else torch.float32
+
         entry = self._models.get(name)
         if entry is None:
             raise KeyError(f"Checkpoint '{name}' not found in registry")
 
         if use_ema and entry.has_ema:
-            return ControlNetModel.from_pretrained(str(entry.path / "controlnet_ema"))
+            return ControlNetModel.from_pretrained(
+                str(entry.path / "controlnet_ema"),
+                torch_dtype=torch_dtype,
+            )
 
         # Fallback: load from training state
         state = self.load(name)
         model = ControlNetModel.from_pretrained(
-            "lllyasviel/control_v11p_sd15_openpose",
+            "CrucibleAI/ControlNetMediaPipeFace",
             subfolder="diffusion_sd15",
+            torch_dtype=torch_dtype,
         )
         key = "ema_controlnet" if use_ema else "controlnet"
         model.load_state_dict(state[key])
@@ -351,7 +364,9 @@ class ModelRegistry:
         for metric in sorted(all_metrics):
             values = [m.metrics[metric] for m in models if metric in m.metrics]
             if values:
-                lines.append(f"  {metric}: {min(values):.4f} — {max(values):.4f}")
+                lines.append(
+                    f"  {metric}: {min(values):.4f} — {max(values):.4f}"
+                )
 
         return "\n".join(lines)