Update test_ai_integration_http.py

test_ai_integration_http.py

@@ -12,12 +12,11 @@ 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,
+    AutoModel,
     AutoConfig
 )
 from virtual_vram import VirtualVRAM
@@ -114,26 +113,21 @@ def test_ai_integration_http():
             logger.info(f"Loading {model_name}")
             
             try:
-                # Load processor with direct configuration
+                # Load processor first
                 processor = AutoProcessor.from_pretrained(
                     model_name,
-                    trust_remote_code=True,
-                    return_tensors="pt"
+                    trust_remote_code=True
                 )
                 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
-                )
+                # Get config and model class from the processor
+                config = processor.config
                 
-                model = AutoModelForCausalLM.from_pretrained(
+                # Load the model as a general vision model
+                model = AutoModel.from_pretrained(
                     model_name,
-                    config=config,
                     trust_remote_code=True,
+                    torch_dtype=torch.float32,  # Use float32 for better compatibility
                     device_map=None  # Don't auto-map devices
                 )
                 status['model_loaded'] = True