Spaces:
Running on Zero
Running on Zero
| """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 | |
| 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() | |