Update app.py

app.py

@@ -1,75 +1,30 @@
-import os
-from pathlib import Path
-import torch
-import numpy as np
 from PIL import Image
-
-# Gradio for UI
 import gradio as gr
 
-# ControlNet aux for OpenPose
+# Import OpenPose from controlnet-aux
 try:
     from controlnet_aux.open_pose import OpenposeDetector
 except ImportError:
-    from controlnet_aux.openpose import OpenposeDetector  # fallback (old versions)
-
-# DepthAnything now comes as its own package
-from depth_anything.dpt import DepthAnything
-
-
-# --------- Wrappers ---------
-class DepthAnythingDetector:
-    def __init__(self, model_type="vitl", device="cpu"):
-        self.device = device
-        self.model = DepthAnything.from_pretrained(model_type).to(device)
-
-    def __call__(self, image: Image.Image):
-        arr = np.array(image.convert("RGB"))
-        tensor = torch.from_numpy(arr).permute(2, 0, 1).float().unsqueeze(0).to(self.device)
-
-        with torch.no_grad():
-            depth = self.model(tensor)[0]
-
-        depth = depth.cpu().numpy()
-        depth = (depth - depth.min()) / (depth.max() - depth.min()) * 255
-        depth = depth.astype(np.uint8)
-        return Image.fromarray(depth)
-
-
-# --------- Predictor Class ---------
-class Predictor:
-    def __init__(self):
-        self.device = "cpu"
-        self.openpose = OpenposeDetector.from_pretrained("lllyasviel/ControlNet")
-        self.depth = DepthAnythingDetector(model_type="vitl", device=self.device)
-
-    def run(self, image: Image.Image, use_openpose: bool = True) -> Image.Image:
-        if use_openpose:
-            result = self.openpose(image)
-        else:
-            result = self.depth(image)
-
-        if isinstance(result, Image.Image):
-            return result
-        else:
-            return Image.fromarray(result)
-
+    from controlnet_aux.openpose import OpenposeDetector  # fallback for older versions
 
-# --------- Gradio UI ---------
-predictor = Predictor()
+# Load the model once
+openpose = OpenposeDetector.from_pretrained("lllyasviel/ControlNet")
 
-def infer(image, method):
-    use_openpose = method == "OpenPose"
-    return predictor.run(image, use_openpose=use_openpose)
+# Gradio inference function
+def generate_pose(image: Image.Image):
+    image = image.convert("RGB")
+    result = openpose(image)
+    if not isinstance(result, Image.Image):
+        result = Image.fromarray(result)
+    return result
 
+# Gradio UI
 demo = gr.Interface(
-    fn=infer,
-    inputs=[
-        gr.Image(type="pil", label="Upload Image"),
-        gr.Radio(["OpenPose", "DepthAnything"], value="OpenPose", label="Method"),
-    ],
-    outputs=gr.Image(type="pil", label="Result"),
-    title="OpenPose + DepthAnything Demo",
+    fn=generate_pose,
+    inputs=gr.Image(type="pil", label="Upload Image"),
+    outputs=gr.Image(type="pil", label="Pose Output"),
+    title="OpenPose Pose Generator",
+    description="Upload an image and generate a pose map using OpenPose."
 )
 
 if __name__ == "__main__":