Spaces:
Running on Zero
Running on Zero
File size: 3,197 Bytes
471169b | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 | import torch
import numpy as np
coefficients = {
"DEFAULT": [-1.12343328e+02, 1.50680483e+02, -5.15023303e+01, 6.24892431e+00, 6.85022158e-02],
}
class TeaCache:
def __init__(self, num_inference_steps, model_name, treshold=0.3, start_step_treshold=0.1, end_step_treshold=0.9):
self.input_bank = []
self.current_step = 0
self.accumulated_distance = 0.0
self.num_inference_steps = num_inference_steps * 2
self.start_step_teacache = int(num_inference_steps * start_step_treshold) * 2
self.end_step_teacache = int(num_inference_steps * end_step_treshold) * 2
self.treshold = treshold # [0.3, 0.5, 0.7, 0.9]
self.coefficients = coefficients[model_name]
self.step_name = "even"
self.init_memory()
def init_memory(self):
self.accumulated_distance = {
"even": 0.0,
"odd": 0.0,
}
self.flow_direction = {
"even": None,
"odd": None,
}
self.previous_modulated_input = {
"even": None,
"odd": None,
}
# print("TEACACHE MEMORY HAS BEEN CREATED")
def check_for_using_cached_value(self, modulated_input):
use_tea_cache = (self.treshold > 0.0) and (self.start_step_teacache <= self.current_step < self.end_step_teacache)
self.step_name = "even" if self.current_step % 2 == 0 else "odd"
use_cached_value = False
if use_tea_cache:
rescale_func = np.poly1d(self.coefficients)
current_disntace = rescale_func(
self.calculate_distance(modulated_input, self.previous_modulated_input[self.step_name])
)
self.accumulated_distance[self.step_name] += current_disntace
if self.accumulated_distance[self.step_name] < self.treshold:
use_cached_value = True
else:
use_cached_value = False
self.accumulated_distance[self.step_name] = 0.0
if self.step_name == "even":
self.input_bank.append(modulated_input.cpu())
self.previous_modulated_input[self.step_name] = modulated_input.clone()
# if use_tea_cache:
# print(f"[ STEP:{self.current_step} | USE CACHED VALUE: {use_cached_value} | ACCUMULATED DISTANCE: {self.accumulated_distance} | CURRENT DISTANCE: {current_disntace} ]")
return use_cached_value
def use_cache(self, hidden_states):
return hidden_states + self.flow_direction[self.step_name].to(device=hidden_states.device)
def calculate_distance(self, previous_tensor, current_tensor):
relative_l1_distance = torch.abs(
previous_tensor - current_tensor
).mean() / torch.abs(previous_tensor).mean()
return relative_l1_distance.to(torch.float32).cpu().item()
def update(self, flow_direction):
self.flow_direction[self.step_name] = flow_direction
self.current_step += 1
if self.current_step == self.num_inference_steps:
self.current_step = 0
self.init_memory() |