tckm.yan/aha
0
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

update interpolate

Browse Source
This commit is contained in:
jhqxxx
2026-03-01 12:44:14 +08:00
parent 6a73ee04c3
commit a61b899741
15 changed files with 598 additions and 486 deletions
Download Patch File Download Diff File Expand all files Collapse all files
README.md
+5
View File
@@ -25,6 +25,11 @@
aha is a high-performance, cross-platform AI inference engine built with Rust and the Candle framework. It brings state-of-the-art AI models to your local machine—no API keys, no cloud dependencies, just pure, fast AI running directly on your hardware.
## Changelog
### 2026-03-01
- update interpolate.rs
### 2026-02-24
- update candle version 0.9.2
### v0.2.0 (2026-02-05)
- Added Qwen3-ASR speech recognition model
src/models/deepseek_ocr/model.rs
+5 -4
View File
@@ -17,9 +17,10 @@ use crate::{
deepseek_ocr::config::{DeepseekOCRConfig, DeepseekV2Config},
},
position_embed::rope::RoPE,
utils::interpolate::{interpolate_bicubic, interpolate_linear_1d},
utils::tensor_utils::{
index_select_2d, interpolate_bicubic, interpolate_linear_1d, masked_scatter_dim0, nonzero,
onehot, prepare_causal_attention_mask, quick_gelu, topk,
index_select_2d, masked_scatter_dim0, nonzero, onehot, prepare_causal_attention_mask,
quick_gelu, topk,
},
};
@@ -480,8 +481,8 @@ impl ImageEncoderViT {
let new_pos_embed = interpolate_bicubic(
&old_pos_embed,
(tgt_size, tgt_size),
Some(true),
Some(false),
Some(true),
)?;
let new_pos_embed = new_pos_embed.permute((0, 2, 3, 1))?;
Ok(new_pos_embed)
@@ -571,8 +572,8 @@ impl CLIPVisionEmbeddings {
let new_pos_embed = interpolate_bicubic(
&old_pos_embed,
(tgt_size, tgt_size),
Some(true),
Some(false),
Some(true),
)?;
let new_pos_embed = new_pos_embed
.permute((0, 2, 3, 1))?
src/models/hunyuan_ocr/model.rs
+3 -4
View File
@@ -11,9 +11,8 @@ use crate::{
hunyuan_ocr::config::{HunYuanVLConfig, HunYuanVLVisionConfig},
},
position_embed::rope::{RoPE, apply_rotary_pos_emb, get_xd_cos_sin},
utils::tensor_utils::{
interpolate_bilinear, masked_scatter_dim0, prepare_causal_attention_mask, split_tensor,
},
utils::interpolate::interpolate_bilinear,
utils::tensor_utils::{masked_scatter_dim0, prepare_causal_attention_mask, split_tensor},
};
pub struct HunYuanVisionPatchEmbed {
@@ -83,7 +82,7 @@ impl HunYuanVisionPatchEmbed {
let grid_h = grid_i.i(1)?.to_scalar::<u32>()? as usize;
let grid_w = grid_i.i(2)?.to_scalar::<u32>()? as usize;
let patch_pos_embed_ =
interpolate_bilinear(&self.patch_pos_embed, (grid_h, grid_w), Some(false))?;
interpolate_bilinear(&self.patch_pos_embed, (grid_h, grid_w), Some(false), None)?;
let patch_pos_embed_ = patch_pos_embed_
.reshape((self.embed_dim, ()))?
.transpose(0, 1)?
src/models/mask_gct/model.rs
+2 -1
View File
@@ -9,7 +9,8 @@ use crate::{
common::{WNConv1d, get_conv1d, get_layer_norm},
mask_gct::config::SemanticCodec,
},
utils::tensor_utils::{interpolate_nearest_1d, l2_normalize},
utils::interpolate::interpolate_nearest_1d,
utils::tensor_utils::l2_normalize,
};
pub struct ConvNeXtBlock {
src/models/paddleocr_vl/model.rs
+4 -3
View File
@@ -16,9 +16,10 @@ use crate::{
},
},
position_embed::rope::{Qwen2_5VLTextRotaryEmbedding, Qwen2_5VisionRotaryEmbedding},
utils::interpolate::interpolate_bilinear,
utils::tensor_utils::{
get_vision_next_indices, interpolate_bilinear, masked_scatter_dim0, nonzero_index,
prepare_causal_attention_mask, zero_index,
get_vision_next_indices, masked_scatter_dim0, nonzero_index, prepare_causal_attention_mask,
zero_index,
},
};
@@ -143,7 +144,7 @@ impl SiglipVisionEmbeddings {
.reshape((1, sqrt_num_positions, sqrt_num_positions, self.embed_dim))?
.permute((0, 3, 1, 2))?;
let patch_pos_embed =
interpolate_bilinear(&patch_pos_embed, (new_height, new_width), Some(false))?;
interpolate_bilinear(&patch_pos_embed, (new_height, new_width), Some(false), None)?;
let patch_pos_embed =
patch_pos_embed
.permute((0, 2, 3, 1))?
src/models/rmbg2_0/model.rs
+19 -20
View File
@@ -10,9 +10,8 @@ use crate::{
Conv2dWithBN, TwoLinearMLP, deform_conv2d_kernel, get_batch_norm, get_conv2d,
get_layer_norm,
},
utils::tensor_utils::{
get_equal_mask, index_select_2d, interpolate_bilinear, split_tensor_with_size,
},
utils::interpolate::interpolate_bilinear,
utils::tensor_utils::{get_equal_mask, index_select_2d, split_tensor_with_size},
};
struct PatchEmbed {
@@ -996,7 +995,7 @@ impl ASPPDeformable {
let x5 = self.global_avg_pool_1.forward(&x5)?;
let x5 = self.global_avg_pool_2.forward_t(&x5, false)?.relu()?;
let (_, _, h, w) = x1.dims4()?;
let x5 = interpolate_bilinear(&x5, (h, w), Some(true))?;
let x5 = interpolate_bilinear(&x5, (h, w), Some(true), None)?;
let xs = Tensor::cat(
&[x1, x_aspp_deforms_0, x_aspp_deforms_1, x_aspp_deforms_2, x5],
1,
@@ -1246,7 +1245,7 @@ impl Decoder {
// let mut outs = vec![];
let patches_batch = self.get_patches_batch(x, x4)?;
let (_, _, x4_h, x4_w) = x4.dims4()?;
let patches_batch = interpolate_bilinear(&patches_batch, (x4_h, x4_w), Some(true))?;
let patches_batch = interpolate_bilinear(&patches_batch, (x4_h, x4_w), Some(true), None)?;
let ipt_blk5_out = self.ipt_blk5.forward(&patches_batch)?;
let x4 = Tensor::cat(&[x4, &ipt_blk5_out], 1)?;
let p4 = self.decoder_block4.forward(&x4)?;
@@ -1257,11 +1256,11 @@ impl Decoder {
let p4 = p4.broadcast_mul(&gdt_attn_4)?;
let (_, _, x3_h, x3_w) = x3.dims4()?;
let p4_inter = interpolate_bilinear(&p4, (x3_h, x3_w), Some(true))?;
let p4_inter = interpolate_bilinear(&p4, (x3_h, x3_w), Some(true), None)?;
let p3_ = self.lateral_block4.forward(x3)?;
let p3_ = p4_inter.add(&p3_)?;
let patches_batch = self.get_patches_batch(x, &p3_)?;
let patches_batch = interpolate_bilinear(&patches_batch, (x3_h, x3_w), Some(true))?;
let patches_batch = interpolate_bilinear(&patches_batch, (x3_h, x3_w), Some(true), None)?;
let ipt_blk4_out = self.ipt_blk4.forward(&patches_batch)?;
let p3_ = Tensor::cat(&[p3_, ipt_blk4_out], 1)?;
let p3 = self.decoder_block3.forward(&p3_)?;
@@ -1272,11 +1271,11 @@ impl Decoder {
let p3 = p3.broadcast_mul(&gdt_attn_3)?;
let (_, _, x2_h, x2_w) = x2.dims4()?;
let p3_inter = interpolate_bilinear(&p3, (x2_h, x2_w), Some(true))?;
let p3_inter = interpolate_bilinear(&p3, (x2_h, x2_w), Some(true), None)?;
let p2_ = self.lateral_block3.forward(x2)?;
let p2_ = p3_inter.add(&p2_)?;
let patches_batch = self.get_patches_batch(x, &p2_)?;
let patches_batch = interpolate_bilinear(&patches_batch, (x2_h, x2_w), Some(true))?;
let patches_batch = interpolate_bilinear(&patches_batch, (x2_h, x2_w), Some(true), None)?;
let ipt_blk3_out = self.ipt_blk3.forward(&patches_batch)?;
let p2_ = Tensor::cat(&[p2_, ipt_blk3_out], 1)?;
let p2 = self.decoder_block2.forward(&p2_)?;
@@ -1287,17 +1286,17 @@ impl Decoder {
let p2 = p2.broadcast_mul(&gdt_attn_2)?;
let (_, _, x1_h, x1_w) = x1.dims4()?;
let p2_inter = interpolate_bilinear(&p2, (x1_h, x1_w), Some(true))?;
let p2_inter = interpolate_bilinear(&p2, (x1_h, x1_w), Some(true), None)?;
let p1_ = self.lateral_block2.forward(x1)?;
let p1_ = p2_inter.add(&p1_)?;
let patches_batch = self.get_patches_batch(x, &p1_)?;
let patches_batch = interpolate_bilinear(&patches_batch, (x1_h, x1_w), Some(true))?;
let patches_batch = interpolate_bilinear(&patches_batch, (x1_h, x1_w), Some(true), None)?;
let ipt_blk2_out = self.ipt_blk2.forward(&patches_batch)?;
let p1_ = Tensor::cat(&[p1_, ipt_blk2_out], 1)?;
let p1_ = self.decoder_block1.forward(&p1_)?;
let (_, _, x_h, x_w) = x.dims4()?;
let p1_ = interpolate_bilinear(&p1_, (x_h, x_w), Some(true))?;
let p1_ = interpolate_bilinear(&p1_, (x_h, x_w), Some(true), None)?;
// let patches_batch = self.get_patches_batch(x, &p1_)?;
// let patches_batch = interpolate_bilinear(&patches_batch, (x_h, x_w), Some(true))?;
let ipt_blk1_out = self.ipt_blk1.forward(x)?;
@@ -1348,27 +1347,27 @@ impl BiRefNet {
};
let (_, _, h, w) = xs.dims4()?;
let cat_xs = interpolate_bilinear(xs, (h / 2, w / 2), Some(true))?;
let cat_xs = interpolate_bilinear(xs, (h / 2, w / 2), Some(true), None)?;
let [ref x1_, ref x2_, ref x3_, ref x4_] = self.bb.forward(&cat_xs)?[..] else {
return Err(anyhow!(format!(
"swintransformer output exactly 3 elements"
)));
};
let (_, _, x1_h, x1_w) = x1.dims4()?;
let x1_ = interpolate_bilinear(x1_, (x1_h, x1_w), Some(true))?;
let x1_ = interpolate_bilinear(x1_, (x1_h, x1_w), Some(true), None)?;
let x1 = Tensor::cat(&[x1, &x1_], 1)?;
let (_, _, x2_h, x2_w) = x2.dims4()?;
let x2_ = interpolate_bilinear(x2_, (x2_h, x2_w), Some(true))?;
let x2_ = interpolate_bilinear(x2_, (x2_h, x2_w), Some(true), None)?;
let x2 = Tensor::cat(&[x2, &x2_], 1)?;
let (_, _, x3_h, x3_w) = x3.dims4()?;
let x3_ = interpolate_bilinear(x3_, (x3_h, x3_w), Some(true))?;
let x3_ = interpolate_bilinear(x3_, (x3_h, x3_w), Some(true), None)?;
let x3 = Tensor::cat(&[x3, &x3_], 1)?;
let (_, _, x4_h, x4_w) = x4.dims4()?;
let x4_ = interpolate_bilinear(x4_, (x4_h, x4_w), Some(true))?;
let x4_ = interpolate_bilinear(x4_, (x4_h, x4_w), Some(true), None)?;
let x4 = Tensor::cat(&[x4, &x4_], 1)?;
let x1_resize = interpolate_bilinear(&x1, (x4_h, x4_w), Some(true))?;
let x2_resize = interpolate_bilinear(&x2, (x4_h, x4_w), Some(true))?;
let x3_resize = interpolate_bilinear(&x3, (x4_h, x4_w), Some(true))?;
let x1_resize = interpolate_bilinear(&x1, (x4_h, x4_w), Some(true), None)?;
let x2_resize = interpolate_bilinear(&x2, (x4_h, x4_w), Some(true), None)?;
let x3_resize = interpolate_bilinear(&x3, (x4_h, x4_w), Some(true), None)?;
let x4 = Tensor::cat(&[x1_resize, x2_resize, x3_resize, x4], 1)?;
let x4 = self.squeeze_module_0.forward(&x4)?;
src/utils/interpolate.rs
+498
View File
@@ -0,0 +1,498 @@
use anyhow::{Result, anyhow};
use candle_core::{DType, Tensor};
fn compute_scale(input_size: usize, output_size: usize, align_corners: bool) -> f32 {
if align_corners && output_size > 1 {
(input_size - 1) as f32 / (output_size - 1) as f32
} else {
input_size as f32 / output_size as f32
}
}
pub fn compute_1d_coords(
input_size: usize,
output_size: usize,
align_corner: Option<bool>,
) -> Result<Vec<f32>> {
if input_size == 0 {
return Err(anyhow!("input_size must be > 0"));
}
if output_size == 0 {
return Err(anyhow!("output_size must be > 0"));
}
if input_size == 1 {
return Ok(vec![0f32; output_size]);
}
let align_corners = align_corner.unwrap_or(false);
let scale = compute_scale(input_size, output_size, align_corners);
if align_corners {
Ok((0..output_size).map(|i| i as f32 * scale).collect())
} else {
Ok((0..output_size)
.map(|i| {
let coord = (i as f32 + 0.5) * scale - 0.5;
coord.clamp(0.0, (input_size - 1) as f32)
})
.collect())
}
}
pub fn interpolate_nearest_1d(t: &Tensor, target_size: usize) -> Result<Tensor> {
// t: [b, channels, features]
if t.rank() != 3 {
return Err(anyhow::anyhow!(
"Input rank must have equal to 3 dimensions"
));
}
let (bs, channels, orig_size) = t.dims3()?;
if orig_size == target_size {
return Ok(t.clone());
}
let input_data = t.to_vec3::<f32>()?;
let mut output_data = vec![vec![vec![0.0f32; target_size]; channels]; bs];
for b in 0..bs {
for c in 0..channels {
for i in 0..target_size {
// Nearest neighbor: round to nearest integer coordinate
let coord = if target_size == 1 {
(orig_size - 1) as f32 / 2.0
} else {
(i as f32 + 0.5) * (orig_size as f32 / target_size as f32) - 0.5
};
let nearest_idx = coord.round() as usize;
let clamped_idx = nearest_idx.clamp(0, orig_size - 1);
output_data[b][c][i] = input_data[b][c][clamped_idx];
}
}
}
let output = Tensor::new(output_data, t.device())?.to_dtype(t.dtype())?;
Ok(output)
}
pub fn interpolate_nearest_2d(input: &Tensor, target_size: (usize, usize)) -> Result<Tensor> {
// t: [batch, channels, height, width]
if input.rank() != 4 {
return Err(anyhow::anyhow!(
"Input tensor must have 4 dimensions [N, C, H, W], got rank {}",
input.rank()
));
}
let (bs, channels, orig_h, orig_w) = input.dims4()?;
let (target_h, target_w) = target_size;
// 如果尺寸相同，直接返回克隆
if orig_h == target_h && orig_w == target_w {
return Ok(input.clone());
}
// 将输入数据转为 Vec 以便索引
let dim0 = bs * channels;
let input_3dim = input.reshape((dim0, orig_h, orig_w))?;
let input_data = input_3dim.to_dtype(DType::F32)?.to_vec3::<f32>()?;
let mut output_data = vec![vec![vec![0.0f32; target_w]; target_h]; dim0];
for c in 0..dim0 {
for i in 0..target_h {
// 计算高度方向的最近邻索引
let coord_h = if target_h == 1 {
(orig_h - 1) as f32 / 2.0
} else {
(i as f32 + 0.5) * (orig_h as f32 / target_h as f32) - 0.5
};
let nearest_h = coord_h.round() as usize;
let clamped_h = nearest_h.clamp(0, orig_h - 1);
for j in 0..target_w {
// 计算宽度方向的最近邻索引
let coord_w = if target_w == 1 {
(orig_w - 1) as f32 / 2.0
} else {
(j as f32 + 0.5) * (orig_w as f32 / target_w as f32) - 0.5
};
let nearest_w = coord_w.round() as usize;
let clamped_w = nearest_w.clamp(0, orig_w - 1);
output_data[c][i][j] = input_data[c][clamped_h][clamped_w];
}
}
}
let output = Tensor::new(output_data, input.device())?
.reshape((bs, channels, target_h, target_w))?
.to_dtype(input.dtype())?
.contiguous()?;
Ok(output)
}
pub fn interpolate_linear_1d(
t: &Tensor,
target_size: usize,
align_corner: Option<bool>,
) -> Result<Tensor> {
// t: [b, channels, features]
if t.rank() != 3 {
return Err(anyhow::anyhow!(
"Input rank must have equal to 3 dimensions"
));
}
let (bs, channels, orig_size) = t.dims3()?;
if orig_size == target_size {
return Ok(t.clone());
}
let coords = compute_1d_coords(orig_size, target_size, align_corner)?;
let input_data = t.to_vec3::<f32>()?;
let mut output_data = vec![vec![vec![0.0f32; target_size]; channels]; bs];
for b in 0..bs {
for c in 0..channels {
for (i, &coord) in coords.iter().enumerate() {
let coord = coord.clamp(0.0, (orig_size - 1) as f32);
let x0 = coord.floor() as usize;
let x1 = (x0 + 1).min(orig_size - 1);
let weight = coord - x0 as f32;
let value0 = input_data[b][c][x0];
let value1 = input_data[b][c][x1];
output_data[b][c][i] = value0 * (1.0 - weight) + value1 * weight;
}
}
}
let output = Tensor::new(output_data, t.device())?.to_dtype(t.dtype())?;
Ok(output)
}
fn antialias_filter(x: f32) -> f32 {
let x = x.abs();
if x < 1.0 { 1.0 - x } else { 0.0 }
}
pub fn interpolate_bilinear(
input: &Tensor,
target_size: (usize, usize),
align_corner: Option<bool>,
antialias: Option<bool>,
) -> Result<Tensor> {
if input.rank() != 4 {
return Err(anyhow::anyhow!(
"Input rank must have equal to 4 dimensions [b, c, h, w]"
));
}
let (_, _, input_height, input_width) = input.dims4()?;
let (target_height, target_width) = target_size;
if input_height == target_height && input_width == target_width {
return Ok(input.clone());
}
let output = if antialias.unwrap_or(false)
&& (target_height < input_height || target_width < input_width)
{
interpolate_bilinear_antialias(input, target_size)?
} else {
interpolate_bilinear_standard(input, target_size, align_corner)?
};
let output = output.to_dtype(input.dtype())?.to_device(input.device())?;
Ok(output)
}
pub fn interpolate_bilinear_standard(
input: &Tensor,
target_size: (usize, usize),
align_corner: Option<bool>,
) -> Result<Tensor> {
let (bs, channels, input_height, input_width) = input.dims4()?;
let (target_height, target_width) = target_size;
// 计算两个维度的采样坐标
let coords_h = compute_1d_coords(input_height, target_height, align_corner)?;
let coords_w = compute_1d_coords(input_width, target_width, align_corner)?;
let dim0 = bs * channels;
let input_3dim = input.reshape((dim0, input_height, input_width))?;
let input_data = input_3dim.to_dtype(DType::F32)?.to_vec3::<f32>()?;
let mut output_data = vec![vec![vec![0.0f32; target_width]; target_height]; dim0];
for c in 0..dim0 {
for (i, &coord_h) in coords_h.iter().enumerate() {
let coord_h = coord_h.clamp(0.0, (input_height - 1) as f32);
let y0 = coord_h.floor() as usize;
let y1 = (y0 + 1).min(input_height - 1);
let dy = coord_h - y0 as f32;
for (j, &coord_w) in coords_w.iter().enumerate() {
let coord_w = coord_w.clamp(0.0, (input_width - 1) as f32);
let x0 = coord_w.floor() as usize;
let x1 = (x0 + 1).min(input_width - 1);
let dx = coord_w - x0 as f32;
let q00 = input_data[c][y0][x0];
let q01 = input_data[c][y0][x1];
let q10 = input_data[c][y1][x0];
let q11 = input_data[c][y1][x1];
output_data[c][i][j] = q00 * (1.0 - dx) * (1.0 - dy)
+ q01 * dx * (1.0 - dy)
+ q10 * (1.0 - dx) * dy
+ q11 * dx * dy;
}
}
}
let output = Tensor::new(output_data, input.device())?
.reshape((bs, channels, target_height, target_width))?
.to_dtype(input.dtype())?
.contiguous()?;
Ok(output)
}
pub fn interpolate_bilinear_antialias(
input: &Tensor,
target_size: (usize, usize),
) -> Result<Tensor> {
let (bs, channels, input_height, input_width) = input.dims4()?;
let (target_height, target_width) = target_size;
let scale_h = input_height as f32 / target_height as f32;
let scale_w = input_width as f32 / target_width as f32;
let dim0 = bs * channels;
let input_3dim = input.reshape((dim0, input_height, input_width))?;
let input_data = input_3dim.to_dtype(DType::F32)?.to_vec3::<f32>()?;
let mut output_data = vec![vec![vec![0.0f32; target_width]; target_height]; dim0];
let support_size = scale_h.max(scale_w);
for c in 0..dim0 {
for out_y in 0..target_height {
let center_y = (out_y as f32 + 0.5) * scale_h - 0.5;
let start_y = (center_y - support_size).max(0.0) as usize;
let end_y = (center_y + support_size).min(input_height as f32 - 1.0) as usize;
for out_x in 0..target_width {
let center_x = (out_x as f32 + 0.5) * scale_w - 0.5;
let start_x = (center_x - support_size).max(0.0) as usize;
let end_x = (center_x + support_size).min(input_width as f32 - 1.0) as usize;
let mut total_weight = 0.0;
let mut weighted_sum = 0.0;
for src_y in start_y..end_y + 1 {
for src_x in start_x..end_x + 1 {
let dist_x = (src_x as f32 - center_x).abs();
let dist_y = (src_y as f32 - center_y).abs();
let weight_x = antialias_filter(dist_x / scale_w);
let weight_y = antialias_filter(dist_y / scale_h);
let weight = weight_x * weight_y;
weighted_sum += input_data[c][src_y][src_x] * weight;
total_weight += weight;
}
}
let result = if total_weight > 0.0 {
weighted_sum / total_weight
} else {
let y = center_y.round().clamp(0.0, (input_height - 1) as f32) as usize;
let x = center_x.round().clamp(0.0, (input_width - 1) as f32) as usize;
input_data[c][y][x]
};
output_data[c][out_y][out_x] = result;
}
}
}
let output = Tensor::new(output_data, input.device())?
.reshape((bs, channels, target_height, target_width))?
.to_dtype(input.dtype())?
.contiguous()?;
Ok(output)
}
fn bicubic_filter(x: f32, a: f32) -> f32 {
let x = x.abs();
if x < 1.0 {
((a + 2.0) * x - (a + 3.0)) * x * x + 1.0
} else if x < 2.0 {
(((x - 5.0) * x + 8.0) * x - 4.0) * a
} else {
0.0
}
}
pub fn interpolate_bicubic_antialias(
input: &Tensor,
target_size: (usize, usize),
) -> Result<Tensor> {
let (bs, channels, input_height, input_width) = input.dims4()?;
let (target_height, target_width) = target_size;
let scale_h = input_height as f32 / target_height as f32;
let scale_w = input_width as f32 / target_width as f32;
// tensor没有to_vec4, 所以把bs和channels先合在一起
let dim0 = bs * channels;
let input_3dim = input.reshape((dim0, input_height, input_width))?;
let input_data = input_3dim.to_dtype(DType::F32)?.to_vec3::<f32>()?;
let mut output_data = vec![vec![vec![0.0f32; target_width]; target_height]; dim0];
let scale = scale_h.max(scale_w);
let support_size = if scale >= 1.0 {
(2.0 * scale).ceil()
} else {
2.0
};
for c in 0..dim0 {
for out_y in 0..target_height {
let center_y = (out_y as f32 + 0.5) * scale_h - 0.5;
let start_y = (center_y - support_size).ceil() as isize;
let end_y = (center_y + support_size).floor() as isize;
for out_x in 0..target_width {
let center_x = (out_x as f32 + 0.5) * scale_w - 0.5;
let start_x = (center_x - support_size).ceil() as isize;
let end_x = (center_x + support_size).floor() as isize;
let mut sum = 0.0;
let mut weight_sum = 0.0;
for iy in start_y..end_y + 1 {
for ix in start_x..end_x + 1 {
if iy >= 0
&& iy < input_height as isize
&& ix >= 0
&& ix < input_width as isize
{
let dx = (ix as f32 - center_x).abs();
let dy = (iy as f32 - center_y).abs();
let wx = bicubic_filter(dx / scale_w.max(1.0), -0.5);
let wy = bicubic_filter(dy / scale_h.max(1.0), -0.5);
let weight = wx * wy;
sum += input_data[c][iy as usize][ix as usize] * weight;
weight_sum += weight;
}
}
}
if weight_sum > 0.0 {
output_data[c][out_y][out_x] = sum / weight_sum;
} else {
let y = center_y.round().clamp(0.0, (input_height - 1) as f32) as usize;
let x = center_x.round().clamp(0.0, (input_width - 1) as f32) as usize;
output_data[c][out_y][out_x] = input_data[c][y][x];
}
}
}
}
let output = Tensor::new(output_data, input.device())?
.reshape((bs, channels, target_height, target_width))?
.to_dtype(input.dtype())?
.contiguous()?;
Ok(output)
}
// 三次卷积函数1
fn cubic_convolution1(x: f64, a: f64) -> f64 {
((a + 2.0) * x - (a + 3.0)) * x * x + 1.0
}
// 三次卷积函数2
fn cubic_convolution2(x: f64, a: f64) -> f64 {
(((x - 5.0) * x + 8.0) * x - 4.0) * a
}
fn get_cubic_coefficients(t: f64, a: f64) -> [f64; 4] {
let coeff0 = cubic_convolution2(t + 1.0, a);
let coeff1 = cubic_convolution1(t, a);
let coeff2 = cubic_convolution1(1.0 - t, a);
let coeff3 = cubic_convolution2(1.0 - t + 1.0, a);
[coeff0, coeff1, coeff2, coeff3]
}
fn cubic_interp1d(x0: f32, x1: f32, x2: f32, x3: f32, t: f64, a: f64) -> f32 {
let coeffs = get_cubic_coefficients(t, a);
x0 * coeffs[0] as f32 + x1 * coeffs[1] as f32 + x2 * coeffs[2] as f32 + x3 * coeffs[3] as f32
}
pub fn interpolate_bicubic_standard(
input: &Tensor,
target_size: (usize, usize),
align_corner: Option<bool>,
) -> Result<Tensor> {
let (bs, channels, input_height, input_width) = input.dims4()?;
let (target_height, target_width) = target_size;
let align_corners = align_corner.unwrap_or(false);
let scale_h = compute_scale(input_height, target_height, align_corners);
let scale_w = compute_scale(input_width, target_width, align_corners);
let dim0 = bs * channels;
let input_3dim = input.reshape((dim0, input_height, input_width))?;
let input_data = input_3dim.to_dtype(DType::F32)?.to_vec3::<f32>()?;
let mut output_data = vec![vec![vec![0.0f32; target_width]; target_height]; dim0];
for c in 0..dim0 {
for out_y in 0..target_height {
let center_y = if align_corners {
out_y as f32 * scale_h
} else {
(out_y as f32 + 0.5) * scale_h - 0.5
}
.clamp(0.0, (input_height - 1) as f32);
let in_y = center_y.floor() as isize;
let t_y = center_y - in_y as f32;
for out_x in 0..target_width {
let center_x = if align_corners {
out_x as f32 * scale_w
} else {
(out_x as f32 + 0.5) * scale_w - 0.5
}
.clamp(0.0, (input_width - 1) as f32);
let in_x: isize = center_x.floor() as isize;
let t_x = center_x - in_x as f32;
let mut coefficients = [0.0; 4];
for k in 0..4 {
let row = (in_y - 1 + k as isize).clamp(0, input_height as isize - 1) as usize;
let x_minus_1 =
input_data[c][row][(in_x - 1).clamp(0, input_width as isize - 1) as usize];
let x_plus_0 =
input_data[c][row][in_x.clamp(0, input_width as isize - 1) as usize];
let x_plus_1 =
input_data[c][row][(in_x + 1).clamp(0, input_width as isize - 1) as usize];
let x_plus_2 =
input_data[c][row][(in_x + 2).clamp(0, input_width as isize - 1) as usize];
coefficients[k] =
cubic_interp1d(x_minus_1, x_plus_0, x_plus_1, x_plus_2, t_x as f64, -0.75);
}
output_data[c][out_y][out_x] = cubic_interp1d(
coefficients[0],
coefficients[1],
coefficients[2],
coefficients[3],
t_y as f64,
-0.75,
);
}
}
}
let output = Tensor::new(output_data, input.device())?
.reshape((bs, channels, target_height, target_width))?
.to_dtype(input.dtype())?
.contiguous()?;
Ok(output)
}
pub fn interpolate_bicubic(
input: &Tensor,
target_size: (usize, usize),
align_corner: Option<bool>,
antialias: Option<bool>,
) -> Result<Tensor> {
if input.rank() != 4 {
return Err(anyhow::anyhow!(
"Input rank must have at least 3 dimensions"
));
}
let (_, _, input_height, input_width) = input.dims4()?;
let (output_height, output_width) = target_size;
if output_height == input_height && output_width == input_width {
return Ok(input.clone());
}
// let input_squeeze = input.squeeze(0)?;
let output = if antialias.unwrap_or(false)
&& (input_height > output_height || input_width > output_width)
{
interpolate_bicubic_antialias(input, target_size)?
} else {
interpolate_bicubic_standard(input, target_size, align_corner)?
};
let output = output.to_dtype(input.dtype())?.to_device(input.device())?;
Ok(output)
}
src/utils/mod.rs
+1
View File
@@ -1,5 +1,6 @@
pub mod audio_utils;
pub mod img_utils;
pub mod interpolate;
pub mod tensor_utils;
pub mod video_utils;
src/utils/tensor_utils.rs
-428
View File
@@ -426,434 +426,6 @@ pub fn mask_index_add(original: &Tensor, mask: &Tensor, add: &Tensor) -> Result<
Ok(xs)
}
pub fn compute_1d_coords(
input_size: usize,
output_size: usize,
align_corner: Option<bool>,
) -> Result<Vec<f32>> {
if input_size == 0 {
return Err(anyhow!("input_size must be > 0"));
}
if output_size == 0 {
return Err(anyhow!("output_size must be > 0"));
}
if input_size == 1 {
return Ok(vec![0f32; output_size]);
}
let align_corners = align_corner.unwrap_or(false);
if align_corners {
let scale = (input_size - 1) as f32 / (output_size - 1) as f32;
Ok((0..output_size).map(|i| i as f32 * scale).collect())
} else {
let scale = input_size as f32 / output_size as f32;
Ok((0..output_size)
.map(|i| {
let coord = (i as f32 + 0.5) * scale - 0.5;
coord.clamp(0.0, (input_size - 1) as f32)
})
.collect())
}
}
pub fn interpolate_linear_1d(
t: &Tensor,
target_size: usize,
align_corner: Option<bool>,
) -> Result<Tensor> {
// t: [b, channels, features]
if t.rank() != 3 {
return Err(anyhow::anyhow!(
"Input rank must have equal to 3 dimensions"
));
}
let shape = t.dims();
let orig_size = shape[shape.len() - 1];
if orig_size == target_size {
return Ok(t.clone());
}
let (bs, channels, _) = t.dims3()?;
let mut output = Tensor::zeros((bs, channels, target_size), t.dtype(), t.device())?;
let coords = compute_1d_coords(orig_size, target_size, align_corner)?;
for b in 0..bs {
for c in 0..channels {
let input_slice = t.i((b, c))?;
let mut out_i = Vec::new();
// for x_out in 0..target_size {
for &coord in coords.iter().take(target_size) {
let coord = if coord < 0.0 { 0.0 } else { coord };
let x0 = coord.floor() as usize;
let x1 = std::cmp::min(x0 + 1, orig_size - 1);
let weight = (coord - x0 as f32) as f64;
let value0 = input_slice.get(x0)?;
let value1 = input_slice.get(x1)?;
let interpolated =
(value0.affine(1.0 - weight, 0.0)? + value1.affine(weight, 0.0)?)?;
out_i.push(interpolated);
}
let out_i = Tensor::stack(&out_i, 0)?.unsqueeze(0)?.unsqueeze(0)?;
output = output.slice_assign(&[(b..b + 1), (c..c + 1), (0..target_size)], &out_i)?;
}
}
output = output.contiguous()?;
Ok(output)
}
pub fn interpolate_nearest_1d(t: &Tensor, target_size: usize) -> Result<Tensor> {
// t: [b, channels, features]
if t.rank() != 3 {
return Err(anyhow::anyhow!(
"Input rank must have equal to 3 dimensions"
));
}
let (bs, channels, orig_size) = t.dims3()?;
if orig_size == target_size {
return Ok(t.clone());
}
let coords = compute_1d_coords(orig_size, target_size, None)?;
let input_data = t.to_vec3::<f32>()?;
let mut output_data = vec![vec![vec![0.0f32; target_size]; channels]; bs];
for b in 0..bs {
for c in 0..channels {
for (i, &coord) in coords.iter().enumerate() {
// Nearest neighbor: round to nearest integer coordinate
let nearest_idx = coord.round() as usize;
let clamped_idx = nearest_idx.min(orig_size - 1);
let value = input_data[b][c][clamped_idx];
output_data[b][c][i] = value;
}
}
}
let output = Tensor::new(output_data, t.device())?.to_dtype(t.dtype())?;
Ok(output)
}
pub fn interpolate_bilinear(
input: &Tensor,
target_size: (usize, usize),
align_corner: Option<bool>,
) -> Result<Tensor> {
// input: [b, channels, height, width]
if input.rank() != 4 {
return Err(anyhow::anyhow!(
"Input rank must have equal to 4 dimensions [b, c, h, w]"
));
}
let (bs, channels, input_height, input_width) = input.dims4()?;
let (target_height, target_width) = target_size;
// If size is the same, return clone
if input_height == target_height && input_width == target_width {
return Ok(input.clone());
}
let align_corners = align_corner.unwrap_or(false);
// Compute scaling factors
let height_scale = if align_corners && target_height > 1 {
(input_height - 1) as f64 / (target_height - 1) as f64
} else {
input_height as f64 / target_height as f64
};
let width_scale = if align_corners && target_width > 1 {
(input_width - 1) as f64 / (target_width - 1) as f64
} else {
input_width as f64 / target_width as f64
};
let dim0 = bs * channels;
let input_3dim = input.reshape((dim0, input_height, input_width))?;
let input_data = input_3dim.to_dtype(DType::F32)?.to_vec3::<f32>()?;
let mut output_data = vec![vec![vec![0.0f32; target_width]; target_height]; dim0];
for c in 0..dim0 {
for out_y in 0..target_height {
let src_y = if align_corners {
out_y as f64 * height_scale
} else {
(out_y as f64 + 0.5) * height_scale - 0.5
};
let src_y = src_y.max(0.0).min((input_height - 1) as f64);
let y0 = src_y.floor() as usize;
let y1 = (y0 + 1).min(input_height - 1);
let dy = (src_y - y0 as f64) as f32;
for out_x in 0..target_width {
let src_x = if align_corners {
out_x as f64 * width_scale
} else {
(out_x as f64 + 0.5) * width_scale - 0.5
};
let src_x = src_x.max(0.0).min((input_width - 1) as f64);
let x0 = src_x.floor() as usize;
let x1 = (x0 + 1).min(input_width - 1);
let q00 = input_data[c][y0][x0];
let q01 = input_data[c][y0][x1];
let q10 = input_data[c][y1][x0];
let q11 = input_data[c][y1][x1];
let dx = (src_x - x0 as f64) as f32;
let interpolated = q00 * (1.0 - dx) * (1.0 - dy)
+ q01 * dx * (1.0 - dy)
+ q10 * (1.0 - dx) * dy
+ q11 * dx * dy;
output_data[c][out_y][out_x] = interpolated;
}
}
}
let output = Tensor::new(output_data, input.device())?
.reshape((bs, channels, target_height, target_width))?
.to_dtype(input.dtype())?;
Ok(output.contiguous()?)
}
fn compute_scale(input_size: usize, output_size: usize, align_corners: bool) -> f64 {
if align_corners && output_size > 1 {
(input_size - 1) as f64 / (output_size - 1) as f64
} else {
input_size as f64 / output_size as f64
}
}
fn bicubic_filter(x: f64) -> f64 {
let a = -0.75;
let x = x.abs();
if x < 1.0 {
((a + 2.0) * x - (a + 3.0)) * x * x + 1.0
} else if x < 2.0 {
(((x - 5.0) * x + 8.0) * x - 4.0) * a
} else {
0.0
}
}
pub fn interpolate_bicubic_antialias(
input: &Tensor,
batch_size: usize,
channels: usize,
input_height: usize,
input_width: usize,
output_height: usize,
output_width: usize,
height_scale: f64,
width_scale: f64,
align_corners: bool,
) -> Result<Tensor> {
// tensor没有to_vec4, 所以把bs和channels先合在一起
let dim0 = batch_size * channels;
let input_3dim = input.reshape((dim0, input_height, input_width))?;
let input_data = input_3dim.to_dtype(DType::F32)?.to_vec3::<f32>()?;
let mut output_data = vec![vec![vec![0.0f32; output_width]; output_height]; dim0];
let support = 2.0 * height_scale.max(width_scale);
for c in 0..dim0 {
for out_y in 0..output_height {
let center_y = if align_corners {
out_y as f64 * height_scale
} else {
(out_y as f64 + 0.5) * height_scale - 0.5
};
let start_y = (center_y - support).ceil() as isize;
let end_y = (center_y + support).floor() as isize;
for out_x in 0..output_width {
let center_x = if align_corners {
out_x as f64 * width_scale
} else {
(out_x as f64 + 0.5) * width_scale - 0.5
};
let mut sum = 0.0;
let mut weight_sum = 0.0;
let start_x = (center_x - support).ceil() as isize;
let end_x = (center_x + support).floor() as isize;
for iy in start_y..end_y {
for ix in start_x..end_x {
if iy >= 0
&& iy < input_height as isize
&& ix >= 0
&& ix < input_width as isize
{
let dx = (ix as f64 - center_x).abs();
let dy = (iy as f64 - center_y).abs();
let wx = bicubic_filter(dx / width_scale.max(1.0));
let wy = bicubic_filter(dy / height_scale.max(1.0));
let weight = (wx * wy) as f32;
sum += input_data[c][iy as usize][ix as usize] * weight;
weight_sum += weight;
}
}
}
if weight_sum > 0.0 {
output_data[c][out_y][out_x] = sum / weight_sum;
} else {
output_data[c][out_y][out_x] = 0.0;
}
}
}
}
let output = Tensor::new(output_data, input.device())?
.reshape((batch_size, channels, output_height, output_width))?
.to_dtype(input.dtype())?;
Ok(output)
}
fn get_cubic_coefficients(t: f64) -> [f64; 4] {
let a = -0.75;
let x1 = t;
let coeff0 = cubic_convolution2(x1 + 1.0, a);
let coeff1 = cubic_convolution1(x1, a);
let x2 = 1.0 - t;
let coeff2 = cubic_convolution1(x2, a);
let coeff3 = cubic_convolution2(x2 + 1.0, a);
[coeff0, coeff1, coeff2, coeff3]
}
// 三次卷积函数1
fn cubic_convolution1(x: f64, a: f64) -> f64 {
((a + 2.0) * x - (a + 3.0)) * x * x + 1.0
}
// 三次卷积函数2
fn cubic_convolution2(x: f64, a: f64) -> f64 {
((a * x - 5.0 * a) * x + 8.0 * a) * x - 4.0 * a
}
fn cubic_interp1d(x0: f32, x1: f32, x2: f32, x3: f32, t: f64) -> f32 {
let coeffs = get_cubic_coefficients(t);
x0 * coeffs[0] as f32 + x1 * coeffs[1] as f32 + x2 * coeffs[2] as f32 + x3 * coeffs[3] as f32
}
pub fn interpolate_bicubic_standard(
input: &Tensor,
batch_size: usize,
channels: usize,
input_height: usize,
input_width: usize,
output_height: usize,
output_width: usize,
height_scale: f64,
width_scale: f64,
align_corners: bool,
) -> Result<Tensor> {
// tensor没有to_vec4, 所以把bs和channels先合在一起
let dim0 = batch_size * channels;
let input_3dim = input.reshape((dim0, input_height, input_width))?;
let input_data = input_3dim.to_dtype(DType::F32)?.to_vec3::<f32>()?;
let mut output_data = vec![vec![vec![0.0f32; output_width]; output_height]; dim0];
for c in 0..dim0 {
for out_y in 0..output_height {
let center_y = if align_corners {
out_y as f64 * height_scale
} else {
(out_y as f64 + 0.5) * height_scale - 0.5
};
let in_y = center_y.floor() as isize;
let t_y = center_y - in_y as f64;
for out_x in 0..output_width {
let center_x = if align_corners {
out_x as f64 * width_scale
} else {
(out_x as f64 + 0.5) * width_scale - 0.5
};
let in_x = center_x.floor() as isize;
let t_x = center_x - in_x as f64;
let mut coefficients = [0.0; 4];
// for k in 0..4 {
for (k, coefficients_k) in coefficients.iter_mut().enumerate() {
let row = (in_y - 1 + k as isize)
.max(0)
.min(input_height as isize - 1) as usize;
let x_minus_1 = input_data[c][row]
[(in_x - 1).max(0).min(input_width as isize - 1) as usize];
let x_plus_0 =
input_data[c][row][in_x.max(0).min(input_width as isize - 1) as usize];
let x_plus_1 = input_data[c][row]
[(in_x + 1).max(0).min(input_width as isize - 1) as usize];
let x_plus_2 = input_data[c][row]
[(in_x + 2).max(0).min(input_width as isize - 1) as usize];
// coefficients[k] = cubic_interp1d(x_minus_1, x_plus_0, x_plus_1, x_plus_2, t_x);
*coefficients_k = cubic_interp1d(x_minus_1, x_plus_0, x_plus_1, x_plus_2, t_x);
}
output_data[c][out_y][out_x] = cubic_interp1d(
coefficients[0],
coefficients[1],
coefficients[2],
coefficients[3],
t_y,
);
}
}
}
let output = Tensor::new(output_data, input.device())?
.reshape((batch_size, channels, output_height, output_width))?
.to_dtype(input.dtype())?;
Ok(output)
}
pub fn interpolate_bicubic(
input: &Tensor,
target_size: (usize, usize),
antialias: Option<bool>,
align_corner: Option<bool>,
) -> Result<Tensor> {
if input.rank() != 4 {
return Err(anyhow::anyhow!(
"Input rank must have at least 3 dimensions"
));
}
// if input.dim(0)? != 1 {
// return Err(anyhow::anyhow!("Input batch_size must be 1"));
// }
let (batch_size, channels, input_height, input_width) = input.dims4()?;
let (output_height, output_width) = target_size;
if output_height == input_height && output_width == input_width {
return Ok(input.clone());
}
let align_corners = match align_corner {
Some(true) => true,
Some(false) => false,
None => false,
};
let height_scale = compute_scale(input_height, output_height, align_corners);
let width_scale = compute_scale(input_width, output_width, align_corners);
// let input_squeeze = input.squeeze(0)?;
let output = if let Some(antialias_) = antialias
&& antialias_
&& (input_height > output_height || input_width > output_width)
{
interpolate_bicubic_antialias(
input,
batch_size,
channels,
input_height,
input_width,
output_height,
output_width,
height_scale,
width_scale,
align_corners,
)?
} else {
interpolate_bicubic_standard(
input,
batch_size,
channels,
input_height,
input_width,
output_height,
output_width,
height_scale,
width_scale,
align_corners,
)?
};
let output = output.to_dtype(input.dtype())?.to_device(input.device())?;
Ok(output)
}
pub fn index_select_2d(t: &Tensor, index: &Tensor) -> Result<Tensor> {
if t.rank() != 2 && index.rank() != 2 {
return Err(anyhow::anyhow!("t and index rank must be equal to 2"));
tests/interpolate.rs
+27
View File
@@ -0,0 +1,27 @@
use aha::utils::interpolate::interpolate_nearest_1d;
use anyhow::Result;
use candle_core::Tensor;
#[test]
fn interpolate_test() -> Result<()> {
// RUST_BACKTRACE=1 cargo test -F cuda --test interpolate interpolate_test -r -- --nocapture
let device = &candle_core::Device::Cpu;
let input = Tensor::arange(0.0f32, 10.0f32, device)?;
let input_1d = input.reshape((1, 1, 10))?;
println!("input_1d: {}", input_1d);
let x_nearest_1d = interpolate_nearest_1d(&input_1d, 20)?;
println!("x_nearest_1d: {}", x_nearest_1d);
// let x_linear_1d = interpolate_linear_1d(&input_1d, 10, Some(true))?;
// println!("x_linear_1d: {}", x_linear_1d);
// let input_2d = input.reshape((1, 1, 10, 10))?;
// println!("input_2d: {}", input_2d);
// let x_nearest_2d = interpolate_nearest_2d(&input_2d, (10, 10))?;
// println!("x_nearest_2d: {}", x_nearest_2d);
// let x_bilinear = interpolate_bilinear(&input_2d, (5, 5), Some(true), Some(false))?;
// println!("x_bilinear: {}", x_bilinear);
// let x_bicubic = interpolate_bicubic(&input_2d, (5, 5), Some(false), Some(true))?;
// // let x_bicubic = interpolate_bicubic_standard(&input_2d, (5, 5), None)?;
// println!("x_bicubic: {}", x_bicubic);
Ok(())
}
tests/messy_test.rs
+29 -21
View File
@@ -5,11 +5,11 @@
// use std::io::{Read, Seek};
// use std::{io::Cursor, time::Instant};
use aha::utils::load_tensor_from_pt;
use aha::utils::interpolate::interpolate_nearest_2d;
// use aha_openai_dive::v1::resources::chat::ChatCompletionParameters;
use anyhow::Result;
// use byteorder::{LittleEndian, ReadBytesExt};
use candle_core::Shape;
use candle_core::Tensor;
// use sentencepiece::SentencePieceProcessor;
// use zip::ZipArchive;
@@ -17,25 +17,33 @@ use candle_core::Shape;
fn messy_test() -> Result<()> {
// RUST_BACKTRACE=1 cargo test -F cuda messy_test -r -- --nocapture
let device = &candle_core::Device::Cpu;
let save_dir: String =
aha::utils::get_default_save_dir().ok_or(anyhow::anyhow!("Failed to get save dir"))?;
let model_path = format!("{}/IndexTeam/IndexTTS-2/", save_dir);
let emo_matrix_path = model_path.clone() + "/feat2.pt";
let t_emo = load_tensor_from_pt(
&emo_matrix_path,
"feat2/data/0",
Shape::from_dims(&[73, 1280]),
device,
)?;
println!("t_emo: {}", t_emo);
let skp_matrix_path = model_path + "/feat1.pt";
let t_skp = load_tensor_from_pt(
&skp_matrix_path,
"feat1/data/0",
Shape::from_dims(&[73, 192]),
device,
)?;
println!("t_skp: {}", t_skp);
let input = Tensor::arange(0.0f32, 25.0f32, device)?.reshape((1, 1, 5, 5))?;
println!("input: {}", input);
let x_nearest = interpolate_nearest_2d(&input, (10, 10))?;
println!("x_nearest: {}", x_nearest);
// let input = Tensor::arange(0.0f32, 25.0f32, device)?.reshape((1, 5, 5))?;
// println!("input: {}", input);
// let x_nearest = interpolate_nearest_1d(&input, 10)?;
// println!("x_nearest: {}", x_nearest);
// let save_dir: String =
// aha::utils::get_default_save_dir().ok_or(anyhow::anyhow!("Failed to get save dir"))?;
// let model_path = format!("{}/IndexTeam/IndexTTS-2/", save_dir);
// let emo_matrix_path = model_path.clone() + "/feat2.pt";
// let t_emo = load_tensor_from_pt(
// &emo_matrix_path,
// "feat2/data/0",
// Shape::from_dims(&[73, 1280]),
// device,
// )?;
// println!("t_emo: {}", t_emo);
// let skp_matrix_path = model_path + "/feat1.pt";
// let t_skp = load_tensor_from_pt(
// &skp_matrix_path,
// "feat1/data/0",
// Shape::from_dims(&[73, 192]),
// device,
// )?;
// println!("t_skp: {}", t_skp);
// let file = File::open(emo_matrix_path)?;
// let mut archive = ZipArchive::new(file)?;
// // 列出所有文件（调试用）
tests/test_deepseek_ocr.rs
+1 -1
View File
@@ -7,7 +7,7 @@ use rocket::futures::StreamExt;
#[test]
fn deepseek_ocr_generate() -> Result<()> {
// RUST_BACKTRACE=1 cargo test -F cuda deepseek_ocr_generate -r -- --nocapture
// RUST_BACKTRACE=1 cargo test -F cuda --test test_deepseek_ocr deepseek_ocr_generate -r -- --nocapture
let message = r#"
{
"model": "deepseek-ocr",
tests/test_hunyuan_ocr.rs
+1 -1
View File
@@ -7,7 +7,7 @@ use rocket::futures::StreamExt;
#[test]
fn hunyuan_ocr_generate() -> Result<()> {
// RUST_BACKTRACE=1 cargo test -F cuda hunyuan_ocr_generate -r -- --nocapture
// RUST_BACKTRACE=1 cargo test -F cuda --test test_hunyuan_ocr hunyuan_ocr_generate -r -- --nocapture
let message = r#"
{
"model": "hunyuan-ocr",
tests/test_paddleocr_vl.rs
+2 -2
View File
@@ -7,7 +7,7 @@ use rocket::futures::StreamExt;
#[test]
fn paddleocr_vl_generate() -> Result<()> {
// RUST_BACKTRACE=1 cargo test -F cuda paddleocr_vl_generate -r -- --nocapture
// RUST_BACKTRACE=1 cargo test -F cuda --test test_paddleocr_vl paddleocr_vl_generate -r -- --nocapture
let message = r#"
{
"model": "paddleocr_vl",
@@ -19,7 +19,7 @@ fn paddleocr_vl_generate() -> Result<()> {
"type": "image",
"image_url":
{
"url": "https://www.qqxiuzi.cn/zh/shouxie-shufa/welcome.png"
"url": "file://./assets/img/ocr_test1.png"
}
},
{
tests/test_rmbg2_0.rs
+1 -1
View File
@@ -6,7 +6,7 @@ use anyhow::Result;
#[test]
fn rmbg2_0_generate() -> Result<()> {
// test with cuda: RUST_BACKTRACE=1 cargo test -F cuda rmbg2_0_generate -r -- --nocapture
// test with cuda: RUST_BACKTRACE=1 cargo test -F cuda --test test_rmbg2_0 rmbg2_0_generate -r -- --nocapture
let save_dir =
aha::utils::get_default_save_dir().ok_or(anyhow::anyhow!("Failed to get save dir"))?;