added building files

pyproject.toml

@@ -0,0 +1,22 @@
+[build-system]
+requires = ["setuptools>=61.0"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "timediffusion"
+version = "0.1"
+authors = [
+  { name="Shishkov Vladislav", email="[email protected]" },
+]
+description = ""
+readme = "README.md"
+requires-python = ">=3.7"
+classifiers = [
+    "Programming Language :: Python :: 3",
+    "License :: OSI Approved :: MIT License",
+    "Operating System :: OS Independent",
+]
+
+[project.urls]
+"Homepage" = "https://github.com/timetoai/TimeDiffusion"
+"Bug Tracker" = "https://github.com/timetoai/TimeDiffusion/issues"

setup.py

@@ -0,0 +1,3 @@
+from setuptools import setup
+
+setup()

src/timediffusion/__init__.py

@@ -0,0 +1,7 @@
+from .timediffusion import TD, TimeDiffusionProjector, TimeDiffusion
+
+__all__ = [
+    "TimeDiffusionProjector",
+    "TimeDiffusion",
+    "TD"
+]

src/timediffusion.py → src/timediffusion/timediffusion.py

@@ -1,3 +1,5 @@
+from typing import Union
+
 from tqdm import tqdm
 
 import numpy as np
@@ -79,8 +81,8 @@ class TemporalBlock(nn.Module):
     adds additional convolutional layer if needed to downsample number of channels
     inspired by https://github.com/locuslab/TCN
     """
-    def __init__(self, n_inputs: int, n_outputs: int, kernel_size: int | tuple[int],
-                  stride: int | tuple[int], dilation: int | tuple[int], padding: int | tuple[int],
+    def __init__(self, n_inputs: int, n_outputs: int, kernel_size: Union[int, tuple[int]],
+                  stride: Union[int, tuple[int]], dilation: Union[int, tuple[int]], padding: Union[int, tuple[int]],
                   dropout: int = 0.2, dims: int = 1, layers: int = 2):
         super().__init__()
 
@@ -130,7 +132,8 @@ class TimeDiffusionProjector(nn.Module):
     convolutional network, used as projector in TD
     consists of temporal blocks with exponentially increasing padding/dilation parameters
     """
-    def __init__(self, input_dims: list | tuple, max_deg_constraint: int = 13, conv_filters: int = 128, base_dropout: float = 0.05):
+    def __init__(self, input_dims: Union[list[int], tuple[int]], max_deg_constraint: int = 13,
+                  conv_filters: int = 128, base_dropout: float = 0.05):
         """
         args:
             `input_dims` - [channels, *dims]
@@ -236,10 +239,10 @@ def dtype(self):
     def device(self):
         return  next(self.model.parameters()).device
 
-    def fit(self, example: np.array | torch.Tensor, mask: None | np.array | torch.Tensor = None,
+    def fit(self, example: Union[np.array, torch.Tensor], mask: Union[None, np.array, torch.Tensor] = None,
             epochs: int = 20, batch_size: int = 2, steps_per_epoch: int = 32,
-            lr: float = 4e-4, distance_loss: str | nn.Module = "MAE",
-            distribution_loss: str | nn.Module = "kl_div", distrib_loss_coef = 1e-2,
+            lr: float = 4e-4, distance_loss: Union[str, nn.Module] = "MAE",
+            distribution_loss: Union[str, nn.Module] = "kl_div", distrib_loss_coef = 1e-2,
             verbose: bool = False, seed=42) -> list[float]:
         """
         simulates diffusion process for model training

tests/test_timediffusion.py

@@ -0,0 +1,8 @@
+import numpy as np
+
+import torch
+from torch import nn
+
+def test_count_params():
+    linear = nn.Linear(100, 100)
+    pass