code/inference_profile.py

"""Checkpoint capability contract for public Ming image inference.

New packages declare their capability in ``transformer/config.json`` via the
``alignment_padding_mode`` / ``multi_frame_output`` pair; the VAE contract is
derived from ``vae/config.json``. Legacy packages that predate the component
metadata are loaded strictly from the root ``inference_profile.json`` during
the compatibility window. The contract is intentionally strict: selecting
behavior from directory names, missing state-dict keys, or task names can
silently load the wrong padding semantics and produce degraded images.
"""

from __future__ import annotations

from dataclasses import dataclass
import json
import math
from pathlib import Path
from typing import Any, Mapping, Optional, Union


PROFILE_FILENAME = "inference_profile.json"
PROFILE_SCHEMA_VERSION = 1

GENERATION_PROFILE = "generation_edit"
LAYER_PROFILE = "layer_decompose"
VALID_PROFILES = {GENERATION_PROFILE, LAYER_PROFILE}

LEARNED_PADDING = "learned"
ZERO_MASKED_PADDING = "zero_masked"
VALID_PADDING_MODES = {LEARNED_PADDING, ZERO_MASKED_PADDING}

VALID_VAE_SAMPLE_MODES = {"sample", "argmax"}
VALID_TASKS = {"text-to-image", "image-edit", "layer-decompose"}

REQUIRED_PROFILE_FIELDS = {
    "schema_version",
    "inference_profile",
    "alignment_padding_mode",
    "multi_frame_output",
    "vae_input_channels",
    "vae_sample_mode",
}


class InferenceProfileError(ValueError):
    """Raised when a checkpoint does not satisfy the inference contract."""


@dataclass(frozen=True)
class SamplingParameters:
    steps: int
    cfg: float


DEFAULT_SAMPLING_PARAMETERS = {
    GENERATION_PROFILE: SamplingParameters(steps=12, cfg=1.0),
    LAYER_PROFILE: SamplingParameters(steps=12, cfg=2.0),
}


