Update app.py

app.py

@@ -1,34 +1,76 @@
-import gradio as gr
+import os
+from pathlib import Path
+import torch
+import numpy as np
 from PIL import Image
-from controlnet_aux.open_pose import OpenposeDetector   # NEW
-from controlnet_aux.depth_anything import DepthAnythingDetector
 
-# Load models
-openpose = OpenposeDetector.from_pretrained("lllyasviel/ControlNet")
-depth = DepthAnythingDetector.from_pretrained("depth-anything/Depth-Anything-V2")
+# Gradio for UI
+import gradio as gr
+
+# ControlNet aux for OpenPose
+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)
 
-def process(image: Image.Image):
-    image = image.convert("RGB")
-    result_openpose = openpose(image)
-    result_depth = depth(image)
 
-    if not isinstance(result_openpose, Image.Image):
-        result_openpose = Image.fromarray(result_openpose)
-    if not isinstance(result_depth, Image.Image):
-        result_depth = Image.fromarray(result_depth)
+# --------- Gradio UI ---------
+predictor = Predictor()
 
-    return result_openpose, result_depth
+def infer(image, method):
+    use_openpose = method == "OpenPose"
+    return predictor.run(image, use_openpose=use_openpose)
 
 demo = gr.Interface(
-    fn=process,
-    inputs=gr.Image(type="pil", label="Upload an image"),
-    outputs=[
-        gr.Image(type="pil", label="OpenPose Result"),
-        gr.Image(type="pil", label="Depth Result"),
+    fn=infer,
+    inputs=[
+        gr.Image(type="pil", label="Upload Image"),
+        gr.Radio(["OpenPose", "DepthAnything"], value="OpenPose", label="Method"),
     ],
-    title="ControlNet Aux Demo",
-    description="Upload an image and run both OpenPose and Depth Anything."
+    outputs=gr.Image(type="pil", label="Result"),
+    title="OpenPose + DepthAnything Demo",
 )
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch(server_name="0.0.0.0", server_port=7860)