(
unet: SDUNet,
motion_modules: SDMotionModel = None,
controlnet: MultiControlNetManager = None,
sample = None,
timestep = None,
encoder_hidden_states = None,
ipadapter_kwargs_list = {},
controlnet_frames = None,
unet_batch_size = 1,
controlnet_batch_size = 1,
cross_frame_attention = False,
tiled=False,
tile_size=64,
tile_stride=32,
device = "cuda",
vram_limit_level = 0,
)
| 5 | |
| 6 | |
| 7 | def lets_dance( |
| 8 | unet: SDUNet, |
| 9 | motion_modules: SDMotionModel = None, |
| 10 | controlnet: MultiControlNetManager = None, |
| 11 | sample = None, |
| 12 | timestep = None, |
| 13 | encoder_hidden_states = None, |
| 14 | ipadapter_kwargs_list = {}, |
| 15 | controlnet_frames = None, |
| 16 | unet_batch_size = 1, |
| 17 | controlnet_batch_size = 1, |
| 18 | cross_frame_attention = False, |
| 19 | tiled=False, |
| 20 | tile_size=64, |
| 21 | tile_stride=32, |
| 22 | device = "cuda", |
| 23 | vram_limit_level = 0, |
| 24 | ): |
| 25 | # 0. Text embedding alignment (only for video processing) |
| 26 | if encoder_hidden_states.shape[0] != sample.shape[0]: |
| 27 | encoder_hidden_states = encoder_hidden_states.repeat(sample.shape[0], 1, 1, 1) |
| 28 | |
| 29 | # 1. ControlNet |
| 30 | # This part will be repeated on overlapping frames if animatediff_batch_size > animatediff_stride. |
| 31 | # I leave it here because I intend to do something interesting on the ControlNets. |
| 32 | controlnet_insert_block_id = 30 |
| 33 | if controlnet is not None and controlnet_frames is not None: |
| 34 | res_stacks = [] |
| 35 | # process controlnet frames with batch |
| 36 | for batch_id in range(0, sample.shape[0], controlnet_batch_size): |
| 37 | batch_id_ = min(batch_id + controlnet_batch_size, sample.shape[0]) |
| 38 | res_stack = controlnet( |
| 39 | sample[batch_id: batch_id_], |
| 40 | timestep, |
| 41 | encoder_hidden_states[batch_id: batch_id_], |
| 42 | controlnet_frames[:, batch_id: batch_id_], |
| 43 | tiled=tiled, tile_size=tile_size, tile_stride=tile_stride |
| 44 | ) |
| 45 | if vram_limit_level >= 1: |
| 46 | res_stack = [res.cpu() for res in res_stack] |
| 47 | res_stacks.append(res_stack) |
| 48 | # concat the residual |
| 49 | additional_res_stack = [] |
| 50 | for i in range(len(res_stacks[0])): |
| 51 | res = torch.concat([res_stack[i] for res_stack in res_stacks], dim=0) |
| 52 | additional_res_stack.append(res) |
| 53 | else: |
| 54 | additional_res_stack = None |
| 55 | |
| 56 | # 2. time |
| 57 | time_emb = unet.time_proj(timestep).to(sample.dtype) |
| 58 | time_emb = unet.time_embedding(time_emb) |
| 59 | |
| 60 | # 3. pre-process |
| 61 | height, width = sample.shape[2], sample.shape[3] |
| 62 | hidden_states = unet.conv_in(sample) |
| 63 | text_emb = encoder_hidden_states |
| 64 | res_stack = [hidden_states.cpu() if vram_limit_level>=1 else hidden_states] |
no test coverage detected