load_model.py

"""
Simple loading script for pixeltext-ai model
"""

from transformers import AutoModel
from PIL import Image

def load_pixeltext_model():
    """Load the pixeltext-ai model properly."""
    
    # Method 1: Direct loading (recommended)
    try:
        from modeling_pixeltext import FixedPaliGemmaOCR
        model = FixedPaliGemmaOCR()
        return model
    except Exception as e:
        print(f"Direct loading failed: {e}")
        
        # Method 2: Fallback to AutoModel
        try:
            model = AutoModel.from_pretrained(
                "BabaK07/pixeltext-ai", 
                trust_remote_code=True
            )
            return model
        except Exception as e2:
            print(f"AutoModel loading failed: {e2}")
            return None

def test_model():
    """Test the loaded model."""
    
    model = load_pixeltext_model()
    if model is None:
        print("❌ Failed to load model")
        return
    
    # Create test image
    from PIL import Image, ImageDraw, ImageFont
    
    img = Image.new('RGB', (400, 200), color='white')
    draw = ImageDraw.Draw(img)
    
    try:
        font = ImageFont.truetype("/System/Library/Fonts/Arial.ttf", 20)
    except:
        font = ImageFont.load_default()
    
    draw.text((20, 50), "Hello World!", fill='black', font=font)
    draw.text((20, 100), "This is a test for pixeltext-ai", fill='blue', font=font)
    
    # Test OCR
    result = model.generate_ocr_text(img)
    
    print("📝 OCR Results:")
    print(f"   Text: {result['text']}")
    print(f"   Confidence: {result['confidence']:.3f}")
    print(f"   Quality: {result['quality']}")
    
    return result

if __name__ == "__main__":
    test_model()