tensorneko 0.3.17__tar.gz → 0.3.19__tar.gz

{tensorneko-0.3.17 → tensorneko-0.3.19}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: tensorneko
-Version: 0.3.17
+Version: 0.3.19
 Summary: Tensor Neural Engine Kompanion. An util library based on PyTorch and PyTorch Lightning.
 Home-page: https://github.com/ControlNet/tensorneko
 Author: ControlNet

{tensorneko-0.3.17 → tensorneko-0.3.19}/src/tensorneko/io/reader.py

@@ -1,6 +1,9 @@
 from __future__ import annotations
-from typing import Type
 
+from pathlib import Path
+from typing import Type, Union, Any
+
+from tensorneko_util.io._path_conversion import _path2str
 from tensorneko_util.io.reader import Reader as BaseReader
 from .weight import WeightReader
 
@@ -28,3 +31,11 @@ class Reader(BaseReader):
             return self._mesh
         else:
             raise ImportError("To use the mesh reader, please install pytorch3d.")
+
+    def __call__(self, path: Union[str, Path], *args, **kwargs) -> Any:
+        """Automatically infer the file type and return the corresponding result. """
+        path = _path2str(path)
+        if path.endswith(".pt") or path.endswith(".pth") or path.endswith(".ckpt") or path.endswith(".safetensors"):
+            return self.weight(path, *args, **kwargs)
+        else:
+            return super().__call__(path, *args, **kwargs)

{tensorneko-0.3.17 → tensorneko-0.3.19}/src/tensorneko/io/weight/weight_reader.py

@@ -1,7 +1,9 @@
-from typing import OrderedDict
+from typing import OrderedDict, Union
+from pathlib import Path
 
 import torch
 
+from tensorneko_util.io._path_conversion import _path2str
 from ...util import Device
 
 
@@ -9,45 +11,48 @@ class WeightReader:
     """WeightReader for read model weights (checkpoints, state_dict, etc)."""
 
     @classmethod
-    def of_pt(cls, path: str, map_location: Device = "cpu") -> OrderedDict[str, torch.Tensor]:
+    def of_pt(cls, path: Union[str, Path], map_location: Device = "cpu") -> OrderedDict[str, torch.Tensor]:
         """
         Reads PyTorch model weights from a `.pt` or `.pth` file.
 
         Args:
-            path (``str``): The path of the `.pt` or `.pth` file.
+            path (``str`` | ``pathlib.Path``): The path of the `.pt` or `.pth` file.
             map_location (:class:`torch.device` | ``str``): The location to load the model weights. Default: "cpu"
 
         Returns:
             :class:`collections.OrderedDict`[``str``, :class:`torch.Tensor`]: The model weights.
         """
+        path = _path2str(path)
         return torch.load(path, map_location=map_location)
 
     @classmethod
-    def of_ckpt(cls, path: str, map_location: Device = "cpu") -> OrderedDict[str, torch.Tensor]:
+    def of_ckpt(cls, path: Union[str, Path], map_location: Device = "cpu") -> OrderedDict[str, torch.Tensor]:
         """
         Reads PyTorch model weights from a `.ckpt` file.
 
         Args:
-            path (``str``): The path of the `.ckpt` file.
+            path (``str`` | ``pathlib.Path``): The path of the `.ckpt` file.
             map_location (:class:`torch.device` | ``str``): The location to load the model weights. Default: "cpu"
 
         Returns:
             :class:`collections.OrderedDict`[``str``, :class:`torch.Tensor`]: The model weights.
         """
+        path = _path2str(path)
         return torch.load(path, map_location=map_location)["state_dict"]
 
     @classmethod
-    def of_safetensors(cls, path: str, map_location: str = "cpu") -> OrderedDict[str, torch.Tensor]:
+    def of_safetensors(cls, path: Union[str, Path], map_location: str = "cpu") -> OrderedDict[str, torch.Tensor]:
         """
         Reads model weights from a `.safetensors` file.
 
         Args:
-            path (``str``): The path of the `.safetensors` file.
+            path (``str`` | ``pathlib.Path``): The path of the `.safetensors` file.
             map_location (``str``): The location to load the model weights. Default: "cpu"
 
         Returns:
             :class:`collections.OrderedDict`[``str``, :class:`torch.Tensor`]: The model weights.
         """
