lt-tensor 0.0.1.dev0__py3-none-any.whl → 0.0.1.dev3__py3-none-any.whl

lt_tensor/__init__.py

@@ -0,0 +1 @@
+__version__ = "0.0.1dev3"

lt_tensor/math_ops.py

@@ -47,66 +47,6 @@ def normalize_tensor(x: Tensor, eps: float = 1e-8) -> Tensor:
     return x / (torch.norm(x, dim=-1, keepdim=True) + eps)
 
 
-def stft(
-    waveform: Tensor,
-    n_fft: int = 512,
-    hop_length: Optional[int] = None,
-    win_length: Optional[int] = None,
-    window_fn: str = "hann",
-    center: bool = True,
-    return_complex: bool = True,
-) -> Tensor:
-    """Performs short-time Fourier transform using PyTorch."""
-    window = (
-        torch.hann_window(win_length or n_fft).to(waveform.device)
-        if window_fn == "hann"
-        else None
-    )
-    return torch.stft(
-        input=waveform,
-        n_fft=n_fft,
-        hop_length=hop_length,
-        win_length=win_length,
-        window=window,
-        center=center,
-        return_complex=return_complex,
-    )
-
-
-def istft(
-    stft_matrix: Tensor,
-    n_fft: int = 512,
-    hop_length: Optional[int] = None,
-    win_length: Optional[int] = None,
-    window_fn: str = "hann",
-    center: bool = True,
-    length: Optional[int] = None,
-) -> Tensor:
-    """Performs inverse short-time Fourier transform using PyTorch."""
-    window = (
-        torch.hann_window(win_length or n_fft).to(stft_matrix.device)
-        if window_fn == "hann"
-        else None
-    )
-    return torch.istft(
-        input=stft_matrix,
-        n_fft=n_fft,
-        hop_length=hop_length,
-        win_length=win_length,
-        window=window,
-        center=center,
-        length=length,
-    )
-
-
-def fft(x: Tensor, norm: Optional[str] = "backward") -> Tensor:
-    """Returns the FFT of a real tensor."""
-    return torch.fft.fft(x, norm=norm)
-
-
-def ifft(x: Tensor, norm: Optional[str] = "backward") -> Tensor:
-    """Returns the inverse FFT of a complex tensor."""
-    return torch.fft.ifft(x, norm=norm)
 
 
 def log_magnitude(stft_complex: Tensor, eps: float = 1e-5) -> Tensor:

lt_tensor/misc_utils.py

