cxr-pneumonia-dreambooth

cxr-pneumonia-dreambooth

Diffuser model finetuned using DreamBooth from stable-diffusion-2-1-base to generate synthetic pneumonia chest x-ray images.

The bones still look a little weird...

---

Generation Prompt

Due to the model being trained using DreamBooth, it is crucial that the correct prompt is used during generation:

A chest xray showing pneumonia infection, lung opacity

---

Training Config

training_config = {
    "method": "full_dreambooth",
    "resolution": 512,
    "train_batch_size": 8,
    "gradient_accumulation_steps": 2,  # Effective batch size: 16
    "learning_rate": 5e-6,  # Lower for full finetuning
    "max_train_steps": 800,
    "train_text_encoder": True,  # Thanks to 80GB A100 in Colab this is possible
    "with_prior_preservation": True,
    "prior_loss_weight": 1.0,
    "num_class_images": 200,
    "mixed_precision": "fp16",
    "gradient_checkpointing": False,
}

---

Validation Images

Train Steps 200

Train Steps 400

Train Steps 600

Train Steps 800