src/main.rs - master - rustwizard/etch

src/main.rs
96 строк3 KB
Rust Wizard
feat: Phase 1 memory optimizations — embedding cache, sequential TE, tiled VAE
16 июн 2026, 16:34
16 июн 2026, 16:3491830c6
#![deny(clippy::unwrap_used)]

mod cache;
mod cli;
mod device;
mod flux;
mod hub;
mod image;
mod logger;
mod lora;
mod pipeline;
mod progress;
mod schedulers;
mod sdxl;
mod vae_tiling;

use cli::Args;

use anyhow::Result;
use candle_core::{DType, Device};
use clap::Parser;
use tracing::info;

fn main() -> Result<()> {
    let args = Args::parse();
    anyhow::ensure!(
        !(args.seed.is_some() && args.seed_range.is_some()),
        "--seed and --seed-range are mutually exclusive"
    );

    let builder = tracing_subscriber::fmt().with_env_filter(
        tracing_subscriber::EnvFilter::try_from_default_env()
            .unwrap_or_else(|_| tracing_subscriber::EnvFilter::new("info")),
    );
    if args.verbose {
        builder.init();
    } else {
        builder.without_time().with_target(false).init();
    }

    let device = device::pick_device(args.cpu);
    info!("Device: {:?}", device);

    let (seeds, seed_source) = if let Some(ref range_str) = args.seed_range {
        (cli::parse_seed_range(range_str)?, "range")
    } else if args.seed.is_some() {
        (vec![args.seed.unwrap_or_else(rand::random)], "explicit")
    } else {
        (vec![rand::random::<u64>()], "auto")
    };
    info!("Generating {} image(s)", seeds.len());

    let dtype: DType = args.dtype.map(DType::from).unwrap_or_else(|| {
        if matches!(device, Device::Cpu) {
            DType::F32
        } else {
            DType::BF16
        }
    });
    info!("Dtype: {:?}", dtype);

    for seed in seeds {
        info!("--- Seed: {seed} ({seed_source}) ---");
        if let Err(e) = device.set_seed(seed) {
            tracing::warn!("Failed to set seed {seed}: {e}");
        }

        let output = cli::output_for_seed(&args.output, seed);
        if let Some(parent) = std::path::Path::new(&output).parent()
            && !parent.as_os_str().is_empty()
        {
            std::fs::create_dir_all(parent)?;
        }
        let iter_args = Args {
            seed: Some(seed),
            output: Some(output),
            ..args.clone()
        };

        let t0 = std::time::Instant::now();
        let result = pipeline::for_model(iter_args.model).run(&iter_args, &device, dtype);
        let out_path = iter_args.output.as_deref().expect("output set above");
        if let Err(e) = result {
            tracing::error!("Seed {seed} failed: {e}");
            if let Err(log_err) = logger::write_log_failure(out_path, &iter_args, seed, &e) {
                tracing::warn!("Failed to write failure log: {log_err}");
            }
            continue;
        }
        let elapsed = t0.elapsed().as_secs_f32();
        info!("Total time: {:.1}s", elapsed);
        logger::write_log_entry(out_path, &iter_args, seed, elapsed)?;
    }

    Ok(())
}