| 117 | |
| 118 | |
| 119 | def lets_dance_xl( |
| 120 | unet: SDXLUNet, |
| 121 | motion_modules: SDXLMotionModel = None, |
| 122 | controlnet: MultiControlNetManager = None, |
| 123 | sample = None, |
| 124 | add_time_id = None, |
| 125 | add_text_embeds = None, |
| 126 | timestep = None, |
| 127 | encoder_hidden_states = None, |
| 128 | ipadapter_kwargs_list = {}, |
| 129 | controlnet_frames = None, |
| 130 | unet_batch_size = 1, |
| 131 | controlnet_batch_size = 1, |
| 132 | cross_frame_attention = False, |
| 133 | tiled=False, |
| 134 | tile_size=64, |
| 135 | tile_stride=32, |
| 136 | device = "cuda", |
| 137 | vram_limit_level = 0, |
| 138 | ): |
| 139 | # 0. Text embedding alignment (only for video processing) |
| 140 | if encoder_hidden_states.shape[0] != sample.shape[0]: |
| 141 | encoder_hidden_states = encoder_hidden_states.repeat(sample.shape[0], 1, 1, 1) |
| 142 | if add_text_embeds.shape[0] != sample.shape[0]: |
| 143 | add_text_embeds = add_text_embeds.repeat(sample.shape[0], 1) |
| 144 | |
| 145 | # 1. ControlNet |
| 146 | controlnet_insert_block_id = 22 |
| 147 | if controlnet is not None and controlnet_frames is not None: |
| 148 | res_stacks = [] |
| 149 | # process controlnet frames with batch |
| 150 | for batch_id in range(0, sample.shape[0], controlnet_batch_size): |
| 151 | batch_id_ = min(batch_id + controlnet_batch_size, sample.shape[0]) |
| 152 | res_stack = controlnet( |
| 153 | sample[batch_id: batch_id_], |
| 154 | timestep, |
| 155 | encoder_hidden_states[batch_id: batch_id_], |
| 156 | controlnet_frames[:, batch_id: batch_id_], |
| 157 | add_time_id=add_time_id, |
| 158 | add_text_embeds=add_text_embeds, |
| 159 | tiled=tiled, tile_size=tile_size, tile_stride=tile_stride, |
| 160 | unet=unet, # for Kolors, some modules in ControlNets will be replaced. |
| 161 | ) |
| 162 | if vram_limit_level >= 1: |
| 163 | res_stack = [res.cpu() for res in res_stack] |
| 164 | res_stacks.append(res_stack) |
| 165 | # concat the residual |
| 166 | additional_res_stack = [] |
| 167 | for i in range(len(res_stacks[0])): |
| 168 | res = torch.concat([res_stack[i] for res_stack in res_stacks], dim=0) |
| 169 | additional_res_stack.append(res) |
| 170 | else: |
| 171 | additional_res_stack = None |
| 172 | |
| 173 | # 2. time |
| 174 | t_emb = unet.time_proj(timestep).to(sample.dtype) |
| 175 | t_emb = unet.time_embedding(t_emb) |
| 176 | |