Rework weight conversions

CONTRIBUTING.md

@@ -1,10 +1,14 @@
-We are happy to accept contributions from everyone. Feel free to browse our [bounty list](https://www.finegrain.ai/bounties) to find a task you would like to work on.
+We are happy to accept contributions from everyone.
+Feel free to browse our [bounty list](https://www.finegrain.ai/bounties) to find a task you would like to work on.
 
 This document describes the process for contributing to Refiners.
 
 ## Licensing
 
-Refiners is a library that is freely available to use and modify under the MIT License. It's essential to exercise caution when using external code, as any code that can affect the licensing of Refiners, including proprietary code, should not be copied and pasted. It's worth noting that some open-source licenses can also be problematic. For instance, we'd need to redistribute an Apache 2.0 license if you're using code from an Apache 2.0 codebase.
+Refiners is a library that is freely available to use and modify under the MIT License.
+It's essential to exercise caution when using external code, as any code that can affect the licensing of Refiners, including proprietary code, should not be copied and pasted.
+It's worth noting that some open-source licenses can also be problematic.
+For instance, we'd need to redistribute an Apache 2.0 license if you're using code from an Apache 2.0 codebase.
 
 ## Design principles
 
@@ -17,73 +21,108 @@ We do not enforce strict rules on the design of the code, but we do have a few g
 
 ## Setting up your environment
 
-We use [Rye](https://rye-up.com/guide/installation/) to manage our development environment. Please follow the instructions on the Rye website to install it.
+We use [Rye](https://rye-up.com/guide/installation/) to manage our development environment.
+Please follow the instructions on the Rye website to install it.
 
 Once Rye is installed, you can clone the repository and run `rye sync` to install the dependencies.
 
-## Linting
+You should regularly update your Rye version by running `rye self update`.
 
-We use the standard integration of [ruff](https://docs.astral.sh/ruff/) in Rye to lint and format our code. You can lint your code by running:
+## Linting
 
+We use the standard integration of [ruff](https://docs.astral.sh/ruff/) in Rye to lint and format our code.
+You can lint your code by running:
 ```bash
 rye fmt
 rye lint --fix
 ```
 
-We also enforce strict type checking with [pyright](https://github.com/microsoft/pyright). You can run the type checker with:
+You should also check for typos, by running:
+```bash
+rye run typos
+```
 
+We also enforce strict type checking with [pyright](https://github.com/microsoft/pyright). You can run the type checker with:
 ```bash
 rye run pyright
 ```
 
-## Running the tests
+## Getting the weights
 
-Running end-to-end tests is pretty compute-intensive, and you must convert all the model weights to the correct format before you can run them.
+Since refiners re-implements models using fluxion, we can't directly use the weights from their original implementations, we need to convert them to the correct format.
+We provide already converted weights for some models on our huggingface organization: https://huggingface.co/refiners
 
-First, install test dependencies with:
+If you need to convert the weights yourself (e.g. for running the tests), you may use the `get_weights` command (c.f. `project.scripts` in `pyproject.toml`) to convert the weights. This command will automatically download all original weights and convert them.
 
-```bash
-rye sync --all-features
-```
+### Contributing a new model conversion
 
-Then, download and convert all the necessary weights. Be aware that this will use around 100 GB of disk space:
+1. Add a new file in the `refiners.conversion.models` module, if necessary.
+2. Instantiate a new `WeightRecipe` object to declare how to translate the keys from the original weights to the refiners keys.
+3. Instantiate a new `Conversion` object to relate the original weights, the converted weights and the recipe.
+4. Update the associated `__init__.py` and `cli.py` files.
+
+Note: The conversion process is not always straightforward, we would prefer if you could use the above steps (since it's the most declarative way to convert the weights),
+although alternative methods are acceptable if they are well documented (you may find some examples in the existing conversion scripts).
+
+## Running the tests
 
+Running end-to-end tests is pretty compute-intensive.
+
+To test everything, you can either:
+- Use `REFINERS_USE_LOCAL_WEIGHTS=1` and manually convert all the model weights to the correct format before running the tests. See the above section to learn how to convert the weights.
+- Use `REFINER_USE_LOCAL_WEIGHTS=0` and automatically download the weights from the huggingface hub. This is the default behavior. Though since some weights aren't available on the hub, some tests may be skipped.
+
+> [!WARNING]
+> GPU tests are notoriously hard to reproduce across different hardware. (see https://pytorch.org/docs/stable/notes/randomness.html)
+> Our tests are designed to work on a single NVIDIA GeForce RTX 3090 24GB, with nvidia drivers v545.23.08 and cuda v12.3.
+> If you have a different GPU/drivers, the tests may fail or give different results.
+
+First, install test dependencies with:
 ```bash
-python scripts/prepare_test_weights.py
+rye sync --all-features
 ```
 
 Some tests require cloning the original implementation of the model as they use `torch.hub.load`:
-
 ```bash
 git clone git@github.com:facebookresearch/dinov2.git tests/repos/dinov2
+git clone git@github.com:microsoft/Swin-Transformer.git tests/repos/Swin-Transformer
 ```
 
 Finally, run the tests:
-
 ```bash
 rye run pytest
 ```
 
 The `-k` option is handy to run a subset of tests that match a given expression, e.g.:
-
 ```bash
-rye run pytest -k diffusion_std_init_image
+rye run pytest -k "diffusion_std_init_image" -v
+rye run pytest -k "not e2e" -v
+rye run pytest -k "e2e" -v
 ```
 
-You can enforce running tests on CPU. Tests that require a GPU will be skipped.
+Some tests require a GPU, and may sometime OOM if you test a lot of models sequentially.
+You can re-run the failed tests with:
+```bash
+pytest -k "e2e" -v --last-failed
+```
 
+You can modify the following environment variables to change the behavior of the tests:
 ```bash
-REFINERS_TEST_DEVICE=cpu rye run pytest
+export REFINERS_USE_LOCAL_WEIGHTS=1  # checkpoints will be loaded from the local hub, instead of the hf hub
+export REFINERS_TEST_DEVICE=cpu  # tests that require a GPU will be skipped
+export REFINERS_HUB_PATH=/path/to/hub  # default is ./tests/weights
+export REFINERS_TEST_DATASETS_DIR=/path/to/datasets  # default it ./tests/datasets
+export REFINERS_TEST_REPOS_DIR=/path/to/repos  # default is ./tests/repos
 ```
 
-You can collect [code coverage](https://github.com/nedbat/coveragepy) data while running tests with, e.g.:
+### Code coverage
 
+You can collect [code coverage](https://github.com/nedbat/coveragepy) data while running tests with, e.g.:
 ```bash
 rye run test-cov
 ```
 
 Then, browse the corresponding HTML report with:
-
 ```bash
 rye run serve-cov-report
 ```

pyproject.toml

@@ -56,6 +56,9 @@ Documentation = "https://refine.rs/"
 Repository = "https://github.com/finegrain-ai/refiners"
 Issues = "https://github.com/finegrain-ai/refiners/issues"
 
+[project.scripts]
+get_weights = "refiners.conversion.cli:main"
+
 [project.optional-dependencies]
 training = [
     "bitsandbytes>=0.41.2.post2",
@@ -75,23 +78,21 @@ test = [
     "transformers>=4.35.2",
     "piq>=0.8.0",
     "torchvision>=0.16.1",
-    # An unofficial Python package for Meta AI's Segment Anything Model:
-    # https://github.com/opengeos/segment-anything
+    # An unofficial Python package for Meta AI's Segment Anything Model: https://github.com/opengeos/segment-anything
     "segment-anything-py>=1.0",
-    # Official Python package for HQ-SAM
+    # Official Python package for HQ-SAM: https://github.com/SysCV/sam-hq
     "segment-anything-hq>=0.3",
-    # HQ-SAM missing dependency:
-    # https://github.com/SysCV/sam-hq/pull/59
+    # HQ-SAM missing dependency: https://github.com/SysCV/sam-hq/pull/59
     "timm>=0.5.0",
     "sentencepiece>=0.2.0",
 ]
 conversion = [
+    "huggingface-hub>=0.25.1",
     "diffusers>=0.26.1",
     "transformers>=4.35.2",
     "segment-anything-py>=1.0",
     "requests>=2.26.0",
     "tqdm>=4.62.3",
-    "gdown>=5.2.0",
 ]
 doc = [
     # required by mkdocs to format the signatures
@@ -129,8 +130,7 @@ allow-direct-references = true
 [tool.rye.scripts]
 serve-docs = "mkdocs serve"
 test-cov = "coverage run -m pytest"
-# Work around for "Couldn't parse" errors due to e.g. opencv-python:
-# https://github.com/nedbat/coveragepy/issues/1653
+# Work around for "Couldn't parse" errors due to e.g. opencv-python: https://github.com/nedbat/coveragepy/issues/1653
 build-html-cov = { cmd = "coverage html", env = { PYTHONWARNINGS = "ignore:Couldn't parse::coverage.report_core" } }
 serve-cov-report = { chain = [
     "build-html-cov",
@@ -164,7 +164,7 @@ combine-as-imports = true
 black = true
 
 [tool.pyright]
-include = ["src/refiners", "tests", "scripts"]
+include = ["src/refiners", "tests"]
 strict = ["*"]
 exclude = ["**/__pycache__", "tests/weights", "tests/repos"]
 reportMissingTypeStubs = "warning"
@@ -185,7 +185,7 @@ exclude_also = [
 ]
 
 [tool.typos.default]
-extend-ignore-identifiers-re = ["NDArray*", "interm", "af000ded"]
+extend-ignore-identifiers-re = ["NDArray*", "interm", "af000ded", "nin"]
 
 [tool.typos.default.extend-words]
 adaptee = "adaptee"  # Common name for an adapter's target

requirements.lock

@@ -29,8 +29,6 @@ attrs==24.2.0
     # via referencing
 babel==2.16.0
     # via mkdocs-material
-beautifulsoup4==4.12.3
-    # via gdown
 bitsandbytes==0.44.0
     # via refiners
 black==24.8.0
@@ -72,7 +70,6 @@ exceptiongroup==1.2.2
 filelock==3.16.1
     # via datasets
     # via diffusers
-    # via gdown
     # via huggingface-hub
     # via torch
     # via transformers
@@ -88,8 +85,6 @@ fsspec==2024.6.1
     # via torch
 future==1.0.0
     # via neptune
-gdown==5.2.0
-    # via refiners
 ghp-import==2.1.0
     # via mkdocs
 gitdb==4.0.11
@@ -285,8 +280,6 @@ pyjwt==2.9.0
 pymdown-extensions==10.10.2
     # via mkdocs-material
     # via mkdocstrings
-pysocks==1.7.1
-    # via requests
 pytest==8.3.3
     # via pytest-rerunfailures
 pytest-rerunfailures==14.0
@@ -328,7 +321,6 @@ requests==2.32.3
     # via bravado-core
     # via datasets
     # via diffusers
-    # via gdown
     # via huggingface-hub
     # via mkdocs-material
     # via neptune
@@ -376,8 +368,6 @@ six==1.16.0
     # via rfc3339-validator
 smmap==5.0.1
     # via gitdb
-soupsieve==2.6
-    # via beautifulsoup4
 swagger-spec-validator==3.0.4
     # via bravado-core
     # via neptune
@@ -406,7 +396,6 @@ torchvision==0.19.1
     # via timm
 tqdm==4.66.5
     # via datasets
-    # via gdown
     # via huggingface-hub
     # via refiners
     # via transformers