Update all files for BitDance-14B-64x-diffusers

bitdance_diffusers/pipeline_bitdance.py

@@ -1,7 +1,5 @@
-from __future__ import annotations
-
 from contextlib import nullcontext
-from typing import List, Optional, Sequence, Tuple, Union
+from typing import List, Optional, Tuple, Union
 
 import torch
 from einops import rearrange
@@ -17,6 +15,13 @@ from .constants import SUPPORTED_IMAGE_SIZES
 PromptType = Union[str, List[str]]
 
 
+def _get_pkv_seq_len(past_key_values) -> int:
+    """Get cached sequence length from past_key_values (supports tuple and DynamicCache)."""
+    if hasattr(past_key_values, "get_seq_length"):
+        return past_key_values.get_seq_length()
+    return past_key_values[0][0].shape[2]
+
+
 class BitDanceDiffusionPipeline(DiffusionPipeline):
     model_cpu_offload_seq = "text_encoder->projector->diffusion_head->autoencoder"
 
@@ -27,7 +32,8 @@ class BitDanceDiffusionPipeline(DiffusionPipeline):
         autoencoder,
         diffusion_head,
         projector,
-        supported_image_sizes: Optional[Sequence[Sequence[int]]] = None,
+        supported_image_sizes: Optional[List[List[int]]] = None,
+        dtype: Optional[torch.dtype] = None,
     ) -> None:
         super().__init__()
         self.register_modules(
@@ -130,6 +136,8 @@ class BitDanceDiffusionPipeline(DiffusionPipeline):
     def _decode_tokens_to_image(self, image_latents: torch.Tensor, image_size: Tuple[int, int], ps: int = 1) -> torch.Tensor:
         h, w = image_size
         image_latents = rearrange(image_latents, "b (h w p1 p2) c -> b c (h p1) (w p2)", h=h // ps, w=w // ps, p1=ps, p2=ps)
+        ae_dtype = next(self.autoencoder.parameters()).dtype
+        image_latents = image_latents.to(dtype=ae_dtype)
         return self.autoencoder.decode(image_latents)
 
     @torch.no_grad()
@@ -183,7 +191,7 @@ class BitDanceDiffusionPipeline(DiffusionPipeline):
             pkv_c = outputs_c.past_key_values
 
             bi_attn_mask = torch.ones(
-                (input_embeds_cond.shape[0], 1, step_width, step_width + pkv_c[0][0].shape[2]),
+                (input_embeds_cond.shape[0], 1, step_width, step_width + _get_pkv_seq_len(pkv_c)),
                 dtype=torch.bool,
                 device=device,
             )
@@ -201,11 +209,16 @@ class BitDanceDiffusionPipeline(DiffusionPipeline):
             if guidance_scale > 1.0 and input_embeds_uncond is not None:
                 outputs_u = model(inputs_embeds=input_embeds_uncond[:, :-step_width, :], use_cache=True)
                 pkv_u = outputs_u.past_key_values
+                bi_attn_mask_u = torch.ones(
+                    (input_embeds_uncond.shape[0], 1, step_width, step_width + _get_pkv_seq_len(pkv_u)),
+                    dtype=torch.bool,
+                    device=device,
+                )
                 outputs_u = model(
                     inputs_embeds=input_embeds_uncond[:, -step_width:, :],
                     past_key_values=pkv_u,
                     use_cache=True,
-                    attention_mask=bi_attn_mask,
+                    attention_mask=bi_attn_mask_u,
                 )
                 pkv_u = outputs_u.past_key_values
                 hidden_u = outputs_u.last_hidden_state[:, -step_width:]
@@ -235,7 +248,7 @@ class BitDanceDiffusionPipeline(DiffusionPipeline):
 
                 model_input = curr_embeds + pos_slice
                 bi_attn_mask = torch.ones(
-                    (model_input.shape[0], 1, model_input.shape[1], model_input.shape[1] + pkv_c[0][0].shape[2]),
+                    (model_input.shape[0], 1, model_input.shape[1], model_input.shape[1] + _get_pkv_seq_len(pkv_c)),
                     dtype=torch.bool,
                     device=device,
                 )
@@ -249,11 +262,16 @@ class BitDanceDiffusionPipeline(DiffusionPipeline):
                 hidden_c = outputs_c.last_hidden_state[:, -step_width:]
 
                 if guidance_scale > 1.0 and hidden_u is not None and pkv_u is not None:
+                    bi_attn_mask_u = torch.ones(
+                        (model_input.shape[0], 1, model_input.shape[1], model_input.shape[1] + _get_pkv_seq_len(pkv_u)),
+                        dtype=torch.bool,
+                        device=device,
+                    )
                     outputs_u = model(
                         inputs_embeds=model_input[num_images_per_prompt:],
                         past_key_values=pkv_u,
                         use_cache=True,
-                        attention_mask=bi_attn_mask[num_images_per_prompt:],
+                        attention_mask=bi_attn_mask_u[num_images_per_prompt:],
                     )
                     pkv_u = outputs_u.past_key_values
                     hidden_u = outputs_u.last_hidden_state[:, -step_width:]