From dcfc2656cb21a2eb608ff0f41c4c15ea24c9dc5d Mon Sep 17 00:00:00 2001
From: wanzihao <1060304770@qq.com>
Date: Mon, 20 Jul 2026 14:35:10 +0800
Subject: [PATCH 1/2] compute image start_idx after audio expansion

---
 lightllm/models/qwen3_omni_moe_thinker/model.py | 7 ++++++-
 1 file changed, 6 insertions(+), 1 deletion(-)

diff --git a/lightllm/models/qwen3_omni_moe_thinker/model.py b/lightllm/models/qwen3_omni_moe_thinker/model.py
index 45fb16f701..a221492126 100644
--- a/lightllm/models/qwen3_omni_moe_thinker/model.py
+++ b/lightllm/models/qwen3_omni_moe_thinker/model.py
@@ -101,7 +101,6 @@ def encode(self, prompt, multimodal_params: MultimodalParams = None, **kwargs):
                     input_ids.extend(origin_ids[: start_idx + 1])
                     token_id = multimodal_params.images[image_id].token_id
                     token_num = multimodal_params.images[image_id].token_num
-                    multimodal_params.images[image_id].start_idx = len(input_ids)
                     input_ids.extend(range(token_id, token_id + token_num))
                     input_ids.append(self.image_end_id)
                     origin_ids = origin_ids[start_idx + 2 :]
@@ -144,6 +143,12 @@ def encode(self, prompt, multimodal_params: MultimodalParams = None, **kwargs):
                 raise ValueError(audio_cnt == audio_id, f"invalid audio tag num: {audio_cnt} vs {audio_id}!")
         input_ids.extend(origin_ids)
 
+        if multimodal_params:
+            image_search_start = 0
+            for image in multimodal_params.images:
+                image.start_idx = input_ids.index(image.token_id, image_search_start)
+                image_search_start = image.start_idx + image.token_num
+
         return input_ids
 
 

From 8ee9a96abd4bf7cf801198f1b187869960b247e9 Mon Sep 17 00:00:00 2001
From: wanzihao <1060304770@qq.com>
Date: Tue, 21 Jul 2026 15:13:00 +0800
Subject: [PATCH 2/2] Simplify

---
 .../models/qwen3_omni_moe_thinker/model.py    | 56 +++++++++----------
 1 file changed, 26 insertions(+), 30 deletions(-)

diff --git a/lightllm/models/qwen3_omni_moe_thinker/model.py b/lightllm/models/qwen3_omni_moe_thinker/model.py
index a221492126..e4ea3ccc34 100644
--- a/lightllm/models/qwen3_omni_moe_thinker/model.py
+++ b/lightllm/models/qwen3_omni_moe_thinker/model.py
@@ -89,34 +89,8 @@ def encode(self, prompt, multimodal_params: MultimodalParams = None, **kwargs):
 
         # <img><image_pad></img> -> <img></img>
         origin_ids = [token for token in origin_ids if token not in (self.image_token_id, self.audio_token_id)]
-        # <img></img> --> <img>id,id+1...id+num</img>
-        input_ids = []
-        image_id = 0
-        while True:
-            try:
-                start_idx = origin_ids.index(self.image_start_id)
-                if start_idx + 1 >= len(origin_ids):
-                    break
-                if origin_ids[start_idx + 1] == self.image_end_id:
-                    input_ids.extend(origin_ids[: start_idx + 1])
-                    token_id = multimodal_params.images[image_id].token_id
-                    token_num = multimodal_params.images[image_id].token_num
-                    input_ids.extend(range(token_id, token_id + token_num))
-                    input_ids.append(self.image_end_id)
-                    origin_ids = origin_ids[start_idx + 2 :]
-                    image_id += 1
-                else:
-                    raise ValueError("image token error")
-            except ValueError:
-                break
-        if multimodal_params:
-            image_cnt = len(multimodal_params.images)
-            if image_cnt != image_id:
-                raise ValueError(image_cnt == image_id, f"invalid image tag num: {image_cnt} vs {image_id}!")
-        input_ids.extend(origin_ids)
 
         # audio
-        origin_ids = input_ids
         input_ids = []
         audio_id = 0
         start_idx = 0
@@ -143,11 +117,33 @@ def encode(self, prompt, multimodal_params: MultimodalParams = None, **kwargs):
                 raise ValueError(audio_cnt == audio_id, f"invalid audio tag num: {audio_cnt} vs {audio_id}!")
         input_ids.extend(origin_ids)
 
+        # <img></img> --> <img>id,id+1...id+num</img>
+        origin_ids = input_ids
+        input_ids = []
+        image_id = 0
+        while True:
+            try:
+                start_idx = origin_ids.index(self.image_start_id)
+                if start_idx + 1 >= len(origin_ids):
+                    break
+                if origin_ids[start_idx + 1] == self.image_end_id:
+                    input_ids.extend(origin_ids[: start_idx + 1])
+                    token_id = multimodal_params.images[image_id].token_id
+                    token_num = multimodal_params.images[image_id].token_num
+                    multimodal_params.images[image_id].start_idx = len(input_ids)
+                    input_ids.extend(range(token_id, token_id + token_num))
+                    input_ids.append(self.image_end_id)
+                    origin_ids = origin_ids[start_idx + 2 :]
+                    image_id += 1
+                else:
+                    raise ValueError("image token error")
+            except ValueError:
+                break
         if multimodal_params:
-            image_search_start = 0
-            for image in multimodal_params.images:
-                image.start_idx = input_ids.index(image.token_id, image_search_start)
-                image_search_start = image.start_idx + image.token_num
+            image_cnt = len(multimodal_params.images)
+            if image_cnt != image_id:
+                raise ValueError(image_cnt == image_id, f"invalid image tag num: {image_cnt} vs {image_id}!")
+        input_ids.extend(origin_ids)
 
         return input_ids