Update tok_demo.py

tok_demo.py

@@ -54,16 +54,17 @@ model = model.to(device)
 model = model.eval()
 model.enable_nest = True
 
-@spaces.GPU(duration=120)
+@spaces.GPU(duration=1200)
 def viz_diff_slots(model, img, nums, cfg=1.0, return_img=False):
     img.to(device)
     n_slots_inf = []
-    for num_slots_to_inference in nums:
+    for i, num_slots_to_inference in enumerate(nums):
         recon_n = model(
             img, sample=True, cfg=cfg,
             inference_with_n_slots=num_slots_to_inference,
         )
         n_slots_inf.append(recon_n)
+        print(f"processing {i} / {len(nums)}")
     return [convert_np(n_slots_inf[i][0]) for i in range(len(n_slots_inf))]
 
 # Removed process_image function as its functionality is now in the update_outputs function