Upload WaveCleanLDM/C6A_FinalArchiveTrueOldCompare/latest/models/VAE.py with huggingface_hub

WaveCleanLDM/C6A_FinalArchiveTrueOldCompare/latest/models/VAE.py

@@ -0,0 +1,1306 @@
+"""
+V5-Aware Analytic-Centered Structured Semantic VAE (AC-SSVAE)
+================================================================
+Stage-2 feature extractor for ICWDS.
+
+Design goals
+------------
+1. Consume WaveSystemSetParserV5 soft wave-system outputs instead of hard labels.
+2. Use an analytic physical descriptor as the prior center, and learn only bounded
+   semantic corrections.
+3. Canonicalize each wave system before shape encoding so the free latent does not
+   waste capacity on location, direction, scale, or energy.
+4. Keep a small causal shape latent. The residual decoder is constructed as
+       Delta(z, s) = F(z, s) - F(0, s)
+   so z=0 exactly returns the semantic base.
+5. Expose deterministic decoding APIs for quantization, intervention, and later
+   30-byte packet design.
+
+Internal semantic vector (9D)
+-----------------------------
+[log_energy,
+ peak_frequency_01,
+ sin_peak_direction,
+ cos_peak_direction,
+ frequency_spread_01,
+ direction_spread_over_pi,
+ f_theta_correlation,
+ frequency_skew,
+ direction_skew]
+
+The packet-facing direction is still one circular quantity. sin/cos are only the
+neural internal representation.
+"""
+from __future__ import annotations
+
+import math
+from dataclasses import dataclass, asdict
+from typing import Dict, Optional, Tuple
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+@dataclass
+class V5AwareACSSVAEConfig:
+    n_freqs: int = 47
+    n_dirs: int = 72
+    n_slots: int = 6
+    semantic_dim: int = 9
+    shape_dim: int = 4
+    context_dim: int = 24
+
+    canonical_h: int = 32
+    canonical_w: int = 48
+    canonical_f_extent: float = 3.0
+    canonical_t_extent: float = 3.0
+
+    encoder_width: int = 48
+    hidden_dim: int = 160
+    decoder_width: int = 64
+
+    # Bounded semantic correction magnitudes around analytic center.
+    # [logE, f, sin/cos correction, spreads, rho, skews]
+    delta_loge: float = 0.35
+    delta_f: float = 0.08
+    delta_angle_rad: float = math.radians(18.0)
+    delta_log_spread: float = 0.45
+    delta_rho_logit: float = 0.75
+    delta_skew: float = 0.90
+
+    min_spread_f: float = 0.010
+    max_spread_f: float = 0.45
+    min_spread_t: float = 0.015
+    max_spread_t: float = 0.95
+    max_abs_rho: float = 0.92
+    max_abs_skew: float = 3.0
+
+    shape_log_residual_scale: float = 2.25
+    energy_log_den: float = math.log1p(47 * 72)
+    eps: float = 1e-6
+
+    def to_dict(self):
+        return asdict(self)
+
+
+class ConvNormAct(nn.Module):
+    def __init__(self, ci: int, co: int, stride: int = 1):
+        super().__init__()
+        self.conv = nn.Conv2d(ci, co, 3, stride=stride, padding=1, bias=False)
+        self.norm = nn.GroupNorm(min(8, co), co)
+        self.act = nn.GELU()
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.act(self.norm(self.conv(x)))
+
+
+class ResBlock(nn.Module):
+    def __init__(self, ch: int):
+        super().__init__()
+        self.net = nn.Sequential(
+            nn.Conv2d(ch, ch, 3, padding=1, bias=False),
+            nn.GroupNorm(min(8, ch), ch),
+            nn.GELU(),
+            nn.Conv2d(ch, ch, 3, padding=1, bias=False),
+            nn.GroupNorm(min(8, ch), ch),
+        )
+        self.act = nn.GELU()
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.act(x + self.net(x))
+
+
+def _wrap_angle(x: torch.Tensor) -> torch.Tensor:
+    return torch.atan2(torch.sin(x), torch.cos(x))
+
+
+def _atanh_safe(x: torch.Tensor, eps: float = 1e-5) -> torch.Tensor:
+    x = x.clamp(-1 + eps, 1 - eps)
+    return 0.5 * (torch.log1p(x) - torch.log1p(-x))
+
+
+def reparameterize(mu: torch.Tensor, logvar: torch.Tensor, stochastic: bool) -> torch.Tensor:
+    if not stochastic:
+        return mu
+    std = torch.exp(0.5 * logvar)
+    return mu + std * torch.randn_like(std)
+
+
+@torch.no_grad()
+def analytic_semantics(
+    energy: torch.Tensor,
+    mask_prob: torch.Tensor,
+    core_prob: Optional[torch.Tensor] = None,
+    eps: float = 1e-6,
+    energy_log_den: Optional[float] = None,
+) -> torch.Tensor:
+    """Compute stable 9D physical descriptors from soft wave-system masks.
+
+    Parameters
+    ----------
+    energy: [B,1,H,W] or [B,H,W], non-negative display-domain energy.
+    mask_prob: [B,K,H,W].
+    core_prob: optional [B,K,H,W]. Core weighting sharpens peak estimation.
+    """
+    if energy.ndim == 3:
+        energy = energy[:, None]
+    B, K, H, W = mask_prob.shape
+    dtype, device = energy.dtype, energy.device
+    if energy_log_den is None:
+        energy_log_den = math.log1p(H * W)
+
+    E = torch.nan_to_num(energy, nan=0.0, posinf=1.0, neginf=0.0).clamp_min(0)
+    M = torch.nan_to_num(mask_prob, nan=0.0).clamp(0, 1)
+    part = E * M
+    mass = part.sum(dim=(-2, -1)).clamp_min(eps)
+    loge = torch.log1p(mass) / float(energy_log_den)
+
+    fgrid = torch.linspace(0, 1, H, device=device, dtype=dtype).view(1, 1, H, 1)
+    theta = torch.linspace(0, 2 * math.pi, W + 1, device=device, dtype=dtype)[:W]
+    tgrid = theta.view(1, 1, 1, W)
+
+    peak_src = part
+    if core_prob is not None:
+        core = torch.nan_to_num(core_prob, nan=0.0).clamp(0, 1)
+        peak_src = part * (0.25 + 0.75 * core)
+
+    # Smooth soft-peak estimator. It is more stable than hard argmax on noisy spectra.
+    pnorm = peak_src / peak_src.amax(dim=(-2, -1), keepdim=True).clamp_min(eps)
+    peak_w = torch.softmax((12.0 * pnorm).flatten(2), dim=-1).view(B, K, H, W)
+    fp = (peak_w * fgrid).sum(dim=(-2, -1))
+    sx = (peak_w * torch.sin(tgrid)).sum(dim=(-2, -1))
+    cx = (peak_w * torch.cos(tgrid)).sum(dim=(-2, -1))
+    theta_p = torch.atan2(sx, cx)
+    sinp = torch.sin(theta_p)
+    cosp = torch.cos(theta_p)
+
+    w = part / mass[:, :, None, None]
+    df = fgrid - fp[:, :, None, None]
+    dt = _wrap_angle(tgrid - theta_p[:, :, None, None]) / math.pi
+
+    sf = torch.sqrt((w * df.square()).sum(dim=(-2, -1)).clamp_min(eps))
+    st = torch.sqrt((w * dt.square()).sum(dim=(-2, -1)).clamp_min(eps))
+
+    zf = df / sf[:, :, None, None].clamp_min(1e-3)
+    zt = dt / st[:, :, None, None].clamp_min(1e-3)
+    rho = (w * zf * zt).sum(dim=(-2, -1)).clamp(-0.95, 0.95)
+    skew_f = (w * zf.pow(3)).sum(dim=(-2, -1)).clamp(-4.0, 4.0)
+    skew_t = (w * zt.pow(3)).sum(dim=(-2, -1)).clamp(-4.0, 4.0)
+
+    return torch.stack([
+        loge, fp, sinp, cosp, sf, st, rho, skew_f, skew_t
+    ], dim=-1)
+
+
+class Canonicalizer(nn.Module):
+    """Differentiably center/scale each slot in frequency-direction coordinates."""
+    def __init__(self, cfg: V5AwareACSSVAEConfig):
+        super().__init__()
+        self.cfg = cfg
+        u = torch.linspace(-cfg.canonical_f_extent, cfg.canonical_f_extent, cfg.canonical_h)
+        v = torch.linspace(-cfg.canonical_t_extent, cfg.canonical_t_extent, cfg.canonical_w)
+        self.register_buffer("u", u.view(1, 1, cfg.canonical_h, 1))
+        self.register_buffer("v", v.view(1, 1, 1, cfg.canonical_w))
+
+    def forward(
+        self,
+        energy: torch.Tensor,
+        mask_prob: torch.Tensor,
+        core_prob: torch.Tensor,
+        support_prob: torch.Tensor,
+        semantics: torch.Tensor,
+    ) -> torch.Tensor:
+        if energy.ndim == 3:
+            energy = energy[:, None]
+        B, K, H, W = mask_prob.shape
+        eps = self.cfg.eps
+
+        E = energy[:, None].expand(B, K, 1, H, W)
+        M = mask_prob[:, :, None]
+        C = core_prob[:, :, None]
+        S = support_prob[:, :, None]
+        part = E * M
+        cpart = E * C
+        spart = E * S
+
+        peak = part.amax(dim=(-2, -1), keepdim=True).clamp_min(eps)
+        channels = torch.cat([
+            part / peak,
+            cpart / peak,
+            spart / peak,
+            M,
+            C,
+            S,
+        ], dim=2).reshape(B * K, 6, H, W)
+
+        sem = semantics.reshape(B * K, -1)
+        fp = sem[:, 1].clamp(0, 1)
+        theta = torch.atan2(sem[:, 2], sem[:, 3]) % (2 * math.pi)
+        sf = sem[:, 4].clamp(self.cfg.min_spread_f, self.cfg.max_spread_f)
+        st = sem[:, 5].clamp(self.cfg.min_spread_t, self.cfg.max_spread_t)
+
+        uf = self.u.to(channels.dtype)
+        vt = self.v.to(channels.dtype)
+        f_sample = fp[:, None, None, None] + sf[:, None, None, None] * uf
+        theta_sample = theta[:, None, None, None] + (math.pi * st[:, None, None, None]) * vt
+        theta_frac = torch.remainder(theta_sample, 2 * math.pi) / (2 * math.pi)
+
+        # Tile direction three times and sample from the middle copy.
+        tiled = torch.cat([channels, channels, channels], dim=-1)
+        y = 2.0 * f_sample - 1.0
+        xpix = (1.0 + theta_frac) * W - 0.5
+        x = 2.0 * xpix / max(3 * W - 1, 1) - 1.0
+        grid = torch.stack([
+            x.expand(-1, -1, self.cfg.canonical_h, self.cfg.canonical_w).squeeze(1),
+            y.expand(-1, -1, self.cfg.canonical_h, self.cfg.canonical_w).squeeze(1),
+        ], dim=-1)
+        out = F.grid_sample(
+            tiled,
+            grid,
+            mode="bilinear",
+            padding_mode="zeros",
+            align_corners=True,
+        )
+        return out.reshape(B, K, 6, self.cfg.canonical_h, self.cfg.canonical_w)
+
+
+class SlotEncoder(nn.Module):
+    def __init__(self, cfg: V5AwareACSSVAEConfig):
+        super().__init__()
+        w = cfg.encoder_width
+        self.image_net = nn.Sequential(
+            ConvNormAct(6, w, 1),
+            ResBlock(w),
+            ConvNormAct(w, w * 2, 2),
+            ResBlock(w * 2),
+            ConvNormAct(w * 2, w * 3, 2),
+            ResBlock(w * 3),
+            ConvNormAct(w * 3, w * 4, 2),
+            nn.AdaptiveAvgPool2d(1),
+            nn.Flatten(),
+        )
+        self.context_net = nn.Sequential(
+            nn.Linear(cfg.context_dim + cfg.semantic_dim + 1, cfg.hidden_dim // 2),
+            nn.LayerNorm(cfg.hidden_dim // 2),
+            nn.GELU(),
+            nn.Linear(cfg.hidden_dim // 2, cfg.hidden_dim // 2),
+            nn.GELU(),
+        )
+        image_dim = w * 4
+        self.fuse = nn.Sequential(
+            nn.Linear(image_dim + cfg.hidden_dim // 2, cfg.hidden_dim),
+            nn.LayerNorm(cfg.hidden_dim),
+            nn.GELU(),
+            nn.Linear(cfg.hidden_dim, cfg.hidden_dim),
+            nn.GELU(),
+        )
+        self.sem_mu = nn.Linear(cfg.hidden_dim, cfg.semantic_dim)
+        self.sem_logvar = nn.Linear(cfg.hidden_dim, cfg.semantic_dim)
+        self.shape_mu = nn.Linear(cfg.hidden_dim, cfg.shape_dim)
+        self.shape_logvar = nn.Linear(cfg.hidden_dim, cfg.shape_dim)
+
+        # Start from the analytic center and nearly-zero free shape code.
+        nn.init.zeros_(self.sem_mu.weight)
+        nn.init.zeros_(self.sem_mu.bias)
+        nn.init.constant_(self.sem_logvar.bias, -4.0)
+        nn.init.normal_(self.shape_mu.weight, std=0.01)
+        nn.init.zeros_(self.shape_mu.bias)
+        nn.init.constant_(self.shape_logvar.bias, -3.0)
+
+    def forward(
+        self,
+        canonical: torch.Tensor,
+        analytic_sem: torch.Tensor,
+        context: torch.Tensor,
+        exist_prob: torch.Tensor,
+    ) -> Dict[str, torch.Tensor]:
+        B, K = canonical.shape[:2]
+        img = self.image_net(canonical.reshape(B * K, *canonical.shape[2:]))
+        ctx_in = torch.cat([
+            context.reshape(B * K, -1),
+            analytic_sem.reshape(B * K, -1),
+            exist_prob.reshape(B * K, 1),
+        ], dim=-1)
+        ctx = self.context_net(ctx_in)
+        h = self.fuse(torch.cat([img, ctx], dim=-1))
+        sem_mu = self.sem_mu(h).reshape(B, K, -1)
+        sem_logvar = self.sem_logvar(h).clamp(-8.0, 3.0).reshape(B, K, -1)
+        shape_mu = self.shape_mu(h).reshape(B, K, -1)
+        shape_logvar = self.shape_logvar(h).clamp(-8.0, 3.0).reshape(B, K, -1)
+        return {
+            "sem_delta_mu": sem_mu,
+            "sem_delta_logvar": sem_logvar,
+            "shape_mu": shape_mu,
+            "shape_logvar": shape_logvar,
+        }
+
+
+class SemanticRenderer(nn.Module):
+    """Parameter-free generalized elliptical renderer from corrected semantics."""
+    def __init__(self, cfg: V5AwareACSSVAEConfig):
+        super().__init__()
+        self.cfg = cfg
+        f = torch.linspace(0, 1, cfg.n_freqs).view(1, 1, cfg.n_freqs, 1)
+        t = torch.linspace(0, 2 * math.pi, cfg.n_dirs + 1)[:cfg.n_dirs].view(1, 1, 1, cfg.n_dirs)
+        self.register_buffer("fgrid", f)
+        self.register_buffer("tgrid", t)
+
+    def forward(self, sem: torch.Tensor) -> torch.Tensor:
+        cfg = self.cfg
+        eps = cfg.eps
+        loge, fp = sem[..., 0], sem[..., 1]
+        theta = torch.atan2(sem[..., 2], sem[..., 3])
+        sf = sem[..., 4].clamp(cfg.min_spread_f, cfg.max_spread_f)
+        st = sem[..., 5].clamp(cfg.min_spread_t, cfg.max_spread_t)
+        rho = sem[..., 6].clamp(-cfg.max_abs_rho, cfg.max_abs_rho)
+        skew_f = sem[..., 7].clamp(-cfg.max_abs_skew, cfg.max_abs_skew)
+        skew_t = sem[..., 8].clamp(-cfg.max_abs_skew, cfg.max_abs_skew)
+
+        df = (self.fgrid - fp[:, :, None, None]) / sf[:, :, None, None].clamp_min(1e-3)
+        dt = _wrap_angle(self.tgrid - theta[:, :, None, None])
+        dt = dt / (math.pi * st[:, :, None, None].clamp_min(1e-3))
+        den = (1.0 - rho.square()).clamp_min(0.08)
+        q = (df.square() + dt.square() - 2.0 * rho[:, :, None, None] * df * dt) / den[:, :, None, None]
+        shape = torch.exp(-0.5 * q.clamp_max(60.0))
+        asym = torch.exp(
+            0.28 * skew_f[:, :, None, None] * torch.tanh(df)
+            + 0.28 * skew_t[:, :, None, None] * torch.tanh(dt)
+        ).clamp(0.15, 6.0)
+        shape = shape * asym
+        shape = shape / shape.sum(dim=(-2, -1), keepdim=True).clamp_min(eps)
+        mass = torch.expm1(loge * cfg.energy_log_den).clamp_min(0.0)
+        return shape * mass[:, :, None, None]
+
+
+class ShapeResidualNet(nn.Module):
+    def __init__(self, cfg: V5AwareACSSVAEConfig):
+        super().__init__()
+        self.cfg = cfg
+        w = cfg.decoder_width
+        self.fc = nn.Sequential(
+            nn.Linear(cfg.semantic_dim + cfg.shape_dim, cfg.hidden_dim),
+            nn.GELU(),
+            nn.Linear(cfg.hidden_dim, w * 6 * 9),
+            nn.GELU(),
+        )
+        self.net = nn.Sequential(
+            ConvNormAct(w, w, 1),
+            ResBlock(w),
+            nn.Upsample(size=(12, 18), mode="bilinear", align_corners=False),
+            ConvNormAct(w, w, 1),
+            ResBlock(w),
+            nn.Upsample(size=(24, 36), mode="bilinear", align_corners=False),
+            ConvNormAct(w, w // 2, 1),
+            ResBlock(w // 2),
+            nn.Upsample(size=(48, 72), mode="bilinear", align_corners=False),
+            ConvNormAct(w // 2, w // 2, 1),
+            nn.Conv2d(w // 2, 1, 3, padding=1),
+        )
+
+    def field(self, sem: torch.Tensor, zshape: torch.Tensor) -> torch.Tensor:
+        B, K = sem.shape[:2]
+        x = torch.cat([sem, zshape], dim=-1).reshape(B * K, -1)
+        h = self.fc(x).view(B * K, self.cfg.decoder_width, 6, 9)
+        out = self.net(h)[:, :, : self.cfg.n_freqs, : self.cfg.n_dirs]
+        return out.reshape(B, K, self.cfg.n_freqs, self.cfg.n_dirs)
+
+    def forward(self, sem: torch.Tensor, zshape: torch.Tensor) -> torch.Tensor:
+        raw = self.field(sem, zshape)
+        zero = self.field(sem, torch.zeros_like(zshape))
+        return self.cfg.shape_log_residual_scale * torch.tanh(raw - zero)
+
+
+class V5AwareStructuredSemanticVAE(nn.Module):
+    def __init__(self, cfg: Optional[V5AwareACSSVAEConfig] = None):
+        super().__init__()
+        self.cfg = cfg or V5AwareACSSVAEConfig()
+        self.canonicalizer = Canonicalizer(self.cfg)
+        self.encoder = SlotEncoder(self.cfg)
+        self.semantic_renderer = SemanticRenderer(self.cfg)
+        self.shape_decoder = ShapeResidualNet(self.cfg)
+
+    def apply_semantic_delta(self, base: torch.Tensor, delta: torch.Tensor) -> torch.Tensor:
+        c = self.cfg
+        out = base.clone()
+        out[..., 0] = (base[..., 0] + c.delta_loge * torch.tanh(delta[..., 0])).clamp_min(0.0)
+        out[..., 1] = (base[..., 1] + c.delta_f * torch.tanh(delta[..., 1])).clamp(0.0, 1.0)
+
+        theta0 = torch.atan2(base[..., 2], base[..., 3])
+        # Use both sin/cos correction channels to form one stable tangent-angle correction.
+        dtheta = c.delta_angle_rad * torch.tanh(0.7071 * (delta[..., 2] - delta[..., 3]))
+        theta = theta0 + dtheta
+        out[..., 2] = torch.sin(theta)
+        out[..., 3] = torch.cos(theta)
+
+        out[..., 4] = (
+            base[..., 4].clamp_min(c.min_spread_f)
+            * torch.exp(c.delta_log_spread * torch.tanh(delta[..., 4]))
+        ).clamp(c.min_spread_f, c.max_spread_f)
+        out[..., 5] = (
+            base[..., 5].clamp_min(c.min_spread_t)
+            * torch.exp(c.delta_log_spread * torch.tanh(delta[..., 5]))
+        ).clamp(c.min_spread_t, c.max_spread_t)
+        out[..., 6] = torch.tanh(
+            _atanh_safe(base[..., 6], c.eps)
+            + c.delta_rho_logit * torch.tanh(delta[..., 6])
+        ).clamp(-c.max_abs_rho, c.max_abs_rho)
+        out[..., 7] = (
+            base[..., 7] + c.delta_skew * torch.tanh(delta[..., 7])
+        ).clamp(-c.max_abs_skew, c.max_abs_skew)
+        out[..., 8] = (
+            base[..., 8] + c.delta_skew * torch.tanh(delta[..., 8])
+        ).clamp(-c.max_abs_skew, c.max_abs_skew)
+        return out
+
+    def decode_from_codes(
+        self,
+        semantics: torch.Tensor,
+        shape_code: torch.Tensor,
+        exist_prob: Optional[torch.Tensor] = None,
+    ) -> Dict[str, torch.Tensor]:
+        base = self.semantic_renderer(semantics)
+        residual_log = self.shape_decoder(semantics, shape_code)
+        # Multiplicative residual in log domain, followed by exact energy renormalization.
+        part_hat = (base + self.cfg.eps) * torch.exp(residual_log)
+        target_mass = torch.expm1(semantics[..., 0] * self.cfg.energy_log_den).clamp_min(0.0)
+        part_hat = part_hat * (
+            target_mass[:, :, None, None]
+            / part_hat.sum(dim=(-2, -1), keepdim=True).clamp_min(self.cfg.eps)
+        )
+        if exist_prob is not None:
+            part_hat = part_hat * exist_prob[:, :, None, None]
+            base = base * exist_prob[:, :, None, None]
+        return {
+            "semantic_base": base,
+            "shape_residual_log": residual_log,
+            "part_hat": part_hat,
+        }
+
+    def forward(
+        self,
+        energy: torch.Tensor,
+        mask_prob: torch.Tensor,
+        core_prob: torch.Tensor,
+        support_prob: torch.Tensor,
+        exist_prob: torch.Tensor,
+        slot_context: torch.Tensor,
+        stochastic: bool = True,
+        shape_code_override: Optional[torch.Tensor] = None,
+        semantic_delta_override: Optional[torch.Tensor] = None,
+    ) -> Dict[str, torch.Tensor]:
+        analytic = analytic_semantics(
+            energy,
+            mask_prob,
+            core_prob=core_prob,
+            eps=self.cfg.eps,
+            energy_log_den=self.cfg.energy_log_den,
+        )
+        canonical = self.canonicalizer(
+            energy,
+            mask_prob,
+            core_prob,
+            support_prob,
+            analytic,
+        )
+        enc = self.encoder(canonical, analytic, slot_context, exist_prob)
+        sem_delta = (
+            semantic_delta_override
+            if semantic_delta_override is not None
+            else reparameterize(enc["sem_delta_mu"], enc["sem_delta_logvar"], stochastic)
+        )
+        zshape = (
+            shape_code_override
+            if shape_code_override is not None
+            else reparameterize(enc["shape_mu"], enc["shape_logvar"], stochastic)
+        )
+        corrected = self.apply_semantic_delta(analytic, sem_delta)
+        dec = self.decode_from_codes(corrected, zshape, exist_prob=exist_prob)
+        return {
+            "analytic_semantics": analytic,
+            "corrected_semantics": corrected,
+            "semantic_delta": sem_delta,
+            "shape_code": zshape,
+            "canonical": canonical,
+            **enc,
+            **dec,
+        }
+
+    @property
+    def token_dim_per_wave(self) -> int:
+        return self.cfg.semantic_dim + self.cfg.shape_dim
+
+    def num_params(self) -> int:
+        return sum(p.numel() for p in self.parameters())
+
+
+def kl_standard_normal(mu: torch.Tensor, logvar: torch.Tensor, weight: Optional[torch.Tensor] = None) -> torch.Tensor:
+    kl = -0.5 * (1.0 + logvar - mu.square() - logvar.exp()).sum(dim=-1)
+    if weight is None:
+        return kl.mean()
+    w = weight.float()
+    return (kl * w).sum() / w.sum().clamp_min(1.0)
+
+
+def covariance_penalty(x: torch.Tensor, weight: Optional[torch.Tensor] = None) -> torch.Tensor:
+    """Off-diagonal covariance penalty for [B,K,D] latent tensors."""
+    D = x.shape[-1]
+    flat = x.reshape(-1, D)
+    if weight is not None:
+        w = weight.reshape(-1).float()
+        keep = w > 0.25
+        flat = flat[keep]
+    if flat.shape[0] < max(4, D):
+        return x.new_zeros(())
+    flat = flat - flat.mean(dim=0, keepdim=True)
+    cov = flat.T @ flat / max(flat.shape[0] - 1, 1)
+    off = cov - torch.diag(torch.diag(cov))
+    return off.square().mean()
+
+
+def semantic_distance(pred: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
+    """Per-slot interpretable semantic distance with circular direction handling."""
+    d = []
+    d.append((pred[..., 0] - target[..., 0]).abs())
+    d.append(2.0 * (pred[..., 1] - target[..., 1]).abs())
+    tp = torch.atan2(pred[..., 2], pred[..., 3])
+    tt = torch.atan2(target[..., 2], target[..., 3])
+    d.append(_wrap_angle(tp - tt).abs() / math.pi)
+    d.append(2.0 * (pred[..., 4] - target[..., 4]).abs())
+    d.append(1.5 * (pred[..., 5] - target[..., 5]).abs())
+    d.append(0.5 * (pred[..., 6] - target[..., 6]).abs())
+    d.append(0.15 * (pred[..., 7] - target[..., 7]).abs())
+    d.append(0.15 * (pred[..., 8] - target[..., 8]).abs())
+    return torch.stack(d, dim=-1).mean(dim=-1)
+
+
+def marginal_l1(pred: torch.Tensor, target: torch.Tensor, eps: float = 1e-6) -> Tuple[torch.Tensor, torch.Tensor]:
+    p = pred.clamp_min(0)
+    t = target.clamp_min(0)
+    pf = p.sum(dim=-1); tf = t.sum(dim=-1)
+    pt = p.sum(dim=-2); tt = t.sum(dim=-2)
+    pf = pf / pf.sum(dim=-1, keepdim=True).clamp_min(eps)
+    tf = tf / tf.sum(dim=-1, keepdim=True).clamp_min(eps)
+    pt = pt / pt.sum(dim=-1, keepdim=True).clamp_min(eps)
+    tt = tt / tt.sum(dim=-1, keepdim=True).clamp_min(eps)
+    return (pf - tf).abs().mean(), (pt - tt).abs().mean()
+
+
+def reconstruction_terms(pred: torch.Tensor, target: torch.Tensor, weight: torch.Tensor, eps: float = 1e-6) -> Dict[str, torch.Tensor]:
+    w = weight[:, :, None, None].float()
+    denom = w.sum().clamp_min(1.0)
+    mse_slot = (pred - target).square().mean(dim=(-2, -1))
+    log_slot = (
+        torch.log1p(20.0 * pred.clamp_min(0))
+        - torch.log1p(20.0 * target.clamp_min(0))
+    ).abs().mean(dim=(-2, -1))
+    mass_p = pred.sum(dim=(-2, -1))
+    mass_t = target.sum(dim=(-2, -1))
+    energy_slot = (mass_p - mass_t).abs() / mass_t.clamp_min(eps)
+    mse = (mse_slot * weight).sum() / denom
+    log_l1 = (log_slot * weight).sum() / denom
+    energy = (energy_slot * weight).sum() / denom
+    freq, direction = marginal_l1(pred * w, target * w, eps=eps)
+    return {
+        "mse": mse,
+        "log_l1": log_l1,
+        "energy": energy,
+        "freq": freq,
+        "direction": direction,
+    }
+
+
+__all__ = [
+    "V5AwareACSSVAEConfig",
+    "V5AwareStructuredSemanticVAE",
+    "analytic_semantics",
+    "kl_standard_normal",
+    "covariance_penalty",
+    "semantic_distance",
+    "reconstruction_terms",
+    "marginal_l1",
+]
+
+# =============================================================================
+# ICWDS Stage-2 hybrid codec components (C3-D / C3-F / C4-B)
+# =============================================================================
+
+NF = 47
+ND = 72
+K_MAX = 6
+
+class PrototypeSelector(nn.Module):
+    def __init__(self, input_dim, n_codes, width=192, dropout=0.10):
+        super().__init__()
+        self.norm = nn.LayerNorm(input_dim)
+        self.fc1 = nn.Linear(input_dim, width)
+        self.fc2 = nn.Linear(width, width)
+        self.fc3 = nn.Linear(width, width)
+        self.drop = nn.Dropout(dropout)
+        self.head = nn.Linear(width, n_codes)
+
+    def forward(self, x):
+        x = self.norm(x)
+        h = F.gelu(self.fc1(x))
+        r = h
+        h = self.drop(F.gelu(self.fc2(h)))
+        h = F.gelu(self.fc3(h) + r)
+        return self.head(h)
+
+class PeriodicConv2d(nn.Module):
+    def __init__(self, ci, co, kernel=3, stride=1, bias=False):
+        super().__init__()
+        self.pad = kernel // 2
+        self.conv = nn.Conv2d(ci, co, kernel_size=kernel, stride=stride, padding=0, bias=bias)
+
+    def forward(self, x):
+        p = self.pad
+        x = F.pad(x, (p, p, 0, 0), mode="circular")
+        x = F.pad(x, (0, 0, p, p), mode="replicate")
+        return self.conv(x)
+
+class PResBlock(nn.Module):
+    def __init__(self, ch):
+        super().__init__()
+        self.c1 = PeriodicConv2d(ch, ch, 3, 1, False)
+        self.n1 = nn.GroupNorm(min(8, ch), ch)
+        self.c2 = PeriodicConv2d(ch, ch, 3, 1, False)
+        self.n2 = nn.GroupNorm(min(8, ch), ch)
+
+    def forward(self, x):
+        h = F.gelu(self.n1(self.c1(x)))
+        h = self.n2(self.c2(h))
+        return F.gelu(x + h)
+
+class NestedBitEncoder(nn.Module):
+    def __init__(self, max_bits, width=32):
+        super().__init__()
+        w = width
+        self.stem = nn.Sequential(
+            PeriodicConv2d(4, w, 3, 2, False),
+            nn.GroupNorm(min(8, w), w), nn.GELU(), PResBlock(w),
+            PeriodicConv2d(w, w*2, 3, 2, False),
+            nn.GroupNorm(min(8, w*2), w*2), nn.GELU(), PResBlock(w*2),
+            PeriodicConv2d(w*2, w*3, 3, 2, False),
+            nn.GroupNorm(min(8, w*3), w*3), nn.GELU(), PResBlock(w*3),
+        )
+        ch = w * 3
+        self.attn = nn.Conv2d(ch, 1, 1)
+        self.head = nn.Sequential(
+            nn.Linear(ch * 3, ch * 2),
+            nn.LayerNorm(ch * 2),
+            nn.GELU(),
+            nn.Linear(ch * 2, max_bits),
+        )
+
+    def forward(self, residual_target, prior, gate):
+        x = torch.stack([residual_target, residual_target.abs(), prior, gate], dim=1)
+        h = self.stem(x)
+        flat = h.flatten(2)
+        attn = torch.softmax(self.attn(h).flatten(2), dim=-1)
+        ap = (flat * attn).sum(dim=-1)
+        mx = F.adaptive_max_pool2d(h, 1).flatten(1)
+        av = F.adaptive_avg_pool2d(h, 1).flatten(1)
+        return self.head(torch.cat([ap, mx, av], dim=-1))
+
+def ste_binary(logits):
+    soft = torch.tanh(logits)
+    hard = torch.where(logits >= 0, torch.ones_like(logits), -torch.ones_like(logits))
+    q = soft + (hard - soft).detach()
+    return q, soft
+
+def prefix_mask(n_bits, max_bits, device, dtype):
+    if int(n_bits) < 0 or int(n_bits) > int(max_bits):
+        raise ValueError((n_bits, max_bits))
+    return (torch.arange(max_bits, device=device) < int(n_bits)).to(dtype=dtype)[None]
+
+class CondPyramid(nn.Module):
+    def __init__(self, width=32):
+        super().__init__()
+        w = width
+        self.c0 = nn.Sequential(
+            PeriodicConv2d(2, w, 3, 1, False),
+            nn.GroupNorm(min(8, w), w), nn.GELU(),
+        )
+        self.c1 = nn.Sequential(
+            PeriodicConv2d(w, w*2, 3, 2, False),
+            nn.GroupNorm(min(8, w*2), w*2), nn.GELU(),
+        )
+        self.c2 = nn.Sequential(
+            PeriodicConv2d(w*2, w*3, 3, 2, False),
+            nn.GroupNorm(min(8, w*3), w*3), nn.GELU(),
+        )
+        self.c3 = nn.Sequential(
+            PeriodicConv2d(w*3, w*3, 3, 2, False),
+            nn.GroupNorm(min(8, w*3), w*3), nn.GELU(),
+        )
+
+    def forward(self, prior, gate):
+        x = torch.stack([prior, gate], dim=1)
+        c0 = self.c0(x)
+        c1 = self.c1(c0)
+        c2 = self.c2(c1)
+        c3 = self.c3(c2)
+        return c0, c1, c2, c3
+
+class CodeFiLM(nn.Module):
+    """Near-identity code injection. Final layer is zero-initialized."""
+    def __init__(self, max_bits, channels):
+        super().__init__()
+        hid = max(48, channels)
+        self.net = nn.Sequential(
+            nn.Linear(max_bits + 1, hid),
+            nn.GELU(),
+            nn.Linear(hid, channels * 2),
+        )
+        nn.init.zeros_(self.net[-1].weight)
+        nn.init.zeros_(self.net[-1].bias)
+
+    def forward(self, h, qbits, bit_fraction):
+        B = qbits.shape[0]
+        bf = qbits.new_full((B, 1), float(bit_fraction))
+        gamma, beta = self.net(torch.cat([qbits, bf], dim=-1)).chunk(2, dim=-1)
+        gamma = 0.20 * torch.tanh(gamma)[:, :, None, None]
+        beta = 0.20 * torch.tanh(beta)[:, :, None, None]
+        return h * (1.0 + gamma) + beta
+
+class DetailAwareResidualDecoder(nn.Module):
+    def __init__(self, max_bits, width=32, branch="intra"):
+        super().__init__()
+        self.max_bits = int(max_bits)
+        self.branch = str(branch)
+        w = width
+        self.cond = CondPyramid(w)
+        self.bit_fc = nn.Sequential(nn.Linear(max_bits, w*3*6*9), nn.GELU())
+        self.b3 = PResBlock(w*3)
+        self.p2 = nn.Conv2d(w*3, w*3, 1)
+        self.b2 = PResBlock(w*3)
+        self.p1 = nn.Conv2d(w*3, w*2, 1)
+        self.b1 = PResBlock(w*2)
+        self.p0 = nn.Conv2d(w*2, w, 1)
+        self.b0 = PResBlock(w)
+        self.out = PeriodicConv2d(w, 1, 3, 1, True)
+        self.scale = nn.Parameter(torch.tensor(0.35))
+
+        # New C3-F modules.
+        self.film3 = CodeFiLM(max_bits, w*3)
+        self.film2 = CodeFiLM(max_bits, w*3)
+        self.film1 = CodeFiLM(max_bits, w*2)
+        self.film0 = CodeFiLM(max_bits, w)
+        self.detail_block = PResBlock(w)
+        self.detail_out = PeriodicConv2d(w, 1, 3, 1, True)
+        nn.init.zeros_(self.detail_out.conv.weight)
+        if self.detail_out.conv.bias is not None:
+            nn.init.zeros_(self.detail_out.conv.bias)
+        self.detail_scale = nn.Parameter(torch.tensor(-2.0))
+
+    def forward(self, qbits, prior, gate, n_bits=None):
+        B = prior.shape[0]
+        if n_bits is None:
+            n_bits = int((qbits.abs().sum(dim=0) > 0).sum().item())
+        frac = float(n_bits) / float(self.max_bits)
+        c0, c1, c2, c3 = self.cond(prior, gate)
+
+        h = self.bit_fc(qbits).view(B, c3.shape[1], 6, 9)
+        h = self.b3(h + c3)
+        h = self.film3(h, qbits, frac)
+
+        h = F.interpolate(h, size=c2.shape[-2:], mode="bilinear", align_corners=False)
+        h = self.b2(self.p2(h) + c2)
+        h = self.film2(h, qbits, frac)
+
+        h = F.interpolate(h, size=c1.shape[-2:], mode="bilinear", align_corners=False)
+        h = self.b1(self.p1(h) + c1)
+        h = self.film1(h, qbits, frac)
+
+        h = F.interpolate(h, size=(48, 72), mode="bilinear", align_corners=False)
+        c0p = F.pad(c0, (0, 0, 0, 1), mode="replicate")
+        h = self.b0(self.p0(h) + c0p)
+        h = self.film0(h, qbits, frac)
+
+        base = self.out(h)[:, 0, :NF, :ND]
+        detail = self.detail_out(self.detail_block(h))[:, 0, :NF, :ND]
+        # High-pass detail residual; direction is periodic, frequency is replicated.
+        dp = F.pad(detail[:, None], (1, 1, 0, 0), mode="circular")
+        dp = F.pad(dp, (0, 0, 1, 1), mode="replicate")
+        blur = F.avg_pool2d(dp, 3, stride=1)[:, 0]
+        detail_hp = detail - blur
+        raw = torch.tanh(base + torch.sigmoid(self.detail_scale) * detail_hp)
+
+        amp = 0.65 * torch.sigmoid(self.scale)
+        spatial = (0.08 + 0.92 * gate) if self.branch == "intra" else (1.0 - 0.72 * gate)
+        return amp * raw * spatial
+
+class AdaptiveDualResidualCodec(nn.Module):
+    def __init__(self, max_bits=24, enc_width=32, dec_width=32):
+        super().__init__()
+        self.max_bits = int(max_bits)
+        self.intra_encoder = NestedBitEncoder(max_bits, enc_width)
+        self.bg_encoder = NestedBitEncoder(max_bits, enc_width)
+        self.intra_decoder = DetailAwareResidualDecoder(max_bits, dec_width, branch="intra")
+        self.bg_decoder = DetailAwareResidualDecoder(max_bits, dec_width, branch="bg")
+
+    def encode(self, prior, gate, intra_target, bg_target):
+        li = self.intra_encoder(intra_target, prior, gate)
+        lb = self.bg_encoder(bg_target, prior, gate)
+        return li, lb
+
+    def decode_logits(self, logits, n_bits, prior, gate, branch, q_override=None):
+        if int(n_bits) == 0:
+            return torch.zeros_like(prior), None, None
+        q, soft = ste_binary(logits)
+        mask = prefix_mask(n_bits, self.max_bits, logits.device, logits.dtype)
+        q = q * mask
+        if q_override is not None:
+            q = q_override * mask
+        dec = self.intra_decoder if branch == "intra" else self.bg_decoder
+        pred = dec(q, prior, gate, n_bits=int(n_bits))
+        return pred, q, soft
+
+    def decode_pair(self, li, lb, prior, gate, intra_bits, bg_bits):
+        ri, qi, si = self.decode_logits(li, intra_bits, prior, gate, "intra")
+        rb, qb, sb = self.decode_logits(lb, bg_bits, prior, gate, "bg")
+        return {"ri": ri, "rb": rb, "qi": qi, "qb": qb, "si": si, "sb": sb}
+
+    def forward(self, prior, gate, intra_target, bg_target, intra_bits, bg_bits):
+        li, lb = self.encode(prior, gate, intra_target, bg_target)
+        out = self.decode_pair(li, lb, prior, gate, intra_bits, bg_bits)
+        out.update({"li": li, "lb": lb})
+        return out
+
+    @property
+    def num_params(self):
+        return sum(p.numel() for p in self.parameters())
+
+class RateConditionedController(nn.Module):
+    def __init__(self, input_dim, n_alloc, width=160, dropout=0.08):
+        super().__init__()
+        self.norm = nn.LayerNorm(input_dim + 1)
+        self.fc1 = nn.Linear(input_dim + 1, width)
+        self.fc2 = nn.Linear(width, width)
+        self.fc3 = nn.Linear(width, width)
+        self.drop = nn.Dropout(dropout)
+        self.head = nn.Linear(width, n_alloc)
+
+    def forward(self, x, lambda_norm):
+        if lambda_norm.ndim == 1:
+            lambda_norm = lambda_norm[:, None]
+        h = self.norm(torch.cat([x, lambda_norm], dim=-1))
+        h = F.gelu(self.fc1(h)); r = h
+        h = self.drop(F.gelu(self.fc2(h)))
+        h = F.gelu(self.fc3(h) + r)
+        return self.head(h)
+
+def branch_scalar_features(x):
+    ax = x.abs()
+    df = (x[:, 1:] - x[:, :-1]).abs().mean(dim=(-2, -1))
+    dd = (torch.roll(x, -1, -1) - x).abs().mean(dim=(-2, -1))
+    lap = torch.zeros_like(x)
+    lap[:, 1:-1] = x[:, 2:] - 2*x[:, 1:-1] + x[:, :-2]
+    return torch.stack([
+        ax.mean(dim=(-2, -1)), ax.std(dim=(-2, -1)), ax.amax(dim=(-2, -1)),
+        df, dd, lap.abs().mean(dim=(-2, -1)),
+    ], dim=-1)
+
+def controller_features(batch, li, lb):
+    intra = branch_scalar_features(batch["intra_target"])
+    bg = branch_scalar_features(batch["bg_struct"])
+    prior = branch_scalar_features(batch["main_prior"])
+    gate = torch.stack([
+        batch["prior_gate"].mean(dim=(-2, -1)),
+        batch["prior_gate"].std(dim=(-2, -1)),
+    ], dim=-1)
+    conf_i = torch.tanh(li).abs(); conf_b = torch.tanh(lb).abs()
+    prefix_points = [4, 8, 12, 16]
+    ci = torch.stack([conf_i[:, :p].mean(-1) for p in prefix_points], dim=-1)
+    cb = torch.stack([conf_b[:, :p].mean(-1) for p in prefix_points], dim=-1)
+    ratio = torch.stack([
+        intra[:, 0] / prior[:, 0].clamp_min(1e-5),
+        bg[:, 0] / prior[:, 0].clamp_min(1e-5),
+    ], dim=-1)
+    return torch.cat([intra, bg, prior, gate, ci, cb, ratio], dim=-1)
+
+
+class ProtocolRuntime:
+    """Runtime view of the frozen C4-A/C4-B exact 30-byte packet protocol."""
+
+    def __init__(self, protocol: dict):
+        self.protocol = protocol
+        self.semantic_fields = list(protocol["semantic_fields"])
+        self.field_specs = list(protocol["field_specs"])
+        self.allocation_menu = [tuple(map(int, a)) for a in protocol["reduced_allocation_menu"]]
+        self.morphology_k = int(protocol["morphology_K"])
+        self.morphology_bits = int(protocol["morphology_bits_per_wave"])
+        self.packet_bytes = int(protocol["packet_bytes"])
+        self.payload_bytes = int(protocol["payload_bytes"])
+        self.crc_bytes = int(protocol["crc_bytes"])
+        self.semantic_profiles = {}
+        for key, bits in protocol["semantic_profiles"].items():
+            left, right = key.split("_")
+            n = int(left[1:])
+            aid = int(right[1:])
+            self.semantic_profiles[(n, aid)] = list(map(int, bits))
+        if self.packet_bytes != 30:
+            raise ValueError(f"Expected exact 30-byte packet protocol, got {self.packet_bytes}")
+
+    @staticmethod
+    def sem_internal_to_packet(sem9: torch.Tensor) -> torch.Tensor:
+        theta = torch.remainder(torch.atan2(sem9[..., 2], sem9[..., 3]), 2 * math.pi) / (2 * math.pi)
+        return torch.stack([
+            sem9[..., 0], sem9[..., 1], theta, sem9[..., 4], sem9[..., 5],
+            sem9[..., 6], sem9[..., 7], sem9[..., 8],
+        ], dim=-1)
+
+    @staticmethod
+    def sem_packet_to_internal(p8: torch.Tensor) -> torch.Tensor:
+        theta = p8[..., 2] * (2 * math.pi)
+        return torch.stack([
+            p8[..., 0], p8[..., 1], torch.sin(theta), torch.cos(theta),
+            p8[..., 3], p8[..., 4], p8[..., 5], p8[..., 6], p8[..., 7],
+        ], dim=-1)
+
+    @staticmethod
+    def _forward_compand(x: torch.Tensor, spec: dict) -> torch.Tensor:
+        kind = spec["kind"]
+        lo, hi = float(spec["lo"]), float(spec["hi"])
+        if kind == "circular":
+            return torch.remainder((x - lo) / (hi - lo), 1.0)
+        if kind == "linear":
+            return ((x - lo) / (hi - lo)).clamp(0, 1)
+        if kind == "log":
+            lx = torch.log(x.clamp_min(lo))
+            l0, l1 = math.log(lo), math.log(hi)
+            return ((lx - l0) / (l1 - l0)).clamp(0, 1)
+        if kind == "asinh":
+            s = float(spec.get("scale", 1.0))
+            a0, a1 = math.asinh(lo / s), math.asinh(hi / s)
+            return ((torch.asinh(x / s) - a0) / (a1 - a0)).clamp(0, 1)
+        raise ValueError(kind)
+
+    @staticmethod
+    def _inverse_compand(y: torch.Tensor, spec: dict) -> torch.Tensor:
+        kind = spec["kind"]
+        lo, hi = float(spec["lo"]), float(spec["hi"])
+        if kind in ("linear", "circular"):
+            x = lo + (hi - lo) * y
+            if kind == "circular":
+                x = lo + torch.remainder(x - lo, hi - lo)
+            return x
+        if kind == "log":
+            l0, l1 = math.log(lo), math.log(hi)
+            return torch.exp(l0 + (l1 - l0) * y)
+        if kind == "asinh":
+            s = float(spec.get("scale", 1.0))
+            a0, a1 = math.asinh(lo / s), math.asinh(hi / s)
+            return s * torch.sinh(a0 + (a1 - a0) * y)
+        raise ValueError(kind)
+
+    def packet_fields_to_unit(self, p8: torch.Tensor) -> torch.Tensor:
+        return torch.stack([
+            self._forward_compand(p8[..., d], self.field_specs[d]) for d in range(8)
+        ], dim=-1)
+
+    def unit_to_packet_fields(self, y8: torch.Tensor) -> torch.Tensor:
+        return torch.stack([
+            self._inverse_compand(y8[..., d], self.field_specs[d]) for d in range(8)
+        ], dim=-1)
+
+    @staticmethod
+    def project_unit_fields(y8: torch.Tensor) -> torch.Tensor:
+        cols = []
+        for d in range(8):
+            v = y8[..., d]
+            cols.append(torch.remainder(v, 1.0) if d == 2 else v.clamp(0, 1))
+        return torch.stack(cols, dim=-1)
+
+    def profile_bits_tensor(self, counts: torch.Tensor, aids: torch.Tensor, device) -> torch.Tensor:
+        rows = []
+        for n, a in zip(counts.detach().cpu().tolist(), aids.detach().cpu().tolist()):
+            rows.append(self.semantic_profiles[(max(1, int(n)), int(a))])
+        return torch.tensor(rows, dtype=torch.long, device=device)
+
+
+class PacketSemanticQATAdapter(nn.Module):
+    """C4-B bounded packet-domain semantic pre-quantization adapter."""
+
+    def __init__(self, width: int = 96, delta_steps: float = 1.25, protocol: ProtocolRuntime | None = None):
+        super().__init__()
+        self.delta_steps = float(delta_steps)
+        self.protocol = protocol
+        self.net = nn.Sequential(
+            nn.LayerNorm(21),
+            nn.Linear(21, width), nn.GELU(),
+            nn.Linear(width, width), nn.GELU(),
+            nn.Linear(width, 8),
+        )
+        nn.init.zeros_(self.net[-1].weight)
+        nn.init.zeros_(self.net[-1].bias)
+
+    def set_protocol(self, protocol: ProtocolRuntime) -> None:
+        self.protocol = protocol
+
+    def _proto(self) -> ProtocolRuntime:
+        if self.protocol is None:
+            raise RuntimeError("PacketSemanticQATAdapter requires a ProtocolRuntime")
+        return self.protocol
+
+    def adjust_unit(self, sem9, counts, aids, active):
+        proto = self._proto()
+        p8 = proto.sem_internal_to_packet(sem9)
+        y = proto.packet_fields_to_unit(p8)
+        bits = proto.profile_bits_tensor(counts, aids, sem9.device)
+        levels = (torch.pow(2.0, bits.float()) - 1.0).clamp_min(1.0)
+        step = 1.0 / levels
+        bit_ctx = (bits.float() / 9.0)[:, None, :].expand(-1, sem9.shape[1], -1)
+        cnt_ctx = (counts.float() / float(K_MAX))[:, None, None].expand(-1, sem9.shape[1], 1)
+        aid_ctx = F.one_hot(aids.long(), num_classes=4).float()[:, None, :].expand(-1, sem9.shape[1], -1)
+        x = torch.cat([y, bit_ctx, cnt_ctx, aid_ctx], dim=-1)
+        raw = self.net(x)
+        delta = self.delta_steps * step[:, None, :] * torch.tanh(raw)
+        y_adj = proto.project_unit_fields(y + delta * active[..., None])
+        return y, y_adj, step, bits
+
+    def fake_quantize(self, sem9, counts, aids, active):
+        proto = self._proto()
+        y, y_adj, step, bits = self.adjust_unit(sem9, counts, aids, active)
+        levels = (torch.pow(2.0, bits.float()) - 1.0).clamp_min(1.0)[:, None, :]
+        hard = torch.round(y_adj * levels) / levels
+        yq = y_adj + (hard - y_adj).detach()
+        semq = proto.sem_packet_to_internal(proto.unit_to_packet_fields(yq))
+        return semq, dict(y=y, y_adj=y_adj, step=step, bits=bits, hard=hard)
+
+    @torch.no_grad()
+    def hard_quantize(self, sem9, counts, aids, active):
+        proto = self._proto()
+        y, y_adj, step, bits = self.adjust_unit(sem9, counts, aids, active)
+        levels = (torch.pow(2.0, bits.float()) - 1.0).clamp_min(1.0)[:, None, :]
+        q = torch.round(y_adj * levels).long().clamp_min(0)
+        hard = q.float() / levels
+        semq = proto.sem_packet_to_internal(proto.unit_to_packet_fields(hard))
+        return q, semq, dict(y=y, y_adj=y_adj, step=step, bits=bits, hard=hard)
+
+    @property
+    def num_params(self):
+        return sum(p.numel() for p in self.parameters())
+
+
+class PacketAwareController(nn.Module):
+    def __init__(self,input_dim,width=160,dropout=0.08,n_alloc=4):
+        super().__init__(); self.norm=nn.LayerNorm(input_dim+1)
+        self.fc1=nn.Linear(input_dim+1,width); self.fc2=nn.Linear(width,width); self.fc3=nn.Linear(width,width)
+        self.drop=nn.Dropout(dropout); self.head=nn.Linear(width,n_alloc)
+    def forward(self,x,lambda_norm):
+        if lambda_norm.ndim==1: lambda_norm=lambda_norm[:,None]
+        h=self.norm(torch.cat([x,lambda_norm],dim=-1)); h=F.gelu(self.fc1(h)); r=h
+        h=self.drop(F.gelu(self.fc2(h))); h=F.gelu(self.fc3(h)+r); return self.head(h)
+
+
+class ICWDSFrontEndCodec(nn.Module):
+    """
+    Consolidated Stage-2 front-end feature codec.
+
+    Historical project naming keeps the archive name ``VAE.pt`` for compatibility,
+    but the deployed object is a hybrid semantic/morphology/residual packet codec:
+      C2 AC-SSVAE backbone + C3-D morphology selector + C3-F dual residual codec
+      + C4-B packet QAT adapter + packet-aware allocation controller.
+    """
+
+    def __init__(
+        self,
+        base_model: V5AwareStructuredSemanticVAE,
+        morph_selector: PrototypeSelector,
+        residual_codec: AdaptiveDualResidualCodec,
+        qat_adapter: PacketSemanticQATAdapter,
+        packet_controller: PacketAwareController,
+        protocol_runtime: ProtocolRuntime,
+        morph_centers: torch.Tensor,
+        selector_x_mean: torch.Tensor,
+        selector_x_std: torch.Tensor,
+        feature_mean: torch.Tensor,
+        feature_std: torch.Tensor,
+        lambdas: list[float],
+    ):
+        super().__init__()
+        self.base_model = base_model
+        self.morph_selector = morph_selector
+        self.residual_codec = residual_codec
+        self.qat_adapter = qat_adapter
+        self.packet_controller = packet_controller
+        self.protocol_runtime = protocol_runtime
+        self.qat_adapter.set_protocol(protocol_runtime)
+        self.register_buffer("morph_centers", morph_centers.float())
+        self.register_buffer("selector_x_mean", selector_x_mean.float())
+        self.register_buffer("selector_x_std", selector_x_std.float().clamp_min(1e-5))
+        self.register_buffer("feature_mean", feature_mean.float())
+        self.register_buffer("feature_std", feature_std.float().clamp_min(1e-5))
+        self.lambdas = [float(x) for x in lambdas]
+
+    @torch.no_grad()
+    def decode_semantic_prior(self, sem9, morph_idx, active):
+        z = self.morph_centers.to(sem9.device)[morph_idx.long()]
+        dec = self.base_model.decode_from_codes(sem9, z, exist_prob=active)
+        return dec["part_hat"].sum(1).clamp(0, 1)
+
+    def quantize_semantics(self, sem9, counts, allocation_ids, active, hard=True):
+        if hard:
+            return self.qat_adapter.hard_quantize(sem9, counts, allocation_ids, active)
+        return self.qat_adapter.fake_quantize(sem9, counts, allocation_ids, active)
+
+    @torch.no_grad()
+    def choose_allocation(self, features: torch.Tensor, mode: str = "quality") -> torch.Tensor:
+        if mode not in {"quality", "comm"}:
+            raise ValueError("mode must be 'quality' or 'comm'")
+        lam = self.lambdas[0] if mode == "quality" else self.lambdas[-1]
+        max_lam = max(self.lambdas)
+        x = (features - self.feature_mean) / self.feature_std
+        ln = x.new_full((len(x), 1), lam / max(max_lam, 1e-12))
+        return self.packet_controller(x, ln).argmax(-1)
+
+    @classmethod
+    def from_archive(cls, archive_path: str, map_location: str | torch.device = "cpu"):
+        bundle = torch.load(archive_path, map_location=map_location, weights_only=False)
+        if bundle.get("format") != "ICWDS_VAE_ARCHIVE_V1":
+            raise ValueError("Not an ICWDS consolidated VAE.pt archive")
+
+        c2 = bundle["c2_checkpoint"]
+        c3d = bundle["c3d_package"]
+        c3f = bundle["c3f_package"]
+        c4b = bundle["c4b_package"]
+
+        cfg_dict = c2.get("config", {})
+        allowed = set(V5AwareACSSVAEConfig.__dataclass_fields__)
+        cfg = V5AwareACSSVAEConfig(**{k: v for k, v in cfg_dict.items() if k in allowed})
+        base_model = V5AwareStructuredSemanticVAE(cfg)
+        base_model.load_state_dict(c2["model"], strict=True)
+
+        morph_k = int(c4b.get("protocol", {}).get("morphology_K", 4))
+        sel_item = c3d["selectors"][morph_k]
+        morph_selector = PrototypeSelector(
+            input_dim=int(c3d["input_dim"]),
+            n_codes=morph_k,
+            width=int(sel_item.get("width", 192)),
+            dropout=float(sel_item.get("dropout", 0.10)),
+        )
+        morph_selector.load_state_dict(sel_item["state_dict"], strict=True)
+
+        c3f_ctrl = c3f.get("control", {})
+        residual_codec = AdaptiveDualResidualCodec(
+            max_bits=int(c3f_ctrl.get("MAX_BITS_PER_BRANCH", 24)),
+            enc_width=int(c3f_ctrl.get("ENC_WIDTH", 32)),
+            dec_width=int(c3f_ctrl.get("DEC_WIDTH", 32)),
+        )
+        residual_codec.load_state_dict(c3f["codec"], strict=True)
+
+        proto = ProtocolRuntime(c4b["protocol"])
+        qa = c4b["qat_adapter_arch"]
+        qat_adapter = PacketSemanticQATAdapter(
+            width=int(qa["width"]),
+            delta_steps=float(qa["delta_steps"]),
+            protocol=proto,
+        )
+        qat_adapter.load_state_dict(c4b["qat_adapter"], strict=True)
+
+        pa = c4b["packet_controller_arch"]
+        packet_controller = PacketAwareController(
+            input_dim=int(pa["input_dim"]),
+            width=int(pa["width"]),
+            dropout=float(pa["dropout"]),
+            n_alloc=int(pa["n_alloc"]),
+        )
+        packet_controller.load_state_dict(c4b["packet_controller"], strict=True)
+
+        model = cls(
+            base_model=base_model,
+            morph_selector=morph_selector,
+            residual_codec=residual_codec,
+            qat_adapter=qat_adapter,
+            packet_controller=packet_controller,
+            protocol_runtime=proto,
+            morph_centers=sel_item["centers"],
+            selector_x_mean=c3d["x_mean"],
+            selector_x_std=c3d["x_std"],
+            feature_mean=c4b["feature_mean"],
+            feature_std=c4b["feature_std"],
+            lambdas=list(c4b["lambdas"]),
+        )
+        return model
+
+
+def build_vae_archive_from_files(
+    c2_checkpoint_path: str,
+    c3d_package_path: str,
+    c3f_package_path: str,
+    c4b_package_path: str,
+    output_path: str = "VAE.pt",
+) -> str:
+    """Consolidate the four validated Stage-2 artifacts into one portable VAE.pt."""
+    def load(path):
+        return torch.load(path, map_location="cpu", weights_only=False)
+
+    bundle = {
+        "format": "ICWDS_VAE_ARCHIVE_V1",
+        "description": "C2 AC-SSVAE + C3-D morphology + C3-F residual codec + C4-B QAT/controller",
+        "c2_checkpoint": load(c2_checkpoint_path),
+        "c3d_package": load(c3d_package_path),
+        "c3f_package": load(c3f_package_path),
+        "c4b_package": load(c4b_package_path),
+    }
+    torch.save(bundle, output_path)
+    return output_path
+
+
+def build_vae_archive_from_hf(
+    output_path: str = "VAE.pt",
+    repo_id: str = "wuff-mann/Wave_compress_system",
+    token: str | None = None,
+) -> str:
+    """Download the four frozen artifacts from HF and build one self-contained VAE.pt."""
+    try:
+        from huggingface_hub import hf_hub_download
+    except ImportError as exc:
+        raise RuntimeError("Install huggingface_hub first") from exc
+
+    paths = {
+        "c2": "WaveSemanticHybridCodec/V5_ACSSVAE_C2_CausalProtectedReal/checkpoints/C2_consolidate/complete.pt",
+        "c3d": "WaveSemanticHybridCodec/V5_ACSSVAE_C3D_ReconstructionAwarePrototypeSelector/C3D_prototype_selector_package.pt",
+        "c3f": "WaveSemanticHybridCodec/V5_ACSSVAE_C3F_AdaptiveDualResidualCodec/C3F_adaptive_dual_residual_codec_package.pt",
+        "c4b": "WaveSemanticHybridCodec/V5_ACSSVAE_C4B_PacketAwareQAT_Controller/C4B_packet_qat_controller_package.pt",
+    }
+    local = {
+        k: hf_hub_download(repo_id=repo_id, filename=v, repo_type="model", token=token)
+        for k, v in paths.items()
+    }
+    return build_vae_archive_from_files(local["c2"], local["c3d"], local["c3f"], local["c4b"], output_path)
+
+
+def load_vae_archive(path: str = "VAE.pt", map_location: str | torch.device = "cpu") -> ICWDSFrontEndCodec:
+    return ICWDSFrontEndCodec.from_archive(path, map_location=map_location)
+
+
+__all__ = [
+    "V5AwareACSSVAEConfig",
+    "V5AwareStructuredSemanticVAE",
+    "PrototypeSelector",
+    "AdaptiveDualResidualCodec",
+    "RateConditionedController",
+    "PacketSemanticQATAdapter",
+    "PacketAwareController",
+    "ProtocolRuntime",
+    "ICWDSFrontEndCodec",
+    "build_vae_archive_from_files",
+    "build_vae_archive_from_hf",
+    "load_vae_archive",
+]
+
+
+if __name__ == "__main__":
+    import argparse
+    parser = argparse.ArgumentParser(description="Build consolidated ICWDS VAE.pt archive")
+    parser.add_argument("--build-from-hf", action="store_true")
+    parser.add_argument("--output", default="VAE.pt")
+    parser.add_argument("--repo", default="wuff-mann/Wave_compress_system")
+    parser.add_argument("--token", default=None)
+    args = parser.parse_args()
+    if args.build_from_hf:
+        out = build_vae_archive_from_hf(args.output, args.repo, args.token)
+        print(out)