Update app.py

app.py

@@ -131,12 +131,38 @@ def run_model_gpu(img):
     masks, flows, _ = model.eval(img, channels = [0,0])
     return masks, flows
 
+@spaces.GPU(duration=60)
+def run_model_gpu60(img):
+    masks, flows, _ = model.eval(img, channels = [0,0])
+    return masks, flows
+
+@spaces.GPU(duration=240)
+def run_model_gpu240(img):
+    masks, flows, _ = model.eval(img, channels = [0,0])
+    return masks, flows
+
+@spaces.GPU(duration=1000)
+def run_model_gpu1000(img):
+    masks, flows, _ = model.eval(img, channels = [0,0])
+    return masks, flows
+
 #@spaces.GPU(duration=10)
-def cellpose_segment(img_pil):
+def cellpose_segment(img_pil, resize = 900):
     img_input = imread(img_pil)
     #img_input = np.array(img_pil)
-    img = image_resize(img_input)
-    masks, flows = run_model_gpu(img)
+    img = image_resize(img_input, resize = resize)
+
+    resize = np.max(img.shape)
+    if resize<1000:
+        masks, flows = run_model_gpu(img)
+    elif resize < 5000:
+        masks, flows = run_model_gpu60(img)
+    elif resize < 20000:
+        masks, flows = run_model_gpu240(img)
+    else:
+        raise ValueError("Image size must be less than 20,000")
+    
+        
     #masks, flows, _ = model.eval(img, channels=[0,0])
     flows = flows[0]
     # masks = np.zeros(img.shape[:2])