app.py

import spaces
from diffusers import StableDiffusionXLPipeline
from diffusers import DiffusionPipeline
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 
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}.jpeg"

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

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


# Start of the Gradio Blocks interface
with gr.Blocks() as demo:
    with gr.Column():
        gr.Markdown("# Image Generation with SSD-1B")
        gr.Markdown("Enter a prompt and (optionally) a negative prompt to generate an image.")
        
        # Input fields for positive and negative prompts
        with gr.Row():
            prompt1 = gr.Textbox(label="Enter prompt")
            negative_prompt = gr.Textbox(label="Enter negative prompt (optional)")
        
        # Button for generating the image
        generate_button1 = gr.Button("Generate Image")

        # Output image display, set to a larger default size
        output_image1 = gr.Image(label="Generated Image")

        # Click event for the generate button
        generate_button1.click(
            generate_and_save_image, 
            inputs=[prompt1, negative_prompt], 
            outputs=output_image1
        )


# Launch the combined Gradio app
demo.launch()