FBCache and comfyUI

Author: Bluear7878

examples/flux.1-dev-IP-adapter.py

@@ -3,6 +3,7 @@
 from diffusers.utils import load_image
 
 from nunchaku import NunchakuFluxTransformer2dModel
+from nunchaku.caching.diffusers_adapters import apply_cache_on_pipe
 from nunchaku.models.IP_adapter.diffusers_adapters import apply_IPA_on_pipe
 from nunchaku.utils import get_precision
 
@@ -18,9 +19,16 @@
     weight_name="ip_adapter.safetensors",
     image_encoder_pretrained_model_name_or_path="openai/clip-vit-large-patch14",
 )
+apply_cache_on_pipe(
+    pipeline,
+    use_double_fb_cache=True,
+    residual_diff_threshold_multi=0.09,
+    residual_diff_threshold_single=0.12,
+)
 
 apply_IPA_on_pipe(pipeline, ip_adapter_scale=1.0, repo_id="XLabs-AI/flux-ip-adapter-v2")
 
+
 IP_image = load_image("https://github.com/ToTheBeginning/PuLID/blob/main/example_inputs/liuyifei.png?raw=true")
 
 image = pipeline(

nunchaku/caching/diffusers_adapters/flux.py

@@ -13,8 +13,13 @@ def apply_cache_on_transformer(
     use_double_fb_cache: bool = False,
     residual_diff_threshold: float = 0.12,
     residual_diff_threshold_multi: float | None = None,
-    residual_diff_threshold_single: float = 0.1,
+    residual_diff_threshold_single: float | None = None,
 ):
+    if not hasattr(transformer, "_original_forward"):
+        transformer._original_forward = transformer.forward
+    if not hasattr(transformer, "_original_blocks"):
+        transformer._original_blocks = transformer.transformer_blocks
+
     if residual_diff_threshold_multi is None:
         residual_diff_threshold_multi = residual_diff_threshold
 
@@ -49,6 +54,9 @@ def new_forward(self, *args, **kwargs):
 
     transformer.forward = new_forward.__get__(transformer)
     transformer._is_cached = True
+    transformer.use_double_fb_cache = use_double_fb_cache
+    transformer.residual_diff_threshold_multi = residual_diff_threshold_multi
+    transformer.residual_diff_threshold_single = residual_diff_threshold_single
 
     return transformer
 

nunchaku/caching/utils.py

@@ -260,7 +260,7 @@ def forward(
         can_use_cache, _ = get_can_use_cache(
             first_hidden_states_residual,
             threshold=self.residual_diff_threshold,
-            parallelized=self.transformer is not None and getattr(self.transformer, "_is_parallelized", False),
+            parallelized=False,
         )
 
         torch._dynamo.graph_break()
@@ -329,7 +329,7 @@ def __init__(
         verbose: bool = False,
     ):
         super().__init__()
-        self.transformer = transformer
+        # self.transformer = transformer
         self.transformer_blocks = transformer.transformer_blocks
         self.single_transformer_blocks = transformer.single_transformer_blocks
 
@@ -486,7 +486,7 @@ def forward(
             hidden_states=hidden_states,
             encoder_hidden_states=encoder_hidden_states,
             threshold=self.residual_diff_threshold_multi,
-            parallelized=(self.transformer is not None and getattr(self.transformer, "_is_parallelized", False)),
+            parallelized=False,
             mode="multi",
             verbose=self.verbose,
             call_remaining_fn=call_remaining_fn,
@@ -515,7 +515,7 @@ def forward(
             hidden_states=cat_hidden_states,
             encoder_hidden_states=None,
             threshold=self.residual_diff_threshold_single,
-            parallelized=(self.transformer is not None and getattr(self.transformer, "_is_parallelized", False)),
+            parallelized=False,
             mode="single",
             verbose=self.verbose,
             call_remaining_fn=call_remaining_fn_single,
@@ -591,8 +591,9 @@ def call_remaining_multi_transformer_blocks(
         controlnet_single_block_samples=None,
         skip_first_layer=False,
         txt_tokens=None,
+        start_idx=1,
     ):
-        start_idx = 1
+        start_idx = start_idx
         original_hidden_states = hidden_states.clone()
         original_encoder_hidden_states = encoder_hidden_states.clone()
 
@@ -627,8 +628,9 @@ def call_remaining_single_transformer_blocks(
         controlnet_single_block_samples=None,
         skip_first_layer=False,
         txt_tokens=None,
+        start_idx=1,
     ):
-        start_idx = 1
+        start_idx = start_idx
         original_hidden_states = hidden_states.clone()
 
         for idx in range(start_idx, num_single_transformer_blocks):

nunchaku/csrc/flux.h

@@ -220,8 +220,6 @@ class QuantizedFluxModel : public ModuleWrapper<FluxModel> { // : public torch::
                              torch::Tensor temb,
                              torch::Tensor rotary_emb_img,
                              torch::Tensor rotary_emb_context,
-                             torch::Tensor k_img,
-                             torch::Tensor v_img,
                              std::optional<torch::Tensor> controlnet_block_samples        = std::nullopt,
                              std::optional<torch::Tensor> controlnet_single_block_samples = std::nullopt) {
         CUDADeviceContext ctx(deviceId);
@@ -233,8 +231,6 @@ class QuantizedFluxModel : public ModuleWrapper<FluxModel> { // : public torch::
         temb                  = temb.contiguous();
         rotary_emb_img        = rotary_emb_img.contiguous();
         rotary_emb_context    = rotary_emb_context.contiguous();
-        k_img                 = k_img.contiguous();
-        v_img                 = v_img.contiguous();
 
         auto &&[hidden_states_, encoder_hidden_states_, ip_query_] = net->forward_ip_adapter(
             idx,
@@ -243,22 +239,13 @@ class QuantizedFluxModel : public ModuleWrapper<FluxModel> { // : public torch::
             from_torch(temb),
             from_torch(rotary_emb_img),
             from_torch(rotary_emb_context),
-            from_torch(k_img),
-            from_torch(v_img),
             controlnet_block_samples.has_value() ? from_torch(controlnet_block_samples.value().contiguous()) : Tensor{},
             controlnet_single_block_samples.has_value()
                 ? from_torch(controlnet_single_block_samples.value().contiguous())
                 : Tensor{});
-        /*
-        auto ip_attn_output_contig =
-        ip_attn_output_.is_contiguous()
-            ? ip_attn_output_
-            : ip_attn_output_.copy(ip_attn_output_.device());
-        */
 
-        hidden_states         = to_torch(hidden_states_);
-        encoder_hidden_states = to_torch(encoder_hidden_states_);
-        // torch::Tensor ip_attn_output = to_torch(ip_attn_output_contig);
+        hidden_states          = to_torch(hidden_states_);
+        encoder_hidden_states  = to_torch(encoder_hidden_states_);
         torch::Tensor ip_query = to_torch(ip_query_);
         Tensor::synchronizeDevice();
 

nunchaku/csrc/pybind.cpp

@@ -57,8 +57,6 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
              py::arg("temb"),
              py::arg("rotary_emb_img"),
              py::arg("rotary_emb_context"),
-             py::arg("k_img"),
-             py::arg("v_img"),
              py::arg("controlnet_block_samples")        = py::none(),
              py::arg("controlnet_single_block_samples") = py::none())
         .def("forward_single_layer", &QuantizedFluxModel::forward_single_layer)

nunchaku/models/IP_adapter/diffusers_adapters/flux.py

@@ -1,11 +1,16 @@
+import functools
+import unittest
+
 from diffusers import DiffusionPipeline, FluxTransformer2DModel
 from torch import nn
 
+from nunchaku.caching.utils import cache_context, create_cache_context
+from nunchaku.models.IP_adapter.utils import undo_all_mods_on_transformer
+
 from ...IP_adapter import utils
 
 
 def apply_IPA_on_transformer(transformer: FluxTransformer2DModel, *, ip_adapter_scale: float = 1.0, repo_id: str):
-
     IPA_transformer_blocks = nn.ModuleList(
         [
             utils.IPA_TransformerBlocks(
@@ -16,19 +21,52 @@ def apply_IPA_on_transformer(transformer: FluxTransformer2DModel, *, ip_adapter_
             )
         ]
     )
+    if getattr(transformer, "_is_cached", False):
+        IPA_transformer_blocks[0].update_residual_diff_threshold(
+            use_double_fb_cache=transformer.use_double_fb_cache,
+            residual_diff_threshold_multi=transformer.residual_diff_threshold_multi,
+            residual_diff_threshold_single=transformer.residual_diff_threshold_single,
+        )
+        undo_all_mods_on_transformer(transformer)
+        if not hasattr(transformer, "_original_forward"):
+            transformer._original_forward = transformer.forward
+        if not hasattr(transformer, "_original_blocks"):
+            transformer._original_blocks = transformer.transformer_blocks
+
     dummy_single_transformer_blocks = nn.ModuleList()
 
     IPA_transformer_blocks[0].load_ip_adapter_weights_per_layer(repo_id=repo_id)
 
     transformer.transformer_blocks = IPA_transformer_blocks
     transformer.single_transformer_blocks = dummy_single_transformer_blocks
+    original_forward = transformer.forward
 
+    @functools.wraps(original_forward)
+    def new_forward(self, *args, **kwargs):
+        with (
+            unittest.mock.patch.object(self, "transformer_blocks", IPA_transformer_blocks),
+            unittest.mock.patch.object(self, "single_transformer_blocks", dummy_single_transformer_blocks),
+        ):
+            return original_forward(*args, **kwargs)
+
+    transformer.forward = new_forward.__get__(transformer)
     transformer._is_IPA = True
 
     return transformer
 
 
 def apply_IPA_on_pipe(pipe: DiffusionPipeline, *, shallow_patch: bool = False, **kwargs):
+    if getattr(pipe, "_is_cached", False):
+        original_call = pipe.__class__.__call__
+
+        @functools.wraps(original_call)
+        def new_call(self, *args, **kwargs):
+            with cache_context(create_cache_context()):
+                return original_call(self, *args, **kwargs)
+
+        pipe.__class__.__call__ = new_call
+        pipe.__class__._is_cached = True
+
     if not shallow_patch:
         apply_IPA_on_transformer(pipe.transformer, **kwargs)