Skip to content

GitLab

Commit f11a2e2a authored by Hyunsung Lee's avatar Hyunsung Lee Committed by muyangli
Browse files

Update utils.py

parent f04b603c
Hide whitespace changes
Inline Side-by-side
Showing with 3 additions and 0 deletions
nunchaku/caching/utils.py
View file @ f11a2e2a
...@@ -245,6 +245,9 @@ class FluxCachedTransformerBlocks(nn.Module): ...@@ -245,6 +245,9 @@ class FluxCachedTransformerBlocks(nn.Module):
self.return_hidden_states_only = return_hidden_states_only self.return_hidden_states_only = return_hidden_states_only
self.verbose = verbose self.verbose = verbose
def update_residual_diff_threshold(self, residual_diff_threshold=0.12):
self.residual_diff_threshold = residual_diff_threshold
def forward(self, hidden_states, encoder_hidden_states, *args, **kwargs): def forward(self, hidden_states, encoder_hidden_states, *args, **kwargs):
batch_size = hidden_states.shape[0] batch_size = hidden_states.shape[0]
if self.residual_diff_threshold <= 0.0 or batch_size > 1: if self.residual_diff_threshold <= 0.0 or batch_size > 1:
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment