import os
import io
from PIL import Image, UnidentifiedImageError
from huggingface_hub import InferenceClient
from dotenv import load_dotenv

load_dotenv()

def test_repro():
    token = os.getenv("HF_API_TOKEN") or os.getenv("HF_TOKEN")
    client = InferenceClient(model="ByteDance/SDXL-Lightning", token=token)
    prompt = "A simple test prompt, historical illustration"
    
    print(f"Requesting image for: {prompt}")
    try:
        # We manually call the API to see the raw result if possible, 
        # but text_to_image is usually what's used.
        result = client.text_to_image(prompt)
        print(f"Result type: {type(result)}")
        
        if hasattr(result, 'load'):
            print("Calling result.load()...")
            result.load()
            print("Success!")
        else:
            print("Result is not a PIL Image (no .load() method)")
            
    except Exception as e:
        print(f"Caught exception: {type(e).__name__}: {e}")
        # Check if it's a data error
        err_msg = str(e).lower()
        is_data_err = any(x in err_msg for x in ["unrecognized data", "unidentifiedimage", "truncated", "data stream"])
        print(f"Is identified as data error? {is_data_err}")

if __name__ == "__main__":
    test_repro()