Update app.py

app.py

@@ -3,10 +3,10 @@ import gradio as gr
 import torch
 from transformers import pipeline, BitsAndBytesConfig
 
-# 1. Setup Auth
+# 1. Setup Auth (No changes)
 token = os.environ.get("HF_TOKEN")
 
-# 2. Configuration for 4-bit Quantization (Saves RAM)
+# 2. Configuration (No changes)
 bnb_config = BitsAndBytesConfig(
     load_in_4bit=True,
     bnb_4bit_compute_dtype=torch.bfloat16
@@ -14,32 +14,31 @@ bnb_config = BitsAndBytesConfig(
 
 print("Loading 3 models into RAM... please wait about 5 minutes.")
 
-# Load Gemma 3 1B (Uses ~1GB RAM)
+# Load models (No changes)
 pipe_1b = pipeline("text-generation", model="google/gemma-3-1b-it", 
                    device_map="cpu", token=token, model_kwargs={"quantization_config": bnb_config})
 
-# Load Gemma 3 4B (Uses ~3.5GB RAM in 4-bit)
 pipe_4b = pipeline("text-generation", model="google/gemma-3-4b-it", 
                    device_map="cpu", token=token, model_kwargs={"quantization_config": bnb_config})
 
-# Load Phi-4 Mini (Uses ~2.5GB RAM in 4-bit)
 pipe_phi = pipeline("text-generation", model="microsoft/Phi-4-mini-instruct", 
                     device_map="cpu", token=token, model_kwargs={"quantization_config": bnb_config})
 
 def battle(message):
-    results = []
-    # Format message for Instruct models
-    msgs = [{"role": "user", "content": message}]
-    
-    # Generate from all three
-    out_1b = pipe_1b(msgs, max_new_tokens=128)[0]['generated_text'][-1]['content']
-    out_4b = pipe_4b(msgs, max_new_tokens=128)[0]['generated_text'][-1]['content']
-    out_phi = pipe_phi(msgs, max_new_tokens=128)[0]['generated_text'][-1]['content']
-    
-    return out_1b, out_4b, out_phi
+    try:
+        msgs = [{"role": "user", "content": message}]
+        
+        # Defensive extraction of text
+        out_1b = pipe_1b(msgs, max_new_tokens=128)[0]['generated_text'][-1]['content']
+        out_4b = pipe_4b(msgs, max_new_tokens=128)[0]['generated_text'][-1]['content']
+        out_phi = pipe_phi(msgs, max_new_tokens=128)[0]['generated_text'][-1]['content']
+        
+        return out_1b, out_4b, out_phi
+    except Exception as e:
+        return f"Error: {str(e)}", "Error", "Error"
 
-# 3. UI Layout
-with gr.Blocks(theme=gr.themes.Soft()) as demo:
+# 3. UI Layout (UPDATED FOR GRADIO 6.0)
+with gr.Blocks() as demo: # Removed theme=gr.themes.Soft() from here
     gr.Markdown("# ⚔️ Triple Model Local Arena")
     gr.Markdown("Running locally on 16GB CPU RAM via 4-bit Quantization.")
     
@@ -61,4 +60,5 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
     btn.click(battle, inputs=txt, outputs=[res_1, res_2, res_3])
 
 if __name__ == "__main__":
-    demo.launch(server_name="0.0.0.0", server_port=7860)
+    # Theme is now passed here in Gradio 6.0
+    demo.launch(server_name="0.0.0.0", server_port=7860, theme=gr.themes.Soft())