| """ |
| Custom diffusers pipeline for ZoomLDM CDM (DiT backbone). |
| """ |
|
|
| import math |
| import json |
| from dataclasses import dataclass |
| from pathlib import Path |
| from typing import Optional |
|
|
| import torch |
| import torch.nn as nn |
| from diffusers import DDIMScheduler, DiffusionPipeline |
| from diffusers.configuration_utils import ConfigMixin, register_to_config |
| from diffusers.models.modeling_utils import ModelMixin |
| from diffusers.utils import BaseOutput |
|
|
|
|
| def _modulate(x: torch.Tensor, shift: torch.Tensor, scale: torch.Tensor) -> torch.Tensor: |
| return x * (1 + scale.unsqueeze(1)) + shift.unsqueeze(1) |
|
|
|
|
| class Attention(nn.Module): |
| """Minimal ViT-style self-attention with timm-compatible parameter names.""" |
|
|
| def __init__(self, dim: int, num_heads: int, qkv_bias: bool = True): |
| super().__init__() |
| if dim % num_heads != 0: |
| raise ValueError(f"dim ({dim}) must be divisible by num_heads ({num_heads})") |
| self.num_heads = num_heads |
| self.head_dim = dim // num_heads |
| self.scale = self.head_dim**-0.5 |
| self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias) |
| self.proj = nn.Linear(dim, dim, bias=True) |
|
|
| def forward(self, x: torch.Tensor) -> torch.Tensor: |
| bsz, tokens, dim = x.shape |
| qkv = self.qkv(x) |
| qkv = qkv.reshape(bsz, tokens, 3, self.num_heads, self.head_dim) |
| qkv = qkv.permute(2, 0, 3, 1, 4) |
| q, k, v = qkv.unbind(0) |
|
|
| attn = (q @ k.transpose(-2, -1)) * self.scale |
| attn = attn.softmax(dim=-1) |
|
|
| x = attn @ v |
| x = x.transpose(1, 2).reshape(bsz, tokens, dim) |
| return self.proj(x) |
|
|
|
|
| class Mlp(nn.Module): |
| """Minimal timm-like MLP block with matching names.""" |
|
|
| def __init__(self, in_features: int, hidden_features: int): |
| super().__init__() |
| self.fc1 = nn.Linear(in_features, hidden_features, bias=True) |
| self.act = nn.GELU(approximate="tanh") |
| self.fc2 = nn.Linear(hidden_features, in_features, bias=True) |
|
|
| def forward(self, x: torch.Tensor) -> torch.Tensor: |
| return self.fc2(self.act(self.fc1(x))) |
|
|
|
|
| class DiTBlock(nn.Module): |
| def __init__(self, hidden_size: int, num_heads: int, mlp_ratio: float = 4.0): |
| super().__init__() |
| self.norm1 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6) |
| self.attn = Attention(hidden_size, num_heads=num_heads, qkv_bias=True) |
| self.norm2 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6) |
| mlp_hidden_dim = int(hidden_size * mlp_ratio) |
| self.mlp = Mlp(in_features=hidden_size, hidden_features=mlp_hidden_dim) |
| self.adaLN_modulation = nn.Sequential(nn.SiLU(), nn.Linear(hidden_size, 6 * hidden_size, bias=True)) |
|
|
| def forward(self, x: torch.Tensor, c: torch.Tensor) -> torch.Tensor: |
| shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.adaLN_modulation(c).chunk(6, dim=1) |
| x = x + gate_msa.unsqueeze(1) * self.attn(_modulate(self.norm1(x), shift_msa, scale_msa)) |
| x = x + gate_mlp.unsqueeze(1) * self.mlp(_modulate(self.norm2(x), shift_mlp, scale_mlp)) |
| return x |
|
|
|
|
| class TimestepEmbedder(nn.Module): |
| def __init__(self, hidden_size: int, frequency_embedding_size: int = 256): |
| super().__init__() |
| self.mlp = nn.Sequential( |
| nn.Linear(frequency_embedding_size, hidden_size, bias=True), |
| nn.SiLU(), |
| nn.Linear(hidden_size, hidden_size, bias=True), |
| ) |
| self.frequency_embedding_size = frequency_embedding_size |
|
|
| @staticmethod |
| def timestep_embedding(t: torch.Tensor, dim: int, max_period: int = 10000) -> torch.Tensor: |
| half = dim // 2 |
| freqs = torch.exp(-math.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32) / half).to( |
| device=t.device |
| ) |
| args = t[:, None].float() * freqs[None] |
| embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1) |
| if dim % 2: |
| embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1) |
| return embedding |
|
|
| def forward(self, t: torch.Tensor) -> torch.Tensor: |
| t_freq = self.timestep_embedding(t, self.frequency_embedding_size) |
| return self.mlp(t_freq) |
|
|
|
|
| class LabelEmbedder(nn.Module): |
| def __init__(self, num_classes: int, hidden_size: int): |
| super().__init__() |
| self.embedding_table = nn.Embedding(num_classes, hidden_size) |
|
|
| def forward(self, labels: torch.Tensor) -> torch.Tensor: |
| return self.embedding_table(labels) |
|
|
|
|
| class FinalLayer(nn.Module): |
| def __init__(self, hidden_size: int, out_channels: int): |
| super().__init__() |
| self.norm_final = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6) |
| self.linear = nn.Linear(hidden_size, out_channels, bias=True) |
| self.adaLN_modulation = nn.Sequential(nn.SiLU(), nn.Linear(hidden_size, 2 * hidden_size, bias=True)) |
|
|
| def forward(self, x: torch.Tensor, c: torch.Tensor) -> torch.Tensor: |
| shift, scale = self.adaLN_modulation(c).chunk(2, dim=1) |
| x = _modulate(self.norm_final(x), shift, scale) |
| return self.linear(x) |
|
|
|
|
| class CDMDiTModel(ModelMixin, ConfigMixin): |
| @register_to_config |
| def __init__( |
| self, |
| num_patches: int = 65, |
| in_channels: int = 512, |
| hidden_size: int = 768, |
| depth: int = 12, |
| num_heads: int = 12, |
| mlp_ratio: float = 4.0, |
| num_classes: int = 8, |
| learn_sigma: bool = True, |
| ): |
| super().__init__() |
| self.learn_sigma = learn_sigma |
| self.in_channels = in_channels |
| self.out_channels = in_channels * 2 if learn_sigma else in_channels |
| self.num_patches = num_patches |
| self.x_embedder = nn.Linear(in_channels, hidden_size) |
| self.t_embedder = TimestepEmbedder(hidden_size) |
| self.y_embedder = LabelEmbedder(num_classes, hidden_size) |
| self.pos_embed = nn.Parameter(torch.zeros(1, num_patches, hidden_size), requires_grad=False) |
| self.blocks = nn.ModuleList([DiTBlock(hidden_size, num_heads, mlp_ratio=mlp_ratio) for _ in range(depth)]) |
| self.final_layer = FinalLayer(hidden_size, self.out_channels) |
|
|
| def forward(self, x: torch.Tensor, t: torch.Tensor, y: torch.Tensor) -> torch.Tensor: |
| |
| x = x.transpose(1, 2) |
| x = self.x_embedder(x) + self.pos_embed |
| t_emb = self.t_embedder(t) |
| y_emb = self.y_embedder(y) |
| c = t_emb + y_emb |
| for block in self.blocks: |
| x = block(x, c) |
| x = self.final_layer(x, c) |
| return x.transpose(1, 2) |
|
|
|
|
| @dataclass |
| class CDMPipelineOutput(BaseOutput): |
| samples: torch.Tensor |
|
|
|
|
| class CDMDiTPipeline(DiffusionPipeline): |
| def __init__(self, scheduler: DDIMScheduler, cdm: Optional[CDMDiTModel] = None): |
| super().__init__() |
| self.register_modules(scheduler=scheduler) |
| self.cdm = cdm |
| self._cdm_root = None |
| scheduler_path = getattr(getattr(scheduler, "config", None), "_name_or_path", None) |
| if scheduler_path: |
| p = Path(scheduler_path) |
| self._cdm_root = p.parent if p.name == "scheduler" else p |
|
|
| @property |
| def device(self) -> torch.device: |
| self._load_cdm_if_needed() |
| return next(self.cdm.parameters()).device |
|
|
| def to(self, *args, **kwargs): |
| self._load_cdm_if_needed() |
| self.cdm.to(*args, **kwargs) |
| return self |
|
|
| def _load_cdm_if_needed(self): |
| if self.cdm is not None: |
| return |
| if self._cdm_root is None: |
| root_from_cfg = self.config.get("_name_or_path", None) |
| if root_from_cfg: |
| self._cdm_root = Path(root_from_cfg) |
| if self._cdm_root is None: |
| raise RuntimeError("Could not infer model root for loading CDM weights.") |
|
|
| cdm_dir = self._cdm_root / "cdm" |
| with open(cdm_dir / "config.json", encoding="utf-8") as f: |
| cfg = json.load(f) |
| cfg.pop("_class_name", None) |
| cfg.pop("_diffusers_version", None) |
|
|
| cdm = CDMDiTModel(**cfg) |
| safetensors_path = cdm_dir / "diffusion_pytorch_model.safetensors" |
| bin_path = cdm_dir / "diffusion_pytorch_model.bin" |
| if safetensors_path.exists(): |
| from safetensors.torch import load_file |
|
|
| state = load_file(str(safetensors_path)) |
| elif bin_path.exists(): |
| try: |
| state = torch.load(bin_path, map_location="cpu", weights_only=True) |
| except TypeError: |
| state = torch.load(bin_path, map_location="cpu") |
| else: |
| raise FileNotFoundError( |
| "No CDM weights found in cdm/ (expected diffusion_pytorch_model.safetensors or .bin)." |
| ) |
| cdm.load_state_dict(state, strict=True) |
| cdm.eval() |
| self.cdm = cdm |
|
|
| @torch.no_grad() |
| def __call__( |
| self, |
| batch_size: int = 1, |
| magnification: Optional[torch.Tensor] = None, |
| num_inference_steps: int = 50, |
| guidance_scale: float = 1.0, |
| num_patches: Optional[int] = None, |
| return_dict: bool = True, |
| ): |
| self._load_cdm_if_needed() |
| device = self.device |
| dtype = next(self.cdm.parameters()).dtype |
|
|
| if magnification is None: |
| magnification = torch.zeros(batch_size, dtype=torch.long, device=device) |
| else: |
| magnification = magnification.to(device=device, dtype=torch.long) |
| if magnification.ndim == 0: |
| magnification = magnification.view(1) |
|
|
| batch_size = int(magnification.shape[0]) |
| tokens = num_patches or self.cdm.config.num_patches |
| channels = self.cdm.config.in_channels |
|
|
| latents = torch.randn((batch_size, channels, tokens), device=device, dtype=dtype) |
| self.scheduler.set_timesteps(num_inference_steps, device=device) |
|
|
| for t in self.progress_bar(self.scheduler.timesteps): |
| model_in = torch.cat([latents, latents], dim=0) |
| t_batch = t.expand(model_in.shape[0]).to(device) |
| y_in = torch.cat([torch.zeros_like(magnification), magnification], dim=0) |
|
|
| model_out = self.cdm(model_in, t_batch, y_in) |
| eps, _sigma = model_out.chunk(2, dim=1) if self.cdm.config.learn_sigma else (model_out, None) |
| eps_uncond, eps_cond = eps.chunk(2, dim=0) |
| eps_guided = eps_uncond + guidance_scale * (eps_cond - eps_uncond) |
|
|
| latents = self.scheduler.step(eps_guided, t, latents).prev_sample |
|
|
| if not return_dict: |
| return (latents,) |
| return CDMPipelineOutput(samples=latents) |
|
|
