The Challenge

Most medical image generation models require multiple high-end GPUs and massive datasets, putting them out of reach for smaller research labs. I wanted to see if I could build a functional text-to-X-ray generation system that could run on a single consumer GPU.

This was my exploration into making medical AI research more accessible - proving you don’t need a supercomputer to experiment with generative models in healthcare.

Technical Implementation

I built a latent diffusion model with three components:

  • VAE: 26.2M parameters to compress X-rays into 8-channel latent space
  • UNet: 39.66M parameters with cross-attention for denoising
  • BioBERT: Fine-tuned text encoder (only 593K trainable parameters)

Total: 151.81M parameters, with just 43.5M trainable - about 10x smaller than comparable models.

The key was operating in latent space (32×32) instead of pixel space (256×256), reducing computation by 8x. I implemented gradient checkpointing, mixed precision training, and gradient accumulation to fit everything in 8GB VRAM.

Training Process

Trained on the Indiana University Chest X-ray dataset (3,301 frontal images):

  • Stage 1: VAE training for 67 epochs until convergence
  • Stage 2: Diffusion model training for 387 epochs
  • Hardware: Single RTX 4060 GPU (8GB VRAM)
  • Total time: ~96 hours (vs 552+ hours for multi-GPU approaches)

The model achieves reasonable metrics (SSIM: 0.82, PSNR: 22.3dB) and generates 256×256 images in 663ms.

What It Can (and Can’t) Do

The model generates plausible-looking chest X-rays from text prompts like:

  • “Normal chest X-ray with clear lungs”
  • “Right lower lobe pneumonia with consolidation”
  • “Cardiomegaly with pulmonary congestion”

But let’s be clear about limitations:

  • No clinical validation - I didn’t have radiologists evaluate the outputs
  • Not for diagnostic use - This is purely a research demonstration
  • Limited resolution - 256×256 isn’t sufficient for clinical detail
  • Single institution data - Trained only on Indiana University dataset

Real Impact

While this isn’t clinically validated, it demonstrates that medical AI research doesn’t require massive resources. The complete codebase, trained models, and detailed documentation are open source, allowing other researchers to:

  • Experiment with medical image generation on consumer hardware
  • Understand optimization techniques for resource-constrained training
  • Build upon the architecture for their own research

Lessons Learned

  • Latent diffusion is incredibly efficient - 8x spatial compression makes consumer GPU training feasible
  • Domain-specific design matters - 8 latent channels worked better than the standard 4 for X-rays
  • Mixed precision + gradient checkpointing + accumulation can reduce memory by 65%
  • You can do meaningful research without A100 GPUs

This project won’t revolutionize radiology, but it might help democratize medical AI research by showing what’s possible with limited resources.

Paper preprint: https://www.preprints.org/manuscript/202506.1783/v1
All code and models: https://github.com/priyam-choksi/cxr-diffusion