


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 				UpperCamelCase_			(				UpperCamelCase    ,						UpperCamelCase    ,						unittest.TestCase):

	"""simple docstring"""





	snake_case__       :						Any					      =    StableDiffusionSAGPipeline
	snake_case__       :						Tuple					      =    TEXT_TO_IMAGE_PARAMS
	snake_case__       :						Tuple					      =    TEXT_TO_IMAGE_BATCH_PARAMS
	snake_case__       :						List[Any]					      =    TEXT_TO_IMAGE_IMAGE_PARAMS
	snake_case__       :						List[str]					      =    TEXT_TO_IMAGE_IMAGE_PARAMS
	snake_case__       :						Tuple					      =    False




	def 						UpperCAmelCase_   (				self    :       List[str]     ) ->			List[str]:
						torch.manual_seed(0     )
						__SCREAMING_SNAKE_CASE		     =       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    ,    )
						__SCREAMING_SNAKE_CASE		     =       DDIMScheduler(
						    beta_start=0.00_085    ,    beta_end=0.012    ,    beta_schedule="scaled_linear"    ,    clip_sample=UpperCAmelCase__    ,    set_alpha_to_one=UpperCAmelCase__    ,    )
						torch.manual_seed(0     )
						__SCREAMING_SNAKE_CASE		     =       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     )
						__SCREAMING_SNAKE_CASE		     =       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    ,    )
						__SCREAMING_SNAKE_CASE		     =       CLIPTextModel(UpperCAmelCase__     )
						__SCREAMING_SNAKE_CASE		     =       CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip"     )

						__SCREAMING_SNAKE_CASE		     =       {
						    "unet": unet,
						    "scheduler": scheduler,
						    "vae": vae,
						    "text_encoder": text_encoder,
						    "tokenizer": tokenizer,
						    "safety_checker": None,
						    "feature_extractor": None,
						}
						return components




	def 						UpperCAmelCase_   (				self    :       str    ,    UpperCAmelCase__    :       Dict    ,    UpperCAmelCase__    :       Optional[Any]=0     ) ->			Dict:
						if str(UpperCAmelCase__     ).startswith("mps"     ):
											__SCREAMING_SNAKE_CASE		     =       torch.manual_seed(UpperCAmelCase__     )
						else:
											__SCREAMING_SNAKE_CASE		     =       torch.Generator(device=UpperCAmelCase__     ).manual_seed(UpperCAmelCase__     )
						__SCREAMING_SNAKE_CASE		     =       {
						    "prompt": ".",
						    "generator": generator,
						    "num_inference_steps": 2,
						    "guidance_scale": 1.0,
						    "sag_scale": 1.0,
						    "output_type": "numpy",
						}
						return inputs




	def 						UpperCAmelCase_   (				self    :       Optional[Any]     ) ->			Dict:
						super().test_inference_batch_single_identical(expected_max_diff=3E-3     )




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

	"""simple docstring"""





	def 						UpperCAmelCase_   (				self    :       int     ) ->			List[Any]:
						# clean up the VRAM after each test
						super().tearDown()
						gc.collect()
						torch.cuda.empty_cache()




	def 						UpperCAmelCase_   (				self    :       List[Any]     ) ->			List[str]:
						__SCREAMING_SNAKE_CASE		     =       StableDiffusionSAGPipeline.from_pretrained("CompVis/stable-diffusion-v1-4"     )
						__SCREAMING_SNAKE_CASE		     =       sag_pipe.to(UpperCAmelCase__     )
						sag_pipe.set_progress_bar_config(disable=UpperCAmelCase__     )

						__SCREAMING_SNAKE_CASE		     =       "."
						__SCREAMING_SNAKE_CASE		     =       torch.manual_seed(0     )
						__SCREAMING_SNAKE_CASE		     =       sag_pipe(
						    [prompt]    ,    generator=UpperCAmelCase__    ,    guidance_scale=7.5    ,    sag_scale=1.0    ,    num_inference_steps=2_0    ,    output_type="np"     )

						__SCREAMING_SNAKE_CASE		     =       output.images

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

						assert image.shape == (1, 5_1_2, 5_1_2, 3)
						__SCREAMING_SNAKE_CASE		     =       np.array([0.1_568, 0.1_738, 0.1_695, 0.1_693, 0.1_507, 0.1_705, 0.1_547, 0.1_751, 0.1_949]     )

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




	def 						UpperCAmelCase_   (				self    :       Dict     ) ->			List[Any]:
						__SCREAMING_SNAKE_CASE		     =       StableDiffusionSAGPipeline.from_pretrained("stabilityai/stable-diffusion-2-1-base"     )
						__SCREAMING_SNAKE_CASE		     =       sag_pipe.to(UpperCAmelCase__     )
						sag_pipe.set_progress_bar_config(disable=UpperCAmelCase__     )

						__SCREAMING_SNAKE_CASE		     =       "."
						__SCREAMING_SNAKE_CASE		     =       torch.manual_seed(0     )
						__SCREAMING_SNAKE_CASE		     =       sag_pipe(
						    [prompt]    ,    generator=UpperCAmelCase__    ,    guidance_scale=7.5    ,    sag_scale=1.0    ,    num_inference_steps=2_0    ,    output_type="np"     )

						__SCREAMING_SNAKE_CASE		     =       output.images

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

						assert image.shape == (1, 5_1_2, 5_1_2, 3)
						__SCREAMING_SNAKE_CASE		     =       np.array([0.3_459, 0.2_876, 0.2_537, 0.3_002, 0.2_671, 0.2_160, 0.3_026, 0.2_262, 0.2_371]     )

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




	def 						UpperCAmelCase_   (				self    :       List[str]     ) ->			Optional[Any]:
						__SCREAMING_SNAKE_CASE		     =       StableDiffusionSAGPipeline.from_pretrained("stabilityai/stable-diffusion-2-1-base"     )
						__SCREAMING_SNAKE_CASE		     =       sag_pipe.to(UpperCAmelCase__     )
						sag_pipe.set_progress_bar_config(disable=UpperCAmelCase__     )

						__SCREAMING_SNAKE_CASE		     =       "."
						__SCREAMING_SNAKE_CASE		     =       torch.manual_seed(0     )
						__SCREAMING_SNAKE_CASE		     =       sag_pipe(
						    [prompt]    ,    width=7_6_8    ,    height=5_1_2    ,    generator=UpperCAmelCase__    ,    guidance_scale=7.5    ,    sag_scale=1.0    ,    num_inference_steps=2_0    ,    output_type="np"    ,    )

						__SCREAMING_SNAKE_CASE		     =       output.images

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




