Update app.py

app.py

@@ -71,33 +71,37 @@ async def api_generate(request: Request):
 # Gradio Interface
 def make_me():
     with gr.Row():
-        with gr.Column(scale=4):
+        with gr.Column(scale=3):
             txt_input = gr.Textbox(
                 label='Your prompt:', 
                 lines=4, 
                 container=False, 
                 elem_id="custom_textbox", 
-                placeholder="Prompt"
+                placeholder="Enter your prompt here..."
             )
-        
-        with gr.Column(scale=1):
-            gen_button = gr.Button('Generate image', elem_id="custom_gen_button")
-            stop_button = gr.Button('Stop', variant='secondary', interactive=False, 
-                                elem_id="custom_stop_button")
-            
-            def on_generate_click():
-                return gr.Button('Generate image', elem_id="custom_gen_button"), gr.Button('Stop', variant='secondary', interactive=True, elem_id="custom_stop_button")
             
-            def on_stop_click():
-                return gr.Button('Generate image', elem_id="custom_gen_button"), gr.Button('Stop', variant='secondary', interactive=False, elem_id="custom_stop_button")
+            with gr.Row():
+                gen_button = gr.Button('Generate', variant='primary', elem_id="custom_gen_button")
+                stop_button = gr.Button('Stop', variant='secondary', interactive=False, 
+                                    elem_id="custom_stop_button")
+                
+                def on_generate_click():
+                    return gr.Button('Generating...', interactive=False, elem_id="custom_gen_button"), gr.Button('Stop', variant='secondary', interactive=True, elem_id="custom_stop_button")
+                
+                def on_stop_click():
+                    return gr.Button('Generate', variant='primary', elem_id="custom_gen_button"), gr.Button('Stop', variant='secondary', interactive=False, elem_id="custom_stop_button")
+                
+                gen_button.click(on_generate_click, inputs=None, outputs=[gen_button, stop_button])
+                stop_button.click(on_stop_click, inputs=None, outputs=[gen_button, stop_button])
             
-            gen_button.click(on_generate_click, inputs=None, outputs=[gen_button, stop_button])
-            stop_button.click(on_stop_click, inputs=None, outputs=[gen_button, stop_button])
-    
-    with gr.Row():
-        with gr.Column():
-            model_dropdown = gr.Dropdown(models, label="Select Model", 
-                                      value=models[0] if models else None)
+            model_dropdown = gr.Dropdown(
+                models, 
+                label="Select Model", 
+                value=models[0] if models else None,
+                interactive=True
+            )
+        
+        with gr.Column(scale=2):
             output_image = gr.Image(
                 label="Generated Image", 
                 width=512,
@@ -107,34 +111,30 @@ def make_me():
                 interactive=False
             )
             
-            # JSON output
-            json_output = gr.JSON(label="API Response")
-            
             def generate_wrapper(model_str, prompt):
                 # Translate prompt to English
                 translated_prompt = translator.translate(prompt)
                 image = gen_image(model_str, translated_prompt)
                 if image is None:
-                    return None, {"error": "Generation failed"}
+                    return None
+                return image
                 
-                base64_str = image_to_base64(image)
-                response = {
-                    "status": "success",
-                    "model": model_str,
-                    "original_prompt": prompt,
-                    "translated_prompt": translated_prompt,
-                    "image_base64": base64_str,
-                    "image_format": "jpeg"
-                }
-                return image, response
-            
-            gen_event = gen_button.click(generate_wrapper, [model_dropdown, txt_input], 
-                                     [output_image, json_output])
-            stop_button.click(on_stop_click, inputs=None, 
-                            outputs=[gen_button, stop_button], cancels=[gen_event])
+            gen_event = gen_button.click(
+                generate_wrapper, 
+                [model_dropdown, txt_input], 
+                output_image
+            )
+            stop_button.click(
+                on_stop_click, 
+                inputs=None, 
+                outputs=[gen_button, stop_button], 
+                cancels=[gen_event]
+            )
 
 # Create Gradio app
-with gr.Blocks(css=css_code) as demo: 
+with gr.Blocks(css=css_code, title="Image Generator") as demo: 
+    gr.Markdown("# Image Generation Tool")
+    gr.Markdown("Enter your prompt and select a model to generate an image")
     make_me()
 
 # Enable queue before mounting