add RMBGv2.0

README.md

@@ -20,6 +20,7 @@
 * Hunyuan-OCR - 腾讯混元光学文字识别模型
 * PaddleOCR-VL - 百度飞桨光学文字识别模型
 * VoxCPM1.5 - 面壁智能语音生成模型1.5版本
+* RMBG2.0 - RMBGv2.0由BRIA AI开发，供非商业用途使用。
 
 ## 计划支持
 我们持续扩展支持的模型列表，欢迎贡献！

src/models/common/mod.rs

@@ -1,8 +1,9 @@
 use anyhow::Result;
 use candle_core::{D, Tensor};
 use candle_nn::{
-    Activation, Conv2d, Conv2dConfig, LayerNorm, LayerNormConfig, Linear, Module, RmsNorm,
+    Activation, BatchNorm, BatchNormConfig, Conv2d, Conv2dConfig, LayerNorm, LayerNormConfig,
-    VarBuilder, conv2d, conv2d_no_bias, layer_norm, linear, linear_no_bias, rms_norm,
+    Linear, Module, RmsNorm, VarBuilder, batch_norm, conv2d, conv2d_no_bias, layer_norm, linear,
+    linear_no_bias, rms_norm,
 };
 
 use crate::{position_embed::rope::apply_rotary_pos_emb, utils::tensor_utils::repeat_kv};
@@ -511,3 +512,111 @@ pub fn get_layer_norm(vb: VarBuilder, eps: f64, dim: usize) -> Result<LayerNorm>
     let norm = layer_norm(dim, ln_config, vb)?;
     Ok(norm)
 }
+
+pub fn get_batch_norm(vb: VarBuilder, eps: f64, dim: usize) -> Result<BatchNorm> {
+    let bn_config = BatchNormConfig {
+        eps,
+        remove_mean: true,
+        affine: true,
+        momentum: 0.1,
+    };
+    let norm = batch_norm(dim, bn_config, vb)?;
+    Ok(norm)
+}
+
+pub fn deform_conv2d_kernel(
+    input: &Tensor,
+    weight: &Tensor,
+    bias: Option<&Tensor>,
+    offset: &Tensor,
+    mask: Option<&Tensor>,
+    stride: usize,
+    padding: usize,
+) -> Result<Tensor> {
+    // 不考虑空洞卷积, bs = 1
+    let (_, in_c, in_h, in_w) = input.dims4()?;
+    let (out_channel, _, ker_h, ker_w) = weight.dims4()?;
+    let out_h = ((in_h + 2 * padding - ker_h) / stride) + 1;
+    let out_w = ((in_w + 2 * padding - ker_w) / stride) + 1;
+
+    let num_kernels = in_c * out_h * out_w;
+    let mask_vec = if let Some(mask) = mask {
+        Some(mask.squeeze(0)?.to_vec3::<f32>()?)
+    } else {
+        None
+    };
+    let offset_vec = offset.squeeze(0)?.to_vec3::<f32>()?;
+    let input_vec = input.squeeze(0)?.to_vec3::<f32>()?;
+    let mut columns_vec = vec![vec![0.0f32; out_h * out_w]; in_c * ker_h * ker_w];
+    for index in 0..num_kernels {
+        let out_x = index % out_w;
+        let out_y = (index / out_w) % out_h;
+        let in_c = index / (out_w * out_h);
+        let out_c = in_c * ker_h * ker_w;
+
+        for i in 0..ker_h {
+            for j in 0..ker_w {
+                let mask_idx = i * ker_w + j;
+                let offset_idx = 2 * mask_idx;
+                let mask_value = if mask.is_some() {
+                    mask_vec.as_ref().unwrap()[mask_idx][out_y][out_x]
+                } else {
+                    1.0
+                };
+                let offset_h = offset_vec[offset_idx][out_y][out_x];
+                let offset_w = offset_vec[offset_idx + 1][out_y][out_x];
+                let y = ((out_y * stride - padding) + i) as f32 + offset_h;
+                let x = ((out_x * stride - padding) + j) as f32 + offset_w;
+                let val = if y <= -1.0 || in_h as f32 <= y || x <= -1.0 || in_w as f32 <= x {
+                    0.0
+                } else {
+                    let h_low = y.floor();
+                    let w_low = x.floor();
+                    let h_high = h_low + 1.0;
+                    let w_high = w_low + 1.0;
+                    let lh = y - h_low;
+                    let lw = x - w_low;
+                    let hh = 1.0 - lh;
+                    let hw = 1.0 - lw;
+                    let w1 = hh * hw;
+                    let w2 = hh * lw;
+                    let w3 = lh * hw;
+                    let w4 = lh * lw;
+                    let v1 = if h_low >= 0.0 && w_low >= 0.0 {
+                        input_vec[in_c][h_low as usize][w_low as usize]
+                    } else {
+                        0.0
+                    };
+                    let v2 = if h_low >= 0.0 && w_high <= (in_w - 1) as f32 {
+                        input_vec[in_c][h_low as usize][w_high as usize]
+                    } else {
+                        0.0
+                    };
+                    let v3 = if h_high <= (in_h - 1) as f32 && w_low >= 0.0 {
+                        input_vec[in_c][h_high as usize][w_low as usize]
+                    } else {
+                        0.0
+                    };
+                    let v4 = if h_high <= (in_h - 1) as f32 && w_high <= (in_w - 1) as f32 {
+                        input_vec[in_c][h_high as usize][w_high as usize]
+                    } else {
+                        0.0
+                    };
+                    w1 * v1 + w2 * v2 + w3 * v3 + w4 * v4
+                };
+                columns_vec[out_c + i * ker_w + j][out_y * out_w + out_x] = mask_value * val;
+            }
+        }
+    }
+
+    let columns = Tensor::new(columns_vec, weight.device())?;
+    let mut out =
+        weight
+            .flatten_from(1)?
+            .matmul(&columns)?
+            .reshape((1, out_channel, out_h, out_w))?;
+    if let Some(bias) = bias {
+        out = out.broadcast_add(bias)?;
+    }
+    Ok(out)
+}

