131 lines
5.0 KiB
Python
131 lines
5.0 KiB
Python
import logging
|
|
import math
|
|
import torch
|
|
|
|
log = logging.getLogger(__name__)
|
|
|
|
|
|
def build_temporal_cost(q_token_idx, Lq, Lk, device, dtype, tokens_per_frame):
|
|
offset = torch.zeros(Lq, Lk, device=device, dtype=dtype)
|
|
query_frames = torch.arange(Lq, device=device, dtype=torch.long) // tokens_per_frame
|
|
for seg in q_token_idx:
|
|
local = seg['local_token_idx'].to(device=device)
|
|
d = (query_frames.float()[:, None] - seg['midpoint']).abs()
|
|
cost = (torch.relu(d - seg['window']) ** 2) / (2 * seg['sigma'] ** 2)
|
|
offset[:, local] = cost.to(offset.dtype)
|
|
return offset
|
|
|
|
|
|
def build_temporal_cost_scaled(q_token_idx, Lq, Lk, device, dtype, latent_frames):
|
|
offset = torch.zeros(Lq, Lk, device=device, dtype=dtype)
|
|
query_frames = torch.arange(Lq, device=device, dtype=torch.float32) * latent_frames / Lq
|
|
for seg in q_token_idx:
|
|
local = seg['local_token_idx'].to(device=device)
|
|
d = (query_frames[:, None] - seg['midpoint']).abs()
|
|
cost = (torch.relu(d - seg['window']) ** 2) / (2 * seg['sigma'] ** 2)
|
|
offset[:, local] = cost.to(offset.dtype)
|
|
return offset
|
|
|
|
|
|
def create_mask_fn(q_token_idx, fallback_tokens_per_frame, latent_frames):
|
|
cache = {}
|
|
max_token_idx = max(int(seg['local_token_idx'].max().item()) for seg in q_token_idx) + 1
|
|
|
|
def mask_fn(q, k, transformer_options):
|
|
Lq, Lk = q.shape[1], k.shape[1]
|
|
if Lq == Lk:
|
|
return None
|
|
cond_or_uncond = transformer_options.get('cond_or_uncond', [])
|
|
if 1 in cond_or_uncond and 0 not in cond_or_uncond:
|
|
return None
|
|
grid_sizes = transformer_options.get('grid_sizes', None)
|
|
video_tpf = int(grid_sizes[1]) * int(grid_sizes[2]) if grid_sizes is not None else fallback_tokens_per_frame
|
|
video_lq = latent_frames * video_tpf
|
|
if Lk == video_lq or Lk < max_token_idx:
|
|
return None
|
|
mode = 'video' if Lq == video_lq else 'scaled'
|
|
key = (Lq, Lk, mode, q.device)
|
|
if key not in cache:
|
|
if mode == 'video':
|
|
cost = build_temporal_cost(q_token_idx, Lq, Lk, q.device, q.dtype, video_tpf)
|
|
else:
|
|
cost = build_temporal_cost_scaled(q_token_idx, Lq, Lk, q.device, q.dtype, latent_frames)
|
|
log.info('[PromptRelay] Built penalty matrix (%s): Lq=%d, Lk=%d, nonzero=%d/%d',
|
|
mode, Lq, Lk, (cost > 0).sum().item(), cost.numel())
|
|
cache[key] = -cost
|
|
return cache[key].to(q.dtype)
|
|
|
|
return mask_fn
|
|
|
|
|
|
def build_segments(token_ranges, segment_lengths, epsilon=1e-3):
|
|
sigma = 1.0 / math.log(1.0 / epsilon) if 0 < epsilon < 1 else 0.1448
|
|
q_token_idx = []
|
|
frame_cursor = 0
|
|
for (tok_start, tok_end), L in zip(token_ranges, segment_lengths):
|
|
if L <= 0:
|
|
frame_cursor += L
|
|
continue
|
|
midpoint = (2 * frame_cursor + L) // 2
|
|
window = max(L // 2 - 2, 0)
|
|
q_token_idx.append({
|
|
'local_token_idx': torch.arange(tok_start, tok_end),
|
|
'midpoint': midpoint,
|
|
'window': window,
|
|
'sigma': sigma,
|
|
})
|
|
frame_cursor += L
|
|
return q_token_idx
|
|
|
|
|
|
def get_raw_tokenizer(clip):
|
|
tokenizer_wrapper = clip.tokenizer
|
|
for attr_name in dir(tokenizer_wrapper):
|
|
if attr_name.startswith('_'):
|
|
continue
|
|
inner = getattr(tokenizer_wrapper, attr_name, None)
|
|
if inner is not None and hasattr(inner, 'tokenizer'):
|
|
return inner.tokenizer
|
|
raise RuntimeError(
|
|
f'Could not find raw tokenizer on CLIP object. '
|
|
f'Known attributes: {[a for a in dir(tokenizer_wrapper) if not a.startswith(chr(95))]}'
|
|
)
|
|
|
|
|
|
def map_token_indices(raw_tokenizer, global_prompt, local_prompts):
|
|
prefixed_locals = [' ' + lp for lp in local_prompts]
|
|
full_prompt = global_prompt + ''.join(prefixed_locals)
|
|
has_eos = getattr(raw_tokenizer, 'add_eos', False)
|
|
eos_adj = 1 if has_eos else 0
|
|
prev_len = len(raw_tokenizer(global_prompt)['input_ids']) - eos_adj
|
|
token_ranges = []
|
|
built = global_prompt
|
|
for plp in prefixed_locals:
|
|
built += plp
|
|
cur_len = len(raw_tokenizer(built)['input_ids']) - eos_adj
|
|
if cur_len <= prev_len:
|
|
raise ValueError(f'Local prompt produced no tokens: {plp.strip()!r}')
|
|
token_ranges.append((prev_len, cur_len))
|
|
prev_len = cur_len
|
|
return full_prompt, token_ranges
|
|
|
|
|
|
def distribute_segment_lengths(num_segments, latent_frames, specified_lengths=None):
|
|
if specified_lengths:
|
|
if len(specified_lengths) != num_segments:
|
|
raise ValueError(
|
|
f'Number of segment_lengths ({len(specified_lengths)}) '
|
|
f'must match number of local prompts ({num_segments})'
|
|
)
|
|
lengths = specified_lengths
|
|
else:
|
|
step = -(-latent_frames // num_segments)
|
|
lengths = [step] * num_segments
|
|
effective = []
|
|
cursor = 0
|
|
for L in lengths:
|
|
end = min(cursor + L, latent_frames)
|
|
effective.append(max(end - cursor, 0))
|
|
cursor = end
|
|
return effective
|