

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   __lowercase      (unittest.TestCase ):
 def 	UpperCamelCase__  (							self					)					->str:



       '''simple docstring'''





       super().tearDown()
       gc.collect()
       torch.cuda.empty_cache()
 @property
 def 	UpperCamelCase__  (							self					)					->Dict:



       '''simple docstring'''





       __lowerCAmelCase					:						int					=					1
       __lowerCAmelCase					:						Optional[int]					=					3
       __lowerCAmelCase					:						Dict					=					(32, 32)

       __lowerCAmelCase					:						Optional[int]					=					floats_tensor((batch_size, num_channels) + sizes      ,							rng=random.Random(0					)					).to(A_					)
       return image
 @property
 def 	UpperCamelCase__  (							self					)					->Tuple:



       '''simple docstring'''





       torch.manual_seed(0					)
       __lowerCAmelCase					:						int					=					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=A_      ,							only_cross_attention=(True, True, False)      ,							num_class_embeds=100      ,							)
       return model
 @property
 def 	UpperCamelCase__  (							self					)					->Dict:



       '''simple docstring'''





       torch.manual_seed(0					)
       __lowerCAmelCase					:						int					=					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 	UpperCamelCase__  (							self					)					->List[Any]:



       '''simple docstring'''





       torch.manual_seed(0					)
       __lowerCAmelCase					:						str					=					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(A_					)
 def 	UpperCamelCase__  (							self					)					->Union[str, Any]:



       '''simple docstring'''





       __lowerCAmelCase					:						Any					=					'''cpu'''  # ensure determinism for the device-dependent torch.Generator
       __lowerCAmelCase					:						Tuple					=					self.dummy_cond_unet_upscale
       __lowerCAmelCase					:						List[Any]					=					DDPMScheduler()
       __lowerCAmelCase					:						Union[str, Any]					=					DDIMScheduler(prediction_type='''v_prediction'''					)
       __lowerCAmelCase					:						Optional[Any]					=					self.dummy_vae
       __lowerCAmelCase					:						Union[str, Any]					=					self.dummy_text_encoder
       __lowerCAmelCase					:						Any					=					CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''					)

       __lowerCAmelCase					:						Optional[int]					=					self.dummy_image.cpu().permute(0      ,							2      ,							3      ,							1					)[0]
       __lowerCAmelCase					:						Optional[Any]					=					Image.fromarray(np.uinta(A_					)					).convert('''RGB'''					).resize((64, 64)					)

       # make sure here that pndm scheduler skips prk
       __lowerCAmelCase					:						str					=					StableDiffusionUpscalePipeline(
           unet=A_      ,							low_res_scheduler=A_      ,							scheduler=A_      ,							vae=A_      ,							text_encoder=A_      ,							tokenizer=A_      ,							max_noise_level=350      ,							)
       __lowerCAmelCase					:						Optional[int]					=					sd_pipe.to(A_					)
       sd_pipe.set_progress_bar_config(disable=A_					)

       __lowerCAmelCase					:						Tuple					=					'''A painting of a squirrel eating a burger'''
       __lowerCAmelCase					:						int					=					torch.Generator(device=A_					).manual_seed(0					)
       __lowerCAmelCase					:						List[str]					=					sd_pipe(
           [prompt]      ,							image=A_      ,							generator=A_      ,							guidance_scale=6.0      ,							noise_level=20      ,							num_inference_steps=2      ,							output_type='''np'''      ,							)

       __lowerCAmelCase					:						Any					=					output.images

       __lowerCAmelCase					:						Any					=					torch.Generator(device=A_					).manual_seed(0					)
       __lowerCAmelCase					:						str					=					sd_pipe(
           [prompt]      ,							image=A_      ,							generator=A_      ,							guidance_scale=6.0      ,							noise_level=20      ,							num_inference_steps=2      ,							output_type='''np'''      ,							return_dict=A_      ,							)[0]

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

       __lowerCAmelCase					:						Optional[int]					=					low_res_image.size[0] * 4
       assert image.shape == (1, expected_height_width, expected_height_width, 3)
       __lowerCAmelCase					:						Any					=					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 	UpperCamelCase__  (							self					)					->List[Any]:



       '''simple docstring'''





       __lowerCAmelCase					:						Optional[Any]					=					'''cpu'''  # ensure determinism for the device-dependent torch.Generator
       __lowerCAmelCase					:						Union[str, Any]					=					self.dummy_cond_unet_upscale
       __lowerCAmelCase					:						List[str]					=					DDPMScheduler()
       __lowerCAmelCase					:						List[str]					=					DDIMScheduler(prediction_type='''v_prediction'''					)
       __lowerCAmelCase					:						int					=					self.dummy_vae
       __lowerCAmelCase					:						Any					=					self.dummy_text_encoder
       __lowerCAmelCase					:						Dict					=					CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''					)

       __lowerCAmelCase					:						Union[str, Any]					=					self.dummy_image.cpu().permute(0      ,							2      ,							3      ,							1					)[0]
       __lowerCAmelCase					:						Dict					=					Image.fromarray(np.uinta(A_					)					).convert('''RGB'''					).resize((64, 64)					)

       # make sure here that pndm scheduler skips prk
       __lowerCAmelCase					:						Optional[Any]					=					StableDiffusionUpscalePipeline(
           unet=A_      ,							low_res_scheduler=A_      ,							scheduler=A_      ,							vae=A_      ,							text_encoder=A_      ,							tokenizer=A_      ,							max_noise_level=350      ,							)
       __lowerCAmelCase					:						str					=					sd_pipe.to(A_					)
       sd_pipe.set_progress_bar_config(disable=A_					)

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

       __lowerCAmelCase					:						List[str]					=					torch.Generator(device=A_					).manual_seed(0					)
       __lowerCAmelCase					:						int					=					sd_pipe(
           [prompt]      ,							image=A_      ,							generator=A_      ,							num_images_per_prompt=2      ,							guidance_scale=6.0      ,							noise_level=20      ,							num_inference_steps=2      ,							output_type='''np'''      ,							)
       __lowerCAmelCase					:						Union[str, Any]					=					output.images
       assert image.shape[0] == 2





 @unittest.skipIf(torch_device != '''cuda'''      ,							'''This test requires a GPU'''					)
 def 	UpperCamelCase__  (							self					)					->List[Any]:



       '''simple docstring'''





       __lowerCAmelCase					:						Dict					=					self.dummy_cond_unet_upscale
       __lowerCAmelCase					:						Any					=					DDPMScheduler()
       __lowerCAmelCase					:						Union[str, Any]					=					DDIMScheduler(prediction_type='''v_prediction'''					)
       __lowerCAmelCase					:						List[Any]					=					self.dummy_vae
       __lowerCAmelCase					:						Dict					=					self.dummy_text_encoder
       __lowerCAmelCase					:						int					=					CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''					)

       __lowerCAmelCase					:						List[Any]					=					self.dummy_image.cpu().permute(0      ,							2      ,							3      ,							1					)[0]
       __lowerCAmelCase					:						int					=					Image.fromarray(np.uinta(A_					)					).convert('''RGB'''					).resize((64, 64)					)

       # put models in fp16, except vae as it overflows in fp16
       __lowerCAmelCase					:						Any					=					unet.half()
       __lowerCAmelCase					:						str					=					text_encoder.half()

       # make sure here that pndm scheduler skips prk
       __lowerCAmelCase					:						Dict					=					StableDiffusionUpscalePipeline(
           unet=A_      ,							low_res_scheduler=A_      ,							scheduler=A_      ,							vae=A_      ,							text_encoder=A_      ,							tokenizer=A_      ,							max_noise_level=350      ,							)
       __lowerCAmelCase					:						List[Any]					=					sd_pipe.to(A_					)
       sd_pipe.set_progress_bar_config(disable=A_					)

       __lowerCAmelCase					:						str					=					'''A painting of a squirrel eating a burger'''
       __lowerCAmelCase					:						Optional[Any]					=					torch.manual_seed(0					)
       __lowerCAmelCase					:						Dict					=					sd_pipe(
           [prompt]      ,							image=A_      ,							generator=A_      ,							num_inference_steps=2      ,							output_type='''np'''      ,							).images

       __lowerCAmelCase					:						List[Any]					=					low_res_image.size[0] * 4
       assert image.shape == (1, expected_height_width, expected_height_width, 3)







