




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        lowerCamelCase__(   __lowerCamelCase					,	__lowerCamelCase					,	unittest.TestCase):
      UpperCAmelCase__ :      str       						=  StableDiffusionParadigmsPipeline
      UpperCAmelCase__ :      Optional[int]       						=  TEXT_TO_IMAGE_PARAMS
      UpperCAmelCase__ :      int       						=  TEXT_TO_IMAGE_BATCH_PARAMS
      UpperCAmelCase__ :      Optional[Any]       						=  TEXT_TO_IMAGE_IMAGE_PARAMS
      UpperCAmelCase__ :      List[str]       						=  TEXT_TO_IMAGE_IMAGE_PARAMS
      def      lowerCAmelCase__		(							self:					Optional[int]  ):
             torch.manual_seed(0  )
             __lowerCamelCase  =		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=UpperCamelCase_       ,			)
             __lowerCamelCase  =		DDIMParallelScheduler(
                 beta_start=0.0_0085       ,			beta_end=0.012       ,			beta_schedule="""scaled_linear"""       ,			clip_sample=UpperCamelCase_       ,			set_alpha_to_one=UpperCamelCase_       ,			)
             torch.manual_seed(0  )
             __lowerCamelCase  =		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=1_28       ,			)
             torch.manual_seed(0  )
             __lowerCamelCase  =		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=10_00       ,			hidden_act="""gelu"""       ,			projection_dim=5_12       ,			)
             __lowerCamelCase  =		CLIPTextModel(UpperCamelCase_  )
             __lowerCamelCase  =		CLIPTokenizer.from_pretrained("""hf-internal-testing/tiny-random-clip"""  )

             __lowerCamelCase  =		{
                 """unet""": unet,
                 """scheduler""": scheduler,
                 """vae""": vae,
                 """text_encoder""": text_encoder,
                 """tokenizer""": tokenizer,
                 """safety_checker""": None,
                 """feature_extractor""": None,
             }
             return components
      def      lowerCAmelCase__		(							self:					List[str]       ,			UpperCamelCase_:					int       ,			UpperCamelCase_:					Optional[Any]=0  ):
             if str(UpperCamelCase_  ).startswith("""mps"""  ):
                    __lowerCamelCase  =		torch.manual_seed(UpperCamelCase_  )
             else:
                    __lowerCamelCase  =		torch.Generator(device=UpperCamelCase_  ).manual_seed(UpperCamelCase_  )
             __lowerCamelCase  =		{
                 """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:					List[str]  ):
             __lowerCamelCase  =		"""cpu"""  # ensure determinism for the device-dependent torch.Generator
             __lowerCamelCase  =		self.get_dummy_components()
             __lowerCamelCase  =		StableDiffusionParadigmsPipeline(**UpperCamelCase_  )
             __lowerCamelCase  =		sd_pipe.to(UpperCamelCase_  )
             sd_pipe.set_progress_bar_config(disable=UpperCamelCase_  )

             __lowerCamelCase  =		self.get_dummy_inputs(UpperCamelCase_  )
             __lowerCamelCase  =		sd_pipe(**UpperCamelCase_  ).images
             __lowerCamelCase  =		image[0, -3:, -3:, -1]
             assert image.shape == (1, 64, 64, 3)

             __lowerCamelCase  =		np.array([0.4773, 0.5417, 0.4723, 0.4925, 0.5631, 0.4752, 0.5240, 0.4935, 0.5023]  )

             assert np.abs(image_slice.flatten() - expected_slice  ).max() < 1E-2
      def      lowerCAmelCase__		(							self:					Optional[int]  ):
             __lowerCamelCase  =		"""cpu"""  # ensure determinism for the device-dependent torch.Generator
             __lowerCamelCase  =		self.get_dummy_components()
             torch.manual_seed(0  )
             __lowerCamelCase  =		DDPMParallelScheduler()
             torch.manual_seed(0  )
             __lowerCamelCase  =		StableDiffusionParadigmsPipeline(**UpperCamelCase_  )
             __lowerCamelCase  =		sd_pipe.to(UpperCamelCase_  )
             sd_pipe.set_progress_bar_config(disable=UpperCamelCase_  )

             __lowerCamelCase  =		self.get_dummy_inputs(UpperCamelCase_  )
             __lowerCamelCase  =		sd_pipe(**UpperCamelCase_  ).images
             __lowerCamelCase  =		image[0, -3:, -3:, -1]
             assert image.shape == (1, 64, 64, 3)

             __lowerCamelCase  =		np.array([0.3573, 0.4420, 0.4960, 0.4799, 0.3796, 0.3879, 0.4819, 0.4365, 0.4468]  )

             assert np.abs(image_slice.flatten() - expected_slice  ).max() < 1E-2
      def      lowerCAmelCase__		(							self:					List[Any]  ):
             super().test_inference_batch_consistent(batch_sizes=[1, 2]  )
      def      lowerCAmelCase__		(							self:					List[Any]  ):
             super().test_inference_batch_single_identical(batch_size=2       ,			expected_max_diff=3E-3  )




      def      lowerCAmelCase__		(							self:					Optional[Any]  ):
             __lowerCamelCase  =		"""cpu"""  # ensure determinism for the device-dependent torch.Generator
             __lowerCamelCase  =		self.get_dummy_components()
             __lowerCamelCase  =		StableDiffusionParadigmsPipeline(**UpperCamelCase_  )
             __lowerCamelCase  =		sd_pipe.to(UpperCamelCase_  )
             sd_pipe.set_progress_bar_config(disable=UpperCamelCase_  )

             __lowerCamelCase  =		self.get_dummy_inputs(UpperCamelCase_  )
             __lowerCamelCase  =		"""french fries"""
             __lowerCamelCase  =		sd_pipe(**UpperCamelCase_       ,			negative_prompt=UpperCamelCase_  )
             __lowerCamelCase  =		output.images
             __lowerCamelCase  =		image[0, -3:, -3:, -1]

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

             __lowerCamelCase  =		np.array([0.4771, 0.5420, 0.4683, 0.4918, 0.5636, 0.4725, 0.5230, 0.4923, 0.5015]  )

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





@slow
@require_torch_gpu
class        lowerCamelCase__(   unittest.TestCase):
      def      lowerCAmelCase__		(							self:					Optional[int]  ):
             super().tearDown()
             gc.collect()
             torch.cuda.empty_cache()
      def      lowerCAmelCase__		(							self:					List[Any]       ,			UpperCamelCase_:					str=0  ):
             __lowerCamelCase  =		torch.Generator(device=UpperCamelCase_  ).manual_seed(UpperCamelCase_  )
             __lowerCamelCase  =		{
                 """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  ):
             __lowerCamelCase  =		"""stabilityai/stable-diffusion-2-base"""
             __lowerCamelCase  =		DDIMParallelScheduler.from_pretrained(UpperCamelCase_       ,			subfolder="""scheduler"""  )
             __lowerCamelCase  =		StableDiffusionParadigmsPipeline.from_pretrained(UpperCamelCase_       ,			scheduler=UpperCamelCase_       ,			safety_checker=UpperCamelCase_  )
             pipe.to(UpperCamelCase_  )
             pipe.set_progress_bar_config(disable=UpperCamelCase_  )
             pipe.enable_attention_slicing()

             __lowerCamelCase  =		self.get_inputs()
             __lowerCamelCase  =		pipe(**UpperCamelCase_  ).images
             __lowerCamelCase  =		image[0, -3:, -3:, -1].flatten()

             assert image.shape == (1, 5_12, 5_12, 3)

             __lowerCamelCase  =		np.array([0.9622, 0.9602, 0.9748, 0.9591, 0.9630, 0.9691, 0.9661, 0.9631, 0.9741]  )

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







