

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      :					Optional[Any]					    =      False





class        lowerCAmelCase__  (					unittest.TestCase			):



       '''simple docstring'''







       def     __UpperCamelCase     (							self     ):







             '''simple docstring'''



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

       @property
       def     __UpperCamelCase     (							self     ):







             '''simple docstring'''



             return 1_2

       @property
       def     __UpperCamelCase     (							self     ):







             '''simple docstring'''



             return 1_2

       @property
       def     __UpperCamelCase     (							self     ):







             '''simple docstring'''



             return 3_2

       @property
       def     __UpperCamelCase     (							self     ):







             '''simple docstring'''



             torch.manual_seed(0     )
             __A    =VQModel(
                 block_out_channels=[3_2, 6_4] ,   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'''



             __A    =CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''     )
             return tokenizer

       @property
       def     __UpperCamelCase     (							self     ):







             '''simple docstring'''



             torch.manual_seed(0     )
             __A    =CLIPTextConfig(
                 bos_token_id=0 ,   eos_token_id=2 ,   hidden_size=self.text_embedder_hidden_size ,   intermediate_size=3_7 ,   layer_norm_eps=1E-05 ,   num_attention_heads=4 ,   num_hidden_layers=5 ,   pad_token_id=1 ,   vocab_size=1_0_0_0 ,   )
             return CLIPTextModel(lowercase__     )

       @property
       def     __UpperCamelCase     (							self     ):







             '''simple docstring'''



             torch.manual_seed(0     )

             __A    =1_2
             __A    =1_2

             __A    ={
                 '''attention_bias''': True,
                 '''cross_attention_dim''': 3_2,
                 '''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''': 3_2,
                 '''sample_size''': width,
                 '''activation_fn''': '''geglu-approximate''',
             }

             __A    =TransformeraDModel(**lowercase__     )
             return model

       def     __UpperCamelCase     (							self     ):







             '''simple docstring'''



             __A    ='''cpu'''

             __A    =self.dummy_vqvae
             __A    =self.dummy_text_encoder
             __A    =self.dummy_tokenizer
             __A    =self.dummy_transformer
             __A    =VQDiffusionScheduler(self.num_embed     )
             __A    =LearnedClassifierFreeSamplingEmbeddings(learnable=lowercase__     )

             __A    =VQDiffusionPipeline(
                 vqvae=lowercase__ ,   text_encoder=lowercase__ ,   tokenizer=lowercase__ ,   transformer=lowercase__ ,   scheduler=lowercase__ ,   learned_classifier_free_sampling_embeddings=lowercase__ ,   )
             __A    =pipe.to(lowercase__     )
             pipe.set_progress_bar_config(disable=lowercase__     )

             __A    ='''teddy bear playing in the pool'''

             __A    =torch.Generator(device=lowercase__     ).manual_seed(0     )
             __A    =pipe([prompt] ,   generator=lowercase__ ,   num_inference_steps=2 ,   output_type='''np'''     )
             __A    =output.images

             __A    =torch.Generator(device=lowercase__     ).manual_seed(0     )
             __A    =pipe(
                 [prompt] ,   generator=lowercase__ ,   output_type='''np''' ,   return_dict=lowercase__ ,   num_inference_steps=2     )[0]

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

             assert image.shape == (1, 2_4, 2_4, 3)

             __A    =np.array([0.6551, 0.6168, 0.5008, 0.5676, 0.5659, 0.4295, 0.6073, 0.5599, 0.4992]     )

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



             __A    ='''cpu'''

             __A    =self.dummy_vqvae
             __A    =self.dummy_text_encoder
             __A    =self.dummy_tokenizer
             __A    =self.dummy_transformer
             __A    =VQDiffusionScheduler(self.num_embed     )
             __A    =LearnedClassifierFreeSamplingEmbeddings(
                 learnable=lowercase__ ,   hidden_size=self.text_embedder_hidden_size ,   length=tokenizer.model_max_length     )

             __A    =VQDiffusionPipeline(
                 vqvae=lowercase__ ,   text_encoder=lowercase__ ,   tokenizer=lowercase__ ,   transformer=lowercase__ ,   scheduler=lowercase__ ,   learned_classifier_free_sampling_embeddings=lowercase__ ,   )
             __A    =pipe.to(lowercase__     )
             pipe.set_progress_bar_config(disable=lowercase__     )

             __A    ='''teddy bear playing in the pool'''

             __A    =torch.Generator(device=lowercase__     ).manual_seed(0     )
             __A    =pipe([prompt] ,   generator=lowercase__ ,   num_inference_steps=2 ,   output_type='''np'''     )
             __A    =output.images

             __A    =torch.Generator(device=lowercase__     ).manual_seed(0     )
             __A    =pipe(
                 [prompt] ,   generator=lowercase__ ,   output_type='''np''' ,   return_dict=lowercase__ ,   num_inference_steps=2     )[0]

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

             assert image.shape == (1, 2_4, 2_4, 3)

             __A    =np.array([0.6693, 0.6075, 0.4959, 0.5701, 0.5583, 0.4333, 0.6171, 0.5684, 0.4988]     )

             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        lowerCAmelCase__  (					unittest.TestCase			):



       '''simple docstring'''







       def     __UpperCamelCase     (							self     ):







             '''simple docstring'''



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


       def     __UpperCamelCase     (							self     ):







             '''simple docstring'''



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

             __A    =VQDiffusionPipeline.from_pretrained('''microsoft/vq-diffusion-ithq'''     )
             __A    =pipeline.to(lowercase__     )
             pipeline.set_progress_bar_config(disable=lowercase__     )

             # requires GPU generator for gumbel softmax
             # don't use GPU generator in tests though
             __A    =torch.Generator(device=lowercase__     ).manual_seed(0     )
             __A    =pipeline(
                 '''teddy bear playing in the pool''' ,   num_images_per_prompt=1 ,   generator=lowercase__ ,   output_type='''np''' ,   )

             __A    =output.images[0]

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





from ...utils import (
    OptionalDependencyNotAvailable,
    is_torch_available,
    is_transformers_available,
    is_transformers_version,
)


try:
   if not (is_transformers_available() and is_torch_available() and is_transformers_version('''>=''', '''4.25.0''')):
      raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
   from ...utils.dummy_torch_and_transformers_objects import (
       VersatileDiffusionDualGuidedPipeline,
       VersatileDiffusionImageVariationPipeline,
       VersatileDiffusionPipeline,
       VersatileDiffusionTextToImagePipeline,
   )
else:
   from .modeling_text_unet import UNetFlatConditionModel
   from .pipeline_versatile_diffusion import VersatileDiffusionPipeline
   from .pipeline_versatile_diffusion_dual_guided import VersatileDiffusionDualGuidedPipeline
   from .pipeline_versatile_diffusion_image_variation import VersatileDiffusionImageVariationPipeline
   from .pipeline_versatile_diffusion_text_to_image import VersatileDiffusionTextToImagePipeline


