storytold/storyteller-ml
2
0
Fork 0
mirror of https://github.com/storytold/storyteller-ml.git synced 2026-10-15 00:03:06 +00:00
Code Issues Packages Projects Releases Wiki Activity

support latest diffusers and longer videos

Browse Source
This commit is contained in:
Justin John
2023-11-05 14:15:00 +05:30
parent e3858289be
commit a05b590b2b
1 changed files with 70 additions and 68 deletions
Download Patch File Download Diff File Expand all files Collapse all files
animation/TokenFlow-WebUI/tokenflow_utils.py
+70 -68
View File
@@ -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