165 lines
6.4 KiB
Python
165 lines
6.4 KiB
Python
# Thin extension of the official diffusers AutoencoderKLWan.
|
|
# All model / tiling / patchify logic comes from diffusers so it stays in sync
|
|
# with upstream. The only thing added here is gradient checkpointing support:
|
|
# the encoder/decoder forwards are monkeypatched with copies of the upstream
|
|
# forwards that add checkpointing branches, and the subclass re-enables
|
|
# _supports_gradient_checkpointing (upstream has it turned off).
|
|
# The checkpointing branches only run cache-free (the chunked-cache path
|
|
# mutates shared lists, which is unsafe under checkpoint recompute), but
|
|
# upstream _encode always goes through the chunked cache — so when
|
|
# checkpointing is active, _encode runs the encoder cache-free over the whole
|
|
# clip in one pass instead. Causal convs make that pass identical to the
|
|
# chunked one; the only divergence is downsample3d WanResample, whose temporal
|
|
# conv upstream applies only via the cache path, so it gets a cache-free
|
|
# forward that reproduces the chunked semantics exactly:
|
|
# out = cat([x[:, :, :1], time_conv(x)], dim=2)
|
|
# (chunk 0 is returned as-is; later chunks run time_conv(cat([prev_last,
|
|
# chunk])), which tiles into exactly these stride-2 windows).
|
|
|
|
import torch
|
|
|
|
from diffusers.models.autoencoders.autoencoder_kl_wan import (
|
|
CACHE_T,
|
|
AutoencoderKLWan as AutoencoderKLWanBase,
|
|
WanDecoder3d,
|
|
WanEncoder3d,
|
|
WanResample,
|
|
patchify,
|
|
)
|
|
|
|
|
|
# copied from diffusers WanEncoder3d.forward with gradient checkpointing added
|
|
def _wan_encoder_forward(self, x, feat_cache=None, feat_idx=[0]):
|
|
use_ckpt = torch.is_grad_enabled() and self.gradient_checkpointing and feat_cache is None
|
|
|
|
if feat_cache is not None:
|
|
idx = feat_idx[0]
|
|
cache_x = x[:, :, -CACHE_T:, :, :].clone()
|
|
if cache_x.shape[2] < 2 and feat_cache[idx] is not None:
|
|
# cache last frame of last two chunk
|
|
cache_x = torch.cat([feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x], dim=2)
|
|
x = self.conv_in(x, feat_cache[idx])
|
|
feat_cache[idx] = cache_x
|
|
feat_idx[0] += 1
|
|
else:
|
|
x = self.conv_in(x)
|
|
|
|
## downsamples
|
|
for layer in self.down_blocks:
|
|
if use_ckpt:
|
|
x = self._gradient_checkpointing_func(layer, x)
|
|
elif feat_cache is not None:
|
|
x = layer(x, feat_cache=feat_cache, feat_idx=feat_idx)
|
|
else:
|
|
x = layer(x)
|
|
|
|
## middle
|
|
if use_ckpt:
|
|
x = self._gradient_checkpointing_func(self.mid_block, x)
|
|
else:
|
|
x = self.mid_block(x, feat_cache=feat_cache, feat_idx=feat_idx)
|
|
|
|
## head
|
|
x = self.norm_out(x)
|
|
x = self.nonlinearity(x)
|
|
if feat_cache is not None:
|
|
idx = feat_idx[0]
|
|
cache_x = x[:, :, -CACHE_T:, :, :].clone()
|
|
if cache_x.shape[2] < 2 and feat_cache[idx] is not None:
|
|
# cache last frame of last two chunk
|
|
cache_x = torch.cat([feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x], dim=2)
|
|
x = self.conv_out(x, feat_cache[idx])
|
|
feat_cache[idx] = cache_x
|
|
feat_idx[0] += 1
|
|
else:
|
|
x = self.conv_out(x)
|
|
|
|
return x
|
|
|
|
|
|
# copied from diffusers WanDecoder3d.forward with gradient checkpointing added
|
|
def _wan_decoder_forward(self, x, feat_cache=None, feat_idx=[0], first_chunk=False):
|
|
use_ckpt = torch.is_grad_enabled() and self.gradient_checkpointing and feat_cache is None
|
|
|
|
## conv1
|
|
if feat_cache is not None:
|
|
idx = feat_idx[0]
|
|
cache_x = x[:, :, -CACHE_T:, :, :].clone()
|
|
if cache_x.shape[2] < 2 and feat_cache[idx] is not None:
|
|
# cache last frame of last two chunk
|
|
cache_x = torch.cat([feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x], dim=2)
|
|
x = self.conv_in(x, feat_cache[idx])
|
|
feat_cache[idx] = cache_x
|
|
feat_idx[0] += 1
|
|
else:
|
|
x = self.conv_in(x)
|
|
|
|
## middle
|
|
if use_ckpt:
|
|
x = self._gradient_checkpointing_func(self.mid_block, x)
|
|
else:
|
|
x = self.mid_block(x, feat_cache=feat_cache, feat_idx=feat_idx)
|
|
|
|
## upsamples
|
|
for up_block in self.up_blocks:
|
|
if use_ckpt:
|
|
x = self._gradient_checkpointing_func(up_block, x, None, [0], first_chunk)
|
|
else:
|
|
x = up_block(x, feat_cache=feat_cache, feat_idx=feat_idx, first_chunk=first_chunk)
|
|
|
|
## head
|
|
x = self.norm_out(x)
|
|
x = self.nonlinearity(x)
|
|
if feat_cache is not None:
|
|
idx = feat_idx[0]
|
|
cache_x = x[:, :, -CACHE_T:, :, :].clone()
|
|
if cache_x.shape[2] < 2 and feat_cache[idx] is not None:
|
|
# cache last frame of last two chunk
|
|
cache_x = torch.cat([feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x], dim=2)
|
|
x = self.conv_out(x, feat_cache[idx])
|
|
feat_cache[idx] = cache_x
|
|
feat_idx[0] += 1
|
|
else:
|
|
x = self.conv_out(x)
|
|
return x
|
|
|
|
|
|
_wan_resample_stock_forward = WanResample.forward
|
|
|
|
|
|
# cache-free downsample3d matching the chunked-cache semantics (see module
|
|
# docstring); upstream's cache-free path skips the temporal conv entirely
|
|
def _wan_resample_forward(self, x, feat_cache=None, feat_idx=[0]):
|
|
if self.mode != "downsample3d" or feat_cache is not None:
|
|
return _wan_resample_stock_forward(self, x, feat_cache=feat_cache, feat_idx=feat_idx)
|
|
b, c, t, h, w = x.size()
|
|
x = x.permute(0, 2, 1, 3, 4).reshape(b * t, c, h, w)
|
|
x = self.resample(x)
|
|
x = x.view(b, t, x.size(1), x.size(2), x.size(3)).permute(0, 2, 1, 3, 4)
|
|
x = torch.cat([x[:, :, :1], self.time_conv(x)], dim=2)
|
|
return x
|
|
|
|
|
|
WanEncoder3d.forward = _wan_encoder_forward
|
|
WanDecoder3d.forward = _wan_decoder_forward
|
|
WanResample.forward = _wan_resample_forward
|
|
|
|
|
|
class AutoencoderKLWan(AutoencoderKLWanBase):
|
|
_supports_gradient_checkpointing = True
|
|
|
|
def _encode(self, x: torch.Tensor):
|
|
# the chunked-cache encode never reaches the checkpointing branches in
|
|
# _wan_encoder_forward (they need feat_cache=None), so run cache-free
|
|
# in one pass when checkpointing is active
|
|
if not (torch.is_grad_enabled() and self.encoder.gradient_checkpointing):
|
|
return super()._encode(x)
|
|
|
|
_, _, num_frame, height, width = x.shape
|
|
if self.config.patch_size is not None:
|
|
x = patchify(x, patch_size=self.config.patch_size)
|
|
if self.use_tiling and (width > self.tile_sample_min_width or height > self.tile_sample_min_height):
|
|
return self.tiled_encode(x)
|
|
out = self.encoder(x)
|
|
return self.quant_conv(out)
|