|
|
--- |
|
|
license: mit |
|
|
tags: |
|
|
- pytorch |
|
|
- diffusers |
|
|
- stable-diffusion |
|
|
- latent-diffusion |
|
|
- medical-imaging |
|
|
- brain-mri |
|
|
- multiple-sclerosis |
|
|
- dataset-conditioning |
|
|
--- |
|
|
|
|
|
# Brain MRI Synthesis with Stable Diffusion (fine-tuned with dataset prompts) |
|
|
|
|
|
This model is a **fine-tuned version of Stable Diffusion v1-4** for **prompt-conditioned synthesis of brain MRI FLAIR slices**. It leverages **latent diffusion** and dataset-specific prompts to generate realistic 256x256 FLAIR scans with control over the source dataset's style or distribution. |
|
|
|
|
|
## 🔍 Prompt Conditioning |
|
|
|
|
|
The model introduces three special prompt tokens corresponding to the dataset of origin. During training, each image was paired with a prompt indicating its source: |
|
|
|
|
|
- `"SHIFTS FLAIR MRI"` |
|
|
- `"VH FLAIR MRI"` |
|
|
- `"WMH2017 FLAIR MRI"` |
|
|
|
|
|
These prompts were added as special tokens to the tokenizer, and their embeddings were fine-tuned alongside the U-Net, enabling dataset-specific synthesis. |
|
|
|
|
|
## 🧠 Training Details |
|
|
|
|
|
- **Base Model:** [`CompVis/stable-diffusion-v1-4`](https://huggingface.co/CompVis/stable-diffusion-v1-4) |
|
|
- **Architecture:** Latent Diffusion with U-Net + ResNet + Attention |
|
|
- **Input resolution (latent):** 32x32 |
|
|
- **Output resolution (decoded):** 256x256 pixels |
|
|
- **Datasets:** SHIFTS, VH, and WMH2017 (FLAIR MRI slices) |
|
|
- **Channels:** 4 latent channels |
|
|
- **Epochs:** 50 |
|
|
- **Batch size:** 8 |
|
|
- **Gradient accumulation:** 4 steps |
|
|
- **Optimizer:** AdamW |
|
|
- Learning Rate: `1.0e-4` |
|
|
- Betas: (0.95, 0.999) |
|
|
- Weight Decay: `1.0e-6` |
|
|
- Epsilon: `1.0e-8` |
|
|
- **LR Scheduler:** Cosine schedule with 500 warm-up steps |
|
|
- **Noise Scheduler:** DDPM with: |
|
|
- `num_train_timesteps`: 1000 |
|
|
- `beta_start`: 0.0001 |
|
|
- `beta_end`: 0.02 |
|
|
- `beta_schedule`: "linear" |
|
|
- **Mixed Precision:** Disabled |
|
|
- **Gradient Clipping:** max norm 1.0 |
|
|
- **Hardware:** NVIDIA A30 GPU with 4 dataloader workers |
|
|
|
|
|
## 🧪 Usage |
|
|
|
|
|
You can use this model via the `diffusers` library for conditional generation: |
|
|
|
|
|
```python |
|
|
from diffusers import DiffusionPipeline |
|
|
import torch |
|
|
|
|
|
# Load the model |
|
|
pipe = DiffusionPipeline.from_pretrained("benetraco/latent_finetuning") |
|
|
pipe.to("cuda") # or "cpu" |
|
|
|
|
|
# Generate a brain MRI image in SHIFTS style |
|
|
prompt = "SHIFTS FLAIR MRI" |
|
|
image = pipe(prompt=prompt, num_inference_steps=50, guidance_scale=2.0).images[0] |
|
|
|
|
|
image.show() |
|
|
|
