39 lines
1.4 KiB
Python
Executable File
39 lines
1.4 KiB
Python
Executable File
import numpy as np
|
|
import torch
|
|
|
|
|
|
def cfg_skip():
|
|
def decorator(func):
|
|
def wrapper(self, x, *args, **kwargs):
|
|
bs = len(x)
|
|
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 = x[bs_half:]
|
|
|
|
new_args = []
|
|
for arg in args:
|
|
if isinstance(arg, (torch.Tensor, list, tuple, np.ndarray)):
|
|
new_args.append(arg[bs_half:])
|
|
else:
|
|
new_args.append(arg)
|
|
|
|
new_kwargs = {}
|
|
for key, content in kwargs.items():
|
|
if isinstance(content, (torch.Tensor, list, tuple, np.ndarray)):
|
|
new_kwargs[key] = content[bs_half:]
|
|
else:
|
|
new_kwargs[key] = content
|
|
else:
|
|
new_x = x
|
|
new_args = args
|
|
new_kwargs = kwargs
|
|
|
|
result = func(self, new_x, *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 |