Update test_ai_integration_http.py

test_ai_integration_http.py

@@ -12,12 +12,13 @@ from typing import Any, Optional
 import torch
 from torch import nn
 import torch.nn.functional as F
-from torch.overrides import TorchFunctionMode
+from torch.utils._python_dispatch import TorchFunctionMode
 from PIL import Image
 from transformers import (
     AutoTokenizer,
     AutoModelForCausalLM,
-    AutoProcessor
+    AutoProcessor,
+    AutoConfig
 )
 from virtual_vram import VirtualVRAM
 from http_storage import HTTPGPUStorage
@@ -113,21 +114,35 @@ def test_ai_integration_http():
             logger.info(f"Loading {model_name}")
             
             try:
+                # Load processor with direct configuration
                 processor = AutoProcessor.from_pretrained(
                     model_name,
-                    trust_remote_code=True
+                    trust_remote_code=True,
+                    return_tensors="pt"
                 )
+                status['processor_loaded'] = True
+                
+                # Load model with vision config
+                from transformers import AutoConfig
+                config = AutoConfig.from_pretrained(
+                    model_name,
+                    trust_remote_code=True,
+                    torch_dtype=torch.float32  # Use float32 for better compatibility
+                )
+                
                 model = AutoModelForCausalLM.from_pretrained(
                     model_name,
-                    trust_remote_code=True
+                    config=config,
+                    trust_remote_code=True,
+                    device_map=None  # Don't auto-map devices
                 )
-                status['processor_loaded'] = True
                 status['model_loaded'] = True
                 
                 # Log model architecture
                 model_size = get_model_size(model)
                 logger.info(f"Model loaded: {model_size/1e9:.2f} GB in parameters")
                 logger.info(f"Model architecture: {model.__class__.__name__}")
+                logger.info(f"Model config type: {type(config).__name__}")
             except Exception as e:
                 logger.error(f"Model loading failed: {str(e)}")
                 raise