We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
Question about flax model output classes
For example here, it is using `from dataclasses import dataclass`
https://github.com/huggingface/diffusers/blob/d8b0e4f433778f9d4555cfab642e43db3aff4833/src/diffusers/models/unet_2d_condition_flax.py#L22-L23

But transformers equivalents use `@flax.struct.dataclass`. For example [here](https://github.com/huggingface/transformers/blob/21f6f58721dd9154357576be6de54eefef1f1818/src/transformers/models/bert/modeling_flax_bert.py#L63-L64)
```py
@flax.struct.dataclass
class FlaxBertForPreTrainingOutput(ModelOutput):
```

The benefit of using `@flax.struct.dataclass` over naive `python dataclass` is that: [`jax.jit` can consume @flax.struct.dataclass](https://github.com/huggingface/transformers/pull/11886)

So the question is: should we use `@flax.struct.dataclass`  on `diffusers` as well ?

--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `src/diffusers/models/unet_2d_condition_flax.py`
Content:
```
1 from dataclasses import dataclass
2 from typing import Tuple, Union
3 
4 import flax.linen as nn
5 import jax
6 import jax.numpy as jnp
7 from flax.core.frozen_dict import FrozenDict
8 
9 from ..configuration_utils import ConfigMixin, flax_register_to_config
10 from ..modeling_flax_utils import FlaxModelMixin
11 from ..utils import BaseOutput
12 from .embeddings_flax import FlaxTimestepEmbedding, FlaxTimesteps
13 from .unet_blocks_flax import (
14     FlaxCrossAttnDownBlock2D,
15     FlaxCrossAttnUpBlock2D,
16     FlaxDownBlock2D,
17     FlaxUNetMidBlock2DCrossAttn,
18     FlaxUpBlock2D,
19 )
20 
21 
22 @dataclass
23 class FlaxUNet2DConditionOutput(BaseOutput):
24     """
25     Args:
26         sample (`jnp.ndarray` of shape `(batch_size, num_channels, height, width)`):
27             Hidden states conditioned on `encoder_hidden_states` input. Output of last layer of model.
28     """
29 
30     sample: jnp.ndarray
31 
32 
33 @flax_register_to_config
34 class FlaxUNet2DConditionModel(nn.Module, FlaxModelMixin, ConfigMixin):
35     r"""
36     FlaxUNet2DConditionModel is a conditional 2D UNet model that takes in a noisy sample, conditional state, and a
37     timestep and returns sample shaped output.
38 
39     This model inherits from [`FlaxModelMixin`]. Check the superclass documentation for the generic methods the library
40     implements for all the models (such as downloading or saving, etc.)
41 
42     Parameters:
43         sample_size (`int`, *optional*): The size of the input sample.
44         in_channels (`int`, *optional*, defaults to 4): The number of channels in the input sample.
45         out_channels (`int`, *optional*, defaults to 4): The number of channels in the output.
46         down_block_types (`Tuple[str]`, *optional*, defaults to `("CrossAttnDownBlock2D", "CrossAttnDownBlock2D", "CrossAttnDownBlock2D", "DownBlock2D")`):
47             The tuple of downsample blocks to use. The corresponding class names will be: "FlaxCrossAttnDownBlock2D",
48             "FlaxCrossAttnDownBlock2D", "FlaxCrossAttnDownBlock2D", "FlaxDownBlock2D"
49         up_block_types (`Tuple[str]`, *optional*, defaults to `("UpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D",)`):
50             The tuple of upsample blocks to use. The corresponding class names will be: "FlaxUpBlock2D",
51             "FlaxCrossAttnUpBlock2D", "FlaxCrossAttnUpBlock2D", "FlaxCrossAttnUpBlock2D"
52         block_out_channels (`Tuple[int]`, *optional*, defaults to `(320, 640, 1280, 1280)`):
53             The tuple of output channels for each block.
54         layers_per_block (`int`, *optional*, defaults to 2): The number of layers per block.
55         attention_head_dim (`int`, *optional*, defaults to 8): The dimension of the attention heads.
56         cross_attention_dim (`int`, *optional*, defaults to 768): The dimension of the cross attention features.
57         dropout (`float`, *optional*, defaults to 0): Dropout probability for down, up and bottleneck blocks.
58     """
59 
60     sample_size: int = 32
61     in_channels: int = 4
62     out_channels: int = 4
63     down_block_types: Tuple[str] = (
64         "CrossAttnDownBlock2D",
65         "CrossAttnDownBlock2D",
66         "CrossAttnDownBlock2D",
67         "DownBlock2D",
68     )
69     up_block_types: Tuple[str] = ("UpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D")
70     block_out_channels: Tuple[int] = (320, 640, 1280, 1280)
71     layers_per_block: int = 2
72     attention_head_dim: int = 8
73     cross_attention_dim: int = 1280
74     dropout: float = 0.0
75     dtype: jnp.dtype = jnp.float32
76 
77     def init_weights(self, rng: jax.random.PRNGKey) -> FrozenDict:
78         # init input tensors
79         sample_shape = (1, self.sample_size, self.sample_size, self.in_channels)
80         sample = jnp.zeros(sample_shape, dtype=jnp.float32)
81         timesteps = jnp.ones((1,), dtype=jnp.int32)
82         encoder_hidden_states = jnp.zeros((1, 1, self.cross_attention_dim), dtype=jnp.float32)
83 
84         params_rng, dropout_rng = jax.random.split(rng)
85         rngs = {"params": params_rng, "dropout": dropout_rng}
86 
87         return self.init(rngs, sample, timesteps, encoder_hidden_states)["params"]
88 
89     def setup(self):
90         block_out_channels = self.block_out_channels
91         time_embed_dim = block_out_channels[0] * 4
92 
93         # input
94         self.conv_in = nn.Conv(
95             block_out_channels[0],
96             kernel_size=(3, 3),
97             strides=(1, 1),
98             padding=((1, 1), (1, 1)),
99             dtype=self.dtype,
100         )
101 
102         # time
103         self.time_proj = FlaxTimesteps(block_out_channels[0])
104         self.time_embedding = FlaxTimestepEmbedding(time_embed_dim, dtype=self.dtype)
105 
106         # down
107         down_blocks = []
108         output_channel = block_out_channels[0]
109         for i, down_block_type in enumerate(self.down_block_types):
110             input_channel = output_channel
111             output_channel = block_out_channels[i]
112             is_final_block = i == len(block_out_channels) - 1
113 
114             if down_block_type == "CrossAttnDownBlock2D":
115                 down_block = FlaxCrossAttnDownBlock2D(
116                     in_channels=input_channel,
117                     out_channels=output_channel,
118                     dropout=self.dropout,
119                     num_layers=self.layers_per_block,
120                     attn_num_head_channels=self.attention_head_dim,
121                     add_downsample=not is_final_block,
122                     dtype=self.dtype,
123                 )
124             else:
125                 down_block = FlaxDownBlock2D(
126                     in_channels=input_channel,
127                     out_channels=output_channel,
128                     dropout=self.dropout,
129                     num_layers=self.layers_per_block,
130                     add_downsample=not is_final_block,
131                     dtype=self.dtype,
132                 )
133 
134             down_blocks.append(down_block)
135         self.down_blocks = down_blocks
136 
137         # mid
138         self.mid_block = FlaxUNetMidBlock2DCrossAttn(
139             in_channels=block_out_channels[-1],
140             dropout=self.dropout,
141             attn_num_head_channels=self.attention_head_dim,
142             dtype=self.dtype,
143         )
144 
145         # up
146         up_blocks = []
147         reversed_block_out_channels = list(reversed(block_out_channels))
148         output_channel = reversed_block_out_channels[0]
149         for i, up_block_type in enumerate(self.up_block_types):
150             prev_output_channel = output_channel
151             output_channel = reversed_block_out_channels[i]
152             input_channel = reversed_block_out_channels[min(i + 1, len(block_out_channels) - 1)]
153 
154             is_final_block = i == len(block_out_channels) - 1
155 
156             if up_block_type == "CrossAttnUpBlock2D":
157                 up_block = FlaxCrossAttnUpBlock2D(
158                     in_channels=input_channel,
159                     out_channels=output_channel,
160                     prev_output_channel=prev_output_channel,
161                     num_layers=self.layers_per_block + 1,
162                     attn_num_head_channels=self.attention_head_dim,
163                     add_upsample=not is_final_block,
164                     dropout=self.dropout,
165                     dtype=self.dtype,
166                 )
167             else:
168                 up_block = FlaxUpBlock2D(
169                     in_channels=input_channel,
170                     out_channels=output_channel,
171                     prev_output_channel=prev_output_channel,
172                     num_layers=self.layers_per_block + 1,
173                     add_upsample=not is_final_block,
174                     dropout=self.dropout,
175                     dtype=self.dtype,
176                 )
177 
178             up_blocks.append(up_block)
179             prev_output_channel = output_channel
180         self.up_blocks = up_blocks
181 
182         # out
183         self.conv_norm_out = nn.GroupNorm(num_groups=32, epsilon=1e-5)
184         self.conv_out = nn.Conv(
185             self.out_channels,
186             kernel_size=(3, 3),
187             strides=(1, 1),
188             padding=((1, 1), (1, 1)),
189             dtype=self.dtype,
190         )
191 
192     def __call__(
193         self,
194         sample,
195         timesteps,
196         encoder_hidden_states,
197         return_dict: bool = True,
198         train: bool = False,
199     ) -> Union[FlaxUNet2DConditionOutput, Tuple]:
200         """r
201         Args:
202             sample (`jnp.ndarray`): (channel, height, width) noisy inputs tensor
203             timestep (`jnp.ndarray` or `float` or `int`): timesteps
204             encoder_hidden_states (`jnp.ndarray`): (channel, height, width) encoder hidden states
205             return_dict (`bool`, *optional*, defaults to `True`):
206                 Whether or not to return a [`models.unet_2d_condition_flax.FlaxUNet2DConditionOutput`] instead of a
207                 plain tuple.
208             train (`bool`, *optional*, defaults to `False`):
209                 Use deterministic functions and disable dropout when not training.
210 
211         Returns:
212             [`~models.unet_2d_condition_flax.FlaxUNet2DConditionOutput`] or `tuple`:
213             [`~models.unet_2d_condition_flax.FlaxUNet2DConditionOutput`] if `return_dict` is True, otherwise a `tuple`.
214             When returning a tuple, the first element is the sample tensor.
215         """
216         # 1. time
217         t_emb = self.time_proj(timesteps)
218         t_emb = self.time_embedding(t_emb)
219 
220         # 2. pre-process
221         sample = self.conv_in(sample)
222 
223         # 3. down
224         down_block_res_samples = (sample,)
225         for down_block in self.down_blocks:
226             if isinstance(down_block, FlaxCrossAttnDownBlock2D):
227                 sample, res_samples = down_block(sample, t_emb, encoder_hidden_states, deterministic=not train)
228             else:
229                 sample, res_samples = down_block(sample, t_emb, deterministic=not train)
230             down_block_res_samples += res_samples
231 
232         # 4. mid
233         sample = self.mid_block(sample, t_emb, encoder_hidden_states, deterministic=not train)
234 
235         # 5. up
236         for up_block in self.up_blocks:
237             res_samples = down_block_res_samples[-(self.layers_per_block + 1) :]
238             down_block_res_samples = down_block_res_samples[: -(self.layers_per_block + 1)]
239             if isinstance(up_block, FlaxCrossAttnUpBlock2D):
240                 sample = up_block(
241                     sample,
242                     temb=t_emb,
243                     encoder_hidden_states=encoder_hidden_states,
244                     res_hidden_states_tuple=res_samples,
245                     deterministic=not train,
246                 )
247             else:
248                 sample = up_block(sample, temb=t_emb, res_hidden_states_tuple=res_samples, deterministic=not train)
249 
250         # 6. post-process
251         sample = self.conv_norm_out(sample)
252         sample = nn.silu(sample)
253         sample = self.conv_out(sample)
254 
255         if not return_dict:
256             return (sample,)
257 
258         return FlaxUNet2DConditionOutput(sample=sample)
259
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

diff --git a/src/diffusers/models/unet_2d_condition_flax.py b/src/diffusers/models/unet_2d_condition_flax.py
--- a/src/diffusers/models/unet_2d_condition_flax.py
+++ b/src/diffusers/models/unet_2d_condition_flax.py
@@ -1,6 +1,6 @@
-from dataclasses import dataclass
 from typing import Tuple, Union
 
+import flax
 import flax.linen as nn
 import jax
 import jax.numpy as jnp
@@ -19,7 +19,7 @@
 )
 
 
-@dataclass
+@flax.struct.dataclass
 class FlaxUNet2DConditionOutput(BaseOutput):
     """
     Args:

{"golden_diff": "diff --git a/src/diffusers/models/unet_2d_condition_flax.py b/src/diffusers/models/unet_2d_condition_flax.py\n--- a/src/diffusers/models/unet_2d_condition_flax.py\n+++ b/src/diffusers/models/unet_2d_condition_flax.py\n@@ -1,6 +1,6 @@\n-from dataclasses import dataclass\n from typing import Tuple, Union\n \n+import flax\n import flax.linen as nn\n import jax\n import jax.numpy as jnp\n@@ -19,7 +19,7 @@\n )\n \n \n-@dataclass\n+@flax.struct.dataclass\n class FlaxUNet2DConditionOutput(BaseOutput):\n     \"\"\"\n     Args:\n", "issue": "Question about flax model output classes\nFor example here, it is using `from dataclasses import dataclass`\r\nhttps://github.com/huggingface/diffusers/blob/d8b0e4f433778f9d4555cfab642e43db3aff4833/src/diffusers/models/unet_2d_condition_flax.py#L22-L23\r\n\r\nBut transformers equivalents use `@flax.struct.dataclass`. For example [here](https://github.com/huggingface/transformers/blob/21f6f58721dd9154357576be6de54eefef1f1818/src/transformers/models/bert/modeling_flax_bert.py#L63-L64)\r\n```py\r\n@flax.struct.dataclass\r\nclass FlaxBertForPreTrainingOutput(ModelOutput):\r\n```\r\n\r\nThe benefit of using `@flax.struct.dataclass` over naive `python dataclass` is that: [`jax.jit` can consume @flax.struct.dataclass](https://github.com/huggingface/transformers/pull/11886)\r\n\r\nSo the question is: should we use `@flax.struct.dataclass`  on `diffusers` as well ?\n", "before_files": [{"content": "from dataclasses import dataclass\nfrom typing import Tuple, Union\n\nimport flax.linen as nn\nimport jax\nimport jax.numpy as jnp\nfrom flax.core.frozen_dict import FrozenDict\n\nfrom ..configuration_utils import ConfigMixin, flax_register_to_config\nfrom ..modeling_flax_utils import FlaxModelMixin\nfrom ..utils import BaseOutput\nfrom .embeddings_flax import FlaxTimestepEmbedding, FlaxTimesteps\nfrom .unet_blocks_flax import (\n    FlaxCrossAttnDownBlock2D,\n    FlaxCrossAttnUpBlock2D,\n    FlaxDownBlock2D,\n    FlaxUNetMidBlock2DCrossAttn,\n    FlaxUpBlock2D,\n)\n\n\n@dataclass\nclass FlaxUNet2DConditionOutput(BaseOutput):\n    \"\"\"\n    Args:\n        sample (`jnp.ndarray` of shape `(batch_size, num_channels, height, width)`):\n            Hidden states conditioned on `encoder_hidden_states` input. Output of last layer of model.\n    \"\"\"\n\n    sample: jnp.ndarray\n\n\n@flax_register_to_config\nclass FlaxUNet2DConditionModel(nn.Module, FlaxModelMixin, ConfigMixin):\n    r\"\"\"\n    FlaxUNet2DConditionModel is a conditional 2D UNet model that takes in a noisy sample, conditional state, and a\n    timestep and returns sample shaped output.\n\n    This model inherits from [`FlaxModelMixin`]. Check the superclass documentation for the generic methods the library\n    implements for all the models (such as downloading or saving, etc.)\n\n    Parameters:\n        sample_size (`int`, *optional*): The size of the input sample.\n        in_channels (`int`, *optional*, defaults to 4): The number of channels in the input sample.\n        out_channels (`int`, *optional*, defaults to 4): The number of channels in the output.\n        down_block_types (`Tuple[str]`, *optional*, defaults to `(\"CrossAttnDownBlock2D\", \"CrossAttnDownBlock2D\", \"CrossAttnDownBlock2D\", \"DownBlock2D\")`):\n            The tuple of downsample blocks to use. The corresponding class names will be: \"FlaxCrossAttnDownBlock2D\",\n            \"FlaxCrossAttnDownBlock2D\", \"FlaxCrossAttnDownBlock2D\", \"FlaxDownBlock2D\"\n        up_block_types (`Tuple[str]`, *optional*, defaults to `(\"UpBlock2D\", \"CrossAttnUpBlock2D\", \"CrossAttnUpBlock2D\", \"CrossAttnUpBlock2D\",)`):\n            The tuple of upsample blocks to use. The corresponding class names will be: \"FlaxUpBlock2D\",\n            \"FlaxCrossAttnUpBlock2D\", \"FlaxCrossAttnUpBlock2D\", \"FlaxCrossAttnUpBlock2D\"\n        block_out_channels (`Tuple[int]`, *optional*, defaults to `(320, 640, 1280, 1280)`):\n            The tuple of output channels for each block.\n        layers_per_block (`int`, *optional*, defaults to 2): The number of layers per block.\n        attention_head_dim (`int`, *optional*, defaults to 8): The dimension of the attention heads.\n        cross_attention_dim (`int`, *optional*, defaults to 768): The dimension of the cross attention features.\n        dropout (`float`, *optional*, defaults to 0): Dropout probability for down, up and bottleneck blocks.\n    \"\"\"\n\n    sample_size: int = 32\n    in_channels: int = 4\n    out_channels: int = 4\n    down_block_types: Tuple[str] = (\n        \"CrossAttnDownBlock2D\",\n        \"CrossAttnDownBlock2D\",\n        \"CrossAttnDownBlock2D\",\n        \"DownBlock2D\",\n    )\n    up_block_types: Tuple[str] = (\"UpBlock2D\", \"CrossAttnUpBlock2D\", \"CrossAttnUpBlock2D\", \"CrossAttnUpBlock2D\")\n    block_out_channels: Tuple[int] = (320, 640, 1280, 1280)\n    layers_per_block: int = 2\n    attention_head_dim: int = 8\n    cross_attention_dim: int = 1280\n    dropout: float = 0.0\n    dtype: jnp.dtype = jnp.float32\n\n    def init_weights(self, rng: jax.random.PRNGKey) -> FrozenDict:\n        # init input tensors\n        sample_shape = (1, self.sample_size, self.sample_size, self.in_channels)\n        sample = jnp.zeros(sample_shape, dtype=jnp.float32)\n        timesteps = jnp.ones((1,), dtype=jnp.int32)\n        encoder_hidden_states = jnp.zeros((1, 1, self.cross_attention_dim), dtype=jnp.float32)\n\n        params_rng, dropout_rng = jax.random.split(rng)\n        rngs = {\"params\": params_rng, \"dropout\": dropout_rng}\n\n        return self.init(rngs, sample, timesteps, encoder_hidden_states)[\"params\"]\n\n    def setup(self):\n        block_out_channels = self.block_out_channels\n        time_embed_dim = block_out_channels[0] * 4\n\n        # input\n        self.conv_in = nn.Conv(\n            block_out_channels[0],\n            kernel_size=(3, 3),\n            strides=(1, 1),\n            padding=((1, 1), (1, 1)),\n            dtype=self.dtype,\n        )\n\n        # time\n        self.time_proj = FlaxTimesteps(block_out_channels[0])\n        self.time_embedding = FlaxTimestepEmbedding(time_embed_dim, dtype=self.dtype)\n\n        # down\n        down_blocks = []\n        output_channel = block_out_channels[0]\n        for i, down_block_type in enumerate(self.down_block_types):\n            input_channel = output_channel\n            output_channel = block_out_channels[i]\n            is_final_block = i == len(block_out_channels) - 1\n\n            if down_block_type == \"CrossAttnDownBlock2D\":\n                down_block = FlaxCrossAttnDownBlock2D(\n                    in_channels=input_channel,\n                    out_channels=output_channel,\n                    dropout=self.dropout,\n                    num_layers=self.layers_per_block,\n                    attn_num_head_channels=self.attention_head_dim,\n                    add_downsample=not is_final_block,\n                    dtype=self.dtype,\n                )\n            else:\n                down_block = FlaxDownBlock2D(\n                    in_channels=input_channel,\n                    out_channels=output_channel,\n                    dropout=self.dropout,\n                    num_layers=self.layers_per_block,\n                    add_downsample=not is_final_block,\n                    dtype=self.dtype,\n                )\n\n            down_blocks.append(down_block)\n        self.down_blocks = down_blocks\n\n        # mid\n        self.mid_block = FlaxUNetMidBlock2DCrossAttn(\n            in_channels=block_out_channels[-1],\n            dropout=self.dropout,\n            attn_num_head_channels=self.attention_head_dim,\n            dtype=self.dtype,\n        )\n\n        # up\n        up_blocks = []\n        reversed_block_out_channels = list(reversed(block_out_channels))\n        output_channel = reversed_block_out_channels[0]\n        for i, up_block_type in enumerate(self.up_block_types):\n            prev_output_channel = output_channel\n            output_channel = reversed_block_out_channels[i]\n            input_channel = reversed_block_out_channels[min(i + 1, len(block_out_channels) - 1)]\n\n            is_final_block = i == len(block_out_channels) - 1\n\n            if up_block_type == \"CrossAttnUpBlock2D\":\n                up_block = FlaxCrossAttnUpBlock2D(\n                    in_channels=input_channel,\n                    out_channels=output_channel,\n                    prev_output_channel=prev_output_channel,\n                    num_layers=self.layers_per_block + 1,\n                    attn_num_head_channels=self.attention_head_dim,\n                    add_upsample=not is_final_block,\n                    dropout=self.dropout,\n                    dtype=self.dtype,\n                )\n            else:\n                up_block = FlaxUpBlock2D(\n                    in_channels=input_channel,\n                    out_channels=output_channel,\n                    prev_output_channel=prev_output_channel,\n                    num_layers=self.layers_per_block + 1,\n                    add_upsample=not is_final_block,\n                    dropout=self.dropout,\n                    dtype=self.dtype,\n                )\n\n            up_blocks.append(up_block)\n            prev_output_channel = output_channel\n        self.up_blocks = up_blocks\n\n        # out\n        self.conv_norm_out = nn.GroupNorm(num_groups=32, epsilon=1e-5)\n        self.conv_out = nn.Conv(\n            self.out_channels,\n            kernel_size=(3, 3),\n            strides=(1, 1),\n            padding=((1, 1), (1, 1)),\n            dtype=self.dtype,\n        )\n\n    def __call__(\n        self,\n        sample,\n        timesteps,\n        encoder_hidden_states,\n        return_dict: bool = True,\n        train: bool = False,\n    ) -> Union[FlaxUNet2DConditionOutput, Tuple]:\n        \"\"\"r\n        Args:\n            sample (`jnp.ndarray`): (channel, height, width) noisy inputs tensor\n            timestep (`jnp.ndarray` or `float` or `int`): timesteps\n            encoder_hidden_states (`jnp.ndarray`): (channel, height, width) encoder hidden states\n            return_dict (`bool`, *optional*, defaults to `True`):\n                Whether or not to return a [`models.unet_2d_condition_flax.FlaxUNet2DConditionOutput`] instead of a\n                plain tuple.\n            train (`bool`, *optional*, defaults to `False`):\n                Use deterministic functions and disable dropout when not training.\n\n        Returns:\n            [`~models.unet_2d_condition_flax.FlaxUNet2DConditionOutput`] or `tuple`:\n            [`~models.unet_2d_condition_flax.FlaxUNet2DConditionOutput`] if `return_dict` is True, otherwise a `tuple`.\n            When returning a tuple, the first element is the sample tensor.\n        \"\"\"\n        # 1. time\n        t_emb = self.time_proj(timesteps)\n        t_emb = self.time_embedding(t_emb)\n\n        # 2. pre-process\n        sample = self.conv_in(sample)\n\n        # 3. down\n        down_block_res_samples = (sample,)\n        for down_block in self.down_blocks:\n            if isinstance(down_block, FlaxCrossAttnDownBlock2D):\n                sample, res_samples = down_block(sample, t_emb, encoder_hidden_states, deterministic=not train)\n            else:\n                sample, res_samples = down_block(sample, t_emb, deterministic=not train)\n            down_block_res_samples += res_samples\n\n        # 4. mid\n        sample = self.mid_block(sample, t_emb, encoder_hidden_states, deterministic=not train)\n\n        # 5. up\n        for up_block in self.up_blocks:\n            res_samples = down_block_res_samples[-(self.layers_per_block + 1) :]\n            down_block_res_samples = down_block_res_samples[: -(self.layers_per_block + 1)]\n            if isinstance(up_block, FlaxCrossAttnUpBlock2D):\n                sample = up_block(\n                    sample,\n                    temb=t_emb,\n                    encoder_hidden_states=encoder_hidden_states,\n                    res_hidden_states_tuple=res_samples,\n                    deterministic=not train,\n                )\n            else:\n                sample = up_block(sample, temb=t_emb, res_hidden_states_tuple=res_samples, deterministic=not train)\n\n        # 6. post-process\n        sample = self.conv_norm_out(sample)\n        sample = nn.silu(sample)\n        sample = self.conv_out(sample)\n\n        if not return_dict:\n            return (sample,)\n\n        return FlaxUNet2DConditionOutput(sample=sample)\n", "path": "src/diffusers/models/unet_2d_condition_flax.py"}], "after_files": [{"content": "from typing import Tuple, Union\n\nimport flax\nimport flax.linen as nn\nimport jax\nimport jax.numpy as jnp\nfrom flax.core.frozen_dict import FrozenDict\n\nfrom ..configuration_utils import ConfigMixin, flax_register_to_config\nfrom ..modeling_flax_utils import FlaxModelMixin\nfrom ..utils import BaseOutput\nfrom .embeddings_flax import FlaxTimestepEmbedding, FlaxTimesteps\nfrom .unet_blocks_flax import (\n    FlaxCrossAttnDownBlock2D,\n    FlaxCrossAttnUpBlock2D,\n    FlaxDownBlock2D,\n    FlaxUNetMidBlock2DCrossAttn,\n    FlaxUpBlock2D,\n)\n\n\n@flax.struct.dataclass\nclass FlaxUNet2DConditionOutput(BaseOutput):\n    \"\"\"\n    Args:\n        sample (`jnp.ndarray` of shape `(batch_size, num_channels, height, width)`):\n            Hidden states conditioned on `encoder_hidden_states` input. Output of last layer of model.\n    \"\"\"\n\n    sample: jnp.ndarray\n\n\n@flax_register_to_config\nclass FlaxUNet2DConditionModel(nn.Module, FlaxModelMixin, ConfigMixin):\n    r\"\"\"\n    FlaxUNet2DConditionModel is a conditional 2D UNet model that takes in a noisy sample, conditional state, and a\n    timestep and returns sample shaped output.\n\n    This model inherits from [`FlaxModelMixin`]. Check the superclass documentation for the generic methods the library\n    implements for all the models (such as downloading or saving, etc.)\n\n    Parameters:\n        sample_size (`int`, *optional*): The size of the input sample.\n        in_channels (`int`, *optional*, defaults to 4): The number of channels in the input sample.\n        out_channels (`int`, *optional*, defaults to 4): The number of channels in the output.\n        down_block_types (`Tuple[str]`, *optional*, defaults to `(\"CrossAttnDownBlock2D\", \"CrossAttnDownBlock2D\", \"CrossAttnDownBlock2D\", \"DownBlock2D\")`):\n            The tuple of downsample blocks to use. The corresponding class names will be: \"FlaxCrossAttnDownBlock2D\",\n            \"FlaxCrossAttnDownBlock2D\", \"FlaxCrossAttnDownBlock2D\", \"FlaxDownBlock2D\"\n        up_block_types (`Tuple[str]`, *optional*, defaults to `(\"UpBlock2D\", \"CrossAttnUpBlock2D\", \"CrossAttnUpBlock2D\", \"CrossAttnUpBlock2D\",)`):\n            The tuple of upsample blocks to use. The corresponding class names will be: \"FlaxUpBlock2D\",\n            \"FlaxCrossAttnUpBlock2D\", \"FlaxCrossAttnUpBlock2D\", \"FlaxCrossAttnUpBlock2D\"\n        block_out_channels (`Tuple[int]`, *optional*, defaults to `(320, 640, 1280, 1280)`):\n            The tuple of output channels for each block.\n        layers_per_block (`int`, *optional*, defaults to 2): The number of layers per block.\n        attention_head_dim (`int`, *optional*, defaults to 8): The dimension of the attention heads.\n        cross_attention_dim (`int`, *optional*, defaults to 768): The dimension of the cross attention features.\n        dropout (`float`, *optional*, defaults to 0): Dropout probability for down, up and bottleneck blocks.\n    \"\"\"\n\n    sample_size: int = 32\n    in_channels: int = 4\n    out_channels: int = 4\n    down_block_types: Tuple[str] = (\n        \"CrossAttnDownBlock2D\",\n        \"CrossAttnDownBlock2D\",\n        \"CrossAttnDownBlock2D\",\n        \"DownBlock2D\",\n    )\n    up_block_types: Tuple[str] = (\"UpBlock2D\", \"CrossAttnUpBlock2D\", \"CrossAttnUpBlock2D\", \"CrossAttnUpBlock2D\")\n    block_out_channels: Tuple[int] = (320, 640, 1280, 1280)\n    layers_per_block: int = 2\n    attention_head_dim: int = 8\n    cross_attention_dim: int = 1280\n    dropout: float = 0.0\n    dtype: jnp.dtype = jnp.float32\n\n    def init_weights(self, rng: jax.random.PRNGKey) -> FrozenDict:\n        # init input tensors\n        sample_shape = (1, self.sample_size, self.sample_size, self.in_channels)\n        sample = jnp.zeros(sample_shape, dtype=jnp.float32)\n        timesteps = jnp.ones((1,), dtype=jnp.int32)\n        encoder_hidden_states = jnp.zeros((1, 1, self.cross_attention_dim), dtype=jnp.float32)\n\n        params_rng, dropout_rng = jax.random.split(rng)\n        rngs = {\"params\": params_rng, \"dropout\": dropout_rng}\n\n        return self.init(rngs, sample, timesteps, encoder_hidden_states)[\"params\"]\n\n    def setup(self):\n        block_out_channels = self.block_out_channels\n        time_embed_dim = block_out_channels[0] * 4\n\n        # input\n        self.conv_in = nn.Conv(\n            block_out_channels[0],\n            kernel_size=(3, 3),\n            strides=(1, 1),\n            padding=((1, 1), (1, 1)),\n            dtype=self.dtype,\n        )\n\n        # time\n        self.time_proj = FlaxTimesteps(block_out_channels[0])\n        self.time_embedding = FlaxTimestepEmbedding(time_embed_dim, dtype=self.dtype)\n\n        # down\n        down_blocks = []\n        output_channel = block_out_channels[0]\n        for i, down_block_type in enumerate(self.down_block_types):\n            input_channel = output_channel\n            output_channel = block_out_channels[i]\n            is_final_block = i == len(block_out_channels) - 1\n\n            if down_block_type == \"CrossAttnDownBlock2D\":\n                down_block = FlaxCrossAttnDownBlock2D(\n                    in_channels=input_channel,\n                    out_channels=output_channel,\n                    dropout=self.dropout,\n                    num_layers=self.layers_per_block,\n                    attn_num_head_channels=self.attention_head_dim,\n                    add_downsample=not is_final_block,\n                    dtype=self.dtype,\n                )\n            else:\n                down_block = FlaxDownBlock2D(\n                    in_channels=input_channel,\n                    out_channels=output_channel,\n                    dropout=self.dropout,\n                    num_layers=self.layers_per_block,\n                    add_downsample=not is_final_block,\n                    dtype=self.dtype,\n                )\n\n            down_blocks.append(down_block)\n        self.down_blocks = down_blocks\n\n        # mid\n        self.mid_block = FlaxUNetMidBlock2DCrossAttn(\n            in_channels=block_out_channels[-1],\n            dropout=self.dropout,\n            attn_num_head_channels=self.attention_head_dim,\n            dtype=self.dtype,\n        )\n\n        # up\n        up_blocks = []\n        reversed_block_out_channels = list(reversed(block_out_channels))\n        output_channel = reversed_block_out_channels[0]\n        for i, up_block_type in enumerate(self.up_block_types):\n            prev_output_channel = output_channel\n            output_channel = reversed_block_out_channels[i]\n            input_channel = reversed_block_out_channels[min(i + 1, len(block_out_channels) - 1)]\n\n            is_final_block = i == len(block_out_channels) - 1\n\n            if up_block_type == \"CrossAttnUpBlock2D\":\n                up_block = FlaxCrossAttnUpBlock2D(\n                    in_channels=input_channel,\n                    out_channels=output_channel,\n                    prev_output_channel=prev_output_channel,\n                    num_layers=self.layers_per_block + 1,\n                    attn_num_head_channels=self.attention_head_dim,\n                    add_upsample=not is_final_block,\n                    dropout=self.dropout,\n                    dtype=self.dtype,\n                )\n            else:\n                up_block = FlaxUpBlock2D(\n                    in_channels=input_channel,\n                    out_channels=output_channel,\n                    prev_output_channel=prev_output_channel,\n                    num_layers=self.layers_per_block + 1,\n                    add_upsample=not is_final_block,\n                    dropout=self.dropout,\n                    dtype=self.dtype,\n                )\n\n            up_blocks.append(up_block)\n            prev_output_channel = output_channel\n        self.up_blocks = up_blocks\n\n        # out\n        self.conv_norm_out = nn.GroupNorm(num_groups=32, epsilon=1e-5)\n        self.conv_out = nn.Conv(\n            self.out_channels,\n            kernel_size=(3, 3),\n            strides=(1, 1),\n            padding=((1, 1), (1, 1)),\n            dtype=self.dtype,\n        )\n\n    def __call__(\n        self,\n        sample,\n        timesteps,\n        encoder_hidden_states,\n        return_dict: bool = True,\n        train: bool = False,\n    ) -> Union[FlaxUNet2DConditionOutput, Tuple]:\n        \"\"\"r\n        Args:\n            sample (`jnp.ndarray`): (channel, height, width) noisy inputs tensor\n            timestep (`jnp.ndarray` or `float` or `int`): timesteps\n            encoder_hidden_states (`jnp.ndarray`): (channel, height, width) encoder hidden states\n            return_dict (`bool`, *optional*, defaults to `True`):\n                Whether or not to return a [`models.unet_2d_condition_flax.FlaxUNet2DConditionOutput`] instead of a\n                plain tuple.\n            train (`bool`, *optional*, defaults to `False`):\n                Use deterministic functions and disable dropout when not training.\n\n        Returns:\n            [`~models.unet_2d_condition_flax.FlaxUNet2DConditionOutput`] or `tuple`:\n            [`~models.unet_2d_condition_flax.FlaxUNet2DConditionOutput`] if `return_dict` is True, otherwise a `tuple`.\n            When returning a tuple, the first element is the sample tensor.\n        \"\"\"\n        # 1. time\n        t_emb = self.time_proj(timesteps)\n        t_emb = self.time_embedding(t_emb)\n\n        # 2. pre-process\n        sample = self.conv_in(sample)\n\n        # 3. down\n        down_block_res_samples = (sample,)\n        for down_block in self.down_blocks:\n            if isinstance(down_block, FlaxCrossAttnDownBlock2D):\n                sample, res_samples = down_block(sample, t_emb, encoder_hidden_states, deterministic=not train)\n            else:\n                sample, res_samples = down_block(sample, t_emb, deterministic=not train)\n            down_block_res_samples += res_samples\n\n        # 4. mid\n        sample = self.mid_block(sample, t_emb, encoder_hidden_states, deterministic=not train)\n\n        # 5. up\n        for up_block in self.up_blocks:\n            res_samples = down_block_res_samples[-(self.layers_per_block + 1) :]\n            down_block_res_samples = down_block_res_samples[: -(self.layers_per_block + 1)]\n            if isinstance(up_block, FlaxCrossAttnUpBlock2D):\n                sample = up_block(\n                    sample,\n                    temb=t_emb,\n                    encoder_hidden_states=encoder_hidden_states,\n                    res_hidden_states_tuple=res_samples,\n                    deterministic=not train,\n                )\n            else:\n                sample = up_block(sample, temb=t_emb, res_hidden_states_tuple=res_samples, deterministic=not train)\n\n        # 6. post-process\n        sample = self.conv_norm_out(sample)\n        sample = nn.silu(sample)\n        sample = self.conv_out(sample)\n\n        if not return_dict:\n            return (sample,)\n\n        return FlaxUNet2DConditionOutput(sample=sample)\n", "path": "src/diffusers/models/unet_2d_condition_flax.py"}]}