test_pipeline.py

#!/usr/bin/env python3
"""
Simple test for the image-text-to-text pipeline setup
"""

import requests
from transformers import pipeline
import asyncio

def test_pipeline_availability():
    """Test if the image-text-to-text pipeline can be initialized"""
    print("🔍 Testing pipeline availability...")
    
    try:
        # Try to initialize the pipeline locally
        print("🚀 Initializing image-text-to-text pipeline...")
        
        # Try with a smaller, more accessible model first
        models_to_try = [
            "Salesforce/blip-image-captioning-base",  # More common model
            "microsoft/git-base-textcaps",  # Alternative model
            "unsloth/gemma-3n-E4B-it-GGUF"  # Original model
        ]
        
        for model_name in models_to_try:
            try:
                print(f"📥 Trying model: {model_name}")
                pipe = pipeline("image-to-text", model=model_name)  # Use image-to-text instead
                print(f"✅ Successfully loaded {model_name}")
                
                # Test with a simple image URL
                test_url = "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/p-blog/candy.JPG"
                print(f"🖼️ Testing with image: {test_url}")
                
                result = pipe(test_url)
                print(f"📝 Result: {result}")
                
                return True, model_name
                
            except Exception as e:
                print(f"❌ Failed to load {model_name}: {e}")
                continue
        
        print("❌ No suitable models could be loaded")
        return False, None
        
    except Exception as e:
        print(f"❌ Pipeline test error: {e}")
        return False, None

def test_backend_models_endpoint():
    """Test the backend models endpoint"""
    print("\n📋 Testing backend models endpoint...")
    
    try:
        response = requests.get("http://localhost:8000/v1/models", timeout=10)
        if response.status_code == 200:
            result = response.json()
            print(f"✅ Available models: {[model['id'] for model in result['data']]}")
            return True
        else:
            print(f"❌ Models endpoint failed: {response.status_code}")
            return False
    except Exception as e:
        print(f"❌ Models endpoint error: {e}")
        return False

def main():
    """Run pipeline tests"""
    print("🧪 Testing Image-Text Pipeline Setup\n")
    
    # Test 1: Check if we can initialize pipelines locally
    success, model_name = test_pipeline_availability()
    
    if success:
        print(f"\n🎉 Pipeline test successful with model: {model_name}")
        print("💡 Recommendation: Update backend_service.py to use this model")
    else:
        print("\n⚠️ Pipeline test failed")
        print("💡 Recommendation: Use image-to-text pipeline instead of image-text-to-text")
    
    # Test 2: Check backend models
    test_backend_models_endpoint()

if __name__ == "__main__":
    main()