



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


_snake_case		=  False



class   UpperCAmelCase_    (       unittest.TestCase):







							def    snake_case__      (   self):

														'''simple docstring'''

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







							@property
							def    snake_case__      (   self):

														'''simple docstring'''

														return 12







							@property
							def    snake_case__      (   self):

														'''simple docstring'''

														return 12







							@property
							def    snake_case__      (   self):

														'''simple docstring'''

														return 32







							@property
							def    snake_case__      (   self):

														'''simple docstring'''

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

														'''simple docstring'''

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







							@property
							def    snake_case__      (   self):

														'''simple docstring'''

														torch.manual_seed(0)
														_lowerCAmelCase     :							Union[str, Any]    =    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    snake_case__      (   self):

														'''simple docstring'''

														torch.manual_seed(0)

														_lowerCAmelCase     :							Any    =    12
														_lowerCAmelCase     :							Dict    =    12

														_lowerCAmelCase     :							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",
														}

														_lowerCAmelCase     :							int    =    TransformeraDModel(**__a)
														return model







							def    snake_case__      (   self):

														'''simple docstring'''

														_lowerCAmelCase     :							str    =    "cpu"

														_lowerCAmelCase     :							Optional[int]    =    self.dummy_vqvae
														_lowerCAmelCase     :							Dict    =    self.dummy_text_encoder
														_lowerCAmelCase     :							Dict    =    self.dummy_tokenizer
														_lowerCAmelCase     :							Tuple    =    self.dummy_transformer
														_lowerCAmelCase     :							Dict    =    VQDiffusionScheduler(self.num_embed)
														_lowerCAmelCase     :							Tuple    =    LearnedClassifierFreeSamplingEmbeddings(learnable=__a)

														_lowerCAmelCase     :							int    =    VQDiffusionPipeline(
														    vqvae=__a,			text_encoder=__a,			tokenizer=__a,			transformer=__a,			scheduler=__a,			learned_classifier_free_sampling_embeddings=__a,			)
														_lowerCAmelCase     :							int    =    pipe.to(__a)
														pipe.set_progress_bar_config(disable=__a)

														_lowerCAmelCase     :							Dict    =    "teddy bear playing in the pool"

														_lowerCAmelCase     :							Tuple    =    torch.Generator(device=__a).manual_seed(0)
														_lowerCAmelCase     :							Union[str, Any]    =    pipe([prompt],			generator=__a,			num_inference_steps=2,			output_type="np")
														_lowerCAmelCase     :							Union[str, Any]    =    output.images

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

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

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

														_lowerCAmelCase     :							Optional[Any]    =    np.array([0.6_551, 0.6_168, 0.5_008, 0.5_676, 0.5_659, 0.4_295, 0.6_073, 0.5_599, 0.4_992])

														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     :							Any    =    "cpu"

														_lowerCAmelCase     :							Union[str, Any]    =    self.dummy_vqvae
														_lowerCAmelCase     :							List[Any]    =    self.dummy_text_encoder
														_lowerCAmelCase     :							Optional[int]    =    self.dummy_tokenizer
														_lowerCAmelCase     :							Optional[int]    =    self.dummy_transformer
														_lowerCAmelCase     :							List[Any]    =    VQDiffusionScheduler(self.num_embed)
														_lowerCAmelCase     :							int    =    LearnedClassifierFreeSamplingEmbeddings(
														    learnable=__a,			hidden_size=self.text_embedder_hidden_size,			length=tokenizer.model_max_length)

														_lowerCAmelCase     :							str    =    VQDiffusionPipeline(
														    vqvae=__a,			text_encoder=__a,			tokenizer=__a,			transformer=__a,			scheduler=__a,			learned_classifier_free_sampling_embeddings=__a,			)
														_lowerCAmelCase     :							Any    =    pipe.to(__a)
														pipe.set_progress_bar_config(disable=__a)

														_lowerCAmelCase     :							int    =    "teddy bear playing in the pool"

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

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

														_lowerCAmelCase     :							Dict    =    image[0, -3:, -3:, -1]
														_lowerCAmelCase     :							Tuple    =    image_from_tuple[0, -3:, -3:, -1]

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

														_lowerCAmelCase     :							Union[str, Any]    =    np.array([0.6_693, 0.6_075, 0.4_959, 0.5_701, 0.5_583, 0.4_333, 0.6_171, 0.5_684, 0.4_988])

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







							def    snake_case__      (   self):

														'''simple docstring'''

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







							def    snake_case__      (   self):

														'''simple docstring'''

														_lowerCAmelCase     :							List[Any]    =    load_numpy(
														    "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"
														    "/vq_diffusion/teddy_bear_pool_classifier_free_sampling.npy")

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

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

														_lowerCAmelCase     :							str    =    output.images[0]

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





