

'''simple docstring'''



import torch
from torch import nn

from ...configuration_utils import ConfigMixin, register_to_config
from ...models import ModelMixin







class 						lowerCAmelCase__   (      a ,					a					):



 """simple docstring"""



 @register_to_config
 def __init__(		self					:       List[str]					,							*,
     __SCREAMING_SNAKE_CASE					:       int = 4					,							__SCREAMING_SNAKE_CASE					:       int = 768					,							__SCREAMING_SNAKE_CASE					:       int					,							__SCREAMING_SNAKE_CASE					:       Any					,							)  -> List[Any]:


      """simple docstring"""




      super().__init__()

      __SCREAMING_SNAKE_CASE     =      nn.Parameter(torch.zeros(__SCREAMING_SNAKE_CASE			)			)

      # parameters for additional clip time embeddings
      __SCREAMING_SNAKE_CASE     =      nn.Linear(__SCREAMING_SNAKE_CASE					,							__SCREAMING_SNAKE_CASE			)
      __SCREAMING_SNAKE_CASE     =      nn.Linear(__SCREAMING_SNAKE_CASE					,							__SCREAMING_SNAKE_CASE			)

      # parameters for encoder hidden states
      __SCREAMING_SNAKE_CASE     =      clip_extra_context_tokens
      __SCREAMING_SNAKE_CASE     =      nn.Linear(
          __SCREAMING_SNAKE_CASE					,							self.clip_extra_context_tokens * cross_attention_dim			)
      __SCREAMING_SNAKE_CASE     =      nn.Linear(__SCREAMING_SNAKE_CASE					,							__SCREAMING_SNAKE_CASE			)
      __SCREAMING_SNAKE_CASE     =      nn.LayerNorm(__SCREAMING_SNAKE_CASE			)


 def       UpperCAmelCase__		(		self					:       List[str]					,							*, __SCREAMING_SNAKE_CASE					:       List[Any]					,							__SCREAMING_SNAKE_CASE					:       Optional[int]					,							__SCREAMING_SNAKE_CASE					:       List[Any]					,							__SCREAMING_SNAKE_CASE					:       Optional[int]			)  -> Any:


      """simple docstring"""




      if do_classifier_free_guidance:
           # Add the classifier free guidance embeddings to the image embeddings
           __SCREAMING_SNAKE_CASE     =      image_embeddings.shape[0]
           __SCREAMING_SNAKE_CASE     =      self.learned_classifier_free_guidance_embeddings.unsqueeze(0			)
           __SCREAMING_SNAKE_CASE     =      classifier_free_guidance_embeddings.expand(
               __SCREAMING_SNAKE_CASE					,							-1			)
           __SCREAMING_SNAKE_CASE     =      torch.cat([classifier_free_guidance_embeddings, image_embeddings]					,							dim=0			)

      # The image embeddings batch size and the text embeddings batch size are equal
      assert image_embeddings.shape[0] == prompt_embeds.shape[0]

      __SCREAMING_SNAKE_CASE     =      prompt_embeds.shape[0]

      # "Specifically, we modify the architecture described in Nichol et al. (2021) by projecting and
      # adding CLIP embeddings to the existing timestep embedding, ...
      __SCREAMING_SNAKE_CASE     =      self.embedding_proj(__SCREAMING_SNAKE_CASE			)
      __SCREAMING_SNAKE_CASE     =      self.clip_image_embeddings_project_to_time_embeddings(__SCREAMING_SNAKE_CASE			)
      __SCREAMING_SNAKE_CASE     =      time_projected_image_embeddings + time_projected_prompt_embeds

      # ... and by projecting CLIP embeddings into four
      # extra tokens of context that are concatenated to the sequence of outputs from the GLIDE text encoder"
      __SCREAMING_SNAKE_CASE     =      self.clip_extra_context_tokens_proj(__SCREAMING_SNAKE_CASE			)
      __SCREAMING_SNAKE_CASE     =      clip_extra_context_tokens.reshape(__SCREAMING_SNAKE_CASE					,							-1					,							self.clip_extra_context_tokens			)
      __SCREAMING_SNAKE_CASE     =      clip_extra_context_tokens.permute(0					,							2					,							1			)

      __SCREAMING_SNAKE_CASE     =      self.encoder_hidden_states_proj(__SCREAMING_SNAKE_CASE			)
      __SCREAMING_SNAKE_CASE     =      self.text_encoder_hidden_states_norm(__SCREAMING_SNAKE_CASE			)
      __SCREAMING_SNAKE_CASE     =      torch.cat([clip_extra_context_tokens, text_encoder_hidden_states]					,							dim=1			)

      return text_encoder_hidden_states, additive_clip_time_embeddings









'''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






