import gradio as gr
import torch
from diffusers import StableDiffusionPipeline, EulerDiscreteScheduler

model_id = "gagong/Traditional-Korean-Painting-Model-v2.0"
scheduler = EulerDiscreteScheduler.from_pretrained(model_id, subfolder="scheduler")
pipe = StableDiffusionPipeline.from_pretrained(model_id, scheduler=scheduler, torch_dtype=torch.float16)
pipe = pipe.to("cuda")

# prompt = "a photo of an astronaut riding a horse on mars"
# image = pipe(prompt).images[0]


def generate_image(prompt):
    if not pipe:
        return "Model not loaded properly"
    try:
        image = pipe(prompt).images[0]
        return image
    except Exception as e:
        print(f"Error generating image: {e}")
        return "Error generating image"

# Gradio 인터페이스 설정
with gr.Blocks() as demo:
    gr.Markdown("# Traditional Korean Painting Generator")
    gr.Markdown("Enter a prompt to generate a traditional Korean painting.")
    
    with gr.Row():
        with gr.Column():
            prompt = gr.Textbox(label="Prompt", placeholder="Describe the scene...")
            generate_btn = gr.Button("Generate")
        with gr.Column():
            output_image = gr.Image(label="Generated Image", type="pil")

    generate_btn.click(fn=generate_image, inputs=prompt, outputs=output_image)

if __name__ == "__main__":
    demo.launch()