modeling_dreamsim.py

# DreamSim (HuggingFace format) — unofficial port.
# Copyright (c) 2026 bigshanedogg. Released under the MIT License (see LICENSE).
#
# Derivative of DreamSim (MIT, (c) 2023 Shobhita Sundaram, Netanel Tamir,
# Stephanie Fu, Richard Zhang — https://github.com/ssundaram21/dreamsim).
# The ViT backbone below is vendored from DINO (Apache-2.0, (c) Meta Platforms).
# Not an official DreamSim release.

"""Self-contained HF modeling for the DreamSim perceptual-similarity ensemble.

Vendors the DreamSim architecture (github ssundaram21/dreamsim, MIT) so the model
loads from ``model.safetensors`` via ``from_pretrained`` with no ``dreamsim`` /
``peft`` dependency. The ViT backbone below is copy-pasted from DINO
(github facebookresearch/dino, Apache-2.0), which DreamSim itself vendors.

Ensemble (all ViT-B/16, patch stride 16), features concatenated → mean/L2-normalize:
  * dino_vitb16      (feat "cls")       : pre-final-norm CLS token          → 768
  * clip_vitb16      (feat "embedding") : post-norm CLS @ proj (QuickGELU)  → 512
  * open_clip_vitb16 (feat "embedding") : post-norm CLS @ proj (GELU)       → 512
Each backbone normalizes the [0,1] input with its own mean/std (DINO→ImageNet,
CLIP/OpenCLIP→OpenAI-CLIP). LoRA is pre-merged into the weights.
"""

import math
from dataclasses import dataclass
from functools import partial
from typing import Optional

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

from .configuration_dreamsim import DreamSimConfig

IMAGENET_MEAN = (0.485, 0.456, 0.406)
IMAGENET_STD = (0.229, 0.224, 0.225)
OPENAI_CLIP_MEAN = (0.48145466, 0.4578275, 0.40821073)
OPENAI_CLIP_STD = (0.26862954, 0.26130258, 0.27577711)


# ── ViT backbone — adapted (WITH MODIFICATIONS) from DINO ─────────────────────
# https://github.com/facebookresearch/dino  •  Apache-2.0  •  (c) Facebook, Inc.
# Modified: vendored into this module, trimmed to the inference path, and restructured
# for `transformers`. See NOTICE and LICENSE.apache-2.0.txt.


class QuickGELU(nn.Module):
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return x * torch.sigmoid(1.702 * x)


class Mlp(nn.Module):
    def __init__(self, in_features, hidden_features=None, out_features=None, act_layer=nn.GELU, drop=0.0):
        super().__init__()
        out_features = out_features or in_features
        hidden_features = hidden_features or in_features
        self.fc1 = nn.Linear(in_features, hidden_features)
        self.act = act_layer()
        self.fc2 = nn.Linear(hidden_features, out_features)
        self.drop = nn.Dropout(drop)

    def forward(self, x):
        return self.drop(self.fc2(self.drop(self.act(self.fc1(x)))))


class Attention(nn.Module):
    def __init__(self, dim, num_heads=8, qkv_bias=False, qk_scale=None, attn_drop=0.0, proj_drop=0.0):
        super().__init__()
        self.num_heads = num_heads
        head_dim = dim // num_heads
        self.scale = qk_scale or head_dim**-0.5
        self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
        self.attn_drop = nn.Dropout(attn_drop)
        self.proj = nn.Linear(dim, dim)
        self.proj_drop = nn.Dropout(proj_drop)

    def forward(self, x):
        B, N, C = x.shape
        qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)
        q, k, v = qkv[0], qkv[1], qkv[2]
        attn = (q @ k.transpose(-2, -1)) * self.scale
        attn = attn.softmax(dim=-1)
        attn = self.attn_drop(attn)
        x = (attn @ v).transpose(1, 2).reshape(B, N, C)
        x = self.proj_drop(self.proj(x))
        return x, attn


class Block(nn.Module):
    def __init__(self, dim, num_heads, mlp_ratio=4.0, qkv_bias=False, qk_scale=None, drop=0.0, attn_drop=0.0,
                 act_layer=nn.GELU, norm_layer=nn.LayerNorm):
        super().__init__()
        self.norm1 = norm_layer(dim)
        self.attn = Attention(dim, num_heads=num_heads, qkv_bias=qkv_bias, qk_scale=qk_scale,
                              attn_drop=attn_drop, proj_drop=drop)
        self.drop_path = nn.Identity()
        self.norm2 = norm_layer(dim)
        self.mlp = Mlp(in_features=dim, hidden_features=int(dim * mlp_ratio), act_layer=act_layer, drop=drop)

    def forward(self, x):
        y, _ = self.attn(self.norm1(x))
        x = x + self.drop_path(y)
        x = x + self.drop_path(self.mlp(self.norm2(x)))
        return x


