[Bug] txt2_img pipe and controlnet pipe share components. Error will be reported when using cache

[todochenxi]

Checklist

• 1. I have searched for related issues and FAQs (https://github.com/mit-han-lab/nunchaku/blob/main/docs/faq.md) but was unable to find a solution.
• 2. The issue persists in the latest version.
• 3. Please note that without environment information and a minimal reproducible example, it will be difficult for us to reproduce and address the issue, which may delay our response.
• 4. If your report is a question rather than a bug, please submit it as a discussion at https://github.com/mit-han-lab/nunchaku/discussions/new/choose. Otherwise, this issue will be closed.
• 5. If this is related to ComfyUI, please report it at https://github.com/mit-han-lab/ComfyUI-nunchaku/issues.
• 6. I will do my best to describe the issue in English.

Describe the Bug

I have provided a reproducible script and several error situations

1. apply cache for pipe and pipe_cn: error msg = "AttributeError: 'NunchakuFluxTransformer2dModel' object has no attribute 'cached_transformer_blocks'. Did you mean: 'single_transformer_blocks'?"
2. apply cache only for pipe or pipe_cn， then use another pipe to infer: error msg = "assert cache_context is not None, "cache_context must be set before"
3. apply cache for pipe_cn don't have acceleration effect
4. In version 2.0, I noticed that the cache setting for pipe will override the setting of pipe_cn, and the cache of pipe_cn must be set, otherwise an error will be reported.

Environment

ubuntu 22.04 python=3.10 torch=2.7 cuda=12.8

Reproduction Steps

import torch
from diffusers import (
    FluxPipeline, 
    FluxInpaintPipeline,
    FluxTransformer2DModel,
    FlowMatchEulerDiscreteScheduler,
    AutoencoderKL,
    FluxControlNetModel,
    FluxControlNetPipeline,
    FluxFillPipeline
)
from diffusers.models import FluxMultiControlNetModel
from transformers import T5EncoderModel, CLIPTextModel, CLIPTokenizer,T5TokenizerFast
from nunchaku.models.transformers.transformer_flux import NunchakuFluxTransformer2dModel
from nunchaku import NunchakuT5EncoderModel
from nunchaku.utils import get_precision
from nunchaku.caching.diffusers_adapters import apply_cache_on_pipe
from diffusers.utils import load_image

bfl_repo = "black-forest-labs/FLUX.1-dev"
dtype =  torch.bfloat16  # or torch.float16, or torch.float32
device = "cuda"  # or "cpu" if you want to run on CPU

scheduler       = FlowMatchEulerDiscreteScheduler.from_pretrained(bfl_repo, subfolder="scheduler", torch_dtype=dtype)
text_encoder    = CLIPTextModel.from_pretrained(bfl_repo, subfolder="text_encoder", torch_dtype=dtype)
text_encoder_2  = T5EncoderModel.from_pretrained(bfl_repo, subfolder="text_encoder_2", torch_dtype=dtype)
tokenizer       = CLIPTokenizer.from_pretrained(bfl_repo, subfolder="tokenizer", torch_dtype=dtype, clean_up_tokenization_spaces=True)
tokenizer_2     = T5TokenizerFast.from_pretrained(bfl_repo, subfolder="tokenizer_2", torch_dtype=dtype, clean_up_tokenization_spaces=True)
vae             = AutoencoderKL.from_pretrained(bfl_repo, subfolder="vae", torch_dtype=dtype)
precision = get_precision()
transformer = NunchakuFluxTransformer2dModel.from_pretrained(
    f"mit-han-lab/nunchaku-flux.1-dev/svdq-{precision}_r32-flux.1-dev.safetensors",
    # offload=True
)
transformer.set_attention_impl("nunchaku-fp16")

# qencoder
text_encoder_2 = NunchakuT5EncoderModel.from_pretrained("mit-han-lab/nunchaku-t5/awq-int4-flux.1-t5xxl.safetensors")
controlnet_union = FluxControlNetModel.from_pretrained("Shakker-Labs/FLUX.1-dev-ControlNet-Union-Pro-2.0", torch_dtype=torch.bfloat16)
controlnet = FluxMultiControlNetModel([controlnet_union])  # we always recommend loading via FluxMultiControlNetModel


params = {
    "scheduler": scheduler,
    "vae": vae,
    "tokenizer": tokenizer,
    "tokenizer_2": tokenizer_2,
    "text_encoder": text_encoder,
    "text_encoder_2": text_encoder_2,
    "transformer": transformer
}
# pipe
pipe = FluxPipeline(**params).to(device, dtype=dtype)
pipe_cn = FluxControlNetPipeline(**params, controlnet=controlnet).to(device, dtype)

# offload
pipe.enable_sequential_cpu_offload(device=device)
pipe_cn.enable_sequential_cpu_offload(device=device)

# cache
apply_cache_on_pipe(
    pipe_cn,
    use_double_fb_cache=True,
    residual_diff_threshold_multi=0.09,
    residual_diff_threshold_single=0.12,
)
apply_cache_on_pipe(
     pipe_cn,
     use_double_fb_cache=False,
     residual_diff_threshold_multi=0,
     residual_diff_threshold_single=0,
 )

params = {
    "prompt": "A bohemian-style female travel blogger with sun-kissed skin and messy beach waves.",
    "height": 1152,
    "width": 768,
    "num_inference_steps": 30,
    "guidance_scale": 3.5,
}

# pipe
txt2img_res = pipe(
    **params,
).images[0]
txt2img_res.save("flux.1-dev-txt2img.jpg")

# pipe_cn
control_iamge = load_image("https://huggingface.co/Shakker-Labs/FLUX.1-dev-ControlNet-Union-Pro/resolve/main/assets/openpose.jpg")
params["control_images"] = [control_iamge]
params["controlnet_conditioning_scale"] = [0.9]
params['control_guidance_end'] = [0.65]
cn_res = pipe_cn(
    **params,
).images[0]
cn_res.save("flux.1-dev-cn-txt2img.jpg")

[lmxyy]

@Bluear7878 could you take a look?

[Bluear7878]

I solved this issue. @todochenxi #452

You can test it with the code below.

import torch
from diffusers import (
    FluxPipeline, 
    FluxInpaintPipeline,
    FluxTransformer2DModel,
    FlowMatchEulerDiscreteScheduler,
    AutoencoderKL,
    FluxControlNetModel,
    FluxControlNetPipeline,
    FluxFillPipeline
)
from diffusers.models import FluxMultiControlNetModel
from transformers import T5EncoderModel, CLIPTextModel, CLIPTokenizer,T5TokenizerFast
from nunchaku.models.transformers.transformer_flux import NunchakuFluxTransformer2dModel
from nunchaku import NunchakuT5EncoderModel
from nunchaku.utils import get_precision
from nunchaku.caching.diffusers_adapters import apply_cache_on_pipe
from diffusers.utils import load_image

bfl_repo = "black-forest-labs/FLUX.1-dev"
dtype =  torch.bfloat16  # or torch.float16, or torch.float32
device = "cuda"  # or "cpu" if you want to run on CPU

scheduler       = FlowMatchEulerDiscreteScheduler.from_pretrained(bfl_repo, subfolder="scheduler", torch_dtype=dtype)
text_encoder    = CLIPTextModel.from_pretrained(bfl_repo, subfolder="text_encoder", torch_dtype=dtype)
text_encoder_2  = T5EncoderModel.from_pretrained(bfl_repo, subfolder="text_encoder_2", torch_dtype=dtype)
tokenizer       = CLIPTokenizer.from_pretrained(bfl_repo, subfolder="tokenizer", torch_dtype=dtype, clean_up_tokenization_spaces=True)
tokenizer_2     = T5TokenizerFast.from_pretrained(bfl_repo, subfolder="tokenizer_2", torch_dtype=dtype, clean_up_tokenization_spaces=True)
vae             = AutoencoderKL.from_pretrained(bfl_repo, subfolder="vae", torch_dtype=dtype)
precision = get_precision()
transformer = NunchakuFluxTransformer2dModel.from_pretrained(
    f"mit-han-lab/nunchaku-flux.1-dev/svdq-{precision}_r32-flux.1-dev.safetensors",
    # offload=True
)
transformer.set_attention_impl("nunchaku-fp16")

# qencoder
text_encoder_2 = NunchakuT5EncoderModel.from_pretrained("mit-han-lab/nunchaku-t5/awq-int4-flux.1-t5xxl.safetensors")
controlnet_union = FluxControlNetModel.from_pretrained("Shakker-Labs/FLUX.1-dev-ControlNet-Union-Pro-2.0", torch_dtype=torch.bfloat16)
controlnet = FluxMultiControlNetModel([controlnet_union])  # we always recommend loading via FluxMultiControlNetModel


params = {
    "scheduler": scheduler,
    "vae": vae,
    "tokenizer": tokenizer,
    "tokenizer_2": tokenizer_2,
    "text_encoder": text_encoder,
    "text_encoder_2": text_encoder_2,
    "transformer": transformer
}
# pipe
pipe = FluxPipeline(**params).to(device, dtype=dtype)
pipe_cn = FluxControlNetPipeline(**params, controlnet=controlnet).to(device, dtype)

# offload
pipe.enable_sequential_cpu_offload(device=device)
pipe_cn.enable_sequential_cpu_offload(device=device)

# cache
apply_cache_on_pipe(
    pipe_cn,
    use_double_fb_cache=True,
    residual_diff_threshold_multi=0.09,
    residual_diff_threshold_single=0.12,
)
apply_cache_on_pipe(
     pipe_cn,
     use_double_fb_cache=False,
     residual_diff_threshold_multi=0,
     residual_diff_threshold_single=0,
 )

params = {
    "prompt": "A bohemian-style female travel blogger with sun-kissed skin and messy beach waves.",
    "height": 1152,
    "width": 768,
    "num_inference_steps": 30,
    "guidance_scale": 3.5,
}

# pipe
txt2img_res = pipe(
    **params,
).images[0]
txt2img_res.save("flux.1-dev-txt2img.jpg")

# cache
apply_cache_on_pipe(
    pipe_cn,
    use_double_fb_cache=True,
    residual_diff_threshold_multi=0.09,
    residual_diff_threshold_single=0.12,
)

# pipe_cn
control_iamge = load_image("https://huggingface.co/Shakker-Labs/FLUX.1-dev-ControlNet-Union-Pro/resolve/main/assets/openpose.jpg")
params["control_image"] = [control_iamge]
params["controlnet_conditioning_scale"] = [0.9]
params['control_guidance_end'] = [0.65]
cn_res = pipe_cn(
    **params,
).images[0]
cn_res.save("flux.1-dev-cn-txt2img.jpg")

[todochenxi]

I solved this issue. @todochenxi #452

You can test it with the code below.

import torch
from diffusers import (
FluxPipeline,
FluxInpaintPipeline,
FluxTransformer2DModel,
FlowMatchEulerDiscreteScheduler,
AutoencoderKL,
FluxControlNetModel,
FluxControlNetPipeline,
FluxFillPipeline
)
from diffusers.models import FluxMultiControlNetModel
from transformers import T5EncoderModel, CLIPTextModel, CLIPTokenizer,T5TokenizerFast
from nunchaku.models.transformers.transformer_flux import NunchakuFluxTransformer2dModel
from nunchaku import NunchakuT5EncoderModel
from nunchaku.utils import get_precision
from nunchaku.caching.diffusers_adapters import apply_cache_on_pipe
from diffusers.utils import load_image

bfl_repo = "black-forest-labs/FLUX.1-dev"
dtype = torch.bfloat16 # or torch.float16, or torch.float32
device = "cuda" # or "cpu" if you want to run on CPU

scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(bfl_repo, subfolder="scheduler", torch_dtype=dtype)
text_encoder = CLIPTextModel.from_pretrained(bfl_repo, subfolder="text_encoder", torch_dtype=dtype)
text_encoder_2 = T5EncoderModel.from_pretrained(bfl_repo, subfolder="text_encoder_2", torch_dtype=dtype)
tokenizer = CLIPTokenizer.from_pretrained(bfl_repo, subfolder="tokenizer", torch_dtype=dtype, clean_up_tokenization_spaces=True)
tokenizer_2 = T5TokenizerFast.from_pretrained(bfl_repo, subfolder="tokenizer_2", torch_dtype=dtype, clean_up_tokenization_spaces=True)
vae = AutoencoderKL.from_pretrained(bfl_repo, subfolder="vae", torch_dtype=dtype)
precision = get_precision()
transformer = NunchakuFluxTransformer2dModel.from_pretrained(
f"mit-han-lab/nunchaku-flux.1-dev/svdq-{precision}_r32-flux.1-dev.safetensors",
# offload=True
)
transformer.set_attention_impl("nunchaku-fp16")

qencoder

text_encoder_2 = NunchakuT5EncoderModel.from_pretrained("mit-han-lab/nunchaku-t5/awq-int4-flux.1-t5xxl.safetensors")
controlnet_union = FluxControlNetModel.from_pretrained("Shakker-Labs/FLUX.1-dev-ControlNet-Union-Pro-2.0", torch_dtype=torch.bfloat16)
controlnet = FluxMultiControlNetModel([controlnet_union]) # we always recommend loading via FluxMultiControlNetModel

params = {
"scheduler": scheduler,
"vae": vae,
"tokenizer": tokenizer,
"tokenizer_2": tokenizer_2,
"text_encoder": text_encoder,
"text_encoder_2": text_encoder_2,
"transformer": transformer
}

pipe

pipe = FluxPipeline(**params).to(device, dtype=dtype)
pipe_cn = FluxControlNetPipeline(**params, controlnet=controlnet).to(device, dtype)

offload

pipe.enable_sequential_cpu_offload(device=device)
pipe_cn.enable_sequential_cpu_offload(device=device)

cache

apply_cache_on_pipe(
pipe_cn,
use_double_fb_cache=True,
residual_diff_threshold_multi=0.09,
residual_diff_threshold_single=0.12,
)
apply_cache_on_pipe(
pipe_cn,
use_double_fb_cache=False,
residual_diff_threshold_multi=0,
residual_diff_threshold_single=0,
)

params = {
"prompt": "A bohemian-style female travel blogger with sun-kissed skin and messy beach waves.",
"height": 1152,
"width": 768,
"num_inference_steps": 30,
"guidance_scale": 3.5,
}

pipe

txt2img_res = pipe(
**params,
).images[0]
txt2img_res.save("flux.1-dev-txt2img.jpg")

cache

apply_cache_on_pipe(
pipe_cn,
use_double_fb_cache=True,
residual_diff_threshold_multi=0.09,
residual_diff_threshold_single=0.12,
)

pipe_cn

control_iamge = load_image("https://huggingface.co/Shakker-Labs/FLUX.1-dev-ControlNet-Union-Pro/resolve/main/assets/openpose.jpg")
params["control_image"] = [control_iamge]
params["controlnet_conditioning_scale"] = [0.9]
params['control_guidance_end'] = [0.65]
cn_res = pipe_cn(
**params,
).images[0]
cn_res.save("flux.1-dev-cn-txt2img.jpg")

Thank you for your help. I found that the error still exists when I tested it. The platform I used is 3090. Due to my negligence, apply cache for pipe_cn was used twice in the reproducible code. In fact, the first time was used for pipe.

[lmxyy]

fixed in #452