Update gradio_sd3.py

gradio_sd3.py

@@ -1,3 +1,4 @@
+import spaces
 import gradio as gr
 import os
 import math
@@ -15,104 +16,102 @@ from src.transformer_sd3_garm import SD3Transformer2DModel as SD3Transformer2DMo
 from src.transformer_sd3_vton import SD3Transformer2DModel as SD3Transformer2DModel_Vton
 import cv2
 import random
+from huggingface_hub import snapshot_download
 
 example_path = os.path.join(os.path.dirname(__file__), 'examples')
 
+fitdit_repo = "BoyuanJiang/FitDiT"
+repo_path = snapshot_download(repo_id=fitdit_repo)
+
+weight_dtype = torch.bfloat16
+device = "cuda"
+transformer_garm = SD3Transformer2DModel_Garm.from_pretrained(os.path.join(repo_path, "transformer_garm"), torch_dtype=weight_dtype)
+transformer_vton = SD3Transformer2DModel_Vton.from_pretrained(os.path.join(repo_path, "transformer_vton"), torch_dtype=weight_dtype)
+pose_guider =  PoseGuider(conditioning_embedding_channels=1536, conditioning_channels=3, block_out_channels=(32, 64, 256, 512))
+pose_guider.load_state_dict(torch.load(os.path.join(repo_path, "pose_guider", "diffusion_pytorch_model.bin")))
+image_encoder_large = CLIPVisionModelWithProjection.from_pretrained("openai/clip-vit-large-patch14", torch_dtype=weight_dtype)
+image_encoder_bigG = CLIPVisionModelWithProjection.from_pretrained("laion/CLIP-ViT-bigG-14-laion2B-39B-b160k", torch_dtype=weight_dtype)
+pose_guider.to(device=device, dtype=weight_dtype)
+image_encoder_large.to(device=device)
+image_encoder_bigG.to(device=device)
+pipeline = StableDiffusion3TryOnPipeline.from_pretrained(repo_path, torch_dtype=weight_dtype, \
+            transformer_garm=transformer_garm, transformer_vton=transformer_vton, pose_guider=pose_guider, \
+            image_encoder_large=image_encoder_large, image_encoder_bigG=image_encoder_bigG)
+pipeline.to(device)
+dwprocessor = DWposeDetector(model_root=repo_path, device=device)
+parsing_model = Parsing(model_root=repo_path, device=device)
+
 
-class FitDiTGenerator:
-    def __init__(self, model_root, offload=False, aggressive_offload=False, device="cuda:0", with_fp16=False):
-        weight_dtype = torch.float16 if with_fp16 else torch.bfloat16
-        transformer_garm = SD3Transformer2DModel_Garm.from_pretrained(os.path.join(model_root, "transformer_garm"), torch_dtype=weight_dtype)
-        transformer_vton = SD3Transformer2DModel_Vton.from_pretrained(os.path.join(model_root, "transformer_vton"), torch_dtype=weight_dtype)
-        pose_guider =  PoseGuider(conditioning_embedding_channels=1536, conditioning_channels=3, block_out_channels=(32, 64, 256, 512))
-        pose_guider.load_state_dict(torch.load(os.path.join(model_root, "pose_guider", "diffusion_pytorch_model.bin")))
-        image_encoder_large = CLIPVisionModelWithProjection.from_pretrained("openai/clip-vit-large-patch14", torch_dtype=weight_dtype)
-        image_encoder_bigG = CLIPVisionModelWithProjection.from_pretrained("laion/CLIP-ViT-bigG-14-laion2B-39B-b160k", torch_dtype=weight_dtype)
-        pose_guider.to(device=device, dtype=weight_dtype)
-        image_encoder_large.to(device=device)
-        image_encoder_bigG.to(device=device)
-        self.pipeline = StableDiffusion3TryOnPipeline.from_pretrained(model_root, torch_dtype=weight_dtype, transformer_garm=transformer_garm, transformer_vton=transformer_vton, pose_guider=pose_guider, image_encoder_large=image_encoder_large, image_encoder_bigG=image_encoder_bigG)
-        self.pipeline.to(device)
-        if offload:
-            self.pipeline.enable_model_cpu_offload()
-            self.dwprocessor = DWposeDetector(model_root=model_root, device='cpu')
-            self.parsing_model = Parsing(model_root=model_root, device='cpu')
-        elif aggressive_offload:
-            self.pipeline.enable_sequential_cpu_offload()
-            self.dwprocessor = DWposeDetector(model_root=model_root, device='cpu')
-            self.parsing_model = Parsing(model_root=model_root, device='cpu')
-        else:
-            self.pipeline.to(device)
-            self.dwprocessor = DWposeDetector(model_root=model_root, device=device)
-            self.parsing_model = Parsing(model_root=model_root, device=device)
-        
-    def generate_mask(self, vton_img, category, offset_top, offset_bottom, offset_left, offset_right):
-        with torch.inference_mode():
-            vton_img = Image.open(vton_img)
-            vton_img_det = resize_image(vton_img)
-            pose_image, keypoints, _, candidate = self.dwprocessor(np.array(vton_img_det)[:,:,::-1])
-            candidate[candidate<0]=0
-            candidate = candidate[0]
 
