| |
| """ |
| 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: |
| 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.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, |
| } |
| 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() |
|
|
