


import inspect
from typing import List, Optional, Union

import torch
from torch import nn
from torch.nn import functional as F
from torchvision import transforms
from transformers import CLIPImageProcessor, CLIPModel, CLIPTextModel, CLIPTokenizer

from diffusers import (
    AutoencoderKL,
    DDIMScheduler,
    DiffusionPipeline,
    DPMSolverMultistepScheduler,
    LMSDiscreteScheduler,
    PNDMScheduler,
    UNetaDConditionModel,
)
from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion import StableDiffusionPipelineOutput







class      _lowercase					(				nn.Module):


       """simple docstring"""




       def __init__(  self				:     str   , __lowerCamelCase				:     Union[str, Any]   , __lowerCamelCase				:     Optional[int]=1.0			):


             '''simple docstring'''



             super().__init__()

             lowerCamelCase__						:							Tuple       =    cut_size
             lowerCamelCase__						:							str       =    cut_power





       def 		lowerCAmelCase	(  self				:     List[str]   , __lowerCamelCase				:     Optional[int]   , __lowerCamelCase				:     Any			):


             '''simple docstring'''



             lowerCamelCase__						:							Any       =    pixel_values.shape[2:4]
             lowerCamelCase__						:							str       =    min(lowercase_   , lowercase_			)
             lowerCamelCase__						:							Union[str, Any]       =    min(lowercase_   , lowercase_   , self.cut_size			)
             lowerCamelCase__						:							Union[str, Any]       =    []
             for _ in range(lowercase_			):
                   lowerCamelCase__						:							List[Any]       =    int(torch.rand([]			) ** self.cut_power * (max_size - min_size) + min_size			)
                   lowerCamelCase__						:							List[Any]       =    torch.randint(0   , sideX - size + 1   , ()			)
                   lowerCamelCase__						:							List[Any]       =    torch.randint(0   , sideY - size + 1   , ()			)
                   lowerCamelCase__						:							Any       =    pixel_values[:, :, offsety : offsety + size, offsetx : offsetx + size]
                   cutouts.append(F.adaptive_avg_poolad(lowercase_   , self.cut_size			)			)
             return torch.cat(lowercase_			)



def  lowercase_						(						_A     :       Optional[Any]       ,   _A     :       Union[str, Any]		):

      """simple docstring"""




      lowerCamelCase__						:							str       =    F.normalize(SCREAMING_SNAKE_CASE__       ,   dim=-1		)
      lowerCamelCase__						:							Tuple       =    F.normalize(SCREAMING_SNAKE_CASE__       ,   dim=-1		)
      return (x - y).norm(dim=-1		).div(2		).arcsin().pow(2		).mul(2		)

def  lowercase_						(						_A     :       Optional[int]       ,   _A     :       Any		):

      """simple docstring"""




      for param in model.parameters():
            lowerCamelCase__						:							Optional[int]       =    value