-            candidate[:, 0]*=vton_img_det.width
-            candidate[:, 1]*=vton_img_det.height
+    
+def generate_mask(vton_img, category, offset_top, offset_bottom, offset_left, offset_right):
+    with torch.inference_mode():
+        vton_img = Image.open(vton_img)
+        vton_img_det = resize_image(vton_img)
+        pose_image, keypoints, _, candidate = dwprocessor(np.array(vton_img_det)[:,:,::-1])
+        candidate[candidate<0]=0
+        candidate = candidate[0]
+
+        candidate[:, 0]*=vton_img_det.width
+        candidate[:, 1]*=vton_img_det.height
 
-            pose_image = pose_image[:,:,::-1] #rgb
-            pose_image = Image.fromarray(pose_image)
-            model_parse, _ = self.parsing_model(vton_img_det)
+        pose_image = pose_image[:,:,::-1] #rgb
+        pose_image = Image.fromarray(pose_image)
+        model_parse, _ = parsing_model(vton_img_det)
 
-            mask, mask_gray = get_mask_location(category, model_parse, \
-                                        candidate, model_parse.width, model_parse.height, \
-                                        offset_top, offset_bottom, offset_left, offset_right)
-            mask = mask.resize(vton_img.size)
-            mask_gray = mask_gray.resize(vton_img.size)
-            mask = mask.convert("L")
-            mask_gray = mask_gray.convert("L")
-            masked_vton_img = Image.composite(mask_gray, vton_img, mask)
+        mask, mask_gray = get_mask_location(category, model_parse, \
+                                    candidate, model_parse.width, model_parse.height, \
+                                    offset_top, offset_bottom, offset_left, offset_right)
+        mask = mask.resize(vton_img.size)
+        mask_gray = mask_gray.resize(vton_img.size)
+        mask = mask.convert("L")
+        mask_gray = mask_gray.convert("L")
+        masked_vton_img = Image.composite(mask_gray, vton_img, mask)
 
-            im = {}
-            im['background'] = np.array(vton_img.convert("RGBA"))
-            im['layers'] = [np.concatenate((np.array(mask_gray.convert("RGB")), np.array(mask)[:,:,np.newaxis]),axis=2)]
-            im['composite'] = np.array(masked_vton_img.convert("RGBA"))
-            
-            return im, pose_image
+        im = {}
+        im['background'] = np.array(vton_img.convert("RGBA"))
+        im['layers'] = [np.concatenate((np.array(mask_gray.convert("RGB")), np.array(mask)[:,:,np.newaxis]),axis=2)]
+        im['composite'] = np.array(masked_vton_img.convert("RGBA"))
+        
+        return im, pose_image
 
