code/tests/test_inference_profile.py

import json
import tempfile
from pathlib import Path
import unittest

from inference_profile import (
    InferenceProfile,
    InferenceProfileError,
    load_checkpoint_capabilities,
    load_inference_profile,
    resolve_model_directory,
)


GENERATION = {
    "schema_version": 1,
    "inference_profile": "generation_edit",
    "alignment_padding_mode": "zero_masked",
    "multi_frame_output": False,
    "vae_input_channels": 4,
    "vae_sample_mode": "argmax",
}

LAYER = {
    "schema_version": 1,
    "inference_profile": "layer_decompose",
    "alignment_padding_mode": "learned",
    "multi_frame_output": True,
    "vae_input_channels": 4,
    "vae_sample_mode": "argmax",
}


class InferenceProfileTest(unittest.TestCase):
    def test_all_profile_fields_are_required(self):
        for field in GENERATION:
            raw = dict(GENERATION)
            raw.pop(field)
            with self.assertRaisesRegex(InferenceProfileError, "missing required fields"):
                InferenceProfile.from_dict(raw)

    def test_profile_rejects_unknown_fields(self):
        raw = dict(GENERATION, inferred_from_directory_name=True)
        with self.assertRaisesRegex(InferenceProfileError, "unsupported fields"):
            InferenceProfile.from_dict(raw)

    def test_generation_profile_task_matrix(self):
        profile = InferenceProfile.from_dict(GENERATION)
        self.assertEqual(profile.resolve_sampling_parameters().steps, 12)
        self.assertEqual(profile.resolve_sampling_parameters().cfg, 1.0)
        profile.validate_task("text-to-image", has_reference_image=False)
        profile.validate_task("image-edit", has_reference_image=True)
        with self.assertRaisesRegex(InferenceProfileError, "layer_decompose checkpoint"):
            profile.validate_task("layer-decompose", has_reference_image=True, num_layers=4)

    def test_layer_profile_task_matrix(self):
        profile = InferenceProfile.from_dict(LAYER)
        self.assertEqual(profile.resolve_sampling_parameters().steps, 12)
        self.assertEqual(profile.resolve_sampling_parameters().cfg, 2.0)
        profile.validate_task("layer-decompose", has_reference_image=True, num_layers=4)
        with self.assertRaisesRegex(InferenceProfileError, "generation_edit checkpoint"):
            profile.validate_task("image-edit", has_reference_image=True)

    def test_generation_profile_requires_qwen_vae_contract(self):
        with self.assertRaisesRegex(InferenceProfileError, "vae_input_channels=4"):
            InferenceProfile.from_dict(dict(GENERATION, vae_input_channels=3))
        with self.assertRaisesRegex(InferenceProfileError, "vae_sample_mode='argmax'"):
            InferenceProfile.from_dict(dict(GENERATION, vae_sample_mode="sample"))

    def test_layer_profile_requires_qwen_vae_contract(self):
        with self.assertRaisesRegex(InferenceProfileError, "vae_sample_mode='argmax'"):
            InferenceProfile.from_dict(dict(LAYER, vae_sample_mode="sample"))

    def test_sampling_overrides_are_validated(self):
        profile = InferenceProfile.from_dict(GENERATION)
        resolved = profile.resolve_sampling_parameters(steps=18, cfg=1.5)
        self.assertEqual((resolved.steps, resolved.cfg), (18, 1.5))
        with self.assertRaisesRegex(InferenceProfileError, "steps"):
            profile.resolve_sampling_parameters(steps=0)
        with self.assertRaisesRegex(InferenceProfileError, "CFG"):
            profile.resolve_sampling_parameters(cfg=float("inf"))

    def test_profile_is_loaded_from_local_model_directory(self):
        with self.subTest("local profile"):
            import tempfile
            from pathlib import Path

            with tempfile.TemporaryDirectory() as directory:
                model_directory = Path(directory)
                (model_directory / "inference_profile.json").write_text(
                    json.dumps(LAYER), encoding="utf-8"
                )
                self.assertEqual(
                    load_inference_profile(model_directory).inference_profile,
                    "layer_decompose",
                )
                self.assertEqual(
                    resolve_model_directory(model_directory), model_directory.resolve()
                )

    def test_missing_profile_is_an_error(self):
        import tempfile

        with tempfile.TemporaryDirectory() as directory:
            with self.assertRaisesRegex(InferenceProfileError, "must contain"):
                load_inference_profile(directory)

    def test_missing_absolute_local_path_is_not_treated_as_hub_id(self):
        import tempfile
        from pathlib import Path

        with tempfile.TemporaryDirectory() as directory:
            with self.assertRaises(FileNotFoundError):
                resolve_model_directory(Path(directory) / "missing")


VAE_QWEN_4CH = {"_class_name": "AutoencoderKLQwenImage", "input_channels": 4}


