





"""simple docstring"""






import gc
import random
import unittest

import numpy as np
import torch
from PIL import Image

from diffusers import (
    DDIMScheduler,
    KandinskyVaaInpaintPipeline,
    KandinskyVaaPriorPipeline,
    UNetaDConditionModel,
    VQModel,
)
from diffusers.utils import floats_tensor, load_image, load_numpy, slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu

from ..test_pipelines_common import PipelineTesterMixin, assert_mean_pixel_difference


enable_full_determinism()



class 	lowerCAmelCase					(      snake_case			,	unittest.TestCase							):
       lowerCAmelCase__						 =  KandinskyVaaInpaintPipeline
       lowerCAmelCase__						 =  ["""image_embeds""", """negative_image_embeds""", """image""", """mask_image"""]
       lowerCAmelCase__						 =  [
           """image_embeds""",
           """negative_image_embeds""",
           """image""",
           """mask_image""",
       ]
       lowerCAmelCase__						 =  [
           """generator""",
           """height""",
           """width""",
           """latents""",
           """guidance_scale""",
           """num_inference_steps""",
           """return_dict""",
           """guidance_scale""",
           """num_images_per_prompt""",
           """output_type""",
           """return_dict""",
       ]
       lowerCAmelCase__						 =  False






       @property
       def 							__A (   self     ):
              return 32






       @property
       def 							__A (   self     ):
              return 32






       @property
       def 							__A (   self     ):
              return self.time_input_dim






       @property
       def 							__A (   self     ):
              return self.time_input_dim * 4






       @property
       def 							__A (   self     ):
              return 1_00






       @property
       def 							__A (   self     ):
              torch.manual_seed(0     )

              _UpperCAmelCase     =     {
                  'in_channels': 9,
                  # Out channels is double in channels because predicts mean and variance
                  'out_channels': 8,
                  'addition_embed_type': 'image',
                  'down_block_types': ('ResnetDownsampleBlock2D', 'SimpleCrossAttnDownBlock2D'),
                  'up_block_types': ('SimpleCrossAttnUpBlock2D', 'ResnetUpsampleBlock2D'),
                  'mid_block_type': 'UNetMidBlock2DSimpleCrossAttn',
                  'block_out_channels': (self.block_out_channels_a, self.block_out_channels_a * 2),
                  'layers_per_block': 1,
                  'encoder_hid_dim': self.text_embedder_hidden_size,
                  'encoder_hid_dim_type': 'image_proj',
                  'cross_attention_dim': self.cross_attention_dim,
                  'attention_head_dim': 4,
                  'resnet_time_scale_shift': 'scale_shift',
                  'class_embed_type': None,
              }

              _UpperCAmelCase     =     UNetaDConditionModel(**a__     )
              return model






       @property
       def 							__A (   self     ):
              return {
                  "block_out_channels": [32, 64],
                  "down_block_types": ["DownEncoderBlock2D", "AttnDownEncoderBlock2D"],
                  "in_channels": 3,
                  "latent_channels": 4,
                  "layers_per_block": 1,
                  "norm_num_groups": 8,
                  "norm_type": "spatial",
                  "num_vq_embeddings": 12,
                  "out_channels": 3,
                  "up_block_types": [
                      "AttnUpDecoderBlock2D",
                      "UpDecoderBlock2D",
                  ],
                  "vq_embed_dim": 4,
              }






       @property
       def 							__A (   self     ):
              torch.manual_seed(0     )
              _UpperCAmelCase     =     VQModel(**self.dummy_movq_kwargs     )
              return model






       def 							__A (   self     ):
              _UpperCAmelCase     =     self.dummy_unet
              _UpperCAmelCase     =     self.dummy_movq

              _UpperCAmelCase     =     DDIMScheduler(
                  num_train_timesteps=10_00    ,  beta_schedule='linear'    ,  beta_start=0.00_085    ,  beta_end=0.012    ,  clip_sample=a__    ,  set_alpha_to_one=a__    ,  steps_offset=1    ,  prediction_type='epsilon'    ,  thresholding=a__    ,  )

              _UpperCAmelCase     =     {
                  'unet': unet,
                  'scheduler': scheduler,
                  'movq': movq,
              }

              return components






       def 							__A (   self    ,  a__    ,  a__=0     ):
              _UpperCAmelCase     =     floats_tensor((1, self.text_embedder_hidden_size)    ,  rng=random.Random(a__     )     ).to(a__     )
              _UpperCAmelCase     =     floats_tensor((1, self.text_embedder_hidden_size)    ,  rng=random.Random(seed + 1     )     ).to(
                  a__     )
              # create init_image
              _UpperCAmelCase     =     floats_tensor((1, 3, 64, 64)    ,  rng=random.Random(a__     )     ).to(a__     )
              _UpperCAmelCase     =     image.cpu().permute(0    ,  2    ,  3    ,  1     )[0]
              _UpperCAmelCase     =     Image.fromarray(np.uinta(a__     )     ).convert('RGB'     ).resize((2_56, 2_56)     )
              # create mask
              _UpperCAmelCase     =     np.ones((64, 64)    ,  dtype=np.floataa     )
              _UpperCAmelCase     =     0

              if str(a__     ).startswith('mps'     ):
                     _UpperCAmelCase     =     torch.manual_seed(a__     )
              else:
                     _UpperCAmelCase     =     torch.Generator(device=a__     ).manual_seed(a__     )
              _UpperCAmelCase     =     {
                  'image': init_image,
                  'mask_image': mask,
                  'image_embeds': image_embeds,
                  'negative_image_embeds': negative_image_embeds,
                  'generator': generator,
                  'height': 64,
                  'width': 64,
                  'num_inference_steps': 2,
                  'guidance_scale': 4.0,
                  'output_type': 'np',
              }
              return inputs






       def 							__A (   self     ):
              _UpperCAmelCase     =     'cpu'

              _UpperCAmelCase     =     self.get_dummy_components()

              _UpperCAmelCase     =     self.pipeline_class(**a__     )
              _UpperCAmelCase     =     pipe.to(a__     )

              pipe.set_progress_bar_config(disable=a__     )

              _UpperCAmelCase     =     pipe(**self.get_dummy_inputs(a__     )     )
              _UpperCAmelCase     =     output.images

              _UpperCAmelCase     =     pipe(
                  **self.get_dummy_inputs(a__     )    ,  return_dict=a__    ,  )[0]

              _UpperCAmelCase     =     image[0, -3:, -3:, -1]
              _UpperCAmelCase     =     image_from_tuple[0, -3:, -3:, -1]

              print(f"""image.shape {image.shape}"""     )

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

              _UpperCAmelCase     =     np.array(
                  [0.50_775_903, 0.49_527_195, 0.48_824_543, 0.50_192_237, 0.48_644_906, 0.49_373_814, 0.4_780_598, 0.47_234_827, 0.48_327_848]     )

              assert (
                  np.abs(image_slice.flatten() - expected_slice     ).max() < 1E-2
              ), f""" expected_slice {expected_slice}, but got {image_slice.flatten()}"""
              assert (
                  np.abs(image_from_tuple_slice.flatten() - expected_slice     ).max() < 1E-2
              ), f""" expected_slice {expected_slice}, but got {image_from_tuple_slice.flatten()}"""






       def 							__A (   self     ):
              super().test_inference_batch_single_identical(expected_max_diff=3E-3     )



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






       def 							__A (   self     ):
              # clean up the VRAM after each test
              super().tearDown()
              gc.collect()
              torch.cuda.empty_cache()






       def 							__A (   self     ):
              _UpperCAmelCase     =     load_numpy(
                  'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'
                  '/kandinskyv22/kandinskyv22_inpaint_cat_with_hat_fp16.npy'     )

              _UpperCAmelCase     =     load_image(
                  'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main' '/kandinsky/cat.png'     )
              _UpperCAmelCase     =     np.ones((7_68, 7_68)    ,  dtype=np.floataa     )
              _UpperCAmelCase     =     0

              _UpperCAmelCase     =     'a hat'

              _UpperCAmelCase     =     KandinskyVaaPriorPipeline.from_pretrained(
                  'kandinsky-community/kandinsky-2-2-prior'    ,  torch_dtype=torch.floataa     )
              pipe_prior.to(a__     )

              _UpperCAmelCase     =     KandinskyVaaInpaintPipeline.from_pretrained(
                  'kandinsky-community/kandinsky-2-2-decoder-inpaint'    ,  torch_dtype=torch.floataa     )
              _UpperCAmelCase     =     pipeline.to(a__     )
              pipeline.set_progress_bar_config(disable=a__     )

              _UpperCAmelCase     =     torch.Generator(device='cpu'     ).manual_seed(0     )
              _UpperCAmelCase    ,							_UpperCAmelCase     =     pipe_prior(
                  a__    ,  generator=a__    ,  num_inference_steps=5    ,  negative_prompt=''    ,  ).to_tuple()

              _UpperCAmelCase     =     pipeline(
                  image=a__    ,  mask_image=a__    ,  image_embeds=a__    ,  negative_image_embeds=a__    ,  generator=a__    ,  num_inference_steps=1_00    ,  height=7_68    ,  width=7_68    ,  output_type='np'    ,  )

              _UpperCAmelCase     =     output.images[0]

              assert image.shape == (7_68, 7_68, 3)

              assert_mean_pixel_difference(a__    ,  a__     )