@dataclass(frozen=True)
class InferenceProfile:
    schema_version: int
    inference_profile: str
    alignment_padding_mode: str
    multi_frame_output: bool
    vae_input_channels: int
    vae_sample_mode: str

    @classmethod
    def from_dict(cls, raw: Mapping[str, Any]) -> "InferenceProfile":
        missing = sorted(REQUIRED_PROFILE_FIELDS.difference(raw))
        if missing:
            raise InferenceProfileError(
                "checkpoint inference profile is missing required fields: "
                + ", ".join(missing)
            )

        unknown = sorted(set(raw).difference(REQUIRED_PROFILE_FIELDS))
        if unknown:
            raise InferenceProfileError(
                "checkpoint inference profile contains unsupported fields: "
                + ", ".join(unknown)
            )

        if type(raw["schema_version"]) is not int:
            raise InferenceProfileError("schema_version must be an integer")
        if type(raw["multi_frame_output"]) is not bool:
            raise InferenceProfileError("multi_frame_output must be a boolean")
        if type(raw["vae_input_channels"]) is not int:
            raise InferenceProfileError("vae_input_channels must be an integer")

        profile = cls(**{key: raw[key] for key in REQUIRED_PROFILE_FIELDS})
        profile.validate()
        return profile

    def validate(self) -> None:
        if self.schema_version != PROFILE_SCHEMA_VERSION:
            raise InferenceProfileError(
                f"unsupported profile schema_version={self.schema_version}; "
                f"expected {PROFILE_SCHEMA_VERSION}"
            )
        if self.inference_profile not in VALID_PROFILES:
            raise InferenceProfileError(
                f"inference_profile must be one of {sorted(VALID_PROFILES)}, "
                f"got {self.inference_profile!r}"
            )
        if self.alignment_padding_mode not in VALID_PADDING_MODES:
            raise InferenceProfileError(
                "alignment_padding_mode must be one of "
                f"{sorted(VALID_PADDING_MODES)}, got {self.alignment_padding_mode!r}"
            )
        if self.vae_input_channels not in (3, 4):
            raise InferenceProfileError(
                f"vae_input_channels must be 3 or 4, got {self.vae_input_channels}"
            )
        if self.vae_sample_mode not in VALID_VAE_SAMPLE_MODES:
            raise InferenceProfileError(
                f"vae_sample_mode must be one of {sorted(VALID_VAE_SAMPLE_MODES)}, "
                f"got {self.vae_sample_mode!r}"
            )

        if self.inference_profile == GENERATION_PROFILE:
            if self.alignment_padding_mode != ZERO_MASKED_PADDING:
                raise InferenceProfileError(
                    "generation_edit checkpoints must use zero_masked alignment padding"
                )
            if self.multi_frame_output:
                raise InferenceProfileError(
                    "generation_edit checkpoints must set multi_frame_output=false"
                )
            if self.vae_input_channels != 4:
                raise InferenceProfileError(
                    "generation_edit checkpoints must declare vae_input_channels=4"
                )
            if self.vae_sample_mode != "argmax":
                raise InferenceProfileError(
                    "generation_edit checkpoints must declare vae_sample_mode='argmax'"
                )
        else:
            if self.alignment_padding_mode != LEARNED_PADDING:
                raise InferenceProfileError(
                    "layer_decompose checkpoints must use learned alignment padding"
                )
            if not self.multi_frame_output:
                raise InferenceProfileError(
                    "layer_decompose checkpoints must set multi_frame_output=true"
                )
            if self.vae_input_channels != 4:
                raise InferenceProfileError(
                    "layer_decompose checkpoints must declare vae_input_channels=4"
                )
            if self.vae_sample_mode != "argmax":
                raise InferenceProfileError(
                    "layer_decompose checkpoints must declare vae_sample_mode='argmax'"
                )

    def validate_task(
        self,
        task: str,
        *,
        has_reference_image: bool,
        num_layers: int = 1,
    ) -> None:
        if task not in VALID_TASKS:
            raise InferenceProfileError(
                f"task must be one of {sorted(VALID_TASKS)}, got {task!r}"
            )
        if num_layers < 1:
            raise InferenceProfileError("num_layers must be at least 1")

        if task == "layer-decompose":
            if self.inference_profile != LAYER_PROFILE:
                raise InferenceProfileError(
                    "layer-decompose requires a layer_decompose checkpoint"
                )
            if not has_reference_image:
                raise InferenceProfileError(
                    "layer-decompose requires an input reference image"
                )
            return

        if self.inference_profile != GENERATION_PROFILE:
            raise InferenceProfileError(
                f"{task} requires a generation_edit checkpoint"
            )
        if num_layers != 1:
            raise InferenceProfileError(
                f"{task} does not support num_layers={num_layers}; expected 1"
            )
        if task == "text-to-image" and has_reference_image:
            raise InferenceProfileError("text-to-image does not accept a reference image")
        if task == "image-edit" and not has_reference_image:
            raise InferenceProfileError("image-edit requires a reference image")

    def resolve_sampling_parameters(
        self,
        *,
        steps: Optional[int] = None,
        cfg: Optional[float] = None,
    ) -> SamplingParameters:
        defaults = DEFAULT_SAMPLING_PARAMETERS[self.inference_profile]
        resolved_steps = defaults.steps if steps is None else steps
        resolved_cfg = defaults.cfg if cfg is None else cfg

        if type(resolved_steps) is not int or resolved_steps < 1:
            raise InferenceProfileError("sampling steps must be an integer >= 1")
        if (
            isinstance(resolved_cfg, bool)
            or not isinstance(resolved_cfg, (int, float))
            or not math.isfinite(float(resolved_cfg))
            or resolved_cfg < 0
        ):
            raise InferenceProfileError("CFG must be a finite number >= 0")

        return SamplingParameters(
            steps=resolved_steps,
            cfg=float(resolved_cfg),
        )