class      _lowercase					(				snake_case_):


       """simple docstring"""




       def __init__(  self				:     Tuple   , __lowerCamelCase				:     Optional[int]   , __lowerCamelCase				:     List[Any]   , __lowerCamelCase				:     List[Any]   , __lowerCamelCase				:     Optional[Any]   , __lowerCamelCase				:     List[Any]   , __lowerCamelCase				:     List[str]   , __lowerCamelCase				:     List[Any]   , ):


             '''simple docstring'''



             super().__init__()
             self.register_modules(
                 vae=lowercase_   , text_encoder=lowercase_   , clip_model=lowercase_   , tokenizer=lowercase_   , unet=lowercase_   , scheduler=lowercase_   , feature_extractor=lowercase_   , )

             lowerCamelCase__						:							Optional[Any]       =    transforms.Normalize(mean=feature_extractor.image_mean   , std=feature_extractor.image_std			)
             lowerCamelCase__						:							Dict       =    (
                 feature_extractor.size
                 if isinstance(feature_extractor.size   , lowercase_			)
                 else feature_extractor.size["""shortest_edge"""]
             )
             lowerCamelCase__						:							str       =    MakeCutouts(self.cut_out_size			)

             set_requires_grad(self.text_encoder   , lowercase_			)
             set_requires_grad(self.clip_model   , lowercase_			)





       def 		lowerCAmelCase	(  self				:     Union[str, Any]   , __lowerCamelCase				:     List[str] = "auto"			):


             '''simple docstring'''



             if slice_size == "auto":
                   # half the attention head size is usually a good trade-off between
                   # speed and memory
                   lowerCamelCase__						:							Any       =    self.unet.config.attention_head_dim // 2
             self.unet.set_attention_slice(lowercase_			)





       def 		lowerCAmelCase	(  self				:     str			):


             '''simple docstring'''



             self.enable_attention_slicing(lowercase_			)





       def 		lowerCAmelCase	(  self				:     Union[str, Any]			):


             '''simple docstring'''



             set_requires_grad(self.vae   , lowercase_			)





       def 		lowerCAmelCase	(  self				:     int			):


             '''simple docstring'''



             set_requires_grad(self.vae   , lowercase_			)





       def 		lowerCAmelCase	(  self				:     Tuple			):


             '''simple docstring'''



             set_requires_grad(self.unet   , lowercase_			)





       def 		lowerCAmelCase	(  self				:     str			):


             '''simple docstring'''



             set_requires_grad(self.unet   , lowercase_			)





       @torch.enable_grad()
       def 		lowerCAmelCase	(  self				:     Tuple   , __lowerCamelCase				:     Tuple   , __lowerCamelCase				:     Dict   , __lowerCamelCase				:     Dict   , __lowerCamelCase				:     int   , __lowerCamelCase				:     List[str]   , __lowerCamelCase				:     List[Any]   , __lowerCamelCase				:     Optional[Any]   , __lowerCamelCase				:     List[Any]   , __lowerCamelCase				:     Tuple=True   , ):


             '''simple docstring'''



             lowerCamelCase__						:							Optional[int]       =    latents.detach().requires_grad_()

             lowerCamelCase__						:							Optional[int]       =    self.scheduler.scale_model_input(lowercase_   , lowercase_			)

             # predict the noise residual
             lowerCamelCase__						:							Dict       =    self.unet(lowercase_   , lowercase_   , encoder_hidden_states=lowercase_			).sample

             if isinstance(self.scheduler   , (PNDMScheduler, DDIMScheduler, DPMSolverMultistepScheduler)			):
                   lowerCamelCase__						:							Any       =    self.scheduler.alphas_cumprod[timestep]
                   lowerCamelCase__						:							Union[str, Any]       =    1 - alpha_prod_t
                   # compute predicted original sample from predicted noise also called
                   # "predicted x_0" of formula (12) from https://arxiv.org/pdf/2010.02502.pdf
                   lowerCamelCase__						:							Tuple       =    (latents - beta_prod_t ** 0.5 * noise_pred) / alpha_prod_t ** 0.5

                   lowerCamelCase__						:							Any       =    torch.sqrt(lowercase_			)
                   lowerCamelCase__						:							Tuple       =    pred_original_sample * (fac) + latents * (1 - fac)
             elif isinstance(self.scheduler   , lowercase_			):
                   lowerCamelCase__						:							Optional[int]       =    self.scheduler.sigmas[index]
                   lowerCamelCase__						:							str       =    latents - sigma * noise_pred
             else:
                   raise ValueError(f"scheduler type {type(self.scheduler			)} not supported"			)

             lowerCamelCase__						:							Dict       =    1 / self.vae.config.scaling_factor * sample
             lowerCamelCase__						:							List[str]       =    self.vae.decode(lowercase_			).sample
             lowerCamelCase__						:							Any       =    (image / 2 + 0.5).clamp(0   , 1			)

             if use_cutouts:
                   lowerCamelCase__						:							Union[str, Any]       =    self.make_cutouts(lowercase_   , lowercase_			)
             else:
                   lowerCamelCase__						:							Union[str, Any]       =    transforms.Resize(self.cut_out_size			)(lowercase_			)
             lowerCamelCase__						:							Any       =    self.normalize(lowercase_			).to(latents.dtype			)

             lowerCamelCase__						:							Union[str, Any]       =    self.clip_model.get_image_features(lowercase_			)
             lowerCamelCase__						:							Dict       =    image_embeddings_clip / image_embeddings_clip.norm(p=2   , dim=-1   , keepdim=lowercase_			)

             if use_cutouts:
                   lowerCamelCase__						:							Optional[int]       =    spherical_dist_loss(lowercase_   , lowercase_			)
                   lowerCamelCase__						:							Optional[int]       =    dists.view([num_cutouts, sample.shape[0], -1]			)
                   lowerCamelCase__						:							int       =    dists.sum(2			).mean(0			).sum() * clip_guidance_scale
             else:
                   lowerCamelCase__						:							List[Any]       =    spherical_dist_loss(lowercase_   , lowercase_			).mean() * clip_guidance_scale

             lowerCamelCase__						:							Optional[Any]       =    -torch.autograd.grad(lowercase_   , lowercase_			)[0]

             if isinstance(self.scheduler   , lowercase_			):
                   lowerCamelCase__						:							int       =    latents.detach() + grads * (sigma**2)
                   lowerCamelCase__						:							Any       =    noise_pred_original
             else:
                   lowerCamelCase__						:							Any       =    noise_pred_original - torch.sqrt(lowercase_			) * grads
             return noise_pred, latents





       @torch.no_grad()
       def __call__(  self				:     Tuple   , __lowerCamelCase				:     Union[str, Any]   , __lowerCamelCase				:     Union[str, Any] = 512   , __lowerCamelCase				:     int = 512   , __lowerCamelCase				:     Union[str, Any] = 50   , __lowerCamelCase				:     Dict = 7.5   , __lowerCamelCase				:     Any = 1   , __lowerCamelCase				:     Any = 0.0   , __lowerCamelCase				:     Optional[Any] = 100   , __lowerCamelCase				:     Optional[Any] = None   , __lowerCamelCase				:     Tuple = 4   , __lowerCamelCase				:     List[str] = True   , __lowerCamelCase				:     str = None   , __lowerCamelCase				:     Optional[Any] = None   , __lowerCamelCase				:     List[str] = "pil"   , __lowerCamelCase				:     Optional[Any] = True   , ):


             '''simple docstring'''



             if isinstance(lowercase_   , lowercase_			):
                   lowerCamelCase__						:							str       =    1
             elif isinstance(lowercase_   , lowercase_			):
                   lowerCamelCase__						:							Optional[Any]       =    len(lowercase_			)
             else:
                   raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(lowercase_			)}"			)

             if height % 8 != 0 or width % 8 != 0:
                   raise ValueError(f"`height` and `width` have to be divisible by 8 but are {height} and {width}."			)

             # get prompt text embeddings
             lowerCamelCase__						:							Tuple       =    self.tokenizer(
                 lowercase_   , padding="max_length"   , max_length=self.tokenizer.model_max_length   , truncation=lowercase_   , return_tensors="pt"   , )
             lowerCamelCase__						:							str       =    self.text_encoder(text_input.input_ids.to(self.device			)			)[0]
             # duplicate text embeddings for each generation per prompt
             lowerCamelCase__						:							Dict       =    text_embeddings.repeat_interleave(lowercase_   , dim=0			)

             if clip_guidance_scale > 0:
                   if clip_prompt is not None:
                         lowerCamelCase__						:							int       =    self.tokenizer(
                             lowercase_   , padding="max_length"   , max_length=self.tokenizer.model_max_length   , truncation=lowercase_   , return_tensors="pt"   , ).input_ids.to(self.device			)
                   else:
                         lowerCamelCase__						:							Tuple       =    text_input.input_ids.to(self.device			)
                   lowerCamelCase__						:							Dict       =    self.clip_model.get_text_features(lowercase_			)
                   lowerCamelCase__						:							Tuple       =    text_embeddings_clip / text_embeddings_clip.norm(p=2   , dim=-1   , keepdim=lowercase_			)
                   # duplicate text embeddings clip for each generation per prompt
                   lowerCamelCase__						:							Any       =    text_embeddings_clip.repeat_interleave(lowercase_   , dim=0			)

             # here `guidance_scale` is defined analog to the guidance weight `w` of equation (2)
             # of the Imagen paper: https://arxiv.org/pdf/2205.11487.pdf . `guidance_scale = 1`
             # corresponds to doing no classifier free guidance.
             lowerCamelCase__						:							Optional[Any]       =    guidance_scale > 1.0
             # get unconditional embeddings for classifier free guidance
             if do_classifier_free_guidance:
                   lowerCamelCase__						:							Optional[Any]       =    text_input.input_ids.shape[-1]
                   lowerCamelCase__						:							Dict       =    self.tokenizer([""]   , padding="max_length"   , max_length=lowercase_   , return_tensors="pt"			)
                   lowerCamelCase__						:							str       =    self.text_encoder(uncond_input.input_ids.to(self.device			)			)[0]
                   # duplicate unconditional embeddings for each generation per prompt
                   lowerCamelCase__						:							Optional[int]       =    uncond_embeddings.repeat_interleave(lowercase_   , dim=0			)

                   # For classifier free guidance, we need to do two forward passes.
                   # Here we concatenate the unconditional and text embeddings into a single batch
                   # to avoid doing two forward passes
                   lowerCamelCase__						:							str       =    torch.cat([uncond_embeddings, text_embeddings]			)

             # get the initial random noise unless the user supplied it

             # Unlike in other pipelines, latents need to be generated in the target device
             # for 1-to-1 results reproducibility with the CompVis implementation.
             # However this currently doesn't work in `mps`.
             lowerCamelCase__						:							Optional[int]       =    (batch_size * num_images_per_prompt, self.unet.config.in_channels, height // 8, width // 8)
             lowerCamelCase__						:							int       =    text_embeddings.dtype
             if latents is None:
                   if self.device.type == "mps":
                         # randn does not work reproducibly on mps
                         lowerCamelCase__						:							str       =    torch.randn(lowercase_   , generator=lowercase_   , device="cpu"   , dtype=lowercase_			).to(
                             self.device			)
                   else:
                         lowerCamelCase__						:							Optional[int]       =    torch.randn(lowercase_   , generator=lowercase_   , device=self.device   , dtype=lowercase_			)
             else:
                   if latents.shape != latents_shape:
                         raise ValueError(f"Unexpected latents shape, got {latents.shape}, expected {latents_shape}"			)
                   lowerCamelCase__						:							Any       =    latents.to(self.device			)

             # set timesteps
             lowerCamelCase__						:							List[str]       =    """offset""" in set(inspect.signature(self.scheduler.set_timesteps			).parameters.keys()			)
             lowerCamelCase__						:							int       =    {}
             if accepts_offset:
                   lowerCamelCase__						:							List[str]       =    1

             self.scheduler.set_timesteps(lowercase_   , **lowercase_			)

             # Some schedulers like PNDM have timesteps as arrays
             # It's more optimized to move all timesteps to correct device beforehand
             lowerCamelCase__						:							Union[str, Any]       =    self.scheduler.timesteps.to(self.device			)

             # scale the initial noise by the standard deviation required by the scheduler
             lowerCamelCase__						:							List[Any]       =    latents * self.scheduler.init_noise_sigma

             # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature
             # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers.
             # eta corresponds to η in DDIM paper: https://arxiv.org/abs/2010.02502
             # and should be between [0, 1]
             lowerCamelCase__						:							List[str]       =    """eta""" in set(inspect.signature(self.scheduler.step			).parameters.keys()			)
             lowerCamelCase__						:							int       =    {}
             if accepts_eta:
                   lowerCamelCase__						:							Any       =    eta

             # check if the scheduler accepts generator
             lowerCamelCase__						:							Union[str, Any]       =    """generator""" in set(inspect.signature(self.scheduler.step			).parameters.keys()			)
             if accepts_generator:
                   lowerCamelCase__						:							Dict       =    generator

             for i, t in enumerate(self.progress_bar(lowercase_			)			):
                   # expand the latents if we are doing classifier free guidance
                   lowerCamelCase__						:							int       =    torch.cat([latents] * 2			) if do_classifier_free_guidance else latents
                   lowerCamelCase__						:							Dict       =    self.scheduler.scale_model_input(lowercase_   , lowercase_			)

                   # predict the noise residual
                   lowerCamelCase__						:							Optional[Any]       =    self.unet(lowercase_   , lowercase_   , encoder_hidden_states=lowercase_			).sample

                   # perform classifier free guidance
                   if do_classifier_free_guidance:
                         lowerCamelCase__						:							int       =    noise_pred.chunk(2			)
                         lowerCamelCase__						:							str       =    noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)

                   # perform clip guidance
                   if clip_guidance_scale > 0:
                         lowerCamelCase__						:							Any       =    (
                             text_embeddings.chunk(2			)[1] if do_classifier_free_guidance else text_embeddings
                         )
                         lowerCamelCase__						:							List[Any]       =    self.cond_fn(
                             lowercase_   , lowercase_   , lowercase_   , lowercase_   , lowercase_   , lowercase_   , lowercase_   , lowercase_   , lowercase_   , )

                   # compute the previous noisy sample x_t -> x_t-1
                   lowerCamelCase__						:							List[Any]       =    self.scheduler.step(lowercase_   , lowercase_   , lowercase_   , **lowercase_			).prev_sample

             # scale and decode the image latents with vae
             lowerCamelCase__						:							Dict       =    1 / self.vae.config.scaling_factor * latents
             lowerCamelCase__						:							Optional[int]       =    self.vae.decode(lowercase_			).sample

             lowerCamelCase__						:							Any       =    (image / 2 + 0.5).clamp(0   , 1			)
             lowerCamelCase__						:							str       =    image.cpu().permute(0   , 2   , 3   , 1			).numpy()

             if output_type == "pil":
                   lowerCamelCase__						:							Tuple       =    self.numpy_to_pil(lowercase_			)

             if not return_dict:
                   return (image, None)

             return StableDiffusionPipelineOutput(images=lowercase_   , nsfw_content_detected=lowercase_			)







