VAIMM - An improved models downloader

Topaz Video AI General
1

First off, I appreciate that Suraj provided the shell scripts for batch-downloading of models.
However I found a number of issues with them:

  1. Slow. Each network (model file) requires a separate launch of ffmpeg, with all the latency that spinning up the VAI engine requires. Took me hours and hundreds of VAI launches to get a few models downloaded .
  2. Unreliable. Sometimes VAI crashed during the matrix downloading, and I had no idea if the model files that were intended for downloading had been downloaded or not. So I had to run the script several times, and pray that I got all the files. I still didn’t know.
  3. Not comprehensive. Even after running the scripts, I noticed that VAI still wanted to download even more files. I attribute this to the specific set of resolutions requested in the script, which do not seem comprehensive.

To address all these issues, I created my own improved version, that is about two orders of magnitude faster (provided you have the bandwidth), downloads all model files required (to my knowledge), and is meticulous about handling failure. Also allows abort and restart from where you left off, without having to worry about you corrupting any model files. It will only download model files you don’t have, so you don’t have to figure that bit out yourself either.

It has a set of additional features as well:

  • Ability to choose only the specific subset of model(s) and versions you are interested in. E.g. prob-3 and amq-13 for onnx16 or any other combination you like.
  • Lists all the models and all versions of them with information about what they do, in an easy to view list.

For more details, view the comprehensive documentation on the PyPi repo, or the github project that’s backing the release. Give it a shot if you like. Testing, feedback and issue reporting are all welcome.

To install:

pip install vaimm

To use, for example:

vaimm download --backend onnx --include prob-3,amq-13 --dir . --cookie "$COOKIE"

Downloading: 100%|████████████████| 54/54 [00:18<00:00,  3.01 files /s, speed=93.12 MiB/s]
Download completed successfully with 1.63 GiB of data fetched.

Links:

Note: the package is automatically built from source in github, and pushed from there to PyPi. So, you can validate the code and the build steps before you install it, if you happen to be as paranoid about third party software as I am.

6 Likes
2

Tip, you may want to reduce the concurrency a bit, since cloudflare may temporarily block you if you go full throttle. By default four parallel downloads are launched, just like the browser default. However launching single-threaded (--threads 1) may be a safer option if you have a very “wide pipe”.

Will take longer, but at least you won’t be temp-blocked by CF. I’ll see if I can fix this, provided CF responds with something useful in the throttling message. If they do I’ll just throttle the client to that level to make everyone happy :slight_smile:

Update: Fixed now in 0.3.0, sufficiently if not optimally.
Default is now single threaded, and if you crank up the thread count, that shouldn’t pose a CF blocking issue since the retry backoff has been tuned to be a lot more aggressive. If you hit the rate limit, a message will be logged in the console making you aware of that fact. It’s an indication that you may want to reduce the thread count a bit, or just ignore it. Up to you, either works.

Unfortunately Cloudflare does not follow the HTTP standard recommendation when rate-limiting, in that they don’t advertise how long the client should wait before retrying. Had they done that, vaimm would have been able to download at the fastest possible speed. But since they don’t, the speed is perhaps only 80-90% of the optimal throughput when you have a 1-10 gbps internet connection.

I could try to infer the optimal throughput in the program, but for now I don’t think the time investment is worth it. Works well enough as-is, so am going to consider the “net-code” done for now.

Another minor improvement in terms of information. It seems topaz may have forgotten to upload a few model file elements. I noticed 3 out of ~600 files were missing on the CDN. Since the goal is to ensure we have all the model files available offline, knowing which ones are missing is still important. For that reason I added a “failed” stats counter on the progress bar, and a final report that will list which specific files were not available from the CDN once the download jobs complete.
(Might be useful for the topaz team, in case these missing files are actually a bug in the model publication process)

To install the latest version if you have an older one, do: pip install --upgrade vaimm

3

Thank you for creating this and posting.
I just find one issue with the code for downloading TensorRT model.
[R] is TensorRT version and [C] is the CUDA compute level of the GPU.
There are some more changes to the json structure coming in the future. Try out the new Alpha and look at the iris-1 json. Message me if you need clarifications.

2 Likes
4

Indeed. PM:ed you to get these sorted as well.

Edit: seems the only change on alpha is a bit of reshuffling of the element order. No semantic changes as far as I can see. At least none which affect this program. No change needed for the upcoming alpha, since the few things I parse from the files haven’t changed.

That just leaves sorting out the tensorrt handling.

5

Currently the [R] option has only 2 possible values. It is 8500 for windows and 8517 for linux.
As far as [C] is concerned it depends on the Nvidia GPU on the machine.
705 for RTX 20 series
806 for RTX 30 series
809 for RTX 40 series

Let me know if you have more questions.

1 Like
6

Now that I got an RTX card on both a linux and a windows machine, I was able to complete the tensorrt download support. New version on pypi.