Build an AI Image Generator with Stable Diffusion Locally

 

Installation

# Create virtual environment (optional)
python -m venv .venv && . .venv/Scripts/activate  # Windows
# source .venv/bin/activate  # Linux/Mac
# Install PyTorch (CPU or GPU)
pip install torch torchvision --index-url https://download.pytorch.org/whl/cpu
# For CUDA GPU:
# pip install torch torchvision --index-url https://download.pytorch.org/whl/cu121
# Install required packages
pip install diffusers transformers accelerate safetensors pillow

Note: You need a Hugging Face account. Accept the model license and login:

pip install huggingface_hub
huggingface-cli login

---

from diffusers import StableDiffusionPipeline
import torch
# Load model (change to "stabilityai/stable-diffusion-2-1" if you want v2.1)
model_id = "runwayml/stable-diffusion-v1-5"
# Load pipeline
pipe = StableDiffusionPipeline.from_pretrained(
    model_id,
    torch_dtype=torch.float16,   # use float16 for GPU
)
# Move to GPU if available
if torch.cuda.is_available():
    pipe = pipe.to("cuda")
else:
    pipe = pipe.to("cpu")
# Prompt
prompt = "A futuristic city at night with neon lights, ultra-realistic"
# Generate image
image = pipe(prompt).images[0]
# Save output
image.save("output.png")
print("Image saved as output.png")

Full Code (Text-to-Image)

# ============================================

# STABLE DIFFUSION LOCAL IMAGE GENERATOR

# Supports: txt2img, img2img, inpainting

# Author: Yusuf Example

# ============================================

import os

import torch

from diffusers import (

    StableDiffusionPipeline,

    StableDiffusionImg2ImgPipeline,

    StableDiffusionInpaintPipeline

)

from PIL import Image

# ====================================================

# GLOBAL SETTINGS

# ====================================================

MODEL_ID = "runwayml/stable-diffusion-v1-5"   # You can replace with stabilityai/stable-diffusion-2-1

DEVICE = "cuda" if torch.cuda.is_available() else "cpu"

# Create output folder

OUTPUT_DIR = "generated_images"

os.makedirs(OUTPUT_DIR, exist_ok=True)

# ====================================================

# LOADERS

# ====================================================

def load_txt2img():

    """Load pipeline for text-to-image generation"""

    pipe = StableDiffusionPipeline.from_pretrained(

        MODEL_ID,

        torch_dtype=torch.float16 if DEVICE == "cuda" else torch.float32,

    )

    return pipe.to(DEVICE)

def load_img2img():

    """Load pipeline for image-to-image generation"""

    pipe = StableDiffusionImg2ImgPipeline.from_pretrained(

        MODEL_ID,

        torch_dtype=torch.float16 if DEVICE == "cuda" else torch.float32,

    )

    return pipe.to(DEVICE)

def load_inpainting():

    """Load pipeline for inpainting (fill missing areas)"""

    pipe = StableDiffusionInpaintPipeline.from_pretrained(

        MODEL_ID,

        torch_dtype=torch.float16 if DEVICE == "cuda" else torch.float32,

    )

    return pipe.to(DEVICE)

# ====================================================

# FUNCTIONS

# ====================================================

def generate_txt2img(prompt: str, filename: str = "txt2img.png"):

    """Generate image from text"""

    pipe = load_txt2img()

    print(">> Generating image from text prompt...")

    image = pipe(prompt).images[0]

    out_path = os.path.join(OUTPUT_DIR, filename)

    image.save(out_path)

    print(f"✅ Saved: {out_path}")

def generate_img2img(prompt: str, init_image_path: str, strength: float = 0.8, filename: str = "img2img.png"):

    """Generate image from text + base image"""

    pipe = load_img2img()

    print(">> Generating image from text + init image...")

    init_image = Image.open(init_image_path).convert("RGB")

    init_image = init_image.resize((512, 512))  # resize for stable diffusion

    image = pipe(prompt=prompt, image=init_image, strength=strength).images[0]

    out_path = os.path.join(OUTPUT_DIR, filename)

    image.save(out_path)

    print(f"✅ Saved: {out_path}")

def generate_inpainting(prompt: str, init_image_path: str, mask_path: str, filename: str = "inpaint.png"):

    """Generate image by inpainting (fill masked areas)"""

    pipe = load_inpainting()

    print(">> Generating inpainting result...")

    init_image = Image.open(init_image_path).convert("RGB")

    mask_image = Image.open(mask_path).convert("RGB")

    init_image = init_image.resize((512, 512))

    mask_image = mask_image.resize((512, 512))

    image = pipe(prompt=prompt, image=init_image, mask_image=mask_image).images[0]

    out_path = os.path.join(OUTPUT_DIR, filename)

    image.save(out_path)

    print(f"✅ Saved: {out_path}")

# ====================================================

# MAIN PROGRAM

# ====================================================

def main():

    print("===============================================")

    print("   LOCAL STABLE DIFFUSION IMAGE GENERATOR")

    print("   Modes: 1) Text-to-Image  2) Image-to-Image  3) Inpainting")

    print("===============================================")

    choice = input("Choose mode (1/2/3): ").strip()

    if choice == "1":

        # Text to Image

        prompt = input("Enter your prompt: ")

        filename = input("Output filename (default=txt2img.png): ").strip() or "txt2img.png"

        generate_txt2img(prompt, filename)

    elif choice == "2":

        # Image to Image

        prompt = input("Enter your prompt: ")

        init_image_path = input("Enter path to initial image: ").strip()

        strength = float(input("Enter strength (0.0 - 1.0, default=0.8): ") or 0.8)

        filename = input("Output filename (default=img2img.png): ").strip() or "img2img.png"

        generate_img2img(prompt, init_image_path, strength, filename)

    elif choice == "3":

        # Inpainting

        prompt = input("Enter your prompt: ")

        init_image_path = input("Enter path to base image: ").strip()

        mask_path = input("Enter path to mask image: ").strip()

        filename = input("Output filename (default=inpaint.png): ").strip() or "inpaint.png"

        generate_inpainting(prompt, init_image_path, mask_path, filename)

    else:

        print("❌ Invalid choice! Please select 1, 2, or 3.")

if __name__ == "__main__":

    main()