xinference 1.6.0.post1__py3-none-any.whl → 1.6.1__py3-none-any.whl

xinference/model/llm/transformers/{ovis2.py → multimodal/ovis2.py}

@@ -1,4 +1,4 @@
-# Copyright 2022-2023 XProbe Inc.
+# Copyright 2022-2025 XProbe Inc.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -12,33 +12,26 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import logging
-import uuid
-from typing import Dict, Iterator, List, Optional, Union
+from threading import Thread
+from typing import Any, Dict, Iterator, List, Tuple
 
 import torch
 from PIL import Image
 
-from ....types import (
-    ChatCompletion,
-    ChatCompletionChunk,
-    ChatCompletionMessage,
-    CompletionChunk,
-)
-from ..llm_family import LLMFamilyV1, LLMSpecV1
-from ..utils import generate_chat_completion, generate_completion_chunk
-from .core import PytorchChatModel, PytorchGenerateConfig
-from .utils import cache_clean
+from ...llm_family import LLMFamilyV1, LLMSpecV1, register_transformer
+from ..core import register_non_default_model
+from .core import PytorchMultiModalModel
 
 logger = logging.getLogger(__name__)
 
 
-class Ovis2ChatModel(PytorchChatModel):
-    def __init__(self, *args, **kwargs):
-        super().__init__(*args, **kwargs)
-        self._tokenizer = None
-        self._model = None
-        self._device = None
-        self._processor = None
+@register_transformer
+@register_non_default_model("Ovis2")
+class Ovis2ChatModel(PytorchMultiModalModel):
+    def __init__(self, *args, **kws):
+        super().__init__(*args, **kws)
+        self._text_tokenizer = None
+        self._visual_tokenizer = None
 
     @classmethod
     def match_json(
@@ -51,127 +44,28 @@ class Ovis2ChatModel(PytorchChatModel):
             return True
         return False
 
-    def load(self):
+    def decide_device(self):
+        pass
+
+    def load_processor(self):
+        pass
+
+    def load_multimodal_model(self):
         from transformers import AutoModelForCausalLM
 
-        # load model
+        kwargs = self.apply_bnb_quantization()
         self._model = AutoModelForCausalLM.from_pretrained(
             self.model_path,
             torch_dtype=torch.bfloat16,
             multimodal_max_length=32768,
             trust_remote_code=True,
+            **kwargs,
         ).cuda()
         self._text_tokenizer = self._model.get_text_tokenizer()
         self._visual_tokenizer = self._model.get_visual_tokenizer()
 
-    @cache_clean
-    def chat(
-        self,
-        messages: List[ChatCompletionMessage],  # type: ignore
-        generate_config: Optional[PytorchGenerateConfig] = None,
-    ) -> Union[ChatCompletion, Iterator[ChatCompletionChunk]]:
-        messages = self._transform_messages(messages)
-
-        generate_config = generate_config if generate_config else {}
-
-        stream = generate_config.get("stream", False) if generate_config else False
-
-        if stream:
-            # raise NotImplementedError("Stream is not supported for Ovis2 model.")
-            it = self._generate_stream(messages, generate_config)
-            return self._to_chat_completion_chunks(it)
-        else:
-            c = self._generate(messages, generate_config)
-            return c
-
-    def _generate(
-        self, messages: List, config: PytorchGenerateConfig = {}
-    ) -> ChatCompletion:
-        input_ids, attention_mask, pixel_values, gen_kwargs = self._generate_chat_data(
-            messages, config
-        )
-
-        # generate output
-        with torch.inference_mode():
-            gen_kwargs.update(
-                dict(
-                    pixel_values=pixel_values,
-                    attention_mask=attention_mask,
-                )
-            )
-
-            output_ids = self._model.generate(
-                input_ids,
-                **gen_kwargs,
-            )[0]
-            output = self._text_tokenizer.decode(output_ids, skip_special_tokens=True)
-        return generate_chat_completion(self.model_uid, output)
-
-    def _generate_stream(
-        self, messages: List, config: PytorchGenerateConfig = {}
-    ) -> Iterator[CompletionChunk]:
-        from threading import Thread
-
-        from transformers import TextIteratorStreamer
-
-        input_ids, attention_mask, pixel_values, gen_kwargs = self._generate_chat_data(
-            messages, config
-        )
-
-        _, inputs_embeds, _, attention_mask = self._model.merge_multimodal(
-            text_input_ids=input_ids,
-            text_attention_masks=attention_mask,
-            text_labels=None,
-            pixel_values=pixel_values,
-            left_padding=True,
-        )
-
-        streamer = TextIteratorStreamer(
-            self._text_tokenizer, timeout=60, skip_prompt=True, skip_special_tokens=True
-        )
-
-        gen_kwargs.update(
-            dict(
-                inputs_embeds=inputs_embeds,
-                attention_mask=attention_mask,
-                streamer=streamer,
-            )
-        )
-
-        inputs_embeds = inputs_embeds.detach()
-        torch.cuda.empty_cache()
-
-        thread = Thread(target=self._model.llm.generate, kwargs=gen_kwargs)
-        thread.start()
-
-        completion_id = str(uuid.uuid1())
-
-        for new_text in streamer:
-            yield generate_completion_chunk(
-                chunk_text=new_text,
-                finish_reason=None,
-                chunk_id=completion_id,
-                model_uid=self.model_uid,
-                prompt_tokens=-1,
-                completion_tokens=-1,
-                total_tokens=-1,
-                has_choice=True,
-                has_content=True,
-            )
-
-        yield generate_completion_chunk(
-            chunk_text=None,
-            finish_reason="stop",
-            chunk_id=completion_id,
-            model_uid=self.model_uid,
-            prompt_tokens=-1,
-            completion_tokens=-1,
-            total_tokens=-1,
-            has_choice=True,
-            has_content=False,
-        )
-
-    def parse_messages_ovis(self, messages: List[Dict]) -> List[Dict]:
+    @staticmethod
+    def _parse_messages_ovis(messages: List[Dict]) -> List[Dict]:
         ovis_msgs = []
         for mess in messages:
             contents = mess["content"]
@@ -189,12 +83,52 @@ class Ovis2ChatModel(PytorchChatModel):
 
         return ovis_msgs
 
-    def _generate_chat_data(
-        self, messages: List[Dict], config: PytorchGenerateConfig = {}
-    ):
+    @staticmethod
+    def _convert_video_tensors_to_pil(video_inputs: List) -> List[Image.Image]:
+        """Convert video tensors to a list of PIL images"""
+        from torchvision import transforms
+
+        to_pil = transforms.ToPILImage()
+        pil_images = []
+
+        for video_tensor_4d in video_inputs:
+            if isinstance(video_tensor_4d, torch.Tensor):
+                # Verify it's a 4D tensor
+                if video_tensor_4d.ndim == 4:
+                    # Iterate through the first dimension (frames) of 4D tensor
+                    for i in range(video_tensor_4d.size(0)):
+                        frame_tensor_3d = video_tensor_4d[
+                            i
+                        ]  # Get 3D frame tensor [C, H, W]
+                        # Ensure tensor is on CPU before conversion
+                        if frame_tensor_3d.is_cuda:
+                            frame_tensor_3d = frame_tensor_3d.cpu()
+                        try:
+                            pil_image = to_pil(frame_tensor_3d)
+                            pil_images.append(pil_image)
+                        except Exception as e:
+                            logger.error(
+                                f"Error converting frame {i} to PIL Image: {e}"
+                            )
+                            # Can choose to skip this frame or handle error differently
+                else:
+                    logger.warning(
+                        f"Expected 4D tensor in video_inputs, but got {video_tensor_4d.ndim}D. Skipping this tensor."
+                    )
+            elif isinstance(video_tensor_4d, Image.Image):
+                # If fetch_video returns Image list, add directly
+                pil_images.append(video_tensor_4d)
+            else:
+                logger.warning(
+                    f"Unexpected type in video_inputs: {type(video_tensor_4d)}. Skipping."
+                )
+
+        return pil_images
+
+    def _generate_chat_data(self, messages: List[Dict]):
         from qwen_vl_utils import process_vision_info
 
-        messages_ovis = self.parse_messages_ovis(messages)
+        messages_ovis = self._parse_messages_ovis(messages)
         max_partition = None
         prompt = messages_ovis[-1]["value"]
 
@@ -246,57 +180,62 @@ class Ovis2ChatModel(PytorchChatModel):
             )
         pixel_values = [pixel_values]
 