import os

from apache_beam.io.filesystems import FileSystems
from apache_beam.pipeline import Pipeline

from .logging import get_logger


lowerCAmelCase					:				Union[str, Any]   	=2 << 20  # 2mb
lowerCAmelCase					:				Union[str, Any]   	=get_logger(__name__)






class 	_a			(      snake_case_  ):






			def 							_snake_case   (    self )	->					Dict:
									lowerCAmelCase	:   Optional[int]					     =					self._options.get_all_options().get("""runner""" )
									return runner in [None, "DirectRunner", "PortableRunner"]







def  _UpperCAmelCase  (				SCREAMING_SNAKE_CASE__			,SCREAMING_SNAKE_CASE__			,SCREAMING_SNAKE_CASE__=False  ):


						'''simple docstring'''




						lowerCAmelCase	:   Union[str, Any]					     =					FileSystems
						if fs.exists(SCREAMING_SNAKE_CASE__  ):
												if force_upload:
																		logger.info(F"""Remote path already exist: {remote_file_path}. Overwriting it as force_upload=True."""  )
												else:
																		logger.info(F"""Remote path already exist: {remote_file_path}. Skipping it as force_upload=False."""  )
																		return
						with fs.create(SCREAMING_SNAKE_CASE__  ) as remote_file:
												with open(SCREAMING_SNAKE_CASE__			,"""rb"""  ) as local_file:
																		lowerCAmelCase	:   Any					     =					local_file.read(SCREAMING_SNAKE_CASE__  )
																		while chunk:
																								remote_file.write(SCREAMING_SNAKE_CASE__  )
																								lowerCAmelCase	:   List[str]					     =					local_file.read(SCREAMING_SNAKE_CASE__  )







def  _UpperCAmelCase  (				SCREAMING_SNAKE_CASE__			,SCREAMING_SNAKE_CASE__			,SCREAMING_SNAKE_CASE__=False  ):


						'''simple docstring'''




						lowerCAmelCase	:   Optional[int]					     =					FileSystems
						if os.path.exists(SCREAMING_SNAKE_CASE__  ):
												if force_download:
																		logger.info(F"""Local path already exist: {remote_file_path}. Overwriting it as force_upload=True."""  )
												else:
																		logger.info(F"""Local path already exist: {remote_file_path}. Skipping it as force_upload=False."""  )
																		return
						with fs.open(SCREAMING_SNAKE_CASE__  ) as remote_file:
												with open(SCREAMING_SNAKE_CASE__			,"""wb"""  ) as local_file:
																		lowerCAmelCase	:   Optional[int]					     =					remote_file.read(SCREAMING_SNAKE_CASE__  )
																		while chunk:
																								local_file.write(SCREAMING_SNAKE_CASE__  )
																								lowerCAmelCase	:   Union[str, Any]					     =					remote_file.read(SCREAMING_SNAKE_CASE__  )

