[Wan2.2 VAE, Qwen-Image-2.1 VAE] Redundant x.clone() in residual up/down blocks keeps an extra activation alive #15040

Describe the bug

WanResidualDownBlock.forward and WanResidualUpBlock.forward in autoencoder_kl_wan.py, and their copies QwenImage21ResidualDownBlock / QwenImage21ResidualUpBlock in autoencoder_kl_qwenimage21.py, start with:

x_copy = x.clone()

x_copy is only read once, at the end, by self.avg_shortcut(x_copy). Nothing in between modifies x in place (the resnets, resample and shortcut layers all return new tensors), so the clone does not protect anything. It only keeps a second activation-sized buffer alive for the whole block, which raises peak memory in both encode and decode.

These are the only four x_copy = x.clone() sites in src/diffusers/models/autoencoders/. The feat_cache[idx] = x.clone() calls elsewhere in these files store tensors in a cache that outlives the call, so they are a different case and not part of this.

Proposed fix: replace x_copy = x.clone() with x_copy = x at the four sites. I opened #15029 with this change before reading the coordination step in the contribution guide, sorry about that. I'm happy to adjust its scope or close it if you'd prefer a different approach.

Reproduction

import torch
from diffusers import AutoencoderKLQwenImage21

vae = AutoencoderKLQwenImage21.from_pretrained(
    "Qwen/Qwen-Image-2.1", subfolder="vae", torch_dtype=torch.bfloat16
).to("cuda").eval()

# latents for a 2048x2048 image
latents = torch.randn(1, 64, 1, 128, 128, device="cuda", dtype=torch.bfloat16)

torch.cuda.reset_peak_memory_stats()
with torch.no_grad():
    vae.decode(latents).sample
print(f"peak allocated: {torch.cuda.max_memory_allocated() / 2**30:.2f} GiB")
# 26.87 GiB on main, 24.62 GiB with the clone removed

The same applies to AutoencoderKLWan with Wan-AI/Wan2.2-TI2V-5B-Diffusers.

Peak allocated memory before and after replacing the clone with x_copy = x, released weights, A100-80GB, torch 2.10. Outputs are bitwise identical (torch.equal) in every case:

Runtime is unchanged.

Logs

System Info

• diffusers: main
• torch 2.10.0, CUDA 13.0, A100-80GB

Who can help?

@sayakpaul @yiyixuxu