



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




 def        lowerCamelCase      (						self				):



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




 @property
 def        lowerCamelCase      (						self				):



    '''simple docstring'''
    __lowerCamelCase													=    1
    __lowerCamelCase													=    3
    __lowerCamelCase													=    (32, 32)

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




 @property
 def        lowerCamelCase      (						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=__UpperCAmelCase	,				only_cross_attention=(True, True, False)	,				num_class_embeds=100	,				)
    return model




 @property
 def        lowerCamelCase      (						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        lowerCamelCase      (						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-0_5	,				num_attention_heads=4	,				num_hidden_layers=5	,				pad_token_id=1	,				vocab_size=1000	,				hidden_act='''gelu'''	,				projection_dim=512	,				)
    return CLIPTextModel(__UpperCAmelCase				)




 def        lowerCamelCase      (						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(__UpperCAmelCase				)				).convert('''RGB'''				).resize((64, 64)				)

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

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

    __lowerCamelCase													=    output.images

    __lowerCamelCase													=    torch.Generator(device=__UpperCAmelCase				).manual_seed(0				)
    __lowerCamelCase													=    sd_pipe(
        [prompt]	,				image=__UpperCAmelCase	,				generator=__UpperCAmelCase	,				guidance_scale=6.0	,				noise_level=20	,				num_inference_steps=2	,				output_type='''np'''	,				return_dict=__UpperCAmelCase	,				)[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        lowerCamelCase      (						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(__UpperCAmelCase				)				).convert('''RGB'''				).resize((64, 64)				)

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

    __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=__UpperCAmelCase				).manual_seed(0				)
    __lowerCamelCase													=    sd_pipe(
        [prompt]	,				image=__UpperCAmelCase	,				generator=__UpperCAmelCase	,				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        lowerCamelCase      (						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(__UpperCAmelCase				)				).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=__UpperCAmelCase	,				low_res_scheduler=__UpperCAmelCase	,				scheduler=__UpperCAmelCase	,				vae=__UpperCAmelCase	,				text_encoder=__UpperCAmelCase	,				tokenizer=__UpperCAmelCase	,				max_noise_level=350	,				)
    __lowerCamelCase													=    sd_pipe.to(__UpperCAmelCase				)
    sd_pipe.set_progress_bar_config(disable=__UpperCAmelCase				)

    __lowerCamelCase													=    '''A painting of a squirrel eating a burger'''
    __lowerCamelCase													=    torch.manual_seed(0				)
    __lowerCamelCase													=    sd_pipe(
        [prompt]	,				image=__UpperCAmelCase	,				generator=__UpperCAmelCase	,				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  __lowerCAmelCase     (	unittest.TestCase   ):




 def        lowerCamelCase      (						self				):



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




 def        lowerCamelCase      (						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(__UpperCAmelCase				)
    pipe.to(__UpperCAmelCase				)
    pipe.set_progress_bar_config(disable=__UpperCAmelCase				)
    pipe.enable_attention_slicing()

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

    __lowerCamelCase													=    torch.manual_seed(0				)
    __lowerCamelCase													=    pipe(
        prompt=__UpperCAmelCase	,				image=__UpperCAmelCase	,				generator=__UpperCAmelCase	,				output_type='''np'''	,				)
    __lowerCamelCase													=    output.images[0]

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




 def        lowerCamelCase      (						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(
        __UpperCAmelCase	,				torch_dtype=torch.floataa	,				)
    pipe.to(__UpperCAmelCase				)
    pipe.set_progress_bar_config(disable=__UpperCAmelCase				)
    pipe.enable_attention_slicing()

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

    __lowerCamelCase													=    torch.manual_seed(0				)
    __lowerCamelCase													=    pipe(
        prompt=__UpperCAmelCase	,				image=__UpperCAmelCase	,				generator=__UpperCAmelCase	,				output_type='''np'''	,				)
    __lowerCamelCase													=    output.images[0]

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




 def        lowerCamelCase      (						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(
        __UpperCAmelCase	,				torch_dtype=torch.floataa	,				)
    pipe.to(__UpperCAmelCase				)
    pipe.set_progress_bar_config(disable=__UpperCAmelCase				)
    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=__UpperCAmelCase	,				image=__UpperCAmelCase	,				generator=__UpperCAmelCase	,				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











import unicodedata
from dataclasses import dataclass
from typing import Optional, Union

import numpy as np

from transformers.data.data_collator import DataCollatorMixin
from transformers.file_utils import PaddingStrategy
from transformers.tokenization_utils_base import PreTrainedTokenizerBase





def 		a__			(		_UpperCamelCase			: Tuple ,_UpperCamelCase			: Any ,_UpperCamelCase			: Union[str, Any] ,_UpperCamelCase			: Any							):
   if isinstance(_UpperCamelCase ,_UpperCamelCase							):
      __lowerCamelCase													=    np.full((len(_UpperCamelCase							), sequence_length, 2) ,_UpperCamelCase							)
   else:
      __lowerCamelCase													=    np.full((len(_UpperCamelCase							), sequence_length) ,_UpperCamelCase							)

   for i, tensor in enumerate(_UpperCamelCase							):
      if padding_side == "right":
         if isinstance(_UpperCamelCase ,_UpperCamelCase							):
            __lowerCamelCase													=    tensor[:sequence_length]
         else:
            __lowerCamelCase													=    tensor[:sequence_length]
      else:
         if isinstance(_UpperCamelCase ,_UpperCamelCase							):
            __lowerCamelCase													=    tensor[:sequence_length]
         else:
            __lowerCamelCase													=    tensor[:sequence_length]

   return out_tensor.tolist()





def 		a__			(		_UpperCamelCase			: Dict							):
   __lowerCamelCase													=    ord(_UpperCamelCase							)
   if (cp >= 33 and cp <= 47) or (cp >= 58 and cp <= 64) or (cp >= 91 and cp <= 96) or (cp >= 1_23 and cp <= 1_26):
      return True
   __lowerCamelCase													=    unicodedata.category(_UpperCamelCase							)
   if cat.startswith('''P'''							):
      return True
   return False





@dataclass
class  __lowerCAmelCase     (	lowerCAmelCase__   ):
 lowerCAmelCase__       					=  42
 lowerCAmelCase__       					=  True
 lowerCAmelCase__       					=  None
 lowerCAmelCase__       					=  None
 lowerCAmelCase__       					=  -1_0_0
 lowerCAmelCase__       					=  "pt"




 def        lowerCamelCase      (						self	,				__UpperCAmelCase				):



    '''simple docstring'''
    import torch

    __lowerCamelCase													=    '''label''' if '''label''' in features[0].keys() else '''labels'''
    __lowerCamelCase													=    [feature[label_name] for feature in features] if label_name in features[0].keys() else None
    __lowerCamelCase													=    self.tokenizer.pad(
        __UpperCAmelCase	,				padding=self.padding	,				max_length=self.max_length	,				pad_to_multiple_of=self.pad_to_multiple_of	,				return_tensors='''pt''' if labels is None else None	,				)

    if labels is None:
       return batch

    __lowerCamelCase													=    torch.tensor(batch['''entity_ids''']				).shape[1]
    __lowerCamelCase													=    self.tokenizer.padding_side
    if padding_side == "right":
       __lowerCamelCase													=    [
           list(__UpperCAmelCase				) + [self.label_pad_token_id] * (sequence_length - len(__UpperCAmelCase				)) for label in labels
       ]
    else:
       __lowerCamelCase													=    [
           [self.label_pad_token_id] * (sequence_length - len(__UpperCAmelCase				)) + list(__UpperCAmelCase				) for label in labels
       ]

    __lowerCamelCase													=    [feature['''ner_tags'''] for feature in features]
    __lowerCamelCase													=    padding_tensor(__UpperCAmelCase	,				-1	,				__UpperCAmelCase	,				__UpperCAmelCase				)
    __lowerCamelCase													=    [feature['''original_entity_spans'''] for feature in features]
    __lowerCamelCase													=    padding_tensor(__UpperCAmelCase	,				(-1, -1)	,				__UpperCAmelCase	,				__UpperCAmelCase				)
    __lowerCamelCase													=    {k: torch.tensor(__UpperCAmelCase	,				dtype=torch.intaa				) for k, v in batch.items()}

    return batch






