| import torch |
| from diffusers import ModelMixin, SchedulerMixin |
| from diffusers.modular_pipelines.modular_pipeline_utils import ComponentSpec, InputParam, OutputParam |
| from diffusers.utils.torch_utils import randn_tensor |
|
|
| from .operations import ResolvableBlock |
|
|
|
|
| class KimodoLatentsStep(ResolvableBlock): |
| model_name = "kimodo" |
|
|
| @property |
| def description(self): |
| return "Draw the initial motion-feature noise for an explicit frame count." |
|
|
| @property |
| def expected_components(self): |
| return self.resolve([ComponentSpec("transformer", ModelMixin)]) |
|
|
| @property |
| def inputs(self): |
| return [ |
| InputParam( |
| "num_frames", type_hint=int, default=120, description="Motion samples at the model frame rate." |
| ), |
| InputParam( |
| "seed", |
| type_hint=int, |
| default=42, |
| description="Seed for a generator on the execution device; ignored when `generator` is given.", |
| ), |
| InputParam.template("generator"), |
| ] |
|
|
| @property |
| def intermediate_outputs(self): |
| return [ |
| OutputParam("latents", type_hint=torch.Tensor, description="Noisy motion features, (1, frames, features).") |
| ] |
|
|
| def __call__(self, components, state): |
| block = self.get_block_state(state) |
| if block.num_frames < 2: |
| raise ValueError("Motion requires at least two frames") |
| device = components._execution_device |
| generator = block.generator |
| if generator is None: |
| generator = torch.Generator(device=device).manual_seed(block.seed) |
| transformer = components.transformer |
| block.latents = randn_tensor( |
| (1, block.num_frames, transformer.motion_rep_dim), |
| generator=generator, |
| device=device, |
| dtype=transformer.dtype, |
| ) |
| self.set_block_state(state, block) |
| return components, state |
|
|
|
|
| class KimodoTimestepsStep(ResolvableBlock): |
| model_name = "kimodo" |
|
|
| @property |
| def description(self): |
| return "Set the Kimodo DDIM schedule for this run." |
|
|
| @property |
| def expected_components(self): |
| return self.resolve([ComponentSpec("scheduler", SchedulerMixin)]) |
|
|
| @property |
| def inputs(self): |
| return [InputParam.template("num_inference_steps", default=100)] |
|
|
| @property |
| def intermediate_outputs(self): |
| return [OutputParam("timesteps", type_hint=torch.Tensor, description="Descending training-step indices.")] |
|
|
| def __call__(self, components, state): |
| block = self.get_block_state(state) |
| components.scheduler.set_timesteps(block.num_inference_steps, device=components._execution_device) |
| block.timesteps = components.scheduler.timesteps |
| self.set_block_state(state, block) |
| return components, state |
|
|
