diffusion: fix wan2.2 ti2v num_frames adjust logic (#13379)

Co-authored-by: adarshxs <adarsh.shirawalmath@gmail.com>
Co-authored-by: Adarsh Shirawalmath <114558126+adarshxs@users.noreply.github.com>

python/sglang/multimodal_gen/configs/pipelines/base.py

@@ -37,7 +37,7 @@ class ModelTaskType(Enum):
     T2I = auto()  # Text to Image
     I2I = auto()  # Image to Image
 
-    def is_image_task(self):
+    def is_image_gen(self):
         return self == ModelTaskType.T2I or self == ModelTaskType.I2I
 
 
@@ -134,12 +134,15 @@ class PipelineConfig:
     def slice_noise_pred(self, noise, latents):
         return noise
 
-    def set_width_and_height(self, width, height, image):
+    def adjust_size(self, width, height, image):
         """
         image: input image
         """
         return width, height
 
+    def adjust_num_frames(self, num_frames):
+        return num_frames
+
     # called in ImageEncodingStage, preprocess the image
     def preprocess_image(self, image, image_processor: VaeImageProcessor):
         return image

python/sglang/multimodal_gen/configs/pipelines/qwen_image.py

@@ -273,7 +273,7 @@ class QwenImageEditPipelineConfig(QwenImagePipelineConfig):
         image = image_processor.resize(image, calculated_height, calculated_width)
         return image
 
-    def set_width_and_height(self, width, height, image):
+    def adjust_size(self, width, height, image):
         image_size = image[0].size if isinstance(image, list) else image.size
         calculated_width, calculated_height, _ = calculate_dimensions(
             1024 * 1024, image_size[0] / image_size[1]

python/sglang/multimodal_gen/configs/pipelines/wan.py

@@ -15,6 +15,9 @@ from sglang.multimodal_gen.configs.models.encoders import (
 )
 from sglang.multimodal_gen.configs.models.vaes import WanVAEConfig
 from sglang.multimodal_gen.configs.pipelines.base import ModelTaskType, PipelineConfig
+from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
+
+logger = init_logger(__name__)
 
 
 def t5_postprocess_text(outputs: BaseEncoderOutput, _text_inputs) -> torch.Tensor:
@@ -33,6 +36,22 @@ def t5_postprocess_text(outputs: BaseEncoderOutput, _text_inputs) -> torch.Tenso
     return prompt_embeds_tensor
 
 
+@dataclass
+class WanI2VCommonConfig(PipelineConfig):
+    # for all wan i2v pipelines
+    def adjust_num_frames(self, num_frames):
+        vae_scale_factor_temporal = self.vae_config.arch_config.scale_factor_temporal
+        if num_frames % vae_scale_factor_temporal != 1:
+            logger.warning(
+                f"`num_frames - 1` has to be divisible by {vae_scale_factor_temporal}. Rounding to the nearest number."
+            )
+            num_frames = (
+                num_frames // vae_scale_factor_temporal * vae_scale_factor_temporal + 1
+            )
+            return num_frames
+        return num_frames
+
+
 @dataclass
 class WanT2V480PConfig(PipelineConfig):
     """Base configuration for Wan T2V 1.3B pipeline architecture."""
@@ -81,7 +100,7 @@ class WanT2V720PConfig(WanT2V480PConfig):
 
 
 @dataclass
-class WanI2V480PConfig(WanT2V480PConfig):
+class WanI2V480PConfig(WanT2V480PConfig, WanI2VCommonConfig):
     """Base configuration for Wan I2V 14B 480P pipeline architecture."""
 
     # WanConfig-specific parameters with defaults
@@ -128,7 +147,7 @@ class FastWan2_1_T2V_480P_Config(WanT2V480PConfig):
 
 
 @dataclass
-class Wan2_2_TI2V_5B_Config(WanT2V480PConfig):
+class Wan2_2_TI2V_5B_Config(WanT2V480PConfig, WanI2VCommonConfig):
     flow_shift: float | None = 5.0
     task_type: ModelTaskType = ModelTaskType.TI2V
     expand_timesteps: bool = True