Update app.py

app.py

@@ -10,7 +10,7 @@ import pandas as pd
 from datetime import datetime
 
 # ----------------------------------------------------------------------
-# Logging setup (MUST be early so logger is defined before any imports that use it)
+# Logging setup
 # ----------------------------------------------------------------------
 logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
 logger = logging.getLogger(__name__)
@@ -110,13 +110,15 @@ from diffusers import StableDiffusionPipeline
 image_pipe = None
 try:
     image_pipe = StableDiffusionPipeline.from_pretrained(
-        "hf-internal-testing/tiny-stable-diffusion-torch"
+        "hf-internal-testing/tiny-stable-diffusion-torch",
+        safety_checker=None
     )
     if not torch.cuda.is_available():
         image_pipe.to("cpu")
     logger.info("Image pipeline loaded.")
 except Exception as e:
     logger.warning(f"Image pipeline load failed (will be disabled): {e}")
+    image_pipe = None
 
 # ----------------------------------------------------------------------
 # Audio transcription (Whisper tiny)
@@ -173,7 +175,6 @@ else:
 
 # ----------------------------------------------------------------------
 # Helper: update risk with feedback (global state – shared across users)
-# For per‑session risk, use gr.State instead of globals.
 # ----------------------------------------------------------------------
 last_task_category = None
 
@@ -193,6 +194,7 @@ async def handle_text(task_type, prompt):
     global last_task_category
     last_task_category = task_type
     try:
+        logger.info(f"Handling text task: {task_type}, prompt: {prompt[:50]}...")
         response, avg_log_prob = generate_with_logprobs(prompt)
         retrieval_score = retriever.get_similarity(prompt)
         event = AIEvent(
@@ -210,7 +212,7 @@ async def handle_text(task_type, prompt):
             prompt=prompt,
             response=response,
             response_length=len(response),
-            confidence=float(np.exp(avg_log_prob)),  # convert to [0,1] scale (approx)
+            confidence=float(np.exp(avg_log_prob)),
             perplexity=None,
             retrieval_scores=[retrieval_score],
             user_feedback=None,
@@ -229,8 +231,8 @@ async def handle_text(task_type, prompt):
             "risk_metrics": risk_metrics
         }
     except Exception as e:
-        logger.error(f"Text task error: {e}")
-        return {"error": str(e)}
+        logger.error(f"Text task error: {e}", exc_info=True)
+        return {"error": str(e), "traceback": traceback.format_exc()}
 
 async def handle_image(prompt, steps):
     """Handle image generation with configurable steps. Returns (image, json_data)."""
@@ -257,9 +259,9 @@ async def handle_image(prompt, steps):
             model_name="tiny-sd",
             model_version="latest",
             prompt=prompt,
-            response="",  # not text
+            response="",
             response_length=0,
-            confidence=1.0 / (gen_time + 1),  # heuristic
+            confidence=1.0 / (gen_time + 1),
             perplexity=None,
             retrieval_scores=[retrieval_score, gen_time],
             user_feedback=None,
@@ -273,8 +275,8 @@ async def handle_image(prompt, steps):
         }
         return image, json_data
     except Exception as e:
-        logger.error(f"Image task error: {e}")
-        return None, {"error": str(e)}
+        logger.error(f"Image task error: {e}", exc_info=True)
+        return None, {"error": str(e), "traceback": traceback.format_exc()}
 
 async def handle_audio(audio_file):
     """Handle audio transcription and quality analysis."""
@@ -282,17 +284,27 @@ async def handle_audio(audio_file):
     last_task_category = "audio"
     if audio_pipe is None:
         return {"error": "Audio model not loaded"}
+    if audio_file is None:
+        return {"error": "No audio file provided"}
     try:
         import librosa
         import soundfile as sf
         import tempfile
+        
+        # Load and process audio
         audio, sr = librosa.load(audio_file, sr=16000)
         with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as tmp:
-            sf.write(tmp.name, audio, sr)
-            result = audio_pipe(tmp.name, return_timestamps=False)
+            tmp_path = tmp.name
+            sf.write(tmp_path, audio, sr)
+        
+        # Transcribe
+        result = audio_pipe(tmp_path, return_timestamps=False)
         text = result["text"]
-        # Whisper does not output log probs easily; use placeholder
-        avg_log_prob = -2.0
+        
+        # Clean up temp file
+        os.unlink(tmp_path)
+        
+        avg_log_prob = -2.0  # Placeholder
         event = AIEvent(
             timestamp=datetime.utcnow(),
             component="audio",
@@ -305,7 +317,7 @@ async def handle_audio(audio_file):
             action_category="audio",
             model_name="whisper-tiny.en",
             model_version="latest",
-            prompt="",  # audio file path
+            prompt="",
             response=text,
             response_length=len(text),
             confidence=float(np.exp(avg_log_prob)),
@@ -322,98 +334,107 @@ async def handle_audio(audio_file):
             "quality_detection": quality_result
         }
     except Exception as e:
