[Model] Migrate MammothModa2-Preview example to standard task example + model_extras (#4691)

Signed-off-by: mo saif <moisf56@gmail.com>
Co-authored-by: Claude <noreply@anthropic.com>

vllm_omni/diffusion/models/mammoth_moda2/pipeline_mammothmoda2_dit.py

@@ -142,6 +142,43 @@ class MammothModa2DiTPipeline(nn.Module, SupportsComponentDiscovery):
             dtype=dtype,
         )
 
+    def _split_ar_conditions(
+        self,
+        *,
+        full_hidden_states: torch.Tensor,
+        full_token_ids: list[int],
+        answer_start_index: int,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        """Split AR-stage hidden states into text / image condition embeds.
+
+        The token ids that distinguish question (text) tokens, generated visual
+        tokens, and multi-modal placeholder tokens are read from the model config
+        (``gen_vocab_start_index`` and the vision placeholder token ids), so the
+        caller no longer needs to pass them. Mirrors the masking the bespoke
+        MammothModa2 example performed via ar2dit.
+        """
+        gen_vocab_start_index = int(self.config.llm_config.gen_vocab_start_index)
+        visual_ids = [
+            int(self.config.image_token_id),
+            int(self.config.video_token_id),
+            int(self.config.vision_start_token_id),
+            int(self.config.vision_end_token_id),
+        ]
+
+        device = full_hidden_states.device
+        token_ids = torch.tensor(full_token_ids, dtype=torch.long, device=device)
+        positions = torch.arange(token_ids.shape[0], device=device)
+        questions_mask = positions < answer_start_index
+        answers_mask = ~questions_mask
+        gen_token_mask = token_ids >= gen_vocab_start_index
+        visual_token_mask = torch.isin(token_ids, torch.tensor(visual_ids, dtype=torch.long, device=device))
+        text_mask = questions_mask & ~(visual_token_mask | gen_token_mask)
+        image_mask = answers_mask & gen_token_mask
+
+        text_cond = full_hidden_states[text_mask].to(dtype=torch.float32).contiguous()
+        image_cond = full_hidden_states[image_mask].to(dtype=torch.float32).contiguous()
+        return text_cond, image_cond
+
     @torch.inference_mode()
     def forward(
         self,
@@ -151,14 +188,35 @@ class MammothModa2DiTPipeline(nn.Module, SupportsComponentDiscovery):
     ) -> OmniOutput:
         runtime_addi = kwargs.get("runtime_additional_information", None)
         info = runtime_addi[0]
-        text_cond = info["text_prompt_embeds"]
-        image_cond = info["image_prompt_embeds"]
+
+        # Sampling knobs are declared in vllm_omni/model_extras/mammothmodal2_preview.py
+        # and routed via extra_body -> sampling_params.extra_args (surfaced here as
+        # ``sampling_extra_args``). Fall back to runtime_additional_information for the
+        # legacy bespoke-example path during the transition.
+        extra_args_list = kwargs.get("sampling_extra_args") or []
+        extra_args = extra_args_list[0] if extra_args_list else {}
+        text_guidance_scale = float(extra_args.get("text_guidance_scale", info["text_guidance_scale"][0]))
+        cfg_range_val = extra_args.get("cfg_range", info["cfg_range"])
+        cfg_range = float(cfg_range_val[0]), float(cfg_range_val[1])
+        num_inference_steps = int(extra_args.get("num_inference_steps", info["num_inference_steps"][0]))
+
         negative_cond = info.get("negative_prompt_embeds")
         negative_attention_mask = info.get("negative_prompt_attention_mask")
         image_hw = info["image_height"][0], info["image_width"][0]
-        text_guidance_scale = info["text_guidance_scale"][0]
-        cfg_range = info["cfg_range"][0], info["cfg_range"][1]
-        num_inference_steps = info["num_inference_steps"][0]
+
+        # Split the AR hidden states into text / image conditions. The token ids that
+        # drive the split are sourced from the model config (see _split_ar_conditions),
+        # formerly supplied by the bespoke example via additional_information. Legacy
+        # fallback: ar2dit may have already produced the split conditions.
+        if "text_prompt_embeds" in info:
+            text_cond = info["text_prompt_embeds"]
+            image_cond = info["image_prompt_embeds"]
+        else:
+            text_cond, image_cond = self._split_ar_conditions(
+                full_hidden_states=info["full_hidden_states"],
+                full_token_ids=info["full_token_ids"],
+                answer_start_index=int(info["answer_start_index"][0]),
+            )
 
         # Move to model device/dtype.
         model_device = next(self.parameters()).device