# CSD (HuggingFace format) — unofficial port. Copyright (c) 2026 bigshanedogg. MIT License.
#
# Self-contained transformers port of the CSD style model from
#   "Measuring Style Similarity in Diffusion Models" (Somepalli et al., 2024)
#   https://github.com/learn2phoenix/CSD  (code: MIT)
# so it loads via AutoModel.from_pretrained(trust_remote_code=True) without the `clip`
# package. The ViT-L/14 vision transformer below is vendored from OpenAI CLIP
#   https://github.com/openai/CLIP  (MIT, (c) 2021 OpenAI) — MODIFIED: trimmed to the
# vision tower, projection removed (folded into the CSD style/content heads).
# The released CSD checkpoint (tomg-group-umd/CSD-ViT-L) is CC-BY-4.0.

from collections import OrderedDict
from typing import Optional

import torch
import torch.nn as nn
from transformers import PretrainedConfig, PreTrainedModel
from transformers.modeling_outputs import ModelOutput


class CSDConfig(PretrainedConfig):
    model_type = "csd"

    def __init__(
        self,
        image_resolution: int = 224,
        patch_size: int = 14,
        width: int = 1024,
        layers: int = 24,
        heads: int = 16,
        embed_dim: int = 768,
        **kwargs,
    ):
        self.image_resolution = image_resolution
        self.patch_size = patch_size
        self.width = width
        self.layers = layers
        self.heads = heads
        self.embed_dim = embed_dim  # style/content projection output dim
        super().__init__(**kwargs)


# ── vendored OpenAI CLIP vision tower (MIT, (c) 2021 OpenAI; MODIFIED) ──────────────
class QuickGELU(nn.Module):
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return x * torch.sigmoid(1.702 * x)


class ResidualAttentionBlock(nn.Module):
    def __init__(self, d_model: int, n_head: int):
        super().__init__()
        self.attn = nn.MultiheadAttention(d_model, n_head)
        self.ln_1 = nn.LayerNorm(d_model)
        self.mlp = nn.Sequential(
            OrderedDict(
                [
                    ("c_fc", nn.Linear(d_model, d_model * 4)),
                    ("gelu", QuickGELU()),
                    ("c_proj", nn.Linear(d_model * 4, d_model)),
                ]
            )
        )
        self.ln_2 = nn.LayerNorm(d_model)

    def attention(self, x: torch.Tensor) -> torch.Tensor:
        return self.attn(x, x, x, need_weights=False, attn_mask=None)[0]

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = x + self.attention(self.ln_1(x))
        x = x + self.mlp(self.ln_2(x))
        return x


class Transformer(nn.Module):
    def __init__(self, width: int, layers: int, heads: int):
        super().__init__()
        self.resblocks = nn.Sequential(*[ResidualAttentionBlock(width, heads) for _ in range(layers)])

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.resblocks(x)


class VisionTransformer(nn.Module):
    def __init__(self, input_resolution: int, patch_size: int, width: int, layers: int, heads: int):
        super().__init__()
        self.conv1 = nn.Conv2d(3, width, kernel_size=patch_size, stride=patch_size, bias=False)
        _scale = width**-0.5
        self.class_embedding = nn.Parameter(_scale * torch.randn(width))
        _num_positions = (input_resolution // patch_size) ** 2 + 1
        self.positional_embedding = nn.Parameter(_scale * torch.randn(_num_positions, width))
        self.ln_pre = nn.LayerNorm(width)
        self.transformer = Transformer(width, layers, heads)
        self.ln_post = nn.LayerNorm(width)
        # NOTE: CSD sets backbone.proj = None and folds projection into last_layer_{style,content}.

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = self.conv1(x)  # (B, width, grid, grid)
        x = x.reshape(x.shape[0], x.shape[1], -1).permute(0, 2, 1)  # (B, grid**2, width)
        _cls = self.class_embedding.to(x.dtype) + torch.zeros(x.shape[0], 1, x.shape[-1], dtype=x.dtype, device=x.device)
        x = torch.cat([_cls, x], dim=1)  # (B, grid**2 + 1, width)
        x = x + self.positional_embedding.to(x.dtype)
        x = self.ln_pre(x)
        x = x.permute(1, 0, 2)  # NLD -> LND
        x = self.transformer(x)
        x = x.permute(1, 0, 2)  # LND -> NLD
        x = self.ln_post(x[:, 0, :])  # take the [CLS] token
        return x


class CSDOutput(ModelOutput):
    embeddings: Optional[torch.FloatTensor] = None  # style embedding (L2-normalized)
    content_embeddings: Optional[torch.FloatTensor] = None
    last_hidden_states: Optional[torch.FloatTensor] = None  # pre-projection ViT feature


class CSDModel(PreTrainedModel):
    """CSD style/content encoder. ``embeddings`` is the L2-normalized style descriptor
    (``feature @ last_layer_style``); the perceptual/style scoring lives in the caller."""

    config_class = CSDConfig

    def __init__(self, config: CSDConfig):
        super().__init__(config)
        self.backbone = VisionTransformer(
            input_resolution=config.image_resolution,
            patch_size=config.patch_size,
            width=config.width,
            layers=config.layers,
            heads=config.heads,
        )
        self.last_layer_style = nn.Parameter(torch.empty(config.width, config.embed_dim))
        self.last_layer_content = nn.Parameter(torch.empty(config.width, config.embed_dim))
        # transformers>=5 sets weight-loading state (e.g. all_tied_weights_keys, read by
        # from_pretrained) in PreTrainedModel.post_init(); call it so the load doesn't
        # AttributeError. from_pretrained overwrites the freshly-inited weights afterward.
        self.post_init()

    def forward(self, pixel_values: torch.Tensor) -> CSDOutput:
        _feature = self.backbone(pixel_values)
        _style = nn.functional.normalize(_feature @ self.last_layer_style, dim=1, p=2)
        _content = nn.functional.normalize(_feature @ self.last_layer_content, dim=1, p=2)
        return CSDOutput(embeddings=_style, content_embeddings=_content, last_hidden_states=_feature)