


'''simple docstring'''





import os
from glob import glob

import imageio
import torch
import torchvision
import wandb
from img_processing import custom_to_pil, loop_post_process, preprocess, preprocess_vqgan
from loaders import load_vqgan
from PIL import Image
from torch import nn

from transformers import CLIPModel, CLIPTokenizerFast
from utils import get_device, get_timestamp, show_pil




class    UpperCAmelCase					:




			def __init__(						self							:		Dict       ,       __magic_name__							:		str = "cpu"       ,       __magic_name__							:		str = "openai/clip-vit-large-patch14"						):







						"""simple docstring"""







						UpperCamelCase         =					device
						UpperCamelCase         =					CLIPTokenizerFast.from_pretrained(__magic_name__						)
						UpperCamelCase         =					[0.48_145_466, 0.4_578_275, 0.40_821_073]
						UpperCamelCase         =					[0.26_862_954, 0.26_130_258, 0.27_577_711]
						UpperCamelCase         =					torchvision.transforms.Normalize(self.image_mean       ,       self.image_std						)
						UpperCamelCase         =					torchvision.transforms.Resize(2_2_4						)
						UpperCamelCase         =					torchvision.transforms.CenterCrop(2_2_4						)




			def   lowerCamelCase_					(						self							:		Optional[Any]       ,       __magic_name__							:		List[str]						):







						"""simple docstring"""







						UpperCamelCase         =					self.resize(__magic_name__						)
						UpperCamelCase         =					self.center_crop(__magic_name__						)
						UpperCamelCase         =					self.normalize(__magic_name__						)
						return images





			def __call__(						self							:		Union[str, Any]       ,       __magic_name__							:		str=None       ,       __magic_name__							:		Any=None       ,       **__magic_name__							:		int						):







						"""simple docstring"""







						UpperCamelCase         =					self.tokenizer(text=__magic_name__       ,       **__magic_name__						)
						UpperCamelCase         =					self.preprocess_img(__magic_name__						)
						UpperCamelCase         =					{key: value.to(self.device						) for (key, value) in encoding.items()}
						return encoding




