Upload zpcodec/model.py with huggingface_hub

41a06ac verified 5 months ago

15.3 kB

"""
ZPCodec: full codec model combining encoder, RVQ, optional repair, and decoder.
Data flow:
    waveform [B, 1, T]
    -> ZPEncoder          -> latent z [B, D, T']
    -> ResidualVQ         -> quantized z_q [B, D, T'], indices, commit_loss
    -> (GE simulator)     -> frame_mask [B, T']   (training only, if use_repair=True)
    -> LatentRepairTransformer -> z_q_post [B, D, T']  (missing frames concealed)
    -> ZPDecoder          -> waveform [B, 1, T_out]
T' = T / hop_length  (hop_length = prod(ratios) = 240 for ratios=[8,5,3,2] -> 15ms/frame)
The repair module is optional (use_repair=False for stage 1 codec-only training).
The GE simulator is optional too: if no GilbertElliottConfig is provided, no
packet loss is simulated and frame_mask is never generated automatically.
"""

import typing as tp
from contextlib import contextmanager

import numpy as np
import torch
import torch.nn as nn
from vector_quantize_pytorch import ResidualVQ

from .components import ZPEncoder, ZPDecoder
from .repair import LatentRepairTransformer
from .GilbertElliot import GilbertElliottConfig, GilbertElliottSimulator


@contextmanager
def temporarily_set(obj, attr: str, value):
    """Context manager that sets obj.attr = value for the duration of the block,
    then restores the original value. Used to toggle quantize_dropout per-batch."""
    original = getattr(obj, attr)
    setattr(obj, attr, value)
    try:
        yield
    finally:
        setattr(obj, attr, original)


