"""Shared fixed-shape helpers for the TripoSplat one-step flow/DiT graph. The neural network definition and checkpoint loader remain in a caller-provided checkout of the official TripoSplat repository. This module owns only the ONNX boundary, two export-only algebraic adaptations, deterministic fixtures, and strict contract/parity utilities. Public graph tensors are always float32. A browser artifact may use float16 weights and compute internally; explicit casts at the graph boundary preserve the official float32 timestep and the host-side float32 Euler sampler state. """ from __future__ import annotations import gc import hashlib import importlib.util import json import subprocess import sys import types from dataclasses import dataclass from pathlib import Path from types import ModuleType from typing import Any, Mapping, Sequence OFFICIAL_REPOSITORY_URL = "https://github.com/VAST-AI-Research/TripoSplat" BATCH_SIZE = 1 LATENT_TOKENS = 8192 LATENT_CHANNELS = 16 CAMERA_CHANNELS = 5 CONDITION_TOKENS = 4101 FEATURE1_CHANNELS = 1280 FEATURE2_CHANNELS = 128 INPUT_NAMES = ("latent", "camera", "t", "feature1", "feature2") OUTPUT_NAMES = ("pred_latent", "pred_camera") LATENT_SHAPE = (BATCH_SIZE, LATENT_TOKENS, LATENT_CHANNELS) CAMERA_SHAPE = (BATCH_SIZE, 1, CAMERA_CHANNELS) TIMESTEP_SHAPE = (BATCH_SIZE,) FEATURE1_SHAPE = (BATCH_SIZE, CONDITION_TOKENS, FEATURE1_CHANNELS) FEATURE2_SHAPE = (BATCH_SIZE, CONDITION_TOKENS, FEATURE2_CHANNELS) PRED_LATENT_SHAPE = LATENT_SHAPE PRED_CAMERA_SHAPE = CAMERA_SHAPE INPUT_SHAPES: Mapping[str, tuple[int, ...]] = { "latent": LATENT_SHAPE, "camera": CAMERA_SHAPE, "t": TIMESTEP_SHAPE, "feature1": FEATURE1_SHAPE, "feature2": FEATURE2_SHAPE, } OUTPUT_SHAPES: Mapping[str, tuple[int, ...]] = { "pred_latent": PRED_LATENT_SHAPE, "pred_camera": PRED_CAMERA_SHAPE, } INTERNAL_PRECISION_METADATA_KEY = "triposplat.internal_precision" PUBLIC_IO_METADATA_VALUE = "float32" @dataclass(frozen=True) class OfficialSource: """Imported official pipeline and model modules from one checkout.""" repository: Path pipeline_module: ModuleType model_module: ModuleType @dataclass(frozen=True) class AdapterMetadata: """Facts recorded after making the official model ONNX-exportable.""" real_rope_modules: int attention_query_chunk: int attention_head_chunk: int attention_head_padding: int qk_norm_padding_tokens: int qk_norm_modules: int stable_rms_norm_modules: int rms_norm_eps: float | None attention_output_chunk: int attention_output_reduction_chunk: int attention_output_modules: int static_position_shape: tuple[int, ...] static_position_dtype: str static_position_sha256: str collapsed_unconditional_context: bool = False def resolved_file(path: Path, description: str) -> Path: result = path.expanduser().resolve() if not result.is_file(): raise FileNotFoundError(f"{description} does not exist or is not a file: {result}") return result def _load_module(name: str, path: Path) -> ModuleType: spec = importlib.util.spec_from_file_location(name, path) if spec is None or spec.loader is None: raise ImportError(f"Unable to create an import specification for {path}") module = importlib.util.module_from_spec(spec) sys.modules[name] = module try: spec.loader.exec_module(module) except Exception: sys.modules.pop(name, None) raise return module def import_official_source(triposplat_repo: Path) -> OfficialSource: """Import official ``model.py`` and ``triposplat.py`` without modifying them. The official pipeline uses ``from model import ...``. A temporary ``model`` module alias is installed only while executing ``triposplat.py``; the previous process module and ``sys.path`` are restored even when an import fails. """ repository = triposplat_repo.expanduser().resolve() model_path = repository / "model.py" pipeline_path = repository / "triposplat.py" if not model_path.is_file() or not pipeline_path.is_file(): raise FileNotFoundError( f"Expected model.py and triposplat.py in {repository}. Pass a clone of " f"{OFFICIAL_REPOSITORY_URL}." ) identity = hashlib.sha256(str(repository).encode("utf-8")).hexdigest()[:16] model_name = f"_triposplat_official_flow_model_{identity}" pipeline_name = f"_triposplat_official_pipeline_{identity}" cached_model = sys.modules.get(model_name) cached_pipeline = sys.modules.get(pipeline_name) if cached_model is not None and cached_pipeline is not None: return OfficialSource(repository, cached_pipeline, cached_model) model_module = cached_model or _load_module(model_name, model_path) previous_model_alias = sys.modules.get("model") inserted_path = str(repository) sys.modules["model"] = model_module sys.path.insert(0, inserted_path) try: pipeline_module = _load_module(pipeline_name, pipeline_path) finally: try: sys.path.remove(inserted_path) except ValueError: pass if previous_model_alias is None: sys.modules.pop("model", None) else: sys.modules["model"] = previous_model_alias if not hasattr(pipeline_module, "load_flow_model"): raise AttributeError( f"Official source {pipeline_path} does not define load_flow_model" ) required_model_symbols = ( "LatentSeqMMFlowModel", "RePo3DRotaryEmbedding", "apply_rotary_emb", ) missing = [name for name in required_model_symbols if not hasattr(model_module, name)] if missing: raise AttributeError( f"Official source {model_path} is missing required symbols: {missing}" ) return OfficialSource(repository, pipeline_module, model_module) def source_revision(repository: Path) -> tuple[str, bool | None]: """Return the checkout commit and whether tracked source files are dirty.""" try: commit = subprocess.run( ["git", "-C", str(repository), "rev-parse", "HEAD"], check=True, capture_output=True, text=True, ).stdout.strip() status = subprocess.run( ["git", "-C", str(repository), "status", "--porcelain", "--", "model.py", "triposplat.py"], check=True, capture_output=True, text=True, ).stdout return commit or "unknown", bool(status.strip()) except (OSError, subprocess.CalledProcessError): return "unknown", None def torch_dtype_for_precision(torch: Any, precision: str) -> Any: if precision == "fp16": return torch.float16 if precision == "fp32": return torch.float32 raise ValueError(f"Unsupported internal precision {precision!r}; expected fp16 or fp32") def choose_torch_device(torch: Any, requested: str) -> Any: if requested == "auto": if torch.cuda.is_available(): requested = "cuda" elif hasattr(torch.backends, "mps") and torch.backends.mps.is_available(): requested = "mps" else: requested = "cpu" if requested == "cuda" and not torch.cuda.is_available(): raise RuntimeError("--device cuda was requested, but PyTorch reports CUDA unavailable") if requested == "mps" and not ( hasattr(torch.backends, "mps") and torch.backends.mps.is_available() ): raise RuntimeError("--device mps was requested, but PyTorch reports MPS unavailable") return torch.device(requested) def synchronize_torch(torch: Any, device: Any) -> None: if device.type == "cuda": torch.cuda.synchronize(device) elif device.type == "mps": torch.mps.synchronize() def release_torch_model(torch: Any, device: Any, *objects: Any) -> None: """Drop large references and ask the selected accelerator to release caches.""" # Deleting a function's local aliases cannot delete the caller's names, but it # does ensure this helper retains no references while collection runs. del objects gc.collect() if device.type == "cuda": torch.cuda.empty_cache() elif device.type == "mps": torch.mps.empty_cache() def validate_official_flow_configuration(model: Any) -> None: expected = { "q_token_length": LATENT_TOKENS, "in_channels": LATENT_CHANNELS, "cam_channels": CAMERA_CHANNELS, "cond_channels": FEATURE1_CHANNELS, "cond2_channels": FEATURE2_CHANNELS, "out_channels": LATENT_CHANNELS, } mismatches = { name: (getattr(model, name, None), value) for name, value in expected.items() if getattr(model, name, None) != value } if mismatches: formatted = ", ".join( f"{name}={observed!r} (expected {wanted!r})" for name, (observed, wanted) in mismatches.items() ) raise ValueError(f"Official flow-model configuration does not match the browser contract: {formatted}") def load_official_flow_model( torch: Any, triposplat_repo: Path, weights: Path, device: Any, internal_precision: str, low_memory_construction: bool = True, ) -> tuple[Any, OfficialSource]: """Load the checkpoint through official ``triposplat.load_flow_model``. For an fp16 artifact, temporarily using fp16 as PyTorch's default construction dtype avoids allocating a second full fp32 parameter set. The official loader, class, strict safetensors state load, and placement logic are still used. Explicitly-dtyped official constants (including the Sobol draw) are unaffected. """ checkpoint = resolved_file(weights, "TripoSplat flow-model safetensors checkpoint") source = import_official_source(triposplat_repo) dtype = torch_dtype_for_precision(torch, internal_precision) previous_dtype = torch.get_default_dtype() original_sobol_engine = torch.quasirandom.SobolEngine class OfficialFloat32SobolEngine: """Keep official Sobol construction/draw float32 during fp16 allocation. SobolEngine builds a large integer-scaled first point using the current default float dtype. Constructing that state in fp16 overflows before ``draw(dtype=float32)`` can help, so both phases must retain the untouched official process default. """ def __init__(self, *engine_args: Any, **engine_kwargs: Any) -> None: construction_dtype = torch.get_default_dtype() try: torch.set_default_dtype(torch.float32) self.engine = original_sobol_engine(*engine_args, **engine_kwargs) finally: torch.set_default_dtype(construction_dtype) def draw( self, n: int = 1, out: Any = None, dtype: Any = None, ) -> Any: return self.engine.draw( n=n, out=out, dtype=torch.float32 if dtype is None else dtype, ) try: if low_memory_construction: torch.set_default_dtype(dtype) torch.quasirandom.SobolEngine = OfficialFloat32SobolEngine model = source.pipeline_module.load_flow_model( str(checkpoint), device=device, dtype=dtype, ) finally: torch.quasirandom.SobolEngine = original_sobol_engine torch.set_default_dtype(previous_dtype) model = model.eval() validate_official_flow_configuration(model) mismatched = [ f"{name}={parameter.dtype}" for name, parameter in model.named_parameters() if parameter.is_floating_point() and parameter.dtype != dtype ] if mismatched: raise RuntimeError( f"Official flow-model parameters did not convert to {dtype}: " + ", ".join(mismatched[:8]) ) return model, source def _real_repo_frequencies(self: Any, hidden_states: Any) -> Any: """Real-valued equivalent of official ``RePo3DRotaryEmbedding.forward``. The final axis stores ``(cos(angle), sin(angle))`` instead of a complex64 number. Every operation used to produce ``angle`` intentionally keeps the ordering of the official source, including ``clamp_mul``'s detached branch. """ import torch h = self.norm(hidden_states) feat = self.act(self.gate_map(h)) * self.content_map(h) out = self.final_map(feat) batch, length, _ = out.shape delta_pos = out.reshape(batch, length, self.num_heads, 3) def clamp_mul_exact(x: Any, frequency: Any) -> Any: frequency_tanh = frequency.tanh() return x * frequency_tanh + x.detach() * (frequency - frequency_tanh) angle_0 = clamp_mul_exact(delta_pos[..., 0].unsqueeze(-1), self.freqs_0) * torch.pi angle_1 = clamp_mul_exact(delta_pos[..., 1].unsqueeze(-1), self.freqs_1) * torch.pi angle_2 = clamp_mul_exact(delta_pos[..., 2].unsqueeze(-1), self.freqs_2) * torch.pi angles = torch.cat((angle_0, angle_1, angle_2), dim=-1).float() return torch.stack((torch.cos(angles), torch.sin(angles)), dim=-1) def _apply_rotary_emb_real(hidden_states: Any, frequencies: Any) -> Any: """Pairwise real multiplication equal to official complex RoPE multiplication.""" import torch pairs = hidden_states.float().reshape(*hidden_states.shape[:-1], -1, 2) real = pairs[..., 0] imaginary = pairs[..., 1] cosine = frequencies[..., 0] sine = frequencies[..., 1] rotated = torch.stack( ( real * cosine - imaginary * sine, real * sine + imaginary * cosine, ), dim=-1, ).reshape(*hidden_states.shape) return rotated.type_as(hidden_states) def _scaled_dot_product_attention_query_chunked( qkv: Any = None, q: Any = None, k: Any = None, v: Any = None, kv: Any = None, *, query_chunk_size: int, collapsed_unconditional_context: bool, head_chunk_size: int, head_padding_size: int, ) -> Any: """Official SDPA split along independent head and query-token axes. Each query row has the same keys, values, scale, and row-wise softmax as the untouched operation. Splitting queries therefore changes no attention dependency while preventing ONNX Runtime WebGPU from materializing the full ``[B,H,Lq,Lk]`` score tensor. The unconditional-only collapsed specialization explicitly evaluates its one-token context and 8,194-token joint attention, adding a constant log-multiplicity bias for the retained context key. A trailing duplicate head shields an ORT WebGPU 1.27 correctness defect that corrupts the final head of the final SDPA group; the duplicate output is discarded. The fixed-shape exporter unrolls all loops. """ import math import torch import torch.nn.functional as functional if not isinstance(query_chunk_size, int) or query_chunk_size <= 0: raise ValueError("query_chunk_size must be a positive integer") if not isinstance(collapsed_unconditional_context, bool): raise TypeError("collapsed_unconditional_context must be a bool") if not isinstance(head_chunk_size, int) or head_chunk_size <= 0: raise ValueError("head_chunk_size must be a positive integer") if not isinstance(head_padding_size, int) or head_padding_size < 0: raise ValueError("head_padding_size must be a non-negative integer") if qkv is not None: q, k, v = qkv.unbind(dim=2) elif kv is not None: k, v = kv.unbind(dim=2) if q is None or k is None or v is None: raise ValueError("Chunked attention requires q, k, and v tensors") collapsed_length = int(q.shape[1]) use_collapsed_attention = ( collapsed_unconditional_context and collapsed_length in {1, LATENT_TOKENS + 2} and int(k.shape[1]) == collapsed_length ) output_dtype = q.dtype q, k, v = ( q.permute(0, 2, 1, 3), k.permute(0, 2, 1, 3), v.permute(0, 2, 1, 3), ) # PyTorch's fused fp16 SDPA uses a higher-precision score/softmax path. Make # that boundary explicit so the decomposed ONNX graph does not ask WebGPU to # carry 24 residual blocks through fp16 softmax accumulation. k_float = k.float() v_float = v.float() if use_collapsed_attention: log_multiplicity = math.log(CONDITION_TOKENS) if collapsed_length == 1: key_bias = torch.full( (1,), log_multiplicity, dtype=torch.float32, device=q.device, ) else: key_bias = torch.cat( ( torch.zeros( LATENT_TOKENS, dtype=torch.float32, device=q.device, ), torch.full( (1,), log_multiplicity, dtype=torch.float32, device=q.device, ), torch.zeros(1, dtype=torch.float32, device=q.device), ) ) key_bias = key_bias.reshape(1, 1, 1, collapsed_length) head_groups = [] for head_start in range(0, q.shape[1], head_chunk_size): head_end = min(head_start + head_chunk_size, q.shape[1]) head_q = q[:, head_start:head_end, :, :] head_k = k_float[:, head_start:head_end, :, :] head_v = v_float[:, head_start:head_end, :, :] real_head_count = head_q.shape[1] if head_padding_size: if head_padding_size > head_q.shape[1]: raise ValueError("head_padding_size cannot exceed the current head group") # Duplicate live input rather than appending a constant-zero head. This # prevents export/runtime optimization from deleting the shielding work. head_q = torch.cat((head_q, head_q[:, :head_padding_size, :, :]), dim=1) head_k = torch.cat((head_k, head_k[:, :head_padding_size, :, :]), dim=1) head_v = torch.cat((head_v, head_v[:, :head_padding_size, :, :]), dim=1) if not use_collapsed_attention: query_chunks = [ functional.scaled_dot_product_attention( head_q[:, :, start : start + query_chunk_size, :].float(), head_k, head_v, ).to(dtype=output_dtype) for start in range(0, q.shape[2], query_chunk_size) ] else: scale = float(head_q.shape[-1]) ** -0.5 query_chunks = [] for query_start in range(0, q.shape[2], query_chunk_size): query = head_q[ :, :, query_start : query_start + query_chunk_size, : ].float() scores = torch.matmul(query, head_k.transpose(-2, -1)) * scale probabilities = torch.softmax(scores + key_bias, dim=-1) output = torch.matmul(probabilities, head_v) query_chunks.append(output.to(dtype=output_dtype)) head_groups.append(torch.cat(query_chunks, dim=2)[:, :real_head_count, :, :]) return torch.cat(head_groups, dim=1).permute(0, 2, 1, 3) def _tensor_sha256(tensor: Any) -> str: array = tensor.detach().to(device="cpu").contiguous().numpy() return hashlib.sha256(memoryview(array).cast("B")).hexdigest() def validate_real_rope_primitives( torch: Any, model: Any, model_module: ModuleType, seed: int = 20260714, ) -> dict[str, float | bool]: """Cheaply gate the real RoPE algebra against untouched official complex ops. This does not replace the validator's full one-call gate. It fails early before a multi-gigabyte export if the official RoPE layout or arithmetic has changed. """ layer = model.noise_repo_layers[0] length = 7 cpu_generator = torch.Generator(device="cpu").manual_seed(seed) hidden = torch.randn( (1, length, int(model.model_channels)), generator=cpu_generator, dtype=torch.float32, ).to(device=model.device, dtype=model.dtype) q = torch.randn( (1, length, int(model.num_heads), int(layer.head_dim)), generator=cpu_generator, dtype=torch.float32, ).to(device=model.device, dtype=model.dtype) with torch.inference_mode(): official_frequencies = layer(hidden) official_rotated = model_module.apply_rotary_emb(q, official_frequencies) real_frequencies = _real_repo_frequencies(layer, hidden) real_rotated = _apply_rotary_emb_real(q, real_frequencies) frequency_reference = torch.view_as_real(official_frequencies).float() frequency_error = float( (frequency_reference - real_frequencies.float()).abs().max().detach().cpu() ) rotation_error = float( (official_rotated.float() - real_rotated.float()).abs().max().detach().cpu() ) tolerance = 2e-3 if model.dtype == torch.float16 else 3e-6 passed = frequency_error <= tolerance and rotation_error <= tolerance return { "passed": passed, "max_frequency_pair_error": frequency_error, "max_rotated_tensor_error": rotation_error, "atol": tolerance, } def _collapsed_unconditional_forward(self: Any, x_t: Any, t: Any, cond: Any) -> Any: """Official forward specialized to all-zero conditioning without source edits.""" import torch import torch.nn.functional as functional if ( len(self.noise_refiner) != 2 or len(self.context_refiner) != 2 or len(self.blocks) != 24 ): raise ValueError( "Collapsed unconditional context requires 2 noise refiners, 2 context " "refiners, and 24 joint blocks" ) if self.cond_embedder2 is None or self.cam_channels is None: raise ValueError( "Collapsed unconditional context requires both condition embedders and camera" ) d = self.dtype z = x_t["latent"].to(d) feat1 = cond["feature1"].to(d) feat2 = cond["feature2"].to(d) if tuple(z.shape) != LATENT_SHAPE: raise ValueError(f"Collapsed latent shape is {tuple(z.shape)}, expected {LATENT_SHAPE}") if tuple(feat1.shape) != FEATURE1_SHAPE: raise ValueError( f"Collapsed feature1 shape is {tuple(feat1.shape)}, expected {FEATURE1_SHAPE}" ) if tuple(feat2.shape) != FEATURE2_SHAPE: raise ValueError( f"Collapsed feature2 shape is {tuple(feat2.shape)}, expected {FEATURE2_SHAPE}" ) if tuple(t.shape) != TIMESTEP_SHAPE: raise ValueError( f"Collapsed timestep shape is {tuple(t.shape)}, expected {TIMESTEP_SHAPE}" ) self.pos_pe = self.pos_pe.to(z.device) h_x = self.input_layer(z) h_cond = self.cond_embedder(feat1) h_cond = h_cond + self.cond_embedder2(feat2) expected_condition_shape = (BATCH_SIZE, CONDITION_TOKENS, int(self.model_channels)) if tuple(h_cond.shape) != expected_condition_shape: raise ValueError( f"Embedded condition shape is {tuple(h_cond.shape)}, expected " f"{expected_condition_shape}" ) # Both public feature tensors have been consumed by their official embedders. # For all-zero features every row is the same learned-bias representative. h_cond = h_cond[:, :1, :] expected_collapsed_shape = (BATCH_SIZE, 1, int(self.model_channels)) if tuple(h_cond.shape) != expected_collapsed_shape: raise ValueError( f"Collapsed condition shape is {tuple(h_cond.shape)}, expected " f"{expected_collapsed_shape}" ) t_emb = self.t_embedder(t) t_mod = self.adaLN_modulation(t_emb) if self.share_mod else t_emb h_x = h_x + self.pos_embedder(self.pos_pe).to(d) for i, block in enumerate(self.noise_refiner): h_x = block(h_x, mod=t_mod, rotary_emb=self.noise_repo_layers[i](h_x)) for i, block in enumerate(self.context_refiner): h_cond = block( h_cond, mod=None, rotary_emb=self.context_repo_layers[i](h_cond), ) cam = x_t.get("camera").to(d) if tuple(cam.shape) != CAMERA_SHAPE: raise ValueError( f"Collapsed camera shape is {tuple(cam.shape)}, expected {CAMERA_SHAPE}" ) h_cam = self.cam_refiner(cam) h = torch.cat([h_x, h_cond], dim=1) h = torch.cat([h, h_cam], dim=1) expected_joint_shape = (BATCH_SIZE, LATENT_TOKENS + 2, int(self.model_channels)) if tuple(h.shape) != expected_joint_shape: raise ValueError( f"Collapsed joint shape is {tuple(h.shape)}, expected {expected_joint_shape}" ) for i, block in enumerate(self.blocks): h = block(h, mod=t_mod, rotary_emb=self.repo_layers[i](h)) h_x = functional.layer_norm( h[:, : z.shape[1]].float(), h.shape[-1:] ).type(d) h_cam = functional.layer_norm( h[:, -cam.shape[1] :].float(), h.shape[-1:] ).type(d) if self.use_shift_table: shift, scale = (self.shift_table + t_emb.unsqueeze(1)).chunk(2, dim=1) h_x = h_x * (1 + scale) + shift h_cam = h_cam * (1 + scale) + shift out = {"latent": self.out_layer(h_x)} out["camera"] = self.cam_out_layer(h_cam) if tuple(out["latent"].shape) != PRED_LATENT_SHAPE: raise ValueError( f"Collapsed latent output shape is {tuple(out['latent'].shape)}, expected " f"{PRED_LATENT_SHAPE}" ) if tuple(out["camera"].shape) != PRED_CAMERA_SHAPE: raise ValueError( f"Collapsed camera output shape is {tuple(out['camera'].shape)}, expected " f"{PRED_CAMERA_SHAPE}" ) return out def adapt_official_flow_for_onnx( torch: Any, model: Any, model_module: ModuleType, attention_query_chunk: int = 256, collapsed_unconditional_context: bool = False, attention_head_chunk: int = 16, attention_head_padding: int = 0, qk_norm_padding_tokens: int = 1, rms_norm_eps: float | None = None, attention_output_chunk: int = 256, attention_output_reduction_chunk: int = 256, ) -> AdapterMetadata: """Apply three one-way, export-only adaptations to an official model instance. 1. Complex RoPE is represented as real ``cos/sin`` pairs and applied with the algebraically identical two-real multiply. 2. The fixed seeded-Sobol absolute position embedding is evaluated once with official code and registered as a buffer. This prevents exporter/runtime constant-folding of large trigonometric arguments from changing its values. 3. SDPA is split on the independent query and head axes so WebGPU never allocates the full multi-gigabyte attention score tensor and never invokes the known-corrupt 16-head ORT WebGPU kernel. When the unconditional-only specialization is enabled, the official forward retains one context token and collapsed self-attention restores the 4,101-key multiplicity exactly. With ``collapsed_unconditional_context=False``, the official forward remains untouched and SDPA retains the canonical functional path. Validators must run the untouched official forward first, then this adapted forward on identical tensors. This function deliberately does not patch files in the source checkout. """ existing = getattr(model, "_triposplat_onnx_adapter_metadata", None) if existing is not None: if existing.attention_query_chunk != attention_query_chunk: raise ValueError( "Flow model was already adapted with attention query chunk " f"{existing.attention_query_chunk}, not {attention_query_chunk}" ) if ( existing.collapsed_unconditional_context != collapsed_unconditional_context ): raise ValueError( "Flow model was already adapted with collapsed unconditional context " f"{existing.collapsed_unconditional_context}, not " f"{collapsed_unconditional_context}" ) if existing.attention_head_chunk != attention_head_chunk: raise ValueError( "Flow model was already adapted with attention head chunk " f"{existing.attention_head_chunk}, not {attention_head_chunk}" ) if existing.attention_head_padding != attention_head_padding: raise ValueError( "Flow model was already adapted with attention head padding " f"{existing.attention_head_padding}, not {attention_head_padding}" ) if existing.qk_norm_padding_tokens != qk_norm_padding_tokens: raise ValueError( "Flow model was already adapted with Q/K norm token padding " f"{existing.qk_norm_padding_tokens}, not {qk_norm_padding_tokens}" ) if existing.rms_norm_eps != rms_norm_eps: raise ValueError( "Flow model was already adapted with RMS norm epsilon " f"{existing.rms_norm_eps}, not {rms_norm_eps}" ) if existing.attention_output_chunk != attention_output_chunk: raise ValueError( "Flow model was already adapted with attention output chunk " f"{existing.attention_output_chunk}, not {attention_output_chunk}" ) if existing.attention_output_reduction_chunk != attention_output_reduction_chunk: raise ValueError( "Flow model was already adapted with attention output reduction chunk " f"{existing.attention_output_reduction_chunk}, not " f"{attention_output_reduction_chunk}" ) return existing if not isinstance(attention_query_chunk, int) or attention_query_chunk <= 0: raise ValueError("attention_query_chunk must be a positive integer") if not isinstance(collapsed_unconditional_context, bool): raise TypeError("collapsed_unconditional_context must be a bool") if not isinstance(attention_head_chunk, int) or attention_head_chunk <= 0: raise ValueError("attention_head_chunk must be a positive integer") if not isinstance(attention_head_padding, int) or attention_head_padding < 0: raise ValueError("attention_head_padding must be a non-negative integer") if attention_head_padding > attention_head_chunk: raise ValueError("attention_head_padding cannot exceed attention_head_chunk") if not isinstance(qk_norm_padding_tokens, int) or qk_norm_padding_tokens <= 0: raise ValueError("qk_norm_padding_tokens must be a positive integer") if rms_norm_eps is not None and ( not isinstance(rms_norm_eps, (int, float)) or rms_norm_eps <= 0 ): raise ValueError("rms_norm_eps must be None or a positive number") if not isinstance(attention_output_chunk, int) or attention_output_chunk <= 0: raise ValueError("attention_output_chunk must be a positive integer") if ( not isinstance(attention_output_reduction_chunk, int) or attention_output_reduction_chunk <= 0 ): raise ValueError("attention_output_reduction_chunk must be a positive integer") if not hasattr(model, "pos_pe") or not hasattr(model, "pos_embedder"): raise AttributeError("Official flow model lacks pos_pe/pos_embedder") with torch.inference_mode(): position_input = model.pos_pe.to(device=model.device, dtype=torch.float32) fixed_position = model.pos_embedder(position_input).to(dtype=model.dtype) expected_position_shape = (BATCH_SIZE, LATENT_TOKENS, int(model.model_channels)) if tuple(fixed_position.shape) != expected_position_shape: raise ValueError( f"Official static position embedding has shape {tuple(fixed_position.shape)}, " f"expected {expected_position_shape}" ) class StaticOfficialPositionEmbedding(torch.nn.Module): def __init__(self, value: Any) -> None: super().__init__() self.register_buffer("value", value.detach().clone(), persistent=True) def forward(self, unused_position_input: Any) -> Any: return self.value model.pos_embedder = StaticOfficialPositionEmbedding(fixed_position).to(model.device) repo_type = model_module.RePo3DRotaryEmbedding replaced = 0 for submodule in model.modules(): if isinstance(submodule, repo_type): submodule.forward = types.MethodType(_real_repo_frequencies, submodule) replaced += 1 expected_replaced = sum( len(getattr(model, name)) for name in ("noise_repo_layers", "context_repo_layers", "repo_layers") ) if replaced != expected_replaced or replaced <= 0: raise RuntimeError( f"Adapted {replaced} RePo modules, expected {expected_replaced}; " "the official architecture may have changed" ) model_module.apply_rotary_emb = _apply_rotary_emb_real model_module.scaled_dot_product_attention = lambda qkv=None, q=None, k=None, v=None, kv=None: ( _scaled_dot_product_attention_query_chunked( qkv=qkv, q=q, k=k, v=v, kv=kv, query_chunk_size=attention_query_chunk, collapsed_unconditional_context=collapsed_unconditional_context, head_chunk_size=attention_head_chunk, head_padding_size=attention_head_padding, ) ) if collapsed_unconditional_context: model.forward = types.MethodType(_collapsed_unconditional_forward, model) # Core AI's TripoSplat conversion identified a true-scale failure in the # original F.normalize formulation: some converters/runtimes omit its eps # clamp for near-zero unconditional Q/K vectors. The mean-square form is # algebraically identical away from epsilon and exposes the stabilization as # primitive ONNX ops that ORT WebGPU preserves. stable_rms_norm_modules = 0 if rms_norm_eps is not None: rms_norm_type = model_module.MultiHeadRMSNorm def stable_rms_norm_forward(self: Any, value: Any) -> Any: original_dtype = value.dtype fp32 = value.float() inverse_rms = torch.rsqrt( (fp32 * fp32).mean(dim=-1, keepdim=True) + rms_norm_eps ) return (fp32 * inverse_rms * self.gamma.float()).to(dtype=original_dtype) for submodule in model.modules(): if isinstance(submodule, rms_norm_type): submodule.forward = types.MethodType(stable_rms_norm_forward, submodule) stable_rms_norm_modules += 1 if stable_rms_norm_modules <= 0: raise RuntimeError("No MultiHeadRMSNorm modules were adapted") class ChunkedAttentionOutputProjection(torch.nn.Module): """Express the token-wise linear map as bounded 1x1 convolutions.""" def __init__(self, projection: Any) -> None: super().__init__() self.projection = projection def forward(self, value: Any) -> Any: import torch.nn.functional as functional outputs = [] output_channels = int(self.projection.weight.shape[0]) original_shape = value.shape token_count = int(value.shape[-2]) first_token = value[..., :1, :].float().unsqueeze(-2) tail_indices = torch.arange( 1, token_count, dtype=torch.int64, device=value.device, ) channels_first = value.transpose(-1, -2).unsqueeze(-1) for start in range(0, output_channels, attention_output_chunk): end = min(start + attention_output_chunk, output_channels) bias = ( None if self.projection.bias is None else self.projection.bias[start:end] ) projected = functional.conv2d( channels_first, self.projection.weight[start:end].unsqueeze(-1).unsqueeze(-1), bias, ) bulk = projected.squeeze(-1).transpose(-1, -2) bulk_tail = torch.index_select(bulk, -2, tail_indices) special_products = ( first_token * self.projection.weight[start:end] .float() .unsqueeze(0) .unsqueeze(0) ).transpose(-1, -2) reduction_rows = int(special_products.shape[-2]) if reduction_rows & (reduction_rows - 1): raise ValueError( "Attention projection input width must be a power of two" ) while reduction_rows > 256: half = reduction_rows // 2 # Keep output channels as the 256-wide innermost dimension. special_products = ( special_products[..., :half, :] + special_products[..., half:reduction_rows, :] ) reduction_rows = half while reduction_rows > 1: half = reduction_rows // 2 padding = torch.zeros( 256 - half, dtype=torch.int64, device=value.device, ) left_indices = torch.cat( ( torch.arange(0, half, dtype=torch.int64, device=value.device), padding, ) ) right_indices = torch.cat( ( torch.arange( half, reduction_rows, dtype=torch.int64, device=value.device, ), padding, ) ) active_mask = torch.cat( ( torch.ones(half, dtype=torch.float32, device=value.device), torch.zeros(256 - half, dtype=torch.float32, device=value.device), ) ).reshape(1, 1, 256, 1) # ORT WebGPU corrupts an accumulator once this dimension reaches # 128 rows (the tail starts at flat element 16,384). Both gathers, # masks, and the Add therefore keep exactly 256 rows, with inactive # lanes explicitly zeroed. special_products = ( torch.index_select(special_products, -2, left_indices) * active_mask + torch.index_select(special_products, -2, right_indices) * active_mask ) reduction_rows = half special = special_products[..., 0, :] if bias is not None: special = special + bias.float().reshape(1, 1, -1) outputs.append( torch.cat( (special.to(dtype=value.dtype), bulk_tail), dim=-2, ) ) return torch.cat(outputs, dim=-1).reshape(*original_shape[:-1], output_channels) class PaddedQkRmsNorm(torch.nn.Module): """Keep real vectors away from ORT WebGPU's final reduction lane.""" def __init__(self, norm: Any) -> None: super().__init__() self.norm = norm def forward(self, value: Any) -> Any: token_count = value.shape[1] if qk_norm_padding_tokens > token_count: raise ValueError("Q/K norm padding cannot exceed the token count") padded = torch.cat( (value, value[:, :qk_norm_padding_tokens, ...]), dim=1, ) return self.norm(padded)[:, :token_count, ...] attention_type = model_module.RopeMultiHeadAttention output_modules = 0 qk_norm_modules = 0 for submodule in model.modules(): if isinstance(submodule, attention_type): if getattr(submodule, "qk_rms_norm", False): submodule.q_norm = PaddedQkRmsNorm(submodule.q_norm) submodule.k_norm = PaddedQkRmsNorm(submodule.k_norm) qk_norm_modules += 2 submodule.out = ChunkedAttentionOutputProjection(submodule.out) output_modules += 1 if output_modules != 28: raise RuntimeError( f"Chunked {output_modules} attention output projections; expected 28" ) if qk_norm_modules != 56: raise RuntimeError(f"Padded {qk_norm_modules} Q/K norm modules; expected 56") metadata = AdapterMetadata( real_rope_modules=replaced, attention_query_chunk=attention_query_chunk, collapsed_unconditional_context=collapsed_unconditional_context, attention_head_chunk=attention_head_chunk, attention_head_padding=attention_head_padding, qk_norm_padding_tokens=qk_norm_padding_tokens, qk_norm_modules=qk_norm_modules, stable_rms_norm_modules=stable_rms_norm_modules, rms_norm_eps=None if rms_norm_eps is None else float(rms_norm_eps), attention_output_chunk=attention_output_chunk, attention_output_reduction_chunk=attention_output_reduction_chunk, attention_output_modules=output_modules, static_position_shape=tuple(int(value) for value in fixed_position.shape), static_position_dtype=str(fixed_position.dtype).removeprefix("torch."), static_position_sha256=_tensor_sha256(fixed_position), ) model._triposplat_onnx_adapter_metadata = metadata return metadata def make_browser_flow_step(torch: Any, model: Any, internal_precision: str) -> Any: """Wrap official dict I/O with fixed float32 browser tensors and explicit casts.""" internal_dtype = torch_dtype_for_precision(torch, internal_precision) class BrowserFlowStep(torch.nn.Module): def __init__(self, official_model: Any) -> None: super().__init__() self.flow_model = official_model def forward( self, latent: Any, camera: Any, t: Any, feature1: Any, feature2: Any, ) -> tuple[Any, Any]: result = self.flow_model( { "latent": latent.to(dtype=internal_dtype), "camera": camera.to(dtype=internal_dtype), }, # Official FlowEulerCfgSampler constructs 1000*t in float32. t.to(dtype=torch.float32), { "feature1": feature1.to(dtype=internal_dtype), "feature2": feature2.to(dtype=internal_dtype), }, ) return ( result["latent"].to(dtype=torch.float32), result["camera"].to(dtype=torch.float32), ) return BrowserFlowStep(model).eval().to(model.device) def deterministic_inputs(seed: int = 20260714, timestep: float = 1000.0) -> dict[str, Any]: """Create stable float32 tensors with encoder-like unit-normal conditioning.""" import numpy as np if not np.isfinite(timestep) or timestep < 0.0 or timestep > 1000.0: raise ValueError(f"timestep must be finite and in [0,1000], got {timestep}") rng = np.random.default_rng(seed) def normal(shape: tuple[int, ...]) -> Any: return np.ascontiguousarray(rng.standard_normal(shape).astype(np.float32)) values = { "latent": normal(LATENT_SHAPE), "camera": normal(CAMERA_SHAPE), "t": np.asarray([timestep], dtype=np.float32), "feature1": normal(FEATURE1_SHAPE), "feature2": normal(FEATURE2_SHAPE), } validate_input_arrays(values) return values def validate_input_arrays(values: Mapping[str, Any]) -> None: import numpy as np missing = [name for name in INPUT_NAMES if name not in values] extra = [name for name in values if name not in INPUT_NAMES] if missing or extra: raise ValueError(f"DiT inputs mismatch: missing={missing}, extra={extra}") for name in INPUT_NAMES: array = np.asarray(values[name]) if tuple(array.shape) != INPUT_SHAPES[name]: raise ValueError( f"{name} has shape {tuple(array.shape)}; expected {INPUT_SHAPES[name]}" ) if array.dtype != np.float32: raise TypeError(f"{name} has dtype {array.dtype}; public contract requires float32") if not np.isfinite(array).all(): raise ValueError(f"{name} contains NaN or infinity") timestep_min = float(np.asarray(values["t"]).min()) timestep_max = float(np.asarray(values["t"]).max()) if timestep_min < 0.0 or timestep_max > 1000.0: raise ValueError( f"t must contain already-scaled official timesteps in [0,1000], " f"observed [{timestep_min:.8g},{timestep_max:.8g}]" ) def load_input_fixture(path: Path) -> tuple[dict[str, Any], dict[str, Any]]: """Load strict public inputs plus optional official prediction references.""" import numpy as np fixture_path = resolved_file(path, "DiT NPZ fixture") with np.load(fixture_path, allow_pickle=False) as fixture: missing = [name for name in INPUT_NAMES if name not in fixture.files] if missing: raise KeyError( f"{fixture_path} is missing {missing}; found {fixture.files}" ) inputs = { name: np.ascontiguousarray(np.asarray(fixture[name], dtype=np.float32)) for name in INPUT_NAMES } expected: dict[str, Any] = {} for name in OUTPUT_NAMES: if name in fixture.files: expected[name] = np.ascontiguousarray( np.asarray(fixture[name], dtype=np.float32) ) if expected and set(expected) != set(OUTPUT_NAMES): raise KeyError( f"{fixture_path} must contain both {list(OUTPUT_NAMES)} or neither; " f"found reference keys {sorted(expected)}" ) validate_input_arrays(inputs) for name, array in expected.items(): if tuple(array.shape) != OUTPUT_SHAPES[name]: raise ValueError( f"Fixture {name} has shape {tuple(array.shape)}; expected {OUTPUT_SHAPES[name]}" ) return inputs, expected def fixed_value_shape(value_info: Any) -> tuple[int | None, ...]: return tuple( int(dimension.dim_value) if dimension.HasField("dim_value") else None for dimension in value_info.type.tensor_type.shape.dim ) def metadata_dict(model_proto: Any) -> dict[str, str]: return {entry.key: entry.value for entry in model_proto.metadata_props} def verify_onnx_contract( onnx: Any, model_proto: Any, require_metadata: bool = True, ) -> str: """Reject names, symbolic axes, dtypes, complex tensors, or metadata drift.""" inputs = {value.name: value for value in model_proto.graph.input} outputs = {value.name: value for value in model_proto.graph.output} if set(inputs) != set(INPUT_NAMES): raise ValueError(f"ONNX inputs are {sorted(inputs)}; expected {list(INPUT_NAMES)}") if set(outputs) != set(OUTPUT_NAMES): raise ValueError(f"ONNX outputs are {sorted(outputs)}; expected {list(OUTPUT_NAMES)}") for name, expected_shape in (*INPUT_SHAPES.items(), *OUTPUT_SHAPES.items()): value = inputs[name] if name in inputs else outputs[name] observed_shape = fixed_value_shape(value) if observed_shape != expected_shape: raise ValueError( f"ONNX {name} shape is {observed_shape}; expected fixed {expected_shape}" ) observed_type = value.type.tensor_type.elem_type if observed_type != onnx.TensorProto.FLOAT: raise ValueError( f"ONNX {name} element type is {observed_type}; public I/O must be float32" ) complex_types = {onnx.TensorProto.COMPLEX64, onnx.TensorProto.COMPLEX128} complex_initializers = [ tensor.name for tensor in model_proto.graph.initializer if tensor.data_type in complex_types ] complex_value_info = [ value.name for value in ( list(model_proto.graph.input) + list(model_proto.graph.output) + list(model_proto.graph.value_info) ) if value.type.HasField("tensor_type") and value.type.tensor_type.elem_type in complex_types ] if complex_initializers or complex_value_info: raise ValueError( "Exported graph still contains complex tensors: " f"initializers={complex_initializers[:8]}, values={complex_value_info[:8]}" ) metadata = metadata_dict(model_proto) precision = metadata.get(INTERNAL_PRECISION_METADATA_KEY, "") if require_metadata and precision not in {"fp16", "fp32"}: raise ValueError( f"ONNX metadata {INTERNAL_PRECISION_METADATA_KEY!r} is {precision!r}; " "expected fp16 or fp32" ) public_io = metadata.get("triposplat.public_io", "") if require_metadata and public_io != PUBLIC_IO_METADATA_VALUE: raise ValueError( f"ONNX metadata triposplat.public_io is {public_io!r}; expected float32" ) return precision or "unknown" def verify_external_data_files(graph_path: Path, model_proto: Any) -> list[Path]: """Validate every external-data location, offset, and byte range.""" graph = graph_path.expanduser().resolve() locations: dict[str, Path] = {} for tensor in model_proto.graph.initializer: if not tensor.external_data: continue info = {entry.key: entry.value for entry in tensor.external_data} location = info.get("location") if not location: raise ValueError(f"External initializer {tensor.name!r} has no location") candidate = (graph.parent / location).resolve() try: candidate.relative_to(graph.parent) except ValueError as exc: raise ValueError( f"External initializer {tensor.name!r} escapes graph directory: {location}" ) from exc if not candidate.is_file(): raise FileNotFoundError( f"External initializer {tensor.name!r} references missing {candidate}" ) offset = int(info.get("offset", "0")) length_text = info.get("length") if offset < 0: raise ValueError(f"External initializer {tensor.name!r} has negative offset") if length_text is not None: length = int(length_text) if length < 0 or offset + length > candidate.stat().st_size: raise ValueError( f"External initializer {tensor.name!r} range " f"[{offset},{offset + length}) exceeds {candidate}" ) locations[location] = candidate return [locations[key] for key in sorted(locations)] def array_summary(array: Any) -> dict[str, float | int]: import numpy as np values = np.asarray(array, dtype=np.float64) return { "count": int(values.size), "min": float(values.min()), "max": float(values.max()), "mean": float(values.mean()), "std": float(values.std()), "l2_norm": float(np.linalg.norm(values.ravel())), } def comparison_metrics( reference: Any, candidate: Any, atol: float, rtol: float, ) -> dict[str, Any]: import numpy as np ref = np.asarray(reference, dtype=np.float64) got = np.asarray(candidate, dtype=np.float64) if ref.shape != got.shape: raise ValueError(f"Comparison shape mismatch: {ref.shape} vs {got.shape}") reference_finite = bool(np.isfinite(ref).all()) candidate_finite = bool(np.isfinite(got).all()) if not reference_finite or not candidate_finite: return { "passed": False, "count": int(ref.size), "reference_finite": reference_finite, "candidate_finite": candidate_finite, } delta = got - ref absolute = np.abs(delta) relative = absolute / np.maximum(np.abs(ref), 1e-6) allowed = atol + rtol * np.abs(ref) within = absolute <= allowed ref_flat = ref.ravel() got_flat = got.ravel() norm_product = float(np.linalg.norm(ref_flat) * np.linalg.norm(got_flat)) cosine = float(np.dot(ref_flat, got_flat) / norm_product) if norm_product else 1.0 worst_flat = int(np.argmax(absolute)) worst_index = tuple(int(value) for value in np.unravel_index(worst_flat, ref.shape)) return { "passed": bool(within.all()), "count": int(ref.size), "reference_finite": True, "candidate_finite": True, "max_absolute_error": float(absolute.max()), "mean_absolute_error": float(absolute.mean()), "rmse": float(np.sqrt(np.mean(delta * delta))), "max_relative_error_at_1e-6_floor": float(relative.max()), "mean_relative_error_at_1e-6_floor": float(relative.mean()), "cosine_similarity": cosine, "fraction_within_tolerance": float(within.mean()), "worst_index": list(worst_index), "worst_reference": float(ref[worst_index]), "worst_candidate": float(got[worst_index]), "worst_allowed_error": float(allowed[worst_index]), } def shifted_flow_schedule(steps: int, shift: float) -> list[tuple[float, float]]: """Return official ``FlowEulerCfgSampler`` (t, t_previous) float64 pairs.""" import numpy as np if not isinstance(steps, int) or steps <= 0: raise ValueError(f"steps must be a positive integer, got {steps}") if not np.isfinite(shift) or shift <= 0: raise ValueError(f"shift must be a positive finite number, got {shift}") linear = np.linspace(1, 0, steps + 1) timesteps = shift * linear / (1 + (shift - 1) * linear) return [ (float(timesteps[index]), float(timesteps[index + 1])) for index in range(steps) ] def write_json(path: Path, value: Any) -> Path: destination = path.expanduser().resolve() destination.parent.mkdir(parents=True, exist_ok=True) destination.write_text(json.dumps(value, indent=2, sort_keys=True) + "\n", encoding="utf-8") return destination def sha256_file(path: Path) -> str: digest = hashlib.sha256() with path.open("rb") as stream: while block := stream.read(8 * 1024 * 1024): digest.update(block) return digest.hexdigest() def coerce_float32_inputs(values: Mapping[str, Any]) -> dict[str, Any]: import numpy as np result = { name: np.ascontiguousarray(np.asarray(values[name], dtype=np.float32)) for name in INPUT_NAMES } validate_input_arrays(result) return result def torch_inputs_from_numpy( torch: Any, values: Mapping[str, Any], device: Any, ) -> tuple[Any, Any, Any, Any, Any]: return tuple( torch.from_numpy(values[name]).to(device=device, dtype=torch.float32) for name in INPUT_NAMES ) # type: ignore[return-value] def output_mapping(outputs: Sequence[Any]) -> dict[str, Any]: if len(outputs) != len(OUTPUT_NAMES): raise ValueError(f"Expected {len(OUTPUT_NAMES)} outputs, got {len(outputs)}") return dict(zip(OUTPUT_NAMES, outputs))