Upload mmalaya_arch.py

mmalaya_arch.py

@@ -299,15 +299,15 @@ class MMAlayaMetaForCausalLM(ABC):
         for x in insert_separator(prompt_chunks, [image_token_index] * (offset + 1)):
             input_ids.extend(x[offset:])
 
-        if return_tensors is not None:
-            if return_tensors == 'pt':
-                return torch.tensor(input_ids, dtype=torch.long)
-            raise ValueError(f'Unsupported tensor type: {return_tensors}')
-        
         # 加载generate stop条件
         stopping_criteria = KeywordsStoppingCriteria([conv.sep2], tokenizer, input_ids)
         # 加载图像
         image_processor = model.get_vision_tower().image_processor
         image_tensor = image_processor(image, return_tensors='pt')['pixel_values'].half().cuda()
 
+        if return_tensors is not None:
+            if return_tensors == 'pt':
+                return torch.tensor(input_ids, dtype=torch.long), image_tensor, stopping_criteria
+            raise ValueError(f'Unsupported tensor type: {return_tensors}')
+        
         return input_ids, image_tensor, stopping_criteria