rlx_sam2/
flow.rs

1// RLX — versatile ML compiler + runtime.
2// Copyright (C) 2026 Eugene Hauptmann, Nataliya Kosmyna.
3//
4// This program is free software: you can redistribute it and/or modify
5// it under the terms of the GNU General Public License as published by
6// the Free Software Foundation, version 3.
7//
8// This program is distributed in the hope that it will be useful,
9// but WITHOUT ANY WARRANTY; without even the implied warranty of
10// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
11// GNU General Public License for more details.
12//
13// You should have received a copy of the GNU General Public License
14// along with this program. If not, see <https://www.gnu.org/licenses/>.
15
16//! Tier-0 SAM2 Hiera image encoder flow.
17
18use anyhow::Result;
19use rlx_flow::BuiltModel;
20
21use super::config::Sam2HieraConfig;
22use super::fpn_neck::FpnNeckWeights;
23use super::preprocess::Sam2PreprocessWeights;
24use rlx_core::flow_util::built_from_hir;
25use rlx_core::weight_map::WeightMap;
26
27#[derive(Debug, Clone)]
28pub struct Sam2ImageEncoderFlow<'a> {
29    cfg: &'a Sam2HieraConfig,
30}
31
32impl<'a> Sam2ImageEncoderFlow<'a> {
33    pub fn new(cfg: &'a Sam2HieraConfig) -> Self {
34        Self { cfg }
35    }
36
37    pub fn build(self, weights: &mut WeightMap) -> Result<Sam2ImageEncoderBuilt> {
38        let (hir, params, preprocess, neck) =
39            super::image_encoder::build_sam2_image_encoder_hir(self.cfg, weights)?;
40        Ok(Sam2ImageEncoderBuilt {
41            model: built_from_hir(hir, params)?,
42            preprocess,
43            neck,
44        })
45    }
46}
47
48pub struct Sam2ImageEncoderBuilt {
49    pub model: BuiltModel,
50    pub preprocess: Sam2PreprocessWeights,
51    pub neck: FpnNeckWeights,
52}
53
54pub fn build_sam2_image_encoder_built(
55    cfg: &Sam2HieraConfig,
56    weights: &mut WeightMap,
57) -> Result<Sam2ImageEncoderBuilt> {
58    Sam2ImageEncoderFlow::new(cfg).build(weights)
59}
rlx_sam2/flow.rs

rlx_sam2/
flow.rs
