Add Pixal3D modular blocks

.gitignore

@@ -0,0 +1 @@
+__pycache__/

LICENSE

@@ -0,0 +1,201 @@
+                                 Apache License
+                           Version 2.0, January 2004
+                        http://www.apache.org/licenses/
+
+   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+   1. Definitions.
+
+      "License" shall mean the terms and conditions for use, reproduction,
+      and distribution as defined by Sections 1 through 9 of this document.
+
+      "Licensor" shall mean the copyright owner or entity authorized by
+      the copyright owner that is granting the License.
+
+      "Legal Entity" shall mean the union of the acting entity and all
+      other entities that control, are controlled by, or are under common
+      control with that entity. For the purposes of this definition,
+      "control" means (i) the power, direct or indirect, to cause the
+      direction or management of such entity, whether by contract or
+      otherwise, or (ii) ownership of fifty percent (50%) or more of the
+      outstanding shares, or (iii) beneficial ownership of such entity.
+
+      "You" (or "Your") shall mean an individual or Legal Entity
+      exercising permissions granted by this License.
+
+      "Source" form shall mean the preferred form for making modifications,
+      including but not limited to software source code, documentation
+      source, and configuration files.
+
+      "Object" form shall mean any form resulting from mechanical
+      transformation or translation of a Source form, including but
+      not limited to compiled object code, generated documentation,
+      and conversions to other media types.
+
+      "Work" shall mean the work of authorship, whether in Source or
+      Object form, made available under the License, as indicated by a
+      copyright notice that is included in or attached to the work
+      (an example is provided in the Appendix below).
+
+      "Derivative Works" shall mean any work, whether in Source or Object
+      form, that is based on (or derived from) the Work and for which the
+      editorial revisions, annotations, elaborations, or other modifications
+      represent, as a whole, an original work of authorship. For the purposes
+      of this License, Derivative Works shall not include works that remain
+      separable from, or merely link (or bind by name) to the interfaces of,
+      the Work and Derivative Works thereof.
+
+      "Contribution" shall mean any work of authorship, including
+      the original version of the Work and any modifications or additions
+      to that Work or Derivative Works thereof, that is intentionally
+      submitted to Licensor for inclusion in the Work by the copyright owner
+      or by an individual or Legal Entity authorized to submit on behalf of
+      the copyright owner. For the purposes of this definition, "submitted"
+      means any form of electronic, verbal, or written communication sent
+      to the Licensor or its representatives, including but not limited to
+      communication on electronic mailing lists, source code control systems,
+      and issue tracking systems that are managed by, or on behalf of, the
+      Licensor for the purpose of discussing and improving the Work, but
+      excluding communication that is conspicuously marked or otherwise
+      designated in writing by the copyright owner as "Not a Contribution."
+
+      "Contributor" shall mean Licensor and any individual or Legal Entity
+      on behalf of whom a Contribution has been received by Licensor and
+      subsequently incorporated within the Work.
+
+   2. Grant of Copyright License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      copyright license to reproduce, prepare Derivative Works of,
+      publicly display, publicly perform, sublicense, and distribute the
+      Work and such Derivative Works in Source or Object form.
+
+   3. Grant of Patent License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      (except as stated in this section) patent license to make, have made,
+      use, offer to sell, sell, import, and otherwise transfer the Work,
+      where such license applies only to those patent claims licensable
+      by such Contributor that are necessarily infringed by their
+      Contribution(s) alone or by combination of their Contribution(s)
+      with the Work to which such Contribution(s) was submitted. If You
+      institute patent litigation against any entity (including a
+      cross-claim or counterclaim in a lawsuit) alleging that the Work
+      or a Contribution incorporated within the Work constitutes direct
+      or contributory patent infringement, then any patent licenses
+      granted to You under this License for that Work shall terminate
+      as of the date such litigation is filed.
+
+   4. Redistribution. You may reproduce and distribute copies of the
+      Work or Derivative Works thereof in any medium, with or without
+      modifications, and in Source or Object form, provided that You
+      meet the following conditions:
+
+      (a) You must give any other recipients of the Work or
+          Derivative Works a copy of this License; and
+
+      (b) You must cause any modified files to carry prominent notices
+          stating that You changed the files; and
+
+      (c) You must retain, in the Source form of any Derivative Works
+          that You distribute, all copyright, patent, trademark, and
+          attribution notices from the Source form of the Work,
+          excluding those notices that do not pertain to any part of
+          the Derivative Works; and
+
+      (d) If the Work includes a "NOTICE" text file as part of its
+          distribution, then any Derivative Works that You distribute must
+          include a readable copy of the attribution notices contained
+          within such NOTICE file, excluding those notices that do not
+          pertain to any part of the Derivative Works, in at least one
+          of the following places: within a NOTICE text file distributed
+          as part of the Derivative Works; within the Source form or
+          documentation, if provided along with the Derivative Works; or,
+          within a display generated by the Derivative Works, if and
+          wherever such third-party notices normally appear. The contents
+          of the NOTICE file are for informational purposes only and
+          do not modify the License. You may add Your own attribution
+          notices within Derivative Works that You distribute, alongside
+          or as an addendum to the NOTICE text from the Work, provided
+          that such additional attribution notices cannot be construed
+          as modifying the License.
+
+      You may add Your own copyright statement to Your modifications and
+      may provide additional or different license terms and conditions
+      for use, reproduction, or distribution of Your modifications, or
+      for any such Derivative Works as a whole, provided Your use,
+      reproduction, and distribution of the Work otherwise complies with
+      the conditions stated in this License.
+
+   5. Submission of Contributions. Unless You explicitly state otherwise,
+      any Contribution intentionally submitted for inclusion in the Work
+      by You to the Licensor shall be under the terms and conditions of
+      this License, without any additional terms or conditions.
+      Notwithstanding the above, nothing herein shall supersede or modify
+      the terms of any separate license agreement you may have executed
+      with Licensor regarding such Contributions.
+
+   6. Trademarks. This License does not grant permission to use the trade
+      names, trademarks, service marks, or product names of the Licensor,
+      except as required for reasonable and customary use in describing the
+      origin of the Work and reproducing the content of the NOTICE file.
+
+   7. Disclaimer of Warranty. Unless required by applicable law or
+      agreed to in writing, Licensor provides the Work (and each
+      Contributor provides its Contributions) on an "AS IS" BASIS,
+      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+      implied, including, without limitation, any warranties or conditions
+      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+      PARTICULAR PURPOSE. You are solely responsible for determining the
+      appropriateness of using or redistributing the Work and assume any
+      risks associated with Your exercise of permissions under this License.
+
+   8. Limitation of Liability. In no event and under no legal theory,
+      whether in tort (including negligence), contract, or otherwise,
+      unless required by applicable law (such as deliberate and grossly
+      negligent acts) or agreed to in writing, shall any Contributor be
+      liable to You for damages, including any direct, indirect, special,
+      incidental, or consequential damages of any character arising as a
+      result of this License or out of the use or inability to use the
+      Work (including but not limited to damages for loss of goodwill,
+      work stoppage, computer failure or malfunction, or any and all
+      other commercial damages or losses), even if such Contributor
+      has been advised of the possibility of such damages.
+
+   9. Accepting Warranty or Additional Liability. While redistributing
+      the Work or Derivative Works thereof, You may choose to offer,
+      and charge a fee for, acceptance of support, warranty, indemnity,
+      or other liability obligations and/or rights consistent with this
+      License. However, in accepting such obligations, You may act only
+      on Your own behalf and on Your sole responsibility, not on behalf
+      of any other Contributor, and only if You agree to indemnify,
+      defend, and hold each Contributor harmless for any liability
+      incurred by, or claims asserted against, such Contributor by reason
+      of your accepting any such warranty or additional liability.
+
+   END OF TERMS AND CONDITIONS
+
+   APPENDIX: How to apply the Apache License to your work.
+
+      To apply the Apache License to your work, attach the following
+      boilerplate notice, with the fields enclosed by brackets "[]"
+      replaced with your own identifying information. (Don't include
+      the brackets!)  The text should be enclosed in the appropriate
+      comment syntax for the file format. We also recommend that a
+      file or class name and description of purpose be included on the
+      same "printed page" as the copyright notice for easier
+      identification within third-party archives.
+
+   Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.

