Fix(BMZ): Relax model output validation kwargs; extract weights and c…

…onfig file following new `spec` and `core` release (#279)

### Description

- **What**: Relaxing the model output validation kwargs, both absolute
and relative tolerance, from the default, `1e-4`, to `1e-2`.
- **Why**: The defaults are pretty strict and some of our uploaded
models are stuck in pending because of slightly mismatching input and
outputs.
- e.g. (Actually maybe absolute tolerance should be put to 0, otherwise
it still might not pass after this PR)
  ```console
  Output and expected output disagree:
    Not equal to tolerance rtol=0.0001, atol=0.00015
  Mismatched elements: 40202 / 1048576 (3.83%)
  Max absolute difference: 0.1965332
  Max relative difference: 0.0003221
  ```
- **How**: In the model description config param, added the new test
kwargs. Additionally, updated `bmz_export` so that the test kwargs in
the model description are used during model testing at export time.

### Changes Made

- **Modified**: Describe existing features or files modified.
  - `create_model_description`: added test_kwargs to config param
- `export_to_bmz`: use test_kwargs in model description for model
testing at export time.

### Related Issues

- Resolves 
  - last checkbox in #278 

EDIT:
This PR also fixes loading from BMZ following an incompatible release of
`bioimageio/core` (`0.7.0`) and `bioimageio/spec` (`0.5.3.5`). The
problem was `load_model_description` no longer unzips the archive file
but only streams the `rdf.yaml` file data. This means we have to now
extract the weights and careamics config from the zip to load them,
which can be done using
`bioimageio.spec._internal.io.resolve_and_extract`

---

**Please ensure your PR meets the following requirements:**

- [x] Code builds and passes tests locally, including doctests
- [ ] New tests have been added (for bug fixes/features)
- [x] Pre-commit passes
- [ ] PR to the documentation exists (for bug fixes / features)

---------

Co-authored-by: Joran Deschamps <[email protected]>

src/careamics/model_io/bioimage/model_description.py

@@ -4,6 +4,7 @@
 from typing import List, Optional, Tuple, Union
 
 import numpy as np
+from bioimageio.spec._internal.io import resolve_and_extract
 from bioimageio.spec.model.v0_5 import (
     ArchitectureFromLibraryDescr,
     Author,
@@ -280,6 +281,16 @@ def create_model_description(
             "https://careamics.github.io/latest/",
         ],
         license="BSD-3-Clause",
+        config={
+            "bioimageio": {
+                "test_kwargs": {
+                    "pytorch_state_dict": {
+                        "absolute_tolerance": 1e-2,
+                        "relative_tolerance": 1e-2,
+                    }
+                }
+            }
+        },
         version="0.1.0",
         weights=weights_descr,
         attachments=[FileDescr(source=config_path)],
@@ -304,15 +315,17 @@ def extract_model_path(model_desc: ModelDescr) -> tuple[Path, Path]:
     """
     if model_desc.weights.pytorch_state_dict is None:
         raise ValueError("No model weights found in model description.")
-    weights_path = model_desc.weights.pytorch_state_dict.download().path
+    weights_path = resolve_and_extract(
+        model_desc.weights.pytorch_state_dict.source
+    ).path
 
     for file in model_desc.attachments:
         file_path = file.source if isinstance(file.source, Path) else file.source.path
         if file_path is None:
             continue
         file_path = Path(file_path)
         if file_path.name == "careamics.yaml":
-            config_path = file.download().path
+            config_path = resolve_and_extract(file.source).path
             break
     else:
         raise ValueError("Configuration file not found.")

src/careamics/model_io/bmz_io.py

@@ -21,7 +21,6 @@
     create_env_text,
     create_model_description,
     extract_model_path,
-    get_unzip_path,
 )
 
 
@@ -185,7 +184,12 @@ def export_to_bmz(
         )
 
         # test model description
-        summary: ValidationSummary = test_model(model_description)
+        test_kwargs = (
+            model_description.config.get("bioimageio", {})
+            .get("test_kwargs", {})
+            .get("pytorch_state_dict", {})
+        )
+        summary: ValidationSummary = test_model(model_description, **test_kwargs)
         if summary.status == "failed":
             raise ValueError(f"Model description test failed: {summary}")
 
@@ -219,14 +223,9 @@ def load_from_bmz(
     # load description, this creates an unzipped folder next to the archive
     model_desc = load_model_description(path)
 
-    # extract relative paths
+    # extract paths
     weights_path, config_path = extract_model_path(model_desc)
 
-    # create folder path and absolute paths
-    unzip_path = get_unzip_path(path)
-    weights_path = unzip_path / weights_path
-    config_path = unzip_path / config_path
-
     # load configuration
     config = load_configuration(config_path)