cache-dit 0.2.33__py3-none-any.whl → 0.2.36__py3-none-any.whl

cache_dit/cache_factory/patch_functors/functor_chroma.py

@@ -1,10 +1,9 @@
-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 (
+    ChromaTransformerBlock,
     ChromaSingleTransformerBlock,
     Transformer2DModelOutput,
 )
@@ -27,24 +26,31 @@ class ChromaPatchFunctor(PatchFunctor):
     def apply(
         self,
         transformer: ChromaTransformer2DModel,
-        blocks: torch.nn.ModuleList = None,
         **kwargs,
     ) -> ChromaTransformer2DModel:
         if hasattr(transformer, "_is_patched"):
             return transformer
 
-        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
+        for index_block, block in enumerate(transformer.transformer_blocks):
+            assert isinstance(block, ChromaTransformerBlock)
+            img_offset = 3 * len(transformer.single_transformer_blocks)
+            txt_offset = img_offset + 6 * len(transformer.transformer_blocks)
+            img_modulation = img_offset + 6 * index_block
+            text_modulation = txt_offset + 6 * index_block
+            block._img_modulation = img_modulation
+            block._text_modulation = text_modulation
+            block.forward = __patch_double_forward__.__get__(block)
+
+        for index_block, block in enumerate(
+            transformer.single_transformer_blocks
+        ):
+            assert isinstance(block, ChromaSingleTransformerBlock)
+            start_idx = 3 * index_block
+            block._start_idx = start_idx
+            block.forward = __patch_single_forward__.__get__(block)
+
+        is_patched = True
 
         cls_name = transformer.__class__.__name__
 
@@ -69,25 +75,123 @@ class ChromaPatchFunctor(PatchFunctor):
         return transformer
 
 
+# Adapted from: https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/transformers/transformer_chroma.py
+def __patch_double_forward__(
+    self: ChromaTransformerBlock,
+    hidden_states: torch.Tensor,
+    encoder_hidden_states: torch.Tensor,
+    pooled_temb: torch.Tensor,
+    image_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
+    attention_mask: Optional[torch.Tensor] = None,
+    joint_attention_kwargs: Optional[Dict[str, Any]] = None,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    # TODO: Fuse controlnet into block forward
+    img_modulation = self._img_modulation
+    text_modulation = self._text_modulation
+    temb = torch.cat(
+        (
+            pooled_temb[:, img_modulation : img_modulation + 6],
+            pooled_temb[:, text_modulation : text_modulation + 6],
+        ),
+        dim=1,
+    )
+
+    temb_img, temb_txt = temb[:, :6], temb[:, 6:]
+    norm_hidden_states, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.norm1(
+        hidden_states, emb=temb_img
+    )
+
+    (
+        norm_encoder_hidden_states,
+        c_gate_msa,
+        c_shift_mlp,
+        c_scale_mlp,
+        c_gate_mlp,
+    ) = self.norm1_context(encoder_hidden_states, emb=temb_txt)
+    joint_attention_kwargs = joint_attention_kwargs or {}
+    if attention_mask is not None:
+        attention_mask = (
+            attention_mask[:, None, None, :] * attention_mask[:, None, :, None]
+        )
+
+    # Attention.
+    attention_outputs = self.attn(
+        hidden_states=norm_hidden_states,
+        encoder_hidden_states=norm_encoder_hidden_states,
+        image_rotary_emb=image_rotary_emb,
+        attention_mask=attention_mask,
+        **joint_attention_kwargs,
+    )
+
+    if len(attention_outputs) == 2:
+        attn_output, context_attn_output = attention_outputs
+    elif len(attention_outputs) == 3:
+        attn_output, context_attn_output, ip_attn_output = attention_outputs
+
+    # Process attention outputs for the `hidden_states`.
+    attn_output = gate_msa.unsqueeze(1) * attn_output
+    hidden_states = hidden_states + attn_output
+
+    norm_hidden_states = self.norm2(hidden_states)
+    norm_hidden_states = (
+        norm_hidden_states * (1 + scale_mlp[:, None]) + shift_mlp[:, None]
+    )
+
+    ff_output = self.ff(norm_hidden_states)
+    ff_output = gate_mlp.unsqueeze(1) * ff_output
+
+    hidden_states = hidden_states + ff_output
+    if len(attention_outputs) == 3:
+        hidden_states = hidden_states + ip_attn_output
+
+    # Process attention outputs for the `encoder_hidden_states`.
+
+    context_attn_output = c_gate_msa.unsqueeze(1) * context_attn_output
+    encoder_hidden_states = encoder_hidden_states + context_attn_output
+
+    norm_encoder_hidden_states = self.norm2_context(encoder_hidden_states)
+    norm_encoder_hidden_states = (
+        norm_encoder_hidden_states * (1 + c_scale_mlp[:, None])
+        + c_shift_mlp[:, None]
+    )
+
+    context_ff_output = self.ff_context(norm_encoder_hidden_states)
+    encoder_hidden_states = (
+        encoder_hidden_states + c_gate_mlp.unsqueeze(1) * context_ff_output
+    )
+    if encoder_hidden_states.dtype == torch.float16:
+        encoder_hidden_states = encoder_hidden_states.clip(-65504, 65504)
+
+    return encoder_hidden_states, hidden_states
+
+
 # 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,
+    pooled_temb: torch.Tensor,
     image_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
+    attention_mask: Optional[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)
+) -> torch.Tensor:
+    # TODO: Fuse controlnet into block forward
+    start_idx = self._start_idx
+    temb = pooled_temb[:, start_idx : start_idx + 3]
 
     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 {}
