tests/test_models.py

from huggingface_hub import InferenceClient
from dotenv import load_dotenv
import os

load_dotenv()
token = os.getenv("HF_API_TOKEN")

models_to_test = [
    "RunDiffusion/Juggernaut-XL-Lightning", # Top choice for quality/speed
    "stabilityai/stable-diffusion-xl-base-1.0", # Reliable baseline
    "playgroundai/playground-v2.5-1024px-aesthetic" # High aesthetic
]

prompt = "A detailed historical copperplate engraving of a steam engine"

print("--- Testing Model Availability ---")
for model in models_to_test:
    print(f"\nTesting: {model}")
    client = InferenceClient(model=model, token=token)
    try:
        image = client.text_to_image(prompt)
        print("SUCCESS! Model is available.")
        image.save(f"test_{model.replace('/', '_')}.png")
    except Exception as e:
        print(f"FAILED: {e}")