import gradio as gr
from diffusers import DiffusionPipeline

# 加载模型到 CPU
model_repo_id = "stabilityai/stable-diffusion-3.5-large"
pipe = DiffusionPipeline.from_pretrained(model_repo_id)
pipe.to("cpu")  # 使用 CPU

# 定义图像生成函数，允许自定义宽度和高度
def generate_image(prompt, width=1024, height=768, guidance_scale=7.5, num_inference_steps=40):
    image = pipe(
        prompt=prompt,
        width=width,
        height=height,
        guidance_scale=guidance_scale,
        num_inference_steps=num_inference_steps
    ).images[0]
    return image

# 设置 Gradio 界面
with gr.Blocks() as demo:
    gr.Markdown(" # Stable Diffusion 3.5 - 自定义宽高比例")
    
    with gr.Row():
        prompt = gr.Textbox(label="Prompt", placeholder="Enter your prompt here")

    # 用户自定义宽度和高度
    width = gr.Slider(label="Width", minimum=512, maximum=1024, step=64, value=1024)
    height = gr.Slider(label="Height", minimum=512, maximum=1024, step=64, value=768)
    
    guidance_scale = gr.Slider(label="Guidance Scale", minimum=0.1, maximum=15.0, step=0.1, value=7.5)
    num_inference_steps = gr.Slider(label="Number of Inference Steps", minimum=1, maximum=50, step=1, value=40)
    
    output_image = gr.Image(label="Generated Image")

    # 点击按钮生成图像
    generate_btn = gr.Button("Generate")
    generate_btn.click(generate_image, inputs=[prompt, width, height, guidance_scale, num_inference_steps], outputs=output_image)

# 启动 Gradio 界面
demo.launch()