





"""simple docstring"""





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

import numpy as np

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


if is_torch_available():
  import torch


__A											=						logging.get_logger(__name__)






class    _lowerCAmelCase  (	_a			):
       """simple docstring"""




       __magic_name__	:Any						    =   ["""pixel_values"""]






       def __init__(       self       ,					__UpperCAmelCase = True       ,					__UpperCAmelCase = None       ,					__UpperCAmelCase = PILImageResampling.BILINEAR       ,					__UpperCAmelCase = True       ,					__UpperCAmelCase = None       ,					__UpperCAmelCase = True       ,					__UpperCAmelCase = 1 / 2_5_5       ,					__UpperCAmelCase = True       ,					__UpperCAmelCase = None       ,					__UpperCAmelCase = None       ,					**__UpperCAmelCase       ,					):



         '''simple docstring'''

         super().__init__(**lowerCAmelCase_ )
         lowerCAmelCase__     :Optional[int]											=				size if size is not None else {'shortest_edge': 2_5_6}
         lowerCAmelCase__     :Optional[int]											=				get_size_dict(lowerCAmelCase_       ,					default_to_square=lowerCAmelCase_ )
         lowerCAmelCase__     :Optional[int]											=				crop_size if crop_size is not None else {'height': 2_2_4, 'width': 2_2_4}
         lowerCAmelCase__     :Optional[int]											=				get_size_dict(lowerCAmelCase_       ,					param_name='crop_size' )
         lowerCAmelCase__     :str											=				do_resize
         lowerCAmelCase__     :str											=				size
         lowerCAmelCase__     :int											=				resample
         lowerCAmelCase__     :List[str]											=				do_center_crop
         lowerCAmelCase__     :List[str]											=				crop_size
         lowerCAmelCase__     :Optional[int]											=				do_rescale
         lowerCAmelCase__     :List[Any]											=				rescale_factor
         lowerCAmelCase__     :Optional[int]											=				do_normalize
         lowerCAmelCase__     :Tuple											=				image_mean if image_mean is not None else IMAGENET_STANDARD_MEAN
         lowerCAmelCase__     :List[str]											=				image_std if image_std is not None else IMAGENET_STANDARD_STD






       def 	snake_case		(       self       ,					__UpperCAmelCase       ,					__UpperCAmelCase       ,					__UpperCAmelCase = PILImageResampling.BICUBIC       ,					__UpperCAmelCase = None       ,					**__UpperCAmelCase       ,					):



         '''simple docstring'''

         lowerCAmelCase__     :List[str]											=				get_size_dict(lowerCAmelCase_       ,					default_to_square=lowerCAmelCase_ )
         if "shortest_edge" not in size:
           raise ValueError(F"The `size` parameter must contain the key `shortest_edge`. Got {size.keys()}" )
         lowerCAmelCase__     :Dict											=				get_resize_output_image_size(lowerCAmelCase_       ,					size=size['shortest_edge']       ,					default_to_square=lowerCAmelCase_ )
         return resize(lowerCAmelCase_       ,					size=lowerCAmelCase_       ,					resample=lowerCAmelCase_       ,					data_format=lowerCAmelCase_       ,					**lowerCAmelCase_ )






       def 	snake_case		(       self       ,					__UpperCAmelCase       ,					__UpperCAmelCase       ,					__UpperCAmelCase = None       ,					**__UpperCAmelCase       ,					):



         '''simple docstring'''

         lowerCAmelCase__     :Union[str, Any]											=				get_size_dict(lowerCAmelCase_ )
         if "height" not in size or "width" not in size:
           raise ValueError(F"The `size` parameter must contain the keys `height` and `width`. Got {size.keys()}" )
         return center_crop(lowerCAmelCase_       ,					size=(size['height'], size['width'])       ,					data_format=lowerCAmelCase_       ,					**lowerCAmelCase_ )






       def 	snake_case		(       self       ,					__UpperCAmelCase       ,					__UpperCAmelCase       ,					__UpperCAmelCase = None       ,					**__UpperCAmelCase ):



         '''simple docstring'''

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






       def 	snake_case		(       self       ,					__UpperCAmelCase       ,					__UpperCAmelCase       ,					__UpperCAmelCase       ,					__UpperCAmelCase = None       ,					**__UpperCAmelCase       ,					):



         '''simple docstring'''

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






       def 	snake_case		(       self       ,					__UpperCAmelCase       ,					__UpperCAmelCase = None       ,					__UpperCAmelCase = None       ,					__UpperCAmelCase = None       ,					__UpperCAmelCase = None       ,					__UpperCAmelCase = None       ,					__UpperCAmelCase = None       ,					__UpperCAmelCase = None       ,					__UpperCAmelCase = None       ,					__UpperCAmelCase = None       ,					__UpperCAmelCase = None       ,					__UpperCAmelCase = None       ,					__UpperCAmelCase = ChannelDimension.FIRST       ,					**__UpperCAmelCase       ,					):



         '''simple docstring'''

         lowerCAmelCase__     :Union[str, Any]											=				do_resize if do_resize is not None else self.do_resize
         lowerCAmelCase__     :Any											=				size if size is not None else self.size
         lowerCAmelCase__     :Optional[int]											=				get_size_dict(lowerCAmelCase_       ,					default_to_square=lowerCAmelCase_ )
         lowerCAmelCase__     :int											=				resample if resample is not None else self.resample
         lowerCAmelCase__     :Tuple											=				do_center_crop if do_center_crop is not None else self.do_center_crop
         lowerCAmelCase__     :Union[str, Any]											=				crop_size if crop_size is not None else self.crop_size
         lowerCAmelCase__     :Union[str, Any]											=				get_size_dict(lowerCAmelCase_       ,					param_name='crop_size' )
         lowerCAmelCase__     :Union[str, Any]											=				do_rescale if do_rescale is not None else self.do_rescale
         lowerCAmelCase__     :List[str]											=				rescale_factor if rescale_factor is not None else self.rescale_factor
         lowerCAmelCase__     :List[Any]											=				do_normalize if do_normalize is not None else self.do_normalize
         lowerCAmelCase__     :Tuple											=				image_mean if image_mean is not None else self.image_mean
         lowerCAmelCase__     :Optional[Any]											=				image_std if image_std is not None else self.image_std

         lowerCAmelCase__     :List[Any]											=				make_list_of_images(lowerCAmelCase_ )

         if not valid_images(lowerCAmelCase_ ):
           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:
           raise ValueError('Size must be specified if do_resize is True.' )

         if do_center_crop and crop_size is None:
           raise ValueError('Crop size must be specified if do_center_crop 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__     :int											=				[to_numpy_array(lowerCAmelCase_ ) for image in images]

         if do_resize:
           lowerCAmelCase__     :int											=				[self.resize(image=lowerCAmelCase_       ,					size=lowerCAmelCase_       ,					resample=lowerCAmelCase_ ) for image in images]

         if do_center_crop:
           lowerCAmelCase__     :List[str]											=				[self.center_crop(image=lowerCAmelCase_       ,					size=lowerCAmelCase_ ) for image in images]

         if do_rescale:
           lowerCAmelCase__     :Optional[Any]											=				[self.rescale(image=lowerCAmelCase_       ,					scale=lowerCAmelCase_ ) for image in images]

         if do_normalize:
           lowerCAmelCase__     :Optional[Any]											=				[self.normalize(image=lowerCAmelCase_       ,					mean=lowerCAmelCase_       ,					std=lowerCAmelCase_ ) for image in images]

         lowerCAmelCase__     :List[str]											=				[to_channel_dimension_format(lowerCAmelCase_       ,					lowerCAmelCase_ ) for image in images]

         lowerCAmelCase__     :Union[str, Any]											=				{'pixel_values': images}
         return BatchFeature(data=lowerCAmelCase_       ,					tensor_type=lowerCAmelCase_ )






       def 	snake_case		(       self       ,					__UpperCAmelCase       ,					__UpperCAmelCase = None ):



         '''simple docstring'''

         lowerCAmelCase__     :Tuple											=				outputs.logits

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

           if is_torch_tensor(lowerCAmelCase_ ):
             lowerCAmelCase__     :str											=				target_sizes.numpy()

           lowerCAmelCase__     :Union[str, Any]											=				[]

           for idx in range(len(lowerCAmelCase_ ) ):
             lowerCAmelCase__     :Tuple											=				torch.nn.functional.interpolate(
                 logits[idx].unsqueeze(dim=0 )       ,					size=target_sizes[idx]       ,					mode='bilinear'       ,					align_corners=lowerCAmelCase_ )
             lowerCAmelCase__     :str											=				resized_logits[0].argmax(dim=0 )
             semantic_segmentation.append(lowerCAmelCase_ )
         else:
           lowerCAmelCase__     :str											=				logits.argmax(dim=1 )
           lowerCAmelCase__     :List[str]											=				[semantic_segmentation[i] for i in range(semantic_segmentation.shape[0] )]

         return semantic_segmentation






'''simple docstring'''


import gc
import unittest

import numpy as np
import torch

from diffusers import StableDiffusionKDiffusionPipeline
from diffusers.utils import slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu


enable_full_determinism()

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







							def   __lowerCAmelCase     (    self			:    int   )     ->				Any:


														"""simple docstring"""


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







							def   __lowerCAmelCase     (    self			:    List[Any]   )     ->				int:


														"""simple docstring"""


														_a      	=							StableDiffusionKDiffusionPipeline.from_pretrained('''CompVis/stable-diffusion-v1-4'''   )
														_a      	=							sd_pipe.to(lowerCAmelCase_   )
														sd_pipe.set_progress_bar_config(disable=lowerCAmelCase_   )

														sd_pipe.set_scheduler('''sample_euler'''   )

														_a      	=							'''A painting of a squirrel eating a burger'''
														_a      	=							torch.manual_seed(0   )
														_a      	=							sd_pipe([prompt]				,     generator=lowerCAmelCase_				,     guidance_scale=9.0				,     num_inference_steps=20				,     output_type='''np'''   )

														_a      	=							output.images

														_a      	=							image[0, -3:, -3:, -1]

														assert image.shape == (1, 5_12, 5_12, 3)
														_a      	=							np.array([0.0_4_4_7, 0.0_4_9_2, 0.0_4_6_8, 0.0_4_0_8, 0.0_3_8_3, 0.0_4_0_8, 0.0_3_5_4, 0.0_3_8_0, 0.0_3_3_9]   )

														assert np.abs(image_slice.flatten() - expected_slice   ).max() < 1e-2







							def   __lowerCAmelCase     (    self			:    Any   )     ->				Optional[Any]:


														"""simple docstring"""


														_a      	=							StableDiffusionKDiffusionPipeline.from_pretrained('''stabilityai/stable-diffusion-2-1-base'''   )
														_a      	=							sd_pipe.to(lowerCAmelCase_   )
														sd_pipe.set_progress_bar_config(disable=lowerCAmelCase_   )

														sd_pipe.set_scheduler('''sample_euler'''   )

														_a      	=							'''A painting of a squirrel eating a burger'''
														_a      	=							torch.manual_seed(0   )
														_a      	=							sd_pipe([prompt]				,     generator=lowerCAmelCase_				,     guidance_scale=9.0				,     num_inference_steps=20				,     output_type='''np'''   )

														_a      	=							output.images

														_a      	=							image[0, -3:, -3:, -1]

														assert image.shape == (1, 5_12, 5_12, 3)
														_a      	=							np.array([0.1_2_3_7, 0.1_3_2_0, 0.1_4_3_8, 0.1_3_5_9, 0.1_3_9_0, 0.1_1_3_2, 0.1_2_7_7, 0.1_1_7_5, 0.1_1_1_2]   )

														assert np.abs(image_slice.flatten() - expected_slice   ).max() < 5e-1







							def   __lowerCAmelCase     (    self			:    Dict   )     ->				Optional[Any]:


														"""simple docstring"""


														_a      	=							StableDiffusionKDiffusionPipeline.from_pretrained('''stabilityai/stable-diffusion-2-1-base'''   )
														_a      	=							sd_pipe.to(lowerCAmelCase_   )
														sd_pipe.set_progress_bar_config(disable=lowerCAmelCase_   )

														sd_pipe.set_scheduler('''sample_dpmpp_2m'''   )

														_a      	=							'''A painting of a squirrel eating a burger'''
														_a      	=							torch.manual_seed(0   )
														_a      	=							sd_pipe(
														    [prompt]				,     generator=lowerCAmelCase_				,     guidance_scale=7.5				,     num_inference_steps=15				,     output_type='''np'''				,     use_karras_sigmas=lowerCAmelCase_				,     )

														_a      	=							output.images

														_a      	=							image[0, -3:, -3:, -1]

														assert image.shape == (1, 5_12, 5_12, 3)
														_a      	=							np.array(
														    [0.1_1_3_8_1_6_8_9, 0.1_2_1_1_2_9_2_1, 0.1_3_8_9_4_5_7, 0.1_2_5_4_9_6_0_6, 0.1_2_4_4_9_6_4, 0.1_0_8_3_1_5_1_7, 0.1_1_5_6_2_8_6_6, 0.1_0_8_6_7_8_1_6, 0.1_0_4_9_9_0_4_8]   )

														assert np.abs(image_slice.flatten() - expected_slice   ).max() < 1e-2