@slow
@require_torch_gpu
class   __lowercase      (unittest.TestCase ):
 def 	UpperCamelCase__  (							self					)					->Tuple:



       '''simple docstring'''





       super().tearDown()
       gc.collect()
       torch.cuda.empty_cache()
 def 	UpperCamelCase__  (							self					)					->Optional[Any]:



       '''simple docstring'''





       __lowerCAmelCase					:						Tuple					=					load_image(
           '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'''
           '''/sd2-upscale/low_res_cat.png'''					)
       __lowerCAmelCase					:						int					=					load_numpy(
           '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/sd2-upscale'''
           '''/upsampled_cat.npy'''					)

       __lowerCAmelCase					:						Dict					=					'''stabilityai/stable-diffusion-x4-upscaler'''
       __lowerCAmelCase					:						Dict					=					StableDiffusionUpscalePipeline.from_pretrained(A_					)
       pipe.to(A_					)
       pipe.set_progress_bar_config(disable=A_					)
       pipe.enable_attention_slicing()

       __lowerCAmelCase					:						Union[str, Any]					=					'''a cat sitting on a park bench'''

       __lowerCAmelCase					:						Optional[int]					=					torch.manual_seed(0					)
       __lowerCAmelCase					:						List[str]					=					pipe(
           prompt=A_      ,							image=A_      ,							generator=A_      ,							output_type='''np'''      ,							)
       __lowerCAmelCase					:						Optional[int]					=					output.images[0]

       assert image.shape == (512, 512, 3)
       assert np.abs(expected_image - image					).max() < 1e-3
 def 	UpperCamelCase__  (							self					)					->Dict:



       '''simple docstring'''





       __lowerCAmelCase					:						int					=					load_image(
           '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'''
           '''/sd2-upscale/low_res_cat.png'''					)
       __lowerCAmelCase					:						Optional[Any]					=					load_numpy(
           '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/sd2-upscale'''
           '''/upsampled_cat_fp16.npy'''					)

       __lowerCAmelCase					:						Optional[int]					=					'''stabilityai/stable-diffusion-x4-upscaler'''
       __lowerCAmelCase					:						int					=					StableDiffusionUpscalePipeline.from_pretrained(
           A_      ,							torch_dtype=torch.floataa      ,							)
       pipe.to(A_					)
       pipe.set_progress_bar_config(disable=A_					)
       pipe.enable_attention_slicing()

       __lowerCAmelCase					:						List[Any]					=					'''a cat sitting on a park bench'''

       __lowerCAmelCase					:						Optional[int]					=					torch.manual_seed(0					)
       __lowerCAmelCase					:						List[Any]					=					pipe(
           prompt=A_      ,							image=A_      ,							generator=A_      ,							output_type='''np'''      ,							)
       __lowerCAmelCase					:						List[str]					=					output.images[0]

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





 def 	UpperCamelCase__  (							self					)					->Union[str, Any]:



       '''simple docstring'''





       torch.cuda.empty_cache()
       torch.cuda.reset_max_memory_allocated()
       torch.cuda.reset_peak_memory_stats()

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

       __lowerCAmelCase					:						Optional[int]					=					'''stabilityai/stable-diffusion-x4-upscaler'''
       __lowerCAmelCase					:						List[str]					=					StableDiffusionUpscalePipeline.from_pretrained(
           A_      ,							torch_dtype=torch.floataa      ,							)
       pipe.to(A_					)
       pipe.set_progress_bar_config(disable=A_					)
       pipe.enable_attention_slicing(1					)
       pipe.enable_sequential_cpu_offload()

       __lowerCAmelCase					:						str					=					'''a cat sitting on a park bench'''

       __lowerCAmelCase					:						int					=					torch.manual_seed(0					)
       __lowerCAmelCase					:						int					=					pipe(
           prompt=A_      ,							image=A_      ,							generator=A_      ,							num_inference_steps=5      ,							output_type='''np'''      ,							)

       __lowerCAmelCase					:						Union[str, Any]					=					torch.cuda.max_memory_allocated()
       # make sure that less than 2.9 GB is allocated
       assert mem_bytes < 2.9 * 10**9





