add deepseek-ocr2 support

src/models/deepseek_ocr/generate.rs

@@ -30,6 +30,7 @@ pub struct DeepseekOCRGenerateModel {
     device: Device,
     size: Vec<u32>,
     model_name: String,
+    version: usize,
 }
 
 impl DeepseekOCRGenerateModel {
@@ -40,13 +41,23 @@ impl DeepseekOCRGenerateModel {
         let cfg_dtype = cfg.language_config.torch_dtype.clone();
         let device = &get_device(device);
         let dtype = get_dtype(dtype, &cfg_dtype);
-        let processor = DeepseekOCRProcessor::new(device, dtype)?;
+        let model_name = std::path::Path::new(path)
+            .file_stem() // 获取文件名主干（不含扩展名）
+            .and_then(|s| s.to_str())
+            .unwrap_or("deepseek-ocr");
+        let version = if model_name.contains("2") || cfg.vision_config.width.qwen2_0_5b.is_some() {
+            2usize
+        } else {
+            1usize
+        };
+        let processor = DeepseekOCRProcessor::new(device, dtype, version)?;
         let eos_token_id = cfg.eos_token_id;
         let bos_token_id = cfg.bos_token_id;
         let model_list = find_type_files(path, "safetensors")?;
         let vb = unsafe { VarBuilder::from_mmaped_safetensors(&model_list, dtype, device)? };
-        let deepseekocr_model = DeepseekOCRModel::new(vb, cfg)?;
+        let deepseekocr_model = DeepseekOCRModel::new(vb, cfg, version)?;
         let size = vec![512u32, 640, 1024, 1280];
+
         Ok(Self {
             tokenizer,
             processor,
@@ -55,7 +66,8 @@ impl DeepseekOCRGenerateModel {
             eos_token_id,
             device: device.clone(),
             size,
-            model_name: "deepseek-ocr".to_string(),
+            model_name: model_name.to_string(),
+            version,
         })
     }
 }
@@ -77,6 +89,8 @@ impl GenerateModel for DeepseekOCRGenerateModel {
         let crop_mode = extract_metadata_value::<bool>(&mes.metadata, "crop_mode").unwrap_or(false);
         let seed = mes.seed.unwrap_or(34562) as u64;
         let mut logit_processor = get_logit_processor(mes.temperature, mes.top_p, None, seed);
+        let base_size = if self.version == 2 { 1024 } else { base_size };
+        let image_size = if self.version == 2 { 768 } else { image_size };
         let (mut input_ids, images_ori, image_crop, images_seq_mask, images_spatial_crop_t) = self
             .processor
             .process_info(&mes, &self.tokenizer, base_size, image_size, crop_mode)?;
@@ -152,6 +166,8 @@ impl GenerateModel for DeepseekOCRGenerateModel {
         let crop_mode = extract_metadata_value::<bool>(&mes.metadata, "crop_mode").unwrap_or(false);
         let seed = mes.seed.unwrap_or(34562) as u64;
         let mut logit_processor = get_logit_processor(mes.temperature, mes.top_p, None, seed);
+        let base_size = if self.version == 2 { 1024 } else { base_size };
+        let image_size = if self.version == 2 { 768 } else { image_size };
         let (mut input_ids, images_ori, image_crop, images_seq_mask, images_spatial_crop_t) = self
             .processor
             .process_info(&mes, &self.tokenizer, base_size, image_size, crop_mode)?;