File size: 1,342 Bytes
ffc54ec | 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 | """Single-image VAE decode without the video frame cache. Built with Qwen.
The upstream image VAE is Wan's causal video decoder specialised to one frame. Its decode loop still
maintains a per-convolution frame cache, cloning the input of every causal convolution for a next
frame that never arrives. For one frame the cache only ever holds its initial values, so calling the
decoder without a cache takes the identical arithmetic path minus those copies. Optionally the
cache-free decoder is compiled (kernel fusion of norm/activation/pad around cuDNN convolutions).
"""
import torch
from diffusers.models.autoencoders.vae import DecoderOutput
def enable_fast_decode(vae, compile=False):
if getattr(vae, '_image21_fast_decode', False):
return vae
original = vae._decode
decoder = vae.decoder
run = torch.compile(decoder, dynamic=False) if compile else decoder
def _decode(z, return_dict=True):
if z.shape[2] != 1 or vae.use_tiling or vae.config.patch_size is not None:
return original(z, return_dict=return_dict)
out = run(vae.post_quant_conv(z), feat_cache=None, first_chunk=True)
out = torch.clamp(out, min=-1.0, max=1.0)
return (out,) if not return_dict else DecoderOutput(sample=out)
vae._decode = _decode
vae._image21_fast_decode = True
return vae
|