lt-tensor 0.0.1a12__py3-none-any.whl → 0.0.1a13__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
@@ -1,185 +0,0 @@
1
- __all__ = [
2
- "Downsample1D",
3
- "Upsample1D",
4
- "DiffusionUNet",
5
- "UNetConvBlock1D",
6
- "UNetUpBlock1D",
7
- "NoisePredictor1D",
8
- ]
9
-
10
- from lt_tensor.torch_commons import *
11
- from lt_tensor.model_base import Model
12
- from lt_tensor.model_zoo.rsd import ResBlock1D
13
- from lt_tensor.misc_utils import log_tensor
14
-
15
- import torch.nn.functional as F
16
-
17
-
18
- class Downsample1D(Model):
19
- def __init__(
20
- self,
21
- in_channels: int,
22
- out_channels: int,
23
- ):
24
- super().__init__()
25
- self.pool = nn.Conv1d(in_channels, out_channels, 4, stride=2, padding=1)
26
-
27
- def forward(self, x):
28
- return self.pool(x)
29
-
30
-
31
- class Upsample1D(Model):
32
- def __init__(
33
- self,
34
- in_channels: int,
35
- out_channels: int,
36
- activation=nn.ReLU(inplace=True),
37
- ):
38
- super().__init__()
39
- self.up = nn.Sequential(
40
- nn.ConvTranspose1d(
41
- in_channels, out_channels, kernel_size=4, stride=2, padding=1
42
- ),
43
- nn.BatchNorm1d(out_channels),
44
- activation,
45
- )
46
-
47
- def forward(self, x):
48
- return self.up(x)
49
-
50
-
51
- class DiffusionUNet(Model):
52
- def __init__(self, in_channels=1, base_channels=64, out_channels=1, depth=4):
53
- super().__init__()
54
-
55
- self.depth = depth
56
- self.encoder_blocks = nn.ModuleList()
57
- self.downsamples = nn.ModuleList()
58
- self.upsamples = nn.ModuleList()
59
- self.decoder_blocks = nn.ModuleList()
60
- # Keep track of channel sizes per layer for skip connections
61
- self.channels = [in_channels] # starting input channel
62
- for i in range(depth):
63
- enc_in = self.channels[-1]
64
- enc_out = base_channels * (2**i)
65
- # Encoder block and downsample
66
- self.encoder_blocks.append(ResBlock1D(enc_in, enc_out))
67
- self.downsamples.append(
68
- Downsample1D(enc_out, enc_out)
69
- ) # halve time, keep channels
70
- self.channels.append(enc_out)
71
- # Bottleneck
72
- bottleneck_ch = self.channels[-1]
73
- self.bottleneck = ResBlock1D(bottleneck_ch, bottleneck_ch)
74
- # Decoder blocks (reverse channel flow)
75
- for i in reversed(range(depth)):
76
- skip_ch = self.channels[i + 1] # from encoder
77
- dec_out = self.channels[i] # match earlier stage's output
78
- self.upsamples.append(Upsample1D(skip_ch, skip_ch))
79
- self.decoder_blocks.append(ResBlock1D(skip_ch * 2, dec_out))
80
- # Final output projection (out_channels)
81
- self.final = nn.Conv1d(in_channels, out_channels, kernel_size=1)
82
-
83
- def forward(self, x: Tensor):
84
- skips = []
85
-
86
- # Encoder
87
- for enc, down in zip(self.encoder_blocks, self.downsamples):
88
- # log_tensor(x, "before enc")
89
- x = enc(x)
90
- skips.append(x)
91
- x = down(x)
92
-
93
- # Bottleneck
94
- x = self.bottleneck(x)
95
-
96
- # Decoder
97
- for up, dec, skip in zip(self.upsamples, self.decoder_blocks, reversed(skips)):
98
- x = up(x)
99
-
100
- # Match lengths via trimming or padding
101
- if x.shape[-1] > skip.shape[-1]:
102
- x = x[..., : skip.shape[-1]]
103
- elif x.shape[-1] < skip.shape[-1]:
104
- diff = skip.shape[-1] - x.shape[-1]
105
- x = F.pad(x, (0, diff))
106
-
107
- x = torch.cat([x, skip], dim=1) # concat on channels
108
- x = dec(x)
109
-
110
- # Final 1x1 conv
111
- return self.final(x)
112
-
113
-
114
- class UNetConvBlock1D(Model):
115
- def __init__(self, in_channels: int, out_channels: int, down: bool = True):
116
- super().__init__()
117
- self.down = down
118
- self.conv = nn.Sequential(
119
- nn.Conv1d(
120
- in_channels,
121
- out_channels,
122
- kernel_size=3,
123
- stride=2 if down else 1,
124
- padding=1,
125
- ),
126
- nn.BatchNorm1d(out_channels),
127
- nn.LeakyReLU(0.2),
128
- nn.Conv1d(out_channels, out_channels, kernel_size=3, padding=1),
129
- nn.BatchNorm1d(out_channels),
130
- nn.LeakyReLU(0.2),
131
- )
132
- self.downsample = (
133
- nn.Conv1d(in_channels, out_channels, kernel_size=1, stride=2 if down else 1)
134
- if in_channels != out_channels
135
- else nn.Identity()
136
- )
137
-
138
- def forward(self, x: torch.Tensor) -> torch.Tensor:
139
- # x: [B, C, T]
140
- residual = self.downsample(x)
141
- return self.conv(x) + residual
142
-
143
-
144
- class UNetUpBlock1D(Model):
145
- def __init__(self, in_channels: int, out_channels: int):
146
- super().__init__()
147
- self.conv = nn.Sequential(
148
- nn.Conv1d(in_channels, out_channels, kernel_size=3, padding=1),
149
- nn.BatchNorm1d(out_channels),
150
- nn.LeakyReLU(0.2),
151
- nn.Conv1d(out_channels, out_channels, kernel_size=3, padding=1),
152
- nn.BatchNorm1d(out_channels),
153
- nn.LeakyReLU(0.2),
154
- )
155
- self.upsample = nn.Upsample(scale_factor=2, mode="nearest")
156
-
157
- def forward(self, x: torch.Tensor, skip: torch.Tensor) -> torch.Tensor:
158
- x = self.upsample(x)
159
- x = torch.cat([x, skip], dim=1) # skip connection
160
- return self.conv(x)
161
-
162
-
163
- class NoisePredictor1D(Model):
164
- def __init__(self, in_channels: int, cond_dim: int = 0, hidden: int = 128):
165
- """
166
- Args:
167
- in_channels: channels of the noisy input [B, C, T]
168
- cond_dim: optional condition vector [B, cond_dim]
169
- """
170
- super().__init__()
171
- self.proj = nn.Linear(cond_dim, hidden) if cond_dim > 0 else None
172
- self.net = nn.Sequential(
173
- nn.Conv1d(in_channels, hidden, kernel_size=3, padding=1),
174
- nn.SiLU(),
175
- nn.Conv1d(hidden, in_channels, kernel_size=3, padding=1),
176
- )
177
-
178
- def forward(self, x: torch.Tensor, cond: Optional[torch.Tensor] = None):
179
- # x: [B, C, T], cond: [B, cond_dim]
180
- if cond is not None:
181
- cond_proj = self.proj(cond).unsqueeze(-1) # [B, hidden, 1]
182
- x = x + cond_proj # simple conditioning
183
- return self.net(x) # [B, C, T]
184
-
185
-