


"""simple docstring"""





import tempfile

import numpy as np
import torch
from transformers import AutoTokenizer, TaEncoderModel

from diffusers import DDPMScheduler, UNetaDConditionModel
from diffusers.models.attention_processor import AttnAddedKVProcessor
from diffusers.pipelines.deepfloyd_if import IFWatermarker
from diffusers.utils.testing_utils import torch_device

from ..test_pipelines_common import to_np



class 						snake_case_:






     def    lowerCamelCase__   (		self     :    Optional[Any] ):
           torch.manual_seed(0 )
           lowerCAmelCase    :		Optional[Any]				=      TaEncoderModel.from_pretrained('''hf-internal-testing/tiny-random-t5''' )

           torch.manual_seed(0 )
           lowerCAmelCase    :		List[str]				=      AutoTokenizer.from_pretrained('''hf-internal-testing/tiny-random-t5''' )

           torch.manual_seed(0 )
           lowerCAmelCase    :		str				=      UNetaDConditionModel(
               sample_size=3_2    , layers_per_block=1    , block_out_channels=[3_2, 6_4]    , down_block_types=[
                   '''ResnetDownsampleBlock2D''',
                   '''SimpleCrossAttnDownBlock2D''',
               ]    , mid_block_type='''UNetMidBlock2DSimpleCrossAttn'''    , up_block_types=['''SimpleCrossAttnUpBlock2D''', '''ResnetUpsampleBlock2D''']    , in_channels=3    , out_channels=6    , cross_attention_dim=3_2    , encoder_hid_dim=3_2    , attention_head_dim=8    , addition_embed_type='''text'''    , addition_embed_type_num_heads=2    , cross_attention_norm='''group_norm'''    , resnet_time_scale_shift='''scale_shift'''    , act_fn='''gelu'''    , )
           unet.set_attn_processor(AttnAddedKVProcessor() )  # For reproducibility tests

           torch.manual_seed(0 )
           lowerCAmelCase    :		int				=      DDPMScheduler(
               num_train_timesteps=1_0_0_0    , beta_schedule='''squaredcos_cap_v2'''    , beta_start=0.0_001    , beta_end=0.02    , thresholding=UpperCamelCase_    , dynamic_thresholding_ratio=0.95    , sample_max_value=1.0    , prediction_type='''epsilon'''    , variance_type='''learned_range'''    , )

           torch.manual_seed(0 )
           lowerCAmelCase    :		Tuple				=      IFWatermarker()

           return {
               "text_encoder": text_encoder,
               "tokenizer": tokenizer,
               "unet": unet,
               "scheduler": scheduler,
               "watermarker": watermarker,
               "safety_checker": None,
               "feature_extractor": None,
           }






     def    lowerCamelCase__   (		self     :    int ):
           torch.manual_seed(0 )
           lowerCAmelCase    :		str				=      TaEncoderModel.from_pretrained('''hf-internal-testing/tiny-random-t5''' )

           torch.manual_seed(0 )
           lowerCAmelCase    :		Any				=      AutoTokenizer.from_pretrained('''hf-internal-testing/tiny-random-t5''' )

           torch.manual_seed(0 )
           lowerCAmelCase    :		Any				=      UNetaDConditionModel(
               sample_size=3_2    , layers_per_block=[1, 2]    , block_out_channels=[3_2, 6_4]    , down_block_types=[
                   '''ResnetDownsampleBlock2D''',
                   '''SimpleCrossAttnDownBlock2D''',
               ]    , mid_block_type='''UNetMidBlock2DSimpleCrossAttn'''    , up_block_types=['''SimpleCrossAttnUpBlock2D''', '''ResnetUpsampleBlock2D''']    , in_channels=6    , out_channels=6    , cross_attention_dim=3_2    , encoder_hid_dim=3_2    , attention_head_dim=8    , addition_embed_type='''text'''    , addition_embed_type_num_heads=2    , cross_attention_norm='''group_norm'''    , resnet_time_scale_shift='''scale_shift'''    , act_fn='''gelu'''    , class_embed_type='''timestep'''    , mid_block_scale_factor=1.414    , time_embedding_act_fn='''gelu'''    , time_embedding_dim=3_2    , )
           unet.set_attn_processor(AttnAddedKVProcessor() )  # For reproducibility tests

           torch.manual_seed(0 )
           lowerCAmelCase    :		Optional[Any]				=      DDPMScheduler(
               num_train_timesteps=1_0_0_0    , beta_schedule='''squaredcos_cap_v2'''    , beta_start=0.0_001    , beta_end=0.02    , thresholding=UpperCamelCase_    , dynamic_thresholding_ratio=0.95    , sample_max_value=1.0    , prediction_type='''epsilon'''    , variance_type='''learned_range'''    , )

           torch.manual_seed(0 )
           lowerCAmelCase    :		Tuple				=      DDPMScheduler(
               num_train_timesteps=1_0_0_0    , beta_schedule='''squaredcos_cap_v2'''    , beta_start=0.0_001    , beta_end=0.02    , )

           torch.manual_seed(0 )
           lowerCAmelCase    :		List[str]				=      IFWatermarker()

           return {
               "text_encoder": text_encoder,
               "tokenizer": tokenizer,
               "unet": unet,
               "scheduler": scheduler,
               "image_noising_scheduler": image_noising_scheduler,
               "watermarker": watermarker,
               "safety_checker": None,
               "feature_extractor": None,
           }






     def    lowerCamelCase__   (		self     :    Tuple ):
           lowerCAmelCase    :		Union[str, Any]				=      self.get_dummy_components()
           lowerCAmelCase    :		str				=      self.pipeline_class(**UpperCamelCase_ )
           pipe.to(UpperCamelCase_ )
           pipe.set_progress_bar_config(disable=UpperCamelCase_ )

           lowerCAmelCase    :		Tuple				=      self.get_dummy_inputs(UpperCamelCase_ )

           lowerCAmelCase    :		Any				=      inputs['''prompt''']
           lowerCAmelCase    :		Optional[Any]				=      inputs['''generator''']
           lowerCAmelCase    :		Tuple				=      inputs['''num_inference_steps''']
           lowerCAmelCase    :		Any				=      inputs['''output_type''']

           if "image" in inputs:
                 lowerCAmelCase    :		str				=      inputs['''image''']
           else:
                 lowerCAmelCase    :		List[str]				=      None

           if "mask_image" in inputs:
                 lowerCAmelCase    :		Optional[Any]				=      inputs['''mask_image''']
           else:
                 lowerCAmelCase    :		List[str]				=      None

           if "original_image" in inputs:
                 lowerCAmelCase    :		Tuple				=      inputs['''original_image''']
           else:
                 lowerCAmelCase    :		Optional[int]				=      None

           lowerCAmelCase,					lowerCAmelCase    :		Optional[Any]				=      pipe.encode_prompt(UpperCamelCase_ )

           # inputs with prompt converted to embeddings
           lowerCAmelCase    :		Dict				=      {
               '''prompt_embeds''': prompt_embeds,
               '''negative_prompt_embeds''': negative_prompt_embeds,
               '''generator''': generator,
               '''num_inference_steps''': num_inference_steps,
               '''output_type''': output_type,
           }

           if image is not None:
                 lowerCAmelCase    :		Optional[Any]				=      image

           if mask_image is not None:
                 lowerCAmelCase    :		Tuple				=      mask_image

           if original_image is not None:
                 lowerCAmelCase    :		Optional[Any]				=      original_image

           # set all optional components to None
           for optional_component in pipe._optional_components:
                 setattr(UpperCamelCase_    , UpperCamelCase_    , UpperCamelCase_ )

           lowerCAmelCase    :		Dict				=      pipe(**UpperCamelCase_ )[0]

           with tempfile.TemporaryDirectory() as tmpdir:
                 pipe.save_pretrained(UpperCamelCase_ )
                 lowerCAmelCase    :		Union[str, Any]				=      self.pipeline_class.from_pretrained(UpperCamelCase_ )
                 pipe_loaded.to(UpperCamelCase_ )
                 pipe_loaded.set_progress_bar_config(disable=UpperCamelCase_ )

           pipe_loaded.unet.set_attn_processor(AttnAddedKVProcessor() )  # For reproducibility tests

           for optional_component in pipe._optional_components:
                 self.assertTrue(
                     getattr(UpperCamelCase_    , UpperCamelCase_ ) is None    , F'''`{optional_component}` did not stay set to None after loading.'''    , )

           lowerCAmelCase    :		int				=      self.get_dummy_inputs(UpperCamelCase_ )

           lowerCAmelCase    :		List[str]				=      inputs['''generator''']
           lowerCAmelCase    :		Dict				=      inputs['''num_inference_steps''']
           lowerCAmelCase    :		Optional[int]				=      inputs['''output_type''']

           # inputs with prompt converted to embeddings
           lowerCAmelCase    :		Union[str, Any]				=      {
               '''prompt_embeds''': prompt_embeds,
               '''negative_prompt_embeds''': negative_prompt_embeds,
               '''generator''': generator,
               '''num_inference_steps''': num_inference_steps,
               '''output_type''': output_type,
           }

           if image is not None:
                 lowerCAmelCase    :		Any				=      image

           if mask_image is not None:
                 lowerCAmelCase    :		Optional[int]				=      mask_image

           if original_image is not None:
                 lowerCAmelCase    :		Any				=      original_image

           lowerCAmelCase    :		List[str]				=      pipe_loaded(**UpperCamelCase_ )[0]

           lowerCAmelCase    :		str				=      np.abs(to_np(UpperCamelCase_ ) - to_np(UpperCamelCase_ ) ).max()
           self.assertLess(UpperCamelCase_    , 1E-4 )






     def    lowerCamelCase__   (		self     :    str ):
           lowerCAmelCase    :		str				=      self.get_dummy_components()
           lowerCAmelCase    :		int				=      self.pipeline_class(**UpperCamelCase_ )
           pipe.to(UpperCamelCase_ )
           pipe.set_progress_bar_config(disable=UpperCamelCase_ )

           lowerCAmelCase    :		Optional[int]				=      self.get_dummy_inputs(UpperCamelCase_ )
           lowerCAmelCase    :		List[Any]				=      pipe(**UpperCamelCase_ )[0]

           with tempfile.TemporaryDirectory() as tmpdir:
                 pipe.save_pretrained(UpperCamelCase_ )
                 lowerCAmelCase    :		Union[str, Any]				=      self.pipeline_class.from_pretrained(UpperCamelCase_ )
                 pipe_loaded.to(UpperCamelCase_ )
                 pipe_loaded.set_progress_bar_config(disable=UpperCamelCase_ )

           pipe_loaded.unet.set_attn_processor(AttnAddedKVProcessor() )  # For reproducibility tests

           lowerCAmelCase    :		Union[str, Any]				=      self.get_dummy_inputs(UpperCamelCase_ )
           lowerCAmelCase    :		int				=      pipe_loaded(**UpperCamelCase_ )[0]

           lowerCAmelCase    :		Union[str, Any]				=      np.abs(to_np(UpperCamelCase_ ) - to_np(UpperCamelCase_ ) ).max()
           self.assertLess(UpperCamelCase_    , 1E-4 )






