Buckets:
| """CSD — Contrastive Style Descriptor, for content-invariant style similarity. | |
| From "Measuring Style Similarity in Diffusion Models" (Somepalli et al., 2024). CLIP | |
| ViT-L/14's visual tower with its projection removed, plus two separately-trained | |
| 1024x768 heads (style and content), finetuned with gradient reversal so that the style | |
| head responds to *rendering* — medium, palette, line quality, texture — rather than to | |
| what is depicted. | |
| This is why it is here. CLIP cosine against a FlexBench style reference is substantially a | |
| content match: those references are frames of real videos containing people, rooms and | |
| furniture, and CLIP cosine between two unrelated natural images already sits near 0.5. On | |
| the 2026-08-10 run that left `style_clip` clearing its shuffled-reference chance level by | |
| only ~0.013-0.036 — below the sample-to-sample scatter — so the style anchors were | |
| degenerate on 5 of 18 task cells and had to be dropped. CSD is the purpose-built | |
| alternative: a style descriptor whose zero is not pinned to shared photographic content. | |
| Weights: `tools/CSD/csd-vit-l.pth`, the authors' released ViT-L checkpoint (mirrored at | |
| huggingface.co/yuxi-liu-wired/CSD, originally Google Drive). The state dict carries a | |
| `module.` prefix from DataParallel training and holds `last_layer_style` / | |
| `last_layer_content` alongside the backbone. | |
| Note the mirror's own README reports that style and content heads embed style about | |
| equally well, i.e. the disentanglement is imperfect. Only the style head is used here. | |
| """ | |
| from __future__ import annotations | |
| from pathlib import Path | |
| import clip | |
| import torch | |
| import torch.nn as nn | |
| import torch.nn.functional as F | |
| DEFAULT_CHECKPOINT = Path(__file__).resolve().parents[2] / "tools/CSD/csd-vit-l.pth" | |
| class CSDStyleEncoder(nn.Module): | |
| """CLIP ViT-L/14 visual backbone (projection removed) + the trained style head.""" | |
| def __init__(self, checkpoint_path: Path | str = DEFAULT_CHECKPOINT, device: str = "cuda"): | |
| super().__init__() | |
| checkpoint_path = Path(checkpoint_path) | |
| if not checkpoint_path.exists(): | |
| raise FileNotFoundError( | |
| f"CSD checkpoint not found: {checkpoint_path}\n" | |
| "Fetch it with:\n" | |
| " curl -L -o tools/CSD/csd-vit-l.pth " | |
| "https://huggingface.co/yuxi-liu-wired/CSD/resolve/main/csd-vit-l.pth" | |
| ) | |
| # clip.load pulls ViT-L-14.pt from ~/.cache/clip; only the visual tower is kept. | |
| clip_model, _ = clip.load("ViT-L/14", device="cpu") | |
| self.backbone = clip_model.visual | |
| self.backbone.proj = None # the style head replaces CLIP's own projection | |
| state = torch.load(checkpoint_path, map_location="cpu", weights_only=False) | |
| state = state.get("model_state_dict", state.get("state_dict", state)) | |
| state = {k.removeprefix("module."): v for k, v in state.items()} | |
| self.last_layer_style = nn.Parameter(state["last_layer_style"]) | |
| backbone_state = { | |
| k.removeprefix("backbone."): v for k, v in state.items() if k.startswith("backbone.") | |
| } | |
| missing, unexpected = self.backbone.load_state_dict(backbone_state, strict=False) | |
| # `proj` is deliberately absent (it was replaced by the style head); anything else | |
| # missing means the checkpoint does not match this architecture and the embeddings | |
| # would be silently wrong, so fail loudly instead. | |
| unexpected = [k for k in unexpected if k != "proj"] | |
| missing = [k for k in missing if k != "proj"] | |
| if missing or unexpected: | |
| raise RuntimeError( | |
| f"CSD backbone state mismatch — missing={missing[:5]} unexpected={unexpected[:5]}" | |
| ) | |
| self.to(device).eval() | |
| self.device = device | |
| def dtype(self): | |
| return self.backbone.conv1.weight.dtype | |
| def forward(self, pixel_values: torch.Tensor) -> torch.Tensor: | |
| """[B, 3, 224, 224] CLIP-normalized -> [B, 768] L2-normalized style embeddings.""" | |
| feature = self.backbone(pixel_values.to(self.device, dtype=self.dtype)) | |
| style = feature @ self.last_layer_style.to(feature.dtype) | |
| return F.normalize(style.float(), dim=-1, p=2) | |
| def load_csd(checkpoint_path: Path | str = DEFAULT_CHECKPOINT, device: str = "cuda") -> CSDStyleEncoder: | |
| return CSDStyleEncoder(checkpoint_path, device) | |
Xet Storage Details
- Size:
- 4.42 kB
- Xet hash:
- 0bd61eba591aba6e96ea8f299a4ce78df4bcebc79ea7c17d820c0cdbc75fb4a9
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.