import math
from typing import Dict, Iterable, List, Optional, Tuple, Union

import numpy as np

from ...image_processing_utils import BaseImageProcessor, BatchFeature, get_size_dict
from ...image_transforms import normalize, rescale, resize, to_channel_dimension_format
from ...image_utils import (
    IMAGENET_STANDARD_MEAN,
    IMAGENET_STANDARD_STD,
    ChannelDimension,
    ImageInput,
    PILImageResampling,
    get_image_size,
    is_torch_available,
    is_torch_tensor,
    make_list_of_images,
    to_numpy_array,
    valid_images,
)
from ...utils import TensorType, is_vision_available, logging


if is_torch_available():
						import torch

if is_vision_available():
						import PIL


_snake_case		=  logging.get_logger(__name__)






def 				A				(					_lowerCamelCase				, _lowerCamelCase				, _lowerCamelCase				, _lowerCamelCase  ):





							'''simple docstring'''




							def constraint_to_multiple_of(_lowerCamelCase				, _lowerCamelCase				, _lowerCamelCase=0				, _lowerCamelCase=None  ):
														_lowerCAmelCase     :							Tuple    =    round(val / multiple  ) * multiple

														if max_val is not None and x > max_val:
																					_lowerCAmelCase     :							Optional[int]    =    math.floor(val / multiple  ) * multiple

														if x < min_val:
																					_lowerCAmelCase     :							List[str]    =    math.ceil(val / multiple  ) * multiple

														return x

							_lowerCAmelCase     :							Union[str, Any]    =    (output_size, output_size) if isinstance(_lowerCamelCase				, _lowerCamelCase  ) else output_size

							_lowerCAmelCase   ,    _lowerCAmelCase     :							Optional[Any]    =    get_image_size(_lowerCamelCase  )
							_lowerCAmelCase   ,    _lowerCAmelCase     :							Any    =    output_size

							# determine new height and width
							_lowerCAmelCase     :							List[Any]    =    output_height / input_height
							_lowerCAmelCase     :							Any    =    output_width / input_width

							if keep_aspect_ratio:
														# scale as little as possible
														if abs(1 - scale_width  ) < abs(1 - scale_height  ):
																					# fit width
																					_lowerCAmelCase     :							Union[str, Any]    =    scale_width
														else:
																					# fit height
																					_lowerCAmelCase     :							Union[str, Any]    =    scale_height

							_lowerCAmelCase     :							List[str]    =    constraint_to_multiple_of(scale_height * input_height				, multiple=_lowerCamelCase  )
							_lowerCAmelCase     :							Dict    =    constraint_to_multiple_of(scale_width * input_width				, multiple=_lowerCamelCase  )

							return (new_height, new_width)



