133 lines
4.1 KiB
Python
Executable File
133 lines
4.1 KiB
Python
Executable File
#!/usr/bin/env python3
|
|
"""Simple test script to test colqwen2 forward pass with a single image."""
|
|
|
|
import os
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
# Add the current directory to path to import leann_multi_vector
|
|
sys.path.insert(0, str(Path(__file__).parent))
|
|
|
|
import torch
|
|
from leann_multi_vector import _embed_images, _ensure_repo_paths_importable, _load_colvision
|
|
from PIL import Image
|
|
|
|
# Ensure repo paths are importable
|
|
_ensure_repo_paths_importable(__file__)
|
|
|
|
# Set environment variable
|
|
os.environ["TOKENIZERS_PARALLELISM"] = "false"
|
|
|
|
|
|
def create_test_image():
|
|
"""Create a simple test image."""
|
|
# Create a simple RGB image (800x600)
|
|
img = Image.new("RGB", (800, 600), color="white")
|
|
return img
|
|
|
|
|
|
def load_test_image_from_file():
|
|
"""Try to load an image from the indexes directory if available."""
|
|
# Try to find an existing image in the indexes directory
|
|
indexes_dir = Path(__file__).parent / "indexes"
|
|
|
|
# Look for images in common locations
|
|
possible_paths = [
|
|
indexes_dir / "vidore_fastplaid" / "images",
|
|
indexes_dir / "colvision_large.leann.images",
|
|
indexes_dir / "colvision.leann.images",
|
|
]
|
|
|
|
for img_dir in possible_paths:
|
|
if img_dir.exists():
|
|
# Find first image file
|
|
for ext in [".png", ".jpg", ".jpeg"]:
|
|
for img_file in img_dir.glob(f"*{ext}"):
|
|
print(f"Loading test image from: {img_file}")
|
|
return Image.open(img_file)
|
|
|
|
return None
|
|
|
|
|
|
def main():
|
|
print("=" * 60)
|
|
print("Testing ColQwen2 Forward Pass")
|
|
print("=" * 60)
|
|
|
|
# Step 1: Load or create test image
|
|
print("\n[Step 1] Loading test image...")
|
|
test_image = load_test_image_from_file()
|
|
if test_image is None:
|
|
print("No existing image found, creating a simple test image...")
|
|
test_image = create_test_image()
|
|
else:
|
|
print(f"✓ Loaded image: {test_image.size} ({test_image.mode})")
|
|
|
|
# Convert to RGB if needed
|
|
if test_image.mode != "RGB":
|
|
test_image = test_image.convert("RGB")
|
|
print(f"✓ Converted to RGB: {test_image.size}")
|
|
|
|
# Step 2: Load model
|
|
print("\n[Step 2] Loading ColQwen2 model...")
|
|
try:
|
|
model_name, model, processor, device_str, device, dtype = _load_colvision("colqwen2")
|
|
print(f"✓ Model loaded: {model_name}")
|
|
print(f"✓ Device: {device_str}, dtype: {dtype}")
|
|
|
|
# Print model info
|
|
if hasattr(model, "device"):
|
|
print(f"✓ Model device: {model.device}")
|
|
if hasattr(model, "dtype"):
|
|
print(f"✓ Model dtype: {model.dtype}")
|
|
|
|
except Exception as e:
|
|
print(f"✗ Error loading model: {e}")
|
|
import traceback
|
|
|
|
traceback.print_exc()
|
|
return
|
|
|
|
# Step 3: Test forward pass
|
|
print("\n[Step 3] Running forward pass...")
|
|
try:
|
|
# Use the _embed_images function which handles batching and forward pass
|
|
images = [test_image]
|
|
print(f"Processing {len(images)} image(s)...")
|
|
|
|
doc_vecs = _embed_images(model, processor, images)
|
|
|
|
print("✓ Forward pass completed!")
|
|
print(f"✓ Number of embeddings: {len(doc_vecs)}")
|
|
|
|
if len(doc_vecs) > 0:
|
|
emb = doc_vecs[0]
|
|
print(f"✓ Embedding shape: {emb.shape}")
|
|
print(f"✓ Embedding dtype: {emb.dtype}")
|
|
print("✓ Embedding stats:")
|
|
print(f" - Min: {emb.min().item():.4f}")
|
|
print(f" - Max: {emb.max().item():.4f}")
|
|
print(f" - Mean: {emb.mean().item():.4f}")
|
|
print(f" - Std: {emb.std().item():.4f}")
|
|
|
|
# Check for NaN or Inf
|
|
if torch.isnan(emb).any():
|
|
print("⚠ Warning: Embedding contains NaN values!")
|
|
if torch.isinf(emb).any():
|
|
print("⚠ Warning: Embedding contains Inf values!")
|
|
|
|
except Exception as e:
|
|
print(f"✗ Error during forward pass: {e}")
|
|
import traceback
|
|
|
|
traceback.print_exc()
|
|
return
|
|
|
|
print("\n" + "=" * 60)
|
|
print("Test completed successfully!")
|
|
print("=" * 60)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|