|
|
import torch |
|
|
from tqdm import tqdm |
|
|
from transformers.cache_utils import Cache, DynamicCache |
|
|
|
|
|
class OmniGenScheduler: |
|
|
def __init__(self, num_steps: int=50, time_shifting_factor: int=1): |
|
|
self.num_steps = num_steps |
|
|
self.time_shift = time_shifting_factor |
|
|
|
|
|
t = torch.linspace(0, 1, num_steps+1) |
|
|
t = t / (t + time_shifting_factor - time_shifting_factor * t) |
|
|
self.sigma = t |
|
|
|
|
|
def crop_kv_cache(self, past_key_values, num_tokens_for_img): |
|
|
crop_past_key_values = () |
|
|
for layer_idx in range(len(past_key_values)): |
|
|
key_states, value_states = past_key_values[layer_idx][:2] |
|
|
crop_past_key_values += ((key_states[..., :-(num_tokens_for_img+1), :], value_states[..., :-(num_tokens_for_img+1), :], ),) |
|
|
return crop_past_key_values |
|
|
|
|
|
|
|
|
def crop_position_ids_for_cache(self, position_ids, num_tokens_for_img): |
|
|
if isinstance(position_ids, list): |
|
|
for i in range(len(position_ids)): |
|
|
position_ids[i] = position_ids[i][:, -(num_tokens_for_img+1):] |
|
|
else: |
|
|
position_ids = position_ids[:, -(num_tokens_for_img+1):] |
|
|
return position_ids |
|
|
|
|
|
def crop_attention_mask_for_cache(self, attention_mask, num_tokens_for_img): |
|
|
if isinstance(attention_mask, list): |
|
|
return [x[..., -(num_tokens_for_img+1):, :] for x in attention_mask] |
|
|
return attention_mask[..., -(num_tokens_for_img+1):, :] |
|
|
|
|
|
def __call__(self, z, func, model_kwargs, use_kv_cache: bool=True): |
|
|
past_key_values = None |
|
|
for i in tqdm(range(self.num_steps)): |
|
|
timesteps = torch.zeros(size=(len(z), )).to(z.device) + self.sigma[i] |
|
|
pred, temp_past_key_values = func(z, timesteps, past_key_values=past_key_values, **model_kwargs) |
|
|
sigma_next = self.sigma[i+1] |
|
|
sigma = self.sigma[i] |
|
|
z = z + (sigma_next - sigma) * pred |
|
|
if i == 0 and use_kv_cache: |
|
|
num_tokens_for_img = z.size(-1)*z.size(-2) // 4 |
|
|
if isinstance(temp_past_key_values, list): |
|
|
past_key_values = [self.crop_kv_cache(x, num_tokens_for_img) for x in temp_past_key_values] |
|
|
model_kwargs['input_ids'] = [None] * len(temp_past_key_values) |
|
|
else: |
|
|
past_key_values = self.crop_kv_cache(temp_past_key_values, num_tokens_for_img) |
|
|
model_kwargs['input_ids'] = None |
|
|
|
|
|
model_kwargs['position_ids'] = self.crop_position_ids_for_cache(model_kwargs['position_ids'], num_tokens_for_img) |
|
|
model_kwargs['attention_mask'] = self.crop_attention_mask_for_cache(model_kwargs['attention_mask'], num_tokens_for_img) |
|
|
return z |
|
|
|
|
|
|