from functools import partial
from typing import Any, Callable, Dict, List, Type, TypeVar, Union, overload

import torch
import torch.nn as nn
import torch.types




def    lowerCamelCase__			(		A__			:   torch.Tensor    ,     A__			:   torch.Tensor    ,     A__			:   bool      ):





       '''simple docstring'''



       if not inplace:
              __lowerCamelCase  =		ma + ma
       else:
              ma += ma

       return ma




def    lowerCamelCase__			(		A__			:   torch.Tensor    ,     A__			:   List[int]      ):





       '''simple docstring'''



       __lowerCamelCase  =		-1 * len(A__      )
       __lowerCamelCase  =		list(range(len(tensor.shape[:zero_index]      )      )      )
       return tensor.permute(first_inds + [zero_index + i for i in inds]      )




def    lowerCamelCase__			(		A__			:   torch.Tensor    ,     A__			:   int      ):





       '''simple docstring'''



       return t.reshape(t.shape[:-no_dims] + (-1,)      )




def    lowerCamelCase__			(		A__			:   torch.Tensor    ,     A__			:   torch.Tensor    ,     A__			:   int    ,     A__			:   float = 1E-4      ):





       '''simple docstring'''



       __lowerCamelCase  =		mask.expand(*value.shape      )
       return torch.sum(mask * value    ,     dim=A__      ) / (eps + torch.sum(A__    ,     dim=A__      ))




