Update app.py

app.py

@@ -1,48 +1,224 @@
 import gradio as gr
 import spaces
 
+import requests
+import torch
+from PIL import Image
+from transformers import MllamaForConditionalGeneration, AutoProcessor
 
-@spaces.GPU
-def diagnose_router(image):
-    result = {
-        "Issue_Description": "Internet connectivity problem - The Internet LED is blinking amber.",
-        "Root_Cause_Analysis": {
-            "LED_Analysis": {
-                "Color": "Amber",
-                "Pattern": "Blinking",
-                "Indicates": "The router is unable to establish a connection with the ISP."
+SCHEMA_DEFINITION= "{
+  "$schema": "http://json-schema.org/draft-04/schema#",
+  "type": "object",
+  "properties": {
+    "Issue_Description": {
+      "type": "string"
+    },
+    "Root_Cause_Analysis": {
+      "type": "object",
+      "properties": {
+        "LED_Analysis": {
+          "type": "object",
+          "properties": {
+            "Color": {
+              "type": "string"
+            },
+            "Pattern": {
+              "type": "string"
+            },
+            "Indicates": {
+              "type": "string"
+            }
+          },
+          "required": [
+            "Color",
+            "Pattern",
+            "Indicates"
+          ]
+        },
+        "Error_Code": {
+          "type": "string"
+        },
+        "Possible_Cause": {
+          "type": "string"
+        }
+      },
+      "required": [
+        "LED_Analysis",
+        "Error_Code",
+        "Possible_Cause"
+      ]
+    },
+    "Step_by_Step_Troubleshooting": {
+      "type": "array",
+      "items": [
+        {
+          "type": "object",
+          "properties": {
+            "Action": {
+              "type": "string"
+            },
+            "Details": {
+              "type": "string"
+            },
+            "Expected Outcome": {
+              "type": "string"
+            }
+          },
+          "required": [
+            "Action",
+            "Details",
+            "Expected Outcome"
+          ]
+        },
+        {
+          "type": "object",
+          "properties": {
+            "Action": {
+              "type": "string"
+            },
+            "Details": {
+              "type": "string"
             },
-            "Error_Code": "WAN PPPoE Timeout",
-            "Possible_Cause": "This error suggests the router cannot authenticate with the ISP's server, likely due to incorrect login credentials or an ISP-side issue."
+            "Expected Outcome": {
+              "type": "string"
+            }
+          },
+          "required": [
+            "Action",
+            "Details",
+            "Expected Outcome"
+          ]
         },
-        "Step_by_Step_Troubleshooting": [
-            {
-                "Action": "Verify Ethernet Cable Connection",
-                "Details": "Ensure that the Ethernet cable is securely plugged into the modem's LAN port and the router's WAN port.",
-                "Expected Outcome": "Ethernet cable is securely connected, and the LED light remains stable."
+        {
+          "type": "object",
+          "properties": {
+            "Action": {
+              "type": "string"
+            },
+            "Details": {
+              "type": "string"
             },
-            {
-                "Action": "Restart Both Modem and Router",
-                "Details": "Unplug both the modem and router from their power sources, wait for 30 seconds, then plug them back in starting with the modem.",
-                "Expected Outcome": "The router should reattempt to establish a connection with the ISP."
+            "Expected Outcome": {
+              "type": "string"
+            }
+          },
+          "required": [
+            "Action",
+            "Details",
+            "Expected Outcome"
+          ]
+        },
+        {
+          "type": "object",
+          "properties": {
+            "Action": {
+              "type": "string"
             },
-            {
-                "Action": "Verify ISP Login Credentials",
-                "Details": "Access the router settings through its web interface and check the PPPoE login credentials (username and password) provided by your ISP.",
-                "Expected Outcome": "Correct credentials should resolve the authentication issue."
+            "Details": {
+              "type": "string"
             },
-            {
-                "Action": "Contact ISP Support",
-                "Details": "If the above steps don't work, contact your ISP to check for any issues on their end such as outages or account problems.",
-                "Expected Outcome": "ISP support will help resolve server authentication issues or provide further guidance."
+            "Expected Outcome": {
+              "type": "string"
             }
-        ],
-        "Recommended_Actions": {
-            "Immediate_Action": "Reconnect cables and restart devices.",
-            "If_Unresolved": "Check ISP login credentials and contact ISP if needed.",
-            "Preventative_Measure": "Ensure router firmware is up-to-date to prevent known issues."
+          },
+          "required": [
+            "Action",
+            "Details",
+            "Expected Outcome"
+          ]
+        }
+      ]
+    },
+    "Recommended_Actions": {
+      "type": "object",
+      "properties": {
+        "Immediate_Action": {
+          "type": "string"
+        },
+        "If_Unresolved": {
+          "type": "string"
+        },
+        "Preventative_Measure": {
+          "type": "string"
         }
+      },
+      "required": [
+        "Immediate_Action",
+        "If_Unresolved",
+        "Preventative_Measure"
+      ]
     }
+  },
+  "required": [
+    "Issue_Description",
+    "Root_Cause_Analysis",
+    "Step_by_Step_Troubleshooting",
+    "Recommended_Actions"
+  ]
+}"
+SYSTEM_INNSTRUCTION="You are a router troubleshooter. Your job is to analyze the provided router image, identify potential issues such as faulty connections, incorrect LED patterns, or error codes, and offer precise troubleshooting steps. Based on your analysis, generate a detailed observation that includes a root cause analysis, step-by-step actions for resolving the issue, and recommended preventive measures. The output must be in JSON format as per the following schema, ensuring users can easily follow and implement the suggested solutions.\n" + SCHEMA_DEFINITION
+
+
+model_id = "meta-llama/Llama-3.2-11B-Vision-Instruct"
+
+model = MllamaForConditionalGeneration.from_pretrained(
+    model_id,
+    torch_dtype=torch.bfloat16,
+    device_map="auto",
+)
+processor = AutoProcessor.from_pretrained(model_id)
+
+
+def extract_json_from_markdown(markdown_text):
+    """Extract JSON or code block from markdown text."""
+    try:
+        # Find the start and end of the code block (with or without "json")
+        start_idx = markdown_text.find('```')
+        end_idx = markdown_text.find('```', start_idx + 3)
+        
+        # If the block starts with '```json', skip the 'json' part
+        if markdown_text[start_idx:start_idx + 7] == '```json':
+            start_idx += len('```json')
+        else:
+            start_idx += len('```')
+
+        # Extract and clean up the code block (json or not)
+        json_str = markdown_text[start_idx:end_idx].strip()
+
+        # Try to load it as JSON
+        return json.loads(json_str)
+    except Exception as e:
+        print(f"Error extracting JSON: {e}")
+        return None
+
+@spaces.GPU
+def diagnose_router(image):
+
+
+    
+
+    messages = [
+         {"role": "system",
+      "content": [{"type": "text",
+                   "text": SYSTEM_INNSTRUCTION}],
+         },
+        {"role": "user", "content": [
+            {"type": "image"},
+            {"type": "text", "text": "Analyze this router issue and provide the diagnosis."}
+        ]}
+    ]
+    input_text = processor.apply_chat_template(messages, add_generation_prompt=True)
+    inputs = processor(image, input_text, return_tensors="pt").to(model.device)
+
+    # Generate the output from the model
+    output = model.generate(**inputs, max_new_tokens=300)
+    markdown_text = processor.decode(output[0])
+
+    # Extract JSON from the markdown text
+    result = extract_json_from_markdown(markdown_text)
+
+    
+    print (result)
     
     # Generate HTML content for structured display
     html_output = f"""