Update app.py

app.py

@@ -255,19 +255,22 @@ with gr.Blocks() as demo:
                         train_batch_size,
                         gradient_accumulation_steps,
                 ):
-                    run_optimization = partial(
-                        run_main,                 
-                        num_tokens=int(num_tokens),
-                        embedding_learning_rate = float(embedding_learning_rate), 
-                        max_emb_train_steps = int(max_emb_train_steps), 
-                        diffusion_model_learning_rate= float(diffusion_model_learning_rate), 
-                        max_diffusion_train_steps = int(max_diffusion_train_steps),
-                        train_batch_size=int(train_batch_size),
-                        gradient_accumulation_steps=int(gradient_accumulation_steps)
-                    )
-                    run_optimization()
-                    print('finish')
-                    return "Optimization finished!"
+                    try:
+                        run_optimization = partial(
+                            run_main,                 
+                            num_tokens=int(num_tokens),
+                            embedding_learning_rate = float(embedding_learning_rate), 
+                            max_emb_train_steps = int(max_emb_train_steps), 
+                            diffusion_model_learning_rate= float(diffusion_model_learning_rate), 
+                            max_diffusion_train_steps = int(max_diffusion_train_steps),
+                            train_batch_size=int(train_batch_size),
+                            gradient_accumulation_steps=int(gradient_accumulation_steps)
+                        )
+                        run_optimization()
+                        print('finish')
+                        return "Optimization finished!"
+                    except:
+                        return "CUDA out of memory, try another picture."
 
                 
                 def immediate_update():