


'''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        lowercase__    (		_snake_case	,	_snake_case	,	unittest.TestCase   ):





							'''simple docstring'''


							A_				: Union[str, Any]       =    StableDiffusionSAGPipeline
							A_				: Any       =    TEXT_TO_IMAGE_PARAMS
							A_				: Union[str, Any]       =    TEXT_TO_IMAGE_BATCH_PARAMS
							A_				: Dict       =    TEXT_TO_IMAGE_IMAGE_PARAMS
							A_				: Optional[Any]       =    TEXT_TO_IMAGE_IMAGE_PARAMS
							A_				: List[Any]       =    False






							def UpperCAmelCase_      (       self    ):
											torch.manual_seed(0    )
											_SCREAMING_SNAKE_CASE					:			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       ,   )
											_SCREAMING_SNAKE_CASE					:			Dict							= DDIMScheduler(
											    beta_start=0.00085       ,   beta_end=0.012       ,   beta_schedule="""scaled_linear"""       ,   clip_sample=_SCREAMING_SNAKE_CASE       ,   set_alpha_to_one=_SCREAMING_SNAKE_CASE       ,   )
											torch.manual_seed(0    )
											_SCREAMING_SNAKE_CASE					:			int							= 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    )
											_SCREAMING_SNAKE_CASE					:			Union[str, Any]							= CLIPTextConfig(
											    bos_token_id=0       ,   eos_token_id=2       ,   hidden_size=32       ,   intermediate_size=37       ,   layer_norm_eps=1e-05       ,   num_attention_heads=4       ,   num_hidden_layers=5       ,   pad_token_id=1       ,   vocab_size=1000       ,   )
											_SCREAMING_SNAKE_CASE					:			List[str]							= CLIPTextModel(_SCREAMING_SNAKE_CASE    )
											_SCREAMING_SNAKE_CASE					:			List[str]							= CLIPTokenizer.from_pretrained("""hf-internal-testing/tiny-random-clip"""    )

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






							def UpperCAmelCase_      (       self       ,   __snake_case       ,   __snake_case=0    ):





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






							def UpperCAmelCase_      (       self    ):
											super().test_inference_batch_single_identical(expected_max_diff=3e-3    )





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





							'''simple docstring'''


							def UpperCAmelCase_      (       self    ):
											super().tearDown()
											gc.collect()
											torch.cuda.empty_cache()






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

											_SCREAMING_SNAKE_CASE					:			List[str]							= """."""
											_SCREAMING_SNAKE_CASE					:			Any							= torch.manual_seed(0    )
											_SCREAMING_SNAKE_CASE					:			Tuple							= sag_pipe(
											    [prompt]       ,   generator=_SCREAMING_SNAKE_CASE       ,   guidance_scale=7.5       ,   sag_scale=1.0       ,   num_inference_steps=20       ,   output_type="""np"""    )

											_SCREAMING_SNAKE_CASE					:			List[str]							= output.images

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

											assert image.shape == (1, 512, 512, 3)
											_SCREAMING_SNAKE_CASE					:			Optional[int]							= 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 UpperCAmelCase_      (       self    ):
											_SCREAMING_SNAKE_CASE					:			Union[str, Any]							= StableDiffusionSAGPipeline.from_pretrained("""stabilityai/stable-diffusion-2-1-base"""    )
											_SCREAMING_SNAKE_CASE					:			int							= sag_pipe.to(_SCREAMING_SNAKE_CASE    )
											sag_pipe.set_progress_bar_config(disable=_SCREAMING_SNAKE_CASE    )

											_SCREAMING_SNAKE_CASE					:			Any							= """."""
											_SCREAMING_SNAKE_CASE					:			Optional[Any]							= torch.manual_seed(0    )
											_SCREAMING_SNAKE_CASE					:			str							= sag_pipe(
											    [prompt]       ,   generator=_SCREAMING_SNAKE_CASE       ,   guidance_scale=7.5       ,   sag_scale=1.0       ,   num_inference_steps=20       ,   output_type="""np"""    )

											_SCREAMING_SNAKE_CASE					:			List[str]							= output.images

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

											assert image.shape == (1, 512, 512, 3)
											_SCREAMING_SNAKE_CASE					:			Tuple							= 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 UpperCAmelCase_      (       self    ):
											_SCREAMING_SNAKE_CASE					:			Union[str, Any]							= StableDiffusionSAGPipeline.from_pretrained("""stabilityai/stable-diffusion-2-1-base"""    )
											_SCREAMING_SNAKE_CASE					:			int							= sag_pipe.to(_SCREAMING_SNAKE_CASE    )
											sag_pipe.set_progress_bar_config(disable=_SCREAMING_SNAKE_CASE    )

											_SCREAMING_SNAKE_CASE					:			str							= """."""
											_SCREAMING_SNAKE_CASE					:			List[str]							= torch.manual_seed(0    )
											_SCREAMING_SNAKE_CASE					:			List[str]							= sag_pipe(
											    [prompt]       ,   width=768       ,   height=512       ,   generator=_SCREAMING_SNAKE_CASE       ,   guidance_scale=7.5       ,   sag_scale=1.0       ,   num_inference_steps=20       ,   output_type="""np"""       ,   )

											_SCREAMING_SNAKE_CASE					:			int							= output.images

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








from typing import Any

def 			lowerCAmelCase__					(     a__:						list     ,    a__:						list     ,    a__:						dict     ,    a__:						dict     ,    a__:						dict     ,    )    ->						list:


							'''simple docstring'''





							_validation(
							    a__     ,    a__     ,    a__     ,    a__     ,    a__     ,    )
							# Creates data structures and fill initial step
							_UpperCAmelCase											=     {}
							_UpperCAmelCase											=     {}
							for state in states_space:
														_UpperCAmelCase											=     observations_space[0]
														_UpperCAmelCase											=     (
														    initial_probabilities[state] * emission_probabilities[state][observation]
														)
														_UpperCAmelCase											=     None

							# Fills the data structure with the probabilities of
							# different transitions and pointers to previous states
							for o in range(1     ,    len(a__			)			):
														_UpperCAmelCase											=     observations_space[o]
														_UpperCAmelCase											=     observations_space[o - 1]
														for state in states_space:
																					# Calculates the argmax for probability function
																					_UpperCAmelCase											=     ''
																					_UpperCAmelCase											=     -1
																					for k_state in states_space:
																												_UpperCAmelCase											=     (
																												    probabilities[(k_state, prior_observation)]
																												    * transition_probabilities[k_state][state]
																												    * emission_probabilities[state][observation]
																												)
																												if probability > max_probability:
																																			_UpperCAmelCase											=     probability
																																			_UpperCAmelCase											=     k_state

            # Update probabilities and pointers dicts
																					_UpperCAmelCase											=     (
																					    probabilities[(arg_max, prior_observation)]
																					    * transition_probabilities[arg_max][state]
																					    * emission_probabilities[state][observation]
																					)

																					_UpperCAmelCase											=     arg_max

    # The final observation
							_UpperCAmelCase											=     observations_space[len(a__			) - 1]

							# argmax for given final observation
							_UpperCAmelCase											=     ''
							_UpperCAmelCase											=     -1
							for k_state in states_space:
														_UpperCAmelCase											=     probabilities[(k_state, final_observation)]
														if probability > max_probability:
																					_UpperCAmelCase											=     probability
																					_UpperCAmelCase											=     k_state
							_UpperCAmelCase											=     arg_max

							# Process pointers backwards
							_UpperCAmelCase											=     last_state
							_UpperCAmelCase											=     []
							for o in range(len(a__			) - 1     ,    -1     ,    -1			):
														result.append(a__			)
														_UpperCAmelCase											=     pointers[previous, observations_space[o]]
							result.reverse()

							return result

def 			lowerCAmelCase__					(     a__:						Any     ,    a__:						Any     ,    a__:						Any     ,    a__:						Any     ,    a__:						Any     ,    )    ->						None:


							'''simple docstring'''





							_validate_not_empty(
							    a__     ,    a__     ,    a__     ,    a__     ,    a__     ,    )
							_validate_lists(a__     ,    a__			)
							_validate_dicts(
							    a__     ,    a__     ,    a__			)

def 			lowerCAmelCase__					(     a__:						Any     ,    a__:						Any     ,    a__:						Any     ,    a__:						Any     ,    a__:						Any     ,    )    ->						None:


							'''simple docstring'''





							if not all(
							    [
							        observations_space,
							        states_space,
							        initial_probabilities,
							        transition_probabilities,
							        emission_probabilities,
							    ]			):
														raise ValueError('There\'s an empty parameter'			)

def 			lowerCAmelCase__					(     a__:						Any     ,    a__:						Any			)    ->						None:


							'''simple docstring'''





							_validate_list(a__     ,    'observations_space'			)
							_validate_list(a__     ,    'states_space'			)

def 			lowerCAmelCase__					(     a__:						Any     ,    a__:						str			)    ->						None:


							'''simple docstring'''





							if not isinstance(_object     ,    a__			):
														_UpperCAmelCase											=     F'''{var_name} must be a list'''
														raise ValueError(a__			)
							else:
														for x in _object:
																					if not isinstance(a__     ,    a__			):
																												_UpperCAmelCase											=     F'''{var_name} must be a list of strings'''
																												raise ValueError(a__			)

def 			lowerCAmelCase__					(     a__:						Any     ,    a__:						Any     ,    a__:						Any     ,    )    ->						None:


							'''simple docstring'''





							_validate_dict(a__     ,    'initial_probabilities'     ,    a__			)
							_validate_nested_dict(a__     ,    'transition_probabilities'			)
							_validate_nested_dict(a__     ,    'emission_probabilities'			)

def 			lowerCAmelCase__					(     a__:						Any     ,    a__:						str			)    ->						None:


							'''simple docstring'''





							_validate_dict(_object     ,    a__     ,    a__			)
							for x in _object.values():
														_validate_dict(a__     ,    a__     ,    a__     ,    a__			)

def 			lowerCAmelCase__					(     a__:						Any     ,    a__:						str     ,    a__:						type     ,    a__:						bool = False			)    ->						None:


							'''simple docstring'''





							if not isinstance(_object     ,    a__			):
														_UpperCAmelCase											=     F'''{var_name} must be a dict'''
														raise ValueError(a__			)
							if not all(isinstance(a__     ,    a__			) for x in _object			):
														_UpperCAmelCase											=     F'''{var_name} all keys must be strings'''
														raise ValueError(a__			)
							if not all(isinstance(a__     ,    a__			) for x in _object.values()			):
														_UpperCAmelCase											=     'nested dictionary ' if nested else ''
														_UpperCAmelCase											=     F'''{var_name} {nested_text}all values must be {value_type.__name__}'''
														raise ValueError(a__			)


if __name__ == "__main__":
		from doctest import testmod

		testmod()
