diff --git a/app.py b/app.py index 525e26b2784f5022b87fbb9071cf02a3f3d312a6..79dca7a2aaef4a8d5ef3a24fd8605cda0f22b9f4 100644 --- a/app.py +++ b/app.py @@ -1,70 +1,70 @@ -"""Supra-Tex entrypoint with the supplied demo presentation layer.""" -import spaces -import gradio as gr -from hf_space import native_build -from hf_space import backend -from hf_space import result_cache -from hf_space import generation_progress -from hf_space import ovoxel_results -import time - -result_cache.prepare_examples() -native_build.start() -# Preload once per server process, independently of visitor sessions. -backend.start_preparation() - - -def deployment_status(): - return {'native_build': native_build.status(), 'models': backend.model_status()} - - -@spaces.GPU(duration=900, size='large') -def generate_on_gpu(input_image, seed, progress, generation_token=None): - mesh, _asset, report = backend.generate(input_image, int(seed), progress=progress, - generation_token=generation_token) - return mesh, report - - -def process_image_generator(input_image, seed=42, generation_token=None, progress=gr.Progress()): - if input_image is None: - raise gr.Error('Please upload an image.') - print('[generation] Explicit Generate request; running released model without cache reuse', flush=True) - token = generation_token if generation_progress.valid(generation_token) else generation_progress.start() - try: - backend.ensure_prepared(token) - generation_progress.update(token, status='queued', label='Waiting for GPU') - mesh, report = generate_on_gpu(input_image, int(seed), progress, token) - generation_progress.update(token, status='saving', label='Saving the 3D asset') - try: - identifier = ovoxel_results.persist(mesh, report) - except Exception as error: - print(f'[result-cache] Could not persist result: {type(error).__name__}: {error}', flush=True) - raise - # The pipeline already rendered the saved asset. Show those views now; - # exporting a GLB belongs to the separate Extract GLB action. - rendered = ovoxel_results.preview_html(identifier) - generation_progress.update(token, status='complete', label='Complete', finished_unix=time.time()) - except Exception: - generation_progress.update(token, status='failed', label='Generation failed', finished_unix=time.time()) - raise - return rendered, identifier - - -@spaces.GPU(duration=600, size='large') -def extract_glb(identifier, extraction_token=None): - return ovoxel_results.extract(identifier, extraction_token) - - -@spaces.GPU(duration=600, size='large') -def recover_preview(identifier): - from hf_space.preview_recovery import recover - return recover(identifier) - - -from hf_space.ui import build_ui, CSS, THEME - -demo = build_ui(process_image_generator, deployment_status, cached_result=result_cache.lookup, extract_glb=extract_glb, recover_preview=recover_preview, - default_image=None, default_model=None, estimated_seconds=390, - example_names=result_cache.EXAMPLE_NAMES, examples_status=result_cache.examples_status) -demo.queue().launch(allowed_paths=['/tmp/supratex-jobs', str(result_cache.ROOT/'assets/previews'), str(result_cache.CACHE),str(ovoxel_results.ROOT), str(result_cache.ROOT/'hf_space/gallery-viewer.bundle.js'), str(result_cache.ROOT/'hf_space/gallery-lights')], - css=CSS, theme=THEME) +"""Supra-Tex entrypoint with the supplied demo presentation layer.""" +import spaces +import gradio as gr +from hf_space import native_build +from hf_space import release_backend as backend +from hf_space import result_cache +from hf_space import generation_progress +from hf_space import ovoxel_results +import time + +result_cache.prepare_examples() +native_build.start() +# Preload once per server process, independently of visitor sessions. +backend.start_preparation() + + +def deployment_status(): + return {'native_build': native_build.status(), 'models': backend.model_status()} + + +@spaces.GPU(duration=900, size='large') +def generate_on_gpu(input_image, seed, progress, generation_token=None): + mesh, _asset, report = backend.generate(input_image, int(seed), progress=progress, + generation_token=generation_token) + return mesh, report + + +def process_image_generator(input_image, seed=42, generation_token=None, progress=gr.Progress()): + if input_image is None: + raise gr.Error('Please upload an image.') + print('[generation] Explicit Generate request; running released model without cache reuse', flush=True) + token = generation_token if generation_progress.valid(generation_token) else generation_progress.start() + try: + backend.ensure_prepared(token) + generation_progress.update(token, status='queued', label='Waiting for GPU') + mesh, report = generate_on_gpu(input_image, int(seed), progress, token) + generation_progress.update(token, status='saving', label='Saving the 3D asset') + try: + identifier = ovoxel_results.persist(mesh, report) + except Exception as error: + print(f'[result-cache] Could not persist result: {type(error).__name__}: {error}', flush=True) + raise + # The pipeline already rendered the saved asset. Show those views now; + # exporting a GLB belongs to the separate Extract GLB action. + rendered = ovoxel_results.preview_html(identifier) + generation_progress.update(token, status='complete', label='Complete', finished_unix=time.time()) + except Exception: + generation_progress.update(token, status='failed', label='Generation failed', finished_unix=time.time()) + raise + return rendered, identifier + + +@spaces.GPU(duration=600, size='large') +def extract_glb(identifier, extraction_token=None): + return ovoxel_results.extract(identifier, extraction_token) + + +@spaces.GPU(duration=600, size='large') +def recover_preview(identifier): + from hf_space.preview_recovery import recover + return recover(identifier) + + +from hf_space.ui import build_ui, CSS, THEME + +demo = build_ui(process_image_generator, deployment_status, cached_result=result_cache.lookup, extract_glb=extract_glb, recover_preview=recover_preview, + default_image=None, default_model=None, estimated_seconds=390, + example_names=result_cache.EXAMPLE_NAMES, examples_status=result_cache.examples_status) +demo.queue().launch(allowed_paths=['/tmp/supratex-jobs', str(result_cache.ROOT/'assets/previews'), str(result_cache.CACHE),str(ovoxel_results.ROOT), str(result_cache.ROOT/'hf_space/gallery-viewer.bundle.js'), str(result_cache.ROOT/'hf_space/gallery-lights')], + css=CSS, theme=THEME) diff --git a/configs/pipeline.json b/configs/pipeline.json new file mode 100644 index 0000000000000000000000000000000000000000..1b9004548ece048f31bf5aaa9c69de04af643d09 --- /dev/null +++ b/configs/pipeline.json @@ -0,0 +1,98 @@ +{ + "format": "supratex_pipeline", + "schema_version": 1, + "models": { + "completion": "completion", + "material": "material" + }, + "external_models": { + "qwen-base": { + "repo_id": "Qwen/Qwen-Image-2.1", + "revision": "d26bb61231c349cf6b7896fa83353113880e1ba3", + "allow_patterns": [ + "processor/*", + "vae/*", + "text_encoder/*", + "transformer/*", + "LICENSE*", + "README.md" + ] + }, + "geometry": { + "repo_id": "TencentARC/Pixal3D", + "revision": "b0cb2e1b794cab9aa0ac38a95d794a4d9337437f", + "allow_patterns": [ + "ckpts/ss_flow_img_dit_1_3B_64_bf16.json", + "ckpts/ss_flow_img_dit_1_3B_64_bf16.safetensors", + "ckpts/ss_dec_conv3d_16l8_fp16.json", + "ckpts/ss_dec_conv3d_16l8_fp16.safetensors", + "ckpts/slat_flow_img2shape_dit_1_3B_512_bf16.json", + "ckpts/slat_flow_img2shape_dit_1_3B_512_bf16.safetensors", + "ckpts/slat_flow_img2shape_dit_1_3B_1024_bf16.json", + "ckpts/slat_flow_img2shape_dit_1_3B_1024_bf16.safetensors", + "ckpts/shape_dec_next_dc_f16c32_fp16.json", + "ckpts/shape_dec_next_dc_f16c32_fp16.safetensors" + ] + }, + "moge2": { + "repo_id": "Ruicheng/moge-2-vitl-normal", + "revision": "cb0e8bbd6b1e243589717c78e750b1ba4c093acf", + "allow_patterns": [ + "model.pt", + "README.md", + "LICENSE*" + ] + }, + "dino": { + "repo_id": "facebook/dinov3-vitl16-pretrain-lvd1689m", + "revision": "ea8dc2863c51be0a264bab82070e3e8836b02d51", + "gated": true, + "allow_patterns": [ + "config.json", + "model.safetensors", + "preprocessor_config.json", + "LICENSE*", + "README.md" + ] + }, + "rmbg2": { + "repo_id": "briaai/RMBG-2.0", + "revision": "5df4c9c76d8170882c34f6986e848ee07fd0ba43", + "gated": true, + "allow_patterns": [ + "config.json", + "model.safetensors", + "BiRefNet_config.py", + "birefnet.py", + "preprocessor_config.json", + "LICENSE*", + "README.md" + ] + }, + "vae": { + "repo_id": "microsoft/TRELLIS.2-4B", + "revision": "af44b45f2e35a493886929c6d786e563ec68364d", + "allow_patterns": [ + "ckpts/shape_enc_next_dc_f16c32_fp16.json", + "ckpts/shape_enc_next_dc_f16c32_fp16.safetensors", + "ckpts/tex_enc_next_dc_f16c32_fp16.json", + "ckpts/tex_enc_next_dc_f16c32_fp16.safetensors", + "ckpts/shape_dec_next_dc_f16c32_fp16.json", + "ckpts/shape_dec_next_dc_f16c32_fp16.safetensors", + "ckpts/tex_dec_next_dc_f16c32_fp16.json", + "ckpts/tex_dec_next_dc_f16c32_fp16.safetensors" + ] + }, + "naf": { + "url": "https://github.com/valeoai/NAF/releases/download/model/naf_release.pth", + "filename": "naf_release.pth" + } + }, + "defaults": { + "seed": 42, + "texture_size": 4096, + "decimation_target": 500000, + "texture_format": "png", + "attention": "auto" + } +} diff --git a/deployment/runtime/SOURCE_PROVENANCE.json b/deployment/runtime/SOURCE_PROVENANCE.json index ad7de677ff449e94e0a9ba72797aaf3d0f3dfad3..ab960ffdcdee487292ce9dfb5086095681fe2ead 100644 --- a/deployment/runtime/SOURCE_PROVENANCE.json +++ b/deployment/runtime/SOURCE_PROVENANCE.json @@ -1,29 +1,26 @@ -{ - "verification": "source file presence and sizes only; no hashes calculated", - "geometry": { - "source": "local Pixal3D/pixal3d source tree", - "git_head": "0b9ff55cbcf5f87aaae9948742cf6388a1a3ecb1", - "note": "Working tree snapshot; HEAD alone does not capture local edits" - }, - "qwen": { - "source": "ckpts/qwen-runtime", - "snapshot_schema": "qwen_pbr_runtime_snapshot_v1", - "snapshot_created_utc": "2026-10-03T23:58:03.610370+00:00", - "note": "Exact staged inference source; upstream git identity is not recorded in staged manifest" - }, - "moge": { - "source": "installed moge 2.0.0 source", - "upstream": "https://github.com/microsoft/MoGe.git", - "recorded_commit": "07444410f1e33f402353b99d6ccd26bd31e469e8" - }, - "utils3d": { - "source": "installed utils3d 0.0.2 Python source", - "upstream": "https://github.com/EasternJournalist/utils3d", - "note": "Installed wheel provenance, no git commit recorded" - }, - "naf": { - "source": "local valeoai_NAF_main Torch Hub source cache", - "upstream": "https://github.com/valeoai/NAF", - "note": "No git metadata in staged torch hub cache; bundled source is the pin" - } -} +{ + "verification": "source file presence and sizes only; no hashes calculated", + "geometry": { + "source": "Bundled Pixal3D runtime based on TencentARC/Pixal3D and Microsoft TRELLIS.2", + "note": "Vendored files are the distribution reference; an exact upstream Git revision is unavailable." + }, + "qwen": { + "source": "Bundled native Qwen-Image-2.1 material runtime", + "note": "Vendored files are the distribution reference; an exact upstream Git revision is unavailable." + }, + "moge": { + "source": "installed moge 2.0.0 source", + "upstream": "https://github.com/microsoft/MoGe.git", + "recorded_commit": "07444410f1e33f402353b99d6ccd26bd31e469e8" + }, + "utils3d": { + "source": "installed utils3d 0.0.2 Python source", + "upstream": "https://github.com/EasternJournalist/utils3d", + "note": "Installed wheel provenance, no git commit recorded" + }, + "naf": { + "source": "Bundled NAF runtime", + "upstream": "https://github.com/valeoai/NAF", + "note": "Bundled source is the distribution reference; an exact upstream Git revision is unavailable." + } +} diff --git a/deployment/runtime/geometry/pixal3d/models/__init__.py b/deployment/runtime/geometry/pixal3d/models/__init__.py index 15efa73b114edbc57ca52e993e2031a8f96da3b9..4e002253d8a7c5c2a8e73701f830bffb3dfaa02b 100644 --- a/deployment/runtime/geometry/pixal3d/models/__init__.py +++ b/deployment/runtime/geometry/pixal3d/models/__init__.py @@ -1,87 +1,78 @@ -import importlib - -__attributes = { - # Sparse Structure - 'SparseStructureEncoder': 'sparse_structure_vae', - 'SparseStructureDecoder': 'sparse_structure_vae', - 'SparseStructureFlowModel': 'sparse_structure_flow', - - # SLat Generation - 'SLatFlowModel': 'structured_latent_flow', - 'ElasticSLatFlowModel': 'structured_latent_flow', - 'V13PSAFlowModel': 'v13psa', - 'V13PSAILVFlowModel': 'v13psailv', - 'V14FXFlowModel': 'v14fx', - - # SC-VAEs - 'SparseUnetVaeEncoder': 'sc_vaes.sparse_unet_vae', - 'SparseUnetVaeDecoder': 'sc_vaes.sparse_unet_vae', - 'FlexiDualGridVaeEncoder': 'sc_vaes.fdg_vae', - 'FlexiDualGridVaeDecoder': 'sc_vaes.fdg_vae' -} - -__submodules = [] - -__all__ = list(__attributes.keys()) + __submodules - -def __getattr__(name): - if name not in globals(): - if name in __attributes: - module_name = __attributes[name] - module = importlib.import_module(f".{module_name}", __name__) - globals()[name] = getattr(module, name) - elif name in __submodules: - module = importlib.import_module(f".{name}", __name__) - globals()[name] = module - else: - raise AttributeError(f"module {__name__} has no attribute {name}") - return globals()[name] - - -def from_pretrained(path: str, **kwargs): - """ - Load a model from a pretrained checkpoint. - - Args: - path: The path to the checkpoint. Can be either local path or a Hugging Face model name. - NOTE: config file and model file should take the name f'{path}.json' and f'{path}.safetensors' respectively. - **kwargs: Additional arguments for the model constructor. - """ - import os - import json - from safetensors.torch import load_file - is_local = os.path.exists(f"{path}.json") and os.path.exists(f"{path}.safetensors") - - if is_local: - config_file = f"{path}.json" - model_file = f"{path}.safetensors" - else: - from huggingface_hub import hf_hub_download - path_parts = path.split('/') - repo_id = f'{path_parts[0]}/{path_parts[1]}' - model_name = '/'.join(path_parts[2:]) - config_file = hf_hub_download(repo_id, f"{model_name}.json") - model_file = hf_hub_download(repo_id, f"{model_name}.safetensors") - - with open(config_file, 'r') as f: - config = json.load(f) - if config['name'] == 'V14FXFlowModel': - raise ValueError( - 'V14FX requires strict native-checkpoint loading through ' - 'pixal3d.pipelines.build_v14fx_inference_model; ' - 'the generic non-strict pretrained loader is unsupported' - ) - model = __getattr__(config['name'])(**config['args'], **kwargs) - model.load_state_dict(load_file(model_file), strict=False) - - return model - - -# For Pylance -if __name__ == '__main__': - from .sparse_structure_vae import SparseStructureEncoder, SparseStructureDecoder - from .sparse_structure_flow import SparseStructureFlowModel - from .structured_latent_flow import SLatFlowModel, ElasticSLatFlowModel - - from .sc_vaes.sparse_unet_vae import SparseUnetVaeEncoder, SparseUnetVaeDecoder - from .sc_vaes.fdg_vae import FlexiDualGridVaeEncoder, FlexiDualGridVaeDecoder +import importlib + +__attributes = { + # Sparse Structure + 'SparseStructureEncoder': 'sparse_structure_vae', + 'SparseStructureDecoder': 'sparse_structure_vae', + 'SparseStructureFlowModel': 'sparse_structure_flow', + + # SLat Generation + 'SLatFlowModel': 'structured_latent_flow', + 'ElasticSLatFlowModel': 'structured_latent_flow', + + # SC-VAEs + 'SparseUnetVaeEncoder': 'sc_vaes.sparse_unet_vae', + 'SparseUnetVaeDecoder': 'sc_vaes.sparse_unet_vae', + 'FlexiDualGridVaeEncoder': 'sc_vaes.fdg_vae', + 'FlexiDualGridVaeDecoder': 'sc_vaes.fdg_vae' +} + +__submodules = [] + +__all__ = list(__attributes.keys()) + __submodules + +def __getattr__(name): + if name not in globals(): + if name in __attributes: + module_name = __attributes[name] + module = importlib.import_module(f".{module_name}", __name__) + globals()[name] = getattr(module, name) + elif name in __submodules: + module = importlib.import_module(f".{name}", __name__) + globals()[name] = module + else: + raise AttributeError(f"module {__name__} has no attribute {name}") + return globals()[name] + + +def from_pretrained(path: str, **kwargs): + """ + Load a model from a pretrained checkpoint. + + Args: + path: The path to the checkpoint. Can be either local path or a Hugging Face model name. + NOTE: config file and model file should take the name f'{path}.json' and f'{path}.safetensors' respectively. + **kwargs: Additional arguments for the model constructor. + """ + import os + import json + from safetensors.torch import load_file + is_local = os.path.exists(f"{path}.json") and os.path.exists(f"{path}.safetensors") + + if is_local: + config_file = f"{path}.json" + model_file = f"{path}.safetensors" + else: + from huggingface_hub import hf_hub_download + path_parts = path.split('/') + repo_id = f'{path_parts[0]}/{path_parts[1]}' + model_name = '/'.join(path_parts[2:]) + config_file = hf_hub_download(repo_id, f"{model_name}.json") + model_file = hf_hub_download(repo_id, f"{model_name}.safetensors") + + with open(config_file, 'r') as f: + config = json.load(f) + model = __getattr__(config['name'])(**config['args'], **kwargs) + model.load_state_dict(load_file(model_file), strict=False) + + return model + + +# For Pylance +if __name__ == '__main__': + from .sparse_structure_vae import SparseStructureEncoder, SparseStructureDecoder + from .sparse_structure_flow import SparseStructureFlowModel + from .structured_latent_flow import SLatFlowModel + + from .sc_vaes.sparse_unet_vae import SparseUnetVaeEncoder, SparseUnetVaeDecoder + from .sc_vaes.fdg_vae import FlexiDualGridVaeEncoder, FlexiDualGridVaeDecoder diff --git a/deployment/runtime/geometry/pixal3d/models/structured_latent_flow.py b/deployment/runtime/geometry/pixal3d/models/structured_latent_flow.py index 50cb834086fc95d81d368fbc6a3fceaf122ebd68..3f2d09bf5009efe7808c0f974ee5a1bc082141d5 100644 --- a/deployment/runtime/geometry/pixal3d/models/structured_latent_flow.py +++ b/deployment/runtime/geometry/pixal3d/models/structured_latent_flow.py @@ -1,1391 +1,1387 @@ -from typing import * -from functools import partial -import torch -import torch.nn as nn -import torch.nn.functional as F -import numpy as np -from ..modules.utils import convert_module_to, manual_cast, str_to_dtype -from ..modules.transformer import AbsolutePositionEmbedder -from ..modules import sparse as sp -from ..modules.sparse.transformer import ModulatedSparseTransformerCrossBlock -from .sparse_structure_flow import TimestepEmbedder -from .sparse_elastic_mixin import SparseTransformerElasticMixin - - -class _PartialZObsInnovationGate(nn.Module): - """Residual, per-observed-token correction for the partial z_obs carrier.""" - - def __init__(self, model_channels: int, message_channels: int, dim: int, time_dim: int, num_blocks: int): - super().__init__() - self.alpha = nn.Parameter(torch.zeros(num_blocks)) - self.embed_y = nn.Linear(message_channels, dim, bias=True) - self.time_proj = nn.Linear(model_channels, time_dim, bias=True) - self.embed_x = nn.Linear(model_channels + time_dim, dim, bias=True) - self.mlp = nn.Sequential( - nn.Linear(dim + 2 + time_dim, dim, bias=True), - nn.SiLU(), - nn.Linear(dim, 1, bias=True), - ) - self.proj = nn.Linear(dim, model_channels, bias=True) - - -_LOCAL_ONE_HOP_OFFSETS = ( - (-1, 0, 0), (0, -1, 0), (0, 0, -1), - (0, 0, 1), (0, 1, 0), (1, 0, 0), -) -_LOCAL_TWO_HOP_OFFSETS = tuple(sorted( - (dx, dy, dz) - for dx in range(-2, 3) - for dy in range(-2, 3) - for dz in range(-2, 3) - if abs(dx) + abs(dy) + abs(dz) <= 2 -)) -assert len(_LOCAL_TWO_HOP_OFFSETS) == 25 - - -def _exact_coordinate_rows( - full_coords: torch.Tensor, - candidate_coords: torch.Tensor, -) -> Tuple[torch.Tensor, torch.Tensor]: - """Return exact full-support rows and a validity mask for candidate coords.""" - assert full_coords.ndim == candidate_coords.ndim == 2 - assert full_coords.shape[1] == candidate_coords.shape[1] - assert full_coords.shape[0] > 0, 'local attention requires nonempty full sparse support' - all_coords = torch.cat([full_coords, candidate_coords], dim=0) - _, inverse = torch.unique(all_coords, dim=0, return_inverse=True) - full_ids = inverse[:full_coords.shape[0]] - candidate_ids = inverse[full_coords.shape[0]:] - assert torch.unique(full_ids).numel() == full_ids.numel(), \ - 'x contains duplicate coordinates' - ordered_ids, ordered_rows = torch.sort(full_ids) - positions = torch.searchsorted(ordered_ids, candidate_ids) - in_range = positions < ordered_ids.numel() - safe_positions = positions.clamp(max=ordered_ids.numel() - 1) - valid = in_range & (ordered_ids[safe_positions] == candidate_ids) - rows = torch.where(valid, ordered_rows[safe_positions], -torch.ones_like(safe_positions)) - return rows, valid - - -def _build_two_hop_observed_topology( - full_coords: torch.Tensor, - observed_rows: torch.Tensor, -) -> Dict[str, torch.Tensor]: - """Pack exact batch-local G(S) distance<=2 observed source/query pairs.""" - assert full_coords.ndim == 2 and full_coords.shape[1] == 4, \ - f'local attention requires [N,4] integer coords, got {tuple(full_coords.shape)}' - assert not torch.is_floating_point(full_coords), 'local attention coordinates must be integer' - assert observed_rows.ndim == 1 and observed_rows.dtype == torch.long - # Validate full-coordinate uniqueness even for an all-absent batch. - _exact_coordinate_rows(full_coords, full_coords.new_empty((0, full_coords.shape[1]))) - if observed_rows.numel() == 0: - empty = torch.empty(0, dtype=torch.long, device=full_coords.device) - return { - 'query_rows': empty, - 'source_indices': empty, - 'offset_indices': empty, - 'segment_indices': empty, - 'neighbour_counts': empty, - } - assert observed_rows.min() >= 0 and observed_rows.max() < full_coords.shape[0], \ - 'observed row lies outside full sparse support' - assert torch.unique(observed_rows).numel() == observed_rows.numel(), \ - 'observed rows must be unique' - - device = full_coords.device - one_hop = torch.tensor( - _LOCAL_ONE_HOP_OFFSETS, dtype=full_coords.dtype, device=device, - ) - source_coords = full_coords[observed_rows] - candidate_hop1 = source_coords[:, None, :].expand(-1, len(_LOCAL_ONE_HOP_OFFSETS), -1).clone() - candidate_hop1[:, :, 1:] += one_hop[None] - hop1_rows, hop1_valid = _exact_coordinate_rows( - full_coords, candidate_hop1.reshape(-1, full_coords.shape[1]), - ) - source_ids_hop1 = torch.arange( - observed_rows.numel(), device=device, dtype=torch.long, - )[:, None].expand(-1, len(_LOCAL_ONE_HOP_OFFSETS)).reshape(-1) - hop1_rows = hop1_rows[hop1_valid] - source_ids_hop1 = source_ids_hop1[hop1_valid] - - if hop1_rows.numel() != 0: - candidate_hop2 = full_coords[hop1_rows][:, None, :].expand( - -1, len(_LOCAL_ONE_HOP_OFFSETS), -1, - ).clone() - candidate_hop2[:, :, 1:] += one_hop[None] - hop2_rows, hop2_valid = _exact_coordinate_rows( - full_coords, candidate_hop2.reshape(-1, full_coords.shape[1]), - ) - source_ids_hop2 = source_ids_hop1[:, None].expand( - -1, len(_LOCAL_ONE_HOP_OFFSETS), - ).reshape(-1) - hop2_rows = hop2_rows[hop2_valid] - source_ids_hop2 = source_ids_hop2[hop2_valid] - else: - hop2_rows = torch.empty(0, dtype=torch.long, device=device) - source_ids_hop2 = torch.empty(0, dtype=torch.long, device=device) - - source_ids = torch.cat([ - torch.arange(observed_rows.numel(), device=device, dtype=torch.long), - source_ids_hop1, - source_ids_hop2, - ]) - query_rows = torch.cat([observed_rows, hop1_rows, hop2_rows]) - # Multiple graph paths to the same query/source pair are one neighbourhood entry. - unique_pairs = torch.unique(torch.stack([query_rows, source_ids], dim=1), dim=0) - query_rows = unique_pairs[:, 0] - source_ids = unique_pairs[:, 1] - query_coords = full_coords[query_rows] - selected_source_coords = source_coords[source_ids] - assert torch.equal(query_coords[:, 0], selected_source_coords[:, 0]), \ - 'local attention constructed a cross-sample edge' - relative = query_coords[:, 1:] - selected_source_coords[:, 1:] - offset_table = torch.tensor( - _LOCAL_TWO_HOP_OFFSETS, dtype=full_coords.dtype, device=device, - ) - offset_matches = (relative[:, None, :] == offset_table[None, :, :]).all(dim=-1) - assert torch.equal(offset_matches.sum(dim=1), torch.ones_like(offset_matches[:, 0], dtype=torch.long)), \ - 'local attention produced an invalid two-hop offset' - offset_indices = offset_matches.to(dtype=torch.long).argmax(dim=1) - - # Canonical per-query real-KV order is lexicographic relative offset, independent - # of either full-support or observed-message row ordering. - order = torch.argsort(query_rows * len(_LOCAL_TWO_HOP_OFFSETS) + offset_indices) - query_rows = query_rows[order] - source_ids = source_ids[order] - offset_indices = offset_indices[order] - active_query_rows, segment_indices, counts = torch.unique_consecutive( - query_rows, return_inverse=True, return_counts=True, - ) - assert counts.max() <= len(_LOCAL_TWO_HOP_OFFSETS), \ - 'local observed neighbourhood exceeds frozen 25-source cap' - return { - 'query_rows': active_query_rows, - 'source_indices': source_ids, - 'offset_indices': offset_indices, - 'segment_indices': segment_indices, - 'neighbour_counts': counts, - } - - -class _PartialZObsLocalAttention(nn.Module): - """Shared packed two-hop observed-innovation attention for V12-A.""" - - def __init__( - self, - model_channels: int, - message_channels: int, - dim: int, - time_dim: int, - num_heads: int, - num_active_blocks: int, - ): - super().__init__() - assert dim == 128 and time_dim == 32 and num_heads == 8 - assert dim % num_heads == 0 - self.dim = dim - self.num_heads = num_heads - self.head_dim = dim // num_heads - self.alpha = nn.Parameter(torch.zeros(num_active_blocks)) - self.time_proj = nn.Linear(model_channels, time_dim, bias=True) - self.embed_x = nn.Linear(model_channels + time_dim, dim, bias=True) - self.embed_y = nn.Linear(message_channels, dim, bias=True) - self.key_proj = nn.Linear(dim, dim, bias=True) - self.value_proj = nn.Linear(dim, dim, bias=True) - self.relative_offset = nn.Parameter(torch.empty(len(_LOCAL_TWO_HOP_OFFSETS), dim)) - self.null_key = nn.Parameter(torch.empty(num_heads, self.head_dim)) - self.out_proj = nn.Linear(dim, model_channels, bias=False) - - @staticmethod - def _rms_norm(value: torch.Tensor) -> torch.Tensor: - return value * torch.rsqrt(value.square().mean(dim=-1, keepdim=True) + 1e-6) - - def forward( - self, - hidden: torch.Tensor, - full_batch: torch.Tensor, - t_base: torch.Tensor, - message: torch.Tensor, - observed_rows: torch.Tensor, - topology: Dict[str, torch.Tensor], - block_index: int, - ) -> Tuple[torch.Tensor, Dict[str, torch.Tensor]]: - assert 0 <= block_index < self.alpha.numel(), 'terminal block has no V12-A residual' - query_rows = topology['query_rows'] - source_indices = topology['source_indices'] - offset_indices = topology['offset_indices'] - segment_indices = topology['segment_indices'] - counts = topology['neighbour_counts'] - assert full_batch.ndim == 1 and full_batch.shape[0] == hidden.shape[0] - assert full_batch.numel() == 0 or ( - full_batch.min() >= 0 and full_batch.max() < t_base.shape[0] - ), 'full sparse batch coordinate is outside timestep batch' - assert message.shape == (observed_rows.numel(), 34) - - if query_rows.numel() == 0: - live_zero = sum(parameter.sum() for parameter in self.parameters()) * 0 - empty_delta = hidden.new_empty((0, hidden.shape[1])) + live_zero - zero = hidden.new_zeros(()) - return empty_delta, { - 'null_attention_mean': zero, - 'real_neighbour_entropy_mean': zero, - 'effective_real_neighbour_count_mean': zero, - 'active_queries': zero, - 'singleton_queries': zero, - 'packed_pairs': zero, - } - - branch_dtype = self.embed_x.weight.dtype - hidden_branch = hidden.to(dtype=branch_dtype) - time = self.time_proj(t_base.to(dtype=branch_dtype)) - x_embed = self.embed_x(torch.cat([ - F.layer_norm(hidden_branch, hidden_branch.shape[-1:], eps=1e-6), - time[full_batch], - ], dim=-1)) - innovation = self.embed_y(message.to(dtype=branch_dtype)) - x_embed[observed_rows] - query = self._rms_norm( - x_embed[query_rows].reshape(-1, self.num_heads, self.head_dim) - ) - real_key = self.key_proj(innovation)[source_indices] - real_key = real_key + self.relative_offset[offset_indices] - real_key = self._rms_norm(real_key.reshape(-1, self.num_heads, self.head_dim)) - real_value = self.value_proj(innovation)[source_indices].reshape( - -1, self.num_heads, self.head_dim, - ) - real_logits = (query[segment_indices] * real_key).sum(dim=-1) / np.sqrt(self.head_dim) - null_key = self._rms_norm(self.null_key) - null_logits = (query * null_key[None]).sum(dim=-1) / np.sqrt(self.head_dim) - - segment_heads = segment_indices[:, None].expand(-1, self.num_heads) - maxima = null_logits.clone() - maxima.scatter_reduce_(0, segment_heads, real_logits, reduce='amax', include_self=True) - real_exp = torch.exp(real_logits - maxima[segment_indices]) - null_exp = torch.exp(null_logits - maxima) - real_sum = torch.zeros_like(null_exp).scatter_add(0, segment_heads, real_exp) - denominator = null_exp + real_sum - real_weight = real_exp / denominator[segment_indices] - null_weight = null_exp / denominator - attended = hidden_branch.new_zeros( - (query_rows.numel(), self.num_heads, self.head_dim), - ) - attended.index_add_( - 0, segment_indices, real_weight[..., None] * real_value, - ) - delta = self.out_proj(attended.reshape(query_rows.numel(), self.dim)) - delta = delta * self.alpha[block_index].to(dtype=delta.dtype) - - # Entropy is conditional on selecting real evidence. Null-vs-real - # acceptance is reported independently as null_attention_mean; mixing - # that mass into this distribution would make singleton entropy nonzero. - real_conditional_weight = real_exp / real_sum[segment_indices].clamp_min( - torch.finfo(real_sum.dtype).tiny - ) - entropy_terms = -real_conditional_weight * torch.log( - real_conditional_weight.clamp_min( - torch.finfo(real_conditional_weight.dtype).tiny - ) - ) - real_entropy = torch.zeros_like(null_weight).scatter_add( - 0, segment_heads, entropy_terms, - ) - return delta, { - 'null_attention_mean': null_weight.mean(), - 'real_neighbour_entropy_mean': real_entropy.mean(), - 'effective_real_neighbour_count_mean': real_entropy.exp().mean(), - 'active_queries': delta.new_tensor(float(query_rows.numel())), - 'singleton_queries': delta.new_tensor(float((counts == 1).sum())), - 'packed_pairs': delta.new_tensor(float(source_indices.numel())), - } - - -class _V12LACrossRope(nn.Module): - """Apply the same 3D rotary basis to independently packed Q and K rows.""" - - def __init__( - self, - head_dim: int, - rope_freq: Tuple[float, float], - reserved_dims: int = 2, - ): - super().__init__() - content_dim = head_dim - reserved_dims - assert content_dim > 0 and content_dim % 6 == 0 - self.head_dim = head_dim - self.content_dim = content_dim - freq_dim = content_dim // 2 // 3 - frequencies = torch.arange(freq_dim, dtype=torch.float32) / freq_dim - frequencies = rope_freq[0] / (rope_freq[1] ** frequencies) - self.register_buffer('frequencies', frequencies, persistent=False) - - def forward(self, value: torch.Tensor, coords: torch.Tensor) -> torch.Tensor: - assert value.ndim == 3 and value.shape[-1] == self.head_dim - assert coords.shape == (value.shape[0], 3) - if value.shape[0] == 0: - return value - phases = torch.outer( - coords.reshape(-1).float(), self.frequencies, - ).reshape(coords.shape[0], self.content_dim // 2) - phases = torch.polar(torch.ones_like(phases), phases).unsqueeze(1) - content = torch.view_as_complex( - value[..., :self.content_dim].float().reshape( - value.shape[0], value.shape[1], -1, 2, - ) - ) - rotated = torch.view_as_real(content * phases).reshape( - value.shape[0], value.shape[1], self.content_dim, - ).to(dtype=value.dtype) - return torch.cat([rotated, value[..., self.content_dim:]], dim=-1) - - -class _PartialZObsLearnedAttention(nn.Module): - """Five-site V12LA observation attention with an exact carrier anchor.""" - - def __init__( - self, - model_channels: int, - num_heads: int, - sites: Sequence[int], - rope_freq: Tuple[float, float], - null_prior: float = 0.5, - ): - super().__init__() - self.sites = tuple(int(site) for site in sites) - assert self.sites == (5, 10, 15, 20, 25), self.sites - assert num_heads == 12 - assert model_channels % num_heads == 0 - assert (model_channels // num_heads - 2) > 0 - assert (model_channels // num_heads - 2) % 6 == 0 - self.model_channels = model_channels - self.num_heads = num_heads - self.head_dim = model_channels // num_heads - self.content_dim = self.head_dim - 2 - self.site_to_slot = {site: slot for slot, site in enumerate(self.sites)} - self.null_prior = float(null_prior) - assert 0.0 < self.null_prior < 1.0, \ - f'learned-attention null prior must lie in (0,1), got {self.null_prior}' - self.null_logit_offset = float( - np.log(self.null_prior) - np.log1p(-self.null_prior) - ) - - # K/V are source-only and shared across all five sites. Q and the - # zero-start deviation/output matrix are site-specific. - self.key_proj = nn.Linear(model_channels, model_channels, bias=False) - self.value_proj = nn.Linear(model_channels, model_channels, bias=False) - self.query_proj = nn.ModuleList([ - nn.Linear(model_channels, model_channels, bias=False) - for _ in self.sites - ]) - self.delta_proj = nn.ModuleList([ - nn.Linear(model_channels, model_channels, bias=False) - for _ in self.sites - ]) - self.null_key = nn.Parameter(torch.empty(num_heads, self.content_dim)) - self.rope = _V12LACrossRope(self.head_dim, rope_freq, reserved_dims=2) - self.reset_parameters() - - def reset_parameters(self) -> None: - nn.init.xavier_uniform_(self.key_proj.weight) - # Attention mixes the already trained V11 carrier basis at warm start; - # the value map may then adapt, but does not begin as a random rewrite. - nn.init.eye_(self.value_proj.weight) - nn.init.xavier_uniform_(self.null_key) - for query_proj in self.query_proj: - nn.init.xavier_uniform_(query_proj.weight) - for delta_proj in self.delta_proj: - nn.init.zeros_(delta_proj.weight) - - @staticmethod - def _rms_norm(value: torch.Tensor) -> torch.Tensor: - normalized = value.float() * torch.rsqrt( - value.float().square().mean(dim=-1, keepdim=True) + 1e-6 - ) - return normalized.to(dtype=value.dtype) - - @staticmethod - def _anchor( - carrier: torch.Tensor, - observed_rows: torch.Tensor, - full_rows: int, - ) -> torch.Tensor: - """Evaluate fixed coordinate/null attention without an additive scatter.""" - assert carrier.ndim == 2 - assert observed_rows.ndim == 1 and observed_rows.dtype == torch.long - assert carrier.shape[0] == observed_rows.numel() - assert full_rows >= observed_rows.numel() - if observed_rows.numel() != 0: - assert observed_rows.min() >= 0 and observed_rows.max() < full_rows - assert torch.unique(observed_rows).numel() == observed_rows.numel() - # Index zero is the exact-zero null value. Every observed row is then - # redirected to its own carrier row; unmatched rows keep selecting null. - values = torch.cat([carrier.new_zeros((1, carrier.shape[1])), carrier], dim=0) - anchor_index = torch.zeros(full_rows, dtype=torch.long, device=carrier.device) - anchor_index[observed_rows] = torch.arange( - 1, observed_rows.numel() + 1, dtype=torch.long, device=carrier.device, - ) - return values[anchor_index] - - def build_cache( - self, - full_coords: torch.Tensor, - full_layout: List[slice], - carrier: torch.Tensor, - observed_rows: torch.Tensor, - ) -> Dict[str, Any]: - """Build the exact anchor and the shared per-sample O+null K/V cache.""" - assert full_coords.ndim == 2 and full_coords.shape[1] == 4 - assert carrier.shape == (observed_rows.numel(), self.model_channels) - assert len(full_layout) > 0 - anchor = self._anchor(carrier, observed_rows, full_coords.shape[0]) - observed_coords = full_coords[observed_rows] - observed_batch = observed_coords[:, 0].long() - assert observed_batch.numel() == 0 or ( - observed_batch.min() >= 0 and observed_batch.max() < len(full_layout) - ) - - branch_dtype = self.key_proj.weight.dtype - carrier_branch = carrier.to(dtype=branch_dtype) - carrier_norm = F.layer_norm( - carrier_branch, (self.model_channels,), eps=1e-6, - ) - real_key = self.key_proj(carrier_norm).reshape( - -1, self.num_heads, self.head_dim, - ) - real_key = self.rope(real_key, observed_coords[:, 1:]) - real_key = self._rms_norm(real_key[..., :self.content_dim]) - real_key = torch.cat([ - real_key, - real_key.new_zeros((real_key.shape[0], self.num_heads, 2)), - ], dim=-1) - real_value = self.value_proj(carrier_branch).reshape( - -1, self.num_heads, self.head_dim, - ) - null_content = self._rms_norm( - self.null_key.to(dtype=real_key.dtype) - ) - - packed_kv = [] - observed_counts = [] - for batch_index in range(len(full_layout)): - source_indices = torch.nonzero( - observed_batch == batch_index, as_tuple=False, - ).flatten() - observed_count = int(source_indices.numel()) - observed_counts.append(observed_count) - # One null competes with O real keys. log(O) cancels its cardinality - # disadvantage; the constant logit offset sets the null mass only in - # the equal-content-logit reference case, not after learned Q/K drift. - null_reserved = real_key.new_zeros((self.num_heads, 2)) - null_reserved[:, 0] = np.sqrt(self.head_dim) * ( - np.log(max(observed_count, 1)) + self.null_logit_offset - ) - null_key = torch.cat([null_content, null_reserved], dim=-1).unsqueeze(0) - null_value = real_value.new_zeros((1, self.num_heads, self.head_dim)) - keys = torch.cat([real_key[source_indices], null_key], dim=0) - values = torch.cat([real_value[source_indices], null_value], dim=0) - packed_kv.append(torch.stack([keys, values], dim=1)) - - kv = sp.VarLenTensor.from_tensor_list(packed_kv) - return { - 'anchor': anchor, - 'kv': kv, - 'observed_rows': observed_rows, - 'observed_counts': tuple(observed_counts), - } - - def forward( - self, - hidden: torch.Tensor, - full_coords: torch.Tensor, - full_layout: List[slice], - cache: Dict[str, Any], - block_index: int, - ) -> Tuple[torch.Tensor, Dict[str, torch.Tensor]]: - assert block_index in self.site_to_slot, block_index - assert hidden.shape == (full_coords.shape[0], self.model_channels) - slot = self.site_to_slot[block_index] - branch_dtype = self.query_proj[slot].weight.dtype - hidden_branch = hidden.to(dtype=branch_dtype) - query = self.query_proj[slot](F.layer_norm( - hidden_branch, (self.model_channels,), eps=1e-6, - )).reshape(-1, self.num_heads, self.head_dim) - query = self.rope(query, full_coords[:, 1:]) - query_content = self._rms_norm(query[..., :self.content_dim]) - query_reserved = query.new_zeros((query.shape[0], self.num_heads, 2)) - query_reserved[..., 0] = 1 - query = torch.cat([query_content, query_reserved], dim=-1) - query_varlen = sp.VarLenTensor(query, full_layout) - attended = sp.sparse_scaled_dot_product_attention( - query_varlen, cache['kv'], - ).feats.reshape(-1, self.model_channels) - - anchor = cache['anchor'].to(dtype=attended.dtype) - deviation_input = attended - anchor - correction = self.delta_proj[slot](deviation_input) - inject = anchor + correction - eps = torch.finfo(inject.dtype).eps - observed_rows = cache['observed_rows'] - observed_mask = torch.zeros( - inject.shape[0], dtype=torch.bool, device=inject.device, - ) - observed_mask[observed_rows] = True - return inject, { - 'correction_to_anchor_norm_ratio': correction.norm() / anchor.norm().clamp_min(eps), - 'correction_to_hidden_norm_ratio': correction.norm() / hidden_branch.norm().clamp_min(eps), - 'observed_correction_l2_norm': correction[observed_mask].float().norm(), - 'unobserved_correction_l2_norm': correction[~observed_mask].float().norm(), - 'attention_output_l2_norm': attended.float().norm(), - 'anchor_l2_norm': anchor.float().norm(), - 'delta_weight_l2_norm': self.delta_proj[slot].weight.float().norm(), - } - -class SLatFlowModel(nn.Module): - """ - Structured Latent Flow Model for 3D generation. - - Supports two conditioning modes: - - "cross": Standard cross-attention with image features - - "proj": View-aligned projection attention with camera-aware features - """ - def __init__( - self, - resolution: int, - in_channels: int, - model_channels: int, - cond_channels: int, - out_channels: int, - num_blocks: int, - num_heads: Optional[int] = None, - num_head_channels: Optional[int] = 64, - mlp_ratio: float = 4, - pe_mode: Literal["ape", "rope"] = "ape", - rope_freq: Tuple[float, float] = (1.0, 10000.0), - dtype: str = 'float32', - use_checkpoint: bool = False, - share_mod: bool = False, - initialization: str = 'vanilla', - qk_rms_norm: bool = False, - qk_rms_norm_cross: bool = False, - image_attn_mode: Literal["cross", "proj", "gated_proj"] = "cross", - proj_in_channels: Optional[int] = None, - vae_in_channels: Optional[int] = None, - zobs_inject: bool = False, - zobs_inject_in_channels: int = 33, - zobs_inject_scale: float = 1.0, - partial_zobs_inject: bool = False, - partial_zobs_in_channels: int = 34, - partial_zobs_inject_position: Literal['post_block', 'pre_mlp'] = 'post_block', - legacy_zobs_zero_input: bool = False, - partial_zobs_innovation_gate: bool = False, - innovation_gate_dim: int = 128, - innovation_gate_time_dim: int = 32, - partial_zobs_innovation_gate_train_only: bool = False, - partial_zobs_local_attention: bool = False, - local_attention_dim: int = 128, - local_attention_time_dim: int = 32, - local_attention_heads: int = 8, - partial_zobs_local_attention_train_only: bool = False, - partial_zobs_learned_attention: bool = False, - learned_attention_sites: Sequence[int] = (5, 10, 15, 20, 25), - learned_attention_heads: int = 12, - learned_attention_null_prior: float = 0.5, - partial_zobs_learned_attention_train_only: bool = False, - **kwargs - ): - super().__init__() - self.resolution = resolution - self.in_channels = in_channels - self.model_channels = model_channels - self.cond_channels = cond_channels - self.out_channels = out_channels - self.num_blocks = num_blocks - self.num_heads = num_heads or model_channels // num_head_channels - self.mlp_ratio = mlp_ratio - self.pe_mode = pe_mode - self.use_checkpoint = use_checkpoint - self.share_mod = share_mod - self.initialization = initialization - self.qk_rms_norm = qk_rms_norm - self.qk_rms_norm_cross = qk_rms_norm_cross - self.image_attn_mode = image_attn_mode - self.proj_in_channels = proj_in_channels - self.vae_in_channels = vae_in_channels - self.dtype = str_to_dtype(dtype) - - self.t_embedder = TimestepEmbedder(model_channels) - if share_mod: - self.adaLN_modulation = nn.Sequential( - nn.SiLU(), - nn.Linear(model_channels, 6 * model_channels, bias=True) - ) - - if pe_mode == "ape": - self.pos_embedder = AbsolutePositionEmbedder(model_channels) - - self.input_layer = sp.SparseLinear(in_channels, model_channels) - # Trained-bridge: single shared per-block injector of [z_obs (x) vis (32) | vis (1)] -> model_channels. - self.zobs_inject = zobs_inject - self.zobs_inject_scale = float(zobs_inject_scale) - self.zobs_injector = sp.SparseLinear(zobs_inject_in_channels, model_channels) if zobs_inject else None - # v8-partial: the observation lives only on its observed support. This adapter is - # deliberately bias-free: absent rows never acquire a learned content/fill value. - self.partial_zobs_inject = partial_zobs_inject - self.partial_zobs_in_channels = partial_zobs_in_channels - assert partial_zobs_inject_position in ('post_block', 'pre_mlp'), \ - f'partial_zobs_inject_position must be post_block or pre_mlp, got {partial_zobs_inject_position!r}' - self.partial_zobs_inject_position = partial_zobs_inject_position - if partial_zobs_inject_position == 'pre_mlp': - assert not any(( - partial_zobs_innovation_gate, - partial_zobs_local_attention, - partial_zobs_learned_attention, - )), 'partial_zobs_inject_position=pre_mlp is incompatible with post-block experimental branches' - self.legacy_zobs_zero_input = legacy_zobs_zero_input - self.partial_zobs_injector = ( - sp.SparseLinear(partial_zobs_in_channels, model_channels, bias=False) - if partial_zobs_inject else None - ) - self.partial_zobs_innovation_gate = partial_zobs_innovation_gate - self.partial_zobs_innovation_gate_train_only = partial_zobs_innovation_gate_train_only - if partial_zobs_innovation_gate_train_only: - assert partial_zobs_innovation_gate, \ - 'partial_zobs_innovation_gate_train_only requires partial_zobs_innovation_gate=true' - if partial_zobs_innovation_gate: - assert partial_zobs_inject, 'partial_zobs_innovation_gate requires partial_zobs_inject=true' - assert partial_zobs_in_channels == 34, \ - f'partial_zobs_innovation_gate requires 34 input channels, got {partial_zobs_in_channels}' - assert innovation_gate_dim == 128, f'innovation_gate_dim is fixed at 128, got {innovation_gate_dim}' - assert innovation_gate_time_dim == 32, \ - f'innovation_gate_time_dim is fixed at 32, got {innovation_gate_time_dim}' - assert model_channels == 1536, f'partial_zobs_innovation_gate requires 1536 model channels, got {model_channels}' - assert num_blocks == 30, f'partial_zobs_innovation_gate requires 30 blocks, got {num_blocks}' - self.innovation_gate = _PartialZObsInnovationGate( - model_channels, partial_zobs_in_channels, innovation_gate_dim, - innovation_gate_time_dim, num_blocks, - ) - else: - self.innovation_gate = None - self.partial_zobs_local_attention = partial_zobs_local_attention - self.partial_zobs_local_attention_train_only = partial_zobs_local_attention_train_only - if partial_zobs_local_attention_train_only: - assert partial_zobs_local_attention, \ - 'partial_zobs_local_attention_train_only requires partial_zobs_local_attention=true' - assert sum(bool(value) for value in ( - partial_zobs_innovation_gate, - partial_zobs_local_attention, - partial_zobs_learned_attention, - )) <= 1, 'partial-zobs experimental attention/gate branches are mutually exclusive' - if partial_zobs_local_attention: - assert partial_zobs_inject, 'partial_zobs_local_attention requires partial_zobs_inject=true' - assert partial_zobs_in_channels == 34, \ - f'partial_zobs_local_attention requires 34 input channels, got {partial_zobs_in_channels}' - assert local_attention_dim == 128, \ - f'local_attention_dim is fixed at 128, got {local_attention_dim}' - assert local_attention_time_dim == 32, \ - f'local_attention_time_dim is fixed at 32, got {local_attention_time_dim}' - assert local_attention_heads == 8, \ - f'local_attention_heads is fixed at 8, got {local_attention_heads}' - assert model_channels == 1536, \ - f'partial_zobs_local_attention requires 1536 model channels, got {model_channels}' - assert num_blocks == 30, \ - f'partial_zobs_local_attention requires 30 blocks, got {num_blocks}' - self.local_obs_attention = _PartialZObsLocalAttention( - model_channels, partial_zobs_in_channels, local_attention_dim, - local_attention_time_dim, local_attention_heads, num_blocks - 1, - ) - else: - self.local_obs_attention = None - - self.partial_zobs_learned_attention = partial_zobs_learned_attention - self.partial_zobs_learned_attention_train_only = \ - partial_zobs_learned_attention_train_only - if partial_zobs_learned_attention_train_only: - assert partial_zobs_learned_attention, \ - 'partial_zobs_learned_attention_train_only requires partial_zobs_learned_attention=true' - if partial_zobs_learned_attention: - assert partial_zobs_inject, \ - 'partial_zobs_learned_attention requires partial_zobs_inject=true' - assert partial_zobs_in_channels == 34, \ - f'partial_zobs_learned_attention requires 34 input channels, got {partial_zobs_in_channels}' - assert model_channels == 1536, \ - f'partial_zobs_learned_attention requires 1536 model channels, got {model_channels}' - assert num_blocks == 30, \ - f'partial_zobs_learned_attention requires 30 blocks, got {num_blocks}' - assert learned_attention_heads == 12, \ - f'partial_zobs_learned_attention requires 12 heads, got {learned_attention_heads}' - assert tuple(learned_attention_sites) == (5, 10, 15, 20, 25), \ - f'V12LA sites are fixed at [5,10,15,20,25], got {learned_attention_sites}' - self.learned_attention = _PartialZObsLearnedAttention( - model_channels, - learned_attention_heads, - learned_attention_sites, - rope_freq, - learned_attention_null_prior, - ) - else: - self.learned_attention = None - - self.blocks = nn.ModuleList([ - ModulatedSparseTransformerCrossBlock( - model_channels, - cond_channels, - num_heads=self.num_heads, - mlp_ratio=self.mlp_ratio, - attn_mode='full', - use_checkpoint=self.use_checkpoint, - use_rope=(pe_mode == "rope"), - rope_freq=rope_freq, - share_mod=self.share_mod, - qk_rms_norm=self.qk_rms_norm, - qk_rms_norm_cross=self.qk_rms_norm_cross, - image_attn_mode=image_attn_mode, - proj_in_channels=proj_in_channels, - vae_in_channels=vae_in_channels, - ) - for _ in range(num_blocks) - ]) - - self.out_layer = sp.SparseLinear(model_channels, out_channels) - - self.initialize_weights() - self.convert_to(self.dtype) - if self.partial_zobs_innovation_gate_train_only: - self._freeze_for_partial_zobs_innovation_gate() - if self.partial_zobs_local_attention_train_only: - self._freeze_for_partial_zobs_local_attention() - if self.partial_zobs_learned_attention_train_only: - self._freeze_for_partial_zobs_learned_attention() - - @property - def device(self) -> torch.device: - """ - Return the device of the model. - """ - return next(self.parameters()).device - - def convert_to(self, dtype: torch.dtype) -> None: - """ - Convert the torso of the model to the specified dtype. - """ - self.dtype = dtype - self.blocks.apply(partial(convert_module_to, dtype=dtype)) - if self.learned_attention is not None: - self.learned_attention.apply(partial(convert_module_to, dtype=dtype)) - - def initialize_weights(self) -> None: - if self.initialization == 'vanilla': - # Initialize transformer layers: - def _basic_init(module): - if isinstance(module, nn.Linear): - torch.nn.init.xavier_uniform_(module.weight) - if module.bias is not None: - nn.init.constant_(module.bias, 0) - self.apply(_basic_init) - - # Initialize timestep embedding MLP: - nn.init.normal_(self.t_embedder.mlp[0].weight, std=0.02) - nn.init.normal_(self.t_embedder.mlp[2].weight, std=0.02) - - # Zero-out adaLN modulation layers in DiT blocks: - if self.share_mod: - nn.init.constant_(self.adaLN_modulation[-1].weight, 0) - nn.init.constant_(self.adaLN_modulation[-1].bias, 0) - else: - for block in self.blocks: - nn.init.constant_(block.adaLN_modulation[-1].weight, 0) - nn.init.constant_(block.adaLN_modulation[-1].bias, 0) - - # Zero-out output layers: - nn.init.constant_(self.out_layer.weight, 0) - nn.init.constant_(self.out_layer.bias, 0) - - elif self.initialization == 'scaled': - # Initialize transformer layers: - def _basic_init(module): - if isinstance(module, nn.Linear): - torch.nn.init.normal_(module.weight, std=np.sqrt(2.0 / (5.0 * self.model_channels))) - if module.bias is not None: - nn.init.constant_(module.bias, 0) - self.apply(_basic_init) - - # Scaled init for to_out and ffn2 - def _scaled_init(module): - if isinstance(module, nn.Linear): - torch.nn.init.normal_(module.weight, std=1.0 / np.sqrt(5 * self.num_blocks * self.model_channels)) - if module.bias is not None: - nn.init.constant_(module.bias, 0) - for block in self.blocks: - block.self_attn.to_out.apply(_scaled_init) - # Handle cross, proj, and gated_proj modes - if self.image_attn_mode in ("proj", "gated_proj"): - block.cross_attn.cross_attn_block.to_out.apply(_scaled_init) - else: - block.cross_attn.to_out.apply(_scaled_init) - block.mlp.mlp[2].apply(_scaled_init) - - # Initialize input layer to make the initial representation have variance 1 - nn.init.normal_(self.input_layer.weight, std=1.0 / np.sqrt(self.in_channels)) - nn.init.zeros_(self.input_layer.bias) - - # Initialize timestep embedding MLP: - nn.init.normal_(self.t_embedder.mlp[0].weight, std=0.02) - nn.init.normal_(self.t_embedder.mlp[2].weight, std=0.02) - - # Zero-out adaLN modulation layers in DiT blocks: - if self.share_mod: - nn.init.constant_(self.adaLN_modulation[-1].weight, 0) - nn.init.constant_(self.adaLN_modulation[-1].bias, 0) - else: - for block in self.blocks: - nn.init.constant_(block.adaLN_modulation[-1].weight, 0) - nn.init.constant_(block.adaLN_modulation[-1].bias, 0) - - # Zero-out output layers: - nn.init.constant_(self.out_layer.weight, 0) - nn.init.constant_(self.out_layer.bias, 0) - - # Trained-bridge: zero-init the z_obs injector LAST (after _basic_init), so step-0 injection == 0. - if getattr(self, 'zobs_injector', None) is not None: - nn.init.zeros_(self.zobs_injector.weight) - nn.init.zeros_(self.zobs_injector.bias) - if getattr(self, 'partial_zobs_injector', None) is not None: - nn.init.zeros_(self.partial_zobs_injector.weight) - if getattr(self, 'innovation_gate', None) is not None: - nn.init.zeros_(self.innovation_gate.alpha) - if getattr(self, 'local_obs_attention', None) is not None: - # V12-A freezes vanilla Xavier initialization independently of the - # scaled V11 backbone initialization mode. Only alpha starts at zero. - for module in self.local_obs_attention.modules(): - if isinstance(module, nn.Linear): - nn.init.xavier_uniform_(module.weight) - if module.bias is not None: - nn.init.zeros_(module.bias) - nn.init.xavier_uniform_(self.local_obs_attention.relative_offset) - nn.init.xavier_uniform_(self.local_obs_attention.null_key) - nn.init.zeros_(self.local_obs_attention.alpha) - if getattr(self, 'learned_attention', None) is not None: - self.learned_attention.reset_parameters() - - def _freeze_for_partial_zobs_innovation_gate(self) -> None: - """Make the V12 gate the sole trainable branch before an optimizer is built.""" - assert self.innovation_gate is not None - for parameter in self.parameters(): - parameter.requires_grad_(False) - for parameter in self.innovation_gate.parameters(): - parameter.requires_grad_(True) - trainable = [name for name, parameter in self.named_parameters() if parameter.requires_grad] - expected = [f'innovation_gate.{name}' for name in self.innovation_gate.state_dict()] - assert trainable == expected, (trainable, expected) - assert sum(parameter.numel() for parameter in self.parameters() if parameter.requires_grad) == 473663 - - def _freeze_for_partial_zobs_local_attention(self) -> None: - """Make the V12-A local attention the sole trainable warm-start branch.""" - assert self.local_obs_attention is not None - for parameter in self.parameters(): - parameter.requires_grad_(False) - for parameter in self.local_obs_attention.parameters(): - parameter.requires_grad_(True) - trainable = [name for name, parameter in self.named_parameters() if parameter.requires_grad] - expected = [f'local_obs_attention.{name}' for name in self.local_obs_attention.state_dict()] - assert trainable == expected, (trainable, expected) - assert len(trainable) == 14 - assert sum(parameter.numel() for parameter in self.parameters() if parameter.requires_grad) == 487485 - - def _freeze_for_partial_zobs_learned_attention(self) -> None: - """Make the V12LA module the sole trainable warm-start branch.""" - assert self.learned_attention is not None - for parameter in self.parameters(): - parameter.requires_grad_(False) - for parameter in self.learned_attention.parameters(): - parameter.requires_grad_(True) - trainable = [name for name, parameter in self.named_parameters() if parameter.requires_grad] - expected = [ - f'learned_attention.{name}' - for name, _ in self.learned_attention.named_parameters() - ] - assert trainable == expected, (trainable, expected) - assert len(trainable) == 13, trainable - - def _partial_zobs_message( - self, - x: sp.SparseTensor, - partial_z_obs: sp.SparseTensor, - partial_z_obs_keep_mask: Optional[torch.Tensor], - ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - """Project a variable-support observation and join its messages onto ``x`` rows. - - The return value stays variable-length: projected messages plus their exact full-row - indices. Neither observation content nor a zero-filled model-channel tensor is - materialized on full support. ``torch.unique(..., dim=0)`` makes the coordinate join - exact and rejects duplicate or non-subset partial coordinates instead of assuming a - row order. - """ - adapter = self.partial_zobs_injector - assert adapter is not None - assert partial_z_obs.feats.ndim == 2 and \ - partial_z_obs.feats.shape[-1] == self.partial_zobs_in_channels, \ - f'partial_z_obs must be [N,{self.partial_zobs_in_channels}], got {tuple(partial_z_obs.feats.shape)}' - assert partial_z_obs.coords.ndim == 2 and partial_z_obs.coords.shape[1] == x.coords.shape[1], \ - f'partial_z_obs coordinate rank mismatch: {tuple(partial_z_obs.coords.shape)} vs {tuple(x.coords.shape)}' - - active = torch.ones(partial_z_obs.feats.shape[0], dtype=torch.bool, device=partial_z_obs.feats.device) - if partial_z_obs_keep_mask is not None: - assert partial_z_obs_keep_mask.ndim == 1 and partial_z_obs_keep_mask.shape[0] == x.shape[0], \ - f'partial_z_obs_keep_mask must be [B={x.shape[0]}], got {tuple(partial_z_obs_keep_mask.shape)}' - keep = partial_z_obs_keep_mask.to(device=partial_z_obs.feats.device, dtype=torch.bool) - active = keep[partial_z_obs.coords[:, 0].long()] - - active_coords = partial_z_obs.coords[active] - active_feats = partial_z_obs.feats[active] - # F.linear on a zero-row tensor is intentional: it keeps the adapter parameters in the - # graph on all-drop batches, so DDP sees a zero (rather than None) gradient. - active_messages = F.linear(active_feats, adapter.weight, None) - - full_rows = torch.empty(0, dtype=torch.long, device=x.feats.device) - if active_coords.shape[0] != 0: - all_coords = torch.cat([x.coords, active_coords], dim=0) - _, inverse = torch.unique(all_coords, dim=0, return_inverse=True) - x_ids = inverse[:x.coords.shape[0]] - partial_ids = inverse[x.coords.shape[0]:] - assert torch.unique(x_ids).numel() == x_ids.numel(), 'x contains duplicate coordinates' - assert torch.unique(partial_ids).numel() == partial_ids.numel(), 'partial_z_obs contains duplicate coordinates' - ordered_x_ids, ordered_x_rows = torch.sort(x_ids) - positions = torch.searchsorted(ordered_x_ids, partial_ids) - assert (positions < ordered_x_ids.numel()).all() and \ - torch.equal(ordered_x_ids[positions], partial_ids), \ - 'partial_z_obs coordinates must be an exact subset of current x coordinates' - full_rows = ordered_x_rows[positions] - - return active_messages, full_rows, active_feats, active_coords[:, 0].long() - - def forward( - self, - x: sp.SparseTensor, - t: torch.Tensor, - cond: Union[torch.Tensor, List[torch.Tensor], Dict[str, Union[torch.Tensor, sp.SparseTensor]], Tuple], - concat_cond: Optional[sp.SparseTensor] = None, - z_obs: Optional[sp.SparseTensor] = None, - partial_z_obs: Optional[sp.SparseTensor] = None, - partial_z_obs_keep_mask: Optional[torch.Tensor] = None, - **kwargs - ) -> sp.SparseTensor: - """ - Forward pass. - - Args: - x: SparseTensor input - t: Timestep tensor [B] - cond: Conditioning tensor. For "cross" mode: list of tensors or tensor. - For "proj" mode: dict {'global': global_cond, 'proj': proj_cond} - or tuple of (global_cond, proj_cond) - concat_cond: Optional concatenation condition - z_obs: Optional 33ch [z_obs(32)|vis(1)] sparse tensor, same coords as x. Feeds ONLY - the zobs_injector (v6 patchcross path: concat_cond is the plain 32ch shape - latent and z_obs never enters the input_layer). When None, the legacy 97ch - concat slice below applies — shipped/v4 behavior is bit-identical. - partial_z_obs: Optional variable-length v8 observation - [normalized_z_obs(32)|coverage(1)|boundary_distance(1)]. - Its coordinates must be a subset of x; it is consumed only by the dedicated, - bias-free partial adapter and is never converted to full-support content. - - Returns: - SparseTensor output - """ - if concat_cond is not None: - x = sp.sparse_cat([x, concat_cond], dim=-1) - - h = self.input_layer(x) - h = manual_cast(h, self.dtype) - - # Trained-bridge: compute the per-block injection vector ONCE (shared zero-init linear, added each block). - # Sources, in priority order: separate z_obs kwarg (v6), else the trailing channels of a - # 97-ch [x_t|shape|z_obs|vis] input (v4). z_obs*vis==0 at vis=0 in both. - zobs_inject = None - _injector = getattr(self, 'zobs_injector', None) - if z_obs is not None: - # Fail fast: without an injector the kwarg would be a silent no-op, and in the 64ch - # patchcross architecture the injector is the ONLY z_obs consumption path. - assert _injector is not None, \ - 'z_obs kwarg passed but the denoiser was built without zobs_inject=true' - # Full coords equality is asserted at data-load; a row-count check catches gross mismatch. - assert z_obs.feats.shape[-1] == 33 and z_obs.feats.shape[0] == x.feats.shape[0], \ - f'z_obs kwarg mismatch: {tuple(z_obs.feats.shape)} vs x rows {x.feats.shape[0]}' - _zo = z_obs.feats[:, :32] - _vis = z_obs.feats[:, 32:33] - zobs_inject = manual_cast(self.zobs_injector(z_obs.replace(torch.cat([_zo * _vis, _vis], dim=-1))), self.dtype) - if not getattr(self, '_zobs_kwarg_logged', False): # in-log proof hook (once per process) - print('[zobs_injector] consuming separate z_obs kwarg (patchcross path)', flush=True) - self._zobs_kwarg_logged = True - elif _injector is not None and x.feats.shape[-1] >= 97: - _zo = x.feats[:, 64:96] - _vis = x.feats[:, 96:97] - zobs_inject = manual_cast(self.zobs_injector(x.replace(torch.cat([_zo * _vis, _vis], dim=-1))), self.dtype) - elif _injector is not None and self.legacy_zobs_zero_input: - # 64ch v8 shell with a legacy injector: preserve the legacy all-zero z_obs mode, - # including its learned per-block bias, without materialising observation content. - # Evaluate one zero row and broadcast the result: this is exactly W*0+b, keeps a - # zero (non-None) weight gradient under DDP, and avoids an N-row 33ch GEMM. - zero_output = F.linear( - x.feats.new_zeros((1, _injector.in_features)), - _injector.weight, - _injector.bias, - ) - zobs_inject = manual_cast( - x.replace(zero_output.expand(x.feats.shape[0], -1)), self.dtype - ) - elif _injector is not None and not getattr(self, '_zobs_none_warned', False): - # Injector present but NO z_obs source (64ch input, kwarg omitted): the trained - # per-block bias is NOT applied — a third mode distinct from "z_obs dropped" (an - # all-zero 33ch tensor still adds the bias each block). A nozobs arm on a patchcross - # model must pass the zero tensor, never omit the kwarg. - print('[zobs_injector] WARNING: injector present but no z_obs source (kwarg omitted, ' - 'input < 97ch) — per-block bias NOT applied; pass an all-zero 33ch z_obs for a ' - 'true nozobs arm', flush=True) - self._zobs_none_warned = True - - partial_zobs_inject = None - partial_zobs_rows = None - partial_zobs_feats = None - partial_zobs_batch = None - if partial_z_obs is not None: - assert self.partial_zobs_injector is not None, \ - 'partial_z_obs passed but denoiser was built without partial_zobs_inject=true' - partial_zobs_inject, partial_zobs_rows, partial_zobs_feats, partial_zobs_batch = self._partial_zobs_message( - x, partial_z_obs, partial_z_obs_keep_mask - ) - partial_zobs_inject = manual_cast(partial_zobs_inject, self.dtype) - elif self.partial_zobs_injector is not None: - raise AssertionError('partial_zobs_inject=true requires a separate partial_z_obs tensor') - - if zobs_inject is not None and self.zobs_inject_scale != 1.0: - zobs_inject = zobs_inject.replace(zobs_inject.feats * self.zobs_inject_scale) - if not getattr(self, '_zobs_inject_scale_logged', False): - print(f'[zobs_injector] output scale={self.zobs_inject_scale:g}', flush=True) - self._zobs_inject_scale_logged = True - - # Keep the pre-AdaLN embedding for the innovation gate; blocks retain their existing - # modulation input exactly. - t_base = self.t_embedder(t) - t_emb = t_base - if self.share_mod: - t_emb = self.adaLN_modulation(t_emb) - t_emb = manual_cast(t_emb, self.dtype) - - innovation_time = None - innovation_y = None - innovation_gate_y = None - if self.innovation_gate is not None and partial_zobs_feats is not None: - # Observation-only terms are shared by all blocks. Empty rows deliberately pass - # through these linears too, so an all-drop batch keeps the full gate graph live. - gate_dtype = self.innovation_gate.embed_y.weight.dtype - innovation_gate_y = partial_zobs_feats.to(dtype=gate_dtype) - innovation_y = self.innovation_gate.embed_y(innovation_gate_y) - innovation_time = self.innovation_gate.time_proj(t_base) - - if self.pe_mode == "ape": - pe = self.pos_embedder(h.coords[:, 1:]) - h = h + manual_cast(pe, self.dtype) - - # Handle different conditioning modes - if self.image_attn_mode == 'proj': - if isinstance(cond, dict): - global_cond = cond['global'] - proj_cond = cond['proj'] - else: - global_cond, proj_cond = cond - if isinstance(global_cond, list): - global_cond = sp.VarLenTensor.from_tensor_list(global_cond) - global_cond = manual_cast(global_cond, self.dtype) - proj_cond = manual_cast(proj_cond, self.dtype) - cond = (global_cond, proj_cond) - elif self.image_attn_mode == 'gated_proj': - global_cond = cond['global'] - if isinstance(global_cond, list): - global_cond = sp.VarLenTensor.from_tensor_list(global_cond) - global_cond = manual_cast(global_cond, self.dtype) - proj_semantic = manual_cast(cond['proj_semantic'], self.dtype) - proj_color = manual_cast(cond['proj_color'], self.dtype) - cond = {'global': global_cond, 'proj_semantic': proj_semantic, 'proj_color': proj_color} - else: - if isinstance(cond, list): - cond = sp.VarLenTensor.from_tensor_list(cond) - cond = manual_cast(cond, self.dtype) - - # Inference-only probe hook. When the attribute is absent (the default, and the - # only state any training path produces) this is None and the loop below is - # byte-identical to the unhooked path: no constructor arg, no config key, and no - # entry in the inference ABI contract. A probe sets it to a length-num_blocks - # sequence to rescale the shared partial message per block. - _partial_scales = getattr(self, 'partial_zobs_layer_scales', None) - # Optional evaluator-only callable: ``fn(gate_value, active_batch, block_index)``. - # It is an instance attribute rather than a constructor/config argument, so no model - # state or training ABI changes; callers can implement neutral, constant, or shuffled - # gate-value probes deterministically. - _gate_value_transform = getattr(self, 'partial_zobs_innovation_gate_value_transform', None) - _innovation_ratio_stats = [] if self.innovation_gate is not None else None - _innovation_gate_mean_stats = [] if self.innovation_gate is not None else None - _innovation_observed_rows = 0 if partial_zobs_rows is None else int(partial_zobs_rows.numel()) - _local_attention = getattr(self, 'local_obs_attention', None) - _local_topology = None - _local_full_batch = None - _local_ratio_base_stats = [] if _local_attention is not None else None - _local_ratio_hidden_stats = [] if _local_attention is not None else None - _local_null_stats = [] if _local_attention is not None else None - _local_entropy_stats = [] if _local_attention is not None else None - _local_effective_neighbour_stats = [] if _local_attention is not None else None - _local_observed_residual_stats = [] if _local_attention is not None else None - _local_unobserved_residual_stats = [] if _local_attention is not None else None - _local_query_is_observed = None - _local_neighbourhood_histogram = None - if _local_attention is not None: - assert partial_zobs_rows is not None and partial_zobs_feats is not None - _local_topology = _build_two_hop_observed_topology(x.coords, partial_zobs_rows) - _local_full_batch = x.coords[:, 0].long() - _active_observed_mask = torch.zeros( - x.coords.shape[0], dtype=torch.bool, device=x.coords.device, - ) - _active_observed_mask[partial_zobs_rows] = True - _local_query_is_observed = _active_observed_mask[ - _local_topology['query_rows'] - ] - _local_neighbourhood_histogram = torch.bincount( - _local_topology['neighbour_counts'], minlength=26, - )[:26] - _local_neighbourhood_histogram[0] = ( - x.coords.shape[0] - _local_topology['query_rows'].numel() - ) - assert int(_local_neighbourhood_histogram.sum()) == x.coords.shape[0] - assert int(_local_neighbourhood_histogram[1:].sum()) == \ - _local_topology['query_rows'].numel() - _learned_attention = getattr(self, 'learned_attention', None) - _learned_cache = None - _learned_stats = [] if _learned_attention is not None else None - if _learned_attention is not None: - assert partial_zobs_inject is not None and partial_zobs_rows is not None - assert _partial_scales is None, \ - 'partial_zobs_layer_scales is not defined for the five-site V12LA operator' - _learned_cache = _learned_attention.build_cache( - x.coords, - x.layout, - partial_zobs_inject, - partial_zobs_rows, - ) - partial_zobs_inject_position = getattr( - self, 'partial_zobs_inject_position', 'post_block', - ) - for _block_index, block in enumerate(self.blocks): - if partial_zobs_inject is not None and partial_zobs_inject_position == 'pre_mlp': - _pre_mlp_scale = 1.0 if _partial_scales is None else float( - _partial_scales[_block_index] - ) - h = block( - h, - t_emb, - cond, - partial_zobs_rows=partial_zobs_rows, - partial_zobs_inject=partial_zobs_inject, - partial_zobs_inject_scale=_pre_mlp_scale, - ) - else: - h = block(h, t_emb, cond) - if zobs_inject is not None: - h = h + zobs_inject - _innovation_ratio = None - _innovation_gate_mean = None - _local_delta = None - _local_stats = None - if _local_attention is not None and _block_index < _local_attention.alpha.numel(): - assert _local_topology is not None and _local_full_batch is not None - _local_delta, _local_stats = _local_attention( - h.feats, - _local_full_batch, - t_base, - partial_zobs_feats, - partial_zobs_rows, - _local_topology, - _block_index, - ) - _local_delta = manual_cast(_local_delta, self.dtype) - _local_ratio_base = _local_delta.norm() / partial_zobs_inject.norm().clamp_min( - torch.finfo(partial_zobs_inject.dtype).eps - ) - _local_hidden = h.feats[_local_topology['query_rows']] - _local_ratio_hidden = _local_delta.norm() / _local_hidden.norm().clamp_min( - torch.finfo(_local_hidden.dtype).eps - ) - assert _local_query_is_observed is not None - _local_observed_residual = _local_delta[ - _local_query_is_observed - ].float().norm() - _local_unobserved_residual = _local_delta[ - ~_local_query_is_observed - ].float().norm() - if _learned_attention is not None and _block_index in _learned_attention.site_to_slot: - assert _learned_cache is not None - _learned_inject, _site_stats = _learned_attention( - h.feats, - h.coords, - h.layout, - _learned_cache, - _block_index, - ) - h = h.replace(h.feats + manual_cast(_learned_inject, self.dtype)) - _learned_stats.append({ - name: value.detach() for name, value in _site_stats.items() - }) - if partial_zobs_inject is not None and _learned_attention is None and \ - partial_zobs_inject_position == 'post_block': - layer_inject = partial_zobs_inject - if self.innovation_gate is not None: - gate = self.innovation_gate - # Every operation intentionally accepts an empty leading dimension. This - # keeps all innovation parameters in the all-drop DDP graph. - assert innovation_gate_y is not None and innovation_y is not None and innovation_time is not None - t_obs = innovation_time[partial_zobs_batch] - h_obs = h.feats[partial_zobs_rows].to(dtype=gate.embed_y.weight.dtype) - x_embed = gate.embed_x(torch.cat([h_obs, t_obs], dim=-1)) - innovation = innovation_y - x_embed - gate_input = torch.cat([ - innovation, - innovation_gate_y[:, 32:34], - t_obs, - ], dim=-1) - gate_value = torch.sigmoid(gate.mlp(gate_input)) - if _gate_value_transform is not None: - gate_value = _gate_value_transform( - gate_value, partial_zobs_batch, _block_index - ) - assert gate_value.shape == (innovation.shape[0], 1), \ - f'gate value transform must return [Nobs,1], got {tuple(gate_value.shape)}' - delta = gate.proj(innovation) * gate_value - delta = delta * gate.alpha[_block_index].to(dtype=delta.dtype) - if _block_index == len(self.blocks) - 1: - # Keep the frozen V11 base carrier at the terminal site, - # but never add a newly learned observed-only shortcut: - # no attention block remains to transport/smooth it into U. - delta = delta * 0 - # Detached telemetry for the training log. This is before the legacy - # layer-scale hook and therefore measures the learned branch itself. - _innovation_ratio = delta.norm() / partial_zobs_inject.norm().clamp_min( - torch.finfo(partial_zobs_inject.dtype).eps - ) - _innovation_gate_mean = ( - gate_value.mean() if gate_value.numel() != 0 else delta.new_zeros(()) - ) - layer_inject = partial_zobs_inject + manual_cast(delta, self.dtype) - if partial_zobs_rows.numel() != 0: - if _partial_scales is None: - h = h.replace(h.feats.index_add( - 0, partial_zobs_rows, layer_inject - )) - else: - h = h.replace(h.feats.index_add( - 0, partial_zobs_rows, layer_inject, - alpha=float(_partial_scales[_block_index]), - )) - else: - # Preserve every empty carrier/gate GEMM in the graph on all-drop batches. - h = h.replace(h.feats + layer_inject.sum() * 0) - if _local_delta is not None: - _local_scale = 1.0 if _partial_scales is None else float( - _partial_scales[_block_index] - ) - if _local_topology['query_rows'].numel() != 0: - h = h.replace(h.feats.index_add( - 0, _local_topology['query_rows'], _local_delta, - alpha=_local_scale, - )) - else: - # The module returns an exact-zero parameter edge for all 14 - # trainable tensors, so DDP observes finite numerical-zero grads. - h = h.replace(h.feats + _local_delta.sum() * _local_scale) - if self.innovation_gate is not None: - if _innovation_ratio is None: - _innovation_ratio = h.feats.new_zeros(()) - _innovation_gate_mean = h.feats.new_zeros(()) - _innovation_ratio_stats.append(_innovation_ratio.detach()) - _innovation_gate_mean_stats.append(_innovation_gate_mean.detach()) - if _local_attention is not None and _block_index < _local_attention.alpha.numel(): - assert _local_stats is not None - _local_ratio_base_stats.append(_local_ratio_base.detach()) - _local_ratio_hidden_stats.append(_local_ratio_hidden.detach()) - _local_null_stats.append(_local_stats['null_attention_mean'].detach()) - _local_entropy_stats.append(_local_stats['real_neighbour_entropy_mean'].detach()) - _local_effective_neighbour_stats.append( - _local_stats['effective_real_neighbour_count_mean'].detach() - ) - _local_observed_residual_stats.append( - _local_observed_residual.detach() - ) - _local_unobserved_residual_stats.append( - _local_unobserved_residual.detach() - ) - - if self.innovation_gate is not None: - # Instance-only telemetry: deliberately neither buffer nor state-dict entry. - self.partial_zobs_innovation_gate_last_stats = { - 'delta_to_base_norm_ratio': torch.stack(_innovation_ratio_stats), - 'gate_mean': torch.stack(_innovation_gate_mean_stats), - 'observed_rows': _innovation_observed_rows, - } - if _local_attention is not None: - # Block 29 is structurally absent: every telemetry vector has length 29. - assert _local_query_is_observed is not None - assert _local_neighbourhood_histogram is not None - _full_support_rows = int(x.coords.shape[0]) - _active_queries = int(_local_topology['query_rows'].numel()) - _empty_queries = _full_support_rows - _active_queries - _singleton_queries = int( - (_local_topology['neighbour_counts'] == 1).sum() - ) - self.partial_zobs_local_attention_last_stats = { - 'delta_to_base_norm_ratio': torch.stack(_local_ratio_base_stats), - 'delta_to_hidden_norm_ratio': torch.stack(_local_ratio_hidden_stats), - 'null_attention_mean': torch.stack(_local_null_stats), - 'real_neighbour_entropy_mean': torch.stack(_local_entropy_stats), - 'effective_real_neighbour_count_mean': torch.stack( - _local_effective_neighbour_stats - ), - 'observed_residual_l2_norm': torch.stack( - _local_observed_residual_stats - ), - 'unobserved_residual_l2_norm': torch.stack( - _local_unobserved_residual_stats - ), - 'observed_rows': _innovation_observed_rows, - 'full_support_rows': _full_support_rows, - 'active_queries': _active_queries, - 'empty_queries': _empty_queries, - 'singleton_queries': _singleton_queries, - 'observed_active_queries': int(_local_query_is_observed.sum()), - 'unobserved_active_queries': int((~_local_query_is_observed).sum()), - 'active_query_fraction': _active_queries / max(_full_support_rows, 1), - 'empty_query_fraction': _empty_queries / max(_full_support_rows, 1), - 'singleton_query_fraction': _singleton_queries / max(_active_queries, 1), - 'neighbourhood_size_histogram': _local_neighbourhood_histogram.detach(), - 'packed_pairs': int(_local_topology['source_indices'].numel()), - } - if _learned_attention is not None: - assert _learned_cache is not None - assert len(_learned_stats) == len(_learned_attention.sites) - self.partial_zobs_learned_attention_last_stats = { - name: torch.stack([site[name] for site in _learned_stats]) - for name in _learned_stats[0] - } - self.partial_zobs_learned_attention_last_stats.update({ - 'sites': _learned_attention.sites, - 'observed_rows': int(partial_zobs_rows.numel()), - 'full_support_rows': int(x.coords.shape[0]), - 'observed_counts': _learned_cache['observed_counts'], - }) - - h = manual_cast(h, x.dtype) - h = h.replace(F.layer_norm(h.feats, h.feats.shape[-1:])) - h = self.out_layer(h) - return h - - -class ElasticSLatFlowModel(SparseTransformerElasticMixin, SLatFlowModel): - """ - SLat Flow Model with elastic memory management. - Used for training with low VRAM. - """ - pass +from typing import * +from functools import partial +import torch +import torch.nn as nn +import torch.nn.functional as F +import numpy as np +from ..modules.utils import convert_module_to, manual_cast, str_to_dtype +from ..modules.transformer import AbsolutePositionEmbedder +from ..modules import sparse as sp +from ..modules.sparse.transformer import ModulatedSparseTransformerCrossBlock +from .sparse_structure_flow import TimestepEmbedder + + +class _PartialZObsInnovationGate(nn.Module): + """Residual, per-observed-token correction for the partial z_obs carrier.""" + + def __init__(self, model_channels: int, message_channels: int, dim: int, time_dim: int, num_blocks: int): + super().__init__() + self.alpha = nn.Parameter(torch.zeros(num_blocks)) + self.embed_y = nn.Linear(message_channels, dim, bias=True) + self.time_proj = nn.Linear(model_channels, time_dim, bias=True) + self.embed_x = nn.Linear(model_channels + time_dim, dim, bias=True) + self.mlp = nn.Sequential( + nn.Linear(dim + 2 + time_dim, dim, bias=True), + nn.SiLU(), + nn.Linear(dim, 1, bias=True), + ) + self.proj = nn.Linear(dim, model_channels, bias=True) + + +_LOCAL_ONE_HOP_OFFSETS = ( + (-1, 0, 0), (0, -1, 0), (0, 0, -1), + (0, 0, 1), (0, 1, 0), (1, 0, 0), +) +_LOCAL_TWO_HOP_OFFSETS = tuple(sorted( + (dx, dy, dz) + for dx in range(-2, 3) + for dy in range(-2, 3) + for dz in range(-2, 3) + if abs(dx) + abs(dy) + abs(dz) <= 2 +)) +assert len(_LOCAL_TWO_HOP_OFFSETS) == 25 + + +def _exact_coordinate_rows( + full_coords: torch.Tensor, + candidate_coords: torch.Tensor, +) -> Tuple[torch.Tensor, torch.Tensor]: + """Return exact full-support rows and a validity mask for candidate coords.""" + assert full_coords.ndim == candidate_coords.ndim == 2 + assert full_coords.shape[1] == candidate_coords.shape[1] + assert full_coords.shape[0] > 0, 'local attention requires nonempty full sparse support' + all_coords = torch.cat([full_coords, candidate_coords], dim=0) + _, inverse = torch.unique(all_coords, dim=0, return_inverse=True) + full_ids = inverse[:full_coords.shape[0]] + candidate_ids = inverse[full_coords.shape[0]:] + assert torch.unique(full_ids).numel() == full_ids.numel(), \ + 'x contains duplicate coordinates' + ordered_ids, ordered_rows = torch.sort(full_ids) + positions = torch.searchsorted(ordered_ids, candidate_ids) + in_range = positions < ordered_ids.numel() + safe_positions = positions.clamp(max=ordered_ids.numel() - 1) + valid = in_range & (ordered_ids[safe_positions] == candidate_ids) + rows = torch.where(valid, ordered_rows[safe_positions], -torch.ones_like(safe_positions)) + return rows, valid + + +def _build_two_hop_observed_topology( + full_coords: torch.Tensor, + observed_rows: torch.Tensor, +) -> Dict[str, torch.Tensor]: + """Pack exact batch-local G(S) distance<=2 observed source/query pairs.""" + assert full_coords.ndim == 2 and full_coords.shape[1] == 4, \ + f'local attention requires [N,4] integer coords, got {tuple(full_coords.shape)}' + assert not torch.is_floating_point(full_coords), 'local attention coordinates must be integer' + assert observed_rows.ndim == 1 and observed_rows.dtype == torch.long + # Validate full-coordinate uniqueness even for an all-absent batch. + _exact_coordinate_rows(full_coords, full_coords.new_empty((0, full_coords.shape[1]))) + if observed_rows.numel() == 0: + empty = torch.empty(0, dtype=torch.long, device=full_coords.device) + return { + 'query_rows': empty, + 'source_indices': empty, + 'offset_indices': empty, + 'segment_indices': empty, + 'neighbour_counts': empty, + } + assert observed_rows.min() >= 0 and observed_rows.max() < full_coords.shape[0], \ + 'observed row lies outside full sparse support' + assert torch.unique(observed_rows).numel() == observed_rows.numel(), \ + 'observed rows must be unique' + + device = full_coords.device + one_hop = torch.tensor( + _LOCAL_ONE_HOP_OFFSETS, dtype=full_coords.dtype, device=device, + ) + source_coords = full_coords[observed_rows] + candidate_hop1 = source_coords[:, None, :].expand(-1, len(_LOCAL_ONE_HOP_OFFSETS), -1).clone() + candidate_hop1[:, :, 1:] += one_hop[None] + hop1_rows, hop1_valid = _exact_coordinate_rows( + full_coords, candidate_hop1.reshape(-1, full_coords.shape[1]), + ) + source_ids_hop1 = torch.arange( + observed_rows.numel(), device=device, dtype=torch.long, + )[:, None].expand(-1, len(_LOCAL_ONE_HOP_OFFSETS)).reshape(-1) + hop1_rows = hop1_rows[hop1_valid] + source_ids_hop1 = source_ids_hop1[hop1_valid] + + if hop1_rows.numel() != 0: + candidate_hop2 = full_coords[hop1_rows][:, None, :].expand( + -1, len(_LOCAL_ONE_HOP_OFFSETS), -1, + ).clone() + candidate_hop2[:, :, 1:] += one_hop[None] + hop2_rows, hop2_valid = _exact_coordinate_rows( + full_coords, candidate_hop2.reshape(-1, full_coords.shape[1]), + ) + source_ids_hop2 = source_ids_hop1[:, None].expand( + -1, len(_LOCAL_ONE_HOP_OFFSETS), + ).reshape(-1) + hop2_rows = hop2_rows[hop2_valid] + source_ids_hop2 = source_ids_hop2[hop2_valid] + else: + hop2_rows = torch.empty(0, dtype=torch.long, device=device) + source_ids_hop2 = torch.empty(0, dtype=torch.long, device=device) + + source_ids = torch.cat([ + torch.arange(observed_rows.numel(), device=device, dtype=torch.long), + source_ids_hop1, + source_ids_hop2, + ]) + query_rows = torch.cat([observed_rows, hop1_rows, hop2_rows]) + # Multiple graph paths to the same query/source pair are one neighbourhood entry. + unique_pairs = torch.unique(torch.stack([query_rows, source_ids], dim=1), dim=0) + query_rows = unique_pairs[:, 0] + source_ids = unique_pairs[:, 1] + query_coords = full_coords[query_rows] + selected_source_coords = source_coords[source_ids] + assert torch.equal(query_coords[:, 0], selected_source_coords[:, 0]), \ + 'local attention constructed a cross-sample edge' + relative = query_coords[:, 1:] - selected_source_coords[:, 1:] + offset_table = torch.tensor( + _LOCAL_TWO_HOP_OFFSETS, dtype=full_coords.dtype, device=device, + ) + offset_matches = (relative[:, None, :] == offset_table[None, :, :]).all(dim=-1) + assert torch.equal(offset_matches.sum(dim=1), torch.ones_like(offset_matches[:, 0], dtype=torch.long)), \ + 'local attention produced an invalid two-hop offset' + offset_indices = offset_matches.to(dtype=torch.long).argmax(dim=1) + + # Canonical per-query real-KV order is lexicographic relative offset, independent + # of either full-support or observed-message row ordering. + order = torch.argsort(query_rows * len(_LOCAL_TWO_HOP_OFFSETS) + offset_indices) + query_rows = query_rows[order] + source_ids = source_ids[order] + offset_indices = offset_indices[order] + active_query_rows, segment_indices, counts = torch.unique_consecutive( + query_rows, return_inverse=True, return_counts=True, + ) + assert counts.max() <= len(_LOCAL_TWO_HOP_OFFSETS), \ + 'local observed neighbourhood exceeds frozen 25-source cap' + return { + 'query_rows': active_query_rows, + 'source_indices': source_ids, + 'offset_indices': offset_indices, + 'segment_indices': segment_indices, + 'neighbour_counts': counts, + } + + +class _PartialZObsLocalAttention(nn.Module): + """Shared packed two-hop observed-innovation attention for V12-A.""" + + def __init__( + self, + model_channels: int, + message_channels: int, + dim: int, + time_dim: int, + num_heads: int, + num_active_blocks: int, + ): + super().__init__() + assert dim == 128 and time_dim == 32 and num_heads == 8 + assert dim % num_heads == 0 + self.dim = dim + self.num_heads = num_heads + self.head_dim = dim // num_heads + self.alpha = nn.Parameter(torch.zeros(num_active_blocks)) + self.time_proj = nn.Linear(model_channels, time_dim, bias=True) + self.embed_x = nn.Linear(model_channels + time_dim, dim, bias=True) + self.embed_y = nn.Linear(message_channels, dim, bias=True) + self.key_proj = nn.Linear(dim, dim, bias=True) + self.value_proj = nn.Linear(dim, dim, bias=True) + self.relative_offset = nn.Parameter(torch.empty(len(_LOCAL_TWO_HOP_OFFSETS), dim)) + self.null_key = nn.Parameter(torch.empty(num_heads, self.head_dim)) + self.out_proj = nn.Linear(dim, model_channels, bias=False) + + @staticmethod + def _rms_norm(value: torch.Tensor) -> torch.Tensor: + return value * torch.rsqrt(value.square().mean(dim=-1, keepdim=True) + 1e-6) + + def forward( + self, + hidden: torch.Tensor, + full_batch: torch.Tensor, + t_base: torch.Tensor, + message: torch.Tensor, + observed_rows: torch.Tensor, + topology: Dict[str, torch.Tensor], + block_index: int, + ) -> Tuple[torch.Tensor, Dict[str, torch.Tensor]]: + assert 0 <= block_index < self.alpha.numel(), 'terminal block has no V12-A residual' + query_rows = topology['query_rows'] + source_indices = topology['source_indices'] + offset_indices = topology['offset_indices'] + segment_indices = topology['segment_indices'] + counts = topology['neighbour_counts'] + assert full_batch.ndim == 1 and full_batch.shape[0] == hidden.shape[0] + assert full_batch.numel() == 0 or ( + full_batch.min() >= 0 and full_batch.max() < t_base.shape[0] + ), 'full sparse batch coordinate is outside timestep batch' + assert message.shape == (observed_rows.numel(), 34) + + if query_rows.numel() == 0: + live_zero = sum(parameter.sum() for parameter in self.parameters()) * 0 + empty_delta = hidden.new_empty((0, hidden.shape[1])) + live_zero + zero = hidden.new_zeros(()) + return empty_delta, { + 'null_attention_mean': zero, + 'real_neighbour_entropy_mean': zero, + 'effective_real_neighbour_count_mean': zero, + 'active_queries': zero, + 'singleton_queries': zero, + 'packed_pairs': zero, + } + + branch_dtype = self.embed_x.weight.dtype + hidden_branch = hidden.to(dtype=branch_dtype) + time = self.time_proj(t_base.to(dtype=branch_dtype)) + x_embed = self.embed_x(torch.cat([ + F.layer_norm(hidden_branch, hidden_branch.shape[-1:], eps=1e-6), + time[full_batch], + ], dim=-1)) + innovation = self.embed_y(message.to(dtype=branch_dtype)) - x_embed[observed_rows] + query = self._rms_norm( + x_embed[query_rows].reshape(-1, self.num_heads, self.head_dim) + ) + real_key = self.key_proj(innovation)[source_indices] + real_key = real_key + self.relative_offset[offset_indices] + real_key = self._rms_norm(real_key.reshape(-1, self.num_heads, self.head_dim)) + real_value = self.value_proj(innovation)[source_indices].reshape( + -1, self.num_heads, self.head_dim, + ) + real_logits = (query[segment_indices] * real_key).sum(dim=-1) / np.sqrt(self.head_dim) + null_key = self._rms_norm(self.null_key) + null_logits = (query * null_key[None]).sum(dim=-1) / np.sqrt(self.head_dim) + + segment_heads = segment_indices[:, None].expand(-1, self.num_heads) + maxima = null_logits.clone() + maxima.scatter_reduce_(0, segment_heads, real_logits, reduce='amax', include_self=True) + real_exp = torch.exp(real_logits - maxima[segment_indices]) + null_exp = torch.exp(null_logits - maxima) + real_sum = torch.zeros_like(null_exp).scatter_add(0, segment_heads, real_exp) + denominator = null_exp + real_sum + real_weight = real_exp / denominator[segment_indices] + null_weight = null_exp / denominator + attended = hidden_branch.new_zeros( + (query_rows.numel(), self.num_heads, self.head_dim), + ) + attended.index_add_( + 0, segment_indices, real_weight[..., None] * real_value, + ) + delta = self.out_proj(attended.reshape(query_rows.numel(), self.dim)) + delta = delta * self.alpha[block_index].to(dtype=delta.dtype) + + # Entropy is conditional on selecting real evidence. Null-vs-real + # acceptance is reported independently as null_attention_mean; mixing + # that mass into this distribution would make singleton entropy nonzero. + real_conditional_weight = real_exp / real_sum[segment_indices].clamp_min( + torch.finfo(real_sum.dtype).tiny + ) + entropy_terms = -real_conditional_weight * torch.log( + real_conditional_weight.clamp_min( + torch.finfo(real_conditional_weight.dtype).tiny + ) + ) + real_entropy = torch.zeros_like(null_weight).scatter_add( + 0, segment_heads, entropy_terms, + ) + return delta, { + 'null_attention_mean': null_weight.mean(), + 'real_neighbour_entropy_mean': real_entropy.mean(), + 'effective_real_neighbour_count_mean': real_entropy.exp().mean(), + 'active_queries': delta.new_tensor(float(query_rows.numel())), + 'singleton_queries': delta.new_tensor(float((counts == 1).sum())), + 'packed_pairs': delta.new_tensor(float(source_indices.numel())), + } + + +class _V12LACrossRope(nn.Module): + """Apply the same 3D rotary basis to independently packed Q and K rows.""" + + def __init__( + self, + head_dim: int, + rope_freq: Tuple[float, float], + reserved_dims: int = 2, + ): + super().__init__() + content_dim = head_dim - reserved_dims + assert content_dim > 0 and content_dim % 6 == 0 + self.head_dim = head_dim + self.content_dim = content_dim + freq_dim = content_dim // 2 // 3 + frequencies = torch.arange(freq_dim, dtype=torch.float32) / freq_dim + frequencies = rope_freq[0] / (rope_freq[1] ** frequencies) + self.register_buffer('frequencies', frequencies, persistent=False) + + def forward(self, value: torch.Tensor, coords: torch.Tensor) -> torch.Tensor: + assert value.ndim == 3 and value.shape[-1] == self.head_dim + assert coords.shape == (value.shape[0], 3) + if value.shape[0] == 0: + return value + phases = torch.outer( + coords.reshape(-1).float(), self.frequencies, + ).reshape(coords.shape[0], self.content_dim // 2) + phases = torch.polar(torch.ones_like(phases), phases).unsqueeze(1) + content = torch.view_as_complex( + value[..., :self.content_dim].float().reshape( + value.shape[0], value.shape[1], -1, 2, + ) + ) + rotated = torch.view_as_real(content * phases).reshape( + value.shape[0], value.shape[1], self.content_dim, + ).to(dtype=value.dtype) + return torch.cat([rotated, value[..., self.content_dim:]], dim=-1) + + +class _PartialZObsLearnedAttention(nn.Module): + """Five-site V12LA observation attention with an exact carrier anchor.""" + + def __init__( + self, + model_channels: int, + num_heads: int, + sites: Sequence[int], + rope_freq: Tuple[float, float], + null_prior: float = 0.5, + ): + super().__init__() + self.sites = tuple(int(site) for site in sites) + assert self.sites == (5, 10, 15, 20, 25), self.sites + assert num_heads == 12 + assert model_channels % num_heads == 0 + assert (model_channels // num_heads - 2) > 0 + assert (model_channels // num_heads - 2) % 6 == 0 + self.model_channels = model_channels + self.num_heads = num_heads + self.head_dim = model_channels // num_heads + self.content_dim = self.head_dim - 2 + self.site_to_slot = {site: slot for slot, site in enumerate(self.sites)} + self.null_prior = float(null_prior) + assert 0.0 < self.null_prior < 1.0, \ + f'learned-attention null prior must lie in (0,1), got {self.null_prior}' + self.null_logit_offset = float( + np.log(self.null_prior) - np.log1p(-self.null_prior) + ) + + # K/V are source-only and shared across all five sites. Q and the + # zero-start deviation/output matrix are site-specific. + self.key_proj = nn.Linear(model_channels, model_channels, bias=False) + self.value_proj = nn.Linear(model_channels, model_channels, bias=False) + self.query_proj = nn.ModuleList([ + nn.Linear(model_channels, model_channels, bias=False) + for _ in self.sites + ]) + self.delta_proj = nn.ModuleList([ + nn.Linear(model_channels, model_channels, bias=False) + for _ in self.sites + ]) + self.null_key = nn.Parameter(torch.empty(num_heads, self.content_dim)) + self.rope = _V12LACrossRope(self.head_dim, rope_freq, reserved_dims=2) + self.reset_parameters() + + def reset_parameters(self) -> None: + nn.init.xavier_uniform_(self.key_proj.weight) + # Attention mixes the already trained V11 carrier basis at warm start; + # the value map may then adapt, but does not begin as a random rewrite. + nn.init.eye_(self.value_proj.weight) + nn.init.xavier_uniform_(self.null_key) + for query_proj in self.query_proj: + nn.init.xavier_uniform_(query_proj.weight) + for delta_proj in self.delta_proj: + nn.init.zeros_(delta_proj.weight) + + @staticmethod + def _rms_norm(value: torch.Tensor) -> torch.Tensor: + normalized = value.float() * torch.rsqrt( + value.float().square().mean(dim=-1, keepdim=True) + 1e-6 + ) + return normalized.to(dtype=value.dtype) + + @staticmethod + def _anchor( + carrier: torch.Tensor, + observed_rows: torch.Tensor, + full_rows: int, + ) -> torch.Tensor: + """Evaluate fixed coordinate/null attention without an additive scatter.""" + assert carrier.ndim == 2 + assert observed_rows.ndim == 1 and observed_rows.dtype == torch.long + assert carrier.shape[0] == observed_rows.numel() + assert full_rows >= observed_rows.numel() + if observed_rows.numel() != 0: + assert observed_rows.min() >= 0 and observed_rows.max() < full_rows + assert torch.unique(observed_rows).numel() == observed_rows.numel() + # Index zero is the exact-zero null value. Every observed row is then + # redirected to its own carrier row; unmatched rows keep selecting null. + values = torch.cat([carrier.new_zeros((1, carrier.shape[1])), carrier], dim=0) + anchor_index = torch.zeros(full_rows, dtype=torch.long, device=carrier.device) + anchor_index[observed_rows] = torch.arange( + 1, observed_rows.numel() + 1, dtype=torch.long, device=carrier.device, + ) + return values[anchor_index] + + def build_cache( + self, + full_coords: torch.Tensor, + full_layout: List[slice], + carrier: torch.Tensor, + observed_rows: torch.Tensor, + ) -> Dict[str, Any]: + """Build the exact anchor and the shared per-sample O+null K/V cache.""" + assert full_coords.ndim == 2 and full_coords.shape[1] == 4 + assert carrier.shape == (observed_rows.numel(), self.model_channels) + assert len(full_layout) > 0 + anchor = self._anchor(carrier, observed_rows, full_coords.shape[0]) + observed_coords = full_coords[observed_rows] + observed_batch = observed_coords[:, 0].long() + assert observed_batch.numel() == 0 or ( + observed_batch.min() >= 0 and observed_batch.max() < len(full_layout) + ) + + branch_dtype = self.key_proj.weight.dtype + carrier_branch = carrier.to(dtype=branch_dtype) + carrier_norm = F.layer_norm( + carrier_branch, (self.model_channels,), eps=1e-6, + ) + real_key = self.key_proj(carrier_norm).reshape( + -1, self.num_heads, self.head_dim, + ) + real_key = self.rope(real_key, observed_coords[:, 1:]) + real_key = self._rms_norm(real_key[..., :self.content_dim]) + real_key = torch.cat([ + real_key, + real_key.new_zeros((real_key.shape[0], self.num_heads, 2)), + ], dim=-1) + real_value = self.value_proj(carrier_branch).reshape( + -1, self.num_heads, self.head_dim, + ) + null_content = self._rms_norm( + self.null_key.to(dtype=real_key.dtype) + ) + + packed_kv = [] + observed_counts = [] + for batch_index in range(len(full_layout)): + source_indices = torch.nonzero( + observed_batch == batch_index, as_tuple=False, + ).flatten() + observed_count = int(source_indices.numel()) + observed_counts.append(observed_count) + # One null competes with O real keys. log(O) cancels its cardinality + # disadvantage; the constant logit offset sets the null mass only in + # the equal-content-logit reference case, not after learned Q/K drift. + null_reserved = real_key.new_zeros((self.num_heads, 2)) + null_reserved[:, 0] = np.sqrt(self.head_dim) * ( + np.log(max(observed_count, 1)) + self.null_logit_offset + ) + null_key = torch.cat([null_content, null_reserved], dim=-1).unsqueeze(0) + null_value = real_value.new_zeros((1, self.num_heads, self.head_dim)) + keys = torch.cat([real_key[source_indices], null_key], dim=0) + values = torch.cat([real_value[source_indices], null_value], dim=0) + packed_kv.append(torch.stack([keys, values], dim=1)) + + kv = sp.VarLenTensor.from_tensor_list(packed_kv) + return { + 'anchor': anchor, + 'kv': kv, + 'observed_rows': observed_rows, + 'observed_counts': tuple(observed_counts), + } + + def forward( + self, + hidden: torch.Tensor, + full_coords: torch.Tensor, + full_layout: List[slice], + cache: Dict[str, Any], + block_index: int, + ) -> Tuple[torch.Tensor, Dict[str, torch.Tensor]]: + assert block_index in self.site_to_slot, block_index + assert hidden.shape == (full_coords.shape[0], self.model_channels) + slot = self.site_to_slot[block_index] + branch_dtype = self.query_proj[slot].weight.dtype + hidden_branch = hidden.to(dtype=branch_dtype) + query = self.query_proj[slot](F.layer_norm( + hidden_branch, (self.model_channels,), eps=1e-6, + )).reshape(-1, self.num_heads, self.head_dim) + query = self.rope(query, full_coords[:, 1:]) + query_content = self._rms_norm(query[..., :self.content_dim]) + query_reserved = query.new_zeros((query.shape[0], self.num_heads, 2)) + query_reserved[..., 0] = 1 + query = torch.cat([query_content, query_reserved], dim=-1) + query_varlen = sp.VarLenTensor(query, full_layout) + attended = sp.sparse_scaled_dot_product_attention( + query_varlen, cache['kv'], + ).feats.reshape(-1, self.model_channels) + + anchor = cache['anchor'].to(dtype=attended.dtype) + deviation_input = attended - anchor + correction = self.delta_proj[slot](deviation_input) + inject = anchor + correction + eps = torch.finfo(inject.dtype).eps + observed_rows = cache['observed_rows'] + observed_mask = torch.zeros( + inject.shape[0], dtype=torch.bool, device=inject.device, + ) + observed_mask[observed_rows] = True + return inject, { + 'correction_to_anchor_norm_ratio': correction.norm() / anchor.norm().clamp_min(eps), + 'correction_to_hidden_norm_ratio': correction.norm() / hidden_branch.norm().clamp_min(eps), + 'observed_correction_l2_norm': correction[observed_mask].float().norm(), + 'unobserved_correction_l2_norm': correction[~observed_mask].float().norm(), + 'attention_output_l2_norm': attended.float().norm(), + 'anchor_l2_norm': anchor.float().norm(), + 'delta_weight_l2_norm': self.delta_proj[slot].weight.float().norm(), + } + +class SLatFlowModel(nn.Module): + """ + Structured Latent Flow Model for 3D generation. + + Supports two conditioning modes: + - "cross": Standard cross-attention with image features + - "proj": View-aligned projection attention with camera-aware features + """ + def __init__( + self, + resolution: int, + in_channels: int, + model_channels: int, + cond_channels: int, + out_channels: int, + num_blocks: int, + num_heads: Optional[int] = None, + num_head_channels: Optional[int] = 64, + mlp_ratio: float = 4, + pe_mode: Literal["ape", "rope"] = "ape", + rope_freq: Tuple[float, float] = (1.0, 10000.0), + dtype: str = 'float32', + use_checkpoint: bool = False, + share_mod: bool = False, + initialization: str = 'vanilla', + qk_rms_norm: bool = False, + qk_rms_norm_cross: bool = False, + image_attn_mode: Literal["cross", "proj", "gated_proj"] = "cross", + proj_in_channels: Optional[int] = None, + vae_in_channels: Optional[int] = None, + zobs_inject: bool = False, + zobs_inject_in_channels: int = 33, + zobs_inject_scale: float = 1.0, + partial_zobs_inject: bool = False, + partial_zobs_in_channels: int = 34, + partial_zobs_inject_position: Literal['post_block', 'pre_mlp'] = 'post_block', + legacy_zobs_zero_input: bool = False, + partial_zobs_innovation_gate: bool = False, + innovation_gate_dim: int = 128, + innovation_gate_time_dim: int = 32, + partial_zobs_innovation_gate_train_only: bool = False, + partial_zobs_local_attention: bool = False, + local_attention_dim: int = 128, + local_attention_time_dim: int = 32, + local_attention_heads: int = 8, + partial_zobs_local_attention_train_only: bool = False, + partial_zobs_learned_attention: bool = False, + learned_attention_sites: Sequence[int] = (5, 10, 15, 20, 25), + learned_attention_heads: int = 12, + learned_attention_null_prior: float = 0.5, + partial_zobs_learned_attention_train_only: bool = False, + **kwargs + ): + super().__init__() + self.resolution = resolution + self.in_channels = in_channels + self.model_channels = model_channels + self.cond_channels = cond_channels + self.out_channels = out_channels + self.num_blocks = num_blocks + self.num_heads = num_heads or model_channels // num_head_channels + self.mlp_ratio = mlp_ratio + self.pe_mode = pe_mode + self.use_checkpoint = use_checkpoint + self.share_mod = share_mod + self.initialization = initialization + self.qk_rms_norm = qk_rms_norm + self.qk_rms_norm_cross = qk_rms_norm_cross + self.image_attn_mode = image_attn_mode + self.proj_in_channels = proj_in_channels + self.vae_in_channels = vae_in_channels + self.dtype = str_to_dtype(dtype) + + self.t_embedder = TimestepEmbedder(model_channels) + if share_mod: + self.adaLN_modulation = nn.Sequential( + nn.SiLU(), + nn.Linear(model_channels, 6 * model_channels, bias=True) + ) + + if pe_mode == "ape": + self.pos_embedder = AbsolutePositionEmbedder(model_channels) + + self.input_layer = sp.SparseLinear(in_channels, model_channels) + # Trained-bridge: single shared per-block injector of [z_obs (x) vis (32) | vis (1)] -> model_channels. + self.zobs_inject = zobs_inject + self.zobs_inject_scale = float(zobs_inject_scale) + self.zobs_injector = sp.SparseLinear(zobs_inject_in_channels, model_channels) if zobs_inject else None + # v8-partial: the observation lives only on its observed support. This adapter is + # deliberately bias-free: absent rows never acquire a learned content/fill value. + self.partial_zobs_inject = partial_zobs_inject + self.partial_zobs_in_channels = partial_zobs_in_channels + assert partial_zobs_inject_position in ('post_block', 'pre_mlp'), \ + f'partial_zobs_inject_position must be post_block or pre_mlp, got {partial_zobs_inject_position!r}' + self.partial_zobs_inject_position = partial_zobs_inject_position + if partial_zobs_inject_position == 'pre_mlp': + assert not any(( + partial_zobs_innovation_gate, + partial_zobs_local_attention, + partial_zobs_learned_attention, + )), 'partial_zobs_inject_position=pre_mlp is incompatible with post-block experimental branches' + self.legacy_zobs_zero_input = legacy_zobs_zero_input + self.partial_zobs_injector = ( + sp.SparseLinear(partial_zobs_in_channels, model_channels, bias=False) + if partial_zobs_inject else None + ) + self.partial_zobs_innovation_gate = partial_zobs_innovation_gate + self.partial_zobs_innovation_gate_train_only = partial_zobs_innovation_gate_train_only + if partial_zobs_innovation_gate_train_only: + assert partial_zobs_innovation_gate, \ + 'partial_zobs_innovation_gate_train_only requires partial_zobs_innovation_gate=true' + if partial_zobs_innovation_gate: + assert partial_zobs_inject, 'partial_zobs_innovation_gate requires partial_zobs_inject=true' + assert partial_zobs_in_channels == 34, \ + f'partial_zobs_innovation_gate requires 34 input channels, got {partial_zobs_in_channels}' + assert innovation_gate_dim == 128, f'innovation_gate_dim is fixed at 128, got {innovation_gate_dim}' + assert innovation_gate_time_dim == 32, \ + f'innovation_gate_time_dim is fixed at 32, got {innovation_gate_time_dim}' + assert model_channels == 1536, f'partial_zobs_innovation_gate requires 1536 model channels, got {model_channels}' + assert num_blocks == 30, f'partial_zobs_innovation_gate requires 30 blocks, got {num_blocks}' + self.innovation_gate = _PartialZObsInnovationGate( + model_channels, partial_zobs_in_channels, innovation_gate_dim, + innovation_gate_time_dim, num_blocks, + ) + else: + self.innovation_gate = None + self.partial_zobs_local_attention = partial_zobs_local_attention + self.partial_zobs_local_attention_train_only = partial_zobs_local_attention_train_only + if partial_zobs_local_attention_train_only: + assert partial_zobs_local_attention, \ + 'partial_zobs_local_attention_train_only requires partial_zobs_local_attention=true' + assert sum(bool(value) for value in ( + partial_zobs_innovation_gate, + partial_zobs_local_attention, + partial_zobs_learned_attention, + )) <= 1, 'partial-zobs experimental attention/gate branches are mutually exclusive' + if partial_zobs_local_attention: + assert partial_zobs_inject, 'partial_zobs_local_attention requires partial_zobs_inject=true' + assert partial_zobs_in_channels == 34, \ + f'partial_zobs_local_attention requires 34 input channels, got {partial_zobs_in_channels}' + assert local_attention_dim == 128, \ + f'local_attention_dim is fixed at 128, got {local_attention_dim}' + assert local_attention_time_dim == 32, \ + f'local_attention_time_dim is fixed at 32, got {local_attention_time_dim}' + assert local_attention_heads == 8, \ + f'local_attention_heads is fixed at 8, got {local_attention_heads}' + assert model_channels == 1536, \ + f'partial_zobs_local_attention requires 1536 model channels, got {model_channels}' + assert num_blocks == 30, \ + f'partial_zobs_local_attention requires 30 blocks, got {num_blocks}' + self.local_obs_attention = _PartialZObsLocalAttention( + model_channels, partial_zobs_in_channels, local_attention_dim, + local_attention_time_dim, local_attention_heads, num_blocks - 1, + ) + else: + self.local_obs_attention = None + + self.partial_zobs_learned_attention = partial_zobs_learned_attention + self.partial_zobs_learned_attention_train_only = \ + partial_zobs_learned_attention_train_only + if partial_zobs_learned_attention_train_only: + assert partial_zobs_learned_attention, \ + 'partial_zobs_learned_attention_train_only requires partial_zobs_learned_attention=true' + if partial_zobs_learned_attention: + assert partial_zobs_inject, \ + 'partial_zobs_learned_attention requires partial_zobs_inject=true' + assert partial_zobs_in_channels == 34, \ + f'partial_zobs_learned_attention requires 34 input channels, got {partial_zobs_in_channels}' + assert model_channels == 1536, \ + f'partial_zobs_learned_attention requires 1536 model channels, got {model_channels}' + assert num_blocks == 30, \ + f'partial_zobs_learned_attention requires 30 blocks, got {num_blocks}' + assert learned_attention_heads == 12, \ + f'partial_zobs_learned_attention requires 12 heads, got {learned_attention_heads}' + assert tuple(learned_attention_sites) == (5, 10, 15, 20, 25), \ + f'V12LA sites are fixed at [5,10,15,20,25], got {learned_attention_sites}' + self.learned_attention = _PartialZObsLearnedAttention( + model_channels, + learned_attention_heads, + learned_attention_sites, + rope_freq, + learned_attention_null_prior, + ) + else: + self.learned_attention = None + + self.blocks = nn.ModuleList([ + ModulatedSparseTransformerCrossBlock( + model_channels, + cond_channels, + num_heads=self.num_heads, + mlp_ratio=self.mlp_ratio, + attn_mode='full', + use_checkpoint=self.use_checkpoint, + use_rope=(pe_mode == "rope"), + rope_freq=rope_freq, + share_mod=self.share_mod, + qk_rms_norm=self.qk_rms_norm, + qk_rms_norm_cross=self.qk_rms_norm_cross, + image_attn_mode=image_attn_mode, + proj_in_channels=proj_in_channels, + vae_in_channels=vae_in_channels, + ) + for _ in range(num_blocks) + ]) + + self.out_layer = sp.SparseLinear(model_channels, out_channels) + + self.initialize_weights() + self.convert_to(self.dtype) + if self.partial_zobs_innovation_gate_train_only: + self._freeze_for_partial_zobs_innovation_gate() + if self.partial_zobs_local_attention_train_only: + self._freeze_for_partial_zobs_local_attention() + if self.partial_zobs_learned_attention_train_only: + self._freeze_for_partial_zobs_learned_attention() + + @property + def device(self) -> torch.device: + """ + Return the device of the model. + """ + return next(self.parameters()).device + + def convert_to(self, dtype: torch.dtype) -> None: + """ + Convert the torso of the model to the specified dtype. + """ + self.dtype = dtype + self.blocks.apply(partial(convert_module_to, dtype=dtype)) + if self.learned_attention is not None: + self.learned_attention.apply(partial(convert_module_to, dtype=dtype)) + + def initialize_weights(self) -> None: + if self.initialization == 'vanilla': + # Initialize transformer layers: + def _basic_init(module): + if isinstance(module, nn.Linear): + torch.nn.init.xavier_uniform_(module.weight) + if module.bias is not None: + nn.init.constant_(module.bias, 0) + self.apply(_basic_init) + + # Initialize timestep embedding MLP: + nn.init.normal_(self.t_embedder.mlp[0].weight, std=0.02) + nn.init.normal_(self.t_embedder.mlp[2].weight, std=0.02) + + # Zero-out adaLN modulation layers in DiT blocks: + if self.share_mod: + nn.init.constant_(self.adaLN_modulation[-1].weight, 0) + nn.init.constant_(self.adaLN_modulation[-1].bias, 0) + else: + for block in self.blocks: + nn.init.constant_(block.adaLN_modulation[-1].weight, 0) + nn.init.constant_(block.adaLN_modulation[-1].bias, 0) + + # Zero-out output layers: + nn.init.constant_(self.out_layer.weight, 0) + nn.init.constant_(self.out_layer.bias, 0) + + elif self.initialization == 'scaled': + # Initialize transformer layers: + def _basic_init(module): + if isinstance(module, nn.Linear): + torch.nn.init.normal_(module.weight, std=np.sqrt(2.0 / (5.0 * self.model_channels))) + if module.bias is not None: + nn.init.constant_(module.bias, 0) + self.apply(_basic_init) + + # Scaled init for to_out and ffn2 + def _scaled_init(module): + if isinstance(module, nn.Linear): + torch.nn.init.normal_(module.weight, std=1.0 / np.sqrt(5 * self.num_blocks * self.model_channels)) + if module.bias is not None: + nn.init.constant_(module.bias, 0) + for block in self.blocks: + block.self_attn.to_out.apply(_scaled_init) + # Handle cross, proj, and gated_proj modes + if self.image_attn_mode in ("proj", "gated_proj"): + block.cross_attn.cross_attn_block.to_out.apply(_scaled_init) + else: + block.cross_attn.to_out.apply(_scaled_init) + block.mlp.mlp[2].apply(_scaled_init) + + # Initialize input layer to make the initial representation have variance 1 + nn.init.normal_(self.input_layer.weight, std=1.0 / np.sqrt(self.in_channels)) + nn.init.zeros_(self.input_layer.bias) + + # Initialize timestep embedding MLP: + nn.init.normal_(self.t_embedder.mlp[0].weight, std=0.02) + nn.init.normal_(self.t_embedder.mlp[2].weight, std=0.02) + + # Zero-out adaLN modulation layers in DiT blocks: + if self.share_mod: + nn.init.constant_(self.adaLN_modulation[-1].weight, 0) + nn.init.constant_(self.adaLN_modulation[-1].bias, 0) + else: + for block in self.blocks: + nn.init.constant_(block.adaLN_modulation[-1].weight, 0) + nn.init.constant_(block.adaLN_modulation[-1].bias, 0) + + # Zero-out output layers: + nn.init.constant_(self.out_layer.weight, 0) + nn.init.constant_(self.out_layer.bias, 0) + + # Trained-bridge: zero-init the z_obs injector LAST (after _basic_init), so step-0 injection == 0. + if getattr(self, 'zobs_injector', None) is not None: + nn.init.zeros_(self.zobs_injector.weight) + nn.init.zeros_(self.zobs_injector.bias) + if getattr(self, 'partial_zobs_injector', None) is not None: + nn.init.zeros_(self.partial_zobs_injector.weight) + if getattr(self, 'innovation_gate', None) is not None: + nn.init.zeros_(self.innovation_gate.alpha) + if getattr(self, 'local_obs_attention', None) is not None: + # V12-A freezes vanilla Xavier initialization independently of the + # scaled V11 backbone initialization mode. Only alpha starts at zero. + for module in self.local_obs_attention.modules(): + if isinstance(module, nn.Linear): + nn.init.xavier_uniform_(module.weight) + if module.bias is not None: + nn.init.zeros_(module.bias) + nn.init.xavier_uniform_(self.local_obs_attention.relative_offset) + nn.init.xavier_uniform_(self.local_obs_attention.null_key) + nn.init.zeros_(self.local_obs_attention.alpha) + if getattr(self, 'learned_attention', None) is not None: + self.learned_attention.reset_parameters() + + def _freeze_for_partial_zobs_innovation_gate(self) -> None: + """Make the V12 gate the sole trainable branch before an optimizer is built.""" + assert self.innovation_gate is not None + for parameter in self.parameters(): + parameter.requires_grad_(False) + for parameter in self.innovation_gate.parameters(): + parameter.requires_grad_(True) + trainable = [name for name, parameter in self.named_parameters() if parameter.requires_grad] + expected = [f'innovation_gate.{name}' for name in self.innovation_gate.state_dict()] + assert trainable == expected, (trainable, expected) + assert sum(parameter.numel() for parameter in self.parameters() if parameter.requires_grad) == 473663 + + def _freeze_for_partial_zobs_local_attention(self) -> None: + """Make the V12-A local attention the sole trainable warm-start branch.""" + assert self.local_obs_attention is not None + for parameter in self.parameters(): + parameter.requires_grad_(False) + for parameter in self.local_obs_attention.parameters(): + parameter.requires_grad_(True) + trainable = [name for name, parameter in self.named_parameters() if parameter.requires_grad] + expected = [f'local_obs_attention.{name}' for name in self.local_obs_attention.state_dict()] + assert trainable == expected, (trainable, expected) + assert len(trainable) == 14 + assert sum(parameter.numel() for parameter in self.parameters() if parameter.requires_grad) == 487485 + + def _freeze_for_partial_zobs_learned_attention(self) -> None: + """Make the V12LA module the sole trainable warm-start branch.""" + assert self.learned_attention is not None + for parameter in self.parameters(): + parameter.requires_grad_(False) + for parameter in self.learned_attention.parameters(): + parameter.requires_grad_(True) + trainable = [name for name, parameter in self.named_parameters() if parameter.requires_grad] + expected = [ + f'learned_attention.{name}' + for name, _ in self.learned_attention.named_parameters() + ] + assert trainable == expected, (trainable, expected) + assert len(trainable) == 13, trainable + + def _partial_zobs_message( + self, + x: sp.SparseTensor, + partial_z_obs: sp.SparseTensor, + partial_z_obs_keep_mask: Optional[torch.Tensor], + ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """Project a variable-support observation and join its messages onto ``x`` rows. + + The return value stays variable-length: projected messages plus their exact full-row + indices. Neither observation content nor a zero-filled model-channel tensor is + materialized on full support. ``torch.unique(..., dim=0)`` makes the coordinate join + exact and rejects duplicate or non-subset partial coordinates instead of assuming a + row order. + """ + adapter = self.partial_zobs_injector + assert adapter is not None + assert partial_z_obs.feats.ndim == 2 and \ + partial_z_obs.feats.shape[-1] == self.partial_zobs_in_channels, \ + f'partial_z_obs must be [N,{self.partial_zobs_in_channels}], got {tuple(partial_z_obs.feats.shape)}' + assert partial_z_obs.coords.ndim == 2 and partial_z_obs.coords.shape[1] == x.coords.shape[1], \ + f'partial_z_obs coordinate rank mismatch: {tuple(partial_z_obs.coords.shape)} vs {tuple(x.coords.shape)}' + + active = torch.ones(partial_z_obs.feats.shape[0], dtype=torch.bool, device=partial_z_obs.feats.device) + if partial_z_obs_keep_mask is not None: + assert partial_z_obs_keep_mask.ndim == 1 and partial_z_obs_keep_mask.shape[0] == x.shape[0], \ + f'partial_z_obs_keep_mask must be [B={x.shape[0]}], got {tuple(partial_z_obs_keep_mask.shape)}' + keep = partial_z_obs_keep_mask.to(device=partial_z_obs.feats.device, dtype=torch.bool) + active = keep[partial_z_obs.coords[:, 0].long()] + + active_coords = partial_z_obs.coords[active] + active_feats = partial_z_obs.feats[active] + # F.linear on a zero-row tensor is intentional: it keeps the adapter parameters in the + # graph on all-drop batches, so DDP sees a zero (rather than None) gradient. + active_messages = F.linear(active_feats, adapter.weight, None) + + full_rows = torch.empty(0, dtype=torch.long, device=x.feats.device) + if active_coords.shape[0] != 0: + all_coords = torch.cat([x.coords, active_coords], dim=0) + _, inverse = torch.unique(all_coords, dim=0, return_inverse=True) + x_ids = inverse[:x.coords.shape[0]] + partial_ids = inverse[x.coords.shape[0]:] + assert torch.unique(x_ids).numel() == x_ids.numel(), 'x contains duplicate coordinates' + assert torch.unique(partial_ids).numel() == partial_ids.numel(), 'partial_z_obs contains duplicate coordinates' + ordered_x_ids, ordered_x_rows = torch.sort(x_ids) + positions = torch.searchsorted(ordered_x_ids, partial_ids) + assert (positions < ordered_x_ids.numel()).all() and \ + torch.equal(ordered_x_ids[positions], partial_ids), \ + 'partial_z_obs coordinates must be an exact subset of current x coordinates' + full_rows = ordered_x_rows[positions] + + return active_messages, full_rows, active_feats, active_coords[:, 0].long() + + def forward( + self, + x: sp.SparseTensor, + t: torch.Tensor, + cond: Union[torch.Tensor, List[torch.Tensor], Dict[str, Union[torch.Tensor, sp.SparseTensor]], Tuple], + concat_cond: Optional[sp.SparseTensor] = None, + z_obs: Optional[sp.SparseTensor] = None, + partial_z_obs: Optional[sp.SparseTensor] = None, + partial_z_obs_keep_mask: Optional[torch.Tensor] = None, + **kwargs + ) -> sp.SparseTensor: + """ + Forward pass. + + Args: + x: SparseTensor input + t: Timestep tensor [B] + cond: Conditioning tensor. For "cross" mode: list of tensors or tensor. + For "proj" mode: dict {'global': global_cond, 'proj': proj_cond} + or tuple of (global_cond, proj_cond) + concat_cond: Optional concatenation condition + z_obs: Optional 33ch [z_obs(32)|vis(1)] sparse tensor, same coords as x. Feeds ONLY + the zobs_injector (v6 patchcross path: concat_cond is the plain 32ch shape + latent and z_obs never enters the input_layer). When None, the legacy 97ch + concat slice below applies — shipped/v4 behavior is bit-identical. + partial_z_obs: Optional variable-length v8 observation + [normalized_z_obs(32)|coverage(1)|boundary_distance(1)]. + Its coordinates must be a subset of x; it is consumed only by the dedicated, + bias-free partial adapter and is never converted to full-support content. + + Returns: + SparseTensor output + """ + if concat_cond is not None: + x = sp.sparse_cat([x, concat_cond], dim=-1) + + h = self.input_layer(x) + h = manual_cast(h, self.dtype) + + # Trained-bridge: compute the per-block injection vector ONCE (shared zero-init linear, added each block). + # Sources, in priority order: separate z_obs kwarg (v6), else the trailing channels of a + # 97-ch [x_t|shape|z_obs|vis] input (v4). z_obs*vis==0 at vis=0 in both. + zobs_inject = None + _injector = getattr(self, 'zobs_injector', None) + if z_obs is not None: + # Fail fast: without an injector the kwarg would be a silent no-op, and in the 64ch + # patchcross architecture the injector is the ONLY z_obs consumption path. + assert _injector is not None, \ + 'z_obs kwarg passed but the denoiser was built without zobs_inject=true' + # Full coords equality is asserted at data-load; a row-count check catches gross mismatch. + assert z_obs.feats.shape[-1] == 33 and z_obs.feats.shape[0] == x.feats.shape[0], \ + f'z_obs kwarg mismatch: {tuple(z_obs.feats.shape)} vs x rows {x.feats.shape[0]}' + _zo = z_obs.feats[:, :32] + _vis = z_obs.feats[:, 32:33] + zobs_inject = manual_cast(self.zobs_injector(z_obs.replace(torch.cat([_zo * _vis, _vis], dim=-1))), self.dtype) + if not getattr(self, '_zobs_kwarg_logged', False): # in-log proof hook (once per process) + print('[zobs_injector] consuming separate z_obs kwarg (patchcross path)', flush=True) + self._zobs_kwarg_logged = True + elif _injector is not None and x.feats.shape[-1] >= 97: + _zo = x.feats[:, 64:96] + _vis = x.feats[:, 96:97] + zobs_inject = manual_cast(self.zobs_injector(x.replace(torch.cat([_zo * _vis, _vis], dim=-1))), self.dtype) + elif _injector is not None and self.legacy_zobs_zero_input: + # 64ch v8 shell with a legacy injector: preserve the legacy all-zero z_obs mode, + # including its learned per-block bias, without materialising observation content. + # Evaluate one zero row and broadcast the result: this is exactly W*0+b, keeps a + # zero (non-None) weight gradient under DDP, and avoids an N-row 33ch GEMM. + zero_output = F.linear( + x.feats.new_zeros((1, _injector.in_features)), + _injector.weight, + _injector.bias, + ) + zobs_inject = manual_cast( + x.replace(zero_output.expand(x.feats.shape[0], -1)), self.dtype + ) + elif _injector is not None and not getattr(self, '_zobs_none_warned', False): + # Injector present but NO z_obs source (64ch input, kwarg omitted): the trained + # per-block bias is NOT applied — a third mode distinct from "z_obs dropped" (an + # all-zero 33ch tensor still adds the bias each block). A nozobs arm on a patchcross + # model must pass the zero tensor, never omit the kwarg. + print('[zobs_injector] WARNING: injector present but no z_obs source (kwarg omitted, ' + 'input < 97ch) — per-block bias NOT applied; pass an all-zero 33ch z_obs for a ' + 'true nozobs arm', flush=True) + self._zobs_none_warned = True + + partial_zobs_inject = None + partial_zobs_rows = None + partial_zobs_feats = None + partial_zobs_batch = None + if partial_z_obs is not None: + assert self.partial_zobs_injector is not None, \ + 'partial_z_obs passed but denoiser was built without partial_zobs_inject=true' + partial_zobs_inject, partial_zobs_rows, partial_zobs_feats, partial_zobs_batch = self._partial_zobs_message( + x, partial_z_obs, partial_z_obs_keep_mask + ) + partial_zobs_inject = manual_cast(partial_zobs_inject, self.dtype) + elif self.partial_zobs_injector is not None: + raise AssertionError('partial_zobs_inject=true requires a separate partial_z_obs tensor') + + if zobs_inject is not None and self.zobs_inject_scale != 1.0: + zobs_inject = zobs_inject.replace(zobs_inject.feats * self.zobs_inject_scale) + if not getattr(self, '_zobs_inject_scale_logged', False): + print(f'[zobs_injector] output scale={self.zobs_inject_scale:g}', flush=True) + self._zobs_inject_scale_logged = True + + # Keep the pre-AdaLN embedding for the innovation gate; blocks retain their existing + # modulation input exactly. + t_base = self.t_embedder(t) + t_emb = t_base + if self.share_mod: + t_emb = self.adaLN_modulation(t_emb) + t_emb = manual_cast(t_emb, self.dtype) + + innovation_time = None + innovation_y = None + innovation_gate_y = None + if self.innovation_gate is not None and partial_zobs_feats is not None: + # Observation-only terms are shared by all blocks. Empty rows deliberately pass + # through these linears too, so an all-drop batch keeps the full gate graph live. + gate_dtype = self.innovation_gate.embed_y.weight.dtype + innovation_gate_y = partial_zobs_feats.to(dtype=gate_dtype) + innovation_y = self.innovation_gate.embed_y(innovation_gate_y) + innovation_time = self.innovation_gate.time_proj(t_base) + + if self.pe_mode == "ape": + pe = self.pos_embedder(h.coords[:, 1:]) + h = h + manual_cast(pe, self.dtype) + + # Handle different conditioning modes + if self.image_attn_mode == 'proj': + if isinstance(cond, dict): + global_cond = cond['global'] + proj_cond = cond['proj'] + else: + global_cond, proj_cond = cond + if isinstance(global_cond, list): + global_cond = sp.VarLenTensor.from_tensor_list(global_cond) + global_cond = manual_cast(global_cond, self.dtype) + proj_cond = manual_cast(proj_cond, self.dtype) + cond = (global_cond, proj_cond) + elif self.image_attn_mode == 'gated_proj': + global_cond = cond['global'] + if isinstance(global_cond, list): + global_cond = sp.VarLenTensor.from_tensor_list(global_cond) + global_cond = manual_cast(global_cond, self.dtype) + proj_semantic = manual_cast(cond['proj_semantic'], self.dtype) + proj_color = manual_cast(cond['proj_color'], self.dtype) + cond = {'global': global_cond, 'proj_semantic': proj_semantic, 'proj_color': proj_color} + else: + if isinstance(cond, list): + cond = sp.VarLenTensor.from_tensor_list(cond) + cond = manual_cast(cond, self.dtype) + + # Inference-only probe hook. When the attribute is absent (the default, and the + # only state any training path produces) this is None and the loop below is + # byte-identical to the unhooked path: no constructor arg, no config key, and no + # entry in the inference ABI contract. A probe sets it to a length-num_blocks + # sequence to rescale the shared partial message per block. + _partial_scales = getattr(self, 'partial_zobs_layer_scales', None) + # Optional evaluator-only callable: ``fn(gate_value, active_batch, block_index)``. + # It is an instance attribute rather than a constructor/config argument, so no model + # state or training ABI changes; callers can implement neutral, constant, or shuffled + # gate-value probes deterministically. + _gate_value_transform = getattr(self, 'partial_zobs_innovation_gate_value_transform', None) + _innovation_ratio_stats = [] if self.innovation_gate is not None else None + _innovation_gate_mean_stats = [] if self.innovation_gate is not None else None + _innovation_observed_rows = 0 if partial_zobs_rows is None else int(partial_zobs_rows.numel()) + _local_attention = getattr(self, 'local_obs_attention', None) + _local_topology = None + _local_full_batch = None + _local_ratio_base_stats = [] if _local_attention is not None else None + _local_ratio_hidden_stats = [] if _local_attention is not None else None + _local_null_stats = [] if _local_attention is not None else None + _local_entropy_stats = [] if _local_attention is not None else None + _local_effective_neighbour_stats = [] if _local_attention is not None else None + _local_observed_residual_stats = [] if _local_attention is not None else None + _local_unobserved_residual_stats = [] if _local_attention is not None else None + _local_query_is_observed = None + _local_neighbourhood_histogram = None + if _local_attention is not None: + assert partial_zobs_rows is not None and partial_zobs_feats is not None + _local_topology = _build_two_hop_observed_topology(x.coords, partial_zobs_rows) + _local_full_batch = x.coords[:, 0].long() + _active_observed_mask = torch.zeros( + x.coords.shape[0], dtype=torch.bool, device=x.coords.device, + ) + _active_observed_mask[partial_zobs_rows] = True + _local_query_is_observed = _active_observed_mask[ + _local_topology['query_rows'] + ] + _local_neighbourhood_histogram = torch.bincount( + _local_topology['neighbour_counts'], minlength=26, + )[:26] + _local_neighbourhood_histogram[0] = ( + x.coords.shape[0] - _local_topology['query_rows'].numel() + ) + assert int(_local_neighbourhood_histogram.sum()) == x.coords.shape[0] + assert int(_local_neighbourhood_histogram[1:].sum()) == \ + _local_topology['query_rows'].numel() + _learned_attention = getattr(self, 'learned_attention', None) + _learned_cache = None + _learned_stats = [] if _learned_attention is not None else None + if _learned_attention is not None: + assert partial_zobs_inject is not None and partial_zobs_rows is not None + assert _partial_scales is None, \ + 'partial_zobs_layer_scales is not defined for the five-site V12LA operator' + _learned_cache = _learned_attention.build_cache( + x.coords, + x.layout, + partial_zobs_inject, + partial_zobs_rows, + ) + partial_zobs_inject_position = getattr( + self, 'partial_zobs_inject_position', 'post_block', + ) + for _block_index, block in enumerate(self.blocks): + if partial_zobs_inject is not None and partial_zobs_inject_position == 'pre_mlp': + _pre_mlp_scale = 1.0 if _partial_scales is None else float( + _partial_scales[_block_index] + ) + h = block( + h, + t_emb, + cond, + partial_zobs_rows=partial_zobs_rows, + partial_zobs_inject=partial_zobs_inject, + partial_zobs_inject_scale=_pre_mlp_scale, + ) + else: + h = block(h, t_emb, cond) + if zobs_inject is not None: + h = h + zobs_inject + _innovation_ratio = None + _innovation_gate_mean = None + _local_delta = None + _local_stats = None + if _local_attention is not None and _block_index < _local_attention.alpha.numel(): + assert _local_topology is not None and _local_full_batch is not None + _local_delta, _local_stats = _local_attention( + h.feats, + _local_full_batch, + t_base, + partial_zobs_feats, + partial_zobs_rows, + _local_topology, + _block_index, + ) + _local_delta = manual_cast(_local_delta, self.dtype) + _local_ratio_base = _local_delta.norm() / partial_zobs_inject.norm().clamp_min( + torch.finfo(partial_zobs_inject.dtype).eps + ) + _local_hidden = h.feats[_local_topology['query_rows']] + _local_ratio_hidden = _local_delta.norm() / _local_hidden.norm().clamp_min( + torch.finfo(_local_hidden.dtype).eps + ) + assert _local_query_is_observed is not None + _local_observed_residual = _local_delta[ + _local_query_is_observed + ].float().norm() + _local_unobserved_residual = _local_delta[ + ~_local_query_is_observed + ].float().norm() + if _learned_attention is not None and _block_index in _learned_attention.site_to_slot: + assert _learned_cache is not None + _learned_inject, _site_stats = _learned_attention( + h.feats, + h.coords, + h.layout, + _learned_cache, + _block_index, + ) + h = h.replace(h.feats + manual_cast(_learned_inject, self.dtype)) + _learned_stats.append({ + name: value.detach() for name, value in _site_stats.items() + }) + if partial_zobs_inject is not None and _learned_attention is None and \ + partial_zobs_inject_position == 'post_block': + layer_inject = partial_zobs_inject + if self.innovation_gate is not None: + gate = self.innovation_gate + # Every operation intentionally accepts an empty leading dimension. This + # keeps all innovation parameters in the all-drop DDP graph. + assert innovation_gate_y is not None and innovation_y is not None and innovation_time is not None + t_obs = innovation_time[partial_zobs_batch] + h_obs = h.feats[partial_zobs_rows].to(dtype=gate.embed_y.weight.dtype) + x_embed = gate.embed_x(torch.cat([h_obs, t_obs], dim=-1)) + innovation = innovation_y - x_embed + gate_input = torch.cat([ + innovation, + innovation_gate_y[:, 32:34], + t_obs, + ], dim=-1) + gate_value = torch.sigmoid(gate.mlp(gate_input)) + if _gate_value_transform is not None: + gate_value = _gate_value_transform( + gate_value, partial_zobs_batch, _block_index + ) + assert gate_value.shape == (innovation.shape[0], 1), \ + f'gate value transform must return [Nobs,1], got {tuple(gate_value.shape)}' + delta = gate.proj(innovation) * gate_value + delta = delta * gate.alpha[_block_index].to(dtype=delta.dtype) + if _block_index == len(self.blocks) - 1: + # Keep the frozen V11 base carrier at the terminal site, + # but never add a newly learned observed-only shortcut: + # no attention block remains to transport/smooth it into U. + delta = delta * 0 + # Detached telemetry for the training log. This is before the legacy + # layer-scale hook and therefore measures the learned branch itself. + _innovation_ratio = delta.norm() / partial_zobs_inject.norm().clamp_min( + torch.finfo(partial_zobs_inject.dtype).eps + ) + _innovation_gate_mean = ( + gate_value.mean() if gate_value.numel() != 0 else delta.new_zeros(()) + ) + layer_inject = partial_zobs_inject + manual_cast(delta, self.dtype) + if partial_zobs_rows.numel() != 0: + if _partial_scales is None: + h = h.replace(h.feats.index_add( + 0, partial_zobs_rows, layer_inject + )) + else: + h = h.replace(h.feats.index_add( + 0, partial_zobs_rows, layer_inject, + alpha=float(_partial_scales[_block_index]), + )) + else: + # Preserve every empty carrier/gate GEMM in the graph on all-drop batches. + h = h.replace(h.feats + layer_inject.sum() * 0) + if _local_delta is not None: + _local_scale = 1.0 if _partial_scales is None else float( + _partial_scales[_block_index] + ) + if _local_topology['query_rows'].numel() != 0: + h = h.replace(h.feats.index_add( + 0, _local_topology['query_rows'], _local_delta, + alpha=_local_scale, + )) + else: + # The module returns an exact-zero parameter edge for all 14 + # trainable tensors, so DDP observes finite numerical-zero grads. + h = h.replace(h.feats + _local_delta.sum() * _local_scale) + if self.innovation_gate is not None: + if _innovation_ratio is None: + _innovation_ratio = h.feats.new_zeros(()) + _innovation_gate_mean = h.feats.new_zeros(()) + _innovation_ratio_stats.append(_innovation_ratio.detach()) + _innovation_gate_mean_stats.append(_innovation_gate_mean.detach()) + if _local_attention is not None and _block_index < _local_attention.alpha.numel(): + assert _local_stats is not None + _local_ratio_base_stats.append(_local_ratio_base.detach()) + _local_ratio_hidden_stats.append(_local_ratio_hidden.detach()) + _local_null_stats.append(_local_stats['null_attention_mean'].detach()) + _local_entropy_stats.append(_local_stats['real_neighbour_entropy_mean'].detach()) + _local_effective_neighbour_stats.append( + _local_stats['effective_real_neighbour_count_mean'].detach() + ) + _local_observed_residual_stats.append( + _local_observed_residual.detach() + ) + _local_unobserved_residual_stats.append( + _local_unobserved_residual.detach() + ) + + if self.innovation_gate is not None: + # Instance-only telemetry: deliberately neither buffer nor state-dict entry. + self.partial_zobs_innovation_gate_last_stats = { + 'delta_to_base_norm_ratio': torch.stack(_innovation_ratio_stats), + 'gate_mean': torch.stack(_innovation_gate_mean_stats), + 'observed_rows': _innovation_observed_rows, + } + if _local_attention is not None: + # Block 29 is structurally absent: every telemetry vector has length 29. + assert _local_query_is_observed is not None + assert _local_neighbourhood_histogram is not None + _full_support_rows = int(x.coords.shape[0]) + _active_queries = int(_local_topology['query_rows'].numel()) + _empty_queries = _full_support_rows - _active_queries + _singleton_queries = int( + (_local_topology['neighbour_counts'] == 1).sum() + ) + self.partial_zobs_local_attention_last_stats = { + 'delta_to_base_norm_ratio': torch.stack(_local_ratio_base_stats), + 'delta_to_hidden_norm_ratio': torch.stack(_local_ratio_hidden_stats), + 'null_attention_mean': torch.stack(_local_null_stats), + 'real_neighbour_entropy_mean': torch.stack(_local_entropy_stats), + 'effective_real_neighbour_count_mean': torch.stack( + _local_effective_neighbour_stats + ), + 'observed_residual_l2_norm': torch.stack( + _local_observed_residual_stats + ), + 'unobserved_residual_l2_norm': torch.stack( + _local_unobserved_residual_stats + ), + 'observed_rows': _innovation_observed_rows, + 'full_support_rows': _full_support_rows, + 'active_queries': _active_queries, + 'empty_queries': _empty_queries, + 'singleton_queries': _singleton_queries, + 'observed_active_queries': int(_local_query_is_observed.sum()), + 'unobserved_active_queries': int((~_local_query_is_observed).sum()), + 'active_query_fraction': _active_queries / max(_full_support_rows, 1), + 'empty_query_fraction': _empty_queries / max(_full_support_rows, 1), + 'singleton_query_fraction': _singleton_queries / max(_active_queries, 1), + 'neighbourhood_size_histogram': _local_neighbourhood_histogram.detach(), + 'packed_pairs': int(_local_topology['source_indices'].numel()), + } + if _learned_attention is not None: + assert _learned_cache is not None + assert len(_learned_stats) == len(_learned_attention.sites) + self.partial_zobs_learned_attention_last_stats = { + name: torch.stack([site[name] for site in _learned_stats]) + for name in _learned_stats[0] + } + self.partial_zobs_learned_attention_last_stats.update({ + 'sites': _learned_attention.sites, + 'observed_rows': int(partial_zobs_rows.numel()), + 'full_support_rows': int(x.coords.shape[0]), + 'observed_counts': _learned_cache['observed_counts'], + }) + + h = manual_cast(h, x.dtype) + h = h.replace(F.layer_norm(h.feats, h.feats.shape[-1:])) + h = self.out_layer(h) + return h + + +# Published Pixal3D geometry configs use this name. The former elastic mixin +# only controlled training memory; inference uses the same SLatFlowModel. +ElasticSLatFlowModel = SLatFlowModel diff --git a/deployment/runtime/geometry/pixal3d/modules/projected_image_feature_extractor.py b/deployment/runtime/geometry/pixal3d/modules/projected_image_feature_extractor.py new file mode 100644 index 0000000000000000000000000000000000000000..d532b1e781b02d44b9bbce0eb8e9fd092469e5de --- /dev/null +++ b/deployment/runtime/geometry/pixal3d/modules/projected_image_feature_extractor.py @@ -0,0 +1,458 @@ +"""Frozen DINOv3 features with camera-aligned projection for Pixal3D inference.""" +from typing import List, Optional, Tuple, Union + +import numpy as np +from PIL import Image +import torch +from torch import nn +import torch.nn.functional as F +from torchvision import transforms + + +def project_points_to_image_batch( + points_3d: torch.Tensor, + transform_matrix: torch.Tensor, + camera_angle_x: torch.Tensor, + resolution: int = 518 +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """ + Project 3D points to 2D image coordinates (batch processing). + + Args: + points_3d: torch.Tensor, shape [N, 3] or [B, N, 3], 3D point coordinates (in [-1, 1] range) + transform_matrix: torch.Tensor, shape [B, 4, 4], camera transformation matrix + camera_angle_x: torch.Tensor, shape [B], horizontal field of view angle (radians) + resolution: int, image resolution, default 518 + + Returns: + points_2d: torch.Tensor, shape [B, N, 2], image coordinates [x, y] + depth: torch.Tensor, shape [B, N], depth values + valid_mask: torch.Tensor, shape [B, N], mask for points within view + """ + device = points_3d.device + B = transform_matrix.shape[0] + + # Ensure inputs are torch.Tensor on correct device + if not isinstance(transform_matrix, torch.Tensor): + transform_matrix = torch.tensor(transform_matrix, dtype=torch.float32, device=device) + if not isinstance(points_3d, torch.Tensor): + points_3d = torch.tensor(points_3d, dtype=torch.float32, device=device) + if not isinstance(camera_angle_x, torch.Tensor): + camera_angle_x = torch.tensor(camera_angle_x, dtype=torch.float32, device=device) + + # Expand points_3d to batch dimension: [N, 3] -> [B, N, 3] + if points_3d.dim() == 2: + points_3d_batch = points_3d.unsqueeze(0).expand(B, -1, -1) + else: + points_3d_batch = points_3d + N = points_3d_batch.shape[1] + + # Add homogeneous coordinates: [B, N, 3] -> [B, N, 4] + ones = torch.ones(B, N, 1, device=device, dtype=points_3d_batch.dtype) + points_homogeneous = torch.cat([points_3d_batch, ones], dim=-1) # [B, N, 4] + + # Compute world to camera transformation matrix + world_to_camera = torch.linalg.inv(transform_matrix.float()).to(transform_matrix.dtype) # linalg.inv requires fp32+ + + # Batch transform to camera coordinate system: [B, N, 4] @ [B, 4, 4]^T -> [B, N, 3] + points_camera = torch.bmm(points_homogeneous, world_to_camera.transpose(-2, -1))[..., :3] # [B, N, 3] + + # Extract camera coordinates + x_cam = points_camera[..., 0] # [B, N] + y_cam = points_camera[..., 1] # [B, N] + z_cam = points_camera[..., 2] # [B, N] + + # Depth value (Z value in camera coordinate system, note Blender camera faces -Z direction) + depth = -z_cam # [B, N] + + # Compute camera intrinsics (batch processing) + sensor_width = 32.0 # mm + focal_length = 16.0 / torch.tan(camera_angle_x / 2.0) # [B] + focal_length_pixels = focal_length * resolution / sensor_width # [B] + + # Expand focal_length_pixels dimension for broadcasting: [B] -> [B, 1] + focal_length_pixels = focal_length_pixels.unsqueeze(1) # [B, 1] + + # Perspective projection to NDC coordinates + x_ndc = focal_length_pixels * x_cam / (-z_cam + 1e-8) # [B, N] + y_ndc = focal_length_pixels * y_cam / (-z_cam + 1e-8) # [B, N] + + # Convert to image coordinates (pixel coordinates) + x_pixel = x_ndc + resolution / 2.0 # [B, N] + y_pixel = -y_ndc + resolution / 2.0 # [B, N], flip Y axis + + # Create validity mask (points within image range and in front of camera) + valid_mask = ( + (x_pixel >= 0) & (x_pixel < resolution) & + (y_pixel >= 0) & (y_pixel < resolution) & + (depth > 0) # In front of camera + ) # [B, N] + + points_2d = torch.stack([x_pixel, y_pixel], dim=-1) # [B, N, 2] + + return points_2d, depth, valid_mask + +def sample_features(fmap: torch.Tensor, queries_ndc: torch.Tensor) -> torch.Tensor: + """ + Sample features from feature map at specified NDC coordinates. + + Args: + fmap: torch.Tensor, shape [B, C, H, W], feature map + queries_ndc: torch.Tensor, shape [B, K, 2], normalized device coordinates + + Returns: + torch.Tensor, shape [B, C, K], sampled features + """ + B, C, H, W = fmap.shape + Bq, K, _ = queries_ndc.shape + assert Bq == B, "Batch size mismatch" + + # grid_sample requires (B, out_h, out_w, 2), here we want K points -> out_h=K, out_w=1 + grid = queries_ndc.view(B, K, 1, 2) # (B, K, 1, 2) + + # Bilinear interpolation, align_corners=False (consistent with [-1,1] pixel center convention) + feat = F.grid_sample( + fmap, grid, mode='bilinear', + align_corners=False, padding_mode='border' # border avoids out-of-bound becoming 0 + ) # (B, C, K, 1) + + return feat.squeeze(-1) # (B, C, K) + +class ProjGrid(nn.Module): + """ + 3D Grid Projection Module. + + Projects a 3D grid of points to 2D image coordinates and samples features + from the image feature map at those locations. + + This is the core module for view-aligned feature extraction. + """ + def __init__(self, grid_resolution: int = 16, image_resolution: int = 518): + super().__init__() + self.grid_resolution = grid_resolution + self.image_resolution = image_resolution + + # Create 3D grid points + one_dim = torch.linspace(-1, 1, grid_resolution) + x, y, z = torch.meshgrid(one_dim, one_dim, one_dim, indexing='ij') + grid_points = torch.stack((x, y, z), dim=-1) + + # Rotation matrix to align with Blender coordinate system + rotation_matrix = torch.tensor([ + [1.0, 0.0, 0.0], + [0.0, 0.0, -1.0], + [0.0, 1.0, 0.0] + ]) + grid_points = torch.matmul(grid_points, rotation_matrix.T) + grid_points = grid_points.reshape(-1, 3) + self.register_buffer('grid_points', grid_points) # [R³, 3] + + # Default front view transformation matrix + front_view_transform_matrix = torch.tensor([ + [1.0, 0.0, 0.0, 0.0], + [0.0, 0.0, -1.0, -2.0], + [0.0, 1.0, 0.0, 0.0], + [0.0, 0.0, 0.0, 1.0] + ]) + self.register_buffer("front_view_transform_matrix", front_view_transform_matrix) + + def forward( + self, + features_map: torch.Tensor, + camera_angle_x: torch.Tensor, + distance: torch.Tensor, + mesh_scale: torch.Tensor, + transform_matrix: Optional[torch.Tensor] = None, + BHWC: bool = True + ) -> torch.Tensor: + """ + Project 3D grid points to image and sample features. + + Args: + features_map: Feature map, shape [B, H, W, C] if BHWC else [B, C, H, W] + camera_angle_x: Camera FOV angle, shape [B] + distance: Camera distance, shape [B] + mesh_scale: Mesh scale factor, shape [B] + transform_matrix: Optional camera transform matrix, shape [B, 4, 4] + BHWC: Whether features_map is in BHWC format + + Returns: + Projected features, shape [B, grid_resolution³, C] + """ + if BHWC: + B, H, W, C = features_map.shape + else: + B, C, H, W = features_map.shape + + grid_points = self.grid_points + grid_points = grid_points.expand(B, -1, -1) + grid_points = grid_points / mesh_scale.unsqueeze(-1).unsqueeze(-1) / 2 # Scale alignment + assert transform_matrix is None, "transform_matrix is not None" + if transform_matrix is None: + transform_matrix = self.front_view_transform_matrix + transform_matrix = transform_matrix.expand(B, -1, -1).clone() + transform_matrix[:, 1, 3] = -distance # Set camera distance + + # Project to image coordinates (simulate Blender projection) + image_points, depth, valid_mask = project_points_to_image_batch( + grid_points, transform_matrix, camera_angle_x, self.image_resolution + ) + + # Normalize to [-1, 1] for grid_sample + image_points_norm = (image_points + 0.5) / self.image_resolution * 2 - 1 + + if BHWC: + features_map = features_map.permute(0, 3, 1, 2) # [B, C, H, W] + + # Sample features from DINOv3 patch feature map + x = sample_features(features_map, image_points_norm) # [B, C, K] + x = x.permute(0, 2, 1) # [B, K, C] + + return x + + +class DinoV3ProjFeatureExtractor(nn.Module): + """ + DINOv3 Feature Extractor with View-Aligned Projection. + + This extractor produces both: + 1. Global features (CLS token + register tokens) in embed_dim + 2. View-aligned projected features (3D grid projected to 2D and sampled) + - Without NAF: [B, R³, embed_dim] + - With NAF: [B, R³, embed_dim * 2] (concat of lr and hr features) + + NOTE: proj_linear has been moved to per-block ProjectAttention / SparseProjectAttention. + This module now outputs raw DINOv3 features for proj (optionally concatenated with NAF-upsampled features). + + Args: + model_name: Name of the pretrained DINOv3 model + image_size: Input image size (default: 512) + grid_resolution: Resolution of the 3D projection grid (default: 16) + use_naf_upsample: Whether to use NAF to upsample features (default: False) + naf_target_size: Target spatial size for NAF upsampling (default: [128, 128]) + """ + def __init__( + self, + model_name: str, + image_size: int = 512, + grid_resolution: int = 16, + use_naf_upsample: bool = False, + naf_target_size: Optional[List[int]] = None, + cls_only: bool = False, + cls_n_tokens: int = 1, + full_tokens: bool = False, + ): + super().__init__() + self.model_name = model_name + self.image_size = image_size + self.grid_resolution = grid_resolution + self.use_naf_upsample = use_naf_upsample + self.cls_only = cls_only + self.cls_n_tokens = cls_n_tokens + self.full_tokens = full_tokens + if naf_target_size is None: + self.naf_target_size = (128, 128) + elif isinstance(naf_target_size, int): + self.naf_target_size = (naf_target_size, naf_target_size) + else: + self.naf_target_size = tuple(naf_target_size) + + # Load DINOv3 model (frozen, no trainable params in this module) + from transformers import DINOv3ViTModel + + self.model = DINOv3ViTModel.from_pretrained(model_name) + self.model.eval() + self.model.requires_grad_(False) + + # Image transform (only normalize, no resize - assume already resized) + self.transform = transforms.Compose([ + transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), + ]) + + # Get patch info + self.patch_size = self.model.config.patch_size + self.patch_number = image_size // self.patch_size + self.embed_dim = self.model.config.hidden_size + + # Projection grid for view-aligned features + self.proj_grid = ProjGrid( + grid_resolution=grid_resolution, + image_resolution=image_size + ) + + # NAF upsampler (frozen, no trainable params) + self.naf_model = None # Lazy-loaded on first use to avoid import if not needed + + # proj_channels: the output dimension of proj features + # Without NAF: embed_dim (e.g. 1024) + # With NAF: embed_dim * 2 (e.g. 2048, concat of lr and hr) + self.proj_channels = self.embed_dim * 2 if use_naf_upsample else self.embed_dim + + # NOTE: proj_linear removed — now lives in each denoiser block's ProjectAttention + + def _load_naf(self): + """Lazy-load pretrained NAF model.""" + if self.naf_model is None: + import torch.hub + device = next(self.model.parameters()).device + self.naf_model = torch.hub.load( + "valeoai/NAF", "naf", pretrained=True, device=device, trust_repo=True + ) + self.naf_model.eval() + self.naf_model.requires_grad_(False) + + def to(self, device): + super().to(device) + self.model.to(device) + self.proj_grid.to(device) + if self.naf_model is not None: + self.naf_model.to(device) + return self + + def cuda(self): + super().cuda() + self.model.cuda() + self.proj_grid.cuda() + if self.naf_model is not None: + self.naf_model.cuda() + return self + + def cpu(self): + super().cpu() + self.model.cpu() + self.proj_grid.cpu() + if self.naf_model is not None: + self.naf_model.cpu() + return self + + def extract_features(self, image: torch.Tensor) -> torch.Tensor: + """Extract features using DINOv3.""" + image = image.to(self.model.embeddings.patch_embeddings.weight.dtype) + hidden_states = self.model.embeddings(image, bool_masked_pos=None) + position_embeddings = self.model.rope_embeddings(image) + + for layer_module in self.model.layer: + hidden_states = layer_module( + hidden_states, + position_embeddings=position_embeddings, + ) + + return F.layer_norm(hidden_states, hidden_states.shape[-1:]) + + def forward( + self, + image: Union[torch.Tensor, List[Image.Image]], + camera_angle_x: Optional[torch.Tensor] = None, + distance: Optional[torch.Tensor] = None, + mesh_scale: Optional[torch.Tensor] = None, + transform_matrix: Optional[torch.Tensor] = None, + ) -> Tuple[torch.Tensor, torch.Tensor]: + """ + Extract view-aligned features from the image. + + Args: + image: Input image tensor [B, C, H, W] or list of PIL images + camera_angle_x: Camera FOV angle in radians [B] + distance: Camera distance [B] + mesh_scale: Mesh scale factor [B] + transform_matrix: Optional camera transform matrix [B, 4, 4] + + Returns: + Tuple of (global_features, proj_features): + - global_features: [B, num_global_tokens, embed_dim] + - proj_features: [B, grid_resolution³, proj_channels] + where proj_channels = embed_dim (no NAF) or embed_dim*2 (with NAF) + """ + # Handle input types + if isinstance(image, torch.Tensor): + assert image.ndim == 4, "Image tensor should be batched (B, C, H, W)" + elif isinstance(image, list): + assert all(isinstance(i, Image.Image) for i in image), "Image list should be list of PIL images" + image = [i.resize((self.image_size, self.image_size), Image.LANCZOS) for i in image] + image = [np.array(i.convert('RGB')).astype(np.float32) / 255 for i in image] + image = [torch.from_numpy(i).permute(2, 0, 1).float() for i in image] + image = torch.stack(image).cuda() + else: + raise ValueError(f"Unsupported type of image: {type(image)}") + + B = image.shape[0] + + # Trained-bridge CLS cross-attn mode: skip ProjGrid + NAF (no camera); return global token(s). + # cls_n_tokens=1 -> CLS only (v4 default); =5 (1 CLS + 4 register) -> faithful warm start of the + # shipped cross-attn (trained on [CLS + register] global tokens; registers are global, NOT patch). + # full_tokens=True (v6 patchcross) -> the UNSLICED [CLS|4 reg|patch...] sequence + # (1029 tokens @ image 512), i.e. the TRELLIS.2 cross-attn context. Lives in this branch + # because the non-cls path below requires camera params for ProjGrid. + if getattr(self, 'cls_only', False): + img = self.transform(image) + with torch.no_grad(): + z = self.extract_features(img) + if getattr(self, 'full_tokens', False): + return z # [B, 1+4+patch_number**2, D=1024] + return z[:, 0:self.cls_n_tokens] # [B, cls_n_tokens, D=1024] + + # Keep a copy of the unnormalized image for NAF guide + if self.use_naf_upsample: + image_for_naf = image.clone() # [B, 3, H, W], in [0, 1] range + + # Apply transform (ImageNet normalization) + image = self.transform(image) + + # Extract DINOv3 features (frozen, no gradients) + with torch.no_grad(): + z = self.extract_features(image) + + # Split into CLS token, register tokens, and patch tokens + z_clstoken = z[:, 0:1] # [B, 1, D] + num_reg = getattr(self.model.config, 'num_register_tokens', 4) + z_regtokens = z[:, 1:1+num_reg] # [B, num_reg, D] + z_patchtokens = z[:, 1+num_reg:] # [B, num_patches, D] + + # Reshape patch tokens to spatial grid: [B, h, w, D] + z_patchtokens_spatial = z_patchtokens.reshape( + B, self.patch_number, self.patch_number, -1 + ) # [B, h, w, D] + + if camera_angle_x is None or distance is None or mesh_scale is None: + raise ValueError("camera_angle_x, distance, and mesh_scale must be provided") + + # --- Low-resolution branch: sample from DINOv3 patch feature map --- + z_proj_lr = self.proj_grid( + z_patchtokens_spatial, + camera_angle_x, + distance, + mesh_scale, + transform_matrix + ) # [B, grid_res³, D] + + # --- High-resolution branch (NAF): upsample then sample --- + if self.use_naf_upsample: + self._load_naf() + # NAF expects: guide [B, 3, H, W], lr_features [B, C, h, w], target_size (H', W') + lr_features_bchw = z_patchtokens_spatial.permute(0, 3, 1, 2) # [B, D, h, w] + hr_features = self.naf_model( + image_for_naf, lr_features_bchw, self.naf_target_size + ) # [B, D, H', W'] + + # Sample from high-res feature map using same projection coordinates + z_proj_hr = self.proj_grid( + hr_features, + camera_angle_x, + distance, + mesh_scale, + transform_matrix, + BHWC=False # hr_features is [B, C, H', W'] + ) # [B, grid_res³, D] + + # Concatenate lr and hr: [B, grid_res³, D*2] + z_proj = torch.cat([z_proj_lr, z_proj_hr], dim=-1) + else: + z_proj = z_proj_lr # [B, grid_res³, D] + + # Combine global tokens + z_global = torch.cat([z_clstoken, z_regtokens], dim=1) # [B, 1+num_reg, D] + + # proj_linear has been moved to per-block ProjectAttention + # z_proj stays in proj_channels, each block will project independently + + return z_global, z_proj diff --git a/deployment/runtime/geometry/pixal3d/pipelines/__init__.py b/deployment/runtime/geometry/pixal3d/pipelines/__init__.py index de01583213347ec0172b5d2ab7c21cd03474136e..b882ce968b6cbc6eee9093e173a14d3a1e26c6ce 100644 --- a/deployment/runtime/geometry/pixal3d/pipelines/__init__.py +++ b/deployment/runtime/geometry/pixal3d/pipelines/__init__.py @@ -1,57 +1,55 @@ -import importlib - -__attributes = { - "Trellis2ImageTo3DPipeline": "trellis2_image_to_3d", - "Trellis2TexturingPipeline": "trellis2_texturing", - "Pixal3DImageTo3DPipeline": "pixal3d_image_to_3d", - "build_v14fx_inference_model": "v14fx", - "sample_v14fx_texture": "v14fx", -} - - -__submodules = ['samplers', 'rembg'] - -__all__ = list(__attributes.keys()) + __submodules - -def __getattr__(name): - if name not in globals(): - if name in __attributes: - module_name = __attributes[name] - module = importlib.import_module(f".{module_name}", __name__) - globals()[name] = getattr(module, name) - elif name in __submodules: - module = importlib.import_module(f".{name}", __name__) - globals()[name] = module - else: - raise AttributeError(f"module {__name__} has no attribute {name}") - return globals()[name] - - -def from_pretrained(path: str): - """ - Load a pipeline from a model folder or a Hugging Face model hub. - - Args: - path: The path to the model. Can be either local path or a Hugging Face model name. - """ - import os - import json - is_local = os.path.exists(f"{path}/pipeline.json") - - if is_local: - config_file = f"{path}/pipeline.json" - else: - from huggingface_hub import hf_hub_download - config_file = hf_hub_download(path, "pipeline.json") - - with open(config_file, 'r') as f: - config = json.load(f) - return globals()[config['name']].from_pretrained(path) - - -# For PyLance -if __name__ == '__main__': - from . import samplers, rembg - from .trellis2_image_to_3d import Trellis2ImageTo3DPipeline - from .trellis2_texturing import Trellis2TexturingPipeline - from .pixal3d_image_to_3d import Pixal3DImageTo3DPipeline +import importlib + +__attributes = { + "Trellis2ImageTo3DPipeline": "trellis2_image_to_3d", + "Trellis2TexturingPipeline": "trellis2_texturing", + "Pixal3DImageTo3DPipeline": "pixal3d_image_to_3d", +} + + +__submodules = ['samplers', 'rembg'] + +__all__ = list(__attributes.keys()) + __submodules + +def __getattr__(name): + if name not in globals(): + if name in __attributes: + module_name = __attributes[name] + module = importlib.import_module(f".{module_name}", __name__) + globals()[name] = getattr(module, name) + elif name in __submodules: + module = importlib.import_module(f".{name}", __name__) + globals()[name] = module + else: + raise AttributeError(f"module {__name__} has no attribute {name}") + return globals()[name] + + +def from_pretrained(path: str): + """ + Load a pipeline from a model folder or a Hugging Face model hub. + + Args: + path: The path to the model. Can be either local path or a Hugging Face model name. + """ + import os + import json + is_local = os.path.exists(f"{path}/pipeline.json") + + if is_local: + config_file = f"{path}/pipeline.json" + else: + from huggingface_hub import hf_hub_download + config_file = hf_hub_download(path, "pipeline.json") + + with open(config_file, 'r') as f: + config = json.load(f) + return globals()[config['name']].from_pretrained(path) + + +# For PyLance +if __name__ == '__main__': + from . import samplers, rembg + from .trellis2_image_to_3d import Trellis2ImageTo3DPipeline + from .trellis2_texturing import Trellis2TexturingPipeline + from .pixal3d_image_to_3d import Pixal3DImageTo3DPipeline diff --git a/deployment/runtime/geometry/pixal3d/pipelines/rembg/RMBG2.py b/deployment/runtime/geometry/pixal3d/pipelines/rembg/RMBG2.py index 1a809226fc3da4607f956f34052d751bb2cf205d..13cfc519256518e04c403d1bb85757e7e92e1218 100644 --- a/deployment/runtime/geometry/pixal3d/pipelines/rembg/RMBG2.py +++ b/deployment/runtime/geometry/pixal3d/pipelines/rembg/RMBG2.py @@ -1,64 +1,59 @@ -from typing import Optional -import os -from pathlib import Path - -from transformers import AutoModelForImageSegmentation -import torch -from torchvision import transforms -from PIL import Image - - -_CY_ROOT = os.environ.get("CY_ROOT", str(Path(__file__).parents[5])) -RMBG2_MODEL_PATH = os.environ.get( - "PIXAL3D_RMBG_MODEL_PATH", - os.path.join(_CY_ROOT, "model_ckpts", "RMBG-2.0"), -) -_REMOTE_MODEL_NAMES = {"ZhengPeng7/BiRefNet", "briaai/RMBG-2.0"} - - -class RMBG2: - def __init__( - self, - model_name: str = RMBG2_MODEL_PATH, - alpha_threshold: Optional[float] = None, - ): - if model_name in _REMOTE_MODEL_NAMES: - model_name = RMBG2_MODEL_PATH - if alpha_threshold is not None and not 0.0 <= alpha_threshold <= 1.0: - raise ValueError(f"alpha_threshold must be in [0, 1], got {alpha_threshold}") - self.model_name = model_name - self.alpha_threshold = alpha_threshold - self.model = AutoModelForImageSegmentation.from_pretrained( - model_name, trust_remote_code=True - ) - self.model.eval() - self.transform_image = transforms.Compose( - [ - transforms.Resize((1024, 1024)), - transforms.ToTensor(), - transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]), - ] - ) - - def to(self, device: str): - self.model.to(device) - - def cuda(self): - self.model.cuda() - - def cpu(self): - self.model.cpu() - - def __call__(self, image: Image.Image) -> Image.Image: - image_size = image.size - input_images = self.transform_image(image).unsqueeze(0).to("cuda") - with torch.no_grad(): - preds = self.model(input_images)[-1].sigmoid().cpu() - pred = preds[0].squeeze() - pred_pil = transforms.ToPILImage()(pred) - mask = pred_pil.resize(image_size) - if self.alpha_threshold is not None: - cutoff = int(round(self.alpha_threshold * 255)) - mask = mask.point([0 if value < cutoff else 255 for value in range(256)]) - image.putalpha(mask) - return image +from typing import Optional +import os + +from transformers import AutoModelForImageSegmentation +import torch +from torchvision import transforms +from PIL import Image + + +# A local RMBG-2.0 or BiRefNet directory can replace the Hub model. +RMBG2_MODEL_PATH = os.environ.get("PIXAL3D_RMBG_MODEL_PATH", "briaai/RMBG-2.0") + + +class RMBG2: + def __init__( + self, + model_name: str = RMBG2_MODEL_PATH, + alpha_threshold: Optional[float] = None, + ): + if model_name == "briaai/RMBG-2.0": + model_name = RMBG2_MODEL_PATH + if alpha_threshold is not None and not 0.0 <= alpha_threshold <= 1.0: + raise ValueError(f"alpha_threshold must be in [0, 1], got {alpha_threshold}") + self.model_name = model_name + self.alpha_threshold = alpha_threshold + self.model = AutoModelForImageSegmentation.from_pretrained( + model_name, trust_remote_code=True + ) + self.model.eval() + self.transform_image = transforms.Compose( + [ + transforms.Resize((1024, 1024)), + transforms.ToTensor(), + transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]), + ] + ) + + def to(self, device: str): + self.model.to(device) + + def cuda(self): + self.model.cuda() + + def cpu(self): + self.model.cpu() + + def __call__(self, image: Image.Image) -> Image.Image: + image_size = image.size + input_images = self.transform_image(image).unsqueeze(0).to("cuda") + with torch.no_grad(): + preds = self.model(input_images)[-1].sigmoid().cpu() + pred = preds[0].squeeze() + pred_pil = transforms.ToPILImage()(pred) + mask = pred_pil.resize(image_size) + if self.alpha_threshold is not None: + cutoff = int(round(self.alpha_threshold * 255)) + mask = mask.point([0 if value < cutoff else 255 for value in range(256)]) + image.putalpha(mask) + return image diff --git a/deployment/runtime/geometry/pixal3d/representations/mesh/base.py b/deployment/runtime/geometry/pixal3d/representations/mesh/base.py index 23bcf6d14ebac5b08893579fb7dd834a11c0d9a1..bff3d8eaae4bf679b8ba099808a83761e2bf2dbb 100644 --- a/deployment/runtime/geometry/pixal3d/representations/mesh/base.py +++ b/deployment/runtime/geometry/pixal3d/representations/mesh/base.py @@ -1,239 +1,239 @@ -from typing import * -import torch -from ..voxel import Voxel -import cumesh -# NOTE: flex_gemm is imported lazily in query_attrs (matching the renderers' pattern), NOT here. -# This module sits in the DataLoader workers' import closure (dataset -> MeshWithVoxel), and -# flex_gemm's triton @autotune decorators query the GPU driver AT IMPORT — in a spawned CPU-only -# worker that intermittently dies with "RuntimeError: 0 active drivers", killing the rank and -# wedging all other ranks in NCCL collectives (observed 2026-07-02 on the 64-rank 512-all run). - - -class Mesh: - def __init__(self, - vertices, - faces, - vertex_attrs=None - ): - self.vertices = vertices.float() - self.faces = faces.int() - self.vertex_attrs = vertex_attrs - - @property - def device(self): - return self.vertices.device - - def to(self, device, non_blocking=False): - return Mesh( - self.vertices.to(device, non_blocking=non_blocking), - self.faces.to(device, non_blocking=non_blocking), - self.vertex_attrs.to(device, non_blocking=non_blocking) if self.vertex_attrs is not None else None, - ) - - def cuda(self, non_blocking=False): - return self.to('cuda', non_blocking=non_blocking) - - def cpu(self): - return self.to('cpu') - - def fill_holes(self, max_hole_perimeter=3e-2): - vertices = self.vertices.clone().cuda().contiguous() - faces = self.faces.clone().cuda().contiguous() - - mesh = cumesh.CuMesh() - mesh.init(vertices, faces) - mesh.get_edges() - mesh.get_boundary_info() - if mesh.num_boundaries == 0: - return - mesh.get_vertex_edge_adjacency() - mesh.get_vertex_boundary_adjacency() - mesh.get_manifold_boundary_adjacency() - mesh.read_manifold_boundary_adjacency() - mesh.get_boundary_connected_components() - mesh.get_boundary_loops() - if mesh.num_boundary_loops == 0: - return - mesh.fill_holes(max_hole_perimeter=max_hole_perimeter) - new_vertices, new_faces = mesh.read() - - self.vertices = new_vertices.to(self.device) - self.faces = new_faces.to(self.device) - - def remove_faces(self, face_mask: torch.Tensor): - vertices = self.vertices.clone().cuda().contiguous() - faces = self.faces.clone().cuda().contiguous() - - mesh = cumesh.CuMesh() - mesh.init(vertices, faces) - mesh.remove_faces(face_mask) - new_vertices, new_faces = mesh.read() - - self.vertices = new_vertices.to(self.device) - self.faces = new_faces.to(self.device) - - def simplify(self, target=1000000, verbose: bool=False, options: dict={}): - vertices = self.vertices.clone().cuda().contiguous() - faces = self.faces.clone().cuda().contiguous() - - mesh = cumesh.CuMesh() - mesh.init(vertices, faces) - mesh.simplify(target, verbose=verbose, options=options) - new_vertices, new_faces = mesh.read() - - self.vertices = new_vertices.to(self.device) - self.faces = new_faces.to(self.device) - - -class TextureFilterMode: - CLOSEST = 0 - LINEAR = 1 - - -class TextureWrapMode: - CLAMP_TO_EDGE = 0 - REPEAT = 1 - MIRRORED_REPEAT = 2 - - -class AlphaMode: - OPAQUE = 0 - MASK = 1 - BLEND = 2 - - -class Texture: - def __init__( - self, - image: torch.Tensor, - filter_mode: TextureFilterMode = TextureFilterMode.LINEAR, - wrap_mode: TextureWrapMode = TextureWrapMode.REPEAT - ): - self.image = image - self.filter_mode = filter_mode - self.wrap_mode = wrap_mode - - def to(self, device, non_blocking=False): - return Texture( - self.image.to(device, non_blocking=non_blocking), - self.filter_mode, - self.wrap_mode, - ) - - -class PbrMaterial: - def __init__( - self, - base_color_texture: Optional[Texture] = None, - base_color_factor: Union[torch.Tensor, List[float]] = [1.0, 1.0, 1.0], - metallic_texture: Optional[Texture] = None, - metallic_factor: float = 1.0, - roughness_texture: Optional[Texture] = None, - roughness_factor: float = 1.0, - alpha_texture: Optional[Texture] = None, - alpha_factor: float = 1.0, - alpha_mode: AlphaMode = AlphaMode.OPAQUE, - alpha_cutoff: float = 0.5, - ): - self.base_color_texture = base_color_texture - self.base_color_factor = torch.tensor(base_color_factor, dtype=torch.float32)[:3] - self.metallic_texture = metallic_texture - self.metallic_factor = metallic_factor - self.roughness_texture = roughness_texture - self.roughness_factor = roughness_factor - self.alpha_texture = alpha_texture - self.alpha_factor = alpha_factor - self.alpha_mode = alpha_mode - self.alpha_cutoff = alpha_cutoff - - def to(self, device, non_blocking=False): - return PbrMaterial( - base_color_texture=self.base_color_texture.to(device, non_blocking=non_blocking) if self.base_color_texture is not None else None, - base_color_factor=self.base_color_factor.to(device, non_blocking=non_blocking), - metallic_texture=self.metallic_texture.to(device, non_blocking=non_blocking) if self.metallic_texture is not None else None, - metallic_factor=self.metallic_factor, - roughness_texture=self.roughness_texture.to(device, non_blocking=non_blocking) if self.roughness_texture is not None else None, - roughness_factor=self.roughness_factor, - alpha_texture=self.alpha_texture.to(device, non_blocking=non_blocking) if self.alpha_texture is not None else None, - alpha_factor=self.alpha_factor, - alpha_mode=self.alpha_mode, - alpha_cutoff=self.alpha_cutoff, - ) - - -class MeshWithPbrMaterial(Mesh): - def __init__(self, - vertices, - faces, - material_ids, - uv_coords, - materials: List[PbrMaterial], - ): - self.vertices = vertices.float() - self.faces = faces.int() - self.material_ids = material_ids # [M] - self.uv_coords = uv_coords # [M, 3, 2] - self.materials = materials - self.layout = { - 'base_color': slice(0, 3), - 'metallic': slice(3, 4), - 'roughness': slice(4, 5), - 'alpha': slice(5, 6), - } - - def to(self, device, non_blocking=False): - return MeshWithPbrMaterial( - self.vertices.to(device, non_blocking=non_blocking), - self.faces.to(device, non_blocking=non_blocking), - self.material_ids.to(device, non_blocking=non_blocking), - self.uv_coords.to(device, non_blocking=non_blocking), - [material.to(device, non_blocking=non_blocking) for material in self.materials], - ) - - -class MeshWithVoxel(Mesh, Voxel): - def __init__(self, - vertices: torch.Tensor, - faces: torch.Tensor, - origin: list, - voxel_size: float, - coords: torch.Tensor, - attrs: torch.Tensor, - voxel_shape: torch.Size, - layout: Dict = {}, - ): - self.vertices = vertices.float() - self.faces = faces.int() - self.origin = torch.tensor(origin, dtype=torch.float32, device=self.device) - self.voxel_size = voxel_size - self.coords = coords - self.attrs = attrs - self.voxel_shape = voxel_shape - self.layout = layout - - def to(self, device, non_blocking=False): - return MeshWithVoxel( - self.vertices.to(device, non_blocking=non_blocking), - self.faces.to(device, non_blocking=non_blocking), - self.origin.tolist(), - self.voxel_size, - self.coords.to(device, non_blocking=non_blocking), - self.attrs.to(device, non_blocking=non_blocking), - self.voxel_shape, - self.layout, - ) - - def query_attrs(self, xyz): - from flex_gemm.ops.grid_sample import grid_sample_3d - grid = ((xyz - self.origin) / self.voxel_size).reshape(1, -1, 3) - vertex_attrs = grid_sample_3d( - self.attrs, - torch.cat([torch.zeros_like(self.coords[..., :1]), self.coords], dim=-1), - self.voxel_shape, - grid, - mode='trilinear' - )[0] - return vertex_attrs - - def query_vertex_attrs(self): - return self.query_attrs(self.vertices) +from typing import * +import torch +from ..voxel import Voxel +import cumesh +# NOTE: flex_gemm is imported lazily in query_attrs (matching the renderers' pattern), NOT here. +# This module sits in the DataLoader workers' import closure (dataset -> MeshWithVoxel), and +# flex_gemm's triton @autotune decorators query the GPU driver AT IMPORT — in a spawned CPU-only +# worker this can raise "RuntimeError: 0 active drivers", killing the rank and +# leaving other ranks blocked in NCCL collectives. + + +class Mesh: + def __init__(self, + vertices, + faces, + vertex_attrs=None + ): + self.vertices = vertices.float() + self.faces = faces.int() + self.vertex_attrs = vertex_attrs + + @property + def device(self): + return self.vertices.device + + def to(self, device, non_blocking=False): + return Mesh( + self.vertices.to(device, non_blocking=non_blocking), + self.faces.to(device, non_blocking=non_blocking), + self.vertex_attrs.to(device, non_blocking=non_blocking) if self.vertex_attrs is not None else None, + ) + + def cuda(self, non_blocking=False): + return self.to('cuda', non_blocking=non_blocking) + + def cpu(self): + return self.to('cpu') + + def fill_holes(self, max_hole_perimeter=3e-2): + vertices = self.vertices.clone().cuda().contiguous() + faces = self.faces.clone().cuda().contiguous() + + mesh = cumesh.CuMesh() + mesh.init(vertices, faces) + mesh.get_edges() + mesh.get_boundary_info() + if mesh.num_boundaries == 0: + return + mesh.get_vertex_edge_adjacency() + mesh.get_vertex_boundary_adjacency() + mesh.get_manifold_boundary_adjacency() + mesh.read_manifold_boundary_adjacency() + mesh.get_boundary_connected_components() + mesh.get_boundary_loops() + if mesh.num_boundary_loops == 0: + return + mesh.fill_holes(max_hole_perimeter=max_hole_perimeter) + new_vertices, new_faces = mesh.read() + + self.vertices = new_vertices.to(self.device) + self.faces = new_faces.to(self.device) + + def remove_faces(self, face_mask: torch.Tensor): + vertices = self.vertices.clone().cuda().contiguous() + faces = self.faces.clone().cuda().contiguous() + + mesh = cumesh.CuMesh() + mesh.init(vertices, faces) + mesh.remove_faces(face_mask) + new_vertices, new_faces = mesh.read() + + self.vertices = new_vertices.to(self.device) + self.faces = new_faces.to(self.device) + + def simplify(self, target=1000000, verbose: bool=False, options: dict={}): + vertices = self.vertices.clone().cuda().contiguous() + faces = self.faces.clone().cuda().contiguous() + + mesh = cumesh.CuMesh() + mesh.init(vertices, faces) + mesh.simplify(target, verbose=verbose, options=options) + new_vertices, new_faces = mesh.read() + + self.vertices = new_vertices.to(self.device) + self.faces = new_faces.to(self.device) + + +class TextureFilterMode: + CLOSEST = 0 + LINEAR = 1 + + +class TextureWrapMode: + CLAMP_TO_EDGE = 0 + REPEAT = 1 + MIRRORED_REPEAT = 2 + + +class AlphaMode: + OPAQUE = 0 + MASK = 1 + BLEND = 2 + + +class Texture: + def __init__( + self, + image: torch.Tensor, + filter_mode: TextureFilterMode = TextureFilterMode.LINEAR, + wrap_mode: TextureWrapMode = TextureWrapMode.REPEAT + ): + self.image = image + self.filter_mode = filter_mode + self.wrap_mode = wrap_mode + + def to(self, device, non_blocking=False): + return Texture( + self.image.to(device, non_blocking=non_blocking), + self.filter_mode, + self.wrap_mode, + ) + + +class PbrMaterial: + def __init__( + self, + base_color_texture: Optional[Texture] = None, + base_color_factor: Union[torch.Tensor, List[float]] = [1.0, 1.0, 1.0], + metallic_texture: Optional[Texture] = None, + metallic_factor: float = 1.0, + roughness_texture: Optional[Texture] = None, + roughness_factor: float = 1.0, + alpha_texture: Optional[Texture] = None, + alpha_factor: float = 1.0, + alpha_mode: AlphaMode = AlphaMode.OPAQUE, + alpha_cutoff: float = 0.5, + ): + self.base_color_texture = base_color_texture + self.base_color_factor = torch.tensor(base_color_factor, dtype=torch.float32)[:3] + self.metallic_texture = metallic_texture + self.metallic_factor = metallic_factor + self.roughness_texture = roughness_texture + self.roughness_factor = roughness_factor + self.alpha_texture = alpha_texture + self.alpha_factor = alpha_factor + self.alpha_mode = alpha_mode + self.alpha_cutoff = alpha_cutoff + + def to(self, device, non_blocking=False): + return PbrMaterial( + base_color_texture=self.base_color_texture.to(device, non_blocking=non_blocking) if self.base_color_texture is not None else None, + base_color_factor=self.base_color_factor.to(device, non_blocking=non_blocking), + metallic_texture=self.metallic_texture.to(device, non_blocking=non_blocking) if self.metallic_texture is not None else None, + metallic_factor=self.metallic_factor, + roughness_texture=self.roughness_texture.to(device, non_blocking=non_blocking) if self.roughness_texture is not None else None, + roughness_factor=self.roughness_factor, + alpha_texture=self.alpha_texture.to(device, non_blocking=non_blocking) if self.alpha_texture is not None else None, + alpha_factor=self.alpha_factor, + alpha_mode=self.alpha_mode, + alpha_cutoff=self.alpha_cutoff, + ) + + +class MeshWithPbrMaterial(Mesh): + def __init__(self, + vertices, + faces, + material_ids, + uv_coords, + materials: List[PbrMaterial], + ): + self.vertices = vertices.float() + self.faces = faces.int() + self.material_ids = material_ids # [M] + self.uv_coords = uv_coords # [M, 3, 2] + self.materials = materials + self.layout = { + 'base_color': slice(0, 3), + 'metallic': slice(3, 4), + 'roughness': slice(4, 5), + 'alpha': slice(5, 6), + } + + def to(self, device, non_blocking=False): + return MeshWithPbrMaterial( + self.vertices.to(device, non_blocking=non_blocking), + self.faces.to(device, non_blocking=non_blocking), + self.material_ids.to(device, non_blocking=non_blocking), + self.uv_coords.to(device, non_blocking=non_blocking), + [material.to(device, non_blocking=non_blocking) for material in self.materials], + ) + + +class MeshWithVoxel(Mesh, Voxel): + def __init__(self, + vertices: torch.Tensor, + faces: torch.Tensor, + origin: list, + voxel_size: float, + coords: torch.Tensor, + attrs: torch.Tensor, + voxel_shape: torch.Size, + layout: Dict = {}, + ): + self.vertices = vertices.float() + self.faces = faces.int() + self.origin = torch.tensor(origin, dtype=torch.float32, device=self.device) + self.voxel_size = voxel_size + self.coords = coords + self.attrs = attrs + self.voxel_shape = voxel_shape + self.layout = layout + + def to(self, device, non_blocking=False): + return MeshWithVoxel( + self.vertices.to(device, non_blocking=non_blocking), + self.faces.to(device, non_blocking=non_blocking), + self.origin.tolist(), + self.voxel_size, + self.coords.to(device, non_blocking=non_blocking), + self.attrs.to(device, non_blocking=non_blocking), + self.voxel_shape, + self.layout, + ) + + def query_attrs(self, xyz): + from flex_gemm.ops.grid_sample import grid_sample_3d + grid = ((xyz - self.origin) / self.voxel_size).reshape(1, -1, 3) + vertex_attrs = grid_sample_3d( + self.attrs, + torch.cat([torch.zeros_like(self.coords[..., :1]), self.coords], dim=-1), + self.voxel_shape, + grid, + mode='trilinear' + )[0] + return vertex_attrs + + def query_vertex_attrs(self): + return self.query_attrs(self.vertices) diff --git a/deployment/runtime/qwen/diffsynth/core/__init__.py b/deployment/runtime/qwen/diffsynth/core/__init__.py index 8b1747471525a8c0d6c17131afd05778faa9bc2a..9b7b3b62be0972179f27c41dd9adebb30d4543e6 100644 --- a/deployment/runtime/qwen/diffsynth/core/__init__.py +++ b/deployment/runtime/qwen/diffsynth/core/__init__.py @@ -1,8 +1,6 @@ -from .attention import * -from .data import * -from .gradient import * -from .quant import * -from .loader import * -from .vram import * -from .device import * -from .offload_training import * +from .attention import * +from .gradient import * +from .quant import * +from .loader import * +from .vram import * +from .device import * diff --git a/deployment/runtime/qwen/diffsynth/core/data/__init__.py b/deployment/runtime/qwen/diffsynth/core/data/__init__.py index d494a277d81eeb2a9575155eb983d8bc3879590a..099e29fd09df53079662eefc4f8e89ae013d0109 100644 --- a/deployment/runtime/qwen/diffsynth/core/data/__init__.py +++ b/deployment/runtime/qwen/diffsynth/core/data/__init__.py @@ -1 +1 @@ -from .unified_dataset import UnifiedDataset +"""Image and media preprocessing operators used by inference pipelines.""" diff --git a/deployment/runtime/qwen/diffsynth/diffusion/__init__.py b/deployment/runtime/qwen/diffsynth/diffusion/__init__.py index f5d9ea1946dcee7a040ef1a21eaa8d1d4f160871..ea8c7d010c8a9c9b0c9198797f2669030693a00e 100644 --- a/deployment/runtime/qwen/diffsynth/diffusion/__init__.py +++ b/deployment/runtime/qwen/diffsynth/diffusion/__init__.py @@ -1,7 +1 @@ -from .flow_match import FlowMatchScheduler, HiDreamO1FlashScheduler, AncestralFlowMatchScheduler -from .training_module import DiffusionTrainingModule -from .logger import ModelLogger -from .runner import launch_training_task, launch_data_process_task -from .parsers import * -from .loss import * -from .dmd2 import * +from .flow_match import FlowMatchScheduler, HiDreamO1FlashScheduler, AncestralFlowMatchScheduler diff --git a/deployment/runtime/qwen/diffsynth/diffusion/base_pipeline.py b/deployment/runtime/qwen/diffsynth/diffusion/base_pipeline.py index 5af48a11085e90ab14254f43a094a5c52012361f..e23b696c0782e7a0a4189ed221708d33e5429ff1 100644 --- a/deployment/runtime/qwen/diffsynth/diffusion/base_pipeline.py +++ b/deployment/runtime/qwen/diffsynth/diffusion/base_pipeline.py @@ -1,542 +1,526 @@ -from PIL import Image -import torch -import numpy as np -from einops import repeat, reduce -from typing import Union -from ..core import AutoTorchModule, AutoWrappedLinear, LoRAHotLoadMixin, load_state_dict, ModelConfig, parse_device_type, enable_vram_management -from ..core.device.npu_compatible_device import get_device_type -from ..utils.lora import GeneralLoRALoader -from ..models.model_loader import ModelPool -from ..utils.controlnet import ControlNetInput -from ..core.device import get_device_name, IS_NPU_AVAILABLE -from .template import load_template_model, load_template_data_processor - - -class PipelineUnit: - def __init__( - self, - seperate_cfg: bool = False, - take_over: bool = False, - input_params: tuple[str] = None, - output_params: tuple[str] = None, - input_params_posi: dict[str, str] = None, - input_params_nega: dict[str, str] = None, - onload_model_names: tuple[str] = None - ): - self.seperate_cfg = seperate_cfg - self.take_over = take_over - self.input_params = input_params - self.output_params = output_params - self.input_params_posi = input_params_posi - self.input_params_nega = input_params_nega - self.onload_model_names = onload_model_names - - def fetch_input_params(self): - params = [] - if self.input_params is not None: - for param in self.input_params: - params.append(param) - if self.input_params_posi is not None: - for _, param in self.input_params_posi.items(): - params.append(param) - if self.input_params_nega is not None: - for _, param in self.input_params_nega.items(): - params.append(param) - params = sorted(list(set(params))) - return params - - def fetch_output_params(self): - params = [] - if self.output_params is not None: - for param in self.output_params: - params.append(param) - return params - - def process(self, pipe, **kwargs) -> dict: - return {} - - def post_process(self, pipe, **kwargs) -> dict: - return {} - - -class BasePipeline(torch.nn.Module): - - def __init__( - self, - device=get_device_type(), torch_dtype=torch.float16, - height_division_factor=64, width_division_factor=64, - time_division_factor=None, time_division_remainder=None, - ): - super().__init__() - # The device and torch_dtype is used for the storage of intermediate variables, not models. - self.device = device - self.torch_dtype = torch_dtype - self.device_type = parse_device_type(device) - # The following parameters are used for shape check. - self.height_division_factor = height_division_factor - self.width_division_factor = width_division_factor - self.time_division_factor = time_division_factor - self.time_division_remainder = time_division_remainder - # VRAM management - self.vram_management_enabled = False - # Pipeline Unit Runner - self.unit_runner = PipelineUnitRunner() - # LoRA Loader - self.lora_loader = GeneralLoRALoader - - - def to(self, *args, **kwargs): - device, dtype, non_blocking, convert_to_format = torch._C._nn._parse_to(*args, **kwargs) - if device is not None: - self.device = device - if dtype is not None: - self.torch_dtype = dtype - super().to(*args, **kwargs) - return self - - - def check_resize_height_width(self, height, width, num_frames=None, verbose=1): - # Shape check - if height % self.height_division_factor != 0: - height = (height + self.height_division_factor - 1) // self.height_division_factor * self.height_division_factor - if verbose > 0: - print(f"height % {self.height_division_factor} != 0. We round it up to {height}.") - if width % self.width_division_factor != 0: - width = (width + self.width_division_factor - 1) // self.width_division_factor * self.width_division_factor - if verbose > 0: - print(f"width % {self.width_division_factor} != 0. We round it up to {width}.") - if num_frames is None: - return height, width - else: - if num_frames % self.time_division_factor != self.time_division_remainder: - num_frames = (num_frames - self.time_division_remainder + self.time_division_factor - 1) // self.time_division_factor * self.time_division_factor + self.time_division_remainder - if verbose > 0: - print(f"num_frames % {self.time_division_factor} != {self.time_division_remainder}. We round it up to {num_frames}.") - return height, width, num_frames - - - def preprocess_image(self, image, torch_dtype=None, device=None, pattern="B C H W", min_value=-1, max_value=1): - # Transform a PIL.Image to torch.Tensor - image = torch.Tensor(np.array(image, dtype=np.float32)) - image = image.to(dtype=torch_dtype or self.torch_dtype, device=device or self.device) - image = image * ((max_value - min_value) / 255) + min_value - image = repeat(image, f"H W C -> {pattern}", **({"B": 1} if "B" in pattern else {})) - return image - - - def preprocess_video(self, video, torch_dtype=None, device=None, pattern="B C T H W", min_value=-1, max_value=1): - # Transform a list of PIL.Image to torch.Tensor - video = [self.preprocess_image(image, torch_dtype=torch_dtype, device=device, min_value=min_value, max_value=max_value) for image in video] - video = torch.stack(video, dim=pattern.index("T") // 2) - return video - - - def vae_output_to_image(self, vae_output, pattern="B C H W", min_value=-1, max_value=1): - # Transform a torch.Tensor to PIL.Image - if pattern != "H W C": - vae_output = reduce(vae_output, f"{pattern} -> H W C", reduction="mean") - image = ((vae_output - min_value) * (255 / (max_value - min_value))).clip(0, 255) - image = image.to(device="cpu", dtype=torch.uint8) - image = Image.fromarray(image.numpy()) - return image - - - def vae_output_to_video(self, vae_output, pattern="B C T H W", min_value=-1, max_value=1): - # Transform a torch.Tensor to list of PIL.Image - if pattern != "T H W C": - vae_output = reduce(vae_output, f"{pattern} -> T H W C", reduction="mean") - video = [self.vae_output_to_image(image, pattern="H W C", min_value=min_value, max_value=max_value) for image in vae_output] - return video - - def output_audio_format_check(self, audio_output): - # output standard foramt: [C, T], output dtype: float() - # remove batch dim - if audio_output.ndim == 3: - audio_output = audio_output.squeeze(0) - return audio_output.float().cpu() - - def load_models_to_device(self, model_names): - if self.vram_management_enabled: - # offload models - for name, model in self.named_children(): - if name not in model_names: - if hasattr(model, "vram_management_enabled") and model.vram_management_enabled: - if hasattr(model, "offload"): - model.offload() - else: - for module in model.modules(): - if hasattr(module, "offload"): - module.offload() - getattr(torch, self.device_type).empty_cache() - # onload models - for name, model in self.named_children(): - if name in model_names: - if hasattr(model, "vram_management_enabled") and model.vram_management_enabled: - if hasattr(model, "onload"): - model.onload() - else: - for module in model.modules(): - if hasattr(module, "onload"): - module.onload() - - - def generate_noise(self, shape, seed=None, rand_device="cpu", rand_torch_dtype=torch.float32, device=None, torch_dtype=None): - # Initialize Gaussian noise - generator = None if seed is None else torch.Generator(rand_device).manual_seed(seed) - noise = torch.randn(shape, generator=generator, device=rand_device, dtype=rand_torch_dtype) - noise = noise.to(dtype=torch_dtype or self.torch_dtype, device=device or self.device) - return noise - - - def get_vram(self): - device = self.device if not IS_NPU_AVAILABLE else get_device_name() - return getattr(torch, self.device_type).mem_get_info(device)[1] / (1024 ** 3) - - def get_module(self, model, name): - if "." in name: - name, suffix = name[:name.index(".")], name[name.index(".") + 1:] - if name.isdigit(): - return self.get_module(model[int(name)], suffix) - else: - return self.get_module(getattr(model, name), suffix) - else: - return getattr(model, name) - - def freeze_except(self, model_names): - self.eval() - self.requires_grad_(False) - for name in model_names: - module = self.get_module(self, name) - if module is None: - print(f"No {name} models in the pipeline. We cannot enable training on the model. If this occurs during the data processing stage, it is normal.") - continue - module.train() - module.requires_grad_(True) - - - def blend_with_mask(self, base, addition, mask): - return base * (1 - mask) + addition * mask - - - def step(self, scheduler, latents, progress_id, noise_pred, input_latents=None, inpaint_mask=None, **kwargs): - timestep = scheduler.timesteps[progress_id] - if inpaint_mask is not None: - noise_pred_expected = scheduler.return_to_timestep(scheduler.timesteps[progress_id], latents, input_latents) - noise_pred = self.blend_with_mask(noise_pred_expected, noise_pred, inpaint_mask) - latents_next = scheduler.step(noise_pred, timestep, latents) - return latents_next - - - def split_pipeline_units(self, model_names: list[str]): - return PipelineUnitGraph().split_pipeline_units(self.units, model_names) - - - def flush_vram_management_device(self, device): - for module in self.modules(): - if isinstance(module, AutoTorchModule): - module.offload_device = device - module.onload_device = device - module.preparing_device = device - module.computation_device = device - - - def check_quant_hot_load(self, module: torch.nn.Module): - # A quantized weight cannot absorb a fused LoRA, so hot-loading is the only option for it. - if getattr(module, "quantize_config", None) is None: - return module - return self.enable_lora_hot_loading(module) - - - def load_lora( - self, - module: torch.nn.Module, - lora_config: Union[ModelConfig, str] = None, - alpha=1, - hotload=None, - state_dict=None, - verbose=1, - ): - module = self.check_quant_hot_load(module) - if state_dict is None: - if isinstance(lora_config, str): - lora = load_state_dict(lora_config, torch_dtype=self.torch_dtype, device=self.device) - else: - lora_config.download_if_necessary() - lora = load_state_dict(lora_config.path, torch_dtype=self.torch_dtype, device=self.device) - else: - lora = state_dict - lora_loader = self.lora_loader(torch_dtype=self.torch_dtype, device=self.device) - lora = lora_loader.convert_state_dict(lora) - if hotload is None: - hotload = hasattr(module, "vram_management_enabled") and getattr(module, "vram_management_enabled") - if hotload: - if not (hasattr(module, "vram_management_enabled") and getattr(module, "vram_management_enabled")): - raise ValueError("VRAM Management is not enabled. LoRA hotloading is not supported.") - updated_num = 0 - for _, module in module.named_modules(): - if isinstance(module, LoRAHotLoadMixin): - name = module.name - lora_a_name = f'{name}.lora_A.weight' - lora_b_name = f'{name}.lora_B.weight' - if lora_a_name in lora and lora_b_name in lora: - updated_num += 1 - module.lora_A_weights.append(lora[lora_a_name] * alpha) - module.lora_B_weights.append(lora[lora_b_name]) - if verbose >= 1: - print(f"{updated_num} tensors are patched by LoRA. You can use `pipe.clear_lora()` to clear all LoRA layers.") - else: - lora_loader.fuse_lora_to_base_model(module, lora, alpha=alpha) - - - def clear_lora(self, verbose=1): - cleared_num = 0 - for name, module in self.named_modules(): - if isinstance(module, LoRAHotLoadMixin): - if hasattr(module, "lora_A_weights"): - if len(module.lora_A_weights) > 0: - cleared_num += 1 - module.lora_A_weights.clear() - if hasattr(module, "lora_B_weights"): - module.lora_B_weights.clear() - if verbose >= 1: - print(f"{cleared_num} LoRA layers are cleared.") - - - def download_and_load_models(self, model_configs: list[ModelConfig] = [], vram_limit: float = None): - model_pool = ModelPool() - for model_config in model_configs: - model_config.download_if_necessary() - vram_config = model_config.vram_config() - vram_config["computation_dtype"] = vram_config["computation_dtype"] or self.torch_dtype - vram_config["computation_device"] = vram_config["computation_device"] or self.device - model_pool.auto_load_model( - model_config.path, - vram_config=vram_config, - vram_limit=vram_limit, - clear_parameters=model_config.clear_parameters, - state_dict=model_config.state_dict, - quantize=model_config.quantize, - ) - return model_pool - - - def check_vram_management_state(self): - vram_management_enabled = False - for module in self.children(): - if hasattr(module, "vram_management_enabled") and getattr(module, "vram_management_enabled"): - vram_management_enabled = True - return vram_management_enabled - - - def cfg_guided_model_fn(self, model_fn, cfg_scale, inputs_shared, inputs_posi, inputs_nega, **inputs_others): - # Positive side forward - if inputs_shared.get("positive_only_lora", None) is not None: - self.load_lora(self.dit, state_dict=inputs_shared["positive_only_lora"], verbose=0) - noise_pred_posi = model_fn(**inputs_posi, **inputs_shared, **inputs_others) - if inputs_shared.get("positive_only_lora", None) is not None: - self.clear_lora(verbose=0) - - if cfg_scale != 1.0: - # Negative side forward - if inputs_shared.get("negative_only_lora", None) is not None: - self.load_lora(self.dit, state_dict=inputs_shared["negative_only_lora"], verbose=0) - noise_pred_nega = model_fn(**inputs_nega, **inputs_shared, **inputs_others) - if inputs_shared.get("negative_only_lora", None) is not None: - self.clear_lora(verbose=0) - - if isinstance(noise_pred_posi, tuple): - # Separately handling different output types of latents, eg. video and audio latents. - noise_pred = tuple( - n_nega + cfg_scale * (n_posi - n_nega) - for n_posi, n_nega in zip(noise_pred_posi, noise_pred_nega) - ) - else: - noise_pred = noise_pred_nega + cfg_scale * (noise_pred_posi - noise_pred_nega) - else: - noise_pred = noise_pred_posi - return noise_pred - - - def load_training_template_model(self, model_config: ModelConfig = None): - if model_config is not None: - model_config.download_if_necessary() - self.template_model = load_template_model(model_config.path, torch_dtype=self.torch_dtype, device=self.device) - self.template_data_processor = load_template_data_processor(model_config.path)() - - - def enable_lora_hot_loading(self, model: torch.nn.Module): - if hasattr(model, "vram_management_enabled") and getattr(model, "vram_management_enabled"): - return model - module_map = {torch.nn.Linear: AutoWrappedLinear} - vram_config = { - "offload_dtype": self.torch_dtype, - "offload_device": self.device, - "onload_dtype": self.torch_dtype, - "onload_device": self.device, - "preparing_dtype": self.torch_dtype, - "preparing_device": self.device, - "computation_dtype": self.torch_dtype, - "computation_device": self.device, - } - model = enable_vram_management(model, module_map, vram_config=vram_config, quantize=getattr(model, "quantize_config", None)) - return model - - def compile_pipeline(self, mode: str = "default", dynamic: bool = True, fullgraph: bool = False, compile_models: list = None, **kwargs): - """ - compile the pipeline with torch.compile. The models that will be compiled are determined by the `compilable_models` attribute of the pipeline. - If a model has `_repeated_blocks` attribute, we will compile these blocks with regional compilation. Otherwise, we will compile the whole model. - See https://docs.pytorch.org/docs/stable/generated/torch.compile.html#torch.compile for details about compilation arguments. - Args: - mode: The compilation mode, which will be passed to `torch.compile`, options are "default", "reduce-overhead", "max-autotune" and "max-autotune-no-cudagraphs. Default to "default". - dynamic: Whether to enable dynamic graph compilation to support dynamic input shapes, which will be passed to `torch.compile`. Default to True (recommended). - fullgraph: Whether to use full graph compilation, which will be passed to `torch.compile`. Default to False (recommended). - compile_models: The list of model names to be compiled. If None, we will compile the models in `pipeline.compilable_models`. Default to None. - **kwargs: Other arguments for `torch.compile`. - """ - compile_models = compile_models or getattr(self, "compilable_models", []) - if len(compile_models) == 0: - print("No compilable models in the pipeline. Skip compilation.") - return - for name in compile_models: - model = getattr(self, name, None) - if model is None: - print(f"Model '{name}' not found in the pipeline.") - continue - repeated_blocks = getattr(model, "_repeated_blocks", None) - # regional compilation for repeated blocks. - if repeated_blocks is not None: - for submod in model.modules(): - if submod.__class__.__name__ in repeated_blocks: - submod.compile(mode=mode, dynamic=dynamic, fullgraph=fullgraph, **kwargs) - # compile the whole model. - else: - model.compile(mode=mode, dynamic=dynamic, fullgraph=fullgraph, **kwargs) - print(f"{name} is compiled with mode={mode}, dynamic={dynamic}, fullgraph={fullgraph}.") - - -class PipelineUnitGraph: - def __init__(self): - pass - - def build_edges(self, units: list[PipelineUnit]): - # Establish dependencies between units - # to search for subsequent related computation units. - last_compute_unit_id = {} - edges = [] - for unit_id, unit in enumerate(units): - for input_param in unit.fetch_input_params(): - if input_param in last_compute_unit_id: - edges.append((last_compute_unit_id[input_param], unit_id)) - for output_param in unit.fetch_output_params(): - last_compute_unit_id[output_param] = unit_id - return edges - - def build_chains(self, units: list[PipelineUnit]): - # Establish updating chains for each variable - # to track their computation process. - params = sum([unit.fetch_input_params() + unit.fetch_output_params() for unit in units], []) - params = sorted(list(set(params))) - chains = {param: [] for param in params} - for unit_id, unit in enumerate(units): - for param in unit.fetch_output_params(): - chains[param].append(unit_id) - return chains - - def search_direct_unit_ids(self, units: list[PipelineUnit], model_names: list[str]): - # Search for units that directly participate in the model's computation. - related_unit_ids = [] - for unit_id, unit in enumerate(units): - for model_name in model_names: - if unit.onload_model_names is not None and model_name in unit.onload_model_names: - related_unit_ids.append(unit_id) - break - return related_unit_ids - - def search_related_unit_ids(self, edges, start_unit_ids, direction="target"): - # Search for subsequent related computation units. - related_unit_ids = [unit_id for unit_id in start_unit_ids] - while True: - neighbors = [] - for source, target in edges: - if direction == "target" and source in related_unit_ids and target not in related_unit_ids: - neighbors.append(target) - elif direction == "source" and source not in related_unit_ids and target in related_unit_ids: - neighbors.append(source) - neighbors = sorted(list(set(neighbors))) - if len(neighbors) == 0: - break - else: - related_unit_ids.extend(neighbors) - related_unit_ids = sorted(list(set(related_unit_ids))) - return related_unit_ids - - def search_updating_unit_ids(self, units: list[PipelineUnit], chains, related_unit_ids): - # If the input parameters of this subgraph are updated outside the subgraph, - # search for the units where these updates occur. - first_compute_unit_id = {} - for unit_id in related_unit_ids: - for param in units[unit_id].fetch_input_params(): - if param not in first_compute_unit_id: - first_compute_unit_id[param] = unit_id - updating_unit_ids = [] - for param in first_compute_unit_id: - unit_id = first_compute_unit_id[param] - chain = chains[param] - if unit_id in chain and chain.index(unit_id) != len(chain) - 1: - for unit_id_ in chain[chain.index(unit_id) + 1:]: - if unit_id_ not in related_unit_ids: - updating_unit_ids.append(unit_id_) - related_unit_ids.extend(updating_unit_ids) - related_unit_ids = sorted(list(set(related_unit_ids))) - return related_unit_ids - - def split_pipeline_units(self, units: list[PipelineUnit], model_names: list[str]): - # Split the computation graph, - # separating all model-related computations. - related_unit_ids = self.search_direct_unit_ids(units, model_names) - edges = self.build_edges(units) - chains = self.build_chains(units) - while True: - num_related_unit_ids = len(related_unit_ids) - related_unit_ids = self.search_related_unit_ids(edges, related_unit_ids, "target") - related_unit_ids = self.search_updating_unit_ids(units, chains, related_unit_ids) - if len(related_unit_ids) == num_related_unit_ids: - break - else: - num_related_unit_ids = len(related_unit_ids) - related_units = [units[i] for i in related_unit_ids] - unrelated_units = [units[i] for i in range(len(units)) if i not in related_unit_ids] - return related_units, unrelated_units - - -class PipelineUnitRunner: - def __init__(self): - pass - - def __call__(self, unit: PipelineUnit, pipe: BasePipeline, inputs_shared: dict, inputs_posi: dict, inputs_nega: dict) -> tuple[dict, dict]: - if unit.take_over: - # Let the pipeline unit take over this function. - inputs_shared, inputs_posi, inputs_nega = unit.process(pipe, inputs_shared=inputs_shared, inputs_posi=inputs_posi, inputs_nega=inputs_nega) - elif unit.seperate_cfg: - # Positive side - processor_inputs = {name: inputs_posi.get(name_) for name, name_ in unit.input_params_posi.items()} - if unit.input_params is not None: - for name in unit.input_params: - processor_inputs[name] = inputs_shared.get(name) - processor_outputs = unit.process(pipe, **processor_inputs) - inputs_posi.update(processor_outputs) - # Negative side - if inputs_shared["cfg_scale"] != 1: - processor_inputs = {name: inputs_nega.get(name_) for name, name_ in unit.input_params_nega.items()} - if unit.input_params is not None: - for name in unit.input_params: - processor_inputs[name] = inputs_shared.get(name) - processor_outputs = unit.process(pipe, **processor_inputs) - inputs_nega.update(processor_outputs) - else: - inputs_nega.update(processor_outputs) - else: - processor_inputs = {name: inputs_shared.get(name) for name in unit.input_params} - processor_outputs = unit.process(pipe, **processor_inputs) - inputs_shared.update(processor_outputs) - return inputs_shared, inputs_posi, inputs_nega +from PIL import Image +import torch +import numpy as np +from einops import repeat, reduce +from typing import Union +from ..core import AutoTorchModule, AutoWrappedLinear, LoRAHotLoadMixin, load_state_dict, ModelConfig, parse_device_type, enable_vram_management +from ..core.device.npu_compatible_device import get_device_type +from ..utils.lora import GeneralLoRALoader +from ..models.model_loader import ModelPool +from ..utils.controlnet import ControlNetInput +from ..core.device import get_device_name, IS_NPU_AVAILABLE + + +class PipelineUnit: + def __init__( + self, + seperate_cfg: bool = False, + take_over: bool = False, + input_params: tuple[str] = None, + output_params: tuple[str] = None, + input_params_posi: dict[str, str] = None, + input_params_nega: dict[str, str] = None, + onload_model_names: tuple[str] = None + ): + self.seperate_cfg = seperate_cfg + self.take_over = take_over + self.input_params = input_params + self.output_params = output_params + self.input_params_posi = input_params_posi + self.input_params_nega = input_params_nega + self.onload_model_names = onload_model_names + + def fetch_input_params(self): + params = [] + if self.input_params is not None: + for param in self.input_params: + params.append(param) + if self.input_params_posi is not None: + for _, param in self.input_params_posi.items(): + params.append(param) + if self.input_params_nega is not None: + for _, param in self.input_params_nega.items(): + params.append(param) + params = sorted(list(set(params))) + return params + + def fetch_output_params(self): + params = [] + if self.output_params is not None: + for param in self.output_params: + params.append(param) + return params + + def process(self, pipe, **kwargs) -> dict: + return {} + + def post_process(self, pipe, **kwargs) -> dict: + return {} + + +class BasePipeline(torch.nn.Module): + + def __init__( + self, + device=get_device_type(), torch_dtype=torch.float16, + height_division_factor=64, width_division_factor=64, + time_division_factor=None, time_division_remainder=None, + ): + super().__init__() + # The device and torch_dtype is used for the storage of intermediate variables, not models. + self.device = device + self.torch_dtype = torch_dtype + self.device_type = parse_device_type(device) + # The following parameters are used for shape check. + self.height_division_factor = height_division_factor + self.width_division_factor = width_division_factor + self.time_division_factor = time_division_factor + self.time_division_remainder = time_division_remainder + # VRAM management + self.vram_management_enabled = False + # Pipeline Unit Runner + self.unit_runner = PipelineUnitRunner() + # LoRA Loader + self.lora_loader = GeneralLoRALoader + + + def to(self, *args, **kwargs): + device, dtype, non_blocking, convert_to_format = torch._C._nn._parse_to(*args, **kwargs) + if device is not None: + self.device = device + if dtype is not None: + self.torch_dtype = dtype + super().to(*args, **kwargs) + return self + + + def check_resize_height_width(self, height, width, num_frames=None, verbose=1): + # Shape check + if height % self.height_division_factor != 0: + height = (height + self.height_division_factor - 1) // self.height_division_factor * self.height_division_factor + if verbose > 0: + print(f"height % {self.height_division_factor} != 0. We round it up to {height}.") + if width % self.width_division_factor != 0: + width = (width + self.width_division_factor - 1) // self.width_division_factor * self.width_division_factor + if verbose > 0: + print(f"width % {self.width_division_factor} != 0. We round it up to {width}.") + if num_frames is None: + return height, width + else: + if num_frames % self.time_division_factor != self.time_division_remainder: + num_frames = (num_frames - self.time_division_remainder + self.time_division_factor - 1) // self.time_division_factor * self.time_division_factor + self.time_division_remainder + if verbose > 0: + print(f"num_frames % {self.time_division_factor} != {self.time_division_remainder}. We round it up to {num_frames}.") + return height, width, num_frames + + + def preprocess_image(self, image, torch_dtype=None, device=None, pattern="B C H W", min_value=-1, max_value=1): + # Transform a PIL.Image to torch.Tensor + image = torch.Tensor(np.array(image, dtype=np.float32)) + image = image.to(dtype=torch_dtype or self.torch_dtype, device=device or self.device) + image = image * ((max_value - min_value) / 255) + min_value + image = repeat(image, f"H W C -> {pattern}", **({"B": 1} if "B" in pattern else {})) + return image + + + def preprocess_video(self, video, torch_dtype=None, device=None, pattern="B C T H W", min_value=-1, max_value=1): + # Transform a list of PIL.Image to torch.Tensor + video = [self.preprocess_image(image, torch_dtype=torch_dtype, device=device, min_value=min_value, max_value=max_value) for image in video] + video = torch.stack(video, dim=pattern.index("T") // 2) + return video + + + def vae_output_to_image(self, vae_output, pattern="B C H W", min_value=-1, max_value=1): + # Transform a torch.Tensor to PIL.Image + if pattern != "H W C": + vae_output = reduce(vae_output, f"{pattern} -> H W C", reduction="mean") + image = ((vae_output - min_value) * (255 / (max_value - min_value))).clip(0, 255) + image = image.to(device="cpu", dtype=torch.uint8) + image = Image.fromarray(image.numpy()) + return image + + + def vae_output_to_video(self, vae_output, pattern="B C T H W", min_value=-1, max_value=1): + # Transform a torch.Tensor to list of PIL.Image + if pattern != "T H W C": + vae_output = reduce(vae_output, f"{pattern} -> T H W C", reduction="mean") + video = [self.vae_output_to_image(image, pattern="H W C", min_value=min_value, max_value=max_value) for image in vae_output] + return video + + def output_audio_format_check(self, audio_output): + # output standard foramt: [C, T], output dtype: float() + # remove batch dim + if audio_output.ndim == 3: + audio_output = audio_output.squeeze(0) + return audio_output.float().cpu() + + def load_models_to_device(self, model_names): + if self.vram_management_enabled: + # offload models + for name, model in self.named_children(): + if name not in model_names: + if hasattr(model, "vram_management_enabled") and model.vram_management_enabled: + if hasattr(model, "offload"): + model.offload() + else: + for module in model.modules(): + if hasattr(module, "offload"): + module.offload() + getattr(torch, self.device_type).empty_cache() + # onload models + for name, model in self.named_children(): + if name in model_names: + if hasattr(model, "vram_management_enabled") and model.vram_management_enabled: + if hasattr(model, "onload"): + model.onload() + else: + for module in model.modules(): + if hasattr(module, "onload"): + module.onload() + + + def generate_noise(self, shape, seed=None, rand_device="cpu", rand_torch_dtype=torch.float32, device=None, torch_dtype=None): + # Initialize Gaussian noise + generator = None if seed is None else torch.Generator(rand_device).manual_seed(seed) + noise = torch.randn(shape, generator=generator, device=rand_device, dtype=rand_torch_dtype) + noise = noise.to(dtype=torch_dtype or self.torch_dtype, device=device or self.device) + return noise + + + def get_vram(self): + device = self.device if not IS_NPU_AVAILABLE else get_device_name() + return getattr(torch, self.device_type).mem_get_info(device)[1] / (1024 ** 3) + + def get_module(self, model, name): + if "." in name: + name, suffix = name[:name.index(".")], name[name.index(".") + 1:] + if name.isdigit(): + return self.get_module(model[int(name)], suffix) + else: + return self.get_module(getattr(model, name), suffix) + else: + return getattr(model, name) + + + + def blend_with_mask(self, base, addition, mask): + return base * (1 - mask) + addition * mask + + + def step(self, scheduler, latents, progress_id, noise_pred, input_latents=None, inpaint_mask=None, **kwargs): + timestep = scheduler.timesteps[progress_id] + if inpaint_mask is not None: + noise_pred_expected = scheduler.return_to_timestep(scheduler.timesteps[progress_id], latents, input_latents) + noise_pred = self.blend_with_mask(noise_pred_expected, noise_pred, inpaint_mask) + latents_next = scheduler.step(noise_pred, timestep, latents) + return latents_next + + + def split_pipeline_units(self, model_names: list[str]): + return PipelineUnitGraph().split_pipeline_units(self.units, model_names) + + + def flush_vram_management_device(self, device): + for module in self.modules(): + if isinstance(module, AutoTorchModule): + module.offload_device = device + module.onload_device = device + module.preparing_device = device + module.computation_device = device + + + def check_quant_hot_load(self, module: torch.nn.Module): + # A quantized weight cannot absorb a fused LoRA, so hot-loading is the only option for it. + if getattr(module, "quantize_config", None) is None: + return module + return self.enable_lora_hot_loading(module) + + + def load_lora( + self, + module: torch.nn.Module, + lora_config: Union[ModelConfig, str] = None, + alpha=1, + hotload=None, + state_dict=None, + verbose=1, + ): + module = self.check_quant_hot_load(module) + if state_dict is None: + if isinstance(lora_config, str): + lora = load_state_dict(lora_config, torch_dtype=self.torch_dtype, device=self.device) + else: + lora_config.download_if_necessary() + lora = load_state_dict(lora_config.path, torch_dtype=self.torch_dtype, device=self.device) + else: + lora = state_dict + lora_loader = self.lora_loader(torch_dtype=self.torch_dtype, device=self.device) + lora = lora_loader.convert_state_dict(lora) + if hotload is None: + hotload = hasattr(module, "vram_management_enabled") and getattr(module, "vram_management_enabled") + if hotload: + if not (hasattr(module, "vram_management_enabled") and getattr(module, "vram_management_enabled")): + raise ValueError("VRAM Management is not enabled. LoRA hotloading is not supported.") + updated_num = 0 + for _, module in module.named_modules(): + if isinstance(module, LoRAHotLoadMixin): + name = module.name + lora_a_name = f'{name}.lora_A.weight' + lora_b_name = f'{name}.lora_B.weight' + if lora_a_name in lora and lora_b_name in lora: + updated_num += 1 + module.lora_A_weights.append(lora[lora_a_name] * alpha) + module.lora_B_weights.append(lora[lora_b_name]) + if verbose >= 1: + print(f"{updated_num} tensors are patched by LoRA. You can use `pipe.clear_lora()` to clear all LoRA layers.") + else: + lora_loader.fuse_lora_to_base_model(module, lora, alpha=alpha) + + + def clear_lora(self, verbose=1): + cleared_num = 0 + for name, module in self.named_modules(): + if isinstance(module, LoRAHotLoadMixin): + if hasattr(module, "lora_A_weights"): + if len(module.lora_A_weights) > 0: + cleared_num += 1 + module.lora_A_weights.clear() + if hasattr(module, "lora_B_weights"): + module.lora_B_weights.clear() + if verbose >= 1: + print(f"{cleared_num} LoRA layers are cleared.") + + + def download_and_load_models(self, model_configs: list[ModelConfig] = [], vram_limit: float = None): + model_pool = ModelPool() + for model_config in model_configs: + model_config.download_if_necessary() + vram_config = model_config.vram_config() + vram_config["computation_dtype"] = vram_config["computation_dtype"] or self.torch_dtype + vram_config["computation_device"] = vram_config["computation_device"] or self.device + model_pool.auto_load_model( + model_config.path, + vram_config=vram_config, + vram_limit=vram_limit, + clear_parameters=model_config.clear_parameters, + state_dict=model_config.state_dict, + quantize=model_config.quantize, + ) + return model_pool + + + def check_vram_management_state(self): + vram_management_enabled = False + for module in self.children(): + if hasattr(module, "vram_management_enabled") and getattr(module, "vram_management_enabled"): + vram_management_enabled = True + return vram_management_enabled + + + def cfg_guided_model_fn(self, model_fn, cfg_scale, inputs_shared, inputs_posi, inputs_nega, **inputs_others): + # Positive side forward + if inputs_shared.get("positive_only_lora", None) is not None: + self.load_lora(self.dit, state_dict=inputs_shared["positive_only_lora"], verbose=0) + noise_pred_posi = model_fn(**inputs_posi, **inputs_shared, **inputs_others) + if inputs_shared.get("positive_only_lora", None) is not None: + self.clear_lora(verbose=0) + + if cfg_scale != 1.0: + # Negative side forward + if inputs_shared.get("negative_only_lora", None) is not None: + self.load_lora(self.dit, state_dict=inputs_shared["negative_only_lora"], verbose=0) + noise_pred_nega = model_fn(**inputs_nega, **inputs_shared, **inputs_others) + if inputs_shared.get("negative_only_lora", None) is not None: + self.clear_lora(verbose=0) + + if isinstance(noise_pred_posi, tuple): + # Separately handling different output types of latents, eg. video and audio latents. + noise_pred = tuple( + n_nega + cfg_scale * (n_posi - n_nega) + for n_posi, n_nega in zip(noise_pred_posi, noise_pred_nega) + ) + else: + noise_pred = noise_pred_nega + cfg_scale * (noise_pred_posi - noise_pred_nega) + else: + noise_pred = noise_pred_posi + return noise_pred + + + + + def enable_lora_hot_loading(self, model: torch.nn.Module): + if hasattr(model, "vram_management_enabled") and getattr(model, "vram_management_enabled"): + return model + module_map = {torch.nn.Linear: AutoWrappedLinear} + vram_config = { + "offload_dtype": self.torch_dtype, + "offload_device": self.device, + "onload_dtype": self.torch_dtype, + "onload_device": self.device, + "preparing_dtype": self.torch_dtype, + "preparing_device": self.device, + "computation_dtype": self.torch_dtype, + "computation_device": self.device, + } + model = enable_vram_management(model, module_map, vram_config=vram_config, quantize=getattr(model, "quantize_config", None)) + return model + + def compile_pipeline(self, mode: str = "default", dynamic: bool = True, fullgraph: bool = False, compile_models: list = None, **kwargs): + """ + compile the pipeline with torch.compile. The models that will be compiled are determined by the `compilable_models` attribute of the pipeline. + If a model has `_repeated_blocks` attribute, we will compile these blocks with regional compilation. Otherwise, we will compile the whole model. + See https://docs.pytorch.org/docs/stable/generated/torch.compile.html#torch.compile for details about compilation arguments. + Args: + mode: The compilation mode, which will be passed to `torch.compile`, options are "default", "reduce-overhead", "max-autotune" and "max-autotune-no-cudagraphs. Default to "default". + dynamic: Whether to enable dynamic graph compilation to support dynamic input shapes, which will be passed to `torch.compile`. Default to True (recommended). + fullgraph: Whether to use full graph compilation, which will be passed to `torch.compile`. Default to False (recommended). + compile_models: The list of model names to be compiled. If None, we will compile the models in `pipeline.compilable_models`. Default to None. + **kwargs: Other arguments for `torch.compile`. + """ + compile_models = compile_models or getattr(self, "compilable_models", []) + if len(compile_models) == 0: + print("No compilable models in the pipeline. Skip compilation.") + return + for name in compile_models: + model = getattr(self, name, None) + if model is None: + print(f"Model '{name}' not found in the pipeline.") + continue + repeated_blocks = getattr(model, "_repeated_blocks", None) + # regional compilation for repeated blocks. + if repeated_blocks is not None: + for submod in model.modules(): + if submod.__class__.__name__ in repeated_blocks: + submod.compile(mode=mode, dynamic=dynamic, fullgraph=fullgraph, **kwargs) + # compile the whole model. + else: + model.compile(mode=mode, dynamic=dynamic, fullgraph=fullgraph, **kwargs) + print(f"{name} is compiled with mode={mode}, dynamic={dynamic}, fullgraph={fullgraph}.") + + +class PipelineUnitGraph: + def __init__(self): + pass + + def build_edges(self, units: list[PipelineUnit]): + # Establish dependencies between units + # to search for subsequent related computation units. + last_compute_unit_id = {} + edges = [] + for unit_id, unit in enumerate(units): + for input_param in unit.fetch_input_params(): + if input_param in last_compute_unit_id: + edges.append((last_compute_unit_id[input_param], unit_id)) + for output_param in unit.fetch_output_params(): + last_compute_unit_id[output_param] = unit_id + return edges + + def build_chains(self, units: list[PipelineUnit]): + # Establish updating chains for each variable + # to track their computation process. + params = sum([unit.fetch_input_params() + unit.fetch_output_params() for unit in units], []) + params = sorted(list(set(params))) + chains = {param: [] for param in params} + for unit_id, unit in enumerate(units): + for param in unit.fetch_output_params(): + chains[param].append(unit_id) + return chains + + def search_direct_unit_ids(self, units: list[PipelineUnit], model_names: list[str]): + # Search for units that directly participate in the model's computation. + related_unit_ids = [] + for unit_id, unit in enumerate(units): + for model_name in model_names: + if unit.onload_model_names is not None and model_name in unit.onload_model_names: + related_unit_ids.append(unit_id) + break + return related_unit_ids + + def search_related_unit_ids(self, edges, start_unit_ids, direction="target"): + # Search for subsequent related computation units. + related_unit_ids = [unit_id for unit_id in start_unit_ids] + while True: + neighbors = [] + for source, target in edges: + if direction == "target" and source in related_unit_ids and target not in related_unit_ids: + neighbors.append(target) + elif direction == "source" and source not in related_unit_ids and target in related_unit_ids: + neighbors.append(source) + neighbors = sorted(list(set(neighbors))) + if len(neighbors) == 0: + break + else: + related_unit_ids.extend(neighbors) + related_unit_ids = sorted(list(set(related_unit_ids))) + return related_unit_ids + + def search_updating_unit_ids(self, units: list[PipelineUnit], chains, related_unit_ids): + # If the input parameters of this subgraph are updated outside the subgraph, + # search for the units where these updates occur. + first_compute_unit_id = {} + for unit_id in related_unit_ids: + for param in units[unit_id].fetch_input_params(): + if param not in first_compute_unit_id: + first_compute_unit_id[param] = unit_id + updating_unit_ids = [] + for param in first_compute_unit_id: + unit_id = first_compute_unit_id[param] + chain = chains[param] + if unit_id in chain and chain.index(unit_id) != len(chain) - 1: + for unit_id_ in chain[chain.index(unit_id) + 1:]: + if unit_id_ not in related_unit_ids: + updating_unit_ids.append(unit_id_) + related_unit_ids.extend(updating_unit_ids) + related_unit_ids = sorted(list(set(related_unit_ids))) + return related_unit_ids + + def split_pipeline_units(self, units: list[PipelineUnit], model_names: list[str]): + # Split the computation graph, + # separating all model-related computations. + related_unit_ids = self.search_direct_unit_ids(units, model_names) + edges = self.build_edges(units) + chains = self.build_chains(units) + while True: + num_related_unit_ids = len(related_unit_ids) + related_unit_ids = self.search_related_unit_ids(edges, related_unit_ids, "target") + related_unit_ids = self.search_updating_unit_ids(units, chains, related_unit_ids) + if len(related_unit_ids) == num_related_unit_ids: + break + else: + num_related_unit_ids = len(related_unit_ids) + related_units = [units[i] for i in related_unit_ids] + unrelated_units = [units[i] for i in range(len(units)) if i not in related_unit_ids] + return related_units, unrelated_units + + +class PipelineUnitRunner: + def __init__(self): + pass + + def __call__(self, unit: PipelineUnit, pipe: BasePipeline, inputs_shared: dict, inputs_posi: dict, inputs_nega: dict) -> tuple[dict, dict]: + if unit.take_over: + # Let the pipeline unit take over this function. + inputs_shared, inputs_posi, inputs_nega = unit.process(pipe, inputs_shared=inputs_shared, inputs_posi=inputs_posi, inputs_nega=inputs_nega) + elif unit.seperate_cfg: + # Positive side + processor_inputs = {name: inputs_posi.get(name_) for name, name_ in unit.input_params_posi.items()} + if unit.input_params is not None: + for name in unit.input_params: + processor_inputs[name] = inputs_shared.get(name) + processor_outputs = unit.process(pipe, **processor_inputs) + inputs_posi.update(processor_outputs) + # Negative side + if inputs_shared["cfg_scale"] != 1: + processor_inputs = {name: inputs_nega.get(name_) for name, name_ in unit.input_params_nega.items()} + if unit.input_params is not None: + for name in unit.input_params: + processor_inputs[name] = inputs_shared.get(name) + processor_outputs = unit.process(pipe, **processor_inputs) + inputs_nega.update(processor_outputs) + else: + inputs_nega.update(processor_outputs) + else: + processor_inputs = {name: inputs_shared.get(name) for name in unit.input_params} + processor_outputs = unit.process(pipe, **processor_inputs) + inputs_shared.update(processor_outputs) + return inputs_shared, inputs_posi, inputs_nega diff --git a/deployment/runtime/qwen/diffsynth/diffusion/template.py b/deployment/runtime/qwen/diffsynth/diffusion/template.py index 8e1cdb18cdaa265f11c1e6ee287c325855185393..fdbd84e8d87d58a125a4025f9da729c510eafe39 100644 --- a/deployment/runtime/qwen/diffsynth/diffusion/template.py +++ b/deployment/runtime/qwen/diffsynth/diffusion/template.py @@ -1,209 +1,200 @@ -import torch, os, importlib, warnings, json, inspect -from typing import Dict, List, Tuple, Union -from ..core import ModelConfig, load_model -from ..core.device.npu_compatible_device import get_device_type -from ..utils.lora.merge import merge_lora - - -KVCache = Dict[str, Tuple[torch.Tensor, torch.Tensor]] - - -class TemplateModel(torch.nn.Module): - def __init__(self): - super().__init__() - - @torch.no_grad() - def process_inputs(self, **kwargs): - return {} - - def forward(self, **kwargs): - raise NotImplementedError() - - -def check_template_model_format(model): - if not hasattr(model, "process_inputs"): - raise NotImplementedError("`process_inputs` is not implemented in the Template model.") - if "kwargs" not in inspect.signature(model.process_inputs).parameters: - raise NotImplementedError("`**kwargs` is not included in `process_inputs`.") - if not hasattr(model, "forward"): - raise NotImplementedError("`forward` is not implemented in the Template model.") - if "kwargs" not in inspect.signature(model.forward).parameters: - raise NotImplementedError("`**kwargs` is not included in `forward`.") - - -def load_template_model(path, torch_dtype=torch.bfloat16, device="cuda", verbose=1, state_dict=None): - spec = importlib.util.spec_from_file_location("template_model", os.path.join(path, "model.py")) - module = importlib.util.module_from_spec(spec) - spec.loader.exec_module(module) - template_model_path = getattr(module, 'TEMPLATE_MODEL_PATH') if hasattr(module, 'TEMPLATE_MODEL_PATH') else None - if template_model_path is not None: - # With `TEMPLATE_MODEL_PATH`, a pretrained model will be loaded. - model = load_model( - model_class=getattr(module, 'TEMPLATE_MODEL'), - config=getattr(module, 'TEMPLATE_MODEL_CONFIG') if hasattr(module, 'TEMPLATE_MODEL_CONFIG') else None, - path=os.path.join(path, getattr(module, 'TEMPLATE_MODEL_PATH')), - torch_dtype=torch_dtype, - device=device, - state_dict=state_dict, - ) - else: - # Without `TEMPLATE_MODEL_PATH`, a randomly initialized model or a non-model module will be loaded. - model = module.TEMPLATE_MODEL() - if hasattr(model, "to"): - model = model.to(dtype=torch_dtype, device=device) - if hasattr(model, "eval"): - model = model.eval() - check_template_model_format(model) - if verbose > 0: - metadata = { - "model_architecture": getattr(module, 'TEMPLATE_MODEL').__name__, - "code_path": os.path.join(path, "model.py"), - "weight_path": template_model_path, - } - print(f"Template model loaded: {json.dumps(metadata, indent=4)}") - return model - - -def load_template_data_processor(path): - spec = importlib.util.spec_from_file_location("template_model", os.path.join(path, "model.py")) - module = importlib.util.module_from_spec(spec) - spec.loader.exec_module(module) - if hasattr(module, 'TEMPLATE_DATA_PROCESSOR'): - processor = getattr(module, 'TEMPLATE_DATA_PROCESSOR') - return processor - else: - return None - - -class TemplatePipeline(torch.nn.Module): - def __init__( - self, - torch_dtype: torch.dtype = torch.bfloat16, - device: Union[str, torch.device] = get_device_type(), - model_configs: list[ModelConfig] = [], - lazy_loading: bool = False, - ): - super().__init__() - self.torch_dtype = torch_dtype - self.device = device - self.model_configs = model_configs - self.lazy_loading = lazy_loading - if lazy_loading: - for model_config in model_configs: - TemplatePipeline.check_vram_config(model_config) - model_config.download_if_necessary() - self.models = None - else: - models = [] - for model_config in model_configs: - TemplatePipeline.check_vram_config(model_config) - model_config.download_if_necessary() - model = load_template_model(model_config.path, torch_dtype=torch_dtype, device=device, state_dict=model_config.state_dict) - models.append(model) - self.models = torch.nn.ModuleList(models) - - def merge_kv_cache(self, kv_cache_list: List[KVCache]) -> KVCache: - names = {} - for kv_cache in kv_cache_list: - for name in kv_cache: - names[name] = None - kv_cache_merged = {} - for name in names: - kv_list = [kv_cache.get(name) for kv_cache in kv_cache_list] - kv_list = [kv for kv in kv_list if kv is not None] - if len(kv_list) > 0: - k = torch.concat([kv[0] for kv in kv_list], dim=1) - v = torch.concat([kv[1] for kv in kv_list], dim=1) - kv_cache_merged[name] = (k, v) - return kv_cache_merged - - def merge_sequential_embeddings(self, embeddings: List[torch.Tensor]): - return torch.concat(embeddings, dim=-2) - - def merge_template_cache(self, template_cache_list): - params = sorted(list(set(sum([list(template_cache.keys()) for template_cache in template_cache_list], [])))) - template_cache_merged = {} - for param in params: - data = [template_cache[param] for template_cache in template_cache_list if param in template_cache] - if param == "kv_cache": - data = self.merge_kv_cache(data) - elif param == "lora": - data = merge_lora(data) - elif param == "text_embedding": - data = self.merge_sequential_embeddings(data) - elif len(data) == 1: - data = data[0] - else: - print(f"Conflict detected: `{param}` appears in the outputs of multiple Template models. Only the first one will be retained.") - data = data[0] - template_cache_merged[param] = data - return template_cache_merged - - @staticmethod - def check_vram_config(model_config: ModelConfig): - params = [ - model_config.offload_device, model_config.offload_dtype, - model_config.onload_device, model_config.onload_dtype, - model_config.preparing_device, model_config.preparing_dtype, - model_config.computation_device, model_config.computation_dtype, - ] - for param in params: - if param is not None: - warnings.warn("TemplatePipeline doesn't support VRAM management. VRAM config will be ignored.") - - @staticmethod - def from_pretrained( - torch_dtype: torch.dtype = torch.bfloat16, - device: Union[str, torch.device] = get_device_type(), - model_configs: list[ModelConfig] = [], - lazy_loading: bool = False, - ): - pipe = TemplatePipeline(torch_dtype, device, model_configs, lazy_loading) - return pipe - - def fetch_model(self, model_id): - if self.lazy_loading: - model_config = self.model_configs[model_id] - model_config.download_if_necessary() - model = load_template_model(model_config.path, torch_dtype=self.torch_dtype, device=self.device, state_dict=model_config.state_dict) - else: - model = self.models[model_id] - return model - - def call_single_side(self, pipe=None, inputs: List[Dict] = None): - model = None - onload_model_id = -1 - template_cache = [] - for i in inputs: - model_id = i.get("model_id", 0) - if model_id != onload_model_id: - model = self.fetch_model(model_id) - onload_model_id = model_id - cache = model.process_inputs(pipe=pipe, **i) - cache = model.forward(pipe=pipe, **cache) - template_cache.append(cache) - template_cache = self.merge_template_cache(template_cache) - return template_cache - - @torch.no_grad() - def __call__( - self, - pipe=None, - template_inputs: List[Dict] = None, - negative_template_inputs: List[Dict] = None, - **kwargs, - ): - template_cache = self.call_single_side(pipe=pipe, inputs=template_inputs or []) - negative_template_cache = self.call_single_side(pipe=pipe, inputs=negative_template_inputs or []) - required_params = list(inspect.signature(pipe.__call__).parameters.keys()) - for param in template_cache: - if param in required_params: - kwargs[param] = template_cache[param] - else: - print(f"`{param}` is not included in the inputs of `{pipe.__class__.__name__}`. This parameter will be ignored.") - for param in negative_template_cache: - if "negative_" + param in required_params: - kwargs["negative_" + param] = negative_template_cache[param] - else: - print(f"`{'negative_' + param}` is not included in the inputs of `{pipe.__class__.__name__}`. This parameter will be ignored.") - return pipe(**kwargs) +import torch, os, importlib, warnings, json, inspect +from typing import Dict, List, Tuple, Union +from ..core import ModelConfig, load_model +from ..core.device.npu_compatible_device import get_device_type +from ..utils.lora.merge import merge_lora + + +KVCache = Dict[str, Tuple[torch.Tensor, torch.Tensor]] + + +class TemplateModel(torch.nn.Module): + def __init__(self): + super().__init__() + + @torch.no_grad() + def process_inputs(self, **kwargs): + return {} + + def forward(self, **kwargs): + raise NotImplementedError() + + +def check_template_model_format(model): + if not hasattr(model, "process_inputs"): + raise NotImplementedError("`process_inputs` is not implemented in the Template model.") + if "kwargs" not in inspect.signature(model.process_inputs).parameters: + raise NotImplementedError("`**kwargs` is not included in `process_inputs`.") + if not hasattr(model, "forward"): + raise NotImplementedError("`forward` is not implemented in the Template model.") + if "kwargs" not in inspect.signature(model.forward).parameters: + raise NotImplementedError("`**kwargs` is not included in `forward`.") + + +def load_template_model(path, torch_dtype=torch.bfloat16, device="cuda", verbose=1, state_dict=None): + spec = importlib.util.spec_from_file_location("template_model", os.path.join(path, "model.py")) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + template_model_path = getattr(module, 'TEMPLATE_MODEL_PATH') if hasattr(module, 'TEMPLATE_MODEL_PATH') else None + if template_model_path is not None: + # With `TEMPLATE_MODEL_PATH`, a pretrained model will be loaded. + model = load_model( + model_class=getattr(module, 'TEMPLATE_MODEL'), + config=getattr(module, 'TEMPLATE_MODEL_CONFIG') if hasattr(module, 'TEMPLATE_MODEL_CONFIG') else None, + path=os.path.join(path, getattr(module, 'TEMPLATE_MODEL_PATH')), + torch_dtype=torch_dtype, + device=device, + state_dict=state_dict, + ) + else: + # Without `TEMPLATE_MODEL_PATH`, a randomly initialized model or a non-model module will be loaded. + model = module.TEMPLATE_MODEL() + if hasattr(model, "to"): + model = model.to(dtype=torch_dtype, device=device) + if hasattr(model, "eval"): + model = model.eval() + check_template_model_format(model) + if verbose > 0: + metadata = { + "model_architecture": getattr(module, 'TEMPLATE_MODEL').__name__, + "code_path": os.path.join(path, "model.py"), + "weight_path": template_model_path, + } + print(f"Template model loaded: {json.dumps(metadata, indent=4)}") + return model + + + + +class TemplatePipeline(torch.nn.Module): + def __init__( + self, + torch_dtype: torch.dtype = torch.bfloat16, + device: Union[str, torch.device] = get_device_type(), + model_configs: list[ModelConfig] = [], + lazy_loading: bool = False, + ): + super().__init__() + self.torch_dtype = torch_dtype + self.device = device + self.model_configs = model_configs + self.lazy_loading = lazy_loading + if lazy_loading: + for model_config in model_configs: + TemplatePipeline.check_vram_config(model_config) + model_config.download_if_necessary() + self.models = None + else: + models = [] + for model_config in model_configs: + TemplatePipeline.check_vram_config(model_config) + model_config.download_if_necessary() + model = load_template_model(model_config.path, torch_dtype=torch_dtype, device=device, state_dict=model_config.state_dict) + models.append(model) + self.models = torch.nn.ModuleList(models) + + def merge_kv_cache(self, kv_cache_list: List[KVCache]) -> KVCache: + names = {} + for kv_cache in kv_cache_list: + for name in kv_cache: + names[name] = None + kv_cache_merged = {} + for name in names: + kv_list = [kv_cache.get(name) for kv_cache in kv_cache_list] + kv_list = [kv for kv in kv_list if kv is not None] + if len(kv_list) > 0: + k = torch.concat([kv[0] for kv in kv_list], dim=1) + v = torch.concat([kv[1] for kv in kv_list], dim=1) + kv_cache_merged[name] = (k, v) + return kv_cache_merged + + def merge_sequential_embeddings(self, embeddings: List[torch.Tensor]): + return torch.concat(embeddings, dim=-2) + + def merge_template_cache(self, template_cache_list): + params = sorted(list(set(sum([list(template_cache.keys()) for template_cache in template_cache_list], [])))) + template_cache_merged = {} + for param in params: + data = [template_cache[param] for template_cache in template_cache_list if param in template_cache] + if param == "kv_cache": + data = self.merge_kv_cache(data) + elif param == "lora": + data = merge_lora(data) + elif param == "text_embedding": + data = self.merge_sequential_embeddings(data) + elif len(data) == 1: + data = data[0] + else: + print(f"Conflict detected: `{param}` appears in the outputs of multiple Template models. Only the first one will be retained.") + data = data[0] + template_cache_merged[param] = data + return template_cache_merged + + @staticmethod + def check_vram_config(model_config: ModelConfig): + params = [ + model_config.offload_device, model_config.offload_dtype, + model_config.onload_device, model_config.onload_dtype, + model_config.preparing_device, model_config.preparing_dtype, + model_config.computation_device, model_config.computation_dtype, + ] + for param in params: + if param is not None: + warnings.warn("TemplatePipeline doesn't support VRAM management. VRAM config will be ignored.") + + @staticmethod + def from_pretrained( + torch_dtype: torch.dtype = torch.bfloat16, + device: Union[str, torch.device] = get_device_type(), + model_configs: list[ModelConfig] = [], + lazy_loading: bool = False, + ): + pipe = TemplatePipeline(torch_dtype, device, model_configs, lazy_loading) + return pipe + + def fetch_model(self, model_id): + if self.lazy_loading: + model_config = self.model_configs[model_id] + model_config.download_if_necessary() + model = load_template_model(model_config.path, torch_dtype=self.torch_dtype, device=self.device, state_dict=model_config.state_dict) + else: + model = self.models[model_id] + return model + + def call_single_side(self, pipe=None, inputs: List[Dict] = None): + model = None + onload_model_id = -1 + template_cache = [] + for i in inputs: + model_id = i.get("model_id", 0) + if model_id != onload_model_id: + model = self.fetch_model(model_id) + onload_model_id = model_id + cache = model.process_inputs(pipe=pipe, **i) + cache = model.forward(pipe=pipe, **cache) + template_cache.append(cache) + template_cache = self.merge_template_cache(template_cache) + return template_cache + + @torch.no_grad() + def __call__( + self, + pipe=None, + template_inputs: List[Dict] = None, + negative_template_inputs: List[Dict] = None, + **kwargs, + ): + template_cache = self.call_single_side(pipe=pipe, inputs=template_inputs or []) + negative_template_cache = self.call_single_side(pipe=pipe, inputs=negative_template_inputs or []) + required_params = list(inspect.signature(pipe.__call__).parameters.keys()) + for param in template_cache: + if param in required_params: + kwargs[param] = template_cache[param] + else: + print(f"`{param}` is not included in the inputs of `{pipe.__class__.__name__}`. This parameter will be ignored.") + for param in negative_template_cache: + if "negative_" + param in required_params: + kwargs["negative_" + param] = negative_template_cache[param] + else: + print(f"`{'negative_' + param}` is not included in the inputs of `{pipe.__class__.__name__}`. This parameter will be ignored.") + return pipe(**kwargs) diff --git a/deployment/runtime/qwen/diffsynth/pipelines/qwen_image_21_pbr.py b/deployment/runtime/qwen/diffsynth/pipelines/qwen_image_21_pbr.py index 7d2db08bdec04da95d6182dba5ac30e20bf7e1a7..0005448a806f6b7bcee63afc7cef5878dfbe351a 100644 --- a/deployment/runtime/qwen/diffsynth/pipelines/qwen_image_21_pbr.py +++ b/deployment/runtime/qwen/diffsynth/pipelines/qwen_image_21_pbr.py @@ -1,98 +1,90 @@ -"""Joint Albedo/MRA latent prediction using the Qwen-Image 2.1 backbone.""" - -import torch -from torch import nn - -from .qwen_image_21 import patchify, unpatchify - - -class QwenImage21PBRModel(nn.Module): - """Predict two separate material maps with a shared noisy target group. - - The modality axis is ordered ``(albedo, mra)``. MRA pixel semantics belong - to the dataset contract; this module does not infer or repack AO/alpha. - Conditions remain clean, and each map retains its own spatial RoPE block. - """ - - modalities = ("albedo", "mra") - - def __init__(self, dit): - super().__init__() - self.dit = dit - channels = dit.img_in.in_features - if channels != dit.out_channels: - raise ValueError("PBR requires equal input and output latent channels") - self.modality_embeddings = nn.Parameter(torch.zeros(2, channels)) - - def forward( - self, - latents, - timestep, - prompt_embeds, - prompt_embeds_mask, - edit_image_pad_mask, - edit_latents=None, - use_gradient_checkpointing=False, - use_gradient_checkpointing_offload=False, - use_flex_attention=False, - kv_cache=None, - gradient_checkpointing_interval=1, - ): - if latents.ndim != 5 or latents.shape[1] != len(self.modalities): - raise ValueError("Expected latents shaped [batch, 2, channels, height, width]") - batch_size, _, channels, height, width = latents.shape - if channels != self.modality_embeddings.shape[1]: - raise ValueError("Target latent channels do not match the DiT") - if height % 2 or width % 2: - raise ValueError("Target latent dimensions must be even (pixel dimensions divisible by 32)") - if prompt_embeds.shape[0] != batch_size or edit_image_pad_mask.shape[0] != batch_size: - raise ValueError("Condition and target batch sizes must match") - if not torch.equal(edit_image_pad_mask, edit_image_pad_mask[:1].expand_as(edit_image_pad_mask)): - raise ValueError("Batch elements must share the same image-token layout") - edit_latents = [] if edit_latents is None else edit_latents - for condition in edit_latents: - if condition.ndim != 4 or condition.shape[:2] != (batch_size, channels): - raise ValueError("Each condition must be shaped [batch, channels, height, width]") - - target_tokens = [] - for index in range(len(self.modalities)): - tokens = patchify(latents[:, index]) - embedding = self.modality_embeddings[index].to(device=tokens.device, dtype=tokens.dtype) - target_tokens.append(tokens + embedding) - hidden_states = torch.cat([*[patchify(image) for image in edit_latents], *target_tokens], dim=1) - target_seq_len = len(self.modalities) * height * width - target_image_pad_mask = torch.ones( - (batch_size, target_seq_len // 4), dtype=torch.bool, device=latents.device - ) - image_pad_mask = torch.cat([edit_image_pad_mask.bool(), target_image_pad_mask], dim=1) - img_shapes = [[ - *[(1, image.shape[2], image.shape[3]) for image in edit_latents], - *[(1, height, width) for _ in self.modalities], - ] for _ in range(batch_size)] - model_output = self.dit( - hidden_states=hidden_states, - timestep=timestep / 1000, - encoder_hidden_states=prompt_embeds, - encoder_hidden_states_mask=prompt_embeds_mask, - img_shapes=img_shapes, - img_mask=image_pad_mask, - kv_cache=kv_cache, - use_gradient_checkpointing=use_gradient_checkpointing, - use_gradient_checkpointing_offload=use_gradient_checkpointing_offload, - use_flex_attention=use_flex_attention, - num_target_images=len(self.modalities), - gradient_checkpointing_interval=gradient_checkpointing_interval, - ) - target_output = model_output[:, -target_seq_len:] - return torch.stack([ - unpatchify(tokens, height, width) - for tokens in target_output.split(height * width, dim=1) - ], dim=1) - - -def pbr_flow_matching_losses(prediction, target): - """Equal-weight latent flow losses, with separate material diagnostics.""" - if prediction.shape != target.shape or prediction.ndim != 5 or prediction.shape[1] != 2: - raise ValueError("Prediction and target must have matching [batch, 2, channels, height, width] shapes") - losses = (prediction.float() - target.float()).square().mean(dim=(0, 2, 3, 4)) - return {"albedo": losses[0], "mra": losses[1], "loss": losses.mean()} +"""Joint Albedo/MRA latent prediction using the Qwen-Image 2.1 backbone.""" + +import torch +from torch import nn + +from .qwen_image_21 import patchify, unpatchify + + +class QwenImage21PBRModel(nn.Module): + """Predict two separate material maps with a shared noisy target group. + + The modality axis is ordered ``(albedo, mra)``. MRA pixel semantics belong + to the dataset contract; this module does not infer or repack AO/alpha. + Conditions remain clean, and each map retains its own spatial RoPE block. + """ + + modalities = ("albedo", "mra") + + def __init__(self, dit): + super().__init__() + self.dit = dit + channels = dit.img_in.in_features + if channels != dit.out_channels: + raise ValueError("PBR requires equal input and output latent channels") + self.modality_embeddings = nn.Parameter(torch.zeros(2, channels)) + + def forward( + self, + latents, + timestep, + prompt_embeds, + prompt_embeds_mask, + edit_image_pad_mask, + edit_latents=None, + use_gradient_checkpointing=False, + use_gradient_checkpointing_offload=False, + use_flex_attention=False, + kv_cache=None, + gradient_checkpointing_interval=1, + ): + if latents.ndim != 5 or latents.shape[1] != len(self.modalities): + raise ValueError("Expected latents shaped [batch, 2, channels, height, width]") + batch_size, _, channels, height, width = latents.shape + if channels != self.modality_embeddings.shape[1]: + raise ValueError("Target latent channels do not match the DiT") + if height % 2 or width % 2: + raise ValueError("Target latent dimensions must be even (pixel dimensions divisible by 32)") + if prompt_embeds.shape[0] != batch_size or edit_image_pad_mask.shape[0] != batch_size: + raise ValueError("Condition and target batch sizes must match") + if not torch.equal(edit_image_pad_mask, edit_image_pad_mask[:1].expand_as(edit_image_pad_mask)): + raise ValueError("Batch elements must share the same image-token layout") + edit_latents = [] if edit_latents is None else edit_latents + for condition in edit_latents: + if condition.ndim != 4 or condition.shape[:2] != (batch_size, channels): + raise ValueError("Each condition must be shaped [batch, channels, height, width]") + + target_tokens = [] + for index in range(len(self.modalities)): + tokens = patchify(latents[:, index]) + embedding = self.modality_embeddings[index].to(device=tokens.device, dtype=tokens.dtype) + target_tokens.append(tokens + embedding) + hidden_states = torch.cat([*[patchify(image) for image in edit_latents], *target_tokens], dim=1) + target_seq_len = len(self.modalities) * height * width + target_image_pad_mask = torch.ones( + (batch_size, target_seq_len // 4), dtype=torch.bool, device=latents.device + ) + image_pad_mask = torch.cat([edit_image_pad_mask.bool(), target_image_pad_mask], dim=1) + img_shapes = [[ + *[(1, image.shape[2], image.shape[3]) for image in edit_latents], + *[(1, height, width) for _ in self.modalities], + ] for _ in range(batch_size)] + model_output = self.dit( + hidden_states=hidden_states, + timestep=timestep / 1000, + encoder_hidden_states=prompt_embeds, + encoder_hidden_states_mask=prompt_embeds_mask, + img_shapes=img_shapes, + img_mask=image_pad_mask, + kv_cache=kv_cache, + use_gradient_checkpointing=use_gradient_checkpointing, + use_gradient_checkpointing_offload=use_gradient_checkpointing_offload, + use_flex_attention=use_flex_attention, + num_target_images=len(self.modalities), + gradient_checkpointing_interval=gradient_checkpointing_interval, + ) + target_output = model_output[:, -target_seq_len:] + return torch.stack([ + unpatchify(tokens, height, width) + for tokens in target_output.split(height * width, dim=1) + ], dim=1) diff --git a/deployment/runtime/qwen/examples/qwen_image_21/pbr/infer.py b/deployment/runtime/qwen/examples/qwen_image_21/pbr/infer.py index 7e971b0ffb43420ca54c9e1d57acde05f83e041a..753ecd9f2f4de857e10d8f5b57a253fddafb681a 100644 --- a/deployment/runtime/qwen/examples/qwen_image_21/pbr/infer.py +++ b/deployment/runtime/qwen/examples/qwen_image_21/pbr/infer.py @@ -1,390 +1,371 @@ -#!/usr/bin/env python3 -"""RGB/RGBA-only endpoint inference. No GT, data manifests, downloads or hashes.""" -from __future__ import annotations -import argparse -from contextlib import nullcontext -import gc -import importlib.util -import json -import math -import os -from pathlib import Path -import time - -import numpy as np -from PIL import Image -import torch -from safetensors import safe_open -from safetensors.torch import save_file - -from train import DEFAULTS, build_model, decode_pair, endpoint_noise -from validate_local import checkpoint_component, explicit_load -from encode_cache import PROMPT -from inference_offload import ( - release_device_cache, sequential_cpu_offload, stage_on_device, - text_encoder_offload_blocks, -) - -# The isolated Qwen environment does not include the application package path. -_timing_helper_path = os.environ.get('SUPRATEX_TIMING_HELPER') -if _timing_helper_path: - _timing_spec = importlib.util.spec_from_file_location('_supratex_timing', _timing_helper_path) - if _timing_spec is None or _timing_spec.loader is None: - raise RuntimeError('could not load the configured Supra-Tex timing helper') - _timing_helper = importlib.util.module_from_spec(_timing_spec) - _timing_spec.loader.exec_module(_timing_helper) - span = _timing_helper.span -else: - span = nullcontext - - -def prepare_image(path, resolution): - if resolution < 32 or resolution % 32: - raise ValueError('resolution must be positive and divisible by 32') - with Image.open(path) as source: - had_alpha = 'A' in source.getbands() or 'transparency' in source.info - original = source.convert('RGBA') - width, height = original.size - scale = min(1., resolution / max(width, height)) - content = (max(1, round(width * scale)), max(1, round(height * scale))) - aligned = tuple(math.ceil(value / 32) * 32 for value in content) - offset = tuple((aligned[i] - content[i]) // 2 + (resolution - aligned[i]) // 2 for i in range(2)) - resized = original.resize(content, Image.Resampling.LANCZOS) if original.size != content else original.copy() - square = Image.new('RGBA', (resolution, resolution), (0, 0, 0, 0)) - square.paste(resized, offset) - transform = dict(source_size=[width, height], content_size=list(content), padded_size=[resolution, resolution], - offset=list(offset), scale=scale, input_had_alpha=had_alpha, - resampling='LANCZOS, same two-stage centered padding as training, no crop or alpha compositing', - coverage='input alpha retained' if had_alpha else 'all-one source coverage; no segmentation inferred') - return original, square, transform - - -def load_checkpoint(path): - path = Path(path) - if path.suffix == '.safetensors': - from safetensors import safe_open - import math - - adapter = json.loads(path.with_name('adapter_config.json').read_text()) - if (not isinstance(adapter, dict) or adapter.get('format') != 'supratex_qwen_material_adapter' - or type(adapter.get('schema_version')) is not int or adapter['schema_version'] != 1 - or adapter.get('adapter_type') != 'qwen_image_21_pbr_endpoint' - or adapter.get('base_model_name_or_path') != 'Qwen/Qwen-Image-2.1' - or adapter.get('weight_file') != path.name - or type(adapter.get('source_step')) is not int or adapter['source_step'] != 100000): - raise ValueError('unsupported Qwen material adapter config') - inference = adapter.get('inference_config') - booleans = {'lora_output_projection', 'omit_all_valid_mask', 'bf16_autocast', 'deterministic_algorithms'} - required = booleans | {'resolution', 'noise_seed', 'lora_rank', 'lora_alpha', 'attention_backend'} - if not isinstance(inference, dict) or set(inference) != required: - raise ValueError('Qwen adapter requires complete inference config') - if (type(inference['resolution']) is not int or inference['resolution'] != 1024 - or type(inference['noise_seed']) is not int or inference['noise_seed'] != 42 - or inference['attention_backend'] != 'torch' - or type(inference['lora_rank']) is not int or inference['lora_rank'] <= 0 - or type(inference['lora_alpha']) not in {int, float} - or not math.isfinite(inference['lora_alpha']) or inference['lora_alpha'] <= 0 - or any(type(inference[key]) is not bool for key in booleans)): - raise ValueError('invalid Qwen adapter inference config') - tensors = adapter.get('tensors') - if not isinstance(tensors, dict) or not tensors: - raise ValueError('Qwen adapter requires a nonempty tensor manifest') - with safe_open(str(path), framework='pt', device='cpu') as reader: - metadata = reader.metadata() or {} - if (metadata.get('format') != 'pt' or metadata.get('adapter_format') != adapter['format'] - or metadata.get('schema_version') != str(adapter['schema_version']) - or set(reader.keys()) != set(tensors)): - raise ValueError('Qwen adapter tensor keys or metadata mismatch') - for name, specification in tensors.items(): - header = reader.get_slice(name) - if (not isinstance(specification, dict) - or not isinstance(specification.get('shape'), list) - or any(type(size) is not int or size < 0 for size in specification['shape']) - or specification['shape'] != header.get_shape() - or specification.get('dtype') not in {'F64', 'F32', 'F16', 'BF16'} - or specification['dtype'] != header.get_dtype()): - raise ValueError(f'Qwen adapter tensor shape or dtype mismatch: {name}') - weights = {name: reader.get_tensor(name) for name in reader.keys()} - config = DEFAULTS | inference | {'model_root': adapter['base_model_name_or_path']} - return config, weights, adapter['source_step'] - marker = json.loads(path.with_suffix('.complete.json').read_text()) - if marker.get('status') != 'complete' or marker.get('bytes') != path.stat().st_size: - raise ValueError('checkpoint completion record is absent or mismatched') - # SupraTex adaptation: public inference weights use restricted loading. - # Legacy trainer checkpoints also contain NumPy uint32 RNG arrays; these - # exact reconstruction types retain compatibility without arbitrary pickle. - from numpy.core.multiarray import _reconstruct - with torch.serialization.safe_globals([_reconstruct, (_reconstruct, 'numpy.core.multiarray._reconstruct'), - (_reconstruct, 'numpy._core.multiarray._reconstruct'), - np.ndarray, np.dtype, type(np.dtype('uint32'))]): - state = torch.load(path, map_location='cpu', weights_only=True) - if state.get('format') != 'qwen_pbr_endpoint_v1' or marker.get('step') != state.get('step'): - raise ValueError('unsupported checkpoint or inconsistent step') - if not isinstance(state.get('config'), dict) or not isinstance(state.get('model'), dict): - raise ValueError('checkpoint needs complete config and trainable state') - config = DEFAULTS | state['config'] - for key in ('model_root', 'resolution', 'lora_rank', 'lora_alpha', 'noise_seed'): - if key not in state['config']: - raise ValueError(f'checkpoint config lacks {key}') - return config, state['model'], state['step'] - - -def strict_load_trainable(model, weights): - expected = {name: parameter for name, parameter in model.named_parameters() if parameter.requires_grad} - if expected.keys() != weights.keys(): - raise ValueError('checkpoint trainable key set differs from constructed model') - with torch.no_grad(): - for name, parameter in expected.items(): - tensor = weights[name] - if tensor.shape != parameter.shape or tensor.dtype != parameter.dtype or not torch.isfinite(tensor).all(): - raise ValueError(f'checkpoint tensor mismatch/nonfinite: {name}') - parameter.copy_(tensor) - model.requires_grad_(False).eval() - - -def canonical_conditions(values, resolution): - """Only retain consumed conditioning; material targets are never loaded.""" - values = dict(values) - layout = values['edit_image_pad_mask'].bool() - prompt = values.get('prompt_embeds') - if prompt is None: - text = values['prompt_text_embeds'] - if layout.ndim != 2 or layout.shape[0] != 1 or text.shape[1] != int((~layout[0]).sum()): - raise ValueError('compact conditioning layout mismatch') - prompt = text.new_zeros((1, layout.shape[1], text.shape[-1])) - prompt[:, ~layout[0]] = text - if prompt.ndim != 3 or prompt.shape[0] != 1 or prompt.shape[2] != 4096 or layout.shape != prompt.shape[:2]: - raise ValueError('invalid prompt/layout shape') - # Model replaces image-placeholder rows from RGB VAE. Canonical zeros allow - # exact native-v2 comparison without changing consumed values. - prompt = prompt.clone() - prompt[:, layout[0]] = 0 - mask = values.get('prompt_embeds_mask') - mask = torch.ones_like(layout) if mask is None else mask.bool() - rgb = values['rgb'] - if rgb.shape != (1, 64, resolution // 16, resolution // 16) or mask.shape != layout.shape: - raise ValueError('invalid RGB latent/mask shape') - result = dict(rgb=rgb, prompt_embeds=prompt, prompt_embeds_mask=mask, edit_image_pad_mask=layout) - if any(value.is_floating_point() and not torch.isfinite(value).all() for value in result.values()): - raise ValueError('nonfinite conditioning') - return result - - -def load_conditions(path, resolution, device): - values = {} - keys = {'rgb', 'prompt_embeds', 'prompt_text_embeds', 'prompt_embeds_mask', 'edit_image_pad_mask'} - with safe_open(str(path), framework='pt', device=str(device)) as reader: - for key in keys & set(reader.keys()): - values[key] = reader.get_tensor(key) - return canonical_conditions(values, resolution) - - -def input_binding(path): - path = Path(os.path.abspath(path)); stat = path.stat() - return dict(path=str(path), size=stat.st_size, mtime_ns=stat.st_mtime_ns) - - -def verify_cache_input(cache, image, resolution, model_references): - metadata = json.loads(Path(cache).with_suffix('.json').read_text()) - binding = metadata.get('input', metadata.get('sources', {}).get('rgb')) - if binding != input_binding(image) or metadata.get('resolution') != resolution: - raise ValueError('cache conditioning is not bound to this exact input/stat/resolution') - if metadata.get('model_root') not in model_references or metadata.get('prompt') != PROMPT: - raise ValueError('cached model reference or prompt differs from checkpoint inference contract') - return metadata['model_root'] - - -def native_conditions(square, model_root, resolution, device, low_vram=False): - from transformers import AutoProcessor - from diffsynth.models.qwen_image_21_vae import QwenImage21VAE - from diffsynth.models.qwen_image_21_text_encoder import QwenImage21TextEncoder - from diffsynth.pipelines.qwen_image_21 import QwenImage21Pipeline, QwenImage21Unit_EditImageEmbedder, QwenImage21Unit_PromptEmbedder - from diffsynth.utils.state_dict_converters.qwen_image_21_text_encoder import QwenImage21TextEncoderStateDictConverter - processor = AutoProcessor.from_pretrained(str(model_root / 'processor'), local_files_only=True) - processor.image_processor.size = {'shortest_edge': resolution ** 2, 'longest_edge': resolution ** 2} - pipe = QwenImage21Pipeline(device=device, torch_dtype=torch.bfloat16) - pipe.processor = processor - loading_device = 'cpu' if low_vram else device - vae = explicit_load(QwenImage21VAE, checkpoint_component(model_root, 'vae'), loading_device, torch.bfloat16) - pipe.vae = vae - with stage_on_device(vae, device, enabled=low_vram), torch.inference_mode(): - edited = QwenImage21Unit_EditImageEmbedder().process(pipe, edit_image=[square], height=resolution, width=resolution, - tiled=False, tile_size=256, tile_stride=192) - if edited['edit_image'][0].size != square.size: - raise ValueError('native edit path changed dimensions') - rgb = edited['edit_latents'][0] - encoder = explicit_load(QwenImage21TextEncoder, checkpoint_component(model_root, 'text_encoder'), loading_device, - torch.bfloat16, QwenImage21TextEncoderStateDictConverter) - pipe.text_encoder = encoder - blocks = text_encoder_offload_blocks(encoder) if low_vram else () - with sequential_cpu_offload(encoder, blocks, device, enabled=low_vram), torch.inference_mode(): - values = QwenImage21Unit_PromptEmbedder().process(pipe, prompt=PROMPT, edit_image=[square]) - conditions = canonical_conditions(dict(values, rgb=rgb), resolution) - pipe.text_encoder = None - del encoder, pipe, values, edited, blocks - gc.collect(); torch.cuda.empty_cache() - return conditions, vae - - -def predict(model, conditions, cfg): - # Shape comes from observed RGB, never a GT or fabricated target tensor. - shape = (conditions['rgb'].shape[0], 2, *conditions['rgb'].shape[1:]) - shape_carrier = conditions['rgb'].new_empty(shape) - noise = endpoint_noise(shape_carrier, cfg['noise_seed']) - mask = conditions['prompt_embeds_mask'] - if cfg['omit_all_valid_mask'] and bool(mask.all()): - mask = None - with torch.inference_mode(), torch.autocast('cuda', dtype=torch.bfloat16, enabled=cfg['bf16_autocast']): - velocity = model(latents=noise, timestep=torch.full((noise.shape[0],), 1000., device=noise.device), - prompt_embeds=conditions['prompt_embeds'], prompt_embeds_mask=mask, - edit_image_pad_mask=conditions['edit_image_pad_mask'], edit_latents=[conditions['rgb']], - use_gradient_checkpointing=False, use_flex_attention=False, kv_cache=None) - prediction = noise.float() - velocity.float() - if not torch.isfinite(prediction).all(): - raise FloatingPointError('nonfinite material latent prediction') - return prediction - - -def predict_and_decode(model, vae, conditions, cfg, low_vram=False): - device = conditions['rgb'].device - try: - with sequential_cpu_offload(model, model.dit.transformer_blocks, device, enabled=low_vram): - latent = predict(model, conditions, cfg) - finally: - if low_vram: - # Qwen stores RoPE tables as an ordinary list, outside Module buffers. - model.dit.pos_embed.freqs = [value.cpu() for value in model.dit.pos_embed.freqs] - release_device_cache(device) - with stage_on_device(vae, device, enabled=low_vram), torch.inference_mode(): - decoded = decode_pair(vae, latent) - return latent, decoded - - -def output_maps(decoded, square, original, transform, restore_size): - if decoded.shape[0:3] != (1, 2, 4) or not torch.isfinite(decoded).all(): - raise ValueError('invalid/nonfinite decoded material pair') - arrays = np.clip(np.rint(decoded[0].detach().float().cpu().permute(0, 2, 3, 1).numpy() * 255), 0, 255).astype(np.uint8) - diagnostics = [Image.fromarray(array[:, :, 3], 'L') for array in arrays] - maps = [] - for modality, array in enumerate(arrays): - color = Image.fromarray(array[:, :, :3], 'RGB') - coverage = square.getchannel('A') - if restore_size: - x, y = transform['offset']; w, h = transform['content_size'] - color = color.crop((x, y, x + w, y + h)) - if color.size != original.size: - color = color.resize(original.size, Image.Resampling.LANCZOS) - coverage = original.getchannel('A') # exact original coverage, no second interpolation - out = np.array(color.convert('RGBA')) - out[:, :, 3] = np.array(coverage) - if modality == 1: - out[:, :, 2] = 255 # placeholder, not material opacity or AO prediction - maps.append(Image.fromarray(out, 'RGBA')) - return maps, diagnostics - - -def main(): - parser = argparse.ArgumentParser(description=__doc__) - parser.add_argument('--checkpoint', required=True, type=Path, - help='Native adapter_model.safetensors with sibling adapter_config.json, or legacy PT with completion marker') - parser.add_argument('--input', required=True, type=Path) - parser.add_argument('--output-dir', required=True, type=Path) - parser.add_argument('--model-root', type=Path, help='Explicit loading location only; checkpoint model reference retained') - parser.add_argument('--conditions-cache', type=Path, help='Read RGB conditioning only; sibling JSON must bind input/model/resolution') - parser.add_argument('--compare-cache', type=Path, help='Native path plus exact conditioning/latent/decoded comparison against cache') - parser.add_argument('--comparison-atol', type=float, default=0.) - parser.add_argument('--save-conditions', action='store_true') - parser.add_argument('--restore-input-size', action='store_true') - parser.add_argument('--save-predicted-coverage', action='store_true') - parser.add_argument('--device', default='cuda:0') - parser.add_argument('--low-vram', action='store_true', default=os.environ.get('SUPRATEX_LOW_VRAM') == '1', - help='CPU stage/block offload; preserves weights and sampling, uses more CPU RAM and transfers (also SUPRATEX_LOW_VRAM=1)') - args = parser.parse_args() - if args.conditions_cache and args.compare_cache: - parser.error('--conditions-cache and --compare-cache are mutually exclusive') - if args.comparison_atol < 0: - parser.error('--comparison-atol must be nonnegative') - with span('qwen.checkpoint_load'): - cfg, weights, step = load_checkpoint(args.checkpoint) - os.environ['DIFFSYNTH_ATTENTION_IMPLEMENTATION'] = cfg['attention_backend'] - os.environ['HF_HUB_OFFLINE'] = '1'; os.environ['TRANSFORMERS_OFFLINE'] = '1' - if cfg['deterministic_algorithms']: - os.environ.setdefault('CUBLAS_WORKSPACE_CONFIG', ':4096:8') - torch.use_deterministic_algorithms(cfg['deterministic_algorithms']) - device = torch.device(args.device) - torch.cuda.set_device(device) - reference_root = cfg['model_root'] - model_root = Path(os.path.abspath(args.model_root or reference_root)) - # An explicit load-location override is an operator-supplied alias, never an - # implicit identity rewrite or a digest-based equivalence claim. - model_references = {reference_root, str(model_root)} - original, square, transform = prepare_image(args.input, cfg['resolution']) - if args.output_dir.exists(): - raise FileExistsError('choose a fresh inference output directory') - args.output_dir.mkdir(parents=True) - started = time.perf_counter() - torch.cuda.reset_peak_memory_stats(device) - with span('qwen.conditions'): - if args.conditions_cache: - verify_cache_input(args.conditions_cache, args.input, cfg['resolution'], model_references) - conditions = load_conditions(args.conditions_cache, cfg['resolution'], device) - from diffsynth.models.qwen_image_21_vae import QwenImage21VAE - vae = explicit_load(QwenImage21VAE, checkpoint_component(model_root, 'vae'), - 'cpu' if args.low_vram else device, torch.bfloat16) - else: - conditions, vae = native_conditions(square, model_root, cfg['resolution'], device, low_vram=args.low_vram) - with span('qwen.transformer_load'): - load_cfg = cfg | {'model_root': str(model_root)} - model = build_model(load_cfg, 'cpu' if args.low_vram else device) - with span('qwen.adapter_load'): - strict_load_trainable(model, weights) - del weights - with span('qwen.predict_decode'): - latent, decoded = predict_and_decode(model, vae, conditions, cfg, low_vram=args.low_vram) - comparison = None - if args.compare_cache: - verify_cache_input(args.compare_cache, args.input, cfg['resolution'], model_references) - cached = load_conditions(args.compare_cache, cfg['resolution'], device) - differences = {key: float((conditions[key].float() - cached[key].float()).abs().max()) for key in conditions} - other_latent, other_decoded = predict_and_decode(model, vae, cached, cfg, low_vram=args.low_vram) - differences.update(latent=float((latent - other_latent).abs().max()), decoded=float((decoded - other_decoded).abs().max())) - comparison = dict(max_abs=differences, atol=args.comparison_atol, passed=all(v <= args.comparison_atol for v in differences.values())) - if not comparison['passed']: - # Discriminate native/cache disagreement from repeated-kernel noise - # while holding conditions, checkpoint and process constant. - repeat_latent, repeat_decoded = predict_and_decode(model, vae, cached, cfg, low_vram=args.low_vram) - comparison['cached_repeat_control_max_abs'] = dict( - latent=float((other_latent-repeat_latent).abs().max()), - decoded=float((other_decoded-repeat_decoded).abs().max())) - (args.output_dir / 'comparison.json').write_text(json.dumps(comparison, indent=2) + '\n') - if not comparison['passed']: - raise ValueError('native/cache comparison exceeded requested tolerance') - with span('qwen.output_write'): - maps, diagnostics = output_maps(decoded, square, original, transform, args.restore_input_size) - for label, image, diag in zip(('albedo', 'mra'), maps, diagnostics): - image.save(args.output_dir / f'{label}.png') - if args.save_predicted_coverage: - diag.save(args.output_dir / f'{label}_predicted_coverage.png') - if args.save_conditions: - save_file({k: v.detach().cpu().contiguous() for k, v in conditions.items()}, str(args.output_dir / 'conditions.safetensors')) - (args.output_dir / 'conditions.json').write_text(json.dumps(dict(input=input_binding(args.input), resolution=cfg['resolution'], - model_root=reference_root, prompt=PROMPT, transform=transform), indent=2) + '\n') - report = dict(status='complete', checkpoint=os.path.abspath(args.checkpoint), checkpoint_step=step, - input=input_binding(args.input), loading_model_root=str(model_root), checkpoint_model_reference=reference_root, - model_override_is_operator_reference_not_hash_verified=bool(args.model_root), config=cfg, transform=transform, - accepted_model_reference_aliases=sorted(model_references), - low_vram=args.low_vram, - memory_policy='cpu_stages_and_sequential_transformer_blocks' if args.low_vram else 'resident_models', - restore_input_size=args.restore_input_size, restoration='crop padding; LANCZOS color interpolation if size changed; exact original coverage' if args.restore_input_size else 'training square canvas', - output_semantics='Albedo sRGB; MRA R=linear Metallic G=linear Roughness B=1 placeholder; A=input coverage', - boundary='Opaque material estimator. Input alpha retained as coverage; no segmentation, material transparency or hidden-layer recovery inferred.', - path='cached_conditioning_only' if args.conditions_cache else 'native_rgb_vae_and_per_image_vl', comparison=comparison, - peak_gpu_allocated_bytes=torch.cuda.max_memory_allocated(device), - peak_gpu_reserved_bytes=torch.cuda.max_memory_reserved(device), - elapsed_seconds=time.perf_counter()-started) - (args.output_dir / 'REPORT.json').write_text(json.dumps(report, indent=2) + '\n') - print(json.dumps(dict(status='complete', output=str(args.output_dir), checkpoint_step=step))) - - -if __name__ == '__main__': - main() +#!/usr/bin/env python3 +"""RGB/RGBA-only endpoint inference. No GT, data manifests, downloads or hashes.""" +from __future__ import annotations +import argparse +import gc +import json +import math +import os +from pathlib import Path +import time + +import numpy as np +from PIL import Image +import torch +from safetensors import safe_open +from safetensors.torch import save_file + +from inference_utils import ( + DEFAULTS, PROMPT, build_model, checkpoint_component, decode_pair, + endpoint_noise, explicit_load, +) +from inference_offload import ( + release_device_cache, sequential_cpu_offload, stage_on_device, + text_encoder_offload_blocks, +) + + +def prepare_image(path, resolution): + if resolution < 32 or resolution % 32: + raise ValueError('resolution must be positive and divisible by 32') + with Image.open(path) as source: + had_alpha = 'A' in source.getbands() or 'transparency' in source.info + original = source.convert('RGBA') + width, height = original.size + scale = min(1., resolution / max(width, height)) + content = (max(1, round(width * scale)), max(1, round(height * scale))) + aligned = tuple(math.ceil(value / 32) * 32 for value in content) + offset = tuple((aligned[i] - content[i]) // 2 + (resolution - aligned[i]) // 2 for i in range(2)) + resized = original.resize(content, Image.Resampling.LANCZOS) if original.size != content else original.copy() + square = Image.new('RGBA', (resolution, resolution), (0, 0, 0, 0)) + square.paste(resized, offset) + transform = dict(source_size=[width, height], content_size=list(content), padded_size=[resolution, resolution], + offset=list(offset), scale=scale, input_had_alpha=had_alpha, + resampling='LANCZOS, two-stage centered padding, no crop or alpha compositing', + coverage='input alpha retained' if had_alpha else 'all-one source coverage; no segmentation inferred') + return original, square, transform + + +def load_checkpoint(path): + path = Path(path) + if path.suffix == '.safetensors': + from safetensors import safe_open + import math + + adapter = json.loads(path.with_name('adapter_config.json').read_text()) + if (not isinstance(adapter, dict) or adapter.get('format') != 'supratex_qwen_material_adapter' + or type(adapter.get('schema_version')) is not int or adapter['schema_version'] != 1 + or adapter.get('adapter_type') != 'qwen_image_21_pbr_endpoint' + or adapter.get('base_model_name_or_path') != 'Qwen/Qwen-Image-2.1' + or adapter.get('weight_file') != path.name + or type(adapter.get('source_step')) is not int or adapter['source_step'] != 100000): + raise ValueError('unsupported Qwen material adapter config') + inference = adapter.get('inference_config') + booleans = {'lora_output_projection', 'omit_all_valid_mask', 'bf16_autocast', 'deterministic_algorithms'} + required = booleans | {'resolution', 'noise_seed', 'lora_rank', 'lora_alpha', 'attention_backend'} + if not isinstance(inference, dict) or set(inference) != required: + raise ValueError('Qwen adapter requires complete inference config') + if (type(inference['resolution']) is not int or inference['resolution'] != 1024 + or type(inference['noise_seed']) is not int or inference['noise_seed'] != 42 + or inference['attention_backend'] != 'torch' + or type(inference['lora_rank']) is not int or inference['lora_rank'] <= 0 + or type(inference['lora_alpha']) not in {int, float} + or not math.isfinite(inference['lora_alpha']) or inference['lora_alpha'] <= 0 + or any(type(inference[key]) is not bool for key in booleans)): + raise ValueError('invalid Qwen adapter inference config') + tensors = adapter.get('tensors') + if not isinstance(tensors, dict) or not tensors: + raise ValueError('Qwen adapter requires a nonempty tensor manifest') + with safe_open(str(path), framework='pt', device='cpu') as reader: + metadata = reader.metadata() or {} + if (metadata.get('format') != 'pt' or metadata.get('adapter_format') != adapter['format'] + or metadata.get('schema_version') != str(adapter['schema_version']) + or set(reader.keys()) != set(tensors)): + raise ValueError('Qwen adapter tensor keys or metadata mismatch') + for name, specification in tensors.items(): + header = reader.get_slice(name) + if (not isinstance(specification, dict) + or not isinstance(specification.get('shape'), list) + or any(type(size) is not int or size < 0 for size in specification['shape']) + or specification['shape'] != header.get_shape() + or specification.get('dtype') not in {'F64', 'F32', 'F16', 'BF16'} + or specification['dtype'] != header.get_dtype()): + raise ValueError(f'Qwen adapter tensor shape or dtype mismatch: {name}') + weights = {name: reader.get_tensor(name) for name in reader.keys()} + config = DEFAULTS | inference | {'model_root': adapter['base_model_name_or_path']} + return config, weights, adapter['source_step'] + marker = json.loads(path.with_suffix('.complete.json').read_text()) + if marker.get('status') != 'complete' or marker.get('bytes') != path.stat().st_size: + raise ValueError('checkpoint completion record is absent or mismatched') + # SupraTex adaptation: public inference weights use restricted loading. + # Legacy trainer checkpoints also contain NumPy uint32 RNG arrays; these + # exact reconstruction types retain compatibility without arbitrary pickle. + from numpy.core.multiarray import _reconstruct + with torch.serialization.safe_globals([_reconstruct, (_reconstruct, 'numpy.core.multiarray._reconstruct'), + (_reconstruct, 'numpy._core.multiarray._reconstruct'), + np.ndarray, np.dtype, type(np.dtype('uint32'))]): + state = torch.load(path, map_location='cpu', weights_only=True) + if state.get('format') != 'qwen_pbr_endpoint_v1' or marker.get('step') != state.get('step'): + raise ValueError('unsupported checkpoint or inconsistent step') + if not isinstance(state.get('config'), dict) or not isinstance(state.get('model'), dict): + raise ValueError('checkpoint needs complete config and trainable state') + config = DEFAULTS | state['config'] + for key in ('model_root', 'resolution', 'lora_rank', 'lora_alpha', 'noise_seed'): + if key not in state['config']: + raise ValueError(f'checkpoint config lacks {key}') + return config, state['model'], state['step'] + + +def strict_load_adapter(model, weights): + expected = {name: parameter for name, parameter in model.named_parameters() if parameter.requires_grad} + if expected.keys() != weights.keys(): + raise ValueError('checkpoint adapter key set differs from constructed model') + with torch.no_grad(): + for name, parameter in expected.items(): + tensor = weights[name] + if tensor.shape != parameter.shape or tensor.dtype != parameter.dtype or not torch.isfinite(tensor).all(): + raise ValueError(f'checkpoint tensor mismatch/nonfinite: {name}') + parameter.copy_(tensor) + model.requires_grad_(False).eval() + + +def canonical_conditions(values, resolution): + """Only retain consumed conditioning; material targets are never loaded.""" + values = dict(values) + layout = values['edit_image_pad_mask'].bool() + prompt = values.get('prompt_embeds') + if prompt is None: + text = values['prompt_text_embeds'] + if layout.ndim != 2 or layout.shape[0] != 1 or text.shape[1] != int((~layout[0]).sum()): + raise ValueError('compact conditioning layout mismatch') + prompt = text.new_zeros((1, layout.shape[1], text.shape[-1])) + prompt[:, ~layout[0]] = text + if prompt.ndim != 3 or prompt.shape[0] != 1 or prompt.shape[2] != 4096 or layout.shape != prompt.shape[:2]: + raise ValueError('invalid prompt/layout shape') + # Model replaces image-placeholder rows from RGB VAE. Canonical zeros allow + # exact native-v2 comparison without changing consumed values. + prompt = prompt.clone() + prompt[:, layout[0]] = 0 + mask = values.get('prompt_embeds_mask') + mask = torch.ones_like(layout) if mask is None else mask.bool() + rgb = values['rgb'] + if rgb.shape != (1, 64, resolution // 16, resolution // 16) or mask.shape != layout.shape: + raise ValueError('invalid RGB latent/mask shape') + result = dict(rgb=rgb, prompt_embeds=prompt, prompt_embeds_mask=mask, edit_image_pad_mask=layout) + if any(value.is_floating_point() and not torch.isfinite(value).all() for value in result.values()): + raise ValueError('nonfinite conditioning') + return result + + +def load_conditions(path, resolution, device): + values = {} + keys = {'rgb', 'prompt_embeds', 'prompt_text_embeds', 'prompt_embeds_mask', 'edit_image_pad_mask'} + with safe_open(str(path), framework='pt', device=str(device)) as reader: + for key in keys & set(reader.keys()): + values[key] = reader.get_tensor(key) + return canonical_conditions(values, resolution) + + +def input_binding(path): + path = Path(os.path.abspath(path)); stat = path.stat() + return dict(path=str(path), size=stat.st_size, mtime_ns=stat.st_mtime_ns) + + +def verify_cache_input(cache, image, resolution, model_references): + metadata = json.loads(Path(cache).with_suffix('.json').read_text()) + binding = metadata.get('input', metadata.get('sources', {}).get('rgb')) + if binding != input_binding(image) or metadata.get('resolution') != resolution: + raise ValueError('cache conditioning is not bound to this exact input/stat/resolution') + if metadata.get('model_root') not in model_references or metadata.get('prompt') != PROMPT: + raise ValueError('cached model reference or prompt differs from checkpoint inference contract') + return metadata['model_root'] + + +def native_conditions(square, model_root, resolution, device, low_vram=False): + from transformers import AutoProcessor + from diffsynth.models.qwen_image_21_vae import QwenImage21VAE + from diffsynth.models.qwen_image_21_text_encoder import QwenImage21TextEncoder + from diffsynth.pipelines.qwen_image_21 import QwenImage21Pipeline, QwenImage21Unit_EditImageEmbedder, QwenImage21Unit_PromptEmbedder + from diffsynth.utils.state_dict_converters.qwen_image_21_text_encoder import QwenImage21TextEncoderStateDictConverter + processor = AutoProcessor.from_pretrained(str(model_root / 'processor'), local_files_only=True) + processor.image_processor.size = {'shortest_edge': resolution ** 2, 'longest_edge': resolution ** 2} + pipe = QwenImage21Pipeline(device=device, torch_dtype=torch.bfloat16) + pipe.processor = processor + loading_device = 'cpu' if low_vram else device + vae = explicit_load(QwenImage21VAE, checkpoint_component(model_root, 'vae'), loading_device, torch.bfloat16) + pipe.vae = vae + with stage_on_device(vae, device, enabled=low_vram), torch.inference_mode(): + edited = QwenImage21Unit_EditImageEmbedder().process(pipe, edit_image=[square], height=resolution, width=resolution, + tiled=False, tile_size=256, tile_stride=192) + if edited['edit_image'][0].size != square.size: + raise ValueError('native edit path changed dimensions') + rgb = edited['edit_latents'][0] + encoder = explicit_load(QwenImage21TextEncoder, checkpoint_component(model_root, 'text_encoder'), loading_device, + torch.bfloat16, QwenImage21TextEncoderStateDictConverter) + pipe.text_encoder = encoder + blocks = text_encoder_offload_blocks(encoder) if low_vram else () + with sequential_cpu_offload(encoder, blocks, device, enabled=low_vram), torch.inference_mode(): + values = QwenImage21Unit_PromptEmbedder().process(pipe, prompt=PROMPT, edit_image=[square]) + conditions = canonical_conditions(dict(values, rgb=rgb), resolution) + pipe.text_encoder = None + del encoder, pipe, values, edited, blocks + gc.collect(); torch.cuda.empty_cache() + return conditions, vae + + +def predict(model, conditions, cfg): + # Shape comes from observed RGB, never a GT or fabricated target tensor. + shape = (conditions['rgb'].shape[0], 2, *conditions['rgb'].shape[1:]) + shape_carrier = conditions['rgb'].new_empty(shape) + noise = endpoint_noise(shape_carrier, cfg['noise_seed']) + mask = conditions['prompt_embeds_mask'] + if cfg['omit_all_valid_mask'] and bool(mask.all()): + mask = None + with torch.inference_mode(), torch.autocast('cuda', dtype=torch.bfloat16, enabled=cfg['bf16_autocast']): + velocity = model(latents=noise, timestep=torch.full((noise.shape[0],), 1000., device=noise.device), + prompt_embeds=conditions['prompt_embeds'], prompt_embeds_mask=mask, + edit_image_pad_mask=conditions['edit_image_pad_mask'], edit_latents=[conditions['rgb']], + use_gradient_checkpointing=False, use_flex_attention=False, kv_cache=None) + prediction = noise.float() - velocity.float() + if not torch.isfinite(prediction).all(): + raise FloatingPointError('nonfinite material latent prediction') + return prediction + + +def predict_and_decode(model, vae, conditions, cfg, low_vram=False): + device = conditions['rgb'].device + try: + with sequential_cpu_offload(model, model.dit.transformer_blocks, device, enabled=low_vram): + latent = predict(model, conditions, cfg) + finally: + if low_vram: + # Qwen stores RoPE tables as an ordinary list, outside Module buffers. + model.dit.pos_embed.freqs = [value.cpu() for value in model.dit.pos_embed.freqs] + release_device_cache(device) + with stage_on_device(vae, device, enabled=low_vram), torch.inference_mode(): + decoded = decode_pair(vae, latent) + return latent, decoded + + +def output_maps(decoded, square, original, transform, restore_size): + if decoded.shape[0:3] != (1, 2, 4) or not torch.isfinite(decoded).all(): + raise ValueError('invalid/nonfinite decoded material pair') + arrays = np.clip(np.rint(decoded[0].detach().float().cpu().permute(0, 2, 3, 1).numpy() * 255), 0, 255).astype(np.uint8) + diagnostics = [Image.fromarray(array[:, :, 3], 'L') for array in arrays] + maps = [] + for modality, array in enumerate(arrays): + color = Image.fromarray(array[:, :, :3], 'RGB') + coverage = square.getchannel('A') + if restore_size: + x, y = transform['offset']; w, h = transform['content_size'] + color = color.crop((x, y, x + w, y + h)) + if color.size != original.size: + color = color.resize(original.size, Image.Resampling.LANCZOS) + coverage = original.getchannel('A') # exact original coverage, no second interpolation + out = np.array(color.convert('RGBA')) + out[:, :, 3] = np.array(coverage) + if modality == 1: + out[:, :, 2] = 255 # placeholder, not material opacity or AO prediction + maps.append(Image.fromarray(out, 'RGBA')) + return maps, diagnostics + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument('--checkpoint', required=True, type=Path, + help='Native adapter_model.safetensors with sibling adapter_config.json, or legacy PT with completion marker') + parser.add_argument('--input', required=True, type=Path) + parser.add_argument('--output-dir', required=True, type=Path) + parser.add_argument('--model-root', type=Path, help='Explicit loading location only; checkpoint model reference retained') + parser.add_argument('--conditions-cache', type=Path, help='Read RGB conditioning only; sibling JSON must bind input/model/resolution') + parser.add_argument('--compare-cache', type=Path, help='Native path plus exact conditioning/latent/decoded comparison against cache') + parser.add_argument('--comparison-atol', type=float, default=0.) + parser.add_argument('--save-conditions', action='store_true') + parser.add_argument('--restore-input-size', action='store_true') + parser.add_argument('--save-predicted-coverage', action='store_true') + parser.add_argument('--device', default='cuda:0') + parser.add_argument('--low-vram', action='store_true', default=os.environ.get('SUPRATEX_LOW_VRAM') == '1', + help='CPU stage/block offload; preserves weights and sampling, uses more CPU RAM and transfers (also SUPRATEX_LOW_VRAM=1)') + args = parser.parse_args() + if args.conditions_cache and args.compare_cache: + parser.error('--conditions-cache and --compare-cache are mutually exclusive') + if args.comparison_atol < 0: + parser.error('--comparison-atol must be nonnegative') + cfg, weights, step = load_checkpoint(args.checkpoint) + os.environ['DIFFSYNTH_ATTENTION_IMPLEMENTATION'] = cfg['attention_backend'] + os.environ['HF_HUB_OFFLINE'] = '1'; os.environ['TRANSFORMERS_OFFLINE'] = '1' + if cfg['deterministic_algorithms']: + os.environ.setdefault('CUBLAS_WORKSPACE_CONFIG', ':4096:8') + torch.use_deterministic_algorithms(cfg['deterministic_algorithms']) + device = torch.device(args.device) + torch.cuda.set_device(device) + reference_root = cfg['model_root'] + model_root = Path(os.path.abspath(args.model_root or reference_root)) + # An explicit load-location override is an operator-supplied alias, never an + # implicit identity rewrite or a digest-based equivalence claim. + model_references = {reference_root, str(model_root)} + original, square, transform = prepare_image(args.input, cfg['resolution']) + if args.output_dir.exists(): + raise FileExistsError('choose a fresh inference output directory') + args.output_dir.mkdir(parents=True) + started = time.perf_counter() + torch.cuda.reset_peak_memory_stats(device) + if args.conditions_cache: + verify_cache_input(args.conditions_cache, args.input, cfg['resolution'], model_references) + conditions = load_conditions(args.conditions_cache, cfg['resolution'], device) + from diffsynth.models.qwen_image_21_vae import QwenImage21VAE + vae = explicit_load(QwenImage21VAE, checkpoint_component(model_root, 'vae'), + 'cpu' if args.low_vram else device, torch.bfloat16) + else: + conditions, vae = native_conditions(square, model_root, cfg['resolution'], device, low_vram=args.low_vram) + load_cfg = cfg | {'model_root': str(model_root)} + model = build_model(load_cfg, 'cpu' if args.low_vram else device) + strict_load_adapter(model, weights) + del weights + latent, decoded = predict_and_decode(model, vae, conditions, cfg, low_vram=args.low_vram) + comparison = None + if args.compare_cache: + verify_cache_input(args.compare_cache, args.input, cfg['resolution'], model_references) + cached = load_conditions(args.compare_cache, cfg['resolution'], device) + differences = {key: float((conditions[key].float() - cached[key].float()).abs().max()) for key in conditions} + other_latent, other_decoded = predict_and_decode(model, vae, cached, cfg, low_vram=args.low_vram) + differences.update(latent=float((latent - other_latent).abs().max()), decoded=float((decoded - other_decoded).abs().max())) + comparison = dict(max_abs=differences, atol=args.comparison_atol, passed=all(v <= args.comparison_atol for v in differences.values())) + if not comparison['passed']: + # Discriminate native/cache disagreement from repeated-kernel noise + # while holding conditions, checkpoint and process constant. + repeat_latent, repeat_decoded = predict_and_decode(model, vae, cached, cfg, low_vram=args.low_vram) + comparison['cached_repeat_control_max_abs'] = dict( + latent=float((other_latent-repeat_latent).abs().max()), + decoded=float((other_decoded-repeat_decoded).abs().max())) + (args.output_dir / 'comparison.json').write_text(json.dumps(comparison, indent=2) + '\n') + if not comparison['passed']: + raise ValueError('native/cache comparison exceeded requested tolerance') + maps, diagnostics = output_maps(decoded, square, original, transform, args.restore_input_size) + for label, image, diag in zip(('albedo', 'mra'), maps, diagnostics): + image.save(args.output_dir / f'{label}.png') + if args.save_predicted_coverage: + diag.save(args.output_dir / f'{label}_predicted_coverage.png') + if args.save_conditions: + save_file({k: v.detach().cpu().contiguous() for k, v in conditions.items()}, str(args.output_dir / 'conditions.safetensors')) + (args.output_dir / 'conditions.json').write_text(json.dumps(dict(input=input_binding(args.input), resolution=cfg['resolution'], + model_root=reference_root, prompt=PROMPT, transform=transform), indent=2) + '\n') + report = dict(status='complete', checkpoint=os.path.abspath(args.checkpoint), checkpoint_step=step, + input=input_binding(args.input), loading_model_root=str(model_root), checkpoint_model_reference=reference_root, + model_override_is_operator_reference_not_hash_verified=bool(args.model_root), config=cfg, transform=transform, + accepted_model_reference_aliases=sorted(model_references), + low_vram=args.low_vram, + memory_policy='cpu_stages_and_sequential_transformer_blocks' if args.low_vram else 'resident_models', + restore_input_size=args.restore_input_size, restoration='crop padding; LANCZOS color interpolation if size changed; exact original coverage' if args.restore_input_size else 'training square canvas', + output_semantics='Albedo sRGB; MRA R=linear Metallic G=linear Roughness B=1 placeholder; A=input coverage', + boundary='Opaque material estimator. Input alpha retained as coverage; no segmentation, material transparency or hidden-layer recovery inferred.', + path='cached_conditioning_only' if args.conditions_cache else 'native_rgb_vae_and_per_image_vl', comparison=comparison, + peak_gpu_allocated_bytes=torch.cuda.max_memory_allocated(device), + peak_gpu_reserved_bytes=torch.cuda.max_memory_reserved(device), + elapsed_seconds=time.perf_counter()-started) + (args.output_dir / 'REPORT.json').write_text(json.dumps(report, indent=2) + '\n') + print(json.dumps(dict(status='complete', output=str(args.output_dir), checkpoint_step=step))) + + +if __name__ == '__main__': + main() diff --git a/deployment/runtime/qwen/examples/qwen_image_21/pbr/inference_utils.py b/deployment/runtime/qwen/examples/qwen_image_21/pbr/inference_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..124771bddb9a681513aa230c6cf1ef5b1def24f6 --- /dev/null +++ b/deployment/runtime/qwen/examples/qwen_image_21/pbr/inference_utils.py @@ -0,0 +1,135 @@ +"""Model construction, local checkpoint loading, and decoding for PBR inference.""" +from __future__ import annotations + +import inspect +import json +from pathlib import Path +import time + +import torch +from safetensors import safe_open +from safetensors.torch import load_file + + +# Only settings consumed by inference; legacy checkpoint fields are not defaults. +DEFAULTS = dict( + lora_rank=32, lora_alpha=32, noise_seed=42, resolution=512, + omit_all_valid_mask=False, bf16_autocast=False, + attention_backend='torch', deterministic_algorithms=False, + lora_output_projection=False, +) + + +PROMPT = ("Generate aligned albedo and metallic/roughness material maps for the object in the input image. " + "Pack MRA with red=metallic, green=roughness, and blue=constant one placeholder. " + "No observed ambient occlusion is available.") + + +def log(message): + print(f"[{time.strftime('%Y-%m-%d %H:%M:%S')}] {message}", flush=True) + + +def read_json(path): + if not path.is_file(): + raise FileNotFoundError(f"Required local file is absent: {path}") + return json.loads(path.read_text(encoding="utf-8")) + + +def checkpoint_component(root, component): + """Resolve an explicit component and verify every indexed shard before GPU work.""" + directory = root / component + config = read_json(directory / "config.json") + if not isinstance(config, dict): + raise ValueError(f"Expected an object config: {directory / 'config.json'}") + indexes = sorted(directory.glob("*.safetensors.index.json")) + if len(indexes) > 1: + raise ValueError(f"Ambiguous safetensors indexes under {directory}") + if indexes: + index = read_json(indexes[0]) + weight_map = index.get("weight_map") + if not isinstance(weight_map, dict) or not weight_map: + raise ValueError(f"Invalid weight_map: {indexes[0]}") + names = sorted(set(weight_map.values())) + if any(Path(name).is_absolute() or ".." in Path(name).parts for name in names): + raise ValueError(f"Indexed shard path leaves component directory: {indexes[0]}") + files = [directory / name for name in names] + else: + pattern = "model*.safetensors" if component == "text_encoder" else "diffusion_pytorch_model*.safetensors" + files = sorted(directory.glob(pattern)) + if not files: + raise FileNotFoundError(f"No local safetensors checkpoint found under {directory}") + for filename in files: + if not filename.is_file() or filename.stat().st_size == 0: + raise FileNotFoundError(f"Missing or empty checkpoint shard: {filename}") + keys = set() + for filename in files: + with safe_open(str(filename), framework="pt", device="cpu") as handle: + shard_keys = set(handle.keys()) + duplicates = keys & shard_keys + if duplicates: + raise ValueError(f"Duplicate state keys between shards: {sorted(duplicates)[:5]}") + keys.update(shard_keys) + if indexes and keys != set(weight_map): + raise ValueError(f"Index and shard state keys disagree under {directory}") + return {"directory": directory, "config": config, "files": files, "tensor_count": len(keys)} + + +def explicit_load(model_class, component, device, dtype, converter=None): + """Instantiate the named class and strictly assign the local tensor state.""" + from diffsynth.core.vram.initialization import skip_model_initialization + + constructor = inspect.signature(model_class.__init__).parameters + config = {key: value for key, value in component["config"].items() if key in constructor and key != "self"} + log(f"Loading {model_class.__name__} explicitly ({len(component['files'])} shards, device={device}, dtype={dtype})") + started = time.perf_counter() + # Set the default dtype to preserve generated float32 text-encoder rotary buffers; + # checkpoint parameters are assigned below and then cast explicitly. + with skip_model_initialization(): + model = model_class(**config) + state = {} + for filename in component["files"]: + state.update(load_file(str(filename), device="cpu")) + if converter is not None: + state = converter(state) + model.load_state_dict(state, strict=True, assign=True) + del state + model = model.to(device=device, dtype=dtype).eval() + if any(parameter.is_meta for parameter in model.parameters()): + raise RuntimeError(f"{model_class.__name__}: parameters remain on meta after strict checkpoint load") + model.requires_grad_(False) + log(f"Loaded {model_class.__name__} in {time.perf_counter() - started:.2f}s") + return model + + +def endpoint_noise(target, seed): + # Draw the fixed endpoint field on CPU before moving it to the model device. + noise = torch.randn(target.shape, generator=torch.Generator(device='cpu').manual_seed(seed), dtype=torch.float32) + return noise.to(device=target.device, dtype=target.dtype) + + +def build_model(cfg, device): + from peft import LoraConfig, inject_adapter_in_model + from diffsynth.models.qwen_image_21_dit import QwenImage21DiT + from diffsynth.pipelines.qwen_image_21_pbr import QwenImage21PBRModel + dit = explicit_load(QwenImage21DiT, checkpoint_component(Path(cfg['model_root']), 'transformer'), device, torch.bfloat16) + if len(dit.transformer_blocks) != 32: + raise ValueError('Expected complete pretrained 32-layer backbone') + lora = LoraConfig(r=cfg['lora_rank'], lora_alpha=cfg['lora_alpha'], lora_dropout=0., bias='none', + target_modules=['attn.to_q', 'attn.to_k', 'attn.to_v', 'attn.to_out.0'] + + (['proj_out'] if cfg['lora_output_projection'] else [])) + dit = inject_adapter_in_model(lora, dit) + model = QwenImage21PBRModel(dit).to(device) + for parameter in model.parameters(): + if parameter.requires_grad: + parameter.data = parameter.data.float() + return model + + +def decode_pair(vae, latents, clamp_output=True): + """Decode the albedo and MRA latent maps in their fixed modality order.""" + decoded = [] + for index in range(2): + z = latents[:, index].to(dtype=torch.bfloat16) + value = vae.decode(z, tiled=False, clamp_output=clamp_output) + decoded.append((value.float() + 1.) / 2.) + return torch.stack(decoded, dim=1) diff --git a/deployment/runtime/qwen/requirements.txt b/deployment/runtime/qwen/requirements.txt index 2f92050dbb9e556b22a49257a8946f3edd3deea0..d00d28fd4e1be7b58965fe4904d77eb54624c9d8 100644 --- a/deployment/runtime/qwen/requirements.txt +++ b/deployment/runtime/qwen/requirements.txt @@ -1,7 +1,6 @@ -transformers==4.57.3 -diffusers==0.37.1 -accelerate==1.13.0 -peft==0.19.1 -ftfy==6.3.1 -pandas==2.3.3 -modelscope==1.34.0 +transformers==4.57.3 +diffusers==0.37.1 +accelerate==1.13.0 +peft==0.19.1 +ftfy==6.3.1 +modelscope==1.34.0 diff --git a/deployment/runtime/requirements-python.txt b/deployment/runtime/requirements-python.txt index b600c2fef26c6500d82b23d7f4f9934c56573d2d..25674b7e36b0e5c4da8e7477c6e841c2c2d6c6c1 100644 --- a/deployment/runtime/requirements-python.txt +++ b/deployment/runtime/requirements-python.txt @@ -1,31 +1,29 @@ -# Recorded from the working Python 3.10 / PyTorch 2.6.0 CUDA 12.4 environment. -# Install torch/torchvision and compiled extensions separately; see DEPENDENCIES.md. -numpy==2.2.6 -Pillow==12.0.0 -scipy==1.15.3 -easydict==1.13 -tqdm==4.67.1 -einops==0.8.2 -safetensors==0.7.0 -huggingface-hub==0.36.2 -trimesh==4.10.1 -opencv-python-headless==4.13.0.92 -imageio==2.37.2 -imageio-ffmpeg==0.6.0 -numba==0.65.1 -transformers==4.57.3 -diffusers==0.37.1 -accelerate==1.13.0 -requests==2.34.2 -timm==1.0.22 -kornia==0.8.2 -plyfile==1.1.3 -zstandard==0.25.0 -ninja==1.13.0 -pandas==2.3.3 -peft==0.19.1 -ftfy==6.3.1 -lpips==0.1.4 -tensorboard==2.20.0 -av==17.1.0 -modelscope==1.34.0 +# Recorded from the working Python 3.10 / PyTorch 2.6.0 CUDA 12.4 environment. +# Install torch/torchvision and compiled extensions separately; see DEPENDENCIES.md. +numpy==2.2.6 +Pillow==12.0.0 +scipy==1.15.3 +easydict==1.13 +tqdm==4.67.1 +einops==0.8.2 +safetensors==0.7.0 +huggingface-hub==0.36.2 +trimesh==4.10.1 +opencv-python-headless==4.13.0.92 +imageio==2.37.2 +imageio-ffmpeg==0.6.0 +numba==0.65.1 +transformers==4.57.3 +diffusers==0.37.1 +accelerate==1.13.0 +requests==2.34.2 +timm==1.0.22 +kornia==0.8.2 +plyfile==1.1.3 +zstandard==0.25.0 +ninja==1.13.0 +pandas==2.3.3 +peft==0.19.1 +ftfy==6.3.1 +av==17.1.0 +modelscope==1.34.0 diff --git a/hf_space/ovoxel_results.py b/hf_space/ovoxel_results.py index 00d94f87b2c53c571533cd41efd5732505946956..7209ac2daf3bd0da61e63f163a501ec352a7bb05 100644 --- a/hf_space/ovoxel_results.py +++ b/hf_space/ovoxel_results.py @@ -1,117 +1,117 @@ -"""Persist final o-voxels and their rendered views in the mounted data Bucket.""" -import json -from pathlib import Path -import re -import shutil -from html import escape -from urllib.parse import quote - -PREVIEW_MODES = ( - ('normal', 'Normal', 'normal.png'), - ('clay', 'Clay', 'clay.png'), - ('base_color', 'Base color', 'basecolor.png'), - ('shaded_studio', 'HDRI Studio', 'hdri_studio.png'), - ('shaded_forest', 'HDRI Forest', 'hdri_forest.png'), - ('shaded_sunset', 'HDRI Sunset', 'hdri_sunset.png'), - ('shaded_courtyard', 'HDRI Courtyard', 'hdri_courtyard.png'), - ('shaded', 'Studio', 'hdri_studio.png'), # Previously saved previews. -) - -ROOT = Path('/data/generated/ovoxel-results-v1') - - -def directory(identifier): - if not isinstance(identifier, str) or re.fullmatch(r'[0-9a-f]{32}', identifier) is None: - raise ValueError('Invalid generated asset identifier') - return ROOT/identifier - - -def persist(source, report): - source = Path(source) - identifier = source.parent.name - destination = directory(identifier) - destination.mkdir(parents=True, exist_ok=True) - for name in ('asset.npz','input.png','report.json','preview.json','timings.jsonl'): - shutil.copyfile(source.parent/name, destination/name) - shutil.copytree(source.parent/'views', destination/'views', dirs_exist_ok=True) - (destination/'ready.json').write_text(json.dumps({'complete':True,'identifier':identifier, - 'project_revision':report['project_revision'],'output_kind':'final_ovoxel'})) - return identifier - - -def preview_html(identifier): - root = directory(identifier) - metadata = json.loads((root/'preview.json').read_text()) - default = metadata.get('default_mode', 'shaded' if 'shaded' in metadata['modes'] else 'shaded_forest') - if default not in metadata['modes']: - default = metadata['modes'][0] - images = [] - for mode in metadata['modes']: - for view in range(metadata['views']): - url = '/gradio_api/file='+quote(str(root/'views'/f'{mode}-{view:02d}.webp'),safe='/') - hidden = '' if mode == default and view == 0 else ' hidden' - images.append(f'{mode.replace(') - buttons = [] - for mode, label, icon in PREVIEW_MODES: - if mode not in metadata['modes']: - continue - # Inline the tiny swatches; preview frames remain individually cached files. - import base64 - icon_path = Path(__file__).resolve().parents[1]/'assets/preview-lighting'/icon - swatch = base64.b64encode(icon_path.read_bytes()).decode('ascii') - active = mode == default - buttons.append(f'') - return (f'
' - '
'+''.join(images)+'
' - '
'+''.join(buttons)+'
' - f'
') - - -def extract(identifier, extraction_token=None): - from hf_space import extraction_progress as progress - from hf_space import generation_progress - from hf_space.timing import observe_stages - token = extraction_token if generation_progress.valid(extraction_token) else progress.start() - try: - progress.stage(token, 'load') - with observe_stages(lambda name: progress.stage(token, name)): - output = _extract(identifier, token) - progress.finish(token) - return output - except Exception: - progress.finish(token, failed=True) - raise - - -def _extract(identifier, token): - """Runs only inside the explicitly requested GPU extraction callback.""" - root = directory(identifier) - if not (root/'ready.json').is_file(): - raise ValueError('The generated o-voxel has not finished saving') - output = root/'asset.glb' - if output.is_file(): - return str(output) - import numpy as np - import torch - from supratex.representations import MeshWithVoxel - from supratex.inference.pipeline import _export, LAYOUT, _runtime - if not torch.cuda.is_available(): - raise RuntimeError('GLB extraction requires CUDA') - with np.load(root/'asset.npz',allow_pickle=False) as state: - resolution = int(state['resolution']) - representation = MeshWithVoxel(torch.from_numpy(state['vertices']).cuda(), - torch.from_numpy(state['faces']).cuda(),state['origin'].tolist(),float(state['voxel_size']), - torch.from_numpy(state['coords']).cuda(),torch.from_numpy(state['attrs']).cuda(), - torch.Size([1,6,resolution,resolution,resolution]),layout=LAYOUT) - import time - started = time.monotonic() - _export(output, representation,4096,500000,'png',_runtime()) - from hf_space import extraction_progress - extraction_progress.stage(token, 'validate') - from hf_space.backend import validate_glb - metadata = validate_glb(output) - (root/'extraction.json').write_text(json.dumps({'status':'complete','device':'cuda', - 'elapsed_seconds':time.monotonic()-started,**metadata},indent=2)) - return str(output) +"""Persist final o-voxels and their rendered views in the mounted data Bucket.""" +import json +from pathlib import Path +import re +import shutil +from html import escape +from urllib.parse import quote + +PREVIEW_MODES = ( + ('normal', 'Normal', 'normal.png'), + ('clay', 'Clay', 'clay.png'), + ('base_color', 'Base color', 'basecolor.png'), + ('shaded_studio', 'HDRI Studio', 'hdri_studio.png'), + ('shaded_forest', 'HDRI Forest', 'hdri_forest.png'), + ('shaded_sunset', 'HDRI Sunset', 'hdri_sunset.png'), + ('shaded_courtyard', 'HDRI Courtyard', 'hdri_courtyard.png'), + ('shaded', 'Studio', 'hdri_studio.png'), # Previously saved previews. +) + +ROOT = Path('/data/generated/ovoxel-results-v1') + + +def directory(identifier): + if not isinstance(identifier, str) or re.fullmatch(r'[0-9a-f]{32}', identifier) is None: + raise ValueError('Invalid generated asset identifier') + return ROOT/identifier + + +def persist(source, report): + source = Path(source) + identifier = source.parent.name + destination = directory(identifier) + destination.mkdir(parents=True, exist_ok=True) + for name in ('asset.npz','input.png','report.json','preview.json','timings.jsonl'): + shutil.copyfile(source.parent/name, destination/name) + shutil.copytree(source.parent/'views', destination/'views', dirs_exist_ok=True) + (destination/'ready.json').write_text(json.dumps({'complete':True,'identifier':identifier, + 'project_revision':report['project_revision'],'output_kind':'final_ovoxel'})) + return identifier + + +def preview_html(identifier): + root = directory(identifier) + metadata = json.loads((root/'preview.json').read_text()) + default = metadata.get('default_mode', 'shaded' if 'shaded' in metadata['modes'] else 'shaded_forest') + if default not in metadata['modes']: + default = metadata['modes'][0] + images = [] + for mode in metadata['modes']: + for view in range(metadata['views']): + url = '/gradio_api/file='+quote(str(root/'views'/f'{mode}-{view:02d}.webp'),safe='/') + hidden = '' if mode == default and view == 0 else ' hidden' + images.append(f'{mode.replace(') + buttons = [] + for mode, label, icon in PREVIEW_MODES: + if mode not in metadata['modes']: + continue + # Inline the tiny swatches; preview frames remain individually cached files. + import base64 + icon_path = Path(__file__).resolve().parents[1]/'assets/preview-lighting'/icon + swatch = base64.b64encode(icon_path.read_bytes()).decode('ascii') + active = mode == default + buttons.append(f'') + return (f'
' + '
'+''.join(images)+'
' + '
'+''.join(buttons)+'
' + f'
') + + +def extract(identifier, extraction_token=None): + from hf_space import extraction_progress as progress + from hf_space import generation_progress + from hf_space.timing import observe_stages + token = extraction_token if generation_progress.valid(extraction_token) else progress.start() + try: + progress.stage(token, 'load') + with observe_stages(lambda name: progress.stage(token, name)): + output = _extract(identifier, token) + progress.finish(token) + return output + except Exception: + progress.finish(token, failed=True) + raise + + +def _extract(identifier, token): + """Runs only inside the explicitly requested GPU extraction callback.""" + root = directory(identifier) + if not (root/'ready.json').is_file(): + raise ValueError('The generated o-voxel has not finished saving') + output = root/'asset.glb' + if output.is_file(): + return str(output) + import numpy as np + import torch + from supratex.representations import MeshWithVoxel + from supratex.inference.pipeline import _export, LAYOUT, _runtime + if not torch.cuda.is_available(): + raise RuntimeError('GLB extraction requires CUDA') + with np.load(root/'asset.npz',allow_pickle=False) as state: + resolution = int(state['resolution']) + representation = MeshWithVoxel(torch.from_numpy(state['vertices']).cuda(), + torch.from_numpy(state['faces']).cuda(),state['origin'].tolist(),float(state['voxel_size']), + torch.from_numpy(state['coords']).cuda(),torch.from_numpy(state['attrs']).cuda(), + torch.Size([1,6,resolution,resolution,resolution]),layout=LAYOUT) + import time + started = time.monotonic() + _export(output, representation,4096,500000,'png',_runtime(), low_vram=True) + from hf_space import extraction_progress + extraction_progress.stage(token, 'validate') + from hf_space.backend import validate_glb + metadata = validate_glb(output) + (root/'extraction.json').write_text(json.dumps({'status':'complete','device':'cuda', + 'elapsed_seconds':time.monotonic()-started,**metadata},indent=2)) + return str(output) diff --git a/hf_space/release_backend.py b/hf_space/release_backend.py new file mode 100644 index 0000000000000000000000000000000000000000..ba5869004aca578290771392b6753ccdceb644d6 --- /dev/null +++ b/hf_space/release_backend.py @@ -0,0 +1,176 @@ +"""Bind the pinned official single-view CLI to the existing Gradio callbacks.""" +from pathlib import Path +import json +import os +import signal +import subprocess +import sys +import threading +import time +import traceback +import uuid + +from hf_space import native_build, generation_progress + +ROOT=Path(__file__).resolve().parents[1] +MODEL_REVISION='1bc02eb969a8ed0f1418c2bbe1e409a8481f834c' +CODE_REVISION='e4dfbbde6a8d76e25deef5aba5d836f03994967d' +CHECKPOINT='denoiser_step0030000.safetensors' +MODEL_ROOT=Path('/data/models') +STATE=Path('/tmp/supratex-release-preparation.json') +READY=Path('/tmp/supratex-latest-ready') +_lock=threading.Lock() +_worker=None + +def model_status(): + return json.loads(STATE.read_text()) if STATE.exists() else {'stage':'pending'} + +def update(**changes): + value=model_status()|changes|{'updated_unix':time.time()} + temporary=STATE.with_suffix('.tmp') + temporary.write_text(json.dumps(value,indent=2)) + temporary.replace(STATE) + print('[release-prepare] '+json.dumps(changes),flush=True) + +def prepare_models(): + try: + update(stage='waiting_native_build',startup_progress={'label':'Preparing runtime','completed':0,'total':0}) + while native_build.status().get('stage') not in ('compiled','failed'): + time.sleep(2) + if native_build.status()['stage']!='compiled': + raise RuntimeError('Native runtime preparation failed') + repository=os.environ.get('SUPRATEX_MODEL_REPO','supra3d/supra-tex') + revision=os.environ.get('SUPRATEX_MODEL_REVISION',MODEL_REVISION) + if (repository,revision)!=('supra3d/supra-tex',MODEL_REVISION): + raise RuntimeError('This source release requires the validated, pinned model revision') + from scripts import prepare_models as prep + from hf_space.download_models import download + plan=prep.build_plan(prep.parse_args(['--repo',repository,'--revision',revision, + '--with-rmbg','--output-dir',str(MODEL_ROOT)]),resolve_remote=True) + operations=sorted(plan['downloads'],key=lambda item:not item.get('gated',False)) + for index,operation in enumerate(operations): + if operation['name']=='project': + operation=dict(operation,destination='project/'+revision) + update(stage='downloading_foundations',project_revision=revision, + startup_progress={'label':'Checking model files: '+operation['name'], + 'completed':index,'total':len(operations),'unit':'groups'}) + download(operation,MODEL_ROOT) + from hf_space.weight_cache import mirror_models,mirror_project + update(stage='staging_local_weights',startup_progress={ + 'label':'Preparing local weight files','completed':0,'total':0}) + project_root=MODEL_ROOT/'project'/revision + sources=[*(MODEL_ROOT/'external').rglob('*.safetensors'),*project_root.rglob('*.safetensors')] + total=sum(path.stat().st_size for path in sources) + copied={} + def file_progress(source,completed,size): + copied[str(source)]=completed + update(startup_progress={'label':'Preparing local weight files','completed':sum(copied.values()), + 'total':total,'unit':'bytes'}) + foundations=mirror_models(MODEL_ROOT,progress=file_progress) + project=mirror_project(project_root,revision,progress=file_progress) + arguments=['--repo',str(project),'--output-dir',str(READY)] + for flag,name in [('qwen-base','qwen-base'),('geometry','geometry'),('dino','dino'), + ('vae','vae'),('moge','moge2'),('rmbg','rmbg2')]: + arguments+=['--'+flag,str(foundations/name)] + arguments+=['--naf-weights',str(foundations/'naf/naf_release.pth')] + config=prep.prepare(prep.build_plan(prep.parse_args(arguments))) + prepared=json.loads(config.read_text()) + if not prepared['checkpoint'].endswith(CHECKPOINT): + raise RuntimeError('The prepared pipeline did not select step-30000 completion weights') + update(stage='ready_released',config=str(config),project_repository=repository, + project_revision=revision,code_revision=CODE_REVISION,weight_origin='released', + actual_training_steps=30000,random_weights_present=False, + checkpoint=CHECKPOINT,model_loading='official_stage_loading', + startup_progress={'label':'Model files ready','completed':1,'total':1,'unit':'groups'}) + except BaseException: + update(stage='failed',error=traceback.format_exc()) + traceback.print_exc() + +def start_preparation(): + global _worker + with _lock: + if model_status().get('stage')=='ready_released':return None + if _worker is None or not _worker.is_alive(): + _worker=threading.Thread(target=prepare_models,daemon=True,name='latest-release-prepare') + _worker.start() + return _worker + +def ensure_prepared(generation_token=None): + worker=start_preparation() + while worker is not None and worker.is_alive(): + state=model_status() + generation_progress.update(generation_token,status='preparing', + label=state.get('startup_progress',{}).get('label','Preparing model files')) + worker.join(timeout=1) + state=model_status() + if state.get('stage')!='ready_released': + raise RuntimeError('Latest-release preparation failed: '+state.get('error','unknown failure')[-2000:]) + +def generate(image,seed,progress=None,timeout_seconds=1800,generation_token=None): + state=model_status() + if state.get('stage')!='ready_released':raise RuntimeError('Model files are not ready') + identifier=uuid.uuid4().hex + job=Path('/tmp/supratex-jobs')/identifier + job.mkdir(parents=True) + image.save(job/'input.png') + output=job/'asset.npz' + command=[sys.executable,str(ROOT/'scripts/run_example.py'),'--config',state['config'], + '--input',str(job/'input.png'),'--output',str(output),'--attention','sdpa','--low-vram', + '--seed',str(int(seed)),'--preview-only'] + environment=os.environ.copy() + environment.update(OMP_NUM_THREADS='4', + SUPRATEX_RUNTIME_PROFILE='torch28',PYTHONUNBUFFERED='1', + SUPRATEX_TIMING_PATH=str(job/'timings.jsonl'),SUPRATEX_TIMING_ORIGIN_NS=str(time.monotonic_ns())) + started=time.monotonic() + generation_progress.update(generation_token,status='running',job=identifier, + generation_started_unix=time.time(),label='Running the latest single-view pipeline') + if progress:progress(.05,'Running the latest single-view pipeline') + process=subprocess.Popen(command,cwd=ROOT,env=environment,stdout=subprocess.PIPE, + stderr=subprocess.STDOUT,text=True,start_new_session=True) + def stop_processes(): + if process.poll() is None: + try:os.killpg(process.pid,signal.SIGKILL) + except ProcessLookupError:pass + timer=threading.Timer(timeout_seconds,stop_processes) + timer.start() + try: + with (job/'inference.log').open('w') as log: + for line in process.stdout: + log.write(line);log.flush() + print(line,end='',flush=True) + if line.startswith('[1/2]'): + generation_progress.update(generation_token,label='Generating geometry and reference materials') + elif line.startswith('[2/2]'): + generation_progress.update(generation_token,label='Completing surface materials and rendering o-voxel previews') + code=process.wait() + finally: + timer.cancel() + if process.poll() is None: + stop_processes();process.wait() + receipt_path=job/'asset_work/run.json' + receipt=json.loads(receipt_path.read_text()) if receipt_path.exists() else {} + if code or receipt.get('status')!='complete' or not output.is_file(): + raise RuntimeError('Latest single-view generation failed: '+(job/'inference.log').read_text(errors='replace')[-5000:]) + import numpy as np + with np.load(output,allow_pickle=False) as asset: + required={'vertices','faces','coords','attrs','resolution','origin','voxel_size'} + if not required<=set(asset.files) or not len(asset['coords']) or not len(asset['faces']): + raise RuntimeError('Generated o-voxel geometry is incomplete') + if asset['attrs'].shape[1]!=6 or not np.isfinite(asset['attrs']).all(): + raise RuntimeError('Generated o-voxel materials are invalid') + voxel_count=len(asset['coords']) + preview=json.loads((job/'preview.json').read_text()) + if len(list((job/'views').glob('*.webp')))!=preview['views']*len(preview['modes']): + raise RuntimeError('The o-voxel preview render set is incomplete') + if list(job.rglob('*.glb')) or preview.get('automatic_glb_export') is not False: + raise RuntimeError('Generation must not export a mesh before selection') + report={'status':'complete','asset_id':identifier,'code_revision':CODE_REVISION, + 'project_revision':MODEL_REVISION,'checkpoint':CHECKPOINT,'weight_origin':'released', + 'elapsed_seconds':time.monotonic()-started,'output_kind':'final_ovoxel','backview':False, + 'refine':False,'model_loading':'official_stage_loading','run':receipt, + 'preview':preview,'voxel_count':voxel_count,'automatic_glb_export':False} + (job/'report.json').write_text(json.dumps(report,indent=2)) + if not (job/'timings.jsonl').exists(): + (job/'timings.jsonl').write_text(json.dumps({'name':'preview_generation','elapsed_seconds':report['elapsed_seconds']})+'\n') + if progress:progress(1.,'Complete') + return str(output),str(output),report diff --git a/hf_space/startup_progress.py b/hf_space/startup_progress.py index 8af897e227e6221ef4cb1fdc11422d050482c766..f8c95146bb5b0d127644362c83ae457338910ea4 100644 --- a/hf_space/startup_progress.py +++ b/hf_space/startup_progress.py @@ -1,23 +1,23 @@ -"""English startup progress, measured from copied bytes and loaded models.""" -import html - - -def render(state): - progress = state.get('models', {}).get('startup_progress', {}) - done = max(0, progress.get('completed', 0)) - total = max(0, progress.get('total', 0)) - label = html.escape(progress.get('label', 'Preparing runtime')) - unit = progress.get('unit') - if total and unit == 'bytes': - detail = f'{done / 1024**3:.1f} / {total / 1024**3:.1f} GB' - elif total: - detail = f'{int(done)} / {int(total)} ' + ('models' if unit == 'models' else 'groups') - else: - detail = '' - value = f' value="{min(100, done / total * 100):.2f}"' if total else '' - return ('
' - '

Loading weights

' - '

Models load automatically in the background when the demo starts. All users share the loaded models.

' - '
Loading progress' + detail + '
' - '' - '
' + label + '
') +"""English startup progress, measured from copied bytes and loaded models.""" +import html + + +def render(state): + progress = state.get('models', {}).get('startup_progress', {}) + done = max(0, progress.get('completed', 0)) + total = max(0, progress.get('total', 0)) + label = html.escape(progress.get('label', 'Preparing runtime')) + unit = progress.get('unit') + if total and unit == 'bytes': + detail = f'{done / 1024**3:.1f} / {total / 1024**3:.1f} GB' + elif total: + detail = f'{int(done)} / {int(total)} ' + ('models' if unit == 'models' else 'groups') + else: + detail = '' + value = f' value="{min(100, done / total * 100):.2f}"' if total else '' + return ('
' + '

Loading weights

' + '

Model files are prepared in the background and cached for all users.

' + '
Loading progress' + detail + '
' + '' + '
' + label + '
') diff --git a/infer.py b/infer.py index 77debb69846ffb1db038d01a418f092219dc30eb..6b7e77b6a9b822719bea4f18d3dc45fbeb244ac6 100644 --- a/infer.py +++ b/infer.py @@ -1,227 +1,350 @@ -#!/usr/bin/env python3 -"""Generate a textured PBR mesh from one RGB/RGBA image.""" -from __future__ import annotations - -import argparse -import json -import math -from pathlib import Path -from hf_space.timing import span - - -def parse_args(argv: list[str] | None = None) -> argparse.Namespace: - parser = argparse.ArgumentParser(description=__doc__) - parser.add_argument("--input", type=Path, - default=Path(__file__).resolve().parent / "assets/pitcher_plant.png", - help="source RGB/RGBA image (default: assets/pitcher_plant.png)") - parser.add_argument("--output", required=True, type=Path, help="new destination .glb file") - parser.add_argument("--preview-only", action="store_true", help="save final o-voxel and GPU multiview snapshots; skip GLB extraction") - parser.add_argument("--work-dir", type=Path, help="fresh intermediate directory (default: _work beside the GLB)") - parser.add_argument("--moge-model", required=True, help="local MoGe-2 checkpoint/directory or HF repo ID") - parser.add_argument("--rmbg-model", help="local RMBG-2 directory or HF repo ID; required for RGB input") - parser.add_argument("--use-input-alpha", action="store_true", help="require and reuse foreground RGBA alpha") - parser.add_argument("--resolution-level", type=int, default=9, help="MoGe-2 resolution level") - parser.add_argument("--geometry-package", required=True, type=Path, help="native or paper geometry source checkout; see docs/geometry.md") - parser.add_argument("--geometry-model", required=True, help="local geometry weights directory or HF repo ID") - parser.add_argument("--geometry-revision", help="geometry weights HF revision") - parser.add_argument("--geometry-initial-mesh-scale", type=float, default=1.0) - parser.add_argument("--geometry-target-padding", type=int, default=3) - parser.add_argument("--geometry-max-iterations", type=int, default=2) - parser.add_argument("--firered-package", type=Path, default=Path(__file__).resolve().parent / "supratex/frontend/material_model") - parser.add_argument("--material-backend", choices=("qwen-image", "firered"), default="qwen-image") - parser.add_argument("--material-base", help="selected backend base model: local directory or HF repo ID") - parser.add_argument("--firered-base", help="legacy FireRed base option; requires --material-backend firered") - parser.add_argument("--qwen-runtime", type=Path, help="native Qwen source snapshot (default: ckpts/qwen-runtime)") - parser.add_argument("--qwen-deps", type=Path, help="Qwen dependency directory (default: ckpts/qwen-deps)") - parser.add_argument("--material-model", help="local backend checkpoint bundle or HF repo ID") - parser.add_argument("--material-revision", help="material checkpoint HF revision") - parser.add_argument("--material-weights", type=Path) - parser.add_argument("--material-spec", type=Path) - parser.add_argument("--material-checkpoint-manifest", type=Path) - parser.add_argument("--pretrained", help="local Supra-Tex completion bundle or HF repo ID") - parser.add_argument("--revision", help="completion bundle HF revision") - parser.add_argument("--model-config", type=Path, help="completion ABI JSON; overrides bundle config") - parser.add_argument("--checkpoint", type=Path, help="completion weights; overrides bundle checkpoint") - parser.add_argument("--trellis-root", type=Path, help="TRELLIS.2 VAE directory (all four components for HF releases)") - parser.add_argument("--cache-dir", type=Path, help="shared HF download cache for all adapters") - parser.add_argument("--local-files-only", action="store_true", help="resolve models locally without Hub downloads") - parser.add_argument("--attention", choices=("auto", "flash_attn_3", "flash_attn", "sdpa", "xformers"), default="auto", - help="inference attention backend; auto selects by GPU and installed extensions") - parser.add_argument("--low-vram", action=argparse.BooleanOptionalAction, default=None, - help="CPU offload for Qwen and 3D stages; auto enabled on GPUs with up to 24 GiB") - parser.add_argument("--seed", type=int, default=42, help="geometry and completion seed (material runner uses its fixed seed 42)") - parser.add_argument("--texture-size", type=int, default=4096) - parser.add_argument("--decimation-target", type=int, default=500000) - parser.add_argument("--texture-format", choices=("png", "webp"), default="png") - parser.add_argument("--gutter", action="store_true", help="enable chart-isolated one-texel base-color gutter padding (default: disabled)") - parser.add_argument("--dual-view", action="store_true", help="add v4 back rendering, Astra/image_gen generation, and dual-view texturing") - parser.add_argument("--back-view-dir", type=Path, help="reuse a prepared, geometry-bound back package instead of generating a new back") - args = parser.parse_args(argv) - if args.back_view_dir is not None and not args.dual_view: - parser.error("--back-view-dir requires --dual-view") - if args.dual_view and (args.texture_size != 4096 or args.texture_format != "png" - or args.decimation_target != 500000 or args.gutter): - parser.error("dual-view reprojection requires 4096 PNG, a 500000-face target, and no --gutter") - material_paths = (args.material_weights, args.material_spec, args.material_checkpoint_manifest) - if args.material_backend == "qwen-image": - if args.firered_base or any(material_paths): - parser.error("FireRed options require --material-backend firered") - args.qwen_runtime = args.qwen_runtime or Path(__file__).resolve().parent / "ckpts/qwen-runtime" - args.qwen_deps = args.qwen_deps or Path(__file__).resolve().parent / "ckpts/qwen-deps" - args.material_base = args.material_base or str(Path(__file__).resolve().parent / "ckpts/qwen-base") - args.material_model = args.material_model or str(Path(__file__).resolve().parent / "ckpts/qwen-material") - else: - if args.qwen_runtime is not None or args.qwen_deps is not None: - parser.error("--qwen-runtime and --qwen-deps require --material-backend qwen-image") - if args.material_base and args.firered_base: - parser.error("provide only one of --material-base and --firered-base") - if not (args.material_base or args.firered_base): - parser.error("FireRed requires --material-base or --firered-base") - if args.material_model and any(material_paths): - parser.error("--material-model cannot be combined with explicit material artifact paths") - if not args.material_model and not all(material_paths): - parser.error("provide --material-model or all three --material-weights, --material-spec, --material-checkpoint-manifest") - if not args.pretrained and not all((args.model_config, args.checkpoint, args.trellis_root)): - parser.error("provide --pretrained or all of --model-config, --checkpoint, and --trellis-root") - if args.seed < 0 or args.texture_size <= 0 or args.decimation_target <= 0: - parser.error("seed must be nonnegative; texture-size and decimation-target must be positive") - if not math.isfinite(args.geometry_initial_mesh_scale) or args.geometry_initial_mesh_scale <= 0: - parser.error("geometry-initial-mesh-scale must be finite and positive") - if not 1 <= args.geometry_target_padding <= 30 or args.geometry_max_iterations < 1: - parser.error("geometry-target-padding must be in [1, 30]; geometry-max-iterations must be positive") - if not 0 <= args.resolution_level <= 9: - parser.error("resolution-level must be in [0, 9]") - return args - - -def _completion_paths(args: argparse.Namespace) -> tuple[Path, Path, Path]: - """Resolve local/Hub artifacts and check missing completion files before GPU work.""" - from supratex.weights import COMPLETION_ALLOW_PATTERNS, resolve_completion_artifacts, resolve_model_dir - root = None - if args.pretrained: - root = resolve_model_dir(args.pretrained, revision=args.revision, cache_dir=args.cache_dir, - local_files_only=args.local_files_only, - allow_patterns=COMPLETION_ALLOW_PATTERNS) - return resolve_completion_artifacts(root, model_config=args.model_config, - checkpoint=args.checkpoint, trellis_root=args.trellis_root) - - -def run(args: argparse.Namespace, *, resident=None) -> Path: - """Run the existing image frontend, then complete and export its generated mesh.""" - image, output = args.input.expanduser().resolve(), args.output.expanduser().resolve() - work = (args.work_dir or output.parent / f"{output.stem}_work").expanduser().resolve() - expected_suffix = '.npz' if args.preview_only else '.glb' - if output.suffix.lower() != expected_suffix: - raise ValueError(f"--output must have a {expected_suffix} suffix") - if output.exists(): - raise FileExistsError(f"refusing to overwrite existing output: {output}") - if work.exists() and (not work.is_dir() or any(work.iterdir())): - raise FileExistsError(f"--work-dir must be a fresh empty directory: {work}") - if not image.is_file(): - raise FileNotFoundError(image) - from PIL import Image - with Image.open(image) as source: - alpha = source.getchannel("A").getextrema() if source.mode == "RGBA" else (255, 255) - has_alpha = alpha[0] < 255 - if args.use_input_alpha and not has_alpha: - raise ValueError("--use-input-alpha requires a non-opaque RGBA input") - if not has_alpha and not args.rmbg_model: - raise ValueError("RGB/opaque input requires --rmbg-model") - source.load() - from supratex.inference.runtime import configure_inference_runtime - runtime = configure_inference_runtime(attention=getattr(args, "attention", "auto"), - low_vram=getattr(args, "low_vram", None), - material_backend=args.material_backend) - # The adapters validate source compatibility and full checkpoint contents. - from supratex.frontend.geometry import validate_source - validate_source(args.geometry_package.expanduser()) - if args.material_backend == "qwen-image": - from supratex.frontend.qwen_image import validate - validate(args.qwen_runtime, args.material_base, args.material_model, deps=args.qwen_deps, - revision=args.material_revision, cache_dir=args.cache_dir, local_files_only=args.local_files_only) - elif not (args.firered_package.expanduser() / "infer.py").is_file(): - raise FileNotFoundError(f"material entrypoint is missing: {args.firered_package / 'infer.py'}") - material_paths = {} - for name in ("material_weights", "material_spec", "material_checkpoint_manifest"): - path = getattr(args, name) - if path is not None: - path = path.expanduser().resolve() - if not path.is_file(): - raise FileNotFoundError(f"missing {name}: {path}") - material_paths[name] = path - if args.dual_view and args.back_view_dir is None: - import shutil - if shutil.which("codex") is None: - raise RuntimeError("automatic back generation requires Codex with gpt-6-astra and built-in image_gen") - config, checkpoint, encoders = _completion_paths(args) - from supratex.frontend.pipeline import run as run_frontend - from supratex.inference import run_completion - work.mkdir(parents=True, exist_ok=True) - handoff = work / "input" - receipt = {"status": "running", "stage": "frontend", "input": str(image), "output": str(output), - "work_dir": str(work), "seed": args.seed, "gutter": args.gutter, - "material_backend": args.material_backend, "material_seed": 42, - "runtime": runtime, - "model_config": str(config), "checkpoint": str(checkpoint), "dual_view": args.dual_view} - def record() -> None: - temporary = work / ".run.json.tmp" - temporary.write_text(json.dumps(receipt, indent=2) + "\n", encoding="utf-8") - temporary.replace(work / "run.json") - record() - try: - print(f"[1/2] Foreground, MoGe-2, optimized geometry, material maps -> {handoff}", flush=True) - with span('frontend'): - result = run_frontend(image, handoff, moge_model=args.moge_model, rmbg_model=args.rmbg_model, - use_input_alpha=args.use_input_alpha, resolution_level=args.resolution_level, - geometry_package=args.geometry_package.expanduser().resolve(), geometry_model=args.geometry_model, - geometry_revision=args.geometry_revision, geometry_seed=args.seed, - geometry_initial_mesh_scale=args.geometry_initial_mesh_scale, - geometry_target_padding=args.geometry_target_padding, - geometry_max_iterations=args.geometry_max_iterations, - firered_package=args.firered_package.expanduser().resolve(), firered_base=args.firered_base, - material_backend=args.material_backend, material_base=args.material_base, - qwen_runtime=args.qwen_runtime if args.material_backend == "qwen-image" else None, - qwen_deps=args.qwen_deps if args.material_backend == "qwen-image" else None, - material_model=args.material_model, **material_paths, - revision=args.material_revision, cache_dir=args.cache_dir, local_files_only=args.local_files_only, - resident=resident, - **({"record_material_provenance": True} if args.dual_view else {})) - optional = {} - if args.dual_view: - receipt.update(stage="back_view", frontend=result) - record() - if args.back_view_dir is None: - from supratex.backview.workflow import build_back_view - back_dir = build_back_view(handoff, work / "back_view", cache_dir=args.cache_dir, - local_files_only=args.local_files_only) - else: - back_dir = args.back_view_dir.expanduser().resolve() - optional = {"back_view_dir": back_dir, "work_dir": work / "dual_view"} - receipt["back_view_dir"] = str(back_dir) - receipt.update(stage="completion", frontend=result) - record() - print(f"[2/2] Surface observations, PBR completion, {'o-voxel preview' if args.preview_only else 'GLB export'} -> {output}", flush=True) - with span('completion'): - run_completion(input_dir=handoff, output=output, model_config=config, checkpoint=checkpoint, - trellis_root=encoders, seed=args.seed, texture_size=args.texture_size, - decimation_target=args.decimation_target, texture_format=args.texture_format, gutter=args.gutter, - resident=resident, preview_only=args.preview_only, **optional) - if not output.is_file() or output.stat().st_size == 0: - raise RuntimeError(f"completion did not emit a nonempty GLB: {output}") - except Exception as error: - receipt.update(status="failed", error=f"{type(error).__name__}: {error}") - record() - raise - receipt.update(status="complete", stage="complete") - record() - print(f"Textured mesh: {output}\nIntermediate files and run receipt: {work}", flush=True) - return output - - -def main() -> None: - with span('pipeline'): - run(parse_args()) - - -if __name__ == "__main__": - main() +#!/usr/bin/env python3 +"""Generate a textured PBR mesh from one RGB/RGBA image.""" +from __future__ import annotations + +import argparse +import json +import math +import os +from pathlib import Path +import sys + + +def _pipeline_overrides(args: argparse.Namespace) -> dict: + """Validate explicit model selections before resolving a release.""" + overrides = {} + for component, attribute in (("completion", "pretrained"), ("material", "material_model"), + ("qwen-base", "material_base"), ("geometry", "geometry_model"), + ("vae", "trellis_root"), ("moge2", "moge_model"), ("rmbg2", "rmbg_model")): + value = getattr(args, attribute) + if value is not None: + if component not in ("completion", "material") and not Path(value).expanduser().exists(): + raise ValueError(f"--pipeline requires a local --{attribute.replace('_', '-')} override") + overrides[component] = value + if args.revision is not None: + if args.pretrained is None: + raise ValueError("--revision requires --pretrained with --pipeline; use --pipeline-revision for the release") + overrides["completion_revision"] = args.revision + if args.material_revision is not None: + if args.material_model is None: + raise ValueError("--material-revision requires --material-model with --pipeline") + overrides["material_revision"] = args.material_revision + return overrides + + +def _prepare_pipeline(args: argparse.Namespace, overrides: dict, descriptor: dict, *, has_alpha: bool) -> dict: + """Prepare the release already selected during argument validation.""" + from scripts.prepare_models import prepare_pipeline + from scripts.run_example import read_config + + source = descriptor["repo_id"] if descriptor["kind"] == "hub" else args.pipeline + revision = descriptor.get("revision") if descriptor["kind"] == "hub" else args.pipeline_revision + prepared = prepare_pipeline(source, output_dir=args.models_dir, revision=revision, + cache_dir=args.cache_dir, local_files_only=args.local_files_only, + with_rmbg=not has_alpha, + overrides=overrides) + return read_config(prepared) + + +def _validate_inference_options(args: argparse.Namespace) -> None: + """Check explicit values first, then call again after release defaults.""" + from scripts.prepare_models import _pipeline + + if args.preview_only and (args.dual_view or args.refine): + raise ValueError("o-voxel preview is single-view and precedes mesh refinement/export") + if args.back_view_dir is not None and not args.dual_view: + raise ValueError("--back-view-dir requires --dual-view") + _pipeline.validate_export_options(args) + if not math.isfinite(args.geometry_initial_mesh_scale) or args.geometry_initial_mesh_scale <= 0: + raise ValueError("geometry-initial-mesh-scale must be finite and positive") + if not 1 <= args.geometry_target_padding <= 30 or args.geometry_max_iterations < 1: + raise ValueError("geometry-target-padding must be in [1, 30]; geometry-max-iterations must be positive") + if not 0 <= args.resolution_level <= 9: + raise ValueError("resolution-level must be in [0, 9]") + + +def _validate_input_output(args: argparse.Namespace, *, require_rmbg: bool = True) -> tuple[Path, Path, Path, bool]: + """Check local inputs and destinations without creating or downloading files.""" + image, output = args.input.expanduser().resolve(), args.output.expanduser().absolute() + expected_suffix = ".npz" if args.preview_only else ".glb" + if output.suffix.lower() != expected_suffix: + raise ValueError(f"--output must have a {expected_suffix} suffix") + if output.exists() or output.is_symlink(): + raise FileExistsError(f"refusing to overwrite existing output: {output}") + output = output.resolve() + work = (args.work_dir or output.parent / f"{output.stem}_work").expanduser().absolute() + if work.is_symlink() and not work.exists(): + raise FileExistsError(f"--work-dir must be a fresh empty directory: {work}") + work = work.resolve() + if work.exists() and (not work.is_dir() or any(work.iterdir())): + raise FileExistsError(f"--work-dir must be a fresh empty directory: {work}") + for destination in (output, work): + for parent in destination.parents: + if parent.exists(): + if not parent.is_dir(): + raise NotADirectoryError(f"destination parent must be a directory: {parent}") + break + if not image.is_file(): + raise FileNotFoundError(f"input image is missing: {image}") + from PIL import Image + with Image.open(image) as source: + low, high = source.getchannel("A").getextrema() if source.mode == "RGBA" else (255, 255) + has_alpha = low < 255 and high > 0 + if args.use_input_alpha and not has_alpha: + raise ValueError("--use-input-alpha requires a non-opaque RGBA input with foreground pixels") + if require_rmbg and not has_alpha and not args.rmbg_model: + raise ValueError("RGB/opaque input requires --rmbg-model") + source.load() + if args.dual_view and args.back_view_dir is None: + import shutil + if shutil.which("codex") is None: + raise RuntimeError("automatic back generation requires Codex with gpt-6-astra and built-in image_gen") + return image, output, work, has_alpha + + +def parse_args(argv: list[str] | None = None) -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__, allow_abbrev=False) + parser.add_argument("--pipeline", help="SUPRA-Tex release directory, pipeline.json, or HF repo ID; prepares and loads all stages") + parser.add_argument("--pipeline-revision", help="HF revision of the complete SUPRA-Tex release") + parser.add_argument("--models-dir", type=Path, help="prepared model directory for --pipeline (default: ckpts/ready)") + parser.add_argument("--input", type=Path, + help="source RGB/RGBA image (default: assets/spiny_whelk_shell.png, or " + "assets/cinema_signage.png with --dual-view)") + parser.add_argument("--output", required=True, type=Path, help="new destination .glb file") + parser.add_argument("--preview-only", action="store_true", help="save/render final o-voxel; export mesh separately") + parser.add_argument("--work-dir", type=Path, help="fresh intermediate directory (default: _work beside the GLB)") + parser.add_argument("--moge-model", help="local MoGe-2 checkpoint/directory or HF repo ID") + parser.add_argument("--rmbg-model", help="local RMBG-2 directory or HF repo ID; required for RGB input") + parser.add_argument("--use-input-alpha", action="store_true", help="require and reuse foreground RGBA alpha") + parser.add_argument("--resolution-level", type=int, default=9, help="MoGe-2 resolution level") + parser.add_argument("--geometry-package", type=Path, help="native or paper geometry source checkout; see docs/geometry.md") + parser.add_argument("--geometry-model", help="local geometry weights directory or HF repo ID") + parser.add_argument("--geometry-revision", help="geometry weights HF revision") + parser.add_argument("--geometry-initial-mesh-scale", type=float, default=1.0) + parser.add_argument("--geometry-scale-search", action="store_true", + help="enable bounded mesh-scale calibration (default: Pixal3D fixed scale)") + parser.add_argument("--parallel-frontend", action="store_true", + help="run geometry and Qwen material concurrently on the first two visible GPUs") + parser.add_argument("--geometry-target-padding", type=int, default=3, help="scale-search target padding") + parser.add_argument("--geometry-max-iterations", type=int, default=2, help="scale-search update limit") + parser.add_argument("--firered-package", type=Path, default=Path(__file__).resolve().parent / "supratex/frontend/material_model") + parser.add_argument("--material-backend", choices=("qwen-image", "firered"), default="qwen-image") + parser.add_argument("--material-base", help="selected backend base model: local directory or HF repo ID") + parser.add_argument("--firered-base", help="legacy FireRed base option; requires --material-backend firered") + parser.add_argument("--qwen-runtime", type=Path, help="native Qwen source snapshot (default: bundled runtime)") + parser.add_argument("--qwen-deps", type=Path, help="optional Qwen dependency overlay; otherwise use the active Python environment") + parser.add_argument("--material-model", help="local backend checkpoint bundle or HF repo ID") + parser.add_argument("--material-revision", help="material checkpoint HF revision") + parser.add_argument("--material-weights", type=Path) + parser.add_argument("--material-spec", type=Path) + parser.add_argument("--material-checkpoint-manifest", type=Path) + parser.add_argument("--pretrained", help="local Supra-Tex completion bundle or HF repo ID") + parser.add_argument("--revision", help="completion bundle HF revision") + parser.add_argument("--model-config", type=Path, help="completion ABI JSON; overrides bundle config") + parser.add_argument("--checkpoint", type=Path, help="completion weights; overrides bundle checkpoint") + parser.add_argument("--trellis-root", type=Path, help="TRELLIS.2 VAE directory (all four components for HF releases)") + parser.add_argument("--cache-dir", type=Path, help="shared HF download cache for all adapters") + parser.add_argument("--local-files-only", action="store_true", help="resolve models locally without Hub downloads") + parser.add_argument("--attention", choices=("auto", "flash_attn_3", "flash_attn", "sdpa", "xformers"), + help="inference attention backend; auto selects by GPU and installed extensions") + parser.add_argument("--low-vram", action=argparse.BooleanOptionalAction, default=None, + help="CPU offload for Qwen and 3D stages; auto enabled on GPUs with up to 24 GiB") + parser.add_argument("--seed", type=int, help="geometry and completion seed (material runner uses its fixed seed 42)") + parser.add_argument("--texture-size", type=int) + parser.add_argument("--decimation-target", type=int) + parser.add_argument("--texture-format", choices=("png", "webp")) + parser.add_argument("--gutter", action="store_true", help="enable chart-isolated one-texel base-color gutter padding (default: disabled); " + "with --refine, the observation-guided gutter is used") + parser.add_argument("--refine", action="store_true", + help="refine the visible texture after export: write the observed reference-view albedo into the atlas " + "and repair its UV gutters from the observation (sharper fine text and detail; needs the 4096 PNG " + "atlas; adds about 2 minutes on an RTX 3090)") + parser.add_argument("--remove-floaters", action=argparse.BooleanOptionalAction, default=True, + help="before completion, drop disconnected geometry pieces that project mostly outside the " + "image foreground (default: on); extra geometry attached to the object is kept") + parser.add_argument("--dual-view", action="store_true", help="generate a back view (needs the external Codex CLI) and texture with both views; includes --refine for both views") + parser.add_argument("--back-view-dir", type=Path, help="reuse a prepared, geometry-bound back package instead of generating a new back") + args = parser.parse_args(argv) + from scripts.run_example import MODEL_KEYS, _pipeline_config, default_input, inference_defaults + if args.input is None: + args.input = default_input(args.dual_view) + if not args.pipeline and (args.pipeline_revision is not None or args.models_dir is not None): + parser.error("--pipeline-revision and --models-dir require --pipeline") + material_paths = (args.material_weights, args.material_spec, args.material_checkpoint_manifest) + if args.pipeline and args.material_backend != "qwen-image": + parser.error("--pipeline requires --material-backend qwen-image") + if args.material_backend == "qwen-image" and (args.firered_base or any(material_paths)): + parser.error("FireRed options require --material-backend firered") + try: + # Reject explicit argument/input errors before fetching any metadata. + _validate_inference_options(args) + if args.pipeline: + overrides = _pipeline_overrides(args) + _, _, _, has_alpha = _validate_input_output(args, require_rmbg=False) + from scripts.prepare_models import _pipeline + manifest, descriptor = _pipeline.resolve_pipeline( + args.pipeline, revision=args.pipeline_revision, cache_dir=args.cache_dir, + local_files_only=args.local_files_only) + defaults = manifest["defaults"] + else: + defaults = inference_defaults() + for key, value in defaults.items(): + if getattr(args, key) is None: + setattr(args, key, value) + # Release defaults can introduce a dual-view conflict; check before downloads. + _validate_inference_options(args) + if args.pipeline: + prepared = _prepare_pipeline(args, overrides, descriptor, has_alpha=has_alpha) + except (OSError, ValueError, RuntimeError) as error: + parser.error(str(error)) + if args.pipeline: + explicit_pretrained = args.pretrained is not None + direct_overrides = {"geometry_package", "qwen_runtime", "qwen_deps", "model_config", "checkpoint"} + for key in MODEL_KEYS: + if prepared.get(key) is not None and (key not in direct_overrides or getattr(args, key) is None): + if explicit_pretrained and key in ("model_config", "checkpoint"): + continue + action = parser._option_string_actions["--" + key.replace("_", "-")] + value = prepared[key] + setattr(args, key, action.type(value) if action.type else value) + args.local_files_only = True + args._pipeline_config = prepared + if not all((args.moge_model, args.geometry_package, args.geometry_model)): + parser.error("provide --pipeline or all of --moge-model, --geometry-package, and --geometry-model") + if args.material_backend == "qwen-image": + args.qwen_runtime = args.qwen_runtime or _pipeline_config().runtime_root(Path(__file__).resolve().parent) / "qwen" + args.material_base = args.material_base or str(Path(__file__).resolve().parent / "ckpts/qwen-base") + args.material_model = args.material_model or str(Path(__file__).resolve().parent / "ckpts/qwen-material") + else: + if args.qwen_runtime is not None or args.qwen_deps is not None: + parser.error("--qwen-runtime and --qwen-deps require --material-backend qwen-image") + if args.material_base and args.firered_base: + parser.error("provide only one of --material-base and --firered-base") + if not (args.material_base or args.firered_base): + parser.error("FireRed requires --material-base or --firered-base") + if args.material_model and any(material_paths): + parser.error("--material-model cannot be combined with explicit material artifact paths") + if not args.material_model and not all(material_paths): + parser.error("provide --material-model or all three --material-weights, --material-spec, --material-checkpoint-manifest") + if not args.pretrained and not all((args.model_config, args.checkpoint, args.trellis_root)): + parser.error("provide --pretrained or all of --model-config, --checkpoint, and --trellis-root") + return args + + +def _completion_paths(args: argparse.Namespace) -> tuple[Path, Path, Path]: + """Resolve local/Hub artifacts and check missing completion files before GPU work.""" + from supratex.weights import COMPLETION_ALLOW_PATTERNS, resolve_completion_artifacts, resolve_model_dir + root = None + if args.pretrained: + root = resolve_model_dir(args.pretrained, revision=args.revision, cache_dir=args.cache_dir, + local_files_only=args.local_files_only, + allow_patterns=COMPLETION_ALLOW_PATTERNS) + return resolve_completion_artifacts(root, model_config=args.model_config, + checkpoint=args.checkpoint, trellis_root=args.trellis_root) + + +def run(args: argparse.Namespace) -> Path: + """Run the existing image frontend, then complete and export its generated mesh.""" + if getattr(args, "_pipeline_config", None) is not None: + from scripts.run_example import runtime_environment + environment = runtime_environment(args._pipeline_config, args.attention) + for key in ("ATTN_BACKEND", "SPARSE_ATTN_BACKEND"): + if key not in environment: + os.environ.pop(key, None) + os.environ.update(environment) + for source in reversed(environment["PYTHONPATH"].split(os.pathsep)): + if source not in sys.path: + sys.path.insert(0, source) + image, output, work, _ = _validate_input_output(args) + from supratex.inference.runtime import configure_inference_runtime + runtime = configure_inference_runtime(attention=getattr(args, "attention", "auto"), + low_vram=getattr(args, "low_vram", None), + material_backend=args.material_backend) + # The adapters validate source compatibility and full checkpoint contents. + from supratex.frontend.geometry import validate_source + validate_source(args.geometry_package.expanduser()) + if args.material_backend == "qwen-image": + from supratex.frontend.qwen_image import validate + validate(args.qwen_runtime, args.material_base, args.material_model, deps=args.qwen_deps, + revision=args.material_revision, cache_dir=args.cache_dir, local_files_only=args.local_files_only) + elif not (args.firered_package.expanduser() / "infer.py").is_file(): + raise FileNotFoundError(f"material entrypoint is missing: {args.firered_package / 'infer.py'}") + material_paths = {} + for name in ("material_weights", "material_spec", "material_checkpoint_manifest"): + path = getattr(args, name) + if path is not None: + path = path.expanduser().resolve() + if not path.is_file(): + raise FileNotFoundError(f"missing {name}: {path}") + material_paths[name] = path + config, checkpoint, encoders = _completion_paths(args) + from supratex.frontend.pipeline import run as run_frontend + from supratex.inference import run_completion + work.mkdir(parents=True, exist_ok=True) + handoff = work / "input" + receipt = {"status": "running", "stage": "frontend", "input": str(image), "output": str(output), + "work_dir": str(work), "seed": args.seed, "gutter": args.gutter, + "material_backend": args.material_backend, "material_seed": 42, + "runtime": runtime, + "model_config": str(config), "checkpoint": str(checkpoint), "dual_view": args.dual_view, + "refine": args.refine, "remove_floaters": args.remove_floaters, "preview_only": args.preview_only} + def record() -> None: + temporary = work / ".run.json.tmp" + temporary.write_text(json.dumps(receipt, indent=2) + "\n", encoding="utf-8") + temporary.replace(work / "run.json") + record() + try: + print(f"[1/2] Foreground, MoGe-2, geometry, material maps -> {handoff}", flush=True) + result = run_frontend(image, handoff, moge_model=args.moge_model, rmbg_model=args.rmbg_model, + use_input_alpha=args.use_input_alpha, resolution_level=args.resolution_level, + geometry_package=args.geometry_package.expanduser().resolve(), geometry_model=args.geometry_model, + geometry_revision=args.geometry_revision, geometry_seed=args.seed, + geometry_initial_mesh_scale=args.geometry_initial_mesh_scale, + geometry_scale_search=args.geometry_scale_search, + parallel_frontend=args.parallel_frontend, + geometry_target_padding=args.geometry_target_padding, + geometry_max_iterations=args.geometry_max_iterations, + firered_package=args.firered_package.expanduser().resolve(), firered_base=args.firered_base, + material_backend=args.material_backend, material_base=args.material_base, + qwen_runtime=args.qwen_runtime if args.material_backend == "qwen-image" else None, + qwen_deps=args.qwen_deps if args.material_backend == "qwen-image" else None, + material_model=args.material_model, **material_paths, + revision=args.material_revision, cache_dir=args.cache_dir, local_files_only=args.local_files_only, + **({"record_material_provenance": True} if args.dual_view else {})) + optional = {"refine": True, "work_dir": work / "refine"} if args.refine else {} + if args.dual_view: + receipt.update(stage="back_view", frontend=result) + record() + if args.back_view_dir is None: + from supratex.backview.workflow import build_back_view + back_dir = build_back_view(handoff, work / "back_view", cache_dir=args.cache_dir, + local_files_only=args.local_files_only) + else: + back_dir = args.back_view_dir.expanduser().resolve() + optional = {"back_view_dir": back_dir, "work_dir": work / "dual_view"} + receipt["back_view_dir"] = str(back_dir) + receipt.update(stage="completion", frontend=result) + record() + print(f"[2/2] Surface observations, PBR completion, {'o-voxel preview' if args.preview_only else 'GLB export'} -> {output}", flush=True) + run_completion(input_dir=handoff, output=output, model_config=config, checkpoint=checkpoint, + trellis_root=encoders, seed=args.seed, texture_size=args.texture_size, + decimation_target=args.decimation_target, texture_format=args.texture_format, gutter=args.gutter, + remove_floaters=args.remove_floaters, preview_only=args.preview_only, **optional) + if not output.is_file() or output.stat().st_size == 0: + raise RuntimeError(f"completion did not emit a nonempty result: {output}") + except Exception as error: + receipt.update(status="failed", error=f"{type(error).__name__}: {error}") + record() + raise + receipt.update(status="complete", stage="complete") + record() + print(f"Textured mesh: {output}\nIntermediate files and run receipt: {work}", flush=True) + return output + + +def main() -> None: + run(parse_args()) + + +if __name__ == "__main__": + main() diff --git a/pyproject.toml b/pyproject.toml index 6a3844105cc54dd50a63dd128f6cc3641ed8fb2e..3a720245f0483d6334aa3c2e4e3673db66bbca8b 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,31 +1,29 @@ -[build-system] -requires = ["setuptools>=68"] -build-backend = "setuptools.build_meta" - -[project] -name = "supra-tex" -version = "0.1.0" -description = "Supra-Tex: super-aligned 3D texture generation" -readme = "README.md" -license = {file = "LICENSE"} -requires-python = ">=3.10" -dependencies = [ - "torch>=2.6", "torchvision>=0.21", "numpy", "Pillow", "scipy", - "easydict", "tqdm", "einops", "safetensors", "huggingface-hub", - "trimesh", "opencv-python-headless", "imageio", "imageio-ffmpeg", "numba", -] - -[project.optional-dependencies] -training = ["pandas>=2", "transformers==4.57.3", "diffusers==0.37.1", "lpips", "plyfile", "tensorboard"] -logging = ["wandb"] -frontend = ["transformers==4.57.3", "diffusers==0.37.1", "accelerate", "requests", "timm"] - -[tool.setuptools] -license-files = ["LICENSE"] - -[tool.setuptools.packages.find] -include = ["supratex*"] - -[tool.setuptools.package-data] -"supratex.backview.harness" = ["prompts/*.txt", "README.md"] - +[build-system] +requires = ["setuptools>=68"] +build-backend = "setuptools.build_meta" + +[project] +name = "supra-tex" +version = "0.1.0" +description = "SUPRA-Tex: 3D Generation with Super-Aligned Texture" +readme = "README.md" +license = {file = "LICENSE"} +requires-python = ">=3.10" +dependencies = [ + "torch>=2.6", "torchvision>=0.21", "numpy", "Pillow", "scipy", + "easydict", "tqdm", "einops", "safetensors", "huggingface-hub", + "trimesh", "opencv-python-headless", "imageio", "imageio-ffmpeg", "numba", +] + +[project.optional-dependencies] +frontend = ["transformers==4.57.3", "diffusers==0.37.1", "accelerate", "requests", "timm"] + +[tool.setuptools] +license-files = ["LICENSE", "THIRD_PARTY_NOTICES.md"] + +[tool.setuptools.packages.find] +include = ["supratex*"] + +[tool.setuptools.package-data] +"supratex.backview.harness" = ["prompts/*.txt", "README.md"] + diff --git a/scripts/check_runtime.py b/scripts/check_runtime.py index f3998ae1c910e622fa2a537329dfdd0693f5404f..75980574b464b74aa28377ec770b9c107b43efef 100644 --- a/scripts/check_runtime.py +++ b/scripts/check_runtime.py @@ -1,179 +1,172 @@ -#!/usr/bin/env python3 -"""Check the source release's real runtime imports and prepared model files.""" -from __future__ import annotations - -import argparse -import json -import os -from pathlib import Path -import subprocess -import sys - -ROOT = Path(__file__).resolve().parents[1] -RUNTIME = ROOT / "deployment/runtime" - - -def runtime_environment(config: dict | None = None) -> dict[str, str]: - env = os.environ.copy() - env.update(PYTHONNOUSERSITE="1", PYTHONDONTWRITEBYTECODE="1", HF_HUB_OFFLINE="1", TRANSFORMERS_OFFLINE="1", - FLEX_GEMM_USE_AUTOTUNE_CACHE="0", FLEX_GEMM_AUTOSAVE_AUTOTUNE_CACHE="0", - PYTHONPATH=os.pathsep.join(str(p) for p in - (ROOT, RUNTIME / "moge", RUNTIME / "utils3d"))) - if config and config.get("torch_home"): - env["TORCH_HOME"] = config["torch_home"] - return env - - -def run_import_check(code: str, *, config: dict | None = None) -> str: - result = subprocess.run([sys.executable, "-c", code], cwd=ROOT, - env=runtime_environment(config), capture_output=True, text=True, timeout=300) - if result.returncode: - raise RuntimeError((result.stderr or result.stdout)[-5000:].strip()) - return result.stdout.strip() - - -def import_checks(attention: str) -> list[tuple[str, str]]: - bootstrap = ("from supratex.inference.runtime import configure_inference_runtime; " - f"configure_inference_runtime(attention={attention!r}, material_backend='qwen-image'); ") - return [ - ("Python, PyTorch and CUDA", """ -import os, sys, torch, torchvision -if sys.version_info[:2] != (3, 10): - raise RuntimeError('The supported installation uses Python 3.10') -profile = os.environ.get('SUPRATEX_RUNTIME_PROFILE', 'recorded') -versions = {'recorded': ('2.6.0', '12.4', '0.21.0'), 'torch28': ('2.8.0', '12.8', '0.23.0')} -if profile not in versions: - raise RuntimeError('Unknown SUPRATEX_RUNTIME_PROFILE: ' + profile) -expected_torch, expected_cuda, expected_vision = versions[profile] -if torch.__version__.split('+')[0] != expected_torch or torch.version.cuda != expected_cuda: - raise RuntimeError('Expected PyTorch/CUDA ' + str((expected_torch, expected_cuda)) + '; got ' + str((torch.__version__, torch.version.cuda))) -if torchvision.__version__.split('+')[0] != expected_vision: - raise RuntimeError('Expected torchvision ' + expected_vision) -if not torch.cuda.is_available(): - raise RuntimeError('CUDA is unavailable; check the NVIDIA driver and GPU access') -x = torch.tensor([1., 2.], device='cuda') -assert (x * x).sum().item() == 5. -torch.cuda.synchronize() -print(torch.cuda.get_device_name(0)) -"""), - ("Native extensions, MoGe and utils3d", """ -import torch, natten, natten.functional -import o_voxel, flex_gemm, cumesh, nvdiffrast.torch -import moge.model.v2, utils3d -if not (hasattr(natten.functional, 'na2d') or - all(hasattr(natten.functional, n) for n in ('na2d_av', 'na2d_qk'))): - raise RuntimeError('NATTEN does not expose the NAF attention API') -assert hasattr(utils3d, 'pt'), 'Bundled utils3d API is required' -print('Native modules and bundled camera utilities imported') -"""), - ("Geometry and NAF implementation", bootstrap + - f"import sys; sys.path.insert(0, {str(RUNTIME / 'geometry')!r}); " - f"sys.path.insert(0, {str(RUNTIME / 'naf')!r}); " - "from pixal3d.pipelines import Pixal3DImageTo3DPipeline; " - "from pixal3d.trainers.flow_matching.mixins.image_conditioned_proj import DinoV3ProjFeatureExtractor; " - "import hubconf; print('Geometry, projection conditioner and NAF imported')"), - ("Qwen models, PBR pipeline and inference helpers", - f"import sys; sys.path.insert(0, {str(RUNTIME / 'qwen')!r}); " - f"sys.path.insert(0, {str(RUNTIME / 'qwen/examples/qwen_image_21/pbr')!r}); " - "import peft; " - "import diffsynth.models.qwen_image_21_dit; " - "import diffsynth.models.qwen_image_21_text_encoder; " - "import diffsynth.models.qwen_image_21_vae; " - "import diffsynth.pipelines.qwen_image_21; " - "import diffsynth.pipelines.qwen_image_21_pbr; " - "from train import build_model, decode_pair, endpoint_noise; " - "from validate_local import checkpoint_component, explicit_load; " - "from encode_cache import PROMPT; " - "import inference_offload; print('Actual Qwen model and helper imports passed')"), - ("Completion implementation", bootstrap + - "from supratex.inference import run_completion; " - "from supratex.models import InterleavedFlowTransformer; " - "print('Completion implementation imported')"), - ] - - -def require_nonempty(*paths: Path) -> None: - for path in paths: - if not path.is_file() or path.stat().st_size == 0: - raise FileNotFoundError(f"Required nonempty file is missing: {path}") - - -def file_checks(config: dict) -> list[tuple[str, object]]: - sys.path.insert(0, str(ROOT)) - from supratex.frontend.geometry import validate_source, validate_model - from supratex.frontend.qwen_image import validate - from supratex.weights import resolve_completion_artifacts - - def completion(): - resolve_completion_artifacts(Path(config["pretrained"]), - model_config=config.get("model_config"), - checkpoint=config.get("checkpoint"), - trellis_root=Path(config["trellis_root"])) - - moge = Path(config["moge_model"]) - checks = [ - ("Bundled geometry source", lambda: validate_source(Path(config["geometry_package"]))), - ("Geometry and DINO model files", lambda: validate_model( - Path(config["geometry_package"]), Path(config["geometry_model"]))), - ("Qwen base and released material adapter", lambda: validate( - Path(config["qwen_runtime"]), config["material_base"], config["material_model"], - deps=Path(config["qwen_deps"]), local_files_only=True)), - ("Released completion and four sparse VAEs", completion), - ("MoGe weights", lambda: require_nonempty(moge / "model.pt" if moge.is_dir() else moge)), - ("NAF local source and weights", lambda: require_nonempty( - Path(config["torch_home"]) / "hub/valeoai_NAF_main/hubconf.py", - Path(config["torch_home"]) / "hub/checkpoints/naf_release.pth")), - ] - if config.get("rmbg_model"): - checks.append(("RGB foreground model", lambda: require_nonempty(*( - Path(config["rmbg_model"]) / name for name in - ("config.json", "model.safetensors", "BiRefNet_config.py", "birefnet.py", "preprocessor_config.json"))))) - return checks - - -def main(argv: list[str] | None = None) -> int: - parser = argparse.ArgumentParser(description=__doc__) - parser.add_argument("--config", type=Path, default=ROOT / "ckpts/ready/models.json") - scope = parser.add_mutually_exclusive_group() - scope.add_argument("--imports-only", action="store_true", help="check installed software without model weights") - scope.add_argument("--files-only", action="store_true", help="validate model layout and headers without CUDA checks") - parser.add_argument("--attention", choices=("auto", "sdpa", "flash_attn", "flash_attn_3", "xformers"), default="sdpa") - args = parser.parse_args(argv) - results = [] - - def check(name, operation): - print(f"[check] {name}", file=sys.stderr, flush=True) - try: - detail = operation() - entry = {"check": name, "status": "passed"} - if isinstance(detail, str) and detail: - entry["detail"] = detail - results.append(entry) - except Exception as error: - results.append({"check": name, "status": "failed", "error": str(error)}) - - config = None - if not args.imports_only: - try: - sys.path.insert(0, str(Path(__file__).resolve().parent)) - from run_example import read_config - config = read_config(args.config) - os.environ.update(HF_HUB_OFFLINE="1", TRANSFORMERS_OFFLINE="1") - os.environ["TORCH_HOME"] = config["torch_home"] - checks = file_checks(config) - except Exception as error: - results.append({"check": "Prepared model configuration", "status": "failed", "error": str(error)}) - else: - for name, operation in checks: - check(name, operation) - if not args.files_only: - for name, code in import_checks(args.attention): - check(name, lambda code=code: run_import_check(code, config=config)) - passed = bool(results) and all(entry["status"] == "passed" for entry in results) - print(json.dumps({"status": "passed" if passed else "failed", "checks": results, - "scope": "Model files and actual runtime imports; full image-to-GLB inference is a separate check"}, indent=2)) - return 0 if passed else 1 - - -if __name__ == "__main__": - raise SystemExit(main()) +#!/usr/bin/env python3 +"""Check the source release's real runtime imports and prepared model files.""" +from __future__ import annotations + +import argparse +import json +import os +from pathlib import Path +import subprocess +import sys + +ROOT = Path(__file__).resolve().parents[1] +RUNTIME = ROOT / "deployment/runtime" + + +def runtime_environment(config: dict | None = None) -> dict[str, str]: + env = os.environ.copy() + env.update(PYTHONNOUSERSITE="1", PYTHONDONTWRITEBYTECODE="1", HF_HUB_OFFLINE="1", TRANSFORMERS_OFFLINE="1", + FLEX_GEMM_USE_AUTOTUNE_CACHE="0", FLEX_GEMM_AUTOSAVE_AUTOTUNE_CACHE="0", + PYTHONPATH=os.pathsep.join(str(p) for p in + (ROOT, RUNTIME / "moge", RUNTIME / "utils3d"))) + if config and config.get("torch_home"): + env["TORCH_HOME"] = config["torch_home"] + return env + + +def run_import_check(code: str, *, config: dict | None = None) -> str: + result = subprocess.run([sys.executable, "-c", code], cwd=ROOT, + env=runtime_environment(config), capture_output=True, text=True, timeout=300) + if result.returncode: + raise RuntimeError((result.stderr or result.stdout)[-5000:].strip()) + return result.stdout.strip() + + +def import_checks(attention: str) -> list[tuple[str, str]]: + bootstrap = ("from supratex.inference.runtime import configure_inference_runtime; " + f"configure_inference_runtime(attention={attention!r}, material_backend='qwen-image'); ") + return [ + ("Python, PyTorch and CUDA", """ +import sys, torch, torchvision +if sys.version_info[:2] != (3, 10): + raise RuntimeError('The supported installation uses Python 3.10') +if torch.__version__.split('+')[0] != '2.6.0' or torch.version.cuda != '12.4': + raise RuntimeError('Expected PyTorch 2.6.0 with CUDA 12.4; got ' + str((torch.__version__, torch.version.cuda))) +if torchvision.__version__.split('+')[0] != '0.21.0': + raise RuntimeError('Expected torchvision 0.21.0') +if not torch.cuda.is_available(): + raise RuntimeError('CUDA is unavailable; check the NVIDIA driver and GPU access') +x = torch.tensor([1., 2.], device='cuda') +assert (x * x).sum().item() == 5. +torch.cuda.synchronize() +print(torch.cuda.get_device_name(0)) +"""), + ("Native extensions, MoGe and utils3d", """ +import torch, natten, natten.functional +import o_voxel, flex_gemm, cumesh, nvdiffrast.torch +import moge.model.v2, utils3d +if not (hasattr(natten.functional, 'na2d') or + all(hasattr(natten.functional, n) for n in ('na2d_av', 'na2d_qk'))): + raise RuntimeError('NATTEN does not expose the NAF attention API') +assert hasattr(utils3d, 'pt'), 'Bundled utils3d API is required' +print('Native modules and bundled camera utilities imported') +"""), + ("Geometry and NAF implementation", bootstrap + + f"import sys; sys.path.insert(0, {str(RUNTIME / 'geometry')!r}); " + f"sys.path.insert(0, {str(RUNTIME / 'naf')!r}); " + "from pixal3d.pipelines import Pixal3DImageTo3DPipeline; " + "from pixal3d.modules.projected_image_feature_extractor import DinoV3ProjFeatureExtractor; " + "import hubconf; print('Geometry, projection conditioner and NAF imported')"), + ("Qwen models, PBR pipeline and inference helpers", + f"import sys; sys.path.insert(0, {str(RUNTIME / 'qwen')!r}); " + f"sys.path.insert(0, {str(RUNTIME / 'qwen/examples/qwen_image_21/pbr')!r}); " + "import peft; " + "import diffsynth.models.qwen_image_21_dit; " + "import diffsynth.models.qwen_image_21_text_encoder; " + "import diffsynth.models.qwen_image_21_vae; " + "import diffsynth.pipelines.qwen_image_21; " + "import diffsynth.pipelines.qwen_image_21_pbr; " + "from inference_utils import build_model, decode_pair, endpoint_noise, checkpoint_component, explicit_load, PROMPT; " + "import inference_offload; print('Actual Qwen model and helper imports passed')"), + ("Completion implementation", bootstrap + + "from supratex.inference import run_completion; " + "from supratex.models import InterleavedFlowTransformer; " + "print('Completion implementation imported')"), + ] + + +def require_nonempty(*paths: Path) -> None: + for path in paths: + if not path.is_file() or path.stat().st_size == 0: + raise FileNotFoundError(f"Required nonempty file is missing: {path}") + + +def file_checks(config: dict) -> list[tuple[str, object]]: + sys.path.insert(0, str(ROOT)) + from supratex.frontend.geometry import validate_source, validate_model + from supratex.frontend.qwen_image import validate + from supratex.weights import resolve_completion_artifacts + + def completion(): + resolve_completion_artifacts(Path(config["pretrained"]), + model_config=config.get("model_config"), + checkpoint=config.get("checkpoint"), + trellis_root=Path(config["trellis_root"])) + + moge = Path(config["moge_model"]) + checks = [ + ("Bundled geometry source", lambda: validate_source(Path(config["geometry_package"]))), + ("Geometry and DINO model files", lambda: validate_model( + Path(config["geometry_package"]), Path(config["geometry_model"]))), + ("Qwen base and released material adapter", lambda: validate( + Path(config["qwen_runtime"]), config["material_base"], config["material_model"], + deps=Path(config["qwen_deps"]) if config.get("qwen_deps") else None, local_files_only=True)), + ("Released completion and four sparse VAEs", completion), + ("MoGe weights", lambda: require_nonempty(moge / "model.pt" if moge.is_dir() else moge)), + ("NAF local source and weights", lambda: require_nonempty( + Path(config["torch_home"]) / "hub/valeoai_NAF_main/hubconf.py", + Path(config["torch_home"]) / "hub/checkpoints/naf_release.pth")), + ] + if config.get("rmbg_model"): + checks.append(("RGB foreground model", lambda: require_nonempty(*( + Path(config["rmbg_model"]) / name for name in + ("config.json", "model.safetensors", "BiRefNet_config.py", "birefnet.py"))))) + return checks + + +def main(argv: list[str] | None = None) -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--config", type=Path, default=ROOT / "ckpts/ready/models.json") + scope = parser.add_mutually_exclusive_group() + scope.add_argument("--imports-only", action="store_true", help="check installed software without model weights") + scope.add_argument("--files-only", action="store_true", help="validate model layout and headers without CUDA checks") + parser.add_argument("--attention", choices=("auto", "sdpa", "flash_attn", "flash_attn_3", "xformers"), default="sdpa") + args = parser.parse_args(argv) + results = [] + + def check(name, operation): + print(f"[check] {name}", file=sys.stderr, flush=True) + try: + detail = operation() + entry = {"check": name, "status": "passed"} + if isinstance(detail, str) and detail: + entry["detail"] = detail + results.append(entry) + except Exception as error: + results.append({"check": name, "status": "failed", "error": str(error)}) + + config = None + if not args.imports_only: + try: + sys.path.insert(0, str(Path(__file__).resolve().parent)) + from run_example import read_config + config = read_config(args.config) + os.environ.update(HF_HUB_OFFLINE="1", TRANSFORMERS_OFFLINE="1") + os.environ["TORCH_HOME"] = config["torch_home"] + checks = file_checks(config) + except Exception as error: + results.append({"check": "Prepared model configuration", "status": "failed", "error": str(error)}) + else: + for name, operation in checks: + check(name, operation) + if not args.files_only: + for name, code in import_checks(args.attention): + check(name, lambda code=code: run_import_check(code, config=config)) + passed = bool(results) and all(entry["status"] == "passed" for entry in results) + print(json.dumps({"status": "passed" if passed else "failed", "checks": results, + "scope": "Model files and actual runtime imports; full image-to-GLB inference is a separate check"}, indent=2)) + return 0 if passed else 1 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/scripts/convert_geometry_weights.py b/scripts/convert_geometry_weights.py new file mode 100644 index 0000000000000000000000000000000000000000..d5325fb198e0b5ab523c21c12ebd83d8eba9099a --- /dev/null +++ b/scripts/convert_geometry_weights.py @@ -0,0 +1,41 @@ +#!/usr/bin/env python3 +"""Write runtime-dtype copies of the prepared native geometry weights. + +The native flow checkpoints store float32 for blocks that run in bfloat16. +A cast copy halves the bytes read on every load (faster on slow disks) and +loads bit-identical tensors. Copies are written beside the originals as +``.runtime.safetensors`` (about 8 GB) and are used automatically. +Run it in the inference environment with a GPU visible; no GPU memory is used. +""" +from __future__ import annotations + +import argparse +import json +from pathlib import Path +import sys + +ROOT = Path(__file__).resolve().parents[1] + + +def main(argv: list[str] | None = None) -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--config", type=Path, default=ROOT / "ckpts/ready/models.json", + help="prepared models.json (default: ckpts/ready/models.json)") + args = parser.parse_args(argv) + config_path = args.config.expanduser().resolve() + config = json.loads(config_path.read_text(encoding="utf-8")) + package = (config_path.parent / config["geometry_package"]).resolve() + model_root = (config_path.parent / config["geometry_model"]).resolve() + sys.path[:0] = [str(package), str(ROOT)] + from supratex.frontend.geometry import NATIVE_GEOMETRY_MODELS + from supratex.frontend.geometry_loading import write_runtime_weights + + stems = json.loads((model_root / "pipeline.json").read_text(encoding="utf-8"))["args"]["models"] + for name in NATIVE_GEOMETRY_MODELS: + written = write_runtime_weights(model_root / stems[name]) + print(f"{name}: {written or 'already compact; skipped'}", flush=True) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/scripts/install_runtime.py b/scripts/install_runtime.py index ee85a440d287bfb2bd5fca3689e6b31449ab6b18..1f0c6738407cb4fdb559206bb02fc3bbca7a2c2b 100644 --- a/scripts/install_runtime.py +++ b/scripts/install_runtime.py @@ -1,276 +1,356 @@ -#!/usr/bin/env python3 -"""Install the recorded Linux/Python 3.10/CUDA 12.4 runtime into its own venv. - -No model weights are downloaded. Native sources come from the recorded public -revisions in deployment/docs/DEPENDENCIES.md. Run --dry-run to inspect commands -without creating directories, invoking subprocesses, or accessing the network. -""" - -from __future__ import annotations - -import argparse -import json -import os -from pathlib import Path -import platform -import re -import shlex -import shutil -import subprocess -import sys - - -ROOT = Path(__file__).resolve().parents[1] -SOURCES = ( - ("CuMesh", "https://github.com/JeffreyXiang/CuMesh.git", "ef40f3ae4612ea3ec934b4151e019dc95c5d276b", ""), - ("FlexGEMM", "https://github.com/JeffreyXiang/FlexGEMM.git", "b1965ad0ae4bc8488d13b5280afb54439ce27359", ""), - ("nvdiffrast", "https://github.com/NVlabs/nvdiffrast.git", "253ac4fcea7de5f396371124af597e6cc957bfae", ""), - ("TRELLIS.2", "https://github.com/microsoft/TRELLIS.2.git", "5565d240c4a494caaf9ece7a554542b76ffa36d3", "o-voxel"), -) -O_VOXEL_FLAGS = " ".join(( - "--extended-lambda", "-U__CUDA_NO_HALF_OPERATORS__", - "-U__CUDA_NO_HALF_CONVERSIONS__", "-U__CUDA_NO_HALF2_OPERATORS__", -)) -CONSTRAINTS = "torch==2.6.0+cu124\ntorchvision==0.21.0+cu124\n" -BUILD_TOOLS = ("pip==25.1.1", "setuptools==80.9.0", "wheel==0.45.1", "packaging==25.0") -# Official https://whl.natten.org/ index links this exact 0.17.5 release asset. -# The former shi-labs.com wheel index has an expired TLS certificate. -NATTEN_WHEEL = ( - "https://github.com/SHI-Labs/NATTEN/releases/download/v0.17.5/" - "natten-0.17.5%2Btorch260cu124-cp310-cp310-linux_x86_64.whl" -) -NATTEN_VERSION = "0.17.5+torch260cu124" -PYTHON_PROBE = ( - "import json,sys; import venv,ensurepip; " - "print(json.dumps({'version':list(sys.version_info[:2])," - "'prefix':sys.prefix,'base_prefix':sys.base_prefix}))" -) -TORCH_PROBE = ( - "import torch,torchvision; " - "assert torch.__version__ == '2.6.0+cu124', torch.__version__; " - "assert torchvision.__version__ == '0.21.0+cu124', torchvision.__version__; " - "assert torch.version.cuda == '12.4', torch.version.cuda; " - "assert torch.cuda.is_available(), 'No usable NVIDIA CUDA GPU/driver'; " - "torch.empty(1, device='cuda'); torch.cuda.synchronize(); " - "print('Torch 2.6.0 / CUDA 12.4 GPU check passed')" -) - - -class InstallError(RuntimeError): - """An actionable installation failure, with no automatic destructive repair.""" - - -def clean_environment(venv: Path) -> dict[str, str]: - """Do not inherit Python overlays or pip destinations from the invoking shell.""" - env = { - key: value for key, value in os.environ.items() - if not key.startswith("PIP_") and key not in { - "PYTHONPATH", "PYTHONHOME", "PYTHONUSERBASE", "VIRTUAL_ENV", - "CONDA_PREFIX", "CONDA_DEFAULT_ENV", - } - } - env.update({ - "PIP_CONFIG_FILE": os.devnull, - "PIP_REQUIRE_VIRTUALENV": "true", - "PIP_DISABLE_PIP_VERSION_CHECK": "1", - "PIP_CACHE_DIR": str(venv / ".cache" / "pip"), - "TORCH_EXTENSIONS_DIR": str(venv / ".cache" / "torch_extensions"), - "ATTN_BACKEND": "sdpa", - "SPARSE_ATTN_BACKEND": "sdpa", - "BUILD_TARGET": "cuda", - }) - return env - - -def capture(command: list[str], env: dict[str, str]) -> str: - try: - return subprocess.check_output(command, env=env, text=True, stderr=subprocess.STDOUT).strip() - except (OSError, subprocess.CalledProcessError) as exc: - detail = getattr(exc, "output", str(exc)) - raise InstallError(f"Command failed: {shlex.join(command)}\n{detail}") from exc - - -def validate_venv_destination(venv: Path) -> None: - if venv in {Path(sys.prefix).resolve(), Path(sys.base_prefix).resolve()}: - raise InstallError("--venv cannot be the invoking Python environment. Choose a separate directory.") - if not venv.exists(): - return - config = venv / "pyvenv.cfg" - if not config.is_file() or not (venv / "bin" / "python").is_file(): - raise InstallError(f"Refusing to overwrite {venv}: existing destination is not a complete venv. Choose a new --venv.") - settings = dict( - line.lower().strip().split("=", 1) - for line in config.read_text().splitlines() if "=" in line - ) - settings = {key.strip(): value.strip() for key, value in settings.items()} - if settings.get("include-system-site-packages") != "false": - raise InstallError(f"{venv} must have include-system-site-packages = false. Choose a new --venv.") - - -def prerequisites(python: str, venv: Path, env: dict[str, str]) -> dict[str, str]: - """Validate before any directory creation, download, or package installation.""" - if platform.system() != "Linux": - raise InstallError("This recorded native runtime requires Linux.") - if platform.machine().lower() not in {"x86_64", "amd64"}: - raise InstallError("The pinned NATTEN 0.17.5 wheel requires Linux x86_64.") - validate_venv_destination(venv) - try: - info = json.loads(capture([python, "-I", "-c", PYTHON_PROBE], env)) - except (ValueError, KeyError) as exc: - raise InstallError("Could not determine the selected Python version.") from exc - if info.get("version") != [3, 10]: - raise InstallError(f"Python 3.10 is required; {python} reports {info.get('version')}. Use --python /path/to/python3.10.") - if Path(info["prefix"]).resolve() == venv: - raise InstallError("--venv cannot be the selected --python environment. Choose a separate directory.") - if not shutil.which("git", path=env.get("PATH")): - raise InstallError("git is missing from PATH; install git before rerunning.") - compiler = shlex.split(env.get("CXX", "c++")) - if not compiler or not shutil.which(compiler[0], path=env.get("PATH")): - raise InstallError("A C++ compiler is required; install g++/c++ or set CXX to a working compiler.") - capture([*compiler, "--version"], env) - cuda_home = env.get("CUDA_HOME") or env.get("CUDA_PATH") - nvcc = str(Path(cuda_home).expanduser() / "bin" / "nvcc") if cuda_home else shutil.which("nvcc", path=env.get("PATH")) - if not nvcc or not Path(nvcc).is_file(): - raise InstallError("CUDA 12.4 nvcc is required. Install the CUDA toolkit and set CUDA_HOME or add its bin directory to PATH.") - version = capture([nvcc, "--version"], env) - if not re.search(r"release\s+12\.4(?:\D|$)", version): - raise InstallError(f"CUDA toolkit 12.4 is required; {nvcc} reports:\n{version}") - result = env.copy() - result["CUDA_HOME"] = str(Path(cuda_home).expanduser().resolve() if cuda_home else Path(nvcc).resolve().parent.parent) - result["PATH"] = os.pathsep.join((str(venv / "bin"), str(Path(result["CUDA_HOME"]) / "bin"), env.get("PATH", ""))) - return result - - -def validate_checkout(path: Path, url: str, revision: str, env: dict[str, str]) -> None: - """Reuse only the exact, clean checkout; never reset or clean user sources.""" - if not (path / ".git").exists(): - raise InstallError(f"Refusing to overwrite {path}: not a Git checkout. Choose a fresh --venv.") - git = ["git", "-C", str(path)] - top = capture([*git, "rev-parse", "--show-toplevel"], env) - origin = capture([*git, "remote", "get-url", "origin"], env) - head = capture([*git, "rev-parse", "HEAD"], env) - dirty = capture([*git, "status", "--porcelain", "--untracked-files=no"], env) - if Path(top).resolve() != path.resolve() or origin.rstrip("/") != url or head != revision or dirty: - raise InstallError( - f"Refusing to alter existing checkout {path}: expected origin {url}, HEAD {revision}, " - "and no tracked changes (including submodules). Repair it manually or choose a fresh --venv." - ) - - -def run(command: list[str], env: dict[str, str], dry_run: bool, extra_env: dict[str, str] | None = None) -> None: - display = ["env", *(f"{key}={value}" for key, value in extra_env.items()), *command] if extra_env else command - print("$ " + shlex.join(display), flush=True) - if not dry_run: - child_env = dict(env, **(extra_env or {})) - try: - subprocess.run(command, env=child_env, check=True) - except (OSError, subprocess.CalledProcessError) as exc: - raise InstallError(f"Installation stopped at: {shlex.join(command)}\n{exc}\nExisting files were kept; no rollback or deletion was attempted.") from exc - - -def install_natten(pip: list[str], env: dict[str, str], dry_run: bool) -> None: - # A direct wheel URL otherwise downloads again on each resumed installation. - # Query only the target venv; the final import checks still verify its ABI. - if dry_run: - print(f"# Skip NATTEN download only if this venv already has distribution version {NATTEN_VERSION}; final import checks still run.") - else: - probe = ( - "from importlib.metadata import version, PackageNotFoundError\n" - "try:\n print(version('natten'))\n" - "except PackageNotFoundError:\n print('')\n" - ) - installed = capture([pip[0], "-I", "-c", probe], env) - if installed == NATTEN_VERSION: - print(f"# NATTEN {NATTEN_VERSION} already installed in target venv; skipping download.") - return - run([*pip, "--no-deps", "--no-index", "--only-binary=:all:", NATTEN_WHEEL], env, dry_run) - - -def install(args: argparse.Namespace) -> None: - venv = Path(args.venv).expanduser().resolve() - python = shutil.which(args.python) or str(Path(args.python).expanduser().absolute()) - env = clean_environment(venv) - managed = venv / "supra-tex-runtime" - sources = managed / "sources" - constraints = managed / "constraints.txt" - target_python = str(venv / "bin" / "python") - pip = [target_python, "-I", "-m", "pip", "install", "--no-input"] - dry = args.dry_run - print(f"{'DRY RUN: no subprocesses, downloads, or writes.' if dry else 'Installing runtime.'}\nSource: {ROOT}\nVirtual environment: {venv}") - print("Prerequisites: Linux x86_64; selected Python 3.10 with venv/ensurepip; git; C++ compiler; CUDA toolkit 12.4 nvcc; usable NVIDIA GPU/driver.") - print("Child processes ignore inherited Python overlays and pip configuration/destinations; caches stay in the venv. Default attention is SDPA; FlashAttention is optional.") - print("Native source pins identify upstream inputs, not verified equivalent binaries. FlexGEMM upstream setup also writes ~/.flex_gemm/autotune_cache.json.") - if not dry: - for required in (ROOT / "deployment/runtime/requirements-python.txt", ROOT / "deployment/runtime/moge/pyproject.toml", ROOT / "deployment/runtime/utils3d/pyproject.toml", ROOT / "scripts/check_runtime.py"): - if not required.is_file(): - raise InstallError(f"Incomplete source release: missing {required}") - env = prerequisites(python, venv, env) - # Check all existing source trees before changing this environment. - for name, url, revision, _ in SOURCES: - path = sources / name - if path.exists(): - validate_checkout(path, url, revision, env) - if not venv.exists(): - run([python, "-I", "-m", "venv", str(venv)], env, dry) - else: - print("# Reuse only an isolated Python 3.10 venv; never clear an existing directory.") - if not dry: - info = json.loads(capture([target_python, "-I", "-c", PYTHON_PROBE], env)) - if info.get("version") != [3, 10] or Path(info["prefix"]).resolve() != venv or info["prefix"] == info["base_prefix"]: - raise InstallError(f"{venv} is not an isolated Python 3.10 environment. Choose a new --venv.") - managed.mkdir(parents=True, exist_ok=True) - sources.mkdir(exist_ok=True) - constraints.write_text(CONSTRAINTS) - print(f"# Write {constraints}: torch==2.6.0+cu124; torchvision==0.21.0+cu124") - run([*pip, *BUILD_TOOLS], env, dry) - run([*pip, "torch==2.6.0+cu124", "torchvision==0.21.0+cu124", "--index-url", "https://download.pytorch.org/whl/cu124"], env, dry) - run([*pip, "--constraint", str(constraints), "-r", str(ROOT / "deployment/runtime/requirements-python.txt")], env, dry) - run([*pip, "--no-build-isolation", "--no-deps", str(ROOT / "deployment/runtime/moge"), str(ROOT / "deployment/runtime/utils3d")], env, dry) - run([*pip, "--no-build-isolation", "--no-deps", "-e", str(ROOT)], env, dry) - install_natten(pip, env, dry) - run([target_python, "-I", "-c", TORCH_PROBE], env, dry) - for name, url, revision, subdir in SOURCES: - path = sources / name - if not path.exists(): - run(["git", "clone", "--no-checkout", url, str(path)], env, dry) - run(["git", "-C", str(path), "checkout", "--detach", revision], env, dry) - else: - print(f"# Verify origin={url}, HEAD={revision}, and clean tracked files: {path}; no reset/clean.") - run(["git", "-C", str(path), "submodule", "update", "--init", "--recursive"], env, dry) - flags = None - if name == "CuMesh": - print("# CuMesh/cubvh: append --extended-lambda for its CUDA device lambdas without patching source.") - flags = {"NVCC_APPEND_FLAGS": " ".join(filter(None, (env.get("NVCC_APPEND_FLAGS"), "--extended-lambda")))} - elif subdir == "o-voxel": - print("# o-voxel: NVCC_APPEND_FLAGS appends the recorded extended-lambda/half flags without patching source. This build still requires destination validation.") - flags = {"NVCC_APPEND_FLAGS": " ".join(filter(None, (env.get("NVCC_APPEND_FLAGS"), O_VOXEL_FLAGS)))} - run([*pip, "--no-build-isolation", "--no-deps", str(path / subdir)], env, dry, flags) - if args.flash_attn: - run([*pip, "--no-build-isolation", "--no-deps", "flash-attn==2.8.3"], env, dry) - run([target_python, "-I", "-c", "import flash_attn; print(flash_attn.__version__)"], env, dry) - run([target_python, "-I", "-m", "pip", "check"], env, dry) - run([target_python, "-I", "-c", TORCH_PROBE], env, dry) - run([target_python, "-I", str(ROOT / "scripts/check_runtime.py"), "--imports-only"], env, dry) - if dry: - print("Dry run complete. Prerequisites, builds, and imports have NOT been validated.") - else: - print(f"Runtime dependency checks passed in {venv}. Model acquisition and an image-to-GLB run are still required for inference acceptance. Use --attention sdpa for the default runtime.") - - -def main(argv: list[str] | None = None) -> int: - parser = argparse.ArgumentParser(description=__doc__) - parser.add_argument("--venv", default=".venv", help="Isolated environment directory (default: .venv relative to the current directory)") - parser.add_argument("--python", default=sys.executable, help="Python 3.10 executable used to create the venv (default: invoking Python)") - parser.add_argument("--dry-run", action="store_true", help="Print the complete plan without subprocesses, network, or filesystem writes") - parser.add_argument("--flash-attn", action="store_true", help="Also build optional FlashAttention 2.8.3; default SDPA needs neither FA2 nor FA3") - args = parser.parse_args(argv) - try: - install(args) - except (InstallError, OSError, ValueError) as exc: - print(f"ERROR: {exc}", file=sys.stderr) - return 1 - return 0 - - -if __name__ == "__main__": - raise SystemExit(main()) +#!/usr/bin/env python3 +"""Install the recorded Linux/Python 3.10/CUDA 12.4 runtime into Conda or a venv. + +No model weights are downloaded. Native sources come from the recorded public +revisions in deployment/docs/DEPENDENCIES.md. Run --dry-run to inspect commands +without creating directories, invoking subprocesses, or accessing the network. +""" + +from __future__ import annotations + +import argparse +import json +import os +from pathlib import Path +import platform +import re +import shlex +import shutil +import subprocess +import sys +import tempfile + + +ROOT = Path(__file__).resolve().parents[1] +SOURCES = ( + ("CuMesh", "https://github.com/JeffreyXiang/CuMesh.git", "ef40f3ae4612ea3ec934b4151e019dc95c5d276b", ""), + ("FlexGEMM", "https://github.com/JeffreyXiang/FlexGEMM.git", "b1965ad0ae4bc8488d13b5280afb54439ce27359", ""), + ("nvdiffrast", "https://github.com/NVlabs/nvdiffrast.git", "253ac4fcea7de5f396371124af597e6cc957bfae", ""), + ("TRELLIS.2", "https://github.com/microsoft/TRELLIS.2.git", "5565d240c4a494caaf9ece7a554542b76ffa36d3", "o-voxel"), +) +O_VOXEL_FLAGS = " ".join(( + "--extended-lambda", "-U__CUDA_NO_HALF_OPERATORS__", + "-U__CUDA_NO_HALF_CONVERSIONS__", "-U__CUDA_NO_HALF2_OPERATORS__", +)) +CONSTRAINTS = "torch==2.6.0+cu124\ntorchvision==0.21.0+cu124\n" +BUILD_TOOLS = ("pip==25.1.1", "setuptools==80.9.0", "wheel==0.45.1", "packaging==25.0") +# Official https://whl.natten.org/ index links this exact 0.17.5 release asset. +# The former shi-labs.com wheel index has an expired TLS certificate. +NATTEN_WHEEL = ( + "https://github.com/SHI-Labs/NATTEN/releases/download/v0.17.5/" + "natten-0.17.5%2Btorch260cu124-cp310-cp310-linux_x86_64.whl" +) +NATTEN_VERSION = "0.17.5+torch260cu124" +PYTHON_PROBE = ( + "import json,sys; " + "print(json.dumps({'version':list(sys.version_info[:2])," + "'prefix':sys.prefix,'base_prefix':sys.base_prefix}))" +) +TORCH_PROBE = ( + "import torch,torchvision; " + "assert torch.__version__ == '2.6.0+cu124', torch.__version__; " + "assert torchvision.__version__ == '0.21.0+cu124', torchvision.__version__; " + "assert torch.version.cuda == '12.4', torch.version.cuda; " + "assert torch.cuda.is_available(), 'No usable NVIDIA CUDA GPU/driver'; " + "torch.empty(1, device='cuda'); torch.cuda.synchronize(); " + "print('Torch 2.6.0 / CUDA 12.4 GPU check passed')" +) + + +class InstallError(RuntimeError): + """An actionable installation failure, with no automatic destructive repair.""" + + +def clean_environment(venv: Path, *, conda: bool = False) -> dict[str, str]: + """Do not inherit Python overlays or pip destinations from the invoking shell.""" + env = { + key: value for key, value in os.environ.items() + if not key.startswith("PIP_") and key not in { + "PYTHONPATH", "PYTHONHOME", "PYTHONUSERBASE", "VIRTUAL_ENV", + "CONDA_PREFIX", "CONDA_DEFAULT_ENV", + } + } + env.update({ + "PIP_CONFIG_FILE": os.devnull, + "PIP_REQUIRE_VIRTUALENV": "false" if conda else "true", + "PIP_DISABLE_PIP_VERSION_CHECK": "1", + "PIP_CACHE_DIR": str(venv / ".cache" / "pip"), + "TORCH_EXTENSIONS_DIR": str(venv / ".cache" / "torch_extensions"), + "FLEX_GEMM_AUTOTUNE_CACHE_PATH": str(venv / ".cache" / "flex_gemm" / "autotune_cache.json"), + "ATTN_BACKEND": "sdpa", + "SPARSE_ATTN_BACKEND": "sdpa", + "BUILD_TARGET": "cuda", + }) + if conda: + env["CONDA_PREFIX"] = str(venv) + if "CONDA_DEFAULT_ENV" in os.environ: + env["CONDA_DEFAULT_ENV"] = os.environ["CONDA_DEFAULT_ENV"] + return env + + +def capture(command: list[str], env: dict[str, str]) -> str: + try: + return subprocess.check_output(command, env=env, text=True, stderr=subprocess.STDOUT).strip() + except (OSError, subprocess.CalledProcessError) as exc: + detail = getattr(exc, "output", str(exc)) + raise InstallError(f"Command failed: {shlex.join(command)}\n{detail}") from exc + + +def validate_venv_destination(venv: Path) -> None: + if venv in {Path(sys.prefix).resolve(), Path(sys.base_prefix).resolve()}: + raise InstallError("--venv cannot be the invoking Python environment. Choose a separate directory.") + if not venv.exists(): + return + config = venv / "pyvenv.cfg" + if not config.is_file() or not (venv / "bin" / "python").is_file(): + raise InstallError(f"Refusing to overwrite {venv}: existing destination is not a complete venv. Choose a new --venv.") + settings = dict( + line.lower().strip().split("=", 1) + for line in config.read_text().splitlines() if "=" in line + ) + settings = {key.strip(): value.strip() for key, value in settings.items()} + if settings.get("include-system-site-packages") != "false": + raise InstallError(f"{venv} must have include-system-site-packages = false. Choose a new --venv.") + + +def active_conda_environment() -> Path: + prefix = os.environ.get("CONDA_PREFIX") + if not prefix: + raise InstallError("Activate a dedicated Conda environment first: conda activate supra-tex.") + destination = Path(prefix).expanduser().resolve() + conda_exe = os.environ.get("CONDA_EXE") + if os.environ.get("CONDA_DEFAULT_ENV") == "base" or ( + conda_exe and Path(conda_exe).resolve().parent.parent == destination + ): + raise InstallError("Create and activate a dedicated Conda environment; do not install into base.") + if (Path(sys.prefix).resolve() != destination + or Path(sys.base_prefix).resolve() != destination): + raise InstallError("Run this script with the active Conda environment's Python, without a nested venv.") + if not (destination / "conda-meta").is_dir() or not (destination / "bin/python").is_file(): + raise InstallError(f"{destination} is not a complete Conda Python environment.") + return destination + + +def validate_target_python(info: dict, destination: Path, *, conda: bool = False) -> None: + valid = info.get("version") == [3, 10] and Path(info["prefix"]).resolve() == destination + if conda: + valid = valid and Path(info["base_prefix"]).resolve() == destination + else: + valid = valid and info["prefix"] != info["base_prefix"] + if not valid: + raise InstallError(f"{destination} is not the expected Python 3.10 {'Conda environment' if conda else 'isolated venv'}.") + + +def prerequisites(python: str, venv: Path, env: dict[str, str], *, conda: bool = False) -> dict[str, str]: + """Validate before any directory creation, download, or package installation.""" + if platform.system() != "Linux": + raise InstallError("This recorded native runtime requires Linux.") + if platform.machine().lower() not in {"x86_64", "amd64"}: + raise InstallError("The pinned NATTEN 0.17.5 wheel requires Linux x86_64.") + if not conda: + validate_venv_destination(venv) + try: + probe = PYTHON_PROBE if conda else "import venv,ensurepip; " + PYTHON_PROBE + info = json.loads(capture([python, "-I", "-c", probe], env)) + except (ValueError, KeyError) as exc: + raise InstallError("Could not determine the selected Python version.") from exc + if info.get("version") != [3, 10]: + hint = "Create a Conda environment with python=3.10." if conda else "Use --python /path/to/python3.10." + raise InstallError(f"Python 3.10 is required; {python} reports {info.get('version')}. {hint}") + if conda: + validate_target_python(info, venv, conda=True) + elif Path(info["prefix"]).resolve() == venv: + raise InstallError("--venv cannot be the selected --python environment. Choose a separate directory.") + if not shutil.which("git", path=env.get("PATH")): + raise InstallError("git is missing from PATH; install git before rerunning.") + compiler = shlex.split(env.get("CXX", "c++")) + if not compiler or not shutil.which(compiler[0], path=env.get("PATH")): + raise InstallError("A C++ compiler is required; install g++/c++ or set CXX to a working compiler.") + capture([*compiler, "--version"], env) + cuda_home = env.get("CUDA_HOME") or env.get("CUDA_PATH") + nvcc = str(Path(cuda_home).expanduser() / "bin" / "nvcc") if cuda_home else shutil.which("nvcc", path=env.get("PATH")) + if not nvcc or not Path(nvcc).is_file(): + raise InstallError("CUDA 12.4 nvcc is required. Install the CUDA toolkit and set CUDA_HOME or add its bin directory to PATH.") + version = capture([nvcc, "--version"], env) + if not re.search(r"release\s+12\.4(?:\D|$)", version): + raise InstallError(f"CUDA toolkit 12.4 is required; {nvcc} reports:\n{version}") + result = env.copy() + result["CUDA_HOME"] = str(Path(cuda_home).expanduser().resolve() if cuda_home else Path(nvcc).resolve().parent.parent) + result["PATH"] = os.pathsep.join((str(venv / "bin"), str(Path(result["CUDA_HOME"]) / "bin"), env.get("PATH", ""))) + return result + + +def validate_checkout(path: Path, url: str, revision: str, env: dict[str, str]) -> None: + """Reuse only the exact, clean checkout; never reset or clean user sources.""" + if not (path / ".git").exists(): + raise InstallError(f"Refusing to overwrite {path}: not a Git checkout. Choose a fresh environment.") + git = ["git", "-C", str(path)] + top = capture([*git, "rev-parse", "--show-toplevel"], env) + origin = capture([*git, "remote", "get-url", "origin"], env) + head = capture([*git, "rev-parse", "HEAD"], env) + dirty = capture([*git, "status", "--porcelain", "--untracked-files=no"], env) + if Path(top).resolve() != path.resolve() or origin.rstrip("/") != url or head != revision or dirty: + raise InstallError( + f"Refusing to alter existing checkout {path}: expected origin {url}, HEAD {revision}, " + "and no tracked changes (including submodules). Repair it manually or choose a fresh environment." + ) + + +def run(command: list[str], env: dict[str, str], dry_run: bool, extra_env: dict[str, str] | None = None) -> None: + display = ["env", *(f"{key}={value}" for key, value in extra_env.items()), *command] if extra_env else command + print("$ " + shlex.join(display), flush=True) + if not dry_run: + child_env = dict(env, **(extra_env or {})) + try: + subprocess.run(command, env=child_env, check=True) + except (OSError, subprocess.CalledProcessError) as exc: + raise InstallError(f"Installation stopped at: {shlex.join(command)}\n{exc}\nExisting files were kept; no rollback or deletion was attempted.") from exc + + +def install_natten(pip: list[str], env: dict[str, str], dry_run: bool) -> None: + # A direct wheel URL otherwise downloads again on each resumed installation. + # Query only the target environment; the final import checks still verify its ABI. + if dry_run: + print(f"# Skip NATTEN download only if this environment already has distribution version {NATTEN_VERSION}; final import checks still run.") + else: + probe = ( + "from importlib.metadata import version, PackageNotFoundError\n" + "try:\n print(version('natten'))\n" + "except PackageNotFoundError:\n print('')\n" + ) + installed = capture([pip[0], "-I", "-c", probe], env) + if installed == NATTEN_VERSION: + print(f"# NATTEN {NATTEN_VERSION} already installed in target environment; skipping download.") + return + run([*pip, "--no-deps", "--no-index", "--only-binary=:all:", NATTEN_WHEEL], env, dry_run) + + +def install_flex_gemm(pip: list[str], source: Path, managed: Path, + env: dict[str, str], dry_run: bool) -> None: + """Redirect the pinned setup.py cache write without altering its checkout.""" + print(f"# FlexGEMM: build a temporary source copy under {managed}; " + f"install and runtime cache: {env['FLEX_GEMM_AUTOTUNE_CACHE_PATH']}.") + command = [*pip, "--no-build-isolation", "--no-deps"] + if dry_run: + run([*command, str(managed / "flex-gemm-build-" / "FlexGEMM")], env, True) + return + old = '''os.makedirs(os.path.expanduser("~/.flex_gemm"), exist_ok=True) +shutil.copyfile( + os.path.join(os.path.dirname(os.path.abspath(__file__)), "autotune_cache.json"), + os.path.expanduser('~/.flex_gemm/autotune_cache.json'), +)''' + new = '''cache_path = os.environ["FLEX_GEMM_AUTOTUNE_CACHE_PATH"] +os.makedirs(os.path.dirname(cache_path), exist_ok=True) +shutil.copyfile( + os.path.join(os.path.dirname(os.path.abspath(__file__)), "autotune_cache.json"), + cache_path, +)''' + setup = (source / "setup.py").read_text() + if setup.count(old) != 1: + raise InstallError("FlexGEMM setup.py does not match the pinned cache-write block; refusing an unverified patch.") + with tempfile.TemporaryDirectory(prefix="flex-gemm-build-", dir=managed) as temporary: + build_source = Path(temporary) / "FlexGEMM" + shutil.copytree(source, build_source, ignore=shutil.ignore_patterns(".git")) + (build_source / "setup.py").write_text(setup.replace(old, new)) + run([*command, str(build_source)], env, False) + + +def install(args: argparse.Namespace) -> None: + venv = active_conda_environment() if args.conda else Path(args.venv).expanduser().resolve() + selected_python = args.python or sys.executable + python = shutil.which(selected_python) or str(Path(selected_python).expanduser().absolute()) + env = clean_environment(venv, conda=args.conda) + managed = venv / "supra-tex-runtime" + sources = managed / "sources" + constraints = managed / "constraints.txt" + target_python = str(venv / "bin" / "python") + if args.conda: + python = target_python + pip = [target_python, "-I", "-m", "pip", "install", "--no-input"] + dry = args.dry_run + print(f"{'DRY RUN: no subprocesses, downloads, or writes.' if dry else 'Installing runtime.'}\nSource: {ROOT}\n{'Conda environment' if args.conda else 'Virtual environment'}: {venv}") + print("Prerequisites: Linux x86_64; Python 3.10; git; C++ compiler; CUDA toolkit 12.4 nvcc; usable NVIDIA GPU/driver.") + print("Child processes ignore inherited Python overlays and pip configuration/destinations; caches stay in the target environment. Without FlashAttention, inference uses SDPA; FlashAttention-2/3 is faster.") + print("Native source pins identify upstream inputs, not verified equivalent binaries. FlexGEMM's install-time cache path is patched in a temporary copy; its original checkout stays unchanged.") + if not dry: + for required in (ROOT / "deployment/runtime/requirements-python.txt", ROOT / "deployment/runtime/moge/pyproject.toml", ROOT / "deployment/runtime/utils3d/pyproject.toml", ROOT / "scripts/check_runtime.py"): + if not required.is_file(): + raise InstallError(f"Incomplete source release: missing {required}") + env = prerequisites(python, venv, env, conda=args.conda) + # Check all existing source trees before changing this environment. + for name, url, revision, _ in SOURCES: + path = sources / name + if path.exists(): + validate_checkout(path, url, revision, env) + if args.conda: + print("# Install into the active Conda environment.") + elif not venv.exists(): + run([python, "-I", "-m", "venv", str(venv)], env, dry) + else: + print("# Reuse only an isolated Python 3.10 venv; never clear an existing directory.") + if not dry: + info = json.loads(capture([target_python, "-I", "-c", PYTHON_PROBE], env)) + validate_target_python(info, venv, conda=args.conda) + managed.mkdir(parents=True, exist_ok=True) + sources.mkdir(exist_ok=True) + constraints.write_text(CONSTRAINTS) + print(f"# Write {constraints}: torch==2.6.0+cu124; torchvision==0.21.0+cu124") + run([*pip, *BUILD_TOOLS], env, dry) + run([*pip, "torch==2.6.0+cu124", "torchvision==0.21.0+cu124", "--index-url", "https://download.pytorch.org/whl/cu124"], env, dry) + run([*pip, "--constraint", str(constraints), "-r", str(ROOT / "deployment/runtime/requirements-python.txt")], env, dry) + run([*pip, "--no-build-isolation", "--no-deps", str(ROOT / "deployment/runtime/moge"), str(ROOT / "deployment/runtime/utils3d")], env, dry) + run([*pip, "--no-build-isolation", "--no-deps", "-e", str(ROOT)], env, dry) + install_natten(pip, env, dry) + run([target_python, "-I", "-c", TORCH_PROBE], env, dry) + for name, url, revision, subdir in SOURCES: + path = sources / name + if not path.exists(): + run(["git", "clone", "--no-checkout", url, str(path)], env, dry) + run(["git", "-C", str(path), "checkout", "--detach", revision], env, dry) + else: + print(f"# Verify origin={url}, HEAD={revision}, and clean tracked files: {path}; no reset/clean.") + run(["git", "-C", str(path), "submodule", "update", "--init", "--recursive"], env, dry) + flags = None + if name == "CuMesh": + print("# CuMesh/cubvh: append --extended-lambda for its CUDA device lambdas without patching source.") + flags = {"NVCC_APPEND_FLAGS": " ".join(filter(None, (env.get("NVCC_APPEND_FLAGS"), "--extended-lambda")))} + elif subdir == "o-voxel": + print("# o-voxel: NVCC_APPEND_FLAGS appends the recorded extended-lambda/half flags without patching source. This build still requires destination validation.") + flags = {"NVCC_APPEND_FLAGS": " ".join(filter(None, (env.get("NVCC_APPEND_FLAGS"), O_VOXEL_FLAGS)))} + if name == "FlexGEMM": + install_flex_gemm(pip, path, managed, env, dry) + else: + run([*pip, "--no-build-isolation", "--no-deps", str(path / subdir)], env, dry, flags) + if args.flash_attn: + run([*pip, "--no-build-isolation", "--no-deps", "flash-attn==2.8.3"], env, dry) + run([target_python, "-I", "-c", "import flash_attn; print(flash_attn.__version__)"], env, dry) + run([target_python, "-I", "-m", "pip", "check"], env, dry) + run([target_python, "-I", "-c", TORCH_PROBE], env, dry) + run([target_python, "-I", str(ROOT / "scripts/check_runtime.py"), "--imports-only"], env, dry) + if dry: + print("Dry run complete. Prerequisites, builds, and imports have NOT been validated.") + else: + print(f"Runtime dependency checks passed in {venv}. Model acquisition and an image-to-GLB run are still required for inference acceptance. FlashAttention-2/3 is faster than SDPA; see docs/installation.md.") + + +def main(argv: list[str] | None = None) -> int: + parser = argparse.ArgumentParser(description=__doc__) + destination = parser.add_mutually_exclusive_group() + destination.add_argument("--conda", action="store_true", help="Install into the active Conda environment (recommended; activate it first)") + destination.add_argument("--venv", default=".venv", help="Create/reuse an isolated venv instead (default without --conda: .venv)") + parser.add_argument("--python", help="Python 3.10 executable used to create a venv (default: invoking Python; unavailable with --conda)") + parser.add_argument("--dry-run", action="store_true", help="Print the complete plan without subprocesses, network, or filesystem writes") + parser.add_argument("--flash-attn", action="store_true", help="Also build optional FlashAttention 2.8.3; default SDPA needs neither FA2 nor FA3") + args = parser.parse_args(argv) + if args.conda and args.python: + parser.error("--python cannot be combined with --conda; use the active Conda environment's Python.") + try: + install(args) + except (InstallError, OSError, ValueError) as exc: + print(f"ERROR: {exc}", file=sys.stderr) + return 1 + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/scripts/prepare_models.py b/scripts/prepare_models.py index 6b464427c118e88285ade34521ade10183aaaea8..e7ded1c63542c706ee481b8bb2b339200443be98 100644 --- a/scripts/prepare_models.py +++ b/scripts/prepare_models.py @@ -1,360 +1,552 @@ -#!/usr/bin/env python3 -"""Download or bind Supra-Tex models, then publish a portable models.json. - -Downloads use the inherited Hugging Face configuration and your existing access. -Local overrides are linked without copying large weights and remain dependencies -of the prepared tree. No tensor data is loaded and no checksums are computed. -""" -from __future__ import annotations - -import argparse -import importlib.util -import json -import os -from pathlib import Path, PurePosixPath -import re -import shutil -import sys -from urllib.parse import unquote, urlsplit -from urllib.request import urlopen - - -ROOT = Path(__file__).resolve().parents[1] -PROJECT_REPO = "supra3d/supra-tex" -NAF_URL = "https://github.com/valeoai/NAF/releases/download/model/naf_release.pth" -COMPLETION_PATTERNS = ["config.json", "partial_normalization.json", "*.safetensors", - "*.safetensors.index.json", "LICENSE*", "README.md"] -MATERIAL_PATTERNS = ["adapter_config.json", "adapter_model.safetensors", "LICENSE*", "Notice", "README.md"] -GEOMETRY_KEYS = ("sparse_structure_flow_model", "sparse_structure_decoder", - "shape_slat_flow_model_512", "shape_slat_flow_model_1024", "shape_slat_decoder") -VAE_STEMS = ("shape_enc_next_dc_f16c32_fp16", "tex_enc_next_dc_f16c32_fp16", - "shape_dec_next_dc_f16c32_fp16", "tex_dec_next_dc_f16c32_fp16") -MARKER = ".prepare-models.json" - - -def parse_args(argv: list[str] | None = None) -> argparse.Namespace: - parser = argparse.ArgumentParser(description=__doc__) - parser.add_argument("--repo", default=PROJECT_REPO, help="project HF repo ID or URL (default: %(default)s)") - parser.add_argument("--revision", help="project repository revision; defaults to its current main") - parser.add_argument("--completion-subdir", default="completion") - parser.add_argument("--material-subdir", default="material") - parser.add_argument("--completion", help="override with standalone completion directory, HF repo ID, or HF URL") - parser.add_argument("--material", help="override with material adapter directory, HF repo ID, or HF URL") - parser.add_argument("--output-dir", type=Path, default=ROOT / "ckpts/ready") - for name in ("qwen-base", "geometry", "dino", "vae", "moge", "rmbg", "naf-weights"): - parser.add_argument("--" + name, type=Path, help="existing local override (no download)") - parser.add_argument("--with-rmbg", action="store_true", help="also prepare RMBG for RGB inputs; default example uses RGBA") - parser.add_argument("--dry-run", action="store_true", help="print the plan without writes, downloads, or model imports") - return parser.parse_args(argv) - - -def relative_name(value: str) -> str: - if (not isinstance(value, str) or not value or "\\" in value or ":" in value - or PurePosixPath(value).is_absolute() or any(part in ("", ".", "..") for part in value.split("/"))): - raise ValueError(f"expected a relative path without traversal: {value!r}") - return value - - -def model_source(value: str, revision: str | None = None) -> dict: - """Normalize local directories and HF URLs without importing the Hub client.""" - local = Path(value).expanduser() - if local.exists(): - if not local.is_dir(): - raise NotADirectoryError(f"model source must be a directory: {local}") - return {"kind": "local", "path": str(local.resolve())} - if value.startswith(("/", "./", "../", "~")) or value in (".", ".."): - raise FileNotFoundError(f"model directory does not exist: {local}") - if value.startswith("https://"): - url = urlsplit(value) - if url.hostname != "huggingface.co" or url.username or url.password or url.query or url.fragment: - raise ValueError("use a public-form https://huggingface.co/namespace/repo URL without credentials or query") - parts = unquote(url.path).strip("/").split("/") - if len(parts) == 4 and parts[2] == "tree": - if revision is not None and revision != parts[3]: - raise ValueError("HF URL revision conflicts with --revision") - revision = parts[3] - parts = parts[:2] - if len(parts) != 2: - raise ValueError("HF URL must select a repository root or /tree/") - value = "/".join(parts) - value = value.removeprefix("hf://") - if (not re.fullmatch(r"[A-Za-z0-9_][A-Za-z0-9_.-]*/[A-Za-z0-9_][A-Za-z0-9_.-]*", value) - or ".." in value or "--" in value or value.endswith((".", "-"))): - raise ValueError("model source must be an existing directory or a namespace/repo HF ID or URL") - return {"kind": "hub", "repo_id": value, "revision": revision} - - -def local_override(value: Path, *, file_allowed: bool = False) -> dict: - path = value.expanduser().resolve() - if not path.exists(): - raise FileNotFoundError(f"local override does not exist: {path}") - if not path.is_dir() and not (file_allowed and path.is_file()): - raise ValueError(f"local override must be {'a directory or file' if file_allowed else 'a directory'}: {path}") - return {"kind": "local", "path": str(path)} - - -def build_plan(args: argparse.Namespace) -> dict: - output = args.output_dir.expanduser().resolve() - manifest = json.loads((ROOT / "deployment/external_models.json").read_text()) - operations, models = [], {} - subdirs = {"completion": relative_name(args.completion_subdir), "material": relative_name(args.material_subdir)} - patterns = {"completion": COMPLETION_PATTERNS, "material": MATERIAL_PATTERNS} - project_needed = [name for name in subdirs if getattr(args, name) is None] - if project_needed: - project = model_source(args.repo, args.revision) - if project["kind"] != "hub": - raise ValueError("--repo must be a HF ID/URL; use --completion and --material for local directories") - operations.append(dict(project, name="project", destination="project", - allow_patterns=[f"{subdirs[name]}/{pattern}" for name in project_needed for pattern in patterns[name]], gated=False)) - for name in subdirs: - value = getattr(args, name) - if value is None: - models[name] = {"kind": "prepared", "path": f"project/{subdirs[name]}"} - else: - source = model_source(value, args.revision) - models[name] = source if source["kind"] == "local" else {"kind": "prepared", "path": name} - if source["kind"] == "hub": - operations.append(dict(source, name=name, destination=name, allow_patterns=patterns[name], gated=False)) - for key, flag in (("qwen-base", "qwen_base"), ("geometry", "geometry"), ("dino", "dino"), - ("vae", "vae"), ("moge2", "moge"), ("rmbg2", "rmbg")): - override = getattr(args, flag) - if key == "rmbg2" and override is None and not args.with_rmbg: - continue - if override is not None: - models[key] = local_override(override, file_allowed=key == "moge2") - else: - entry = manifest[key] - destination = f"external/{key}" - operations.append({"name": key, "kind": "hub", "repo_id": entry["repo_id"], - "revision": entry["revision"], "allow_patterns": entry["allow_patterns"], - "gated": entry.get("gated", False), "destination": destination}) - models[key] = {"kind": "prepared", "path": destination} - if args.naf_weights is not None: - models["naf"] = local_override(args.naf_weights, file_allowed=True) - else: - operations.append({"name": "naf", "kind": "url", "url": NAF_URL, "destination": "external/naf/naf_release.pth"}) - models["naf"] = {"kind": "prepared", "path": "external/naf/naf_release.pth"} - return {"schema_version": 1, "output_dir": str(output), "downloads": operations, "models": models, - "bindings": {"geometry": "geometry/pipeline.json + five config/weight pairs + dino", - "pipeline_source": "local geometry pipeline.json or deployment/runtime/geometry/pipeline.json", - "vae": "all four encoder/decoder config/weight pairs under upstream ckpts/", - "naf_source": "torch/hub/valeoai_NAF_main -> deployment/runtime/naf", - "naf_weights": "torch/hub/checkpoints/naf_release.pth", - "config": "models.json (published only after all checks)"}, - "access": "Gated repositories require your existing upstream permission and local HF authentication. HF_ENDPOINT is inherited.", - "local_override_note": "Local overrides remain external dependencies; their absolute paths/links are intentionally machine-local."} - - -def require_file(path: Path) -> Path: - if not path.is_file() or path.stat().st_size == 0: - raise FileNotFoundError(f"required model file is missing or empty: {path}") - return path - - -def check_index(path: Path) -> None: - index = json.loads(require_file(path).read_text()) - mapping = index.get("weight_map") - if not isinstance(mapping, dict) or not mapping: - raise ValueError(f"checkpoint index requires a nonempty weight_map: {path}") - for filename in set(mapping.values()): - relative_name(filename) - if not filename.endswith(".safetensors"): - raise ValueError(f"index shard must be .safetensors: {filename}") - require_file(path.parent / filename) - - -def check_safe_weights(directory: Path, single: str = "model.safetensors") -> Path: - index = directory / (single + ".index.json") - weights = directory / single - if index.is_file() and weights.is_file(): - raise ValueError(f"ambiguous single and sharded weights: {directory}") - if index.is_file(): - check_index(index) - return index - return require_file(weights) - - -def validate_inputs(paths: dict[str, Path]) -> tuple[dict, Path]: - # Load the path-only helper directly: package __init__ imports GPU-facing - # modules, and resolve_completion_artifacts performs deeper ABI validation. - spec = importlib.util.spec_from_file_location("_supratex_prepare_weights", ROOT / "supratex/weights.py") - module = importlib.util.module_from_spec(spec) - spec.loader.exec_module(module) - for key in ("completion", "material"): - module.resolve_model_dir(paths[key], local_files_only=True) - completion = json.loads(require_file(paths["completion"] / "config.json").read_text()) - if completion.get("format") != "supratex_completion" or completion.get("schema_version") != 1: - raise ValueError("completion requires the standalone Supra-Tex config.json (schema_version 1)") - require_file(paths["completion"] / "partial_normalization.json") - checkpoint = check_safe_weights(paths["completion"]) - adapter = json.loads(require_file(paths["material"] / "adapter_config.json").read_text()) - if adapter.get("format") != "supratex_qwen_material_adapter" or adapter.get("weight_file") != "adapter_model.safetensors": - raise ValueError("material requires the standalone Supra-Tex adapter_config.json and adapter_model.safetensors") - require_file(paths["material"] / "adapter_model.safetensors") - for name in ("preprocessor_config.json", "tokenizer_config.json", "tokenizer.json", "chat_template.jinja"): - require_file(paths["qwen-base"] / "processor" / name) - for name in ("vae", "text_encoder", "transformer"): - component = paths["qwen-base"] / name - require_file(component / "config.json") - weights = list(component.glob("*.safetensors")) - if not weights: - raise FileNotFoundError(f"Qwen foundation component has no weights: {component}") - for weight in weights: - require_file(weight) - for index in component.glob("*.safetensors.index.json"): - check_index(index) - for name in ("config.json", "preprocessor_config.json"): - require_file(paths["dino"] / name) - check_safe_weights(paths["dino"]) - if "rmbg2" in paths: - for name in ("config.json", "model.safetensors", "preprocessor_config.json", "BiRefNet_config.py", "birefnet.py"): - require_file(paths["rmbg2"] / name) - paths["moge2"] = require_file(paths["moge2"] / "model.pt" if paths["moge2"].is_dir() else paths["moge2"]) - paths["naf"] = require_file(paths["naf"] / "naf_release.pth" if paths["naf"].is_dir() else paths["naf"]) - vae = paths["vae"] / "ckpts" if (paths["vae"] / "ckpts").is_dir() else paths["vae"] - for stem in VAE_STEMS: - for suffix in (".json", ".safetensors"): - require_file(vae / (stem + suffix)) - pipeline_file = paths["geometry"] / "pipeline.json" - if not pipeline_file.is_file(): - pipeline_file = ROOT / "deployment/runtime/geometry/pipeline.json" - pipeline = json.loads(require_file(pipeline_file).read_text()) - for key in GEOMETRY_KEYS: - stem = relative_name(pipeline["args"]["models"][key]) - for suffix in (".json", ".safetensors"): - require_file(paths["geometry"] / (stem + suffix)) - return pipeline, checkpoint - - -def write_once(path: Path, text: str) -> None: - """Accept identical products on retry, never overwrite an existing file.""" - if path.exists() or path.is_symlink(): - if path.is_symlink() or not path.is_file() or path.read_text() != text: - raise FileExistsError(f"refusing to overwrite conflicting output: {path}") - return - path.parent.mkdir(parents=True, exist_ok=True) - with path.open("x", encoding="utf-8") as stream: - stream.write(text) - - -def json_text(value: dict) -> str: - return json.dumps(value, indent=2) + "\n" - - -def link_once(destination: Path, source: Path, *, relative: bool) -> None: - target = os.path.relpath(source, destination.parent) if relative else str(source) - if destination.is_symlink() and os.readlink(destination) == target: - return - if destination.exists() or destination.is_symlink(): - raise FileExistsError(f"refusing to overwrite conflicting output: {destination}") - destination.parent.mkdir(parents=True, exist_ok=True) - destination.symlink_to(target, target_is_directory=source.is_dir()) - - -def download(operation: dict, output: Path) -> None: - destination = output / operation["destination"] - if operation["kind"] == "url": - if destination.exists(): - require_file(destination) - return - destination.parent.mkdir(parents=True, exist_ok=True) - temporary = destination.with_suffix(".pth.part") - try: - with urlopen(operation["url"], timeout=60) as response, temporary.open("wb") as target: - expected_size = response.headers.get("Content-Length") - shutil.copyfileobj(response, target) - require_file(temporary) - if expected_size is not None and temporary.stat().st_size != int(expected_size): - raise ValueError("NAF response ended before its declared Content-Length") - # A same-directory hard link publishes without replacing a file. - os.link(temporary, destination) - except Exception as error: - raise RuntimeError("NAF download failed; retry or supply --naf-weights with the official naf_release.pth") from error - finally: - temporary.unlink(missing_ok=True) - return - from huggingface_hub import snapshot_download - try: - snapshot_download(repo_id=operation["repo_id"], revision=operation["revision"], - allow_patterns=operation["allow_patterns"], local_dir=str(destination), - cache_dir=str(output / ".cache/huggingface")) - except Exception as error: - status = getattr(getattr(error, "response", None), "status_code", None) - access = operation.get("gated") or status in (401, 403) or "GatedRepo" in type(error).__name__ - detail = ("This upstream repository requires existing access permission and your own local HF authentication. " - "Accept its publisher's terms and authenticate with your normal HF tooling, or provide a local override." - if access else "Check the repository/revision and your network or local HF authentication, then retry.") - # Do not interpolate exception text: upstream URLs can contain secrets. - raise RuntimeError(f"Download failed for {operation['repo_id']}. {detail}") from None - - -def prepare(plan: dict) -> Path: - output = Path(plan["output_dir"]) - marker = output / MARKER - identity = {key: value for key, value in plan.items() if key != "output_dir"} - if output.exists() and not output.is_dir(): - raise NotADirectoryError(output) - if output.is_dir() and any(output.iterdir()) and not marker.is_file(): - raise FileExistsError(f"output directory is not owned by prepare_models.py; choose an empty --output-dir: {output}") - write_once(marker, json_text(identity)) - # Check publisher-controlled access before downloading the large public bases. - for operation in sorted(plan["downloads"], key=lambda item: not item.get("gated", False)): - download(operation, output) - paths = {name: Path(source["path"]) if source["kind"] == "local" else output / source["path"] - for name, source in plan["models"].items()} - pipeline, checkpoint = validate_inputs(paths) - geometry = output / "geometry" - for key in GEOMETRY_KEYS: - stem = pipeline["args"]["models"][key] - write_once(geometry / (stem + ".json"), (paths["geometry"] / (stem + ".json")).read_text()) - link_once(geometry / (stem + ".safetensors"), paths["geometry"] / (stem + ".safetensors"), - relative=plan["models"]["geometry"]["kind"] != "local") - write_once(geometry / "pipeline.json", json_text(pipeline)) - link_once(geometry / "dino", paths["dino"], relative=plan["models"]["dino"]["kind"] != "local") - runtime = ROOT / "deployment/runtime" - require_file(runtime / "naf/hubconf.py") - link_once(output / "torch/hub/valeoai_NAF_main", runtime / "naf", relative=True) - link_once(output / "torch/hub/checkpoints/naf_release.pth", paths["naf"], relative=plan["models"]["naf"]["kind"] != "local") - - def config_path(path: Path, local: bool = False) -> str: - return str(path) if local else os.path.relpath(path, output) - - def model_path(name: str) -> str: - return config_path(paths[name], plan["models"][name]["kind"] == "local") - - local_names = [key for key, value in plan["models"].items() if value["kind"] == "local"] - config = {"moge_model": model_path("moge2"), "rmbg_model": model_path("rmbg2") if "rmbg2" in paths else None, - "geometry_package": config_path(runtime / "geometry"), "geometry_model": "geometry", - "qwen_runtime": config_path(runtime / "qwen"), "qwen_deps": config_path(runtime / "qwen-deps"), - "material_base": model_path("qwen-base"), "material_model": model_path("material"), - "pretrained": model_path("completion"), "torch_home": "torch", "trellis_root": model_path("vae"), - "model_config": config_path(paths["completion"] / "config.json", "completion" in local_names), - "checkpoint": config_path(checkpoint, "completion" in local_names), - "_meta": {"schema_version": 1, "local_overrides": local_names, - "portability": "Relative paths resolve from models.json. Move with the source package; local overrides remain external dependencies.", - "validation": "File presence and shard inventory only; run scripts/check_runtime.py before inference."}} - result = output / "models.json" - write_once(result, json_text(config)) - return result - - -def main(argv: list[str] | None = None) -> int: - args = parse_args(argv) - try: - plan = build_plan(args) - if args.dry_run: - print(json_text(plan), end="") - else: - result = prepare(plan) - print(f"Prepared {result}") - if any(value["kind"] == "local" for value in plan["models"].values()): - print("Local overrides remain external dependencies; keep their original paths available.") - print(f"Next: python scripts/check_runtime.py --config {result}") - except (OSError, ValueError, KeyError, RuntimeError) as error: - print(f"error: {error}", file=sys.stderr) - return 2 - return 0 - - -if __name__ == "__main__": - raise SystemExit(main()) +#!/usr/bin/env python3 +"""Download or bind Supra-Tex models, then publish a portable models.json. + +Downloads use the inherited Hugging Face configuration and your existing access. +Local overrides are linked without copying large weights and remain dependencies +of the prepared tree. Cached Hub snapshots are also linked when using an explicit +cache or preparing offline; keep that cache available. No tensor data is loaded +and no checksums are computed. +""" +from __future__ import annotations + +import argparse +import copy +import filecmp +import fnmatch +import importlib.util +import json +import os +from pathlib import Path +import shutil +import sys +import tempfile +from urllib.request import urlopen + + +ROOT = Path(__file__).resolve().parents[1] +PROJECT_REPO = "supra3d/supra-tex" +MATERIAL_PATTERNS = ["adapter_config.json", "adapter_model.safetensors", "LICENSE*", "Notice", "README.md"] +GEOMETRY_KEYS = ("sparse_structure_flow_model", "sparse_structure_decoder", + "shape_slat_flow_model_512", "shape_slat_flow_model_1024", "shape_slat_decoder") +VAE_STEMS = ("shape_enc_next_dc_f16c32_fp16", "tex_enc_next_dc_f16c32_fp16", + "shape_dec_next_dc_f16c32_fp16", "tex_dec_next_dc_f16c32_fp16") +MARKER = ".prepare-models.json" + +# Keep setup/help usable without importing the GPU-facing supratex package. +_spec = importlib.util.spec_from_file_location("_supratex_pipeline_config", ROOT / "supratex/pipeline_config.py") +_pipeline = importlib.util.module_from_spec(_spec) +_spec.loader.exec_module(_pipeline) +_weights_spec = importlib.util.spec_from_file_location("_supratex_prepare_weights", ROOT / "supratex/weights.py") +_weights = importlib.util.module_from_spec(_weights_spec) +_weights_spec.loader.exec_module(_weights) +DEFAULT_COMPLETION_CHECKPOINT = _weights.DEFAULT_COMPLETION_CHECKPOINT +# The release keeps an older sharded model beside the default checkpoint; fetch +# only the default weights. +COMPLETION_PATTERNS = ["config.json", "partial_normalization.json", DEFAULT_COMPLETION_CHECKPOINT, + "LICENSE*", "README.md"] +relative_name = _pipeline.relative_name +model_source = _pipeline.model_source + + +def parse_args(argv: list[str] | None = None) -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--repo", default=PROJECT_REPO, help="pipeline release directory or HF repo ID/URL (default: %(default)s)") + parser.add_argument("--revision", help="project repository revision; defaults to its current main") + parser.add_argument("--pipeline", type=Path, help="explicit local pipeline.json; otherwise read it from --repo") + parser.add_argument("--completion-subdir", help="override the completion directory declared by pipeline.json") + parser.add_argument("--material-subdir", help="override the material directory declared by pipeline.json") + parser.add_argument("--completion", help="override with standalone completion directory, HF repo ID, or HF URL") + parser.add_argument("--material", help="override with material adapter directory, HF repo ID, or HF URL") + parser.add_argument("--completion-revision", help="revision for a separate --completion repository") + parser.add_argument("--material-revision", help="revision for a separate --material repository") + parser.add_argument("--output-dir", type=Path, default=ROOT / "ckpts/ready") + parser.add_argument("--cache-dir", type=Path, help="Hugging Face cache for manifests and model snapshots; prepared files link to this cache, so keep it available") + for name in ("qwen-base", "geometry", "dino", "vae", "moge", "rmbg", "naf-weights"): + parser.add_argument("--" + name, type=Path, help="existing local override (no download)") + parser.add_argument("--with-rmbg", action="store_true", help="also prepare RMBG for RGB inputs; default example uses RGBA") + parser.add_argument("--local-files-only", action="store_true", help="use local files and cached Hub snapshots; never download") + parser.add_argument("--dry-run", action="store_true", help="offline plan; remote repositories use the bundled pipeline preview") + return parser.parse_args(argv) + + +def local_override(value: Path, *, file_allowed: bool = False) -> dict: + path = value.expanduser().resolve() + if not path.exists(): + raise FileNotFoundError(f"local override does not exist: {path}") + if not path.is_dir() and not (file_allowed and path.is_file()): + raise ValueError(f"local override must be {'a directory or file' if file_allowed else 'a directory'}: {path}") + return {"kind": "local", "path": str(path)} + + +def build_plan(args: argparse.Namespace, *, resolve_remote: bool = False) -> dict: + """Build a plan from one pipeline; offline previews never fetch metadata.""" + output = args.output_dir.expanduser().resolve() + project = model_source(args.repo, args.revision) + project_needed = [name for name in ("completion", "material") if getattr(args, name) is None] + if args.pipeline is not None: + pipeline = _pipeline.load_pipeline(args.pipeline) + manifest_source = "explicit local pipeline" + elif project["kind"] == "local": + pipeline, project = _pipeline.resolve_pipeline(args.repo) + manifest_source = "local release" + elif resolve_remote and (project_needed or args.repo != PROJECT_REPO or args.revision): + pipeline, project = _pipeline.resolve_pipeline( + args.repo, revision=args.revision, cache_dir=args.cache_dir, local_files_only=args.local_files_only) + manifest_source = "Hub pipeline at selected revision" + else: + pipeline = _pipeline.load_pipeline(ROOT / "configs/pipeline.json") + manifest_source = "bundled preview; Hub pipeline is resolved before preparation" if project_needed else "bundled pipeline" + manifest = pipeline["external_models"] + operations, models = [], {} + subdirs = {name: relative_name(getattr(args, name + "_subdir") or pipeline["models"][name]) + for name in ("completion", "material")} + patterns = {"completion": COMPLETION_PATTERNS, "material": MATERIAL_PATTERNS} + if project_needed and project["kind"] == "hub": + operations.append(dict(project, name="project", destination="project", + allow_patterns=["pipeline.json", "README.md", ".gitattributes"] + + [f"{subdirs[name]}/{pattern}" for name in project_needed for pattern in patterns[name]])) + for name in subdirs: + value = getattr(args, name) + if value is None: + if project["kind"] == "local": + models[name] = local_override(Path(project["path"]) / subdirs[name]) + else: + models[name] = {"kind": "prepared", "path": f"project/{subdirs[name]}"} + else: + source = model_source(value, getattr(args, name + "_revision")) + models[name] = source if source["kind"] == "local" else {"kind": "prepared", "path": name} + if source["kind"] == "hub": + operations.append(dict(source, name=name, destination=name, allow_patterns=patterns[name])) + for key, flag in (("qwen-base", "qwen_base"), ("geometry", "geometry"), ("dino", "dino"), + ("vae", "vae"), ("moge2", "moge"), ("rmbg2", "rmbg")): + override = getattr(args, flag) + if key == "rmbg2" and override is None and not args.with_rmbg: + continue + if override is not None: + models[key] = local_override(override, file_allowed=key == "moge2") + else: + entry = manifest[key] + destination = f"external/{key}" + operations.append({"name": key, "kind": "hub", "repo_id": entry["repo_id"], + "revision": entry["revision"], "allow_patterns": entry["allow_patterns"], + "gated": entry.get("gated", False), "destination": destination}) + models[key] = {"kind": "prepared", "path": destination} + if args.naf_weights is not None: + models["naf"] = local_override(args.naf_weights, file_allowed=True) + else: + entry = manifest["naf"] + destination = "external/naf/" + entry["filename"] + operations.append({"name": "naf", "kind": "url", "url": entry["url"], "destination": destination}) + models["naf"] = {"kind": "prepared", "path": destination} + for operation in operations: + if args.local_files_only: + operation["local_files_only"] = True + if args.cache_dir is not None and operation["kind"] == "hub": + operation["cache_dir"] = str(args.cache_dir.expanduser().resolve()) + return {"schema_version": 1, "output_dir": str(output), "downloads": operations, "models": models, + "cache_dir": str(args.cache_dir.expanduser().resolve()) if args.cache_dir is not None else None, + "local_files_only": args.local_files_only, + "pipeline": pipeline, "pipeline_source": manifest_source, + "bindings": {"geometry": "geometry/pipeline.json + five config/weight pairs + dino", + "pipeline_source": "local geometry pipeline.json or deployment/runtime/geometry/pipeline.json", + "vae": "all four encoder/decoder config/weight pairs under upstream ckpts/", + "naf_source": "torch/hub/valeoai_NAF_main -> deployment/runtime/naf", + "naf_weights": "torch/hub/checkpoints/naf_release.pth", + "config": "models.json (generated paths and selected pipeline defaults)"}, + "access": "Gated repositories require your existing upstream permission and local HF authentication. HF_ENDPOINT is inherited.", + "local_override_note": "Local overrides remain external dependencies; their absolute paths/links are intentionally machine-local."} + + +def require_file(path: Path) -> Path: + if not path.is_file() or path.stat().st_size == 0: + raise FileNotFoundError(f"required model file is missing or empty: {path}") + return path + + +def check_index(path: Path) -> None: + index = json.loads(require_file(path).read_text()) + mapping = index.get("weight_map") + if not isinstance(mapping, dict) or not mapping: + raise ValueError(f"checkpoint index requires a nonempty weight_map: {path}") + for filename in set(mapping.values()): + relative_name(filename) + if not filename.endswith(".safetensors"): + raise ValueError(f"index shard must be .safetensors: {filename}") + require_file(path.parent / filename) + + +def check_safe_weights(directory: Path, single: str = "model.safetensors") -> Path: + index = directory / (single + ".index.json") + weights = directory / single + if index.is_file() and weights.is_file(): + raise ValueError(f"ambiguous single and sharded weights: {directory}") + if index.is_file(): + check_index(index) + return index + return require_file(weights) + + +def completion_checkpoint(directory: Path) -> Path: + """Prefer the default completion weights; fall back to model.safetensors or its shards.""" + default = directory / DEFAULT_COMPLETION_CHECKPOINT + return default if default.is_file() else check_safe_weights(directory) + + +def validate_inputs(paths: dict[str, Path]) -> tuple[dict, Path]: + # Load the path-only helper directly: package __init__ imports GPU-facing + # modules, and resolve_completion_artifacts performs deeper ABI validation. + for key in ("completion", "material"): + _weights.resolve_model_dir(paths[key], local_files_only=True) + completion = json.loads(require_file(paths["completion"] / "config.json").read_text()) + if completion.get("format") != "supratex_completion" or completion.get("schema_version") != 1: + raise ValueError("completion requires the standalone Supra-Tex config.json (schema_version 1)") + require_file(paths["completion"] / "partial_normalization.json") + checkpoint = completion_checkpoint(paths["completion"]) + adapter = json.loads(require_file(paths["material"] / "adapter_config.json").read_text()) + if adapter.get("format") != "supratex_qwen_material_adapter" or adapter.get("weight_file") != "adapter_model.safetensors": + raise ValueError("material requires the standalone Supra-Tex adapter_config.json and adapter_model.safetensors") + require_file(paths["material"] / "adapter_model.safetensors") + for name in ("preprocessor_config.json", "tokenizer_config.json", "tokenizer.json", "chat_template.jinja"): + require_file(paths["qwen-base"] / "processor" / name) + for name in ("vae", "text_encoder", "transformer"): + component = paths["qwen-base"] / name + require_file(component / "config.json") + weights = list(component.glob("*.safetensors")) + if not weights: + raise FileNotFoundError(f"Qwen foundation component has no weights: {component}") + for weight in weights: + require_file(weight) + for index in component.glob("*.safetensors.index.json"): + check_index(index) + # The geometry extractors normalize images directly and load only DINO's + # model/config; a Transformers image-processor file is not required. + require_file(paths["dino"] / "config.json") + check_safe_weights(paths["dino"]) + if "rmbg2" in paths: + # The foreground adapter normalizes images itself, so RMBG-2.0 and the + # public BiRefNet release share this layout without an image processor. + for name in ("config.json", "model.safetensors", "BiRefNet_config.py", "birefnet.py"): + require_file(paths["rmbg2"] / name) + paths["moge2"] = require_file(paths["moge2"] / "model.pt" if paths["moge2"].is_dir() else paths["moge2"]) + paths["naf"] = require_file(paths["naf"] / "naf_release.pth" if paths["naf"].is_dir() else paths["naf"]) + vae = paths["vae"] / "ckpts" if (paths["vae"] / "ckpts").is_dir() else paths["vae"] + for stem in VAE_STEMS: + for suffix in (".json", ".safetensors"): + require_file(vae / (stem + suffix)) + pipeline_file = paths["geometry"] / "pipeline.json" + if not pipeline_file.is_file(): + pipeline_file = _pipeline.runtime_root(ROOT) / "geometry/pipeline.json" + pipeline = json.loads(require_file(pipeline_file).read_text()) + for key in GEOMETRY_KEYS: + stem = relative_name(pipeline["args"]["models"][key]) + for suffix in (".json", ".safetensors"): + require_file(paths["geometry"] / (stem + suffix)) + return pipeline, checkpoint + + +def write_once(path: Path, text: str) -> None: + """Accept identical products on retry, never overwrite an existing file.""" + if path.exists() or path.is_symlink(): + if path.is_symlink() or not path.is_file() or path.read_text() != text: + raise FileExistsError(f"refusing to overwrite conflicting output: {path}") + return + path.parent.mkdir(parents=True, exist_ok=True) + with path.open("x", encoding="utf-8") as stream: + stream.write(text) + + +def json_text(value: dict) -> str: + return json.dumps(value, indent=2) + "\n" + + +def link_once(destination: Path, source: Path, *, relative: bool) -> None: + target = os.path.relpath(source, destination.parent) if relative else str(source) + if destination.is_symlink() and os.readlink(destination) == target: + return + if destination.exists() or destination.is_symlink(): + raise FileExistsError(f"refusing to overwrite conflicting output: {destination}") + destination.parent.mkdir(parents=True, exist_ok=True) + destination.symlink_to(target, target_is_directory=source.is_dir()) + + +def bind_cached_snapshot(destination: Path, cached: Path, patterns: list[str]) -> None: + """Fill missing cached files only after checking every existing destination.""" + if not destination.exists() and not destination.is_symlink(): + link_once(destination, cached, relative=False) + return + if destination.is_symlink(): + if destination.resolve() == cached.resolve(): + return + if not destination.is_dir(): + raise FileExistsError(f"refusing to overwrite conflicting output: {destination}") + missing = [] + for source in sorted(cached.rglob("*")): + relative = source.relative_to(cached) + if not source.is_file() or not any(fnmatch.fnmatch(relative.as_posix(), pattern) for pattern in patterns): + continue + target = destination / relative + parent = target.parent + while parent != destination: + if parent.is_symlink() or (parent.exists() and not parent.is_dir()): + raise FileExistsError(f"refusing to modify conflicting output: {parent}") + parent = parent.parent + if target.exists() or target.is_symlink(): + if not target.is_file() or (not source.samefile(target) and not filecmp.cmp(source, target, shallow=False)): + raise FileExistsError(f"refusing to overwrite conflicting output: {target}") + else: + if destination.is_symlink(): + raise FileExistsError(f"refusing to modify a conflicting cache binding: {destination}") + missing.append((target, source)) + for target, source in missing: + link_once(target, source, relative=False) + + +def download(operation: dict, output: Path) -> None: + destination = output / operation["destination"] + if operation["kind"] == "url": + if destination.exists(): + require_file(destination) + return + if operation.get("local_files_only"): + raise FileNotFoundError("NAF weights are unavailable offline; supply --naf-weights or prepare them online first") + destination.parent.mkdir(parents=True, exist_ok=True) + temporary = destination.with_suffix(".pth.part") + try: + with urlopen(operation["url"], timeout=60) as response, temporary.open("wb") as target: + expected_size = response.headers.get("Content-Length") + shutil.copyfileobj(response, target) + require_file(temporary) + if expected_size is not None and temporary.stat().st_size != int(expected_size): + raise ValueError("NAF response ended before its declared Content-Length") + # A same-directory hard link publishes without replacing a file. + os.link(temporary, destination) + except Exception as error: + raise RuntimeError("NAF download failed; retry or supply --naf-weights with the official naf_release.pth") from error + finally: + temporary.unlink(missing_ok=True) + return + from huggingface_hub import snapshot_download + try: + if operation.get("local_files_only"): + try: + cached = snapshot_download(repo_id=operation["repo_id"], revision=operation["revision"], + allow_patterns=operation["allow_patterns"], + cache_dir=operation.get("cache_dir"), local_files_only=True) + except Exception: + if destination.is_dir(): + # Online local_dir downloads may have no shared cache entry. + # The caller validates all required files before publishing. + return + raise + elif operation.get("cache_dir") is not None: + # local_dir bypasses the shared cache in huggingface_hub, so an + # explicit cache must use a snapshot and bind its files afterwards. + cached = snapshot_download(repo_id=operation["repo_id"], revision=operation["revision"], + allow_patterns=operation["allow_patterns"], + cache_dir=operation["cache_dir"]) + else: + snapshot_download(repo_id=operation["repo_id"], revision=operation["revision"], + allow_patterns=operation["allow_patterns"], local_dir=str(destination), + cache_dir=operation.get("cache_dir")) + except Exception as error: + status = getattr(getattr(error, "response", None), "status_code", None) + access = operation.get("gated") or status in (401, 403) or "GatedRepo" in type(error).__name__ + detail = ("This upstream repository requires existing access permission and your own local HF authentication. " + "Accept its publisher's terms and authenticate with your normal HF tooling, or provide a local override." + if access else "Check the repository/revision and your network or local HF authentication, then retry.") + # Do not interpolate exception text: upstream URLs can contain secrets. + raise RuntimeError(f"Download failed for {operation['repo_id']}. {detail}") from None + if operation.get("local_files_only") or operation.get("cache_dir") is not None: + # Keep local conflict errors distinct from remote access failures. + bind_cached_snapshot(destination, Path(cached), operation["allow_patterns"]) + + +def plan_identity(plan: dict) -> dict: + """Transport/cache choices are not part of the prepared model identity.""" + identity = {key: copy.deepcopy(value) for key, value in plan.items() + if key not in {"output_dir", "pipeline_source", "cache_dir", "local_files_only"}} + identity["downloads"] = [{key: value for key, value in item.items() + if key not in {"local_files_only", "cache_dir"}} + for item in plan["downloads"]] + return identity + + +def without_rmbg(identity: dict) -> dict: + identity = copy.deepcopy(identity) + identity["models"].pop("rmbg2", None) + identity["downloads"] = [item for item in identity["downloads"] if item["name"] != "rmbg2"] + return identity + + +def replace_owned_json(path: Path, value: dict, previous: dict, *, check_only: bool = False) -> None: + """Atomically extend a generated file, rejecting user edits and symlinks.""" + text = json_text(value) + if not path.exists() and not path.is_symlink(): + if not check_only: + write_once(path, text) + return + if path.is_symlink() or not path.is_file() or path.read_text() not in {json_text(previous), text}: + raise FileExistsError(f"refusing to overwrite conflicting output: {path}") + if check_only or path.read_text() == text: + return + with tempfile.NamedTemporaryFile(mode="w", dir=path.parent, prefix="." + path.name + ".", delete=False) as stream: + temporary = Path(stream.name) + try: + stream.write(text) + stream.close() + os.replace(temporary, path) + finally: + temporary.unlink(missing_ok=True) + + +def prepare(plan: dict) -> Path: + plan = copy.deepcopy(plan) + output = Path(plan["output_dir"]) + marker = output / MARKER + if output.exists() and not output.is_dir(): + raise NotADirectoryError(output) + if output.is_dir() and any(output.iterdir()) and not marker.is_file(): + raise FileExistsError(f"output directory is not owned by prepare_models.py; choose an empty --output-dir: {output}") + previous = None + identity = plan_identity(plan) + if marker.exists(): + if marker.is_symlink(): + raise FileExistsError(f"refusing to overwrite conflicting output: {marker}") + previous = json.loads(marker.read_text()) + if without_rmbg(previous) != without_rmbg(identity): + raise FileExistsError(f"refusing to overwrite conflicting output: {marker}") + if "rmbg2" in previous["models"]: + old_downloads = [item for item in previous["downloads"] if item["name"] == "rmbg2"] + if "rmbg2" in identity["models"]: + new_downloads = [item for item in identity["downloads"] if item["name"] == "rmbg2"] + if identity["models"]["rmbg2"] != previous["models"]["rmbg2"] or old_downloads != new_downloads: + raise FileExistsError(f"refusing to overwrite conflicting output: {marker}") + else: + # Keep an existing optional dependency when switching to RGBA. + plan["models"] = {**previous["models"], **plan["models"]} + for operation in old_downloads: + operation = copy.deepcopy(operation) + if plan.get("local_files_only") or any(item.get("local_files_only") for item in plan["downloads"]): + operation["local_files_only"] = True + if plan.get("cache_dir") is not None: + operation["cache_dir"] = plan["cache_dir"] + position = next(index for index, item in enumerate(previous["downloads"]) if item["name"] == "rmbg2") + plan["downloads"].insert(position, operation) + identity = plan_identity(plan) + else: + write_once(marker, json_text(identity)) + paths = {name: Path(source["path"]) if source["kind"] == "local" else output / source["path"] + for name, source in plan["models"].items()} + existing_inputs_complete = False + if any(operation.get("local_files_only") for operation in plan["downloads"]): + try: + validate_inputs(dict(paths)) + existing_inputs_complete = True + except (OSError, ValueError, KeyError): + pass + # Check publisher-controlled access before downloading the large public bases. + for operation in sorted(plan["downloads"], key=lambda item: not item.get("gated", False)): + if operation.get("local_files_only") and existing_inputs_complete: + continue + download(operation, output) + pipeline, checkpoint = validate_inputs(paths) + geometry = output / "geometry" + for key in GEOMETRY_KEYS: + stem = pipeline["args"]["models"][key] + write_once(geometry / (stem + ".json"), (paths["geometry"] / (stem + ".json")).read_text()) + link_once(geometry / (stem + ".safetensors"), paths["geometry"] / (stem + ".safetensors"), + relative=plan["models"]["geometry"]["kind"] != "local") + write_once(geometry / "pipeline.json", json_text(pipeline)) + link_once(geometry / "dino", paths["dino"], relative=plan["models"]["dino"]["kind"] != "local") + runtime = _pipeline.runtime_root(ROOT) + require_file(runtime / "naf/hubconf.py") + link_once(output / "torch/hub/valeoai_NAF_main", runtime / "naf", relative=True) + link_once(output / "torch/hub/checkpoints/naf_release.pth", paths["naf"], relative=plan["models"]["naf"]["kind"] != "local") + + def config_path(path: Path, local: bool = False) -> str: + return str(path) if local else os.path.relpath(path, output) + + def model_path(name: str) -> str: + return config_path(paths[name], plan["models"][name]["kind"] == "local") + + local_names = [key for key, value in plan["models"].items() if value["kind"] == "local"] + config = {"moge_model": model_path("moge2"), "rmbg_model": model_path("rmbg2") if "rmbg2" in paths else None, + "geometry_package": config_path(runtime / "geometry"), "geometry_model": "geometry", + "qwen_runtime": config_path(runtime / "qwen"), + "material_base": model_path("qwen-base"), "material_model": model_path("material"), + "pretrained": model_path("completion"), "torch_home": "torch", "trellis_root": model_path("vae"), + "model_config": config_path(paths["completion"] / "config.json", "completion" in local_names), + "checkpoint": config_path(checkpoint, "completion" in local_names), + "defaults": plan["pipeline"]["defaults"], + "_meta": {"schema_version": 1, "local_overrides": local_names, + "portability": "Relative paths resolve from models.json. Move with the source package; local overrides remain external dependencies.", + "validation": "File presence and shard inventory only; run scripts/check_runtime.py before inference."}} + result = output / "models.json" + if "rmbg2" in identity["models"]: + old_config = copy.deepcopy(config) + old_config["rmbg_model"] = None + old_config["_meta"]["local_overrides"] = [name for name in old_config["_meta"]["local_overrides"] if name != "rmbg2"] + # Check for user edits before recording the extension. Recording the + # dependency first lets either input type recover an interrupted publish. + replace_owned_json(result, config, old_config, check_only=True) + if previous is not None and previous != identity: + replace_owned_json(marker, identity, previous) + replace_owned_json(result, config, old_config) + else: + write_once(result, json_text(config)) + return result + + +def prepare_pipeline(source: str | Path, *, output_dir: str | Path | None = None, + revision: str | None = None, cache_dir: str | Path | None = None, + local_files_only: bool = False, + with_rmbg: bool = False, overrides: dict | None = None) -> Path: + """Resolve a release and generate the runtime bindings used by inference.""" + source = Path(source).expanduser().absolute() if isinstance(source, Path) else source + # A manifest filename denotes its release root as well as the selected config. + local = Path(source).expanduser() + manifest_file = local.is_file() or (not str(source).startswith(("https://", "hf://")) and local.suffix == ".json") + argv = ["--repo", str(local.absolute().parent if manifest_file else source)] + if manifest_file: + argv.extend(("--pipeline", str(local))) + if output_dir is not None: + argv.extend(("--output-dir", str(output_dir))) + if revision is not None: + argv.extend(("--revision", revision)) + if cache_dir is not None: + argv.extend(("--cache-dir", str(cache_dir))) + if local_files_only: + argv.append("--local-files-only") + if with_rmbg: + argv.append("--with-rmbg") + flags = {"completion": "completion", "material": "material", "qwen-base": "qwen-base", + "geometry": "geometry", "dino": "dino", "vae": "vae", "moge2": "moge", + "rmbg2": "rmbg", "naf": "naf-weights", + "completion_revision": "completion-revision", "material_revision": "material-revision"} + for key, value in (overrides or {}).items(): + if key not in flags: + raise ValueError(f"unknown pipeline model override: {key}") + argv.extend(("--" + flags[key], str(value))) + return prepare(build_plan(parse_args(argv), resolve_remote=True)) + + +def main(argv: list[str] | None = None) -> int: + args = parse_args(argv) + try: + plan = build_plan(args, resolve_remote=not args.dry_run) + if args.dry_run: + print(json_text(plan), end="") + else: + result = prepare(plan) + print(f"Prepared {result}") + if any(value["kind"] == "local" for value in plan["models"].values()): + print("Local overrides remain external dependencies; keep their original paths available.") + print(f"Next: python scripts/check_runtime.py --config {result}") + except (OSError, ValueError, KeyError, RuntimeError) as error: + print(f"error: {error}", file=sys.stderr) + return 2 + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/scripts/run_example.py b/scripts/run_example.py index c50ae715417eeca7ae6a3be77e817a6ea8d211df..a7a936d480d5e61bf9a0f627cdb3cd39b701c02b 100644 --- a/scripts/run_example.py +++ b/scripts/run_example.py @@ -1,116 +1,156 @@ -#!/usr/bin/env python3 -"""Run the bundled RGBA example with prepared models and an isolated source path.""" -from __future__ import annotations - -import argparse -import json -import os -from pathlib import Path -import shlex -import subprocess -import sys - -ROOT = Path(__file__).resolve().parents[1] -MODEL_KEYS = ("moge_model", "rmbg_model", "geometry_package", "geometry_model", "qwen_runtime", - "qwen_deps", "material_base", "material_model", "pretrained", "model_config", "checkpoint", "trellis_root") -PATH_KEYS = (*MODEL_KEYS, "torch_home") - - -def read_config(path: Path) -> dict: - """Resolve paths relative to models.json, preserving metadata for preflight.""" - path = Path(path).expanduser().resolve() - config = json.loads(path.read_text()) - if not isinstance(config, dict): - raise ValueError("models.json must contain an object") - for key in PATH_KEYS: - value = config.get(key) - if value is None and key in ("rmbg_model", "model_config", "checkpoint"): - continue - if not isinstance(value, str) or not value: - raise ValueError(f"models.json requires a nonempty {key} path") - resolved = Path(value).expanduser() - config[key] = str(resolved if resolved.is_absolute() else (path.parent / resolved).resolve()) - return config - - -def parse_args(argv: list[str] | None = None) -> tuple[argparse.Namespace, list[str]]: - parser = argparse.ArgumentParser(description=__doc__) - parser.add_argument("--config", type=Path, default=ROOT / "ckpts/ready/models.json") - parser.add_argument("--input", type=Path, help="default: bundled assets/pitcher_plant.png") - parser.add_argument("--output", type=Path, required=True, help="destination .glb") - parser.add_argument("--use-input-alpha", action="store_true", help="explicitly require the input foreground alpha") - parser.add_argument("--attention", choices=("auto", "flash_attn_3", "flash_attn", "sdpa", "xformers"), default="auto") - parser.add_argument("--low-vram", action=argparse.BooleanOptionalAction, default=None) - parser.add_argument("--print-command", action="store_true", help="print without creating files or launching inference") - args, extra = parser.parse_known_args(argv) - if extra[:1] == ["--"]: - extra = extra[1:] - protected = {"--" + key.replace("_", "-") for key in MODEL_KEYS} - protected.update(("--material-backend", "--local-files-only", "--cache-dir", "--revision", "--material-revision", "--geometry-revision")) - if any(value.split("=", 1)[0] in protected for value in extra): - parser.error("model paths and offline settings come from --config; additional arguments cannot override them") - return args, extra - - -def has_foreground_alpha(path: Path) -> bool: - from PIL import Image - with Image.open(path) as image: - # The inference foreground contract accepts only an actual RGBA - # input; converting LA or palette transparency here would hide a - # mode that the inference process itself rejects. - if image.mode != "RGBA": - return False - low, high = image.getchannel("A").getextrema() - return low < 255 and high > 0 - - -def build_command(args: argparse.Namespace, extra: list[str], config: dict) -> tuple[list[str], dict[str, str]]: - image = (args.input or ROOT / "assets/pitcher_plant.png").expanduser().resolve() - if not image.is_file(): - raise FileNotFoundError(f"input image is missing: {image}") - use_alpha = args.use_input_alpha or args.input is None or has_foreground_alpha(image) - if not use_alpha and not config.get("rmbg_model"): - raise ValueError("RGB or opaque input requires RMBG: prepare with --with-rmbg or --rmbg, then use that models.json") - command = [sys.executable, str(ROOT / "infer.py"), "--input", str(image), "--output", - str(args.output.expanduser().resolve()), "--material-backend", "qwen-image", "--local-files-only", - "--attention", args.attention] - if use_alpha: - command.append("--use-input-alpha") - if args.low_vram is not None: - command.append("--low-vram" if args.low_vram else "--no-low-vram") - for key in MODEL_KEYS: - if config.get(key): - command.extend(("--" + key.replace("_", "-"), config[key])) - command.extend(extra) - environment = os.environ.copy() - environment.update(TORCH_HOME=config["torch_home"], HF_HUB_OFFLINE="1", TRANSFORMERS_OFFLINE="1", - HF_DATASETS_OFFLINE="1", PYTHONNOUSERSITE="1", - PYTHONPATH=os.pathsep.join(str(path) for path in - (ROOT, ROOT / "deployment/runtime/moge", ROOT / "deployment/runtime/utils3d"))) - if args.attention == "auto": - environment.pop("ATTN_BACKEND", None) - environment.pop("SPARSE_ATTN_BACKEND", None) - else: - environment.update(ATTN_BACKEND=args.attention, SPARSE_ATTN_BACKEND=args.attention) - environment.setdefault("OMP_NUM_THREADS", "4") - environment.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True") - return command, environment - - -def main(argv: list[str] | None = None) -> int: - args, extra = parse_args(argv) - try: - command, environment = build_command(args, extra, read_config(args.config)) - if args.print_command: - print(shlex.join(command)) - return 0 - args.output.expanduser().resolve().parent.mkdir(parents=True, exist_ok=True) - subprocess.run(command, cwd=ROOT, env=environment, check=True) - except (OSError, ValueError, subprocess.CalledProcessError) as error: - print(f"error: {error}", file=sys.stderr) - return 2 - return 0 - - -if __name__ == "__main__": - raise SystemExit(main()) +#!/usr/bin/env python3 +"""Run the bundled RGBA example with prepared models and an isolated source path.""" +from __future__ import annotations + +import argparse +import importlib.util +import json +import os +from pathlib import Path +import shlex +import subprocess +import sys + +ROOT = Path(__file__).resolve().parents[1] +MODEL_KEYS = ("moge_model", "rmbg_model", "geometry_package", "geometry_model", "qwen_runtime", + "qwen_deps", "material_base", "material_model", "pretrained", "model_config", "checkpoint", "trellis_root") +PATH_KEYS = (*MODEL_KEYS, "torch_home") + + +def _pipeline_config(): + """Load the lightweight shared helpers without importing GPU model modules.""" + spec = importlib.util.spec_from_file_location("supratex_pipeline_config", ROOT / "supratex/pipeline_config.py") + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +def inference_defaults(config: dict | None = None) -> dict: + """Read defaults from the pipeline without importing GPU model modules.""" + module = _pipeline_config() + defaults = module.default_pipeline()["defaults"] + configured = (config or {}).get("defaults", {}) + if not isinstance(configured, dict): + raise ValueError("models.json defaults must contain an object") + defaults.update(configured) + return module.validate_defaults(defaults) + + +def read_config(path: Path) -> dict: + """Resolve paths relative to models.json, preserving metadata for preflight.""" + path = Path(path).expanduser().resolve() + config = json.loads(path.read_text()) + if not isinstance(config, dict): + raise ValueError("models.json must contain an object") + for key in PATH_KEYS: + value = config.get(key) + if value is None and key in ("rmbg_model", "qwen_deps", "model_config", "checkpoint"): + continue + if not isinstance(value, str) or not value: + raise ValueError(f"models.json requires a nonempty {key} path") + resolved = Path(value).expanduser() + # Normalize without following links: cache-backed files link to + # opaque blobs, and weight validation needs the logical file name. + config[key] = os.path.normpath(resolved if resolved.is_absolute() else path.parent / resolved) + return config + + +def parse_args(argv: list[str] | None = None) -> tuple[argparse.Namespace, list[str]]: + parser = argparse.ArgumentParser(description=__doc__, allow_abbrev=False) + parser.add_argument("--config", type=Path, default=ROOT / "ckpts/ready/models.json") + parser.add_argument("--input", type=Path, help="default: assets/spiny_whelk_shell.png, or " + "assets/cinema_signage.png with --dual-view") + parser.add_argument("--output", type=Path, required=True, help="destination .glb") + parser.add_argument("--use-input-alpha", action="store_true", help="explicitly require the input foreground alpha") + parser.add_argument("--attention", choices=("auto", "flash_attn_3", "flash_attn", "sdpa", "xformers"), + help="override the prepared pipeline's attention default") + parser.add_argument("--low-vram", action=argparse.BooleanOptionalAction, default=None) + parser.add_argument("--print-command", action="store_true", help="print without creating files or launching inference") + args, extra = parser.parse_known_args(argv) + if extra[:1] == ["--"]: + extra = extra[1:] + protected = {"--" + key.replace("_", "-") for key in MODEL_KEYS} + protected.update(("--material-backend", "--local-files-only", "--cache-dir", "--revision", "--material-revision", "--geometry-revision", + "--pipeline", "--pipeline-revision", "--models-dir")) + if any(value.split("=", 1)[0] in protected for value in extra): + parser.error("model paths and offline settings come from --config; additional arguments cannot override them") + return args, extra + + +def has_foreground_alpha(path: Path) -> bool: + from PIL import Image + with Image.open(path) as image: + # The inference foreground contract accepts only an actual RGBA + # input; converting LA or palette transparency here would hide a + # mode that the inference process itself rejects. + if image.mode != "RGBA": + return False + low, high = image.getchannel("A").getextrema() + return low < 255 and high > 0 + + +def default_input(dual_view: bool) -> Path: + """Bundled example: the whelk shell, or the cinema signage for dual-view.""" + return ROOT / "assets" / ("cinema_signage.png" if dual_view else "spiny_whelk_shell.png") + + +def build_command(args: argparse.Namespace, extra: list[str], config: dict) -> tuple[list[str], dict[str, str]]: + defaults = inference_defaults(config) + attention = args.attention or defaults["attention"] + image = (args.input or default_input("--dual-view" in extra)).expanduser().resolve() + if not image.is_file(): + raise FileNotFoundError(f"input image is missing: {image}") + use_alpha = args.use_input_alpha or has_foreground_alpha(image) + if not use_alpha and not config.get("rmbg_model"): + raise ValueError("RGB or opaque input requires RMBG: prepare with --with-rmbg or --rmbg, then use that models.json") + command = [sys.executable, str(ROOT / "infer.py"), "--input", str(image), "--output", + str(args.output.expanduser().resolve()), "--material-backend", "qwen-image", "--local-files-only", + "--attention", attention] + for key in ("seed", "texture_size", "decimation_target", "texture_format"): + command.extend(("--" + key.replace("_", "-"), str(defaults[key]))) + if use_alpha: + command.append("--use-input-alpha") + if args.low_vram is not None: + command.append("--low-vram" if args.low_vram else "--no-low-vram") + for key in MODEL_KEYS: + if config.get(key): + command.extend(("--" + key.replace("_", "-"), config[key])) + command.extend(extra) + return command, runtime_environment(config, attention) + + +def runtime_environment(config: dict, attention: str) -> dict[str, str]: + """Use prepared caches and bundled sources consistently in both entrypoints.""" + runtime = _pipeline_config().runtime_root(ROOT) + environment = os.environ.copy() + environment.update(TORCH_HOME=config["torch_home"], HF_HUB_OFFLINE="1", TRANSFORMERS_OFFLINE="1", + HF_DATASETS_OFFLINE="1", PYTHONNOUSERSITE="1", + PYTHONPATH=os.pathsep.join(str(path) for path in + (ROOT, runtime / "moge", runtime / "utils3d"))) + if attention == "auto": + environment.pop("ATTN_BACKEND", None) + environment.pop("SPARSE_ATTN_BACKEND", None) + else: + environment.update(ATTN_BACKEND=attention, SPARSE_ATTN_BACKEND=attention) + environment.setdefault("OMP_NUM_THREADS", "4") + environment.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True") + return environment + + +def main(argv: list[str] | None = None) -> int: + args, extra = parse_args(argv) + try: + command, environment = build_command(args, extra, read_config(args.config)) + if args.print_command: + print(shlex.join(command)) + return 0 + args.output.expanduser().resolve().parent.mkdir(parents=True, exist_ok=True) + subprocess.run(command, cwd=ROOT, env=environment, check=True) + except (OSError, ValueError, subprocess.CalledProcessError) as error: + print(f"error: {error}", file=sys.stderr) + return 2 + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/supratex/frontend/camera.py b/supratex/frontend/camera.py index 5ca8a95dbf6be49b3d52268a7e08e6944672f225..832674b93cc64526bc801c63c67fbe750a5da03d 100644 --- a/supratex/frontend/camera.py +++ b/supratex/frontend/camera.py @@ -1,79 +1,76 @@ -"""MoGe-2 camera adapter for the canonical Supra-Tex camera JSON contract.""" -from __future__ import annotations - -import argparse -import importlib -import json -import math -from pathlib import Path - -import numpy as np -from PIL import Image - - -def camera_from_intrinsics(intrinsics: np.ndarray, width: int, height: int, *, resolution_level: int) -> dict: - """Translate normalized MoGe intrinsics to the texture-stage camera fields.""" - intrinsics = np.asarray(intrinsics, dtype=np.float64) - if intrinsics.shape != (3, 3): - raise ValueError("MoGe intrinsics must be 3x3") - fx = float(intrinsics[0, 0]) * width - if not math.isfinite(fx) or fx <= 0.0: - raise ValueError("MoGe returned an invalid focal length") - camera_angle_x = 2.0 * math.atan(width / (2.0 * fx)) - focal_pixels = 16.0 / math.tan(camera_angle_x / 2.0) * 512.0 / 32.0 - distance = focal_pixels * -0.5 / -256.0 - if not math.isfinite(distance) or distance <= 0.0: - raise ValueError("MoGe-derived camera distance is invalid") - return {"camera_angle_x": float(camera_angle_x), "distance": float(distance), "mesh_scale": 1.0, - "image_width": int(width), "image_height": int(height), "camera_estimator": "MoGe-2", - "moge_resolution_level": int(resolution_level)} - - -def estimate(image: Path, model_path: Path, module_name: str = "moge.model.v2", resolution_level: int = 9, - device: str = "cuda", *, model=None) -> dict: - if not 0 <= resolution_level <= 9: - raise ValueError("resolution_level must be in [0, 9]") - module = importlib.import_module(module_name) - model_cls = getattr(module, "MoGeModel", None) - if model_cls is None: - raise RuntimeError(f"{module_name} does not expose MoGeModel") - import torch - checkpoint = model_path / "model.pt" if model_path.is_dir() else model_path - if not checkpoint.is_file(): - raise FileNotFoundError(checkpoint) - resident = model is not None - if model is None: - model = model_cls.from_pretrained(str(checkpoint)).to(device).eval() - tensor = result = None - try: - rgb = Image.open(image).convert("RGB") - width, height = rgb.size - tensor = torch.from_numpy(np.asarray(rgb, dtype=np.float32) / 255.0).permute(2, 0, 1).to(device) - with torch.inference_mode(): - result = model.infer(tensor, resolution_level=resolution_level, force_projection=True, apply_mask=True) - return camera_from_intrinsics(result["intrinsics"].squeeze().float().cpu().numpy(), width, height, - resolution_level=resolution_level) - finally: - if not resident: - model.cpu() - del tensor, result, model - if str(device).startswith("cuda"): - torch.cuda.empty_cache() - - -def main() -> None: - parser = argparse.ArgumentParser(description="Estimate canonical Supra-Tex camera parameters with MoGe-2.") - parser.add_argument("--image", type=Path, required=True) - parser.add_argument("--model", type=Path, required=True) - parser.add_argument("--module", default="moge.model.v2") - parser.add_argument("--resolution-level", type=int, default=9, - help="MoGe-2 detail level from 0 (fastest) to 9 (highest).") - parser.add_argument("--device", default="cuda") - parser.add_argument("--output", type=Path, required=True) - args = parser.parse_args() - args.output.parent.mkdir(parents=True, exist_ok=True) - args.output.write_text(json.dumps(estimate(args.image, args.model, args.module, args.resolution_level, args.device), indent=2) + "\n") - - -if __name__ == "__main__": - main() +"""MoGe-2 camera adapter for the canonical Supra-Tex camera JSON contract.""" +from __future__ import annotations + +import argparse +import importlib +import json +import math +from pathlib import Path + +import numpy as np +from PIL import Image + + +def camera_from_intrinsics(intrinsics: np.ndarray, width: int, height: int, *, resolution_level: int) -> dict: + """Translate normalized MoGe intrinsics to the texture-stage camera fields.""" + intrinsics = np.asarray(intrinsics, dtype=np.float64) + if intrinsics.shape != (3, 3): + raise ValueError("MoGe intrinsics must be 3x3") + fx = float(intrinsics[0, 0]) * width + if not math.isfinite(fx) or fx <= 0.0: + raise ValueError("MoGe returned an invalid focal length") + camera_angle_x = 2.0 * math.atan(width / (2.0 * fx)) + focal_pixels = 16.0 / math.tan(camera_angle_x / 2.0) * 512.0 / 32.0 + distance = focal_pixels * -0.5 / -256.0 + if not math.isfinite(distance) or distance <= 0.0: + raise ValueError("MoGe-derived camera distance is invalid") + return {"camera_angle_x": float(camera_angle_x), "distance": float(distance), "mesh_scale": 1.0, + "image_width": int(width), "image_height": int(height), "camera_estimator": "MoGe-2", + "moge_resolution_level": int(resolution_level)} + + +def estimate(image: Path, model_path: Path, module_name: str = "moge.model.v2", resolution_level: int = 9, + device: str = "cuda") -> dict: + if not 0 <= resolution_level <= 9: + raise ValueError("resolution_level must be in [0, 9]") + module = importlib.import_module(module_name) + model_cls = getattr(module, "MoGeModel", None) + if model_cls is None: + raise RuntimeError(f"{module_name} does not expose MoGeModel") + import torch + checkpoint = model_path / "model.pt" if model_path.is_dir() else model_path + if not checkpoint.is_file(): + raise FileNotFoundError(checkpoint) + model = model_cls.from_pretrained(str(checkpoint)).to(device).eval() + tensor = result = None + try: + rgb = Image.open(image).convert("RGB") + width, height = rgb.size + tensor = torch.from_numpy(np.asarray(rgb, dtype=np.float32) / 255.0).permute(2, 0, 1).to(device) + with torch.inference_mode(): + result = model.infer(tensor, resolution_level=resolution_level, force_projection=True, apply_mask=True) + return camera_from_intrinsics(result["intrinsics"].squeeze().float().cpu().numpy(), width, height, + resolution_level=resolution_level) + finally: + model.cpu() + del tensor, result, model + if str(device).startswith("cuda"): + torch.cuda.empty_cache() + + +def main() -> None: + parser = argparse.ArgumentParser(description="Estimate canonical Supra-Tex camera parameters with MoGe-2.") + parser.add_argument("--image", type=Path, required=True) + parser.add_argument("--model", type=Path, required=True) + parser.add_argument("--module", default="moge.model.v2") + parser.add_argument("--resolution-level", type=int, default=9, + help="MoGe-2 detail level from 0 (fastest) to 9 (highest).") + parser.add_argument("--device", default="cuda") + parser.add_argument("--output", type=Path, required=True) + args = parser.parse_args() + args.output.parent.mkdir(parents=True, exist_ok=True) + args.output.write_text(json.dumps(estimate(args.image, args.model, args.module, args.resolution_level, args.device), indent=2) + "\n") + + +if __name__ == "__main__": + main() diff --git a/supratex/frontend/firered.py b/supratex/frontend/firered.py index b20ad6d7c4faf2ed88615fd76bbe4c53b501e498..7b6cd2fdf5f786b7fe2d63526335b887c0eaa76b 100644 --- a/supratex/frontend/firered.py +++ b/supratex/frontend/firered.py @@ -1,126 +1,135 @@ -"""FireRed Image-to-Material Translation adapter. - -The supplied package remains the sole FireRed implementation: this adapter -creates its documented JSONL item and runs only its public ``infer.py``. -""" -from __future__ import annotations - -import argparse -import json -import shutil -import subprocess -import sys -from pathlib import Path - - -def item_for(condition: Path, key: str) -> dict[str, object]: - """Build the public package record without inventing material supervision.""" - condition = condition.resolve() - return { - "key": key, "dataset": "user", "sha": key, "asset_id": key, "view": 0, "cond": str(condition), - # The public package validates these schema slots but prediction only - # consumes ``cond``. Point them at the supplied condition so there is - # no fabricated or private GT material file in an inference request. - "albedo_gt": str(condition), "mra_gt": str(condition), "normal_gt": str(condition), "normal_valid": str(condition), - "instr_albedo": "Convert this rendered image into its albedo (base color) map: remove all lighting, shading, shadows and highlights, keep only the intrinsic surface colors on a black background.", - "instr_mra": "Convert this rendered image into its material property map: encode metallic in the red channel, roughness in the green channel, and material alpha in the blue channel, on a black background.", - "instr_normal": "Convert this rendered image into its camera-space surface normal map: encode the normal direction at every pixel as RGB with x pointing right, y pointing up and z toward the camera, on a flat gray background.", "split": "eval", - } - - -def write_manifest(condition: Path, path: Path, *, key: str = "input") -> Path: - if not condition.is_file(): - raise FileNotFoundError(condition) - path.parent.mkdir(parents=True, exist_ok=True) - path.write_text(json.dumps(item_for(condition, key), sort_keys=True) + "\n", encoding="utf-8") - return path - - -def run(package: Path, base_model: str | Path, condition: Path, output: Path, manifest: Path | None = None, - *, key: str = "input", weights: Path | None = None, spec: Path | None = None, - checkpoint_manifest: Path | None = None, material_model: str | Path | None = None, - revision: str | None = None, cache_dir: Path | None = None, - local_files_only: bool = False) -> tuple[Path, Path]: - if not key or Path(key).name != key or key in {".", ".."} or "\\" in key: - raise ValueError("key must be a single safe filename component") - from supratex.weights import resolve_model_dir - if material_model is not None: - if any(value is not None for value in (weights, spec, checkpoint_manifest)): - raise ValueError("--material-model cannot be combined with explicit material artifact paths") - root = resolve_model_dir(material_model, revision=revision, cache_dir=cache_dir, - local_files_only=local_files_only, - allow_patterns=["full_transformer.safetensors", "joint_contract.json", "checkpoint_complete.json"]) - weights, spec, checkpoint_manifest = (root / name for name in - ("full_transformer.safetensors", "joint_contract.json", "checkpoint_complete.json")) - for path in (weights, spec, checkpoint_manifest): - if not path.is_file(): - raise FileNotFoundError(path) - base_model = resolve_model_dir(base_model, cache_dir=cache_dir, local_files_only=local_files_only) - package, condition, output = (p.resolve() for p in (package, condition, output)) - if manifest is not None: - manifest = manifest.resolve() - entry = package / "infer.py" - if not entry.is_file(): - raise FileNotFoundError(f"FireRed public entrypoint is missing: {entry}") - if not base_model.exists(): - raise FileNotFoundError(base_model) - output.mkdir(parents=True, exist_ok=True) - targets = [output / f"{key}_{channel}.png" for channel in ("albedo", "mra", "normal")] - targets += [output / name for name in ("albedo.png", "mra.png", "firered_albedo.png", "firered_mra.png")] - if any(path.exists() for path in targets): - raise FileExistsError("material outputs already exist; select a fresh output directory") - manifest = manifest or output / "firered_input.jsonl" - if manifest == output / "firered_input.jsonl": - write_manifest(condition, manifest, key=key) - elif not manifest.is_file(): - raise FileNotFoundError(manifest) - command = [sys.executable, str(entry), "--base-model", str(base_model), "--items-jsonl", str(manifest), "--out-dir", str(output)] - bundled = package.resolve() == (Path(__file__).resolve().parent / "material_model").resolve() - if bundled: - if not all((weights, spec, checkpoint_manifest)): - raise ValueError("bundled material model requires --weights, --spec, and --checkpoint-manifest") - command += ["--weights", str(weights), "--spec", str(spec), "--checkpoint-manifest", str(checkpoint_manifest)] - subprocess.run(command, check=True) - albedo, mra = output / f"{key}_albedo.png", output / f"{key}_mra.png" - if not albedo.is_file() or not mra.is_file(): - raise FileNotFoundError("FireRed did not emit the expected albedo and MRA maps") - from PIL import Image - for generated in (albedo, mra): - with Image.open(generated) as image: - if image.mode != "RGB" or image.size != (1024, 1024): - raise ValueError(f"material prediction must be RGB 1024x1024: {generated}") - canonical_albedo, canonical_mra = output / "firered_albedo.png", output / "firered_mra.png" - shutil.copyfile(albedo, canonical_albedo) - shutil.copyfile(mra, canonical_mra) - shutil.copyfile(albedo, output / "albedo.png") - shutil.copyfile(mra, output / "mra.png") - return canonical_albedo, canonical_mra - - -def main() -> None: - parser = argparse.ArgumentParser(description="Run FireRed Image-to-Material Translation through its packaged public entrypoint.") - parser.add_argument("--package", type=Path, default=Path(__file__).resolve().parent / "material_model", - help="Bundled material-model package or a compatible package containing infer.py") - parser.add_argument("--base-model", required=True, help="local FireRed base directory or Hugging Face namespace/repo") - parser.add_argument("--material-model", help="local material-model directory or Hugging Face namespace/repo") - parser.add_argument("--revision", help="material-model Hugging Face revision") - parser.add_argument("--cache-dir", type=Path) - parser.add_argument("--local-files-only", action="store_true") - parser.add_argument("--condition", type=Path, required=True) - parser.add_argument("--manifest", type=Path, - help="Optional package-contract JSONL; omitted builds a validation-only local record.") - parser.add_argument("--key", default="input") - parser.add_argument("--weights", type=Path, help="External factorized material-model weights for the bundled package.") - parser.add_argument("--spec", type=Path, help="External factorized material-model checkpoint contract.") - parser.add_argument("--checkpoint-manifest", type=Path, help="External material-model completion manifest.") - parser.add_argument("--output", type=Path, required=True) - args = parser.parse_args() - albedo, mra = run(args.package, args.base_model, args.condition, args.output, args.manifest, key=args.key, - weights=args.weights, spec=args.spec, checkpoint_manifest=args.checkpoint_manifest, - material_model=args.material_model, revision=args.revision, cache_dir=args.cache_dir, - local_files_only=args.local_files_only) - print(f"albedo={albedo}\nmra={mra}") - - -if __name__ == "__main__": - main() +"""FireRed Image-to-Material Translation adapter. + +The supplied package remains the sole FireRed implementation: this adapter +creates its documented JSONL item, runs its public ``infer.py``, and fixes +the decoded MR blue channel to opaque for the 3D material encoder. +""" +from __future__ import annotations + +import argparse +import json +import shutil +import subprocess +import sys +from pathlib import Path + + +def item_for(condition: Path, key: str) -> dict[str, object]: + """Build the public package record without inventing material supervision.""" + condition = condition.resolve() + return { + "key": key, "dataset": "user", "sha": key, "asset_id": key, "view": 0, "cond": str(condition), + # The public package validates these schema slots but prediction only + # consumes ``cond``. Point them at the supplied condition so there is + # no fabricated or private GT material file in an inference request. + "albedo_gt": str(condition), "mra_gt": str(condition), "normal_gt": str(condition), "normal_valid": str(condition), + "instr_albedo": "Convert this rendered image into its albedo (base color) map: remove all lighting, shading, shadows and highlights, keep only the intrinsic surface colors on a black background.", + # Retain the checkpoint's native prompt; run() replaces its predicted + # material alpha with the constant opaque channel used by completion. + "instr_mra": "Convert this rendered image into its material property map: encode metallic in the red channel, roughness in the green channel, and material alpha in the blue channel, on a black background.", + "instr_normal": "Convert this rendered image into its camera-space surface normal map: encode the normal direction at every pixel as RGB with x pointing right, y pointing up and z toward the camera, on a flat gray background.", "split": "eval", + } + + +def write_manifest(condition: Path, path: Path, *, key: str = "input") -> Path: + if not condition.is_file(): + raise FileNotFoundError(condition) + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text(json.dumps(item_for(condition, key), sort_keys=True) + "\n", encoding="utf-8") + return path + + +def run(package: Path, base_model: str | Path, condition: Path, output: Path, manifest: Path | None = None, + *, key: str = "input", weights: Path | None = None, spec: Path | None = None, + checkpoint_manifest: Path | None = None, material_model: str | Path | None = None, + revision: str | None = None, cache_dir: Path | None = None, + local_files_only: bool = False) -> tuple[Path, Path]: + if not key or Path(key).name != key or key in {".", ".."} or "\\" in key: + raise ValueError("key must be a single safe filename component") + from supratex.weights import resolve_model_dir + if material_model is not None: + if any(value is not None for value in (weights, spec, checkpoint_manifest)): + raise ValueError("--material-model cannot be combined with explicit material artifact paths") + root = resolve_model_dir(material_model, revision=revision, cache_dir=cache_dir, + local_files_only=local_files_only, + allow_patterns=["full_transformer.safetensors", "joint_contract.json", "checkpoint_complete.json"]) + weights, spec, checkpoint_manifest = (root / name for name in + ("full_transformer.safetensors", "joint_contract.json", "checkpoint_complete.json")) + for path in (weights, spec, checkpoint_manifest): + if not path.is_file(): + raise FileNotFoundError(path) + base_model = resolve_model_dir(base_model, cache_dir=cache_dir, local_files_only=local_files_only) + package, condition, output = (p.resolve() for p in (package, condition, output)) + if manifest is not None: + manifest = manifest.resolve() + entry = package / "infer.py" + if not entry.is_file(): + raise FileNotFoundError(f"FireRed public entrypoint is missing: {entry}") + if not base_model.exists(): + raise FileNotFoundError(base_model) + output.mkdir(parents=True, exist_ok=True) + targets = [output / f"{key}_{channel}.png" for channel in ("albedo", "mra", "normal")] + targets += [output / name for name in ("albedo.png", "mra.png", "firered_albedo.png", "firered_mra.png")] + if any(path.exists() for path in targets): + raise FileExistsError("material outputs already exist; select a fresh output directory") + manifest = manifest or output / "firered_input.jsonl" + if manifest == output / "firered_input.jsonl": + write_manifest(condition, manifest, key=key) + elif not manifest.is_file(): + raise FileNotFoundError(manifest) + command = [sys.executable, str(entry), "--base-model", str(base_model), "--items-jsonl", str(manifest), "--out-dir", str(output)] + bundled = package.resolve() == (Path(__file__).resolve().parent / "material_model").resolve() + if bundled: + if not all((weights, spec, checkpoint_manifest)): + raise ValueError("bundled material model requires --weights, --spec, and --checkpoint-manifest") + command += ["--weights", str(weights), "--spec", str(spec), "--checkpoint-manifest", str(checkpoint_manifest)] + subprocess.run(command, check=True) + albedo, mra = output / f"{key}_albedo.png", output / f"{key}_mra.png" + if not albedo.is_file() or not mra.is_file(): + raise FileNotFoundError("FireRed did not emit the expected albedo and MRA maps") + from PIL import Image + for generated in (albedo, mra): + with Image.open(generated) as image: + if image.mode != "RGB" or image.size != (1024, 1024): + raise ValueError(f"material prediction must be RGB 1024x1024: {generated}") + # Completion consumes (albedo RGB, metallic, roughness, opacity). Its MR + # placeholder is constant 1; foreground coverage comes from a separate mask. + with Image.open(mra) as image: + fixed_mra = Image.merge("RGB", (image.getchannel("R"), image.getchannel("G"), + Image.new("L", image.size, 255))) + fixed_mra.save(mra) + canonical_albedo, canonical_mra = output / "firered_albedo.png", output / "firered_mra.png" + shutil.copyfile(albedo, canonical_albedo) + shutil.copyfile(mra, canonical_mra) + shutil.copyfile(albedo, output / "albedo.png") + shutil.copyfile(mra, output / "mra.png") + return canonical_albedo, canonical_mra + + +def main() -> None: + parser = argparse.ArgumentParser(description="Run FireRed Image-to-Material Translation through its packaged public entrypoint.") + parser.add_argument("--package", type=Path, default=Path(__file__).resolve().parent / "material_model", + help="Bundled material-model package or a compatible package containing infer.py") + parser.add_argument("--base-model", required=True, help="local FireRed base directory or Hugging Face namespace/repo") + parser.add_argument("--material-model", help="local material-model directory or Hugging Face namespace/repo") + parser.add_argument("--revision", help="material-model Hugging Face revision") + parser.add_argument("--cache-dir", type=Path) + parser.add_argument("--local-files-only", action="store_true") + parser.add_argument("--condition", type=Path, required=True) + parser.add_argument("--manifest", type=Path, + help="Optional package-contract JSONL; omitted builds a validation-only local record.") + parser.add_argument("--key", default="input") + parser.add_argument("--weights", type=Path, help="External factorized material-model weights for the bundled package.") + parser.add_argument("--spec", type=Path, help="External factorized material-model checkpoint contract.") + parser.add_argument("--checkpoint-manifest", type=Path, help="External material-model completion manifest.") + parser.add_argument("--output", type=Path, required=True) + args = parser.parse_args() + albedo, mra = run(args.package, args.base_model, args.condition, args.output, args.manifest, key=args.key, + weights=args.weights, spec=args.spec, checkpoint_manifest=args.checkpoint_manifest, + material_model=args.material_model, revision=args.revision, cache_dir=args.cache_dir, + local_files_only=args.local_files_only) + print(f"albedo={albedo}\nmra={mra}") + + +if __name__ == "__main__": + main() diff --git a/supratex/frontend/geometry.py b/supratex/frontend/geometry.py index 6d3e5101fed818b22ed37507ab16f55c6360b4eb..e467a643ccc3c4ad9d3401344b03321696ed8e7f 100644 --- a/supratex/frontend/geometry.py +++ b/supratex/frontend/geometry.py @@ -1,169 +1,174 @@ -"""Run native or paper Pixal3D geometry with bounded mesh-scale calibration.""" -from __future__ import annotations - -import argparse -import json -import math -import os -import shutil -import subprocess -import sys -from pathlib import Path - -from supratex.weights import resolve_model_dir - - -def validate_source(package: Path) -> None: - """Check the source checkout before loading any image-stage models.""" - if (package / "pixal3d/pipelines/pixal3d_image_to_3d.py").is_file(): - return - for name in ("pixal3dpipeline.py", "pixal3dpipeline2stage.py"): - if not (package / name).is_file(): - raise FileNotFoundError(f"Pixal3D two-stage source is missing: {package / name}") - - -NATIVE_GEOMETRY_MODELS = ("sparse_structure_flow_model", "sparse_structure_decoder", - "shape_slat_flow_model_512", "shape_slat_flow_model_1024", "shape_slat_decoder") - - -def validate_model(package: Path, model_root: Path) -> str: - """Require complete local artifacts; native model stems must stay in the root.""" - native = (package / "pixal3d/pipelines/pixal3d_image_to_3d.py").is_file() - if native: - config = model_root / "pipeline.json" - if not config.is_file(): - raise FileNotFoundError(f"geometry model artifact is missing: {config}") - models = json.loads(config.read_text())["args"]["models"] - required = [] - for name in NATIVE_GEOMETRY_MODELS: - stem = (model_root / models[name]).resolve() - if not stem.is_relative_to(model_root.resolve()): - raise ValueError(f"geometry model must be inside model root: {name}") - required.extend(Path(str(stem) + suffix) for suffix in (".json", ".safetensors")) - dino = model_root / "dino" - required.append(dino / "config.json") - index = dino / "model.safetensors.index.json" - if index.is_file(): - for shard in set(json.loads(index.read_text())["weight_map"].values()): - path = (dino / shard).resolve() - if not path.is_relative_to(dino.resolve()): - raise ValueError("DINO weight shard must be inside its model directory") - required.append(path) - else: - required.append(dino / "model.safetensors") - else: - required = [] - for stage in ("dense", "sparse512", "sparse1024"): - for component in ("dit", "vae", "conditioner", "scheduler"): - required.append(model_root / stage / component / "config.json") - if component in ("dit", "vae"): - required.append(model_root / stage / component / "model.safetensors") - for path in required: - if not path.is_file(): - raise FileNotFoundError(f"geometry model artifact is missing: {path}") - return "native" if native else "paper" - - -def _validate_inputs(package: Path, input_dir: Path) -> None: - validate_source(package) - for name in ("condition_rgb.png", "condition_rgba.png", "camera.json"): - if not (input_dir / name).is_file(): - raise FileNotFoundError(f"geometry input is missing: {input_dir / name}") - from PIL import Image - with Image.open(input_dir / "condition_rgb.png") as rgb, Image.open(input_dir / "condition_rgba.png") as rgba: - if rgb.mode != "RGB" or rgba.mode != "RGBA" or rgb.size != rgba.size or rgb.width != rgb.height: - raise ValueError("geometry conditions must be matching square RGB and RGBA images") - if rgba.getchannel("A").getbbox() is None: - raise ValueError("geometry condition alpha must contain foreground") - rgb.load() - rgba.load() - try: - camera = json.loads((input_dir / "camera.json").read_text(encoding="utf-8")) - if not isinstance(camera, dict): - raise ValueError("camera must be an object") - if any(key in camera for key in ("extrinsic", "intrinsic", "projection_intrinsic")): - raise ValueError("geometry requires the canonical camera, without explicit matrices") - distance, fov = float(camera["distance"]), float(camera["camera_angle_x"]) - if not math.isfinite(distance) or distance <= 0 or not math.isfinite(fov) or not 0 < fov < math.pi: - raise ValueError("invalid distance or field of view") - except (KeyError, TypeError, ValueError) as error: - raise ValueError("camera.json must provide finite positive distance and camera_angle_x in (0, pi), without explicit matrices") from error - - -def run(package: Path, model_source: str | Path, input_dir: Path, output_dir: Path, *, seed: int = 42, - initial_mesh_scale: float = 1.0, target_padding: int = 3, max_iterations: int = 2, - cache_dir: Path | None = None, revision: str | None = None, - local_files_only: bool = False, resident=None) -> dict[str, str]: - """Generate a final mesh using the adapter's bounded scale calibration. - - The external implementation runs in a separate interpreter so its package - imports and GPU resources do not share the material frontend's process. - """ - package, input_dir, output_dir = (Path(path).expanduser().resolve() for path in (package, input_dir, output_dir)) - if output_dir.exists() and (not output_dir.is_dir() or any(output_dir.iterdir())): - raise FileExistsError("geometry output must be a fresh empty directory") - _validate_inputs(package, input_dir) - if not math.isfinite(initial_mesh_scale) or initial_mesh_scale <= 0: - raise ValueError("initial_mesh_scale must be finite and positive") - if isinstance(target_padding, bool) or not isinstance(target_padding, int) or not 1 <= target_padding <= 30: - raise ValueError("target_padding must be an integer between 1 and 30 for the 64-cell grid") - if isinstance(max_iterations, bool) or not isinstance(max_iterations, int) or max_iterations < 1: - raise ValueError("max_iterations must be a positive integer") - if isinstance(seed, bool) or not isinstance(seed, int) or seed < 0: - raise ValueError("seed must be a nonnegative integer") - model_root = resolve_model_dir(model_source, revision=revision, cache_dir=cache_dir, - local_files_only=local_files_only) - # The upstream local loader permits missing weight files. Require them here - # so an incomplete checkpoint can never silently run random initial weights. - validate_model(package, model_root) - output_dir.mkdir(parents=True, exist_ok=True) - runner = Path(__file__).resolve().with_name("geometry_runner.py") - command = [sys.executable, str(runner), "--package", str(package), "--model-root", str(model_root), - "--input-dir", str(input_dir), "--output-dir", str(output_dir), "--seed", str(seed), - "--initial-mesh-scale", str(initial_mesh_scale), "--target-padding", str(target_padding), - "--max-iterations", str(max_iterations)] - environment = os.environ.copy() - source_root = Path(__file__).resolve().parents[2] - environment["PYTHONPATH"] = os.pathsep.join((str(package), str(source_root))) + (os.pathsep + environment["PYTHONPATH"] if environment.get("PYTHONPATH") else "") - if local_files_only: - environment["HF_HUB_OFFLINE"] = "1" - environment["TRANSFORMERS_OFFLINE"] = "1" - if resident is not None: - resident.run_geometry(package=package, model_root=model_root, input_dir=input_dir, - output_dir=output_dir, seed=seed, initial_mesh_scale=initial_mesh_scale, - target_padding=target_padding, max_iterations=max_iterations) - else: - with (output_dir / "geometry.log").open("x", encoding="utf-8") as log: - subprocess.run(command, check=True, cwd=package, env=environment, stdout=log, stderr=subprocess.STDOUT) - results = {"geometry": output_dir / "geometry.pt", "camera": output_dir / "camera.json", - "receipt": output_dir / "geometry_receipt.json"} - for path in results.values(): - if not path.is_file(): - raise FileNotFoundError(f"geometry runner did not emit expected output: {path}") - for name in ("condition_rgb.png", "condition_rgba.png"): - shutil.copyfile(input_dir / name, output_dir / name) - return {key: str(path) for key, path in results.items()} - - -def main() -> None: - parser = argparse.ArgumentParser(description="Generate the final mesh with native or paper Pixal3D geometry and bounded scale calibration.") - parser.add_argument("--package", type=Path, required=True, help="Native source checkout (pixal3d/pipelines/) or paper source (pixal3dpipeline2stage.py)") - parser.add_argument("--model", required=True, help="Local geometry-model directory or Hugging Face namespace/repo") - parser.add_argument("--input-dir", type=Path, required=True, help="Condition RGB/RGBA and MoGe camera directory") - parser.add_argument("--output-dir", type=Path, required=True, help="Fresh empty geometry output directory") - parser.add_argument("--seed", type=int, default=42) - parser.add_argument("--initial-mesh-scale", type=float, default=1.0) - parser.add_argument("--target-padding", type=int, default=3) - parser.add_argument("--max-iterations", type=int, default=2) - parser.add_argument("--cache-dir", type=Path) - parser.add_argument("--revision") - parser.add_argument("--local-files-only", action="store_true") - args = parser.parse_args() - print(json.dumps(run(args.package, args.model, args.input_dir, args.output_dir, seed=args.seed, - initial_mesh_scale=args.initial_mesh_scale, target_padding=args.target_padding, - max_iterations=args.max_iterations, cache_dir=args.cache_dir, revision=args.revision, - local_files_only=args.local_files_only), indent=2)) - - -if __name__ == "__main__": - main() +"""Run native or paper Pixal3D geometry with bounded mesh-scale calibration.""" +from __future__ import annotations + +import argparse +import json +import math +import os +import shutil +import subprocess +import sys +from pathlib import Path + +from supratex.weights import resolve_model_dir + + +def validate_source(package: Path) -> None: + """Check the source checkout before loading any image-stage models.""" + if (package / "pixal3d/pipelines/pixal3d_image_to_3d.py").is_file(): + return + for name in ("pixal3dpipeline.py", "pixal3dpipeline2stage.py"): + if not (package / name).is_file(): + raise FileNotFoundError(f"Pixal3D two-stage source is missing: {package / name}") + + +NATIVE_GEOMETRY_MODELS = ("sparse_structure_flow_model", "sparse_structure_decoder", + "shape_slat_flow_model_512", "shape_slat_flow_model_1024", "shape_slat_decoder") + + +def validate_model(package: Path, model_root: Path) -> str: + """Require complete local artifacts; native model stems must stay in the root.""" + native = (package / "pixal3d/pipelines/pixal3d_image_to_3d.py").is_file() + if native: + config = model_root / "pipeline.json" + if not config.is_file(): + raise FileNotFoundError(f"geometry model artifact is missing: {config}") + models = json.loads(config.read_text())["args"]["models"] + required = [] + for name in NATIVE_GEOMETRY_MODELS: + stem = (model_root / models[name]).resolve() + if not stem.is_relative_to(model_root.resolve()): + raise ValueError(f"geometry model must be inside model root: {name}") + required.extend(Path(str(stem) + suffix) for suffix in (".json", ".safetensors")) + dino = model_root / "dino" + required.append(dino / "config.json") + index = dino / "model.safetensors.index.json" + if index.is_file(): + weight_map = json.loads(index.read_text())["weight_map"] + if not isinstance(weight_map, dict) or not weight_map: + raise ValueError("DINO checkpoint index requires a nonempty weight_map") + for shard in weight_map.values(): + if (not isinstance(shard, str) or not shard or "\\" in shard or ":" in shard + or Path(shard).is_absolute() or ".." in Path(shard).parts): + raise ValueError("DINO weight shard must be inside its model directory") + # HF snapshot files point to shared cache blobs outside the + # snapshot. Validate the indexed name, then follow its link + # only when checking that the required file exists. + required.append(dino / shard) + else: + required.append(dino / "model.safetensors") + else: + required = [] + for stage in ("dense", "sparse512", "sparse1024"): + for component in ("dit", "vae", "conditioner", "scheduler"): + required.append(model_root / stage / component / "config.json") + if component in ("dit", "vae"): + required.append(model_root / stage / component / "model.safetensors") + for path in required: + if not path.is_file(): + raise FileNotFoundError(f"geometry model artifact is missing: {path}") + return "native" if native else "paper" + + +def _validate_inputs(package: Path, input_dir: Path) -> None: + validate_source(package) + for name in ("condition_rgb.png", "condition_rgba.png", "camera.json"): + if not (input_dir / name).is_file(): + raise FileNotFoundError(f"geometry input is missing: {input_dir / name}") + from PIL import Image + with Image.open(input_dir / "condition_rgb.png") as rgb, Image.open(input_dir / "condition_rgba.png") as rgba: + if rgb.mode != "RGB" or rgba.mode != "RGBA" or rgb.size != rgba.size or rgb.width != rgb.height: + raise ValueError("geometry conditions must be matching square RGB and RGBA images") + if rgba.getchannel("A").getbbox() is None: + raise ValueError("geometry condition alpha must contain foreground") + rgb.load() + rgba.load() + try: + camera = json.loads((input_dir / "camera.json").read_text(encoding="utf-8")) + if not isinstance(camera, dict): + raise ValueError("camera must be an object") + if any(key in camera for key in ("extrinsic", "intrinsic", "projection_intrinsic")): + raise ValueError("geometry requires the canonical camera, without explicit matrices") + distance, fov = float(camera["distance"]), float(camera["camera_angle_x"]) + if not math.isfinite(distance) or distance <= 0 or not math.isfinite(fov) or not 0 < fov < math.pi: + raise ValueError("invalid distance or field of view") + except (KeyError, TypeError, ValueError) as error: + raise ValueError("camera.json must provide finite positive distance and camera_angle_x in (0, pi), without explicit matrices") from error + + +def run(package: Path, model_source: str | Path, input_dir: Path, output_dir: Path, *, seed: int = 42, + initial_mesh_scale: float = 1.0, scale_search: bool = False, target_padding: int = 3, max_iterations: int = 2, + cuda_visible_devices: str | None = None, + cache_dir: Path | None = None, revision: str | None = None, + local_files_only: bool = False) -> dict[str, str]: + """Generate a final mesh using the adapter's bounded scale calibration. + + The external implementation runs in a separate interpreter so its package + imports and GPU resources do not share the material frontend's process. + """ + package, input_dir, output_dir = (Path(path).expanduser().resolve() for path in (package, input_dir, output_dir)) + if output_dir.exists() and (not output_dir.is_dir() or any(output_dir.iterdir())): + raise FileExistsError("geometry output must be a fresh empty directory") + _validate_inputs(package, input_dir) + if not math.isfinite(initial_mesh_scale) or initial_mesh_scale <= 0: + raise ValueError("initial_mesh_scale must be finite and positive") + if not 1 <= target_padding <= 30: + raise ValueError("target_padding must be between 1 and 30 for the 64-cell grid") + if max_iterations < 1: + raise ValueError("max_iterations must be positive") + if seed < 0: + raise ValueError("seed must be nonnegative") + model_root = resolve_model_dir(model_source, revision=revision, cache_dir=cache_dir, + local_files_only=local_files_only) + # The upstream local loader permits missing weight files. Require them here + # so an incomplete checkpoint can never silently run random initial weights. + validate_model(package, model_root) + output_dir.mkdir(parents=True, exist_ok=True) + runner = Path(__file__).resolve().with_name("geometry_runner.py") + command = [sys.executable, str(runner), "--package", str(package), "--model-root", str(model_root), + "--input-dir", str(input_dir), "--output-dir", str(output_dir), "--seed", str(seed), + "--initial-mesh-scale", str(initial_mesh_scale), "--target-padding", str(target_padding), + "--max-iterations", str(max_iterations), *(["--scale-search"] if scale_search else [])] + environment = os.environ.copy() + source_root = Path(__file__).resolve().parents[2] + environment["PYTHONPATH"] = os.pathsep.join((str(package), str(source_root))) + (os.pathsep + environment["PYTHONPATH"] if environment.get("PYTHONPATH") else "") + if cuda_visible_devices is not None: + environment["CUDA_VISIBLE_DEVICES"] = cuda_visible_devices + if local_files_only: + environment["HF_HUB_OFFLINE"] = "1" + environment["TRANSFORMERS_OFFLINE"] = "1" + with (output_dir / "geometry.log").open("x", encoding="utf-8") as log: + subprocess.run(command, check=True, cwd=package, env=environment, stdout=log, stderr=subprocess.STDOUT) + results = {"geometry": output_dir / "geometry.pt", "camera": output_dir / "camera.json", + "receipt": output_dir / "geometry_receipt.json"} + for path in results.values(): + if not path.is_file(): + raise FileNotFoundError(f"geometry runner did not emit expected output: {path}") + for name in ("condition_rgb.png", "condition_rgba.png"): + shutil.copyfile(input_dir / name, output_dir / name) + return {key: str(path) for key, path in results.items()} + + +def main() -> None: + parser = argparse.ArgumentParser(description="Generate the final mesh with native or paper Pixal3D geometry and bounded scale calibration.") + parser.add_argument("--package", type=Path, required=True, help="Native source checkout (pixal3d/pipelines/) or paper source (pixal3dpipeline2stage.py)") + parser.add_argument("--model", required=True, help="Local geometry-model directory or Hugging Face namespace/repo") + parser.add_argument("--input-dir", type=Path, required=True, help="Condition RGB/RGBA and MoGe camera directory") + parser.add_argument("--output-dir", type=Path, required=True, help="Fresh empty geometry output directory") + parser.add_argument("--seed", type=int, default=42) + parser.add_argument("--initial-mesh-scale", type=float, default=1.0) + parser.add_argument("--scale-search", action="store_true", help="enable bounded mesh-scale calibration (default: Pixal3D fixed scale)") + parser.add_argument("--target-padding", type=int, default=3) + parser.add_argument("--max-iterations", type=int, default=2) + parser.add_argument("--cache-dir", type=Path) + parser.add_argument("--revision") + parser.add_argument("--local-files-only", action="store_true") + args = parser.parse_args() + print(json.dumps(run(args.package, args.model, args.input_dir, args.output_dir, seed=args.seed, + initial_mesh_scale=args.initial_mesh_scale, scale_search=args.scale_search, target_padding=args.target_padding, + max_iterations=args.max_iterations, cache_dir=args.cache_dir, revision=args.revision, + local_files_only=args.local_files_only), indent=2)) + + +if __name__ == "__main__": + main() diff --git a/supratex/frontend/geometry_loading.py b/supratex/frontend/geometry_loading.py index 42365782a95d809c3bfc6f348efede537f8e79ce..574b55d53ace1de890972842fde8e87510300290 100644 --- a/supratex/frontend/geometry_loading.py +++ b/supratex/frontend/geometry_loading.py @@ -1,69 +1,128 @@ -"""Strict local checkpoint loading without initializing discarded geometry weights.""" - -import json -from pathlib import Path - -import torch - - -_MODEL_NAMES = { - "SparseStructureFlowModel", "SparseStructureDecoder", - "ElasticSLatFlowModel", "FlexiDualGridVaeDecoder", -} - - -def _restore_position_frequencies(model): - """Recreate small native tensors omitted from state_dict, on the CPU.""" - for module in model.modules(): - frequencies = vars(module).get("freqs") - if not isinstance(frequencies, torch.Tensor) or not frequencies.is_meta: - continue - cls = type(module) - with torch.device("cpu"): - if (cls.__name__ in {"RotaryPositionEmbedder", "SparseRotaryPositionEmbedder"} - and cls.__module__.startswith("pixal3d.modules.") - and cls.__module__.endswith(".attention.rope")): - original = cls(module.head_dim, dim=module.dim, rope_freq=module.rope_freq) - elif (cls.__name__ == "AbsolutePositionEmbedder" - and cls.__module__ == "pixal3d.modules.transformer.blocks"): - original = cls(module.channels, in_channels=module.in_channels) - else: - raise ValueError(f"Unsupported geometry meta frequency tensor: {cls.__module__}.{cls.__name__}") - module.freqs = original.freqs - - -def load_native_geometry_model(path: Path): - """Load a reviewed local Pixal3D network with its constructor's dtype ABI. - - Native flow checkpoints store float32 even for bfloat16 transformer blocks; - decoder checkpoints can instead store float16 for float32 boundary layers. - Cast to each meta-shell tensor's dtype before assignment, exactly as native - copy-based load_state_dict does. Unknown/missing weights fail before return. - """ - from pixal3d import models - from safetensors.torch import load_file - - path = Path(path) - config = json.loads(Path(str(path) + ".json").read_text(encoding="utf-8")) - if config["name"] not in _MODEL_NAMES: - raise ValueError(f"Unsupported native geometry model: {config['name']}") - with torch.device("meta"): - model = getattr(models, config["name"])(**config["args"]) - expected = model.state_dict() - state = load_file(str(path) + ".safetensors", device="cpu") - missing, unexpected = expected.keys() - state.keys(), state.keys() - expected.keys() - if missing or unexpected: - raise ValueError(f"Geometry checkpoint key mismatch: missing={sorted(missing)}, unexpected={sorted(unexpected)}") - for name, target in expected.items(): - value = state[name] - if value.shape != target.shape: - raise ValueError(f"Geometry checkpoint shape mismatch for {name}: {tuple(value.shape)} != {tuple(target.shape)}") - state[name] = value.to(dtype=target.dtype) - model.load_state_dict(state, strict=True, assign=True) - _restore_position_frequencies(model) - for name, module in model.named_modules(): - tensors = list(module._parameters.items()) + list(module._buffers.items()) + list(vars(module).items()) - for key, value in tensors: - if isinstance(value, torch.Tensor) and value.is_meta: - raise ValueError(f"Unmaterialized geometry tensor: {name}.{key}") - return model.eval() +"""Strict local checkpoint loading without initializing discarded geometry weights.""" + +import json +from pathlib import Path + +import torch + + +_MODEL_NAMES = { + "SparseStructureFlowModel", "SparseStructureDecoder", + "ElasticSLatFlowModel", "FlexiDualGridVaeDecoder", +} + + +def _restore_position_frequencies(model): + """Recreate small native tensors omitted from state_dict, on the CPU.""" + for module in model.modules(): + frequencies = vars(module).get("freqs") + if not isinstance(frequencies, torch.Tensor) or not frequencies.is_meta: + continue + cls = type(module) + with torch.device("cpu"): + if (cls.__name__ in {"RotaryPositionEmbedder", "SparseRotaryPositionEmbedder"} + and cls.__module__.startswith("pixal3d.modules.") + and cls.__module__.endswith(".attention.rope")): + original = cls(module.head_dim, dim=module.dim, rope_freq=module.rope_freq) + elif (cls.__name__ == "AbsolutePositionEmbedder" + and cls.__module__ == "pixal3d.modules.transformer.blocks"): + original = cls(module.channels, in_channels=module.in_channels) + else: + raise ValueError(f"Unsupported geometry meta frequency tensor: {cls.__module__}.{cls.__name__}") + module.freqs = original.freqs + + +RUNTIME_SUFFIX = ".runtime.safetensors" + + +def _meta_model(path: Path): + from pixal3d import models + + config = json.loads(Path(str(path) + ".json").read_text(encoding="utf-8")) + if config["name"] not in _MODEL_NAMES: + raise ValueError(f"Unsupported native geometry model: {config['name']}") + with torch.device("meta"): + return getattr(models, config["name"])(**config["args"]) + + +def _source_identity(path: Path) -> dict[str, str]: + stat = Path(str(path) + ".safetensors").resolve().stat() + return {"source_size": str(stat.st_size), "source_mtime_ns": str(stat.st_mtime_ns)} + + +def _runtime_state(path: Path, expected: dict) -> dict: + """Read checkpoint tensors and cast them to the meta shell's dtypes.""" + from safetensors.torch import load_file + + state = load_file(str(path) + ".safetensors", device="cpu") + missing, unexpected = expected.keys() - state.keys(), state.keys() - expected.keys() + if missing or unexpected: + raise ValueError(f"Geometry checkpoint key mismatch: missing={sorted(missing)}, unexpected={sorted(unexpected)}") + for name, target in expected.items(): + value = state[name] + if value.shape != target.shape: + raise ValueError(f"Geometry checkpoint shape mismatch for {name}: {tuple(value.shape)} != {tuple(target.shape)}") + state[name] = value.to(dtype=target.dtype) + return state + + +def write_runtime_weights(path: Path) -> Path | None: + """Save the cast state beside the checkpoint when that copy is smaller. + + Native flow checkpoints store float32 for bfloat16 blocks, so the cast copy + halves the bytes read per load and loads bit-identical tensors. Decoders + that store float16 for float32 layers would grow, so they are skipped (None). + """ + from safetensors.torch import save_file + + path = Path(path) + expected = _meta_model(path).state_dict() + runtime_bytes = sum(target.numel() * target.element_size() for target in expected.values()) + # Only worth a second file when it is clearly smaller than the original. + if runtime_bytes > .9 * Path(str(path) + ".safetensors").resolve().stat().st_size: + return None + target = Path(str(path) + RUNTIME_SUFFIX) + temporary = target.with_name(target.name + ".tmp") + save_file({name: tensor.contiguous() for name, tensor in _runtime_state(path, expected).items()}, + str(temporary), metadata=_source_identity(path)) + temporary.replace(target) + return target + + +def _runtime_copy(path: Path) -> Path | None: + """Use a runtime-dtype copy only if it was made from the current checkpoint.""" + from safetensors import safe_open + + runtime = Path(str(path) + RUNTIME_SUFFIX) + if not runtime.is_file(): + return None + with safe_open(str(runtime), framework="pt", device="cpu") as handle: + if handle.metadata() != _source_identity(path): + print(f"Ignoring stale runtime geometry weights: {runtime}", flush=True) + return None + return runtime + + +def load_native_geometry_model(path: Path): + """Load a reviewed local Pixal3D network with its constructor's dtype ABI. + + Native flow checkpoints store float32 even for bfloat16 transformer blocks; + decoder checkpoints can instead store float16 for float32 boundary layers. + Cast to each meta-shell tensor's dtype before assignment, exactly as native + copy-based load_state_dict does. A matching runtime-dtype copy written by + scripts/convert_geometry_weights.py is read instead when present. + Unknown/missing weights fail before return. + """ + path = Path(path) + model = _meta_model(path) + expected = model.state_dict() + runtime = _runtime_copy(path) + state = _runtime_state(Path(str(runtime).removesuffix(".safetensors")) if runtime else path, expected) + model.load_state_dict(state, strict=True, assign=True) + _restore_position_frequencies(model) + for name, module in model.named_modules(): + tensors = list(module._parameters.items()) + list(module._buffers.items()) + list(vars(module).items()) + for key, value in tensors: + if isinstance(value, torch.Tensor) and value.is_meta: + raise ValueError(f"Unmaterialized geometry tensor: {name}.{key}") + return model.eval() diff --git a/supratex/frontend/geometry_runner.py b/supratex/frontend/geometry_runner.py index 72be70eb2d15d8c831b6fda1aec67e5bc2d062a8..cc71abfa95fe48a3e9ed5789d172ac14fe6c3302 100644 --- a/supratex/frontend/geometry_runner.py +++ b/supratex/frontend/geometry_runner.py @@ -1,450 +1,434 @@ -"""Isolated adapter for native and paper Pixal3D geometry generators. - -The prepared demo camera owns FOV and distance. This adapter calibrates scale -using the main coarse geometry model, then publishes the final mesh/camera in -the SUPRA-Tex canonical frame. -""" -from __future__ import annotations - -import argparse -import copy -import importlib -import json -import math -from pathlib import Path -import sys - -from hf_space.timing import span - - -def calibrate_scale(probe, initial_scale: float = 1.0, *, target_padding: int = 3, - max_iterations: int = 2) -> tuple[float, dict]: - """Keep the closest safe occupancy, with a bounded number of scale updates. - - ``probe(scale)`` returns single-sample [N,4] decoded grid indices. Invalid - occupancy ends the search; inference/runtime errors propagate to the caller. - """ - import torch - - if not math.isfinite(initial_scale) or initial_scale <= 0: - raise ValueError("initial_scale must be finite and positive") - if isinstance(target_padding, bool) or not isinstance(target_padding, int) or not 1 <= target_padding <= 30: - raise ValueError("target_padding must be an integer in [1, 30]") - if isinstance(max_iterations, bool) or not isinstance(max_iterations, int) or max_iterations < 1: - raise ValueError("max_iterations must be a positive integer") - max_iterations = min(max_iterations, 2) - lower, upper = max(1, target_padding - 1), target_padding + 1 - scale = best = float(initial_scale) - best_padding, best_error = None, math.inf - report = {"algorithm": "bounded_main_dense_padding", "initial_mesh_scale": scale, - "target_padding": target_padding, "padding_tolerance": [lower, upper], - "scale_bounds": [scale / 2, scale * 2], "max_iterations": max_iterations, "probes": []} - reason = "max_iterations" - for iteration in range(max_iterations + 1): - coords = probe(scale) - valid = (isinstance(coords, torch.Tensor) and coords.ndim == 2 and coords.shape[1] == 4 - and coords.shape[0] > 0 and not coords.is_complex()) - if valid: - valid = bool(torch.isfinite(coords).all() and (coords[:, 0] == 0).all() - and (coords[:, 1:] >= 0).all() and (coords[:, 1:] <= 63).all() - and (coords == coords.float().round()).all()) - state = {"mesh_scale": scale, "valid": valid} - report["probes"].append(state) - if not valid: - reason = "invalid_occupancy" - break - xyz = coords[:, 1:].float() - padding = min(float(xyz.min()), 63 - float(xyz.max())) - offset = float((xyz - 31.5).abs().max()) - safe, error = padding >= lower, abs(padding - target_padding) - state.update(padding=padding, max_offset=offset, occupied_voxels=len(coords), safe=safe) - if safe and error < best_error: - best, best_padding, best_error = scale, padding, error - if lower <= padding <= upper: - reason = "target_reached" - break - if iteration == max_iterations: - break - proposed = scale * (31.5 - target_padding) / offset - if not math.isfinite(proposed) or not report["scale_bounds"][0] <= proposed <= report["scale_bounds"][1]: - reason = "scale_out_of_bounds" - break - if math.isclose(proposed, scale, rel_tol=1e-6): - reason = "scale_stalled" - break - scale = proposed - report.update(selected_mesh_scale=best, selected_padding=best_padding, - converged=best_padding is not None and lower <= best_padding <= upper, - used_baseline=best == initial_scale, stop_reason=reason) - return best, report - - -def traced_pipeline(base_class, camera: dict, receipt: dict): - """Customize scale calibration while keeping the upstream generation path.""" - fov = float(camera["camera_angle_x"]) - if not math.isfinite(fov) or not 0 < fov < math.pi: - raise ValueError("camera_angle_x must be finite and between zero and pi") - distance = float(camera["distance"]) - if not math.isfinite(distance) or distance <= 0: - raise ValueError("camera distance must be finite and positive") - - class GeometryPipeline(base_class): - _optimizing_scale = False - - def estimate_fov(self, image_path): - return fov - - def _optimize_mesh_scale(self, image_tensor, camera_angle_x_tensor, distance_tensor, - initial_mesh_scale, dense_steps, dense_guidance_scale, dense_seed, - dense_threshold, target_padding=3, max_iterations=2): - import torch - - device = image_tensor.device - cuda_devices = [device.index if device.index is not None else torch.cuda.current_device()] if device.type == "cuda" else [] - receipt["upstream_initial_generation_camera"] = { - "camera_angle_x": float(camera_angle_x_tensor.item()), - "distance": float(distance_tensor.item()), "mesh_scale": float(initial_mesh_scale)} - # The upstream infer() re-derives distance using its 518px/padding - # convention. Mutate its shared camera tensor before any probe so - # final dense and both sparse stages retain the prepared demo camera. - distance_tensor.fill_(distance) - receipt["camera_distance_policy"] = "preserve_prepared_demo_camera" - receipt["initial_generation_camera"] = { - "camera_angle_x": float(camera_angle_x_tensor.item()), - "distance": float(distance_tensor.item()), "mesh_scale": float(initial_mesh_scale)} - - def probe(scale): - receipt["dense_probe_calls"] = receipt.get("dense_probe_calls", 0) + 1 - original_scheduler = self.dense_scheduler - try: - # Probe state must not leak into another probe or final generation. - self.dense_scheduler = copy.deepcopy(original_scheduler) - with torch.random.fork_rng(devices=cuda_devices): - latents = self.infer_dense( - image_tensor, camera_angle_x_tensor, distance_tensor, - torch.tensor([scale], device=device, dtype=torch.float32), - dense_steps, dense_guidance_scale, dense_seed) - with torch.autocast(device.type, dtype=torch.float16, enabled=device.type == "cuda"): - return self.dense_vae.decode_mesh(latents, mc_threshold=dense_threshold, return_index=True)[0] - finally: - self.dense_scheduler = original_scheduler - - self._optimizing_scale = True - try: - with span("geometry.scale_search"): - result, report = calibrate_scale(probe, initial_mesh_scale, target_padding=target_padding, - max_iterations=max_iterations) - finally: - self._optimizing_scale = False - receipt["scale_search"] = report - receipt["optimized_mesh_scale"] = float(result) - print(f"[mesh_scale] selected={result:.6f}, stop={report['stop_reason']}, probes={len(report['probes'])}", flush=True) - return result - - def infer_dense(self, image, camera_angle_x, distance, mesh_scale, *args, **kwargs): - if not math.isclose(float(distance.item()), float(camera["distance"]), rel_tol=1e-5): - raise RuntimeError("geometry generation did not retain the prepared demo camera distance") - if not self._optimizing_scale: - receipt["final_generation_camera"] = { - "camera_angle_x": float(camera_angle_x.item()), - "distance": float(distance.item()), - "mesh_scale": float(mesh_scale.item()), - } - with span("geometry.dense_sampling"): - return super().infer_dense(image, camera_angle_x, distance, mesh_scale, *args, **kwargs) - - return GeometryPipeline - - -def canonical_handoff(mesh, camera: dict, receipt: dict): - """Undo upstream export rotation and scale while preserving projection.""" - import numpy as np - import torch - - try: - selected = receipt["final_generation_camera"] - scale = float(selected["mesh_scale"]) - optimized = float(receipt["optimized_mesh_scale"]) - except (KeyError, TypeError, ValueError) as error: - raise RuntimeError("geometry source did not expose the two-stage scale/camera hooks") from error - if not math.isfinite(scale) or scale <= 0 or not math.isclose(scale, optimized, rel_tol=1e-5): - raise RuntimeError("final geometry camera does not match optimized mesh_scale") - if not math.isfinite(selected["distance"]) or selected["distance"] <= 0: - raise RuntimeError("final geometry distance must be finite and positive") - if not math.isclose(selected["distance"], float(camera["distance"]), rel_tol=1e-5): - raise RuntimeError("final geometry did not retain the prepared demo camera distance") - if not math.isclose(selected["camera_angle_x"], float(camera["camera_angle_x"]), rel_tol=1e-5): - raise RuntimeError("geometry generation did not retain the supplied MoGe-2 FOV") - # Upstream exports R @ raw_xyz * (0.5 / scale), with R mapping - # (x,y,z) -> (x,-z,y). Canonical points are raw_xyz * 0.5. - rotation = np.array([[1., 0., 0.], [0., 0., -1.], [0., 1., 0.]]) - vertices = np.asarray(mesh.vertices, dtype=np.float32) @ rotation * scale - faces = np.asarray(mesh.faces) - if vertices.ndim != 2 or vertices.shape[1:] != (3,) or not len(vertices) or not np.isfinite(vertices).all(): - raise ValueError("geometry generator returned invalid vertices") - if faces.ndim != 2 or faces.shape[1:] != (3,) or not len(faces) or not np.issubdtype(faces.dtype, np.integer): - raise ValueError("geometry generator returned invalid faces") - if faces.min() < 0 or faces.max() >= len(vertices): - raise ValueError("geometry generator returned out-of-range face indices") - if float(np.abs(vertices).max()) > 0.5001: - raise ValueError("converted geometry exceeds the canonical [-0.5, 0.5] volume") - effective_distance = selected["distance"] * scale - if effective_distance <= float(vertices[:, 2].max()): - raise ValueError("optimized camera lies inside or behind the geometry") - final_camera = dict(camera) - final_camera.update(selected) - final_camera["moge_initial_camera"] = {key: camera[key] for key in - ("camera_angle_x", "distance", "mesh_scale") if key in camera} - final_camera["geometry_scale_optimization"] = { - "initial_mesh_scale": receipt["initial_mesh_scale"], - "optimized_mesh_scale": scale, - "effective_camera_distance": effective_distance, - "objective": "dense_grid_boundary_padding", - } - geometry = {"verts": torch.from_numpy(vertices.copy()).float(), - "faces": torch.from_numpy(faces.astype(np.int64, copy=True)).long()} - return geometry, final_camera - - -def build_native_pipeline(model_root: Path): - """Load only native geometry networks and their projection conditioners.""" - import torch - local_naf = Path(torch.hub.get_dir()) / "valeoai_NAF_main" - for artifact in (local_naf / "hubconf.py", local_naf.parent / "checkpoints/naf_release.pth"): - if not artifact.is_file(): - raise FileNotFoundError(f"native geometry requires the local NAF cache: {artifact}") - from pixal3d.pipelines import Pixal3DImageTo3DPipeline, samplers - from pixal3d.trainers.flow_matching.mixins.image_conditioned_proj import DinoV3ProjFeatureExtractor - from supratex.frontend.geometry_loading import load_native_geometry_model - - class GeometryOnlyPipeline(Pixal3DImageTo3DPipeline): - model_names_to_load = ["sparse_structure_flow_model", "sparse_structure_decoder", - "shape_slat_flow_model_512", "shape_slat_flow_model_1024", "shape_slat_decoder"] - - # Conditions already include foreground removal. Base loading avoids creating - # another background-removal model or any texture networks/samplers. - args = json.loads((model_root / "pipeline.json").read_text(encoding="utf-8"))["args"] - networks = {name: load_native_geometry_model(model_root / checkpoint) - for name, checkpoint in args["models"].items() - if name in GeometryOnlyPipeline.model_names_to_load} - missing = set(GeometryOnlyPipeline.model_names_to_load) - networks.keys() - if missing: - raise ValueError(f"Native geometry config is missing models: {sorted(missing)}") - pipe = GeometryOnlyPipeline(networks) - pipe._pretrained_args = args - for name in ("sparse_structure", "shape_slat"): - spec = args[f"{name}_sampler"] - setattr(pipe, f"{name}_sampler", getattr(samplers, spec["name"])(**spec["args"])) - setattr(pipe, f"{name}_sampler_params", spec["params"]) - pipe.shape_slat_normalization = args["shape_slat_normalization"] - original_load = torch.hub.load - - def local_hub_load(repo_or_dir, model, *args, **kwargs): - if repo_or_dir == "valeoai/NAF": - kwargs.pop("trust_repo", None) - kwargs["source"] = "local" - repo_or_dir = str(local_naf) - return original_load(repo_or_dir, model, *args, **kwargs) - - try: - torch.hub.load = local_hub_load - for name, size, grid, use_naf in (("ss", 512, 16, False), - ("shape_512", 512, 32, True), - ("shape_1024", 1024, 64, True)): - conditioner = DinoV3ProjFeatureExtractor( - model_name=str(model_root / "dino"), image_size=size, grid_resolution=grid, - use_naf_upsample=use_naf, naf_target_size=512).eval() - if use_naf: - conditioner._load_naf() - setattr(pipe, f"image_cond_model_{name}", conditioner) - finally: - torch.hub.load = original_load - pipe._device = torch.device("cuda") - pipe.low_vram = True - return pipe - - -def infer_native(pipe, image, camera: dict, receipt: dict, *, seed: int, - initial_mesh_scale: float, target_padding: int, max_iterations: int): - """Calibrate native 64-cell occupancy, then run the demo's shape cascade.""" - import torch - - device = torch.device(pipe.device) - cuda_devices = [device.index if device.index is not None else torch.cuda.current_device()] if device.type == "cuda" else [] - sampler_ss = {"steps": 12, "guidance_strength": 7.5, "guidance_rescale": .7, "rescale_t": 5.} - sampler_shape = {"steps": 12, "guidance_strength": 7.5, "guidance_rescale": .5, "rescale_t": 3.} - condition = {key: float(camera[key]) for key in ("camera_angle_x", "distance")} - receipt["initial_generation_camera"] = {**condition, "mesh_scale": initial_mesh_scale} - receipt["camera_distance_policy"] = "preserve_prepared_demo_camera" - - def clear_cache(): - if device.type == "cuda": - torch.cuda.empty_cache() - - def probe(scale): - original_sampler = pipe.sparse_structure_sampler - receipt["dense_probe_calls"] += 1 - try: - pipe.sparse_structure_sampler = copy.deepcopy(original_sampler) - with torch.random.fork_rng(devices=cuda_devices): - torch.default_generator.manual_seed(seed) - if cuda_devices: - with torch.cuda.device(device): - torch.cuda.manual_seed(seed) - with span("geometry.conditioning_probe"): - cond = pipe.get_proj_cond_ss([image], **condition, mesh_scale=scale) - with span("geometry.sampling_probe"): - coords = pipe.sample_sparse_structure(cond, resolution=64, sampler_params=sampler_ss) - del cond - clear_cache() - return coords - finally: - pipe.sparse_structure_sampler = original_sampler - - with span("geometry.scale_search"): - selected, report = calibrate_scale(probe, initial_mesh_scale, target_padding=target_padding, - max_iterations=max_iterations) - condition["mesh_scale"] = selected - receipt.update(scale_search=report, optimized_mesh_scale=selected, final_generation_camera=dict(condition), - uses_dense_check_model=False, final_sparse_structure_resolution=32, resolution=1024) - print(f"[mesh_scale] selected={selected:.6f}, stop={report['stop_reason']}, probes={len(report['probes'])}", flush=True) - torch.manual_seed(seed) - with span("geometry.conditioning_sparse"): - cond = pipe.get_proj_cond_ss([image], **condition) - with span("geometry.sampling_sparse"): - coords = pipe.sample_sparse_structure(cond, resolution=32, sampler_params=sampler_ss) - del cond - clear_cache() - with span("geometry.conditioning_512"): - cond = pipe.get_proj_cond_shape(pipe.image_cond_model_shape_512, [image], coords, **condition) - with span("geometry.sampling_512"): - slat = pipe.sample_shape_slat(cond, pipe.models["shape_slat_flow_model_512"], coords, sampler_shape) - del cond, coords - clear_cache() - with span("geometry.upsample"): - decoder = pipe.models["shape_slat_decoder"] - decoder.to(device) - decoder.low_vram = True - try: - hr_coords = decoder.upsample(slat, upsample_times=4) - finally: - decoder.cpu() - decoder.low_vram = False - # Match native 1024-cascade projection-grid quantization. - coords = torch.cat([hr_coords[:, :1], ((hr_coords[:, 1:] + .5) / 512 * 63).round().int()], dim=1).unique(dim=0) - del slat, hr_coords - clear_cache() - with span("geometry.conditioning_1024"): - cond = pipe.get_proj_cond_shape(pipe.image_cond_model_shape_1024, [image], coords, - grid_resolution_override=64, **condition) - with span("geometry.sampling_1024"): - slat = pipe.sample_shape_slat(cond, pipe.models["shape_slat_flow_model_1024"], coords, sampler_shape) - del cond, coords - clear_cache() - with span("geometry.decode"): - meshes, _ = pipe.decode_shape_slat(slat, 1024) - if len(meshes) != 1 or not meshes[0].vertices.numel() or not meshes[0].faces.numel(): - raise RuntimeError("geometry decoder returned an empty mesh") - meshes[0].fill_holes() - mesh = meshes[0] - vertices, faces = mesh.vertices.detach().float().cpu(), mesh.faces.detach().long().cpu() - if (vertices.ndim != 2 or vertices.shape[1] != 3 or not torch.isfinite(vertices).all() - or float(vertices.abs().max()) > .5001 or faces.ndim != 2 or faces.shape[1] != 3 - or faces.min() < 0 or faces.max() >= len(vertices)): - raise ValueError("native geometry decoder returned invalid canonical mesh") - if condition["distance"] * selected <= float(vertices[:, 2].max()): - raise ValueError("optimized camera lies inside or behind the geometry") - final_camera = {**camera, **condition, - "moge_initial_camera": {key: camera[key] for key in ("camera_angle_x", "distance", "mesh_scale") if key in camera}, - "geometry_scale_optimization": {"initial_mesh_scale": initial_mesh_scale, - "optimized_mesh_scale": selected, - "effective_camera_distance": condition["distance"] * selected, - "objective": "dense_grid_boundary_padding"}} - return {"verts": vertices, "faces": faces, "resolution": 1024}, final_camera - - -def run(package: Path, model_root: Path, input_dir: Path, output_dir: Path, *, seed: int = 42, - initial_mesh_scale: float = 1.0, target_padding: int = 3, max_iterations: int = 2, pipeline=None): - import torch - - for name in ("geometry.pt", "camera.json", "geometry_receipt.json"): - if (output_dir / name).exists(): - raise FileExistsError(f"geometry output already exists: {output_dir / name}") - if not torch.cuda.is_available(): - raise RuntimeError("geometry generation requires CUDA") - from supratex.inference.runtime import configure_inference_runtime - runtime = configure_inference_runtime() - camera = json.loads((input_dir / "camera.json").read_text(encoding="utf-8")) - if any(key in camera for key in ("extrinsic", "intrinsic", "projection_intrinsic")): - raise ValueError("geometry generation requires a canonical scalar camera") - sys.path.insert(0, str(package.resolve())) - native = (package / "pixal3d/pipelines/pixal3d_image_to_3d.py").is_file() - kernel_config = None - if native: - from supratex.inference.kernel_config import configure_sparse_kernels - kernel_config = configure_sparse_kernels() - module = importlib.import_module("pixal3d.pipelines.pixal3d_image_to_3d" if native else "pixal3dpipeline2stage") - if not Path(module.__file__).resolve().is_relative_to(package.resolve()): - raise RuntimeError("geometry pipeline was imported outside the selected source package") - receipt = {"schema": "supratex.geometry.v2", "generator": "Pixal3D paper two-stage", - "runtime": runtime, - "source_package": str(package.resolve()), "model_root": str(model_root.resolve()), - "camera_estimator": camera.get("camera_estimator", "supplied camera"), - "initial_mesh_scale": initial_mesh_scale, "target_padding": target_padding, - "max_iterations": max_iterations, "seed": seed, "dense_probe_calls": 0, - "objective": "dense_grid_boundary_padding", "geometry_frame": "canonical [-0.5,0.5]", - "scope": "paper geometry generation with prepared demo camera and bounded main-model scale calibration"} - if native: - from PIL import Image - from supratex.frontend.geometry_acceleration import NativeGeometryAcceleration - receipt.update(generator="Pixal3D native geometry cascade", - sparse_kernel_config=kernel_config, - scope="native geometry generation with prepared demo camera and bounded main-model scale calibration") - if pipeline is None: - with span("geometry.model_load"): - pipeline = build_native_pipeline(model_root) - else: - receipt['model_loading'] = 'startup_pool' - with span("geometry.inference"), Image.open(input_dir / "condition_rgb.png") as image, torch.inference_mode(), NativeGeometryAcceleration(pipeline) as acceleration: - geometry, final_camera = infer_native(pipeline, image.convert("RGB"), camera, receipt, seed=seed, - initial_mesh_scale=initial_mesh_scale, target_padding=target_padding, - max_iterations=max_iterations) - receipt["conditioning_acceleration"] = dict(acceleration.stats) - else: - pipeline_class = traced_pipeline(module.Pixal3DPipeline2Stage, camera, receipt) - with span("geometry.model_load"): - pipeline = pipeline_class.from_pretrained(ckpt_dir=str(model_root), repo_id=None, - use_moge=False, use_dense_check=False) - receipt["uses_dense_check_model"] = bool(pipeline.has_dense_check) - with span("geometry.inference"), torch.inference_mode(): - mesh = pipeline.infer(image=str(input_dir / "condition_rgb.png"), dense_seed=seed, - sparse_seed=seed, mesh_scale=initial_mesh_scale, - optimize_mesh_scale=True, target_padding=target_padding, - max_optim_iterations=max_iterations) - geometry, final_camera = canonical_handoff(mesh, camera, receipt) - with span("geometry.output_write"): - receipt["vertices"] = int(geometry["verts"].shape[0]) - receipt["faces"] = int(geometry["faces"].shape[0]) - output_dir.mkdir(parents=True, exist_ok=True) - torch.save(geometry, output_dir / "geometry.pt") - (output_dir / "camera.json").write_text(json.dumps(final_camera, indent=2) + "\n", encoding="utf-8") - # Publish the receipt last, after the complete final mesh/camera handoff. - (output_dir / "geometry_receipt.json").write_text(json.dumps(receipt, indent=2) + "\n", encoding="utf-8") - - -def main(): - parser = argparse.ArgumentParser(description=__doc__) - for name in ("package", "model-root", "input-dir", "output-dir"): - parser.add_argument(f"--{name}", type=Path, required=True) - parser.add_argument("--seed", type=int, default=42) - parser.add_argument("--initial-mesh-scale", type=float, default=1.0) - parser.add_argument("--target-padding", type=int, default=3) - parser.add_argument("--max-iterations", type=int, default=2) - args = parser.parse_args() - run(args.package, args.model_root, args.input_dir, args.output_dir, seed=args.seed, - initial_mesh_scale=args.initial_mesh_scale, target_padding=args.target_padding, - max_iterations=args.max_iterations) - - -if __name__ == "__main__": - main() +"""Isolated adapter for native and paper Pixal3D geometry generators. + +The prepared demo camera owns FOV and distance. This adapter calibrates scale +using the main coarse geometry model, then publishes the final mesh/camera in +the SUPRA-Tex canonical frame. +""" +from __future__ import annotations + +import argparse +import copy +import importlib +import json +import math +from pathlib import Path +import sys + + +def calibrate_scale(probe, initial_scale: float = 1.0, *, target_padding: int = 3, + max_iterations: int = 2) -> tuple[float, dict]: + """Keep the closest safe occupancy, with a bounded number of scale updates. + + ``probe(scale)`` returns single-sample [N,4] decoded grid indices. Invalid + occupancy ends the search; inference/runtime errors propagate to the caller. + """ + import torch + + # geometry.run validates these options before launching this runner. + max_iterations = min(max_iterations, 2) + lower, upper = max(1, target_padding - 1), target_padding + 1 + scale = best = float(initial_scale) + best_padding, best_error = None, math.inf + report = {"algorithm": "bounded_main_dense_padding", "initial_mesh_scale": scale, + "target_padding": target_padding, "padding_tolerance": [lower, upper], + "scale_bounds": [scale / 2, scale * 2], "max_iterations": max_iterations, "probes": []} + reason = "max_iterations" + for iteration in range(max_iterations + 1): + coords = probe(scale) + valid = (isinstance(coords, torch.Tensor) and coords.ndim == 2 and coords.shape[1] == 4 + and coords.shape[0] > 0 and not coords.is_complex()) + if valid: + valid = bool(torch.isfinite(coords).all() and (coords[:, 0] == 0).all() + and (coords[:, 1:] >= 0).all() and (coords[:, 1:] <= 63).all() + and (coords == coords.float().round()).all()) + state = {"mesh_scale": scale, "valid": valid} + report["probes"].append(state) + if not valid: + reason = "invalid_occupancy" + break + xyz = coords[:, 1:].float() + padding = min(float(xyz.min()), 63 - float(xyz.max())) + offset = float((xyz - 31.5).abs().max()) + safe, error = padding >= lower, abs(padding - target_padding) + state.update(padding=padding, max_offset=offset, occupied_voxels=len(coords), safe=safe) + if safe and error < best_error: + best, best_padding, best_error = scale, padding, error + if lower <= padding <= upper: + reason = "target_reached" + break + if iteration == max_iterations: + break + proposed = scale * (31.5 - target_padding) / offset + if not math.isfinite(proposed) or not report["scale_bounds"][0] <= proposed <= report["scale_bounds"][1]: + reason = "scale_out_of_bounds" + break + if math.isclose(proposed, scale, rel_tol=1e-6): + reason = "scale_stalled" + break + scale = proposed + report.update(selected_mesh_scale=best, selected_padding=best_padding, + converged=best_padding is not None and lower <= best_padding <= upper, + used_baseline=best == initial_scale, stop_reason=reason) + return best, report + + +def traced_pipeline(base_class, camera: dict, receipt: dict, *, scale_search: bool = False): + """Customize scale calibration while keeping the upstream generation path.""" + fov = float(camera["camera_angle_x"]) + if not math.isfinite(fov) or not 0 < fov < math.pi: + raise ValueError("camera_angle_x must be finite and between zero and pi") + distance = float(camera["distance"]) + if not math.isfinite(distance) or distance <= 0: + raise ValueError("camera distance must be finite and positive") + + class GeometryPipeline(base_class): + _optimizing_scale = False + + def estimate_fov(self, image_path): + return fov + + def _optimize_mesh_scale(self, image_tensor, camera_angle_x_tensor, distance_tensor, + initial_mesh_scale, dense_steps, dense_guidance_scale, dense_seed, + dense_threshold, target_padding=3, max_iterations=2): + import torch + + device = image_tensor.device + cuda_devices = [device.index if device.index is not None else torch.cuda.current_device()] if device.type == "cuda" else [] + receipt["upstream_initial_generation_camera"] = { + "camera_angle_x": float(camera_angle_x_tensor.item()), + "distance": float(distance_tensor.item()), "mesh_scale": float(initial_mesh_scale)} + # The upstream infer() re-derives distance using its 518px/padding + # convention. Mutate its shared camera tensor before any probe so + # final dense and both sparse stages retain the prepared demo camera. + distance_tensor.fill_(distance) + receipt["camera_distance_policy"] = "preserve_prepared_demo_camera" + receipt["initial_generation_camera"] = { + "camera_angle_x": float(camera_angle_x_tensor.item()), + "distance": float(distance_tensor.item()), "mesh_scale": float(initial_mesh_scale)} + + def probe(scale): + receipt["dense_probe_calls"] = receipt.get("dense_probe_calls", 0) + 1 + original_scheduler = self.dense_scheduler + try: + # Probe state must not leak into another probe or final generation. + self.dense_scheduler = copy.deepcopy(original_scheduler) + with torch.random.fork_rng(devices=cuda_devices): + latents = self.infer_dense( + image_tensor, camera_angle_x_tensor, distance_tensor, + torch.tensor([scale], device=device, dtype=torch.float32), + dense_steps, dense_guidance_scale, dense_seed) + with torch.autocast(device.type, dtype=torch.float16, enabled=device.type == "cuda"): + return self.dense_vae.decode_mesh(latents, mc_threshold=dense_threshold, return_index=True)[0] + finally: + self.dense_scheduler = original_scheduler + + if not scale_search: + result, report = float(initial_mesh_scale), {"enabled": False, "probes": [], "stop_reason": "disabled"} + else: + self._optimizing_scale = True + try: + result, report = calibrate_scale(probe, initial_mesh_scale, target_padding=target_padding, + max_iterations=max_iterations) + finally: + self._optimizing_scale = False + receipt["scale_search"] = report + receipt["optimized_mesh_scale"] = float(result) + print(f"[mesh_scale] selected={result:.6f}, stop={report['stop_reason']}, probes={len(report['probes'])}", flush=True) + return result + + def infer_dense(self, image, camera_angle_x, distance, mesh_scale, *args, **kwargs): + if not math.isclose(float(distance.item()), float(camera["distance"]), rel_tol=1e-5): + raise RuntimeError("geometry generation did not retain the prepared demo camera distance") + if not self._optimizing_scale: + receipt["final_generation_camera"] = { + "camera_angle_x": float(camera_angle_x.item()), + "distance": float(distance.item()), + "mesh_scale": float(mesh_scale.item()), + } + return super().infer_dense(image, camera_angle_x, distance, mesh_scale, *args, **kwargs) + + return GeometryPipeline + + +def canonical_handoff(mesh, camera: dict, receipt: dict): + """Undo upstream export rotation and scale while preserving projection.""" + import numpy as np + import torch + + try: + selected = receipt["final_generation_camera"] + scale = float(selected["mesh_scale"]) + optimized = float(receipt["optimized_mesh_scale"]) + except (KeyError, TypeError, ValueError) as error: + raise RuntimeError("geometry source did not expose the two-stage scale/camera hooks") from error + if not math.isfinite(scale) or scale <= 0 or not math.isclose(scale, optimized, rel_tol=1e-5): + raise RuntimeError("final geometry camera does not match optimized mesh_scale") + if not math.isfinite(selected["distance"]) or selected["distance"] <= 0: + raise RuntimeError("final geometry distance must be finite and positive") + if not math.isclose(selected["distance"], float(camera["distance"]), rel_tol=1e-5): + raise RuntimeError("final geometry did not retain the prepared demo camera distance") + if not math.isclose(selected["camera_angle_x"], float(camera["camera_angle_x"]), rel_tol=1e-5): + raise RuntimeError("geometry generation did not retain the supplied MoGe-2 FOV") + # Upstream exports R @ raw_xyz * (0.5 / scale), with R mapping + # (x,y,z) -> (x,-z,y). Canonical points are raw_xyz * 0.5. + rotation = np.array([[1., 0., 0.], [0., 0., -1.], [0., 1., 0.]]) + vertices = np.asarray(mesh.vertices, dtype=np.float32) @ rotation * scale + faces = np.asarray(mesh.faces) + if vertices.ndim != 2 or vertices.shape[1:] != (3,) or not len(vertices) or not np.isfinite(vertices).all(): + raise ValueError("geometry generator returned invalid vertices") + if faces.ndim != 2 or faces.shape[1:] != (3,) or not len(faces) or not np.issubdtype(faces.dtype, np.integer): + raise ValueError("geometry generator returned invalid faces") + if faces.min() < 0 or faces.max() >= len(vertices): + raise ValueError("geometry generator returned out-of-range face indices") + if float(np.abs(vertices).max()) > 0.5001: + raise ValueError("converted geometry exceeds the canonical [-0.5, 0.5] volume") + effective_distance = selected["distance"] * scale + if effective_distance <= float(vertices[:, 2].max()): + raise ValueError("optimized camera lies inside or behind the geometry") + final_camera = dict(camera) + final_camera.update(selected) + final_camera["moge_initial_camera"] = {key: camera[key] for key in + ("camera_angle_x", "distance", "mesh_scale") if key in camera} + final_camera["geometry_scale_optimization"] = { + "initial_mesh_scale": receipt["initial_mesh_scale"], + "optimized_mesh_scale": scale, + "effective_camera_distance": effective_distance, + "objective": "dense_grid_boundary_padding", + } + geometry = {"verts": torch.from_numpy(vertices.copy()).float(), + "faces": torch.from_numpy(faces.astype(np.int64, copy=True)).long()} + return geometry, final_camera + + +def build_native_pipeline(model_root: Path): + """Load only native geometry networks and their projection conditioners.""" + import torch + local_naf = Path(torch.hub.get_dir()) / "valeoai_NAF_main" + for artifact in (local_naf / "hubconf.py", local_naf.parent / "checkpoints/naf_release.pth"): + if not artifact.is_file(): + raise FileNotFoundError(f"native geometry requires the local NAF cache: {artifact}") + from pixal3d.pipelines import Pixal3DImageTo3DPipeline, samplers + from pixal3d.modules.projected_image_feature_extractor import DinoV3ProjFeatureExtractor + from supratex.frontend.geometry_loading import load_native_geometry_model + + class GeometryOnlyPipeline(Pixal3DImageTo3DPipeline): + model_names_to_load = ["sparse_structure_flow_model", "sparse_structure_decoder", + "shape_slat_flow_model_512", "shape_slat_flow_model_1024", "shape_slat_decoder"] + + # Conditions already include foreground removal. Base loading avoids creating + # another background-removal model or any texture networks/samplers. + args = json.loads((model_root / "pipeline.json").read_text(encoding="utf-8"))["args"] + networks = {name: load_native_geometry_model(model_root / checkpoint) + for name, checkpoint in args["models"].items() + if name in GeometryOnlyPipeline.model_names_to_load} + missing = set(GeometryOnlyPipeline.model_names_to_load) - networks.keys() + if missing: + raise ValueError(f"Native geometry config is missing models: {sorted(missing)}") + pipe = GeometryOnlyPipeline(networks) + pipe._pretrained_args = args + for name in ("sparse_structure", "shape_slat"): + spec = args[f"{name}_sampler"] + setattr(pipe, f"{name}_sampler", getattr(samplers, spec["name"])(**spec["args"])) + setattr(pipe, f"{name}_sampler_params", spec["params"]) + pipe.shape_slat_normalization = args["shape_slat_normalization"] + original_load = torch.hub.load + + def local_hub_load(repo_or_dir, model, *args, **kwargs): + if repo_or_dir == "valeoai/NAF": + kwargs.pop("trust_repo", None) + kwargs["source"] = "local" + repo_or_dir = str(local_naf) + return original_load(repo_or_dir, model, *args, **kwargs) + + try: + torch.hub.load = local_hub_load + for name, size, grid, use_naf in (("ss", 512, 16, False), + ("shape_512", 512, 32, True), + ("shape_1024", 1024, 64, True)): + conditioner = DinoV3ProjFeatureExtractor( + model_name=str(model_root / "dino"), image_size=size, grid_resolution=grid, + use_naf_upsample=use_naf, naf_target_size=512).eval() + if use_naf: + conditioner._load_naf() + setattr(pipe, f"image_cond_model_{name}", conditioner) + finally: + torch.hub.load = original_load + pipe._device = torch.device("cuda") + pipe.low_vram = True + return pipe + + +def infer_native(pipe, image, camera: dict, receipt: dict, *, seed: int, + initial_mesh_scale: float, scale_search: bool, target_padding: int, max_iterations: int): + """Optionally calibrate native 64-cell occupancy, then run the Pixal3D shape cascade.""" + import torch + + device = torch.device(pipe.device) + cuda_devices = [device.index if device.index is not None else torch.cuda.current_device()] if device.type == "cuda" else [] + sampler_ss = {"steps": 12, "guidance_strength": 7.5, "guidance_rescale": .7, "rescale_t": 5.} + sampler_shape = {"steps": 12, "guidance_strength": 7.5, "guidance_rescale": .5, "rescale_t": 3.} + condition = {key: float(camera[key]) for key in ("camera_angle_x", "distance")} + receipt["initial_generation_camera"] = {**condition, "mesh_scale": initial_mesh_scale} + receipt["camera_distance_policy"] = "preserve_prepared_demo_camera" + + def clear_cache(): + if device.type == "cuda": + torch.cuda.empty_cache() + + def probe(scale): + original_sampler = pipe.sparse_structure_sampler + receipt["dense_probe_calls"] += 1 + try: + pipe.sparse_structure_sampler = copy.deepcopy(original_sampler) + with torch.random.fork_rng(devices=cuda_devices): + torch.default_generator.manual_seed(seed) + if cuda_devices: + with torch.cuda.device(device): + torch.cuda.manual_seed(seed) + cond = pipe.get_proj_cond_ss([image], **condition, mesh_scale=scale) + coords = pipe.sample_sparse_structure(cond, resolution=64, sampler_params=sampler_ss) + del cond + clear_cache() + return coords + finally: + pipe.sparse_structure_sampler = original_sampler + + if scale_search: + selected, report = calibrate_scale(probe, initial_mesh_scale, target_padding=target_padding, + max_iterations=max_iterations) + else: + # Default: keep Pixal3D's fixed mesh scale without dense probes. + selected, report = initial_mesh_scale, {"enabled": False, "probes": [], "stop_reason": "disabled"} + condition["mesh_scale"] = selected + receipt.update(scale_search=report, optimized_mesh_scale=selected, final_generation_camera=dict(condition), + uses_dense_check_model=False, final_sparse_structure_resolution=32, resolution=1024) + print(f"[mesh_scale] selected={selected:.6f}, stop={report['stop_reason']}, probes={len(report['probes'])}", flush=True) + torch.manual_seed(seed) + cond = pipe.get_proj_cond_ss([image], **condition) + coords = pipe.sample_sparse_structure(cond, resolution=32, sampler_params=sampler_ss) + del cond + clear_cache() + cond = pipe.get_proj_cond_shape(pipe.image_cond_model_shape_512, [image], coords, **condition) + slat = pipe.sample_shape_slat(cond, pipe.models["shape_slat_flow_model_512"], coords, sampler_shape) + del cond, coords + clear_cache() + decoder = pipe.models["shape_slat_decoder"] + decoder.to(device) + decoder.low_vram = True + try: + hr_coords = decoder.upsample(slat, upsample_times=4) + finally: + decoder.cpu() + decoder.low_vram = False + # Match native 1024-cascade projection-grid quantization. + coords = torch.cat([hr_coords[:, :1], ((hr_coords[:, 1:] + .5) / 512 * 63).round().int()], dim=1).unique(dim=0) + del slat, hr_coords + clear_cache() + cond = pipe.get_proj_cond_shape(pipe.image_cond_model_shape_1024, [image], coords, + grid_resolution_override=64, **condition) + slat = pipe.sample_shape_slat(cond, pipe.models["shape_slat_flow_model_1024"], coords, sampler_shape) + del cond, coords + clear_cache() + meshes, _ = pipe.decode_shape_slat(slat, 1024) + if len(meshes) != 1 or not meshes[0].vertices.numel() or not meshes[0].faces.numel(): + raise RuntimeError("geometry decoder returned an empty mesh") + meshes[0].fill_holes() + mesh = meshes[0] + vertices, faces = mesh.vertices.detach().float().cpu(), mesh.faces.detach().long().cpu() + if (vertices.ndim != 2 or vertices.shape[1] != 3 or not torch.isfinite(vertices).all() + or float(vertices.abs().max()) > .5001 or faces.ndim != 2 or faces.shape[1] != 3 + or faces.min() < 0 or faces.max() >= len(vertices)): + raise ValueError("native geometry decoder returned invalid canonical mesh") + if condition["distance"] * selected <= float(vertices[:, 2].max()): + raise ValueError("optimized camera lies inside or behind the geometry") + final_camera = {**camera, **condition, + "moge_initial_camera": {key: camera[key] for key in ("camera_angle_x", "distance", "mesh_scale") if key in camera}, + "geometry_scale_optimization": {"initial_mesh_scale": initial_mesh_scale, + "optimized_mesh_scale": selected, + "effective_camera_distance": condition["distance"] * selected, + "objective": "dense_grid_boundary_padding"}} + return {"verts": vertices, "faces": faces, "resolution": 1024}, final_camera + + +def run(package: Path, model_root: Path, input_dir: Path, output_dir: Path, *, seed: int = 42, + initial_mesh_scale: float = 1.0, scale_search: bool = False, target_padding: int = 3, + max_iterations: int = 2): + import torch + + for name in ("geometry.pt", "camera.json", "geometry_receipt.json"): + if (output_dir / name).exists(): + raise FileExistsError(f"geometry output already exists: {output_dir / name}") + if not torch.cuda.is_available(): + raise RuntimeError("geometry generation requires CUDA") + from supratex.inference.runtime import configure_inference_runtime + runtime = configure_inference_runtime() + camera = json.loads((input_dir / "camera.json").read_text(encoding="utf-8")) + if any(key in camera for key in ("extrinsic", "intrinsic", "projection_intrinsic")): + raise ValueError("geometry generation requires a canonical scalar camera") + sys.path.insert(0, str(package.resolve())) + native = (package / "pixal3d/pipelines/pixal3d_image_to_3d.py").is_file() + kernel_config = None + if native: + from supratex.inference.kernel_config import configure_sparse_kernels + kernel_config = configure_sparse_kernels() + module = importlib.import_module("pixal3d.pipelines.pixal3d_image_to_3d" if native else "pixal3dpipeline2stage") + if not Path(module.__file__).resolve().is_relative_to(package.resolve()): + raise RuntimeError("geometry pipeline was imported outside the selected source package") + receipt = {"schema": "supratex.geometry.v2", "generator": "Pixal3D paper two-stage", + "runtime": runtime, + "source_package": str(package.resolve()), "model_root": str(model_root.resolve()), + "camera_estimator": camera.get("camera_estimator", "supplied camera"), + "initial_mesh_scale": initial_mesh_scale, "scale_search": scale_search, + "target_padding": target_padding, + "max_iterations": max_iterations, "seed": seed, "dense_probe_calls": 0, + "objective": "dense_grid_boundary_padding", "geometry_frame": "canonical [-0.5,0.5]", + "scope": "paper geometry generation with prepared demo camera; optional bounded scale calibration"} + if native: + from PIL import Image + from supratex.frontend.geometry_acceleration import NativeGeometryAcceleration + receipt.update(generator="Pixal3D native geometry cascade", + sparse_kernel_config=kernel_config, + scope="native geometry generation with prepared demo camera; optional bounded scale calibration") + pipeline = build_native_pipeline(model_root) + with Image.open(input_dir / "condition_rgb.png") as image, torch.inference_mode(), NativeGeometryAcceleration(pipeline) as acceleration: + geometry, final_camera = infer_native(pipeline, image.convert("RGB"), camera, receipt, seed=seed, + initial_mesh_scale=initial_mesh_scale, scale_search=scale_search, + target_padding=target_padding, max_iterations=max_iterations) + receipt["conditioning_acceleration"] = dict(acceleration.stats) + else: + pipeline_class = traced_pipeline(module.Pixal3DPipeline2Stage, camera, receipt, scale_search=scale_search) + pipeline = pipeline_class.from_pretrained(ckpt_dir=str(model_root), repo_id=None, + use_moge=False, use_dense_check=False) + receipt["uses_dense_check_model"] = bool(pipeline.has_dense_check) + with torch.inference_mode(): + mesh = pipeline.infer(image=str(input_dir / "condition_rgb.png"), dense_seed=seed, + sparse_seed=seed, mesh_scale=initial_mesh_scale, + optimize_mesh_scale=True, target_padding=target_padding, + max_optim_iterations=max_iterations) + geometry, final_camera = canonical_handoff(mesh, camera, receipt) + receipt["vertices"] = int(geometry["verts"].shape[0]) + receipt["faces"] = int(geometry["faces"].shape[0]) + output_dir.mkdir(parents=True, exist_ok=True) + torch.save(geometry, output_dir / "geometry.pt") + (output_dir / "camera.json").write_text(json.dumps(final_camera, indent=2) + "\n", encoding="utf-8") + # Publish the receipt last, after the complete final mesh/camera handoff. + (output_dir / "geometry_receipt.json").write_text(json.dumps(receipt, indent=2) + "\n", encoding="utf-8") + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + for name in ("package", "model-root", "input-dir", "output-dir"): + parser.add_argument(f"--{name}", type=Path, required=True) + parser.add_argument("--seed", type=int, default=42) + parser.add_argument("--initial-mesh-scale", type=float, default=1.0) + parser.add_argument("--scale-search", action="store_true") + parser.add_argument("--target-padding", type=int, default=3) + parser.add_argument("--max-iterations", type=int, default=2) + args = parser.parse_args() + run(args.package, args.model_root, args.input_dir, args.output_dir, seed=args.seed, + initial_mesh_scale=args.initial_mesh_scale, scale_search=args.scale_search, + target_padding=args.target_padding, max_iterations=args.max_iterations) + + +if __name__ == "__main__": + main() diff --git a/supratex/frontend/material_model/runtime/pbr/conditioning.py b/supratex/frontend/material_model/runtime/pbr/conditioning.py new file mode 100644 index 0000000000000000000000000000000000000000..78bd228b6b5d3e995dc402a7c34e530a2fe78022 --- /dev/null +++ b/supratex/frontend/material_model/runtime/pbr/conditioning.py @@ -0,0 +1,178 @@ +# Adapted for SUPRA-Tex distribution (2026): packaged local inference runtime. +# Upstream FireRed portions remain under Apache-2.0; see THIRD_PARTY_NOTICES.md. +# Copyright (c) 2025 FireRed-Image-Edit. All rights reserved. +"""Image preprocessing and prompt conditioning for material inference.""" +from __future__ import annotations + +import math +from typing import List, Tuple + +from PIL import Image +import torch +from torchvision import transforms + + +DEFAULT_SYSTEM_PROMPT = "Describe the key features of the input image (color, shape, size, texture, objects, background), then explain how the user's text instruction should alter or modify the image. Generate a new image that meets the user's requirements while maintaining consistency with the original input where appropriate." + +DEFAULT_SYSTEM_PROMPT_T2I = "Describe the image by detailing the color, shape, size, texture, quantity, text, spatial relationships of the objects and background:" + +ASPECT_RATIO_512 = { + '0.25': [256.0, 1024.0], '0.26': [256.0, 992.0], '0.27': [256.0, 960.0], '0.28': [256.0, 928.0], + '0.32': [288.0, 896.0], '0.33': [288.0, 864.0], '0.35': [288.0, 832.0], '0.4': [320.0, 800.0], + '0.42': [320.0, 768.0], '0.48': [352.0, 736.0], '0.5': [352.0, 704.0], '0.52': [352.0, 672.0], + '0.57': [384.0, 672.0], '0.6': [384.0, 640.0], '0.68': [416.0, 608.0], '0.72': [416.0, 576.0], + '0.78': [448.0, 576.0], '0.82': [448.0, 544.0], '0.88': [480.0, 544.0], '0.94': [480.0, 512.0], + '1.0': [512.0, 512.0], '1.07': [512.0, 480.0], '1.13': [544.0, 480.0], '1.21': [544.0, 448.0], + '1.29': [576.0, 448.0], '1.38': [576.0, 416.0], '1.46': [608.0, 416.0], '1.67': [640.0, 384.0], + '1.75': [672.0, 384.0], '2.0': [704.0, 352.0], '2.09': [736.0, 352.0], '2.4': [768.0, 320.0], + '2.5': [800.0, 320.0], '2.89': [832.0, 288.0], '3.0': [864.0, 288.0], '3.11': [896.0, 288.0], + '3.62': [928.0, 256.0], '3.75': [960.0, 256.0], '3.88': [992.0, 256.0], '4.0': [1024.0, 256.0] +} + +CONDITION_IMAGE_SIZE = 384 * 384 + + +def calculate_shift( + image_seq_len, + base_seq_len: int = 256, + max_seq_len: int = 4096, + base_shift: float = 0.5, + max_shift: float = 1.15, +): + """根据图像序列长度计算 flow 的 shift 系数（线性插值）。""" + m = (max_shift - base_shift) / (max_seq_len - base_seq_len) + b = base_shift - m * base_seq_len + mu = image_seq_len * m + b + return mu + +def get_embedding(source_exist, instructions, source_images, processor, text_encoder): + def _extract_valid_hidden_states( + hidden_states: torch.Tensor, + attention_mask: torch.Tensor, + ): + batch_size = hidden_states.shape[0] + valid_hidden_states = [] + for i in range(batch_size): + valid_indices = attention_mask[i].nonzero(as_tuple=True)[0] + valid_hidden_states.append(hidden_states[i][valid_indices]) + return valid_hidden_states + + with torch.no_grad(): + conversations = [] + if not source_exist: + source_images = [[] for _ in range(len(instructions))] + for instruction, source_images_sample in zip(instructions, source_images): + system_prompt = DEFAULT_SYSTEM_PROMPT_T2I if not source_exist else DEFAULT_SYSTEM_PROMPT + messages = [ + {"role": "system", "content": [{"type": "text", "text": system_prompt}]}, + {"role": "user", "content": []}, + ] + for condition_image in source_images_sample: + messages[1]["content"].append({"type": "image", "image": condition_image}) + messages[1]["content"].append({"type": "text", "text": instruction}) + conversations.append(messages) + + batch_inputs = processor.apply_chat_template( + conversations, + add_generation_prompt=True, + padding=True, + tokenize=True, + return_tensors="pt", + return_dict=True, + ).to(text_encoder.device) + + encoder_hidden_states = text_encoder( + input_ids=batch_inputs["input_ids"], + attention_mask=batch_inputs["attention_mask"], + pixel_values=batch_inputs["pixel_values"] if source_exist else None, + image_grid_thw=batch_inputs["image_grid_thw"] if source_exist else None, + output_hidden_states=True, + ) + + hidden_states = encoder_hidden_states.hidden_states[-1] + split_hidden_states = _extract_valid_hidden_states(hidden_states, batch_inputs["attention_mask"]) + start_idx = 64 if source_exist else 34 + split_hidden_states = [e[start_idx:] for e in split_hidden_states] + attn_mask_list = [torch.ones(e.size(0), dtype=torch.long, device=e.device) for e in split_hidden_states] + max_seq_len = max([e.size(0) for e in split_hidden_states]) + prompt_embeds = torch.stack( + [torch.cat([u, u.new_zeros(max_seq_len - u.size(0), u.size(1))]) for u in split_hidden_states] + ) + encoder_attention_mask = torch.stack( + [torch.cat([u, u.new_zeros(max_seq_len - u.size(0))]) for u in attn_mask_list] + ) + + return prompt_embeds, encoder_attention_mask + +def get_closest_ratio(height: float, width: float, ratios: dict = ASPECT_RATIO_512) -> Tuple[List[float], float]: + aspect_ratio = height / width + closest_ratio = min(ratios.keys(), key=lambda r: abs(float(r) - aspect_ratio)) + return ratios[closest_ratio], float(closest_ratio) + +def calculate_dimensions(target_area: float, ratio: float) -> Tuple[int, int]: + width = math.sqrt(target_area * ratio) + height = width / ratio + width = round(width / 32) * 32 + height = round(height / 32) * 32 + return int(width), int(height) + +def resize_source_images_for_condition( + source_images_list: List[List[Image.Image]], + ref_height: int, + ref_width: int, + image_sample_size: int, + condition_image_size: int = CONDITION_IMAGE_SIZE, +) -> List[List[Image.Image]]: + """ + 对送入 text_encoder 的 source 图像做 Resize + CenterCrop，再缩放到 condition_image_size 对应尺寸。 + 参考尺寸为 (ref_height, ref_width)。source_images_list: 每样本多张 source 图的 list of list of PIL。 + 返回相同结构的 list of list of PIL Image。 + """ + if not source_images_list: + return source_images_list + return [ + apply_condition_transform_to_images(sample, ref_height, ref_width, image_sample_size, condition_image_size) + for sample in source_images_list + ] + +def apply_condition_transform_to_images( + images: List[Image.Image], + ref_height: int, + ref_width: int, + image_sample_size: int, + condition_image_size: int = CONDITION_IMAGE_SIZE, +) -> List[Image.Image]: + """ + 对多张图像应用统一的 Resize+CenterCrop+最终缩放。参考尺寸 (ref_height, ref_width)。 + 返回与 images 等长的 list of PIL Image。 + """ + if not images: + return [] + h, w = float(ref_height), float(ref_width) + aspect_ratio_sample_size = { + k: [x / 512 * image_sample_size for x in ASPECT_RATIO_512[k]] + for k in ASPECT_RATIO_512 + } + closest_size, _ = get_closest_ratio(h, w, ratios=aspect_ratio_sample_size) + closest_size = [int(x / 16) * 16 for x in closest_size] + closest_size = list(map(int, closest_size)) + if closest_size[0] / h > closest_size[1] / w: + resize_size = (closest_size[0], int(w * closest_size[0] / h)) + else: + resize_size = (int(h * closest_size[1] / w), closest_size[1]) + + source_transform = transforms.Compose([ + transforms.Resize(resize_size, interpolation=transforms.InterpolationMode.BILINEAR), + transforms.CenterCrop(closest_size), + ]) + + result = [] + for im in images: + if im.mode != "RGB": + im = im.convert("RGB") + cropped = source_transform(im) + width_c, height_c = cropped.size + aspect_ratio = width_c / height_c + new_width, new_height = calculate_dimensions(condition_image_size, aspect_ratio) + result.append(cropped.resize((new_width, new_height), Image.Resampling.LANCZOS)) + return result diff --git a/supratex/frontend/material_model/runtime/pbr/model_loading.py b/supratex/frontend/material_model/runtime/pbr/model_loading.py new file mode 100644 index 0000000000000000000000000000000000000000..9966fea60dfb1e5db40322ebf9b7ee5cf9741853 --- /dev/null +++ b/supratex/frontend/material_model/runtime/pbr/model_loading.py @@ -0,0 +1,136 @@ +# Adapted for SUPRA-Tex distribution (2026): packaged local inference runtime. +# Upstream FireRed portions remain under Apache-2.0; see THIRD_PARTY_NOTICES.md. +# Copyright (c) 2025 FireRed-Image-Edit. All rights reserved. +"""Strict checkpoint loading for the FireRed material inference model.""" +from __future__ import annotations + +import logging +from pathlib import Path + +import torch + +logger = logging.getLogger(__name__) +JOINT_MODULE_PREFIX = "modality_joint_blocks." + + +def _require_exact_loading_info(loading_info, source: str) -> None: + problems = { + name: list(loading_info.get(name, [])) + for name in ("missing_keys", "unexpected_keys", "mismatched_keys", "error_msgs") + if loading_info.get(name) + } + if problems: + summary = ", ".join(f"{name}={len(values)}" for name, values in problems.items()) + raise RuntimeError(f"{source} transformer load is not exact: {summary}") + +def _load_base_transformer(args, weight_dtype, *, use_layer3d_rope: bool = False): + from .qwen_pbr_factorized import ( + PRODUCTION_JOINT_HEADS, + PRODUCTION_JOINT_MODALITIES, + PRODUCTION_JOINT_RANK, + QwenImagePBRFactorizedTransformer2DModel, + ) + + loaded = QwenImagePBRFactorizedTransformer2DModel.from_pretrained( + args.pretrained_model_name_or_path, + subfolder="transformer", + torch_dtype=weight_dtype, + joint_sites=(), + joint_modalities=PRODUCTION_JOINT_MODALITIES, + joint_rank=PRODUCTION_JOINT_RANK, + joint_heads=PRODUCTION_JOINT_HEADS, + use_layer3d_rope=use_layer3d_rope, + output_loading_info=True, + ) + if not isinstance(loaded, tuple) or len(loaded) != 2: + raise RuntimeError("custom transformer loader did not return loading information") + transformer, loading_info = loaded + _require_exact_loading_info(loading_info, "base") + if tuple(transformer.joint_sites): + raise RuntimeError("base transformer must load without joint modules") + return transformer + +def _resolve_full_transformer_checkpoint(path: str | Path) -> Path: + checkpoint = Path(path) + if checkpoint.is_dir(): + checkpoint = checkpoint / "diffusion_pytorch_model.safetensors" + if not checkpoint.is_file(): + raise FileNotFoundError(f"full transformer checkpoint not found: {checkpoint}") + return checkpoint + +def _load_checkpoint_state(path: Path): + if path.suffix == ".safetensors": + from safetensors.torch import load_file + + state = load_file(str(path), device="cpu") + else: + state = torch.load(str(path), map_location="cpu", weights_only=True) + if isinstance(state, dict) and "state_dict" in state: + state = state["state_dict"] + if not isinstance(state, dict): + raise ValueError(f"transformer checkpoint must contain a state dict: {path}") + return state + +def _load_exact_full_transformer( + transformer, + checkpoint_path: str | Path, + *, + allow_joint_sidecar: bool = False, +) -> None: + path = _resolve_full_transformer_checkpoint(checkpoint_path) + state = _load_checkpoint_state(path) + if allow_joint_sidecar: + state = { + name: value + for name, value in state.items() + if not name.startswith(JOINT_MODULE_PREFIX) + } + result = transformer.load_state_dict(state, strict=False) + if result.missing_keys or result.unexpected_keys: + raise RuntimeError( + "full transformer checkpoint load is not exact: " + f"missing_keys={len(result.missing_keys)}, " + f"unexpected_keys={len(result.unexpected_keys)}" + ) + logger.info( "Loaded exact full transformer checkpoint: %s", path) + +def _resolve_joint_adapter_checkpoint(path: str | Path) -> Path: + checkpoint = Path(path) + if checkpoint.is_dir(): + checkpoint = checkpoint / "joint_adapter.safetensors" + if not checkpoint.is_file(): + raise FileNotFoundError(f"joint adapter checkpoint not found: {checkpoint}") + return checkpoint + +def _load_exact_joint_adapter(transformer, checkpoint_path: str | Path) -> None: + path = _resolve_joint_adapter_checkpoint(checkpoint_path) + state = _load_checkpoint_state(path) + expected = { + name: parameter + for name, parameter in transformer.named_parameters() + if name.startswith(JOINT_MODULE_PREFIX) + } + actual_keys = set(state) + expected_keys = set(expected) + if actual_keys != expected_keys: + missing = sorted(expected_keys - actual_keys) + unexpected = sorted(actual_keys - expected_keys) + raise RuntimeError( + "joint adapter checkpoint is not exact: " + f"missing={missing[:3]}, unexpected={unexpected[:3]}" + ) + mismatched = [ + name + for name, parameter in expected.items() + if not isinstance(state[name], torch.Tensor) + or tuple(state[name].shape) != tuple(parameter.shape) + ] + if mismatched: + raise RuntimeError(f"joint adapter checkpoint shape mismatch: {mismatched[:3]}") + result = transformer.load_state_dict(state, strict=False) + if result.unexpected_keys: + raise RuntimeError( + "joint adapter checkpoint has unexpected keys after load: " + f"{list(result.unexpected_keys)[:3]}" + ) + logger.info( "Loaded exact joint adapter checkpoint: %s", path) diff --git a/supratex/frontend/material_model/runtime/pbr/predict_joint_pbr.py b/supratex/frontend/material_model/runtime/pbr/predict_joint_pbr.py index 143e2735d7b5888e1a1a64e1bd5fe68c53443d63..323949f7ff4b775d1bccbc0a83a8d8dfb58ee1d2 100644 --- a/supratex/frontend/material_model/runtime/pbr/predict_joint_pbr.py +++ b/supratex/frontend/material_model/runtime/pbr/predict_joint_pbr.py @@ -1,595 +1,592 @@ -#!/usr/bin/env python -# Adapted for SUPRA-Tex distribution (2026): packaged local inference runtime. -# Upstream FireRed portions remain under Apache-2.0; see LICENSE. -"""Joint albedo/MRA/normal inference for the factorized PBR checkpoint. - -The stock QwenImageEditPlusPipeline rejects prompt batches larger than one, so -it cannot exercise modality-axis attention. This runner keeps the accepted -pipeline preprocessing/scheduler/decoder semantics, but folds one physical -item into the fixed ``albedo, mra, normal`` batch and advances all three streams -together. Joint checkpoints are trained only at sigma=1, therefore this -runner intentionally accepts one-step, CFG-free inference. The checkpoint -contract selects the recorded 512/1024 resolution for v2/v3. -""" - -from __future__ import annotations - -import argparse -import importlib -import json -import os -import sys -import time -from pathlib import Path -from types import SimpleNamespace - - -REPO = Path(__file__).absolute().parent.parent -TRAIN_ROOT = REPO / "train" -MODALITIES = ("albedo", "mra", "normal") -JOINT_ROUTES = ("aligned", "self_slot", "blockdiag") - - -def log(message: str) -> None: - print(f"[predict_joint_pbr {time.strftime('%H:%M:%S')}] {message}", flush=True) - - -def output_path(out_dir: Path, item: dict, modality: str) -> Path: - return out_dir / f"{item['key']}_{modality}.png" - - -def load_items(path: Path, limit: int | None) -> list[dict]: - required = { - "key", - "dataset", - "sha", - "asset_id", - "view", - "cond", - "albedo_gt", - "mra_gt", - "normal_gt", - "normal_valid", - "instr_albedo", - "instr_mra", - "instr_normal", - "split", - } - items = [] - seen = set() - with path.open() as stream: - for line_number, raw_line in enumerate(stream, 1): - if not raw_line.strip(): - continue - item = json.loads(raw_line) - missing = required - set(item) - if missing: - raise ValueError(f"{path}:{line_number}: missing fields {sorted(missing)}") - key = item["key"] - if not isinstance(key, str) or not key or Path(key).name != key or key in {".", ".."} or "\\" in key: - raise ValueError(f"{path}:{line_number}: key must be a single safe filename component") - if item["key"] in seen: - raise ValueError(f"{path}:{line_number}: duplicate key {item['key']}") - if item["split"] != "eval": - raise ValueError(f"{path}:{line_number}: expected split=eval") - seen.add(item["key"]) - items.append(item) - if limit is not None: - items = items[:limit] - if not items: - raise ValueError(f"no evaluation items in {path}") - return items - - -def validate_checkpoint(checkpoint: Path, route: str) -> SimpleNamespace: - manifest_path = checkpoint / "checkpoint_complete.json" - contract_path = checkpoint / "joint_contract.json" - run_contract_path = checkpoint.parent / "RUN_CONTRACT.json" - adapter_path = checkpoint / "joint_adapter.safetensors" - full_transformer_path = checkpoint / "diffusion_pytorch_model.safetensors" - for path in (manifest_path, contract_path): - if not path.is_file(): - raise FileNotFoundError(path) - manifest = json.loads(manifest_path.read_text()) - contract = json.loads(contract_path.read_text()) - common_expected = { - "sites": [29, 39, 49, 59], - "modalities": list(MODALITIES), - "rank": 256, - "heads": 4, - "training_route": "aligned", - } - schema = contract.get("schema") - if schema == "pbr_joint_checkpoint_v1": - expected = { - **common_expected, - "schema": schema, - "trainable_tensor_count": 60, - "trainable_parameter_count": 7_385_088, - } - input_size = 512 - elif schema in ("pbr_joint_checkpoint_v2", "pbr_joint_checkpoint_v3"): - if not run_contract_path.is_file(): - raise FileNotFoundError(run_contract_path) - run_contract = json.loads(run_contract_path.read_text()) - input_size = run_contract.get("resolution") - if input_size not in (512, 1024): - raise ValueError( - f"unsupported joint-v2/v3 run resolution: {input_size!r}" - ) - expected = { - **common_expected, - "schema": schema, - "joint_tensor_count": 60, - "joint_parameter_count": 7_385_088, - "trainable_modules": ["modality_joint_blocks."], - } - backbone_low_lr = [ - f"transformer_blocks.{index}." for index in range(29, 50) - ] - if schema == "pbr_joint_checkpoint_v2": - expected.update( - low_lr_modules=backbone_low_lr, - low_lr_ratio=0.0025, - ) - if contract.get("fresh_joint_w_up_init_std") not in (0.0, 1e-4): - expected["fresh_joint_w_up_init_std"] = "0.0 or 1e-4" - residual_mode = "raw" - else: - low_lr_modules = contract.get("low_lr_modules") - if low_lr_modules == []: - expected.update(low_lr_modules=[], low_lr_ratio=None) - elif low_lr_modules == backbone_low_lr: - expected.update( - low_lr_modules=backbone_low_lr, - low_lr_ratio=0.0025, - ) - else: - expected["low_lr_modules"] = "[] or transformer blocks 29 through 49" - residual_mode = contract.get("residual_mode") - if residual_mode != "rms_scaled": - expected["residual_mode"] = "rms_scaled" - if contract.get("fresh_joint_w_up_init_std") not in (0.0, 3e-3): - expected["fresh_joint_w_up_init_std"] = "0.0 or 0.003" - if contract.get("source_joint_adapter_action") not in ( - "fresh", "loaded", "reset" - ): - expected["source_joint_adapter_action"] = "fresh, loaded, or reset" - else: - raise ValueError(f"unsupported joint checkpoint schema: {schema!r}") - bad = { - key: (contract.get(key), value) - for key, value in expected.items() - if contract.get(key) != value - } - if bad: - raise ValueError(f"joint checkpoint contract mismatch: {bad}") - if manifest.get("joint_contract") != contract: - raise ValueError("checkpoint manifest is not bound to joint_contract.json") - if schema == "pbr_joint_checkpoint_v1": - if (manifest.get("global_step"), manifest.get("world_size"), manifest.get("seed")) \ - != (2000, 40, 42): - raise ValueError(f"joint checkpoint step/world/seed mismatch: {manifest}") - elif ( - manifest.get("format_version") != 2 - or not isinstance(manifest.get("global_step"), int) - or manifest["global_step"] <= 0 - ): - raise ValueError(f"joint-v2/v3 checkpoint manifest mismatch: {manifest}") - if route in JOINT_ROUTES and schema == "pbr_joint_checkpoint_v1": - if not adapter_path.is_file(): - raise FileNotFoundError(adapter_path) - if route in JOINT_ROUTES and schema in ( - "pbr_joint_checkpoint_v2", "pbr_joint_checkpoint_v3" - ): - if not full_transformer_path.is_file(): - raise FileNotFoundError(full_transformer_path) - return SimpleNamespace( - schema=schema, - input_size=input_size, - adapter_path=adapter_path if schema == "pbr_joint_checkpoint_v1" else None, - full_transformer_path=( - full_transformer_path - if schema in ("pbr_joint_checkpoint_v2", "pbr_joint_checkpoint_v3") - else None - ), - residual_mode="raw" if schema != "pbr_joint_checkpoint_v3" else residual_mode, - ) - - -def validate_payloads(items: list[dict]) -> None: - fields = ("cond", "albedo_gt", "mra_gt", "normal_gt", "normal_valid") - for item in items: - missing = [item[field] for field in fields if not Path(item[field]).is_file()] - if missing: - raise FileNotFoundError( - f"{item['key']}: missing payload(s): {missing[:3]}" - ) - - -def load_source_rgb(path: str, size: int): - from PIL import Image - - image = Image.open(path) - if image.mode != "RGB": - rgba = image.convert("RGBA") - background = Image.new("RGBA", rgba.size, (0, 0, 0, 255)) - image = Image.alpha_composite(background, rgba).convert("RGB") - if image.size != (size, size): - image = image.resize((size, size), Image.Resampling.BICUBIC) - return image - - -def _flash_attention_available(module_name: str) -> bool: - try: - module = importlib.import_module(module_name) - except (ImportError, OSError, RuntimeError): - return False - return all( - callable(getattr(module, name, None)) - for name in ("flash_attn_func", "flash_attn_varlen_func") - ) - - -def select_attention_backends(torch, device) -> tuple[str, str]: - """Resolve inference attention without depending on the parent CLI package. - - The image transformer supports FA3 and native PyTorch SDPA. FA2 can still - accelerate the text encoder when the image transformer uses native SDPA. - An explicit SDPA selection applies to both components. - """ - selected = os.environ.get("ATTN_BACKEND", "auto") - if selected == "sdpa": - return "native", "sdpa" - supported = ("auto", "flash_attn_3", "flash_attn", "xformers", "naive", "flash_attn_4") - if selected not in supported: - raise ValueError(f"unsupported ATTN_BACKEND for FireRed: {selected!r}") - - device = torch.device(device) - capability = ( - torch.cuda.get_device_capability(device) - if device.type == "cuda" and torch.cuda.is_available() - else (0, 0) - ) - use_fa3 = ( - selected in ("auto", "flash_attn_3") - and capability[0] == 9 - and _flash_attention_available("flash_attn_interface") - ) - if selected == "flash_attn_3" and not use_fa3: - raise RuntimeError( - "ATTN_BACKEND=flash_attn_3 requires Hopper (SM90) and an importable " - "flash_attn_interface extension; select sdpa for native attention" - ) - use_fa2 = capability[0] in (8, 9) and _flash_attention_available("flash_attn") - if selected == "flash_attn" and not use_fa2: - raise RuntimeError( - "ATTN_BACKEND=flash_attn requires an Ampere/Ada/Hopper GPU and an " - "importable flash_attn extension; select sdpa for native attention" - ) - return "fa3" if use_fa3 else "native", "flash_attention_2" if use_fa2 else "sdpa" - - -def load_pipeline(args, device): - import torch - from diffusers import FlowMatchEulerDiscreteScheduler, QwenImageEditPlusPipeline - - image_backend, text_backend = select_attention_backends(torch, device) - sys.path.insert(0, str(TRAIN_ROOT)) - from src.joint_model_provider import ( - _load_base_transformer, - _load_exact_full_transformer, - _load_exact_joint_adapter, - ) - from src.models import ( - PBR_ATTENTION_FA3, - PBR_ATTENTION_NATIVE, - PRODUCTION_JOINT_HEADS, - PRODUCTION_JOINT_MODALITIES, - PRODUCTION_JOINT_RANK, - PRODUCTION_JOINT_SITES, - ) - - model_args = SimpleNamespace(pretrained_model_name_or_path=str(args.base_model)) - transformer = _load_base_transformer(model_args, torch.bfloat16) - if args.route == "base" or args.joint_spec.schema == "pbr_joint_checkpoint_v1": - _load_exact_full_transformer( - transformer, - args.transformer_path, - allow_joint_sidecar=args.route == "base", - ) - if args.route in JOINT_ROUTES: - transformer.attach_joint_modules( - sites=PRODUCTION_JOINT_SITES, - modalities=PRODUCTION_JOINT_MODALITIES, - rank=PRODUCTION_JOINT_RANK, - heads=PRODUCTION_JOINT_HEADS, - residual_mode=args.joint_spec.residual_mode, - ) - if args.joint_spec.schema == "pbr_joint_checkpoint_v1": - _load_exact_joint_adapter(transformer, args.joint_spec.adapter_path) - else: - _load_exact_full_transformer( - transformer, args.joint_spec.full_transformer_path - ) - transformer.set_pbr_attention_backend( - PBR_ATTENTION_FA3 if image_backend == "fa3" else PBR_ATTENTION_NATIVE, - device=device, - dtype=torch.bfloat16, - ) - transformer.to(dtype=torch.bfloat16) - transformer.requires_grad_(False) - transformer.eval() - - pipe = QwenImageEditPlusPipeline.from_pretrained( - str(args.base_model), - transformer=transformer, - torch_dtype=torch.bfloat16, - local_files_only=True, - ) - scheduler_config = dict(pipe.scheduler.config) - scheduler_config["shift_terminal"] = None - pipe.scheduler = FlowMatchEulerDiscreteScheduler.from_config(scheduler_config) - pipe.to(device) - pipe.text_encoder.set_attn_implementation(text_backend) - pipe.set_progress_bar_config(disable=True) - log(f"attention image={image_backend} text={text_backend}") - return pipe - - -def make_generators(torch, device, seed: int): - # Keep per-modality noise independent, deterministic, and invariant to GPU - # sharding. All routes use these exact three streams. - offsets = (104_729, 209_759, 314_159) - return [ - torch.Generator(device=device).manual_seed(seed + offset) - for offset in offsets - ] - - -def encode_source(pipe, source, device, dtype, seed: int): - import numpy as np - import torch - - pixels = torch.from_numpy( - np.asarray(source, dtype=np.float32).copy() - ).permute(2, 0, 1).unsqueeze(0) - pixels = (pixels / 127.5 - 1.0).to(device=device, dtype=dtype).unsqueeze(2) - generator = torch.Generator(device=device).manual_seed(seed) - source_latents = pipe._encode_vae_image(pixels, generator=generator) - _, channels, _, height, width = source_latents.shape - return pipe._pack_latents(source_latents, 1, channels, height, width) - - -def encode_prompts(pipe, item: dict, source, input_size: int, device, dtype): - sys.path.insert(0, str(TRAIN_ROOT)) - from src.forward_step import get_embedding - from src.utils.image_utils import resize_source_images_for_condition - - prompts = [item[f"instr_{modality}"] for modality in MODALITIES] - source_lists = resize_source_images_for_condition( - [[source] for _ in MODALITIES], - ref_height=input_size, - ref_width=input_size, - image_sample_size=input_size, - ) - embeds, mask = get_embedding( - True, - prompts, - source_lists, - pipe.processor, - pipe.text_encoder, - ) - return embeds.to(device=device, dtype=dtype), mask.to(device=device) - - -def generate_triplet(pipe, item: dict, args, device): - import numpy as np - import torch - - sys.path.insert(0, str(TRAIN_ROOT)) - from src.forward_step import calculate_shift - - source = load_source_rgb(item["cond"], args.input_size) - dtype = pipe.transformer.dtype - prompt_embeds, prompt_mask = encode_prompts( - pipe, item, source, args.input_size, device, dtype - ) - source_latents = encode_source(pipe, source, device, dtype, args.seed) - source_latents = source_latents.expand(len(MODALITIES), -1, -1) - - channels = pipe.transformer.config.in_channels // 4 - generators = make_generators(torch, device, args.seed) - latents, _ = pipe.prepare_latents( - images=None, - batch_size=len(MODALITIES), - num_channels_latents=channels, - height=args.input_size, - width=args.input_size, - dtype=dtype, - device=device, - generator=generators, - latents=None, - ) - target_seq_len = latents.shape[1] - grid = args.input_size // pipe.vae_scale_factor // 2 - img_shapes = [[(1, grid, grid), (1, grid, grid)] for _ in MODALITIES] - - sigmas = np.linspace(1.0, 1.0, 1) - mu = calculate_shift( - target_seq_len, - pipe.scheduler.config.get("base_image_seq_len", 256), - pipe.scheduler.config.get("max_image_seq_len", 4096), - pipe.scheduler.config.get("base_shift", 0.5), - pipe.scheduler.config.get("max_shift", 1.15), - ) - pipe.scheduler.set_timesteps(sigmas=sigmas, device=device, mu=mu) - pipe.scheduler.set_begin_index(0) - - with torch.inference_mode(): - for timestep in pipe.scheduler.timesteps: - hidden_states = torch.cat([latents, source_latents], dim=1) - model_timestep = timestep.expand(len(MODALITIES)).to(dtype) / 1000 - transformer_kwargs = { - "hidden_states": hidden_states, - "timestep": model_timestep, - "guidance": None, - "encoder_hidden_states_mask": prompt_mask, - "encoder_hidden_states": prompt_embeds, - "img_shapes": img_shapes, - "attention_kwargs": {}, - "return_dict": False, - } - if args.route in JOINT_ROUTES: - transformer_kwargs.update( - pbr_joint_batch_size=1, - pbr_joint_target_seq_len=target_seq_len, - pbr_joint_route=args.route, - ) - with pipe.transformer.cache_context("cond"): - prediction = pipe.transformer(**transformer_kwargs)[0] - prediction = prediction[:, :target_seq_len] - latents = pipe.scheduler.step( - prediction, timestep, latents, return_dict=False - )[0] - - unpacked = pipe._unpack_latents( - latents, args.input_size, args.input_size, pipe.vae_scale_factor - ).to(pipe.vae.dtype) - mean = torch.tensor(pipe.vae.config.latents_mean).view( - 1, pipe.vae.config.z_dim, 1, 1, 1 - ).to(unpacked.device, unpacked.dtype) - std = 1.0 / torch.tensor(pipe.vae.config.latents_std).view( - 1, pipe.vae.config.z_dim, 1, 1, 1 - ).to(unpacked.device, unpacked.dtype) - unpacked = unpacked / std + mean - decoded = pipe.vae.decode(unpacked, return_dict=False)[0][:, :, 0] - images = pipe.image_processor.postprocess(decoded, output_type="pil") - if len(images) != len(MODALITIES): - raise RuntimeError(f"decoder returned {len(images)} images, expected 3") - return images - - -def save_triplet(out_dir: Path, item: dict, images) -> None: - for modality, image in zip(MODALITIES, images): - destination = output_path(out_dir, item, modality) - temporary = destination.with_name(destination.name + f".tmp.{os.getpid()}.png") - image.save(temporary, format="PNG") - os.replace(temporary, destination) - - -def main() -> None: - parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter) - parser.add_argument("--items-jsonl", type=Path, required=True) - parser.add_argument("--out-dir", type=Path, required=True) - parser.add_argument("--base-model", type=Path, required=True) - parser.add_argument("--transformer-path", type=Path, required=True) - parser.add_argument("--joint-checkpoint", type=Path, required=True) - parser.add_argument("--route", choices=("base",) + JOINT_ROUTES, required=True) - parser.add_argument("--input-size", type=int, default=512) - parser.add_argument("--steps", type=int, default=1) - parser.add_argument("--cfg", type=float, default=1.0) - parser.add_argument("--seed", type=int, default=42) - parser.add_argument("--limit", type=int, default=None) - parser.add_argument("--num-shards", type=int, default=1) - parser.add_argument("--shard-index", type=int, default=0) - parser.add_argument("--dry-run", action="store_true") - args = parser.parse_args() - - if args.limit is not None and args.limit <= 0: - parser.error("--limit must be positive") - if args.num_shards < 1 or not 0 <= args.shard_index < args.num_shards: - parser.error("require num_shards >= 1 and 0 <= shard_index < num_shards") - - for path in (args.base_model / "model_index.json", args.transformer_path): - if not path.is_file(): - raise FileNotFoundError(path) - args.joint_spec = validate_checkpoint(args.joint_checkpoint, args.route) - if ( - args.input_size != args.joint_spec.input_size - or args.steps != 1 - or args.cfg != 1.0 - ): - parser.error( - f"{args.joint_spec.schema} evaluation is pinned to " - f"--input-size {args.joint_spec.input_size} --steps 1 --cfg 1" - ) - items = load_items(args.items_jsonl, args.limit) - validate_payloads(items) - shard_items = items[args.shard_index::args.num_shards] - args.out_dir.mkdir(parents=True, exist_ok=True) - pending = [ - item for item in shard_items - if not all(output_path(args.out_dir, item, modality).stat().st_size > 0 - if output_path(args.out_dir, item, modality).is_file() else False - for modality in MODALITIES) - ] - log( - f"route={args.route} shard={args.shard_index}/{args.num_shards} " - f"items={len(shard_items)} pending={len(pending)} seed={args.seed} " - f"protocol={args.input_size}/s1/cfg1 schema={args.joint_spec.schema}" - ) - if args.dry_run: - for item in shard_items[:10]: - log(f"DRY {item['key']} -> {args.out_dir}") - log("dry run complete; no model or GPU was loaded") - return - if not pending: - log("nothing to do") - return - - os.environ.setdefault("HF_HUB_OFFLINE", "1") - os.environ.setdefault("TOKENIZERS_PARALLELISM", "false") - import torch - - if not torch.cuda.is_available(): - raise RuntimeError("CUDA is required") - device = torch.device("cuda") - torch.manual_seed(args.seed) - log(f"loading model on {torch.cuda.get_device_name(device)}") - pipe = load_pipeline(args, device) - log("model ready") - - started = time.time() - for index, item in enumerate(pending, 1): - images = generate_triplet(pipe, item, args, device) - save_triplet(args.out_dir, item, images) - if index == 1 or index % 10 == 0 or index == len(pending): - elapsed = time.time() - started - rate = elapsed / index - eta = rate * (len(pending) - index) - log( - f"{index}/{len(pending)} key={item['key']} " - f"avg={rate:.2f}s/item eta={eta / 60:.1f}min" - ) - - manifest = { - "schema": "joint_pbr_prediction_shard_v1", - "route": args.route, - "resolution": args.input_size, - "steps": args.steps, - "cfg": args.cfg, - "seed": args.seed, - "items_jsonl": str(args.items_jsonl.absolute()), - "joint_checkpoint": str(args.joint_checkpoint.absolute()), - "joint_checkpoint_schema": args.joint_spec.schema, - "base_transformer": str(args.transformer_path.absolute()), - "shard_index": args.shard_index, - "num_shards": args.num_shards, - "n_items": len(shard_items), - "n_outputs": len(shard_items) * len(MODALITIES), - "elapsed_seconds": time.time() - started, - "peak_cuda_allocated_bytes": torch.cuda.max_memory_allocated(), - "peak_cuda_reserved_bytes": torch.cuda.max_memory_reserved(), - } - manifest_path = args.out_dir / ( - f"prediction_shard_{args.shard_index:02d}_of_{args.num_shards:02d}.json" - ) - temporary = manifest_path.with_name(manifest_path.name + f".tmp.{os.getpid()}") - temporary.write_text(json.dumps(manifest, indent=2) + "\n") - os.replace(temporary, manifest_path) - log(f"DONE outputs={manifest['n_outputs']} manifest={manifest_path}") - - -if __name__ == "__main__": - main() +#!/usr/bin/env python +# Adapted for SUPRA-Tex distribution (2026): packaged local inference runtime. +# Upstream FireRed portions remain under Apache-2.0; see THIRD_PARTY_NOTICES.md. +"""Joint albedo/MRA/normal inference for the factorized PBR checkpoint. + +The stock QwenImageEditPlusPipeline rejects prompt batches larger than one, so +it cannot exercise modality-axis attention. This runner keeps the accepted +pipeline preprocessing/scheduler/decoder semantics, but folds one physical +item into the fixed ``albedo, mra, normal`` batch and advances all three streams +together. Joint checkpoints are trained only at sigma=1, therefore this +runner intentionally accepts one-step, CFG-free inference. The checkpoint +contract selects the recorded 512/1024 resolution for v2/v3. +""" + +from __future__ import annotations + +import argparse +import importlib +import json +import os +import sys +import time +from pathlib import Path +from types import SimpleNamespace + + +REPO = Path(__file__).absolute().parent.parent +sys.path.insert(0, str(REPO)) +MODALITIES = ("albedo", "mra", "normal") +JOINT_ROUTES = ("aligned", "self_slot", "blockdiag") + + +def log(message: str) -> None: + print(f"[predict_joint_pbr {time.strftime('%H:%M:%S')}] {message}", flush=True) + + +def output_path(out_dir: Path, item: dict, modality: str) -> Path: + return out_dir / f"{item['key']}_{modality}.png" + + +def load_items(path: Path, limit: int | None) -> list[dict]: + required = { + "key", + "dataset", + "sha", + "asset_id", + "view", + "cond", + "albedo_gt", + "mra_gt", + "normal_gt", + "normal_valid", + "instr_albedo", + "instr_mra", + "instr_normal", + "split", + } + items = [] + seen = set() + with path.open() as stream: + for line_number, raw_line in enumerate(stream, 1): + if not raw_line.strip(): + continue + item = json.loads(raw_line) + missing = required - set(item) + if missing: + raise ValueError(f"{path}:{line_number}: missing fields {sorted(missing)}") + key = item["key"] + if not isinstance(key, str) or not key or Path(key).name != key or key in {".", ".."} or "\\" in key: + raise ValueError(f"{path}:{line_number}: key must be a single safe filename component") + if item["key"] in seen: + raise ValueError(f"{path}:{line_number}: duplicate key {item['key']}") + if item["split"] != "eval": + raise ValueError(f"{path}:{line_number}: expected split=eval") + seen.add(item["key"]) + items.append(item) + if limit is not None: + items = items[:limit] + if not items: + raise ValueError(f"no evaluation items in {path}") + return items + + +def validate_checkpoint(checkpoint: Path, route: str) -> SimpleNamespace: + manifest_path = checkpoint / "checkpoint_complete.json" + contract_path = checkpoint / "joint_contract.json" + run_contract_path = checkpoint.parent / "RUN_CONTRACT.json" + adapter_path = checkpoint / "joint_adapter.safetensors" + full_transformer_path = checkpoint / "diffusion_pytorch_model.safetensors" + for path in (manifest_path, contract_path): + if not path.is_file(): + raise FileNotFoundError(path) + manifest = json.loads(manifest_path.read_text()) + contract = json.loads(contract_path.read_text()) + common_expected = { + "sites": [29, 39, 49, 59], + "modalities": list(MODALITIES), + "rank": 256, + "heads": 4, + "training_route": "aligned", + } + schema = contract.get("schema") + if schema == "pbr_joint_checkpoint_v1": + expected = { + **common_expected, + "schema": schema, + "trainable_tensor_count": 60, + "trainable_parameter_count": 7_385_088, + } + input_size = 512 + elif schema in ("pbr_joint_checkpoint_v2", "pbr_joint_checkpoint_v3"): + if not run_contract_path.is_file(): + raise FileNotFoundError(run_contract_path) + run_contract = json.loads(run_contract_path.read_text()) + input_size = run_contract.get("resolution") + if input_size not in (512, 1024): + raise ValueError( + f"unsupported joint-v2/v3 run resolution: {input_size!r}" + ) + expected = { + **common_expected, + "schema": schema, + "joint_tensor_count": 60, + "joint_parameter_count": 7_385_088, + "trainable_modules": ["modality_joint_blocks."], + } + backbone_low_lr = [ + f"transformer_blocks.{index}." for index in range(29, 50) + ] + if schema == "pbr_joint_checkpoint_v2": + expected.update( + low_lr_modules=backbone_low_lr, + low_lr_ratio=0.0025, + ) + if contract.get("fresh_joint_w_up_init_std") not in (0.0, 1e-4): + expected["fresh_joint_w_up_init_std"] = "0.0 or 1e-4" + residual_mode = "raw" + else: + low_lr_modules = contract.get("low_lr_modules") + if low_lr_modules == []: + expected.update(low_lr_modules=[], low_lr_ratio=None) + elif low_lr_modules == backbone_low_lr: + expected.update( + low_lr_modules=backbone_low_lr, + low_lr_ratio=0.0025, + ) + else: + expected["low_lr_modules"] = "[] or transformer blocks 29 through 49" + residual_mode = contract.get("residual_mode") + if residual_mode != "rms_scaled": + expected["residual_mode"] = "rms_scaled" + if contract.get("fresh_joint_w_up_init_std") not in (0.0, 3e-3): + expected["fresh_joint_w_up_init_std"] = "0.0 or 0.003" + if contract.get("source_joint_adapter_action") not in ( + "fresh", "loaded", "reset" + ): + expected["source_joint_adapter_action"] = "fresh, loaded, or reset" + else: + raise ValueError(f"unsupported joint checkpoint schema: {schema!r}") + bad = { + key: (contract.get(key), value) + for key, value in expected.items() + if contract.get(key) != value + } + if bad: + raise ValueError(f"joint checkpoint contract mismatch: {bad}") + if manifest.get("joint_contract") != contract: + raise ValueError("checkpoint manifest is not bound to joint_contract.json") + if schema == "pbr_joint_checkpoint_v1": + if (manifest.get("global_step"), manifest.get("world_size"), manifest.get("seed")) \ + != (2000, 40, 42): + raise ValueError(f"joint checkpoint step/world/seed mismatch: {manifest}") + elif ( + manifest.get("format_version") != 2 + or not isinstance(manifest.get("global_step"), int) + or manifest["global_step"] <= 0 + ): + raise ValueError(f"joint-v2/v3 checkpoint manifest mismatch: {manifest}") + if route in JOINT_ROUTES and schema == "pbr_joint_checkpoint_v1": + if not adapter_path.is_file(): + raise FileNotFoundError(adapter_path) + if route in JOINT_ROUTES and schema in ( + "pbr_joint_checkpoint_v2", "pbr_joint_checkpoint_v3" + ): + if not full_transformer_path.is_file(): + raise FileNotFoundError(full_transformer_path) + return SimpleNamespace( + schema=schema, + input_size=input_size, + adapter_path=adapter_path if schema == "pbr_joint_checkpoint_v1" else None, + full_transformer_path=( + full_transformer_path + if schema in ("pbr_joint_checkpoint_v2", "pbr_joint_checkpoint_v3") + else None + ), + residual_mode="raw" if schema != "pbr_joint_checkpoint_v3" else residual_mode, + ) + + +def validate_payloads(items: list[dict]) -> None: + fields = ("cond", "albedo_gt", "mra_gt", "normal_gt", "normal_valid") + for item in items: + missing = [item[field] for field in fields if not Path(item[field]).is_file()] + if missing: + raise FileNotFoundError( + f"{item['key']}: missing payload(s): {missing[:3]}" + ) + + +def load_source_rgb(path: str, size: int): + from PIL import Image + + image = Image.open(path) + if image.mode != "RGB": + rgba = image.convert("RGBA") + background = Image.new("RGBA", rgba.size, (0, 0, 0, 255)) + image = Image.alpha_composite(background, rgba).convert("RGB") + if image.size != (size, size): + image = image.resize((size, size), Image.Resampling.BICUBIC) + return image + + +def _flash_attention_available(module_name: str) -> bool: + try: + module = importlib.import_module(module_name) + except (ImportError, OSError, RuntimeError): + return False + return all( + callable(getattr(module, name, None)) + for name in ("flash_attn_func", "flash_attn_varlen_func") + ) + + +def select_attention_backends(torch, device) -> tuple[str, str]: + """Resolve inference attention without depending on the parent CLI package. + + The image transformer supports FA3 and native PyTorch SDPA. FA2 can still + accelerate the text encoder when the image transformer uses native SDPA. + An explicit SDPA selection applies to both components. + """ + selected = os.environ.get("ATTN_BACKEND", "auto") + if selected == "sdpa": + return "native", "sdpa" + supported = ("auto", "flash_attn_3", "flash_attn", "xformers", "naive", "flash_attn_4") + if selected not in supported: + raise ValueError(f"unsupported ATTN_BACKEND for FireRed: {selected!r}") + + device = torch.device(device) + capability = ( + torch.cuda.get_device_capability(device) + if device.type == "cuda" and torch.cuda.is_available() + else (0, 0) + ) + use_fa3 = ( + selected in ("auto", "flash_attn_3") + and capability[0] == 9 + and _flash_attention_available("flash_attn_interface") + ) + if selected == "flash_attn_3" and not use_fa3: + raise RuntimeError( + "ATTN_BACKEND=flash_attn_3 requires Hopper (SM90) and an importable " + "flash_attn_interface extension; select sdpa for native attention" + ) + use_fa2 = capability[0] in (8, 9) and _flash_attention_available("flash_attn") + if selected == "flash_attn" and not use_fa2: + raise RuntimeError( + "ATTN_BACKEND=flash_attn requires an Ampere/Ada/Hopper GPU and an " + "importable flash_attn extension; select sdpa for native attention" + ) + return "fa3" if use_fa3 else "native", "flash_attention_2" if use_fa2 else "sdpa" + + +def load_pipeline(args, device): + import torch + from diffusers import FlowMatchEulerDiscreteScheduler, QwenImageEditPlusPipeline + + image_backend, text_backend = select_attention_backends(torch, device) + from pbr.model_loading import ( + _load_base_transformer, + _load_exact_full_transformer, + _load_exact_joint_adapter, + ) + from pbr.qwen_pbr_factorized import ( + PBR_ATTENTION_FA3, + PBR_ATTENTION_NATIVE, + PRODUCTION_JOINT_HEADS, + PRODUCTION_JOINT_MODALITIES, + PRODUCTION_JOINT_RANK, + PRODUCTION_JOINT_SITES, + ) + + model_args = SimpleNamespace(pretrained_model_name_or_path=str(args.base_model)) + transformer = _load_base_transformer(model_args, torch.bfloat16) + if args.route == "base" or args.joint_spec.schema == "pbr_joint_checkpoint_v1": + _load_exact_full_transformer( + transformer, + args.transformer_path, + allow_joint_sidecar=args.route == "base", + ) + if args.route in JOINT_ROUTES: + transformer.attach_joint_modules( + sites=PRODUCTION_JOINT_SITES, + modalities=PRODUCTION_JOINT_MODALITIES, + rank=PRODUCTION_JOINT_RANK, + heads=PRODUCTION_JOINT_HEADS, + residual_mode=args.joint_spec.residual_mode, + ) + if args.joint_spec.schema == "pbr_joint_checkpoint_v1": + _load_exact_joint_adapter(transformer, args.joint_spec.adapter_path) + else: + _load_exact_full_transformer( + transformer, args.joint_spec.full_transformer_path + ) + transformer.set_pbr_attention_backend( + PBR_ATTENTION_FA3 if image_backend == "fa3" else PBR_ATTENTION_NATIVE, + device=device, + dtype=torch.bfloat16, + ) + transformer.to(dtype=torch.bfloat16) + transformer.requires_grad_(False) + transformer.eval() + + pipe = QwenImageEditPlusPipeline.from_pretrained( + str(args.base_model), + transformer=transformer, + torch_dtype=torch.bfloat16, + local_files_only=True, + ) + scheduler_config = dict(pipe.scheduler.config) + scheduler_config["shift_terminal"] = None + pipe.scheduler = FlowMatchEulerDiscreteScheduler.from_config(scheduler_config) + pipe.to(device) + pipe.text_encoder.set_attn_implementation(text_backend) + pipe.set_progress_bar_config(disable=True) + log(f"attention image={image_backend} text={text_backend}") + return pipe + + +def make_generators(torch, device, seed: int): + # Keep per-modality noise independent, deterministic, and invariant to GPU + # sharding. All routes use these exact three streams. + offsets = (104_729, 209_759, 314_159) + return [ + torch.Generator(device=device).manual_seed(seed + offset) + for offset in offsets + ] + + +def encode_source(pipe, source, device, dtype, seed: int): + import numpy as np + import torch + + pixels = torch.from_numpy( + np.asarray(source, dtype=np.float32).copy() + ).permute(2, 0, 1).unsqueeze(0) + pixels = (pixels / 127.5 - 1.0).to(device=device, dtype=dtype).unsqueeze(2) + generator = torch.Generator(device=device).manual_seed(seed) + source_latents = pipe._encode_vae_image(pixels, generator=generator) + _, channels, _, height, width = source_latents.shape + return pipe._pack_latents(source_latents, 1, channels, height, width) + + +def encode_prompts(pipe, item: dict, source, input_size: int, device, dtype): + from pbr.conditioning import get_embedding + from pbr.conditioning import resize_source_images_for_condition + + prompts = [item[f"instr_{modality}"] for modality in MODALITIES] + source_lists = resize_source_images_for_condition( + [[source] for _ in MODALITIES], + ref_height=input_size, + ref_width=input_size, + image_sample_size=input_size, + ) + embeds, mask = get_embedding( + True, + prompts, + source_lists, + pipe.processor, + pipe.text_encoder, + ) + return embeds.to(device=device, dtype=dtype), mask.to(device=device) + + +def generate_triplet(pipe, item: dict, args, device): + import numpy as np + import torch + + from pbr.conditioning import calculate_shift + + source = load_source_rgb(item["cond"], args.input_size) + dtype = pipe.transformer.dtype + prompt_embeds, prompt_mask = encode_prompts( + pipe, item, source, args.input_size, device, dtype + ) + source_latents = encode_source(pipe, source, device, dtype, args.seed) + source_latents = source_latents.expand(len(MODALITIES), -1, -1) + + channels = pipe.transformer.config.in_channels // 4 + generators = make_generators(torch, device, args.seed) + latents, _ = pipe.prepare_latents( + images=None, + batch_size=len(MODALITIES), + num_channels_latents=channels, + height=args.input_size, + width=args.input_size, + dtype=dtype, + device=device, + generator=generators, + latents=None, + ) + target_seq_len = latents.shape[1] + grid = args.input_size // pipe.vae_scale_factor // 2 + img_shapes = [[(1, grid, grid), (1, grid, grid)] for _ in MODALITIES] + + sigmas = np.linspace(1.0, 1.0, 1) + mu = calculate_shift( + target_seq_len, + pipe.scheduler.config.get("base_image_seq_len", 256), + pipe.scheduler.config.get("max_image_seq_len", 4096), + pipe.scheduler.config.get("base_shift", 0.5), + pipe.scheduler.config.get("max_shift", 1.15), + ) + pipe.scheduler.set_timesteps(sigmas=sigmas, device=device, mu=mu) + pipe.scheduler.set_begin_index(0) + + with torch.inference_mode(): + for timestep in pipe.scheduler.timesteps: + hidden_states = torch.cat([latents, source_latents], dim=1) + model_timestep = timestep.expand(len(MODALITIES)).to(dtype) / 1000 + transformer_kwargs = { + "hidden_states": hidden_states, + "timestep": model_timestep, + "guidance": None, + "encoder_hidden_states_mask": prompt_mask, + "encoder_hidden_states": prompt_embeds, + "img_shapes": img_shapes, + "attention_kwargs": {}, + "return_dict": False, + } + if args.route in JOINT_ROUTES: + transformer_kwargs.update( + pbr_joint_batch_size=1, + pbr_joint_target_seq_len=target_seq_len, + pbr_joint_route=args.route, + ) + with pipe.transformer.cache_context("cond"): + prediction = pipe.transformer(**transformer_kwargs)[0] + prediction = prediction[:, :target_seq_len] + latents = pipe.scheduler.step( + prediction, timestep, latents, return_dict=False + )[0] + + unpacked = pipe._unpack_latents( + latents, args.input_size, args.input_size, pipe.vae_scale_factor + ).to(pipe.vae.dtype) + mean = torch.tensor(pipe.vae.config.latents_mean).view( + 1, pipe.vae.config.z_dim, 1, 1, 1 + ).to(unpacked.device, unpacked.dtype) + std = 1.0 / torch.tensor(pipe.vae.config.latents_std).view( + 1, pipe.vae.config.z_dim, 1, 1, 1 + ).to(unpacked.device, unpacked.dtype) + unpacked = unpacked / std + mean + decoded = pipe.vae.decode(unpacked, return_dict=False)[0][:, :, 0] + images = pipe.image_processor.postprocess(decoded, output_type="pil") + if len(images) != len(MODALITIES): + raise RuntimeError(f"decoder returned {len(images)} images, expected 3") + return images + + +def save_triplet(out_dir: Path, item: dict, images) -> None: + for modality, image in zip(MODALITIES, images): + destination = output_path(out_dir, item, modality) + temporary = destination.with_name(destination.name + f".tmp.{os.getpid()}.png") + image.save(temporary, format="PNG") + os.replace(temporary, destination) + + +def main() -> None: + parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter) + parser.add_argument("--items-jsonl", type=Path, required=True) + parser.add_argument("--out-dir", type=Path, required=True) + parser.add_argument("--base-model", type=Path, required=True) + parser.add_argument("--transformer-path", type=Path, required=True) + parser.add_argument("--joint-checkpoint", type=Path, required=True) + parser.add_argument("--route", choices=("base",) + JOINT_ROUTES, required=True) + parser.add_argument("--input-size", type=int, default=512) + parser.add_argument("--steps", type=int, default=1) + parser.add_argument("--cfg", type=float, default=1.0) + parser.add_argument("--seed", type=int, default=42) + parser.add_argument("--limit", type=int, default=None) + parser.add_argument("--num-shards", type=int, default=1) + parser.add_argument("--shard-index", type=int, default=0) + parser.add_argument("--dry-run", action="store_true") + args = parser.parse_args() + + if args.limit is not None and args.limit <= 0: + parser.error("--limit must be positive") + if args.num_shards < 1 or not 0 <= args.shard_index < args.num_shards: + parser.error("require num_shards >= 1 and 0 <= shard_index < num_shards") + + for path in (args.base_model / "model_index.json", args.transformer_path): + if not path.is_file(): + raise FileNotFoundError(path) + args.joint_spec = validate_checkpoint(args.joint_checkpoint, args.route) + if ( + args.input_size != args.joint_spec.input_size + or args.steps != 1 + or args.cfg != 1.0 + ): + parser.error( + f"{args.joint_spec.schema} evaluation is pinned to " + f"--input-size {args.joint_spec.input_size} --steps 1 --cfg 1" + ) + items = load_items(args.items_jsonl, args.limit) + validate_payloads(items) + shard_items = items[args.shard_index::args.num_shards] + args.out_dir.mkdir(parents=True, exist_ok=True) + pending = [ + item for item in shard_items + if not all(output_path(args.out_dir, item, modality).stat().st_size > 0 + if output_path(args.out_dir, item, modality).is_file() else False + for modality in MODALITIES) + ] + log( + f"route={args.route} shard={args.shard_index}/{args.num_shards} " + f"items={len(shard_items)} pending={len(pending)} seed={args.seed} " + f"protocol={args.input_size}/s1/cfg1 schema={args.joint_spec.schema}" + ) + if args.dry_run: + for item in shard_items[:10]: + log(f"DRY {item['key']} -> {args.out_dir}") + log("dry run complete; no model or GPU was loaded") + return + if not pending: + log("nothing to do") + return + + os.environ.setdefault("HF_HUB_OFFLINE", "1") + os.environ.setdefault("TOKENIZERS_PARALLELISM", "false") + import torch + + if not torch.cuda.is_available(): + raise RuntimeError("CUDA is required") + device = torch.device("cuda") + torch.manual_seed(args.seed) + log(f"loading model on {torch.cuda.get_device_name(device)}") + pipe = load_pipeline(args, device) + log("model ready") + + started = time.time() + for index, item in enumerate(pending, 1): + images = generate_triplet(pipe, item, args, device) + save_triplet(args.out_dir, item, images) + if index == 1 or index % 10 == 0 or index == len(pending): + elapsed = time.time() - started + rate = elapsed / index + eta = rate * (len(pending) - index) + log( + f"{index}/{len(pending)} key={item['key']} " + f"avg={rate:.2f}s/item eta={eta / 60:.1f}min" + ) + + manifest = { + "schema": "joint_pbr_prediction_shard_v1", + "route": args.route, + "resolution": args.input_size, + "steps": args.steps, + "cfg": args.cfg, + "seed": args.seed, + "items_jsonl": str(args.items_jsonl.absolute()), + "joint_checkpoint": str(args.joint_checkpoint.absolute()), + "joint_checkpoint_schema": args.joint_spec.schema, + "base_transformer": str(args.transformer_path.absolute()), + "shard_index": args.shard_index, + "num_shards": args.num_shards, + "n_items": len(shard_items), + "n_outputs": len(shard_items) * len(MODALITIES), + "elapsed_seconds": time.time() - started, + "peak_cuda_allocated_bytes": torch.cuda.max_memory_allocated(), + "peak_cuda_reserved_bytes": torch.cuda.max_memory_reserved(), + } + manifest_path = args.out_dir / ( + f"prediction_shard_{args.shard_index:02d}_of_{args.num_shards:02d}.json" + ) + temporary = manifest_path.with_name(manifest_path.name + f".tmp.{os.getpid()}") + temporary.write_text(json.dumps(manifest, indent=2) + "\n") + os.replace(temporary, manifest_path) + log(f"DONE outputs={manifest['n_outputs']} manifest={manifest_path}") + + +if __name__ == "__main__": + main() diff --git a/supratex/frontend/material_model/runtime/pbr/qwen_pbr_factorized.py b/supratex/frontend/material_model/runtime/pbr/qwen_pbr_factorized.py new file mode 100644 index 0000000000000000000000000000000000000000..54179f777b90253a121456ad75a3eb7084bf7e7b --- /dev/null +++ b/supratex/frontend/material_model/runtime/pbr/qwen_pbr_factorized.py @@ -0,0 +1,853 @@ +# Adapted for SUPRA-Tex distribution (2026): packaged local inference runtime. +# Upstream FireRed portions remain under Apache-2.0; see THIRD_PARTY_NOTICES.md. +"""Factorized modality-axis attention for joint RGB-to-PBR prediction. + +The stock Qwen image transformer continues to process each modality as an +independent folded-batch item. Small registered residual modules then mix the +three modalities at equal latent-grid positions after selected complete Qwen +blocks. The final up projection is zero initialized, so attaching the modules +is an exact no-op before training. +""" + +from __future__ import annotations + +from functools import lru_cache +from math import prod +from typing import Any, Sequence + +import numpy as np +import torch +from torch import nn + +from diffusers.configuration_utils import register_to_config +from diffusers.models.modeling_outputs import Transformer2DModelOutput +from diffusers.models.transformers.transformer_qwenimage import ( + QwenDoubleStreamAttnProcessor2_0, + QwenImageTransformer2DModel, + apply_rotary_emb_qwen, + compute_text_seq_len_from_mask, +) +from diffusers.utils import apply_lora_scale, deprecate + + +PRODUCTION_JOINT_MODALITIES = ("albedo", "mra", "normal") +PRODUCTION_JOINT_SITES = (29, 39, 49, 59) +PRODUCTION_JOINT_RANK = 256 +PRODUCTION_JOINT_HEADS = 4 + +JOINT_RESIDUAL_RAW = "raw" +JOINT_RESIDUAL_RMS_SCALED = "rms_scaled" +JOINT_RESIDUAL_MODES = (JOINT_RESIDUAL_RAW, JOINT_RESIDUAL_RMS_SCALED) + +PBR_ATTENTION_NATIVE = "native" +PBR_ATTENTION_FA3 = "fa3" +PBR_ATTENTION_BACKENDS = (PBR_ATTENTION_NATIVE, PBR_ATTENTION_FA3) + +ROUTE_ALIGNED = 0 +ROUTE_SELF_SLOT = 1 +ROUTE_BLOCKDIAG = 2 +_ROUTE_NAME_TO_CODE = { + "aligned": ROUTE_ALIGNED, + "self_slot": ROUTE_SELF_SLOT, + "blockdiag": ROUTE_BLOCKDIAG, +} + + +@lru_cache(maxsize=1) +def _load_flash_attention_3_ops(): + try: + from flash_attn_interface import flash_attn_func, flash_attn_varlen_func + except (ImportError, OSError, RuntimeError) as error: + raise RuntimeError( + "FireRed image attention requires the t2h FlashAttention-3 " + "flash_attn_interface extension" + ) from error + return flash_attn_func, flash_attn_varlen_func + + +@lru_cache(maxsize=16) +def validate_fa3_runtime(device: torch.device | str, dtype: torch.dtype) -> None: + """Fail closed unless the selected worker can execute Hopper BF16/FP16 FA3.""" + + device = torch.device(device) + if device.type != "cuda" or not torch.cuda.is_available(): + raise RuntimeError("FireRed FA3 image attention requires a CUDA device") + if dtype not in (torch.bfloat16, torch.float16): + raise RuntimeError( + "FireRed FA3 image attention requires bfloat16 or float16, " + f"got {dtype}" + ) + capability = torch.cuda.get_device_capability(device) + if capability < (9, 0): + raise RuntimeError( + "FireRed FA3 image attention requires Hopper compute capability >= 9.0, " + f"got {capability[0]}.{capability[1]}" + ) + _load_flash_attention_3_ops() + + +def prepare_fa3_key_mask_plan( + attention_mask: torch.Tensor, + query_seq_len: int, +) -> dict[str, torch.Tensor | int]: + """Pack arbitrary valid K/V positions once for all Qwen transformer blocks. + + Qwen joint attention lays tokens out as ``[text, image]``. Because text is + padded before the image suffix, its key mask is not necessarily a valid + prefix. Diffusers' generic FA3-varlen backend slices ``K/V[:valid_len]`` + and would therefore drop real image tokens. This plan gathers the true + boolean-mask positions while leaving every query in its original RoPE + position. + """ + + if attention_mask.dtype != torch.bool or attention_mask.ndim != 2: + raise ValueError( + "FA3 Qwen attention_mask must be a 2D bool tensor, " + f"got shape={tuple(attention_mask.shape)} dtype={attention_mask.dtype}" + ) + batch_size, key_seq_len = attention_mask.shape + query_seq_len = int(query_seq_len) + if batch_size <= 0 or key_seq_len <= 0 or query_seq_len <= 0: + raise ValueError("FA3 Qwen attention dimensions must be positive") + + key_lengths = attention_mask.sum(dim=1, dtype=torch.int32) + if not bool((key_lengths > 0).all().item()): + raise ValueError("every FA3 Qwen sample must retain at least one key token") + key_indices = attention_mask.reshape(-1).nonzero(as_tuple=False).flatten() + cu_seqlens_k = torch.zeros( + batch_size + 1, + dtype=torch.int32, + device=attention_mask.device, + ) + cu_seqlens_k[1:] = key_lengths.cumsum(dim=0) + cu_seqlens_q = torch.arange( + 0, + (batch_size + 1) * query_seq_len, + query_seq_len, + dtype=torch.int32, + device=attention_mask.device, + ) + return { + "key_indices": key_indices, + "cu_seqlens_q": cu_seqlens_q, + "cu_seqlens_k": cu_seqlens_k, + "max_seqlen_q": query_seq_len, + "max_seqlen_k": int(key_lengths.max().item()), + } + + +def flash_attention_3_with_key_mask( + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + *, + attention_mask: torch.Tensor | None, + mask_plan: dict[str, torch.Tensor | int] | None = None, +) -> torch.Tensor: + """Run exact non-causal FA3 attention with Qwen's arbitrary key mask.""" + + if query.ndim != 4 or key.ndim != 4 or value.ndim != 4: + raise ValueError("FA3 Qwen Q/K/V must have shape [B, S, H, D]") + if query.shape != key.shape or key.shape != value.shape: + raise ValueError( + "FA3 Qwen requires equal Q/K/V shapes, got " + f"{tuple(query.shape)}, {tuple(key.shape)}, {tuple(value.shape)}" + ) + if query.device != key.device or key.device != value.device: + raise ValueError("FA3 Qwen Q/K/V must be on the same device") + if query.dtype != key.dtype or key.dtype != value.dtype: + raise ValueError("FA3 Qwen Q/K/V must have the same dtype") + validate_fa3_runtime(query.device, query.dtype) + flash_attn_func, flash_attn_varlen_func = _load_flash_attention_3_ops() + + query = query.contiguous() + key = key.contiguous() + value = value.contiguous() + if attention_mask is None: + return flash_attn_func( + query, + key, + value, + softmax_scale=None, + causal=False, + ) + + if attention_mask.shape != query.shape[:2]: + raise ValueError( + "FA3 Qwen key mask must match [B, S], got " + f"{tuple(attention_mask.shape)} for {tuple(query.shape[:2])}" + ) + if attention_mask.device != query.device: + raise ValueError("FA3 Qwen key mask must be on the Q/K/V device") + if mask_plan is None: + mask_plan = prepare_fa3_key_mask_plan(attention_mask, query.shape[1]) + + batch_size, seq_len, heads, head_dim = query.shape + key_indices = mask_plan["key_indices"] + if not isinstance(key_indices, torch.Tensor) or key_indices.device != query.device: + raise ValueError("invalid FA3 key_indices mask plan") + flat_key = key.reshape(batch_size * seq_len, heads, head_dim) + flat_value = value.reshape(batch_size * seq_len, heads, head_dim) + packed_key = flat_key.index_select(0, key_indices) + packed_value = flat_value.index_select(0, key_indices) + packed_query = query.reshape(batch_size * seq_len, heads, head_dim) + output = flash_attn_varlen_func( + packed_query, + packed_key, + packed_value, + mask_plan["cu_seqlens_q"], + mask_plan["cu_seqlens_k"], + int(mask_plan["max_seqlen_q"]), + int(mask_plan["max_seqlen_k"]), + softmax_scale=None, + causal=False, + ) + return output.reshape(batch_size, seq_len, heads, head_dim) + + +class QwenDoubleStreamAttnProcessorFA3: + """Qwen double-stream processor using mask-correct Hopper FA3 kernels.""" + + _attention_backend = "_flash_3" + _parallel_config = None + + def __init__(self) -> None: + _load_flash_attention_3_ops() + + def __call__( + self, + attn, + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor | None = None, + encoder_hidden_states_mask: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + image_rotary_emb: torch.Tensor | None = None, + fa3_key_indices: torch.Tensor | None = None, + fa3_cu_seqlens_q: torch.Tensor | None = None, + fa3_cu_seqlens_k: torch.Tensor | None = None, + fa3_max_seqlen_q: int | None = None, + fa3_max_seqlen_k: int | None = None, + ) -> tuple[torch.Tensor, torch.Tensor]: + if encoder_hidden_states is None: + raise ValueError( + "QwenDoubleStreamAttnProcessorFA3 requires encoder_hidden_states" + ) + + seq_txt = encoder_hidden_states.shape[1] + img_query = attn.to_q(hidden_states) + img_key = attn.to_k(hidden_states) + img_value = attn.to_v(hidden_states) + txt_query = attn.add_q_proj(encoder_hidden_states) + txt_key = attn.add_k_proj(encoder_hidden_states) + txt_value = attn.add_v_proj(encoder_hidden_states) + + img_query = img_query.unflatten(-1, (attn.heads, -1)) + img_key = img_key.unflatten(-1, (attn.heads, -1)) + img_value = img_value.unflatten(-1, (attn.heads, -1)) + txt_query = txt_query.unflatten(-1, (attn.heads, -1)) + txt_key = txt_key.unflatten(-1, (attn.heads, -1)) + txt_value = txt_value.unflatten(-1, (attn.heads, -1)) + + if attn.norm_q is not None: + img_query = attn.norm_q(img_query) + if attn.norm_k is not None: + img_key = attn.norm_k(img_key) + if attn.norm_added_q is not None: + txt_query = attn.norm_added_q(txt_query) + if attn.norm_added_k is not None: + txt_key = attn.norm_added_k(txt_key) + + if image_rotary_emb is not None: + img_freqs, txt_freqs = image_rotary_emb + img_query = apply_rotary_emb_qwen(img_query, img_freqs, use_real=False) + img_key = apply_rotary_emb_qwen(img_key, img_freqs, use_real=False) + txt_query = apply_rotary_emb_qwen(txt_query, txt_freqs, use_real=False) + txt_key = apply_rotary_emb_qwen(txt_key, txt_freqs, use_real=False) + + joint_query = torch.cat([txt_query, img_query], dim=1) + joint_key = torch.cat([txt_key, img_key], dim=1) + joint_value = torch.cat([txt_value, img_value], dim=1) + mask_plan = None + if attention_mask is not None: + required_plan = ( + fa3_key_indices, + fa3_cu_seqlens_q, + fa3_cu_seqlens_k, + fa3_max_seqlen_q, + fa3_max_seqlen_k, + ) + if any(item is None for item in required_plan): + raise ValueError("FA3 Qwen attention mask plan is incomplete") + mask_plan = { + "key_indices": fa3_key_indices, + "cu_seqlens_q": fa3_cu_seqlens_q, + "cu_seqlens_k": fa3_cu_seqlens_k, + "max_seqlen_q": fa3_max_seqlen_q, + "max_seqlen_k": fa3_max_seqlen_k, + } + joint_hidden_states = flash_attention_3_with_key_mask( + joint_query, + joint_key, + joint_value, + attention_mask=attention_mask, + mask_plan=mask_plan, + ) + joint_hidden_states = joint_hidden_states.flatten(2, 3).to(joint_query.dtype) + + txt_attn_output = joint_hidden_states[:, :seq_txt, :] + img_attn_output = joint_hidden_states[:, seq_txt:, :] + img_attn_output = attn.to_out[0](img_attn_output.contiguous()) + if len(attn.to_out) > 1: + img_attn_output = attn.to_out[1](img_attn_output) + txt_attn_output = attn.to_add_out(txt_attn_output.contiguous()) + return img_attn_output, txt_attn_output + + +def _normalize_modalities(modalities: Sequence[str]) -> tuple[str, ...]: + result = tuple(str(item) for item in modalities) + if result != PRODUCTION_JOINT_MODALITIES: + raise ValueError( + "the first factorized PBR model fixes modality order to " + f"{PRODUCTION_JOINT_MODALITIES}, got {result}" + ) + return result + + +def _normalize_sites(sites: Sequence[int], num_layers: int) -> tuple[int, ...]: + result = tuple(int(site) for site in sites) + if len(set(result)) != len(result): + raise ValueError(f"joint sites must be unique, got {result}") + if any(site < 0 or site >= num_layers for site in result): + raise ValueError(f"joint sites must be within [0, {num_layers - 1}], got {result}") + return result + + +def _normalize_residual_mode(residual_mode: str) -> str: + result = str(residual_mode) + if result not in JOINT_RESIDUAL_MODES: + raise ValueError( + f"joint residual mode must be one of {JOINT_RESIDUAL_MODES}, got {result!r}" + ) + return result + + +class ModalityAxisJointAttention(nn.Module): + """Mix equal-position modality tokens while keeping objects/tokens separate. + + Inputs use object-major folded-batch order: ``[B*M, L, D]`` where M is the + fixed modality count. Attention is over the donor-modality axis only. The + checkpointable forward accepts tensors exclusively. + """ + + def __init__( + self, + hidden_dim: int, + rank: int = PRODUCTION_JOINT_RANK, + heads: int = PRODUCTION_JOINT_HEADS, + modalities: Sequence[str] = PRODUCTION_JOINT_MODALITIES, + residual_mode: str = JOINT_RESIDUAL_RAW, + ) -> None: + super().__init__() + self.hidden_dim = int(hidden_dim) + self.rank = int(rank) + self.heads = int(heads) + self.modalities = _normalize_modalities(modalities) + self.num_modalities = len(self.modalities) + self.residual_mode = _normalize_residual_mode(residual_mode) + + if self.hidden_dim <= 0 or self.rank <= 0 or self.heads <= 0: + raise ValueError("hidden_dim, rank, and heads must all be positive") + if self.rank % self.heads != 0: + raise ValueError(f"rank={self.rank} must be divisible by heads={self.heads}") + self.head_dim = self.rank // self.heads + self.scale = self.head_dim**-0.5 + + self.norm = nn.LayerNorm(self.hidden_dim) + self.w_down = nn.Linear(self.hidden_dim, self.rank) + self.modality_embedding = nn.Parameter(torch.empty(self.num_modalities, self.rank)) + self.to_q = nn.Linear(self.rank, self.rank) + self.to_k = nn.Linear(self.rank, self.rank) + self.to_v = nn.Linear(self.rank, self.rank) + self.to_out = nn.Linear(self.rank, self.rank) + self.w_up = nn.Linear(self.rank, self.hidden_dim) + + nn.init.normal_(self.modality_embedding, mean=0.0, std=0.02) + nn.init.zeros_(self.w_up.weight) + nn.init.zeros_(self.w_up.bias) + + @staticmethod + def _route_value(route_code: torch.Tensor) -> int: + if not isinstance(route_code, torch.Tensor) or route_code.numel() != 1: + raise ValueError("route_code must be a scalar tensor") + route = int(route_code.detach().item()) + if route not in (ROUTE_ALIGNED, ROUTE_SELF_SLOT, ROUTE_BLOCKDIAG): + raise ValueError(f"unknown joint route code: {route}") + return route + + def forward( + self, + hidden_states: torch.Tensor, + route_code: torch.Tensor, + modality_attention_mask: torch.Tensor, + ) -> torch.Tensor: + """Return a modality-mixed residual update for target-prefix tokens.""" + + if hidden_states.ndim != 3: + raise ValueError( + f"hidden_states must have shape [B*M, L, D], got {tuple(hidden_states.shape)}" + ) + folded_batch, seq_len, hidden_dim = hidden_states.shape + if hidden_dim != self.hidden_dim: + raise ValueError(f"expected hidden dim {self.hidden_dim}, got {hidden_dim}") + if folded_batch % self.num_modalities != 0: + raise ValueError( + f"folded batch {folded_batch} is not divisible by M={self.num_modalities}" + ) + if seq_len <= 0: + raise ValueError("the target token sequence must be non-empty") + + route = self._route_value(route_code) + if route == ROUTE_BLOCKDIAG and self.training: + raise ValueError("blockdiag is an evaluation-only intervention") + + if not isinstance(modality_attention_mask, torch.Tensor): + raise ValueError("modality_attention_mask must be a tensor") + if tuple(modality_attention_mask.shape) != (self.num_modalities, self.num_modalities): + raise ValueError( + "modality_attention_mask must have shape " + f"[{self.num_modalities}, {self.num_modalities}], got " + f"{tuple(modality_attention_mask.shape)}" + ) + attention_mask = modality_attention_mask.to(device=hidden_states.device, dtype=torch.bool) + if not bool(attention_mask.any(dim=-1).all().item()): + raise ValueError("every modality query must have at least one permitted donor") + if route == ROUTE_BLOCKDIAG: + expected = torch.eye(self.num_modalities, dtype=torch.bool, device=hidden_states.device) + if not torch.equal(attention_mask, expected): + raise ValueError("blockdiag route requires an identity modality attention mask") + + batch_size = folded_batch // self.num_modalities + # Object-major fold: flat index = object_index * M + modality_index. + x = hidden_states.reshape(batch_size, self.num_modalities, seq_len, hidden_dim) + x = x.permute(0, 2, 1, 3) # [B, L, M, D] + z = self.w_down(self.norm(x)) # [B, L, M, R] + + modality_embedding = self.modality_embedding.to(dtype=z.dtype) + query_input = z + modality_embedding.view(1, 1, self.num_modalities, self.rank) + + if route == ROUTE_SELF_SLOT: + # Each query supplies its own content to all donor slots. Donor + # modality embeddings keep the slots distinct before K/V. + donor_input = z.unsqueeze(3).expand(-1, -1, -1, self.num_modalities, -1) + else: + # Every query sees the real aligned modality tokens at this object + # and latent-grid position. Blockdiag changes only the edge mask. + donor_input = z.unsqueeze(2).expand(-1, -1, self.num_modalities, -1, -1) + donor_input = donor_input + modality_embedding.view( + 1, 1, 1, self.num_modalities, self.rank + ) + + query = self.to_q(query_input).reshape( + batch_size, seq_len, self.num_modalities, self.heads, self.head_dim + ) + query = query.permute(0, 1, 3, 2, 4) # [B, L, H, Q, dh] + + key = self.to_k(donor_input).reshape( + batch_size, + seq_len, + self.num_modalities, + self.num_modalities, + self.heads, + self.head_dim, + ) + value = self.to_v(donor_input).reshape_as(key) + key = key.permute(0, 1, 4, 2, 3, 5) # [B, L, H, Q, K, dh] + value = value.permute(0, 1, 4, 2, 3, 5) + + logits = (query.unsqueeze(-2) * key).sum(dim=-1) * self.scale + mask = attention_mask.view(1, 1, 1, self.num_modalities, self.num_modalities) + logits = logits.masked_fill(~mask, torch.finfo(logits.dtype).min) + weights = torch.softmax(logits.float(), dim=-1).to(dtype=logits.dtype) + context = (weights.unsqueeze(-1) * value).sum(dim=-2) + context = context.permute(0, 1, 3, 2, 4).reshape( + batch_size, seq_len, self.num_modalities, self.rank + ) + delta = self.w_up(self.to_out(context)) + if self.residual_mode == JOINT_RESIDUAL_RMS_SCALED: + # Qwen image residual states can reach 1e6--1e7 in BF16 even though + # every following block consumes a LayerNorm-normalized view. A + # raw adapter update therefore disappears at the residual add. In + # the repaired parameterization, the adapter predicts a normalized + # update and this detached RMS converts it back to the host unit. + token_rms = torch.linalg.vector_norm( + x, dim=-1, keepdim=True, dtype=torch.float32 + ) * (hidden_dim**-0.5) + delta = delta * token_rms.detach().to(dtype=delta.dtype) + output = x + delta + return output.permute(0, 2, 1, 3).reshape(folded_batch, seq_len, hidden_dim) + + +class QwenImagePBRFactorizedTransformer2DModel(QwenImageTransformer2DModel): + """Qwen image transformer with registered factorized PBR joint modules.""" + + @register_to_config + def __init__( + self, + patch_size: int = 2, + in_channels: int = 64, + out_channels: int | None = 16, + num_layers: int = 60, + attention_head_dim: int = 128, + num_attention_heads: int = 24, + joint_attention_dim: int = 3584, + guidance_embeds: bool = False, + axes_dims_rope: tuple[int, int, int] = (16, 56, 56), + zero_cond_t: bool = False, + use_additional_t_cond: bool = False, + use_layer3d_rope: bool = False, + joint_sites: Sequence[int] = (), + joint_modalities: Sequence[str] = PRODUCTION_JOINT_MODALITIES, + joint_rank: int = PRODUCTION_JOINT_RANK, + joint_heads: int = PRODUCTION_JOINT_HEADS, + joint_residual_mode: str = JOINT_RESIDUAL_RAW, + ) -> None: + super().__init__( + patch_size=patch_size, + in_channels=in_channels, + out_channels=out_channels, + num_layers=num_layers, + attention_head_dim=attention_head_dim, + num_attention_heads=num_attention_heads, + joint_attention_dim=joint_attention_dim, + guidance_embeds=guidance_embeds, + axes_dims_rope=axes_dims_rope, + zero_cond_t=zero_cond_t, + use_additional_t_cond=use_additional_t_cond, + use_layer3d_rope=use_layer3d_rope, + ) + self._pbr_attention_backend = PBR_ATTENTION_NATIVE + self.modality_joint_blocks = nn.ModuleDict() + self.joint_modalities = _normalize_modalities(joint_modalities) + self.joint_rank = int(joint_rank) + self.joint_heads = int(joint_heads) + self.joint_residual_mode = _normalize_residual_mode(joint_residual_mode) + + normalized_sites = _normalize_sites(joint_sites, len(self.transformer_blocks)) + if normalized_sites: + self._attach_joint_modules( + sites=normalized_sites, + modalities=self.joint_modalities, + rank=self.joint_rank, + heads=self.joint_heads, + residual_mode=self.joint_residual_mode, + update_config=True, + ) + + @property + def joint_sites(self) -> tuple[int, ...]: + return tuple(int(site) for site in self.modality_joint_blocks.keys()) + + def _attach_joint_modules( + self, + sites: Sequence[int], + modalities: Sequence[str], + rank: int, + heads: int, + residual_mode: str, + *, + update_config: bool, + ) -> None: + if len(self.modality_joint_blocks) != 0: + raise RuntimeError(f"joint modules are already attached at sites {self.joint_sites}") + normalized_sites = _normalize_sites(sites, len(self.transformer_blocks)) + if not normalized_sites: + raise ValueError("attach_joint_modules requires at least one site") + normalized_modalities = _normalize_modalities(modalities) + rank = int(rank) + heads = int(heads) + residual_mode = _normalize_residual_mode(residual_mode) + if rank <= 0 or heads <= 0 or rank % heads != 0: + raise ValueError(f"joint rank={rank} must be positive and divisible by heads={heads}") + + self.modality_joint_blocks.update( + { + str(site): ModalityAxisJointAttention( + hidden_dim=self.inner_dim, + rank=rank, + heads=heads, + modalities=normalized_modalities, + residual_mode=residual_mode, + ) + for site in normalized_sites + } + ) + self.joint_modalities = normalized_modalities + self.joint_rank = rank + self.joint_heads = heads + self.joint_residual_mode = residual_mode + if update_config: + self.register_to_config( + architecture_type="QwenImagePBRFactorized", + joint_sites=list(normalized_sites), + joint_modalities=list(normalized_modalities), + joint_rank=rank, + joint_heads=heads, + joint_residual_mode=residual_mode, + supported_joint_routes=["self_slot", "aligned"], + blockdiag_eval_only=True, + ) + + def attach_joint_modules( + self, + sites: Sequence[int] = PRODUCTION_JOINT_SITES, + modalities: Sequence[str] = PRODUCTION_JOINT_MODALITIES, + rank: int = PRODUCTION_JOINT_RANK, + heads: int = PRODUCTION_JOINT_HEADS, + residual_mode: str = JOINT_RESIDUAL_RAW, + ) -> None: + """Attach initialized modules after an exact base-only load.""" + + self._attach_joint_modules( + sites=sites, + modalities=modalities, + rank=rank, + heads=heads, + residual_mode=residual_mode, + update_config=True, + ) + + def set_pbr_attention_backend( + self, + backend: str, + *, + device: torch.device | str | None = None, + dtype: torch.dtype | None = None, + ) -> None: + """Select the Qwen image-transformer attention kernel for this runtime.""" + + backend = str(backend) + if backend not in PBR_ATTENTION_BACKENDS: + raise ValueError( + f"pbr attention backend must be one of {PBR_ATTENTION_BACKENDS}, " + f"got {backend!r}" + ) + if backend == PBR_ATTENTION_FA3: + if device is None or dtype is None: + raise ValueError("FA3 backend selection requires runtime device and dtype") + validate_fa3_runtime(device, dtype) + processor_factory = QwenDoubleStreamAttnProcessorFA3 + else: + processor_factory = QwenDoubleStreamAttnProcessor2_0 + + for block in self.transformer_blocks: + block.attn.set_processor(processor_factory()) + self._pbr_attention_backend = backend + + @staticmethod + def _runtime_int(value: int | torch.Tensor | None, name: str) -> int: + if isinstance(value, torch.Tensor): + if value.numel() != 1: + raise ValueError(f"{name} must be a scalar") + value = int(value.detach().item()) + if not isinstance(value, int): + raise ValueError(f"{name} must be an integer") + return value + + @staticmethod + def _runtime_route(value: str | int | torch.Tensor | None, device: torch.device) -> torch.Tensor: + if isinstance(value, str): + if value not in _ROUTE_NAME_TO_CODE: + raise ValueError(f"unknown pbr_joint_route: {value!r}") + code = _ROUTE_NAME_TO_CODE[value] + elif isinstance(value, torch.Tensor): + if value.numel() != 1: + raise ValueError("pbr_joint_route tensor must be scalar") + code = int(value.detach().item()) + elif isinstance(value, int): + code = value + else: + raise ValueError("pbr_joint_route must be self_slot, aligned, or blockdiag") + if code not in (ROUTE_ALIGNED, ROUTE_SELF_SLOT, ROUTE_BLOCKDIAG): + raise ValueError(f"unknown pbr_joint_route code: {code}") + return torch.tensor(code, dtype=torch.int64, device=device) + + @apply_lora_scale("attention_kwargs") + def forward( + self, + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor | None = None, + encoder_hidden_states_mask: torch.Tensor | None = None, + timestep: torch.LongTensor | None = None, + img_shapes: list[tuple[int, int, int]] | None = None, + txt_seq_lens: list[int] | None = None, + guidance: torch.Tensor | None = None, + attention_kwargs: dict[str, Any] | None = None, + controlnet_block_samples=None, + additional_t_cond=None, + return_dict: bool = True, + pbr_joint_batch_size: int | torch.Tensor | None = None, + pbr_joint_target_seq_len: int | torch.Tensor | None = None, + pbr_joint_route: str | int | torch.Tensor | None = None, + ) -> torch.Tensor | Transformer2DModelOutput: + """Run stock Qwen blocks and insert target-only modality adapters.""" + + has_joint_modules = len(self.modality_joint_blocks) > 0 + runtime_values = (pbr_joint_batch_size, pbr_joint_target_seq_len, pbr_joint_route) + if has_joint_modules and any(value is None for value in runtime_values): + raise ValueError( + "joint models require pbr_joint_batch_size, pbr_joint_target_seq_len, " + "and pbr_joint_route" + ) + if not has_joint_modules and any(value is not None for value in runtime_values): + raise ValueError("joint runtime arguments were provided before joint modules were attached") + + if txt_seq_lens is not None: + deprecate( + "txt_seq_lens", + "0.39.0", + "Passing `txt_seq_lens` is deprecated and will be removed in version 0.39.0. " + "Please use `encoder_hidden_states_mask` instead. " + "The mask-based approach is more flexible and supports variable-length sequences.", + standard_warn=False, + ) + + if has_joint_modules: + batch_size = self._runtime_int(pbr_joint_batch_size, "pbr_joint_batch_size") + target_seq_len = self._runtime_int( + pbr_joint_target_seq_len, "pbr_joint_target_seq_len" + ) + if batch_size <= 0: + raise ValueError("pbr_joint_batch_size must be positive") + expected_folded_batch = batch_size * len(self.joint_modalities) + if hidden_states.shape[0] != expected_folded_batch: + raise ValueError( + "object-major folded batch mismatch: expected " + f"B*M={expected_folded_batch}, got {hidden_states.shape[0]}" + ) + if target_seq_len <= 0 or target_seq_len > hidden_states.shape[1]: + raise ValueError( + f"pbr_joint_target_seq_len must be in [1, {hidden_states.shape[1]}], " + f"got {target_seq_len}" + ) + route_code = self._runtime_route(pbr_joint_route, hidden_states.device) + if int(route_code.item()) == ROUTE_BLOCKDIAG: + modality_attention_mask = torch.eye( + len(self.joint_modalities), dtype=torch.bool, device=hidden_states.device + ) + else: + modality_attention_mask = torch.ones( + (len(self.joint_modalities), len(self.joint_modalities)), + dtype=torch.bool, + device=hidden_states.device, + ) + else: + target_seq_len = 0 + route_code = None + modality_attention_mask = None + + hidden_states = self.img_in(hidden_states) + timestep = timestep.to(hidden_states.dtype) + + if self.zero_cond_t: + timestep = torch.cat([timestep, timestep * 0], dim=0) + modulate_index = torch.tensor( + [[0] * prod(sample[0]) + [1] * sum(prod(shape) for shape in sample[1:]) for sample in img_shapes], + device=timestep.device, + dtype=torch.int, + ) + else: + modulate_index = None + + encoder_hidden_states = self.txt_norm(encoder_hidden_states) + encoder_hidden_states = self.txt_in(encoder_hidden_states) + text_seq_len, _, encoder_hidden_states_mask = compute_text_seq_len_from_mask( + encoder_hidden_states, encoder_hidden_states_mask + ) + + if guidance is not None: + guidance = guidance.to(hidden_states.dtype) * 1000 + + temb = ( + self.time_text_embed(timestep, hidden_states, additional_t_cond) + if guidance is None + else self.time_text_embed(timestep, guidance, hidden_states, additional_t_cond) + ) + image_rotary_emb = self.pos_embed( + img_shapes, max_txt_seq_len=text_seq_len, device=hidden_states.device + ) + + block_attention_kwargs = attention_kwargs.copy() if attention_kwargs is not None else {} + if encoder_hidden_states_mask is not None: + folded_batch, image_seq_len = hidden_states.shape[:2] + image_mask = torch.ones( + (folded_batch, image_seq_len), dtype=torch.bool, device=hidden_states.device + ) + joint_attention_mask = torch.cat([encoder_hidden_states_mask, image_mask], dim=1) + block_attention_kwargs["attention_mask"] = joint_attention_mask + if self._pbr_attention_backend == PBR_ATTENTION_FA3: + reserved = { + "fa3_key_indices", + "fa3_cu_seqlens_q", + "fa3_cu_seqlens_k", + "fa3_max_seqlen_q", + "fa3_max_seqlen_k", + } + conflict = reserved.intersection(block_attention_kwargs) + if conflict: + raise ValueError(f"reserved FA3 attention kwargs supplied: {sorted(conflict)}") + mask_plan = prepare_fa3_key_mask_plan( + joint_attention_mask, + query_seq_len=joint_attention_mask.shape[1], + ) + block_attention_kwargs.update( + { + "fa3_key_indices": mask_plan["key_indices"], + "fa3_cu_seqlens_q": mask_plan["cu_seqlens_q"], + "fa3_cu_seqlens_k": mask_plan["cu_seqlens_k"], + "fa3_max_seqlen_q": mask_plan["max_seqlen_q"], + "fa3_max_seqlen_k": mask_plan["max_seqlen_k"], + } + ) + + for index_block, block in enumerate(self.transformer_blocks): + encoder_hidden_states, hidden_states = block( + hidden_states=hidden_states, + encoder_hidden_states=encoder_hidden_states, + encoder_hidden_states_mask=None, + temb=temb, + image_rotary_emb=image_rotary_emb, + joint_attention_kwargs=block_attention_kwargs, + modulate_index=modulate_index, + ) + + mixer_key = str(index_block) + if mixer_key in self.modality_joint_blocks: + mixer = self.modality_joint_blocks[mixer_key] + target_hidden = hidden_states[:, :target_seq_len] + target_hidden = mixer( + target_hidden, + route_code, + modality_attention_mask, + ) + if target_seq_len == hidden_states.shape[1]: + hidden_states = target_hidden + else: + # The adapter never reads or writes the common source suffix. + hidden_states = torch.cat( + [target_hidden, hidden_states[:, target_seq_len:]], dim=1 + ) + + if controlnet_block_samples is not None: + interval_control = int( + np.ceil(len(self.transformer_blocks) / len(controlnet_block_samples)) + ) + hidden_states = hidden_states + controlnet_block_samples[index_block // interval_control] + + if self.zero_cond_t: + temb = temb.chunk(2, dim=0)[0] + hidden_states = self.norm_out(hidden_states, temb) + output = self.proj_out(hidden_states) + + if not return_dict: + return (output,) + return Transformer2DModelOutput(sample=output) diff --git a/supratex/frontend/pipeline.py b/supratex/frontend/pipeline.py index 42bb4a7343eb764fd045d9589629fba9a8efb8ec..4cfa0e09c32b03de358edea583937ffbcb03f723 100644 --- a/supratex/frontend/pipeline.py +++ b/supratex/frontend/pipeline.py @@ -1,265 +1,291 @@ -"""Runnable Image-to-Material Translation handoff builder.""" -from __future__ import annotations - -import argparse -import hashlib -import json -from pathlib import Path - -from hf_space.timing import span - -from .camera import estimate -from .firered import run as run_firered -from .preprocess import RMBG2, prepare_image -from .geometry import run as run_geometry -from .qwen_image import run as run_qwen, validate as validate_qwen - - -def predict_material_maps(rgba: Path, rgb: Path, output: Path, *, - material_backend: str, material_base: str | Path, - material_model: str | Path | None = None, - qwen_runtime: Path | None = None, qwen_deps: Path | None = None, - firered_package: Path | None = None, firered_manifest: Path | None = None, - material_weights: Path | None = None, material_spec: Path | None = None, - material_checkpoint_manifest: Path | None = None, - revision: str | None = None, cache_dir: Path | None = None, - local_files_only: bool = False, device: str = "cuda", key: str = "input", - receipt_path: Path | None = None, resident=None) -> tuple[Path, Path]: - """Predict either view with the same material backend and artifact selection.""" - if receipt_path is not None and firered_manifest is not None: - raise ValueError("material provenance does not support a custom FireRed manifest; " - "omit it to bind the current condition and use matching prompts for both views") - if receipt_path is not None: - from supratex.weights import resolve_model_dir - material_base = resolve_model_dir(material_base, cache_dir=cache_dir, local_files_only=local_files_only) - if material_model is not None: - patterns = (["checkpoint-0100000.pt", "checkpoint-0100000.complete.json"] - if material_backend == "qwen-image" else - ["full_transformer.safetensors", "joint_contract.json", "checkpoint_complete.json"]) - material_model = resolve_model_dir(material_model, revision=revision, cache_dir=cache_dir, - local_files_only=local_files_only, allow_patterns=patterns) - from supratex.backview.material_identity import artifact_identity - before_artifacts = artifact_identity({"base": material_base, "model": material_model, - "weights": material_weights, "spec": material_spec, - "checkpoint_manifest": material_checkpoint_manifest}) - if material_backend == "qwen-image": - albedo, mra = run_qwen(qwen_runtime, material_base, material_model, rgba, output, - deps=qwen_deps, cache_dir=cache_dir, revision=revision, - local_files_only=local_files_only, device=device, resident=resident) - elif material_backend == "firered": - albedo, mra = run_firered(firered_package, material_base, rgb, output, - firered_manifest, key=key, weights=material_weights, - spec=material_spec, checkpoint_manifest=material_checkpoint_manifest, - material_model=material_model, revision=revision, - cache_dir=cache_dir, local_files_only=local_files_only) - else: - raise ValueError("material backend must be qwen-image or firered") - if receipt_path is not None: - def identity(value): - if value is None: - return None - path = Path(value).expanduser() - return str(path.resolve()) if path.exists() else str(value) - signature = {"backend": material_backend, "base": identity(material_base), - "model": identity(material_model), "revision": revision, - "runtime": identity(qwen_runtime if material_backend == "qwen-image" else firered_package), - "deps": identity(qwen_deps) if material_backend == "qwen-image" else None, - "weights": identity(material_weights), "spec": identity(material_spec), - "checkpoint_manifest": identity(material_checkpoint_manifest), "seed": 42} - from supratex.backview.material_identity import artifact_identity - signature["artifacts"] = artifact_identity(signature) - if signature["artifacts"] != before_artifacts: - raise RuntimeError("material artifacts changed during prediction; discard this view and use a new output") - # Image digests bind the receipt to the exact maps consumed downstream. - paths = {"condition_rgba.png": rgba, "condition_rgb.png": rgb, - "albedo.png": albedo, "mra.png": mra} - receipt = {"status": "complete", "backend": material_backend, "model_signature": signature, - "files": {name: hashlib.sha256(Path(path).read_bytes()).hexdigest() - for name, path in paths.items()}} - receipt_path = Path(receipt_path) - if receipt_path.exists(): - raise FileExistsError(receipt_path) - receipt_path.write_text(json.dumps(receipt, indent=2) + "\n", encoding="utf-8") - return albedo, mra - - -def run(input_image: Path, output: Path, *, moge_model: str | Path, - firered_package: Path | None = None, firered_base: str | Path | None = None, - material_backend: str = "qwen-image", material_base: str | Path | None = None, - qwen_runtime: Path | None = None, qwen_deps: Path | None = None, - firered_manifest: Path | None = None, material_weights: Path | None = None, material_spec: Path | None = None, - material_checkpoint_manifest: Path | None = None, rmbg_model: str | Path | None = None, use_input_alpha: bool = False, - device: str = "cuda", moge_module: str = "moge.model.v2", resolution_level: int = 9, - key: str = "input", material_model: str | Path | None = None, revision: str | None = None, - cache_dir: Path | None = None, local_files_only: bool = False, - geometry_package: Path | None = None, geometry_model: str | Path | None = None, - geometry_revision: str | None = None, geometry_initial_mesh_scale: float = 1.0, - geometry_target_padding: int = 3, geometry_max_iterations: int = 2, - geometry_seed: int = 42, record_material_provenance: bool = False, - attention: str = "auto", low_vram: bool | None = None, resident=None) -> dict[str, str]: - """Build the condition, canonical camera, and material-map handoff in one directory.""" - with span("frontend.validation"): - if record_material_provenance and firered_manifest is not None: - raise ValueError("material provenance does not support a custom FireRed manifest; " - "omit it to bind the current condition and use matching prompts for both views") - if output.exists() and any(output.iterdir()): - raise FileExistsError("frontend output must be a fresh empty directory") - if (geometry_package is None) != (geometry_model is None): - raise ValueError("--geometry-package and --geometry-model must be supplied together") - if material_backend not in {"qwen-image", "firered"}: - raise ValueError("--material-backend must be qwen-image or firered") - from supratex.inference.runtime import configure_inference_runtime - configure_inference_runtime(attention=attention, low_vram=low_vram, device=device, - material_backend=material_backend) - if material_base is not None and firered_base is not None: - raise ValueError("--material-base and --firered-base cannot be combined") - from supratex.weights import resolve_model_dir - if material_model is not None and any(value is not None for value in - (material_weights, material_spec, material_checkpoint_manifest)): - raise ValueError("--material-model cannot be combined with explicit material artifact paths") - if material_backend == "qwen-image": - if any(value is not None for value in (firered_base, firered_manifest, material_weights, - material_spec, material_checkpoint_manifest)): - raise ValueError("FireRed base, manifest and explicit artifacts require --material-backend firered") - checkpoints = Path(__file__).resolve().parents[2] / "ckpts" - qwen_runtime, material_base, checkpoint, qwen_deps = validate_qwen( - qwen_runtime or checkpoints / "qwen-runtime", material_base or checkpoints / "qwen-base", - material_model or checkpoints / "qwen-material", deps=qwen_deps or checkpoints / "qwen-deps", - cache_dir=cache_dir, revision=revision, local_files_only=local_files_only) - material_model = checkpoint.parent - else: - if qwen_runtime is not None or qwen_deps is not None: - raise ValueError("--qwen-runtime and --qwen-deps require --material-backend qwen-image") - material_base = material_base or firered_base - if material_base is None: - raise ValueError("FireRed requires --material-base or --firered-base") - firered_package = firered_package or Path(__file__).resolve().parent / "material_model" - if not Path(moge_model).expanduser().is_file(): - moge_model = resolve_model_dir(moge_model, cache_dir=cache_dir, local_files_only=local_files_only, - allow_patterns=["model.pt"]) - else: - moge_model = Path(moge_model).expanduser() - if rmbg_model is not None: - rmbg_model = resolve_model_dir(rmbg_model, cache_dir=cache_dir, local_files_only=local_files_only) - output.mkdir(parents=True, exist_ok=True) - with span("frontend.preprocess"): - if resident is not None and not use_input_alpha: - import torch - with resident.pool.use('background', device=device) as background: - estimator = RMBG2(rmbg_model, device=device) - estimator._model, estimator._torch = background, torch - rgba, rgb, foreground_source = prepare_image(input_image, output_dir=output, - use_input_alpha=False, background_model=estimator) - else: - estimator = RMBG2(rmbg_model, device=device) if rmbg_model and not use_input_alpha else None - try: - rgba, rgb, foreground_source = prepare_image(input_image, output_dir=output, use_input_alpha=use_input_alpha, - background_model=estimator) - finally: - if estimator is not None: - estimator.close() - with span("frontend.camera"): - if resident is not None: - with resident.pool.use('camera', device=device) as camera_model: - camera = estimate(rgb, moge_model, moge_module, resolution_level, device, model=camera_model) - else: - camera = estimate(rgb, moge_model, moge_module, resolution_level, device) - camera["foreground_source"] = foreground_source - camera_path = output / "camera.json" - camera_path.write_text(json.dumps(camera, indent=2) + "\n", encoding="utf-8") - with span("frontend.geometry"): - geometry_result = {} - if geometry_package is not None: - geometry_result = run_geometry(geometry_package, geometry_model, output, output / "geometry", - seed=geometry_seed, initial_mesh_scale=geometry_initial_mesh_scale, - target_padding=geometry_target_padding, - max_iterations=geometry_max_iterations, cache_dir=cache_dir, - revision=geometry_revision, local_files_only=local_files_only, resident=resident) - geometry_path = output / "geometry.pt" - geometry_path.write_bytes(Path(geometry_result["geometry"]).read_bytes()) - camera_path.write_bytes(Path(geometry_result["camera"]).read_bytes()) - if record_material_provenance: - (output / "geometry_receipt.json").write_bytes(Path(geometry_result["receipt"]).read_bytes()) - with span("frontend.material"): - material_dir = output / "material" - albedo, mra = predict_material_maps( - rgba, rgb, material_dir, material_backend=material_backend, material_base=material_base, - material_model=material_model, qwen_runtime=qwen_runtime, qwen_deps=qwen_deps, - firered_package=firered_package, firered_manifest=firered_manifest, - material_weights=material_weights, material_spec=material_spec, - material_checkpoint_manifest=material_checkpoint_manifest, revision=revision, - cache_dir=cache_dir, local_files_only=local_files_only, device=device, key=key, - receipt_path=output / "material_receipt.json" if record_material_provenance else None, resident=resident) - with span("frontend.handoff"): - # The downstream texture stage consumes a single canonical work directory. - final_albedo, final_mra = output / "albedo.png", output / "mra.png" - final_albedo.write_bytes(albedo.read_bytes()) - final_mra.write_bytes(mra.read_bytes()) - if material_backend == "firered": - (output / "firered_albedo.png").write_bytes(albedo.read_bytes()) - (output / "firered_mra.png").write_bytes(mra.read_bytes()) - result = {"condition_rgba": str(rgba), "condition_rgb": str(rgb), "camera": str(camera_path), - "albedo": str(final_albedo), "mra": str(final_mra), "foreground_source": foreground_source, - "material_backend": material_backend} - if geometry_result: - result.update(geometry=str(geometry_path), geometry_receipt=geometry_result["receipt"]) - return result - - -def main() -> None: - parser = argparse.ArgumentParser(description="Run the complete Supra-Tex Image-to-Material Translation frontend.") - parser.add_argument("--input", type=Path, required=True) - parser.add_argument("--output", type=Path, required=True) - parser.add_argument("--moge-model", required=True, help="local MoGe checkpoint/directory or Hugging Face namespace/repo") - parser.add_argument("--moge-module", default="moge.model.v2") - parser.add_argument("--resolution-level", type=int, default=9) - parser.add_argument("--rmbg-model", help="local RMBG-2 directory or Hugging Face namespace/repo") - parser.add_argument("--use-input-alpha", action="store_true") - parser.add_argument("--firered-package", type=Path, default=Path(__file__).resolve().parent / "material_model") - parser.add_argument("--material-backend", choices=("qwen-image", "firered"), default="qwen-image") - parser.add_argument("--material-base", help="selected backend base directory or Hugging Face namespace/repo") - parser.add_argument("--qwen-runtime", type=Path, help="native Qwen source snapshot (default: ckpts/qwen-runtime)") - parser.add_argument("--qwen-deps", type=Path, help="native Qwen dependency directory (default: ckpts/qwen-deps)") - parser.add_argument("--firered-base", help="FireRed-only alias for --material-base") - parser.add_argument("--material-model", help="local material-model directory or Hugging Face namespace/repo") - parser.add_argument("--revision", help="material-model Hugging Face revision") - parser.add_argument("--cache-dir", type=Path) - parser.add_argument("--local-files-only", action="store_true") - parser.add_argument("--geometry-package", type=Path, help="source checkout for the original two-stage geometry generator") - parser.add_argument("--geometry-model", help="local geometry model directory or Hugging Face namespace/repo") - parser.add_argument("--geometry-revision", help="geometry model Hugging Face revision") - parser.add_argument("--geometry-initial-mesh-scale", type=float, default=1.0) - parser.add_argument("--geometry-target-padding", type=int, default=3) - parser.add_argument("--geometry-max-iterations", type=int, default=2) - parser.add_argument("--firered-manifest", type=Path) - parser.add_argument("--material-weights", type=Path) - parser.add_argument("--material-spec", type=Path) - parser.add_argument("--material-checkpoint-manifest", type=Path) - parser.add_argument("--key", default="input") - parser.add_argument("--record-material-provenance", action="store_true", help="write material bindings for optional dual-view inference") - parser.add_argument("--device", default="cuda") - parser.add_argument("--attention", choices=("auto", "flash_attn_3", "flash_attn", "sdpa", "xformers"), default="auto") - parser.add_argument("--low-vram", action=argparse.BooleanOptionalAction, default=None, - help="CPU offload for Qwen/3D; auto enabled on GPUs with up to 24 GiB") - args = parser.parse_args() - print(json.dumps(run(args.input, args.output, moge_model=args.moge_model, firered_package=args.firered_package, - material_backend=args.material_backend, material_base=args.material_base, - qwen_runtime=args.qwen_runtime, qwen_deps=args.qwen_deps, - firered_base=args.firered_base, firered_manifest=args.firered_manifest, - material_weights=args.material_weights, material_spec=args.material_spec, - material_checkpoint_manifest=args.material_checkpoint_manifest, - material_model=args.material_model, revision=args.revision, cache_dir=args.cache_dir, - local_files_only=args.local_files_only, - geometry_package=args.geometry_package, geometry_model=args.geometry_model, - geometry_revision=args.geometry_revision, - geometry_initial_mesh_scale=args.geometry_initial_mesh_scale, - geometry_target_padding=args.geometry_target_padding, - geometry_max_iterations=args.geometry_max_iterations, - rmbg_model=args.rmbg_model, use_input_alpha=args.use_input_alpha, device=args.device, - moge_module=args.moge_module, resolution_level=args.resolution_level, key=args.key, - record_material_provenance=args.record_material_provenance, - attention=args.attention, low_vram=args.low_vram), indent=2)) - - -if __name__ == "__main__": - main() +"""Runnable Image-to-Material Translation handoff builder.""" +from __future__ import annotations + +import argparse +import hashlib +import json +from pathlib import Path + +from .camera import estimate +from .firered import run as run_firered +from .preprocess import RMBG2, prepare_image +from .geometry import run as run_geometry +from .qwen_image import run as run_qwen, validate as validate_qwen + + +def predict_material_maps(rgba: Path, rgb: Path, output: Path, *, + material_backend: str, material_base: str | Path, + material_model: str | Path | None = None, + qwen_runtime: Path | None = None, qwen_deps: Path | None = None, + firered_package: Path | None = None, firered_manifest: Path | None = None, + material_weights: Path | None = None, material_spec: Path | None = None, + material_checkpoint_manifest: Path | None = None, + revision: str | None = None, cache_dir: Path | None = None, + local_files_only: bool = False, device: str = "cuda", key: str = "input", + receipt_path: Path | None = None, + cuda_visible_devices: str | None = None) -> tuple[Path, Path]: + """Predict either view with the same material backend and artifact selection.""" + if receipt_path is not None and firered_manifest is not None: + raise ValueError("material provenance does not support a custom FireRed manifest; " + "omit it to bind the current condition and use matching prompts for both views") + if receipt_path is not None: + from supratex.weights import resolve_model_dir + material_base = resolve_model_dir(material_base, cache_dir=cache_dir, local_files_only=local_files_only) + if material_model is not None: + patterns = (["adapter_model.safetensors", "adapter_config.json", + "checkpoint-0100000.pt", "checkpoint-0100000.complete.json"] + if material_backend == "qwen-image" else + ["full_transformer.safetensors", "joint_contract.json", "checkpoint_complete.json"]) + material_model = resolve_model_dir(material_model, revision=revision, cache_dir=cache_dir, + local_files_only=local_files_only, allow_patterns=patterns) + from supratex.backview.material_identity import artifact_identity + before_artifacts = artifact_identity({"base": material_base, "model": material_model, + "weights": material_weights, "spec": material_spec, + "checkpoint_manifest": material_checkpoint_manifest}) + if material_backend == "qwen-image": + albedo, mra = run_qwen(qwen_runtime, material_base, material_model, rgba, output, + deps=qwen_deps, cache_dir=cache_dir, revision=revision, + local_files_only=local_files_only, device=device, + cuda_visible_devices=cuda_visible_devices) + elif material_backend == "firered": + albedo, mra = run_firered(firered_package, material_base, rgb, output, + firered_manifest, key=key, weights=material_weights, + spec=material_spec, checkpoint_manifest=material_checkpoint_manifest, + material_model=material_model, revision=revision, + cache_dir=cache_dir, local_files_only=local_files_only) + else: + raise ValueError("material backend must be qwen-image or firered") + if receipt_path is not None: + def identity(value): + if value is None: + return None + path = Path(value).expanduser() + return str(path.resolve()) if path.exists() else str(value) + signature = {"backend": material_backend, "base": identity(material_base), + "model": identity(material_model), "revision": revision, + "runtime": identity(qwen_runtime if material_backend == "qwen-image" else firered_package), + "deps": identity(qwen_deps) if material_backend == "qwen-image" else None, + "weights": identity(material_weights), "spec": identity(material_spec), + "checkpoint_manifest": identity(material_checkpoint_manifest), "seed": 42} + from supratex.backview.material_identity import artifact_identity + signature["artifacts"] = artifact_identity(signature) + if signature["artifacts"] != before_artifacts: + raise RuntimeError("material artifacts changed during prediction; discard this view and use a new output") + # Image digests bind the receipt to the exact maps consumed downstream. + paths = {"condition_rgba.png": rgba, "condition_rgb.png": rgb, + "albedo.png": albedo, "mra.png": mra} + receipt = {"status": "complete", "backend": material_backend, "model_signature": signature, + "files": {name: hashlib.sha256(Path(path).read_bytes()).hexdigest() + for name, path in paths.items()}} + receipt_path = Path(receipt_path) + if receipt_path.exists(): + raise FileExistsError(receipt_path) + receipt_path.write_text(json.dumps(receipt, indent=2) + "\n", encoding="utf-8") + return albedo, mra + + +def _split_frontend_gpus(material_backend: str, device: str) -> tuple[str, str]: + """Pick the first two visible GPUs for geometry and Qwen material.""" + import os + import torch + + if material_backend != "qwen-image" or device != "cuda" or torch.cuda.device_count() < 2: + raise ValueError("--parallel-frontend requires the qwen-image backend, device 'cuda', " + "geometry generation, and at least two visible GPUs") + visible = os.environ.get("CUDA_VISIBLE_DEVICES") + names = [name.strip() for name in visible.split(",")] if visible else [str(index) for index in range(torch.cuda.device_count())] + return names[0], names[1] + + +def run(input_image: Path, output: Path, *, moge_model: str | Path, + firered_package: Path | None = None, firered_base: str | Path | None = None, + material_backend: str = "qwen-image", material_base: str | Path | None = None, + qwen_runtime: Path | None = None, qwen_deps: Path | None = None, + firered_manifest: Path | None = None, material_weights: Path | None = None, material_spec: Path | None = None, + material_checkpoint_manifest: Path | None = None, rmbg_model: str | Path | None = None, use_input_alpha: bool = False, + device: str = "cuda", moge_module: str = "moge.model.v2", resolution_level: int = 9, + key: str = "input", material_model: str | Path | None = None, revision: str | None = None, + cache_dir: Path | None = None, local_files_only: bool = False, + geometry_package: Path | None = None, geometry_model: str | Path | None = None, + geometry_revision: str | None = None, geometry_initial_mesh_scale: float = 1.0, + geometry_scale_search: bool = False, geometry_target_padding: int = 3, geometry_max_iterations: int = 2, + geometry_seed: int = 42, record_material_provenance: bool = False, + attention: str = "auto", low_vram: bool | None = None, + parallel_frontend: bool = False) -> dict[str, str]: + """Build the condition, canonical camera, and material-map handoff in one directory.""" + if record_material_provenance and firered_manifest is not None: + raise ValueError("material provenance does not support a custom FireRed manifest; " + "omit it to bind the current condition and use matching prompts for both views") + if output.exists() and any(output.iterdir()): + raise FileExistsError("frontend output must be a fresh empty directory") + if (geometry_package is None) != (geometry_model is None): + raise ValueError("--geometry-package and --geometry-model must be supplied together") + if material_backend not in {"qwen-image", "firered"}: + raise ValueError("--material-backend must be qwen-image or firered") + if parallel_frontend and geometry_package is None: + raise ValueError("--parallel-frontend requires geometry generation") + gpus = _split_frontend_gpus(material_backend, device) if parallel_frontend else None + from supratex.inference.runtime import configure_inference_runtime + configure_inference_runtime(attention=attention, low_vram=low_vram, device=device, + material_backend=material_backend) + if material_base is not None and firered_base is not None: + raise ValueError("--material-base and --firered-base cannot be combined") + from supratex.weights import resolve_model_dir + if material_model is not None and any(value is not None for value in + (material_weights, material_spec, material_checkpoint_manifest)): + raise ValueError("--material-model cannot be combined with explicit material artifact paths") + if material_backend == "qwen-image": + if any(value is not None for value in (firered_base, firered_manifest, material_weights, + material_spec, material_checkpoint_manifest)): + raise ValueError("FireRed base, manifest and explicit artifacts require --material-backend firered") + from supratex.pipeline_config import runtime_root + root = Path(__file__).resolve().parents[2] + checkpoints = root / "ckpts" + qwen_runtime, material_base, checkpoint, qwen_deps = validate_qwen( + qwen_runtime or runtime_root(root) / "qwen", material_base or checkpoints / "qwen-base", + material_model or checkpoints / "qwen-material", deps=qwen_deps, + cache_dir=cache_dir, revision=revision, local_files_only=local_files_only) + material_model = checkpoint.parent + else: + if qwen_runtime is not None or qwen_deps is not None: + raise ValueError("--qwen-runtime and --qwen-deps require --material-backend qwen-image") + material_base = material_base or firered_base + if material_base is None: + raise ValueError("FireRed requires --material-base or --firered-base") + firered_package = firered_package or Path(__file__).resolve().parent / "material_model" + if not Path(moge_model).expanduser().is_file(): + moge_model = resolve_model_dir(moge_model, cache_dir=cache_dir, local_files_only=local_files_only, + allow_patterns=["model.pt"]) + else: + moge_model = Path(moge_model).expanduser() + if rmbg_model is not None: + rmbg_model = resolve_model_dir(rmbg_model, cache_dir=cache_dir, local_files_only=local_files_only) + output.mkdir(parents=True, exist_ok=True) + estimator = RMBG2(rmbg_model, device=device) if rmbg_model else None + try: + rgba, rgb, foreground_source = prepare_image(input_image, output_dir=output, use_input_alpha=use_input_alpha, + background_model=estimator) + finally: + if estimator is not None: + estimator.close() + camera = estimate(rgb, moge_model, moge_module, resolution_level, device) + camera["foreground_source"] = foreground_source + camera_path = output / "camera.json" + camera_path.write_text(json.dumps(camera, indent=2) + "\n", encoding="utf-8") + # Material maps depend only on the foreground condition, not on geometry. With + # --parallel-frontend both subprocesses run concurrently, one GPU each. + geometry_result, geometry_job = {}, None + if geometry_package is not None: + def geometry_call() -> dict: + return run_geometry(geometry_package, geometry_model, output, output / "geometry", + seed=geometry_seed, initial_mesh_scale=geometry_initial_mesh_scale, + scale_search=geometry_scale_search, target_padding=geometry_target_padding, + max_iterations=geometry_max_iterations, cache_dir=cache_dir, + revision=geometry_revision, local_files_only=local_files_only, + cuda_visible_devices=gpus[0] if gpus else None) + if gpus: + from concurrent.futures import ThreadPoolExecutor + executor = ThreadPoolExecutor(max_workers=1) + geometry_job = executor.submit(geometry_call) + executor.shutdown(wait=False) + print(f"Frontend: geometry on GPU {gpus[0]}, material on GPU {gpus[1]} in parallel", flush=True) + else: + geometry_result = geometry_call() + material_dir = output / "material" + try: + albedo, mra = predict_material_maps( + rgba, rgb, material_dir, material_backend=material_backend, material_base=material_base, + material_model=material_model, qwen_runtime=qwen_runtime, qwen_deps=qwen_deps, + firered_package=firered_package, firered_manifest=firered_manifest, + material_weights=material_weights, material_spec=material_spec, + material_checkpoint_manifest=material_checkpoint_manifest, revision=revision, + cache_dir=cache_dir, local_files_only=local_files_only, device=device, key=key, + receipt_path=output / "material_receipt.json" if record_material_provenance else None, + cuda_visible_devices=gpus[1] if gpus else None) + finally: + if geometry_job is not None: + # Always wait, so a material failure never leaves geometry running. + geometry_result = geometry_job.result() + if geometry_package is not None: + geometry_path = output / "geometry.pt" + geometry_path.write_bytes(Path(geometry_result["geometry"]).read_bytes()) + camera_path.write_bytes(Path(geometry_result["camera"]).read_bytes()) + if record_material_provenance: + (output / "geometry_receipt.json").write_bytes(Path(geometry_result["receipt"]).read_bytes()) + # The downstream texture stage consumes a single canonical work directory. + final_albedo, final_mra = output / "albedo.png", output / "mra.png" + final_albedo.write_bytes(albedo.read_bytes()) + final_mra.write_bytes(mra.read_bytes()) + if material_backend == "firered": + (output / "firered_albedo.png").write_bytes(albedo.read_bytes()) + (output / "firered_mra.png").write_bytes(mra.read_bytes()) + result = {"condition_rgba": str(rgba), "condition_rgb": str(rgb), "camera": str(camera_path), + "albedo": str(final_albedo), "mra": str(final_mra), "foreground_source": foreground_source, + "material_backend": material_backend} + if geometry_result: + result.update(geometry=str(geometry_path), geometry_receipt=geometry_result["receipt"]) + return result + + +def main() -> None: + parser = argparse.ArgumentParser(description="Run the complete Supra-Tex Image-to-Material Translation frontend.") + parser.add_argument("--input", type=Path, required=True) + parser.add_argument("--output", type=Path, required=True) + parser.add_argument("--moge-model", required=True, help="local MoGe checkpoint/directory or Hugging Face namespace/repo") + parser.add_argument("--moge-module", default="moge.model.v2") + parser.add_argument("--resolution-level", type=int, default=9) + parser.add_argument("--rmbg-model", help="local RMBG-2 directory or Hugging Face namespace/repo") + parser.add_argument("--use-input-alpha", action="store_true") + parser.add_argument("--firered-package", type=Path, default=Path(__file__).resolve().parent / "material_model") + parser.add_argument("--material-backend", choices=("qwen-image", "firered"), default="qwen-image") + parser.add_argument("--material-base", help="selected backend base directory or Hugging Face namespace/repo") + parser.add_argument("--qwen-runtime", type=Path, help="native Qwen source snapshot (default: bundled runtime)") + parser.add_argument("--qwen-deps", type=Path, help="optional Qwen dependency overlay; otherwise use the active Python environment") + parser.add_argument("--firered-base", help="FireRed-only alias for --material-base") + parser.add_argument("--material-model", help="local material-model directory or Hugging Face namespace/repo") + parser.add_argument("--revision", help="material-model Hugging Face revision") + parser.add_argument("--cache-dir", type=Path) + parser.add_argument("--local-files-only", action="store_true") + parser.add_argument("--geometry-package", type=Path, help="source checkout for the original two-stage geometry generator") + parser.add_argument("--geometry-model", help="local geometry model directory or Hugging Face namespace/repo") + parser.add_argument("--geometry-revision", help="geometry model Hugging Face revision") + parser.add_argument("--geometry-initial-mesh-scale", type=float, default=1.0) + parser.add_argument("--geometry-scale-search", action="store_true", help="enable bounded mesh-scale calibration (default: Pixal3D fixed scale)") + parser.add_argument("--parallel-frontend", action="store_true", + help="run geometry and Qwen material concurrently on the first two visible GPUs") + parser.add_argument("--geometry-target-padding", type=int, default=3) + parser.add_argument("--geometry-max-iterations", type=int, default=2) + parser.add_argument("--firered-manifest", type=Path) + parser.add_argument("--material-weights", type=Path) + parser.add_argument("--material-spec", type=Path) + parser.add_argument("--material-checkpoint-manifest", type=Path) + parser.add_argument("--key", default="input") + parser.add_argument("--record-material-provenance", action="store_true", help="write material bindings for optional dual-view inference") + parser.add_argument("--device", default="cuda") + parser.add_argument("--attention", choices=("auto", "flash_attn_3", "flash_attn", "sdpa", "xformers"), default="auto") + parser.add_argument("--low-vram", action=argparse.BooleanOptionalAction, default=None, + help="CPU offload for Qwen/3D; auto enabled on GPUs with up to 24 GiB") + args = parser.parse_args() + print(json.dumps(run(args.input, args.output, moge_model=args.moge_model, firered_package=args.firered_package, + material_backend=args.material_backend, material_base=args.material_base, + qwen_runtime=args.qwen_runtime, qwen_deps=args.qwen_deps, + firered_base=args.firered_base, firered_manifest=args.firered_manifest, + material_weights=args.material_weights, material_spec=args.material_spec, + material_checkpoint_manifest=args.material_checkpoint_manifest, + material_model=args.material_model, revision=args.revision, cache_dir=args.cache_dir, + local_files_only=args.local_files_only, + geometry_package=args.geometry_package, geometry_model=args.geometry_model, + geometry_revision=args.geometry_revision, + geometry_initial_mesh_scale=args.geometry_initial_mesh_scale, + geometry_scale_search=args.geometry_scale_search, + parallel_frontend=args.parallel_frontend, + geometry_target_padding=args.geometry_target_padding, + geometry_max_iterations=args.geometry_max_iterations, + rmbg_model=args.rmbg_model, use_input_alpha=args.use_input_alpha, device=args.device, + moge_module=args.moge_module, resolution_level=args.resolution_level, key=args.key, + record_material_provenance=args.record_material_provenance, + attention=args.attention, low_vram=args.low_vram), indent=2)) + + +if __name__ == "__main__": + main() diff --git a/supratex/frontend/preprocess.py b/supratex/frontend/preprocess.py index c625c203cd9a5fe6a4c161222c0bb3b343ac852f..1492c32f9eff73b057a8170e12df05faf7d5d1bb 100644 --- a/supratex/frontend/preprocess.py +++ b/supratex/frontend/preprocess.py @@ -1,128 +1,128 @@ -"""Foreground preparation for the Supra-Tex Image-to-Material Translation stage.""" -from __future__ import annotations - -import argparse -from pathlib import Path -from typing import Protocol - -import numpy as np -from PIL import Image - - -class AlphaEstimator(Protocol): - def __call__(self, image: Image.Image) -> Image.Image: ... - - -class RMBG2: - """Local RMBG-2 foreground adapter using its documented Transformers form.""" - - def __init__(self, model: str | Path, *, device: str = "cuda", alpha_threshold: float | None = None): - self.model_path = Path(model) - if not self.model_path.is_dir(): - raise FileNotFoundError(f"RMBG-2 model directory does not exist: {self.model_path}") - if alpha_threshold is not None and not 0.0 <= alpha_threshold <= 1.0: - raise ValueError("alpha_threshold must be in [0, 1]") - self.device, self.alpha_threshold, self._model = device, alpha_threshold, None - - def _load(self): - if self._model is None: - try: - from transformers import AutoModelForImageSegmentation - import torch - except ImportError as error: # pragma: no cover - optional runtime - raise RuntimeError("RMBG-2 requires torch and transformers") from error - self._torch = torch - self._model = AutoModelForImageSegmentation.from_pretrained( - str(self.model_path), trust_remote_code=True, local_files_only=True - ).to(self.device).eval() - return self._model - - def close(self): - if self._model is not None: - self._model.cpu() - self._model = None - if str(self.device).startswith("cuda"): - self._torch.cuda.empty_cache() - - def __call__(self, image: Image.Image) -> Image.Image: - model, torch = self._load(), self._torch - rgb = image.convert("RGB") - resized = rgb.resize((1024, 1024), Image.Resampling.LANCZOS) - value = np.asarray(resized, dtype=np.float32) / 255.0 - value = (value - np.asarray((0.485, 0.456, 0.406), np.float32)) / np.asarray((0.229, 0.224, 0.225), np.float32) - tensor = torch.from_numpy(value).permute(2, 0, 1).unsqueeze(0).to(self.device) - with torch.inference_mode(): - prediction = model(tensor) - logits = prediction[-1] if isinstance(prediction, (tuple, list)) else prediction.logits - alpha = logits[0].squeeze().sigmoid().float().cpu().numpy() - mask = Image.fromarray(np.clip(alpha * 255.0, 0, 255).astype(np.uint8), mode="L") - mask = mask.resize(rgb.size, Image.Resampling.LANCZOS) - if self.alpha_threshold is not None: - cutoff = round(self.alpha_threshold * 255) - mask = mask.point(lambda value: 255 if value >= cutoff else 0) - rgb.putalpha(mask) - return rgb - - -def _crop_square(image: Image.Image) -> Image.Image: - alpha = np.asarray(image.convert("RGBA"), dtype=np.uint8)[..., 3] - foreground = np.argwhere(alpha > 0.8 * 255) - if not foreground.size: - raise ValueError("foreground alpha has no pixels above the 0.8 threshold") - x0, y0 = foreground[:, 1].min(), foreground[:, 0].min() - x1, y1 = foreground[:, 1].max(), foreground[:, 0].max() - side = int(max(x1 - x0, y1 - y0) * 1.1) - if side <= 0: - raise ValueError("foreground crop has zero extent") - cx, cy = (x0 + x1) / 2, (y0 + y1) / 2 - left, top = int(cx - side // 2), int(cy - side // 2) - return image.crop((left, top, left + side, top + side)) - - -def prepare_image(source: str | Path, *, output_dir: str | Path, use_input_alpha: bool = False, - already_cropped: bool = False, background_model: AlphaEstimator | None = None) -> tuple[Path, Path, str]: - """Write aligned condition RGBA/RGB frames and return their paths and provenance.""" - image = Image.open(source) - rgba_present = image.mode == "RGBA" and np.any(np.asarray(image)[..., 3] != 255) - if use_input_alpha and not rgba_present: - raise ValueError("--use-input-alpha requires a non-opaque RGBA input") - if not already_cropped: - scale = min(1.0, 1024.0 / max(image.size)) - if scale < 1: - image = image.resize((int(image.width * scale), int(image.height * scale)), Image.Resampling.LANCZOS) - if rgba_present: - rgba, provenance = image.convert("RGBA"), "input_alpha" - else: - if background_model is None: - raise ValueError("no foreground alpha supplied; pass an explicit RMBG-2 model") - rgba, provenance = background_model(image.convert("RGB")).convert("RGBA"), "rmbg2" - if not already_cropped: - rgba = _crop_square(rgba) - data = np.asarray(rgba, dtype=np.float32) / 255.0 - rgb = Image.fromarray((np.clip(data[..., :3] * data[..., 3:4], 0, 1) * 255).astype(np.uint8)) - out = Path(output_dir) - out.mkdir(parents=True, exist_ok=True) - rgba_path, rgb_path = out / "condition_rgba.png", out / "condition_rgb.png" - rgba.save(rgba_path) - rgb.save(rgb_path) - return rgba_path, rgb_path, provenance + ("_precropped" if already_cropped else "") - - -def main() -> None: - parser = argparse.ArgumentParser(description="Prepare an Image-to-Material Translation condition image.") - parser.add_argument("--input", type=Path, required=True) - parser.add_argument("--output-dir", type=Path, required=True) - parser.add_argument("--use-input-alpha", action="store_true") - parser.add_argument("--already-cropped", action="store_true") - parser.add_argument("--rmbg-model", type=Path, help="Local RMBG-2 model directory for opaque images.") - parser.add_argument("--device", default="cuda") - parser.add_argument("--alpha-threshold", type=float) - args = parser.parse_args() - estimator = RMBG2(args.rmbg_model, device=args.device, alpha_threshold=args.alpha_threshold) if args.rmbg_model else None - rgba, rgb, provenance = prepare_image(args.input, output_dir=args.output_dir, use_input_alpha=args.use_input_alpha, - already_cropped=args.already_cropped, background_model=estimator) - print(f"condition_rgba={rgba}\ncondition_rgb={rgb}\nforeground_source={provenance}") - - -if __name__ == "__main__": - main() +"""Foreground preparation for the Supra-Tex Image-to-Material Translation stage.""" +from __future__ import annotations + +import argparse +from pathlib import Path +from typing import Protocol + +import numpy as np +from PIL import Image + + +class AlphaEstimator(Protocol): + def __call__(self, image: Image.Image) -> Image.Image: ... + + +class RMBG2: + """Local RMBG-2 foreground adapter using its documented Transformers form.""" + + def __init__(self, model: str | Path, *, device: str = "cuda", alpha_threshold: float | None = None): + self.model_path = Path(model) + if not self.model_path.is_dir(): + raise FileNotFoundError(f"RMBG-2 model directory does not exist: {self.model_path}") + if alpha_threshold is not None and not 0.0 <= alpha_threshold <= 1.0: + raise ValueError("alpha_threshold must be in [0, 1]") + self.device, self.alpha_threshold, self._model = device, alpha_threshold, None + + def _load(self): + if self._model is None: + try: + from transformers import AutoModelForImageSegmentation + import torch + except ImportError as error: # pragma: no cover - optional runtime + raise RuntimeError("RMBG-2 requires torch and transformers") from error + self._torch = torch + self._model = AutoModelForImageSegmentation.from_pretrained( + str(self.model_path), trust_remote_code=True, local_files_only=True + ).to(self.device).eval() + return self._model + + def close(self): + if self._model is not None: + self._model.cpu() + self._model = None + if str(self.device).startswith("cuda"): + self._torch.cuda.empty_cache() + + def __call__(self, image: Image.Image) -> Image.Image: + model, torch = self._load(), self._torch + rgb = image.convert("RGB") + resized = rgb.resize((1024, 1024), Image.Resampling.LANCZOS) + value = np.asarray(resized, dtype=np.float32) / 255.0 + value = (value - np.asarray((0.485, 0.456, 0.406), np.float32)) / np.asarray((0.229, 0.224, 0.225), np.float32) + tensor = torch.from_numpy(value).permute(2, 0, 1).unsqueeze(0).to(self.device) + with torch.inference_mode(): + prediction = model(tensor) + logits = prediction[-1] if isinstance(prediction, (tuple, list)) else prediction.logits + alpha = logits[0].squeeze().sigmoid().float().cpu().numpy() + mask = Image.fromarray(np.clip(alpha * 255.0, 0, 255).astype(np.uint8), mode="L") + mask = mask.resize(rgb.size, Image.Resampling.LANCZOS) + if self.alpha_threshold is not None: + cutoff = round(self.alpha_threshold * 255) + mask = mask.point(lambda value: 255 if value >= cutoff else 0) + rgb.putalpha(mask) + return rgb + + +def _crop_square(image: Image.Image) -> Image.Image: + alpha = np.asarray(image.convert("RGBA"), dtype=np.uint8)[..., 3] + foreground = np.argwhere(alpha > 0.8 * 255) + if not foreground.size: + raise ValueError("foreground alpha has no pixels above the 0.8 threshold") + x0, y0 = foreground[:, 1].min(), foreground[:, 0].min() + x1, y1 = foreground[:, 1].max(), foreground[:, 0].max() + side = int(max(x1 - x0, y1 - y0) * 1.1) + if side <= 0: + raise ValueError("foreground crop has zero extent") + cx, cy = (x0 + x1) / 2, (y0 + y1) / 2 + # Same float box as Pixal3D's preprocess_image; PIL rounds it identically. + return image.crop((cx - side // 2, cy - side // 2, cx + side // 2, cy + side // 2)) + + +def prepare_image(source: str | Path, *, output_dir: str | Path, use_input_alpha: bool = False, + already_cropped: bool = False, background_model: AlphaEstimator | None = None) -> tuple[Path, Path, str]: + """Write aligned condition RGBA/RGB frames and return their paths and provenance.""" + image = Image.open(source) + rgba_present = image.mode == "RGBA" and np.any(np.asarray(image)[..., 3] != 255) + if use_input_alpha and not rgba_present: + raise ValueError("--use-input-alpha requires a non-opaque RGBA input") + if not already_cropped: + scale = min(1.0, 1024.0 / max(image.size)) + if scale < 1: + image = image.resize((int(image.width * scale), int(image.height * scale)), Image.Resampling.LANCZOS) + if rgba_present: + rgba, provenance = image.convert("RGBA"), "input_alpha" + else: + if background_model is None: + raise ValueError("no foreground alpha supplied; pass an explicit RMBG-2 model") + rgba, provenance = background_model(image.convert("RGB")).convert("RGBA"), "rmbg2" + if not already_cropped: + rgba = _crop_square(rgba) + data = np.asarray(rgba, dtype=np.float32) / 255.0 + rgb = Image.fromarray((np.clip(data[..., :3] * data[..., 3:4], 0, 1) * 255).astype(np.uint8)) + out = Path(output_dir) + out.mkdir(parents=True, exist_ok=True) + rgba_path, rgb_path = out / "condition_rgba.png", out / "condition_rgb.png" + rgba.save(rgba_path) + rgb.save(rgb_path) + return rgba_path, rgb_path, provenance + ("_precropped" if already_cropped else "") + + +def main() -> None: + parser = argparse.ArgumentParser(description="Prepare an Image-to-Material Translation condition image.") + parser.add_argument("--input", type=Path, required=True) + parser.add_argument("--output-dir", type=Path, required=True) + parser.add_argument("--use-input-alpha", action="store_true") + parser.add_argument("--already-cropped", action="store_true") + parser.add_argument("--rmbg-model", type=Path, help="Local RMBG-2 model directory for opaque images.") + parser.add_argument("--device", default="cuda") + parser.add_argument("--alpha-threshold", type=float) + args = parser.parse_args() + estimator = RMBG2(args.rmbg_model, device=args.device, alpha_threshold=args.alpha_threshold) if args.rmbg_model else None + rgba, rgb, provenance = prepare_image(args.input, output_dir=args.output_dir, use_input_alpha=args.use_input_alpha, + already_cropped=args.already_cropped, background_model=estimator) + print(f"condition_rgba={rgba}\ncondition_rgb={rgb}\nforeground_source={provenance}") + + +if __name__ == "__main__": + main() diff --git a/supratex/frontend/qwen_image.py b/supratex/frontend/qwen_image.py index 2fb9edde8a8a2b1171d01c98fed317c28f6fe6fe..0e1d57c14ec6fedbcd8e7d5666260262d2b71e6b 100644 --- a/supratex/frontend/qwen_image.py +++ b/supratex/frontend/qwen_image.py @@ -1,186 +1,180 @@ -"""Qwen Image 2.1 material adapter inference through its native CLI.""" -from __future__ import annotations - -import json -import math -import os -from pathlib import Path -import subprocess -import sys - -from supratex.weights import resolve_model_dir - - -def validate_adapter(checkpoint: Path) -> dict: - """Validate a native adapter config and tensor headers without loading tensors.""" - from safetensors import safe_open - - checkpoint = Path(checkpoint) - config = json.loads(checkpoint.with_name("adapter_config.json").read_text()) - if (not isinstance(config, dict) or config.get("format") != "supratex_qwen_material_adapter" - or type(config.get("schema_version")) is not int or config["schema_version"] != 1 - or config.get("adapter_type") != "qwen_image_21_pbr_endpoint" - or config.get("base_model_name_or_path") != "Qwen/Qwen-Image-2.1" - or config.get("weight_file") != checkpoint.name - or type(config.get("source_step")) is not int or config["source_step"] != 100000): - raise ValueError("unsupported Qwen material adapter config") - inference = config.get("inference_config") - booleans = {"lora_output_projection", "omit_all_valid_mask", "bf16_autocast", "deterministic_algorithms"} - required = booleans | {"resolution", "noise_seed", "lora_rank", "lora_alpha", "attention_backend"} - if not isinstance(inference, dict) or set(inference) != required: - raise ValueError("Qwen adapter requires complete inference config") - if (type(inference["resolution"]) is not int or inference["resolution"] != 1024 - or type(inference["noise_seed"]) is not int or inference["noise_seed"] != 42 - or inference["attention_backend"] != "torch" - or type(inference["lora_rank"]) is not int or inference["lora_rank"] <= 0 - or type(inference["lora_alpha"]) not in {int, float} - or not math.isfinite(inference["lora_alpha"]) or inference["lora_alpha"] <= 0 - or any(type(inference[key]) is not bool for key in booleans)): - raise ValueError("invalid Qwen adapter inference config") - tensors = config.get("tensors") - if not isinstance(tensors, dict) or not tensors: - raise ValueError("Qwen adapter requires a nonempty tensor manifest") - with safe_open(str(checkpoint), framework="pt", device="cpu") as reader: - metadata = reader.metadata() or {} - if (metadata.get("format") != "pt" or metadata.get("adapter_format") != config["format"] - or metadata.get("schema_version") != str(config["schema_version"]) - or set(reader.keys()) != set(tensors)): - raise ValueError("Qwen adapter tensor keys or metadata mismatch") - for name, specification in tensors.items(): - header = reader.get_slice(name) - if (not isinstance(specification, dict) - or not isinstance(specification.get("shape"), list) - or any(type(size) is not int or size < 0 for size in specification["shape"]) - or specification["shape"] != header.get_shape() - or specification.get("dtype") not in {"F64", "F32", "F16", "BF16"} - or specification["dtype"] != header.get_dtype()): - raise ValueError(f"Qwen adapter tensor shape or dtype mismatch: {name}") - return config - - -def validate(runtime: Path, base_model: str | Path, model_source: str | Path, *, - deps: Path | None = None, cache_dir: Path | None = None, - revision: str | None = None, local_files_only: bool = False - ) -> tuple[Path, Path, Path, Path | None]: - """Check local artifacts without importing the runtime or loading tensors.""" - runtime = Path(runtime).expanduser().resolve() - for name in ("infer.py", "train.py", "validate_local.py", "encode_cache.py"): - path = runtime / "examples/qwen_image_21/pbr" / name - if not path.is_file(): - raise FileNotFoundError(f"Qwen native runtime is incomplete: {path}") - if not (runtime / "diffsynth").is_dir(): - raise FileNotFoundError(runtime / "diffsynth") - if deps is not None: - deps = Path(deps).expanduser().resolve() - if not deps.is_dir(): - raise NotADirectoryError(deps) - root = resolve_model_dir(model_source, revision=revision, cache_dir=cache_dir, - local_files_only=local_files_only, - allow_patterns=["adapter_model.safetensors", "adapter_config.json", - "checkpoint-0100000.pt", "checkpoint-0100000.complete.json"]) - checkpoint = root / "adapter_model.safetensors" - if checkpoint.exists() or (root / "adapter_config.json").exists(): - validate_adapter(checkpoint) - else: - checkpoint = root / "checkpoint-0100000.pt" - marker = json.loads(checkpoint.with_suffix(".complete.json").read_text()) - if (marker.get("status") != "complete" or marker.get("step") != 100000 - or marker.get("bytes") != checkpoint.stat().st_size or checkpoint.stat().st_size == 0): - raise ValueError("Qwen original checkpoint requires a matching complete record at step 100000") - base = resolve_model_dir(base_model, cache_dir=cache_dir, local_files_only=local_files_only) - for name in ("preprocessor_config.json", "tokenizer_config.json", "tokenizer.json", "chat_template.jinja"): - if not (base / "processor" / name).is_file(): - raise FileNotFoundError(base / "processor" / name) - for component in ("vae", "text_encoder", "transformer"): - directory = base / component - if not (directory / "config.json").is_file(): - raise FileNotFoundError(directory / "config.json") - weights = list(directory.glob("*.safetensors")) - if not weights or any(path.stat().st_size == 0 for path in weights): - raise ValueError(f"Qwen base component has missing or empty weights: {directory}") - for index in directory.glob("*.safetensors.index.json"): - shards = set(json.loads(index.read_text()).get("weight_map", {}).values()) - if not shards or any(not (directory / name).is_file() for name in shards): - raise ValueError(f"Qwen base component has missing indexed shards: {index}") - return runtime, base, checkpoint, deps - - -def run(runtime: Path, base_model: str | Path, model_source: str | Path, - condition: Path, output: Path, *, deps: Path | None = None, - cache_dir: Path | None = None, revision: str | None = None, - local_files_only: bool = False, device: str = "cuda:0", resident=None) -> tuple[Path, Path]: - """Preserve the frontend canvas and translate native MR into completion MRA.""" - from PIL import Image, ImageChops - - # The native CLI calls torch.cuda.set_device(), which requires an index. - device = "cuda:0" if device == "cuda" else device - output = Path(output).resolve() - targets = [output / name for name in ("native", "input_rgba.png", "inference.log", - "albedo.png", "mra.png", "qwen_receipt.json")] - if any(path.exists() for path in targets): - raise FileExistsError("Qwen material outputs already exist; select a fresh output directory") - runtime, base, checkpoint, deps = validate( - runtime, base_model, model_source, deps=deps, cache_dir=cache_dir, - revision=revision, local_files_only=local_files_only) - condition = Path(condition).resolve() - with Image.open(condition) as image: - if image.mode != "RGBA" or image.width != image.height: - raise ValueError("Qwen condition must be the frontend's square straight RGBA canvas") - source_size = image.size - square = image.copy() if image.size == (1024, 1024) else image.resize((1024, 1024), Image.Resampling.LANCZOS) - output.mkdir(parents=True, exist_ok=True) - input_path, native = output / "input_rgba.png", output / "native" - square.save(input_path) - env = os.environ.copy() - env.update(PYTHONPATH=os.pathsep.join(str(path) for path in (runtime, deps) if path is not None), - PYTHONNOUSERSITE="1", PYTHONDONTWRITEBYTECODE="1", PYTHONUNBUFFERED="1", - HF_HUB_OFFLINE="1", TRANSFORMERS_OFFLINE="1", TOKENIZERS_PARALLELISM="false") - command = [sys.executable, str(runtime / "examples/qwen_image_21/pbr/infer.py"), - "--checkpoint", str(checkpoint), "--model-root", str(base), - "--input", str(input_path), "--output-dir", str(native), "--device", device] - if resident is not None: - import torch - deterministic = torch.are_deterministic_algorithms_enabled() - warn_only = torch.is_deterministic_algorithms_warn_only_enabled() - try: - resident.run_material(input_path, native, device=device) - finally: - torch.use_deterministic_algorithms(deterministic, warn_only=warn_only) - else: - with (output / "inference.log").open("w", encoding="utf-8") as log: - subprocess.run(command, check=True, env=env, stdout=log, stderr=subprocess.STDOUT, cwd=runtime) - report = json.loads((native / "REPORT.json").read_text()) - config = report.get("config", {}) - if (report.get("status") != "complete" or report.get("checkpoint_step") != 100000 - or config.get("resolution") != 1024 or config.get("noise_seed") != 42 - or config.get("attention_backend") != "torch"): - raise ValueError("Qwen native report does not match the original 100k inference contract") - coverage = square.getchannel("A") - maps = [] - for name in ("albedo", "mra"): - with Image.open(native / f"{name}.png") as image: - if image.mode != "RGBA" or image.size != (1024, 1024): - raise ValueError(f"Qwen native {name} must be RGBA 1024x1024") - if ImageChops.difference(image.getchannel("A"), coverage).getbbox() is not None: - raise ValueError(f"Qwen native {name} coverage differs from its input") - if name == "mra": - if image.getchannel("B").getextrema() != (255, 255): - raise ValueError("Qwen native MR.B must be the constant 255 placeholder") - maps.append(Image.merge("RGB", (image.getchannel("R"), image.getchannel("G"), coverage))) - else: - maps.append(image.convert("RGB")) - albedo, mra = output / "albedo.png", output / "mra.png" - for image, path in zip(maps, (albedo, mra)): - image.save(path) - receipt = { - "backend": "qwen-image", "checkpoint": str(checkpoint), "checkpoint_step": 100000, - "base_model": str(base), "runtime": str(runtime), "condition": str(condition), - "source_size": list(source_size), "canvas_size": [1024, 1024], - "preprocessing": "square canvas resized once with LANCZOS; straight RGBA; no crop or compositing", - "native_report": str(native / "REPORT.json"), - "albedo": "sRGB; no alpha premultiplication", - "mra": "R=linear metallic, G=linear roughness, B=input coverage; no predicted material transparency", - "noise_seed": 42, "attention_backend": "torch", "status": "complete", - } - (output / "qwen_receipt.json").write_text(json.dumps(receipt, indent=2) + "\n") - return albedo, mra +"""Qwen Image 2.1 material adapter inference through its native CLI.""" +from __future__ import annotations + +import json +import math +import os +from pathlib import Path +import subprocess +import sys + +from supratex.weights import resolve_model_dir + + +def validate_adapter(checkpoint: Path) -> dict: + """Validate a native adapter config and tensor headers without loading tensors.""" + from safetensors import safe_open + + checkpoint = Path(checkpoint) + config = json.loads(checkpoint.with_name("adapter_config.json").read_text()) + if (not isinstance(config, dict) or config.get("format") != "supratex_qwen_material_adapter" + or type(config.get("schema_version")) is not int or config["schema_version"] != 1 + or config.get("adapter_type") != "qwen_image_21_pbr_endpoint" + or config.get("base_model_name_or_path") != "Qwen/Qwen-Image-2.1" + or config.get("weight_file") != checkpoint.name + or type(config.get("source_step")) is not int or config["source_step"] != 100000): + raise ValueError("unsupported Qwen material adapter config") + inference = config.get("inference_config") + booleans = {"lora_output_projection", "omit_all_valid_mask", "bf16_autocast", "deterministic_algorithms"} + required = booleans | {"resolution", "noise_seed", "lora_rank", "lora_alpha", "attention_backend"} + if not isinstance(inference, dict) or set(inference) != required: + raise ValueError("Qwen adapter requires complete inference config") + if (type(inference["resolution"]) is not int or inference["resolution"] != 1024 + or type(inference["noise_seed"]) is not int or inference["noise_seed"] != 42 + or inference["attention_backend"] != "torch" + or type(inference["lora_rank"]) is not int or inference["lora_rank"] <= 0 + or type(inference["lora_alpha"]) not in {int, float} + or not math.isfinite(inference["lora_alpha"]) or inference["lora_alpha"] <= 0 + or any(type(inference[key]) is not bool for key in booleans)): + raise ValueError("invalid Qwen adapter inference config") + tensors = config.get("tensors") + if not isinstance(tensors, dict) or not tensors: + raise ValueError("Qwen adapter requires a nonempty tensor manifest") + with safe_open(str(checkpoint), framework="pt", device="cpu") as reader: + metadata = reader.metadata() or {} + if (metadata.get("format") != "pt" or metadata.get("adapter_format") != config["format"] + or metadata.get("schema_version") != str(config["schema_version"]) + or set(reader.keys()) != set(tensors)): + raise ValueError("Qwen adapter tensor keys or metadata mismatch") + for name, specification in tensors.items(): + header = reader.get_slice(name) + if (not isinstance(specification, dict) + or not isinstance(specification.get("shape"), list) + or any(type(size) is not int or size < 0 for size in specification["shape"]) + or specification["shape"] != header.get_shape() + or specification.get("dtype") not in {"F64", "F32", "F16", "BF16"} + or specification["dtype"] != header.get_dtype()): + raise ValueError(f"Qwen adapter tensor shape or dtype mismatch: {name}") + return config + + +def validate(runtime: Path, base_model: str | Path, model_source: str | Path, *, + deps: Path | None = None, cache_dir: Path | None = None, + revision: str | None = None, local_files_only: bool = False + ) -> tuple[Path, Path, Path, Path | None]: + """Check local artifacts without importing the runtime or loading tensors.""" + runtime = Path(runtime).expanduser().resolve() + for name in ("infer.py", "inference_utils.py", "inference_offload.py"): + path = runtime / "examples/qwen_image_21/pbr" / name + if not path.is_file(): + raise FileNotFoundError(f"Qwen native runtime is incomplete: {path}") + if not (runtime / "diffsynth").is_dir(): + raise FileNotFoundError(runtime / "diffsynth") + if deps is not None: + deps = Path(deps).expanduser().resolve() + if not deps.is_dir(): + raise NotADirectoryError(deps) + root = resolve_model_dir(model_source, revision=revision, cache_dir=cache_dir, + local_files_only=local_files_only, + allow_patterns=["adapter_model.safetensors", "adapter_config.json", + "checkpoint-0100000.pt", "checkpoint-0100000.complete.json"]) + checkpoint = root / "adapter_model.safetensors" + if checkpoint.exists() or (root / "adapter_config.json").exists(): + validate_adapter(checkpoint) + else: + checkpoint = root / "checkpoint-0100000.pt" + marker = json.loads(checkpoint.with_suffix(".complete.json").read_text()) + if (marker.get("status") != "complete" or marker.get("step") != 100000 + or marker.get("bytes") != checkpoint.stat().st_size or checkpoint.stat().st_size == 0): + raise ValueError("Qwen original checkpoint requires a matching complete record at step 100000") + base = resolve_model_dir(base_model, cache_dir=cache_dir, local_files_only=local_files_only) + for name in ("preprocessor_config.json", "tokenizer_config.json", "tokenizer.json", "chat_template.jinja"): + if not (base / "processor" / name).is_file(): + raise FileNotFoundError(base / "processor" / name) + for component in ("vae", "text_encoder", "transformer"): + directory = base / component + if not (directory / "config.json").is_file(): + raise FileNotFoundError(directory / "config.json") + weights = list(directory.glob("*.safetensors")) + if not weights or any(path.stat().st_size == 0 for path in weights): + raise ValueError(f"Qwen base component has missing or empty weights: {directory}") + for index in directory.glob("*.safetensors.index.json"): + shards = set(json.loads(index.read_text()).get("weight_map", {}).values()) + if not shards or any(not (directory / name).is_file() for name in shards): + raise ValueError(f"Qwen base component has missing indexed shards: {index}") + return runtime, base, checkpoint, deps + + +def run(runtime: Path, base_model: str | Path, model_source: str | Path, + condition: Path, output: Path, *, deps: Path | None = None, + cache_dir: Path | None = None, revision: str | None = None, + local_files_only: bool = False, device: str = "cuda:0", + cuda_visible_devices: str | None = None) -> tuple[Path, Path]: + """Preserve the frontend canvas and translate native MR into completion MRA.""" + from PIL import Image, ImageChops + + # The native CLI calls torch.cuda.set_device(), which requires an index. + device = "cuda:0" if device == "cuda" else device + output = Path(output).resolve() + targets = [output / name for name in ("native", "input_rgba.png", "inference.log", + "albedo.png", "mra.png", "qwen_receipt.json")] + if any(path.exists() for path in targets): + raise FileExistsError("Qwen material outputs already exist; select a fresh output directory") + runtime, base, checkpoint, deps = validate( + runtime, base_model, model_source, deps=deps, cache_dir=cache_dir, + revision=revision, local_files_only=local_files_only) + condition = Path(condition).resolve() + with Image.open(condition) as image: + if image.mode != "RGBA" or image.width != image.height: + raise ValueError("Qwen condition must be the frontend's square straight RGBA canvas") + source_size = image.size + square = image.copy() if image.size == (1024, 1024) else image.resize((1024, 1024), Image.Resampling.LANCZOS) + output.mkdir(parents=True, exist_ok=True) + input_path, native = output / "input_rgba.png", output / "native" + square.save(input_path) + env = os.environ.copy() + env.update(PYTHONPATH=os.pathsep.join(str(path) for path in (runtime, deps) if path is not None), + PYTHONNOUSERSITE="1", PYTHONDONTWRITEBYTECODE="1", PYTHONUNBUFFERED="1", + HF_HUB_OFFLINE="1", TRANSFORMERS_OFFLINE="1", TOKENIZERS_PARALLELISM="false") + if cuda_visible_devices is not None: + env["CUDA_VISIBLE_DEVICES"] = cuda_visible_devices + command = [sys.executable, str(runtime / "examples/qwen_image_21/pbr/infer.py"), + "--checkpoint", str(checkpoint), "--model-root", str(base), + "--input", str(input_path), "--output-dir", str(native), "--device", device] + with (output / "inference.log").open("w", encoding="utf-8") as log: + subprocess.run(command, check=True, env=env, stdout=log, stderr=subprocess.STDOUT, cwd=runtime) + report = json.loads((native / "REPORT.json").read_text()) + config = report.get("config", {}) + if (report.get("status") != "complete" or report.get("checkpoint_step") != 100000 + or config.get("resolution") != 1024 or config.get("noise_seed") != 42 + or config.get("attention_backend") != "torch"): + raise ValueError("Qwen native report does not match the original 100k inference contract") + coverage = square.getchannel("A") + maps = [] + for name in ("albedo", "mra"): + with Image.open(native / f"{name}.png") as image: + if image.mode != "RGBA" or image.size != (1024, 1024): + raise ValueError(f"Qwen native {name} must be RGBA 1024x1024") + if ImageChops.difference(image.getchannel("A"), coverage).getbbox() is not None: + raise ValueError(f"Qwen native {name} coverage differs from its input") + if name == "mra": + if image.getchannel("B").getextrema() != (255, 255): + raise ValueError("Qwen native MR.B must be the constant 255 placeholder") + maps.append(Image.merge("RGB", (image.getchannel("R"), image.getchannel("G"), coverage))) + else: + maps.append(image.convert("RGB")) + albedo, mra = output / "albedo.png", output / "mra.png" + for image, path in zip(maps, (albedo, mra)): + image.save(path) + receipt = { + "backend": "qwen-image", "checkpoint": str(checkpoint), "checkpoint_step": 100000, + "base_model": str(base), "runtime": str(runtime), "condition": str(condition), + "source_size": list(source_size), "canvas_size": [1024, 1024], + "preprocessing": "square canvas resized once with LANCZOS; straight RGBA; no crop or compositing", + "native_report": str(native / "REPORT.json"), + "albedo": "sRGB; no alpha premultiplication", + "mra": "R=linear metallic, G=linear roughness, B=input coverage; no predicted material transparency", + "noise_seed": 42, "attention_backend": "torch", "status": "complete", + } + (output / "qwen_receipt.json").write_text(json.dumps(receipt, indent=2) + "\n") + return albedo, mra diff --git a/supratex/inference/_kernel_config_h800.py b/supratex/inference/_kernel_config_h800.py index 605c8092bb52910ed31d966646cba8f0a78d08b9..93d150f8901460e95437699e2905d70e5fc33ebe 100644 --- a/supratex/inference/_kernel_config_h800.py +++ b/supratex/inference/_kernel_config_h800.py @@ -1,1091 +1,1091 @@ -"""Measured H800 forward-convolution launch choices for kernel_config. - -Captured from accepted 2026-10-10 geometry/completion inference and the matching -encoder profile. Accepted end-to-end entries take precedence on conflicts. -Only measured shape buckets are present; all other keys retain autotuning. -This Python data module is included automatically in source/wheel packages. -""" - -CONFIG = {'NVIDIA H800': {'flex_gemm.kernels.triton.spconv.sparse_submanifold_conv_fwd_masked_implicit_gemm.sparse_submanifold_conv_fwd_masked_implicit_gemm_kernel': {"(10, 1024, 1024, 27, True, 'torch.float16', 'torch.float16', 'torch.float16', 'torch.uint32', 'torch.int64', 'torch.float16')": {'kwargs': {'B1': 128, - 'B2': 256, - 'BK': 64}, - 'maxnreg': None, - 'num_buffers_warp_spec': 0, - 'num_consumer_groups': 0, - 'num_ctas': 1, - 'num_stages': 3, - 'num_warps': 8, - 'pre_hook': None, - 'reg_dec_producer': 0, - 'reg_inc_consumer': 0}, - "(10, 1024, 4096, 27, True, 'torch.float16', 'torch.float16', 'torch.float16', 'torch.uint32', 'torch.int64', 'torch.float16')": {'kwargs': {'B1': 128, - 'B2': 256, - 'BK': 64}, - 'maxnreg': None, - 'num_buffers_warp_spec': 0, - 'num_consumer_groups': 0, - 'num_ctas': 1, - 'num_stages': 3, - 'num_warps': 8, - 'pre_hook': None, - 'reg_dec_producer': 0, - 'reg_inc_consumer': 0}, - "(12, 1024, 1024, 27, True, 'torch.float16', 'torch.float16', 'torch.float16', 'torch.uint32', 'torch.int64', 'torch.float16')": {'kwargs': {'B1': 128, - 'B2': 256, - 'BK': 64}, - 'maxnreg': None, - 'num_buffers_warp_spec': 0, - 'num_consumer_groups': 0, - 'num_ctas': 1, - 'num_stages': 3, - 'num_warps': 8, - 'pre_hook': None, - 'reg_dec_producer': 0, - 'reg_inc_consumer': 0}, - "(13, 1024, 1024, 27, True, 'torch.float16', 'torch.float16', 'torch.float16', 'torch.uint32', 'torch.int64', 'torch.float16')": {'kwargs': {'B1': 128, - 'B2': 256, - 'BK': 64}, - 'maxnreg': None, - 'num_buffers_warp_spec': 0, - 'num_consumer_groups': 0, - 'num_ctas': 1, - 'num_stages': 3, - 'num_warps': 8, - 'pre_hook': None, - 'reg_dec_producer': 0, - 'reg_inc_consumer': 0}, - "(13, 1024, 4096, 27, True, 'torch.float16', 'torch.float16', 'torch.float16', 'torch.uint32', 'torch.int64', 'torch.float16')": {'kwargs': {'B1': 128, - 'B2': 256, - 'BK': 64}, - 'maxnreg': None, - 'num_buffers_warp_spec': 0, - 'num_consumer_groups': 0, - 'num_ctas': 1, - 'num_stages': 3, - 'num_warps': 8, - 'pre_hook': None, - 'reg_dec_producer': 0, - 'reg_inc_consumer': 0}, - "(13, 512, 128, 27, True, 'torch.float16', 'torch.float16', 'torch.float16', 'torch.uint32', 'torch.int64', 'torch.float16')": {'kwargs': {'B1': 64, - 'B2': 128, - 'BK': 32}, - 'maxnreg': None, - 'num_buffers_warp_spec': 0, - 'num_consumer_groups': 0, - 'num_ctas': 1, - 'num_stages': 4, - 'num_warps': 4, - 'pre_hook': None, - 'reg_dec_producer': 0, - 'reg_inc_consumer': 0}, - "(13, 512, 2048, 27, True, 'torch.float16', 'torch.float16', 'torch.float16', 'torch.uint32', 'torch.int64', 'torch.float16')": {'kwargs': {'B1': 128, - 'B2': 256, - 'BK': 64}, - 'maxnreg': None, - 'num_buffers_warp_spec': 0, - 'num_consumer_groups': 0, - 'num_ctas': 1, - 'num_stages': 3, - 'num_warps': 8, - 'pre_hook': None, - 'reg_dec_producer': 0, - 'reg_inc_consumer': 0}, - "(13, 512, 512, 27, True, 'torch.float16', 'torch.float16', 'torch.float16', 'torch.uint32', 'torch.int64', 'torch.float16')": {'kwargs': {'B1': 128, - 'B2': 256, - 'BK': 64}, - 'maxnreg': None, - 'num_buffers_warp_spec': 0, - 'num_consumer_groups': 0, - 'num_ctas': 1, - 'num_stages': 3, - 'num_warps': 8, - 'pre_hook': None, - 'reg_dec_producer': 0, - 'reg_inc_consumer': 0}, - "(15, 256, 1024, 27, True, 'torch.float16', 'torch.float16', 'torch.float16', 'torch.uint32', 'torch.int64', 'torch.float16')": {'kwargs': {'B1': 128, - 'B2': 256, - 'BK': 64}, - 'maxnreg': None, - 'num_buffers_warp_spec': 0, - 'num_consumer_groups': 0, - 'num_ctas': 1, - 'num_stages': 3, - 'num_warps': 8, - 'pre_hook': None, - 'reg_dec_producer': 0, - 'reg_inc_consumer': 0}, - "(15, 256, 256, 27, True, 'torch.float16', 'torch.float16', 'torch.float16', 'torch.uint32', 'torch.int64', 'torch.float16')": {'kwargs': {'B1': 128, - 'B2': 256, - 'BK': 64}, - 'maxnreg': None, - 'num_buffers_warp_spec': 0, - 'num_consumer_groups': 0, - 'num_ctas': 1, - 'num_stages': 3, - 'num_warps': 8, - 'pre_hook': None, - 'reg_dec_producer': 0, - 'reg_inc_consumer': 0}, - "(15, 256, 64, 27, True, 'torch.float16', 'torch.float16', 'torch.float16', 'torch.uint32', 'torch.int64', 'torch.float16')": {'kwargs': {'B1': 128, - 'B2': 64, - 'BK': 32}, - 'maxnreg': None, - 'num_buffers_warp_spec': 0, - 'num_consumer_groups': 0, - 'num_ctas': 1, - 'num_stages': 4, - 'num_warps': 4, - 'pre_hook': None, - 'reg_dec_producer': 0, - 'reg_inc_consumer': 0}, - "(15, 512, 128, 27, True, 'torch.float16', 'torch.float16', 'torch.float16', 'torch.uint32', 'torch.int64', 'torch.float16')": {'kwargs': {'B1': 64, - 'B2': 128, - 'BK': 32}, - 'maxnreg': None, - 'num_buffers_warp_spec': 0, - 'num_consumer_groups': 0, - 'num_ctas': 1, - 'num_stages': 4, - 'num_warps': 4, - 'pre_hook': None, - 'reg_dec_producer': 0, - 'reg_inc_consumer': 0}, - "(15, 512, 2048, 27, True, 'torch.float16', 'torch.float16', 'torch.float16', 'torch.uint32', 'torch.int64', 'torch.float16')": {'kwargs': {'B1': 128, - 'B2': 256, - 'BK': 64}, - 'maxnreg': None, - 'num_buffers_warp_spec': 0, - 'num_consumer_groups': 0, - 'num_ctas': 1, - 'num_stages': 3, - 'num_warps': 8, - 'pre_hook': None, - 'reg_dec_producer': 0, - 'reg_inc_consumer': 0}, - "(15, 512, 512, 27, True, 'torch.float16', 'torch.float16', 'torch.float16', 'torch.uint32', 'torch.int64', 'torch.float16')": {'kwargs': {'B1': 128, - 'B2': 256, - 'BK': 64}, - 'maxnreg': None, - 'num_buffers_warp_spec': 0, - 'num_consumer_groups': 0, - 'num_ctas': 1, - 'num_stages': 3, - 'num_warps': 8, - 'pre_hook': None, - 'reg_dec_producer': 0, - 'reg_inc_consumer': 0}, - "(17, 128, 128, 27, True, 'torch.float16', 'torch.float16', 'torch.float16', 'torch.uint32', 'torch.int64', 'torch.float16')": {'kwargs': {'B1': 64, - 'B2': 128, - 'BK': 32}, - 'maxnreg': None, - 'num_buffers_warp_spec': 0, - 'num_consumer_groups': 0, - 'num_ctas': 1, - 'num_stages': 4, - 'num_warps': 4, - 'pre_hook': None, - 'reg_dec_producer': 0, - 'reg_inc_consumer': 0}, - "(17, 128, 32, 27, True, 'torch.float16', 'torch.float16', 'torch.float16', 'torch.uint32', 'torch.int64', 'torch.float16')": {'kwargs': {'B1': 64, - 'B2': 32, - 'BK': 32}, - 'maxnreg': None, - 'num_buffers_warp_spec': 0, - 'num_consumer_groups': 0, - 'num_ctas': 1, - 'num_stages': 5, - 'num_warps': 2, - 'pre_hook': None, - 'reg_dec_producer': 0, - 'reg_inc_consumer': 0}, - "(17, 128, 512, 27, True, 'torch.float16', 'torch.float16', 'torch.float16', 'torch.uint32', 'torch.int64', 'torch.float16')": {'kwargs': {'B1': 128, - 'B2': 256, - 'BK': 64}, - 'maxnreg': None, - 'num_buffers_warp_spec': 0, - 'num_consumer_groups': 0, - 'num_ctas': 1, - 'num_stages': 3, - 'num_warps': 8, - 'pre_hook': None, - 'reg_dec_producer': 0, - 'reg_inc_consumer': 0}, - "(17, 256, 1024, 27, True, 'torch.float16', 'torch.float16', 'torch.float16', 'torch.uint32', 'torch.int64', 'torch.float16')": {'kwargs': {'B1': 128, - 'B2': 256, - 'BK': 64}, - 'maxnreg': None, - 'num_buffers_warp_spec': 0, - 'num_consumer_groups': 0, - 'num_ctas': 1, - 'num_stages': 3, - 'num_warps': 8, - 'pre_hook': None, - 'reg_dec_producer': 0, - 'reg_inc_consumer': 0}, - "(17, 256, 256, 27, True, 'torch.float16', 'torch.float16', 'torch.float16', 'torch.uint32', 'torch.int64', 'torch.float16')": {'kwargs': {'B1': 128, - 'B2': 256, - 'BK': 64}, - 'maxnreg': None, - 'num_buffers_warp_spec': 0, - 'num_consumer_groups': 0, - 'num_ctas': 1, - 'num_stages': 3, - 'num_warps': 8, - 'pre_hook': None, - 'reg_dec_producer': 0, - 'reg_inc_consumer': 0}, - "(17, 256, 64, 27, True, 'torch.float16', 'torch.float16', 'torch.float16', 'torch.uint32', 'torch.int64', 'torch.float16')": {'kwargs': {'B1': 128, - 'B2': 64, - 'BK': 32}, - 'maxnreg': None, - 'num_buffers_warp_spec': 0, - 'num_consumer_groups': 0, - 'num_ctas': 1, - 'num_stages': 4, - 'num_warps': 4, - 'pre_hook': None, - 'reg_dec_producer': 0, - 'reg_inc_consumer': 0}, - "(18, 64, 16, 27, True, 'torch.float16', 'torch.float16', 'torch.float16', 'torch.uint32', 'torch.int64', 'torch.float16')": {'kwargs': {'B1': 64, - 'B2': 32, - 'BK': 32}, - 'maxnreg': None, - 'num_buffers_warp_spec': 0, - 'num_consumer_groups': 0, - 'num_ctas': 1, - 'num_stages': 5, - 'num_warps': 2, - 'pre_hook': None, - 'reg_dec_producer': 0, - 'reg_inc_consumer': 0}, - "(19, 128, 128, 27, True, 'torch.float16', 'torch.float16', 'torch.float16', 'torch.uint32', 'torch.int64', 'torch.float16')": {'kwargs': {'B1': 64, - 'B2': 128, - 'BK': 32}, - 'maxnreg': None, - 'num_buffers_warp_spec': 0, - 'num_consumer_groups': 0, - 'num_ctas': 1, - 'num_stages': 4, - 'num_warps': 4, - 'pre_hook': None, - 'reg_dec_producer': 0, - 'reg_inc_consumer': 0}, - "(19, 128, 32, 27, True, 'torch.float16', 'torch.float16', 'torch.float16', 'torch.uint32', 'torch.int64', 'torch.float16')": {'kwargs': {'B1': 64, - 'B2': 32, - 'BK': 32}, - 'maxnreg': None, - 'num_buffers_warp_spec': 0, - 'num_consumer_groups': 0, - 'num_ctas': 1, - 'num_stages': 5, - 'num_warps': 2, - 'pre_hook': None, - 'reg_dec_producer': 0, - 'reg_inc_consumer': 0}, - "(19, 128, 512, 27, True, 'torch.float16', 'torch.float16', 'torch.float16', 'torch.uint32', 'torch.int64', 'torch.float16')": {'kwargs': {'B1': 128, - 'B2': 256, - 'BK': 64}, - 'maxnreg': None, - 'num_buffers_warp_spec': 0, - 'num_consumer_groups': 0, - 'num_ctas': 1, - 'num_stages': 3, - 'num_warps': 8, - 'pre_hook': None, - 'reg_dec_producer': 0, - 'reg_inc_consumer': 0}, - "(19, 64, 64, 27, True, 'torch.float16', 'torch.float16', 'torch.float16', 'torch.uint32', 'torch.int64', 'torch.float16')": {'kwargs': {'B1': 128, - 'B2': 64, - 'BK': 32}, - 'maxnreg': None, - 'num_buffers_warp_spec': 0, - 'num_consumer_groups': 0, - 'num_ctas': 1, - 'num_stages': 4, - 'num_warps': 4, - 'pre_hook': None, - 'reg_dec_producer': 0, - 'reg_inc_consumer': 0}, - "(21, 64, 16, 27, True, 'torch.float16', 'torch.float16', 'torch.float16', 'torch.uint32', 'torch.int64', 'torch.float16')": {'kwargs': {'B1': 64, - 'B2': 32, - 'BK': 32}, - 'maxnreg': None, - 'num_buffers_warp_spec': 0, - 'num_consumer_groups': 0, - 'num_ctas': 1, - 'num_stages': 5, - 'num_warps': 2, - 'pre_hook': None, - 'reg_dec_producer': 0, - 'reg_inc_consumer': 0}, - "(21, 64, 64, 27, True, 'torch.float16', 'torch.float16', 'torch.float16', 'torch.uint32', 'torch.int64', 'torch.float16')": {'kwargs': {'B1': 32, - 'B2': 64, - 'BK': 32}, - 'maxnreg': None, - 'num_buffers_warp_spec': 0, - 'num_consumer_groups': 0, - 'num_ctas': 1, - 'num_stages': 5, - 'num_warps': 2, - 'pre_hook': None, - 'reg_dec_producer': 0, - 'reg_inc_consumer': 0}}, - 'flex_gemm.kernels.triton.spconv.sparse_submanifold_conv_fwd_masked_implicit_gemm_splitk.sparse_submanifold_conv_fwd_masked_implicit_gemm_kernel': {"(10, 1024, 1024, 27, True, 'torch.float16', 'torch.float16', 'torch.float16', 'torch.uint32', 'torch.int64', 'torch.float16')": {'kwargs': {'B1': 128, - 'B2': 256, - 'BK': 64}, - 'maxnreg': None, - 'num_buffers_warp_spec': 0, - 'num_consumer_groups': 0, - 'num_ctas': 1, - 'num_stages': 3, - 'num_warps': 8, - 'pre_hook': None, - 'reg_dec_producer': 0, - 'reg_inc_consumer': 0}, - "(10, 1024, 4096, 27, True, 'torch.float16', 'torch.float16', 'torch.float16', 'torch.uint32', 'torch.int64', 'torch.float16')": {'kwargs': {'B1': 128, - 'B2': 256, - 'BK': 64}, - 'maxnreg': None, - 'num_buffers_warp_spec': 0, - 'num_consumer_groups': 0, - 'num_ctas': 1, - 'num_stages': 3, - 'num_warps': 8, - 'pre_hook': None, - 'reg_dec_producer': 0, - 'reg_inc_consumer': 0}, - "(12, 1024, 1024, 27, True, 'torch.float16', 'torch.float16', 'torch.float16', 'torch.uint32', 'torch.int64', 'torch.float16')": {'kwargs': {'B1': 128, - 'B2': 256, - 'BK': 64}, - 'maxnreg': None, - 'num_buffers_warp_spec': 0, - 'num_consumer_groups': 0, - 'num_ctas': 1, - 'num_stages': 3, - 'num_warps': 8, - 'pre_hook': None, - 'reg_dec_producer': 0, - 'reg_inc_consumer': 0}, - "(13, 1024, 1024, 27, True, 'torch.float16', 'torch.float16', 'torch.float16', 'torch.uint32', 'torch.int64', 'torch.float16')": {'kwargs': {'B1': 128, - 'B2': 256, - 'BK': 64}, - 'maxnreg': None, - 'num_buffers_warp_spec': 0, - 'num_consumer_groups': 0, - 'num_ctas': 1, - 'num_stages': 3, - 'num_warps': 8, - 'pre_hook': None, - 'reg_dec_producer': 0, - 'reg_inc_consumer': 0}, - "(13, 1024, 4096, 27, True, 'torch.float16', 'torch.float16', 'torch.float16', 'torch.uint32', 'torch.int64', 'torch.float16')": {'kwargs': {'B1': 128, - 'B2': 256, - 'BK': 64}, - 'maxnreg': None, - 'num_buffers_warp_spec': 0, - 'num_consumer_groups': 0, - 'num_ctas': 1, - 'num_stages': 3, - 'num_warps': 8, - 'pre_hook': None, - 'reg_dec_producer': 0, - 'reg_inc_consumer': 0}, - "(13, 512, 128, 27, True, 'torch.float16', 'torch.float16', 'torch.float16', 'torch.uint32', 'torch.int64', 'torch.float16')": {'kwargs': {'B1': 64, - 'B2': 128, - 'BK': 32}, - 'maxnreg': None, - 'num_buffers_warp_spec': 0, - 'num_consumer_groups': 0, - 'num_ctas': 1, - 'num_stages': 4, - 'num_warps': 4, - 'pre_hook': None, - 'reg_dec_producer': 0, - 'reg_inc_consumer': 0}, - "(13, 512, 2048, 27, True, 'torch.float16', 'torch.float16', 'torch.float16', 'torch.uint32', 'torch.int64', 'torch.float16')": {'kwargs': {'B1': 128, - 'B2': 256, - 'BK': 64}, - 'maxnreg': None, - 'num_buffers_warp_spec': 0, - 'num_consumer_groups': 0, - 'num_ctas': 1, - 'num_stages': 3, - 'num_warps': 8, - 'pre_hook': None, - 'reg_dec_producer': 0, - 'reg_inc_consumer': 0}, - "(13, 512, 512, 27, True, 'torch.float16', 'torch.float16', 'torch.float16', 'torch.uint32', 'torch.int64', 'torch.float16')": {'kwargs': {'B1': 128, - 'B2': 256, - 'BK': 64}, - 'maxnreg': None, - 'num_buffers_warp_spec': 0, - 'num_consumer_groups': 0, - 'num_ctas': 1, - 'num_stages': 3, - 'num_warps': 8, - 'pre_hook': None, - 'reg_dec_producer': 0, - 'reg_inc_consumer': 0}, - "(15, 256, 1024, 27, True, 'torch.float16', 'torch.float16', 'torch.float16', 'torch.uint32', 'torch.int64', 'torch.float16')": {'kwargs': {'B1': 128, - 'B2': 256, - 'BK': 64}, - 'maxnreg': None, - 'num_buffers_warp_spec': 0, - 'num_consumer_groups': 0, - 'num_ctas': 1, - 'num_stages': 3, - 'num_warps': 8, - 'pre_hook': None, - 'reg_dec_producer': 0, - 'reg_inc_consumer': 0}, - "(15, 256, 256, 27, True, 'torch.float16', 'torch.float16', 'torch.float16', 'torch.uint32', 'torch.int64', 'torch.float16')": {'kwargs': {'B1': 128, - 'B2': 256, - 'BK': 64}, - 'maxnreg': None, - 'num_buffers_warp_spec': 0, - 'num_consumer_groups': 0, - 'num_ctas': 1, - 'num_stages': 3, - 'num_warps': 8, - 'pre_hook': None, - 'reg_dec_producer': 0, - 'reg_inc_consumer': 0}, - "(15, 256, 64, 27, True, 'torch.float16', 'torch.float16', 'torch.float16', 'torch.uint32', 'torch.int64', 'torch.float16')": {'kwargs': {'B1': 128, - 'B2': 64, - 'BK': 32}, - 'maxnreg': None, - 'num_buffers_warp_spec': 0, - 'num_consumer_groups': 0, - 'num_ctas': 1, - 'num_stages': 4, - 'num_warps': 4, - 'pre_hook': None, - 'reg_dec_producer': 0, - 'reg_inc_consumer': 0}, - "(15, 512, 128, 27, True, 'torch.float16', 'torch.float16', 'torch.float16', 'torch.uint32', 'torch.int64', 'torch.float16')": {'kwargs': {'B1': 64, - 'B2': 128, - 'BK': 32}, - 'maxnreg': None, - 'num_buffers_warp_spec': 0, - 'num_consumer_groups': 0, - 'num_ctas': 1, - 'num_stages': 4, - 'num_warps': 4, - 'pre_hook': None, - 'reg_dec_producer': 0, - 'reg_inc_consumer': 0}, - "(15, 512, 2048, 27, True, 'torch.float16', 'torch.float16', 'torch.float16', 'torch.uint32', 'torch.int64', 'torch.float16')": {'kwargs': {'B1': 128, - 'B2': 256, - 'BK': 64}, - 'maxnreg': None, - 'num_buffers_warp_spec': 0, - 'num_consumer_groups': 0, - 'num_ctas': 1, - 'num_stages': 3, - 'num_warps': 8, - 'pre_hook': None, - 'reg_dec_producer': 0, - 'reg_inc_consumer': 0}, - "(15, 512, 512, 27, True, 'torch.float16', 'torch.float16', 'torch.float16', 'torch.uint32', 'torch.int64', 'torch.float16')": {'kwargs': {'B1': 128, - 'B2': 256, - 'BK': 64}, - 'maxnreg': None, - 'num_buffers_warp_spec': 0, - 'num_consumer_groups': 0, - 'num_ctas': 1, - 'num_stages': 3, - 'num_warps': 8, - 'pre_hook': None, - 'reg_dec_producer': 0, - 'reg_inc_consumer': 0}, - "(17, 128, 128, 27, True, 'torch.float16', 'torch.float16', 'torch.float16', 'torch.uint32', 'torch.int64', 'torch.float16')": {'kwargs': {'B1': 64, - 'B2': 128, - 'BK': 32}, - 'maxnreg': None, - 'num_buffers_warp_spec': 0, - 'num_consumer_groups': 0, - 'num_ctas': 1, - 'num_stages': 4, - 'num_warps': 4, - 'pre_hook': None, - 'reg_dec_producer': 0, - 'reg_inc_consumer': 0}, - "(17, 128, 32, 27, True, 'torch.float16', 'torch.float16', 'torch.float16', 'torch.uint32', 'torch.int64', 'torch.float16')": {'kwargs': {'B1': 64, - 'B2': 32, - 'BK': 32}, - 'maxnreg': None, - 'num_buffers_warp_spec': 0, - 'num_consumer_groups': 0, - 'num_ctas': 1, - 'num_stages': 5, - 'num_warps': 2, - 'pre_hook': None, - 'reg_dec_producer': 0, - 'reg_inc_consumer': 0}, - "(17, 128, 512, 27, True, 'torch.float16', 'torch.float16', 'torch.float16', 'torch.uint32', 'torch.int64', 'torch.float16')": {'kwargs': {'B1': 128, - 'B2': 256, - 'BK': 64}, - 'maxnreg': None, - 'num_buffers_warp_spec': 0, - 'num_consumer_groups': 0, - 'num_ctas': 1, - 'num_stages': 3, - 'num_warps': 8, - 'pre_hook': None, - 'reg_dec_producer': 0, - 'reg_inc_consumer': 0}, - "(17, 256, 1024, 27, True, 'torch.float16', 'torch.float16', 'torch.float16', 'torch.uint32', 'torch.int64', 'torch.float16')": {'kwargs': {'B1': 128, - 'B2': 256, - 'BK': 64}, - 'maxnreg': None, - 'num_buffers_warp_spec': 0, - 'num_consumer_groups': 0, - 'num_ctas': 1, - 'num_stages': 3, - 'num_warps': 8, - 'pre_hook': None, - 'reg_dec_producer': 0, - 'reg_inc_consumer': 0}, - "(17, 256, 256, 27, True, 'torch.float16', 'torch.float16', 'torch.float16', 'torch.uint32', 'torch.int64', 'torch.float16')": {'kwargs': {'B1': 128, - 'B2': 256, - 'BK': 64}, - 'maxnreg': None, - 'num_buffers_warp_spec': 0, - 'num_consumer_groups': 0, - 'num_ctas': 1, - 'num_stages': 3, - 'num_warps': 8, - 'pre_hook': None, - 'reg_dec_producer': 0, - 'reg_inc_consumer': 0}, - "(17, 256, 64, 27, True, 'torch.float16', 'torch.float16', 'torch.float16', 'torch.uint32', 'torch.int64', 'torch.float16')": {'kwargs': {'B1': 128, - 'B2': 64, - 'BK': 32}, - 'maxnreg': None, - 'num_buffers_warp_spec': 0, - 'num_consumer_groups': 0, - 'num_ctas': 1, - 'num_stages': 4, - 'num_warps': 4, - 'pre_hook': None, - 'reg_dec_producer': 0, - 'reg_inc_consumer': 0}, - "(18, 64, 16, 27, True, 'torch.float16', 'torch.float16', 'torch.float16', 'torch.uint32', 'torch.int64', 'torch.float16')": {'kwargs': {'B1': 64, - 'B2': 32, - 'BK': 32}, - 'maxnreg': None, - 'num_buffers_warp_spec': 0, - 'num_consumer_groups': 0, - 'num_ctas': 1, - 'num_stages': 5, - 'num_warps': 2, - 'pre_hook': None, - 'reg_dec_producer': 0, - 'reg_inc_consumer': 0}, - "(19, 128, 128, 27, True, 'torch.float16', 'torch.float16', 'torch.float16', 'torch.uint32', 'torch.int64', 'torch.float16')": {'kwargs': {'B1': 64, - 'B2': 128, - 'BK': 32}, - 'maxnreg': None, - 'num_buffers_warp_spec': 0, - 'num_consumer_groups': 0, - 'num_ctas': 1, - 'num_stages': 4, - 'num_warps': 4, - 'pre_hook': None, - 'reg_dec_producer': 0, - 'reg_inc_consumer': 0}, - "(19, 128, 32, 27, True, 'torch.float16', 'torch.float16', 'torch.float16', 'torch.uint32', 'torch.int64', 'torch.float16')": {'kwargs': {'B1': 64, - 'B2': 32, - 'BK': 32}, - 'maxnreg': None, - 'num_buffers_warp_spec': 0, - 'num_consumer_groups': 0, - 'num_ctas': 1, - 'num_stages': 5, - 'num_warps': 2, - 'pre_hook': None, - 'reg_dec_producer': 0, - 'reg_inc_consumer': 0}, - "(19, 128, 512, 27, True, 'torch.float16', 'torch.float16', 'torch.float16', 'torch.uint32', 'torch.int64', 'torch.float16')": {'kwargs': {'B1': 128, - 'B2': 256, - 'BK': 64}, - 'maxnreg': None, - 'num_buffers_warp_spec': 0, - 'num_consumer_groups': 0, - 'num_ctas': 1, - 'num_stages': 3, - 'num_warps': 8, - 'pre_hook': None, - 'reg_dec_producer': 0, - 'reg_inc_consumer': 0}, - "(19, 64, 64, 27, True, 'torch.float16', 'torch.float16', 'torch.float16', 'torch.uint32', 'torch.int64', 'torch.float16')": {'kwargs': {'B1': 128, - 'B2': 64, - 'BK': 32}, - 'maxnreg': None, - 'num_buffers_warp_spec': 0, - 'num_consumer_groups': 0, - 'num_ctas': 1, - 'num_stages': 4, - 'num_warps': 4, - 'pre_hook': None, - 'reg_dec_producer': 0, - 'reg_inc_consumer': 0}, - "(21, 64, 16, 27, True, 'torch.float16', 'torch.float16', 'torch.float16', 'torch.uint32', 'torch.int64', 'torch.float16')": {'kwargs': {'B1': 64, - 'B2': 32, - 'BK': 32}, - 'maxnreg': None, - 'num_buffers_warp_spec': 0, - 'num_consumer_groups': 0, - 'num_ctas': 1, - 'num_stages': 5, - 'num_warps': 2, - 'pre_hook': None, - 'reg_dec_producer': 0, - 'reg_inc_consumer': 0}, - "(21, 64, 64, 27, True, 'torch.float16', 'torch.float16', 'torch.float16', 'torch.uint32', 'torch.int64', 'torch.float16')": {'kwargs': {'B1': 32, - 'B2': 64, - 'BK': 32}, - 'maxnreg': None, - 'num_buffers_warp_spec': 0, - 'num_consumer_groups': 0, - 'num_ctas': 1, - 'num_stages': 5, - 'num_warps': 2, - 'pre_hook': None, - 'reg_dec_producer': 0, - 'reg_inc_consumer': 0}}, - 'flex_gemm.kernels.triton.spconv.sparse_submanifold_conv_fwd_masked_implicit_gemm_splitk.sparse_submanifold_conv_fwd_masked_implicit_gemm_splitk': {'(2^10, 1024, 1024, 27)': {'SPLITK': 8}, - '(2^10, 1024, 4096, 27)': {'SPLITK': 4}, - '(2^12, 1024, 1024, 27)': {'SPLITK': 8}, - '(2^13, 1024, 1024, 27)': {'SPLITK': 4}, - '(2^13, 1024, 4096, 27)': {'SPLITK': 1}, - '(2^13, 512, 128, 27)': {'SPLITK': 1}, - '(2^13, 512, 2048, 27)': {'SPLITK': 1}, - '(2^13, 512, 512, 27)': {'SPLITK': 4}, - '(2^15, 256, 1024, 27)': {'SPLITK': 1}, - '(2^15, 256, 256, 27)': {'SPLITK': 1}, - '(2^15, 256, 64, 27)': {'SPLITK': 1}, - '(2^15, 512, 128, 27)': {'SPLITK': 1}, - '(2^15, 512, 2048, 27)': {'SPLITK': 1}, - '(2^15, 512, 512, 27)': {'SPLITK': 1}, - '(2^17, 128, 128, 27)': {'SPLITK': 1}, - '(2^17, 128, 32, 27)': {'SPLITK': 1}, - '(2^17, 128, 512, 27)': {'SPLITK': 1}, - '(2^17, 256, 1024, 27)': {'SPLITK': 1}, - '(2^17, 256, 256, 27)': {'SPLITK': 1}, - '(2^17, 256, 64, 27)': {'SPLITK': 1}, - '(2^18, 64, 16, 27)': {'SPLITK': 1}, - '(2^19, 128, 128, 27)': {'SPLITK': 1}, - '(2^19, 128, 32, 27)': {'SPLITK': 1}, - '(2^19, 128, 512, 27)': {'SPLITK': 1}, - '(2^19, 64, 64, 27)': {'SPLITK': 1}, - '(2^21, 64, 16, 27)': {'SPLITK': 1}, - '(2^21, 64, 64, 27)': {'SPLITK': 1}}, - 'flex_gemm.kernels.triton.spconv.sparse_submanifold_conv_fwd_masked_implicit_gemm_splitk.sparse_submanifold_conv_fwd_masked_implicit_gemm_splitk_kernel': {"(10, 1024, 1024, 27, 16, True, 'torch.float16', 'torch.float16', 'torch.float16', 'torch.uint32', 'torch.int64', 'torch.float32')": {'kwargs': {'B1': 128, - 'B2': 256, - 'BK': 64}, - 'maxnreg': None, - 'num_buffers_warp_spec': 0, - 'num_consumer_groups': 0, - 'num_ctas': 1, - 'num_stages': 3, - 'num_warps': 8, - 'pre_hook': None, - 'reg_dec_producer': 0, - 'reg_inc_consumer': 0}, - "(10, 1024, 1024, 27, 2, True, 'torch.float16', 'torch.float16', 'torch.float16', 'torch.uint32', 'torch.int64', 'torch.float32')": {'kwargs': {'B1': 128, - 'B2': 256, - 'BK': 64}, - 'maxnreg': None, - 'num_buffers_warp_spec': 0, - 'num_consumer_groups': 0, - 'num_ctas': 1, - 'num_stages': 3, - 'num_warps': 8, - 'pre_hook': None, - 'reg_dec_producer': 0, - 'reg_inc_consumer': 0}, - "(10, 1024, 1024, 27, 32, True, 'torch.float16', 'torch.float16', 'torch.float16', 'torch.uint32', 'torch.int64', 'torch.float32')": {'kwargs': {'B1': 128, - 'B2': 256, - 'BK': 64}, - 'maxnreg': None, - 'num_buffers_warp_spec': 0, - 'num_consumer_groups': 0, - 'num_ctas': 1, - 'num_stages': 3, - 'num_warps': 8, - 'pre_hook': None, - 'reg_dec_producer': 0, - 'reg_inc_consumer': 0}, - "(10, 1024, 1024, 27, 4, True, 'torch.float16', 'torch.float16', 'torch.float16', 'torch.uint32', 'torch.int64', 'torch.float32')": {'kwargs': {'B1': 128, - 'B2': 256, - 'BK': 64}, - 'maxnreg': None, - 'num_buffers_warp_spec': 0, - 'num_consumer_groups': 0, - 'num_ctas': 1, - 'num_stages': 3, - 'num_warps': 8, - 'pre_hook': None, - 'reg_dec_producer': 0, - 'reg_inc_consumer': 0}, - "(10, 1024, 1024, 27, 8, True, 'torch.float16', 'torch.float16', 'torch.float16', 'torch.uint32', 'torch.int64', 'torch.float32')": {'kwargs': {'B1': 128, - 'B2': 256, - 'BK': 64}, - 'maxnreg': None, - 'num_buffers_warp_spec': 0, - 'num_consumer_groups': 0, - 'num_ctas': 1, - 'num_stages': 3, - 'num_warps': 8, - 'pre_hook': None, - 'reg_dec_producer': 0, - 'reg_inc_consumer': 0}, - "(10, 1024, 4096, 27, 2, True, 'torch.float16', 'torch.float16', 'torch.float16', 'torch.uint32', 'torch.int64', 'torch.float32')": {'kwargs': {'B1': 128, - 'B2': 256, - 'BK': 64}, - 'maxnreg': None, - 'num_buffers_warp_spec': 0, - 'num_consumer_groups': 0, - 'num_ctas': 1, - 'num_stages': 3, - 'num_warps': 8, - 'pre_hook': None, - 'reg_dec_producer': 0, - 'reg_inc_consumer': 0}, - "(10, 1024, 4096, 27, 4, True, 'torch.float16', 'torch.float16', 'torch.float16', 'torch.uint32', 'torch.int64', 'torch.float32')": {'kwargs': {'B1': 128, - 'B2': 256, - 'BK': 64}, - 'maxnreg': None, - 'num_buffers_warp_spec': 0, - 'num_consumer_groups': 0, - 'num_ctas': 1, - 'num_stages': 3, - 'num_warps': 8, - 'pre_hook': None, - 'reg_dec_producer': 0, - 'reg_inc_consumer': 0}, - "(10, 1024, 4096, 27, 8, True, 'torch.float16', 'torch.float16', 'torch.float16', 'torch.uint32', 'torch.int64', 'torch.float32')": {'kwargs': {'B1': 128, - 'B2': 256, - 'BK': 64}, - 'maxnreg': None, - 'num_buffers_warp_spec': 0, - 'num_consumer_groups': 0, - 'num_ctas': 1, - 'num_stages': 3, - 'num_warps': 8, - 'pre_hook': None, - 'reg_dec_producer': 0, - 'reg_inc_consumer': 0}, - "(12, 1024, 1024, 27, 2, True, 'torch.float16', 'torch.float16', 'torch.float16', 'torch.uint32', 'torch.int64', 'torch.float32')": {'kwargs': {'B1': 128, - 'B2': 256, - 'BK': 64}, - 'maxnreg': None, - 'num_buffers_warp_spec': 0, - 'num_consumer_groups': 0, - 'num_ctas': 1, - 'num_stages': 3, - 'num_warps': 8, - 'pre_hook': None, - 'reg_dec_producer': 0, - 'reg_inc_consumer': 0}, - "(12, 1024, 1024, 27, 4, True, 'torch.float16', 'torch.float16', 'torch.float16', 'torch.uint32', 'torch.int64', 'torch.float32')": {'kwargs': {'B1': 128, - 'B2': 256, - 'BK': 64}, - 'maxnreg': None, - 'num_buffers_warp_spec': 0, - 'num_consumer_groups': 0, - 'num_ctas': 1, - 'num_stages': 3, - 'num_warps': 8, - 'pre_hook': None, - 'reg_dec_producer': 0, - 'reg_inc_consumer': 0}, - "(12, 1024, 1024, 27, 8, True, 'torch.float16', 'torch.float16', 'torch.float16', 'torch.uint32', 'torch.int64', 'torch.float32')": {'kwargs': {'B1': 128, - 'B2': 256, - 'BK': 64}, - 'maxnreg': None, - 'num_buffers_warp_spec': 0, - 'num_consumer_groups': 0, - 'num_ctas': 1, - 'num_stages': 3, - 'num_warps': 8, - 'pre_hook': None, - 'reg_dec_producer': 0, - 'reg_inc_consumer': 0}, - "(13, 1024, 1024, 27, 2, True, 'torch.float16', 'torch.float16', 'torch.float16', 'torch.uint32', 'torch.int64', 'torch.float32')": {'kwargs': {'B1': 128, - 'B2': 256, - 'BK': 64}, - 'maxnreg': None, - 'num_buffers_warp_spec': 0, - 'num_consumer_groups': 0, - 'num_ctas': 1, - 'num_stages': 3, - 'num_warps': 8, - 'pre_hook': None, - 'reg_dec_producer': 0, - 'reg_inc_consumer': 0}, - "(13, 1024, 1024, 27, 4, True, 'torch.float16', 'torch.float16', 'torch.float16', 'torch.uint32', 'torch.int64', 'torch.float32')": {'kwargs': {'B1': 128, - 'B2': 256, - 'BK': 64}, - 'maxnreg': None, - 'num_buffers_warp_spec': 0, - 'num_consumer_groups': 0, - 'num_ctas': 1, - 'num_stages': 3, - 'num_warps': 8, - 'pre_hook': None, - 'reg_dec_producer': 0, - 'reg_inc_consumer': 0}, - "(13, 512, 128, 27, 16, True, 'torch.float16', 'torch.float16', 'torch.float16', 'torch.uint32', 'torch.int64', 'torch.float32')": {'kwargs': {'B1': 128, - 'B2': 128, - 'BK': 32}, - 'maxnreg': None, - 'num_buffers_warp_spec': 0, - 'num_consumer_groups': 0, - 'num_ctas': 1, - 'num_stages': 4, - 'num_warps': 4, - 'pre_hook': None, - 'reg_dec_producer': 0, - 'reg_inc_consumer': 0}, - "(13, 512, 128, 27, 2, True, 'torch.float16', 'torch.float16', 'torch.float16', 'torch.uint32', 'torch.int64', 'torch.float32')": {'kwargs': {'B1': 128, - 'B2': 128, - 'BK': 32}, - 'maxnreg': None, - 'num_buffers_warp_spec': 0, - 'num_consumer_groups': 0, - 'num_ctas': 1, - 'num_stages': 4, - 'num_warps': 4, - 'pre_hook': None, - 'reg_dec_producer': 0, - 'reg_inc_consumer': 0}, - "(13, 512, 128, 27, 32, True, 'torch.float16', 'torch.float16', 'torch.float16', 'torch.uint32', 'torch.int64', 'torch.float32')": {'kwargs': {'B1': 128, - 'B2': 64, - 'BK': 32}, - 'maxnreg': None, - 'num_buffers_warp_spec': 0, - 'num_consumer_groups': 0, - 'num_ctas': 1, - 'num_stages': 4, - 'num_warps': 4, - 'pre_hook': None, - 'reg_dec_producer': 0, - 'reg_inc_consumer': 0}, - "(13, 512, 128, 27, 4, True, 'torch.float16', 'torch.float16', 'torch.float16', 'torch.uint32', 'torch.int64', 'torch.float32')": {'kwargs': {'B1': 128, - 'B2': 128, - 'BK': 32}, - 'maxnreg': None, - 'num_buffers_warp_spec': 0, - 'num_consumer_groups': 0, - 'num_ctas': 1, - 'num_stages': 4, - 'num_warps': 4, - 'pre_hook': None, - 'reg_dec_producer': 0, - 'reg_inc_consumer': 0}, - "(13, 512, 128, 27, 8, True, 'torch.float16', 'torch.float16', 'torch.float16', 'torch.uint32', 'torch.int64', 'torch.float32')": {'kwargs': {'B1': 128, - 'B2': 128, - 'BK': 32}, - 'maxnreg': None, - 'num_buffers_warp_spec': 0, - 'num_consumer_groups': 0, - 'num_ctas': 1, - 'num_stages': 4, - 'num_warps': 4, - 'pre_hook': None, - 'reg_dec_producer': 0, - 'reg_inc_consumer': 0}, - "(13, 512, 2048, 27, 2, True, 'torch.float16', 'torch.float16', 'torch.float16', 'torch.uint32', 'torch.int64', 'torch.float32')": {'kwargs': {'B1': 128, - 'B2': 256, - 'BK': 64}, - 'maxnreg': None, - 'num_buffers_warp_spec': 0, - 'num_consumer_groups': 0, - 'num_ctas': 1, - 'num_stages': 3, - 'num_warps': 8, - 'pre_hook': None, - 'reg_dec_producer': 0, - 'reg_inc_consumer': 0}, - "(13, 512, 512, 27, 2, True, 'torch.float16', 'torch.float16', 'torch.float16', 'torch.uint32', 'torch.int64', 'torch.float32')": {'kwargs': {'B1': 128, - 'B2': 256, - 'BK': 64}, - 'maxnreg': None, - 'num_buffers_warp_spec': 0, - 'num_consumer_groups': 0, - 'num_ctas': 1, - 'num_stages': 3, - 'num_warps': 8, - 'pre_hook': None, - 'reg_dec_producer': 0, - 'reg_inc_consumer': 0}, - "(13, 512, 512, 27, 4, True, 'torch.float16', 'torch.float16', 'torch.float16', 'torch.uint32', 'torch.int64', 'torch.float32')": {'kwargs': {'B1': 128, - 'B2': 256, - 'BK': 64}, - 'maxnreg': None, - 'num_buffers_warp_spec': 0, - 'num_consumer_groups': 0, - 'num_ctas': 1, - 'num_stages': 3, - 'num_warps': 8, - 'pre_hook': None, - 'reg_dec_producer': 0, - 'reg_inc_consumer': 0}, - "(13, 512, 512, 27, 8, True, 'torch.float16', 'torch.float16', 'torch.float16', 'torch.uint32', 'torch.int64', 'torch.float32')": {'kwargs': {'B1': 128, - 'B2': 256, - 'BK': 64}, - 'maxnreg': None, - 'num_buffers_warp_spec': 0, - 'num_consumer_groups': 0, - 'num_ctas': 1, - 'num_stages': 3, - 'num_warps': 8, - 'pre_hook': None, - 'reg_dec_producer': 0, - 'reg_inc_consumer': 0}, - "(15, 256, 256, 27, 2, True, 'torch.float16', 'torch.float16', 'torch.float16', 'torch.uint32', 'torch.int64', 'torch.float32')": {'kwargs': {'B1': 128, - 'B2': 256, - 'BK': 64}, - 'maxnreg': None, - 'num_buffers_warp_spec': 0, - 'num_consumer_groups': 0, - 'num_ctas': 1, - 'num_stages': 3, - 'num_warps': 8, - 'pre_hook': None, - 'reg_dec_producer': 0, - 'reg_inc_consumer': 0}, - "(15, 256, 256, 27, 4, True, 'torch.float16', 'torch.float16', 'torch.float16', 'torch.uint32', 'torch.int64', 'torch.float32')": {'kwargs': {'B1': 128, - 'B2': 256, - 'BK': 64}, - 'maxnreg': None, - 'num_buffers_warp_spec': 0, - 'num_consumer_groups': 0, - 'num_ctas': 1, - 'num_stages': 3, - 'num_warps': 8, - 'pre_hook': None, - 'reg_dec_producer': 0, - 'reg_inc_consumer': 0}, - "(15, 256, 64, 27, 2, True, 'torch.float16', 'torch.float16', 'torch.float16', 'torch.uint32', 'torch.int64', 'torch.float32')": {'kwargs': {'B1': 128, - 'B2': 64, - 'BK': 32}, - 'maxnreg': None, - 'num_buffers_warp_spec': 0, - 'num_consumer_groups': 0, - 'num_ctas': 1, - 'num_stages': 4, - 'num_warps': 4, - 'pre_hook': None, - 'reg_dec_producer': 0, - 'reg_inc_consumer': 0}, - "(15, 256, 64, 27, 4, True, 'torch.float16', 'torch.float16', 'torch.float16', 'torch.uint32', 'torch.int64', 'torch.float32')": {'kwargs': {'B1': 128, - 'B2': 64, - 'BK': 32}, - 'maxnreg': None, - 'num_buffers_warp_spec': 0, - 'num_consumer_groups': 0, - 'num_ctas': 1, - 'num_stages': 4, - 'num_warps': 4, - 'pre_hook': None, - 'reg_dec_producer': 0, - 'reg_inc_consumer': 0}, - "(15, 256, 64, 27, 8, True, 'torch.float16', 'torch.float16', 'torch.float16', 'torch.uint32', 'torch.int64', 'torch.float32')": {'kwargs': {'B1': 128, - 'B2': 64, - 'BK': 32}, - 'maxnreg': None, - 'num_buffers_warp_spec': 0, - 'num_consumer_groups': 0, - 'num_ctas': 1, - 'num_stages': 4, - 'num_warps': 4, - 'pre_hook': None, - 'reg_dec_producer': 0, - 'reg_inc_consumer': 0}, - "(15, 512, 128, 27, 2, True, 'torch.float16', 'torch.float16', 'torch.float16', 'torch.uint32', 'torch.int64', 'torch.float32')": {'kwargs': {'B1': 128, - 'B2': 64, - 'BK': 32}, - 'maxnreg': None, - 'num_buffers_warp_spec': 0, - 'num_consumer_groups': 0, - 'num_ctas': 1, - 'num_stages': 4, - 'num_warps': 4, - 'pre_hook': None, - 'reg_dec_producer': 0, - 'reg_inc_consumer': 0}, - "(15, 512, 128, 27, 4, True, 'torch.float16', 'torch.float16', 'torch.float16', 'torch.uint32', 'torch.int64', 'torch.float32')": {'kwargs': {'B1': 128, - 'B2': 64, - 'BK': 32}, - 'maxnreg': None, - 'num_buffers_warp_spec': 0, - 'num_consumer_groups': 0, - 'num_ctas': 1, - 'num_stages': 4, - 'num_warps': 4, - 'pre_hook': None, - 'reg_dec_producer': 0, - 'reg_inc_consumer': 0}, - "(15, 512, 128, 27, 8, True, 'torch.float16', 'torch.float16', 'torch.float16', 'torch.uint32', 'torch.int64', 'torch.float32')": {'kwargs': {'B1': 128, - 'B2': 64, - 'BK': 32}, - 'maxnreg': None, - 'num_buffers_warp_spec': 0, - 'num_consumer_groups': 0, - 'num_ctas': 1, - 'num_stages': 4, - 'num_warps': 4, - 'pre_hook': None, - 'reg_dec_producer': 0, - 'reg_inc_consumer': 0}, - "(15, 512, 512, 27, 2, True, 'torch.float16', 'torch.float16', 'torch.float16', 'torch.uint32', 'torch.int64', 'torch.float32')": {'kwargs': {'B1': 128, - 'B2': 256, - 'BK': 64}, - 'maxnreg': None, - 'num_buffers_warp_spec': 0, - 'num_consumer_groups': 0, - 'num_ctas': 1, - 'num_stages': 3, - 'num_warps': 8, - 'pre_hook': None, - 'reg_dec_producer': 0, - 'reg_inc_consumer': 0}, - "(17, 128, 128, 27, 2, True, 'torch.float16', 'torch.float16', 'torch.float16', 'torch.uint32', 'torch.int64', 'torch.float32')": {'kwargs': {'B1': 128, - 'B2': 64, - 'BK': 32}, - 'maxnreg': None, - 'num_buffers_warp_spec': 0, - 'num_consumer_groups': 0, - 'num_ctas': 1, - 'num_stages': 4, - 'num_warps': 4, - 'pre_hook': None, - 'reg_dec_producer': 0, - 'reg_inc_consumer': 0}, - "(17, 128, 32, 27, 2, True, 'torch.float16', 'torch.float16', 'torch.float16', 'torch.uint32', 'torch.int64', 'torch.float32')": {'kwargs': {'B1': 64, - 'B2': 32, - 'BK': 32}, - 'maxnreg': None, - 'num_buffers_warp_spec': 0, - 'num_consumer_groups': 0, - 'num_ctas': 1, - 'num_stages': 5, - 'num_warps': 2, - 'pre_hook': None, - 'reg_dec_producer': 0, - 'reg_inc_consumer': 0}, - "(17, 256, 64, 27, 2, True, 'torch.float16', 'torch.float16', 'torch.float16', 'torch.uint32', 'torch.int64', 'torch.float32')": {'kwargs': {'B1': 128, - 'B2': 64, - 'BK': 32}, - 'maxnreg': None, - 'num_buffers_warp_spec': 0, - 'num_consumer_groups': 0, - 'num_ctas': 1, - 'num_stages': 4, - 'num_warps': 4, - 'pre_hook': None, - 'reg_dec_producer': 0, - 'reg_inc_consumer': 0}}}} +"""Measured H800 forward-convolution launch choices for kernel_config. + +Profiled for geometry/completion inference and encoding. End-to-end inference +entries take precedence on conflicts. +Only measured shape buckets are present; all other keys retain autotuning. +This Python data module is included automatically in source/wheel packages. +""" + +CONFIG = {'NVIDIA H800': {'flex_gemm.kernels.triton.spconv.sparse_submanifold_conv_fwd_masked_implicit_gemm.sparse_submanifold_conv_fwd_masked_implicit_gemm_kernel': {"(10, 1024, 1024, 27, True, 'torch.float16', 'torch.float16', 'torch.float16', 'torch.uint32', 'torch.int64', 'torch.float16')": {'kwargs': {'B1': 128, + 'B2': 256, + 'BK': 64}, + 'maxnreg': None, + 'num_buffers_warp_spec': 0, + 'num_consumer_groups': 0, + 'num_ctas': 1, + 'num_stages': 3, + 'num_warps': 8, + 'pre_hook': None, + 'reg_dec_producer': 0, + 'reg_inc_consumer': 0}, + "(10, 1024, 4096, 27, True, 'torch.float16', 'torch.float16', 'torch.float16', 'torch.uint32', 'torch.int64', 'torch.float16')": {'kwargs': {'B1': 128, + 'B2': 256, + 'BK': 64}, + 'maxnreg': None, + 'num_buffers_warp_spec': 0, + 'num_consumer_groups': 0, + 'num_ctas': 1, + 'num_stages': 3, + 'num_warps': 8, + 'pre_hook': None, + 'reg_dec_producer': 0, + 'reg_inc_consumer': 0}, + "(12, 1024, 1024, 27, True, 'torch.float16', 'torch.float16', 'torch.float16', 'torch.uint32', 'torch.int64', 'torch.float16')": {'kwargs': {'B1': 128, + 'B2': 256, + 'BK': 64}, + 'maxnreg': None, + 'num_buffers_warp_spec': 0, + 'num_consumer_groups': 0, + 'num_ctas': 1, + 'num_stages': 3, + 'num_warps': 8, + 'pre_hook': None, + 'reg_dec_producer': 0, + 'reg_inc_consumer': 0}, + "(13, 1024, 1024, 27, True, 'torch.float16', 'torch.float16', 'torch.float16', 'torch.uint32', 'torch.int64', 'torch.float16')": {'kwargs': {'B1': 128, + 'B2': 256, + 'BK': 64}, + 'maxnreg': None, + 'num_buffers_warp_spec': 0, + 'num_consumer_groups': 0, + 'num_ctas': 1, + 'num_stages': 3, + 'num_warps': 8, + 'pre_hook': None, + 'reg_dec_producer': 0, + 'reg_inc_consumer': 0}, + "(13, 1024, 4096, 27, True, 'torch.float16', 'torch.float16', 'torch.float16', 'torch.uint32', 'torch.int64', 'torch.float16')": {'kwargs': {'B1': 128, + 'B2': 256, + 'BK': 64}, + 'maxnreg': None, + 'num_buffers_warp_spec': 0, + 'num_consumer_groups': 0, + 'num_ctas': 1, + 'num_stages': 3, + 'num_warps': 8, + 'pre_hook': None, + 'reg_dec_producer': 0, + 'reg_inc_consumer': 0}, + "(13, 512, 128, 27, True, 'torch.float16', 'torch.float16', 'torch.float16', 'torch.uint32', 'torch.int64', 'torch.float16')": {'kwargs': {'B1': 64, + 'B2': 128, + 'BK': 32}, + 'maxnreg': None, + 'num_buffers_warp_spec': 0, + 'num_consumer_groups': 0, + 'num_ctas': 1, + 'num_stages': 4, + 'num_warps': 4, + 'pre_hook': None, + 'reg_dec_producer': 0, + 'reg_inc_consumer': 0}, + "(13, 512, 2048, 27, True, 'torch.float16', 'torch.float16', 'torch.float16', 'torch.uint32', 'torch.int64', 'torch.float16')": {'kwargs': {'B1': 128, + 'B2': 256, + 'BK': 64}, + 'maxnreg': None, + 'num_buffers_warp_spec': 0, + 'num_consumer_groups': 0, + 'num_ctas': 1, + 'num_stages': 3, + 'num_warps': 8, + 'pre_hook': None, + 'reg_dec_producer': 0, + 'reg_inc_consumer': 0}, + "(13, 512, 512, 27, True, 'torch.float16', 'torch.float16', 'torch.float16', 'torch.uint32', 'torch.int64', 'torch.float16')": {'kwargs': {'B1': 128, + 'B2': 256, + 'BK': 64}, + 'maxnreg': None, + 'num_buffers_warp_spec': 0, + 'num_consumer_groups': 0, + 'num_ctas': 1, + 'num_stages': 3, + 'num_warps': 8, + 'pre_hook': None, + 'reg_dec_producer': 0, + 'reg_inc_consumer': 0}, + "(15, 256, 1024, 27, True, 'torch.float16', 'torch.float16', 'torch.float16', 'torch.uint32', 'torch.int64', 'torch.float16')": {'kwargs': {'B1': 128, + 'B2': 256, + 'BK': 64}, + 'maxnreg': None, + 'num_buffers_warp_spec': 0, + 'num_consumer_groups': 0, + 'num_ctas': 1, + 'num_stages': 3, + 'num_warps': 8, + 'pre_hook': None, + 'reg_dec_producer': 0, + 'reg_inc_consumer': 0}, + "(15, 256, 256, 27, True, 'torch.float16', 'torch.float16', 'torch.float16', 'torch.uint32', 'torch.int64', 'torch.float16')": {'kwargs': {'B1': 128, + 'B2': 256, + 'BK': 64}, + 'maxnreg': None, + 'num_buffers_warp_spec': 0, + 'num_consumer_groups': 0, + 'num_ctas': 1, + 'num_stages': 3, + 'num_warps': 8, + 'pre_hook': None, + 'reg_dec_producer': 0, + 'reg_inc_consumer': 0}, + "(15, 256, 64, 27, True, 'torch.float16', 'torch.float16', 'torch.float16', 'torch.uint32', 'torch.int64', 'torch.float16')": {'kwargs': {'B1': 128, + 'B2': 64, + 'BK': 32}, + 'maxnreg': None, + 'num_buffers_warp_spec': 0, + 'num_consumer_groups': 0, + 'num_ctas': 1, + 'num_stages': 4, + 'num_warps': 4, + 'pre_hook': None, + 'reg_dec_producer': 0, + 'reg_inc_consumer': 0}, + "(15, 512, 128, 27, True, 'torch.float16', 'torch.float16', 'torch.float16', 'torch.uint32', 'torch.int64', 'torch.float16')": {'kwargs': {'B1': 64, + 'B2': 128, + 'BK': 32}, + 'maxnreg': None, + 'num_buffers_warp_spec': 0, + 'num_consumer_groups': 0, + 'num_ctas': 1, + 'num_stages': 4, + 'num_warps': 4, + 'pre_hook': None, + 'reg_dec_producer': 0, + 'reg_inc_consumer': 0}, + "(15, 512, 2048, 27, True, 'torch.float16', 'torch.float16', 'torch.float16', 'torch.uint32', 'torch.int64', 'torch.float16')": {'kwargs': {'B1': 128, + 'B2': 256, + 'BK': 64}, + 'maxnreg': None, + 'num_buffers_warp_spec': 0, + 'num_consumer_groups': 0, + 'num_ctas': 1, + 'num_stages': 3, + 'num_warps': 8, + 'pre_hook': None, + 'reg_dec_producer': 0, + 'reg_inc_consumer': 0}, + "(15, 512, 512, 27, True, 'torch.float16', 'torch.float16', 'torch.float16', 'torch.uint32', 'torch.int64', 'torch.float16')": {'kwargs': {'B1': 128, + 'B2': 256, + 'BK': 64}, + 'maxnreg': None, + 'num_buffers_warp_spec': 0, + 'num_consumer_groups': 0, + 'num_ctas': 1, + 'num_stages': 3, + 'num_warps': 8, + 'pre_hook': None, + 'reg_dec_producer': 0, + 'reg_inc_consumer': 0}, + "(17, 128, 128, 27, True, 'torch.float16', 'torch.float16', 'torch.float16', 'torch.uint32', 'torch.int64', 'torch.float16')": {'kwargs': {'B1': 64, + 'B2': 128, + 'BK': 32}, + 'maxnreg': None, + 'num_buffers_warp_spec': 0, + 'num_consumer_groups': 0, + 'num_ctas': 1, + 'num_stages': 4, + 'num_warps': 4, + 'pre_hook': None, + 'reg_dec_producer': 0, + 'reg_inc_consumer': 0}, + "(17, 128, 32, 27, True, 'torch.float16', 'torch.float16', 'torch.float16', 'torch.uint32', 'torch.int64', 'torch.float16')": {'kwargs': {'B1': 64, + 'B2': 32, + 'BK': 32}, + 'maxnreg': None, + 'num_buffers_warp_spec': 0, + 'num_consumer_groups': 0, + 'num_ctas': 1, + 'num_stages': 5, + 'num_warps': 2, + 'pre_hook': None, + 'reg_dec_producer': 0, + 'reg_inc_consumer': 0}, + "(17, 128, 512, 27, True, 'torch.float16', 'torch.float16', 'torch.float16', 'torch.uint32', 'torch.int64', 'torch.float16')": {'kwargs': {'B1': 128, + 'B2': 256, + 'BK': 64}, + 'maxnreg': None, + 'num_buffers_warp_spec': 0, + 'num_consumer_groups': 0, + 'num_ctas': 1, + 'num_stages': 3, + 'num_warps': 8, + 'pre_hook': None, + 'reg_dec_producer': 0, + 'reg_inc_consumer': 0}, + "(17, 256, 1024, 27, True, 'torch.float16', 'torch.float16', 'torch.float16', 'torch.uint32', 'torch.int64', 'torch.float16')": {'kwargs': {'B1': 128, + 'B2': 256, + 'BK': 64}, + 'maxnreg': None, + 'num_buffers_warp_spec': 0, + 'num_consumer_groups': 0, + 'num_ctas': 1, + 'num_stages': 3, + 'num_warps': 8, + 'pre_hook': None, + 'reg_dec_producer': 0, + 'reg_inc_consumer': 0}, + "(17, 256, 256, 27, True, 'torch.float16', 'torch.float16', 'torch.float16', 'torch.uint32', 'torch.int64', 'torch.float16')": {'kwargs': {'B1': 128, + 'B2': 256, + 'BK': 64}, + 'maxnreg': None, + 'num_buffers_warp_spec': 0, + 'num_consumer_groups': 0, + 'num_ctas': 1, + 'num_stages': 3, + 'num_warps': 8, + 'pre_hook': None, + 'reg_dec_producer': 0, + 'reg_inc_consumer': 0}, + "(17, 256, 64, 27, True, 'torch.float16', 'torch.float16', 'torch.float16', 'torch.uint32', 'torch.int64', 'torch.float16')": {'kwargs': {'B1': 128, + 'B2': 64, + 'BK': 32}, + 'maxnreg': None, + 'num_buffers_warp_spec': 0, + 'num_consumer_groups': 0, + 'num_ctas': 1, + 'num_stages': 4, + 'num_warps': 4, + 'pre_hook': None, + 'reg_dec_producer': 0, + 'reg_inc_consumer': 0}, + "(18, 64, 16, 27, True, 'torch.float16', 'torch.float16', 'torch.float16', 'torch.uint32', 'torch.int64', 'torch.float16')": {'kwargs': {'B1': 64, + 'B2': 32, + 'BK': 32}, + 'maxnreg': None, + 'num_buffers_warp_spec': 0, + 'num_consumer_groups': 0, + 'num_ctas': 1, + 'num_stages': 5, + 'num_warps': 2, + 'pre_hook': None, + 'reg_dec_producer': 0, + 'reg_inc_consumer': 0}, + "(19, 128, 128, 27, True, 'torch.float16', 'torch.float16', 'torch.float16', 'torch.uint32', 'torch.int64', 'torch.float16')": {'kwargs': {'B1': 64, + 'B2': 128, + 'BK': 32}, + 'maxnreg': None, + 'num_buffers_warp_spec': 0, + 'num_consumer_groups': 0, + 'num_ctas': 1, + 'num_stages': 4, + 'num_warps': 4, + 'pre_hook': None, + 'reg_dec_producer': 0, + 'reg_inc_consumer': 0}, + "(19, 128, 32, 27, True, 'torch.float16', 'torch.float16', 'torch.float16', 'torch.uint32', 'torch.int64', 'torch.float16')": {'kwargs': {'B1': 64, + 'B2': 32, + 'BK': 32}, + 'maxnreg': None, + 'num_buffers_warp_spec': 0, + 'num_consumer_groups': 0, + 'num_ctas': 1, + 'num_stages': 5, + 'num_warps': 2, + 'pre_hook': None, + 'reg_dec_producer': 0, + 'reg_inc_consumer': 0}, + "(19, 128, 512, 27, True, 'torch.float16', 'torch.float16', 'torch.float16', 'torch.uint32', 'torch.int64', 'torch.float16')": {'kwargs': {'B1': 128, + 'B2': 256, + 'BK': 64}, + 'maxnreg': None, + 'num_buffers_warp_spec': 0, + 'num_consumer_groups': 0, + 'num_ctas': 1, + 'num_stages': 3, + 'num_warps': 8, + 'pre_hook': None, + 'reg_dec_producer': 0, + 'reg_inc_consumer': 0}, + "(19, 64, 64, 27, True, 'torch.float16', 'torch.float16', 'torch.float16', 'torch.uint32', 'torch.int64', 'torch.float16')": {'kwargs': {'B1': 128, + 'B2': 64, + 'BK': 32}, + 'maxnreg': None, + 'num_buffers_warp_spec': 0, + 'num_consumer_groups': 0, + 'num_ctas': 1, + 'num_stages': 4, + 'num_warps': 4, + 'pre_hook': None, + 'reg_dec_producer': 0, + 'reg_inc_consumer': 0}, + "(21, 64, 16, 27, True, 'torch.float16', 'torch.float16', 'torch.float16', 'torch.uint32', 'torch.int64', 'torch.float16')": {'kwargs': {'B1': 64, + 'B2': 32, + 'BK': 32}, + 'maxnreg': None, + 'num_buffers_warp_spec': 0, + 'num_consumer_groups': 0, + 'num_ctas': 1, + 'num_stages': 5, + 'num_warps': 2, + 'pre_hook': None, + 'reg_dec_producer': 0, + 'reg_inc_consumer': 0}, + "(21, 64, 64, 27, True, 'torch.float16', 'torch.float16', 'torch.float16', 'torch.uint32', 'torch.int64', 'torch.float16')": {'kwargs': {'B1': 32, + 'B2': 64, + 'BK': 32}, + 'maxnreg': None, + 'num_buffers_warp_spec': 0, + 'num_consumer_groups': 0, + 'num_ctas': 1, + 'num_stages': 5, + 'num_warps': 2, + 'pre_hook': None, + 'reg_dec_producer': 0, + 'reg_inc_consumer': 0}}, + 'flex_gemm.kernels.triton.spconv.sparse_submanifold_conv_fwd_masked_implicit_gemm_splitk.sparse_submanifold_conv_fwd_masked_implicit_gemm_kernel': {"(10, 1024, 1024, 27, True, 'torch.float16', 'torch.float16', 'torch.float16', 'torch.uint32', 'torch.int64', 'torch.float16')": {'kwargs': {'B1': 128, + 'B2': 256, + 'BK': 64}, + 'maxnreg': None, + 'num_buffers_warp_spec': 0, + 'num_consumer_groups': 0, + 'num_ctas': 1, + 'num_stages': 3, + 'num_warps': 8, + 'pre_hook': None, + 'reg_dec_producer': 0, + 'reg_inc_consumer': 0}, + "(10, 1024, 4096, 27, True, 'torch.float16', 'torch.float16', 'torch.float16', 'torch.uint32', 'torch.int64', 'torch.float16')": {'kwargs': {'B1': 128, + 'B2': 256, + 'BK': 64}, + 'maxnreg': None, + 'num_buffers_warp_spec': 0, + 'num_consumer_groups': 0, + 'num_ctas': 1, + 'num_stages': 3, + 'num_warps': 8, + 'pre_hook': None, + 'reg_dec_producer': 0, + 'reg_inc_consumer': 0}, + "(12, 1024, 1024, 27, True, 'torch.float16', 'torch.float16', 'torch.float16', 'torch.uint32', 'torch.int64', 'torch.float16')": {'kwargs': {'B1': 128, + 'B2': 256, + 'BK': 64}, + 'maxnreg': None, + 'num_buffers_warp_spec': 0, + 'num_consumer_groups': 0, + 'num_ctas': 1, + 'num_stages': 3, + 'num_warps': 8, + 'pre_hook': None, + 'reg_dec_producer': 0, + 'reg_inc_consumer': 0}, + "(13, 1024, 1024, 27, True, 'torch.float16', 'torch.float16', 'torch.float16', 'torch.uint32', 'torch.int64', 'torch.float16')": {'kwargs': {'B1': 128, + 'B2': 256, + 'BK': 64}, + 'maxnreg': None, + 'num_buffers_warp_spec': 0, + 'num_consumer_groups': 0, + 'num_ctas': 1, + 'num_stages': 3, + 'num_warps': 8, + 'pre_hook': None, + 'reg_dec_producer': 0, + 'reg_inc_consumer': 0}, + "(13, 1024, 4096, 27, True, 'torch.float16', 'torch.float16', 'torch.float16', 'torch.uint32', 'torch.int64', 'torch.float16')": {'kwargs': {'B1': 128, + 'B2': 256, + 'BK': 64}, + 'maxnreg': None, + 'num_buffers_warp_spec': 0, + 'num_consumer_groups': 0, + 'num_ctas': 1, + 'num_stages': 3, + 'num_warps': 8, + 'pre_hook': None, + 'reg_dec_producer': 0, + 'reg_inc_consumer': 0}, + "(13, 512, 128, 27, True, 'torch.float16', 'torch.float16', 'torch.float16', 'torch.uint32', 'torch.int64', 'torch.float16')": {'kwargs': {'B1': 64, + 'B2': 128, + 'BK': 32}, + 'maxnreg': None, + 'num_buffers_warp_spec': 0, + 'num_consumer_groups': 0, + 'num_ctas': 1, + 'num_stages': 4, + 'num_warps': 4, + 'pre_hook': None, + 'reg_dec_producer': 0, + 'reg_inc_consumer': 0}, + "(13, 512, 2048, 27, True, 'torch.float16', 'torch.float16', 'torch.float16', 'torch.uint32', 'torch.int64', 'torch.float16')": {'kwargs': {'B1': 128, + 'B2': 256, + 'BK': 64}, + 'maxnreg': None, + 'num_buffers_warp_spec': 0, + 'num_consumer_groups': 0, + 'num_ctas': 1, + 'num_stages': 3, + 'num_warps': 8, + 'pre_hook': None, + 'reg_dec_producer': 0, + 'reg_inc_consumer': 0}, + "(13, 512, 512, 27, True, 'torch.float16', 'torch.float16', 'torch.float16', 'torch.uint32', 'torch.int64', 'torch.float16')": {'kwargs': {'B1': 128, + 'B2': 256, + 'BK': 64}, + 'maxnreg': None, + 'num_buffers_warp_spec': 0, + 'num_consumer_groups': 0, + 'num_ctas': 1, + 'num_stages': 3, + 'num_warps': 8, + 'pre_hook': None, + 'reg_dec_producer': 0, + 'reg_inc_consumer': 0}, + "(15, 256, 1024, 27, True, 'torch.float16', 'torch.float16', 'torch.float16', 'torch.uint32', 'torch.int64', 'torch.float16')": {'kwargs': {'B1': 128, + 'B2': 256, + 'BK': 64}, + 'maxnreg': None, + 'num_buffers_warp_spec': 0, + 'num_consumer_groups': 0, + 'num_ctas': 1, + 'num_stages': 3, + 'num_warps': 8, + 'pre_hook': None, + 'reg_dec_producer': 0, + 'reg_inc_consumer': 0}, + "(15, 256, 256, 27, True, 'torch.float16', 'torch.float16', 'torch.float16', 'torch.uint32', 'torch.int64', 'torch.float16')": {'kwargs': {'B1': 128, + 'B2': 256, + 'BK': 64}, + 'maxnreg': None, + 'num_buffers_warp_spec': 0, + 'num_consumer_groups': 0, + 'num_ctas': 1, + 'num_stages': 3, + 'num_warps': 8, + 'pre_hook': None, + 'reg_dec_producer': 0, + 'reg_inc_consumer': 0}, + "(15, 256, 64, 27, True, 'torch.float16', 'torch.float16', 'torch.float16', 'torch.uint32', 'torch.int64', 'torch.float16')": {'kwargs': {'B1': 128, + 'B2': 64, + 'BK': 32}, + 'maxnreg': None, + 'num_buffers_warp_spec': 0, + 'num_consumer_groups': 0, + 'num_ctas': 1, + 'num_stages': 4, + 'num_warps': 4, + 'pre_hook': None, + 'reg_dec_producer': 0, + 'reg_inc_consumer': 0}, + "(15, 512, 128, 27, True, 'torch.float16', 'torch.float16', 'torch.float16', 'torch.uint32', 'torch.int64', 'torch.float16')": {'kwargs': {'B1': 64, + 'B2': 128, + 'BK': 32}, + 'maxnreg': None, + 'num_buffers_warp_spec': 0, + 'num_consumer_groups': 0, + 'num_ctas': 1, + 'num_stages': 4, + 'num_warps': 4, + 'pre_hook': None, + 'reg_dec_producer': 0, + 'reg_inc_consumer': 0}, + "(15, 512, 2048, 27, True, 'torch.float16', 'torch.float16', 'torch.float16', 'torch.uint32', 'torch.int64', 'torch.float16')": {'kwargs': {'B1': 128, + 'B2': 256, + 'BK': 64}, + 'maxnreg': None, + 'num_buffers_warp_spec': 0, + 'num_consumer_groups': 0, + 'num_ctas': 1, + 'num_stages': 3, + 'num_warps': 8, + 'pre_hook': None, + 'reg_dec_producer': 0, + 'reg_inc_consumer': 0}, + "(15, 512, 512, 27, True, 'torch.float16', 'torch.float16', 'torch.float16', 'torch.uint32', 'torch.int64', 'torch.float16')": {'kwargs': {'B1': 128, + 'B2': 256, + 'BK': 64}, + 'maxnreg': None, + 'num_buffers_warp_spec': 0, + 'num_consumer_groups': 0, + 'num_ctas': 1, + 'num_stages': 3, + 'num_warps': 8, + 'pre_hook': None, + 'reg_dec_producer': 0, + 'reg_inc_consumer': 0}, + "(17, 128, 128, 27, True, 'torch.float16', 'torch.float16', 'torch.float16', 'torch.uint32', 'torch.int64', 'torch.float16')": {'kwargs': {'B1': 64, + 'B2': 128, + 'BK': 32}, + 'maxnreg': None, + 'num_buffers_warp_spec': 0, + 'num_consumer_groups': 0, + 'num_ctas': 1, + 'num_stages': 4, + 'num_warps': 4, + 'pre_hook': None, + 'reg_dec_producer': 0, + 'reg_inc_consumer': 0}, + "(17, 128, 32, 27, True, 'torch.float16', 'torch.float16', 'torch.float16', 'torch.uint32', 'torch.int64', 'torch.float16')": {'kwargs': {'B1': 64, + 'B2': 32, + 'BK': 32}, + 'maxnreg': None, + 'num_buffers_warp_spec': 0, + 'num_consumer_groups': 0, + 'num_ctas': 1, + 'num_stages': 5, + 'num_warps': 2, + 'pre_hook': None, + 'reg_dec_producer': 0, + 'reg_inc_consumer': 0}, + "(17, 128, 512, 27, True, 'torch.float16', 'torch.float16', 'torch.float16', 'torch.uint32', 'torch.int64', 'torch.float16')": {'kwargs': {'B1': 128, + 'B2': 256, + 'BK': 64}, + 'maxnreg': None, + 'num_buffers_warp_spec': 0, + 'num_consumer_groups': 0, + 'num_ctas': 1, + 'num_stages': 3, + 'num_warps': 8, + 'pre_hook': None, + 'reg_dec_producer': 0, + 'reg_inc_consumer': 0}, + "(17, 256, 1024, 27, True, 'torch.float16', 'torch.float16', 'torch.float16', 'torch.uint32', 'torch.int64', 'torch.float16')": {'kwargs': {'B1': 128, + 'B2': 256, + 'BK': 64}, + 'maxnreg': None, + 'num_buffers_warp_spec': 0, + 'num_consumer_groups': 0, + 'num_ctas': 1, + 'num_stages': 3, + 'num_warps': 8, + 'pre_hook': None, + 'reg_dec_producer': 0, + 'reg_inc_consumer': 0}, + "(17, 256, 256, 27, True, 'torch.float16', 'torch.float16', 'torch.float16', 'torch.uint32', 'torch.int64', 'torch.float16')": {'kwargs': {'B1': 128, + 'B2': 256, + 'BK': 64}, + 'maxnreg': None, + 'num_buffers_warp_spec': 0, + 'num_consumer_groups': 0, + 'num_ctas': 1, + 'num_stages': 3, + 'num_warps': 8, + 'pre_hook': None, + 'reg_dec_producer': 0, + 'reg_inc_consumer': 0}, + "(17, 256, 64, 27, True, 'torch.float16', 'torch.float16', 'torch.float16', 'torch.uint32', 'torch.int64', 'torch.float16')": {'kwargs': {'B1': 128, + 'B2': 64, + 'BK': 32}, + 'maxnreg': None, + 'num_buffers_warp_spec': 0, + 'num_consumer_groups': 0, + 'num_ctas': 1, + 'num_stages': 4, + 'num_warps': 4, + 'pre_hook': None, + 'reg_dec_producer': 0, + 'reg_inc_consumer': 0}, + "(18, 64, 16, 27, True, 'torch.float16', 'torch.float16', 'torch.float16', 'torch.uint32', 'torch.int64', 'torch.float16')": {'kwargs': {'B1': 64, + 'B2': 32, + 'BK': 32}, + 'maxnreg': None, + 'num_buffers_warp_spec': 0, + 'num_consumer_groups': 0, + 'num_ctas': 1, + 'num_stages': 5, + 'num_warps': 2, + 'pre_hook': None, + 'reg_dec_producer': 0, + 'reg_inc_consumer': 0}, + "(19, 128, 128, 27, True, 'torch.float16', 'torch.float16', 'torch.float16', 'torch.uint32', 'torch.int64', 'torch.float16')": {'kwargs': {'B1': 64, + 'B2': 128, + 'BK': 32}, + 'maxnreg': None, + 'num_buffers_warp_spec': 0, + 'num_consumer_groups': 0, + 'num_ctas': 1, + 'num_stages': 4, + 'num_warps': 4, + 'pre_hook': None, + 'reg_dec_producer': 0, + 'reg_inc_consumer': 0}, + "(19, 128, 32, 27, True, 'torch.float16', 'torch.float16', 'torch.float16', 'torch.uint32', 'torch.int64', 'torch.float16')": {'kwargs': {'B1': 64, + 'B2': 32, + 'BK': 32}, + 'maxnreg': None, + 'num_buffers_warp_spec': 0, + 'num_consumer_groups': 0, + 'num_ctas': 1, + 'num_stages': 5, + 'num_warps': 2, + 'pre_hook': None, + 'reg_dec_producer': 0, + 'reg_inc_consumer': 0}, + "(19, 128, 512, 27, True, 'torch.float16', 'torch.float16', 'torch.float16', 'torch.uint32', 'torch.int64', 'torch.float16')": {'kwargs': {'B1': 128, + 'B2': 256, + 'BK': 64}, + 'maxnreg': None, + 'num_buffers_warp_spec': 0, + 'num_consumer_groups': 0, + 'num_ctas': 1, + 'num_stages': 3, + 'num_warps': 8, + 'pre_hook': None, + 'reg_dec_producer': 0, + 'reg_inc_consumer': 0}, + "(19, 64, 64, 27, True, 'torch.float16', 'torch.float16', 'torch.float16', 'torch.uint32', 'torch.int64', 'torch.float16')": {'kwargs': {'B1': 128, + 'B2': 64, + 'BK': 32}, + 'maxnreg': None, + 'num_buffers_warp_spec': 0, + 'num_consumer_groups': 0, + 'num_ctas': 1, + 'num_stages': 4, + 'num_warps': 4, + 'pre_hook': None, + 'reg_dec_producer': 0, + 'reg_inc_consumer': 0}, + "(21, 64, 16, 27, True, 'torch.float16', 'torch.float16', 'torch.float16', 'torch.uint32', 'torch.int64', 'torch.float16')": {'kwargs': {'B1': 64, + 'B2': 32, + 'BK': 32}, + 'maxnreg': None, + 'num_buffers_warp_spec': 0, + 'num_consumer_groups': 0, + 'num_ctas': 1, + 'num_stages': 5, + 'num_warps': 2, + 'pre_hook': None, + 'reg_dec_producer': 0, + 'reg_inc_consumer': 0}, + "(21, 64, 64, 27, True, 'torch.float16', 'torch.float16', 'torch.float16', 'torch.uint32', 'torch.int64', 'torch.float16')": {'kwargs': {'B1': 32, + 'B2': 64, + 'BK': 32}, + 'maxnreg': None, + 'num_buffers_warp_spec': 0, + 'num_consumer_groups': 0, + 'num_ctas': 1, + 'num_stages': 5, + 'num_warps': 2, + 'pre_hook': None, + 'reg_dec_producer': 0, + 'reg_inc_consumer': 0}}, + 'flex_gemm.kernels.triton.spconv.sparse_submanifold_conv_fwd_masked_implicit_gemm_splitk.sparse_submanifold_conv_fwd_masked_implicit_gemm_splitk': {'(2^10, 1024, 1024, 27)': {'SPLITK': 8}, + '(2^10, 1024, 4096, 27)': {'SPLITK': 4}, + '(2^12, 1024, 1024, 27)': {'SPLITK': 8}, + '(2^13, 1024, 1024, 27)': {'SPLITK': 4}, + '(2^13, 1024, 4096, 27)': {'SPLITK': 1}, + '(2^13, 512, 128, 27)': {'SPLITK': 1}, + '(2^13, 512, 2048, 27)': {'SPLITK': 1}, + '(2^13, 512, 512, 27)': {'SPLITK': 4}, + '(2^15, 256, 1024, 27)': {'SPLITK': 1}, + '(2^15, 256, 256, 27)': {'SPLITK': 1}, + '(2^15, 256, 64, 27)': {'SPLITK': 1}, + '(2^15, 512, 128, 27)': {'SPLITK': 1}, + '(2^15, 512, 2048, 27)': {'SPLITK': 1}, + '(2^15, 512, 512, 27)': {'SPLITK': 1}, + '(2^17, 128, 128, 27)': {'SPLITK': 1}, + '(2^17, 128, 32, 27)': {'SPLITK': 1}, + '(2^17, 128, 512, 27)': {'SPLITK': 1}, + '(2^17, 256, 1024, 27)': {'SPLITK': 1}, + '(2^17, 256, 256, 27)': {'SPLITK': 1}, + '(2^17, 256, 64, 27)': {'SPLITK': 1}, + '(2^18, 64, 16, 27)': {'SPLITK': 1}, + '(2^19, 128, 128, 27)': {'SPLITK': 1}, + '(2^19, 128, 32, 27)': {'SPLITK': 1}, + '(2^19, 128, 512, 27)': {'SPLITK': 1}, + '(2^19, 64, 64, 27)': {'SPLITK': 1}, + '(2^21, 64, 16, 27)': {'SPLITK': 1}, + '(2^21, 64, 64, 27)': {'SPLITK': 1}}, + 'flex_gemm.kernels.triton.spconv.sparse_submanifold_conv_fwd_masked_implicit_gemm_splitk.sparse_submanifold_conv_fwd_masked_implicit_gemm_splitk_kernel': {"(10, 1024, 1024, 27, 16, True, 'torch.float16', 'torch.float16', 'torch.float16', 'torch.uint32', 'torch.int64', 'torch.float32')": {'kwargs': {'B1': 128, + 'B2': 256, + 'BK': 64}, + 'maxnreg': None, + 'num_buffers_warp_spec': 0, + 'num_consumer_groups': 0, + 'num_ctas': 1, + 'num_stages': 3, + 'num_warps': 8, + 'pre_hook': None, + 'reg_dec_producer': 0, + 'reg_inc_consumer': 0}, + "(10, 1024, 1024, 27, 2, True, 'torch.float16', 'torch.float16', 'torch.float16', 'torch.uint32', 'torch.int64', 'torch.float32')": {'kwargs': {'B1': 128, + 'B2': 256, + 'BK': 64}, + 'maxnreg': None, + 'num_buffers_warp_spec': 0, + 'num_consumer_groups': 0, + 'num_ctas': 1, + 'num_stages': 3, + 'num_warps': 8, + 'pre_hook': None, + 'reg_dec_producer': 0, + 'reg_inc_consumer': 0}, + "(10, 1024, 1024, 27, 32, True, 'torch.float16', 'torch.float16', 'torch.float16', 'torch.uint32', 'torch.int64', 'torch.float32')": {'kwargs': {'B1': 128, + 'B2': 256, + 'BK': 64}, + 'maxnreg': None, + 'num_buffers_warp_spec': 0, + 'num_consumer_groups': 0, + 'num_ctas': 1, + 'num_stages': 3, + 'num_warps': 8, + 'pre_hook': None, + 'reg_dec_producer': 0, + 'reg_inc_consumer': 0}, + "(10, 1024, 1024, 27, 4, True, 'torch.float16', 'torch.float16', 'torch.float16', 'torch.uint32', 'torch.int64', 'torch.float32')": {'kwargs': {'B1': 128, + 'B2': 256, + 'BK': 64}, + 'maxnreg': None, + 'num_buffers_warp_spec': 0, + 'num_consumer_groups': 0, + 'num_ctas': 1, + 'num_stages': 3, + 'num_warps': 8, + 'pre_hook': None, + 'reg_dec_producer': 0, + 'reg_inc_consumer': 0}, + "(10, 1024, 1024, 27, 8, True, 'torch.float16', 'torch.float16', 'torch.float16', 'torch.uint32', 'torch.int64', 'torch.float32')": {'kwargs': {'B1': 128, + 'B2': 256, + 'BK': 64}, + 'maxnreg': None, + 'num_buffers_warp_spec': 0, + 'num_consumer_groups': 0, + 'num_ctas': 1, + 'num_stages': 3, + 'num_warps': 8, + 'pre_hook': None, + 'reg_dec_producer': 0, + 'reg_inc_consumer': 0}, + "(10, 1024, 4096, 27, 2, True, 'torch.float16', 'torch.float16', 'torch.float16', 'torch.uint32', 'torch.int64', 'torch.float32')": {'kwargs': {'B1': 128, + 'B2': 256, + 'BK': 64}, + 'maxnreg': None, + 'num_buffers_warp_spec': 0, + 'num_consumer_groups': 0, + 'num_ctas': 1, + 'num_stages': 3, + 'num_warps': 8, + 'pre_hook': None, + 'reg_dec_producer': 0, + 'reg_inc_consumer': 0}, + "(10, 1024, 4096, 27, 4, True, 'torch.float16', 'torch.float16', 'torch.float16', 'torch.uint32', 'torch.int64', 'torch.float32')": {'kwargs': {'B1': 128, + 'B2': 256, + 'BK': 64}, + 'maxnreg': None, + 'num_buffers_warp_spec': 0, + 'num_consumer_groups': 0, + 'num_ctas': 1, + 'num_stages': 3, + 'num_warps': 8, + 'pre_hook': None, + 'reg_dec_producer': 0, + 'reg_inc_consumer': 0}, + "(10, 1024, 4096, 27, 8, True, 'torch.float16', 'torch.float16', 'torch.float16', 'torch.uint32', 'torch.int64', 'torch.float32')": {'kwargs': {'B1': 128, + 'B2': 256, + 'BK': 64}, + 'maxnreg': None, + 'num_buffers_warp_spec': 0, + 'num_consumer_groups': 0, + 'num_ctas': 1, + 'num_stages': 3, + 'num_warps': 8, + 'pre_hook': None, + 'reg_dec_producer': 0, + 'reg_inc_consumer': 0}, + "(12, 1024, 1024, 27, 2, True, 'torch.float16', 'torch.float16', 'torch.float16', 'torch.uint32', 'torch.int64', 'torch.float32')": {'kwargs': {'B1': 128, + 'B2': 256, + 'BK': 64}, + 'maxnreg': None, + 'num_buffers_warp_spec': 0, + 'num_consumer_groups': 0, + 'num_ctas': 1, + 'num_stages': 3, + 'num_warps': 8, + 'pre_hook': None, + 'reg_dec_producer': 0, + 'reg_inc_consumer': 0}, + "(12, 1024, 1024, 27, 4, True, 'torch.float16', 'torch.float16', 'torch.float16', 'torch.uint32', 'torch.int64', 'torch.float32')": {'kwargs': {'B1': 128, + 'B2': 256, + 'BK': 64}, + 'maxnreg': None, + 'num_buffers_warp_spec': 0, + 'num_consumer_groups': 0, + 'num_ctas': 1, + 'num_stages': 3, + 'num_warps': 8, + 'pre_hook': None, + 'reg_dec_producer': 0, + 'reg_inc_consumer': 0}, + "(12, 1024, 1024, 27, 8, True, 'torch.float16', 'torch.float16', 'torch.float16', 'torch.uint32', 'torch.int64', 'torch.float32')": {'kwargs': {'B1': 128, + 'B2': 256, + 'BK': 64}, + 'maxnreg': None, + 'num_buffers_warp_spec': 0, + 'num_consumer_groups': 0, + 'num_ctas': 1, + 'num_stages': 3, + 'num_warps': 8, + 'pre_hook': None, + 'reg_dec_producer': 0, + 'reg_inc_consumer': 0}, + "(13, 1024, 1024, 27, 2, True, 'torch.float16', 'torch.float16', 'torch.float16', 'torch.uint32', 'torch.int64', 'torch.float32')": {'kwargs': {'B1': 128, + 'B2': 256, + 'BK': 64}, + 'maxnreg': None, + 'num_buffers_warp_spec': 0, + 'num_consumer_groups': 0, + 'num_ctas': 1, + 'num_stages': 3, + 'num_warps': 8, + 'pre_hook': None, + 'reg_dec_producer': 0, + 'reg_inc_consumer': 0}, + "(13, 1024, 1024, 27, 4, True, 'torch.float16', 'torch.float16', 'torch.float16', 'torch.uint32', 'torch.int64', 'torch.float32')": {'kwargs': {'B1': 128, + 'B2': 256, + 'BK': 64}, + 'maxnreg': None, + 'num_buffers_warp_spec': 0, + 'num_consumer_groups': 0, + 'num_ctas': 1, + 'num_stages': 3, + 'num_warps': 8, + 'pre_hook': None, + 'reg_dec_producer': 0, + 'reg_inc_consumer': 0}, + "(13, 512, 128, 27, 16, True, 'torch.float16', 'torch.float16', 'torch.float16', 'torch.uint32', 'torch.int64', 'torch.float32')": {'kwargs': {'B1': 128, + 'B2': 128, + 'BK': 32}, + 'maxnreg': None, + 'num_buffers_warp_spec': 0, + 'num_consumer_groups': 0, + 'num_ctas': 1, + 'num_stages': 4, + 'num_warps': 4, + 'pre_hook': None, + 'reg_dec_producer': 0, + 'reg_inc_consumer': 0}, + "(13, 512, 128, 27, 2, True, 'torch.float16', 'torch.float16', 'torch.float16', 'torch.uint32', 'torch.int64', 'torch.float32')": {'kwargs': {'B1': 128, + 'B2': 128, + 'BK': 32}, + 'maxnreg': None, + 'num_buffers_warp_spec': 0, + 'num_consumer_groups': 0, + 'num_ctas': 1, + 'num_stages': 4, + 'num_warps': 4, + 'pre_hook': None, + 'reg_dec_producer': 0, + 'reg_inc_consumer': 0}, + "(13, 512, 128, 27, 32, True, 'torch.float16', 'torch.float16', 'torch.float16', 'torch.uint32', 'torch.int64', 'torch.float32')": {'kwargs': {'B1': 128, + 'B2': 64, + 'BK': 32}, + 'maxnreg': None, + 'num_buffers_warp_spec': 0, + 'num_consumer_groups': 0, + 'num_ctas': 1, + 'num_stages': 4, + 'num_warps': 4, + 'pre_hook': None, + 'reg_dec_producer': 0, + 'reg_inc_consumer': 0}, + "(13, 512, 128, 27, 4, True, 'torch.float16', 'torch.float16', 'torch.float16', 'torch.uint32', 'torch.int64', 'torch.float32')": {'kwargs': {'B1': 128, + 'B2': 128, + 'BK': 32}, + 'maxnreg': None, + 'num_buffers_warp_spec': 0, + 'num_consumer_groups': 0, + 'num_ctas': 1, + 'num_stages': 4, + 'num_warps': 4, + 'pre_hook': None, + 'reg_dec_producer': 0, + 'reg_inc_consumer': 0}, + "(13, 512, 128, 27, 8, True, 'torch.float16', 'torch.float16', 'torch.float16', 'torch.uint32', 'torch.int64', 'torch.float32')": {'kwargs': {'B1': 128, + 'B2': 128, + 'BK': 32}, + 'maxnreg': None, + 'num_buffers_warp_spec': 0, + 'num_consumer_groups': 0, + 'num_ctas': 1, + 'num_stages': 4, + 'num_warps': 4, + 'pre_hook': None, + 'reg_dec_producer': 0, + 'reg_inc_consumer': 0}, + "(13, 512, 2048, 27, 2, True, 'torch.float16', 'torch.float16', 'torch.float16', 'torch.uint32', 'torch.int64', 'torch.float32')": {'kwargs': {'B1': 128, + 'B2': 256, + 'BK': 64}, + 'maxnreg': None, + 'num_buffers_warp_spec': 0, + 'num_consumer_groups': 0, + 'num_ctas': 1, + 'num_stages': 3, + 'num_warps': 8, + 'pre_hook': None, + 'reg_dec_producer': 0, + 'reg_inc_consumer': 0}, + "(13, 512, 512, 27, 2, True, 'torch.float16', 'torch.float16', 'torch.float16', 'torch.uint32', 'torch.int64', 'torch.float32')": {'kwargs': {'B1': 128, + 'B2': 256, + 'BK': 64}, + 'maxnreg': None, + 'num_buffers_warp_spec': 0, + 'num_consumer_groups': 0, + 'num_ctas': 1, + 'num_stages': 3, + 'num_warps': 8, + 'pre_hook': None, + 'reg_dec_producer': 0, + 'reg_inc_consumer': 0}, + "(13, 512, 512, 27, 4, True, 'torch.float16', 'torch.float16', 'torch.float16', 'torch.uint32', 'torch.int64', 'torch.float32')": {'kwargs': {'B1': 128, + 'B2': 256, + 'BK': 64}, + 'maxnreg': None, + 'num_buffers_warp_spec': 0, + 'num_consumer_groups': 0, + 'num_ctas': 1, + 'num_stages': 3, + 'num_warps': 8, + 'pre_hook': None, + 'reg_dec_producer': 0, + 'reg_inc_consumer': 0}, + "(13, 512, 512, 27, 8, True, 'torch.float16', 'torch.float16', 'torch.float16', 'torch.uint32', 'torch.int64', 'torch.float32')": {'kwargs': {'B1': 128, + 'B2': 256, + 'BK': 64}, + 'maxnreg': None, + 'num_buffers_warp_spec': 0, + 'num_consumer_groups': 0, + 'num_ctas': 1, + 'num_stages': 3, + 'num_warps': 8, + 'pre_hook': None, + 'reg_dec_producer': 0, + 'reg_inc_consumer': 0}, + "(15, 256, 256, 27, 2, True, 'torch.float16', 'torch.float16', 'torch.float16', 'torch.uint32', 'torch.int64', 'torch.float32')": {'kwargs': {'B1': 128, + 'B2': 256, + 'BK': 64}, + 'maxnreg': None, + 'num_buffers_warp_spec': 0, + 'num_consumer_groups': 0, + 'num_ctas': 1, + 'num_stages': 3, + 'num_warps': 8, + 'pre_hook': None, + 'reg_dec_producer': 0, + 'reg_inc_consumer': 0}, + "(15, 256, 256, 27, 4, True, 'torch.float16', 'torch.float16', 'torch.float16', 'torch.uint32', 'torch.int64', 'torch.float32')": {'kwargs': {'B1': 128, + 'B2': 256, + 'BK': 64}, + 'maxnreg': None, + 'num_buffers_warp_spec': 0, + 'num_consumer_groups': 0, + 'num_ctas': 1, + 'num_stages': 3, + 'num_warps': 8, + 'pre_hook': None, + 'reg_dec_producer': 0, + 'reg_inc_consumer': 0}, + "(15, 256, 64, 27, 2, True, 'torch.float16', 'torch.float16', 'torch.float16', 'torch.uint32', 'torch.int64', 'torch.float32')": {'kwargs': {'B1': 128, + 'B2': 64, + 'BK': 32}, + 'maxnreg': None, + 'num_buffers_warp_spec': 0, + 'num_consumer_groups': 0, + 'num_ctas': 1, + 'num_stages': 4, + 'num_warps': 4, + 'pre_hook': None, + 'reg_dec_producer': 0, + 'reg_inc_consumer': 0}, + "(15, 256, 64, 27, 4, True, 'torch.float16', 'torch.float16', 'torch.float16', 'torch.uint32', 'torch.int64', 'torch.float32')": {'kwargs': {'B1': 128, + 'B2': 64, + 'BK': 32}, + 'maxnreg': None, + 'num_buffers_warp_spec': 0, + 'num_consumer_groups': 0, + 'num_ctas': 1, + 'num_stages': 4, + 'num_warps': 4, + 'pre_hook': None, + 'reg_dec_producer': 0, + 'reg_inc_consumer': 0}, + "(15, 256, 64, 27, 8, True, 'torch.float16', 'torch.float16', 'torch.float16', 'torch.uint32', 'torch.int64', 'torch.float32')": {'kwargs': {'B1': 128, + 'B2': 64, + 'BK': 32}, + 'maxnreg': None, + 'num_buffers_warp_spec': 0, + 'num_consumer_groups': 0, + 'num_ctas': 1, + 'num_stages': 4, + 'num_warps': 4, + 'pre_hook': None, + 'reg_dec_producer': 0, + 'reg_inc_consumer': 0}, + "(15, 512, 128, 27, 2, True, 'torch.float16', 'torch.float16', 'torch.float16', 'torch.uint32', 'torch.int64', 'torch.float32')": {'kwargs': {'B1': 128, + 'B2': 64, + 'BK': 32}, + 'maxnreg': None, + 'num_buffers_warp_spec': 0, + 'num_consumer_groups': 0, + 'num_ctas': 1, + 'num_stages': 4, + 'num_warps': 4, + 'pre_hook': None, + 'reg_dec_producer': 0, + 'reg_inc_consumer': 0}, + "(15, 512, 128, 27, 4, True, 'torch.float16', 'torch.float16', 'torch.float16', 'torch.uint32', 'torch.int64', 'torch.float32')": {'kwargs': {'B1': 128, + 'B2': 64, + 'BK': 32}, + 'maxnreg': None, + 'num_buffers_warp_spec': 0, + 'num_consumer_groups': 0, + 'num_ctas': 1, + 'num_stages': 4, + 'num_warps': 4, + 'pre_hook': None, + 'reg_dec_producer': 0, + 'reg_inc_consumer': 0}, + "(15, 512, 128, 27, 8, True, 'torch.float16', 'torch.float16', 'torch.float16', 'torch.uint32', 'torch.int64', 'torch.float32')": {'kwargs': {'B1': 128, + 'B2': 64, + 'BK': 32}, + 'maxnreg': None, + 'num_buffers_warp_spec': 0, + 'num_consumer_groups': 0, + 'num_ctas': 1, + 'num_stages': 4, + 'num_warps': 4, + 'pre_hook': None, + 'reg_dec_producer': 0, + 'reg_inc_consumer': 0}, + "(15, 512, 512, 27, 2, True, 'torch.float16', 'torch.float16', 'torch.float16', 'torch.uint32', 'torch.int64', 'torch.float32')": {'kwargs': {'B1': 128, + 'B2': 256, + 'BK': 64}, + 'maxnreg': None, + 'num_buffers_warp_spec': 0, + 'num_consumer_groups': 0, + 'num_ctas': 1, + 'num_stages': 3, + 'num_warps': 8, + 'pre_hook': None, + 'reg_dec_producer': 0, + 'reg_inc_consumer': 0}, + "(17, 128, 128, 27, 2, True, 'torch.float16', 'torch.float16', 'torch.float16', 'torch.uint32', 'torch.int64', 'torch.float32')": {'kwargs': {'B1': 128, + 'B2': 64, + 'BK': 32}, + 'maxnreg': None, + 'num_buffers_warp_spec': 0, + 'num_consumer_groups': 0, + 'num_ctas': 1, + 'num_stages': 4, + 'num_warps': 4, + 'pre_hook': None, + 'reg_dec_producer': 0, + 'reg_inc_consumer': 0}, + "(17, 128, 32, 27, 2, True, 'torch.float16', 'torch.float16', 'torch.float16', 'torch.uint32', 'torch.int64', 'torch.float32')": {'kwargs': {'B1': 64, + 'B2': 32, + 'BK': 32}, + 'maxnreg': None, + 'num_buffers_warp_spec': 0, + 'num_consumer_groups': 0, + 'num_ctas': 1, + 'num_stages': 5, + 'num_warps': 2, + 'pre_hook': None, + 'reg_dec_producer': 0, + 'reg_inc_consumer': 0}, + "(17, 256, 64, 27, 2, True, 'torch.float16', 'torch.float16', 'torch.float16', 'torch.uint32', 'torch.int64', 'torch.float32')": {'kwargs': {'B1': 128, + 'B2': 64, + 'BK': 32}, + 'maxnreg': None, + 'num_buffers_warp_spec': 0, + 'num_consumer_groups': 0, + 'num_ctas': 1, + 'num_stages': 4, + 'num_warps': 4, + 'pre_hook': None, + 'reg_dec_producer': 0, + 'reg_inc_consumer': 0}}}} diff --git a/supratex/inference/geometry_cleanup.py b/supratex/inference/geometry_cleanup.py new file mode 100644 index 0000000000000000000000000000000000000000..47b6de8419c849cd9d3a5bcaad9e028436f79b1f --- /dev/null +++ b/supratex/inference/geometry_cleanup.py @@ -0,0 +1,114 @@ +"""Optional removal of floating geometry that the reference image does not show. + +Generated geometry can contain disconnected pieces outside the object in the +image. A piece is removed when it is not the largest connected component and +most of its faces project outside the image foreground, which is dilated to +tolerate ordinary misalignment. Parts hidden behind the foreground, such as +the back of the object, still project inside it and are kept, as are parts of +any size that the image shows. Extra geometry connected to the object is not +affected. +""" +from __future__ import annotations + +import math +from typing import Any, Mapping + +import numpy as np +import torch + +FOREGROUND_ALPHA = 128 +# Misalignment tolerance as a fraction of the image width (16 px at 1024). +MARGIN = 16 / 1024 +MIN_INSIDE = 0.5 +# Vertices closer than this are welded, so UV-seam duplicates stay connected. +WELD_STEP = 1e-6 +KEY_BITS = 21 + + +def _weld(vertices: torch.Tensor) -> torch.Tensor: + """Map each vertex to the index of its quantized position.""" + low = vertices.amin(0) + step = max(WELD_STEP, float((vertices.amax(0) - low).max()) / (2 ** KEY_BITS - 1)) + q = torch.round((vertices.double() - low) / step).long() + key = (q[:, 0] << (2 * KEY_BITS)) | (q[:, 1] << KEY_BITS) | q[:, 2] + return torch.unique(key, return_inverse=True)[1] + + +def _components(vertices: torch.Tensor, faces: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + """Label faces by connected component with parallel union-find. + + Each round hooks every root to the smallest root among its edges, then + compresses paths by pointer jumping; components halve per round. + """ + welded = _weld(vertices)[faces] + count = int(welded.max()) + 1 + first = torch.cat((welded[:, 0], welded[:, 1])) + second = torch.cat((welded[:, 1], welded[:, 2])) + parent = torch.arange(count, device=faces.device) + while True: + a, b = parent[first], parent[second] + differ = a != b + if not bool(differ.any()): + break + low, high = torch.minimum(a[differ], b[differ]), torch.maximum(a[differ], b[differ]) + parent.scatter_reduce_(0, high, low, reduce="amin") + while True: + jumped = parent[parent] + if torch.equal(jumped, parent): + break + parent = jumped + first, second = first[differ], second[differ] + face_labels = torch.unique(parent[welded[:, 0]], return_inverse=True)[1] + return face_labels, torch.bincount(face_labels) + + +def _project(points: torch.Tensor, camera: Mapping[str, Any]) -> tuple[torch.Tensor, torch.Tensor]: + """Normalized image coordinates (y down) and camera depth, matching completion.""" + if "extrinsic" in camera: + extrinsic = torch.as_tensor(camera["extrinsic"], dtype=points.dtype, device=points.device) + intrinsic = torch.as_tensor(camera["intrinsic"], dtype=points.dtype, device=points.device) + local = points @ extrinsic[:3, :3].T + extrinsic[:3, 3] + uvw = local @ intrinsic.T + return uvw[:, :2] / uvw[:, 2:3].clamp(min=1e-8), local[:, 2] + distance = float(camera["distance"]) * float(camera.get("mesh_scale", 1.0)) + focal = 0.5 / math.tan(float(camera["camera_angle_x"]) / 2) + depth = distance - points[:, 2] + safe = depth.clamp(min=1e-8) + return torch.stack((points[:, 0] / safe * focal + 0.5, -points[:, 1] / safe * focal + 0.5), dim=1), depth + + +def _foreground(alpha: np.ndarray) -> np.ndarray: + """Foreground mask dilated by the misalignment margin.""" + from scipy.ndimage import distance_transform_edt + + radius = max(1, round(MARGIN * alpha.shape[1])) + return distance_transform_edt(alpha < FOREGROUND_ALPHA) <= radius + + +def remove_floaters(vertices: torch.Tensor, faces: torch.Tensor, camera: Mapping[str, Any], + alpha: np.ndarray, device: str | None = None) -> tuple[torch.Tensor, torch.Tensor, dict[str, Any]]: + """Drop disconnected components that project mostly outside the image foreground.""" + device = device or ("cuda" if torch.cuda.is_available() else "cpu") + source_device = vertices.device + vertices, faces = vertices.to(device), faces.to(device).long() + face_labels, counts = _components(vertices, faces) + report = {"components": int(len(counts)), "removed_components": 0, "removed_faces": 0, "faces": int(len(faces))} + floating = None + if len(counts) > 1: + height, width = alpha.shape + foreground = torch.from_numpy(_foreground(alpha)).to(device) + uv, depth = _project(vertices[faces].mean(dim=1), camera) + x = torch.floor(uv[:, 0] * width).long() + y = torch.floor(uv[:, 1] * height).long() + valid = (depth > 0) & (x >= 0) & (x < width) & (y >= 0) & (y < height) + inside = torch.zeros_like(valid) + inside[valid] = foreground[y[valid], x[valid]] + inside_count = torch.bincount(face_labels, weights=inside.double(), minlength=len(counts)) + floating = inside_count < MIN_INSIDE * counts + floating[torch.argmax(counts)] = False + if floating is None or not bool(floating.any()): + return vertices.to(source_device), faces.to(source_device), report + kept_faces = faces[~floating[face_labels]] + used, remap = torch.unique(kept_faces, return_inverse=True) + report.update(removed_components=int(floating.sum()), removed_faces=int(counts[floating].sum())) + return vertices[used].to(source_device), remap.reshape(kept_faces.shape).to(source_device), report diff --git a/supratex/inference/native_observation.py b/supratex/inference/native_observation.py index 882ffab3b3772c0c0cadb3806490a6f900e83a04..5b5a6e632fb4a2f9e523fb9fe486d70b0860b84b 100644 --- a/supratex/inference/native_observation.py +++ b/supratex/inference/native_observation.py @@ -1,364 +1,413 @@ -"""Visible-surface material observations for Supra-Tex inference. - -This module owns the public pixel-first DDA visibility contract. It has no -dataset, scheduler, or training dependency. -""" -from __future__ import annotations - -import math -from typing import Any - -import numpy as np -import torch - -RESOLUTION = 1024 -LATENT_RESOLUTION = 64 - - -def build_occupancy_index(coords: torch.Tensor) -> tuple[np.ndarray, np.ndarray, np.ndarray]: - """Return an exact packed occupancy bitset and sorted row map.""" - values = coords.detach().cpu().numpy().astype(np.int64, copy=False) - if values.ndim != 2 or values.shape[1] != 3: - raise ValueError("surface coordinates must have shape [N,3]") - if values.size and (values.min() < 0 or values.max() >= RESOLUTION): - raise ValueError("surface coordinates are outside the material volume") - keys = (values[:, 0] * RESOLUTION + values[:, 1]) * RESOLUTION + values[:, 2] - order = np.argsort(keys) - ordered = keys[order] - if ordered.size > 1 and np.any(ordered[1:] == ordered[:-1]): - raise ValueError("surface coordinates contain duplicate cells") - bits = np.zeros((RESOLUTION ** 3 + 7) // 8, dtype=np.uint8) - if keys.size: - np.bitwise_or.at(bits, keys >> 3, np.left_shift(np.uint8(1), (keys & 7).astype(np.uint8))) - return bits, ordered, order.astype(np.int64, copy=False) - - -_DDA_BATCH_KERNEL: Any | None = None - -def _dda_batch_kernel() -> Any: - """Lazily compile the CPU DDA kernel so mesh-depth mode needs no Numba.""" - global _DDA_BATCH_KERNEL - if _DDA_BATCH_KERNEL is not None: - return _DDA_BATCH_KERNEL - try: - import numba - except ImportError as error: - raise RuntimeError("pixel_dda visibility requires numba") from error - - @numba.njit(parallel=True, cache=True) - def kernel( - ray_origin: np.ndarray, - ray_directions: np.ndarray, - occupancy_bits: np.ndarray, - grid_min: float, - grid_max: float, - resolution: int, - ) -> tuple[np.ndarray, np.ndarray]: - ray_count = ray_directions.shape[0] - owners = np.full((ray_count, 3), -1, dtype=np.int32) - reasons = np.full(ray_count, 3, dtype=np.uint8) - # Per-ray count of multi-axis (diagonal) steps. Written only at - # ray_index, so prange stays race-free; summed by the caller. - ties = np.zeros(ray_count, dtype=np.int64) - cell_size = (grid_max - grid_min) / resolution - slab_epsilon = 1e-14 - tie_epsilon = 1e-12 - - for ray_index in numba.prange(ray_count): - direction = ray_directions[ray_index] - valid_ray = True - for axis in range(3): - if not np.isfinite(ray_origin[axis]) or not np.isfinite(direction[axis]): - valid_ray = False - if not valid_ray: - continue - if ( - abs(direction[0]) < slab_epsilon - and abs(direction[1]) < slab_epsilon - and abs(direction[2]) < slab_epsilon - ): - continue - - t_entry = -np.inf - t_exit = np.inf - intersects = True - for axis in range(3): - component = direction[axis] - if abs(component) < slab_epsilon: - if ray_origin[axis] < grid_min or ray_origin[axis] > grid_max: - intersects = False - break - else: - t1 = (grid_min - ray_origin[axis]) / component - t2 = (grid_max - ray_origin[axis]) / component - if t1 > t2: - temporary = t1 - t1 = t2 - t2 = temporary - if t1 > t_entry: - t_entry = t1 - if t2 < t_exit: - t_exit = t2 - if not intersects or t_entry > t_exit or t_exit < 0.0: - reasons[ray_index] = 2 - continue - - t_start = max(t_entry, 0.0) - # Enter the half-open grid interior deterministically. This avoids - # floor(1024) for a ray entering exactly through the +AABB face. - entry_epsilon = max(1e-12, abs(t_start) * 1e-12) - point = ray_origin + (t_start + entry_epsilon) * direction - current = np.empty(3, dtype=np.int32) - step = np.zeros(3, dtype=np.int32) - t_delta = np.full(3, np.inf, dtype=np.float64) - t_max = np.full(3, np.inf, dtype=np.float64) - for axis in range(3): - cell = int(math.floor((point[axis] - grid_min) / cell_size)) - if cell < 0: - cell = 0 - elif cell >= resolution: - cell = resolution - 1 - current[axis] = cell - component = direction[axis] - if component > slab_epsilon: - step[axis] = 1 - boundary = grid_min + (cell + 1) * cell_size - t_max[axis] = (boundary - ray_origin[axis]) / component - t_delta[axis] = cell_size / component - elif component < -slab_epsilon: - step[axis] = -1 - boundary = grid_min + cell * cell_size - t_max[axis] = (boundary - ray_origin[axis]) / component - t_delta[axis] = -cell_size / component - - max_steps = 3 * resolution + 3 - for _ in range(max_steps): - key = ( - (np.int64(current[0]) * resolution + current[1]) * resolution - + current[2] - ) - byte_value = occupancy_bits[key >> 3] - if byte_value & (1 << (key & 7)): - owners[ray_index, 0] = current[0] - owners[ray_index, 1] = current[1] - owners[ray_index, 2] = current[2] - reasons[ray_index] = 0 - break - - next_t = min(t_max[0], t_max[1], t_max[2]) - if not np.isfinite(next_t) or next_t > t_exit + tie_epsilon: - reasons[ray_index] = 1 - break - tolerance = tie_epsilon * max(1.0, abs(next_t)) - in_bounds = True - # Step all axes crossed at the same ray parameter. This is the - # half-open tie rule for edge/corner hits and is deterministic. - crossed = 0 - for axis in range(3): - if abs(t_max[axis] - next_t) <= tolerance: - crossed += 1 - current[axis] += step[axis] - t_max[axis] += t_delta[axis] - if current[axis] < 0 or current[axis] >= resolution: - in_bounds = False - # A multi-axis step is a diagonal move, so it can skip cells that - # share only an edge or corner. The rule is frozen and fixture - # verified; count the events so the magnitude is known over a - # full corpus (~2e12 rays) instead of assumed negligible. - if crossed > 1: - ties[ray_index] += 1 - if not in_bounds: - reasons[ray_index] = 1 - break - - return owners, reasons, ties - - _DDA_BATCH_KERNEL = kernel - return _DDA_BATCH_KERNEL - - -def pixel_first_occupied(camera: dict[str, Any], occupancy: np.ndarray, trusted: torch.Tensor, depth_edge_rtol: float = 0.04) -> torch.Tensor: - """Map trusted camera pixels to their first occupied volume cell.""" - if "extrinsic" in camera: - return pixel_first_occupied_cv(camera, occupancy, trusted, depth_edge_rtol) - if trusted.shape != (RESOLUTION, RESOLUTION): - raise ValueError("trusted alpha mask must be 1024 by 1024") - if "camera_angle_x" not in camera or "distance" not in camera: - raise ValueError("camera must supply camera_angle_x and distance") - ys, xs = torch.nonzero(trusted, as_tuple=True) - x, y = xs.cpu().numpy(), ys.cpu().numpy() - focal = 16.0 / math.tan(float(camera["camera_angle_x"]) / 2.0) * RESOLUTION / 32.0 - directions = np.column_stack(((x + .5 - RESOLUTION / 2) / focal, -(y + .5 - RESOLUTION / 2) / focal, -np.ones_like(x))).astype(np.float64) - origin = np.array([0.0, 0.0, float(camera["distance"]) * float(camera.get("mesh_scale", 1.0))], dtype=np.float64) - owners, reasons, _ties = _dda_batch_kernel()(origin, directions, occupancy, -0.5, 0.5, RESOLUTION) - hit_mask = reasons == 0 - hit = owners[hit_mask] - if not hit.size: - raise RuntimeError("no trusted pixels intersect the material volume") - # Preserve the production first-hit-depth admission gate. It removes - # pixels at local internal-occlusion jumps before their owner cells become - # observations; applying it after unique-owner reduction would be wrong. - import torch.nn.functional as F - depth = np.full((RESOLUTION, RESOLUTION), np.nan, dtype=np.float32) - hit_depth = origin[2] - ((hit[:, 2].astype(np.float64) + .5) / RESOLUTION - .5) - depth[y[hit_mask], x[hit_mask]] = hit_depth.astype(np.float32) - values = torch.from_numpy(depth) - valid = torch.from_numpy(~np.isnan(depth)) - masked = torch.where(valid, values, torch.nan) - local_max = F.max_pool2d(masked.nan_to_num(-torch.inf)[None, None], 3, 1, 1)[0, 0] - local_min = -F.max_pool2d((-masked).nan_to_num(-torch.inf)[None, None], 3, 1, 1)[0, 0] - relative = ((local_max - local_min) / masked).nan_to_num(nan=0., posinf=0., neginf=0.) - retained = hit_mask.copy() - retained[hit_mask] &= ~((relative.numpy() > depth_edge_rtol)[y[hit_mask], x[hit_mask]]) - hit = owners[retained] - if not hit.size: - raise RuntimeError("depth-edge admission removed every first-hit observation") - keys = np.unique((hit[:, 0].astype(np.int64) * RESOLUTION + hit[:, 1]) * RESOLUTION + hit[:, 2]) - return torch.from_numpy(np.stack((keys // (RESOLUTION * RESOLUTION), (keys // RESOLUTION) % RESOLUTION, keys % RESOLUTION), axis=1).astype(np.int32)) - - -def camera_rays(camera: dict[str, Any], ys: np.ndarray, xs: np.ndarray) -> tuple[np.ndarray, np.ndarray, np.ndarray]: - """Build pixel-center rays from a complete OpenCV world-to-camera pose.""" - extrinsic = np.asarray(camera["extrinsic"], dtype=np.float64) - intrinsic = np.asarray(camera.get("projection_intrinsic", camera.get("intrinsic")), dtype=np.float64) - if extrinsic.shape != (4, 4) or intrinsic.shape != (3, 3): - raise ValueError("DDA camera requires 4x4 extrinsic and normalized 3x3 intrinsic") - if not np.isfinite(extrinsic).all() or not np.isfinite(intrinsic).all(): - raise ValueError("DDA camera matrices must be finite") - rotation = extrinsic[:3, :3] - if (not np.allclose(rotation @ rotation.T, np.eye(3), atol=1e-6, rtol=0) - or not np.isclose(np.linalg.det(rotation), 1, atol=1e-6) - or not np.allclose(extrinsic[3], [0, 0, 0, 1], atol=1e-8, rtol=0)): - raise ValueError("DDA camera must contain a rigid world-to-camera transform") - if (intrinsic[0, 0] <= 0 or intrinsic[1, 1] <= 0 - or not np.allclose(intrinsic[[0, 1, 2, 2], [1, 0, 0, 1]], 0, atol=1e-10) - or intrinsic[2, 2] != 1): - raise ValueError("DDA requires positive focal lengths and an unskewed normalized intrinsic") - eye = np.ascontiguousarray(-rotation.T @ extrinsic[:3, 3]) - x = (xs.astype(np.float64) + .5 - intrinsic[0, 2] * RESOLUTION) / (intrinsic[0, 0] * RESOLUTION) - y = (ys.astype(np.float64) + .5 - intrinsic[1, 2] * RESOLUTION) / (intrinsic[1, 1] * RESOLUTION) - directions = np.ascontiguousarray(np.stack((x, y, np.ones_like(x)), axis=1) @ rotation) - return eye, directions, rotation - - -def pixel_hits_cv(camera: dict[str, Any], occupancy: np.ndarray, mask: np.ndarray) -> np.ndarray: - """Return the raw first-hit silhouette, without trusted-mask depth filtering.""" - if mask.shape != (RESOLUTION, RESOLUTION) or mask.dtype != np.bool_: - raise ValueError("DDA pixel mask must be a boolean observation-resolution image") - ys, xs = np.nonzero(mask) - eye, directions, _ = camera_rays(camera, ys, xs) - _, reasons, _ = _dda_batch_kernel()(eye, directions, occupancy, -.5, .5, RESOLUTION) - hits = np.zeros_like(mask) - hits[ys[reasons == 0], xs[reasons == 0]] = True - return hits - - -def pixel_first_occupied_cv(camera: dict[str, Any], occupancy: np.ndarray, trusted: torch.Tensor, - depth_edge_rtol: float = .04) -> torch.Tensor: - """Apply the production first-hit and depth-edge rules to a full CV camera.""" - import torch.nn.functional as F - - if trusted.shape != (RESOLUTION, RESOLUTION) or trusted.dtype != torch.bool: - raise ValueError("trusted alpha must be a boolean observation-resolution image") - if not math.isfinite(depth_edge_rtol) or depth_edge_rtol < 0: - raise ValueError("depth-edge tolerance must be finite and nonnegative") - ys, xs = np.nonzero(trusted.detach().cpu().numpy()) - if not ys.size: - raise RuntimeError("trusted alpha gate leaves no camera rays") - eye, directions, rotation = camera_rays(camera, ys, xs) - owners, reasons, _ = _dda_batch_kernel()(eye, directions, occupancy, -.5, .5, RESOLUTION) - hits = reasons == 0 - if not hits.any(): - raise RuntimeError("no trusted pixels intersect the material volume") - centers = (owners[hits].astype(np.float64) + .5) / RESOLUTION - .5 - hit_depth = ((centers - eye) @ rotation.T)[:, 2] - if not np.isfinite(hit_depth).all() or not (hit_depth > 0).all(): - raise RuntimeError("first-hit owners must lie in front of the camera") - depth = np.full((RESOLUTION, RESOLUTION), np.nan, dtype=np.float32) - depth[ys[hits], xs[hits]] = hit_depth.astype(np.float32) - masked = torch.from_numpy(depth) - local_max = F.max_pool2d(masked.nan_to_num(-torch.inf)[None, None], 3, 1, 1)[0, 0] - local_min = -F.max_pool2d((-masked).nan_to_num(-torch.inf)[None, None], 3, 1, 1)[0, 0] - relative = ((local_max - local_min) / masked).nan_to_num(nan=0., posinf=0., neginf=0.) - retained = hits & ~(relative.numpy() > depth_edge_rtol)[ys, xs] - kept = owners[retained].astype(np.int64, copy=False) - if not kept.size: - raise RuntimeError("depth-edge admission removed every first-hit observation") - keys = np.unique((kept[:, 0] * RESOLUTION + kept[:, 1]) * RESOLUTION + kept[:, 2]) - return torch.from_numpy(np.stack((keys // (RESOLUTION * RESOLUTION), - keys // RESOLUTION % RESOLUTION, - keys % RESOLUTION), axis=1).astype(np.int32)) - - -def owner_rows(coords: torch.Tensor, owners: torch.Tensor, ordered: np.ndarray, order: np.ndarray) -> torch.Tensor: - values = owners.cpu().numpy().astype(np.int64, copy=False) - keys = (values[:, 0] * RESOLUTION + values[:, 1]) * RESOLUTION + values[:, 2] - positions = np.searchsorted(ordered, keys) - if positions.size and (np.any(positions >= ordered.size) or np.any(ordered[positions] != keys)): - raise RuntimeError("visible owner is absent from the surface support") - result = torch.zeros(coords.shape[0], dtype=torch.bool, device=coords.device) - if positions.size: - result[torch.as_tensor(order[positions], device=coords.device)] = True - return result - - -def boundary_distance_to_unobserved(surface_counts: torch.Tensor, observed_counts: torch.Tensor) -> torch.Tensor: - """Exact sparse six-neighbour distance used by the observation message.""" - active = (surface_counts > 0).reshape(LATENT_RESOLUTION, LATENT_RESOLUTION, LATENT_RESOLUTION) - observed = (observed_counts > 0).reshape_as(active) - distance = torch.full(active.shape, -1, dtype=torch.int16, device=active.device) - frontier = active & ~observed - distance[frontier] = 0 - step = 0 - while bool(frontier.any()): - expanded = torch.zeros_like(frontier) - expanded[1:] |= frontier[:-1]; expanded[:-1] |= frontier[1:] - expanded[:, 1:] |= frontier[:, :-1]; expanded[:, :-1] |= frontier[:, 1:] - expanded[:, :, 1:] |= frontier[:, :, :-1]; expanded[:, :, :-1] |= frontier[:, :, 1:] - frontier = expanded & active & (distance < 0) - if not bool(frontier.any()): - break - step += 1 - distance[frontier] = step - return distance.reshape(-1) - - -def encode_partial(encoder: Any, coords: torch.Tensor, visible: torch.Tensor, rgb: torch.Tensor, mra: torch.Tensor, shape_coords: torch.Tensor, normalization: dict[str, Any], sparse: Any, device: str) -> Any: - """Encode visible PBR rows into the frozen 34-channel observation message.""" - selected = coords[visible] - features = torch.cat((rgb[visible], mra[visible]), dim=1).clamp(0, 1).float() * 2.0 - 1.0 - encoded = encoder(sparse.SparseTensor(features, torch.cat((torch.zeros_like(selected[:, :1]), selected), dim=1).int()), sample_posterior=False) - raw = encoded.feats.float() - if raw.ndim != 2 or raw.shape[1] != 32: - raise RuntimeError("material encoder must emit 32 channels") - blocks = coords // (RESOLUTION // LATENT_RESOLUTION) - keys = (blocks[:, 0] * LATENT_RESOLUTION + blocks[:, 1]) * LATENT_RESOLUTION + blocks[:, 2] - all_counts = torch.bincount(keys, minlength=LATENT_RESOLUTION ** 3) - observed_counts = torch.bincount(keys[visible], minlength=LATENT_RESOLUTION ** 3) - xyz = encoded.coords[:, 1:].long() - latent_keys = (xyz[:, 0] * LATENT_RESOLUTION + xyz[:, 1]) * LATENT_RESOLUTION + xyz[:, 2] - # The bounded grid already has exact occupancy counts. Enumerate occupied - # keys instead of sorting millions of three-column surface coordinates. - # Linear-key order is the same lexicographic order as unique(..., dim=0). - expected_observed = torch.nonzero(observed_counts, as_tuple=False).flatten() - if (bool(((xyz < 0) | (xyz >= LATENT_RESOLUTION)).any()) - or not torch.equal(torch.sort(latent_keys).values, expected_observed)): - raise RuntimeError("encoded observation support differs from floor(visible_surface/16)") - coverage = (observed_counts[latent_keys].float() / all_counts[latent_keys].float()).unsqueeze(1) - expected_keys = torch.nonzero(all_counts, as_tuple=False).flatten() - shape_xyz = shape_coords.long() - shape_keys = (shape_xyz[:, 0] * LATENT_RESOLUTION + shape_xyz[:, 1]) * LATENT_RESOLUTION + shape_xyz[:, 2] - sorted_shape = torch.sort(shape_keys).values - if expected_keys.shape[0] != shape_keys.shape[0] or not torch.equal(expected_keys, sorted_shape): - raise RuntimeError("pooled material support does not exactly match the shape latent support") - positions = torch.searchsorted(sorted_shape, latent_keys) - if bool((positions >= sorted_shape.shape[0]).any()) or not torch.equal(sorted_shape[positions], latent_keys): - raise RuntimeError("encoded observation support is not a subset of full material support") - boundary = boundary_distance_to_unobserved(all_counts, observed_counts)[latent_keys].to(dtype=coverage.dtype).unsqueeze(1) - if bool((boundary == 0).any()) or bool((boundary < -1).any()): - raise RuntimeError("invalid observed material boundary distance") - mean = torch.as_tensor(normalization["mean"], dtype=raw.dtype, device=device).reshape(1, 32) - std = torch.as_tensor(normalization["std"], dtype=raw.dtype, device=device).reshape(1, 32) - if mean.shape[1] != 32 or std.shape[1] != 32 or not bool(torch.isfinite(std).all()) or not bool((std > 0).all()): - raise ValueError("partial observation normalization must provide 32 finite positive standard deviations") - return sparse.SparseTensor(torch.cat(((raw - mean) / std, coverage, boundary), dim=1), encoded.coords.int()) +"""Visible-surface material observations for Supra-Tex inference. + +This module owns the public pixel-first DDA visibility contract. It has no +dataset, scheduler, or training dependency. +""" +from __future__ import annotations + +import math +from typing import Any + +import numpy as np +import torch + +RESOLUTION = 1024 +LATENT_RESOLUTION = 64 + + +def build_occupancy_index(coords: torch.Tensor) -> tuple[np.ndarray, np.ndarray, np.ndarray]: + """Return an exact packed occupancy bitset and sorted row map.""" + values = coords.detach().cpu().numpy().astype(np.int64, copy=False) + if values.ndim != 2 or values.shape[1] != 3: + raise ValueError("surface coordinates must have shape [N,3]") + if values.size and (values.min() < 0 or values.max() >= RESOLUTION): + raise ValueError("surface coordinates are outside the material volume") + keys = (values[:, 0] * RESOLUTION + values[:, 1]) * RESOLUTION + values[:, 2] + order = np.argsort(keys) + ordered = keys[order] + if ordered.size > 1 and np.any(ordered[1:] == ordered[:-1]): + raise ValueError("surface coordinates contain duplicate cells") + bits = np.zeros((RESOLUTION ** 3 + 7) // 8, dtype=np.uint8) + if keys.size: + np.bitwise_or.at(bits, keys >> 3, np.left_shift(np.uint8(1), (keys & 7).astype(np.uint8))) + return bits, ordered, order.astype(np.int64, copy=False) + + +_DDA_BATCH_KERNEL: Any | None = None + +def _dda_batch_kernel() -> Any: + """Lazily compile the CPU DDA kernel so mesh-depth mode needs no Numba.""" + global _DDA_BATCH_KERNEL + if _DDA_BATCH_KERNEL is not None: + return _DDA_BATCH_KERNEL + try: + import numba + except ImportError as error: + raise RuntimeError("pixel_dda visibility requires numba") from error + + @numba.njit(parallel=True, cache=True) + def kernel( + ray_origin: np.ndarray, + ray_directions: np.ndarray, + occupancy_bits: np.ndarray, + grid_min: float, + grid_max: float, + resolution: int, + ) -> tuple[np.ndarray, np.ndarray]: + ray_count = ray_directions.shape[0] + owners = np.full((ray_count, 3), -1, dtype=np.int32) + reasons = np.full(ray_count, 3, dtype=np.uint8) + # Per-ray count of multi-axis (diagonal) steps. Written only at + # ray_index, so prange stays race-free; summed by the caller. + ties = np.zeros(ray_count, dtype=np.int64) + cell_size = (grid_max - grid_min) / resolution + slab_epsilon = 1e-14 + tie_epsilon = 1e-12 + + for ray_index in numba.prange(ray_count): + direction = ray_directions[ray_index] + valid_ray = True + for axis in range(3): + if not np.isfinite(ray_origin[axis]) or not np.isfinite(direction[axis]): + valid_ray = False + if not valid_ray: + continue + if ( + abs(direction[0]) < slab_epsilon + and abs(direction[1]) < slab_epsilon + and abs(direction[2]) < slab_epsilon + ): + continue + + t_entry = -np.inf + t_exit = np.inf + intersects = True + for axis in range(3): + component = direction[axis] + if abs(component) < slab_epsilon: + if ray_origin[axis] < grid_min or ray_origin[axis] > grid_max: + intersects = False + break + else: + t1 = (grid_min - ray_origin[axis]) / component + t2 = (grid_max - ray_origin[axis]) / component + if t1 > t2: + temporary = t1 + t1 = t2 + t2 = temporary + if t1 > t_entry: + t_entry = t1 + if t2 < t_exit: + t_exit = t2 + if not intersects or t_entry > t_exit or t_exit < 0.0: + reasons[ray_index] = 2 + continue + + t_start = max(t_entry, 0.0) + # Enter the half-open grid interior deterministically. This avoids + # floor(1024) for a ray entering exactly through the +AABB face. + entry_epsilon = max(1e-12, abs(t_start) * 1e-12) + point = ray_origin + (t_start + entry_epsilon) * direction + current = np.empty(3, dtype=np.int32) + step = np.zeros(3, dtype=np.int32) + t_delta = np.full(3, np.inf, dtype=np.float64) + t_max = np.full(3, np.inf, dtype=np.float64) + for axis in range(3): + cell = int(math.floor((point[axis] - grid_min) / cell_size)) + if cell < 0: + cell = 0 + elif cell >= resolution: + cell = resolution - 1 + current[axis] = cell + component = direction[axis] + if component > slab_epsilon: + step[axis] = 1 + boundary = grid_min + (cell + 1) * cell_size + t_max[axis] = (boundary - ray_origin[axis]) / component + t_delta[axis] = cell_size / component + elif component < -slab_epsilon: + step[axis] = -1 + boundary = grid_min + cell * cell_size + t_max[axis] = (boundary - ray_origin[axis]) / component + t_delta[axis] = -cell_size / component + + max_steps = 3 * resolution + 3 + for _ in range(max_steps): + key = ( + (np.int64(current[0]) * resolution + current[1]) * resolution + + current[2] + ) + byte_value = occupancy_bits[key >> 3] + if byte_value & (1 << (key & 7)): + owners[ray_index, 0] = current[0] + owners[ray_index, 1] = current[1] + owners[ray_index, 2] = current[2] + reasons[ray_index] = 0 + break + + next_t = min(t_max[0], t_max[1], t_max[2]) + if not np.isfinite(next_t) or next_t > t_exit + tie_epsilon: + reasons[ray_index] = 1 + break + tolerance = tie_epsilon * max(1.0, abs(next_t)) + in_bounds = True + # Step all axes crossed at the same ray parameter. This is the + # half-open tie rule for edge/corner hits and is deterministic. + crossed = 0 + for axis in range(3): + if abs(t_max[axis] - next_t) <= tolerance: + crossed += 1 + current[axis] += step[axis] + t_max[axis] += t_delta[axis] + if current[axis] < 0 or current[axis] >= resolution: + in_bounds = False + # A multi-axis step is a diagonal move, so it can skip cells that + # share only an edge or corner. The rule is frozen and fixture + # verified; count the events so the magnitude is known over a + # full corpus (~2e12 rays) instead of assumed negligible. + if crossed > 1: + ties[ray_index] += 1 + if not in_bounds: + reasons[ray_index] = 1 + break + + return owners, reasons, ties + + _DDA_BATCH_KERNEL = kernel + return _DDA_BATCH_KERNEL + + +# Optional admission rules, off by default. On well-aligned demo inputs they +# removed 3-10% of trusted pixels without a visible change, so the released +# behavior keeps only the depth-discontinuity rule. +SILHOUETTE_MARGIN = 0 +GRAZING_COS = 0.0 + + +def _admitted_pixels(depth: np.ndarray, missed: np.ndarray, focal: tuple[float, float], depth_edge_rtol: float, + silhouette_margin: int = SILHOUETTE_MARGIN, grazing_cos: float = GRAZING_COS, + stats: dict | None = None) -> np.ndarray: + """Return pixels whose first-hit depth is a reliable observation. + + ``depth`` is the first-hit view depth (NaN without a trusted hit) and + ``missed`` marks trusted rays that left the volume, i.e. the image + foreground outside the geometry silhouette. A local depth discontinuity + always rejects a pixel. Two optional rules can also reject pixels near the + geometry silhouette (``silhouette_margin`` pixels), where image/geometry + misalignment lands colors on the wrong surface, and at grazing incidence + (cosine below ``grazing_cos``, from the screen-space depth slope), where one + pixel stretches over a large surface area. + """ + import torch.nn.functional as F + + masked = torch.from_numpy(depth) + local_max = F.max_pool2d(masked.nan_to_num(-torch.inf)[None, None], 3, 1, 1)[0, 0] + local_min = -F.max_pool2d((-masked).nan_to_num(-torch.inf)[None, None], 3, 1, 1)[0, 0] + relative = ((local_max - local_min) / masked).nan_to_num(nan=0., posinf=0., neginf=0.) + edge = relative.numpy() > depth_edge_rtol + near_silhouette = np.zeros_like(edge) + if silhouette_margin > 0 and missed.any(): + size = 2 * silhouette_margin + 1 + near_silhouette = F.max_pool2d(torch.from_numpy(missed).float()[None, None], size, 1, silhouette_margin)[0, 0].numpy() > 0 + grazing = np.zeros_like(edge) + if grazing_cos > 0: + # Central differences in pixels; a missing neighbour leaves NaN and is + # not rejected here. Dividing by the pixel footprint (depth / focal) + # turns the depth change into the surface slope along each image axis. + dz_dx = np.full_like(depth, np.nan) + dz_dy = np.full_like(depth, np.nan) + dz_dx[:, 1:-1] = (depth[:, 2:] - depth[:, :-2]) * .5 + dz_dy[1:-1, :] = (depth[2:, :] - depth[:-2, :]) * .5 + with np.errstate(invalid="ignore", divide="ignore"): + slope2 = (dz_dx * np.float32(focal[0]) / depth) ** 2 + (dz_dy * np.float32(focal[1]) / depth) ** 2 + grazing = slope2 > np.float32(1 / grazing_cos ** 2 - 1) + hit = ~np.isnan(depth) + admitted = hit & ~edge & ~near_silhouette & ~grazing + if stats is not None: + stats.update(trusted_rays=int(hit.sum() + missed.sum()), first_hits=int(hit.sum()), + rejected_depth_edge=int((hit & edge).sum()), + rejected_silhouette=int((hit & ~edge & near_silhouette).sum()), + rejected_grazing=int((hit & ~edge & ~near_silhouette & grazing).sum()), + admitted=int(admitted.sum())) + return admitted + + +def pixel_first_occupied(camera: dict[str, Any], occupancy: np.ndarray, trusted: torch.Tensor, depth_edge_rtol: float = 0.04, + stats: dict | None = None) -> torch.Tensor: + """Map trusted camera pixels to their first occupied volume cell.""" + if "extrinsic" in camera: + return pixel_first_occupied_cv(camera, occupancy, trusted, depth_edge_rtol, stats=stats) + if trusted.shape != (RESOLUTION, RESOLUTION): + raise ValueError("trusted alpha mask must be 1024 by 1024") + if "camera_angle_x" not in camera or "distance" not in camera: + raise ValueError("camera must supply camera_angle_x and distance") + ys, xs = torch.nonzero(trusted, as_tuple=True) + x, y = xs.cpu().numpy(), ys.cpu().numpy() + focal = 16.0 / math.tan(float(camera["camera_angle_x"]) / 2.0) * RESOLUTION / 32.0 + directions = np.column_stack(((x + .5 - RESOLUTION / 2) / focal, -(y + .5 - RESOLUTION / 2) / focal, -np.ones_like(x))).astype(np.float64) + origin = np.array([0.0, 0.0, float(camera["distance"]) * float(camera.get("mesh_scale", 1.0))], dtype=np.float64) + owners, reasons, _ties = _dda_batch_kernel()(origin, directions, occupancy, -0.5, 0.5, RESOLUTION) + hit_mask = reasons == 0 + hit = owners[hit_mask] + if not hit.size: + raise RuntimeError("no trusted pixels intersect the material volume") + # Admission runs per pixel, before owner cells are deduplicated; applying + # it after unique-owner reduction would be wrong. + depth = np.full((RESOLUTION, RESOLUTION), np.nan, dtype=np.float32) + hit_depth = origin[2] - ((hit[:, 2].astype(np.float64) + .5) / RESOLUTION - .5) + depth[y[hit_mask], x[hit_mask]] = hit_depth.astype(np.float32) + missed = np.zeros((RESOLUTION, RESOLUTION), dtype=bool) + missed[y[~hit_mask], x[~hit_mask]] = True + admitted = _admitted_pixels(depth, missed, (focal, focal), depth_edge_rtol, stats=stats) + retained = hit_mask.copy() + retained[hit_mask] &= admitted[y[hit_mask], x[hit_mask]] + hit = owners[retained] + if not hit.size: + raise RuntimeError("observation admission removed every first-hit observation") + keys = np.unique((hit[:, 0].astype(np.int64) * RESOLUTION + hit[:, 1]) * RESOLUTION + hit[:, 2]) + return torch.from_numpy(np.stack((keys // (RESOLUTION * RESOLUTION), (keys // RESOLUTION) % RESOLUTION, keys % RESOLUTION), axis=1).astype(np.int32)) + + +def camera_rays(camera: dict[str, Any], ys: np.ndarray, xs: np.ndarray) -> tuple[np.ndarray, np.ndarray, np.ndarray]: + """Build pixel-center rays from a complete OpenCV world-to-camera pose.""" + extrinsic = np.asarray(camera["extrinsic"], dtype=np.float64) + intrinsic = np.asarray(camera.get("projection_intrinsic", camera.get("intrinsic")), dtype=np.float64) + if extrinsic.shape != (4, 4) or intrinsic.shape != (3, 3): + raise ValueError("DDA camera requires 4x4 extrinsic and normalized 3x3 intrinsic") + if not np.isfinite(extrinsic).all() or not np.isfinite(intrinsic).all(): + raise ValueError("DDA camera matrices must be finite") + rotation = extrinsic[:3, :3] + if (not np.allclose(rotation @ rotation.T, np.eye(3), atol=1e-6, rtol=0) + or not np.isclose(np.linalg.det(rotation), 1, atol=1e-6) + or not np.allclose(extrinsic[3], [0, 0, 0, 1], atol=1e-8, rtol=0)): + raise ValueError("DDA camera must contain a rigid world-to-camera transform") + if (intrinsic[0, 0] <= 0 or intrinsic[1, 1] <= 0 + or not np.allclose(intrinsic[[0, 1, 2, 2], [1, 0, 0, 1]], 0, atol=1e-10) + or intrinsic[2, 2] != 1): + raise ValueError("DDA requires positive focal lengths and an unskewed normalized intrinsic") + eye = np.ascontiguousarray(-rotation.T @ extrinsic[:3, 3]) + x = (xs.astype(np.float64) + .5 - intrinsic[0, 2] * RESOLUTION) / (intrinsic[0, 0] * RESOLUTION) + y = (ys.astype(np.float64) + .5 - intrinsic[1, 2] * RESOLUTION) / (intrinsic[1, 1] * RESOLUTION) + directions = np.ascontiguousarray(np.stack((x, y, np.ones_like(x)), axis=1) @ rotation) + return eye, directions, rotation + + +def pixel_hits_cv(camera: dict[str, Any], occupancy: np.ndarray, mask: np.ndarray) -> np.ndarray: + """Return the raw first-hit silhouette, without trusted-mask depth filtering.""" + if mask.shape != (RESOLUTION, RESOLUTION) or mask.dtype != np.bool_: + raise ValueError("DDA pixel mask must be a boolean observation-resolution image") + ys, xs = np.nonzero(mask) + eye, directions, _ = camera_rays(camera, ys, xs) + _, reasons, _ = _dda_batch_kernel()(eye, directions, occupancy, -.5, .5, RESOLUTION) + hits = np.zeros_like(mask) + hits[ys[reasons == 0], xs[reasons == 0]] = True + return hits + + +def pixel_first_occupied_cv(camera: dict[str, Any], occupancy: np.ndarray, trusted: torch.Tensor, + depth_edge_rtol: float = .04, stats: dict | None = None) -> torch.Tensor: + """Apply the production first-hit admission rules to a full CV camera.""" + if trusted.shape != (RESOLUTION, RESOLUTION) or trusted.dtype != torch.bool: + raise ValueError("trusted alpha must be a boolean observation-resolution image") + if not math.isfinite(depth_edge_rtol) or depth_edge_rtol < 0: + raise ValueError("depth-edge tolerance must be finite and nonnegative") + ys, xs = np.nonzero(trusted.detach().cpu().numpy()) + if not ys.size: + raise RuntimeError("trusted alpha gate leaves no camera rays") + eye, directions, rotation = camera_rays(camera, ys, xs) + owners, reasons, _ = _dda_batch_kernel()(eye, directions, occupancy, -.5, .5, RESOLUTION) + hits = reasons == 0 + if not hits.any(): + raise RuntimeError("no trusted pixels intersect the material volume") + centers = (owners[hits].astype(np.float64) + .5) / RESOLUTION - .5 + hit_depth = ((centers - eye) @ rotation.T)[:, 2] + if not np.isfinite(hit_depth).all() or not (hit_depth > 0).all(): + raise RuntimeError("first-hit owners must lie in front of the camera") + depth = np.full((RESOLUTION, RESOLUTION), np.nan, dtype=np.float32) + depth[ys[hits], xs[hits]] = hit_depth.astype(np.float32) + missed = np.zeros((RESOLUTION, RESOLUTION), dtype=bool) + missed[ys[~hits], xs[~hits]] = True + intrinsic = np.asarray(camera.get("projection_intrinsic", camera.get("intrinsic")), dtype=np.float64) + focal = (float(intrinsic[0, 0]) * RESOLUTION, float(intrinsic[1, 1]) * RESOLUTION) + retained = hits & _admitted_pixels(depth, missed, focal, depth_edge_rtol, stats=stats)[ys, xs] + kept = owners[retained].astype(np.int64, copy=False) + if not kept.size: + raise RuntimeError("observation admission removed every first-hit observation") + keys = np.unique((kept[:, 0] * RESOLUTION + kept[:, 1]) * RESOLUTION + kept[:, 2]) + return torch.from_numpy(np.stack((keys // (RESOLUTION * RESOLUTION), + keys // RESOLUTION % RESOLUTION, + keys % RESOLUTION), axis=1).astype(np.int32)) + + +def owner_rows(coords: torch.Tensor, owners: torch.Tensor, ordered: np.ndarray, order: np.ndarray) -> torch.Tensor: + values = owners.cpu().numpy().astype(np.int64, copy=False) + keys = (values[:, 0] * RESOLUTION + values[:, 1]) * RESOLUTION + values[:, 2] + positions = np.searchsorted(ordered, keys) + if positions.size and (np.any(positions >= ordered.size) or np.any(ordered[positions] != keys)): + raise RuntimeError("visible owner is absent from the surface support") + result = torch.zeros(coords.shape[0], dtype=torch.bool, device=coords.device) + if positions.size: + result[torch.as_tensor(order[positions], device=coords.device)] = True + return result + + +def boundary_distance_to_unobserved(surface_counts: torch.Tensor, observed_counts: torch.Tensor) -> torch.Tensor: + """Exact sparse six-neighbour distance used by the observation message.""" + active = (surface_counts > 0).reshape(LATENT_RESOLUTION, LATENT_RESOLUTION, LATENT_RESOLUTION) + observed = (observed_counts > 0).reshape_as(active) + distance = torch.full(active.shape, -1, dtype=torch.int16, device=active.device) + frontier = active & ~observed + distance[frontier] = 0 + step = 0 + while bool(frontier.any()): + expanded = torch.zeros_like(frontier) + expanded[1:] |= frontier[:-1]; expanded[:-1] |= frontier[1:] + expanded[:, 1:] |= frontier[:, :-1]; expanded[:, :-1] |= frontier[:, 1:] + expanded[:, :, 1:] |= frontier[:, :, :-1]; expanded[:, :, :-1] |= frontier[:, :, 1:] + frontier = expanded & active & (distance < 0) + if not bool(frontier.any()): + break + step += 1 + distance[frontier] = step + return distance.reshape(-1) + + +def encode_partial(encoder: Any, coords: torch.Tensor, visible: torch.Tensor, rgb: torch.Tensor, mra: torch.Tensor, shape_coords: torch.Tensor, normalization: dict[str, Any], sparse: Any, device: str) -> Any: + """Encode visible PBR rows into the frozen 34-channel observation message.""" + selected = coords[visible] + features = torch.cat((rgb[visible], mra[visible]), dim=1).clamp(0, 1).float() * 2.0 - 1.0 + encoded = encoder(sparse.SparseTensor(features, torch.cat((torch.zeros_like(selected[:, :1]), selected), dim=1).int()), sample_posterior=False) + raw = encoded.feats.float() + if raw.ndim != 2 or raw.shape[1] != 32: + raise RuntimeError("material encoder must emit 32 channels") + blocks = coords // (RESOLUTION // LATENT_RESOLUTION) + keys = (blocks[:, 0] * LATENT_RESOLUTION + blocks[:, 1]) * LATENT_RESOLUTION + blocks[:, 2] + all_counts = torch.bincount(keys, minlength=LATENT_RESOLUTION ** 3) + observed_counts = torch.bincount(keys[visible], minlength=LATENT_RESOLUTION ** 3) + xyz = encoded.coords[:, 1:].long() + latent_keys = (xyz[:, 0] * LATENT_RESOLUTION + xyz[:, 1]) * LATENT_RESOLUTION + xyz[:, 2] + # The bounded grid already has exact occupancy counts. Enumerate occupied + # keys instead of sorting millions of three-column surface coordinates. + # Linear-key order is the same lexicographic order as unique(..., dim=0). + expected_observed = torch.nonzero(observed_counts, as_tuple=False).flatten() + if (bool(((xyz < 0) | (xyz >= LATENT_RESOLUTION)).any()) + or not torch.equal(torch.sort(latent_keys).values, expected_observed)): + raise RuntimeError("encoded observation support differs from floor(visible_surface/16)") + coverage = (observed_counts[latent_keys].float() / all_counts[latent_keys].float()).unsqueeze(1) + expected_keys = torch.nonzero(all_counts, as_tuple=False).flatten() + shape_xyz = shape_coords.long() + shape_keys = (shape_xyz[:, 0] * LATENT_RESOLUTION + shape_xyz[:, 1]) * LATENT_RESOLUTION + shape_xyz[:, 2] + sorted_shape = torch.sort(shape_keys).values + if expected_keys.shape[0] != shape_keys.shape[0] or not torch.equal(expected_keys, sorted_shape): + raise RuntimeError("pooled material support does not exactly match the shape latent support") + positions = torch.searchsorted(sorted_shape, latent_keys) + if bool((positions >= sorted_shape.shape[0]).any()) or not torch.equal(sorted_shape[positions], latent_keys): + raise RuntimeError("encoded observation support is not a subset of full material support") + boundary = boundary_distance_to_unobserved(all_counts, observed_counts)[latent_keys].to(dtype=coverage.dtype).unsqueeze(1) + if bool((boundary == 0).any()) or bool((boundary < -1).any()): + raise RuntimeError("invalid observed material boundary distance") + mean = torch.as_tensor(normalization["mean"], dtype=raw.dtype, device=device).reshape(1, 32) + std = torch.as_tensor(normalization["std"], dtype=raw.dtype, device=device).reshape(1, 32) + if mean.shape[1] != 32 or std.shape[1] != 32 or not bool(torch.isfinite(std).all()) or not bool((std > 0).all()): + raise ValueError("partial observation normalization must provide 32 finite positive standard deviations") + return sparse.SparseTensor(torch.cat(((raw - mean) / std, coverage, boundary), dim=1), encoded.coords.int()) diff --git a/supratex/inference/pipeline.py b/supratex/inference/pipeline.py index 8098ff1e34f719b699112b551bec576aa079faf2..84256b200248054d156be93964ea6e3e386b7177 100644 --- a/supratex/inference/pipeline.py +++ b/supratex/inference/pipeline.py @@ -1,552 +1,700 @@ -"""Observation-guided PBR completion and GLB export. - -The public boundary begins with a mesh/camera/RGBA observation plus albedo and -MRA maps. Geometry prediction and 2D map prediction intentionally remain -outside this inference-only package. -""" -from __future__ import annotations - -import json -import math -import os -from contextlib import contextmanager -from pathlib import Path -from typing import Any, Mapping - -from hf_space.timing import span - -from . import native_observation -from .checkpoint import load_completion_state_dict, validate_completion_checkpoint - -RESOLUTION = native_observation.RESOLUTION -LAYOUT = {"base_color": slice(0, 3), "metallic": slice(3, 4), "roughness": slice(4, 5), "alpha": slice(5, 6)} - - -def _file(path: str | Path, label: str) -> Path: - value = Path(path) - if not value.is_file(): - raise FileNotFoundError(f"missing {label}: {value}") - return value - - -def _normalization(value: Any, label: str) -> dict[str, list[float]]: - if not isinstance(value, Mapping): - raise ValueError(f"{label} normalization must be an object") - result: dict[str, list[float]] = {} - for key in ("mean", "std"): - raw = value.get(key) - if not isinstance(raw, list) or len(raw) != 32: - raise ValueError(f"{label} normalization requires 32-value {key}") - result[key] = [float(item) for item in raw] - if not all(math.isfinite(item) for item in result["mean"] + result["std"]) or any(item <= 0 for item in result["std"]): - raise ValueError(f"{label} normalization standard deviations must be positive") - return result - - -def load_contract(model_config: str | Path, checkpoint: str | Path) -> tuple[dict[str, Any], dict[str, Any], Path]: - """Load the public model ABI without inspecting or hashing weight bytes.""" - config_path, checkpoint_path = _file(model_config, "model config"), _file(checkpoint, "checkpoint") - try: - config = json.loads(config_path.read_text(encoding="utf-8")) - dataset = config["dataset"]["args"] - model = config["models"]["denoiser"] - args = model["args"] - except (OSError, KeyError, TypeError, json.JSONDecodeError) as error: - raise ValueError("model config does not satisfy the Supra-Tex inference ABI") from error - if config.get("format") == "supratex_completion" and config.get("schema_version") != 1: - raise ValueError("unsupported Supra-Tex completion config schema_version") - if model.get("name") != "InterleavedFlowTransformer": - raise ValueError("model config must select InterleavedFlowTransformer") - if int(dataset.get("resolution", -1)) != RESOLUTION or dataset.get("attrs") != ["base_color", "metallic", "roughness", "alpha"]: - raise ValueError("model config must describe 1024 base-color/metallic/roughness/alpha inputs") - if (args.get("in_channels"), args.get("out_channels")) != (64, 32) or args.get("cond_channels") != 0: - raise ValueError("model config has an incompatible material-flow shell") - if args.get("num_blocks") != 60 or args.get("base_block_policy") != "alternating_u_x": - raise ValueError("model config has an incompatible interleaved-flow policy") - if args.get("rope_freq") not in ([1, 10000], (1, 10000)) or args.get("dtype") != "float32": - raise ValueError("model config has incompatible positional or numeric settings") - if dataset.get("bridge_z_obs_args", {}).get("partial_support_mode") != "partial_messages": - raise ValueError("model config requires partial observation messages") - _normalization(dataset.get("shape_slat_normalization"), "shape") - material = _normalization(dataset.get("pbr_slat_normalization"), "material") - for key in ("pbr_mean", "pbr_std"): - expected = material["mean" if key.endswith("mean") else "std"] - if [float(value) for value in args.get(key, [])] != expected: - raise ValueError(f"model {key} must exactly match material normalization") - return config, dataset, checkpoint_path - - -def _stem(path: str | Path, label: str) -> Path: - stem = Path(path) - _file(Path(f"{stem}.json"), f"{label} config") - _file(Path(f"{stem}.safetensors"), f"{label} weights") - return stem - - -def load_completion_contract(model_config: str | Path, checkpoint: str | Path) -> tuple[dict[str, Any], dict[str, Any], Path, dict[str, list[float]]]: - """Validate the model ABI and observation statistics without loading weights.""" - config, dataset, checkpoint_path = load_contract(model_config, checkpoint) - validate_completion_checkpoint(checkpoint_path) - path = Path(dataset["bridge_z_obs_args"]["partial_z_obs_normalization_path"]) - if not path.is_absolute(): - path = Path(model_config).expanduser().absolute().parent / path - partial_stats = _normalization(json.loads(_file(path, "partial normalization").read_text(encoding="utf-8")), "partial observation") - model_args = config["models"]["denoiser"]["args"] - for key in ("partial_mean", "partial_std"): - expected = partial_stats["mean" if key.endswith("mean") else "std"] - if [float(value) for value in model_args.get(key, [])] != expected: - raise ValueError(f"model {key} must exactly match partial-observation normalization") - return config, dataset, checkpoint_path, partial_stats - - -def _completion_vae_root(trellis_root: str | Path) -> Path: - root = Path(trellis_root).expanduser() - return root / "ckpts" if (root / "ckpts").is_dir() else root - - -def completion_decoder_stems(config: Mapping[str, Any], dataset: Mapping[str, Any], - model_config: str | Path, trellis_root: str | Path) -> dict[str, Path]: - """Resolve release VAE dependencies or preserve legacy config-relative paths.""" - base = Path(model_config).expanduser().absolute().parent - standalone = config.get("format") == "supratex_completion" - if standalone: - base = _completion_vae_root(trellis_root) - stems = {} - for key in ("pretrained_shape_slat_dec", "pretrained_pbr_slat_dec"): - value = dataset.get(key) - if not isinstance(value, str) or not value: - raise ValueError(f"model config requires {key}") - path = Path(value).expanduser() - if standalone and (path.name != value or value in (".", "..") or "\\" in value): - raise ValueError(f"standalone completion {key} must be a bare sparse VAE stem") - stems[key] = path if path.is_absolute() else base / path - return stems - - -def validate_completion_artifacts(model_config: str | Path, checkpoint: str | Path, - trellis_root: str | Path) -> tuple[dict[str, Any], dict[str, Any], Path, dict[str, list[float]]]: - """Check completion metadata, safe-file headers, and all VAE files before GPU work.""" - contract = load_completion_contract(model_config, checkpoint) - config, dataset, _, _ = contract - ckpts = _completion_vae_root(trellis_root) - stems = [ckpts / "shape_enc_next_dc_f16c32_fp16", ckpts / "tex_enc_next_dc_f16c32_fp16"] - stems.extend(completion_decoder_stems(config, dataset, model_config, trellis_root).values()) - for stem in stems: - _stem(stem, "completion artifact") - return contract - - -def _load_case(input_dir: str | Path, torch: Any, image: Any) -> dict[str, Any]: - root = Path(input_dir) - required = {name: _file(root / name, name) for name in ("geometry.pt", "camera.json", "condition_rgba.png", "albedo.png", "mra.png")} - geometry = torch.load(required["geometry.pt"], map_location="cpu", weights_only=True) - try: - vertices, faces = geometry["verts"].float().contiguous(), geometry["faces"].long().contiguous() - except (KeyError, AttributeError) as error: - raise ValueError("geometry.pt must contain tensor keys verts and faces") from error - if vertices.ndim != 2 or vertices.shape[1] != 3 or faces.ndim != 2 or faces.shape[1] != 3: - raise ValueError("geometry.pt must contain verts[N,3] and faces[M,3]") - rgba = image.open(required["condition_rgba.png"]).convert("RGBA") - albedo, mra = image.open(required["albedo.png"]).convert("RGB"), image.open(required["mra.png"]).convert("RGB") - if rgba.size[0] != rgba.size[1]: - raise ValueError("condition_rgba.png must be square") - if albedo.size != (RESOLUTION, RESOLUTION) or mra.size != (RESOLUTION, RESOLUTION): - raise ValueError("albedo.png and mra.png must be 1024 by 1024 RGB maps") - try: - camera = json.loads(required["camera.json"].read_text(encoding="utf-8")) - except json.JSONDecodeError as error: - raise ValueError("camera.json must be valid JSON") from error - if "camera_angle_x" not in camera or "distance" not in camera: - raise ValueError("camera.json must provide canonical camera_angle_x and distance") - return {"vertices": vertices, "faces": faces, "rgba": rgba, "albedo": albedo, "mra": mra, "camera": camera} - - -def _runtime() -> dict[str, Any]: - from .runtime import configure_inference_runtime - configure_inference_runtime() - from .kernel_config import configure_sparse_kernels - kernel_config = configure_sparse_kernels() - if kernel_config["enabled"]: - print(f"Sparse kernels: {kernel_config['preset']}", flush=True) - import numpy as np - import o_voxel - import torch - import torch.nn.functional as F - import utils3d - from PIL import Image, ImageFilter - from supratex import models - from supratex.modules import sparse - from supratex.representations import MeshWithVoxel - from .mesh_export import to_glb - return locals() - - -def _camera(spec: Mapping[str, Any], device: str, torch: Any, utils3d: Any) -> tuple[Any, Any]: - if "extrinsic" in spec and ("projection_intrinsic" in spec or "intrinsic" in spec): - intrinsic = spec.get("projection_intrinsic", spec.get("intrinsic")) - external, projection = torch.as_tensor(spec["extrinsic"], dtype=torch.float32, device=device), torch.as_tensor(intrinsic, dtype=torch.float32, device=device) - if external.shape != (4, 4) or projection.shape != (3, 3): raise ValueError("camera matrices must be 4x4 and 3x3") - return external, projection - if "distance" not in spec or "camera_angle_x" not in spec: raise ValueError("camera must supply distance and camera_angle_x") - eye = torch.tensor([0., 0., float(spec["distance"]) * float(spec.get("mesh_scale", 1.0))], device=device) - return utils3d.torch.extrinsics_look_at(eye, torch.zeros(3, device=device), torch.tensor([0., 1., 0.], device=device)), utils3d.torch.intrinsics_from_fov_xy(torch.tensor(float(spec["camera_angle_x"]), device=device), torch.tensor(float(spec["camera_angle_x"]), device=device)) - - -def _shape(vertices: Any, faces: Any, encoder: Any, runtime: Mapping[str, Any], device: str) -> tuple[Any, Any]: - torch, o_voxel, sparse = runtime["torch"], runtime["o_voxel"], runtime["sparse"] - with span("voxelization"): - indices, dual_vertices, intersected = o_voxel.convert.mesh_to_flexible_dual_grid(vertices=vertices.float().cpu().contiguous(), faces=faces.long().cpu().contiguous(), grid_size=RESOLUTION, aabb=[[-.5, -.5, -.5], [.5, .5, .5]], face_weight=1., boundary_weight=.2, regularization_weight=1e-2, timing=False) - offsets = (dual_vertices.float() * RESOLUTION - indices.float()).clamp(0, 1) - coords = torch.cat((torch.zeros_like(indices[:, :1]), indices), dim=1).int().to(device) - vertices_sparse = sparse.SparseTensor(offsets.to(device), coords) - with span("shape_vae"): - with torch.no_grad(): latent = encoder(vertices_sparse, vertices_sparse.replace(intersected.bool().to(device))) - return latent, sparse.SparseTensor(torch.zeros(coords.shape[0], 6, device=device), coords) - - -def _partial(surface: Any, shape: Any, case: Mapping[str, Any], external: Any, projection: Any, encoder: Any, normalization: Mapping[str, Any], runtime: Mapping[str, Any], device: str) -> Any: - torch, F, np, image, image_filter, sparse = runtime["torch"], runtime["F"], runtime["np"], runtime["Image"], runtime["ImageFilter"], runtime["sparse"] - coords = surface.coords[:, 1:].long(); points = (coords.float() + .5) / RESOLUTION - .5 - camera = (torch.cat((points, torch.ones_like(points[:, :1])), dim=1) @ external.T)[:, :3] - uvw = camera @ projection.T; uv = uvw[:, :2] / uvw[:, 2:3].clamp(min=1e-8) - grid = torch.stack((uv[:, 0] * 2 - 1, uv[:, 1] * 2 - 1), dim=1).view(1, 1, -1, 2) - def sampled(source: Any) -> Any: - data = np.asarray(source.convert("RGB").resize((RESOLUTION, RESOLUTION), image.Resampling.LANCZOS), dtype=np.float32).copy() / 255. - return F.grid_sample(torch.from_numpy(data).permute(2, 0, 1).to(device)[None], grid, mode="bilinear", align_corners=False, padding_mode="border")[0, :, 0].T.clamp(0, 1) - alpha = np.asarray(case["rgba"].getchannel("A"), dtype=np.uint8) - mask = image.fromarray(np.where(alpha >= 254, 255, 0).astype(np.uint8)).resize((RESOLUTION, RESOLUTION), image.Resampling.NEAREST).filter(image_filter.MinFilter(5)) - trusted = torch.from_numpy(np.asarray(mask, dtype=np.uint8).copy() == 255) - bits, ordered, order = native_observation.build_occupancy_index(coords) - owners = native_observation.pixel_first_occupied(dict(case["camera"]), bits, trusted) - visible = native_observation.owner_rows(coords, owners, ordered, order) - if not bool(visible.any()): raise RuntimeError("the trusted alpha gate leaves no observed material surface") - with torch.no_grad(): - return native_observation.encode_partial(encoder, coords, visible, sampled(case["albedo"]), sampled(case["mra"]), shape.coords[:, 1:], dict(normalization), sparse, device) - - -def _export(output: Path, representation: Any, texture_size: int, decimation_target: int, texture_format: str, runtime: Mapping[str, Any], *, gutter: bool = False, low_vram: bool = False) -> None: - with span("export_prepare"): - if output.suffix.lower() != ".glb": raise ValueError("output must have a .glb suffix") - output.parent.mkdir(parents=True, exist_ok=True) - lock = output.with_name(f".{output.name}.publish.lock") - try: - handle = lock.open("x") - except FileExistsError as error: - raise FileExistsError(f"another process is publishing {output}") from error - handle.close() - if output.exists(): - lock.unlink(missing_ok=True) - raise FileExistsError(f"refusing to overwrite existing output: {output}") - temporary = output.with_name(f".{output.stem}.tmp.{os.getpid()}.glb") - try: - scene = runtime["to_glb"](vertices=representation.vertices, faces=representation.faces, attr_volume=representation.attrs, coords=representation.coords, attr_layout=LAYOUT, grid_size=RESOLUTION, aabb=[[-.5, -.5, -.5], [.5, .5, .5]], decimation_target=decimation_target, texture_size=texture_size, remesh=True, remesh_band=1, remesh_project=0, use_tqdm=True, **({"low_vram": True} if low_vram else {})) - if gutter: - with span("gutter"): - from .gutter import apply_gutter - print(f"Gutter postprocess: {apply_gutter(scene)}", flush=True) - with span("file_write"): - scene.apply_transform(runtime["np"].asarray([[-1, 0, 0, 0], [0, 0, -1, 0], [0, -1, 0, 0], [0, 0, 0, 1]], dtype=runtime["np"].float64)) - scene.export(str(temporary), extension_webp=texture_format == "webp") - if output.exists(): - raise FileExistsError(f"refusing to overwrite existing output: {output}") - os.replace(temporary, output) - finally: - with span("export_cleanup"): - temporary.unlink(missing_ok=True) - lock.unlink(missing_ok=True) - - -def _stage_sparse(value: Any, device: str, *, keep_topology: bool = False) -> Any: - """Transfer sparse data without retaining backend caches on its old device. - - SparseTensor.to()/cpu() share spatial/backend caches. Encoder subdivision - maps are semantic: the decoders use them to reproduce the input surface, - overriding predicted subdivisions. Preserve those maps when requested; - convolution and attention caches are safe to recompute on the new device. - """ - def transfer(item: Any) -> Any: - if isinstance(item, dict): - return {key: transfer(child) for key, child in item.items()} - if isinstance(item, (list, tuple)): - return type(item)(transfer(child) for child in item) - if hasattr(item, "detach") and hasattr(item, "to"): - return item.detach().to(device) - return item - staged = type(value)(value.feats.detach().to(device), value.coords.detach().to(device), - shape=value._shape, scale=value._scale) - if keep_topology: - staged._spatial_cache = { - scale: {name: transfer(item) for name, item in cache.items() - if name.startswith(("channel2spatial_", "upsample_")) or name == "shape"} - for scale, cache in value._spatial_cache.items() - } - return staged - - -@contextmanager -def _completion_stage(model: Any, torch: Any, device: str, *, layer_offload: bool = False): - """Own one model's residency, including cleanup after failed inference.""" - previous = getattr(model, "low_vram", False) - try: - if layer_offload: - model.low_vram = True - else: - model.to(device) - model.eval() - with torch.no_grad(): - yield model - finally: - model.to("cpu") - if layer_offload: - model.low_vram = previous - torch.cuda.empty_cache() - - -def _run_completion_low_vram(*, case: Mapping[str, Any], back_case: Mapping[str, Any] | None, - ckpts: Path, config: Mapping[str, Any], dataset: Mapping[str, Any], - checkpoint_path: Path, partial_stats: Mapping[str, Any], resolved: Any, - runtime: Mapping[str, Any], device: str, seed: int, output: str | Path, - texture_size: int, decimation_target: int, texture_format: str, gutter: bool) -> dict[str, Any]: - """Serialize completion stages and keep only their CPU handoffs alive.""" - torch, models = runtime["torch"], runtime["models"] - print("Completion low VRAM: sequential models, CPU handoffs, decoder layer offload", flush=True) - shape_encoder = models.from_pretrained(str(_stem(ckpts / "shape_enc_next_dc_f16c32_fp16", "shape encoder"))) - with _completion_stage(shape_encoder, torch, device): - shape, surface = _shape(case["vertices"], case["faces"], shape_encoder, runtime, device) - shape_cpu = _stage_sparse(shape, "cpu", keep_topology=True) - shape = _stage_sparse(shape, device) - del shape_encoder - texture_encoder = models.from_pretrained(str(_stem(ckpts / "tex_enc_next_dc_f16c32_fp16", "material encoder"))) - dual_stats = {} - with _completion_stage(texture_encoder, torch, device): - if back_case is None: - external, projection = _camera(case["camera"], device, torch, runtime["utils3d"]) - observation = _partial(surface, shape, case, external, projection, texture_encoder, partial_stats, runtime, device) - del external, projection - else: - from .dual_view import encode_two_views - observation, dual_stats = encode_two_views(surface, shape, case, back_case, texture_encoder, - partial_stats, runtime, device) - del texture_encoder, surface - torch.cuda.empty_cache() - - from supratex.pipelines.observation_guided_completion import build_completion_model, sample_material_latents - state = load_completion_state_dict(checkpoint_path, torch_module=torch) - model = build_completion_model(config, state, assign_weights=True) - del state - with _completion_stage(model, torch, device): - shape_condition = shape.replace(feats=(shape.feats - torch.as_tensor(dataset["shape_slat_normalization"]["mean"], device=device)[None]) / torch.as_tensor(dataset["shape_slat_normalization"]["std"], device=device)[None]) - noise = shape.replace(feats=torch.randn((shape.coords.shape[0], 32), generator=torch.Generator(device=device).manual_seed(int(seed)), device=device)) - print("Observation cache: disabled (low VRAM)", flush=True) - sampled = sample_material_latents(model, noise, concat_cond=shape_condition, partial_z_obs=observation, steps=25, rescale_t=3., sigma_min=1e-5, use_observation_cache=False).samples - if not torch.equal(sampled.coords, noise.coords) or sampled.feats.shape != noise.feats.shape: - raise RuntimeError("material sampler changed sparse support or feature shape") - generated = shape.replace(feats=sampled.feats.float() * torch.as_tensor(dataset["pbr_slat_normalization"]["std"], device=device)[None] + torch.as_tensor(dataset["pbr_slat_normalization"]["mean"], device=device)[None]) - generated_cpu = _stage_sparse(generated, "cpu") - del model, shape, shape_condition, noise, sampled, generated, observation - torch.cuda.empty_cache() - - # Reuse the encoder's exact support in both decoders. Do not carry GPU - # convolution caches or the unused decoded mesh into material decoding. - shape_decoder = models.from_pretrained(str(_stem(resolved(dataset["pretrained_shape_slat_dec"]), "shape decoder"))) - shape_decoder.set_resolution(RESOLUTION) - with _completion_stage(shape_decoder, torch, device, layer_offload=True): - shape = _stage_sparse(shape_cpu, device, keep_topology=True) - decoded_shape, guide = shape_decoder(shape, return_subs=True, return_mesh=False) - guide_cpu = [_stage_sparse(value, "cpu") for value in guide] - del shape_decoder, shape, decoded_shape, guide - generated_cpu._spatial_cache = shape_cpu._spatial_cache - del shape_cpu - torch.cuda.empty_cache() - - material_decoder = models.from_pretrained(str(_stem(resolved(dataset["pretrained_pbr_slat_dec"]), "material decoder"))) - with _completion_stage(material_decoder, torch, device, layer_offload=True): - generated = _stage_sparse(generated_cpu, device, keep_topology=True) - guide = [_stage_sparse(value, device) for value in guide_cpu] - decoded = material_decoder(generated, guide_subs=guide) * .5 + .5 - if decoded.feats.shape[1] != 6: raise RuntimeError("material decoder must emit six PBR channels") - decoded_cpu = _stage_sparse(decoded, "cpu") - del material_decoder, generated, generated_cpu, guide, guide_cpu, decoded - torch.cuda.empty_cache() - representation = runtime["MeshWithVoxel"](case["vertices"], case["faces"], [-.5, -.5, -.5], 1. / RESOLUTION, decoded_cpu.coords[:, 1:].int(), decoded_cpu.feats.float().clamp(0, 1), torch.Size([1, 6, RESOLUTION, RESOLUTION, RESOLUTION]), layout=LAYOUT) - del decoded_cpu - _export(Path(output), representation, texture_size, decimation_target, texture_format, runtime, gutter=gutter, low_vram=True) - return dual_stats - - -def _run_completion(*, input_dir: str | Path, output: str | Path, model_config: str | Path, checkpoint: str | Path, trellis_root: str | Path, seed: int = 42, device: str = "cuda", texture_size: int = 4096, decimation_target: int = 500000, texture_format: str = "png", gutter: bool = False, back_case: Mapping[str, Any] | None = None, resident=None, preview_only=False) -> dict[str, Any]: - """Complete PBR texture from public input maps and write one new GLB.""" - with span("validation"): - if device != "cuda": raise ValueError("public material completion currently requires device='cuda'") - if texture_size <= 0 or decimation_target <= 0 or texture_format not in {"png", "webp"}: raise ValueError("invalid GLB export settings") - config, dataset, checkpoint_path, partial_stats = validate_completion_artifacts(model_config, checkpoint, trellis_root) - decoder_stems = completion_decoder_stems(config, dataset, model_config, trellis_root) - decoder_paths = {dataset[key]: path for key, path in decoder_stems.items()} - def resolved(value: str) -> Path: - return decoder_paths[value] - with span("runtime"): - runtime = _runtime(); torch, models = runtime["torch"], runtime["models"] - if not torch.cuda.is_available(): raise RuntimeError("CUDA is unavailable") - with span("case_load"): - case = _load_case(input_dir, torch, runtime["Image"]); ckpts = _completion_vae_root(trellis_root) - if resident is None and os.environ.get("SUPRATEX_LOW_VRAM", "0") == "1": - return _run_completion_low_vram(case=case, back_case=back_case, ckpts=ckpts, config=config, - dataset=dataset, checkpoint_path=checkpoint_path, partial_stats=partial_stats, - resolved=resolved, runtime=runtime, device=device, seed=seed, output=output, - texture_size=texture_size, decimation_target=decimation_target, - texture_format=texture_format, gutter=gutter) - with span("encoder_load"): - if resident is not None: - shape_encoder = resident.completion_model('shape_encoder').to(device).eval() - texture_encoder = resident.completion_model('texture_encoder').to(device).eval() - else: - shape_encoder = models.from_pretrained(str(_stem(ckpts / "shape_enc_next_dc_f16c32_fp16", "shape encoder"))).to(device).eval() - texture_encoder = models.from_pretrained(str(_stem(ckpts / "tex_enc_next_dc_f16c32_fp16", "material encoder"))).to(device).eval() - with span("shape_encode"): - external, projection = _camera(case["camera"], device, torch, runtime["utils3d"]) - shape, surface = _shape(case["vertices"], case["faces"], shape_encoder, runtime, device) - shape_encoder.to("cpu"); torch.cuda.empty_cache() - with span("observation_encode"): - dual_stats = {} - if back_case is None: - observation = _partial(surface, shape, case, external, projection, texture_encoder, partial_stats, runtime, device) - else: - from .dual_view import encode_two_views - observation, dual_stats = encode_two_views(surface, shape, case, back_case, texture_encoder, - partial_stats, runtime, device) - texture_encoder.to("cpu"); torch.cuda.empty_cache() - with span("model_load"): - from supratex.pipelines.observation_guided_completion import build_completion_model, sample_material_latents - if resident is not None: - model = resident.completion_model('transformer').to(device).eval() - else: - state = load_completion_state_dict(checkpoint_path, torch_module=torch) - model = build_completion_model(config, state, assign_weights=True).to(device).eval() - del state - with span("conditioning"): - shape_condition = shape.replace(feats=(shape.feats - torch.as_tensor(dataset["shape_slat_normalization"]["mean"], device=device)[None]) / torch.as_tensor(dataset["shape_slat_normalization"]["std"], device=device)[None]) - noise = shape.replace(feats=torch.randn((shape.coords.shape[0], 32), generator=torch.Generator(device=device).manual_seed(int(seed)), device=device)) - # Cache fixed observation work within this sample only. Retain ample room - # for attention activations; low-memory devices keep the uncached path. - cache_bytes = 2 * (len(model.blocks) + len(model.refiners)) * observation.feats.shape[0] * model.model_channels * 2 - free_bytes, _ = torch.cuda.mem_get_info(noise.feats.device) - use_cache = free_bytes >= cache_bytes * 1.25 + 8 * 1024 ** 3 - print(f"Observation cache: {'enabled' if use_cache else 'disabled (memory headroom)'}", flush=True) - with span("sample"): - with span("sample_25_steps"): - sampled = sample_material_latents(model, noise, concat_cond=shape_condition, partial_z_obs=observation, steps=25, rescale_t=3., sigma_min=1e-5, use_observation_cache=use_cache).samples - if not torch.equal(sampled.coords, noise.coords) or sampled.feats.shape != noise.feats.shape: - raise RuntimeError("material sampler changed sparse support or feature shape") - generated = shape.replace(feats=sampled.feats.float() * torch.as_tensor(dataset["pbr_slat_normalization"]["std"], device=device)[None] + torch.as_tensor(dataset["pbr_slat_normalization"]["mean"], device=device)[None]) - model.to("cpu"); torch.cuda.empty_cache() - with span("decoder_load"): - shape_stem, material_stem = _stem(resolved(dataset["pretrained_shape_slat_dec"]), "shape decoder"), _stem(resolved(dataset["pretrained_pbr_slat_dec"]), "material decoder") - if resident is not None: - shape_decoder = resident.completion_model('shape_decoder').to(device).eval() - material_decoder = resident.completion_model('material_decoder').to(device).eval() - else: - shape_decoder, material_decoder = models.from_pretrained(str(shape_stem)).to(device).eval(), models.from_pretrained(str(material_stem)).to(device).eval() - shape_decoder.set_resolution(RESOLUTION) - with span("decode"): - with torch.no_grad(): _, guide = shape_decoder(shape, return_subs=True); decoded = material_decoder(generated, guide_subs=guide) * .5 + .5 - if decoded.feats.shape[1] != 6: raise RuntimeError("material decoder must emit six PBR channels") - representation = runtime["MeshWithVoxel"](case["vertices"].to(device), case["faces"].to(device), [-.5, -.5, -.5], 1. / RESOLUTION, decoded.coords[:, 1:].int(), decoded.feats.float().clamp(0, 1), torch.Size([1, 6, RESOLUTION, RESOLUTION, RESOLUTION]), layout=LAYOUT) - if resident is not None: - shape_decoder.cpu(); material_decoder.cpu(); torch.cuda.empty_cache() - if preview_only: - from hf_space.ovoxel_preview import save_and_render - save_and_render(Path(output), representation) - else: - with span("export"): - _export(Path(output), representation, texture_size, decimation_target, texture_format, runtime, gutter=gutter) - return dual_stats - - -def run_completion(*, input_dir: str | Path, output: str | Path, model_config: str | Path, - checkpoint: str | Path, trellis_root: str | Path, seed: int = 42, - device: str = "cuda", texture_size: int = 4096, - decimation_target: int = 500000, texture_format: str = "png", gutter: bool = False, - back_view_dir: str | Path | None = None, work_dir: str | Path | None = None, resident=None, preview_only=False) -> None: - """Complete one view by default, or validate and publish a dual-view result.""" - arguments = dict(input_dir=input_dir, output=output, model_config=model_config, - checkpoint=checkpoint, trellis_root=trellis_root, seed=seed, device=device, - texture_size=texture_size, decimation_target=decimation_target, - texture_format=texture_format, gutter=gutter, resident=resident, preview_only=preview_only) - if preview_only and back_view_dir is not None: - raise ValueError('Preview mode currently requires single-view completion') - if back_view_dir is None: - if work_dir is not None: - raise ValueError("completion --work-dir is reserved for dual-view inference") - _run_completion(**arguments) - return - from .dual_view import CAMERA_PROTOCOL, PROTOCOL, digest, load_back_case, write_json - import shutil - - output = Path(output).expanduser().resolve() - work = Path(work_dir).expanduser().resolve() if work_dir else output.with_name(output.stem + "_dual_view") - if output.suffix.lower() != ".glb": - raise ValueError("output must have a .glb suffix") - if output.exists(): - raise FileExistsError(f"refusing to overwrite existing output: {output}") - if work.exists(): - raise FileExistsError(f"dual-view work directory must be new: {work}") - if texture_size != 4096 or texture_format != "png" or decimation_target != 500000 or gutter: - raise ValueError("dual-view reprojection requires a 4096 PNG atlas, 500000-face target, and no --gutter") - front_root, back_root = Path(input_dir).expanduser().resolve(), Path(back_view_dir).expanduser().resolve() - back = load_back_case(front_root, back_root) - work.mkdir(parents=True) - receipt = {"protocol": PROTOCOL, "camera_protocol": CAMERA_PROTOCOL, "view": "dual", - "status": "running", "stage": "texture", "front_input": str(front_root), - "back_input": str(back_root), "output": str(output), - "checkpoint": str(Path(checkpoint).expanduser().resolve()), - "model_config": str(Path(model_config).expanduser().resolve()), - "seed": seed, "sampler": "Euler", "steps": 25, - "geometry_sha256": digest(front_root / "geometry.pt"), - "cameras_sha256": digest(back_root / "cameras.json"), - "front_material_receipt_sha256": digest(front_root / "material_receipt.json"), - "back_material_receipt_sha256": digest(back_root / "material_receipt.json")} - write_json(work / "receipt.json", receipt) - try: - native = work / "preds/mesh.glb" - arguments.update(input_dir=front_root, output=native, back_case=back) - stats = _run_completion(**arguments) - # Reprojection runs in separate CUDA workers. Release the completion - # allocator's inactive blocks before those processes create contexts. - import gc - import torch - gc.collect() - torch.cuda.empty_cache() - write_json(work / "receipt_texture.json", {**receipt, "status": "success", - "dual_stats": stats, "sampling_calls": 1, "mesh": str(native), "mesh_sha256": digest(native)}) - receipt.update(stage="reprojection") - write_json(work / "receipt.json", receipt) - from .reprojection import run_dual_view, validate_dual_view - final = run_dual_view(native, front_root, back_root, back["cameras"], work) - validation = validate_dual_view(work, raise_on_invalid=True) - # Model and reprojection work can be long-running. Revalidate the - # original inputs and generation evidence before publishing the mesh. - load_back_case(front_root, back_root) - bindings = {"geometry_sha256": front_root / "geometry.pt", - "cameras_sha256": back_root / "cameras.json", - "front_material_receipt_sha256": front_root / "material_receipt.json", - "back_material_receipt_sha256": back_root / "material_receipt.json"} - for key, path in bindings.items(): - if receipt[key] != digest(path): - raise RuntimeError(f"input changed during dual-view completion: {key}") - output.parent.mkdir(parents=True, exist_ok=True) - temporary = output.with_name(f".{output.name}.dual-view.{os.getpid()}.tmp") - try: - with temporary.open("xb") as target, final.open("rb") as source: - shutil.copyfileobj(source, target) - # Atomic publication that cannot replace another writer's output. - os.link(temporary, output) - finally: - temporary.unlink(missing_ok=True) - receipt.update(status="valid", stage="complete", mesh=str(final), mesh_sha256=digest(final), - output_sha256=digest(output), dual_stats=stats, - postprocess_validation=validation, stitched_19_phase_checked=False) - write_json(work / "receipt.json", receipt) - except Exception as error: - receipt.update(status="failed", error=f"{type(error).__name__}: {error}") - write_json(work / "receipt.json", receipt) - raise +"""Observation-guided PBR completion and GLB export. + +The public boundary begins with a mesh/camera/RGBA observation plus albedo and +MRA maps. Geometry prediction and 2D map prediction intentionally remain +outside this inference-only package. +""" +from __future__ import annotations + +import json +import math +import os +import sys +from contextlib import contextmanager +from pathlib import Path +from typing import Any, Mapping + +from . import native_observation +from .checkpoint import load_completion_state_dict, validate_completion_checkpoint + +RESOLUTION = native_observation.RESOLUTION +LAYOUT = {"base_color": slice(0, 3), "metallic": slice(3, 4), "roughness": slice(4, 5), "alpha": slice(5, 6)} +# Released completion-model ABI. It is checked once, when the config is loaded. +MODEL_ABI = {"in_channels": 64, "out_channels": 32, "cond_channels": 0, "model_channels": 1536, + "num_heads": 12, "num_blocks": 60, "base_block_policy": "alternating_u_x", + "v11_init_policy": "odd_copy_even_identity", "direct_modulation_broadcast": True, + "rope_freq": [1, 10000], "dtype": "float32"} + + +def _file(path: str | Path, label: str) -> Path: + value = Path(path) + if not value.is_file(): + raise FileNotFoundError(f"missing {label}: {value}") + return value + + +def _normalization(value: Any, label: str) -> dict[str, list[float]]: + if not isinstance(value, Mapping): + raise ValueError(f"{label} normalization must be an object") + result: dict[str, list[float]] = {} + for key in ("mean", "std"): + raw = value.get(key) + if not isinstance(raw, list) or len(raw) != 32: + raise ValueError(f"{label} normalization requires 32-value {key}") + result[key] = [float(item) for item in raw] + if not all(math.isfinite(item) for item in result["mean"] + result["std"]) or any(item <= 0 for item in result["std"]): + raise ValueError(f"{label} normalization standard deviations must be positive") + return result + + +def load_contract(model_config: str | Path, checkpoint: str | Path) -> tuple[dict[str, Any], dict[str, Any], Path]: + """Load the public model ABI without inspecting or hashing weight bytes.""" + config_path, checkpoint_path = _file(model_config, "model config"), _file(checkpoint, "checkpoint") + try: + config = json.loads(config_path.read_text(encoding="utf-8")) + dataset = config["dataset"]["args"] + model = config["models"]["denoiser"] + args = model["args"] + except (OSError, KeyError, TypeError, json.JSONDecodeError) as error: + raise ValueError("model config does not satisfy the Supra-Tex inference ABI") from error + if config.get("format") == "supratex_completion" and config.get("schema_version") != 1: + raise ValueError("unsupported Supra-Tex completion config schema_version") + if model.get("name") != "InterleavedFlowTransformer": + raise ValueError("model config must select InterleavedFlowTransformer") + if int(dataset.get("resolution", -1)) != RESOLUTION or dataset.get("attrs") != ["base_color", "metallic", "roughness", "alpha"]: + raise ValueError("model config must describe 1024 base-color/metallic/roughness/alpha inputs") + mismatched = {key: args.get(key) for key, expected in MODEL_ABI.items() if args.get(key) != expected} + if mismatched: + raise ValueError(f"model config does not match the released completion ABI: {mismatched}") + if dataset.get("bridge_z_obs_args", {}).get("partial_support_mode") != "partial_messages": + raise ValueError("model config requires partial observation messages") + _normalization(dataset.get("shape_slat_normalization"), "shape") + material = _normalization(dataset.get("pbr_slat_normalization"), "material") + for key in ("pbr_mean", "pbr_std"): + expected = material["mean" if key.endswith("mean") else "std"] + if [float(value) for value in args.get(key, [])] != expected: + raise ValueError(f"model {key} must exactly match material normalization") + return config, dataset, checkpoint_path + + +def _stem(path: str | Path, label: str) -> Path: + stem = Path(path) + _file(Path(f"{stem}.json"), f"{label} config") + _file(Path(f"{stem}.safetensors"), f"{label} weights") + return stem + + +def load_completion_contract(model_config: str | Path, checkpoint: str | Path) -> tuple[dict[str, Any], dict[str, Any], Path, dict[str, list[float]]]: + """Validate the model ABI and observation statistics without loading weights.""" + config, dataset, checkpoint_path = load_contract(model_config, checkpoint) + validate_completion_checkpoint(checkpoint_path) + path = Path(dataset["bridge_z_obs_args"]["partial_z_obs_normalization_path"]) + if not path.is_absolute(): + path = Path(model_config).expanduser().absolute().parent / path + partial_stats = _normalization(json.loads(_file(path, "partial normalization").read_text(encoding="utf-8")), "partial observation") + model_args = config["models"]["denoiser"]["args"] + for key in ("partial_mean", "partial_std"): + expected = partial_stats["mean" if key.endswith("mean") else "std"] + if [float(value) for value in model_args.get(key, [])] != expected: + raise ValueError(f"model {key} must exactly match partial-observation normalization") + return config, dataset, checkpoint_path, partial_stats + + +def _completion_vae_root(trellis_root: str | Path) -> Path: + root = Path(trellis_root).expanduser() + return root / "ckpts" if (root / "ckpts").is_dir() else root + + +def completion_decoder_stems(config: Mapping[str, Any], dataset: Mapping[str, Any], + model_config: str | Path, trellis_root: str | Path) -> dict[str, Path]: + """Resolve release VAE dependencies or preserve legacy config-relative paths.""" + base = Path(model_config).expanduser().absolute().parent + standalone = config.get("format") == "supratex_completion" + if standalone: + base = _completion_vae_root(trellis_root) + stems = {} + for key in ("pretrained_shape_slat_dec", "pretrained_pbr_slat_dec"): + value = dataset.get(key) + if not isinstance(value, str) or not value: + raise ValueError(f"model config requires {key}") + path = Path(value).expanduser() + if standalone and (path.name != value or value in (".", "..") or "\\" in value): + raise ValueError(f"standalone completion {key} must be a bare sparse VAE stem") + stems[key] = path if path.is_absolute() else base / path + return stems + + +def validate_completion_artifacts(model_config: str | Path, checkpoint: str | Path, + trellis_root: str | Path) -> tuple[dict[str, Any], dict[str, Any], Path, dict[str, list[float]]]: + """Check completion metadata, safe-file headers, and all VAE files before GPU work.""" + contract = load_completion_contract(model_config, checkpoint) + config, dataset, _, _ = contract + ckpts = _completion_vae_root(trellis_root) + stems = [ckpts / "shape_enc_next_dc_f16c32_fp16", ckpts / "tex_enc_next_dc_f16c32_fp16"] + stems.extend(completion_decoder_stems(config, dataset, model_config, trellis_root).values()) + for stem in stems: + _stem(stem, "completion artifact") + return contract + + +def _load_vae(models: Any, stem: Path) -> Any: + """Build a sparse VAE on the meta device so weights are not randomly initialized first. + + Checkpoint tensors are cast to the constructor's dtypes, exactly as the + copy-based loader does. Partial checkpoints or constructor-made tensors that + a meta shell cannot materialize fall back to models.from_pretrained. + """ + import torch + from safetensors.torch import load_file + + config = json.loads(Path(f"{stem}.json").read_text(encoding="utf-8")) + with torch.device("meta"): + model = getattr(models, config["name"])(**config["args"]) + expected = model.state_dict() + state = load_file(f"{stem}.safetensors", device="cpu") + if state.keys() != expected.keys() or any(state[name].shape != target.shape for name, target in expected.items()): + return models.from_pretrained(str(stem)) + model.load_state_dict({name: state[name].to(dtype=target.dtype) for name, target in expected.items()}, + strict=True, assign=True) + if any(isinstance(value, torch.Tensor) and value.is_meta + for module in model.modules() for value in (*module._parameters.values(), *module._buffers.values(), *vars(module).values())): + return models.from_pretrained(str(stem)) + return model + + +def _load_geometry(root: Path, torch: Any) -> tuple[Any, Any]: + """Read geometry.pt, or a compressed geometry.npz with the same verts/faces arrays.""" + if not (root / "geometry.pt").is_file() and (root / "geometry.npz").is_file(): + import numpy as np + with np.load(root / "geometry.npz") as arrays: + try: + return torch.from_numpy(arrays["verts"]).float().contiguous(), torch.from_numpy(arrays["faces"]).long().contiguous() + except KeyError as error: + raise ValueError("geometry.npz must contain arrays verts and faces") from error + geometry = torch.load(_file(root / "geometry.pt", "geometry.pt or geometry.npz"), map_location="cpu", weights_only=True) + try: + return geometry["verts"].float().contiguous(), geometry["faces"].long().contiguous() + except (KeyError, AttributeError) as error: + raise ValueError("geometry.pt must contain tensor keys verts and faces") from error + + +def _load_case(input_dir: str | Path, torch: Any, image: Any, remove_floaters: bool = False) -> dict[str, Any]: + root = Path(input_dir) + required = {name: _file(root / name, name) for name in ("camera.json", "condition_rgba.png", "albedo.png", "mra.png")} + vertices, faces = _load_geometry(root, torch) + if vertices.ndim != 2 or vertices.shape[1] != 3 or faces.ndim != 2 or faces.shape[1] != 3: + raise ValueError("geometry must contain verts[N,3] and faces[M,3]") + rgba = image.open(required["condition_rgba.png"]).convert("RGBA") + albedo, mra = image.open(required["albedo.png"]).convert("RGB"), image.open(required["mra.png"]).convert("RGB") + if rgba.size[0] != rgba.size[1]: + raise ValueError("condition_rgba.png must be square") + if albedo.size != (RESOLUTION, RESOLUTION) or mra.size != (RESOLUTION, RESOLUTION): + raise ValueError("albedo.png and mra.png must be 1024 by 1024 RGB maps") + try: + camera = json.loads(required["camera.json"].read_text(encoding="utf-8")) + except json.JSONDecodeError as error: + raise ValueError("camera.json must be valid JSON") from error + if "camera_angle_x" not in camera or "distance" not in camera: + raise ValueError("camera.json must provide canonical camera_angle_x and distance") + if remove_floaters: + import numpy as np + from .geometry_cleanup import remove_floaters as clean + vertices, faces, report = clean(vertices, faces, camera, np.asarray(rgba.getchannel("A"))) + vertices, faces = vertices.contiguous(), faces.contiguous() + print(f"Floaters: removed {report['removed_components']} of {report['components']} components " + f"({report['removed_faces']} of {report['faces']} faces)", flush=True) + return {"vertices": vertices, "faces": faces, "rgba": rgba, "albedo": albedo, "mra": mra, "camera": camera} + + +def _runtime() -> dict[str, Any]: + from .runtime import configure_inference_runtime + configure_inference_runtime() + from .kernel_config import configure_sparse_kernels + kernel_config = configure_sparse_kernels() + if kernel_config["enabled"]: + print(f"Sparse kernels: {kernel_config['preset']}", flush=True) + import numpy as np + import o_voxel + import torch + import torch.nn.functional as F + import utils3d + from PIL import Image, ImageFilter + from supratex import models + from supratex.modules import sparse + from supratex.representations import MeshWithVoxel + from .mesh_export import to_glb + return locals() + + +def _camera(spec: Mapping[str, Any], device: str, torch: Any, utils3d: Any) -> tuple[Any, Any]: + if "extrinsic" in spec and ("projection_intrinsic" in spec or "intrinsic" in spec): + intrinsic = spec.get("projection_intrinsic", spec.get("intrinsic")) + external, projection = torch.as_tensor(spec["extrinsic"], dtype=torch.float32, device=device), torch.as_tensor(intrinsic, dtype=torch.float32, device=device) + if external.shape != (4, 4) or projection.shape != (3, 3): raise ValueError("camera matrices must be 4x4 and 3x3") + return external, projection + if "distance" not in spec or "camera_angle_x" not in spec: raise ValueError("camera must supply distance and camera_angle_x") + eye = torch.tensor([0., 0., float(spec["distance"]) * float(spec.get("mesh_scale", 1.0))], device=device) + return utils3d.torch.extrinsics_look_at(eye, torch.zeros(3, device=device), torch.tensor([0., 1., 0.], device=device)), utils3d.torch.intrinsics_from_fov_xy(torch.tensor(float(spec["camera_angle_x"]), device=device), torch.tensor(float(spec["camera_angle_x"]), device=device)) + + +def _shape(vertices: Any, faces: Any, encoder: Any, runtime: Mapping[str, Any], device: str) -> tuple[Any, Any]: + torch, o_voxel, sparse = runtime["torch"], runtime["o_voxel"], runtime["sparse"] + indices, dual_vertices, intersected = o_voxel.convert.mesh_to_flexible_dual_grid(vertices=vertices.float().cpu().contiguous(), faces=faces.long().cpu().contiguous(), grid_size=RESOLUTION, aabb=[[-.5, -.5, -.5], [.5, .5, .5]], face_weight=1., boundary_weight=.2, regularization_weight=1e-2, timing=False) + offsets = (dual_vertices.float() * RESOLUTION - indices.float()).clamp(0, 1) + coords = torch.cat((torch.zeros_like(indices[:, :1]), indices), dim=1).int().to(device) + vertices_sparse = sparse.SparseTensor(offsets.to(device), coords) + with torch.no_grad(): latent = encoder(vertices_sparse, vertices_sparse.replace(intersected.bool().to(device))) + return latent, sparse.SparseTensor(torch.zeros(coords.shape[0], 6, device=device), coords) + + +ALIGNMENT_WARNING_IOU = .9 + + +def _alignment_report(uv: Any, alpha: Any, admission: Mapping[str, int], np: Any, image: Any, image_filter: Any) -> dict[str, float]: + """Compare the projected geometry silhouette with the image foreground and print the result.""" + pixels = (uv.detach().cpu().numpy() * RESOLUTION).astype(np.int64) + inside = (pixels >= 0).all(1) & (pixels < RESOLUTION).all(1) + splat = np.zeros((RESOLUTION, RESOLUTION), dtype=np.uint8) + splat[pixels[inside, 1], pixels[inside, 0]] = 255 + # Surface cells project to about one pixel each; a 3x3 closing fills the gaps. + silhouette = np.asarray(image.fromarray(splat).filter(image_filter.MaxFilter(3)).filter(image_filter.MinFilter(3))) > 0 + foreground = np.asarray(image.fromarray(alpha).resize((RESOLUTION, RESOLUTION), image.Resampling.NEAREST)) > 127 + overlap = float((silhouette & foreground).sum()) + report = {"silhouette_iou": overlap / max(float((silhouette | foreground).sum()), 1.), + "geometry_inside_image": overlap / max(float(silhouette.sum()), 1.), + "image_covered_by_geometry": overlap / max(float(foreground.sum()), 1.), **admission} + rejected = [f"missed geometry {admission['trusted_rays'] - admission['first_hits']}", + f"depth edge {admission['rejected_depth_edge']}"] + rejected += [f"{label} {admission[key]}" for key, label in (("rejected_silhouette", "silhouette margin"), + ("rejected_grazing", "grazing")) if admission[key]] + print(f"Observation: silhouette IoU {report['silhouette_iou']:.3f} " + f"(geometry inside image {report['geometry_inside_image']:.3f}, image covered {report['image_covered_by_geometry']:.3f}); " + f"admitted {admission['admitted']} of {admission['trusted_rays']} trusted pixels ({', '.join(rejected)})", flush=True) + if report["silhouette_iou"] < ALIGNMENT_WARNING_IOU: + print(f"Warning: the image and geometry silhouettes disagree (IoU {report['silhouette_iou']:.3f} < " + f"{ALIGNMENT_WARNING_IOU}); mismatched regions rely on completion rather than the image.", flush=True) + return report + + +def _partial(surface: Any, shape: Any, case: Mapping[str, Any], external: Any, projection: Any, encoder: Any, normalization: Mapping[str, Any], runtime: Mapping[str, Any], device: str) -> Any: + torch, F, np, image, image_filter, sparse = runtime["torch"], runtime["F"], runtime["np"], runtime["Image"], runtime["ImageFilter"], runtime["sparse"] + coords = surface.coords[:, 1:].long(); points = (coords.float() + .5) / RESOLUTION - .5 + camera = (torch.cat((points, torch.ones_like(points[:, :1])), dim=1) @ external.T)[:, :3] + uvw = camera @ projection.T; uv = uvw[:, :2] / uvw[:, 2:3].clamp(min=1e-8) + grid = torch.stack((uv[:, 0] * 2 - 1, uv[:, 1] * 2 - 1), dim=1).view(1, 1, -1, 2) + def sampled(source: Any) -> Any: + data = np.asarray(source.convert("RGB").resize((RESOLUTION, RESOLUTION), image.Resampling.LANCZOS), dtype=np.float32).copy() / 255. + return F.grid_sample(torch.from_numpy(data).permute(2, 0, 1).to(device)[None], grid, mode="bilinear", align_corners=False, padding_mode="border")[0, :, 0].T.clamp(0, 1) + alpha = np.asarray(case["rgba"].getchannel("A"), dtype=np.uint8) + mask = image.fromarray(np.where(alpha >= 254, 255, 0).astype(np.uint8)).resize((RESOLUTION, RESOLUTION), image.Resampling.NEAREST).filter(image_filter.MinFilter(5)) + trusted = torch.from_numpy(np.asarray(mask, dtype=np.uint8).copy() == 255) + bits, ordered, order = native_observation.build_occupancy_index(coords) + admission: dict[str, int] = {} + owners = native_observation.pixel_first_occupied(dict(case["camera"]), bits, trusted, stats=admission) + _alignment_report(uv, alpha, admission, np, image, image_filter) + visible = native_observation.owner_rows(coords, owners, ordered, order) + if not bool(visible.any()): raise RuntimeError("the trusted alpha gate leaves no observed material surface") + with torch.no_grad(): + return native_observation.encode_partial(encoder, coords, visible, sampled(case["albedo"]), sampled(case["mra"]), shape.coords[:, 1:], dict(normalization), sparse, device) + + +def _export(output: Path, representation: Any, texture_size: int, decimation_target: int, texture_format: str, runtime: Mapping[str, Any], *, gutter: bool = False, low_vram: bool = False) -> None: + if output.suffix.lower() != ".glb": raise ValueError("output must have a .glb suffix") + output.parent.mkdir(parents=True, exist_ok=True) + lock = output.with_name(f".{output.name}.publish.lock") + try: + handle = lock.open("x") + except FileExistsError as error: + raise FileExistsError(f"another process is publishing {output}") from error + handle.close() + if output.exists(): + lock.unlink(missing_ok=True) + raise FileExistsError(f"refusing to overwrite existing output: {output}") + temporary = output.with_name(f".{output.stem}.tmp.{os.getpid()}.glb") + try: + scene = runtime["to_glb"](vertices=representation.vertices, faces=representation.faces, attr_volume=representation.attrs, coords=representation.coords, attr_layout=LAYOUT, grid_size=RESOLUTION, aabb=[[-.5, -.5, -.5], [.5, .5, .5]], decimation_target=decimation_target, texture_size=texture_size, remesh=True, remesh_band=1, remesh_project=0, use_tqdm=True, **({"low_vram": True} if low_vram else {})) + if gutter: + from .gutter import apply_gutter + print(f"Gutter postprocess: {apply_gutter(scene)}", flush=True) + scene.apply_transform(runtime["np"].asarray([[-1, 0, 0, 0], [0, 0, -1, 0], [0, -1, 0, 0], [0, 0, 0, 1]], dtype=runtime["np"].float64)) + scene.export(str(temporary), extension_webp=texture_format == "webp") + if output.exists(): + raise FileExistsError(f"refusing to overwrite existing output: {output}") + os.replace(temporary, output) + finally: + temporary.unlink(missing_ok=True) + lock.unlink(missing_ok=True) + + +def _stage_sparse(value: Any, device: str, *, keep_topology: bool = False) -> Any: + """Transfer sparse data without retaining backend caches on its old device. + + SparseTensor.to()/cpu() share spatial/backend caches. Encoder subdivision + maps are semantic: the decoders use them to reproduce the input surface, + overriding predicted subdivisions. Preserve those maps when requested; + convolution and attention caches are safe to recompute on the new device. + """ + def transfer(item: Any) -> Any: + if isinstance(item, dict): + return {key: transfer(child) for key, child in item.items()} + if isinstance(item, (list, tuple)): + return type(item)(transfer(child) for child in item) + if hasattr(item, "detach") and hasattr(item, "to"): + return item.detach().to(device) + return item + staged = type(value)(value.feats.detach().to(device), value.coords.detach().to(device), + shape=value._shape, scale=value._scale) + if keep_topology: + staged._spatial_cache = { + scale: {name: transfer(item) for name, item in cache.items() + if name.startswith(("channel2spatial_", "upsample_")) or name == "shape"} + for scale, cache in value._spatial_cache.items() + } + return staged + + +@contextmanager +def _completion_stage(model: Any, torch: Any, device: str, *, layer_offload: bool = False): + """Own one model's residency, including cleanup after failed inference.""" + previous = getattr(model, "low_vram", False) + try: + if layer_offload: + model.low_vram = True + else: + model.to(device) + model.eval() + with torch.no_grad(): + yield model + finally: + model.to("cpu") + if layer_offload: + model.low_vram = previous + torch.cuda.empty_cache() + + +def _run_completion_low_vram(*, case: Mapping[str, Any], back_case: Mapping[str, Any] | None, + ckpts: Path, config: Mapping[str, Any], dataset: Mapping[str, Any], + checkpoint_path: Path, partial_stats: Mapping[str, Any], resolved: Any, + runtime: Mapping[str, Any], device: str, seed: int, output: str | Path, + texture_size: int, decimation_target: int, texture_format: str, gutter: bool, preview_only: bool = False) -> dict[str, Any]: + """Serialize completion stages and keep only their CPU handoffs alive.""" + torch, models = runtime["torch"], runtime["models"] + print("Completion low VRAM: sequential models, CPU handoffs, decoder layer offload", flush=True) + shape_encoder = _load_vae(models, _stem(ckpts / "shape_enc_next_dc_f16c32_fp16", "shape encoder")) + with _completion_stage(shape_encoder, torch, device): + shape, surface = _shape(case["vertices"], case["faces"], shape_encoder, runtime, device) + shape_cpu = _stage_sparse(shape, "cpu", keep_topology=True) + shape = _stage_sparse(shape, device) + del shape_encoder + texture_encoder = _load_vae(models, _stem(ckpts / "tex_enc_next_dc_f16c32_fp16", "material encoder")) + dual_stats = {} + with _completion_stage(texture_encoder, torch, device): + if back_case is None: + external, projection = _camera(case["camera"], device, torch, runtime["utils3d"]) + observation = _partial(surface, shape, case, external, projection, texture_encoder, partial_stats, runtime, device) + del external, projection + else: + from .dual_view import encode_two_views + observation, dual_stats = encode_two_views(surface, shape, case, back_case, texture_encoder, + partial_stats, runtime, device) + del texture_encoder, surface + torch.cuda.empty_cache() + + from supratex.pipelines.observation_guided_completion import build_completion_model, sample_material_latents + state = load_completion_state_dict(checkpoint_path, torch_module=torch) + model = build_completion_model(config, state, assign_weights=True) + del state + with _completion_stage(model, torch, device): + shape_condition = shape.replace(feats=(shape.feats - torch.as_tensor(dataset["shape_slat_normalization"]["mean"], device=device)[None]) / torch.as_tensor(dataset["shape_slat_normalization"]["std"], device=device)[None]) + noise = shape.replace(feats=torch.randn((shape.coords.shape[0], 32), generator=torch.Generator(device=device).manual_seed(int(seed)), device=device)) + print("Observation cache: disabled (low VRAM)", flush=True) + sampled = sample_material_latents(model, noise, concat_cond=shape_condition, partial_z_obs=observation, steps=25, rescale_t=3., sigma_min=1e-5, use_observation_cache=False).samples + if not torch.equal(sampled.coords, noise.coords) or sampled.feats.shape != noise.feats.shape: + raise RuntimeError("material sampler changed sparse support or feature shape") + generated = shape.replace(feats=sampled.feats.float() * torch.as_tensor(dataset["pbr_slat_normalization"]["std"], device=device)[None] + torch.as_tensor(dataset["pbr_slat_normalization"]["mean"], device=device)[None]) + generated_cpu = _stage_sparse(generated, "cpu") + del model, shape, shape_condition, noise, sampled, generated, observation + torch.cuda.empty_cache() + + # Reuse the encoder's exact support in both decoders. Do not carry GPU + # convolution caches or the unused decoded mesh into material decoding. + shape_decoder = _load_vae(models, _stem(resolved(dataset["pretrained_shape_slat_dec"]), "shape decoder")) + shape_decoder.set_resolution(RESOLUTION) + with _completion_stage(shape_decoder, torch, device, layer_offload=True): + shape = _stage_sparse(shape_cpu, device, keep_topology=True) + decoded_shape, guide = shape_decoder(shape, return_subs=True, return_mesh=False) + guide_cpu = [_stage_sparse(value, "cpu") for value in guide] + del shape_decoder, shape, decoded_shape, guide + generated_cpu._spatial_cache = shape_cpu._spatial_cache + del shape_cpu + torch.cuda.empty_cache() + + material_decoder = _load_vae(models, _stem(resolved(dataset["pretrained_pbr_slat_dec"]), "material decoder")) + with _completion_stage(material_decoder, torch, device, layer_offload=True): + generated = _stage_sparse(generated_cpu, device, keep_topology=True) + guide = [_stage_sparse(value, device) for value in guide_cpu] + decoded = material_decoder(generated, guide_subs=guide) * .5 + .5 + if decoded.feats.shape[1] != 6: raise RuntimeError("material decoder must emit six PBR channels") + decoded_cpu = _stage_sparse(decoded, "cpu") + del material_decoder, generated, generated_cpu, guide, guide_cpu, decoded + torch.cuda.empty_cache() + representation = runtime["MeshWithVoxel"](case["vertices"], case["faces"], [-.5, -.5, -.5], 1. / RESOLUTION, decoded_cpu.coords[:, 1:].int(), decoded_cpu.feats.float().clamp(0, 1), torch.Size([1, 6, RESOLUTION, RESOLUTION, RESOLUTION]), layout=LAYOUT) + del decoded_cpu + if preview_only: + from hf_space.ovoxel_preview import save_and_render + save_and_render(Path(output), representation.to(device)) + else: + _export(Path(output), representation, texture_size, decimation_target, texture_format, runtime, gutter=gutter, low_vram=True) + return dual_stats + + +def _run_completion(*, input_dir: str | Path, output: str | Path, model_config: str | Path, checkpoint: str | Path, trellis_root: str | Path, seed: int = 42, device: str = "cuda", texture_size: int = 4096, decimation_target: int = 500000, texture_format: str = "png", gutter: bool = False, back_case: Mapping[str, Any] | None = None, remove_floaters: bool = False, preview_only: bool = False) -> dict[str, Any]: + """Complete PBR texture from public input maps and write one new GLB.""" + if device != "cuda": raise ValueError("public material completion currently requires device='cuda'") + if texture_size <= 0 or decimation_target <= 0 or texture_format not in {"png", "webp"}: raise ValueError("invalid GLB export settings") + config, dataset, checkpoint_path, partial_stats = validate_completion_artifacts(model_config, checkpoint, trellis_root) + decoder_stems = completion_decoder_stems(config, dataset, model_config, trellis_root) + decoder_paths = {dataset[key]: path for key, path in decoder_stems.items()} + def resolved(value: str) -> Path: + return decoder_paths[value] + runtime = _runtime(); torch, models = runtime["torch"], runtime["models"] + if not torch.cuda.is_available(): raise RuntimeError("CUDA is unavailable") + case = _load_case(input_dir, torch, runtime["Image"], remove_floaters); ckpts = _completion_vae_root(trellis_root) + if os.environ.get("SUPRATEX_LOW_VRAM", "0") == "1": + return _run_completion_low_vram(case=case, back_case=back_case, ckpts=ckpts, config=config, + dataset=dataset, checkpoint_path=checkpoint_path, partial_stats=partial_stats, + resolved=resolved, runtime=runtime, device=device, seed=seed, output=output, + texture_size=texture_size, decimation_target=decimation_target, + texture_format=texture_format, gutter=gutter, preview_only=preview_only) + shape_encoder = _load_vae(models, _stem(ckpts / "shape_enc_next_dc_f16c32_fp16", "shape encoder")).to(device).eval() + texture_encoder = _load_vae(models, _stem(ckpts / "tex_enc_next_dc_f16c32_fp16", "material encoder")).to(device).eval() + external, projection = _camera(case["camera"], device, torch, runtime["utils3d"]) + shape, surface = _shape(case["vertices"], case["faces"], shape_encoder, runtime, device) + shape_encoder.to("cpu"); torch.cuda.empty_cache() + dual_stats = {} + if back_case is None: + observation = _partial(surface, shape, case, external, projection, texture_encoder, partial_stats, runtime, device) + else: + from .dual_view import encode_two_views + observation, dual_stats = encode_two_views(surface, shape, case, back_case, texture_encoder, + partial_stats, runtime, device) + texture_encoder.to("cpu"); torch.cuda.empty_cache() + from supratex.pipelines.observation_guided_completion import build_completion_model, sample_material_latents + state = load_completion_state_dict(checkpoint_path, torch_module=torch) + model = build_completion_model(config, state, assign_weights=True).to(device).eval() + del state + shape_condition = shape.replace(feats=(shape.feats - torch.as_tensor(dataset["shape_slat_normalization"]["mean"], device=device)[None]) / torch.as_tensor(dataset["shape_slat_normalization"]["std"], device=device)[None]) + noise = shape.replace(feats=torch.randn((shape.coords.shape[0], 32), generator=torch.Generator(device=device).manual_seed(int(seed)), device=device)) + # Cache fixed observation work within this sample only. Retain ample room + # for attention activations; low-memory devices keep the uncached path. + cache_bytes = 2 * (len(model.blocks) + len(model.refiners)) * observation.feats.shape[0] * model.model_channels * 2 + free_bytes, _ = torch.cuda.mem_get_info(noise.feats.device) + use_cache = free_bytes >= cache_bytes * 1.25 + 8 * 1024 ** 3 + print(f"Observation cache: {'enabled' if use_cache else 'disabled (memory headroom)'}", flush=True) + sampled = sample_material_latents(model, noise, concat_cond=shape_condition, partial_z_obs=observation, steps=25, rescale_t=3., sigma_min=1e-5, use_observation_cache=use_cache).samples + if not torch.equal(sampled.coords, noise.coords) or sampled.feats.shape != noise.feats.shape: + raise RuntimeError("material sampler changed sparse support or feature shape") + generated = shape.replace(feats=sampled.feats.float() * torch.as_tensor(dataset["pbr_slat_normalization"]["std"], device=device)[None] + torch.as_tensor(dataset["pbr_slat_normalization"]["mean"], device=device)[None]) + model.to("cpu"); torch.cuda.empty_cache() + shape_stem, material_stem = _stem(resolved(dataset["pretrained_shape_slat_dec"]), "shape decoder"), _stem(resolved(dataset["pretrained_pbr_slat_dec"]), "material decoder") + shape_decoder, material_decoder = _load_vae(models, shape_stem).to(device).eval(), _load_vae(models, material_stem).to(device).eval() + shape_decoder.set_resolution(RESOLUTION) + with torch.no_grad(): _, guide = shape_decoder(shape, return_subs=True); decoded = material_decoder(generated, guide_subs=guide) * .5 + .5 + if decoded.feats.shape[1] != 6: raise RuntimeError("material decoder must emit six PBR channels") + representation = runtime["MeshWithVoxel"](case["vertices"].to(device), case["faces"].to(device), [-.5, -.5, -.5], 1. / RESOLUTION, decoded.coords[:, 1:].int(), decoded.feats.float().clamp(0, 1), torch.Size([1, 6, RESOLUTION, RESOLUTION, RESOLUTION]), layout=LAYOUT) + if preview_only: + from hf_space.ovoxel_preview import save_and_render + save_and_render(Path(output), representation.to(device)) + else: + _export(Path(output), representation, texture_size, decimation_target, texture_format, runtime, gutter=gutter) + return dual_stats + + +def _require_reprojection_atlas(texture_size: int, texture_format: str) -> None: + from supratex.pipeline_config import REFINE_ATLAS + if {"texture_size": texture_size, "texture_format": texture_format} != REFINE_ATLAS: + raise ValueError("--refine and --dual-view require a 4096 PNG atlas") + + +def _publish(final: Path, output: Path, tag: str) -> None: + """Copy a finished mesh to ``output`` without replacing another writer's file.""" + import shutil + output.parent.mkdir(parents=True, exist_ok=True) + temporary = output.with_name(f".{output.name}.{tag}.{os.getpid()}.tmp") + try: + with temporary.open("xb") as target, final.open("rb") as source: + shutil.copyfileobj(source, target) + os.link(temporary, output) + finally: + temporary.unlink(missing_ok=True) + + +def _publish_unrefined(native: Path, output: Path, gutter: bool, error: Exception) -> None: + """Publish the completed mesh without refinement after a refinement failure.""" + print(f"WARNING: texture refinement failed and was skipped ({type(error).__name__}: {error}); " + f"writing the unrefined result to {output}", file=sys.stderr, flush=True) + if not gutter: + _publish(native, output, "unrefined") + return + # Refinement normally supplies the gutter; apply the export-time one instead. + import tempfile + import trimesh + from .gutter import apply_gutter + scene = trimesh.load(native, force="scene") + for mesh in scene.geometry.values(): + apply_gutter(mesh) + with tempfile.TemporaryDirectory(dir=output.parent) as directory: + padded = Path(directory) / "mesh.glb" + scene.export(str(padded)) + _publish(padded, output, "unrefined") + + +def _run_front_reprojection(arguments: dict[str, Any], work_dir: str | Path | None) -> None: + """Complete one view, then reproject its reference observation into the atlas.""" + from .dual_view import digest, write_json + + output = Path(arguments["output"]).expanduser().resolve() + work = Path(work_dir).expanduser().resolve() if work_dir else output.with_name(output.stem + "_refine") + if output.suffix.lower() != ".glb": + raise ValueError("output must have a .glb suffix") + if output.exists(): + raise FileExistsError(f"refusing to overwrite existing output: {output}") + if work.exists(): + raise FileExistsError(f"--refine work directory must be new: {work}") + _require_reprojection_atlas(arguments["texture_size"], arguments["texture_format"]) + source = Path(arguments["input_dir"]).expanduser().resolve() + work.mkdir(parents=True) + # Reprojection always ends with its observation-guided one-texel gutter, so + # the export-time nearest-color gutter (--gutter) is not applied first. + receipt = {"profile": "supratex.front_refinement", "view": "front", "status": "running", "stage": "texture", + "input": str(source), "output": str(output), "seed": arguments["seed"], + "decimation_target": arguments["decimation_target"], "gutter": "continuous (reprojection)", + "checkpoint": str(Path(arguments["checkpoint"]).expanduser().resolve()), + "model_config": str(Path(arguments["model_config"]).expanduser().resolve())} + write_json(work / "receipt.json", receipt) + try: + native = work / "preds/mesh.glb" + _run_completion(**{**arguments, "input_dir": source, "output": native, "gutter": False}) + # Reprojection runs in separate CUDA workers; release cached blocks first. + import gc + import torch + gc.collect() + torch.cuda.empty_cache() + receipt.update(stage="reprojection", native_mesh=str(native), native_mesh_sha256=digest(native)) + write_json(work / "receipt.json", receipt) + from .reprojection import run_front_view + try: + final = run_front_view(native, source, work / "reproject") + except Exception as error: + _publish_unrefined(native, output, arguments["gutter"], error) + receipt.update(status="valid", stage="complete", refine="skipped", + refine_error=f"{type(error).__name__}: {error}", mesh=str(native), + gutter="export-time" if arguments["gutter"] else False, + output_sha256=digest(output)) + write_json(work / "receipt.json", receipt) + return + _publish(final, output, "reproject") + receipt.update(status="valid", stage="complete", refine="applied", mesh=str(final), + output_sha256=digest(output)) + write_json(work / "receipt.json", receipt) + except Exception as error: + receipt.update(status="failed", error=f"{type(error).__name__}: {error}") + write_json(work / "receipt.json", receipt) + raise + + +def run_completion(*, input_dir: str | Path, output: str | Path, model_config: str | Path, + checkpoint: str | Path, trellis_root: str | Path, seed: int = 42, + device: str = "cuda", texture_size: int = 4096, + decimation_target: int = 500000, texture_format: str = "png", gutter: bool = False, + back_view_dir: str | Path | None = None, work_dir: str | Path | None = None, + refine: bool = False, remove_floaters: bool = False, preview_only: bool = False) -> None: + """Complete one view by default, optionally refining it, or publish a refined dual-view result.""" + arguments = dict(input_dir=input_dir, output=output, model_config=model_config, + checkpoint=checkpoint, trellis_root=trellis_root, seed=seed, device=device, + texture_size=texture_size, decimation_target=decimation_target, + texture_format=texture_format, gutter=gutter, remove_floaters=remove_floaters, preview_only=preview_only) + if preview_only and (back_view_dir is not None or refine): + raise ValueError("o-voxel preview must precede single-view mesh export") + # Dual-view refines both views itself, so --refine needs no separate pass there. + if refine and back_view_dir is None: + _run_front_reprojection(arguments, work_dir) + return + if back_view_dir is None: + if work_dir is not None: + raise ValueError("completion --work-dir is reserved for dual-view inference or --refine") + _run_completion(**arguments) + return + from .dual_view import CAMERA_PROTOCOL, PROTOCOL, digest, load_back_case, write_json + + output = Path(output).expanduser().resolve() + work = Path(work_dir).expanduser().resolve() if work_dir else output.with_name(output.stem + "_dual_view") + if output.suffix.lower() != ".glb": + raise ValueError("output must have a .glb suffix") + if output.exists(): + raise FileExistsError(f"refusing to overwrite existing output: {output}") + if work.exists(): + raise FileExistsError(f"dual-view work directory must be new: {work}") + _require_reprojection_atlas(texture_size, texture_format) + front_root, back_root = Path(input_dir).expanduser().resolve(), Path(back_view_dir).expanduser().resolve() + back = load_back_case(front_root, back_root) + work.mkdir(parents=True) + receipt = {"protocol": PROTOCOL, "camera_protocol": CAMERA_PROTOCOL, "view": "dual", + "status": "running", "stage": "texture", "front_input": str(front_root), + "back_input": str(back_root), "output": str(output), + "checkpoint": str(Path(checkpoint).expanduser().resolve()), + "model_config": str(Path(model_config).expanduser().resolve()), + "seed": seed, "sampler": "Euler", "steps": 25, + "decimation_target": decimation_target, "gutter": "continuous (reprojection)", + "geometry_sha256": digest(front_root / "geometry.pt"), + "cameras_sha256": digest(back_root / "cameras.json"), + "front_material_receipt_sha256": digest(front_root / "material_receipt.json"), + "back_material_receipt_sha256": digest(back_root / "material_receipt.json")} + write_json(work / "receipt.json", receipt) + try: + native = work / "preds/mesh.glb" + # Reprojection ends with its observation-guided gutter, so skip the export-time one. + arguments.update(input_dir=front_root, output=native, back_case=back, gutter=False) + stats = _run_completion(**arguments) + # Reprojection runs in separate CUDA workers. Release the completion + # allocator's inactive blocks before those processes create contexts. + import gc + import torch + gc.collect() + torch.cuda.empty_cache() + write_json(work / "receipt_texture.json", {**receipt, "status": "success", + "dual_stats": stats, "sampling_calls": 1, "mesh": str(native), "mesh_sha256": digest(native)}) + receipt.update(stage="reprojection") + write_json(work / "receipt.json", receipt) + from .reprojection import run_dual_view, validate_dual_view + try: + final = run_dual_view(native, front_root, back_root, back["cameras"], work) + validation = validate_dual_view(work, raise_on_invalid=True) + except Exception as error: + final, validation, refine_error = None, None, error + # Model and reprojection work can be long-running. Revalidate the + # original inputs and generation evidence before publishing the mesh. + load_back_case(front_root, back_root) + bindings = {"geometry_sha256": front_root / "geometry.pt", + "cameras_sha256": back_root / "cameras.json", + "front_material_receipt_sha256": front_root / "material_receipt.json", + "back_material_receipt_sha256": back_root / "material_receipt.json"} + for key, path in bindings.items(): + if receipt[key] != digest(path): + raise RuntimeError(f"input changed during dual-view completion: {key}") + if final is None: + _publish_unrefined(native, output, gutter, refine_error) + receipt.update(status="valid", stage="complete", refine="skipped", + refine_error=f"{type(refine_error).__name__}: {refine_error}", mesh=str(native), + gutter="export-time" if gutter else False, + mesh_sha256=digest(native), output_sha256=digest(output), dual_stats=stats) + else: + _publish(final, output, "dual-view") + receipt.update(status="valid", stage="complete", refine="applied", mesh=str(final), + mesh_sha256=digest(final), output_sha256=digest(output), dual_stats=stats, + postprocess_validation=validation, stitched_19_phase_checked=False) + write_json(work / "receipt.json", receipt) + except Exception as error: + receipt.update(status="failed", error=f"{type(error).__name__}: {error}") + write_json(work / "receipt.json", receipt) + raise diff --git a/supratex/inference/reprojection/__init__.py b/supratex/inference/reprojection/__init__.py index 3f040d463a96f6cd229679df0238a4bd58d54dd8..fc156e56b4a29d143cbda8b71d54c8868d33989c 100644 --- a/supratex/inference/reprojection/__init__.py +++ b/supratex/inference/reprojection/__init__.py @@ -1,286 +1,329 @@ -"""Dual-view atlas reprojection with measured candidate acceptance gates. - -Both view candidates undergo sixteen train and three held-out phase renders. -The final atlas is stitched with exact front priority and binary preservation; -the stitched result is not itself subjected to another nineteen render phases. -""" -from __future__ import annotations - -import json -import shutil -import subprocess -import sys -from pathlib import Path -from typing import Any - -from . import _glb -from ._common import ( - bind_file, check_binding, read_json, validate_candidate, write_json, -) -from ._transform import ( - atlas, back_to_front, camera_spec, front_protection, stitch, - transform_glb, transformed_payload, validate_stitch, -) - -__all__ = ["run_dual_view", "validate_dual_view"] -PROFILE = "supratex.dual_view.reproject_continuous" - - -def _prepare_inputs(source: Path, destination: Path, camera: dict[str, Any]) -> dict[str, Any]: - """Align observation alpha and RGB once, preserving native albedo/MRA maps.""" - from PIL import Image - - names = ("condition_rgb.png", "condition_rgba.png", "albedo.png", "mra.png") - bindings = {name: bind_file(source / name) for name in names} - if (source / "camera.json").is_file(): - bindings["camera.json"] = bind_file(source / "camera.json") - destination.mkdir(parents=True) - with Image.open(source / "condition_rgb.png") as rgb, Image.open(source / "condition_rgba.png") as rgba: - if rgb.size != rgba.size: - raise ValueError("condition RGB and RGBA dimensions must agree") - resized = rgb.convert("RGB").resize((1024, 1024), Image.Resampling.LANCZOS) - alpha = rgba.convert("RGBA").getchannel("A").point(lambda value: 255 if value >= 254 else 0) - alpha = alpha.resize((1024, 1024), Image.Resampling.NEAREST) - resized.save(destination / "condition_rgb.png") - aligned = resized.copy() - aligned.putalpha(alpha) - aligned.save(destination / "condition_rgba.png") - for name in ("albedo.png", "mra.png"): - with Image.open(source / name) as image: - if image.size != (1024, 1024): - raise ValueError(f"{name} must have native dimensions 1024 by 1024") - image.convert("RGB").save(destination / name) - write_json(destination / "camera.json", camera) - return bindings - - -def _command(module: str, arguments: list[str], log: Path) -> None: - command = [sys.executable, "-m", f"supratex.inference.reprojection.{module}", *arguments] - with log.open("x", encoding="utf-8") as stream: - stream.write(json.dumps(command) + "\n") - stream.flush() - # Share the caller's process group so cancellation reaches GPU workers. - process = subprocess.Popen(command, stdout=stream, stderr=subprocess.STDOUT) - try: - result = process.wait() - except BaseException: - process.terminate() - process.wait() - raise - if result: - raise RuntimeError(f"reprojection worker failed with exit code {result}; see {log}") - - -def _run_side(baseline: Path, source: Path, output: Path, camera: dict[str, Any]) -> Path: - """Publish one view only after the B and continuous candidates pass.""" - output.mkdir(parents=True) - baseline_binding = bind_file(baseline) - code = {path.name: bind_file(path) for path in Path(__file__).parent.glob("*.py")} - inputs = output / "inputs_1024" - source_bindings = _prepare_inputs(source, inputs, camera) - input_bindings = {path.name: bind_file(path) for path in inputs.iterdir() if path.is_file()} - renderer = Path(__file__).with_name("_renderer.py").resolve() - project, continuous = output / "B", output / "continuous" - _command("_project", [ - "--arm", "B", "--baseline-glb", str(baseline), - "--source-albedo", str(inputs / "albedo.png"), - "--source-rgba", str(inputs / "condition_rgba.png"), - "--camera", str(inputs / "camera.json"), "--evaluator", str(renderer), - "--output", str(project), - ], output / "B.log") - validate_candidate(project, continuous=False) - manifest = output / "continuous_inputs.json" - write_json(manifest, [{"label": "observation", "baseline_glb": str(baseline), - "work": str(inputs), "b_output": str(project)}]) - _command("_continuous", [ - "--manifest", str(manifest), "--case", "observation", - "--evaluator", str(renderer), "--output", str(continuous), - ], output / "continuous.log") - validate_candidate(continuous, continuous=True) - for name, record in source_bindings.items(): - check_binding(record, source / name) - check_binding(baseline_binding, baseline) - for name, record in code.items(): - check_binding(record, Path(__file__).with_name(name)) - final = output / "mesh.glb" - shutil.copyfile(continuous / "candidate" / "mesh.glb", final) - artifacts = {str(path.relative_to(output)): bind_file(path) for path in ( - project / "PLAN.json", project / "RECEIPT.json", project / "RESULT.json", - project / "COMMON_DOMAIN.npz", project / "candidate" / "mesh.glb", - continuous / "PLAN.json", continuous / "RECEIPT.json", continuous / "RESULT.json", - continuous / "GUTTER_DIAGNOSTICS.npz", continuous / "candidate" / "mesh.glb", - manifest, final, - )} - receipt = { - "profile": PROFILE, "status": "valid", "baseline": baseline_binding, - "source_inputs": source_bindings, "aligned_inputs": input_bindings, - "code": code, "artifacts": artifacts, "mesh": str(final), - "gates": {"accepted": True, "B_accepted_all_19": True, "continuous_accepted_all_19": True}, - } - write_json(output / "receipt.json", receipt) - _validate_side(output) - return final - - -def _validate_side(output: Path) -> dict[str, Any]: - import numpy as np - - receipt = read_json(output / "receipt.json") - if receipt.get("profile") != PROFILE or receipt.get("status") != "valid": - raise ValueError("invalid reprojection receipt profile or status") - if not all(receipt.get("gates", {}).get(key) is True for key in ("accepted", "B_accepted_all_19", "continuous_accepted_all_19")): - raise ValueError("reprojection side is missing mandatory gates") - baseline = check_binding(receipt["baseline"]) - required_inputs = {"condition_rgb.png", "condition_rgba.png", "albedo.png", "mra.png"} - if not required_inputs.issubset(receipt.get("source_inputs", {})): - raise ValueError("reprojection side has incomplete source bindings") - if set(receipt.get("aligned_inputs", {})) != required_inputs | {"camera.json"}: - raise ValueError("reprojection side has incomplete aligned input bindings") - for record in receipt["source_inputs"].values(): - check_binding(record) - for name, record in receipt["aligned_inputs"].items(): - check_binding(record, output / "inputs_1024" / name) - expected_code = {path.name for path in Path(__file__).parent.glob("*.py")} - if set(receipt.get("code", {})) != expected_code: - raise ValueError("reprojection code bindings are incomplete") - for name, record in receipt["code"].items(): - check_binding(record, Path(__file__).with_name(name)) - expected_artifacts = { - "B/PLAN.json", "B/RECEIPT.json", "B/RESULT.json", "B/COMMON_DOMAIN.npz", "B/candidate/mesh.glb", - "continuous/PLAN.json", "continuous/RECEIPT.json", "continuous/RESULT.json", - "continuous/GUTTER_DIAGNOSTICS.npz", "continuous/candidate/mesh.glb", "continuous_inputs.json", "mesh.glb", - } - if set(receipt.get("artifacts", {})) != expected_artifacts: - raise ValueError("reprojection artifact bindings are incomplete") - for name, record in receipt["artifacts"].items(): - check_binding(record, output / name) - if Path(receipt.get("mesh", "")).resolve() != (output / "mesh.glb").resolve(): - raise ValueError("reprojection output mesh path changed") - for name, continuous in (("B", False), ("continuous", True)): - validate_candidate(output / name, continuous=continuous) - bplan = read_json(output / "B" / "PLAN.json") - for key, path in {"baseline_glb": baseline, "source_albedo": output / "inputs_1024" / "albedo.png", - "source_rgba": output / "inputs_1024" / "condition_rgba.png", - "camera": output / "inputs_1024" / "camera.json"}.items(): - if Path(bplan.get(key, "")).resolve() != path.resolve(): - raise ValueError(f"B plan has an incorrect {key} binding") - base_doc, base_binary, base_view, _, _, base_pixels = atlas(baseline) - for name in ("B/candidate/mesh.glb", "continuous/candidate/mesh.glb", "mesh.glb"): - candidate = output / name - _glb._nonbase_identity_check(base_doc, base_binary, candidate, base_view) - pixels = atlas(candidate)[-1] - if pixels.shape != base_pixels.shape or (pixels.shape[-1] == 4 and not np.array_equal(pixels[..., 3], base_pixels[..., 3])): - raise ValueError("candidate altered atlas dimensions or alpha") - if bind_file(output / "mesh.glb")["sha256"] != bind_file(output / "continuous/candidate/mesh.glb")["sha256"]: - raise ValueError("side output differs from its accepted continuous candidate") - return receipt - - -def run_dual_view(native_glb: Path, front_input_dir: Path, back_input_dir: Path, - cameras: dict[str, Any], output_dir: Path) -> Path: - """Run both measured reprojections and publish the stitched final GLB. - - The caller owns geometry, generated-image, and material-model provenance. - This function binds the provided model-independent maps, camera pack, native - mesh, operator code, candidate gates, and final atlas in its own receipt. - """ - native, front, back = (Path(path).resolve(strict=True) for path in (native_glb, front_input_dir, back_input_dir)) - output = Path(output_dir).resolve() - raw = cameras.get("raw", cameras) - front_camera, back_camera = camera_spec(raw, "front"), camera_spec(raw, "back") - for path in (output / "post", output / "post_back"): - if path.exists(): - raise FileExistsError(f"dual-view reprojection requires a new destination: {path}") - output.mkdir(parents=True, exist_ok=True) - final_dir = output / "post_back" - final_dir.mkdir() - try: - write_json(final_dir / "cameras.json", raw) - front_mesh = _run_side(native, front, output / "post", front_camera) - rotation, translation, distance = back_to_front(raw["back"]) - rotated = final_dir / "rotated_baseline.glb" - transform_glb(front_mesh, rotated, rotation, translation) - canonical_back = { - "extrinsic": [[1, 0, 0, 0], [0, -1, 0, 0], [0, 0, -1, distance], [0, 0, 0, 1]], - "projection_intrinsic": back_camera["projection_intrinsic"], "distance": distance, - "camera_angle_x": back_camera["camera_angle_x"], "mesh_scale": 1.0, - } - back_mesh = _run_side(rotated, back, final_dir / "reproject", canonical_back) - protected = front_protection(output / "post" / "B" / "COMMON_DOMAIN.npz") - final = final_dir / "mesh.glb" - stitch_gates = stitch(front_mesh, back_mesh, protected, final) - _validate_side(output / "post") - _validate_side(final_dir / "reproject") - receipt = { - "profile": PROFILE, "status": "valid", "view": "dual", "mesh": str(final), - "gates": {"accepted": True, "front_candidates_accepted_all_19": True, - "back_candidates_accepted_all_19": True, **stitch_gates}, - "front_post_receipt": bind_file(output / "post" / "receipt.json"), - "back_post_receipt": bind_file(final_dir / "reproject" / "receipt.json"), - "camera_pack": bind_file(final_dir / "cameras.json"), "final_mesh": bind_file(final), - "front_protection": bind_file(output / "post" / "B" / "COMMON_DOMAIN.npz"), - "back_transform": {"rotation": rotation.tolist(), "translation": translation.tolist(), "distance": distance}, - "validation_scope": "Both view candidates passed 19-phase gates; the stitched atlas has exact front priority and preserved non-base content.", - } - write_json(final_dir / "receipt.json", receipt) - validate_dual_view(output, raise_on_invalid=True) - return final - except BaseException as error: - failure = final_dir / "FAILED.json" - if not failure.exists(): - write_json(failure, {"profile": PROFILE, "status": "failed", "error": f"{type(error).__name__}: {error}"}) - raise - - -def validate_dual_view(output_dir: Path, *, raise_on_invalid: bool = False) -> dict[str, Any]: - """Re-read the complete postprocess receipt chain and content bindings.""" - try: - output = Path(output_dir).resolve() - receipt = read_json(output / "post_back" / "receipt.json") - if (output / "post_back" / "FAILED.json").exists(): - raise ValueError("dual-view reprojection has a recorded failure") - if receipt.get("profile") != PROFILE or receipt.get("status") != "valid" or receipt.get("view") != "dual": - raise ValueError("invalid dual-view reprojection identity or status") - gates = receipt.get("gates", {}) - required = ("accepted", "front_candidates_accepted_all_19", "back_candidates_accepted_all_19", - "front_priority_exact", "alpha_exact", "position_exact", "normal_exact", "uv_exact") - if not all(gates.get(key) is True for key in required) or gates.get("final_19_phase_rendered") is not False: - raise ValueError("dual-view reprojection gates or validation scope are incomplete") - check_binding(receipt["front_post_receipt"], output / "post" / "receipt.json") - check_binding(receipt["back_post_receipt"], output / "post_back" / "reproject" / "receipt.json") - check_binding(receipt["front_protection"], output / "post" / "B" / "COMMON_DOMAIN.npz") - camera_path = check_binding(receipt["camera_pack"], output / "post_back" / "cameras.json") - final = check_binding(receipt["final_mesh"], output / "post_back" / "mesh.glb") - if Path(receipt.get("mesh", "")).resolve() != final: - raise ValueError("dual-view final mesh path is inconsistent") - cameras = read_json(camera_path) - front_camera = camera_spec(cameras, "front") - back_camera = camera_spec(cameras, "back") - rotation, translation, distance = back_to_front(cameras["back"]) - expected_transform = {"rotation": rotation.tolist(), "translation": translation.tolist(), "distance": distance} - if receipt.get("back_transform") != expected_transform: - raise ValueError("back transform differs from the camera pack") - if read_json(output / "post" / "inputs_1024" / "camera.json") != front_camera: - raise ValueError("front postprocess did not use the bound front camera") - expected_back = { - "extrinsic": [[1, 0, 0, 0], [0, -1, 0, 0], [0, 0, -1, distance], [0, 0, 0, 1]], - "projection_intrinsic": back_camera["projection_intrinsic"], "distance": distance, - "camera_angle_x": back_camera["camera_angle_x"], "mesh_scale": 1.0, - } - if read_json(output / "post_back" / "reproject" / "inputs_1024" / "camera.json") != expected_back: - raise ValueError("back postprocess did not use the canonicalized back camera") - expected_document, expected_binary = transformed_payload(output / "post" / "mesh.glb", rotation, translation) - actual_document, actual_binary = _glb._read_glb(output / "post_back" / "rotated_baseline.glb") - if actual_document != expected_document or actual_binary != expected_binary: - raise ValueError("temporary back mesh differs from the required rigid transform") - _validate_side(output / "post") - _validate_side(output / "post_back" / "reproject") - protected = front_protection(output / "post" / "B" / "COMMON_DOMAIN.npz") - actual = validate_stitch(output / "post" / "mesh.glb", output / "post_back" / "reproject" / "mesh.glb", final, protected) - if any(gates.get(key) != value for key, value in actual.items()): - raise ValueError("stitched preservation evidence changed") - return {"valid": True, "reason": "Both candidate chains and the stitched final mesh are current", "receipt": receipt} - except (OSError, ValueError, KeyError, TypeError, RuntimeError) as error: - if raise_on_invalid: - raise - return {"valid": False, "reason": str(error)} +"""Atlas reprojection with measured candidate acceptance gates. + +Each reprojected view writes observed albedo into visible atlas texels, then +repairs one-texel UV gutters from the observed albedo gradient. Candidates +are re-rendered at nineteen validation phases (a 4x4 sub-pixel grid and three +off-grid offsets) and must pass every gate. Dual-view texturing +stitches front and back with exact front priority and binary preservation; the +stitched result is not itself subjected to another nineteen render phases. +Single-view texturing can reproject its reference view alone. +""" +from __future__ import annotations + +import json +import shutil +import subprocess +import sys +from pathlib import Path +from typing import Any + +from . import _glb +from ._common import ( + bind_file, check_binding, read_json, validate_candidate, write_json, +) +from ._transform import ( + atlas, back_to_front, camera_spec, front_protection, stitch, + transform_glb, transformed_payload, validate_stitch, +) + +__all__ = ["run_dual_view", "run_front_view", "validate_dual_view"] +PROFILE = "supratex.dual_view.reproject_continuous" + + +def _prepare_inputs(source: Path, destination: Path, camera: dict[str, Any]) -> dict[str, Any]: + """Align observation alpha and RGB once, preserving native albedo/MRA maps.""" + from PIL import Image + + names = ("condition_rgb.png", "condition_rgba.png", "albedo.png", "mra.png") + bindings = {name: bind_file(source / name) for name in names} + if (source / "camera.json").is_file(): + bindings["camera.json"] = bind_file(source / "camera.json") + destination.mkdir(parents=True) + with Image.open(source / "condition_rgb.png") as rgb, Image.open(source / "condition_rgba.png") as rgba: + if rgb.size != rgba.size: + raise ValueError("condition RGB and RGBA dimensions must agree") + resized = rgb.convert("RGB").resize((1024, 1024), Image.Resampling.LANCZOS) + alpha = rgba.convert("RGBA").getchannel("A").point(lambda value: 255 if value >= 254 else 0) + alpha = alpha.resize((1024, 1024), Image.Resampling.NEAREST) + resized.save(destination / "condition_rgb.png") + aligned = resized.copy() + aligned.putalpha(alpha) + aligned.save(destination / "condition_rgba.png") + for name in ("albedo.png", "mra.png"): + with Image.open(source / name) as image: + if image.size != (1024, 1024): + raise ValueError(f"{name} must have native dimensions 1024 by 1024") + image.convert("RGB").save(destination / name) + write_json(destination / "camera.json", camera) + return bindings + + +def _command(module: str, arguments: list[str], log: Path) -> None: + command = [sys.executable, "-m", f"supratex.inference.reprojection.{module}", *arguments] + with log.open("x", encoding="utf-8") as stream: + stream.write(json.dumps(command) + "\n") + stream.flush() + # Share the caller's process group so cancellation reaches GPU workers. + process = subprocess.Popen(command, stdout=stream, stderr=subprocess.STDOUT) + try: + result = process.wait() + except BaseException: + process.terminate() + process.wait() + raise + if result: + raise RuntimeError(f"reprojection worker failed with exit code {result}; see {log}") + + +def _run_side(baseline: Path, source: Path, output: Path, camera: dict[str, Any]) -> Path: + """Publish one view only after the projection and continuous candidates pass.""" + output.mkdir(parents=True) + baseline_binding = bind_file(baseline) + code = {path.name: bind_file(path) for path in Path(__file__).parent.glob("*.py")} + inputs = output / "inputs_1024" + source_bindings = _prepare_inputs(source, inputs, camera) + input_bindings = {path.name: bind_file(path) for path in inputs.iterdir() if path.is_file()} + renderer = Path(__file__).with_name("_renderer.py").resolve() + project, continuous = output / "project", output / "continuous" + _command("_project", [ + "--baseline-glb", str(baseline), + "--source-albedo", str(inputs / "albedo.png"), + "--source-rgba", str(inputs / "condition_rgba.png"), + "--camera", str(inputs / "camera.json"), "--renderer", str(renderer), + "--output", str(project), + ], output / "project.log") + validate_candidate(project, continuous=False) + manifest = output / "continuous_inputs.json" + write_json(manifest, [{"label": "observation", "baseline_glb": str(baseline), + "work": str(inputs), "project_output": str(project)}]) + _command("_continuous", [ + "--manifest", str(manifest), "--case", "observation", + "--renderer", str(renderer), "--output", str(continuous), + ], output / "continuous.log") + validate_candidate(continuous, continuous=True) + for name, record in source_bindings.items(): + check_binding(record, source / name) + check_binding(baseline_binding, baseline) + for name, record in code.items(): + check_binding(record, Path(__file__).with_name(name)) + final = output / "mesh.glb" + shutil.copyfile(continuous / "candidate" / "mesh.glb", final) + artifacts = {str(path.relative_to(output)): bind_file(path) for path in ( + project / "PLAN.json", project / "RECEIPT.json", project / "RESULT.json", + project / "COMMON_DOMAIN.npz", project / "candidate" / "mesh.glb", + continuous / "PLAN.json", continuous / "RECEIPT.json", continuous / "RESULT.json", + continuous / "GUTTER_DIAGNOSTICS.npz", continuous / "candidate" / "mesh.glb", + manifest, final, + )} + receipt = { + "profile": PROFILE, "status": "valid", "baseline": baseline_binding, + "source_inputs": source_bindings, "aligned_inputs": input_bindings, + "code": code, "artifacts": artifacts, "mesh": str(final), + "gates": {"accepted": True, "project_accepted_all_19": True, "continuous_accepted_all_19": True}, + } + write_json(output / "receipt.json", receipt) + _validate_side(output) + return final + + +def _validate_side(output: Path) -> dict[str, Any]: + import numpy as np + + receipt = read_json(output / "receipt.json") + if receipt.get("profile") != PROFILE or receipt.get("status") != "valid": + raise ValueError("invalid reprojection receipt profile or status") + if not all(receipt.get("gates", {}).get(key) is True for key in ("accepted", "project_accepted_all_19", "continuous_accepted_all_19")): + raise ValueError("reprojection side is missing mandatory gates") + baseline = check_binding(receipt["baseline"]) + required_inputs = {"condition_rgb.png", "condition_rgba.png", "albedo.png", "mra.png"} + if not required_inputs.issubset(receipt.get("source_inputs", {})): + raise ValueError("reprojection side has incomplete source bindings") + if set(receipt.get("aligned_inputs", {})) != required_inputs | {"camera.json"}: + raise ValueError("reprojection side has incomplete aligned input bindings") + for record in receipt["source_inputs"].values(): + check_binding(record) + for name, record in receipt["aligned_inputs"].items(): + check_binding(record, output / "inputs_1024" / name) + expected_code = {path.name for path in Path(__file__).parent.glob("*.py")} + if set(receipt.get("code", {})) != expected_code: + raise ValueError("reprojection code bindings are incomplete") + for name, record in receipt["code"].items(): + check_binding(record, Path(__file__).with_name(name)) + expected_artifacts = { + "project/PLAN.json", "project/RECEIPT.json", "project/RESULT.json", "project/COMMON_DOMAIN.npz", + "project/candidate/mesh.glb", + "continuous/PLAN.json", "continuous/RECEIPT.json", "continuous/RESULT.json", + "continuous/GUTTER_DIAGNOSTICS.npz", "continuous/candidate/mesh.glb", "continuous_inputs.json", "mesh.glb", + } + if set(receipt.get("artifacts", {})) != expected_artifacts: + raise ValueError("reprojection artifact bindings are incomplete") + for name, record in receipt["artifacts"].items(): + check_binding(record, output / name) + if Path(receipt.get("mesh", "")).resolve() != (output / "mesh.glb").resolve(): + raise ValueError("reprojection output mesh path changed") + for name, continuous in (("project", False), ("continuous", True)): + validate_candidate(output / name, continuous=continuous) + project_plan = read_json(output / "project" / "PLAN.json") + expected_plan = { + "baseline_glb": baseline, + "source_albedo": output / "inputs_1024" / "albedo.png", + "source_rgba": output / "inputs_1024" / "condition_rgba.png", + "camera": output / "inputs_1024" / "camera.json", + } + for key, path in expected_plan.items(): + if Path(project_plan.get(key, "")).resolve() != path.resolve(): + raise ValueError(f"projection plan has an incorrect {key} binding") + base_doc, base_binary, base_view, _, _, base_pixels = atlas(baseline) + for name in ("project/candidate/mesh.glb", "continuous/candidate/mesh.glb", "mesh.glb"): + candidate = output / name + _glb._nonbase_identity_check(base_doc, base_binary, candidate, base_view) + pixels = atlas(candidate)[-1] + if pixels.shape != base_pixels.shape or (pixels.shape[-1] == 4 and not np.array_equal(pixels[..., 3], base_pixels[..., 3])): + raise ValueError("candidate altered atlas dimensions or alpha") + if bind_file(output / "mesh.glb")["sha256"] != bind_file(output / "continuous/candidate/mesh.glb")["sha256"]: + raise ValueError("side output differs from its accepted continuous candidate") + return receipt + + +def run_front_view(native_glb: Path, input_dir: Path, output_dir: Path) -> Path: + """Reproject the single reference view onto a single-view result. + + This is the front half of dual-view texturing with the canonical reference + camera from ``camera.json``. Inputs are staged under ``output_dir`` so their + bindings stay valid; a missing ``condition_rgb.png`` is derived from the + RGBA observation with black premultiplication, as the frontend writes it. + """ + import numpy as np + from PIL import Image + + native, source = (Path(path).resolve(strict=True) for path in (native_glb, input_dir)) + output = Path(output_dir).resolve() + if output.exists(): + raise FileExistsError(f"front reprojection requires a new destination: {output}") + camera = read_json(source / "camera.json") + try: + camera = {"camera_angle_x": float(camera["camera_angle_x"]), "distance": float(camera["distance"]), + "mesh_scale": float(camera.get("mesh_scale", 1.0))} + except (KeyError, TypeError, ValueError) as error: + raise ValueError("front reprojection requires a canonical camera with camera_angle_x and distance") from error + staged = output / "source" + staged.mkdir(parents=True) + for name in ("condition_rgba.png", "albedo.png", "mra.png", "camera.json"): + shutil.copyfile(source / name, staged / name) + if (source / "condition_rgb.png").is_file(): + shutil.copyfile(source / "condition_rgb.png", staged / "condition_rgb.png") + else: + with Image.open(source / "condition_rgba.png") as image: + rgba = np.asarray(image.convert("RGBA"), dtype=np.float32) / 255. + Image.fromarray((np.clip(rgba[..., :3] * rgba[..., 3:4], 0, 1) * 255).astype(np.uint8)).save(staged / "condition_rgb.png") + return _run_side(native, staged, output / "front", camera) + + +def run_dual_view(native_glb: Path, front_input_dir: Path, back_input_dir: Path, + cameras: dict[str, Any], output_dir: Path) -> Path: + """Run both measured reprojections and publish the stitched final GLB. + + The caller owns geometry, generated-image, and material-model provenance. + This function binds the provided model-independent maps, camera pack, native + mesh, reprojection code, candidate gates, and final atlas in its own receipt. + """ + native, front, back = (Path(path).resolve(strict=True) for path in (native_glb, front_input_dir, back_input_dir)) + output = Path(output_dir).resolve() + raw = cameras.get("raw", cameras) + front_camera, back_camera = camera_spec(raw, "front"), camera_spec(raw, "back") + for path in (output / "post", output / "post_back"): + if path.exists(): + raise FileExistsError(f"dual-view reprojection requires a new destination: {path}") + output.mkdir(parents=True, exist_ok=True) + final_dir = output / "post_back" + final_dir.mkdir() + try: + write_json(final_dir / "cameras.json", raw) + front_mesh = _run_side(native, front, output / "post", front_camera) + rotation, translation, distance = back_to_front(raw["back"]) + rotated = final_dir / "rotated_baseline.glb" + transform_glb(front_mesh, rotated, rotation, translation) + canonical_back = { + "extrinsic": [[1, 0, 0, 0], [0, -1, 0, 0], [0, 0, -1, distance], [0, 0, 0, 1]], + "projection_intrinsic": back_camera["projection_intrinsic"], "distance": distance, + "camera_angle_x": back_camera["camera_angle_x"], "mesh_scale": 1.0, + } + back_mesh = _run_side(rotated, back, final_dir / "reproject", canonical_back) + protected = front_protection(output / "post" / "project" / "COMMON_DOMAIN.npz") + final = final_dir / "mesh.glb" + stitch_gates = stitch(front_mesh, back_mesh, protected, final) + _validate_side(output / "post") + _validate_side(final_dir / "reproject") + receipt = { + "profile": PROFILE, "status": "valid", "view": "dual", "mesh": str(final), + "gates": {"accepted": True, "front_candidates_accepted_all_19": True, + "back_candidates_accepted_all_19": True, **stitch_gates}, + "front_post_receipt": bind_file(output / "post" / "receipt.json"), + "back_post_receipt": bind_file(final_dir / "reproject" / "receipt.json"), + "camera_pack": bind_file(final_dir / "cameras.json"), "final_mesh": bind_file(final), + "front_protection": bind_file(output / "post" / "project" / "COMMON_DOMAIN.npz"), + "back_transform": {"rotation": rotation.tolist(), "translation": translation.tolist(), "distance": distance}, + "validation_scope": "Both view candidates passed 19-phase gates; the stitched atlas has exact front priority and preserved non-base content.", + } + write_json(final_dir / "receipt.json", receipt) + validate_dual_view(output, raise_on_invalid=True) + return final + except BaseException as error: + failure = final_dir / "FAILED.json" + if not failure.exists(): + write_json(failure, {"profile": PROFILE, "status": "failed", "error": f"{type(error).__name__}: {error}"}) + raise + + +def validate_dual_view(output_dir: Path, *, raise_on_invalid: bool = False) -> dict[str, Any]: + """Re-read the complete postprocess receipt chain and content bindings.""" + try: + output = Path(output_dir).resolve() + receipt = read_json(output / "post_back" / "receipt.json") + if (output / "post_back" / "FAILED.json").exists(): + raise ValueError("dual-view reprojection has a recorded failure") + if receipt.get("profile") != PROFILE or receipt.get("status") != "valid" or receipt.get("view") != "dual": + raise ValueError("invalid dual-view reprojection identity or status") + gates = receipt.get("gates", {}) + required = ("accepted", "front_candidates_accepted_all_19", "back_candidates_accepted_all_19", + "front_priority_exact", "alpha_exact", "position_exact", "normal_exact", "uv_exact") + if not all(gates.get(key) is True for key in required) or gates.get("final_19_phase_rendered") is not False: + raise ValueError("dual-view reprojection gates or validation scope are incomplete") + check_binding(receipt["front_post_receipt"], output / "post" / "receipt.json") + check_binding(receipt["back_post_receipt"], output / "post_back" / "reproject" / "receipt.json") + check_binding(receipt["front_protection"], output / "post" / "project" / "COMMON_DOMAIN.npz") + camera_path = check_binding(receipt["camera_pack"], output / "post_back" / "cameras.json") + final = check_binding(receipt["final_mesh"], output / "post_back" / "mesh.glb") + if Path(receipt.get("mesh", "")).resolve() != final: + raise ValueError("dual-view final mesh path is inconsistent") + cameras = read_json(camera_path) + front_camera = camera_spec(cameras, "front") + back_camera = camera_spec(cameras, "back") + rotation, translation, distance = back_to_front(cameras["back"]) + expected_transform = {"rotation": rotation.tolist(), "translation": translation.tolist(), "distance": distance} + if receipt.get("back_transform") != expected_transform: + raise ValueError("back transform differs from the camera pack") + if read_json(output / "post" / "inputs_1024" / "camera.json") != front_camera: + raise ValueError("front postprocess did not use the bound front camera") + expected_back = { + "extrinsic": [[1, 0, 0, 0], [0, -1, 0, 0], [0, 0, -1, distance], [0, 0, 0, 1]], + "projection_intrinsic": back_camera["projection_intrinsic"], "distance": distance, + "camera_angle_x": back_camera["camera_angle_x"], "mesh_scale": 1.0, + } + if read_json(output / "post_back" / "reproject" / "inputs_1024" / "camera.json") != expected_back: + raise ValueError("back postprocess did not use the canonicalized back camera") + expected_document, expected_binary = transformed_payload(output / "post" / "mesh.glb", rotation, translation) + actual_document, actual_binary = _glb._read_glb(output / "post_back" / "rotated_baseline.glb") + if actual_document != expected_document or actual_binary != expected_binary: + raise ValueError("temporary back mesh differs from the required rigid transform") + _validate_side(output / "post") + _validate_side(output / "post_back" / "reproject") + protected = front_protection(output / "post" / "project" / "COMMON_DOMAIN.npz") + actual = validate_stitch(output / "post" / "mesh.glb", output / "post_back" / "reproject" / "mesh.glb", final, protected) + if any(gates.get(key) != value for key, value in actual.items()): + raise ValueError("stitched preservation evidence changed") + return {"valid": True, "reason": "Both candidate chains and the stitched final mesh are current", "receipt": receipt} + except (OSError, ValueError, KeyError, TypeError, RuntimeError) as error: + if raise_on_invalid: + raise + return {"valid": False, "reason": str(error)} diff --git a/supratex/inference/reprojection/_common.py b/supratex/inference/reprojection/_common.py index d8468613f5f548cbaa7b334f15ce6065d1b49dd1..777d927c437df670fd9fe6ec1e43429fbd2cbde2 100644 --- a/supratex/inference/reprojection/_common.py +++ b/supratex/inference/reprojection/_common.py @@ -1,115 +1,117 @@ -"""File bindings and measured acceptance checks shared by atlas operators.""" -from __future__ import annotations - -import hashlib -import json -import math -import os -from pathlib import Path -from typing import Any, Mapping - -OUTSIDE_COLOR_MAX_ABS = 5 / 255.0 -REPLAY_MAX_ABS = 1e-5 -TRAIN = tuple((x / 4, y / 4) for x in range(4) for y in range(4)) -HELD = ((.125, .375), (.375, .625), (.875, .125)) -IDENTITY_GATES = ( - "materials_exact", "mesh_structure_exact", "node_structure_exact", - "nonbase_buffer_views_exact", -) - - -def outside_color_passes(value: float) -> bool: - return math.isfinite(value) and 0 <= value <= OUTSIDE_COLOR_MAX_ABS - - -def write_json(path: Path, value: Any) -> None: - """Publish one new JSON artifact without overwriting an earlier result.""" - if path.exists(): - raise FileExistsError(f"refusing to overwrite {path}") - path.parent.mkdir(parents=True, exist_ok=True) - temporary = path.with_name(f".{path.name}.{os.getpid()}.tmp") - temporary.write_text(json.dumps(value, indent=2, sort_keys=True, allow_nan=False) + "\n", encoding="utf-8") - os.replace(temporary, path) - - -def read_json(path: Path) -> dict[str, Any]: - value = json.loads(path.read_text(encoding="utf-8")) - if not isinstance(value, dict): - raise ValueError(f"expected a JSON object: {path}") - return value - - -def bind_file(path: Path) -> dict[str, Any]: - path = path.resolve(strict=True) - if not path.is_file(): - raise ValueError(f"expected a regular file: {path}") - digest = hashlib.sha256() - with path.open("rb") as stream: - for block in iter(lambda: stream.read(1024 * 1024), b""): - digest.update(block) - return {"path": str(path), "bytes": path.stat().st_size, "sha256": digest.hexdigest()} - - -def check_binding(record: Mapping[str, Any], path: Path | None = None) -> Path: - bound = Path(record["path"]) - if path is not None and bound.resolve() != path.resolve(): - raise ValueError(f"artifact path binding changed: {path}") - if dict(record) != bind_file(bound): - raise ValueError(f"artifact content binding changed: {bound}") - return bound - - -def load_renderer(path: Path): - """Only the renderer shipped in this package may establish these gates.""" - from . import _renderer - if path.resolve() != Path(_renderer.__file__).resolve(): - raise ValueError("reprojection requires the bundled SupraTex renderer") - return _renderer - - -def validate_candidate(directory: Path, *, continuous: bool) -> dict[str, Any]: - """Read every phase receipt and validate numerical gates, not just status.""" - receipt = read_json(directory / "RECEIPT.json") - if read_json(directory / "RESULT.json") != receipt: - raise ValueError("candidate result and receipt disagree") - if receipt.get("status") != "valid": - raise ValueError(f"candidate was not accepted: {directory}") - gates = receipt.get("gates", {}) - if gates.get("accepted") is not True or gates.get("all_19_phase_renders") is not True: - raise ValueError("candidate lacks complete acceptance gates") - if not all(gates.get("nonbase_identity", {}).get(key) is True for key in IDENTITY_GATES): - raise ValueError("candidate lacks non-base preservation gates") - if continuous: - if not all(gates.get(key) is True for key in ("UV_covered_exact", "B_active_exact", "blocked_exact", "alpha_exact")): - raise ValueError("continuous candidate lacks atlas preservation gates") - elif receipt.get("atlas", {}).get("alpha_exact") is not True: - raise ValueError("projection candidate did not preserve alpha") - candidate = directory / "candidate" / "mesh.glb" - if Path(receipt.get("candidate_glb", "")).resolve() != candidate.resolve() or not candidate.is_file(): - raise ValueError("candidate mesh binding is missing or incorrect") - for split, expected in (("train", TRAIN), ("heldout", HELD)): - rows = receipt.get("phase_results", {}).get(split, []) - if len(rows) != len(expected): - raise ValueError(f"incomplete {split} phase coverage") - for phase, row in zip(expected, rows): - if row.get("phase") != list(phase) or row.get("passes") is not True: - raise ValueError("phase coordinates or acceptance changed") - replay = float(row["quantized_atlas_replay_max"]) - if not math.isfinite(replay) or not 0 <= replay <= REPLAY_MAX_ABS: - raise ValueError("renderer replay gate failed") - if row.get("outside_source_hardtrust_max_abs_threshold") != OUTSIDE_COLOR_MAX_ABS: - raise ValueError("outside-source color threshold changed") - if not outside_color_passes(float(row["outside_source_hardtrust_delta_max_abs"])): - raise ValueError("outside-source color gate failed") - if row.get("mask_exact") is not True or row.get("mra_exact") is not True: - raise ValueError("phase mask or material preservation gate is missing") - folder = directory / "renders" / split / f"x{phase[0]:.3f}_y{phase[1]:.3f}" - if read_json(folder / "PHASE_RECEIPT.json") != row: - raise ValueError("phase detail differs from candidate receipt") - evidence = row.get("evidence", {}) - names = ("repair", "B") if continuous else ("candidate", "baseline") - if set(evidence) != set(names): - raise ValueError("phase lacks its two rendered evidence arrays") - for name in names: - check_binding(evidence[name], folder / f"{name}.npz") - return receipt +"""File bindings and measured acceptance checks shared by atlas operators.""" +from __future__ import annotations + +import hashlib +import json +import math +import os +from pathlib import Path +from typing import Any, Mapping + +OUTSIDE_COLOR_MAX_ABS = 5 / 255.0 +REPLAY_MAX_ABS = 1e-5 +# Validation render phases, as sub-pixel camera offsets: a 4x4 grid plus three +# off-grid points. Both operators render every candidate at all nineteen. +GRID_PHASES = tuple((x / 4, y / 4) for x in range(4) for y in range(4)) +OFFGRID_PHASES = ((.125, .375), (.375, .625), (.875, .125)) +IDENTITY_GATES = ( + "materials_exact", "mesh_structure_exact", "node_structure_exact", + "nonbase_buffer_views_exact", +) + + +def outside_color_passes(value: float) -> bool: + return math.isfinite(value) and 0 <= value <= OUTSIDE_COLOR_MAX_ABS + + +def write_json(path: Path, value: Any) -> None: + """Publish one new JSON artifact without overwriting an earlier result.""" + if path.exists(): + raise FileExistsError(f"refusing to overwrite {path}") + path.parent.mkdir(parents=True, exist_ok=True) + temporary = path.with_name(f".{path.name}.{os.getpid()}.tmp") + temporary.write_text(json.dumps(value, indent=2, sort_keys=True, allow_nan=False) + "\n", encoding="utf-8") + os.replace(temporary, path) + + +def read_json(path: Path) -> dict[str, Any]: + value = json.loads(path.read_text(encoding="utf-8")) + if not isinstance(value, dict): + raise ValueError(f"expected a JSON object: {path}") + return value + + +def bind_file(path: Path) -> dict[str, Any]: + path = path.resolve(strict=True) + if not path.is_file(): + raise ValueError(f"expected a regular file: {path}") + digest = hashlib.sha256() + with path.open("rb") as stream: + for block in iter(lambda: stream.read(1024 * 1024), b""): + digest.update(block) + return {"path": str(path), "bytes": path.stat().st_size, "sha256": digest.hexdigest()} + + +def check_binding(record: Mapping[str, Any], path: Path | None = None) -> Path: + bound = Path(record["path"]) + if path is not None and bound.resolve() != path.resolve(): + raise ValueError(f"artifact path binding changed: {path}") + if dict(record) != bind_file(bound): + raise ValueError(f"artifact content binding changed: {bound}") + return bound + + +def load_renderer(path: Path): + """Only the renderer shipped in this package may establish these gates.""" + from . import _renderer + if path.resolve() != Path(_renderer.__file__).resolve(): + raise ValueError("reprojection requires the bundled SupraTex renderer") + return _renderer + + +def validate_candidate(directory: Path, *, continuous: bool) -> dict[str, Any]: + """Read every phase receipt and validate numerical gates, not just status.""" + receipt = read_json(directory / "RECEIPT.json") + if read_json(directory / "RESULT.json") != receipt: + raise ValueError("candidate result and receipt disagree") + if receipt.get("status") != "valid": + raise ValueError(f"candidate was not accepted: {directory}") + gates = receipt.get("gates", {}) + if gates.get("accepted") is not True or gates.get("all_19_phase_renders") is not True: + raise ValueError("candidate lacks complete acceptance gates") + if not all(gates.get("nonbase_identity", {}).get(key) is True for key in IDENTITY_GATES): + raise ValueError("candidate lacks non-base preservation gates") + if continuous: + if not all(gates.get(key) is True for key in ("UV_covered_exact", "project_active_exact", "blocked_exact", "alpha_exact")): + raise ValueError("continuous candidate lacks atlas preservation gates") + elif receipt.get("atlas", {}).get("alpha_exact") is not True: + raise ValueError("projection candidate did not preserve alpha") + candidate = directory / "candidate" / "mesh.glb" + if Path(receipt.get("candidate_glb", "")).resolve() != candidate.resolve() or not candidate.is_file(): + raise ValueError("candidate mesh binding is missing or incorrect") + for split, expected in (("grid", GRID_PHASES), ("offgrid", OFFGRID_PHASES)): + rows = receipt.get("phase_results", {}).get(split, []) + if len(rows) != len(expected): + raise ValueError(f"incomplete {split} phase coverage") + for phase, row in zip(expected, rows): + if row.get("phase") != list(phase) or row.get("passes") is not True: + raise ValueError("phase coordinates or acceptance changed") + replay = float(row["quantized_atlas_replay_max"]) + if not math.isfinite(replay) or not 0 <= replay <= REPLAY_MAX_ABS: + raise ValueError("renderer replay gate failed") + if row.get("outside_source_hardtrust_max_abs_threshold") != OUTSIDE_COLOR_MAX_ABS: + raise ValueError("outside-source color threshold changed") + if not outside_color_passes(float(row["outside_source_hardtrust_delta_max_abs"])): + raise ValueError("outside-source color gate failed") + if row.get("mask_exact") is not True or row.get("mra_exact") is not True: + raise ValueError("phase mask or material preservation gate is missing") + folder = directory / "renders" / split / f"x{phase[0]:.3f}_y{phase[1]:.3f}" + if read_json(folder / "PHASE_RECEIPT.json") != row: + raise ValueError("phase detail differs from candidate receipt") + evidence = row.get("evidence", {}) + names = ("repair", "project") if continuous else ("candidate", "baseline") + if set(evidence) != set(names): + raise ValueError("phase lacks its two rendered evidence arrays") + for name in names: + check_binding(evidence[name], folder / f"{name}.npz") + return receipt diff --git a/supratex/inference/reprojection/_continuous.py b/supratex/inference/reprojection/_continuous.py index 608b96e2cac15f7abfdf4a1b3524683a6896bbb1..6a43f419751a384a3058ae2c5a3665d722d827bf 100644 --- a/supratex/inference/reprojection/_continuous.py +++ b/supratex/inference/reprojection/_continuous.py @@ -1,389 +1,514 @@ -"""Repair one-texel UV gutters with observed continuous albedo gradients.""" -from __future__ import annotations -import argparse, copy, io, json, os, shutil, sys, time -from pathlib import Path -from ._common import OUTSIDE_COLOR_MAX_ABS, outside_color_passes, load_renderer, bind_file -from . import _glb as fixed -SIZE = 4096 -RES = 1024 -TRAIN = tuple(((x / 4, y / 4) for x in range(4) for y in range(4))) -HELD = ((0.125, 0.375), (0.375, 0.625), (0.875, 0.125)) -REJECT_UNSAFE = 1 -REJECT_CHART = 2 -REJECT_DEGENERATE_UV = 4 -REJECT_SOURCE = 8 - -def write_once(path, value): - if path.exists(): - raise FileExistsError(path) - path.parent.mkdir(parents=True, exist_ok=True) - path.write_text(json.dumps(value, indent=2, sort_keys=True, allow_nan=False) + '\n') - -def require(path, label): - path = Path(path).absolute() - if not path.is_file(): - raise FileNotFoundError(f'{label}: {path}') - return path - -def selected_case(manifest, label): - found = [x for x in json.loads(require(manifest, 'manifest').read_text()) if x.get('label') == label] - if len(found) != 1: - raise ValueError(f'expected one manifest case for {label!r}') - return found[0] - -def assert_skew_barycentric_oracle(np): - """Catch row/column and signed-area errors on a non-axis-aligned triangle.""" - tri = np.asarray(((0.11, 0.17), (0.83, 0.29), (0.26, 0.91)), np.float64) - expected = np.asarray((0.19, 0.34, 0.47), np.float64) - point = (tri * expected[:, None]).sum(0) - basis = tri[:2] - tri[2] - solved = np.linalg.solve(basis.T, point - tri[2]) - recovered = np.asarray((solved[0], solved[1], 1 - solved.sum())) - if not np.allclose(recovered, expected, rtol=0, atol=1e-12): - raise RuntimeError('skew barycentric oracle failed') - -def main(): - ap = argparse.ArgumentParser(description=__doc__) - ap.add_argument('--manifest', type=Path, required=True) - ap.add_argument('--case', required=True) - ap.add_argument('--evaluator', type=Path, required=True) - ap.add_argument('--output', type=Path, required=True) - ap.add_argument('--gutter-radius', type=int, choices=(1,), default=1, help='Only one-texel gutter repair is supported.') - a = ap.parse_args() - if a.output.exists(): - raise FileExistsError(a.output) - case = selected_case(a.manifest, a.case) - evaluator = require(a.evaluator, 'evaluator') - (work, bdir) = (Path(case['work']).absolute(), Path(case['b_output']).absolute()) - baseline = require(case['baseline_glb'], 'baseline GLB') - bglb = require(bdir / 'candidate/mesh.glb', 'accepted B GLB') - camera_path = require(work / 'camera.json', 'camera') - rgba_path = require(work / 'condition_rgba.png', 'condition RGBA') - albedo_path = require(work / 'albedo.png', 'predicted albedo') - domain_path = require(bdir / 'COMMON_DOMAIN.npz', 'B common domain') - bplan = json.loads(require(bdir / 'PLAN.json', 'B PLAN').read_text()) - bresult = json.loads(require(bdir / 'RESULT.json', 'B RESULT').read_text()) - expected = {'arm': 'B', 'baseline_glb': str(baseline), 'camera': str(camera_path), 'source_albedo': str(albedo_path), 'source_rgba': str(rgba_path), 'evaluator': str(evaluator)} - if [k for (k, v) in expected.items() if bplan.get(k) != v]: - raise RuntimeError('B PLAN binding failed') - if bresult.get('status') != 'valid' or not bresult.get('gates', {}).get('accepted'): - raise RuntimeError('B source is not an accepted B result') - if Path(bresult.get('candidate_glb', '')).absolute() != bglb: - raise RuntimeError('B result candidate binding failed') - identity = bresult.get('gates', {}).get('nonbase_identity', {}) - if not all((identity.get(k) for k in ('materials_exact', 'mesh_structure_exact', 'node_structure_exact', 'nonbase_buffer_views_exact'))): - raise RuntimeError('B non-base identity gate failed') - phases = bresult.get('phase_results', {}).get('train', []) + bresult.get('phase_results', {}).get('heldout', []) - if not bresult.get('gates', {}).get('all_19_phase_renders') or not bresult.get('atlas', {}).get('alpha_exact') or len(phases) != 19 or (not all((x.get('passes') for x in phases))): - raise RuntimeError('B alpha or 19-phase receipt binding failed') - a.output.mkdir(parents=True) - shutil.copyfile(Path(__file__), a.output / 'CODE_SNAPSHOT.py') - write_once(a.output / 'PLAN.json', {'schema': 'supratex.reprojection.continuous', 'case': a.case, 'manifest': str(a.manifest.absolute()), 'accepted_b_glb': str(bglb), 'baseline_glb': str(baseline), 'camera': str(camera_path), 'evaluator': str(evaluator), 'source_albedo': str(albedo_path), 'source_rgba': str(rgba_path), 'common_domain': str(domain_path), 'train_phases': TRAIN, 'heldout_phases': HELD, 'gutter_radius': a.gutter_radius, 'operator': 'anchor each gutter to its accepted-B nearest covered donor RGB, then add the continuous predicted-albedo projection gradient from donor centre to extrapolated gutter centre', 'preservation': 'every UV-covered texel, blocked ID, alpha, mesh and all non-base GLB content remain exact', 'source_guard': 'requires positive camera depth, a finite non-degenerate UV face, bounded extrapolated barycentrics, and all four native source bilinear footprint pixels hard-trusted for both donor and extrapolated gutter samples', 'unsafe': 'all bilinear atlas taps of non-hard-trusted pixels from all 19 declared phases are excluded before writes', 'limits': 'Construction uses train and held-out phases for safety; held-out render is a validation phase, not independent selection.'}) - os.environ['FLEX_GEMM_USE_AUTOTUNE_CACHE'] = '0' - os.environ['FLEX_GEMM_AUTOSAVE_AUTOTUNE_CACHE'] = '0' - import cv2 - import numpy as np - import torch - from PIL import Image - assert_skew_barycentric_oracle(np) - torch.backends.cuda.matmul.allow_tf32 = False - torch.manual_seed(42) - started = time.monotonic() - rm = load_renderer(evaluator) - rt = rm._runtime_imports() - dr = rt['dr'] - camera = json.loads(camera_path.read_text()) - (bdoc, bbin) = fixed._read_glb(bglb) - fixed._assert_fixed_scene(bdoc) - (ii, vi) = fixed._base_color_image_ref(bdoc) - (bimg, _, off, n) = fixed._decode_base_image(bdoc, bbin, ii, vi) - before = np.asarray(bimg).copy() - (basedoc, basebin) = fixed._read_glb(baseline) - fixed._assert_fixed_scene(basedoc) - (baseii, basevi) = fixed._base_color_image_ref(basedoc) - (baseimg, _, _, _) = fixed._decode_base_image(basedoc, basebin, baseii, basevi) - base_before = np.asarray(baseimg).copy() - if before.shape[:2] != (SIZE, SIZE): - raise ValueError('requires 4096 atlas') - if base_before.shape != before.shape: - raise RuntimeError('baseline/B baseColor image shape mismatch') - with np.load(domain_path, allow_pickle=False) as z: - (active_ids, blocked_ids) = (z['active_texel_ids'].astype(np.int64), z['blocked_texel_ids'].astype(np.int64)) - if len(active_ids) != int(bresult.get('atlas', {}).get('active_texels', -1)) or len(blocked_ids) != int(bresult.get('atlas', {}).get('blocked_texels', -1)) or len(np.unique(active_ids)) != len(active_ids) or (len(np.unique(blocked_ids)) != len(blocked_ids)): - raise RuntimeError('B common-domain count or uniqueness binding failed') - (active, blocked) = (np.zeros(SIZE * SIZE, bool), np.zeros(SIZE * SIZE, bool)) - (active[active_ids], blocked[blocked_ids]) = (True, True) - if np.any(blocked & ~active): - raise RuntimeError('blocked IDs must be in B active domain') - (vertices, faces_np, _, uv_np, _, _, factor, _, _) = rt['_load_textured_mesh'](bglb, rt) - if not np.allclose(factor, 1): - raise RuntimeError('unexpected baseColorFactor') - exported = torch.as_tensor(vertices, dtype=torch.float32, device='cuda') - native = exported * torch.tensor((-1.0, 1.0, -1.0), device='cuda') - faces = torch.as_tensor(faces_np, dtype=torch.int32, device='cuda').contiguous() - uv = torch.as_tensor(uv_np, dtype=torch.float32, device='cuda') - uv = torch.stack((uv[:, 0], 1 - uv[:, 1]), -1).contiguous() - rgba = np.asarray(Image.open(rgba_path).convert('RGBA')) - hard = cv2.erode((rgba[..., 3] >= 254).astype('uint8'), np.ones((5, 5), np.uint8)).astype(bool) - source = np.asarray(Image.open(albedo_path).convert('RGB'), np.float32) / 255 - if hard.shape != (RES, RES) or source.shape[:2] != (RES, RES): - raise ValueError('source inputs must be aligned 1024') - hard_t = torch.as_tensor(hard, device='cuda') - src = torch.as_tensor(source, device='cuda').permute(2, 0, 1)[None] - trust = np.clip(cv2.distanceTransform(hard.astype('uint8'), cv2.DIST_L2, 3) / 4.0, 0, 1) - trust_t = torch.as_tensor(trust, dtype=torch.float32, device='cuda')[None, None] - (E, K) = rm._native_camera(camera, 'cuda', rt) - original = rm._exported_glb_camera - - def pc(cam, device, runtime): - (ex, inn) = original(cam, device, runtime) - (dx, dy) = cam.get('front_atlas_phase', (0.0, 0.0)) - inn = inn.clone() - inn[0, 2] += dx / RES - inn[1, 2] += dy / RES - return (ex, inn) - rm._exported_glb_camera = pc - - def footprint(rows, uvs): - tex = uvs[0].reshape(-1, 2)[rows] * SIZE - 0.5 - (xy, frac) = (tex.floor().long(), tex - tex.floor()) - (ids, weights) = ([], []) - for (dx, dy) in ((0, 0), (1, 0), (0, 1), (1, 1)): - ids.append((xy[:, 1] + dy) % SIZE * SIZE + (xy[:, 0] + dx) % SIZE) - weights.append((frac[:, 0] if dx else 1 - frac[:, 0]) * (frac[:, 1] if dy else 1 - frac[:, 1])) - return (torch.stack(ids, 1), torch.stack(weights, 1)) - - def operation(phase): - cam = copy.deepcopy(camera) - cam['front_atlas_phase'] = phase - (ee, kk) = pc(cam, 'cuda', rt) - dist = rm._camera_distance(cam) - proj = rt['intrinsics_to_projection'](kk, max(0.2, dist - 1), dist + 1) - h = torch.cat((exported, torch.ones_like(exported[:, :1])), 1) - (rast, _) = dr.rasterize(dr.RasterizeCudaContext(device='cuda'), (h @ (proj @ ee).T)[None].contiguous(), faces, (RES, RES)) - (uvs, _) = dr.interpolate(uv[None], rast, faces) - (pts, _) = dr.interpolate(native[None], rast, faces) - mask = (rast[0, ..., -1] > 0).cpu().numpy() - rowsall = torch.as_tensor(np.flatnonzero(mask.reshape(-1)), device='cuda') - xyz = pts[0].reshape(-1, 3)[rowsall] - q = torch.cat((xyz, torch.ones_like(xyz[:, :1])), 1) @ E.T - q = q[:, :3] @ K.T - suv = q[:, :2] / q[:, 2:3].clamp_min(1e-08) - near = torch.floor(suv * RES).long() - inside = (near[:, 0] >= 0) & (near[:, 0] < RES) & (near[:, 1] >= 0) & (near[:, 1] < RES) & (q[:, 2] > 0) - trusted = torch.zeros(len(rowsall), dtype=torch.bool, device='cuda') - trusted[inside] = hard_t[near[inside, 1], near[inside, 0]] - rows = rowsall[trusted] - (idx, weights) = footprint(rows, uvs) - unsafe_screen = np.zeros(RES * RES, np.uint8) - unsafe_screen[rowsall[~trusted].cpu().numpy()] = 1 - unsafe_screen = cv2.dilate(unsafe_screen.reshape(RES, RES), np.ones((5, 5), np.uint8)) & mask - unsafe_rows = torch.as_tensor(np.flatnonzero(unsafe_screen.reshape(-1)), device='cuda') - unsafe = footprint(unsafe_rows, uvs)[0] if len(unsafe_rows) else torch.empty((0, 4), dtype=torch.long, device='cuda') - target = rt['F'].grid_sample(src, (suv[trusted] * 2 - 1).reshape(1, 1, -1, 2), align_corners=False, padding_mode='zeros')[0, :, 0].T - screen = np.zeros(RES * RES, bool) - screen[rows.cpu().numpy()] = True - return {'phase': phase, 'idx': idx, 'w': weights, 'unsafe': unsafe, 'rows': rows, 'screen': screen.reshape(RES, RES), 'target': target} - operations = [operation(p) for p in TRAIN + HELD] - unsafe_ids = torch.unique(torch.cat([x['unsafe'].reshape(-1) for x in operations if x['unsafe'].numel()])) if any((x['unsafe'].numel() for x in operations)) else torch.empty(0, dtype=torch.long, device='cuda') - unsafe = np.zeros(SIZE * SIZE, bool) - unsafe[unsafe_ids.cpu().numpy()] = True - uvclip = torch.cat((uv * 2 - 1, torch.zeros_like(uv[:, :1]), torch.ones_like(uv[:, :1])), 1)[None] - (urast, _) = dr.rasterize(dr.RasterizeCudaContext(device='cuda'), uvclip, faces, (SIZE, SIZE)) - covered = (urast[0, ..., -1] > 0).cpu().numpy() - (fcpu, uvcpu) = (faces.cpu().numpy(), uv.cpu().numpy()) - parent = np.arange(len(fcpu), dtype=np.int32) - - def root(i): - while parent[i] != i: - parent[i] = parent[parent[i]] - i = parent[i] - return i - - def join(i, j): - (i, j) = (root(i), root(j)) - if i != j: - parent[j] = i - edges = {} - for (fi, tri) in enumerate(fcpu): - for (u, v) in ((tri[0], tri[1]), (tri[1], tri[2]), (tri[2], tri[0])): - key = (int(min(u, v)), int(max(u, v))) - if key in edges: - join(fi, edges[key]) - else: - edges[key] = fi - face_chart = np.asarray([root(i) for i in range(len(fcpu))], np.int32) - (_, face_chart) = np.unique(face_chart, return_inverse=True) - primitive = urast[0, ..., -1].long().cpu().numpy() - 1 - chart = np.full((SIZE, SIZE), -1, np.int32) - chart[covered] = face_chart[primitive[covered]] - core = (active & ~blocked).reshape(SIZE, SIZE) & covered - if np.any(active.reshape(SIZE, SIZE) & ~covered): - raise RuntimeError('B active domain includes uncovered texels') - base = ~covered & ~blocked.reshape(SIZE, SIZE) - rejection = np.zeros((SIZE, SIZE), np.uint8) - rejection[base & unsafe.reshape(SIZE, SIZE)] |= REJECT_UNSAFE - candidate = np.zeros_like(core) - (donor_y, donor_x) = (np.full(core.shape, -1, np.int32), np.full(core.shape, -1, np.int32)) - donor_chart = np.full_like(chart, -1) - for (dy, dx) in ((-1, 0), (0, -1), (0, 1), (1, 0), (-1, -1), (-1, 1), (1, -1), (1, 1)): - shifted = np.zeros_like(core) - (sy, sx) = (slice(max(0, -dy), SIZE - max(0, dy)), slice(max(0, -dx), SIZE - max(0, dx))) - (ty, tx) = (slice(max(0, dy), SIZE - max(0, -dy)), slice(max(0, dx), SIZE - max(0, -dx))) - shifted[ty, tx] = core[sy, sx] - add = shifted & base & ~unsafe.reshape(SIZE, SIZE) - (ay, ax) = np.nonzero(add) - labels = chart[ay - dy, ax - dx] - for (y, x, label) in zip(ay.tolist(), ax.tolist(), labels.tolist()): - if not candidate[y, x]: - candidate[y, x] = True - donor_y[y, x] = y - dy - donor_x[y, x] = x - dx - donor_chart[y, x] = label - elif donor_chart[y, x] != label: - donor_chart[y, x] = -1 - prechart = candidate.copy() - for (dy, dx) in ((-1, -1), (-1, 0), (-1, 1), (0, -1), (0, 1), (1, -1), (1, 0), (1, 1)): - neigh = np.full_like(chart, -1) - neigh[max(0, dy):SIZE - max(0, -dy), max(0, dx):SIZE - max(0, -dx)] = chart[max(0, -dy):SIZE - max(0, dy), max(0, -dx):SIZE - max(0, dx)] - candidate &= (neigh < 0) | (neigh == donor_chart) - candidate &= donor_chart >= 0 - rejection[prechart & ~candidate] |= REJECT_CHART - (cy, cx) = np.nonzero(candidate) - did = donor_y[cy, cx].astype(np.int64) * SIZE + donor_x[cy, cx] - dfi = primitive.reshape(-1)[did] - tri_uv = uvcpu[fcpu[dfi]].astype(np.float64) - point_uv = np.stack(((cx + 0.5) / SIZE, (cy + 0.5) / SIZE), 1).astype(np.float64) - basis = tri_uv[:, :2] - tri_uv[:, 2:3] - det = basis[:, 0, 0] * basis[:, 1, 1] - basis[:, 0, 1] * basis[:, 1, 0] - finite_face = np.isfinite(tri_uv).all((1, 2)) & (np.abs(det) > 1e-12) - bary = np.full((len(cy), 3), np.nan, np.float32) - good = finite_face - delta = point_uv[good] - tri_uv[good, 2] - bary[good, :2] = np.linalg.solve(np.swapaxes(basis[good], 1, 2), delta[..., None])[..., 0] - bary[good, 2] = 1 - bary[good, 0] - bary[good, 1] - good &= np.isfinite(bary).all(1) & (np.abs(bary).max(1) <= 8.0) - rejection[cy[~good], cx[~good]] |= REJECT_DEGENERATE_UV - tri_native = vertices[fcpu[dfi]].astype(np.float64) * np.array((-1.0, 1.0, -1.0), np.float64) - xyz = (tri_native * bary[:, :, None]).sum(1) - donor_uv = np.stack(((donor_x[cy, cx] + 0.5) / SIZE, (donor_y[cy, cx] + 0.5) / SIZE), 1) - donor_delta = donor_uv - tri_uv[:, 2] - donor_bary = np.full((len(cy), 3), np.nan, np.float32) - donor_bary[good, :2] = np.linalg.solve(np.swapaxes(basis[good], 1, 2), donor_delta[good, ..., None])[..., 0] - donor_bary[good, 2] = 1 - donor_bary[good, 0] - donor_bary[good, 1] - good &= np.isfinite(donor_bary).all(1) & (np.abs(donor_bary).max(1) <= 8.0) - donor_contained = (donor_bary >= -1e-05).all(1) & (donor_bary <= 1 + 1e-05).all(1) - good &= donor_contained - rejection[cy[~donor_contained], cx[~donor_contained]] |= REJECT_DEGENERATE_UV - reconstructed_donor_uv = (tri_uv * donor_bary[:, :, None]).sum(1) - donor_uv_error = np.abs(reconstructed_donor_uv - donor_uv).max(1) - good &= donor_uv_error <= 1e-06 - rejection[cy[~good], cx[~good]] |= REJECT_DEGENERATE_UV - donor_xyz = (tri_native * donor_bary[:, :, None]).sum(1) - (atlas_native, _) = dr.interpolate(native[None], urast, faces) - donor_surface = atlas_native[0, torch.as_tensor(donor_y[cy, cx], device='cuda'), torch.as_tensor(donor_x[cy, cx], device='cuda')].cpu().numpy() - del atlas_native - donor_surface_error = np.abs(donor_surface - donor_xyz).max(1) - surface_checked = good.copy() - surface_ok = donor_surface_error <= 1e-05 - good &= surface_ok - rejection[cy[~surface_ok], cx[~surface_ok]] |= REJECT_DEGENERATE_UV - - def project_native(points): - points_t = torch.as_tensor(points, dtype=torch.float32, device='cuda') - q = torch.cat((points_t, torch.ones_like(points_t[:, :1])), 1) @ E.T - q = q[:, :3] @ K.T - return ((q[:, :2] / q[:, 2:3]).cpu().numpy(), q[:, 2].cpu().numpy()) - (source_uv, source_depth) = project_native(xyz) - (donor_source_uv, donor_depth) = project_native(donor_xyz) - cell = np.floor(np.where(np.isfinite(source_uv), source_uv, -2.0) * RES - 0.5).astype(np.int64) - donor_cell = np.floor(np.where(np.isfinite(donor_source_uv), donor_source_uv, -2.0) * RES - 0.5).astype(np.int64) - footprint_ok = good & np.isfinite(source_uv).all(1) & np.isfinite(donor_source_uv).all(1) & (source_depth > 0) & (donor_depth > 0) - footprint_ok &= (cell[:, 0] >= 0) & (cell[:, 0] + 1 < RES) & (cell[:, 1] >= 0) & (cell[:, 1] + 1 < RES) - footprint_ok &= (donor_cell[:, 0] >= 0) & (donor_cell[:, 0] + 1 < RES) & (donor_cell[:, 1] >= 0) & (donor_cell[:, 1] + 1 < RES) - fi = np.flatnonzero(footprint_ok) - if len(fi): - footprint_ok[fi] &= hard[cell[fi, 1], cell[fi, 0]] & hard[cell[fi, 1], cell[fi, 0] + 1] & hard[cell[fi, 1] + 1, cell[fi, 0]] & hard[cell[fi, 1] + 1, cell[fi, 0] + 1] & hard[donor_cell[fi, 1], donor_cell[fi, 0]] & hard[donor_cell[fi, 1], donor_cell[fi, 0] + 1] & hard[donor_cell[fi, 1] + 1, donor_cell[fi, 0]] & hard[donor_cell[fi, 1] + 1, donor_cell[fi, 0] + 1] - donor_strength = rt['F'].grid_sample(trust_t, (torch.as_tensor(donor_source_uv, dtype=torch.float32, device='cuda') * 2 - 1).reshape(1, 1, -1, 2), align_corners=False, padding_mode='zeros')[0, 0, 0].cpu().numpy() - endpoint_strength = rt['F'].grid_sample(trust_t, (torch.as_tensor(source_uv, dtype=torch.float32, device='cuda') * 2 - 1).reshape(1, 1, -1, 2), align_corners=False, padding_mode='zeros')[0, 0, 0].cpu().numpy() - footprint_ok &= (donor_strength >= 1 - 1e-06) & (endpoint_strength >= 1 - 1e-06) - sample_uv = torch.as_tensor(source_uv[footprint_ok], dtype=torch.float32, device='cuda') - donor_sample_uv = torch.as_tensor(donor_source_uv[footprint_ok], dtype=torch.float32, device='cuda') - sampled = rt['F'].grid_sample(src, (sample_uv * 2 - 1).reshape(1, 1, -1, 2), align_corners=False, padding_mode='zeros')[0, :, 0].T.cpu().numpy() - donor_sampled = rt['F'].grid_sample(src, (donor_sample_uv * 2 - 1).reshape(1, 1, -1, 2), align_corners=False, padding_mode='zeros')[0, :, 0].T.cpu().numpy() - source_rows = np.flatnonzero(footprint_ok) - donor_base = base_before[donor_y[cy[source_rows], cx[source_rows]], donor_x[cy[source_rows], cx[source_rows]], :3].astype(np.float32) / 255 - donor_actual = before[donor_y[cy[source_rows], cx[source_rows]], donor_x[cy[source_rows], cx[source_rows]], :3].astype(np.float32) / 255 - donor_rebuilt = donor_base + donor_strength[source_rows, None] * (donor_sampled - donor_base) - donor_b_source_byte_error = np.full(len(cy), np.inf, np.float32) - donor_b_source_byte_error[source_rows] = np.abs(donor_rebuilt - donor_actual).max(1) * 255 - donor_b_source_ok = donor_b_source_byte_error <= 1.0 - rejection[cy[~footprint_ok | footprint_ok & ~donor_b_source_ok], cx[~footprint_ok | footprint_ok & ~donor_b_source_ok]] |= REJECT_SOURCE - continuous_ok = footprint_ok & donor_b_source_ok - eligible = np.zeros_like(base) - eligible[cy[continuous_ok], cx[continuous_ok]] = True - after = before.copy() - after[cy, cx, :3] = before[donor_y[cy, cx], donor_x[cy, cx], :3] - if not np.array_equal(after[cy, cx, :3], before[donor_y[cy, cx], donor_x[cy, cx], :3]): - raise RuntimeError('accepted-B donor byte fallback failed') - (ey, ex) = np.nonzero(eligible) - anchor = before[donor_y[ey, ex], donor_x[ey, ex], :3].astype(np.float32) / 255 - selected_rows = np.flatnonzero(continuous_ok) - sample_indices = np.searchsorted(source_rows, selected_rows) - after[ey, ex, :3] = np.clip((anchor + sampled[sample_indices] - donor_sampled[sample_indices]) * 255 + 0.5, 0, 255).astype(np.uint8) - if not np.array_equal(after[covered], before[covered]): - raise RuntimeError('UV covered preservation failed') - if after.shape[-1] == 4 and (not np.array_equal(after[..., 3], before[..., 3])): - raise RuntimeError('alpha changed') - if not np.array_equal(after.reshape(-1, after.shape[-1])[active | blocked | unsafe], before.reshape(-1, before.shape[-1])[active | blocked | unsafe]): - raise RuntimeError('active, blocked, or unsafe texels changed') - eligible_ids = np.flatnonzero(eligible.reshape(-1)) - np.savez_compressed(a.output / 'GUTTER_DIAGNOSTICS.npz', covered=covered, unsafe=unsafe.reshape(SIZE, SIZE), candidate=candidate, eligible=eligible, rejection=rejection, eligible_texel_ids=eligible_ids, donor_texel_ids=did[continuous_ok], donor_charts=donor_chart.reshape(-1)[eligible_ids], source_uv=source_uv[continuous_ok], donor_source_uv=donor_source_uv[continuous_ok], barycentric=bary[continuous_ok], donor_barycentric=donor_bary[continuous_ok], donor_uv_reconstruction_error=donor_uv_error[continuous_ok], donor_surface_error=donor_surface_error[continuous_ok], donor_strength=donor_strength[continuous_ok], endpoint_strength=endpoint_strength[continuous_ok], donor_b_source_byte_error=donor_b_source_byte_error[continuous_ok], gutter_radius=np.asarray(a.gutter_radius, dtype=np.int32)) - blob = io.BytesIO() - Image.fromarray(after).save(blob, format='PNG') - odoc = copy.deepcopy(bdoc) - obin = fixed._replace_image_payload(odoc, bbin, vi, off, n, blob.getvalue()) - glb = a.output / 'candidate/mesh.glb' - glb.parent.mkdir() - fixed._write_glb(glb, odoc, obin) - preservation = fixed._nonbase_identity_check(bdoc, bbin, glb, vi) - - def render(o, split): - p = o['phase'] - folder = a.output / 'renders' / split / f'x{p[0]:.3f}_y{p[1]:.3f}' - folder.mkdir(parents=True) - cam = copy.deepcopy(camera) - cam['front_atlas_phase'] = p - (ca, bb) = (rm.render_glb(glb, cam, folder / 'repair.png', resolution=RES, rt=rt), rm.render_glb(bglb, cam, folder / 'B.png', resolution=RES, rt=rt)) - with np.load(ca['attrs_npz']) as z: - (ci, cm, cmra) = (z['albedo'], z['mask'].astype(bool), z['mra']) - with np.load(bb['attrs_npz']) as z: - (bi, bm, bmra) = (z['albedo'], z['mask'].astype(bool), z['mra']) - if not np.array_equal(cm, bm) or not np.array_equal(cmra, bmra): - raise RuntimeError('mask/MRA changed') - target = np.zeros_like(ci) - target.reshape(-1, 3)[o['rows'].cpu().numpy()] = o['target'].cpu().numpy() - sel = o['screen'] & cm - quant = torch.as_tensor(after[..., :3], dtype=torch.float32, device='cuda').reshape(-1, 3) / 255 - replay = (quant[o['idx']] * o['w'][..., None]).sum(1).cpu().numpy() - actual = ci.reshape(-1, 3)[o['rows'].cpu().numpy()] - outside = cm & ~o['screen'] - outdelta = float(np.abs(ci - bi)[outside].max()) if outside.any() else 0.0 - changed = cm & (np.abs(ci - bi).max(2) > 0) - r = {'phase': p, 'hardtrust_pixels': int(sel.sum()), 'repair_vs_source': rm._masked_metrics(ci, target, sel), 'B_vs_source': rm._masked_metrics(bi, target, sel), 'repair_vs_B_changed_pixels': int(changed.sum()), 'repair_vs_B': rm._masked_metrics(ci, bi, changed) if changed.any() else {'valid_pixels': 0}, 'quantized_atlas_replay_max': float(np.abs(replay - actual).max()), 'outside_source_hardtrust_max_abs_threshold': OUTSIDE_COLOR_MAX_ABS, 'outside_source_hardtrust_delta_max_abs': outdelta} - r['passes'] = bool(r['quantized_atlas_replay_max'] <= 1e-05 and outside_color_passes(outdelta)) - r.update(mask_exact=True, mra_exact=True, evidence={'repair': bind_file(Path(ca['attrs_npz'])), 'B': bind_file(Path(bb['attrs_npz']))}) - write_once(folder / 'PHASE_RECEIPT.json', r) - return r - results = [render(x, 'train' if i < len(TRAIN) else 'heldout') for (i, x) in enumerate(operations)] - changed = np.any(after[..., :3] != before[..., :3], axis=-1) - accepted = all((x['passes'] for x in results)) - gate_ok = accepted and all((preservation.get(k) for k in ('materials_exact', 'mesh_structure_exact', 'node_structure_exact', 'nonbase_buffer_views_exact'))) - receipt = {'status': 'valid' if gate_ok else 'invalid', 'candidate_glb': str(glb), 'seconds': time.monotonic() - started, 'gates': {'accepted': gate_ok, 'UV_covered_exact': True, 'B_active_exact': True, 'blocked_exact': True, 'alpha_exact': True, 'nonbase_identity': preservation, 'all_19_phase_renders': len(results) == 19}, 'atlas': {'UV_covered_texels': int(covered.sum()), 'UV_face_charts': int(face_chart.max() + 1), 'B_active_texels': int(active.sum()), 'blocked_texels': int(blocked.sum()), 'unsafe_all_phase_texels': int(unsafe.sum()), 'gutter_radius': a.gutter_radius, 'nearest_donor_uncovered_gutter_texels': int(candidate.sum()), 'continuous_eligible_uncovered_gutter_texels': int(eligible.sum()), 'changed_texels': int(changed.sum()), 'rejected_unsafe': int((rejection & REJECT_UNSAFE != 0).sum()), 'rejected_chart': int((rejection & REJECT_CHART != 0).sum()), 'rejected_degenerate_uv': int((rejection & REJECT_DEGENERATE_UV != 0).sum()), 'rejected_source_footprint_or_feather': int((rejection & REJECT_SOURCE != 0).sum()), 'donor_surface_oracle': {'threshold': 1e-05, 'checked': int(surface_checked.sum()), 'rejected': int((surface_checked & ~surface_ok).sum()), 'pre_filter_max': float(donor_surface_error[surface_checked].max()) if surface_checked.any() else 0.0, 'pre_filter_p99': float(np.quantile(donor_surface_error[surface_checked], 0.99)) if surface_checked.any() else 0.0, 'accepted_max': float(donor_surface_error[footprint_ok].max()) if footprint_ok.any() else 0.0}}, 'phase_results': {'train': results[:len(TRAIN)], 'heldout': results[len(TRAIN):]}, 'diagnostics': str(a.output / 'GUTTER_DIAGNOSTICS.npz'), 'code_snapshot': str(a.output / 'CODE_SNAPSHOT.py'), 'limits': 'Fixed-shape observed-source continuous-gutter candidate. It does not establish novel-view quality or model improvement.'} - write_once(a.output / 'RECEIPT.json', receipt) - write_once(a.output / 'RESULT.json', receipt) - if not gate_ok: - raise RuntimeError('acceptance failed') - print(json.dumps({'status': 'valid', 'candidate': str(glb), 'changed_texels': int(changed.sum())})) -if __name__ == '__main__': - main() +"""Repair one-texel UV gutters with observed continuous albedo gradients. + +Each uncovered gutter texel next to an unblocked projected texel is anchored +to that nearest covered donor's projected RGB, then shifted by the predicted +albedo difference between the donor centre and the gutter centre extrapolated +on the donor's UV face. A gutter is written only if both samples have positive +camera depth, a finite non-degenerate UV face, bounded barycentrics, and fully +hard-trusted bilinear source footprints, and if its bilinear atlas taps avoid +untrusted pixels in all nineteen phases. Every UV-covered texel, blocked texel, +alpha, and all non-base GLB content stay exact. +""" +from __future__ import annotations +import argparse, copy, io, json, os, shutil, sys, time +from pathlib import Path +from ._common import GRID_PHASES, OFFGRID_PHASES, OUTSIDE_COLOR_MAX_ABS, IDENTITY_GATES, outside_color_passes, load_renderer, bind_file +from . import _glb as fixed +from ._project import SCHEMA as PROJECT_SCHEMA +SCHEMA = 'supratex.reprojection.continuous' +SIZE = 4096 +RES = 1024 +REJECT_UNSAFE = 1 +REJECT_CHART = 2 +REJECT_DEGENERATE_UV = 4 +REJECT_SOURCE = 8 + +def write_once(path, value): + if path.exists(): + raise FileExistsError(path) + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text(json.dumps(value, indent=2, sort_keys=True, allow_nan=False) + '\n') + +def require(path, label): + path = Path(path).absolute() + if not path.is_file(): + raise FileNotFoundError(f'{label}: {path}') + return path + +def selected_case(manifest, label): + found = [x for x in json.loads(require(manifest, 'manifest').read_text()) if x.get('label') == label] + if len(found) != 1: + raise ValueError(f'expected one manifest case for {label!r}') + return found[0] + +def assert_skew_barycentric_oracle(np): + """Catch row/column and signed-area errors on a non-axis-aligned triangle.""" + tri = np.asarray(((0.11, 0.17), (0.83, 0.29), (0.26, 0.91)), np.float64) + expected = np.asarray((0.19, 0.34, 0.47), np.float64) + point = (tri * expected[:, None]).sum(0) + basis = tri[:2] - tri[2] + solved = np.linalg.solve(basis.T, point - tri[2]) + recovered = np.asarray((solved[0], solved[1], 1 - solved.sum())) + if not np.allclose(recovered, expected, rtol=0, atol=1e-12): + raise RuntimeError('skew barycentric oracle failed') + +def main(): + ap = argparse.ArgumentParser(description=__doc__) + ap.add_argument('--manifest', type=Path, required=True) + ap.add_argument('--case', required=True) + ap.add_argument('--renderer', type=Path, required=True) + ap.add_argument('--output', type=Path, required=True) + ap.add_argument('--gutter-radius', type=int, choices=(1,), default=1, help='Only one-texel gutter repair is supported.') + a = ap.parse_args() + if a.output.exists(): + raise FileExistsError(a.output) + case = selected_case(a.manifest, a.case) + renderer = require(a.renderer, 'renderer') + (work, pdir) = (Path(case['work']).absolute(), Path(case['project_output']).absolute()) + baseline = require(case['baseline_glb'], 'baseline GLB') + pglb = require(pdir / 'candidate/mesh.glb', 'accepted projection GLB') + camera_path = require(work / 'camera.json', 'camera') + rgba_path = require(work / 'condition_rgba.png', 'condition RGBA') + albedo_path = require(work / 'albedo.png', 'predicted albedo') + domain_path = require(pdir / 'COMMON_DOMAIN.npz', 'projection common domain') + pplan = json.loads(require(pdir / 'PLAN.json', 'projection PLAN').read_text()) + presult = json.loads(require(pdir / 'RESULT.json', 'projection RESULT').read_text()) + expected = { + 'schema': PROJECT_SCHEMA, + 'baseline_glb': str(baseline), + 'camera': str(camera_path), + 'source_albedo': str(albedo_path), + 'source_rgba': str(rgba_path), + 'renderer': str(renderer), + } + if [k for (k, v) in expected.items() if pplan.get(k) != v]: + raise RuntimeError('projection PLAN binding failed') + if presult.get('status') != 'valid' or not presult.get('gates', {}).get('accepted'): + raise RuntimeError('projection source is not an accepted projection result') + if Path(presult.get('candidate_glb', '')).absolute() != pglb: + raise RuntimeError('projection result candidate binding failed') + identity = presult.get('gates', {}).get('nonbase_identity', {}) + if not all((identity.get(k) for k in IDENTITY_GATES)): + raise RuntimeError('projection non-base identity gate failed') + phase_results = presult.get('phase_results', {}) + phases = phase_results.get('grid', []) + phase_results.get('offgrid', []) + if (not presult.get('gates', {}).get('all_19_phase_renders') + or not presult.get('atlas', {}).get('alpha_exact') + or len(phases) != 19 + or not all((x.get('passes') for x in phases))): + raise RuntimeError('projection alpha or 19-phase receipt binding failed') + a.output.mkdir(parents=True) + shutil.copyfile(Path(__file__), a.output / 'CODE_SNAPSHOT.py') + write_once(a.output / 'PLAN.json', { + 'schema': SCHEMA, + 'case': a.case, + 'manifest': str(a.manifest.absolute()), + 'project_glb': str(pglb), + 'baseline_glb': str(baseline), + 'camera': str(camera_path), + 'renderer': str(renderer), + 'source_albedo': str(albedo_path), + 'source_rgba': str(rgba_path), + 'common_domain': str(domain_path), + 'grid_phases': GRID_PHASES, + 'offgrid_phases': OFFGRID_PHASES, + 'gutter_radius': a.gutter_radius, + }) + os.environ['FLEX_GEMM_USE_AUTOTUNE_CACHE'] = '0' + os.environ['FLEX_GEMM_AUTOSAVE_AUTOTUNE_CACHE'] = '0' + import cv2 + import numpy as np + import torch + from PIL import Image + assert_skew_barycentric_oracle(np) + torch.backends.cuda.matmul.allow_tf32 = False + torch.manual_seed(42) + started = time.monotonic() + rm = load_renderer(renderer) + rt = rm._runtime_imports() + dr = rt['dr'] + camera = json.loads(camera_path.read_text()) + (pdoc, pbin) = fixed._read_glb(pglb) + fixed._assert_fixed_scene(pdoc) + (ii, vi) = fixed._base_color_image_ref(pdoc) + (pimg, _, off, n) = fixed._decode_base_image(pdoc, pbin, ii, vi) + before = np.asarray(pimg).copy() + (basedoc, basebin) = fixed._read_glb(baseline) + fixed._assert_fixed_scene(basedoc) + (baseii, basevi) = fixed._base_color_image_ref(basedoc) + (baseimg, _, _, _) = fixed._decode_base_image(basedoc, basebin, baseii, basevi) + base_before = np.asarray(baseimg).copy() + if before.shape[:2] != (SIZE, SIZE): + raise ValueError('requires 4096 atlas') + if base_before.shape != before.shape: + raise RuntimeError('baseline/projection baseColor image shape mismatch') + with np.load(domain_path, allow_pickle=False) as z: + (active_ids, blocked_ids) = (z['active_texel_ids'].astype(np.int64), z['blocked_texel_ids'].astype(np.int64)) + project_atlas = presult.get('atlas', {}) + if (len(active_ids) != int(project_atlas.get('active_texels', -1)) + or len(blocked_ids) != int(project_atlas.get('blocked_texels', -1)) + or len(np.unique(active_ids)) != len(active_ids) + or len(np.unique(blocked_ids)) != len(blocked_ids)): + raise RuntimeError('projection common-domain count or uniqueness binding failed') + (active, blocked) = (np.zeros(SIZE * SIZE, bool), np.zeros(SIZE * SIZE, bool)) + (active[active_ids], blocked[blocked_ids]) = (True, True) + if np.any(blocked & ~active): + raise RuntimeError('blocked IDs must be in the projection active domain') + (vertices, faces_np, _, uv_np, _, _, factor, _, _) = rt['_load_textured_mesh'](pglb, rt) + if not np.allclose(factor, 1): + raise RuntimeError('unexpected baseColorFactor') + exported = torch.as_tensor(vertices, dtype=torch.float32, device='cuda') + native = exported * torch.tensor((-1.0, 1.0, -1.0), device='cuda') + faces = torch.as_tensor(faces_np, dtype=torch.int32, device='cuda').contiguous() + uv = torch.as_tensor(uv_np, dtype=torch.float32, device='cuda') + uv = torch.stack((uv[:, 0], 1 - uv[:, 1]), -1).contiguous() + rgba = np.asarray(Image.open(rgba_path).convert('RGBA')) + hard = cv2.erode((rgba[..., 3] >= 254).astype('uint8'), np.ones((5, 5), np.uint8)).astype(bool) + source = np.asarray(Image.open(albedo_path).convert('RGB'), np.float32) / 255 + if hard.shape != (RES, RES) or source.shape[:2] != (RES, RES): + raise ValueError('source inputs must be aligned 1024') + hard_t = torch.as_tensor(hard, device='cuda') + src = torch.as_tensor(source, device='cuda').permute(2, 0, 1)[None] + trust = np.clip(cv2.distanceTransform(hard.astype('uint8'), cv2.DIST_L2, 3) / 4.0, 0, 1) + trust_t = torch.as_tensor(trust, dtype=torch.float32, device='cuda')[None, None] + (E, K) = rm._native_camera(camera, 'cuda', rt) + original = rm._exported_glb_camera + + def pc(cam, device, runtime): + (ex, inn) = original(cam, device, runtime) + (dx, dy) = cam.get('front_atlas_phase', (0.0, 0.0)) + inn = inn.clone() + inn[0, 2] += dx / RES + inn[1, 2] += dy / RES + return (ex, inn) + rm._exported_glb_camera = pc + + def footprint(rows, uvs): + tex = uvs[0].reshape(-1, 2)[rows] * SIZE - 0.5 + (xy, frac) = (tex.floor().long(), tex - tex.floor()) + (ids, weights) = ([], []) + for (dx, dy) in ((0, 0), (1, 0), (0, 1), (1, 1)): + ids.append((xy[:, 1] + dy) % SIZE * SIZE + (xy[:, 0] + dx) % SIZE) + weights.append((frac[:, 0] if dx else 1 - frac[:, 0]) * (frac[:, 1] if dy else 1 - frac[:, 1])) + return (torch.stack(ids, 1), torch.stack(weights, 1)) + + def operation(phase): + cam = copy.deepcopy(camera) + cam['front_atlas_phase'] = phase + (ee, kk) = pc(cam, 'cuda', rt) + dist = rm._camera_distance(cam) + proj = rt['intrinsics_to_projection'](kk, max(0.2, dist - 1), dist + 1) + h = torch.cat((exported, torch.ones_like(exported[:, :1])), 1) + (rast, _) = dr.rasterize(dr.RasterizeCudaContext(device='cuda'), (h @ (proj @ ee).T)[None].contiguous(), faces, (RES, RES)) + (uvs, _) = dr.interpolate(uv[None], rast, faces) + (pts, _) = dr.interpolate(native[None], rast, faces) + mask = (rast[0, ..., -1] > 0).cpu().numpy() + rowsall = torch.as_tensor(np.flatnonzero(mask.reshape(-1)), device='cuda') + xyz = pts[0].reshape(-1, 3)[rowsall] + q = torch.cat((xyz, torch.ones_like(xyz[:, :1])), 1) @ E.T + q = q[:, :3] @ K.T + suv = q[:, :2] / q[:, 2:3].clamp_min(1e-08) + near = torch.floor(suv * RES).long() + inside = (near[:, 0] >= 0) & (near[:, 0] < RES) & (near[:, 1] >= 0) & (near[:, 1] < RES) & (q[:, 2] > 0) + trusted = torch.zeros(len(rowsall), dtype=torch.bool, device='cuda') + trusted[inside] = hard_t[near[inside, 1], near[inside, 0]] + rows = rowsall[trusted] + (idx, weights) = footprint(rows, uvs) + unsafe_screen = np.zeros(RES * RES, np.uint8) + unsafe_screen[rowsall[~trusted].cpu().numpy()] = 1 + unsafe_screen = cv2.dilate(unsafe_screen.reshape(RES, RES), np.ones((5, 5), np.uint8)) & mask + unsafe_rows = torch.as_tensor(np.flatnonzero(unsafe_screen.reshape(-1)), device='cuda') + unsafe = footprint(unsafe_rows, uvs)[0] if len(unsafe_rows) else torch.empty((0, 4), dtype=torch.long, device='cuda') + target = rt['F'].grid_sample(src, (suv[trusted] * 2 - 1).reshape(1, 1, -1, 2), align_corners=False, padding_mode='zeros')[0, :, 0].T + screen = np.zeros(RES * RES, bool) + screen[rows.cpu().numpy()] = True + return { + 'phase': phase, + 'idx': idx, + 'w': weights, + 'unsafe': unsafe, + 'rows': rows, + 'screen': screen.reshape(RES, RES), + 'target': target, + } + operations = [operation(p) for p in GRID_PHASES + OFFGRID_PHASES] + if any(x['unsafe'].numel() for x in operations): + unsafe_ids = torch.unique(torch.cat([x['unsafe'].reshape(-1) for x in operations if x['unsafe'].numel()])) + else: + unsafe_ids = torch.empty(0, dtype=torch.long, device='cuda') + unsafe = np.zeros(SIZE * SIZE, bool) + unsafe[unsafe_ids.cpu().numpy()] = True + uvclip = torch.cat((uv * 2 - 1, torch.zeros_like(uv[:, :1]), torch.ones_like(uv[:, :1])), 1)[None] + (urast, _) = dr.rasterize(dr.RasterizeCudaContext(device='cuda'), uvclip, faces, (SIZE, SIZE)) + covered = (urast[0, ..., -1] > 0).cpu().numpy() + (fcpu, uvcpu) = (faces.cpu().numpy(), uv.cpu().numpy()) + parent = np.arange(len(fcpu), dtype=np.int32) + + def root(i): + while parent[i] != i: + parent[i] = parent[parent[i]] + i = parent[i] + return i + + def join(i, j): + (i, j) = (root(i), root(j)) + if i != j: + parent[j] = i + edges = {} + for (fi, tri) in enumerate(fcpu): + for (u, v) in ((tri[0], tri[1]), (tri[1], tri[2]), (tri[2], tri[0])): + key = (int(min(u, v)), int(max(u, v))) + if key in edges: + join(fi, edges[key]) + else: + edges[key] = fi + face_chart = np.asarray([root(i) for i in range(len(fcpu))], np.int32) + (_, face_chart) = np.unique(face_chart, return_inverse=True) + primitive = urast[0, ..., -1].long().cpu().numpy() - 1 + chart = np.full((SIZE, SIZE), -1, np.int32) + chart[covered] = face_chart[primitive[covered]] + core = (active & ~blocked).reshape(SIZE, SIZE) & covered + if np.any(active.reshape(SIZE, SIZE) & ~covered): + raise RuntimeError('projection active domain includes uncovered texels') + base = ~covered & ~blocked.reshape(SIZE, SIZE) + rejection = np.zeros((SIZE, SIZE), np.uint8) + rejection[base & unsafe.reshape(SIZE, SIZE)] |= REJECT_UNSAFE + candidate = np.zeros_like(core) + (donor_y, donor_x) = (np.full(core.shape, -1, np.int32), np.full(core.shape, -1, np.int32)) + donor_chart = np.full_like(chart, -1) + for (dy, dx) in ((-1, 0), (0, -1), (0, 1), (1, 0), (-1, -1), (-1, 1), (1, -1), (1, 1)): + shifted = np.zeros_like(core) + (sy, sx) = (slice(max(0, -dy), SIZE - max(0, dy)), slice(max(0, -dx), SIZE - max(0, dx))) + (ty, tx) = (slice(max(0, dy), SIZE - max(0, -dy)), slice(max(0, dx), SIZE - max(0, -dx))) + shifted[ty, tx] = core[sy, sx] + add = shifted & base & ~unsafe.reshape(SIZE, SIZE) + (ay, ax) = np.nonzero(add) + labels = chart[ay - dy, ax - dx] + for (y, x, label) in zip(ay.tolist(), ax.tolist(), labels.tolist()): + if not candidate[y, x]: + candidate[y, x] = True + donor_y[y, x] = y - dy + donor_x[y, x] = x - dx + donor_chart[y, x] = label + elif donor_chart[y, x] != label: + donor_chart[y, x] = -1 + prechart = candidate.copy() + for (dy, dx) in ((-1, -1), (-1, 0), (-1, 1), (0, -1), (0, 1), (1, -1), (1, 0), (1, 1)): + neigh = np.full_like(chart, -1) + neigh[max(0, dy):SIZE - max(0, -dy), max(0, dx):SIZE - max(0, -dx)] = chart[max(0, -dy):SIZE - max(0, dy), max(0, -dx):SIZE - max(0, dx)] + candidate &= (neigh < 0) | (neigh == donor_chart) + candidate &= donor_chart >= 0 + rejection[prechart & ~candidate] |= REJECT_CHART + (cy, cx) = np.nonzero(candidate) + did = donor_y[cy, cx].astype(np.int64) * SIZE + donor_x[cy, cx] + dfi = primitive.reshape(-1)[did] + tri_uv = uvcpu[fcpu[dfi]].astype(np.float64) + point_uv = np.stack(((cx + 0.5) / SIZE, (cy + 0.5) / SIZE), 1).astype(np.float64) + basis = tri_uv[:, :2] - tri_uv[:, 2:3] + det = basis[:, 0, 0] * basis[:, 1, 1] - basis[:, 0, 1] * basis[:, 1, 0] + finite_face = np.isfinite(tri_uv).all((1, 2)) & (np.abs(det) > 1e-12) + bary = np.full((len(cy), 3), np.nan, np.float32) + good = finite_face + delta = point_uv[good] - tri_uv[good, 2] + bary[good, :2] = np.linalg.solve(np.swapaxes(basis[good], 1, 2), delta[..., None])[..., 0] + bary[good, 2] = 1 - bary[good, 0] - bary[good, 1] + good &= np.isfinite(bary).all(1) & (np.abs(bary).max(1) <= 8.0) + rejection[cy[~good], cx[~good]] |= REJECT_DEGENERATE_UV + tri_native = vertices[fcpu[dfi]].astype(np.float64) * np.array((-1.0, 1.0, -1.0), np.float64) + xyz = (tri_native * bary[:, :, None]).sum(1) + donor_uv = np.stack(((donor_x[cy, cx] + 0.5) / SIZE, (donor_y[cy, cx] + 0.5) / SIZE), 1) + donor_delta = donor_uv - tri_uv[:, 2] + donor_bary = np.full((len(cy), 3), np.nan, np.float32) + donor_bary[good, :2] = np.linalg.solve(np.swapaxes(basis[good], 1, 2), donor_delta[good, ..., None])[..., 0] + donor_bary[good, 2] = 1 - donor_bary[good, 0] - donor_bary[good, 1] + good &= np.isfinite(donor_bary).all(1) & (np.abs(donor_bary).max(1) <= 8.0) + donor_contained = (donor_bary >= -1e-05).all(1) & (donor_bary <= 1 + 1e-05).all(1) + good &= donor_contained + rejection[cy[~donor_contained], cx[~donor_contained]] |= REJECT_DEGENERATE_UV + reconstructed_donor_uv = (tri_uv * donor_bary[:, :, None]).sum(1) + donor_uv_error = np.abs(reconstructed_donor_uv - donor_uv).max(1) + good &= donor_uv_error <= 1e-06 + rejection[cy[~good], cx[~good]] |= REJECT_DEGENERATE_UV + donor_xyz = (tri_native * donor_bary[:, :, None]).sum(1) + (atlas_native, _) = dr.interpolate(native[None], urast, faces) + donor_surface = atlas_native[0, torch.as_tensor(donor_y[cy, cx], device='cuda'), torch.as_tensor(donor_x[cy, cx], device='cuda')].cpu().numpy() + del atlas_native + donor_surface_error = np.abs(donor_surface - donor_xyz).max(1) + surface_checked = good.copy() + surface_ok = donor_surface_error <= 1e-05 + good &= surface_ok + rejection[cy[~surface_ok], cx[~surface_ok]] |= REJECT_DEGENERATE_UV + + def project_native(points): + points_t = torch.as_tensor(points, dtype=torch.float32, device='cuda') + q = torch.cat((points_t, torch.ones_like(points_t[:, :1])), 1) @ E.T + q = q[:, :3] @ K.T + return ((q[:, :2] / q[:, 2:3]).cpu().numpy(), q[:, 2].cpu().numpy()) + (source_uv, source_depth) = project_native(xyz) + (donor_source_uv, donor_depth) = project_native(donor_xyz) + cell = np.floor(np.where(np.isfinite(source_uv), source_uv, -2.0) * RES - 0.5).astype(np.int64) + donor_cell = np.floor(np.where(np.isfinite(donor_source_uv), donor_source_uv, -2.0) * RES - 0.5).astype(np.int64) + footprint_ok = good & np.isfinite(source_uv).all(1) & np.isfinite(donor_source_uv).all(1) & (source_depth > 0) & (donor_depth > 0) + footprint_ok &= (cell[:, 0] >= 0) & (cell[:, 0] + 1 < RES) & (cell[:, 1] >= 0) & (cell[:, 1] + 1 < RES) + footprint_ok &= (donor_cell[:, 0] >= 0) & (donor_cell[:, 0] + 1 < RES) & (donor_cell[:, 1] >= 0) & (donor_cell[:, 1] + 1 < RES) + fi = np.flatnonzero(footprint_ok) + if len(fi): + footprint_ok[fi] &= (hard[cell[fi, 1], cell[fi, 0]] & hard[cell[fi, 1], cell[fi, 0] + 1] + & hard[cell[fi, 1] + 1, cell[fi, 0]] & hard[cell[fi, 1] + 1, cell[fi, 0] + 1] + & hard[donor_cell[fi, 1], donor_cell[fi, 0]] & hard[donor_cell[fi, 1], donor_cell[fi, 0] + 1] + & hard[donor_cell[fi, 1] + 1, donor_cell[fi, 0]] & hard[donor_cell[fi, 1] + 1, donor_cell[fi, 0] + 1]) + donor_strength = rt['F'].grid_sample(trust_t, (torch.as_tensor(donor_source_uv, dtype=torch.float32, device='cuda') * 2 - 1).reshape(1, 1, -1, 2), align_corners=False, padding_mode='zeros')[0, 0, 0].cpu().numpy() + endpoint_strength = rt['F'].grid_sample(trust_t, (torch.as_tensor(source_uv, dtype=torch.float32, device='cuda') * 2 - 1).reshape(1, 1, -1, 2), align_corners=False, padding_mode='zeros')[0, 0, 0].cpu().numpy() + footprint_ok &= (donor_strength >= 1 - 1e-06) & (endpoint_strength >= 1 - 1e-06) + sample_uv = torch.as_tensor(source_uv[footprint_ok], dtype=torch.float32, device='cuda') + donor_sample_uv = torch.as_tensor(donor_source_uv[footprint_ok], dtype=torch.float32, device='cuda') + sampled = rt['F'].grid_sample(src, (sample_uv * 2 - 1).reshape(1, 1, -1, 2), align_corners=False, padding_mode='zeros')[0, :, 0].T.cpu().numpy() + donor_sampled = rt['F'].grid_sample(src, (donor_sample_uv * 2 - 1).reshape(1, 1, -1, 2), align_corners=False, padding_mode='zeros')[0, :, 0].T.cpu().numpy() + source_rows = np.flatnonzero(footprint_ok) + donor_base = base_before[donor_y[cy[source_rows], cx[source_rows]], donor_x[cy[source_rows], cx[source_rows]], :3].astype(np.float32) / 255 + donor_actual = before[donor_y[cy[source_rows], cx[source_rows]], donor_x[cy[source_rows], cx[source_rows]], :3].astype(np.float32) / 255 + donor_rebuilt = donor_base + donor_strength[source_rows, None] * (donor_sampled - donor_base) + donor_project_byte_error = np.full(len(cy), np.inf, np.float32) + donor_project_byte_error[source_rows] = np.abs(donor_rebuilt - donor_actual).max(1) * 255 + donor_project_ok = donor_project_byte_error <= 1.0 + source_rejected = ~footprint_ok | footprint_ok & ~donor_project_ok + rejection[cy[source_rejected], cx[source_rejected]] |= REJECT_SOURCE + continuous_ok = footprint_ok & donor_project_ok + eligible = np.zeros_like(base) + eligible[cy[continuous_ok], cx[continuous_ok]] = True + after = before.copy() + after[cy, cx, :3] = before[donor_y[cy, cx], donor_x[cy, cx], :3] + if not np.array_equal(after[cy, cx, :3], before[donor_y[cy, cx], donor_x[cy, cx], :3]): + raise RuntimeError('projection donor byte fallback failed') + (ey, ex) = np.nonzero(eligible) + anchor = before[donor_y[ey, ex], donor_x[ey, ex], :3].astype(np.float32) / 255 + selected_rows = np.flatnonzero(continuous_ok) + sample_indices = np.searchsorted(source_rows, selected_rows) + after[ey, ex, :3] = np.clip((anchor + sampled[sample_indices] - donor_sampled[sample_indices]) * 255 + 0.5, 0, 255).astype(np.uint8) + if not np.array_equal(after[covered], before[covered]): + raise RuntimeError('UV covered preservation failed') + if after.shape[-1] == 4 and (not np.array_equal(after[..., 3], before[..., 3])): + raise RuntimeError('alpha changed') + if not np.array_equal(after.reshape(-1, after.shape[-1])[active | blocked | unsafe], before.reshape(-1, before.shape[-1])[active | blocked | unsafe]): + raise RuntimeError('active, blocked, or unsafe texels changed') + eligible_ids = np.flatnonzero(eligible.reshape(-1)) + np.savez_compressed( + a.output / 'GUTTER_DIAGNOSTICS.npz', + covered=covered, + unsafe=unsafe.reshape(SIZE, SIZE), + candidate=candidate, + eligible=eligible, + rejection=rejection, + eligible_texel_ids=eligible_ids, + donor_texel_ids=did[continuous_ok], + donor_charts=donor_chart.reshape(-1)[eligible_ids], + source_uv=source_uv[continuous_ok], + donor_source_uv=donor_source_uv[continuous_ok], + barycentric=bary[continuous_ok], + donor_barycentric=donor_bary[continuous_ok], + donor_uv_reconstruction_error=donor_uv_error[continuous_ok], + donor_surface_error=donor_surface_error[continuous_ok], + donor_strength=donor_strength[continuous_ok], + endpoint_strength=endpoint_strength[continuous_ok], + donor_project_byte_error=donor_project_byte_error[continuous_ok], + gutter_radius=np.asarray(a.gutter_radius, dtype=np.int32), + ) + blob = io.BytesIO() + Image.fromarray(after).save(blob, format='PNG') + odoc = copy.deepcopy(pdoc) + obin = fixed._replace_image_payload(odoc, pbin, vi, off, n, blob.getvalue()) + glb = a.output / 'candidate/mesh.glb' + glb.parent.mkdir() + fixed._write_glb(glb, odoc, obin) + preservation = fixed._nonbase_identity_check(pdoc, pbin, glb, vi) + + def render(o, split): + p = o['phase'] + folder = a.output / 'renders' / split / f'x{p[0]:.3f}_y{p[1]:.3f}' + folder.mkdir(parents=True) + cam = copy.deepcopy(camera) + cam['front_atlas_phase'] = p + ca = rm.render_glb(glb, cam, folder / 'repair.png', resolution=RES, rt=rt) + pa = rm.render_glb(pglb, cam, folder / 'project.png', resolution=RES, rt=rt) + with np.load(ca['attrs_npz']) as z: + (ci, cm, cmra) = (z['albedo'], z['mask'].astype(bool), z['mra']) + with np.load(pa['attrs_npz']) as z: + (bi, bm, bmra) = (z['albedo'], z['mask'].astype(bool), z['mra']) + if not np.array_equal(cm, bm) or not np.array_equal(cmra, bmra): + raise RuntimeError('mask/MRA changed') + target = np.zeros_like(ci) + target.reshape(-1, 3)[o['rows'].cpu().numpy()] = o['target'].cpu().numpy() + sel = o['screen'] & cm + quant = torch.as_tensor(after[..., :3], dtype=torch.float32, device='cuda').reshape(-1, 3) / 255 + replay = (quant[o['idx']] * o['w'][..., None]).sum(1).cpu().numpy() + actual = ci.reshape(-1, 3)[o['rows'].cpu().numpy()] + outside = cm & ~o['screen'] + outdelta = float(np.abs(ci - bi)[outside].max()) if outside.any() else 0.0 + changed = cm & (np.abs(ci - bi).max(2) > 0) + r = { + 'phase': p, + 'hardtrust_pixels': int(sel.sum()), + 'repair_vs_source': rm._masked_metrics(ci, target, sel), + 'project_vs_source': rm._masked_metrics(bi, target, sel), + 'repair_vs_project_changed_pixels': int(changed.sum()), + 'repair_vs_project': rm._masked_metrics(ci, bi, changed) if changed.any() else {'valid_pixels': 0}, + 'quantized_atlas_replay_max': float(np.abs(replay - actual).max()), + 'outside_source_hardtrust_max_abs_threshold': OUTSIDE_COLOR_MAX_ABS, + 'outside_source_hardtrust_delta_max_abs': outdelta, + } + r['passes'] = bool(r['quantized_atlas_replay_max'] <= 1e-05 and outside_color_passes(outdelta)) + r.update(mask_exact=True, mra_exact=True, evidence={'repair': bind_file(Path(ca['attrs_npz'])), 'project': bind_file(Path(pa['attrs_npz']))}) + write_once(folder / 'PHASE_RECEIPT.json', r) + return r + results = [render(x, 'grid' if i < len(GRID_PHASES) else 'offgrid') for (i, x) in enumerate(operations)] + changed = np.any(after[..., :3] != before[..., :3], axis=-1) + accepted = all((x['passes'] for x in results)) + gate_ok = accepted and all((preservation.get(k) for k in IDENTITY_GATES)) + receipt = { + 'status': 'valid' if gate_ok else 'invalid', + 'schema': SCHEMA, + 'candidate_glb': str(glb), + 'seconds': time.monotonic() - started, + 'gates': { + 'accepted': gate_ok, + 'UV_covered_exact': True, + 'project_active_exact': True, + 'blocked_exact': True, + 'alpha_exact': True, + 'nonbase_identity': preservation, + 'all_19_phase_renders': len(results) == 19, + }, + 'atlas': { + 'UV_covered_texels': int(covered.sum()), + 'UV_face_charts': int(face_chart.max() + 1), + 'project_active_texels': int(active.sum()), + 'blocked_texels': int(blocked.sum()), + 'unsafe_all_phase_texels': int(unsafe.sum()), + 'gutter_radius': a.gutter_radius, + 'nearest_donor_uncovered_gutter_texels': int(candidate.sum()), + 'continuous_eligible_uncovered_gutter_texels': int(eligible.sum()), + 'changed_texels': int(changed.sum()), + 'rejected_unsafe': int((rejection & REJECT_UNSAFE != 0).sum()), + 'rejected_chart': int((rejection & REJECT_CHART != 0).sum()), + 'rejected_degenerate_uv': int((rejection & REJECT_DEGENERATE_UV != 0).sum()), + 'rejected_source_footprint_or_feather': int((rejection & REJECT_SOURCE != 0).sum()), + 'donor_surface_oracle': { + 'threshold': 1e-05, + 'checked': int(surface_checked.sum()), + 'rejected': int((surface_checked & ~surface_ok).sum()), + 'pre_filter_max': float(donor_surface_error[surface_checked].max()) if surface_checked.any() else 0.0, + 'pre_filter_p99': float(np.quantile(donor_surface_error[surface_checked], 0.99)) if surface_checked.any() else 0.0, + 'accepted_max': float(donor_surface_error[footprint_ok].max()) if footprint_ok.any() else 0.0, + }, + }, + 'phase_results': {'grid': results[:len(GRID_PHASES)], 'offgrid': results[len(GRID_PHASES):]}, + 'diagnostics': str(a.output / 'GUTTER_DIAGNOSTICS.npz'), + 'code_snapshot': str(a.output / 'CODE_SNAPSHOT.py'), + } + write_once(a.output / 'RECEIPT.json', receipt) + write_once(a.output / 'RESULT.json', receipt) + if not gate_ok: + raise RuntimeError('acceptance failed') + print(json.dumps({'status': 'valid', 'candidate': str(glb), 'changed_texels': int(changed.sum())})) +if __name__ == '__main__': + main() diff --git a/supratex/inference/reprojection/_project.py b/supratex/inference/reprojection/_project.py index 42198d3c7f7c8d147fc51907386f5785f8e8edbc..e30946de839ea5f071349113aafc1a9ef0e570ab 100644 --- a/supratex/inference/reprojection/_project.py +++ b/supratex/inference/reprojection/_project.py @@ -1,227 +1,283 @@ -"""Project observed albedo at UV texel centres with measured 19-phase gates. -""" -from __future__ import annotations -import argparse, copy, io, json, os, sys, time -from pathlib import Path -from ._common import OUTSIDE_COLOR_MAX_ABS, outside_color_passes, load_renderer, bind_file -from . import _glb as fixed -SIZE = 4096 -RES = 1024 -TRAIN = tuple(((x / 4, y / 4) for x in range(4) for y in range(4))) -HELD = ((0.125, 0.375), (0.375, 0.625), (0.875, 0.125)) - -def put(p, x): - if p.exists(): - raise FileExistsError(p) - p.write_text(json.dumps(x, indent=2, sort_keys=True, allow_nan=False) + '\n') - -def main(): - ap = argparse.ArgumentParser(description=__doc__) - ap.add_argument('--arm', choices=('B',), required=True) - ap.add_argument('--baseline-glb', type=Path, required=True) - ap.add_argument('--source-albedo', type=Path, required=True) - ap.add_argument('--source-rgba', type=Path, required=True) - ap.add_argument('--camera', type=Path, required=True) - ap.add_argument('--evaluator', type=Path, required=True) - ap.add_argument('--output', type=Path, required=True) - a = ap.parse_args() - if a.output.exists(): - raise FileExistsError(a.output) - for p in (a.baseline_glb, a.source_albedo, a.source_rgba, a.camera, a.evaluator): - if not p.is_file(): - raise FileNotFoundError(p) - a.output.mkdir(parents=True) - put(a.output / 'PLAN.json', {'arm': a.arm, 'baseline_glb': str(a.baseline_glb.absolute()), 'source_albedo': str(a.source_albedo.absolute()), 'source_rgba': str(a.source_rgba.absolute()), 'camera': str(a.camera.absolute()), 'evaluator': str(a.evaluator.absolute()), 'size': SIZE, 'resolution': RES, 'train_phases': TRAIN, 'heldout_phases': HELD, 'source': 'predicted albedo only; no GT read', 'support': 'common sixteen-phase union, intersected with covered source-hard-trusted UV-triangle centres; alpha>=254 source support eroded 2 pixels; two-pixel unsafe-surface footprint blacklist', 'B': 'atlas-centre UV-triangle final-surface projection then bilinear source albedo sample mixed as T0+s*(source-T0)', 'validation': 'actual GLB rerender at 16 train + 3 heldout phases'}) - os.environ['FLEX_GEMM_USE_AUTOTUNE_CACHE'] = '0' - os.environ['FLEX_GEMM_AUTOSAVE_AUTOTUNE_CACHE'] = '0' - import cv2, numpy as np, torch - from PIL import Image - torch.backends.cuda.matmul.allow_tf32 = False - torch.manual_seed(42) - began = time.monotonic() - rm = load_renderer(a.evaluator) - rt = rm._runtime_imports() - dr = rt['dr'] - camera = json.loads(a.camera.read_text()) - (doc, bin0) = fixed._read_glb(a.baseline_glb) - fixed._assert_fixed_scene(doc) - (ii, vi) = fixed._base_color_image_ref(doc) - (im, _, off, n) = fixed._decode_base_image(doc, bin0, ii, vi) - pix = np.asarray(im).copy() - if pix.shape[:2] != (SIZE, SIZE): - raise ValueError('requires 4096 atlas') - vals = rt['_load_textured_mesh'](a.baseline_glb, rt) - (vertices, faces, _, uv, _, _, factor, _, _) = vals - if not np.allclose(factor, 1): - raise RuntimeError('baseColorFactor') - exported = torch.as_tensor(vertices, dtype=torch.float32, device='cuda') - native = exported * torch.tensor((-1.0, 1.0, -1.0), device='cuda') - faces = torch.as_tensor(faces, dtype=torch.int32, device='cuda').contiguous() - uv = torch.as_tensor(uv, dtype=torch.float32, device='cuda') - uv = torch.stack((uv[:, 0], 1 - uv[:, 1]), -1).contiguous() - base = torch.as_tensor(pix[..., :3].copy(), dtype=torch.float32, device='cuda').reshape(-1, 3) / 255 - rgba = np.asarray(Image.open(a.source_rgba).convert('RGBA')) - hard = cv2.erode((rgba[..., 3] >= 254).astype('uint8'), np.ones((5, 5), np.uint8)).astype(bool) - albedo = np.asarray(Image.open(a.source_albedo).convert('RGB'), np.float32) / 255 - if albedo.shape[:2] != (RES, RES) or hard.shape != (RES, RES): - raise ValueError('aligned source must be 1024') - trust = np.clip(cv2.distanceTransform(hard.astype('uint8'), cv2.DIST_L2, 3) / 4.0, 0, 1) - src = torch.as_tensor(albedo, device='cuda').permute(2, 0, 1)[None] - trust_t = torch.as_tensor(trust, dtype=torch.float32, device='cuda')[None, None] - hard_t = torch.as_tensor(hard, device='cuda') - (E, K) = rm._native_camera(camera, 'cuda', rt) - original = rm._exported_glb_camera - - def pc(cam, device, runtime): - (ex, inn) = original(cam, device, runtime) - (dx, dy) = cam.get('front_atlas_phase', (0.0, 0.0)) - inn = inn.clone() - inn[0, 2] += dx / RES - inn[1, 2] += dy / RES - return (ex, inn) - rm._exported_glb_camera = pc - - def op(phase, need_all=False): - cam = copy.deepcopy(camera) - cam['front_atlas_phase'] = phase - (ee, kk) = pc(cam, 'cuda', rt) - dist = rm._camera_distance(cam) - proj = rt['intrinsics_to_projection'](kk, max(0.2, dist - 1), dist + 1) - h = torch.cat((exported, torch.ones_like(exported[:, :1])), 1) - (rast, _) = dr.rasterize(dr.RasterizeCudaContext(device='cuda'), (h @ (proj @ ee).T)[None].contiguous(), faces, (RES, RES)) - (uvs, _) = dr.interpolate(uv[None], rast, faces) - (points, _) = dr.interpolate(native[None], rast, faces) - mask = (rast[0, ..., -1] > 0).cpu().numpy() - rowsall = torch.as_tensor(np.flatnonzero(mask.reshape(-1)), device='cuda') - xyz = points[0].reshape(-1, 3)[rowsall] - q = torch.cat((xyz, torch.ones_like(xyz[:, :1])), 1) @ E.T - q = q[:, :3] @ K.T - suv = q[:, :2] / q[:, 2:3].clamp_min(1e-08) - near = torch.floor(suv * RES).long() - inside = (near[:, 0] >= 0) & (near[:, 0] < RES) & (near[:, 1] >= 0) & (near[:, 1] < RES) & (q[:, 2] > 0) - trusted = torch.zeros(len(rowsall), dtype=torch.bool, device='cuda') - trusted[inside] = hard_t[near[inside, 1], near[inside, 0]] - rows = rowsall[trusted] - tsuv = suv[trusted] - - def fp(rr): - t = uvs[0].reshape(-1, 2)[rr] * SIZE - 0.5 - xy = t.floor().long() - f = t - xy - ids = [] - ws = [] - for (dx, dy) in ((0, 0), (1, 0), (0, 1), (1, 1)): - ids.append((xy[:, 1] + dy) % SIZE * SIZE + (xy[:, 0] + dx) % SIZE) - ws.append((f[:, 0] if dx else 1 - f[:, 0]) * (f[:, 1] if dy else 1 - f[:, 1])) - return (torch.stack(ids, 1), torch.stack(ws, 1)) - (idx, w) = fp(rows) - unsafe = np.zeros(RES * RES, np.uint8) - unsafe[rowsall[~trusted].cpu().numpy()] = 1 - unsafe = cv2.dilate(unsafe.reshape(RES, RES), np.ones((5, 5), np.uint8)) & mask - ur = torch.as_tensor(np.flatnonzero(unsafe.reshape(-1)), device='cuda') - outside = fp(ur)[0] if len(ur) else torch.empty((0, 4), device='cuda', dtype=torch.long) - target = rt['F'].grid_sample(src, (tsuv * 2 - 1).reshape(1, 1, -1, 2), align_corners=False, padding_mode='zeros')[0, :, 0].T - strength = rt['F'].grid_sample(trust_t, (tsuv * 2 - 1).reshape(1, 1, -1, 2), align_corners=False, padding_mode='zeros')[0, 0, 0] - screen = np.zeros(RES * RES, bool) - screen[rows.cpu().numpy()] = True - return {'phase': phase, 'idx': idx, 'w': w, 'outside': outside, 'target': target, 'strength': strength, 'rows': rows, 'screen': screen.reshape(RES, RES), 'uvs': uvs, 'points': points, 'rast': rast, 'trusted': trusted, 'rowsall': rowsall, 'suv': suv} - ops = [op(p) for p in TRAIN] - active = torch.unique(torch.cat([o['idx'].reshape(-1) for o in ops])) - (rastuv, _) = dr.rasterize(dr.RasterizeCudaContext(device='cuda'), torch.cat((uv * 2 - 1, torch.zeros_like(uv[:, :1]), torch.ones_like(uv[:, :1])), 1)[None], faces, (SIZE, SIZE)) - (pxyz, _) = dr.interpolate(native[None], rastuv, faces) - xy = torch.stack((active // SIZE, active % SIZE), 1) - covered = rastuv[0, xy[:, 0], xy[:, 1], -1] > 0 - xyz = pxyz[0, xy[:, 0], xy[:, 1]] - q = torch.cat((xyz, torch.ones_like(xyz[:, :1])), 1) @ E.T - q = q[:, :3] @ K.T - centre_suv = q[:, :2] / q[:, 2:3].clamp_min(1e-08) - ne = torch.floor(centre_suv * RES).long() - inside = (ne[:, 0] >= 0) & (ne[:, 0] < RES) & (ne[:, 1] >= 0) & (ne[:, 1] < RES) & (q[:, 2] > 0) - centre_trust = torch.zeros(len(active), dtype=torch.bool, device='cuda') - centre_trust[inside] = hard_t[ne[inside, 1], ne[inside, 0]] - centre_keep = covered & centre_trust - active = active[centre_keep] - centre_suv = centre_suv[centre_keep] - look = torch.full((SIZE * SIZE,), -1, dtype=torch.long, device='cuda') - look[active] = torch.arange(len(active), device='cuda') - outids = torch.unique(torch.cat([o['outside'].reshape(-1) for o in ops if o['outside'].numel()])) if any((o['outside'].numel() for o in ops)) else torch.empty(0, dtype=torch.long, device='cuda') - blocked = torch.zeros(len(active), dtype=torch.bool, device='cuda') - if len(outids): - where = torch.searchsorted(active, outids) - hit = (where < len(active)) & (active[where.clamp_max(len(active) - 1)] == outids) - blocked[where[hit]] = True - if active.numel() == 0 or not bool((~blocked).any()): - raise RuntimeError('no unblocked trusted atlas texels remain') - np.savez_compressed(a.output / 'COMMON_DOMAIN.npz', active_texel_ids=active.cpu().numpy(), blocked_texel_ids=active[blocked].cpu().numpy()) - for o in ops: - compact = look[o['idx']] - keep = (compact >= 0).all(1) - for key in ('idx', 'w'): - o[key] = o[key][keep] - for key in ('target', 'strength', 'rows'): - o[key] = o[key][keep] - o['compact'] = compact[keep] - o['allowed'] = (~blocked[o['compact']]).float() - o['base'] = (base[o['idx']] * o['w'][..., None]).sum(1) - o['screen'] = np.zeros((RES, RES), bool) - o['screen'].reshape(-1)[o['rows'].cpu().numpy()] = True - mode = 'uv_texel_centre_projection' - corrected = base[active].clone() - s = centre_suv - sample = rt['F'].grid_sample(src, (s * 2 - 1).reshape(1, 1, -1, 2), align_corners=False, padding_mode='zeros')[0, :, 0].T - strength = rt['F'].grid_sample(trust_t, (s * 2 - 1).reshape(1, 1, -1, 2), align_corners=False, padding_mode='zeros')[0, 0, 0] - take = ~blocked - corrected[take] = base[active][take] + strength[take, None] * (sample[take] - base[active][take]) - out = pix.copy() - out.reshape(-1, out.shape[-1])[active.cpu().numpy(), :3] = (corrected.cpu().numpy() * 255 + 0.5).astype(np.uint8) - blob = io.BytesIO() - Image.fromarray(out).save(blob, format='PNG') - od = copy.deepcopy(doc) - ob = fixed._replace_image_payload(od, bin0, vi, off, n, blob.getvalue()) - candidate = a.output / 'candidate' - candidate.mkdir() - glb = candidate / 'mesh.glb' - fixed._write_glb(glb, od, ob) - pres = fixed._nonbase_identity_check(doc, bin0, glb, vi) - - def render(p, split): - folder = a.output / 'renders' / split / f'x{p[0]:.3f}_y{p[1]:.3f}' - folder.mkdir(parents=True) - cam = copy.deepcopy(camera) - cam['front_atlas_phase'] = p - o = op(p) - ca = rm.render_glb(glb, cam, folder / 'candidate.png', resolution=RES, rt=rt) - ba = rm.render_glb(a.baseline_glb, cam, folder / 'baseline.png', resolution=RES, rt=rt) - with np.load(ca['attrs_npz']) as z: - (ci, cm, cmra) = (z['albedo'], z['mask'].astype(bool), z['mra']) - with np.load(ba['attrs_npz']) as z: - (bi, bm, bmra) = (z['albedo'], z['mask'].astype(bool), z['mra']) - if not np.array_equal(cm, bm) or not np.array_equal(cmra, bmra): - raise RuntimeError('mask/MRA changed') - target = np.zeros_like(ci) - target.reshape(-1, 3)[o['rows'].cpu().numpy()] = o['target'].cpu().numpy() - sel = o['screen'] & cm - Image.fromarray((target * 255 + 0.5).astype(np.uint8)).save(folder / 'target_projected.png') - quant = torch.as_tensor(out[..., :3], dtype=torch.float32, device='cuda').reshape(-1, 3) / 255 - replay = (quant[o['idx']] * o['w'][..., None]).sum(1).cpu().numpy() - actual = ci.reshape(-1, 3)[o['rows'].cpu().numpy()] - mx = float(np.abs(replay - actual).max()) - outside = cm & ~o['screen'] - oo = float(np.abs(ci - bi)[outside].max()) if outside.any() else 0.0 - r = {'phase': p, 'pixels': int(sel.sum()), 'candidate_vs_source': rm._masked_metrics(ci, target, sel), 'baseline_vs_source': rm._masked_metrics(bi, target, sel), 'quantized_atlas_replay_max': mx, 'outside_source_hardtrust_max_abs_threshold': OUTSIDE_COLOR_MAX_ABS, 'outside_source_hardtrust_delta_max_abs': oo, 'passes': bool(mx <= 1e-05 and outside_color_passes(oo))} - r.update(mask_exact=True, mra_exact=True, evidence={'candidate': bind_file(Path(ca['attrs_npz'])), 'baseline': bind_file(Path(ba['attrs_npz']))}) - put(folder / 'PHASE_RECEIPT.json', r) - return r - train = [render(x, 'train') for x in TRAIN] - held = [render(x, 'heldout') for x in HELD] - changed = int(np.any(out[..., :3] != pix[..., :3], axis=-1).sum()) - conv = True - ok = conv and all((x['passes'] for x in train + held)) - result = {'status': 'valid' if ok else 'invalid', 'arm': a.arm, 'mode': mode, 'seconds': time.monotonic() - began, 'candidate_glb': str(glb), 'gates': {'accepted': ok, 'common_support_and_protection': True, 'nonbase_identity': pres, 'all_19_phase_renders': len(train) + len(held) == 19}, 'atlas': {'active_texels': int(len(active)), 'blocked_texels': int(blocked.sum()), 'changed_texels': changed, 'alpha_exact': bool(np.array_equal(out[..., 3], pix[..., 3])) if out.shape[-1] == 4 else True}, 'phase_results': {'train': train, 'heldout': held}, 'limits': 'Fixed-shape, observed-source front control. It does not establish novel-view quality or a model improvement.'} - put(a.output / 'RECEIPT.json', result) - put(a.output / 'RESULT.json', result) - if not ok: - raise RuntimeError('acceptance failed') - print(json.dumps({'status': 'valid', 'candidate': str(glb)})) -if __name__ == '__main__': - main() +"""Project observed albedo at UV texel centres with measured 19-phase gates. + +Support is the union of atlas texels touched by the sixteen grid phases, +intersected with UV-covered texel centres whose source projection lands on +hard-trusted observation pixels (alpha >= 254, eroded by two pixels). Bilinear +atlas taps of pixels near untrusted observation regions are blocked. Each kept +texel T0 becomes T0 + s * (source - T0), where source is the bilinear predicted +albedo at the texel centre's projection and s is the observation trust. The +candidate GLB is then re-rendered at all grid and off-grid phases. +""" +from __future__ import annotations +import argparse, copy, io, json, os, sys, time +from pathlib import Path +from ._common import GRID_PHASES, OFFGRID_PHASES, OUTSIDE_COLOR_MAX_ABS, outside_color_passes, load_renderer, bind_file +from . import _glb as fixed +SCHEMA = 'supratex.reprojection.project' +SIZE = 4096 +RES = 1024 + +def put(p, x): + if p.exists(): + raise FileExistsError(p) + p.write_text(json.dumps(x, indent=2, sort_keys=True, allow_nan=False) + '\n') + +def main(): + ap = argparse.ArgumentParser(description=__doc__) + ap.add_argument('--baseline-glb', type=Path, required=True) + ap.add_argument('--source-albedo', type=Path, required=True) + ap.add_argument('--source-rgba', type=Path, required=True) + ap.add_argument('--camera', type=Path, required=True) + ap.add_argument('--renderer', type=Path, required=True) + ap.add_argument('--output', type=Path, required=True) + a = ap.parse_args() + if a.output.exists(): + raise FileExistsError(a.output) + for p in (a.baseline_glb, a.source_albedo, a.source_rgba, a.camera, a.renderer): + if not p.is_file(): + raise FileNotFoundError(p) + a.output.mkdir(parents=True) + put(a.output / 'PLAN.json', { + 'schema': SCHEMA, + 'baseline_glb': str(a.baseline_glb.absolute()), + 'source_albedo': str(a.source_albedo.absolute()), + 'source_rgba': str(a.source_rgba.absolute()), + 'camera': str(a.camera.absolute()), + 'renderer': str(a.renderer.absolute()), + 'size': SIZE, + 'resolution': RES, + 'grid_phases': GRID_PHASES, + 'offgrid_phases': OFFGRID_PHASES, + }) + os.environ['FLEX_GEMM_USE_AUTOTUNE_CACHE'] = '0' + os.environ['FLEX_GEMM_AUTOSAVE_AUTOTUNE_CACHE'] = '0' + import cv2, numpy as np, torch + from PIL import Image + torch.backends.cuda.matmul.allow_tf32 = False + torch.manual_seed(42) + began = time.monotonic() + rm = load_renderer(a.renderer) + rt = rm._runtime_imports() + dr = rt['dr'] + camera = json.loads(a.camera.read_text()) + (doc, bin0) = fixed._read_glb(a.baseline_glb) + fixed._assert_fixed_scene(doc) + (ii, vi) = fixed._base_color_image_ref(doc) + (im, _, off, n) = fixed._decode_base_image(doc, bin0, ii, vi) + pix = np.asarray(im).copy() + if pix.shape[:2] != (SIZE, SIZE): + raise ValueError('requires 4096 atlas') + vals = rt['_load_textured_mesh'](a.baseline_glb, rt) + (vertices, faces, _, uv, _, _, factor, _, _) = vals + if not np.allclose(factor, 1): + raise RuntimeError('baseColorFactor') + exported = torch.as_tensor(vertices, dtype=torch.float32, device='cuda') + native = exported * torch.tensor((-1.0, 1.0, -1.0), device='cuda') + faces = torch.as_tensor(faces, dtype=torch.int32, device='cuda').contiguous() + uv = torch.as_tensor(uv, dtype=torch.float32, device='cuda') + uv = torch.stack((uv[:, 0], 1 - uv[:, 1]), -1).contiguous() + base = torch.as_tensor(pix[..., :3].copy(), dtype=torch.float32, device='cuda').reshape(-1, 3) / 255 + rgba = np.asarray(Image.open(a.source_rgba).convert('RGBA')) + hard = cv2.erode((rgba[..., 3] >= 254).astype('uint8'), np.ones((5, 5), np.uint8)).astype(bool) + albedo = np.asarray(Image.open(a.source_albedo).convert('RGB'), np.float32) / 255 + if albedo.shape[:2] != (RES, RES) or hard.shape != (RES, RES): + raise ValueError('aligned source must be 1024') + trust = np.clip(cv2.distanceTransform(hard.astype('uint8'), cv2.DIST_L2, 3) / 4.0, 0, 1) + src = torch.as_tensor(albedo, device='cuda').permute(2, 0, 1)[None] + trust_t = torch.as_tensor(trust, dtype=torch.float32, device='cuda')[None, None] + hard_t = torch.as_tensor(hard, device='cuda') + (E, K) = rm._native_camera(camera, 'cuda', rt) + original = rm._exported_glb_camera + + def pc(cam, device, runtime): + (ex, inn) = original(cam, device, runtime) + (dx, dy) = cam.get('front_atlas_phase', (0.0, 0.0)) + inn = inn.clone() + inn[0, 2] += dx / RES + inn[1, 2] += dy / RES + return (ex, inn) + rm._exported_glb_camera = pc + + def op(phase): + cam = copy.deepcopy(camera) + cam['front_atlas_phase'] = phase + (ee, kk) = pc(cam, 'cuda', rt) + dist = rm._camera_distance(cam) + proj = rt['intrinsics_to_projection'](kk, max(0.2, dist - 1), dist + 1) + h = torch.cat((exported, torch.ones_like(exported[:, :1])), 1) + (rast, _) = dr.rasterize(dr.RasterizeCudaContext(device='cuda'), (h @ (proj @ ee).T)[None].contiguous(), faces, (RES, RES)) + (uvs, _) = dr.interpolate(uv[None], rast, faces) + (points, _) = dr.interpolate(native[None], rast, faces) + mask = (rast[0, ..., -1] > 0).cpu().numpy() + rowsall = torch.as_tensor(np.flatnonzero(mask.reshape(-1)), device='cuda') + xyz = points[0].reshape(-1, 3)[rowsall] + q = torch.cat((xyz, torch.ones_like(xyz[:, :1])), 1) @ E.T + q = q[:, :3] @ K.T + suv = q[:, :2] / q[:, 2:3].clamp_min(1e-08) + near = torch.floor(suv * RES).long() + inside = (near[:, 0] >= 0) & (near[:, 0] < RES) & (near[:, 1] >= 0) & (near[:, 1] < RES) & (q[:, 2] > 0) + trusted = torch.zeros(len(rowsall), dtype=torch.bool, device='cuda') + trusted[inside] = hard_t[near[inside, 1], near[inside, 0]] + rows = rowsall[trusted] + tsuv = suv[trusted] + + def fp(rr): + t = uvs[0].reshape(-1, 2)[rr] * SIZE - 0.5 + xy = t.floor().long() + f = t - xy + ids = [] + ws = [] + for (dx, dy) in ((0, 0), (1, 0), (0, 1), (1, 1)): + ids.append((xy[:, 1] + dy) % SIZE * SIZE + (xy[:, 0] + dx) % SIZE) + ws.append((f[:, 0] if dx else 1 - f[:, 0]) * (f[:, 1] if dy else 1 - f[:, 1])) + return (torch.stack(ids, 1), torch.stack(ws, 1)) + (idx, w) = fp(rows) + unsafe = np.zeros(RES * RES, np.uint8) + unsafe[rowsall[~trusted].cpu().numpy()] = 1 + unsafe = cv2.dilate(unsafe.reshape(RES, RES), np.ones((5, 5), np.uint8)) & mask + ur = torch.as_tensor(np.flatnonzero(unsafe.reshape(-1)), device='cuda') + outside = fp(ur)[0] if len(ur) else torch.empty((0, 4), device='cuda', dtype=torch.long) + target = rt['F'].grid_sample(src, (tsuv * 2 - 1).reshape(1, 1, -1, 2), align_corners=False, padding_mode='zeros')[0, :, 0].T + strength = rt['F'].grid_sample(trust_t, (tsuv * 2 - 1).reshape(1, 1, -1, 2), align_corners=False, padding_mode='zeros')[0, 0, 0] + screen = np.zeros(RES * RES, bool) + screen[rows.cpu().numpy()] = True + return { + 'phase': phase, + 'idx': idx, + 'w': w, + 'outside': outside, + 'target': target, + 'strength': strength, + 'rows': rows, + 'screen': screen.reshape(RES, RES), + } + ops = [op(p) for p in GRID_PHASES] + active = torch.unique(torch.cat([o['idx'].reshape(-1) for o in ops])) + (rastuv, _) = dr.rasterize(dr.RasterizeCudaContext(device='cuda'), torch.cat((uv * 2 - 1, torch.zeros_like(uv[:, :1]), torch.ones_like(uv[:, :1])), 1)[None], faces, (SIZE, SIZE)) + (pxyz, _) = dr.interpolate(native[None], rastuv, faces) + xy = torch.stack((active // SIZE, active % SIZE), 1) + covered = rastuv[0, xy[:, 0], xy[:, 1], -1] > 0 + xyz = pxyz[0, xy[:, 0], xy[:, 1]] + q = torch.cat((xyz, torch.ones_like(xyz[:, :1])), 1) @ E.T + q = q[:, :3] @ K.T + centre_suv = q[:, :2] / q[:, 2:3].clamp_min(1e-08) + ne = torch.floor(centre_suv * RES).long() + inside = (ne[:, 0] >= 0) & (ne[:, 0] < RES) & (ne[:, 1] >= 0) & (ne[:, 1] < RES) & (q[:, 2] > 0) + centre_trust = torch.zeros(len(active), dtype=torch.bool, device='cuda') + centre_trust[inside] = hard_t[ne[inside, 1], ne[inside, 0]] + centre_keep = covered & centre_trust + active = active[centre_keep] + centre_suv = centre_suv[centre_keep] + look = torch.full((SIZE * SIZE,), -1, dtype=torch.long, device='cuda') + look[active] = torch.arange(len(active), device='cuda') + if any(o['outside'].numel() for o in ops): + outids = torch.unique(torch.cat([o['outside'].reshape(-1) for o in ops if o['outside'].numel()])) + else: + outids = torch.empty(0, dtype=torch.long, device='cuda') + blocked = torch.zeros(len(active), dtype=torch.bool, device='cuda') + if len(outids): + where = torch.searchsorted(active, outids) + hit = (where < len(active)) & (active[where.clamp_max(len(active) - 1)] == outids) + blocked[where[hit]] = True + if active.numel() == 0 or not bool((~blocked).any()): + raise RuntimeError('no unblocked trusted atlas texels remain') + np.savez_compressed(a.output / 'COMMON_DOMAIN.npz', active_texel_ids=active.cpu().numpy(), blocked_texel_ids=active[blocked].cpu().numpy()) + for o in ops: + compact = look[o['idx']] + keep = (compact >= 0).all(1) + for key in ('idx', 'w'): + o[key] = o[key][keep] + for key in ('target', 'strength', 'rows'): + o[key] = o[key][keep] + o['compact'] = compact[keep] + o['allowed'] = (~blocked[o['compact']]).float() + o['base'] = (base[o['idx']] * o['w'][..., None]).sum(1) + o['screen'] = np.zeros((RES, RES), bool) + o['screen'].reshape(-1)[o['rows'].cpu().numpy()] = True + mode = 'uv_texel_centre_projection' + corrected = base[active].clone() + s = centre_suv + sample = rt['F'].grid_sample(src, (s * 2 - 1).reshape(1, 1, -1, 2), align_corners=False, padding_mode='zeros')[0, :, 0].T + strength = rt['F'].grid_sample(trust_t, (s * 2 - 1).reshape(1, 1, -1, 2), align_corners=False, padding_mode='zeros')[0, 0, 0] + take = ~blocked + corrected[take] = base[active][take] + strength[take, None] * (sample[take] - base[active][take]) + out = pix.copy() + out.reshape(-1, out.shape[-1])[active.cpu().numpy(), :3] = (corrected.cpu().numpy() * 255 + 0.5).astype(np.uint8) + blob = io.BytesIO() + Image.fromarray(out).save(blob, format='PNG') + od = copy.deepcopy(doc) + ob = fixed._replace_image_payload(od, bin0, vi, off, n, blob.getvalue()) + candidate = a.output / 'candidate' + candidate.mkdir() + glb = candidate / 'mesh.glb' + fixed._write_glb(glb, od, ob) + pres = fixed._nonbase_identity_check(doc, bin0, glb, vi) + + def render(p, split): + folder = a.output / 'renders' / split / f'x{p[0]:.3f}_y{p[1]:.3f}' + folder.mkdir(parents=True) + cam = copy.deepcopy(camera) + cam['front_atlas_phase'] = p + o = op(p) + ca = rm.render_glb(glb, cam, folder / 'candidate.png', resolution=RES, rt=rt) + ba = rm.render_glb(a.baseline_glb, cam, folder / 'baseline.png', resolution=RES, rt=rt) + with np.load(ca['attrs_npz']) as z: + (ci, cm, cmra) = (z['albedo'], z['mask'].astype(bool), z['mra']) + with np.load(ba['attrs_npz']) as z: + (bi, bm, bmra) = (z['albedo'], z['mask'].astype(bool), z['mra']) + if not np.array_equal(cm, bm) or not np.array_equal(cmra, bmra): + raise RuntimeError('mask/MRA changed') + target = np.zeros_like(ci) + target.reshape(-1, 3)[o['rows'].cpu().numpy()] = o['target'].cpu().numpy() + sel = o['screen'] & cm + Image.fromarray((target * 255 + 0.5).astype(np.uint8)).save(folder / 'target_projected.png', compress_level=1) + quant = torch.as_tensor(out[..., :3], dtype=torch.float32, device='cuda').reshape(-1, 3) / 255 + replay = (quant[o['idx']] * o['w'][..., None]).sum(1).cpu().numpy() + actual = ci.reshape(-1, 3)[o['rows'].cpu().numpy()] + mx = float(np.abs(replay - actual).max()) + outside = cm & ~o['screen'] + oo = float(np.abs(ci - bi)[outside].max()) if outside.any() else 0.0 + r = { + 'phase': p, + 'pixels': int(sel.sum()), + 'candidate_vs_source': rm._masked_metrics(ci, target, sel), + 'baseline_vs_source': rm._masked_metrics(bi, target, sel), + 'quantized_atlas_replay_max': mx, + 'outside_source_hardtrust_max_abs_threshold': OUTSIDE_COLOR_MAX_ABS, + 'outside_source_hardtrust_delta_max_abs': oo, + 'passes': bool(mx <= 1e-05 and outside_color_passes(oo)), + } + r.update(mask_exact=True, mra_exact=True, evidence={'candidate': bind_file(Path(ca['attrs_npz'])), 'baseline': bind_file(Path(ba['attrs_npz']))}) + put(folder / 'PHASE_RECEIPT.json', r) + return r + grid = [render(x, 'grid') for x in GRID_PHASES] + offgrid = [render(x, 'offgrid') for x in OFFGRID_PHASES] + changed = int(np.any(out[..., :3] != pix[..., :3], axis=-1).sum()) + ok = all((x['passes'] for x in grid + offgrid)) + result = { + 'status': 'valid' if ok else 'invalid', + 'schema': SCHEMA, + 'mode': mode, + 'seconds': time.monotonic() - began, + 'candidate_glb': str(glb), + 'gates': { + 'accepted': ok, + 'common_support_and_protection': True, + 'nonbase_identity': pres, + 'all_19_phase_renders': len(grid) + len(offgrid) == 19, + }, + 'atlas': { + 'active_texels': int(len(active)), + 'blocked_texels': int(blocked.sum()), + 'changed_texels': changed, + 'alpha_exact': bool(np.array_equal(out[..., 3], pix[..., 3])) if out.shape[-1] == 4 else True, + }, + 'phase_results': {'grid': grid, 'offgrid': offgrid}, + } + put(a.output / 'RECEIPT.json', result) + put(a.output / 'RESULT.json', result) + if not ok: + raise RuntimeError('acceptance failed') + print(json.dumps({'status': 'valid', 'candidate': str(glb)})) +if __name__ == '__main__': + main() diff --git a/supratex/inference/reprojection/_renderer.py b/supratex/inference/reprojection/_renderer.py index 787d05d3d17ba80403a0c91e73a44ec18cbfd3ab..7263430f59191ce18e56540e5dc33d61e5028d70 100644 --- a/supratex/inference/reprojection/_renderer.py +++ b/supratex/inference/reprojection/_renderer.py @@ -1,170 +1,207 @@ -"""Unlit full-camera rendering for measured atlas preservation gates.""" -from __future__ import annotations -import math -from pathlib import Path -from typing import Any, Mapping -DEFAULT_RESOLUTION = 1024 - -def _load_textured_mesh(glb: Path, rt: dict[str, Any]) -> tuple[Any, Any, Any, Any, Any, Any, Any, float, float]: - (np, trimesh) = (rt['np'], rt['trimesh']) - loaded = trimesh.load(glb, process=False) - mesh = next(iter(loaded.geometry.values())) if isinstance(loaded, trimesh.Scene) else loaded - vertices = np.asarray(mesh.vertices, dtype=np.float32) - faces = np.asarray(mesh.faces, dtype=np.int32) - normals = np.asarray(mesh.vertex_normals, dtype=np.float32) - visual = mesh.visual - if vertices.size == 0 or faces.size == 0 or normals.shape != vertices.shape: - raise ValueError('GLB mesh is empty or lacks usable vertex normals') - if not hasattr(visual, 'uv') or visual.uv is None: - raise ValueError('GLB mesh lacks UV coordinates') - material = visual.material - image = getattr(material, 'baseColorTexture', None) - if image is None: - raise ValueError('GLB mesh lacks a baseColorTexture') - metallic_roughness_image = getattr(material, 'metallicRoughnessTexture', None) - if metallic_roughness_image is None: - raise ValueError('GLB mesh lacks metallicRoughnessTexture; refusing to label textured render as material-aware') - texture = np.asarray(image.convert('RGB'), dtype=np.float32) / 255.0 - metallic_roughness = np.asarray(metallic_roughness_image.convert('RGB'), dtype=np.float32) / 255.0 - raw_factor = getattr(material, 'baseColorFactor', None) - factor = np.asarray([1, 1, 1, 1] if raw_factor is None else raw_factor, dtype=np.float32)[:3] - if factor.max(initial=1.0) > 1.0: - factor /= 255.0 - metallic = getattr(material, 'metallicFactor', None) - roughness = getattr(material, 'roughnessFactor', None) - return (vertices, faces, normals, np.asarray(visual.uv, dtype=np.float32), texture.astype(np.float32), metallic_roughness.astype(np.float32), factor.astype(np.float32), float(1 if metallic is None else metallic), float(1 if roughness is None else roughness)) - -def _srgb_factor_torch(srgb: Any, factor: Any, torch: Any) -> Any: - """Apply a glTF baseColorFactor in linear light and return sRGB samples.""" - linear = torch.where(srgb <= 0.04045, srgb / 12.92, ((srgb + 0.055) / 1.055).pow(2.4)) - linear = linear * torch.as_tensor(factor, dtype=srgb.dtype, device=srgb.device).view(1, 1, 1, 3) - return torch.where(linear <= 0.0031308, linear * 12.92, 1.055 * linear.clamp_min(0).pow(1.0 / 2.4) - 0.055).clamp(0.0, 1.0) - -def _runtime_imports() -> dict[str, Any]: - """Import CUDA dependencies only when rendering is actually requested.""" - import numpy as np - import torch - import torch.nn.functional as F - import trimesh - import nvdiffrast.torch as dr - import utils3d - from PIL import Image - from supratex.renderers.mesh_renderer import intrinsics_to_projection - from ._renderer import _load_textured_mesh - return locals() - -def _check_resolution(resolution: int) -> None: - if resolution <= 0: - raise ValueError('resolution must be positive') - -def _camera_distance(camera: Mapping[str, Any]) -> float: - try: - distance = float(camera['distance']) * float(camera.get('mesh_scale', 1.0)) - except (KeyError, TypeError, ValueError) as error: - raise ValueError('camera requires canonical distance and optional mesh_scale') from error - if not math.isfinite(distance) or distance <= 0: - raise ValueError('camera distance * mesh_scale must be finite and positive') - return distance - -def _native_camera(camera: Mapping[str, Any], device: str, rt: Mapping[str, Any]) -> tuple[Any, Any]: - """Return the baseline camera, accepting a full CV camera when supplied.""" - if 'extrinsic' in camera and ('projection_intrinsic' in camera or 'intrinsic' in camera): - intrinsic_key = 'projection_intrinsic' if 'projection_intrinsic' in camera else 'intrinsic' - extrinsic = rt['torch'].as_tensor(camera['extrinsic'], dtype=rt['torch'].float32, device=device) - intrinsic = rt['torch'].as_tensor(camera[intrinsic_key], dtype=rt['torch'].float32, device=device) - if extrinsic.shape != (4, 4) or intrinsic.shape != (3, 3): - raise ValueError('full camera requires 4x4 extrinsic and 3x3 intrinsic') - return (extrinsic, intrinsic) - distance = _camera_distance(camera) - try: - fov = float(camera['camera_angle_x']) - except (KeyError, TypeError, ValueError) as error: - raise ValueError('canonical camera requires camera_angle_x') from error - if not 0 < fov < math.pi: - raise ValueError('camera_angle_x must be in (0, pi)') - (torch, utils3d) = (rt['torch'], rt['utils3d']) - eye = torch.tensor((0.0, 0.0, distance), dtype=torch.float32, device=device) - target = torch.zeros(3, dtype=torch.float32, device=device) - up = torch.tensor((0.0, 1.0, 0.0), dtype=torch.float32, device=device) - extrinsic = utils3d.torch.extrinsics_look_at(eye, target, up) - fov_t = torch.tensor(fov, dtype=torch.float32, device=device) - return (extrinsic, utils3d.torch.intrinsics_from_fov_xy(fov_t, fov_t)) - -def _exported_glb_camera(camera: Mapping[str, Any], device: str, rt: Mapping[str, Any]): - """Map exported vertices into the native camera, including off-axis poses.""" - (extrinsic, intrinsic) = _native_camera(camera, device, rt) - sign = rt['torch'].diag(rt['torch'].tensor([-1.0, 1.0, -1.0, 1.0], device=device)) - return (extrinsic @ sign, intrinsic) - -def _write_outputs(out: Path, albedo: Any, mra: Any, mask: Any, rt: Mapping[str, Any], save_attrs: bool) -> dict[str, str]: - (np, Image) = (rt['np'], rt['Image']) - out.parent.mkdir(parents=True, exist_ok=True) - albedo_u8 = (np.clip(albedo, 0.0, 1.0) * 255.0 + 0.5).astype(np.uint8) - mask_u8 = (np.asarray(mask, dtype=np.float32) >= 0.5).astype(np.uint8) * 255 - Image.fromarray(albedo_u8, mode='RGB').save(out) - mask_path = out.with_name(out.stem + '_mask.png') - Image.fromarray(mask_u8, mode='L').save(mask_path) - result = {'png': str(out), 'mask_png': str(mask_path)} - if save_attrs: - attrs_path = out.with_suffix('.npz') - np.savez_compressed(attrs_path, albedo=np.asarray(albedo, dtype=np.float32), mra=np.asarray(mra, dtype=np.float32), mask=np.asarray(mask, dtype=np.uint8)) - result['attrs_npz'] = str(attrs_path) - return result - -def render_glb(glb: str | Path, camera: Mapping[str, Any], out: str | Path, resolution: int=DEFAULT_RESOLUTION, *, device: str='cuda', rt: Mapping[str, Any] | None=None, save_attrs: bool=True) -> dict[str, str]: - """Render an exported GLB through the supplied native-frame camera. - - Sample base color and MRA directly with nvdiffrast, without lighting. - """ - _check_resolution(resolution) - runtime = _runtime_imports() if rt is None else dict(rt) - distance = _camera_distance(camera) - (extrinsic, intrinsic) = _exported_glb_camera(camera, device, runtime) - values = runtime['_load_textured_mesh'](Path(glb), runtime) - (vertices, faces, _normals, uv, texture, mr_texture, factor, metallic_factor, roughness_factor) = values - (torch, dr) = (runtime['torch'], runtime['dr']) - projection = runtime['intrinsics_to_projection'](intrinsic, max(0.2, distance - 1.0), distance + 1.0) - vertices_t = torch.as_tensor(vertices, dtype=torch.float32, device=device) - faces_t = torch.as_tensor(faces, dtype=torch.int32, device=device).contiguous() - uv_t = torch.as_tensor(uv, dtype=torch.float32, device=device) - uv_t = torch.stack((uv_t[:, 0], 1.0 - uv_t[:, 1]), dim=-1).contiguous() - vertices_h = torch.cat((vertices_t, torch.ones_like(vertices_t[:, :1])), dim=-1) - clip = vertices_h @ (projection @ extrinsic).T - glctx = dr.RasterizeCudaContext(device=device) - (raster, _) = dr.rasterize(glctx, clip[None].contiguous(), faces_t, (resolution, resolution)) - mask = raster[..., -1] > 0 - (uv_pixel, _) = dr.interpolate(uv_t[None], raster, faces_t) - base = dr.texture(torch.as_tensor(texture, dtype=torch.float32, device=device)[None].contiguous(), uv_pixel, filter_mode='linear') - base = _srgb_factor_torch(base, factor, torch) - mra = dr.texture(torch.as_tensor(mr_texture, dtype=torch.float32, device=device)[None].contiguous(), uv_pixel, filter_mode='linear') - mra = torch.stack((mra[..., 2] * metallic_factor, mra[..., 1] * roughness_factor, torch.ones_like(mra[..., 0])), dim=-1).clamp(0.0, 1.0) - base = torch.where(mask[..., None], base, torch.zeros_like(base)) - mra = torch.where(mask[..., None], mra, torch.zeros_like(mra)) - return _write_outputs(Path(out), base[0].detach().cpu().numpy(), mra[0].detach().cpu().numpy(), mask[0].detach().cpu().numpy(), runtime, save_attrs) - -def _as_rgb(value: Any, name: str): - import numpy as np - array = np.asarray(value, dtype=np.float32) - if array.ndim != 3 or array.shape[-1] != 3 or (not np.isfinite(array).all()) or (array.min() < 0) or (array.max() > 1): - raise ValueError(f'{name} must be finite HxWx3 in [0,1]') - return array - -def _masked_metrics(prediction: Any, target: Any, mask: Any) -> dict[str, float | int | None]: - import numpy as np - (pred, truth, valid) = (_as_rgb(prediction, 'prediction'), _as_rgb(target, 'target'), np.asarray(mask, dtype=bool)) - if pred.shape != truth.shape or valid.shape != pred.shape[:2]: - raise ValueError('prediction, target, and mask shapes must agree') - count = int(valid.sum()) - if count == 0: - return {'valid_pixels': 0, 'masked_mae': None, 'psnr_db': None, 'gradient_vector_mae': None} - delta = pred[valid] - truth[valid] - mse = float(np.mean(delta * delta)) - psnr = 100.0 if mse <= 1e-12 else min(100.0, -10.0 * math.log10(mse)) - - def sobel(image: Any): - padded = np.pad(image, ((1, 1), (1, 1), (0, 0)), mode='edge') - dx = padded[:-2, 2:] + 2 * padded[1:-1, 2:] + padded[2:, 2:] - padded[:-2, :-2] - 2 * padded[1:-1, :-2] - padded[2:, :-2] - dy = padded[2:, :-2] + 2 * padded[2:, 1:-1] + padded[2:, 2:] - padded[:-2, :-2] - 2 * padded[:-2, 1:-1] - padded[:-2, 2:] - return np.stack((dx, dy), axis=-1) - grad_delta = sobel(pred)[valid] - sobel(truth)[valid] - return {'valid_pixels': count, 'masked_mae': float(np.abs(delta).mean()), 'psnr_db': float(psnr), 'gradient_vector_mae': float(np.abs(grad_delta).mean())} +"""Unlit full-camera rendering for measured atlas preservation gates.""" +from __future__ import annotations +import math +from pathlib import Path +from typing import Any, Mapping +DEFAULT_RESOLUTION = 1024 +# Phase gates render the same baseline and candidate many times; reloading a GLB +# decodes two 4096-square textures and recomputes normals each time. Keep the +# last two meshes' render tensors per process, keyed by file version. +_MESH_CACHE: dict[tuple, dict[str, Any]] = {} +_RASTER_CONTEXTS: dict[str, Any] = {} + +def _load_textured_mesh(glb: Path, rt: dict[str, Any]) -> tuple[Any, Any, Any, Any, Any, Any, Any, float, float]: + (np, trimesh) = (rt['np'], rt['trimesh']) + loaded = trimesh.load(glb, process=False) + mesh = next(iter(loaded.geometry.values())) if isinstance(loaded, trimesh.Scene) else loaded + vertices = np.asarray(mesh.vertices, dtype=np.float32) + faces = np.asarray(mesh.faces, dtype=np.int32) + normals = np.asarray(mesh.vertex_normals, dtype=np.float32) + visual = mesh.visual + if vertices.size == 0 or faces.size == 0 or normals.shape != vertices.shape: + raise ValueError('GLB mesh is empty or lacks usable vertex normals') + if not hasattr(visual, 'uv') or visual.uv is None: + raise ValueError('GLB mesh lacks UV coordinates') + material = visual.material + image = getattr(material, 'baseColorTexture', None) + if image is None: + raise ValueError('GLB mesh lacks a baseColorTexture') + metallic_roughness_image = getattr(material, 'metallicRoughnessTexture', None) + if metallic_roughness_image is None: + raise ValueError('GLB mesh lacks metallicRoughnessTexture; refusing to label textured render as material-aware') + texture = np.asarray(image.convert('RGB'), dtype=np.float32) / 255.0 + metallic_roughness = np.asarray(metallic_roughness_image.convert('RGB'), dtype=np.float32) / 255.0 + raw_factor = getattr(material, 'baseColorFactor', None) + factor = np.asarray([1, 1, 1, 1] if raw_factor is None else raw_factor, dtype=np.float32)[:3] + if factor.max(initial=1.0) > 1.0: + factor /= 255.0 + metallic = getattr(material, 'metallicFactor', None) + roughness = getattr(material, 'roughnessFactor', None) + return (vertices, faces, normals, np.asarray(visual.uv, dtype=np.float32), texture.astype(np.float32), metallic_roughness.astype(np.float32), factor.astype(np.float32), float(1 if metallic is None else metallic), float(1 if roughness is None else roughness)) + +def _srgb_factor_torch(srgb: Any, factor: Any, torch: Any) -> Any: + """Apply a glTF baseColorFactor in linear light and return sRGB samples.""" + linear = torch.where(srgb <= 0.04045, srgb / 12.92, ((srgb + 0.055) / 1.055).pow(2.4)) + linear = linear * torch.as_tensor(factor, dtype=srgb.dtype, device=srgb.device).view(1, 1, 1, 3) + return torch.where(linear <= 0.0031308, linear * 12.92, 1.055 * linear.clamp_min(0).pow(1.0 / 2.4) - 0.055).clamp(0.0, 1.0) + +def _runtime_imports() -> dict[str, Any]: + """Import CUDA dependencies only when rendering is actually requested.""" + import numpy as np + import torch + import torch.nn.functional as F + import trimesh + import nvdiffrast.torch as dr + import utils3d + from PIL import Image + from supratex.renderers.mesh_renderer import intrinsics_to_projection + from ._renderer import _load_textured_mesh + return locals() + +def _check_resolution(resolution: int) -> None: + if resolution <= 0: + raise ValueError('resolution must be positive') + +def _camera_distance(camera: Mapping[str, Any]) -> float: + try: + distance = float(camera['distance']) * float(camera.get('mesh_scale', 1.0)) + except (KeyError, TypeError, ValueError) as error: + raise ValueError('camera requires canonical distance and optional mesh_scale') from error + if not math.isfinite(distance) or distance <= 0: + raise ValueError('camera distance * mesh_scale must be finite and positive') + return distance + +def _native_camera(camera: Mapping[str, Any], device: str, rt: Mapping[str, Any]) -> tuple[Any, Any]: + """Return the baseline camera, accepting a full CV camera when supplied.""" + if 'extrinsic' in camera and ('projection_intrinsic' in camera or 'intrinsic' in camera): + intrinsic_key = 'projection_intrinsic' if 'projection_intrinsic' in camera else 'intrinsic' + extrinsic = rt['torch'].as_tensor(camera['extrinsic'], dtype=rt['torch'].float32, device=device) + intrinsic = rt['torch'].as_tensor(camera[intrinsic_key], dtype=rt['torch'].float32, device=device) + if extrinsic.shape != (4, 4) or intrinsic.shape != (3, 3): + raise ValueError('full camera requires 4x4 extrinsic and 3x3 intrinsic') + return (extrinsic, intrinsic) + distance = _camera_distance(camera) + try: + fov = float(camera['camera_angle_x']) + except (KeyError, TypeError, ValueError) as error: + raise ValueError('canonical camera requires camera_angle_x') from error + if not 0 < fov < math.pi: + raise ValueError('camera_angle_x must be in (0, pi)') + (torch, utils3d) = (rt['torch'], rt['utils3d']) + eye = torch.tensor((0.0, 0.0, distance), dtype=torch.float32, device=device) + target = torch.zeros(3, dtype=torch.float32, device=device) + up = torch.tensor((0.0, 1.0, 0.0), dtype=torch.float32, device=device) + extrinsic = utils3d.torch.extrinsics_look_at(eye, target, up) + fov_t = torch.tensor(fov, dtype=torch.float32, device=device) + return (extrinsic, utils3d.torch.intrinsics_from_fov_xy(fov_t, fov_t)) + +def _exported_glb_camera(camera: Mapping[str, Any], device: str, rt: Mapping[str, Any]): + """Map exported vertices into the native camera, including off-axis poses.""" + (extrinsic, intrinsic) = _native_camera(camera, device, rt) + sign = rt['torch'].diag(rt['torch'].tensor([-1.0, 1.0, -1.0, 1.0], device=device)) + return (extrinsic @ sign, intrinsic) + +def _write_outputs(out: Path, albedo: Any, mra: Any, mask: Any, rt: Mapping[str, Any], save_attrs: bool) -> dict[str, str]: + (np, Image) = (rt['np'], rt['Image']) + out.parent.mkdir(parents=True, exist_ok=True) + albedo_u8 = (np.clip(albedo, 0.0, 1.0) * 255.0 + 0.5).astype(np.uint8) + mask_u8 = (np.asarray(mask, dtype=np.float32) >= 0.5).astype(np.uint8) * 255 + # Diagnostic previews and evidence arrays use fast compression; the stored + # values are unchanged and receipts bind whatever bytes are written. + Image.fromarray(albedo_u8, mode='RGB').save(out, compress_level=1) + mask_path = out.with_name(out.stem + '_mask.png') + Image.fromarray(mask_u8, mode='L').save(mask_path, compress_level=1) + result = {'png': str(out), 'mask_png': str(mask_path)} + if save_attrs: + attrs_path = out.with_suffix('.npz') + _save_npz(attrs_path, np, albedo=np.asarray(albedo, dtype=np.float32), mra=np.asarray(mra, dtype=np.float32), mask=np.asarray(mask, dtype=np.uint8)) + result['attrs_npz'] = str(attrs_path) + return result + +def _save_npz(path: Path, np: Any, **arrays: Any) -> None: + """Write a standard compressed .npz with zlib level 1 (numpy uses level 6).""" + import zipfile + with zipfile.ZipFile(path, 'w', zipfile.ZIP_DEFLATED, compresslevel=1) as archive: + for name, value in arrays.items(): + with archive.open(f'{name}.npy', 'w', force_zip64=True) as stream: + np.lib.format.write_array(stream, np.asanyarray(value), allow_pickle=False) + +def _mesh_tensors(glb: Path, device: str, runtime: Mapping[str, Any]) -> dict[str, Any]: + """Load a GLB once per file version and keep its render inputs on the device.""" + path = glb.resolve() + stat = path.stat() + key = (str(path), stat.st_mtime_ns, stat.st_size, device, id(runtime['_load_textured_mesh'])) + cached = _MESH_CACHE.get(key) + if cached is None: + torch = runtime['torch'] + (vertices, faces, _normals, uv, texture, mr_texture, factor, metallic_factor, roughness_factor) = runtime['_load_textured_mesh'](glb, runtime) + uv_t = torch.as_tensor(uv, dtype=torch.float32, device=device) + cached = { + 'vertices': torch.as_tensor(vertices, dtype=torch.float32, device=device), + 'faces': torch.as_tensor(faces, dtype=torch.int32, device=device).contiguous(), + 'uv': torch.stack((uv_t[:, 0], 1.0 - uv_t[:, 1]), dim=-1).contiguous(), + 'texture': torch.as_tensor(texture, dtype=torch.float32, device=device)[None].contiguous(), + 'mr_texture': torch.as_tensor(mr_texture, dtype=torch.float32, device=device)[None].contiguous(), + 'factor': factor, 'metallic_factor': metallic_factor, 'roughness_factor': roughness_factor, + } + while len(_MESH_CACHE) >= 2: + _MESH_CACHE.pop(next(iter(_MESH_CACHE))) + _MESH_CACHE[key] = cached + return cached + +def render_glb(glb: str | Path, camera: Mapping[str, Any], out: str | Path, resolution: int=DEFAULT_RESOLUTION, *, device: str='cuda', rt: Mapping[str, Any] | None=None, save_attrs: bool=True) -> dict[str, str]: + """Render an exported GLB through the supplied native-frame camera. + + Sample base color and MRA directly with nvdiffrast, without lighting. + """ + _check_resolution(resolution) + runtime = _runtime_imports() if rt is None else dict(rt) + distance = _camera_distance(camera) + (extrinsic, intrinsic) = _exported_glb_camera(camera, device, runtime) + mesh = _mesh_tensors(Path(glb), device, runtime) + (torch, dr) = (runtime['torch'], runtime['dr']) + (vertices_t, faces_t, uv_t) = (mesh['vertices'], mesh['faces'], mesh['uv']) + (factor, metallic_factor, roughness_factor) = (mesh['factor'], mesh['metallic_factor'], mesh['roughness_factor']) + projection = runtime['intrinsics_to_projection'](intrinsic, max(0.2, distance - 1.0), distance + 1.0) + vertices_h = torch.cat((vertices_t, torch.ones_like(vertices_t[:, :1])), dim=-1) + clip = vertices_h @ (projection @ extrinsic).T + glctx = _RASTER_CONTEXTS.get(device) + if glctx is None: + glctx = _RASTER_CONTEXTS[device] = dr.RasterizeCudaContext(device=device) + (raster, _) = dr.rasterize(glctx, clip[None].contiguous(), faces_t, (resolution, resolution)) + mask = raster[..., -1] > 0 + (uv_pixel, _) = dr.interpolate(uv_t[None], raster, faces_t) + base = dr.texture(mesh['texture'], uv_pixel, filter_mode='linear') + base = _srgb_factor_torch(base, factor, torch) + mra = dr.texture(mesh['mr_texture'], uv_pixel, filter_mode='linear') + mra = torch.stack((mra[..., 2] * metallic_factor, mra[..., 1] * roughness_factor, torch.ones_like(mra[..., 0])), dim=-1).clamp(0.0, 1.0) + base = torch.where(mask[..., None], base, torch.zeros_like(base)) + mra = torch.where(mask[..., None], mra, torch.zeros_like(mra)) + return _write_outputs(Path(out), base[0].detach().cpu().numpy(), mra[0].detach().cpu().numpy(), mask[0].detach().cpu().numpy(), runtime, save_attrs) + +def _as_rgb(value: Any, name: str): + import numpy as np + array = np.asarray(value, dtype=np.float32) + if array.ndim != 3 or array.shape[-1] != 3 or (not np.isfinite(array).all()) or (array.min() < 0) or (array.max() > 1): + raise ValueError(f'{name} must be finite HxWx3 in [0,1]') + return array + +def _masked_metrics(prediction: Any, target: Any, mask: Any) -> dict[str, float | int | None]: + import numpy as np + (pred, truth, valid) = (_as_rgb(prediction, 'prediction'), _as_rgb(target, 'target'), np.asarray(mask, dtype=bool)) + if pred.shape != truth.shape or valid.shape != pred.shape[:2]: + raise ValueError('prediction, target, and mask shapes must agree') + count = int(valid.sum()) + if count == 0: + return {'valid_pixels': 0, 'masked_mae': None, 'psnr_db': None, 'gradient_vector_mae': None} + delta = pred[valid] - truth[valid] + mse = float(np.mean(delta * delta)) + psnr = 100.0 if mse <= 1e-12 else min(100.0, -10.0 * math.log10(mse)) + + def sobel(image: Any): + padded = np.pad(image, ((1, 1), (1, 1), (0, 0)), mode='edge') + dx = padded[:-2, 2:] + 2 * padded[1:-1, 2:] + padded[2:, 2:] - padded[:-2, :-2] - 2 * padded[1:-1, :-2] - padded[2:, :-2] + dy = padded[2:, :-2] + 2 * padded[2:, 1:-1] + padded[2:, 2:] - padded[:-2, :-2] - 2 * padded[:-2, 1:-1] - padded[:-2, 2:] + return np.stack((dx, dy), axis=-1) + grad_delta = sobel(pred)[valid] - sobel(truth)[valid] + return {'valid_pixels': count, 'masked_mae': float(np.abs(delta).mean()), 'psnr_db': float(psnr), 'gradient_vector_mae': float(np.abs(grad_delta).mean())} diff --git a/supratex/inference/reprojection/_transform.py b/supratex/inference/reprojection/_transform.py index de9a310df63db902408ef8e3a05c051c13b23192..2867f3dded33114b6e475b559eac18d7a715e348 100644 --- a/supratex/inference/reprojection/_transform.py +++ b/supratex/inference/reprojection/_transform.py @@ -1,166 +1,166 @@ -"""Temporary back-camera transforms and front-priority atlas stitching.""" -from __future__ import annotations - -import copy -import io -import math -from pathlib import Path -from typing import Any - -import numpy as np -from PIL import Image - -from . import _glb - -CAMERA_PROTOCOL = "texture_eval_wild.back_view_camera.v4" -CAMERA_SCHEMA = "texture_eval_wild.bank.back_normal_pack.v3" - - -def camera_spec(cameras: dict[str, Any], name: str) -> dict[str, Any]: - """Validate a current camera pack and return a full native-frame camera.""" - if cameras.get("protocol") != CAMERA_PROTOCOL or cameras.get("schema") != CAMERA_SCHEMA: - raise ValueError("reprojection requires a current camera v4 pack") - view = cameras[name] - if not str(view.get("convention", "")).startswith("OpenCV world-to-camera"): - raise ValueError("camera must use OpenCV world-to-camera coordinates") - if list(view.get("resolution", [])) != [1024, 1024]: - raise ValueError("reprojection camera must have a 1024 square image") - extrinsic = np.asarray(view["extrinsics_w2c"], dtype=np.float64) - intrinsic = np.asarray(view["intrinsics_normalized"], dtype=np.float64) - eye = np.asarray(view["eye"], dtype=np.float64) - look_at = np.asarray(view["look_at"], dtype=np.float64) - if extrinsic.shape != (4, 4) or intrinsic.shape != (3, 3) or eye.shape != (3,) or look_at.shape != (3,): - raise ValueError("invalid camera array dimensions") - if not all(np.isfinite(x).all() for x in (extrinsic, intrinsic, eye, look_at)): - raise ValueError("camera arrays must be finite") - rotation = extrinsic[:3, :3] - if not np.allclose(extrinsic[3], [0, 0, 0, 1], atol=1e-6) or not np.allclose(rotation @ rotation.T, np.eye(3), atol=1e-5) or not np.isclose(np.linalg.det(rotation), 1, atol=1e-5): - raise ValueError("camera must contain a proper rigid transform") - distance = float(np.linalg.norm(eye - look_at)) - if distance <= 0 or not np.allclose(-rotation.T @ extrinsic[:3, 3], eye, atol=1e-5): - raise ValueError("camera eye binding is inconsistent") - if not np.allclose(rotation @ (look_at - eye) / distance, [0, 0, 1], atol=1e-5): - raise ValueError("camera optical direction differs from its look-at point") - focal = float(intrinsic[0, 0]) - if focal <= 0 or not np.allclose(intrinsic, [[focal, 0, .5], [0, focal, .5], [0, 0, 1]], atol=1e-5): - raise ValueError("camera requires square centered perspective intrinsics") - return {"extrinsic": extrinsic.tolist(), "projection_intrinsic": intrinsic.tolist(), - "distance": distance, "mesh_scale": 1.0, "camera_angle_x": 2 * math.atan(.5 / focal)} - - -def back_to_front(view: dict[str, Any]) -> tuple[np.ndarray, np.ndarray, float]: - """Map an off-axis back camera onto a canonical +Z camera in GLB space.""" - extrinsic = np.asarray(view["extrinsics_w2c"], dtype=np.float64) - eye, target = np.asarray(view["eye"], dtype=np.float64), np.asarray(view["look_at"], dtype=np.float64) - distance = float(np.linalg.norm(eye - target)) - front = np.diag([1., -1., -1.]) - rotation = front.T @ extrinsic[:3, :3] - translation = front.T @ (extrinsic[:3, 3] - [0., 0., distance]) - if not np.allclose(rotation @ eye + translation, [0, 0, distance], atol=1e-5) or not np.allclose(rotation @ target + translation, 0, atol=1e-5): - raise ValueError("back-camera transform did not recover the canonical pose") - sign = np.diag([-1., 1., -1.]) - return sign @ rotation @ sign, sign @ translation, distance - - -def attribute_bytes(path: Path, name: str) -> bytes: - document, binary = _glb._read_glb(path) - _glb._assert_fixed_scene(document) - primitive = document["meshes"][0]["primitives"][0] - accessor = document["accessors"][primitive["attributes"][name]] - view = document["bufferViews"][accessor["bufferView"]] - width = {"VEC2": 2, "VEC3": 3}.get(accessor.get("type")) - if accessor.get("componentType") != 5126 or width is None or accessor.get("sparse"): - raise ValueError(f"{name} must be a dense float32 vector attribute") - stride = width * 4 - if int(view.get("byteStride", stride)) != stride: - raise ValueError(f"{name} must be stored without interleaving") - start = int(view.get("byteOffset", 0)) + int(accessor.get("byteOffset", 0)) - end = start + int(accessor["count"]) * stride - if start < 0 or end > len(binary) or end > int(view.get("byteOffset", 0)) + int(view["byteLength"]): - raise ValueError(f"{name} attribute exceeds its binary bounds") - return binary[start:end] - - -def transformed_payload(source: Path, rotation: np.ndarray, translation: np.ndarray): - """Return a temporary transform without changing other document content.""" - document, binary = _glb._read_glb(source) - _glb._assert_fixed_scene(document) - edited = bytearray(binary) - primitive = document["meshes"][0]["primitives"][0] - for name in ("POSITION", "NORMAL"): - raw = attribute_bytes(source, name) - accessor = document["accessors"][primitive["attributes"][name]] - view = document["bufferViews"][accessor["bufferView"]] - start = int(view.get("byteOffset", 0)) + int(accessor.get("byteOffset", 0)) - values = np.frombuffer(raw, dtype=" None: - """Change only temporary POSITION/NORMAL data and POSITION bounds.""" - document, binary = transformed_payload(source, rotation, translation) - _glb._write_glb(destination, document, binary) - - -def atlas(path: Path): - document, binary = _glb._read_glb(path) - _glb._assert_fixed_scene(document) - image_index, view_index = _glb._base_color_image_ref(document) - image, _, offset, length = _glb._decode_base_image(document, binary, image_index, view_index) - return document, binary, view_index, offset, length, np.asarray(image).copy() - - -def front_protection(domain: Path) -> np.ndarray: - with np.load(domain, allow_pickle=False) as data: - active, blocked = data["active_texel_ids"], data["blocked_texel_ids"] - for values in (active, blocked): - if values.ndim != 1 or not np.issubdtype(values.dtype, np.integer) or np.any(values < 0) or np.any(values >= 4096 ** 2) or len(np.unique(values)) != len(values): - raise ValueError("invalid atlas protection domain") - if not np.isin(blocked, active).all(): - raise ValueError("blocked atlas texels must belong to the active domain") - protected = np.setdiff1d(active, blocked) - if not protected.size: - raise ValueError("front postprocess has no protected trusted atlas texels") - return protected - - -def validate_stitch(front: Path, back: Path, final: Path, protected: np.ndarray) -> dict[str, Any]: - fdoc, fbin, fview, _, _, fpixels = atlas(front) - _, _, _, _, _, bpixels = atlas(back) - _, _, _, _, _, pixels = atlas(final) - if fpixels.shape != bpixels.shape or pixels.shape != fpixels.shape: - raise ValueError("front, back, and final atlas shapes differ") - expected = bpixels.copy() - expected.reshape(-1, expected.shape[-1])[protected, :3] = fpixels.reshape(-1, fpixels.shape[-1])[protected, :3] - if expected.shape[-1] == 4: - expected[..., 3] = fpixels[..., 3] - if not np.array_equal(pixels, expected): - raise ValueError("final atlas does not follow exact front-priority stitching") - for name in ("POSITION", "NORMAL", "TEXCOORD_0"): - if attribute_bytes(front, name) != attribute_bytes(final, name): - raise ValueError(f"final mesh changed {name}") - identity = _glb._nonbase_identity_check(fdoc, fbin, final, fview) - return {"front_priority_exact": True, "alpha_exact": True, "position_exact": True, - "normal_exact": True, "uv_exact": True, "nonbase_identity": identity, - "protected_texels": int(protected.size), "final_19_phase_rendered": False} - - -def stitch(front: Path, back: Path, protected: np.ndarray, destination: Path) -> dict[str, Any]: - document, binary, view, offset, length, front_pixels = atlas(front) - _, _, _, _, _, back_pixels = atlas(back) - if front_pixels.shape != back_pixels.shape: - raise ValueError("front and back atlas shapes differ") - merged = back_pixels.copy() - merged.reshape(-1, merged.shape[-1])[protected, :3] = front_pixels.reshape(-1, front_pixels.shape[-1])[protected, :3] - if merged.shape[-1] == 4: - merged[..., 3] = front_pixels[..., 3] - payload = io.BytesIO() - Image.fromarray(merged).save(payload, format="PNG") - output_document = copy.deepcopy(document) - output_binary = _glb._replace_image_payload(output_document, binary, view, offset, length, payload.getvalue()) - _glb._write_glb(destination, output_document, output_binary) - return validate_stitch(front, back, destination, protected) +"""Temporary back-camera transforms and front-priority atlas stitching.""" +from __future__ import annotations + +import copy +import io +import math +from pathlib import Path +from typing import Any + +import numpy as np +from PIL import Image + +from . import _glb + +CAMERA_PROTOCOL = "supratex.back_view_camera.v4" +CAMERA_SCHEMA = "supratex.back_normal_pack.v3" + + +def camera_spec(cameras: dict[str, Any], name: str) -> dict[str, Any]: + """Validate a current camera pack and return a full native-frame camera.""" + if cameras.get("protocol") != CAMERA_PROTOCOL or cameras.get("schema") != CAMERA_SCHEMA: + raise ValueError("reprojection requires a current camera v4 pack") + view = cameras[name] + if not str(view.get("convention", "")).startswith("OpenCV world-to-camera"): + raise ValueError("camera must use OpenCV world-to-camera coordinates") + if list(view.get("resolution", [])) != [1024, 1024]: + raise ValueError("reprojection camera must have a 1024 square image") + extrinsic = np.asarray(view["extrinsics_w2c"], dtype=np.float64) + intrinsic = np.asarray(view["intrinsics_normalized"], dtype=np.float64) + eye = np.asarray(view["eye"], dtype=np.float64) + look_at = np.asarray(view["look_at"], dtype=np.float64) + if extrinsic.shape != (4, 4) or intrinsic.shape != (3, 3) or eye.shape != (3,) or look_at.shape != (3,): + raise ValueError("invalid camera array dimensions") + if not all(np.isfinite(x).all() for x in (extrinsic, intrinsic, eye, look_at)): + raise ValueError("camera arrays must be finite") + rotation = extrinsic[:3, :3] + if not np.allclose(extrinsic[3], [0, 0, 0, 1], atol=1e-6) or not np.allclose(rotation @ rotation.T, np.eye(3), atol=1e-5) or not np.isclose(np.linalg.det(rotation), 1, atol=1e-5): + raise ValueError("camera must contain a proper rigid transform") + distance = float(np.linalg.norm(eye - look_at)) + if distance <= 0 or not np.allclose(-rotation.T @ extrinsic[:3, 3], eye, atol=1e-5): + raise ValueError("camera eye binding is inconsistent") + if not np.allclose(rotation @ (look_at - eye) / distance, [0, 0, 1], atol=1e-5): + raise ValueError("camera optical direction differs from its look-at point") + focal = float(intrinsic[0, 0]) + if focal <= 0 or not np.allclose(intrinsic, [[focal, 0, .5], [0, focal, .5], [0, 0, 1]], atol=1e-5): + raise ValueError("camera requires square centered perspective intrinsics") + return {"extrinsic": extrinsic.tolist(), "projection_intrinsic": intrinsic.tolist(), + "distance": distance, "mesh_scale": 1.0, "camera_angle_x": 2 * math.atan(.5 / focal)} + + +def back_to_front(view: dict[str, Any]) -> tuple[np.ndarray, np.ndarray, float]: + """Map an off-axis back camera onto a canonical +Z camera in GLB space.""" + extrinsic = np.asarray(view["extrinsics_w2c"], dtype=np.float64) + eye, target = np.asarray(view["eye"], dtype=np.float64), np.asarray(view["look_at"], dtype=np.float64) + distance = float(np.linalg.norm(eye - target)) + front = np.diag([1., -1., -1.]) + rotation = front.T @ extrinsic[:3, :3] + translation = front.T @ (extrinsic[:3, 3] - [0., 0., distance]) + if not np.allclose(rotation @ eye + translation, [0, 0, distance], atol=1e-5) or not np.allclose(rotation @ target + translation, 0, atol=1e-5): + raise ValueError("back-camera transform did not recover the canonical pose") + sign = np.diag([-1., 1., -1.]) + return sign @ rotation @ sign, sign @ translation, distance + + +def attribute_bytes(path: Path, name: str) -> bytes: + document, binary = _glb._read_glb(path) + _glb._assert_fixed_scene(document) + primitive = document["meshes"][0]["primitives"][0] + accessor = document["accessors"][primitive["attributes"][name]] + view = document["bufferViews"][accessor["bufferView"]] + width = {"VEC2": 2, "VEC3": 3}.get(accessor.get("type")) + if accessor.get("componentType") != 5126 or width is None or accessor.get("sparse"): + raise ValueError(f"{name} must be a dense float32 vector attribute") + stride = width * 4 + if int(view.get("byteStride", stride)) != stride: + raise ValueError(f"{name} must be stored without interleaving") + start = int(view.get("byteOffset", 0)) + int(accessor.get("byteOffset", 0)) + end = start + int(accessor["count"]) * stride + if start < 0 or end > len(binary) or end > int(view.get("byteOffset", 0)) + int(view["byteLength"]): + raise ValueError(f"{name} attribute exceeds its binary bounds") + return binary[start:end] + + +def transformed_payload(source: Path, rotation: np.ndarray, translation: np.ndarray): + """Return a temporary transform without changing other document content.""" + document, binary = _glb._read_glb(source) + _glb._assert_fixed_scene(document) + edited = bytearray(binary) + primitive = document["meshes"][0]["primitives"][0] + for name in ("POSITION", "NORMAL"): + raw = attribute_bytes(source, name) + accessor = document["accessors"][primitive["attributes"][name]] + view = document["bufferViews"][accessor["bufferView"]] + start = int(view.get("byteOffset", 0)) + int(accessor.get("byteOffset", 0)) + values = np.frombuffer(raw, dtype=" None: + """Change only temporary POSITION/NORMAL data and POSITION bounds.""" + document, binary = transformed_payload(source, rotation, translation) + _glb._write_glb(destination, document, binary) + + +def atlas(path: Path): + document, binary = _glb._read_glb(path) + _glb._assert_fixed_scene(document) + image_index, view_index = _glb._base_color_image_ref(document) + image, _, offset, length = _glb._decode_base_image(document, binary, image_index, view_index) + return document, binary, view_index, offset, length, np.asarray(image).copy() + + +def front_protection(domain: Path) -> np.ndarray: + with np.load(domain, allow_pickle=False) as data: + active, blocked = data["active_texel_ids"], data["blocked_texel_ids"] + for values in (active, blocked): + if values.ndim != 1 or not np.issubdtype(values.dtype, np.integer) or np.any(values < 0) or np.any(values >= 4096 ** 2) or len(np.unique(values)) != len(values): + raise ValueError("invalid atlas protection domain") + if not np.isin(blocked, active).all(): + raise ValueError("blocked atlas texels must belong to the active domain") + protected = np.setdiff1d(active, blocked) + if not protected.size: + raise ValueError("front postprocess has no protected trusted atlas texels") + return protected + + +def validate_stitch(front: Path, back: Path, final: Path, protected: np.ndarray) -> dict[str, Any]: + fdoc, fbin, fview, _, _, fpixels = atlas(front) + _, _, _, _, _, bpixels = atlas(back) + _, _, _, _, _, pixels = atlas(final) + if fpixels.shape != bpixels.shape or pixels.shape != fpixels.shape: + raise ValueError("front, back, and final atlas shapes differ") + expected = bpixels.copy() + expected.reshape(-1, expected.shape[-1])[protected, :3] = fpixels.reshape(-1, fpixels.shape[-1])[protected, :3] + if expected.shape[-1] == 4: + expected[..., 3] = fpixels[..., 3] + if not np.array_equal(pixels, expected): + raise ValueError("final atlas does not follow exact front-priority stitching") + for name in ("POSITION", "NORMAL", "TEXCOORD_0"): + if attribute_bytes(front, name) != attribute_bytes(final, name): + raise ValueError(f"final mesh changed {name}") + identity = _glb._nonbase_identity_check(fdoc, fbin, final, fview) + return {"front_priority_exact": True, "alpha_exact": True, "position_exact": True, + "normal_exact": True, "uv_exact": True, "nonbase_identity": identity, + "protected_texels": int(protected.size), "final_19_phase_rendered": False} + + +def stitch(front: Path, back: Path, protected: np.ndarray, destination: Path) -> dict[str, Any]: + document, binary, view, offset, length, front_pixels = atlas(front) + _, _, _, _, _, back_pixels = atlas(back) + if front_pixels.shape != back_pixels.shape: + raise ValueError("front and back atlas shapes differ") + merged = back_pixels.copy() + merged.reshape(-1, merged.shape[-1])[protected, :3] = front_pixels.reshape(-1, front_pixels.shape[-1])[protected, :3] + if merged.shape[-1] == 4: + merged[..., 3] = front_pixels[..., 3] + payload = io.BytesIO() + Image.fromarray(merged).save(payload, format="PNG") + output_document = copy.deepcopy(document) + output_binary = _glb._replace_image_payload(output_document, binary, view, offset, length, payload.getvalue()) + _glb._write_glb(destination, output_document, output_binary) + return validate_stitch(front, back, destination, protected) diff --git a/supratex/models/efficient_attention.py b/supratex/models/efficient_attention.py index 3fa19feabaf26be12a507edbf4016d8ee494c995..b735699e883db85230453b227d28c82e2672afb2 100644 --- a/supratex/models/efficient_attention.py +++ b/supratex/models/efficient_attention.py @@ -1,103 +1,99 @@ -"""Opt-in modulation and empty-cache fixes for the frozen observation attention-O model. - -Both policies are deliberately plain attributes rather than module state, so a -observation attention-O checkpoint loads strictly without an ABI migration. Their defaults -keep the inherited execution path literal. -""" -from __future__ import annotations - -from typing import Any, List, Optional, Tuple - -import torch - -from ..modules import sparse as sp -from .observation_attention import _ObservationCache, _PSABlock -from .observation_stream import ObservationStreamModel - - -class EfficientObservationAttentionModel(ObservationStreamModel): - """observation attention-O with independently opt-in execution-only policies. - - ``direct_modulation_broadcast`` avoids the per-token modulation gather. - At B=1 it retains the singleton modulation row for PyTorch broadcasting; - at B>1 it expands each sample's own singleton row and concatenates those - pieces, preserving ragged sample boundaries. The changed BF16 reduction - path is mathematically equivalent but is not a bit-exact gradient claim. - - ``safe_empty_cache`` repairs only the cached all-empty-O representation. - It is default-off because the inherited literal policy is frozen; enable it - for an eval cache that may be populated by an all-dropped observation call. - """ - def __init__(self, direct_modulation_broadcast: bool = False, - safe_empty_cache: bool = False, **kwargs: Any): - if not isinstance(direct_modulation_broadcast, bool): - raise ValueError('direct_modulation_broadcast must be a bool') - if not isinstance(safe_empty_cache, bool): - raise ValueError('safe_empty_cache must be a bool') - super().__init__(**kwargs) - self.direct_modulation_broadcast = direct_modulation_broadcast - self.safe_empty_cache = safe_empty_cache - - def _cache_signature(self, x: sp.SparseTensor, shape: sp.SparseTensor, - observed: sp.SparseTensor, keep: Optional[torch.Tensor] - ) -> Tuple[Any, ...]: - return super()._cache_signature(x, shape, observed, keep) + ( - ('fast_policy', self.direct_modulation_broadcast, self.safe_empty_cache), - ) - - def _attention(self, block: _PSABlock, o: torch.Tensor, o_coords: torch.Tensor, - o_layout: List[slice], x: torch.Tensor, x_tensor: sp.SparseTensor, - cached_kv: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, **kwargs: Any): - if self.safe_empty_cache and cached_kv is not None and cached_kv[0].ndim == 2: - head_dim = self.model_channels // self.num_heads - cached_kv = tuple(item.reshape(0, self.num_heads, head_dim) for item in cached_kv) - return super()._attention(block, o, o_coords, o_layout, x, x_tensor, cached_kv, **kwargs) - - @staticmethod - def _direct_modulation(value: torch.Tensor, layout: List[slice]) -> torch.Tensor: - """Return token-aligned modulation without an indexed batch gather.""" - if len(layout) == 1: - return value - parts = [value[index:index + 1].expand(item.stop - item.start, -1) - for index, item in enumerate(layout)] - return torch.cat(parts, dim=0) - - def _block(self, block: _PSABlock, o: torch.Tensor, o_coords: torch.Tensor, - o_layout: List[slice], x: torch.Tensor, x_tensor: sp.SparseTensor, - mod_x: torch.Tensor, mod_o: torch.Tensor, terminal: bool, - cache: Optional[_ObservationCache] = None, index: int = 0 - ) -> Tuple[torch.Tensor, torch.Tensor]: - if not self.direct_modulation_broadcast: - return super()._block(block, o, o_coords, o_layout, x, x_tensor, mod_x, mod_o, - terminal, cache, index) - - sx, scx, gx, smx, scm, gmx = block._modulation(mod_x) - so, sco, go, smo, scmo, gmo = block._modulation(mod_o) - x_layout = x_tensor.layout - xo = block.norm1(x) * (1 + self._direct_modulation(scx, x_layout)) + self._direct_modulation(sx, x_layout) - oo = (block.norm1(o) * (1 + self._direct_modulation(sco, o_layout)) + self._direct_modulation(so, o_layout) - if len(o) else o) - cached_kv = None if cache is None else cache.kv[index] - refresh_o = bool(len(o)) and not terminal and index % self.o_refresh_stride == 0 and cached_kv is None - compute_o_attention = refresh_o or ( - terminal and not self.terminal_o_kv_only and bool(len(o)) and cached_kv is None - ) - ao, ax, built_kv = self._attention( - block, oo, o_coords, o_layout, xo, x_tensor, - compute_o_attention=compute_o_attention, cached_kv=cached_kv, - ) - if cache is not None and built_kv is not None: - cache.record(index, *built_kv) - x = x + ax * self._direct_modulation(gx, x_layout) - if refresh_o: - assert ao is not None - o = o + ao * self._direct_modulation(go, o_layout) - o_mlp = sp.VarLenTensor( - block.norm2(o) * (1 + self._direct_modulation(scmo, o_layout)) + self._direct_modulation(smo, o_layout), - o_layout, - ) - o = o + block.mlp(o_mlp).feats * self._direct_modulation(gmo, o_layout) - x = x + block.mlp.mlp(x_tensor.replace( - block.norm2(x) * (1 + self._direct_modulation(scm, x_layout)) + self._direct_modulation(smx, x_layout), - )).feats * self._direct_modulation(gmx, x_layout) - return o, x +"""Opt-in modulation and empty-cache fixes for the frozen observation attention-O model. + +Both policies are deliberately plain attributes rather than module state, so a +observation attention-O checkpoint loads strictly without an ABI migration. Their defaults +keep the inherited execution path literal. +""" +from __future__ import annotations + +from typing import Any, List, Optional, Tuple + +import torch + +from ..modules import sparse as sp +from .observation_attention import _ObservationCache, _PSABlock +from .observation_stream import ObservationStreamModel + + +class EfficientObservationAttentionModel(ObservationStreamModel): + """observation attention-O with independently opt-in execution-only policies. + + ``direct_modulation_broadcast`` avoids the per-token modulation gather. + At B=1 it retains the singleton modulation row for PyTorch broadcasting; + at B>1 it expands each sample's own singleton row and concatenates those + pieces, preserving ragged sample boundaries. The changed BF16 reduction + path is mathematically equivalent but is not a bit-exact gradient claim. + + ``safe_empty_cache`` repairs only the cached all-empty-O representation. + It is default-off because the inherited literal policy is frozen; enable it + for an eval cache that may be populated by an all-dropped observation call. + """ + def __init__(self, direct_modulation_broadcast: bool = False, + safe_empty_cache: bool = False, **kwargs: Any): + super().__init__(**kwargs) + self.direct_modulation_broadcast = direct_modulation_broadcast + self.safe_empty_cache = safe_empty_cache + + def _cache_signature(self, x: sp.SparseTensor, shape: sp.SparseTensor, + observed: sp.SparseTensor, keep: Optional[torch.Tensor] + ) -> Tuple[Any, ...]: + return super()._cache_signature(x, shape, observed, keep) + ( + ('fast_policy', self.direct_modulation_broadcast, self.safe_empty_cache), + ) + + def _attention(self, block: _PSABlock, o: torch.Tensor, o_coords: torch.Tensor, + o_layout: List[slice], x: torch.Tensor, x_tensor: sp.SparseTensor, + cached_kv: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, **kwargs: Any): + if self.safe_empty_cache and cached_kv is not None and cached_kv[0].ndim == 2: + head_dim = self.model_channels // self.num_heads + cached_kv = tuple(item.reshape(0, self.num_heads, head_dim) for item in cached_kv) + return super()._attention(block, o, o_coords, o_layout, x, x_tensor, cached_kv, **kwargs) + + @staticmethod + def _direct_modulation(value: torch.Tensor, layout: List[slice]) -> torch.Tensor: + """Return token-aligned modulation without an indexed batch gather.""" + if len(layout) == 1: + return value + parts = [value[index:index + 1].expand(item.stop - item.start, -1) + for index, item in enumerate(layout)] + return torch.cat(parts, dim=0) + + def _block(self, block: _PSABlock, o: torch.Tensor, o_coords: torch.Tensor, + o_layout: List[slice], x: torch.Tensor, x_tensor: sp.SparseTensor, + mod_x: torch.Tensor, mod_o: torch.Tensor, terminal: bool, + cache: Optional[_ObservationCache] = None, index: int = 0 + ) -> Tuple[torch.Tensor, torch.Tensor]: + if not self.direct_modulation_broadcast: + return super()._block(block, o, o_coords, o_layout, x, x_tensor, mod_x, mod_o, + terminal, cache, index) + + sx, scx, gx, smx, scm, gmx = block._modulation(mod_x) + so, sco, go, smo, scmo, gmo = block._modulation(mod_o) + x_layout = x_tensor.layout + xo = block.norm1(x) * (1 + self._direct_modulation(scx, x_layout)) + self._direct_modulation(sx, x_layout) + oo = (block.norm1(o) * (1 + self._direct_modulation(sco, o_layout)) + self._direct_modulation(so, o_layout) + if len(o) else o) + cached_kv = None if cache is None else cache.kv[index] + refresh_o = bool(len(o)) and not terminal and index % self.o_refresh_stride == 0 and cached_kv is None + compute_o_attention = refresh_o or ( + terminal and not self.terminal_o_kv_only and bool(len(o)) and cached_kv is None + ) + ao, ax, built_kv = self._attention( + block, oo, o_coords, o_layout, xo, x_tensor, + compute_o_attention=compute_o_attention, cached_kv=cached_kv, + ) + if cache is not None and built_kv is not None: + cache.record(index, *built_kv) + x = x + ax * self._direct_modulation(gx, x_layout) + if refresh_o: + assert ao is not None + o = o + ao * self._direct_modulation(go, o_layout) + o_mlp = sp.VarLenTensor( + block.norm2(o) * (1 + self._direct_modulation(scmo, o_layout)) + self._direct_modulation(smo, o_layout), + o_layout, + ) + o = o + block.mlp(o_mlp).feats * self._direct_modulation(gmo, o_layout) + x = x + block.mlp.mlp(x_tensor.replace( + block.norm2(x) * (1 + self._direct_modulation(scm, x_layout)) + self._direct_modulation(smx, x_layout), + )).feats * self._direct_modulation(gmx, x_layout) + return o, x diff --git a/supratex/models/interleaved_backbone.py b/supratex/models/interleaved_backbone.py index 04e4aa182856e7cba7ff104df223a8e4361af10f..1bc91d0eaf543b083fa98a06b8b93d79c95b537b 100644 --- a/supratex/models/interleaved_backbone.py +++ b/supratex/models/interleaved_backbone.py @@ -1,245 +1,243 @@ -"""Packed joint O/X attention for the fixed-stride observation attention-ILV policy. - -This is a distinct architecture with the same observation attention parameter layout and V11 -loader interface but a different block mapping, not an equivalent optimization. Observation -state is static at step zero, then depends on current X after its first joint -refresh at a nonterminal even block, so cross-denoising O K/V caches are -deliberately unsupported. -""" -from __future__ import annotations - -import math -import re -from typing import Any, Dict, List, Optional, Tuple - -import torch -from torch import nn - -from ..modules import sparse as sp -from ..modules.sparse.attention.full_attn import sparse_scaled_dot_product_attention -from .observation_attention import _ObservationCache, _PSABlock -from .efficient_attention import EfficientObservationAttentionModel - - -class InterleavedBackbone(EfficientObservationAttentionModel): - """observation attention Fast with fixed stride-two packed joint O/X refreshes. - - Even nonterminal blocks run one varlen attention over each sample's - ``[O; X]`` sequence and split its output back into O and X rows. Odd - blocks retain the directed X query over that same sample-local K/V set, - without updating O or its MLP. This class never materializes or consumes - an observation cache. - """ - def __init__(self, direct_modulation_broadcast: bool = False, - safe_empty_cache: bool = False, terminal_o_kv_only: bool = True, - o_refresh_stride: int = 2, compact_observation_cache: bool = False, - hot_observation_cache: bool = False, use_observation_cache: bool = False, - v11_init_policy: str = 'odd_copy_even_identity', - **kwargs: Any): - if v11_init_policy != 'odd_copy_even_identity': - raise ValueError("interleaved backbone fixes v11_init_policy='odd_copy_even_identity'") - if not isinstance(terminal_o_kv_only, bool) or not terminal_o_kv_only: - raise ValueError('interleaved backbone fixes terminal_o_kv_only=True') - if isinstance(o_refresh_stride, bool) or not isinstance(o_refresh_stride, int) or o_refresh_stride != 2: - raise ValueError('interleaved backbone fixes o_refresh_stride=2') - for name, value in ( - ('compact_observation_cache', compact_observation_cache), - ('hot_observation_cache', hot_observation_cache), - ('use_observation_cache', use_observation_cache), - ): - if not isinstance(value, bool): - raise ValueError(f'{name} must be a bool') - if value: - raise ValueError(f'interleaved backbone does not support {name}=True') - # This is deliberately not module state: observation attention checkpoints retain their - # state_dict ABI, while a load has one frozen, reviewable expansion rule. - self.v11_init_policy = v11_init_policy - super().__init__( - direct_modulation_broadcast=direct_modulation_broadcast, - safe_empty_cache=safe_empty_cache, - terminal_o_kv_only=True, - o_refresh_stride=2, - **kwargs, - ) - - def initialize_weights(self) -> None: - """Initialize active odd blocks and identity-inserted even blocks. - - The observation attention base initializes its former tail as identity blocks. ILV - instead places the active V11 copies at odd positions, so restore the - ordinary residual-output initialization for odd tail blocks and zero - exactly the two residual outputs of every even inserted block. - """ - super().initialize_weights() - output_std = math.sqrt(2.0 / (5.0 * self.model_channels)) - for index, block in enumerate(self.blocks): - if index % 2: - nn.init.normal_(block.self_attn.to_out.weight, std=output_std) - nn.init.zeros_(block.self_attn.to_out.bias) - nn.init.normal_(block.mlp.mlp[2].weight, std=output_std) - nn.init.zeros_(block.mlp.mlp[2].bias) - else: - nn.init.zeros_(block.self_attn.to_out.weight) - nn.init.zeros_(block.self_attn.to_out.bias) - nn.init.zeros_(block.mlp.mlp[2].weight) - nn.init.zeros_(block.mlp.mlp[2].bias) - - def migrate_v11_state_dict(self, state_dict: Dict[str, torch.Tensor]) -> Dict[str, List[str]]: - """Strictly expand V11's 30 blocks into ILV identity/copy pairs. - - Parent validation remains the sole source classifier: it rejects every - unknown tensor, shape mismatch, and V11 block beyond the frozen source - depth before it writes. Once that validation succeeds, this method - remaps directly from the untouched source dictionary, preventing an - in-place 0 -> 1, 1 -> 3 overlap from corrupting a later source block. - """ - if self.num_blocks != 60: - raise RuntimeError('interleaved backbone V11 migration requires the formal 60-block target') - inventory = super().migrate_v11_state_dict(state_dict) - target = self.state_dict() - remapped: Dict[str, torch.Tensor] = {} - seeded: List[str] = [] - zeroed: List[str] = [] - - for source_key in inventory['copied']: - match = re.match(r'^blocks\.(\d+)\.(.*)$', source_key) - if match is None: - continue - source_index, suffix = int(match.group(1)), match.group(2) - if source_index >= 30: - # The parent classifier has already made this impossible; keep - # this local invariant explicit beside the ILV index arithmetic. - raise RuntimeError(f'unclassified V11 tensor outside frozen 30-block source: {source_key}') - suffix = suffix.replace('norm3.', 'norm2.') - for destination_index in (2 * source_index, 2 * source_index + 1): - destination = f'blocks.{destination_index}.{suffix}' - if destination not in target: - raise RuntimeError(f'ILV destination is missing: {destination}') - remapped[destination] = state_dict[source_key] - - for index in range(0, self.num_blocks, 2): - for suffix in ( - 'self_attn.to_out.weight', 'self_attn.to_out.bias', - 'mlp.mlp.2.weight', 'mlp.mlp.2.bias', - ): - destination = f'blocks.{index}.{suffix}' - remapped[destination] = torch.zeros_like(target[destination]) - zeroed.append(destination) - seeded.extend( - key for key in target if key.startswith(f'blocks.{index}.') and key not in zeroed - ) - - self.load_state_dict(remapped, strict=False) - loaded = set(remapped) - loaded.update(key for key in inventory['copied'] if not key.startswith('blocks.')) - return { - 'copied': inventory['copied'], - 'discarded': inventory['discarded'], - 'seeded': sorted(seeded), - 'zeroed': sorted(zeroed), - 'new': sorted(set(target) - loaded), - } - - def load_v11_state_dict(self, state_dict: Dict[str, torch.Tensor]) -> Dict[str, List[str]]: - return self.migrate_v11_state_dict(state_dict) - - def _cache_signature(self, x: sp.SparseTensor, shape: sp.SparseTensor, - observed: sp.SparseTensor, keep: Optional[torch.Tensor] - ) -> Tuple[Any, ...]: - """Tag the class should a caller inspect identity without creating a cache.""" - return super()._cache_signature(x, shape, observed, keep) + (('psailv_cache_disabled',),) - - def observation_cache(self, x: sp.SparseTensor, concat_cond: sp.SparseTensor, - partial_z_obs: sp.SparseTensor, - partial_z_obs_keep_mask: Optional[torch.Tensor] = None) -> _ObservationCache: - raise RuntimeError('interleaved backbone disables observation_cache; packed joint attention requires uncached O K/V') - - @staticmethod - def _joined_layout(o_layout: List[slice], x_layout: List[slice]) -> List[slice]: - start = 0 - result = [] - for os, xs in zip(o_layout, x_layout): - stop = start + (os.stop - os.start) + (xs.stop - xs.start) - result.append(slice(start, stop)) - start = stop - return result - - def _attention(self, block: _PSABlock, o: torch.Tensor, o_coords: torch.Tensor, - o_layout: List[slice], x: torch.Tensor, x_tensor: sp.SparseTensor, - cached_kv: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, *, - compute_o_attention: Optional[bool] = None, **_: Any - ) -> Tuple[Optional[torch.Tensor], torch.Tensor, - Optional[Tuple[torch.Tensor, torch.Tensor]]]: - """Run the fixed ILV attention policy without a baseline O->O fallback.""" - if cached_kv is not None: - raise RuntimeError('interleaved backbone disables observation_cache; cached K/V consumption is unsupported') - if compute_o_attention is None: - raise RuntimeError('interleaved backbone requires an explicit O-attention policy from its block scheduler') - - qx, kx, vx = self._project_qkv(block, x, x_tensor.coords, x_tensor.layout) - if len(o): - qo, ko, vo = self._project_qkv(block, o, o_coords, o_layout) - else: - qo = ko = vo = None - - joint_layout = self._joined_layout(o_layout, x_tensor.layout) - joint_k = [] - joint_v = [] - joint_q = [] - for os, xs in zip(o_layout, x_tensor.layout): - if len(o): - joint_k.append(torch.cat((ko.feats[os], kx.feats[xs]), dim=0)) - joint_v.append(torch.cat((vo.feats[os], vx.feats[xs]), dim=0)) - if compute_o_attention: - joint_q.append(torch.cat((qo.feats[os], qx.feats[xs]), dim=0)) - else: - joint_k.append(kx.feats[xs]) - joint_v.append(vx.feats[xs]) - if compute_o_attention: - joint_q.append(qx.feats[xs]) - kv = sp.VarLenTensor(torch.stack((torch.cat(joint_k, dim=0), torch.cat(joint_v, dim=0)), dim=1), joint_layout) - - if not compute_o_attention: - x_out = sparse_scaled_dot_product_attention(qx, kv).feats - return None, block.self_attn.to_out(x_out.reshape(len(x_out), self.model_channels)), None - - joint_q_tensor = sp.VarLenTensor(torch.cat(joint_q, dim=0), joint_layout) - joint_out = sparse_scaled_dot_product_attention(joint_q_tensor, kv).feats - o_parts, x_parts = [], [] - start = 0 - for os, xs in zip(o_layout, x_tensor.layout): - o_count, x_count = os.stop - os.start, xs.stop - xs.start - o_parts.append(joint_out[start:start + o_count]) - x_parts.append(joint_out[start + o_count:start + o_count + x_count]) - start += o_count + x_count - o_out = torch.cat(o_parts, dim=0) - x_out = torch.cat(x_parts, dim=0) - return (block.self_attn.to_out(o_out.reshape(len(o_out), self.model_channels)), - block.self_attn.to_out(x_out.reshape(len(x_out), self.model_channels)), None) - - def _block(self, block: _PSABlock, o: torch.Tensor, o_coords: torch.Tensor, - o_layout: List[slice], x: torch.Tensor, x_tensor: sp.SparseTensor, - mod_x: torch.Tensor, mod_o: torch.Tensor, terminal: bool, - cache: Optional[_ObservationCache] = None, index: int = 0 - ) -> Tuple[torch.Tensor, torch.Tensor]: - if cache is not None: - raise RuntimeError('interleaved backbone disables observation_cache; cache consumption is unsupported') - return super()._block(block, o, o_coords, o_layout, x, x_tensor, mod_x, mod_o, - terminal, None, index) - - def forward(self, x: sp.SparseTensor, t: torch.Tensor, cond: Any = None, - concat_cond: Optional[sp.SparseTensor] = None, - partial_z_obs: Optional[sp.SparseTensor] = None, - partial_z_obs_keep_mask: Optional[torch.Tensor] = None, - observation_cache: Optional[_ObservationCache] = None, **kwargs: Any) -> sp.SparseTensor: - if observation_cache is not None: - raise RuntimeError('interleaved backbone disables observation_cache; pass observation_cache=None') - for name in ('use_observation_cache', 'compact_observation_cache', 'hot_observation_cache'): - if name in kwargs: - value = kwargs.pop(name) - if value is not False: - raise RuntimeError(f'interleaved backbone disables {name}; O depends on current X') - if self.o_refresh_stride != 2 or self.terminal_o_kv_only is not True: - raise RuntimeError('interleaved backbone requires its fixed stride2/terminal-KV-only policy') - return super().forward(x, t, cond, concat_cond, partial_z_obs, - partial_z_obs_keep_mask, None, **kwargs) +"""Packed joint O/X attention for the fixed-stride observation attention-ILV policy. + +This is a distinct architecture with the same observation attention parameter layout and V11 +loader interface but a different block mapping, not an equivalent optimization. Observation +state is static at step zero, then depends on current X after its first joint +refresh at a nonterminal even block, so cross-denoising O K/V caches are +deliberately unsupported. +""" +from __future__ import annotations + +import math +import re +from typing import Any, Dict, List, Optional, Tuple + +import torch +from torch import nn + +from ..modules import sparse as sp +from ..modules.sparse.attention.full_attn import sparse_scaled_dot_product_attention +from .observation_attention import _ObservationCache, _PSABlock +from .efficient_attention import EfficientObservationAttentionModel + + +class InterleavedBackbone(EfficientObservationAttentionModel): + """observation attention Fast with fixed stride-two packed joint O/X refreshes. + + Even nonterminal blocks run one varlen attention over each sample's + ``[O; X]`` sequence and split its output back into O and X rows. Odd + blocks retain the directed X query over that same sample-local K/V set, + without updating O or its MLP. This class never materializes or consumes + an observation cache. + """ + def __init__(self, direct_modulation_broadcast: bool = False, + safe_empty_cache: bool = False, terminal_o_kv_only: bool = True, + o_refresh_stride: int = 2, compact_observation_cache: bool = False, + hot_observation_cache: bool = False, use_observation_cache: bool = False, + v11_init_policy: str = 'odd_copy_even_identity', + **kwargs: Any): + if v11_init_policy != 'odd_copy_even_identity': + raise ValueError("interleaved backbone fixes v11_init_policy='odd_copy_even_identity'") + if not terminal_o_kv_only: + raise ValueError('interleaved backbone fixes terminal_o_kv_only=True') + if o_refresh_stride != 2: + raise ValueError('interleaved backbone fixes o_refresh_stride=2') + for name, value in ( + ('compact_observation_cache', compact_observation_cache), + ('hot_observation_cache', hot_observation_cache), + ('use_observation_cache', use_observation_cache), + ): + if value: + raise ValueError(f'interleaved backbone does not support {name}=True') + # This is deliberately not module state: observation attention checkpoints retain their + # state_dict ABI, while a load has one frozen, reviewable expansion rule. + self.v11_init_policy = v11_init_policy + super().__init__( + direct_modulation_broadcast=direct_modulation_broadcast, + safe_empty_cache=safe_empty_cache, + terminal_o_kv_only=True, + o_refresh_stride=2, + **kwargs, + ) + + def initialize_weights(self) -> None: + """Initialize active odd blocks and identity-inserted even blocks. + + The observation attention base initializes its former tail as identity blocks. ILV + instead places the active V11 copies at odd positions, so restore the + ordinary residual-output initialization for odd tail blocks and zero + exactly the two residual outputs of every even inserted block. + """ + super().initialize_weights() + output_std = math.sqrt(2.0 / (5.0 * self.model_channels)) + for index, block in enumerate(self.blocks): + if index % 2: + nn.init.normal_(block.self_attn.to_out.weight, std=output_std) + nn.init.zeros_(block.self_attn.to_out.bias) + nn.init.normal_(block.mlp.mlp[2].weight, std=output_std) + nn.init.zeros_(block.mlp.mlp[2].bias) + else: + nn.init.zeros_(block.self_attn.to_out.weight) + nn.init.zeros_(block.self_attn.to_out.bias) + nn.init.zeros_(block.mlp.mlp[2].weight) + nn.init.zeros_(block.mlp.mlp[2].bias) + + def migrate_v11_state_dict(self, state_dict: Dict[str, torch.Tensor]) -> Dict[str, List[str]]: + """Strictly expand V11's 30 blocks into ILV identity/copy pairs. + + Parent validation remains the sole source classifier: it rejects every + unknown tensor, shape mismatch, and V11 block beyond the frozen source + depth before it writes. Once that validation succeeds, this method + remaps directly from the untouched source dictionary, preventing an + in-place 0 -> 1, 1 -> 3 overlap from corrupting a later source block. + """ + if self.num_blocks != 60: + raise RuntimeError('interleaved backbone V11 migration requires the formal 60-block target') + inventory = super().migrate_v11_state_dict(state_dict) + target = self.state_dict() + remapped: Dict[str, torch.Tensor] = {} + seeded: List[str] = [] + zeroed: List[str] = [] + + for source_key in inventory['copied']: + match = re.match(r'^blocks\.(\d+)\.(.*)$', source_key) + if match is None: + continue + source_index, suffix = int(match.group(1)), match.group(2) + if source_index >= 30: + # The parent classifier has already made this impossible; keep + # this local invariant explicit beside the ILV index arithmetic. + raise RuntimeError(f'unclassified V11 tensor outside frozen 30-block source: {source_key}') + suffix = suffix.replace('norm3.', 'norm2.') + for destination_index in (2 * source_index, 2 * source_index + 1): + destination = f'blocks.{destination_index}.{suffix}' + if destination not in target: + raise RuntimeError(f'ILV destination is missing: {destination}') + remapped[destination] = state_dict[source_key] + + for index in range(0, self.num_blocks, 2): + for suffix in ( + 'self_attn.to_out.weight', 'self_attn.to_out.bias', + 'mlp.mlp.2.weight', 'mlp.mlp.2.bias', + ): + destination = f'blocks.{index}.{suffix}' + remapped[destination] = torch.zeros_like(target[destination]) + zeroed.append(destination) + seeded.extend( + key for key in target if key.startswith(f'blocks.{index}.') and key not in zeroed + ) + + self.load_state_dict(remapped, strict=False) + loaded = set(remapped) + loaded.update(key for key in inventory['copied'] if not key.startswith('blocks.')) + return { + 'copied': inventory['copied'], + 'discarded': inventory['discarded'], + 'seeded': sorted(seeded), + 'zeroed': sorted(zeroed), + 'new': sorted(set(target) - loaded), + } + + def load_v11_state_dict(self, state_dict: Dict[str, torch.Tensor]) -> Dict[str, List[str]]: + return self.migrate_v11_state_dict(state_dict) + + def _cache_signature(self, x: sp.SparseTensor, shape: sp.SparseTensor, + observed: sp.SparseTensor, keep: Optional[torch.Tensor] + ) -> Tuple[Any, ...]: + """Tag the class should a caller inspect identity without creating a cache.""" + return super()._cache_signature(x, shape, observed, keep) + (('psailv_cache_disabled',),) + + def observation_cache(self, x: sp.SparseTensor, concat_cond: sp.SparseTensor, + partial_z_obs: sp.SparseTensor, + partial_z_obs_keep_mask: Optional[torch.Tensor] = None) -> _ObservationCache: + raise RuntimeError('interleaved backbone disables observation_cache; packed joint attention requires uncached O K/V') + + @staticmethod + def _joined_layout(o_layout: List[slice], x_layout: List[slice]) -> List[slice]: + start = 0 + result = [] + for os, xs in zip(o_layout, x_layout): + stop = start + (os.stop - os.start) + (xs.stop - xs.start) + result.append(slice(start, stop)) + start = stop + return result + + def _attention(self, block: _PSABlock, o: torch.Tensor, o_coords: torch.Tensor, + o_layout: List[slice], x: torch.Tensor, x_tensor: sp.SparseTensor, + cached_kv: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, *, + compute_o_attention: Optional[bool] = None, **_: Any + ) -> Tuple[Optional[torch.Tensor], torch.Tensor, + Optional[Tuple[torch.Tensor, torch.Tensor]]]: + """Run the fixed ILV attention policy without a baseline O->O fallback.""" + if cached_kv is not None: + raise RuntimeError('interleaved backbone disables observation_cache; cached K/V consumption is unsupported') + if compute_o_attention is None: + raise RuntimeError('interleaved backbone requires an explicit O-attention policy from its block scheduler') + + qx, kx, vx = self._project_qkv(block, x, x_tensor.coords, x_tensor.layout) + if len(o): + qo, ko, vo = self._project_qkv(block, o, o_coords, o_layout) + else: + qo = ko = vo = None + + joint_layout = self._joined_layout(o_layout, x_tensor.layout) + joint_k = [] + joint_v = [] + joint_q = [] + for os, xs in zip(o_layout, x_tensor.layout): + if len(o): + joint_k.append(torch.cat((ko.feats[os], kx.feats[xs]), dim=0)) + joint_v.append(torch.cat((vo.feats[os], vx.feats[xs]), dim=0)) + if compute_o_attention: + joint_q.append(torch.cat((qo.feats[os], qx.feats[xs]), dim=0)) + else: + joint_k.append(kx.feats[xs]) + joint_v.append(vx.feats[xs]) + if compute_o_attention: + joint_q.append(qx.feats[xs]) + kv = sp.VarLenTensor(torch.stack((torch.cat(joint_k, dim=0), torch.cat(joint_v, dim=0)), dim=1), joint_layout) + + if not compute_o_attention: + x_out = sparse_scaled_dot_product_attention(qx, kv).feats + return None, block.self_attn.to_out(x_out.reshape(len(x_out), self.model_channels)), None + + joint_q_tensor = sp.VarLenTensor(torch.cat(joint_q, dim=0), joint_layout) + joint_out = sparse_scaled_dot_product_attention(joint_q_tensor, kv).feats + o_parts, x_parts = [], [] + start = 0 + for os, xs in zip(o_layout, x_tensor.layout): + o_count, x_count = os.stop - os.start, xs.stop - xs.start + o_parts.append(joint_out[start:start + o_count]) + x_parts.append(joint_out[start + o_count:start + o_count + x_count]) + start += o_count + x_count + o_out = torch.cat(o_parts, dim=0) + x_out = torch.cat(x_parts, dim=0) + return (block.self_attn.to_out(o_out.reshape(len(o_out), self.model_channels)), + block.self_attn.to_out(x_out.reshape(len(x_out), self.model_channels)), None) + + def _block(self, block: _PSABlock, o: torch.Tensor, o_coords: torch.Tensor, + o_layout: List[slice], x: torch.Tensor, x_tensor: sp.SparseTensor, + mod_x: torch.Tensor, mod_o: torch.Tensor, terminal: bool, + cache: Optional[_ObservationCache] = None, index: int = 0 + ) -> Tuple[torch.Tensor, torch.Tensor]: + if cache is not None: + raise RuntimeError('interleaved backbone disables observation_cache; cache consumption is unsupported') + return super()._block(block, o, o_coords, o_layout, x, x_tensor, mod_x, mod_o, + terminal, None, index) + + def forward(self, x: sp.SparseTensor, t: torch.Tensor, cond: Any = None, + concat_cond: Optional[sp.SparseTensor] = None, + partial_z_obs: Optional[sp.SparseTensor] = None, + partial_z_obs_keep_mask: Optional[torch.Tensor] = None, + observation_cache: Optional[_ObservationCache] = None, **kwargs: Any) -> sp.SparseTensor: + if observation_cache is not None: + raise RuntimeError('interleaved backbone disables observation_cache; pass observation_cache=None') + for name in ('use_observation_cache', 'compact_observation_cache', 'hot_observation_cache'): + if name in kwargs: + value = kwargs.pop(name) + if value is not False: + raise RuntimeError(f'interleaved backbone disables {name}; O depends on current X') + if self.o_refresh_stride != 2 or self.terminal_o_kv_only is not True: + raise RuntimeError('interleaved backbone requires its fixed stride2/terminal-KV-only policy') + return super().forward(x, t, cond, concat_cond, partial_z_obs, + partial_z_obs_keep_mask, None, **kwargs) diff --git a/supratex/models/interleaved_flow_transformer.py b/supratex/models/interleaved_flow_transformer.py index 4172e8408b960898c4457ac0d02203eab5f18f74..5efd3d8ae300e9fbde6d19f0c6cf497624bf5a96 100644 --- a/supratex/models/interleaved_flow_transformer.py +++ b/supratex/models/interleaved_flow_transformer.py @@ -1,452 +1,441 @@ -"""Interleaved flow transformer: alternate full-X and U-only updates with static observations. - -The observation attention-ILV weights remain a 60-block backbone. O is always read-only; -the production policy uses the original even blocks for U-only updates and -the original odd blocks for full-X updates. Explicitly configured, -identity-seeded post-block refiners remain available for engineering tests. -No O-update schedule or transition phase exists. -""" -from __future__ import annotations - -import copy -from dataclasses import dataclass -from typing import Any, Dict, Iterable, List, Optional, Sequence, Tuple - -import torch -from torch import nn -import torch.nn.functional as F -import torch.utils.checkpoint - -from ..modules import sparse as sp -from ..modules.sparse.attention.full_attn import sparse_scaled_dot_product_attention -from ..modules.utils import convert_module_to, manual_cast -from .observation_attention import _PSABlock, _join_exact_rows, _layouts_from_batch, _live_zero -from .interleaved_backbone import InterleavedBackbone - - -@dataclass -class _PreparedObservation: - hidden: torch.Tensor - coords: torch.Tensor - layout: List[slice] - urows: torch.Tensor - ucoords: torch.Tensor - ubatch: torch.Tensor - u_layout: List[slice] - - -class _StaticObservationCache: - """Invocation-local, eval-only cache for static O preparation and layer K/V.""" - def __init__(self, model: 'InterleavedFlowTransformer', x: sp.SparseTensor, shape: sp.SparseTensor, - observed: sp.SparseTensor, keep: Optional[torch.Tensor]): - if model.training or torch.is_grad_enabled(): - raise RuntimeError('observation_cache is inference-only; training cannot retain detached O K/V') - tensors = { - 'x.coords': x.coords, 'concat_cond.coords': shape.coords, 'concat_cond.feats': shape.feats, - 'partial_z_obs.coords': observed.coords, 'partial_z_obs.feats': observed.feats, - } - if keep is not None: - tensors['partial_z_obs_keep_mask'] = keep - inference = next((name for name, tensor in tensors.items() if torch.is_inference(tensor)), None) - if inference is not None: - raise RuntimeError( - f'observation_cache cannot safely validate inference-mode tensor {inference}; ' - 'pass regular tensors under torch.no_grad()' - ) - model_tensor = next( - (name for name, tensor in [*model.named_parameters(), *model.named_buffers()] - if torch.is_inference(tensor)), - None, - ) - if model_tensor is not None: - raise RuntimeError( - f'observation_cache cannot safely validate inference-mode model tensor {model_tensor}; ' - 'construct the model outside torch.inference_mode()' - ) - self.model = model - self.signature = model._cache_signature(x, shape, observed, keep) - self.kv: Dict[str, Tuple[torch.Tensor, torch.Tensor]] = {} - self.prepared: Optional[_PreparedObservation] = None - self.zero_time: Optional[torch.Tensor] = None - self.zero_time_signature: Optional[Tuple[Any, ...]] = None - self.bytes = 0 - self.closed = False - - @property - def preprocessing_bytes(self) -> int: - """Resident preparation tensors, separate from the existing K/V byte count.""" - tensors = [] if self.prepared is None else [ - self.prepared.hidden, self.prepared.coords, self.prepared.urows, - self.prepared.ucoords, self.prepared.ubatch, - ] - if self.zero_time is not None: - tensors.append(self.zero_time) - return sum(tensor.numel() * tensor.element_size() for tensor in tensors) - - def validate(self, x: sp.SparseTensor, shape: sp.SparseTensor, observed: sp.SparseTensor, - keep: Optional[torch.Tensor]) -> None: - if self.closed or self.model.training or torch.is_grad_enabled(): - raise RuntimeError('observation_cache is closed or used outside eval/no_grad') - if self.signature != self.model._cache_signature(x, shape, observed, keep): - raise RuntimeError('observation_cache inputs, model weights, dtype, device, or policy changed; rebuild the context') - - def get(self, key: str) -> Optional[Tuple[torch.Tensor, torch.Tensor]]: - return self.kv.get(key) - - def record(self, key: str, k: torch.Tensor, v: torch.Tensor) -> None: - if key not in self.kv: - self.kv[key] = (k.detach(), v.detach()) - self.bytes += (k.numel() * k.element_size()) + (v.numel() * v.element_size()) - - def close(self) -> None: - self.kv.clear() - self.prepared = None - self.zero_time = None - self.zero_time_signature = None - self.closed = True - - def __enter__(self) -> '_StaticObservationCache': - return self - - def __exit__(self, *_: Any) -> None: - self.close() - - -class InterleavedFlowTransformer(InterleavedBackbone): - """Inherited backbone with permanently static O and optional U refiners. - - O receives no Q projection, attention output, or MLP execution. It still - receives its layer-specific normalization/modulation and K/V projection so - every X layer attends over that layer's ``[O; X]`` keys and values. - """ - def __init__(self, refinement_indices: Sequence[int] = (), - base_block_policy: str = 'alternating_u_x', - allow_test_config: bool = False, - direct_modulation_broadcast: bool = True, **kwargs: Any): - retired = {'o_update_scale', 'o_transition_steps'} & kwargs.keys() - if retired: - raise ValueError(f'Interleaved flow transformer O is always static; removed options: {sorted(retired)}') - if not isinstance(base_block_policy, str) or base_block_policy not in {'alternating_u_x', 'all_x'}: - raise ValueError("base_block_policy must be 'alternating_u_x' or 'all_x'") - if not isinstance(allow_test_config, bool): - raise ValueError('allow_test_config must be a bool') - if not isinstance(direct_modulation_broadcast, bool): - raise ValueError('direct_modulation_broadcast must be a bool') - super().__init__(direct_modulation_broadcast=direct_modulation_broadcast, **kwargs) - if self.num_blocks % 2: - raise ValueError('Interleaved flow transformer requires an even number of base blocks') - if self.num_blocks != 60 and not allow_test_config: - raise ValueError('non-60-block Interleaved flow transformer test configurations require allow_test_config=True') - self.base_block_policy = base_block_policy - indices = self._validate_refinement_indices(refinement_indices) - self.refinement_indices = tuple(indices) - self.refiners = nn.ModuleDict({str(index): self._new_refiner(index) for index in indices}) - - def _validate_refinement_indices(self, indices: Sequence[int]) -> Tuple[int, ...]: - if isinstance(indices, (str, bytes)): - raise ValueError('refinement_indices must be an iterable of base-block indices') - try: - result = tuple(indices) - except TypeError as error: - raise ValueError('refinement_indices must be an iterable of base-block indices') from error - if any(isinstance(index, bool) or not isinstance(index, int) for index in result): - raise ValueError('refinement_indices must contain integer base-block indices') - if tuple(sorted(set(result))) != result or any(index < 0 or index >= self.num_blocks for index in result): - raise ValueError('refinement_indices must be unique, sorted, valid zero-based base-block indices') - return result - - def _new_refiner(self, index: int) -> _PSABlock: - # The attention module does not retain a public rope-frequency field; - # a structural clone is therefore the exact construction contract. - block = copy.deepcopy(self.blocks[index]) - with torch.no_grad(): - nn.init.zeros_(block.self_attn.to_out.weight) - nn.init.zeros_(block.self_attn.to_out.bias) - nn.init.zeros_(block.mlp.mlp[2].weight) - nn.init.zeros_(block.mlp.mlp[2].bias) - return block - - def convert_to(self, dtype: torch.dtype) -> None: - """Keep post-parent refiners on the same explicit inference dtype.""" - super().convert_to(dtype) - if hasattr(self, 'refiners'): - self.refiners.apply(lambda module: convert_module_to(module, dtype=dtype)) - - def _cache_signature(self, x: sp.SparseTensor, shape: sp.SparseTensor, observed: sp.SparseTensor, - keep: Optional[torch.Tensor]) -> Tuple[Any, ...]: - """Extend the parent identity signature with static-O cache policy.""" - return super()._cache_signature(x, shape, observed, keep) + ( - ('interleaved_flow_transformer_static_o_kv', self.base_block_policy, self.direct_modulation_broadcast, self.dtype, - torch.is_autocast_enabled('cuda'), torch.get_autocast_dtype('cuda'), - torch.is_autocast_enabled('cpu'), torch.get_autocast_dtype('cpu')), - ) - - def observation_cache(self, x: sp.SparseTensor, concat_cond: sp.SparseTensor, - partial_z_obs: sp.SparseTensor, - partial_z_obs_keep_mask: Optional[torch.Tensor] = None) -> _StaticObservationCache: - """Create an invocation-local cache for static O preparation and K/V.""" - return _StaticObservationCache(self, x, concat_cond, partial_z_obs, partial_z_obs_keep_mask) - - def _prepare_observation(self, x: sp.SparseTensor, shape: sp.SparseTensor, - observed: sp.SparseTensor, keep: Optional[torch.Tensor]) -> _PreparedObservation: - """Prepare the time-independent O stream and the complementary U support.""" - if not torch.equal(shape.coords, x.coords): - raise AssertionError('Interleaved flow transformer requires coordinate-aligned PBR32 x and shape concat_cond') - batch_size = x.shape[0] - keep = torch.ones(batch_size, device=x.device, dtype=torch.bool) if keep is None else keep.to(x.device, torch.bool) - if keep.shape != (batch_size,): - raise AssertionError(f'partial_z_obs_keep_mask must be [B={batch_size}]') - active = keep[observed.coords[:, 0].long()] - ocoords, ofeats = observed.coords[active], observed.feats[active] - rows = _join_exact_rows(x.coords, ocoords) - o_layout = _layouts_from_batch(ocoords[:, 0], batch_size) - raw = ofeats[:, :32].float() * self.partial_std + self.partial_mean - opbr = (raw - self.pbr_mean) / self.pbr_std - oinput = torch.cat((opbr.to(ofeats.dtype), shape.feats[rows]), -1) - # Preserve the parent addition/cast ordering under autocast. - o = self.input_layer(sp.VarLenTensor(oinput, o_layout)).feats - o = o + self.observation_metadata(ofeats[:, 32:34].to(oinput.dtype)) + self.observation_role.to(oinput.dtype) - o = manual_cast(o, self.dtype) - observed_rows = torch.zeros(len(x.feats), device=x.device, dtype=torch.bool) - observed_rows[rows] = True - urows = torch.nonzero(~observed_rows, as_tuple=False).flatten() - ucoords = x.coords[urows] - ubatch = ucoords[:, 0].long() - u_layout = _layouts_from_batch(ubatch, batch_size) - return _PreparedObservation(o, ocoords, o_layout, urows, ucoords, ubatch, u_layout) - - def migrate_v11_state_dict(self, state_dict: Dict[str, torch.Tensor]) -> Dict[str, List[str]]: - raise RuntimeError('Interleaved flow transformer requires a interleaved backbone parent state; V11 migration is not a valid parent path') - - def load_v11_state_dict(self, state_dict: Dict[str, torch.Tensor]) -> Dict[str, List[str]]: - raise RuntimeError('Interleaved flow transformer requires a interleaved backbone parent state; V11 migration is not a valid parent path') - - def load_interleaved_backbone_state_dict(self, state_dict: Dict[str, torch.Tensor]) -> Dict[str, List[str]]: - """Strictly seed Interleaved flow transformer from an exact interleaved backbone state before writing. - - Base keys must be complete and shape-exact. Refiners inherit the - matching base block, then their two residual outputs are zeroed, making - each newly inserted operation exactly identity at migration time. - """ - target = self.state_dict() - base = {key for key in target if not key.startswith('refiners.')} - source_keys = set(state_dict) - missing = sorted(base - source_keys) - unexpected = sorted(source_keys - base) - non_tensors = sorted(key for key in base & source_keys if not isinstance(state_dict[key], torch.Tensor)) - mismatched = sorted(key for key in base & source_keys if isinstance(state_dict[key], torch.Tensor) - and (target[key].shape != state_dict[key].shape or target[key].dtype != state_dict[key].dtype)) - if missing or unexpected or non_tensors or mismatched: - raise RuntimeError('invalid interleaved backbone parent state before write: ' - f'missing={missing}, unexpected={unexpected}, non_tensors={non_tensors}, ' - f'shape_or_dtype_mismatch={mismatched}') - updates = {key: state_dict[key] for key in base} - seeded: List[str] = [] - zeroed: List[str] = [] - for index, refiner in self.refiners.items(): - for suffix, value in refiner.state_dict().items(): - destination = f'refiners.{index}.{suffix}' - source = f'blocks.{index}.{suffix}' - if source not in state_dict or target[destination].shape != state_dict[source].shape: - raise RuntimeError(f'invalid refiner seed source before write: {source}') - updates[destination] = state_dict[source] - seeded.append(destination) - for suffix in ('self_attn.to_out.weight', 'self_attn.to_out.bias', - 'mlp.mlp.2.weight', 'mlp.mlp.2.bias'): - destination = f'refiners.{index}.{suffix}' - updates[destination] = torch.zeros_like(target[destination]) - zeroed.append(destination) - self.load_state_dict(updates, strict=True) - return {'inherited': sorted(base), 'new': sorted(set(target) - base), - 'seeded': sorted(seeded), 'zeroed': sorted(zeroed)} - - def _mod(self, value: torch.Tensor, layout: List[slice], batch: torch.Tensor) -> torch.Tensor: - return self._direct_modulation(value, layout) if self.direct_modulation_broadcast else value[batch] - - @staticmethod - def _obatch(o: torch.Tensor, layout: List[slice]) -> torch.Tensor: - if not len(o): - return torch.empty(0, device=o.device, dtype=torch.long) - return torch.cat([torch.full((item.stop - item.start,), index, device=o.device, dtype=torch.long) - for index, item in enumerate(layout)]) - - def _qkv(self, block: _PSABlock, hidden: torch.Tensor, coords: torch.Tensor, - layout: List[slice], which: Iterable[int]) -> Dict[int, sp.VarLenTensor]: - """Project exactly requested Q/K/V slices; static O never materializes Q.""" - heads, dim = self.num_heads, self.model_channels // self.num_heads - selected = tuple(which) - if selected == (0, 1, 2): - q, k, v = self._project_qkv(block, hidden, coords, layout) - return {0: q, 1: k, 2: v} - weight = block.self_attn.to_qkv.weight.reshape(3, self.model_channels, self.model_channels) - bias = block.self_attn.to_qkv.bias.reshape(3, self.model_channels) - result: Dict[int, sp.VarLenTensor] = {} - if selected == (1, 2): - # One sliced K/V projection is the actual static-O saving; two - # independent linears would preserve semantics but lose the budget. - packed = F.linear(hidden, weight[1:].reshape(2 * self.model_channels, self.model_channels), - bias[1:].reshape(2 * self.model_channels)).reshape(-1, 2, heads, dim) - token = sp.VarLenTensor(packed, layout) - result[1], result[2] = token.unbind(1) - for item in selected: - if item in result: - continue - feats = F.linear(hidden, weight[item], bias[item]).reshape(-1, heads, dim) - result[item] = sp.VarLenTensor(feats, layout) - if 0 in result and block.self_attn.qk_rms_norm: - result[0] = block.self_attn.q_rms_norm(result[0]) - if 1 in result and block.self_attn.qk_rms_norm: - result[1] = block.self_attn.k_rms_norm(result[1]) - if len(coords) and (0 in result or 1 in result): - q = result.get(0); k = result.get(1) - if q is None: - dummy = sp.VarLenTensor(torch.zeros_like(result[1].feats), layout) - _, k = block.self_attn.rope(sp.SparseTensor(dummy.feats, coords), sp.SparseTensor(k.feats, coords)) - result[1] = sp.VarLenTensor(k.feats, layout) - elif k is None: - q, _ = block.self_attn.rope(sp.SparseTensor(q.feats, coords), sp.SparseTensor(torch.zeros_like(q.feats), coords)) - result[0] = sp.VarLenTensor(q.feats, layout) - else: - q, k = block.self_attn.rope(sp.SparseTensor(q.feats, coords), sp.SparseTensor(k.feats, coords)) - result[0], result[1] = sp.VarLenTensor(q.feats, layout), sp.VarLenTensor(k.feats, layout) - return result - - def _x_attention(self, block: _PSABlock, o: torch.Tensor, ocoords: torch.Tensor, o_layout: List[slice], - x: torch.Tensor, x_tensor: sp.SparseTensor, - cache: Optional[_StaticObservationCache] = None, cache_key: Optional[str] = None) -> torch.Tensor: - xqkv = self._qkv(block, x, x_tensor.coords, x_tensor.layout, (0, 1, 2)) - cached_kv = None if cache is None or cache_key is None else cache.get(cache_key) - if cached_kv is not None: - oqkv = {1: sp.VarLenTensor(cached_kv[0], o_layout), 2: sp.VarLenTensor(cached_kv[1], o_layout)} - else: - oqkv = self._qkv(block, o, ocoords, o_layout, (1, 2)) if len(o) else {} - if cache is not None and cache_key is not None and len(o): - cache.record(cache_key, oqkv[1].feats, oqkv[2].feats) - keys: List[torch.Tensor] = []; values: List[torch.Tensor] = []; layout: List[slice] = []; start = 0 - for os, xs in zip(o_layout, x_tensor.layout): - key = torch.cat(([oqkv[1].feats[os], xqkv[1].feats[xs]]) if len(o) else [xqkv[1].feats[xs]], 0) - value = torch.cat(([oqkv[2].feats[os], xqkv[2].feats[xs]]) if len(o) else [xqkv[2].feats[xs]], 0) - keys.append(key); values.append(value); layout.append(slice(start, start + len(key))); start += len(key) - kv = sp.VarLenTensor(torch.stack((torch.cat(keys), torch.cat(values)), dim=1), layout) - xout = sparse_scaled_dot_product_attention(xqkv[0], kv).feats - return block.self_attn.to_out(xout.reshape(len(x_tensor.feats), self.model_channels)) - - def _base_block(self, block: _PSABlock, o: torch.Tensor, ocoords: torch.Tensor, o_layout: List[slice], - h: torch.Tensor, x: sp.SparseTensor, tx: torch.Tensor, to: torch.Tensor, - cache: Optional[_StaticObservationCache] = None, cache_key: Optional[str] = None - ) -> Tuple[torch.Tensor, torch.Tensor]: - sx, scx, gx, smx, scm, gmx = block._modulation(tx) - xo = block.norm1(h) * (1 + self._mod(scx, x.layout, x.batch_boardcast_map)) + self._mod(sx, x.layout, x.batch_boardcast_map) - cached_kv = None if cache is None or cache_key is None else cache.get(cache_key) - if cached_kv is None: - so, sco, _go, _smo, _scmo, _gmo = block._modulation(to) - obatch = self._obatch(o, o_layout) if not self.direct_modulation_broadcast else torch.empty(0, device=o.device, dtype=torch.long) - oo = block.norm1(o) * (1 + self._mod(sco, o_layout, obatch)) + self._mod(so, o_layout, obatch) if len(o) else o - else: - oo = o - ax = self._x_attention(block, oo, ocoords, o_layout, xo, x, cache, cache_key) - h = h + ax * self._mod(gx, x.layout, x.batch_boardcast_map) - h = h + block.mlp.mlp(x.replace(block.norm2(h) * (1 + self._mod(scm, x.layout, x.batch_boardcast_map)) + self._mod(smx, x.layout, x.batch_boardcast_map))).feats * self._mod(gmx, x.layout, x.batch_boardcast_map) - return o, h - - def _refine(self, block: _PSABlock, o: torch.Tensor, ocoords: torch.Tensor, o_layout: List[slice], h: torch.Tensor, - x: sp.SparseTensor, tx: torch.Tensor, to: torch.Tensor, urows: torch.Tensor, u_layout: List[slice], - cache: Optional[_StaticObservationCache] = None, cache_key: Optional[str] = None) -> torch.Tensor: - if not len(urows): - return h + _live_zero(block.parameters(), h) - sx, scx, gx, smx, scm, gmx = block._modulation(tx) - xo = block.norm1(h) * (1 + self._mod(scx, x.layout, x.batch_boardcast_map)) + self._mod(sx, x.layout, x.batch_boardcast_map) - cached_kv = None if cache is None or cache_key is None else cache.get(cache_key) - if cached_kv is None: - so, sco, _go, _smo, _scmo, _gmo = block._modulation(to) - obatch = self._obatch(o, o_layout) if not self.direct_modulation_broadcast else torch.empty(0, device=o.device, dtype=torch.long) - oo = block.norm1(o) * (1 + self._mod(sco, o_layout, obatch)) + self._mod(so, o_layout, obatch) if len(o) else o - else: - oo = o - u = h[urows] - if cache is not None and cache.prepared is not None: - ucoords, ubatch = cache.prepared.ucoords, cache.prepared.ubatch - else: - ucoords, ubatch = x.coords[urows], x.coords[urows, 0].long() - uo = xo[urows] - # U is Q-only; K/V deliberately include O plus every current X row. - uq = self._qkv(block, uo, ucoords, u_layout, (0,))[0] - xkv = self._qkv(block, xo, x.coords, x.layout, (1, 2)) - cached_kv = None if cache is None or cache_key is None else cache.get(cache_key) - if cached_kv is not None: - okv = {1: sp.VarLenTensor(cached_kv[0], o_layout), 2: sp.VarLenTensor(cached_kv[1], o_layout)} - else: - okv = self._qkv(block, oo, ocoords, o_layout, (1, 2)) if len(o) else {} - if cache is not None and cache_key is not None and len(o): - cache.record(cache_key, okv[1].feats, okv[2].feats) - parts: List[torch.Tensor] = []; values: List[torch.Tensor] = []; layout: List[slice] = []; start = 0 - for os, xs in zip(o_layout, x.layout): - key = torch.cat(([okv[1].feats[os], xkv[1].feats[xs]]) if len(o) else [xkv[1].feats[xs]], 0) - value = torch.cat(([okv[2].feats[os], xkv[2].feats[xs]]) if len(o) else [xkv[2].feats[xs]], 0) - parts.append(key); values.append(value); count = key.shape[0]; layout.append(slice(start, start + count)); start += count - kv = sp.VarLenTensor(torch.stack((torch.cat(parts), torch.cat(values)), 1), layout) - u = u + block.self_attn.to_out(sparse_scaled_dot_product_attention(uq, kv).feats.reshape(len(u), self.model_channels)) * self._mod(gx, u_layout, ubatch) - u = u + block.mlp.mlp(sp.SparseTensor(block.norm2(u) * (1 + self._mod(scm, u_layout, ubatch)) + self._mod(smx, u_layout, ubatch), ucoords, shape=x.shape)).feats * self._mod(gmx, u_layout, ubatch) - h = h.clone(); h[urows] = u - return h - - def forward(self, x: sp.SparseTensor, t: torch.Tensor, cond: Any = None, concat_cond: Optional[sp.SparseTensor] = None, - partial_z_obs: Optional[sp.SparseTensor] = None, partial_z_obs_keep_mask: Optional[torch.Tensor] = None, - observation_cache: Optional[_StaticObservationCache] = None, **kwargs: Any) -> sp.SparseTensor: - retired = {'o_update_scale', 'o_transition_steps'} & kwargs.keys() - if retired: - raise ValueError(f'Interleaved flow transformer O is always static; removed options: {sorted(retired)}') - if observation_cache is not None: - if not isinstance(observation_cache, _StaticObservationCache) or observation_cache.model is not self: - raise RuntimeError('observation_cache belongs to a different Interleaved flow transformer model') - observation_cache.validate(x, concat_cond, partial_z_obs, partial_z_obs_keep_mask) - if kwargs: - disabled = next((name for name in ('use_observation_cache', 'compact_observation_cache', 'hot_observation_cache') if name in kwargs and kwargs[name] is not False), None) - if disabled is not None: raise RuntimeError(f'Interleaved flow transformer disables {disabled}') - if concat_cond is None or x.feats.shape[1] != 32 or concat_cond.feats.shape != x.feats.shape: - raise AssertionError('Interleaved flow transformer requires coordinate-aligned PBR32 x and shape concat_cond') - batch_size = x.shape[0] - if t.shape != (batch_size,): raise AssertionError(f't must be [B={batch_size}]') - if partial_z_obs is None or partial_z_obs.feats.ndim != 2 or partial_z_obs.feats.shape[1] != 34: - raise AssertionError('Interleaved flow transformer requires 34-channel partial_z_obs (possibly empty)') - prepared = None if observation_cache is None else observation_cache.prepared - if prepared is None: - prepared = self._prepare_observation(x, concat_cond, partial_z_obs, partial_z_obs_keep_mask) - if observation_cache is not None: - observation_cache.prepared = prepared - o, ocoords, o_layout = prepared.hidden, prepared.coords, prepared.layout - urows, u_layout = prepared.urows, prepared.u_layout - h = self.input_layer(x.replace(torch.cat((x.feats, concat_cond.feats), -1))).feats - h = manual_cast(h, self.dtype) - tx = manual_cast(self.adaLN_modulation(self.t_embedder(t)), h.dtype) - zero_time_signature = (t.dtype, t.device, h.dtype) - if observation_cache is not None and observation_cache.zero_time_signature == zero_time_signature: - to = observation_cache.zero_time - else: - to = manual_cast(self.adaLN_modulation(self.t_embedder(torch.zeros_like(t))), h.dtype) - if observation_cache is not None: - observation_cache.zero_time = to - observation_cache.zero_time_signature = zero_time_signature - def run_group(observation: torch.Tensor, generated: torch.Tensor, start: int, stop: int) -> Tuple[torch.Tensor, torch.Tensor]: - for index in range(start, stop): - if self.base_block_policy == 'alternating_u_x' and not index % 2: - generated = self._refine(self.blocks[index], observation, ocoords, o_layout, - generated, x, tx, to, urows, u_layout, - observation_cache, f'blocks.{index}') - else: - observation, generated = self._base_block(self.blocks[index], observation, ocoords, o_layout, - generated, x, tx, to, - observation_cache, f'blocks.{index}') - key = str(index) - if key in self.refiners: - generated = self._refine(self.refiners[key], observation, ocoords, o_layout, - generated, x, tx, to, urows, u_layout, - observation_cache, f'refiners.{key}') - return observation, generated - for start in range(0, len(self.blocks), self.checkpoint_group_size): - stop = min(start + self.checkpoint_group_size, len(self.blocks)) - if self.use_checkpoint and self.training: - o, h = torch.utils.checkpoint.checkpoint(lambda a, b, begin=start, end=stop: run_group(a, b, begin, end), o, h, use_reentrant=False) - else: o, h = run_group(o, h, start, stop) - if not len(o): h = h + _live_zero(self.observation_metadata.parameters(), h) + self.observation_role.sum().to(h.dtype) * 0 - return self._finalize(x, h) +"""Interleaved flow transformer: alternate full-X and U-only updates with static observations. + +The observation attention-ILV weights remain a 60-block backbone. O is always read-only; +the production policy uses the original even blocks for U-only updates and +the original odd blocks for full-X updates. Explicitly configured, +identity-seeded post-block refiners remain available for engineering tests. +No O-update schedule or transition phase exists. +""" +from __future__ import annotations + +import copy +from dataclasses import dataclass +from typing import Any, Dict, Iterable, List, Optional, Sequence, Tuple + +import torch +from torch import nn +import torch.nn.functional as F +import torch.utils.checkpoint + +from ..modules import sparse as sp +from ..modules.sparse.attention.full_attn import sparse_scaled_dot_product_attention +from ..modules.utils import convert_module_to, manual_cast +from .observation_attention import _PSABlock, _join_exact_rows, _layouts_from_batch, _live_zero +from .interleaved_backbone import InterleavedBackbone + + +@dataclass +class _PreparedObservation: + hidden: torch.Tensor + coords: torch.Tensor + layout: List[slice] + urows: torch.Tensor + ucoords: torch.Tensor + ubatch: torch.Tensor + u_layout: List[slice] + + +class _StaticObservationCache: + """Invocation-local, eval-only cache for static O preparation and layer K/V.""" + def __init__(self, model: 'InterleavedFlowTransformer', x: sp.SparseTensor, shape: sp.SparseTensor, + observed: sp.SparseTensor, keep: Optional[torch.Tensor]): + if model.training or torch.is_grad_enabled(): + raise RuntimeError('observation_cache is inference-only; training cannot retain detached O K/V') + tensors = { + 'x.coords': x.coords, 'concat_cond.coords': shape.coords, 'concat_cond.feats': shape.feats, + 'partial_z_obs.coords': observed.coords, 'partial_z_obs.feats': observed.feats, + } + if keep is not None: + tensors['partial_z_obs_keep_mask'] = keep + inference = next((name for name, tensor in tensors.items() if torch.is_inference(tensor)), None) + if inference is not None: + raise RuntimeError( + f'observation_cache cannot safely validate inference-mode tensor {inference}; ' + 'pass regular tensors under torch.no_grad()' + ) + model_tensor = next( + (name for name, tensor in [*model.named_parameters(), *model.named_buffers()] + if torch.is_inference(tensor)), + None, + ) + if model_tensor is not None: + raise RuntimeError( + f'observation_cache cannot safely validate inference-mode model tensor {model_tensor}; ' + 'construct the model outside torch.inference_mode()' + ) + self.model = model + self.signature = model._cache_signature(x, shape, observed, keep) + self.kv: Dict[str, Tuple[torch.Tensor, torch.Tensor]] = {} + self.prepared: Optional[_PreparedObservation] = None + self.zero_time: Optional[torch.Tensor] = None + self.zero_time_signature: Optional[Tuple[Any, ...]] = None + self.bytes = 0 + self.closed = False + + @property + def preprocessing_bytes(self) -> int: + """Resident preparation tensors, separate from the existing K/V byte count.""" + tensors = [] if self.prepared is None else [ + self.prepared.hidden, self.prepared.coords, self.prepared.urows, + self.prepared.ucoords, self.prepared.ubatch, + ] + if self.zero_time is not None: + tensors.append(self.zero_time) + return sum(tensor.numel() * tensor.element_size() for tensor in tensors) + + def validate(self, x: sp.SparseTensor, shape: sp.SparseTensor, observed: sp.SparseTensor, + keep: Optional[torch.Tensor]) -> None: + if self.closed or self.model.training or torch.is_grad_enabled(): + raise RuntimeError('observation_cache is closed or used outside eval/no_grad') + if self.signature != self.model._cache_signature(x, shape, observed, keep): + raise RuntimeError('observation_cache inputs, model weights, dtype, device, or policy changed; rebuild the context') + + def get(self, key: str) -> Optional[Tuple[torch.Tensor, torch.Tensor]]: + return self.kv.get(key) + + def record(self, key: str, k: torch.Tensor, v: torch.Tensor) -> None: + if key not in self.kv: + self.kv[key] = (k.detach(), v.detach()) + self.bytes += (k.numel() * k.element_size()) + (v.numel() * v.element_size()) + + def close(self) -> None: + self.kv.clear() + self.prepared = None + self.zero_time = None + self.zero_time_signature = None + self.closed = True + + def __enter__(self) -> '_StaticObservationCache': + return self + + def __exit__(self, *_: Any) -> None: + self.close() + + +class InterleavedFlowTransformer(InterleavedBackbone): + """Inherited backbone with permanently static O and optional U refiners. + + O receives no Q projection, attention output, or MLP execution. It still + receives its layer-specific normalization/modulation and K/V projection so + every X layer attends over that layer's ``[O; X]`` keys and values. + """ + def __init__(self, refinement_indices: Sequence[int] = (), + base_block_policy: str = 'alternating_u_x', + allow_test_config: bool = False, + direct_modulation_broadcast: bool = True, **kwargs: Any): + retired = {'o_update_scale', 'o_transition_steps'} & kwargs.keys() + if retired: + raise ValueError(f'Interleaved flow transformer O is always static; removed options: {sorted(retired)}') + if base_block_policy not in {'alternating_u_x', 'all_x'}: + raise ValueError("base_block_policy must be 'alternating_u_x' or 'all_x'") + super().__init__(direct_modulation_broadcast=direct_modulation_broadcast, **kwargs) + if self.num_blocks % 2: + raise ValueError('Interleaved flow transformer requires an even number of base blocks') + if self.num_blocks != 60 and not allow_test_config: + raise ValueError('non-60-block Interleaved flow transformer test configurations require allow_test_config=True') + self.base_block_policy = base_block_policy + indices = self._validate_refinement_indices(refinement_indices) + self.refinement_indices = tuple(indices) + self.refiners = nn.ModuleDict({str(index): self._new_refiner(index) for index in indices}) + + def _validate_refinement_indices(self, indices: Sequence[int]) -> Tuple[int, ...]: + result = tuple(indices) + if tuple(sorted(set(result))) != result or any(index < 0 or index >= self.num_blocks for index in result): + raise ValueError('refinement_indices must be unique, sorted, valid zero-based base-block indices') + return result + + def _new_refiner(self, index: int) -> _PSABlock: + # The attention module does not retain a public rope-frequency field; + # a structural clone is therefore the exact construction contract. + block = copy.deepcopy(self.blocks[index]) + with torch.no_grad(): + nn.init.zeros_(block.self_attn.to_out.weight) + nn.init.zeros_(block.self_attn.to_out.bias) + nn.init.zeros_(block.mlp.mlp[2].weight) + nn.init.zeros_(block.mlp.mlp[2].bias) + return block + + def convert_to(self, dtype: torch.dtype) -> None: + """Keep post-parent refiners on the same explicit inference dtype.""" + super().convert_to(dtype) + if hasattr(self, 'refiners'): + self.refiners.apply(lambda module: convert_module_to(module, dtype=dtype)) + + def _cache_signature(self, x: sp.SparseTensor, shape: sp.SparseTensor, observed: sp.SparseTensor, + keep: Optional[torch.Tensor]) -> Tuple[Any, ...]: + """Extend the parent identity signature with static-O cache policy.""" + return super()._cache_signature(x, shape, observed, keep) + ( + ('interleaved_flow_transformer_static_o_kv', self.base_block_policy, self.direct_modulation_broadcast, self.dtype, + torch.is_autocast_enabled('cuda'), torch.get_autocast_dtype('cuda'), + torch.is_autocast_enabled('cpu'), torch.get_autocast_dtype('cpu')), + ) + + def observation_cache(self, x: sp.SparseTensor, concat_cond: sp.SparseTensor, + partial_z_obs: sp.SparseTensor, + partial_z_obs_keep_mask: Optional[torch.Tensor] = None) -> _StaticObservationCache: + """Create an invocation-local cache for static O preparation and K/V.""" + return _StaticObservationCache(self, x, concat_cond, partial_z_obs, partial_z_obs_keep_mask) + + def _prepare_observation(self, x: sp.SparseTensor, shape: sp.SparseTensor, + observed: sp.SparseTensor, keep: Optional[torch.Tensor]) -> _PreparedObservation: + """Prepare the time-independent O stream and the complementary U support.""" + if not torch.equal(shape.coords, x.coords): + raise AssertionError('Interleaved flow transformer requires coordinate-aligned PBR32 x and shape concat_cond') + batch_size = x.shape[0] + keep = torch.ones(batch_size, device=x.device, dtype=torch.bool) if keep is None else keep.to(x.device, torch.bool) + if keep.shape != (batch_size,): + raise AssertionError(f'partial_z_obs_keep_mask must be [B={batch_size}]') + active = keep[observed.coords[:, 0].long()] + ocoords, ofeats = observed.coords[active], observed.feats[active] + rows = _join_exact_rows(x.coords, ocoords) + o_layout = _layouts_from_batch(ocoords[:, 0], batch_size) + raw = ofeats[:, :32].float() * self.partial_std + self.partial_mean + opbr = (raw - self.pbr_mean) / self.pbr_std + oinput = torch.cat((opbr.to(ofeats.dtype), shape.feats[rows]), -1) + # Preserve the parent addition/cast ordering under autocast. + o = self.input_layer(sp.VarLenTensor(oinput, o_layout)).feats + o = o + self.observation_metadata(ofeats[:, 32:34].to(oinput.dtype)) + self.observation_role.to(oinput.dtype) + o = manual_cast(o, self.dtype) + observed_rows = torch.zeros(len(x.feats), device=x.device, dtype=torch.bool) + observed_rows[rows] = True + urows = torch.nonzero(~observed_rows, as_tuple=False).flatten() + ucoords = x.coords[urows] + ubatch = ucoords[:, 0].long() + u_layout = _layouts_from_batch(ubatch, batch_size) + return _PreparedObservation(o, ocoords, o_layout, urows, ucoords, ubatch, u_layout) + + def migrate_v11_state_dict(self, state_dict: Dict[str, torch.Tensor]) -> Dict[str, List[str]]: + raise RuntimeError('Interleaved flow transformer requires a interleaved backbone parent state; V11 migration is not a valid parent path') + + def load_v11_state_dict(self, state_dict: Dict[str, torch.Tensor]) -> Dict[str, List[str]]: + raise RuntimeError('Interleaved flow transformer requires a interleaved backbone parent state; V11 migration is not a valid parent path') + + def load_interleaved_backbone_state_dict(self, state_dict: Dict[str, torch.Tensor]) -> Dict[str, List[str]]: + """Strictly seed Interleaved flow transformer from an exact interleaved backbone state before writing. + + Base keys must be complete and shape-exact. Refiners inherit the + matching base block, then their two residual outputs are zeroed, making + each newly inserted operation exactly identity at migration time. + """ + target = self.state_dict() + base = {key for key in target if not key.startswith('refiners.')} + source_keys = set(state_dict) + missing = sorted(base - source_keys) + unexpected = sorted(source_keys - base) + non_tensors = sorted(key for key in base & source_keys if not isinstance(state_dict[key], torch.Tensor)) + mismatched = sorted(key for key in base & source_keys if isinstance(state_dict[key], torch.Tensor) + and (target[key].shape != state_dict[key].shape or target[key].dtype != state_dict[key].dtype)) + if missing or unexpected or non_tensors or mismatched: + raise RuntimeError('invalid interleaved backbone parent state before write: ' + f'missing={missing}, unexpected={unexpected}, non_tensors={non_tensors}, ' + f'shape_or_dtype_mismatch={mismatched}') + updates = {key: state_dict[key] for key in base} + seeded: List[str] = [] + zeroed: List[str] = [] + for index, refiner in self.refiners.items(): + for suffix, value in refiner.state_dict().items(): + destination = f'refiners.{index}.{suffix}' + source = f'blocks.{index}.{suffix}' + if source not in state_dict or target[destination].shape != state_dict[source].shape: + raise RuntimeError(f'invalid refiner seed source before write: {source}') + updates[destination] = state_dict[source] + seeded.append(destination) + for suffix in ('self_attn.to_out.weight', 'self_attn.to_out.bias', + 'mlp.mlp.2.weight', 'mlp.mlp.2.bias'): + destination = f'refiners.{index}.{suffix}' + updates[destination] = torch.zeros_like(target[destination]) + zeroed.append(destination) + self.load_state_dict(updates, strict=True) + return {'inherited': sorted(base), 'new': sorted(set(target) - base), + 'seeded': sorted(seeded), 'zeroed': sorted(zeroed)} + + def _mod(self, value: torch.Tensor, layout: List[slice], batch: torch.Tensor) -> torch.Tensor: + return self._direct_modulation(value, layout) if self.direct_modulation_broadcast else value[batch] + + @staticmethod + def _obatch(o: torch.Tensor, layout: List[slice]) -> torch.Tensor: + if not len(o): + return torch.empty(0, device=o.device, dtype=torch.long) + return torch.cat([torch.full((item.stop - item.start,), index, device=o.device, dtype=torch.long) + for index, item in enumerate(layout)]) + + def _qkv(self, block: _PSABlock, hidden: torch.Tensor, coords: torch.Tensor, + layout: List[slice], which: Iterable[int]) -> Dict[int, sp.VarLenTensor]: + """Project exactly requested Q/K/V slices; static O never materializes Q.""" + heads, dim = self.num_heads, self.model_channels // self.num_heads + selected = tuple(which) + if selected == (0, 1, 2): + q, k, v = self._project_qkv(block, hidden, coords, layout) + return {0: q, 1: k, 2: v} + weight = block.self_attn.to_qkv.weight.reshape(3, self.model_channels, self.model_channels) + bias = block.self_attn.to_qkv.bias.reshape(3, self.model_channels) + result: Dict[int, sp.VarLenTensor] = {} + if selected == (1, 2): + # One sliced K/V projection is the actual static-O saving; two + # independent linears would preserve semantics but lose the budget. + packed = F.linear(hidden, weight[1:].reshape(2 * self.model_channels, self.model_channels), + bias[1:].reshape(2 * self.model_channels)).reshape(-1, 2, heads, dim) + token = sp.VarLenTensor(packed, layout) + result[1], result[2] = token.unbind(1) + for item in selected: + if item in result: + continue + feats = F.linear(hidden, weight[item], bias[item]).reshape(-1, heads, dim) + result[item] = sp.VarLenTensor(feats, layout) + if 0 in result and block.self_attn.qk_rms_norm: + result[0] = block.self_attn.q_rms_norm(result[0]) + if 1 in result and block.self_attn.qk_rms_norm: + result[1] = block.self_attn.k_rms_norm(result[1]) + if len(coords) and (0 in result or 1 in result): + q = result.get(0); k = result.get(1) + if q is None: + dummy = sp.VarLenTensor(torch.zeros_like(result[1].feats), layout) + _, k = block.self_attn.rope(sp.SparseTensor(dummy.feats, coords), sp.SparseTensor(k.feats, coords)) + result[1] = sp.VarLenTensor(k.feats, layout) + elif k is None: + q, _ = block.self_attn.rope(sp.SparseTensor(q.feats, coords), sp.SparseTensor(torch.zeros_like(q.feats), coords)) + result[0] = sp.VarLenTensor(q.feats, layout) + else: + q, k = block.self_attn.rope(sp.SparseTensor(q.feats, coords), sp.SparseTensor(k.feats, coords)) + result[0], result[1] = sp.VarLenTensor(q.feats, layout), sp.VarLenTensor(k.feats, layout) + return result + + def _x_attention(self, block: _PSABlock, o: torch.Tensor, ocoords: torch.Tensor, o_layout: List[slice], + x: torch.Tensor, x_tensor: sp.SparseTensor, + cache: Optional[_StaticObservationCache] = None, cache_key: Optional[str] = None) -> torch.Tensor: + xqkv = self._qkv(block, x, x_tensor.coords, x_tensor.layout, (0, 1, 2)) + cached_kv = None if cache is None or cache_key is None else cache.get(cache_key) + if cached_kv is not None: + oqkv = {1: sp.VarLenTensor(cached_kv[0], o_layout), 2: sp.VarLenTensor(cached_kv[1], o_layout)} + else: + oqkv = self._qkv(block, o, ocoords, o_layout, (1, 2)) if len(o) else {} + if cache is not None and cache_key is not None and len(o): + cache.record(cache_key, oqkv[1].feats, oqkv[2].feats) + keys: List[torch.Tensor] = []; values: List[torch.Tensor] = []; layout: List[slice] = []; start = 0 + for os, xs in zip(o_layout, x_tensor.layout): + key = torch.cat(([oqkv[1].feats[os], xqkv[1].feats[xs]]) if len(o) else [xqkv[1].feats[xs]], 0) + value = torch.cat(([oqkv[2].feats[os], xqkv[2].feats[xs]]) if len(o) else [xqkv[2].feats[xs]], 0) + keys.append(key); values.append(value); layout.append(slice(start, start + len(key))); start += len(key) + kv = sp.VarLenTensor(torch.stack((torch.cat(keys), torch.cat(values)), dim=1), layout) + xout = sparse_scaled_dot_product_attention(xqkv[0], kv).feats + return block.self_attn.to_out(xout.reshape(len(x_tensor.feats), self.model_channels)) + + def _base_block(self, block: _PSABlock, o: torch.Tensor, ocoords: torch.Tensor, o_layout: List[slice], + h: torch.Tensor, x: sp.SparseTensor, tx: torch.Tensor, to: torch.Tensor, + cache: Optional[_StaticObservationCache] = None, cache_key: Optional[str] = None + ) -> Tuple[torch.Tensor, torch.Tensor]: + sx, scx, gx, smx, scm, gmx = block._modulation(tx) + xo = block.norm1(h) * (1 + self._mod(scx, x.layout, x.batch_boardcast_map)) + self._mod(sx, x.layout, x.batch_boardcast_map) + cached_kv = None if cache is None or cache_key is None else cache.get(cache_key) + if cached_kv is None: + so, sco, _go, _smo, _scmo, _gmo = block._modulation(to) + obatch = self._obatch(o, o_layout) if not self.direct_modulation_broadcast else torch.empty(0, device=o.device, dtype=torch.long) + oo = block.norm1(o) * (1 + self._mod(sco, o_layout, obatch)) + self._mod(so, o_layout, obatch) if len(o) else o + else: + oo = o + ax = self._x_attention(block, oo, ocoords, o_layout, xo, x, cache, cache_key) + h = h + ax * self._mod(gx, x.layout, x.batch_boardcast_map) + h = h + block.mlp.mlp(x.replace(block.norm2(h) * (1 + self._mod(scm, x.layout, x.batch_boardcast_map)) + self._mod(smx, x.layout, x.batch_boardcast_map))).feats * self._mod(gmx, x.layout, x.batch_boardcast_map) + return o, h + + def _refine(self, block: _PSABlock, o: torch.Tensor, ocoords: torch.Tensor, o_layout: List[slice], h: torch.Tensor, + x: sp.SparseTensor, tx: torch.Tensor, to: torch.Tensor, urows: torch.Tensor, u_layout: List[slice], + cache: Optional[_StaticObservationCache] = None, cache_key: Optional[str] = None) -> torch.Tensor: + if not len(urows): + return h + _live_zero(block.parameters(), h) + sx, scx, gx, smx, scm, gmx = block._modulation(tx) + xo = block.norm1(h) * (1 + self._mod(scx, x.layout, x.batch_boardcast_map)) + self._mod(sx, x.layout, x.batch_boardcast_map) + cached_kv = None if cache is None or cache_key is None else cache.get(cache_key) + if cached_kv is None: + so, sco, _go, _smo, _scmo, _gmo = block._modulation(to) + obatch = self._obatch(o, o_layout) if not self.direct_modulation_broadcast else torch.empty(0, device=o.device, dtype=torch.long) + oo = block.norm1(o) * (1 + self._mod(sco, o_layout, obatch)) + self._mod(so, o_layout, obatch) if len(o) else o + else: + oo = o + u = h[urows] + if cache is not None and cache.prepared is not None: + ucoords, ubatch = cache.prepared.ucoords, cache.prepared.ubatch + else: + ucoords, ubatch = x.coords[urows], x.coords[urows, 0].long() + uo = xo[urows] + # U is Q-only; K/V deliberately include O plus every current X row. + uq = self._qkv(block, uo, ucoords, u_layout, (0,))[0] + xkv = self._qkv(block, xo, x.coords, x.layout, (1, 2)) + cached_kv = None if cache is None or cache_key is None else cache.get(cache_key) + if cached_kv is not None: + okv = {1: sp.VarLenTensor(cached_kv[0], o_layout), 2: sp.VarLenTensor(cached_kv[1], o_layout)} + else: + okv = self._qkv(block, oo, ocoords, o_layout, (1, 2)) if len(o) else {} + if cache is not None and cache_key is not None and len(o): + cache.record(cache_key, okv[1].feats, okv[2].feats) + parts: List[torch.Tensor] = []; values: List[torch.Tensor] = []; layout: List[slice] = []; start = 0 + for os, xs in zip(o_layout, x.layout): + key = torch.cat(([okv[1].feats[os], xkv[1].feats[xs]]) if len(o) else [xkv[1].feats[xs]], 0) + value = torch.cat(([okv[2].feats[os], xkv[2].feats[xs]]) if len(o) else [xkv[2].feats[xs]], 0) + parts.append(key); values.append(value); count = key.shape[0]; layout.append(slice(start, start + count)); start += count + kv = sp.VarLenTensor(torch.stack((torch.cat(parts), torch.cat(values)), 1), layout) + u = u + block.self_attn.to_out(sparse_scaled_dot_product_attention(uq, kv).feats.reshape(len(u), self.model_channels)) * self._mod(gx, u_layout, ubatch) + u = u + block.mlp.mlp(sp.SparseTensor(block.norm2(u) * (1 + self._mod(scm, u_layout, ubatch)) + self._mod(smx, u_layout, ubatch), ucoords, shape=x.shape)).feats * self._mod(gmx, u_layout, ubatch) + h = h.clone(); h[urows] = u + return h + + def forward(self, x: sp.SparseTensor, t: torch.Tensor, cond: Any = None, concat_cond: Optional[sp.SparseTensor] = None, + partial_z_obs: Optional[sp.SparseTensor] = None, partial_z_obs_keep_mask: Optional[torch.Tensor] = None, + observation_cache: Optional[_StaticObservationCache] = None, **kwargs: Any) -> sp.SparseTensor: + retired = {'o_update_scale', 'o_transition_steps'} & kwargs.keys() + if retired: + raise ValueError(f'Interleaved flow transformer O is always static; removed options: {sorted(retired)}') + if observation_cache is not None: + if not isinstance(observation_cache, _StaticObservationCache) or observation_cache.model is not self: + raise RuntimeError('observation_cache belongs to a different Interleaved flow transformer model') + observation_cache.validate(x, concat_cond, partial_z_obs, partial_z_obs_keep_mask) + if kwargs: + disabled = next((name for name in ('use_observation_cache', 'compact_observation_cache', 'hot_observation_cache') if name in kwargs and kwargs[name] is not False), None) + if disabled is not None: raise RuntimeError(f'Interleaved flow transformer disables {disabled}') + if concat_cond is None or x.feats.shape[1] != 32 or concat_cond.feats.shape != x.feats.shape: + raise AssertionError('Interleaved flow transformer requires coordinate-aligned PBR32 x and shape concat_cond') + batch_size = x.shape[0] + if t.shape != (batch_size,): raise AssertionError(f't must be [B={batch_size}]') + if partial_z_obs is None or partial_z_obs.feats.ndim != 2 or partial_z_obs.feats.shape[1] != 34: + raise AssertionError('Interleaved flow transformer requires 34-channel partial_z_obs (possibly empty)') + prepared = None if observation_cache is None else observation_cache.prepared + if prepared is None: + prepared = self._prepare_observation(x, concat_cond, partial_z_obs, partial_z_obs_keep_mask) + if observation_cache is not None: + observation_cache.prepared = prepared + o, ocoords, o_layout = prepared.hidden, prepared.coords, prepared.layout + urows, u_layout = prepared.urows, prepared.u_layout + h = self.input_layer(x.replace(torch.cat((x.feats, concat_cond.feats), -1))).feats + h = manual_cast(h, self.dtype) + tx = manual_cast(self.adaLN_modulation(self.t_embedder(t)), h.dtype) + zero_time_signature = (t.dtype, t.device, h.dtype) + if observation_cache is not None and observation_cache.zero_time_signature == zero_time_signature: + to = observation_cache.zero_time + else: + to = manual_cast(self.adaLN_modulation(self.t_embedder(torch.zeros_like(t))), h.dtype) + if observation_cache is not None: + observation_cache.zero_time = to + observation_cache.zero_time_signature = zero_time_signature + def run_group(observation: torch.Tensor, generated: torch.Tensor, start: int, stop: int) -> Tuple[torch.Tensor, torch.Tensor]: + for index in range(start, stop): + if self.base_block_policy == 'alternating_u_x' and not index % 2: + generated = self._refine(self.blocks[index], observation, ocoords, o_layout, + generated, x, tx, to, urows, u_layout, + observation_cache, f'blocks.{index}') + else: + observation, generated = self._base_block(self.blocks[index], observation, ocoords, o_layout, + generated, x, tx, to, + observation_cache, f'blocks.{index}') + key = str(index) + if key in self.refiners: + generated = self._refine(self.refiners[key], observation, ocoords, o_layout, + generated, x, tx, to, urows, u_layout, + observation_cache, f'refiners.{key}') + return observation, generated + for start in range(0, len(self.blocks), self.checkpoint_group_size): + stop = min(start + self.checkpoint_group_size, len(self.blocks)) + if self.use_checkpoint and self.training: + o, h = torch.utils.checkpoint.checkpoint(lambda a, b, begin=start, end=stop: run_group(a, b, begin, end), o, h, use_reentrant=False) + else: o, h = run_group(o, h, start, stop) + if not len(o): h = h + _live_zero(self.observation_metadata.parameters(), h) + self.observation_role.sum().to(h.dtype) * 0 + return self._finalize(x, h) diff --git a/supratex/models/observation_stream.py b/supratex/models/observation_stream.py index 279dfce5f939c55d717f7e445ae3c4f6d0ae1b14..f29cea191c10a1ffd04db377329b79cdba22a7a8 100644 --- a/supratex/models/observation_stream.py +++ b/supratex/models/observation_stream.py @@ -1,158 +1,156 @@ -"""Optional O-stream refresh policy for the frozen observation attention architecture. - -The default policy deliberately delegates to :class:`ObservationAttentionModel` so it -remains bit-for-bit on the established path. The opt-in policy retains an -observation K/V projection at every layer because every X layer attends to its -own [O; X] keys and values, while avoiding O updates on selected layers. -""" -from __future__ import annotations - -from typing import Any, List, Optional, Tuple - -import torch - -from ..modules import sparse as sp -from ..modules.sparse.attention.full_attn import sparse_scaled_dot_product_attention -from .observation_attention import ( - ObservationAttentionModel, - _ObservationCache, - _PSABlock, -) - - -class ObservationStreamModel(ObservationAttentionModel): - """observation attention with opt-in sparse O-stream updates. - - ``o_refresh_stride`` refreshes O at blocks ``0, stride, 2 * stride, ...`` - before the terminal block. All blocks still project the current O state - to their own K/V for X's unchanged joint ``[O; X]`` self-attention. When - terminal relief is disabled, the inherited terminal O attention is retained - so the stride experiment changes no other policy. - """ - def __init__(self, terminal_o_kv_only: bool = False, o_refresh_stride: int = 1, - **kwargs: Any): - if not isinstance(terminal_o_kv_only, bool): - raise ValueError('terminal_o_kv_only must be a bool') - if isinstance(o_refresh_stride, bool) or not isinstance(o_refresh_stride, int) or o_refresh_stride < 1: - raise ValueError('o_refresh_stride must be a positive integer (bool is not allowed)') - super().__init__(**kwargs) - # Plain policy attributes intentionally add no checkpoint state. - self.terminal_o_kv_only = terminal_o_kv_only - self.o_refresh_stride = o_refresh_stride - - def _cache_signature(self, x: sp.SparseTensor, shape: sp.SparseTensor, observed: sp.SparseTensor, - keep: Optional[torch.Tensor]) -> Tuple[Any, ...]: - return super()._cache_signature(x, shape, observed, keep) + ( - ('o_refresh_policy', self.terminal_o_kv_only, self.o_refresh_stride), - ) - - def _project_qkv(self, block: _PSABlock, hidden: torch.Tensor, coords: torch.Tensor, - layout: List[slice]) -> Tuple[sp.VarLenTensor, sp.VarLenTensor, sp.VarLenTensor]: - """Retain the parent's full QKV projection for numerical/checkpoint parity.""" - packed = block.self_attn.to_qkv(hidden).reshape( - -1, 3, self.num_heads, self.model_channels // self.num_heads, - ) - token = sp.VarLenTensor(packed, layout) - q, k, v = token.unbind(1) - if block.self_attn.qk_rms_norm: - q, k = block.self_attn.q_rms_norm(q), block.self_attn.k_rms_norm(k) - if len(coords): - q_sparse, k_sparse = sp.SparseTensor(q.feats, coords), sp.SparseTensor(k.feats, coords) - q_sparse, k_sparse = block.self_attn.rope(q_sparse, k_sparse) - q, k = sp.VarLenTensor(q_sparse.feats, layout), sp.VarLenTensor(k_sparse.feats, layout) - return q, k, v - - def _attention(self, block: _PSABlock, o: torch.Tensor, o_coords: torch.Tensor, o_layout: List[slice], - x: torch.Tensor, x_tensor: sp.SparseTensor, - cached_kv: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, *, - compute_o_attention: Optional[bool] = None - ) -> Tuple[Optional[torch.Tensor], torch.Tensor, Optional[Tuple[torch.Tensor, torch.Tensor]]]: - """Materialize O->O only for a refresh or inherited terminal behavior. - - We deliberately still use ``to_qkv`` for O on K/V-only layers. Slicing - its weight would be a new numerical path and is not necessary to retain - distinct per-layer O K/V for the unchanged X attention. - """ - # Preserve the exact inherited helper for the parent _block call path. - # Besides keeping default parity literal, this accepts its positional - # ``cached_kv`` call convention. - if compute_o_attention is None: - return super()._attention(block, o, o_coords, o_layout, x, x_tensor, cached_kv) - qx, kx, vx = self._project_qkv(block, x, x_tensor.coords, x_tensor.layout) - if cached_kv is not None: - ko, vo = cached_kv - built_kv = None - elif len(o): - qo, ko, vo = self._project_qkv(block, o, o_coords, o_layout) - built_kv = (ko.feats, vo.feats) - else: - qo = None - ko = vo = o - built_kv = (o, o) - - if not len(o) and cached_kv is None: - kv = sp.VarLenTensor(torch.stack((kx.feats, vx.feats), dim=1), x_tensor.layout) - else: - kv_feats, kv_layout = [], [] - start = 0 - for os, xs in zip(o_layout, x_tensor.layout): - joined_k = torch.cat((ko[os], kx.feats[xs]), dim=0) if cached_kv is not None else torch.cat((ko.feats[os], kx.feats[xs]), dim=0) - joined_v = torch.cat((vo[os], vx.feats[xs]), dim=0) if cached_kv is not None else torch.cat((vo.feats[os], vx.feats[xs]), dim=0) - kv_feats.append(torch.stack((joined_k, joined_v), dim=1)) - kv_layout.append(slice(start, start + len(joined_k))) - start += len(joined_k) - kv = sp.VarLenTensor(torch.cat(kv_feats, dim=0), kv_layout) - x_out = sparse_scaled_dot_product_attention(qx, kv).feats - o_out = None - if compute_o_attention: - # A cache is eval/no-grad and therefore never permits an O update. - assert cached_kv is None and qo is not None - o_out = block.self_attn.to_out( - sparse_scaled_dot_product_attention(qo, ko, vo).feats.reshape(len(o), self.model_channels), - ) - return o_out, block.self_attn.to_out(x_out.reshape(len(x_out), self.model_channels)), built_kv - - @staticmethod - def _obatch(o: torch.Tensor, o_layout: List[slice]) -> torch.Tensor: - if not len(o): - return torch.empty(0, device=o.device, dtype=torch.long) - return torch.cat([ - torch.full((item.stop - item.start,), index, device=o.device, dtype=torch.long) - for index, item in enumerate(o_layout) - ]) - - def _block(self, block: _PSABlock, o: torch.Tensor, o_coords: torch.Tensor, o_layout: List[slice], - x: torch.Tensor, x_tensor: sp.SparseTensor, mod_x: torch.Tensor, mod_o: torch.Tensor, terminal: bool, - cache: Optional[_ObservationCache] = None, index: int = 0) -> Tuple[torch.Tensor, torch.Tensor]: - # Default O semantics remain the literal parent block, including its - # terminal dead O work. The parent forward dynamically calls this. - if self.o_refresh_stride == 1 and (not self.terminal_o_kv_only or not terminal): - return super()._block(block, o, o_coords, o_layout, x, x_tensor, mod_x, mod_o, terminal, cache, index) - sx, scx, gx, smx, scm, gmx = block._modulation(mod_x) - so, sco, go, smo, scmo, gmo = block._modulation(mod_o) - obatch = self._obatch(o, o_layout) - xo = block.norm1(x) * (1 + scx[x_tensor.batch_boardcast_map]) + sx[x_tensor.batch_boardcast_map] - oo = block.norm1(o) * (1 + sco[obatch]) + so[obatch] if len(o) else o - cached_kv = None if cache is None else cache.kv[index] - refresh_o = bool(len(o)) and not terminal and index % self.o_refresh_stride == 0 and cached_kv is None - # With the terminal relief disabled, retain the inherited dead terminal - # O attention so a stride-only experiment changes no other policy. - compute_o_attention = refresh_o or ( - terminal and not self.terminal_o_kv_only and bool(len(o)) and cached_kv is None - ) - ao, ax, built_kv = self._attention( - block, oo, o_coords, o_layout, xo, x_tensor, - compute_o_attention=compute_o_attention, cached_kv=cached_kv, - ) - if cache is not None and built_kv is not None: - cache.record(index, *built_kv) - x = x + ax * gx[x_tensor.batch_boardcast_map] - if refresh_o: - assert ao is not None - o = o + ao * go[obatch] - o_mlp = sp.VarLenTensor(block.norm2(o) * (1 + scmo[obatch]) + smo[obatch], o_layout) - o = o + block.mlp(o_mlp).feats * gmo[obatch] - x = x + block.mlp.mlp( - x_tensor.replace(block.norm2(x) * (1 + scm[x_tensor.batch_boardcast_map]) + smx[x_tensor.batch_boardcast_map]), - ).feats * gmx[x_tensor.batch_boardcast_map] - return o, x +"""Optional O-stream refresh policy for the frozen observation attention architecture. + +The default policy deliberately delegates to :class:`ObservationAttentionModel` so it +remains bit-for-bit on the established path. The opt-in policy retains an +observation K/V projection at every layer because every X layer attends to its +own [O; X] keys and values, while avoiding O updates on selected layers. +""" +from __future__ import annotations + +from typing import Any, List, Optional, Tuple + +import torch + +from ..modules import sparse as sp +from ..modules.sparse.attention.full_attn import sparse_scaled_dot_product_attention +from .observation_attention import ( + ObservationAttentionModel, + _ObservationCache, + _PSABlock, +) + + +class ObservationStreamModel(ObservationAttentionModel): + """observation attention with opt-in sparse O-stream updates. + + ``o_refresh_stride`` refreshes O at blocks ``0, stride, 2 * stride, ...`` + before the terminal block. All blocks still project the current O state + to their own K/V for X's unchanged joint ``[O; X]`` self-attention. When + terminal relief is disabled, the inherited terminal O attention is retained + so the stride experiment changes no other policy. + """ + def __init__(self, terminal_o_kv_only: bool = False, o_refresh_stride: int = 1, + **kwargs: Any): + if o_refresh_stride < 1: + raise ValueError('o_refresh_stride must be positive') + super().__init__(**kwargs) + # Plain policy attributes intentionally add no checkpoint state. + self.terminal_o_kv_only = terminal_o_kv_only + self.o_refresh_stride = o_refresh_stride + + def _cache_signature(self, x: sp.SparseTensor, shape: sp.SparseTensor, observed: sp.SparseTensor, + keep: Optional[torch.Tensor]) -> Tuple[Any, ...]: + return super()._cache_signature(x, shape, observed, keep) + ( + ('o_refresh_policy', self.terminal_o_kv_only, self.o_refresh_stride), + ) + + def _project_qkv(self, block: _PSABlock, hidden: torch.Tensor, coords: torch.Tensor, + layout: List[slice]) -> Tuple[sp.VarLenTensor, sp.VarLenTensor, sp.VarLenTensor]: + """Retain the parent's full QKV projection for numerical/checkpoint parity.""" + packed = block.self_attn.to_qkv(hidden).reshape( + -1, 3, self.num_heads, self.model_channels // self.num_heads, + ) + token = sp.VarLenTensor(packed, layout) + q, k, v = token.unbind(1) + if block.self_attn.qk_rms_norm: + q, k = block.self_attn.q_rms_norm(q), block.self_attn.k_rms_norm(k) + if len(coords): + q_sparse, k_sparse = sp.SparseTensor(q.feats, coords), sp.SparseTensor(k.feats, coords) + q_sparse, k_sparse = block.self_attn.rope(q_sparse, k_sparse) + q, k = sp.VarLenTensor(q_sparse.feats, layout), sp.VarLenTensor(k_sparse.feats, layout) + return q, k, v + + def _attention(self, block: _PSABlock, o: torch.Tensor, o_coords: torch.Tensor, o_layout: List[slice], + x: torch.Tensor, x_tensor: sp.SparseTensor, + cached_kv: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, *, + compute_o_attention: Optional[bool] = None + ) -> Tuple[Optional[torch.Tensor], torch.Tensor, Optional[Tuple[torch.Tensor, torch.Tensor]]]: + """Materialize O->O only for a refresh or inherited terminal behavior. + + We deliberately still use ``to_qkv`` for O on K/V-only layers. Slicing + its weight would be a new numerical path and is not necessary to retain + distinct per-layer O K/V for the unchanged X attention. + """ + # Preserve the exact inherited helper for the parent _block call path. + # Besides keeping default parity literal, this accepts its positional + # ``cached_kv`` call convention. + if compute_o_attention is None: + return super()._attention(block, o, o_coords, o_layout, x, x_tensor, cached_kv) + qx, kx, vx = self._project_qkv(block, x, x_tensor.coords, x_tensor.layout) + if cached_kv is not None: + ko, vo = cached_kv + built_kv = None + elif len(o): + qo, ko, vo = self._project_qkv(block, o, o_coords, o_layout) + built_kv = (ko.feats, vo.feats) + else: + qo = None + ko = vo = o + built_kv = (o, o) + + if not len(o) and cached_kv is None: + kv = sp.VarLenTensor(torch.stack((kx.feats, vx.feats), dim=1), x_tensor.layout) + else: + kv_feats, kv_layout = [], [] + start = 0 + for os, xs in zip(o_layout, x_tensor.layout): + joined_k = torch.cat((ko[os], kx.feats[xs]), dim=0) if cached_kv is not None else torch.cat((ko.feats[os], kx.feats[xs]), dim=0) + joined_v = torch.cat((vo[os], vx.feats[xs]), dim=0) if cached_kv is not None else torch.cat((vo.feats[os], vx.feats[xs]), dim=0) + kv_feats.append(torch.stack((joined_k, joined_v), dim=1)) + kv_layout.append(slice(start, start + len(joined_k))) + start += len(joined_k) + kv = sp.VarLenTensor(torch.cat(kv_feats, dim=0), kv_layout) + x_out = sparse_scaled_dot_product_attention(qx, kv).feats + o_out = None + if compute_o_attention: + # A cache is eval/no-grad and therefore never permits an O update. + assert cached_kv is None and qo is not None + o_out = block.self_attn.to_out( + sparse_scaled_dot_product_attention(qo, ko, vo).feats.reshape(len(o), self.model_channels), + ) + return o_out, block.self_attn.to_out(x_out.reshape(len(x_out), self.model_channels)), built_kv + + @staticmethod + def _obatch(o: torch.Tensor, o_layout: List[slice]) -> torch.Tensor: + if not len(o): + return torch.empty(0, device=o.device, dtype=torch.long) + return torch.cat([ + torch.full((item.stop - item.start,), index, device=o.device, dtype=torch.long) + for index, item in enumerate(o_layout) + ]) + + def _block(self, block: _PSABlock, o: torch.Tensor, o_coords: torch.Tensor, o_layout: List[slice], + x: torch.Tensor, x_tensor: sp.SparseTensor, mod_x: torch.Tensor, mod_o: torch.Tensor, terminal: bool, + cache: Optional[_ObservationCache] = None, index: int = 0) -> Tuple[torch.Tensor, torch.Tensor]: + # Default O semantics remain the literal parent block, including its + # terminal dead O work. The parent forward dynamically calls this. + if self.o_refresh_stride == 1 and (not self.terminal_o_kv_only or not terminal): + return super()._block(block, o, o_coords, o_layout, x, x_tensor, mod_x, mod_o, terminal, cache, index) + sx, scx, gx, smx, scm, gmx = block._modulation(mod_x) + so, sco, go, smo, scmo, gmo = block._modulation(mod_o) + obatch = self._obatch(o, o_layout) + xo = block.norm1(x) * (1 + scx[x_tensor.batch_boardcast_map]) + sx[x_tensor.batch_boardcast_map] + oo = block.norm1(o) * (1 + sco[obatch]) + so[obatch] if len(o) else o + cached_kv = None if cache is None else cache.kv[index] + refresh_o = bool(len(o)) and not terminal and index % self.o_refresh_stride == 0 and cached_kv is None + # With the terminal relief disabled, retain the inherited dead terminal + # O attention so a stride-only experiment changes no other policy. + compute_o_attention = refresh_o or ( + terminal and not self.terminal_o_kv_only and bool(len(o)) and cached_kv is None + ) + ao, ax, built_kv = self._attention( + block, oo, o_coords, o_layout, xo, x_tensor, + compute_o_attention=compute_o_attention, cached_kv=cached_kv, + ) + if cache is not None and built_kv is not None: + cache.record(index, *built_kv) + x = x + ax * gx[x_tensor.batch_boardcast_map] + if refresh_o: + assert ao is not None + o = o + ao * go[obatch] + o_mlp = sp.VarLenTensor(block.norm2(o) * (1 + scmo[obatch]) + smo[obatch], o_layout) + o = o + block.mlp(o_mlp).feats * gmo[obatch] + x = x + block.mlp.mlp( + x_tensor.replace(block.norm2(x) * (1 + scm[x_tensor.batch_boardcast_map]) + smx[x_tensor.batch_boardcast_map]), + ).feats * gmx[x_tensor.batch_boardcast_map] + return o, x diff --git a/supratex/models/structured_latent_flow.py b/supratex/models/structured_latent_flow.py index 50cb834086fc95d81d368fbc6a3fceaf122ebd68..329995ca7954a3493532d0178d62aab57d851f38 100644 --- a/supratex/models/structured_latent_flow.py +++ b/supratex/models/structured_latent_flow.py @@ -1,1391 +1,1386 @@ -from typing import * -from functools import partial -import torch -import torch.nn as nn -import torch.nn.functional as F -import numpy as np -from ..modules.utils import convert_module_to, manual_cast, str_to_dtype -from ..modules.transformer import AbsolutePositionEmbedder -from ..modules import sparse as sp -from ..modules.sparse.transformer import ModulatedSparseTransformerCrossBlock -from .sparse_structure_flow import TimestepEmbedder -from .sparse_elastic_mixin import SparseTransformerElasticMixin - - -class _PartialZObsInnovationGate(nn.Module): - """Residual, per-observed-token correction for the partial z_obs carrier.""" - - def __init__(self, model_channels: int, message_channels: int, dim: int, time_dim: int, num_blocks: int): - super().__init__() - self.alpha = nn.Parameter(torch.zeros(num_blocks)) - self.embed_y = nn.Linear(message_channels, dim, bias=True) - self.time_proj = nn.Linear(model_channels, time_dim, bias=True) - self.embed_x = nn.Linear(model_channels + time_dim, dim, bias=True) - self.mlp = nn.Sequential( - nn.Linear(dim + 2 + time_dim, dim, bias=True), - nn.SiLU(), - nn.Linear(dim, 1, bias=True), - ) - self.proj = nn.Linear(dim, model_channels, bias=True) - - -_LOCAL_ONE_HOP_OFFSETS = ( - (-1, 0, 0), (0, -1, 0), (0, 0, -1), - (0, 0, 1), (0, 1, 0), (1, 0, 0), -) -_LOCAL_TWO_HOP_OFFSETS = tuple(sorted( - (dx, dy, dz) - for dx in range(-2, 3) - for dy in range(-2, 3) - for dz in range(-2, 3) - if abs(dx) + abs(dy) + abs(dz) <= 2 -)) -assert len(_LOCAL_TWO_HOP_OFFSETS) == 25 - - -def _exact_coordinate_rows( - full_coords: torch.Tensor, - candidate_coords: torch.Tensor, -) -> Tuple[torch.Tensor, torch.Tensor]: - """Return exact full-support rows and a validity mask for candidate coords.""" - assert full_coords.ndim == candidate_coords.ndim == 2 - assert full_coords.shape[1] == candidate_coords.shape[1] - assert full_coords.shape[0] > 0, 'local attention requires nonempty full sparse support' - all_coords = torch.cat([full_coords, candidate_coords], dim=0) - _, inverse = torch.unique(all_coords, dim=0, return_inverse=True) - full_ids = inverse[:full_coords.shape[0]] - candidate_ids = inverse[full_coords.shape[0]:] - assert torch.unique(full_ids).numel() == full_ids.numel(), \ - 'x contains duplicate coordinates' - ordered_ids, ordered_rows = torch.sort(full_ids) - positions = torch.searchsorted(ordered_ids, candidate_ids) - in_range = positions < ordered_ids.numel() - safe_positions = positions.clamp(max=ordered_ids.numel() - 1) - valid = in_range & (ordered_ids[safe_positions] == candidate_ids) - rows = torch.where(valid, ordered_rows[safe_positions], -torch.ones_like(safe_positions)) - return rows, valid - - -def _build_two_hop_observed_topology( - full_coords: torch.Tensor, - observed_rows: torch.Tensor, -) -> Dict[str, torch.Tensor]: - """Pack exact batch-local G(S) distance<=2 observed source/query pairs.""" - assert full_coords.ndim == 2 and full_coords.shape[1] == 4, \ - f'local attention requires [N,4] integer coords, got {tuple(full_coords.shape)}' - assert not torch.is_floating_point(full_coords), 'local attention coordinates must be integer' - assert observed_rows.ndim == 1 and observed_rows.dtype == torch.long - # Validate full-coordinate uniqueness even for an all-absent batch. - _exact_coordinate_rows(full_coords, full_coords.new_empty((0, full_coords.shape[1]))) - if observed_rows.numel() == 0: - empty = torch.empty(0, dtype=torch.long, device=full_coords.device) - return { - 'query_rows': empty, - 'source_indices': empty, - 'offset_indices': empty, - 'segment_indices': empty, - 'neighbour_counts': empty, - } - assert observed_rows.min() >= 0 and observed_rows.max() < full_coords.shape[0], \ - 'observed row lies outside full sparse support' - assert torch.unique(observed_rows).numel() == observed_rows.numel(), \ - 'observed rows must be unique' - - device = full_coords.device - one_hop = torch.tensor( - _LOCAL_ONE_HOP_OFFSETS, dtype=full_coords.dtype, device=device, - ) - source_coords = full_coords[observed_rows] - candidate_hop1 = source_coords[:, None, :].expand(-1, len(_LOCAL_ONE_HOP_OFFSETS), -1).clone() - candidate_hop1[:, :, 1:] += one_hop[None] - hop1_rows, hop1_valid = _exact_coordinate_rows( - full_coords, candidate_hop1.reshape(-1, full_coords.shape[1]), - ) - source_ids_hop1 = torch.arange( - observed_rows.numel(), device=device, dtype=torch.long, - )[:, None].expand(-1, len(_LOCAL_ONE_HOP_OFFSETS)).reshape(-1) - hop1_rows = hop1_rows[hop1_valid] - source_ids_hop1 = source_ids_hop1[hop1_valid] - - if hop1_rows.numel() != 0: - candidate_hop2 = full_coords[hop1_rows][:, None, :].expand( - -1, len(_LOCAL_ONE_HOP_OFFSETS), -1, - ).clone() - candidate_hop2[:, :, 1:] += one_hop[None] - hop2_rows, hop2_valid = _exact_coordinate_rows( - full_coords, candidate_hop2.reshape(-1, full_coords.shape[1]), - ) - source_ids_hop2 = source_ids_hop1[:, None].expand( - -1, len(_LOCAL_ONE_HOP_OFFSETS), - ).reshape(-1) - hop2_rows = hop2_rows[hop2_valid] - source_ids_hop2 = source_ids_hop2[hop2_valid] - else: - hop2_rows = torch.empty(0, dtype=torch.long, device=device) - source_ids_hop2 = torch.empty(0, dtype=torch.long, device=device) - - source_ids = torch.cat([ - torch.arange(observed_rows.numel(), device=device, dtype=torch.long), - source_ids_hop1, - source_ids_hop2, - ]) - query_rows = torch.cat([observed_rows, hop1_rows, hop2_rows]) - # Multiple graph paths to the same query/source pair are one neighbourhood entry. - unique_pairs = torch.unique(torch.stack([query_rows, source_ids], dim=1), dim=0) - query_rows = unique_pairs[:, 0] - source_ids = unique_pairs[:, 1] - query_coords = full_coords[query_rows] - selected_source_coords = source_coords[source_ids] - assert torch.equal(query_coords[:, 0], selected_source_coords[:, 0]), \ - 'local attention constructed a cross-sample edge' - relative = query_coords[:, 1:] - selected_source_coords[:, 1:] - offset_table = torch.tensor( - _LOCAL_TWO_HOP_OFFSETS, dtype=full_coords.dtype, device=device, - ) - offset_matches = (relative[:, None, :] == offset_table[None, :, :]).all(dim=-1) - assert torch.equal(offset_matches.sum(dim=1), torch.ones_like(offset_matches[:, 0], dtype=torch.long)), \ - 'local attention produced an invalid two-hop offset' - offset_indices = offset_matches.to(dtype=torch.long).argmax(dim=1) - - # Canonical per-query real-KV order is lexicographic relative offset, independent - # of either full-support or observed-message row ordering. - order = torch.argsort(query_rows * len(_LOCAL_TWO_HOP_OFFSETS) + offset_indices) - query_rows = query_rows[order] - source_ids = source_ids[order] - offset_indices = offset_indices[order] - active_query_rows, segment_indices, counts = torch.unique_consecutive( - query_rows, return_inverse=True, return_counts=True, - ) - assert counts.max() <= len(_LOCAL_TWO_HOP_OFFSETS), \ - 'local observed neighbourhood exceeds frozen 25-source cap' - return { - 'query_rows': active_query_rows, - 'source_indices': source_ids, - 'offset_indices': offset_indices, - 'segment_indices': segment_indices, - 'neighbour_counts': counts, - } - - -class _PartialZObsLocalAttention(nn.Module): - """Shared packed two-hop observed-innovation attention for V12-A.""" - - def __init__( - self, - model_channels: int, - message_channels: int, - dim: int, - time_dim: int, - num_heads: int, - num_active_blocks: int, - ): - super().__init__() - assert dim == 128 and time_dim == 32 and num_heads == 8 - assert dim % num_heads == 0 - self.dim = dim - self.num_heads = num_heads - self.head_dim = dim // num_heads - self.alpha = nn.Parameter(torch.zeros(num_active_blocks)) - self.time_proj = nn.Linear(model_channels, time_dim, bias=True) - self.embed_x = nn.Linear(model_channels + time_dim, dim, bias=True) - self.embed_y = nn.Linear(message_channels, dim, bias=True) - self.key_proj = nn.Linear(dim, dim, bias=True) - self.value_proj = nn.Linear(dim, dim, bias=True) - self.relative_offset = nn.Parameter(torch.empty(len(_LOCAL_TWO_HOP_OFFSETS), dim)) - self.null_key = nn.Parameter(torch.empty(num_heads, self.head_dim)) - self.out_proj = nn.Linear(dim, model_channels, bias=False) - - @staticmethod - def _rms_norm(value: torch.Tensor) -> torch.Tensor: - return value * torch.rsqrt(value.square().mean(dim=-1, keepdim=True) + 1e-6) - - def forward( - self, - hidden: torch.Tensor, - full_batch: torch.Tensor, - t_base: torch.Tensor, - message: torch.Tensor, - observed_rows: torch.Tensor, - topology: Dict[str, torch.Tensor], - block_index: int, - ) -> Tuple[torch.Tensor, Dict[str, torch.Tensor]]: - assert 0 <= block_index < self.alpha.numel(), 'terminal block has no V12-A residual' - query_rows = topology['query_rows'] - source_indices = topology['source_indices'] - offset_indices = topology['offset_indices'] - segment_indices = topology['segment_indices'] - counts = topology['neighbour_counts'] - assert full_batch.ndim == 1 and full_batch.shape[0] == hidden.shape[0] - assert full_batch.numel() == 0 or ( - full_batch.min() >= 0 and full_batch.max() < t_base.shape[0] - ), 'full sparse batch coordinate is outside timestep batch' - assert message.shape == (observed_rows.numel(), 34) - - if query_rows.numel() == 0: - live_zero = sum(parameter.sum() for parameter in self.parameters()) * 0 - empty_delta = hidden.new_empty((0, hidden.shape[1])) + live_zero - zero = hidden.new_zeros(()) - return empty_delta, { - 'null_attention_mean': zero, - 'real_neighbour_entropy_mean': zero, - 'effective_real_neighbour_count_mean': zero, - 'active_queries': zero, - 'singleton_queries': zero, - 'packed_pairs': zero, - } - - branch_dtype = self.embed_x.weight.dtype - hidden_branch = hidden.to(dtype=branch_dtype) - time = self.time_proj(t_base.to(dtype=branch_dtype)) - x_embed = self.embed_x(torch.cat([ - F.layer_norm(hidden_branch, hidden_branch.shape[-1:], eps=1e-6), - time[full_batch], - ], dim=-1)) - innovation = self.embed_y(message.to(dtype=branch_dtype)) - x_embed[observed_rows] - query = self._rms_norm( - x_embed[query_rows].reshape(-1, self.num_heads, self.head_dim) - ) - real_key = self.key_proj(innovation)[source_indices] - real_key = real_key + self.relative_offset[offset_indices] - real_key = self._rms_norm(real_key.reshape(-1, self.num_heads, self.head_dim)) - real_value = self.value_proj(innovation)[source_indices].reshape( - -1, self.num_heads, self.head_dim, - ) - real_logits = (query[segment_indices] * real_key).sum(dim=-1) / np.sqrt(self.head_dim) - null_key = self._rms_norm(self.null_key) - null_logits = (query * null_key[None]).sum(dim=-1) / np.sqrt(self.head_dim) - - segment_heads = segment_indices[:, None].expand(-1, self.num_heads) - maxima = null_logits.clone() - maxima.scatter_reduce_(0, segment_heads, real_logits, reduce='amax', include_self=True) - real_exp = torch.exp(real_logits - maxima[segment_indices]) - null_exp = torch.exp(null_logits - maxima) - real_sum = torch.zeros_like(null_exp).scatter_add(0, segment_heads, real_exp) - denominator = null_exp + real_sum - real_weight = real_exp / denominator[segment_indices] - null_weight = null_exp / denominator - attended = hidden_branch.new_zeros( - (query_rows.numel(), self.num_heads, self.head_dim), - ) - attended.index_add_( - 0, segment_indices, real_weight[..., None] * real_value, - ) - delta = self.out_proj(attended.reshape(query_rows.numel(), self.dim)) - delta = delta * self.alpha[block_index].to(dtype=delta.dtype) - - # Entropy is conditional on selecting real evidence. Null-vs-real - # acceptance is reported independently as null_attention_mean; mixing - # that mass into this distribution would make singleton entropy nonzero. - real_conditional_weight = real_exp / real_sum[segment_indices].clamp_min( - torch.finfo(real_sum.dtype).tiny - ) - entropy_terms = -real_conditional_weight * torch.log( - real_conditional_weight.clamp_min( - torch.finfo(real_conditional_weight.dtype).tiny - ) - ) - real_entropy = torch.zeros_like(null_weight).scatter_add( - 0, segment_heads, entropy_terms, - ) - return delta, { - 'null_attention_mean': null_weight.mean(), - 'real_neighbour_entropy_mean': real_entropy.mean(), - 'effective_real_neighbour_count_mean': real_entropy.exp().mean(), - 'active_queries': delta.new_tensor(float(query_rows.numel())), - 'singleton_queries': delta.new_tensor(float((counts == 1).sum())), - 'packed_pairs': delta.new_tensor(float(source_indices.numel())), - } - - -class _V12LACrossRope(nn.Module): - """Apply the same 3D rotary basis to independently packed Q and K rows.""" - - def __init__( - self, - head_dim: int, - rope_freq: Tuple[float, float], - reserved_dims: int = 2, - ): - super().__init__() - content_dim = head_dim - reserved_dims - assert content_dim > 0 and content_dim % 6 == 0 - self.head_dim = head_dim - self.content_dim = content_dim - freq_dim = content_dim // 2 // 3 - frequencies = torch.arange(freq_dim, dtype=torch.float32) / freq_dim - frequencies = rope_freq[0] / (rope_freq[1] ** frequencies) - self.register_buffer('frequencies', frequencies, persistent=False) - - def forward(self, value: torch.Tensor, coords: torch.Tensor) -> torch.Tensor: - assert value.ndim == 3 and value.shape[-1] == self.head_dim - assert coords.shape == (value.shape[0], 3) - if value.shape[0] == 0: - return value - phases = torch.outer( - coords.reshape(-1).float(), self.frequencies, - ).reshape(coords.shape[0], self.content_dim // 2) - phases = torch.polar(torch.ones_like(phases), phases).unsqueeze(1) - content = torch.view_as_complex( - value[..., :self.content_dim].float().reshape( - value.shape[0], value.shape[1], -1, 2, - ) - ) - rotated = torch.view_as_real(content * phases).reshape( - value.shape[0], value.shape[1], self.content_dim, - ).to(dtype=value.dtype) - return torch.cat([rotated, value[..., self.content_dim:]], dim=-1) - - -class _PartialZObsLearnedAttention(nn.Module): - """Five-site V12LA observation attention with an exact carrier anchor.""" - - def __init__( - self, - model_channels: int, - num_heads: int, - sites: Sequence[int], - rope_freq: Tuple[float, float], - null_prior: float = 0.5, - ): - super().__init__() - self.sites = tuple(int(site) for site in sites) - assert self.sites == (5, 10, 15, 20, 25), self.sites - assert num_heads == 12 - assert model_channels % num_heads == 0 - assert (model_channels // num_heads - 2) > 0 - assert (model_channels // num_heads - 2) % 6 == 0 - self.model_channels = model_channels - self.num_heads = num_heads - self.head_dim = model_channels // num_heads - self.content_dim = self.head_dim - 2 - self.site_to_slot = {site: slot for slot, site in enumerate(self.sites)} - self.null_prior = float(null_prior) - assert 0.0 < self.null_prior < 1.0, \ - f'learned-attention null prior must lie in (0,1), got {self.null_prior}' - self.null_logit_offset = float( - np.log(self.null_prior) - np.log1p(-self.null_prior) - ) - - # K/V are source-only and shared across all five sites. Q and the - # zero-start deviation/output matrix are site-specific. - self.key_proj = nn.Linear(model_channels, model_channels, bias=False) - self.value_proj = nn.Linear(model_channels, model_channels, bias=False) - self.query_proj = nn.ModuleList([ - nn.Linear(model_channels, model_channels, bias=False) - for _ in self.sites - ]) - self.delta_proj = nn.ModuleList([ - nn.Linear(model_channels, model_channels, bias=False) - for _ in self.sites - ]) - self.null_key = nn.Parameter(torch.empty(num_heads, self.content_dim)) - self.rope = _V12LACrossRope(self.head_dim, rope_freq, reserved_dims=2) - self.reset_parameters() - - def reset_parameters(self) -> None: - nn.init.xavier_uniform_(self.key_proj.weight) - # Attention mixes the already trained V11 carrier basis at warm start; - # the value map may then adapt, but does not begin as a random rewrite. - nn.init.eye_(self.value_proj.weight) - nn.init.xavier_uniform_(self.null_key) - for query_proj in self.query_proj: - nn.init.xavier_uniform_(query_proj.weight) - for delta_proj in self.delta_proj: - nn.init.zeros_(delta_proj.weight) - - @staticmethod - def _rms_norm(value: torch.Tensor) -> torch.Tensor: - normalized = value.float() * torch.rsqrt( - value.float().square().mean(dim=-1, keepdim=True) + 1e-6 - ) - return normalized.to(dtype=value.dtype) - - @staticmethod - def _anchor( - carrier: torch.Tensor, - observed_rows: torch.Tensor, - full_rows: int, - ) -> torch.Tensor: - """Evaluate fixed coordinate/null attention without an additive scatter.""" - assert carrier.ndim == 2 - assert observed_rows.ndim == 1 and observed_rows.dtype == torch.long - assert carrier.shape[0] == observed_rows.numel() - assert full_rows >= observed_rows.numel() - if observed_rows.numel() != 0: - assert observed_rows.min() >= 0 and observed_rows.max() < full_rows - assert torch.unique(observed_rows).numel() == observed_rows.numel() - # Index zero is the exact-zero null value. Every observed row is then - # redirected to its own carrier row; unmatched rows keep selecting null. - values = torch.cat([carrier.new_zeros((1, carrier.shape[1])), carrier], dim=0) - anchor_index = torch.zeros(full_rows, dtype=torch.long, device=carrier.device) - anchor_index[observed_rows] = torch.arange( - 1, observed_rows.numel() + 1, dtype=torch.long, device=carrier.device, - ) - return values[anchor_index] - - def build_cache( - self, - full_coords: torch.Tensor, - full_layout: List[slice], - carrier: torch.Tensor, - observed_rows: torch.Tensor, - ) -> Dict[str, Any]: - """Build the exact anchor and the shared per-sample O+null K/V cache.""" - assert full_coords.ndim == 2 and full_coords.shape[1] == 4 - assert carrier.shape == (observed_rows.numel(), self.model_channels) - assert len(full_layout) > 0 - anchor = self._anchor(carrier, observed_rows, full_coords.shape[0]) - observed_coords = full_coords[observed_rows] - observed_batch = observed_coords[:, 0].long() - assert observed_batch.numel() == 0 or ( - observed_batch.min() >= 0 and observed_batch.max() < len(full_layout) - ) - - branch_dtype = self.key_proj.weight.dtype - carrier_branch = carrier.to(dtype=branch_dtype) - carrier_norm = F.layer_norm( - carrier_branch, (self.model_channels,), eps=1e-6, - ) - real_key = self.key_proj(carrier_norm).reshape( - -1, self.num_heads, self.head_dim, - ) - real_key = self.rope(real_key, observed_coords[:, 1:]) - real_key = self._rms_norm(real_key[..., :self.content_dim]) - real_key = torch.cat([ - real_key, - real_key.new_zeros((real_key.shape[0], self.num_heads, 2)), - ], dim=-1) - real_value = self.value_proj(carrier_branch).reshape( - -1, self.num_heads, self.head_dim, - ) - null_content = self._rms_norm( - self.null_key.to(dtype=real_key.dtype) - ) - - packed_kv = [] - observed_counts = [] - for batch_index in range(len(full_layout)): - source_indices = torch.nonzero( - observed_batch == batch_index, as_tuple=False, - ).flatten() - observed_count = int(source_indices.numel()) - observed_counts.append(observed_count) - # One null competes with O real keys. log(O) cancels its cardinality - # disadvantage; the constant logit offset sets the null mass only in - # the equal-content-logit reference case, not after learned Q/K drift. - null_reserved = real_key.new_zeros((self.num_heads, 2)) - null_reserved[:, 0] = np.sqrt(self.head_dim) * ( - np.log(max(observed_count, 1)) + self.null_logit_offset - ) - null_key = torch.cat([null_content, null_reserved], dim=-1).unsqueeze(0) - null_value = real_value.new_zeros((1, self.num_heads, self.head_dim)) - keys = torch.cat([real_key[source_indices], null_key], dim=0) - values = torch.cat([real_value[source_indices], null_value], dim=0) - packed_kv.append(torch.stack([keys, values], dim=1)) - - kv = sp.VarLenTensor.from_tensor_list(packed_kv) - return { - 'anchor': anchor, - 'kv': kv, - 'observed_rows': observed_rows, - 'observed_counts': tuple(observed_counts), - } - - def forward( - self, - hidden: torch.Tensor, - full_coords: torch.Tensor, - full_layout: List[slice], - cache: Dict[str, Any], - block_index: int, - ) -> Tuple[torch.Tensor, Dict[str, torch.Tensor]]: - assert block_index in self.site_to_slot, block_index - assert hidden.shape == (full_coords.shape[0], self.model_channels) - slot = self.site_to_slot[block_index] - branch_dtype = self.query_proj[slot].weight.dtype - hidden_branch = hidden.to(dtype=branch_dtype) - query = self.query_proj[slot](F.layer_norm( - hidden_branch, (self.model_channels,), eps=1e-6, - )).reshape(-1, self.num_heads, self.head_dim) - query = self.rope(query, full_coords[:, 1:]) - query_content = self._rms_norm(query[..., :self.content_dim]) - query_reserved = query.new_zeros((query.shape[0], self.num_heads, 2)) - query_reserved[..., 0] = 1 - query = torch.cat([query_content, query_reserved], dim=-1) - query_varlen = sp.VarLenTensor(query, full_layout) - attended = sp.sparse_scaled_dot_product_attention( - query_varlen, cache['kv'], - ).feats.reshape(-1, self.model_channels) - - anchor = cache['anchor'].to(dtype=attended.dtype) - deviation_input = attended - anchor - correction = self.delta_proj[slot](deviation_input) - inject = anchor + correction - eps = torch.finfo(inject.dtype).eps - observed_rows = cache['observed_rows'] - observed_mask = torch.zeros( - inject.shape[0], dtype=torch.bool, device=inject.device, - ) - observed_mask[observed_rows] = True - return inject, { - 'correction_to_anchor_norm_ratio': correction.norm() / anchor.norm().clamp_min(eps), - 'correction_to_hidden_norm_ratio': correction.norm() / hidden_branch.norm().clamp_min(eps), - 'observed_correction_l2_norm': correction[observed_mask].float().norm(), - 'unobserved_correction_l2_norm': correction[~observed_mask].float().norm(), - 'attention_output_l2_norm': attended.float().norm(), - 'anchor_l2_norm': anchor.float().norm(), - 'delta_weight_l2_norm': self.delta_proj[slot].weight.float().norm(), - } - -class SLatFlowModel(nn.Module): - """ - Structured Latent Flow Model for 3D generation. - - Supports two conditioning modes: - - "cross": Standard cross-attention with image features - - "proj": View-aligned projection attention with camera-aware features - """ - def __init__( - self, - resolution: int, - in_channels: int, - model_channels: int, - cond_channels: int, - out_channels: int, - num_blocks: int, - num_heads: Optional[int] = None, - num_head_channels: Optional[int] = 64, - mlp_ratio: float = 4, - pe_mode: Literal["ape", "rope"] = "ape", - rope_freq: Tuple[float, float] = (1.0, 10000.0), - dtype: str = 'float32', - use_checkpoint: bool = False, - share_mod: bool = False, - initialization: str = 'vanilla', - qk_rms_norm: bool = False, - qk_rms_norm_cross: bool = False, - image_attn_mode: Literal["cross", "proj", "gated_proj"] = "cross", - proj_in_channels: Optional[int] = None, - vae_in_channels: Optional[int] = None, - zobs_inject: bool = False, - zobs_inject_in_channels: int = 33, - zobs_inject_scale: float = 1.0, - partial_zobs_inject: bool = False, - partial_zobs_in_channels: int = 34, - partial_zobs_inject_position: Literal['post_block', 'pre_mlp'] = 'post_block', - legacy_zobs_zero_input: bool = False, - partial_zobs_innovation_gate: bool = False, - innovation_gate_dim: int = 128, - innovation_gate_time_dim: int = 32, - partial_zobs_innovation_gate_train_only: bool = False, - partial_zobs_local_attention: bool = False, - local_attention_dim: int = 128, - local_attention_time_dim: int = 32, - local_attention_heads: int = 8, - partial_zobs_local_attention_train_only: bool = False, - partial_zobs_learned_attention: bool = False, - learned_attention_sites: Sequence[int] = (5, 10, 15, 20, 25), - learned_attention_heads: int = 12, - learned_attention_null_prior: float = 0.5, - partial_zobs_learned_attention_train_only: bool = False, - **kwargs - ): - super().__init__() - self.resolution = resolution - self.in_channels = in_channels - self.model_channels = model_channels - self.cond_channels = cond_channels - self.out_channels = out_channels - self.num_blocks = num_blocks - self.num_heads = num_heads or model_channels // num_head_channels - self.mlp_ratio = mlp_ratio - self.pe_mode = pe_mode - self.use_checkpoint = use_checkpoint - self.share_mod = share_mod - self.initialization = initialization - self.qk_rms_norm = qk_rms_norm - self.qk_rms_norm_cross = qk_rms_norm_cross - self.image_attn_mode = image_attn_mode - self.proj_in_channels = proj_in_channels - self.vae_in_channels = vae_in_channels - self.dtype = str_to_dtype(dtype) - - self.t_embedder = TimestepEmbedder(model_channels) - if share_mod: - self.adaLN_modulation = nn.Sequential( - nn.SiLU(), - nn.Linear(model_channels, 6 * model_channels, bias=True) - ) - - if pe_mode == "ape": - self.pos_embedder = AbsolutePositionEmbedder(model_channels) - - self.input_layer = sp.SparseLinear(in_channels, model_channels) - # Trained-bridge: single shared per-block injector of [z_obs (x) vis (32) | vis (1)] -> model_channels. - self.zobs_inject = zobs_inject - self.zobs_inject_scale = float(zobs_inject_scale) - self.zobs_injector = sp.SparseLinear(zobs_inject_in_channels, model_channels) if zobs_inject else None - # v8-partial: the observation lives only on its observed support. This adapter is - # deliberately bias-free: absent rows never acquire a learned content/fill value. - self.partial_zobs_inject = partial_zobs_inject - self.partial_zobs_in_channels = partial_zobs_in_channels - assert partial_zobs_inject_position in ('post_block', 'pre_mlp'), \ - f'partial_zobs_inject_position must be post_block or pre_mlp, got {partial_zobs_inject_position!r}' - self.partial_zobs_inject_position = partial_zobs_inject_position - if partial_zobs_inject_position == 'pre_mlp': - assert not any(( - partial_zobs_innovation_gate, - partial_zobs_local_attention, - partial_zobs_learned_attention, - )), 'partial_zobs_inject_position=pre_mlp is incompatible with post-block experimental branches' - self.legacy_zobs_zero_input = legacy_zobs_zero_input - self.partial_zobs_injector = ( - sp.SparseLinear(partial_zobs_in_channels, model_channels, bias=False) - if partial_zobs_inject else None - ) - self.partial_zobs_innovation_gate = partial_zobs_innovation_gate - self.partial_zobs_innovation_gate_train_only = partial_zobs_innovation_gate_train_only - if partial_zobs_innovation_gate_train_only: - assert partial_zobs_innovation_gate, \ - 'partial_zobs_innovation_gate_train_only requires partial_zobs_innovation_gate=true' - if partial_zobs_innovation_gate: - assert partial_zobs_inject, 'partial_zobs_innovation_gate requires partial_zobs_inject=true' - assert partial_zobs_in_channels == 34, \ - f'partial_zobs_innovation_gate requires 34 input channels, got {partial_zobs_in_channels}' - assert innovation_gate_dim == 128, f'innovation_gate_dim is fixed at 128, got {innovation_gate_dim}' - assert innovation_gate_time_dim == 32, \ - f'innovation_gate_time_dim is fixed at 32, got {innovation_gate_time_dim}' - assert model_channels == 1536, f'partial_zobs_innovation_gate requires 1536 model channels, got {model_channels}' - assert num_blocks == 30, f'partial_zobs_innovation_gate requires 30 blocks, got {num_blocks}' - self.innovation_gate = _PartialZObsInnovationGate( - model_channels, partial_zobs_in_channels, innovation_gate_dim, - innovation_gate_time_dim, num_blocks, - ) - else: - self.innovation_gate = None - self.partial_zobs_local_attention = partial_zobs_local_attention - self.partial_zobs_local_attention_train_only = partial_zobs_local_attention_train_only - if partial_zobs_local_attention_train_only: - assert partial_zobs_local_attention, \ - 'partial_zobs_local_attention_train_only requires partial_zobs_local_attention=true' - assert sum(bool(value) for value in ( - partial_zobs_innovation_gate, - partial_zobs_local_attention, - partial_zobs_learned_attention, - )) <= 1, 'partial-zobs experimental attention/gate branches are mutually exclusive' - if partial_zobs_local_attention: - assert partial_zobs_inject, 'partial_zobs_local_attention requires partial_zobs_inject=true' - assert partial_zobs_in_channels == 34, \ - f'partial_zobs_local_attention requires 34 input channels, got {partial_zobs_in_channels}' - assert local_attention_dim == 128, \ - f'local_attention_dim is fixed at 128, got {local_attention_dim}' - assert local_attention_time_dim == 32, \ - f'local_attention_time_dim is fixed at 32, got {local_attention_time_dim}' - assert local_attention_heads == 8, \ - f'local_attention_heads is fixed at 8, got {local_attention_heads}' - assert model_channels == 1536, \ - f'partial_zobs_local_attention requires 1536 model channels, got {model_channels}' - assert num_blocks == 30, \ - f'partial_zobs_local_attention requires 30 blocks, got {num_blocks}' - self.local_obs_attention = _PartialZObsLocalAttention( - model_channels, partial_zobs_in_channels, local_attention_dim, - local_attention_time_dim, local_attention_heads, num_blocks - 1, - ) - else: - self.local_obs_attention = None - - self.partial_zobs_learned_attention = partial_zobs_learned_attention - self.partial_zobs_learned_attention_train_only = \ - partial_zobs_learned_attention_train_only - if partial_zobs_learned_attention_train_only: - assert partial_zobs_learned_attention, \ - 'partial_zobs_learned_attention_train_only requires partial_zobs_learned_attention=true' - if partial_zobs_learned_attention: - assert partial_zobs_inject, \ - 'partial_zobs_learned_attention requires partial_zobs_inject=true' - assert partial_zobs_in_channels == 34, \ - f'partial_zobs_learned_attention requires 34 input channels, got {partial_zobs_in_channels}' - assert model_channels == 1536, \ - f'partial_zobs_learned_attention requires 1536 model channels, got {model_channels}' - assert num_blocks == 30, \ - f'partial_zobs_learned_attention requires 30 blocks, got {num_blocks}' - assert learned_attention_heads == 12, \ - f'partial_zobs_learned_attention requires 12 heads, got {learned_attention_heads}' - assert tuple(learned_attention_sites) == (5, 10, 15, 20, 25), \ - f'V12LA sites are fixed at [5,10,15,20,25], got {learned_attention_sites}' - self.learned_attention = _PartialZObsLearnedAttention( - model_channels, - learned_attention_heads, - learned_attention_sites, - rope_freq, - learned_attention_null_prior, - ) - else: - self.learned_attention = None - - self.blocks = nn.ModuleList([ - ModulatedSparseTransformerCrossBlock( - model_channels, - cond_channels, - num_heads=self.num_heads, - mlp_ratio=self.mlp_ratio, - attn_mode='full', - use_checkpoint=self.use_checkpoint, - use_rope=(pe_mode == "rope"), - rope_freq=rope_freq, - share_mod=self.share_mod, - qk_rms_norm=self.qk_rms_norm, - qk_rms_norm_cross=self.qk_rms_norm_cross, - image_attn_mode=image_attn_mode, - proj_in_channels=proj_in_channels, - vae_in_channels=vae_in_channels, - ) - for _ in range(num_blocks) - ]) - - self.out_layer = sp.SparseLinear(model_channels, out_channels) - - self.initialize_weights() - self.convert_to(self.dtype) - if self.partial_zobs_innovation_gate_train_only: - self._freeze_for_partial_zobs_innovation_gate() - if self.partial_zobs_local_attention_train_only: - self._freeze_for_partial_zobs_local_attention() - if self.partial_zobs_learned_attention_train_only: - self._freeze_for_partial_zobs_learned_attention() - - @property - def device(self) -> torch.device: - """ - Return the device of the model. - """ - return next(self.parameters()).device - - def convert_to(self, dtype: torch.dtype) -> None: - """ - Convert the torso of the model to the specified dtype. - """ - self.dtype = dtype - self.blocks.apply(partial(convert_module_to, dtype=dtype)) - if self.learned_attention is not None: - self.learned_attention.apply(partial(convert_module_to, dtype=dtype)) - - def initialize_weights(self) -> None: - if self.initialization == 'vanilla': - # Initialize transformer layers: - def _basic_init(module): - if isinstance(module, nn.Linear): - torch.nn.init.xavier_uniform_(module.weight) - if module.bias is not None: - nn.init.constant_(module.bias, 0) - self.apply(_basic_init) - - # Initialize timestep embedding MLP: - nn.init.normal_(self.t_embedder.mlp[0].weight, std=0.02) - nn.init.normal_(self.t_embedder.mlp[2].weight, std=0.02) - - # Zero-out adaLN modulation layers in DiT blocks: - if self.share_mod: - nn.init.constant_(self.adaLN_modulation[-1].weight, 0) - nn.init.constant_(self.adaLN_modulation[-1].bias, 0) - else: - for block in self.blocks: - nn.init.constant_(block.adaLN_modulation[-1].weight, 0) - nn.init.constant_(block.adaLN_modulation[-1].bias, 0) - - # Zero-out output layers: - nn.init.constant_(self.out_layer.weight, 0) - nn.init.constant_(self.out_layer.bias, 0) - - elif self.initialization == 'scaled': - # Initialize transformer layers: - def _basic_init(module): - if isinstance(module, nn.Linear): - torch.nn.init.normal_(module.weight, std=np.sqrt(2.0 / (5.0 * self.model_channels))) - if module.bias is not None: - nn.init.constant_(module.bias, 0) - self.apply(_basic_init) - - # Scaled init for to_out and ffn2 - def _scaled_init(module): - if isinstance(module, nn.Linear): - torch.nn.init.normal_(module.weight, std=1.0 / np.sqrt(5 * self.num_blocks * self.model_channels)) - if module.bias is not None: - nn.init.constant_(module.bias, 0) - for block in self.blocks: - block.self_attn.to_out.apply(_scaled_init) - # Handle cross, proj, and gated_proj modes - if self.image_attn_mode in ("proj", "gated_proj"): - block.cross_attn.cross_attn_block.to_out.apply(_scaled_init) - else: - block.cross_attn.to_out.apply(_scaled_init) - block.mlp.mlp[2].apply(_scaled_init) - - # Initialize input layer to make the initial representation have variance 1 - nn.init.normal_(self.input_layer.weight, std=1.0 / np.sqrt(self.in_channels)) - nn.init.zeros_(self.input_layer.bias) - - # Initialize timestep embedding MLP: - nn.init.normal_(self.t_embedder.mlp[0].weight, std=0.02) - nn.init.normal_(self.t_embedder.mlp[2].weight, std=0.02) - - # Zero-out adaLN modulation layers in DiT blocks: - if self.share_mod: - nn.init.constant_(self.adaLN_modulation[-1].weight, 0) - nn.init.constant_(self.adaLN_modulation[-1].bias, 0) - else: - for block in self.blocks: - nn.init.constant_(block.adaLN_modulation[-1].weight, 0) - nn.init.constant_(block.adaLN_modulation[-1].bias, 0) - - # Zero-out output layers: - nn.init.constant_(self.out_layer.weight, 0) - nn.init.constant_(self.out_layer.bias, 0) - - # Trained-bridge: zero-init the z_obs injector LAST (after _basic_init), so step-0 injection == 0. - if getattr(self, 'zobs_injector', None) is not None: - nn.init.zeros_(self.zobs_injector.weight) - nn.init.zeros_(self.zobs_injector.bias) - if getattr(self, 'partial_zobs_injector', None) is not None: - nn.init.zeros_(self.partial_zobs_injector.weight) - if getattr(self, 'innovation_gate', None) is not None: - nn.init.zeros_(self.innovation_gate.alpha) - if getattr(self, 'local_obs_attention', None) is not None: - # V12-A freezes vanilla Xavier initialization independently of the - # scaled V11 backbone initialization mode. Only alpha starts at zero. - for module in self.local_obs_attention.modules(): - if isinstance(module, nn.Linear): - nn.init.xavier_uniform_(module.weight) - if module.bias is not None: - nn.init.zeros_(module.bias) - nn.init.xavier_uniform_(self.local_obs_attention.relative_offset) - nn.init.xavier_uniform_(self.local_obs_attention.null_key) - nn.init.zeros_(self.local_obs_attention.alpha) - if getattr(self, 'learned_attention', None) is not None: - self.learned_attention.reset_parameters() - - def _freeze_for_partial_zobs_innovation_gate(self) -> None: - """Make the V12 gate the sole trainable branch before an optimizer is built.""" - assert self.innovation_gate is not None - for parameter in self.parameters(): - parameter.requires_grad_(False) - for parameter in self.innovation_gate.parameters(): - parameter.requires_grad_(True) - trainable = [name for name, parameter in self.named_parameters() if parameter.requires_grad] - expected = [f'innovation_gate.{name}' for name in self.innovation_gate.state_dict()] - assert trainable == expected, (trainable, expected) - assert sum(parameter.numel() for parameter in self.parameters() if parameter.requires_grad) == 473663 - - def _freeze_for_partial_zobs_local_attention(self) -> None: - """Make the V12-A local attention the sole trainable warm-start branch.""" - assert self.local_obs_attention is not None - for parameter in self.parameters(): - parameter.requires_grad_(False) - for parameter in self.local_obs_attention.parameters(): - parameter.requires_grad_(True) - trainable = [name for name, parameter in self.named_parameters() if parameter.requires_grad] - expected = [f'local_obs_attention.{name}' for name in self.local_obs_attention.state_dict()] - assert trainable == expected, (trainable, expected) - assert len(trainable) == 14 - assert sum(parameter.numel() for parameter in self.parameters() if parameter.requires_grad) == 487485 - - def _freeze_for_partial_zobs_learned_attention(self) -> None: - """Make the V12LA module the sole trainable warm-start branch.""" - assert self.learned_attention is not None - for parameter in self.parameters(): - parameter.requires_grad_(False) - for parameter in self.learned_attention.parameters(): - parameter.requires_grad_(True) - trainable = [name for name, parameter in self.named_parameters() if parameter.requires_grad] - expected = [ - f'learned_attention.{name}' - for name, _ in self.learned_attention.named_parameters() - ] - assert trainable == expected, (trainable, expected) - assert len(trainable) == 13, trainable - - def _partial_zobs_message( - self, - x: sp.SparseTensor, - partial_z_obs: sp.SparseTensor, - partial_z_obs_keep_mask: Optional[torch.Tensor], - ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - """Project a variable-support observation and join its messages onto ``x`` rows. - - The return value stays variable-length: projected messages plus their exact full-row - indices. Neither observation content nor a zero-filled model-channel tensor is - materialized on full support. ``torch.unique(..., dim=0)`` makes the coordinate join - exact and rejects duplicate or non-subset partial coordinates instead of assuming a - row order. - """ - adapter = self.partial_zobs_injector - assert adapter is not None - assert partial_z_obs.feats.ndim == 2 and \ - partial_z_obs.feats.shape[-1] == self.partial_zobs_in_channels, \ - f'partial_z_obs must be [N,{self.partial_zobs_in_channels}], got {tuple(partial_z_obs.feats.shape)}' - assert partial_z_obs.coords.ndim == 2 and partial_z_obs.coords.shape[1] == x.coords.shape[1], \ - f'partial_z_obs coordinate rank mismatch: {tuple(partial_z_obs.coords.shape)} vs {tuple(x.coords.shape)}' - - active = torch.ones(partial_z_obs.feats.shape[0], dtype=torch.bool, device=partial_z_obs.feats.device) - if partial_z_obs_keep_mask is not None: - assert partial_z_obs_keep_mask.ndim == 1 and partial_z_obs_keep_mask.shape[0] == x.shape[0], \ - f'partial_z_obs_keep_mask must be [B={x.shape[0]}], got {tuple(partial_z_obs_keep_mask.shape)}' - keep = partial_z_obs_keep_mask.to(device=partial_z_obs.feats.device, dtype=torch.bool) - active = keep[partial_z_obs.coords[:, 0].long()] - - active_coords = partial_z_obs.coords[active] - active_feats = partial_z_obs.feats[active] - # F.linear on a zero-row tensor is intentional: it keeps the adapter parameters in the - # graph on all-drop batches, so DDP sees a zero (rather than None) gradient. - active_messages = F.linear(active_feats, adapter.weight, None) - - full_rows = torch.empty(0, dtype=torch.long, device=x.feats.device) - if active_coords.shape[0] != 0: - all_coords = torch.cat([x.coords, active_coords], dim=0) - _, inverse = torch.unique(all_coords, dim=0, return_inverse=True) - x_ids = inverse[:x.coords.shape[0]] - partial_ids = inverse[x.coords.shape[0]:] - assert torch.unique(x_ids).numel() == x_ids.numel(), 'x contains duplicate coordinates' - assert torch.unique(partial_ids).numel() == partial_ids.numel(), 'partial_z_obs contains duplicate coordinates' - ordered_x_ids, ordered_x_rows = torch.sort(x_ids) - positions = torch.searchsorted(ordered_x_ids, partial_ids) - assert (positions < ordered_x_ids.numel()).all() and \ - torch.equal(ordered_x_ids[positions], partial_ids), \ - 'partial_z_obs coordinates must be an exact subset of current x coordinates' - full_rows = ordered_x_rows[positions] - - return active_messages, full_rows, active_feats, active_coords[:, 0].long() - - def forward( - self, - x: sp.SparseTensor, - t: torch.Tensor, - cond: Union[torch.Tensor, List[torch.Tensor], Dict[str, Union[torch.Tensor, sp.SparseTensor]], Tuple], - concat_cond: Optional[sp.SparseTensor] = None, - z_obs: Optional[sp.SparseTensor] = None, - partial_z_obs: Optional[sp.SparseTensor] = None, - partial_z_obs_keep_mask: Optional[torch.Tensor] = None, - **kwargs - ) -> sp.SparseTensor: - """ - Forward pass. - - Args: - x: SparseTensor input - t: Timestep tensor [B] - cond: Conditioning tensor. For "cross" mode: list of tensors or tensor. - For "proj" mode: dict {'global': global_cond, 'proj': proj_cond} - or tuple of (global_cond, proj_cond) - concat_cond: Optional concatenation condition - z_obs: Optional 33ch [z_obs(32)|vis(1)] sparse tensor, same coords as x. Feeds ONLY - the zobs_injector (v6 patchcross path: concat_cond is the plain 32ch shape - latent and z_obs never enters the input_layer). When None, the legacy 97ch - concat slice below applies — shipped/v4 behavior is bit-identical. - partial_z_obs: Optional variable-length v8 observation - [normalized_z_obs(32)|coverage(1)|boundary_distance(1)]. - Its coordinates must be a subset of x; it is consumed only by the dedicated, - bias-free partial adapter and is never converted to full-support content. - - Returns: - SparseTensor output - """ - if concat_cond is not None: - x = sp.sparse_cat([x, concat_cond], dim=-1) - - h = self.input_layer(x) - h = manual_cast(h, self.dtype) - - # Trained-bridge: compute the per-block injection vector ONCE (shared zero-init linear, added each block). - # Sources, in priority order: separate z_obs kwarg (v6), else the trailing channels of a - # 97-ch [x_t|shape|z_obs|vis] input (v4). z_obs*vis==0 at vis=0 in both. - zobs_inject = None - _injector = getattr(self, 'zobs_injector', None) - if z_obs is not None: - # Fail fast: without an injector the kwarg would be a silent no-op, and in the 64ch - # patchcross architecture the injector is the ONLY z_obs consumption path. - assert _injector is not None, \ - 'z_obs kwarg passed but the denoiser was built without zobs_inject=true' - # Full coords equality is asserted at data-load; a row-count check catches gross mismatch. - assert z_obs.feats.shape[-1] == 33 and z_obs.feats.shape[0] == x.feats.shape[0], \ - f'z_obs kwarg mismatch: {tuple(z_obs.feats.shape)} vs x rows {x.feats.shape[0]}' - _zo = z_obs.feats[:, :32] - _vis = z_obs.feats[:, 32:33] - zobs_inject = manual_cast(self.zobs_injector(z_obs.replace(torch.cat([_zo * _vis, _vis], dim=-1))), self.dtype) - if not getattr(self, '_zobs_kwarg_logged', False): # in-log proof hook (once per process) - print('[zobs_injector] consuming separate z_obs kwarg (patchcross path)', flush=True) - self._zobs_kwarg_logged = True - elif _injector is not None and x.feats.shape[-1] >= 97: - _zo = x.feats[:, 64:96] - _vis = x.feats[:, 96:97] - zobs_inject = manual_cast(self.zobs_injector(x.replace(torch.cat([_zo * _vis, _vis], dim=-1))), self.dtype) - elif _injector is not None and self.legacy_zobs_zero_input: - # 64ch v8 shell with a legacy injector: preserve the legacy all-zero z_obs mode, - # including its learned per-block bias, without materialising observation content. - # Evaluate one zero row and broadcast the result: this is exactly W*0+b, keeps a - # zero (non-None) weight gradient under DDP, and avoids an N-row 33ch GEMM. - zero_output = F.linear( - x.feats.new_zeros((1, _injector.in_features)), - _injector.weight, - _injector.bias, - ) - zobs_inject = manual_cast( - x.replace(zero_output.expand(x.feats.shape[0], -1)), self.dtype - ) - elif _injector is not None and not getattr(self, '_zobs_none_warned', False): - # Injector present but NO z_obs source (64ch input, kwarg omitted): the trained - # per-block bias is NOT applied — a third mode distinct from "z_obs dropped" (an - # all-zero 33ch tensor still adds the bias each block). A nozobs arm on a patchcross - # model must pass the zero tensor, never omit the kwarg. - print('[zobs_injector] WARNING: injector present but no z_obs source (kwarg omitted, ' - 'input < 97ch) — per-block bias NOT applied; pass an all-zero 33ch z_obs for a ' - 'true nozobs arm', flush=True) - self._zobs_none_warned = True - - partial_zobs_inject = None - partial_zobs_rows = None - partial_zobs_feats = None - partial_zobs_batch = None - if partial_z_obs is not None: - assert self.partial_zobs_injector is not None, \ - 'partial_z_obs passed but denoiser was built without partial_zobs_inject=true' - partial_zobs_inject, partial_zobs_rows, partial_zobs_feats, partial_zobs_batch = self._partial_zobs_message( - x, partial_z_obs, partial_z_obs_keep_mask - ) - partial_zobs_inject = manual_cast(partial_zobs_inject, self.dtype) - elif self.partial_zobs_injector is not None: - raise AssertionError('partial_zobs_inject=true requires a separate partial_z_obs tensor') - - if zobs_inject is not None and self.zobs_inject_scale != 1.0: - zobs_inject = zobs_inject.replace(zobs_inject.feats * self.zobs_inject_scale) - if not getattr(self, '_zobs_inject_scale_logged', False): - print(f'[zobs_injector] output scale={self.zobs_inject_scale:g}', flush=True) - self._zobs_inject_scale_logged = True - - # Keep the pre-AdaLN embedding for the innovation gate; blocks retain their existing - # modulation input exactly. - t_base = self.t_embedder(t) - t_emb = t_base - if self.share_mod: - t_emb = self.adaLN_modulation(t_emb) - t_emb = manual_cast(t_emb, self.dtype) - - innovation_time = None - innovation_y = None - innovation_gate_y = None - if self.innovation_gate is not None and partial_zobs_feats is not None: - # Observation-only terms are shared by all blocks. Empty rows deliberately pass - # through these linears too, so an all-drop batch keeps the full gate graph live. - gate_dtype = self.innovation_gate.embed_y.weight.dtype - innovation_gate_y = partial_zobs_feats.to(dtype=gate_dtype) - innovation_y = self.innovation_gate.embed_y(innovation_gate_y) - innovation_time = self.innovation_gate.time_proj(t_base) - - if self.pe_mode == "ape": - pe = self.pos_embedder(h.coords[:, 1:]) - h = h + manual_cast(pe, self.dtype) - - # Handle different conditioning modes - if self.image_attn_mode == 'proj': - if isinstance(cond, dict): - global_cond = cond['global'] - proj_cond = cond['proj'] - else: - global_cond, proj_cond = cond - if isinstance(global_cond, list): - global_cond = sp.VarLenTensor.from_tensor_list(global_cond) - global_cond = manual_cast(global_cond, self.dtype) - proj_cond = manual_cast(proj_cond, self.dtype) - cond = (global_cond, proj_cond) - elif self.image_attn_mode == 'gated_proj': - global_cond = cond['global'] - if isinstance(global_cond, list): - global_cond = sp.VarLenTensor.from_tensor_list(global_cond) - global_cond = manual_cast(global_cond, self.dtype) - proj_semantic = manual_cast(cond['proj_semantic'], self.dtype) - proj_color = manual_cast(cond['proj_color'], self.dtype) - cond = {'global': global_cond, 'proj_semantic': proj_semantic, 'proj_color': proj_color} - else: - if isinstance(cond, list): - cond = sp.VarLenTensor.from_tensor_list(cond) - cond = manual_cast(cond, self.dtype) - - # Inference-only probe hook. When the attribute is absent (the default, and the - # only state any training path produces) this is None and the loop below is - # byte-identical to the unhooked path: no constructor arg, no config key, and no - # entry in the inference ABI contract. A probe sets it to a length-num_blocks - # sequence to rescale the shared partial message per block. - _partial_scales = getattr(self, 'partial_zobs_layer_scales', None) - # Optional evaluator-only callable: ``fn(gate_value, active_batch, block_index)``. - # It is an instance attribute rather than a constructor/config argument, so no model - # state or training ABI changes; callers can implement neutral, constant, or shuffled - # gate-value probes deterministically. - _gate_value_transform = getattr(self, 'partial_zobs_innovation_gate_value_transform', None) - _innovation_ratio_stats = [] if self.innovation_gate is not None else None - _innovation_gate_mean_stats = [] if self.innovation_gate is not None else None - _innovation_observed_rows = 0 if partial_zobs_rows is None else int(partial_zobs_rows.numel()) - _local_attention = getattr(self, 'local_obs_attention', None) - _local_topology = None - _local_full_batch = None - _local_ratio_base_stats = [] if _local_attention is not None else None - _local_ratio_hidden_stats = [] if _local_attention is not None else None - _local_null_stats = [] if _local_attention is not None else None - _local_entropy_stats = [] if _local_attention is not None else None - _local_effective_neighbour_stats = [] if _local_attention is not None else None - _local_observed_residual_stats = [] if _local_attention is not None else None - _local_unobserved_residual_stats = [] if _local_attention is not None else None - _local_query_is_observed = None - _local_neighbourhood_histogram = None - if _local_attention is not None: - assert partial_zobs_rows is not None and partial_zobs_feats is not None - _local_topology = _build_two_hop_observed_topology(x.coords, partial_zobs_rows) - _local_full_batch = x.coords[:, 0].long() - _active_observed_mask = torch.zeros( - x.coords.shape[0], dtype=torch.bool, device=x.coords.device, - ) - _active_observed_mask[partial_zobs_rows] = True - _local_query_is_observed = _active_observed_mask[ - _local_topology['query_rows'] - ] - _local_neighbourhood_histogram = torch.bincount( - _local_topology['neighbour_counts'], minlength=26, - )[:26] - _local_neighbourhood_histogram[0] = ( - x.coords.shape[0] - _local_topology['query_rows'].numel() - ) - assert int(_local_neighbourhood_histogram.sum()) == x.coords.shape[0] - assert int(_local_neighbourhood_histogram[1:].sum()) == \ - _local_topology['query_rows'].numel() - _learned_attention = getattr(self, 'learned_attention', None) - _learned_cache = None - _learned_stats = [] if _learned_attention is not None else None - if _learned_attention is not None: - assert partial_zobs_inject is not None and partial_zobs_rows is not None - assert _partial_scales is None, \ - 'partial_zobs_layer_scales is not defined for the five-site V12LA operator' - _learned_cache = _learned_attention.build_cache( - x.coords, - x.layout, - partial_zobs_inject, - partial_zobs_rows, - ) - partial_zobs_inject_position = getattr( - self, 'partial_zobs_inject_position', 'post_block', - ) - for _block_index, block in enumerate(self.blocks): - if partial_zobs_inject is not None and partial_zobs_inject_position == 'pre_mlp': - _pre_mlp_scale = 1.0 if _partial_scales is None else float( - _partial_scales[_block_index] - ) - h = block( - h, - t_emb, - cond, - partial_zobs_rows=partial_zobs_rows, - partial_zobs_inject=partial_zobs_inject, - partial_zobs_inject_scale=_pre_mlp_scale, - ) - else: - h = block(h, t_emb, cond) - if zobs_inject is not None: - h = h + zobs_inject - _innovation_ratio = None - _innovation_gate_mean = None - _local_delta = None - _local_stats = None - if _local_attention is not None and _block_index < _local_attention.alpha.numel(): - assert _local_topology is not None and _local_full_batch is not None - _local_delta, _local_stats = _local_attention( - h.feats, - _local_full_batch, - t_base, - partial_zobs_feats, - partial_zobs_rows, - _local_topology, - _block_index, - ) - _local_delta = manual_cast(_local_delta, self.dtype) - _local_ratio_base = _local_delta.norm() / partial_zobs_inject.norm().clamp_min( - torch.finfo(partial_zobs_inject.dtype).eps - ) - _local_hidden = h.feats[_local_topology['query_rows']] - _local_ratio_hidden = _local_delta.norm() / _local_hidden.norm().clamp_min( - torch.finfo(_local_hidden.dtype).eps - ) - assert _local_query_is_observed is not None - _local_observed_residual = _local_delta[ - _local_query_is_observed - ].float().norm() - _local_unobserved_residual = _local_delta[ - ~_local_query_is_observed - ].float().norm() - if _learned_attention is not None and _block_index in _learned_attention.site_to_slot: - assert _learned_cache is not None - _learned_inject, _site_stats = _learned_attention( - h.feats, - h.coords, - h.layout, - _learned_cache, - _block_index, - ) - h = h.replace(h.feats + manual_cast(_learned_inject, self.dtype)) - _learned_stats.append({ - name: value.detach() for name, value in _site_stats.items() - }) - if partial_zobs_inject is not None and _learned_attention is None and \ - partial_zobs_inject_position == 'post_block': - layer_inject = partial_zobs_inject - if self.innovation_gate is not None: - gate = self.innovation_gate - # Every operation intentionally accepts an empty leading dimension. This - # keeps all innovation parameters in the all-drop DDP graph. - assert innovation_gate_y is not None and innovation_y is not None and innovation_time is not None - t_obs = innovation_time[partial_zobs_batch] - h_obs = h.feats[partial_zobs_rows].to(dtype=gate.embed_y.weight.dtype) - x_embed = gate.embed_x(torch.cat([h_obs, t_obs], dim=-1)) - innovation = innovation_y - x_embed - gate_input = torch.cat([ - innovation, - innovation_gate_y[:, 32:34], - t_obs, - ], dim=-1) - gate_value = torch.sigmoid(gate.mlp(gate_input)) - if _gate_value_transform is not None: - gate_value = _gate_value_transform( - gate_value, partial_zobs_batch, _block_index - ) - assert gate_value.shape == (innovation.shape[0], 1), \ - f'gate value transform must return [Nobs,1], got {tuple(gate_value.shape)}' - delta = gate.proj(innovation) * gate_value - delta = delta * gate.alpha[_block_index].to(dtype=delta.dtype) - if _block_index == len(self.blocks) - 1: - # Keep the frozen V11 base carrier at the terminal site, - # but never add a newly learned observed-only shortcut: - # no attention block remains to transport/smooth it into U. - delta = delta * 0 - # Detached telemetry for the training log. This is before the legacy - # layer-scale hook and therefore measures the learned branch itself. - _innovation_ratio = delta.norm() / partial_zobs_inject.norm().clamp_min( - torch.finfo(partial_zobs_inject.dtype).eps - ) - _innovation_gate_mean = ( - gate_value.mean() if gate_value.numel() != 0 else delta.new_zeros(()) - ) - layer_inject = partial_zobs_inject + manual_cast(delta, self.dtype) - if partial_zobs_rows.numel() != 0: - if _partial_scales is None: - h = h.replace(h.feats.index_add( - 0, partial_zobs_rows, layer_inject - )) - else: - h = h.replace(h.feats.index_add( - 0, partial_zobs_rows, layer_inject, - alpha=float(_partial_scales[_block_index]), - )) - else: - # Preserve every empty carrier/gate GEMM in the graph on all-drop batches. - h = h.replace(h.feats + layer_inject.sum() * 0) - if _local_delta is not None: - _local_scale = 1.0 if _partial_scales is None else float( - _partial_scales[_block_index] - ) - if _local_topology['query_rows'].numel() != 0: - h = h.replace(h.feats.index_add( - 0, _local_topology['query_rows'], _local_delta, - alpha=_local_scale, - )) - else: - # The module returns an exact-zero parameter edge for all 14 - # trainable tensors, so DDP observes finite numerical-zero grads. - h = h.replace(h.feats + _local_delta.sum() * _local_scale) - if self.innovation_gate is not None: - if _innovation_ratio is None: - _innovation_ratio = h.feats.new_zeros(()) - _innovation_gate_mean = h.feats.new_zeros(()) - _innovation_ratio_stats.append(_innovation_ratio.detach()) - _innovation_gate_mean_stats.append(_innovation_gate_mean.detach()) - if _local_attention is not None and _block_index < _local_attention.alpha.numel(): - assert _local_stats is not None - _local_ratio_base_stats.append(_local_ratio_base.detach()) - _local_ratio_hidden_stats.append(_local_ratio_hidden.detach()) - _local_null_stats.append(_local_stats['null_attention_mean'].detach()) - _local_entropy_stats.append(_local_stats['real_neighbour_entropy_mean'].detach()) - _local_effective_neighbour_stats.append( - _local_stats['effective_real_neighbour_count_mean'].detach() - ) - _local_observed_residual_stats.append( - _local_observed_residual.detach() - ) - _local_unobserved_residual_stats.append( - _local_unobserved_residual.detach() - ) - - if self.innovation_gate is not None: - # Instance-only telemetry: deliberately neither buffer nor state-dict entry. - self.partial_zobs_innovation_gate_last_stats = { - 'delta_to_base_norm_ratio': torch.stack(_innovation_ratio_stats), - 'gate_mean': torch.stack(_innovation_gate_mean_stats), - 'observed_rows': _innovation_observed_rows, - } - if _local_attention is not None: - # Block 29 is structurally absent: every telemetry vector has length 29. - assert _local_query_is_observed is not None - assert _local_neighbourhood_histogram is not None - _full_support_rows = int(x.coords.shape[0]) - _active_queries = int(_local_topology['query_rows'].numel()) - _empty_queries = _full_support_rows - _active_queries - _singleton_queries = int( - (_local_topology['neighbour_counts'] == 1).sum() - ) - self.partial_zobs_local_attention_last_stats = { - 'delta_to_base_norm_ratio': torch.stack(_local_ratio_base_stats), - 'delta_to_hidden_norm_ratio': torch.stack(_local_ratio_hidden_stats), - 'null_attention_mean': torch.stack(_local_null_stats), - 'real_neighbour_entropy_mean': torch.stack(_local_entropy_stats), - 'effective_real_neighbour_count_mean': torch.stack( - _local_effective_neighbour_stats - ), - 'observed_residual_l2_norm': torch.stack( - _local_observed_residual_stats - ), - 'unobserved_residual_l2_norm': torch.stack( - _local_unobserved_residual_stats - ), - 'observed_rows': _innovation_observed_rows, - 'full_support_rows': _full_support_rows, - 'active_queries': _active_queries, - 'empty_queries': _empty_queries, - 'singleton_queries': _singleton_queries, - 'observed_active_queries': int(_local_query_is_observed.sum()), - 'unobserved_active_queries': int((~_local_query_is_observed).sum()), - 'active_query_fraction': _active_queries / max(_full_support_rows, 1), - 'empty_query_fraction': _empty_queries / max(_full_support_rows, 1), - 'singleton_query_fraction': _singleton_queries / max(_active_queries, 1), - 'neighbourhood_size_histogram': _local_neighbourhood_histogram.detach(), - 'packed_pairs': int(_local_topology['source_indices'].numel()), - } - if _learned_attention is not None: - assert _learned_cache is not None - assert len(_learned_stats) == len(_learned_attention.sites) - self.partial_zobs_learned_attention_last_stats = { - name: torch.stack([site[name] for site in _learned_stats]) - for name in _learned_stats[0] - } - self.partial_zobs_learned_attention_last_stats.update({ - 'sites': _learned_attention.sites, - 'observed_rows': int(partial_zobs_rows.numel()), - 'full_support_rows': int(x.coords.shape[0]), - 'observed_counts': _learned_cache['observed_counts'], - }) - - h = manual_cast(h, x.dtype) - h = h.replace(F.layer_norm(h.feats, h.feats.shape[-1:])) - h = self.out_layer(h) - return h - - -class ElasticSLatFlowModel(SparseTransformerElasticMixin, SLatFlowModel): - """ - SLat Flow Model with elastic memory management. - Used for training with low VRAM. - """ - pass +from typing import * +from functools import partial +import torch +import torch.nn as nn +import torch.nn.functional as F +import numpy as np +from ..modules.utils import convert_module_to, manual_cast, str_to_dtype +from ..modules.transformer import AbsolutePositionEmbedder +from ..modules import sparse as sp +from ..modules.sparse.transformer import ModulatedSparseTransformerCrossBlock +from .sparse_structure_flow import TimestepEmbedder + + +class _PartialZObsInnovationGate(nn.Module): + """Residual, per-observed-token correction for the partial z_obs carrier.""" + + def __init__(self, model_channels: int, message_channels: int, dim: int, time_dim: int, num_blocks: int): + super().__init__() + self.alpha = nn.Parameter(torch.zeros(num_blocks)) + self.embed_y = nn.Linear(message_channels, dim, bias=True) + self.time_proj = nn.Linear(model_channels, time_dim, bias=True) + self.embed_x = nn.Linear(model_channels + time_dim, dim, bias=True) + self.mlp = nn.Sequential( + nn.Linear(dim + 2 + time_dim, dim, bias=True), + nn.SiLU(), + nn.Linear(dim, 1, bias=True), + ) + self.proj = nn.Linear(dim, model_channels, bias=True) + + +_LOCAL_ONE_HOP_OFFSETS = ( + (-1, 0, 0), (0, -1, 0), (0, 0, -1), + (0, 0, 1), (0, 1, 0), (1, 0, 0), +) +_LOCAL_TWO_HOP_OFFSETS = tuple(sorted( + (dx, dy, dz) + for dx in range(-2, 3) + for dy in range(-2, 3) + for dz in range(-2, 3) + if abs(dx) + abs(dy) + abs(dz) <= 2 +)) +assert len(_LOCAL_TWO_HOP_OFFSETS) == 25 + + +def _exact_coordinate_rows( + full_coords: torch.Tensor, + candidate_coords: torch.Tensor, +) -> Tuple[torch.Tensor, torch.Tensor]: + """Return exact full-support rows and a validity mask for candidate coords.""" + assert full_coords.ndim == candidate_coords.ndim == 2 + assert full_coords.shape[1] == candidate_coords.shape[1] + assert full_coords.shape[0] > 0, 'local attention requires nonempty full sparse support' + all_coords = torch.cat([full_coords, candidate_coords], dim=0) + _, inverse = torch.unique(all_coords, dim=0, return_inverse=True) + full_ids = inverse[:full_coords.shape[0]] + candidate_ids = inverse[full_coords.shape[0]:] + assert torch.unique(full_ids).numel() == full_ids.numel(), \ + 'x contains duplicate coordinates' + ordered_ids, ordered_rows = torch.sort(full_ids) + positions = torch.searchsorted(ordered_ids, candidate_ids) + in_range = positions < ordered_ids.numel() + safe_positions = positions.clamp(max=ordered_ids.numel() - 1) + valid = in_range & (ordered_ids[safe_positions] == candidate_ids) + rows = torch.where(valid, ordered_rows[safe_positions], -torch.ones_like(safe_positions)) + return rows, valid + + +def _build_two_hop_observed_topology( + full_coords: torch.Tensor, + observed_rows: torch.Tensor, +) -> Dict[str, torch.Tensor]: + """Pack exact batch-local G(S) distance<=2 observed source/query pairs.""" + assert full_coords.ndim == 2 and full_coords.shape[1] == 4, \ + f'local attention requires [N,4] integer coords, got {tuple(full_coords.shape)}' + assert not torch.is_floating_point(full_coords), 'local attention coordinates must be integer' + assert observed_rows.ndim == 1 and observed_rows.dtype == torch.long + # Validate full-coordinate uniqueness even for an all-absent batch. + _exact_coordinate_rows(full_coords, full_coords.new_empty((0, full_coords.shape[1]))) + if observed_rows.numel() == 0: + empty = torch.empty(0, dtype=torch.long, device=full_coords.device) + return { + 'query_rows': empty, + 'source_indices': empty, + 'offset_indices': empty, + 'segment_indices': empty, + 'neighbour_counts': empty, + } + assert observed_rows.min() >= 0 and observed_rows.max() < full_coords.shape[0], \ + 'observed row lies outside full sparse support' + assert torch.unique(observed_rows).numel() == observed_rows.numel(), \ + 'observed rows must be unique' + + device = full_coords.device + one_hop = torch.tensor( + _LOCAL_ONE_HOP_OFFSETS, dtype=full_coords.dtype, device=device, + ) + source_coords = full_coords[observed_rows] + candidate_hop1 = source_coords[:, None, :].expand(-1, len(_LOCAL_ONE_HOP_OFFSETS), -1).clone() + candidate_hop1[:, :, 1:] += one_hop[None] + hop1_rows, hop1_valid = _exact_coordinate_rows( + full_coords, candidate_hop1.reshape(-1, full_coords.shape[1]), + ) + source_ids_hop1 = torch.arange( + observed_rows.numel(), device=device, dtype=torch.long, + )[:, None].expand(-1, len(_LOCAL_ONE_HOP_OFFSETS)).reshape(-1) + hop1_rows = hop1_rows[hop1_valid] + source_ids_hop1 = source_ids_hop1[hop1_valid] + + if hop1_rows.numel() != 0: + candidate_hop2 = full_coords[hop1_rows][:, None, :].expand( + -1, len(_LOCAL_ONE_HOP_OFFSETS), -1, + ).clone() + candidate_hop2[:, :, 1:] += one_hop[None] + hop2_rows, hop2_valid = _exact_coordinate_rows( + full_coords, candidate_hop2.reshape(-1, full_coords.shape[1]), + ) + source_ids_hop2 = source_ids_hop1[:, None].expand( + -1, len(_LOCAL_ONE_HOP_OFFSETS), + ).reshape(-1) + hop2_rows = hop2_rows[hop2_valid] + source_ids_hop2 = source_ids_hop2[hop2_valid] + else: + hop2_rows = torch.empty(0, dtype=torch.long, device=device) + source_ids_hop2 = torch.empty(0, dtype=torch.long, device=device) + + source_ids = torch.cat([ + torch.arange(observed_rows.numel(), device=device, dtype=torch.long), + source_ids_hop1, + source_ids_hop2, + ]) + query_rows = torch.cat([observed_rows, hop1_rows, hop2_rows]) + # Multiple graph paths to the same query/source pair are one neighbourhood entry. + unique_pairs = torch.unique(torch.stack([query_rows, source_ids], dim=1), dim=0) + query_rows = unique_pairs[:, 0] + source_ids = unique_pairs[:, 1] + query_coords = full_coords[query_rows] + selected_source_coords = source_coords[source_ids] + assert torch.equal(query_coords[:, 0], selected_source_coords[:, 0]), \ + 'local attention constructed a cross-sample edge' + relative = query_coords[:, 1:] - selected_source_coords[:, 1:] + offset_table = torch.tensor( + _LOCAL_TWO_HOP_OFFSETS, dtype=full_coords.dtype, device=device, + ) + offset_matches = (relative[:, None, :] == offset_table[None, :, :]).all(dim=-1) + assert torch.equal(offset_matches.sum(dim=1), torch.ones_like(offset_matches[:, 0], dtype=torch.long)), \ + 'local attention produced an invalid two-hop offset' + offset_indices = offset_matches.to(dtype=torch.long).argmax(dim=1) + + # Canonical per-query real-KV order is lexicographic relative offset, independent + # of either full-support or observed-message row ordering. + order = torch.argsort(query_rows * len(_LOCAL_TWO_HOP_OFFSETS) + offset_indices) + query_rows = query_rows[order] + source_ids = source_ids[order] + offset_indices = offset_indices[order] + active_query_rows, segment_indices, counts = torch.unique_consecutive( + query_rows, return_inverse=True, return_counts=True, + ) + assert counts.max() <= len(_LOCAL_TWO_HOP_OFFSETS), \ + 'local observed neighbourhood exceeds frozen 25-source cap' + return { + 'query_rows': active_query_rows, + 'source_indices': source_ids, + 'offset_indices': offset_indices, + 'segment_indices': segment_indices, + 'neighbour_counts': counts, + } + + +class _PartialZObsLocalAttention(nn.Module): + """Shared packed two-hop observed-innovation attention for V12-A.""" + + def __init__( + self, + model_channels: int, + message_channels: int, + dim: int, + time_dim: int, + num_heads: int, + num_active_blocks: int, + ): + super().__init__() + assert dim == 128 and time_dim == 32 and num_heads == 8 + assert dim % num_heads == 0 + self.dim = dim + self.num_heads = num_heads + self.head_dim = dim // num_heads + self.alpha = nn.Parameter(torch.zeros(num_active_blocks)) + self.time_proj = nn.Linear(model_channels, time_dim, bias=True) + self.embed_x = nn.Linear(model_channels + time_dim, dim, bias=True) + self.embed_y = nn.Linear(message_channels, dim, bias=True) + self.key_proj = nn.Linear(dim, dim, bias=True) + self.value_proj = nn.Linear(dim, dim, bias=True) + self.relative_offset = nn.Parameter(torch.empty(len(_LOCAL_TWO_HOP_OFFSETS), dim)) + self.null_key = nn.Parameter(torch.empty(num_heads, self.head_dim)) + self.out_proj = nn.Linear(dim, model_channels, bias=False) + + @staticmethod + def _rms_norm(value: torch.Tensor) -> torch.Tensor: + return value * torch.rsqrt(value.square().mean(dim=-1, keepdim=True) + 1e-6) + + def forward( + self, + hidden: torch.Tensor, + full_batch: torch.Tensor, + t_base: torch.Tensor, + message: torch.Tensor, + observed_rows: torch.Tensor, + topology: Dict[str, torch.Tensor], + block_index: int, + ) -> Tuple[torch.Tensor, Dict[str, torch.Tensor]]: + assert 0 <= block_index < self.alpha.numel(), 'terminal block has no V12-A residual' + query_rows = topology['query_rows'] + source_indices = topology['source_indices'] + offset_indices = topology['offset_indices'] + segment_indices = topology['segment_indices'] + counts = topology['neighbour_counts'] + assert full_batch.ndim == 1 and full_batch.shape[0] == hidden.shape[0] + assert full_batch.numel() == 0 or ( + full_batch.min() >= 0 and full_batch.max() < t_base.shape[0] + ), 'full sparse batch coordinate is outside timestep batch' + assert message.shape == (observed_rows.numel(), 34) + + if query_rows.numel() == 0: + live_zero = sum(parameter.sum() for parameter in self.parameters()) * 0 + empty_delta = hidden.new_empty((0, hidden.shape[1])) + live_zero + zero = hidden.new_zeros(()) + return empty_delta, { + 'null_attention_mean': zero, + 'real_neighbour_entropy_mean': zero, + 'effective_real_neighbour_count_mean': zero, + 'active_queries': zero, + 'singleton_queries': zero, + 'packed_pairs': zero, + } + + branch_dtype = self.embed_x.weight.dtype + hidden_branch = hidden.to(dtype=branch_dtype) + time = self.time_proj(t_base.to(dtype=branch_dtype)) + x_embed = self.embed_x(torch.cat([ + F.layer_norm(hidden_branch, hidden_branch.shape[-1:], eps=1e-6), + time[full_batch], + ], dim=-1)) + innovation = self.embed_y(message.to(dtype=branch_dtype)) - x_embed[observed_rows] + query = self._rms_norm( + x_embed[query_rows].reshape(-1, self.num_heads, self.head_dim) + ) + real_key = self.key_proj(innovation)[source_indices] + real_key = real_key + self.relative_offset[offset_indices] + real_key = self._rms_norm(real_key.reshape(-1, self.num_heads, self.head_dim)) + real_value = self.value_proj(innovation)[source_indices].reshape( + -1, self.num_heads, self.head_dim, + ) + real_logits = (query[segment_indices] * real_key).sum(dim=-1) / np.sqrt(self.head_dim) + null_key = self._rms_norm(self.null_key) + null_logits = (query * null_key[None]).sum(dim=-1) / np.sqrt(self.head_dim) + + segment_heads = segment_indices[:, None].expand(-1, self.num_heads) + maxima = null_logits.clone() + maxima.scatter_reduce_(0, segment_heads, real_logits, reduce='amax', include_self=True) + real_exp = torch.exp(real_logits - maxima[segment_indices]) + null_exp = torch.exp(null_logits - maxima) + real_sum = torch.zeros_like(null_exp).scatter_add(0, segment_heads, real_exp) + denominator = null_exp + real_sum + real_weight = real_exp / denominator[segment_indices] + null_weight = null_exp / denominator + attended = hidden_branch.new_zeros( + (query_rows.numel(), self.num_heads, self.head_dim), + ) + attended.index_add_( + 0, segment_indices, real_weight[..., None] * real_value, + ) + delta = self.out_proj(attended.reshape(query_rows.numel(), self.dim)) + delta = delta * self.alpha[block_index].to(dtype=delta.dtype) + + # Entropy is conditional on selecting real evidence. Null-vs-real + # acceptance is reported independently as null_attention_mean; mixing + # that mass into this distribution would make singleton entropy nonzero. + real_conditional_weight = real_exp / real_sum[segment_indices].clamp_min( + torch.finfo(real_sum.dtype).tiny + ) + entropy_terms = -real_conditional_weight * torch.log( + real_conditional_weight.clamp_min( + torch.finfo(real_conditional_weight.dtype).tiny + ) + ) + real_entropy = torch.zeros_like(null_weight).scatter_add( + 0, segment_heads, entropy_terms, + ) + return delta, { + 'null_attention_mean': null_weight.mean(), + 'real_neighbour_entropy_mean': real_entropy.mean(), + 'effective_real_neighbour_count_mean': real_entropy.exp().mean(), + 'active_queries': delta.new_tensor(float(query_rows.numel())), + 'singleton_queries': delta.new_tensor(float((counts == 1).sum())), + 'packed_pairs': delta.new_tensor(float(source_indices.numel())), + } + + +class _V12LACrossRope(nn.Module): + """Apply the same 3D rotary basis to independently packed Q and K rows.""" + + def __init__( + self, + head_dim: int, + rope_freq: Tuple[float, float], + reserved_dims: int = 2, + ): + super().__init__() + content_dim = head_dim - reserved_dims + assert content_dim > 0 and content_dim % 6 == 0 + self.head_dim = head_dim + self.content_dim = content_dim + freq_dim = content_dim // 2 // 3 + frequencies = torch.arange(freq_dim, dtype=torch.float32) / freq_dim + frequencies = rope_freq[0] / (rope_freq[1] ** frequencies) + self.register_buffer('frequencies', frequencies, persistent=False) + + def forward(self, value: torch.Tensor, coords: torch.Tensor) -> torch.Tensor: + assert value.ndim == 3 and value.shape[-1] == self.head_dim + assert coords.shape == (value.shape[0], 3) + if value.shape[0] == 0: + return value + phases = torch.outer( + coords.reshape(-1).float(), self.frequencies, + ).reshape(coords.shape[0], self.content_dim // 2) + phases = torch.polar(torch.ones_like(phases), phases).unsqueeze(1) + content = torch.view_as_complex( + value[..., :self.content_dim].float().reshape( + value.shape[0], value.shape[1], -1, 2, + ) + ) + rotated = torch.view_as_real(content * phases).reshape( + value.shape[0], value.shape[1], self.content_dim, + ).to(dtype=value.dtype) + return torch.cat([rotated, value[..., self.content_dim:]], dim=-1) + + +class _PartialZObsLearnedAttention(nn.Module): + """Five-site V12LA observation attention with an exact carrier anchor.""" + + def __init__( + self, + model_channels: int, + num_heads: int, + sites: Sequence[int], + rope_freq: Tuple[float, float], + null_prior: float = 0.5, + ): + super().__init__() + self.sites = tuple(int(site) for site in sites) + assert self.sites == (5, 10, 15, 20, 25), self.sites + assert num_heads == 12 + assert model_channels % num_heads == 0 + assert (model_channels // num_heads - 2) > 0 + assert (model_channels // num_heads - 2) % 6 == 0 + self.model_channels = model_channels + self.num_heads = num_heads + self.head_dim = model_channels // num_heads + self.content_dim = self.head_dim - 2 + self.site_to_slot = {site: slot for slot, site in enumerate(self.sites)} + self.null_prior = float(null_prior) + assert 0.0 < self.null_prior < 1.0, \ + f'learned-attention null prior must lie in (0,1), got {self.null_prior}' + self.null_logit_offset = float( + np.log(self.null_prior) - np.log1p(-self.null_prior) + ) + + # K/V are source-only and shared across all five sites. Q and the + # zero-start deviation/output matrix are site-specific. + self.key_proj = nn.Linear(model_channels, model_channels, bias=False) + self.value_proj = nn.Linear(model_channels, model_channels, bias=False) + self.query_proj = nn.ModuleList([ + nn.Linear(model_channels, model_channels, bias=False) + for _ in self.sites + ]) + self.delta_proj = nn.ModuleList([ + nn.Linear(model_channels, model_channels, bias=False) + for _ in self.sites + ]) + self.null_key = nn.Parameter(torch.empty(num_heads, self.content_dim)) + self.rope = _V12LACrossRope(self.head_dim, rope_freq, reserved_dims=2) + self.reset_parameters() + + def reset_parameters(self) -> None: + nn.init.xavier_uniform_(self.key_proj.weight) + # Attention mixes the already trained V11 carrier basis at warm start; + # the value map may then adapt, but does not begin as a random rewrite. + nn.init.eye_(self.value_proj.weight) + nn.init.xavier_uniform_(self.null_key) + for query_proj in self.query_proj: + nn.init.xavier_uniform_(query_proj.weight) + for delta_proj in self.delta_proj: + nn.init.zeros_(delta_proj.weight) + + @staticmethod + def _rms_norm(value: torch.Tensor) -> torch.Tensor: + normalized = value.float() * torch.rsqrt( + value.float().square().mean(dim=-1, keepdim=True) + 1e-6 + ) + return normalized.to(dtype=value.dtype) + + @staticmethod + def _anchor( + carrier: torch.Tensor, + observed_rows: torch.Tensor, + full_rows: int, + ) -> torch.Tensor: + """Evaluate fixed coordinate/null attention without an additive scatter.""" + assert carrier.ndim == 2 + assert observed_rows.ndim == 1 and observed_rows.dtype == torch.long + assert carrier.shape[0] == observed_rows.numel() + assert full_rows >= observed_rows.numel() + if observed_rows.numel() != 0: + assert observed_rows.min() >= 0 and observed_rows.max() < full_rows + assert torch.unique(observed_rows).numel() == observed_rows.numel() + # Index zero is the exact-zero null value. Every observed row is then + # redirected to its own carrier row; unmatched rows keep selecting null. + values = torch.cat([carrier.new_zeros((1, carrier.shape[1])), carrier], dim=0) + anchor_index = torch.zeros(full_rows, dtype=torch.long, device=carrier.device) + anchor_index[observed_rows] = torch.arange( + 1, observed_rows.numel() + 1, dtype=torch.long, device=carrier.device, + ) + return values[anchor_index] + + def build_cache( + self, + full_coords: torch.Tensor, + full_layout: List[slice], + carrier: torch.Tensor, + observed_rows: torch.Tensor, + ) -> Dict[str, Any]: + """Build the exact anchor and the shared per-sample O+null K/V cache.""" + assert full_coords.ndim == 2 and full_coords.shape[1] == 4 + assert carrier.shape == (observed_rows.numel(), self.model_channels) + assert len(full_layout) > 0 + anchor = self._anchor(carrier, observed_rows, full_coords.shape[0]) + observed_coords = full_coords[observed_rows] + observed_batch = observed_coords[:, 0].long() + assert observed_batch.numel() == 0 or ( + observed_batch.min() >= 0 and observed_batch.max() < len(full_layout) + ) + + branch_dtype = self.key_proj.weight.dtype + carrier_branch = carrier.to(dtype=branch_dtype) + carrier_norm = F.layer_norm( + carrier_branch, (self.model_channels,), eps=1e-6, + ) + real_key = self.key_proj(carrier_norm).reshape( + -1, self.num_heads, self.head_dim, + ) + real_key = self.rope(real_key, observed_coords[:, 1:]) + real_key = self._rms_norm(real_key[..., :self.content_dim]) + real_key = torch.cat([ + real_key, + real_key.new_zeros((real_key.shape[0], self.num_heads, 2)), + ], dim=-1) + real_value = self.value_proj(carrier_branch).reshape( + -1, self.num_heads, self.head_dim, + ) + null_content = self._rms_norm( + self.null_key.to(dtype=real_key.dtype) + ) + + packed_kv = [] + observed_counts = [] + for batch_index in range(len(full_layout)): + source_indices = torch.nonzero( + observed_batch == batch_index, as_tuple=False, + ).flatten() + observed_count = int(source_indices.numel()) + observed_counts.append(observed_count) + # One null competes with O real keys. log(O) cancels its cardinality + # disadvantage; the constant logit offset sets the null mass only in + # the equal-content-logit reference case, not after learned Q/K drift. + null_reserved = real_key.new_zeros((self.num_heads, 2)) + null_reserved[:, 0] = np.sqrt(self.head_dim) * ( + np.log(max(observed_count, 1)) + self.null_logit_offset + ) + null_key = torch.cat([null_content, null_reserved], dim=-1).unsqueeze(0) + null_value = real_value.new_zeros((1, self.num_heads, self.head_dim)) + keys = torch.cat([real_key[source_indices], null_key], dim=0) + values = torch.cat([real_value[source_indices], null_value], dim=0) + packed_kv.append(torch.stack([keys, values], dim=1)) + + kv = sp.VarLenTensor.from_tensor_list(packed_kv) + return { + 'anchor': anchor, + 'kv': kv, + 'observed_rows': observed_rows, + 'observed_counts': tuple(observed_counts), + } + + def forward( + self, + hidden: torch.Tensor, + full_coords: torch.Tensor, + full_layout: List[slice], + cache: Dict[str, Any], + block_index: int, + ) -> Tuple[torch.Tensor, Dict[str, torch.Tensor]]: + assert block_index in self.site_to_slot, block_index + assert hidden.shape == (full_coords.shape[0], self.model_channels) + slot = self.site_to_slot[block_index] + branch_dtype = self.query_proj[slot].weight.dtype + hidden_branch = hidden.to(dtype=branch_dtype) + query = self.query_proj[slot](F.layer_norm( + hidden_branch, (self.model_channels,), eps=1e-6, + )).reshape(-1, self.num_heads, self.head_dim) + query = self.rope(query, full_coords[:, 1:]) + query_content = self._rms_norm(query[..., :self.content_dim]) + query_reserved = query.new_zeros((query.shape[0], self.num_heads, 2)) + query_reserved[..., 0] = 1 + query = torch.cat([query_content, query_reserved], dim=-1) + query_varlen = sp.VarLenTensor(query, full_layout) + attended = sp.sparse_scaled_dot_product_attention( + query_varlen, cache['kv'], + ).feats.reshape(-1, self.model_channels) + + anchor = cache['anchor'].to(dtype=attended.dtype) + deviation_input = attended - anchor + correction = self.delta_proj[slot](deviation_input) + inject = anchor + correction + eps = torch.finfo(inject.dtype).eps + observed_rows = cache['observed_rows'] + observed_mask = torch.zeros( + inject.shape[0], dtype=torch.bool, device=inject.device, + ) + observed_mask[observed_rows] = True + return inject, { + 'correction_to_anchor_norm_ratio': correction.norm() / anchor.norm().clamp_min(eps), + 'correction_to_hidden_norm_ratio': correction.norm() / hidden_branch.norm().clamp_min(eps), + 'observed_correction_l2_norm': correction[observed_mask].float().norm(), + 'unobserved_correction_l2_norm': correction[~observed_mask].float().norm(), + 'attention_output_l2_norm': attended.float().norm(), + 'anchor_l2_norm': anchor.float().norm(), + 'delta_weight_l2_norm': self.delta_proj[slot].weight.float().norm(), + } + +class SLatFlowModel(nn.Module): + """ + Structured Latent Flow Model for 3D generation. + + Supports two conditioning modes: + - "cross": Standard cross-attention with image features + - "proj": View-aligned projection attention with camera-aware features + """ + def __init__( + self, + resolution: int, + in_channels: int, + model_channels: int, + cond_channels: int, + out_channels: int, + num_blocks: int, + num_heads: Optional[int] = None, + num_head_channels: Optional[int] = 64, + mlp_ratio: float = 4, + pe_mode: Literal["ape", "rope"] = "ape", + rope_freq: Tuple[float, float] = (1.0, 10000.0), + dtype: str = 'float32', + use_checkpoint: bool = False, + share_mod: bool = False, + initialization: str = 'vanilla', + qk_rms_norm: bool = False, + qk_rms_norm_cross: bool = False, + image_attn_mode: Literal["cross", "proj", "gated_proj"] = "cross", + proj_in_channels: Optional[int] = None, + vae_in_channels: Optional[int] = None, + zobs_inject: bool = False, + zobs_inject_in_channels: int = 33, + zobs_inject_scale: float = 1.0, + partial_zobs_inject: bool = False, + partial_zobs_in_channels: int = 34, + partial_zobs_inject_position: Literal['post_block', 'pre_mlp'] = 'post_block', + legacy_zobs_zero_input: bool = False, + partial_zobs_innovation_gate: bool = False, + innovation_gate_dim: int = 128, + innovation_gate_time_dim: int = 32, + partial_zobs_innovation_gate_train_only: bool = False, + partial_zobs_local_attention: bool = False, + local_attention_dim: int = 128, + local_attention_time_dim: int = 32, + local_attention_heads: int = 8, + partial_zobs_local_attention_train_only: bool = False, + partial_zobs_learned_attention: bool = False, + learned_attention_sites: Sequence[int] = (5, 10, 15, 20, 25), + learned_attention_heads: int = 12, + learned_attention_null_prior: float = 0.5, + partial_zobs_learned_attention_train_only: bool = False, + **kwargs + ): + super().__init__() + self.resolution = resolution + self.in_channels = in_channels + self.model_channels = model_channels + self.cond_channels = cond_channels + self.out_channels = out_channels + self.num_blocks = num_blocks + self.num_heads = num_heads or model_channels // num_head_channels + self.mlp_ratio = mlp_ratio + self.pe_mode = pe_mode + self.use_checkpoint = use_checkpoint + self.share_mod = share_mod + self.initialization = initialization + self.qk_rms_norm = qk_rms_norm + self.qk_rms_norm_cross = qk_rms_norm_cross + self.image_attn_mode = image_attn_mode + self.proj_in_channels = proj_in_channels + self.vae_in_channels = vae_in_channels + self.dtype = str_to_dtype(dtype) + + self.t_embedder = TimestepEmbedder(model_channels) + if share_mod: + self.adaLN_modulation = nn.Sequential( + nn.SiLU(), + nn.Linear(model_channels, 6 * model_channels, bias=True) + ) + + if pe_mode == "ape": + self.pos_embedder = AbsolutePositionEmbedder(model_channels) + + self.input_layer = sp.SparseLinear(in_channels, model_channels) + # Trained-bridge: single shared per-block injector of [z_obs (x) vis (32) | vis (1)] -> model_channels. + self.zobs_inject = zobs_inject + self.zobs_inject_scale = float(zobs_inject_scale) + self.zobs_injector = sp.SparseLinear(zobs_inject_in_channels, model_channels) if zobs_inject else None + # v8-partial: the observation lives only on its observed support. This adapter is + # deliberately bias-free: absent rows never acquire a learned content/fill value. + self.partial_zobs_inject = partial_zobs_inject + self.partial_zobs_in_channels = partial_zobs_in_channels + assert partial_zobs_inject_position in ('post_block', 'pre_mlp'), \ + f'partial_zobs_inject_position must be post_block or pre_mlp, got {partial_zobs_inject_position!r}' + self.partial_zobs_inject_position = partial_zobs_inject_position + if partial_zobs_inject_position == 'pre_mlp': + assert not any(( + partial_zobs_innovation_gate, + partial_zobs_local_attention, + partial_zobs_learned_attention, + )), 'partial_zobs_inject_position=pre_mlp is incompatible with post-block experimental branches' + self.legacy_zobs_zero_input = legacy_zobs_zero_input + self.partial_zobs_injector = ( + sp.SparseLinear(partial_zobs_in_channels, model_channels, bias=False) + if partial_zobs_inject else None + ) + self.partial_zobs_innovation_gate = partial_zobs_innovation_gate + self.partial_zobs_innovation_gate_train_only = partial_zobs_innovation_gate_train_only + if partial_zobs_innovation_gate_train_only: + assert partial_zobs_innovation_gate, \ + 'partial_zobs_innovation_gate_train_only requires partial_zobs_innovation_gate=true' + if partial_zobs_innovation_gate: + assert partial_zobs_inject, 'partial_zobs_innovation_gate requires partial_zobs_inject=true' + assert partial_zobs_in_channels == 34, \ + f'partial_zobs_innovation_gate requires 34 input channels, got {partial_zobs_in_channels}' + assert innovation_gate_dim == 128, f'innovation_gate_dim is fixed at 128, got {innovation_gate_dim}' + assert innovation_gate_time_dim == 32, \ + f'innovation_gate_time_dim is fixed at 32, got {innovation_gate_time_dim}' + assert model_channels == 1536, f'partial_zobs_innovation_gate requires 1536 model channels, got {model_channels}' + assert num_blocks == 30, f'partial_zobs_innovation_gate requires 30 blocks, got {num_blocks}' + self.innovation_gate = _PartialZObsInnovationGate( + model_channels, partial_zobs_in_channels, innovation_gate_dim, + innovation_gate_time_dim, num_blocks, + ) + else: + self.innovation_gate = None + self.partial_zobs_local_attention = partial_zobs_local_attention + self.partial_zobs_local_attention_train_only = partial_zobs_local_attention_train_only + if partial_zobs_local_attention_train_only: + assert partial_zobs_local_attention, \ + 'partial_zobs_local_attention_train_only requires partial_zobs_local_attention=true' + assert sum(bool(value) for value in ( + partial_zobs_innovation_gate, + partial_zobs_local_attention, + partial_zobs_learned_attention, + )) <= 1, 'partial-zobs experimental attention/gate branches are mutually exclusive' + if partial_zobs_local_attention: + assert partial_zobs_inject, 'partial_zobs_local_attention requires partial_zobs_inject=true' + assert partial_zobs_in_channels == 34, \ + f'partial_zobs_local_attention requires 34 input channels, got {partial_zobs_in_channels}' + assert local_attention_dim == 128, \ + f'local_attention_dim is fixed at 128, got {local_attention_dim}' + assert local_attention_time_dim == 32, \ + f'local_attention_time_dim is fixed at 32, got {local_attention_time_dim}' + assert local_attention_heads == 8, \ + f'local_attention_heads is fixed at 8, got {local_attention_heads}' + assert model_channels == 1536, \ + f'partial_zobs_local_attention requires 1536 model channels, got {model_channels}' + assert num_blocks == 30, \ + f'partial_zobs_local_attention requires 30 blocks, got {num_blocks}' + self.local_obs_attention = _PartialZObsLocalAttention( + model_channels, partial_zobs_in_channels, local_attention_dim, + local_attention_time_dim, local_attention_heads, num_blocks - 1, + ) + else: + self.local_obs_attention = None + + self.partial_zobs_learned_attention = partial_zobs_learned_attention + self.partial_zobs_learned_attention_train_only = \ + partial_zobs_learned_attention_train_only + if partial_zobs_learned_attention_train_only: + assert partial_zobs_learned_attention, \ + 'partial_zobs_learned_attention_train_only requires partial_zobs_learned_attention=true' + if partial_zobs_learned_attention: + assert partial_zobs_inject, \ + 'partial_zobs_learned_attention requires partial_zobs_inject=true' + assert partial_zobs_in_channels == 34, \ + f'partial_zobs_learned_attention requires 34 input channels, got {partial_zobs_in_channels}' + assert model_channels == 1536, \ + f'partial_zobs_learned_attention requires 1536 model channels, got {model_channels}' + assert num_blocks == 30, \ + f'partial_zobs_learned_attention requires 30 blocks, got {num_blocks}' + assert learned_attention_heads == 12, \ + f'partial_zobs_learned_attention requires 12 heads, got {learned_attention_heads}' + assert tuple(learned_attention_sites) == (5, 10, 15, 20, 25), \ + f'V12LA sites are fixed at [5,10,15,20,25], got {learned_attention_sites}' + self.learned_attention = _PartialZObsLearnedAttention( + model_channels, + learned_attention_heads, + learned_attention_sites, + rope_freq, + learned_attention_null_prior, + ) + else: + self.learned_attention = None + + self.blocks = nn.ModuleList([ + ModulatedSparseTransformerCrossBlock( + model_channels, + cond_channels, + num_heads=self.num_heads, + mlp_ratio=self.mlp_ratio, + attn_mode='full', + use_checkpoint=self.use_checkpoint, + use_rope=(pe_mode == "rope"), + rope_freq=rope_freq, + share_mod=self.share_mod, + qk_rms_norm=self.qk_rms_norm, + qk_rms_norm_cross=self.qk_rms_norm_cross, + image_attn_mode=image_attn_mode, + proj_in_channels=proj_in_channels, + vae_in_channels=vae_in_channels, + ) + for _ in range(num_blocks) + ]) + + self.out_layer = sp.SparseLinear(model_channels, out_channels) + + self.initialize_weights() + self.convert_to(self.dtype) + if self.partial_zobs_innovation_gate_train_only: + self._freeze_for_partial_zobs_innovation_gate() + if self.partial_zobs_local_attention_train_only: + self._freeze_for_partial_zobs_local_attention() + if self.partial_zobs_learned_attention_train_only: + self._freeze_for_partial_zobs_learned_attention() + + @property + def device(self) -> torch.device: + """ + Return the device of the model. + """ + return next(self.parameters()).device + + def convert_to(self, dtype: torch.dtype) -> None: + """ + Convert the torso of the model to the specified dtype. + """ + self.dtype = dtype + self.blocks.apply(partial(convert_module_to, dtype=dtype)) + if self.learned_attention is not None: + self.learned_attention.apply(partial(convert_module_to, dtype=dtype)) + + def initialize_weights(self) -> None: + if self.initialization == 'vanilla': + # Initialize transformer layers: + def _basic_init(module): + if isinstance(module, nn.Linear): + torch.nn.init.xavier_uniform_(module.weight) + if module.bias is not None: + nn.init.constant_(module.bias, 0) + self.apply(_basic_init) + + # Initialize timestep embedding MLP: + nn.init.normal_(self.t_embedder.mlp[0].weight, std=0.02) + nn.init.normal_(self.t_embedder.mlp[2].weight, std=0.02) + + # Zero-out adaLN modulation layers in DiT blocks: + if self.share_mod: + nn.init.constant_(self.adaLN_modulation[-1].weight, 0) + nn.init.constant_(self.adaLN_modulation[-1].bias, 0) + else: + for block in self.blocks: + nn.init.constant_(block.adaLN_modulation[-1].weight, 0) + nn.init.constant_(block.adaLN_modulation[-1].bias, 0) + + # Zero-out output layers: + nn.init.constant_(self.out_layer.weight, 0) + nn.init.constant_(self.out_layer.bias, 0) + + elif self.initialization == 'scaled': + # Initialize transformer layers: + def _basic_init(module): + if isinstance(module, nn.Linear): + torch.nn.init.normal_(module.weight, std=np.sqrt(2.0 / (5.0 * self.model_channels))) + if module.bias is not None: + nn.init.constant_(module.bias, 0) + self.apply(_basic_init) + + # Scaled init for to_out and ffn2 + def _scaled_init(module): + if isinstance(module, nn.Linear): + torch.nn.init.normal_(module.weight, std=1.0 / np.sqrt(5 * self.num_blocks * self.model_channels)) + if module.bias is not None: + nn.init.constant_(module.bias, 0) + for block in self.blocks: + block.self_attn.to_out.apply(_scaled_init) + # Handle cross, proj, and gated_proj modes + if self.image_attn_mode in ("proj", "gated_proj"): + block.cross_attn.cross_attn_block.to_out.apply(_scaled_init) + else: + block.cross_attn.to_out.apply(_scaled_init) + block.mlp.mlp[2].apply(_scaled_init) + + # Initialize input layer to make the initial representation have variance 1 + nn.init.normal_(self.input_layer.weight, std=1.0 / np.sqrt(self.in_channels)) + nn.init.zeros_(self.input_layer.bias) + + # Initialize timestep embedding MLP: + nn.init.normal_(self.t_embedder.mlp[0].weight, std=0.02) + nn.init.normal_(self.t_embedder.mlp[2].weight, std=0.02) + + # Zero-out adaLN modulation layers in DiT blocks: + if self.share_mod: + nn.init.constant_(self.adaLN_modulation[-1].weight, 0) + nn.init.constant_(self.adaLN_modulation[-1].bias, 0) + else: + for block in self.blocks: + nn.init.constant_(block.adaLN_modulation[-1].weight, 0) + nn.init.constant_(block.adaLN_modulation[-1].bias, 0) + + # Zero-out output layers: + nn.init.constant_(self.out_layer.weight, 0) + nn.init.constant_(self.out_layer.bias, 0) + + # Trained-bridge: zero-init the z_obs injector LAST (after _basic_init), so step-0 injection == 0. + if getattr(self, 'zobs_injector', None) is not None: + nn.init.zeros_(self.zobs_injector.weight) + nn.init.zeros_(self.zobs_injector.bias) + if getattr(self, 'partial_zobs_injector', None) is not None: + nn.init.zeros_(self.partial_zobs_injector.weight) + if getattr(self, 'innovation_gate', None) is not None: + nn.init.zeros_(self.innovation_gate.alpha) + if getattr(self, 'local_obs_attention', None) is not None: + # V12-A freezes vanilla Xavier initialization independently of the + # scaled V11 backbone initialization mode. Only alpha starts at zero. + for module in self.local_obs_attention.modules(): + if isinstance(module, nn.Linear): + nn.init.xavier_uniform_(module.weight) + if module.bias is not None: + nn.init.zeros_(module.bias) + nn.init.xavier_uniform_(self.local_obs_attention.relative_offset) + nn.init.xavier_uniform_(self.local_obs_attention.null_key) + nn.init.zeros_(self.local_obs_attention.alpha) + if getattr(self, 'learned_attention', None) is not None: + self.learned_attention.reset_parameters() + + def _freeze_for_partial_zobs_innovation_gate(self) -> None: + """Make the V12 gate the sole trainable branch before an optimizer is built.""" + assert self.innovation_gate is not None + for parameter in self.parameters(): + parameter.requires_grad_(False) + for parameter in self.innovation_gate.parameters(): + parameter.requires_grad_(True) + trainable = [name for name, parameter in self.named_parameters() if parameter.requires_grad] + expected = [f'innovation_gate.{name}' for name in self.innovation_gate.state_dict()] + assert trainable == expected, (trainable, expected) + assert sum(parameter.numel() for parameter in self.parameters() if parameter.requires_grad) == 473663 + + def _freeze_for_partial_zobs_local_attention(self) -> None: + """Make the V12-A local attention the sole trainable warm-start branch.""" + assert self.local_obs_attention is not None + for parameter in self.parameters(): + parameter.requires_grad_(False) + for parameter in self.local_obs_attention.parameters(): + parameter.requires_grad_(True) + trainable = [name for name, parameter in self.named_parameters() if parameter.requires_grad] + expected = [f'local_obs_attention.{name}' for name in self.local_obs_attention.state_dict()] + assert trainable == expected, (trainable, expected) + assert len(trainable) == 14 + assert sum(parameter.numel() for parameter in self.parameters() if parameter.requires_grad) == 487485 + + def _freeze_for_partial_zobs_learned_attention(self) -> None: + """Make the V12LA module the sole trainable warm-start branch.""" + assert self.learned_attention is not None + for parameter in self.parameters(): + parameter.requires_grad_(False) + for parameter in self.learned_attention.parameters(): + parameter.requires_grad_(True) + trainable = [name for name, parameter in self.named_parameters() if parameter.requires_grad] + expected = [ + f'learned_attention.{name}' + for name, _ in self.learned_attention.named_parameters() + ] + assert trainable == expected, (trainable, expected) + assert len(trainable) == 13, trainable + + def _partial_zobs_message( + self, + x: sp.SparseTensor, + partial_z_obs: sp.SparseTensor, + partial_z_obs_keep_mask: Optional[torch.Tensor], + ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """Project a variable-support observation and join its messages onto ``x`` rows. + + The return value stays variable-length: projected messages plus their exact full-row + indices. Neither observation content nor a zero-filled model-channel tensor is + materialized on full support. ``torch.unique(..., dim=0)`` makes the coordinate join + exact and rejects duplicate or non-subset partial coordinates instead of assuming a + row order. + """ + adapter = self.partial_zobs_injector + assert adapter is not None + assert partial_z_obs.feats.ndim == 2 and \ + partial_z_obs.feats.shape[-1] == self.partial_zobs_in_channels, \ + f'partial_z_obs must be [N,{self.partial_zobs_in_channels}], got {tuple(partial_z_obs.feats.shape)}' + assert partial_z_obs.coords.ndim == 2 and partial_z_obs.coords.shape[1] == x.coords.shape[1], \ + f'partial_z_obs coordinate rank mismatch: {tuple(partial_z_obs.coords.shape)} vs {tuple(x.coords.shape)}' + + active = torch.ones(partial_z_obs.feats.shape[0], dtype=torch.bool, device=partial_z_obs.feats.device) + if partial_z_obs_keep_mask is not None: + assert partial_z_obs_keep_mask.ndim == 1 and partial_z_obs_keep_mask.shape[0] == x.shape[0], \ + f'partial_z_obs_keep_mask must be [B={x.shape[0]}], got {tuple(partial_z_obs_keep_mask.shape)}' + keep = partial_z_obs_keep_mask.to(device=partial_z_obs.feats.device, dtype=torch.bool) + active = keep[partial_z_obs.coords[:, 0].long()] + + active_coords = partial_z_obs.coords[active] + active_feats = partial_z_obs.feats[active] + # F.linear on a zero-row tensor is intentional: it keeps the adapter parameters in the + # graph on all-drop batches, so DDP sees a zero (rather than None) gradient. + active_messages = F.linear(active_feats, adapter.weight, None) + + full_rows = torch.empty(0, dtype=torch.long, device=x.feats.device) + if active_coords.shape[0] != 0: + all_coords = torch.cat([x.coords, active_coords], dim=0) + _, inverse = torch.unique(all_coords, dim=0, return_inverse=True) + x_ids = inverse[:x.coords.shape[0]] + partial_ids = inverse[x.coords.shape[0]:] + assert torch.unique(x_ids).numel() == x_ids.numel(), 'x contains duplicate coordinates' + assert torch.unique(partial_ids).numel() == partial_ids.numel(), 'partial_z_obs contains duplicate coordinates' + ordered_x_ids, ordered_x_rows = torch.sort(x_ids) + positions = torch.searchsorted(ordered_x_ids, partial_ids) + assert (positions < ordered_x_ids.numel()).all() and \ + torch.equal(ordered_x_ids[positions], partial_ids), \ + 'partial_z_obs coordinates must be an exact subset of current x coordinates' + full_rows = ordered_x_rows[positions] + + return active_messages, full_rows, active_feats, active_coords[:, 0].long() + + def forward( + self, + x: sp.SparseTensor, + t: torch.Tensor, + cond: Union[torch.Tensor, List[torch.Tensor], Dict[str, Union[torch.Tensor, sp.SparseTensor]], Tuple], + concat_cond: Optional[sp.SparseTensor] = None, + z_obs: Optional[sp.SparseTensor] = None, + partial_z_obs: Optional[sp.SparseTensor] = None, + partial_z_obs_keep_mask: Optional[torch.Tensor] = None, + **kwargs + ) -> sp.SparseTensor: + """ + Forward pass. + + Args: + x: SparseTensor input + t: Timestep tensor [B] + cond: Conditioning tensor. For "cross" mode: list of tensors or tensor. + For "proj" mode: dict {'global': global_cond, 'proj': proj_cond} + or tuple of (global_cond, proj_cond) + concat_cond: Optional concatenation condition + z_obs: Optional 33ch [z_obs(32)|vis(1)] sparse tensor, same coords as x. Feeds ONLY + the zobs_injector (v6 patchcross path: concat_cond is the plain 32ch shape + latent and z_obs never enters the input_layer). When None, the legacy 97ch + concat slice below applies — shipped/v4 behavior is bit-identical. + partial_z_obs: Optional variable-length v8 observation + [normalized_z_obs(32)|coverage(1)|boundary_distance(1)]. + Its coordinates must be a subset of x; it is consumed only by the dedicated, + bias-free partial adapter and is never converted to full-support content. + + Returns: + SparseTensor output + """ + if concat_cond is not None: + x = sp.sparse_cat([x, concat_cond], dim=-1) + + h = self.input_layer(x) + h = manual_cast(h, self.dtype) + + # Trained-bridge: compute the per-block injection vector ONCE (shared zero-init linear, added each block). + # Sources, in priority order: separate z_obs kwarg (v6), else the trailing channels of a + # 97-ch [x_t|shape|z_obs|vis] input (v4). z_obs*vis==0 at vis=0 in both. + zobs_inject = None + _injector = getattr(self, 'zobs_injector', None) + if z_obs is not None: + # Fail fast: without an injector the kwarg would be a silent no-op, and in the 64ch + # patchcross architecture the injector is the ONLY z_obs consumption path. + assert _injector is not None, \ + 'z_obs kwarg passed but the denoiser was built without zobs_inject=true' + # Full coords equality is asserted at data-load; a row-count check catches gross mismatch. + assert z_obs.feats.shape[-1] == 33 and z_obs.feats.shape[0] == x.feats.shape[0], \ + f'z_obs kwarg mismatch: {tuple(z_obs.feats.shape)} vs x rows {x.feats.shape[0]}' + _zo = z_obs.feats[:, :32] + _vis = z_obs.feats[:, 32:33] + zobs_inject = manual_cast(self.zobs_injector(z_obs.replace(torch.cat([_zo * _vis, _vis], dim=-1))), self.dtype) + if not getattr(self, '_zobs_kwarg_logged', False): # in-log proof hook (once per process) + print('[zobs_injector] consuming separate z_obs kwarg (patchcross path)', flush=True) + self._zobs_kwarg_logged = True + elif _injector is not None and x.feats.shape[-1] >= 97: + _zo = x.feats[:, 64:96] + _vis = x.feats[:, 96:97] + zobs_inject = manual_cast(self.zobs_injector(x.replace(torch.cat([_zo * _vis, _vis], dim=-1))), self.dtype) + elif _injector is not None and self.legacy_zobs_zero_input: + # 64ch v8 shell with a legacy injector: preserve the legacy all-zero z_obs mode, + # including its learned per-block bias, without materialising observation content. + # Evaluate one zero row and broadcast the result: this is exactly W*0+b, keeps a + # zero (non-None) weight gradient under DDP, and avoids an N-row 33ch GEMM. + zero_output = F.linear( + x.feats.new_zeros((1, _injector.in_features)), + _injector.weight, + _injector.bias, + ) + zobs_inject = manual_cast( + x.replace(zero_output.expand(x.feats.shape[0], -1)), self.dtype + ) + elif _injector is not None and not getattr(self, '_zobs_none_warned', False): + # Injector present but NO z_obs source (64ch input, kwarg omitted): the trained + # per-block bias is NOT applied — a third mode distinct from "z_obs dropped" (an + # all-zero 33ch tensor still adds the bias each block). A nozobs arm on a patchcross + # model must pass the zero tensor, never omit the kwarg. + print('[zobs_injector] WARNING: injector present but no z_obs source (kwarg omitted, ' + 'input < 97ch) — per-block bias NOT applied; pass an all-zero 33ch z_obs for a ' + 'true nozobs arm', flush=True) + self._zobs_none_warned = True + + partial_zobs_inject = None + partial_zobs_rows = None + partial_zobs_feats = None + partial_zobs_batch = None + if partial_z_obs is not None: + assert self.partial_zobs_injector is not None, \ + 'partial_z_obs passed but denoiser was built without partial_zobs_inject=true' + partial_zobs_inject, partial_zobs_rows, partial_zobs_feats, partial_zobs_batch = self._partial_zobs_message( + x, partial_z_obs, partial_z_obs_keep_mask + ) + partial_zobs_inject = manual_cast(partial_zobs_inject, self.dtype) + elif self.partial_zobs_injector is not None: + raise AssertionError('partial_zobs_inject=true requires a separate partial_z_obs tensor') + + if zobs_inject is not None and self.zobs_inject_scale != 1.0: + zobs_inject = zobs_inject.replace(zobs_inject.feats * self.zobs_inject_scale) + if not getattr(self, '_zobs_inject_scale_logged', False): + print(f'[zobs_injector] output scale={self.zobs_inject_scale:g}', flush=True) + self._zobs_inject_scale_logged = True + + # Keep the pre-AdaLN embedding for the innovation gate; blocks retain their existing + # modulation input exactly. + t_base = self.t_embedder(t) + t_emb = t_base + if self.share_mod: + t_emb = self.adaLN_modulation(t_emb) + t_emb = manual_cast(t_emb, self.dtype) + + innovation_time = None + innovation_y = None + innovation_gate_y = None + if self.innovation_gate is not None and partial_zobs_feats is not None: + # Observation-only terms are shared by all blocks. Empty rows deliberately pass + # through these linears too, so an all-drop batch keeps the full gate graph live. + gate_dtype = self.innovation_gate.embed_y.weight.dtype + innovation_gate_y = partial_zobs_feats.to(dtype=gate_dtype) + innovation_y = self.innovation_gate.embed_y(innovation_gate_y) + innovation_time = self.innovation_gate.time_proj(t_base) + + if self.pe_mode == "ape": + pe = self.pos_embedder(h.coords[:, 1:]) + h = h + manual_cast(pe, self.dtype) + + # Handle different conditioning modes + if self.image_attn_mode == 'proj': + if isinstance(cond, dict): + global_cond = cond['global'] + proj_cond = cond['proj'] + else: + global_cond, proj_cond = cond + if isinstance(global_cond, list): + global_cond = sp.VarLenTensor.from_tensor_list(global_cond) + global_cond = manual_cast(global_cond, self.dtype) + proj_cond = manual_cast(proj_cond, self.dtype) + cond = (global_cond, proj_cond) + elif self.image_attn_mode == 'gated_proj': + global_cond = cond['global'] + if isinstance(global_cond, list): + global_cond = sp.VarLenTensor.from_tensor_list(global_cond) + global_cond = manual_cast(global_cond, self.dtype) + proj_semantic = manual_cast(cond['proj_semantic'], self.dtype) + proj_color = manual_cast(cond['proj_color'], self.dtype) + cond = {'global': global_cond, 'proj_semantic': proj_semantic, 'proj_color': proj_color} + else: + if isinstance(cond, list): + cond = sp.VarLenTensor.from_tensor_list(cond) + cond = manual_cast(cond, self.dtype) + + # Inference-only probe hook. When the attribute is absent (the default, and the + # only state any training path produces) this is None and the loop below is + # byte-identical to the unhooked path: no constructor arg, no config key, and no + # entry in the inference ABI contract. A probe sets it to a length-num_blocks + # sequence to rescale the shared partial message per block. + _partial_scales = getattr(self, 'partial_zobs_layer_scales', None) + # Optional evaluator-only callable: ``fn(gate_value, active_batch, block_index)``. + # It is an instance attribute rather than a constructor/config argument, so no model + # state or training ABI changes; callers can implement neutral, constant, or shuffled + # gate-value probes deterministically. + _gate_value_transform = getattr(self, 'partial_zobs_innovation_gate_value_transform', None) + _innovation_ratio_stats = [] if self.innovation_gate is not None else None + _innovation_gate_mean_stats = [] if self.innovation_gate is not None else None + _innovation_observed_rows = 0 if partial_zobs_rows is None else int(partial_zobs_rows.numel()) + _local_attention = getattr(self, 'local_obs_attention', None) + _local_topology = None + _local_full_batch = None + _local_ratio_base_stats = [] if _local_attention is not None else None + _local_ratio_hidden_stats = [] if _local_attention is not None else None + _local_null_stats = [] if _local_attention is not None else None + _local_entropy_stats = [] if _local_attention is not None else None + _local_effective_neighbour_stats = [] if _local_attention is not None else None + _local_observed_residual_stats = [] if _local_attention is not None else None + _local_unobserved_residual_stats = [] if _local_attention is not None else None + _local_query_is_observed = None + _local_neighbourhood_histogram = None + if _local_attention is not None: + assert partial_zobs_rows is not None and partial_zobs_feats is not None + _local_topology = _build_two_hop_observed_topology(x.coords, partial_zobs_rows) + _local_full_batch = x.coords[:, 0].long() + _active_observed_mask = torch.zeros( + x.coords.shape[0], dtype=torch.bool, device=x.coords.device, + ) + _active_observed_mask[partial_zobs_rows] = True + _local_query_is_observed = _active_observed_mask[ + _local_topology['query_rows'] + ] + _local_neighbourhood_histogram = torch.bincount( + _local_topology['neighbour_counts'], minlength=26, + )[:26] + _local_neighbourhood_histogram[0] = ( + x.coords.shape[0] - _local_topology['query_rows'].numel() + ) + assert int(_local_neighbourhood_histogram.sum()) == x.coords.shape[0] + assert int(_local_neighbourhood_histogram[1:].sum()) == \ + _local_topology['query_rows'].numel() + _learned_attention = getattr(self, 'learned_attention', None) + _learned_cache = None + _learned_stats = [] if _learned_attention is not None else None + if _learned_attention is not None: + assert partial_zobs_inject is not None and partial_zobs_rows is not None + assert _partial_scales is None, \ + 'partial_zobs_layer_scales is not defined for the five-site V12LA operator' + _learned_cache = _learned_attention.build_cache( + x.coords, + x.layout, + partial_zobs_inject, + partial_zobs_rows, + ) + partial_zobs_inject_position = getattr( + self, 'partial_zobs_inject_position', 'post_block', + ) + for _block_index, block in enumerate(self.blocks): + if partial_zobs_inject is not None and partial_zobs_inject_position == 'pre_mlp': + _pre_mlp_scale = 1.0 if _partial_scales is None else float( + _partial_scales[_block_index] + ) + h = block( + h, + t_emb, + cond, + partial_zobs_rows=partial_zobs_rows, + partial_zobs_inject=partial_zobs_inject, + partial_zobs_inject_scale=_pre_mlp_scale, + ) + else: + h = block(h, t_emb, cond) + if zobs_inject is not None: + h = h + zobs_inject + _innovation_ratio = None + _innovation_gate_mean = None + _local_delta = None + _local_stats = None + if _local_attention is not None and _block_index < _local_attention.alpha.numel(): + assert _local_topology is not None and _local_full_batch is not None + _local_delta, _local_stats = _local_attention( + h.feats, + _local_full_batch, + t_base, + partial_zobs_feats, + partial_zobs_rows, + _local_topology, + _block_index, + ) + _local_delta = manual_cast(_local_delta, self.dtype) + _local_ratio_base = _local_delta.norm() / partial_zobs_inject.norm().clamp_min( + torch.finfo(partial_zobs_inject.dtype).eps + ) + _local_hidden = h.feats[_local_topology['query_rows']] + _local_ratio_hidden = _local_delta.norm() / _local_hidden.norm().clamp_min( + torch.finfo(_local_hidden.dtype).eps + ) + assert _local_query_is_observed is not None + _local_observed_residual = _local_delta[ + _local_query_is_observed + ].float().norm() + _local_unobserved_residual = _local_delta[ + ~_local_query_is_observed + ].float().norm() + if _learned_attention is not None and _block_index in _learned_attention.site_to_slot: + assert _learned_cache is not None + _learned_inject, _site_stats = _learned_attention( + h.feats, + h.coords, + h.layout, + _learned_cache, + _block_index, + ) + h = h.replace(h.feats + manual_cast(_learned_inject, self.dtype)) + _learned_stats.append({ + name: value.detach() for name, value in _site_stats.items() + }) + if partial_zobs_inject is not None and _learned_attention is None and \ + partial_zobs_inject_position == 'post_block': + layer_inject = partial_zobs_inject + if self.innovation_gate is not None: + gate = self.innovation_gate + # Every operation intentionally accepts an empty leading dimension. This + # keeps all innovation parameters in the all-drop DDP graph. + assert innovation_gate_y is not None and innovation_y is not None and innovation_time is not None + t_obs = innovation_time[partial_zobs_batch] + h_obs = h.feats[partial_zobs_rows].to(dtype=gate.embed_y.weight.dtype) + x_embed = gate.embed_x(torch.cat([h_obs, t_obs], dim=-1)) + innovation = innovation_y - x_embed + gate_input = torch.cat([ + innovation, + innovation_gate_y[:, 32:34], + t_obs, + ], dim=-1) + gate_value = torch.sigmoid(gate.mlp(gate_input)) + if _gate_value_transform is not None: + gate_value = _gate_value_transform( + gate_value, partial_zobs_batch, _block_index + ) + assert gate_value.shape == (innovation.shape[0], 1), \ + f'gate value transform must return [Nobs,1], got {tuple(gate_value.shape)}' + delta = gate.proj(innovation) * gate_value + delta = delta * gate.alpha[_block_index].to(dtype=delta.dtype) + if _block_index == len(self.blocks) - 1: + # Keep the frozen V11 base carrier at the terminal site, + # but never add a newly learned observed-only shortcut: + # no attention block remains to transport/smooth it into U. + delta = delta * 0 + # Detached telemetry for the training log. This is before the legacy + # layer-scale hook and therefore measures the learned branch itself. + _innovation_ratio = delta.norm() / partial_zobs_inject.norm().clamp_min( + torch.finfo(partial_zobs_inject.dtype).eps + ) + _innovation_gate_mean = ( + gate_value.mean() if gate_value.numel() != 0 else delta.new_zeros(()) + ) + layer_inject = partial_zobs_inject + manual_cast(delta, self.dtype) + if partial_zobs_rows.numel() != 0: + if _partial_scales is None: + h = h.replace(h.feats.index_add( + 0, partial_zobs_rows, layer_inject + )) + else: + h = h.replace(h.feats.index_add( + 0, partial_zobs_rows, layer_inject, + alpha=float(_partial_scales[_block_index]), + )) + else: + # Preserve every empty carrier/gate GEMM in the graph on all-drop batches. + h = h.replace(h.feats + layer_inject.sum() * 0) + if _local_delta is not None: + _local_scale = 1.0 if _partial_scales is None else float( + _partial_scales[_block_index] + ) + if _local_topology['query_rows'].numel() != 0: + h = h.replace(h.feats.index_add( + 0, _local_topology['query_rows'], _local_delta, + alpha=_local_scale, + )) + else: + # The module returns an exact-zero parameter edge for all 14 + # trainable tensors, so DDP observes finite numerical-zero grads. + h = h.replace(h.feats + _local_delta.sum() * _local_scale) + if self.innovation_gate is not None: + if _innovation_ratio is None: + _innovation_ratio = h.feats.new_zeros(()) + _innovation_gate_mean = h.feats.new_zeros(()) + _innovation_ratio_stats.append(_innovation_ratio.detach()) + _innovation_gate_mean_stats.append(_innovation_gate_mean.detach()) + if _local_attention is not None and _block_index < _local_attention.alpha.numel(): + assert _local_stats is not None + _local_ratio_base_stats.append(_local_ratio_base.detach()) + _local_ratio_hidden_stats.append(_local_ratio_hidden.detach()) + _local_null_stats.append(_local_stats['null_attention_mean'].detach()) + _local_entropy_stats.append(_local_stats['real_neighbour_entropy_mean'].detach()) + _local_effective_neighbour_stats.append( + _local_stats['effective_real_neighbour_count_mean'].detach() + ) + _local_observed_residual_stats.append( + _local_observed_residual.detach() + ) + _local_unobserved_residual_stats.append( + _local_unobserved_residual.detach() + ) + + if self.innovation_gate is not None: + # Instance-only telemetry: deliberately neither buffer nor state-dict entry. + self.partial_zobs_innovation_gate_last_stats = { + 'delta_to_base_norm_ratio': torch.stack(_innovation_ratio_stats), + 'gate_mean': torch.stack(_innovation_gate_mean_stats), + 'observed_rows': _innovation_observed_rows, + } + if _local_attention is not None: + # Block 29 is structurally absent: every telemetry vector has length 29. + assert _local_query_is_observed is not None + assert _local_neighbourhood_histogram is not None + _full_support_rows = int(x.coords.shape[0]) + _active_queries = int(_local_topology['query_rows'].numel()) + _empty_queries = _full_support_rows - _active_queries + _singleton_queries = int( + (_local_topology['neighbour_counts'] == 1).sum() + ) + self.partial_zobs_local_attention_last_stats = { + 'delta_to_base_norm_ratio': torch.stack(_local_ratio_base_stats), + 'delta_to_hidden_norm_ratio': torch.stack(_local_ratio_hidden_stats), + 'null_attention_mean': torch.stack(_local_null_stats), + 'real_neighbour_entropy_mean': torch.stack(_local_entropy_stats), + 'effective_real_neighbour_count_mean': torch.stack( + _local_effective_neighbour_stats + ), + 'observed_residual_l2_norm': torch.stack( + _local_observed_residual_stats + ), + 'unobserved_residual_l2_norm': torch.stack( + _local_unobserved_residual_stats + ), + 'observed_rows': _innovation_observed_rows, + 'full_support_rows': _full_support_rows, + 'active_queries': _active_queries, + 'empty_queries': _empty_queries, + 'singleton_queries': _singleton_queries, + 'observed_active_queries': int(_local_query_is_observed.sum()), + 'unobserved_active_queries': int((~_local_query_is_observed).sum()), + 'active_query_fraction': _active_queries / max(_full_support_rows, 1), + 'empty_query_fraction': _empty_queries / max(_full_support_rows, 1), + 'singleton_query_fraction': _singleton_queries / max(_active_queries, 1), + 'neighbourhood_size_histogram': _local_neighbourhood_histogram.detach(), + 'packed_pairs': int(_local_topology['source_indices'].numel()), + } + if _learned_attention is not None: + assert _learned_cache is not None + assert len(_learned_stats) == len(_learned_attention.sites) + self.partial_zobs_learned_attention_last_stats = { + name: torch.stack([site[name] for site in _learned_stats]) + for name in _learned_stats[0] + } + self.partial_zobs_learned_attention_last_stats.update({ + 'sites': _learned_attention.sites, + 'observed_rows': int(partial_zobs_rows.numel()), + 'full_support_rows': int(x.coords.shape[0]), + 'observed_counts': _learned_cache['observed_counts'], + }) + + h = manual_cast(h, x.dtype) + h = h.replace(F.layer_norm(h.feats, h.feats.shape[-1:])) + h = self.out_layer(h) + return h + + +# Existing checkpoint configs use this name; inference uses the same base model. +ElasticSLatFlowModel = SLatFlowModel diff --git a/supratex/pipeline_config.py b/supratex/pipeline_config.py new file mode 100644 index 0000000000000000000000000000000000000000..8f5d76ff24e56054ebf1055d34b84762ec94087b --- /dev/null +++ b/supratex/pipeline_config.py @@ -0,0 +1,192 @@ +"""Read the SUPRA-Tex release manifest without importing model runtimes.""" +from __future__ import annotations + +import json +import os +from pathlib import Path, PurePosixPath +import re +from urllib.parse import unquote, urlsplit + + +ROOT = Path(__file__).resolve().parents[1] + + +def runtime_root(root: Path) -> Path: + """Locate the single runtime tree in a source checkout or legacy bundle.""" + source = root / "deployment/runtime" + bundled = root / "runtime" + return bundled if not source.is_dir() and bundled.is_dir() else source + + +def relative_name(value: str) -> str: + if (not isinstance(value, str) or not value or "\\" in value or ":" in value + or PurePosixPath(value).is_absolute() + or any(part in ("", ".", "..") for part in value.split("/")) + or any(ord(char) < 32 for char in value)): + raise ValueError(f"expected a relative path without traversal: {value!r}") + return value + + +def model_source(value: str | os.PathLike, revision: str | None = None) -> dict: + """Normalize local directories and public-form HF IDs/URLs without networking.""" + local_only = isinstance(value, os.PathLike) + value = os.fspath(value) + if not isinstance(value, str) or not value: + raise ValueError("model source must be a nonempty path or namespace/repo HF ID") + if revision is not None and (not isinstance(revision, str) or not revision.strip()): + raise ValueError("model revision must be a nonempty string") + local = Path(value).expanduser() + if local.exists(): + if not local.is_dir(): + raise NotADirectoryError(f"model source must be a directory: {local}") + return {"kind": "local", "path": str(local.resolve())} + if local_only or value.startswith(("/", "./", "../", "~")) or value in (".", ".."): + raise FileNotFoundError(f"model directory does not exist: {local}") + if value.startswith("https://"): + url = urlsplit(value) + if url.netloc != "huggingface.co" or url.query or url.fragment: + raise ValueError("use a public-form https://huggingface.co/namespace/repo URL without credentials or query") + parts = unquote(url.path).strip("/").split("/") + if len(parts) == 4 and parts[2] == "tree": + if revision is not None and revision != parts[3]: + raise ValueError("HF URL revision conflicts with --revision") + revision = parts[3] + parts = parts[:2] + if len(parts) != 2: + raise ValueError("HF URL must select a repository root or /tree/") + value = "/".join(parts) + value = value.removeprefix("hf://") + if (not re.fullmatch(r"[A-Za-z0-9_][A-Za-z0-9_.-]*/[A-Za-z0-9_][A-Za-z0-9_.-]*", value) + or ".." in value or "--" in value or value.endswith((".", "-"))): + raise ValueError("model source must be an existing directory or a namespace/repo HF ID or URL") + return {"kind": "hub", "repo_id": value, "revision": revision} + + +def _keys(value: object, required: set[str], name: str, optional: set[str] | None = None) -> None: + if not isinstance(value, dict): + raise ValueError(f"pipeline {name} must be an object") + missing = required - value.keys() + unknown = value.keys() - required - (optional or set()) + if missing or unknown: + raise ValueError(f"pipeline {name} has missing keys {sorted(missing)} or unsupported keys {sorted(unknown)}") + + +def load_pipeline(path: str | os.PathLike) -> dict: + """Load and validate the data-only schema; component configs stay separate.""" + path = Path(path).expanduser() + if not path.is_file(): + raise FileNotFoundError(f"pipeline manifest does not exist: {path}") + manifest = json.loads(path.read_text(encoding="utf-8")) + _keys(manifest, {"format", "schema_version", "models", "external_models", "defaults"}, "root") + if (manifest["format"] != "supratex_pipeline" or type(manifest["schema_version"]) is not int + or manifest["schema_version"] != 1): + raise ValueError("pipeline requires format 'supratex_pipeline' and schema_version 1") + _keys(manifest["models"], {"completion", "material"}, "models") + for value in manifest["models"].values(): + relative_name(value) + external = manifest["external_models"] + _keys(external, {"qwen-base", "geometry", "moge2", "dino", "rmbg2", "vae", "naf"}, "external_models") + for name, entry in external.items(): + if name == "naf": + _keys(entry, {"url", "filename"}, "external_models.naf") + filename = relative_name(entry["filename"]) + if "/" in filename: + raise ValueError("pipeline NAF filename must be a single filename") + if not isinstance(entry["url"], str): + raise ValueError("pipeline NAF url must be an HTTPS URL") + url = urlsplit(entry["url"]) + if (url.scheme != "https" or not url.hostname or url.username or url.password + or url.query or url.fragment or not url.path + or any(char.isspace() for char in entry["url"])): + raise ValueError("pipeline NAF url must be an HTTPS URL without credentials or query") + continue + _keys(entry, {"repo_id", "revision", "allow_patterns"}, f"external_models.{name}", {"gated"}) + repo_id = entry["repo_id"] + if not isinstance(repo_id, str): + raise ValueError(f"pipeline {name} repo_id must be a namespace/repo HF ID") + source = model_source("hf://" + repo_id, entry["revision"]) + if source["kind"] != "hub" or source["repo_id"] != repo_id or entry["revision"] is None: + raise ValueError(f"pipeline {name} requires an HF repo_id and a nonempty revision") + patterns = entry["allow_patterns"] + if not isinstance(patterns, list) or not patterns: + raise ValueError(f"pipeline {name} allow_patterns must be a nonempty list") + for pattern in patterns: + relative_name(pattern) + if "gated" in entry and type(entry["gated"]) is not bool: + raise ValueError(f"pipeline {name} gated must be a boolean") + validate_defaults(manifest["defaults"]) + return manifest + + +def validate_defaults(defaults: dict) -> dict: + """Validate a complete defaults object, including resolved runtime configs.""" + _keys(defaults, {"seed", "texture_size", "decimation_target", "texture_format", "attention"}, "defaults") + for name in ("seed", "texture_size", "decimation_target"): + if type(defaults[name]) is not int or defaults[name] < (0 if name == "seed" else 1): + raise ValueError(f"pipeline default {name} must be {'nonnegative' if name == 'seed' else 'positive'} integer") + if defaults["texture_format"] not in ("png", "webp"): + raise ValueError("pipeline default texture_format must be png or webp") + if defaults["attention"] not in ("auto", "flash_attn_3", "flash_attn", "sdpa", "xformers"): + raise ValueError("pipeline default attention is unsupported") + return defaults + + +# Refinement reads and writes a 4096-square PNG base-color atlas. +REFINE_ATLAS = {"texture_size": 4096, "texture_format": "png"} + + +def validate_export_options(args) -> None: + """Check export options shared by infer.py and texturing.py; None means not yet resolved.""" + # Dual-view always refines both views, so it implies --refine. + if (args.dual_view or getattr(args, "refine", False)) and any( + getattr(args, key) not in (None, value) for key, value in REFINE_ATLAS.items()): + raise ValueError("--refine and --dual-view require the 4096 PNG atlas") + if ((args.seed is not None and args.seed < 0) + or (args.texture_size is not None and args.texture_size <= 0) + or (args.decimation_target is not None and args.decimation_target <= 0)): + raise ValueError("seed must be nonnegative; texture-size and decimation-target must be positive") + + +def default_pipeline() -> dict: + """Read the repository's canonical defaults, without maintaining a copy.""" + return load_pipeline(ROOT / "configs/pipeline.json") + + +def resolve_pipeline( + source: str | os.PathLike, + *, + revision: str | None = None, + cache_dir: str | os.PathLike | None = None, + local_files_only: bool = False, +) -> tuple[dict, dict]: + """Return a manifest and its local release root or pinned Hub descriptor. + + Hub resolution downloads only ``pipeline.json``. The returned Hub revision + follows its cache snapshot so subsequent component downloads use that same + commit even if the requested branch changes in the meantime. + """ + value = os.fspath(source) + if not isinstance(value, str) or not value: + raise ValueError("pipeline source must be a path or a namespace/repo HF ID") + path = Path(value).expanduser() + if path.is_file() or (not value.startswith(("https://", "hf://")) and path.suffix == ".json"): + # Keep the logical path: an exported manifest can be a symlink, and its + # component paths must still be relative to the selected release root. + path = path.absolute() + return load_pipeline(path), {"kind": "local", "path": str(path.parent)} + descriptor = model_source(source, revision) + if descriptor["kind"] == "local": + return load_pipeline(Path(descriptor["path"]) / "pipeline.json"), descriptor + from huggingface_hub import hf_hub_download + + try: + path = Path(hf_hub_download(repo_id=descriptor["repo_id"], filename="pipeline.json", + revision=descriptor["revision"], cache_dir=cache_dir, + local_files_only=local_files_only)) + except Exception: + detail = "a cached manifest is required in offline mode" if local_files_only else "check the repository, revision and your access" + raise RuntimeError(f"Cannot load pipeline.json from {descriptor['repo_id']}; {detail}") from None + manifest = load_pipeline(path) + if path.parent.parent.name == "snapshots" and re.fullmatch(r"[0-9a-fA-F]{40}", path.parent.name): + descriptor["revision"] = path.parent.name + return manifest, descriptor diff --git a/supratex/pipelines/observation_guided_completion.py b/supratex/pipelines/observation_guided_completion.py index 5b6d60450214bc1d9e3b858e5929ab523b28269a..ce701f59d1b1e52a756e1bfb140ae41ca2423155 100644 --- a/supratex/pipelines/observation_guided_completion.py +++ b/supratex/pipelines/observation_guided_completion.py @@ -1,160 +1,96 @@ -"""Strict native-checkpoint inference for static-O Interleaved flow transformer texture flow.""" -from __future__ import annotations - -from contextlib import nullcontext -from copy import deepcopy -import math -from typing import Any, Mapping - -import torch - -from ..models.interleaved_flow_transformer import InterleavedFlowTransformer -from .samplers.flow_euler import FlowEulerSampler - - -_POLICY = { - 'num_blocks': 60, - 'v11_init_policy': 'odd_copy_even_identity', - 'base_block_policy': 'alternating_u_x', - 'direct_modulation_broadcast': True, -} -_RETIRED_O_KWARGS = frozenset({'o_update_scale', 'o_transition_steps'}) -_NORMALIZATION_KEYS = ('partial_mean', 'partial_std', 'pbr_mean', 'pbr_std') - - -def build_completion_model( - config: Mapping[str, Any], state_dict: Mapping[str, torch.Tensor], *, - allow_test_tiny: bool = False, - assign_weights: bool = False, -) -> InterleavedFlowTransformer: - """Strict-load a native static-O Interleaved flow transformer state. - - Parent migration belongs to training. Native Interleaved flow transformer state is always static - O; legacy transition config and state packets are intentionally rejected. - """ - try: - spec = config['models']['denoiser'] - name, raw_args = spec['name'], spec['args'] - except (KeyError, TypeError) as error: - raise ValueError('Interleaved flow transformer config requires models.denoiser.name and .args') from error - if name != 'InterleavedFlowTransformer' or not isinstance(raw_args, Mapping): - raise ValueError('Interleaved flow transformer inference requires InterleavedFlowTransformer config') - if not isinstance(allow_test_tiny, bool): - raise ValueError('allow_test_tiny must be a bool') - if not isinstance(assign_weights, bool): - raise ValueError('assign_weights must be a bool') - args = deepcopy(dict(raw_args)) - retired = sorted(_RETIRED_O_KWARGS & set(args)) - if retired: - raise ValueError(f'Interleaved flow transformer no longer supports retired model args: {retired}') - for key, expected in _POLICY.items(): - if args.get(key) != expected: - raise ValueError(f'Interleaved flow transformer inference requires {key}={expected!r}') - if args.get('rope_freq') not in ([1, 10000], (1, 10000)): - raise ValueError('Interleaved flow transformer inference requires rope_freq=[1, 10000]') - if args.get('dtype') != 'float32': - raise ValueError('Interleaved flow transformer inference requires dtype="float32"') - if args.get('_test_tiny', False) and not allow_test_tiny: - raise ValueError('Interleaved flow transformer inference rejects _test_tiny without allow_test_tiny=True') - if not isinstance(state_dict, Mapping): - raise TypeError('state_dict must be a mapping for strict Interleaved flow transformer loading') - # Validate config values outside the meta context so they remain available - # for exact comparison before checkpoint buffers replace model buffers. - normalization = {} - for key in _NORMALIZATION_KEYS: - try: - value = torch.as_tensor(args[key], dtype=torch.float32, device='cpu') - except (KeyError, TypeError, ValueError, RuntimeError) as error: - raise ValueError(f'Invalid config normalization {key}: expected a finite 32-vector') from error - if value.shape != (32,) or not bool(torch.isfinite(value).all()): - raise ValueError(f'Invalid config normalization {key}: expected a finite 32-vector') - if key.endswith('_std') and not bool((value > 0).all()): - raise ValueError(f'Invalid config normalization {key}: standard deviations must be positive') - normalization[key] = value - # The inference loader already places checkpoint tensors on the target - # device. A meta shell avoids allocating/initializing a second full model - # and copying its weights through CPU memory before moving it back to CUDA. - with torch.device('meta') if assign_weights else nullcontext(): - model = InterleavedFlowTransformer(**args) - target = model.state_dict() - missing = sorted(set(target) - set(state_dict)) - unexpected = sorted(set(state_dict) - set(target)) - invalid = sorted( - key for key in set(target) & set(state_dict) - if not isinstance(state_dict[key], torch.Tensor) - or state_dict[key].is_meta - or state_dict[key].shape != target[key].shape - or state_dict[key].dtype != target[key].dtype - ) - if missing or unexpected or invalid: - raise RuntimeError( - 'Invalid native Interleaved flow transformer checkpoint before load: ' - f'missing={missing}, unexpected={unexpected}, type_shape_dtype_mismatch={invalid}' - ) - mismatched_normalization = [ - key for key, expected in normalization.items() - if not torch.equal(state_dict[key].detach().cpu(), expected) - ] - if mismatched_normalization: - raise RuntimeError( - 'Invalid native Interleaved flow transformer checkpoint before load: ' - f'normalization_mismatch={mismatched_normalization}; ' - 'buffers must match the finite config values with positive standard deviations' - ) - model.load_state_dict(state_dict, strict=True, assign=assign_weights) - if assign_weights: - from ..modules.sparse.attention.rope import SparseRotaryPositionEmbedder - for module in model.modules(): - if isinstance(module, SparseRotaryPositionEmbedder): - # RoPE frequencies are a plain attribute, absent from the - # checkpoint and nn.Module.to(); materialize them explicitly. - frequencies = torch.arange(module.freq_dim, dtype=torch.float32, device='cpu') / module.freq_dim - module.freqs = module.rope_freq[0] / (module.rope_freq[1] ** frequencies) - return model.eval() - - -def sample_material_latents( - model: InterleavedFlowTransformer, - noise: Any, - *, - concat_cond: Any, - partial_z_obs: Any, - partial_z_obs_keep_mask: Any = None, - steps: int = 25, - rescale_t: float = 3.0, - sigma_min: float = 1e-5, - use_observation_cache: bool = False, - verbose: bool = False, -) -> Any: - """Sample static-O Interleaved flow transformer with native Euler settings and an opt-in O K/V cache.""" - if not isinstance(model, InterleavedFlowTransformer): - raise TypeError('Interleaved flow transformer sampling requires InterleavedFlowTransformer') - if model.training: - raise RuntimeError('Interleaved flow transformer sampling requires an eval model') - for key, expected in _POLICY.items(): - if getattr(model, key) != expected: - raise RuntimeError(f'Interleaved flow transformer sampling requires {key}={expected!r}') - if not isinstance(use_observation_cache, bool): - raise ValueError('use_observation_cache must be a bool') - if isinstance(steps, bool) or not isinstance(steps, int) or steps < 1: - raise ValueError('steps must be a positive integer') - if isinstance(rescale_t, bool) or not isinstance(rescale_t, (int, float)) \ - or not math.isfinite(rescale_t) or rescale_t <= 0: - raise ValueError('rescale_t must be finite and positive') - if isinstance(sigma_min, bool) or not isinstance(sigma_min, (int, float)) \ - or not math.isfinite(sigma_min) or not 0 < sigma_min < 1: - raise ValueError('sigma_min must be finite and in (0, 1)') - sampler = FlowEulerSampler(sigma_min=sigma_min) - autocast = (torch.autocast('cuda', dtype=torch.bfloat16) - if getattr(noise, 'device', torch.device('cpu')).type == 'cuda' else nullcontext()) - with torch.no_grad(), autocast: - cache_context = (model.observation_cache(noise, concat_cond, partial_z_obs, - partial_z_obs_keep_mask) - if use_observation_cache else nullcontext(None)) - with cache_context as observation_cache: - return sampler.sample( - model, noise, cond=None, steps=steps, rescale_t=rescale_t, - verbose=verbose, concat_cond=concat_cond, partial_z_obs=partial_z_obs, - partial_z_obs_keep_mask=partial_z_obs_keep_mask, - observation_cache=observation_cache, - ) +"""Strict native-checkpoint inference for static-O Interleaved flow transformer texture flow.""" +from __future__ import annotations + +from contextlib import nullcontext +from typing import Any, Mapping + +import torch + +from ..models.interleaved_flow_transformer import InterleavedFlowTransformer +from .samplers.flow_euler import FlowEulerSampler + + +_NORMALIZATION_KEYS = ('partial_mean', 'partial_std', 'pbr_mean', 'pbr_std') + + +def build_completion_model( + config: Mapping[str, Any], state_dict: Mapping[str, torch.Tensor], *, + assign_weights: bool = False, +) -> InterleavedFlowTransformer: + """Strict-load a native static-O Interleaved flow transformer state. + + ``config`` must already have passed ``supratex.inference.pipeline.load_contract``, + which checks the released model ABI and normalization statistics. + """ + args = dict(config['models']['denoiser']['args']) + # The inference loader already places checkpoint tensors on the target + # device. A meta shell avoids allocating/initializing a second full model + # and copying its weights through CPU memory before moving it back to CUDA. + with torch.device('meta') if assign_weights else nullcontext(): + model = InterleavedFlowTransformer(**args) + target = model.state_dict() + missing = sorted(set(target) - set(state_dict)) + unexpected = sorted(set(state_dict) - set(target)) + invalid = sorted( + key for key in set(target) & set(state_dict) + if not isinstance(state_dict[key], torch.Tensor) + or state_dict[key].is_meta + or state_dict[key].shape != target[key].shape + or state_dict[key].dtype != target[key].dtype + ) + if missing or unexpected or invalid: + raise RuntimeError( + 'Invalid native Interleaved flow transformer checkpoint before load: ' + f'missing={missing}, unexpected={unexpected}, type_shape_dtype_mismatch={invalid}' + ) + # Normalization buffers are stored in the checkpoint; they must equal the + # config statistics that the rest of the pipeline uses. + mismatched_normalization = [ + key for key in _NORMALIZATION_KEYS + if not torch.equal(state_dict[key].detach().cpu(), torch.as_tensor(args[key], dtype=torch.float32)) + ] + if mismatched_normalization: + raise RuntimeError( + 'Invalid native Interleaved flow transformer checkpoint before load: ' + f'normalization_mismatch={mismatched_normalization}' + ) + model.load_state_dict(state_dict, strict=True, assign=assign_weights) + if assign_weights: + from ..modules.sparse.attention.rope import SparseRotaryPositionEmbedder + for module in model.modules(): + if isinstance(module, SparseRotaryPositionEmbedder): + # RoPE frequencies are a plain attribute, absent from the + # checkpoint and nn.Module.to(); materialize them explicitly. + frequencies = torch.arange(module.freq_dim, dtype=torch.float32, device='cpu') / module.freq_dim + module.freqs = module.rope_freq[0] / (module.rope_freq[1] ** frequencies) + return model.eval() + + +def sample_material_latents( + model: InterleavedFlowTransformer, + noise: Any, + *, + concat_cond: Any, + partial_z_obs: Any, + partial_z_obs_keep_mask: Any = None, + steps: int = 25, + rescale_t: float = 3.0, + sigma_min: float = 1e-5, + use_observation_cache: bool = False, + verbose: bool = False, +) -> Any: + """Sample static-O Interleaved flow transformer with native Euler settings and an opt-in O K/V cache.""" + sampler = FlowEulerSampler(sigma_min=sigma_min) + autocast = (torch.autocast('cuda', dtype=torch.bfloat16) + if getattr(noise, 'device', torch.device('cpu')).type == 'cuda' else nullcontext()) + with torch.no_grad(), autocast: + cache_context = (model.observation_cache(noise, concat_cond, partial_z_obs, + partial_z_obs_keep_mask) + if use_observation_cache else nullcontext(None)) + with cache_context as observation_cache: + return sampler.sample( + model, noise, cond=None, steps=steps, rescale_t=rescale_t, + verbose=verbose, concat_cond=concat_cond, partial_z_obs=partial_z_obs, + partial_z_obs_keep_mask=partial_z_obs_keep_mask, + observation_cache=observation_cache, + ) diff --git a/supratex/renderers/pbr_mesh_renderer.py b/supratex/renderers/pbr_mesh_renderer.py index 14ded300c861e163e15c1bff4c3dd7ee27a08e7c..d4ce1e9f937fab58a5e0e62b8002638937de1080 100644 --- a/supratex/renderers/pbr_mesh_renderer.py +++ b/supratex/renderers/pbr_mesh_renderer.py @@ -1,521 +1,520 @@ -from typing import * -import torch -from easydict import EasyDict as edict -import numpy as np -import utils3d -from ..representations.mesh import Mesh, MeshWithVoxel, MeshWithPbrMaterial, TextureFilterMode, AlphaMode, TextureWrapMode -import torch.nn.functional as F - - -def cube_to_dir(s, x, y): - if s == 0: rx, ry, rz = torch.ones_like(x), -x, -y - elif s == 1: rx, ry, rz = -torch.ones_like(x), x, -y - elif s == 2: rx, ry, rz = x, y, torch.ones_like(x) - elif s == 3: rx, ry, rz = x, -y, -torch.ones_like(x) - elif s == 4: rx, ry, rz = x, torch.ones_like(x), -y - elif s == 5: rx, ry, rz = -x, -torch.ones_like(x), -y - return torch.stack((rx, ry, rz), dim=-1) - - -def latlong_to_cubemap(latlong_map, res): - if 'dr' not in globals(): - import nvdiffrast.torch as dr - cubemap = torch.zeros(6, res[0], res[1], latlong_map.shape[-1], dtype=torch.float32, device='cuda') - for s in range(6): - gy, gx = torch.meshgrid(torch.linspace(-1.0 + 1.0 / res[0], 1.0 - 1.0 / res[0], res[0], device='cuda'), - torch.linspace(-1.0 + 1.0 / res[1], 1.0 - 1.0 / res[1], res[1], device='cuda'), - indexing='ij') - v = F.normalize(cube_to_dir(s, gx, gy), dim=-1) - - tu = torch.atan2(v[..., 0:1], -v[..., 2:3]) / (2 * np.pi) + 0.5 - tv = torch.acos(torch.clamp(v[..., 1:2], min=-1, max=1)) / np.pi - texcoord = torch.cat((tu, tv), dim=-1) - - cubemap[s, ...] = dr.texture(latlong_map[None, ...], texcoord[None, ...], filter_mode='linear')[0] - return cubemap - - -class EnvMap: - def __init__(self, image: torch.Tensor): - self.image = image - - @property - def _backend(self): - if not hasattr(self, '_nvdiffrec_envlight'): - if 'EnvironmentLight' not in globals(): - from nvdiffrec_render.light import EnvironmentLight - cubemap = latlong_to_cubemap(self.image, [512, 512]) - self._nvdiffrec_envlight = EnvironmentLight(cubemap) - self._nvdiffrec_envlight.build_mips() - return self._nvdiffrec_envlight - - def shade(self, gb_pos, gb_normal, kd, ks, view_pos, specular=True): - return self._backend.shade(gb_pos, gb_normal, kd, ks, view_pos, specular) - - def sample(self, directions: torch.Tensor): - if 'dr' not in globals(): - import nvdiffrast.torch as dr - return dr.texture( - self._backend.base.unsqueeze(0), - directions.unsqueeze(0), - boundary_mode='cube', - )[0] - - -def intrinsics_to_projection( - intrinsics: torch.Tensor, - near: float, - far: float, - ) -> torch.Tensor: - """ - OpenCV intrinsics to OpenGL perspective matrix - - Args: - intrinsics (torch.Tensor): [3, 3] OpenCV intrinsics matrix - near (float): near plane to clip - far (float): far plane to clip - Returns: - (torch.Tensor): [4, 4] OpenGL perspective matrix - """ - fx, fy = intrinsics[0, 0], intrinsics[1, 1] - cx, cy = intrinsics[0, 2], intrinsics[1, 2] - ret = torch.zeros((4, 4), dtype=intrinsics.dtype, device=intrinsics.device) - ret[0, 0] = 2 * fx - ret[1, 1] = 2 * fy - ret[0, 2] = 2 * cx - 1 - ret[1, 2] = - 2 * cy + 1 - ret[2, 2] = (far + near) / (far - near) - ret[2, 3] = 2 * near * far / (near - far) - ret[3, 2] = 1. - return ret - - -def screen_space_ambient_occlusion( - depth: torch.Tensor, - normal: torch.Tensor, - perspective: torch.Tensor, - radius: float = 0.1, - bias: float = 1e-6, - samples: int = 64, - intensity: float = 1.0, -) -> torch.Tensor: - """ - Screen space ambient occlusion (SSAO) - - Args: - depth (torch.Tensor): [H, W, 1] depth image - normal (torch.Tensor): [H, W, 3] normal image - perspective (torch.Tensor): [4, 4] camera projection matrix - radius (float): radius of the SSAO kernel - bias (float): bias to avoid self-occlusion - samples (int): number of samples to use for the SSAO kernel - intensity (float): intensity of the SSAO effect - Returns: - (torch.Tensor): [H, W, 1] SSAO image - """ - device = depth.device - H, W, _ = depth.shape - - fx = perspective[0, 0] - fy = perspective[1, 1] - cx = perspective[0, 2] - cy = perspective[1, 2] - - y_grid, x_grid = torch.meshgrid( - (torch.arange(H, device=device) + 0.5) / H * 2 - 1, - (torch.arange(W, device=device) + 0.5) / W * 2 - 1, - indexing='ij' - ) - x_view = (x_grid.float() - cx) * depth[..., 0] / fx - y_view = (y_grid.float() - cy) * depth[..., 0] / fy - view_pos = torch.stack([x_view, y_view, depth[..., 0]], dim=-1) # [H, W, 3] - - depth_feat = depth.permute(2, 0, 1).unsqueeze(0) - occlusion = torch.zeros((H, W), device=device) - - # start sampling - for _ in range(samples): - # sample normal distribution, if inside, flip the sign - rnd_vec = torch.randn(H, W, 3, device=device) - rnd_vec = F.normalize(rnd_vec, p=2, dim=-1) - dot_val = torch.sum(rnd_vec * normal, dim=-1, keepdim=True) - sample_dir = torch.sign(dot_val) * rnd_vec - scale = torch.rand(H, W, 1, device=device) - scale = scale * scale - sample_pos = view_pos + sample_dir * radius * scale - sample_z = sample_pos[..., 2] - - # project to screen space - z_safe = torch.clamp(sample_pos[..., 2], min=1e-5) - proj_u = (sample_pos[..., 0] * fx / z_safe) + cx - proj_v = (sample_pos[..., 1] * fy / z_safe) + cy - grid = torch.stack([proj_u, proj_v], dim=-1).unsqueeze(0) - geo_z = F.grid_sample(depth_feat, grid, mode='nearest', padding_mode='border').squeeze() - range_check = torch.abs(geo_z - sample_z) < radius - is_occluded = (geo_z <= sample_z - bias) & range_check - occlusion += is_occluded.float() - - f_occ = occlusion / samples * intensity - f_occ = torch.clamp(f_occ, 0.0, 1.0) - - return f_occ.unsqueeze(-1) - - -def aces_tonemapping(x: torch.Tensor) -> torch.Tensor: - """ - Applies ACES tone mapping curve to an HDR image tensor. - Input: x - HDR tensor, shape (..., 3), range [0, +inf) - Output: LDR tensor, same shape, range [0, 1] - - NOTE: This causes bright base_color objects to over-expose (appear white) - compared to Blender's standard sRGB display transform. Use linear_to_srgb() - instead for renders_cond alignment. See pbr_color_diff_debug_guide.md #5. - """ - a = 2.51 - b = 0.03 - c = 2.43 - d = 0.59 - e = 0.14 - - # Apply the ACES fitted curve - mapped = (x * (a * x + b)) / (x * (c * x + d) + e) - - # Clamp to [0, 1] for display or saving - return torch.clamp(mapped, 0.0, 1.0) - - -def gamma_correction(x: torch.Tensor, gamma: float = 2.2) -> torch.Tensor: - """ - Applies gamma correction to an HDR image tensor. - """ - return torch.clamp(x ** (1.0 / gamma), 0.0, 1.0) - - -def linear_to_srgb(x: torch.Tensor) -> torch.Tensor: - """ - Standard linear-to-sRGB conversion matching Blender's sRGB display transform. - - Applies the official sRGB EOTF inverse: - - For values <= 0.0031308: sRGB = 12.92 * linear - - For values > 0.0031308: sRGB = 1.055 * linear^(1/2.4) - 0.055 - - Input: x - linear HDR tensor, clamped to [0, +inf) - Output: sRGB tensor, range [0, 1] - """ - x = torch.clamp(x, 0.0) - low = 12.92 * x - high = 1.055 * torch.pow(torch.clamp(x, min=0.0031308), 1.0 / 2.4) - 0.055 - return torch.clamp(torch.where(x <= 0.0031308, low, high), 0.0, 1.0) - - -class PbrMeshRenderer: - """ - Renderer for the PBR mesh. - - Args: - rendering_options (dict): Rendering options. - """ - def __init__(self, rendering_options={}, device='cuda'): - if 'dr' not in globals(): - import nvdiffrast.torch as dr - - self.rendering_options = edict({ - "resolution": None, - "near": None, - "far": None, - "ssaa": 1, - "peel_layers": 8, - }) - self.rendering_options.update(rendering_options) - self.glctx = dr.RasterizeCudaContext(device=device) - self.device=device - - def render( - self, - mesh : Mesh, - extrinsics: torch.Tensor, - intrinsics: torch.Tensor, - envmap : Union[EnvMap, Dict[str, EnvMap]], - use_envmap_bg : bool = False, - transformation : Optional[torch.Tensor] = None - ) -> edict: - """ - Render the mesh. - - Args: - mesh : meshmodel - extrinsics (torch.Tensor): (4, 4) camera extrinsics - intrinsics (torch.Tensor): (3, 3) camera intrinsics - envmap (Union[EnvMap, Dict[str, EnvMap]]): environment map or a dictionary of environment maps - use_envmap_bg (bool): whether to use envmap as background - transformation (torch.Tensor): (4, 4) transformation matrix - - Returns: - edict based on return_types containing: - shaded (torch.Tensor): [3, H, W] shaded color image - normal (torch.Tensor): [3, H, W] normal image - base_color (torch.Tensor): [3, H, W] base color image - metallic (torch.Tensor): [H, W] metallic image - roughness (torch.Tensor): [H, W] roughness image - """ - if 'dr' not in globals(): - import nvdiffrast.torch as dr - - if not isinstance(envmap, dict): - envmap = {'' : envmap} - num_envmaps = len(envmap) - - resolution = self.rendering_options["resolution"] - near = self.rendering_options["near"] - far = self.rendering_options["far"] - ssaa = self.rendering_options["ssaa"] - - if mesh.vertices.shape[0] == 0 or mesh.faces.shape[0] == 0 or \ - not torch.isfinite(mesh.vertices).all() or \ - mesh.faces.max() >= mesh.vertices.shape[0] or mesh.faces.min() < 0: - if mesh.vertices.shape[0] > 0 and not torch.isfinite(mesh.vertices).all(): - print(f"[PbrMeshRenderer] WARNING: mesh vertices contain NaN/Inf, returning blank image.") - if mesh.faces.shape[0] > 0 and mesh.vertices.shape[0] > 0 and \ - (mesh.faces.max() >= mesh.vertices.shape[0] or mesh.faces.min() < 0): - print(f"[PbrMeshRenderer] WARNING: mesh faces contain out-of-bound indices " - f"(max={mesh.faces.max().item()}, min={mesh.faces.min().item()}, " - f"num_vertices={mesh.vertices.shape[0]}), returning blank image.") - out_dict = edict( - normal=torch.zeros((3, resolution, resolution), dtype=torch.float32, device=self.device), - mask=torch.zeros((resolution, resolution), dtype=torch.float32, device=self.device), - base_color=torch.zeros((3, resolution, resolution), dtype=torch.float32, device=self.device), - metallic=torch.zeros((resolution, resolution), dtype=torch.float32, device=self.device), - roughness=torch.zeros((resolution, resolution), dtype=torch.float32, device=self.device), - alpha=torch.zeros((resolution, resolution), dtype=torch.float32, device=self.device), - clay=torch.zeros((resolution, resolution), dtype=torch.float32, device=self.device), - ) - for i, k in enumerate(envmap.keys()): - shaded_key = f"shaded_{k}" if k != '' else "shaded" - out_dict[shaded_key] = torch.zeros((3, resolution, resolution), dtype=torch.float32, device=self.device) - return out_dict - - rays_o, rays_d = utils3d.torch.get_image_rays( - extrinsics, intrinsics, resolution * ssaa, resolution * ssaa - ) - - perspective = intrinsics_to_projection(intrinsics, near, far) - - full_proj = (perspective @ extrinsics).unsqueeze(0) - extrinsics = extrinsics.unsqueeze(0) - - vertices = mesh.vertices.unsqueeze(0) - vertices_orig = vertices.clone() - vertices_homo = torch.cat([vertices, torch.ones_like(vertices[..., :1])], dim=-1) - if transformation is not None: - vertices_homo = torch.bmm(vertices_homo, transformation.unsqueeze(0).transpose(-1, -2)) - vertices = vertices_homo[..., :3].contiguous() - vertices_camera = torch.bmm(vertices_homo, extrinsics.transpose(-1, -2)) - vertices_clip = torch.bmm(vertices_homo, full_proj.transpose(-1, -2)) - faces = mesh.faces - - v0 = vertices[0, mesh.faces[:, 0], :3] - v1 = vertices[0, mesh.faces[:, 1], :3] - v2 = vertices[0, mesh.faces[:, 2], :3] - e0 = v1 - v0 - e1 = v2 - v0 - face_normal = torch.cross(e0, e1, dim=1) - face_normal = F.normalize(face_normal, dim=1) - - out_dict = edict() - shaded = torch.zeros((num_envmaps, resolution * ssaa, resolution * ssaa, 3), dtype=torch.float32, device=self.device) - depth = torch.full((resolution * ssaa, resolution * ssaa, 1), 1e10, dtype=torch.float32, device=self.device) - normal = torch.zeros((resolution * ssaa, resolution * ssaa, 3), dtype=torch.float32, device=self.device) - max_w = torch.zeros((resolution * ssaa, resolution * ssaa, 1), dtype=torch.float32, device=self.device) - alpha = torch.zeros((resolution * ssaa, resolution * ssaa, 1), dtype=torch.float32, device=self.device) - with dr.DepthPeeler(self.glctx, vertices_clip, faces, (resolution * ssaa, resolution * ssaa)) as peeler: - for _ in range(self.rendering_options["peel_layers"]): - rast, rast_db = peeler.rasterize_next_layer() - - # Pos - pos = dr.interpolate(vertices, rast, faces)[0][0] - - # Depth - gb_depth = dr.interpolate(vertices_camera[..., 2:3].contiguous(), rast, faces)[0][0] - - # Normal - gb_normal = dr.interpolate(face_normal.unsqueeze(0), rast, torch.arange(face_normal.shape[0], dtype=torch.int, device=self.device).unsqueeze(1).repeat(1, 3).contiguous())[0][0] - gb_normal = torch.where( - torch.sum(gb_normal * (pos - rays_o), dim=-1, keepdim=True) > 0, - -gb_normal, - gb_normal - ) - gb_cam_normal = (extrinsics[..., :3, :3].reshape(1, 1, 3, 3) @ gb_normal.unsqueeze(-1)).squeeze(-1) - if _ == 0: - out_dict.normal = -gb_cam_normal * 0.5 + 0.5 - mask = (rast[0, ..., -1:] > 0).float() - out_dict.mask = mask - - # PBR attributes - if isinstance(mesh, MeshWithVoxel): - if 'grid_sample_3d' not in globals(): - from flex_gemm.ops.grid_sample import grid_sample_3d - mask = rast[..., -1:] > 0 - xyz = dr.interpolate(vertices_orig, rast, faces)[0] - xyz = ((xyz - mesh.origin) / mesh.voxel_size).reshape(1, -1, 3) - img = grid_sample_3d( - mesh.attrs, - torch.cat([torch.zeros_like(mesh.coords[..., :1]), mesh.coords], dim=-1), - mesh.voxel_shape, - xyz, - mode='trilinear' - ) - img = img.reshape(1, resolution * ssaa, resolution * ssaa, mesh.attrs.shape[-1]) * mask - gb_basecolor = img[0, ..., mesh.layout['base_color']] - gb_metallic = img[0, ..., mesh.layout['metallic']] - gb_roughness = img[0, ..., mesh.layout['roughness']] - gb_alpha = img[0, ..., mesh.layout['alpha']] - elif isinstance(mesh, MeshWithPbrMaterial): - tri_id = rast[0, :, :, -1:] - mask = tri_id > 0 - uv_coords = mesh.uv_coords.reshape(1, -1, 2) - texc, texd = dr.interpolate( - uv_coords, - rast, - torch.arange(mesh.uv_coords.shape[0] * 3, dtype=torch.int, device=self.device).reshape(-1, 3), - rast_db=rast_db, - diff_attrs='all' - ) - # Fix problematic texture coordinates - texc = torch.nan_to_num(texc, nan=0.0, posinf=1e3, neginf=-1e3) - texc = torch.clamp(texc, min=-1e3, max=1e3) - texd = torch.nan_to_num(texd, nan=0.0, posinf=1e3, neginf=-1e3) - texd = torch.clamp(texd, min=-1e3, max=1e3) - mid = mesh.material_ids[(tri_id - 1).long()] - gb_basecolor = torch.zeros((resolution * ssaa, resolution * ssaa, 3), dtype=torch.float32, device=self.device) - gb_metallic = torch.zeros((resolution * ssaa, resolution * ssaa, 1), dtype=torch.float32, device=self.device) - gb_roughness = torch.zeros((resolution * ssaa, resolution * ssaa, 1), dtype=torch.float32, device=self.device) - gb_alpha = torch.zeros((resolution * ssaa, resolution * ssaa, 1), dtype=torch.float32, device=self.device) - for id, mat in enumerate(mesh.materials): - mat_mask = (mid == id).float() * mask.float() - mat_texc = texc * mat_mask - mat_texd = texd * mat_mask - - if mat.base_color_texture is not None: - bc = dr.texture( - mat.base_color_texture.image.unsqueeze(0), - mat_texc, - mat_texd, - filter_mode='linear-mipmap-linear' if mat.base_color_texture.filter_mode == TextureFilterMode.LINEAR else 'nearest', - boundary_mode='clamp' if mat.base_color_texture.wrap_mode == TextureWrapMode.CLAMP_TO_EDGE else 'wrap' - )[0] - gb_basecolor += bc * mat.base_color_factor * mat_mask - else: - gb_basecolor += mat.base_color_factor * mat_mask - - if mat.metallic_texture is not None: - m = dr.texture( - mat.metallic_texture.image.unsqueeze(0), - mat_texc, - mat_texd, - filter_mode='linear-mipmap-linear' if mat.metallic_texture.filter_mode == TextureFilterMode.LINEAR else 'nearest', - boundary_mode='clamp' if mat.metallic_texture.wrap_mode == TextureWrapMode.CLAMP_TO_EDGE else 'wrap' - )[0] - gb_metallic += m * mat.metallic_factor * mat_mask - else: - gb_metallic += mat.metallic_factor * mat_mask - - if mat.roughness_texture is not None: - r = dr.texture( - mat.roughness_texture.image.unsqueeze(0), - mat_texc, - mat_texd, - filter_mode='linear-mipmap-linear' if mat.roughness_texture.filter_mode == TextureFilterMode.LINEAR else 'nearest', - boundary_mode='clamp' if mat.roughness_texture.wrap_mode == TextureWrapMode.CLAMP_TO_EDGE else 'wrap' - )[0] - gb_roughness += r * mat.roughness_factor * mat_mask - else: - gb_roughness += mat.roughness_factor * mat_mask - - if mat.alpha_mode == AlphaMode.OPAQUE: - gb_alpha += 1.0 * mat_mask - else: - if mat.alpha_texture is not None: - a = dr.texture( - mat.alpha_texture.image.unsqueeze(0), - mat_texc, - mat_texd, - filter_mode='linear-mipmap-linear' if mat.alpha_texture.filter_mode == TextureFilterMode.LINEAR else 'nearest', - boundary_mode='clamp' if mat.alpha_texture.wrap_mode == TextureWrapMode.CLAMP_TO_EDGE else 'wrap' - )[0] - if mat.alpha_mode == AlphaMode.MASK: - gb_alpha += (a * mat.alpha_factor > mat.alpha_cutoff).float() * mat_mask - elif mat.alpha_mode == AlphaMode.BLEND: - gb_alpha += a * mat.alpha_factor * mat_mask - else: - if mat.alpha_mode == AlphaMode.MASK: - gb_alpha += (mat.alpha_factor > mat.alpha_cutoff).float() * mat_mask - elif mat.alpha_mode == AlphaMode.BLEND: - gb_alpha += mat.alpha_factor * mat_mask - if _ == 0: - out_dict.base_color = gb_basecolor - out_dict.metallic = gb_metallic - out_dict.roughness = gb_roughness - out_dict.alpha = gb_alpha - - # Shading - #TODO - gb_basecolor = torch.clamp(gb_basecolor, 0.0, 1.0) ** 2.2 - # gb_basecolor = torch.clamp(gb_basecolor, 0.0, 1.0) - gb_metallic = torch.clamp(gb_metallic, 0.0, 1.0) - gb_roughness = torch.clamp(gb_roughness, 0.0, 1.0) - gb_alpha = torch.clamp(gb_alpha, 0.0, 1.0) - gb_orm = torch.cat([ - torch.zeros_like(gb_metallic), - gb_roughness, - gb_metallic, - ], dim=-1) - gb_shaded = torch.stack([ - e.shade( - pos.unsqueeze(0), - gb_normal.unsqueeze(0), - gb_basecolor.unsqueeze(0), - gb_orm.unsqueeze(0), - rays_o, - specular=True, - )[0] - for e in envmap.values() - ], dim=0) - - # Compositing - w = (1 - alpha) * gb_alpha - depth = torch.where(w > max_w, gb_depth, depth) - normal = torch.where(w > max_w, gb_cam_normal, normal) - max_w = torch.maximum(max_w, w) - shaded += w * gb_shaded - alpha += w - - # Ambient occulusion - f_occ = screen_space_ambient_occlusion( - depth, normal, perspective, intensity=1.5 - ) - shaded *= (1 - f_occ) - out_dict.clay = (1 - f_occ) - - # Background - if use_envmap_bg: - bg = torch.stack([e.sample(rays_d) for e in envmap.values()], dim=0) - shaded += (1 - alpha) * bg - - for i, k in enumerate(envmap.keys()): - shaded_key = f"shaded_{k}" if k != '' else "shaded" - out_dict[shaded_key] = shaded[i] - - # SSAA - for k in out_dict.keys(): - if ssaa > 1: - out_dict[k] = F.interpolate(out_dict[k].unsqueeze(0).permute(0, 3, 1, 2), (resolution, resolution), mode='bilinear', align_corners=False, antialias=True) - else: - out_dict[k] = out_dict[k].permute(2, 0, 1) - out_dict[k] = out_dict[k].squeeze() - - # Post processing: linear → sRGB (matches Blender's display transform) - for k in envmap.keys(): - shaded_key = f"shaded_{k}" if k != '' else "shaded" - out_dict[shaded_key] = linear_to_srgb(out_dict[shaded_key]) - - return out_dict +from typing import * +import torch +from easydict import EasyDict as edict +import numpy as np +import utils3d +from ..representations.mesh import Mesh, MeshWithVoxel, MeshWithPbrMaterial, TextureFilterMode, AlphaMode, TextureWrapMode +import torch.nn.functional as F + + +def cube_to_dir(s, x, y): + if s == 0: rx, ry, rz = torch.ones_like(x), -x, -y + elif s == 1: rx, ry, rz = -torch.ones_like(x), x, -y + elif s == 2: rx, ry, rz = x, y, torch.ones_like(x) + elif s == 3: rx, ry, rz = x, -y, -torch.ones_like(x) + elif s == 4: rx, ry, rz = x, torch.ones_like(x), -y + elif s == 5: rx, ry, rz = -x, -torch.ones_like(x), -y + return torch.stack((rx, ry, rz), dim=-1) + + +def latlong_to_cubemap(latlong_map, res): + if 'dr' not in globals(): + import nvdiffrast.torch as dr + cubemap = torch.zeros(6, res[0], res[1], latlong_map.shape[-1], dtype=torch.float32, device='cuda') + for s in range(6): + gy, gx = torch.meshgrid(torch.linspace(-1.0 + 1.0 / res[0], 1.0 - 1.0 / res[0], res[0], device='cuda'), + torch.linspace(-1.0 + 1.0 / res[1], 1.0 - 1.0 / res[1], res[1], device='cuda'), + indexing='ij') + v = F.normalize(cube_to_dir(s, gx, gy), dim=-1) + + tu = torch.atan2(v[..., 0:1], -v[..., 2:3]) / (2 * np.pi) + 0.5 + tv = torch.acos(torch.clamp(v[..., 1:2], min=-1, max=1)) / np.pi + texcoord = torch.cat((tu, tv), dim=-1) + + cubemap[s, ...] = dr.texture(latlong_map[None, ...], texcoord[None, ...], filter_mode='linear')[0] + return cubemap + + +class EnvMap: + def __init__(self, image: torch.Tensor): + self.image = image + + @property + def _backend(self): + if not hasattr(self, '_nvdiffrec_envlight'): + if 'EnvironmentLight' not in globals(): + from nvdiffrec_render.light import EnvironmentLight + cubemap = latlong_to_cubemap(self.image, [512, 512]) + self._nvdiffrec_envlight = EnvironmentLight(cubemap) + self._nvdiffrec_envlight.build_mips() + return self._nvdiffrec_envlight + + def shade(self, gb_pos, gb_normal, kd, ks, view_pos, specular=True): + return self._backend.shade(gb_pos, gb_normal, kd, ks, view_pos, specular) + + def sample(self, directions: torch.Tensor): + if 'dr' not in globals(): + import nvdiffrast.torch as dr + return dr.texture( + self._backend.base.unsqueeze(0), + directions.unsqueeze(0), + boundary_mode='cube', + )[0] + + +def intrinsics_to_projection( + intrinsics: torch.Tensor, + near: float, + far: float, + ) -> torch.Tensor: + """ + OpenCV intrinsics to OpenGL perspective matrix + + Args: + intrinsics (torch.Tensor): [3, 3] OpenCV intrinsics matrix + near (float): near plane to clip + far (float): far plane to clip + Returns: + (torch.Tensor): [4, 4] OpenGL perspective matrix + """ + fx, fy = intrinsics[0, 0], intrinsics[1, 1] + cx, cy = intrinsics[0, 2], intrinsics[1, 2] + ret = torch.zeros((4, 4), dtype=intrinsics.dtype, device=intrinsics.device) + ret[0, 0] = 2 * fx + ret[1, 1] = 2 * fy + ret[0, 2] = 2 * cx - 1 + ret[1, 2] = - 2 * cy + 1 + ret[2, 2] = (far + near) / (far - near) + ret[2, 3] = 2 * near * far / (near - far) + ret[3, 2] = 1. + return ret + + +def screen_space_ambient_occlusion( + depth: torch.Tensor, + normal: torch.Tensor, + perspective: torch.Tensor, + radius: float = 0.1, + bias: float = 1e-6, + samples: int = 64, + intensity: float = 1.0, +) -> torch.Tensor: + """ + Screen space ambient occlusion (SSAO) + + Args: + depth (torch.Tensor): [H, W, 1] depth image + normal (torch.Tensor): [H, W, 3] normal image + perspective (torch.Tensor): [4, 4] camera projection matrix + radius (float): radius of the SSAO kernel + bias (float): bias to avoid self-occlusion + samples (int): number of samples to use for the SSAO kernel + intensity (float): intensity of the SSAO effect + Returns: + (torch.Tensor): [H, W, 1] SSAO image + """ + device = depth.device + H, W, _ = depth.shape + + fx = perspective[0, 0] + fy = perspective[1, 1] + cx = perspective[0, 2] + cy = perspective[1, 2] + + y_grid, x_grid = torch.meshgrid( + (torch.arange(H, device=device) + 0.5) / H * 2 - 1, + (torch.arange(W, device=device) + 0.5) / W * 2 - 1, + indexing='ij' + ) + x_view = (x_grid.float() - cx) * depth[..., 0] / fx + y_view = (y_grid.float() - cy) * depth[..., 0] / fy + view_pos = torch.stack([x_view, y_view, depth[..., 0]], dim=-1) # [H, W, 3] + + depth_feat = depth.permute(2, 0, 1).unsqueeze(0) + occlusion = torch.zeros((H, W), device=device) + + # start sampling + for _ in range(samples): + # sample normal distribution, if inside, flip the sign + rnd_vec = torch.randn(H, W, 3, device=device) + rnd_vec = F.normalize(rnd_vec, p=2, dim=-1) + dot_val = torch.sum(rnd_vec * normal, dim=-1, keepdim=True) + sample_dir = torch.sign(dot_val) * rnd_vec + scale = torch.rand(H, W, 1, device=device) + scale = scale * scale + sample_pos = view_pos + sample_dir * radius * scale + sample_z = sample_pos[..., 2] + + # project to screen space + z_safe = torch.clamp(sample_pos[..., 2], min=1e-5) + proj_u = (sample_pos[..., 0] * fx / z_safe) + cx + proj_v = (sample_pos[..., 1] * fy / z_safe) + cy + grid = torch.stack([proj_u, proj_v], dim=-1).unsqueeze(0) + geo_z = F.grid_sample(depth_feat, grid, mode='nearest', padding_mode='border').squeeze() + range_check = torch.abs(geo_z - sample_z) < radius + is_occluded = (geo_z <= sample_z - bias) & range_check + occlusion += is_occluded.float() + + f_occ = occlusion / samples * intensity + f_occ = torch.clamp(f_occ, 0.0, 1.0) + + return f_occ.unsqueeze(-1) + + +def aces_tonemapping(x: torch.Tensor) -> torch.Tensor: + """ + Applies ACES tone mapping curve to an HDR image tensor. + Input: x - HDR tensor, shape (..., 3), range [0, +inf) + Output: LDR tensor, same shape, range [0, 1] + + NOTE: This causes bright base_color objects to over-expose (appear white) + compared to Blender's standard sRGB display transform. Use linear_to_srgb() + instead for renders_cond alignment. See pbr_color_diff_debug_guide.md #5. + """ + a = 2.51 + b = 0.03 + c = 2.43 + d = 0.59 + e = 0.14 + + # Apply the ACES fitted curve + mapped = (x * (a * x + b)) / (x * (c * x + d) + e) + + # Clamp to [0, 1] for display or saving + return torch.clamp(mapped, 0.0, 1.0) + + +def gamma_correction(x: torch.Tensor, gamma: float = 2.2) -> torch.Tensor: + """ + Applies gamma correction to an HDR image tensor. + """ + return torch.clamp(x ** (1.0 / gamma), 0.0, 1.0) + + +def linear_to_srgb(x: torch.Tensor) -> torch.Tensor: + """ + Standard linear-to-sRGB conversion matching Blender's sRGB display transform. + + Applies the official sRGB EOTF inverse: + - For values <= 0.0031308: sRGB = 12.92 * linear + - For values > 0.0031308: sRGB = 1.055 * linear^(1/2.4) - 0.055 + + Input: x - linear HDR tensor, clamped to [0, +inf) + Output: sRGB tensor, range [0, 1] + """ + x = torch.clamp(x, 0.0) + low = 12.92 * x + high = 1.055 * torch.pow(torch.clamp(x, min=0.0031308), 1.0 / 2.4) - 0.055 + return torch.clamp(torch.where(x <= 0.0031308, low, high), 0.0, 1.0) + + +class PbrMeshRenderer: + """ + Renderer for the PBR mesh. + + Args: + rendering_options (dict): Rendering options. + """ + def __init__(self, rendering_options={}, device='cuda'): + if 'dr' not in globals(): + import nvdiffrast.torch as dr + + self.rendering_options = edict({ + "resolution": None, + "near": None, + "far": None, + "ssaa": 1, + "peel_layers": 8, + }) + self.rendering_options.update(rendering_options) + self.glctx = dr.RasterizeCudaContext(device=device) + self.device=device + + def render( + self, + mesh : Mesh, + extrinsics: torch.Tensor, + intrinsics: torch.Tensor, + envmap : Union[EnvMap, Dict[str, EnvMap]], + use_envmap_bg : bool = False, + transformation : Optional[torch.Tensor] = None + ) -> edict: + """ + Render the mesh. + + Args: + mesh : meshmodel + extrinsics (torch.Tensor): (4, 4) camera extrinsics + intrinsics (torch.Tensor): (3, 3) camera intrinsics + envmap (Union[EnvMap, Dict[str, EnvMap]]): environment map or a dictionary of environment maps + use_envmap_bg (bool): whether to use envmap as background + transformation (torch.Tensor): (4, 4) transformation matrix + + Returns: + edict based on return_types containing: + shaded (torch.Tensor): [3, H, W] shaded color image + normal (torch.Tensor): [3, H, W] normal image + base_color (torch.Tensor): [3, H, W] base color image + metallic (torch.Tensor): [H, W] metallic image + roughness (torch.Tensor): [H, W] roughness image + """ + if 'dr' not in globals(): + import nvdiffrast.torch as dr + + if not isinstance(envmap, dict): + envmap = {'' : envmap} + num_envmaps = len(envmap) + + resolution = self.rendering_options["resolution"] + near = self.rendering_options["near"] + far = self.rendering_options["far"] + ssaa = self.rendering_options["ssaa"] + + if mesh.vertices.shape[0] == 0 or mesh.faces.shape[0] == 0 or \ + not torch.isfinite(mesh.vertices).all() or \ + mesh.faces.max() >= mesh.vertices.shape[0] or mesh.faces.min() < 0: + if mesh.vertices.shape[0] > 0 and not torch.isfinite(mesh.vertices).all(): + print(f"[PbrMeshRenderer] WARNING: mesh vertices contain NaN/Inf, returning blank image.") + if mesh.faces.shape[0] > 0 and mesh.vertices.shape[0] > 0 and \ + (mesh.faces.max() >= mesh.vertices.shape[0] or mesh.faces.min() < 0): + print(f"[PbrMeshRenderer] WARNING: mesh faces contain out-of-bound indices " + f"(max={mesh.faces.max().item()}, min={mesh.faces.min().item()}, " + f"num_vertices={mesh.vertices.shape[0]}), returning blank image.") + out_dict = edict( + normal=torch.zeros((3, resolution, resolution), dtype=torch.float32, device=self.device), + mask=torch.zeros((resolution, resolution), dtype=torch.float32, device=self.device), + base_color=torch.zeros((3, resolution, resolution), dtype=torch.float32, device=self.device), + metallic=torch.zeros((resolution, resolution), dtype=torch.float32, device=self.device), + roughness=torch.zeros((resolution, resolution), dtype=torch.float32, device=self.device), + alpha=torch.zeros((resolution, resolution), dtype=torch.float32, device=self.device), + clay=torch.zeros((resolution, resolution), dtype=torch.float32, device=self.device), + ) + for i, k in enumerate(envmap.keys()): + shaded_key = f"shaded_{k}" if k != '' else "shaded" + out_dict[shaded_key] = torch.zeros((3, resolution, resolution), dtype=torch.float32, device=self.device) + return out_dict + + rays_o, rays_d = utils3d.torch.get_image_rays( + extrinsics, intrinsics, resolution * ssaa, resolution * ssaa + ) + + perspective = intrinsics_to_projection(intrinsics, near, far) + + full_proj = (perspective @ extrinsics).unsqueeze(0) + extrinsics = extrinsics.unsqueeze(0) + + vertices = mesh.vertices.unsqueeze(0) + vertices_orig = vertices.clone() + vertices_homo = torch.cat([vertices, torch.ones_like(vertices[..., :1])], dim=-1) + if transformation is not None: + vertices_homo = torch.bmm(vertices_homo, transformation.unsqueeze(0).transpose(-1, -2)) + vertices = vertices_homo[..., :3].contiguous() + vertices_camera = torch.bmm(vertices_homo, extrinsics.transpose(-1, -2)) + vertices_clip = torch.bmm(vertices_homo, full_proj.transpose(-1, -2)) + faces = mesh.faces + + v0 = vertices[0, mesh.faces[:, 0], :3] + v1 = vertices[0, mesh.faces[:, 1], :3] + v2 = vertices[0, mesh.faces[:, 2], :3] + e0 = v1 - v0 + e1 = v2 - v0 + face_normal = torch.cross(e0, e1, dim=1) + face_normal = F.normalize(face_normal, dim=1) + + out_dict = edict() + shaded = torch.zeros((num_envmaps, resolution * ssaa, resolution * ssaa, 3), dtype=torch.float32, device=self.device) + depth = torch.full((resolution * ssaa, resolution * ssaa, 1), 1e10, dtype=torch.float32, device=self.device) + normal = torch.zeros((resolution * ssaa, resolution * ssaa, 3), dtype=torch.float32, device=self.device) + max_w = torch.zeros((resolution * ssaa, resolution * ssaa, 1), dtype=torch.float32, device=self.device) + alpha = torch.zeros((resolution * ssaa, resolution * ssaa, 1), dtype=torch.float32, device=self.device) + with dr.DepthPeeler(self.glctx, vertices_clip, faces, (resolution * ssaa, resolution * ssaa)) as peeler: + for _ in range(self.rendering_options["peel_layers"]): + rast, rast_db = peeler.rasterize_next_layer() + + # Pos + pos = dr.interpolate(vertices, rast, faces)[0][0] + + # Depth + gb_depth = dr.interpolate(vertices_camera[..., 2:3].contiguous(), rast, faces)[0][0] + + # Normal + gb_normal = dr.interpolate(face_normal.unsqueeze(0), rast, torch.arange(face_normal.shape[0], dtype=torch.int, device=self.device).unsqueeze(1).repeat(1, 3).contiguous())[0][0] + gb_normal = torch.where( + torch.sum(gb_normal * (pos - rays_o), dim=-1, keepdim=True) > 0, + -gb_normal, + gb_normal + ) + gb_cam_normal = (extrinsics[..., :3, :3].reshape(1, 1, 3, 3) @ gb_normal.unsqueeze(-1)).squeeze(-1) + if _ == 0: + out_dict.normal = -gb_cam_normal * 0.5 + 0.5 + mask = (rast[0, ..., -1:] > 0).float() + out_dict.mask = mask + + # PBR attributes + if isinstance(mesh, MeshWithVoxel): + if 'grid_sample_3d' not in globals(): + from flex_gemm.ops.grid_sample import grid_sample_3d + mask = rast[..., -1:] > 0 + xyz = dr.interpolate(vertices_orig, rast, faces)[0] + xyz = ((xyz - mesh.origin) / mesh.voxel_size).reshape(1, -1, 3) + img = grid_sample_3d( + mesh.attrs, + torch.cat([torch.zeros_like(mesh.coords[..., :1]), mesh.coords], dim=-1), + mesh.voxel_shape, + xyz, + mode='trilinear' + ) + img = img.reshape(1, resolution * ssaa, resolution * ssaa, mesh.attrs.shape[-1]) * mask + gb_basecolor = img[0, ..., mesh.layout['base_color']] + gb_metallic = img[0, ..., mesh.layout['metallic']] + gb_roughness = img[0, ..., mesh.layout['roughness']] + gb_alpha = img[0, ..., mesh.layout['alpha']] + elif isinstance(mesh, MeshWithPbrMaterial): + tri_id = rast[0, :, :, -1:] + mask = tri_id > 0 + uv_coords = mesh.uv_coords.reshape(1, -1, 2) + texc, texd = dr.interpolate( + uv_coords, + rast, + torch.arange(mesh.uv_coords.shape[0] * 3, dtype=torch.int, device=self.device).reshape(-1, 3), + rast_db=rast_db, + diff_attrs='all' + ) + # Fix problematic texture coordinates + texc = torch.nan_to_num(texc, nan=0.0, posinf=1e3, neginf=-1e3) + texc = torch.clamp(texc, min=-1e3, max=1e3) + texd = torch.nan_to_num(texd, nan=0.0, posinf=1e3, neginf=-1e3) + texd = torch.clamp(texd, min=-1e3, max=1e3) + mid = mesh.material_ids[(tri_id - 1).long()] + gb_basecolor = torch.zeros((resolution * ssaa, resolution * ssaa, 3), dtype=torch.float32, device=self.device) + gb_metallic = torch.zeros((resolution * ssaa, resolution * ssaa, 1), dtype=torch.float32, device=self.device) + gb_roughness = torch.zeros((resolution * ssaa, resolution * ssaa, 1), dtype=torch.float32, device=self.device) + gb_alpha = torch.zeros((resolution * ssaa, resolution * ssaa, 1), dtype=torch.float32, device=self.device) + for id, mat in enumerate(mesh.materials): + mat_mask = (mid == id).float() * mask.float() + mat_texc = texc * mat_mask + mat_texd = texd * mat_mask + + if mat.base_color_texture is not None: + bc = dr.texture( + mat.base_color_texture.image.unsqueeze(0), + mat_texc, + mat_texd, + filter_mode='linear-mipmap-linear' if mat.base_color_texture.filter_mode == TextureFilterMode.LINEAR else 'nearest', + boundary_mode='clamp' if mat.base_color_texture.wrap_mode == TextureWrapMode.CLAMP_TO_EDGE else 'wrap' + )[0] + gb_basecolor += bc * mat.base_color_factor * mat_mask + else: + gb_basecolor += mat.base_color_factor * mat_mask + + if mat.metallic_texture is not None: + m = dr.texture( + mat.metallic_texture.image.unsqueeze(0), + mat_texc, + mat_texd, + filter_mode='linear-mipmap-linear' if mat.metallic_texture.filter_mode == TextureFilterMode.LINEAR else 'nearest', + boundary_mode='clamp' if mat.metallic_texture.wrap_mode == TextureWrapMode.CLAMP_TO_EDGE else 'wrap' + )[0] + gb_metallic += m * mat.metallic_factor * mat_mask + else: + gb_metallic += mat.metallic_factor * mat_mask + + if mat.roughness_texture is not None: + r = dr.texture( + mat.roughness_texture.image.unsqueeze(0), + mat_texc, + mat_texd, + filter_mode='linear-mipmap-linear' if mat.roughness_texture.filter_mode == TextureFilterMode.LINEAR else 'nearest', + boundary_mode='clamp' if mat.roughness_texture.wrap_mode == TextureWrapMode.CLAMP_TO_EDGE else 'wrap' + )[0] + gb_roughness += r * mat.roughness_factor * mat_mask + else: + gb_roughness += mat.roughness_factor * mat_mask + + if mat.alpha_mode == AlphaMode.OPAQUE: + gb_alpha += 1.0 * mat_mask + else: + if mat.alpha_texture is not None: + a = dr.texture( + mat.alpha_texture.image.unsqueeze(0), + mat_texc, + mat_texd, + filter_mode='linear-mipmap-linear' if mat.alpha_texture.filter_mode == TextureFilterMode.LINEAR else 'nearest', + boundary_mode='clamp' if mat.alpha_texture.wrap_mode == TextureWrapMode.CLAMP_TO_EDGE else 'wrap' + )[0] + if mat.alpha_mode == AlphaMode.MASK: + gb_alpha += (a * mat.alpha_factor > mat.alpha_cutoff).float() * mat_mask + elif mat.alpha_mode == AlphaMode.BLEND: + gb_alpha += a * mat.alpha_factor * mat_mask + else: + if mat.alpha_mode == AlphaMode.MASK: + gb_alpha += (mat.alpha_factor > mat.alpha_cutoff).float() * mat_mask + elif mat.alpha_mode == AlphaMode.BLEND: + gb_alpha += mat.alpha_factor * mat_mask + if _ == 0: + out_dict.base_color = gb_basecolor + out_dict.metallic = gb_metallic + out_dict.roughness = gb_roughness + out_dict.alpha = gb_alpha + + # Shading + gb_basecolor = torch.clamp(gb_basecolor, 0.0, 1.0) ** 2.2 + # gb_basecolor = torch.clamp(gb_basecolor, 0.0, 1.0) + gb_metallic = torch.clamp(gb_metallic, 0.0, 1.0) + gb_roughness = torch.clamp(gb_roughness, 0.0, 1.0) + gb_alpha = torch.clamp(gb_alpha, 0.0, 1.0) + gb_orm = torch.cat([ + torch.zeros_like(gb_metallic), + gb_roughness, + gb_metallic, + ], dim=-1) + gb_shaded = torch.stack([ + e.shade( + pos.unsqueeze(0), + gb_normal.unsqueeze(0), + gb_basecolor.unsqueeze(0), + gb_orm.unsqueeze(0), + rays_o, + specular=True, + )[0] + for e in envmap.values() + ], dim=0) + + # Compositing + w = (1 - alpha) * gb_alpha + depth = torch.where(w > max_w, gb_depth, depth) + normal = torch.where(w > max_w, gb_cam_normal, normal) + max_w = torch.maximum(max_w, w) + shaded += w * gb_shaded + alpha += w + + # Ambient occulusion + f_occ = screen_space_ambient_occlusion( + depth, normal, perspective, intensity=1.5 + ) + shaded *= (1 - f_occ) + out_dict.clay = (1 - f_occ) + + # Background + if use_envmap_bg: + bg = torch.stack([e.sample(rays_d) for e in envmap.values()], dim=0) + shaded += (1 - alpha) * bg + + for i, k in enumerate(envmap.keys()): + shaded_key = f"shaded_{k}" if k != '' else "shaded" + out_dict[shaded_key] = shaded[i] + + # SSAA + for k in out_dict.keys(): + if ssaa > 1: + out_dict[k] = F.interpolate(out_dict[k].unsqueeze(0).permute(0, 3, 1, 2), (resolution, resolution), mode='bilinear', align_corners=False, antialias=True) + else: + out_dict[k] = out_dict[k].permute(2, 0, 1) + out_dict[k] = out_dict[k].squeeze() + + # Post processing: linear → sRGB (matches Blender's display transform) + for k in envmap.keys(): + shaded_key = f"shaded_{k}" if k != '' else "shaded" + out_dict[shaded_key] = linear_to_srgb(out_dict[shaded_key]) + + return out_dict diff --git a/supratex/representations/mesh/base.py b/supratex/representations/mesh/base.py index 23bcf6d14ebac5b08893579fb7dd834a11c0d9a1..2b3ea5109d0da4d664a7360c1bfaec615fec7a6d 100644 --- a/supratex/representations/mesh/base.py +++ b/supratex/representations/mesh/base.py @@ -1,239 +1,239 @@ -from typing import * -import torch -from ..voxel import Voxel -import cumesh -# NOTE: flex_gemm is imported lazily in query_attrs (matching the renderers' pattern), NOT here. -# This module sits in the DataLoader workers' import closure (dataset -> MeshWithVoxel), and -# flex_gemm's triton @autotune decorators query the GPU driver AT IMPORT — in a spawned CPU-only -# worker that intermittently dies with "RuntimeError: 0 active drivers", killing the rank and -# wedging all other ranks in NCCL collectives (observed 2026-07-02 on the 64-rank 512-all run). - - -class Mesh: - def __init__(self, - vertices, - faces, - vertex_attrs=None - ): - self.vertices = vertices.float() - self.faces = faces.int() - self.vertex_attrs = vertex_attrs - - @property - def device(self): - return self.vertices.device - - def to(self, device, non_blocking=False): - return Mesh( - self.vertices.to(device, non_blocking=non_blocking), - self.faces.to(device, non_blocking=non_blocking), - self.vertex_attrs.to(device, non_blocking=non_blocking) if self.vertex_attrs is not None else None, - ) - - def cuda(self, non_blocking=False): - return self.to('cuda', non_blocking=non_blocking) - - def cpu(self): - return self.to('cpu') - - def fill_holes(self, max_hole_perimeter=3e-2): - vertices = self.vertices.clone().cuda().contiguous() - faces = self.faces.clone().cuda().contiguous() - - mesh = cumesh.CuMesh() - mesh.init(vertices, faces) - mesh.get_edges() - mesh.get_boundary_info() - if mesh.num_boundaries == 0: - return - mesh.get_vertex_edge_adjacency() - mesh.get_vertex_boundary_adjacency() - mesh.get_manifold_boundary_adjacency() - mesh.read_manifold_boundary_adjacency() - mesh.get_boundary_connected_components() - mesh.get_boundary_loops() - if mesh.num_boundary_loops == 0: - return - mesh.fill_holes(max_hole_perimeter=max_hole_perimeter) - new_vertices, new_faces = mesh.read() - - self.vertices = new_vertices.to(self.device) - self.faces = new_faces.to(self.device) - - def remove_faces(self, face_mask: torch.Tensor): - vertices = self.vertices.clone().cuda().contiguous() - faces = self.faces.clone().cuda().contiguous() - - mesh = cumesh.CuMesh() - mesh.init(vertices, faces) - mesh.remove_faces(face_mask) - new_vertices, new_faces = mesh.read() - - self.vertices = new_vertices.to(self.device) - self.faces = new_faces.to(self.device) - - def simplify(self, target=1000000, verbose: bool=False, options: dict={}): - vertices = self.vertices.clone().cuda().contiguous() - faces = self.faces.clone().cuda().contiguous() - - mesh = cumesh.CuMesh() - mesh.init(vertices, faces) - mesh.simplify(target, verbose=verbose, options=options) - new_vertices, new_faces = mesh.read() - - self.vertices = new_vertices.to(self.device) - self.faces = new_faces.to(self.device) - - -class TextureFilterMode: - CLOSEST = 0 - LINEAR = 1 - - -class TextureWrapMode: - CLAMP_TO_EDGE = 0 - REPEAT = 1 - MIRRORED_REPEAT = 2 - - -class AlphaMode: - OPAQUE = 0 - MASK = 1 - BLEND = 2 - - -class Texture: - def __init__( - self, - image: torch.Tensor, - filter_mode: TextureFilterMode = TextureFilterMode.LINEAR, - wrap_mode: TextureWrapMode = TextureWrapMode.REPEAT - ): - self.image = image - self.filter_mode = filter_mode - self.wrap_mode = wrap_mode - - def to(self, device, non_blocking=False): - return Texture( - self.image.to(device, non_blocking=non_blocking), - self.filter_mode, - self.wrap_mode, - ) - - -class PbrMaterial: - def __init__( - self, - base_color_texture: Optional[Texture] = None, - base_color_factor: Union[torch.Tensor, List[float]] = [1.0, 1.0, 1.0], - metallic_texture: Optional[Texture] = None, - metallic_factor: float = 1.0, - roughness_texture: Optional[Texture] = None, - roughness_factor: float = 1.0, - alpha_texture: Optional[Texture] = None, - alpha_factor: float = 1.0, - alpha_mode: AlphaMode = AlphaMode.OPAQUE, - alpha_cutoff: float = 0.5, - ): - self.base_color_texture = base_color_texture - self.base_color_factor = torch.tensor(base_color_factor, dtype=torch.float32)[:3] - self.metallic_texture = metallic_texture - self.metallic_factor = metallic_factor - self.roughness_texture = roughness_texture - self.roughness_factor = roughness_factor - self.alpha_texture = alpha_texture - self.alpha_factor = alpha_factor - self.alpha_mode = alpha_mode - self.alpha_cutoff = alpha_cutoff - - def to(self, device, non_blocking=False): - return PbrMaterial( - base_color_texture=self.base_color_texture.to(device, non_blocking=non_blocking) if self.base_color_texture is not None else None, - base_color_factor=self.base_color_factor.to(device, non_blocking=non_blocking), - metallic_texture=self.metallic_texture.to(device, non_blocking=non_blocking) if self.metallic_texture is not None else None, - metallic_factor=self.metallic_factor, - roughness_texture=self.roughness_texture.to(device, non_blocking=non_blocking) if self.roughness_texture is not None else None, - roughness_factor=self.roughness_factor, - alpha_texture=self.alpha_texture.to(device, non_blocking=non_blocking) if self.alpha_texture is not None else None, - alpha_factor=self.alpha_factor, - alpha_mode=self.alpha_mode, - alpha_cutoff=self.alpha_cutoff, - ) - - -class MeshWithPbrMaterial(Mesh): - def __init__(self, - vertices, - faces, - material_ids, - uv_coords, - materials: List[PbrMaterial], - ): - self.vertices = vertices.float() - self.faces = faces.int() - self.material_ids = material_ids # [M] - self.uv_coords = uv_coords # [M, 3, 2] - self.materials = materials - self.layout = { - 'base_color': slice(0, 3), - 'metallic': slice(3, 4), - 'roughness': slice(4, 5), - 'alpha': slice(5, 6), - } - - def to(self, device, non_blocking=False): - return MeshWithPbrMaterial( - self.vertices.to(device, non_blocking=non_blocking), - self.faces.to(device, non_blocking=non_blocking), - self.material_ids.to(device, non_blocking=non_blocking), - self.uv_coords.to(device, non_blocking=non_blocking), - [material.to(device, non_blocking=non_blocking) for material in self.materials], - ) - - -class MeshWithVoxel(Mesh, Voxel): - def __init__(self, - vertices: torch.Tensor, - faces: torch.Tensor, - origin: list, - voxel_size: float, - coords: torch.Tensor, - attrs: torch.Tensor, - voxel_shape: torch.Size, - layout: Dict = {}, - ): - self.vertices = vertices.float() - self.faces = faces.int() - self.origin = torch.tensor(origin, dtype=torch.float32, device=self.device) - self.voxel_size = voxel_size - self.coords = coords - self.attrs = attrs - self.voxel_shape = voxel_shape - self.layout = layout - - def to(self, device, non_blocking=False): - return MeshWithVoxel( - self.vertices.to(device, non_blocking=non_blocking), - self.faces.to(device, non_blocking=non_blocking), - self.origin.tolist(), - self.voxel_size, - self.coords.to(device, non_blocking=non_blocking), - self.attrs.to(device, non_blocking=non_blocking), - self.voxel_shape, - self.layout, - ) - - def query_attrs(self, xyz): - from flex_gemm.ops.grid_sample import grid_sample_3d - grid = ((xyz - self.origin) / self.voxel_size).reshape(1, -1, 3) - vertex_attrs = grid_sample_3d( - self.attrs, - torch.cat([torch.zeros_like(self.coords[..., :1]), self.coords], dim=-1), - self.voxel_shape, - grid, - mode='trilinear' - )[0] - return vertex_attrs - - def query_vertex_attrs(self): - return self.query_attrs(self.vertices) +from typing import * +import torch +from ..voxel import Voxel +import cumesh +# NOTE: flex_gemm is imported lazily in query_attrs (matching the renderers' pattern), NOT here. +# This module sits in the DataLoader workers' import closure (dataset -> MeshWithVoxel), and +# flex_gemm's triton @autotune decorators query the GPU driver AT IMPORT — in a spawned CPU-only +# worker that intermittently dies with "RuntimeError: 0 active drivers", killing the rank and +# blocking all other ranks in NCCL collectives. + + +class Mesh: + def __init__(self, + vertices, + faces, + vertex_attrs=None + ): + self.vertices = vertices.float() + self.faces = faces.int() + self.vertex_attrs = vertex_attrs + + @property + def device(self): + return self.vertices.device + + def to(self, device, non_blocking=False): + return Mesh( + self.vertices.to(device, non_blocking=non_blocking), + self.faces.to(device, non_blocking=non_blocking), + self.vertex_attrs.to(device, non_blocking=non_blocking) if self.vertex_attrs is not None else None, + ) + + def cuda(self, non_blocking=False): + return self.to('cuda', non_blocking=non_blocking) + + def cpu(self): + return self.to('cpu') + + def fill_holes(self, max_hole_perimeter=3e-2): + vertices = self.vertices.clone().cuda().contiguous() + faces = self.faces.clone().cuda().contiguous() + + mesh = cumesh.CuMesh() + mesh.init(vertices, faces) + mesh.get_edges() + mesh.get_boundary_info() + if mesh.num_boundaries == 0: + return + mesh.get_vertex_edge_adjacency() + mesh.get_vertex_boundary_adjacency() + mesh.get_manifold_boundary_adjacency() + mesh.read_manifold_boundary_adjacency() + mesh.get_boundary_connected_components() + mesh.get_boundary_loops() + if mesh.num_boundary_loops == 0: + return + mesh.fill_holes(max_hole_perimeter=max_hole_perimeter) + new_vertices, new_faces = mesh.read() + + self.vertices = new_vertices.to(self.device) + self.faces = new_faces.to(self.device) + + def remove_faces(self, face_mask: torch.Tensor): + vertices = self.vertices.clone().cuda().contiguous() + faces = self.faces.clone().cuda().contiguous() + + mesh = cumesh.CuMesh() + mesh.init(vertices, faces) + mesh.remove_faces(face_mask) + new_vertices, new_faces = mesh.read() + + self.vertices = new_vertices.to(self.device) + self.faces = new_faces.to(self.device) + + def simplify(self, target=1000000, verbose: bool=False, options: dict={}): + vertices = self.vertices.clone().cuda().contiguous() + faces = self.faces.clone().cuda().contiguous() + + mesh = cumesh.CuMesh() + mesh.init(vertices, faces) + mesh.simplify(target, verbose=verbose, options=options) + new_vertices, new_faces = mesh.read() + + self.vertices = new_vertices.to(self.device) + self.faces = new_faces.to(self.device) + + +class TextureFilterMode: + CLOSEST = 0 + LINEAR = 1 + + +class TextureWrapMode: + CLAMP_TO_EDGE = 0 + REPEAT = 1 + MIRRORED_REPEAT = 2 + + +class AlphaMode: + OPAQUE = 0 + MASK = 1 + BLEND = 2 + + +class Texture: + def __init__( + self, + image: torch.Tensor, + filter_mode: TextureFilterMode = TextureFilterMode.LINEAR, + wrap_mode: TextureWrapMode = TextureWrapMode.REPEAT + ): + self.image = image + self.filter_mode = filter_mode + self.wrap_mode = wrap_mode + + def to(self, device, non_blocking=False): + return Texture( + self.image.to(device, non_blocking=non_blocking), + self.filter_mode, + self.wrap_mode, + ) + + +class PbrMaterial: + def __init__( + self, + base_color_texture: Optional[Texture] = None, + base_color_factor: Union[torch.Tensor, List[float]] = [1.0, 1.0, 1.0], + metallic_texture: Optional[Texture] = None, + metallic_factor: float = 1.0, + roughness_texture: Optional[Texture] = None, + roughness_factor: float = 1.0, + alpha_texture: Optional[Texture] = None, + alpha_factor: float = 1.0, + alpha_mode: AlphaMode = AlphaMode.OPAQUE, + alpha_cutoff: float = 0.5, + ): + self.base_color_texture = base_color_texture + self.base_color_factor = torch.tensor(base_color_factor, dtype=torch.float32)[:3] + self.metallic_texture = metallic_texture + self.metallic_factor = metallic_factor + self.roughness_texture = roughness_texture + self.roughness_factor = roughness_factor + self.alpha_texture = alpha_texture + self.alpha_factor = alpha_factor + self.alpha_mode = alpha_mode + self.alpha_cutoff = alpha_cutoff + + def to(self, device, non_blocking=False): + return PbrMaterial( + base_color_texture=self.base_color_texture.to(device, non_blocking=non_blocking) if self.base_color_texture is not None else None, + base_color_factor=self.base_color_factor.to(device, non_blocking=non_blocking), + metallic_texture=self.metallic_texture.to(device, non_blocking=non_blocking) if self.metallic_texture is not None else None, + metallic_factor=self.metallic_factor, + roughness_texture=self.roughness_texture.to(device, non_blocking=non_blocking) if self.roughness_texture is not None else None, + roughness_factor=self.roughness_factor, + alpha_texture=self.alpha_texture.to(device, non_blocking=non_blocking) if self.alpha_texture is not None else None, + alpha_factor=self.alpha_factor, + alpha_mode=self.alpha_mode, + alpha_cutoff=self.alpha_cutoff, + ) + + +class MeshWithPbrMaterial(Mesh): + def __init__(self, + vertices, + faces, + material_ids, + uv_coords, + materials: List[PbrMaterial], + ): + self.vertices = vertices.float() + self.faces = faces.int() + self.material_ids = material_ids # [M] + self.uv_coords = uv_coords # [M, 3, 2] + self.materials = materials + self.layout = { + 'base_color': slice(0, 3), + 'metallic': slice(3, 4), + 'roughness': slice(4, 5), + 'alpha': slice(5, 6), + } + + def to(self, device, non_blocking=False): + return MeshWithPbrMaterial( + self.vertices.to(device, non_blocking=non_blocking), + self.faces.to(device, non_blocking=non_blocking), + self.material_ids.to(device, non_blocking=non_blocking), + self.uv_coords.to(device, non_blocking=non_blocking), + [material.to(device, non_blocking=non_blocking) for material in self.materials], + ) + + +class MeshWithVoxel(Mesh, Voxel): + def __init__(self, + vertices: torch.Tensor, + faces: torch.Tensor, + origin: list, + voxel_size: float, + coords: torch.Tensor, + attrs: torch.Tensor, + voxel_shape: torch.Size, + layout: Dict = {}, + ): + self.vertices = vertices.float() + self.faces = faces.int() + self.origin = torch.tensor(origin, dtype=torch.float32, device=self.device) + self.voxel_size = voxel_size + self.coords = coords + self.attrs = attrs + self.voxel_shape = voxel_shape + self.layout = layout + + def to(self, device, non_blocking=False): + return MeshWithVoxel( + self.vertices.to(device, non_blocking=non_blocking), + self.faces.to(device, non_blocking=non_blocking), + self.origin.tolist(), + self.voxel_size, + self.coords.to(device, non_blocking=non_blocking), + self.attrs.to(device, non_blocking=non_blocking), + self.voxel_shape, + self.layout, + ) + + def query_attrs(self, xyz): + from flex_gemm.ops.grid_sample import grid_sample_3d + grid = ((xyz - self.origin) / self.voxel_size).reshape(1, -1, 3) + vertex_attrs = grid_sample_3d( + self.attrs, + torch.cat([torch.zeros_like(self.coords[..., :1]), self.coords], dim=-1), + self.voxel_shape, + grid, + mode='trilinear' + )[0] + return vertex_attrs + + def query_vertex_attrs(self): + return self.query_attrs(self.vertices) diff --git a/supratex/weights.py b/supratex/weights.py index f001d992846b65e190221dee202904145989a59f..9e575553ac6718f176fd58e30c19d5896839e9cd 100644 --- a/supratex/weights.py +++ b/supratex/weights.py @@ -1,114 +1,120 @@ -"""Resolve model directories from local storage or the Hugging Face Hub.""" - -from pathlib import Path - - -COMPLETION_ALLOW_PATTERNS = [ - "config.json", "partial_normalization.json", "*.safetensors", "*.safetensors.index.json", - "configs/*.json", "weights/material_completion.pt", "weights/sparse_encoders/*.json", - "weights/scvae/*.json", "vae/*.json", "vae/ckpts/*.json", -] - - -def resolve_completion_artifacts( - root: str | Path | None = None, - *, - model_config: str | Path | None = None, - checkpoint: str | Path | None = None, - trellis_root: str | Path | None = None, -) -> tuple[Path, Path, Path]: - """Resolve and preflight a local completion release plus its sparse VAEs. - - Standalone releases contain root config/normalization and safe weights; - their separately obtained VAEs live at ``trellis_root`` (or ``root/vae``). - The original configs/weights bundle layout remains supported unchanged. - """ - if root is not None: - root = Path(root).expanduser().resolve() - if not root.is_dir(): - raise NotADirectoryError(f"completion bundle is not a directory: {root}") - standalone = (root / "config.json").is_file() - if model_config is None: - model_config = root / ("config.json" if standalone else "configs/inference_model.json") - if checkpoint is None: - if standalone: - candidates = [root / name for name in ("model.safetensors", "model.safetensors.index.json")] - existing = [path for path in candidates if path.is_file()] - if len(existing) > 1: - raise ValueError("completion bundle has both single and sharded weights; select --checkpoint explicitly") - if not existing: - raise FileNotFoundError(f"missing completion checkpoint: expected model.safetensors or model.safetensors.index.json in {root}") - checkpoint = existing[0] - else: - checkpoint = root / "weights/material_completion.pt" - if trellis_root is None: - trellis_root = root / ("vae" if standalone else "weights/sparse_encoders") - if model_config is None or checkpoint is None or trellis_root is None: - raise ValueError("provide a completion bundle or model_config, checkpoint, and trellis_root") - # Snapshot files are symlinks to opaque cache blobs. Keep their logical - # directory so config-relative metadata and index-relative shards resolve. - paths = tuple(Path(path).expanduser().absolute() for path in (model_config, checkpoint, trellis_root)) - from .inference.pipeline import validate_completion_artifacts - validate_completion_artifacts(*paths) - return paths - - -def resolve_model_dir( - source: str | Path, - *, - revision: str | None = None, - cache_dir: str | Path | None = None, - local_files_only: bool = False, - allow_patterns: list[str] | None = None, -) -> Path: - """Return an existing directory for a local path or ``namespace/repo``. - - Local directories take precedence over Hub IDs. Use ``hf://namespace/repo`` - to identify a Hub source explicitly; use a ``Path`` or an explicit path - prefix to ensure a missing local directory never triggers a download. - """ - if not isinstance(source, (str, Path)): - raise TypeError("Model source must be a local path or a Hugging Face repo ID") - if source == "": - raise ValueError("Invalid model source: expected a nonempty path or namespace/repo") - - local_path = Path(source).expanduser() - if local_path.exists(): - if not local_path.is_dir(): - raise NotADirectoryError(f"Model source is not a directory: {local_path}") - return local_path.resolve() - - if isinstance(source, Path) or local_path.is_absolute() or source.startswith( - ("./", "../", "~/") - ) or source in (".", "..", "~"): - raise FileNotFoundError(f"Model directory does not exist: {local_path}") - - repo_id = source.removeprefix("hf://") - # Keep Hub imports off the local loading path (and lightweight CLI help). - from huggingface_hub import snapshot_download - from huggingface_hub.utils import HFValidationError, validate_repo_id - - try: - validate_repo_id(repo_id) - if repo_id.count("/") != 1: - raise ValueError("expected namespace/repo") - except (HFValidationError, ValueError) as exc: - raise ValueError( - f"Invalid model source {source!r}: expected an existing local directory " - "or a Hugging Face repo ID in namespace/repo form" - ) from exc - - downloaded = Path( - snapshot_download( - repo_id=repo_id, - revision=revision, - cache_dir=cache_dir, - local_files_only=local_files_only, - allow_patterns=allow_patterns, - ) - ).expanduser() - if not downloaded.exists(): - raise FileNotFoundError(f"Downloaded model directory does not exist: {downloaded}") - if not downloaded.is_dir(): - raise NotADirectoryError(f"Downloaded model source is not a directory: {downloaded}") - return downloaded.resolve() +"""Resolve model directories from local storage or the Hugging Face Hub.""" + +from pathlib import Path + + +# Default completion weights inside a standalone release (RAW step 30000). The +# release also keeps the earlier sharded model, which is used only as a fallback. +DEFAULT_COMPLETION_CHECKPOINT = "denoiser_step0030000.safetensors" +COMPLETION_ALLOW_PATTERNS = [ + "config.json", "partial_normalization.json", DEFAULT_COMPLETION_CHECKPOINT, + "configs/*.json", "weights/material_completion.pt", "weights/sparse_encoders/*.json", + "weights/scvae/*.json", "vae/*.json", "vae/ckpts/*.json", +] + + +def resolve_completion_artifacts( + root: str | Path | None = None, + *, + model_config: str | Path | None = None, + checkpoint: str | Path | None = None, + trellis_root: str | Path | None = None, +) -> tuple[Path, Path, Path]: + """Resolve and preflight a local completion release plus its sparse VAEs. + + Standalone releases contain root config/normalization and safe weights; + their separately obtained VAEs live at ``trellis_root`` (or ``root/vae``). + The original configs/weights bundle layout remains supported unchanged. + """ + if root is not None: + root = Path(root).expanduser().resolve() + if not root.is_dir(): + raise NotADirectoryError(f"completion bundle is not a directory: {root}") + standalone = (root / "config.json").is_file() + if model_config is None: + model_config = root / ("config.json" if standalone else "configs/inference_model.json") + if checkpoint is None: + if standalone and (root / DEFAULT_COMPLETION_CHECKPOINT).is_file(): + checkpoint = root / DEFAULT_COMPLETION_CHECKPOINT + elif standalone: + candidates = [root / name for name in ("model.safetensors", "model.safetensors.index.json")] + existing = [path for path in candidates if path.is_file()] + if len(existing) > 1: + raise ValueError("completion bundle has both single and sharded weights; select --checkpoint explicitly") + if not existing: + raise FileNotFoundError(f"missing completion checkpoint: expected {DEFAULT_COMPLETION_CHECKPOINT}, " + f"model.safetensors, or model.safetensors.index.json in {root}") + checkpoint = existing[0] + else: + checkpoint = root / "weights/material_completion.pt" + if trellis_root is None: + trellis_root = root / ("vae" if standalone else "weights/sparse_encoders") + if model_config is None or checkpoint is None or trellis_root is None: + raise ValueError("provide a completion bundle or model_config, checkpoint, and trellis_root") + # Snapshot files are symlinks to opaque cache blobs. Keep their logical + # directory so config-relative metadata and index-relative shards resolve. + paths = tuple(Path(path).expanduser().absolute() for path in (model_config, checkpoint, trellis_root)) + from .inference.pipeline import validate_completion_artifacts + validate_completion_artifacts(*paths) + return paths + + +def resolve_model_dir( + source: str | Path, + *, + revision: str | None = None, + cache_dir: str | Path | None = None, + local_files_only: bool = False, + allow_patterns: list[str] | None = None, +) -> Path: + """Return an existing directory for a local path or ``namespace/repo``. + + Local directories take precedence over Hub IDs. Use ``hf://namespace/repo`` + to identify a Hub source explicitly; use a ``Path`` or an explicit path + prefix to ensure a missing local directory never triggers a download. + """ + if not isinstance(source, (str, Path)): + raise TypeError("Model source must be a local path or a Hugging Face repo ID") + if source == "": + raise ValueError("Invalid model source: expected a nonempty path or namespace/repo") + + local_path = Path(source).expanduser() + if local_path.exists(): + if not local_path.is_dir(): + raise NotADirectoryError(f"Model source is not a directory: {local_path}") + return local_path.resolve() + + if isinstance(source, Path) or local_path.is_absolute() or source.startswith( + ("./", "../", "~/") + ) or source in (".", "..", "~"): + raise FileNotFoundError(f"Model directory does not exist: {local_path}") + + repo_id = source.removeprefix("hf://") + # Keep Hub imports off the local loading path (and lightweight CLI help). + from huggingface_hub import snapshot_download + from huggingface_hub.utils import HFValidationError, validate_repo_id + + try: + validate_repo_id(repo_id) + if repo_id.count("/") != 1: + raise ValueError("expected namespace/repo") + except (HFValidationError, ValueError) as exc: + raise ValueError( + f"Invalid model source {source!r}: expected an existing local directory " + "or a Hugging Face repo ID in namespace/repo form" + ) from exc + + downloaded = Path( + snapshot_download( + repo_id=repo_id, + revision=revision, + cache_dir=cache_dir, + local_files_only=local_files_only, + allow_patterns=allow_patterns, + ) + ).expanduser() + if not downloaded.exists(): + raise FileNotFoundError(f"Downloaded model directory does not exist: {downloaded}") + if not downloaded.is_dir(): + raise NotADirectoryError(f"Downloaded model source is not a directory: {downloaded}") + return downloaded.resolve() diff --git a/texturing.py b/texturing.py index 3f0b20a7ea4617a789db92d81215d6d289267fbc..499a54d38d7a3c9148b686db25b07958eaced57a 100644 --- a/texturing.py +++ b/texturing.py @@ -1,62 +1,192 @@ -#!/usr/bin/env python3 -"""Texture a supplied mesh with Supra-Tex PBR materials.""" -from __future__ import annotations - -import argparse -from pathlib import Path - - -def parse_args() -> argparse.Namespace: - parser = argparse.ArgumentParser(description="Complete a PBR GLB from a mesh, camera, RGBA image, and material maps.") - parser.add_argument("--input-dir", required=True, type=Path, help="directory containing geometry.pt, camera.json, condition_rgba.png, albedo.png, and mra.png") - parser.add_argument("--output", required=True, type=Path, help="new destination GLB; an existing file is never replaced") - parser.add_argument("--pretrained", help="local model bundle directory or Hugging Face namespace/repo") - parser.add_argument("--revision", help="Hugging Face model revision") - parser.add_argument("--cache-dir", type=Path, help="Hugging Face download cache") - parser.add_argument("--local-files-only", action="store_true", help="use cached model files without network access") - parser.add_argument("--attention", choices=("auto", "flash_attn_3", "flash_attn", "sdpa", "xformers"), default="auto") - parser.add_argument("--low-vram", action=argparse.BooleanOptionalAction, default=None, - help="CPU offload; auto enabled on GPUs with up to 24 GiB") - parser.add_argument("--model-config", type=Path, help="public Supra-Tex model ABI JSON; overrides bundle config") - parser.add_argument("--checkpoint", type=Path, help="public material-completion checkpoint; overrides bundle weights") - parser.add_argument("--trellis-root", type=Path, help="TRELLIS.2 VAE directory (all four components for HF releases)") - parser.add_argument("--seed", type=int, default=42) - parser.add_argument("--texture-size", type=int, default=4096) - parser.add_argument("--decimation-target", type=int, default=500000) - parser.add_argument("--texture-format", choices=("png", "webp"), default="png") - parser.add_argument("--gutter", action="store_true", help="enable chart-isolated one-texel base-color gutter padding (default: disabled)") - parser.add_argument("--dual-view", action="store_true", help="enable optional front/back observation fusion and validated reprojection") - parser.add_argument("--back-view-dir", type=Path, help="prepared v4 back-view package with generated RGBA, material maps and receipts") - parser.add_argument("--work-dir", type=Path, help="new directory for dual-view intermediate meshes and receipts") - args = parser.parse_args() - if not args.pretrained and not all((args.model_config, args.checkpoint, args.trellis_root)): - parser.error("provide --pretrained or all of --model-config, --checkpoint, and --trellis-root") - if args.dual_view != (args.back_view_dir is not None): - parser.error("--dual-view and --back-view-dir must be provided together") - if args.work_dir is not None and not args.dual_view: - parser.error("--work-dir requires --dual-view") - if args.dual_view and (args.texture_size != 4096 or args.texture_format != "png" - or args.decimation_target != 500000 or args.gutter): - parser.error("dual-view reprojection requires 4096 PNG, a 500000-face target, and no --gutter") - return args - - -def main() -> None: - args = parse_args() - from supratex.inference.runtime import configure_inference_runtime - configure_inference_runtime(attention=args.attention, low_vram=args.low_vram) - if args.pretrained: - from supratex.weights import COMPLETION_ALLOW_PATTERNS, resolve_completion_artifacts, resolve_model_dir - root = resolve_model_dir(args.pretrained, revision=args.revision, cache_dir=args.cache_dir, - local_files_only=args.local_files_only, - allow_patterns=COMPLETION_ALLOW_PATTERNS) - args.model_config, args.checkpoint, args.trellis_root = resolve_completion_artifacts( - root, model_config=args.model_config, checkpoint=args.checkpoint, - trellis_root=args.trellis_root) - from supratex.inference import run_completion - optional = {"back_view_dir": args.back_view_dir, "work_dir": args.work_dir} if args.dual_view else {} - run_completion(input_dir=args.input_dir, output=args.output, model_config=args.model_config, checkpoint=args.checkpoint, trellis_root=args.trellis_root, seed=args.seed, texture_size=args.texture_size, decimation_target=args.decimation_target, texture_format=args.texture_format, gutter=args.gutter, **optional) - - -if __name__ == "__main__": - main() +#!/usr/bin/env python3 +"""Texture a supplied mesh with Supra-Tex PBR materials.""" +from __future__ import annotations + +import argparse +from pathlib import Path + + +def _validate_options(args: argparse.Namespace) -> None: + from scripts.prepare_models import _pipeline + + if args.dual_view != (args.back_view_dir is not None): + raise ValueError("--dual-view and --back-view-dir must be provided together") + if args.work_dir is not None and not (args.dual_view or args.refine): + raise ValueError("--work-dir requires --dual-view or --refine") + _pipeline.validate_export_options(args) + + +def _validate_pipeline_paths(args: argparse.Namespace) -> None: + """Reject local input/output errors before fetching a manifest or weights.""" + from scripts.prepare_models import require_file + + args.input_dir = args.input_dir.expanduser().absolute() + args.output = args.output.expanduser().absolute() + if args.work_dir is not None: + args.work_dir = args.work_dir.expanduser().absolute() + for name in ("camera.json", "condition_rgba.png", "albedo.png", "mra.png"): + require_file(args.input_dir / name) + if not any((args.input_dir / name).is_file() for name in ("geometry.pt", "geometry.npz")): + raise FileNotFoundError(f"missing geometry.pt or geometry.npz in {args.input_dir}") + for name in ("model_config", "checkpoint"): + if getattr(args, name) is not None: + require_file(getattr(args, name).expanduser()) + for name in ("trellis_root", "back_view_dir"): + if getattr(args, name) is not None and not getattr(args, name).expanduser().is_dir(): + raise NotADirectoryError(f"--{name.replace('_', '-')} must be an existing directory") + output = args.output.expanduser().absolute() + if output.suffix.lower() != ".glb": + raise ValueError("--output must have a .glb suffix") + if output.exists() or output.is_symlink(): + raise FileExistsError(f"refusing to overwrite existing output: {output}") + if args.work_dir is not None: + work = args.work_dir.expanduser().absolute() + if work.exists() or work.is_symlink(): + raise FileExistsError(f"--work-dir must be a new directory: {work}") + for destination in (output, args.work_dir): + if destination is not None: + for parent in destination.expanduser().absolute().parents: + if parent.exists(): + if not parent.is_dir(): + raise NotADirectoryError(f"destination parent must be a directory: {parent}") + break + + +def _prepare_pipeline(args: argparse.Namespace, manifest: dict, source: dict, + completion: dict | None) -> None: + """Resolve only completion and VAE weights, keeping the Hub cache reusable.""" + from scripts.prepare_models import COMPLETION_PATTERNS, VAE_STEMS, completion_checkpoint, require_file + + def snapshot(descriptor: dict, patterns: list[str]) -> Path: + if descriptor["kind"] == "local": + return Path(descriptor["path"]) + from huggingface_hub import snapshot_download + try: + return Path(snapshot_download( + repo_id=descriptor["repo_id"], revision=descriptor.get("revision"), + allow_patterns=patterns, + cache_dir=str(args.cache_dir.expanduser()) if args.cache_dir is not None else None, + local_files_only=args.local_files_only)) + except Exception: + detail = "the required weights must already be cached" if args.local_files_only else "check the revision, network and your access" + raise RuntimeError(f"Cannot load weights from {descriptor['repo_id']}; {detail}") from None + + if completion is not None: + # Preserve standalone/legacy bundle overrides accepted by --pretrained. + root = snapshot(completion, COMPLETION_PATTERNS + ["configs/*.json", "weights/material_completion.pt", + "weights/sparse_encoders/*.json", "weights/scvae/*.json"]) + args.pretrained = str(root) + elif args.model_config is None or args.checkpoint is None: + subdir = manifest["models"]["completion"] + root = snapshot(source, [f"{subdir}/{pattern}" for pattern in COMPLETION_PATTERNS]) / subdir + args.pretrained = str(root) + else: + root = None + if root is not None: + standalone = (root / "config.json").is_file() + if args.model_config is None: + require_file(root / ("config.json" if standalone else "configs/inference_model.json")) + if standalone: + require_file(root / "partial_normalization.json") + if args.checkpoint is None: + if standalone: + completion_checkpoint(root) + else: + require_file(root / "weights/material_completion.pt") + if args.trellis_root is None: + entry = manifest["external_models"]["vae"] + args.trellis_root = snapshot({"kind": "hub", **entry}, entry["allow_patterns"]) + vae = args.trellis_root.expanduser() + vae = vae / "ckpts" if (vae / "ckpts").is_dir() else vae + for stem in VAE_STEMS: + for suffix in (".json", ".safetensors"): + require_file(vae / (stem + suffix)) + args.local_files_only = True + + +def parse_args(argv: list[str] | None = None) -> argparse.Namespace: + parser = argparse.ArgumentParser(description="Complete a PBR GLB from a mesh, camera, RGBA image, and material maps.", + allow_abbrev=False) + parser.add_argument("--pipeline", help="SUPRA-Tex release directory, pipeline.json, or HF repo ID; loads completion and VAE weights") + parser.add_argument("--pipeline-revision", help="HF revision of the complete SUPRA-Tex release") + parser.add_argument("--input-dir", type=Path, default=Path(__file__).resolve().parent / "assets/examples/pitcher_plant", + help="directory containing geometry.pt or geometry.npz, camera.json, condition_rgba.png, " + "albedo.png, and mra.png (default: bundled assets/examples/pitcher_plant)") + parser.add_argument("--output", required=True, type=Path, help="new destination GLB; an existing file is never replaced") + parser.add_argument("--pretrained", help="local model bundle directory or Hugging Face namespace/repo") + parser.add_argument("--revision", help="Hugging Face model revision") + parser.add_argument("--cache-dir", type=Path, help="Hugging Face download cache") + parser.add_argument("--local-files-only", action="store_true", help="use cached model files without network access") + parser.add_argument("--attention", choices=("auto", "flash_attn_3", "flash_attn", "sdpa", "xformers")) + parser.add_argument("--low-vram", action=argparse.BooleanOptionalAction, default=None, + help="CPU offload; auto enabled on GPUs with up to 24 GiB") + parser.add_argument("--model-config", type=Path, help="public Supra-Tex model ABI JSON; overrides bundle config") + parser.add_argument("--checkpoint", type=Path, help="public material-completion checkpoint; overrides bundle weights") + parser.add_argument("--trellis-root", type=Path, help="TRELLIS.2 VAE directory (all four components for HF releases)") + parser.add_argument("--seed", type=int) + parser.add_argument("--texture-size", type=int) + parser.add_argument("--decimation-target", type=int) + parser.add_argument("--texture-format", choices=("png", "webp")) + parser.add_argument("--gutter", action="store_true", help="enable chart-isolated one-texel base-color gutter padding (default: disabled); " + "with --refine, the observation-guided gutter is used") + parser.add_argument("--refine", action="store_true", + help="refine the visible texture after export: write the observed reference-view albedo into the atlas " + "and repair its UV gutters from the observation (sharper fine text and detail; needs the 4096 PNG " + "atlas; adds about 2 minutes on an RTX 3090)") + parser.add_argument("--remove-floaters", action="store_true", + help="before completion, drop disconnected geometry pieces that project mostly outside the " + "image foreground; extra geometry attached to the object is kept") + parser.add_argument("--dual-view", action="store_true", help="add front/back observation fusion from a prepared back view; includes --refine for both views") + parser.add_argument("--back-view-dir", type=Path, help="prepared v4 back-view package with generated RGBA, material maps and receipts") + parser.add_argument("--work-dir", type=Path, help="new directory for dual-view intermediate meshes and receipts") + args = parser.parse_args(argv) + if not args.pipeline and args.pipeline_revision is not None: + parser.error("--pipeline-revision requires --pipeline") + if args.pipeline and args.revision is not None and args.pretrained is None: + parser.error("--revision requires --pretrained with --pipeline; use --pipeline-revision for the release") + try: + _validate_options(args) + if args.pipeline: + from scripts.prepare_models import _pipeline + completion = _pipeline.model_source(args.pretrained, args.revision) if args.pretrained is not None else None + _validate_pipeline_paths(args) + manifest, source = _pipeline.resolve_pipeline( + args.pipeline, revision=args.pipeline_revision, cache_dir=args.cache_dir, + local_files_only=args.local_files_only) + defaults = manifest["defaults"] + else: + from scripts.run_example import inference_defaults + defaults = inference_defaults() + for key, value in defaults.items(): + if getattr(args, key) is None: + setattr(args, key, value) + _validate_options(args) + if args.pipeline: + _prepare_pipeline(args, manifest, source, completion) + except (OSError, ValueError, RuntimeError) as error: + parser.error(str(error)) + if not args.pretrained and not all((args.model_config, args.checkpoint, args.trellis_root)): + parser.error("provide --pipeline, --pretrained, or all of --model-config, --checkpoint, and --trellis-root") + return args + + +def main() -> None: + args = parse_args() + from supratex.inference.runtime import configure_inference_runtime + configure_inference_runtime(attention=args.attention, low_vram=args.low_vram) + if args.pretrained: + from supratex.weights import COMPLETION_ALLOW_PATTERNS, resolve_completion_artifacts, resolve_model_dir + root = resolve_model_dir(args.pretrained, revision=args.revision, cache_dir=args.cache_dir, + local_files_only=args.local_files_only, + allow_patterns=COMPLETION_ALLOW_PATTERNS) + args.model_config, args.checkpoint, args.trellis_root = resolve_completion_artifacts( + root, model_config=args.model_config, checkpoint=args.checkpoint, + trellis_root=args.trellis_root) + from supratex.inference import run_completion + optional = ({"back_view_dir": args.back_view_dir, "work_dir": args.work_dir} if args.dual_view else + {"refine": True, "work_dir": args.work_dir} if args.refine else {}) + run_completion(input_dir=args.input_dir, output=args.output, model_config=args.model_config, checkpoint=args.checkpoint, trellis_root=args.trellis_root, seed=args.seed, texture_size=args.texture_size, decimation_target=args.decimation_target, texture_format=args.texture_format, gutter=args.gutter, remove_floaters=args.remove_floaters, **optional) + + +if __name__ == "__main__": + main()