lazy load seg; dynamic duration

app.py

@@ -93,7 +93,7 @@ generated_object_map = {}
 
 # Prepare models
 ## Grounding SAM
-sam2_predictor = SAM2ImagePredictor.from_pretrained("facebook/sam2-hiera-large")
+sam2_predictor = None # SAM2ImagePredictor.from_pretrained("facebook/sam2-hiera-large")
 
 ############## 3D-Fixer model
 model_dir = 'HorizonRobotics/3D-Fixer'
@@ -217,6 +217,10 @@ def run_segmentation(
     rgb_image = image_prompts["image"].convert("RGB")
 
     global work_space
+    global sam2_predictor
+    if sam2_predictor is None:
+        # lazy initialization
+        sam2_predictor = SAM2ImagePredictor.from_pretrained("facebook/sam2-hiera-large")
 
     # pre-process the layers and get the xyxy boxes of each layer
     if len(image_prompts["points"]) == 0:
@@ -399,8 +403,15 @@ def export_scene_glb(trimeshes, work_space, scene_name):
 
     return scene_path
 
+def get_duration(rgb_image, seg_image, seed, randomize_seed,
+                 num_inference_steps, guidance_scale, cfg_interval_start,
+                 cfg_interval_end, t_rescale):
+    instance_labels = np.unique(np.array(seg_image).reshape(-1, 3), axis=0)
+    step_duration = 15.0
+    return instance_labels.shape[0] * step_duration
+
 @torch.no_grad()
-@spaces.GPU(duration=600)
+@spaces.GPU(duration=get_duration)
 def run_generation(
     rgb_image: Any,
     seg_image: Union[str, Image.Image],