tensorneko 0.3.8__py3-none-any.whl → 0.3.9__py3-none-any.whl

tensorneko/__init__.py

@@ -1,11 +1,11 @@
 import os.path
 
-from tensorneko_util import io
 from . import backend
 from . import callback
 from . import dataset
 from . import debug
 from . import evaluation
+from . import io
 from . import layer
 from . import module
 from . import notebook
@@ -13,10 +13,13 @@ from . import optim
 from . import preprocess
 from . import util
 from . import visualization
+from .io import read, write
 from .neko_model import NekoModel
 from .neko_module import NekoModule
 from .neko_trainer import NekoTrainer
 
+__version__ = io.read.text(os.path.join(util.get_tensorneko_path(), "version.txt"))
+
 __all__ = [
     "callback",
     "dataset",
@@ -33,7 +36,9 @@ __all__ = [
     "debug",
     "NekoModel",
     "NekoTrainer",
-    "NekoModule"
+    "NekoModule",
+    "read",
+    "write",
 ]
 
-__version__ = io.read.text(os.path.join(util.get_tensorneko_path(), "version.txt"))
+

tensorneko/backend/__init__.py

@@ -1,8 +1,10 @@
-from tensorneko_util.backend import parallel, run_blocking, VisualLib, AudioLib
+from tensorneko_util.backend import parallel, run_blocking, VisualLib, AudioLib, import_tqdm_auto, import_tqdm
 
 __all__ = [
     "parallel",
     "run_blocking",
     "VisualLib",
     "AudioLib",
+    "import_tqdm_auto",
+    "import_tqdm",
 ]

tensorneko/evaluation/fid.py

@@ -12,12 +12,13 @@ from torchmetrics.image.fid import FrechetInceptionDistance
 from torchvision.transforms.functional import resize
 
 from tensorneko_util.backend import VisualLib
-from tensorneko_util.backend._tqdm import import_tqdm_auto
+from tensorneko_util.backend.tqdm import import_tqdm_auto
 
 try:
     from typing import Literal
     TypeOption = Literal["video", "image"]
 except ImportError:
+    Literal = None
     TypeOption = str
 
 

tensorneko/io/reader.py

@@ -2,6 +2,7 @@ from __future__ import annotations
 from typing import Type
 
 from tensorneko_util.io.reader import Reader as BaseReader
+from .weight import WeightReader
 
 try:
     from .mesh import MeshReader
@@ -15,6 +16,7 @@ class Reader(BaseReader):
 
     def __init__(self):
         super().__init__()
+        self.weight = WeightReader
         self._mesh = None
 
     @property

tensorneko/io/weight/__init__.py

@@ -0,0 +1,2 @@
+from .weight_reader import WeightReader
+from .weight_writer import WeightWriter

tensorneko/io/weight/weight_reader.py

@@ -0,0 +1,81 @@
+from typing import OrderedDict
+
+import torch
+
+from ...util import Device
+
+
+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]:
+        """
+        Reads PyTorch model weights from a `.pt` or `.pth` file.
+
+        Args:
+            path (``str``): 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.
+        """
+        return torch.load(path, map_location=map_location)
+
+    @classmethod
+    def of_ckpt(cls, path: str, 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.
+            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.
+        """
+        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]:
+        """
+        Reads model weights from a `.safetensors` file.
+
+        Args:
+            path (``str``): 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.
+        """
+        import safetensors
+        from collections import OrderedDict
+        tensors = OrderedDict()
+        with safetensors.safe_open(path, framework="pt", device=map_location) as f:
+            for key in f.keys():
+                tensors[key] = f.get_tensor(key)
+        return tensors
+
+    @classmethod
+    def of(cls, path: str, map_location: Device = "cpu") -> OrderedDict[str, torch.Tensor]:
+        """
+        Reads model weights from a file.
+
+        Args:
+            path (``str``): 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.
+        """
+
+        if path.endswith(".pt") or path.endswith(".pth"):
+            return cls.of_pt(path, map_location)
+        elif path.endswith(".ckpt"):
+            return cls.of_ckpt(path, map_location)
+        elif path.endswith(".safetensors"):
+            if isinstance(map_location, torch.device):
+                map_location = str(map_location)
+            return cls.of_safetensors(path, map_location)
+        else:
+            raise ValueError("Unknown file type. Supported types: .pt, .pth, .ckpt, .safetensors")

tensorneko/io/weight/weight_writer.py

@@ -0,0 +1,48 @@
+from typing import Dict
+
+import torch
+
+
+class WeightWriter:
+    """WeightWriter for write model weights (checkpoints, state_dict, etc)."""
+
+    @classmethod
+    def to_pt(cls, path: str, 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.
+            weights (:class:`collections.OrderedDict`[``str``, :class:`torch.Tensor`]): The model weights.
+        """
+        torch.save(weights, path)
+
+    @classmethod
+    def to_safetensors(cls, path: str, weights: Dict[str, torch.Tensor]) -> None:
+        """
+        Writes model weights to a `.safetensors` file.
+
+        Args:
+            path (``str``): The path of the `.safetensors` file.
+            weights (:class:`collections.OrderedDict`[``str``, :class:`torch.Tensor`]): The model weights.
+        """
+        import safetensors.torch
+        safetensors.torch.save_file(weights, path)
+
+    @classmethod
+    def to(cls, path: str, weights: Dict[str, torch.Tensor]) -> None:
+        """
+        Writes model weights to a file.
+
+        Args:
+            path (``str``): The path of the file.
+            weights (:class:`collections.OrderedDict`[``str``, :class:`torch.Tensor`]): The model weights.
+        """
+        file_type = path.split(".")[-1]
+
+        if file_type == "pt":
+            cls.to_pt(path, weights)
+        elif file_type == "safetensors":
+            cls.to_safetensors(path, weights)
+        else:
+            raise ValueError("Unknown file type. Supported types: .pt, .safetensors")

tensorneko/io/writer.py

@@ -3,6 +3,7 @@ from __future__ import annotations
 from typing import Type
 
 from tensorneko_util.io.writer import Writer as BaseWriter
+from .weight import WeightWriter
 
 try:
     from .mesh import MeshWriter
@@ -16,6 +17,7 @@ class Writer(BaseWriter):
 
     def __init__(self):
         super().__init__()
+        self.weight = WeightWriter
         self._mesh = None
 
     @property

tensorneko/layer/noise.py

@@ -22,7 +22,7 @@ class GaussianNoise(NekoModule):
         from https://discuss.pytorch.org/t/writing-a-simple-gaussian-noise-layer-in-pytorch/4694/3
     """
 
-    def __init__(self, sigma=0.1, device: Union[Device, str] = "cuda"):
+    def __init__(self, sigma=0.1, device: Device = "cuda"):
         super().__init__()
         self.sigma = sigma
         self.noise = torch.tensor(0.).to(device)

tensorneko/util/type.py

@@ -11,7 +11,7 @@ ModuleFactory = Union[Callable[[], Module], Callable[[int], Module]]
 """The module builder type of ``() -> torch.nn.Module | (int) -> torch.nn.Module``"""
 
 
-Device = device
+Device = Union[str, device]
 """Device type of :class:`torch.device`"""
 
 

tensorneko/version.txt

@@ -1 +1 @@
-0.3.8
+0.3.9

{tensorneko-0.3.8.dist-info → tensorneko-0.3.9.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: tensorneko
-Version: 0.3.8
+Version: 0.3.9
 Summary: Tensor Neural Engine Kompanion. An util library based on PyTorch and PyTorch Lightning.
 Home-page: https://github.com/ControlNet/tensorneko
 Author: ControlNet
@@ -33,7 +33,7 @@ Requires-Dist: pillow >=8.1
 Requires-Dist: av >=8.0.3
 Requires-Dist: numpy >=1.20.1
 Requires-Dist: einops >=0.3.0
-Requires-Dist: tensorneko-util ==0.3.8
+Requires-Dist: tensorneko-util ==0.3.9
 Requires-Dist: pysoundfile >=0.9.0 ; platform_system == "Windows"
 
 <h1 style="text-align: center">TensorNeko</h1> 

{tensorneko-0.3.8.dist-info → tensorneko-0.3.9.dist-info}/RECORD

@@ -1,15 +1,15 @@
-tensorneko/__init__.py,sha256=VPPK00Kduwi84QHnZKBZm8kBRdnPAji6f7J-adYAp_Y,770
+tensorneko/__init__.py,sha256=lLVC4StQ5q9OFJTceNrJbj-CIFHFMRvFQyGYJTjuRM4,812
 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=L_j9rGHiSDgBzj7QgRwf9tGYLHGVnIvOsy9NcC-vNsU,5
+tensorneko/version.txt,sha256=JCuzYNF_l66JsKuMd57Z4WJC2x1-36y0IfHrRwznDqM,5
 tensorneko/arch/__init__.py,sha256=w4lTUeyBIZelrnSjlBFWUF0erzOmBFl9FqeWQuSOyKs,248
 tensorneko/arch/auto_encoder.py,sha256=j6PWWyaNYaYNtw_zZ9ikzhCASqe9viXR3JGBIXSK92Y,2137
 tensorneko/arch/binary_classifier.py,sha256=1MkEbReXKLdDksRG5Rsife40grJk08EVDcNKp54Xvb4,2316
 tensorneko/arch/gan.py,sha256=bHh8s9UxHNSj9N6C6SllAaVJ61JhbY4ltwmQ-a1543w,7215
 tensorneko/arch/vqvae.py,sha256=_HJoDVBonB4PHymaV9JaAecS-NJZ1jfrKAhoqcLzf0g,3750
 tensorneko/arch/wgan.py,sha256=Oj1vZPXzhgTPHGGXGmbStgylskjjzulr6tYxj7E0zpQ,4722
-tensorneko/backend/__init__.py,sha256=hWEoQo3Bh7Fn7cfaJznA3k1sHWAUIJ9KAbDxOVjySxU,164
+tensorneko/backend/__init__.py,sha256=ppJhb1MC_WK6XLK2fX8x0z7Cn-8CkF97x67XBLoRCys,238
 tensorneko/callback/__init__.py,sha256=H1jOTsSYm9c4sxtcV9_uzumXZ95b4gbiftKqimNER7s,556
 tensorneko/callback/display_metrics_callback.py,sha256=qzhHcb68B7o9byfD1ZqEitSVkrwsSGOF-u59_Ip9dEg,318
 tensorneko/callback/earlystop_lr.py,sha256=d0G3NHYi-tNNCdaB7Rt7vVceI8CXOGT4zBTJYNciyug,1195
@@ -28,17 +28,20 @@ tensorneko/dataset/sampler/sequential_iter_sampler.py,sha256=DxBwSoWjYlq6kA6g-54
 tensorneko/debug/__init__.py,sha256=ZMfU3qquhMhl6EgPzM7Yuvvv0PWy3cR39UjPrrSmQcs,163
 tensorneko/evaluation/__init__.py,sha256=jW8dh1JRMpx3npjTp7wJLzz-IxFZTBh7F-Ztfoep9xs,296
 tensorneko/evaluation/enum.py,sha256=s3P8XAobku-as4in5vh6BanvVW5Ccwnff0t124lVFFg,137
-tensorneko/evaluation/fid.py,sha256=mDsgh7Ge7K8KrOLeWnSEVzzKfdCK0cI9TAWJJd5eqcQ,5550
+tensorneko/evaluation/fid.py,sha256=fNuE1CEp2rPXbaZfI0E1CspluInzFlUdKc8XZEexUME,5568
 tensorneko/evaluation/iou.py,sha256=phEmOWQ3cnWW377WeSHCoB8mGkHLHMHCl8_LL0IX3JA,2914
 tensorneko/evaluation/psnr.py,sha256=DeKxvY_xxawWMXHY0z3Nvbsi4dR57OUV4hjtUoCINXc,3757
 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=KB4xpdHKaqtEQXj2EOVB21Ev3ODPiQZFjNadZOipCMU,705
-tensorneko/io/writer.py,sha256=BHNtzROUY3AImx1QwVxbtZXuxMIfQq3WUI5PU1jeCpM,708
+tensorneko/io/reader.py,sha256=MSEfmzbRTk69qawTPrYruaKvQ8TXSyu5ZmFY7xVn-aU,773
+tensorneko/io/writer.py,sha256=MbO878ob24WSGKqjxyK-yRqp7xAiJBEtMxOGkOg_vOE,776
 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/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
@@ -47,7 +50,7 @@ tensorneko/layer/conv.py,sha256=c0XLcf2hRvyv6hMj4irIAYM0JyPzpv7tJBIW0Dl_sw8,5904
 tensorneko/layer/linear.py,sha256=NKIzSmSodk0wftu-f-6aSPazJ3X-eduRkZdLMMBjN_8,3283
 tensorneko/layer/log.py,sha256=UN7xgfzC7tF4P856taRtVgVHjc5CTTZqXYODvw6DhKQ,716
 tensorneko/layer/masked_conv2d.py,sha256=4DWjQ_Wc8lncYfk2Iur_M5V0mByvaQ2JZlOGRm29XLE,4796
-tensorneko/layer/noise.py,sha256=rc2Q1V1tT4KUszecOs82NeffrmDPOT9gELumo_tgdrU,1281
+tensorneko/layer/noise.py,sha256=bbZyrIZUk41cKBuCj3BWiQ_XTIRY9YXyDrlBavQh2Lo,1269
 tensorneko/layer/patching.py,sha256=FYwqVXbOtiLgT5Cy63zISK9bTzMIKG4U8hO-n14liiM,6585
 tensorneko/layer/positional_embedding.py,sha256=2_r1sZQPYNEdBVLDByNV85gZDofuV8IZSIjSNoKJKhw,4283
 tensorneko/layer/reshape.py,sha256=7GmpjQkir2eW2CkWuiUK648C4kKHUV81EqNiCzk2QSY,813
@@ -75,14 +78,14 @@ tensorneko/util/dispatched_misc.py,sha256=_0Go7XezdYB7bpMnCs1MDD_6mPNoWP5qt8DoKu
 tensorneko/util/misc.py,sha256=LEvACtGDOX43iK86A8-Cek0S9rbXFR0AtTP1edE3XDI,4701
 tensorneko/util/reproducibility.py,sha256=sw1vVi7VOnmzQYUocI5x9yKeZoHHiA4A5ja136XolrI,2102
 tensorneko/util/string_getter.py,sha256=Cq2mDYr3q758xJ9OBTwLDf-b6EMSYwlnNB0-kfsElfs,2491
-tensorneko/util/type.py,sha256=IaLpRQ5l8Ci6FZaGRohIb1ygrnJ3NTalomxDbhz68VM,716
+tensorneko/util/type.py,sha256=-dWknUu7RM4pFm7f3spgSlCIXpHA11RS36ol8uUFJgU,728
 tensorneko/visualization/__init__.py,sha256=PuNMhLz3oosY39AmKUr0biIgjfc_G_rQzp960me08Fg,626
 tensorneko/visualization/log_graph.py,sha256=NvOwWVc_petXWYdgaHosPFLa43sHBeacbYcfNtdRQg4,1511
 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.8.dist-info/LICENSE,sha256=Vd75kwgJpVuMnCRBWasQzceMlXt4YQL13ikBLy8G5h0,1067
-tensorneko-0.3.8.dist-info/METADATA,sha256=1i0yJukLciQ3ASW4rpphTBK_dEcTtnWpEn64DTCpk1s,18892
-tensorneko-0.3.8.dist-info/WHEEL,sha256=oiQVh_5PnQM0E3gPdiz09WCNmwiHDMaGer_elqB3coM,92
-tensorneko-0.3.8.dist-info/top_level.txt,sha256=sZHwlP0iyk7_zHuhRHzSBkdY9yEgyC48f6UVuZ6CvqE,11
-tensorneko-0.3.8.dist-info/RECORD,,
+tensorneko-0.3.9.dist-info/LICENSE,sha256=Vd75kwgJpVuMnCRBWasQzceMlXt4YQL13ikBLy8G5h0,1067
+tensorneko-0.3.9.dist-info/METADATA,sha256=QIW5lNPHf27Np7FwVhn3HoWeD-AQCe_uid7RApT9R-k,18892
+tensorneko-0.3.9.dist-info/WHEEL,sha256=oiQVh_5PnQM0E3gPdiz09WCNmwiHDMaGer_elqB3coM,92
+tensorneko-0.3.9.dist-info/top_level.txt,sha256=sZHwlP0iyk7_zHuhRHzSBkdY9yEgyC48f6UVuZ6CvqE,11
+tensorneko-0.3.9.dist-info/RECORD,,