# 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)