@@ -23,7 +23,7 @@ def log_tensor(
         _b = 20 + len(title.strip())
     print(f"shape: {item.shape}")
     print(f"dtype: {item.dtype}")
-    if not print_details:
+    if print_details:
         print(f"ndim: {item.ndim}")
         if isinstance(item, Tensor):
             print(f"device: {item.device}")

lt_tensor/transform.py

@@ -28,6 +28,68 @@ def to_mel_spectrogram(
     return mel_spectrogram(waveform)
 
 
+def stft(
+    waveform: Tensor,
+    n_fft: int = 512,
+    hop_length: Optional[int] = None,
+    win_length: Optional[int] = None,
+    window_fn: str = "hann",
+    center: bool = True,
+    return_complex: bool = True,
+) -> Tensor:
+    """Performs short-time Fourier transform using PyTorch."""
+    window = (
+        torch.hann_window(win_length or n_fft).to(waveform.device)
+        if window_fn == "hann"
+        else None
+    )
+    return torch.stft(
+        input=waveform,
+        n_fft=n_fft,
+        hop_length=hop_length,
+        win_length=win_length,
+        window=window,
+        center=center,
+        return_complex=return_complex,
+    )
+
+
+def istft(
+    stft_matrix: Tensor,
+    n_fft: int = 512,
+    hop_length: Optional[int] = None,
+    win_length: Optional[int] = None,
+    window_fn: str = "hann",
+    center: bool = True,
+    length: Optional[int] = None,
+) -> Tensor:
+    """Performs inverse short-time Fourier transform using PyTorch."""
+    window = (
+        torch.hann_window(win_length or n_fft).to(stft_matrix.device)
+        if window_fn == "hann"
+        else None
+    )
+    return torch.istft(
+        input=stft_matrix,
+        n_fft=n_fft,
+        hop_length=hop_length,
+        win_length=win_length,
+        window=window,
+        center=center,
+        length=length,
+    )
+
+
+def fft(x: Tensor, norm: Optional[str] = "backward") -> Tensor:
+    """Returns the FFT of a real tensor."""
+    return torch.fft.fft(x, norm=norm)
+
+
+def ifft(x: Tensor, norm: Optional[str] = "backward") -> Tensor:
+    """Returns the inverse FFT of a complex tensor."""
+    return torch.fft.ifft(x, norm=norm)
+
+
 def to_log_mel_spectrogram(
     waveform: torch.Tensor, sample_rate: int = 22050, eps: float = 1e-9, **kwargs
 ) -> torch.Tensor:
@@ -111,3 +173,157 @@ def get_sinusoidal_embedding(timesteps: torch.Tensor, dim: int) -> torch.Tensor:
     emb = timesteps[:, None].float() * emb[None, :]  # [B, half_dim]
     emb = torch.cat((emb.sin(), emb.cos()), dim=-1)  # [B, dim]
     return emb
+
+
+def _generate_window(
+    M: int, alpha: float = 0.5, device: Optional[DeviceType] = None
+) -> Tensor:
+    if M < 1:
+        raise ValueError("Window length M must be >= 1.")
+    if M == 1:
+        return torch.ones(1, device=device)
+
+    n = torch.arange(M, dtype=torch.float32, device=device)
+    window = alpha - (1.0 - alpha) * torch.cos(2.0 * math.pi * n / (M - 1))
+    return window
+
+
+def pad_center(tensor: torch.Tensor, size: int, axis: int = -1) -> torch.Tensor:
+    n = tensor.shape[axis]
+    if size < n:
+        raise ValueError(f"Target size ({size}) must be at least input size ({n})")
+
+    lpad = (size - n) // 2
+    rpad = size - n - lpad
+
+    pad = [0] * (2 * tensor.ndim)
+    pad[2 * axis + 1] = rpad
+    pad[2 * axis] = lpad
+
+    return F.pad(tensor, pad, mode="constant", value=0)
+
+
+def normalize(
+    S: torch.Tensor,
+    norm: float = float("inf"),
+    axis: int = 0,
+    threshold: float = 1e-10,
+    fill: bool = False,
+) -> torch.Tensor:
+    mag = S.abs().float()
+
+    if norm is None:
+        return S
+
+    elif norm == float("inf"):
+        length = mag.max(dim=axis, keepdim=True).values
+
+    elif norm == float("-inf"):
+        length = mag.min(dim=axis, keepdim=True).values
+
+    elif norm == 0:
+        length = (mag > 0).sum(dim=axis, keepdim=True).float()
+
+    elif norm > 0:
+        length = (mag**norm).sum(dim=axis, keepdim=True) ** (1.0 / norm)
+
+    else:
+        raise ValueError(f"Unsupported norm: {norm}")
+
+    small_idx = length < threshold
+    length = length.clone()
+    if fill:
+        length[small_idx] = float("nan")
+        Snorm = S / length
+        Snorm[Snorm != Snorm] = 1.0  # replace nan with fill_norm (default 1.0)
+    else:
+        length[small_idx] = float("inf")
+        Snorm = S / length
+
+    return Snorm
+
+
+def window_sumsquare(
+    window_spec: Union[str, int, float, Callable, List[Any], Tuple[Any, ...]],
+    n_frames: int,
+    hop_length: int = 300,
+    win_length: int = 1200,
+    n_fft: int = 2048,
+    dtype: torch.dtype = torch.float32,
+    norm: Optional[Union[int, float]] = None,
+    device: Optional[torch.device] = "cpu",
+):
+    if win_length is None:
+        win_length = n_fft
+
+    total_length = n_fft + hop_length * (n_frames - 1)
+    x = torch.zeros(total_length, dtype=dtype, device=device)
+
+    # Get the window (from scipy for now)
+    win = _generate_window(window_spec, win_length, fftbins=True)
+    win = torch.tensor(win, dtype=dtype, device=device)
+
+    # Normalize and square
+    win_sq = normalize(win, norm=norm, axis=0) ** 2
+    win_sq = pad_center(win_sq, size=n_fft, axis=0)
+
+    # Accumulate squared windows
+    for i in range(n_frames):
+        sample = i * hop_length
+        end = min(total_length, sample + n_fft)
+        length = end - sample
+        x[sample:end] += win_sq[:length]
+
+    return x
+
+
+
+def inverse_transform(
+    spec: Tensor,
+    phase: Tensor,
+    window: Optional[Tensor] = None,
+    n_fft: int = 2048,
+    hop_length: int = 300,
+    win_length: int = 1200,
+    length: Optional[Any] = None,
+):
+    if window is None:
+        window = _generate_window(win_length)
+    return torch.istft(
+        spec * torch.exp(phase * 1j),
+        n_fft,
+        hop_length,
+        win_length,
+        window=window,
+        length=length,
+    )
+
+
+def stft_istft_rebuild(
+    input_data: Tensor,
+    window: Optional[Tensor] = None,
+    n_fft: int = 2048,
+    hop_length: int = 300,
+    win_length: int = 1200,
+):
+    """
+    Perform STFT followed by ISTFT reconstruction using magnitude and phase.
+    """
+    if window is None:
+        window = _generate_window(win_length)
+    st = torch.stft(
+        input_data,
+        n_fft,
+        hop_length,
+        win_length,
+        window=window,
+        return_complex=True,
+    )
+    return torch.istft(
+        torch.abs(st) * torch.exp(1j * torch.angle(st)),
+        n_fft,
+        hop_length,
+        win_length,
+        window=window,
+        length=input_data.shape[-1],
+    ).squeeze(0)

{lt_tensor-0.0.1.dev0.dist-info → lt_tensor-0.0.1.dev3.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: lt-tensor
-Version: 0.0.1.dev0
+Version: 0.0.1.dev3
 Summary: General utilities for PyTorch and others. Built for general use.
 Home-page: https://github.com/gr1336/lt-tensor/
 Author: gr1336

{lt_tensor-0.0.1.dev0.dist-info → lt_tensor-0.0.1.dev3.dist-info}/RECORD

@@ -1,11 +1,11 @@
-lt_tensor/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+lt_tensor/__init__.py,sha256=pUB05ZkgkpP10ivzwoWdbq_HCxw-iOsbf6m8eFtx-YM,26
 lt_tensor/_basics.py,sha256=XS2OrvyzboUTKURGMU1fmbn2gFq2779HW8xz05fm4x8,8181
 lt_tensor/_torch_commons.py,sha256=_2Eck-MsQ46PxW5ku7NJvNSL5vg54_4GkLCqdzFevwA,402
 lt_tensor/lr_schedulers.py,sha256=oLYw2X78KSdlOD0pwPO0lsBj1xqLOAzT8rnRBnCO19o,3560
-lt_tensor/math_ops.py,sha256=ajHjt5xmS0dU5vdFYV0QE9BPxTp4ymXVwspO6u6Pwaw,3481
-lt_tensor/misc_utils.py,sha256=EaQ__986n5J-oAnMXmTp9hXxCDk6NrgvGlU_C05M7G4,19193
+lt_tensor/math_ops.py,sha256=4jUTtS1ZwJnh1AaFGhMuKTcakjwLQ846AKcBoDEwudI,1862
+lt_tensor/misc_utils.py,sha256=Zo9p-iwRLNwVBk5kRP5sPCw7IiOBST9hW7LBK9JKCVs,19189
 lt_tensor/monotonic_align.py,sha256=LhBd8p1xdBzg6jQrQX1j7b4PNeYGwIqM24zcU-pHOLE,2239
-lt_tensor/transform.py,sha256=ucHy-nuibhoH75brNBPxXFVahAd_NHr27vXLgLlkNGk,3552
+lt_tensor/transform.py,sha256=CbTzqRUzl-U-2KMHtYUgiliU_61H3cag8a-N526nzdg,9124
 lt_tensor/model_zoo/__init__.py,sha256=sIidI3gSoxTq2OTenqADiFlmd9EH8rS_mzjLePdNHqc,77
 lt_tensor/model_zoo/basic.py,sha256=Ccjg26-gnqbkOUg0aNnLhZmsfvMwqTxXVXimfC5qnl0,1862
 lt_tensor/model_zoo/residual.py,sha256=iViJ5MrIY1vUbPB0ZiOjh62wfQagbPulMvbZPTXp-OU,7370
@@ -14,8 +14,8 @@ lt_tensor/model_zoo/diffusion/models.py,sha256=LRf5B2MPic4Dwfvg2PxG61KFtqvATV3wO
 lt_tensor/model_zoo/transformer_models/__init__.py,sha256=NOvCP-EySkNHLbGOxFMURp0AztvJQSNo9J_0hsEzLto,130
 lt_tensor/model_zoo/transformer_models/models.py,sha256=Xh2nq83w0qMX-Co7EmaQb_auq7Fi7kyULklc0_yq3oo,4088
 lt_tensor/model_zoo/transformer_models/positional_encoders.py,sha256=ilJDKAonnU0973BT_7gT5ke_8PbWHMxsr8EvNQHAcUY,3529
-lt_tensor-0.0.1.dev0.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
-lt_tensor-0.0.1.dev0.dist-info/METADATA,sha256=2hs2eLrpyXhkLdWdzBre_trMJsizLgC0wLPnE4eHVCE,1055
-lt_tensor-0.0.1.dev0.dist-info/WHEEL,sha256=zaaOINJESkSfm_4HQVc5ssNzHCPXhJm0kEUakpsEHaU,91
-lt_tensor-0.0.1.dev0.dist-info/top_level.txt,sha256=35FuhFeXnUyvHWdbVHGPh0hS8euofafnJ_GJAVSF4Kk,10
-lt_tensor-0.0.1.dev0.dist-info/RECORD,,
+lt_tensor-0.0.1.dev3.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
+lt_tensor-0.0.1.dev3.dist-info/METADATA,sha256=cIBisih0DD2ZXdcHxjEUofVmPG8L0BxhIiym7AIkVSo,1055
+lt_tensor-0.0.1.dev3.dist-info/WHEEL,sha256=zaaOINJESkSfm_4HQVc5ssNzHCPXhJm0kEUakpsEHaU,91
+lt_tensor-0.0.1.dev3.dist-info/top_level.txt,sha256=35FuhFeXnUyvHWdbVHGPh0hS8euofafnJ_GJAVSF4Kk,10
+lt_tensor-0.0.1.dev3.dist-info/RECORD,,