Update app.py

app.py

@@ -56,9 +56,9 @@ def load_models():
         traceback.print_exc()
         raise
 
-def generate_tryon(person_img, cloth_img):
+def generate_tryon(person_img, cloth_img, garment_category):  # Added parameter
     print("="*50, file=sys.stderr)
-    print(f"Received - Person: {type(person_img)}, Cloth: {type(cloth_img)}", file=sys.stderr)
+    print(f"Received - Person: {type(person_img)}, Cloth: {type(cloth_img)}, Category: {garment_category}", file=sys.stderr)
     
     if person_img is None or cloth_img is None:
         raise gr.Error("Both images required!")
@@ -68,7 +68,6 @@ def generate_tryon(person_img, cloth_img):
         
         load_models()
         
-        # CRITICAL: Clear cache and create fresh copies
         gc.collect()
         if torch.cuda.is_available():
             torch.cuda.empty_cache()
@@ -81,7 +80,8 @@ def generate_tryon(person_img, cloth_img):
         person_img = resize_and_crop(person_img, (target_width, target_height))
         cloth_img = resize_and_padding(cloth_img, (target_width, target_height))
         
-        mask = automasker(person_img, "upper")['mask']
+        # CHANGED: Use selected category instead of hardcoded "upper"
+        mask = automasker(person_img, garment_category)['mask']
         gc.collect()
         if torch.cuda.is_available():
             torch.cuda.empty_cache()
@@ -92,7 +92,7 @@ def generate_tryon(person_img, cloth_img):
             mask=mask,
             num_inference_steps=50,
             guidance_scale=2.5,
-            seed=None,  # CHANGED: Random seed for different results
+            seed=None,
             height=target_height,
             width=target_width
         )[0]
@@ -105,16 +105,22 @@ def generate_tryon(person_img, cloth_img):
         traceback.print_exc()
         raise gr.Error(str(e))
 
-# CRITICAL: Use gr.Interface instead of gr.Blocks for better API support
+# Updated Interface with category dropdown
 demo = gr.Interface(
     fn=generate_tryon,
     inputs=[
-        gr.Image(label="Person Image", type="pil"),  # Changed from "filepath" to "pil"
-        gr.Image(label="Garment Image", type="pil")  # Changed from "filepath" to "pil"
+        gr.Image(label="Person Image", type="pil"),
+        gr.Image(label="Garment Image", type="pil"),
+        gr.Dropdown(
+            choices=["upper", "lower", "overall"],
+            value="upper",
+            label="Garment Category",
+            info="Select: upper (tops), lower (pants/skirts), overall (dresses/full outfits)"
+        )
     ],
-    outputs=gr.Image(label="Result", type="pil"),  # Changed from "filepath" to "pil"
+    outputs=gr.Image(label="Result", type="pil"),
     title="Try-Space Virtual Try-On",
-    description="Upload person and garment images. Processing takes 2-3 minutes on GPU T4.",
+    description="Upload person and garment images. Select category. Processing takes 2-3 minutes on GPU T4.",
     api_name="generate_tryon",
     allow_flagging="never"
 )