diff --git a/tests/model_tests/diffusion/diff_model_builders.py b/tests/model_tests/diffusion/diff_model_builders.py index 394d7a3e1..c5eadecdb 100644 --- a/tests/model_tests/diffusion/diff_model_builders.py +++ b/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", diff --git a/tests/model_tests/diffusion/model_settings.py b/tests/model_tests/diffusion/model_settings.py index e671522cd..f24564f62 100644 --- a/tests/model_tests/diffusion/model_settings.py +++ b/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, diff --git a/tests/model_tests/diffusion/test_alignment.py b/tests/model_tests/diffusion/test_alignment.py index 6c35b231a..9d4b5a1ad 100644 --- a/tests/model_tests/diffusion/test_alignment.py +++ b/tests/model_tests/diffusion/test_alignment.py @@ -47,7 +47,6 @@ EXCLUDED_MODELS = [ "LancePipeline", "MingImagePipeline", "InternVLAA1Pipeline", - "LongCatImageEditPipeline", "StableDiffusion3Pipeline", "HunyuanImage3ForCausalMM", "ErnieImagePipeline", diff --git a/vllm_omni/diffusion/models/longcat_image/pipeline_longcat_image_edit.py b/vllm_omni/diffusion/models/longcat_image/pipeline_longcat_image_edit.py index 8622ccc8a..ac08b8def 100644 --- a/vllm_omni/diffusion/models/longcat_image/pipeline_longcat_image_edit.py +++ b/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: