From d1de2a67403261418c21e75f85038684aaee829d Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E6=BA=90=E6=96=87=E9=9B=A8?= <41315874+fumiama@users.noreply.github.com> Date: Thu, 18 Apr 2024 18:21:01 +0900 Subject: [PATCH] feat: add cmdline `-c` & optimize log printing --- README.md | 31 ++++++++++++++++++++++++++- cfg.go | 60 +++++++++++++++++++++++++++++++++-------------------- http.go | 4 ++-- main.go | 9 ++++---- progress.go | 42 +++++++++++++++++++++++++++++++++++++ 5 files changed, 116 insertions(+), 30 deletions(-) create mode 100644 progress.go diff --git a/README.md b/README.md index 2efa401..58e4053 100644 --- a/README.md +++ b/README.md @@ -16,10 +16,39 @@ rvcmd packs/general/latest ```bash rvcmd tools/ffmpeg ``` - +## Customized Download +### Download ffmpeg Tools & Latest Intel Pack +1. Write and save the following `cust.yaml`. + ```yaml + BaseURL: https://huggingface.co/lj1995/VoiceConversionWebUI/resolve/main + Targets: + - Refer: tools/ffmpeg + - Refer: packs/intel/latest + ``` +2. Run `rvcmd` in the same folder. + ```bash + rvcmd -c cust + ``` +### Download other Repositories +> Use [Stable Diffusion v1-5](https://huggingface.co/runwayml/stable-diffusion-v1-5) as the example. +1. Write and save the following `cust.yaml`. + ```yaml + BaseURL: https://huggingface.co/runwayml/stable-diffusion-v1-5/resolve/main + Targets: + - Folder: sd1.5 # the folder you want to download into + Copy: # files to download + - v1-5-pruned-emaonly.ckpt + - v1-5-pruned-emaonly.safetensors + - vae/diffusion_pytorch_model.bin + ``` +2. Run `rvcmd` in the same folder. + ```bash + rvcmd -c cust + ``` ## Full Usage ```bash Usage: rvcmd [-notrs] [-dns dns.yaml] 'target/to/download' + -c use custom yaml instruction -dns string custom dns.yaml -notrs diff --git a/cfg.go b/cfg.go index 129dd26..ba2b7da 100644 --- a/cfg.go +++ b/cfg.go @@ -9,6 +9,7 @@ import ( "net/http" "os" "runtime" + "strconv" "strings" "sync" "time" @@ -32,9 +33,12 @@ var cfg = func() *zip.Reader { return r }() -func readconfig(path string) (c config, err error) { +func readconfig(path string, usecust bool) (c config, err error) { fname := path + ".yaml" f, err := cfg.Open(fname) + if usecust && err != nil { + f, err = os.Open(fname) + } if err != nil { err = errors.Wrap(err, "invalid path") return @@ -61,81 +65,91 @@ type targets struct { Arch string `yaml:"Arch"` } -func (c *config) download(path string) error { +func (c *config) download(path, prefix string, usecust bool) error { for i, t := range c.Targets { if t.Refer != "" { refp := path[:strings.LastIndex(path, "/")+1] + t.Refer - logrus.Infof("#%d refer to target '%s'", i+1, refp) - refcfg, err := readconfig(refp) + logrus.Infof("#%s%d refer to target '%s'.", prefix, i+1, refp) + refcfg, err := readconfig(refp, usecust) if err != nil { return err } - err = refcfg.download(refp) + err = refcfg.download(refp, prefix+strconv.Itoa(i+1)+".", usecust) if err != nil { return err } continue } if t.OS != "" && t.OS != runtime.GOOS { - logrus.Warnf("#%d target required OS: %s but you are %s, skip.", i+1, t.OS, runtime.GOOS) + logrus.Warnf("#%s%d target required OS: %s but you are %s, skip.", prefix, i+1, t.OS, runtime.GOOS) continue } if t.Arch != "" && t.Arch != runtime.GOARCH { - logrus.Warnf("#%d target required Arch: %s but you are %s, skip.", i+1, t.Arch, runtime.GOARCH) + logrus.Warnf("#%s%d target required Arch: %s but you are %s, skip.", prefix, i+1, t.Arch, runtime.GOARCH) continue } err := os.MkdirAll(t.Folder, 0755) if err != nil { - return errors.Wrap(err, fmt.Sprintf("#%d make target folder '%s'", i+1, t.Folder)) + return errors.Wrap(err, fmt.Sprintf("#%s%d make target folder '%s'", prefix, i+1, t.Folder)) } - logrus.Infof("#%d open target folder '%s'", i+1, t.Folder) + logrus.Infof("#%s%d open target folder '%s'.", prefix, i+1, t.Folder) if len(t.Copy) == 0 { - logrus.Warningf("#%d empty copy target", i+1) + logrus.Warningf("#%s%d empty copy target.", prefix, i+1) continue } wg := sync.WaitGroup{} wg.Add(len(t.Copy)) - logrus.Infof("#%d download copy: %v", i+1, t.Copy) - for i, cp := range t.Copy { - go func(i int, cp string) { + logrus.Infof("#%s%d download copy: '%v'.", prefix, i+1, t.Copy) + for j, cp := range t.Copy { + go func(i int, cp, prefix string) { defer wg.Done() + if strings.Contains(cp, "/") { // have innner folder + infldr := t.Folder + "/" + cp[:strings.LastIndex(cp, "/")] + err := os.MkdirAll(infldr, 0755) + if err != nil { + logrus.Errorf("#%s%d make target inner folder '%s' err: %v", prefix, i+1, t.Folder, err) + return + } + logrus.Infof("#%s%d make target inner folder '%s'.", prefix, i+1, t.Folder) + } sleep := time.Millisecond * 100 * time.Duration(i) if sleep > time.Millisecond { time.Sleep(sleep) } req, err := http.NewRequest("GET", c.BaseURL+"/"+cp, nil) if err != nil { - logrus.Errorln("new request to", cp, "err:", err) + logrus.Errorf("#%s%d new request to %s err: %v", prefix, i+1, cp, err) return } - logrus.Infoln("get:", req.URL) + logrus.Infof("#%s%d get: %s", prefix, i+1, req.URL) req.Header.Add("user-agent", ua) resp, err := cli.Do(req) if err != nil { - logrus.Errorln("get", req.URL, "err:", err) + logrus.Errorf("#%s%d get %s err: %v", prefix, i+1, req.URL, err) return } defer resp.Body.Close() if resp.StatusCode != http.StatusOK { err := errors.New(fmt.Sprintf("HTTP %d %s", resp.StatusCode, resp.Status)) - logrus.Errorln("get", req.URL, "err:", err) + logrus.Errorf("#%s%d get %s err: %v", prefix, i+1, req.URL, err) return } fname := t.Folder + "/" + cp f, err := os.Create(fname) if err != nil { - logrus.Errorln("create file", fname, "err:", err) + logrus.Errorf("#%s%d create file %s err: %v", prefix, i+1, fname, err) return } - logrus.Infoln("writing file", fname) + logrus.Infof("#%s%d writing file %s", prefix, i+1, fname) defer f.Close() - _, err = io.Copy(f, resp.Body) + pm := newmeter(fmt.Sprintf("#%s%d", prefix, i+1), fname, int(resp.ContentLength)) + _, err = io.Copy(io.MultiWriter(f, &pm), resp.Body) if err != nil { - logrus.Errorln("download file", fname, "err:", err) + logrus.Errorf("#%s%d download file %s err: %v", prefix, i+1, fname, err) return } - logrus.Infoln("downloaded file", fname) - }(i, cp) + logrus.Infof("#%s%d finished download %s", prefix, i+1, fname) + }(j, cp, fmt.Sprintf("%s%d.", prefix, i+1)) } wg.Wait() } diff --git a/http.go b/http.go index 8dde730..ba5a26b 100644 --- a/http.go +++ b/http.go @@ -17,7 +17,7 @@ import ( ) var ( - ErrEmptyHostAddress = errors.New("empty host addr") + errEmptyHostAddress = errors.New("empty host addr") ) var httpdialer = net.Dialer{ @@ -57,7 +57,7 @@ var cli = http.Client{ lookupTable.Set(host, addrs) } if len(addr) == 0 { - return nil, ErrEmptyHostAddress + return nil, errEmptyHostAddress } var tlsConn *tls.Conn for _, a := range addrs { diff --git a/main.go b/main.go index 2449dcf..42ddf65 100644 --- a/main.go +++ b/main.go @@ -21,10 +21,11 @@ const ua = "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 ( func main() { logrus.Infoln("RVC Models Downloader start at", time.Now().Local().Format(time.DateTime+" (MST)")) - logrus.Infof("operating system: %s, architecture: %s\n", runtime.GOOS, runtime.GOARCH) + logrus.Infof("operating system: %s, architecture: %s", runtime.GOOS, runtime.GOARCH) logrus.Infoln("can use ipv6:", canUseIPv6.Get()) ntrs := flag.Bool("notrs", false, "use standard TLS client") dnsf := flag.String("dns", "", "custom dns.yaml") + cust := flag.Bool("c", false, "use custom yaml instruction") flag.Parse() args := flag.Args() if len(args) != 1 { @@ -58,15 +59,15 @@ func main() { } fmt.Println("custom dns file added") } - usercfg, err := readconfig(args[0]) + usercfg, err := readconfig(args[0], *cust) if err != nil { logrus.Errorln(err) return } - err = usercfg.download(args[0]) + err = usercfg.download(args[0], "", *cust) if err != nil { logrus.Errorln(err) return } - logrus.Info("download tasks finished") + logrus.Info("all download tasks finished.") } diff --git a/progress.go b/progress.go new file mode 100644 index 0000000..3578bec --- /dev/null +++ b/progress.go @@ -0,0 +1,42 @@ +package main + +import ( + "errors" + "io" + + "github.com/sirupsen/logrus" +) + +var ( + errZeroMeterSize = errors.New("zero meter size") +) + +type progressmeter struct { + prefix string + name string + size int + prgs int + lstp int + io.Writer +} + +func newmeter(prefix, name string, size int) (pm progressmeter) { + pm.prefix = prefix + pm.name = name + pm.size = size + return +} + +func (pm *progressmeter) Write(p []byte) (n int, err error) { + if pm.size == 0 { + return 0, errZeroMeterSize + } + pm.prgs += len(p) + percent := pm.prgs * 100 / pm.size + if percent == pm.lstp { + return len(p), nil + } + logrus.Infof("%s [%2d%%] %s\t(%d/%dMB)", pm.prefix, percent, pm.name, pm.prgs/1024/1024, pm.size/1024/1024) + pm.lstp = percent + return len(p), nil +}