stable garment app

app.py

@@ -1,18 +1,185 @@
+# adapted from https://huggingface.co/spaces/HumanAIGC/OutfitAnyone/blob/main/app.py
+import torch
+import spaces
 import gradio as gr
-from transformers import pipeline
+from PIL import Image
+import numpy as np
+from torchvision import transforms
 
-pipeline = pipeline(task="image-classification", model="julien-c/hotdog-not-hotdog")
+from transformers import CLIPTextModel, CLIPTokenizer
 
-def predict(input_img):
-    predictions = pipeline(input_img)
-    return input_img, {p["label"]: p["score"] for p in predictions} 
+from diffusers import UniPCMultistepScheduler
+from diffusers import AutoencoderKL
+from diffusers import StableDiffusionPipeline
+from diffusers.loaders import LoraLoaderMixin
 
-gradio_app = gr.Interface(
-    predict,
-    inputs=gr.Image(label="Select hot dog candidate", sources=['upload', 'webcam'], type="pil"),
-    outputs=[gr.Image(label="Processed Image"), gr.Label(label="Result", num_top_classes=2)],
-    title="Hot Dog? Or Not?",
-)
+from stablegarment.models import AppearanceEncoderModel,ControlNetModel
+from stablegarment.piplines import StableGarmentPipeline,StableGarmentControlNetPipeline
 
