main

gradio_app.py

@@ -111,10 +111,11 @@ def load_target_model(selected_model):
         lora_model.eval()
 
         logger.info("Models loaded successfully.")
+        return "Models loaded successfully."
 
     except Exception as e:
         logger.error(f"Error loading models: {e}")
-        raise
+        return f"Error loading models: {e}"
 
 # Image pre-processing (resize and padding)
 class ResizeWithPadding:
@@ -296,6 +297,9 @@ with gr.Blocks() as demo:
         # Load Model Button
         load_button = gr.Button("Load Model")
 
+        # Status message box
+        status_box = gr.Textbox(label="Status", placeholder="Model loading status", interactive=False, value="Model not loaded")
+
         # Run Button
         run_button = gr.Button("Generate Image")
 
@@ -303,7 +307,7 @@ with gr.Blocks() as demo:
         result_image = gr.Image(label="Generated Image")
 
     # Load model button action
-    load_button.click(fn=load_target_model, inputs=[recraft_model], outputs=[])
+    load_button.click(fn=load_target_model, inputs=[recraft_model], outputs=[status_box])
 
     # Run Button 
     run_button.click(fn=infer, inputs=[prompt, sample_image, frame_num, seed], outputs=[result_image])