diff --git a/python/sglang/multimodal_gen/configs/sample/sampling_params.py b/python/sglang/multimodal_gen/configs/sample/sampling_params.py index 36b19ac0a..746d38bdd 100644 --- a/python/sglang/multimodal_gen/configs/sample/sampling_params.py +++ b/python/sglang/multimodal_gen/configs/sample/sampling_params.py @@ -156,7 +156,7 @@ class SamplingParams: suppress_logs: bool = False return_file_paths_only: bool = True - enable_sequence_shard: bool = False + enable_sequence_shard: bool | None = None def _set_output_file_ext(self): # add extension if needed @@ -390,10 +390,20 @@ class SamplingParams: ) logger.warning(error_msg) + pipeline_name_lower = server_args.pipeline_config.__class__.__name__.lower() + + if "wan" in pipeline_name_lower and ( + self.enable_sequence_shard is None or self.enable_sequence_shard + ): + self.enable_sequence_shard = True + logger.debug("Automatically enabled enable_sequence_shard") + else: + self.enable_sequence_shard = False + if self.enable_sequence_shard: self.adjust_frames = False logger.info( - f"Sequence dimension shard is enabled, disabling frame adjustment" + f"Sequence dimension shard is enabled, disabling frame adjustment for better performance" ) if pipeline_config.task_type.is_image_gen(): @@ -402,52 +412,66 @@ class SamplingParams: logger.debug(f"Setting `num_frames` to 1 for image generation model") self.num_frames = 1 - elif self.adjust_frames: + else: + # mandatory frame adjusting logic, mod # NOTE: We must apply adjust_num_frames BEFORE the SP alignment logic below. # If we apply it after, adjust_num_frames might modify the frame count # and break the divisibility constraint (alignment) required by num_gpus. + original_num_frames = self.num_frames self.num_frames = server_args.pipeline_config.adjust_num_frames( - self.num_frames + original_num_frames + ) + logger.info( + "Adjusting number of frames from %s to %s based on model", + original_num_frames, + self.num_frames, ) - # Adjust number of frames based on number of GPUs for video task - use_temporal_scaling_frames = ( - pipeline_config.vae_config.use_temporal_scaling_frames - ) - num_frames = self.num_frames - num_gpus = server_args.num_gpus - temporal_scale_factor = ( - pipeline_config.vae_config.arch_config.temporal_compression_ratio - ) - - if use_temporal_scaling_frames: - orig_latent_num_frames = (num_frames - 1) // temporal_scale_factor + 1 - - if orig_latent_num_frames % server_args.num_gpus != 0: - # Adjust latent frames to be divisible by number of GPUs - if self.num_frames_round_down: - # Ensure we have at least 1 batch per GPU - new_latent_num_frames = ( - max(1, (orig_latent_num_frames // num_gpus)) * num_gpus - ) - else: - new_latent_num_frames = ( - math.ceil(orig_latent_num_frames / num_gpus) * num_gpus - ) + if self.adjust_frames: + # Adjust number of frames based on number of GPUs for video task + use_temporal_scaling_frames = ( + pipeline_config.vae_config.use_temporal_scaling_frames + ) + num_frames = self.num_frames + num_gpus = server_args.num_gpus + temporal_scale_factor = ( + pipeline_config.vae_config.arch_config.temporal_compression_ratio + ) if use_temporal_scaling_frames: - # Convert back to number of frames, ensuring num_frames-1 is a multiple of temporal_scale_factor - new_num_frames = ( - new_latent_num_frames - 1 - ) * temporal_scale_factor + 1 + orig_latent_num_frames = ( + num_frames - 1 + ) // temporal_scale_factor + 1 + else: + orig_latent_num_frames = num_frames - logger.info( - "Adjusting number of frames from %s to %s based on number of GPUs (%s)", - self.num_frames, - new_num_frames, - server_args.num_gpus, - ) - self.num_frames = new_num_frames + if orig_latent_num_frames % server_args.num_gpus != 0: + # Adjust latent frames to be divisible by number of GPUs + if self.num_frames_round_down: + # Ensure we have at least 1 batch per GPU + new_latent_num_frames = ( + max(1, (orig_latent_num_frames // num_gpus)) * num_gpus + ) + else: + new_latent_num_frames = ( + math.ceil(orig_latent_num_frames / num_gpus) * num_gpus + ) + + if use_temporal_scaling_frames: + # Convert back to number of frames, ensuring num_frames-1 is a multiple of temporal_scale_factor + new_num_frames = ( + new_latent_num_frames - 1 + ) * temporal_scale_factor + 1 + else: + new_num_frames = new_latent_num_frames + + logger.info( + "Adjusting number of frames from %s to %s based on number of GPUs (%s)", + self.num_frames, + new_num_frames, + server_args.num_gpus, + ) + self.num_frames = new_num_frames if not server_args.comfyui_mode: self._set_output_file_name() diff --git a/python/sglang/multimodal_gen/configs/sample/wan.py b/python/sglang/multimodal_gen/configs/sample/wan.py index efe86c4b1..2c405b2f0 100644 --- a/python/sglang/multimodal_gen/configs/sample/wan.py +++ b/python/sglang/multimodal_gen/configs/sample/wan.py @@ -239,8 +239,8 @@ class Wan2_2_T2V_A14B_SamplingParam(Wan2_2_Base_SamplingParams): guidance_scale_2: float = 3.0 # low_noise num_inference_steps: int = 40 fps: int = 16 - # NOTE(will): default boundary timestep is tracked by PipelineConfig, but - # can be overridden during sampling + + num_frames: int = 81 # Wan2.2 T2V A14B supported resolutions supported_resolutions: list[tuple[int, int]] | None = field( @@ -259,8 +259,8 @@ class Wan2_2_I2V_A14B_SamplingParam(Wan2_2_Base_SamplingParams): guidance_scale_2: float = 3.5 # low_noise num_inference_steps: int = 40 fps: int = 16 - # NOTE(will): default boundary timestep is tracked by PipelineConfig, but - # can be overridden during sampling + + num_frames: int = 81 # Wan2.2 I2V A14B supported resolutions supported_resolutions: list[tuple[int, int]] | None = field(