From 66ba4e1ee7975bdeeb0dec6f49385e363a686907 Mon Sep 17 00:00:00 2001 From: kijai <40791699+kijai@users.noreply.github.com> Date: Mon, 28 Oct 2024 22:52:30 +0200 Subject: [PATCH] fix fastercache start step --- custom_cogvideox_transformer_3d.py | 32 +++++++++++------------------- 1 file changed, 12 insertions(+), 20 deletions(-) diff --git a/custom_cogvideox_transformer_3d.py b/custom_cogvideox_transformer_3d.py index cb34513..d4bc137 100644 --- a/custom_cogvideox_transformer_3d.py +++ b/custom_cogvideox_transformer_3d.py @@ -318,36 +318,28 @@ class CogVideoXBlock(nn.Module): del h, fuser #fastercache - if fastercache_counter >= fastercache_start_step + 3 and fastercache_counter%3!=0 and self.cached_hidden_states[-1].shape[0]>=norm_hidden_states.shape[0]: + B = norm_hidden_states.shape[0] + if fastercache_counter >= fastercache_start_step + 3 and fastercache_counter%3!=0 and self.cached_hidden_states[-1].shape[0] >= B: attn_hidden_states = ( - self.cached_hidden_states[1][:norm_hidden_states.shape[0]] + - (self.cached_hidden_states[1][:norm_hidden_states.shape[0]] - - self.cached_hidden_states[0][:norm_hidden_states.shape[0]]) + self.cached_hidden_states[1][:B] + + (self.cached_hidden_states[1][:B] - self.cached_hidden_states[0][:B]) * 0.3 ).to(norm_hidden_states.device, non_blocking=True) attn_encoder_hidden_states = ( - self.cached_encoder_hidden_states[1][:norm_hidden_states.shape[0]] + - (self.cached_encoder_hidden_states[1][:norm_hidden_states.shape[0]] - - self.cached_encoder_hidden_states[0][:norm_hidden_states.shape[0]]) - *0.3 + self.cached_encoder_hidden_states[1][:B] + + (self.cached_encoder_hidden_states[1][:B] - self.cached_encoder_hidden_states[0][:B]) + * 0.3 ).to(norm_hidden_states.device, non_blocking=True) - else: attn_hidden_states, attn_encoder_hidden_states = self.attn1( hidden_states=norm_hidden_states, encoder_hidden_states=norm_encoder_hidden_states, image_rotary_emb=image_rotary_emb, ) - if fastercache_counter==fastercache_start_step: - self.cached_hidden_states = [ - attn_hidden_states.to(fastercache_device), - attn_hidden_states.to(fastercache_device) - ] - self.cached_encoder_hidden_states = [ - attn_encoder_hidden_states.to(fastercache_device), - attn_encoder_hidden_states.to(fastercache_device) - ] - elif fastercache_counter>fastercache_start_step: + if fastercache_counter == fastercache_start_step: + self.cached_hidden_states = [attn_hidden_states.to(fastercache_device), attn_hidden_states.to(fastercache_device)] + self.cached_encoder_hidden_states = [attn_encoder_hidden_states.to(fastercache_device), attn_encoder_hidden_states.to(fastercache_device)] + elif fastercache_counter > fastercache_start_step: self.cached_hidden_states[-1].copy_(attn_hidden_states.to(fastercache_device)) self.cached_encoder_hidden_states[-1].copy_(attn_encoder_hidden_states.to(fastercache_device)) @@ -769,7 +761,7 @@ class CogVideoXTransformer3DModel(ModelMixin, ConfigMixin): output = hidden_states.reshape(batch_size, num_frames, height // p, width // p, -1, p, p) output = output.permute(0, 1, 4, 2, 5, 3, 6).flatten(5, 6).flatten(3, 4) - if self.fastercache_counter>=16: + if self.fastercache_counter >= self.fastercache_start_step + 1: (bb, tt, cc, hh, ww) = output.shape cond = rearrange(output[0:1].float(), "B T C H W -> (B T) C H W", B=bb//2, C=cc, T=tt, H=hh, W=ww) uncond = rearrange(output[1:2].float(), "B T C H W -> (B T) C H W", B=bb//2, C=cc, T=tt, H=hh, W=ww)