



"""simple docstring"""







import gc
import random
import unittest

import numpy as np
import torch
from transformers import XLMRobertaTokenizerFast

from diffusers import DDIMScheduler, KandinskyPipeline, KandinskyPriorPipeline, UNetaDConditionModel, VQModel
from diffusers.pipelines.kandinsky.text_encoder import MCLIPConfig, MultilingualCLIP
from diffusers.utils import floats_tensor, 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  UpperCamelCase__  (      lowercase_,   unittest.TestCase    ):




      """simple docstring"""




      _SCREAMING_SNAKE_CASE    							=	KandinskyPipeline
      _SCREAMING_SNAKE_CASE    							=	[
          """prompt""",
          """image_embeds""",
          """negative_image_embeds""",
      ]
      _SCREAMING_SNAKE_CASE    							=	["""prompt""", """negative_prompt""", """image_embeds""", """negative_image_embeds"""]
      _SCREAMING_SNAKE_CASE    							=	[
          """generator""",
          """height""",
          """width""",
          """latents""",
          """guidance_scale""",
          """negative_prompt""",
          """num_inference_steps""",
          """return_dict""",
          """guidance_scale""",
          """num_images_per_prompt""",
          """output_type""",
          """return_dict""",
      ]
      _SCREAMING_SNAKE_CASE    							=	False



      @property
      def 	SCREAMING_SNAKE_CASE__					( self							: List[Any]			):
        return 3_2



      @property
      def 	SCREAMING_SNAKE_CASE__					( self							: Tuple			):
        return 3_2



      @property
      def 	SCREAMING_SNAKE_CASE__					( self							: List[Any]			):
        return self.time_input_dim



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



      @property
      def 	SCREAMING_SNAKE_CASE__					( self							: Dict			):
        return 1_0_0



      @property
      def 	SCREAMING_SNAKE_CASE__					( self							: int			):
        lowerCAmelCase_  :							Any								=					XLMRobertaTokenizerFast.from_pretrained('YiYiXu/tiny-random-mclip-base'			)
        return tokenizer



      @property
      def 	SCREAMING_SNAKE_CASE__					( self							: Any			):
        torch.manual_seed(0			)
        lowerCAmelCase_  :							Optional[Any]								=					MCLIPConfig(
            numDims=self.cross_attention_dim       ,     transformerDimensions=self.text_embedder_hidden_size       ,     hidden_size=self.text_embedder_hidden_size       ,     intermediate_size=3_7       ,     num_attention_heads=4       ,     num_hidden_layers=5       ,     vocab_size=1_0_0_5       ,     )

        lowerCAmelCase_  :							Dict								=					MultilingualCLIP(SCREAMING_SNAKE_CASE_			)
        lowerCAmelCase_  :							Union[str, Any]								=					text_encoder.eval()

        return text_encoder



      @property
      def 	SCREAMING_SNAKE_CASE__					( self							: Tuple			):
        torch.manual_seed(0			)

        lowerCAmelCase_  :							Union[str, Any]								=					{
            'in_channels': 4,
            # Out channels is double in channels because predicts mean and variance
            'out_channels': 8,
            'addition_embed_type': 'text_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': 'text_image_proj',
            'cross_attention_dim': self.cross_attention_dim,
            'attention_head_dim': 4,
            'resnet_time_scale_shift': 'scale_shift',
            'class_embed_type': None,
        }

        lowerCAmelCase_  :							Optional[Any]								=					UNetaDConditionModel(**SCREAMING_SNAKE_CASE_			)
        return model



      @property
      def 	SCREAMING_SNAKE_CASE__					( self							: str			):
        return {
            "block_out_channels": [3_2, 6_4],
            "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": 1_2,
            "out_channels": 3,
            "up_block_types": [
                "AttnUpDecoderBlock2D",
                "UpDecoderBlock2D",
            ],
            "vq_embed_dim": 4,
        }



      @property
      def 	SCREAMING_SNAKE_CASE__					( self							: int			):
        torch.manual_seed(0			)
        lowerCAmelCase_  :							Dict								=					VQModel(**self.dummy_movq_kwargs			)
        return model



      def 	SCREAMING_SNAKE_CASE__					( self							: Dict			):
        lowerCAmelCase_  :							str								=					self.dummy_text_encoder
        lowerCAmelCase_  :							List[Any]								=					self.dummy_tokenizer
        lowerCAmelCase_  :							Any								=					self.dummy_unet
        lowerCAmelCase_  :							int								=					self.dummy_movq

        lowerCAmelCase_  :							int								=					DDIMScheduler(
            num_train_timesteps=1_0_0_0       ,     beta_schedule='linear'       ,     beta_start=0.0_00_85       ,     beta_end=0.0_12       ,     clip_sample=SCREAMING_SNAKE_CASE_       ,     set_alpha_to_one=SCREAMING_SNAKE_CASE_       ,     steps_offset=1       ,     prediction_type='epsilon'       ,     thresholding=SCREAMING_SNAKE_CASE_       ,     )

        lowerCAmelCase_  :							List[str]								=					{
            'text_encoder': text_encoder,
            'tokenizer': tokenizer,
            'unet': unet,
            'scheduler': scheduler,
            'movq': movq,
        }
        return components



      def 	SCREAMING_SNAKE_CASE__					( self							: Any       ,     SCREAMING_SNAKE_CASE_							: Dict       ,     SCREAMING_SNAKE_CASE_							: List[Any]=0			):
        lowerCAmelCase_  :							Any								=					floats_tensor((1, self.cross_attention_dim)       ,     rng=random.Random(SCREAMING_SNAKE_CASE_			)			).to(SCREAMING_SNAKE_CASE_			)
        lowerCAmelCase_  :							List[str]								=					floats_tensor((1, self.cross_attention_dim)       ,     rng=random.Random(seed + 1			)			).to(SCREAMING_SNAKE_CASE_			)
        if str(SCREAMING_SNAKE_CASE_			).startswith('mps'			):
          lowerCAmelCase_  :							Dict								=					torch.manual_seed(SCREAMING_SNAKE_CASE_			)
        else:
          lowerCAmelCase_  :							Tuple								=					torch.Generator(device=SCREAMING_SNAKE_CASE_			).manual_seed(SCREAMING_SNAKE_CASE_			)
        lowerCAmelCase_  :							Union[str, Any]								=					{
            'prompt': 'horse',
            'image_embeds': image_embeds,
            'negative_image_embeds': negative_image_embeds,
            'generator': generator,
            'height': 6_4,
            'width': 6_4,
            'guidance_scale': 4.0,
            'num_inference_steps': 2,
            'output_type': 'np',
        }
        return inputs





      def 	SCREAMING_SNAKE_CASE__					( self							: Optional[Any]			):
        lowerCAmelCase_  :							List[Any]								=					'cpu'

        lowerCAmelCase_  :							Optional[Any]								=					self.get_dummy_components()

        lowerCAmelCase_  :							str								=					self.pipeline_class(**SCREAMING_SNAKE_CASE_			)
        lowerCAmelCase_  :							Union[str, Any]								=					pipe.to(SCREAMING_SNAKE_CASE_			)

        pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE_			)

        lowerCAmelCase_  :							Dict								=					pipe(**self.get_dummy_inputs(SCREAMING_SNAKE_CASE_			)			)
        lowerCAmelCase_  :							Optional[Any]								=					output.images

        lowerCAmelCase_  :							int								=					pipe(
            **self.get_dummy_inputs(SCREAMING_SNAKE_CASE_			)       ,     return_dict=SCREAMING_SNAKE_CASE_       ,     )[0]

        lowerCAmelCase_  :							Any								=					image[0, -3:, -3:, -1]
        lowerCAmelCase_  :							str								=					image_from_tuple[0, -3:, -3:, -1]

        assert image.shape == (1, 6_4, 6_4, 3)

        lowerCAmelCase_  :							Tuple								=					np.array(
            [0.32_86_63, 1.0, 0.23_21_68_73, 1.0, 0.92_71_75_64, 0.4_63_90_46, 0.96_89_47_77, 0.31_71_33_78, 0.6_29_39_53]			)

        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()}"







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




      """simple docstring"""




      def 	SCREAMING_SNAKE_CASE__					( self							: Optional[Any]			):
        # clean up the VRAM after each test
        super().tearDown()
        gc.collect()
        torch.cuda.empty_cache()





      def 	SCREAMING_SNAKE_CASE__					( self							: int			):
        lowerCAmelCase_  :							Dict								=					load_numpy(
            'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'
            '/kandinsky/kandinsky_text2img_cat_fp16.npy'			)

        lowerCAmelCase_  :							Tuple								=					KandinskyPriorPipeline.from_pretrained(
            'kandinsky-community/kandinsky-2-1-prior'       ,     torch_dtype=torch.floataa			)
        pipe_prior.to(SCREAMING_SNAKE_CASE_			)

        lowerCAmelCase_  :							Dict								=					KandinskyPipeline.from_pretrained('kandinsky-community/kandinsky-2-1'       ,     torch_dtype=torch.floataa			)
        lowerCAmelCase_  :							Union[str, Any]								=					pipeline.to(SCREAMING_SNAKE_CASE_			)
        pipeline.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE_			)

        lowerCAmelCase_  :							str								=					'red cat, 4k photo'

        lowerCAmelCase_  :							Optional[Any]								=					torch.Generator(device='cuda'			).manual_seed(0			)
        lowerCAmelCase_		,lowerCAmelCase_  :							Optional[int]								=					pipe_prior(
            SCREAMING_SNAKE_CASE_       ,     generator=SCREAMING_SNAKE_CASE_       ,     num_inference_steps=5       ,     negative_prompt=''       ,     ).to_tuple()

        lowerCAmelCase_  :							Union[str, Any]								=					torch.Generator(device='cuda'			).manual_seed(0			)
        lowerCAmelCase_  :							Optional[int]								=					pipeline(
            SCREAMING_SNAKE_CASE_       ,     image_embeds=SCREAMING_SNAKE_CASE_       ,     negative_image_embeds=SCREAMING_SNAKE_CASE_       ,     generator=SCREAMING_SNAKE_CASE_       ,     num_inference_steps=1_0_0       ,     output_type='np'       ,     )

        lowerCAmelCase_  :							Tuple								=					output.images[0]

        assert image.shape == (5_1_2, 5_1_2, 3)

        assert_mean_pixel_difference(SCREAMING_SNAKE_CASE_       ,     SCREAMING_SNAKE_CASE_			)










"""simple docstring"""







def 							UpperCamelCase_       (       lowerCAmelCase__ :							int = 400_0000		)		->  int:

  """simple docstring"""




  if n <= 1:
    return 0
  lowerCAmelCase_  :							Union[str, Any]								=					0
  lowerCAmelCase_  :							Tuple								=					2
  lowerCAmelCase_  :							List[Any]								=					0
  while 4 * b + a <= n:
    lowerCAmelCase_		,lowerCAmelCase_  :							str								=					b, 4 * b + a
    count += a
  return count + b


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





