Update app.py

app.py

@@ -59,8 +59,8 @@ def parse_args():
         ),
     )
     parser.add_argument(
-        "--repaint", 
-        action="store_true", 
+        "--repaint",
+        action="store_true",
         help="Whether to repaint the result image with the original background."
     )
     parser.add_argument(
@@ -79,11 +79,11 @@ def parse_args():
         choices=["no", "fp16", "bf16"],
         help=(
             "Whether to use mixed precision. Choose between fp16 and bf16 (bfloat16). Bf16 requires PyTorch >="
-            " 1.10.and an Nvidia Ampere GPU.  Default to the value of accelerate config of the current system or the"
+            " 1.10 and an Nvidia Ampere GPU.  Default to the value of accelerate config of the current system or the"
             " flag passed with the `accelerate.launch` command. Use this argument to override the accelerate config."
         ),
     )
-    
+
     args = parser.parse_args()
     env_local_rank = int(os.environ.get("LOCAL_RANK", -1))
     if env_local_rank != -1 and env_local_rank != args.local_rank:
@@ -103,6 +103,7 @@ def image_grid(imgs, rows, cols):
 
 args = parse_args()
 repo_path = snapshot_download(repo_id=args.resume_path)
+
 # Pipeline
 pipeline = CatVTONPipeline(
     base_ckpt=args.base_model_path,
@@ -112,12 +113,18 @@ pipeline = CatVTONPipeline(
     use_tf32=args.allow_tf32,
     device='cuda'
 )
+
 # AutoMasker
-mask_processor = VaeImageProcessor(vae_scale_factor=8, do_normalize=False, do_binarize=True, do_convert_grayscale=True)
+mask_processor = VaeImageProcessor(
+    vae_scale_factor=8,
+    do_normalize=False,
+    do_binarize=True,
+    do_convert_grayscale=True
+)
 automasker = AutoMasker(
     densepose_ckpt=os.path.join(repo_path, "DensePose"),
     schp_ckpt=os.path.join(repo_path, "SCHP"),
-    device='cuda', 
+    device='cuda',
 )
 
 @spaces.GPU(duration=120)
@@ -130,8 +137,11 @@ def submit_function(
     seed,
     show_type
 ):
+    # person_image 객체에서 background와 layers[0]을 분리
     person_image, mask = person_image["background"], person_image["layers"][0]
     mask = Image.open(mask).convert("L")
+
+    # 만약 마스크가 전부 0(검정)이면 None 처리
     if len(np.unique(np.array(mask))) == 1:
         mask = None
     else:
@@ -153,8 +163,8 @@ def submit_function(
     cloth_image = Image.open(cloth_image).convert("RGB")
     person_image = resize_and_crop(person_image, (args.width, args.height))
     cloth_image = resize_and_padding(cloth_image, (args.width, args.height))
-    
-    # Process mask
+
+    # If user didn't draw a mask
     if mask is not None:
         mask = resize_and_crop(mask, (args.width, args.height))
     else:
@@ -173,11 +183,12 @@ def submit_function(
         guidance_scale=guidance_scale,
         generator=generator
     )[0]
-    
-    # Post-process
+
+    # Post-process & Save
     masked_person = vis_mask(person_image, mask)
     save_result_image = image_grid([person_image, masked_person, cloth_image, result_image], 1, 4)
     save_result_image.save(result_save_path)
+
     if show_type == "result only":
         return result_image
     else:
@@ -187,17 +198,18 @@ def submit_function(
             conditions = image_grid([person_image, cloth_image], 2, 1)
         else:
             condition_width = width // 3
-            conditions = image_grid([person_image, masked_person , cloth_image], 3, 1)
+            conditions = image_grid([person_image, masked_person, cloth_image], 3, 1)
+
         conditions = conditions.resize((condition_width, height), Image.NEAREST)
         new_result_image = Image.new("RGB", (width + condition_width + 5, height))
         new_result_image.paste(conditions, (0, 0))
         new_result_image.paste(result_image, (condition_width + 5, 0))
-    return new_result_image
+        return new_result_image
 
 def person_example_fn(image_path):
     return image_path
 
-# Custom CSS for enhanced visual appeal
+# Custom CSS
 css = """
 footer {visibility: hidden}
 
@@ -322,10 +334,11 @@ def app_gradio():
             Transform your look with AI-powered virtual clothing try-on!
             """
         )
-        
+
         with gr.Row():
             with gr.Column(scale=1, min_width=350):
-                with gr.Box():
+                # (수정) 아래 라인을 gr.Box() -> gr.Group()으로 변경
+                with gr.Group():
                     gr.Markdown("### 📸 Upload Images")
                     with gr.Row():
                         image_path = gr.Image(
@@ -371,7 +384,7 @@ def app_gradio():
                     </div>
                     """
                 )
-                
+
                 with gr.Accordion("⚙️ Advanced Settings", open=False):
                     num_inference_steps = gr.Slider(
                         label="Quality Level",
@@ -411,14 +424,13 @@ def app_gradio():
                     elem_classes="image-container"
                 )
                 with gr.Row():
-                    # Photo Examples
                     root_path = "resource/demo/example"
                     with gr.Column():
                         gr.Markdown("#### 👤 Model Examples")
                         men_exm = gr.Examples(
                             examples=[
-                                os.path.join(root_path, "person", "men", _)
-                                for _ in os.listdir(os.path.join(root_path, "person", "men"))
+                                os.path.join(root_path, "person", "men", file)
+                                for file in os.listdir(os.path.join(root_path, "person", "men"))
                             ],
                             examples_per_page=4,
                             inputs=image_path,
@@ -427,8 +439,8 @@ def app_gradio():
                         )
                         women_exm = gr.Examples(
                             examples=[
-                                os.path.join(root_path, "person", "women", _)
-                                for _ in os.listdir(os.path.join(root_path, "person", "women"))
+                                os.path.join(root_path, "person", "women", file)
+                                for file in os.listdir(os.path.join(root_path, "person", "women"))
                             ],
                             examples_per_page=4,
                             inputs=image_path,
@@ -438,13 +450,13 @@ def app_gradio():
                         gr.Markdown(
                             '<div class="info-text">Model examples courtesy of <a href="https://huggingface.co/spaces/levihsu/OOTDiffusion">OOTDiffusion</a> and <a href="https://www.outfitanyone.org">OutfitAnyone</a></div>'
                         )
-                    
+
                     with gr.Column():
                         gr.Markdown("#### 👕 Clothing Examples")
                         condition_upper_exm = gr.Examples(
                             examples=[
-                                os.path.join(root_path, "condition", "upper", _)
-                                for _ in os.listdir(os.path.join(root_path, "condition", "upper"))
+                                os.path.join(root_path, "condition", "upper", file)
+                                for file in os.listdir(os.path.join(root_path, "condition", "upper"))
                             ],
                             examples_per_page=4,
                             inputs=cloth_image,
@@ -453,8 +465,8 @@ def app_gradio():
                         )
                         condition_overall_exm = gr.Examples(
                             examples=[
-                                os.path.join(root_path, "condition", "overall", _)
-                                for _ in os.listdir(os.path.join(root_path, "condition", "overall"))
+                                os.path.join(root_path, "condition", "overall", file)
+                                for file in os.listdir(os.path.join(root_path, "condition", "overall"))
                             ],
                             examples_per_page=4,
                             inputs=cloth_image,
@@ -463,8 +475,8 @@ def app_gradio():
                         )
                         condition_person_exm = gr.Examples(
                             examples=[
-                                os.path.join(root_path, "condition", "person", _)
-                                for _ in os.listdir(os.path.join(root_path, "condition", "person"))
+                                os.path.join(root_path, "condition", "person", file)
+                                for file in os.listdir(os.path.join(root_path, "condition", "person"))
                             ],
                             examples_per_page=4,
                             inputs=cloth_image,
@@ -494,7 +506,7 @@ def app_gradio():
                 ],
                 result_image,
             )
-            
+
         gr.Markdown(
             """
             ### 💡 Tips & Instructions
@@ -507,8 +519,8 @@ def app_gradio():
             For best results, use clear, front-facing images with good lighting.
             """
         )
-        
+
     demo.queue().launch(share=True, show_error=True)
 
 if __name__ == "__main__":
-    app_gradio()
+    app_gradio()