-    def process(self, vton_img, garm_img, pre_mask, pose_image, n_steps, image_scale, seed, num_images_per_prompt, resolution):
-        assert resolution in ["768x1024", "1152x1536", "1536x2048"]
-        new_width, new_height = resolution.split("x")
-        new_width = int(new_width)
-        new_height = int(new_height)
-        with torch.inference_mode():
-            garm_img = Image.open(garm_img)
-            vton_img = Image.open(vton_img)
+@spaces.GPU
+def process(vton_img, garm_img, pre_mask, pose_image, n_steps, image_scale, seed, num_images_per_prompt, resolution):
+    assert resolution in ["768x1024", "1152x1536", "1536x2048"]
+    new_width, new_height = resolution.split("x")
+    new_width = int(new_width)
+    new_height = int(new_height)
+    with torch.inference_mode():
+        garm_img = Image.open(garm_img)
+        vton_img = Image.open(vton_img)
 
-            model_image_size = vton_img.size
-            garm_img, _, _ = pad_and_resize(garm_img, new_width=new_width, new_height=new_height)
-            vton_img, pad_w, pad_h = pad_and_resize(vton_img, new_width=new_width, new_height=new_height)
+        model_image_size = vton_img.size
+        garm_img, _, _ = pad_and_resize(garm_img, new_width=new_width, new_height=new_height)
+        vton_img, pad_w, pad_h = pad_and_resize(vton_img, new_width=new_width, new_height=new_height)
 
-            mask = pre_mask["layers"][0][:,:,3]
-            mask = Image.fromarray(mask)
-            mask, _, _ = pad_and_resize(mask, new_width=new_width, new_height=new_height, pad_color=(0,0,0))
-            mask = mask.convert("L")
-            pose_image = Image.fromarray(pose_image)
-            pose_image, _, _ = pad_and_resize(pose_image, new_width=new_width, new_height=new_height, pad_color=(0,0,0))
-            if seed==-1:
-                seed = random.randint(0, 2147483647)
-            res = self.pipeline(
-                height=new_height,
-                width=new_width,
-                guidance_scale=image_scale,
-                num_inference_steps=n_steps,
-                generator=torch.Generator("cpu").manual_seed(seed),
-                cloth_image=garm_img,
-                model_image=vton_img,
-                mask=mask,
-                pose_image=pose_image,
-                num_images_per_prompt=num_images_per_prompt
-            ).images
-            for idx in range(len(res)):
-                res[idx] = unpad_and_resize(res[idx], pad_w, pad_h, model_image_size[0], model_image_size[1])
-            return res
+        mask = pre_mask["layers"][0][:,:,3]
+        mask = Image.fromarray(mask)
+        mask, _, _ = pad_and_resize(mask, new_width=new_width, new_height=new_height, pad_color=(0,0,0))
+        mask = mask.convert("L")
+        pose_image = Image.fromarray(pose_image)
+        pose_image, _, _ = pad_and_resize(pose_image, new_width=new_width, new_height=new_height, pad_color=(0,0,0))
+        if seed==-1:
+            seed = random.randint(0, 2147483647)
+        res = pipeline(
+            height=new_height,
+            width=new_width,
+            guidance_scale=image_scale,
+            num_inference_steps=n_steps,
+            generator=torch.Generator("cpu").manual_seed(seed),
+            cloth_image=garm_img,
+            model_image=vton_img,
+            mask=mask,
+            pose_image=pose_image,
+            num_images_per_prompt=num_images_per_prompt
+        ).images
+        for idx in range(len(res)):
+            res[idx] = unpad_and_resize(res[idx], pad_w, pad_h, model_image_size[0], model_image_size[1])
+        return res
 
 
 def pad_and_resize(im, new_width=768, new_height=1024, pad_color=(255, 255, 255), mode=Image.LANCZOS):
