app.py

import gradio as gr
from PIL import Image
from diffusers import AutoPipelineForImage2Image, AutoPipelineForText2Image
import torch
import math

# Setup for torch devices and pipelines based on the environment
device = "cuda" if torch.cuda.is_available() else "cpu"
torch_dtype = torch.float16 if torch.cuda.is_available() else torch.float32

# Load the pipelines
i2i_pipe = AutoPipelineForImage2Image.from_pretrained(
    "stabilityai/sdxl-turbo", torch_dtype=torch_dtype
).to(device)

t2i_pipe = AutoPipelineForText2Image.from_pretrained(
    "stabilityai/sdxl-turbo", torch_dtype=torch_dtype
).to(device)

# Prompts for the dropdown menu
prompts = {
    "Island": "fantasy parchment ((map)) of magical islands, elegant, ocean, sea, waves, cartography, vintage, crosshatch, seafaring",
    "City": "city map, urban layout, streets, zones, panoramic, detailed",
    "Continent": "continental map, broad geographic features, mountains, rivers, topography",
}

# Modifier list
modifiers = [
    "Rocky",
    "Cliffs",
    "Lush",
    "Sandy",
    "Urban",
    "Greek",
    "Edo Period",
    "Cyberpunk",
    "Snowy",
    "Dense Forest",
    "Barren",
    "Spooky",
]


def resize_crop(image, size=512):
    # Convert the image to 'RGB' if it's not
    image = image.convert("RGB")
    # Get dimensions of the original image
    w, h = image.size

    # Check if the dimensions are too small
    if w < size or h < size:
        scale = max(size / w, size / h)
        new_w, new_h = int(w * scale), int(h * scale)
        image = image.resize((new_w, new_h), Image.BICUBIC)

    crop_x = (image.width - size) // 2
    crop_y = (image.height - size) // 2
    image = image.crop((crop_x, crop_y, crop_x + size, crop_y + size))
    return image


def predict(init_image, landscape_type, modifier_checks, strength, steps, seed=1231231):
    try:
        generator = torch.Generator(device).manual_seed(seed)
        prompt = prompts[landscape_type]

        # Append selected modifiers to the prompt as a comma-separated string
        if modifier_checks:
            prompt += ", " + ", ".join(modifier_checks)

        # Ensure a minimum number of steps
        min_steps = 2
        if steps < min_steps:
            print(
                f"Number of steps provided ({steps}) is less than minimum required ({min_steps}). Adjusting to minimum."
            )
            steps = min_steps

        if init_image is not None:
            init_image = resize_crop(init_image)
            if init_image.size[0] == 0 or init_image.size[1] == 0:
                raise ValueError("Processed image has zero size.")
            print(
                f"Processed image size: {init_image.size}, using {steps} steps for processing."
            )  # Debugging line

            result = i2i_pipe(
                prompt=prompt,
                image=init_image,
                strength=strength,
                num_inference_steps=steps,
                guidance_scale=0.0,
                generator=generator,
            ).images[0]
        else:
            result = t2i_pipe(
                prompt=prompt,
                num_inference_steps=steps,
                guidance_scale=0.0,
                generator=generator,
                width=512,
                height=512,
            ).images[0]
        return result
    except Exception as e:
        print("Error during prediction:", str(e))
        raise e


# Setup Gradio interface
interface = gr.Interface(
    fn=predict,
    inputs=[
        gr.Image(
            sources=["upload", "webcam", "clipboard"],
            type="pil",
            label="Upload Image (Optional)",
        ),
        gr.Dropdown(
            choices=list(prompts.keys()),
            value="Island",
            label="Choose a Landscape Type",
        ),
        gr.CheckboxGroup(label="Add Modifiers", choices=modifiers),
        gr.Slider(minimum=0.1, maximum=1.0, step=0.1, value=0.5, label="Strength"),
        gr.Slider(minimum=1, maximum=10, step=1, value=5, label="Number of Steps"),
        gr.Number(label="Seed", value=1231231),
    ],
    outputs=gr.Image(type="pil"),
    title="Stable Diffusion Image Generation",
    description="Upload an image or select a landscape type to generate new images. Select any modifiers as desired.",
)

interface.launch()