scripts/02_teacher_smoke.py

"""Teacher smoke test: load Flux2KleinPipeline, generate a couple of 512/4-step images.
Validates the inference path (SDPA on A100) and the __call__ signature before surgery."""
import os
import time

import torch

from flux2distill.model_utils import load_pipeline

os.makedirs("outputs/teacher_smoke", exist_ok=True)

t0 = time.time()
pipe = load_pipeline(device="cuda")
print(f"pipeline loaded in {time.time()-t0:.1f}s; transformer dtype={next(pipe.transformer.parameters()).dtype}")
print("scheduler:", type(pipe.scheduler).__name__)

prompts = [
    'a vintage bookshop storefront with a wooden sign that reads "THE OPEN PAGE"',
    "a serene mountain lake at sunrise reflecting snow-capped peaks, mist over the water",
]

torch.cuda.synchronize()
t0 = time.time()
gen = torch.Generator(device="cuda").manual_seed(0)
out = pipe(prompt=prompts, num_inference_steps=4, guidance_scale=1.0,
           height=512, width=512, generator=gen)
torch.cuda.synchronize()
dt = time.time() - t0
imgs = out.images
for i, im in enumerate(imgs):
    im.save(f"outputs/teacher_smoke/teacher_{i}.png")
print(f"generated {len(imgs)} imgs in {dt:.2f}s ({dt/len(imgs):.2f}s/img) at 512/4steps")
print(f"peak VRAM: {torch.cuda.max_memory_allocated()/1e9:.1f}GB")
print("saved to outputs/teacher_smoke/")