Spaces:
Running
on
Zero
Running
on
Zero
File size: 1,850 Bytes
4cd7f21 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 |
import functools
import unittest
import torch
from qwenimage.models.pipeline_qwenimage_edit_plus import QwenImageEditPlusPipeline
from qwenimage.models.transformer_qwenimage import QwenImageTransformer2DModel
from para_attn.first_block_cache import utils
def apply_cache_on_transformer(
transformer: QwenImageTransformer2DModel,
):
if getattr(transformer, "_is_cached", False):
return transformer
cached_transformer_blocks = torch.nn.ModuleList(
[
utils.CachedTransformerBlocks(
transformer.transformer_blocks,
transformer=transformer,
return_hidden_states_first=False,
)
]
)
original_forward = transformer.forward
@functools.wraps(transformer.__class__.forward)
def new_forward(
self,
*args,
**kwargs,
):
with unittest.mock.patch.object(
self,
"transformer_blocks",
cached_transformer_blocks,
):
return original_forward(
*args,
**kwargs,
)
transformer.forward = new_forward.__get__(transformer)
transformer._is_cached = True
return transformer
def apply_cache_on_pipe(
pipe: QwenImageEditPlusPipeline,
*,
shallow_patch: bool = False,
**kwargs,
):
if not getattr(pipe, "_is_cached", False):
original_call = pipe.__class__.__call__
@functools.wraps(original_call)
def new_call(self, *args, **kwargs_):
with utils.cache_context(utils.create_cache_context(**kwargs)):
return original_call(self, *args, **kwargs_)
pipe.__class__.__call__ = new_call
pipe.__class__._is_cached = True
if not shallow_patch:
apply_cache_on_transformer(pipe.transformer)
return pipe
|