import gradio as gr
from diffusers import StableDiffusionPipeline

# 加載 Hugging Face 上的 SANA 模型
model_id = "Efficient-Large-Model/Sana_1600M_1024px"
pipeline = StableDiffusionPipeline.from_pretrained(model_id)
pipeline.to("cuda")  # 如果有 GPU，加速推理

# 定義圖像生成邏輯
def generate_image(prompt, steps=50, guidance_scale=7.5):
    image = pipeline(prompt, num_inference_steps=steps, guidance_scale=guidance_scale).images[0]
    return image

# 構建 Gradio 介面
with gr.Blocks() as demo:
    gr.Markdown("## SANA Image Generator")
    with gr.Row():
        prompt = gr.Textbox(label="Prompt", placeholder="Describe your image here...", lines=2)
    with gr.Row():
        steps = gr.Slider(10, 100, value=50, step=1, label="Steps")
        guidance_scale = gr.Slider(1.0, 20.0, value=7.5, step=0.5, label="Guidance Scale")
    with gr.Row():
        generate_button = gr.Button("Generate")
    with gr.Row():
        output_image = gr.Image(label="Generated Image")

    generate_button.click(
        fn=generate_image,
        inputs=[prompt, steps, guidance_scale],
        outputs=output_image,
    )

# 啟動 Gradio 應用
demo.launch()