src/models/mod.rs

@@ -5,8 +5,8 @@ pub mod minicpm4;
 pub mod paddleocr_vl;
 pub mod qwen2_5vl;
 pub mod qwen3vl;
-pub mod voxcpm;
 pub mod rmbg2_0;
+pub mod voxcpm;
 
 use aha_openai_dive::v1::resources::chat::{
     ChatCompletionChunkResponse, ChatCompletionParameters, ChatCompletionResponse,

src/models/rmbg2_0/generate.rs

@@ -0,0 +1,80 @@
+use aha_openai_dive::v1::resources::chat::ChatCompletionParameters;
+use anyhow::Result;
+use candle_core::{DType, Device, Tensor};
+use candle_nn::VarBuilder;
+use image::{Rgba, RgbaImage};
+
+use crate::{
+    models::rmbg2_0::model::BiRefNet,
+    utils::{
+        find_type_files, get_device, get_dtype,
+        img_utils::{extract_images, float_tensor_to_dynamic_image, img_transform_with_resize},
+    },
+};
+
+pub struct RMBG2_0 {
+    model: BiRefNet,
+    h: u32,
+    w: u32,
+    img_mean: Tensor,
+    img_std: Tensor,
+    device: Device,
+    dtype: DType,
+}
+
+impl RMBG2_0 {
+    pub fn init(path: &str, device: Option<&Device>, dtype: Option<DType>) -> Result<Self> {
+        let device = get_device(device);
+        let dtype = get_dtype(dtype, "float32");
+        let model_list = find_type_files(path, "safetensors")?;
+        let vb = unsafe { VarBuilder::from_mmaped_safetensors(&model_list, dtype, &device)? };
+        let model = BiRefNet::new(vb)?;
+        let img_mean =
+            Tensor::from_slice(&[0.485, 0.456, 0.406], (3, 1, 1), &device)?.to_dtype(dtype)?;
+        let img_std =
+            Tensor::from_slice(&[0.229, 0.224, 0.225], (3, 1, 1), &device)?.to_dtype(dtype)?;
+        Ok(Self {
+            model,
+            h: 1024,
+            w: 1024,
+            img_mean,
+            img_std,
+            device,
+            dtype,
+        })
+    }
+
+    pub fn generate(&self, mes: ChatCompletionParameters) -> Result<Vec<RgbaImage>> {
+        let imgs = extract_images(&mes)?;
+        let mut rmbg_png = vec![];
+        for img in imgs {
+            let height = img.height();
+            let width = img.width();
+            let img_tensor = img_transform_with_resize(
+                &img,
+                self.h,
+                self.w,
+                &self.img_mean,
+                &self.img_std,
+                &self.device,
+                self.dtype,
+            )?
+            .unsqueeze(0)?;
+            let rmbg_img = self.model.forward(&img_tensor)?.squeeze(0)?;
+            let alpha_img = float_tensor_to_dynamic_image(&rmbg_img)?;
+            let alpha_img =
+                alpha_img.resize_exact(width, height, image::imageops::FilterType::CatmullRom);
+            let alpha_gray = alpha_img.to_luma8();
+            let mut rgba_img = RgbaImage::new(width, height);
+
+            // 遍历像素并组合
+            for (x, y, pixel) in img.to_rgb8().enumerate_pixels() {
+                let alpha_value = alpha_gray.get_pixel(x, y).0[0];
+                let rgba_pixel = Rgba([pixel.0[0], pixel.0[1], pixel.0[2], alpha_value]);
+                rgba_img.put_pixel(x, y, rgba_pixel);
+            }
+            rmbg_png.push(rgba_img);
+        }
+        Ok(rmbg_png)
+    }
+}

src/models/rmbg2_0/mod.rs

@@ -1 +1,2 @@
+pub mod generate;
 pub mod model;

src/utils/img_utils.rs

@@ -285,3 +285,40 @@ pub fn img_smart_resize(
     }
     Ok((h_bar, w_bar))
 }
