diff options
| -rw-r--r-- | generate.py | 14 | ||||
| -rw-r--r-- | test_generate.py | 29 |
2 files changed, 42 insertions, 1 deletions
diff --git a/generate.py b/generate.py index 5eed627..eb20dc1 100644 --- a/generate.py +++ b/generate.py @@ -24,12 +24,24 @@ from pathlib import Path import torch from diffusers import StableDiffusionXLPipeline, AutoencoderKL -MODEL = "stabilityai/stable-diffusion-xl-base-1.0" +MODELS = { + "sdxl": "stabilityai/stable-diffusion-xl-base-1.0", # default, general + "realvis": "SG161222/RealVisXL_V5.0", # photoreal, permissive + "pony": "AstraliteHeart/pony-diffusion-v6-xl", # unrestricted, versatile +} +DEFAULT_MODEL = "sdxl" # fp16-fix VAE avoids the fp32 upcast that OOMs a 12GB card at decode time. VAE = "madebyollin/sdxl-vae-fp16-fix" OUT_DIR = Path("/out") +def resolve_model(key): + if key not in MODELS: + valid = ", ".join(sorted(MODELS)) + raise SystemExit(f"unknown model '{key}'. valid: {valid}") + return MODELS[key] + + def build_pipe(): if not torch.xpu.is_available(): # Fail loud and specific. A silent CPU fallback would run but take diff --git a/test_generate.py b/test_generate.py new file mode 100644 index 0000000..c2425d5 --- /dev/null +++ b/test_generate.py @@ -0,0 +1,29 @@ +import generate + + +def test_registry_integrity(): + assert generate.DEFAULT_MODEL in generate.MODELS + for key, repo in generate.MODELS.items(): + assert isinstance(key, str) and key + assert isinstance(repo, str) and "/" in repo + + +def test_resolve_known_key(): + assert generate.resolve_model("sdxl") == generate.MODELS["sdxl"] + assert generate.resolve_model("realvis") == generate.MODELS["realvis"] + + +def test_resolve_unknown_key_raises_listing_valid(): + try: + generate.resolve_model("nope") + except SystemExit as e: + assert "sdxl" in str(e) and "realvis" in str(e) and "pony" in str(e) + else: + assert False, "expected SystemExit for unknown key" + + +if __name__ == "__main__": + test_registry_integrity() + test_resolve_known_key() + test_resolve_unknown_key_raises_listing_valid() + print("ok") |