import os
import gradio as gr
import numpy as np
from PIL import Image
from generate_shadow_main import Shadow_Diffusion
shadow_diffuser = Shadow_Diffusion()
shadow_prompt_dict = {"软阴影": "a product with soft natural shadow, white background",
                      "硬阴影": "a product with hard natural shadow, white background",
                      "悬浮阴影": "a product with floating natural shadow, white background",}
def generate_shadow(image_rgba, shadow_type, padding_rate, denoise_strength, num_inference_steps,
                    controlnet_conditioning_scale, cfg, seed):
    img_pil = Image.fromarray(image_rgba[..., :-1])
    mask_pil = Image.fromarray(np.repeat(image_rgba[..., -1:], 3, -1))
    shadow_prompt = shadow_prompt_dict[shadow_type]
    print(padding_rate, denoise_strength, num_inference_steps, controlnet_conditioning_scale, cfg, seed)
    _, _, _, shadow_result_pil = shadow_diffuser.generate_shadow(img_pil,
                                                                 mask_pil,
                                                                 shadow_prompt,
                                                                 padding_rate=float(padding_rate),
                                                                 denoise_strength=float(denoise_strength),
                                                                 num_inference_steps=int(num_inference_steps),
                                                                 controlnet_conditioning_scale=float(controlnet_conditioning_scale),
                                                                 cfg=float(cfg),
                                                                 seed=int(seed))
    return np.array(shadow_result_pil)


with gr.Blocks() as demo:
    gr.Markdown("# 💡基于Diffusion的白底阴影生成 \n"
                "请确保上传带有透明通道的RGBA图像作为输入")
    with gr.Row():
        with gr.Column():
            rgba = gr.Image(image_mode="RGBA", label="输入商品图(RGBA)")
            gr.Examples(label="示例图片", inputs=[rgba],
                        examples=[os.path.join("./test_images", n) for n in os.listdir("./test_images")])
        with gr.Column():
            shadow_output = gr.Image(image_mode="RGB", label="阴影生成结果")

    with gr.Row():
        shadow_type = gr.Radio(["软阴影", "硬阴影", "悬浮阴影"], value="硬阴影", label="阴影类型")
        generate_btn = gr.Button(value="生成阴影")

    with gr.Accordion(label="其他参数>>", open=False) as sku_accordion:
        padding_rate = gr.Slider(0, 0.99, value=0.4, step=0.1, label="白边填充比例")
        denoise_strength = gr.Slider(0, 0.99, value=1.0, step=0.01, label="去噪程度")
        num_inference_steps = gr.Slider(10, 50, value=20, step=1, label="推理步数")
        controlnet_conditioning_scale = gr.Slider(0, 1.0, value=0.5, step=0.01, label="控制强度(ControlNet)")
        step_num = gr.Slider(1, 50, value=20, step=1, label="推理步数")
        cfg = gr.Slider(0, 20, value=5, step=0.5, label="CFG")
        seed = gr.Slider(-1, 99999999999999, value=42, step=0.01, label="随机种子")

    generate_btn.click(generate_shadow, inputs=[rgba, shadow_type, padding_rate, denoise_strength, num_inference_steps,
                                                controlnet_conditioning_scale, cfg, seed], outputs=shadow_output)

demo.queue().launch(server_name='[::]', share=True)