Make motion_encode_batch_size configurable in pipeline __call__

Author: dg845

src/diffusers/models/transformers/transformer_wan_animate.py

@@ -1049,7 +1049,7 @@ def __init__(
         face_encoder_hidden_dim: int = 1024,
         face_encoder_num_heads: int = 4,
         inject_face_latents_blocks: int = 5,
-        motion_encoder_batch_size: Optional[int] = 8,
+        motion_encoder_batch_size: int = 8,
     ) -> None:
         super().__init__()
 
@@ -1135,8 +1135,6 @@ def __init__(
 
         self.gradient_checkpointing = False
 
-        self.motion_encoder_batch_size = motion_encoder_batch_size
-
     def forward(
         self,
         hidden_states: torch.Tensor,
@@ -1168,8 +1166,9 @@ def forward(
             face_pixel_values (`torch.Tensor` of shape `(B, C', S, H', W')`):
                 Face video in pixel space (not latent space). Typically C' = 3 and H' and W' are the height/width of
                 the face video in pixels. Here S is the inference segment length, usually set to 77.
-            motion_encode_batch_size (`int`, *optional*, defaults to `8`):
-                The batch size for batched encoding of the face video via the motion encoder.
+            motion_encode_batch_size (`int`, *optional*):
+                The batch size for batched encoding of the face video via the motion encoder. Will default to
+                `self.config.motion_encoder_batch_size` if not set.
             return_dict (`bool`, *optional*, defaults to `True`):
                 Whether to return the output as a dict or tuple.
         """
@@ -1233,7 +1232,7 @@ def forward(
 
         # Extract motion features using motion encoder
         # Perform batched motion encoder inference to allow trading off inference speed for memory usage
-        motion_encode_batch_size = motion_encode_batch_size or self.motion_encoder_batch_size
+        motion_encode_batch_size = motion_encode_batch_size or self.config.motion_encoder_batch_size
         face_batches = torch.split(face_pixel_values, motion_encode_batch_size)
         motion_vec_batches = []
         for face_batch in face_batches:

src/diffusers/pipelines/wan/pipeline_wan_animate.py

@@ -787,6 +787,7 @@ def __call__(
         num_inference_steps: int = 20,
         mode: str = "animate",
         prev_segment_conditioning_frames: int = 1,
+        motion_encode_batch_size: Optional[int] = None,
         guidance_scale: float = 1.0,
         num_videos_per_prompt: Optional[int] = 1,
         generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
@@ -832,6 +833,10 @@ def __call__(
             prev_segment_conditioning_frames (`int`, defaults to `1`):
                 The number of frames from the previous video segment to be used for temporal guidance. Recommended to
                 be 1 or 5. In general, should be 4N + 1, where N is a non-negative integer.
+            motion_encode_batch_size (`int`, *optional*):
+                The batch size for batched encoding of the face video via the motion encoder. This allows trading off
+                inference speed for lower memory usage by setting a smaller batch size. Will default to
+                `self.transformer.config.motion_encoder_batch_size` if not set.
             height (`int`, defaults to `720`):
                 The height of the generated video.
             width (`int`, defaults to `1280`):
@@ -1127,6 +1132,7 @@ def __call__(
                             encoder_hidden_states_image=image_embeds,
                             pose_hidden_states=pose_latents,
                             face_pixel_values=face_video_segment,
+                            motion_encode_batch_size=motion_encode_batch_size,
                             attention_kwargs=attention_kwargs,
                             return_dict=False,
                         )[0]
@@ -1142,6 +1148,7 @@ def __call__(
                                 encoder_hidden_states_image=image_embeds,
                                 pose_hidden_states=pose_latents,
                                 face_pixel_values=face_pixel_values_uncond,
+                                motion_encode_batch_size=motion_encode_batch_size,
                                 attention_kwargs=attention_kwargs,
                                 return_dict=False,
                             )[0]