Skip to main content

rlx_sam/
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 SAM v1 image encoder flow.
17
18use anyhow::Result;
19use rlx_flow::BuiltModel;
20
21use super::config::SamEncoderConfig;
22use super::preprocess::SamPreprocessWeights;
23use rlx_core::flow_util::built_from_hir;
24use rlx_core::weight_map::WeightMap;
25
26#[derive(Debug, Clone)]
27pub struct SamEncoderFlow<'a> {
28    cfg: &'a SamEncoderConfig,
29}
30
31impl<'a> SamEncoderFlow<'a> {
32    pub fn new(cfg: &'a SamEncoderConfig) -> Self {
33        Self { cfg }
34    }
35
36    pub fn build(self, weights: &mut WeightMap) -> Result<SamEncoderBuilt> {
37        let (hir, params, preprocess) =
38            super::image_encoder::build_sam_encoder_hir(self.cfg, weights)?;
39        Ok(SamEncoderBuilt {
40            model: built_from_hir(hir, params)?,
41            preprocess,
42        })
43    }
44}
45
46pub struct SamEncoderBuilt {
47    pub model: BuiltModel,
48    pub preprocess: SamPreprocessWeights,
49}
50
51pub fn build_sam_encoder_built(
52    cfg: &SamEncoderConfig,
53    weights: &mut WeightMap,
54) -> Result<SamEncoderBuilt> {
55    SamEncoderFlow::new(cfg).build(weights)
56}