@@ -186,11 +185,10 @@ HEADER = """
 </div>
 <br>
 FitDiT is designed for high-fidelity virtual try-on using Diffusion Transformers (DiT). It can only be used for <b>Non-commercial Use</b>.<br>
-If you like our work, please star <a href="https://github.com/BoyuanJiang/FitDiT" style="color: blue; text-decoration: underline;">our github repository</a>.
+If you like our work, please star <a href="https://github.com/BoyuanJiang/FitDiT" style="color: blue; text-decoration: underline;">our github repository</a>. A <b>ComfyUI version</b> of FitDiT is available <a href="https://github.com/BoyuanJiang/FitDiT/tree/FitDiT-ComfyUI" style="color: blue; text-decoration: underline;">here</a>.
 """
 
-def create_demo(model_path, device, offload, aggressive_offload, with_fp16):
-    generator = FitDiTGenerator(model_path, offload, aggressive_offload, device, with_fp16)
+def create_demo():            
     with gr.Blocks(title="FitDiT") as demo:
         gr.Markdown(HEADER)
         with gr.Row():
@@ -264,7 +262,7 @@ def create_demo(model_path, device, offload, aggressive_offload, with_fp16):
                     inputs=garm_img,
                     examples_per_page=7,
                     examples=[
-                        os.path.join(example_path, 'garment/12.png'),
+                        os.path.join(example_path, 'garment/12.jpg'),
                         os.path.join(example_path, 'garment/0012.jpg'),
                         os.path.join(example_path, 'garment/0047.jpg'),
                         os.path.join(example_path, 'garment/0049.jpg'),
@@ -291,25 +289,17 @@ def create_demo(model_path, device, offload, aggressive_offload, with_fp16):
                     ])
             with gr.Column():
                 category = gr.Dropdown(label="Garment category", choices=["Upper-body", "Lower-body", "Dresses"], value="Upper-body")
-                resolution = gr.Dropdown(label="Try-on resolution", choices=["768x1024", "1152x1536", "1536x2048"], value="1152x1536")
+                resolution = gr.Dropdown(label="Try-on resolution", choices=["768x1024", "1152x1536", "1536x2048"], value="768x1024")
             with gr.Column():
                 run_mask_button = gr.Button(value="Step1: Run Mask")
                 run_button = gr.Button(value="Step2: Run Try-on")
 
         ips1 = [vton_img, category, offset_top, offset_bottom, offset_left, offset_right]
         ips2 = [vton_img, garm_img, masked_vton_img, pose_image, n_steps, image_scale, seed, num_images_per_prompt, resolution]
-        run_mask_button.click(fn=generator.generate_mask, inputs=ips1, outputs=[masked_vton_img, pose_image])
-        run_button.click(fn=generator.process, inputs=ips2, outputs=[result_gallery])
+        run_mask_button.click(fn=generate_mask, inputs=ips1, outputs=[masked_vton_img, pose_image])
+        run_button.click(fn=process, inputs=ips2, outputs=[result_gallery])
     return demo
 
 if __name__ == "__main__":
-    import argparse
-    parser = argparse.ArgumentParser(description="FitDiT")
-    parser.add_argument("--model_path", type=str, default="BoyuanJiang/FitDiT", required=True, help="The path of FitDiT model.")
-    parser.add_argument("--device", type=str, default="cuda:0", help="Device to use")
-    parser.add_argument("--fp16", action="store_true", help="Load model with fp16, default is bf16")
-    parser.add_argument("--offload", action="store_true", help="Offload model to CPU when not in use.")
-    parser.add_argument("--aggressive_offload", action="store_true", help="Offload model more aggressively to CPU when not in use.")
-    args = parser.parse_args()
-    demo = create_demo(args.model_path, args.device, args.offload, args.aggressive_offload, args.fp16)
-    demo.launch(share=True)
+    demo = create_demo()
+    demo.launch()