mod host;
use std::env;
use std::path::Path;
use std::sync::Arc;
use anyhow::{Context, Result, anyhow, ensure};
use image::{DynamicImage, GrayImage, RgbImage};
use imageproc::contours::{BorderType, find_contours_with_threshold};
use koharu_runtime::package::{Package, libtorch::Libtorch, loading::preload};
use koharu_torch::{Cuda, Device, Kind, Tensor};
pub(crate) use vm::{CallOutcome, Value, VmResult};
use vm::{Program, compile_source};
use crate::host::TorchHostRuntime;
pub use crate::host::{ScriptRunner, ScriptTextOutput, TorchScriptRunner};
pub struct LamaRustScript {
device: Device,
program: Arc<Program>,
runtime: TorchHostRuntime,
}
impl LamaRustScript {
pub async fn new(device: Device) -> Result<Self> {
preload_libtorch()
.await
.context("failed to initialize LibTorch runtime")?;
let compiled = compile_source(include_str!("../scripts/lama.rss"))
.map_err(|err| anyhow!("failed to compile LaMa RustScript: {err}"))?;
Ok(Self {
device,
program: Arc::new(compiled.program),
runtime: TorchHostRuntime::new(device),
})
}
pub fn inference(
&self,
weights_path: impl AsRef<Path>,
image: &DynamicImage,
mask: &GrayImage,
) -> Result<RgbImage> {
let weights_path = weights_path
.as_ref()
.to_str()
.context("weights path is not valid UTF-8")?
.to_owned();
let image = image.to_rgb8();
ensure!(
image.dimensions() == mask.dimensions(),
"image and mask dimensions differ: image={:?}, mask={:?}",
image.dimensions(),
mask.dimensions()
);
ensure!(
image.width() > 0 && image.height() > 0,
"image dimensions must be non-zero"
);
koharu_torch::no_grad(|| {
if image.width().max(image.height()) > 800 {
let boxes = boxes_from_mask(mask);
if boxes.is_empty() {
return Ok(image);
}
let image_tensor = Tensor::from_slice(image.as_raw())
.view([i64::from(image.height()), i64::from(image.width()), 3])
.to_device(self.device);
let mask_tensor = Tensor::from_slice(mask.as_raw())
.view([i64::from(mask.height()), i64::from(mask.width())])
.to_device(self.device);
let image_width = image.width();
let image_height = image.height();
let mut result = image;
for bounding_box in boxes {
let [left, top, right, bottom] =
crop_box(image_width, image_height, bounding_box);
let crop_result = self.forward(
&weights_path,
image_tensor
.narrow(0, i64::from(top), i64::from(bottom - top))
.narrow(1, i64::from(left), i64::from(right - left)),
mask_tensor
.narrow(0, i64::from(top), i64::from(bottom - top))
.narrow(1, i64::from(left), i64::from(right - left)),
)?;
let source_stride = crop_result.width() as usize * 3;
let target_stride = image_width as usize * 3;
for y in 0..crop_result.height() as usize {
let source_start = y * source_stride;
let target_start = (top as usize + y) * target_stride + left as usize * 3;
result.as_mut()[target_start..target_start + source_stride]
.copy_from_slice(
&crop_result.as_raw()[source_start..source_start + source_stride],
);
}
}
Ok(result)
} else {
let image_tensor = Tensor::from_slice(image.as_raw())
.view([i64::from(image.height()), i64::from(image.width()), 3])
.to_device(self.device);
let mask_tensor = Tensor::from_slice(mask.as_raw())
.view([i64::from(mask.height()), i64::from(mask.width())])
.to_device(self.device);
self.forward(&weights_path, image_tensor, mask_tensor)
}
})
}
fn forward(&self, weights_path: &str, image: Tensor, mask: Tensor) -> Result<RgbImage> {
let height = image.size()[0] as u32;
let width = image.size()[1] as u32;
let image = (image
.permute([2, 0, 1])
.unsqueeze(0)
.to_kind(Kind::Float)
.contiguous())
/ 255.0;
let mask = mask
.gt(0.0)
.unsqueeze(0)
.unsqueeze(0)
.to_kind(Kind::Float)
.contiguous();
let image = symmetric_pad(image, width, height);
let mask = symmetric_pad(mask, width, height);
let output = self.runtime.run(
Arc::clone(&self.program),
image,
mask,
vec![weights_path.to_owned()],
)?;
tensor_to_rgb_image(&output, width, height)
}
}
pub async fn preload_libtorch() -> Result<()> {
let libtorch = Libtorch::for_current_target()?;
let dylibs = libtorch.dylibs()?.collect::<Vec<_>>();
let lib_dir = libtorch.resolve().await?.join("libtorch").join("lib");
prepend_dll_search_path(&lib_dir)?;
for dylib in dylibs {
preload(lib_dir.join(dylib))?;
}
Ok(())
}
#[cfg(windows)]
fn prepend_dll_search_path(path: &Path) -> Result<()> {
let current = env::var_os("PATH").unwrap_or_default();
if env::split_paths(¤t).any(|item| item == path) {
return Ok(());
}
let mut paths = Vec::with_capacity(1);
paths.push(path.to_owned());
paths.extend(env::split_paths(¤t));
let joined = env::join_paths(paths).context("failed to update DLL search path")?;
unsafe {
env::set_var("PATH", joined);
}
Ok(())
}
#[cfg(not(windows))]
fn prepend_dll_search_path(_path: &Path) -> Result<()> {
Ok(())
}
pub fn parse_device(value: &str) -> Result<Device> {
match value.to_ascii_lowercase().as_str() {
"cpu" => Ok(Device::Cpu),
"cuda" => Ok(Device::Cuda(0)),
"mps" => Ok(Device::Mps),
"vulkan" => Ok(Device::Vulkan),
value => value
.strip_prefix("cuda:")
.context("device must be cpu, cuda, cuda:N, mps, or vulkan")?
.parse::<usize>()
.map(Device::Cuda)
.context("invalid CUDA device index"),
}
}
pub fn auto_device() -> Device {
if Cuda::is_available() && Cuda::device_count() > 0 {
Device::Cuda(0)
} else {
Device::Cpu
}
}
pub fn resolve_device(value: Option<&str>) -> Result<Device> {
match value {
Some(value) => parse_device(value),
None => Ok(auto_device()),
}
}
fn boxes_from_mask(mask: &GrayImage) -> Vec<[u32; 4]> {
let width = mask.width();
let mut left = width;
let mut top = mask.height();
let mut right = 0;
let mut bottom = 0;
for y in 0..mask.height() {
let row = &mask.as_raw()[y as usize * width as usize..(y + 1) as usize * width as usize];
let Some(row_left) = row.iter().position(|value| *value > 127) else {
continue;
};
let row_right = row
.iter()
.rposition(|value| *value > 127)
.expect("masked row must have a right edge");
left = left.min(row_left as u32);
top = top.min(y);
right = right.max(row_right as u32 + 1);
bottom = y + 1;
}
if right <= left || bottom <= top {
return Vec::new();
}
let cropped_width = right - left;
let cropped_height = bottom - top;
let padded_width = cropped_width + 2;
let mut padded = GrayImage::new(padded_width, cropped_height + 2);
for y in 0..cropped_height as usize {
let source_start = (top as usize + y) * width as usize + left as usize;
let target_start = (y + 1) * padded_width as usize + 1;
padded.as_mut()[target_start..target_start + cropped_width as usize]
.copy_from_slice(&mask.as_raw()[source_start..source_start + cropped_width as usize]);
}
find_contours_with_threshold::<u32>(&padded, 127)
.into_iter()
.filter(|contour| contour.border_type == BorderType::Outer && contour.parent.is_none())
.filter_map(|contour| {
let mut points = contour.points.into_iter();
let first = points.next()?;
let mut contour_left = first.x;
let mut contour_top = first.y;
let mut contour_right = first.x;
let mut contour_bottom = first.y;
for point in points {
contour_left = contour_left.min(point.x);
contour_top = contour_top.min(point.y);
contour_right = contour_right.max(point.x);
contour_bottom = contour_bottom.max(point.y);
}
Some([
left + contour_left.saturating_sub(1),
top + contour_top.saturating_sub(1),
(left + contour_right).min(mask.width()),
(top + contour_bottom).min(mask.height()),
])
})
.filter(|[left, top, right, bottom]| right > left && bottom > top)
.collect()
}
fn crop_box(image_width: u32, image_height: u32, [left, top, right, bottom]: [u32; 4]) -> [u32; 4] {
let image_width = i64::from(image_width);
let image_height = i64::from(image_height);
let crop_width = i64::from(right - left) + 256;
let crop_height = i64::from(bottom - top) + 256;
let center_x = (i64::from(left) + i64::from(right)) / 2;
let center_y = (i64::from(top) + i64::from(bottom)) / 2;
let raw_left = center_x - crop_width / 2;
let raw_right = center_x + crop_width / 2;
let raw_top = center_y - crop_height / 2;
let raw_bottom = center_y + crop_height / 2;
let mut left = raw_left.max(0);
let mut right = raw_right.min(image_width);
let mut top = raw_top.max(0);
let mut bottom = raw_bottom.min(image_height);
if raw_left < 0 {
right += -raw_left;
}
if raw_right > image_width {
left -= raw_right - image_width;
}
if raw_top < 0 {
bottom += -raw_top;
}
if raw_bottom > image_height {
top -= raw_bottom - image_height;
}
[
left.clamp(0, image_width) as u32,
top.clamp(0, image_height) as u32,
right.clamp(0, image_width) as u32,
bottom.clamp(0, image_height) as u32,
]
}
fn tensor_to_rgb_image(tensor: &Tensor, width: u32, height: u32) -> Result<RgbImage> {
let tensor = tensor
.narrow(2, 0, i64::from(height))
.narrow(3, 0, i64::from(width))
.squeeze_dim(0)
.permute([1, 2, 0])
.clamp(0.0, 1.0)
* 255.0;
let tensor = tensor
.to_kind(Kind::Uint8)
.contiguous()
.to_device(Device::Cpu)
.view([-1]);
let rgb = Vec::<u8>::try_from(&tensor)?;
RgbImage::from_raw(width, height, rgb).context("failed to convert LaMa tensor to RGB image")
}
fn symmetric_pad(tensor: Tensor, width: u32, height: u32) -> Tensor {
let mut tensor = tensor;
let padded_height = ceil_modulo(height, 8);
if padded_height != height {
let indices = (0..padded_height)
.map(|index| i64::from(symmetric_index(index, height)))
.collect::<Vec<_>>();
tensor = tensor.index_select(2, &Tensor::from_slice(&indices).to_device(tensor.device()));
}
let padded_width = ceil_modulo(width, 8);
if padded_width != width {
let indices = (0..padded_width)
.map(|index| i64::from(symmetric_index(index, width)))
.collect::<Vec<_>>();
tensor = tensor.index_select(3, &Tensor::from_slice(&indices).to_device(tensor.device()));
}
tensor
}
fn ceil_modulo(value: u32, modulo: u32) -> u32 {
if value.is_multiple_of(modulo) {
value
} else {
(value / modulo + 1) * modulo
}
}
fn symmetric_index(index: u32, len: u32) -> u32 {
let index = index % (len * 2);
if index < len {
index
} else {
len * 2 - index - 1
}
}