

'''simple docstring'''



import unittest

import numpy as np
import torch
from torch import nn
from transformers import (
    CLIPImageProcessor,
    CLIPTextConfig,
    CLIPTextModelWithProjection,
    CLIPTokenizer,
    CLIPVisionConfig,
    CLIPVisionModelWithProjection,
)

from diffusers import KandinskyVaaPriorPipeline, PriorTransformer, UnCLIPScheduler
from diffusers.utils import torch_device
from diffusers.utils.testing_utils import enable_full_determinism, skip_mps

from ..test_pipelines_common import PipelineTesterMixin


enable_full_determinism()







class 						lowerCAmelCase__   (      a ,					unittest.TestCase					):



 """simple docstring"""



 lowerCAmelCase__  					=      KandinskyVaaPriorPipeline
 lowerCAmelCase__  					=      ["prompt"]
 lowerCAmelCase__  					=      ["prompt", "negative_prompt"]
 lowerCAmelCase__  					=      [
     "num_images_per_prompt",
     "generator",
     "num_inference_steps",
     "latents",
     "negative_prompt",
     "guidance_scale",
     "output_type",
     "return_dict",
 ]
 lowerCAmelCase__  					=      False
 @property
 def       UpperCAmelCase__		(		self					:       Union[str, Any]			)  -> Tuple:


      """simple docstring"""




      return 32
 @property
 def       UpperCAmelCase__		(		self					:       Optional[int]			)  -> Dict:


      """simple docstring"""




      return 32
 @property
 def       UpperCAmelCase__		(		self					:       List[Any]			)  -> str:


      """simple docstring"""




      return self.time_input_dim
 @property
 def       UpperCAmelCase__		(		self					:       Optional[int]			)  -> List[Any]:


      """simple docstring"""




      return self.time_input_dim * 4
 @property
 def       UpperCAmelCase__		(		self					:       int			)  -> Union[str, Any]:


      """simple docstring"""




      return 100
 @property
 def       UpperCAmelCase__		(		self					:       Union[str, Any]			)  -> Optional[int]:


      """simple docstring"""




      __SCREAMING_SNAKE_CASE     =      CLIPTokenizer.from_pretrained("""hf-internal-testing/tiny-random-clip"""			)
      return tokenizer
 @property
 def       UpperCAmelCase__		(		self					:       Union[str, Any]			)  -> Union[str, Any]:


      """simple docstring"""




      torch.manual_seed(0			)
      __SCREAMING_SNAKE_CASE     =      CLIPTextConfig(
          bos_token_id=0					,							eos_token_id=2					,							hidden_size=self.text_embedder_hidden_size					,							projection_dim=self.text_embedder_hidden_size					,							intermediate_size=37					,							layer_norm_eps=1E-05					,							num_attention_heads=4					,							num_hidden_layers=5					,							pad_token_id=1					,							vocab_size=1_000					,							)
      return CLIPTextModelWithProjection(__SCREAMING_SNAKE_CASE			)
 @property
 def       UpperCAmelCase__		(		self					:       List[Any]			)  -> List[str]:


      """simple docstring"""




      torch.manual_seed(0			)

      __SCREAMING_SNAKE_CASE     =      {
          """num_attention_heads""": 2,
          """attention_head_dim""": 12,
          """embedding_dim""": self.text_embedder_hidden_size,
          """num_layers""": 1,
      }

      __SCREAMING_SNAKE_CASE     =      PriorTransformer(**__SCREAMING_SNAKE_CASE			)
      # clip_std and clip_mean is initialized to be 0 so PriorTransformer.post_process_latents will always return 0 - set clip_std to be 1 so it won't return 0
      __SCREAMING_SNAKE_CASE     =      nn.Parameter(torch.ones(model.clip_std.shape			)			)
      return model
 @property
 def       UpperCAmelCase__		(		self					:       Tuple			)  -> Optional[int]:


      """simple docstring"""




      torch.manual_seed(0			)
      __SCREAMING_SNAKE_CASE     =      CLIPVisionConfig(
          hidden_size=self.text_embedder_hidden_size					,							image_size=224					,							projection_dim=self.text_embedder_hidden_size					,							intermediate_size=37					,							num_attention_heads=4					,							num_channels=3					,							num_hidden_layers=5					,							patch_size=14					,							)

      __SCREAMING_SNAKE_CASE     =      CLIPVisionModelWithProjection(__SCREAMING_SNAKE_CASE			)
      return model
 @property
 def       UpperCAmelCase__		(		self					:       Dict			)  -> str:


      """simple docstring"""




      __SCREAMING_SNAKE_CASE     =      CLIPImageProcessor(
          crop_size=224					,							do_center_crop=__SCREAMING_SNAKE_CASE					,							do_normalize=__SCREAMING_SNAKE_CASE					,							do_resize=__SCREAMING_SNAKE_CASE					,							image_mean=[0.48145466, 0.4578275, 0.40821073]					,							image_std=[0.26862954, 0.26130258, 0.27577711]					,							resample=3					,							size=224					,							)

      return image_processor
 def       UpperCAmelCase__		(		self					:       Dict			)  -> Optional[Any]:


      """simple docstring"""




      __SCREAMING_SNAKE_CASE     =      self.dummy_prior
      __SCREAMING_SNAKE_CASE     =      self.dummy_image_encoder
      __SCREAMING_SNAKE_CASE     =      self.dummy_text_encoder
      __SCREAMING_SNAKE_CASE     =      self.dummy_tokenizer
      __SCREAMING_SNAKE_CASE     =      self.dummy_image_processor

      __SCREAMING_SNAKE_CASE     =      UnCLIPScheduler(
          variance_type="""fixed_small_log"""					,							prediction_type="""sample"""					,							num_train_timesteps=1_000					,							clip_sample=__SCREAMING_SNAKE_CASE					,							clip_sample_range=10.0					,							)

      __SCREAMING_SNAKE_CASE     =      {
          """prior""": prior,
          """image_encoder""": image_encoder,
          """text_encoder""": text_encoder,
          """tokenizer""": tokenizer,
          """scheduler""": scheduler,
          """image_processor""": image_processor,
      }

      return components
 def       UpperCAmelCase__		(		self					:       Optional[int]					,							__SCREAMING_SNAKE_CASE					:       List[str]					,							__SCREAMING_SNAKE_CASE					:       Union[str, Any]=0			)  -> int:


      """simple docstring"""




      if str(__SCREAMING_SNAKE_CASE			).startswith("""mps"""			):
           __SCREAMING_SNAKE_CASE     =      torch.manual_seed(__SCREAMING_SNAKE_CASE			)
      else:
           __SCREAMING_SNAKE_CASE     =      torch.Generator(device=__SCREAMING_SNAKE_CASE			).manual_seed(__SCREAMING_SNAKE_CASE			)
      __SCREAMING_SNAKE_CASE     =      {
          """prompt""": """horse""",
          """generator""": generator,
          """guidance_scale""": 4.0,
          """num_inference_steps""": 2,
          """output_type""": """np""",
      }
      return inputs
 def       UpperCAmelCase__		(		self					:       Optional[Any]			)  -> int:


      """simple docstring"""




      __SCREAMING_SNAKE_CASE     =      """cpu"""

      __SCREAMING_SNAKE_CASE     =      self.get_dummy_components()

      __SCREAMING_SNAKE_CASE     =      self.pipeline_class(**__SCREAMING_SNAKE_CASE			)
      __SCREAMING_SNAKE_CASE     =      pipe.to(__SCREAMING_SNAKE_CASE			)

      pipe.set_progress_bar_config(disable=__SCREAMING_SNAKE_CASE			)

      __SCREAMING_SNAKE_CASE     =      pipe(**self.get_dummy_inputs(__SCREAMING_SNAKE_CASE			)			)
      __SCREAMING_SNAKE_CASE     =      output.image_embeds

      __SCREAMING_SNAKE_CASE     =      pipe(
          **self.get_dummy_inputs(__SCREAMING_SNAKE_CASE			)					,							return_dict=__SCREAMING_SNAKE_CASE					,							)[0]

      __SCREAMING_SNAKE_CASE     =      image[0, -10:]
      __SCREAMING_SNAKE_CASE     =      image_from_tuple[0, -10:]

      assert image.shape == (1, 32)

      __SCREAMING_SNAKE_CASE     =      np.array(
          [-0.0532, 1.7120, 0.3656, -1.0852, -0.8946, -1.1756, 0.4348, 0.2482, 0.5146, -0.1156]			)

      assert np.abs(image_slice.flatten() - expected_slice			).max() < 1E-2
      assert np.abs(image_from_tuple_slice.flatten() - expected_slice			).max() < 1E-2
 @skip_mps
 def       UpperCAmelCase__		(		self					:       List[Any]			)  -> Optional[Any]:


      """simple docstring"""




      __SCREAMING_SNAKE_CASE     =      torch_device == """cpu"""
      __SCREAMING_SNAKE_CASE     =      True
      __SCREAMING_SNAKE_CASE     =      False

      self._test_inference_batch_single_identical(
          test_max_difference=__SCREAMING_SNAKE_CASE					,							relax_max_difference=__SCREAMING_SNAKE_CASE					,							test_mean_pixel_difference=__SCREAMING_SNAKE_CASE					,							)


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


      """simple docstring"""




      __SCREAMING_SNAKE_CASE     =      torch_device == """cpu"""
      __SCREAMING_SNAKE_CASE     =      False

      self._test_attention_slicing_forward_pass(
          test_max_difference=__SCREAMING_SNAKE_CASE					,							test_mean_pixel_difference=__SCREAMING_SNAKE_CASE					,							)









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






