add RMBGv2.0

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,
-    VarBuilder, conv2d, conv2d_no_bias, layer_norm, linear, linear_no_bias, rms_norm,
+    Activation, BatchNorm, BatchNormConfig, Conv2d, Conv2dConfig, LayerNorm, LayerNormConfig,
+    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 model;
+pub mod generate;
+pub mod model;