adjusted model loading

app.py

@@ -28,37 +28,30 @@ model_3n_id = os.getenv("MODEL_3N_ID", "google/gemma-3n-E4B-it")
 MAX_VIDEO_SIZE = 100 * 1024 * 1024  # 100 MB
 MAX_IMAGE_SIZE = 10 * 1024 * 1024   # 10 MB
 
-# Global variables to hold models (loaded lazily)
-input_processor = None
-model_12 = None
-model_3n = None
-
-def load_models():
-    """Load models lazily when needed."""
-    global input_processor, model_12, model_3n
-    
-    # Skip model loading during testing
-    if os.getenv("SKIP_MODEL_LOADING") == "true" or "pytest" in os.getenv("_", ""):
-        return
-    
-    if input_processor is None:
-        input_processor = Gemma3Processor.from_pretrained(model_12_id)
-    
-    if model_12 is None:
-        model_12 = Gemma3ForConditionalGeneration.from_pretrained(
-            model_12_id,
-            torch_dtype=torch.bfloat16,
-            device_map="auto",
-            attn_implementation="eager",
-        )
-    
-    if model_3n is None:
-        model_3n = Gemma3nForConditionalGeneration.from_pretrained(
-            model_3n_id,
-            torch_dtype=torch.bfloat16,
-            device_map="auto",
-            attn_implementation="eager",
-        )
+# Skip model loading during tests
+SKIP_MODEL_LOADING = os.getenv("SKIP_MODEL_LOADING", "false").lower() == "true"
+
+if not SKIP_MODEL_LOADING:
+    input_processor = Gemma3Processor.from_pretrained(model_12_id)
+
+    model_12 = Gemma3ForConditionalGeneration.from_pretrained(
+        model_12_id,
+        torch_dtype=torch.bfloat16,
+        device_map="auto",
+        attn_implementation="eager",
+    )
+
+    model_3n = Gemma3nForConditionalGeneration.from_pretrained(
+        model_3n_id,
+        torch_dtype=torch.bfloat16,
+        device_map="auto",
+        attn_implementation="eager",
+    )
+else:
+    # Mock objects for testing
+    input_processor = None
+    model_12 = None
+    model_3n = None
 
 
 def check_file_size(file_path: str) -> bool:
@@ -229,9 +222,6 @@ def run(
     repetition_penalty: float,
 ) -> Iterator[str]:
 
-    # Load models only when needed (not during testing)
-    load_models()
-
     # Define preset system prompts
     preset_prompts = {
         "General Assistant": "You are a helpful AI assistant capable of analyzing images, videos, and PDF documents. Provide clear, accurate, and helpful responses to user queries.",
@@ -259,6 +249,11 @@ def run(
     )
 
     selected_model = model_12 if model_choice == "Gemma 3 12B" else model_3n
+    
+    # If models are skipped (during testing), return a mock response
+    if SKIP_MODEL_LOADING or selected_model is None or input_processor is None:
+        yield "Mock response for testing - models are not loaded."
+        return
 
     messages = []
     if system_prompt:

tests/conftest.py

@@ -0,0 +1,12 @@
+import os
+import pytest
+
+# Set environment variable before any imports
+os.environ["SKIP_MODEL_LOADING"] = "true"
+
+@pytest.fixture(autouse=True)
+def skip_model_loading():
+    """Automatically set environment variable to skip model loading for all tests."""
+    # Environment variable is already set above
+    yield
+    # Keep the variable set for the entire test session