Spaces:
Runtime error
Runtime error
| from tqdm import tqdm | |
| from typing import Optional, Dict, Any, Tuple, List | |
| import gc | |
| import torch | |
| from transformers.cache_utils import Cache, DynamicCache, OffloadedCache | |
| class OmniGenCache(DynamicCache): | |
| def __init__(self, | |
| num_tokens_for_img: int, offload_kv_cache: bool=False) -> None: | |
| if not torch.cuda.is_available(): | |
| raise RuntimeError("OffloadedCache can only be used with a GPU") | |
| super().__init__() | |
| self.original_device = [] | |
| self.prefetch_stream = torch.cuda.Stream() | |
| self.num_tokens_for_img = num_tokens_for_img | |
| self.offload_kv_cache = offload_kv_cache | |
| def prefetch_layer(self, layer_idx: int): | |
| "Starts prefetching the next layer cache" | |
| if layer_idx < len(self): | |
| with torch.cuda.stream(self.prefetch_stream): | |
| # Prefetch next layer tensors to GPU | |
| device = self.original_device[layer_idx] | |
| self.key_cache[layer_idx] = self.key_cache[layer_idx].to(device, non_blocking=True) | |
| self.value_cache[layer_idx] = self.value_cache[layer_idx].to(device, non_blocking=True) | |
| def evict_previous_layer(self, layer_idx: int): | |
| "Moves the previous layer cache to the CPU" | |
| if len(self) > 2: | |
| # We do it on the default stream so it occurs after all earlier computations on these tensors are done | |
| if layer_idx == 0: | |
| prev_layer_idx = -1 | |
| else: | |
| prev_layer_idx = (layer_idx - 1) % len(self) | |
| self.key_cache[prev_layer_idx] = self.key_cache[prev_layer_idx].to("cpu", non_blocking=True) | |
| self.value_cache[prev_layer_idx] = self.value_cache[prev_layer_idx].to("cpu", non_blocking=True) | |
| def __getitem__(self, layer_idx: int) -> List[Tuple[torch.Tensor]]: | |
| "Gets the cache for this layer to the device. Prefetches the next and evicts the previous layer." | |
| if layer_idx < len(self): | |
| if self.offload_kv_cache: | |
| # Evict the previous layer if necessary | |
| torch.cuda.current_stream().synchronize() | |
| self.evict_previous_layer(layer_idx) | |
| # Load current layer cache to its original device if not already there | |
| original_device = self.original_device[layer_idx] | |
| # self.prefetch_stream.synchronize(original_device) | |
| torch.cuda.synchronize(self.prefetch_stream) | |
| key_tensor = self.key_cache[layer_idx] | |
| value_tensor = self.value_cache[layer_idx] | |
| # Prefetch the next layer | |
| self.prefetch_layer((layer_idx + 1) % len(self)) | |
| else: | |
| key_tensor = self.key_cache[layer_idx] | |
| value_tensor = self.value_cache[layer_idx] | |
| return (key_tensor, value_tensor) | |
| else: | |
| raise KeyError(f"Cache only has {len(self)} layers, attempted to access layer with index {layer_idx}") | |
| def update( | |
| self, | |
| key_states: torch.Tensor, | |
| value_states: torch.Tensor, | |
| layer_idx: int, | |
| cache_kwargs: Optional[Dict[str, Any]] = None, | |
| ) -> Tuple[torch.Tensor, torch.Tensor]: | |
| """ | |
| Updates the cache with the new `key_states` and `value_states` for the layer `layer_idx`. | |
| Parameters: | |
| key_states (`torch.Tensor`): | |
| The new key states to cache. | |
| value_states (`torch.Tensor`): | |
| The new value states to cache. | |
| layer_idx (`int`): | |
| The index of the layer to cache the states for. | |
| cache_kwargs (`Dict[str, Any]`, `optional`): | |
| Additional arguments for the cache subclass. No additional arguments are used in `OffloadedCache`. | |
| Return: | |
| A tuple containing the updated key and value states. | |
| """ | |
| # Update the cache | |
| if len(self.key_cache) < layer_idx: | |
| raise ValueError("OffloadedCache does not support model usage where layers are skipped. Use DynamicCache.") | |
| elif len(self.key_cache) == layer_idx: | |
| # only cache the states for condition tokens | |
| key_states = key_states[..., :-(self.num_tokens_for_img+1), :] | |
| value_states = value_states[..., :-(self.num_tokens_for_img+1), :] | |
| # Update the number of seen tokens | |
| if layer_idx == 0: | |
| self._seen_tokens += key_states.shape[-2] | |
| self.key_cache.append(key_states) | |
| self.value_cache.append(value_states) | |
| self.original_device.append(key_states.device) | |
| if self.offload_kv_cache: | |
| self.evict_previous_layer(layer_idx) | |
| return self.key_cache[layer_idx], self.value_cache[layer_idx] | |
| else: | |
| # only cache the states for condition tokens | |
| key_tensor, value_tensor = self[layer_idx] | |
| k = torch.cat([key_tensor, key_states], dim=-2) | |
| v = torch.cat([value_tensor, value_states], dim=-2) | |
| return k, v | |
| 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): | |
| # return | |
| 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 | |
| return DynamicCache.from_legacy_cache(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 crop_cache(self, cache, num_tokens_for_img): | |
| for i in range(len(cache.key_cache)): | |
| cache.key_cache[i] = cache.key_cache[i][..., :-(num_tokens_for_img+1), :] | |
| cache.value_cache[i] = cache.value_cache[i][..., :-(num_tokens_for_img+1), :] | |
| return cache | |
| def __call__(self, z, func, model_kwargs, use_kv_cache: bool=True, offload_kv_cache: bool=True): | |
| num_tokens_for_img = z.size(-1)*z.size(-2) // 4 | |
| if isinstance(model_kwargs['input_ids'], list): | |
| cache = [OmniGenCache(num_tokens_for_img, offload_kv_cache) for _ in range(len(model_kwargs['input_ids']))] if use_kv_cache else None | |
| else: | |
| cache = OmniGenCache(num_tokens_for_img, offload_kv_cache) if use_kv_cache else None | |
| results = {} | |
| for i in tqdm(range(self.num_steps)): | |
| timesteps = torch.zeros(size=(len(z), )).to(z.device) + self.sigma[i] | |
| pred, cache = func(z, timesteps, past_key_values=cache, **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(cache, list): | |
| model_kwargs['input_ids'] = [None] * len(cache) | |
| else: | |
| 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) | |
| del cache | |
| torch.cuda.empty_cache() | |
| gc.collect() | |
| return z | |