support gpu
This commit is contained in:
@@ -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);
|
||||
|
||||
Reference in New Issue
Block a user