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.- 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
pip install diffusers transformers accelerate torchThe Code
Save this as generate_image.py.
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
# 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