Fix zoedepth (#2010)

nod-ai · Dec 5, 2023 · db0c53a · db0c53a
1 parent ce9ce3a
commit db0c53a
Showing 1 changed file with 9 additions and 3 deletions.
diff --git a/apps/stable_diffusion/src/utils/stencils/zoe/__init__.py b/apps/stable_diffusion/src/utils/stencils/zoe/__init__.py
@@ -30,9 +30,15 @@ def __init__(self):
             pretrained=False,
             force_reload=False,
         )
-        model.load_state_dict(
-            torch.load(modelpath, map_location=model.device)["model"]
-        )
+
+        # Hack to fix the ZoeDepth import issue
+        model_keys = model.state_dict().keys()
+        loaded_dict = torch.load(modelpath, map_location=model.device)["model"]
+        loaded_keys = loaded_dict.keys()
+        for key in loaded_keys - model_keys:
+            loaded_dict.pop(key)
+
+        model.load_state_dict(loaded_dict)
         model.eval()
         self.model = model