EquiFashionModel

app.py

@@ -0,0 +1,263 @@
+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
+)
+
+##########################################
+# Download model files from HF Hub
+##########################################
+
+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")
+
+##########################################
+# Initialize model on GPU once
+##########################################
+
+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 images
+##########################################
+
+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)]
+
+##########################################
+# Utility functions
+##########################################
+
+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")
+
+##########################################
+# HLG PROCESS 
+##########################################
+
+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]
+
+##########################################
+# LLE PROCESS 
+##########################################
+
+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)]
+
+##########################################
+# Send result to attribute editor
+##########################################
+
+def result2input(images):
+    return {"background": images[-1], "layers": None, "composite": None}
+
+##########################################
+# FULL UI
+##########################################
+
+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()