+        path = _path2str(path)
         import safetensors
         from collections import OrderedDict
         tensors = OrderedDict()
@@ -57,18 +62,18 @@ class WeightReader:
         return tensors
 
     @classmethod
-    def of(cls, path: str, map_location: Device = "cpu") -> OrderedDict[str, torch.Tensor]:
+    def of(cls, path: Union[str, Path], map_location: Device = "cpu") -> OrderedDict[str, torch.Tensor]:
         """
         Reads model weights from a file.
 
         Args:
-            path (``str``): The path of the file.
+            path (``str`` | ``pathlib.Path``): The path of the file.
             map_location (:class:`torch.device` | ``str``): The location to load the model weights. Default: "cpu"
 
         Returns:
             :class:`collections.OrderedDict`[``str``, :class:`torch.Tensor`]: The model weights.
         """
-
+        path = _path2str(path)
         if path.endswith(".pt") or path.endswith(".pth"):
             return cls.of_pt(path, map_location)
         elif path.endswith(".ckpt"):
@@ -79,3 +84,8 @@ class WeightReader:
             return cls.of_safetensors(path, map_location)
         else:
             raise ValueError("Unknown file type. Supported types: .pt, .pth, .ckpt, .safetensors")
+
+    def __new__(cls, path: Union[str, Path], map_location: Device = "cpu") -> OrderedDict[str, torch.Tensor]:
+        """Alias of :meth:`~tensorneko.io.weight.weight_reader.WeightReader.of`."""
+        path = _path2str(path)
+        return cls.of(path, map_location)

{tensorneko-0.3.17 → tensorneko-0.3.19}/src/tensorneko/io/weight/weight_writer.py

@@ -1,48 +1,59 @@
-from typing import Dict
+from typing import Dict, Union
+from pathlib import Path
 
 import torch
 
+from tensorneko_util.io._path_conversion import _path2str
+from ...util import Device
 
 class WeightWriter:
     """WeightWriter for write model weights (checkpoints, state_dict, etc)."""
 
     @classmethod
-    def to_pt(cls, path: str, weights: Dict[str, torch.Tensor]) -> None:
+    def to_pt(cls, path: Union[str, Path], weights: Dict[str, torch.Tensor]) -> None:
         """
         Writes PyTorch model weights to a `.pt` or `.pth` file.
 
         Args:
-            path (``str``): The path of the `.pt` or `.pth` file.
+            path (``str`` | ``pathlib.Path``): The path of the `.pt` or `.pth` file.
             weights (:class:`collections.OrderedDict`[``str``, :class:`torch.Tensor`]): The model weights.
         """
+        path = _path2str(path)
         torch.save(weights, path)
 
     @classmethod
-    def to_safetensors(cls, path: str, weights: Dict[str, torch.Tensor]) -> None:
+    def to_safetensors(cls, path: Union[str, Path], weights: Dict[str, torch.Tensor]) -> None:
         """
         Writes model weights to a `.safetensors` file.
 
         Args:
-            path (``str``): The path of the `.safetensors` file.
+            path (``str`` | ``pathlib.Path``): The path of the `.safetensors` file.
             weights (:class:`collections.OrderedDict`[``str``, :class:`torch.Tensor`]): The model weights.
         """
+        path = _path2str(path)
         import safetensors.torch
         safetensors.torch.save_file(weights, path)
 
     @classmethod
-    def to(cls, path: str, weights: Dict[str, torch.Tensor]) -> None:
+    def to(cls, path: Union[str, Path], weights: Dict[str, torch.Tensor]) -> None:
         """
         Writes model weights to a file.
 
         Args:
-            path (``str``): The path of the file.
+            path (``str`` | ``pathlib.Path``): The path of the file.
             weights (:class:`collections.OrderedDict`[``str``, :class:`torch.Tensor`]): The model weights.
         """
+        path = _path2str(path)
         file_type = path.split(".")[-1]
 
-        if file_type == "pt":
+        if file_type in ("pt", "pth"):
             cls.to_pt(path, weights)
         elif file_type == "safetensors":
             cls.to_safetensors(path, weights)
         else:
             raise ValueError("Unknown file type. Supported types: .pt, .safetensors")
