tensorneko 0.3.8__tar.gz → 0.3.9__tar.gz

{tensorneko-0.3.8/src/tensorneko.egg-info → tensorneko-0.3.9}/PKG-INFO

@@ -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

{tensorneko-0.3.8 → tensorneko-0.3.9}/src/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-0.3.8 → tensorneko-0.3.9}/src/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-0.3.8 → tensorneko-0.3.9}/src/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-0.3.8 → tensorneko-0.3.9}/src/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-0.3.9/src/tensorneko/io/weight/__init__.py

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

tensorneko-0.3.9/src/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-0.3.9/src/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-0.3.8 → tensorneko-0.3.9}/src/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-0.3.8 → tensorneko-0.3.9}/src/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-0.3.8 → tensorneko-0.3.9}/src/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-0.3.9/src/tensorneko/version.txt

@@ -0,0 +1 @@
+0.3.9

{tensorneko-0.3.8 → tensorneko-0.3.9/src/tensorneko.egg-info}/PKG-INFO

@@ -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

{tensorneko-0.3.8 → tensorneko-0.3.9}/src/tensorneko.egg-info/SOURCES.txt

@@ -47,6 +47,9 @@ src/tensorneko/io/writer.py
 src/tensorneko/io/mesh/__init__.py
 src/tensorneko/io/mesh/mesh_reader.py
 src/tensorneko/io/mesh/mesh_writer.py
+src/tensorneko/io/weight/__init__.py
+src/tensorneko/io/weight/weight_reader.py
+src/tensorneko/io/weight/weight_writer.py
 src/tensorneko/layer/__init__.py
 src/tensorneko/layer/aggregation.py
 src/tensorneko/layer/attention.py

{tensorneko-0.3.8 → tensorneko-0.3.9}/src/tensorneko.egg-info/requires.txt

@@ -8,7 +8,7 @@ pillow>=8.1
 av>=8.0.3
 numpy>=1.20.1
 einops>=0.3.0
-tensorneko_util==0.3.8
+tensorneko_util==0.3.9
 
 [:platform_system == "Windows"]
 pysoundfile>=0.9.0

tensorneko-0.3.8/src/tensorneko/version.txt

@@ -1 +0,0 @@
-0.3.8