optimize gpu allocation

app.py

@@ -85,7 +85,6 @@ downscaled_outputs = ['./images/ncut_0_small.jpg', './images/ncut_1_small.jpg',
 
 example_items = downscaled_images[:3] + downscaled_outputs[:3]
 
-@spaces.GPU(duration=30)
 def main_fn(
     images,
     model_name="SAM(sam_vit_b)",
@@ -102,8 +101,6 @@ def main_fn(
     n_neighbors=500,
     min_dist=0.1,
 ):
-    if len(images) is None:
-        return [], example_items
     
     if perplexity >= num_sample_tsne or n_neighbors >= num_sample_tsne:
         # raise gr.Error("Perplexity must be less than the number of samples for t-SNE.")
@@ -138,7 +135,66 @@ def main_fn(
     rgb = dont_use_too_much_green(rgb)
     return to_pil_images(rgb), []
 
+@spaces.GPU(duration=10)
+def quick_run(*args, **kwargs):
+    return main_fn(*args, **kwargs)
+
+@spaces.GPU(duration=30)
+def long_run(*args, **kwargs):
+    return main_fn(*args, **kwargs)
+
+@spaces.GPU(duration=120)
+def super_duper_long_run(*args, **kwargs):
+    return main_fn(*args, **kwargs)
 
+def run_fn(
+    images,
+    model_name="SAM(sam_vit_b)",
+    layer=-1,
+    num_eig=100,
+    node_type="block",
+    affinity_focal_gamma=0.3,
+    num_sample_ncut=10000,
+    knn_ncut=10,
+    embedding_method="UMAP",
+    num_sample_tsne=1000,
+    knn_tsne=10,
+    perplexity=500,
+    n_neighbors=500,
+    min_dist=0.1,
+):
+    if images is None:
+        return [], example_items
+    
+    kwargs = {
+        "images": images,
+        "model_name": model_name,
+        "layer": layer,
+        "num_eig": num_eig,
+        "node_type": node_type,
+        "affinity_focal_gamma": affinity_focal_gamma,
+        "num_sample_ncut": num_sample_ncut,
+        "knn_ncut": knn_ncut,
+        "embedding_method": embedding_method,
+        "num_sample_tsne": num_sample_tsne,
+        "knn_tsne": knn_tsne,
+        "perplexity": perplexity,
+        "n_neighbors": n_neighbors,
+        "min_dist": min_dist,
+    }
+    num_images = len(images)
+    if num_images > 100:
+        return super_duper_long_run(images, **kwargs)
+    if num_images > 20:
+        return long_run(images, **kwargs)
+    if embedding_method == "UMAP":
+        return long_run(images, **kwargs)
+    if perplexity >= 250:
+        return long_run(images, **kwargs)
+    if num_sample_tsne >= 500:
+        return long_run(images, **kwargs)
+    return quick_run(images, **kwargs)
+    
 with gr.Blocks() as demo:
 
     with gr.Row():