Update app.py

Fixed side extension

app.py

@@ -6,10 +6,9 @@ Upload any image and predict where it was taken using Vision-Language Models
 
 import gradio as gr
 from PIL import Image
-from transformers import AutoProcessor, AutoModelForImageTextToText
+from transformers import AutoProcessor, Qwen2VLForConditionalGeneration
 import torch
 import re
-import math
 from dataclasses import dataclass
 
 # ============================================================================
@@ -61,7 +60,6 @@ def parse_response(text: str) -> ParsedResponse:
     if not text:
         return ParsedResponse(None, None, None, None, text, False)
     
-    # Parse key-value lines
     key_pattern = re.compile(
         r'^\s*(?:[-*+\u2022]\s*)?(?P<key>[A-Za-z][A-Za-z0-9\s\-/_.]*?)\s*:\s*(?P<value>.+)$'
     )
@@ -97,7 +95,6 @@ def parse_response(text: str) -> ParsedResponse:
                     except ValueError:
                         pass
     
-    # Build coords if available
     coords = None
     if "lat" in parsed and "lon" in parsed:
         try:
@@ -131,110 +128,114 @@ def load_model():
     """Load model once on startup"""
     global model, processor
     if model is None:
-        print(f"Loading model: {MODEL_NAME}")
-        processor = AutoProcessor.from_pretrained(MODEL_NAME)
-        model = AutoModelForImageTextToText.from_pretrained(
-            MODEL_NAME,
-            torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
-            device_map="auto" if torch.cuda.is_available() else "cpu"
-        )
-        print("Model loaded successfully!")
+        print(f"🔄 Loading model: {MODEL_NAME}")
+        try:
+            processor = AutoProcessor.from_pretrained(MODEL_NAME, trust_remote_code=True)
+            model = Qwen2VLForConditionalGeneration.from_pretrained(
+                MODEL_NAME,
+                torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
+                device_map="auto",
+                trust_remote_code=True
+            )
+            print("✅ Model loaded successfully!")
+        except Exception as e:
+            print(f"❌ Error loading model: {e}")
+            raise
 
 def predict_location(image):
     """Predict geolocation from an image"""
-    if image is None:
-        return "Please upload an image.", ""
-    
-    # Ensure model is loaded
-    load_model()
-    
-    # Convert to PIL if needed
-    if not isinstance(image, Image.Image):
-        image = Image.fromarray(image).convert("RGB")
-    else:
-        image = image.convert("RGB")
-    
-    # Prepare prompt
-    messages = [
-        {
-            "role": "user",
-            "content": [
-                {"type": "image"},
-                {"type": "text", "text": PROMPT_TEMPLATE}
-            ]
-        }
-    ]
-    
-    # Process inputs
-    text = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
-    inputs = processor(text=[text], images=[image], return_tensors="pt", padding=True)
-    
-    # Move to device
-    inputs = {k: v.to(model.device) for k, v in inputs.items()}
-    
-    # Generate
-    with torch.no_grad():
-        output_ids = model.generate(
-            **inputs,
-            max_new_tokens=256,
-            do_sample=False,
-        )
-    
-    # Decode
-    generated_ids = output_ids[0][inputs['input_ids'].shape[1]:]
-    response = processor.decode(generated_ids, skip_special_tokens=True).strip()
-    
-    # Parse
-    parsed = parse_response(response)
-    
-    # Format output
-    output = f"""
-## 🤖 Raw Model Response:
-```
-{response}
-```
+    try:
+        if image is None:
+            return "⚠️ Please upload an image.", ""
+        
+        load_model()
+        
+        print("📸 Processing image...")
+        
+        if not isinstance(image, Image.Image):
+            image = Image.fromarray(image).convert("RGB")
+        else:
+            image = image.convert("RGB")
+        
+        messages = [
+            {
+                "role": "user",
+                "content": [
+                    {"type": "image", "image": image},
+                    {"type": "text", "text": PROMPT_TEMPLATE}
+                ]
+            }
+        ]
+        
+        text = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+        inputs = processor(text=[text], images=[image], return_tensors="pt", padding=True)
+        inputs = {k: v.to(model.device) for k, v in inputs.items()}
+        
+        print("🤖 Generating prediction...")
+        
+        with torch.no_grad():
+            output_ids = model.generate(**inputs, max_new_tokens=256, do_sample=False)
+        
+        generated_ids = output_ids[0][inputs['input_ids'].shape[1]:]
+        response = processor.decode(generated_ids, skip_special_tokens=True).strip()
+        
+        print(f"✅ Response generated")
+        
+        parsed = parse_response(response)
+        
+        output = f"""
+## 🤖 AI Prediction
 
