Update modeling_moss_vl.py

modeling_moss_vl.py

@@ -14,6 +14,7 @@
 # limitations under the License.
 """PyTorch MossVL model - Qwen3VL Vision + Text with Cross Attention"""
 
+import copy
 from dataclasses import dataclass
 import queue
 import threading
@@ -2160,6 +2161,7 @@ class MossVLForConditionalGeneration(MossVLPreTrainedModel, GenerationMixin):
         super().__init__(config)
         self.model = MossVLModel(config)
         self.lm_head = nn.Linear(config.text_config.hidden_size, config.text_config.vocab_size, bias=False)
+        self._offline_processor_lock = threading.RLock()
 
         self.post_init()
 
@@ -2459,7 +2461,65 @@ class MossVLForConditionalGeneration(MossVLPreTrainedModel, GenerationMixin):
 
         return [{"role": "user", "content": content}]
 
-    def _offline_prepare_input_text(self, processor, messages: List[Dict[str, Any]]) -> str:
+    @staticmethod
+    def _offline_extract_content_parts(content: Any) -> Tuple[str, List[Any], List[Any]]:
+        if isinstance(content, str):
+            return content, [], []
+        if not isinstance(content, list):
+            return (str(content) if content else ""), [], []
+
+        text_parts: List[str] = []
+        images: List[Any] = []
+        videos: List[Any] = []
+        for item in content:
+            if isinstance(item, dict):
+                if item.get("type") == "image" or "image" in item or "image_url" in item:
+                    image = item.get("image") or item.get("image_url")
+                    if image is not None:
+                        images.append(image)
+                elif item.get("type") == "video" or "video" in item or "video_path" in item:
+                    video = item.get("video") or item.get("video_path")
+                    if video is not None:
+                        videos.append(video)
+
+                if "text" in item and item["text"] is not None:
+                    text_parts.append(str(item["text"]))
+            elif isinstance(item, str):
+                text_parts.append(item)
+
+        return "".join(text_parts), images, videos
+
+    @staticmethod
+    def _offline_resolve_use_template(query: Dict[str, Any]) -> bool:
+        return bool(query.get("use_template", False))
+
+    def _offline_prepare_input_text(
+        self,
+        processor,
+        messages: List[Dict[str, Any]],
+        use_template: bool,
+    ) -> str:
+        if not use_template:
+            if any(isinstance(message, dict) and message.get("role") == "system" for message in messages):
+                raise ValueError("system messages require use_template=True")
+
+            parts = ["<|im_start|>"]
+            for message in messages:
+                role = message.get("role", "user") if isinstance(message, dict) else "user"
+                content = message.get("content", "") if isinstance(message, dict) else message
+                text, msg_images, msg_videos = self._offline_extract_content_parts(content)
+
+                if role == "user":
+                    media_tokens = ""
+                    if msg_images:
+                        media_tokens += "<|image|>" * len(msg_images)
+                    if msg_videos:
+                        media_tokens += "<|video|>" * len(msg_videos)
+                    parts.append(f"{media_tokens}{text}")
+                else:
+                    parts.append(f"{text}<|im_end|>")
+            return "".join(parts)
+
         processed_messages = []
         for message in messages:
             message_copy = dict(message)
@@ -2528,7 +2588,11 @@ class MossVLForConditionalGeneration(MossVLPreTrainedModel, GenerationMixin):
 
     def _offline_prepare_inputs(self, processor, query: Dict[str, Any]):
         messages = self._offline_prepare_messages(processor, query)
