use _ensure_local for arch_file

Author: FynnBe

bioimageio/core/build_spec/build_model.py

@@ -61,16 +61,10 @@ def _get_pytorch_state_dict_weight_kwargs(architecture, model_kwargs, root):
     tmp_archtecture = None
     weight_kwargs = {"kwargs": model_kwargs} if model_kwargs else {}
     if ":" in architecture:
-        arch_file, callable_name = architecture.replace("::", ":").split(":")
-
-        # this goes haywire if we pass an absolute path, so need to copt to a tmp relative path
-        if os.path.isabs(arch_file):
-            tmp_archtecture = Path("this_model_architecture.py")
-            copyfile(arch_file, root / tmp_archtecture)
-            arch = ImportableSourceFile(callable_name, tmp_archtecture)
-        else:
-            arch = ImportableSourceFile(callable_name, Path(arch_file))
-
+        # note: path itself might include : for absolute paths in windows
+        *arch_file_parts, callable_name = architecture.replace("::", ":").split(":")
+        arch_file = _ensure_local(":".join(arch_file_parts), root)
+        arch = ImportableSourceFile(callable_name, arch_file)
         arch_hash = _get_hash(root / arch.source_file)
         weight_kwargs["architecture_sha256"] = arch_hash
     else:
@@ -123,30 +117,21 @@ def _get_weights(
         if tensorflow_version is None:
             raise ValueError("tensorflow_version needs to be passed for building a keras model")
         weights = model_spec.raw_nodes.KerasHdf5WeightsEntry(
-            source=weight_source,
-            sha256=weight_hash,
-            tensorflow_version=tensorflow_version,
-            **attachments,
+            source=weight_source, sha256=weight_hash, tensorflow_version=tensorflow_version, **attachments
         )
 
     elif weight_type == "tensorflow_saved_model_bundle":
         if tensorflow_version is None:
             raise ValueError("tensorflow_version needs to be passed for building a tensorflow model")
         weights = model_spec.raw_nodes.TensorflowSavedModelBundleWeightsEntry(
-            source=weight_source,
-            sha256=weight_hash,
-            tensorflow_version=tensorflow_version,
-            **attachments,
+            source=weight_source, sha256=weight_hash, tensorflow_version=tensorflow_version, **attachments
         )
 
     elif weight_type == "tensorflow_js":
         if tensorflow_version is None:
             raise ValueError("tensorflow_version needs to be passed for building a tensorflow_js model")
         weights = model_spec.raw_nodes.TensorflowJsWeightsEntry(
-            source=weight_source,
-            sha256=weight_hash,
-            tensorflow_version=tensorflow_version,
-            **attachments,
+            source=weight_source, sha256=weight_hash, tensorflow_version=tensorflow_version, **attachments
         )
 
     elif weight_type in weight_types:
@@ -519,9 +504,8 @@ def _ensure_local_or_url(source: Union[Path, URI, str, list], root: Path) -> Uni
         return [_ensure_local_or_url(s, root) for s in source]
 
     local_source = resolve_local_source(source, root)
-    local_source = resolve_local_source(
-        local_source, root, None if isinstance(local_source, URI) else root / local_source.name
-    )
+    if not isinstance(local_source, URI):
+        local_source = resolve_local_source(local_source, root, root / local_source.name)
     return local_source.relative_to(root)
 
 
@@ -654,6 +638,7 @@ def build_model(
             Only requred for models with onnx weight format.
         weight_kwargs: additional keyword arguments for this weight type.
     """
+    assert architecture is None or isinstance(architecture, str)
     if root is None:
         root = "."
     root = Path(root)