tensorneko 0.3.17__py3-none-any.whl → 0.3.19__py3-none-any.whl

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/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/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/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/notebook/__init__.py

@@ -3,5 +3,5 @@ from tensorneko_util.notebook import ipython_available
 __all__ = []
 
 if ipython_available:
-    from tensorneko_util.notebook import display
-    __all__.append("display")
+    from tensorneko_util.notebook import display, animation
+    __all__.extend(["display", "animation"])

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/version.txt

@@ -1 +1 @@
-0.3.17
+0.3.19

{tensorneko-0.3.17.dist-info → tensorneko-0.3.19.dist-info}/METADATA

@@ -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
@@ -27,7 +27,7 @@ Requires-Dist: av (>=8.0.3)
 Requires-Dist: einops (>=0.3.0)
 Requires-Dist: numpy (>=1.20.1)
 Requires-Dist: pillow (>=8.1)
-Requires-Dist: tensorneko-util (==0.3.17)
+Requires-Dist: tensorneko-util (==0.3.19)
 Requires-Dist: torch (>=1.9.0)
 Requires-Dist: torchaudio (>=0.9.0)
 Requires-Dist: torchmetrics (>=0.7.3)

{tensorneko-0.3.17.dist-info → tensorneko-0.3.19.dist-info}/RECORD

@@ -2,7 +2,7 @@ tensorneko/__init__.py,sha256=uh1HNn1sNpX1bbOqAE_kNJfrH4eMtEzus0hO-Fh9tEw,990
 tensorneko/neko_model.py,sha256=QTbdOAg9ki0ix6mDL_Qu8Wmd5WJOoUFF3M1SXEp3KGc,10551
 tensorneko/neko_module.py,sha256=qELXvguSjWo_NvcRQibiFl0Qauzd9JWLSnT4dbGNS3Y,1473
 tensorneko/neko_trainer.py,sha256=JC8qoKSZ5ngz3grf3S0SjvIFVktDIP_GExth5aFfbGA,10074
-tensorneko/version.txt,sha256=sOiNkmvwnZw9dn4Od6oRM8jDNFGeaDoTmEVv7sd2w2k,6
+tensorneko/version.txt,sha256=xV2o_U6zdvujab1PTmKLrZ6dQjt1fW9CyXhxWiQJBpo,6
 tensorneko/arch/__init__.py,sha256=w4lTUeyBIZelrnSjlBFWUF0erzOmBFl9FqeWQuSOyKs,248
 tensorneko/arch/auto_encoder.py,sha256=j6PWWyaNYaYNtw_zZ9ikzhCASqe9viXR3JGBIXSK92Y,2137
 tensorneko/arch/binary_classifier.py,sha256=1MkEbReXKLdDksRG5Rsife40grJk08EVDcNKp54Xvb4,2316
@@ -34,14 +34,14 @@ tensorneko/evaluation/psnr.py,sha256=DeKxvY_xxawWMXHY0z3Nvbsi4dR57OUV4hjtUoCINXc
 tensorneko/evaluation/secs.py,sha256=D710GgcSxQgbGyPcWlC5ffF5n1GselLrUr5aA5Vq7oE,1622
 tensorneko/evaluation/ssim.py,sha256=6vPS4VQqoKxHOG49lChH51KxwNo07B4XHdhLub5DEPU,3758
 tensorneko/io/__init__.py,sha256=QEyA0mOC-BlKKskYYbDYttYWWRjCeh73lX-yKAUGNik,213
-tensorneko/io/reader.py,sha256=MSEfmzbRTk69qawTPrYruaKvQ8TXSyu5ZmFY7xVn-aU,773
-tensorneko/io/writer.py,sha256=MbO878ob24WSGKqjxyK-yRqp7xAiJBEtMxOGkOg_vOE,776
+tensorneko/io/reader.py,sha256=DSeTGLh84sFYwCwJmNTr-fGWkluudCbf7je29t0Z2U8,1303
+tensorneko/io/writer.py,sha256=BR_1h-wXekBdctXymJBU44HoWsKxPhbbh6N3AKYNkjE,1292
 tensorneko/io/mesh/__init__.py,sha256=cdR5QWNUgPaoU_fFcJO9sx7PeJy7pTlEvusjaivP1ok,72
 tensorneko/io/mesh/mesh_reader.py,sha256=ErUv9nBMARu-eR-uHlEhrb4bH0yl0cHiKoF9GSJ569A,184
 tensorneko/io/mesh/mesh_writer.py,sha256=d_lBhN2JEhaY79mwZALr-ylp0ZtyJmItCTUn_A3F4q0,184
 tensorneko/io/weight/__init__.py,sha256=zlUTKTYL7uhOxgyR3VpoR408pYjG0XHOGihGDL6mHyc,79
-tensorneko/io/weight/weight_reader.py,sha256=iDtphQsPQXzfKU_xaxvHJUEfealXHIcyoD2adXLpgg0,3079
-tensorneko/io/weight/weight_writer.py,sha256=06SCNPw3vDQ-6MLSpb4oR1AuOQ1slSpOextLoAl4PTA,1609
+tensorneko/io/weight/weight_reader.py,sha256=3W9MyULuYtY6rjBrQvJ0F5Lu0Yb0Wvha0aSebZis6eU,3688
+tensorneko/io/weight/weight_writer.py,sha256=yjtLXi5negrodfirEHzbjHrnZ2ALa6GZETBc1UxA5t4,2165
 tensorneko/layer/__init__.py,sha256=HHCBzwR-8UEVsrUKz8j_dhQSVGRASg0flLc_pyG1JN0,1120
 tensorneko/layer/aggregation.py,sha256=ykH6u-NLJx4Yesu_BLWa6T-vWIYFzJXJV1Txrbz3mPE,1177
 tensorneko/layer/attention.py,sha256=sNC6gZgZaeHqUAfzrPh0Vefp7T6nc1gFvrreuOL7wUg,6360