-        input_text = self._offline_prepare_input_text(processor, messages)
+        input_text = self._offline_prepare_input_text(
+            processor,
+            messages,
+            use_template=self._offline_resolve_use_template(query),
+        )
         all_images, all_videos = self._offline_collect_media(messages)
         media_kwargs = dict(query.get("media_kwargs") or {})
         processor_kwargs = self._offline_build_processor_kwargs(
@@ -2543,25 +2607,26 @@ class MossVLForConditionalGeneration(MossVLPreTrainedModel, GenerationMixin):
         modified_multi_image = False
         modified_video = False
 
-        try:
-            multi_image_max_pixels = media_kwargs.get("multi_image_max_pixels")
-            if multi_image_max_pixels is not None and image_proc is not None:
-                orig_multi_image_max_pixels = getattr(image_proc, "multi_image_max_pixels", None)
-                image_proc.multi_image_max_pixels = multi_image_max_pixels
-                modified_multi_image = True
-
-            video_max_pixels = media_kwargs.get("video_max_pixels")
-            if video_max_pixels is not None and video_proc is not None:
-                orig_video_max_pixels = getattr(video_proc, "video_max_pixels", None)
-                video_proc.video_max_pixels = video_max_pixels
-                modified_video = True
-
-            inputs = processor(**processor_kwargs)
-        finally:
-            if modified_multi_image and image_proc is not None:
-                image_proc.multi_image_max_pixels = orig_multi_image_max_pixels
-            if modified_video and video_proc is not None:
-                video_proc.video_max_pixels = orig_video_max_pixels
+        with self._offline_processor_lock:
+            try:
+                multi_image_max_pixels = media_kwargs.get("multi_image_max_pixels")
+                if multi_image_max_pixels is not None and image_proc is not None:
+                    orig_multi_image_max_pixels = getattr(image_proc, "multi_image_max_pixels", None)
+                    image_proc.multi_image_max_pixels = multi_image_max_pixels
+                    modified_multi_image = True
+
+                video_max_pixels = media_kwargs.get("video_max_pixels")
+                if video_max_pixels is not None and video_proc is not None:
+                    orig_video_max_pixels = getattr(video_proc, "video_max_pixels", None)
+                    video_proc.video_max_pixels = video_max_pixels
+                    modified_video = True
+
+                inputs = processor(**processor_kwargs)
+            finally:
+                if modified_multi_image and image_proc is not None:
+                    image_proc.multi_image_max_pixels = orig_multi_image_max_pixels
+                if modified_video and video_proc is not None:
+                    video_proc.video_max_pixels = orig_video_max_pixels
 
         text_device = self.get_input_embeddings().weight.device
         vision_device = self.visual.patch_embed.proj.weight.device
@@ -2773,7 +2838,13 @@ class MossVLForConditionalGeneration(MossVLPreTrainedModel, GenerationMixin):
                 current_session,
             )
             working_messages_list.append(working_messages)
-            input_texts.append(self._offline_prepare_input_text(processor, working_messages))
+            input_texts.append(
+                self._offline_prepare_input_text(
+                    processor,
+                    working_messages,
+                    use_template=self._offline_resolve_use_template(query),
+                )
+            )
 
             all_images, all_videos = self._offline_collect_media(working_messages)
             all_images_per_query.append(all_images)
@@ -2798,31 +2869,32 @@ class MossVLForConditionalGeneration(MossVLPreTrainedModel, GenerationMixin):
         modified_video = False
         orig_padding_side = None
 
