ApoorvBrooklyn/stable-diffusion-implementation

1---2language:3  - en4tags:5  - stable-diffusion6  - pytorch7  - text-to-image8  - image-to-image9  - diffusion-models10  - computer-vision11  - generative-ai12  - deep-learning13  - neural-networks14license: mit15library_name: pytorch16pipeline_tag: text-to-image17base_model: stable-diffusion-v1-518model-index:19  - name: pytorch-stable-diffusion20    results:21      - task:22          type: text-to-image23          name: Text-to-Image Generation24        dataset:25          type: custom26          name: Stable Diffusion v1.527        metrics:28          - type: inference_steps29            value: 5030          - type: cfg_scale31            value: 832          - type: image_size33            value: 512x51234---35 36# PyTorch Stable Diffusion Implementation37 38A complete, from-scratch PyTorch implementation of Stable Diffusion v1.5, featuring both text-to-image and image-to-image generation capabilities. This project demonstrates the inner workings of diffusion models by implementing all components without relying on pre-built libraries.39 40## 🚀 Features41 42- **Text-to-Image Generation**: Create high-quality images from text descriptions43- **Image-to-Image Generation**: Transform existing images using text prompts44- **Complete Implementation**: All components built from scratch in PyTorch45- **Flexible Sampling**: Configurable inference steps and CFG scale46- **Model Compatibility**: Support for various fine-tuned Stable Diffusion models47- **Clean Architecture**: Modular design with separate components for each part of the pipeline48 49## 🏗️ Architecture50 51This implementation includes all the core components of Stable Diffusion:52 53- **CLIP Text Encoder**: Processes text prompts into embeddings54- **VAE Encoder/Decoder**: Handles image compression and reconstruction55- **U-Net Diffusion Model**: Core denoising network with attention mechanisms56- **DDPM Sampler**: Implements the denoising diffusion probabilistic model57- **Pipeline Orchestration**: Coordinates all components for generation58 59## 📁 Project Structure60 61```62├── main/63│   ├── attention.py      # Multi-head attention implementation64│   ├── clip.py           # CLIP text encoder65│   ├── ddpm.py           # DDPM sampling algorithm66│   ├── decoder.py        # VAE decoder for image reconstruction67│   ├── diffusion.py      # U-Net diffusion model68│   ├── encoder.py        # VAE encoder for image compression69│   ├── model_converter.py # Converts checkpoint files to PyTorch format70│   ├── model_loader.py   # Loads and manages model weights71│   ├── pipeline.py       # Main generation pipeline72│   └── demo.py           # Example usage and demonstration73├── data/                 # Model weights and tokenizer files74└── images/               # Input/output images75```76 77## 🛠️ Installation78 79### Prerequisites80 81- Python 3.8+82- PyTorch 1.12+83- Transformers library84- PIL (Pillow)85- NumPy86- tqdm87 88### Setup89 901. **Clone the repository:**91   ```bash92   git clone https://github.com/https://github.com/ApoorvBrooklyn/Stable-Diffusion93   cd pytorch-stable-diffusion94   ```95 962. **Create virtual environment:**97   ```bash98   python -m venv venv99   source venv/bin/activate  # On Windows: venv\Scripts\activate100   ```101 1023. **Install dependencies:**103   ```bash104   pip install torch torchvision torchaudio105   pip install transformers pillow numpy tqdm106   ```107 1084. **Download required model files:**109   - Download `vocab.json` and `merges.txt` from [Stable Diffusion v1.5 tokenizer](https://huggingface.co/ApoorvBrooklyn/stable-diffusion-implementation/tree/main/data)110   - Download `v1-5-pruned-emaonly.ckpt` from [Stable Diffusion v1.5](https://huggingface.co/ApoorvBrooklyn/stable-diffusion-implementation/tree/main/data)111   - Place all files in the `data/` folder112 113## 🎯 Usage114 115### Basic Text-to-Image Generation116 117```python118import model_loader119import pipeline120from transformers import CLIPTokenizer121 122# Initialize tokenizer and load models123tokenizer = CLIPTokenizer("data/vocab.json", merges_file="data/merges.txt")124models = model_loader.preload_models_from_standard_weights("data/v1-5-pruned-emaonly.ckpt", "cpu")125 126# Generate image from text127output_image = pipeline.generate(128    prompt="A beautiful sunset over mountains, highly detailed, 8k resolution",129    uncond_prompt="",  # Negative prompt130    do_cfg=True,131    cfg_scale=8,132    sampler_name="ddpm",133    n_inference_steps=50,134    seed=42,135    models=models,136    device="cpu",137    tokenizer=tokenizer138)139```140 141### Image-to-Image Generation142 143```python144from PIL import Image145 146# Load input image147input_image = Image.open("images/input.jpg")148 149# Generate transformed image150output_image = pipeline.generate(151    prompt="Transform this into a watercolor painting",152    input_image=input_image,153    strength=0.8,  # Controls how much to change the input154    # ... other parameters155)156```157 158### Advanced Configuration159 160- **CFG Scale**: Controls how closely the image follows the prompt (1-14)161- **Inference Steps**: More steps = higher quality but slower generation162- **Strength**: For image-to-image, controls transformation intensity (0-1)163- **Seed**: Set for reproducible results164 165## 🔧 Model Conversion166 167The `model_converter.py` script converts Stable Diffusion checkpoint files to PyTorch format:168 169```bash170python main/model_converter.py --checkpoint_path data/v1-5-pruned-emaonly.ckpt --output_dir converted_models/171```172 173## 🎨 Supported Models174 175This implementation is compatible with:176- **Stable Diffusion v1.5**: Base model177- **Fine-tuned Models**: Any SD v1.5 compatible checkpoint178- **Custom Models**: Models trained on specific datasets or styles179 180### Tested Fine-tuned Models:181- **InkPunk Diffusion**: Artistic ink-style images182- **Illustration Diffusion**: Hollie Mengert's illustration style183 184## 🚀 Performance Tips185 186- **Device Selection**: Use CUDA for GPU acceleration, MPS for Apple Silicon187- **Batch Processing**: Process multiple prompts simultaneously188- **Memory Management**: Use `idle_device="cpu"` to free GPU memory189- **Optimization**: Adjust inference steps based on quality vs. speed needs190 191## 🔬 Technical Details192 193### Diffusion Process194- Implements DDPM (Denoising Diffusion Probabilistic Models)195- Uses U-Net architecture with cross-attention for text conditioning196- VAE handles 512x512 image compression to 64x64 latents197 198### Attention Mechanisms199- Multi-head self-attention in U-Net200- Cross-attention between text embeddings and image features201- Efficient attention implementation for memory optimization202 203### Sampling204- Configurable number of denoising steps205- Classifier-free guidance (CFG) for prompt adherence206- Deterministic generation with seed control207 208## 🤝 Contributing209 210Contributions are welcome! Please feel free to submit pull requests or open issues for:211- Bug fixes212- Performance improvements213- New sampling algorithms214- Additional model support215- Documentation improvements216 217## 📄 License218 219This project is licensed under the MIT License - see the [LICENSE](LICENSE) file for details.220 221## 🙏 Acknowledgments222 223- **Stability AI** for the original Stable Diffusion model224- **OpenAI** for the CLIP architecture225- **CompVis** for the VAE implementation226- **Hugging Face** for the transformers library227 228## 📚 References229 230- [High-Resolution Image Synthesis with Latent Diffusion Models](https://arxiv.org/abs/2112.10752)231- [Denoising Diffusion Probabilistic Models](https://arxiv.org/abs/2006.11239)232- [Learning Transferable Visual Models From Natural Language Supervision](https://arxiv.org/abs/2103.00020)233 234## 📞 Support235 236If you encounter any issues or have questions:237- Open an issue on GitHub238- Check the existing documentation239- Review the demo code for examples240 241---242 243**Note**: This is a research and educational implementation. For production use, consider using the official Stable Diffusion implementations or cloud-based APIs.