Update app.py

app.py

@@ -60,6 +60,7 @@ def start_over(gallery_state, loaded_model_setup):
     return gallery_state, None, None, gr.update(visible=False), loaded_model_setup
 
 def setup_model(prompt, model, seed, num_iterations, enable_hps, hps_w, enable_imagereward, imgrw_w, enable_pickscore, pcks_w, enable_clip, clip_w, learning_rate, progress=gr.Progress(track_tqdm=True)):
+    gr.Info(f"Loading {model} model ...")
     if prompt is None or prompt == "":
         raise gr.Error("You forgot to provide a prompt !")
         
@@ -113,6 +114,8 @@ def generate_image(setup_args, num_iterations):
     torch.cuda.empty_cache()  # Free up cached memory
     gc.collect()
 
+    gr.Info(f"Executing iterations task ...")
+
     args = setup_args[0]
     trainer = setup_args[1]
     device = setup_args[2]
@@ -201,11 +204,11 @@ def generate_image(setup_args, num_iterations):
 
     except torch.cuda.OutOfMemoryError as e:
         print(f"Global CUDA Out of Memory Error: {e}")
-        yield (None, "CUDA out of memory.", None)
+        yield (None, f"{e}", None)
     except RuntimeError as e:
         if 'out of memory' in str(e):
             print(f"Runtime Error: {e}")
-            yield (None, "CUDA out of memory.", None)
+            yield (None, f"{e}", None)
         else:
             yield (None, f"An error occurred: {str(e)}", None)
     except Exception as e: