手机建网站需要多少钱旅游最新资讯
FreeU: Free Lunch in Diffusion U-Net
摘要
作者研究了 U-Net 架构对去噪过程的关键贡献,并发现其主干部分主要在去噪方面发挥作用,而其跳跃连接主要是向解码器模块引入高频特征,这使得网络忽略了主干部分的语义信息。基于这一发现,我们提出了一种简单却有效的方法-- “FreeU”,它无需额外训练或微调就能提升生成质量。我们的核心思路是从策略上对源自 U-Net 跳跃连接和主干特征图的贡献进行重新加权,以充分利用 U-Net 架构中这两个组件的优势。在图像和视频生成任务上取得的良好结果表明, FreeU 方法可以很容易地集成到现有的扩散模型中,例如稳定扩散(Stable Diffusion)、DreamBooth、ModelScope、Rerender 和 ReVersion 等,只需几行代码就能提升生成质量。
试验表明,如果把decoder阶段的全部backbone都放大,会导致oversmoothed texture。为了缓解这种情况,只在decoder的前两个阶段使用,放大backbone,并且缩小skip features。skip features需要进行FFT和IFFT,详见函数 fourier_filter代码。
完整的stable diffusion1.5的UNet结构可参考UNet2DConditionModel
SDXL效果对比
参数,来自于FreeU
SD1.4: (will be updated soon)
b1: 1.3, b2: 1.4, s1: 0.9, s2: 0.2SD1.5: (will be updated soon)
b1: 1.5, b2: 1.6, s1: 0.9, s2: 0.2SD2.1
b1: 1.1, b2: 1.2, s1: 0.9, s2: 0.2
b1: 1.4, b2: 1.6, s1: 0.9, s2: 0.2SDXL
b1: 1.3, b2: 1.4, s1: 0.9, s2: 0.2 SDXL resultsRange for More Parameters
When trying additional parameters, consider the following ranges:b1: 1 ≤ b1 ≤ 1.2
b2: 1.2 ≤ b2 ≤ 1.6
s1: s1 ≤ 1
s2: s2 ≤ 1
代码
使用方法
import torch
from diffusers import DiffusionPipelinepipeline = DiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.float16,
).to("cuda")
pipeline.enable_freeu(s1=0.9, s2=0.2, b1=1.3, b2=1.4) ##add
generator = torch.Generator(device="cpu").manual_seed(13)
prompt = "A squirrel eating a burger"
image = pipeline(prompt, generator=generator).images[0]
image
FreeU函数(来自于diffusers)
def apply_freeu(resolution_idx: int, hidden_states: "torch.Tensor", res_hidden_states: "torch.Tensor", **freeu_kwargs
) -> Tuple["torch.Tensor", "torch.Tensor"]:"""Applies the FreeU mechanism as introduced in https://arxiv.org/abs/2309.11497. Adapted from the official code repository: https://github.com/ChenyangSi/FreeU.Args:resolution_idx (`int`): Integer denoting the UNet block where FreeU is being applied.hidden_states (`torch.Tensor`): Inputs to the underlying block.res_hidden_states (`torch.Tensor`): Features from the skip block corresponding to the underlying block.s1 (`float`): Scaling factor for stage 1 to attenuate the contributions of the skip features.s2 (`float`): Scaling factor for stage 2 to attenuate the contributions of the skip features.b1 (`float`): Scaling factor for stage 1 to amplify the contributions of backbone features.b2 (`float`): Scaling factor for stage 2 to amplify the contributions of backbone features."""if resolution_idx == 0:num_half_channels = hidden_states.shape[1] // 2hidden_states[:, :num_half_channels] = hidden_states[:, :num_half_channels] * freeu_kwargs["b1"]res_hidden_states = fourier_filter(res_hidden_states, threshold=1, scale=freeu_kwargs["s1"])if resolution_idx == 1:num_half_channels = hidden_states.shape[1] // 2hidden_states[:, :num_half_channels] = hidden_states[:, :num_half_channels] * freeu_kwargs["b2"]res_hidden_states = fourier_filter(res_hidden_states, threshold=1, scale=freeu_kwargs["s2"])return hidden_states, res_hidden_states
def fourier_filter(x_in: "torch.Tensor", threshold: int, scale: int) -> "torch.Tensor":"""Fourier filter as introduced in FreeU (https://arxiv.org/abs/2309.11497).This version of the method comes from here:https://github.com/huggingface/diffusers/pull/5164#issuecomment-1732638706"""x = x_inB, C, H, W = x.shape# Non-power of 2 images must be float32if (W & (W - 1)) != 0 or (H & (H - 1)) != 0:x = x.to(dtype=torch.float32)# fftn does not support bfloat16elif x.dtype == torch.bfloat16:x = x.to(dtype=torch.float32)# FFTx_freq = fftn(x, dim=(-2, -1))x_freq = fftshift(x_freq, dim=(-2, -1))B, C, H, W = x_freq.shapemask = torch.ones((B, C, H, W), device=x.device)crow, ccol = H // 2, W // 2mask[..., crow - threshold : crow + threshold, ccol - threshold : ccol + threshold] = scalex_freq = x_freq * mask# IFFTx_freq = ifftshift(x_freq, dim=(-2, -1))x_filtered = ifftn(x_freq, dim=(-2, -1)).realreturn x_filtered.to(dtype=x_in.dtype)