




import gc
import unittest

import numpy as np
import torch
from transformers import CLIPTextConfig, CLIPTextModel, CLIPTokenizer

from diffusers import (
    AutoencoderKL,
    DDIMScheduler,
    EulerAncestralDiscreteScheduler,
    PNDMScheduler,
    StableDiffusionModelEditingPipeline,
    UNetaDConditionModel,
)
from diffusers.utils import slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu, skip_mps

from ..pipeline_params import TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_IMAGE_PARAMS, TEXT_TO_IMAGE_PARAMS
from ..test_pipelines_common import PipelineKarrasSchedulerTesterMixin, PipelineLatentTesterMixin, PipelineTesterMixin


enable_full_determinism()
@skip_mps
class      _UpperCamelCase (			lowerCAmelCase		,    lowerCAmelCase		,    lowerCAmelCase		,    unittest.TestCase						):
      UpperCAmelCase_            =						StableDiffusionModelEditingPipeline
      UpperCAmelCase_            =						TEXT_TO_IMAGE_PARAMS
      UpperCAmelCase_            =						TEXT_TO_IMAGE_BATCH_PARAMS
      UpperCAmelCase_            =						TEXT_TO_IMAGE_IMAGE_PARAMS
      UpperCAmelCase_            =						TEXT_TO_IMAGE_IMAGE_PARAMS
      def 		UpperCAmelCase_  (     self				:Dict )       ->     str:
       torch.manual_seed(0 )
       UpperCAmelCase__								=  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	, )
       UpperCAmelCase__								=  DDIMScheduler()
       torch.manual_seed(0 )
       UpperCAmelCase__								=  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	, )
       torch.manual_seed(0 )
       UpperCAmelCase__								=  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=1000	, )
       UpperCAmelCase__								=  CLIPTextModel(lowerCamelCase )
       UpperCAmelCase__								=  CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip" )

       UpperCAmelCase__								=  {
           "unet": unet,
           "scheduler": scheduler,
           "vae": vae,
           "text_encoder": text_encoder,
           "tokenizer": tokenizer,
           "safety_checker": None,
           "feature_extractor": None,
       }
       return components
      def 		UpperCAmelCase_  (     self				:str	, lowerCamelCase				:Union[str, Any]	, lowerCamelCase				:Any=0 )       ->     int:
       UpperCAmelCase__								=  torch.manual_seed(lowerCamelCase )
       UpperCAmelCase__								=  {
           "prompt": "A field of roses",
           "generator": generator,
           # Setting height and width to None to prevent OOMs on CPU.
           "height": None,
           "width": None,
           "num_inference_steps": 2,
           "guidance_scale": 6.0,
           "output_type": "numpy",
       }
       return inputs
      def 		UpperCAmelCase_  (     self				:Optional[int] )       ->     List[str]:
       UpperCAmelCase__								=  "cpu"  # ensure determinism for the device-dependent torch.Generator
       UpperCAmelCase__								=  self.get_dummy_components()
       UpperCAmelCase__								=  StableDiffusionModelEditingPipeline(**lowerCamelCase )
       UpperCAmelCase__								=  sd_pipe.to(lowerCamelCase )
       sd_pipe.set_progress_bar_config(disable=lowerCamelCase )

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

       UpperCAmelCase__								=  np.array([0.47_55, 0.51_32, 0.49_76, 0.39_04, 0.35_54, 0.47_65, 0.51_39, 0.51_58, 0.48_89] )

       assert np.abs(image_slice.flatten() - expected_slice ).max() < 1e-2
      def 		UpperCAmelCase_  (     self				:Union[str, Any] )       ->     Dict:
       UpperCAmelCase__								=  "cpu"  # ensure determinism for the device-dependent torch.Generator
       UpperCAmelCase__								=  self.get_dummy_components()
       UpperCAmelCase__								=  StableDiffusionModelEditingPipeline(**lowerCamelCase )
       UpperCAmelCase__								=  sd_pipe.to(lowerCamelCase )
       sd_pipe.set_progress_bar_config(disable=lowerCamelCase )

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

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

       UpperCAmelCase__								=  np.array([0.49_92, 0.51_01, 0.50_04, 0.39_49, 0.36_04, 0.47_35, 0.52_16, 0.52_04, 0.49_13] )

       assert np.abs(image_slice.flatten() - expected_slice ).max() < 1e-2
      def 		UpperCAmelCase_  (     self				:Tuple )       ->     Optional[int]:
       UpperCAmelCase__								=  "cpu"  # ensure determinism for the device-dependent torch.Generator
       UpperCAmelCase__								=  self.get_dummy_components()
       UpperCAmelCase__								=  EulerAncestralDiscreteScheduler(
           beta_start=0.0_00_85	, beta_end=0.0_12	, beta_schedule="scaled_linear" )
       UpperCAmelCase__								=  StableDiffusionModelEditingPipeline(**lowerCamelCase )
       UpperCAmelCase__								=  sd_pipe.to(lowerCamelCase )
       sd_pipe.set_progress_bar_config(disable=lowerCamelCase )

       UpperCAmelCase__								=  self.get_dummy_inputs(lowerCamelCase )
       UpperCAmelCase__								=  sd_pipe(**lowerCamelCase ).images
       UpperCAmelCase__								=  image[0, -3:, -3:, -1]

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

       UpperCAmelCase__								=  np.array([0.47_47, 0.53_72, 0.47_79, 0.49_82, 0.55_43, 0.48_16, 0.52_38, 0.49_04, 0.50_27] )

       assert np.abs(image_slice.flatten() - expected_slice ).max() < 1e-2
      def 		UpperCAmelCase_  (     self				:Any )       ->     Dict:
       UpperCAmelCase__								=  "cpu"  # ensure determinism for the device-dependent torch.Generator
       UpperCAmelCase__								=  self.get_dummy_components()
       UpperCAmelCase__								=  PNDMScheduler()
       UpperCAmelCase__								=  StableDiffusionModelEditingPipeline(**lowerCamelCase )
       UpperCAmelCase__								=  sd_pipe.to(lowerCamelCase )
       sd_pipe.set_progress_bar_config(disable=lowerCamelCase )

       UpperCAmelCase__								=  self.get_dummy_inputs(lowerCamelCase )
       # the pipeline does not expect pndm so test if it raises error.
       with self.assertRaises(lowerCamelCase ):
        UpperCAmelCase__								=  sd_pipe(**lowerCamelCase ).images
      def 		UpperCAmelCase_  (     self				:Union[str, Any] )       ->     Tuple:
       super().test_inference_batch_single_identical(expected_max_diff=5e-3 )







      def 		UpperCAmelCase_  (     self				:int )       ->     List[Any]:
       super().test_attention_slicing_forward_pass(expected_max_diff=5e-3 )



