support gpu

This commit is contained in:
Julian Freeman
2026-01-19 15:35:53 -04:00
parent 330a62027c
commit 2c89ef42b3
6 changed files with 61 additions and 27 deletions

View File

@@ -1,5 +1,5 @@
use crate::ort_ops;
use image::{Rgba, RgbaImage};
use ort::session::{Session, builder::GraphOptimizationLevel};
use ort::value::Value;
pub fn run_lama_inpainting(
@@ -8,14 +8,8 @@ pub fn run_lama_inpainting(
mask_image: &image::GrayImage,
) -> Result<RgbaImage, String> {
// 1. Initialize Session
let mut session = Session::builder()
.map_err(|e| format!("Failed to create session builder: {}", e))?
.with_optimization_level(GraphOptimizationLevel::Level3)
.map_err(|e| format!("Failed to set opt level: {}", e))?
.with_intra_threads(4)
.map_err(|e| format!("Failed to set threads: {}", e))?
.commit_from_file(model_path)
.map_err(|e| format!("Failed to load model from {:?}: {}", model_path, e))?;
let mut session = ort_ops::create_session(model_path)
.map_err(|e| format!("Failed to create ORT session for LAMA: {}", e))?;
// 2. Preprocess
let target_size = (512, 512);