
import gc
import unittest

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

from diffusers import (
    AutoencoderKL,
    DDIMParallelScheduler,
    DDPMParallelScheduler,
    StableDiffusionParadigmsPipeline,
    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	(   __lowerCamelCase						, __lowerCamelCase						, unittest.TestCase					):
					a__: str			  =				StableDiffusionParadigmsPipeline
					a__: Optional[int]			  =				TEXT_TO_IMAGE_PARAMS
					a__: Dict			  =				TEXT_TO_IMAGE_BATCH_PARAMS
					a__: Union[str, Any]			  =				TEXT_TO_IMAGE_IMAGE_PARAMS
					a__: Tuple			  =				TEXT_TO_IMAGE_IMAGE_PARAMS

					def        _lowerCAmelCase	(			self      :					Union[str, Any]					):
											torch.manual_seed(0					)
											lowercase       :    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 ,      attention_head_dim=(2, 4) ,      use_linear_projection=lowerCAmelCase ,      )
											lowercase       :    Optional[int]											=  DDIMParallelScheduler(
											    beta_start=0.0_00_85 ,      beta_end=0.0_12 ,      beta_schedule='''scaled_linear''' ,      clip_sample=lowerCAmelCase ,      set_alpha_to_one=lowerCAmelCase ,      )
											torch.manual_seed(0					)
											lowercase       :    Optional[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 ,      sample_size=128 ,      )
											torch.manual_seed(0					)
											lowercase       :    int											=  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 ,      hidden_act='''gelu''' ,      projection_dim=512 ,      )
											lowercase       :    List[Any]											=  CLIPTextModel(lowerCAmelCase					)
											lowercase       :    Tuple											=  CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''					)

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

					def        _lowerCAmelCase	(			self      :					Union[str, Any] ,      lowerCAmelCase      :					List[Any] ,      lowerCAmelCase      :					List[Any]=0					):
											if str(lowerCAmelCase					).startswith('''mps'''					):
																	lowercase       :    Tuple											=  torch.manual_seed(lowerCAmelCase					)
											else:
																	lowercase       :    List[str]											=  torch.Generator(device=lowerCAmelCase					).manual_seed(lowerCAmelCase					)
											lowercase       :    Dict											=  {
											    '''prompt''': '''a photograph of an astronaut riding a horse''',
											    '''generator''': generator,
											    '''num_inference_steps''': 10,
											    '''guidance_scale''': 6.0,
											    '''output_type''': '''numpy''',
											    '''parallel''': 3,
											    '''debug''': True,
											}
											return inputs

					def        _lowerCAmelCase	(			self      :					Any					):
											lowercase       :    str											=  '''cpu'''  # ensure determinism for the device-dependent torch.Generator
											lowercase       :    int											=  self.get_dummy_components()
											lowercase       :    Tuple											=  StableDiffusionParadigmsPipeline(**lowerCAmelCase					)
											lowercase       :    List[Any]											=  sd_pipe.to(lowerCAmelCase					)
											sd_pipe.set_progress_bar_config(disable=lowerCAmelCase					)

											lowercase       :    Optional[int]											=  self.get_dummy_inputs(lowerCAmelCase					)
											lowercase       :    int											=  sd_pipe(**lowerCAmelCase					).images
											lowercase       :    Dict											=  image[0, -3:, -3:, -1]
											assert image.shape == (1, 64, 64, 3)

											lowercase       :    Dict											=  np.array([0.47_73, 0.54_17, 0.47_23, 0.49_25, 0.56_31, 0.47_52, 0.52_40, 0.49_35, 0.50_23]					)

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

					def        _lowerCAmelCase	(			self      :					List[Any]					):
											lowercase       :    Tuple											=  '''cpu'''  # ensure determinism for the device-dependent torch.Generator
											lowercase       :    List[str]											=  self.get_dummy_components()
											torch.manual_seed(0					)
											lowercase       :    List[Any]											=  DDPMParallelScheduler()
											torch.manual_seed(0					)
											lowercase       :    Tuple											=  StableDiffusionParadigmsPipeline(**lowerCAmelCase					)
											lowercase       :    Dict											=  sd_pipe.to(lowerCAmelCase					)
											sd_pipe.set_progress_bar_config(disable=lowerCAmelCase					)

											lowercase       :    str											=  self.get_dummy_inputs(lowerCAmelCase					)
											lowercase       :    Tuple											=  sd_pipe(**lowerCAmelCase					).images
											lowercase       :    Optional[int]											=  image[0, -3:, -3:, -1]
											assert image.shape == (1, 64, 64, 3)

											lowercase       :    str											=  np.array([0.35_73, 0.44_20, 0.49_60, 0.47_99, 0.37_96, 0.38_79, 0.48_19, 0.43_65, 0.44_68]					)

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

					def        _lowerCAmelCase	(			self      :					Optional[int]					):
											super().test_inference_batch_consistent(batch_sizes=[1, 2]					)

					def        _lowerCAmelCase	(			self      :					int					):
											super().test_inference_batch_single_identical(batch_size=2 ,      expected_max_diff=3E-3					)




					def        _lowerCAmelCase	(			self      :					Tuple					):
											lowercase       :    List[Any]											=  '''cpu'''  # ensure determinism for the device-dependent torch.Generator
											lowercase       :    Optional[Any]											=  self.get_dummy_components()
											lowercase       :    str											=  StableDiffusionParadigmsPipeline(**lowerCAmelCase					)
											lowercase       :    List[Any]											=  sd_pipe.to(lowerCAmelCase					)
											sd_pipe.set_progress_bar_config(disable=lowerCAmelCase					)

											lowercase       :    Union[str, Any]											=  self.get_dummy_inputs(lowerCAmelCase					)
											lowercase       :    Union[str, Any]											=  '''french fries'''
											lowercase       :    Tuple											=  sd_pipe(**lowerCAmelCase ,      negative_prompt=lowerCAmelCase					)
											lowercase       :    Union[str, Any]											=  output.images
											lowercase       :    List[str]											=  image[0, -3:, -3:, -1]

											assert image.shape == (1, 64, 64, 3)

											lowercase       :    Optional[int]											=  np.array([0.47_71, 0.54_20, 0.46_83, 0.49_18, 0.56_36, 0.47_25, 0.52_30, 0.49_23, 0.50_15]					)

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







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

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

					def        _lowerCAmelCase	(			self      :					Dict ,      lowerCAmelCase      :					Dict=0					):
											lowercase       :    Optional[Any]											=  torch.Generator(device=lowerCAmelCase					).manual_seed(lowerCAmelCase					)
											lowercase       :    List[Any]											=  {
											    '''prompt''': '''a photograph of an astronaut riding a horse''',
											    '''generator''': generator,
											    '''num_inference_steps''': 10,
											    '''guidance_scale''': 7.5,
											    '''output_type''': '''numpy''',
											    '''parallel''': 3,
											    '''debug''': True,
											}
											return inputs




					def        _lowerCAmelCase	(			self      :					Tuple					):
											lowercase       :    Tuple											=  '''stabilityai/stable-diffusion-2-base'''
											lowercase       :    Optional[int]											=  DDIMParallelScheduler.from_pretrained(lowerCAmelCase ,      subfolder='''scheduler'''					)
											lowercase       :    Optional[Any]											=  StableDiffusionParadigmsPipeline.from_pretrained(lowerCAmelCase ,      scheduler=lowerCAmelCase ,      safety_checker=lowerCAmelCase					)
											pipe.to(lowerCAmelCase					)
											pipe.set_progress_bar_config(disable=lowerCAmelCase					)
											pipe.enable_attention_slicing()

											lowercase       :    List[str]											=  self.get_inputs()
											lowercase       :    int											=  pipe(**lowerCAmelCase					).images
											lowercase       :    int											=  image[0, -3:, -3:, -1].flatten()

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

											lowercase       :    Optional[int]											=  np.array([0.96_22, 0.96_02, 0.97_48, 0.95_91, 0.96_30, 0.96_91, 0.96_61, 0.96_31, 0.97_41]					)

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




from itertools import combinations_with_replacement

def        lowerCamelCase_      (			UpperCAmelCase_							:				int = 1_00       ):
						lowercase       :    list[int]											=  [*list(range(1							,			21       )       ), 25]
						lowercase       :    list[int]											=  [2 * x for x in range(1							,			21       )] + [50]
						lowercase       :    list[int]											=  [3 * x for x in range(1							,			21       )]
						lowercase       :    list[int]											=  singles + doubles + triples + [0]

						lowercase       :    int											=  0
						lowercase       :    int
						lowercase       :    int
						lowercase       :    int
						lowercase       :    int

						for double in doubles:
												for throwa, throwa in combinations_with_replacement(UpperCAmelCase_							,			2       ):
																		lowercase       :    Union[str, Any]											=  double + throwa + throwa
																		if checkout_total < limit:
																								num_checkouts += 1

						return num_checkouts


if __name__ == "__main__":
				print(F'{solution() = }')


