





import gc
import unittest

import numpy as np
import torch
from transformers import CLIPTextConfig, CLIPTextModel, CLIPTokenizer

from diffusers import TransformeraDModel, VQDiffusionPipeline, VQDiffusionScheduler, VQModel
from diffusers.pipelines.vq_diffusion.pipeline_vq_diffusion import LearnedClassifierFreeSamplingEmbeddings
from diffusers.utils import load_numpy, slow, torch_device
from diffusers.utils.testing_utils import require_torch_gpu


__lowerCamelCase :  int										=  False



class       A__     ( unittest.TestCase    ):





      def  __UpperCamelCase(				self     ):
       '''simple docstring'''




       super().tearDown()
       gc.collect()
       torch.cuda.empty_cache()





      @property
      def  __UpperCamelCase(				self     ):
       '''simple docstring'''




       return 12





      @property
      def  __UpperCamelCase(				self     ):
       '''simple docstring'''




       return 12





      @property
      def  __UpperCamelCase(				self     ):
       '''simple docstring'''




       return 32





      @property
      def  __UpperCamelCase(				self     ):
       '''simple docstring'''




       torch.manual_seed(0     )
       UpperCamelCase			:      List[str] 			=     VQModel(
           block_out_channels=[32, 64] ,							in_channels=3 ,							out_channels=3 ,							down_block_types=["DownEncoderBlock2D", "DownEncoderBlock2D"] ,							up_block_types=["UpDecoderBlock2D", "UpDecoderBlock2D"] ,							latent_channels=3 ,							num_vq_embeddings=self.num_embed ,							vq_embed_dim=3 ,							)
       return model





      @property
      def  __UpperCamelCase(				self     ):
       '''simple docstring'''




       UpperCamelCase			:      Tuple 			=     CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip"     )
       return tokenizer





      @property
      def  __UpperCamelCase(				self     ):
       '''simple docstring'''




       torch.manual_seed(0     )
       UpperCamelCase			:      Dict 			=     CLIPTextConfig(
           bos_token_id=0 ,							eos_token_id=2 ,							hidden_size=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=1000 ,							)
       return CLIPTextModel(A_     )





      @property
      def  __UpperCamelCase(				self     ):
       '''simple docstring'''




       torch.manual_seed(0     )

       UpperCamelCase			:      Optional[int] 			=     12
       UpperCamelCase			:      Optional[int] 			=     12

       UpperCamelCase			:      Dict 			=     {
           "attention_bias": True,
           "cross_attention_dim": 32,
           "attention_head_dim": height * width,
           "num_attention_heads": 1,
           "num_vector_embeds": self.num_embed,
           "num_embeds_ada_norm": self.num_embeds_ada_norm,
           "norm_num_groups": 32,
           "sample_size": width,
           "activation_fn": "geglu-approximate",
       }

       UpperCamelCase			:      Any 			=     TransformeraDModel(**A_     )
       return model





      def  __UpperCamelCase(				self     ):
       '''simple docstring'''




       UpperCamelCase			:      Optional[Any] 			=     "cpu"

       UpperCamelCase			:      Union[str, Any] 			=     self.dummy_vqvae
       UpperCamelCase			:      Dict 			=     self.dummy_text_encoder
       UpperCamelCase			:      List[str] 			=     self.dummy_tokenizer
       UpperCamelCase			:      str 			=     self.dummy_transformer
       UpperCamelCase			:      Optional[int] 			=     VQDiffusionScheduler(self.num_embed     )
       UpperCamelCase			:      Any 			=     LearnedClassifierFreeSamplingEmbeddings(learnable=A_     )

       UpperCamelCase			:      Any 			=     VQDiffusionPipeline(
           vqvae=A_ ,							text_encoder=A_ ,							tokenizer=A_ ,							transformer=A_ ,							scheduler=A_ ,							learned_classifier_free_sampling_embeddings=A_ ,							)
       UpperCamelCase			:      Any 			=     pipe.to(A_     )
       pipe.set_progress_bar_config(disable=A_     )

       UpperCamelCase			:      Union[str, Any] 			=     "teddy bear playing in the pool"

       UpperCamelCase			:      int 			=     torch.Generator(device=A_     ).manual_seed(0     )
       UpperCamelCase			:      List[str] 			=     pipe([prompt] ,							generator=A_ ,							num_inference_steps=2 ,							output_type="np"     )
       UpperCamelCase			:      List[str] 			=     output.images

       UpperCamelCase			:      Union[str, Any] 			=     torch.Generator(device=A_     ).manual_seed(0     )
       UpperCamelCase			:      str 			=     pipe(
           [prompt] ,							generator=A_ ,							output_type="np" ,							return_dict=A_ ,							num_inference_steps=2     )[0]

       UpperCamelCase			:      Optional[int] 			=     image[0, -3:, -3:, -1]
       UpperCamelCase			:      Optional[Any] 			=     image_from_tuple[0, -3:, -3:, -1]

       assert image.shape == (1, 24, 24, 3)

       UpperCamelCase			:      Tuple 			=     np.array([0.65_51, 0.61_68, 0.50_08, 0.56_76, 0.56_59, 0.42_95, 0.60_73, 0.55_99, 0.49_92]     )

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





      def  __UpperCamelCase(				self     ):
       '''simple docstring'''




       UpperCamelCase			:      str 			=     "cpu"

       UpperCamelCase			:      Tuple 			=     self.dummy_vqvae
       UpperCamelCase			:      Dict 			=     self.dummy_text_encoder
       UpperCamelCase			:      Union[str, Any] 			=     self.dummy_tokenizer
       UpperCamelCase			:      Dict 			=     self.dummy_transformer
       UpperCamelCase			:      List[Any] 			=     VQDiffusionScheduler(self.num_embed     )
       UpperCamelCase			:      Tuple 			=     LearnedClassifierFreeSamplingEmbeddings(
           learnable=A_ ,							hidden_size=self.text_embedder_hidden_size ,							length=tokenizer.model_max_length     )

       UpperCamelCase			:      List[str] 			=     VQDiffusionPipeline(
           vqvae=A_ ,							text_encoder=A_ ,							tokenizer=A_ ,							transformer=A_ ,							scheduler=A_ ,							learned_classifier_free_sampling_embeddings=A_ ,							)
       UpperCamelCase			:      int 			=     pipe.to(A_     )
       pipe.set_progress_bar_config(disable=A_     )

       UpperCamelCase			:      List[str] 			=     "teddy bear playing in the pool"

       UpperCamelCase			:      List[str] 			=     torch.Generator(device=A_     ).manual_seed(0     )
       UpperCamelCase			:      List[str] 			=     pipe([prompt] ,							generator=A_ ,							num_inference_steps=2 ,							output_type="np"     )
       UpperCamelCase			:      Any 			=     output.images

       UpperCamelCase			:      List[Any] 			=     torch.Generator(device=A_     ).manual_seed(0     )
       UpperCamelCase			:      int 			=     pipe(
           [prompt] ,							generator=A_ ,							output_type="np" ,							return_dict=A_ ,							num_inference_steps=2     )[0]

       UpperCamelCase			:      Any 			=     image[0, -3:, -3:, -1]
       UpperCamelCase			:      Tuple 			=     image_from_tuple[0, -3:, -3:, -1]

       assert image.shape == (1, 24, 24, 3)

       UpperCamelCase			:      Optional[Any] 			=     np.array([0.66_93, 0.60_75, 0.49_59, 0.57_01, 0.55_83, 0.43_33, 0.61_71, 0.56_84, 0.49_88]     )

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