def    lowerCamelCase__			(		A__			:   torch.Tensor    ,     A__			:   torch.types.Number = 2.3_125    ,     A__			:   torch.types.Number = 21.6_875    ,     A__			:   int = 64      ):





       '''simple docstring'''



       __lowerCamelCase  =		torch.linspace(A__    ,     A__    ,     no_bins - 1    ,     device=pts.device      )
       __lowerCamelCase  =		torch.sqrt(torch.sum((pts.unsqueeze(-2      ) - pts.unsqueeze(-3      )) ** 2    ,     dim=-1      )      )
       return torch.bucketize(A__    ,     A__      )




def    lowerCamelCase__			(		A__			:   Callable[[list], Any]    ,     A__			:   List[dict]      ):





       '''simple docstring'''



       __lowerCamelCase  =		dicts[0]
       __lowerCamelCase  =		{}
       for k, v in first.items():
              __lowerCamelCase  =		[d[k] for d in dicts]
              if isinstance(A__    ,     A__      ):
                     __lowerCamelCase  =		dict_multimap(A__    ,     A__      )
              else:
                     __lowerCamelCase  =		fn(A__      )

       return new_dict




def    lowerCamelCase__			(		A__			:   torch.Tensor    ,     A__			:   torch.Tensor      ):





       '''simple docstring'''



       __lowerCamelCase  =		v_bins.view(((1,) * len(x.shape      )) + (len(A__      ),)      )
       __lowerCamelCase  =		x[..., None] - reshaped_bins
       __lowerCamelCase  =		torch.argmin(torch.abs(A__      )    ,     dim=-1      )
       return nn.functional.one_hot(A__    ,     num_classes=len(A__      )      ).float()




def    lowerCamelCase__			(		A__			:   torch.Tensor    ,     A__			:   torch.Tensor    ,     A__			:   int = 0    ,     A__			:   int = 0      ):





       '''simple docstring'''



       __lowerCamelCase  =		[]
       for i, s in enumerate(data.shape[:no_batch_dims]      ):
              __lowerCamelCase  =		torch.arange(A__      )
              __lowerCamelCase  =		r.view(*(*((1,) * i), -1, *((1,) * (len(inds.shape      ) - i - 1)))      )
              ranges.append(A__      )

       __lowerCamelCase  =		[slice(A__      ) for _ in range(len(data.shape      ) - no_batch_dims      )]
       __lowerCamelCase  =		inds
       ranges.extend(A__      )
       # Matt note: Editing this to get around the behaviour of using a list as an array index changing
       # in recent Numpy versions
       return data[tuple(A__      )]


UpperCAmelCase_   							=      TypeVar('T')




def    lowerCamelCase__			(		A__			:   Callable[[T], Any]    ,     A__			:   Dict[Any, Union[dict, list, tuple, T]]    ,     A__			:   Type[T]      ):





       '''simple docstring'''



       __lowerCamelCase  =		{}
       for k, v in dic.items():
              if isinstance(A__    ,     A__      ):
                     __lowerCamelCase  =		dict_map(A__    ,     A__    ,     A__      )
              else:
                     __lowerCamelCase  =		tree_map(A__    ,     A__    ,     A__      )

       return new_dict




@overload
def    lowerCamelCase__			(		A__			:   Callable[[T], Any]    ,     A__			:   T    ,     A__			:   Type[T]      ):





       '''simple docstring'''



       ...




@overload
def    lowerCamelCase__			(		A__			:   Callable[[T], Any]    ,     A__			:   dict    ,     A__			:   Type[T]      ):





       '''simple docstring'''



       ...




@overload
def    lowerCamelCase__			(		A__			:   Callable[[T], Any]    ,     A__			:   list    ,     A__			:   Type[T]      ):





       '''simple docstring'''



       ...




@overload
def    lowerCamelCase__			(		A__			:   Callable[[T], Any]    ,     A__			:   tuple    ,     A__			:   Type[T]      ):





       '''simple docstring'''



       ...




def    lowerCamelCase__			(		A__			:   Dict    ,     A__			:   List[str]    ,     A__			:   List[str]      ):





       '''simple docstring'''



       if isinstance(A__    ,     A__      ):
              return dict_map(A__    ,     A__    ,     A__      )
       elif isinstance(A__    ,     A__      ):
              return [tree_map(A__    ,     A__    ,     A__      ) for x in tree]
       elif isinstance(A__    ,     A__      ):
              return tuple(tree_map(A__    ,     A__    ,     A__      ) for x in tree      )
       elif isinstance(A__    ,     A__      ):
              return fn(A__      )
       else:
              print(type(A__      )      )
              raise ValueError("""Not supported"""      )




UpperCAmelCase_   							=      partial(tree_map, leaf_type=torch.Tensor)

