feat: introduce cache-dit to nunchaku (#763)

Co-authored-by: yanjun.qiu <yanjun.qiu@vipshop.com>

Author: DefTruth

examples/v1/flux.1-dev-cache-dit.py

@@ -0,0 +1,31 @@
+import cache_dit
+import torch
+from cache_dit import DBCacheConfig
+from diffusers import FluxPipeline
+
+from nunchaku.models.transformers.transformer_flux_v2 import NunchakuFluxTransformer2DModelV2
+from nunchaku.utils import get_precision
+
+precision = get_precision()  # auto-detect your precision is 'int4' or 'fp4' based on your GPU
+transformer = NunchakuFluxTransformer2DModelV2.from_pretrained(
+    f"nunchaku-tech/nunchaku-flux.1-dev/svdq-{precision}_r32-flux.1-dev.safetensors"
+)
+pipeline = FluxPipeline.from_pretrained(
+    "black-forest-labs/FLUX.1-dev", transformer=transformer, torch_dtype=torch.bfloat16
+).to("cuda")
+
+
+# Please check https://github.com/vipshop/cache-dit for more details about the parameters.
+cache_dit.enable_cache(
+    pipeline,
+    cache_config=DBCacheConfig(
+        Fn_compute_blocks=1,
+        Bn_compute_blocks=0,
+        residual_diff_threshold=0.12,
+    ),
+)
+
+image = pipeline("A cat holding a sign that says hello world", num_inference_steps=50, guidance_scale=3.5).images[0]
+image.save(f"flux.1-dev-cache-dit-{precision}.png")
+
+cache_dit.summary(pipeline)

examples/v1/qwen-image-cache-dit.py

@@ -0,0 +1,59 @@
+import cache_dit
+import torch
+from cache_dit import DBCacheConfig
+from diffusers import QwenImagePipeline
+
+from nunchaku.models.transformers.transformer_qwenimage import NunchakuQwenImageTransformer2DModel
+from nunchaku.utils import get_gpu_memory, get_precision
+
+rank = 32  # you can also use rank=128 model to improve the quality
+
+# Load the model
+transformer = NunchakuQwenImageTransformer2DModel.from_pretrained(
+    f"nunchaku-tech/nunchaku-qwen-image/svdq-{get_precision()}_r{rank}-qwen-image.safetensors"
+)
+
+# currently, you need to use this pipeline to offload the model to CPU
+pipe = QwenImagePipeline.from_pretrained("Qwen/Qwen-Image", transformer=transformer, torch_dtype=torch.bfloat16)
+
+# Please check https://github.com/vipshop/cache-dit for more details about the parameters.
+
+cache_dit.enable_cache(
+    pipe,
+    cache_config=DBCacheConfig(
+        Fn_compute_blocks=8,
+        Bn_compute_blocks=0,
+        residual_diff_threshold=0.12,
+    ),
+)
+
+
+if get_gpu_memory() > 18:
+    pipe.enable_model_cpu_offload()
+else:
+    # use per-layer offloading for low VRAM. This only requires 3-4GB of VRAM.
+    transformer.set_offload(
+        True, use_pin_memory=False, num_blocks_on_gpu=1
+    )  # increase num_blocks_on_gpu if you have more VRAM
+    pipe._exclude_from_cpu_offload.append("transformer")
+    pipe.enable_sequential_cpu_offload()
+
+positive_magic = {
+    "en": "Ultra HD, 4K, cinematic composition.",  # for english prompt,
+    "zh": "超清，4K，电影级构图",  # for chinese prompt,
+}
+
+# Generate image
+prompt = """Bookstore window display. A sign displays “New Arrivals This Week”. Below, a shelf tag with the text “Best-Selling Novels Here”. To the side, a colorful poster advertises “Author Meet And Greet on Saturday” with a central portrait of the author. There are four books on the bookshelf, namely “The light between worlds” “When stars are scattered” “The slient patient” “The night circus”"""
+negative_prompt = " "  # using an empty string if you do not have specific concept to remove
+
+image = pipe(
+    prompt=prompt + positive_magic["en"],
+    negative_prompt=negative_prompt,
+    width=1664,
+    height=928,
+    num_inference_steps=50,
+    true_cfg_scale=4.0,
+).images[0]
+
+image.save(f"qwen-image-cache-dit-r{rank}.png")