cache-dit 0.3.2__py3-none-any.whl → 1.0.14__py3-none-any.whl

cache_dit/caching/patch_functors/functor_qwen_image_controlnet.py

@@ -0,0 +1,263 @@
+import torch
+import numpy as np
+from typing import Tuple, Optional, Dict, Any, Union, List
+from diffusers import QwenImageTransformer2DModel
+from diffusers.models.transformers.transformer_qwenimage import (
+    QwenImageTransformerBlock,
+    Transformer2DModelOutput,
+)
+from diffusers.utils import (
+    USE_PEFT_BACKEND,
+    scale_lora_layers,
+    unscale_lora_layers,
+)
+from cache_dit.caching.patch_functors.functor_base import (
+    PatchFunctor,
+)
+from cache_dit.logger import init_logger
+
+logger = init_logger(__name__)
+
+
+class QwenImageControlNetPatchFunctor(PatchFunctor):
+
+    def apply(
+        self,
+        transformer: QwenImageTransformer2DModel,
+        **kwargs,
+    ) -> QwenImageTransformer2DModel:
+        if hasattr(transformer, "_is_patched"):
+            return transformer
+
+        is_patched = False
+
+        _index_block = 0
+        _num_blocks = len(transformer.transformer_blocks)
+        for block in transformer.transformer_blocks:
+            assert isinstance(block, QwenImageTransformerBlock)
+            block._index_block = _index_block
+            block._num_blocks = _num_blocks
+            block.forward = __patch_block_forward__.__get__(block)
+            _index_block += 1
+
+        is_patched = True
+        cls_name = transformer.__class__.__name__
+
+        if is_patched:
+            logger.warning(f"Patched {cls_name} for cache-dit.")
+            assert not getattr(transformer, "_is_parallelized", False), (
+                "Please call `cache_dit.enable_cache` before Parallelize, "
+                "the __patch_transformer_forward__ will overwrite the "
+                "parallized forward and cause a downgrade of performance."
+            )
+            transformer.forward = __patch_transformer_forward__.__get__(
+                transformer
+            )
+
+        transformer._is_patched = is_patched  # True or False
+
+        logger.info(
+            f"Applied {self.__class__.__name__} for {cls_name}, "
+            f"Patch: {is_patched}."
+        )
+
+        return transformer
+
+
+def __patch_block_forward__(
+    self: QwenImageTransformerBlock,
+    hidden_states: torch.Tensor,
+    encoder_hidden_states: torch.Tensor,
+    encoder_hidden_states_mask: torch.Tensor,
+    temb: torch.Tensor,
+    image_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
+    joint_attention_kwargs: Optional[Dict[str, Any]] = None,
+    controlnet_block_samples: Optional[List[torch.Tensor]] = None,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    # Get modulation parameters for both streams
+    img_mod_params = self.img_mod(temb)  # [B, 6*dim]
+    txt_mod_params = self.txt_mod(temb)  # [B, 6*dim]
+
+    # Split modulation parameters for norm1 and norm2
+    img_mod1, img_mod2 = img_mod_params.chunk(2, dim=-1)  # Each [B, 3*dim]
+    txt_mod1, txt_mod2 = txt_mod_params.chunk(2, dim=-1)  # Each [B, 3*dim]
+
+    # Process image stream - norm1 + modulation
+    img_normed = self.img_norm1(hidden_states)
+    img_modulated, img_gate1 = self._modulate(img_normed, img_mod1)
+
+    # Process text stream - norm1 + modulation
+    txt_normed = self.txt_norm1(encoder_hidden_states)
+    txt_modulated, txt_gate1 = self._modulate(txt_normed, txt_mod1)
+
+    # Use QwenAttnProcessor2_0 for joint attention computation
+    # This directly implements the DoubleStreamLayerMegatron logic:
+    # 1. Computes QKV for both streams
+    # 2. Applies QK normalization and RoPE
+    # 3. Concatenates and runs joint attention
+    # 4. Splits results back to separate streams
+    joint_attention_kwargs = joint_attention_kwargs or {}
+    attn_output = self.attn(
+        hidden_states=img_modulated,  # Image stream (will be processed as "sample")
+        encoder_hidden_states=txt_modulated,  # Text stream (will be processed as "context")
+        encoder_hidden_states_mask=encoder_hidden_states_mask,
+        image_rotary_emb=image_rotary_emb,
+        **joint_attention_kwargs,
+    )
+
+    # QwenAttnProcessor2_0 returns (img_output, txt_output) when encoder_hidden_states is provided
+    img_attn_output, txt_attn_output = attn_output
+
+    # Apply attention gates and add residual (like in Megatron)
+    hidden_states = hidden_states + img_gate1 * img_attn_output
+    encoder_hidden_states = encoder_hidden_states + txt_gate1 * txt_attn_output
+
+    # Process image stream - norm2 + MLP
+    img_normed2 = self.img_norm2(hidden_states)
+    img_modulated2, img_gate2 = self._modulate(img_normed2, img_mod2)
+    img_mlp_output = self.img_mlp(img_modulated2)
+    hidden_states = hidden_states + img_gate2 * img_mlp_output
+
+    # Process text stream - norm2 + MLP
+    txt_normed2 = self.txt_norm2(encoder_hidden_states)
+    txt_modulated2, txt_gate2 = self._modulate(txt_normed2, txt_mod2)
+    txt_mlp_output = self.txt_mlp(txt_modulated2)
+    encoder_hidden_states = encoder_hidden_states + txt_gate2 * txt_mlp_output
+
+    # Clip to prevent overflow for fp16
+    if encoder_hidden_states.dtype == torch.float16:
+        encoder_hidden_states = encoder_hidden_states.clip(-65504, 65504)
+    if hidden_states.dtype == torch.float16:
+        hidden_states = hidden_states.clip(-65504, 65504)
+
+    if controlnet_block_samples is not None:
+        # Add ControlNet conditioning
+        num_blocks = self._num_blocks
+        index_block = self._index_block
+        interval_control = num_blocks / len(controlnet_block_samples)
+        interval_control = int(np.ceil(interval_control))
+        hidden_states = (
+            hidden_states
+            + controlnet_block_samples[index_block // interval_control]
+        )
+
+    return encoder_hidden_states, hidden_states
+
+
+def __patch_transformer_forward__(
+    self: QwenImageTransformer2DModel,
+    hidden_states: torch.Tensor,
+    encoder_hidden_states: torch.Tensor = None,
+    encoder_hidden_states_mask: torch.Tensor = None,
+    timestep: torch.LongTensor = None,
+    img_shapes: Optional[List[Tuple[int, int, int]]] = None,
+    txt_seq_lens: Optional[List[int]] = None,
+    guidance: torch.Tensor = None,  # TODO: this should probably be removed
+    attention_kwargs: Optional[Dict[str, Any]] = None,
+    controlnet_block_samples=None,
+    return_dict: bool = True,
+) -> Union[torch.Tensor, Transformer2DModelOutput]:
+    """
+    The [`QwenTransformer2DModel`] forward method.
+
+    Args:
+        hidden_states (`torch.Tensor` of shape `(batch_size, image_sequence_length, in_channels)`):
+            Input `hidden_states`.
+        encoder_hidden_states (`torch.Tensor` of shape `(batch_size, text_sequence_length, joint_attention_dim)`):
+            Conditional embeddings (embeddings computed from the input conditions such as prompts) to use.
+        encoder_hidden_states_mask (`torch.Tensor` of shape `(batch_size, text_sequence_length)`):
+            Mask of the input conditions.
+        timestep ( `torch.LongTensor`):
+            Used to indicate denoising step.
+        attention_kwargs (`dict`, *optional*):
+            A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under
+            `self.processor` in
+            [diffusers.models.attention_processor](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py).
+        return_dict (`bool`, *optional*, defaults to `True`):
+            Whether or not to return a [`~models.transformer_2d.Transformer2DModelOutput`] instead of a plain
+            tuple.
+
+    Returns:
+        If `return_dict` is True, an [`~models.transformer_2d.Transformer2DModelOutput`] is returned, otherwise a
+        `tuple` where the first element is the sample tensor.
+    """
+    if attention_kwargs is not None:
+        attention_kwargs = attention_kwargs.copy()
+        lora_scale = attention_kwargs.pop("scale", 1.0)
+    else:
+        lora_scale = 1.0
+
+    if USE_PEFT_BACKEND:
+        # weight the lora layers by setting `lora_scale` for each PEFT layer
+        scale_lora_layers(self, lora_scale)
+    else:
+        if (
+            attention_kwargs is not None
+            and attention_kwargs.get("scale", None) is not None
+        ):
+            logger.warning(
+                "Passing `scale` via `joint_attention_kwargs` when not using the PEFT backend is ineffective."
+            )
+
+    hidden_states = self.img_in(hidden_states)
+
+    timestep = timestep.to(hidden_states.dtype)
+    encoder_hidden_states = self.txt_norm(encoder_hidden_states)
+    encoder_hidden_states = self.txt_in(encoder_hidden_states)
+
+    if guidance is not None:
+        guidance = guidance.to(hidden_states.dtype) * 1000
+
+    temb = (
+        self.time_text_embed(timestep, hidden_states)
+        if guidance is None
+        else self.time_text_embed(timestep, guidance, hidden_states)
+    )
+
+    image_rotary_emb = self.pos_embed(
+        img_shapes, txt_seq_lens, device=hidden_states.device
+    )
+
+    for index_block, block in enumerate(self.transformer_blocks):
+        if torch.is_grad_enabled() and self.gradient_checkpointing:
+            encoder_hidden_states, hidden_states = (
+                self._gradient_checkpointing_func(
+                    block,
+                    hidden_states,
+                    encoder_hidden_states,
+                    encoder_hidden_states_mask,
+                    temb,
+                    image_rotary_emb,
+                    controlnet_block_samples,
+                )
+            )
+
+        else:
+            encoder_hidden_states, hidden_states = block(
+                hidden_states=hidden_states,
+                encoder_hidden_states=encoder_hidden_states,
+                encoder_hidden_states_mask=encoder_hidden_states_mask,
+                temb=temb,
+                image_rotary_emb=image_rotary_emb,
+                controlnet_block_samples=controlnet_block_samples,
+                joint_attention_kwargs=attention_kwargs,
+            )
+
+        # # controlnet residual
+        # if controlnet_block_samples is not None:
+        #     interval_control = len(self.transformer_blocks) / len(controlnet_block_samples)
+        #     interval_control = int(np.ceil(interval_control))
+        #     hidden_states = hidden_states + controlnet_block_samples[index_block // interval_control]
+
+    # Use only the image part (hidden_states) from the dual-stream blocks
+    hidden_states = self.norm_out(hidden_states, temb)
+    output = self.proj_out(hidden_states)
+
+    if USE_PEFT_BACKEND:
+        # remove `lora_scale` from each PEFT layer
+        unscale_lora_layers(self, lora_scale)
+
+    if not return_dict:
+        return (output,)
+
+    return Transformer2DModelOutput(sample=output)

cache_dit/caching/utils.py

@@ -0,0 +1,68 @@
+import yaml
+
+
+def load_cache_options_from_yaml(yaml_file_path):
+    try:
+        with open(yaml_file_path, "r") as f:
+            kwargs: dict = yaml.safe_load(f)
+
+        required_keys = [
+            "residual_diff_threshold",
+        ]
+        for key in required_keys:
+            if key not in kwargs:
+                raise ValueError(
+                    f"Configuration file missing required item: {key}"
+                )
+
+        cache_context_kwargs = {}
+        if kwargs.get("enable_taylorseer", False):
+            from cache_dit.caching.cache_contexts.calibrators import (
+                TaylorSeerCalibratorConfig,
+            )
+
+            cache_context_kwargs["calibrator_config"] = (
+                TaylorSeerCalibratorConfig(
+                    enable_calibrator=kwargs.pop("enable_taylorseer"),
+                    enable_encoder_calibrator=kwargs.pop(
+                        "enable_encoder_taylorseer", False
+                    ),
+                    calibrator_cache_type=kwargs.pop(
+                        "taylorseer_cache_type", "residual"
+                    ),
+                    taylorseer_order=kwargs.pop("taylorseer_order", 1),
+                )
+            )
+
+        if "cache_type" not in kwargs:
+            from cache_dit.caching.cache_contexts import BasicCacheConfig
+
+            cache_context_kwargs["cache_config"] = BasicCacheConfig()
+            cache_context_kwargs["cache_config"].update(**kwargs)
+        else:
+            cache_type = kwargs.pop("cache_type")
+            if cache_type == "DBCache":
+                from cache_dit.caching.cache_contexts import DBCacheConfig
+
+                cache_context_kwargs["cache_config"] = DBCacheConfig()
+                cache_context_kwargs["cache_config"].update(**kwargs)
+            elif cache_type == "DBPrune":
+                from cache_dit.caching.cache_contexts import DBPruneConfig
+
+                cache_context_kwargs["cache_config"] = DBPruneConfig()
+                cache_context_kwargs["cache_config"].update(**kwargs)
+            else:
+                raise ValueError(f"Unsupported cache_type: {cache_type}.")
+
+        return cache_context_kwargs
+
+    except FileNotFoundError:
+        raise FileNotFoundError(
+            f"Configuration file not found: {yaml_file_path}"
+        )
+    except yaml.YAMLError as e:
+        raise yaml.YAMLError(f"YAML file parsing error: {str(e)}")
+
+
+def load_options(path: str):
+    return load_cache_options_from_yaml(path)

cache_dit/metrics/__init__.py

@@ -1,3 +1,14 @@
+try:
+    import ImageReward
+    import lpips
+    import skimage
+    import scipy
+except ImportError:
+    raise ImportError(
+        "Metrics functionality requires the 'metrics' extra dependencies. "
+        "Install with:\npip install cache-dit[metrics]"
+    )
+
 from cache_dit.metrics.metrics import compute_psnr
 from cache_dit.metrics.metrics import compute_ssim
 from cache_dit.metrics.metrics import compute_mse

cache_dit/metrics/metrics.py

@@ -646,6 +646,7 @@ def entrypoint():
                     not os.path.exists(img_test),
                 )
             ):