@slow
@require_torch_gpu
class       A__     ( unittest.TestCase    ):





      def  __UpperCamelCase(				self     ):
       '''simple docstring'''




       super().tearDown()
       gc.collect()
       torch.cuda.empty_cache()





      def  __UpperCamelCase(				self     ):
       '''simple docstring'''




       UpperCamelCase			:      int 			=     load_numpy(
           "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"
           "/vq_diffusion/teddy_bear_pool_classifier_free_sampling.npy"     )

       UpperCamelCase			:      List[str] 			=     VQDiffusionPipeline.from_pretrained("microsoft/vq-diffusion-ithq"     )
       UpperCamelCase			:      List[str] 			=     pipeline.to(A_     )
       pipeline.set_progress_bar_config(disable=A_     )

       # requires GPU generator for gumbel softmax
       # don't use GPU generator in tests though
       UpperCamelCase			:      int 			=     torch.Generator(device=A_     ).manual_seed(0     )
       UpperCamelCase			:      List[Any] 			=     pipeline(
           "teddy bear playing in the pool" ,							num_images_per_prompt=1 ,							generator=A_ ,							output_type="np" ,							)

       UpperCamelCase			:      List[Any] 			=     output.images[0]

       assert image.shape == (256, 256, 3)
       assert np.abs(expected_image - image     ).max() < 2.0







from math import loga



def 	A_       (  _lowerCAmelCase  )		->				int:
 if a < 0:
  raise ValueError("Input value must be a positive integer"  )
 elif isinstance(_lowerCAmelCase							,     _lowerCAmelCase  ):
  raise TypeError("Input value must be a 'int' type"  )
 return 0 if (a == 0) else int(loga(a & -a  )  )


if __name__ == "__main__":
    import doctest

    doctest.testmod()
