tests/test_model.py

import sys
import os
import torch
from PIL import Image
from transformers import BlipProcessor, BlipForConditionalGeneration

# Add parent directory to path to import app if needed, 
# but here we just test the model loading directly as per the plan.
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

def test_model_loading():
    print("Testing model loading...")
    model_path = "./blip-model"
    try:
        processor = BlipProcessor.from_pretrained(model_path)
        model = BlipForConditionalGeneration.from_pretrained(model_path)
        print("Model loaded successfully.")
        
        # Check device
        device = "cuda" if torch.cuda.is_available() else "cpu"
        print(f"Device available: {device}")
        model.to(device)
        
        # Test generation with a dummy image
        print("Testing generation with dummy image...")
        raw_image = Image.new('RGB', (384, 384), color = 'red')
        inputs = processor(raw_image, return_tensors="pt").to(device)
        out = model.generate(**inputs)
        caption = processor.decode(out[0], skip_special_tokens=True)
        print(f"Generated dummy caption: {caption}")
        print("Verification successful!")
        
    except Exception as e:
        print(f"Error during verification: {e}")
        sys.exit(1)

if __name__ == "__main__":
    test_model_loading()