def load_inference_profile(model_directory: Union[str, Path]) -> InferenceProfile:
    """Legacy parser: strict root ``inference_profile.json`` (see module doc)."""

    model_directory = Path(model_directory)
    profile_path = model_directory / PROFILE_FILENAME
    if not profile_path.is_file():
        raise InferenceProfileError(
            f"checkpoint must contain {PROFILE_FILENAME}: {profile_path}"
        )
    try:
        raw = json.loads(profile_path.read_text(encoding="utf-8"))
    except (OSError, json.JSONDecodeError) as error:
        raise InferenceProfileError(
            f"failed to read checkpoint inference profile {profile_path}: {error}"
        ) from error
    if not isinstance(raw, dict):
        raise InferenceProfileError("checkpoint inference profile must be a JSON object")
    return InferenceProfile.from_dict(raw)


TRANSFORMER_CONFIG_FILENAME = "transformer/config.json"
VAE_CONFIG_FILENAME = "vae/config.json"
CAPABILITY_FIELDS = ("alignment_padding_mode", "multi_frame_output")
QWEN_VAE_CLASS_NAME = "AutoencoderKLQwenImage"

# The only valid capability pairs and the runtime family they select.
CAPABILITY_PROFILES = {
    (ZERO_MASKED_PADDING, False): GENERATION_PROFILE,
    (LEARNED_PADDING, True): LAYER_PROFILE,
}


def _read_component_config(model_directory: Path, relative: str) -> Mapping[str, Any]:
    config_path = model_directory / relative
    if not config_path.is_file():
        raise InferenceProfileError(
            f"checkpoint is missing component config {relative}: {config_path}"
        )
    try:
        raw = json.loads(config_path.read_text(encoding="utf-8"))
    except (OSError, json.JSONDecodeError) as error:
        raise InferenceProfileError(
            f"failed to read component config {config_path}: {error}"
        ) from error
    if not isinstance(raw, dict):
        raise InferenceProfileError(f"component config {relative} must be a JSON object")
    return raw


def _derive_vae_contract(model_directory: Path) -> tuple[int, str]:
    """(vae_input_channels, vae_sample_mode) from the VAE component config."""

    config = _read_component_config(model_directory, VAE_CONFIG_FILENAME)
    class_name = config.get("_class_name")
    if class_name != QWEN_VAE_CLASS_NAME:
        raise InferenceProfileError(
            f"unsupported VAE contract: {VAE_CONFIG_FILENAME} _class_name must be "
            f"{QWEN_VAE_CLASS_NAME!r} for argmax reference encoding, got "
            f"{class_name!r}"
        )
    declared = [
        config[key]
        for key in ("input_channels", "in_channels")
        if key in config
    ]
    if not declared:
        raise InferenceProfileError(
            f"{VAE_CONFIG_FILENAME} must declare input_channels or in_channels"
        )
    if len(set(declared)) != 1:
        raise InferenceProfileError(
            f"{VAE_CONFIG_FILENAME} input_channels and in_channels disagree: "
            f"{declared}"
        )
    channels = declared[0]
    if type(channels) is not int or channels != 4:
        raise InferenceProfileError(
            "the supported public families require a 4-channel VAE, got "
            f"{channels!r} in {VAE_CONFIG_FILENAME}"
        )
    return channels, "argmax"