-        try:
-            multi_image_max_pixels = media_kwargs.get("multi_image_max_pixels")
-            if multi_image_max_pixels is not None and image_proc is not None:
-                orig_multi_image_max_pixels = getattr(image_proc, "multi_image_max_pixels", None)
-                image_proc.multi_image_max_pixels = multi_image_max_pixels
-                modified_multi_image = True
-
-            video_max_pixels = media_kwargs.get("video_max_pixels")
-            if video_max_pixels is not None and video_proc is not None:
-                orig_video_max_pixels = getattr(video_proc, "video_max_pixels", None)
-                video_proc.video_max_pixels = video_max_pixels
-                modified_video = True
-
-            if tokenizer is not None and hasattr(tokenizer, "padding_side"):
-                orig_padding_side = tokenizer.padding_side
-                tokenizer.padding_side = "left"
-
-            inputs = processor(**processor_kwargs)
-        finally:
-            if modified_multi_image and image_proc is not None:
-                image_proc.multi_image_max_pixels = orig_multi_image_max_pixels
-            if modified_video and video_proc is not None:
-                video_proc.video_max_pixels = orig_video_max_pixels
-            if tokenizer is not None and orig_padding_side is not None:
-                tokenizer.padding_side = orig_padding_side
+        with self._offline_processor_lock:
+            try:
+                multi_image_max_pixels = media_kwargs.get("multi_image_max_pixels")
+                if multi_image_max_pixels is not None and image_proc is not None:
+                    orig_multi_image_max_pixels = getattr(image_proc, "multi_image_max_pixels", None)
+                    image_proc.multi_image_max_pixels = multi_image_max_pixels
+                    modified_multi_image = True
+
+                video_max_pixels = media_kwargs.get("video_max_pixels")
+                if video_max_pixels is not None and video_proc is not None:
+                    orig_video_max_pixels = getattr(video_proc, "video_max_pixels", None)
+                    video_proc.video_max_pixels = video_max_pixels
+                    modified_video = True
+
+                if tokenizer is not None and hasattr(tokenizer, "padding_side"):
+                    orig_padding_side = tokenizer.padding_side
+                    tokenizer.padding_side = "left"
+
+                inputs = processor(**processor_kwargs)
+            finally:
+                if modified_multi_image and image_proc is not None:
+                    image_proc.multi_image_max_pixels = orig_multi_image_max_pixels
+                if modified_video and video_proc is not None:
+                    video_proc.video_max_pixels = orig_video_max_pixels
+                if tokenizer is not None and orig_padding_side is not None:
+                    tokenizer.padding_side = orig_padding_side
 
         text_device = self.get_input_embeddings().weight.device
         vision_device = self.visual.patch_embed.proj.weight.device
@@ -2905,7 +2977,11 @@ class MossVLForConditionalGeneration(MossVLPreTrainedModel, GenerationMixin):
         for index, (query, session_state) in enumerate(zip(prepared_queries, session_states)):
             current_session = [] if query.get("reset_session") or query.get("clear_history") else session_state
             working_messages = self._offline_build_session_messages(processor, query, current_session)
-            input_text = self._offline_prepare_input_text(processor, working_messages)
+            input_text = self._offline_prepare_input_text(
+                processor,
+                working_messages,
+                use_template=self._offline_resolve_use_template(query),
+            )
 
             if tokenizer is not None:
                 token_ids = tokenizer(input_text, add_special_tokens=False)["input_ids"]
@@ -2972,6 +3048,210 @@ class MossVLForConditionalGeneration(MossVLPreTrainedModel, GenerationMixin):
         new_tokens = outputs[0][inputs["input_ids"].shape[1]:]
         return processor.decode(new_tokens, skip_special_tokens=True)
 