-        logger.error(f"Audio task error: {e}")
-        return {"error": str(e)}
+        logger.error(f"Audio task error: {e}", exc_info=True)
+        return {"error": str(e), "traceback": traceback.format_exc()}
 
 async def read_iot_sensors(fault_type, history_state):
     """Read simulated IoT sensors, run diagnostics, predict failure, and return updated plot data."""
     global last_task_category
     last_task_category = "iot"
-    iot_sim.set_fault(fault_type if fault_type != "none" else None)
-    data = iot_sim.read()
-    history_state.append(data)
-    if len(history_state) > 100:
-        history_state.pop(0)
-
-    # Create IoTEvent with valid component name
-    event = IoTEvent(
-        timestamp=datetime.utcnow(),
-        component="robotic-arm",
-        service_mesh="factory",
-        latency_p99=0,
-        error_rate=0.0,
-        throughput=1,
-        cpu_util=None,
-        memory_util=None,
-        temperature=data['temperature'],
-        vibration=data['vibration'],
-        motor_current=data['motor_current'],
-        position_error=data['position_error']
-    )
-    diag_result = await robotics_diagnostician.analyze(event)
-
-    # Simple failure prediction
-    prediction = None
-    if len(history_state) >= 5:
-        temps = [h['temperature'] for h in history_state[-5:]]
-        x = np.arange(len(temps))
-        slope, intercept = np.polyfit(x, temps, 1)
-        next_temp = slope * len(temps) + intercept
-        if slope > 0.1:
-            time_to_threshold = (40.0 - next_temp) / slope if slope > 0 else None
-            prediction = {
-                "predicted_temperature": next_temp,
-                "time_to_overheat_min": time_to_threshold
-            }
-
-    # Prepare temperature history for plotting as DataFrame
-    temp_history = [h['temperature'] for h in history_state[-20:]]
-    df = pd.DataFrame({
-        "index": list(range(len(temp_history))),
-        "temperature": temp_history
-    })
-
-    return data, diag_result, prediction, df, history_state
-
-# ========== NEW: Infrastructure Reliability Handler ==========
+    try:
+        iot_sim.set_fault(fault_type if fault_type != "none" else None)
+        data = iot_sim.read()
+        history_state.append(data)
+        if len(history_state) > 100:
+            history_state.pop(0)
+
+        # Create IoTEvent
+        event = IoTEvent(
+            timestamp=datetime.utcnow(),
+            component="robotic-arm",
+            service_mesh="factory",
+            latency_p99=0,
+            error_rate=0.0,
+            throughput=1,
+            cpu_util=None,
+            memory_util=None,
+            temperature=data['temperature'],
+            vibration=data['vibration'],
+            motor_current=data['motor_current'],
+            position_error=data['position_error']
+        )
+        diag_result = await robotics_diagnostician.analyze(event)
+
+        # Simple failure prediction
+        prediction = None
+        if len(history_state) >= 5:
+            temps = [h['temperature'] for h in history_state[-5:]]
+            x = np.arange(len(temps))
+            slope, intercept = np.polyfit(x, temps, 1)
+            next_temp = slope * len(temps) + intercept
+            if slope > 0.1:
+                time_to_threshold = (40.0 - next_temp) / slope if slope > 0 else None
+                prediction = {
+                    "predicted_temperature": float(next_temp),
+                    "time_to_overheat_min": float(time_to_threshold) if time_to_threshold else None
+                }
+
+        # Prepare temperature history for plotting
+        temp_history = [h['temperature'] for h in history_state[-20:]]
+        df = pd.DataFrame({
+            "index": list(range(len(temp_history))),
+            "temperature": temp_history
+        })
+
+        return data, diag_result, prediction, df, history_state
+    except Exception as e:
+        logger.error(f"IoT task error: {e}", exc_info=True)
+        return {"error": str(e)}, {"error": str(e)}, {"error": str(e)}, pd.DataFrame({"index": [], "temperature": []}), history_state
+
+# ========== Infrastructure Reliability Handler ==========
 async def handle_infra(fault_type, session_state):
     """Run infrastructure reliability analysis."""
     if not INFRA_DEPS_AVAILABLE:
         return {"error": "Infrastructure modules not installed (see logs)"}, session_state
 
