File size: 4,218 Bytes
bee9822 | 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 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 | import logging
import torch
__all__ = ["freq_mix_temporal", "freq_mix_spatial"]
def _real_fft_energy(t_fft: torch.Tensor, band_end: int | None = None) -> torch.Tensor:
"""Return the energy of a real FFT tensor, accounting for mirrored bins."""
if band_end is not None:
t_fft = t_fft[:, :band_end, ...]
dc_energy = t_fft[:, 0, ...].norm() ** 2
mirrored_energy = t_fft[:, 1:, ...].norm() ** 2
return dc_energy + 2 * mirrored_energy
def _temporal_high_band_scale(mixed_t: torch.Tensor, alpha: int, gamma: float) -> torch.Tensor:
total_energy = _real_fft_energy(mixed_t)
low_energy = _real_fft_energy(mixed_t, band_end=alpha)
high_energy = torch.clamp(total_energy - low_energy, min=1e-8)
target_high_energy = total_energy - (low_energy / (gamma**2))
return torch.sqrt(torch.clamp(target_high_energy / high_energy, min=0.0))
def _frequency_radius_grid(latents: torch.Tensor, fft_dims: tuple[int, ...]) -> torch.Tensor:
grids = [torch.linspace(-1, 1, latents.shape[d], device=latents.device) for d in fft_dims]
mesh = torch.meshgrid(*grids, indexing="ij")
rr = torch.zeros_like(mesh[0])
for grid in mesh:
rr = rr + grid**2
return torch.sqrt(rr)
def freq_mix_temporal(l1, l2, gamma=30.0, alpha=3, **kwargs):
"""Mix temporal frequency magnitude from ``l1`` with phase from ``l2``."""
l1, l2 = l1[0], l2[0]
l1_f, l2_f = l1.float(), l2.float()
fft1_t = torch.fft.rfft(l1_f, dim=1, norm='ortho')
fft2_t = torch.fft.rfft(l2_f, dim=1, norm='ortho')
magnitude1_t = torch.abs(fft1_t)
phase2_t = torch.angle(fft2_t)
if alpha > 0:
alpha = int(alpha)
mixed_t = torch.polar(magnitude1_t, phase2_t)
mixed_t[:, alpha:] = fft1_t[:, alpha:]
high_band_scale = _temporal_high_band_scale(mixed_t, alpha, gamma)
temporal_scale = torch.empty(mixed_t.shape[1], device=mixed_t.device, dtype=mixed_t.real.dtype)
temporal_scale[:alpha] = 1.0 / gamma
temporal_scale[alpha:] = high_band_scale
mixed_t_final = mixed_t * temporal_scale[None, :, None, None]
logging.info("beta term: %f", high_band_scale)
logging.info(f'l1_f norm: {l1_f.norm()}\t{l1.norm()}')
else:
mixed_t_final = fft1_t.clone()
combined_latents_t = torch.fft.irfft(mixed_t_final, dim=1, n=l1_f.shape[1], norm='ortho')
return [combined_latents_t.to(l1.dtype)]
def freq_mix_spatial(latents_hi, latents_lo, alpha, gamma, dims=("t", "h", "w"), **kwargs):
"""
Replace LOW-FREQUENCY PHASE of latents_hi with latents_lo
"""
assert latents_hi.shape == latents_lo.shape
device = latents_hi.device
dim_map = {
"t": 1,
"h": 2,
"w": 3,
}
fft_dims = tuple(dim_map[d] for d in dims)
fft_hi = torch.fft.fftn(latents_hi, dim=fft_dims, norm='ortho')
fft_lo = torch.fft.fftn(latents_lo, dim=fft_dims, norm='ortho')
fft_hi = torch.fft.fftshift(fft_hi, dim=fft_dims)
fft_lo = torch.fft.fftshift(fft_lo, dim=fft_dims)
# frequency grid
rr = _frequency_radius_grid(latents_hi, fft_dims)
cutoff = rr.max() / (2 ** alpha)
low_mask = (rr < cutoff).float()
high_mask = 1.0 - low_mask
shape = [1] * latents_hi.ndim
for i, d in enumerate(fft_dims):
shape[d] = low_mask.shape[i]
low_mask = low_mask.reshape(shape)
high_mask = high_mask.reshape(shape)
mag_hi = torch.abs(fft_hi)
phase_hi = torch.angle(fft_hi)
mag_lo = torch.abs(fft_lo)
phase_lo = torch.angle(fft_lo)
# swap phase only
phase_mix = phase_lo * low_mask + phase_hi * high_mask
fft_mix = mag_hi * torch.exp(1j * phase_mix)
# energy over 2D spatial freq bins
power = (torch.abs(fft_mix) ** 2)
total_energy = power.sum()
low_energy = (power * low_mask).sum()
high_energy = (power * high_mask).sum().clamp(min=1e-12)
high_band_scale = torch.sqrt((total_energy - (low_energy / (gamma ** 2))) / high_energy)
scale = (low_mask / gamma) + (high_mask * high_band_scale)
fft_mix = fft_mix * scale
fft_mix = torch.fft.ifftshift(fft_mix, dim=fft_dims)
out = torch.fft.ifftn(fft_mix, dim=fft_dims, norm='ortho').real
return out
|