def build_component_package(
    root: Path,
    transformer_extra=None,
    vae=None,
    profile=None,
) -> None:
    (root / "transformer").mkdir(parents=True, exist_ok=True)
    (root / "vae").mkdir(parents=True, exist_ok=True)
    transformer = {"_class_name": "DiffusionTransformer", "dim": 4}
    transformer.update(transformer_extra or {})
    (root / "transformer" / "config.json").write_text(json.dumps(transformer))
    (root / "vae" / "config.json").write_text(json.dumps(vae or VAE_QWEN_4CH))
    if profile is not None:
        (root / "inference_profile.json").write_text(json.dumps(profile))


class ComponentCapabilityTest(unittest.TestCase):
    def capabilities(self, **kwargs):
        with tempfile.TemporaryDirectory() as directory:
            root = Path(directory)
            build_component_package(root, **kwargs)
            return load_checkpoint_capabilities(root)

    def test_generation_pair_without_profile(self):
        profile = self.capabilities(
            transformer_extra={
                "alignment_padding_mode": "zero_masked",
                "multi_frame_output": False,
            }
        )
        self.assertEqual(profile.inference_profile, "generation_edit")
        self.assertEqual(profile.vae_input_channels, 4)
        self.assertEqual(profile.vae_sample_mode, "argmax")
        self.assertEqual(profile.resolve_sampling_parameters().steps, 12)
        self.assertEqual(profile.resolve_sampling_parameters().cfg, 1.0)

    def test_layer_pair_without_profile(self):
        profile = self.capabilities(
            transformer_extra={
                "alignment_padding_mode": "learned",
                "multi_frame_output": True,
            }
        )
        self.assertEqual(profile.inference_profile, "layer_decompose")
        self.assertEqual(profile.resolve_sampling_parameters().cfg, 2.0)

    def test_partial_pair_is_rejected(self):
        for field in ("alignment_padding_mode", "multi_frame_output"):
            with self.subTest(only=field):
                pair = {
                    "alignment_padding_mode": "zero_masked",
                    "multi_frame_output": False,
                }
                with self.assertRaisesRegex(InferenceProfileError, "declared together"):
                    self.capabilities(transformer_extra={field: pair[field]})

    def test_every_invalid_pair_is_rejected(self):
        invalid = [
            {"alignment_padding_mode": "zero_masked", "multi_frame_output": True},
            {"alignment_padding_mode": "learned", "multi_frame_output": False},
            {"alignment_padding_mode": "zero", "multi_frame_output": False},
            {"alignment_padding_mode": "learned", "multi_frame_output": 1},
            {"alignment_padding_mode": None, "multi_frame_output": False},
        ]
        for pair in invalid:
            with self.subTest(pair=pair):
                with self.assertRaises(InferenceProfileError):
                    self.capabilities(transformer_extra=pair)

    def test_vae_class_must_be_qwen(self):
        with self.assertRaisesRegex(InferenceProfileError, "AutoencoderKLQwenImage"):
            self.capabilities(
                transformer_extra={
                    "alignment_padding_mode": "zero_masked",
                    "multi_frame_output": False,
                },
                vae={"_class_name": "AutoencoderKL", "in_channels": 4},
            )

    def test_vae_channel_fields_must_agree_and_be_four(self):
        base = {"alignment_padding_mode": "zero_masked", "multi_frame_output": False}
        with self.assertRaisesRegex(InferenceProfileError, "disagree"):
            self.capabilities(
                transformer_extra=base,
                vae=dict(VAE_QWEN_4CH, in_channels=16),
            )
        with self.assertRaisesRegex(InferenceProfileError, "4-channel"):
            self.capabilities(
                transformer_extra=base,
                vae={"_class_name": "AutoencoderKLQwenImage", "input_channels": 3},
            )
        with self.assertRaisesRegex(InferenceProfileError, "must declare"):
            self.capabilities(
                transformer_extra=base,
                vae={"_class_name": "AutoencoderKLQwenImage"},
            )
        # in_channels alone is accepted when it agrees by itself.
        profile = self.capabilities(
            transformer_extra=base,
            vae={"_class_name": "AutoencoderKLQwenImage", "in_channels": 4},
        )
        self.assertEqual(profile.vae_input_channels, 4)

    def test_legacy_profile_fallback(self):
        profile = self.capabilities(profile=LAYER)
        self.assertEqual(profile.inference_profile, "layer_decompose")
        self.assertEqual(profile.alignment_padding_mode, "learned")

    def test_legacy_fallback_requires_profile_file(self):
        with tempfile.TemporaryDirectory() as directory:
            root = Path(directory)
            build_component_package(root)
            with self.assertRaisesRegex(InferenceProfileError, "must contain"):
                load_checkpoint_capabilities(root)

    def test_component_fields_and_agreeing_legacy_profile(self):
        profile = self.capabilities(
            transformer_extra={
                "alignment_padding_mode": "zero_masked",
                "multi_frame_output": False,
            },
            profile=GENERATION,
        )
        self.assertEqual(profile.inference_profile, "generation_edit")

    def test_component_fields_disagreeing_with_legacy_profile(self):
        with self.assertRaisesRegex(InferenceProfileError, "disagree"):
            self.capabilities(
                transformer_extra={
                    "alignment_padding_mode": "learned",
                    "multi_frame_output": True,
                },
                profile=GENERATION,
            )


if __name__ == "__main__":
    unittest.main()