Update qwen2-vl inference (transformers instead of swift)

app.py

@@ -1,6 +1,9 @@
+import base64
 import hashlib
+import io
 import os
 from pathlib import Path
+from threading import Thread
 from typing import Iterator, Optional, List, Union
 
 import gradio as gr
@@ -8,6 +11,7 @@ import spaces
 import torch
 from PIL import Image
 from pydantic import BaseModel
+from qwen_vl_utils import process_vision_info
 from swift.llm import (
     ModelType,
     get_model_tokenizer,
@@ -19,6 +23,9 @@ from swift.llm import (
 from transformers import (
     Qwen2VLForConditionalGeneration,
     PreTrainedTokenizer,
+    Qwen2VLProcessor,
+    TextIteratorStreamer,
+    AutoTokenizer,
 )
 
 MAX_MAX_NEW_TOKENS = 2048
@@ -42,6 +49,19 @@ this demo is governed by the original [license](https://huggingface.co/meta-llam
 """
 
 
+def convert_image_to_text(image: Image) -> str:
+    # This is also how OpenAI encodes images: https://platform.openai.com/docs/guides/vision
+    with io.BytesIO() as output:
+        image.save(output, format="PNG")
+        data = output.getvalue()
+    return base64.b64encode(data).decode("utf-8")
+
+
+def convert_text_to_image(text: str) -> Image:
+    data = base64.b64decode(text.encode("utf-8"))
+    return Image.open(io.BytesIO(data))
+
+
 def save_image(image: Image.Image, folder: str) -> str:
     image_hash = hashlib.md5(image.tobytes()).hexdigest()
     path = Path(folder, f"{image_hash}.png")
@@ -139,12 +159,122 @@ class SwiftQwenModel(EvalModel):
             yield delta
 
 
+class QwenModel(EvalModel):
+    path: str = "models/qwen"
+    engine: str = "Qwen/Qwen2-VL-7B-Instruct"
+    model: Optional[Qwen2VLForConditionalGeneration] = None
+    processor: Optional[Qwen2VLProcessor] = None
+    tokenizer: Optional[AutoTokenizer]
+    device: str = "cuda"
+    image_size: int = 768
+    lora_path: str = ""
+
+    def load(self):
+        if self.model is None:
+            path = self.path if os.path.exists(self.path) else self.engine
+            print(dict(load_path=path))
+            # noinspection PyTypeChecker
+            self.model = Qwen2VLForConditionalGeneration.from_pretrained(
+                path, torch_dtype="auto", device_map="auto"
+            )
+            self.tokenizer = AutoTokenizer.from_pretrained(self.engine)
+
+            if self.lora_path:
+                print("Loading LORA from", self.lora_path)
+                self.model.load_adapter(self.lora_path)
+
+            self.model = self.model.to(self.device).eval()
+            self.processor = Qwen2VLProcessor.from_pretrained(self.engine)
+            torch.manual_seed(0)
+            torch.cuda.manual_seed_all(0)
+
+    def make_messages(self, inputs: List[Union[str, Image.Image]]) -> List[dict]:
+        text = "\n\n".join([x for x in inputs if isinstance(x, str)])
+        content = [
+            dict(
+                type="image",
+                image=f"data:image;base64,{convert_image_to_text(resize_image(x, self.image_size))}",
+            )
+            for x in inputs
+            if isinstance(x, Image.Image)
+        ]
+        content.append(dict(type="text", text=text))
+        return [dict(role="user", content=content)]
+
+    def run(self, inputs: List[Union[str, Image.Image]]) -> str:
+        self.load()
+        messages = self.make_messages(inputs)
+        text = self.processor.apply_chat_template(
+            messages, tokenize=False, add_generation_prompt=True
+        )
+
+        image_inputs, video_inputs = process_vision_info(messages)
+        # noinspection PyTypeChecker
+        model_inputs = self.processor(
+            text=[text],
+            images=image_inputs,
+            videos=video_inputs,
+            padding=True,
+            return_tensors="pt",
+        ).to(self.device)
+
+        with torch.inference_mode():
+            generated_ids = self.model.generate(
+                **model_inputs, max_new_tokens=self.max_output_tokens
+            )
+
+        generated_ids_trimmed = [
+            out_ids[len(in_ids) :]
+            for in_ids, out_ids in zip(model_inputs.input_ids, generated_ids)
+        ]
+        output_text = self.processor.batch_decode(
+            generated_ids_trimmed,
+            skip_special_tokens=True,
+            clean_up_tokenization_spaces=False,
+        )
+        return output_text[0]
+
+    def run_stream(self, inputs: List[Union[str, Image.Image]]) -> Iterator[str]:
+        self.load()
+        messages = self.make_messages(inputs)
+        text = self.processor.apply_chat_template(
+            messages, tokenize=False, add_generation_prompt=True
+        )
+
+        image_inputs, video_inputs = process_vision_info(messages)
+        # noinspection PyTypeChecker
+        model_inputs = self.processor(
+            text=[text],
+            images=image_inputs,
+            videos=video_inputs,
+            padding=True,
+            return_tensors="pt",
+        ).to(self.device)
+
+        streamer = TextIteratorStreamer(
+            self.tokenizer,
+            timeout=10.0,
+            skip_prompt=True,
+            skip_special_tokens=True,
+        )
+
+        generate_kwargs = dict(**model_inputs)
+        t = Thread(target=self.model.generate, kwargs=generate_kwargs)
+        t.start()
+
+        outputs = []
+        for text in streamer:
+            outputs.append(text)
+            yield "".join(outputs)
+
+
 if not torch.cuda.is_available():
     DESCRIPTION += "\n<p>Running on CPU 🥶 This demo does not work on CPU.</p>"
 
 
 if torch.cuda.is_available():
-    model = SwiftQwenModel()
+    model = QwenModel()
+    model.load()
 
 
 @spaces.GPU
@@ -158,10 +288,8 @@ def generate(
     top_k: int = 50,
     repetition_penalty: float = 1.2,
 ) -> Iterator[str]:
-    outputs = []
-    for text in model.run_stream(inputs=[message]):
-        outputs.append(text)
-        yield "".join(outputs)
+    for text in model.run_stream([message]):
+        yield text
 
 
 chat_interface = gr.ChatInterface(