





'''simple docstring'''





import gc
import unittest

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

from diffusers import (
    AutoencoderKL,
    DDIMScheduler,
    StableDiffusionSAGPipeline,
    UNetaDConditionModel,
)
from diffusers.utils import slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu

from ..pipeline_params import TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_IMAGE_PARAMS, TEXT_TO_IMAGE_PARAMS
from ..test_pipelines_common import PipelineLatentTesterMixin, PipelineTesterMixin


enable_full_determinism()



class      _SCREAMING_SNAKE_CASE		(			lowerCamelCase__					,							lowerCamelCase__					,							unittest.TestCase					):







 '''simple docstring'''
 lowercase_	=    StableDiffusionSAGPipeline
 lowercase_	=    TEXT_TO_IMAGE_PARAMS
 lowercase_	=    TEXT_TO_IMAGE_BATCH_PARAMS
 lowercase_	=    TEXT_TO_IMAGE_IMAGE_PARAMS
 lowercase_	=    TEXT_TO_IMAGE_IMAGE_PARAMS
 lowercase_	=    False

 def 							SCREAMING_SNAKE_CASE_      (self       :     int) ->List[str]:



   '''simple docstring'''



   torch.manual_seed(0)
   lowerCamelCase__:   List[Any]    						=UNetaDConditionModel(
       block_out_channels=(32, 64) ,				layers_per_block=2 ,				sample_size=32 ,				in_channels=4 ,				out_channels=4 ,				down_block_types=("DownBlock2D", "CrossAttnDownBlock2D") ,				up_block_types=("CrossAttnUpBlock2D", "UpBlock2D") ,				cross_attention_dim=32 ,				)
   lowerCamelCase__:   Optional[Any]    						=DDIMScheduler(
       beta_start=0.0_0085 ,				beta_end=0.012 ,				beta_schedule="scaled_linear" ,				clip_sample=__lowerCamelCase ,				set_alpha_to_one=__lowerCamelCase ,				)
   torch.manual_seed(0)
   lowerCamelCase__:   List[Any]    						=AutoencoderKL(
       block_out_channels=[32, 64] ,				in_channels=3 ,				out_channels=3 ,				down_block_types=["DownEncoderBlock2D", "DownEncoderBlock2D"] ,				up_block_types=["UpDecoderBlock2D", "UpDecoderBlock2D"] ,				latent_channels=4 ,				)
   torch.manual_seed(0)
   lowerCamelCase__:   Union[str, Any]    						=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=1_000 ,				)
   lowerCamelCase__:   int    						=CLIPTextModel(__lowerCamelCase)
   lowerCamelCase__:   Optional[int]    						=CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip")

   lowerCamelCase__:   str    						={
       '''unet''': unet,
       '''scheduler''': scheduler,
       '''vae''': vae,
       '''text_encoder''': text_encoder,
       '''tokenizer''': tokenizer,
       '''safety_checker''': None,
       '''feature_extractor''': None,
   }
   return components

 def 							SCREAMING_SNAKE_CASE_      (self       :     Optional[int] ,				UpperCAmelCase_       :     List[Any] ,				UpperCAmelCase_       :     Any=0) ->List[str]:



   '''simple docstring'''



   if str(__lowerCamelCase).startswith("mps"):
     lowerCamelCase__:   Union[str, Any]    						=torch.manual_seed(__lowerCamelCase)
   else:
     lowerCamelCase__:   Any    						=torch.Generator(device=__lowerCamelCase).manual_seed(__lowerCamelCase)
   lowerCamelCase__:   str    						={
       '''prompt''': '''.''',
       '''generator''': generator,
       '''num_inference_steps''': 2,
       '''guidance_scale''': 1.0,
       '''sag_scale''': 1.0,
       '''output_type''': '''numpy''',
   }
   return inputs







 def 							SCREAMING_SNAKE_CASE_      (self       :     Dict) ->List[str]:



   '''simple docstring'''



   super().test_inference_batch_single_identical(expected_max_diff=3E-3)



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







 '''simple docstring'''

 def 							SCREAMING_SNAKE_CASE_      (self       :     List[str]) ->List[Any]:



   '''simple docstring'''



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

 def 							SCREAMING_SNAKE_CASE_      (self       :     List[Any]) ->Any:



   '''simple docstring'''



   lowerCamelCase__:   Dict    						=StableDiffusionSAGPipeline.from_pretrained("CompVis/stable-diffusion-v1-4")
   lowerCamelCase__:   Optional[int]    						=sag_pipe.to(__lowerCamelCase)
   sag_pipe.set_progress_bar_config(disable=__lowerCamelCase)

   lowerCamelCase__:   Union[str, Any]    						='''.'''
   lowerCamelCase__:   Dict    						=torch.manual_seed(0)
   lowerCamelCase__:   Dict    						=sag_pipe(
       [prompt] ,				generator=__lowerCamelCase ,				guidance_scale=7.5 ,				sag_scale=1.0 ,				num_inference_steps=20 ,				output_type="np")

   lowerCamelCase__:   List[Any]    						=output.images

   lowerCamelCase__:   int    						=image[0, -3:, -3:, -1]

   assert image.shape == (1, 512, 512, 3)
   lowerCamelCase__:   Any    						=np.array([0.1568, 0.1738, 0.1695, 0.1693, 0.1507, 0.1705, 0.1547, 0.1751, 0.1949])

   assert np.abs(image_slice.flatten() - expected_slice).max() < 5E-2

 def 							SCREAMING_SNAKE_CASE_      (self       :     Any) ->Tuple:



   '''simple docstring'''



   lowerCamelCase__:   int    						=StableDiffusionSAGPipeline.from_pretrained("stabilityai/stable-diffusion-2-1-base")
   lowerCamelCase__:   List[Any]    						=sag_pipe.to(__lowerCamelCase)
   sag_pipe.set_progress_bar_config(disable=__lowerCamelCase)

   lowerCamelCase__:   int    						='''.'''
   lowerCamelCase__:   Tuple    						=torch.manual_seed(0)
   lowerCamelCase__:   Any    						=sag_pipe(
       [prompt] ,				generator=__lowerCamelCase ,				guidance_scale=7.5 ,				sag_scale=1.0 ,				num_inference_steps=20 ,				output_type="np")

   lowerCamelCase__:   Dict    						=output.images

   lowerCamelCase__:   Any    						=image[0, -3:, -3:, -1]

   assert image.shape == (1, 512, 512, 3)
   lowerCamelCase__:   Optional[Any]    						=np.array([0.3459, 0.2876, 0.2537, 0.3002, 0.2671, 0.2160, 0.3026, 0.2262, 0.2371])

   assert np.abs(image_slice.flatten() - expected_slice).max() < 5E-2







 def 							SCREAMING_SNAKE_CASE_      (self       :     Union[str, Any]) ->Optional[int]:



   '''simple docstring'''



   lowerCamelCase__:   Any    						=StableDiffusionSAGPipeline.from_pretrained("stabilityai/stable-diffusion-2-1-base")
   lowerCamelCase__:   Optional[int]    						=sag_pipe.to(__lowerCamelCase)
   sag_pipe.set_progress_bar_config(disable=__lowerCamelCase)

   lowerCamelCase__:   Optional[int]    						='''.'''
   lowerCamelCase__:   int    						=torch.manual_seed(0)
   lowerCamelCase__:   Optional[int]    						=sag_pipe(
       [prompt] ,				width=768 ,				height=512 ,				generator=__lowerCamelCase ,				guidance_scale=7.5 ,				sag_scale=1.0 ,				num_inference_steps=20 ,				output_type="np" ,				)

   lowerCamelCase__:   List[str]    						=output.images

   assert image.shape == (1, 512, 768, 3)








