dockerfile

app.py

@@ -2,15 +2,36 @@ import gradio as gr
 from PIL import Image
 import os
 import sys
-from llava_inference import LLaVAHelper
+import logging
+
+# Set up logging
+logging.basicConfig(
+    level=logging.INFO,
+    format='%(asctime)s - %(levelname)s - %(message)s',
+    handlers=[logging.StreamHandler(sys.stdout)]
+)
+logger = logging.getLogger(__name__)
+
+logger.info("Starting UK Public Transport Assistant app")
 
 # Add error handling for module imports
 try:
-    model = LLaVAHelper()
+    from llava_inference import LLaVAHelper
+    logger.info("Successfully imported LLaVAHelper")
 except Exception as e:
-    print(f"Failed to initialize LLaVA model: {e}")
-    # Continue execution to show error in the UI
+    logger.error(f"Failed to import LLaVAHelper: {e}")
+    logger.error("Stack trace:", exc_info=True)
     model = None
+else:
+    # Initialize model
+    try:
+        logger.info("Initializing LLaVA model...")
+        model = LLaVAHelper()
+        logger.info("LLaVA model initialized successfully")
+    except Exception as e:
+        logger.error(f"Failed to initialize LLaVA model: {e}")
+        logger.error("Stack trace:", exc_info=True)
+        model = None
 
 def answer_question(image, question):
     if model is None:

dockerfile

@@ -0,0 +1,27 @@
+FROM python:3.10-slim
+
+WORKDIR /app
+
+# Install system dependencies
+RUN apt-get update && apt-get install -y \
+    build-essential \
+    git \
+    && rm -rf /var/lib/apt/lists/*
+
+# Copy requirements first for better caching
+COPY requirements.txt .
+
+# Install Python dependencies
+RUN pip install --no-cache-dir -r requirements.txt
+
+# Copy application code
+COPY . .
+
+# Make sure the cache directory exists
+RUN mkdir -p ./model_cache
+
+# Expose port for Gradio
+EXPOSE 7860
+
+# Start the application
+CMD ["python", "app.py"]

llava_inference.py

@@ -1,45 +1,115 @@
-from llava.model.builder import load_pretrained_model
-from llava.mm_utils import process_images, tokenizer_image_token
-from transformers import AutoTokenizer, AutoConfig
-import torch
-import requests
-from PIL import Image
-from io import BytesIO
+import sys
+import logging
 import os
 
+# Configure logging
+logging.basicConfig(
+    level=logging.INFO,
+    format='%(asctime)s - %(levelname)s - %(message)s',
+    handlers=[logging.StreamHandler(sys.stdout)]
+)
+logger = logging.getLogger(__name__)
+
+# First try to import from llava
+try:
+    from llava.model.builder import load_pretrained_model
+    from llava.mm_utils import process_images, tokenizer_image_token
+    logger.info("Successfully imported llava modules")
+except ImportError as e:
+    logger.error(f"Failed to import llava modules: {e}")
+    sys.exit(1)
+
+# Then import other dependencies
+try:
+    from transformers import AutoTokenizer, AutoConfig
+    import torch
+    import requests
+    from PIL import Image
+    from io import BytesIO
+    logger.info("Successfully imported other required modules")
+except ImportError as e:
+    logger.error(f"Failed to import dependency: {e}")
+    sys.exit(1)
+
 class LLaVAHelper:
     def __init__(self, model_name="llava-hf/llava-1.5-7b-hf"):
+        """
+        Initialize the LLaVA model for image-text processing
+        """
+        logger.info(f"Initializing LLaVAHelper with model: {model_name}")
+        
         # Create cache directory if it doesn't exist
         os.makedirs("./model_cache", exist_ok=True)
+        logger.info("Created model cache directory")
         
-        # First, try loading just the config to ensure the model is valid
+        # Try loading just the config to ensure the model is valid
         try:
             AutoConfig.from_pretrained(model_name)
+            logger.info(f"Successfully loaded config for {model_name}")
         except Exception as e:
-            print(f"Error loading model config: {e}")
+            logger.warning(f"Error loading model config: {e}")
             # Try a different model version as fallback
             model_name = "llava-hf/llava-1.5-13b-hf"
-            print(f"Trying alternative model: {model_name}")
+            logger.info(f"Trying alternative model: {model_name}")
 
         try:
             # Use specific tokenizer class to avoid issues
+            logger.info("Loading tokenizer...")
             self.tokenizer = AutoTokenizer.from_pretrained(
                 model_name,
                 cache_dir="./model_cache",
                 use_fast=False,  # Use the Python implementation instead of the Rust one
-                legacy=True
+                trust_remote_code=True
             )
+            logger.info("Tokenizer loaded successfully")
             
-            # Load model with same cache directory - fix parameter order
-            self.model, self.image_processor, _ = load_pretrained_model(
-                model_path=model_name,
-                model_name=model_name,  # Add the missing model_name parameter
-                model_base=None,
-                cache_dir="./model_cache",
-                load_8bit=False,
-                load_4bit=False,
-                device_map="auto"
-            )
+            # Inspect the load_pretrained_model function to understand its parameters
+            import inspect
+            logger.info(f"load_pretrained_model signature: {inspect.signature(load_pretrained_model)}")
+            
+            # Try loading with different parameter combinations
+            logger.info("Loading model...")
+            try:
+                # First attempt - standard parameter order
+                self.model, self.image_processor, _ = load_pretrained_model(
+                    model_path=model_name,
+                    model_base=None,
+                    cache_dir="./model_cache",
+                )
+            except Exception as e1:
+                logger.warning(f"First attempt to load model failed: {e1}")
+                try:
+                    # Second attempt - try with model_name parameter
+                    self.model, self.image_processor, _ = load_pretrained_model(
+                        model_name=model_name,
+                        model_path=model_name,
+                        model_base=None,
+                        cache_dir="./model_cache",
+                    )
+                except Exception as e2:
+                    logger.warning(f"Second attempt to load model failed: {e2}")
+                    # Third attempt - minimal parameters
+                    self.model, self.image_processor, _ = load_pretrained_model(
+                        model_name,
+                        None,
+                        "./model_cache",
+                    )
+            
+            logger.info("Model loaded successfully")
+            self.model.eval()
+            
+            # Move model to appropriate device
+            self.device = "cuda" if torch.cuda.is_available() else "cpu"
+            logger.info(f"Using device: {self.device}")
+            if self.device == "cpu":
+                # If using CPU, make sure model is in the right place
+                self.model = self.model.to(self.device)
+            
+            logger.info(f"Model successfully loaded on {self.device}")
+        except Exception as e:
+            logger.error(f"Detailed initialization error: {e}")
+            logger.error("Stack trace:", exc_info=True)
+            raise
             
             self.model.eval()