+
+    if attention_mask is not None:
+        attention_mask = (
+            attention_mask[:, None, None, :] * attention_mask[:, None, :, None]
+        )
+
     attn_output = self.attn(
         hidden_states=norm_hidden_states,
         image_rotary_emb=image_rotary_emb,
+        attention_mask=attention_mask,
         **joint_attention_kwargs,
     )
 
@@ -98,11 +202,7 @@ def __patch_single_forward__(
     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
+    return hidden_states
 
 
 # Adapted from: https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/transformers/transformer_chroma.py
@@ -174,24 +274,13 @@ def __patch_transformer_forward__(
         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,
+                    pooled_temb,
                     image_rotary_emb,
                     attention_mask,
                 )
@@ -201,12 +290,13 @@ def __patch_transformer_forward__(
             encoder_hidden_states, hidden_states = block(
                 hidden_states=hidden_states,
                 encoder_hidden_states=encoder_hidden_states,
-                temb=temb,
+                pooled_temb=pooled_temb,
                 image_rotary_emb=image_rotary_emb,
                 attention_mask=attention_mask,
                 joint_attention_kwargs=joint_attention_kwargs,
             )
 
+        # TODO: Fuse controlnet into block forward
         # controlnet residual
         if controlnet_block_samples is not None:
             interval_control = len(self.transformer_blocks) / len(
@@ -227,43 +317,43 @@ def __patch_transformer_forward__(
                     + controlnet_block_samples[index_block // interval_control]
                 )
 
+    hidden_states = torch.cat([encoder_hidden_states, hidden_states], dim=1)
+
     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,
-                )
+            hidden_states = self._gradient_checkpointing_func(
+                block,
+                hidden_states,
+                pooled_temb,
+                image_rotary_emb,
+                attention_mask,
+                joint_attention_kwargs,
             )
 
         else:
-            encoder_hidden_states, hidden_states = block(
+            hidden_states = block(
                 hidden_states=hidden_states,
-                encoder_hidden_states=encoder_hidden_states,
-                temb=temb,
+                pooled_temb=pooled_temb,
                 image_rotary_emb=image_rotary_emb,
                 attention_mask=attention_mask,
                 joint_attention_kwargs=joint_attention_kwargs,
             )
 
+        # TODO: Fuse controlnet into block forward
         # 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
+            hidden_states[:, encoder_hidden_states.shape[1] :, ...] = (
+                hidden_states[:, encoder_hidden_states.shape[1] :, ...]
                 + controlnet_single_block_samples[
                     index_block // interval_control
                 ]
             )
 
+    hidden_states = hidden_states[:, encoder_hidden_states.shape[1] :, ...]
     temb = pooled_temb[:, -2:]
     hidden_states = self.norm_out(hidden_states, temb)
     output = self.proj_out(hidden_states)

cache_dit/cache_factory/patch_functors/functor_dit.py

@@ -0,0 +1,130 @@
+import torch
+import torch.nn.functional as F
+
+from typing import Optional, Dict, Any
+from diffusers.models.transformers.dit_transformer_2d import (
+    DiTTransformer2DModel,
+    Transformer2DModelOutput,
+)
+from cache_dit.cache_factory.patch_functors.functor_base import (
+    PatchFunctor,
+)
+from cache_dit.logger import init_logger
+
+logger = init_logger(__name__)
+
+
+class DiTPatchFunctor(PatchFunctor):
+
+    def apply(
+        self,
+        transformer: DiTTransformer2DModel,
+        **kwargs,
+    ) -> DiTTransformer2DModel:
+        if hasattr(transformer, "_is_patched"):
+            return transformer
+
+        is_patched = False
+
+        transformer._norm1_emb = transformer.transformer_blocks[0].norm1.emb
+
+        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_transformer_forward__(
+    self: DiTTransformer2DModel,
+    hidden_states: torch.Tensor,
+    timestep: Optional[torch.LongTensor] = None,
+    class_labels: Optional[torch.LongTensor] = None,
+    cross_attention_kwargs: Dict[str, Any] = None,
+    return_dict: bool = True,
+):
+    height, width = (
+        hidden_states.shape[-2] // self.patch_size,
+        hidden_states.shape[-1] // self.patch_size,
+    )
+    hidden_states = self.pos_embed(hidden_states)
+
+    # 2. Blocks
+    for block in self.transformer_blocks:
+        if torch.is_grad_enabled() and self.gradient_checkpointing:
+            hidden_states = self._gradient_checkpointing_func(
+                block,
+                hidden_states,
+                None,
+                None,
+                None,
+                timestep,
+                cross_attention_kwargs,
+                class_labels,
+            )
+        else:
+            hidden_states = block(
+                hidden_states,
+                attention_mask=None,
+                encoder_hidden_states=None,
+                encoder_attention_mask=None,
+                timestep=timestep,
+                cross_attention_kwargs=cross_attention_kwargs,
+                class_labels=class_labels,
+            )
+
+    # 3. Output
+    # conditioning = self.transformer_blocks[0].norm1.emb(timestep, class_labels, hidden_dtype=hidden_states.dtype)
+    conditioning = self._norm1_emb(
+        timestep, class_labels, hidden_dtype=hidden_states.dtype
+    )
+    shift, scale = self.proj_out_1(F.silu(conditioning)).chunk(2, dim=1)
+    hidden_states = (
+        self.norm_out(hidden_states) * (1 + scale[:, None]) + shift[:, None]
+    )
+    hidden_states = self.proj_out_2(hidden_states)
+
+    # unpatchify
+    height = width = int(hidden_states.shape[1] ** 0.5)
+    hidden_states = hidden_states.reshape(
+        shape=(
+            -1,
+            height,
+            width,
+            self.patch_size,
+            self.patch_size,
+            self.out_channels,
+        )
+    )
+    hidden_states = torch.einsum("nhwpqc->nchpwq", hidden_states)
+    output = hidden_states.reshape(
+        shape=(
+            -1,
+            self.out_channels,
+            height * self.patch_size,
+            width * self.patch_size,
+        )
+    )
+
+    if not return_dict:
+        return (output,)
+
+    return Transformer2DModelOutput(sample=output)

cache_dit/metrics/clip_score.py

@@ -0,0 +1,135 @@
+import os
+import re
+import pathlib
+import numpy as np
+from PIL import Image
+from tqdm import tqdm
+import torch
+from transformers import CLIPProcessor, CLIPModel
+
+from typing import Tuple, Union
+from cache_dit.metrics.config import _IMAGE_EXTENSIONS
+from cache_dit.metrics.config import get_metrics_verbose
+from cache_dit.logger import init_logger
+
+logger = init_logger(__name__)
+
+
+DISABLE_VERBOSE = not get_metrics_verbose()
+
+
+class CLIPScore:
+    def __init__(
+        self,
+        device="cuda" if torch.cuda.is_available() else "cpu",
+        clip_model_path: str = None,
+    ):
+        self.device = device
+        if clip_model_path is None:
+            clip_model_path = os.environ.get(
+                "CLIP_MODEL_DIR", "laion/CLIP-ViT-g-14-laion2B-s12B-b42K"
+            )
+
+        # Load models
+        self.clip_model = CLIPModel.from_pretrained(clip_model_path)
+        self.clip_model = self.clip_model.to(device)  # type: ignore
+        self.clip_processor = CLIPProcessor.from_pretrained(clip_model_path)
+
+    @torch.no_grad()
+    def compute_clip_score(
+        self,
+        img: Image.Image | np.ndarray,
+        prompt: str,
+    ) -> float:
+        if isinstance(img, Image.Image):
+            img_pil = img.convert("RGB")
+        elif isinstance(img, np.ndarray):
+            img_pil = Image.fromarray(img).convert("RGB")
+        else:
+            img_pil = Image.open(img).convert("RGB")
+        with torch.no_grad():
+            inputs = self.clip_processor(
+                text=prompt,
+                images=img_pil,
+                return_tensors="pt",
+                padding=True,
+                truncation=True,
+            ).to(self.device)
+            outputs = self.clip_model(**inputs)
+        return outputs.logits_per_image.item()
+
+
+clip_score_instance: CLIPScore = None
+
+
+def compute_clip_score_img(
+    img: Image.Image | np.ndarray | str,
+    prompt: str,
+    clip_model_path: str = None,
+) -> float:
+    global clip_score_instance
+    if clip_score_instance is None:
+        clip_score_instance = CLIPScore(clip_model_path=clip_model_path)
+    assert clip_score_instance is not None
+    return clip_score_instance.compute_clip_score(img, prompt)
+
+
+def compute_clip_score(
+    img_dir: Image.Image | np.ndarray | str,
+    prompts: str | list[str],
+    clip_model_path: str = None,
+) -> Union[Tuple[float, int], Tuple[None, None]]:
+    if not os.path.isdir(img_dir) or (
+        not isinstance(prompts, list) and not os.path.isfile(prompts)
+    ):
+        return (
+            compute_clip_score_img(
+                img_dir,
+                prompts,
+                clip_model_path=clip_model_path,
+            ),
+            1,
+        )
+
+    # compute dir metric
+    def natural_sort_key(filename):
+        match = re.search(r"(\d+)\D*$", filename)
+        return int(match.group(1)) if match else filename
+
+    img_dir: pathlib.Path = pathlib.Path(img_dir)
+    img_files = [
+        file
+        for ext in _IMAGE_EXTENSIONS
+        for file in img_dir.rglob("*.{}".format(ext))
+    ]
+    img_files = [file.as_posix() for file in img_files]
+    img_files = sorted(img_files, key=natural_sort_key)
+
+    if os.path.isfile(prompts):
+        """Load prompts from file"""
+        with open(prompts, "r", encoding="utf-8") as f:
+            prompts_load = [line.strip() for line in f.readlines()]
+        prompts = prompts_load.copy()
+
+    vaild_len = min(len(img_files), len(prompts))
+    img_files = img_files[:vaild_len]
+    prompts = prompts[:vaild_len]
+
+    clip_scores = []
+
+    for img_file, prompt in tqdm(
+        zip(img_files, prompts),
+        total=vaild_len,
+        disable=not get_metrics_verbose(),
+    ):
+        clip_scores.append(
+            compute_clip_score_img(
+                img_file,
+                prompt,
+                clip_model_path=clip_model_path,
+            )
+        )
+
+    if vaild_len > 0:
+        return np.mean(clip_scores), vaild_len
+    return None, None

cache_dit/metrics/fid.py

@@ -1,6 +1,8 @@
 import os
 import cv2
 import pathlib
+import warnings
+
 import numpy as np
 from PIL import Image
 from tqdm import tqdm
@@ -8,13 +10,21 @@ from scipy import linalg
 import torch
 import torchvision.transforms as TF
 from torch.nn.functional import adaptive_avg_pool2d
+
+from typing import Tuple, Union
 from cache_dit.metrics.inception import InceptionV3
 from cache_dit.metrics.config import _IMAGE_EXTENSIONS
 from cache_dit.metrics.config import _VIDEO_EXTENSIONS
+from cache_dit.metrics.config import get_metrics_verbose
+from cache_dit.utils import disable_print
 from cache_dit.logger import init_logger
 
+warnings.filterwarnings("ignore")
+
 logger = init_logger(__name__)
 
+DISABLE_VERBOSE = not get_metrics_verbose()
+
 
 # Adapted from: https://github.com/mseitzer/pytorch-fid/blob/master/src/pytorch_fid/fid_score.py
 class ImagePathDataset(torch.utils.data.Dataset):
@@ -496,3 +506,35 @@ class FrechetInceptionDistance:
             return [], [], 0
 
         return video_true_frames, video_test_frames, valid_frames
+
+
+fid_instance: FrechetInceptionDistance = None
+
+
+def compute_fid(
+    image_true: np.ndarray | str,
+    image_test: np.ndarray | str,
+) -> Union[Tuple[float, int], Tuple[None, None]]:
+    global fid_instance
+    if fid_instance is None:
+        with disable_print():
+            fid_instance = FrechetInceptionDistance(
+                disable_tqdm=not get_metrics_verbose(),
+            )
+    assert fid_instance is not None
+    return fid_instance.compute_fid(image_true, image_test)
+
+
+def compute_video_fid(
+    # file or dir
+    video_true: str,
+    video_test: str,
+) -> Union[Tuple[float, int], Tuple[None, None]]:
+    global fid_instance
+    if fid_instance is None:
+        with disable_print():
+            fid_instance = FrechetInceptionDistance(
+                disable_tqdm=not get_metrics_verbose(),
+            )
+    assert fid_instance is not None
+    return fid_instance.compute_fid(video_true, video_test)