def _profile_from_components(
    model_directory: Path,
    alignment_padding_mode: Any,
    multi_frame_output: Any,
) -> InferenceProfile:
    if type(alignment_padding_mode) is not str:
        raise InferenceProfileError(
            f"{TRANSFORMER_CONFIG_FILENAME} alignment_padding_mode must be a "
            f"string, got {alignment_padding_mode!r}"
        )
    if type(multi_frame_output) is not bool:
        raise InferenceProfileError(
            f"{TRANSFORMER_CONFIG_FILENAME} multi_frame_output must be a "
            f"boolean, got {multi_frame_output!r}"
        )
    pair = (alignment_padding_mode, multi_frame_output)
    if pair not in CAPABILITY_PROFILES:
        raise InferenceProfileError(
            f"unsupported capability pair {pair!r} in "
            f"{TRANSFORMER_CONFIG_FILENAME}; expected one of "
            f"{sorted(CAPABILITY_PROFILES)}"
        )
    channels, sample_mode = _derive_vae_contract(model_directory)
    profile = InferenceProfile(
        schema_version=PROFILE_SCHEMA_VERSION,
        inference_profile=CAPABILITY_PROFILES[pair],
        alignment_padding_mode=alignment_padding_mode,
        multi_frame_output=multi_frame_output,
        vae_input_channels=channels,
        vae_sample_mode=sample_mode,
    )
    profile.validate()
    return profile


def load_checkpoint_capabilities(
    model_directory: Union[str, Path],
) -> InferenceProfile:
    """Derive the runtime capability from component configs.

    New packages declare ``alignment_padding_mode`` and
    ``multi_frame_output`` in ``transformer/config.json``; legacy packages
    without them fall back to the strict root ``inference_profile.json``.
    A missing ``transformer/config.json`` means both fields are absent, so
    the legacy path applies. A partial pair is a hard error. When both the
    component fields and the legacy file exist, the component metadata is
    authoritative and the legacy file may only agree with it.
    """

    model_directory = Path(model_directory)
    transformer_path = model_directory / TRANSFORMER_CONFIG_FILENAME
    transformer_config = (
        _read_component_config(model_directory, TRANSFORMER_CONFIG_FILENAME)
        if transformer_path.is_file()
        else {}
    )
    present = [field for field in CAPABILITY_FIELDS if field in transformer_config]
    if len(present) == 1:
        raise InferenceProfileError(
            f"{TRANSFORMER_CONFIG_FILENAME} carries only {present[0]!r}; "
            "alignment_padding_mode and multi_frame_output must be declared "
            "together"
        )
    if not present:
        return load_inference_profile(model_directory)

    profile = _profile_from_components(
        model_directory,
        transformer_config["alignment_padding_mode"],
        transformer_config["multi_frame_output"],
    )
    legacy_path = model_directory / PROFILE_FILENAME
    if legacy_path.is_file():
        legacy = load_inference_profile(model_directory)
        if legacy != profile:
            raise InferenceProfileError(
                f"component configs disagree with legacy {PROFILE_FILENAME}: "
                f"transformer/vae derive {profile!r} but the profile declares "
                f"{legacy!r}"
            )
    return profile


def resolve_model_directory(
    model_name_or_path: Union[str, Path],
    *,
    revision: Optional[str] = None,
    cache_dir: Optional[Union[str, Path]] = None,
    local_files_only: bool = False,
    token: Optional[Union[str, bool]] = None,
) -> Path:
    """Resolve a local directory or materialize a Hugging Face Hub snapshot."""

    candidate = Path(model_name_or_path).expanduser()
    if candidate.is_dir():
        return candidate.resolve()
    if candidate.exists():
        raise ValueError(f"model path must be a directory: {candidate}")
    if candidate.is_absolute():
        raise FileNotFoundError(f"local model directory does not exist: {candidate}")

    try:
        from huggingface_hub import snapshot_download
    except ImportError as error:
        raise RuntimeError(
            "huggingface_hub is required when --model is a Hub repository ID"
        ) from error

    snapshot_path = snapshot_download(
        repo_id=str(model_name_or_path),
        revision=revision,
        cache_dir=str(cache_dir) if cache_dir is not None else None,
        local_files_only=local_files_only,
        token=token,
    )
    return Path(snapshot_path).resolve()