class    UpperCAmelCase					(       nn.Module    ):




			def __init__(						self							:		Tuple       ,       __magic_name__							:		Union[str, Any]=1_0       ,       __magic_name__							:		Union[str, Any]=0.01       ,       __magic_name__							:		str=None       ,       __magic_name__							:		Any=None       ,       __magic_name__							:		List[Any]=None       ,       __magic_name__							:		Optional[Any]=None       ,       __magic_name__							:		Tuple=None       ,       __magic_name__							:		List[Any]=None       ,       __magic_name__							:		Tuple=False       ,       __magic_name__							:		List[str]=True       ,       __magic_name__							:		Union[str, Any]="image"       ,       __magic_name__							:		Union[str, Any]=True       ,       __magic_name__							:		List[str]=False       ,       __magic_name__							:		Any=False       ,       __magic_name__							:		List[Any]=False       ,       ):







						"""simple docstring"""







						super().__init__()
						UpperCamelCase         =					None
						UpperCamelCase         =					device if device else get_device()
						if vqgan:
									UpperCamelCase         =					vqgan
						else:
									UpperCamelCase         =					load_vqgan(self.device       ,       conf_path=__magic_name__       ,       ckpt_path=__magic_name__						)
						self.vqgan.eval()
						if clip:
									UpperCamelCase         =					clip
						else:
									UpperCamelCase         =					CLIPModel.from_pretrained("""openai/clip-vit-base-patch32"""						)
						self.clip.to(self.device						)
						UpperCamelCase         =					ProcessorGradientFlow(device=self.device						)

						UpperCamelCase         =					iterations
						UpperCamelCase         =					lr
						UpperCamelCase         =					log
						UpperCamelCase         =					make_grid
						UpperCamelCase         =					return_val
						UpperCamelCase         =					quantize
						UpperCamelCase         =					self.vqgan.decoder.z_shape




			def   lowerCamelCase_					(						self							:		int       ,       __magic_name__							:		List[Any]=None       ,       __magic_name__							:		Any=None       ,       __magic_name__							:		Tuple=5       ,       __magic_name__							:		Any=True						):







						"""simple docstring"""







						UpperCamelCase         =					[]
						if output_path is None:
									UpperCamelCase         =					"""./animation.gif"""
						if input_path is None:
									UpperCamelCase         =					self.save_path
						UpperCamelCase         =					sorted(glob(input_path + """/*"""						)						)
						if not len(__magic_name__						):
									raise ValueError(
									    """No images found in save path, aborting (did you pass save_intermediate=True to the generate"""
									    """ function?)"""						)
						if len(__magic_name__						) == 1:
									print("""Only one image found in save path, (did you pass save_intermediate=True to the generate function?)"""						)
						UpperCamelCase         =					total_duration / len(__magic_name__						)
						UpperCamelCase         =					[frame_duration] * len(__magic_name__						)
						if extend_frames:
									UpperCamelCase         =					1.5
									UpperCamelCase         =					3
						for file_name in paths:
									if file_name.endswith(""".png"""						):
												images.append(imageio.imread(__magic_name__						)						)
						imageio.mimsave(__magic_name__       ,       __magic_name__       ,       duration=__magic_name__						)
						print(F'gif saved to {output_path}'						)




			def   lowerCamelCase_					(						self							:		int       ,       __magic_name__							:		Optional[int]=None       ,       __magic_name__							:		List[Any]=None						):







						"""simple docstring"""







						if not (path or img):
									raise ValueError("""Input either path or tensor"""						)
						if img is not None:
									raise NotImplementedError
						UpperCamelCase         =					preprocess(Image.open(__magic_name__						)       ,       target_image_size=2_5_6						).to(self.device						)
						UpperCamelCase         =					preprocess_vqgan(__magic_name__						)
						UpperCamelCase  ,  *UpperCamelCase         =					self.vqgan.encode(__magic_name__						)
						return z




			def   lowerCamelCase_					(						self							:		Tuple       ,       __magic_name__							:		Dict						):







						"""simple docstring"""







						UpperCamelCase         =					self.latent.detach().requires_grad_()
						UpperCamelCase         =					base_latent + transform_vector
						if self.quantize:
									UpperCamelCase  ,  *UpperCamelCase         =					self.vqgan.quantize(__magic_name__						)
						else:
									UpperCamelCase         =					trans_latent
						return self.vqgan.decode(__magic_name__						)




			def   lowerCamelCase_					(						self							:		Union[str, Any]       ,       __magic_name__							:		Tuple       ,       __magic_name__							:		Optional[int]       ,       __magic_name__							:		Optional[Any]=None						):







						"""simple docstring"""







						UpperCamelCase         =					self.clip_preprocessor(text=__magic_name__       ,       images=__magic_name__       ,       return_tensors="""pt"""       ,       padding=__magic_name__						)
						UpperCamelCase         =					self.clip(**__magic_name__						)
						UpperCamelCase         =					clip_outputs.logits_per_image
						if weights is not None:
									UpperCamelCase         =					similarity_logits * weights
						return similarity_logits.sum()




			def   lowerCamelCase_					(						self							:		Any       ,       __magic_name__							:		Dict       ,       __magic_name__							:		List[Any]       ,       __magic_name__							:		str						):







						"""simple docstring"""







						UpperCamelCase         =					self._get_clip_similarity(pos_prompts["""prompts"""]       ,       __magic_name__       ,       weights=(1 / pos_prompts["""weights"""])						)
						if neg_prompts:
									UpperCamelCase         =					self._get_clip_similarity(neg_prompts["""prompts"""]       ,       __magic_name__       ,       weights=neg_prompts["""weights"""]						)
						else:
									UpperCamelCase         =					torch.tensor([1]       ,       device=self.device						)
						UpperCamelCase         =					-torch.log(__magic_name__						) + torch.log(__magic_name__						)
						return loss




			def   lowerCamelCase_					(						self							:		Dict       ,       __magic_name__							:		int       ,       __magic_name__							:		int       ,       __magic_name__							:		List[Any]						):







						"""simple docstring"""







						UpperCamelCase         =					torch.randn_like(self.latent       ,       requires_grad=__magic_name__       ,       device=self.device						)
						UpperCamelCase         =					torch.optim.Adam([vector]       ,       lr=self.lr						)

						for i in range(self.iterations						):
									optim.zero_grad()
									UpperCamelCase         =					self._add_vector(__magic_name__						)
									UpperCamelCase         =					loop_post_process(__magic_name__						)
									UpperCamelCase         =					self._get_CLIP_loss(__magic_name__       ,       __magic_name__       ,       __magic_name__						)
									print("""CLIP loss"""       ,       __magic_name__						)
									if self.log:
												wandb.log({"""CLIP Loss""": clip_loss}						)
									clip_loss.backward(retain_graph=__magic_name__						)
									optim.step()
									if self.return_val == "image":
												yield custom_to_pil(transformed_img[0]						)
									else:
												yield vector




			def   lowerCamelCase_					(						self							:		Optional[int]       ,       __magic_name__							:		List[Any]       ,       __magic_name__							:		Union[str, Any]       ,       __magic_name__							:		List[Any]						):







						"""simple docstring"""







						wandb.init(reinit=__magic_name__       ,       project="""face-editor"""						)
						wandb.config.update({"""Positive Prompts""": positive_prompts}						)
						wandb.config.update({"""Negative Prompts""": negative_prompts}						)
						wandb.config.update({"""lr""": self.lr, """iterations""": self.iterations}						)
						if image_path:
									UpperCamelCase         =					Image.open(__magic_name__						)
									UpperCamelCase         =					image.resize((2_5_6, 2_5_6)						)
									wandb.log("""Original Image"""       ,       wandb.Image(__magic_name__						)						)




			def   lowerCamelCase_					(						self							:		Optional[Any]       ,       __magic_name__							:		List[Any]						):







						"""simple docstring"""







						if not prompts:
									return []
						UpperCamelCase         =					[]
						UpperCamelCase         =					[]
						if isinstance(__magic_name__       ,       __magic_name__						):
									UpperCamelCase         =					[prompt.strip() for prompt in prompts.split("""|"""						)]
						for prompt in prompts:
									if isinstance(__magic_name__       ,       (tuple, list)						):
												UpperCamelCase         =					prompt[0]
												UpperCamelCase         =					float(prompt[1]						)
									elif ":" in prompt:
												UpperCamelCase  ,  UpperCamelCase         =					prompt.split(""":"""						)
												UpperCamelCase         =					float(__magic_name__						)
									else:
												UpperCamelCase         =					prompt
												UpperCamelCase         =					1.0
									processed_prompts.append(__magic_name__						)
									weights.append(__magic_name__						)
						return {
						    "prompts": processed_prompts,
						    "weights": torch.tensor(__magic_name__       ,       device=self.device						),
						}





			def   lowerCamelCase_					(						self							:		Any       ,       __magic_name__							:		Optional[Any]       ,       __magic_name__							:		List[Any]=None       ,       __magic_name__							:		Optional[Any]=None       ,       __magic_name__							:		Optional[int]=True       ,       __magic_name__							:		Optional[Any]=False       ,       __magic_name__							:		Any=True       ,       __magic_name__							:		Dict=True       ,       __magic_name__							:		Dict=None       ,       ):







						"""simple docstring"""







						if image_path:
									UpperCamelCase         =					self._get_latent(__magic_name__						)
						else:
									UpperCamelCase         =					torch.randn(self.latent_dim       ,       device=self.device						)
						if self.log:
									self._init_logging(__magic_name__       ,       __magic_name__       ,       __magic_name__						)

						assert pos_prompts, "You must provide at least one positive prompt."
						UpperCamelCase         =					self.process_prompts(__magic_name__						)
						UpperCamelCase         =					self.process_prompts(__magic_name__						)
						if save_final and save_path is None:
									UpperCamelCase         =					os.path.join("""./outputs/"""       ,       """_""".join(pos_prompts["""prompts"""]						)						)
						if not os.path.exists(__magic_name__						):
									os.makedirs(__magic_name__						)
						else:
									UpperCamelCase         =					save_path + """_""" + get_timestamp()
									os.makedirs(__magic_name__						)
						UpperCamelCase         =					save_path

						UpperCamelCase         =					self.vqgan.decode(self.latent						)[0]
						if show_intermediate:
									print("""Original Image"""						)
									show_pil(custom_to_pil(__magic_name__						)						)

						UpperCamelCase         =					loop_post_process(__magic_name__						)
						for iter, transformed_img in enumerate(self._optimize_CLIP(__magic_name__       ,       __magic_name__       ,       __magic_name__						)						):
									if show_intermediate:
												show_pil(__magic_name__						)
									if save_intermediate:
												transformed_img.save(os.path.join(self.save_path       ,       F'iter_{iter:03d}.png'						)						)
									if self.log:
												wandb.log({"""Image""": wandb.Image(__magic_name__						)}						)
						if show_final:
									show_pil(__magic_name__						)
						if save_final:
									transformed_img.save(os.path.join(self.save_path       ,       F'iter_{iter:03d}_final.png'						)						)








