




'''simple docstring'''
import tempfile
import unittest

import numpy as np
import torch

from diffusers import DiffusionPipeline
from diffusers.models.attention_processor import Attention, AttnAddedKVProcessor






class 			SCREAMING_SNAKE_CASE(							unittest.TestCase	):

       def        snake_case__     (       self   ,						lowerCamelCase__ = False	)       ->   Optional[Any]:

           """simple docstring"""



           __lowercase     							=		10

           if only_cross_attention:
               __lowercase     							=		12
           else:
               # when only cross attention is not set, the cross attention dim must be the same as the query dim
               __lowercase     							=		query_dim

           return {
               "query_dim": query_dim,
               "cross_attention_dim": cross_attention_dim,
               "heads": 2,
               "dim_head": 4,
               "added_kv_proj_dim": 6,
               "norm_num_groups": 1,
               "only_cross_attention": only_cross_attention,
               "processor": AttnAddedKVProcessor(),
           }

       def        snake_case__     (       self   ,						lowerCamelCase__   ,						lowerCamelCase__	)       ->   List[str]:

           """simple docstring"""



           __lowercase     							=		2

           __lowercase     							=		torch.rand(SCREAMING_SNAKE_CASE_   ,						SCREAMING_SNAKE_CASE_   ,						3   ,						2	)
           __lowercase     							=		torch.rand(SCREAMING_SNAKE_CASE_   ,						4   ,						SCREAMING_SNAKE_CASE_	)
           __lowercase     							=		None

           return {
               "hidden_states": hidden_states,
               "encoder_hidden_states": encoder_hidden_states,
               "attention_mask": attention_mask,
           }




       def        snake_case__     (       self	)       ->   Optional[Any]:

           """simple docstring"""



           torch.manual_seed(0	)

           __lowercase     							=		self.get_constructor_arguments(only_cross_attention=SCREAMING_SNAKE_CASE_	)
           __lowercase     							=		Attention(**SCREAMING_SNAKE_CASE_	)

           self.assertTrue(attn.to_k is not None	)
           self.assertTrue(attn.to_v is not None	)

           __lowercase     							=		self.get_forward_arguments(
               query_dim=constructor_args["""query_dim"""]   ,						added_kv_proj_dim=constructor_args["""added_kv_proj_dim"""]	)

           __lowercase     							=		attn(**SCREAMING_SNAKE_CASE_	)

           # only self attention

           torch.manual_seed(0	)

           __lowercase     							=		self.get_constructor_arguments(only_cross_attention=SCREAMING_SNAKE_CASE_	)
           __lowercase     							=		Attention(**SCREAMING_SNAKE_CASE_	)

           self.assertTrue(attn.to_k is None	)
           self.assertTrue(attn.to_v is None	)

           __lowercase     							=		self.get_forward_arguments(
               query_dim=constructor_args["""query_dim"""]   ,						added_kv_proj_dim=constructor_args["""added_kv_proj_dim"""]	)

           __lowercase     							=		attn(**SCREAMING_SNAKE_CASE_	)

           self.assertTrue((only_cross_attn_out != self_and_cross_attn_out).all()	)






class 			SCREAMING_SNAKE_CASE(							unittest.TestCase	):

       def        snake_case__     (       self	)       ->   Optional[int]:

           """simple docstring"""



           __lowercase     							=		DiffusionPipeline.from_pretrained("""hf-internal-testing/tiny-stable-diffusion-pipe"""   ,						safety_checker=SCREAMING_SNAKE_CASE_	)

           __lowercase     							=		pipe(
               """foo"""   ,						num_inference_steps=2   ,						generator=torch.Generator("""cpu"""	).manual_seed(0	)   ,						output_type="""np"""   ,						).images

           # the initial conversion succeeds
           __lowercase     							=		DiffusionPipeline.from_pretrained(
               """hf-internal-testing/tiny-stable-diffusion-pipe"""   ,						device_map="""sequential"""   ,						safety_checker=SCREAMING_SNAKE_CASE_	)

           __lowercase     							=		pipe(
               """foo"""   ,						num_inference_steps=2   ,						generator=torch.Generator("""cpu"""	).manual_seed(0	)   ,						output_type="""np"""   ,						).images

           with tempfile.TemporaryDirectory() as tmpdir:
               # save the converted model
               pipe.save_pretrained(SCREAMING_SNAKE_CASE_	)

               # can also load the converted weights
               __lowercase     							=		DiffusionPipeline.from_pretrained(SCREAMING_SNAKE_CASE_   ,						device_map="""sequential"""   ,						safety_checker=SCREAMING_SNAKE_CASE_	)

           __lowercase     							=		pipe(
               """foo"""   ,						num_inference_steps=2   ,						generator=torch.Generator("""cpu"""	).manual_seed(0	)   ,						output_type="""np"""   ,						).images

           self.assertTrue(np.allclose(SCREAMING_SNAKE_CASE_   ,						SCREAMING_SNAKE_CASE_   ,						atol=1E-5	)	)
           self.assertTrue(np.allclose(SCREAMING_SNAKE_CASE_   ,						SCREAMING_SNAKE_CASE_   ,						atol=1E-5	)	)










'''simple docstring'''


from __future__ import annotations

import random






class 						lowerCAmelCase_	:







  '''simple docstring'''

  UpperCamelCase_			:      Optional[int]										=				6



  def __init__( self				:							Dict	)			->		Dict:







       '''simple docstring'''



       A:		Union[str, Any]          =		list(range(1					,  Dice.NUM_SIDES + 1	)	)




  def 					_snake_case     ( self				:							Optional[Any]	)			->		Any:







       '''simple docstring'''



       return random.choice(self.sides	)






def 					SCREAMING_SNAKE_CASE(   __lowercase		,				__lowercase = 2	)			-> list[float]:
     A:		Union[str, Any]          =		[Dice() for i in range(__lowercase	)]
     A:		Union[str, Any]          =		[0] * (len(__lowercase	) * Dice.NUM_SIDES + 1)
     for _ in range(__lowercase	):
          count_of_sum[sum(dice.roll() for dice in dices	)] += 1
     A:		Dict          =		[round((count * 1_0_0) / num_throws		,				2	) for count in count_of_sum]
     return probability[num_dice:]  # remove probability of sums that never appear


if __name__ == "__main__":
     import doctest

     doctest.testmod()






