update

app.py

@@ -13,26 +13,16 @@ DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
 print(f"Using device: {DEVICE}")
 
 # Model cache
-_openpose_detector = None
-_dwpose_detector = None
+_processors = {}
 
 
-def get_openpose_detector():
-    """Get or create OpenPose detector."""
-    global _openpose_detector
-    if _openpose_detector is None:
-        from controlnet_aux import OpenposeDetector
-        _openpose_detector = OpenposeDetector.from_pretrained("lllyasviel/Annotators")
-    return _openpose_detector
-
-
-def get_dwpose_detector():
-    """Get or create DWPose detector."""
-    global _dwpose_detector
-    if _dwpose_detector is None:
-        from controlnet_aux import DWposeDetector
-        _dwpose_detector = DWposeDetector.from_pretrained("yolox_l.onnx", "dw-ll_ucoco_384.onnx")
-    return _dwpose_detector
+def get_processor(processor_id):
+    """Get or create a processor by ID."""
+    global _processors
+    if processor_id not in _processors:
+        from controlnet_aux.processor import Processor
+        _processors[processor_id] = Processor(processor_id)
+    return _processors[processor_id]
 
 
 def detect_pose(image, model_type, detect_hand, detect_face, detect_resolution):
@@ -49,50 +39,46 @@ def detect_pose(image, model_type, detect_hand, detect_face, detect_resolution):
         if image.mode != "RGB":
             image = image.convert("RGB")
         
-        # Process based on model type
+        # Resize to detect_resolution while maintaining aspect ratio
+        original_size = image.size
+        ratio = detect_resolution / max(original_size)
+        new_size = (int(original_size[0] * ratio), int(original_size[1] * ratio))
+        image_resized = image.resize(new_size, Image.Resampling.LANCZOS)
+        
+        # Map model type to processor ID
         if model_type == "DWPose":
-            detector = get_dwpose_detector()
-            result = detector(
-                image,
-                detect_resolution=detect_resolution,
-                image_resolution=detect_resolution,
-                include_hand=detect_hand,
-                include_face=detect_face,
-                include_body=True,
-                output_type="pil"
-            )
+            processor_id = "dwpose"
         elif model_type == "OpenPose (Full)":
-            detector = get_openpose_detector()
-            result = detector(
-                image,
-                detect_resolution=detect_resolution,
-                hand_and_face=True,
-                output_type="pil"
-            )
+            processor_id = "openpose_full"
         elif model_type == "OpenPose (Face Only)":
-            detector = get_openpose_detector()
-            result = detector(
-                image,
-                detect_resolution=detect_resolution,
-                include_body=False,
-                include_hand=False,
-                include_face=True,
-                output_type="pil"
-            )
+            processor_id = "openpose_faceonly"
+        elif model_type == "OpenPose (Hand)":
+            processor_id = "openpose_hand"
         else:
-            # Basic OpenPose
-            detector = get_openpose_detector()
-            result = detector(
-                image,
-                detect_resolution=detect_resolution,
-                hand_and_face=detect_hand and detect_face,
-                output_type="pil"
-            )
+            # Basic OpenPose with options
+            if detect_hand and detect_face:
+                processor_id = "openpose_full"
+            elif detect_face:
+                processor_id = "openpose_face"
+            elif detect_hand:
+                processor_id = "openpose_hand"
+            else:
+                processor_id = "openpose"
+        
+        # Get processor and process
+        processor = get_processor(processor_id)
+        result = processor(image_resized, to_pil=True)
+        
+        # Resize result back to original size if needed
+        if result is not None and result.size != original_size:
+            result = result.resize(original_size, Image.Resampling.LANCZOS)
         
         return result
             
     except Exception as e:
         print(f"Error during processing: {str(e)}")
+        import traceback
+        traceback.print_exc()
         return None
 
 
@@ -118,14 +104,14 @@ with gr.Blocks(
             
             model_type = gr.Dropdown(
                 label="🤖 Model",
-                choices=["DWPose", "OpenPose", "OpenPose (Full)", "OpenPose (Face Only)"],
-                value="DWPose",
-                info="DWPose is recommended for better accuracy"
+                choices=["DWPose", "OpenPose", "OpenPose (Full)", "OpenPose (Face Only)", "OpenPose (Hand)"],
+                value="OpenPose (Full)",
+                info="OpenPose (Full) is recommended for body+face+hand detection"
             )
             
             with gr.Row():
-                detect_hand = gr.Checkbox(label="👆 Detect Hands", value=True)
-                detect_face = gr.Checkbox(label="😊 Detect Face", value=True)
+                detect_hand = gr.Checkbox(label="👆 Detect Hands", value=True, info="For basic OpenPose mode")
+                detect_face = gr.Checkbox(label="😊 Detect Face", value=True, info="For basic OpenPose mode")
             
             detect_resolution = gr.Slider(
                 label="📏 Detection Resolution",
@@ -144,10 +130,14 @@ with gr.Blocks(
     gr.Markdown(
         """
         ### 📌 Tips
-        - **DWPose** is recommended for best accuracy, especially for hands
-        - **OpenPose (Full)** detects body, face, and hands together
+        - **OpenPose (Full)** detects body, face, and hands together - recommended for most use cases
+        - **DWPose** requires additional dependencies and may not work in all environments
         - Higher **Detection Resolution** improves accuracy but increases processing time
         - The output image can be directly used with ControlNet OpenPose models
+        
+        ### ⚠️ Note
+        - Hand/Face checkboxes only apply when using basic "OpenPose" mode
+        - For full functionality, use "OpenPose (Full)" which includes everything
         """
     )