update

Author: DN6

src/diffusers/pipelines/hunyuan_video/pipeline_hunyuan_video_image2video.py

@@ -24,6 +24,7 @@
     CLIPTokenizer,
     LlamaTokenizerFast,
     LlavaForConditionalGeneration,
+    LlavaProcessor,
 )
 
 from ...callbacks import MultiPipelineCallbacks, PipelineCallback
@@ -100,85 +101,6 @@
 }
 
 
-def _merge_input_ids_with_image_features(
-    image_features, inputs_embeds, input_ids, attention_mask, image_token_index, pad_token_id
-):
-    num_images, num_image_patches, embed_dim = image_features.shape
-    batch_size, sequence_length = input_ids.shape
-    special_image_token_mask = input_ids == image_token_index
-    # 1. Create a mask to know where special image tokens are
-    special_image_token_mask = input_ids == image_token_index
-    num_special_image_tokens = torch.sum(special_image_token_mask, dim=-1)
-    batch_indices, non_image_indices = torch.where(input_ids != image_token_index)
-    max_embed_dim = (num_special_image_tokens.max() * (num_image_patches - 1)) + sequence_length
-    batch_indices, non_image_indices = torch.where(input_ids != image_token_index)
-
-    # 2. Compute the positions where text should be written
-    # Calculate new positions for text tokens in merged image-text sequence.
-    # `special_image_token_mask` identifies image tokens. Each image token will be replaced by `nb_text_tokens_per_images - 1` text tokens.
-    # `torch.cumsum` computes how each image token shifts subsequent text token positions.
-    # - 1 to adjust for zero-based indexing, as `cumsum` inherently increases indices by one.
-    new_token_positions = torch.cumsum((special_image_token_mask * (num_image_patches - 1) + 1), -1) - 1
-    nb_image_pad = max_embed_dim - 1 - new_token_positions[:, -1]
-    if left_padding:
-        new_token_positions += nb_image_pad[:, None]  # offset for left padding
-    text_to_overwrite = new_token_positions[batch_indices, non_image_indices]
-
-    # 3. Create the full embedding, already padded to the maximum position
-    final_embedding = torch.zeros(
-        batch_size, max_embed_dim, embed_dim, dtype=inputs_embeds.dtype, device=inputs_embeds.device
-    )
-    final_attention_mask = torch.zeros(
-        batch_size, max_embed_dim, dtype=attention_mask.dtype, device=inputs_embeds.device
-    )
-    # In case the Vision model or the Language model has been offloaded to CPU, we need to manually
-    # set the corresponding tensors into their correct target device.
-    target_device = inputs_embeds.device
-    batch_indices, non_image_indices, text_to_overwrite = (
-        batch_indices.to(target_device),
-        non_image_indices.to(target_device),
-        text_to_overwrite.to(target_device),
-    )
-    attention_mask = attention_mask.to(target_device)
-
-    # 4. Fill the embeddings based on the mask. If we have ["hey" "<image>", "how", "are"]
-    # we need to index copy on [0, 577, 578, 579] for the text and [1:576] for the image features
-    final_embedding[batch_indices, text_to_overwrite] = inputs_embeds[batch_indices, non_image_indices]
-    final_attention_mask[batch_indices, text_to_overwrite] = attention_mask[batch_indices, non_image_indices]
-
-    # 5. Fill the embeddings corresponding to the images. Anything that is not `text_positions` needs filling (#29835)
-    image_to_overwrite = torch.full((batch_size, max_embed_dim), True, dtype=torch.bool, device=inputs_embeds.device)
-    image_to_overwrite[batch_indices, text_to_overwrite] = False
-    if left_padding:
-        image_to_overwrite &= image_to_overwrite.cumsum(-1) - 1 >= nb_image_pad[:, None].to(target_device)
-    else:
-        mask = torch.ones_like(image_to_overwrite, dtype=torch.bool).cumsum(-1) - 1
-        padding_mask = mask <= new_token_positions[:, -1:].to(target_device)
-        image_to_overwrite &= padding_mask
-
-    if image_to_overwrite.sum() != image_features.shape[:-1].numel():
-        raise ValueError(
-            f"The input provided to the model are wrong. The number of image tokens is {torch.sum(special_image_token_mask)} while"
-            f" the number of image given to the model is {num_images}. This prevents correct indexing and breaks batch generation."
-        )
-
-    batch_indices, pad_indices = torch.where(input_ids == pad_token_id)
-    final_attention_mask |= image_to_overwrite
-    position_ids = (final_attention_mask.cumsum(-1) - 1).masked_fill_((final_attention_mask == 0), 1)
-
-    # 6. Mask out the embedding at padding positions, as we later use the past_key_value value to determine the non-attended tokens.
-    batch_indices, pad_indices = torch.where(input_ids == pad_token_id)
-    indices_to_mask = new_token_positions[batch_indices, pad_indices]
-
-    final_embedding[batch_indices, indices_to_mask] = 0
-
-    return final_embedding, final_attention_mask, position_ids
-
-
-def _text_encoder_custom_forward():
-    return
-
-
 # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps
 def retrieve_timesteps(
     scheduler,
@@ -310,6 +232,13 @@ def __init__(
         self.vae_scale_factor_temporal = self.vae.temporal_compression_ratio if getattr(self, "vae", None) else 4
         self.vae_scale_factor_spatial = self.vae.spatial_compression_ratio if getattr(self, "vae", None) else 8
         self.video_processor = VideoProcessor(vae_scale_factor=self.vae_scale_factor_spatial)
+        self.llava_processor = LlavaProcessor(
+            self.image_processor,
+            self.tokenizer,
+            patch_size=self.text_encoder.config.vision_config.patch_size,
+            vision_feature_select_strategy=self.text_encoder.config.vision_feature_select_strategy,
+            num_additional_image_tokens=1,
+        )
 
     def _get_llama_prompt_embeds(
         self,
@@ -358,30 +287,25 @@ def _get_llama_prompt_embeds(
         prompt_attention_mask = text_inputs.attention_mask.to(device=device)
 
         image_embeds = self.image_processor(image, return_tensors="pt").pixel_values.to(device)
-        input_embeds = self.text_encoder.get_input_embeddings()(text_input_ids)
-
-        inputs_embeds, attention_mask, position_ids = _merge_input_ids_with_image_features(
-            image_embeds,
-            input_embeds,
-            text_input_ids,
-            prompt_attention_mask,
-            self.text_encoder.config.image_token_index,
-            self.text_encoder.pad_token_id,
-        )
+        inputs = self.llava_processor(
+            text=prompt,
+            images=image,
+            # max_length=max_sequence_length,
+            padding="max_length",
+            truncation=True,
+            return_length=False,
+            return_overflowing_tokens=False,
+            return_attention_mask=True,
+            return_tensors="pt",
+        ).to(device)
+
+        text_input_ids = inputs["input_ids"]
+        prompt_attention_mask = inputs["attention_mask"]
 
-        prompt_embeds = self.text_encoder.language_model(
-            attention_mask=attention_mask,
-            position_ids=position_ids,
-            inputs_embeds=inputs_embeds,
-        ).hidden_states[-(num_hidden_layers_to_skip + 1)]
-        """
         prompt_embeds = self.text_encoder(
-            input_ids=text_input_ids,
-            attention_mask=prompt_attention_mask,
-            pixel_values=image_embeds,
+            **inputs,
             output_hidden_states=True,
         ).hidden_states[-(num_hidden_layers_to_skip + 1)]
-        """
         prompt_embeds = prompt_embeds.to(dtype=dtype)
 
         image_emb_len = prompt_template.get("image_emb_len", 576)