+
+pub fn img_transform_with_resize(
+    img: &DynamicImage,
+    h: u32,
+    w: u32,
+    mean: &Tensor,
+    std: &Tensor,
+    device: &Device,
+    dtype: DType,
+) -> Result<Tensor> {
+    let img_resize = img.resize_exact(w, h, imageops::FilterType::CatmullRom);
+    let img_tensor = img_transform(&img_resize, mean, std, device, dtype)?;
+    Ok(img_tensor)
+}
+
+pub fn float_tensor_to_dynamic_image(tensor: &Tensor) -> Result<DynamicImage> {
+    let tensor = tensor.affine(255.0, 0.0)?.clamp(0.0, 255.0)?;
+    let tensor_u8 = tensor.to_dtype(DType::U8)?.to_device(&Device::Cpu)?;
+    let (c, h, w) = tensor_u8.dims3()?;
+    match c {
+        1 => {
+            let tensor_u8 = tensor_u8.reshape((h, w))?;
+            let data: Vec<u8> = tensor_u8.flatten_all()?.to_vec1()?;
+            let img = ImageBuffer::from_raw(w as u32, h as u32, data)
+                .ok_or_else(|| anyhow!("Failed to create image buffer"))?;
+            Ok(DynamicImage::ImageLuma8(img))
+        }
+        3 => {
+            let tensor_u8 = tensor_u8.permute((1, 2, 0))?;
+            let data: Vec<u8> = tensor_u8.flatten_all()?.to_vec1()?;
+            let img = ImageBuffer::from_raw(w as u32, h as u32, data)
+                .ok_or_else(|| anyhow!("Failed to create image buffer"))?;
+            Ok(DynamicImage::ImageRgb8(img))
+        }
+        _ => Err(anyhow!(format!("Unsupported number of channels: {}", c))),
+    }
+}

src/utils/tensor_utils.rs

@@ -51,6 +51,9 @@ pub fn repeat_kv(xs: Tensor, n_rep: usize) -> Result<Tensor> {
 }
 
 pub fn split_tensor<D: Dim>(t: &Tensor, splits: &[usize], dim: D) -> Result<Vec<Tensor>> {
+    // 按给定长度切分tensor
+    // 例： t:(25), splits: [5, 10, 5, 5] dim: 0,
+    // 返回vec len=4, 其中tensor维度分别是:(5), (10), (5), (5)
     let dim = dim.to_index(t.shape(), "split")?;
     let mut split_res = Vec::new();
     let mut index = 0;
@@ -61,6 +64,28 @@ pub fn split_tensor<D: Dim>(t: &Tensor, splits: &[usize], dim: D) -> Result<Vec<
     Ok(split_res)
 }
 