"""simple docstring"""







from math import asin, atan, cos, radians, sin, sqrt, tan

a__       :							List[str]        =  6_37_81_37.0
a__       :							Tuple        =  6_35_67_52.31_42_45
a__       :							str        =  6_3_7_8_1_3_7




def      UpperCAmelCase__ (lowerCAmelCase_						,		lowerCAmelCase_						,		lowerCAmelCase_						,		lowerCAmelCase_     ):




					'''simple docstring'''


					__SCREAMING_SNAKE_CASE		     =       (AXIS_A - AXIS_B) / AXIS_A
					__SCREAMING_SNAKE_CASE		     =       atan((1 - flattening) * tan(radians(lowerCAmelCase_     )     )     )
					__SCREAMING_SNAKE_CASE		     =       atan((1 - flattening) * tan(radians(lowerCAmelCase_     )     )     )
					__SCREAMING_SNAKE_CASE		     =       radians(lowerCAmelCase_     )
					__SCREAMING_SNAKE_CASE		     =       radians(lowerCAmelCase_     )
					# Equation
					__SCREAMING_SNAKE_CASE		     =       sin((phi_a - phi_a) / 2     )
					__SCREAMING_SNAKE_CASE		     =       sin((lambda_a - lambda_a) / 2     )
					# Square both values
					sin_sq_phi *= sin_sq_phi
					sin_sq_lambda *= sin_sq_lambda
					__SCREAMING_SNAKE_CASE		     =       sqrt(sin_sq_phi + (cos(lowerCAmelCase_     ) * cos(lowerCAmelCase_     ) * sin_sq_lambda)     )
					return 2 * RADIUS * asin(lowerCAmelCase_     )


if __name__ == "__main__":
		import doctest

		doctest.testmod()
