import gradio as gr
from diffusers import StableDiffusionPipeline
import torch

# โหลดโมเดล
pipe = StableDiffusionPipeline.from_pretrained(
    "stabilityai/stable-diffusion-2",
    torch_dtype=torch.float16
).to("cuda")

# ฟังก์ชันที่รันเมื่อผู้ใช้กรอก prompt
def generate_image(prompt):
    image = pipe(prompt).images[0]
    return image

# สร้าง Gradio UI
with gr.Blocks() as demo:
    gr.Markdown("## 🧠 Text-to-Image with Stable Diffusion 2")
    with gr.Row():
        prompt_input = gr.Textbox(label="Enter your prompt", placeholder="A fantasy landscape with waterfalls...")
        output_image = gr.Image(label="Generated Image")
    generate_btn = gr.Button("Generate")

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

# ให้ Gradio แสดง UI
demo.launch()