class ZPCodec(nn.Module):
    """
    Full codec: encoder -> RVQ -> (repair) -> decoder.
    Stage 1 (codec pre-training): use_repair=False, no GilbertElliottConfig.
    Stage 2 (repair training):    use_repair=True, GilbertElliottConfig provided.
    Stage 3 (joint fine-tuning):  use_repair=True, GE curriculum via set_gilbert_elliott_config().
    """
    def __init__(
        self,
        channels: int = 1,
        dimension: int = 128,
        n_filters: int = 32,
        ratios: tp.List[int] = [8, 5, 3, 2],
        norm: str = 'weight_norm',
        causal: bool = True,
        num_quantizers: int = 9,
        codebook_size: int = 1024,
        sample_rate: int = 16000,
        # --- Repair module ---
        use_repair: bool = False,
        repair_hidden_dim: int = 256,
        repair_num_layers: int = 4,
        repair_num_heads: int = 4,
        repair_ffn_mult: int = 2,
        repair_past: int = 8,
        repair_future: int = 2,
        repair_two_pass: bool = True,
        # --- Packet loss simulation ---
        gilbert_elliott_config: tp.Optional[GilbertElliottConfig] = None,
    ):
        super().__init__()
        self.encoder = ZPEncoder(
            channels=channels,
            dimension=dimension,
            n_filters=n_filters,
            ratios=ratios,
            norm=norm,
            causal=causal,
        )
        self.rvq = ResidualVQ(
            dim=dimension,
            num_quantizers=num_quantizers,
            codebook_size=codebook_size,
            kmeans_init=True,
            kmeans_iters=10, 
            use_cosine_sim=True, # prop to improved RVQGAN's paper 
            threshold_ema_dead_code=2,
            quantize_dropout=True,
            quantize_dropout_cutoff_index=5,   # first 5 quantizers are always active -
            # theoretically with 5 quant active we can switch to 3kbps. But this was not my focus for that project...
            quantize_dropout_multiple_of=1,
        )
        self.decoder = ZPDecoder(
            channels=channels,
            dimension=dimension,
            n_filters=n_filters,
            ratios=ratios,
            norm=norm,
            causal=causal,
        )

        self.sample_rate = sample_rate
        self.hop_length = int(np.prod(ratios))   # 240 for ratios=[8,5,3,2]

        self.use_repair = use_repair
        self.repair_two_pass = repair_two_pass
        if use_repair:
            self.repair = LatentRepairTransformer(
                latent_dim=dimension,
                hidden_dim=repair_hidden_dim,
                num_layers=repair_num_layers,
                num_heads=repair_num_heads,
                ffn_mult=repair_ffn_mult,
                past=repair_past,
                future=repair_future,
            )
        else:
            self.repair = None

        self.ge_simulator: tp.Optional[GilbertElliottSimulator] = None
        if gilbert_elliott_config is not None:
            self.set_gilbert_elliott_config(gilbert_elliott_config)

    
    # Runtime configuration of the packet-loss simulator
    def set_gilbert_elliott_config(self, config: GilbertElliottConfig) -> None:
        """Replace the GE simulator at runtime. Called between training stages
        to apply a harder loss curriculum without reloading the model."""
        self.ge_simulator = GilbertElliottSimulator(
            config=config,
            sample_rate=self.sample_rate,
            hop_length=self.hop_length,
        )

    def sample_frame_mask(
        self,
        batch_size: int,
        num_frames: int,
        device: tp.Optional[torch.device] = None,
        seed: tp.Optional[int] = None,
    ) -> torch.Tensor:
        """Expose the GE simulator directly. Useful when the same mask needs to
        be reused across multiple points (e.g. logging, loss weighting)."""
        assert self.ge_simulator is not None, (
            "No GilbertElliottConfig configured. Call set_gilbert_elliott_config() first."
        )
        return self.ge_simulator.sample_frame_mask(
            batch_size, num_frames, device=device, seed=seed
        )

    # Encoding
    def _encode_raw(self, x: torch.Tensor):
        """Encode waveform to quantized latent. Returns (z, z_q, indices, commit_loss).
        quantize_dropout is randomly toggled per-call during training to teach
        the decoder to handle a variable number of active quantizers (bitrate scalability)."""
        z = self.encoder(x)                  # [B, D, T']
        z_seq = z.permute(0, 2, 1)           # [B, T', D]  — RVQ expects (B, T, D)
        use_dropout = self.training and (torch.rand(1).item() < 0.5) # dropout applied only 50% of the time, this improve the 
        # quality at full kbps. Citing the improved RVQGAN paper.
        with temporarily_set(self.rvq, 'quantize_dropout', use_dropout):
            z_q, indices, commit_loss = self.rvq(z_seq)
        z_q = z_q.permute(0, 2, 1)          # [B, D, T']
        return z, z_q, indices, commit_loss

    
    # Repair
    def _apply_repair(
        self,
        z_q: torch.Tensor,
        frame_mask: torch.Tensor,
    ) -> torch.Tensor:
        """Run the repair transformer and selectively substitute only missing frames.
        z_q:        [B, D, T']
        frame_mask: [B, T']   1 = received, 0 = missing
        The transformer outputs a full [B, D, T'] tensor, but received frames are
        kept as-is from z_q — only positions where frame_mask == 0 are replaced.
        This means z_q_post == z_q on received frames by construction, which is
        important for latent_repair_loss (the mask isolates the useful gradient).
        Two-pass mode (repair_two_pass=True): mimics streaming deployment where
        previous repair estimates are already in the buffer when estimating frame t.
        See LatentRepairTransformer.forward_two_pass for the full explanation.
        """
        assert self.repair is not None, "use_repair=False, repair not initialised"

        z_seq = z_q.permute(0, 2, 1)                      # [B, T', D]

        if self.repair_two_pass:
            z_repaired = self.repair.forward_two_pass(z_seq, frame_mask)
        else:
            # Single-pass fallback
            z_seq_filled = self.repair.fill_missing(z_seq, frame_mask)
            z_repaired = self.repair(z_seq_filled, frame_mask)

        # Selective substitution: keep received frames from z_q, replace missing ones
        m = frame_mask.unsqueeze(-1).to(z_seq.dtype)      # [B, T', 1]
        z_out = z_seq * m + z_repaired * (1.0 - m)
        return z_out.permute(0, 2, 1)                     # [B, D, T']

    def _get_frame_mask(
        self,
        z_q: torch.Tensor,
        frame_mask: tp.Optional[torch.Tensor],
    ) -> torch.Tensor:
        """Return the provided frame_mask, or sample one from the GE simulator."""
        if frame_mask is not None:
            return frame_mask
        assert self.ge_simulator is not None, (
            "use_repair=True but no GilbertElliottConfig configured. "
            "Call set_gilbert_elliott_config() before training."
        )
        B, _, T_prime = z_q.shape
        return self.ge_simulator.sample_frame_mask(B, T_prime, device=z_q.device)

    # Public encode / decode API
    def encode(self, x: torch.Tensor) -> tp.Tuple[torch.Tensor, torch.Tensor]:
        """Encode waveform to (z_q, indices). x: [B, 1, T]"""
        _, z_q, indices, _ = self._encode_raw(x)
        return z_q, indices

    def decode(
        self,
        z_q: torch.Tensor,
        frame_mask: tp.Optional[torch.Tensor] = None,
    ) -> torch.Tensor:
        """Decode quantized latent to waveform.
        z_q:        [B, D, T']
        frame_mask: [B, T']  optional; if provided and use_repair=True, runs repair first.
        """
        if self.use_repair and frame_mask is not None:
            z_q = self._apply_repair(z_q, frame_mask)
        return self.decoder(z_q)

    # Training forward
    def forward(
        self,
        x: torch.Tensor,
        frame_mask: tp.Optional[torch.Tensor] = None,
        return_intermediates: bool = False,
    ):
        """
        x:                    [B, 1, T]
        frame_mask:           [B, T'] optional. If use_repair=True and None,
                              sampled automatically from the GE simulator.
        return_intermediates: if True, also returns z_q pre/post repair and the
                              effective frame_mask — required by latent_repair_loss
                              and ZPCodecTrainer.forward_codec during training.
        Returns:
          return_intermediates=False:  (x_hat, commit_loss)
          return_intermediates=True:   (x_hat, commit_loss, z_q_pre, z_q_post, frame_mask)
              When use_repair=False: z_q_pre == z_q_post and frame_mask == None.
        """
        _, z_q_pre, _, commit_loss = self._encode_raw(x)
        commit_loss = commit_loss.mean()

        if self.use_repair:
            frame_mask = self._get_frame_mask(z_q_pre, frame_mask)
            z_q_post = self._apply_repair(z_q_pre, frame_mask)
        else:
            z_q_post = z_q_pre
            frame_mask = None

        x_hat = self.decoder(z_q_post)

        if return_intermediates:
            return x_hat, commit_loss, z_q_pre, z_q_post, frame_mask
        return x_hat, commit_loss

    # ------------------------------------------------------------------
    # from_pretrained — load from Hugging Face Hub or local path
    # ------------------------------------------------------------------
    @classmethod
    def from_pretrained(
        cls,
        model_id: str,
        device: str = "cpu",
        filename: str = "zpcodec_weights.pt",
        **hf_kwargs,
    ) -> "ZPCodec":
        """
        Load ZPCodec from a Hugging Face Hub repo or a local file.
        Args:
            model_id: HF repo id (e.g. "yourname/zpcodec") OR a local path
                      to a .pt file OR a local directory containing filename.
            device:   "cpu" | "cuda" | "cuda:0" etc.
            filename: name of the weights file inside the HF repo.
            **hf_kwargs: forwarded to huggingface_hub.hf_hub_download
                         (e.g. revision="main", token="hf_...").
        Returns:
            ZPCodec in eval mode.
        Examples:
            # From Hugging Face Hub
            model = ZPCodec.from_pretrained("yourname/zpcodec")
            # From a local .pt file
            model = ZPCodec.from_pretrained("./zpcodec_weights.pt")
            # With explicit device
            model = ZPCodec.from_pretrained("yourname/zpcodec", device="cuda")
        """
        import os
        import torch

        # Resolve checkpoint path: local file, local dir, or HF Hub
        if os.path.isfile(model_id):
            ckpt_path = model_id
        elif os.path.isdir(model_id):
            ckpt_path = os.path.join(model_id, filename)
            if not os.path.isfile(ckpt_path):
                raise FileNotFoundError(
                    f"No '{filename}' found in directory '{model_id}'"
                )
        else:
            # Treat as a Hugging Face Hub repo id
            try:
                from huggingface_hub import hf_hub_download
            except ImportError:
                raise ImportError(
                    "huggingface_hub is required to download from the Hub.\n"
                    "Install with:  pip install huggingface_hub"
                )
            ckpt_path = hf_hub_download(
                repo_id=model_id, filename=filename, **hf_kwargs
            )

        ckpt = torch.load(ckpt_path, map_location="cpu", weights_only=False)

        # Support both clean checkpoints (with 'config' key) and raw
        # full-trainer checkpoints (with 'args' key) for backward compat
        if "config" in ckpt:
            cfg = ckpt["config"]
            state_dict = ckpt["model_state_dict"]
        elif "args" in ckpt and "trainer" in ckpt:
            # Full trainer checkpoint — extract codec weights and config
            args = ckpt["args"]
            state_dict = {
                k[len("codec."):]: v
                for k, v in ckpt["trainer"].items()
                if k.startswith("codec.")
            }
            cfg = {
                "channels": 1, "dimension": args["dimension"],
                "n_filters": args["n_filters"], "ratios": [8, 5, 3, 2],
                "norm": "weight_norm", "causal": True,
                "num_quantizers": args["num_quantizers"],
                "codebook_size": args["codebook_size"], "sample_rate": 16000,
                "use_repair": True,
                "repair_hidden_dim": args["repair_hidden_dim"],
                "repair_num_layers": args["repair_num_layers"],
                "repair_num_heads": args["repair_num_heads"],
                "repair_ffn_mult": args["repair_ffn_mult"],
                "repair_past": args["repair_past"],
                "repair_future": args["repair_future"],
                "repair_two_pass": True,
            }
        else:
            raise ValueError(
                "Unrecognised checkpoint format. "
                "Expected keys: 'config'+'model_state_dict' or 'args'+'trainer'."
            )

        model = cls(**cfg)
        missing, unexpected = model.load_state_dict(state_dict, strict=True)
        if missing:
            raise RuntimeError(f"Missing keys: {missing[:5]}")
        if unexpected:
            raise RuntimeError(f"Unexpected keys: {unexpected[:5]}")

        n_params = sum(p.numel() for p in model.parameters()) / 1e6
        info = ckpt.get("training_info", {})
        stoi = info.get("best_val_stoi", ckpt.get("best_val_metric", "?"))
        print(f"✓ ZPCodec loaded — {n_params:.1f}M params | best val STOI: {stoi}")

        model = model.to(device)
        model.eval()
        return model

	"""
	ZPCodec: full codec model combining encoder, RVQ, optional repair, and decoder.

	Data flow:
	waveform [B, 1, T]
	-> ZPEncoder -> latent z [B, D, T']
	-> ResidualVQ -> quantized z_q [B, D, T'], indices, commit_loss
	-> (GE simulator) -> frame_mask [B, T'] (training only, if use_repair=True)
	-> LatentRepairTransformer -> z_q_post [B, D, T'] (missing frames concealed)
	-> ZPDecoder -> waveform [B, 1, T_out]

	T' = T / hop_length (hop_length = prod(ratios) = 240 for ratios=[8,5,3,2] -> 15ms/frame)

	The repair module is optional (use_repair=False for stage 1 codec-only training).
	The GE simulator is optional too: if no GilbertElliottConfig is provided, no
	packet loss is simulated and frame_mask is never generated automatically.
	"""

	import typing as tp
	from contextlib import contextmanager

	import numpy as np
	import torch
	import torch.nn as nn
	from vector_quantize_pytorch import ResidualVQ

	from .components import ZPEncoder, ZPDecoder
	from .repair import LatentRepairTransformer
	from .GilbertElliot import GilbertElliottConfig, GilbertElliottSimulator


	@contextmanager
	def temporarily_set(obj, attr: str, value):
	"""Context manager that sets obj.attr = value for the duration of the block,
	then restores the original value. Used to toggle quantize_dropout per-batch."""
	original = getattr(obj, attr)
	setattr(obj, attr, value)
	try:
	yield
	finally:
	setattr(obj, attr, original)


	class ZPCodec(nn.Module):
	"""
	Full codec: encoder -> RVQ -> (repair) -> decoder.

	Stage 1 (codec pre-training): use_repair=False, no GilbertElliottConfig.
	Stage 2 (repair training): use_repair=True, GilbertElliottConfig provided.
	Stage 3 (joint fine-tuning): use_repair=True, GE curriculum via set_gilbert_elliott_config().
	"""
	def __init__(
	self,
	channels: int = 1,
	dimension: int = 128,
	n_filters: int = 32,
	ratios: tp.List[int] = [8, 5, 3, 2],
	norm: str = 'weight_norm',
	causal: bool = True,
	num_quantizers: int = 9,
	codebook_size: int = 1024,
	sample_rate: int = 16000,
	# --- Repair module ---
	use_repair: bool = False,
	repair_hidden_dim: int = 256,
	repair_num_layers: int = 4,
	repair_num_heads: int = 4,
	repair_ffn_mult: int = 2,
	repair_past: int = 8,
	repair_future: int = 2,
	repair_two_pass: bool = True,
	# --- Packet loss simulation ---
	gilbert_elliott_config: tp.Optional[GilbertElliottConfig] = None,
	):
	super().__init__()
	self.encoder = ZPEncoder(
	channels=channels,
	dimension=dimension,
	n_filters=n_filters,
	ratios=ratios,
	norm=norm,
	causal=causal,
	)
	self.rvq = ResidualVQ(
	dim=dimension,
	num_quantizers=num_quantizers,
	codebook_size=codebook_size,
	kmeans_init=True,
	kmeans_iters=10,
	use_cosine_sim=True, # prop to improved RVQGAN's paper
	threshold_ema_dead_code=2,
	quantize_dropout=True,
	quantize_dropout_cutoff_index=5, # first 5 quantizers are always active -
	# theoretically with 5 quant active we can switch to 3kbps. But this was not my focus for that project...
	quantize_dropout_multiple_of=1,
	)
	self.decoder = ZPDecoder(
	channels=channels,
	dimension=dimension,
	n_filters=n_filters,
	ratios=ratios,
	norm=norm,
	causal=causal,
	)

	self.sample_rate = sample_rate
	self.hop_length = int(np.prod(ratios)) # 240 for ratios=[8,5,3,2]

	self.use_repair = use_repair
	self.repair_two_pass = repair_two_pass
	if use_repair:
	self.repair = LatentRepairTransformer(
	latent_dim=dimension,
	hidden_dim=repair_hidden_dim,
	num_layers=repair_num_layers,
	num_heads=repair_num_heads,
	ffn_mult=repair_ffn_mult,
	past=repair_past,
	future=repair_future,
	)
	else:
	self.repair = None

	self.ge_simulator: tp.Optional[GilbertElliottSimulator] = None
	if gilbert_elliott_config is not None:
	self.set_gilbert_elliott_config(gilbert_elliott_config)


	# Runtime configuration of the packet-loss simulator
	def set_gilbert_elliott_config(self, config: GilbertElliottConfig) -> None:
	"""Replace the GE simulator at runtime. Called between training stages
	to apply a harder loss curriculum without reloading the model."""
	self.ge_simulator = GilbertElliottSimulator(
	config=config,
	sample_rate=self.sample_rate,
	hop_length=self.hop_length,
	)

	def sample_frame_mask(
	self,
	batch_size: int,
	num_frames: int,
	device: tp.Optional[torch.device] = None,
	seed: tp.Optional[int] = None,
	) -> torch.Tensor:
	"""Expose the GE simulator directly. Useful when the same mask needs to
	be reused across multiple points (e.g. logging, loss weighting)."""
	assert self.ge_simulator is not None, (
	"No GilbertElliottConfig configured. Call set_gilbert_elliott_config() first."
	)
	return self.ge_simulator.sample_frame_mask(
	batch_size, num_frames, device=device, seed=seed
	)

	# Encoding
	def _encode_raw(self, x: torch.Tensor):
	"""Encode waveform to quantized latent. Returns (z, z_q, indices, commit_loss).
	quantize_dropout is randomly toggled per-call during training to teach
	the decoder to handle a variable number of active quantizers (bitrate scalability)."""
	z = self.encoder(x) # [B, D, T']
	z_seq = z.permute(0, 2, 1) # [B, T', D] — RVQ expects (B, T, D)
	use_dropout = self.training and (torch.rand(1).item() < 0.5) # dropout applied only 50% of the time, this improve the
	# quality at full kbps. Citing the improved RVQGAN paper.
	with temporarily_set(self.rvq, 'quantize_dropout', use_dropout):
	z_q, indices, commit_loss = self.rvq(z_seq)
	z_q = z_q.permute(0, 2, 1) # [B, D, T']
	return z, z_q, indices, commit_loss


	# Repair
	def _apply_repair(
	self,
	z_q: torch.Tensor,
	frame_mask: torch.Tensor,
	) -> torch.Tensor:
	"""Run the repair transformer and selectively substitute only missing frames.

	z_q: [B, D, T']
	frame_mask: [B, T'] 1 = received, 0 = missing

	The transformer outputs a full [B, D, T'] tensor, but received frames are
	kept as-is from z_q — only positions where frame_mask == 0 are replaced.
	This means z_q_post == z_q on received frames by construction, which is
	important for latent_repair_loss (the mask isolates the useful gradient).

	Two-pass mode (repair_two_pass=True): mimics streaming deployment where
	previous repair estimates are already in the buffer when estimating frame t.
	See LatentRepairTransformer.forward_two_pass for the full explanation.
	"""
	assert self.repair is not None, "use_repair=False, repair not initialised"

	z_seq = z_q.permute(0, 2, 1) # [B, T', D]

	if self.repair_two_pass:
	z_repaired = self.repair.forward_two_pass(z_seq, frame_mask)
	else:
	# Single-pass fallback
	z_seq_filled = self.repair.fill_missing(z_seq, frame_mask)
	z_repaired = self.repair(z_seq_filled, frame_mask)

	# Selective substitution: keep received frames from z_q, replace missing ones
	m = frame_mask.unsqueeze(-1).to(z_seq.dtype) # [B, T', 1]
	z_out = z_seq * m + z_repaired * (1.0 - m)
	return z_out.permute(0, 2, 1) # [B, D, T']

	def _get_frame_mask(
	self,
	z_q: torch.Tensor,
	frame_mask: tp.Optional[torch.Tensor],
	) -> torch.Tensor:
	"""Return the provided frame_mask, or sample one from the GE simulator."""
	if frame_mask is not None:
	return frame_mask
	assert self.ge_simulator is not None, (
	"use_repair=True but no GilbertElliottConfig configured. "
	"Call set_gilbert_elliott_config() before training."
	)
	B, _, T_prime = z_q.shape
	return self.ge_simulator.sample_frame_mask(B, T_prime, device=z_q.device)

	# Public encode / decode API
	def encode(self, x: torch.Tensor) -> tp.Tuple[torch.Tensor, torch.Tensor]:
	"""Encode waveform to (z_q, indices). x: [B, 1, T]"""
	_, z_q, indices, _ = self._encode_raw(x)
	return z_q, indices

	def decode(
	self,
	z_q: torch.Tensor,
	frame_mask: tp.Optional[torch.Tensor] = None,
	) -> torch.Tensor:
	"""Decode quantized latent to waveform.
	z_q: [B, D, T']
	frame_mask: [B, T'] optional; if provided and use_repair=True, runs repair first.
	"""
	if self.use_repair and frame_mask is not None:
	z_q = self._apply_repair(z_q, frame_mask)
	return self.decoder(z_q)

	# Training forward
	def forward(
	self,
	x: torch.Tensor,
	frame_mask: tp.Optional[torch.Tensor] = None,
	return_intermediates: bool = False,
	):
	"""
	x: [B, 1, T]
	frame_mask: [B, T'] optional. If use_repair=True and None,
	sampled automatically from the GE simulator.
	return_intermediates: if True, also returns z_q pre/post repair and the
	effective frame_mask — required by latent_repair_loss
	and ZPCodecTrainer.forward_codec during training.

	Returns:
	return_intermediates=False: (x_hat, commit_loss)
	return_intermediates=True: (x_hat, commit_loss, z_q_pre, z_q_post, frame_mask)
	When use_repair=False: z_q_pre == z_q_post and frame_mask == None.
	"""
	_, z_q_pre, _, commit_loss = self._encode_raw(x)
	commit_loss = commit_loss.mean()

	if self.use_repair:
	frame_mask = self._get_frame_mask(z_q_pre, frame_mask)
	z_q_post = self._apply_repair(z_q_pre, frame_mask)
	else:
	z_q_post = z_q_pre
	frame_mask = None

	x_hat = self.decoder(z_q_post)

	if return_intermediates:
	return x_hat, commit_loss, z_q_pre, z_q_post, frame_mask
	return x_hat, commit_loss

	# ------------------------------------------------------------------
	# from_pretrained — load from Hugging Face Hub or local path
	# ------------------------------------------------------------------
	@classmethod
	def from_pretrained(
	cls,
	model_id: str,
	device: str = "cpu",
	filename: str = "zpcodec_weights.pt",
	**hf_kwargs,
	) -> "ZPCodec":
	"""
	Load ZPCodec from a Hugging Face Hub repo or a local file.

	Args:
	model_id: HF repo id (e.g. "yourname/zpcodec") OR a local path
	to a .pt file OR a local directory containing filename.
	device: "cpu" \| "cuda" \| "cuda:0" etc.
	filename: name of the weights file inside the HF repo.
	**hf_kwargs: forwarded to huggingface_hub.hf_hub_download
	(e.g. revision="main", token="hf_...").

	Returns:
	ZPCodec in eval mode.

	Examples:
	# From Hugging Face Hub
	model = ZPCodec.from_pretrained("yourname/zpcodec")

	# From a local .pt file
	model = ZPCodec.from_pretrained("./zpcodec_weights.pt")

	# With explicit device
	model = ZPCodec.from_pretrained("yourname/zpcodec", device="cuda")
	"""
	import os
	import torch

	# Resolve checkpoint path: local file, local dir, or HF Hub
	if os.path.isfile(model_id):
	ckpt_path = model_id
	elif os.path.isdir(model_id):
	ckpt_path = os.path.join(model_id, filename)
	if not os.path.isfile(ckpt_path):
	raise FileNotFoundError(
	f"No '{filename}' found in directory '{model_id}'"
	)
	else:
	# Treat as a Hugging Face Hub repo id
	try:
	from huggingface_hub import hf_hub_download
	except ImportError:
	raise ImportError(
	"huggingface_hub is required to download from the Hub.\n"
	"Install with: pip install huggingface_hub"
	)
	ckpt_path = hf_hub_download(
	repo_id=model_id, filename=filename, **hf_kwargs
	)

	ckpt = torch.load(ckpt_path, map_location="cpu", weights_only=False)

	# Support both clean checkpoints (with 'config' key) and raw
	# full-trainer checkpoints (with 'args' key) for backward compat
	if "config" in ckpt:
	cfg = ckpt["config"]
	state_dict = ckpt["model_state_dict"]
	elif "args" in ckpt and "trainer" in ckpt:
	# Full trainer checkpoint — extract codec weights and config
	args = ckpt["args"]
	state_dict = {
	k[len("codec."):]: v
	for k, v in ckpt["trainer"].items()
	if k.startswith("codec.")
	}
	cfg = {
	"channels": 1, "dimension": args["dimension"],
	"n_filters": args["n_filters"], "ratios": [8, 5, 3, 2],
	"norm": "weight_norm", "causal": True,
	"num_quantizers": args["num_quantizers"],
	"codebook_size": args["codebook_size"], "sample_rate": 16000,
	"use_repair": True,
	"repair_hidden_dim": args["repair_hidden_dim"],
	"repair_num_layers": args["repair_num_layers"],
	"repair_num_heads": args["repair_num_heads"],
	"repair_ffn_mult": args["repair_ffn_mult"],
	"repair_past": args["repair_past"],
	"repair_future": args["repair_future"],
	"repair_two_pass": True,
	}
	else:
	raise ValueError(
	"Unrecognised checkpoint format. "
	"Expected keys: 'config'+'model_state_dict' or 'args'+'trainer'."
	)

	model = cls(**cfg)
	missing, unexpected = model.load_state_dict(state_dict, strict=True)
	if missing:
	raise RuntimeError(f"Missing keys: {missing[:5]}")
	if unexpected:
	raise RuntimeError(f"Unexpected keys: {unexpected[:5]}")

	n_params = sum(p.numel() for p in model.parameters()) / 1e6
	info = ckpt.get("training_info", {})
	stoi = info.get("best_val_stoi", ckpt.get("best_val_metric", "?"))
	print(f"✓ ZPCodec loaded — {n_params:.1f}M params \| best val STOI: {stoi}")

	model = model.to(device)
	model.eval()
	return model