class   UpperCAmelCase_    (       a):
							lowerCamelCase__   			=		['pixel_values']







							def __init__(   self,			__a = True,			__a = None,			__a = PILImageResampling.BILINEAR,			__a = False,			__a = 1,			__a = True,			__a = 1 / 255,			__a = True,			__a = None,			__a = None,			**__a,			):

														'''simple docstring'''

														super().__init__(**__a)
														_lowerCAmelCase     :							Any    =    size if size is not None else {"height": 384, "width": 384}
														_lowerCAmelCase     :							Optional[int]    =    get_size_dict(__a)
														_lowerCAmelCase     :							Optional[Any]    =    do_resize
														_lowerCAmelCase     :							Dict    =    size
														_lowerCAmelCase     :							Any    =    keep_aspect_ratio
														_lowerCAmelCase     :							str    =    ensure_multiple_of
														_lowerCAmelCase     :							str    =    resample
														_lowerCAmelCase     :							Dict    =    do_rescale
														_lowerCAmelCase     :							Optional[int]    =    rescale_factor
														_lowerCAmelCase     :							Dict    =    do_normalize
														_lowerCAmelCase     :							Dict    =    image_mean if image_mean is not None else IMAGENET_STANDARD_MEAN
														_lowerCAmelCase     :							int    =    image_std if image_std is not None else IMAGENET_STANDARD_STD







							def    snake_case__      (   self,			__a,			__a,			__a = False,			__a = 1,			__a = PILImageResampling.BICUBIC,			__a = None,			**__a,			):

														'''simple docstring'''

														_lowerCAmelCase     :							List[Any]    =    get_size_dict(__a)
														if "height" not in size or "width" not in size:
																					raise ValueError(f"The size dictionary must contain the keys 'height' and 'width'. Got {size.keys()}")
														_lowerCAmelCase     :							List[Any]    =    get_resize_output_image_size(
														    __a,			output_size=(size["height"], size["width"]),			keep_aspect_ratio=__a,			multiple=__a,			)
														return resize(__a,			size=__a,			resample=__a,			data_format=__a,			**__a)







							def    snake_case__      (   self,			__a,			__a,			__a = None,			**__a,			):

														'''simple docstring'''

														return rescale(__a,			scale=__a,			data_format=__a,			**__a)







							def    snake_case__      (   self,			__a,			__a,			__a,			__a = None,			**__a,			):

														'''simple docstring'''

														return normalize(__a,			mean=__a,			std=__a,			data_format=__a,			**__a)







							def    snake_case__      (   self,			__a,			__a = None,			__a = None,			__a = None,			__a = None,			__a = None,			__a = None,			__a = None,			__a = None,			__a = None,			__a = None,			__a = None,			__a = ChannelDimension.FIRST,			**__a,			):

														'''simple docstring'''

														_lowerCAmelCase     :							int    =    do_resize if do_resize is not None else self.do_resize
														_lowerCAmelCase     :							List[Any]    =    size if size is not None else self.size
														_lowerCAmelCase     :							str    =    get_size_dict(__a)
														_lowerCAmelCase     :							Dict    =    keep_aspect_ratio if keep_aspect_ratio is not None else self.keep_aspect_ratio
														_lowerCAmelCase     :							Any    =    ensure_multiple_of if ensure_multiple_of is not None else self.ensure_multiple_of
														_lowerCAmelCase     :							int    =    resample if resample is not None else self.resample
														_lowerCAmelCase     :							Union[str, Any]    =    do_rescale if do_rescale is not None else self.do_rescale
														_lowerCAmelCase     :							Tuple    =    rescale_factor if rescale_factor is not None else self.rescale_factor
														_lowerCAmelCase     :							List[str]    =    do_normalize if do_normalize is not None else self.do_normalize
														_lowerCAmelCase     :							Dict    =    image_mean if image_mean is not None else self.image_mean
														_lowerCAmelCase     :							List[str]    =    image_std if image_std is not None else self.image_std

														_lowerCAmelCase     :							Optional[Any]    =    make_list_of_images(__a)

														if not valid_images(__a):
																					raise ValueError(
																					    "Invalid image type. Must be of type PIL.Image.Image, numpy.ndarray, "
																					    "torch.Tensor, tf.Tensor or jax.ndarray.")

														if do_resize and size is None or resample is None:
																					raise ValueError("Size and resample must be specified if do_resize is True.")

														if do_rescale and rescale_factor is None:
																					raise ValueError("Rescale factor must be specified if do_rescale is True.")

														if do_normalize and (image_mean is None or image_std is None):
																					raise ValueError("Image mean and std must be specified if do_normalize is True.")

														# All transformations expect numpy arrays.
														_lowerCAmelCase     :							List[Any]    =    [to_numpy_array(__a) for image in images]

														if do_resize:
																					_lowerCAmelCase     :							Any    =    [self.resize(image=__a,			size=__a,			resample=__a) for image in images]

														if do_rescale:
																					_lowerCAmelCase     :							List[str]    =    [self.rescale(image=__a,			scale=__a) for image in images]

														if do_normalize:
																					_lowerCAmelCase     :							Dict    =    [self.normalize(image=__a,			mean=__a,			std=__a) for image in images]

														_lowerCAmelCase     :							List[str]    =    [to_channel_dimension_format(__a,			__a) for image in images]

														_lowerCAmelCase     :							Optional[Any]    =    {"pixel_values": images}
														return BatchFeature(data=__a,			tensor_type=__a)







							def    snake_case__      (   self,			__a,			__a = None):

														'''simple docstring'''

														_lowerCAmelCase     :							Optional[Any]    =    outputs.logits

														# Resize logits and compute semantic segmentation maps
														if target_sizes is not None:
																					if len(__a) != len(__a):
																												raise ValueError(
																												    "Make sure that you pass in as many target sizes as the batch dimension of the logits")

																					if is_torch_tensor(__a):
																												_lowerCAmelCase     :							List[Any]    =    target_sizes.numpy()

																					_lowerCAmelCase     :							Dict    =    []

																					for idx in range(len(__a)):
																												_lowerCAmelCase     :							int    =    torch.nn.functional.interpolate(
																												    logits[idx].unsqueeze(dim=0),			size=target_sizes[idx],			mode="bilinear",			align_corners=__a)
																												_lowerCAmelCase     :							int    =    resized_logits[0].argmax(dim=0)
																												semantic_segmentation.append(__a)
														else:
																					_lowerCAmelCase     :							Dict    =    logits.argmax(dim=1)
																					_lowerCAmelCase     :							str    =    [semantic_segmentation[i] for i in range(semantic_segmentation.shape[0])]

														return semantic_segmentation
