kijai
ab06ac2f64
Turns out I never had the right way of using EchoShot and was always just accidentally using my existing rope splitting... which just worked with the EchoShot weights, this commit adds proper EchoShot and it's used when their prompt format is used, the previous behaviour can be restored by using my splitting method with " | " between the prompts. EchoShot example has been updated to reflect that.
104 lines
3.6 KiB
Python
104 lines
3.6 KiB
Python
import torch
|
|
from comfy.model_management import get_autocast_device, get_torch_device
|
|
|
|
@torch.autocast(device_type=get_autocast_device(get_torch_device()), enabled=False)
|
|
@torch.compiler.disable()
|
|
def rope_apply_z(x, grid_sizes, freqs, inner_t, shift=6):
|
|
n, c = x.size(2), x.size(3) // 2
|
|
|
|
# loop over samples
|
|
output = []
|
|
for i, (f, h, w) in enumerate(grid_sizes.tolist()):
|
|
seq_len = f * h * w
|
|
|
|
# precompute multipliers
|
|
x_i = torch.view_as_complex(
|
|
x[i, :seq_len].to(torch.float64).reshape(seq_len, n, -1, 2)
|
|
)
|
|
start_ind = [sum(inner_t[i][:_]) for _ in range(len(inner_t[i]))]
|
|
end_ind = [sum(inner_t[i][:_+1]) for _ in range(len(inner_t[i]))]
|
|
|
|
freq_select = []
|
|
for shot_ind, (s, e) in enumerate(zip(start_ind, end_ind)):
|
|
freq_select += [shot_ind * shift] * (e - s)
|
|
shot_freqs = freqs[freq_select]
|
|
|
|
freqs_i = shot_freqs.view(f, 1, 1, -1).expand(f, h, w, -1).reshape(seq_len, 1, -1)
|
|
|
|
# apply rotary embedding
|
|
x_i = torch.view_as_real(x_i * freqs_i).flatten(2)
|
|
x_i = torch.cat([x_i, x[i, seq_len:]])
|
|
|
|
# append to collection
|
|
output.append(x_i)
|
|
return torch.stack(output).float()
|
|
|
|
|
|
@torch.autocast(device_type=get_autocast_device(get_torch_device()), enabled=False)
|
|
@torch.compiler.disable()
|
|
def rope_apply_c(x, freqs, inner_c, shift=6):
|
|
|
|
b, s, n, c = x.size(0), x.size(1), x.size(2), x.size(3) // 2
|
|
|
|
# loop over samples
|
|
output = []
|
|
for i in range(b):
|
|
|
|
# precompute multipliers
|
|
x_i = torch.view_as_complex(
|
|
x[i].to(torch.float64).reshape(s, n, -1, 2)
|
|
)
|
|
|
|
freq_select = []
|
|
for shot_ind, c_len in enumerate(inner_c[i]):
|
|
freq_select += [shot_ind * shift] * c_len
|
|
freq_select += [shot_ind+10] * (s-len(freq_select)) # extra suppression for the empty token
|
|
shot_freqs = freqs[freq_select]
|
|
|
|
freqs_i = shot_freqs.view(s, 1, -1)
|
|
|
|
# apply rotary embedding
|
|
x_i = torch.view_as_real(x_i * freqs_i).flatten(2)
|
|
|
|
# append to collection
|
|
output.append(x_i)
|
|
return torch.stack(output).float()
|
|
|
|
@torch.autocast(device_type=get_autocast_device(get_torch_device()), enabled=False)
|
|
@torch.compiler.disable()
|
|
def rope_apply_echoshot(x, grid_sizes, freqs, inner_t, shift=4):
|
|
n, c = x.size(2), x.size(3) // 2
|
|
|
|
# split freqs
|
|
freqs = freqs.split([c - 2 * (c // 3), c // 3, c // 3], dim=1)
|
|
|
|
# loop over samples
|
|
output = []
|
|
for i, (f, h, w) in enumerate(grid_sizes.tolist()):
|
|
seq_len = f * h * w
|
|
|
|
# precompute multipliers
|
|
x_i = torch.view_as_complex(
|
|
x[i, :seq_len].to(torch.float64).reshape(seq_len, n, -1, 2)
|
|
)
|
|
start_ind = [sum(inner_t[i][:_]) for _ in range(len(inner_t[i]))]
|
|
end_ind = [sum(inner_t[i][:_+1]) for _ in range(len(inner_t[i]))]
|
|
freq_select = []
|
|
for shot_ind, (s, e) in enumerate(zip(start_ind, end_ind)):
|
|
freq_select += list(range(shot_ind * shift + s, shot_ind * shift + e))
|
|
t_freqs = freqs[0][freq_select]
|
|
|
|
freqs_i = torch.cat([
|
|
# freqs[0][:f].view(f, 1, 1, -1).expand(f, h, w, -1),
|
|
t_freqs.view(f, 1, 1, -1).expand(f, h, w, -1), ###
|
|
freqs[1][:h].view(1, h, 1, -1).expand(f, h, w, -1),
|
|
freqs[2][:w].view(1, 1, w, -1).expand(f, h, w, -1)
|
|
], dim=-1).reshape(seq_len, 1, -1)
|
|
|
|
# apply rotary embedding
|
|
x_i = torch.view_as_real(x_i * freqs_i).flatten(2)
|
|
x_i = torch.cat([x_i, x[i, seq_len:]])
|
|
|
|
# append to collection
|
|
output.append(x_i)
|
|
return torch.stack(output).float() |