Update app.py

app.py

@@ -1,5 +1,4 @@
 import os
-os.system("pip install opencv-python")
 # os.system("pip uninstall -y gradio")
 # os.system("pip install gradio==3.41.0")
 
@@ -206,7 +205,6 @@ with gr.Blocks() as demo:
                 
                 segment_button  = gr.Button("1.1 Run segmentation")
                 
-                text_button  = gr.Button("Waiting 1.1 to complete",visible = False)
 
 
                 flag = gr.State(False)
@@ -226,7 +224,6 @@ with gr.Blocks() as demo:
         
 
 
-        #canvas.upload(image_change, inputs=[], outputs=[text_button])
                 
                 
     with gr.Tab(label="2 Optimization"):
@@ -238,11 +235,11 @@ with gr.Blocks() as demo:
                 gr.Markdown("""<p style="text-align: center; font-size: 20px">Optimization settings (SD)</p>""")
                 num_tokens = gr.Number(value="5", label="num tokens to represent each object", interactive= True)
                 num_tokens_global = num_tokens
-                embedding_learning_rate = gr.Textbox(value="0.0001", label="Embedding optimization: Learning rate", interactive= True )
-                max_emb_train_steps =  gr.Number(value="200", label="embedding optimization: Training steps", interactive= True )
+                embedding_learning_rate = gr.Textbox(value="0.00005", label="Embedding optimization: Learning rate", interactive= True )
+                max_emb_train_steps =  gr.Number(value="100", label="embedding optimization: Training steps", interactive= True )
                 
-                diffusion_model_learning_rate = gr.Textbox(value="0.00005", label="UNet Optimization: Learning rate", interactive= True )
-                max_diffusion_train_steps = gr.Number(value="200", label="UNet Optimization: Learning rate: Training steps", interactive= True )
+                diffusion_model_learning_rate = gr.Textbox(value="0.00002", label="UNet Optimization: Learning rate", interactive= True )
+                max_diffusion_train_steps = gr.Number(value="100", label="UNet Optimization: Learning rate: Training steps", interactive= True )
                 
                 train_batch_size = gr.Number(value="5", label="Batch size", interactive= True )
                 gradient_accumulation_steps=gr.Number(value="5", label="Gradient accumulation", interactive= True )
@@ -261,7 +258,7 @@ with gr.Blocks() as demo:
                         train_batch_size,
                         gradient_accumulation_steps,
                 ):
-                    if 1:
+                    try:
                         run_optimization = partial(
                             run_main,  
                             mask_np_list=mask_np_list, 
@@ -278,8 +275,8 @@ with gr.Blocks() as demo:
                         run_optimization()
                         print('finish')
                         return "Optimization finished!"
-                    #except:
-                    #    return "CUDA out of memory, use a smaller batch size or try another picture."
+                    except:
+                        return "CUDA out of memory, use a smaller batch size or try another picture."
 
                 
                 def immediate_update():
@@ -378,8 +375,6 @@ with gr.Blocks() as demo:
             outputs=[slider2]
         )
 
-        #text_button.click(load_image_ui, [false] , 
-        #                [image_loaded, segmentation,  mask_np_list, mask_label_list, canvas, slider, slider2] )
 
         segment_button.click(run_segmentation_wrapper, 
                 [canvas] ,