hf-vision/chest-xray-pneumonia
Viewer • Updated • 5.86k • 1.22k • 9
How to use chimbiwide/cxr-normal-dreambooth with Diffusers:
pip install -U diffusers transformers accelerate
import torch
from diffusers import DiffusionPipeline
# switch to "mps" for apple devices
pipe = DiffusionPipeline.from_pretrained("chimbiwide/cxr-normal-dreambooth", dtype=torch.bfloat16, device_map="cuda")
prompt = "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k"
image = pipe(prompt).images[0]Diffuser model finetuned using DreamBooth from stable-diffusion-2-1-base to generate synthetic normal chest x-ray images.
The bones still look a little weird...
Due to the model being trained using DreamBooth, it is crucial that the correct prompt is used during generation:
A chest xray of healthy normal lungs, clear lung fields
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,
}
Train Steps 200
Base model
Manojb/stable-diffusion-2-1-base