add mdiblock / outblock architecture

Author: natolambert

src/diffusers/models/unet_1d.py

@@ -60,7 +60,7 @@ def __init__(
         out_channels: int = 14,
         down_block_types: Tuple[str] = ("DownResnetBlock1D", "DownResnetBlock1D", "DownResnetBlock1D"),
         up_block_types: Tuple[str] = ("UpResnetBlock1D", "UpResnetBlock1D"),
-        mid_block_types: Tuple[str] = ("MidResTemporalBlock1D", "MidResTemporalBlock1D"),
+        mid_block_type: Tuple[str] = "MidResTemporalBlock1D",
         out_block_type: str = "OutConv1DBlock",
         block_out_channels: Tuple[int] = (32, 128, 256),
         act_fn: str = "mish",
@@ -79,7 +79,9 @@ def __init__(
         )
 
         self.down_blocks = nn.ModuleList([])
+        self.mid_block = None
         self.up_blocks = nn.ModuleList([])
+        self.out_block = None
         mid_dim = block_out_channels[-1]
 
         # down
@@ -101,25 +103,15 @@ def __init__(
             self.down_blocks.append(down_block)
 
         # mid
-        self.mid_blocks = nn.ModuleList([])
-        for i, mid_block_type in enumerate(mid_block_types):
-            if always_downsample:
-                mid_block = get_mid_block(
-                    mid_block_type,
-                    in_channels=mid_dim // (i + 1),
-                    out_channels=mid_dim // ((i + 1) * 2),
-                    embed_dim=block_out_channels[0],
-                    add_downsample=True,
-                )
-            else:
-                mid_block = get_mid_block(
-                    mid_block_type,
-                    in_channels=mid_dim,
-                    out_channels=mid_dim,
-                    embed_dim=block_out_channels[0],
-                    add_downsample=False,
-                )
-            self.mid_blocks.append(mid_block)
+        self.mid_block = get_mid_block(
+            mid_block_type,
+            in_channels=mid_dim,
+            out_channels=mid_dim,
+            embed_dim=block_out_channels[0],
+            num_layers=layers_per_block,
+            add_downsample=always_downsample,
+        )
+
         # up
         reversed_block_out_channels = list(reversed(block_out_channels))
         for i, up_block_type in enumerate(up_block_types):
@@ -184,15 +176,16 @@ def forward(
             down_block_res_samples.append(res_samples[0])
 
         # 3. mid
-        for mid_block in self.mid_blocks:
-            sample = mid_block(sample, temb)
+        if self.mid_block:
+            sample = self.mid_block(sample, temb)
 
         # 4. up
         for up_block in self.up_blocks:
             sample = up_block(hidden_states=sample, res_hidden_states=down_block_res_samples.pop(), temb=temb)
 
         # 5. post-process
-        sample = self.out_block(sample, temb)
+        if self.out_block:
+            sample = self.out_block(sample, temb)
 
         if not return_dict:
             return (sample,)

src/diffusers/models/unet_1d_blocks.py

@@ -47,7 +47,7 @@ def __init__(
         if groups_out is None:
             groups_out = groups
 
-        # there will always be at least one resenet
+        # there will always be at least one resnet
         resnets = [ResidualTemporalBlock1D(in_channels, out_channels, embed_dim=temb_channels)]
 
         for _ in range(num_layers):
@@ -111,7 +111,7 @@ def __init__(
         if groups_out is None:
             groups_out = groups
 
-        # there will always be at least one resenet
+        # there will always be at least one resnet
         resnets = [ResidualTemporalBlock1D(2 * in_channels, out_channels, embed_dim=temb_channels)]
 
         for _ in range(num_layers):
@@ -174,22 +174,60 @@ class UpBlock1DNoSkip(nn.Module):
 
 
 class MidResTemporalBlock1D(nn.Module):
-    def __init__(self, in_channels, out_channels, embed_dim, add_downsample):
+    def __init__(
+        self,
+        in_channels,
+        out_channels,
+        embed_dim,
+        num_layers: int = 1,
+        add_downsample: bool = False,
+        add_upsample: bool = False,
+        non_linearity=None,
+    ):
         super().__init__()
         self.in_channels = in_channels
         self.out_channels = out_channels
         self.add_downsample = add_downsample
-        self.resnet = ResidualTemporalBlock1D(in_channels, out_channels, embed_dim=embed_dim)
 
+        # there will always be at least one resnet
+        resnets = [ResidualTemporalBlock1D(in_channels, out_channels, embed_dim=embed_dim)]
+
+        for _ in range(num_layers):
+            resnets.append(ResidualTemporalBlock1D(out_channels, out_channels, embed_dim=embed_dim))
+
+        self.resnets = nn.ModuleList(resnets)
+
+        if non_linearity == "swish":
+            self.nonlinearity = lambda x: F.silu(x)
+        elif non_linearity == "mish":
+            self.nonlinearity = nn.Mish()
+        elif non_linearity == "silu":
+            self.nonlinearity = nn.SiLU()
+        else:
+            self.nonlinearity = None
+
+        self.upsample = None
+        if add_downsample:
+            self.upsample = Downsample1D(out_channels, use_conv=True)
+
+        self.downsample = None
         if add_downsample:
             self.downsample = Downsample1D(out_channels, use_conv=True)
-        else:
-            self.downsample = nn.Identity()
 
-    def forward(self, sample, temb):
-        sample = self.resnet(sample, temb)
-        sample = self.downsample(sample)
-        return sample
+        if self.upsample and self.downsample:
+            raise ValueError("Block cannot downsample and upsample")
+
+    def forward(self, hidden_states, temb):
+        hidden_states = self.resnets[0](hidden_states, temb)
+        for resnet in self.resnets[1:]:
+            hidden_states = resnet(hidden_states, temb)
+
+        if self.upsample:
+            hidden_states = self.upsample(hidden_states)
+        if self.downsample:
+            self.downsample = self.downsample(hidden_states)
+
+        return hidden_states
 
 
 class OutConv1DBlock(nn.Module):
@@ -203,14 +241,14 @@ def __init__(self, num_groups_out, out_channels, embed_dim, act_fn):
             self.final_conv1d_act = nn.Mish()
         self.final_conv1d_2 = nn.Conv1d(embed_dim, out_channels, 1)
 
-    def forward(self, sample, t):
-        sample = self.final_conv1d_1(sample)
-        sample = rearrange_dims(sample)
-        sample = self.final_conv1d_gn(sample)
-        sample = rearrange_dims(sample)
-        sample = self.final_conv1d_act(sample)
-        sample = self.final_conv1d_2(sample)
-        return sample
+    def forward(self, hidden_states, temb=None):
+        hidden_states = self.final_conv1d_1(hidden_states)
+        hidden_states = rearrange_dims(hidden_states)
+        hidden_states = self.final_conv1d_gn(hidden_states)
+        hidden_states = rearrange_dims(hidden_states)
+        hidden_states = self.final_conv1d_act(hidden_states)
+        hidden_states = self.final_conv1d_2(hidden_states)
+        return hidden_states
 
 
 class OutValueFunctionBlock(nn.Module):
@@ -224,13 +262,13 @@ def __init__(self, fc_dim, embed_dim):
             ]
         )
 
-    def forward(self, sample, t):
-        sample = sample.view(sample.shape[0], -1)
-        sample = torch.cat((sample, t), dim=-1)
+    def forward(self, hidden_states, temb):
+        hidden_states = hidden_states.view(hidden_states.shape[0], -1)
+        hidden_states = torch.cat((hidden_states, temb), dim=-1)
         for layer in self.final_block:
-            sample = layer(sample)
+            hidden_states = layer(hidden_states)
 
-        return sample
+        return hidden_states
 
 
 def get_down_block(down_block_type, num_layers, in_channels, out_channels, temb_channels, add_downsample):
@@ -260,9 +298,15 @@ def get_up_block(up_block_type, num_layers, in_channels, out_channels, temb_chan
     raise ValueError(f"{up_block_type} does not exist.")
 
 
-def get_mid_block(mid_block_type, in_channels, out_channels, embed_dim, add_downsample):
+def get_mid_block(mid_block_type, num_layers, in_channels, out_channels, embed_dim, add_downsample):
     if mid_block_type == "MidResTemporalBlock1D":
-        return MidResTemporalBlock1D(in_channels, out_channels, embed_dim, add_downsample)
+        return MidResTemporalBlock1D(
+            num_layers=num_layers,
+            in_channels=in_channels,
+            out_channels=out_channels,
+            embed_dim=embed_dim,
+            add_downsample=add_downsample,
+        )
     raise ValueError(f"{mid_block_type} does not exist.")