[Bugfix] Honor height/width overrides in LongCatImageEditPipeline (#6222)

Signed-off-by: Yifan Zong <yzong@redhat.com>

tests/model_tests/diffusion/diff_model_builders.py

@@ -200,6 +200,22 @@ def tiny_longcat_image_builder() -> str:
     )
 
 
+def tiny_longcat_image_edit_builder() -> str:
+    return build_tiny_from_configs(
+        "LongCatImageEditPipeline",
+        "meituan-longcat/LongCat-Image-Edit",
+        transform={
+            "text_encoder": _shrink_qwen_text_encoder_config,
+            "transformer": partial(
+                _shrink_dit_rope_config,
+                num_single_layers=2,
+                default_axes_dims_rope=[16, 56, 56],
+                joint_attention_dim=64,
+            ),
+        },
+    )
+
+
 def tiny_flux_builder() -> str:
     return build_tiny_from_configs(
         "FluxPipeline",

tests/model_tests/diffusion/model_settings.py

@@ -70,6 +70,16 @@ DIFFUSION_TEST_SETTINGS = {
             [DiffusionAccs.CFG_PARALLEL, DiffusionAccs.TENSOR_PARALLEL, DiffusionAccs.CPU_OFFLOAD],
         ],
     ),
+    "LongCatImageEditPipeline": DiffusionModelTestOpts(
+        model="meituan-longcat/LongCat-Image-Edit",
+        builder=diff_model_builders.tiny_longcat_image_edit_builder,
+        supported_tasks=[DiffusionTasks.IMAGE_TO_IMAGE],
+        extra_test_groups=[
+            [DiffusionAccs.TEA_CACHE],
+            [DiffusionAccs.SEQUENCE_PARALLEL, DiffusionAccs.CACHE_DIT, DiffusionAccs.LAYERWISE_OFFLOAD],
+            [DiffusionAccs.CFG_PARALLEL, DiffusionAccs.TENSOR_PARALLEL, DiffusionAccs.CPU_OFFLOAD],
+        ],
+    ),
     "FluxPipeline": DiffusionModelTestOpts(
         model="black-forest-labs/FLUX.1-schnell",
         builder=diff_model_builders.tiny_flux_builder,

tests/model_tests/diffusion/test_alignment.py

@@ -47,7 +47,6 @@ EXCLUDED_MODELS = [
     "LancePipeline",
     "MingImagePipeline",
     "InternVLAA1Pipeline",
-    "LongCatImageEditPipeline",
     "StableDiffusion3Pipeline",
     "HunyuanImage3ForCausalMM",
     "ErnieImagePipeline",

vllm_omni/diffusion/models/longcat_image/pipeline_longcat_image_edit.py

@@ -37,6 +37,7 @@ from vllm_omni.diffusion.models.longcat_image.longcat_image_transformer import (
 from vllm_omni.diffusion.models.longcat_image.pipeline_longcat_image import calculate_shift
 from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin
 from vllm_omni.diffusion.request import OmniDiffusionRequest
+from vllm_omni.diffusion.utils.size_utils import normalize_min_aligned_size
 from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
 from vllm_omni.inputs.data import OmniTextPrompt
 from vllm_omni.model_executor.model_loader.weight_utils import (
@@ -93,17 +94,16 @@ def get_longcat_image_edit_pre_process_func(
         height = request.sampling_params.height or calculated_height
         width = request.sampling_params.width or calculated_width
 
-        # Store calculated dimensions in request
-        prompt["additional_information"]["calculated_height"] = calculated_height
-        prompt["additional_information"]["calculated_width"] = calculated_width
+        # Ensure dimensions are multiples of vae_scale_factor * 2
+        height, width = normalize_min_aligned_size(height, width, vae_scale_factor * 2)
         request.sampling_params.height = height
         request.sampling_params.width = width
 
         # Preprocess image
         if image is not None and not (isinstance(image, torch.Tensor) and image.size(1) == latent_channels):
-            image = image_processor.resize(image, calculated_height, calculated_width)
-            prompt_image = image_processor.resize(image, calculated_height // 2, calculated_width // 2)
-            image = image_processor.preprocess(image, calculated_height, calculated_width)
+            image = image_processor.resize(image, height, width)
+            prompt_image = image_processor.resize(image, height // 2, width // 2)
+            image = image_processor.preprocess(image, height, width)
 
             # Store preprocessed image and prompt image in request
             prompt["additional_information"]["preprocessed_image"] = image
@@ -575,15 +575,13 @@ class LongCatImageEditPipeline(
         ):
             prompt_image = additional_information.get("prompt_image")
             image = additional_information.get("preprocessed_image")
-            calculated_height = additional_information.get("calculated_height", height)
-            calculated_width = additional_information.get("calculated_width", width)
         else:
             raise RuntimeError("Missing preprocess image that should have been created by the preprocess function.")
 
         self.check_inputs(
             prompt,
-            calculated_height,
-            calculated_width,
+            height,
+            width,
             negative_prompt=negative_prompt,
             prompt_embeds=prompt_embeds,
             negative_prompt_embeds=negative_prompt_embeds,
@@ -610,8 +608,8 @@ class LongCatImageEditPipeline(
             image,
             batch_size * num_images_per_prompt,
             num_channels_latents,
-            calculated_height,
-            calculated_width,
+            height,
+            width,
             prompt_embeds.dtype,
             prompt_embeds.shape[1],
             device,
@@ -692,7 +690,7 @@ class LongCatImageEditPipeline(
         if output_type == "latent":
             image = latents
         else:
-            latents = self._unpack_latents(latents, calculated_height, calculated_width, self.vae_scale_factor)
+            latents = self._unpack_latents(latents, height, width, self.vae_scale_factor)
             latents = (latents / self.vae.config.scaling_factor) + self.vae.config.shift_factor
 
             if latents.dtype != self.vae.dtype: