"""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