"""simple docstring"""






from itertools import product

def       __lowerCamelCase    (       SCREAMING_SNAKE_CASE,SCREAMING_SNAKE_CASE							)		->       list[int]:






       """simple docstring"""




       _UpperCAmelCase     =     sides_number
       _UpperCAmelCase     =     max_face_number * dice_number
       _UpperCAmelCase     =     [0] * (max_total + 1)

       _UpperCAmelCase     =     1
       _UpperCAmelCase     =     range(SCREAMING_SNAKE_CASE,max_face_number + 1							)
       for dice_numbers in product(SCREAMING_SNAKE_CASE,repeat=SCREAMING_SNAKE_CASE							):
              _UpperCAmelCase     =     sum(SCREAMING_SNAKE_CASE							)
              totals_frequencies[total] += 1

       return totals_frequencies

def       __lowerCamelCase    (       )		->       float:






       """simple docstring"""




       _UpperCAmelCase     =     total_frequency_distribution(
           sides_number=4,dice_number=9							)
       _UpperCAmelCase     =     total_frequency_distribution(
           sides_number=6,dice_number=6							)

       _UpperCAmelCase     =     0
       _UpperCAmelCase     =     9
       _UpperCAmelCase     =     4 * 9
       _UpperCAmelCase     =     6
       for peter_total in range(SCREAMING_SNAKE_CASE,max_peter_total + 1							):
              peter_wins_count += peter_totals_frequencies[peter_total] * sum(
                  colin_totals_frequencies[min_colin_total:peter_total]							)

       _UpperCAmelCase     =     (4**9) * (6**6)
       _UpperCAmelCase     =     peter_wins_count / total_games_number

       _UpperCAmelCase     =     round(SCREAMING_SNAKE_CASE,ndigits=7							)

       return rounded_peter_win_probability


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






