Update app.py

app.py

@@ -44,6 +44,18 @@ def load_models():
             device='cuda'
         )
         
+        # CRITICAL FIX: Disable NSFW filter
+        if hasattr(pipeline, 'safety_checker'):
+            pipeline.safety_checker = None
+            print("✅ NSFW filter disabled", file=sys.stderr)
+        
+        # Alternative: Replace safety checker with dummy function
+        if hasattr(pipeline, 'run_safety_checker'):
+            def dummy_safety_checker(images, *args, **kwargs):
+                return images, [False] * len(images)
+            pipeline.run_safety_checker = dummy_safety_checker
+            print("✅ Safety checker replaced with dummy", file=sys.stderr)
+        
         automasker = AutoMasker(
             densepose_ckpt=os.path.join(repo_path, "DensePose"),
             schp_ckpt=os.path.join(repo_path, "SCHP"),
@@ -53,21 +65,20 @@ def load_models():
         print("✅ Models loaded!", file=sys.stderr)
         
     except Exception as e:
-        print(f"❌ Error: {e}", file=sys.stderr)
+        print(f"❌ Error loading models: {e}", file=sys.stderr)
         traceback.print_exc()
         raise
 
 def generate_tryon(person_img, cloth_img):
+    """Generate virtual try-on"""
+    
     print("="*50, file=sys.stderr)
-    print(f"Received - Person: {type(person_img)}, Cloth: {type(cloth_img)}", file=sys.stderr)
+    print(f"API CALL - Person: {type(person_img)}, Cloth: {type(cloth_img)}", file=sys.stderr)
     
     if person_img is None or cloth_img is None:
         raise gr.Error("Both images required!")
     
     try:
-        # person_img and cloth_img are already PIL Images, no conversion needed
-        print("Images received as PIL", file=sys.stderr)
-        
         load_models()
         
         target_height = 1024
@@ -78,6 +89,7 @@ def generate_tryon(person_img, cloth_img):
         mask = automasker(person_img, "upper")['mask']
         gc.collect()
         
+        print("Running pipeline...", file=sys.stderr)
         result = pipeline(
             image=person_img,
             condition_image=cloth_img,
@@ -89,7 +101,7 @@ def generate_tryon(person_img, cloth_img):
             width=target_width
         )[0]
         
-        print("✅ Success!", file=sys.stderr)
+        print(f"✅ Success! Result size: {result.size}", file=sys.stderr)
         return result
         
     except Exception as e:
@@ -97,14 +109,13 @@ 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
 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")
     ],
-    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.",
     api_name="generate_tryon",