Spaces:
Paused
Paused
| 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="hfd_100epochs.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 (unchanged) | |
| ########################################## | |
| 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 (unchanged except device fix) | |
| ########################################## | |
| 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 (unchanged except device fix) | |
| ########################################## | |
| 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 (unchanged) | |
| ########################################## | |
| 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=1, value=1, 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=100, 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() | |