all tests pass

natolambert · natolambert · commit a2dd559e12dd · 2022-10-25T17:10:00.000-07:00
diff --git a/src/diffusers/models/unet_1d.py b/src/diffusers/models/unet_1d.py
@@ -44,7 +44,7 @@ class UNet1DModel(ModelMixin, ConfigMixin):
     implements for all the model (such as downloading or saving, etc.)
 
     Parameters:
-        sample_size (`int`, *optionl*): Default length of sample. Should be adaptable at runtime.
+        sample_size (`int`, *optional*): Default length of sample. Should be adaptable at runtime.
         in_channels (`int`, *optional*, defaults to 2): Number of channels in the input sample.
         out_channels (`int`, *optional*, defaults to 2): Number of channels in the output.
         time_embedding_type (`str`, *optional*, defaults to `"fourier"`): Type of time embedding to use.
@@ -78,7 +78,7 @@ def __init__(
         down_block_types: Tuple[str] = ("DownResnetBlock1D", "DownResnetBlock1D", "DownResnetBlock1D"),
         up_block_types: Tuple[str] = ("UpResnetBlock1D", "UpResnetBlock1D"),
         mid_block_type: Tuple[str] = "MidResTemporalBlock1D",
-        out_block_type: str = "OutConv1DBlock",
+        out_block_type: str = None,
         block_out_channels: Tuple[int] = (32, 128, 256),
         act_fn: str = "mish",
         norm_num_groups: int = 8,
@@ -211,7 +211,7 @@ def forward(
             timesteps = timesteps[None].to(sample.device)
 
         timestep_embed = self.time_proj(timesteps)
-        if self.time_mlp:
+        if self.config.use_timestep_embedding:
             timestep_embed = self.time_mlp(timestep_embed)
         else:
             timestep_embed = timestep_embed[..., None]
diff --git a/src/diffusers/models/unet_1d_blocks.py b/src/diffusers/models/unet_1d_blocks.py
@@ -312,7 +312,7 @@ def __init__(self, kernel="linear", pad_mode="reflect"):
         self.pad = kernel_1d.shape[0] // 2 - 1
         self.register_buffer("kernel", kernel_1d)
 
-    def forward(self, hidden_states):
+    def forward(self, hidden_states, temb=None):
         hidden_states = F.pad(hidden_states, ((self.pad + 1) // 2,) * 2, self.pad_mode)
         weight = hidden_states.new_zeros([hidden_states.shape[1], hidden_states.shape[1], self.kernel.shape[0]])
         indices = torch.arange(hidden_states.shape[1], device=hidden_states.device)
@@ -441,7 +441,7 @@ def __init__(self, mid_channels, in_channels, out_channels=None):
         self.attentions = nn.ModuleList(attentions)
         self.resnets = nn.ModuleList(resnets)
 
-    def forward(self, hidden_states):
+    def forward(self, hidden_states, temb=None):
         hidden_states = self.down(hidden_states)
         for attn, resnet in zip(self.attentions, self.resnets):
             hidden_states = resnet(hidden_states)
@@ -546,7 +546,7 @@ def __init__(self, in_channels, out_channels, mid_channels=None):
         self.resnets = nn.ModuleList(resnets)
         self.up = Upsample1d(kernel="cubic")
 
-    def forward(self, hidden_states, res_hidden_states_tuple):
+    def forward(self, hidden_states, res_hidden_states_tuple, temb=None):
         res_hidden_states = res_hidden_states_tuple[-1]
         hidden_states = torch.cat([hidden_states, res_hidden_states], dim=1)
 
@@ -573,7 +573,7 @@ def __init__(self, in_channels, out_channels, mid_channels=None):
         self.resnets = nn.ModuleList(resnets)
         self.up = Upsample1d(kernel="cubic")
 
-    def forward(self, hidden_states, res_hidden_states_tuple):
+    def forward(self, hidden_states, res_hidden_states_tuple, temb=None):
         res_hidden_states = res_hidden_states_tuple[-1]
         hidden_states = torch.cat([hidden_states, res_hidden_states], dim=1)
 
@@ -598,7 +598,7 @@ def __init__(self, in_channels, out_channels, mid_channels=None):
 
         self.resnets = nn.ModuleList(resnets)
 
-    def forward(self, hidden_states, res_hidden_states_tuple):
+    def forward(self, hidden_states, res_hidden_states_tuple, temb=None):
         res_hidden_states = res_hidden_states_tuple[-1]
         hidden_states = torch.cat([hidden_states, res_hidden_states], dim=1)
 
diff --git a/tests/test_models_unet_1d.py b/tests/test_models_unet_1d.py
@@ -60,6 +60,8 @@ def prepare_init_args_and_inputs_for_common(self):
             "in_channels": 14,
             "out_channels": 14,
             "time_embedding_type": "positional",
+            "use_timestep_embedding": True,
+            "out_block_type": "OutConv1DBlock",
         }
         inputs_dict = self.dummy_input
         return init_dict, inputs_dict
@@ -159,6 +161,7 @@ def prepare_init_args_and_inputs_for_common(self):
             "block_out_channels": [32, 64, 128, 256],
             "layers_per_block": 1,
             "always_downsample": True,
+            "use_timestep_embedding": True
         }
         inputs_dict = self.dummy_input
         return init_dict, inputs_dict
