




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




class       __lowerCAmelCase				(							nn.Module):
       _a           =  42
       _a           =  jnp.floataa






       def     SCREAMING_SNAKE_CASE			(      self:    int ):
         lowercase			:Union[str, Any]							=  nn.Conv(
             self.out_channels				,    kernel_size=(3, 3)				,    strides=(1, 1)				,    padding=((1, 1), (1, 1))				,    dtype=self.dtype				,    )






       def __call__(      self:    Any				,    _lowerCAmelCase:    Union[str, Any] ):
         lowercase		,			lowercase		,			lowercase		,			lowercase			:List[str]							=  hidden_states.shape
         lowercase			:Optional[Any]							=  jax.image.resize(
             __SCREAMING_SNAKE_CASE				,    shape=(batch, height * 2, width * 2, channels)				,    method="nearest"				,    )
         lowercase			:Union[str, Any]							=  self.conv(__SCREAMING_SNAKE_CASE )
         return hidden_states







class       __lowerCAmelCase				(							nn.Module):
       _a           =  42
       _a           =  jnp.floataa






       def     SCREAMING_SNAKE_CASE			(      self:    Any ):
         lowercase			:str							=  nn.Conv(
             self.out_channels				,    kernel_size=(3, 3)				,    strides=(2, 2)				,    padding=((1, 1), (1, 1))				,    dtype=self.dtype				,    )






       def __call__(      self:    Any				,    _lowerCAmelCase:    Optional[int] ):
         lowercase			:int							=  self.conv(__SCREAMING_SNAKE_CASE )
         return hidden_states







class       __lowerCAmelCase				(							nn.Module):
       _a           =  42
       _a           =  None
       _a           =  0.0
       _a           =  None
       _a           =  jnp.floataa






       def     SCREAMING_SNAKE_CASE			(      self:    Union[str, Any] ):
         lowercase			:Union[str, Any]							=  self.in_channels if self.out_channels is None else self.out_channels

         lowercase			:Optional[Any]							=  nn.GroupNorm(num_groups=32				,    epsilon=1e-5 )
         lowercase			:Optional[int]							=  nn.Conv(
             __SCREAMING_SNAKE_CASE				,    kernel_size=(3, 3)				,    strides=(1, 1)				,    padding=((1, 1), (1, 1))				,    dtype=self.dtype				,    )

         lowercase			:str							=  nn.Dense(__SCREAMING_SNAKE_CASE				,    dtype=self.dtype )

         lowercase			:Union[str, Any]							=  nn.GroupNorm(num_groups=32				,    epsilon=1e-5 )
         lowercase			:Any							=  nn.Dropout(self.dropout_prob )
         lowercase			:Tuple							=  nn.Conv(
             __SCREAMING_SNAKE_CASE				,    kernel_size=(3, 3)				,    strides=(1, 1)				,    padding=((1, 1), (1, 1))				,    dtype=self.dtype				,    )

         lowercase			:Optional[int]							=  self.in_channels != out_channels if self.use_nin_shortcut is None else self.use_nin_shortcut

         lowercase			:Tuple							=  None
         if use_nin_shortcut:
           lowercase			:int							=  nn.Conv(
               __SCREAMING_SNAKE_CASE				,    kernel_size=(1, 1)				,    strides=(1, 1)				,    padding="VALID"				,    dtype=self.dtype				,    )






       def __call__(      self:    List[str]				,    _lowerCAmelCase:    str				,    _lowerCAmelCase:    List[str]				,    _lowerCAmelCase:    Union[str, Any]=True ):
         lowercase			:List[Any]							=  hidden_states
         lowercase			:Dict							=  self.norma(__SCREAMING_SNAKE_CASE )
         lowercase			:int							=  nn.swish(__SCREAMING_SNAKE_CASE )
         lowercase			:Tuple							=  self.conva(__SCREAMING_SNAKE_CASE )

         lowercase			:List[Any]							=  self.time_emb_proj(nn.swish(__SCREAMING_SNAKE_CASE ) )
         lowercase			:Tuple							=  jnp.expand_dims(jnp.expand_dims(__SCREAMING_SNAKE_CASE				,    1 )				,    1 )
         lowercase			:Dict							=  hidden_states + temb

         lowercase			:Optional[int]							=  self.norma(__SCREAMING_SNAKE_CASE )
         lowercase			:Union[str, Any]							=  nn.swish(__SCREAMING_SNAKE_CASE )
         lowercase			:List[Any]							=  self.dropout(__SCREAMING_SNAKE_CASE				,    __SCREAMING_SNAKE_CASE )
         lowercase			:Any							=  self.conva(__SCREAMING_SNAKE_CASE )

         if self.conv_shortcut is not None:
           lowercase			:Optional[int]							=  self.conv_shortcut(__SCREAMING_SNAKE_CASE )

         return hidden_states + residual






