From a05b590b2b1abeaddbc62bf44eac54c8dda5c0d0 Mon Sep 17 00:00:00 2001 From: Justin John Date: Sun, 5 Nov 2023 14:15:00 +0530 Subject: [PATCH] support latest diffusers and longer videos --- animation/TokenFlow-WebUI/tokenflow_utils.py | 138 ++++++++++--------- 1 file changed, 70 insertions(+), 68 deletions(-) diff --git a/animation/TokenFlow-WebUI/tokenflow_utils.py b/animation/TokenFlow-WebUI/tokenflow_utils.py index 6b76ec6..f828b9a 100644 --- a/animation/TokenFlow-WebUI/tokenflow_utils.py +++ b/animation/TokenFlow-WebUI/tokenflow_utils.py @@ -48,7 +48,7 @@ def load_source_latents_t(t, latents_path): def register_conv_injection(model, injection_schedule): def conv_forward(self): - def forward(input_tensor, temb): + def forward(input_tensor, temb,scale=None): hidden_states = input_tensor hidden_states = self.norm1(hidden_states) @@ -103,6 +103,7 @@ def register_conv_injection(model, injection_schedule): conv_module.forward = conv_forward(conv_module) setattr(conv_module, 'injection_schedule', injection_schedule) + def register_extended_attention_pnp(model, injection_schedule): def sa_forward(self): to_out = self.to_out @@ -110,11 +111,72 @@ def register_extended_attention_pnp(model, injection_schedule): to_out = self.to_out[0] else: to_out = self.to_out + + def forward_original(q, k, v): + n_frames, seq_len, dim = q.shape + h = self.heads + head_dim = dim // h + + q = self.head_to_batch_dim(q).reshape(n_frames, h, seq_len, head_dim) + k = self.head_to_batch_dim(k).reshape(n_frames, h, seq_len, head_dim) + v = self.head_to_batch_dim(v).reshape(n_frames, h, seq_len, head_dim) - def forward(x, encoder_hidden_states=None): + out_all = [] + + for frame in range(n_frames): + out = [] + for j in range(h): + sim = torch.matmul(q[frame, j], k[frame, j].transpose(-1, -2)) * self.scale # (seq_len, seq_len) + out.append(torch.matmul(sim.softmax(dim=-1), v[frame, j])) # h * (seq_len, head_dim) + + out = torch.cat(out, dim=0).reshape(-1, seq_len, head_dim) # (h, seq_len, head_dim) + out_all.append(out) # n_frames * (h, seq_len, head_dim) + + out = torch.cat(out_all, dim=0) # (n_frames * h, seq_len, head_dim) + out = self.batch_to_head_dim(out) # (n_frames, seq_len, h * head_dim) + return out + + def forward_extended(q, k, v): + n_frames, seq_len, dim = q.shape + h = self.heads + head_dim = dim // h + + q = self.head_to_batch_dim(q).reshape(n_frames, h, seq_len, head_dim) + k = self.head_to_batch_dim(k).reshape(n_frames, h, seq_len, head_dim) + v = self.head_to_batch_dim(v).reshape(n_frames, h, seq_len, head_dim) + + out_all = [] + window_size = 3 + + for frame in range(n_frames): + out = [] + # sliding window to improve speed. + window = range(max(0, frame-window_size // 2), min(n_frames, frame+window_size//2+1)) + + for j in range(h): + sim_all = [] + + for kframe in window: + sim_all.append(torch.matmul(q[frame, j], k[kframe, j].transpose(-1, -2)) * self.scale) # window * (seq_len, seq_len) + + sim_all = torch.cat(sim_all).reshape(len(window), seq_len, seq_len).transpose(0, 1) # (seq_len, window, seq_len) + sim_all = sim_all.reshape(seq_len, len(window) * seq_len) # (seq_len, window * seq_len) + out.append(torch.matmul(sim_all.softmax(dim=-1), v[window, j].reshape(len(window) * seq_len, head_dim))) # h * (seq_len, head_dim) + + out = torch.cat(out, dim=0).reshape(-1, seq_len, head_dim) # (h, seq_len, head_dim) + out_all.append(out) # n_frames * (h, seq_len, head_dim) + + out = torch.cat(out_all, dim=0) # (n_frames * h, seq_len, head_dim) + out = self.batch_to_head_dim(out) # (n_frames, seq_len, h * head_dim) + + return out + + + def forward(x, encoder_hidden_states=None, attention_mask=None, scale=None): batch_size, sequence_length, dim = x.shape h = self.heads n_frames = batch_size // 3 + is_cross = encoder_hidden_states is not None encoder_hidden_states = encoder_hidden_states if is_cross else x q = self.to_q(x) @@ -129,72 +191,11 @@ def register_extended_attention_pnp(model, injection_schedule): q[2 * n_frames:] = q[:n_frames] k[2 * n_frames:] = k[:n_frames] - k_source = k[:n_frames] - k_uncond = k[n_frames:2 * n_frames].reshape(1, n_frames * sequence_length, -1).repeat(n_frames, 1, 1) - k_cond = k[2 * n_frames:].reshape(1, n_frames * sequence_length, -1).repeat(n_frames, 1, 1) - - v_source = v[:n_frames] - v_uncond = v[n_frames:2 * n_frames].reshape(1, n_frames * sequence_length, -1).repeat(n_frames, 1, 1) - v_cond = v[2 * n_frames:].reshape(1, n_frames * sequence_length, -1).repeat(n_frames, 1, 1) - - q_source = self.head_to_batch_dim(q[:n_frames]) - q_uncond = self.head_to_batch_dim(q[n_frames:2 * n_frames]) - q_cond = self.head_to_batch_dim(q[2 * n_frames:]) - k_source = self.head_to_batch_dim(k_source) - k_uncond = self.head_to_batch_dim(k_uncond) - k_cond = self.head_to_batch_dim(k_cond) - v_source = self.head_to_batch_dim(v_source) - v_uncond = self.head_to_batch_dim(v_uncond) - v_cond = self.head_to_batch_dim(v_cond) - - - q_src = q_source.view(n_frames, h, sequence_length, dim // h) - k_src = k_source.view(n_frames, h, sequence_length, dim // h) - v_src = v_source.view(n_frames, h, sequence_length, dim // h) - q_uncond = q_uncond.view(n_frames, h, sequence_length, dim // h) - k_uncond = k_uncond.view(n_frames, h, sequence_length * n_frames, dim // h) - v_uncond = v_uncond.view(n_frames, h, sequence_length * n_frames, dim // h) - q_cond = q_cond.view(n_frames, h, sequence_length, dim // h) - k_cond = k_cond.view(n_frames, h, sequence_length * n_frames, dim // h) - v_cond = v_cond.view(n_frames, h, sequence_length * n_frames, dim // h) - - out_source_all = [] - out_uncond_all = [] - out_cond_all = [] - - single_batch = n_frames<=12 - b = n_frames if single_batch else 1 - - for frame in range(0, n_frames, b): - out_source = [] - out_uncond = [] - out_cond = [] - for j in range(h): - sim_source_b = torch.bmm(q_src[frame: frame+ b, j], k_src[frame: frame+ b, j].transpose(-1, -2)) * self.scale - sim_uncond_b = torch.bmm(q_uncond[frame: frame+ b, j], k_uncond[frame: frame+ b, j].transpose(-1, -2)) * self.scale - sim_cond = torch.bmm(q_cond[frame: frame+ b, j], k_cond[frame: frame+ b, j].transpose(-1, -2)) * self.scale - - out_source.append(torch.bmm(sim_source_b.softmax(dim=-1), v_src[frame: frame+ b, j])) - out_uncond.append(torch.bmm(sim_uncond_b.softmax(dim=-1), v_uncond[frame: frame+ b, j])) - out_cond.append(torch.bmm(sim_cond.softmax(dim=-1), v_cond[frame: frame+ b, j])) - - out_source = torch.cat(out_source, dim=0) - out_uncond = torch.cat(out_uncond, dim=0) - out_cond = torch.cat(out_cond, dim=0) - if single_batch: - out_source = out_source.view(h, n_frames,sequence_length, dim // h).permute(1, 0, 2, 3).reshape(h * n_frames, sequence_length, -1) - out_uncond = out_uncond.view(h, n_frames,sequence_length, dim // h).permute(1, 0, 2, 3).reshape(h * n_frames, sequence_length, -1) - out_cond = out_cond.view(h, n_frames,sequence_length, dim // h).permute(1, 0, 2, 3).reshape(h * n_frames, sequence_length, -1) - out_source_all.append(out_source) - out_uncond_all.append(out_uncond) - out_cond_all.append(out_cond) - - out_source = torch.cat(out_source_all, dim=0) - out_uncond = torch.cat(out_uncond_all, dim=0) - out_cond = torch.cat(out_cond_all, dim=0) - - out = torch.cat([out_source, out_uncond, out_cond], dim=0) - out = self.batch_to_head_dim(out) + out_source = forward_original(q[:n_frames], k[:n_frames], v[:n_frames]) + out_uncond = forward_extended(q[n_frames:2 * n_frames], k[n_frames:2 * n_frames], v[n_frames:2 * n_frames]) + out_cond = forward_extended(q[2 * n_frames:], k[2 * n_frames:], v[2 * n_frames:]) + + out = torch.cat([out_source, out_uncond, out_cond], dim=0) # (3 * n_frames, seq_len, dim) return to_out(out) @@ -306,6 +307,7 @@ def make_tokenflow_attention_block(block_class: Type[torch.nn.Module]) -> Type[t timestep=None, cross_attention_kwargs=None, class_labels=None, + scale=None, ) -> torch.Tensor: batch_size, sequence_length, dim = hidden_states.shape