import gradio as gr
from diffusers import StableDiffusionPipeline
import torch
import spaces

@spaces.GPU
def generate_image(prompt):
    pipeline = StableDiffusionPipeline.from_pretrained("Mitsua/vroid-diffusion-test-unconditional")
    pipeline.to(torch.device("cuda" if torch.cuda.is_available() else "cpu"))
    
    image = pipeline(prompt).images[0]
    return image

with gr.Blocks() as demo:
    gr.Markdown("# VRoid Diffusion Test Unconditional")
    with gr.Row():
        prompt = gr.Textbox(label="Prompt", placeholder="Enter your prompt here")
    
    with gr.Row():
        generate_button = gr.Button("Generate Image")
    
    output_image = gr.Image(label="Generated Image")
    
    generate_button.click(fn=generate_image, inputs=prompt, outputs=output_image)
    
if __name__ == "__main__":
    demo.launch()