'''simple docstring'''



import unittest

import numpy as np
import torch

from diffusers import PNDMPipeline, PNDMScheduler, UNetaDModel
from diffusers.utils.testing_utils import enable_full_determinism, require_torch, slow, torch_device


enable_full_determinism()







class 						lowerCAmelCase__   (      unittest.TestCase					):



 """simple docstring"""



 @property
 def       UpperCAmelCase__		(		self					:       List[str]			)  -> Optional[int]:


      """simple docstring"""




      torch.manual_seed(0			)
      __SCREAMING_SNAKE_CASE     =      UNetaDModel(
          block_out_channels=(32, 64)					,							layers_per_block=2					,							sample_size=32					,							in_channels=3					,							out_channels=3					,							down_block_types=("""DownBlock2D""", """AttnDownBlock2D""")					,							up_block_types=("""AttnUpBlock2D""", """UpBlock2D""")					,							)
      return model


 def       UpperCAmelCase__		(		self					:       List[Any]			)  -> List[str]:


      """simple docstring"""




      __SCREAMING_SNAKE_CASE     =      self.dummy_uncond_unet
      __SCREAMING_SNAKE_CASE     =      PNDMScheduler()

      __SCREAMING_SNAKE_CASE     =      PNDMPipeline(unet=__SCREAMING_SNAKE_CASE					,							scheduler=__SCREAMING_SNAKE_CASE			)
      pndm.to(__SCREAMING_SNAKE_CASE			)
      pndm.set_progress_bar_config(disable=__SCREAMING_SNAKE_CASE			)

      __SCREAMING_SNAKE_CASE     =      torch.manual_seed(0			)
      __SCREAMING_SNAKE_CASE     =      pndm(generator=__SCREAMING_SNAKE_CASE					,							num_inference_steps=20					,							output_type="""numpy"""			).images

      __SCREAMING_SNAKE_CASE     =      torch.manual_seed(0			)
      __SCREAMING_SNAKE_CASE     =      pndm(generator=__SCREAMING_SNAKE_CASE					,							num_inference_steps=20					,							output_type="""numpy"""					,							return_dict=__SCREAMING_SNAKE_CASE			)[0]

      __SCREAMING_SNAKE_CASE     =      image[0, -3:, -3:, -1]
      __SCREAMING_SNAKE_CASE     =      image_from_tuple[0, -3:, -3:, -1]

      assert image.shape == (1, 32, 32, 3)
      __SCREAMING_SNAKE_CASE     =      np.array([1.0, 1.0, 0.0, 1.0, 0.0, 1.0, 0.0, 0.0, 0.0]			)

      assert np.abs(image_slice.flatten() - expected_slice			).max() < 1E-2
      assert np.abs(image_from_tuple_slice.flatten() - expected_slice			).max() < 1E-2







@slow
@require_torch
class 						lowerCAmelCase__   (      unittest.TestCase					):



 """simple docstring"""



 def       UpperCAmelCase__		(		self					:       Optional[int]			)  -> List[Any]:


      """simple docstring"""




      __SCREAMING_SNAKE_CASE     =      """google/ddpm-cifar10-32"""

      __SCREAMING_SNAKE_CASE     =      UNetaDModel.from_pretrained(__SCREAMING_SNAKE_CASE			)
      __SCREAMING_SNAKE_CASE     =      PNDMScheduler()

      __SCREAMING_SNAKE_CASE     =      PNDMPipeline(unet=__SCREAMING_SNAKE_CASE					,							scheduler=__SCREAMING_SNAKE_CASE			)
      pndm.to(__SCREAMING_SNAKE_CASE			)
      pndm.set_progress_bar_config(disable=__SCREAMING_SNAKE_CASE			)
      __SCREAMING_SNAKE_CASE     =      torch.manual_seed(0			)
      __SCREAMING_SNAKE_CASE     =      pndm(generator=__SCREAMING_SNAKE_CASE					,							output_type="""numpy"""			).images

      __SCREAMING_SNAKE_CASE     =      image[0, -3:, -3:, -1]

      assert image.shape == (1, 32, 32, 3)
      __SCREAMING_SNAKE_CASE     =      np.array([0.1564, 0.14645, 0.1406, 0.14715, 0.12425, 0.14045, 0.13115, 0.12175, 0.125]			)

      assert np.abs(image_slice.flatten() - expected_slice			).max() < 1E-2