----
+**📍 Location Details:**
+- **City:** {parsed.city or "Unknown"}
+- **Region:** {parsed.region or "Unknown"}
+- **Country:** {parsed.country or "Unknown"}
+- **Coordinates:** {f"{parsed.coords.lat:.6f}°, {parsed.coords.lon:.6f}°" if parsed.coords else "Not found"}
 
-## 📍 Parsed Prediction:
+---
 
-**City:** {parsed.city or "Not provided"}  
-**Region:** {parsed.region or "Not provided"}  
-**Country:** {parsed.country or "Not provided"}  
-**Coordinates:** {f"{parsed.coords.lat:.6f}, {parsed.coords.lon:.6f}" if parsed.coords else "Not provided"}  
-**Format Valid:** {"✅ Yes" if parsed.format_valid else "❌ No"}
+## 🔍 Raw Response:
+```
+{response}
+```
 """
-    
-    # Create map embed
-    map_html = ""
-    if parsed.coords:
-        map_html = f"""
-        <div style="margin-top: 20px;">
-            <iframe 
-                width="100%" 
-                height="450" 
-                frameborder="0" 
-                scrolling="no" 
-                marginheight="0" 
-                marginwidth="0" 
-                src="https://www.openstreetmap.org/export/embed.html?bbox={parsed.coords.lon-0.1},{parsed.coords.lat-0.1},{parsed.coords.lon+0.1},{parsed.coords.lat+0.1}&marker={parsed.coords.lat},{parsed.coords.lon}"
-                style="border: 2px solid #ddd; border-radius: 8px;">
-            </iframe>
-            <div style="margin-top: 10px; text-align: center;">
-                <a href="https://www.google.com/maps?q={parsed.coords.lat},{parsed.coords.lon}" target="_blank" style="margin: 0 10px; color: #4285f4; text-decoration: none; font-weight: bold;">
-                    🗺️ View on Google Maps
-                </a>
-                <span style="color: #666;">|</span>
-                <a href="https://www.openstreetmap.org/?mlat={parsed.coords.lat}&mlon={parsed.coords.lon}#map=12/{parsed.coords.lat}/{parsed.coords.lon}" target="_blank" style="margin: 0 10px; color: #7ebc6f; text-decoration: none; font-weight: bold;">
-                    🌍 View on OpenStreetMap
-                </a>
+        
+        map_html = ""
+        if parsed.coords:
+            map_html = f"""
+            <div style="margin-top: 20px;">
+                <iframe 
+                    width="100%" 
+                    height="450" 
+                    frameborder="0" 
+                    scrolling="no" 
+                    marginheight="0" 
+                    marginwidth="0" 
+                    src="https://www.openstreetmap.org/export/embed.html?bbox={parsed.coords.lon-0.1},{parsed.coords.lat-0.1},{parsed.coords.lon+0.1},{parsed.coords.lat+0.1}&marker={parsed.coords.lat},{parsed.coords.lon}"
+                    style="border: 2px solid #ddd; border-radius: 8px;">
+                </iframe>
+                <div style="margin-top: 10px; text-align: center;">
+                    <a href="https://www.google.com/maps?q={parsed.coords.lat},{parsed.coords.lon}" target="_blank" style="margin: 0 10px; color: #4285f4; text-decoration: none; font-weight: bold;">
+                        🗺️ Google Maps
+                    </a>
+                    <span style="color: #666;">|</span>
+                    <a href="https://www.openstreetmap.org/?mlat={parsed.coords.lat}&mlon={parsed.coords.lon}#map=12/{parsed.coords.lat}/{parsed.coords.lon}" target="_blank" style="margin: 0 10px; color: #7ebc6f; text-decoration: none; font-weight: bold;">
+                        🌍 OpenStreetMap
+                    </a>
+                </div>
             </div>
-        </div>
-        """
-    else:
-        map_html = "<div style='text-align: center; padding: 20px; color: #666;'>No valid coordinates found</div>"
-    
-    return output, map_html
+            """
+        else:
+            map_html = "<div style='text-align: center; padding: 20px; color: #666;'>❌ No valid coordinates found</div>"
+        
+        return output, map_html
+        
+    except Exception as e:
+        error_msg = f"❌ Error: {str(e)}"
+        print(error_msg)
+        import traceback
+        traceback.print_exc()
+        return error_msg, ""
 
 # ============================================================================
 # Gradio Interface
@@ -245,12 +246,7 @@ with gr.Blocks(title="GeoVLM - AI Geolocation", theme=gr.themes.Soft()) as demo:
         """
         # 🌍 GeoVLM - AI-Powered Geolocation
         
