import gradio as gr
from diffusers import StableDiffusionPipeline
import torch

# 모델 로드
model_id = "stabilityai/stable-diffusion-2"  # 또는 최신 버전 사용 가능
pipe = StableDiffusionPipeline.from_pretrained(
    model_id,
    torch_dtype=torch.float16,
    revision="fp16",
    use_auth_token=True  # HF_TOKEN 필요 시
)
pipe.to("cuda")

# 생성 함수
def generate_image(prompt):
    image = pipe(prompt, guidance_scale=7.5).images[0]
    return image

# Gradio 인터페이스 구성
with gr.Blocks() as demo:
    gr.Markdown("# 🖼️ Text to Image Generator")
    with gr.Row():
        prompt_input = gr.Textbox(label="Prompt", placeholder="Enter a description...")
        generate_button = gr.Button("Generate")
    output_image = gr.Image(label="Generated Image")

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

# 실행
demo.launch()