+    @staticmethod
+    def _offline_capture_processor_attrs(target, overrides: Optional[Dict[str, Any]]) -> Optional[Dict[str, Any]]:
+        if target is None or not overrides:
+            return None
+        return {name: copy.deepcopy(getattr(target, name)) for name in overrides}
+
+    @staticmethod
+    def _offline_apply_processor_attrs(target, overrides: Optional[Dict[str, Any]]) -> None:
+        if target is None or not overrides:
+            return
+        for name, value in overrides.items():
+            setattr(target, name, copy.deepcopy(value))
+
+    @staticmethod
+    def _offline_restore_processor_attrs(target, snapshot: Optional[Dict[str, Any]]) -> None:
+        if target is None or snapshot is None:
+            return
+        for name, value in snapshot.items():
+            setattr(target, name, copy.deepcopy(value))
+
+    def _offline_generate_one_with_processor_overrides(
+        self,
+        processor,
+        query: Dict[str, Any],
+        image_processor_overrides: Optional[Dict[str, Any]] = None,
+        video_processor_overrides: Optional[Dict[str, Any]] = None,
+    ) -> str:
+        image_proc = getattr(processor, "image_processor", None)
+        video_proc = getattr(processor, "video_processor", None)
+        image_snapshot = self._offline_capture_processor_attrs(image_proc, image_processor_overrides)
+        video_snapshot = self._offline_capture_processor_attrs(video_proc, video_processor_overrides)
+
+        with self._offline_processor_lock:
+            try:
+                self._offline_apply_processor_attrs(image_proc, image_processor_overrides)
+                self._offline_apply_processor_attrs(video_proc, video_processor_overrides)
+                return self._offline_generate_one(processor, query)
+            finally:
+                self._offline_restore_processor_attrs(image_proc, image_snapshot)
+                self._offline_restore_processor_attrs(video_proc, video_snapshot)
+
+    def offline_image_generate(
+        self,
+        processor,
+        prompt: str = "",
+        image: Any = None,
+        *,
+        shortest_edge: int = 4096,
+        longest_edge: int = 16777216,
+        multi_image_max_pixels: int = 201326592,
+        patch_size: int = 16,
+        temporal_patch_size: int = 1,
+        merge_size: int = 2,
+        image_mean: Optional[Union[List[float], Tuple[float, ...]]] = (0.5, 0.5, 0.5),
+        image_std: Optional[Union[List[float], Tuple[float, ...]]] = (0.5, 0.5, 0.5),
+        max_new_tokens: int = 1024,
+        temperature: float = 1.0,
+        top_k: int = 50,
+        top_p: float = 1.0,
+        repetition_penalty: float = 1.0,
+        do_sample: bool = False,
+        vision_chunked_length: int = 64,
+        use_template: bool = False,
+        thinking_mode: Optional[str] = None,
+        system_prompt_type: Optional[str] = None,
+        system_prompt: Optional[str] = None,
+    ) -> str:
+        """
+        Single-image offline generation with explicit image preprocessor defaults.
+
+        The default values mirror `preprocessor_config.json` so README examples can
+        surface the full image preprocessing setup without requiring a batch wrapper.
+        """
+        if image is None:
+            raise ValueError("`image` is required.")
+        query: Dict[str, Any] = {
+            "prompt": prompt,
+            "images": [image],
+            "videos": [],
+            "media_kwargs": {
+                "min_pixels": shortest_edge,
+                "max_pixels": longest_edge,
+                "multi_image_max_pixels": multi_image_max_pixels,
+            },
+            "generate_kwargs": {
+                "max_new_tokens": max_new_tokens,
+                "temperature": temperature,
+                "top_k": top_k,
+                "top_p": top_p,
+                "repetition_penalty": repetition_penalty,
+                "do_sample": do_sample,
+                "vision_chunked_length": vision_chunked_length,
+            },
+            "use_template": use_template,
+        }
+        if thinking_mode is not None:
+            query["thinking_mode"] = thinking_mode
+        if system_prompt_type is not None:
+            query["system_prompt_type"] = system_prompt_type
+        if system_prompt is not None:
+            query["system_prompt"] = system_prompt
+
+        image_processor_overrides = {
+            "size": {"shortest_edge": shortest_edge, "longest_edge": longest_edge},
+            "multi_image_max_pixels": multi_image_max_pixels,
+            "patch_size": patch_size,
+            "temporal_patch_size": temporal_patch_size,
+            "merge_size": merge_size,
+            "image_mean": list(image_mean) if image_mean is not None else None,
+            "image_std": list(image_std) if image_std is not None else None,
+        }
+        return self._offline_generate_one_with_processor_overrides(
+            processor,
+            query,
+            image_processor_overrides=image_processor_overrides,
+        )
+
+    def offline_video_generate(
+        self,
+        processor,
+        prompt: str = "",
+        video: Any = None,
+        *,
+        shortest_edge: int = 4096,
+        longest_edge: int = 16777216,
+        video_max_pixels: int = 201326592,
+        patch_size: int = 16,
+        temporal_patch_size: int = 1,
+        merge_size: int = 2,
+        video_fps: float = 1.0,
+        min_frames: int = 1,
+        max_frames: int = 256,
+        num_extract_threads: int = 4,
+        image_mean: Optional[Union[List[float], Tuple[float, ...]]] = (0.5, 0.5, 0.5),
+        image_std: Optional[Union[List[float], Tuple[float, ...]]] = (0.5, 0.5, 0.5),
+        max_new_tokens: int = 1024,
+        temperature: float = 1.0,
+        top_k: int = 50,
+        top_p: float = 1.0,
+        repetition_penalty: float = 1.0,
+        do_sample: bool = False,
+        vision_chunked_length: int = 64,
+        use_template: bool = False,
+        thinking_mode: Optional[str] = None,
+        system_prompt_type: Optional[str] = None,
+        system_prompt: Optional[str] = None,
+    ) -> str:
+        """
+        Single-video offline generation with explicit video preprocessor defaults.
+
+        The default values mirror `video_preprocessor_config.json` so README examples
+        can show a standalone video entry point with the effective preprocessing knobs.
+        """
+        if video is None:
+            raise ValueError("`video` is required.")
+        query: Dict[str, Any] = {
+            "prompt": prompt,
+            "images": [],
+            "videos": [video],
+            "media_kwargs": {
+                "min_pixels": shortest_edge,
+                "max_pixels": longest_edge,
+                "video_max_pixels": video_max_pixels,
+                "video_fps": video_fps,
+                "min_frames": min_frames,
+                "max_frames": max_frames,
+            },
+            "generate_kwargs": {
+                "max_new_tokens": max_new_tokens,
+                "temperature": temperature,
+                "top_k": top_k,
+                "top_p": top_p,
+                "repetition_penalty": repetition_penalty,
+                "do_sample": do_sample,
+                "vision_chunked_length": vision_chunked_length,
+            },
+            "use_template": use_template,
+        }
+        if thinking_mode is not None:
+            query["thinking_mode"] = thinking_mode
+        if system_prompt_type is not None:
+            query["system_prompt_type"] = system_prompt_type
+        if system_prompt is not None:
+            query["system_prompt"] = system_prompt
+
+        video_processor_overrides = {
+            "size": {"shortest_edge": shortest_edge, "longest_edge": longest_edge},
+            "video_max_pixels": video_max_pixels,
+            "patch_size": patch_size,
+            "temporal_patch_size": temporal_patch_size,
+            "merge_size": merge_size,
+            "video_fps": video_fps,
+            "min_frames": min_frames,
+            "max_frames": max_frames,
+            "num_extract_threads": num_extract_threads,
+            "image_mean": list(image_mean) if image_mean is not None else None,
+            "image_std": list(image_std) if image_std is not None else None,
+        }
+        return self._offline_generate_one_with_processor_overrides(
+            processor,
+            query,
+            video_processor_overrides=video_processor_overrides,
+        )
+
     def offline_generate(
         self,
         processor,
@@ -2990,6 +3270,8 @@ class MossVLForConditionalGeneration(MossVLPreTrainedModel, GenerationMixin):
         - `prompt` / `messages`
         - `images` / `videos`
         - `media_kwargs` / `generate_kwargs`
+        - `use_template` to switch between backend-style pretrain prompting
+          (`False`, default for base) and tokenizer chat template prompting (`True`)
         - `thinking_mode` (`no_thinking` or `deep_thinking`, plus compatible aliases)
         - `system_prompt_type` (`text_image` or `video`, plus compatible aliases)
         - `system_prompt` for an explicit override