import gradio as gr
from diffusers import StableDiffusionImg2ImgPipeline
import torch
from PIL import Image

# 模型名称
MODEL_NAME = "rsshekhawat/Qwen-Edit-3DChibi-LoRA"

# 加载管线（pipeline），使用 float16 提升速度
pipe = StableDiffusionImg2ImgPipeline.from_pretrained(
    MODEL_NAME,
    torch_dtype=torch.float16
)

# 将模型加载到 GPU（可选 "cpu"）
pipe.to("cuda")

# 处理函数
def process(img, prompt, strength):
    if img is None:
        return None
    
    # 转换为 RGB
    img = img.convert("RGB")
    
    # 使用 pipeline 生成图像
    output = pipe(prompt=prompt, init_image=img, strength=strength).images[0]
    
    return output

# Gradio 前端
with gr.Blocks() as demo:
    gr.Markdown("## 🧸 Qwen 3D Chibi Image Edit Demo")

    with gr.Row():
        with gr.Column():
            img_input = gr.Image(label="Upload Image", type="pil")
            prompt = gr.Textbox(label="Prompt", value="cute 3d chibi anime style")
            strength = gr.Slider(0.1, 1.0, value=0.6, label="Edit Strength")
            btn = gr.Button("Generate")
        with gr.Column():
            img_output = gr.Image(label="Output")

    btn.click(process, inputs=[img_input, prompt, strength], outputs=img_output)

# 启动 Gradio App
demo.launch()