Update app.py

app.py

@@ -4,30 +4,37 @@ from utils import process_image, process_video, generate_references
 import tempfile
 import os
 import gc
+from PIL import Image
 
 loader = ModelLoader()
 
-def chat(message, history, file=None):
-    # Security: Basic sanitization (expand as needed)
+def chat(message: str, history, file=None):
     if not message and not file:
         return "Please provide a message or upload a file.", history
 
-    # Process uploaded file
     context = ""
+    file_path = None
+
     if file:
-        if isinstance(file, str):  # Gradio path
-            ext = os.path.splitext(file)[1].lower()
-            if ext in ['.jpg', '.jpeg', '.png', '.webp']:
-                img = Image.open(file)
-                context = f"Image description: {process_image(img)}\n\n"
-            elif ext in ['.mp4', '.avi', '.mov']:
-                context = f"Video analysis:\n{process_video(file)}\n\n"
-            else:
-                context = "File uploaded. Describe its content or ask questions about it.\n\n"
+        file_path = file if isinstance(file, str) else file.name
+        ext = os.path.splitext(file_path)[1].lower()
+        
+        if ext in ['.jpg', '.jpeg', '.png', '.webp', '.gif']:
+            try:
+                img = Image.open(file_path)
+                context = f"**Image Description**: {process_image(img)}\n\n"
+            except:
+                context = "**Image uploaded.** Please ask questions about it.\n\n"
+        elif ext in ['.mp4', '.avi', '.mov', '.webm']:
+            try:
+                context = f"**Video Analysis**:\n{process_video(file_path)}\n\n"
+            except:
+                context = "**Video uploaded.** Summarized from key frames.\n\n"
+        else:
+            context = "**File uploaded.** Describe its content if needed.\n\n"
 
     full_prompt = f"{context}User: {message}\n\nAssistant:"
 
-    # Generate response
     try:
         model, tokenizer = loader.load_text_model()
         inputs = tokenizer(full_prompt, return_tensors="pt").to(loader.device)
@@ -38,83 +45,103 @@ def chat(message, history, file=None):
                 max_new_tokens=512,
                 temperature=0.7,
                 do_sample=True,
-                pad_token_id=tokenizer.eos_token_id
+                pad_token_id=tokenizer.eos_token_id,
+                repetition_penalty=1.1
             )
         
-        response = tokenizer.decode(outputs[0][inputs.input_ids.shape[1]:], skip_special_tokens=True)
+        response = tokenizer.decode(outputs[0][inputs.input_ids.shape[1]:], 
+                                  skip_special_tokens=True)
         
-        # Add references
         refs = generate_references(message[:100])
-        final_response = f"{response.strip()}\n\n**Sources/References:**\n{refs}"
+        final_response = f"{response.strip()}\n\n**Sources / References:**\n{refs}"
         
     except Exception as e:
-        final_response = f"I processed your input: {str(e)[:200]}. Please try again or rephrase."
-    
-    # Privacy: Clear memory aggressively
+        final_response = f"Error processing request: {str(e)[:150]}\nPlease try again."
+
+    # Privacy: Aggressive cleanup
     gc.collect()
     if torch.cuda.is_available():
         torch.cuda.empty_cache()
-    
-    return final_response, history + [[{"text": message, "files": [file] if file else None}, final_response]]
 
-# Gradio UI
-with gr.Blocks(title="Private Personal AI Chatbot", theme=gr.themes.Soft()) as demo:
-    gr.Markdown("# 🛡️ Private Personal AI Chatbot\nFully local/open-source • No data stored • No external APIs")
+    return final_response
+
+
+# ==================== Gradio 6.x UI ====================
+
+with gr.Blocks() as demo:
+    gr.Markdown("# 🛡️ Private Personal AI Chatbot\n**Fully local • No data stored • Open-source**")
     
     chatbot = gr.Chatbot(
-        height=600,
-        show_copy_button=True,
-        avatar_images=(None, "🧠"),
+        height=650,
+        avatar_images=(None, "🧠"),           # Still supported in Gradio 6
+        bubble_full_width=False,
+        show_copy_button=False,               # Will be replaced with buttons=
+        buttons=["copy", "copy_all"]          # New Gradio 6 way
     )
     
     with gr.Row():
         msg = gr.Textbox(
-            placeholder="Ask anything... (text + image/video supported)",
+            placeholder="Type your message here... (supports image/video uploads)",
             scale=4,
-            container=False
+            container=False,
+            autofocus=True
         )
-        submit = gr.Button("Send", variant="primary")
+        submit_btn = gr.Button("Send", variant="primary", scale=1)
     
     with gr.Row():
         file_upload = gr.File(
-            label="Upload Image or Video (max ~400MB)",
-            file_types=["image", "video"],
+            label="📎 Upload Image or Video",
+            file_types=["image", ".mp4", ".mov", ".avi", ".webm"],
             scale=1
         )
         clear_btn = gr.Button("Clear Chat", variant="stop")
-    
+
+    # Examples
     gr.Examples(
         examples=[
-            ["What is in this image?", "example_image.jpg"],
-            ["Summarize this video", "example_video.mp4"],
-            ["Explain quantum computing simply"],
+            ["Explain how transformers work in simple terms"],
+            ["What is visible in this image?", "example.jpg"],
         ],
-        inputs=[msg, file_upload]
+        inputs=[msg, file_upload],
+        label="Try these examples"
     )
+
+    gr.Markdown("**Privacy**: Everything runs in memory. Nothing is logged or stored permanently.")
+
+    # ===================== Event Handlers =====================
     
-    gr.Markdown("**Privacy Note**: Everything runs in-memory. No conversations are saved.")
+    def respond(message, history, file):
+        if not history:
+            history = []
+            
+        bot_response = chat(message, history, file)
+        history.append([message, bot_response])
+        return "", history, None   # Clear message + return history + clear file
 
-    # Chat logic
-    def user_input(message, history, file):
-        if file:
-            # Display file in chat
-            return "", history + [[{"text": message, "files": [file]}, None]]
-        return "", history + [[message, None]]
+    # Submit via button or Enter
+    submit_btn.click(
+        respond,
+        inputs=[msg, chatbot, file_upload],
+        outputs=[msg, chatbot, file_upload]
+    )
 
     msg.submit(
-        user_input, [msg, chatbot, file_upload], [msg, chatbot]
-    ).then(
-        chat, [msg, chatbot, file_upload], [chatbot, chatbot]  # Note: Gradio handles multimodal display
+        respond,
+        inputs=[msg, chatbot, file_upload],
+        outputs=[msg, chatbot, file_upload]
     )
 
-    submit.click(
-        user_input, [msg, chatbot, file_upload], [msg, chatbot]
-    ).then(
-        chat, [msg, chatbot, file_upload], [chatbot, chatbot]
+    clear_btn.click(
+        lambda: [], 
+        outputs=chatbot,
+        queue=False
     )
 
-    clear_btn.click(lambda: None, None, chatbot, queue=False)
 
-# For HF Spaces
 if __name__ == "__main__":
-    demo.queue(max_size=10).launch()
+    demo.launch(
+        server_name="0.0.0.0",
+        server_port=7860,
+        theme=gr.themes.Soft(),           # Moved here for Gradio 6.x
+        share=False
+    )