-        gen_kwargs = dict(
-            max_new_tokens=config.get("max_tokens", 1024),
+        return input_ids, attention_mask, pixel_values
+
+    def build_generate_kwargs(
+        self,
+        generate_config: Dict,
+    ) -> Dict[str, Any]:
+        return dict(
+            max_new_tokens=generate_config.get("max_tokens", 1024),
             do_sample=False,
             top_p=None,
             top_k=None,
-            temperature=config.get("temperature", None),
+            temperature=generate_config.get("temperature", None),
             repetition_penalty=None,
             eos_token_id=self._model.generation_config.eos_token_id,
             pad_token_id=self._text_tokenizer.pad_token_id,
             use_cache=True,
         )
 
-        return input_ids, attention_mask, pixel_values, gen_kwargs
+    def build_inputs_from_messages(
+        self,
+        messages: List[Dict],
+        generate_config: Dict,
+    ):
+        msgs = self._transform_messages(messages)
+        input_ids, attention_mask, pixel_values = self._generate_chat_data(msgs)
+        _, inputs_embeds, _, attention_mask = self._model.merge_multimodal(
+            text_input_ids=input_ids,
+            text_attention_masks=attention_mask,
+            text_labels=None,
+            pixel_values=pixel_values,
+            left_padding=True,
+        )
+        inputs_embeds = inputs_embeds.detach()
+        torch.cuda.empty_cache()
+        return dict(
+            input_ids=input_ids,
+            inputs_embeds=inputs_embeds,
+            attention_mask=attention_mask,
+        )
 
-    def _convert_video_tensors_to_pil(self, video_inputs: List) -> List[Image.Image]:
-        """Convert video tensors to a list of PIL images"""
-        from torchvision import transforms
+    def build_streaming_iter(
+        self,
+        messages: List[Dict],
+        generate_config: Dict,
+    ) -> Tuple[Iterator, int]:
+        from transformers import TextIteratorStreamer
 
-        to_pil = transforms.ToPILImage()
-        pil_images = []
+        streamer = TextIteratorStreamer(
+            self._text_tokenizer, timeout=60, skip_prompt=True, skip_special_tokens=True
+        )
+        config = self.build_generate_kwargs(generate_config)
+        inputs = self.build_inputs_from_messages(messages, generate_config)
+        input_ids = inputs.pop("input_ids")
 
-        for video_tensor_4d in video_inputs:
-            if isinstance(video_tensor_4d, torch.Tensor):
-                # Verify it's a 4D tensor
-                if video_tensor_4d.ndim == 4:
-                    # Iterate through the first dimension (frames) of 4D tensor
-                    for i in range(video_tensor_4d.size(0)):
-                        frame_tensor_3d = video_tensor_4d[
-                            i
-                        ]  # Get 3D frame tensor [C, H, W]
-                        # Ensure tensor is on CPU before conversion
-                        if frame_tensor_3d.is_cuda:
-                            frame_tensor_3d = frame_tensor_3d.cpu()
-                        try:
-                            pil_image = to_pil(frame_tensor_3d)
-                            pil_images.append(pil_image)
-                        except Exception as e:
-                            logger.error(
-                                f"Error converting frame {i} to PIL Image: {e}"
-                            )
-                            # Can choose to skip this frame or handle error differently
-                else:
-                    logger.warning(
-                        f"Expected 4D tensor in video_inputs, but got {video_tensor_4d.ndim}D. Skipping this tensor."
-                    )
-            elif isinstance(video_tensor_4d, Image.Image):
-                # If fetch_video returns Image list, add directly
-                pil_images.append(video_tensor_4d)
-            else:
-                logger.warning(
-                    f"Unexpected type in video_inputs: {type(video_tensor_4d)}. Skipping."
-                )
+        gen_kwargs = dict(**inputs, **config, streamer=streamer)
 
-        return pil_images
+        thread = Thread(target=self._model.llm.generate, kwargs=gen_kwargs)
+        thread.start()
+        return streamer, len(input_ids[0])

xinference/model/llm/transformers/{qwen-omni.py → multimodal/qwen-omni.py}

@@ -11,49 +11,36 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-
 import base64
 import importlib.util
 import io
 import logging
-import sys
 import time
 import uuid
-from typing import Dict, Iterator, List, Optional, Union
+from threading import Thread
+from typing import Any, Dict, Iterator, List, Optional, Tuple
 
-from ....model.utils import select_device
-from ....types import (
+from .....model.utils import select_device
+from .....types import (
     ChatCompletion,
     ChatCompletionAudio,
     ChatCompletionChoice,
-    ChatCompletionChunk,
-    ChatCompletionMessage,
-    CompletionChunk,
     CompletionUsage,
 )
-from ..llm_family import LLMFamilyV1, LLMSpecV1, register_transformer
-from ..utils import generate_completion_chunk
-from .core import PytorchChatModel, PytorchGenerateConfig, register_non_default_model
-from .utils import cache_clean
+from ...llm_family import LLMFamilyV1, LLMSpecV1, register_transformer
+from ..core import PytorchGenerateConfig, register_non_default_model
+from .core import PytorchMultiModalModel
 
 logger = logging.getLogger(__name__)
 
-DEFAULT_SYSTEM_PROMPT = (
-    "You are Qwen, a virtual human developed by the Qwen Team, Alibaba Group, "
-    "capable of perceiving auditory and visual inputs, as well as generating text and speech."
-)
-
 
 @register_transformer
 @register_non_default_model("qwen2.5-omni")
-class Qwen2_5OmniChatModel(PytorchChatModel):
-    def __init__(self, *args, **kwargs):
-        super().__init__(*args, **kwargs)
-
-        self._tokenizer = None
-        self._model = None
-        self._device = None
-        self._processor = None
+class Qwen2_5OmniChatModel(PytorchMultiModalModel):
+    DEFAULT_SYSTEM_PROMPT = (
+        "You are Qwen, a virtual human developed by the Qwen Team, Alibaba Group, "
+        "capable of perceiving auditory and visual inputs, as well as generating text and speech."
+    )
 
     @classmethod
     def match_json(
@@ -66,23 +53,24 @@ class Qwen2_5OmniChatModel(PytorchChatModel):
             return True
         return False
 
-    def load(self):
-        logger.debug(
-            "Try to load model, current python: %s, sys path: %s",
-            sys.executable,
-            sys.path,
-        )
-
-        from transformers import (
-            Qwen2_5OmniForConditionalGeneration,
-            Qwen2_5OmniProcessor,
-        )
-
+    def decide_device(self):
         device = self._pytorch_model_config.get("device", "auto")
         device = select_device(device)
         self._device = device
+
+    def load_processor(self):
+        from transformers import Qwen2_5OmniProcessor
+
+        self._processor = Qwen2_5OmniProcessor.from_pretrained(
+            self.model_path, trust_remote_code=True
+        )
+        self._tokenizer = self._processor.tokenizer
+
+    def load_multimodal_model(self):
+        from transformers import Qwen2_5OmniForConditionalGeneration
+
         # for multiple GPU, set back to auto to make multiple devices work
-        device = "auto" if device == "cuda" else device
+        device = "auto" if self._device == "cuda" else self._device
         flash_attn_installed = importlib.util.find_spec("flash_attn") is not None
         kwargs = (
             {}
@@ -92,10 +80,6 @@ class Qwen2_5OmniChatModel(PytorchChatModel):
         kwargs = self.apply_bnb_quantization(kwargs)
         logger.debug("Loading model with extra kwargs: %s", kwargs)
 
-        self._processor = Qwen2_5OmniProcessor.from_pretrained(
-            self.model_path, trust_remote_code=True
-        )
-        self._tokenizer = self._processor.tokenizer
         self._model = Qwen2_5OmniForConditionalGeneration.from_pretrained(
             self.model_path,
             torch_dtype="auto",
@@ -104,28 +88,9 @@ class Qwen2_5OmniChatModel(PytorchChatModel):
             **kwargs,
         )
 
-    @cache_clean
-    def chat(
-        self,
-        messages: List[Dict],
-        generate_config: Optional[PytorchGenerateConfig] = None,
-    ) -> Union[ChatCompletion, Iterator[ChatCompletionChunk]]:
-        messages = self._transform_messages(messages)
-
-        generate_config = generate_config if generate_config else {}
-
-        stream = generate_config.get("stream", False) if generate_config else False
-
-        if stream:
-            it = self._generate_stream(messages, generate_config)
-            return self._to_chat_completion_chunks(it)
-        else:
-            c = self._generate(messages, generate_config)
-            return c
-
     def _transform_messages(
         self,
-        messages: Union[List[ChatCompletionMessage], List[dict]],
+        messages: List[dict],  # type: ignore
     ):
         messages = super()._transform_messages(messages)
         if messages[0]["role"] != "system":
@@ -133,23 +98,24 @@ class Qwen2_5OmniChatModel(PytorchChatModel):
                 0,
                 {
                     "role": "system",
-                    "content": [{"type": "text", "text": DEFAULT_SYSTEM_PROMPT}],  # type: ignore
+                    "content": [{"type": "text", "text": self.DEFAULT_SYSTEM_PROMPT}],  # type: ignore
                 },
             )
         else:
             logger.debug("Force to set system prompt")
-            messages[0]["content"] = [{"type": "text", "text": DEFAULT_SYSTEM_PROMPT}]  # type: ignore
+            messages[0]["content"] = [{"type": "text", "text": self.DEFAULT_SYSTEM_PROMPT}]  # type: ignore
         return messages
 
-    def _generate(
-        self, messages: List, config: PytorchGenerateConfig = {}
-    ) -> ChatCompletion:
-        import soundfile as sf
+    def build_inputs_from_messages(
+        self,
+        messages: List[Dict],
+        generate_config: Dict,
+    ):
         from qwen_omni_utils import process_mm_info
 
-        use_audio_in_video = config.get("use_audio_in_video", True)
-        voice = config.get("voice", "Chelsie")
+        use_audio_in_video = generate_config.get("use_audio_in_video", True)
 
+        messages = self._transform_messages(messages)
         text = self._processor.apply_chat_template(
             messages, tokenize=False, add_generation_prompt=True
         )
@@ -169,15 +135,54 @@ class Qwen2_5OmniChatModel(PytorchChatModel):
             use_audio_in_video=use_audio_in_video,
         )
         inputs = inputs.to(self._device)
+        return inputs
 
-        # Inference: Generation of the output
-        generated_ids, audio = self._model.generate(
-            **inputs,
-            speaker=voice,
-            max_new_tokens=config.get("max_tokens", 512),
-            temperature=config.get("temperature", 1),
-            use_audio_in_video=use_audio_in_video,
+    def build_generate_kwargs(
+        self,
+        generate_config: Dict,
+    ) -> Dict[str, Any]:
+        voice = generate_config.get("voice", "Chelsie")
+        return {
+            "max_new_tokens": generate_config.get("max_tokens", 512),
+            "temperature": generate_config.get("temperature", 1),
+            "speaker": voice,
+        }
+
+    def build_streaming_iter(
+        self,
+        messages: List[Dict],
+        generate_config: Dict,
+    ) -> Tuple[Iterator, int]:
+        from transformers import TextIteratorStreamer
+
+        tokenizer = self._tokenizer
+        streamer = TextIteratorStreamer(
+            tokenizer, timeout=60.0, skip_prompt=True, skip_special_tokens=True
         )
+
+        config = self.build_generate_kwargs(generate_config)
+        inputs = self.build_inputs_from_messages(messages, generate_config)
+        gen_kwargs = dict(**inputs, **config, streamer=streamer)
+        thread = Thread(target=self._model.generate, kwargs=gen_kwargs)
+        thread.start()
+        return streamer, len(inputs.input_ids[0])
+
+    def generate_non_streaming(
+        self,
+        messages: List[Dict],
+        generate_config: Optional[PytorchGenerateConfig] = None,
+    ) -> ChatCompletion:
+        """
+        Special case for qwen2.5-omni, since it has audio output
+        """
+        import soundfile as sf
+
+        generate_config = generate_config if generate_config else {}  # type: ignore
+        config = self.build_generate_kwargs(generate_config)  # type: ignore
+        inputs = self.build_inputs_from_messages(messages, generate_config)  # type: ignore
+        use_audio_in_video = generate_config.get("use_audio_in_video", True)
+        gen_kwargs = dict(**inputs, **config, use_audio_in_video=use_audio_in_video)
+        generated_ids, audio = self._model.generate(**gen_kwargs)
         generated_ids_trimmed = [
             out_ids[len(in_ids) :]
             for in_ids, out_ids in zip(inputs.input_ids, generated_ids)
@@ -223,93 +228,3 @@ class Qwen2_5OmniChatModel(PytorchChatModel):
                 prompt_tokens=-1, completion_tokens=-1, total_tokens=-1
             ),
         )
-
-    def _generate_stream(
-        self, messages: List, config: PytorchGenerateConfig = {}
-    ) -> Iterator[CompletionChunk]:
-        from threading import Thread
-
-        from qwen_omni_utils import process_mm_info
-        from transformers import TextIteratorStreamer
-
-        use_audio_in_video = config.get("use_audio_in_video", True)
-        voice = config.get("voice", "Chelsie")
-
-        text = self._processor.apply_chat_template(
-            messages, tokenize=False, add_generation_prompt=True
-        )
-        audios, images, videos = process_mm_info(
-            messages, use_audio_in_video=use_audio_in_video
-        )
-        logger.debug(
-            "Text, audio, image, video: %s, %s, %s, %s", text, audios, images, videos
-        )
-        inputs = self._processor(
-            text=text,
-            images=images,
-            audio=audios,
-            videos=videos,
-            padding=True,
-            return_tensors="pt",
-            use_audio_in_video=use_audio_in_video,
-        )
-        inputs = inputs.to(self._device)
-
-        tokenizer = self._tokenizer
-        streamer = TextIteratorStreamer(
-            tokenizer, timeout=60.0, skip_prompt=True, skip_special_tokens=True
-        )
-
-        # TODO(xuye): Cannot find a way to streaming output,
-        # will implement it when it's supported
-
-        gen_kwargs = {
-            "max_new_tokens": config.get("max_tokens", 512),
-            "temperature": config.get("temperature", 1),
-            "streamer": streamer,
-            "speaker": voice,
-            **inputs,
-        }
-        error = None
-
-        def model_generate():
-            try:
-                return self._model.generate(**gen_kwargs)
-            except Exception:
-                nonlocal error
-                error = sys.exc_info()
-                streamer.end()
-                raise
-
-        thread = Thread(target=model_generate)
-        thread.start()
-
-        completion_id = str(uuid.uuid1())
-        for new_text in streamer:
-            yield generate_completion_chunk(
-                chunk_text=new_text,
-                finish_reason=None,
-                chunk_id=completion_id,
-                model_uid=self.model_uid,
-                prompt_tokens=-1,
-                completion_tokens=-1,
-                total_tokens=-1,
-                has_choice=True,
-                has_content=True,
-            )
-
-        if error:
-            _, err, tb = error  # type: ignore
-            raise err.with_traceback(tb)
-
-        yield generate_completion_chunk(
-            chunk_text=None,
-            finish_reason="stop",
-            chunk_id=completion_id,
-            model_uid=self.model_uid,
-            prompt_tokens=-1,
-            completion_tokens=-1,
-            total_tokens=-1,
-            has_choice=True,
-            has_content=False,
-        )