Update bridge/latent_bridge.py

bridge/latent_bridge.py

@@ -9,9 +9,9 @@ from diffusers import ModelMixin
 class LatentBridge(ModelMixin):
     def __init__(self, in_ch: int = 4, out_ch: int = 16):
         super().__init__()
-        # dec: 4→16, enc: 16→4
-        self.dec = nn.Conv2d(in_ch, out_ch, kernel_size=1)
+        # enc: 16→4 channels, dec: 4→16 channels
         self.enc = nn.Conv2d(out_ch, in_ch, kernel_size=1)
+        self.dec = nn.Conv2d(in_ch, out_ch, kernel_size=1)
 
     def forward(self, x):
         return x
@@ -20,31 +20,40 @@ class LatentBridge(ModelMixin):
     def load_config(cls, pretrained_model_name_or_path: str, **kwargs):
         sub = kwargs.get("subfolder", "")
         base = os.path.join(pretrained_model_name_or_path, sub) if sub else pretrained_model_name_or_path
-        with open(os.path.join(base, "config.json"), "r") as f:
+        cfg_file = os.path.join(base, "config.json")
+        with open(cfg_file, "r") as f:
             cfg = json.load(f)
         cfg.pop("_class_name", None)
         return cls(**cfg)
 
     @classmethod
     def from_pretrained(cls, pretrained_model_name_or_path: str, **kwargs):
-        # 1) Build instance from config
+        # 1) Build the instance from config
         bridge = cls.load_config(pretrained_model_name_or_path, **kwargs)
 
-        # 2) Load raw weights
+        # 2) Load raw state dict
         sub = kwargs.get("subfolder", "")
         base = os.path.join(pretrained_model_name_or_path, sub) if sub else pretrained_model_name_or_path
         raw_sd = torch.load(os.path.join(base, "pytorch_model.bin"), map_location="cpu")
 
-        # 3) Remap keys
-        remapped = {}
+        # 3) Remap keys explicitly
+        remapped_sd = {
+            # map to16.* → dec.*
+            k.replace("to16", "dec", 1): v if k.startswith("to16") else
+            # map to4.*  → enc.*
+            v
+            for k, v in raw_sd.items()
+        }
+        # But dictionary comprehension above only covers to16; let's do full loop:
+        final_sd = {}
         for k, v in raw_sd.items():
-            if k.startswith("to4."):
-                remapped[k.replace("to4", "enc", 1)] = v
-            elif k.startswith("to16."):
-                remapped[k.replace("to16", "dec", 1)] = v
+            if k.startswith("to16."):
+                final_sd[k.replace("to16", "dec", 1)] = v
+            elif k.startswith("to4."):
+                final_sd[k.replace("to4", "enc", 1)] = v
             else:
-                remapped[k] = v
+                final_sd[k] = v
 
-        # 4) **Load** the remapped dict into your layers
-        bridge.load_state_dict(remapped)
+        # 4) Load into your layers
+        bridge.load_state_dict(final_sd)
         return bridge