Update infer.py

infer.py

@@ -6,6 +6,7 @@ from diffusers.utils import check_min_version
 from pipeline import LotusGPipeline, LotusDPipeline
 from utils.image_utils import colorize_depth_map
 from contextlib import nullcontext
+import spaces  # Import the spaces module for ZeroGPU
 
 check_min_version('0.28.0.dev0')
 
@@ -18,21 +19,10 @@ def load_models(task_name, device):
         model_d = 'jingheya/lotus-normal-d-v1-0'
 
     dtype = torch.float16
-    pipe_g = LotusGPipeline.from_pretrained(
-        model_g,
-        torch_dtype=dtype,
-    )
-    pipe_d = LotusDPipeline.from_pretrained(
-        model_d,
-        torch_dtype=dtype,
-    )
-    pipe_g.to(device)
-    pipe_d.to(device)
-    pipe_g.set_progress_bar_config(disable=True)
-    pipe_d.set_progress_bar_config(disable=True)
-    logging.info(f"Successfully loaded pipelines from {model_g} and {model_d}.")
-    return pipe_g, pipe_d
+    # Models will be loaded inside the GPU-decorated function
+    return model_g, model_d, dtype
 
+@spaces.GPU
 def infer_pipe(pipe, image, task_name, seed, device):
     if seed is None:
         generator = None
@@ -42,7 +32,7 @@ def infer_pipe(pipe, image, task_name, seed, device):
     if torch.backends.mps.is_available():
         autocast_ctx = nullcontext()
     else:
-        autocast_ctx = torch.autocast(pipe.device.type)
+        autocast_ctx = torch.autocast(device_type='cuda', dtype=torch.float16)
 
     with torch.no_grad():
         with autocast_ctx:
@@ -77,6 +67,22 @@ def infer_pipe(pipe, image, task_name, seed, device):
 
     return output_color
 
-def lotus(image, task_name, seed, device, pipe_g, pipe_d):
+@spaces.GPU
+def lotus(image, task_name, seed, device, model_g, model_d, dtype):
+    # Load models inside the GPU-decorated function
+    pipe_g = LotusGPipeline.from_pretrained(
+        model_g,
+        torch_dtype=dtype,
+    )
+    pipe_d = LotusDPipeline.from_pretrained(
+        model_d,
+        torch_dtype=dtype,
+    )
+    pipe_g.to(device)
+    pipe_d.to(device)
+    pipe_g.set_progress_bar_config(disable=True)
+    pipe_d.set_progress_bar_config(disable=True)
+    logging.info(f"Successfully loaded pipelines from {model_g} and {model_d}.")
+
     output_d = infer_pipe(pipe_d, image, task_name, seed, device)
     return output_d  # Only returning depth outputs for this application