File size: 6,023 Bytes
99ddb81 213477f 99ddb81 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 | # 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)
|