from diffusers import StableDiffusionXLPipeline
import torch
import gradio as gr

# Load the Stable Diffusion XL model
pipeline = StableDiffusionXLPipeline.from_pretrained(
    "stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.float16, variant="fp16", use_safetensors=True
).to("cuda")

def generate_image(prompt, prompt_2):
    # Generate the image based on the prompts
    image = pipeline(prompt=prompt, prompt_2=prompt_2).images[0]
    
    # Convert the PIL image to a format that Gradio can display
    return image

# Define the Gradio interface
iface = gr.Interface(fn=generate_image,
                     inputs=[gr.Textbox(label="Prompt 1: Astronaut in a jungle, cold color palette, muted colors, detailed, 8k"),
                             gr.Textbox(label="Prompt 2: Van Gogh painting")],
                     outputs="image",
                     title="Stable Diffusion XL Image Generator",
                     description="Generate images with Stable Diffusion XL based on two prompts.")

# Launch the Gradio app
iface.launch()