class PatchEmbed(nn.Module):
    def __init__(self, img_size=224, patch_size=16, in_chans=3, embed_dim=768):
        super().__init__()
        self.img_size = img_size
        self.patch_size = patch_size
        self.num_patches = (img_size // patch_size) * (img_size // patch_size)
        self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=patch_size)

    def forward(self, x):
        return self.proj(x).flatten(2).transpose(1, 2)


class VisionTransformer(nn.Module):
    def __init__(self, img_size=224, patch_size=16, in_chans=3, embed_dim=768, depth=12, num_heads=12,
                 mlp_ratio=4.0, qkv_bias=True, norm_layer=None, act_layer=nn.GELU):
        super().__init__()
        norm_layer = norm_layer or partial(nn.LayerNorm, eps=1e-6)
        self.num_features = self.embed_dim = embed_dim
        self.patch_embed = PatchEmbed(img_size=img_size, patch_size=patch_size, in_chans=in_chans, embed_dim=embed_dim)
        num_patches = self.patch_embed.num_patches
        self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))
        self.pos_embed = nn.Parameter(torch.zeros(1, num_patches + 1, embed_dim))
        self.pos_drop = nn.Dropout(p=0.0)
        self.blocks = nn.ModuleList([
            Block(dim=embed_dim, num_heads=num_heads, mlp_ratio=mlp_ratio, qkv_bias=qkv_bias,
                  act_layer=act_layer, norm_layer=norm_layer)
            for _ in range(depth)
        ])
        self.norm = norm_layer(embed_dim)
        self.head = nn.Identity()

    def interpolate_pos_encoding(self, x, w, h):
        npatch = x.shape[1] - 1
        N = self.pos_embed.shape[1] - 1
        if npatch == N and w == h:
            return self.pos_embed
        class_pos_embed = self.pos_embed[:, 0]
        patch_pos_embed = self.pos_embed[:, 1:]
        dim = x.shape[-1]
        w0, h0 = w // self.patch_embed.patch_size + 0.1, h // self.patch_embed.patch_size + 0.1
        patch_pos_embed = nn.functional.interpolate(
            patch_pos_embed.reshape(1, int(math.sqrt(N)), int(math.sqrt(N)), dim).permute(0, 3, 1, 2),
            scale_factor=(w0 / math.sqrt(N), h0 / math.sqrt(N)), mode="bicubic",
        )
        patch_pos_embed = patch_pos_embed.permute(0, 2, 3, 1).view(1, -1, dim)
        return torch.cat((class_pos_embed.unsqueeze(0), patch_pos_embed), dim=1)

    def prepare_tokens(self, x):
        B, nc, w, h = x.shape
        x = self.patch_embed(x)
        cls_tokens = self.cls_token.expand(B, -1, -1)
        x = torch.cat((cls_tokens, x), dim=1)
        x = x + self.interpolate_pos_encoding(x, w, h)
        return self.pos_drop(x)

    def forward(self, x, apply_norm=True):
        x = self.prepare_tokens(x)
        for blk in self.blocks:
            x = blk(x)
        if apply_norm:
            x = self.norm(x)
        return x[:, 0]


class DINOHead(nn.Module):
    # Present to receive dino's projection weights (unused for the ensemble's cls feature).
    def __init__(self, in_dim, out_dim, hidden_dim=2048, bottleneck_dim=256):
        super().__init__()
        self.mlp = nn.Sequential(
            nn.Linear(in_dim, hidden_dim), nn.GELU(),
            nn.Linear(hidden_dim, hidden_dim), nn.GELU(),
            nn.Linear(hidden_dim, bottleneck_dim),
        )
        self.last_layer = nn.utils.weight_norm(nn.Linear(bottleneck_dim, out_dim, bias=False))

    def forward(self, x):
        x = nn.functional.normalize(self.mlp(x), dim=-1, p=2)
        return self.last_layer(x)


def _vit_base(act_layer=nn.GELU, norm_eps=1e-6):
    return VisionTransformer(
        patch_size=16, embed_dim=768, depth=12, num_heads=12, mlp_ratio=4, qkv_bias=True,
        norm_layer=partial(nn.LayerNorm, eps=norm_eps), act_layer=act_layer,
    )