@slow
@require_torch_gpu
class      _UpperCamelCase (			unittest.TestCase						):
      def 		UpperCAmelCase_  (     self				:List[Any] )       ->     int:
       super().tearDown()
       gc.collect()
       torch.cuda.empty_cache()
      def 		UpperCAmelCase_  (     self				:List[Any]	, lowerCamelCase				:List[Any]=0 )       ->     Optional[Any]:
       UpperCAmelCase__								=  torch.manual_seed(lowerCamelCase )
       UpperCAmelCase__								=  {
           "prompt": "A field of roses",
           "generator": generator,
           "num_inference_steps": 3,
           "guidance_scale": 7.5,
           "output_type": "numpy",
       }
       return inputs
      def 		UpperCAmelCase_  (     self				:Optional[int] )       ->     int:
       UpperCAmelCase__								=  "CompVis/stable-diffusion-v1-4"
       UpperCAmelCase__								=  StableDiffusionModelEditingPipeline.from_pretrained(lowerCamelCase	, safety_checker=lowerCamelCase )
       pipe.to(lowerCamelCase )
       pipe.set_progress_bar_config(disable=lowerCamelCase )
       pipe.enable_attention_slicing()

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

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

       UpperCAmelCase__								=  np.array(
           [0.6_74_94_96, 0.6_38_64_53, 0.51_44_32_67, 0.66_09_49_05, 0.61_92_12_15, 0.5_49_13_32, 0.5_74_44_17, 0.58_07_51_06, 0.5_17_46_58] )

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

       # make sure image changes after editing
       pipe.edit_model("A pack of roses"	, "A pack of blue roses" )

       UpperCAmelCase__								=  pipe(**lowerCamelCase ).images
       UpperCAmelCase__								=  image[0, -3:, -3:, -1].flatten()

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

       assert np.abs(expected_slice - image_slice ).max() > 1e-1







      def 		UpperCAmelCase_  (     self				:Optional[int] )       ->     Optional[Any]:
       torch.cuda.empty_cache()
       torch.cuda.reset_max_memory_allocated()
       torch.cuda.reset_peak_memory_stats()

       UpperCAmelCase__								=  "CompVis/stable-diffusion-v1-4"
       UpperCAmelCase__								=  DDIMScheduler.from_pretrained(lowerCamelCase	, subfolder="scheduler" )
       UpperCAmelCase__								=  StableDiffusionModelEditingPipeline.from_pretrained(
           lowerCamelCase	, scheduler=lowerCamelCase	, safety_checker=lowerCamelCase )
       UpperCAmelCase__								=  pipe.to(lowerCamelCase )
       pipe.set_progress_bar_config(disable=lowerCamelCase )
       pipe.enable_attention_slicing(1 )
       pipe.enable_sequential_cpu_offload()

       UpperCAmelCase__								=  self.get_inputs()
       UpperCAmelCase__								=  pipe(**lowerCamelCase )

       UpperCAmelCase__								=  torch.cuda.max_memory_allocated()
       # make sure that less than 4.4 GB is allocated
       assert mem_bytes < 4.4 * 10**9