from typing import Dict, List, Optional, 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_vision_available, logging


if is_vision_available():
     import PIL


__A             =       logging.get_logger(__name__)



class      _SCREAMING_SNAKE_CASE		(			__SCREAMING_SNAKE_CASE					):







 '''simple docstring'''
 lowercase_	=    ["pixel_values"]

 def __init__(self       :     List[str] ,				UpperCAmelCase_       :     bool = True ,				UpperCAmelCase_       :     Dict[str, int] = None ,				UpperCAmelCase_       :     float = None ,				UpperCAmelCase_       :     PILImageResampling = PILImageResampling.BILINEAR ,				UpperCAmelCase_       :     bool = True ,				UpperCAmelCase_       :     Union[int, float] = 1 / 255 ,				UpperCAmelCase_       :     bool = True ,				UpperCAmelCase_       :     Optional[Union[float, List[float]]] = None ,				UpperCAmelCase_       :     Optional[Union[float, List[float]]] = None ,				**UpperCAmelCase_       :     List[str] ,				) ->None:



   '''simple docstring'''



   super().__init__(**UpperCAmelCase_)
   lowerCamelCase__:   int    						=size if size is not None else {"shortest_edge": 384}
   lowerCamelCase__:   Tuple    						=get_size_dict(UpperCAmelCase_ ,				default_to_square=UpperCAmelCase_)

   lowerCamelCase__:   Union[str, Any]    						=do_resize
   lowerCamelCase__:   Union[str, Any]    						=size
   # Default value set here for backwards compatibility where the value in config is None
   lowerCamelCase__:   str    						=crop_pct if crop_pct is not None else 224 / 256
   lowerCamelCase__:   Optional[int]    						=resample
   lowerCamelCase__:   str    						=do_rescale
   lowerCamelCase__:   List[Any]    						=rescale_factor
   lowerCamelCase__:   Dict    						=do_normalize
   lowerCamelCase__:   List[str]    						=image_mean if image_mean is not None else IMAGENET_STANDARD_MEAN
   lowerCamelCase__:   Dict    						=image_std if image_std is not None else IMAGENET_STANDARD_STD

 def 							SCREAMING_SNAKE_CASE_      (self       :     Optional[int] ,				UpperCAmelCase_       :     np.ndarray ,				UpperCAmelCase_       :     Dict[str, int] ,				UpperCAmelCase_       :     float ,				UpperCAmelCase_       :     PILImageResampling = PILImageResampling.BICUBIC ,				UpperCAmelCase_       :     Optional[Union[str, ChannelDimension]] = None ,				**UpperCAmelCase_       :     str ,				) ->np.ndarray:



   '''simple docstring'''



   lowerCamelCase__:   List[str]    						=get_size_dict(UpperCAmelCase_ ,				default_to_square=UpperCAmelCase_)
   if "shortest_edge" not in size:
     raise ValueError(F"""Size dictionary must contain 'shortest_edge' key. Got {size.keys()}""")
   lowerCamelCase__:   Tuple    						=size["shortest_edge"]

   if shortest_edge < 384:
     # maintain same ratio, resizing shortest edge to shortest_edge/crop_pct
     lowerCamelCase__:   List[str]    						=int(shortest_edge / crop_pct)
     lowerCamelCase__:   List[str]    						=get_resize_output_image_size(UpperCAmelCase_ ,				size=UpperCAmelCase_ ,				default_to_square=UpperCAmelCase_)
     lowerCamelCase__:   List[str]    						=resize(image=UpperCAmelCase_ ,				size=UpperCAmelCase_ ,				resample=UpperCAmelCase_ ,				data_format=UpperCAmelCase_ ,				**UpperCAmelCase_)
     # then crop to (shortest_edge, shortest_edge)
     return center_crop(image=UpperCAmelCase_ ,				size=(shortest_edge, shortest_edge) ,				data_format=UpperCAmelCase_ ,				**UpperCAmelCase_)
   else:
     # warping (no cropping) when evaluated at 384 or larger
     return resize(
         UpperCAmelCase_ ,				size=(shortest_edge, shortest_edge) ,				resample=UpperCAmelCase_ ,				data_format=UpperCAmelCase_ ,				**UpperCAmelCase_)

 def 							SCREAMING_SNAKE_CASE_      (self       :     Optional[int] ,				UpperCAmelCase_       :     np.ndarray ,				UpperCAmelCase_       :     Union[int, float] ,				UpperCAmelCase_       :     Optional[Union[str, ChannelDimension]] = None ,				**UpperCAmelCase_       :     str ,				) ->Tuple:



   '''simple docstring'''



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

 def 							SCREAMING_SNAKE_CASE_      (self       :     Dict ,				UpperCAmelCase_       :     np.ndarray ,				UpperCAmelCase_       :     Union[float, List[float]] ,				UpperCAmelCase_       :     Union[float, List[float]] ,				UpperCAmelCase_       :     Optional[Union[str, ChannelDimension]] = None ,				**UpperCAmelCase_       :     Optional[Any] ,				) ->np.ndarray:



   '''simple docstring'''



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







 def 							SCREAMING_SNAKE_CASE_      (self       :     Dict ,				UpperCAmelCase_       :     ImageInput ,				UpperCAmelCase_       :     bool = None ,				UpperCAmelCase_       :     Dict[str, int] = None ,				UpperCAmelCase_       :     float = None ,				UpperCAmelCase_       :     PILImageResampling = None ,				UpperCAmelCase_       :     bool = None ,				UpperCAmelCase_       :     float = None ,				UpperCAmelCase_       :     bool = None ,				UpperCAmelCase_       :     Optional[Union[float, List[float]]] = None ,				UpperCAmelCase_       :     Optional[Union[float, List[float]]] = None ,				UpperCAmelCase_       :     Optional[Union[str, TensorType]] = None ,				UpperCAmelCase_       :     ChannelDimension = ChannelDimension.FIRST ,				**UpperCAmelCase_       :     Dict ,				) ->PIL.Image.Image:



   '''simple docstring'''



   lowerCamelCase__:   List[Any]    						=do_resize if do_resize is not None else self.do_resize
   lowerCamelCase__:   List[Any]    						=crop_pct if crop_pct is not None else self.crop_pct
   lowerCamelCase__:   str    						=resample if resample is not None else self.resample
   lowerCamelCase__:   str    						=do_rescale if do_rescale is not None else self.do_rescale
   lowerCamelCase__:   Optional[Any]    						=rescale_factor if rescale_factor is not None else self.rescale_factor
   lowerCamelCase__:   int    						=do_normalize if do_normalize is not None else self.do_normalize
   lowerCamelCase__:   Any    						=image_mean if image_mean is not None else self.image_mean
   lowerCamelCase__:   Any    						=image_std if image_std is not None else self.image_std

   lowerCamelCase__:   List[str]    						=size if size is not None else self.size
   lowerCamelCase__:   List[str]    						=get_size_dict(UpperCAmelCase_ ,				default_to_square=UpperCAmelCase_)

   lowerCamelCase__:   Any    						=make_list_of_images(UpperCAmelCase_)

   if not valid_images(UpperCAmelCase_):
     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_resize and size["shortest_edge"] < 384 and crop_pct is None:
     raise ValueError("crop_pct must be specified if size < 384.")

   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(UpperCAmelCase_) for image in images]

   if do_resize:
     lowerCamelCase__:   Dict    						=[self.resize(image=UpperCAmelCase_ ,				size=UpperCAmelCase_ ,				crop_pct=UpperCAmelCase_ ,				resample=UpperCAmelCase_) for image in images]

   if do_rescale:
     lowerCamelCase__:   Union[str, Any]    						=[self.rescale(image=UpperCAmelCase_ ,				scale=UpperCAmelCase_) for image in images]

   if do_normalize:
     lowerCamelCase__:   Tuple    						=[self.normalize(image=UpperCAmelCase_ ,				mean=UpperCAmelCase_ ,				std=UpperCAmelCase_) for image in images]

   lowerCamelCase__:   Optional[Any]    						=[to_channel_dimension_format(UpperCAmelCase_ ,				UpperCAmelCase_) for image in images]

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

