



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 	__snake_case	(     _UpperCAmelCase  ,					_UpperCAmelCase  ,					unittest.TestCase):
  snake_case__    :       Tuple        =						StableDiffusionSAGPipeline
  snake_case__    :       Optional[int]        =						TEXT_TO_IMAGE_PARAMS
  snake_case__    :       Union[str, Any]        =						TEXT_TO_IMAGE_BATCH_PARAMS
  snake_case__    :       Optional[Any]        =						TEXT_TO_IMAGE_IMAGE_PARAMS
  snake_case__    :       str        =						TEXT_TO_IMAGE_IMAGE_PARAMS
  snake_case__    :       Optional[Any]        =						False




  def 						SCREAMING_SNAKE_CASE    (			self					:						List[str]		):






         """simple docstring"""






         torch.manual_seed(0		)
         _lowerCamelCase						:    str         =			UNetaDConditionModel(
             block_out_channels=(3_2, 6_4)						, layers_per_block=2						, sample_size=3_2						, in_channels=4						, out_channels=4						, down_block_types=('''DownBlock2D''', '''CrossAttnDownBlock2D''')						, up_block_types=('''CrossAttnUpBlock2D''', '''UpBlock2D''')						, cross_attention_dim=3_2						, )
         _lowerCamelCase						:    Any         =			DDIMScheduler(
             beta_start=0.0_00_85						, beta_end=0.0_12						, beta_schedule='''scaled_linear'''						, clip_sample=lowercase_						, set_alpha_to_one=lowercase_						, )
         torch.manual_seed(0		)
         _lowerCamelCase						:    Dict         =			AutoencoderKL(
             block_out_channels=[3_2, 6_4]						, in_channels=3						, out_channels=3						, down_block_types=['''DownEncoderBlock2D''', '''DownEncoderBlock2D''']						, up_block_types=['''UpDecoderBlock2D''', '''UpDecoderBlock2D''']						, latent_channels=4						, )
         torch.manual_seed(0		)
         _lowerCamelCase						:    Any         =			CLIPTextConfig(
             bos_token_id=0						, eos_token_id=2						, hidden_size=3_2						, intermediate_size=3_7						, layer_norm_eps=1E-05						, num_attention_heads=4						, num_hidden_layers=5						, pad_token_id=1						, vocab_size=1_0_0_0						, )
         _lowerCamelCase						:    Dict         =			CLIPTextModel(lowercase_		)
         _lowerCamelCase						:    int         =			CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''		)

         _lowerCamelCase						:    Dict         =			{
             """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					:						List[str]						, __lowerCAmelCase					:						Tuple						, __lowerCAmelCase					:						List[str]=0		):






         """simple docstring"""






         if str(lowercase_		).startswith('''mps'''		):
                _lowerCamelCase						:    Tuple         =			torch.manual_seed(lowercase_		)
         else:
                _lowerCamelCase						:    str         =			torch.Generator(device=lowercase_		).manual_seed(lowercase_		)
         _lowerCamelCase						:    Any         =			{
             """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					:						List[Any]		):






         """simple docstring"""






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






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




  def 						SCREAMING_SNAKE_CASE    (			self					:						List[str]		):






         """simple docstring"""






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




  def 						SCREAMING_SNAKE_CASE    (			self					:						List[str]		):






         """simple docstring"""






         _lowerCamelCase						:    str         =			StableDiffusionSAGPipeline.from_pretrained('''CompVis/stable-diffusion-v1-4'''		)
         _lowerCamelCase						:    Optional[int]         =			sag_pipe.to(lowercase_		)
         sag_pipe.set_progress_bar_config(disable=lowercase_		)

         _lowerCamelCase						:    str         =			"""."""
         _lowerCamelCase						:    List[Any]         =			torch.manual_seed(0		)
         _lowerCamelCase						:    int         =			sag_pipe(
             [prompt]						, generator=lowercase_						, guidance_scale=7.5						, sag_scale=1.0						, num_inference_steps=2_0						, output_type='''np'''		)

         _lowerCamelCase						:    Dict         =			output.images

         _lowerCamelCase						:    Union[str, Any]         =			image[0, -3:, -3:, -1]

         assert image.shape == (1, 5_1_2, 5_1_2, 3)
         _lowerCamelCase						:    List[str]         =			np.array([0.15_68, 0.17_38, 0.16_95, 0.16_93, 0.15_07, 0.17_05, 0.15_47, 0.17_51, 0.19_49]		)

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




  def 						SCREAMING_SNAKE_CASE    (			self					:						Dict		):






         """simple docstring"""






         _lowerCamelCase						:    Tuple         =			StableDiffusionSAGPipeline.from_pretrained('''stabilityai/stable-diffusion-2-1-base'''		)
         _lowerCamelCase						:    Tuple         =			sag_pipe.to(lowercase_		)
         sag_pipe.set_progress_bar_config(disable=lowercase_		)

         _lowerCamelCase						:    str         =			"""."""
         _lowerCamelCase						:    Any         =			torch.manual_seed(0		)
         _lowerCamelCase						:    Optional[int]         =			sag_pipe(
             [prompt]						, generator=lowercase_						, guidance_scale=7.5						, sag_scale=1.0						, num_inference_steps=2_0						, output_type='''np'''		)

         _lowerCamelCase						:    str         =			output.images

         _lowerCamelCase						:    Optional[Any]         =			image[0, -3:, -3:, -1]

         assert image.shape == (1, 5_1_2, 5_1_2, 3)
         _lowerCamelCase						:    Any         =			np.array([0.34_59, 0.28_76, 0.25_37, 0.30_02, 0.26_71, 0.21_60, 0.30_26, 0.22_62, 0.23_71]		)

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




  def 						SCREAMING_SNAKE_CASE    (			self					:						str		):






         """simple docstring"""






         _lowerCamelCase						:    List[str]         =			StableDiffusionSAGPipeline.from_pretrained('''stabilityai/stable-diffusion-2-1-base'''		)
         _lowerCamelCase						:    Union[str, Any]         =			sag_pipe.to(lowercase_		)
         sag_pipe.set_progress_bar_config(disable=lowercase_		)

         _lowerCamelCase						:    Optional[Any]         =			"""."""
         _lowerCamelCase						:    List[str]         =			torch.manual_seed(0		)
         _lowerCamelCase						:    int         =			sag_pipe(
             [prompt]						, width=7_6_8						, height=5_1_2						, generator=lowercase_						, guidance_scale=7.5						, sag_scale=1.0						, num_inference_steps=2_0						, output_type='''np'''						, )

         _lowerCamelCase						:    int         =			output.images

         assert image.shape == (1, 5_1_2, 7_6_8, 3)







"""simple docstring"""






import math






def 	snake_case_  (				A_ :		int						):


       '''simple docstring'''





       return math.sqrt(A_						) * math.sqrt(A_						) == num






def 	snake_case_  (				A_ :		int						):


       '''simple docstring'''





       _lowerCamelCase						:    Dict         =			0
       _lowerCamelCase						:    Tuple         =			n
       while left <= right:
              _lowerCamelCase						:    List[str]         =			(left + right) // 2
              if mid**2 == n:
                     return True
              elif mid**2 > n:
                     _lowerCamelCase						:    int         =			mid - 1
              else:
                     _lowerCamelCase						:    Optional[Any]         =			mid + 1
       return False


if __name__ == "__main__":
   import doctest

   doctest.testmod()


