Update bridge/latent_bridge.py

bridge/latent_bridge.py

@@ -18,43 +18,33 @@ class LatentBridge(ModelMixin):
 
     @classmethod
     def load_config(cls, pretrained_model_name_or_path: str, **kwargs):
-        # Determine the folder (honor subfolder kw if provided)
         sub = kwargs.get("subfolder", "")
         base = os.path.join(pretrained_model_name_or_path, sub) if sub else pretrained_model_name_or_path
-
-        # Load and parse config.json
-        config_file = os.path.join(base, "config.json")
-        with open(config_file, "r") as f:
+        with open(os.path.join(base, "config.json"), "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):
-        # Resolve base folder for both config & weights
-        sub = kwargs.get("subfolder", "")
-        base = os.path.join(pretrained_model_name_or_path, sub) if sub else pretrained_model_name_or_path
-
-        # 1) Instantiate from config
+        # 1) Build instance from config
         bridge = cls.load_config(pretrained_model_name_or_path, **kwargs)
 
-        # 2) Load original state dict
-        bin_path = os.path.join(base, "pytorch_model.bin")
-        state_dict = torch.load(bin_path, map_location="cpu")
+        # 2) Load raw weights
+        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: 
-        #    to4.* → enc.*  (16→4 conv)
-        #    to16.* → dec.* (4→16 conv)
+        # 3) Remap keys
         remapped = {}
-        for k, v in state_dict.items():
+        for k, v in raw_sd.items():
             if k.startswith("to4."):
-                new_key = k.replace("to4", "enc", 1)
+                remapped[k.replace("to4", "enc", 1)] = v
             elif k.startswith("to16."):
-                new_key = k.replace("to16", "dec", 1)
+                remapped[k.replace("to16", "dec", 1)] = v
             else:
-                new_key = k
-            remapped[new_key] = v
+                remapped[k] = v
 
-        # 4) Load into your bridge
+        # 4) **Load** the remapped dict into your layers
         bridge.load_state_dict(remapped)
         return bridge