()
| 3 | |
| 4 | |
| 5 | def cfg_skip(): |
| 6 | def decorator(func): |
| 7 | def wrapper(self, x, *args, **kwargs): |
| 8 | bs = len(x) |
| 9 | if bs >= 2 and self.cfg_skip_ratio is not None and self.current_steps >= self.num_inference_steps * (1 - self.cfg_skip_ratio): |
| 10 | bs_half = int(bs // 2) |
| 11 | |
| 12 | new_x = x[bs_half:] |
| 13 | |
| 14 | new_args = [] |
| 15 | for arg in args: |
| 16 | if isinstance(arg, (torch.Tensor, list, tuple, np.ndarray)): |
| 17 | new_args.append(arg[bs_half:]) |
| 18 | else: |
| 19 | new_args.append(arg) |
| 20 | |
| 21 | new_kwargs = {} |
| 22 | for key, content in kwargs.items(): |
| 23 | if isinstance(content, (torch.Tensor, list, tuple, np.ndarray)): |
| 24 | new_kwargs[key] = content[bs_half:] |
| 25 | else: |
| 26 | new_kwargs[key] = content |
| 27 | else: |
| 28 | new_x = x |
| 29 | new_args = args |
| 30 | new_kwargs = kwargs |
| 31 | |
| 32 | result = func(self, new_x, *new_args, **new_kwargs) |
| 33 | |
| 34 | if bs >= 2 and self.cfg_skip_ratio is not None and self.current_steps >= self.num_inference_steps * (1 - self.cfg_skip_ratio): |
| 35 | result = torch.cat([result, result], dim=0) |
| 36 | |
| 37 | return result |
| 38 | return wrapper |
| 39 | return decorator |
nothing calls this directly
no outgoing calls
no test coverage detected