+import os
+from os.path import join as opj
+
+device = "cuda" if torch.cuda.is_available() else "cpu"
+torch_dtype = torch.float16 if "cuda"==device else torch.float32
+height = 512
+width = 384
+
+base_model_path = "SG161222/Realistic_Vision_V4.0_noVAE"
+vae = AutoencoderKL.from_pretrained("stabilityai/sd-vae-ft-mse").to(dtype=torch_dtype,device=device)
+scheduler = UniPCMultistepScheduler.from_pretrained("runwayml/stable-diffusion-v1-5",subfolder="scheduler")
+
+pretrained_garment_encoder_path = "StableGarment_text2img"
+garment_encoder = AppearanceEncoderModel.from_pretrained(pretrained_garment_encoder_path,torch_dtype=torch_dtype,subfolder="garment_encoder")
+garment_encoder = garment_encoder.to(device=device,dtype=torch_dtype)
+
+pipeline_t2i = StableGarmentPipeline.from_pretrained(base_model_path, vae=vae, torch_dtype=torch_dtype, variant="fp16").to(device=device)
+# pipeline = StableDiffusionPipeline.from_pretrained("SG161222/Realistic_Vision_V4.0_noVAE", vae=vae, torch_dtype=torch_dtype, variant="fp16").to(device=device)
+pipeline_t2i.scheduler = scheduler
+
+pipeline_tryon = None
+'''
+# not ready
+pretrained_model_path = "part_module_controlnet_imp2"
+controlnet = ControlNetModel.from_pretrained(pretrained_model_path,subfolder="controlnet")
+text_encoder = CLIPTextModel.from_pretrained(base_model_path, subfolder='text_encoder')
+tokenizer = CLIPTokenizer.from_pretrained(base_model_path, subfolder='tokenizer')
+pipeline_tryon = StableGarmentControlNetPipeline(
+    vae,
+    text_encoder, 
+    tokenizer,
+    pipeline_t2i.unet,
+    controlnet,
+    scheduler,
+).to(device=device,dtype=torch_dtype)
+'''
+
+
+def prepare_controlnet_inputs(agn_mask_list,densepose_list):
+    for i,agn_mask_img in enumerate(agn_mask_list):
+        agn_mask_img = np.array(agn_mask_img.convert("L"))
+        agn_mask_img = np.expand_dims(agn_mask_img, axis=-1)
+        agn_mask_img = (agn_mask_img >= 128).astype(np.float32)  # 0 or 1
+        agn_mask_list[i] = 1. - agn_mask_img
+    densepose_list = [np.array(img)/255. for img in densepose_list]
+    controlnet_inputs = []
+    for mask,pose in zip(agn_mask_list,densepose_list):
+        controlnet_inputs.append(torch.tensor(np.concatenate([mask, pose], axis=-1)).permute(2,0,1))
+    controlnet_inputs = torch.stack(controlnet_inputs)
+    return controlnet_inputs
+
+@spaces.GPU(enable_queue=True)
+def tryon(prompt,init_image,garment_top,garment_down,):
+    basename = os.path.splitext(os.path.basename(init_image))[0]
+    image_agn = Image.open(opj(parse_dir,basename+"_agn.jpg")).resize((width,height))
+    image_agn_mask = Image.open(opj(parse_dir,basename+"_mask.png")).resize((width,height))
+    densepose_image = Image.open(opj(parse_dir,basename+"_densepose.png")).resize((width,height))
+    garment_top = Image.open(garment_top).resize((width,height))
+
+    garment_images = [garment_top,]
+    prompt = [prompt,]
+    cloth_prompt = ["",]
+    controlnet_condition = prepare_controlnet_inputs([image_agn_mask],[densepose_image])
+
+    images = pipeline_tryon(prompt, negative_prompt="",cloth_prompt=cloth_prompt, # negative_cloth_prompt = n_prompt,
+                  height=height,width=width,num_inference_steps=25,guidance_scale=1.5,eta=0.0,
+                  controlnet_condition=controlnet_condition,reference_image=garment_images, 
+                  garment_encoder=garment_encoder,condition_extra=image_agn,
+                  generator=None,).images
+    return images[0]
+
+@spaces.GPU(enable_queue=True)
+def text2image(prompt,init_image,garment_top,garment_down,style_fidelity=1.):
+
+    garment_top = Image.open(garment_top).resize((width,height))
+    garment_top = transforms.CenterCrop((height,width))(transforms.Resize(max(height, width))(garment_top))
+
+    garment_images = [garment_top,]
+    prompt = [prompt,]
+    cloth_prompt = ["",]
+    n_prompt = "nsfw, unsaturated, abnormal, unnatural, artifact"
+    negative_prompt = [n_prompt]
+    images = pipeline_t2i(prompt,negative_prompt=negative_prompt,cloth_prompt=cloth_prompt,height=height,width=width,
+                    num_inference_steps=30,guidance_scale=4,num_images_per_prompt=1,style_fidelity=style_fidelity,
+                    garment_encoder=garment_encoder,garment_image=garment_images,).images
+    return images[0]
+
+# def text2image(prompt,init_image,garment_top,garment_down,):
+#     return pipeline(prompt).images[0]
+
+def infer(prompt,init_image,garment_top,garment_down,t2i_only,style_fidelity):
+    if t2i_only:
+        return text2image(prompt,init_image,garment_top,garment_down,style_fidelity)
+    else:
+        return tryon(prompt,init_image,garment_top,garment_down)
+
+init_state,prompt_state = None,""
+t2i_only_state = True
+def set_mode(t2i_only,person_condition,prompt):
+    global init_state, prompt_state, t2i_only_state
+    t2i_only_state = not t2i_only_state
+    init_state, prompt_state =  person_condition or init_state, prompt_state or prompt
+    if t2i_only:
+        return [gr.Image(sources='clipboard', type="filepath", label="model",value=None, interactive=False),
+                gr.Textbox(placeholder="", label="prompt(for t2i)", value=prompt_state, interactive=True),
+                ]
+    else:
+        return [gr.Image(sources='clipboard', type="filepath", label="model",value=init_state, interactive=False),
+                gr.Textbox(placeholder="", label="prompt(for t2i)", value="", interactive=False),
+                ]
+
+def example_fn(inputs,):
+    if t2i_only_state:
+        return gr.Image(sources='clipboard', type="filepath", label="model", value=None, interactive=False)
+    return gr.Image(sources='clipboard', type="filepath", label="model",value=inputs, interactive=False)
+
+gr.set_static_paths(paths=["assets/images/model"])
+model_dir = opj(os.path.dirname(__file__), "assets/images/model")
+garment_dir = opj(os.path.dirname(__file__), "assets/images/garment")
+parse_dir = opj(os.path.dirname(__file__), "assets/images/image_parse")
+
+model = opj(model_dir, "13987_00.jpg")
+all_person = [opj(model_dir,fname) for fname in os.listdir(model_dir) if fname.endswith(".jpg")]
+with gr.Blocks(css = ".output-image, .input-image, .image-preview {height: 400px !important} ", ) as gradio_app:
+    gr.Markdown("# StableGarment")
+    with gr.Row():
+        with gr.Column():
+            init_image = gr.Image(sources='clipboard', type="filepath", label="model", value=None, interactive=False)
+            example = gr.Examples(inputs=gr.Image(visible=False), #init_image,
+                                  examples_per_page=4,
+                                  examples=all_person,
+                                  run_on_click=True,
+                                  outputs=init_image,
+                                  fn=example_fn,)
+        with gr.Column():
+            with gr.Row():
+                images_top = [opj(garment_dir,fname) for fname in os.listdir(garment_dir) if fname.endswith(".jpg")]
+                garment_top = gr.Image(sources='upload', type="filepath", label="top garment",value=images_top[0]) # ,interactive=False
+                example_top = gr.Examples(inputs=garment_top,
+                                            examples_per_page=4,
+                                            examples=images_top)
+                images_down = []
+                garment_down = gr.Image(sources='upload', type="filepath", label="lower garment",interactive=False, visible=False)
+                example_down = gr.Examples(inputs=garment_down,
+                                            examples_per_page=4,
+                                            examples=images_down)
+            prompt = gr.Textbox(placeholder="", label="prompt(for t2i)",) # interactive=False
+            with gr.Row():
+                t2i_only = gr.Checkbox(label="t2i with garment", info="Only text and garment.", elem_id="t2i_switch", value=True, interactive=False,)
+                run_button = gr.Button(value="Run")
+                style_fidelity = gr.Slider(0, 1, value=1, label="fidelity(for t2i)") # , info=""
+                t2i_only.change(fn=set_mode,inputs=[t2i_only,init_image,prompt],outputs=[init_image,prompt,])
+        with gr.Column():
+            gallery = gr.Image()
+            run_button.click(fn=infer, 
+                            inputs=[
+                                    prompt,
+                                    init_image,
+                                    garment_top,
+                                    garment_down,
+                                    t2i_only,
+                                    style_fidelity,
+                                    ], 
+                            outputs=[gallery],)
+    
 if __name__ == "__main__":
-    gradio_app.launch()
+    gradio_app.launch()

requirements.txt

@@ -1,2 +1,8 @@
 transformers
-torch
+torch
+torchvision
+diffusers
+spaces
+Pillow
+numpy
+git+https://$ACCESS_TOKEN@github.com/logn-2024/StableGarment.git