"""simple docstring"""





import math


def    _snake_case  (  ):
      lowerCAmelCase    :		Union[str, Any]				=      input('''Enter message: '''	)
      lowerCAmelCase    :		Optional[int]				=      int(input(f'''Enter key [2-{len(_snake_case	) - 1}]: '''	)	)
      lowerCAmelCase    :		str				=      input('''Encryption/Decryption [e/d]: '''	)

      if mode.lower().startswith('''e'''	):
            lowerCAmelCase    :		Any				=      encrypt_message(_snake_case       ,      _snake_case	)
      elif mode.lower().startswith('''d'''	):
            lowerCAmelCase    :		Union[str, Any]				=      decrypt_message(_snake_case       ,      _snake_case	)

      # Append pipe symbol (vertical bar) to identify spaces at the end.
      print(f'''Output:\n{text + "|"}'''	)


def    _snake_case  (  _snake_case :	int       ,      _snake_case :	str	):
      lowerCAmelCase    :		Optional[Any]				=      [''''''] * key
      for col in range(_snake_case	):
            lowerCAmelCase    :		Optional[Any]				=      col
            while pointer < len(_snake_case	):
                  cipher_text[col] += message[pointer]
                  pointer += key
      return "".join(_snake_case	)


def    _snake_case  (  _snake_case :	int       ,      _snake_case :	str	):
      lowerCAmelCase    :		Union[str, Any]				=      math.ceil(len(_snake_case	) / key	)
      lowerCAmelCase    :		str				=      key
      lowerCAmelCase    :		Any				=      (num_cols * num_rows) - len(_snake_case	)
      lowerCAmelCase    :		Dict				=      [''''''] * num_cols
      lowerCAmelCase    :		int				=      0
      lowerCAmelCase    :		int				=      0

      for symbol in message:
            plain_text[col] += symbol
            col += 1

            if (
                (col == num_cols)
                or (col == num_cols - 1)
                and (row >= num_rows - num_shaded_boxes)
            ):
                  lowerCAmelCase    :		int				=      0
                  row += 1

      return "".join(_snake_case	)


if __name__ == "__main__":
 import doctest

 doctest.testmod()
 main()


