File size: 3,496 Bytes
d048c6f | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 | 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)
|