sd with federated learning

README.md

@@ -15,59 +15,64 @@ pipeline_tag: image-to-image
 
 # Conditional Latent Diffusion Model for Retinal Future-State Synthesis
 
-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).
-
-## Files
-
-| 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`. |
+Trained model weights for predicting two-year follow-up retinal fundus images from baseline photographs and clinical metadata.
 
 ## Model Description
 
-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.
+This model adapts Stable Diffusion 1.5 for longitudinal retinal image prediction. It consists of two components:
 
-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).
+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
 
-Held-out test set (n = 110 pairs), field-of-view-masked metrics.
+| Metric | Value |
+|--------|-------|
+| SSIM | 0.762 |
+| PSNR | 17.26 dB |
+| LPIPS | 0.379 |
+| FID | 107.28 |
 
-| 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 |
+Evaluated on 110 held-out test image pairs.
 
-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.
+## Qualitative Results
 
-## Usage
+![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.
+
+## Training Dynamics
+
+![Training dynamics](fig_training_dynamics.png)
 
-The full sampler is in the code repository (`src/inference/diffusion_sampler.py`). It loads the two checkpoints as follows:
+(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.
+
+## Guidance Scale Sweep
+
+![Guidance scale](fig_guidance_scale.png)
+
+SSIM peaks at guidance scale 7.5, while FID increases monotonically with stronger guidance, reflecting the fidelity-diversity tradeoff.
+
+## Usage
 
 ```python
-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]
+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
 ```
 
-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)
@@ -79,7 +84,7 @@ The checkpoint exposes `unet_state_dict`, `clinical_encoder_state_dict`, `ema_un
 @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},
-  note={Manuscript under review},
+  journal={Computers in Biology and Medicine (under review)},
   year={2026}
 }
 ```

diffusion_best.pt

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

vae_best.pt

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