Update app.py

app.py

@@ -42,28 +42,47 @@ def chat_logic(message, image_file, history, mode):
     
     # 3. Process Request
     try:
-        response_text, active_module = chimera.process_request(
+        response_data, active_module = chimera.process_request(
             message or "", 
             simple_history, 
             mode, 
             image_file
         )
     except Exception as e:
-        response_text = f"Processing Error: {str(e)}"
+        response_data = f"Processing Error: {str(e)}"
         active_module = "ERR"
     
-    # 4. Format Output
-    final_response = f"**[{active_module} Active]**\n\n{response_text}"
+    # 4. Handle response data (could be text or tuple with image)
+    if isinstance(response_data, tuple):
+        # Image generation response (text, image_path)
+        response_text, image_path = response_data
+        final_response = f"**[{active_module} Active]**\n\n{response_text}"
+    else:
+        # Regular text response
+        response_text = response_data
+        image_path = None
+        final_response = f"**[{active_module} Active]**\n\n{response_text}"
     
     # 5. Create user message
     if image_file:
-        user_msg = f"[Uploaded Image] {message or 'Analyze this image'}"
+        user_msg = f"🖼️ [Image Uploaded]\n\n{message or 'Analyze this image'}"
     else:
         user_msg = message
     
-    # 6. Append to History (Dictionary Format for Gradio 6.x)
+    # 6. Append to History
     history.append({"role": "user", "content": user_msg})
-    history.append({"role": "assistant", "content": final_response})
+    
+    # If there's an image to display, include it in the response
+    if image_path:
+        history.append({
+            "role": "assistant", 
+            "content": {
+                "text": final_response,
+                "files": [image_path]
+            }
+        })
+    else:
+        history.append({"role": "assistant", "content": final_response})
     
     return history, "", None
 
@@ -76,27 +95,44 @@ body {
 .gradio-container {
     font-family: 'IBM Plex Mono', monospace;
 }
+#chatbot {
+    border-radius: 10px;
+}
 """
 
-with gr.Blocks(title="Axon God Mode") as demo:
+with gr.Blocks(title="⚡ AXON: GOD MODE") as demo:
     gr.Markdown("# ⚡ AXON: GOD MODE")
     gr.Markdown("*> Modules: VIM (Vision) | NET (Web) | IGM (Art) | ASM (Code)*")
     
     with gr.Row():
-        chatbot = gr.Chatbot(height=500, elem_id="chatbot")
+        chatbot = gr.Chatbot(
+            height=500, 
+            elem_id="chatbot",
+            avatar_images=(None, "🤖"),
+            show_copy_button=True
+        )
     
     with gr.Row():
         with gr.Column(scale=4):
-            msg = gr.Textbox(placeholder="Ask anything, or upload an image...", show_label=False)
-            btn_upload = gr.Image(type="filepath", label="Upload for Vision (VIM)", height=70)
+            msg = gr.Textbox(
+                placeholder="Ask anything, or upload an image...", 
+                show_label=False,
+                container=False
+            )
+            btn_upload = gr.Image(
+                type="filepath", 
+                label="📸 Upload for Vision (VIM)", 
+                height=100
+            )
         
         with gr.Column(scale=1):
             mode = gr.Dropdown(
                 choices=["Auto", "ASM (Code)", "IGM (Generate Image)", "NET (Search)", "VIM (Vision)"],
                 value="Auto",
-                show_label=False
+                label="Mode",
+                show_label=True
             )
-            submit = gr.Button("🚀 EXECUTE", variant="primary")
+            submit = gr.Button("🚀 EXECUTE", variant="primary", size="lg")
     
     # Event Handlers
     submit.click(