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