import gradio as gr
from PIL import Image
import torch
from diffusers import DiffusionPipeline
from utils import generate_canny

# 加载模型
pipe = DiffusionPipeline.from_pretrained(
    "stabilityai/stable-diffusion-3.5-large-controlnet-canny",
     
    torch_dtype=torch.float16,
    variant="fp16"
).to("cuda")

# 推理函数
def infer(image: Image.Image, prompt: str, neg_prompt: str = "bad face, blurry, text, watermark") -> Image.Image:
    control = generate_canny(image)
    result = pipe(
        prompt=prompt,
        negative_prompt=neg_prompt,
        controlnet_conditioning_image=control,
        num_inference_steps=30,
        guidance_scale=9.0,
    ).images[0]
    return result

# Gradio 界面
demo = gr.Interface(
    fn=infer,
    inputs=[
        gr.Image(type="pil", label="上传原图（自动转边缘图）"),
        gr.Textbox(lines=2, label="Prompt（描述服装/风格）", value="a woman wearing a spring coat, photorealistic, soft lighting, 8k"),
        gr.Textbox(lines=1, label="Negative Prompt", value="bad face, cropped, lowres, extra limbs")
    ],
    outputs=gr.Image(type="pil", label="生成结果"),
    title="👗 ControlNet-Canny 高质量穿搭图生成器（SD3.5）",
    description="上传图像 + 输入Prompt，即可生成结构一致的高清穿搭图（Stable Diffusion 3.5 + ControlNet）",
    examples=[["example.jpg", "a fashion photo of a woman wearing a trench coat, full body, 8k"]],
)

demo.launch()