diff --git a/src/maxdiffusion/pipelines/wan/wan_vace_pipeline_2_1.py b/src/maxdiffusion/pipelines/wan/wan_vace_pipeline_2_1.py
index 868c8fa9e..aeaf7b617 100644
--- a/src/maxdiffusion/pipelines/wan/wan_vace_pipeline_2_1.py
+++ b/src/maxdiffusion/pipelines/wan/wan_vace_pipeline_2_1.py
@@ -562,7 +562,6 @@ def __call__(
       if prompt is not None and isinstance(prompt, str):
         prompt = [prompt]
 
-      batch_size = len(prompt)
       if num_videos_per_prompt != 1:
         raise ValueError("Generating multiple videos per prompt is not yet supported. This may be supported in the future.")
 
@@ -573,6 +572,7 @@ def __call__(
           prompt_embeds=prompt_embeds,
           negative_prompt_embeds=negative_prompt_embeds,
       )
+      batch_size = len(prompt_embeds)
 
       transformer_dtype = self.transformer.proj_out.bias.dtype
       vace_layers = self.transformer.config.vace_layers