Source code for simvx.graphics.assets.ktx2_loader

"""Pure-Python KTX2 (Khronos Texture 2.0) parser for block-compressed textures.

Parses a ``.ktx2`` container into a :class:`KTX2Texture` whose field set mirrors
:class:`~simvx.graphics.assets.dds_loader.DDSTexture` exactly, so the texture
manager treats both identically (sibling types, no shim, no unification).

KTX2 is the easy container: it stores the target ``VkFormat`` as an int in the
header, so there is no DXGI translation table to maintain. Block sizing reuses
:func:`~simvx.graphics.assets.dds_loader.vk_format_block_size`.

Scope: 2D single-array BC1-BC7 with ``supercompressionScheme`` of 0 (none) or 2
(Zstd). ``vkFormat == 0`` with a UASTC LDR Data Format Descriptor is transcoded
to a device-chosen block target (BC7 / ASTC-4x4 / ETC2) via the native
basis_universal transcoder when available, then handed to the same
compressed-upload path as explicit-BC textures. ETC1S /
BasisLZ (``scheme == 1``) remains out of scope and raises a ``ValueError`` naming
basis_universal. Cubemaps, arrays, 3D, and ZLIB supercompression also raise
``ValueError`` (the caller converts that to a one-time WARNING + a -1
"couldn't resolve" sentinel).

Pure stdlib (``struct`` + ``pathlib``) plus a LAZY ``zstandard`` import inside
the Zstd branch and a LAZY native-transcoder import inside the UASTC branch, so
this module imports cleanly with neither the GPU, the optional ``zstandard``
package, nor the native transcoder extension present.
"""

import logging
import struct
from dataclasses import dataclass
from pathlib import Path

from .dds_loader import vk_format_block_size

__all__ = ["KTX2Texture", "KTX2_MAGIC", "UASTCSource", "load_ktx2", "load_ktx2_uastc_source"]

log = logging.getLogger(__name__)

# 12-byte KTX2 file identifier (spec 2.0).
KTX2_MAGIC = b"\xab\x4b\x54\x58\x20\x32\x30\xbb\x0d\x0a\x1a\x0a"

# supercompressionScheme values.
_SC_NONE = 0
_SC_BASISLZ = 1
_SC_ZSTD = 2
_SC_ZLIB = 3

# Header (17 u32) + level-index start. The level index begins at byte 112:
# 12 (magic) + 68 (17xu32 header) + 32 (index block: 4xu32 + 2xu64) = 112.
_HEADER_OFFSET = 12
_INDEX_OFFSET = 80
_LEVEL_INDEX_OFFSET = 112

# Data Format Descriptor: colorModel (KHR_DF_MODEL_*) and transferFunction
# (KHR_DF_TRANSFER_*) live in the 4th u32 of the descriptor block, which itself
# follows the 4-byte dfdTotalSize. So colorModel = byte at dfdByteOffset+12,
# transferFunction = byte at dfdByteOffset+14.
_DFD_BITS_OFFSET = 12
_KHR_DF_MODEL_ETC1S = 163
_KHR_DF_MODEL_UASTC = 166
_KHR_DF_TRANSFER_SRGB = 2

# UASTC transcode targets -> (UNORM VkFormat name, SRGB VkFormat name). Resolved
# lazily by name (vulkan only touched at call time, so this module imports
# cleanly with no GPU). All three families pack 16 B/block, so block_size is a
# constant 16 for every UASTC transcode target. ``TextureManager`` picks which
# target via the device probe; the loader just transcodes to the requested one.
_UASTC_TARGET_FORMATS: dict[str, tuple[str, str]] = {
    "bc7": ("VK_FORMAT_BC7_UNORM_BLOCK", "VK_FORMAT_BC7_SRGB_BLOCK"),
    "astc4x4": ("VK_FORMAT_ASTC_4x4_UNORM_BLOCK", "VK_FORMAT_ASTC_4x4_SRGB_BLOCK"),
    "etc2": ("VK_FORMAT_ETC2_R8G8B8A8_UNORM_BLOCK", "VK_FORMAT_ETC2_R8G8B8A8_SRGB_BLOCK"),
}

