cache-dit 0.2.24__py3-none-any.whl → 0.2.26__py3-none-any.whl

cache_dit/cache_factory/patch_functors/functor_chroma.py

@@ -0,0 +1,273 @@
+import inspect
+
+import torch
+import numpy as np
+from typing import Tuple, Optional, Dict, Any, Union
+from diffusers import ChromaTransformer2DModel
+from diffusers.models.transformers.transformer_chroma import (
+    ChromaSingleTransformerBlock,
+    Transformer2DModelOutput,
+)
+from diffusers.utils import (
+    USE_PEFT_BACKEND,
+    scale_lora_layers,
+    unscale_lora_layers,
+)
+
+from cache_dit.cache_factory.patch_functors.functor_base import (
+    PatchFunctor,
+)
+from cache_dit.logger import init_logger
+
+logger = init_logger(__name__)
+
+
+class ChromaPatchFunctor(PatchFunctor):
+
+    def apply(
+        self,
+        transformer: ChromaTransformer2DModel,
+        blocks: torch.nn.ModuleList = None,
+        **kwargs,
+    ) -> ChromaTransformer2DModel:
+        if blocks is None:
+            blocks = transformer.single_transformer_blocks
+
+        is_patched = False
+        for block in blocks:
+            if isinstance(block, ChromaSingleTransformerBlock):
+                forward_parameters = inspect.signature(
+                    block.forward
+                ).parameters.keys()
+                if "encoder_hidden_states" not in forward_parameters:
+                    block.forward = __patch_single_forward__.__get__(block)
+                    is_patched = True
+
+        if is_patched:
+            logger.warning("Patched Chroma 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 = True
+
+        cls_name = transformer.__class__.__name__
+        logger.info(
+            f"Applied {self.__class__.__name__} for {cls_name}, "
+            f"Patch: {is_patched}."
+        )
+
+        return transformer
+
+
+# adapted from: https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/transformers/transformer_chroma.py
+def __patch_single_forward__(
+    self: ChromaSingleTransformerBlock,  # Almost same as FluxSingleTransformerBlock
+    hidden_states: torch.Tensor,
+    encoder_hidden_states: torch.Tensor,
+    temb: torch.Tensor,
+    image_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
+    joint_attention_kwargs: Optional[Dict[str, Any]] = None,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    text_seq_len = encoder_hidden_states.shape[1]
+    hidden_states = torch.cat([encoder_hidden_states, hidden_states], dim=1)
+
+    residual = hidden_states
+    norm_hidden_states, gate = self.norm(hidden_states, emb=temb)
+    mlp_hidden_states = self.act_mlp(self.proj_mlp(norm_hidden_states))
+    joint_attention_kwargs = joint_attention_kwargs or {}
+    attn_output = self.attn(
+        hidden_states=norm_hidden_states,
+        image_rotary_emb=image_rotary_emb,
+        **joint_attention_kwargs,
+    )
+
+    hidden_states = torch.cat([attn_output, mlp_hidden_states], dim=2)
+    gate = gate.unsqueeze(1)
+    hidden_states = gate * self.proj_out(hidden_states)
+    hidden_states = residual + hidden_states
+    if hidden_states.dtype == torch.float16:
+        hidden_states = hidden_states.clip(-65504, 65504)
+
+    encoder_hidden_states, hidden_states = (
+        hidden_states[:, :text_seq_len],
+        hidden_states[:, text_seq_len:],
+    )
+    return encoder_hidden_states, hidden_states
+
+
+# Adapted from: https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/transformers/transformer_chroma.py
+def __patch_transformer_forward__(
+    self: ChromaTransformer2DModel,
+    hidden_states: torch.Tensor,
+    encoder_hidden_states: torch.Tensor = None,
+    timestep: torch.LongTensor = None,
+    img_ids: torch.Tensor = None,
+    txt_ids: torch.Tensor = None,
+    attention_mask: torch.Tensor = None,
+    joint_attention_kwargs: Optional[Dict[str, Any]] = None,
+    controlnet_block_samples=None,
+    controlnet_single_block_samples=None,
+    return_dict: bool = True,
+    controlnet_blocks_repeat: bool = False,
+) -> Union[torch.Tensor, Transformer2DModelOutput]:
+    if joint_attention_kwargs is not None:
+        joint_attention_kwargs = joint_attention_kwargs.copy()
+        lora_scale = joint_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 (
+            joint_attention_kwargs is not None
+            and joint_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.x_embedder(hidden_states)
+
+    timestep = timestep.to(hidden_states.dtype) * 1000
+
+    input_vec = self.time_text_embed(timestep)
+    pooled_temb = self.distilled_guidance_layer(input_vec)
+
+    encoder_hidden_states = self.context_embedder(encoder_hidden_states)
+
+    if txt_ids.ndim == 3:
+        logger.warning(
+            "Passing `txt_ids` 3d torch.Tensor is deprecated."
+            "Please remove the batch dimension and pass it as a 2d torch Tensor"
+        )
+        txt_ids = txt_ids[0]
+    if img_ids.ndim == 3:
+        logger.warning(
+            "Passing `img_ids` 3d torch.Tensor is deprecated."
+            "Please remove the batch dimension and pass it as a 2d torch Tensor"
+        )
+        img_ids = img_ids[0]
+
+    ids = torch.cat((txt_ids, img_ids), dim=0)
+    image_rotary_emb = self.pos_embed(ids)
+
+    if (
+        joint_attention_kwargs is not None
+        and "ip_adapter_image_embeds" in joint_attention_kwargs
+    ):
+        ip_adapter_image_embeds = joint_attention_kwargs.pop(
+            "ip_adapter_image_embeds"
+        )
+        ip_hidden_states = self.encoder_hid_proj(ip_adapter_image_embeds)
+        joint_attention_kwargs.update({"ip_hidden_states": ip_hidden_states})
+
+    for index_block, block in enumerate(self.transformer_blocks):
+        img_offset = 3 * len(self.single_transformer_blocks)
+        txt_offset = img_offset + 6 * len(self.transformer_blocks)
+        img_modulation = img_offset + 6 * index_block
+        text_modulation = txt_offset + 6 * index_block
+        temb = torch.cat(
+            (
+                pooled_temb[:, img_modulation : img_modulation + 6],
+                pooled_temb[:, text_modulation : text_modulation + 6],
+            ),
+            dim=1,
+        )
+        if torch.is_grad_enabled() and self.gradient_checkpointing:
+            encoder_hidden_states, hidden_states = (
+                self._gradient_checkpointing_func(
+                    block,
+                    hidden_states,
+                    encoder_hidden_states,
+                    temb,
+                    image_rotary_emb,
+                    attention_mask,
+                )
+            )
+
+        else:
+            encoder_hidden_states, hidden_states = block(
+                hidden_states=hidden_states,
+                encoder_hidden_states=encoder_hidden_states,
+                temb=temb,
+                image_rotary_emb=image_rotary_emb,
+                attention_mask=attention_mask,
+                joint_attention_kwargs=joint_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))
+            # For Xlabs ControlNet.
+            if controlnet_blocks_repeat:
+                hidden_states = (
+                    hidden_states
+                    + controlnet_block_samples[
+                        index_block % len(controlnet_block_samples)
+                    ]
+                )
+            else:
+                hidden_states = (
+                    hidden_states
+                    + controlnet_block_samples[index_block // interval_control]
+                )
+
+    for index_block, block in enumerate(self.single_transformer_blocks):
+        start_idx = 3 * index_block
+        temb = pooled_temb[:, start_idx : start_idx + 3]
+        if torch.is_grad_enabled() and self.gradient_checkpointing:
+            encoder_hidden_states, hidden_states = (
+                self._gradient_checkpointing_func(
+                    block,
+                    hidden_states,
+                    encoder_hidden_states,
+                    temb,
+                    image_rotary_emb,
+                )
+            )
+
+        else:
+            encoder_hidden_states, hidden_states = block(
+                hidden_states=hidden_states,
+                encoder_hidden_states=encoder_hidden_states,
+                temb=temb,
+                image_rotary_emb=image_rotary_emb,
+                attention_mask=attention_mask,
+                joint_attention_kwargs=joint_attention_kwargs,
+            )
+
+        # controlnet residual
+        if controlnet_single_block_samples is not None:
+            interval_control = len(self.single_transformer_blocks) / len(
+                controlnet_single_block_samples
+            )
+            interval_control = int(np.ceil(interval_control))
+            hidden_states = (
+                hidden_states
+                + controlnet_single_block_samples[
+                    index_block // interval_control
+                ]
+            )
+
+    temb = pooled_temb[:, -2:]
+    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/cache_factory/{patch/flux.py → patch_functors/functor_flux.py}

@@ -14,12 +14,56 @@ from diffusers.utils import (
     unscale_lora_layers,
 )
 
-
+from cache_dit.cache_factory.patch_functors.functor_base import (
+    PatchFunctor,
+)
 from cache_dit.logger import init_logger
 
 logger = init_logger(__name__)
 
 
+class FluxPatchFunctor(PatchFunctor):
+
+    def apply(
+        self,
+        transformer: FluxTransformer2DModel,
+        blocks: torch.nn.ModuleList = None,
+        **kwargs,
+    ) -> FluxTransformer2DModel:
+        if blocks is None:
+            blocks = transformer.single_transformer_blocks
+
+        is_patched = False
+        for block in blocks:
+            if isinstance(block, FluxSingleTransformerBlock):
+                forward_parameters = inspect.signature(
+                    block.forward
+                ).parameters.keys()
+                if "encoder_hidden_states" not in forward_parameters:
+                    block.forward = __patch_single_forward__.__get__(block)
+                    is_patched = True
+
+        if is_patched:
+            logger.warning("Patched Flux 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 = True
+
+        cls_name = transformer.__class__.__name__
+        logger.info(
+            f"Applied {self.__class__.__name__} for {cls_name}, "
+            f"Patch: {is_patched}."
+        )
+
+        return transformer
+
+
 # copy from: https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/transformers/transformer_flux.py#L380
 def __patch_single_forward__(
     self: FluxSingleTransformerBlock,
@@ -217,33 +261,3 @@ def __patch_transformer_forward__(
         return (output,)
 
     return Transformer2DModelOutput(sample=output)
-
-
-def maybe_patch_flux_transformer(
-    transformer: FluxTransformer2DModel,
-    blocks: torch.nn.ModuleList = None,
-) -> FluxTransformer2DModel:
-    if blocks is None:
-        blocks = transformer.single_transformer_blocks
-
-    is_patched = False
-    for block in blocks:
-        if isinstance(block, FluxSingleTransformerBlock):
-            forward_parameters = inspect.signature(
-                block.forward
-            ).parameters.keys()
-            if "encoder_hidden_states" not in forward_parameters:
-                block.forward = __patch_single_forward__.__get__(block)
-                is_patched = True
-
-    if is_patched:
-        logger.warning("Patched Flux for cache-dit.")
-        assert not getattr(transformer, "_is_parallelized", False), (
-            "Please call apply_cache_on_pipe 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 = True
-
-    return transformer

cache_dit/compile/utils.py

@@ -24,7 +24,7 @@ def epilogue_prologue_fusion_enabled(**kwargs) -> bool:
 
 
 def set_compile_configs(
-    descent_tuning: bool = True,
+    descent_tuning: bool = False,
     cuda_graphs: bool = False,
     force_disable_compile_caches: bool = False,
     use_fast_math: bool = False,

cache_dit/quantize/__init__.py

@@ -0,0 +1 @@
+from cache_dit.quantize.quantize_interface import quantize

cache_dit/quantize/quantize_ao.py

@@ -0,0 +1,196 @@
+import gc
+import time
+import torch
+from typing import Callable, Optional, List
+from cache_dit.logger import init_logger
+
+logger = init_logger(__name__)
+
+
+def quantize_ao(
+    module: torch.nn.Module,
+    quant_type: str = "fp8_w8a8_dq",
+    exclude_layers: List[str] = [
+        "embedder",
+        "embed",
+    ],
+    filter_fn: Optional[Callable] = None,
+    # paramters for fp8 quantization
+    per_row: bool = True,
+    **kwargs,
+) -> torch.nn.Module:
+    # Apply FP8 DQ for module and skip any `embed` modules
+    # by default to avoid non-trivial precision downgrade. Please
+    # set `exclude_layers` as `[]` if you don't want this behavior.
+    assert isinstance(module, torch.nn.Module)
+
+    quant_type = quant_type.lower()
+    assert quant_type in (
+        "fp8_w8a8_dq",
+        "fp8_w8a16_wo",
+        "int8_w8a8_dq",
+        "int8_w8a16_wo",
+        "int4_w4a8_dq",
+        "int4_w4a4_dq",
+        "int4_w4a16_wo",
+    ), f"{quant_type} is not supported for torchao backend now!"
+
+    if "fp8" in quant_type:
+        assert torch.cuda.get_device_capability() >= (
+            8,
+            9,
+        ), "FP8 is not supported for current device."
+
+    num_quant_linear = 0
+    num_skip_linear = 0
+    num_linear_layers = 0
+    num_layers = 0
+
+    # Ensure bfloat16 for per_row
+    def _filter_fn(m: torch.nn.Module, name: str) -> bool:
+        nonlocal num_quant_linear, num_skip_linear, num_linear_layers, num_layers
+        num_layers += 1
+        if isinstance(m, torch.nn.Linear):
+            num_linear_layers += 1
+            for exclude_name in exclude_layers:
+                if exclude_name in name:
+                    logger.info(
+                        f"Skip Quantization: {name} -> "
+                        f"pattern<{exclude_name}>"
+                    )
+
+                    num_skip_linear += 1
+                    return False
+
+            if (
+                per_row
+                and m.weight.dtype != torch.bfloat16
+                and quant_type == "fp8_w8a8_dq"
+            ):
+                logger.info(
+                    f"Skip Quantization: {name} -> "
+                    f"pattern<dtype({m.weight.dtype})!=bfloat16>"
+                )
+
+                num_skip_linear += 1
+                return False
+
+            num_quant_linear += 1
+            return True
+
+        return False
+
+    def _quantization_fn():
+        try:
+            if quant_type == "fp8_w8a8_dq":
+                from torchao.quantization import (
+                    float8_dynamic_activation_float8_weight,
+                    PerTensor,
+                    PerRow,
+                )
+
+                quantization_fn = float8_dynamic_activation_float8_weight(
+                    weight_dtype=kwargs.get(
+                        "weight_dtype",
+                        torch.float8_e4m3fn,
+                    ),
+                    activation_dtype=kwargs.get(
+                        "activation_dtype",
+                        torch.float8_e4m3fn,
+                    ),
+                    granularity=(
+                        ((PerRow(), PerRow()))
+                        if per_row
+                        else ((PerTensor(), PerTensor()))
+                    ),
+                )
+
+            elif quant_type == "fp8_w8a16_wo":
+                from torchao.quantization import float8_weight_only
+
+                quantization_fn = float8_weight_only(
+                    weight_dtype=kwargs.get(
+                        "weight_dtype",
+                        torch.float8_e4m3fn,
+                    ),
+                )
+
+            elif quant_type == "int8_w8a8_dq":
+                from torchao.quantization import (
+                    int8_dynamic_activation_int8_weight,
+                )
+
+                quantization_fn = int8_dynamic_activation_int8_weight()
+
+            elif quant_type == "int8_w8a16_wo":
+                from torchao.quantization import int8_weight_only
+
+                quantization_fn = int8_weight_only(
+                    # group_size is None -> per_channel, else per group
+                    group_size=kwargs.get("group_size", None),
+                )
+
+            elif quant_type == "int4_w4a8_dq":
+                from torchao.quantization import (
+                    int8_dynamic_activation_int4_weight,
+                )
+
+                quantization_fn = int8_dynamic_activation_int4_weight(
+                    group_size=kwargs.get("group_size", 32),
+                )
+
+            elif quant_type == "int4_w4a4_dq":
+                from torchao.quantization import (
+                    int4_dynamic_activation_int4_weight,
+                )
+
+                quantization_fn = int4_dynamic_activation_int4_weight()
+
+            elif quant_type == "int4_w4a16_wo":
+                from torchao.quantization import int4_weight_only
+
+                quantization_fn = int4_weight_only(
+                    group_size=kwargs.get("group_size", 32),
+                )
+
+            else:
+                raise ValueError(
+                    f"quant_type: {quant_type} is not supported now!"
+                )
+
+        except ImportError as e:
+            e.msg += (
+                f"{quant_type} is not supported in torchao backend now! "
+                "Please upgrade the torchao library."
+            )
+            raise e
+
+        return quantization_fn
+
+    from torchao.quantization import quantize_
+
+    quantize_(
+        module,
+        _quantization_fn(),
+        filter_fn=_filter_fn if filter_fn is None else filter_fn,
+        device=kwargs.get("device", None),
+    )
+
+    force_empty_cache()
+
+    logger.info(
+        f"Quantized Linear Layers: {num_quant_linear:>5}\n"
+        f"Skipped   Linear Layers: {num_skip_linear:>5}\n"
+        f"Total     Linear Layers: {num_linear_layers:>5}\n"
+        f"Total     (all)  Layers: {num_layers:>5}"
+    )
+    return module
+
+
+def force_empty_cache():
+    time.sleep(1)
+    gc.collect()
+    torch.cuda.empty_cache()
+    time.sleep(1)
+    gc.collect()
+    torch.cuda.empty_cache()

cache_dit/quantize/quantize_interface.py

@@ -0,0 +1,46 @@
+import torch
+from typing import Callable, Optional, List
+from cache_dit.logger import init_logger
+
+logger = init_logger(__name__)
+
+
+def quantize(
+    module: torch.nn.Module,
+    quant_type: str = "fp8_w8a8_dq",
+    backend: str = "ao",
+    exclude_layers: List[str] = [
+        "embedder",
+        "embed",
+    ],
+    filter_fn: Optional[Callable] = None,
+    # only for fp8_w8a8_dq
+    per_row: bool = True,
+    **kwargs,
+) -> torch.nn.Module:
+    assert isinstance(module, torch.nn.Module)
+
+    if backend.lower() in ("ao", "torchao"):
+        from cache_dit.quantize.quantize_ao import quantize_ao
+
+        quant_type = quant_type.lower()
+        assert quant_type in (
+            "fp8_w8a8_dq",
+            "fp8_w8a16_wo",
+            "int8_w8a8_dq",
+            "int8_w8a16_wo",
+            "int4_w4a8_dq",
+            "int4_w4a4_dq",
+            "int4_w4a16_wo",
+        ), f"{quant_type} is not supported for torchao backend now!"
+
+        return quantize_ao(
+            module,
+            quant_type=quant_type,
+            per_row=per_row,
+            exclude_layers=exclude_layers,
+            filter_fn=filter_fn,
+            **kwargs,
+        )
+    else:
+        raise ValueError(f"backend: {backend} is not supported now!")