Use validated latest single-view inference and step-30000 completion weights

Pin upstream e4dfbbd and model release 1bc02eb. Retain rendered o-voxel previews and explicit later GLB export. Backview and refine disabled. Keep current UI layout. Validated on L4: 112 preview frames in 399.8s without GLB, separate 4096px PBR GLB export in 48.4s.

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)

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"
+  }
+}

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."
+  }
+}

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

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

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

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

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)

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 *

deployment/runtime/qwen/diffsynth/core/data/__init__.py

@@ -1 +1 @@
-from .unified_dataset import UnifiedDataset
+"""Image and media preprocessing operators used by inference pipelines."""

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

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

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)

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)

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

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)

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

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

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'<img src="{url}" data-mode="{mode}" data-view="{view}" alt="{mode.replace("_"," ")} view {view+1}"{hidden}>')
-    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'<button type="button" data-mode="{mode}" class="{"active" if active else ""}" '
-            f'title="{label}" aria-label="{label}" aria-pressed="{str(active).lower()}">'
-            f'<img src="data:image/png;base64,{swatch}" alt=""><span class="st-mode-name">{label}</span></button>')
-    return (f'<div class="st-ovoxel-preview" data-mode="{escape(default, quote=True)}" data-view="0">'
-        '<div class="st-ovoxel-frames">'+''.join(images)+'</div>'
-        '<div class="st-ovoxel-modes" role="group" aria-label="Render mode">'+''.join(buttons)+'</div>'
-        f'<label class="st-ovoxel-angle">View angle <input type="range" min="0" max="{metadata["views"]-1}" value="0" step="1" aria-label="View angle"></label></div>')
-
-
-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'<img src="{url}" data-mode="{mode}" data-view="{view}" alt="{mode.replace("_"," ")} view {view+1}"{hidden}>')
+    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'<button type="button" data-mode="{mode}" class="{"active" if active else ""}" '
+            f'title="{label}" aria-label="{label}" aria-pressed="{str(active).lower()}">'
+            f'<img src="data:image/png;base64,{swatch}" alt=""><span class="st-mode-name">{label}</span></button>')
+    return (f'<div class="st-ovoxel-preview" data-mode="{escape(default, quote=True)}" data-view="0">'
+        '<div class="st-ovoxel-frames">'+''.join(images)+'</div>'
+        '<div class="st-ovoxel-modes" role="group" aria-label="Render mode">'+''.join(buttons)+'</div>'
+        f'<label class="st-ovoxel-angle">View angle <input type="range" min="0" max="{metadata["views"]-1}" value="0" step="1" aria-label="View angle"></label></div>')
+
+
+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)

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

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 ('<div class="st-startup" role="status" aria-live="polite">'
-            '<h3>Loading weights</h3>'
-            '<p class="st-startup-note">Models load automatically in the background when the demo starts. All users share the loaded models.</p>'
-            '<div class="st-startup-heading"><span>Loading progress</span><span>' + detail + '</span></div>'
-            '<progress max="100"' + value + ' aria-label="Weight loading progress"></progress>'
-            '<div class="st-startup-stage">' + label + '</div></div>')
+"""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 ('<div class="st-startup" role="status" aria-live="polite">'
+            '<h3>Loading weights</h3>'
+            '<p class="st-startup-note">Model files are prepared in the background and cached for all users.</p>'
+            '<div class="st-startup-heading"><span>Loading progress</span><span>' + detail + '</span></div>'
+            '<progress max="100"' + value + ' aria-label="Weight loading progress"></progress>'
+            '<div class="st-startup-stage">' + label + '</div></div>')

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: <output stem>_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: <output stem>_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()

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"]
+

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

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
+``<stem>.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())

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-<temporary>" / "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())

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/<revision>")
-        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())

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

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

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

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

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

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

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

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)

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

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)

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

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

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

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

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

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

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

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

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

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

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())}