| 216 | return AttnBlock(in_dims) if using_sa else nn.Identity() |
| 217 | |
| 218 | class Encoder(nn.Module): |
| 219 | def __init__( |
| 220 | self, *, ch=32, ch_mult=(1, 2, 4, 8), num_res_blocks=2, |
| 221 | dropout=0.0, in_dims=3, z_channels, double_z=False, |
| 222 | using_sa=True, using_mid_sa=True, |
| 223 | patchwise_cfg: dict | None = None |
| 224 | ): |
| 225 | super().__init__() |
| 226 | self.ch = ch |
| 227 | self.num_resolutions = len(ch_mult) |
| 228 | self.num_res_blocks = num_res_blocks |
| 229 | self.in_dims = in_dims |
| 230 | self.patchwise_cfg = patchwise_cfg or {'enable': False} |
| 231 | |
| 232 | # Grouped conv setup |
| 233 | self.use_grouped = self.patchwise_cfg.get('enable', False) |
| 234 | self.grouped_depth = self.patchwise_cfg.get('grouped_depth', 2) if self.use_grouped else 0 |
| 235 | |
| 236 | # conv_in: if grouped, input is 3*D_embed; else original in_dims |
| 237 | actual_in_dims = in_dims if not self.use_grouped else (3 * self.patchwise_cfg.get('d_embed', 64)) |
| 238 | self.conv_in = torch.nn.Conv1d( |
| 239 | actual_in_dims, self.ch, |
| 240 | kernel_size=3, stride=1, padding=1, |
| 241 | groups=3 if self.use_grouped else 1 |
| 242 | ) |
| 243 | |
| 244 | in_ch_mult = (1,) + tuple(ch_mult) |
| 245 | self.down = nn.ModuleList() |
| 246 | for i_level in range(self.num_resolutions): |
| 247 | block = nn.ModuleList() |
| 248 | attn = nn.ModuleList() |
| 249 | block_in = ch * in_ch_mult[i_level] |
| 250 | block_out = ch * ch_mult[i_level] |
| 251 | for i_block in range(self.num_res_blocks): |
| 252 | # Use grouped conv for early layers if patchwise enabled |
| 253 | use_groups_here = self.use_grouped and i_level < self.grouped_depth |
| 254 | if use_groups_here: |
| 255 | # Custom grouped ResnetBlock - simplified, just group the convs |
| 256 | block.append(GroupedResnetBlock( |
| 257 | in_dims=block_in, out_channels=block_out, |
| 258 | dropout=dropout, groups=3 |
| 259 | )) |
| 260 | else: |
| 261 | block.append(ResnetBlock(in_dims=block_in, out_channels=block_out, dropout=dropout)) |
| 262 | block_in = block_out |
| 263 | if i_level == self.num_resolutions - 1 and using_sa: |
| 264 | attn.append(make_attn(block_in, using_sa=True)) |
| 265 | down = nn.Module() |
| 266 | down.block = block |
| 267 | down.attn = attn |
| 268 | if i_level != self.num_resolutions - 1: |
| 269 | down.downsample = Downsample1D_2x(block_in) |
| 270 | self.down.append(down) |
| 271 | |
| 272 | self.mid = nn.Module() |
| 273 | self.mid.block_1 = ResnetBlock(in_dims=block_in, out_channels=block_in, dropout=dropout) |
| 274 | self.mid.attn_1 = make_attn(block_in, using_sa=using_mid_sa) |
| 275 | self.mid.block_2 = ResnetBlock(in_dims=block_in, out_channels=block_in, dropout=dropout) |