


'''simple docstring'''

import gc
import random
import unittest

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

from diffusers import AutoencoderKL, DDIMScheduler, DDPMScheduler, StableDiffusionUpscalePipeline, UNetaDConditionModel
from diffusers.utils import floats_tensor, load_image, load_numpy, slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu


enable_full_determinism()





class   snake_case					(					unittest.TestCase    ):





  """simple docstring"""




  def 						snake_case	(					self  ):







   """simple docstring"""
   super().tearDown()
   gc.collect()
   torch.cuda.empty_cache()






  @property
  def 						snake_case	(					self  ):







   """simple docstring"""
   lowerCamelCase_  				=						1
   lowerCamelCase_  				=						3
   lowerCamelCase_  				=						(32, 32)

   lowerCamelCase_  				=						floats_tensor((batch_size, num_channels) + sizes    ,						rng=random.Random(0  )  ).to(lowercase__  )
   return image






  @property
  def 						snake_case	(					self  ):







   """simple docstring"""
   torch.manual_seed(0  )
   lowerCamelCase_  				=						UNetaDConditionModel(
       block_out_channels=(32, 32, 64)    ,						layers_per_block=2    ,						sample_size=32    ,						in_channels=7    ,						out_channels=4    ,						down_block_types=("DownBlock2D", "CrossAttnDownBlock2D", "CrossAttnDownBlock2D")    ,						up_block_types=("CrossAttnUpBlock2D", "CrossAttnUpBlock2D", "UpBlock2D")    ,						cross_attention_dim=32    ,						attention_head_dim=8    ,						use_linear_projection=lowercase__    ,						only_cross_attention=(True, True, False)    ,						num_class_embeds=100    ,						)
   return model






  @property
  def 						snake_case	(					self  ):







   """simple docstring"""
   torch.manual_seed(0  )
   lowerCamelCase_  				=						AutoencoderKL(
       block_out_channels=[32, 32, 64]    ,						in_channels=3    ,						out_channels=3    ,						down_block_types=["DownEncoderBlock2D", "DownEncoderBlock2D", "DownEncoderBlock2D"]    ,						up_block_types=["UpDecoderBlock2D", "UpDecoderBlock2D", "UpDecoderBlock2D"]    ,						latent_channels=4    ,						)
   return model






  @property
  def 						snake_case	(					self  ):







   """simple docstring"""
   torch.manual_seed(0  )
   lowerCamelCase_  				=						CLIPTextConfig(
       bos_token_id=0    ,						eos_token_id=2    ,						hidden_size=32    ,						intermediate_size=37    ,						layer_norm_eps=1e-05    ,						num_attention_heads=4    ,						num_hidden_layers=5    ,						pad_token_id=1    ,						vocab_size=1000    ,						hidden_act="gelu"    ,						projection_dim=512    ,						)
   return CLIPTextModel(lowercase__  )






  def 						snake_case	(					self  ):







   """simple docstring"""
   lowerCamelCase_  				=						'''cpu'''  # ensure determinism for the device-dependent torch.Generator
   lowerCamelCase_  				=						self.dummy_cond_unet_upscale
   lowerCamelCase_  				=						DDPMScheduler()
   lowerCamelCase_  				=						DDIMScheduler(prediction_type="v_prediction"  )
   lowerCamelCase_  				=						self.dummy_vae
   lowerCamelCase_  				=						self.dummy_text_encoder
   lowerCamelCase_  				=						CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip"  )

   lowerCamelCase_  				=						self.dummy_image.cpu().permute(0    ,						2    ,						3    ,						1  )[0]
   lowerCamelCase_  				=						Image.fromarray(np.uinta(lowercase__  )  ).convert("RGB"  ).resize((64, 64)  )

   # make sure here that pndm scheduler skips prk
   lowerCamelCase_  				=						StableDiffusionUpscalePipeline(
       unet=lowercase__    ,						low_res_scheduler=lowercase__    ,						scheduler=lowercase__    ,						vae=lowercase__    ,						text_encoder=lowercase__    ,						tokenizer=lowercase__    ,						max_noise_level=350    ,						)
   lowerCamelCase_  				=						sd_pipe.to(lowercase__  )
   sd_pipe.set_progress_bar_config(disable=lowercase__  )

   lowerCamelCase_  				=						'''A painting of a squirrel eating a burger'''
   lowerCamelCase_  				=						torch.Generator(device=lowercase__  ).manual_seed(0  )
   lowerCamelCase_  				=						sd_pipe(
       [prompt]    ,						image=lowercase__    ,						generator=lowercase__    ,						guidance_scale=6.0    ,						noise_level=20    ,						num_inference_steps=2    ,						output_type="np"    ,						)

   lowerCamelCase_  				=						output.images

   lowerCamelCase_  				=						torch.Generator(device=lowercase__  ).manual_seed(0  )
   lowerCamelCase_  				=						sd_pipe(
       [prompt]    ,						image=lowercase__    ,						generator=lowercase__    ,						guidance_scale=6.0    ,						noise_level=20    ,						num_inference_steps=2    ,						output_type="np"    ,						return_dict=lowercase__    ,						)[0]

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

   lowerCamelCase_  				=						low_res_image.size[0] * 4
   assert image.shape == (1, expected_height_width, expected_height_width, 3)
   lowerCamelCase_  				=						np.array([0.3_113, 0.3_910, 0.4_272, 0.4_859, 0.5_061, 0.4_652, 0.5_362, 0.5_715, 0.5_661]  )

   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 						snake_case	(					self  ):







   """simple docstring"""
   lowerCamelCase_  				=						'''cpu'''  # ensure determinism for the device-dependent torch.Generator
   lowerCamelCase_  				=						self.dummy_cond_unet_upscale
   lowerCamelCase_  				=						DDPMScheduler()
   lowerCamelCase_  				=						DDIMScheduler(prediction_type="v_prediction"  )
   lowerCamelCase_  				=						self.dummy_vae
   lowerCamelCase_  				=						self.dummy_text_encoder
   lowerCamelCase_  				=						CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip"  )

   lowerCamelCase_  				=						self.dummy_image.cpu().permute(0    ,						2    ,						3    ,						1  )[0]
   lowerCamelCase_  				=						Image.fromarray(np.uinta(lowercase__  )  ).convert("RGB"  ).resize((64, 64)  )

   # make sure here that pndm scheduler skips prk
   lowerCamelCase_  				=						StableDiffusionUpscalePipeline(
       unet=lowercase__    ,						low_res_scheduler=lowercase__    ,						scheduler=lowercase__    ,						vae=lowercase__    ,						text_encoder=lowercase__    ,						tokenizer=lowercase__    ,						max_noise_level=350    ,						)
   lowerCamelCase_  				=						sd_pipe.to(lowercase__  )
   sd_pipe.set_progress_bar_config(disable=lowercase__  )

   lowerCamelCase_  				=						'''A painting of a squirrel eating a burger'''
   lowerCamelCase_  				=						sd_pipe(
       2 * [prompt]    ,						image=2 * [low_res_image]    ,						guidance_scale=6.0    ,						noise_level=20    ,						num_inference_steps=2    ,						output_type="np"    ,						)
   lowerCamelCase_  				=						output.images
   assert image.shape[0] == 2

   lowerCamelCase_  				=						torch.Generator(device=lowercase__  ).manual_seed(0  )
   lowerCamelCase_  				=						sd_pipe(
       [prompt]    ,						image=lowercase__    ,						generator=lowercase__    ,						num_images_per_prompt=2    ,						guidance_scale=6.0    ,						noise_level=20    ,						num_inference_steps=2    ,						output_type="np"    ,						)
   lowerCamelCase_  				=						output.images
   assert image.shape[0] == 2







  @unittest.skipIf(torch_device != "cuda"    ,						"This test requires a GPU"  )
  def 						snake_case	(					self  ):







   """simple docstring"""
   lowerCamelCase_  				=						self.dummy_cond_unet_upscale
   lowerCamelCase_  				=						DDPMScheduler()
   lowerCamelCase_  				=						DDIMScheduler(prediction_type="v_prediction"  )
   lowerCamelCase_  				=						self.dummy_vae
   lowerCamelCase_  				=						self.dummy_text_encoder
   lowerCamelCase_  				=						CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip"  )

   lowerCamelCase_  				=						self.dummy_image.cpu().permute(0    ,						2    ,						3    ,						1  )[0]
   lowerCamelCase_  				=						Image.fromarray(np.uinta(lowercase__  )  ).convert("RGB"  ).resize((64, 64)  )

   # put models in fp16, except vae as it overflows in fp16
   lowerCamelCase_  				=						unet.half()
   lowerCamelCase_  				=						text_encoder.half()

   # make sure here that pndm scheduler skips prk
   lowerCamelCase_  				=						StableDiffusionUpscalePipeline(
       unet=lowercase__    ,						low_res_scheduler=lowercase__    ,						scheduler=lowercase__    ,						vae=lowercase__    ,						text_encoder=lowercase__    ,						tokenizer=lowercase__    ,						max_noise_level=350    ,						)
   lowerCamelCase_  				=						sd_pipe.to(lowercase__  )
   sd_pipe.set_progress_bar_config(disable=lowercase__  )

   lowerCamelCase_  				=						'''A painting of a squirrel eating a burger'''
   lowerCamelCase_  				=						torch.manual_seed(0  )
   lowerCamelCase_  				=						sd_pipe(
       [prompt]    ,						image=lowercase__    ,						generator=lowercase__    ,						num_inference_steps=2    ,						output_type="np"    ,						).images

   lowerCamelCase_  				=						low_res_image.size[0] * 4
   assert image.shape == (1, expected_height_width, expected_height_width, 3)





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





  """simple docstring"""




  def 						snake_case	(					self  ):







   """simple docstring"""
   super().tearDown()
   gc.collect()
   torch.cuda.empty_cache()






  def 						snake_case	(					self  ):







   """simple docstring"""
   lowerCamelCase_  				=						load_image(
       "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"
       "/sd2-upscale/low_res_cat.png"  )
   lowerCamelCase_  				=						load_numpy(
       "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/sd2-upscale"
       "/upsampled_cat.npy"  )

   lowerCamelCase_  				=						'''stabilityai/stable-diffusion-x4-upscaler'''
   lowerCamelCase_  				=						StableDiffusionUpscalePipeline.from_pretrained(lowercase__  )
   pipe.to(lowercase__  )
   pipe.set_progress_bar_config(disable=lowercase__  )
   pipe.enable_attention_slicing()

   lowerCamelCase_  				=						'''a cat sitting on a park bench'''

   lowerCamelCase_  				=						torch.manual_seed(0  )
   lowerCamelCase_  				=						pipe(
       prompt=lowercase__    ,						image=lowercase__    ,						generator=lowercase__    ,						output_type="np"    ,						)
   lowerCamelCase_  				=						output.images[0]

   assert image.shape == (512, 512, 3)
   assert np.abs(expected_image - image  ).max() < 1e-3






  def 						snake_case	(					self  ):







   """simple docstring"""
   lowerCamelCase_  				=						load_image(
       "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"
       "/sd2-upscale/low_res_cat.png"  )
   lowerCamelCase_  				=						load_numpy(
       "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/sd2-upscale"
       "/upsampled_cat_fp16.npy"  )

   lowerCamelCase_  				=						'''stabilityai/stable-diffusion-x4-upscaler'''
   lowerCamelCase_  				=						StableDiffusionUpscalePipeline.from_pretrained(
       lowercase__    ,						torch_dtype=torch.floataa    ,						)
   pipe.to(lowercase__  )
   pipe.set_progress_bar_config(disable=lowercase__  )
   pipe.enable_attention_slicing()

   lowerCamelCase_  				=						'''a cat sitting on a park bench'''

   lowerCamelCase_  				=						torch.manual_seed(0  )
   lowerCamelCase_  				=						pipe(
       prompt=lowercase__    ,						image=lowercase__    ,						generator=lowercase__    ,						output_type="np"    ,						)
   lowerCamelCase_  				=						output.images[0]

   assert image.shape == (512, 512, 3)
   assert np.abs(expected_image - image  ).max() < 5e-1







  def 						snake_case	(					self  ):







   """simple docstring"""
   torch.cuda.empty_cache()
   torch.cuda.reset_max_memory_allocated()
   torch.cuda.reset_peak_memory_stats()

   lowerCamelCase_  				=						load_image(
       "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"
       "/sd2-upscale/low_res_cat.png"  )

   lowerCamelCase_  				=						'''stabilityai/stable-diffusion-x4-upscaler'''
   lowerCamelCase_  				=						StableDiffusionUpscalePipeline.from_pretrained(
       lowercase__    ,						torch_dtype=torch.floataa    ,						)
   pipe.to(lowercase__  )
   pipe.set_progress_bar_config(disable=lowercase__  )
   pipe.enable_attention_slicing(1  )
   pipe.enable_sequential_cpu_offload()

   lowerCamelCase_  				=						'''a cat sitting on a park bench'''

   lowerCamelCase_  				=						torch.manual_seed(0  )
   lowerCamelCase_  				=						pipe(
       prompt=lowercase__    ,						image=lowercase__    ,						generator=lowercase__    ,						num_inference_steps=5    ,						output_type="np"    ,						)

   lowerCamelCase_  				=						torch.cuda.max_memory_allocated()
   # make sure that less than 2.9 GB is allocated
   assert mem_bytes < 2.9 * 10**9







