Update model card: zero-SNR/v-pred, corrected results, w=1

README.md

@@ -15,64 +15,59 @@ pipeline_tag: image-to-image
 
 # Conditional Latent Diffusion Model for Retinal Future-State Synthesis
 
-Trained model weights for predicting two-year follow-up retinal fundus images from baseline photographs and clinical metadata.
+Trained weights for predicting two-year follow-up retinal fundus images from a baseline photograph and a seven-variable clinical profile (age, sex, glycated hemoglobin, fasting glucose, DR severity grade, hypertension status, follow-up interval).
 
-## Model Description
-
-This model adapts Stable Diffusion 1.5 for longitudinal retinal image prediction. It consists of two components:
-
-1. **Fine-tuned VAE** (`vae_best.pt`, 320 MB): SD 1.5 VAE encoder/decoder fine-tuned on retinal fundus images with L1 + SSIM + LPIPS + KL loss. Achieves SSIM 0.954 on reconstruction.
-
-2. **Conditional U-Net** (`diffusion_best.pt`, 13 GB): 860M-parameter denoising U-Net with 15-channel input (4 noisy latent + 4 baseline latent + 7 clinical feature maps). Trained for 500 epochs with cosine LR schedule, EMA, and classifier-free guidance dropout.
-
-## Performance
-
-| Metric | Value |
-|--------|-------|
-| SSIM | 0.762 |
-| PSNR | 17.26 dB |
-| LPIPS | 0.379 |
-| FID | 107.28 |
+## Files
 
-Evaluated on 110 held-out test image pairs.
+| File | Size | Description |
+|------|------|-------------|
+| `diffusion_v3_zerosnr_vpred_epoch500.pt` | ~6.9 GB | Conditional denoising U-Net (860M params, 15-channel input) plus the clinical encoder. Trained for 500 epochs (seed 42) under a zero-terminal-SNR schedule with the velocity-prediction objective. Contains both raw and EMA weights; the reported numbers use the raw (non-EMA) weights. Optimizer state is stripped. |
+| `vae_finetuned.pt` | ~335 MB | SD 1.5 VAE fine-tuned on retinal fundus images (reconstruction SSIM 0.954). Stored under `model_state_dict`. |
 
-## Qualitative Results
-
-![Qualitative comparison](fig_qualitative_comparison.png)
-
-Each row shows a different test patient. Columns: baseline fundus, ground-truth follow-up, our prediction, Regression U-Net, and Pix2Pix. Our diffusion model generates sharper, more realistic retinal textures compared to deterministic baselines.
+## Model Description
 
-## Training Dynamics
+The U-Net takes a 15-channel latent input formed by concatenating the noisy target latent (4), the encoded baseline latent (4), and a per-feature clinical map (7). It is initialized from Stable Diffusion 1.5 and trained on eye-corrected, registered baseline/follow-up pairs. The zero-terminal-SNR schedule with velocity prediction removes a color-space artifact mode and reduces the low-quality prediction rate from 26.4% to 5.5%. Optimal inference uses DDIM with guidance w = 1 (no classifier-free guidance); larger guidance over-exposes the output.
 
-![Training dynamics](fig_training_dynamics.png)
+From the single trained model two estimators span the distortion-perception frontier: a single sample (best perceptual quality) and the posterior mean over K = 12 samples (best structural similarity).
 
-(a, b) Stage 1: VAE fine-tuning over 50 epochs reaching SSIM 0.954. (c-f) Stage 2: U-Net training over 500 epochs with cosine LR schedule and warmup.
+## Performance
 
-## Guidance Scale Sweep
+Held-out test set (n = 110 pairs), field-of-view-masked metrics.
 
-![Guidance scale](fig_guidance_scale.png)
+| Estimator | SSIM | PSNR (dB) | LPIPS | FID |
+|-----------|:----:|:---------:|:-----:|:---:|
+| Single sample (w = 1) | 0.791 | 21.60 | 0.123 | 33.2 |
+| Posterior mean (K = 12) | 0.809 | 21.39 | 0.175 | 103.3 |
 
-SSIM peaks at guidance scale 7.5, while FID increases monotonically with stronger guidance, reflecting the fidelity-diversity tradeoff.
+Five-seed means: single-sample SSIM 0.781 +/- 0.006, FID 32.5 +/- 0.6. No baseline retrained on the same corrected data significantly outperforms these on any metric under paired Wilcoxon testing.
 
 ## Usage
 
+The full sampler is in the code repository (`src/inference/diffusion_sampler.py`). It loads the two checkpoints as follows:
+
 ```python
-import torch
-from diffusers import AutoencoderKL, UNet2DConditionModel
-
-# Load VAE
-vae = AutoencoderKL.from_pretrained("runwayml/stable-diffusion-v1-5", subfolder="vae")
-vae_state = torch.load("vae_best.pt", map_location="cpu")
-if "model_state_dict" in vae_state:
-    vae_state = vae_state["model_state_dict"]
-vae.load_state_dict(vae_state, strict=False)
-
-# Load U-Net (requires modified conv_in for 15 input channels)
-unet = UNet2DConditionModel.from_pretrained("runwayml/stable-diffusion-v1-5", subfolder="unet")
-# ... modify conv_in and load checkpoint
-# See full inference code at the GitHub repository
+from src.inference.diffusion_sampler import load_model, sample
+
+unet, clin, vae = load_model(
+    "diffusion_v3_zerosnr_vpred_epoch500.pt",
+    "vae_finetuned.pt",
+    device="cuda",
+    use_ema=False,           # raw weights reproduce the reported numbers
+)
+
+pred = sample(
+    unet, clin, vae,
+    baseline,                # (1, 3, 512, 512) in [-1, 1]
+    clinical,                # (1, 7) standardized clinical vector
+    guidance_scale=1.0,
+    num_steps=50,
+    prediction_type="v_prediction",
+    zero_snr=True,
+)[0]
 ```
 
+The checkpoint exposes `unet_state_dict`, `clinical_encoder_state_dict`, `ema_unet_state_dict`, `ema_clinical_state_dict`, and `config`; the VAE exposes `model_state_dict`.
+
 ## Links
 
 - **Code:** [github.com/Usama1002/retinal-diffusion](https://github.com/Usama1002/retinal-diffusion)
@@ -84,7 +79,7 @@ unet = UNet2DConditionModel.from_pretrained("runwayml/stable-diffusion-v1-5", su
 @article{usama2026retinal,
   title={Conditional Latent Diffusion for Predictive Retinal Fundus Image Synthesis from Baseline Imaging and Clinical Metadata},
   author={Usama, Muhammad and Pazo, Emmanuel Eric and Li, Xiaorong and Liu, Juping},
-  journal={Computers in Biology and Medicine (under review)},
+  note={Manuscript under review},
   year={2026}
 }
 ```