""" 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)