# Per-target native transcode fn name on ``basis_transcoder``.
_UASTC_TARGET_FN: dict[str, str] = {
    "bc7": "transcode_uastc_to_bc7",
    "astc4x4": "transcode_uastc_to_astc4x4",
    "etc2": "transcode_uastc_to_etc2",
}


[docs] @dataclass(frozen=True) class KTX2Texture: """A parsed block-compressed KTX2 texture. Field set and order DELIBERATELY match :class:`~simvx.graphics.assets.dds_loader.DDSTexture` so the texture manager can feed either to the same compressed-upload path. Attributes: vk_format: ``VkFormat`` integer (a ``VK_FORMAT_BC*`` constant). width / height: top-level mip dimensions in texels. block_size: bytes per 4x4 block (8 for BC1/BC4, 16 otherwise). mips: tightly-packed block bytes per mip level, level 0 (largest) first. """ vk_format: int width: int height: int block_size: int mips: list[bytes]
[docs] @dataclass(frozen=True) class UASTCSource: """Raw UASTC LDR 4x4 source levels, untranscoded. The web backend ships these bytes to the browser, where the basis_universal WebAssembly transcoder picks the GPU's preferred block format (BC7 / ASTC / ETC2) per-device, or decodes to RGBA8 when no block family is present. The desktop loader transcodes UASTC->BC7 at parse time instead (one fixed target), which is why the web path needs the raw levels rather than ``KTX2Texture``. Attributes: width / height: top-level mip dimensions in texels. srgb: True when the DFD transfer function is sRGB (colour textures). mips: tightly-packed UASTC block bytes per mip level, level 0 first. """ width: int height: int srgb: bool mips: list[bytes]
[docs] def load_ktx2(source: str | Path | bytes, *, target: str = "bc7") -> KTX2Texture: """Parse a KTX2 file path or raw bytes into a :class:`KTX2Texture`. ``target`` selects the UASTC transcode target for ``vkFormat == 0`` files: ``"bc7"`` (default, preserves all existing callers byte-for-byte), ``"astc4x4"``, or ``"etc2"``. It is ignored for explicit-format (BC) KTX2s, which carry their own VkFormat. ``TextureManager`` picks ``target`` from the device's compressed-texture caps; the loader stays device-agnostic. Raises ``ValueError`` (never crashes) for: bad magic, truncation, a cubemap / array / 3D container, an explicit non-BC format, BasisLZ / ZLIB supercompression, or a missing ``zstandard`` package on a Zstd file. The caller converts that to a one-time WARNING + -1 sentinel. """ data = Path(source).read_bytes() if isinstance(source, str | Path) else bytes(source) if data[:12] != KTX2_MAGIC: raise ValueError("Not a KTX2 file: bad magic") if len(data) < _LEVEL_INDEX_OFFSET: raise ValueError("KTX2 file truncated: header/index incomplete") ( vk_format, _type_size, width, height, depth, layer_count, face_count, level_count, scheme, ) = struct.unpack_from("<9I", data, _HEADER_OFFSET) if depth != 0: raise ValueError("KTX2 3D textures are not supported in this phase") if layer_count > 1: raise ValueError("KTX2 texture arrays are not supported in this phase") if face_count == 6: raise ValueError("KTX2 cubemap textures are not supported in this phase") if face_count not in (0, 1): raise ValueError(f"KTX2 unexpected faceCount {face_count}") if scheme == _SC_ZLIB: raise ValueError("KTX2 ZLIB supercompression is not supported (phase T2b)") if scheme not in (_SC_NONE, _SC_ZSTD, _SC_BASISLZ): raise ValueError(f"KTX2 unsupported supercompressionScheme {scheme}") if vk_format == 0: # vkFormat 0: basis-universal (UASTC or ETC1S/BasisLZ). The DFD colorModel # disambiguates. UASTC LDR is transcoded to the requested ``target``; # ETC1S/BasisLZ stays out of scope and raises (the caller -> warn + -1). return _load_uastc(data, width, height, level_count, scheme, target) if scheme == _SC_BASISLZ: raise ValueError("KTX2 BasisLZ supercompression needs Basis Universal transcoding (phase T2b)") try: block_size = vk_format_block_size(vk_format) except ValueError: raise ValueError(f"KTX2 vkFormat {vk_format} is not a supported BC block format") from None levels = level_count or 1 mips = _read_levels(data, levels, scheme) log.debug("KTX2 parsed: %dx%d fmt=%d block=%d mips=%d scheme=%d", width, height, vk_format, block_size, len(mips), scheme) return KTX2Texture(vk_format=vk_format, width=width, height=height, block_size=block_size, mips=mips)
[docs] def load_ktx2_uastc_source(source: str | Path | bytes) -> UASTCSource: """Parse a UASTC LDR KTX2 into raw (untranscoded) :class:`UASTCSource` levels. For the web backend, which ships UASTC to the browser instead of transcoding to a single fixed block format on the CPU. Raises ``ValueError`` (the same contract as :func:`load_ktx2`) for anything that is not a plain or Zstd-supercompressed UASTC LDR 4x4 texture: explicit-BC vkFormats, ETC1S / BasisLZ, cubemaps / arrays / 3D, ZLIB, or a missing ``zstandard`` package. """ data = Path(source).read_bytes() if isinstance(source, str | Path) else bytes(source) if data[:12] != KTX2_MAGIC: raise ValueError("Not a KTX2 file: bad magic") if len(data) < _LEVEL_INDEX_OFFSET: raise ValueError("KTX2 file truncated: header/index incomplete") ( vk_format, _type_size, width, height, depth, layer_count, face_count, level_count, scheme, ) = struct.unpack_from("<9I", data, _HEADER_OFFSET) if depth != 0: raise ValueError("KTX2 3D textures are not supported in this phase") if layer_count > 1: raise ValueError("KTX2 texture arrays are not supported in this phase") if face_count == 6: raise ValueError("KTX2 cubemap textures are not supported in this phase") if scheme == _SC_ZLIB: raise ValueError("KTX2 ZLIB supercompression is not supported (phase T2b)") if vk_format != 0: raise ValueError("KTX2 has an explicit vkFormat (not UASTC); use load_ktx2") color_model, is_srgb = _dfd_color_model_and_srgb(data) if color_model == _KHR_DF_MODEL_ETC1S: raise ValueError("KTX2 ETC1S/BasisLZ needs Basis Universal transcoding (phase T2b, out of scope)") if color_model != _KHR_DF_MODEL_UASTC: raise ValueError(f"KTX2 vkFormat 0 unknown DFD colorModel {color_model} (expected UASTC 166)") if scheme == _SC_BASISLZ: raise ValueError("KTX2 UASTC with BasisLZ supercompression is invalid") levels = level_count or 1 mips = _read_levels(data, levels, scheme) log.debug("KTX2 UASTC source: %dx%d srgb=%s mips=%d", width, height, is_srgb, len(mips)) return UASTCSource(width=width, height=height, srgb=is_srgb, mips=mips)
def _read_levels(data: bytes, levels: int, scheme: int) -> list[bytes]: """Read + (Zstd-)decompress every mip level into raw block bytes. Shared by the explicit-BC and UASTC paths: for UASTC the resulting bytes are the raw packed UASTC blocks ready for transcoding. """ index_end = _LEVEL_INDEX_OFFSET + levels * 24 if len(data) < index_end: raise ValueError("KTX2 file truncated: level index incomplete") mips: list[bytes] = [] for i in range(levels): off, comp_len, raw_len = struct.unpack_from("<3Q", data, _LEVEL_INDEX_OFFSET + i * 24) chunk = data[off:off + comp_len] if len(chunk) < comp_len: raise ValueError(f"KTX2 data truncated at level {i}: need {comp_len} bytes, have {len(chunk)}") if scheme == _SC_ZSTD: try: import zstandard # noqa: PLC0415 (lazy: optional dep, only on a Zstd file) except ImportError: raise ValueError("KTX2 Zstd supercompression needs the 'zstandard' package") from None chunk = zstandard.ZstdDecompressor().decompress(chunk, max_output_size=raw_len) if len(chunk) != raw_len: raise ValueError(f"KTX2 level {i} zstd size mismatch: got {len(chunk)}, expected {raw_len}") mips.append(chunk) return mips def _dfd_color_model_and_srgb(data: bytes) -> tuple[int, bool]: """Parse the DFD colorModel byte and the sRGB transfer flag. The DFD byte offset/length live in the index block at ``_INDEX_OFFSET`` (dfdByteOffset, dfdByteLength as the first two u32). colorModel and transferFunction sit in the 4th u32 of the descriptor block, which follows the 4-byte ``dfdTotalSize``. """ dfd_offset, dfd_length = struct.unpack_from("<2I", data, _INDEX_OFFSET) if dfd_offset == 0 or dfd_length < _DFD_BITS_OFFSET + 4: raise ValueError("KTX2 vkFormat 0 with no/short Data Format Descriptor; cannot identify basis mode") if len(data) < dfd_offset + _DFD_BITS_OFFSET + 4: raise ValueError("KTX2 file truncated: Data Format Descriptor incomplete") (dfd_bits,) = struct.unpack_from("<I", data, dfd_offset + _DFD_BITS_OFFSET) color_model = dfd_bits & 0xFF transfer_func = (dfd_bits >> 16) & 0xFF return color_model, transfer_func == _KHR_DF_TRANSFER_SRGB def _load_uastc(data: bytes, width: int, height: int, level_count: int, scheme: int, target: str) -> KTX2Texture: """Transcode a single-source UASTC LDR KTX2 to a ``target`` :class:`KTX2Texture`. Detects UASTC vs ETC1S via the DFD colorModel. ETC1S/BasisLZ is out of scope and raises. UASTC levels are transcoded to ``target`` (``"bc7"`` / ``"astc4x4"`` / ``"etc2"``, all 16 B/block) via the native basis_universal transcoder; if that extension is unavailable, raises with an actionable message so the caller degrades to warn + -1 (never a crash). The loader is device-agnostic: ``TextureManager`` decides which ``target`` is legal on this GPU and threads it in. The (target, is_srgb) pair selects the VkFormat name; the matching native fn does the transcode. """ color_model, is_srgb = _dfd_color_model_and_srgb(data) if color_model == _KHR_DF_MODEL_ETC1S: raise ValueError("KTX2 ETC1S/BasisLZ needs Basis Universal transcoding (out of scope)") if color_model != _KHR_DF_MODEL_UASTC: raise ValueError(f"KTX2 vkFormat 0 unknown DFD colorModel {color_model} (expected UASTC 166)") if scheme == _SC_BASISLZ: # BasisLZ supercompression is only ever ETC1S; a UASTC DFD with BasisLZ is malformed. raise ValueError("KTX2 UASTC with BasisLZ supercompression is invalid") if target not in _UASTC_TARGET_FORMATS: raise ValueError(f"KTX2 UASTC unknown transcode target {target!r} (expected bc7/astc4x4/etc2)") from .._native import basis_transcoder # noqa: PLC0415 (lazy: native ext, only on a UASTC file) if not basis_transcoder.is_available(): basis_transcoder.ensure_built() if not basis_transcoder.is_available(): raise ValueError( "KTX2 UASTC needs the native basis transcoder: run `simvx build-textures` with a " "C++ compiler, or pre-transcode to BC7 with toktx" ) levels = level_count or 1 raw_levels = _read_levels(data, levels, scheme) import vulkan as vk # noqa: PLC0415 (call-time only, never at import) unorm_name, srgb_name = _UASTC_TARGET_FORMATS[target] vk_format = int(getattr(vk, srgb_name if is_srgb else unorm_name)) transcode = getattr(basis_transcoder, _UASTC_TARGET_FN[target]) mips: list[bytes] = [] for i, uastc_level in enumerate(raw_levels): lw = max(1, width >> i) lh = max(1, height >> i) mips.append(transcode(uastc_level, lw, lh)) log.debug("KTX2 UASTC->%s: %dx%d srgb=%s mips=%d -> fmt=%d", target, width, height, is_srgb, len(mips), vk_format) return KTX2Texture(vk_format=vk_format, width=width, height=height, block_size=16, mips=mips)