Add files using upload-large-folder tool

README.md

@@ -0,0 +1,91 @@
+---
+license: apache-2.0
+library_name: diffusers
+pipeline_tag: unconditional-image-generation
+tags:
+  - zoomldm
+  - cdm
+  - dit
+  - histopathology
+  - brca
+  - custom-pipeline
+---
+
+# BiliSakura/ZoomLDM-CDM-brca
+
+Diffusers-style wrapped **CDM (DiT)** checkpoint for BRCA, converted from ZoomLDM `cdm_dit` training outputs.
+
+## Model Description
+
+- **Architecture:** DiT-B style conditioning diffusion model (CDM)
+- **Domain:** BRCA conditioning space used by ZoomLDM
+- **Output:** conditioning tokens/embeddings (`(B, 512, 65)`)
+- **Format:** custom diffusers pipeline (`pipeline.py`)
+
+## Intended Use
+
+Use this model to sample BRCA conditioning embeddings that can be consumed by downstream ZoomLDM workflows.
+
+## Out-of-Scope Use
+
+- Not a complete pixel-space generator by itself.
+- Not intended for clinical or diagnostic use.
+- Not validated for non-BRCA domains without adaptation.
+
+## Files
+
+- `pipeline.py`: custom `DiffusionPipeline` implementation (`CDMDiTPipeline`)
+- `model_index.json`: diffusers metadata
+- `cdm/`: active model weights/config used by pipeline
+- `scheduler/`: DDIM scheduler config
+- `model_raw.safetensors`: non-EMA training weights (optional)
+- `optimizer.pt`: optimizer state (optional)
+- `config.json`: conversion metadata
+
+## Usage
+
+```python
+import torch
+from diffusers import DiffusionPipeline
+
+pipe = DiffusionPipeline.from_pretrained(
+    "BiliSakura/ZoomLDM-CDM-brca",
+    custom_pipeline="pipelin.py",
+    trust_remote_code=True,
+).to("cuda")
+
+out = pipe(
+    batch_size=2,
+    magnification=torch.tensor([0, 0], device="cuda"),  # class labels 0..7
+    num_inference_steps=50,
+    guidance_scale=1.0,
+)
+
+samples = out.samples  # (B, 512, 65)
+```
+
+## Limitations
+
+- Produces conditioning embeddings, not final images.
+- Requires correct class/magnification label conventions.
+- Inherits data biases and quality limits from the original training data.
+
+## Citation
+
+```bibtex
+@InProceedings{Yellapragada_2025_CVPR,
+  author = {Yellapragada, Srikar and Graikos, Alexandros and Triaridis, Kostas and Prasanna, Prateek and Gupta, Rajarsi and Saltz, Joel and Samaras, Dimitris},
+  title = {ZoomLDM: Latent Diffusion Model for Multi-scale Image Generation},
+  booktitle = {Proceedings of the Computer Vision and Pattern Recognition Conference (CVPR)},
+  month = {June},
+  year = {2025},
+  pages = {23453-23463}
+}
+
+@inproceedings{Peebles2023DiT,
+  title={Scalable Diffusion Models with Transformers},
+  author={Peebles, William and Xie, Saining},
+  booktitle={Proceedings of the IEEE/CVF International Conference on Computer Vision (ICCV)},
+  year={2023}
+}
+```

cdm/config.json

@@ -0,0 +1,12 @@
+{
+  "_class_name": "CDMDiTModel",
+  "_diffusers_version": "0.30.0",
+  "num_patches": 65,
+  "in_channels": 512,
+  "hidden_size": 768,
+  "depth": 12,
+  "num_heads": 12,
+  "mlp_ratio": 4.0,
+  "num_classes": 8,
+  "learn_sigma": true
+}

cdm/diffusion_pytorch_model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f593642111dd86a2c03994be5dafd1d5c8f170a7bce0c1338849e5b171fb6043
+size 523000344

model_index.json

@@ -0,0 +1,8 @@
+{
+  "_class_name": "CDMDiTPipeline",
+  "_diffusers_version": "0.30.0",
+  "scheduler": [
+    "diffusers",
+    "DDIMScheduler"
+  ]
+}

pipeline.py

@@ -0,0 +1,272 @@
+"""
+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)  # 3, B, H, T, D
+        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: (B, C, T), output: (B, out_channels, T)
+        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)

scheduler/scheduler_config.json

@@ -0,0 +1,19 @@
+{
+  "_class_name": "DDIMScheduler",
+  "_diffusers_version": "0.36.0",
+  "beta_end": 0.012,
+  "beta_schedule": "scaled_linear",
+  "beta_start": 0.00085,
+  "clip_sample": false,
+  "clip_sample_range": 1.0,
+  "dynamic_thresholding_ratio": 0.995,
+  "num_train_timesteps": 1000,
+  "prediction_type": "epsilon",
+  "rescale_betas_zero_snr": false,
+  "sample_max_value": 1.0,
+  "set_alpha_to_one": false,
+  "steps_offset": 1,
+  "thresholding": false,
+  "timestep_spacing": "leading",
+  "trained_betas": null
+}