Upload app.py

app.py

@@ -670,7 +670,7 @@ def check_hallucination(text):
 def image_brain_state(prompt):
     """Generate brain state visualization with adaptive coloring"""
     if not prompt or len(prompt.strip()) == 0:
-        return None, None, "Please enter a prompt."
+        return "<p>Please enter a prompt.</p>", ""
 
     try:
         g = load_guardrail()
@@ -696,13 +696,38 @@ def image_brain_state(prompt):
             img_user = decoder.decode(z_user).squeeze().numpy()
             img_normal = decoder.decode(z_normal).squeeze().numpy()
 
-        # Create comparison image
-        comparison = create_brain_comparison(
+        # Create comparison image and encode as base64
+        img_path = create_brain_comparison(
             img_normal, img_user, ttfs_normal, ttfs_user, deviation)
 
-        # Generate heartbeat audio
+        import base64
+        with open(img_path, 'rb') as f:
+            img_b64 = base64.b64encode(f.read()).decode('utf-8')
+        os.unlink(img_path)
+
+        # Generate heartbeat audio and encode as base64
         audio_mode = 'attack' if is_attack else 'normal'
-        heartbeat = generate_heartbeat_wav(mode=audio_mode, duration=3.0)
+        wav_path = generate_heartbeat_wav(mode=audio_mode, duration=3.0)
+        audio_html = ''
+        if wav_path:
+            with open(wav_path, 'rb') as f:
+                wav_b64 = base64.b64encode(f.read()).decode('utf-8')
+            os.unlink(wav_path)
+            audio_label = '🚨 Arrhythmia Detected' if is_attack else '💚 Steady Heartbeat'
+            audio_html = f'''
+            <div style="margin-top:12px;">
+                <p style="color:{'#ff4444' if is_attack else '#44cc44'};font-weight:bold;">{audio_label}</p>
+                <audio controls style="width:100%;">
+                    <source src="data:audio/wav;base64,{wav_b64}" type="audio/wav">
+                </audio>
+            </div>'''
+
+        # Build HTML output
+        html_output = f'''
+        <div style="background:#0a0a0a;border-radius:12px;padding:16px;text-align:center;">
+            <img src="data:image/png;base64,{img_b64}" style="max-width:100%;border-radius:8px;" />
+            {audio_html}
+        </div>'''
 
         # Summary
         status_emoji = "🚨" if is_attack else "✅"
@@ -717,15 +742,13 @@ def image_brain_state(prompt):
 | **Risk Score** | {user_result['risk_score']:.2f} |
 
 **Brain State Distance (L2):** {np.linalg.norm(user_state - normal_state):.3f}
-
-{'🔊 **Arrhythmia detected** — Listen to the AI heartbeat alarm!' if is_attack else '🔊 **Steady heartbeat** — AI is calm and stable.'}
 """
 
-        return comparison, heartbeat, summary
+        return html_output, summary
 
     except Exception as e:
         import traceback
-        return None, None, f"Error: {str(e)}\n{traceback.format_exc()}"
+        return f"<p style='color:red;'>Error: {str(e)}</p>", f"```\n{traceback.format_exc()}\n```"
 
 
 # ============================================================
@@ -913,10 +936,7 @@ with gr.Blocks(
                     brain_submit = gr.Button("🧠 Visualize Brain State", variant="primary")
             
             with gr.Row():
-                with gr.Column(scale=2):
-                    brain_image = gr.Image(label="Brain State Comparison (Blue=Normal | Orange=Scar | Red=Attack)")
-                with gr.Column(scale=1):
-                    brain_audio = gr.Audio(label="AI Heartbeat 🔊")
+                brain_output_html = gr.HTML(label="Brain State Visualization")
             
             brain_summary = gr.Markdown(label="Analysis")
             
@@ -925,12 +945,12 @@ with gr.Blocks(
             brain_submit.click(
                 fn=image_brain_state,
                 inputs=brain_input,
-                outputs=[brain_image, brain_audio, brain_summary]
+                outputs=[brain_output_html, brain_summary]
             )
             brain_input.submit(
                 fn=image_brain_state,
                 inputs=brain_input,
-                outputs=[brain_image, brain_audio, brain_summary]
+                outputs=[brain_output_html, brain_summary]
             )
     
     gr.Markdown("""