Skip to main content

zsfm/
main.rs

1mod chronos;
2mod common;
3mod flowstate;
4mod generic;
5mod lag_llama;
6mod mitra;
7mod moirai;
8mod moirai2;
9mod moment;
10mod sundial;
11mod tabdpt;
12mod tabfm;
13mod tabicl;
14mod tabpfn;
15mod timesfm;
16mod tirex;
17mod toto;
18mod ttm;
19
20use clap::{Parser, Subcommand};
21
22#[derive(Parser)]
23#[command(
24    name = "zsfm",
25    about = "Zero-shot foundation models — GGUF conversion + inference CLI"
26)]
27struct Cli {
28    /// Silence status lines on stderr (JSON on stdout is never silenced).
29    /// Same as setting `ZSFM_QUIET=1`.
30    #[arg(long, global = true)]
31    quiet: bool,
32    /// Verbose status on stderr (default already logs to stderr; reserved for
33    /// extra per-tensor detail in future).
34    #[arg(long, global = true)]
35    verbose: bool,
36    #[command(subcommand)]
37    model: ModelCommand,
38}
39
40#[derive(Subcommand)]
41enum ModelCommand {
42    /// Toto-2 (Datadog) zero-shot time-series forecaster.
43    Toto {
44        #[command(subcommand)]
45        command: toto::Command,
46    },
47    /// TabFM (google/tabfm-1.0.0-pytorch) zero-shot tabular classifier/regressor.
48    Tabfm {
49        #[command(subcommand)]
50        command: tabfm::Command,
51    },
52    /// Mitra (autogluon/mitra-classifier, autogluon/mitra-regressor) zero-shot tabular
53    /// foundation model.
54    Mitra {
55        #[command(subcommand)]
56        command: mitra::Command,
57    },
58    /// TabDPT (Layer6/TabDPT) zero-shot tabular foundation model.
59    Tabdpt {
60        #[command(subcommand)]
61        command: tabdpt::Command,
62    },
63    /// TabICL (jingang/TabICL, v2) zero-shot tabular foundation model (classification only).
64    Tabicl {
65        #[command(subcommand)]
66        command: tabicl::Command,
67    },
68    /// TabPFN-3 (Prior-Labs/tabpfn_3) zero-shot tabular foundation model (classification only).
69    /// NON-COMMERCIAL WEIGHTS LICENSE — see `zsfm tabpfn infer --help`.
70    Tabpfn {
71        #[command(subcommand)]
72        command: tabpfn::Command,
73    },
74    /// Chronos-2 (Amazon) zero-shot time-series forecaster.
75    Chronos {
76        #[command(subcommand)]
77        command: chronos::Command,
78    },
79    /// TimesFM 2.5 200M (Google) zero-shot time-series forecaster.
80    Timesfm {
81        #[command(subcommand)]
82        command: timesfm::Command,
83    },
84    /// Sundial / Timer v3 (thuml) zero-shot time-series forecaster.
85    Sundial {
86        #[command(subcommand)]
87        command: sundial::Command,
88    },
89    /// TinyTimeMixer / TTM-R2 (IBM Granite) zero-shot time-series forecaster.
90    Ttm {
91        #[command(subcommand)]
92        command: ttm::Command,
93    },
94    /// Lag-Llama (ServiceNow) zero-shot time-series forecaster.
95    LagLlama {
96        #[command(subcommand)]
97        command: lag_llama::Command,
98    },
99    /// MOMENT-1-large (CMU) zero-shot time-series forecaster.
100    Moment {
101        #[command(subcommand)]
102        command: moment::Command,
103    },
104    /// Moirai-1.0-R-large (Salesforce) zero-shot time-series forecaster.
105    Moirai {
106        #[command(subcommand)]
107        command: moirai::Command,
108    },
109    /// Moirai-2.0-R-small (Salesforce) zero-shot time-series forecaster.
110    Moirai2 {
111        #[command(subcommand)]
112        command: moirai2::Command,
113    },
114    /// FlowState-R1 (IBM Granite) zero-shot time-series forecaster.
115    Flowstate {
116        #[command(subcommand)]
117        command: flowstate::Command,
118    },
119    /// TiRex (NX-AI) zero-shot time-series forecaster.
120    Tirex {
121        #[command(subcommand)]
122        command: tirex::Command,
123    },
124    /// Convert ANY checkpoint (safetensors / PyTorch pickle / npy / npz / ONNX /
125    /// HDF5 / Keras / GGUF) to GGUF, tensor names passed through. GGUF input
126    /// re-quantizes.
127    Convert(generic::ConvertArgs),
128    /// List tensors (and optionally values) of any supported checkpoint format.
129    Inspect(generic::InspectArgs),
130    /// Download a HuggingFace repo's checkpoint in whatever format it's
131    /// published in (safetensors / pickle / npy / npz / ONNX / HDF5 / Keras /
132    /// GGUF), without converting.
133    Pull(generic::PullArgs),
134}
135
136#[tokio::main]
137async fn main() -> anyhow::Result<()> {
138    let cli = Cli::parse();
139    if cli.quiet {
140        // SAFETY: single-threaded startup before any async work spawns threads.
141        unsafe { std::env::set_var("ZSFM_QUIET", "1") };
142    }
143    match cli.model {
144        ModelCommand::Toto { command } => toto::run(command).await,
145        ModelCommand::Tabfm { command } => tabfm::run(command).await,
146        ModelCommand::Mitra { command } => mitra::run(command).await,
147        ModelCommand::Tabdpt { command } => tabdpt::run(command).await,
148        ModelCommand::Tabicl { command } => tabicl::run(command).await,
149        ModelCommand::Tabpfn { command } => tabpfn::run(command).await,
150        ModelCommand::Chronos { command } => chronos::run(command).await,
151        ModelCommand::Timesfm { command } => timesfm::run(command).await,
152        ModelCommand::Sundial { command } => sundial::run(command).await,
153        ModelCommand::Ttm { command } => ttm::run(command).await,
154        ModelCommand::LagLlama { command } => lag_llama::run(command).await,
155        ModelCommand::Moment { command } => moment::run(command).await,
156        ModelCommand::Moirai { command } => moirai::run(command).await,
157        ModelCommand::Moirai2 { command } => moirai2::run(command).await,
158        ModelCommand::Flowstate { command } => flowstate::run(command).await,
159        ModelCommand::Tirex { command } => tirex::run(command).await,
160        ModelCommand::Convert(args) => generic::run_convert(args).await,
161        ModelCommand::Inspect(args) => generic::run_inspect(args),
162        ModelCommand::Pull(args) => generic::run_pull(args).await,
163    }
164}