LICENSE-PIXAL3D

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 Tencent.
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

NOTICE

@@ -0,0 +1,6 @@
+The camera projection (projection.py) and the latent normalization statistics (normalization.py) are adapted from:
+Pixal3D, Copyright (c) 2026 Tencent, MIT (see LICENSE-PIXAL3D).
+https://github.com/TencentARC/Pixal3D
+Reference revision: f7cf384
+The remaining block code is original and licensed under Apache-2.0 (see LICENSE).
+The package contains no learned weights. Converted checkpoint licensing is separate.

README.md

@@ -0,0 +1,83 @@
+---
+license: apache-2.0
+library_name: diffusers
+pipeline_tag: image-to-3d
+base_model:
+- TencentARC/Pixal3D
+tags:
+- pixal3d
+- modular-diffusers
+- image-to-3d
+---
+# Pixal3D Modular
+
+Modular Diffusers blocks for [Pixal3D](https://github.com/TencentARC/Pixal3D), from a prepared image to sparse geometry and material attributes. Model code and converted weights are packaged separately in `Pixal3D`.
+
+## Usage
+
+Requires PyTorch, Diffusers with Modular Pipelines, Transformers with DINOv3 support, Accelerate, Safetensors, NumPy and Pillow. Tested with Diffusers `0.41.0` and Transformers `5.20.2`.
+
+```python
+import math
+
+import torch
+from diffusers import ModularPipelineBlocks
+from diffusers.utils import load_image
+
+
+blocks = ModularPipelineBlocks.from_pretrained(
+    "OzzyGT/Pixal3D_Modular",
+    trust_remote_code=True,
+    components_repo="OzzyGT/Pixal3D",
+    trust_components_code=True,
+)
+pipe = blocks.init_pipeline()
+flow_names = (
+    "sparse_structure_flow_model",
+    "shape_slat_flow_model_512",
+    "shape_slat_flow_model_1024",
+    "tex_slat_flow_model_1024",
+)
+dtypes = {spec.name: torch.float32 for spec in blocks.expected_components}
+dtypes.update({name: torch.bfloat16 for name in (*flow_names, "image_encoder")})
+pipe.load_components(dtype=dtypes)
+pipe.to("cuda")
+
+image = load_image("https://huggingface.co/datasets/OzzyGT/diffusers-examples/resolve/main/Pixal3D/Pixal3D_jester_input.png")
+result = pipe(
+    image=image,
+    camera_angle_x=math.radians(33.85),
+    distance=1.643,
+    mesh_scale=1.0,
+    resolution=1024,
+    num_inference_steps=12,
+    generator=torch.Generator().manual_seed(42),
+    output=["geometry", "materials", "actual_resolution"],
+)
+geometry = result["geometry"]
+materials = result["materials"]
+```
+
+The bf16 model does not fit fully on a 32 GB GPU; use CPU offloading or the SDNQ [INT8](https://huggingface.co/OzzyGT/Pixal3D_sdnq_dynamic_8bit) and [INT4](https://huggingface.co/OzzyGT/Pixal3D_sdnq_dynamic_4bit) versions.
+
+The example uses `trust_remote_code=True`, which loads these blocks, and `trust_components_code=True`, which allows the model repository's custom classes.
+
+> **Note:** both flags run Python code downloaded from the Hub: `trust_remote_code=True` runs the code of the blocks repository, and `trust_components_code=True` runs the code shipped in the `components_repo`. Only use them with repositories you trust, and review their code before passing different ones.
+
+Input is one prepared RGB/RGBA image. RGBA is composited on black, then resized; there is no automatic foreground extraction or crop. Camera parameters must correspond to that prepared image.
+
+## Outputs
+
+- `geometry`: integer `(batch, x, y, z)` cell coordinates, vertex offsets, three edge-intersection flags and quad weights for flexible-dual-grid extraction.
+- `materials`: matching cell coordinates and six attributes, with layout `base_color[0:3]`, `metallic[3:4]`, `roughness[4:5]`, `alpha[5:6]`.
+- `actual_resolution`: effective grid resolution. A 1536 request can fall back in steps of 128 to 1024 when the token limit is exceeded.
+
+Coordinates retain the original model convention (Y-up in the reference mesh viewer). Grid positions become normalized object coordinates through `(cell + vertex_offset) / actual_resolution - 0.5`; these are not meters. No game-engine axis conversion or mesh export is performed. Mesh extraction, hole filling and texture baking are separate operations for the consumer.
+
+The blocks expose image encoding, NAF upsampling, projection, denoising, normalization and decoding as separate stages. The initial workflow supports one view and one sample. The four flow stages use their own scheduler/guider components, with the shape stages sharing settings. Guidance rescaling follows Pixal3D's clean-sample calculation.
+
+## Licensing
+
+The blocks are [Apache-2.0](LICENSE), with the camera projection and latent statistics adapted from [TencentARC/Pixal3D](https://github.com/TencentARC/Pixal3D) under [MIT](LICENSE-PIXAL3D). See [NOTICE](NOTICE). This repository contains no model weights. Model components retain their own licenses: Pixal3D MIT, NAF Apache-2.0, and DINOv3's separate terms. See the model repository's license files and provenance notice. No background-removal or preview-rendering dependencies are included.
+
+Validation: pretrained flow/structure components passed FP32 reference checks. NAF matched native NATTEN at 64, 512 and 1024 feature sizes (maximum error 0.0000141). Both sparse decoders matched the original FlexGEMM FP32 implementation on a latent patch, including identical subdivision coordinates (maximum error 0.000168). Full 12-step generation passed at 1024 and 1536 on RTX 5090. The tested mixed-precision setup uses BF16 for DINOv3 and the four flow models, with FP32 NAF and decoders. Synchronous two-block group offloading is verified; asynchronous streaming and other GPUs are not covered by these checks.

decoders.py

@@ -0,0 +1,207 @@
+import torch
+from diffusers import ModelMixin
+from diffusers.modular_pipelines.modular_pipeline_utils import (
+    ComponentSpec,
+    InputParam,
+    OutputParam,
+)
+from torch.nn import functional as F
+
+from .resolvable import ResolvableBlock
+
+
+class Pixal3DStructureDecodeStep(ResolvableBlock):
+    model_name = "pixal3d"
+
+    @property
+    def description(self):
+        return "Decode structure latents and extract occupied sparse coordinates."
+
+    @property
+    def expected_components(self):
+        return self.resolve([ComponentSpec("sparse_structure_decoder", ModelMixin)])
+
+    @property
+    def inputs(self):
+        return [
+            InputParam("latents", required=True, type_hint=torch.Tensor),
+            InputParam(
+                "structure_resolution",
+                default=32,
+                type_hint=int,
+                description="Occupancy grid size; the decoded 64^3 grid is max-pooled down to it.",
+            ),
+        ]
+
+    @property
+    def intermediate_outputs(self):
+        return [OutputParam("coords", type_hint=torch.Tensor)]
+
+    @torch.no_grad()
+    def __call__(self, components, state):
+        b = self.get_block_state(state)
+        model = components.sparse_structure_decoder
+        occupancy = model(b.latents.to(model.device)) > 0
+        if occupancy.shape[2] != b.structure_resolution:
+            if occupancy.shape[2] % b.structure_resolution:
+                raise ValueError("Structure resolution must divide the decoded grid size")
+            ratio = occupancy.shape[2] // b.structure_resolution
+            occupancy = F.max_pool3d(occupancy.float(), ratio, ratio) > 0.5
+        b.coords = torch.argwhere(occupancy)[:, [0, 2, 3, 4]].int()
+        self.set_block_state(state, b)
+        return components, state
+
+
+class Pixal3DShapeUpsampleStep(ResolvableBlock):
+    model_name = "pixal3d"
+
+    @property
+    def description(self):
+        return "Use the shape decoder to generate the high-resolution sparse coordinate set."
+
+    @property
+    def expected_components(self):
+        return self.resolve([ComponentSpec("shape_slat_decoder", ModelMixin)])
+
+    @property
+    def inputs(self):
+        return [
+            InputParam("shape_latents", required=True, type_hint=torch.Tensor),
+            InputParam("coords", required=True, type_hint=torch.Tensor),
+            InputParam("resolution", default=1024, type_hint=int, description="Output grid resolution: 1024 or 1536."),
+            InputParam(
+                "max_num_tokens",
+                default=49152,
+                type_hint=int,
+                description="Token budget; 1536 falls back in steps of 128 until the sparse grid fits.",
+            ),
+        ]
+
+    @property
+    def intermediate_outputs(self):
+        return [
+            OutputParam("coords", type_hint=torch.Tensor),
+            OutputParam("actual_resolution", type_hint=int),
+        ]
+
+    @torch.no_grad()
+    def __call__(self, components, state):
+        b = self.get_block_state(state)
+        if b.resolution not in (1024, 1536):
+            raise ValueError("Supported resolutions are 1024 and 1536")
+        model = components.shape_slat_decoder
+        coords = model(
+            b.shape_latents[0].to(model.device),
+            b.coords.to(model.device),
+            upsample_times=4,
+        )["coords"]
+        resolution = b.resolution
+        while True:
+            grid = resolution // 16
+            quantized = torch.cat(
+                [
+                    coords[:, :1],
+                    ((coords[:, 1:] + 0.5) / 512 * (grid - 1)).round().int(),
+                ],
+                1,
+            ).unique(dim=0)
+            if len(quantized) < b.max_num_tokens or resolution == 1024:
+                break
+            resolution -= 128
+        b.coords = quantized
+        b.actual_resolution = resolution
+        self.set_block_state(state, b)
+        return components, state
+
+
+class Pixal3DShapeDecodeStep(ResolvableBlock):
+    model_name = "pixal3d"
+
+    @property
+    def description(self):
+        return "Decode shape latents into a flexible dual grid and texture subdivision guides."
+
+    @property
+    def expected_components(self):
+        return self.resolve([ComponentSpec("shape_slat_decoder", ModelMixin)])
+
+    @property
+    def inputs(self):
+        return [
+            InputParam("shape_latents", required=True, type_hint=torch.Tensor),
+            InputParam("coords", required=True, type_hint=torch.Tensor),
+        ]
+
+    @property
+    def intermediate_outputs(self):
+        return [
+            OutputParam("geometry", type_hint=dict),
+            OutputParam("subdivisions", type_hint=list),
+        ]
+
+    @torch.no_grad()
+    def __call__(self, components, state):
+        b = self.get_block_state(state)
+        model = components.shape_slat_decoder
+        decoded = model(
+            b.shape_latents[0].to(model.device),
+            b.coords.to(model.device),
+            return_subs=True,
+        )
+        features = decoded["features"]
+        b.geometry = {
+            "coords": decoded["coords"],
+            "vertex_offsets": 2 * features[:, :3].sigmoid() - 0.5,
+            "intersected": features[:, 3:6] > 0,
+            "quad_weights": F.softplus(features[:, 6:7]),
+        }
+        b.subdivisions = decoded["subdivisions"]
+        self.set_block_state(state, b)
+        return components, state
+
+
+class Pixal3DTextureDecodeStep(ResolvableBlock):
+    model_name = "pixal3d"
+
+    @property
+    def description(self):
+        return "Decode material attributes on the shape subdivision grid, without texture baking."
+
+    @property
+    def expected_components(self):
+        return self.resolve([ComponentSpec("tex_slat_decoder", ModelMixin)])
+
+    @property
+    def inputs(self):
+        return [
+            InputParam("texture_latents", required=True, type_hint=torch.Tensor),
+            InputParam("coords", required=True, type_hint=torch.Tensor),
+            InputParam("subdivisions", required=True, type_hint=list),
+        ]
+
+    @property
+    def intermediate_outputs(self):
+        return [OutputParam("materials", type_hint=dict)]
+
+    @torch.no_grad()
+    def __call__(self, components, state):
+        b = self.get_block_state(state)
+        model = components.tex_slat_decoder
+        guides = [{k: v.to(model.device) for k, v in sub.items()} for sub in b.subdivisions]
+        decoded = model(
+            b.texture_latents[0].to(model.device),
+            b.coords.to(model.device),
+            guide_subs=guides,
+        )
+        b.materials = {
+            "coords": decoded["coords"],
+            "attributes": decoded["features"] * 0.5 + 0.5,
+            "layout": {
+                "base_color": [0, 3],
+                "metallic": [3, 4],
+                "roughness": [4, 5],
+                "alpha": [5, 6],
+            },
+        }
+        self.set_block_state(state, b)
+        return components, state

denoise.py

@@ -0,0 +1,261 @@
+from dataclasses import dataclass, field
+
+import torch
+from diffusers import ModelMixin, SchedulerMixin
+from diffusers.modular_pipelines import LoopSequentialPipelineBlocks
+from diffusers.modular_pipelines.modular_pipeline_utils import ComponentSpec, InputParam, OutputParam
+from diffusers.utils.torch_utils import randn_tensor
+
+from .guidance import Pixal3DGuidance
+from .resolvable import ResolvableBlock
+
+
+@dataclass(frozen=True)
+class Pixal3DStage:
+    """Components and guidance settings of one flow stage."""
+
+    transformer: str
+    scheduler: str
+    guider: str
+    guider_config: dict = field(default_factory=dict)
+    sparse: bool = True
+    concat_cond: bool = False
+
+
+STRUCTURE_STAGE = Pixal3DStage(
+    "sparse_structure_flow_model",
+    "structure_scheduler",
+    "structure_guider",
+    {"guidance_scale": 7.5, "guidance_rescale": 0.7, "guidance_interval": (0.6, 1.0)},
+    sparse=False,
+)
+SHAPE_GUIDANCE = {"guidance_scale": 7.5, "guidance_rescale": 0.5, "guidance_interval": (0.6, 1.0)}
+SHAPE_LOW_STAGE = Pixal3DStage("shape_slat_flow_model_512", "shape_scheduler", "shape_guider", SHAPE_GUIDANCE)
+SHAPE_HIGH_STAGE = Pixal3DStage("shape_slat_flow_model_1024", "shape_scheduler", "shape_guider", SHAPE_GUIDANCE)
+TEXTURE_STAGE = Pixal3DStage(
+    "tex_slat_flow_model_1024",
+    "texture_scheduler",
+    "texture_guider",
+    {"guidance_scale": 1.0, "guidance_rescale": 0.0, "guidance_interval": (0.6, 0.9)},
+    concat_cond=True,
+)
+
+
+class Pixal3DPrepareLatentsStep(ResolvableBlock):
+    model_name = "pixal3d"
+
+    def __init__(self, stage: Pixal3DStage):
+        self.stage = stage
+        super().__init__()
+
+    @property
+    def description(self):
+        return "Draw the initial noise for a dense structure volume or for the sparse shape/texture tokens."
+
+    @property
+    def expected_components(self):
+        return self.resolve([ComponentSpec(self.stage.transformer, ModelMixin)])
+
+    @property
+    def inputs(self):
+        inputs = [InputParam.template("generator")]
+        if self.stage.sparse:
+            inputs.append(
+                InputParam("coords", required=True, type_hint=torch.Tensor, description="Sparse (N, 4) coordinates.")
+            )
+        if self.stage.concat_cond:
+            inputs.append(
+                InputParam("concat_cond", required=True, type_hint=torch.Tensor, description="Shape latents.")
+            )
+        return inputs
+
+    @property
+    def intermediate_outputs(self):
+        return [OutputParam("latents", type_hint=torch.Tensor, description="Initial noise for this stage.")]
+
+    @torch.no_grad()
+    def __call__(self, components, state):
+        b = self.get_block_state(state)
+        config = getattr(components, self.stage.transformer).config
+        if not self.stage.sparse:
+            shape = (1, config.in_channels, *([config.resolution] * 3))
+        else:
+            if len(b.coords) == 0:
+                raise ValueError("Nonempty sparse coordinates are required")
+            channels = config.in_channels - (b.concat_cond.shape[-1] if self.stage.concat_cond else 0)
+            shape = (1, len(b.coords), channels)
+        b.latents = randn_tensor(
+            shape, generator=b.generator, device=components._execution_device, dtype=torch.float32
+        )
+        self.set_block_state(state, b)
+        return components, state
+
+
+class Pixal3DSetTimestepsStep(ResolvableBlock):
+    model_name = "pixal3d"
+
+    def __init__(self, stage: Pixal3DStage):
+        self.stage = stage
+        super().__init__()
+
+    @property
+    def description(self):
+        return "Set the shifted flow-matching sigmas for this stage."
+
+    @property
+    def expected_components(self):
+        return self.resolve([ComponentSpec(self.stage.scheduler, SchedulerMixin)])
+
+    @property
+    def inputs(self):
+        return [InputParam.template("num_inference_steps", default=12)]
+
+    @property
+    def intermediate_outputs(self):
+        return [OutputParam("timesteps", type_hint=list, description="Sigmas in [0, 1], one per step.")]
+
+    def __call__(self, components, state):
+        b = self.get_block_state(state)
+        scheduler = getattr(components, self.stage.scheduler)
+        scheduler.set_timesteps(b.num_inference_steps)
+        b.timesteps = scheduler.timesteps
+        self.set_block_state(state, b)
+        return components, state
+
+
+class Pixal3DLoopDenoiser(ResolvableBlock):
+    model_name = "pixal3d"
+    # Conditional batches use the image features; the unconditional batch uses zeros of the same shape.
+    _guider_input_fields = {
+        "global_features": ("global_features", "negative_global_features"),
+        "projected_features": ("projected_features", "negative_projected_features"),
+        "latents": "latents",
+    }
+
+    def __init__(self, stage: Pixal3DStage):
+        self.stage = stage
+        super().__init__()
+
+    @property
+    def description(self):
+        return "Step within the denoising loop: run the flow model for each guider batch and combine them."
+
+    @property
+    def expected_components(self):
+        return self.resolve(
+            [
+                ComponentSpec(self.stage.transformer, ModelMixin),
+                ComponentSpec(
+                    self.stage.guider,
+                    Pixal3DGuidance,
+                    default_creation_method="from_config",
+                    config=self.stage.guider_config,
+                ),
+            ]
+        )
+
+    @property
+    def inputs(self):
+        return [
+            InputParam("global_features", required=True, type_hint=torch.Tensor, description="DINOv3 global tokens."),
+            InputParam(
+                "projected_features", required=True, type_hint=torch.Tensor, description="Camera-projected features."
+            ),
+        ]
+
+    @torch.no_grad()
+    def __call__(self, components, block_state, i, t):
+        model = getattr(components, self.stage.transformer)
+        guider = getattr(components, self.stage.guider)
+        guider.set_state(step=i, num_inference_steps=len(block_state.timesteps), timestep=t)
+        batches = guider.prepare_inputs_from_block_state(block_state, self._guider_input_fields)
+        latents = block_state.latents
+        timestep = torch.full((latents.shape[0],), 1000 * t, device=latents.device, dtype=torch.float32)
+        for batch in batches:
+            guider.prepare_models(model)
+            batch.noise_pred = model(
+                latents,
+                timestep,
+                batch.global_features,
+                batch.projected_features,
+                coords=block_state.model_coords,
+                concat_cond=block_state.model_concat_cond,
+            )
+            guider.cleanup_models(model)
+        block_state.noise_pred = guider(batches).pred
+        return components, block_state
+
+
+class Pixal3DLoopAfterDenoiser(ResolvableBlock):
+    model_name = "pixal3d"
+
+    def __init__(self, stage: Pixal3DStage):
+        self.stage = stage
+        super().__init__()
+
+    @property
+    def description(self):
+        return "Step within the denoising loop: take one Euler step with the guided velocity."
+
+    @property
+    def expected_components(self):
+        return self.resolve([ComponentSpec(self.stage.scheduler, SchedulerMixin)])
+
+    @torch.no_grad()
+    def __call__(self, components, block_state, i, t):
+        scheduler = getattr(components, self.stage.scheduler)
+        prediction = block_state.noise_pred
+        block_state.latents = scheduler.step(prediction, t, block_state.latents.to(prediction.device)).prev_sample
+        return components, block_state
+
+
+class Pixal3DDenoiseLoop(LoopSequentialPipelineBlocks):
+    model_name = "pixal3d"
+
+    def __init__(self, stage: Pixal3DStage):
+        self.stage = stage
+        self.block_classes = [Pixal3DLoopDenoiser(stage), Pixal3DLoopAfterDenoiser(stage)]
+        self.block_names = ["denoiser", "after_denoiser"]
+        super().__init__()
+
+    @property
+    def description(self):
+        return (
+            "Iteratively denoise one stage's latents over `timesteps`. Each iteration runs:\n"
+            " - `Pixal3DLoopDenoiser`: flow model per guider batch, then guidance\n"
+            " - `Pixal3DLoopAfterDenoiser`: Euler step"
+        )
+
+    @property
+    def loop_inputs(self):
+        inputs = [
+            InputParam("timesteps", required=True, type_hint=list, description="From the set-timesteps step."),
+            InputParam("latents", required=True, type_hint=torch.Tensor, description="From the prepare-latents step."),
+        ]
+        if self.stage.sparse:
+            inputs.append(InputParam("coords", required=True, type_hint=torch.Tensor))
+        if self.stage.concat_cond:
+            inputs.append(InputParam("concat_cond", required=True, type_hint=torch.Tensor))
+        return inputs
+
+    @property
+    def loop_intermediate_outputs(self):
+        return [OutputParam("latents", type_hint=torch.Tensor, description="Denoised latents of this stage.")]
+
+    @torch.no_grad()
+    def __call__(self, components, state):
+        b = self.get_block_state(state)
+        device = components._execution_device
+        b.latents = b.latents.to(device)
+        b.global_features = b.global_features.to(device)
+        b.projected_features = b.projected_features.to(device)
+        b.negative_global_features = torch.zeros_like(b.global_features)
+        b.negative_projected_features = torch.zeros_like(b.projected_features)
+        b.model_coords = b.coords[None, :, 1:].to(device) if self.stage.sparse else None
+        b.model_concat_cond = b.concat_cond.to(device) if self.stage.concat_cond else None
+        with self.progress_bar(total=len(b.timesteps)) as progress_bar:
+            for i, t in enumerate(b.timesteps):
+                components, b = self.loop_step(components, b, i=i, t=t)
+                progress_bar.update()
+        self.set_block_state(state, b)
+        return components, state

encoders.py

@@ -0,0 +1,209 @@
+import numpy as np
+import torch
+from diffusers import ModelMixin
+from diffusers.modular_pipelines.modular_pipeline_utils import ComponentSpec, InputParam, OutputParam
+from PIL import Image
+from torch.nn import functional as F
+from transformers import DINOv3ViTModel
+
+from .projection import ProjGrid
+from .resolvable import ResolvableBlock
+
+
+class Pixal3DImageEncodeStep(ResolvableBlock):
+    model_name = "pixal3d"
+
+    def __init__(self, image_size: int):
+        self.image_size = image_size
+        super().__init__()
+
+    @property
+    def description(self):
+        return f"Encode the prepared image at {self.image_size}px into DINOv3 global tokens and patch features."
+
+    @property
+    def expected_components(self):
+        return self.resolve([ComponentSpec("image_encoder", DINOv3ViTModel)])
+
+    @property
+    def inputs(self):
+        return [
+            InputParam(
+                "image",
+                required=True,
+                type_hint=Image.Image,
+                description="Prepared RGB or RGBA image; RGBA is composited on black. No background removal or crop.",
+            ),
+        ]
+
+    @property
+    def intermediate_outputs(self):
+        return [
+            OutputParam("image_tensor", type_hint=torch.Tensor, description="Resized image, B 3 H W in [0, 1]."),
+            OutputParam("global_features", type_hint=torch.Tensor, description="CLS and register tokens."),
+            OutputParam("patch_features", type_hint=torch.Tensor, description="Patch features, B C h w."),
+        ]
+
+    @torch.no_grad()
+    def __call__(self, components, state):
+        block = self.get_block_state(state)
+        image = block.image
+        if not isinstance(image, Image.Image):
+            raise TypeError("image must be a prepared PIL image")
+        if image.mode == "RGBA":
+            background = Image.new("RGBA", image.size, (0, 0, 0, 255))
+            image = Image.alpha_composite(background, image).convert("RGB")
+        image = image.convert("RGB").resize((self.image_size, self.image_size), Image.Resampling.LANCZOS)
+        pixels = torch.from_numpy(np.array(image).astype(np.float32) / 255).permute(2, 0, 1)[None]
+        model = components.image_encoder
+        pixels = pixels.to(model.device)
+        mean = pixels.new_tensor([0.485, 0.456, 0.406])[None, :, None, None]
+        std = pixels.new_tensor([0.229, 0.224, 0.225])[None, :, None, None]
+        # The reference reads the last layer before the encoder's final norm and applies an unlearned layer norm.
+        normalized = (pixels - mean) / std
+        raw = model(normalized.to(model.dtype), output_hidden_states=True).hidden_states[-1]
+        features = F.layer_norm(raw, raw.shape[-1:])
+        global_count = 1 + model.config.num_register_tokens
+        patches = self.image_size // model.config.patch_size
+        block.image_tensor = pixels
+        block.global_features = features[:, :global_count]
+        block.patch_features = (
+            features[:, global_count:].reshape(1, patches, patches, -1).permute(0, 3, 1, 2).contiguous()
+        )
+        self.set_block_state(state, block)
+        return components, state
+
+
+class Pixal3DUpsampleFeaturesStep(ResolvableBlock):
+    model_name = "pixal3d"
+
+    def __init__(self, feature_size: int):
+        self.feature_size = feature_size
+        super().__init__()
+
+    @property
+    def description(self):
+        return f"Upsample the patch features to {self.feature_size}x{self.feature_size} with NAF."
+
+    @property
+    def expected_components(self):
+        return self.resolve([ComponentSpec("feature_upsampler", ModelMixin)])
+
+    @property
+    def inputs(self):
+        return [
+            InputParam("image_tensor", required=True, type_hint=torch.Tensor),
+            InputParam("patch_features", required=True, type_hint=torch.Tensor),
+        ]
+
+    @property
+    def intermediate_outputs(self):
+        return [OutputParam("upsampled_features", type_hint=torch.Tensor, description="NAF-upsampled features.")]
+
+    @torch.no_grad()
+    def __call__(self, components, state):
+        b = self.get_block_state(state)
+        model = components.feature_upsampler
+        b.upsampled_features = model(
+            b.image_tensor.to(device=model.device, dtype=model.dtype),
+            b.patch_features.to(device=model.device, dtype=model.dtype),
+            (self.feature_size, self.feature_size),
+        )
+        self.set_block_state(state, b)
+        return components, state
+
+
+class Pixal3DProjectFeaturesStep(ResolvableBlock):
+    """Project image features onto the stage's 3D grid.
+
+    Args:
+        image_size: Resolution the features were encoded at.
+        grid_resolution: Projection grid size. `None` derives it from `actual_resolution // 16`.
+        grid_from_model: Read the grid size from this component's `config.resolution` instead (dense stage).
+        sparse: Sample only at `coords` and concatenate the NAF-upsampled features.
+    """
+
+    model_name = "pixal3d"
+
+    def __init__(self, image_size: int, grid_resolution: int | None = None, grid_from_model=None, sparse=True):
+        self.image_size = image_size
+        self.grid_resolution = grid_resolution
+        self.grid_from_model = grid_from_model
+        self.sparse = sparse
+        self._grids = {}
+        super().__init__()
+
+    @property
+    def description(self):
+        target = "selected sparse coordinates" if self.sparse else "a dense grid"
+        return f"Project {self.image_size}px image features onto {target}."
+
+    @property
+    def expected_components(self):
+        return self.resolve([ComponentSpec(self.grid_from_model, ModelMixin)] if self.grid_from_model else [])
+
+    @property
+    def inputs(self):
+        inputs = [
+            InputParam("patch_features", required=True, type_hint=torch.Tensor),
+            InputParam(
+                "camera_angle_x", required=True, type_hint=float, description="Horizontal field of view, radians."
+            ),
+            InputParam(
+                "distance",
+                required=True,
+                type_hint=float,
+                description="Camera distance in the normalized source coordinate system.",
+            ),
+            InputParam(
+                "mesh_scale", default=1.0, type_hint=float, description="Object scale; 1.0 fills the unit box."
+            ),
+        ]
+        if self.sparse:
+            inputs += [
+                InputParam("upsampled_features", required=True, type_hint=torch.Tensor),
+                InputParam("coords", required=True, type_hint=torch.Tensor, description="Sparse (N, 4) coordinates."),
+            ]
+        if self.grid_resolution is None and self.grid_from_model is None:
+            inputs.append(InputParam("actual_resolution", required=True, type_hint=int))
+        return inputs
+
+    @property
+    def intermediate_outputs(self):
+        return [OutputParam("projected_features", type_hint=torch.Tensor, description="B N C projected features.")]
+
+    def _grid(self, grid_resolution, device):
+        key = (grid_resolution, device)
+        if key not in self._grids:
+            self._grids[key] = ProjGrid(grid_resolution, self.image_size).to(device)
+        return self._grids[key]
+
+    @torch.no_grad()
+    def __call__(self, components, state):
+        b = self.get_block_state(state)
+        if not 0 < b.camera_angle_x < np.pi or b.distance <= 0 or b.mesh_scale <= 0:
+            raise ValueError("FOV must be in (0, pi), distance and mesh_scale must be positive")
+        if self.grid_from_model:
+            grid_resolution = getattr(components, self.grid_from_model).config.resolution
+        elif self.grid_resolution is None:
+            grid_resolution = b.actual_resolution // 16
+        else:
+            grid_resolution = self.grid_resolution
+        device = b.patch_features.device
+        grid = self._grid(grid_resolution, device)
+        camera = [torch.tensor([v], device=device) for v in (b.camera_angle_x, b.distance, b.mesh_scale)]
+        indices = None
+        if self.sparse:
+            coords = b.coords.long().to(device)
+            if coords.ndim != 2 or coords.shape[1] != 4 or (coords[:, 0] != 0).any():
+                raise ValueError("Coordinates must have shape (N,4) with batch index 0")
+            if ((coords[:, 1:] < 0) | (coords[:, 1:] >= grid_resolution)).any():
+                raise ValueError("Coordinates fall outside the projection grid")
+            indices = (coords[:, 1] * grid_resolution + coords[:, 2]) * grid_resolution + coords[:, 3]
+        projected = grid(b.patch_features.float(), *camera, grid_indices=indices)
+        if self.sparse:
+            upsampled = grid(b.upsampled_features.to(device).float(), *camera, grid_indices=indices)
+            projected = torch.cat([projected, upsampled], -1)
+        b.projected_features = projected
+        self.set_block_state(state, b)
+        return components, state

guidance.py

@@ -0,0 +1,82 @@
+from diffusers.configuration_utils import register_to_config
+from diffusers.guiders.guider_utils import BaseGuidance, GuiderOutput
+
+
+class Pixal3DGuidance(BaseGuidance):
+    """Classifier-free guidance with Pixal3D's sigma interval and clean-sample (x0) rescaling.
+
+    Guidance is active while the current sigma lies inside `guidance_interval`; outside it, or with
+    `guidance_scale=1`, only the conditional prediction is computed. The rescale matches the standard
+    deviation of the guided x0 estimate to the conditional one, which needs the current `latents`.
+    """
+
+    _input_predictions = ["pred_cond", "pred_uncond"]
+
+    @register_to_config
+    def __init__(
+        self,
+        guidance_scale=7.5,
+        guidance_rescale=0.7,
+        guidance_interval=(0.6, 1.0),
+        sigma_min=1e-5,
+        enabled=True,
+    ):
+        super().__init__(enabled=enabled)
+
+    @property
+    def _scale(self):
+        sigma = self._timestep
+        low, high = self.config.guidance_interval
+        if self._enabled and sigma is not None and low <= sigma <= high:
+            return self.config.guidance_scale
+        return 1.0
+
+    @property
+    def is_conditional(self):
+        return self._count_prepared == 1
+
+    @property
+    def num_conditions(self):
+        return 1 if self._scale in (0, 1) else 2
+
+    def _predictions(self):
+        # guidance_scale=0 is purely unconditional; 1 (or outside the interval) is purely conditional.
+        if self._scale == 0:
+            return ["pred_uncond"]
+        return self._input_predictions[: self.num_conditions]
+
+    def prepare_inputs(self, data):
+        return [self._prepare_batch(data, i, name) for i, name in self._identified_predictions()]
+
+    def prepare_inputs_from_block_state(self, data, input_fields):
+        return [
+            self._prepare_batch_from_block_state(input_fields, data, i, name)
+            for i, name in self._identified_predictions()
+        ]
+
+    def _identified_predictions(self):
+        return [(self._input_predictions.index(name), name) for name in self._predictions()]
+
+    def __call__(self, data):
+        predictions = {getattr(d, self._identifier_key): d.noise_pred for d in data}
+        return self.forward(**predictions, latents=data[0].latents)
+
+    def forward(self, pred_cond=None, pred_uncond=None, latents=None):
+        if pred_cond is None:
+            return GuiderOutput(pred=pred_uncond)
+        if pred_uncond is None:
+            return GuiderOutput(pred=pred_cond)
+        latents = latents.to(pred_cond.device)
+        pred_uncond = pred_uncond.to(pred_cond.device)
+        scale = self._scale
+        result = scale * pred_cond + (1 - scale) * pred_uncond
+        if self.config.guidance_rescale > 0:
+            sigma_min = self.config.sigma_min
+            sigma = sigma_min + (1 - sigma_min) * self._timestep
+            clean_pos = (1 - sigma_min) * latents - sigma * pred_cond
+            clean_cfg = (1 - sigma_min) * latents - sigma * result
+            dims = list(range(1, latents.ndim))
+            rescaled = clean_cfg * (clean_pos.std(dim=dims, keepdim=True) / clean_cfg.std(dim=dims, keepdim=True))
+            clean = self.config.guidance_rescale * rescaled + (1 - self.config.guidance_rescale) * clean_cfg
+            result = ((1 - sigma_min) * latents - clean) / sigma
+        return GuiderOutput(pred=result, pred_cond=pred_cond, pred_uncond=pred_uncond)

modular_blocks_pixal3d.py

@@ -0,0 +1,183 @@
+import copy
+from pathlib import Path
+
+from diffusers.modular_pipelines import ModularPipeline, SequentialPipelineBlocks
+from diffusers.modular_pipelines.modular_pipeline_utils import InsertableDict, OutputParam
+from diffusers.utils.dynamic_modules_utils import get_class_from_dynamic_module
+from huggingface_hub import file_exists
+from huggingface_hub.errors import HFValidationError
+from transformers import DINOv3ViTModel
+
+from .decoders import (
+    Pixal3DShapeDecodeStep,
+    Pixal3DShapeUpsampleStep,
+    Pixal3DStructureDecodeStep,
+    Pixal3DTextureDecodeStep,
+)
+from .denoise import (
+    SHAPE_HIGH_STAGE,
+    SHAPE_LOW_STAGE,
+    STRUCTURE_STAGE,
+    TEXTURE_STAGE,
+    Pixal3DDenoiseLoop,
+    Pixal3DPrepareLatentsStep,
+    Pixal3DSetTimestepsStep,
+)
+from .encoders import Pixal3DImageEncodeStep, Pixal3DProjectFeaturesStep, Pixal3DUpsampleFeaturesStep
+from .guidance import Pixal3DGuidance  # noqa: F401
+from .normalization import SHAPE_MEAN, SHAPE_STD, TEXTURE_MEAN, TEXTURE_STD  # noqa: F401
+
+# Expose transitive dependencies to the Hub custom-code collector.
+from .projection import ProjGrid  # noqa: F401
+from .resolvable import ResolvableBlock  # noqa: F401
+from .stages import Pixal3DShapeNormalizeStep, Pixal3DTextureConditionStep, Pixal3DTextureNormalizeStep
+
+
+def _stage(prefix, stage):
+    return [
+        (f"{prefix}_prepare_latents", Pixal3DPrepareLatentsStep(stage)),
+        (f"{prefix}_set_timesteps", Pixal3DSetTimestepsStep(stage)),
+        (f"{prefix}_denoise", Pixal3DDenoiseLoop(stage)),
+    ]
+
+
+Pixal3DBlockSequence = InsertableDict(
+    [
+        ("structure_image_encoder", Pixal3DImageEncodeStep(image_size=512)),
+        (
+            "structure_projection",
+            Pixal3DProjectFeaturesStep(image_size=512, grid_from_model=STRUCTURE_STAGE.transformer, sparse=False),
+        ),
+        *_stage("structure", STRUCTURE_STAGE),
+        ("structure_decode", Pixal3DStructureDecodeStep()),
+        ("shape_low_features", Pixal3DUpsampleFeaturesStep(feature_size=512)),
+        ("shape_low_projection", Pixal3DProjectFeaturesStep(image_size=512, grid_resolution=32)),
+        *_stage("shape_low", SHAPE_LOW_STAGE),
+        ("shape_low_normalize", Pixal3DShapeNormalizeStep()),
+        ("shape_upsample", Pixal3DShapeUpsampleStep()),
+        ("shape_high_image_encoder", Pixal3DImageEncodeStep(image_size=1024)),
+        ("shape_high_features", Pixal3DUpsampleFeaturesStep(feature_size=512)),
+        ("shape_high_projection", Pixal3DProjectFeaturesStep(image_size=1024)),
+        *_stage("shape_high", SHAPE_HIGH_STAGE),
+        ("shape_high_normalize", Pixal3DShapeNormalizeStep()),
+        ("texture_condition", Pixal3DTextureConditionStep()),
+        ("texture_features", Pixal3DUpsampleFeaturesStep(feature_size=1024)),
+        ("texture_projection", Pixal3DProjectFeaturesStep(image_size=1024)),
+        *_stage("texture", TEXTURE_STAGE),
+        ("texture_normalize", Pixal3DTextureNormalizeStep()),
+        ("shape_decode", Pixal3DShapeDecodeStep()),
+        ("texture_decode", Pixal3DTextureDecodeStep()),
+    ]
+)
+
+
+def _component_source(entry, name, repo, revision):
+    """Return (module, class, repo, subfolder, revision) for one `modular_model_index.json` entry.
+
+    The model repository's own index names a module file (or `transformers`); its components load from `repo`, so a
+    local copy stays local even if the index records Hub ids. An index written by `pipe.save_pretrained()` stores
+    dotted module paths into the dynamic-module cache (or null with a `type_hint`) and the location of each component.
+    """
+    module, cls = entry[0], entry[1]
+    spec = entry[2] if len(entry) > 2 and entry[2] else {}
+    if module is not None and "." not in module:
+        return module, cls, repo, spec.get("subfolder") or name, revision
+    if module is None:
+        module, cls = spec["type_hint"]
+    return (
+        module.rsplit(".", 1)[-1],
+        cls,
+        spec.get("pretrained_model_name_or_path") or repo,
+        spec.get("subfolder") or name,
+        spec.get("revision") or revision,
+    )
+
+
+def _has_file(repo, filename, revision):
+    if Path(str(repo)).is_dir():
+        return (Path(str(repo)) / filename).is_file()
+    try:
+        return file_exists(str(repo), filename, revision=revision)
+    except (HFValidationError, OSError):
+        return False
+
+
+def _code_repo(module, candidates):
+    """First (repo, revision) holding `<module>.py`.
+
+    A saved pipeline's components point at the save directory, which holds weights but not the model code.
+    """
+    for repo, revision in candidates:
+        if _has_file(repo, f"{module}.py", revision):
+            return repo, revision
+    raise OSError(
+        f"Could not find {module}.py in {[str(repo) for repo, _ in candidates]}. A pipeline saved with "
+        "`save_pretrained` keeps only weights; copy the model repository's .py files into it."
+    )
+
+
+def _leaves(block):
+    if not block.sub_blocks:
+        yield block
+    for child in block.sub_blocks.values():
+        yield from _leaves(child)
+
+
+class Pixal3DBlocks(SequentialPipelineBlocks):
+    model_name = "pixal3d"
+    block_classes = Pixal3DBlockSequence.values()
+    block_names = Pixal3DBlockSequence.keys()
+
+    def __init__(self, components_repo=None, components_revision=None, trust_components_code=False):
+        super().__init__()
+        self.components_repo = components_repo
+        self.components_revision = components_revision
+        if components_repo is None:
+            return
+        if not trust_components_code:
+            raise ValueError("Set trust_components_code=True to load the model repository code")
+        index = ModularPipeline.load_config(components_repo, revision=components_revision)
+        names = [s.name for s in super().expected_components if s.default_creation_method == "from_pretrained"]
+        sources = {name: _component_source(index[name], name, components_repo, components_revision) for name in names}
+        # Code can live in any repo the index references, the given repo first.
+        candidates = [(components_repo, components_revision)] + [(r, v) for _, _, r, _, v in sources.values()]
+        candidates = list(dict.fromkeys(candidates))
+        overrides = {}
+        for name, (module, cls, repo, subfolder, revision) in sources.items():
+            if name == "image_encoder":
+                type_hint = DINOv3ViTModel
+            else:
+                code_repo, code_revision = _code_repo(module, candidates)
+                type_hint = get_class_from_dynamic_module(
+                    code_repo,
+                    module_file=f"{module}.py",
+                    class_name=cls,
+                    revision=code_revision,
+                    trust_remote_code=True,
+                )
+            overrides[name] = dict(
+                type_hint=type_hint, pretrained_model_name_or_path=str(repo), subfolder=subfolder, revision=revision
+            )
+        # The block instances are shared module-level objects; give this pipeline its own copies to carry the specs.
+        self.sub_blocks = copy.deepcopy(self.sub_blocks)
+        for leaf in _leaves(self):
+            leaf.component_overrides = overrides
+
+    @property
+    def description(self):
+        return "Pixal3D single-image generation to flexible-dual-grid geometry and voxel material attributes."
+
+    @property
+    def outputs(self):
+        return [
+            OutputParam(
+                "geometry", type_hint=dict, description="Flexible dual grid: coords, offsets, flags, weights."
+            ),
+            OutputParam("materials", type_hint=dict, description="Per-cell base color, metallic, roughness, alpha."),
+            OutputParam(
+                "actual_resolution", type_hint=int, description="Effective grid resolution (1536 can fall back)."
+            ),
+            OutputParam("shape_latents", type_hint=object),
+            OutputParam("texture_latents", type_hint=object),
+            OutputParam("coords", type_hint=object),
+        ]

modular_config.json

@@ -0,0 +1,7 @@
+{
+  "_class_name": "Pixal3DBlocks",
+  "_diffusers_version": "0.41.0.dev0",
+  "auto_map": {
+    "ModularPipelineBlocks": "modular_blocks_pixal3d.Pixal3DBlocks"
+  }
+}

normalization.py

@@ -0,0 +1,136 @@
+SHAPE_MEAN = [
+    0.781296,
+    0.018091,
+    -0.495192,
+    -0.558457,
+    1.06053,
+    0.093252,
+    1.518149,
+    -0.933218,
+    -0.732996,
+    2.604095,
+    -0.118341,
+    -2.143904,
+    0.495076,
+    -2.179512,
+    -2.130751,
+    -0.996944,
+    0.261421,
+    -2.217463,
+    1.260067,
+    -0.150213,
+    3.790713,
+    1.481266,
+    -1.046058,
+    -1.523667,
+    -0.059621,
+    2.22078,
+    1.621212,
+    0.87723,
+    0.567247,
+    -3.175944,
+    -3.186688,
+    1.578665,
+]
+SHAPE_STD = [
+    5.972266,
+    4.706852,
+    5.44501,
+    5.209927,
+    5.32022,
+    4.547237,
+    5.020802,
+    5.444004,
+    5.226681,
+    5.683095,
+    4.831436,
+    5.286469,
+    5.652043,
+    5.367606,
+    5.525084,
+    4.730578,
+    4.805265,
+    5.124013,
+    5.530808,
+    5.619001,
+    5.10393,
+    5.41767,
+    5.269677,
+    5.547194,
+    5.634698,
+    5.235274,
+    6.110351,
+    5.511298,
+    6.237273,
+    4.879207,
+    5.347008,
+    5.405691,
+]
+TEXTURE_MEAN = [
+    3.501659,
+    2.212398,
+    2.226094,
+    0.251093,
+    -0.026248,
+    -0.687364,
+    0.439898,
+    -0.928075,
+    0.029398,
+    -0.339596,
+    -0.869527,
+    1.038479,
+    -0.972385,
+    0.126042,
+    -1.129303,
+    0.455149,
+    -1.209521,
+    2.069067,
+    0.544735,
+    2.569128,
+    -0.323407,
+    2.293,
+    -1.925608,
+    -1.217717,
+    1.213905,
+    0.971588,
+    -0.023631,
+    0.10675,
+    2.021786,
+    0.250524,
+    -0.662387,
+    -0.768862,
+]
+TEXTURE_STD = [
+    2.665652,
+    2.743913,
+    2.765121,
+    2.595319,
+    3.037293,
+    2.291316,
+    2.144656,
+    2.911822,
+    2.969419,
+    2.501689,
+    2.154811,
+    3.163343,
+    2.621215,
+    2.381943,
+    3.186697,
+    3.021588,
+    2.295916,
+    3.234985,
+    3.233086,
+    2.26014,
+    2.874801,
+    2.810596,
+    3.29272,
+    2.674999,
+    2.680878,
+    2.372054,
+    2.451546,
+    2.353556,
+    2.995195,
+    2.379849,
+    2.786195,
+    2.77519,
+]

projection.py

@@ -0,0 +1,87 @@
+# Adapted from TencentARC/Pixal3D, Copyright (c) 2026 Tencent, MIT (see LICENSE-PIXAL3D).
+
+import torch
+from torch import nn
+from torch.nn import functional as F
+
+
+def project_points_to_image_batch(
+    points_3d: torch.Tensor,
+    transform_matrix: torch.Tensor,
+    camera_angle_x: torch.Tensor,
+    resolution: int,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    """Project 3D points to pixel coordinates with a pinhole camera (32 mm sensor width).
+
+    Args:
+        points_3d: B N 3 points.
+        transform_matrix: B 4 4 camera-to-world matrices.
+        camera_angle_x: B horizontal fields of view, radians.
+        resolution: Square image resolution in pixels.
+
+    Returns:
+        Pixel coordinates (B N 2), depth (B N) and a mask of points inside the view (B N).
+    """
+    batch, count = points_3d.shape[:2]
+    ones = torch.ones(batch, count, 1, device=points_3d.device, dtype=points_3d.dtype)
+    points_homogeneous = torch.cat([points_3d, ones], dim=-1)
+    world_to_camera = torch.linalg.inv(transform_matrix.float()).to(transform_matrix.dtype)
+    points_camera = torch.bmm(points_homogeneous, world_to_camera.transpose(-2, -1))[..., :3]
+    x_cam, y_cam, z_cam = points_camera.unbind(-1)
+    depth = -z_cam
+    sensor_width = 32.0
+    focal_length = 16.0 / torch.tan(camera_angle_x / 2.0)
+    focal_length_pixels = (focal_length * resolution / sensor_width).unsqueeze(1)
+    x_ndc = focal_length_pixels * x_cam / (-z_cam + 1e-08)
+    y_ndc = focal_length_pixels * y_cam / (-z_cam + 1e-08)
+    x_pixel = x_ndc + resolution / 2.0
+    y_pixel = -y_ndc + resolution / 2.0
+    valid_mask = (x_pixel >= 0) & (x_pixel < resolution) & (y_pixel >= 0) & (y_pixel < resolution) & (depth > 0)
+    return torch.stack([x_pixel, y_pixel], dim=-1), depth, valid_mask
+
+
+class ProjGrid(nn.Module):
+    """Project a 3D grid (or selected grid cells) into the front view and sample image features there."""
+
+    def __init__(self, grid_resolution: int, image_resolution: int):
+        super().__init__()
+        self.image_resolution = image_resolution
+        one_dim = torch.linspace(-1, 1, grid_resolution)
+        x, y, z = torch.meshgrid(one_dim, one_dim, one_dim, indexing="ij")
+        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(torch.stack((x, y, z), dim=-1), rotation_matrix.T).reshape(-1, 3)
+        self.register_buffer("grid_points", grid_points)
+        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,
+        grid_indices: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        """Sample a B C H W feature map at the projected grid points; returns B N C."""
+        batch = features_map.shape[0]
+        grid_points = self.grid_points
+        if grid_indices is not None:
+            grid_points = grid_points[grid_indices]
+        grid_points = grid_points.expand(batch, -1, -1)
+        grid_points = grid_points / mesh_scale.unsqueeze(-1).unsqueeze(-1) / 2
+        transform_matrix = self.front_view_transform_matrix.expand(batch, -1, -1).clone()
+        transform_matrix[:, 1, 3] = -distance
+        image_points, _, _ = project_points_to_image_batch(
+            grid_points, transform_matrix, camera_angle_x, self.image_resolution
+        )
+        image_points_norm = (image_points + 0.5) / self.image_resolution * 2 - 1
+        sample_grid = image_points_norm.view(batch, -1, 1, 2)
+        sampled = F.grid_sample(features_map, sample_grid, mode="bilinear", align_corners=False, padding_mode="border")
+        return sampled.squeeze(-1).permute(0, 2, 1)

resolvable.py

@@ -0,0 +1,19 @@
+from dataclasses import replace
+
+from diffusers.modular_pipelines import ModularPipelineBlocks
+
+
+class ResolvableBlock(ModularPipelineBlocks):
+    """Leaf block whose generic component specs can carry the model repository's resolved classes and location.
+
+    `Pixal3DBlocks(components_repo=...)` fills `component_overrides`, so a sub-block popped out of the pipeline still
+    loads the right custom classes from the right repository.
+    """
+
+    component_overrides = {}
+
+    def resolve(self, specs):
+        return [
+            replace(spec, **self.component_overrides[spec.name]) if spec.name in self.component_overrides else spec
+            for spec in specs
+        ]

stages.py

@@ -0,0 +1,75 @@
+import torch
+from diffusers.modular_pipelines.modular_pipeline_utils import InputParam, OutputParam
+
+from .normalization import SHAPE_MEAN, SHAPE_STD, TEXTURE_MEAN, TEXTURE_STD
+from .resolvable import ResolvableBlock
+
+
+class Pixal3DShapeNormalizeStep(ResolvableBlock):
+    model_name = "pixal3d"
+
+    @property
+    def description(self):
+        return "Restore shape latents to the decoder training distribution."
+
+    @property
+    def inputs(self):
+        return [InputParam("latents", required=True, type_hint=torch.Tensor, description="Denoised shape latents.")]
+
+    @property
+    def intermediate_outputs(self):
+        return [OutputParam("shape_latents", type_hint=torch.Tensor, description="Shape latents for the decoder.")]
+
+    def __call__(self, components, state):
+        b = self.get_block_state(state)
+        b.shape_latents = b.latents * b.latents.new_tensor(SHAPE_STD) + b.latents.new_tensor(SHAPE_MEAN)
+        self.set_block_state(state, b)
+        return components, state
+
+
+class Pixal3DTextureConditionStep(ResolvableBlock):
+    model_name = "pixal3d"
+
+    @property
+    def description(self):
+        return "Normalize the shape latents into the texture flow model's concatenated condition."
+
+    @property
+    def inputs(self):
+        return [InputParam("shape_latents", required=True, type_hint=torch.Tensor)]
+
+    @property
+    def intermediate_outputs(self):
+        return [OutputParam("concat_cond", type_hint=torch.Tensor, description="Normalized shape latents.")]
+
+    def __call__(self, components, state):
+        b = self.get_block_state(state)
+        b.concat_cond = (b.shape_latents - b.shape_latents.new_tensor(SHAPE_MEAN)) / b.shape_latents.new_tensor(
+            SHAPE_STD
+        )
+        self.set_block_state(state, b)
+        return components, state
+
+
+class Pixal3DTextureNormalizeStep(ResolvableBlock):
+    model_name = "pixal3d"
+
+    @property
+    def description(self):
+        return "Restore material latents to the decoder training distribution."
+
+    @property
+    def inputs(self):
+        return [InputParam("latents", required=True, type_hint=torch.Tensor, description="Denoised material latents.")]
+
+    @property
+    def intermediate_outputs(self):
+        return [
+            OutputParam("texture_latents", type_hint=torch.Tensor, description="Material latents for the decoder.")
+        ]
+
+    def __call__(self, components, state):
+        b = self.get_block_state(state)
+        b.texture_latents = b.latents * b.latents.new_tensor(TEXTURE_STD) + b.latents.new_tensor(TEXTURE_MEAN)
+        self.set_block_state(state, b)
+        return components, state