| """Unit tests for Imager-1-Pro.""" |
| import os |
| import torch |
| import pytest |
| from PIL import Image |
|
|
| from src.config import ImagerPro, MODEL_CONFIG, load_model, predict |
|
|
|
|
| def test_config(): |
| assert MODEL_CONFIG["architecture"] == "convnext_base" |
| assert MODEL_CONFIG["num_params"] == 89_000_000 |
| assert MODEL_CONFIG["input_size"] == 224 |
|
|
|
|
| def test_architecture_output_shape(): |
| model = ImagerPro(pretrained=False) |
| model.eval() |
| dummy = torch.randn(1, 3, 224, 224) |
| with torch.no_grad(): |
| out = model(dummy) |
| assert out.shape == (1, 1), f"Expected (1,1), got {out.shape}" |
|
|
|
|
| def test_architecture_param_count(): |
| model = ImagerPro(pretrained=False) |
| n = sum(p.numel() for p in model.parameters()) |
| |
| assert 85_000_000 < n < 92_000_000, f"Unexpected param count: {n}" |
|
|
|
|
| @pytest.mark.skipif( |
| not os.path.exists("model/model_imager-1-pro.pth"), |
| reason="Weights not present", |
| ) |
| def test_predict_no_nan(): |
| """Critical test: model must not output NaN.""" |
| model = load_model("model/model_imager-1-pro.pth") |
| img = Image.new("RGB", (512, 512), color=(128, 128, 128)) |
| result = predict(img, model=model) |
| assert result["probability"] == result["probability"], "Probability is NaN!" |
| assert 0.0 <= result["probability"] <= 1.0 |
| assert result["label"] in ("Real", "AI-Generated") |
|
|