'''simple docstring'''

import unittest

import numpy as np

from transformers import DistilBertConfig, is_flax_available
from transformers.testing_utils import require_flax, slow

from ...test_modeling_flax_common import FlaxModelTesterMixin, ids_tensor, random_attention_mask


if is_flax_available():
     import jax.numpy as jnp

     from transformers.models.distilbert.modeling_flax_distilbert import (
         FlaxDistilBertForMaskedLM,
         FlaxDistilBertForMultipleChoice,
         FlaxDistilBertForQuestionAnswering,
         FlaxDistilBertForSequenceClassification,
         FlaxDistilBertForTokenClassification,
         FlaxDistilBertModel,
     )





class   snake_case					(					unittest.TestCase    ):





  """simple docstring"""




  def __init__(					self    ,						UpperCamelCase    ,						UpperCamelCase=13    ,						UpperCamelCase=7    ,						UpperCamelCase=True    ,						UpperCamelCase=True    ,						UpperCamelCase=True    ,						UpperCamelCase=True    ,						UpperCamelCase=99    ,						UpperCamelCase=32    ,						UpperCamelCase=5    ,						UpperCamelCase=4    ,						UpperCamelCase=37    ,						UpperCamelCase="gelu"    ,						UpperCamelCase=0.1    ,						UpperCamelCase=0.1    ,						UpperCamelCase=512    ,						UpperCamelCase=16    ,						UpperCamelCase=2    ,						UpperCamelCase=0.02    ,						UpperCamelCase=4    ,						):







   """simple docstring"""
   lowerCamelCase_  				=						parent
   lowerCamelCase_  				=						batch_size
   lowerCamelCase_  				=						seq_length
   lowerCamelCase_  				=						is_training
   lowerCamelCase_  				=						use_attention_mask
   lowerCamelCase_  				=						use_token_type_ids
   lowerCamelCase_  				=						use_labels
   lowerCamelCase_  				=						vocab_size
   lowerCamelCase_  				=						hidden_size
   lowerCamelCase_  				=						num_hidden_layers
   lowerCamelCase_  				=						num_attention_heads
   lowerCamelCase_  				=						intermediate_size
   lowerCamelCase_  				=						hidden_act
   lowerCamelCase_  				=						hidden_dropout_prob
   lowerCamelCase_  				=						attention_probs_dropout_prob
   lowerCamelCase_  				=						max_position_embeddings
   lowerCamelCase_  				=						type_vocab_size
   lowerCamelCase_  				=						type_sequence_label_size
   lowerCamelCase_  				=						initializer_range
   lowerCamelCase_  				=						num_choices






  def 						snake_case	(					self  ):







   """simple docstring"""
   lowerCamelCase_  				=						ids_tensor([self.batch_size, self.seq_length]    ,						self.vocab_size  )

   lowerCamelCase_  				=						None
   if self.use_attention_mask:
    lowerCamelCase_  				=						random_attention_mask([self.batch_size, self.seq_length]  )

   lowerCamelCase_  				=						DistilBertConfig(
       vocab_size=self.vocab_size    ,						dim=self.hidden_size    ,						n_layers=self.num_hidden_layers    ,						n_heads=self.num_attention_heads    ,						hidden_dim=self.intermediate_size    ,						hidden_act=self.hidden_act    ,						dropout=self.hidden_dropout_prob    ,						attention_dropout=self.attention_probs_dropout_prob    ,						max_position_embeddings=self.max_position_embeddings    ,						initializer_range=self.initializer_range    ,						tie_weights_=UpperCamelCase    ,						)

   return config, input_ids, attention_mask







  def 						snake_case	(					self  ):







   """simple docstring"""
   lowerCamelCase_  				=						self.prepare_config_and_inputs()
   lowerCamelCase_				,lowerCamelCase_				,lowerCamelCase_  				=						config_and_inputs
   lowerCamelCase_  				=						{"input_ids": input_ids, "attention_mask": attention_mask}
   return config, inputs_dict