# ── DreamSim ensemble ─────────────────────────────────────────────────────────


class _Extractor(nn.Module):
    """One backbone (``model``) + its projection (``proj``). Matches the upstream
    ViTExtractor module names so the merged safetensors load 1:1."""

    def __init__(self, model_type: str):
        super().__init__()
        if model_type == "dino_vitb16":
            self.model = _vit_base(act_layer=nn.GELU, norm_eps=1e-6)
            self.proj = DINOHead(768, 2048)
        elif model_type in ("clip_vitb16", "open_clip_vitb16"):
            _act = QuickGELU if model_type == "clip_vitb16" else nn.GELU
            self.model = _vit_base(act_layer=_act, norm_eps=1e-5)
            self.model.pos_drop = nn.LayerNorm(self.model.embed_dim, eps=1e-5)  # loaders swap Dropout→LayerNorm
            self.proj = nn.Parameter(torch.zeros(768, 512))  # clip/open_clip embedding projection (raw tensor)
        else:
            raise ValueError(f"unsupported DreamSim backbone: {model_type}")


class DreamSimOutput(ModelOutput):
    embeddings: Optional[torch.Tensor] = None
    last_hidden_states: Optional[torch.Tensor] = None


DreamSimOutput = dataclass(DreamSimOutput)


class DreamSimModel(PreTrainedModel):
    config_class = DreamSimConfig
    main_input_name = "pixel_values"

    def __init__(self, config: DreamSimConfig):
        super().__init__(config)
        self._model_types = config.model_types.split(",")
        self._feat_types = config.feat_types.split(",")
        self.normalize_embeds = config.normalize_embeds
        self.extractor_list = nn.ModuleList([_Extractor(_m) for _m in self._model_types])
        self.mlp = nn.Identity()  # ensemble uses LoRA → Identity MLP head
        # Per-backbone input normalization stats — plain Python, NOT buffers. Non-persistent
        # buffers are left uninitialized by meta-device from_pretrained (transformers 5.x),
        # which silently corrupts normalization; build the tensors inline in _extract_one.
        self._pixel_mean = [OPENAI_CLIP_MEAN if "clip" in _m else IMAGENET_MEAN for _m in self._model_types]
        self._pixel_std = [OPENAI_CLIP_STD if "clip" in _m else IMAGENET_STD for _m in self._model_types]
        self.post_init()

    def _extract_one(self, index: int, pixel_values: torch.Tensor) -> torch.Tensor:
        _extractor = self.extractor_list[index]
        _feat_type = self._feat_types[index]
        _mean = torch.tensor(self._pixel_mean[index], device=pixel_values.device, dtype=pixel_values.dtype).view(1, 3, 1, 1)
        _std = torch.tensor(self._pixel_std[index], device=pixel_values.device, dtype=pixel_values.dtype).view(1, 3, 1, 1)
        _x = (pixel_values - _mean) / _std
        if _feat_type == "cls":
            # DINO: CLS token of the last block's output (pre-final-norm).
            return _extractor.model(_x, apply_norm=False)
        # CLIP / OpenCLIP: post-norm CLS token projected by ``proj``.
        return _extractor.model(_x, apply_norm=True) @ _extractor.proj

    def forward(self, pixel_values: torch.Tensor, **kwargs) -> DreamSimOutput:
        _feats = [self._extract_one(_i, pixel_values) for _i in range(len(self.extractor_list))]
        _concat = torch.cat(_feats, dim=-1)
        _embeddings = self.mlp(_concat)
        if self.normalize_embeds:
            _embeddings = self._normalize_embedding(_embeddings)
        return DreamSimOutput(embeddings=_embeddings, last_hidden_states=_concat)

    @staticmethod
    def _normalize_embedding(embed: torch.Tensor) -> torch.Tensor:
        # Subtract per-sample mean, divide by per-sample L2 norm (upstream normalize_embeds).
        embed = (embed.T - torch.mean(embed, dim=1)).T
        return (embed.T / torch.norm(embed, dim=1)).T

    @torch.no_grad()
    def compute_distance(self, pixel_values_a: torch.Tensor, pixel_values_b: torch.Tensor) -> torch.Tensor:
        """Perceptual distance ``1 - cos`` between two preprocessed image batches."""
        _a = self.forward(pixel_values_a).embeddings
        _b = self.forward(pixel_values_b).embeddings
        return 1 - nn.functional.cosine_similarity(_a, _b, dim=-1)