-        Upload any image and let AI predict where it was taken using vision-language models!
-        
-        ### How it works:
-        - Analyzes visual features: architecture, vegetation, road signs, landscape
-        - Uses state-of-the-art vision-language models (Qwen2-VL)
-        - Predicts city, region, country, and GPS coordinates
+        Upload any image and let AI predict where it was taken!
         
         **Powered by [vlm-gym](https://github.com/sdan/vlm-gym)** | Model: Qwen2-VL-2B-Instruct
         """
@@ -258,11 +254,7 @@ with gr.Blocks(title="GeoVLM - AI Geolocation", theme=gr.themes.Soft()) as demo:
     
     with gr.Row():
         with gr.Column(scale=1):
-            image_input = gr.Image(
-                type="pil",
-                label="📸 Upload Image",
-                height=400
-            )
+            image_input = gr.Image(type="pil", label="📸 Upload Image", height=400)
             predict_btn = gr.Button("🔍 Predict Location", variant="primary", size="lg")
             
             gr.Markdown(
@@ -276,34 +268,29 @@ with gr.Blocks(title="GeoVLM - AI Geolocation", theme=gr.themes.Soft()) as demo:
             )
         
         with gr.Column(scale=1):
-            output_text = gr.Markdown(label="Results")
-            map_output = gr.HTML(label="Map")
+            output_text = gr.Markdown(label="📊 Results")
+            map_output = gr.HTML(label="🗺️ Map Location")
     
     gr.Markdown(
         """
         ---
         ### 🎯 Use Cases:
-        - **OSINT Research** - Verify photo locations for investigations
+        - **OSINT Research** - Verify photo locations
         - **GeoGuessr Training** - Practice location identification
-        - **Education** - Learn about geographic features and cultures
-        - **Travel Planning** - Identify interesting locations from photos
+        - **Education** - Learn world geography
+        - **Travel** - Discover interesting places
         
         ---
         
-        **Note:** This is a demo. Predictions may not always be accurate. Use responsibly for educational and research purposes.
+        **Note:** Predictions take 2-5 minutes on CPU. Accuracy varies by location.
         
-        Built with ❤️ using [Gradio](https://gradio.app) and [Hugging Face Transformers](https://huggingface.co/transformers)
+        Built by [Vance Poitier](https://www.linkedin.com/in/vance-poitier/) | Based on [vlm-gym](https://github.com/sdan/vlm-gym)
         """
     )
     
-    # Event handlers
-    predict_btn.click(
-        fn=predict_location,
-        inputs=image_input,
-        outputs=[output_text, map_output]
-    )
+    predict_btn.click(fn=predict_location, inputs=image_input, outputs=[output_text, map_output])
 
 if __name__ == "__main__":
     print("🚀 Starting GeoVLM...")
     load_model()
-    demo.launch()
+    demo.launch(server_name="0.0.0.0", server_port=7860)