| from functools import wraps |
| import torch |
| import diffusers |
|
|
| |
|
|
|
|
| |
| original_fourier_filter = diffusers.utils.torch_utils.fourier_filter |
| @wraps(diffusers.utils.torch_utils.fourier_filter) |
| def fourier_filter(x_in, threshold, scale): |
| return_dtype = x_in.dtype |
| return original_fourier_filter(x_in.to(dtype=torch.float32), threshold, scale).to(dtype=return_dtype) |
|
|
|
|
| |
| class FluxPosEmbed(torch.nn.Module): |
| def __init__(self, theta: int, axes_dim): |
| super().__init__() |
| self.theta = theta |
| self.axes_dim = axes_dim |
|
|
| def forward(self, ids: torch.Tensor) -> torch.Tensor: |
| n_axes = ids.shape[-1] |
| cos_out = [] |
| sin_out = [] |
| pos = ids.float() |
| for i in range(n_axes): |
| cos, sin = diffusers.models.embeddings.get_1d_rotary_pos_embed( |
| self.axes_dim[i], |
| pos[:, i], |
| theta=self.theta, |
| repeat_interleave_real=True, |
| use_real=True, |
| freqs_dtype=torch.float32, |
| ) |
| cos_out.append(cos) |
| sin_out.append(sin) |
| freqs_cos = torch.cat(cos_out, dim=-1).to(ids.device) |
| freqs_sin = torch.cat(sin_out, dim=-1).to(ids.device) |
| return freqs_cos, freqs_sin |
|
|
|
|
| def ipex_diffusers(device_supports_fp64=False, can_allocate_plus_4gb=False): |
| diffusers.utils.torch_utils.fourier_filter = fourier_filter |
| if not device_supports_fp64: |
| diffusers.models.embeddings.FluxPosEmbed = FluxPosEmbed |
|
|