from collections import OrderedDict
from typing import Mapping

from packaging import version

from ...configuration_utils import PretrainedConfig
from ...onnx import OnnxConfig
from ...utils import logging
from ..auto import CONFIG_MAPPING


_UpperCamelCase 			=		logging.get_logger(__name__)

_UpperCamelCase 			=		{
    "microsoft/table-transformer-detection": (
        "https://huggingface.co/microsoft/table-transformer-detection/resolve/main/config.json"
    ),
}







class   __lowercase      (_UpperCAmelCase ):
 _UpperCamelCase			  =  """table-transformer"""
 _UpperCamelCase			  =  ["""past_key_values"""]
 _UpperCamelCase			  =  {
     """hidden_size""": """d_model""",
     """num_attention_heads""": """encoder_attention_heads""",
 }
 def __init__(							self      ,							A_=True      ,							A_=None      ,							A_=3      ,							A_=100      ,							A_=6      ,							A_=2048      ,							A_=8      ,							A_=6      ,							A_=2048      ,							A_=8      ,							A_=0.0      ,							A_=0.0      ,							A_=True      ,							A_="relu"      ,							A_=256      ,							A_=0.1      ,							A_=0.0      ,							A_=0.0      ,							A_=0.02      ,							A_=1.0      ,							A_=False      ,							A_="sine"      ,							A_="resnet50"      ,							A_=True      ,							A_=False      ,							A_=1      ,							A_=5      ,							A_=2      ,							A_=1      ,							A_=1      ,							A_=5      ,							A_=2      ,							A_=0.1      ,							**A_      ,							)					->Any:



       '''simple docstring'''





       if backbone_config is not None and use_timm_backbone:
             raise ValueError('''You can\'t specify both `backbone_config` and `use_timm_backbone`.'''					)

       if not use_timm_backbone:
             if backbone_config is None:
                   logger.info('''`backbone_config` is `None`. Initializing the config with the default `ResNet` backbone.'''					)
                   __lowerCAmelCase					:						Optional[Any]					=					CONFIG_MAPPING['''resnet'''](out_features=['''stage4''']					)
             elif isinstance(A_      ,							A_					):
                   __lowerCAmelCase					:						int					=					backbone_config.get('''model_type'''					)
                   __lowerCAmelCase					:						List[str]					=					CONFIG_MAPPING[backbone_model_type]
                   __lowerCAmelCase					:						Any					=					config_class.from_dict(A_					)
             # set timm attributes to None
             __lowerCAmelCase,							__lowerCAmelCase,							__lowerCAmelCase					:						List[str]					=					None, None, None

       __lowerCAmelCase					:						Tuple					=					use_timm_backbone
       __lowerCAmelCase					:						Optional[Any]					=					backbone_config
       __lowerCAmelCase					:						List[str]					=					num_channels
       __lowerCAmelCase					:						Tuple					=					num_queries
       __lowerCAmelCase					:						int					=					d_model
       __lowerCAmelCase					:						List[Any]					=					encoder_ffn_dim
       __lowerCAmelCase					:						Optional[int]					=					encoder_layers
       __lowerCAmelCase					:						List[str]					=					encoder_attention_heads
       __lowerCAmelCase					:						str					=					decoder_ffn_dim
       __lowerCAmelCase					:						Union[str, Any]					=					decoder_layers
       __lowerCAmelCase					:						Any					=					decoder_attention_heads
       __lowerCAmelCase					:						Optional[int]					=					dropout
       __lowerCAmelCase					:						Any					=					attention_dropout
       __lowerCAmelCase					:						Tuple					=					activation_dropout
       __lowerCAmelCase					:						Optional[Any]					=					activation_function
       __lowerCAmelCase					:						List[str]					=					init_std
       __lowerCAmelCase					:						Tuple					=					init_xavier_std
       __lowerCAmelCase					:						Any					=					encoder_layerdrop
       __lowerCAmelCase					:						List[Any]					=					decoder_layerdrop
       __lowerCAmelCase					:						Optional[Any]					=					encoder_layers
       __lowerCAmelCase					:						Optional[Any]					=					auxiliary_loss
       __lowerCAmelCase					:						Optional[Any]					=					position_embedding_type
       __lowerCAmelCase					:						Tuple					=					backbone
       __lowerCAmelCase					:						Any					=					use_pretrained_backbone
       __lowerCAmelCase					:						int					=					dilation
       # Hungarian matcher
       __lowerCAmelCase					:						Dict					=					class_cost
       __lowerCAmelCase					:						List[str]					=					bbox_cost
       __lowerCAmelCase					:						int					=					giou_cost
       # Loss coefficients
       __lowerCAmelCase					:						Optional[Any]					=					mask_loss_coefficient
       __lowerCAmelCase					:						Tuple					=					dice_loss_coefficient
       __lowerCAmelCase					:						int					=					bbox_loss_coefficient
       __lowerCAmelCase					:						List[Any]					=					giou_loss_coefficient
       __lowerCAmelCase					:						int					=					eos_coefficient
       super().__init__(is_encoder_decoder=A_      ,							**A_					)
 @property
 def 	UpperCamelCase__  (							self					)					->int:



       '''simple docstring'''





       return self.encoder_attention_heads





 @property
 def 	UpperCamelCase__  (							self					)					->int:



       '''simple docstring'''





       return self.d_model







class   __lowercase      (_UpperCAmelCase ):
 _UpperCamelCase			  =  version.parse("""1.11""" )
 @property
 def 	UpperCamelCase__  (							self					)					->Mapping[str, Mapping[int, str]]:



       '''simple docstring'''





       return OrderedDict(
           [
               ('''pixel_values''', {0: '''batch''', 1: '''num_channels''', 2: '''height''', 3: '''width'''}),
               ('''pixel_mask''', {0: '''batch'''}),
           ]					)
 @property
 def 	UpperCamelCase__  (							self					)					->float:



       '''simple docstring'''





       return 1e-5





 @property
 def 	UpperCamelCase__  (							self					)					->int:



       '''simple docstring'''





       return 12


