Spaces:
Runtime error
Runtime error
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") | |
# 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='jpeg/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.", | |
layout="vertical" | |
) | |
# Launch the Gradio app | |
iface.launch() | |