+pub fn split_tensor_with_size<D: Dim>(
+    t: &Tensor,
+    splits_size: usize,
+    dim: D,
+) -> Result<Vec<Tensor>> {
+    // 按给定size切分tensor
+    // 例： t:(25), splits: 5 dim: 0,
+    // 返回vec len=5, 其中tensor维度分别是:(5), (5), (5), (5), (5)
+    let dim = dim.to_index(t.shape(), "split")?;
+    let mut split_res = Vec::new();
+    let dim_size = t.dim(dim)?;
+    assert_eq!(
+        dim_size % splits_size,
+        0,
+        "input tensor dim size % splits_size must be equal to 0"
+    );
+    for split in (0..dim_size).step_by(splits_size) {
+        split_res.push(t.narrow(dim, split, splits_size)?);
+    }
+    Ok(split_res)
+}
+
 pub fn safe_arg_sort_last_dim(t: &Tensor, ascending: bool) -> Result<Tensor> {
     // tensor在GPU上时，维度超过1024， arg_sort_last_dim方法会报错
     // 所以维度大于1024时，放到CPU上处理
@@ -230,7 +255,8 @@ pub fn get_not_equal_mask(input_ids: &Tensor, token_ids: u32) -> Result<Tensor>
 }
 
 pub fn get_equal_mask(input_ids: &Tensor, token_ids: u32) -> Result<Tensor> {
-    let image_token_id_tensor = Tensor::new(vec![token_ids], input_ids.device())?;
+    let image_token_id_tensor =
+        Tensor::new(vec![token_ids], input_ids.device())?.to_dtype(input_ids.dtype())?;
     let mask = input_ids
         .broadcast_eq(&image_token_id_tensor)?
         .to_dtype(candle_core::DType::U32)?;

tests/messy_test.rs

@@ -1,22 +1,62 @@
-use std::{path::PathBuf, str::FromStr};
-
 use anyhow::Result;
+use candle_core::Tensor;
 
 #[test]
 fn messy_test() -> Result<()> {
     // RUST_BACKTRACE=1 cargo test -F cuda messy_test -r -- --nocapture
-    let path_str = "file://./assets/img/ocr_test1.png";
+    let device = &candle_core::Device::Cpu;
-    let path = url::Url::from_str(path_str)?;
+    let x = Tensor::arange(0.0, 9.0, device)?;
-    let path = path.to_file_path();
+    println!("x: {}", x);
-    let path = match path {
+    let x = x
-        Ok(path) => path,
+        .unsqueeze(0)?
-        Err(_) => {
+        .unsqueeze(0)?
-            let mut path = path_str.to_owned();
+        .broadcast_as((5, 5, 9))?
-            path = path.split_off(7);
+        .reshape((5, 5, 3, 3))?;
-            PathBuf::from(path)
+    println!("x: {}", x);
-        }
+    let x = x.permute((0, 2, 1, 3))?;
-    };
+    println!("x: {}", x);
-    println!("to file path: {:?}", path);
+    let x = x.reshape((15, 15))?;
+    println!("x: {}", x);
+    // let xs = Tensor::rand(0.0, 5.0, (1, 1, 3, 3), device)?;
+    // println!("xs: {}", xs);
+    // let xs = xs.pad_with_zeros(3, 2, 2)?
+    //             .pad_with_zeros(2, 2, 2)?;
+    // println!("xs: {}", xs);
+    // let xs = Tensor::arange(0.0, 25.0, device)?;
+    // println!("xs: {}", xs);
+    // let splits = split_tensor_with_size(&xs, 5, 0)?;
+    // for v in splits {
+    //     println!("v: {}", v);
+    // }
+    // let xs = Tensor::arange(0.0, 25.0, device)?.broadcast_as((1, 1, 5, 5))?;
+    // println!("xs: {}", xs);
+    // let xs = xs.avg_pool2d(5)?;
+    // println!("xs: {}", xs);
+    // let xs = Tensor::rand(0.0, 1.0, (1, 4, 4, 2), device)?;
+    // println!("xs: {}", xs);
+    // let shape = Shape::from_dims(&[1, 2, 2, 2, 2, 2]);
+    // let xs = xs.reshape(shape)?;
+    // println!("xs: {}", xs);
+    // let x0 = xs.i((.., .., 0, .., 0, ..))?;
+    // let x1 = xs.i((.., .., 1, .., 0, ..))?;
+    // let x2 = xs.i((.., .., 0, .., 1, ..))?;
+    // let x3 = xs.i((.., .., 1, .., 1, ..))?;
+    // let xs = Tensor::cat(&[x0, x1, x2, x3], D::Minus1)?;
+    // println!("xs: {}", xs);
+    // let xs = xs.reshape((1, (), 4 * 2))?;
+    // println!("xs: {}", xs);
+    // let path_str = "file://./assets/img/ocr_test1.png";
+    // let path = url::Url::from_str(path_str)?;
+    // let path = path.to_file_path();
+    // let path = match path {
+    //     Ok(path) => path,
+    //     Err(_) => {
+    //         let mut path = path_str.to_owned();
+    //         path = path.split_off(7);
+    //         PathBuf::from(path)
+    //     }
+    // };
+    // println!("to file path: {:?}", path);
 
     // let device = &candle_core::Device::Cpu;
     // let t = Tensor::arange(0.0f32, 40.0, device)?.broadcast_as((1, 1, 40, 40))?;

tests/test_rmbg2_0.rs

@@ -0,0 +1,47 @@
+use std::time::Instant;
+
+use aha::models::rmbg2_0::generate::RMBG2_0;
+use aha_openai_dive::v1::resources::chat::ChatCompletionParameters;
+use anyhow::Result;
+
+#[test]
+fn rmbg2_0_generate() -> Result<()> {
+    // test with cuda: RUST_BACKTRACE=1 cargo test -F cuda rmbg2_0_generate -r -- --nocapture
+
+    let model_path = "/home/jhq/huggingface_model/AI-ModelScope/RMBG-2.0/";
+
+    let message = r#"
+    {
+        "model": "rmbg2.0",
+        "messages": [
+            {
+                "role": "user",
+                "content": [ 
+                    {
+                        "type": "image",
+                        "image_url": 
+                        {
+                            "url": "file://./assets/img/gougou.jpg"
+                        }
+                    }
+                ]
+            }
+        ]
+    }
+    "#;
+    let mes: ChatCompletionParameters = serde_json::from_str(message)?;
+    let i_start = Instant::now();
+    let model = RMBG2_0::init(model_path, None, None)?;
+    let i_duration = i_start.elapsed();
+    println!("Time elapsed in load model is: {:?}", i_duration);
+
+    let i_start = Instant::now();
+    let result = model.generate(mes)?;
+    for (i, img) in result.iter().enumerate() {
+        let _ = img.save(format!("rmbg_{i}.png"));
+    }
+    let i_duration = i_start.elapsed();
+    println!("Time elapsed in generate is: {:?}", i_duration);
+
+    Ok(())
+}

tests/test_voxcpm.rs

@@ -19,7 +19,7 @@ fn voxcpm_generate() -> Result<()> {
     let i_start = Instant::now();
     // let generate = voxcpm_generate.generate_simple("太阳当空照，花儿对我笑，小鸟说早早早".to_string())?;
     let generate = voxcpm_generate.generate(
-        "太阳当空照，花儿对我笑，小鸟说早早早".to_string(),
+        "VoxCPM is an innovative end-to-end TTS model from ModelBest, designed to generate highly realistic speech.".to_string(),
         Some("啥子小师叔，打狗还要看主人，你再要继续，我，就是你的对手".to_string()),
         Some("./assets/audio/voice_01.wav".to_string()),
         // Some("一定被灰太狼给吃了，我已经为他准备好了花圈了".to_string()),

tests/test_voxcpm1_5.rs

@@ -19,7 +19,7 @@ fn voxcpm1_5_generate() -> Result<()> {
     let i_start = Instant::now();
     // let generate = voxcpm_generate.generate_simple("太阳当空照，花儿对我笑，小鸟说早早早".to_string())?;
     let generate = voxcpm_generate.generate(
-        "太阳当空照，花儿对我笑，小鸟说早早早".to_string(),
+        "VoxCPM is an innovative end-to-end TTS model from ModelBest, designed to generate highly realistic speech.".to_string(),
         Some("啥子小师叔，打狗还要看主人，你再要继续，我就是你的对手".to_string()),
         Some("./assets/audio/voice_01.wav".to_string()),
         // Some("一定被灰太狼给吃了，我已经为他准备好了花圈了".to_string()),