Spaces:
Runtime error
Runtime error
File size: 3,785 Bytes
0608f20 | 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 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 | import inspect
import numpy as np
import torch
def cfg_skip():
def decorator(func):
def wrapper(self, *args, **kwargs):
if torch.is_grad_enabled():
return func(self, *args, **kwargs)
if 'hidden_states' in kwargs and kwargs['hidden_states'] is not None:
main_input = kwargs['hidden_states']
elif 'x' in kwargs and kwargs['x'] is not None:
main_input = kwargs['x']
elif len(args) > 0:
main_input = args[0]
else:
raise ValueError("No input tensor found in args or kwargs")
bs = len(main_input)
if bs >= 2 and self.cfg_skip_ratio is not None and self.current_steps >= self.num_inference_steps * (1 - self.cfg_skip_ratio):
bs_half = int(bs // 2)
new_x = main_input[bs_half:]
new_args = [
arg[bs_half:] if
isinstance(arg,
(torch.Tensor, list, tuple, np.ndarray)) and
len(arg) == bs else arg for arg in args
]
new_kwargs = {
k: (v[bs_half:] if
isinstance(v,
(torch.Tensor, list, tuple,
np.ndarray)) and len(v) == bs else v
) for k, v in kwargs.items()
}
else:
new_x = main_input
new_args = args
new_kwargs = kwargs
sig = inspect.signature(func)
new_bs = len(new_x)
new_bs_half = int(new_bs // 2)
if new_bs >= 2:
# cond
args_i = [
arg[new_bs_half:] if
isinstance(arg,
(torch.Tensor, list, tuple, np.ndarray)) and
len(arg) == new_bs else arg for arg in new_args
]
kwargs_i = {
k: (v[new_bs_half:] if
isinstance(v,
(torch.Tensor, list, tuple,
np.ndarray)) and len(v) == new_bs else v
) for k, v in new_kwargs.items()
}
if 'cond_flag' in sig.parameters:
kwargs_i["cond_flag"] = True
cond_out = func(self, *args_i, **kwargs_i)
# uncond
uncond_args_i = [
arg[:new_bs_half] if
isinstance(arg,
(torch.Tensor, list, tuple, np.ndarray)) and
len(arg) == new_bs else arg for arg in new_args
]
uncond_kwargs_i = {
k: (v[:new_bs_half] if
isinstance(v,
(torch.Tensor, list, tuple,
np.ndarray)) and len(v) == new_bs else v
) for k, v in new_kwargs.items()
}
if 'cond_flag' in sig.parameters:
uncond_kwargs_i["cond_flag"] = False
uncond_out = func(self, *uncond_args_i,
**uncond_kwargs_i)
result = torch.cat([uncond_out, cond_out], dim=0)
else:
result = func(self, *new_args, **new_kwargs)
if bs >= 2 and self.cfg_skip_ratio is not None and self.current_steps >= self.num_inference_steps * (1 - self.cfg_skip_ratio):
result = torch.cat([result, result], dim=0)
return result
return wrapper
return decorator |