rlx-sam 0.2.4

Segment Anything Model (SAM v1) for RLX
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64

// RLX — versatile ML compiler + runtime.
// Copyright (C) 2026 Eugene Hauptmann, Nataliya Kosmyna.
//
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU General Public License as published by
// the Free Software Foundation, version 3.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
// GNU General Public License for more details.
//
// You should have received a copy of the GNU General Public License
// along with this program. If not, see <https://www.gnu.org/licenses/>.

//! SAM v1 prompt-encoder mask downscale (IR).

use super::config::SAM_EMBED_HW;
use super::prompt_encoder::PromptEncoderWeights;
use anyhow::Result;
use rlx_flow::CompileProfile;
use rlx_runtime::Device;
use rlx_sam_ir::mask_prompt_ir::{MaskDownscaleWeights, PromptMaskCompiled};

pub struct SamPromptMaskCompiled(PromptMaskCompiled);

impl SamPromptMaskCompiled {
    pub fn compile(w: &PromptEncoderWeights, device: Device) -> Result<Self> {
        Self::compile_with_profile(w, device, &CompileProfile::sam_encoder())
    }

    pub fn compile_with_profile(
        w: &PromptEncoderWeights,
        device: Device,
        profile: &CompileProfile,
    ) -> Result<Self> {
        let in_h = 4 * SAM_EMBED_HW;
        let md = mask_weights(w);
        Ok(Self(PromptMaskCompiled::compile_with_profile(
            md, in_h, in_h, device, profile,
        )?))
    }

    pub fn run(&mut self, mask: &[f32], in_h: usize, in_w: usize) -> Result<Vec<f32>> {
        self.0.run(mask, in_h, in_w)
    }
}

fn mask_weights<'a>(w: &'a PromptEncoderWeights) -> MaskDownscaleWeights<'a> {
    MaskDownscaleWeights {
        mask_in_chans: w.mask_in_chans,
        embed_dim: w.embed_dim,
        mask_conv1_w: &w.mask_conv1_w,
        mask_conv1_b: &w.mask_conv1_b,
        mask_ln1_g: &w.mask_ln1_g,
        mask_ln1_b: &w.mask_ln1_b,
        mask_conv2_w: &w.mask_conv2_w,
        mask_conv2_b: &w.mask_conv2_b,
        mask_ln2_g: &w.mask_ln2_g,
        mask_ln2_b: &w.mask_ln2_b,
        mask_conv3_w: &w.mask_conv3_w,
        mask_conv3_b: &w.mask_conv3_b,
    }
}