File size: 5,156 Bytes
9ed958c 903e660 9ed958c 903e660 9ed958c 903e660 9ed958c 903e660 9ed958c 903e660 9ed958c 903e660 9ed958c 903e660 9ed958c 903e660 9ed958c 903e660 9ed958c 903e660 9ed958c 903e660 9ed958c 903e660 9ed958c 903e660 9ed958c 903e660 9ed958c 903e660 9ed958c 903e660 9ed958c 903e660 9ed958c 903e660 9ed958c 903e660 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 | 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"
# Text-conditioned and unconditional batches; the unconditional prompt features are zeros.
_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
|