display model loading status separately

app.py

@@ -58,8 +58,8 @@ 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, learning_rate, hps_w, imgrw_w, pcks_w, clip_w, progress=gr.Progress(track_tqdm=True)):
-    if prompt is None:
-        raise gr.Error("You forgot the prompt !")
+    if prompt is None or prompt == "":
+        raise gr.Error("You forgot to provide a prompt !")
         
     """Clear CUDA memory before starting the training."""
     torch.cuda.empty_cache()  # Free up cached memory
@@ -86,10 +86,14 @@ def setup_model(prompt, model, seed, num_iterations, learning_rate, hps_w, imgrw
         args.enable_multi_apply= True
         args.multi_step_model = "flux"
 
-    args, trainer, device, dtype, shape, enable_grad, multi_apply_fn, settings = setup(args)
-    loaded_setup = [args, trainer, device, dtype, shape, enable_grad, multi_apply_fn, settings]
-
-    return None, loaded_setup
+    try:
+        args, trainer, device, dtype, shape, enable_grad, multi_apply_fn, settings = setup(args)
+        loaded_setup = [args, trainer, device, dtype, shape, enable_grad, multi_apply_fn, settings]
+        return f"{model} model loaded succesfully !", loaded_setup
+    
+    except Exception as e:
+        print(f"Unexpected Error: {e}")
+        return f"Something went wrong with {model} loading", None
 
 def generate_image(setup_args, num_iterations):
     torch.cuda.empty_cache()  # Free up cached memory
@@ -198,7 +202,19 @@ def show_gallery_output(gallery_state):
 title="# ReNO: Enhancing One-step Text-to-Image Models through Reward-based Noise Optimization"
 description="Enter a prompt to generate an image using ReNO. Adjust the model and parameters as needed."
 
-with gr.Blocks(analytics_enabled=False) as demo:
+css="""
+#model-status-id{
+    height: 126px;
+}
+#model-status-id .progress-text{
+    font-size: 10px!important;
+}
+#model-status-id .progress-level-inner{
+    font-size: 8px!important;
+}
+"""
+
+with gr.Blocks(css=css, analytics_enabled=False) as demo:
     loaded_model_setup = gr.State()
     gallery_state = gr.State()
     with gr.Column():
@@ -221,7 +237,6 @@ with gr.Blocks(analytics_enabled=False) as demo:
                 with gr.Row():
                     chosen_model = gr.Dropdown(["sd-turbo", "sdxl-turbo", "pixart", "hyper-sd", "flux"], label="Model", value="sd-turbo")
                     seed = gr.Number(label="seed", value=0)
-                    model_status = gr.Textbox(label="model status", visible=False)
                 
                 with gr.Row():
                     n_iter = gr.Slider(minimum=10, maximum=100, step=10, value=50, label="Number of Iterations")
@@ -249,6 +264,7 @@ with gr.Blocks(analytics_enabled=False) as demo:
                 )
             
             with gr.Column():
+                model_status = gr.Textbox(label="model status", visible=True, elem_id="model-status-id")
                 output_image = gr.Image(type="filepath", label="Best Generated Image")
                 status = gr.Textbox(label="Status")
                 iter_gallery = gr.Gallery(label="Iterations", columns=4, visible=False)
@@ -260,7 +276,7 @@ with gr.Blocks(analytics_enabled=False) as demo:
     ).then(
         fn = setup_model,
         inputs = [prompt, chosen_model, seed, n_iter, hps_w, imgrw_w, pcks_w, clip_w, learning_rate],
-        outputs = [output_image, loaded_model_setup]  # Load the new setup into the state
+        outputs = [model_status, loaded_model_setup]  # Load the new setup into the state
     ).then(
         fn = generate_image,
         inputs = [loaded_model_setup, n_iter],