Update app.py

app.py

@@ -61,12 +61,12 @@ from transformers import Mask2FormerImageProcessor, Mask2FormerForUniversalSegme
 # Load Mask2Former
 processor = Mask2FormerImageProcessor.from_pretrained("facebook/mask2former-swin-large-ade-semantic")
 model = Mask2FormerForUniversalSegmentation.from_pretrained("facebook/mask2former-swin-large-ade-semantic")
-model = model.cuda()
+model = model.to(DEVICE)
 
 # Load Depth model
 depth_image_processor = AutoImageProcessor.from_pretrained("LiheYoung/depth-anything-large-hf", torch_dtype=torch.float16)
 depth_model = AutoModelForDepthEstimation.from_pretrained("LiheYoung/depth-anything-large-hf", torch_dtype=torch.float16)
-depth_model = depth_model.cuda()
+depth_model = depth_model.to(DEVICE)
 
 def to_rgb(color: str) -> tuple:
     return tuple(int(color[i:i+2], 16) for i in (1, 3, 5))