| from collections.abc import Callable |
|
|
| import torch |
| from diffusers import ModelMixin, SchedulerMixin |
| from diffusers.configuration_utils import FrozenDict |
| from diffusers.modular_pipelines import LoopSequentialPipelineBlocks |
| from diffusers.modular_pipelines.modular_pipeline_utils import ComponentSpec, InputParam, OutputParam |
|
|
| from .guidance import KimodoGuidance |
| from .operations import ResolvableBlock, checkpoint |
|
|
|
|
| class KimodoLoopDenoiser(ResolvableBlock): |
| model_name = "kimodo" |
| |
| _guider_input_fields = {"encoder_hidden_states": ("prompt_embeds", "negative_prompt_embeds")} |
|
|
| @property |
| def description(self): |
| return "Step within the denoising loop: run the transformer for each guider batch, then apply guidance." |
|
|
| @property |
| def expected_components(self): |
| return self.resolve( |
| [ |
| ComponentSpec("transformer", ModelMixin), |
| ComponentSpec( |
| "guider", KimodoGuidance, config=FrozenDict(text_scale=2.0), default_creation_method="from_config" |
| ), |
| ] |
| ) |
|
|
| @property |
| def inputs(self): |
| return [ |
| InputParam("prompt_embeds", required=True, type_hint=torch.Tensor, description="LLM2Vec prompt features.") |
| ] |
|
|
| @property |
| def intermediate_outputs(self): |
| return [OutputParam("noise_pred", type_hint=torch.Tensor, description="Guided clean-motion prediction.")] |
|
|
| @torch.no_grad() |
| def __call__(self, components, block, i, t): |
| latents = block.latents |
| guider = components.guider |
| guider.set_state(step=i, num_inference_steps=len(block.timesteps), timestep=t) |
| batches = guider.prepare_inputs_from_block_state(block, self._guider_input_fields) |
| timestep = t.reshape(1).repeat(latents.shape[0]) |
| heading = latents.new_zeros(latents.shape[0]) |
| for batch in batches: |
| guider.prepare_models(components.transformer) |
| batch.noise_pred = components.transformer( |
| hidden_states=latents, |
| timestep=timestep, |
| encoder_hidden_states=batch.encoder_hidden_states, |
| first_heading_angle=heading, |
| ).sample |
| guider.cleanup_models(components.transformer) |
| block.noise_pred = guider(batches).pred |
| return components, block |
|
|
|
|
| class KimodoLoopSchedulerStep(ResolvableBlock): |
| model_name = "kimodo" |
|
|
| @property |
| def description(self): |
| return "Step within the denoising loop: apply one DDIM update to the motion features." |
|
|
| @property |
| def expected_components(self): |
| return self.resolve([ComponentSpec("scheduler", SchedulerMixin)]) |
|
|
| @property |
| def inputs(self): |
| return [InputParam("noise_pred", required=True, type_hint=torch.Tensor)] |
|
|
| @property |
| def intermediate_outputs(self): |
| return [OutputParam("latents", type_hint=torch.Tensor, description="Updated normalized motion features.")] |
|
|
| def __call__(self, components, block, i, t): |
| block.latents = components.scheduler.step(block.noise_pred, t, block.latents).prev_sample |
| return components, block |
|
|
|
|
| class KimodoDenoiseStep(LoopSequentialPipelineBlocks): |
| model_name = "kimodo" |
| block_classes = [KimodoLoopDenoiser, KimodoLoopSchedulerStep] |
| block_names = ["denoiser", "scheduler_step"] |
|
|
| @property |
| def description(self): |
| return ( |
| "Iteratively denoise the motion features over `timesteps`. Each iteration runs:\n" |
| " - `KimodoLoopDenoiser`: transformer per guider batch, then guidance\n" |
| " - `KimodoLoopSchedulerStep`: DDIM update" |
| ) |
|
|
| @property |
| def loop_inputs(self): |
| return [ |
| InputParam("latents", type_hint=torch.Tensor, required=True, description="From the latents step."), |
| InputParam("timesteps", type_hint=torch.Tensor, required=True, description="From the timesteps step."), |
| InputParam( |
| "progress_callback", |
| type_hint=Callable, |
| default=None, |
| description="Called after each step with (step, total, stage).", |
| ), |
| InputParam("cancel_event", default=None, description="Optional threading.Event that cancels the run."), |
| ] |
|
|
| @property |
| def loop_intermediate_outputs(self): |
| return [OutputParam("latents", type_hint=torch.Tensor, description="Denoised normalized motion features.")] |
|
|
| @torch.no_grad() |
| def __call__(self, components, state): |
| block = self.get_block_state(state) |
| block.negative_prompt_embeds = torch.zeros_like(block.prompt_embeds) |
| for i, t in enumerate(block.timesteps): |
| checkpoint(block.cancel_event) |
| components, block = self.loop_step(components, block, i=i, t=t) |
| if block.progress_callback is not None: |
| block.progress_callback(i + 1, len(block.timesteps), "Denoising") |
| checkpoint(block.cancel_event) |
| self.set_block_state(state, block) |
| return components, state |
|
|
