sdxl_generate.py

#!/usr/bin/env python3
"""
Script to generate 32x32 pixel art hippo images using Stable Diffusion XL
in batches, and save them into a "hippo_images" directory.
"""

import os
import torch
from diffusers import StableDiffusionXLPipeline
from PIL import Image
from torchvision import transforms
from tqdm import tqdm
import time


def setup_pipeline():
    """
    Initialize the SDXL pipeline.
    Downloads the model and casts parameters to fp16.
    Moves the pipeline to GPU if available.
    """
    model_id = "stabilityai/stable-diffusion-xl-base-1.0"

    pipeline = StableDiffusionXLPipeline.from_pretrained(
        model_id,
        torch_dtype=torch.float16,
        variant="fp16",
        use_safetensors=True
    )

    if torch.cuda.is_available():
        pipeline = pipeline.to("cuda")
    else:
        print("CUDA not available, using CPU")

    return pipeline


def generate_images(pipeline, prompt, output_size=1024, batch_size=8):
    """
    Generate a batch of images for the provided prompt.

    Args:
        pipeline: The initialized SDXL pipeline.
        prompt (str): Text prompt.
        output_size (int): Height and width (in pixels) of generated image.
        batch_size (int): Number of images to generate in a single call.

    Returns:
        List of PIL.Images generated by the model.
    """
    results = pipeline(
        [prompt] * batch_size,  # Repeat prompt for batch generation
        num_inference_steps=30,
        guidance_scale=7.5,
        height=output_size,
        width=output_size
    )
    return results.images


def process_image(image, target_size=32):
    """
    Downscale the image to target_size x target_size.
    Nearest neighbor interpolation helps keep sharp, pixelated edges.

    Args:
        image (PIL.Image): The generated high-resolution image.
        target_size (int): The desired output resolution (default 32).

    Returns:
        A downsampled PIL.Image.
    """
    transform = transforms.Compose([
        transforms.Resize((target_size, target_size),
                          interpolation=transforms.InterpolationMode.NEAREST)
    ])
    return transform(image)


def save_image(image, filename, output_dir="hippo_face_images"):
    """
    Save the image to disk under the specified output directory.

    Args:
        image (PIL.Image): Image to save.
        filename (str): The name of the file.
        output_dir (str): Directory to save the image.
    """
    os.makedirs(output_dir, exist_ok=True)
    image.save(os.path.join(output_dir, filename))


def main():
    # Initialize pipeline
    print("Initializing SDXL pipeline...")
    pipeline = setup_pipeline()

    # Define the prompt.
    prompt = ("portrait of a hippo, only the face, in the style of Stardew Valley, pixel art, "
              "cute, 2d game art, minimal pastel background, no frame, flat color background")

    # Set up generation parameters.
    num_images = 50000      # Total number of images you want to generate
    batch_size = 16       # Number of images per batch
    image_count = 0       # Global image counter

    # Calculate number of batches
    num_batches = (num_images + batch_size - 1) // batch_size

    print(f"Generating {num_images} images in batches of {batch_size} ...")

    # Create global progress bar
    with tqdm(total=num_images, desc="Total Progress", unit="image") as pbar:
        # Loop over batches
        for i in range(0, num_images, batch_size):
            current_batch_size = min(batch_size, num_images - i)

            # Generate a batch of images (each 1024x1024)
            images = generate_images(
                pipeline, prompt, output_size=1024, batch_size=current_batch_size)

            for img in images:
                # Downscale to 32x32 pixel art image
                downscaled_img = process_image(img, target_size=64)
                image_count += 1
                filename = f"hippo_{image_count:04d}_{time.time()}.png"
                save_image(downscaled_img, filename,
                           output_dir="hippo_images_face")

                # Update progress bar
                pbar.update(1)

    print("\nDataset generation complete! Check the 'hippo_images' directory.")


if __name__ == "__main__":
    main()