(
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",
animatediff_batch_size=16,
animatediff_stride=8,
)
| 12 | |
| 13 | |
| 14 | def lets_dance_with_long_video( |
| 15 | unet: SDUNet, |
| 16 | motion_modules: SDMotionModel = None, |
| 17 | controlnet: MultiControlNetManager = None, |
| 18 | sample = None, |
| 19 | timestep = None, |
| 20 | encoder_hidden_states = None, |
| 21 | ipadapter_kwargs_list = {}, |
| 22 | controlnet_frames = None, |
| 23 | unet_batch_size = 1, |
| 24 | controlnet_batch_size = 1, |
| 25 | cross_frame_attention = False, |
| 26 | tiled=False, |
| 27 | tile_size=64, |
| 28 | tile_stride=32, |
| 29 | device="cuda", |
| 30 | animatediff_batch_size=16, |
| 31 | animatediff_stride=8, |
| 32 | ): |
| 33 | num_frames = sample.shape[0] |
| 34 | hidden_states_output = [(torch.zeros(sample[0].shape, dtype=sample[0].dtype), 0) for i in range(num_frames)] |
| 35 | |
| 36 | for batch_id in range(0, num_frames, animatediff_stride): |
| 37 | batch_id_ = min(batch_id + animatediff_batch_size, num_frames) |
| 38 | |
| 39 | # process this batch |
| 40 | hidden_states_batch = lets_dance( |
| 41 | unet, motion_modules, controlnet, |
| 42 | sample[batch_id: batch_id_].to(device), |
| 43 | timestep, |
| 44 | encoder_hidden_states, |
| 45 | ipadapter_kwargs_list=ipadapter_kwargs_list, |
| 46 | controlnet_frames=controlnet_frames[:, batch_id: batch_id_].to(device) if controlnet_frames is not None else None, |
| 47 | unet_batch_size=unet_batch_size, controlnet_batch_size=controlnet_batch_size, |
| 48 | cross_frame_attention=cross_frame_attention, |
| 49 | tiled=tiled, tile_size=tile_size, tile_stride=tile_stride, device=device |
| 50 | ).cpu() |
| 51 | |
| 52 | # update hidden_states |
| 53 | for i, hidden_states_updated in zip(range(batch_id, batch_id_), hidden_states_batch): |
| 54 | bias = max(1 - abs(i - (batch_id + batch_id_ - 1) / 2) / ((batch_id_ - batch_id - 1 + 1e-2) / 2), 1e-2) |
| 55 | hidden_states, num = hidden_states_output[i] |
| 56 | hidden_states = hidden_states * (num / (num + bias)) + hidden_states_updated * (bias / (num + bias)) |
| 57 | hidden_states_output[i] = (hidden_states, num + bias) |
| 58 | |
| 59 | if batch_id_ == num_frames: |
| 60 | break |
| 61 | |
| 62 | # output |
| 63 | hidden_states = torch.stack([h for h, _ in hidden_states_output]) |
| 64 | return hidden_states |
| 65 | |
| 66 | |
| 67 |
no test coverage detected