Add loading ckpt logic from state_dict and not directly from model class

model_blocks/controlnet.py

@@ -22,7 +22,7 @@ class ControlNet(nn.Module):
     """
 
     def __init__(
-        self, device, model_config, trained_ckpt_path=None, model_locked=True
+        self, device, model_config, model_ckpt=None, model_locked=True
     ) -> None:
         super().__init__()
 
@@ -30,14 +30,23 @@ class ControlNet(nn.Module):
         self.model = UNet(model_config)
         self.model_locked = model_locked
 
-        if trained_ckpt_path is not None:
-            print("Loading Checkpoint")
-            self.model = torch.load(trained_ckpt_path).to(device)
+        # Load weights for the trained model
+        if model_ckpt is not None and device is not None:
+            print("Loading Trained Diffusion Model")
+            self.model.load_state_dict(
+                torch.load(model_ckpt, map_location=device), strict=True
+            )
 
-        # False the upblocks (Decoder blocks) from the DDPM and uses only the encoder
+        # ControlNet Copy of Trained DDPM
+        # use_up = False removes the upblocks(decoder layers) from DDPM Unet
         self.control_copy = UNet(model_config, use_up=False)
-        if trained_ckpt_path is not None:
-            self.control_copy.load_state_dict(self.model.state_dict(), strict=False)
+        # Load same weights as the trained model
+
+        if model_ckpt is not None and device is not None:
+            print("Loading Control Diffusion Model")
+            self.control_copy.load_state_dict(
+                torch.load(model_ckpt, map_location=device), strict=False
+            )
 
         # Hint Block for ControlNet
         # Stack of Conv Activation and Zero Convolution at the end

pyproject.toml

@@ -9,5 +9,6 @@ dependencies = [
     "pandas>=2.2.3",
     "torch>=2.6.0",
     "torchvision>=0.21.0",
+    "tqdm>=4.67.1",
     "wandb>=0.19.9",
 ]

requirements.txt

@@ -45,6 +45,7 @@ smmap==5.0.2
 sympy==1.13.1
 torch==2.6.0
 torchvision==0.21.0
+tqdm==4.67.1
 triton==3.2.0
 typing-extensions==4.13.0
 typing-inspection==0.4.0

uv.lock

@@ -76,6 +76,7 @@ dependencies = [
     { name = "pandas" },
     { name = "torch" },
     { name = "torchvision" },
+    { name = "tqdm" },
     { name = "wandb" },
 ]
 
@@ -85,6 +86,7 @@ requires-dist = [
     { name = "pandas", specifier = ">=2.2.3" },
     { name = "torch", specifier = ">=2.6.0" },
     { name = "torchvision", specifier = ">=0.21.0" },
+    { name = "tqdm", specifier = ">=4.67.1" },
     { name = "wandb", specifier = ">=0.19.9" },
 ]
 
@@ -666,6 +668,18 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/ed/b4/fc60e3bc003879d3de842baea258fffc3586f4b49cd435a5ba1e09c33315/torchvision-0.21.0-cp313-cp313-win_amd64.whl", hash = "sha256:9147f5e096a9270684e3befdee350f3cacafd48e0c54ab195f45790a9c146d67", size = 1560519 },
 ]
 
+[[package]]
+name = "tqdm"
+version = "4.67.1"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+    { name = "colorama", marker = "sys_platform == 'win32'" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/a8/4b/29b4ef32e036bb34e4ab51796dd745cdba7ed47ad142a9f4a1eb8e0c744d/tqdm-4.67.1.tar.gz", hash = "sha256:f8aef9c52c08c13a65f30ea34f4e5aac3fd1a34959879d7e59e63027286627f2", size = 169737 }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/d0/30/dc54f88dd4a2b5dc8a0279bdd7270e735851848b762aeb1c1184ed1f6b14/tqdm-4.67.1-py3-none-any.whl", hash = "sha256:26445eca388f82e72884e0d580d5464cd801a3ea01e63e5601bdff9ba6a48de2", size = 78540 },
+]
+
 [[package]]
 name = "triton"
 version = "3.2.0"