# Stable Diffusion Image Generator

This guide demonstrates how to generate images from text prompts using the Stable Diffusion model via the Hugging Face `diffusers` library.

**Modules Used:**
*   `diffusers`: The go-to library for state-of-the-art pretrained diffusion models.
*   `torch`: PyTorch, the deep learning framework backend.
*   [[programming/python/modules/argparse-module|argparse]]: To handle command-line arguments.

## Prerequisites

*   **GPU**: A GPU with CUDA support (NVIDIA) is highly recommended. Running on CPU is possible but significantly slower.

## Installation

```bash
pip install diffusers transformers accelerate torch
```

## The Code

Save this as `generate_image.py`.

```python
from diffusers import StableDiffusionPipeline
import torch
import argparse
import sys

def generate(prompt, output_file, steps, device_type):
    print(f"Initializing Stable Diffusion pipeline on {device_type}...")
    
    # Using v1-5 as it is a stable and widely compatible baseline
    model_id = "runwayml/stable-diffusion-v1-5"
    
    try:
        # Use float16 for GPU to save memory/speed, float32 for CPU
        dtype = torch.float16 if device_type == "cuda" else torch.float32
        
        pipe = StableDiffusionPipeline.from_pretrained(
            model_id, 
            torch_dtype=dtype,
            use_safetensors=True
        )
        
        pipe = pipe.to(device_type)
        
        # Optimization for lower VRAM (optional)
        if device_type == "cuda":
            # Offloads parts of the model to CPU when not in use to save VRAM
            pipe.enable_model_cpu_offload()
        
        print(f"Generating image for: '{prompt}'")
        # The pipe returns a specialized output class, .images[0] gets the PIL image
        image = pipe(prompt, num_inference_steps=steps).images[0]
        
        image.save(output_file)
        print(f"Success! Image saved to '{output_file}'")
        
    except Exception as e:
        print(f"Error: {e}")
        if device_type == "cuda":
            print("If you ran out of memory, try adding --cpu or closing other GPU apps.")

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Stable Diffusion Image Generator")
    parser.add_argument("prompt", help="Text description of the image to generate")
    parser.add_argument("-o", "--output", default="generated.png", help="Output filename (default: generated.png)")
    parser.add_argument("--steps", type=int, default=50, help="Inference steps (higher = better quality, slower. Default: 50)")
    parser.add_argument("--cpu", action="store_true", help="Force CPU usage (Very slow)")
    
    args = parser.parse_args()
    
    device = "cuda" if torch.cuda.is_available() and not args.cpu else "cpu"
    
    if device == "cpu":
        print("Warning: Running on CPU. This will be slow.")
        
    generate(args.prompt, args.output, args.steps, device)
```

## Usage

```bash
# Generate an image on GPU (default)
python generate_image.py "A futuristic city on Mars, cyberpunk style, high detail"

# Generate with specific filename and steps
python generate_image.py "A cute cat in a spacesuit" -o cat_space.png --steps 30
```

[[programming/python/python]]