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

add hunyuan_ocr

Browse Source
This commit is contained in:
jhqxxx
2025-12-03 17:21:01 +08:00
parent 7d72cb3baf
commit 697484cf23
29 changed files with 1756 additions and 190 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/utils/img_utils.rs
+38
View File
@@ -9,6 +9,8 @@ use base64::{Engine, engine::general_purpose};
use candle_core::{DType, Device, Tensor};
use image::{DynamicImage, ImageBuffer, ImageReader, Rgb, RgbImage, imageops};
use crate::utils::{ceil_by_factor, floor_by_factor, round_by_factor};
pub fn load_image_from_url(url: &str) -> Result<DynamicImage> {
tokio::task::block_in_place(|| {
let response = reqwest::blocking::get(url)
@@ -237,3 +239,39 @@ pub fn img_transform(
.to_dtype(dtype)?;
Ok(img_tensor)
}
pub fn img_smart_resize(
img_h: u32,
img_w: u32,
factor: u32,
min_pixels: u32,
max_pixels: u32,
) -> Result<(u32, u32)> {
if std::cmp::max(img_h, img_w) / std::cmp::min(img_h, img_w) > 200 {
return Err(anyhow!(format!(
"absolute aspect ratio mush be smaller than {}, got {}",
200,
std::cmp::max(img_h, img_w) / std::cmp::min(img_h, img_w)
)));
}
let image_factor = factor;
let mut h_bar = std::cmp::max(image_factor, round_by_factor(img_h, image_factor));
let mut w_bar = std::cmp::max(image_factor, round_by_factor(img_w, image_factor));
if h_bar * w_bar > max_pixels {
let beta = ((img_h * img_w) as f32 / max_pixels as f32).sqrt();
h_bar = std::cmp::max(
image_factor,
floor_by_factor(img_h as f32 / beta, image_factor),
);
w_bar = std::cmp::max(
image_factor,
floor_by_factor(img_w as f32 / beta, image_factor),
);
} else if h_bar * w_bar < min_pixels {
let beta = (min_pixels as f32 / (img_h * img_w) as f32).sqrt();
h_bar = ceil_by_factor(img_h as f32 * beta, image_factor);
w_bar = ceil_by_factor(img_w as f32 * beta, image_factor);
}
Ok((h_bar, w_bar))
}
src/utils/mod.rs
+40 -1
View File
@@ -3,6 +3,8 @@ pub mod img_utils;
pub mod tensor_utils;
pub mod video_utils;
use std::process::Command;
use aha_openai_dive::v1::resources::{
chat::{
ChatCompletionChoice, ChatCompletionChunkChoice, ChatCompletionChunkResponse,
@@ -31,6 +33,33 @@ pub fn get_device(device: Option<&Device>) -> Device {
}
}
pub fn get_gpu_sm_arch() -> Result<f32> {
let output = Command::new("nvidia-smi")
.arg("--query-gpu=compute_cap")
.arg("--format=csv,noheader")
.output()
.map_err(|e| anyhow::anyhow!(format!("Failed to execute nvidia-smi: {}", e)))?;
if !output.status.success() {
return Err(anyhow::anyhow!(format!(
"nvidia-smi failed with status: {}\nError: {}",
output.status,
String::from_utf8_lossy(&output.stderr)
)));
}
let output_str = String::from_utf8_lossy(&output.stdout);
let output_str = output_str.trim();
let sm_float = match output_str.parse::<f32>() {
Ok(num) => num,
Err(_) => {
return Err(anyhow::anyhow!(format!(
"gpr sm arch: {} parse float32 error",
output_str
)));
}
};
Ok(sm_float)
}
pub fn get_dtype(dtype: Option<DType>, cfg_dtype: &str) -> DType {
match dtype {
Some(d) => d,
@@ -41,7 +70,16 @@ pub fn get_dtype(dtype: Option<DType>, cfg_dtype: &str) -> DType {
"float32" | "float" => DType::F32,
"float64" | "double" => DType::F64,
"float16" => DType::F16,
"bfloat16" => DType::BF16,
"bfloat16" => {
let arch = get_gpu_sm_arch();
match arch {
Err(_) => DType::F16,
Ok(a) => {
// nvidia显卡sm架构>=8.0的才支持BF16
if a >= 8.0 { DType::BF16 } else { DType::F16 }
}
}
}
"uint8" => DType::U8,
"int8" | "int16" | "int32" | "int64" => DType::I64,
_ => DType::F32,
@@ -257,6 +295,7 @@ pub fn get_logit_processor(
top_k: Option<usize>,
seed: u64,
) -> LogitsProcessor {
let temperature = temperature.and_then(|v| if v < 1e-7 { None } else { Some(v) });
match top_k {
None => LogitsProcessor::new(
seed,
src/utils/tensor_utils.rs
+97 -17
View File
@@ -221,6 +221,14 @@ pub fn masked_scatter_dim0(original: &Tensor, replace: &Tensor, mask: &Tensor) -
Ok(original)
}
pub fn get_not_equal_mask(input_ids: &Tensor, token_ids: u32) -> Result<Tensor> {
let image_token_id_tensor = Tensor::new(vec![token_ids], input_ids.device())?;
let mask = input_ids
.broadcast_ne(&image_token_id_tensor)?
.to_dtype(candle_core::DType::U32)?;
Ok(mask)
}
pub fn get_equal_mask(input_ids: &Tensor, token_ids: u32) -> Result<Tensor> {
let image_token_id_tensor = Tensor::new(vec![token_ids], input_ids.device())?;
let mask = input_ids
@@ -229,10 +237,16 @@ pub fn get_equal_mask(input_ids: &Tensor, token_ids: u32) -> Result<Tensor> {
Ok(mask)
}
pub fn get_vision_next_indices(input_ids: &Tensor, token_id: u32) -> Result<Tensor> {
pub fn get_eq_indices(input_ids: &Tensor, token_id: u32) -> Result<Tensor> {
// input_ids -> shape: (seq_len)
let mask = get_equal_mask(input_ids, token_id)?;
let indices = nonzero_index(&mask)?;
Ok(indices)
}
pub fn get_vision_next_indices(input_ids: &Tensor, token_id: u32) -> Result<Tensor> {
// input_ids -> shape: (seq_len)
let indices = get_eq_indices(input_ids, token_id)?;
let indices = indices.broadcast_add(&Tensor::new(vec![1u32], input_ids.device())?)?;
Ok(indices)
}
@@ -384,9 +398,9 @@ pub fn interpolate_linear_1d(
align_corner: Option<bool>,
) -> Result<Tensor> {
// t: [b, channels, features]
if t.rank() < 3 {
if t.rank() != 3 {
return Err(anyhow::anyhow!(
"Input rank must have at least 3 dimensions"
"Input rank must have equal to 3 dimensions"
));
}
let shape = t.dims();
@@ -394,19 +408,13 @@ pub fn interpolate_linear_1d(
if orig_size == target_size {
return Ok(t.clone());
}
let mut reshaped = t.clone();
if shape.len() > 3 {
let bs = shape[0];
let channels = shape[1..shape.len() - 1].iter().product::<usize>();
reshaped = reshaped.reshape((bs, channels, orig_size))?;
}
let (bs, channels, _) = reshaped.dims3()?;
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 = reshaped.i((b, c))?;
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) {
@@ -424,16 +432,88 @@ pub fn interpolate_linear_1d(
output = output.slice_assign(&[(b..b + 1), (c..c + 1), (0..target_size)], &out_i)?;
}
}
if shape.len() != 3 {
let mut new_shape = shape.to_vec();
let last_dim = new_shape.len() - 1;
new_shape[last_dim] = target_size;
output = output.reshape(new_shape)?
}
output = output.contiguous()?;
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