s

dkt/pipelines/pipeline.py

@@ -956,26 +956,9 @@ class DKTPipeline:
 
 
 
-    @spaces.GPU(duration=30)
+    @spaces.GPU()
     @torch.inference_mode()
-    def moge_infer(self, input_image):
-        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-        
-        # Ensure model and input are on the same device
-        model_device = next(self.moge_pipe.parameters()).device
-        logger.info(f'moge_infer: model device={model_device}, target device={device}')
-        print(f'moge_infer: model device={model_device}, target device={device}')
-        
-        if model_device != device:
-            logger.info(f'Moving MoGe model from {model_device} to {device}')
-            print(f'Moving MoGe model from {model_device} to {device}')
-            self.moge_pipe = self.moge_pipe.to(device)
-        
-        input_image = input_image.to(device)
-        logger.info(f'moge_infer: input device={input_image.device}')
-        print(f'moge_infer: input device={input_image.device}')
-        
-
+    def moge_infer(self, input_image):        
         return self.moge_pipe.infer(input_image)