

import flax.linen as nn
import jax
import jax.numpy as jnp




class      UpperCamelCase (     nn.Module   ):
    __UpperCamelCase         =42
    __UpperCamelCase         =jnp.floataa

    def UpperCamelCase			(			self :      int    ):

     """simple docstring"""






     SCREAMING_SNAKE_CASE								=			nn.Conv(
         self.out_channels							,    kernel_size=(3, 3)							,    strides=(1, 1)							,    padding=((1, 1), (1, 1))							,    dtype=self.dtype							,    )




    def __call__(			self :      List[Any]							,    snake_case__ :      int    ):

     """simple docstring"""






     SCREAMING_SNAKE_CASE						,	SCREAMING_SNAKE_CASE						,	SCREAMING_SNAKE_CASE						,	SCREAMING_SNAKE_CASE								=			hidden_states.shape
     SCREAMING_SNAKE_CASE								=			jax.image.resize(
         snake_case__							,    shape=(batch, height * 2, width * 2, channels)							,    method='nearest'							,    )
     SCREAMING_SNAKE_CASE								=			self.conv(snake_case__    )
     return hidden_states




class      UpperCamelCase (     nn.Module   ):
    __UpperCamelCase         =42
    __UpperCamelCase         =jnp.floataa

    def UpperCamelCase			(			self :      str    ):

     """simple docstring"""






     SCREAMING_SNAKE_CASE								=			nn.Conv(
         self.out_channels							,    kernel_size=(3, 3)							,    strides=(2, 2)							,    padding=((1, 1), (1, 1))							,    dtype=self.dtype							,    )




    def __call__(			self :      str							,    snake_case__ :      Union[str, Any]    ):

     """simple docstring"""






     SCREAMING_SNAKE_CASE								=			self.conv(snake_case__    )
     return hidden_states




class      UpperCamelCase (     nn.Module   ):
    __UpperCamelCase         =42
    __UpperCamelCase         =None
    __UpperCamelCase         =0.0
    __UpperCamelCase         =None
    __UpperCamelCase         =jnp.floataa

    def UpperCamelCase			(			self :      List[str]    ):

     """simple docstring"""






     SCREAMING_SNAKE_CASE								=			self.in_channels if self.out_channels is None else self.out_channels

     SCREAMING_SNAKE_CASE								=			nn.GroupNorm(num_groups=3_2							,    epsilon=1E-5    )
     SCREAMING_SNAKE_CASE								=			nn.Conv(
         snake_case__							,    kernel_size=(3, 3)							,    strides=(1, 1)							,    padding=((1, 1), (1, 1))							,    dtype=self.dtype							,    )

     SCREAMING_SNAKE_CASE								=			nn.Dense(snake_case__							,    dtype=self.dtype    )

     SCREAMING_SNAKE_CASE								=			nn.GroupNorm(num_groups=3_2							,    epsilon=1E-5    )
     SCREAMING_SNAKE_CASE								=			nn.Dropout(self.dropout_prob    )
     SCREAMING_SNAKE_CASE								=			nn.Conv(
         snake_case__							,    kernel_size=(3, 3)							,    strides=(1, 1)							,    padding=((1, 1), (1, 1))							,    dtype=self.dtype							,    )

     SCREAMING_SNAKE_CASE								=			self.in_channels != out_channels if self.use_nin_shortcut is None else self.use_nin_shortcut

     SCREAMING_SNAKE_CASE								=			None
     if use_nin_shortcut:
      SCREAMING_SNAKE_CASE								=			nn.Conv(
          snake_case__							,    kernel_size=(1, 1)							,    strides=(1, 1)							,    padding='VALID'							,    dtype=self.dtype							,    )




    def __call__(			self :      List[Any]							,    snake_case__ :      Optional[Any]							,    snake_case__ :      Union[str, Any]							,    snake_case__ :      List[Any]=True    ):

     """simple docstring"""






     SCREAMING_SNAKE_CASE								=			hidden_states
     SCREAMING_SNAKE_CASE								=			self.norma(snake_case__    )
     SCREAMING_SNAKE_CASE								=			nn.swish(snake_case__    )
     SCREAMING_SNAKE_CASE								=			self.conva(snake_case__    )

     SCREAMING_SNAKE_CASE								=			self.time_emb_proj(nn.swish(snake_case__    )    )
     SCREAMING_SNAKE_CASE								=			jnp.expand_dims(jnp.expand_dims(snake_case__							,    1    )							,    1    )
     SCREAMING_SNAKE_CASE								=			hidden_states + temb

     SCREAMING_SNAKE_CASE								=			self.norma(snake_case__    )
     SCREAMING_SNAKE_CASE								=			nn.swish(snake_case__    )
     SCREAMING_SNAKE_CASE								=			self.dropout(snake_case__							,    snake_case__    )
     SCREAMING_SNAKE_CASE								=			self.conva(snake_case__    )

     if self.conv_shortcut is not None:
      SCREAMING_SNAKE_CASE								=			self.conv_shortcut(snake_case__    )

     return hidden_states + residual




# limitations under the License.

# NOTE: This file is deprecated and will be removed in a future version.
# It only exists so that temporarely `from diffusers.pipelines import DiffusionPipeline` works

from .pipelines import DiffusionPipeline, ImagePipelineOutput  # noqa: F401
from .utils import deprecate


deprecate(
    "pipelines_utils",
    "0.22.0",
    "Importing `DiffusionPipeline` or `ImagePipelineOutput` from diffusers.pipeline_utils is deprecated. Please import from diffusers.pipelines.pipeline_utils instead.",
    standard_warn=False,
    stacklevel=3,
)

