Skip to main content

rlx_pkg/
from_dduf.rs

1// RLX — versatile ML compiler + runtime.
2// Copyright (C) 2026 Eugene Hauptmann, Nataliya Kosmyna.
3// SPDX-License-Identifier: MIT OR Apache-2.0
4
5//! Import HuggingFace DDUF (`.dduf`) into an `.rlxp` package.
6
7use crate::tier::StorageTier;
8use crate::write::{ContainerKind, PackedWeight, WriteOptions, write_package};
9use anyhow::{Context, Result};
10use rlx_dduf::visit_f32_tensors;
11use rlx_ir::{DType, Graph, Shape};
12use std::path::Path;
13
14/// Options for [`dduf_to_rlxp`].
15#[derive(Debug, Clone)]
16pub struct DdufImportOptions {
17    pub container: ContainerKind,
18    pub include_graph: bool,
19    pub compress_sidecars: bool,
20    pub auto_tier: bool,
21}
22
23impl Default for DdufImportOptions {
24    fn default() -> Self {
25        Self {
26            container: ContainerKind::Flat,
27            include_graph: false,
28            compress_sidecars: true,
29            auto_tier: true,
30        }
31    }
32}
33
34/// Convert `model.dduf` → `.rlxp` (weights + optional stub graph).
35///
36/// Streams one safetensors ZIP member at a time so peak memory is roughly one
37/// member + the growing `PackedWeight` list (not the full ZIP decoded at once).
38pub fn dduf_to_rlxp(
39    dduf_path: impl AsRef<Path>,
40    out: impl AsRef<Path>,
41    opts: &DdufImportOptions,
42) -> Result<()> {
43    let dduf_path = dduf_path.as_ref();
44    let mut weights = Vec::new();
45    let meta = visit_f32_tensors(dduf_path, |t| {
46        let mut bytes = Vec::with_capacity(t.data.len() * 4);
47        for x in &t.data {
48            bytes.extend_from_slice(&x.to_le_bytes());
49        }
50        weights.push(PackedWeight {
51            name: t.name,
52            shape: t.shape,
53            scheme: "f32".into(),
54            layout: "row_major".into(),
55            data: bytes,
56            rank: None,
57            tier: StorageTier::Hot,
58        });
59        Ok(())
60    })
61    .with_context(|| format!("open dduf {}", dduf_path.display()))?;
62
63    let mut sidecars = Vec::new();
64    if let Some(idx) = &meta.model_index {
65        sidecars.push((
66            "model_index.json".into(),
67            "application/json".into(),
68            serde_json::to_vec_pretty(idx)?,
69        ));
70    }
71    for (comp, cfg) in &meta.configs {
72        sidecars.push((
73            format!("{comp}/config.json"),
74            "application/json".into(),
75            serde_json::to_vec_pretty(cfg)?,
76        ));
77    }
78
79    let mut wopts = WriteOptions {
80        name: dduf_path
81            .file_stem()
82            .and_then(|s| s.to_str())
83            .unwrap_or("dduf")
84            .into(),
85        producer: Some("rlx-pkg/dduf_import".into()),
86        container: opts.container,
87        compress_sidecars: opts.compress_sidecars,
88        include_graph: opts.include_graph,
89        write_checksums: true,
90        ..WriteOptions::default()
91    };
92    wopts.sidecars.extend(sidecars);
93
94    if opts.auto_tier {
95        crate::auto_tier::apply_auto_tier(&mut weights, &Default::default());
96    }
97
98    let graph = if opts.include_graph {
99        let mut g = Graph::new(&wopts.name);
100        let s = Shape::new(&[1], DType::F32);
101        let x = g.input("x", s);
102        g.set_outputs(vec![x]);
103        g
104    } else {
105        Graph::new(&wopts.name)
106    };
107
108    write_package(out.as_ref(), &graph, &weights, &wopts)
109        .with_context(|| format!("write {}", out.as_ref().display()))?;
110    Ok(())
111}