+                logger.error(f"Not exist: {img_true} or {img_test}, skip.")
                 return
             # img_true and img_test can be files or dirs
             img_true_info = os.path.basename(img_true)
@@ -684,6 +685,7 @@ def entrypoint():
                     not os.path.exists(img_test),  # dir
                 )
             ):
+                logger.error(f"Not exist: {prompt_true} or {img_test}, skip.")
                 return
 
             # img_true and img_test can be files or dirs
@@ -714,6 +716,7 @@ def entrypoint():
                     not os.path.exists(video_test),
                 )
             ):
+                logger.error(f"Not exist: {video_true} or {video_test}, skip.")
                 return
 
             # video_true and video_test can be files or dirs

cache_dit/parallelism/__init__.py

@@ -0,0 +1,3 @@
+from cache_dit.parallelism.parallel_backend import ParallelismBackend
+from cache_dit.parallelism.parallel_config import ParallelismConfig
+from cache_dit.parallelism.parallel_interface import enable_parallelism

cache_dit/parallelism/backends/native_diffusers/__init__.py

@@ -0,0 +1,6 @@
+from cache_dit.parallelism.backends.native_diffusers.context_parallelism import (
+    ContextParallelismPlannerRegister,
+)
+from cache_dit.parallelism.backends.native_diffusers.parallel_difffusers import (
+    maybe_enable_parallelism,
+)

