[Perf] Support TeaCache for MammothModa2 (#5357)
Signed-off-by: nagisa-kun <1434936049@qq.com> Signed-off-by: nagisa <1434936049@qq.com>
This commit is contained in:
1 parent
c9bee166fb
commit
c9c478dd07
12 files changed
+684
-15
No files matched your search
@@ -743,7 +743,11 @@ class Transformer2DModel(ModelMixin, ConfigMixin):
|
||||
ar_image_hidden_states: torch.Tensor | None = None,
|
||||
ar_image_attention_mask: torch.Tensor | None = None,
|
||||
return_dict: bool = False,
|
||||
teacache_branch: str | None = None,
|
||||
) -> torch.Tensor:
|
||||
# Consumed by TeaCacheHook/extractor before this plain forward runs.
|
||||
del teacache_branch
|
||||
|
||||
batch_size, height, width = self._validate_inputs(
|
||||
hidden_states, text_hidden_states, text_attention_mask, ref_image_hidden_states, return_dict
|
||||
)
|
||||
|
||||
Reference in new issue
Block a user