1use 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#[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
34pub 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}