92 lines
3.3 KiB
Python
92 lines
3.3 KiB
Python
import importlib.metadata
|
|
import torch
|
|
import logging
|
|
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
|
|
log = logging.getLogger(__name__)
|
|
|
|
def check_diffusers_version():
|
|
try:
|
|
version = importlib.metadata.version('diffusers')
|
|
required_version = '0.31.0'
|
|
if version < required_version:
|
|
raise AssertionError(f"diffusers version {version} is installed, but version {required_version} or higher is required.")
|
|
except importlib.metadata.PackageNotFoundError:
|
|
raise AssertionError("diffusers is not installed.")
|
|
|
|
def print_memory(device):
|
|
memory = torch.cuda.memory_allocated(device) / 1024**3
|
|
max_memory = torch.cuda.max_memory_allocated(device) / 1024**3
|
|
max_reserved = torch.cuda.max_memory_reserved(device) / 1024**3
|
|
log.info(f"-------------------------------")
|
|
log.info(f"Allocated memory: {memory=:.3f} GB")
|
|
log.info(f"Max allocated memory: {max_memory=:.3f} GB")
|
|
log.info(f"Max reserved memory: {max_reserved=:.3f} GB")
|
|
log.info(f"-------------------------------")
|
|
#memory_summary = torch.cuda.memory_summary(device=device, abbreviated=False)
|
|
#log.info(f"Memory Summary:\n{memory_summary}")
|
|
|
|
def optimized_scale(positive_flat, negative_flat):
|
|
|
|
# Calculate dot production
|
|
dot_product = torch.sum(positive_flat * negative_flat, dim=1, keepdim=True)
|
|
|
|
# Squared norm of uncondition
|
|
squared_norm = torch.sum(negative_flat ** 2, dim=1, keepdim=True) + 1e-8
|
|
|
|
# st_star = v_cond^T * v_uncond / ||v_uncond||^2
|
|
st_star = dot_product / squared_norm
|
|
|
|
return st_star
|
|
|
|
# Code based on https://github.com/WikiChao/FreSca (MIT License)
|
|
import torch
|
|
import torch.fft as fft
|
|
|
|
def fourier_filter(x, scale_low=1.0, scale_high=1.5, freq_cutoff=20):
|
|
"""
|
|
Apply frequency-dependent scaling to an image tensor using Fourier transforms.
|
|
|
|
Parameters:
|
|
x: Input tensor of shape (B, C, H, W)
|
|
scale_low: Scaling factor for low-frequency components (default: 1.0)
|
|
scale_high: Scaling factor for high-frequency components (default: 1.5)
|
|
freq_cutoff: Number of frequency indices around center to consider as low-frequency (default: 20)
|
|
|
|
Returns:
|
|
x_filtered: Filtered version of x in spatial domain with frequency-specific scaling applied.
|
|
"""
|
|
# Preserve input dtype and device
|
|
dtype, device = x.dtype, x.device
|
|
|
|
# Convert to float32 for FFT computations
|
|
x = x.to(torch.float32)
|
|
|
|
# 1) Apply FFT and shift low frequencies to center
|
|
x_freq = fft.fftn(x, dim=(-2, -1))
|
|
x_freq = fft.fftshift(x_freq, dim=(-2, -1))
|
|
|
|
# 2) Create a mask to scale frequencies differently
|
|
B, C, T, H, W = x_freq.shape
|
|
crow, ccol = H // 2, W // 2
|
|
|
|
# Initialize mask with high-frequency scaling factor
|
|
mask = torch.ones((B, C, T, H, W), device=device) * scale_high
|
|
|
|
# Apply low-frequency scaling factor to center region
|
|
mask[
|
|
...,
|
|
crow - freq_cutoff : crow + freq_cutoff,
|
|
ccol - freq_cutoff : ccol + freq_cutoff,
|
|
] = scale_low
|
|
|
|
# 3) Apply frequency-specific scaling
|
|
x_freq = x_freq * mask
|
|
|
|
# 4) Convert back to spatial domain
|
|
x_freq = fft.ifftshift(x_freq, dim=(-2, -1))
|
|
x_filtered = fft.ifftn(x_freq, dim=(-2, -1)).real
|
|
|
|
# 5) Restore original dtype
|
|
x_filtered = x_filtered.to(dtype)
|
|
|
|
return x_filtered |