+
+    def __new__(cls, path: Union[str, Path], weights: Dict[str, torch.Tensor]) -> None:
+        """Alias of :meth:`~tensorneko.io.weight.weight_writer.WeightWriter.to`."""
+        path = _path2str(path)
+        return cls.to(path, weights)

{tensorneko-0.3.17 → tensorneko-0.3.19}/src/tensorneko/io/writer.py

@@ -1,7 +1,9 @@
 from __future__ import annotations
 
-from typing import Type
+from pathlib import Path
+from typing import Type, Union
 
+from tensorneko_util.io._path_conversion import _path2str
 from tensorneko_util.io.writer import Writer as BaseWriter
 from .weight import WeightWriter
 
@@ -29,3 +31,12 @@ class Writer(BaseWriter):
             return self._mesh
         else:
             raise ImportError("To use the mesh writer, please install pytorch3d.")
+
+    def __call__(self, path: Union[str, Path], obj, *args, **kwargs):
+        """Automatically infer the file type and return the corresponding result. """
+        path = _path2str(path)
+
+        if path.endswith(".pt") or path.endswith(".pth") or path.endswith(".ckpt") or path.endswith(".safetensors"):
+            return self.weight(path, obj)
+        else:
+            return super().__call__(path, obj, *args, **kwargs)

tensorneko-0.3.19/src/tensorneko/notebook/__init__.py

@@ -0,0 +1,7 @@
+from tensorneko_util.notebook import ipython_available
+
+__all__ = []
+
+if ipython_available:
+    from tensorneko_util.notebook import display, animation
+    __all__.extend(["display", "animation"])

{tensorneko-0.3.17 → tensorneko-0.3.19}/src/tensorneko/util/__init__.py

@@ -4,7 +4,8 @@ from tensorneko_util.util import dispatch, AverageMeter, tensorneko_util_path
 from tensorneko_util.util.fp import Seq, AbstractSeq, curry, F, Stream, return_option, Option, Monad, Eval, _, __
 from tensorneko_util.util import ref, Timer, Singleton
 from tensorneko_util.util.eventbus import Event, EventBus, EventHandler, subscribe
-from tensorneko_util.util import download_file, download_file_thread, download_files_thread, WindowMerger, Registry
+from tensorneko_util.util import download_file, download_file_thread, download_files_thread, WindowMerger, Registry, \
+    MultiLayerIndexer
 from . import type
 from .configuration import Configuration
 from .misc import reduce_dict_by, summarize_dict_by, with_printed_shape, is_bad_num, count_parameters, compose, \
@@ -78,5 +79,6 @@ __all__ = [
     "download_files_thread",
     "WindowMerger",
     "Registry",
-    "run_gc"
+    "run_gc",
+    "MultiLayerIndexer"
 ]

tensorneko-0.3.19/src/tensorneko/version.txt

@@ -0,0 +1 @@
+0.3.19

{tensorneko-0.3.17 → tensorneko-0.3.19}/src/tensorneko.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: tensorneko
-Version: 0.3.17
+Version: 0.3.19
 Summary: Tensor Neural Engine Kompanion. An util library based on PyTorch and PyTorch Lightning.
 Home-page: https://github.com/ControlNet/tensorneko
 Author: ControlNet

{tensorneko-0.3.17 → tensorneko-0.3.19}/src/tensorneko.egg-info/requires.txt

@@ -2,7 +2,7 @@ av>=8.0.3
 einops>=0.3.0
 numpy>=1.20.1
 pillow>=8.1
-tensorneko_util==0.3.17
+tensorneko_util==0.3.19
 torch>=1.9.0
 torchaudio>=0.9.0
 torchmetrics>=0.7.3

tensorneko-0.3.17/src/tensorneko/notebook/__init__.py

@@ -1,7 +0,0 @@
-from tensorneko_util.notebook import ipython_available
-
-__all__ = []
-
-if ipython_available:
-    from tensorneko_util.notebook import display
-    __all__.append("display")

tensorneko-0.3.17/src/tensorneko/version.txt

@@ -1 +0,0 @@
-0.3.17