Update test_ai_integration_http.py

test_ai_integration_http.py

@@ -113,32 +113,46 @@ def test_ai_integration_http():
             logger.info(f"Loading {model_name}")
             
             try:
-                # Load processor first
-                processor = AutoProcessor.from_pretrained(
-                    model_name,
-                    trust_remote_code=True
-                )
-                status['processor_loaded'] = True
+                # Disable transformers logging temporarily
+                import logging
+                transformers_logger = logging.getLogger("transformers")
+                original_level = transformers_logger.level
+                transformers_logger.setLevel(logging.ERROR)
                 
-                # Load model directly - let it handle config internally
-                model = AutoModel.from_pretrained(
-                    model_name,
-                    trust_remote_code=True,
-                    torch_dtype=torch.float32,  # Use float32 for better compatibility
-                    device_map=None,  # Don't auto-map devices
-                    ignore_mismatched_sizes=True  # Handle any size mismatches
-                )
-                status['model_loaded'] = True
-                
-                logger.info(f"Processor type: {type(processor).__name__}")
-                logger.info(f"Model type: {type(model).__name__}")
+                try:
+                    # Load processor first
+                    processor = AutoProcessor.from_pretrained(
+                        model_name,
+                        trust_remote_code=True
+                    )
+                    status['processor_loaded'] = True
+                    
+                    # Import the specific model class
+                    from transformers.models.florence.modeling_florence import Florence2Model
+                    
+                    # Load model directly with specific class
+                    model = Florence2Model.from_pretrained(
+                        model_name,
+                        trust_remote_code=True,
+                        torch_dtype=torch.float32,
+                        device_map=None,
+                        ignore_mismatched_sizes=True
+                    )
+                    status['model_loaded'] = True
+                    
+                    # Log model details
+                    logger.info(f"Processor type: {type(processor).__name__}")
+                    logger.info(f"Model type: {type(model).__name__}")
+                    
+                    # 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__}")
                 
-                # Log model architecture and details
-                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__}")
-                if hasattr(model, 'config'):
-                    logger.info(f"Model config type: {type(model.config).__name__}")
+                finally:
+                    # Restore original logging level
+                    transformers_logger.setLevel(original_level)
+                    
             except Exception as e:
                 logger.error(f"Model loading failed: {str(e)}")
                 raise