Add adaln-ddpm-p02 (PixCell-GE-B-p02)
Browse files- adaln-ddpm-p02/README.md +20 -0
- adaln-ddpm-p02/model.pt +3 -0
- adaln-ddpm-p02/training_config.yaml +114 -0
adaln-ddpm-p02/README.md
ADDED
|
@@ -0,0 +1,20 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
---
|
| 2 |
+
license: apache-2.0
|
| 3 |
+
tags:
|
| 4 |
+
- histopathology
|
| 5 |
+
- diffusion
|
| 6 |
+
- spatial-transcriptomics
|
| 7 |
+
- icml-2026-sd4h-workshop
|
| 8 |
+
---
|
| 9 |
+
|
| 10 |
+
# PixCell-GE-B-p02
|
| 11 |
+
|
| 12 |
+
EMA-only inference weights for the **PixCell-GE-B-p02** row reported in the
|
| 13 |
+
ICML 2026 SD4H workshop submission *Transcriptomics-Conditioned Virtual Tissue
|
| 14 |
+
Synthesis via Diffusion Transformers*.
|
| 15 |
+
|
| 16 |
+
- **Source checkpoint**: `step_0545000_ema.pt`
|
| 17 |
+
- **Architecture**: see `training_config.yaml` in this folder.
|
| 18 |
+
- **License**: Apache-2.0.
|
| 19 |
+
|
| 20 |
+
See the umbrella repo README at `stmdit-anon/stmdit-checkpoints` for usage.
|
adaln-ddpm-p02/model.pt
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:fb3997b8e6d1603ef5972101b32d32cda00ebfa946ed90e4a34d86c38405a4ed
|
| 3 |
+
size 531422045
|
adaln-ddpm-p02/training_config.yaml
ADDED
|
@@ -0,0 +1,114 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# Training Configuration - PixArtGE-B with CancerFoundation Encoder
|
| 2 |
+
# ==================================================================
|
| 3 |
+
# Base variant (130M params) with CancerFoundation gene expression encoder.
|
| 4 |
+
#
|
| 5 |
+
# Usage:
|
| 6 |
+
# run-training training.yaml --lightning
|
| 7 |
+
|
| 8 |
+
output_dir: "/cluster/work/grlab/projects/projects2025-virtual-tissue-gen/scratch/10x_TuPro/PixCell-GE/training/pixart-ge-cf-B-dropout-p02"
|
| 9 |
+
device: "cuda"
|
| 10 |
+
|
| 11 |
+
# ============================================================================
|
| 12 |
+
# MODEL
|
| 13 |
+
# ============================================================================
|
| 14 |
+
|
| 15 |
+
model:
|
| 16 |
+
type: "pixart_ge"
|
| 17 |
+
variant: "B" # 130M params: depth=12, hidden=768, heads=12
|
| 18 |
+
ge_encoder_type: "cancerfoundation"
|
| 19 |
+
ge_hidden_dim: 512
|
| 20 |
+
cf_model_dir: "/cluster/home/pvlachas/leomed-home/pretrained_model_weights/cancer-foundation"
|
| 21 |
+
cf_freeze_backbone: true
|
| 22 |
+
|
| 23 |
+
# ============================================================================
|
| 24 |
+
# DATA
|
| 25 |
+
# ============================================================================
|
| 26 |
+
|
| 27 |
+
data:
|
| 28 |
+
# Dataset-centric path
|
| 29 |
+
features_dir: "/cluster/work/grlab/projects/projects2025-virtual-tissue-gen/scratch/10x_TuPro/feat-extraction/features_train"
|
| 30 |
+
load_gene_expression: true
|
| 31 |
+
load_gene_expression_binned: true
|
| 32 |
+
num_workers: 8
|
| 33 |
+
pin_memory: true
|
| 34 |
+
val_split: 0.1
|
| 35 |
+
|
| 36 |
+
# ============================================================================
|
| 37 |
+
# DIFFUSION
|
| 38 |
+
# ============================================================================
|
| 39 |
+
|
| 40 |
+
diffusion:
|
| 41 |
+
timesteps: 1000
|
| 42 |
+
beta_schedule: "linear"
|
| 43 |
+
image_size: 256
|
| 44 |
+
latent_size: 32
|
| 45 |
+
|
| 46 |
+
# ============================================================================
|
| 47 |
+
# TRAINING
|
| 48 |
+
# ============================================================================
|
| 49 |
+
|
| 50 |
+
training:
|
| 51 |
+
batch_size: 32
|
| 52 |
+
batch_size_val: 32
|
| 53 |
+
gradient_accumulation_steps: 4 # effective batch = 128 (matches original PixCell)
|
| 54 |
+
num_epochs: 1000
|
| 55 |
+
seed: 42
|
| 56 |
+
gradient_clip: 0.01
|
| 57 |
+
ema_rate: 0.9999
|
| 58 |
+
|
| 59 |
+
optimizer:
|
| 60 |
+
lr: 2e-5 # matches original PixCell (was 1e-4)
|
| 61 |
+
weight_decay: 0.01
|
| 62 |
+
betas: [0.9, 0.999]
|
| 63 |
+
|
| 64 |
+
scheduler:
|
| 65 |
+
warmup_steps: 1000
|
| 66 |
+
min_lr_ratio: 0.1
|
| 67 |
+
|
| 68 |
+
classifier_free_guidance:
|
| 69 |
+
conditioning_schedule:
|
| 70 |
+
- mask: [uni, ge] # full conditioning (UNI + GE active)
|
| 71 |
+
weight: 64
|
| 72 |
+
- mask: [ge] # GE only (UNI dropped)
|
| 73 |
+
weight: 16
|
| 74 |
+
- mask: [uni] # UNI only (GE dropped)
|
| 75 |
+
weight: 16
|
| 76 |
+
- mask: [] # unconditional (both dropped)
|
| 77 |
+
weight: 4
|
| 78 |
+
|
| 79 |
+
modality_monitor:
|
| 80 |
+
enabled: true
|
| 81 |
+
diagnostic_freq: 10
|
| 82 |
+
diagnostic_batch_size: 64
|
| 83 |
+
|
| 84 |
+
# ============================================================================
|
| 85 |
+
# DISTRIBUTED
|
| 86 |
+
# ============================================================================
|
| 87 |
+
|
| 88 |
+
distributed:
|
| 89 |
+
precision: "32"
|
| 90 |
+
|
| 91 |
+
# ============================================================================
|
| 92 |
+
# CHECKPOINT
|
| 93 |
+
# ============================================================================
|
| 94 |
+
|
| 95 |
+
checkpoint:
|
| 96 |
+
save_every: 1000
|
| 97 |
+
resume: null
|
| 98 |
+
|
| 99 |
+
# ============================================================================
|
| 100 |
+
# LOGGING
|
| 101 |
+
# ============================================================================
|
| 102 |
+
|
| 103 |
+
logging:
|
| 104 |
+
log_every: 100
|
| 105 |
+
validate_every: 0
|
| 106 |
+
gpu_monitor: true
|
| 107 |
+
gpu_monitor_interval: 60.0
|
| 108 |
+
# Periodic image sampling (disabled; set vae_path + epochs/steps to enable)
|
| 109 |
+
sample_every_epochs: 10
|
| 110 |
+
sample_every_steps: 0
|
| 111 |
+
num_samples: 16
|
| 112 |
+
sample_guidance_scale: 3.0
|
| 113 |
+
sample_num_steps: 20
|
| 114 |
+
sample_vae_path: "/cluster/home/pvlachas/leomed-home/pretrained_model_weights/stability-ai-stable-diffusion-3-5-large/models--stabilityai--stable-diffusion-3.5-large/snapshots/ceddf0a7fdf2064ea28e2213e3b84e4afa170a0f/vae"
|