[Bugfix] Fix LongCat TeaCache CFG negative-branch guidance kwargs (#6181)

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

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",

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,

tests/model_tests/diffusion/test_alignment.py

@@ -40,7 +40,6 @@ EXCLUDED_MODELS = [
     "WanS2VPipeline",
     "WanT2VDMD2Pipeline",
     "WanI2VDMD2Pipeline",
-    "LongCatImagePipeline",
     "LongCatVideoAvatarPipeline",
     "BagelPipeline",
     "BooguImagePipeline",

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.

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,

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,