import torch
from diffusers import StableDiffusionPipeline
import gradio as gr

# 모델 로드 (로컬에 없을 경우 허깅페이스에서 자동 다운로드됨)
pipe = StableDiffusionPipeline.from_pretrained(
    "runwayml/stable-diffusion-v1-5",
    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
    revision="fp16" if torch.cuda.is_available() else None,
    use_safetensors=True
)

# GPU 사용 가능하면 이동
if torch.cuda.is_available():
    pipe = pipe.to("cuda")
else:
    pipe = pipe.to("cpu")

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

# Gradio 인터페이스
demo = gr.Interface(
    fn=generate_image,
    inputs=gr.Textbox(label="텍스트를 입력하세요", placeholder="예: A fantasy landscape with flying whales"),
    outputs=gr.Image(label="생성된 이미지"),
    title="🖼️ Text to Image Generator (Stable Diffusion)",
    description="텍스트를 입력하면 인공지능이 이미지를 생성합니다. 모델: runwayml/stable-diffusion-v1-5"
)

# 실행
if __name__ == "__main__":
    demo.launch()