from huggingface_hub import InferenceClient
from dotenv import load_dotenv
import os

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

prompt = "A detailed historical illustration of roman mercenaries"

configs = [
    {
        "model": "ByteDance/SDXL-Lightning",
        "params": {"guidance_scale": 0.0, "num_inference_steps": 4, "width":1024, "height":1024},
        "name": "lightning_g0"
    },
    {
        "model": "ByteDance/SDXL-Lightning",
        "params": {"guidance_scale": 1.0, "num_inference_steps": 4, "width":1024, "height":1024},
        "name": "lightning_g1"
    },
    {
        "model": "stabilityai/sdxl-turbo",
        "params": {"guidance_scale": 0.0, "num_inference_steps": 2, "width":512, "height":512}, # Turbo is 512 native usually? No 1024 is fine but safe 512 test
        "name": "turbo_g0"
    }
]

print("--- Testing Generation Params ---")
for conf in configs:
    print(f"\nTesting: {conf['name']} ({conf['model']})")
    client = InferenceClient(model=conf['model'], token=token)
    try:
        image = client.text_to_image(prompt, **conf['params'])
        fname = f"test_{conf['name']}.png"
        image.save(fname)
        print(f"SUCCESS! Saved {fname}")
    except Exception as e:
        print(f"FAILED: {e}")