-    # Create a new simulator per session (or reuse from state)
-    if "sim" not in session_state or session_state["sim"] is None:
-        session_state["sim"] = InfraSimulator()
-    sim = session_state["sim"]
-
-    # Inject fault
-    sim.set_fault(fault_type if fault_type != "none" else None)
-    components = sim.read_state()
-
-    # Update graph
-    infra_graph.update_from_state(components)
-
-    # Run Bayesian inference (mock for now; in reality would use Pyro)
-    bayesian_risk = {"switch_failure": 0.1, "server_failure": 0.05}
-
-    # Run GNN prediction (mock if PyG not available)
-    predictions = {"at_risk": ["server-1"] if fault_type != "none" else []}
-
-    # Run ProbLog (via python-problog)
-    logic_explanations = "ProbLog output: ..."  # Replace with actual ProbLog call
-
-    # Ontology reasoning
-    ontology_result = ontology.classify("server") if ontology else {"inferred": [], "consistent": True}
-
-    # Combine results
-    output = {
-        "topology": components,
-        "bayesian_risk": bayesian_risk,
-        "gnn_predictions": predictions,
-        "logic_explanations": logic_explanations,
-        "ontology": ontology_result
-    }
-    return output, session_state
+    try:
+        # Create a new simulator per session (or reuse from state)
+        if "sim" not in session_state or session_state["sim"] is None:
+            session_state["sim"] = InfraSimulator()
+        sim = session_state["sim"]
+
+        # Inject fault
+        sim.set_fault(fault_type if fault_type != "none" else None)
+        components = sim.read_state()
+
+        # Update graph
+        if infra_graph:
+            infra_graph.update_from_state(components)
+
+        # Run Bayesian inference (mock for now)
+        bayesian_risk = {"switch_failure": 0.1, "server_failure": 0.05}
+
+        # Run GNN prediction (mock)
+        predictions = {"at_risk": ["server-1"] if fault_type != "none" else []}
+
+        # Run ProbLog (mock)
+        logic_explanations = "ProbLog output: ..."
+
+        # Ontology reasoning
+        ontology_result = ontology.classify("server") if ontology else {"inferred": [], "consistent": True}
+
+        # Combine results
+        output = {
+            "topology": components,
+            "bayesian_risk": bayesian_risk,
+            "gnn_predictions": predictions,
+            "logic_explanations": logic_explanations,
+            "ontology": ontology_result
+        }
+        return output, session_state
+    except Exception as e:
+        logger.error(f"Infra task error: {e}", exc_info=True)
+        return {"error": str(e), "traceback": traceback.format_exc()}, session_state
 
 # ----------------------------------------------------------------------
 # Gradio UI
@@ -432,14 +453,14 @@ with gr.Blocks(title="ARF v4 – AI Reliability Lab", theme="soft") as demo:
         # Tab 2: Image Generation
         with gr.TabItem("Image Generation"):
             img_prompt = gr.Textbox(label="Prompt", value="A cat wearing a hat")
-            img_steps = gr.Slider(1, 10, value=2, step=1, label="Inference Steps (higher = better quality, slower)")
+            img_steps = gr.Slider(1, 10, value=2, step=1, label="Inference Steps")
             img_btn = gr.Button("Generate")
             img_output = gr.Image(label="Generated Image")
             img_json = gr.JSON(label="Analysis")
 
         # Tab 3: Audio Transcription
         with gr.TabItem("Audio Transcription"):
-            gr.Markdown("Click the microphone to record, or upload a file. Try the sample: [Sample Audio](https://huggingface.co/datasets/Narsil/asr_dummy/resolve/main/1.flac)")
+            gr.Markdown("Upload an audio file to transcribe")
             audio_input = gr.Audio(type="filepath", label="Upload audio file")
             audio_btn = gr.Button("Transcribe")
             audio_output = gr.JSON(label="Analysis")
@@ -465,17 +486,16 @@ with gr.Blocks(title="ARF v4 – AI Reliability Lab", theme="soft") as demo:
                 with gr.Column():
                     pred_display = gr.JSON(label="Failure Prediction")
             with gr.Row():
-                with gr.Column(scale=1, min_width=600):
-                    temp_plot = gr.LinePlot(
-                        label="Temperature History (last 20 readings)",
-                        x="index",
-                        y="temperature"
-                    )
+                temp_plot = gr.LinePlot(
+                    label="Temperature History",
+                    x="index",
+                    y="temperature"
+                )
 
-        # ========== NEW: Infrastructure Reliability Tab ==========
+        # Tab 5: Infrastructure Reliability
         with gr.TabItem("Infrastructure Reliability"):
-            gr.Markdown("### Neuro‑Symbolic Infrastructure Monitoring (Bayesian + Graph + Logic)")
-            infra_state = gr.State(value={})  # per‑session state
+            gr.Markdown("### Neuro‑Symbolic Infrastructure Monitoring")
+            infra_state = gr.State(value={})
 
             with gr.Row():
                 with gr.Column():
@@ -488,7 +508,7 @@ with gr.Blocks(title="ARF v4 – AI Reliability Lab", theme="soft") as demo:
                 with gr.Column():
                     infra_output = gr.JSON(label="Analysis Results")
 
-        # Tab 5: Enterprise
+        # Tab 6: Enterprise
         with gr.TabItem("Enterprise"):
             gr.Markdown("""
             ## 🚀 ARF Enterprise – Governed Execution for Autonomous Infrastructure
@@ -506,9 +526,6 @@ with gr.Blocks(title="ARF v4 – AI Reliability Lab", theme="soft") as demo:
             ### Get Started
             - 📅 [Book a Demo](https://calendly.com/petter2025us/30min)
             - 📧 [Contact Sales](mailto:petter2025us@outlook.com)
-            - 📄 [Download Datasheet](#) (coming soon)
-
-            *Already using ARF OSS? Upgrade seamlessly – same core, governed execution.*
             """)
 
     # Feedback row