@require_flax
class   snake_case					(					lowercase    ,		unittest.TestCase    ):





  """simple docstring"""




  _lowerCamelCase						=						(
      (
          FlaxDistilBertModel,
          FlaxDistilBertForMaskedLM,
          FlaxDistilBertForMultipleChoice,
          FlaxDistilBertForQuestionAnswering,
          FlaxDistilBertForSequenceClassification,
          FlaxDistilBertForTokenClassification,
          FlaxDistilBertForQuestionAnswering,
      )
      if is_flax_available()
      else ()
  )






  def 						snake_case	(					self  ):







   """simple docstring"""
   lowerCamelCase_  				=						FlaxDistilBertModelTester(self  )







  @slow
  def 						snake_case	(					self  ):







   """simple docstring"""
   for model_class_name in self.all_model_classes:
    lowerCamelCase_  				=						model_class_name.from_pretrained("distilbert-base-uncased"  )
    lowerCamelCase_  				=						model(np.ones((1, 1)  )  )
    self.assertIsNotNone(UpperCamelCase  )





@require_flax
class   snake_case					(					unittest.TestCase    ):





  """simple docstring"""




  @slow
  def 						snake_case	(					self  ):







   """simple docstring"""
   lowerCamelCase_  				=						FlaxDistilBertModel.from_pretrained("distilbert-base-uncased"  )
   lowerCamelCase_  				=						np.array([[0, 345, 232, 328, 740, 140, 1695, 69, 6078, 1588, 2]]  )
   lowerCamelCase_  				=						np.array([[0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]]  )
   lowerCamelCase_  				=						model(UpperCamelCase    ,						attention_mask=UpperCamelCase  )[0]
   lowerCamelCase_  				=						(1, 11, 768)
   self.assertEqual(output.shape    ,						UpperCamelCase  )
   lowerCamelCase_  				=						np.array([[[-0.1_639, 0.3_299, 0.1_648], [-0.1_746, 0.3_289, 0.1_710], [-0.1_884, 0.3_357, 0.1_810]]]  )

   self.assertTrue(jnp.allclose(output[:, 1:4, 1:4]    ,						UpperCamelCase    ,						atol=1e-4  )  )



