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."
        )