| |
| |
| |
|
|
| from dataclasses import replace |
|
|
| from diffusers.modular_pipelines import ModularPipeline, SequentialPipelineBlocks |
| from diffusers.modular_pipelines.modular_pipeline_utils import OutputParam |
| from diffusers.utils.dynamic_modules_utils import get_class_from_dynamic_module |
|
|
| from .decoders import MotionBricksDecodeStep, MotionBricksPostprocessStep |
| from .encoders import MotionBricksGuidesStep |
| from .generation import MotionBricksPoseTokensStep, MotionBricksRootStep |
|
|
| |
| from .motion_utils import encode_guides |
|
|
|
|
| class MotionBricksBlocks(SequentialPipelineBlocks): |
| """Motion in-betweening with NVIDIA MotionBricks on the SOMA skeleton. |
| |
| Every component loads from `components_repo`, one subfolder per component, and so does its code, which requires |
| `trust_components_code=True`. |
| """ |
|
|
| model_name = "motionbricks" |
| block_classes = [ |
| MotionBricksGuidesStep(), |
| MotionBricksRootStep(), |
| MotionBricksPoseTokensStep(), |
| MotionBricksDecodeStep(), |
| MotionBricksPostprocessStep(), |
| ] |
| block_names = ["guides", "root", "pose_tokens", "decode", "postprocess"] |
|
|
| _requirements = {"diffusers": ">=0.41.0"} |
|
|
| 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 |
| self.component_types = {} |
| if components_repo is None: |
| return |
| if not trust_components_code: |
| raise ValueError("Set `trust_components_code=True` to load the model code from the components repository") |
| index = ModularPipeline.load_config(str(components_repo), revision=components_revision) |
| for spec in super().expected_components: |
| module_name, class_name, _ = index[spec.name] |
| self.component_types[spec.name] = get_class_from_dynamic_module( |
| str(components_repo), |
| subfolder=spec.name, |
| module_file=f"{module_name}.py", |
| class_name=class_name, |
| revision=components_revision, |
| trust_remote_code=True, |
| ) |
|
|
| @property |
| def expected_components(self): |
| if self.components_repo is None: |
| return super().expected_components |
| return [ |
| replace( |
| spec, |
| type_hint=self.component_types.get(spec.name, spec.type_hint), |
| pretrained_model_name_or_path=str(self.components_repo), |
| subfolder=spec.name, |
| revision=self.components_revision, |
| ) |
| for spec in super().expected_components |
| ] |
|
|
| @property |
| def description(self) -> str: |
| return ( |
| "Generates the transition from the last poses of a context motion to target poses, choosing its length " |
| "among the allowed ones." |
| ) |
|
|
| @property |
| def outputs(self): |
| return [ |
| OutputParam( |
| "motion", |
| type_hint=list, |
| description="One SOMA30 motion dict per batch item, 30 FPS, meters, +Y up.", |
| ), |
| OutputParam( |
| "placed_target", |
| type_hint=list, |
| description="`target` as used, one dict per batch item, moved by `target_offset` and `target_turn`.", |
| ), |
| ] |
|
|
