File size: 1,422 Bytes
87520e2 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 | 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()
|