import os
import random
import sys

from . import cryptomath_module as cryptomath
from . import rabin_miller

__snake_case            =   3






def 							_lowercase    (		SCREAMING_SNAKE_CASE_   :		int	):





			"""simple docstring"""




			print("""Generating primitive root of p"""	)
			while True:
						UpperCamelCase         =					random.randrange(3       ,  SCREAMING_SNAKE_CASE_	)
						if pow(SCREAMING_SNAKE_CASE_       ,  2       ,  SCREAMING_SNAKE_CASE_	) == 1:
									continue
						if pow(SCREAMING_SNAKE_CASE_       ,  SCREAMING_SNAKE_CASE_       ,  SCREAMING_SNAKE_CASE_	) == 1:
									continue
						return g






def 							_lowercase    (		SCREAMING_SNAKE_CASE_   :		int	):





			"""simple docstring"""




			print("""Generating prime p..."""	)
			UpperCamelCase         =					rabin_miller.generate_large_prime(SCREAMING_SNAKE_CASE_	)  # select large prime number.
			UpperCamelCase         =					primitive_root(SCREAMING_SNAKE_CASE_	)  # one primitive root on modulo p.
			UpperCamelCase         =					random.randrange(3       ,  SCREAMING_SNAKE_CASE_	)  # private_key -> have to be greater than 2 for safety.
			UpperCamelCase         =					cryptomath.find_mod_inverse(pow(SCREAMING_SNAKE_CASE_       ,  SCREAMING_SNAKE_CASE_       ,  SCREAMING_SNAKE_CASE_	)       ,  SCREAMING_SNAKE_CASE_	)

			UpperCamelCase         =					(key_size, e_a, e_a, p)
			UpperCamelCase         =					(key_size, d)

			return public_key, private_key






