Update app.py

app.py

@@ -20,44 +20,41 @@ if torch.cuda.is_available():
     model = model.to('cuda')
 
 def analyze_image(image, prompt):
-    try:
-        messages = [
-            {"role": "user", "content": [
-                {"type": "image"},
-                {"type": "text", "text": prompt}
-            ]}
-        ]
-        input_text = processor.apply_chat_template(messages, add_generation_prompt=True)
-        
-        inputs = processor(
-            image,
-            input_text,
-            add_special_tokens=False,
-            return_tensors="pt"
-        ).to(model.device)
-        
-        # # Separate inputs for generate method
-        # generate_inputs = {
-        #     k: v for k, v in inputs.items() 
-        #     if k not in ['pixel_values', 'aspect_ratio_ids', 'aspect_ratio_mask']
-        # }
-        
-        with torch.no_grad():
-            output = model.generate(**inputs, max_new_tokens=100)
-        print(processor.decode(output[0]))
-        full_response = processor.decode(output[0])
     
-        
-        json_match = re.search(r'\{.*\}', full_response, re.DOTALL)
-        if json_match:
-            json_str = json_match.group(0)
-            try:
-                return json.loads(json_str)
-            except json.JSONDecodeError:
-                return {"error": "Failed to parse model output as JSON", "raw_output": json_str}
-        else:
-            return {"error": "No JSON found in the response", "raw_output": full_response}
+    messages = [
+        {"role": "user", "content": [
+            {"type": "image"},
+            {"type": "text", "text": prompt}
+        ]}
+    ]
+    input_text = processor.apply_chat_template(messages, add_generation_prompt=True)
+    
+    inputs = processor(
+        image,
+        input_text,
+        add_special_tokens=False,
+        return_tensors="pt"
+    ).to(model.device)
+    
+    # # Separate inputs for generate method
+    # generate_inputs = {
+    #     k: v for k, v in inputs.items() 
+    #     if k not in ['pixel_values', 'aspect_ratio_ids', 'aspect_ratio_mask']
+    # }
+    
+    with torch.no_grad():
+        output = model.generate(**inputs, max_new_tokens=100)
+    print(processor.decode(output[0]))
+    full_response = processor.decode(output[0])
 
+    
+    json_match = re.search(r'\{.*\}', full_response, re.DOTALL)
+    if json_match:
+        json_str = json_match.group(0)
+            return json.loads(json_str)
+    else:
+        return full_response
+    
 
 default_prompt = """Analyze this image and determine if it contains a data logger. A data logger is typically a small, black electronic device used to monitor and record data over time, such as voltage, temperature, or current, via external sensors.