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