mold-ai 0.1.0

Local AI image generation CLI — FLUX, SDXL, SD3.5, Z-Image diffusion models on your GPU
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
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134

use anyhow::Result;
use colored::Colorize;
use mold_core::config::Config;
use mold_core::download::DownloadError;
use mold_core::manifest::{find_manifest, known_manifests, resolve_model_name};

use crate::output::status;
use crate::AlreadyReported;

/// Download a model and write its config. Returns the updated Config.
pub async fn pull_and_configure(model: &str) -> Result<Config> {
    let canonical = resolve_model_name(model);

    // Pre-flight: print status and validate manifest exists (for CLI-specific error formatting)
    let manifest = match find_manifest(&canonical) {
        Some(m) => m,
        None => {
            print_unknown_model_error(model);
            return Err(AlreadyReported.into());
        }
    };

    let shared_gb = match manifest.family.as_str() {
        "sd15" => mold_core::manifest::SHARED_SD15_COMPONENTS_GB,
        "sdxl" => mold_core::manifest::SHARED_SDXL_COMPONENTS_GB,
        "z-image" => mold_core::manifest::SHARED_ZIMAGE_COMPONENTS_GB,
        _ => mold_core::manifest::SHARED_COMPONENTS_GB,
    };
    let total_gb = manifest.size_gb + shared_gb;
    status!(
        "{} Pulling {} ({:.1}GB transformer, {:.1}GB total with shared components)",
        "●".cyan(),
        manifest.name.bold(),
        manifest.size_gb,
        total_gb,
    );
    status!(
        "  {}",
        crate::output::colorize_description(&manifest.description)
    );
    status!("");

    // Delegate to core pull_and_configure
    let (config, _paths) = mold_core::download::pull_and_configure(model)
        .await
        .map_err(|e| -> anyhow::Error {
            match e {
                DownloadError::UnknownModel { .. } => {
                    print_unknown_model_error(model);
                }
                DownloadError::Unauthorized { repo, .. } => {
                    eprintln!();
                    eprintln!("{} Authentication required for {repo}", "✗".red().bold());
                    eprintln!();
                    eprintln!("  1. Create a token at: https://huggingface.co/settings/tokens");
                    eprintln!("     (select at least \"Read\" access)");
                    eprintln!("  2. Set: export HF_TOKEN=hf_...");
                    eprintln!("     Or run: huggingface-cli login");
                    eprintln!("  3. Retry: mold pull {}", canonical);
                    if std::env::var("HF_TOKEN").is_ok() {
                        eprintln!();
                        eprintln!(
                            "  {} HF_TOKEN is set but was rejected — it may be invalid or expired.",
                            "!".yellow().bold()
                        );
                    }
                }
                DownloadError::GatedModel { .. } => {
                    eprintln!();
                    eprintln!(
                        "{} This model requires access approval on HuggingFace.",
                        "✗".red().bold()
                    );
                    eprintln!();

                    let gated_repo = manifest
                        .files
                        .iter()
                        .find(|f| f.gated)
                        .map(|f| f.hf_repo.as_str())
                        .unwrap_or("the model repository");

                    eprintln!("  1. Visit: https://huggingface.co/{gated_repo}");
                    eprintln!("  2. Accept the license agreement");
                    eprintln!("  3. Create a token at: https://huggingface.co/settings/tokens");
                    eprintln!("  4. Set: export HF_TOKEN=hf_...");
                    eprintln!("  5. Retry: mold pull {}", canonical);
                }
                other => {
                    eprintln!();
                    eprintln!("{} Download failed: {other}", "✗".red().bold());
                }
            }
            AlreadyReported.into()
        })?;

    status!("");
    status!("{} {} is ready!", "✓".green().bold(), canonical.bold());

    Ok(config)
}

fn print_unknown_model_error(model: &str) {
    eprintln!("{} Unknown model: {}", "✗".red().bold(), model.bold());
    eprintln!();
    eprintln!("Available models:");
    let all = known_manifests();
    let nw = all.iter().map(|m| m.name.len()).max().unwrap_or(4) + 2;
    for m in &all {
        eprintln!(
            "  {:<nw$} {:>5.1}GB  {}",
            m.name.bold(),
            m.size_gb,
            crate::output::colorize_description(&m.description),
            nw = nw,
        );
    }
    eprintln!();
    eprintln!(
        "{}",
        format!(
            "Sizes are transformer only. First pull also downloads {:.1}GB of shared components (T5, CLIP, VAE).",
            mold_core::manifest::SHARED_COMPONENTS_GB
        )
        .dimmed(),
    );
    eprintln!("Usage: mold pull <model>");
}

pub async fn run(model: &str) -> Result<()> {
    pull_and_configure(model).await?;
    status!("  mold run \"your prompt\"");
    Ok(())
}