support latest diffusers and longer videos

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]
+            out_source = forward_original(q[:n_frames], k[:n_frames], v[:n_frames])
-            k_uncond = k[n_frames:2 * n_frames].reshape(1, n_frames * sequence_length, -1).repeat(n_frames, 1, 1)
+            out_uncond = forward_extended(q[n_frames:2 * n_frames], k[n_frames:2 * n_frames], v[n_frames:2 * n_frames])
-            k_cond = k[2 * n_frames:].reshape(1, n_frames * sequence_length, -1).repeat(n_frames, 1, 1)
+            out_cond = forward_extended(q[2 * n_frames:], k[2 * n_frames:], v[2 * n_frames:])
-
+                            
-            v_source = v[:n_frames]
+            out = torch.cat([out_source, out_uncond, out_cond], dim=0) # (3 * n_frames, seq_len, dim)
-            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)
 
             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