@@ -64,7 +64,7 @@ tensorneko/module/mlp.py,sha256=AFN6xmvlrNWOflLqVl-zVkoOJRZpYxYB4bnI10JG5CU,3361
 tensorneko/module/residual.py,sha256=S59TqiiD_310HQ3a6s3r49XY_7Dc4RGxONQtSvzEfN0,2958
 tensorneko/module/transformer.py,sha256=h4NvH3zGa0rZt0bv6e8VM31SimbQKRcocSR42zJYVoY,7602
 tensorneko/msg/__init__.py,sha256=GHrHjzw__0DcPBHBN6GzHrD8PD_7CwOZPRullOaZyW0,71
-tensorneko/notebook/__init__.py,sha256=Pgz4aTJg5_3zTzBIaML4LHAMOFgrBOHw3o00ZuvNuQU,171
+tensorneko/notebook/__init__.py,sha256=4cCi3ZyaX48hLDvJQqW0G3a4z_vdzmh_jtJ-Jzil4SM,197
 tensorneko/optim/__init__.py,sha256=89XjYQICij8SkrW5iryfZgmbxcTDxA3hhVZTgR4588o,33
 tensorneko/optim/lr_scheduler/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 tensorneko/preprocess/__init__.py,sha256=0Z0eA3_I2wphyyZlzZYRrx2muWTF0QMFq2Y-jh8oVKU,808
@@ -73,7 +73,7 @@ tensorneko/preprocess/enum.py,sha256=Wp5qFaUjea5XU4o3N0WxUd-qfzI-m5vr4ZWSqWjELb4
 tensorneko/preprocess/pad.py,sha256=b4IbbhGNRotZ7weZcKA7hfDqSixPo5KjM6khnqzaeUA,3238
 tensorneko/preprocess/resize.py,sha256=hitMlzVnN6n_8nEJwxy4C4ErZrTwpM86QGnYewsrmf8,3469
 tensorneko/preprocess/face_detector/__init__.py,sha256=_ktIfUZqGTX0hk7RBgKf-zHwG2n9KRH4RS7rjuOI8Bo,262
-tensorneko/util/__init__.py,sha256=1ygIyTqZm1ROCsgQJ9_df5TOezyHx08stMDGjtyMxhE,2208
+tensorneko/util/__init__.py,sha256=39G34a2k5ktVtBAh4N4RMsePEak5rzPDhbNNdXo-Ye4,2258
 tensorneko/util/configuration.py,sha256=xXeAjDh1FCNTmSPwDdkL-uH-ULfzFF6Fg0LT7gsZ6nQ,2510
 tensorneko/util/dispatched_misc.py,sha256=_0Go7XezdYB7bpMnCs1MDD_6mPNoWP5qt8DoKuPxynI,997
 tensorneko/util/gc.py,sha256=P3bOZ-2VUNyswnfVz5xfj__ecTSAHpu_kLp2wFcpb6M,185
@@ -86,8 +86,8 @@ tensorneko/visualization/log_graph.py,sha256=NvOwWVc_petXWYdgaHosPFLa43sHBeacbYc
 tensorneko/visualization/matplotlib.py,sha256=xs9Ssc44ojZX65QU8-fftA7Ug_pBuZ3TBtM8vETNq9w,1568
 tensorneko/visualization/image_browser/__init__.py,sha256=AtykhAE3bXQS6SOWbeYFeeUE9ts9XOFMvrL31z0LoMg,63
 tensorneko/visualization/watcher/__init__.py,sha256=Nq752qIYvfRUZ8VctKQRSqhxh5KmFbWcqPfZlijVx6s,379
-tensorneko-0.3.17.dist-info/LICENSE,sha256=Vd75kwgJpVuMnCRBWasQzceMlXt4YQL13ikBLy8G5h0,1067
-tensorneko-0.3.17.dist-info/METADATA,sha256=cezBL0TH1cdY0Sct7GCGLK9jW0ckUXZIL7PRScTckBg,19998
-tensorneko-0.3.17.dist-info/WHEEL,sha256=g4nMs7d-Xl9-xC9XovUrsDHGXt-FT0E17Yqo92DEfvY,92
-tensorneko-0.3.17.dist-info/top_level.txt,sha256=sZHwlP0iyk7_zHuhRHzSBkdY9yEgyC48f6UVuZ6CvqE,11
-tensorneko-0.3.17.dist-info/RECORD,,
+tensorneko-0.3.19.dist-info/LICENSE,sha256=Vd75kwgJpVuMnCRBWasQzceMlXt4YQL13ikBLy8G5h0,1067
+tensorneko-0.3.19.dist-info/METADATA,sha256=9M81K0ENnhym0S-qyKQTlQzSADnpOZIszNWJ66q4IuM,19998
+tensorneko-0.3.19.dist-info/WHEEL,sha256=g4nMs7d-Xl9-xC9XovUrsDHGXt-FT0E17Yqo92DEfvY,92
+tensorneko-0.3.19.dist-info/top_level.txt,sha256=sZHwlP0iyk7_zHuhRHzSBkdY9yEgyC48f6UVuZ6CvqE,11
+tensorneko-0.3.19.dist-info/RECORD,,