Update app.py

app.py

@@ -142,26 +142,26 @@ def get_depth_image(
     image = Image.fromarray((image * 255.0).clip(0, 255).astype(np.uint8))
     return image
 
-# ---------------------
-# ------- DETR --------
-# ---------------------
-@torch.inference_mode()
-@spaces.GPU
-def run_detr(image:Image, confidence_threshold: float = 0.8):
+# # ---------------------
+# # ------- DETR --------
+# # ---------------------
+# @torch.inference_mode()
+# @spaces.GPU
+# def run_detr(image:Image, confidence_threshold: float = 0.8):
 
-    target = {'image_id': 0, 'annotations': []}
-    encoding = detr_processor(images=image, annotations=target, return_tensors="pt")
-    pixel_values = encoding["pixel_values"].to(device)
-    outputs = detr_model(pixel_values=pixel_values, pixel_mask=None)
+#     target = {'image_id': 0, 'annotations': []}
+#     encoding = detr_processor(images=image, annotations=target, return_tensors="pt")
+#     pixel_values = encoding["pixel_values"].to(device)
+#     outputs = detr_model(pixel_values=pixel_values, pixel_mask=None)
     
-    width, height = image.size
-    postprocessed_outputs = detr_processor.post_process_object_detection(outputs,
-                                                                    target_sizes=[(height, width)],
-                                                                    threshold=confidence_threshold)
-    outputs = postprocessed_outputs[0]
-    # dict{scores, logits, labels, boxes}
-    outputs = outputs['boxes']
-    return outputs
+#     width, height = image.size
+#     postprocessed_outputs = detr_processor.post_process_object_detection(outputs,
+#                                                                     target_sizes=[(height, width)],
+#                                                                     threshold=confidence_threshold)
+#     outputs = postprocessed_outputs[0]
+#     # dict{scores, logits, labels, boxes}
+#     outputs = outputs['boxes']
+#     return outputs
 
 def resize_dimensions(dimensions, target_size):
     """ 
@@ -415,11 +415,11 @@ depth_estimator = depth_estimator.to(device)
 # ---------------------------------------
 # Load Detr Model
 # ---------------------------------------
-detr_model = DetrForObjectDetection.from_pretrained("facebook-detr-50", 
-                                                    # id2label={v:k for k,v in CLASS2ID.items()}, 
-                                                    use_safetensors=True)
-detr_processor = DetrImageProcessor.from_pretrained("facebook-detr-50", 
-                                                    use_safetensors=True)
+# detr_model = DetrForObjectDetection.from_pretrained("facebook-detr-50", 
+#                                                     # id2label={v:k for k,v in CLASS2ID.items()}, 
+#                                                     use_safetensors=True)
+# detr_processor = DetrImageProcessor.from_pretrained("facebook-detr-50", 
+#                                                     use_safetensors=True)
 
 def main():
     model = ControlNetDepthDesignModelMulti()