diff --git a/tests/model_tests/diffusion/diff_model_builders.py b/tests/model_tests/diffusion/diff_model_builders.py index e47ea1895..51881ddf3 100644 --- a/tests/model_tests/diffusion/diff_model_builders.py +++ b/tests/model_tests/diffusion/diff_model_builders.py @@ -54,14 +54,18 @@ def _shrink_flux_t5_text_encoder(config: dict) -> dict: def _shrink_qwen_text_encoder_config(config: dict, hidden_size: int = 64) -> dict: - text_config = config["text_config"] + """Shrink a Qwen2.5-VL text encoder (nested Qwen-Image or flat LongCat layout).""" + # Nested: text fields live under text_config, with top-level mirrors. + # Flat: text fields live at the top level. + text_config = config["text_config"] if "text_config" in config else config old_head_dim = text_config["hidden_size"] / text_config["num_attention_heads"] text_config["num_hidden_layers"] = 2 text_config["intermediate_size"] = 64 text_config["hidden_size"] = hidden_size text_config["num_attention_heads"] = 2 text_config["num_key_value_heads"] = 2 - text_config["layer_types"] = text_config["layer_types"][:2] + if "layer_types" in text_config: + text_config["layer_types"] = text_config["layer_types"][:2] factor = old_head_dim / (hidden_size / 2) mrope_section = text_config["rope_scaling"]["mrope_section"] text_config["rope_scaling"]["mrope_section"] = [round(d / factor) for d in mrope_section] @@ -105,6 +109,22 @@ def tiny_qwen_image_edit_plus_builder() -> str: ) +def tiny_longcat_image_builder() -> str: + return build_tiny_from_configs( + "LongCatImagePipeline", + "meituan-longcat/LongCat-Image", + 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 cbb37729a..399a3d765 100644 --- a/tests/model_tests/diffusion/model_settings.py +++ b/tests/model_tests/diffusion/model_settings.py @@ -54,6 +54,16 @@ DIFFUSION_TEST_SETTINGS = { builder=diff_model_builders.tiny_qwen_image_builder, supported_tasks=[DiffusionTasks.TEXT_TO_IMAGE], ), + "LongCatImagePipeline": DiffusionModelTestOpts( + model="meituan-longcat/LongCat-Image", + builder=diff_model_builders.tiny_longcat_image_builder, + supported_tasks=[DiffusionTasks.TEXT_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 1b40b1e5f..89913a073 100644 --- a/tests/model_tests/diffusion/test_alignment.py +++ b/tests/model_tests/diffusion/test_alignment.py @@ -40,7 +40,6 @@ EXCLUDED_MODELS = [ "WanS2VPipeline", "WanT2VDMD2Pipeline", "WanI2VDMD2Pipeline", - "LongCatImagePipeline", "LongCatVideoAvatarPipeline", "BagelPipeline", "BooguImagePipeline", diff --git a/vllm_omni/diffusion/cache/teacache/extractors.py b/vllm_omni/diffusion/cache/teacache/extractors.py index dea36fff5..db4427958 100644 --- a/vllm_omni/diffusion/cache/teacache/extractors.py +++ b/vllm_omni/diffusion/cache/teacache/extractors.py @@ -723,11 +723,11 @@ def extract_flux2_klein_context( def extract_longcat_context( module: nn.Module, # LongCatImageTransformer2DModel hidden_states, - timestep, - guidance, - encoder_hidden_states, - txt_ids, - img_ids, + encoder_hidden_states=None, + timestep=None, + img_ids=None, + txt_ids=None, + guidance=None, **kwargs, ) -> CacheContext: """Extract the cache context for LongCat Image. diff --git a/vllm_omni/diffusion/models/longcat_image/pipeline_longcat_image.py b/vllm_omni/diffusion/models/longcat_image/pipeline_longcat_image.py index b95f59228..112426002 100644 --- a/vllm_omni/diffusion/models/longcat_image/pipeline_longcat_image.py +++ b/vllm_omni/diffusion/models/longcat_image/pipeline_longcat_image.py @@ -233,7 +233,7 @@ class LongCatImagePipeline(nn.Module, CFGParallelMixin, DiffusionPipelineProfile # See ``hub_prefetch.py`` for the transformers v5 multi-worker subfolder # race; prefetch the whole component set before any from_pretrained. - longcat_subfolders = ["scheduler", "text_encoder", "tokenizer", "vae"] + longcat_subfolders = ["scheduler", "text_encoder", "text_processor", "tokenizer", "vae"] prefetch_subfolders(model, longcat_subfolders, local_files_only=local_files_only) self.scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained( @@ -248,7 +248,7 @@ class LongCatImagePipeline(nn.Module, CFGParallelMixin, DiffusionPipelineProfile local_files_only=local_files_only, ) self.text_processor = Qwen2VLProcessor.from_pretrained( - model, subfolder="tokenizer", local_files_only=local_files_only + model, subfolder="text_processor", local_files_only=local_files_only ) self.vae = from_pretrained_with_prefetch( AutoencoderKL.from_pretrained, @@ -627,6 +627,7 @@ class LongCatImagePipeline(nn.Module, CFGParallelMixin, DiffusionPipelineProfile negative_kwargs = { "hidden_states": latents, "timestep": timestep / 1000, + "guidance": guidance, "encoder_hidden_states": negative_prompt_embeds, "txt_ids": negative_text_ids, "img_ids": latent_image_ids, 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 4821425b0..8622ccc8a 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 @@ -667,6 +667,7 @@ class LongCatImageEditPipeline( negative_kwargs = { "hidden_states": latent_model_input, "timestep": timestep / 1000, + "guidance": guidance, "encoder_hidden_states": negative_prompt_embeds, "txt_ids": negative_text_ids, "img_ids": latent_image_ids,