#!/usr/bin/env python3 """ Per-scene autoencoder for compressing f_rest* attributes. Provides train / encode / decode commands. """ from __future__ import annotations import argparse import io import json import subprocess import sys from functools import partial from itertools import cycle from pathlib import Path from typing import Dict, Iterable, List, Tuple, Sequence import numpy as np import torch from torch import nn from torch.utils.data import DataLoader, Dataset import shutil from tqdm.auto import tqdm try: import zstandard as zstd except ImportError: # pragma: no cover - optional dependency zstd = None REPO_ROOT = Path(__file__).resolve().parents[1] if str(REPO_ROOT) not in sys.path: sys.path.insert(0, str(REPO_ROOT)) from tools.compress_point_cloud import _load_ply, _write_ply from tools.entropy_pack_point_cloud import ( compress_point_cloud as entropy_compress, decompress_point_cloud as entropy_decompress, ) from tools import restfree_quant_pack ZSTD_MAGIC = b"\x28\xb5\x2f\xfd" ANSI_BLUE = "\033[34m" ANSI_RESET = "\033[0m" def _require_zstd() -> None: if zstd is None: raise RuntimeError("zstandard module not available. Please install 'zstandard' to use npz_zstd compression.") def _write_latents_pack(path: Path, payload: Dict[str, object], compression: str) -> None: buffer = io.BytesIO() np.savez_compressed(buffer, **payload) data = buffer.getvalue() if compression == "npz": path.write_bytes(data) elif compression == "npz_zstd": _require_zstd() compressor = zstd.ZstdCompressor(level=10) path.write_bytes(compressor.compress(data)) else: raise ValueError(f"Unsupported latents compression '{compression}'") def _load_latents_pack(path: Path): with path.open("rb") as fh: magic = fh.read(4) if magic == ZSTD_MAGIC: _require_zstd() compressed = path.read_bytes() decompressed = zstd.ZstdDecompressor().decompress(compressed) return np.load(io.BytesIO(decompressed), allow_pickle=False) return np.load(path, allow_pickle=False) def _format_mb_colored(value: float) -> str: return f"{ANSI_BLUE}{value:.2f} MB{ANSI_RESET}" def _select_uint_dtype(bits: int): if bits <= 8: return np.uint8 if bits <= 16: return np.uint16 if bits <= 32: return np.uint32 raise ValueError("quant_bits must be <=32") def _initialise_codebook(data: torch.Tensor, k: int, generator: torch.Generator) -> torch.Tensor: perm = torch.randperm(data.shape[0], generator=generator, device=data.device) return data[perm[:k]].clone() def _assign_points(data: torch.Tensor, centers: torch.Tensor, chunk_size: int) -> torch.Tensor: assignments = [] n = data.shape[0] for start in range(0, n, chunk_size): end = min(start + chunk_size, n) batch = data[start:end] dists = torch.cdist(batch, centers) assignments.append(torch.argmin(dists, dim=1)) return torch.cat(assignments) def _update_centers( data: torch.Tensor, assignments: torch.Tensor, centers: torch.Tensor, chunk_size: int, generator: torch.Generator, ) -> torch.Tensor: k = centers.shape[0] new_centers = torch.zeros_like(centers) counts = torch.zeros(k, device=data.device, dtype=torch.float32) n = data.shape[0] for start in range(0, n, chunk_size): end = min(start + chunk_size, n) batch = data[start:end] idx = assignments[start:end] new_centers.index_add_(0, idx, batch) counts.index_add_(0, idx, torch.ones_like(idx, dtype=torch.float32)) nonzero = counts > 0 if nonzero.any(): new_centers[nonzero] /= counts[nonzero].unsqueeze(1) zero_mask = ~nonzero if zero_mask.any(): replace = data[torch.randint(0, n, (int(zero_mask.sum().item()),), device=data.device, generator=generator)] new_centers[zero_mask] = replace return new_centers def _kmeans_codebook( data: np.ndarray, *, k: int, iters: int, chunk_size: int, device: torch.device, seed: int, ) -> Tuple[np.ndarray, np.ndarray]: torch.manual_seed(seed) generator = torch.Generator(device=device) generator.manual_seed(seed) tensor = torch.from_numpy(data).to(device) centers = _initialise_codebook(tensor, k, generator) assignments = torch.zeros(tensor.shape[0], dtype=torch.long, device=device) for _ in range(iters): assignments = _assign_points(tensor, centers, chunk_size) centers = _update_centers(tensor, assignments, centers, chunk_size, generator) assignments = _assign_points(tensor, centers, chunk_size) return centers.cpu().numpy().astype(np.float32), assignments.cpu().numpy() def _resolve_ply(path: Path) -> Path: path = Path(path) if path.is_dir(): path = path / "point_cloud.ply" if not path.exists(): raise FileNotFoundError(f"PLY not found: {path}") return path def _iteration_dir(model_path: Path, iteration: int) -> Path: return model_path / "point_cloud" / f"iteration_{iteration}" def _strip_rest_fields(source_ply: Path, dest_ply: Path, rest_names: List[str]) -> None: arrays, names = _load_ply(source_ply) filtered = [name for name in names if name not in rest_names] data = {name: arrays[name] for name in filtered} _write_ply(dest_ply, data, filtered) print(f"[autoencoder-auto] Wrote rest-free PLY to {dest_ply}") def _human_mb(path: Path) -> float: return path.stat().st_size / (1024 * 1024) def _stack_fields(arrays: Dict[str, np.ndarray], names: Sequence[str]) -> np.ndarray: if not names: return np.empty((len(next(iter(arrays.values()))), 0), dtype=np.float32) return np.stack([arrays[name] for name in names], axis=1).astype(np.float32) def _decode_latent_arrays( latents_path: Path, checkpoint: Path, *, chunk: int, device: str, ) -> Tuple[np.ndarray, List[str]]: with _load_latents_pack(latents_path) as latents_pack: meta = json.loads(str(latents_pack["meta"])) if meta.get("quant_bits"): shape = tuple(meta.get("latents_shape") or (meta["count"], meta["latent_dim"])) if meta.get("bitpack_format") == "bitplane_little": latents_q = _bitunpack_latents( latents_pack["latents_bitpack"], shape, meta["quant_bits"], ) else: latents_q = latents_pack["latents_q"] latents = _dequantise_latents( latents_q, latents_pack["latents_min"], latents_pack["latents_max"], meta["quant_bits"], ) else: latents = latents_pack["latents"] model, target_mean, target_std, _, _, target_names, _ = _load_checkpoint(checkpoint, device) target_mean_t = torch.from_numpy(target_mean).to(device) target_std_t = torch.from_numpy(target_std).to(device) reconstructed = np.empty((latents.shape[0], len(target_names)), dtype=np.float32) chunk_size = latents.shape[0] if chunk <= 0 else chunk with torch.no_grad(): for start in range(0, latents.shape[0], chunk_size): end = start + chunk_size batch = torch.from_numpy(latents[start:end].astype(np.float32)).to(device) rec = model.decode_batch(batch) * target_std_t + target_mean_t reconstructed[start:end] = rec.cpu().numpy() return reconstructed, target_names def _expand_field_specs(arrays: Dict[str, np.ndarray], specs: Sequence[str]) -> List[str]: expanded: List[str] = [] for spec in specs: if spec in arrays: expanded.append(spec) continue matches = [name for name in arrays if name.startswith(f"{spec}_")] if not matches: raise KeyError(f"Field spec '{spec}' not found in PLY attributes.") expanded.extend(sorted(matches, key=lambda x: int(x.split("_")[-1]) if x.split("_")[-1].isdigit() else x)) return expanded class LatentDataset(Dataset): def __init__( self, target: np.ndarray, cond: np.ndarray, target_mean: np.ndarray, target_std: np.ndarray, cond_mean: np.ndarray, cond_std: np.ndarray, ) -> None: self.target = torch.from_numpy((target - target_mean) / target_std) if cond.size: self.cond = torch.from_numpy((cond - cond_mean) / cond_std) else: self.cond = torch.empty((target.shape[0], 0), dtype=torch.float32) def __len__(self) -> int: return self.target.shape[0] def __getitem__(self, idx: int) -> torch.Tensor: if self.cond.numel(): return torch.cat([self.target[idx], self.cond[idx]], dim=0), self.target[idx] return self.target[idx], self.target[idx] class RestAutoencoder(nn.Module): def __init__(self, input_dim: int, latent_dim: int, hidden_dim: int, num_layers: int, dropout: float = 0.0, output_dim: int | None = None): super().__init__() self.encoder = self._build_mlp(input_dim, latent_dim, hidden_dim, num_layers, dropout, final_activation=False) out_dim = input_dim if output_dim is None else output_dim self.decoder = self._build_mlp(latent_dim, out_dim, hidden_dim, num_layers, dropout, final_activation=False) @staticmethod def _build_mlp(inp: int, out: int, hidden: int, layers: int, dropout: float, final_activation: bool) -> nn.Sequential: blocks: List[nn.Module] = [] last = inp for _ in range(layers): blocks.append(nn.Linear(last, hidden)) blocks.append(nn.GELU()) if dropout > 0: blocks.append(nn.Dropout(dropout)) last = hidden blocks.append(nn.Linear(last, out)) if final_activation: blocks.append(nn.GELU()) return nn.Sequential(*blocks) def forward(self, inputs: torch.Tensor) -> torch.Tensor: z = self.encoder(inputs) recon = self.decoder(z) return recon def encode_batch(self, x: torch.Tensor) -> torch.Tensor: return self.encoder(x) def decode_batch(self, z: torch.Tensor) -> torch.Tensor: return self.decoder(z) def _quantise_latents(values: np.ndarray, bits: int) -> Tuple[np.ndarray, np.ndarray, np.ndarray]: mins = values.min(axis=0) maxs = values.max(axis=0) spans = np.where(maxs > mins, maxs - mins, 1.0) levels = (1 << bits) - 1 scaled = np.round((values - mins) / spans * levels) dtype = _select_uint_dtype(bits) q = scaled.clip(0, levels).astype(dtype) return q, mins.astype(np.float32), maxs.astype(np.float32) def _dequantise_latents(q: np.ndarray, mins: np.ndarray, maxs: np.ndarray, bits: int) -> np.ndarray: levels = (1 << bits) - 1 spans = np.where(maxs > mins, maxs - mins, 1.0) return mins + (q.astype(np.float32) / levels) * spans def _bitpack_latents(q: np.ndarray, bits: int) -> np.ndarray: flat = q.reshape(-1).astype(np.uint32) count = flat.size bytes_per_plane = (count + 7) // 8 packed = np.empty((bits, bytes_per_plane), dtype=np.uint8) for b in range(bits): plane = np.packbits(((flat >> b) & 1).astype(np.uint8), bitorder="little") if plane.size < bytes_per_plane: plane = np.pad(plane, (0, bytes_per_plane - plane.size), constant_values=0) packed[b] = plane return packed.reshape(-1) def _bitunpack_latents(buffer: np.ndarray, shape: Sequence[int], bits: int) -> np.ndarray: count = int(np.prod(shape)) bytes_per_plane = (count + 7) // 8 expected = bytes_per_plane * bits if buffer.size != expected: raise ValueError( f"Bitpacked latents size mismatch: expected {expected} bytes, got {buffer.size}" ) flat = np.zeros(count, dtype=np.uint32) for b in range(bits): start = b * bytes_per_plane plane = buffer[start : start + bytes_per_plane] bits_arr = np.unpackbits(plane, bitorder="little")[:count].astype(np.uint32) flat |= bits_arr << b dtype = _select_uint_dtype(bits) return flat.astype(dtype).reshape(shape) def train_autoencoder( input_path: Path, output_dir: Path, *, latent_dim: int, hidden_dim: int, num_layers: int, dropout: float, target_fields: List[str], condition_fields: List[str], batch_size: int, steps: int, lr: float, device: str, log_every: int, seed: int, checkpoint_dtype: str, ) -> None: arrays, names = _load_ply(input_path) target_names = _expand_field_specs(arrays, target_fields) if not target_names: raise RuntimeError("No target fields selected for autoencoder.") rest = _stack_fields(arrays, target_names) for field in condition_fields: if field not in arrays: raise KeyError(f"Condition field '{field}' not found in PLY.") cond = _stack_fields(arrays, condition_fields) torch.manual_seed(seed) np.random.seed(seed) target_mean = rest.mean(axis=0, keepdims=True) target_std = rest.std(axis=0, keepdims=True) target_std = np.where(target_std < 1e-6, 1.0, target_std) if cond.size: cond_mean = cond.mean(axis=0, keepdims=True) cond_std = cond.std(axis=0, keepdims=True) cond_std = np.where(cond_std < 1e-6, 1.0, cond_std) else: cond_mean = np.zeros((1, 0), dtype=np.float32) cond_std = np.ones((1, 0), dtype=np.float32) dataset = LatentDataset(rest, cond, target_mean, target_std, cond_mean, cond_std) loader = DataLoader(dataset, batch_size=batch_size, shuffle=True, drop_last=True) iterator = cycle(loader) input_dim = rest.shape[1] + cond.shape[1] model = RestAutoencoder(input_dim, latent_dim, hidden_dim, num_layers, dropout, output_dim=rest.shape[1]).to(device) optimizer = torch.optim.AdamW(model.parameters(), lr=lr) criterion = nn.MSELoss() output_dir.mkdir(parents=True, exist_ok=True) model.train() progress = tqdm(range(1, steps + 1), desc="[autoencoder] training", leave=True) for step in progress: inputs, targets = next(iterator) inputs = inputs.to(device) targets = targets.to(device) optimizer.zero_grad(set_to_none=True) recon = model(inputs) loss = criterion(recon, targets) loss.backward() optimizer.step() progress.set_postfix(loss=f"{loss.item():.6f}") # Progress bar already shows loss; suppress per-step prints for cleanliness. progress.close() raw_state = {k: v.detach().cpu() for k, v in model.state_dict().items()} if checkpoint_dtype == "float16": save_state = {k: v.half() for k, v in raw_state.items()} elif checkpoint_dtype == "float32": save_state = raw_state else: raise ValueError(f"Unsupported checkpoint dtype '{checkpoint_dtype}'") def _norm_array(arr: np.ndarray) -> np.ndarray: if checkpoint_dtype == "float16": return arr.astype(np.float16) return arr.astype(np.float32) ckpt = { "state_dict": save_state, "config": { "input_dim": input_dim, "latent_dim": latent_dim, "hidden_dim": hidden_dim, "num_layers": num_layers, "dropout": dropout, }, "normalizer": { "target_mean": _norm_array(target_mean), "target_std": _norm_array(target_std), "cond_mean": _norm_array(cond_mean), "cond_std": _norm_array(cond_std), }, "target_names": target_names, "condition_fields": condition_fields, "steps": steps, "checkpoint_dtype": checkpoint_dtype, } ckpt_path = output_dir / "rest_autoencoder.pt" torch.save(ckpt, ckpt_path) print(f"[autoencoder] Saved checkpoint to {ckpt_path}") def _torch_load_trusted(path: Path, map_location): try: return torch.load(path, map_location=map_location, weights_only=False) except TypeError: return torch.load(path, map_location=map_location) def _load_checkpoint(path: Path, device: str): ckpt = _torch_load_trusted(path, map_location=device) cfg = ckpt["config"] ckpt_dtype = ckpt.get("checkpoint_dtype", "float32") state_dict = ckpt["state_dict"] if ckpt_dtype == "float16": state_dict = {k: v.to(torch.float32) for k, v in state_dict.items()} model = RestAutoencoder( cfg["input_dim"], cfg["latent_dim"], cfg["hidden_dim"], cfg["num_layers"], cfg.get("dropout", 0.0), output_dim=ckpt["normalizer"]["target_mean"].shape[1], ).to(device) model.load_state_dict(state_dict) model.eval() norm = ckpt["normalizer"] target_mean = norm["target_mean"].astype(np.float32) target_std = norm["target_std"].astype(np.float32) cond_mean = norm["cond_mean"].astype(np.float32) cond_std = norm["cond_std"].astype(np.float32) target_names = ckpt["target_names"] cond_fields = ckpt.get("condition_fields", []) return model, target_mean, target_std, cond_mean, cond_std, target_names, cond_fields def encode_rest( input_path: Path, checkpoint: Path, output_path: Path, *, dtype: str, chunk: int, device: str, quant_bits: int | None, latents_compression: str = "npz", bitpack_latents: bool = False, ) -> None: input_path = _resolve_ply(input_path) arrays, names = _load_ply(input_path) model, target_mean, target_std, cond_mean, cond_std, target_names, cond_fields = _load_checkpoint(checkpoint, device) rest = _stack_fields(arrays, target_names) cond = _stack_fields(arrays, cond_fields) dtype_map = { "float32": np.float32, "float16": np.float16, "bfloat16": np.float32, # PyTorch lacks BF16 numpy } if dtype not in dtype_map: raise ValueError(f"Unsupported dtype {dtype}") target_dtype = dtype_map[dtype] rest_mean_t = torch.from_numpy(target_mean).to(device) rest_std_t = torch.from_numpy(target_std).to(device) cond_mean_t = torch.from_numpy(cond_mean).to(device) if cond.size else None cond_std_t = torch.from_numpy(cond_std).to(device) if cond.size else None latents_list: List[np.ndarray] = [] chunk_size = rest.shape[0] if chunk <= 0 else chunk with torch.no_grad(): for start in range(0, rest.shape[0], chunk_size): end = start + chunk_size rest_batch = torch.from_numpy(rest[start:end]).to(device) rest_norm = (rest_batch - rest_mean_t) / rest_std_t if cond.size: cond_batch = torch.from_numpy(cond[start:end]).to(device) cond_norm = (cond_batch - cond_mean_t) / cond_std_t inputs = torch.cat([rest_norm, cond_norm], dim=1) else: inputs = rest_norm z = model.encode_batch(inputs).cpu().numpy() latents_list.append(z.astype(target_dtype)) latents = np.concatenate(latents_list, axis=0) meta = { "checkpoint": str(checkpoint), "input_ply": str(input_path), "latent_dim": latents.shape[1], "dtype": dtype, "count": latents.shape[0], "target_names": target_names, "condition_fields": cond_fields, "latents_shape": latents.shape, } if quant_bits is not None: q, lat_min, lat_max = _quantise_latents(latents.astype(np.float32), quant_bits) meta["quant_bits"] = quant_bits if bitpack_latents: packed = _bitpack_latents(q, quant_bits) meta["bitpack_format"] = "bitplane_little" payload = { "latents_bitpack": packed, "latents_min": lat_min, "latents_max": lat_max, "meta": np.array(json.dumps(meta)), } else: payload = { "latents_q": q, "latents_min": lat_min, "latents_max": lat_max, "meta": np.array(json.dumps(meta)), } else: meta["quant_bits"] = None payload = { "latents": latents, "meta": np.array(json.dumps(meta)), } _write_latents_pack(output_path, payload, latents_compression) print(f"[autoencoder] Encoded latents saved to {output_path}") def decode_rest( latents_path: Path, checkpoint: Path, source_ply: Path, output_ply: Path, *, chunk: int, device: str, ) -> None: with _load_latents_pack(latents_path) as latents_pack: meta = json.loads(str(latents_pack["meta"])) if meta.get("quant_bits"): shape = tuple(meta.get("latents_shape") or (meta["count"], meta["latent_dim"])) if meta.get("bitpack_format") == "bitplane_little": latents_q = _bitunpack_latents( latents_pack["latents_bitpack"], shape, meta["quant_bits"], ) else: latents_q = latents_pack["latents_q"] latents = _dequantise_latents( latents_q, latents_pack["latents_min"], latents_pack["latents_max"], meta["quant_bits"], ) else: latents = latents_pack["latents"] model, target_mean, target_std, _, _, target_names, _ = _load_checkpoint(checkpoint, device) source_ply = _resolve_ply(source_ply) arrays, names = _load_ply(source_ply) count = latents.shape[0] names_out = list(names) for name in target_names: if name not in arrays: arrays[name] = np.zeros(count, dtype=np.float32) names_out.append(name) if latents.shape[0] != arrays[target_names[0]].shape[0]: raise ValueError("Latent count mismatch with source PLY.") target_mean_t = torch.from_numpy(target_mean).to(device) target_std_t = torch.from_numpy(target_std).to(device) reconstructed = np.empty((latents.shape[0], len(target_names)), dtype=np.float32) chunk_size = latents.shape[0] if chunk <= 0 else chunk with torch.no_grad(): for start in range(0, latents.shape[0], chunk_size): end = start + chunk_size batch = torch.from_numpy(latents[start:end].astype(np.float32)).to(device) rec = model.decode_batch(batch) * target_std_t + target_mean_t reconstructed[start:end] = rec.cpu().numpy() for idx, name in enumerate(target_names): arrays[name] = reconstructed[:, idx] _write_ply(output_ply, arrays, names_out) print(f"[autoencoder] Reconstructed PLY saved to {output_ply}") def _fit_residual_codebook( arrays: Dict[str, np.ndarray], latents_path: Path, checkpoint: Path, *, chunk: int, device: str, condition_fields: List[str], condition_weight: float, codebook_size: int, iters: int, seed: int, chunk_size: int, output_dir: Path, ) -> Tuple[Dict[str, Path], List[str]]: rest_hat, rest_names = _decode_latent_arrays(latents_path, checkpoint, chunk=chunk, device=device) rest_original = _stack_fields(arrays, rest_names) residual = rest_original - rest_hat residual_mean = residual.mean(axis=0, keepdims=True) residual_std = residual.std(axis=0, keepdims=True) residual_std = np.where(residual_std < 1e-6, 1.0, residual_std) residual_norm = (residual - residual_mean) / residual_std cond = _stack_fields(arrays, condition_fields) if cond.size: cond_mean = cond.mean(axis=0, keepdims=True) cond_std = cond.std(axis=0, keepdims=True) cond_std = np.where(cond_std < 1e-6, 1.0, cond_std) cond_norm = (cond - cond_mean) / cond_std fused = np.concatenate([residual_norm, condition_weight * cond_norm], axis=1) else: fused = residual_norm codebook_norm, assignments = _kmeans_codebook( fused, k=codebook_size, iters=iters, chunk_size=fused.shape[0] if chunk_size <= 0 else chunk_size, device=torch.device(device), seed=seed, ) rng = np.random.default_rng(seed) codebook = np.zeros((codebook_size, residual.shape[1]), dtype=np.float32) for idx in range(codebook_size): mask = assignments == idx if mask.any(): codebook[idx] = residual[mask].mean(axis=0) else: codebook[idx] = residual[rng.integers(0, residual.shape[0])] dtype = _select_uint_dtype(int(np.ceil(np.log2(max(codebook_size, 1)))) if codebook_size > 0 else 1) assign_quant = assignments.astype(dtype) codebook_path = output_dir / "residual_codebook.npy" assignments_path = output_dir / "residual_assignments.npy" np.save(codebook_path, codebook) np.save(assignments_path, assign_quant) return {"codebook": codebook_path, "assignments": assignments_path}, rest_names def _apply_residual_to_ply(ply_path: Path, rest_names: List[str], codebook_path: Path, assignments_path: Path) -> None: if not codebook_path.exists() or not assignments_path.exists(): return codebook = np.load(codebook_path) assignments = np.load(assignments_path) arrays, names = _load_ply(ply_path) residual = codebook[assignments] for idx, name in enumerate(rest_names): arrays[name] = arrays[name] + residual[:, idx] _write_ply(ply_path, arrays, names) def _write_compressed_bundle( iteration_dir: Path, *, target_free_ply: Path, latents_path: Path, checkpoint_path: Path, quantize_groups: str, xyz_bits: int, scale_bits: int, opacity_bits: int, rotation_bits: int, dynamic_bits: int, dynamic_mu: float, residual_info: Dict[str, Path] | None = None, ) -> None: compressed_dir = iteration_dir / "compressed" if compressed_dir.exists(): shutil.rmtree(compressed_dir) compressed_dir.mkdir(parents=True, exist_ok=True) pack_dir = compressed_dir / "restfree" restfree_quant_pack.pack_restfree( input_ply=target_free_ply, output_dir=pack_dir, xyz_bits=xyz_bits, scale_bits=scale_bits, opacity_bits=opacity_bits, rotation_bits=rotation_bits, dynamic_bits=dynamic_bits, dynamic_mu=dynamic_mu, quantize_groups=quantize_groups, ) shutil.copy2(latents_path, compressed_dir / "latents.npz") shutil.copy2(checkpoint_path, compressed_dir / "rest_autoencoder.pt") if residual_info: shutil.copy2(residual_info["codebook"], compressed_dir / "residual_codebook.npy") shutil.copy2(residual_info["assignments"], compressed_dir / "residual_assignments.npy") manifest = { "target_free_source": str(target_free_ply), "quantize_groups": quantize_groups, "xyz_bits": xyz_bits, "scale_bits": scale_bits, "opacity_bits": opacity_bits, "rotation_bits": rotation_bits, "dynamic_bits": dynamic_bits, "dynamic_mu": dynamic_mu, } (compressed_dir / "restfree_manifest.json").write_text(json.dumps(manifest, indent=2)) print(f"[autoencoder-auto] Wrote compressed assets to {compressed_dir}") def run_auto_pipeline(args) -> None: model_path = args.model_path src_dir = _iteration_dir(model_path, args.source_iteration) src_ply = _resolve_ply(src_dir) checkpoint_dir = model_path / "autoencoder" / f"iteration_{args.iteration}" checkpoint_dir.mkdir(parents=True, exist_ok=True) print(f"[autoencoder-auto] Training on {src_ply} ...") train_autoencoder( src_ply, checkpoint_dir, latent_dim=args.latent_dim, hidden_dim=args.hidden_dim, num_layers=args.num_layers, dropout=args.dropout, target_fields=args.target_fields, condition_fields=args.condition_fields, batch_size=args.batch_size, steps=args.steps, lr=args.lr, device=args.device, log_every=args.log_every, seed=args.seed, checkpoint_dtype=args.checkpoint_dtype, ) checkpoint = checkpoint_dir / "rest_autoencoder.pt" latents_path = checkpoint_dir / "latents.npz" print("[autoencoder-auto] Encoding latents ...") encode_rest( src_ply, checkpoint, latents_path, dtype=args.dtype, chunk=args.chunk, device=args.device, quant_bits=args.quant_bits, latents_compression=args.latents_compression, bitpack_latents=args.bitpack_latents, ) residual_info = None residual_rest_names: List[str] = [] if args.residual_codebook_size > 0: arrays_src, _ = _load_ply(src_ply) residual_info, residual_rest_names = _fit_residual_codebook( arrays_src, latents_path=latents_path, checkpoint=checkpoint, chunk=args.chunk, device=args.device, condition_fields=args.condition_fields, condition_weight=args.residual_condition_weight, codebook_size=args.residual_codebook_size, iters=args.residual_iters, seed=args.seed + 123, chunk_size=args.chunk, output_dir=checkpoint_dir, ) model_state_names = ("deformation.pth", "FDHash.pth") target_dir = _iteration_dir(model_path, args.iteration) target_dir.mkdir(parents=True, exist_ok=True) target_ply = target_dir / "point_cloud.ply" decode_source = target_ply if target_ply.exists() and not args.overwrite: print(f"[autoencoder-auto] Target PLY exists at {target_ply}; reusing source without overwrite.") decode_source = src_ply else: print("[autoencoder-auto] Decoding into target iteration ...") decode_rest( latents_path, checkpoint, src_ply, target_ply, chunk=args.chunk, device=args.device, ) if residual_info: _apply_residual_to_ply(target_ply, residual_rest_names, residual_info["codebook"], residual_info["assignments"]) copied_model_state = False for state_name in model_state_names: src_state = src_dir / state_name dst_state = target_dir / state_name if src_state.exists(): shutil.copy2(src_state, dst_state) print(f"[autoencoder-auto] Copied model checkpoint to {dst_state}") copied_model_state = True if not copied_model_state: expected = ", ".join(str(src_dir / name) for name in model_state_names) print(f"[autoencoder-auto] Warning: no model checkpoint found; expected one of: {expected}") ckpt_meta = _torch_load_trusted(checkpoint, map_location="cpu") target_names = ckpt_meta["target_names"] target_free_path = checkpoint_dir / "point_cloud_without_targets.ply" _strip_rest_fields(decode_source, target_free_path, target_names) if args.write_compressed: _write_compressed_bundle( target_dir, target_free_ply=target_free_path, latents_path=latents_path, checkpoint_path=checkpoint, quantize_groups=args.restfree_groups, xyz_bits=args.restfree_xyz_bits, scale_bits=args.restfree_scale_bits, opacity_bits=args.restfree_opacity_bits, rotation_bits=args.restfree_rotation_bits, dynamic_bits=args.restfree_dynamic_bits, dynamic_mu=args.restfree_dynamic_mu, residual_info=residual_info, ) share_artifact_path = target_free_path share_label = "target-free PLY" entropy_npz = None if args.entropy_share: entropy_npz = checkpoint_dir / "target_free_entropy.npz" entropy_compress( target_free_path, entropy_npz, pos_bits=args.entropy_pos_bits, feature_bits=args.entropy_feature_bits, rest_bits=args.entropy_rest_bits, ) share_artifact_path = entropy_npz share_label = "target-free entropy NPZ" original_ply_size = _human_mb(src_ply) share_latents = _human_mb(latents_path) share_ckpt = _human_mb(checkpoint) share_ply = _human_mb(share_artifact_path) residual_codebook_mb = _human_mb(residual_info["codebook"]) if residual_info else 0.0 residual_assign_mb = _human_mb(residual_info["assignments"]) if residual_info else 0.0 total_share = share_latents + share_ckpt + share_ply + residual_codebook_mb + residual_assign_mb reduction = 100.0 * (1 - total_share / original_ply_size) print("[autoencoder-auto] Files to share:") print(f" - {share_label}: {share_artifact_path} ({_format_mb_colored(share_ply)})") print(f" - latents: {latents_path} ({_format_mb_colored(share_latents)})") print(f" - checkpoint: {checkpoint} ({_format_mb_colored(share_ckpt)})") if residual_info: print(f" - residual codebook: {residual_info['codebook']} ({_format_mb_colored(residual_codebook_mb)})") print(f" - residual assignments: {residual_info['assignments']} ({_format_mb_colored(residual_assign_mb)})") print( f" => Combined share size {_format_mb_colored(total_share)} vs original {_format_mb_colored(original_ply_size)} ({reduction:.1f}% smaller)" ) deploy_iter = args.deploy_iteration or args.iteration deploy_dir = _iteration_dir(model_path, deploy_iter) deploy_dir.mkdir(parents=True, exist_ok=True) deploy_ply = deploy_dir / "point_cloud.ply" deploy_source = target_free_path temp_entropy_ply = None if args.entropy_share: temp_entropy_ply = checkpoint_dir / "target_free_entropy_decompressed.ply" entropy_decompress(entropy_npz, temp_entropy_ply) deploy_source = temp_entropy_ply if deploy_ply.exists() and not args.overwrite: print(f"[autoencoder-auto] Deploy PLY exists at {deploy_ply}; skipping decode (use --overwrite to refresh).") else: decode_rest( latents_path, checkpoint, deploy_source, deploy_ply, chunk=args.chunk, device=args.device, ) if residual_info: _apply_residual_to_ply(deploy_ply, residual_rest_names, residual_info["codebook"], residual_info["assignments"]) copied_model_state = False for state_name in model_state_names: src_state = src_dir / state_name dst_state = deploy_dir / state_name if src_state.exists(): shutil.copy2(src_state, dst_state) print(f"[autoencoder-auto] Copied model checkpoint to {dst_state}") copied_model_state = True if not copied_model_state: expected = ", ".join(str(src_dir / name) for name in model_state_names) print(f"[autoencoder-auto] Warning: no model checkpoint found; expected one of: {expected}") if temp_entropy_ply is not None: temp_entropy_ply.unlink(missing_ok=True) if args.run_render: required = [args.dataset_path, args.config_path] if any(v is None for v in required): raise ValueError("--run-render requires --dataset-path and --config-path") render_cmd = [ sys.executable, str(REPO_ROOT / "render.py"), "-s", str(args.dataset_path), "--configs", str(args.config_path), "--model_path", str(model_path), "--iteration", str(deploy_iter), "--skip_train", "--skip_video", ] print("[autoencoder-auto] Running render.py ...") subprocess.run(render_cmd, check=True) if not args.skip_metrics: metrics_cmd = [ sys.executable, str(REPO_ROOT / "metrics.py"), "-m", str(model_path), "--iteration", str(deploy_iter), ] print("[autoencoder-auto] Running metrics.py ...") subprocess.run(metrics_cmd, check=True) print(f"[autoencoder-auto] Finished. Iteration {args.iteration} ready at {target_dir}; deployed to iteration {deploy_iter}") def build_parser() -> argparse.ArgumentParser: parser = argparse.ArgumentParser(description="Autoencoder-based rest compression.") sub = parser.add_subparsers(dest="command", required=True) train = sub.add_parser("train", help="Train autoencoder on f_rest*.") train.add_argument("input", type=Path) train.add_argument("--output-dir", type=Path, required=True) train.add_argument("--latent-dim", type=int, default=12) train.add_argument("--hidden-dim", type=int, default=256) train.add_argument("--num-layers", type=int, default=2) train.add_argument("--dropout", type=float, default=0.0) train.add_argument( "--condition-fields", type=str, nargs="*", default=[], help="Additional PLY attributes to concatenate as conditioning inputs.", ) train.add_argument( "--target-fields", type=str, nargs="*", default=["f_rest"], help="Attribute groups to encode (prefix like f_rest, f_dc, or exact field names).", ) train.add_argument("--batch-size", type=int, default=8192) train.add_argument("--steps", type=int, default=20000) train.add_argument("--lr", type=float, default=1e-3) train.add_argument("--device", type=str, default="cuda") train.add_argument("--log-every", type=int, default=100) train.add_argument("--seed", type=int, default=42) train.add_argument( "--checkpoint-dtype", type=str, choices=["float32", "float16"], default="float32", help="Precision used when serialising the autoencoder checkpoint.", ) encode = sub.add_parser("encode", help="Encode rest vectors into latents.") encode.add_argument("input", type=Path) encode.add_argument("--checkpoint", type=Path, required=True) encode.add_argument("--output", type=Path, required=True) encode.add_argument("--dtype", type=str, default="float16", choices=["float32", "float16"]) encode.add_argument("--chunk", type=int, default=65536) encode.add_argument("--device", type=str, default="cuda") encode.add_argument("--quant-bits", type=int, default=None) encode.add_argument( "--latents-compression", type=str, choices=["npz", "npz_zstd"], default="npz", help="Container format for latents.npz (npz_zstd wraps the NPZ in Zstandard).", ) encode.add_argument( "--bitpack-latents", action="store_true", help="Store quantised latents in a bit-packed representation.", ) decode = sub.add_parser("decode", help="Decode latents and rebuild PLY.") decode.add_argument("latents", type=Path) decode.add_argument("--checkpoint", type=Path, required=True) decode.add_argument("--source", type=Path, required=True, help="Original PLY or iteration directory for non-rest attributes.") decode.add_argument("--output", type=Path, required=True) decode.add_argument("--chunk", type=int, default=65536) decode.add_argument("--device", type=str, default="cuda") decode.add_argument("--residual-dir", type=Path, default=None, help="Directory containing residual_codebook.npy and residual_assignments.npy.") auto = sub.add_parser("auto", help="Train + encode + decode directly into a target iteration.") auto.add_argument("--model-path", type=Path, required=True) auto.add_argument("--source-iteration", type=int, required=True) auto.add_argument("--iteration", type=int, required=True, help="Target iteration to write decoded PLY.") auto.add_argument("--latent-dim", type=int, default=12) auto.add_argument("--hidden-dim", type=int, default=256) auto.add_argument("--num-layers", type=int, default=2) auto.add_argument("--dropout", type=float, default=0.0) auto.add_argument("--condition-fields", type=str, nargs="*", default=[]) auto.add_argument("--target-fields", type=str, nargs="*", default=["f_rest"]) auto.add_argument("--batch-size", type=int, default=8192) auto.add_argument("--steps", type=int, default=20000) auto.add_argument("--lr", type=float, default=1e-3) auto.add_argument("--device", type=str, default="cuda") auto.add_argument("--log-every", type=int, default=100) auto.add_argument("--seed", type=int, default=42) auto.add_argument("--dtype", type=str, default="float16", choices=["float32", "float16"]) auto.add_argument("--chunk", type=int, default=65536, help="Chunk size for encode/decode.") auto.add_argument("--quant-bits", type=int, default=None) auto.add_argument( "--checkpoint-dtype", type=str, choices=["float32", "float16"], default="float32", help="Precision used to store rest_autoencoder.pt.", ) auto.add_argument( "--latents-compression", type=str, choices=["npz", "npz_zstd"], default="npz", help="Container format for latents artifacts.", ) auto.add_argument( "--bitpack-latents", action="store_true", help="Store quantised latents in a bit-packed representation.", ) auto.add_argument("--overwrite", action="store_true", help="Overwrite existing target iteration PLY.") auto.add_argument("--residual-codebook-size", type=int, default=256) auto.add_argument("--residual-iters", type=int, default=20) auto.add_argument("--residual-condition-weight", type=float, default=0.2) auto.add_argument("--deploy-iteration", type=int, default=None, help="Iteration to publish final PLY (default: --iteration).") auto.add_argument("--entropy-share", action="store_true", help="Compress target-free PLY with entropy coding for sharing.") auto.add_argument("--entropy-pos-bits", type=int, default=18) auto.add_argument("--entropy-feature-bits", type=int, default=16) auto.add_argument("--entropy-rest-bits", type=int, default=16) auto.add_argument("--run-render", action="store_true", help="After deploying, run render.py and metrics.py.") auto.add_argument("--dataset-path", type=Path) auto.add_argument("--config-path", type=Path) auto.add_argument("--skip-metrics", action="store_true", help="Skip metrics.py invocation when --run-render is set.") auto.add_argument("--write-compressed", action="store_true", help="Emit compressed/ bundle (latents + restfree pack) for --compressed rendering.") auto.add_argument("--restfree-groups", type=str, default="scale,opacity,rotation,dynamic", help="Comma-separated attribute groups to restfree-quantise for compressed output.") auto.add_argument("--restfree-xyz-bits", type=int, default=14) auto.add_argument("--restfree-scale-bits", type=int, default=12) auto.add_argument("--restfree-opacity-bits", type=int, default=12) auto.add_argument("--restfree-rotation-bits", type=int, default=12) auto.add_argument("--restfree-dynamic-bits", type=int, default=16) auto.add_argument("--restfree-dynamic-mu", type=float, default=255.0, help="μ-law parameter for dynamic logits during compressed packing.") return parser def main() -> None: parser = build_parser() args = parser.parse_args() if args.command == "train": input_ply = _resolve_ply(args.input) train_autoencoder( input_ply, args.output_dir, latent_dim=args.latent_dim, hidden_dim=args.hidden_dim, num_layers=args.num_layers, dropout=args.dropout, target_fields=args.target_fields, condition_fields=args.condition_fields, batch_size=args.batch_size, steps=args.steps, lr=args.lr, device=args.device, log_every=args.log_every, seed=args.seed, checkpoint_dtype=args.checkpoint_dtype, ) elif args.command == "encode": encode_rest( args.input, args.checkpoint, args.output, dtype=args.dtype, chunk=args.chunk, device=args.device, quant_bits=args.quant_bits, latents_compression=args.latents_compression, bitpack_latents=args.bitpack_latents, ) elif args.command == "decode": decode_rest( args.latents, args.checkpoint, args.source, args.output, chunk=args.chunk, device=args.device, ) if args.residual_dir: codebook_path = args.residual_dir / "residual_codebook.npy" assignments_path = args.residual_dir / "residual_assignments.npy" if codebook_path.exists() and assignments_path.exists(): _, _, _, _, _, target_names, _ = _load_checkpoint(args.checkpoint, args.device) _apply_residual_to_ply(args.output, target_names, codebook_path, assignments_path) elif args.command == "auto": run_auto_pipeline(args) else: parser.error(f"Unsupported command {args.command}") if __name__ == "__main__": main()