Clean up weights in cwd in add_weights

Author: constantinpape

bioimageio/core/build_spec/add_weights.py

@@ -63,6 +63,13 @@ def add_weights(
     except Exception as e:
         raise e
     finally:
+        # clean up tmp files
+        os.remove(weight_out)
         if tmp_arch is not None:
             os.remove(tmp_arch)
+        # for some reason the weights are also copied to the cwd.
+        # not sure why this happens, but it needs to be cleaned up, unless these are the input weigths
+        weights_cwd = Path(os.path.split(weight_uri)[1])
+        if weights_cwd.exists() and weights_cwd.absolute() != Path(weight_uri).absolute():
+            os.remove(weights_cwd)
     return model

tests/build_spec/test_add_weights.py

@@ -1,3 +1,4 @@
+import os
 from bioimageio.core import export_resource_package, load_raw_resource_description, load_resource_description
 from bioimageio.core.resource_tests import test_model as _test_model
 
@@ -33,6 +34,9 @@ def _test_add_weights(model, tmp_path, base_weights, added_weights, **kwargs):
     test_res = _test_model(out_path)
     assert test_res["error"] is None
 
+    # make sure the weights were cleaned from the cwd
+    assert not os.path.exists(os.path.split(weight_path)[1])
+
 
 def test_add_torchscript(unet2d_nuclei_broad_model, tmp_path):
     _test_add_weights(unet2d_nuclei_broad_model, tmp_path, "pytorch_state_dict", "torchscript")