diff --git a/README.md b/README.md index 8c65f49..17e6c57 100644 --- a/README.md +++ b/README.md @@ -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 diff --git a/src/models/deepseek_ocr/model.rs b/src/models/deepseek_ocr/model.rs index 6c9639f..90e6c4e 100644 --- a/src/models/deepseek_ocr/model.rs +++ b/src/models/deepseek_ocr/model.rs @@ -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))? diff --git a/src/models/hunyuan_ocr/model.rs b/src/models/hunyuan_ocr/model.rs index c9d4b60..f00a3c7 100644 --- a/src/models/hunyuan_ocr/model.rs +++ b/src/models/hunyuan_ocr/model.rs @@ -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::()? as usize; let grid_w = grid_i.i(2)?.to_scalar::()? 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)? diff --git a/src/models/mask_gct/model.rs b/src/models/mask_gct/model.rs index 4745968..51c66c5 100644 --- a/src/models/mask_gct/model.rs +++ b/src/models/mask_gct/model.rs @@ -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 { diff --git a/src/models/paddleocr_vl/model.rs b/src/models/paddleocr_vl/model.rs index 50847e8..f8b319e 100644 --- a/src/models/paddleocr_vl/model.rs +++ b/src/models/paddleocr_vl/model.rs @@ -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))? diff --git a/src/models/rmbg2_0/model.rs b/src/models/rmbg2_0/model.rs index b3af2d2..a3447bf 100644 --- a/src/models/rmbg2_0/model.rs +++ b/src/models/rmbg2_0/model.rs @@ -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)?; diff --git a/src/utils/interpolate.rs b/src/utils/interpolate.rs new file mode 100644 index 0000000..6897b00 --- /dev/null +++ b/src/utils/interpolate.rs @@ -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, +) -> Result> { + 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 { + // 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::()?; + 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 { + // 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::()?; + 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, +) -> Result { + // 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::()?; + 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, + antialias: Option, +) -> Result { + 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, +) -> Result { + 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::()?; + 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 { + 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::()?; + 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 { + 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::()?; + 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, +) -> Result { + 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::()?; + 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, + antialias: Option, +) -> Result { + 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) +} diff --git a/src/utils/mod.rs b/src/utils/mod.rs index 0fdd24c..031950a 100644 --- a/src/utils/mod.rs +++ b/src/utils/mod.rs @@ -1,5 +1,6 @@ pub mod audio_utils; pub mod img_utils; +pub mod interpolate; pub mod tensor_utils; pub mod video_utils; diff --git a/src/utils/tensor_utils.rs b/src/utils/tensor_utils.rs index 8cbc96b..a0bb812 100644 --- a/src/utils/tensor_utils.rs +++ b/src/utils/tensor_utils.rs @@ -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, -) -> Result> { - 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, -) -> Result { - // 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 { - // 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::()?; - 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, -) -> Result { - // 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::()?; - 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没有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::()?; - 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没有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::()?; - 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, - align_corner: Option, -) -> Result { - 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 { if t.rank() != 2 && index.rank() != 2 { return Err(anyhow::anyhow!("t and index rank must be equal to 2")); diff --git a/tests/interpolate.rs b/tests/interpolate.rs new file mode 100644 index 0000000..a8ea879 --- /dev/null +++ b/tests/interpolate.rs @@ -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(()) +} diff --git a/tests/messy_test.rs b/tests/messy_test.rs index 3b3cf3b..b6faf89 100644 --- a/tests/messy_test.rs +++ b/tests/messy_test.rs @@ -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)?; // // 列出所有文件（调试用） diff --git a/tests/test_deepseek_ocr.rs b/tests/test_deepseek_ocr.rs index a7b0a24..d716ce0 100644 --- a/tests/test_deepseek_ocr.rs +++ b/tests/test_deepseek_ocr.rs @@ -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", diff --git a/tests/test_hunyuan_ocr.rs b/tests/test_hunyuan_ocr.rs index fa42c4d..eee1f59 100644 --- a/tests/test_hunyuan_ocr.rs +++ b/tests/test_hunyuan_ocr.rs @@ -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", diff --git a/tests/test_paddleocr_vl.rs b/tests/test_paddleocr_vl.rs index dbab38e..0262810 100644 --- a/tests/test_paddleocr_vl.rs +++ b/tests/test_paddleocr_vl.rs @@ -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" } }, { diff --git a/tests/test_rmbg2_0.rs b/tests/test_rmbg2_0.rs index 346d8ab..892b1ab 100644 --- a/tests/test_rmbg2_0.rs +++ b/tests/test_rmbg2_0.rs @@ -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"))?;