Spaces:
Running on Zero
Running on Zero
File size: 6,008 Bytes
076b9c5 107af9c 076b9c5 107af9c 076b9c5 | 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 | """Decoder refiner (Sec. 3.1, Appendix B.2): the Stable Diffusion 1.5 VAE decoder with fusion modules that inject
features of the blurred input into its up-blocks, recovering structure that VAE decoding loses.
SD1.5 and SD2.1 share the VAE encoder and quant_conv, so the refiner decodes the latents of the SD2.1-based Stage I.
The input enters as (a) the VAE-encoder features at the resolutions of the decoder stages (conv_in, down-blocks 0 and 1,
matched by spatial size) and (b) a guidance map computed from the image and its high-frequency magnitude
(Laplacian-pyramid residual), shared by every fusion module. Each fusion module adds a zero-initialised residual to the
decoder feature, so training starts from the plain decoder. Parameter names follow the released checkpoint."""
import torch
import torch.nn.functional as F
from torch import nn
from torchvision.transforms.functional import gaussian_blur
from .. import paths
from .._compat import patch_huggingface_hub
from ..layers import NAFBlock
patch_huggingface_hub() # before diffusers is imported
def high_frequency_residual(image, kernel_size=5, sigma=1.0):
"""image minus its blurred 2x down-up-sampled version (one Laplacian-pyramid level)."""
low = gaussian_blur(image, kernel_size=kernel_size, sigma=sigma)[:, :, ::2, ::2]
up = gaussian_blur(F.interpolate(low, size=image.shape[-2:], mode="bilinear", align_corners=False),
kernel_size=kernel_size, sigma=sigma)
return image - up
class FeatureFusion(nn.Module):
"""Decoder feature + fusion of (encoder feature, decoder feature, guidance map)."""
def __init__(self, decoder_channels, encoder_channels, guidance_channels, depth):
super().__init__()
self.in_conv = nn.Conv2d(encoder_channels + decoder_channels, decoder_channels, 3, 1, 1)
self.dm_blocks = nn.ModuleList([NAFBlock(decoder_channels + guidance_channels) for _ in range(depth)])
self.out_conv = nn.Conv2d(decoder_channels + guidance_channels, decoder_channels, 3, 1, 1)
def forward(self, encoder_feat, decoder_feat, guidance):
if guidance.shape[-2:] != decoder_feat.shape[-2:]:
guidance = F.interpolate(guidance, size=decoder_feat.shape[-2:], mode="bilinear", align_corners=False)
fused = torch.cat((self.in_conv(torch.cat((encoder_feat, decoder_feat), dim=1)), guidance), dim=1)
for block in self.dm_blocks:
fused = block(fused)
return self.out_conv(fused) + decoder_feat
def _last_resnet_channels(block):
return int(block.resnets[-1].out_channels)
class DecoderRefiner(nn.Module):
ENCODER_STAGES = 2 # down-blocks whose features the decoder can match (the decoder runs at H/4, H/2, H, H)
def __init__(self, vae, depth=2, guidance_channels=32):
super().__init__()
self.vae = vae
decoder_channels = [_last_resnet_channels(b) for b in vae.decoder.up_blocks]
encoder_channels = [int(vae.encoder.conv_in.out_channels)] + [_last_resnet_channels(b) for b in vae.encoder.down_blocks]
n = len(decoder_channels)
self.fusion_blocks = nn.ModuleList([
FeatureFusion(dec, encoder_channels[max(0, n - 2 - i)], guidance_channels, depth)
for i, dec in enumerate(decoder_channels)])
self.map_init = nn.Sequential(nn.Conv2d(3 + 1, guidance_channels, 3, 1, 1), # image + high frequency
*[NAFBlock(guidance_channels) for _ in range(depth)])
for block in self.fusion_blocks:
nn.init.zeros_(block.out_conv.weight)
nn.init.zeros_(block.out_conv.bias)
def encoder_features(self, image):
"""VAE-encoder features at the resolutions of the decoder stages: conv_in and the first down-blocks."""
with torch.no_grad():
h = self.vae.encoder.conv_in(image)
feats = [h.detach()]
for block in self.vae.encoder.down_blocks[:self.ENCODER_STAGES]:
h = block(h)
feats.append(h.detach())
return feats
@staticmethod
def _matching_feature(decoder_feat, encoder_features):
for feat in reversed(encoder_features):
if feat.shape[0] == decoder_feat.shape[0] and feat.shape[-2:] == decoder_feat.shape[-2:]:
return feat
return None
def forward(self, image, latent):
"""image: the blurred input in [-1, 1]; latent: the restored latent (unscaled). Returns the image in [-1, 1]."""
encoder_features = self.encoder_features(image)
high_freq = high_frequency_residual(image).abs().mean(dim=1, keepdim=True)
guidance = self.map_init(torch.cat((image, high_freq), dim=1))
h = self.vae.decoder.mid_block(self.vae.decoder.conv_in(self.vae.post_quant_conv(latent)))
for block, fusion in zip(self.vae.decoder.up_blocks, self.fusion_blocks):
h = block(h)
feat = self._matching_feature(h, encoder_features)
if feat is not None:
h = fusion(feat, h, guidance)
return self.vae.decoder.conv_out(self.vae.decoder.conv_act(self.vae.decoder.conv_norm_out(h)))
def build_decoder_refiner(sd15=paths.SD15):
from diffusers import AutoencoderKL
return DecoderRefiner(AutoencoderKL.from_pretrained(str(sd15), subfolder="vae"))
def load_decoder_refiner(path=paths.REFINER, sd15=paths.SD15, device="cuda", dtype=torch.bfloat16):
"""The released refiner, in bf16 as trained. Its checkpoint stores every parameter except the VAE encoder, which
comes from SD1.5 (only its first two down-blocks run)."""
refiner = build_decoder_refiner(sd15)
res = refiner.load_state_dict(torch.load(str(path), map_location="cpu"), strict=False)
missing = [k for k in res.missing_keys if not k.startswith(("vae.encoder", "vae.quant_conv", "vae.post_quant_conv"))]
assert not missing and not res.unexpected_keys, (missing[:5], res.unexpected_keys[:5])
return refiner.to(device, dtype=dtype).eval()
|