def 							_lowercase    (		SCREAMING_SNAKE_CASE_   :		str       ,  SCREAMING_SNAKE_CASE_   :		int	):





			"""simple docstring"""




			if os.path.exists(f'{name}_pubkey.txt'	) or os.path.exists(f'{name}_privkey.txt'	):
						print("""\nWARNING:"""	)
						print(
						    f'"{name}_pubkey.txt" or "{name}_privkey.txt" already exists. \n'
						    """Use a different name or delete these files and re-run this program."""	)
						sys.exit()

			UpperCamelCase  ,  UpperCamelCase         =					generate_key(SCREAMING_SNAKE_CASE_	)
			print(f'\nWriting public key to file {name}_pubkey.txt...'	)
			with open(f'{name}_pubkey.txt'       ,  """w"""	) as fo:
						fo.write(f'{public_key[0]},{public_key[1]},{public_key[2]},{public_key[3]}'	)

			print(f'Writing private key to file {name}_privkey.txt...'	)
			with open(f'{name}_privkey.txt'       ,  """w"""	) as fo:
						fo.write(f'{private_key[0]},{private_key[1]}'	)






def 							_lowercase    (		):





			"""simple docstring"""




			print("""Making key files..."""	)
			make_key_files("""elgamal"""       ,  2_048	)
			print("""Key files generation successful"""	)


if __name__ == "__main__":
			main()




