Update app.py

app.py

@@ -240,11 +240,7 @@ def get_mask_sam_process(
     print("MODEL LOADED")
 
     # set predictor 
-    if torch.cuda.is_available():
-        inference_state["device"] = 'cuda'
-    #     predictor = build_sam2_video_predictor(model_cfg, sam2_checkpoint)
-    else:
-        inference_state["device"] = 'cpu'
+    
     predictor = build_sam2_video_predictor(model_cfg, sam2_checkpoint, device='cpu')
         
     print("PREDICTOR READY")
@@ -267,6 +263,13 @@ def get_mask_sam_process(
     else:
         inference_state = stored_inference_state
 
+
+    if torch.cuda.is_available():
+        inference_state["device"] = 'cuda'
+    #     predictor = build_sam2_video_predictor(model_cfg, sam2_checkpoint)
+    else:
+        inference_state["device"] = 'cpu'
+        
     # segment and track one object
     # predictor.reset_state(inference_state) # if any previous tracking, reset
     
@@ -329,6 +332,9 @@ def propagate_to_all(video_in, checkpoint, stored_inference_state, stored_frame_
     #### PROPAGATION ####
     sam2_checkpoint, model_cfg = load_model(checkpoint)
     # set predictor 
+   
+    inference_state = stored_inference_state
+
     if torch.cuda.is_available():
         inference_state["device"] = 'cuda'
         predictor = build_sam2_video_predictor(model_cfg, sam2_checkpoint)
@@ -337,7 +343,7 @@ def propagate_to_all(video_in, checkpoint, stored_inference_state, stored_frame_
         predictor = build_sam2_video_predictor(model_cfg, sam2_checkpoint, device='cpu')
     
 
-    inference_state = stored_inference_state
+    
     frame_names = stored_frame_names
     video_dir = video_frames_dir