<fix> add tp ip attn enhancement control in forward.

app.py

@@ -264,6 +264,7 @@ def process_image_and_text(condition_image, target_prompt, condition_image_promp
         frame_gap=48,
         mixup=True,
         mixup_num_imgs=2,
+        enhance_tp=task in ['subject_driven', 'style_transfer'],
     ).frames
 
     gen_img = gen_img[:, 0:1, :, :, :]

models/hyvideo/transformer_hunyuan_video_i2v.py

@@ -64,6 +64,7 @@ class HunyuanVideoAttnProcessor2_0:
         encoder_hidden_states: Optional[torch.Tensor] = None,
         attention_mask: Optional[torch.Tensor] = None,
         image_rotary_emb: Optional[torch.Tensor] = None,
+        enhance_tp: bool = False,
     ) -> torch.Tensor:
         if attn.add_q_proj is None and encoder_hidden_states is not None:
             hidden_states = torch.cat([hidden_states, encoder_hidden_states], dim=1)
@@ -154,7 +155,7 @@ class HunyuanVideoAttnProcessor2_0:
                 k_lens = torch.tensor([sum([u[seg_start[seg]:seg_end[seg]].long().sum().item() for seg in segs]) for u in valid_indices for segs in k_segs], 
                                         dtype=torch.int32, device=valid_indices.device)
                 query = torch.cat([u[i:j][v[i:j]] for u,v in zip(query, valid_indices) for i,j in zip(seg_start, seg_end)], dim=0)
-                if self.inference_subject_driven:
+                if self.inference_subject_driven or enhance_tp:
                     key = torch.cat([torch.cat([ torch.cat([u[seg_start[seg]:seg_end[seg]][v[seg_start[seg]:seg_end[seg]]][:144], u[seg_start[seg]:seg_end[seg]][v[seg_start[seg]:seg_end[seg]]][144:] + 0.6 * u[seg_start[seg]:seg_end[seg]][v[seg_start[seg]:seg_end[seg]]][144:].abs().mean()], dim=0) if segs == [0, 1, 2] and seg == 2 else u[seg_start[seg]:seg_end[seg]][v[seg_start[seg]:seg_end[seg]]] for seg in segs], dim=0) \
                                         for u,v in zip(key, valid_indices) for segs in k_segs], dim=0)
                 else:
@@ -756,6 +757,7 @@ class HunyuanVideoTokenReplaceSingleTransformerBlock(nn.Module):
         image_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
         token_replace_emb: torch.Tensor = None,
         num_tokens: int = None,
+        enhance_tp: bool = False,
     ) -> torch.Tensor:
         text_seq_length = encoder_hidden_states.shape[1]
         hidden_states = torch.cat([hidden_states, encoder_hidden_states], dim=1)
@@ -777,6 +779,7 @@ class HunyuanVideoTokenReplaceSingleTransformerBlock(nn.Module):
             encoder_hidden_states=norm_encoder_hidden_states,
             attention_mask=attention_mask,
             image_rotary_emb=image_rotary_emb,
+            enhance_tp=enhance_tp,
         )
         attn_output = torch.cat([attn_output, context_attn_output], dim=1)
 
@@ -841,6 +844,7 @@ class HunyuanVideoTokenReplaceTransformerBlock(nn.Module):
         freqs_cis: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
         token_replace_emb: torch.Tensor = None,
         num_tokens: int = None,
+        enhance_tp: bool = False,
     ) -> Tuple[torch.Tensor, torch.Tensor]:
         # 1. Input normalization
         (
@@ -864,6 +868,7 @@ class HunyuanVideoTokenReplaceTransformerBlock(nn.Module):
             encoder_hidden_states=norm_encoder_hidden_states,
             attention_mask=attention_mask,
             image_rotary_emb=freqs_cis,
+            enhance_tp=enhance_tp,
         )
 
         # 3. Modulation and residual connection
@@ -1109,6 +1114,7 @@ class HunyuanVideoTransformer3DModel(ModelMixin, ConfigMixin, PeftAdapterMixin,
         attention_kwargs: Optional[Dict[str, Any]] = None,
         return_dict: bool = True,
         frame_gap: Union[int, None] = None,
+        enhance_tp: bool = False,
     ) -> Union[torch.Tensor, Dict[str, torch.Tensor]]:
         if attention_kwargs is not None:
             attention_kwargs = attention_kwargs.copy()
@@ -1181,6 +1187,7 @@ class HunyuanVideoTransformer3DModel(ModelMixin, ConfigMixin, PeftAdapterMixin,
                     image_rotary_emb,
                     token_replace_emb,
                     first_frame_num_tokens,
+                    enhance_tp,
                 )
 
             for block in self.single_transformer_blocks:
@@ -1193,6 +1200,7 @@ class HunyuanVideoTransformer3DModel(ModelMixin, ConfigMixin, PeftAdapterMixin,
                     image_rotary_emb,
                     token_replace_emb,
                     first_frame_num_tokens,
+                    enhance_tp,
                 )
 
         else:
@@ -1205,6 +1213,7 @@ class HunyuanVideoTransformer3DModel(ModelMixin, ConfigMixin, PeftAdapterMixin,
                     image_rotary_emb,
                     token_replace_emb,
                     first_frame_num_tokens,
+                    enhance_tp,
                 )
 
             for block in self.single_transformer_blocks:
@@ -1216,6 +1225,7 @@ class HunyuanVideoTransformer3DModel(ModelMixin, ConfigMixin, PeftAdapterMixin,
                     image_rotary_emb,
                     token_replace_emb,
                     first_frame_num_tokens,
+                    enhance_tp,
                 )
 
         # 5. Output projection

pipelines/pipeline_hunyuan_video_i2v.py

@@ -649,6 +649,7 @@ class HunyuanVideoImageToVideoPipeline(DiffusionPipeline, HunyuanVideoLoraLoader
         frame_gap: Union[int, None] = None,
         mixup: bool = False,
         mixup_num_imgs: Union[int, None] = None,
+        enhance_tp: bool = False,
     ):
         r"""
         The call function to the pipeline for generation.
@@ -899,6 +900,7 @@ class HunyuanVideoImageToVideoPipeline(DiffusionPipeline, HunyuanVideoLoraLoader
                     attention_kwargs=attention_kwargs,
                     return_dict=False,
                     frame_gap=int(frame_gap / 4) if frame_gap is not None else frame_gap,
+                    enhance_tp=enhance_tp,
                 )[0]
 
                 if do_true_cfg: