from __future__ import annotations from contextlib import contextmanager from dataclasses import asdict, dataclass import json import math from pathlib import Path from typing import Any, Mapping from runtime_compat import apply_runtime_compat apply_runtime_compat() PACKAGE_ROOT = Path(__file__).resolve().parents[1] # the demo root (data/ sits directly under it) import torch import torch.nn.functional as F DEFAULT_CODEC_FAMILY = "encodec" DEFAULT_CODEC_MODEL_ID = "facebook/encodec_32khz" DEFAULT_DAC_CODEC_MODEL_ID = "descript/dac_44khz" DEFAULT_SHAME_CODEC_MODEL_ID = "stabilityai/SAME-L" DEFAULT_SAO_CODEC_MODEL_ID = "stabilityai/stable-audio-open-1.0" DEFAULT_ENCODEC_BANDWIDTH = 2.2 SAO_SAMPLE_RATE = 44100 SAO_HOP_LENGTH = 2048 SAO_LATENT_DIM = 64 # The Stable Audio Open 1.0 VAE (its pretransform autoencoder, vae_model.ckpt). # It is stereo; mono audio is encoded as dual mono and decoded to the channel mean. SAO_VAE_CONFIG: dict[str, Any] = { "model_type": "autoencoder", "sample_size": 2_097_152, "sample_rate": SAO_SAMPLE_RATE, "audio_channels": 2, "model": { "encoder": { "type": "oobleck", "requires_grad": False, "config": { "in_channels": 2, "channels": 128, "c_mults": [1, 2, 4, 8, 16], "strides": [2, 4, 4, 8, 8], "latent_dim": 2 * SAO_LATENT_DIM, "use_snake": True, }, }, "decoder": { "type": "oobleck", "requires_grad": False, "config": { "out_channels": 2, "channels": 128, "c_mults": [1, 2, 4, 8, 16], "strides": [2, 4, 4, 8, 8], "latent_dim": SAO_LATENT_DIM, "use_snake": True, "final_tanh": False, }, }, "bottleneck": {"type": "vae"}, "latent_dim": SAO_LATENT_DIM, "downsampling_ratio": SAO_HOP_LENGTH, "io_channels": 2, }, } LEGACY_TARGET_LAYOUT = "framewise_sum_t128" DEFAULT_TARGET_LAYOUT = "framewise_sum" PCA_TARGET_LAYOUT = "framewise_pca" PRESNAP_TARGET_LAYOUT = "dac_projected_latents_presnap" PRESNAP_LATENT_SPACE = PRESNAP_TARGET_LAYOUT @dataclass(frozen=True) class AudioCodecMetadata: codec_family: str codec_model_id: str codec_sample_rate: int codec_audio_channels: int codec_frame_rate: float codec_codebook_size: int codec_num_codebooks: int codec_target_dim: int encodec_bandwidth: float | None = None dac_num_quantizers: int | None = None codec_hop_length: int | None = None def to_dict(self) -> dict[str, Any]: payload = asdict(self) if self.encodec_bandwidth is None: payload.pop("encodec_bandwidth", None) if self.dac_num_quantizers is None: payload.pop("dac_num_quantizers", None) if self.codec_hop_length is None: payload.pop("codec_hop_length", None) return payload def _legacy_dac_hop_length(mapping: Mapping[str, Any]) -> int | None: codec_family = infer_codec_family( codec_family=None if mapping.get("codec_family") is None else str(mapping.get("codec_family")), codec_model_id=None if mapping.get("codec_model_id") is None else str(mapping.get("codec_model_id")), ) codec_model_id = str(mapping.get("codec_model_id", "")).strip().lower() if codec_family == "dac" and codec_model_id == DEFAULT_DAC_CODEC_MODEL_ID: sample_rate = int(mapping.get("codec_sample_rate", mapping.get("sample_rate", 0)) or 0) if int(sample_rate) == 44100: return 512 return None def _resolve_codec_hop_length_from_mapping(mapping: Mapping[str, Any]) -> int | None: value = mapping.get("codec_hop_length", mapping.get("hop_length")) try: hop = int(value) except (TypeError, ValueError): hop = int(_legacy_dac_hop_length(mapping) or 0) if int(hop) <= 0: return None return int(hop) def resolve_device(device: str = "auto") -> str: device = str(device).strip().lower() if device in {"", "auto"}: return "cuda:0" if torch.cuda.is_available() else "cpu" return str(device) @contextmanager def _temporary_full_context_encodec(codec_model: Any): config = getattr(codec_model, "config", None) if config is None: yield return old_chunk_length_s = getattr(config, "chunk_length_s", None) old_overlap = getattr(config, "overlap", None) try: if hasattr(config, "chunk_length_s"): config.chunk_length_s = None if hasattr(config, "overlap"): config.overlap = None yield finally: if hasattr(config, "chunk_length_s"): config.chunk_length_s = old_chunk_length_s if hasattr(config, "overlap"): config.overlap = old_overlap def infer_codec_family(*, codec_family: str | None = None, codec_model_id: str | None = None) -> str: family_eff = str(codec_family or "").strip().lower() if family_eff in {"encodec", "dac", "shame", "sao"}: return family_eff model_id = str(codec_model_id or "").strip().lower() if model_id.startswith("stabilityai/same") or model_id == "same-l": return "shame" if "stable-audio-open" in model_id: return "sao" if "dac" in model_id: return "dac" return "encodec" def _read_json(path: Path) -> dict[str, Any]: return dict(json.loads(path.read_text(encoding="utf-8"))) def _maybe_read_json(path: Path) -> dict[str, Any] | None: if not path.is_file(): return None try: return _read_json(path) except Exception: return None def _audio_codec_meta_from_mapping(mapping: Mapping[str, Any] | None) -> AudioCodecMetadata | None: if mapping is None: return None if "codec_family" not in mapping and "encodec_model_id" in mapping: return AudioCodecMetadata( codec_family="encodec", codec_model_id=str(mapping["encodec_model_id"]).strip(), codec_sample_rate=int(mapping.get("codec_sample_rate", mapping.get("sample_rate", 32000))), codec_audio_channels=int(mapping.get("codec_audio_channels", mapping.get("audio_channels", 1))), codec_frame_rate=float(mapping.get("codec_frame_rate", mapping.get("frame_rate", 50.0))), codec_codebook_size=int(mapping.get("codec_codebook_size", mapping.get("codebook_size", 2048))), codec_num_codebooks=int(mapping.get("codec_num_codebooks", mapping.get("num_codebooks", 4))), codec_target_dim=int(mapping.get("codec_target_dim", mapping.get("target_dim", 128))), encodec_bandwidth=float(mapping.get("encodec_bandwidth", mapping.get("bandwidth", DEFAULT_ENCODEC_BANDWIDTH))), codec_hop_length=_resolve_codec_hop_length_from_mapping(mapping), ) if "codec_family" not in mapping or "codec_model_id" not in mapping: return None codec_hop_length = _resolve_codec_hop_length_from_mapping(mapping) return AudioCodecMetadata( codec_family=str(mapping["codec_family"]).strip().lower(), codec_model_id=str(mapping["codec_model_id"]).strip(), codec_sample_rate=int(mapping["codec_sample_rate"]), codec_audio_channels=int(mapping["codec_audio_channels"]), codec_frame_rate=float(mapping["codec_frame_rate"]), codec_codebook_size=int(mapping["codec_codebook_size"]), codec_num_codebooks=int(mapping["codec_num_codebooks"]), codec_target_dim=int(mapping["codec_target_dim"]), encodec_bandwidth=( None if mapping.get("encodec_bandwidth") is None else float(mapping["encodec_bandwidth"]) ), dac_num_quantizers=( None if mapping.get("dac_num_quantizers") is None else int(mapping["dac_num_quantizers"]) ), codec_hop_length=codec_hop_length, ) def resolve_codec_num_codebooks_for_encodec( *, frame_rate: float, codebook_size: int, bandwidth: float, max_num_codebooks: int, ) -> int: bits_per_codebook = math.log2(float(codebook_size)) raw = int(math.floor((float(bandwidth) * 1000.0) / (float(bits_per_codebook) * float(frame_rate)))) return int(max(1, min(int(max_num_codebooks), int(raw)))) def _resolve_encodec_codebook_size(model: Any) -> int: layers = list(getattr(getattr(model, "quantizer", None), "layers", []) or []) if not layers: raise RuntimeError("encodec model has no quantizer layers") embed = torch.as_tensor(layers[0].codebook.embed) return int(embed.shape[0]) def _resolve_encodec_target_dim(model: Any) -> int: layers = list(getattr(getattr(model, "quantizer", None), "layers", []) or []) if not layers: raise RuntimeError("encodec model has no quantizer layers") embed = torch.as_tensor(layers[0].codebook.embed) return int(embed.shape[-1]) def _build_encodec_metadata( model: Any, *, codec_model_id: str, bandwidth: float | None = None, ) -> AudioCodecMetadata: config = getattr(model, "config", None) if config is None: raise RuntimeError("encodec model is missing config") sample_rate = int(getattr(config, "sampling_rate")) audio_channels = int(getattr(config, "audio_channels", 1)) if int(audio_channels) != 1: raise ValueError( f"phase-1 codec support is mono-only; {codec_model_id} reports audio_channels={audio_channels}" ) frame_rate = float(getattr(config, "frame_rate")) codebook_size = int(getattr(config, "codebook_size", _resolve_encodec_codebook_size(model))) max_num_codebooks = int(len(list(getattr(getattr(model, "quantizer", None), "layers", []) or []))) if int(max_num_codebooks) <= 0: raise RuntimeError("encodec model has no quantizer layers") target_bandwidths = list(getattr(config, "target_bandwidths", []) or []) bandwidth_eff = ( float(bandwidth) if bandwidth is not None else float(target_bandwidths[-1] if target_bandwidths else DEFAULT_ENCODEC_BANDWIDTH) ) num_codebooks = resolve_codec_num_codebooks_for_encodec( frame_rate=float(frame_rate), codebook_size=int(codebook_size), bandwidth=float(bandwidth_eff), max_num_codebooks=int(max_num_codebooks), ) return AudioCodecMetadata( codec_family="encodec", codec_model_id=str(codec_model_id), codec_sample_rate=int(sample_rate), codec_audio_channels=int(audio_channels), codec_frame_rate=float(frame_rate), codec_codebook_size=int(codebook_size), codec_num_codebooks=int(num_codebooks), codec_target_dim=int(_resolve_encodec_target_dim(model)), encodec_bandwidth=float(bandwidth_eff), ) def _build_dac_metadata( model: Any, *, codec_model_id: str, num_quantizers: int | None = None, ) -> AudioCodecMetadata: config = getattr(model, "config", None) if config is None: raise RuntimeError("dac model is missing config") sample_rate = int(getattr(config, "sampling_rate")) audio_channels = int(getattr(config, "audio_channels", 1) or 1) if int(audio_channels) != 1: raise ValueError( f"phase-1 codec support is mono-only; {codec_model_id} reports audio_channels={audio_channels}" ) frame_rate = float(getattr(config, "frame_rate")) hop_length = int(getattr(config, "hop_length", 0) or 0) if int(hop_length) <= 0 and str(codec_model_id).strip().lower() == DEFAULT_DAC_CODEC_MODEL_ID: if int(sample_rate) == 44100: hop_length = 512 if int(hop_length) > 0: frame_rate = float(sample_rate) / float(hop_length) max_codebooks = int(getattr(config, "n_codebooks")) num_quantizers_eff = int(max(1, min(int(max_codebooks), int(num_quantizers or max_codebooks)))) return AudioCodecMetadata( codec_family="dac", codec_model_id=str(codec_model_id), codec_sample_rate=int(sample_rate), codec_audio_channels=int(audio_channels), codec_frame_rate=float(frame_rate), codec_codebook_size=int(getattr(config, "codebook_size")), codec_num_codebooks=int(num_quantizers_eff), codec_target_dim=int(getattr(config, "hidden_size")), dac_num_quantizers=int(num_quantizers_eff), codec_hop_length=int(hop_length) if int(hop_length) > 0 else None, ) def sao_codec_metadata(codec_model_id: str = DEFAULT_SAO_CODEC_MODEL_ID) -> AudioCodecMetadata: return AudioCodecMetadata( codec_family="sao", codec_model_id=str(codec_model_id), codec_sample_rate=int(SAO_SAMPLE_RATE), codec_audio_channels=1, codec_frame_rate=float(SAO_SAMPLE_RATE) / float(SAO_HOP_LENGTH), codec_codebook_size=0, codec_num_codebooks=0, codec_target_dim=int(SAO_LATENT_DIM), codec_hop_length=int(SAO_HOP_LENGTH), ) def _load_sao_vae(model_id: str, device: str) -> Any: try: from huggingface_hub import hf_hub_download from stable_audio_tools.models.factory import create_model_from_config from stable_audio_tools.models.utils import copy_state_dict, load_ckpt_state_dict except ModuleNotFoundError as exc: # pragma: no cover raise RuntimeError( "The SAO codec needs stable_audio_tools and huggingface_hub in the active Python environment." ) from exc try: checkpoint_path = hf_hub_download(model_id, "vae_model.ckpt", local_files_only=True) except Exception: checkpoint_path = hf_hub_download(model_id, "vae_model.ckpt") model = create_model_from_config(SAO_VAE_CONFIG) copy_state_dict(model, load_ckpt_state_dict(checkpoint_path)) model = model.to(device).eval() if (int(model.sample_rate), int(model.latent_dim), int(model.downsampling_ratio)) != ( int(SAO_SAMPLE_RATE), int(SAO_LATENT_DIM), int(SAO_HOP_LENGTH), ): raise RuntimeError("loaded SAO VAE does not match the expected 44.1 kHz / 64-d / 2048-hop contract") return model @torch.no_grad() def encode_audio_to_sao_latents(codec_model: Any, audio_bt: torch.Tensor) -> torch.Tensor: """Mono [B,N] audio -> SAO posterior means [B, ceil(N / hop), 64].""" audio = torch.as_tensor(audio_bt, dtype=torch.float32) if int(audio.dim()) != 2: raise ValueError(f"audio_bt must be [B,N], got {tuple(audio.shape)}") num_frames = int(math.ceil(int(audio.shape[-1]) / int(SAO_HOP_LENGTH))) stereo = F.pad(audio.unsqueeze(1).repeat(1, 2, 1), (0, num_frames * int(SAO_HOP_LENGTH) - int(audio.shape[-1]))) posterior = codec_model.encoder(stereo) return posterior.chunk(2, dim=1)[0][..., :num_frames].transpose(1, 2).contiguous() def attach_audio_codec_metadata(model: Any, metadata: AudioCodecMetadata) -> Any: setattr(model, "_pca_diffusion_codec_metadata", metadata.to_dict()) return model def get_audio_codec_metadata(model: Any) -> AudioCodecMetadata | None: attached = getattr(model, "_pca_diffusion_codec_metadata", None) meta = _audio_codec_meta_from_mapping(attached if isinstance(attached, Mapping) else None) if meta is not None: return meta config = getattr(model, "config", None) if config is None: return None model_type = str(getattr(config, "model_type", "")).strip().lower() if model_type == "dac": try: return _build_dac_metadata(model, codec_model_id="descript/dac_44khz") except Exception: return None try: return _build_encodec_metadata(model, codec_model_id=DEFAULT_CODEC_MODEL_ID) except Exception: return None def load_audio_codec_model( *, codec_family: str | None = None, codec_model_id: str | None = None, device: str = "auto", encodec_bandwidth: float | None = None, dac_num_quantizers: int | None = None, metadata: AudioCodecMetadata | Mapping[str, Any] | None = None, ) -> tuple[Any, str, AudioCodecMetadata]: metadata_mapping = dict(metadata) if isinstance(metadata, Mapping) else None if metadata is not None: meta_obj = metadata if isinstance(metadata, AudioCodecMetadata) else _audio_codec_meta_from_mapping(metadata) if meta_obj is None: raise ValueError("metadata must be AudioCodecMetadata or a mapping with codec_* keys") codec_family = str(meta_obj.codec_family) codec_model_id = str(meta_obj.codec_model_id) encodec_bandwidth = meta_obj.encodec_bandwidth dac_num_quantizers = meta_obj.dac_num_quantizers family = infer_codec_family(codec_family=codec_family, codec_model_id=codec_model_id) model_id = str( codec_model_id or ( DEFAULT_DAC_CODEC_MODEL_ID if family == "dac" else DEFAULT_SHAME_CODEC_MODEL_ID if family == "shame" else DEFAULT_SAO_CODEC_MODEL_ID if family == "sao" else DEFAULT_CODEC_MODEL_ID ) ).strip() model_path = Path(model_id).expanduser() if model_id and not model_path.is_absolute(): packaged_model_path = PACKAGE_ROOT / model_path if packaged_model_path.exists(): model_id = str(packaged_model_path) resolved_device = resolve_device(device) torch_device = torch.device(resolved_device) if torch_device.type == "cuda" and torch_device.index is not None: torch.cuda.set_device(torch_device) def _load_pretrained_model(model_cls: Any) -> Any: try: return model_cls.from_pretrained(model_id, local_files_only=True) except Exception as local_exc: try: return model_cls.from_pretrained(model_id) except Exception: raise local_exc if family == "encodec": try: from transformers import EncodecModel except Exception as exc: # pragma: no cover raise RuntimeError( "Could not import transformers.EncodecModel. Check the local transformers installation." ) from exc model = _load_pretrained_model(EncodecModel).to(resolved_device).eval() meta_obj = _build_encodec_metadata(model, codec_model_id=model_id, bandwidth=encodec_bandwidth) elif family == "dac": try: from transformers import DacModel except Exception as exc: # pragma: no cover raise RuntimeError( "Could not import transformers.DacModel. Check the local transformers installation." ) from exc model = _load_pretrained_model(DacModel).to(resolved_device).eval() meta_obj = _build_dac_metadata(model, codec_model_id=model_id, num_quantizers=dac_num_quantizers) elif family == "shame": try: from SHAME.shame_same import ( build_shame_codec_metadata, load_same_model, same_latent_dim, ) except ModuleNotFoundError as exc: # pragma: no cover raise RuntimeError( "SHAME/SAME codec metadata was requested, but stable_audio_tools is not available in this Python. " "Install stable_audio_tools in the active Python environment." ) from exc same_model_id = str((metadata_mapping or {}).get("same_model_id") or model_id or DEFAULT_SHAME_CODEC_MODEL_ID) try: model, same_model_config, resolved_device, same_use_half = load_same_model( model_id=same_model_id, device=resolved_device, half="auto", ) except ModuleNotFoundError as exc: # pragma: no cover if str(getattr(exc, "name", "")) != "stable_audio_tools": raise raise RuntimeError( "SHAME/SAME codec metadata was requested, but stable_audio_tools is not available in this Python. " "Install stable_audio_tools in the active Python environment." ) from exc if metadata is not None: meta_obj = metadata if isinstance(metadata, AudioCodecMetadata) else _audio_codec_meta_from_mapping(metadata) if meta_obj is None: raise ValueError("metadata must be AudioCodecMetadata or a mapping with codec_* keys") loaded_dim = int(same_latent_dim(model, same_model_config, fallback=int(meta_obj.codec_target_dim))) if int(meta_obj.codec_target_dim) != int(loaded_dim): raise ValueError( f"SAME latent dim mismatch: metadata={int(meta_obj.codec_target_dim)} loaded_model={int(loaded_dim)}" ) else: nested_model_cfg = same_model_config.get("model") if isinstance(same_model_config, Mapping) else None nested_latent_dim = ( int(nested_model_cfg["latent_dim"]) if isinstance(nested_model_cfg, Mapping) and nested_model_cfg.get("latent_dim") is not None else None ) latent_dim = int(same_latent_dim(model, same_model_config, fallback=nested_latent_dim)) meta_obj = _audio_codec_meta_from_mapping( build_shame_codec_metadata( model=model, model_config=same_model_config, model_id=same_model_id, latent_dim=int(latent_dim), ) ) if meta_obj is None: raise RuntimeError("could not build SHAME codec metadata") setattr(model, "_pca_diffusion_same_model_config", dict(same_model_config or {})) setattr(model, "_pca_diffusion_same_device", str(resolved_device)) setattr(model, "_pca_diffusion_same_use_half", bool(same_use_half)) elif family == "sao": model = _load_sao_vae(model_id, resolved_device) meta_obj = sao_codec_metadata(model_id) else: # pragma: no cover raise ValueError(f"unsupported codec_family={family!r}") for param in model.parameters(): param.requires_grad_(False) attach_audio_codec_metadata(model, meta_obj) return model, resolved_device, meta_obj def load_encodec_model(model_id: str = DEFAULT_CODEC_MODEL_ID, device: str = "auto") -> tuple[Any, str]: model, resolved_device, _metadata = load_audio_codec_model( codec_family="encodec", codec_model_id=str(model_id), device=device, ) return model, resolved_device def _resolve_codec_family_from_model(model: Any, metadata: AudioCodecMetadata | None = None) -> str: if metadata is not None: return str(metadata.codec_family) config = getattr(model, "config", None) if str(getattr(config, "model_type", "")).strip().lower() == "dac": return "dac" if callable(getattr(model, "decode_audio", None)) and callable(getattr(model, "encode_audio", None)): return "shame" return "encodec" def resolve_audio_codec_sample_rate(model: Any, default: int = 32000) -> int: meta_obj = get_audio_codec_metadata(model) if meta_obj is not None and int(meta_obj.codec_sample_rate) > 0: return int(meta_obj.codec_sample_rate) config = getattr(model, "config", None) value = getattr(config, "sampling_rate", None) if value is None: return int(default) return int(value) def extract_codebook_embeddings( codec_model: Any, *, device: torch.device | str, dtype: torch.dtype = torch.float32, metadata: AudioCodecMetadata | Mapping[str, Any] | None = None, ) -> torch.Tensor: meta_obj = ( metadata if isinstance(metadata, AudioCodecMetadata) else _audio_codec_meta_from_mapping(metadata if isinstance(metadata, Mapping) else None) ) or get_audio_codec_metadata(codec_model) family = _resolve_codec_family_from_model(codec_model, meta_obj) if family in {"shame", "sao"}: raise ValueError(f"{family} is a continuous latent codec and does not expose RVQ codebook embeddings") if family == "dac": quantizers = list(getattr(getattr(codec_model, "quantizer", None), "quantizers", []) or []) if not quantizers: raise RuntimeError("dac model has no quantizer modules") num_codebooks = int(meta_obj.codec_num_codebooks if meta_obj is not None else len(quantizers)) embeds: list[torch.Tensor] = [] for quantizer in list(quantizers[: int(num_codebooks)]): base_embed = torch.as_tensor( quantizer.codebook.weight.detach(), dtype=dtype, device=device, ).transpose(0, 1).unsqueeze(0) projected = quantizer.out_proj(base_embed).squeeze(0).transpose(0, 1).contiguous() embeds.append(projected) return torch.stack(embeds, dim=0).contiguous() layers = list(getattr(getattr(codec_model, "quantizer", None), "layers", []) or []) if not layers: raise RuntimeError("encodec model has no quantizer layers") num_codebooks = int(meta_obj.codec_num_codebooks if meta_obj is not None else len(layers)) embeds = [ layer.codebook.embed.detach().to(device=device, dtype=dtype).contiguous() for layer in list(layers[: int(num_codebooks)]) ] return torch.stack(embeds, dim=0).contiguous() def extract_dac_latent_subspace_basis( codec_model: Any, *, device: torch.device | str, dtype: torch.dtype = torch.float32, metadata: AudioCodecMetadata | Mapping[str, Any] | None = None, ) -> torch.Tensor: diagnostics = compute_dac_latent_subspace_diagnostics( codec_model, device=device, dtype=dtype, metadata=metadata, ) basis_dq = torch.as_tensor(diagnostics["basis_dq"], dtype=dtype, device=device) rank = int(diagnostics["rank"]) if int(rank) <= 0: raise RuntimeError("DAC latent subspace basis rank is zero") return basis_dq[:, : int(rank)].contiguous() def compute_dac_latent_subspace_diagnostics( codec_model: Any, *, device: torch.device | str, dtype: torch.dtype = torch.float32, metadata: AudioCodecMetadata | Mapping[str, Any] | None = None, ) -> dict[str, Any]: meta_obj = ( metadata if isinstance(metadata, AudioCodecMetadata) else _audio_codec_meta_from_mapping(metadata if isinstance(metadata, Mapping) else None) ) or get_audio_codec_metadata(codec_model) family = _resolve_codec_family_from_model(codec_model, meta_obj) if family != "dac": raise ValueError(f"DAC latent subspace basis requires codec_family='dac', got {family!r}") quantizers = list(getattr(getattr(codec_model, "quantizer", None), "quantizers", []) or []) if not quantizers: raise RuntimeError("dac model has no quantizer modules") num_codebooks = int(meta_obj.codec_num_codebooks if meta_obj is not None else len(quantizers)) mats = [] for quantizer in list(quantizers[: int(num_codebooks)]): weight = torch.as_tensor( quantizer.out_proj.weight.detach(), dtype=dtype, device=device, ).squeeze(-1) if int(weight.dim()) != 2: raise RuntimeError(f"unexpected DAC out_proj weight shape: {tuple(weight.shape)}") mats.append(weight.contiguous()) stacked = torch.cat(mats, dim=1).contiguous() stacked_f32 = stacked.to(dtype=torch.float32) singular_values = torch.linalg.svdvals(stacked_f32) max_sv = float(singular_values.max().item()) if int(singular_values.numel()) > 0 else 0.0 rank_tol = float(max_sv * float(max(stacked_f32.shape)) * float(torch.finfo(torch.float32).eps)) rank = int(torch.count_nonzero(singular_values > float(rank_tol)).item()) if int(rank) <= 0: raise RuntimeError("DAC latent subspace basis rank is zero") basis_dq, _r = torch.linalg.qr(stacked, mode="reduced") first_discarded = ( float(singular_values[int(rank)].item()) if int(rank) < int(singular_values.numel()) else 0.0 ) return { "basis_dq": basis_dq[:, : int(rank)].contiguous(), "rank": int(rank), "rank_tolerance": float(rank_tol), "matrix_shape": [int(dim) for dim in stacked_f32.shape], "singular_value_max": float(max_sv), "singular_value_min_retained": float(singular_values[int(rank) - 1].item()), "singular_value_first_discarded": float(first_discarded), } def _flatten_encodec_audio_codes(audio_codes_fbqt: torch.Tensor) -> torch.Tensor: if int(audio_codes_fbqt.dim()) != 4: raise RuntimeError(f"unexpected EnCodec code tensor shape: {tuple(audio_codes_fbqt.shape)}") return ( audio_codes_fbqt.permute(1, 2, 0, 3) .contiguous() .view(int(audio_codes_fbqt.shape[1]), int(audio_codes_fbqt.shape[2]), -1) ) def encode_audio_batch_to_codes( codec_model: Any, audio_bct: torch.Tensor, *, device: torch.device | str, metadata: AudioCodecMetadata | Mapping[str, Any] | None = None, ) -> torch.Tensor: audio = torch.as_tensor(audio_bct, dtype=torch.float32, device=device) if int(audio.dim()) == 2: audio = audio.unsqueeze(0) if int(audio.dim()) != 3: raise ValueError(f"audio_bct must be [B,C,T] or [C,T], got {tuple(audio.shape)}") meta_obj = ( metadata if isinstance(metadata, AudioCodecMetadata) else _audio_codec_meta_from_mapping(metadata if isinstance(metadata, Mapping) else None) ) or get_audio_codec_metadata(codec_model) family = _resolve_codec_family_from_model(codec_model, meta_obj) if family in {"shame", "sao"}: raise ValueError(f"{family} is a continuous latent codec and does not expose discrete audio codes") if family == "dac": encoded = codec_model.encode( audio, n_quantizers=None if meta_obj is None else meta_obj.dac_num_quantizers, ) audio_codes = getattr(encoded, "audio_codes", None) if audio_codes is None and isinstance(encoded, tuple) and len(encoded) >= 3: audio_codes = encoded[2] if audio_codes is None: raise RuntimeError("DAC encode did not return audio_codes") audio_codes = torch.as_tensor(audio_codes, dtype=torch.long, device=device) if int(audio_codes.dim()) != 3: raise RuntimeError(f"unexpected DAC code tensor shape: {tuple(audio_codes.shape)}") return audio_codes.contiguous() padding_mask = torch.ones_like(audio, dtype=torch.bool, device=device) with _temporary_full_context_encodec(codec_model): encoded = codec_model.encode( audio, padding_mask=padding_mask, bandwidth=None if meta_obj is None else meta_obj.encodec_bandwidth, ) audio_codes = getattr(encoded, "audio_codes", None) if audio_codes is None and isinstance(encoded, tuple) and len(encoded) >= 1: audio_codes = encoded[0] if audio_codes is None: raise RuntimeError("Encodec encode did not return audio_codes") audio_codes = torch.as_tensor(audio_codes, dtype=torch.long, device=device) return _flatten_encodec_audio_codes(audio_codes).contiguous() def encode_audio_chunks_to_codes( codec_model: Any, chunks_ct: list[torch.Tensor], *, device: torch.device | str, metadata: AudioCodecMetadata | Mapping[str, Any] | None = None, ) -> list[torch.Tensor]: if not chunks_ct: return [] lengths = [int(torch.as_tensor(chunk).shape[-1]) for chunk in chunks_ct] expected_len = int(lengths[0]) if any(int(length) != int(expected_len) for length in lengths): raise ValueError("all chunks_ct must share the same sample length") batch = torch.stack( [torch.as_tensor(chunk, dtype=torch.float32) for chunk in chunks_ct], dim=0, ).contiguous() codes_bct = encode_audio_batch_to_codes( codec_model, batch, device=device, metadata=metadata, ) return [ codes_bct[int(idx)].detach().to(device="cpu", dtype=torch.long).contiguous() for idx in range(int(codes_bct.shape[0])) ] def token_ids_to_codebook_embeddings( token_ids_bct: torch.Tensor, codebook_embed_ckd: torch.Tensor, *, valid_bt: torch.Tensor | None = None, ) -> torch.Tensor: if int(token_ids_bct.dim()) != 3: raise ValueError(f"token_ids_bct must be [B,C,T], got {tuple(token_ids_bct.shape)}") if int(codebook_embed_ckd.dim()) != 3: raise ValueError(f"codebook_embed_ckd must be [C,K,D], got {tuple(codebook_embed_ckd.shape)}") batch_size, num_codebooks, _num_frames = [int(x) for x in list(token_ids_bct.shape)] if int(codebook_embed_ckd.shape[0]) != int(num_codebooks): raise ValueError( f"codebook embeddings must match codebook count, got {tuple(codebook_embed_ckd.shape)} " f"for token shape {tuple(token_ids_bct.shape)}" ) embeds = [] for codebook_idx in range(int(num_codebooks)): embed_kd = codebook_embed_ckd[int(codebook_idx)].to(device=token_ids_bct.device, dtype=torch.float32) token_idx_bt = token_ids_bct[:, int(codebook_idx), :].to(device=token_ids_bct.device, dtype=torch.long) token_idx_bt = token_idx_bt.clamp(min=0, max=int(embed_kd.shape[0]) - 1) embeds.append(F.embedding(token_idx_bt, embed_kd).unsqueeze(1)) out = torch.cat(embeds, dim=1).contiguous() if valid_bt is not None: valid_eff = valid_bt.to(device=out.device, dtype=torch.bool) out = out.masked_fill(~valid_eff[:, None, :, None], 0.0) return out.contiguous() def rvq_sum_latents( codebook_latents_bctd: torch.Tensor, *, valid_bt: torch.Tensor | None = None, ) -> torch.Tensor: if int(codebook_latents_bctd.dim()) != 4: raise ValueError( f"codebook_latents_bctd must be [B,C,T,D], got {tuple(codebook_latents_bctd.shape)}" ) summed = codebook_latents_bctd.to(dtype=torch.float32).sum(dim=1) if valid_bt is not None: summed = summed.masked_fill(~valid_bt.to(device=summed.device, dtype=torch.bool)[:, :, None], 0.0) return summed.contiguous() def summed_frame_latents_from_code_ids( token_ids_ct: torch.Tensor, codebook_embed_ckd: torch.Tensor, ) -> torch.Tensor: if int(token_ids_ct.dim()) != 2: raise ValueError(f"token_ids_ct must be [C,T], got {tuple(token_ids_ct.shape)}") token_ids_bct = token_ids_ct.unsqueeze(0).to(dtype=torch.long) valid_bt = torch.ones((1, int(token_ids_ct.shape[1])), device=token_ids_ct.device, dtype=torch.bool) latents_bctd = token_ids_to_codebook_embeddings( token_ids_bct, codebook_embed_ckd, valid_bt=valid_bt, ) return rvq_sum_latents(latents_bctd, valid_bt=valid_bt)[0].to(dtype=torch.float32).contiguous() def pooled_segment_latent_from_code_ids( token_ids_ct: torch.Tensor, codebook_embed_ckd: torch.Tensor, ) -> torch.Tensor: return summed_frame_latents_from_code_ids(token_ids_ct, codebook_embed_ckd).sum(dim=0).to(dtype=torch.float32) def _is_cudnn_engine_failure(exc: RuntimeError) -> bool: message = str(exc).lower() return ( "unable to find an engine" in message or "cudnn" in message and "engine" in message ) def _decode_dac_with_cudnn_fallback(codec_model: Any, **kwargs: Any) -> torch.Tensor: try: return codec_model.decode(**kwargs).audio_values except RuntimeError as exc: if not _is_cudnn_engine_failure(exc): raise with torch.backends.cudnn.flags(enabled=False): return codec_model.decode(**kwargs).audio_values def decode_codes_to_audio_b1t( codec_model: Any, codes_bct: torch.Tensor, *, device: torch.device | str, metadata: AudioCodecMetadata | Mapping[str, Any] | None = None, ) -> torch.Tensor: codes = torch.as_tensor(codes_bct, dtype=torch.long, device=device) if int(codes.dim()) == 2: codes = codes.unsqueeze(0) meta_obj = ( metadata if isinstance(metadata, AudioCodecMetadata) else _audio_codec_meta_from_mapping(metadata if isinstance(metadata, Mapping) else None) ) or get_audio_codec_metadata(codec_model) family = _resolve_codec_family_from_model(codec_model, meta_obj) if family == "dac": codebook_embed_ckd = extract_codebook_embeddings( codec_model, device=device, metadata=meta_obj, ) valid_bt = codes.ge(0).all(dim=1).to(dtype=torch.bool) latents_bctd = token_ids_to_codebook_embeddings( codes.clamp_min(0), codebook_embed_ckd, valid_bt=valid_bt, ) quantized_latent = rvq_sum_latents(latents_bctd, valid_bt=valid_bt) return decode_quantized_latent_to_audio(codec_model, quantized_latent) if family in {"shame", "sao"}: raise ValueError(f"{family} is a continuous latent codec and does not expose discrete audio codes") audio_codes = codes.unsqueeze(0).contiguous() decoded = codec_model.decode(audio_codes, [None], padding_mask=None).audio_values audio = decoded.to(device=device, dtype=torch.float32) if int(audio.dim()) == 2: audio = audio.unsqueeze(1) if int(audio.dim()) != 3: raise RuntimeError(f"unexpected decoded audio shape: {tuple(audio.shape)}") return audio.contiguous() def decode_quantized_latent_to_audio( codec_model: Any, latent_btd: torch.Tensor, ) -> torch.Tensor: latent = torch.as_tensor(latent_btd, dtype=torch.float32) if int(latent.dim()) != 3: raise ValueError(f"latent_btd must be [B,T,D], got {tuple(latent.shape)}") meta_obj = get_audio_codec_metadata(codec_model) family = _resolve_codec_family_from_model(codec_model, meta_obj) z_q = latent.transpose(1, 2).contiguous() if family == "sao": audio = codec_model.decode(z_q.to(device=next(codec_model.parameters()).device)) return audio.to(device=latent.device, dtype=torch.float32).mean(dim=1, keepdim=True).contiguous() if family == "dac": audio = _decode_dac_with_cudnn_fallback(codec_model, quantized_representation=z_q) audio = torch.as_tensor(audio, dtype=torch.float32, device=latent.device) if int(audio.dim()) == 2: audio = audio.unsqueeze(1) return audio.contiguous() if family == "shame": try: from SHAME.shame_same import decode_same_latent_to_audio except ModuleNotFoundError as exc: # pragma: no cover raise RuntimeError( "SHAME/SAME decode requires stable_audio_tools. " "Install stable_audio_tools in the active Python environment." ) from exc same_model_config = dict(getattr(codec_model, "_pca_diffusion_same_model_config", {}) or {}) same_device = str(getattr(codec_model, "_pca_diffusion_same_device", "") or "") if not same_device: try: same_device = str(next(codec_model.parameters()).device) except Exception: same_device = str(latent.device) same_use_half = bool(getattr(codec_model, "_pca_diffusion_same_use_half", False)) audio = decode_same_latent_to_audio( codec_model, latent, model_config=same_model_config, device=same_device, use_half=bool(same_use_half), ) return audio.to(device=latent.device, dtype=torch.float32).contiguous() return codec_model.decoder(z_q) def load_target_pca_basis( source: str | Path | Mapping[str, Any], *, device: torch.device | str = "cpu", dtype: torch.dtype = torch.float32, ) -> dict[str, Any]: if isinstance(source, (str, Path)): payload = dict( torch.load(Path(source).expanduser().resolve(), map_location="cpu", weights_only=False) ) elif isinstance(source, Mapping): payload = dict(source) else: raise TypeError(f"unsupported PCA basis source: {type(source).__name__}") mean = torch.as_tensor(payload["mean"], dtype=dtype, device=device).view(-1).contiguous() components = torch.as_tensor(payload["components"], dtype=dtype, device=device).contiguous() if int(components.dim()) != 2: raise ValueError(f"PCA basis components must be [K,D], got {tuple(components.shape)}") if int(components.shape[-1]) != int(mean.shape[0]): raise ValueError( f"PCA basis mean/components mismatch: mean={tuple(mean.shape)} components={tuple(components.shape)}" ) out: dict[str, Any] = {**payload} out["mean"] = mean out["components"] = components out["k"] = int(payload.get("k", int(components.shape[0]))) out["target_dim"] = int(payload.get("target_dim", int(components.shape[0]))) out["full_target_dim"] = int(payload.get("full_target_dim", int(components.shape[-1]))) out["target_layout"] = str(payload.get("target_layout", PCA_TARGET_LAYOUT)).strip().lower() return out def project_latent_to_pca( latent_btd: torch.Tensor, pca_basis: Mapping[str, Any], ) -> torch.Tensor: latent = torch.as_tensor(latent_btd, dtype=torch.float32) basis = load_target_pca_basis( pca_basis, device=latent.device, dtype=latent.dtype, ) mean = torch.as_tensor(basis["mean"], dtype=latent.dtype, device=latent.device) components = torch.as_tensor(basis["components"], dtype=latent.dtype, device=latent.device) if int(latent.shape[-1]) != int(components.shape[-1]): raise ValueError( f"latent/PCA basis mismatch: latent={tuple(latent.shape)} components={tuple(components.shape)}" ) return torch.matmul(latent - mean.view(*([1] * (int(latent.dim()) - 1)), -1), components.transpose(0, 1)) def reconstruct_latent_from_pca( latent_btd: torch.Tensor, pca_basis: Mapping[str, Any] | None, ) -> torch.Tensor: latent = torch.as_tensor(latent_btd, dtype=torch.float32) if pca_basis is None: return latent.contiguous() basis = load_target_pca_basis( pca_basis, device=latent.device, dtype=latent.dtype, ) mean = torch.as_tensor(basis["mean"], dtype=latent.dtype, device=latent.device) components = torch.as_tensor(basis["components"], dtype=latent.dtype, device=latent.device) if int(latent.shape[-1]) != int(components.shape[0]): raise ValueError( f"latent/PCA basis mismatch: latent={tuple(latent.shape)} components={tuple(components.shape)}" ) return ( torch.matmul(latent, components) + mean.view(*([1] * (int(latent.dim()) - 1)), -1) ).contiguous() @contextmanager def _temporary_full_context_encodec(codec_model: Any): config = getattr(codec_model, "config", None) if config is None: yield return old_chunk_length_s = getattr(config, "chunk_length_s", None) old_overlap = getattr(config, "overlap", None) if hasattr(config, "chunk_length_s"): config.chunk_length_s = None if hasattr(config, "overlap"): config.overlap = None try: yield finally: if hasattr(config, "chunk_length_s"): config.chunk_length_s = old_chunk_length_s if hasattr(config, "overlap"): config.overlap = old_overlap @torch.inference_mode() def encode_audio_to_codes_bct( codec_model: Any, audio_bct: torch.Tensor, *, device: torch.device | str, metadata: AudioCodecMetadata | Mapping[str, Any] | None = None, ) -> torch.Tensor: return encode_audio_batch_to_codes( codec_model, audio_bct, device=device, metadata=metadata, ) def is_presnap_latent_basis(basis: Mapping[str, Any] | None) -> bool: return isinstance(basis, Mapping) and str(basis.get("latent_space", "")) == PRESNAP_LATENT_SPACE def snap_presnap_latent_to_codes_bct(codec_model: Any, latent_btd: torch.Tensor) -> torch.Tensor: """Snap DAC projected_latents [B,T,Q*d] to RVQ codes, one codebook_dim chunk per quantizer.""" latent = torch.as_tensor(latent_btd, dtype=torch.float32) if int(latent.dim()) != 3: raise ValueError(f"latent_btd must be [B,T,D], got {tuple(latent.shape)}") latent_bdt = latent.transpose(1, 2).contiguous() codes: list[torch.Tensor] = [] offset = 0 for quantizer in list(getattr(getattr(codec_model, "quantizer", None), "quantizers", []) or []): width = int(quantizer.codebook_dim) if offset + width > int(latent_bdt.shape[1]): break _embed, indices = quantizer.decode_latents(latent_bdt[:, offset : offset + width, :]) codes.append(indices) offset += width if not codes or offset != int(latent_bdt.shape[1]): raise ValueError(f"pre-snap latent dim {int(latent_bdt.shape[1])} does not split into DAC codebook chunks") return torch.stack(codes, dim=1).to(dtype=torch.long).contiguous() def requantize_latent_to_codes_bct( codec_model: Any, latent_btd: torch.Tensor, *, device: torch.device | str, metadata: AudioCodecMetadata | Mapping[str, Any] | None = None, target_pca_basis: Mapping[str, Any] | None = None, target_layout: str | None = None, ) -> torch.Tensor: layout = str(target_layout or "").strip().lower() if layout == PRESNAP_TARGET_LAYOUT or is_presnap_latent_basis(target_pca_basis): return snap_presnap_latent_to_codes_bct( codec_model, torch.as_tensor(latent_btd, dtype=torch.float32, device=device), ) latent = reconstruct_latent_from_pca( torch.as_tensor(latent_btd, dtype=torch.float32, device=device), target_pca_basis, ) if int(latent.dim()) != 3: raise ValueError(f"latent_btd must be [B,T,D], got {tuple(latent.shape)}") meta_obj = ( metadata if isinstance(metadata, AudioCodecMetadata) else _audio_codec_meta_from_mapping(metadata if isinstance(metadata, Mapping) else None) ) or get_audio_codec_metadata(codec_model) family = _resolve_codec_family_from_model(codec_model, meta_obj) if family in {"shame", "sao"}: raise ValueError(f"{family} is a continuous latent codec and cannot be requantized to codes") z_q = latent.transpose(1, 2).contiguous() if family == "dac": _zq, audio_codes, _projected, _commitment, _codebook = codec_model.quantizer( z_q, None if meta_obj is None else meta_obj.dac_num_quantizers, ) if int(audio_codes.dim()) != 3: raise RuntimeError(f"unexpected DAC quantizer output shape: {tuple(audio_codes.shape)}") return audio_codes.contiguous().to(dtype=torch.long) bandwidth = None if meta_obj is None else meta_obj.encodec_bandwidth codes_cbt = codec_model.quantizer.encode(z_q, bandwidth=bandwidth) if int(codes_cbt.dim()) != 3: raise RuntimeError(f"unexpected quantizer.encode output shape: {tuple(codes_cbt.shape)}") if int(codes_cbt.shape[1]) != int(latent.shape[0]): raise RuntimeError( f"quantizer.encode batch mismatch: latent batch={int(latent.shape[0])}, codes={tuple(codes_cbt.shape)}" ) return codes_cbt.permute(1, 0, 2).contiguous().to(dtype=torch.long) def legacy_source_cache_codec_metadata() -> AudioCodecMetadata: return AudioCodecMetadata( codec_family="encodec", codec_model_id=DEFAULT_CODEC_MODEL_ID, codec_sample_rate=32000, codec_audio_channels=1, codec_frame_rate=50.0, codec_codebook_size=2048, codec_num_codebooks=4, codec_target_dim=128, encodec_bandwidth=DEFAULT_ENCODEC_BANDWIDTH, ) def resolve_codec_metadata_from_cache_config(cache_root: str | Path) -> AudioCodecMetadata: cache_root_path = Path(cache_root).expanduser().resolve() config_payload = _maybe_read_json(cache_root_path / "config.json") meta_obj = _audio_codec_meta_from_mapping(config_payload) if meta_obj is not None: return meta_obj return legacy_source_cache_codec_metadata() def resolve_codec_metadata_from_payload( payload: Mapping[str, Any], *, fallback: AudioCodecMetadata | None = None, ) -> AudioCodecMetadata: meta_obj = _audio_codec_meta_from_mapping(payload) if meta_obj is not None: return meta_obj if "codec_metadata" in payload and isinstance(payload["codec_metadata"], Mapping): meta_obj = _audio_codec_meta_from_mapping(payload["codec_metadata"]) if meta_obj is not None: return meta_obj if fallback is not None: return fallback return legacy_source_cache_codec_metadata() def normalize_target_payload( payload: Mapping[str, Any], *, fallback_target_dim: int | None = None, ) -> tuple[torch.Tensor, torch.Tensor, int]: target_layout = str( payload.get( "target_layout", PCA_TARGET_LAYOUT if payload.get("target_pc_tk") is not None else DEFAULT_TARGET_LAYOUT, ) or DEFAULT_TARGET_LAYOUT ).strip().lower() if target_layout in {PCA_TARGET_LAYOUT, PRESNAP_TARGET_LAYOUT}: target_pc_tk = payload.get("target_pc_tk") target_pc_pool_k = payload.get("target_pc_pool_k") if target_pc_tk is None: raise KeyError(f"cache payload is missing target_pc_tk for {target_layout} targets") target_pc_tk_t = torch.as_tensor(target_pc_tk, dtype=torch.float32).contiguous() if target_pc_pool_k is None: target_pc_pool_k_t = target_pc_tk_t.sum(dim=0).to(dtype=torch.float32).contiguous() else: target_pc_pool_k_t = torch.as_tensor(target_pc_pool_k, dtype=torch.float32).contiguous() if int(target_pc_tk_t.dim()) != 2: raise RuntimeError(f"expected target_pc_tk [T,K], got {tuple(target_pc_tk_t.shape)}") if int(target_pc_pool_k_t.dim()) != 1: raise RuntimeError(f"expected target_pc_pool_k [K], got {tuple(target_pc_pool_k_t.shape)}") target_dim = int(payload.get("target_dim", fallback_target_dim or int(target_pc_tk_t.shape[-1]))) if int(target_pc_tk_t.shape[-1]) != int(target_dim): raise RuntimeError( f"target_dim mismatch: target_pc_tk has {target_pc_tk_t.shape[-1]} dims vs target_dim={target_dim}" ) if int(target_pc_pool_k_t.shape[0]) != int(target_dim): raise RuntimeError( f"target_pc_pool_k mismatch: target_pc_pool_k has {target_pc_pool_k_t.shape[0]} dims vs target_dim={target_dim}" ) return target_pc_tk_t, target_pc_pool_k_t, int(target_dim) target_sum_td = payload.get("target_sum_td", payload.get("target_sum_t128")) target_sum_pool_d = payload.get("target_sum_pool_d", payload.get("target_sum_pool_128")) if target_sum_td is None or target_sum_pool_d is None: missing = [] if target_sum_td is None: missing.append("target_sum_td") if target_sum_pool_d is None: missing.append("target_sum_pool_d") raise KeyError(f"cache payload is missing target tensors {missing}") target_sum_td_t = torch.as_tensor(target_sum_td, dtype=torch.float32).contiguous() target_sum_pool_d_t = torch.as_tensor(target_sum_pool_d, dtype=torch.float32).contiguous() if int(target_sum_td_t.dim()) != 2: raise RuntimeError(f"expected target_sum_td [T,D], got {tuple(target_sum_td_t.shape)}") if int(target_sum_pool_d_t.dim()) != 1: raise RuntimeError(f"expected target_sum_pool_d [D], got {tuple(target_sum_pool_d_t.shape)}") target_dim = int(payload.get("target_dim", fallback_target_dim or int(target_sum_td_t.shape[-1]))) if int(target_sum_td_t.shape[-1]) != int(target_dim): raise RuntimeError( f"target_dim mismatch: target_sum_td has {target_sum_td_t.shape[-1]} dims vs target_dim={target_dim}" ) if int(target_sum_pool_d_t.shape[0]) != int(target_dim): raise RuntimeError( f"target_sum_pool_d mismatch: target_sum_pool_d has {target_sum_pool_d_t.shape[0]} dims vs target_dim={target_dim}" ) return target_sum_td_t, target_sum_pool_d_t, int(target_dim) def resolve_target_dim_from_cache_config(cache_root: str | Path, default: int = 128) -> int: cache_root_path = Path(cache_root).expanduser().resolve() config_payload = _maybe_read_json(cache_root_path / "config.json") or {} if "target_dim" in config_payload: return int(config_payload["target_dim"]) return int(default) def resolve_target_layout_from_cache_config( cache_root: str | Path, default: str = DEFAULT_TARGET_LAYOUT, ) -> str: cache_root_path = Path(cache_root).expanduser().resolve() config_payload = _maybe_read_json(cache_root_path / "config.json") or {} if str(config_payload.get("latent_space") or "").strip().lower() == PRESNAP_LATENT_SPACE: return PRESNAP_TARGET_LAYOUT layout = str(config_payload.get("target_layout") or "").strip().lower() return layout or str(default).strip().lower() def resolve_target_pca_basis_path_from_cache_config(cache_root: str | Path) -> Path | None: cache_root_path = Path(cache_root).expanduser().resolve() config_payload = _maybe_read_json(cache_root_path / "config.json") or {} if str(config_payload.get("latent_space") or "").strip().lower() == PRESNAP_LATENT_SPACE: return None rel_path = str(config_payload.get("pca_basis_path") or "").strip() if rel_path: candidate = (cache_root_path / rel_path).resolve() if candidate.is_file(): return candidate fallback = (cache_root_path / "pca_basis.pt").resolve() if fallback.is_file(): return fallback return None