musicgpt 0.3.1

Generate music samples from natural language prompt locally with your own computer
use std::fmt::Write;
use std::ops::{Deref, DerefMut};
use std::time::Duration;

use indicatif::{MultiProgress, ProgressBar, ProgressDrawTarget, ProgressState, ProgressStyle};

pub struct LoadingBarFactor;

pub struct Bar(ProgressBar);

impl Bar {
    pub fn update_elapsed_total(&self, elapsed: usize, total: usize) {
        self.0.set_length(total as u64);
        self.0.set_position(elapsed as u64);
    }
    
    pub fn into_update_callback(self) ->  Box<dyn Fn(usize, usize) + Send + Sync + 'static> {
        Box::new(move |el, t|  self.update_elapsed_total(el, t))
    }
}

impl Deref for Bar {
    type Target = ProgressBar;

    fn deref(&self) -> &Self::Target {
        &self.0
    }
}

impl DerefMut for Bar {
    fn deref_mut(&mut self) -> &mut Self::Target {
        &mut self.0
    }
}

pub struct MultiBar(MultiProgress);

impl MultiBar {
    pub fn add(&self, bar: Bar) -> Bar {
        Bar(self.0.add(bar.0))
    }
}

impl Deref for MultiBar {
    type Target = MultiProgress;

    fn deref(&self) -> &Self::Target {
        &self.0
    }
}

impl DerefMut for MultiBar {
    fn deref_mut(&mut self) -> &mut Self::Target {
        &mut self.0
    }
}

impl LoadingBarFactor {
    pub fn multi() -> MultiBar {
        MultiBar(MultiProgress::new())
    }

    pub fn spinner(msg: impl Into<String>) -> Bar {
        let pb = ProgressBar::new_spinner();
        pb.enable_steady_tick(Duration::from_millis(120));
        pb.set_style(ProgressStyle::with_template("{spinner:.blue} {msg}").unwrap());
        pb.set_message(msg.into());
        Bar(pb)
    }

    pub fn bar(prefix: impl Into<String>) -> Bar {
        Self::fixed_bar(prefix, 1)
    }

    pub fn fixed_bar(prefix: impl Into<String>, len: usize) -> Bar {
        let pb = ProgressBar::new(len as u64);
        pb.set_style(
            ProgressStyle::with_template(
                &(prefix.into()
                    + " {spinner:.green} [{elapsed_precise}] [{wide_bar:.cyan/blue}] ({eta})"),
            )
            .unwrap()
            .with_key("eta", |state: &ProgressState, w: &mut dyn Write| {
                write!(w, "{:.1}s", state.eta().as_secs_f64()).unwrap()
            })
            .progress_chars("#>-"),
        );
        Bar(pb)
    }

    pub fn download_bar(file: &str) -> Bar {
        const NAME_LEN: usize = 32;
        const NAME_SHIFT_INTERVAL: usize = 300;
        let pb = ProgressBar::with_draw_target(None, ProgressDrawTarget::stderr());
        let file_string = file.to_string();
        pb.set_style(
            ProgressStyle::with_template(
                "{file:>32} {spinner:.green} [{wide_bar:.cyan/blue}] {bytes}/{total_bytes}",
            )
            .unwrap()
            .with_key("file", move |state: &ProgressState, w: &mut dyn Write| {
                if file_string.len() > NAME_LEN {
                    let el = state.elapsed().as_millis() as usize;
                    let offset = (el / NAME_SHIFT_INTERVAL) % (file_string.len() - NAME_LEN + 1);
                    let view = &file_string[offset..offset + NAME_LEN];
                    write!(w, "{view: >w$}", w = NAME_LEN).unwrap();
                } else {
                    write!(w, "{file_string: >w$}", w = NAME_LEN).unwrap();
                }
            })
            .with_key("eta", |state: &ProgressState, w: &mut dyn Write| {
                write!(w, "{:.1}s", state.eta().as_secs_f64()).unwrap()
            })
            .progress_chars("#>-"),
        );
        Bar(pb)
    }
}