Build and download pipe models without GPU

src/smc/inference.py

@@ -22,7 +22,10 @@ from typing import List
 MIN_GPU_DURATION = 60
 
 
-model_load_lock = threading.Lock()
+pipe_build_lock = threading.Lock()
+pipe_load_lock = threading.Lock()
+reward_model_load_lock = threading.Lock()
+
 
 def build_pipe(device):
     model_path = "Collov-Labs/Monetico"
@@ -38,13 +41,8 @@ def build_pipe(device):
         device=device,
     )
     pipe = Pipeline(vq_model, tokenizer=tokenizer, text_encoder=text_encoder, transformer=model, scheduler=scheduler_new)
-    pipe.to(device)
     return pipe
 
-def build_pipe_threadsafe(device):
-    with model_load_lock:
-        return build_pipe(device)
-
 @dataclass
 class InferenceOutput:
     images: List[Image.Image]
@@ -61,16 +59,26 @@ class PretrainedInferenceConfig:
     steps: int = 48
     num_batches: int = 4
 
-def _get_pretrained_duration(config: PretrainedInferenceConfig, device='cpu') -> int:
+def infer_pretrained(config: PretrainedInferenceConfig, device='cpu'):
+    with pipe_build_lock:
+        pipe = build_pipe(device)
+    return infer_pretrained_with_pipe(config, pipe, device=device)
+
+def _get_pretrained_duration(config: PretrainedInferenceConfig, pipe: Pipeline, device='cpu') -> int:
+    setup_duration = 30.0
     step_duration = 1.0
-    total_duration = math.ceil(step_duration * config.steps)
+    total_duration = math.ceil(setup_duration + step_duration * config.steps)
     return max(total_duration, MIN_GPU_DURATION)
 
 @spaces.GPU(duration=_get_pretrained_duration)
-def infer_pretrained(config: PretrainedInferenceConfig, device='cpu'):
-    pipe = build_pipe_threadsafe(device)
+def infer_pretrained_with_pipe(config: PretrainedInferenceConfig, pipe: Pipeline, device='cpu'):
+    if isinstance(device, str):
+        device = torch.device(device)
+    with pipe_load_lock:
+        pipe = pipe.to(device)
     reward_bias = 5.0
-    reward_fn, reward_name = rewards.ImageReward_Fk_Steering(device=device, bias=reward_bias), "image_reward_plus_5"
+    with reward_model_load_lock:
+        reward_fn, reward_name = rewards.ImageReward_Fk_Steering(device=device, bias=reward_bias), "image_reward_plus_5"
     image_reward_fn = lambda images: reward_fn(
         images, 
         [config.prompt] * len(images)
@@ -142,18 +150,26 @@ def _get_batch_size_based_on_gpu_mem_smc_grad(device, phi):
         batch_p = 1
     return batch_p
 
-def _get_smc_grad_duration(config: SMCGradInferenceConfig, device='cpu') -> int:
+def infer_smc_grad(config: SMCGradInferenceConfig, device='cpu'):
+    with pipe_build_lock:
+        pipe = build_pipe(device)
+    return infer_smc_grad_with_pipe(config, pipe, device=device)
+
+def _get_smc_grad_duration(config: SMCGradInferenceConfig, pipe: Pipeline, device='cpu') -> int:
+    setup_duration = 30.0
     step_duration = 5.0
-    total_duration = math.ceil(step_duration * config.steps)
+    total_duration = math.ceil(setup_duration + step_duration * config.steps)
     return max(total_duration, MIN_GPU_DURATION)
 
 @spaces.GPU(duration=_get_smc_grad_duration)
-def infer_smc_grad(config: SMCGradInferenceConfig, device='cpu'):
+def infer_smc_grad_with_pipe(config: SMCGradInferenceConfig, pipe: Pipeline, device='cpu'):
     if isinstance(device, str):
         device = torch.device(device)
-    pipe = build_pipe_threadsafe(device)
+    with pipe_load_lock:
+        pipe = pipe.to(device)
     reward_bias = 5.0
-    reward_fn, reward_name = rewards.ImageReward_Fk_Steering(device=device, bias=reward_bias), "image_reward_plus_5"
+    with reward_model_load_lock:
+        reward_fn, reward_name = rewards.ImageReward_Fk_Steering(device=device, bias=reward_bias), "image_reward_plus_5"
     image_reward_fn = lambda images: reward_fn(
         images, 
         [config.prompt] * len(images)