Make write_safetensors work with shared tensors

Author: mittagessen

kraken/models/__init__.py

@@ -2,3 +2,4 @@
 from .writers import *  # NOQA
 from .loaders import *  # NOQA
 from .utils import *  # NOQA
+from .convert import *  # NOQA

kraken/models/convert.py

@@ -0,0 +1,71 @@
+import importlib
+
+from pathlib import Path
+from collections.abc import Iterable
+from typing import TYPE_CHECKING, Union
+from kraken.models.loaders import load_models
+
+if TYPE_CHECKING:
+    from os import PathLike
+
+__all__ = ['convert_models']
+
+
+def convert_models(paths: Iterable[Union[str, 'PathLike']],
+                   output: Union[str, 'PathLike'],
+                   weights_format: str = 'safetensors') -> 'PathLike':
+    """
+    Converts the models in a set of checkpoint or weights files into a single
+    output weights file.
+
+    It accepts checkpoints and weights files interchangeably for all supported
+    formats and models.
+
+    This function has a number of uses:
+
+        * it can be used to convert checkpoints into weights.
+
+          convert_models(['model.ckpt'], 'model.safetensors')
+
+        * it can be used to convert multiple related models into a single
+          weights file for joint inference:
+
+          convert_models(['blla_line.ckpt', 'blla_region.ckpt'], 'model.safetensors')
+
+        * it can convert models between coreml and safetensors formats:
+
+          convert_models(['blla.mlmodel'], 'blla.safetensors')
+
+    Args:
+        paths: Paths to checkpoint or weights files.
+        output: Output path to the combined/converted file. The actual output
+                path may be modified.
+        weights_format: Serialization format to write the weights to.
+
+    Returns:
+        The path the actual weights file was written to.
+    """
+    try:
+        (entry_point,) = importlib.metadata.entry_points(group='kraken.writers', name=weights_format)
+        writer = entry_point.load()
+    except ValueError:
+        raise ValueError('No writer for format {weights_format} found.')
+
+    def _find_module(path):
+        for entry_point in importlib.metadata.entry_points(group='kraken.lightning_modules'):
+            module = entry_point.load()
+            try:
+                return module.load_from_checkpoint(path)
+            except ValueError:
+                continue
+        raise ValueError(f'No lightning module found for checkpoint {path}')
+
+    models = []
+    for ckpt in paths:
+        ckpt = Path(ckpt)
+        if ckpt.suffix == '.ckpt':
+            models.append(_find_module(ckpt).net)
+        else:
+            models.extend(load_models(ckpt))
+
+    return writer(models, output)

kraken/models/utils.py

@@ -1,18 +1,15 @@
 import importlib
 
-from pathlib import Path
-from collections.abc import Iterable
-from typing import Union, TYPE_CHECKING
-from kraken.models.base import BaseModel
-from kraken.models.loaders import load_models
+from typing import TYPE_CHECKING
 
-__all__ = ['create_model', 'convert_models']
+
+__all__ = ['create_model']
 
 if TYPE_CHECKING:
-    from os import PathLike
+    from .base import BaseModel
 
 
-def create_model(name, *args, **kwargs) -> BaseModel:
+def create_model(name, *args, **kwargs) -> 'BaseModel':
     """
     Constructs an empty model from the model registry.
     """
@@ -26,63 +23,3 @@ def create_model(name, *args, **kwargs) -> BaseModel:
 
     cls = entry_point.load()
     return cls(*args, **kwargs)
-
-
-def convert_models(paths: Iterable[Union[str, 'PathLike']],
-                   output: Union[str, 'PathLike'],
-                   weights_format: str = 'safetensors') -> 'PathLike':
-    """
-    Converts the models in a set of checkpoint or weights files into a single
-    output weights file.
-
-    It accepts checkpoints and weights files interchangeably for all supported
-    formats and models.
-
-    This function has a number of uses:
-
-        * it can be used to convert checkpoints into weights.
-
-          convert_models(['model.ckpt'], 'model.safetensors')
-
-        * it can be used to convert multiple related models into a single
-          weights file for joint inference:
-
-          convert_models(['blla_line.ckpt', 'blla_region.ckpt'], 'model.safetensors')
-
-        * it can convert models between coreml and safetensors formats:
-
-          convert_models(['blla.mlmodel'], 'blla.safetensors')
-
-    Args:
-        paths: Paths to checkpoint or weights files.
-        output: Output path to the combined/converted file. The actual output
-                path may be modified.
-        weights_format: Serialization format to write the weights to.
-
-    Returns:
-        The path the actual weights file was written to.
-    """
-    try:
-        (entry_point,) = importlib.metadata.entry_points(group='kraken.writers', name=weights_format)
-        writer = entry_point.load()
-    except ValueError:
-        raise ValueError('No writer for format {weights_format} found.')
-
-    def _find_module(path):
-        for entry_point in importlib.metadata.entry_points(group='kraken.lightning_modules'):
-            module = entry_point.load()
-            try:
-                return module.load_from_checkpoint(path)
-            except ValueError:
-                continue
-        raise ValueError(f'No lightning module found for checkpoint {path}')
-
-    models = []
-    for ckpt in paths:
-        ckpt = Path(ckpt)
-        if ckpt.suffix == '.ckpt':
-            models.append(_find_module(ckpt).net)
-        else:
-            models.extend(load_models(ckpt))
-
-    return writer(models, output)

kraken/models/writers.py

@@ -9,6 +9,7 @@
 import logging
 import importlib
 
+from torch import nn
 from os import PathLike
 from pathlib import Path
 from typing import Union, TYPE_CHECKING
@@ -44,21 +45,17 @@ def write_safetensors(objs: list[BaseModel], path: Union[str, PathLike]) -> Path
     """
     Writes a set of models as a safetensors.
     """
-    from safetensors.torch import save_file
+    from safetensors.torch import save_model
     # assign unique prefixes to each model in model list
-    prefixes = {str(uuid.uuid4()): model for model in objs}
-    metadatas = {k: {'_kraken_min_version': v._kraken_min_version,
-                     '_tasks': v.model_type,
-                     '_model': v.__class__.__name__,
-                     **v.user_metadata} for k, v in prefixes.items()}
-
-    weights = {}
-    for prefix, model in prefixes.items():
-        for name in (state_dict := model.state_dict()):
-            weights[f'{prefix}.{name}'] = state_dict[name]
-    save_file(weights,
-              filename=path,
-              metadata={'kraken_meta': json.dumps(metadatas)})
+    prefixes = nn.ModuleDict({str(uuid.uuid4()): model for model in objs})
+    metadata = {k: {'_kraken_min_version': v._kraken_min_version,
+                    '_tasks': v.model_type,
+                    '_model': v.__class__.__name__,
+                    **v.user_metadata} for k, v in prefixes.items()}
+
+    save_model(prefixes,
+               filename=path,
+               metadata={'kraken_meta': json.dumps(metadata)})
     return Path(path)