|
|
import os |
|
|
import cv2 |
|
|
import einops |
|
|
import gradio as gr |
|
|
import numpy as np |
|
|
import torch |
|
|
import random |
|
|
from huggingface_hub import hf_hub_download |
|
|
|
|
|
from pytorch_lightning import seed_everything |
|
|
from utils.resize import resize_image, HWC3 |
|
|
from cldm.model import create_model, load_state_dict |
|
|
from cldm.ddim_lle import DDIMSampler as DDIMSampler_LLE |
|
|
from cldm.ddim_hlg import DDIMSampler as DDIMSampler_HLG |
|
|
from automation_pose_mask.openpose import OpenposeDetector |
|
|
from automation_pose_mask.auto_mask import MaskDetector |
|
|
from PIL import Image |
|
|
from rembg import remove |
|
|
|
|
|
from utils.config import ( |
|
|
model_yaml, |
|
|
category_dict, |
|
|
attribute_dict |
|
|
) |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
MODEL_REPO = "NguyenDinhHieu/EquiFashionModel" |
|
|
|
|
|
openpose_body_model_path = hf_hub_download(MODEL_REPO, filename="body_pose_model.pth") |
|
|
openpose_hand_model_path = hf_hub_download(MODEL_REPO, filename="hand_pose_model.pth") |
|
|
sam_model_path = hf_hub_download(MODEL_REPO, filename="open_clip_pytorch_model.bin") |
|
|
my_model_path = hf_hub_download(MODEL_REPO, filename="eqf_final.ckpt") |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
device = "cuda" if torch.cuda.is_available() else "cpu" |
|
|
|
|
|
apply_openpose = OpenposeDetector( |
|
|
body_model_path=openpose_body_model_path, |
|
|
hand_model_path=openpose_hand_model_path |
|
|
) |
|
|
|
|
|
apply_mask = MaskDetector(sam_model_path=sam_model_path) |
|
|
|
|
|
model = create_model(model_yaml).to(device) |
|
|
model.load_state_dict(load_state_dict(my_model_path, location=device)) |
|
|
model.eval() |
|
|
|
|
|
hlg_sampler = DDIMSampler_HLG(model) |
|
|
lle_sampler = DDIMSampler_LLE(model) |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
example_path = os.path.join(os.path.dirname(__file__), "preselected_images") |
|
|
example_image_list = [os.path.join(example_path, x) for x in os.listdir(example_path)] |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
def pil_to_binary_mask(pil_image, threshold=0): |
|
|
np_image = np.array(pil_image) |
|
|
grayscale_image = Image.fromarray(np_image).convert("L") |
|
|
binary_mask = (np.array(grayscale_image) > threshold).astype(np.uint8) * 255 |
|
|
return Image.fromarray(binary_mask) |
|
|
|
|
|
def add_white_background(image): |
|
|
image = image.convert("RGBA") |
|
|
white_bg = Image.new("RGBA", image.size, "WHITE") |
|
|
white_bg.paste(image, (0, 0), image) |
|
|
return white_bg.convert("RGB") |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
def hlg_process(hlg_prompt, input_image, category, a_prompt, n_prompt, |
|
|
num_samples, image_resolution, detect_resolution, ddim_steps, |
|
|
guess_mode, strength, scale, seed, eta): |
|
|
|
|
|
with torch.no_grad(): |
|
|
input_image = HWC3(input_image) |
|
|
detected_map, _ = apply_openpose(resize_image(input_image, detect_resolution)) |
|
|
detected_map = HWC3(detected_map) |
|
|
img = resize_image(input_image, image_resolution) |
|
|
H, W, C = img.shape |
|
|
|
|
|
detected_map = cv2.resize(detected_map, (W, H), interpolation=cv2.INTER_NEAREST) |
|
|
|
|
|
control = torch.from_numpy(detected_map).float().to(device) / 255.0 |
|
|
control = torch.stack([control for _ in range(num_samples)], dim=0) |
|
|
control = einops.rearrange(control, 'b h w c -> b c h w') |
|
|
|
|
|
if seed == -1: |
|
|
seed = random.randint(0, 4294967294) |
|
|
seed_everything(seed) |
|
|
|
|
|
cond = { |
|
|
"c_concat": [control], |
|
|
"c_crossattn": [model.get_learned_conditioning([hlg_prompt + ', ' + a_prompt] * num_samples)] |
|
|
} |
|
|
un_cond = { |
|
|
"c_concat": None if guess_mode else [control], |
|
|
"c_crossattn": [model.get_learned_conditioning([n_prompt] * num_samples)] |
|
|
} |
|
|
shape = (4, H // 8, W // 8) |
|
|
|
|
|
model.control_scales = ([strength] * 13) |
|
|
|
|
|
samples, _ = hlg_sampler.sample(ddim_steps, num_samples, shape, cond, |
|
|
verbose=False, eta=eta, |
|
|
unconditional_guidance_scale=scale, |
|
|
unconditional_conditioning=un_cond) |
|
|
|
|
|
x_samples = model.decode_first_stage(samples) |
|
|
x_samples = (einops.rearrange(x_samples, 'b c h w -> b h w c') |
|
|
* 127.5 + 127.5).cpu().numpy().clip(0, 255).astype(np.uint8) |
|
|
|
|
|
results = [Image.fromarray(x_samples[i]) for i in range(num_samples)] |
|
|
return [add_white_background(remove(img)) for img in results] |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
def lle_process(lle_prompt, dict_img_mask, category, a_prompt, n_prompt, |
|
|
num_samples, image_resolution, detect_resolution, ddim_steps, |
|
|
guess_mode, strength, scale, seed, eta, attribute, selection_mode): |
|
|
|
|
|
input_image = dict_img_mask["background"].convert("RGB") |
|
|
input_image = HWC3(np.array(input_image)) |
|
|
|
|
|
detected_map, keypoints = apply_openpose(resize_image(input_image, detect_resolution)) |
|
|
detected_map = HWC3(detected_map) |
|
|
|
|
|
if selection_mode == "Automatically recognize": |
|
|
mask = apply_mask(resize_image(input_image, detect_resolution), keypoints, |
|
|
category=category, attribute=attribute, sam_mode=True) |
|
|
else: |
|
|
mask = pil_to_binary_mask(dict_img_mask['layers'][0].convert("RGB")) |
|
|
|
|
|
if mask is not None: |
|
|
mask = torch.from_numpy(np.array(mask.convert("L"))).float().to(device) / 255.0 |
|
|
mask = mask.unsqueeze(0).unsqueeze(0) |
|
|
|
|
|
img = resize_image(input_image, image_resolution) |
|
|
H, W, C = img.shape |
|
|
|
|
|
init_img = torch.from_numpy(img).float().to(device) / 127.0 - 1.0 |
|
|
init_img = einops.rearrange(init_img, 'h w c -> 1 c h w') |
|
|
init_img = torch.stack([init_img] * num_samples) |
|
|
|
|
|
detected_map = cv2.resize(detected_map, (W, H), interpolation=cv2.INTER_NEAREST) |
|
|
control = torch.from_numpy(detected_map).float().to(device) / 255.0 |
|
|
control = torch.stack([control]*num_samples) |
|
|
control = einops.rearrange(control, 'b h w c -> b c h w') |
|
|
|
|
|
if seed == -1: |
|
|
seed = random.randint(0, 4294967294) |
|
|
seed_everything(seed) |
|
|
|
|
|
cond = {"c_concat": [control], "c_crossattn": [model.get_learned_conditioning([lle_prompt + ', ' + a_prompt] * num_samples)]} |
|
|
un_cond = {"c_concat": None if guess_mode else [control], "c_crossattn": [model.get_learned_conditioning([n_prompt] * num_samples)]} |
|
|
shape = (4, H // 8, W // 8) |
|
|
|
|
|
samples, _ = lle_sampler.sample(ddim_steps, num_samples, shape, cond, |
|
|
verbose=False, eta=eta, |
|
|
unconditional_guidance_scale=scale, |
|
|
unconditional_conditioning=un_cond, |
|
|
init_img=init_img, mask=mask, |
|
|
english_attribute=attribute) |
|
|
|
|
|
x_samples = model.decode_first_stage(samples) |
|
|
x_samples = (einops.rearrange(x_samples, 'b c h w -> b h w c') |
|
|
* 127.5 + 127.5).cpu().numpy().clip(0, 255).astype(np.uint8) |
|
|
|
|
|
return [Image.fromarray(x_samples[i]) for i in range(num_samples)] |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
def result2input(images): |
|
|
return {"background": images[-1], "layers": None, "composite": None} |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
def create_hfddm(): |
|
|
with gr.Blocks().queue() as app: |
|
|
|
|
|
category = gr.Radio(list(category_dict.values()), value=list(category_dict.values())[0], label="Clothing Category") |
|
|
|
|
|
with gr.Row(): |
|
|
with gr.Column(): |
|
|
with gr.Tab("Draft Design"): |
|
|
hlg_prompt = gr.Textbox(label="High-level design prompt") |
|
|
hlg_input_image = gr.Image(sources=("upload", "webcam"), type="numpy", value=example_image_list[0], label="Reference pose") |
|
|
gr.Examples(inputs=hlg_input_image, examples=example_image_list) |
|
|
hlg_run = gr.Button("Generate") |
|
|
|
|
|
with gr.Tab("Attribute Editing"): |
|
|
lle_prompt = gr.Textbox(label="Attribute prompt") |
|
|
lle_input_image = gr.ImageEditor(sources='upload', type="pil", label="Edit regions", value=example_image_list[0]) |
|
|
gr.Examples(inputs=lle_input_image, examples=example_image_list) |
|
|
selection_mode = gr.Radio(["Automatically recognize", "User interface"], label="Mask Selection", value="Automatically recognize") |
|
|
|
|
|
current_tab = {} |
|
|
lle_run = {} |
|
|
for tab_elem in attribute_dict.values(): |
|
|
with gr.Tab(tab_elem): |
|
|
current_tab[tab_elem] = gr.Label(value=tab_elem, visible=False) |
|
|
lle_run[tab_elem] = gr.Button("Generate") |
|
|
|
|
|
with gr.Column(): |
|
|
result_gallery = gr.Gallery(label="Result", show_label=False, elem_id="gallery", selected_index=0, interactive=False) |
|
|
send2llg = gr.Button("Send to Attribute Editing") |
|
|
|
|
|
with gr.Accordion("Advanced Options", open=False): |
|
|
num_samples = gr.Slider(label="Images", minimum=1, maximum=4, value=2, step=1) |
|
|
image_resolution = gr.Slider(label="Resolution", minimum=256, maximum=768, value=512, step=64) |
|
|
strength = gr.Slider(label="Control Strength", minimum=0.0, maximum=2.0, value=1.0, step=0.01) |
|
|
guess_mode = gr.Checkbox(label='Guess Mode', value=False) |
|
|
detect_resolution = gr.Slider(label="Pose Detection Resolution", minimum=128, maximum=1024, value=512, step=1) |
|
|
ddim_steps = gr.Slider(label="Steps", minimum=1, maximum=100, value=10, step=1, visible=False) |
|
|
scale = gr.Slider(label="Guidance Scale", minimum=0.1, maximum=30.0, value=9.0, step=0.1) |
|
|
seed = gr.Slider(label="Seed", minimum=-1, maximum=4294967294, value=11, step=1) |
|
|
eta = gr.Number(label="ETA (DDIM)", value=0.0) |
|
|
a_prompt = gr.Textbox(label="Added Prompt", value='best quality, extremely detailed, masterpiece, 8k, white background') |
|
|
n_prompt = gr.Textbox(label="Negative Prompt", value='worst quality, low quality, bad anatomy, watermark, signature, blurry') |
|
|
|
|
|
hlg_run.click(fn=hlg_process, inputs=[hlg_prompt, hlg_input_image, category, a_prompt, n_prompt, |
|
|
num_samples, image_resolution, detect_resolution, ddim_steps, |
|
|
guess_mode, strength, scale, seed, eta], outputs=[result_gallery]) |
|
|
|
|
|
for tab_elem in attribute_dict.values(): |
|
|
lle_run[tab_elem].click(fn=lle_process, inputs=[lle_prompt, lle_input_image, category, a_prompt, n_prompt, |
|
|
num_samples, image_resolution, detect_resolution, |
|
|
ddim_steps, guess_mode, strength, scale, seed, eta, |
|
|
current_tab[tab_elem], selection_mode], outputs=[result_gallery]) |
|
|
|
|
|
send2llg.click(fn=result2input, inputs=result_gallery, outputs=lle_input_image) |
|
|
|
|
|
return app |
|
|
|
|
|
hfddm_block = create_hfddm() |
|
|
|
|
|
demo = gr.Blocks(title="AI Fashion Design", theme=gr.themes.Monochrome(secondary_hue="orange", neutral_hue="gray")).queue() |
|
|
|
|
|
with demo: |
|
|
gr.Markdown("# **AI Fashion Design** 👗") |
|
|
with gr.Tab("Fashion Design"): |
|
|
hfddm_block.render() |
|
|
|
|
|
demo.launch() |
|
|
|
