mirror of
https://github.com/storytold/craft-ai-integrations.git
synced 2026-11-02 18:06:49 +00:00
feat(ai-engine-burn): InSPyReNet for Remove background, built in only (#112)
* feat(ai-engine-burn): InSPyReNet for image.matte (Remove background) A second model on the Burn engine: `inspyrenet` answers `image.matte` from the picture alone, a grayscale PNG at the picture's size (white is the subject, soft along hair and fur). `remove-background` has had no local model; this is one. - Network: InSPyReNet with the Swin-B backbone, written by hand in Burn after the MIT reference code. Module names follow the PyTorch keys, so the published .safetensors loads by name through Burn's PyTorch adapter. A load that leaves a weight unset, fails one, or finds a foreign tensor is refused in plain words. - Any picture size: full resolution up to 1280 px on the long side; above that a reflect-padded square reduced with antialiased Lanczos-3, the matte enlarged back and refined at full resolution with a colour guided filter on the tiles near the edge (the Gyre approach), then a 0.2 fringe shrink. - Weights: 1038lab/inspyrenet inspyrenet.safetensors at a pinned revision, checked tensor for tensor against the author's ckpt_base.pth. Checked on one Windows machine (release): every stage against transparent-background's PyTorch code (worst relative difference 3.2e-5 on the processor, 2.4e-4 on the card, no pixel changing side); the shared conformance suite; on a real PNG the mask is byte-identical to the same port in PhotoCraft. Four new deliberate-bug cases in `be`, all caught. Not done: a catalogue recipe (a recipe names a ComfyUI workflow today), clean timings, macOS and Linux runs, a run through an app. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * feat(ai-features): a recipe that is built in only, InSPyReNet for Remove background `LocalRecipe::comfyui` (default true) marks a recipe ComfyUI runs. With false it is only its built-in routes: never a ComfyUI setup option, plan, install offer or connection, and engine advice does not list ComfyUI for a tool without a ComfyUI recipe. `Feature::comfyui_recipes()` gives those recipes; `engine_recipe`, `plan_native`, `setup_offer_native` and `enable_native` are unchanged. `remove-background` gets one such recipe: InSPyReNet on Burn, the pinned `inspyrenet.safetensors` (1038lab/inspyrenet @ 0f594663, 367,120,444 bytes) in the folder `inspyrenet`, so Set up AI offers one checked download and no ComfyUI. Tests: the catalogue, setup options, plans and offers (`ai-features/tests/built_in_only.rs`), registration and advice (`ai-client/tests/built_in_only.rs`), and through `AiClient` on Burn (`ai-client/tests/burn_matte.rs`, the real run ignored). Mutation group `bo`: 8 of 8 caught. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * fix(tools): retarget two `lc` mutation cases at the ComfyUI recipe lookups `unrunnable-recipe-accepted` and `second-recipe-takes-plain-id` searched for the `feature.locals` lookups in `ai-client/src/local.rs`, which now go through `Feature::comfyui_recipes()`. Same mutations, new targets; group `lc`: 21 caught, one known equivalent. Every case in `GROUPS` again has exactly one target (640 checked). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * test(ai-engine-burn): InSPyReNet runner checks around a stand-in, readout without panic CI has no graphics card, only a software renderer (lavapipe, WARP), and the engine takes it. There the tiny network was too slow for the conformance suite's poll limit, and on lavapipe the readout failed. Unoptimised, Burn's processor backend takes minutes for the same network. - The loaded model runs the pipeline around a `Predict`: the network, or in tests a stand-in, as the SAM 2.1 tests here use a fake backend. The conformance suite and the mask test use the stand-in (debug: the crate's unit tests in about 18 s). - The tiny network still runs through load and run on the processor in an ignored release test (about 4 s). The runner's device choice is a field, so tests can name the processor. - The prediction is read with `try_into_data`, so a device that fails while running is an error, not a guarded panic. - docs/BURN-ENGINE.md: the test split and that a software renderer counts as a card. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * feat(ai-engine-burn): InSPyReNet moves to the processor when the card fails A card can answer the engine's probe and still fail on the network: on CI a software renderer (lavapipe) failed while reading the answer back. The first failure on a card that is not a cancel now builds the network again on the processor from the same file, runs that picture again there with a notice, and stays on the processor (`inspyrenet/fallback.rs`). A network loaded on the processor in the first place has nothing to move to. Tests with stand-ins: a working card is kept; a failing one moves once and stays; a cancel is not a failure; a rebuild that fails is the answer. Mutation group `be`: three new cases, all seven InSPyReNet cases caught. A real photo on the card still runs on the card. docs/BURN-ENGINE.md: the move, and the client run's picture is 2848 x 4272, not 3101 x 1613. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * docs(ai-engine-burn): the move to the processor stays only after a rebuild that worked The rebuild is tried once. If it fails, its error is the answer and later pictures go to the card again, as `a_rebuild_that_fails_is_the_answer` already pins. The module docs and docs/BURN-ENGINE.md said the model stays on the processor without that condition. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * feat(ai-engine-burn): retry a failed InSPyReNet load on the card on the processor The move to the processor covered failures while predicting, not a load on the card that fails. `load_or_processor` now retries such a load on the processor, unless the file itself is not the checkpoint (`Unsupported`), which is refused at once rather than read a second time. A network loaded on the processor in the first place is loaded once. Tests: a card that loads keeps the network; one that fails loads on the processor; a bad file is not read twice; on the processor there is one load. Mutation group `be`: two new cases, all nine InSPyReNet cases caught; every case in the checker has one target (645). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * style(ai-engine-burn): method references in the guided filter, one name for the fallback path Three SonarCloud findings on #112: two closures in `guided.rs` become method references, and the mutation cases name `inspyrenet/fallback.rs` once. No change in behaviour. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> --------- Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com> Co-authored-by: imperator <imperator.maximus@gmail.com>
This commit is contained in:
imperator
@@ -11,7 +11,8 @@ impl AiClient {
|
||||
/// The local engines `feature` can run on, best first, judged for `profile` (see
|
||||
/// `craft_ai_comfy_env::system::detect`). An engine this app has no adapter for is listed but
|
||||
/// never suggested. Every built-in engine some usable recipe lists is offered (each runs its own
|
||||
/// recipe, [`Feature::engine_recipe`]), then ComfyUI; the memory limit is the first usable
|
||||
/// recipe, [`Feature::engine_recipe`]), then ComfyUI when a usable recipe runs on it; the
|
||||
/// memory limit is the first usable
|
||||
/// recipe's, the one [`ToolDesk::enable_local`](crate::ToolDesk::enable_local) registers. A tool
|
||||
/// with no usable recipe has no advice.
|
||||
pub fn engine_advice(&self, feature: &Feature, profile: &SystemProfile) -> Vec<EngineAdvice> {
|
||||
@@ -19,7 +20,9 @@ impl AiClient {
|
||||
return Vec::new();
|
||||
};
|
||||
let mut providers: Vec<&str> = feature.engines();
|
||||
providers.push("comfyui");
|
||||
if feature.comfyui_recipes().any(|r| r.usable) {
|
||||
providers.push("comfyui");
|
||||
}
|
||||
let offers: Vec<EngineOffer> = providers
|
||||
.into_iter()
|
||||
.map(|p| EngineOffer {
|
||||
|
||||
@@ -69,15 +69,18 @@ pub(crate) fn is_recipe_connection(connection: &Connection, feature: &str) -> bo
|
||||
}
|
||||
|
||||
impl ToolDesk {
|
||||
/// Registers the first local model of `feature` that can run, served by the ComfyUI at
|
||||
/// `endpoint` (a loopback address), and binds the tool to it if nothing is chosen yet.
|
||||
/// Registers the first local model of `feature` that ComfyUI can run, served by the ComfyUI
|
||||
/// at `endpoint` (a loopback address), and binds the tool to it if nothing is chosen yet.
|
||||
pub fn enable_local(&self, feature: &Feature, endpoint: &str) -> Result<LocalEnabled, AiError> {
|
||||
let recipe = feature.locals.iter().find(|r| r.usable).ok_or_else(|| {
|
||||
AiError::unsupported(format!(
|
||||
"{} has no local model that can run yet",
|
||||
feature.title
|
||||
))
|
||||
})?;
|
||||
let recipe = feature
|
||||
.comfyui_recipes()
|
||||
.find(|r| r.usable)
|
||||
.ok_or_else(|| {
|
||||
AiError::unsupported(format!(
|
||||
"{} has no local model that can run yet",
|
||||
feature.title
|
||||
))
|
||||
})?;
|
||||
self.enable_recipe(feature, recipe, endpoint)
|
||||
}
|
||||
|
||||
@@ -92,8 +95,7 @@ impl ToolDesk {
|
||||
endpoint: &str,
|
||||
) -> Result<LocalEnabled, AiError> {
|
||||
let recipe = feature
|
||||
.locals
|
||||
.iter()
|
||||
.comfyui_recipes()
|
||||
.find(|r| r.workflow == workflow && r.usable)
|
||||
.ok_or_else(|| {
|
||||
AiError::unsupported(format!(
|
||||
@@ -129,8 +131,8 @@ impl ToolDesk {
|
||||
&& held.models.iter().any(|m| m.id.as_str() == recipe.workflow)
|
||||
}
|
||||
None => feature
|
||||
.locals
|
||||
.first()
|
||||
.comfyui_recipes()
|
||||
.next()
|
||||
.is_some_and(|first| first.workflow == recipe.workflow),
|
||||
};
|
||||
let id = if primary {
|
||||
|
||||
@@ -0,0 +1,125 @@
|
||||
//! A tool whose only local model is built in: Remove background on InSPyReNet, read by the Burn
|
||||
//! engine. The launcher registers the built-in route, never a ComfyUI one, and engine advice does
|
||||
//! not offer ComfyUI for it.
|
||||
use std::sync::Arc;
|
||||
|
||||
use craft_ai_client::{AiClient, ProviderFactory, ToolDesk};
|
||||
use craft_ai_core::{AiError, AppId, ErrorKind, StoreFile};
|
||||
use craft_ai_features::{Feature, Gpu, GpuVendor, Os, SystemProfile};
|
||||
use craft_ai_native::secrets::MemoryStore;
|
||||
use craft_ai_native::shared::SharedStore;
|
||||
|
||||
fn remove_background() -> Feature {
|
||||
craft_ai_features::builtin()
|
||||
.into_iter()
|
||||
.find(|f| f.id == "remove-background")
|
||||
.unwrap()
|
||||
}
|
||||
|
||||
fn windows_nvidia() -> SystemProfile {
|
||||
SystemProfile {
|
||||
os: Os::Windows,
|
||||
gpus: vec![Gpu {
|
||||
vendor: GpuVendor::Nvidia,
|
||||
name: "RTX 3090 Ti".to_owned(),
|
||||
vram_mb: Some(24_576),
|
||||
}],
|
||||
comfyui_found: true,
|
||||
}
|
||||
}
|
||||
|
||||
struct World {
|
||||
dir: tempfile::TempDir,
|
||||
shared: SharedStore,
|
||||
}
|
||||
|
||||
impl World {
|
||||
fn new() -> Self {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let shared = SharedStore::new(dir.path().join("connections.json"));
|
||||
shared.replace(&StoreFile::default()).unwrap();
|
||||
Self { dir, shared }
|
||||
}
|
||||
fn client(&self, factory: ProviderFactory) -> AiClient {
|
||||
AiClient::new(
|
||||
AppId::new("photocraft").unwrap(),
|
||||
self.shared.clone(),
|
||||
Arc::new(MemoryStore::new()),
|
||||
)
|
||||
.with_factory(factory)
|
||||
}
|
||||
fn desk(&self) -> ToolDesk {
|
||||
ToolDesk::new(
|
||||
self.client(ProviderFactory::standard()),
|
||||
self.shared.clone(),
|
||||
None,
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
/// A factory with no built-in engine, whatever cargo features are on.
|
||||
fn without_engines() -> ProviderFactory {
|
||||
let mut factory = ProviderFactory::standard();
|
||||
for engine in ["rten", "candle", "burn"] {
|
||||
factory.unregister(engine);
|
||||
}
|
||||
factory
|
||||
}
|
||||
|
||||
/// The plain built-in call registers InSPyReNet on Burn with the file it was given and binds the
|
||||
/// tool; asking for a ComfyUI model is refused and writes nothing.
|
||||
#[test]
|
||||
fn the_built_in_route_is_registered_and_comfyui_is_refused() {
|
||||
let w = World::new();
|
||||
let refused = w
|
||||
.desk()
|
||||
.enable_local(&remove_background(), "http://127.0.0.1:8188")
|
||||
.err()
|
||||
.unwrap();
|
||||
assert_eq!(refused.kind, ErrorKind::Unsupported);
|
||||
let refused = w
|
||||
.desk()
|
||||
.enable_local_recipe(
|
||||
&remove_background(),
|
||||
"inspyrenet-matte",
|
||||
"http://127.0.0.1:8188",
|
||||
)
|
||||
.err()
|
||||
.unwrap();
|
||||
assert_eq!(refused.kind, ErrorKind::Unsupported);
|
||||
assert!(w.shared.load().unwrap().connections.is_empty());
|
||||
|
||||
let file = w.dir.path().join("inspyrenet.safetensors");
|
||||
std::fs::write(&file, b"not really weights").unwrap();
|
||||
let done = w.desk().enable_native(&remove_background(), &file).unwrap();
|
||||
assert!(done.bound);
|
||||
assert_eq!(done.connection.as_str(), "local-remove-background-burn");
|
||||
assert_eq!(done.model.as_str(), "inspyrenet");
|
||||
let store = w.shared.load().unwrap();
|
||||
assert_eq!(store.connections.len(), 1);
|
||||
assert_eq!(store.connections[0].provider, "burn");
|
||||
}
|
||||
|
||||
/// Advice lists Burn alone: never ComfyUI, which has no recipe for this tool. An app with the
|
||||
/// engine on a graphics card it vouches for is advised it.
|
||||
#[test]
|
||||
fn engine_advice_offers_burn_and_never_comfyui() {
|
||||
let w = World::new();
|
||||
let advice = w
|
||||
.client(without_engines())
|
||||
.engine_advice(&remove_background(), &windows_nvidia());
|
||||
let providers: Vec<&str> = advice.iter().map(|a| a.provider.as_str()).collect();
|
||||
assert_eq!(providers, ["burn"]);
|
||||
assert!(!advice[0].usable && !advice[0].recommended);
|
||||
|
||||
let mut factory = without_engines();
|
||||
factory.register("burn", |_, _| {
|
||||
Err(AiError::unsupported("advice never builds a provider"))
|
||||
});
|
||||
factory.declare_accelerators("burn", &[GpuVendor::Nvidia]);
|
||||
let advice = w
|
||||
.client(factory)
|
||||
.engine_advice(&remove_background(), &windows_nvidia());
|
||||
assert_eq!(advice.len(), 1);
|
||||
assert!(advice[0].usable && advice[0].recommended, "{advice:?}");
|
||||
}
|
||||
@@ -0,0 +1,136 @@
|
||||
//! Remove background on the Burn engine through the client: the catalogue's built-in route, the
|
||||
//! shared store, the cargo feature and the InSPyReNet runner. An app asks for `image.matte` only.
|
||||
//!
|
||||
//! The first test needs no model file. The second (ignored) needs the real ones:
|
||||
//! `CRAFT_TEST_INSPYRENET_WEIGHTS` (`inspyrenet.safetensors`) and `CRAFT_TEST_INSPYRENET_IMAGE` (a
|
||||
//! PNG). Run:
|
||||
//! `cargo test -p craft-ai-client --features burn --release --test burn_matte -- --ignored`.
|
||||
#![cfg(feature = "burn")]
|
||||
|
||||
use std::path::Path;
|
||||
use std::sync::Arc;
|
||||
use std::time::Duration;
|
||||
|
||||
use craft_ai_client::{AiClient, ToolDesk};
|
||||
use craft_ai_core::{
|
||||
AiError, AppId, ConsentLedger, Event, InputRole, JobState, Payload, Request, StoreFile,
|
||||
};
|
||||
use craft_ai_features::Feature;
|
||||
use craft_ai_native::secrets::MemoryStore;
|
||||
use craft_ai_native::shared::SharedStore;
|
||||
|
||||
fn remove_background() -> Feature {
|
||||
craft_ai_features::builtin()
|
||||
.into_iter()
|
||||
.find(|f| f.id == "remove-background")
|
||||
.unwrap()
|
||||
}
|
||||
|
||||
struct World {
|
||||
_dir: tempfile::TempDir,
|
||||
shared: SharedStore,
|
||||
}
|
||||
|
||||
impl World {
|
||||
fn new() -> Self {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let shared = SharedStore::new(dir.path().join("connections.json"));
|
||||
shared.replace(&StoreFile::default()).unwrap();
|
||||
Self { _dir: dir, shared }
|
||||
}
|
||||
|
||||
fn client(&self, app: &str) -> AiClient {
|
||||
AiClient::new(
|
||||
AppId::new(app).unwrap(),
|
||||
self.shared.clone(),
|
||||
Arc::new(MemoryStore::new()),
|
||||
)
|
||||
}
|
||||
|
||||
/// The launcher's part: set up the tool's built-in route with `weights` as its file.
|
||||
fn set_up(&self, weights: &Path) {
|
||||
ToolDesk::new(self.client("launcher"), self.shared.clone(), None)
|
||||
.enable_native(&remove_background(), weights)
|
||||
.unwrap();
|
||||
}
|
||||
}
|
||||
|
||||
fn request(picture: Vec<u8>) -> Request {
|
||||
Request::new(remove_background().capability.clone()).with_input(
|
||||
InputRole::Source,
|
||||
Payload::inline("image/png", picture).unwrap(),
|
||||
)
|
||||
}
|
||||
|
||||
/// Runs `request` for `app` to its end: the outputs, or the refusal from wherever it came.
|
||||
fn run(app: &AiClient, request: Request) -> Result<Vec<Payload>, AiError> {
|
||||
let role = remove_background().role().unwrap();
|
||||
let plan = app.plan(&role, request)?;
|
||||
let done = app
|
||||
.start(&plan, &ConsentLedger::default())?
|
||||
.wait(Duration::from_millis(50), Duration::from_secs(600))?;
|
||||
match done.state {
|
||||
JobState::Succeeded => Ok(done
|
||||
.events
|
||||
.into_iter()
|
||||
.filter_map(|e| match e {
|
||||
Event::Output(o) => Some(o.payload),
|
||||
_ => None,
|
||||
})
|
||||
.collect()),
|
||||
JobState::Failed { error } => Err(error),
|
||||
other => panic!("unexpected end {other:?}"),
|
||||
}
|
||||
}
|
||||
|
||||
/// A file that is not the checkpoint is refused in plain words that point to Set up AI.
|
||||
#[test]
|
||||
fn a_file_that_is_not_the_checkpoint_is_refused_in_words() {
|
||||
let w = World::new();
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let weights = dir.path().join("inspyrenet.safetensors");
|
||||
std::fs::write(&weights, b"not the model").unwrap();
|
||||
w.set_up(&weights);
|
||||
let error = run(&w.client("photocraft"), request(vec![1, 2, 3])).unwrap_err();
|
||||
assert!(error.message.contains("Set up AI"), "{}", error.message);
|
||||
}
|
||||
|
||||
/// Two apps ask only for Remove background and each gets one grayscale mask at the picture's size
|
||||
/// with both subject and background in it.
|
||||
#[test]
|
||||
#[ignore = "needs CRAFT_TEST_INSPYRENET_WEIGHTS and CRAFT_TEST_INSPYRENET_IMAGE; see the module docs"]
|
||||
fn two_apps_get_a_matte_from_the_real_model() {
|
||||
let (Some(weights), Some(image)) = (
|
||||
std::env::var_os("CRAFT_TEST_INSPYRENET_WEIGHTS"),
|
||||
std::env::var_os("CRAFT_TEST_INSPYRENET_IMAGE"),
|
||||
) else {
|
||||
eprintln!("skipped: set CRAFT_TEST_INSPYRENET_WEIGHTS and CRAFT_TEST_INSPYRENET_IMAGE");
|
||||
return;
|
||||
};
|
||||
let picture = std::fs::read(image).unwrap();
|
||||
let size = png::Decoder::new(std::io::Cursor::new(&picture))
|
||||
.read_info()
|
||||
.unwrap()
|
||||
.info()
|
||||
.size();
|
||||
let w = World::new();
|
||||
w.set_up(Path::new(&weights));
|
||||
for app in ["photocraft", "lightcraft"] {
|
||||
let started = std::time::Instant::now();
|
||||
let outputs = run(&w.client(app), request(picture.clone())).unwrap();
|
||||
eprintln!("{app}: {:.1?}", started.elapsed());
|
||||
assert_eq!(outputs.len(), 1);
|
||||
let Payload::Inline { data, .. } = &outputs[0] else {
|
||||
panic!("expected the mask itself");
|
||||
};
|
||||
let mut reader = png::Decoder::new(std::io::Cursor::new(data))
|
||||
.read_info()
|
||||
.unwrap();
|
||||
let mut pixels = vec![0; reader.output_buffer_size().unwrap()];
|
||||
let frame = reader.next_frame(&mut pixels).unwrap();
|
||||
assert_eq!((frame.width, frame.height), size);
|
||||
assert_eq!(frame.color_type, png::ColorType::Grayscale);
|
||||
let subject = pixels.iter().filter(|v| **v > 127).count() as f64 / pixels.len() as f64;
|
||||
assert!((0.02..0.98).contains(&subject), "subject {subject}");
|
||||
}
|
||||
}
|
||||
@@ -348,14 +348,15 @@ fn the_next_recipe_that_can_run_supplies_the_built_in_runtime() {
|
||||
assert_eq!(w.shared.load().unwrap().connections.len(), 1);
|
||||
}
|
||||
|
||||
/// Every other tool refuses a built-in runtime and writes nothing.
|
||||
/// Every tool other than Select subject and Remove background refuses a built-in runtime and
|
||||
/// writes nothing.
|
||||
#[test]
|
||||
fn a_tool_without_a_built_in_runtime_is_refused() {
|
||||
let w = World::new();
|
||||
let file = w.weights("sam3.safetensors");
|
||||
let others = craft_ai_features::builtin()
|
||||
.into_iter()
|
||||
.filter(|f| f.id != "select-subject");
|
||||
.filter(|f| f.id != "select-subject" && f.id != "remove-background");
|
||||
for f in others {
|
||||
assert!(
|
||||
w.desk("launcher").enable_native(&f, &file).is_err(),
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
[package]
|
||||
name = "craft-ai-engine-burn"
|
||||
description = "SAM 2.1 on the Burn engine for Craft AI: an optional graphics-card runner for image.segment, reading weights prepared once on this computer. Native only."
|
||||
description = "The Burn engine for Craft AI: optional graphics-card runners for image.segment (SAM 2.1, weights prepared once on this computer) and image.matte (InSPyReNet). Native only."
|
||||
version.workspace = true
|
||||
edition.workspace = true
|
||||
license.workspace = true
|
||||
|
||||
@@ -92,11 +92,16 @@ pub const DEVICE_ENV: &str = "CRAFT_AI_BURN_DEVICE";
|
||||
|
||||
/// The graphics card when one answers, else the processor. `CRAFT_AI_BURN_DEVICE=cpu` forces the
|
||||
/// processor.
|
||||
fn pick_device() -> Device {
|
||||
pub(crate) fn pick_device() -> Device {
|
||||
let forced_processor = std::env::var_os(DEVICE_ENV).is_some_and(|v| v == "cpu");
|
||||
if !forced_processor && graphics_card_answers() {
|
||||
return WgpuDevice::default().into();
|
||||
}
|
||||
processor()
|
||||
}
|
||||
|
||||
/// The processor, through Burn's `ndarray` backend.
|
||||
pub(crate) fn processor() -> Device {
|
||||
#[allow(deprecated)]
|
||||
let cpu = NdArrayDevice::Cpu;
|
||||
cpu.into()
|
||||
|
||||
@@ -0,0 +1,238 @@
|
||||
//! Moving to the processor when the graphics card cannot run the network.
|
||||
//!
|
||||
//! A card can answer the engine's start-up probe and still fail on the network: on a build server
|
||||
//! a software renderer (lavapipe) did, while reading the answer back. So the first failure that is
|
||||
//! not a cancel builds the network again on the processor and runs that picture again there. Once
|
||||
//! that rebuild has worked the model stays on the processor. The move is tried once: if the rebuild
|
||||
//! fails, its error is the answer and later pictures go to the card again.
|
||||
|
||||
use craft_ai_core::{AiError, ErrorKind};
|
||||
use craft_ai_engine::RunContext;
|
||||
|
||||
use super::Predict;
|
||||
use super::pipeline::Input;
|
||||
|
||||
/// Loads the network with `card` when `on_card`, else with `processor` alone. A card load that
|
||||
/// fails for any reason but the file is retried with `processor`; a file that is not the checkpoint
|
||||
/// ([`ErrorKind::Unsupported`]) is refused at once instead of being read a second time. Returns the
|
||||
/// network and whether it is on the card.
|
||||
pub(super) fn load_or_processor<T>(
|
||||
on_card: bool,
|
||||
card: impl FnOnce() -> Result<T, AiError>,
|
||||
processor: impl FnOnce() -> Result<T, AiError>,
|
||||
) -> Result<(T, bool), AiError> {
|
||||
if !on_card {
|
||||
return Ok((processor()?, false));
|
||||
}
|
||||
match card() {
|
||||
Ok(network) => Ok((network, true)),
|
||||
Err(error) if error.kind == ErrorKind::Unsupported => Err(error),
|
||||
Err(_) => Ok((processor()?, false)),
|
||||
}
|
||||
}
|
||||
|
||||
/// Builds the network again on the processor. Called at most once.
|
||||
pub(super) type Rebuild = Box<dyn FnOnce() -> Result<Box<dyn Predict>, AiError> + Send>;
|
||||
|
||||
/// The network on a graphics card, and the way to build it on the processor instead.
|
||||
pub(super) struct WithFallback {
|
||||
current: Box<dyn Predict>,
|
||||
/// `None` once the network has moved to the processor.
|
||||
rebuild: Option<Rebuild>,
|
||||
}
|
||||
|
||||
impl WithFallback {
|
||||
pub(super) fn new(on_card: Box<dyn Predict>, rebuild: Rebuild) -> Self {
|
||||
Self {
|
||||
current: on_card,
|
||||
rebuild: Some(rebuild),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl Predict for WithFallback {
|
||||
fn predict(&mut self, input: &Input, ctx: &RunContext) -> Result<Vec<f32>, AiError> {
|
||||
let failed = match self.current.predict(input, ctx) {
|
||||
Err(error) if error.kind != ErrorKind::Cancelled => error,
|
||||
answer => return answer,
|
||||
};
|
||||
let Some(rebuild) = self.rebuild.take() else {
|
||||
return Err(failed);
|
||||
};
|
||||
ctx.notice(
|
||||
"The graphics card could not run the background remover, so it continues on the processor. This is slower.",
|
||||
);
|
||||
self.current = rebuild()?;
|
||||
self.current.predict(input, ctx)
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use std::sync::Arc;
|
||||
use std::sync::atomic::{AtomicUsize, Ordering};
|
||||
|
||||
use super::*;
|
||||
|
||||
/// Answers every pixel with its number, counting its calls.
|
||||
struct Fixed(f32, Arc<AtomicUsize>);
|
||||
|
||||
impl Predict for Fixed {
|
||||
fn predict(&mut self, input: &Input, _: &RunContext) -> Result<Vec<f32>, AiError> {
|
||||
self.1.fetch_add(1, Ordering::SeqCst);
|
||||
Ok(vec![self.0; input.width * input.height])
|
||||
}
|
||||
}
|
||||
|
||||
/// Fails every call with an error of `kind`, counting its calls.
|
||||
struct Failing(ErrorKind, Arc<AtomicUsize>);
|
||||
|
||||
impl Predict for Failing {
|
||||
fn predict(&mut self, _: &Input, _: &RunContext) -> Result<Vec<f32>, AiError> {
|
||||
self.1.fetch_add(1, Ordering::SeqCst);
|
||||
Err(AiError::new(self.0, "the card failed"))
|
||||
}
|
||||
}
|
||||
|
||||
fn input() -> Input {
|
||||
Input {
|
||||
width: 4,
|
||||
height: 2,
|
||||
px: vec![[0.5; 3]; 8],
|
||||
}
|
||||
}
|
||||
|
||||
/// A model on a card that fails with `kind`, whose rebuild answers 1.0. Returns the model and
|
||||
/// the counts of card calls, rebuilds and processor calls.
|
||||
fn failing_card(kind: ErrorKind) -> (WithFallback, [Arc<AtomicUsize>; 3]) {
|
||||
let counts = [(); 3].map(|_| Arc::new(AtomicUsize::new(0)));
|
||||
let [card, rebuilds, processor] = counts.clone();
|
||||
let model = WithFallback::new(
|
||||
Box::new(Failing(kind, card)),
|
||||
Box::new(move || {
|
||||
rebuilds.fetch_add(1, Ordering::SeqCst);
|
||||
Ok(Box::new(Fixed(1.0, processor)) as Box<dyn Predict>)
|
||||
}),
|
||||
);
|
||||
(model, counts)
|
||||
}
|
||||
|
||||
fn count(c: &Arc<AtomicUsize>) -> usize {
|
||||
c.load(Ordering::SeqCst)
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn a_card_that_works_is_kept() {
|
||||
let calls = Arc::new(AtomicUsize::new(0));
|
||||
let mut model = WithFallback::new(
|
||||
Box::new(Fixed(0.25, calls.clone())),
|
||||
Box::new(|| Err(AiError::invalid("no rebuild expected"))),
|
||||
);
|
||||
for _ in 0..2 {
|
||||
let answer = model.predict(&input(), &RunContext::detached()).unwrap();
|
||||
assert_eq!(answer, vec![0.25; 8]);
|
||||
}
|
||||
assert_eq!(count(&calls), 2);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn a_failing_card_moves_to_the_processor_once_and_stays_there() {
|
||||
let (mut model, [card, rebuilds, processor]) = failing_card(ErrorKind::Internal);
|
||||
let ctx = RunContext::detached();
|
||||
assert_eq!(model.predict(&input(), &ctx).unwrap(), vec![1.0; 8]);
|
||||
assert_eq!(model.predict(&input(), &ctx).unwrap(), vec![1.0; 8]);
|
||||
assert_eq!(
|
||||
(count(&card), count(&rebuilds), count(&processor)),
|
||||
(1, 1, 2)
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn a_cancel_is_not_a_failure() {
|
||||
let (mut model, [card, rebuilds, _]) = failing_card(ErrorKind::Cancelled);
|
||||
let error = model
|
||||
.predict(&input(), &RunContext::detached())
|
||||
.unwrap_err();
|
||||
assert_eq!(error.kind, ErrorKind::Cancelled);
|
||||
assert_eq!((count(&card), count(&rebuilds)), (1, 0));
|
||||
}
|
||||
|
||||
/// Loads that answer `Ok(name)` or fail with `kind`, recording which ran.
|
||||
fn load(
|
||||
name: &'static str,
|
||||
fails: Option<ErrorKind>,
|
||||
ran: &std::sync::Mutex<Vec<&'static str>>,
|
||||
) -> Result<&'static str, AiError> {
|
||||
ran.lock().unwrap().push(name);
|
||||
match fails {
|
||||
Some(kind) => Err(AiError::new(kind, "failed")),
|
||||
None => Ok(name),
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn a_card_that_loads_keeps_the_network() {
|
||||
let ran = std::sync::Mutex::new(Vec::new());
|
||||
let got = load_or_processor(
|
||||
true,
|
||||
|| load("card", None, &ran),
|
||||
|| load("cpu", None, &ran),
|
||||
);
|
||||
assert_eq!(got.unwrap(), ("card", true));
|
||||
assert_eq!(*ran.lock().unwrap(), ["card"]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn a_card_that_fails_to_load_loads_on_the_processor() {
|
||||
let ran = std::sync::Mutex::new(Vec::new());
|
||||
let got = load_or_processor(
|
||||
true,
|
||||
|| load("card", Some(ErrorKind::Provider), &ran),
|
||||
|| load("cpu", None, &ran),
|
||||
);
|
||||
assert_eq!(got.unwrap(), ("cpu", false));
|
||||
assert_eq!(*ran.lock().unwrap(), ["card", "cpu"]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn a_file_that_is_not_the_checkpoint_is_not_read_twice() {
|
||||
let ran = std::sync::Mutex::new(Vec::new());
|
||||
let got = load_or_processor(
|
||||
true,
|
||||
|| load("card", Some(ErrorKind::Unsupported), &ran),
|
||||
|| load("cpu", None, &ran),
|
||||
);
|
||||
assert_eq!(got.unwrap_err().kind, ErrorKind::Unsupported);
|
||||
assert_eq!(*ran.lock().unwrap(), ["card"]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn on_the_processor_there_is_one_load() {
|
||||
let ran = std::sync::Mutex::new(Vec::new());
|
||||
let got = load_or_processor(
|
||||
false,
|
||||
|| load("card", None, &ran),
|
||||
|| load("cpu", Some(ErrorKind::Provider), &ran),
|
||||
);
|
||||
assert_eq!(got.unwrap_err().kind, ErrorKind::Provider);
|
||||
assert_eq!(*ran.lock().unwrap(), ["cpu"]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn a_rebuild_that_fails_is_the_answer() {
|
||||
let card = Arc::new(AtomicUsize::new(0));
|
||||
let mut model = WithFallback::new(
|
||||
Box::new(Failing(ErrorKind::Provider, card)),
|
||||
Box::new(|| Err(AiError::unsupported("the checkpoint went away"))),
|
||||
);
|
||||
let error = model
|
||||
.predict(&input(), &RunContext::detached())
|
||||
.unwrap_err();
|
||||
assert_eq!(error.kind, ErrorKind::Unsupported);
|
||||
// Nothing is left to try: the card's own failure from now on.
|
||||
let again = model
|
||||
.predict(&input(), &RunContext::detached())
|
||||
.unwrap_err();
|
||||
assert_eq!(again.kind, ErrorKind::Provider);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,286 @@
|
||||
//! InSPyReNet on the Burn engine for `image.matte` (Remove background): the subject's soft alpha
|
||||
//! at the picture's size, from the picture alone.
|
||||
//!
|
||||
//! - **Network**: InSPyReNet with the Swin-B backbone (T. Kim et al., "Revisiting Image Pyramid
|
||||
//! Structure for High Resolution Salient Object Detection", ACCV 2022), written by hand in Burn
|
||||
//! after the MIT-licensed reference code (`net/`). It reads the published checkpoint as
|
||||
//! `.safetensors` and is checked against PyTorch, not generated from an ONNX export.
|
||||
//! - **Any size**: the network is trained near one megapixel. `pipeline/` runs it at full
|
||||
//! resolution up to 1280 pixels on the long side, and above that on a reduced square whose matte
|
||||
//! is enlarged back and refined at full resolution with a colour guided filter.
|
||||
//! - **Answer**: one grayscale PNG at the picture's size, white where the subject is, with soft
|
||||
//! values along its edge (hair, fur).
|
||||
|
||||
mod fallback;
|
||||
mod net;
|
||||
mod pipeline;
|
||||
mod weights;
|
||||
|
||||
use std::path::Path;
|
||||
|
||||
use burn::prelude::*;
|
||||
use burn::tensor::TensorData;
|
||||
use craft_ai_core::{
|
||||
AiError, CapabilityId, CapabilitySpec, ErrorKind, InputRole, ModelId, ModelInfo, Output,
|
||||
Payload, Request,
|
||||
};
|
||||
use craft_ai_engine::{LoadedModel, ModelRunner, RunContext};
|
||||
use craft_ai_engine_rten::picture::{decode_picture, encode_mask};
|
||||
use craft_ai_features::GpuVendor;
|
||||
|
||||
use crate::backend::{pick_device, processor};
|
||||
use fallback::{WithFallback, load_or_processor};
|
||||
use net::InSPyReNet;
|
||||
use net::swin::SwinConfig;
|
||||
use pipeline::{Input, MAX_RES, Picture, SHRINK};
|
||||
|
||||
/// The model id as the catalogue writes it (`NativeRoute.model`).
|
||||
pub const MODEL: &str = "inspyrenet";
|
||||
|
||||
/// The capability this runner answers.
|
||||
const CAPABILITY: &str = "image.matte";
|
||||
|
||||
/// ImageNet channel means and standard deviations: the network's input normalisation.
|
||||
const MEAN: [f32; 3] = [0.485, 0.456, 0.406];
|
||||
const STD: [f32; 3] = [0.229, 0.224, 0.225];
|
||||
|
||||
/// Stack for building and running the network. Burn recurses deeply through the module tree, more
|
||||
/// than a default thread has on Windows.
|
||||
const STACK: usize = 64 << 20;
|
||||
|
||||
/// The card makers Burn's graphics layer reaches, as for SAM 2.1 on this engine.
|
||||
const CARDS: [GpuVendor; 4] = [
|
||||
GpuVendor::Nvidia,
|
||||
GpuVendor::Amd,
|
||||
GpuVendor::Intel,
|
||||
GpuVendor::Apple,
|
||||
];
|
||||
|
||||
/// InSPyReNet on the Burn engine. It uses a graphics card when one answers, else the processor.
|
||||
pub struct InspyrenetRunner {
|
||||
/// The backbone size: Swin-B for the published checkpoint, a tiny one in tests.
|
||||
backbone: SwinConfig,
|
||||
/// Where a load puts the network: the graphics card when one answers, the processor in tests
|
||||
/// (a build server's software renderer answers too, and is far too slow for them).
|
||||
device: fn() -> Device,
|
||||
}
|
||||
|
||||
impl InspyrenetRunner {
|
||||
/// The runner for the published Swin-B checkpoint.
|
||||
pub const fn new() -> Self {
|
||||
Self {
|
||||
backbone: SwinConfig::BASE,
|
||||
device: pick_device,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl Default for InspyrenetRunner {
|
||||
fn default() -> Self {
|
||||
Self::new()
|
||||
}
|
||||
}
|
||||
|
||||
/// Runs `work` on a thread with [`STACK`] bytes of stack. A panic in it (Burn and the graphics
|
||||
/// layer panic on what they cannot handle) comes back as an error a person can read.
|
||||
fn on_big_stack<T: Send>(work: impl FnOnce() -> Result<T, AiError> + Send) -> Result<T, AiError> {
|
||||
let stopped = || {
|
||||
AiError::new(
|
||||
ErrorKind::Provider,
|
||||
"The background remover stopped unexpectedly. Try again, or choose another in Set up AI.",
|
||||
)
|
||||
};
|
||||
std::thread::scope(|scope| {
|
||||
let worker = std::thread::Builder::new()
|
||||
.name("craft-ai-inspyrenet".into())
|
||||
.stack_size(STACK)
|
||||
.spawn_scoped(scope, work)
|
||||
.map_err(|_| stopped())?;
|
||||
worker.join().unwrap_or_else(|_| Err(stopped()))
|
||||
})
|
||||
}
|
||||
|
||||
impl ModelRunner for InspyrenetRunner {
|
||||
fn engine(&self) -> &str {
|
||||
crate::runner::ENGINE
|
||||
}
|
||||
|
||||
/// `image.matte` from one required picture and nothing else.
|
||||
fn model(&self) -> ModelInfo {
|
||||
let mut info = ModelInfo::new(model_id(), "InSPyReNet, graphics card");
|
||||
if let Ok(capability) = CapabilityId::parse(CAPABILITY) {
|
||||
info = info.with(
|
||||
CapabilitySpec::new(capability)
|
||||
.accepting(InputRole::Source, 1)
|
||||
.requiring(InputRole::Source),
|
||||
);
|
||||
}
|
||||
info
|
||||
}
|
||||
|
||||
fn accelerators(&self) -> &[GpuVendor] {
|
||||
&CARDS
|
||||
}
|
||||
|
||||
/// Reads the checkpoint (`inspyrenet.safetensors`) or, in tests, a saved copy (`.bpk`).
|
||||
fn load(&self, weights: &Path) -> Result<Box<dyn LoadedModel>, AiError> {
|
||||
if !weights.is_file() {
|
||||
return Err(AiError::unsupported(
|
||||
"The InSPyReNet checkpoint is missing. Open Set up AI to download it.",
|
||||
));
|
||||
}
|
||||
let backbone = self.backbone;
|
||||
let device = (self.device)();
|
||||
let (network, on_card) = load_or_processor(
|
||||
device != processor(),
|
||||
|| load_network(backbone, weights, device),
|
||||
|| load_network(backbone, weights, processor()),
|
||||
)?;
|
||||
let mut predict: Box<dyn Predict> = Box::new(network);
|
||||
if on_card {
|
||||
let path = weights.to_path_buf();
|
||||
predict = Box::new(WithFallback::new(
|
||||
predict,
|
||||
Box::new(move || {
|
||||
let network = load_network(backbone, &path, processor())?;
|
||||
Ok(Box::new(network) as Box<dyn Predict>)
|
||||
}),
|
||||
));
|
||||
}
|
||||
Ok(Box::new(Loaded::new(predict)))
|
||||
}
|
||||
}
|
||||
|
||||
/// Builds the network on `device` and fills it from the file at `path`, on a thread with a large
|
||||
/// stack.
|
||||
fn load_network(backbone: SwinConfig, path: &Path, device: Device) -> Result<Network, AiError> {
|
||||
let path = path.to_path_buf();
|
||||
on_big_stack(move || {
|
||||
let net = weights::fill(InSPyReNet::with_backbone(backbone, &device), &path)?;
|
||||
Ok(Network {
|
||||
net: Box::new(net),
|
||||
device,
|
||||
})
|
||||
})
|
||||
}
|
||||
|
||||
/// The model id, from a literal that always parses.
|
||||
#[allow(clippy::expect_used)] // `inspyrenet` is a valid id: a literal that cannot fail
|
||||
fn model_id() -> ModelId {
|
||||
ModelId::new(MODEL).expect("inspyrenet is a valid model id")
|
||||
}
|
||||
|
||||
/// What the pipeline asks for a picture: the foreground probability of `input`, one value per
|
||||
/// pixel. The network answers it; tests give a stand-in, as the SAM 2.1 tests here do.
|
||||
trait Predict: Send {
|
||||
fn predict(&mut self, input: &Input, ctx: &RunContext) -> Result<Vec<f32>, AiError>;
|
||||
}
|
||||
|
||||
/// InSPyReNet with its weights, on the device they were loaded to.
|
||||
struct Network {
|
||||
/// Boxed: the module tree is a few hundred kilobytes by value.
|
||||
net: Box<InSPyReNet>,
|
||||
device: Device,
|
||||
}
|
||||
|
||||
impl Network {
|
||||
/// `input` as the network's normalised `[1, 3, H, W]` tensor.
|
||||
fn tensor(&self, input: &Input) -> Tensor<4> {
|
||||
let (w, h) = (input.width, input.height);
|
||||
let mut planar = vec![0.0f32; 3 * w * h];
|
||||
for (i, p) in input.px.iter().enumerate() {
|
||||
for c in 0..3 {
|
||||
if let Some(slot) = planar.get_mut(c * w * h + i) {
|
||||
*slot = (p[c] - MEAN[c]) / STD[c];
|
||||
}
|
||||
}
|
||||
}
|
||||
Tensor::from_data(TensorData::new(planar, [1, 3, h, w]), &self.device)
|
||||
}
|
||||
}
|
||||
|
||||
impl Predict for Network {
|
||||
/// The network's foreground probability for `input`, one value per pixel.
|
||||
fn predict(&mut self, input: &Input, ctx: &RunContext) -> Result<Vec<f32>, AiError> {
|
||||
let ctx = ctx.clone();
|
||||
on_big_stack(|| {
|
||||
let cancelled = || ctx.is_cancelled();
|
||||
let pred = self
|
||||
.net
|
||||
.forward(self.tensor(input), &cancelled)
|
||||
.ok_or_else(AiError::cancelled)?;
|
||||
// A device that fails while running reports it here; `into_data` would panic.
|
||||
pred.try_into_data()
|
||||
.ok()
|
||||
.and_then(|data| data.try_to_vec::<f32>().ok())
|
||||
.ok_or_else(|| {
|
||||
AiError::new(
|
||||
ErrorKind::Internal,
|
||||
"The background remover's answer was unreadable.",
|
||||
)
|
||||
})
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
/// A loaded model: the any-size pipeline around whatever predicts.
|
||||
struct Loaded {
|
||||
predict: Box<dyn Predict>,
|
||||
}
|
||||
|
||||
impl Loaded {
|
||||
fn new(predict: Box<dyn Predict>) -> Self {
|
||||
Self { predict }
|
||||
}
|
||||
}
|
||||
|
||||
impl LoadedModel for Loaded {
|
||||
/// Decodes the picture, runs the any-size pipeline around the network and returns the matte as
|
||||
/// a grayscale PNG at the picture's size.
|
||||
fn run(&mut self, request: &Request, ctx: &RunContext) -> Result<Vec<Output>, AiError> {
|
||||
let Some(source) = request
|
||||
.inputs
|
||||
.iter()
|
||||
.find(|i| i.role == InputRole::Source)
|
||||
.map(|i| &i.payload)
|
||||
else {
|
||||
return Err(AiError::invalid("A picture is needed."));
|
||||
};
|
||||
let Payload::Inline { data, .. } = source else {
|
||||
return Err(AiError::unsupported(
|
||||
"The background remover needs the picture itself, not a link to it.",
|
||||
));
|
||||
};
|
||||
let rgb = decode_picture(data)?;
|
||||
ctx.progress(5);
|
||||
ctx.check_cancelled()?;
|
||||
let picture = Picture {
|
||||
width: rgb.width,
|
||||
height: rgb.height,
|
||||
bytes: &rgb.bytes,
|
||||
};
|
||||
ctx.notice(
|
||||
"Finding the subject. This takes a few seconds, longer without a graphics card.",
|
||||
);
|
||||
let mask = pipeline::matte(
|
||||
picture,
|
||||
MAX_RES,
|
||||
SHRINK,
|
||||
&mut |input| self.predict.predict(input, ctx),
|
||||
ctx,
|
||||
)?;
|
||||
ctx.check_cancelled()?;
|
||||
let alpha: Vec<f32> = mask.iter().map(|v| *v as f32 / 255.0).collect();
|
||||
let png = encode_mask(&alpha, rgb.width, rgb.height)?;
|
||||
ctx.progress(95);
|
||||
Ok(vec![Output {
|
||||
payload: Payload::inline("image/png", png)?,
|
||||
label: None,
|
||||
}])
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod parity;
|
||||
#[cfg(test)]
|
||||
mod tests;
|
||||
@@ -0,0 +1,345 @@
|
||||
//! InSPyReNet (T. Kim, K. Kim, J. Lee, D. Cha, J. Lee and D. Kim, "Revisiting Image Pyramid
|
||||
//! Structure for High Resolution Salient Object Detection", ACCV 2022), ported from the
|
||||
//! MIT-licensed reference implementation (github.com/plemeri/InSPyReNet, as packaged by
|
||||
//! transparent-background). Module and field names follow the PyTorch code, so its checkpoints
|
||||
//! load by key.
|
||||
//!
|
||||
//! The network predicts a saliency pyramid: a coarse map from the decoder at 1/8 resolution,
|
||||
//! refined twice by scale-invariant context attention (SICA) through a Laplacian pyramid up to
|
||||
//! the input size.
|
||||
|
||||
#![allow(non_snake_case)] // `Hattn` / `Wattn`, as named in the checkpoints
|
||||
|
||||
use burn::module::{Module, Param};
|
||||
use burn::tensor::activation::{relu, sigmoid, softmax};
|
||||
use burn::tensor::{Device, Tensor};
|
||||
|
||||
pub mod nn;
|
||||
pub mod swin;
|
||||
|
||||
use nn::{ConvBn, ImagePyramid, SelfAttention, resize};
|
||||
use swin::{Swin, SwinConfig};
|
||||
|
||||
/// Channels of the decoder and attention features.
|
||||
pub const DEPTH: usize = 64;
|
||||
/// The size SICA pools its context at, at stage 0 (InSPyReNet's `base_size`); stage `s` uses
|
||||
/// `BASE_SIZE >> s`.
|
||||
pub const BASE_SIZE: usize = 1024;
|
||||
|
||||
/// `PAA_kernel`: a 1 × 1 projection, a 1 × r and an r × 1 convolution, axial attention along
|
||||
/// both axes, then a 3 × 3 convolution dilated by r.
|
||||
#[derive(Module, Debug)]
|
||||
pub struct PaaKernel {
|
||||
pub conv0: ConvBn,
|
||||
pub conv1: ConvBn,
|
||||
pub conv2: ConvBn,
|
||||
pub conv3: ConvBn,
|
||||
pub Hattn: SelfAttention,
|
||||
pub Wattn: SelfAttention,
|
||||
}
|
||||
|
||||
impl PaaKernel {
|
||||
fn new(c_in: usize, c_out: usize, r: usize, dev: &Device) -> Self {
|
||||
Self {
|
||||
conv0: ConvBn::new(c_in, c_out, [1, 1], 1, true, false, dev),
|
||||
conv1: ConvBn::new(c_out, c_out, [1, r], 1, true, false, dev),
|
||||
conv2: ConvBn::new(c_out, c_out, [r, 1], 1, true, false, dev),
|
||||
conv3: ConvBn::new(c_out, c_out, [3, 3], r, true, false, dev),
|
||||
Hattn: SelfAttention::new(c_out, true, dev),
|
||||
Wattn: SelfAttention::new(c_out, false, dev),
|
||||
}
|
||||
}
|
||||
|
||||
fn forward(&self, x: Tensor<4>) -> Tensor<4> {
|
||||
let x = self
|
||||
.conv2
|
||||
.forward(self.conv1.forward(self.conv0.forward(x)));
|
||||
self.conv3
|
||||
.forward(self.Hattn.forward(x.clone()) + self.Wattn.forward(x))
|
||||
}
|
||||
}
|
||||
|
||||
/// `PAA_e`: the context module on each backbone output.
|
||||
#[derive(Module, Debug)]
|
||||
pub struct PaaE {
|
||||
pub branch0: ConvBn,
|
||||
pub branch1: PaaKernel,
|
||||
pub branch2: PaaKernel,
|
||||
pub branch3: PaaKernel,
|
||||
pub conv_cat: ConvBn,
|
||||
pub conv_res: ConvBn,
|
||||
}
|
||||
|
||||
impl PaaE {
|
||||
fn new(c_in: usize, c_out: usize, dev: &Device) -> Self {
|
||||
Self {
|
||||
branch0: ConvBn::new(c_in, c_out, [1, 1], 1, true, false, dev),
|
||||
branch1: PaaKernel::new(c_in, c_out, 3, dev),
|
||||
branch2: PaaKernel::new(c_in, c_out, 5, dev),
|
||||
branch3: PaaKernel::new(c_in, c_out, 7, dev),
|
||||
conv_cat: ConvBn::square(4 * c_out, c_out, 3, false, dev),
|
||||
conv_res: ConvBn::new(c_in, c_out, [1, 1], 1, true, false, dev),
|
||||
}
|
||||
}
|
||||
|
||||
fn forward(&self, x: Tensor<4>) -> Tensor<4> {
|
||||
let cat = Tensor::cat(
|
||||
vec![
|
||||
self.branch0.forward(x.clone()),
|
||||
self.branch1.forward(x.clone()),
|
||||
self.branch2.forward(x.clone()),
|
||||
self.branch3.forward(x.clone()),
|
||||
],
|
||||
1,
|
||||
);
|
||||
relu(self.conv_cat.forward(cat) + self.conv_res.forward(x))
|
||||
}
|
||||
}
|
||||
|
||||
/// `PAA_d`: the decoder that turns the three deepest context features into the coarse map.
|
||||
#[derive(Module, Debug)]
|
||||
pub struct PaaD {
|
||||
pub conv1: ConvBn,
|
||||
pub conv2: ConvBn,
|
||||
pub conv3: ConvBn,
|
||||
pub conv4: ConvBn,
|
||||
pub conv5: ConvBn,
|
||||
pub Hattn: SelfAttention,
|
||||
pub Wattn: SelfAttention,
|
||||
}
|
||||
|
||||
impl PaaD {
|
||||
fn new(c_in: usize, depth: usize, dev: &Device) -> Self {
|
||||
Self {
|
||||
conv1: ConvBn::square(c_in, depth, 3, false, dev),
|
||||
conv2: ConvBn::square(depth, depth, 3, false, dev),
|
||||
conv3: ConvBn::square(depth, depth, 3, false, dev),
|
||||
conv4: ConvBn::square(depth, depth, 3, false, dev),
|
||||
conv5: ConvBn::new(depth, 1, [3, 3], 1, false, false, dev),
|
||||
Hattn: SelfAttention::new(depth, true, dev),
|
||||
Wattn: SelfAttention::new(depth, false, dev),
|
||||
}
|
||||
}
|
||||
|
||||
/// `(features, coarse map)` from `[x3, x4, x5]` (finest first), at `x3`'s size.
|
||||
fn forward(&self, x3: Tensor<4>, x4: Tensor<4>, x5: Tensor<4>) -> (Tensor<4>, Tensor<4>) {
|
||||
let [_, _, h, w] = x3.dims();
|
||||
let fx = Tensor::cat(
|
||||
vec![resize(x5, [h, w], true), resize(x4, [h, w], true), x3],
|
||||
1,
|
||||
);
|
||||
let fx = self.conv1.forward(fx);
|
||||
let fx = self
|
||||
.conv2
|
||||
.forward(self.Hattn.forward(fx.clone()) + self.Wattn.forward(fx));
|
||||
let fx = self.conv4.forward(self.conv3.forward(fx));
|
||||
let out = self.conv5.forward(fx.clone());
|
||||
(fx, out)
|
||||
}
|
||||
}
|
||||
|
||||
/// `SICA`: scale-invariant context attention. It splits the coarser map into foreground,
|
||||
/// background and uncertain areas (plus the Laplacian's, with `lmap_in`), pools one context
|
||||
/// vector per area at a fixed size, and lets every pixel attend to them.
|
||||
#[derive(Module, Debug)]
|
||||
pub struct Sica {
|
||||
pub conv_query: Vec<ConvBn>,
|
||||
pub conv_key: Vec<ConvBn>,
|
||||
pub conv_value: Vec<ConvBn>,
|
||||
pub conv_out1: ConvBn,
|
||||
pub conv_out2: ConvBn,
|
||||
pub conv_out3: ConvBn,
|
||||
pub conv_out4: ConvBn,
|
||||
pub threshold: Param<Tensor<1>>,
|
||||
pub lthreshold: Option<Param<Tensor<1>>>,
|
||||
pub depth: usize,
|
||||
/// The side the context is pooled at.
|
||||
pub stage_size: usize,
|
||||
}
|
||||
|
||||
impl Sica {
|
||||
fn new(c_in: usize, depth: usize, stage: usize, lmap_in: bool, dev: &Device) -> Self {
|
||||
let pair = |k: usize| {
|
||||
vec![
|
||||
ConvBn::square(c_in, depth, k, true, dev),
|
||||
ConvBn::square(depth, depth, k, true, dev),
|
||||
]
|
||||
};
|
||||
let half = || Param::from_tensor(Tensor::from_floats([0.5], dev));
|
||||
Self {
|
||||
conv_query: pair(3),
|
||||
conv_key: pair(1),
|
||||
conv_value: pair(1),
|
||||
conv_out1: ConvBn::square(depth, depth, 3, true, dev),
|
||||
conv_out2: ConvBn::square(c_in + depth, depth, 3, true, dev),
|
||||
conv_out3: ConvBn::square(depth, depth, 3, true, dev),
|
||||
conv_out4: ConvBn::square(depth, 1, 1, false, dev),
|
||||
threshold: half(),
|
||||
lthreshold: lmap_in.then(half),
|
||||
depth,
|
||||
stage_size: BASE_SIZE >> stage,
|
||||
}
|
||||
}
|
||||
|
||||
fn seq(convs: &[ConvBn], x: Tensor<4>) -> Tensor<4> {
|
||||
convs.iter().fold(x, |x, c| c.forward(x))
|
||||
}
|
||||
|
||||
/// `(features, Laplacian detail)` for `x`, guided by the coarser saliency `smap` and, with
|
||||
/// `lmap_in`, the coarser Laplacian `lmap`.
|
||||
fn forward(
|
||||
&self,
|
||||
x: Tensor<4>,
|
||||
smap: Tensor<4>,
|
||||
lmap: Option<Tensor<4>>,
|
||||
) -> (Tensor<4>, Tensor<4>) {
|
||||
let [b, c, h, w] = x.dims();
|
||||
let t = self.threshold.val().reshape([1, 1, 1, 1]);
|
||||
let p = sigmoid(resize(smap, [h, w], false)) - t.clone();
|
||||
let mut prob = vec![
|
||||
p.clone().clamp(0.0, 1.0),
|
||||
p.clone().neg().clamp(0.0, 1.0),
|
||||
t - p.abs(),
|
||||
];
|
||||
if let (Some(lt), Some(lmap)) = (&self.lthreshold, lmap) {
|
||||
let lp = sigmoid(resize(lmap, [h, w], false)) - lt.val().reshape([1, 1, 1, 1]);
|
||||
prob.push(lp.clone().clamp(0.0, 1.0));
|
||||
prob.push(lp.neg().clamp(0.0, 1.0));
|
||||
}
|
||||
let ctx = prob.len();
|
||||
let (s, d) = (self.stage_size, self.depth);
|
||||
// PyTorch views these as (b, s², −1) and (b, ctx, s²): reshapes in memory order.
|
||||
let f = resize(x.clone(), [s, s], false).reshape([b, s * s, c]);
|
||||
let prob = resize(Tensor::cat(prob, 1), [s, s], false).reshape([b, ctx, s * s]);
|
||||
let context = prob.matmul(f).swap_dims(1, 2).unsqueeze_dim::<4>(3); // [b, c, ctx, 1]
|
||||
let query = Self::seq(&self.conv_query, x.clone())
|
||||
.reshape([b, d, h * w])
|
||||
.swap_dims(1, 2);
|
||||
let key = Self::seq(&self.conv_key, context.clone()).reshape([b, d, ctx]);
|
||||
let value = Self::seq(&self.conv_value, context)
|
||||
.reshape([b, d, ctx])
|
||||
.swap_dims(1, 2);
|
||||
let sim = softmax(query.matmul(key) * (d as f64).powf(-0.5), 2);
|
||||
let context = self
|
||||
.conv_out1
|
||||
.forward(sim.matmul(value).swap_dims(1, 2).reshape([b, d, h, w]));
|
||||
let x = self
|
||||
.conv_out3
|
||||
.forward(self.conv_out2.forward(Tensor::cat(vec![x, context], 1)));
|
||||
let out = self.conv_out4.forward(x.clone());
|
||||
(x, out)
|
||||
}
|
||||
}
|
||||
|
||||
/// Everything the network computes on the way, for parity tests against the reference. Only the
|
||||
/// finest saliency is used outside them.
|
||||
pub struct Parts {
|
||||
/// The patch embedding and the four Swin stages.
|
||||
#[cfg_attr(not(test), allow(dead_code))]
|
||||
pub backbone: [Tensor<4>; 5],
|
||||
/// Saliency logits, coarse (1/8) to fine (full size).
|
||||
pub saliency: [Tensor<4>; 4],
|
||||
/// Laplacian details at 1/4, 1/2 and full size.
|
||||
#[cfg_attr(not(test), allow(dead_code))]
|
||||
pub laplacian: [Tensor<4>; 3],
|
||||
}
|
||||
|
||||
/// InSPyReNet with the Swin-B backbone (`InSPyReNet_SwinB`).
|
||||
#[derive(Module, Debug)]
|
||||
pub struct InSPyReNet {
|
||||
pub backbone: Swin,
|
||||
pub context1: PaaE,
|
||||
pub context2: PaaE,
|
||||
pub context3: PaaE,
|
||||
pub context4: PaaE,
|
||||
pub context5: PaaE,
|
||||
pub decoder: PaaD,
|
||||
pub attention0: Sica,
|
||||
pub attention1: Sica,
|
||||
pub attention2: Sica,
|
||||
}
|
||||
|
||||
impl InSPyReNet {
|
||||
/// The Swin-B architecture with placeholder weights (load a checkpoint before use).
|
||||
#[cfg(test)]
|
||||
pub fn new(dev: &Device) -> Self {
|
||||
Self::with_backbone(SwinConfig::BASE, dev)
|
||||
}
|
||||
|
||||
/// The same network over a Swin backbone of another size (tests use a tiny one).
|
||||
pub fn with_backbone(cfg: SwinConfig, dev: &Device) -> Self {
|
||||
let e = cfg.embed_dim;
|
||||
Self {
|
||||
backbone: Swin::new(cfg, dev),
|
||||
context1: PaaE::new(e, DEPTH, dev),
|
||||
context2: PaaE::new(e, DEPTH, dev),
|
||||
context3: PaaE::new(2 * e, DEPTH, dev),
|
||||
context4: PaaE::new(4 * e, DEPTH, dev),
|
||||
context5: PaaE::new(8 * e, DEPTH, dev),
|
||||
decoder: PaaD::new(3 * DEPTH, DEPTH, dev),
|
||||
attention0: Sica::new(DEPTH, DEPTH, 0, true, dev),
|
||||
attention1: Sica::new(2 * DEPTH, DEPTH, 1, true, dev),
|
||||
attention2: Sica::new(2 * DEPTH, DEPTH, 2, false, dev),
|
||||
}
|
||||
}
|
||||
|
||||
/// The saliency pyramid of a normalised `[1, 3, H, W]` input (`H`, `W` multiples of 32):
|
||||
/// PyTorch's `forward_inspyre`. `None` once `cancelled` says yes (asked between stages).
|
||||
pub fn forward_parts(&self, x: Tensor<4>, cancelled: &dyn Fn() -> bool) -> Option<Parts> {
|
||||
let [_, _, h, w] = x.dims();
|
||||
let pyramid = ImagePyramid::new(&x.device());
|
||||
let backbone = self.backbone.forward(x, cancelled)?;
|
||||
if cancelled() {
|
||||
return None;
|
||||
}
|
||||
let [b1, b2, b3, b4, b5] = backbone.clone();
|
||||
let x1 = self.context1.forward(b1);
|
||||
let x2 = self.context2.forward(b2);
|
||||
let x3 = self.context3.forward(b3);
|
||||
let x4 = self.context4.forward(b4);
|
||||
let x5 = self.context5.forward(b5);
|
||||
if cancelled() {
|
||||
return None;
|
||||
}
|
||||
let (f3, d3) = self.decoder.forward(x3, x4, x5);
|
||||
|
||||
let f3 = resize(f3, [h / 4, w / 4], false);
|
||||
let (f2, p2) = self
|
||||
.attention2
|
||||
.forward(Tensor::cat(vec![x2, f3], 1), d3.clone(), None);
|
||||
let d2 = pyramid.reconstruct(d3.clone(), p2.clone());
|
||||
if cancelled() {
|
||||
return None;
|
||||
}
|
||||
|
||||
let x1 = resize(x1, [h / 2, w / 2], false);
|
||||
let f2 = resize(f2, [h / 2, w / 2], false);
|
||||
let (f1, p1) =
|
||||
self.attention1
|
||||
.forward(Tensor::cat(vec![x1, f2], 1), d2.clone(), Some(p2.clone()));
|
||||
let d1 = pyramid.reconstruct(d2.clone(), p1.clone());
|
||||
if cancelled() {
|
||||
return None;
|
||||
}
|
||||
|
||||
let f1 = resize(f1, [h, w], false);
|
||||
let (_, p0) = self.attention0.forward(f1, d1.clone(), Some(p1.clone()));
|
||||
let d0 = pyramid.reconstruct(d1.clone(), p0.clone());
|
||||
Some(Parts {
|
||||
backbone,
|
||||
saliency: [d3, d2, d1, d0],
|
||||
laplacian: [p2, p1, p0],
|
||||
})
|
||||
}
|
||||
|
||||
/// Foreground probability `[1, 1, H, W]`: the finest saliency through a sigmoid, stretched to
|
||||
/// `0..=1` (PyTorch's `forward_inference` without the LR/HR blending, as Gyre runs it).
|
||||
pub fn forward(&self, x: Tensor<4>, cancelled: &dyn Fn() -> bool) -> Option<Tensor<4>> {
|
||||
let parts = self.forward_parts(x, cancelled)?;
|
||||
let [_, _, _, d0] = parts.saliency;
|
||||
let pred = sigmoid(d0);
|
||||
let (lo, hi) = (pred.clone().min(), pred.clone().max());
|
||||
let range = hi - lo.clone() + 1e-8;
|
||||
Some((pred - lo.reshape([1, 1, 1, 1])) / range.reshape([1, 1, 1, 1]))
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,246 @@
|
||||
//! Building blocks shared by the networks, named after their PyTorch counterparts so a PyTorch
|
||||
//! checkpoint loads by key without renaming: `weight` / `bias` on the norms and PyTorch's `same`
|
||||
//! padding on the convolutions. Linear layers use Burn's own `[in, out]` layout: Burn's PyTorch
|
||||
//! reader transposes PyTorch's `[out, in]` weights as it loads them.
|
||||
|
||||
use burn::module::{Module, Param};
|
||||
use burn::nn::conv::{Conv2d, Conv2dConfig};
|
||||
use burn::nn::{Linear, LinearConfig, PaddingConfig2d};
|
||||
use burn::tensor::activation::{relu, softmax};
|
||||
use burn::tensor::module::{conv2d, interpolate};
|
||||
use burn::tensor::ops::{ConvOptions, InterpolateMode, InterpolateOptions, PadMode};
|
||||
use burn::tensor::{Device, Tensor, TensorData};
|
||||
|
||||
/// PyTorch's default normalisation epsilon.
|
||||
const EPS: f64 = 1e-5;
|
||||
|
||||
/// A linear layer (Burn's `[in, out]` weight layout).
|
||||
pub fn linear(d_in: usize, d_out: usize, bias: bool, dev: &Device) -> Linear {
|
||||
LinearConfig::new(d_in, d_out).with_bias(bias).init(dev)
|
||||
}
|
||||
|
||||
/// `nn.LayerNorm` over the last dimension.
|
||||
#[derive(Module, Debug)]
|
||||
pub struct LayerNorm {
|
||||
pub weight: Param<Tensor<1>>,
|
||||
pub bias: Param<Tensor<1>>,
|
||||
}
|
||||
|
||||
impl LayerNorm {
|
||||
pub fn new(d: usize, dev: &Device) -> Self {
|
||||
Self {
|
||||
weight: Param::from_tensor(Tensor::ones([d], dev)),
|
||||
bias: Param::from_tensor(Tensor::zeros([d], dev)),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn forward<const D: usize>(&self, x: Tensor<D>) -> Tensor<D> {
|
||||
let mean = x.clone().mean_dim(D - 1);
|
||||
let centred = x - mean;
|
||||
let var = centred.clone().square().mean_dim(D - 1);
|
||||
let norm = centred / (var + EPS).sqrt();
|
||||
let shape = |t: Tensor<1>| {
|
||||
let mut dims = [1usize; D];
|
||||
dims[D - 1] = t.dims()[0];
|
||||
t.reshape(dims)
|
||||
};
|
||||
norm * shape(self.weight.val()) + shape(self.bias.val())
|
||||
}
|
||||
}
|
||||
|
||||
/// `nn.BatchNorm2d` in evaluation mode (running statistics).
|
||||
#[derive(Module, Debug)]
|
||||
pub struct BatchNorm {
|
||||
pub weight: Param<Tensor<1>>,
|
||||
pub bias: Param<Tensor<1>>,
|
||||
pub running_mean: Param<Tensor<1>>,
|
||||
pub running_var: Param<Tensor<1>>,
|
||||
}
|
||||
|
||||
impl BatchNorm {
|
||||
pub fn new(c: usize, dev: &Device) -> Self {
|
||||
Self {
|
||||
weight: Param::from_tensor(Tensor::ones([c], dev)),
|
||||
bias: Param::from_tensor(Tensor::zeros([c], dev)),
|
||||
running_mean: Param::from_tensor(Tensor::zeros([c], dev)),
|
||||
running_var: Param::from_tensor(Tensor::ones([c], dev)),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn forward(&self, x: Tensor<4>) -> Tensor<4> {
|
||||
let c = self.weight.dims()[0];
|
||||
let scale = self.weight.val() / (self.running_var.val() + EPS).sqrt();
|
||||
let shift = self.bias.val() - self.running_mean.val() * scale.clone();
|
||||
x * scale.reshape([1, c, 1, 1]) + shift.reshape([1, c, 1, 1])
|
||||
}
|
||||
}
|
||||
|
||||
/// InSPyReNet's `Conv2d` wrapper: a bias-free convolution with PyTorch `same` padding
|
||||
/// (`dilation · (k − 1) / 2` per side), then batch norm and ReLU when enabled.
|
||||
#[derive(Module, Debug)]
|
||||
pub struct ConvBn {
|
||||
pub conv: Conv2d,
|
||||
pub bn: Option<BatchNorm>,
|
||||
pub relu: bool,
|
||||
}
|
||||
|
||||
impl ConvBn {
|
||||
pub fn new(
|
||||
c_in: usize,
|
||||
c_out: usize,
|
||||
kernel: [usize; 2],
|
||||
dilation: usize,
|
||||
bn: bool,
|
||||
relu: bool,
|
||||
dev: &Device,
|
||||
) -> Self {
|
||||
let (ph, pw) = (
|
||||
dilation * (kernel[0] - 1) / 2,
|
||||
dilation * (kernel[1] - 1) / 2,
|
||||
);
|
||||
let conv = Conv2dConfig::new([c_in, c_out], kernel)
|
||||
.with_dilation([dilation, dilation])
|
||||
.with_padding(PaddingConfig2d::Explicit(ph, pw, ph, pw))
|
||||
.with_bias(false)
|
||||
.init(dev);
|
||||
Self {
|
||||
conv,
|
||||
bn: bn.then(|| BatchNorm::new(c_out, dev)),
|
||||
relu,
|
||||
}
|
||||
}
|
||||
|
||||
/// A square `k × k` convolution with batch norm.
|
||||
pub fn square(c_in: usize, c_out: usize, k: usize, relu: bool, dev: &Device) -> Self {
|
||||
Self::new(c_in, c_out, [k, k], 1, true, relu, dev)
|
||||
}
|
||||
|
||||
pub fn forward(&self, x: Tensor<4>) -> Tensor<4> {
|
||||
let mut x = self.conv.forward(x);
|
||||
if let Some(bn) = &self.bn {
|
||||
x = bn.forward(x);
|
||||
}
|
||||
if self.relu { relu(x) } else { x }
|
||||
}
|
||||
}
|
||||
|
||||
/// InSPyReNet's axial `SelfAttention`. Its query, key and value maps are *reshaped* (not
|
||||
/// transposed) to `(b, −1, axis)`, which mixes rows and columns in memory order; the port does
|
||||
/// the same so it computes what the trained weights expect.
|
||||
#[derive(Module, Debug)]
|
||||
pub struct SelfAttention {
|
||||
pub query_conv: ConvBn,
|
||||
pub key_conv: ConvBn,
|
||||
pub value_conv: ConvBn,
|
||||
pub gamma: Param<Tensor<1>>,
|
||||
/// Attend along the height (`"h"`) rather than the width (`"w"`).
|
||||
pub along_height: bool,
|
||||
}
|
||||
|
||||
impl SelfAttention {
|
||||
pub fn new(c: usize, along_height: bool, dev: &Device) -> Self {
|
||||
Self {
|
||||
query_conv: ConvBn::new(c, c / 8, [1, 1], 1, true, false, dev),
|
||||
key_conv: ConvBn::new(c, c / 8, [1, 1], 1, true, false, dev),
|
||||
value_conv: ConvBn::new(c, c, [1, 1], 1, true, false, dev),
|
||||
gamma: Param::from_tensor(Tensor::zeros([1], dev)),
|
||||
along_height,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn forward(&self, x: Tensor<4>) -> Tensor<4> {
|
||||
let [b, c, h, w] = x.dims();
|
||||
let axis = if self.along_height { h } else { w };
|
||||
let flat = |t: Tensor<4>| {
|
||||
let [_, tc, th, tw] = t.dims();
|
||||
t.reshape([b, tc * th * tw / axis, axis])
|
||||
};
|
||||
let query = flat(self.query_conv.forward(x.clone())).swap_dims(1, 2);
|
||||
let key = flat(self.key_conv.forward(x.clone()));
|
||||
let attention = softmax(query.matmul(key), 2);
|
||||
let value = flat(self.value_conv.forward(x.clone()));
|
||||
let out = value
|
||||
.matmul(attention.swap_dims(1, 2))
|
||||
.reshape([b, c, h, w]);
|
||||
out * self.gamma.val().reshape([1, 1, 1, 1]) + x
|
||||
}
|
||||
}
|
||||
|
||||
/// Bilinear resize of `[b, c, h, w]` to `size`, with PyTorch's `align_corners` convention.
|
||||
pub fn resize(x: Tensor<4>, size: [usize; 2], align_corners: bool) -> Tensor<4> {
|
||||
let [_, _, h, w] = x.dims();
|
||||
if [h, w] == size {
|
||||
// The identity under either convention.
|
||||
return x;
|
||||
}
|
||||
interpolate(
|
||||
x,
|
||||
InterpolateOptions::new(InterpolateMode::Bilinear)
|
||||
.with_align_corners(align_corners)
|
||||
.with_output_size(size),
|
||||
)
|
||||
}
|
||||
|
||||
/// `F.pixel_shuffle`: `[b, c·r², h, w]` → `[b, c, h·r, w·r]`.
|
||||
pub fn pixel_shuffle(x: Tensor<4>, r: usize) -> Tensor<4> {
|
||||
let [b, crr, h, w] = x.dims();
|
||||
let c = crr / (r * r);
|
||||
x.reshape([b, c, r, r, h, w])
|
||||
.permute([0, 1, 4, 2, 5, 3])
|
||||
.reshape([b, c, h * r, w * r])
|
||||
}
|
||||
|
||||
/// The Laplacian image pyramid InSPyReNet reconstructs its saliency with: a 7 × 7 Gaussian with
|
||||
/// σ = 1 (OpenCV's `getGaussianKernel`, normalised), on one channel.
|
||||
pub struct ImagePyramid {
|
||||
kernel: Tensor<4>,
|
||||
}
|
||||
|
||||
impl ImagePyramid {
|
||||
pub const KSIZE: usize = 7;
|
||||
|
||||
pub fn new(dev: &Device) -> Self {
|
||||
let k = Self::KSIZE;
|
||||
let g: Vec<f64> = (0..k)
|
||||
.map(|i| (-((i as f64 - (k - 1) as f64 / 2.0).powi(2)) / 2.0).exp())
|
||||
.collect();
|
||||
let sum: f64 = g.iter().sum();
|
||||
let g: Vec<f64> = g.iter().map(|v| v / sum).collect();
|
||||
let outer: Vec<f32> = g
|
||||
.iter()
|
||||
.flat_map(|a| g.iter().map(move |b| (a * b) as f32))
|
||||
.collect();
|
||||
Self {
|
||||
kernel: Tensor::from_data(TensorData::new(outer, [1, 1, k, k]), dev),
|
||||
}
|
||||
}
|
||||
|
||||
/// Doubles the size: each sample lands on the even pixel of its 2 × 2 cell, then the
|
||||
/// Gaussian (×4, for the three zeros) fills the rest.
|
||||
pub fn expand(&self, x: Tensor<4>) -> Tensor<4> {
|
||||
let z = Tensor::zeros_like(&x);
|
||||
let x = pixel_shuffle(Tensor::cat(vec![x, z.clone(), z.clone(), z], 1), 2);
|
||||
let p = Self::KSIZE / 2;
|
||||
let x = x.pad([(p, p), (p, p)], PadMode::Reflect);
|
||||
conv2d(
|
||||
x,
|
||||
self.kernel.clone() * 4.0,
|
||||
None,
|
||||
ConvOptions::new([1, 1], [0, 0], [1, 1], 1),
|
||||
)
|
||||
}
|
||||
|
||||
/// The next finer level: `x` expanded plus the Laplacian detail.
|
||||
pub fn reconstruct(&self, x: Tensor<4>, laplacian: Tensor<4>) -> Tensor<4> {
|
||||
let expanded = self.expand(x);
|
||||
let [_, _, h, w] = expanded.dims();
|
||||
// PyTorch resizes the detail to the expanded size (align_corners = True); at equal
|
||||
// sizes that is the identity.
|
||||
let laplacian = if laplacian.dims()[2..] == [h, w] {
|
||||
laplacian
|
||||
} else {
|
||||
resize(laplacian, [h, w], true)
|
||||
};
|
||||
expanded + laplacian
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,408 @@
|
||||
//! The Swin Transformer backbone (Z. Liu et al., "Swin Transformer: Hierarchical Vision
|
||||
//! Transformer using Shifted Windows", ICCV 2021), in the detection variant InSPyReNet uses:
|
||||
//! windows padded to the window size, a shift in every second block, a LayerNorm on each stage
|
||||
//! output, and the patch embedding as an extra first output.
|
||||
|
||||
use burn::module::{Module, Param};
|
||||
use burn::nn::Linear;
|
||||
use burn::nn::conv::{Conv2d, Conv2dConfig};
|
||||
use burn::tensor::activation::{gelu, softmax};
|
||||
use burn::tensor::{Device, Int, Tensor, TensorData};
|
||||
|
||||
use super::nn::{LayerNorm, linear};
|
||||
|
||||
/// Swin-B hyperparameters.
|
||||
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
|
||||
pub struct SwinConfig {
|
||||
pub embed_dim: usize,
|
||||
pub depths: [usize; 4],
|
||||
pub heads: [usize; 4],
|
||||
pub window: usize,
|
||||
pub patch: usize,
|
||||
}
|
||||
|
||||
impl SwinConfig {
|
||||
/// Swin-B, window 12 (InSPyReNet's backbone).
|
||||
pub const BASE: Self = Self {
|
||||
embed_dim: 128,
|
||||
depths: [2, 2, 18, 2],
|
||||
heads: [4, 8, 16, 32],
|
||||
window: 12,
|
||||
patch: 4,
|
||||
};
|
||||
}
|
||||
|
||||
#[derive(Module, Debug)]
|
||||
pub struct WindowAttention {
|
||||
pub qkv: Linear,
|
||||
pub proj: Linear,
|
||||
/// `[(2·window − 1)², heads]`.
|
||||
pub relative_position_bias_table: Param<Tensor<2>>,
|
||||
pub heads: usize,
|
||||
}
|
||||
|
||||
#[derive(Module, Debug)]
|
||||
pub struct Mlp {
|
||||
pub fc1: Linear,
|
||||
pub fc2: Linear,
|
||||
}
|
||||
|
||||
#[derive(Module, Debug)]
|
||||
pub struct Block {
|
||||
pub norm1: LayerNorm,
|
||||
pub attn: WindowAttention,
|
||||
pub norm2: LayerNorm,
|
||||
pub mlp: Mlp,
|
||||
}
|
||||
|
||||
#[derive(Module, Debug)]
|
||||
pub struct PatchMerging {
|
||||
pub norm: LayerNorm,
|
||||
pub reduction: Linear,
|
||||
}
|
||||
|
||||
#[derive(Module, Debug)]
|
||||
pub struct Stage {
|
||||
pub blocks: Vec<Block>,
|
||||
pub downsample: Option<PatchMerging>,
|
||||
}
|
||||
|
||||
#[derive(Module, Debug)]
|
||||
pub struct PatchEmbed {
|
||||
pub proj: Conv2d,
|
||||
pub norm: LayerNorm,
|
||||
}
|
||||
|
||||
#[derive(Module, Debug)]
|
||||
pub struct Swin {
|
||||
pub patch_embed: PatchEmbed,
|
||||
pub layers: Vec<Stage>,
|
||||
pub norm0: LayerNorm,
|
||||
pub norm1: LayerNorm,
|
||||
pub norm2: LayerNorm,
|
||||
pub norm3: LayerNorm,
|
||||
pub window: usize,
|
||||
}
|
||||
|
||||
impl Swin {
|
||||
pub fn new(cfg: SwinConfig, dev: &Device) -> Self {
|
||||
let w = cfg.window;
|
||||
let layers = (0..4)
|
||||
.map(|i| {
|
||||
let c = cfg.embed_dim << i;
|
||||
let heads = cfg.heads[i];
|
||||
let blocks = (0..cfg.depths[i])
|
||||
.map(|_| Block {
|
||||
norm1: LayerNorm::new(c, dev),
|
||||
attn: WindowAttention {
|
||||
qkv: linear(c, 3 * c, true, dev),
|
||||
proj: linear(c, c, true, dev),
|
||||
relative_position_bias_table: Param::from_tensor(Tensor::zeros(
|
||||
[(2 * w - 1) * (2 * w - 1), heads],
|
||||
dev,
|
||||
)),
|
||||
heads,
|
||||
},
|
||||
norm2: LayerNorm::new(c, dev),
|
||||
mlp: Mlp {
|
||||
fc1: linear(c, 4 * c, true, dev),
|
||||
fc2: linear(4 * c, c, true, dev),
|
||||
},
|
||||
})
|
||||
.collect();
|
||||
let downsample = (i < 3).then(|| PatchMerging {
|
||||
norm: LayerNorm::new(4 * c, dev),
|
||||
reduction: linear(4 * c, 2 * c, false, dev),
|
||||
});
|
||||
Stage { blocks, downsample }
|
||||
})
|
||||
.collect();
|
||||
let e = cfg.embed_dim;
|
||||
Self {
|
||||
patch_embed: PatchEmbed {
|
||||
proj: Conv2dConfig::new([3, e], [cfg.patch, cfg.patch])
|
||||
.with_stride([cfg.patch, cfg.patch])
|
||||
.init(dev),
|
||||
norm: LayerNorm::new(e, dev),
|
||||
},
|
||||
layers,
|
||||
norm0: LayerNorm::new(e, dev),
|
||||
norm1: LayerNorm::new(2 * e, dev),
|
||||
norm2: LayerNorm::new(4 * e, dev),
|
||||
norm3: LayerNorm::new(8 * e, dev),
|
||||
window: w,
|
||||
}
|
||||
}
|
||||
|
||||
/// `x`: `[1, 3, H, W]` with `H` and `W` multiples of 32. Returns the patch embedding and the
|
||||
/// four normalised stage outputs, each `[1, C, h, w]`. `cancelled` is asked between stages;
|
||||
/// `None` when it said yes.
|
||||
pub fn forward(&self, x: Tensor<4>, cancelled: &dyn Fn() -> bool) -> Option<[Tensor<4>; 5]> {
|
||||
let dev = x.device();
|
||||
let x = self.patch_embed.proj.forward(x); // [b, C, h, w]
|
||||
let embed = self.patch_embed.norm.forward(x.permute([0, 2, 3, 1])); // [b, h, w, C]
|
||||
let first = embed.clone().permute([0, 3, 1, 2]);
|
||||
let rel = rel_index(self.window, &dev);
|
||||
let norms = [&self.norm0, &self.norm1, &self.norm2, &self.norm3];
|
||||
let mut outs = Vec::with_capacity(4);
|
||||
let mut x = embed;
|
||||
for (stage, norm) in self.layers.iter().zip(norms) {
|
||||
if cancelled() {
|
||||
return None;
|
||||
}
|
||||
let [_, sh, sw, _] = x.dims();
|
||||
let mask = shift_mask(self.window, sh, sw, &dev);
|
||||
for (i, blk) in stage.blocks.iter().enumerate() {
|
||||
let shift = if i % 2 == 1 { self.window / 2 } else { 0 };
|
||||
x = block(
|
||||
blk,
|
||||
x,
|
||||
self.window,
|
||||
shift,
|
||||
&rel,
|
||||
(shift > 0).then_some(&mask),
|
||||
);
|
||||
}
|
||||
outs.push(norm.forward(x.clone()).permute([0, 3, 1, 2]));
|
||||
if let Some(m) = &stage.downsample {
|
||||
x = merge(m, x);
|
||||
}
|
||||
}
|
||||
let mut it = outs.into_iter();
|
||||
Some([first, it.next()?, it.next()?, it.next()?, it.next()?])
|
||||
}
|
||||
}
|
||||
|
||||
/// Relative position index of a `w × w` window, flattened to `[N·N]`.
|
||||
fn rel_index(w: usize, dev: &Device) -> Tensor<1, Int> {
|
||||
let n = w * w;
|
||||
let wi = w as i64;
|
||||
let mut idx = Vec::with_capacity(n * n);
|
||||
for i in 0..n {
|
||||
for j in 0..n {
|
||||
let (hi, ci, hj, cj) = (
|
||||
(i / w) as i64,
|
||||
(i % w) as i64,
|
||||
(j / w) as i64,
|
||||
(j % w) as i64,
|
||||
);
|
||||
idx.push((hi - hj + wi - 1) * (2 * wi - 1) + (ci - cj + wi - 1));
|
||||
}
|
||||
}
|
||||
Tensor::from_data(TensorData::new(idx, [n * n]), dev)
|
||||
}
|
||||
|
||||
fn padded(n: usize, w: usize) -> usize {
|
||||
n.div_ceil(w) * w
|
||||
}
|
||||
|
||||
/// The shifted-window attention mask of an `h × w` grid: `[nW, N, N]`, −100 between tokens of
|
||||
/// different regions.
|
||||
fn shift_mask(win: usize, h: usize, w: usize, dev: &Device) -> Tensor<3> {
|
||||
let (hp, wp, s) = (padded(h, win), padded(w, win), win / 2);
|
||||
let region = |v: usize, len: usize| {
|
||||
if v + win < len {
|
||||
0
|
||||
} else if v + s < len {
|
||||
1
|
||||
} else {
|
||||
2
|
||||
}
|
||||
};
|
||||
let label = |y: usize, x: usize| region(y, hp) * 3 + region(x, wp);
|
||||
let (nwh, nww, n) = (hp / win, wp / win, win * win);
|
||||
let mut m = Vec::with_capacity(nwh * nww * n * n);
|
||||
for wy in 0..nwh {
|
||||
for wx in 0..nww {
|
||||
let lab: Vec<usize> = (0..n)
|
||||
.map(|t| label(wy * win + t / win, wx * win + t % win))
|
||||
.collect();
|
||||
for a in &lab {
|
||||
for b in &lab {
|
||||
m.push(if a == b { 0.0f32 } else { -100.0 });
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
Tensor::from_data(TensorData::new(m, [nwh * nww, n, n]), dev)
|
||||
}
|
||||
|
||||
/// `[B, Hp, Wp, C]` → `[B·nW, win², C]`.
|
||||
fn partition(x: Tensor<4>, win: usize) -> Tensor<3> {
|
||||
let [b, h, w, c] = x.dims();
|
||||
x.reshape([b, h / win, win, w / win, win, c])
|
||||
.permute([0, 1, 3, 2, 4, 5])
|
||||
.reshape([b * (h / win) * (w / win), win * win, c])
|
||||
}
|
||||
|
||||
/// Inverse of [`partition`].
|
||||
fn reverse(x: Tensor<3>, win: usize, b: usize, h: usize, w: usize) -> Tensor<4> {
|
||||
let c = x.dims()[2];
|
||||
x.reshape([b, h / win, w / win, win, win, c])
|
||||
.permute([0, 1, 3, 2, 4, 5])
|
||||
.reshape([b, h, w, c])
|
||||
}
|
||||
|
||||
/// Zero-pads dims 1 and 2 of `[B, H, W, C]` at the bottom and right to multiples of `win`.
|
||||
fn pad(x: Tensor<4>, win: usize) -> Tensor<4> {
|
||||
let [b, h, w, c] = x.dims();
|
||||
let dev = x.device();
|
||||
let (hp, wp) = (padded(h, win), padded(w, win));
|
||||
let x = if wp > w {
|
||||
Tensor::cat(vec![x, Tensor::zeros([b, h, wp - w, c], &dev)], 2)
|
||||
} else {
|
||||
x
|
||||
};
|
||||
if hp > h {
|
||||
Tensor::cat(vec![x, Tensor::zeros([b, hp - h, wp, c], &dev)], 1)
|
||||
} else {
|
||||
x
|
||||
}
|
||||
}
|
||||
|
||||
fn block(
|
||||
blk: &Block,
|
||||
x: Tensor<4>,
|
||||
win: usize,
|
||||
shift: usize,
|
||||
rel: &Tensor<1, Int>,
|
||||
mask: Option<&Tensor<3>>,
|
||||
) -> Tensor<4> {
|
||||
let [b, h, w, c] = x.dims();
|
||||
let heads = blk.attn.heads;
|
||||
let shortcut = x.clone();
|
||||
let mut y = pad(blk.norm1.forward(x), win);
|
||||
let (hp, wp) = (padded(h, win), padded(w, win));
|
||||
let s = shift as i64;
|
||||
if shift > 0 {
|
||||
y = y.roll(&[-s, -s], &[1, 2]);
|
||||
}
|
||||
let windows = partition(y, win);
|
||||
let [bw, n, _] = windows.dims();
|
||||
let d = c / heads;
|
||||
let qkv = blk
|
||||
.attn
|
||||
.qkv
|
||||
.forward(windows)
|
||||
.reshape([bw, n, 3, heads, d])
|
||||
.permute([2, 0, 3, 1, 4]);
|
||||
let take = |i: usize| qkv.clone().narrow(0, i, 1).reshape([bw, heads, n, d]);
|
||||
let (q, k, v) = (take(0) * (d as f64).powf(-0.5), take(1), take(2));
|
||||
let mut attn = q.matmul(k.swap_dims(2, 3));
|
||||
let bias = blk
|
||||
.attn
|
||||
.relative_position_bias_table
|
||||
.val()
|
||||
.select(0, rel.clone())
|
||||
.reshape([n, n, heads])
|
||||
.permute([2, 0, 1])
|
||||
.unsqueeze_dim::<4>(0);
|
||||
attn = attn + bias;
|
||||
if let Some(m) = mask {
|
||||
let nw = m.dims()[0];
|
||||
attn = (attn.reshape([bw / nw, nw, heads, n, n])
|
||||
+ m.clone().unsqueeze_dim::<4>(1).unsqueeze_dim::<5>(0))
|
||||
.reshape([bw, heads, n, n]);
|
||||
}
|
||||
let out = softmax(attn, 3)
|
||||
.matmul(v)
|
||||
.swap_dims(1, 2)
|
||||
.reshape([bw, n, c]);
|
||||
let mut y = reverse(blk.attn.proj.forward(out), win, b, hp, wp);
|
||||
if shift > 0 {
|
||||
y = y.roll(&[s, s], &[1, 2]);
|
||||
}
|
||||
let x = shortcut + y.narrow(1, 0, h).narrow(2, 0, w);
|
||||
x.clone()
|
||||
+ blk
|
||||
.mlp
|
||||
.fc2
|
||||
.forward(gelu(blk.mlp.fc1.forward(blk.norm2.forward(x))))
|
||||
}
|
||||
|
||||
/// Patch merging: the 2 × 2 neighbours concatenated in PyTorch's order (`x[0::2, 0::2]`,
|
||||
/// `x[1::2, 0::2]`, `x[0::2, 1::2]`, `x[1::2, 1::2]`), normalised and projected to 2C. Odd sizes
|
||||
/// are zero-padded first.
|
||||
fn merge(m: &PatchMerging, x: Tensor<4>) -> Tensor<4> {
|
||||
let [b, h, w, c] = x.dims();
|
||||
let dev = x.device();
|
||||
let x = if w % 2 == 1 {
|
||||
Tensor::cat(vec![x, Tensor::zeros([b, h, 1, c], &dev)], 2)
|
||||
} else {
|
||||
x
|
||||
};
|
||||
let x = if h % 2 == 1 {
|
||||
Tensor::cat(vec![x, Tensor::zeros([b, 1, w + w % 2, c], &dev)], 1)
|
||||
} else {
|
||||
x
|
||||
};
|
||||
let [b, h, w, c] = x.dims();
|
||||
let x = x
|
||||
.reshape([b, h / 2, 2, w / 2, 2, c])
|
||||
.permute([0, 1, 3, 4, 2, 5])
|
||||
.reshape([b, h / 2, w / 2, 4 * c]);
|
||||
m.reduction.forward(m.norm.forward(x))
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
fn max_diff<const D: usize>(a: Tensor<D>, b: Tensor<D>) -> f32 {
|
||||
(a - b).abs().max().into_scalar::<f32>()
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn windows_and_rolls_are_exact_inverses() {
|
||||
let dev = Device::flex();
|
||||
let x = Tensor::<4>::random([1, 36, 48, 8], burn::tensor::Distribution::Default, &dev);
|
||||
assert_eq!(
|
||||
max_diff(x.clone(), reverse(partition(x.clone(), 12), 12, 1, 36, 48)),
|
||||
0.0
|
||||
);
|
||||
assert_eq!(
|
||||
max_diff(
|
||||
x.clone(),
|
||||
x.clone()
|
||||
.roll(&[-6i64, -6], &[1, 2])
|
||||
.roll(&[6i64, 6], &[1, 2])
|
||||
),
|
||||
0.0
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn merge_takes_neighbours_in_pytorch_order() {
|
||||
let dev = Device::flex();
|
||||
// A 2 × 2 grid of one channel holding 0, 1 (top row) and 2, 3 (bottom row).
|
||||
let x = Tensor::<4>::from_data(
|
||||
TensorData::new(vec![0.0f32, 1.0, 2.0, 3.0], [1, 2, 2, 1]),
|
||||
&dev,
|
||||
);
|
||||
let cat = x
|
||||
.reshape([1, 1, 2, 1, 2, 1])
|
||||
.permute([0, 1, 3, 4, 2, 5])
|
||||
.reshape([1, 1, 1, 4]);
|
||||
// x[0::2,0::2] = 0, x[1::2,0::2] = 2, x[0::2,1::2] = 1, x[1::2,1::2] = 3.
|
||||
assert_eq!(
|
||||
cat.into_data().try_to_vec::<f32>().unwrap(),
|
||||
[0.0, 2.0, 1.0, 3.0]
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn shift_mask_regions_match_the_reference_slices() {
|
||||
// 24 × 24 with window 12, shift 6: the last window row and column hold three regions.
|
||||
let m = shift_mask(12, 24, 24, &Device::flex())
|
||||
.into_data()
|
||||
.try_to_vec::<f32>()
|
||||
.unwrap();
|
||||
let n = 144;
|
||||
// Window 0 is one region: no masking.
|
||||
assert!(m[..n * n].iter().all(|v| *v == 0.0));
|
||||
// Window 3 (bottom right) mixes four regions: its first and last token differ.
|
||||
let w3 = &m[3 * n * n..4 * n * n];
|
||||
assert_eq!(w3[n - 1], -100.0);
|
||||
assert_eq!(w3[0], 0.0);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,160 @@
|
||||
//! Parity with the PyTorch reference implementation of InSPyReNet.
|
||||
//!
|
||||
//! Opt-in: it needs the published checkpoint and a reference file (input, every intermediate
|
||||
//! stage and the prediction, computed once by PyTorch on a synthetic 384² input with
|
||||
//! `tools/burn/inspyrenet_reference.py`):
|
||||
//!
|
||||
//! ```text
|
||||
//! CRAFT_TEST_INSPYRENET_WEIGHTS=…/inspyrenet.safetensors \
|
||||
//! CRAFT_TEST_INSPYRENET_REFERENCE=…/inspyrenet-ref-384.safetensors \
|
||||
//! cargo test --release -p craft-ai-engine-burn -j 8 --lib parity -- --ignored --nocapture
|
||||
//! ```
|
||||
//!
|
||||
//! Each stage is compared on its own, so a mismatch points at the module that causes it.
|
||||
|
||||
use std::collections::HashMap;
|
||||
use std::path::PathBuf;
|
||||
|
||||
use burn::backend::wgpu::WgpuDevice;
|
||||
use burn::prelude::*;
|
||||
use burn::tensor::TensorData;
|
||||
|
||||
use super::net::InSPyReNet;
|
||||
use super::{on_big_stack, weights};
|
||||
|
||||
struct Ref {
|
||||
shape: Vec<usize>,
|
||||
data: Vec<f32>,
|
||||
}
|
||||
|
||||
/// A float32 safetensors file, read without a dependency.
|
||||
fn read_reference(path: &str) -> HashMap<String, Ref> {
|
||||
let bytes = std::fs::read(path).unwrap_or_else(|e| panic!("reading {path}: {e}"));
|
||||
let n = u64::from_le_bytes(bytes[..8].try_into().unwrap()) as usize;
|
||||
let header: serde_json::Value = serde_json::from_slice(&bytes[8..8 + n]).unwrap();
|
||||
let data = &bytes[8 + n..];
|
||||
let mut out = HashMap::new();
|
||||
for (k, v) in header.as_object().unwrap() {
|
||||
if k == "__metadata__" {
|
||||
continue;
|
||||
}
|
||||
assert_eq!(v["dtype"], "F32", "{k}");
|
||||
let shape = v["shape"]
|
||||
.as_array()
|
||||
.unwrap()
|
||||
.iter()
|
||||
.map(|d| d.as_u64().unwrap() as usize)
|
||||
.collect();
|
||||
let a = v["data_offsets"][0].as_u64().unwrap() as usize;
|
||||
let b = v["data_offsets"][1].as_u64().unwrap() as usize;
|
||||
let floats = data[a..b]
|
||||
.as_chunks::<4>()
|
||||
.0
|
||||
.iter()
|
||||
.map(|c| f32::from_le_bytes(*c))
|
||||
.collect();
|
||||
out.insert(
|
||||
k.clone(),
|
||||
Ref {
|
||||
shape,
|
||||
data: floats,
|
||||
},
|
||||
);
|
||||
}
|
||||
out
|
||||
}
|
||||
|
||||
/// `(max |Δ|, max |reference|)`, printed with the mean difference.
|
||||
fn compare(name: &str, got: Tensor<4>, want: &Ref) -> (f32, f32) {
|
||||
assert_eq!(got.dims().to_vec(), want.shape, "{name}: shape");
|
||||
let got = got.into_data().try_to_vec::<f32>().unwrap();
|
||||
let (mut max, mut sum, mut scale) = (0.0f32, 0.0f64, 0.0f32);
|
||||
for (a, b) in got.iter().zip(&want.data) {
|
||||
let d = (a - b).abs();
|
||||
max = max.max(d);
|
||||
sum += d as f64;
|
||||
scale = scale.max(b.abs());
|
||||
}
|
||||
let mean = sum / got.len().max(1) as f64;
|
||||
println!("{name:<12} max |Δ| {max:<12.3e} mean |Δ| {mean:<12.3e} max |ref| {scale:.3}");
|
||||
(max, scale)
|
||||
}
|
||||
|
||||
fn run(device: Device) {
|
||||
let (Some(weights_path), Ok(reference)) = (
|
||||
std::env::var_os("CRAFT_TEST_INSPYRENET_WEIGHTS").map(PathBuf::from),
|
||||
std::env::var("CRAFT_TEST_INSPYRENET_REFERENCE"),
|
||||
) else {
|
||||
panic!("set CRAFT_TEST_INSPYRENET_WEIGHTS and CRAFT_TEST_INSPYRENET_REFERENCE");
|
||||
};
|
||||
let r = read_reference(&reference);
|
||||
on_big_stack(move || {
|
||||
let net = weights::fill(InSPyReNet::new(&device), &weights_path)?;
|
||||
let input = &r["input"];
|
||||
let x = Tensor::<4>::from_data(
|
||||
TensorData::new(input.data.clone(), input.shape.clone()),
|
||||
&device,
|
||||
);
|
||||
let parts = net.forward_parts(x.clone(), &|| false).unwrap();
|
||||
let pred = net.forward(x, &|| false).unwrap();
|
||||
|
||||
let mut worst_rel = 0.0f32;
|
||||
let stages = parts
|
||||
.backbone
|
||||
.into_iter()
|
||||
.enumerate()
|
||||
.map(|(i, t)| (format!("backbone.{i}"), t))
|
||||
.chain(
|
||||
parts
|
||||
.laplacian
|
||||
.into_iter()
|
||||
.enumerate()
|
||||
.map(|(i, t)| (format!("laplacian.{i}"), t)),
|
||||
)
|
||||
.chain(
|
||||
parts
|
||||
.saliency
|
||||
.into_iter()
|
||||
.enumerate()
|
||||
.map(|(i, t)| (format!("saliency.{i}"), t)),
|
||||
);
|
||||
for (name, t) in stages {
|
||||
let (max, scale) = compare(&name, t, &r[&name]);
|
||||
worst_rel = worst_rel.max(max / scale.max(1e-6));
|
||||
}
|
||||
let want = &r["pred"];
|
||||
let got = pred.into_data().try_to_vec::<f32>().unwrap();
|
||||
let (mut max, mut flips) = (0.0f32, 0usize);
|
||||
for (a, b) in got.iter().zip(&want.data) {
|
||||
max = max.max((a - b).abs());
|
||||
flips += usize::from((*a >= 0.5) != (*b >= 0.5));
|
||||
}
|
||||
println!(
|
||||
"pred max |Δ| {max:.3e}, {flips} of {} pixels change side at 0.5",
|
||||
got.len()
|
||||
);
|
||||
println!("worst relative difference over the stages: {worst_rel:.3e}");
|
||||
assert!(
|
||||
worst_rel < 5e-3,
|
||||
"a stage drifts from PyTorch by {worst_rel:.3e} of its range"
|
||||
);
|
||||
assert!(max < 1e-2, "the prediction differs by up to {max}");
|
||||
assert!(flips * 10_000 <= got.len(), "{flips} pixels change side");
|
||||
Ok(())
|
||||
})
|
||||
.unwrap();
|
||||
}
|
||||
|
||||
#[test]
|
||||
#[ignore = "needs CRAFT_TEST_INSPYRENET_WEIGHTS and CRAFT_TEST_INSPYRENET_REFERENCE"]
|
||||
fn matches_pytorch_on_the_processor() {
|
||||
#[allow(deprecated)]
|
||||
let cpu = burn::backend::ndarray::NdArrayDevice::Cpu;
|
||||
run(cpu.into());
|
||||
}
|
||||
|
||||
#[test]
|
||||
#[ignore = "needs CRAFT_TEST_INSPYRENET_WEIGHTS, CRAFT_TEST_INSPYRENET_REFERENCE and a graphics card"]
|
||||
fn matches_pytorch_on_a_graphics_card() {
|
||||
run(WgpuDevice::default().into());
|
||||
}
|
||||
@@ -0,0 +1,259 @@
|
||||
//! The colour guided filter (K. He, J. Sun and X. Tang, "Guided Image Filtering", PAMI 2013) that
|
||||
//! brings fine detail back into an enlarged matte, run only near the matte's edge.
|
||||
//!
|
||||
//! The filter only changes pixels within `2r` of a non-constant part of the matte, so it runs tile
|
||||
//! by tile on tiles near the edge only and gives the same result as filtering the whole picture:
|
||||
//! memory follows the edge length, not the area.
|
||||
|
||||
use craft_ai_core::AiError;
|
||||
use craft_ai_engine::RunContext;
|
||||
|
||||
use super::Picture;
|
||||
|
||||
/// Output tile size. At least `2r`, so a tile's neighbours cover its halo.
|
||||
const TILE: usize = 256;
|
||||
|
||||
/// `work` on every item, spread over the processors with scoped threads, answers in item order. A
|
||||
/// thread that cannot start or that panics leaves `None` for its items.
|
||||
fn in_parallel<T: Sync, R: Send>(
|
||||
items: &[T],
|
||||
work: &(dyn Fn(&T) -> Option<R> + Sync),
|
||||
) -> Vec<Option<R>> {
|
||||
let threads = std::thread::available_parallelism()
|
||||
.map_or(1, std::num::NonZeroUsize::get)
|
||||
.clamp(1, 32);
|
||||
let per = items.len().div_ceil(threads).max(1);
|
||||
std::thread::scope(|scope| {
|
||||
let handles: Vec<_> = items
|
||||
.chunks(per)
|
||||
.map(|chunk| {
|
||||
let run = move || chunk.iter().map(work).collect::<Vec<_>>();
|
||||
(
|
||||
chunk.len(),
|
||||
std::thread::Builder::new().spawn_scoped(scope, run),
|
||||
)
|
||||
})
|
||||
.collect();
|
||||
let mut out = Vec::with_capacity(items.len());
|
||||
for (len, handle) in handles {
|
||||
match handle.map(std::thread::ScopedJoinHandle::join) {
|
||||
Ok(Ok(part)) => out.extend(part),
|
||||
_ => out.extend((0..len).map(|_| None)),
|
||||
}
|
||||
}
|
||||
out
|
||||
})
|
||||
}
|
||||
|
||||
/// Mean of `src` (`w × h`) over the `(2r + 1)²` window around each pixel, clipped at the edges.
|
||||
pub(super) fn box_mean(src: &[f32], w: usize, h: usize, r: usize) -> Vec<f32> {
|
||||
let mut tmp = vec![0.0f32; w * h];
|
||||
let mut pre = vec![0.0f64; w + 1];
|
||||
for (row, out) in src
|
||||
.chunks_exact(w.max(1))
|
||||
.zip(tmp.chunks_exact_mut(w.max(1)))
|
||||
{
|
||||
for (x, v) in row.iter().enumerate() {
|
||||
pre[x + 1] = pre[x] + *v as f64;
|
||||
}
|
||||
for (x, o) in out.iter_mut().enumerate() {
|
||||
let (lo, hi) = (x.saturating_sub(r), (x + r + 1).min(w));
|
||||
*o = ((pre[hi] - pre[lo]) / (hi - lo) as f64) as f32;
|
||||
}
|
||||
}
|
||||
let mut col = vec![0.0f64; (h + 1) * w];
|
||||
for y in 0..h {
|
||||
for x in 0..w {
|
||||
col[(y + 1) * w + x] = col[y * w + x] + tmp[y * w + x] as f64;
|
||||
}
|
||||
}
|
||||
let mut out = vec![0.0f32; w * h];
|
||||
for y in 0..h {
|
||||
let (lo, hi) = (y.saturating_sub(r), (y + r + 1).min(h));
|
||||
let n = (hi - lo) as f64;
|
||||
for x in 0..w {
|
||||
out[y * w + x] = ((col[hi * w + x] - col[lo * w + x]) / n) as f32;
|
||||
}
|
||||
}
|
||||
out
|
||||
}
|
||||
|
||||
/// Guided filter with a colour guide (He et al., Eq. 19–21: `a_k = (Σ_k + εU)⁻¹ cov_k(I, p)`).
|
||||
/// `guide` and `p` are both `w × h`.
|
||||
pub(super) fn filter(
|
||||
guide: &[[f32; 3]],
|
||||
p: &[f32],
|
||||
w: usize,
|
||||
h: usize,
|
||||
r: usize,
|
||||
eps: f32,
|
||||
) -> Vec<f32> {
|
||||
let n = w * h;
|
||||
let ch = |c: usize| -> Vec<f32> { guide.iter().map(|q| q[c]).collect() };
|
||||
let chans = [ch(0), ch(1), ch(2)];
|
||||
let mean = |v: &[f32]| box_mean(v, w, h, r);
|
||||
let product =
|
||||
|a: &[f32], b: &[f32]| -> Vec<f32> { a.iter().zip(b).map(|(x, y)| x * y).collect() };
|
||||
let m: Vec<Vec<f32>> = chans.iter().map(|c| mean(c)).collect();
|
||||
let mp = mean(p);
|
||||
let mip: Vec<Vec<f32>> = chans.iter().map(|c| mean(&product(c, p))).collect();
|
||||
let pairs = [(0, 0), (0, 1), (0, 2), (1, 1), (1, 2), (2, 2)];
|
||||
let mii: Vec<Vec<f32>> = pairs
|
||||
.iter()
|
||||
.map(|&(a, b)| mean(&product(&chans[a], &chans[b])))
|
||||
.collect();
|
||||
let mut av = [vec![0.0f32; n], vec![0.0f32; n], vec![0.0f32; n]];
|
||||
let mut bv = vec![0.0f32; n];
|
||||
let e = eps as f64;
|
||||
for i in 0..n {
|
||||
let mu = [m[0][i] as f64, m[1][i] as f64, m[2][i] as f64];
|
||||
let s = |k: usize, a: usize, b: usize| mii[k][i] as f64 - mu[a] * mu[b];
|
||||
let (rr, rg, rb) = (s(0, 0, 0) + e, s(1, 0, 1), s(2, 0, 2));
|
||||
let (gg, gb, bb) = (s(3, 1, 1) + e, s(4, 1, 2), s(5, 2, 2) + e);
|
||||
let cov = [
|
||||
mip[0][i] as f64 - mu[0] * mp[i] as f64,
|
||||
mip[1][i] as f64 - mu[1] * mp[i] as f64,
|
||||
mip[2][i] as f64 - mu[2] * mp[i] as f64,
|
||||
];
|
||||
// Inverse of the symmetric 3×3 matrix.
|
||||
let i00 = gg * bb - gb * gb;
|
||||
let i01 = gb * rb - rg * bb;
|
||||
let i02 = rg * gb - gg * rb;
|
||||
let i11 = rr * bb - rb * rb;
|
||||
let i12 = rb * rg - rr * gb;
|
||||
let i22 = rr * gg - rg * rg;
|
||||
let det = rr * i00 + rg * i01 + rb * i02;
|
||||
let (a0, a1, a2) = if det.abs() < 1e-30 {
|
||||
(0.0, 0.0, 0.0)
|
||||
} else {
|
||||
(
|
||||
(i00 * cov[0] + i01 * cov[1] + i02 * cov[2]) / det,
|
||||
(i01 * cov[0] + i11 * cov[1] + i12 * cov[2]) / det,
|
||||
(i02 * cov[0] + i12 * cov[1] + i22 * cov[2]) / det,
|
||||
)
|
||||
};
|
||||
av[0][i] = a0 as f32;
|
||||
av[1][i] = a1 as f32;
|
||||
av[2][i] = a2 as f32;
|
||||
bv[i] = (mp[i] as f64 - a0 * mu[0] - a1 * mu[1] - a2 * mu[2]) as f32;
|
||||
}
|
||||
let ma: Vec<Vec<f32>> = av.iter().map(|a| mean(a)).collect();
|
||||
let mb = mean(&bv);
|
||||
(0..n)
|
||||
.map(|i| ma[0][i] * chans[0][i] + ma[1][i] * chans[1][i] + ma[2][i] * chans[2][i] + mb[i])
|
||||
.collect()
|
||||
}
|
||||
|
||||
/// A rectangle of the picture, `x0..x1` × `y0..y1`.
|
||||
#[derive(Clone, Copy)]
|
||||
struct Area {
|
||||
x0: usize,
|
||||
y0: usize,
|
||||
x1: usize,
|
||||
y1: usize,
|
||||
}
|
||||
|
||||
impl Area {
|
||||
fn width(&self) -> usize {
|
||||
self.x1 - self.x0
|
||||
}
|
||||
}
|
||||
|
||||
/// `mask = max(guided_filter(picture, mask), mask)`, computed only on tiles within `2r` of a
|
||||
/// non-constant part of the mask (elsewhere the filter returns its input exactly).
|
||||
pub(super) fn refine(
|
||||
picture: Picture<'_>,
|
||||
mask: &mut [u8],
|
||||
r: usize,
|
||||
eps: f32,
|
||||
ctx: &RunContext,
|
||||
) -> Result<(), AiError> {
|
||||
let (w, h) = (picture.width, picture.height);
|
||||
let halo = 2 * r;
|
||||
let (tw, th) = (w.div_ceil(TILE), h.div_ceil(TILE));
|
||||
let tile = |tx: usize, ty: usize| Area {
|
||||
x0: tx * TILE,
|
||||
y0: ty * TILE,
|
||||
x1: ((tx + 1) * TILE).min(w),
|
||||
y1: ((ty + 1) * TILE).min(h),
|
||||
};
|
||||
// Per-tile value range; a tile needs filtering when the range over its neighbourhood (which
|
||||
// covers the 2r halo, as 2r ≤ TILE) is not a single value.
|
||||
let mut range = vec![(u8::MAX, u8::MIN); tw * th];
|
||||
for ty in 0..th {
|
||||
for tx in 0..tw {
|
||||
let t = tile(tx, ty);
|
||||
let mut rg = (u8::MAX, u8::MIN);
|
||||
for y in t.y0..t.y1 {
|
||||
for v in mask.get(y * w + t.x0..y * w + t.x1).unwrap_or(&[]) {
|
||||
rg = (rg.0.min(*v), rg.1.max(*v));
|
||||
}
|
||||
}
|
||||
range[ty * tw + tx] = rg;
|
||||
}
|
||||
}
|
||||
let mut active = Vec::new();
|
||||
for ty in 0..th {
|
||||
for tx in 0..tw {
|
||||
let mut rg = (u8::MAX, u8::MIN);
|
||||
for ny in ty.saturating_sub(1)..=(ty + 1).min(th - 1) {
|
||||
for nx in tx.saturating_sub(1)..=(tx + 1).min(tw - 1) {
|
||||
let q = range[ny * tw + nx];
|
||||
rg = (rg.0.min(q.0), rg.1.max(q.1));
|
||||
}
|
||||
}
|
||||
if rg.0 != rg.1 {
|
||||
active.push(tile(tx, ty));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Each tile reads the mask as it was before any tile is written, so tiles are independent and
|
||||
// run on every processor at once.
|
||||
let src: &[u8] = mask;
|
||||
let one = |t: &Area| -> Option<Vec<u8>> {
|
||||
if ctx.is_cancelled() {
|
||||
return None;
|
||||
}
|
||||
let rd = Area {
|
||||
x0: t.x0.saturating_sub(halo),
|
||||
y0: t.y0.saturating_sub(halo),
|
||||
x1: (t.x1 + halo).min(w),
|
||||
y1: (t.y1 + halo).min(h),
|
||||
};
|
||||
let (rw, rh) = (rd.width(), rd.y1 - rd.y0);
|
||||
let guide = picture.rgb(rd.x0, rd.y0, rd.x1, rd.y1);
|
||||
let mut pm = Vec::with_capacity(rw * rh);
|
||||
for y in rd.y0..rd.y1 {
|
||||
pm.extend(
|
||||
src.get(y * w + rd.x0..y * w + rd.x1)?
|
||||
.iter()
|
||||
.map(|v| *v as f32 / 255.0),
|
||||
);
|
||||
}
|
||||
let g = filter(&guide, &pm, rw, rh, r, eps);
|
||||
let mut out = Vec::with_capacity(t.width() * (t.y1 - t.y0));
|
||||
for y in t.y0..t.y1 {
|
||||
let at = (y - rd.y0) * rw + (t.x0 - rd.x0);
|
||||
let gs = g.get(at..at + t.width())?;
|
||||
let ps = pm.get(at..at + t.width())?;
|
||||
out.extend(
|
||||
gs.iter()
|
||||
.zip(ps)
|
||||
.map(|(gv, pv)| (gv.clamp(0.0, 1.0).max(*pv) * 255.0).round() as u8),
|
||||
);
|
||||
}
|
||||
Some(out)
|
||||
};
|
||||
let results = in_parallel(&active, &one);
|
||||
ctx.check_cancelled()?;
|
||||
for (t, out) in active.iter().zip(results) {
|
||||
let Some(out) = out else { continue };
|
||||
for (y, row) in (t.y0..t.y1).zip(out.chunks_exact(t.width().max(1))) {
|
||||
if let Some(dst) = mask.get_mut(y * w + t.x0..y * w + t.x1) {
|
||||
dst.copy_from_slice(row);
|
||||
}
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
@@ -0,0 +1,199 @@
|
||||
//! A matting network trained near one megapixel, made usable on a picture of any size (the
|
||||
//! approach of the Gyre background remover):
|
||||
//!
|
||||
//! 1. Up to [`MAX_RES`] pixels on the long side the picture runs at full resolution,
|
||||
//! reflect-padded to a multiple of [`ALIGN`] (centred).
|
||||
//! 2. Larger pictures are reflect-padded to a square (at the left and top) and reduced with
|
||||
//! antialiased Lanczos-3 so the square is [`MAX_RES`]²; the matte is enlarged back with
|
||||
//! Lanczos-3 and cropped. The square is never built: source rows stream through in strips, and
|
||||
//! the enlargement runs row by row.
|
||||
//! 3. When that scale is below [`GUIDED_BELOW`], the matte is refined with a colour guided filter
|
||||
//! (K. He, J. Sun and X. Tang, "Guided Image Filtering", PAMI 2013) guided by the
|
||||
//! full-resolution picture, keeping `max(guided, enlarged)`: fine detail (hair, fur) the small
|
||||
//! input lost comes back without eroding the body of the matte.
|
||||
//! 4. A shrink, `(m − s) / (1 − s)`, tightens the soft fringe.
|
||||
|
||||
mod guided;
|
||||
mod resize;
|
||||
|
||||
use craft_ai_core::AiError;
|
||||
use craft_ai_engine::RunContext;
|
||||
|
||||
pub use resize::Input;
|
||||
|
||||
/// Longest side the network sees (pixels).
|
||||
pub const MAX_RES: usize = 1280;
|
||||
/// Input sides are padded to a multiple of this (the network's coarsest stride).
|
||||
pub const ALIGN: usize = 32;
|
||||
/// Below this working scale the matte is refined at full resolution by the guided filter.
|
||||
const GUIDED_BELOW: f64 = 0.5;
|
||||
/// Guided-filter window radius (full-resolution pixels).
|
||||
const GUIDED_RADIUS: usize = 32;
|
||||
/// Guided-filter regulariser. Gyre used 1e-8 with float64 statistics; the box means here are
|
||||
/// `f32`, so a slightly larger value keeps flat colour areas well-conditioned while still
|
||||
/// following edges a fraction of an 8-bit level apart.
|
||||
const GUIDED_EPS: f32 = 1e-6;
|
||||
/// The fringe shrink (Gyre's default).
|
||||
pub const SHRINK: f32 = 0.2;
|
||||
|
||||
/// An 8-bit RGB picture, row by row, as `craft-ai-engine-rten` decodes it.
|
||||
#[derive(Clone, Copy)]
|
||||
pub struct Picture<'a> {
|
||||
pub width: usize,
|
||||
pub height: usize,
|
||||
pub bytes: &'a [u8],
|
||||
}
|
||||
|
||||
impl Picture<'_> {
|
||||
/// Rows `y0..y1`, columns `x0..x1` as RGB in `0..=1`. Outside the picture reads black.
|
||||
fn rgb(&self, x0: usize, y0: usize, x1: usize, y1: usize) -> Vec<[f32; 3]> {
|
||||
let mut out = Vec::with_capacity(x1.saturating_sub(x0) * y1.saturating_sub(y0));
|
||||
for y in y0..y1 {
|
||||
for x in x0..x1 {
|
||||
let at = (y * self.width + x) * 3;
|
||||
out.push(match self.bytes.get(at..at + 3) {
|
||||
Some(p) => [
|
||||
p[0] as f32 / 255.0,
|
||||
p[1] as f32 / 255.0,
|
||||
p[2] as f32 / 255.0,
|
||||
],
|
||||
None => [0.0; 3],
|
||||
});
|
||||
}
|
||||
}
|
||||
out
|
||||
}
|
||||
}
|
||||
|
||||
/// How a picture maps to the network input.
|
||||
#[derive(Clone, Copy, Debug, PartialEq)]
|
||||
pub struct Plan {
|
||||
/// Network pixels per picture pixel (`1.0` = full resolution).
|
||||
pub scale: f64,
|
||||
/// The padded network input.
|
||||
pub in_w: usize,
|
||||
pub in_h: usize,
|
||||
/// Padding before the picture (left, top), in picture pixels.
|
||||
pub left: usize,
|
||||
pub top: usize,
|
||||
}
|
||||
|
||||
impl Plan {
|
||||
/// Whether the guided filter refines the enlarged matte.
|
||||
fn guided(&self) -> bool {
|
||||
self.scale < GUIDED_BELOW
|
||||
}
|
||||
}
|
||||
|
||||
/// The input layout for a `w × h` picture (`None` if it is empty). `max_res` is rounded down to a
|
||||
/// multiple of [`ALIGN`], and at least [`ALIGN`].
|
||||
pub fn plan(w: usize, h: usize, max_res: usize) -> Option<Plan> {
|
||||
if w == 0 || h == 0 {
|
||||
return None;
|
||||
}
|
||||
let max_res = (max_res.max(ALIGN) / ALIGN) * ALIGN;
|
||||
let long = w.max(h);
|
||||
if long > max_res {
|
||||
Some(Plan {
|
||||
scale: max_res as f64 / long as f64,
|
||||
in_w: max_res,
|
||||
in_h: max_res,
|
||||
left: long - w,
|
||||
top: long - h,
|
||||
})
|
||||
} else {
|
||||
let (in_w, in_h) = (w.div_ceil(ALIGN) * ALIGN, h.div_ceil(ALIGN) * ALIGN);
|
||||
Some(Plan {
|
||||
scale: 1.0,
|
||||
in_w,
|
||||
in_h,
|
||||
left: (in_w - w) / 2,
|
||||
top: (in_h - h) / 2,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
/// 256-entry table of the fringe shrink `(m − s) / (1 − s)`.
|
||||
fn shrink_lut(shrink: f32) -> [u8; 256] {
|
||||
let s = if shrink.is_finite() {
|
||||
shrink.clamp(0.0, 0.99)
|
||||
} else {
|
||||
0.0
|
||||
};
|
||||
std::array::from_fn(|v| {
|
||||
((((v as f32 / 255.0) - s) / (1.0 - s)).clamp(0.0, 1.0) * 255.0).round() as u8
|
||||
})
|
||||
}
|
||||
|
||||
/// The subject's coverage of `picture`, one 8-bit value per pixel, row by row.
|
||||
///
|
||||
/// `predict` runs the network on a padded input whose sides are multiples of [`ALIGN`] and
|
||||
/// returns one foreground probability per input pixel. Progress is reported from 10 to 90;
|
||||
/// Cancel is checked between stages.
|
||||
pub fn matte(
|
||||
picture: Picture<'_>,
|
||||
max_res: usize,
|
||||
shrink: f32,
|
||||
predict: &mut dyn FnMut(&Input) -> Result<Vec<f32>, AiError>,
|
||||
ctx: &RunContext,
|
||||
) -> Result<Vec<u8>, AiError> {
|
||||
let (w, h) = (picture.width, picture.height);
|
||||
let n = w
|
||||
.checked_mul(h)
|
||||
.filter(|n| picture.bytes.len() == n * 3)
|
||||
.ok_or_else(|| AiError::invalid("The picture could not be read. Use a PNG or a JPEG."))?;
|
||||
let Some(p) = plan(w, h, max_res) else {
|
||||
return Ok(Vec::new());
|
||||
};
|
||||
|
||||
ctx.check_cancelled()?;
|
||||
let input = if p.scale >= 1.0 {
|
||||
resize::pad_input(picture, &p)
|
||||
} else {
|
||||
resize::reduce_input(picture, &p, ctx)?
|
||||
};
|
||||
ctx.progress(10);
|
||||
|
||||
let mut pred = predict(&input)?;
|
||||
if pred.len() != p.in_w * p.in_h {
|
||||
return Err(AiError::new(
|
||||
craft_ai_core::ErrorKind::Internal,
|
||||
"The background remover returned a matte of the wrong size.",
|
||||
));
|
||||
}
|
||||
// NaN is background; out-of-range values (infinities included) clamp.
|
||||
for v in &mut pred {
|
||||
*v = if v.is_nan() { 0.0 } else { v.clamp(0.0, 1.0) };
|
||||
}
|
||||
ctx.check_cancelled()?;
|
||||
ctx.progress(75);
|
||||
|
||||
let mut mask = vec![0u8; n];
|
||||
if p.scale >= 1.0 {
|
||||
for (y, row) in mask.chunks_exact_mut(w).enumerate() {
|
||||
let at = (y + p.top) * p.in_w + p.left;
|
||||
for (o, v) in row.iter_mut().zip(pred.get(at..at + w).unwrap_or(&[])) {
|
||||
*o = (v * 255.0).round() as u8;
|
||||
}
|
||||
}
|
||||
} else {
|
||||
resize::expand_matte(&pred, &p, w, &mut mask);
|
||||
}
|
||||
ctx.check_cancelled()?;
|
||||
|
||||
if p.guided() {
|
||||
guided::refine(picture, &mut mask, GUIDED_RADIUS, GUIDED_EPS, ctx)?;
|
||||
}
|
||||
ctx.progress(90);
|
||||
|
||||
if shrink > 0.0 {
|
||||
let lut = shrink_lut(shrink);
|
||||
for m in &mut mask {
|
||||
*m = lut[*m as usize];
|
||||
}
|
||||
}
|
||||
Ok(mask)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests;
|
||||
@@ -0,0 +1,182 @@
|
||||
//! Getting the picture into the network's input and the matte back out: reflect padding and
|
||||
//! antialiased Lanczos-3, as Gyre's resize (ResizeRight) does.
|
||||
|
||||
use craft_ai_core::AiError;
|
||||
use craft_ai_engine::RunContext;
|
||||
|
||||
use super::{Picture, Plan};
|
||||
|
||||
/// Source rows read at once while reducing a large picture.
|
||||
const STRIP: usize = 64;
|
||||
|
||||
/// The network input: RGB in `0..=1`, row by row.
|
||||
pub struct Input {
|
||||
pub width: usize,
|
||||
pub height: usize,
|
||||
pub px: Vec<[f32; 3]>,
|
||||
}
|
||||
|
||||
/// Mirror index `i` into `0..n` without repeating the edge sample (PyTorch's `reflect`), for any
|
||||
/// offset: padding wider than the picture keeps bouncing.
|
||||
pub(super) fn reflect(i: i64, n: usize) -> usize {
|
||||
if n <= 1 {
|
||||
return 0;
|
||||
}
|
||||
let period = 2 * (n as i64 - 1);
|
||||
let m = i.rem_euclid(period);
|
||||
(if m < n as i64 { m } else { period - m }) as usize
|
||||
}
|
||||
|
||||
/// The Lanczos-3 kernel.
|
||||
fn lanczos(x: f64) -> f64 {
|
||||
let x = x.abs();
|
||||
if x < 1e-9 {
|
||||
1.0
|
||||
} else if x < 3.0 {
|
||||
let px = std::f64::consts::PI * x;
|
||||
3.0 * px.sin() * (px / 3.0).sin() / (px * px)
|
||||
} else {
|
||||
0.0
|
||||
}
|
||||
}
|
||||
|
||||
/// Lanczos-3 taps of output sample `o` on an axis scaled by `scale` (output / input): the first
|
||||
/// source index and the normalised weights. Pixel centres align, and the kernel widens by
|
||||
/// 1/scale when reducing (antialiasing).
|
||||
pub(super) fn lanczos_taps(o: usize, scale: f64) -> (i64, Vec<f32>) {
|
||||
let stretch = if scale < 1.0 { 1.0 / scale } else { 1.0 };
|
||||
let u = (o as f64 + 0.5) / scale - 0.5;
|
||||
let lo = (u - 3.0 * stretch).ceil() as i64;
|
||||
let hi = (u + 3.0 * stretch).floor() as i64;
|
||||
let w: Vec<f64> = (lo..=hi)
|
||||
.map(|i| lanczos((i as f64 - u) / stretch))
|
||||
.collect();
|
||||
let sum: f64 = w.iter().sum();
|
||||
let k = if sum.abs() > 1e-12 { 1.0 / sum } else { 1.0 };
|
||||
(lo, w.into_iter().map(|v| (v * k) as f32).collect())
|
||||
}
|
||||
|
||||
/// `picture` reflect-padded into the network input (the full-resolution path).
|
||||
pub(super) fn pad_input(picture: Picture<'_>, p: &Plan) -> Input {
|
||||
let (w, h) = (picture.width, picture.height);
|
||||
let all = picture.rgb(0, 0, w, h);
|
||||
let mut px = Vec::with_capacity(p.in_w * p.in_h);
|
||||
for y in 0..p.in_h {
|
||||
let sy = reflect(y as i64 - p.top as i64, h);
|
||||
for x in 0..p.in_w {
|
||||
let sx = reflect(x as i64 - p.left as i64, w);
|
||||
px.push(all.get(sy * w + sx).copied().unwrap_or([0.0; 3]));
|
||||
}
|
||||
}
|
||||
Input {
|
||||
width: p.in_w,
|
||||
height: p.in_h,
|
||||
px,
|
||||
}
|
||||
}
|
||||
|
||||
/// The network input of a large picture: reflect-padded to a square (at the left and top) and
|
||||
/// reduced to `in_w × in_h` with antialiased Lanczos-3. Source rows stream through in strips;
|
||||
/// each is filtered horizontally and added to the output rows it feeds.
|
||||
pub(super) fn reduce_input(
|
||||
picture: Picture<'_>,
|
||||
p: &Plan,
|
||||
ctx: &RunContext,
|
||||
) -> Result<Input, AiError> {
|
||||
let (w, h) = (picture.width, picture.height);
|
||||
// A tap past the square reflects at the square's edge first (the resize pads the padded
|
||||
// square), then lands in the picture through the padding's own reflection.
|
||||
let source = |i: i64, pad: usize, n: usize| reflect(reflect(i, pad + n) as i64 - pad as i64, n);
|
||||
let cols: Vec<Vec<(usize, f32)>> = (0..p.in_w)
|
||||
.map(|o| {
|
||||
let (lo, ws) = lanczos_taps(o, p.scale);
|
||||
ws.into_iter()
|
||||
.enumerate()
|
||||
.map(|(i, wt)| (source(lo + i as i64, p.left, w), wt))
|
||||
.collect()
|
||||
})
|
||||
.collect();
|
||||
// For each source row, the output rows it feeds and with what weight.
|
||||
let mut feeds: Vec<Vec<(usize, f32)>> = vec![Vec::new(); h];
|
||||
for oy in 0..p.in_h {
|
||||
let (lo, ws) = lanczos_taps(oy, p.scale);
|
||||
for (i, wt) in ws.into_iter().enumerate() {
|
||||
if let Some(f) = feeds.get_mut(source(lo + i as i64, p.top, h)) {
|
||||
f.push((oy, wt));
|
||||
}
|
||||
}
|
||||
}
|
||||
let mut out = vec![[0.0f32; 3]; p.in_w * p.in_h];
|
||||
for y0 in (0..h).step_by(STRIP) {
|
||||
ctx.check_cancelled()?;
|
||||
let y1 = (y0 + STRIP).min(h);
|
||||
let strip = picture.rgb(0, y0, w, y1);
|
||||
for (src, f) in strip.chunks_exact(w).zip(feeds.get(y0..y1).unwrap_or(&[])) {
|
||||
let row: Vec<[f32; 3]> = cols
|
||||
.iter()
|
||||
.map(|taps| {
|
||||
let mut acc = [0.0f32; 3];
|
||||
for (c, wt) in taps {
|
||||
let s = src.get(*c).copied().unwrap_or([0.0; 3]);
|
||||
for k in 0..3 {
|
||||
acc[k] += wt * s[k];
|
||||
}
|
||||
}
|
||||
acc
|
||||
})
|
||||
.collect();
|
||||
for &(oy, wt) in f {
|
||||
let dst = out
|
||||
.get_mut(oy * p.in_w..(oy + 1) * p.in_w)
|
||||
.unwrap_or(&mut []);
|
||||
for (d, s) in dst.iter_mut().zip(&row) {
|
||||
for k in 0..3 {
|
||||
d[k] += wt * s[k];
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
ctx.progress((10 * y1 / h.max(1)) as u8);
|
||||
}
|
||||
for px in &mut out {
|
||||
*px = px.map(|v| v.clamp(0.0, 1.0));
|
||||
}
|
||||
Ok(Input {
|
||||
width: p.in_w,
|
||||
height: p.in_h,
|
||||
px: out,
|
||||
})
|
||||
}
|
||||
|
||||
/// The full-resolution `w`-wide matte of a large picture: the prediction enlarged back to the
|
||||
/// padded square with Lanczos-3 and cropped, one output row at a time (vertical taps first, so
|
||||
/// nothing larger than a prediction row is held per row).
|
||||
pub(super) fn expand_matte(pred: &[f32], p: &Plan, w: usize, mask: &mut [u8]) {
|
||||
let up = 1.0 / p.scale;
|
||||
let cols: Vec<(i64, Vec<f32>)> = (0..w).map(|x| lanczos_taps(x + p.left, up)).collect();
|
||||
let mut line = vec![0.0f32; p.in_w];
|
||||
for (y, row) in mask.chunks_exact_mut(w.max(1)).enumerate() {
|
||||
let (lo, wy) = lanczos_taps(y + p.top, up);
|
||||
line.fill(0.0);
|
||||
for (i, wt) in wy.iter().enumerate() {
|
||||
let r = reflect(lo + i as i64, p.in_h);
|
||||
let src = pred.get(r * p.in_w..(r + 1) * p.in_w).unwrap_or(&[]);
|
||||
for (l, s) in line.iter_mut().zip(src) {
|
||||
*l += wt * s;
|
||||
}
|
||||
}
|
||||
for (o, (lo, wx)) in row.iter_mut().zip(&cols) {
|
||||
let v: f32 = wx
|
||||
.iter()
|
||||
.enumerate()
|
||||
.map(|(i, wt)| {
|
||||
wt * line
|
||||
.get(reflect(lo + i as i64, p.in_w))
|
||||
.copied()
|
||||
.unwrap_or(0.0)
|
||||
})
|
||||
.sum();
|
||||
*o = (v.clamp(0.0, 1.0) * 255.0).round() as u8;
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,337 @@
|
||||
use std::cell::RefCell;
|
||||
|
||||
use craft_ai_core::{AiError, ErrorKind};
|
||||
use craft_ai_engine::RunContext;
|
||||
|
||||
use super::guided::{filter, refine};
|
||||
use super::resize::{Input, expand_matte, lanczos_taps, reduce_input, reflect};
|
||||
use super::*;
|
||||
|
||||
/// A disc of `inside` colour on an `outside` background, and which pixels are inside.
|
||||
fn disc(w: usize, h: usize, cx: f32, cy: f32, rad: f32) -> (Vec<u8>, Vec<bool>) {
|
||||
let mut bytes = Vec::with_capacity(w * h * 3);
|
||||
let mut truth = Vec::with_capacity(w * h);
|
||||
for y in 0..h {
|
||||
for x in 0..w {
|
||||
let (dx, dy) = (x as f32 + 0.5 - cx, y as f32 + 0.5 - cy);
|
||||
let inside = dx * dx + dy * dy <= rad * rad;
|
||||
bytes.extend_from_slice(if inside {
|
||||
&[230, 76, 51]
|
||||
} else {
|
||||
&[51, 102, 178]
|
||||
});
|
||||
truth.push(inside);
|
||||
}
|
||||
}
|
||||
(bytes, truth)
|
||||
}
|
||||
|
||||
fn picture(bytes: &[u8], width: usize, height: usize) -> Picture<'_> {
|
||||
Picture {
|
||||
width,
|
||||
height,
|
||||
bytes,
|
||||
}
|
||||
}
|
||||
|
||||
/// Intersection over union of the coverage above one half and `truth`.
|
||||
fn iou(mask: &[u8], truth: &[bool]) -> f64 {
|
||||
let (mut both, mut any) = (0usize, 0usize);
|
||||
for (m, t) in mask.iter().zip(truth) {
|
||||
let m = *m > 127;
|
||||
both += usize::from(m && *t);
|
||||
any += usize::from(m || *t);
|
||||
}
|
||||
both as f64 / any.max(1) as f64
|
||||
}
|
||||
|
||||
/// A stand-in network: foreground where the red channel is high. It records the input sizes.
|
||||
struct Red(RefCell<Vec<(usize, usize)>>);
|
||||
|
||||
impl Red {
|
||||
fn new() -> Self {
|
||||
Self(RefCell::new(Vec::new()))
|
||||
}
|
||||
|
||||
fn predict(&self, input: &Input) -> Result<Vec<f32>, AiError> {
|
||||
assert_eq!(input.width % ALIGN, 0);
|
||||
assert_eq!(input.height % ALIGN, 0);
|
||||
self.0.borrow_mut().push((input.width, input.height));
|
||||
Ok(input
|
||||
.px
|
||||
.iter()
|
||||
.map(|p| if p[0] > 0.6 { 1.0 } else { 0.0 })
|
||||
.collect())
|
||||
}
|
||||
}
|
||||
|
||||
fn run(pic: Picture<'_>, shrink: f32, model: &Red) -> Result<Vec<u8>, AiError> {
|
||||
matte(
|
||||
pic,
|
||||
MAX_RES,
|
||||
shrink,
|
||||
&mut |i| model.predict(i),
|
||||
&RunContext::detached(),
|
||||
)
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn plan_pads_small_pictures_to_the_stride_centred() {
|
||||
let p = plan(100, 50, MAX_RES).unwrap();
|
||||
assert_eq!(
|
||||
(p.scale, p.in_w, p.in_h, p.left, p.top),
|
||||
(1.0, 128, 64, 14, 7)
|
||||
);
|
||||
assert!(!p.guided());
|
||||
let p = plan(64, 1280, MAX_RES).unwrap();
|
||||
assert_eq!((p.in_w, p.in_h, p.left, p.top), (64, 1280, 0, 0));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn plan_pads_large_pictures_to_a_square_at_the_left_and_top() {
|
||||
let p = plan(4000, 2000, MAX_RES).unwrap();
|
||||
assert_eq!(
|
||||
(p.scale, p.in_w, p.in_h, p.left, p.top),
|
||||
(0.32, 1280, 1280, 0, 2000)
|
||||
);
|
||||
assert!(p.guided());
|
||||
assert!(!plan(2000, 1000, MAX_RES).unwrap().guided());
|
||||
assert_eq!(plan(100_000, 3, MAX_RES).unwrap().top, 99_997);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn plan_rejects_empty_and_bounds_max_res() {
|
||||
assert_eq!(plan(0, 10, MAX_RES), None);
|
||||
assert_eq!(plan(10, 0, MAX_RES), None);
|
||||
assert_eq!(plan(5000, 5000, 0).unwrap().in_w, ALIGN);
|
||||
assert_eq!(plan(5000, 10, 1000).unwrap().in_w, 992);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn reflect_mirrors_without_repeating_the_edge() {
|
||||
let r: Vec<usize> = (-4..8).map(|i| reflect(i, 3)).collect();
|
||||
assert_eq!(r, [0, 1, 2, 1, 0, 1, 2, 1, 0, 1, 2, 1]);
|
||||
assert_eq!(reflect(-7, 1), 0);
|
||||
assert_eq!(reflect(i64::MIN / 2, 5), reflect(i64::MIN / 2 + 8, 5));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn lanczos_taps_are_normalised_and_exact_at_scale_one() {
|
||||
for scale in [0.07, 0.32, 0.5, 0.9, 1.7, 4.0] {
|
||||
for o in [0, 1, 17, 200] {
|
||||
let (_, w) = lanczos_taps(o, scale);
|
||||
assert!(
|
||||
(w.iter().sum::<f32>() - 1.0).abs() < 1e-5,
|
||||
"scale {scale} o {o}"
|
||||
);
|
||||
}
|
||||
}
|
||||
assert!(lanczos_taps(10, 0.25).1.len() >= 23);
|
||||
assert!(lanczos_taps(10, 3.0).1.len() <= 7);
|
||||
let (lo, w) = lanczos_taps(5, 1.0);
|
||||
let one: Vec<(i64, f32)> = (lo..).zip(w).filter(|(_, v)| v.abs() > 1e-6).collect();
|
||||
assert_eq!(one, [(5, 1.0)]);
|
||||
}
|
||||
|
||||
/// `src` (`sw × sh`, `ch` channels) resized by applying the taps directly, reflecting at its
|
||||
/// edges: what the streamed versions must equal.
|
||||
fn naive_resize(
|
||||
src: &[f32],
|
||||
sw: usize,
|
||||
sh: usize,
|
||||
ch: usize,
|
||||
ow: usize,
|
||||
oh: usize,
|
||||
scale: f64,
|
||||
) -> Vec<f32> {
|
||||
let mut out = vec![0.0f32; ow * oh * ch];
|
||||
for oy in 0..oh {
|
||||
let (ly, wy) = lanczos_taps(oy, scale);
|
||||
for ox in 0..ow {
|
||||
let (lx, wx) = lanczos_taps(ox, scale);
|
||||
for (j, a) in wy.iter().enumerate() {
|
||||
for (i, b) in wx.iter().enumerate() {
|
||||
let s = reflect(ly + j as i64, sh) * sw + reflect(lx + i as i64, sw);
|
||||
for c in 0..ch {
|
||||
out[(oy * ow + ox) * ch + c] += a * b * src[s * ch + c];
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
out
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn streamed_reduction_matches_the_padded_square() {
|
||||
for (w, h) in [(300usize, 180usize), (300, 40)] {
|
||||
let bytes: Vec<u8> = (0..w * h)
|
||||
.flat_map(|i| {
|
||||
let (x, y) = (i % w, i / w);
|
||||
[
|
||||
(x * 255 / w) as u8,
|
||||
(y * 255 / h) as u8,
|
||||
((x * 7 + y * 13) % 17 * 15) as u8,
|
||||
]
|
||||
})
|
||||
.collect();
|
||||
let pic = picture(&bytes, w, h);
|
||||
let p = plan(w, h, 64).unwrap();
|
||||
let got = reduce_input(pic, &p, &RunContext::detached()).unwrap();
|
||||
let side = w.max(h);
|
||||
let square: Vec<f32> = (0..side * side)
|
||||
.flat_map(|i| {
|
||||
let sx = reflect((i % side) as i64 - p.left as i64, w);
|
||||
let sy = reflect((i / side) as i64 - p.top as i64, h);
|
||||
let at = (sy * w + sx) * 3;
|
||||
[bytes[at], bytes[at + 1], bytes[at + 2]].map(|v| v as f32 / 255.0)
|
||||
})
|
||||
.collect();
|
||||
let want = naive_resize(&square, side, side, 3, 64, 64, p.scale);
|
||||
let worst = got
|
||||
.px
|
||||
.iter()
|
||||
.flatten()
|
||||
.zip(&want)
|
||||
.map(|(a, b)| (a - b.clamp(0.0, 1.0)).abs())
|
||||
.fold(0.0f32, f32::max);
|
||||
assert!(worst < 1e-4, "{w}×{h}: {worst}");
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn row_by_row_enlargement_matches_the_padded_square() {
|
||||
let (w, h) = (300usize, 180usize);
|
||||
let p = plan(w, h, 64).unwrap();
|
||||
let pred: Vec<f32> = (0..64 * 64)
|
||||
.map(|i| (((i % 64) as f32 - 30.0) / 6.0).tanh() * 0.5 + 0.5)
|
||||
.collect();
|
||||
let mut got = vec![0u8; w * h];
|
||||
expand_matte(&pred, &p, w, &mut got);
|
||||
let side = w.max(h);
|
||||
let full = naive_resize(&pred, 64, 64, 1, side, side, 1.0 / p.scale);
|
||||
let worst = (0..w * h)
|
||||
.map(|i| {
|
||||
let (x, y) = (i % w, i / w);
|
||||
let want = (full[(y + p.top) * side + x + p.left].clamp(0.0, 1.0) * 255.0).round();
|
||||
(got[i] as i32 - want as i32).abs()
|
||||
})
|
||||
.max()
|
||||
.unwrap();
|
||||
assert!(worst <= 1, "{worst}");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn shrink_maps_the_fringe() {
|
||||
assert_eq!(shrink_lut(0.0), std::array::from_fn(|v| v as u8));
|
||||
let l = shrink_lut(0.2);
|
||||
assert_eq!((l[0], l[51], l[153], l[255]), (0, 0, 128, 255));
|
||||
assert_eq!(shrink_lut(1.0)[255], 255);
|
||||
assert_eq!(shrink_lut(f32::NAN), shrink_lut(0.0));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn a_small_picture_runs_at_full_resolution() {
|
||||
let (bytes, truth) = disc(200, 150, 100.0, 75.0, 40.0);
|
||||
let model = Red::new();
|
||||
let mask = run(picture(&bytes, 200, 150), 0.0, &model).unwrap();
|
||||
assert_eq!(*model.0.borrow(), [(224, 160)]);
|
||||
assert!(iou(&mask, &truth) > 0.98, "iou {}", iou(&mask, &truth));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn a_large_picture_is_reduced_and_refined_at_full_resolution() {
|
||||
// 2800 × 1400: working scale 0.457, so the guided filter runs.
|
||||
let (w, h) = (2800usize, 1400usize);
|
||||
let (bytes, truth) = disc(w, h, 1400.0, 700.0, 500.0);
|
||||
let model = Red::new();
|
||||
let raw = run(picture(&bytes, w, h), 0.0, &model).unwrap();
|
||||
assert_eq!(*model.0.borrow(), [(1280, 1280)]);
|
||||
// The guided filter snaps the blurry enlarged edge to the picture's hard edge.
|
||||
assert!(iou(&raw, &truth) > 0.995, "iou {}", iou(&raw, &truth));
|
||||
// The default shrink removes the faint haze the filter windows leave outside the edge.
|
||||
let mask = run(picture(&bytes, w, h), SHRINK, &model).unwrap();
|
||||
assert_eq!(mask[700 * w + 1400], 255);
|
||||
assert_eq!(mask[190 * w + 1400], 0);
|
||||
assert!(iou(&mask, &truth) > 0.995);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn tiled_refinement_matches_the_whole_picture() {
|
||||
let (w, h) = (700usize, 560usize);
|
||||
let (bytes, _) = disc(w, h, 330.0, 290.0, 200.0);
|
||||
let mut mask: Vec<u8> = (0..w * h)
|
||||
.map(|i| {
|
||||
let (x, y) = ((i % w) as f32, (i / w) as f32);
|
||||
let d = ((x - 340.0).powi(2) + (y - 280.0).powi(2)).sqrt();
|
||||
(((210.0 - d) / 12.0).clamp(0.0, 1.0) * 255.0).round() as u8
|
||||
})
|
||||
.collect();
|
||||
let pm: Vec<f32> = mask.iter().map(|v| *v as f32 / 255.0).collect();
|
||||
let pic = picture(&bytes, w, h);
|
||||
let guide = pic.rgb(0, 0, w, h);
|
||||
let whole: Vec<u8> = filter(&guide, &pm, w, h, 32, GUIDED_EPS)
|
||||
.iter()
|
||||
.zip(&pm)
|
||||
.map(|(g, p)| (g.clamp(0.0, 1.0).max(*p) * 255.0).round() as u8)
|
||||
.collect();
|
||||
refine(pic, &mut mask, 32, GUIDED_EPS, &RunContext::detached()).unwrap();
|
||||
let worst = mask
|
||||
.iter()
|
||||
.zip(&whole)
|
||||
.map(|(a, b)| (*a as i32 - *b as i32).abs())
|
||||
.max()
|
||||
.unwrap();
|
||||
assert!(worst <= 1, "tiles differ from the whole picture by {worst}");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn a_wrong_answer_from_the_network_is_an_error_and_odd_values_clamp() {
|
||||
let (bytes, _) = disc(64, 48, 32.0, 24.0, 12.0);
|
||||
let pic = picture(&bytes, 64, 48);
|
||||
let ctx = RunContext::detached();
|
||||
let answer = |v: Vec<f32>| matte(pic, MAX_RES, 0.0, &mut |_| Ok(v.clone()), &ctx);
|
||||
assert_eq!(answer(vec![0.5; 7]).unwrap_err().kind, ErrorKind::Internal);
|
||||
assert_eq!(answer(Vec::new()).unwrap_err().kind, ErrorKind::Internal);
|
||||
let failed = matte(
|
||||
pic,
|
||||
MAX_RES,
|
||||
0.0,
|
||||
&mut |_| Err(AiError::invalid("no")),
|
||||
&ctx,
|
||||
);
|
||||
assert_eq!(failed.unwrap_err().message, "no");
|
||||
// NaN is background, +inf foreground. Input 64×64 holds the picture 8 rows down.
|
||||
let mut odd = vec![f32::NAN; 64 * 64];
|
||||
odd[16 * 64 + 10] = f32::INFINITY;
|
||||
odd[17 * 64 + 10] = -3.0;
|
||||
let mask = answer(odd).unwrap();
|
||||
assert_eq!(mask.iter().filter(|v| **v > 0).count(), 1);
|
||||
assert_eq!(mask[8 * 64 + 10], 255);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn bad_and_empty_pictures() {
|
||||
let ctx = RunContext::detached();
|
||||
let model = Red::new();
|
||||
// Bytes that do not match the stated size are refused, never indexed past.
|
||||
let short = matte(
|
||||
picture(&[1, 2, 3], 4, 4),
|
||||
MAX_RES,
|
||||
0.0,
|
||||
&mut |i| model.predict(i),
|
||||
&ctx,
|
||||
);
|
||||
assert_eq!(short.unwrap_err().kind, ErrorKind::InvalidRequest);
|
||||
let empty = matte(
|
||||
picture(&[], 0, 5),
|
||||
MAX_RES,
|
||||
0.0,
|
||||
&mut |i| model.predict(i),
|
||||
&ctx,
|
||||
)
|
||||
.unwrap();
|
||||
assert!(empty.is_empty());
|
||||
let one = run(picture(&[230, 50, 50], 1, 1), SHRINK, &model).unwrap();
|
||||
assert_eq!(one, [255]);
|
||||
}
|
||||
@@ -0,0 +1,287 @@
|
||||
use std::path::{Path, PathBuf};
|
||||
use std::sync::Arc;
|
||||
|
||||
use burn::prelude::*;
|
||||
use burn_store::{BurnpackStore, ModuleSnapshot};
|
||||
use craft_ai_core::{AiError, CapabilityId, ErrorKind, InputRole, ModelInfo, Payload, Request};
|
||||
use craft_ai_engine::testing::check_runner;
|
||||
use craft_ai_engine::{LoadedModel, ModelRunner, RunContext};
|
||||
use craft_ai_engine_rten::picture::{decode_picture, encode_mask};
|
||||
use craft_ai_features::GpuVendor;
|
||||
|
||||
use super::net::InSPyReNet;
|
||||
use super::net::swin::SwinConfig;
|
||||
use super::pipeline::Input;
|
||||
use super::{InspyrenetRunner, Loaded, Predict, on_big_stack, weights};
|
||||
|
||||
/// A Swin backbone small enough to build, save and run in a test.
|
||||
const TINY: SwinConfig = SwinConfig {
|
||||
embed_dim: 8,
|
||||
depths: [1, 1, 1, 1],
|
||||
heads: [1, 1, 1, 1],
|
||||
window: 12,
|
||||
patch: 4,
|
||||
};
|
||||
|
||||
#[allow(deprecated)]
|
||||
fn processor() -> Device {
|
||||
burn::backend::ndarray::NdArrayDevice::Cpu.into()
|
||||
}
|
||||
|
||||
/// A tiny network with placeholder weights, saved as a Burn copy in `dir`.
|
||||
fn tiny_weights(dir: &Path) -> PathBuf {
|
||||
let path = dir.join("tiny.bpk");
|
||||
let target = path.clone();
|
||||
on_big_stack(move || {
|
||||
let net = InSPyReNet::with_backbone(TINY, &processor());
|
||||
net.save_into(&mut BurnpackStore::from_file(&target))
|
||||
.map_err(|_| craft_ai_core::AiError::invalid("save"))
|
||||
})
|
||||
.unwrap();
|
||||
path
|
||||
}
|
||||
|
||||
/// The tiny network's runner, on the processor whatever graphics the computer has.
|
||||
fn tiny_runner() -> Arc<dyn ModelRunner> {
|
||||
Arc::new(InspyrenetRunner {
|
||||
backbone: TINY,
|
||||
device: processor,
|
||||
})
|
||||
}
|
||||
|
||||
/// A request for `image.matte` with a gray `w × h` PNG.
|
||||
fn request(w: usize, h: usize) -> Request {
|
||||
let alpha: Vec<f32> = (0..w * h).map(|i| (i % 7) as f32 / 7.0).collect();
|
||||
let png = encode_mask(&alpha, w, h).unwrap();
|
||||
Request::new(CapabilityId::parse("image.matte").unwrap()).with_input(
|
||||
InputRole::Source,
|
||||
Payload::inline("image/png", png).unwrap(),
|
||||
)
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn the_runner_describes_an_image_matte_from_one_picture() {
|
||||
let info = InspyrenetRunner::new().model();
|
||||
assert_eq!(info.id.as_str(), "inspyrenet");
|
||||
assert_eq!(InspyrenetRunner::new().engine(), "burn");
|
||||
let spec = info.capabilities.first().unwrap();
|
||||
assert_eq!(spec.capability.as_str(), "image.matte");
|
||||
assert!(
|
||||
crate::runners()
|
||||
.iter()
|
||||
.any(|r| r.model().id.as_str() == "inspyrenet")
|
||||
);
|
||||
}
|
||||
|
||||
/// Stands in for the network: the left half of the network's input is the subject.
|
||||
struct LeftHalf;
|
||||
|
||||
impl Predict for LeftHalf {
|
||||
fn predict(&mut self, input: &Input, _: &RunContext) -> Result<Vec<f32>, AiError> {
|
||||
Ok((0..input.width * input.height)
|
||||
.map(|i| {
|
||||
if i % input.width < input.width / 2 {
|
||||
1.0
|
||||
} else {
|
||||
0.0
|
||||
}
|
||||
})
|
||||
.collect())
|
||||
}
|
||||
}
|
||||
|
||||
/// The real runner's description and request handling around [`LeftHalf`], whatever file is
|
||||
/// named: the network itself is far too slow for these tests in an unoptimised build.
|
||||
struct StandIn;
|
||||
|
||||
impl ModelRunner for StandIn {
|
||||
fn engine(&self) -> &str {
|
||||
crate::runner::ENGINE
|
||||
}
|
||||
fn model(&self) -> ModelInfo {
|
||||
InspyrenetRunner::new().model()
|
||||
}
|
||||
fn accelerators(&self) -> &[GpuVendor] {
|
||||
&super::CARDS
|
||||
}
|
||||
fn load(&self, _: &Path) -> Result<Box<dyn LoadedModel>, AiError> {
|
||||
Ok(Box::new(Loaded::new(Box::new(LeftHalf))))
|
||||
}
|
||||
}
|
||||
|
||||
/// The mask PNG of a run's one output, after checking that there is exactly one.
|
||||
fn one_mask(out: &[craft_ai_core::Output]) -> craft_ai_engine_rten::picture::Rgb {
|
||||
assert_eq!(out.len(), 1);
|
||||
let Payload::Inline { media_type, data } = &out[0].payload else {
|
||||
panic!("not inline");
|
||||
};
|
||||
assert_eq!(media_type.as_str(), "image/png");
|
||||
decode_picture(data).unwrap()
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn passes_the_shared_engine_checks() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let weights = dir.path().join("inspyrenet.safetensors");
|
||||
std::fs::write(&weights, b"x").unwrap();
|
||||
let bad = Request::new(CapabilityId::parse("image.matte").unwrap());
|
||||
check_runner(
|
||||
Arc::new(StandIn),
|
||||
&weights,
|
||||
&request(40, 30),
|
||||
&bad,
|
||||
ErrorKind::InvalidRequest,
|
||||
)
|
||||
.unwrap();
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn answers_one_mask_at_the_pictures_size() {
|
||||
let mut model = StandIn.load(Path::new("unused")).unwrap();
|
||||
let mask = one_mask(
|
||||
&model
|
||||
.run(&request(50, 20), &RunContext::detached())
|
||||
.unwrap(),
|
||||
);
|
||||
assert_eq!((mask.width, mask.height), (50, 20));
|
||||
// The subject is the left half: white there, black on the right after the fringe shrink.
|
||||
let at = |x: usize| mask.bytes[3 * x];
|
||||
assert!(at(5) > 200 && at(45) < 50, "{} {}", at(5), at(45));
|
||||
}
|
||||
|
||||
/// The tiny network itself, through load and run on the processor. Ignored because Burn's
|
||||
/// processor backend is very slow without optimisation; run it in release.
|
||||
#[test]
|
||||
#[ignore = "slow in an unoptimised build: cargo test --release -p craft-ai-engine-burn --lib tiny_network -- --ignored"]
|
||||
fn the_tiny_network_answers_one_mask_at_the_pictures_size() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let mut model = tiny_runner().load(&tiny_weights(dir.path())).unwrap();
|
||||
let mask = one_mask(
|
||||
&model
|
||||
.run(&request(50, 20), &RunContext::detached())
|
||||
.unwrap(),
|
||||
);
|
||||
assert_eq!((mask.width, mask.height), (50, 20));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn files_that_are_not_the_checkpoint_are_refused_in_words() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let runner = InspyrenetRunner {
|
||||
device: processor,
|
||||
..InspyrenetRunner::new()
|
||||
};
|
||||
let missing = runner
|
||||
.load(&dir.path().join("inspyrenet.safetensors"))
|
||||
.err()
|
||||
.unwrap();
|
||||
assert_eq!(missing.kind, ErrorKind::Unsupported);
|
||||
for (name, bytes) in [
|
||||
("junk.pth", b"not a checkpoint".as_slice()),
|
||||
("cut.bpk", b"NRUB and then nothing".as_slice()),
|
||||
] {
|
||||
let path = dir.path().join(name);
|
||||
std::fs::write(&path, bytes).unwrap();
|
||||
let refused = runner.load(&path).err().unwrap();
|
||||
assert_eq!(refused.kind, ErrorKind::Unsupported, "{name}");
|
||||
}
|
||||
// The tiny copy does not fit the Swin-B network: refused, not half loaded.
|
||||
let tiny = tiny_weights(dir.path());
|
||||
assert_eq!(
|
||||
runner.load(&tiny).err().unwrap().kind,
|
||||
ErrorKind::Unsupported
|
||||
);
|
||||
}
|
||||
|
||||
/// The published checkpoint, for the tests that need the real weights.
|
||||
fn checkpoint() -> Option<PathBuf> {
|
||||
std::env::var_os("CRAFT_TEST_INSPYRENET_WEIGHTS").map(PathBuf::from)
|
||||
}
|
||||
|
||||
#[test]
|
||||
#[ignore = "needs CRAFT_TEST_INSPYRENET_WEIGHTS (inspyrenet.safetensors)"]
|
||||
fn the_published_checkpoint_fills_every_weight() {
|
||||
let Some(path) = checkpoint() else {
|
||||
return;
|
||||
};
|
||||
let values = on_big_stack(move || {
|
||||
let device = processor();
|
||||
let net = weights::fill(weights::untrained(&device), &path)?;
|
||||
let x = Tensor::<4>::from_data(
|
||||
burn::tensor::TensorData::new(vec![0.1f32; 3 * 64 * 64], [1, 3, 64, 64]),
|
||||
&device,
|
||||
);
|
||||
let out = net
|
||||
.forward(x, &|| false)
|
||||
.ok_or_else(craft_ai_core::AiError::cancelled)?;
|
||||
out.into_data()
|
||||
.try_to_vec::<f32>()
|
||||
.map_err(|_| craft_ai_core::AiError::invalid("unreadable"))
|
||||
})
|
||||
.unwrap();
|
||||
assert_eq!(values.len(), 64 * 64);
|
||||
assert!(values.iter().all(|v| v.is_finite()));
|
||||
}
|
||||
|
||||
/// The whole runner on a real photo, as an app calls it: the picture in, the mask PNG out, written
|
||||
/// to `CRAFT_TEST_INSPYRENET_OUT` for a look. Prints the times.
|
||||
#[test]
|
||||
#[ignore = "needs CRAFT_TEST_INSPYRENET_WEIGHTS, CRAFT_TEST_INSPYRENET_IMAGE and CRAFT_TEST_INSPYRENET_OUT"]
|
||||
fn removes_the_background_of_a_real_photo() {
|
||||
let (Some(weights), Some(image), Some(out)) = (
|
||||
checkpoint(),
|
||||
std::env::var_os("CRAFT_TEST_INSPYRENET_IMAGE"),
|
||||
std::env::var_os("CRAFT_TEST_INSPYRENET_OUT"),
|
||||
) else {
|
||||
return;
|
||||
};
|
||||
let picture = std::fs::read(image).unwrap();
|
||||
let request = Request::new(CapabilityId::parse("image.matte").unwrap()).with_input(
|
||||
InputRole::Source,
|
||||
Payload::inline("image/jpeg", picture).unwrap(),
|
||||
);
|
||||
let started = std::time::Instant::now();
|
||||
let mut model = InspyrenetRunner::new().load(&weights).unwrap();
|
||||
println!("load {:.2?}", started.elapsed());
|
||||
for round in 0..2 {
|
||||
let started = std::time::Instant::now();
|
||||
let answer = model.run(&request, &RunContext::detached()).unwrap();
|
||||
println!("run {round}: {:.2?}", started.elapsed());
|
||||
let Payload::Inline { data, .. } = &answer[0].payload else {
|
||||
panic!("not inline");
|
||||
};
|
||||
std::fs::write(&out, data).unwrap();
|
||||
}
|
||||
}
|
||||
|
||||
/// A load with nothing missing, no error and only recomputed buffers left over is complete; one
|
||||
/// that misses a weight, had an error, or left a tensor this network does not have is not.
|
||||
#[test]
|
||||
fn a_load_that_is_not_whole_is_refused() {
|
||||
use burn_store::{ApplyError, ApplyResult};
|
||||
let clean = ApplyResult {
|
||||
applied: vec!["backbone.norm0.weight".into()],
|
||||
skipped: Vec::new(),
|
||||
missing: Vec::new(),
|
||||
unused: vec![
|
||||
"backbone.layers.0.blocks.0.attn.relative_position_index".into(),
|
||||
"context1.branch0.conv0.bn.num_batches_tracked".into(),
|
||||
],
|
||||
errors: Vec::new(),
|
||||
};
|
||||
assert!(weights::complete(&clean));
|
||||
let mut missing = clean.clone();
|
||||
missing
|
||||
.missing
|
||||
.push(("decoder.conv1.conv.weight".into(), "Struct:Conv2d".into()));
|
||||
assert!(!weights::complete(&missing));
|
||||
let mut failed = clean.clone();
|
||||
failed.errors.push(ApplyError::LoadError {
|
||||
path: "backbone.norm0.weight".into(),
|
||||
message: "unreadable".into(),
|
||||
});
|
||||
assert!(!weights::complete(&failed));
|
||||
let mut foreign = clean.clone();
|
||||
foreign.unused.push("head.classifier.weight".into());
|
||||
assert!(!weights::complete(&foreign));
|
||||
}
|
||||
@@ -0,0 +1,66 @@
|
||||
//! Filling the network from its checkpoint: the published InSPyReNet Swin-B (Plus Ultra) weights
|
||||
//! as a `.safetensors` file, read with Burn's PyTorch adapter (PyTorch's `[out, in]` linear
|
||||
//! weights are transposed as they load). A saved Burn copy (`.bpk`) is read too.
|
||||
|
||||
use std::io::Read;
|
||||
use std::path::Path;
|
||||
|
||||
use burn_store::{
|
||||
ApplyResult, BurnpackStore, ModuleSnapshot, PyTorchToBurnAdapter, SafetensorsStore,
|
||||
};
|
||||
use craft_ai_core::AiError;
|
||||
|
||||
use super::net::InSPyReNet;
|
||||
|
||||
/// The first bytes of a saved Burn weight file.
|
||||
const BURNPACK: &[u8; 4] = b"NRUB";
|
||||
|
||||
/// Buffers the checkpoint carries that the port computes itself.
|
||||
fn recomputed(key: &str) -> bool {
|
||||
key.ends_with(".relative_position_index") || key.ends_with(".num_batches_tracked")
|
||||
}
|
||||
|
||||
/// Said when a file cannot be used: damaged, cut short, or another model's.
|
||||
pub(super) fn unreadable() -> AiError {
|
||||
AiError::unsupported(
|
||||
"This file is not the InSPyReNet checkpoint Craft AI expects. Open Set up AI to download it again.",
|
||||
)
|
||||
}
|
||||
|
||||
/// Whether every weight was set, none failed, and nothing in the file was left over apart from the
|
||||
/// buffers the port recomputes.
|
||||
pub(super) fn complete(result: &ApplyResult) -> bool {
|
||||
result.is_success() && result.missing.is_empty() && result.unused.iter().all(|k| recomputed(k))
|
||||
}
|
||||
|
||||
/// Whether the file starts like a saved Burn weight file.
|
||||
fn is_burnpack(path: &Path) -> bool {
|
||||
let mut head = [0u8; 4];
|
||||
std::fs::File::open(path)
|
||||
.and_then(|mut f| f.read_exact(&mut head))
|
||||
.is_ok()
|
||||
&& &head == BURNPACK
|
||||
}
|
||||
|
||||
/// Fills `net` from the file at `path`, refusing one that is damaged, leaves any weight unset or
|
||||
/// holds tensors this network does not have.
|
||||
pub(super) fn fill(mut net: InSPyReNet, path: &Path) -> Result<InSPyReNet, AiError> {
|
||||
let applied = if is_burnpack(path) {
|
||||
net.load_from(&mut BurnpackStore::from_file(path))
|
||||
.map_err(|_| unreadable())?
|
||||
} else {
|
||||
let mut store = SafetensorsStore::from_file(path).with_from_adapter(PyTorchToBurnAdapter);
|
||||
net.load_from(&mut store).map_err(|_| unreadable())?
|
||||
};
|
||||
if complete(&applied) {
|
||||
Ok(net)
|
||||
} else {
|
||||
Err(unreadable())
|
||||
}
|
||||
}
|
||||
|
||||
/// The network with placeholder weights on `device`, built the way a load builds it.
|
||||
#[cfg(test)]
|
||||
pub(super) fn untrained(device: &burn::prelude::Device) -> InSPyReNet {
|
||||
InSPyReNet::new(device)
|
||||
}
|
||||
@@ -1,6 +1,10 @@
|
||||
//! SAM 2.1 on the Burn engine for Craft AI: an optional runner for the `image.segment` capability
|
||||
//! that uses a graphics card when one answers and the processor when none does. The default
|
||||
//! engine for this model stays rten (`craft-ai-engine-rten`); this one is a choice.
|
||||
//! The Burn engine for Craft AI: optional runners that use a graphics card when one answers and
|
||||
//! the processor when none does.
|
||||
//!
|
||||
//! - **SAM 2.1** for `image.segment`. The default engine for this model stays rten
|
||||
//! (`craft-ai-engine-rten`); this one is a choice. The rest of this page is about it.
|
||||
//! - **InSPyReNet** for `image.matte` (Remove background): the subject's soft alpha from the
|
||||
//! picture alone, at any picture size. See the `inspyrenet` module and `docs/BURN-ENGINE.md`.
|
||||
//!
|
||||
//! An app does not use this crate directly. It puts [`runners`] in an
|
||||
//! [`EngineRegistry`](craft_ai_engine::EngineRegistry), or turns on the `burn` feature of
|
||||
@@ -31,6 +35,7 @@
|
||||
extern crate alloc;
|
||||
|
||||
mod backend;
|
||||
mod inspyrenet;
|
||||
// Generated code: its lints are the generator's business, and its forward passes can panic on a
|
||||
// shape they were not generated for, so the runner calls them only through `backend::guarded`.
|
||||
#[allow(
|
||||
@@ -47,5 +52,6 @@ mod model;
|
||||
mod runner;
|
||||
|
||||
pub use backend::{DECODER_FILE, ENCODER_FILE, graphics_card_answers};
|
||||
pub use inspyrenet::{InspyrenetRunner, MODEL as INSPYRENET_MODEL};
|
||||
pub use model::ARCHITECTURE_REVISION;
|
||||
pub use runner::{ENGINE, MODEL, Sam21Runner, runners};
|
||||
|
||||
@@ -35,7 +35,10 @@ const CARDS: [GpuVendor; 4] = [
|
||||
|
||||
/// Every runner this crate offers, for an `EngineRegistry`.
|
||||
pub fn runners() -> Vec<Arc<dyn ModelRunner>> {
|
||||
vec![Arc::new(Sam21Runner)]
|
||||
vec![
|
||||
Arc::new(Sam21Runner),
|
||||
Arc::new(crate::inspyrenet::InspyrenetRunner::new()),
|
||||
]
|
||||
}
|
||||
|
||||
/// SAM 2.1 on the Burn engine. It uses a graphics card when one answers, else the processor.
|
||||
|
||||
@@ -132,7 +132,12 @@ fn the_runner_is_the_burn_sam21_the_catalogue_names() {
|
||||
assert_eq!(spec.requires, vec![InputRole::Source]);
|
||||
assert!(spec.accepts.contains_key(&InputRole::Prompt));
|
||||
assert!(runner.accelerators().contains(&GpuVendor::Nvidia));
|
||||
assert_eq!(runners().len(), 1);
|
||||
// SAM 2.1 for `image.segment`, InSPyReNet for `image.matte`.
|
||||
let ids: Vec<String> = runners()
|
||||
.iter()
|
||||
.map(|r| r.model().id.as_str().to_owned())
|
||||
.collect();
|
||||
assert_eq!(ids, ["sam21", "inspyrenet"]);
|
||||
}
|
||||
|
||||
/// The downloaded ONNX file is not what this engine reads, and it is not the supported export
|
||||
|
||||
@@ -0,0 +1,54 @@
|
||||
//! InSPyReNet (Swin-B, Plus Ultra) for Remove background, run by the built-in Burn engine only.
|
||||
|
||||
use crate::recipe::{Licence, LocalRecipe, ModelFile, NativeRoute, Source};
|
||||
|
||||
/// The author's checkpoint as a `.safetensors` file at a pinned revision. Its 1,625 tensors were
|
||||
/// checked one by one against the author's own `ckpt_base.pth` (plemeri/InSPyReNet on Hugging
|
||||
/// Face). The engine reads this file as it is: no conversion, nothing hosted by us.
|
||||
fn inspyrenet_files() -> Vec<ModelFile> {
|
||||
const FILE: &str = "inspyrenet.safetensors";
|
||||
vec![ModelFile {
|
||||
folder: "inspyrenet".to_owned(),
|
||||
filename: FILE.to_owned(),
|
||||
size_bytes: 367_120_444,
|
||||
sha256: "5c8ab4541f9a960b376a80362b07f1d96262e689c501365e6323faa958de47ea".to_owned(),
|
||||
source: Source::HuggingFace {
|
||||
repo: "1038lab/inspyrenet".to_owned(),
|
||||
file: FILE.to_owned(),
|
||||
revision: Some("0f5946638dcbc7dbc1e56306b823959958ac0ae0".to_owned()),
|
||||
},
|
||||
licence: Licence {
|
||||
name: "MIT, as declared by the author of InSPyReNet".to_owned(),
|
||||
url: Some("https://github.com/plemeri/InSPyReNet/blob/main/LICENSE".to_owned()),
|
||||
accept_required: false,
|
||||
note: Some(
|
||||
"This file is a format conversion of the author's checkpoint, identical tensor for tensor. The publisher of the conversion states no licence of its own."
|
||||
.to_owned(),
|
||||
),
|
||||
},
|
||||
}]
|
||||
}
|
||||
|
||||
/// InSPyReNet for Remove background. ComfyUI does not run it: a node pack and a workflow would be
|
||||
/// far more to set up than the one file the built-in engine reads.
|
||||
pub(super) fn inspyrenet() -> LocalRecipe {
|
||||
LocalRecipe {
|
||||
label: "InSPyReNet".to_owned(),
|
||||
usable: true,
|
||||
workflow: "inspyrenet-matte".to_owned(),
|
||||
comfyui: false,
|
||||
models: Vec::new(),
|
||||
packs: Vec::new(),
|
||||
min_vram_mb: None,
|
||||
native: vec![NativeRoute {
|
||||
provider: "burn".to_owned(),
|
||||
model: "inspyrenet".to_owned(),
|
||||
label: "InSPyReNet, built in (any graphics card)".to_owned(),
|
||||
models: inspyrenet_files(),
|
||||
note: None,
|
||||
// Seconds a picture on the processor, not measured as quick (docs/BURN-ENGINE.md).
|
||||
fast_on_processor: false,
|
||||
}],
|
||||
components: Vec::new(),
|
||||
}
|
||||
}
|
||||
@@ -9,10 +9,12 @@ use crate::feature::{AppStatus, AppUse, Feature};
|
||||
use crate::recipe::CloudOption;
|
||||
|
||||
mod florence2;
|
||||
mod inspyrenet;
|
||||
mod sam21;
|
||||
mod sam3;
|
||||
use craft_ai_core::{CapabilityId, InputRole};
|
||||
pub use florence2::FLORENCE2_LARGE_DIR;
|
||||
use inspyrenet::inspyrenet;
|
||||
use sam3::sam3;
|
||||
use sam21::sam21;
|
||||
|
||||
@@ -81,8 +83,8 @@ fn try_builtin() -> Option<Vec<Feature>> {
|
||||
)?;
|
||||
background.inputs = vec![InputRole::Source];
|
||||
background.uses = uses(&[("PhotoCraft", "Remove Background")]);
|
||||
// No local recipe and no hosted route yet: no adapter here matts an image, and no model file
|
||||
// has been checked against its real source. The screen says so instead of offering a guess.
|
||||
// One model, built in only (Burn), and no hosted route yet.
|
||||
background.locals = vec![inspyrenet()];
|
||||
|
||||
let mut remove = feature(
|
||||
"remove-object",
|
||||
|
||||
@@ -76,6 +76,7 @@ pub(super) fn sam21() -> LocalRecipe {
|
||||
label: "SAM 2.1".to_owned(),
|
||||
usable: true,
|
||||
workflow: "sam21-segment".to_owned(),
|
||||
comfyui: true,
|
||||
models,
|
||||
packs: vec![
|
||||
NodePack {
|
||||
|
||||
@@ -54,6 +54,7 @@ pub(super) fn sam3() -> LocalRecipe {
|
||||
// PhotoCraft runs it today (Object Selection, `Sam3Session` through ComfyUI).
|
||||
usable: true,
|
||||
workflow: "sam3-segment".to_owned(),
|
||||
comfyui: true,
|
||||
models: vec![ModelFile {
|
||||
folder: "sam3".to_owned(),
|
||||
filename: "sam3.safetensors".to_owned(),
|
||||
|
||||
@@ -57,12 +57,18 @@ impl Feature {
|
||||
})
|
||||
}
|
||||
|
||||
/// The first recipe apps can run. The ComfyUI route, a plain built-in Install offer and
|
||||
/// engine advice's hardware limits use it.
|
||||
/// The first recipe apps can run. A plain built-in Install offer and engine advice's hardware
|
||||
/// limits use it.
|
||||
pub fn first_usable_recipe(&self) -> Option<&LocalRecipe> {
|
||||
self.locals.iter().find(|r| r.usable)
|
||||
}
|
||||
|
||||
/// The recipes ComfyUI runs, in catalogue order. A recipe that is built in only
|
||||
/// ([`LocalRecipe::comfyui`] false) is never one of them.
|
||||
pub fn comfyui_recipes(&self) -> impl Iterator<Item = &LocalRecipe> {
|
||||
self.locals.iter().filter(|r| r.comfyui)
|
||||
}
|
||||
|
||||
/// The recipe the built-in engine `provider` acts on: the first one apps can run that lists
|
||||
/// that engine. Planning, setup options, the Install offer for that engine, file setup and
|
||||
/// registration all use it, so an option the picker shows is never one that setup then
|
||||
|
||||
@@ -49,7 +49,7 @@ pub struct SetupOption {
|
||||
pub size_bytes: Option<u64>,
|
||||
}
|
||||
|
||||
/// Every local model first, one per recipe in catalogue order, then every tool's built-in
|
||||
/// Every local model first, one per recipe ComfyUI runs in catalogue order, then every tool's built-in
|
||||
/// runtimes (one per runtime, from its first usable recipe), then hosted services in the order the
|
||||
/// features first name them.
|
||||
pub fn setup_options(features: &[Feature]) -> Vec<SetupOption> {
|
||||
@@ -57,7 +57,7 @@ pub fn setup_options(features: &[Feature]) -> Vec<SetupOption> {
|
||||
let mut natives = Vec::new();
|
||||
let mut services: Vec<SetupOption> = Vec::new();
|
||||
for feature in features {
|
||||
for recipe in &feature.locals {
|
||||
for recipe in feature.comfyui_recipes() {
|
||||
locals.push(SetupOption {
|
||||
id: format!("local:{}:{}", feature.id, recipe.workflow),
|
||||
title: recipe.label.clone(),
|
||||
|
||||
@@ -90,23 +90,22 @@ pub struct Plan {
|
||||
pub blockers: Vec<Blocker>,
|
||||
}
|
||||
|
||||
/// Works out what `feature`'s first local recipe still needs on a computer described by
|
||||
/// `inventory`. A feature with no local recipe gets [`State::NoLocalRoute`] and an empty
|
||||
/// [`Plan::recipe`].
|
||||
/// Works out what `feature`'s first ComfyUI recipe still needs on a computer described by
|
||||
/// `inventory`. A feature with no ComfyUI recipe gets [`State::NoLocalRoute`] and an empty
|
||||
/// [`Plan::recipe`]; a built-in route is planned by [`plan_native`].
|
||||
pub fn plan(feature: &Feature, inventory: &Inventory) -> Plan {
|
||||
match feature.locals.first() {
|
||||
match feature.comfyui_recipes().next() {
|
||||
Some(recipe) => plan_recipe(feature, recipe, inventory),
|
||||
None => empty(feature, String::new()),
|
||||
}
|
||||
}
|
||||
|
||||
/// Works out what every one of `feature`'s local recipes still needs, one [`Plan`] each, in
|
||||
/// Works out what every one of `feature`'s ComfyUI recipes still needs, one [`Plan`] each, in
|
||||
/// catalogue order. Each names its recipe in [`Plan::recipe`], so several models for one tool
|
||||
/// can be shown and installed side by side without one standing in for another.
|
||||
pub fn plans(feature: &Feature, inventory: &Inventory) -> Vec<Plan> {
|
||||
feature
|
||||
.locals
|
||||
.iter()
|
||||
.comfyui_recipes()
|
||||
.map(|recipe| plan_recipe(feature, recipe, inventory))
|
||||
.collect()
|
||||
}
|
||||
|
||||
@@ -142,7 +142,8 @@ impl NativeRoute {
|
||||
}
|
||||
}
|
||||
|
||||
/// Running a feature on this computer through ComfyUI.
|
||||
/// Running a feature on this computer: through ComfyUI, through the built-in runtimes in
|
||||
/// [`LocalRecipe::native`], or both.
|
||||
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
|
||||
pub struct LocalRecipe {
|
||||
/// What the model is called, for people: "SAM 3".
|
||||
@@ -150,8 +151,14 @@ pub struct LocalRecipe {
|
||||
/// Apps can run it today. False when the files can be downloaded but no workflow runs them yet.
|
||||
#[serde(default = "yes")]
|
||||
pub usable: bool,
|
||||
/// Id of the workflow template the adapter runs.
|
||||
/// Id of the workflow template the adapter runs. For a recipe ComfyUI does not run, only the
|
||||
/// recipe's id.
|
||||
pub workflow: String,
|
||||
/// ComfyUI runs this recipe. False for a model only the built-in runtimes in
|
||||
/// [`LocalRecipe::native`] run: it is then never a ComfyUI setup option, plan or connection,
|
||||
/// and its `models` and `packs` are empty.
|
||||
#[serde(default = "yes")]
|
||||
pub comfyui: bool,
|
||||
pub models: Vec<ModelFile>,
|
||||
pub packs: Vec<NodePack>,
|
||||
/// Graphics memory the workflow needs, when known.
|
||||
|
||||
@@ -273,7 +273,7 @@ fn offer_for(
|
||||
None => feature.first_usable_recipe()?,
|
||||
}
|
||||
} else {
|
||||
feature.locals.first().filter(|r| r.usable)?
|
||||
feature.comfyui_recipes().next().filter(|r| r.usable)?
|
||||
};
|
||||
// The built-in route compares the runtime's own files: the named one, or else the first
|
||||
// runtime that names files at all.
|
||||
|
||||
@@ -0,0 +1,93 @@
|
||||
//! A recipe only a built-in runtime runs: Remove background on InSPyReNet, read by the Burn engine.
|
||||
//! It is offered and planned as that runtime alone, never as a ComfyUI model, install or plan.
|
||||
use craft_ai_features::{
|
||||
Feature, Inventory, Os, Route, SetupRoute, Source, State, SystemProfile, plan, plan_native,
|
||||
plans, setup_offer, setup_offer_native, setup_options,
|
||||
};
|
||||
|
||||
fn remove_background() -> Feature {
|
||||
craft_ai_features::builtin()
|
||||
.into_iter()
|
||||
.find(|f| f.id == "remove-background")
|
||||
.unwrap()
|
||||
}
|
||||
|
||||
fn windows() -> SystemProfile {
|
||||
SystemProfile {
|
||||
os: Os::Windows,
|
||||
gpus: Vec::new(),
|
||||
comfyui_found: false,
|
||||
}
|
||||
}
|
||||
|
||||
/// Plenty of space, ComfyUI present, nothing downloaded.
|
||||
fn inventory() -> Inventory {
|
||||
Inventory {
|
||||
comfyui_found: true,
|
||||
free_disk_bytes: Some(500_000_000_000),
|
||||
..Inventory::default()
|
||||
}
|
||||
}
|
||||
|
||||
/// The one recipe is built in only, names the Burn runner's model, and pins the published file.
|
||||
#[test]
|
||||
fn remove_background_runs_inspyrenet_on_burn_from_one_pinned_file() {
|
||||
let feature = remove_background();
|
||||
assert_eq!(feature.locals.len(), 1);
|
||||
let recipe = &feature.locals[0];
|
||||
assert!(recipe.usable && !recipe.comfyui);
|
||||
assert!(recipe.models.is_empty() && recipe.packs.is_empty());
|
||||
assert_eq!(feature.engines(), ["burn"]);
|
||||
assert_eq!(feature.comfyui_recipes().count(), 0);
|
||||
let route = &recipe.native[0];
|
||||
assert_eq!(route.model, "inspyrenet");
|
||||
let weights = route.weights().unwrap();
|
||||
assert_eq!(weights.filename, "inspyrenet.safetensors");
|
||||
assert_eq!(weights.size_bytes, 367_120_444);
|
||||
assert_eq!(
|
||||
weights.sha256,
|
||||
"5c8ab4541f9a960b376a80362b07f1d96262e689c501365e6323faa958de47ea"
|
||||
);
|
||||
assert!(!weights.licence.accept_required);
|
||||
match &weights.source {
|
||||
Source::HuggingFace { revision, .. } => {
|
||||
assert_eq!(revision.as_deref().map(str::len), Some(40), "pinned");
|
||||
}
|
||||
Source::Url { .. } => panic!("expected a pinned Hugging Face file"),
|
||||
}
|
||||
}
|
||||
|
||||
/// Set up AI lists the built-in route and no ComfyUI model for the tool.
|
||||
#[test]
|
||||
fn only_the_built_in_route_is_offered() {
|
||||
let options = setup_options(&[remove_background()]);
|
||||
assert_eq!(options.len(), 1, "{options:?}");
|
||||
assert_eq!(options[0].id, "native:remove-background:burn");
|
||||
assert!(matches!(&options[0].route, Route::Native { model, .. } if model == "inspyrenet"));
|
||||
assert_eq!(options[0].size_bytes, Some(367_120_444));
|
||||
assert!(!setup_options(&craft_ai_features::builtin()).iter().any(
|
||||
|o| matches!(&o.route, Route::Local { feature, .. } if feature == "remove-background")
|
||||
));
|
||||
}
|
||||
|
||||
/// No ComfyUI plan or install offer exists for it; the built-in plan and offer download the file.
|
||||
#[test]
|
||||
fn it_is_planned_and_installed_as_the_built_in_route_only() {
|
||||
let feature = remove_background();
|
||||
let inventory = inventory();
|
||||
assert_eq!(plan(&feature, &inventory).state, State::NoLocalRoute);
|
||||
assert!(plans(&feature, &inventory).is_empty());
|
||||
// Not even on a computer without ComfyUI, where a ComfyUI offer would start with installing it.
|
||||
let no_comfyui = Inventory {
|
||||
comfyui_found: false,
|
||||
..inventory.clone()
|
||||
};
|
||||
assert!(setup_offer(&feature, &windows(), &no_comfyui, SetupRoute::ComfyUi).is_none());
|
||||
|
||||
let native = plan_native(&feature, "burn", &inventory);
|
||||
assert_eq!(native.state, State::NeedsInstall);
|
||||
assert_eq!(native.recipe, "inspyrenet-matte");
|
||||
assert_eq!(native.download_bytes, 367_120_444);
|
||||
let offer = setup_offer_native(&feature, "burn", &windows(), &inventory).unwrap();
|
||||
assert_eq!(offer.download_bytes, 367_120_444);
|
||||
}
|
||||
@@ -308,8 +308,8 @@ fn each_photocraft_tool_has_one_shared_tool_with_the_right_job_and_inputs() {
|
||||
|
||||
#[test]
|
||||
fn a_tool_without_a_checked_route_offers_none() {
|
||||
// Remove background has no verified local recipe and no hosted adapter that matts an image.
|
||||
let f = feature("remove-background");
|
||||
// Remove objects has no verified local recipe and no hosted adapter that fills a masked area.
|
||||
let f = feature("remove-object");
|
||||
assert!(f.locals.is_empty());
|
||||
assert!(f.cloud.is_empty());
|
||||
}
|
||||
|
||||
+98
-1
@@ -1,4 +1,4 @@
|
||||
# The Burn engine (SAM 2.1, optional)
|
||||
# The Burn engine (SAM 2.1 and InSPyReNet, optional)
|
||||
|
||||
`craft-ai-engine-burn` runs SAM 2.1 tiny for `image.segment` through Burn, on a graphics card when
|
||||
one answers and on the processor when none does. It is a choice beside rten, never the default:
|
||||
@@ -91,3 +91,100 @@ python tools/mutation/mutation_check.py be
|
||||
without tracking every value that passes between the halves. No file is 700 lines or more.
|
||||
- Progress and Cancel during the first preparation: `load` has neither, so the app shows its own
|
||||
busy state until the first selection returns.
|
||||
|
||||
## InSPyReNet for `image.matte`
|
||||
|
||||
The second model on this engine: `inspyrenet` (engine `burn`), the subject's soft alpha for
|
||||
Remove background, from the picture alone. It is not SAM: no clicks, no box, no prompt.
|
||||
|
||||
| Piece | Where |
|
||||
|---|---|
|
||||
| Runner, request and answer | `src/inspyrenet/mod.rs`: a PNG or JPEG `source` in, one grayscale PNG at the picture's size out (white is the subject, soft along hair and fur) |
|
||||
| The network | `src/inspyrenet/net/`: InSPyReNet with the Swin-B backbone (T. Kim et al., ACCV 2022), **written by hand** in Burn after the MIT-licensed reference code (plemeri/InSPyReNet, as packaged by transparent-background). Module names follow the PyTorch keys, so the checkpoint loads by name |
|
||||
| Weights | `src/inspyrenet/weights.rs`: the `.safetensors` checkpoint read with Burn's PyTorch adapter (linear weights transposed as they load); a load that leaves any weight unset, fails one, or finds a tensor this network does not have is refused in plain words |
|
||||
| Any picture size | `src/inspyrenet/pipeline/`: full resolution up to 1280 pixels on the long side (reflect-padded to a multiple of 32); above that, reflect-padded to a square, reduced with antialiased Lanczos-3 to 1280², the matte enlarged back and, below half scale, refined at full resolution with a colour guided filter (He et al., PAMI 2013) on the tiles near the edge; then a 0.2 fringe shrink. The approach of the Gyre background remover |
|
||||
|
||||
Unlike SAM 2.1 here, nothing is generated and nothing is prepared on the computer: the runner
|
||||
reads the checkpoint file as downloaded.
|
||||
|
||||
### The weights file
|
||||
|
||||
`inspyrenet.safetensors` from [1038lab/inspyrenet](https://huggingface.co/1038lab/inspyrenet) at
|
||||
revision `0f5946638dcbc7dbc1e56306b823959958ac0ae0`: 367,120,444 bytes, SHA-256
|
||||
`5c8ab4541f9a960b376a80362b07f1d96262e689c501365e6323faa958de47ea`. It is a format conversion
|
||||
with no model card of its own. On 29/10/2026 it was checked tensor for tensor (all 1,625 identical)
|
||||
against the author's own checkpoint, `ckpt_base.pth` (InSPyReNet Swin-B Plus Ultra, MIT) at
|
||||
[plemeri/InSPyReNet](https://huggingface.co/plemeri/InSPyReNet) revision
|
||||
`d94c2baaa4d023ab018c6f97be6ef37548e3bd1f`: 367,520,613 bytes, SHA-256
|
||||
`0a6fe2a73ab0532d6d0b8d82849a9760a226df719e3063d09b4149ece6f80fcd`, the file transparent-background
|
||||
downloads. The author's file is not used directly because it is a pickle (`.pth`), and a local
|
||||
connection takes only `.safetensors`, `.onnx` or `.bpk` weights.
|
||||
|
||||
### In the catalogue
|
||||
|
||||
`remove-background` has one recipe, InSPyReNet (`crates/ai-features/src/catalogue/inspyrenet.rs`),
|
||||
and it is **built in only**: `LocalRecipe::comfyui` is false, so it is never a ComfyUI setup option,
|
||||
plan, install offer or connection, and engine advice does not offer ComfyUI for the tool. Its one
|
||||
route is `burn` / `inspyrenet` with the pinned file above in the folder `inspyrenet`, so Set up AI
|
||||
offers one 367 MB download, checked by size and SHA-256, and records that file on the connection.
|
||||
A node pack and a workflow would be far more to set up for one file and one tool.
|
||||
|
||||
### What was run
|
||||
|
||||
On one Windows machine (i5-12600KF, RTX 3090 Ti), release:
|
||||
|
||||
- **Parity with PyTorch** (`src/inspyrenet/parity.rs`): every stage of the network (patch
|
||||
embedding, four Swin stages, the saliency and Laplacian pyramids, the prediction) against
|
||||
transparent-background's own code on a synthetic 384² input. Worst relative difference over the
|
||||
stages 3.2e-5 on the processor and 2.4e-4 on the graphics card; no pixel of the prediction
|
||||
changes side at 0.5.
|
||||
- **The shared conformance suite** (`check_runner`) on the runner's request handling and pipeline
|
||||
around a stand-in for the network, as the SAM 2.1 tests here do. The network itself, a tiny one
|
||||
saved as `.bpk`, runs through load and run in an ignored test on the processor, in release (about
|
||||
4 s). Unoptimised, Burn's processor backend takes minutes for it, and a build server's software
|
||||
renderer is too slow and, on Linux (lavapipe), fails.
|
||||
- **The whole runner on a real photo**, against the same port inside PhotoCraft (which runs it
|
||||
through the same pipeline): on a 3101 × 1613 PNG the two masks are identical, byte for byte. On a
|
||||
JPEG they differ by 0.003 on average (IoU 0.994), from the two JPEG decoders.
|
||||
- **Through the shared client** (`crates/ai-client/tests/burn_matte.rs`): the launcher sets up the
|
||||
catalogue's built-in route with the real file, then two app ids ask `AiClient` for Remove
|
||||
background only and each gets one grayscale mask at the picture's size (a 2848 × 4272 PNG). A
|
||||
file that is not the checkpoint is refused in words that point to Set up AI.
|
||||
- Not measured cleanly: the time per picture. During these runs another program held about 20 GB of
|
||||
the card's memory and the PhotoCraft port slowed down just as much (about 20 s a picture instead
|
||||
of about 2.5 s). No run on macOS or Linux, and no real app has mounted it yet.
|
||||
|
||||
```sh
|
||||
# The PyTorch reference, once (Python with torch, safetensors and transparent-background)
|
||||
python tools/burn/inspyrenet_reference.py <inspyrenet.safetensors> <inspyrenet-ref-384.safetensors>
|
||||
# Then the parity and weight tests, release
|
||||
CRAFT_TEST_INSPYRENET_WEIGHTS=<inspyrenet.safetensors> CRAFT_TEST_INSPYRENET_REFERENCE=<inspyrenet-ref-384.safetensors> cargo test --release -p craft-ai-engine-burn -j 8 --lib inspyrenet -- --include-ignored --nocapture
|
||||
# One photo through the runner, the mask written out
|
||||
CRAFT_TEST_INSPYRENET_WEIGHTS=<inspyrenet.safetensors> CRAFT_TEST_INSPYRENET_IMAGE=<photo> CRAFT_TEST_INSPYRENET_OUT=<mask.png> cargo test --release -p craft-ai-engine-burn -j 8 --lib removes_the_background -- --ignored --nocapture
|
||||
# Two apps through the shared client, from the catalogue's route
|
||||
CRAFT_TEST_INSPYRENET_WEIGHTS=<inspyrenet.safetensors> CRAFT_TEST_INSPYRENET_IMAGE=<photo.png> cargo test --release -p craft-ai-client --features burn --test burn_matte -- --ignored --nocapture
|
||||
# The tiny network through load and run, on the processor
|
||||
cargo test --release -p craft-ai-engine-burn --lib tiny_network -- --ignored
|
||||
python tools/mutation/mutation_check.py be inspyrenet
|
||||
python tools/mutation/mutation_check.py bo
|
||||
```
|
||||
|
||||
### Limits to know
|
||||
|
||||
- **Graphics card advice.** The route is not marked quick on the processor, so engine advice
|
||||
calls it slow on a computer without a card the engine can use.
|
||||
- **A software renderer counts as a card.** The engine takes the graphics card when one answers,
|
||||
and a software renderer (lavapipe, WARP) answers too. On CI the network was far slower there and
|
||||
failed on lavapipe. So the first failure on a card that is not a cancel builds the network again
|
||||
on the processor from the same file and runs that picture again there with a notice. Once that
|
||||
rebuild has worked, the model stays on the processor. The move is tried once: if the rebuild
|
||||
fails, its error is the answer and later pictures go to the card again. A load on the card that
|
||||
fails is likewise retried on the processor, unless the file itself is not the checkpoint
|
||||
(`src/inspyrenet/fallback.rs`). SAM 2.1 on this engine makes the same choice of
|
||||
device without that move. `CRAFT_AI_BURN_DEVICE=cpu` forces the processor from the start.
|
||||
- **The processor is slow.** The fallback runs on Burn's `ndarray` backend, which this crate
|
||||
builds without its `multi-threads` and `simd` features, so it uses one core. The same network on
|
||||
one core took about 230 s a picture in PhotoCraft (its `flex` backend), about 35 s with four
|
||||
threads and SIMD. The graphics card is the point of this runner.
|
||||
- The parity reference is made by `tools/burn/inspyrenet_reference.py`, which needs Python with
|
||||
PyTorch and transparent-background.
|
||||
|
||||
+1
-1
@@ -16,7 +16,7 @@ No real hosted service, download host or repository was contacted by these tests
|
||||
| ArtCraft Omni API key | `craft-ai-artcraft` | Prompt-only image and video. Account or subscription access has no adapter: no official programmatic route is documented. An API key is not evidence that a subscription covers the requests. |
|
||||
| Hosted media copies | `craft-ai-media` | https only (loopback only for a connection marked local), every redirect re-checked, private addresses refused, size bounded, file type taken from the bytes. A failed copy keeps the original link with a notice. A busy host is retried with a bounded, doubling pause. Not built: retries for provider API calls, expiry tracking, resume, and an input-side upload cache. |
|
||||
| Generic ComfyUI workflows | `craft-ai-comfyui` | Maps a prompt field; `image.inpaint` needs mapped `LoadImage` and `LoadImageMask` fields. A workflow is offered as a model only after `/object_info` confirms its nodes and checkpoint names; submission repeats the check. Output stays inline (32 MiB) unless the app gives a directory (1 GiB, local `file:` handle, app owns cleanup). Targeted cancel is off by default because an older server may treat it as a global interrupt. Imported image.segment workflows map a description and source to a checked, source-sized grayscale mask plus the shared result. Empty masks mean no match; source uploads are remembered for selection jobs. This is not calibrated absence detection: the example Florence graph can find a box for an absent subject. No input file handles or general parameter mapping yet. |
|
||||
| Built-in engines | `craft-ai-engine`, `craft-ai-engine-candle`, `craft-ai-engine-rten`, `craft-ai-engine-burn` | Contract, registry and provider are tested with a fake runner and the shared conformance suite. SAM 3 on Candle (point and text prompts, processor, or Metal on macOS) ran on a real photograph and passes the same suite on Meta's checkpoint; the live tests skip without `CRAFT_SAM3_WEIGHTS`. It reads the Hugging Face checkpoint layout with tokenizer files, not the ComfyUI file. Pictures go in as PNG or JPEG (a JPEG's EXIF turn is applied). No Nvidia (CUDA) build, and the launcher does not download the checkpoint yet. SAM 2.1 on rten (processor only) reads the publisher's two `.onnx` files, pinned by size and SHA-256, with click and box prompts; a description is refused in plain words. Pictures are bounded by side and pixel count. With the real files in release on one Windows machine: a mask at the picture's size, and a hover through the ordinary job path with a median of about 35 ms after a first look of about 0.7 s (one picture, not a benchmark). SAM 2.1 is the first recipe for Select subject, so the plain built-in offer is SAM 2.1 on rten. No macOS or Linux run, and no app has run it through the shared binding yet. SAM 2.1 on Burn (`craft-ai-engine-burn`, cargo feature `burn`, optional, after rten in the recipe) reads weights that `craft-ai-burn-setup` prepares once on the computer; on one Windows machine with a graphics card its mask matched rten's (best-mask IoU 1.0, one click and one box, one picture, not a benchmark). Given the downloaded `encoder.onnx`, the runner prepares the saved copy itself on first load (in `burn-prepared/` beside the model folder; no progress and no Cancel during that load). No app has run it through the shared binding yet, and there has been no live run on macOS or Linux. See `docs/BURN-ENGINE.md`. |
|
||||
| Built-in engines | `craft-ai-engine`, `craft-ai-engine-candle`, `craft-ai-engine-rten`, `craft-ai-engine-burn` | Contract, registry and provider are tested with a fake runner and the shared conformance suite. SAM 3 on Candle (point and text prompts, processor, or Metal on macOS) ran on a real photograph and passes the same suite on Meta's checkpoint; the live tests skip without `CRAFT_SAM3_WEIGHTS`. It reads the Hugging Face checkpoint layout with tokenizer files, not the ComfyUI file. Pictures go in as PNG or JPEG (a JPEG's EXIF turn is applied). No Nvidia (CUDA) build, and the launcher does not download the checkpoint yet. SAM 2.1 on rten (processor only) reads the publisher's two `.onnx` files, pinned by size and SHA-256, with click and box prompts; a description is refused in plain words. Pictures are bounded by side and pixel count. With the real files in release on one Windows machine: a mask at the picture's size, and a hover through the ordinary job path with a median of about 35 ms after a first look of about 0.7 s (one picture, not a benchmark). SAM 2.1 is the first recipe for Select subject, so the plain built-in offer is SAM 2.1 on rten. No macOS or Linux run, and no app has run it through the shared binding yet. SAM 2.1 on Burn (`craft-ai-engine-burn`, cargo feature `burn`, optional, after rten in the recipe) reads weights that `craft-ai-burn-setup` prepares once on the computer; on one Windows machine with a graphics card its mask matched rten's (best-mask IoU 1.0, one click and one box, one picture, not a benchmark). Given the downloaded `encoder.onnx`, the runner prepares the saved copy itself on first load (in `burn-prepared/` beside the model folder; no progress and no Cancel during that load). No app has run it through the shared binding yet, and there has been no live run on macOS or Linux. InSPyReNet on Burn (`inspyrenet`, same crate and feature) answers `image.matte` from the picture alone with a grayscale PNG at the picture's size, at any picture size. Its hand-written network reads the published `.safetensors` checkpoint (checked tensor for tensor against the author's `ckpt_base.pth`); on one Windows machine every stage matched transparent-background's PyTorch code (worst relative difference 3.2e-5 on the processor, 2.4e-4 on the card, no pixel changing side), it passes the shared conformance suite, and on a real PNG its mask was byte-identical to the same port inside PhotoCraft. Time per picture was not measured cleanly (the card was shared during the runs). `remove-background` lists it as its one recipe, built in only (`LocalRecipe::comfyui` false): Set up AI offers the Burn route and its pinned file, never a ComfyUI model, and scripted tests cover the offer, the plan, the registration and the advice, and an ignored test ran the real file from that route through `AiClient` for two app ids; the file has not been downloaded through Set up AI, and no real app has mounted it. See `docs/BURN-ENGINE.md`. |
|
||||
| First-run setup flow | `craft-ai-client`, `craft-ai-egui`, `craft-ai-features` | One start button on `SetupSession` (checks complete files on disk, otherwise downloads; `blocker` says why it does nothing), a licence tick only where a file's licence needs one, `FolderDesk::ensure_default` for the suggested folder, `ToolDesk::adopt` for checked files only, a try-it sample that clicks for built-in selection models and fails a selection that chose nothing, status and picker that agree when a model's file has gone, and engine advice that does not call a measured-quick built-in engine (rten, one machine) slow. Scripted tests and a fake provider only: no real download, no model run through it, no app has mounted it, no macOS or Linux run. |
|
||||
| Shared egui widgets | `craft-ai-egui` | Source and tests migrated to egui 0.37 with the reviewed storytold egui family at `6d518fc799681571963842f95b55a6b68431b15a` and winit at `a517b3846ee8457d84ee4975110bfda6797c44dc`. The current headless candidate passed 126 tests with 0 failures and 4 ignored doctests under managed job `egui037-ai-final-ci-20261025-1331`, with formatting that left the worktree unchanged and strict all-target clippy for `craft-ai-egui` and `craft-ai-engine` on Linux. Five deterministic kittest `wgpu`/snapshot goldens are committed and passed in that run. The historical 108-test checkpoint at `egui037-ai-regression-mutation-20261024-2205` also recorded a pure-background RGB(27,27,27) to RGB(30,30,30) change failing the expected snapshot with 75,486 changed pixels, and the two compiled connection-panel lifecycle mutants caught before the source and golden hashes were restored. A locked consumer-smoke test build and its one draw test passed with one `craft-ai-core` identity and the pinned egui family; that rendererless root does not use the winit patch. No app imports the widgets. Native X11/Wayland/Windows behavior, IME/clipboard/focus, physical GPU qualification and performance measurements remain open. Full remote workspace/platform CI also remains open. |
|
||||
| Declarative adapter | `craft-ai-declarative` | Request/response and queue flows from a data file; prompt-only inputs. No launcher screen to add a spec yet. |
|
||||
|
||||
@@ -0,0 +1,59 @@
|
||||
"""PyTorch reference for the parity test in crates/ai-engine-burn/src/inspyrenet/parity.rs: the
|
||||
input, every stage and the prediction of InSPyReNet Swin-B, run by transparent-background's own
|
||||
code (pip install transparent-background torch safetensors) on a synthetic 384² picture.
|
||||
|
||||
usage: python tools/burn/inspyrenet_reference.py <ckpt_base.pth or inspyrenet.safetensors> <out.safetensors>
|
||||
"""
|
||||
import sys
|
||||
import types
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
# The Swin code imports three helpers from timm; stand-ins keep timm (and its version) out of it.
|
||||
layers = types.ModuleType("timm.models.layers")
|
||||
layers.DropPath = lambda *a, **k: nn.Identity()
|
||||
layers.to_2tuple = lambda x: x if isinstance(x, tuple) else (x, x)
|
||||
layers.trunc_normal_ = lambda t, *a, **k: t
|
||||
timm = types.ModuleType("timm")
|
||||
models = types.ModuleType("timm.models")
|
||||
models.layers = layers
|
||||
timm.models = models
|
||||
sys.modules.update({"timm": timm, "timm.models": models, "timm.models.layers": layers})
|
||||
|
||||
from safetensors.torch import save_file # noqa: E402
|
||||
from transparent_background.InSPyReNet import InSPyReNet_SwinB # noqa: E402
|
||||
|
||||
ckpt, out = sys.argv[1], sys.argv[2]
|
||||
model = InSPyReNet_SwinB(depth=64, pretrained=False, threshold=None, base_size=[1024, 1024])
|
||||
if ckpt.endswith(".safetensors"):
|
||||
from safetensors.torch import load_file
|
||||
|
||||
state = load_file(ckpt)
|
||||
else:
|
||||
state = torch.load(ckpt, map_location="cpu", weights_only=True)
|
||||
model.load_state_dict(state, strict=True)
|
||||
model.eval()
|
||||
|
||||
n = 384
|
||||
ys, xs = torch.meshgrid(torch.linspace(0, 1, n), torch.linspace(0, 1, n), indexing="ij")
|
||||
disc = (((xs - 0.55) ** 2 + (ys - 0.5) ** 2) < 0.08).float()
|
||||
g = torch.Generator().manual_seed(0)
|
||||
rgb = torch.stack([0.2 + 0.6 * disc, 0.3 + 0.3 * ys, 0.7 - 0.4 * disc * xs]) + 0.03 * torch.randn(3, n, n, generator=g)
|
||||
rgb = rgb.clamp(0, 1)
|
||||
mean = torch.tensor([0.485, 0.456, 0.406]).view(3, 1, 1)
|
||||
std = torch.tensor([0.229, 0.224, 0.225]).view(3, 1, 1)
|
||||
x = ((rgb - mean) / std).unsqueeze(0).contiguous()
|
||||
|
||||
with torch.no_grad():
|
||||
backbone = model.backbone(x)
|
||||
o = model.forward_inspyre(x)
|
||||
pred = torch.sigmoid(o["saliency"][3])
|
||||
pred = (pred - pred.min()) / (pred.max() - pred.min() + 1e-8)
|
||||
|
||||
tensors = {"input": x, "pred": pred}
|
||||
tensors.update({f"backbone.{i}": t for i, t in enumerate(backbone)})
|
||||
tensors.update({f"saliency.{i}": t for i, t in enumerate(o["saliency"])})
|
||||
tensors.update({f"laplacian.{i}": t for i, t in enumerate(o["laplacian"])})
|
||||
save_file({k: v.float().contiguous() for k, v in tensors.items()}, out)
|
||||
print({k: tuple(v.shape) for k, v in tensors.items()})
|
||||
@@ -14,6 +14,7 @@ honouring `CARGO_BUILD_JOBS` from the build resource.
|
||||
| `mutation_check.py rt` | rten engine picture handling: bounded PNG and JPEG decoding, mask encoding, resampling. |
|
||||
| `mutation_check.py bp` | Local Burn weight preparation (`ai-burn-setup`): source hashes, cache key, damaged saved weights, cancel. |
|
||||
| `mutation_check.py be` | Burn engine: saved-weights loading, the panic guard, answer checks, file names, its catalogue route, the `.bpk` binding rule and the shared architecture revision. |
|
||||
| `mutation_check.py bo` | A recipe only built-in runtimes run is never offered, planned, installed, registered or advised as a ComfyUI model. |
|
||||
| `mutation_check.py fr` | First-run setup flow: the licence tick only where needed, the one start button, the suggested folder, adopting checked files, the route-aware try-it sample, status and picker when a file has gone, route notes and engine advice. Cases in `mutation_first_run.py`. |
|
||||
| `mutation_check.py gs` | Generic selection: scratch sinks, the refusal of clicks and mapped selection workflows. |
|
||||
| `mutation_check.py gs` | Generic selection workflows: the scratch sinks, the refusal of clicks and mapped selection workflows. |
|
||||
|
||||
@@ -0,0 +1,21 @@
|
||||
"""bo: a recipe only built-in runtimes run (Remove background on InSPyReNet, Burn) is never treated
|
||||
as a ComfyUI model: not offered, planned, installed, registered or advised as one.
|
||||
|
||||
Loaded by mutation_check.py. Each case breaks one rule and names the test that should notice.
|
||||
Format: (name, file, old, new, package, test target, test filter).
|
||||
"""
|
||||
F = 'crates/ai-features/src/'
|
||||
CLS = 'crates/ai-client/src/'
|
||||
FE = 'craft-ai-features'
|
||||
CL = 'craft-ai-client'
|
||||
|
||||
BO = [
|
||||
('recipe-run-by-comfyui', F + 'catalogue/inspyrenet.rs', 'comfyui: false,', 'comfyui: true,', FE, 'built_in_only', 'runs_inspyrenet_on_burn'),
|
||||
('comfyui-option-for-built-in-recipe', F + 'options.rs', 'for recipe in feature.comfyui_recipes() {', 'for recipe in &feature.locals {', FE, 'built_in_only', 'only_the_built_in_route'),
|
||||
('comfyui-plan-for-built-in-recipe', F + 'plan.rs', ' match feature.comfyui_recipes().next() {', ' match feature.locals.first() {', FE, 'built_in_only', 'planned_and_installed'),
|
||||
('comfyui-plans-for-built-in-recipe', F + 'plan.rs', ' .comfyui_recipes()\n .map(|recipe| plan_recipe', ' .locals\n .iter()\n .map(|recipe| plan_recipe', FE, 'built_in_only', 'planned_and_installed'),
|
||||
('comfyui-offer-for-built-in-recipe', F + 'setup.rs', 'feature.comfyui_recipes().next().filter(|r| r.usable)?', 'feature.locals.first().filter(|r| r.usable)?', FE, 'built_in_only', 'planned_and_installed'),
|
||||
('comfyui-connection-for-built-in-recipe', CLS + 'local.rs', ' .comfyui_recipes()\n .find(|r| r.usable)', ' .locals\n .iter()\n .find(|r| r.usable)', CL, 'built_in_only', 'comfyui_is_refused'),
|
||||
('named-comfyui-connection-for-built-in-recipe', CLS + 'local.rs', ' .comfyui_recipes()\n .find(|r| r.workflow == workflow && r.usable)', ' .locals\n .iter()\n .find(|r| r.workflow == workflow && r.usable)', CL, 'built_in_only', 'comfyui_is_refused'),
|
||||
('comfyui-advised-without-recipe', CLS + 'engine.rs', 'if feature.comfyui_recipes().any(|r| r.usable) {', 'if true {', CL, 'built_in_only', 'never_comfyui'),
|
||||
]
|
||||
@@ -1,6 +1,7 @@
|
||||
"""be: the Burn engine's own rules (saved-weights loading, the panic guard, the answer checks, the
|
||||
file names it accepts), the catalogue route it is listed by, the binding's `.bpk` rule and the
|
||||
shared revision.
|
||||
shared revision; for InSPyReNet, the complete-weights rule, the matte size, the picture size and
|
||||
the move to the processor when the card fails.
|
||||
|
||||
Not reached here, because they need the real model, a graphics card or a graph the converter
|
||||
panics on (the ignored live tests in `crates/ai-engine-burn/tests/live.rs` and
|
||||
@@ -12,6 +13,7 @@ Case shape: name, file, old, new, package, target, test filter.
|
||||
"""
|
||||
B = 'crates/ai-engine-burn/src/'
|
||||
PKG = 'craft-ai-engine-burn'
|
||||
FALLBACK = 'inspyrenet/fallback.rs'
|
||||
|
||||
|
||||
def lib(name, file, old, new, test):
|
||||
@@ -54,6 +56,30 @@ BE = [
|
||||
lib('burn-damaged-copy-not-explained', 'runner/prepare.rs',
|
||||
'if error.message == craft_ai_burn_setup::DAMAGED_COPY {', 'if false {',
|
||||
'a_damaged_saved_copy_is_explained'),
|
||||
lib('inspyrenet-incomplete-weights-accepted', 'inspyrenet/weights.rs',
|
||||
'result.is_success() && result.missing.is_empty() && result.unused.iter().all(|k| recomputed(k))',
|
||||
'true', 'a_load_that_is_not_whole'),
|
||||
lib('inspyrenet-foreign-tensors-accepted', 'inspyrenet/weights.rs',
|
||||
'result.unused.iter().all(|k| recomputed(k))', 'true', 'a_load_that_is_not_whole'),
|
||||
lib('inspyrenet-matte-size-unchecked', 'inspyrenet/pipeline/mod.rs',
|
||||
'if pred.len() != p.in_w * p.in_h {', 'if false {', 'a_wrong_answer_from_the_network'),
|
||||
lib('inspyrenet-picture-bytes-unchecked', 'inspyrenet/pipeline/mod.rs',
|
||||
'.filter(|n| picture.bytes.len() == n * 3)', '.filter(|_| true)', 'bad_and_empty_pictures'),
|
||||
lib('inspyrenet-cancel-moves-to-processor', FALLBACK,
|
||||
'Err(error) if error.kind != ErrorKind::Cancelled => error,', 'Err(error) => error,',
|
||||
'a_cancel_is_not_a_failure'),
|
||||
lib('inspyrenet-never-moves-to-processor', FALLBACK,
|
||||
'let Some(rebuild) = self.rebuild.take() else {', 'let Some(rebuild) = None::<Rebuild> else {',
|
||||
'moves_to_the_processor'),
|
||||
lib('inspyrenet-processor-not-kept', FALLBACK,
|
||||
' self.current = rebuild()?;\n self.current.predict(input, ctx)',
|
||||
' rebuild()?.predict(input, ctx)', 'moves_to_the_processor'),
|
||||
lib('inspyrenet-card-load-not-retried', FALLBACK,
|
||||
'Err(_) => Ok((processor()?, false)),', 'Err(error) => Err(error),',
|
||||
'fails_to_load_loads_on_the_processor'),
|
||||
lib('inspyrenet-bad-file-read-twice', FALLBACK,
|
||||
'Err(error) if error.kind == ErrorKind::Unsupported => Err(error),',
|
||||
'Err(error) if false => Err(error),', 'not_read_twice'),
|
||||
('burn-revision-bumped-in-the-architecture', B + 'model/mod.rs',
|
||||
'onnx-0.22.0-v1";', 'onnx-0.22.0-v2";', PKG, 'revision', 'the_architecture_and_the_saved'),
|
||||
('burn-revision-bumped-in-the-cache-key', 'crates/ai-burn-setup/src/sam21.rs',
|
||||
|
||||
@@ -411,10 +411,10 @@ LC = [
|
||||
('non-local-address-accepted', LOCAL_RS, "if !endpoint.is_loopback() || trimmed.contains(['?', '#']) {", "if trimmed.contains(['?', '#']) {", CL, 'local', ''),
|
||||
('query-in-address-accepted', LOCAL_RS, "if !endpoint.is_loopback() || trimmed.contains(['?', '#']) {", 'if !endpoint.is_loopback() {', CL, 'local', ''),
|
||||
('trailing-slash-kept', LOCAL_RS, "input.trim_end_matches('/')", "input.trim_end_matches('~')", CL, 'local', ''),
|
||||
('unrunnable-recipe-accepted', LOCAL_RS, 'feature.locals.iter().find(|r| r.usable)', 'feature.locals.iter().find(|_| true)', CL, 'local', ''),
|
||||
('unrunnable-recipe-accepted', LOCAL_RS, ' .comfyui_recipes()\n .find(|r| r.usable)\n .ok_or_else', ' .comfyui_recipes()\n .find(|_| true)\n .ok_or_else', CL, 'local', ''),
|
||||
('unrunnable-named-recipe-accepted', LOCAL_RS, '.find(|r| r.workflow == workflow && r.usable)', '.find(|r| r.workflow == workflow)', CL, 'local_recipes', 'cannot_run_yet'),
|
||||
('named-recipe-name-ignored', LOCAL_RS, '.find(|r| r.workflow == workflow && r.usable)', '.find(|r| r.usable)', CL, 'local_recipes', ''),
|
||||
('second-recipe-takes-plain-id', LOCAL_RS, ' None => feature\n .locals\n .first()\n .is_some_and(|first| first.workflow == recipe.workflow),', ' None => true,', CL, 'local_recipes', 'own_connection'),
|
||||
('second-recipe-takes-plain-id', LOCAL_RS, ' None => feature\n .comfyui_recipes()\n .next()\n .is_some_and(|first| first.workflow == recipe.workflow),', ' None => true,', CL, 'local_recipes', 'own_connection'),
|
||||
('recipe-id-shares-runtime-namespace', LOCAL_RS, 'format!("local-{}-recipe-{}", feature.id, recipe.workflow)', 'format!("local-{}-{}", feature.id, recipe.workflow)', CL, 'local_recipes', 'built_in_runtime_never'),
|
||||
('picker-ignores-which-recipe-is-installed', PICK_G, '&& (!is_recipe_connection(c, feature) || m.id.as_str() == recipe)', '&& (!is_recipe_connection(c, feature) || true)', CL, 'local_recipes', 'keeps_the_setup_link'),
|
||||
('picker-recipe-connection-never-counts', PICK_G, '&& (!is_recipe_connection(c, feature) || m.id.as_str() == recipe)', '&& (!is_recipe_connection(c, feature) || false)', CL, 'local_recipes', 'setup_link'),
|
||||
@@ -492,6 +492,7 @@ from mutation_generic_uploads import GU_CASES # noqa: E402
|
||||
from mutation_segment import SEG_CASES # noqa: E402
|
||||
from mutation_burn_preparation import BP # noqa: E402
|
||||
from mutation_burn_engine import BE # noqa: E402
|
||||
from mutation_built_in_only import BO # noqa: E402
|
||||
from mutation_first_run import FR # noqa: E402
|
||||
|
||||
PS_RS = C + 'ai-comfy-env/src/source/pack/status.rs'
|
||||
@@ -554,7 +555,7 @@ S21 = [
|
||||
('selections-share-a-sink-name', SAM21, 'NEXT_SELECTION.fetch_add(1, Ordering::Relaxed)', '0', CU, 'sam21_session', 'its_own_sink_name'),
|
||||
]
|
||||
|
||||
GROUPS = {'fr': FR, 'be': BE, 'bp': BP, 'rt': RTEN_CASES, 'gu': GU_CASES, 'gm': GM_CASES, 'gs': GS_CASES, 'seg': SEG_CASES, 'nw': NW, 'nq': NQ, 'ns': NS, 'nf': NF, 'ce': CE, 'cg': CG, 'sf': SF, 's21': S21, 'ca': CA, 'lg': LG, 'ps': PS, 'ct': CT, 'u2': U2, 'nr': NR, 'u1': U1, 'u3': U3, 'u4': U4, 'p4': P4, 'p10': P10, 'v1': V1, 'v2': V2, 'v3': V3, 'sd': SD, 'vx': VX, 'lc': LC, 'lr': LR, 'wf': WF, 'eg': EG_CASES}
|
||||
GROUPS = {'fr': FR, 'be': BE, 'bo': BO, 'bp': BP, 'rt': RTEN_CASES, 'gu': GU_CASES, 'gm': GM_CASES, 'gs': GS_CASES, 'seg': SEG_CASES, 'nw': NW, 'nq': NQ, 'ns': NS, 'nf': NF, 'ce': CE, 'cg': CG, 'sf': SF, 's21': S21, 'ca': CA, 'lg': LG, 'ps': PS, 'ct': CT, 'u2': U2, 'nr': NR, 'u1': U1, 'u3': U3, 'u4': U4, 'p4': P4, 'p10': P10, 'v1': V1, 'v2': V2, 'v3': V3, 'sd': SD, 'vx': VX, 'lc': LC, 'lr': LR, 'wf': WF, 'eg': EG_CASES}
|
||||
# Mutants whose effect is also refused by a second, independent check on this machine. They are
|
||||
# recorded as equivalent here, with the reason, and are never counted as caught.
|
||||
EQUIVALENT = {
|
||||
|
||||
Reference in New Issue
Block a user