import spaces
from diffusers import StableDiffusionXLPipeline
from pydantic import BaseModel
from PIL import Image
import gradio as gr
import torch
import uuid
import io
import os


# Load your model
pipe = StableDiffusionXLPipeline.from_pretrained(
    "segmind/SSD-1B", 
    torch_dtype=torch.float16, 
    use_safetensors=True, 
    variant="fp16"
)
pipe.to("cuda:0")

@spaces.GPU  # Apply the GPU decorator
def generate_and_save_image(prompt, negative_prompt=''):
    # Generate image using the provided prompts
    image = pipe(prompt=prompt, negative_prompt=negative_prompt).images[0]

    # Generate a unique UUID for the filename
    unique_id = str(uuid.uuid4())
    image_path = f"generated_images/{unique_id}.png"

    # Save generated image locally
    os.makedirs('generated_images', exist_ok=True)
    image.save(image_path, format='JPG')

    # Return the path of the saved image to display in Gradio interface
    return image_path

# Define the Gradio interface
iface = gr.Interface(
    fn=generate_and_save_image,
    inputs=[
        gr.Textbox(label="Enter prompt", placeholder="Type your prompt here..."),
        gr.Textbox(label="Enter negative prompt (optional)", placeholder="Type your negative prompt here...")
    ],
    outputs=gr.Image(type="filepath", label="Generated Image"),
    title="Image Generation with Stable Diffusion XL",
    description="Enter a prompt and (optionally) a negative prompt to generate an image."
)

# Launch the Gradio app
iface.launch()