optimize logging

app.py

@@ -25,7 +25,7 @@ def compute_ncut(
     n_neighbors=150,
     min_dist=0.1,
 ):        
-
+    logging_str = ""
     start = time.time()
     eigvecs, eigvals = NCUT(
         num_eig=num_eig,
@@ -34,7 +34,8 @@ def compute_ncut(
         affinity_focal_gamma=affinity_focal_gamma,
         knn=knn_ncut,
     ).fit_transform(features.reshape(-1, features.shape[-1]))
-    print(f"NCUT time: {time.time() - start:.2f}s")
+    # print(f"NCUT time: {time.time() - start:.2f}s")
+    logging_str += f"NCUT time: {time.time() - start:.2f}s\n"
     
     start = time.time()
     if embedding_method == "UMAP":
@@ -44,7 +45,8 @@ def compute_ncut(
             min_dist=min_dist,
             device="cuda" if torch.cuda.is_available() else "cpu",
         )
-        print(f"UMAP time: {time.time() - start:.2f}s")
+        # print(f"UMAP time: {time.time() - start:.2f}s")
+        logging_str += f"UMAP time: {time.time() - start:.2f}s\n"
     elif embedding_method == "t-SNE":    
         X_3d, rgb = rgb_from_tsne_3d(
             eigvecs,
@@ -53,12 +55,13 @@ def compute_ncut(
             knn=knn_tsne,
             device="cuda" if torch.cuda.is_available() else "cpu",
         )
-        print(f"t-SNE time: {time.time() - start:.2f}s")
+        # print(f"t-SNE time: {time.time() - start:.2f}s")
+        logging_str += f"t-SNE time: {time.time() - start:.2f}s\n"
     else:
         raise ValueError(f"Embedding method {embedding_method} not supported.")
     
     rgb = rgb.reshape(features.shape[:3] + (3,))
-    return rgb
+    return rgb, logging_str
 
 
 def dont_use_too_much_green(image_rgb):
@@ -85,7 +88,7 @@ downscaled_outputs = ['./images/ncut_0_small.jpg', './images/ncut_1_small.jpg',
 
 example_items = downscaled_images[:3] + downscaled_outputs[:3]
 
-def main_fn(
+def ncut_run(
     images,
     model_name="SAM(sam_vit_b)",
     layer=-1,
@@ -101,7 +104,7 @@ def main_fn(
     n_neighbors=500,
     min_dist=0.1,
 ):
-    
+    logging_str = ""
     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.")
         gr.Warning("Perplexity/n_neighbors must be less than the number of samples.\n" f"Setting to {num_sample_tsne-1}.")
@@ -117,9 +120,10 @@ def main_fn(
     features = extract_features(
         images, model_name=model_name, node_type=node_type, layer=layer
     )
-    print(f"Feature extraction time (gpu): {time.time() - start:.2f}s")
+    # print(f"Feature extraction time (gpu): {time.time() - start:.2f}s")
+    logging_str += f"Backbone time: {time.time() - start:.2f}s\n"
     
-    rgb = compute_ncut(
+    rgb, _logging_str = compute_ncut(
         features,
         num_eig=num_eig,
         num_sample_ncut=num_sample_ncut,
@@ -132,20 +136,21 @@ def main_fn(
         n_neighbors=n_neighbors,
         min_dist=min_dist,
     )
+    logging_str += _logging_str
     rgb = dont_use_too_much_green(rgb)
-    return to_pil_images(rgb), []
+    return to_pil_images(rgb), [], logging_str
 
 @spaces.GPU(duration=10)
 def quick_run(*args, **kwargs):
-    return main_fn(*args, **kwargs)
+    return ncut_run(*args, **kwargs)
 
 @spaces.GPU(duration=30)
 def long_run(*args, **kwargs):
-    return main_fn(*args, **kwargs)
+    return ncut_run(*args, **kwargs)
 
 @spaces.GPU(duration=120)
 def super_duper_long_run(*args, **kwargs):
-    return main_fn(*args, **kwargs)
+    return ncut_run(*args, **kwargs)
 
 def run_fn(
     images,
@@ -227,6 +232,9 @@ with gr.Blocks() as demo:
                 n_neighbors_slider = gr.Slider(10, 500, step=10, label="n_neighbors (UMAP)", value=150, elem_id="n_neighbors")
                 min_dist_slider = gr.Slider(0.1, 1, step=0.1, label="min_dist (UMAP)", value=0.1, elem_id="min_dist")
 
+            # logging text box
+            logging_text = gr.Textbox("logging output", label="Logging", elem_id="logging", type="text", placeholder="Logging information", default="")
+            
     def load_default_images():
         return default_images, default_outputs, []
 
@@ -236,14 +244,14 @@ with gr.Blocks() as demo:
     load_images_button.click(load_default_images, outputs=[input_gallery, output_gallery, example_gallery])
     clear_images_button.click(empty_input_and_output, outputs=[input_gallery, output_gallery, example_gallery])
     submit_button.click(
-        main_fn,
+        run_fn,
         inputs=[
             input_gallery, model_dropdown, layer_slider, num_eig_slider, node_type_dropdown, 
             affinity_focal_gamma_slider, num_sample_ncut_slider, knn_ncut_slider, 
             embedding_method_dropdown, num_sample_tsne_slider, knn_tsne_slider, 
             perplexity_slider, n_neighbors_slider, min_dist_slider
         ],
-        outputs=[output_gallery, example_gallery]
+        outputs=[output_gallery, example_gallery, logging_text]
     )