Create modeling_flux2.py

transformer/modeling_flux2.py

@@ -0,0 +1,57 @@
+import torch
+import torch.nn as nn
+from diffusers import ModelMixin, ConfigMixin
+
+
+class Flux2Transformer2DModel(ModelMixin, ConfigMixin):
+    config_name = "config.json"
+
+    def __init__(self, **kwargs):
+        super().__init__()
+
+        # Store config
+        self.register_to_config(**kwargs)
+
+        # Internal storage (safe for arbitrary keys)
+        self._sd = {}
+
+    # -----------------------------
+    # LOAD: accept ANY state dict
+    # -----------------------------
+    def load_state_dict(self, state_dict, strict=False):
+        """
+        Store raw tensors exactly as-is.
+        No validation, no structure assumptions.
+        """
+        self._sd = {}
+
+        for k, v in state_dict.items():
+            if isinstance(v, torch.Tensor):
+                self._sd[k] = v.contiguous()
+            else:
+                self._sd[k] = v
+
+        # Pretend everything matched
+        return torch.nn.modules.module._IncompatibleKeys([], [])
+
+    # -----------------------------
+    # SAVE: return original weights
+    # -----------------------------
+    def state_dict(self, *args, **kwargs):
+        return dict(self._sd)
+
+    # -----------------------------
+    # Required by diffusers internals
+    # -----------------------------
+    def _convert_deprecated_attention_blocks(self, state_dict):
+        # Diffusers sometimes calls this
+        return state_dict
+
+    # -----------------------------
+    # Forward (dummy)
+    # -----------------------------
+    def forward(self, *args, **kwargs):
+        raise RuntimeError(
+            "Flux2Transformer2DModel is a stub loader. "
+            "It cannot run inference."
+        )