Update infer.py

infer.py

@@ -10,7 +10,13 @@ import spaces  # Import the spaces module for ZeroGPU
 
 check_min_version('0.28.0.dev0')
 
+# Global variables to store the models
+pipe_g = None
+pipe_d = None
+
+@spaces.GPU
 def load_models(task_name, device):
+    global pipe_g, pipe_d  # Use global variables to store the models
     if task_name == 'depth':
         model_g = 'jingheya/lotus-depth-g-v1-0'
         model_d = 'jingheya/lotus-depth-d-v1-1'
@@ -19,10 +25,20 @@ def load_models(task_name, device):
         model_d = 'jingheya/lotus-normal-d-v1-0'
 
     dtype = torch.float16
-    # Models will be loaded inside the GPU-decorated function
-    return model_g, model_d, dtype
+    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}.")
 
-@spaces.GPU
 def infer_pipe(pipe, image, task_name, seed, device):
     if seed is None:
         generator = None
@@ -67,22 +83,7 @@ def infer_pipe(pipe, image, task_name, seed, device):
 
     return output_color
 
-@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}.")
-
+def lotus(image, task_name, seed, device):
+    global pipe_g, pipe_d  # Access the global models
     output_d = infer_pipe(pipe_d, image, task_name, seed, device)
     return output_d  # Only returning depth outputs for this application