from __future__ import annotations

from string import ascii_letters



def   lowerCAmelCase						(						_lowerCAmelCase							:				str  ,							_lowerCAmelCase							:				int  ,							_lowerCAmelCase							:				str | None = None		):




 """simple docstring"""







 UpperCAmelCase__								=  alphabet or ascii_letters

 # The final result string
 UpperCAmelCase__								=  ""

 for character in input_string:
  if character not in alpha:
   # Append without encryption if character is not in the alphabet
   result += character
  else:
   # Get the index of the new key and make sure it isn't too large
   UpperCAmelCase__								=  (alpha.index(_lowerCAmelCase		) + key) % len(_lowerCAmelCase		)

   # Append the encoded character to the alphabet
   result += alpha[new_key]

 return result



def   lowerCAmelCase						(						_lowerCAmelCase							:				str  ,							_lowerCAmelCase							:				int  ,							_lowerCAmelCase							:				str | None = None		):




 """simple docstring"""







 key *= -1

 return encrypt(_lowerCAmelCase  ,							_lowerCAmelCase  ,							_lowerCAmelCase		)



def   lowerCAmelCase						(						_lowerCAmelCase							:				str  ,							_lowerCAmelCase							:				str | None = None		):




 """simple docstring"""







 UpperCAmelCase__								=  alphabet or ascii_letters

 # To store data on all the combinations
 UpperCAmelCase__								=  {}

 # Cycle through each combination
 for key in range(1  ,							len(_lowerCAmelCase		) + 1		):
  # Decrypt the message and store the result in the data
  UpperCAmelCase__								=  decrypt(_lowerCAmelCase  ,							_lowerCAmelCase  ,							_lowerCAmelCase		)

 return brute_force_data


if __name__ == "__main__":
  while True:
    print(F'''\n{"-" * 1_0}\n Menu\n{"-" * 1_0}''')
    print(*["1.Encrypt", "2.Decrypt", "3.BruteForce", "4.Quit"], sep="\n")

    # get user input
    _lowerCAmelCase  :				List[Any]						=				input("\nWhat would you like to do?: ").strip() or "4"

    # run functions based on what the user chose
    if choice not in ("1", "2", "3", "4"):
      print("Invalid choice, please enter a valid choice")
    elif choice == "1":
      _lowerCAmelCase  :				Any						=				input("Please enter the string to be encrypted: ")
      _lowerCAmelCase  :				str						=				int(input("Please enter off-set: ").strip())

      print(encrypt(input_string, key))
    elif choice == "2":
      _lowerCAmelCase  :				Optional[Any]						=				input("Please enter the string to be decrypted: ")
      _lowerCAmelCase  :				Any						=				int(input("Please enter off-set: ").strip())

      print(decrypt(input_string, key))
    elif choice == "3":
      _lowerCAmelCase  :				int						=				input("Please enter the string to be decrypted: ")
      _lowerCAmelCase  :				Dict						=				brute_force(input_string)

      for key, value in brute_force_data.items():
        print(F'''Key: {key} | Message: {value}''')

    elif choice == "4":
      print("Goodbye.")
      break





