Fix: Prevent crash when trying to load SigLIP/CLIP encoders with transformers >= 5.6.0

transformer_lens/model_bridge/generalized_components/clip_vision_encoder.py

@@ -175,3 +175,17 @@ def forward(
             output = self.hook_out(output)
 
         return output
+
+    def set_original_component(self, original_component: torch.nn.Module) -> None:
+        """Set the original component that this bridge wraps.
+        Note that CLIPVisionModel used to wrap a inner object as .vision_model before
+        transformers version 5.6.0, but after that it must directly be used.
+        This is a temporary hack to fix that till the transformers version is bumped.
+
+        Args:
+            original_component: The original transformer component to wrap
+        """
+        if not hasattr(original_component, "vision_model"):
+            # We should bypass any pytorch module registration.
+            object.__setattr__(original_component, "vision_model", original_component)
+        super().set_original_component(original_component)

transformer_lens/model_bridge/generalized_components/siglip_vision_encoder.py

@@ -115,8 +115,8 @@ def __init__(
         """
         # All submodule names are resolved relative to the parent's
         # original_component (a SiglipVisionModel) by setup_submodules().
-        # SiglipVisionModel wraps SiglipVisionTransformer as .vision_model,
-        # so all paths go through vision_model.*.
+        # SiglipVisionModel wraps a SiglipVisionTransformer as .vision_model till
+        # transformers version 5.6.0
         # post_layernorm is nn.LayerNorm; NormalizationBridge introspects the
         # wrapped module so the RMSNorm-LM config (Gemma 3, LLaVA) doesn't leak.
         default_submodules = {
@@ -171,3 +171,17 @@ def forward(
             output = self.hook_out(output)
 
         return output
+
+    def set_original_component(self, original_component: torch.nn.Module) -> None:
+        """Set the original component that this bridge wraps.
+        Note that SiglipVisionModel used to wrap a inner object as .vision_model before
+        transformers version 5.6.0, but after that it must directly be used.
+        This is a temporary hack to fix that till the transformers version is bumped.
+
+        Args:
+            original_component: The original transformer component to wrap
+        """
+        if not hasattr(original_component, "vision_model"):
+            # We should bypass any pytorch module registration.
+            object.__setattr__(original_component, "vision_model", original_component)
+        super().set_original_component(original_component)