vae/__init__.py

"""Local ROI-VAE compression package.

This package was previously named `compression`, but was renamed to `vae` to
avoid colliding with optional third-party imports (e.g. via Gradio/fsspec).

It intentionally uses lazy attribute loading to keep import-time overhead low.
"""

from __future__ import annotations

from importlib import import_module
from typing import TYPE_CHECKING


__all__ = [
    "TIC",
    "ModifiedTIC",
    "load_checkpoint",
    "compute_padding",
    "compress_image",
    "highlight_roi",
    "create_comparison_grid",
    "RSTB",
    "CausalAttentionModule",
]


_EXPORTS: dict[str, tuple[str, str]] = {
    "TIC": (".tic_model", "TIC"),
    "ModifiedTIC": (".roi_tic", "ModifiedTIC"),
    "load_checkpoint": (".roi_tic", "load_checkpoint"),
    "compute_padding": (".utils", "compute_padding"),
    "compress_image": (".utils", "compress_image"),
    "highlight_roi": (".visualization", "highlight_roi"),
    "create_comparison_grid": (".visualization", "create_comparison_grid"),
    "RSTB": (".RSTB", "RSTB"),
    "CausalAttentionModule": (".RSTB", "CausalAttentionModule"),
}


def __getattr__(name: str):
    if name not in _EXPORTS:
        raise AttributeError(f"module {__name__!r} has no attribute {name!r}")

    module_name, attr_name = _EXPORTS[name]
    module = import_module(module_name, package=__name__)
    value = getattr(module, attr_name)
    globals()[name] = value
    return value


if TYPE_CHECKING:
    from .RSTB import CausalAttentionModule, RSTB
    from .roi_tic import ModifiedTIC, load_checkpoint
    from .tic_model import TIC
    from .utils import compress_image, compute_padding
    from .visualization import create_comparison_grid, highlight_roi