init

app.py

@@ -75,6 +75,7 @@ def get_segmentation_pipeline(
 
 
 @torch.inference_mode()
+@spaces.GPU
 def segment_image(
         image: Image,
         image_processor: AutoImageProcessor,
@@ -114,11 +115,12 @@ def get_depth_pipeline():
     feature_extractor = AutoImageProcessor.from_pretrained("LiheYoung/depth-anything-large-hf",
                                                            torch_dtype=dtype)
     depth_estimator = AutoModelForDepthEstimation.from_pretrained("LiheYoung/depth-anything-large-hf",
-                                                                  torch_dtype=dtype).to(device)
+                                                                  torch_dtype=dtype)
     return feature_extractor, depth_estimator
 
 
 @torch.inference_mode()
+@spaces.GPU
 def get_depth_image(
         image: Image,
         feature_extractor: AutoImageProcessor,
@@ -209,6 +211,7 @@ class ControlNetDepthDesignModelMulti:
         
         self.seg_image_processor, self.image_segmentor = get_segmentation_pipeline()
         self.depth_feature_extractor, self.depth_estimator = get_depth_pipeline()
+        self.depth_estimator = self.depth_estimator.to(device)
         
     @spaces.GPU
     def generate_design(self, empty_room_image: Image, prompt: str, guidance_scale: int = 10, num_steps: int = 50, strength: float =0.9, img_size: int = 640) -> Image:
@@ -223,6 +226,8 @@ class ControlNetDepthDesignModelMulti:
                            If the size is not the same the submission will fail.
         """
         print(prompt)
+        print(self.depth_estimator.device)
+        print(self.image_segmentor.device)
         flush()
         self.generator = torch.Generator(device=device).manual_seed(self.seed)