cache_dit/parallelism/backends/native_diffusers/context_parallelism/__init__.py

@@ -0,0 +1,164 @@
+import torch
+from typing import Optional
+
+from diffusers.models.modeling_utils import ModelMixin
+from cache_dit.parallelism.parallel_backend import ParallelismBackend
+from cache_dit.parallelism.parallel_config import ParallelismConfig
+from cache_dit.logger import init_logger
+from ..utils import (
+    native_diffusers_parallelism_available,
+    ContextParallelConfig,
+)
+from .attention import maybe_resigter_native_attention_backend
+from .cp_planners import *
+
+try:
+    maybe_resigter_native_attention_backend()
+except ImportError as e:
+    raise ImportError(e)
+
+logger = init_logger(__name__)
+
+
+def maybe_enable_context_parallelism(
+    transformer: torch.nn.Module,
+    parallelism_config: Optional[ParallelismConfig],
+) -> torch.nn.Module:
+    assert isinstance(transformer, ModelMixin), (
+        "transformer must be an instance of diffusers' ModelMixin, "
+        f"but got {type(transformer)}"
+    )
+    if parallelism_config is None:
+        return transformer
+
+    assert isinstance(parallelism_config, ParallelismConfig), (
+        "parallelism_config must be an instance of ParallelismConfig"
+        f" but got {type(parallelism_config)}"
+    )
+
+    if (
+        parallelism_config.backend == ParallelismBackend.NATIVE_DIFFUSER
+        and native_diffusers_parallelism_available()
+    ):
+        cp_config = None
+        if (
+            parallelism_config.ulysses_size is not None
+            or parallelism_config.ring_size is not None
+        ):
+            cp_config = ContextParallelConfig(
+                ulysses_degree=parallelism_config.ulysses_size,
+                ring_degree=parallelism_config.ring_size,
+            )
+        if cp_config is not None:
+            attention_backend = parallelism_config.parallel_kwargs.get(
+                "attention_backend", None
+            )
+            if hasattr(transformer, "enable_parallelism"):
+                if hasattr(transformer, "set_attention_backend"):
+                    # native, _native_cudnn, flash, etc.
+                    if attention_backend is None:
+                        # Now only _native_cudnn is supported for parallelism
+                        # issue: https://github.com/huggingface/diffusers/pull/12443
+                        transformer.set_attention_backend("_native_cudnn")
+                        logger.warning(
+                            "attention_backend is None, set default attention backend "
+                            "to _native_cudnn for parallelism because of the issue: "
+                            "https://github.com/huggingface/diffusers/pull/12443"
+                        )
+                    else:
+                        transformer.set_attention_backend(attention_backend)
+                        logger.info(
+                            "Found attention_backend from config, set attention "
+                            f"backend to: {attention_backend}"
+                        )
+                # Prefer custom cp_plan if provided
+                cp_plan = parallelism_config.parallel_kwargs.get(
+                    "cp_plan", None
+                )
+                if cp_plan is not None:
+                    logger.info(
+                        f"Using custom context parallelism plan: {cp_plan}"
+                    )
+                else:
+                    # Try get context parallelism plan from register if not provided
+                    extra_parallel_kwargs = {}
+                    if parallelism_config.parallel_kwargs is not None:
+                        extra_parallel_kwargs = (
+                            parallelism_config.parallel_kwargs
+                        )
+                    cp_plan = ContextParallelismPlannerRegister.get_planner(
+                        transformer
+                    )().apply(transformer=transformer, **extra_parallel_kwargs)
+
+                transformer.enable_parallelism(
+                    config=cp_config, cp_plan=cp_plan
+                )
+                _maybe_patch_native_parallel_config(transformer)
+            else:
+                raise ValueError(
+                    f"{transformer.__class__.__name__} does not support context parallelism."
+                )
+
+    return transformer
+
+
+def _maybe_patch_native_parallel_config(
+    transformer: torch.nn.Module,
+) -> torch.nn.Module:
+
+    cls_name = transformer.__class__.__name__
+    if not cls_name.startswith("Nunchaku"):
+        return transformer
+
+    from diffusers import FluxTransformer2DModel, QwenImageTransformer2DModel
+
+    try:
+        from nunchaku.models.transformers.transformer_flux_v2 import (
+            NunchakuFluxTransformer2DModelV2,
+            NunchakuFluxAttention,
+            NunchakuFluxFA2Processor,
+        )
+        from nunchaku.models.transformers.transformer_qwenimage import (
+            NunchakuQwenAttention,
+            NunchakuQwenImageNaiveFA2Processor,
+            NunchakuQwenImageTransformer2DModel,
+        )
+    except ImportError:
+        raise ImportError(
+            "NunchakuFluxTransformer2DModelV2 or NunchakuQwenImageTransformer2DModel "
+            "requires the 'nunchaku' package. Please install nunchaku before using "
+            "the context parallelism for nunchaku 4-bits models."
+        )
+    assert isinstance(
+        transformer,
+        (
+            NunchakuFluxTransformer2DModelV2,
+            FluxTransformer2DModel,
+        ),
+    ) or isinstance(
+        transformer,
+        (
+            NunchakuQwenImageTransformer2DModel,
+            QwenImageTransformer2DModel,
+        ),
+    ), (
+        "transformer must be an instance of NunchakuFluxTransformer2DModelV2 "
+        f"or NunchakuQwenImageTransformer2DModel, but got {type(transformer)}"
+    )
+    config = transformer._parallel_config
+
+    attention_classes = (
+        NunchakuFluxAttention,
+        NunchakuFluxFA2Processor,
+        NunchakuQwenAttention,
+        NunchakuQwenImageNaiveFA2Processor,
+    )
+    for module in transformer.modules():
+        if not isinstance(module, attention_classes):
+            continue
+        processor = getattr(module, "processor", None)
+        if processor is None or not hasattr(processor, "_parallel_config"):
+            continue
+        processor._parallel_config = config
+
+    return transformer

cache_dit/parallelism/backends/native_diffusers/context_parallelism/attention/__init__.py

@@ -0,0 +1,4 @@
+def maybe_resigter_native_attention_backend():
+    """Maybe re-register native attention backend to enable context parallelism."""
+    # Import custom attention backend ensuring registration
+    from ._attention_dispatch import _native_attention