Repository: huggingface/diffusers

Issue:
Update UNet1D Tests to not use Mish
**Is your feature request related to a problem? Please describe.**
Currently the `UNet1DTests` heavily rely on mish because they were based on a re-implementation of a model. This is not needed generally, so it would be good to add another 1d model just for tests so we can delete a lot of the tests that look like this.

```
    @unittest.skipIf(torch_device == "mps", "mish op not supported in MPS")
    def test_determinism(self):
        super().test_determinism()
```



Discussion:
Hi @natolambert, wondering if I could take up this issue. And do we essentially want _some_ 1D model that can run without requiring mish? Thanks!
Yeah @thehimalayanleo that's exactly right. It's preferable that it has an application, but not needed.

The RL tests can stay as is, it's just that a lot of the `UNet1D` tests were added in the RL PR, so there can be some quality of life improvements there 🤗.
Thank you, @natolambert. I will take a look over the weekend and submit a PR.
The activation function as well should be an optional param until Mish is supported on M1. I am unable to run training even with MPS fallback enabled 

Resolving patch:
diff --git a/src/diffusers/models/unet_1d_blocks.py b/src/diffusers/models/unet_1d_blocks.py
--- a/src/diffusers/models/unet_1d_blocks.py
+++ b/src/diffusers/models/unet_1d_blocks.py
@@ -235,12 +235,12 @@ def forward(self, hidden_states, temb=None):
 
 
 class OutValueFunctionBlock(nn.Module):
-    def __init__(self, fc_dim, embed_dim):
+    def __init__(self, fc_dim, embed_dim, act_fn="mish"):
         super().__init__()
         self.final_block = nn.ModuleList(
             [
                 nn.Linear(fc_dim + embed_dim, fc_dim // 2),
-                nn.Mish(),
+                get_activation(act_fn),
                 nn.Linear(fc_dim // 2, 1),
             ]
         )
@@ -652,5 +652,5 @@ def get_out_block(*, out_block_type, num_groups_out, embed_dim, out_channels, ac
     if out_block_type == "OutConv1DBlock":
         return OutConv1DBlock(num_groups_out, out_channels, embed_dim, act_fn)
     elif out_block_type == "ValueFunction":
-        return OutValueFunctionBlock(fc_dim, embed_dim)
+        return OutValueFunctionBlock(fc_dim, embed_dim, act_fn)
     return None


Test patch:
diff --git a/tests/models/test_models_unet_1d.py b/tests/models/test_models_unet_1d.py
--- a/tests/models/test_models_unet_1d.py
+++ b/tests/models/test_models_unet_1d.py
@@ -52,27 +52,21 @@ def test_ema_training(self):
     def test_training(self):
         pass
 
-    @unittest.skipIf(torch_device == "mps", "mish op not supported in MPS")
     def test_determinism(self):
         super().test_determinism()
 
-    @unittest.skipIf(torch_device == "mps", "mish op not supported in MPS")
     def test_outputs_equivalence(self):
         super().test_outputs_equivalence()
 
-    @unittest.skipIf(torch_device == "mps", "mish op not supported in MPS")
     def test_from_save_pretrained(self):
         super().test_from_save_pretrained()
 
-    @unittest.skipIf(torch_device == "mps", "mish op not supported in MPS")
     def test_from_save_pretrained_variant(self):
         super().test_from_save_pretrained_variant()
 
-    @unittest.skipIf(torch_device == "mps", "mish op not supported in MPS")
     def test_model_from_pretrained(self):
         super().test_model_from_pretrained()
 
-    @unittest.skipIf(torch_device == "mps", "mish op not supported in MPS")
     def test_output(self):
         super().test_output()
 
@@ -89,12 +83,11 @@ def prepare_init_args_and_inputs_for_common(self):
             "mid_block_type": "MidResTemporalBlock1D",
             "down_block_types": ("DownResnetBlock1D", "DownResnetBlock1D", "DownResnetBlock1D", "DownResnetBlock1D"),
             "up_block_types": ("UpResnetBlock1D", "UpResnetBlock1D", "UpResnetBlock1D"),
-            "act_fn": "mish",
+            "act_fn": "swish",
         }
         inputs_dict = self.dummy_input
         return init_dict, inputs_dict
 
-    @unittest.skipIf(torch_device == "mps", "mish op not supported in MPS")
     def test_from_pretrained_hub(self):
         model, loading_info = UNet1DModel.from_pretrained(
             "bglick13/hopper-medium-v2-value-function-hor32", output_loading_info=True, subfolder="unet"
@@ -107,7 +100,6 @@ def test_from_pretrained_hub(self):
 
         assert image is not None, "Make sure output is not None"
 
-    @unittest.skipIf(torch_device == "mps", "mish op not supported in MPS")
     def test_output_pretrained(self):
         model = UNet1DModel.from_pretrained("bglick13/hopper-medium-v2-value-function-hor32", subfolder="unet")
         torch.manual_seed(0)
@@ -177,27 +169,21 @@ def input_shape(self):
     def output_shape(self):
         return (4, 14, 1)
 
-    @unittest.skipIf(torch_device == "mps", "mish op not supported in MPS")
     def test_determinism(self):
         super().test_determinism()
 
-    @unittest.skipIf(torch_device == "mps", "mish op not supported in MPS")
     def test_outputs_equivalence(self):
         super().test_outputs_equivalence()
 
-    @unittest.skipIf(torch_device == "mps", "mish op not supported in MPS")
     def test_from_save_pretrained(self):
         super().test_from_save_pretrained()
 
-    @unittest.skipIf(torch_device == "mps", "mish op not supported in MPS")
     def test_from_save_pretrained_variant(self):
         super().test_from_save_pretrained_variant()
 
-    @unittest.skipIf(torch_device == "mps", "mish op not supported in MPS")
     def test_model_from_pretrained(self):
         super().test_model_from_pretrained()
 
-    @unittest.skipIf(torch_device == "mps", "mish op not supported in MPS")
     def test_output(self):
         # UNetRL is a value-function is different output shape
         init_dict, inputs_dict = self.prepare_init_args_and_inputs_for_common()
@@ -241,7 +227,6 @@ def prepare_init_args_and_inputs_for_common(self):
         inputs_dict = self.dummy_input
         return init_dict, inputs_dict
 
-    @unittest.skipIf(torch_device == "mps", "mish op not supported in MPS")
     def test_from_pretrained_hub(self):
         value_function, vf_loading_info = UNet1DModel.from_pretrained(
             "bglick13/hopper-medium-v2-value-function-hor32", output_loading_info=True, subfolder="value_function"
@@ -254,7 +239,6 @@ def test_from_pretrained_hub(self):
 
         assert image is not None, "Make sure output is not None"
 
-    @unittest.skipIf(torch_device == "mps", "mish op not supported in MPS")
     def test_output_pretrained(self):
         value_function, vf_loading_info = UNet1DModel.from_pretrained(
             "bglick13/hopper-medium-v2-value-function-hor32", output_loading_info=True, subfolder="value_function"