




import inspect
from typing import Optional, Union

import numpy as np
import PIL
import torch
from torch.nn import functional as F
from torchvision import transforms
from transformers import CLIPFeatureExtractor, CLIPModel, CLIPTextModel, CLIPTokenizer

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

def     _SCREAMING_SNAKE_CASE    (							lowercase			: str			,							lowercase			: Union[str, Any]			,							lowercase			: Optional[int]  ):


      '''simple docstring'''



      if isinstance(lowercase			,							torch.Tensor  ):
            return image
      elif isinstance(lowercase			,							PIL.Image.Image  ):
            lowerCamelCase_	  =					[image]

      if isinstance(image[0]			,							PIL.Image.Image  ):
            lowerCamelCase_	  =					[np.array(i.resize((w, h)			,							resample=PIL_INTERPOLATION['lanczos']  )  )[None, :] for i in image]
            lowerCamelCase_	  =					np.concatenate(lowercase			,							axis=0  )
            lowerCamelCase_	  =					np.array(lowercase  ).astype(np.floataa  ) / 255.0
            lowerCamelCase_	  =					image.transpose(0			,							3			,							1			,							2  )
            lowerCamelCase_	  =					2.0 * image - 1.0
            lowerCamelCase_	  =					torch.from_numpy(lowercase  )
      elif isinstance(image[0]			,							torch.Tensor  ):
            lowerCamelCase_	  =					torch.cat(lowercase			,							dim=0  )
      return image

def     _SCREAMING_SNAKE_CASE    (							lowercase			: int			,							lowercase			: Dict			,							lowercase			: Union[str, Any]			,							lowercase			: Optional[int]=0.9995  ):


      '''simple docstring'''



      if not isinstance(lowercase			,							np.ndarray  ):
            lowerCamelCase_	  =					True
            lowerCamelCase_	  =					va.device
            lowerCamelCase_	  =					va.cpu().numpy()
            lowerCamelCase_	  =					va.cpu().numpy()

      lowerCamelCase_	  =					np.sum(va * va / (np.linalg.norm(lowercase  ) * np.linalg.norm(lowercase  ))  )
      if np.abs(lowercase  ) > DOT_THRESHOLD:
            lowerCamelCase_	  =					(1 - t) * va + t * va
      else:
            lowerCamelCase_	  =					np.arccos(lowercase  )
            lowerCamelCase_	  =					np.sin(lowercase  )
            lowerCamelCase_	  =					theta_a * t
            lowerCamelCase_	  =					np.sin(lowercase  )
            lowerCamelCase_	  =					np.sin(theta_a - theta_t  ) / sin_theta_a
            lowerCamelCase_	  =					sin_theta_t / sin_theta_a
            lowerCamelCase_	  =					sa * va + sa * va

      if inputs_are_torch:
            lowerCamelCase_	  =					torch.from_numpy(lowercase  ).to(lowercase  )

      return va

def     _SCREAMING_SNAKE_CASE    (							lowercase			: Optional[int]			,							lowercase			: int  ):


      '''simple docstring'''



      lowerCamelCase_	  =					F.normalize(lowercase			,							dim=-1  )
      lowerCamelCase_	  =					F.normalize(lowercase			,							dim=-1  )
      return (x - y).norm(dim=-1  ).div(2  ).arcsin().pow(2  ).mul(2  )

def     _SCREAMING_SNAKE_CASE    (							lowercase			: Optional[Any]			,							lowercase			: Union[str, Any]  ):


      '''simple docstring'''



      for param in model.parameters():
            lowerCamelCase_	  =					value







class 						A(		UpperCamelCase     ):



 '''simple docstring'''


 def __init__(       self      :						Tuple		,							A_      :						AutoencoderKL		,							A_      :						CLIPTextModel		,							A_      :						CLIPModel		,							A_      :						CLIPTokenizer		,							A_      :						UNetaDConditionModel		,							A_      :						Union[PNDMScheduler, LMSDiscreteScheduler, DDIMScheduler, DPMSolverMultistepScheduler]		,							A_      :						CLIPFeatureExtractor		,							A_      :						Optional[int]=None		,							A_      :						Tuple=None		,							A_      :						Any=None		,							)     ->			List[Any]:

       """simple docstring"""






       super().__init__()
       self.register_modules(
           vae=A_		,							text_encoder=A_		,							clip_model=A_		,							tokenizer=A_		,							unet=A_		,							scheduler=A_		,							feature_extractor=A_		,							coca_model=A_		,							coca_tokenizer=A_		,							coca_transform=A_		,							)
       lowerCamelCase_	  =					(
           feature_extractor.size
           if isinstance(feature_extractor.size		,							A_						)
           else feature_extractor.size['shortest_edge']
       )
       lowerCamelCase_	  =					transforms.Normalize(mean=feature_extractor.image_mean		,							std=feature_extractor.image_std						)
       set_requires_grad(self.text_encoder		,							A_						)
       set_requires_grad(self.clip_model		,							A_						)

 def        a__				(       self      :						Union[str, Any]		,							A_      :						Optional[Union[str, int]] = "auto"						)     ->			str:

       """simple docstring"""






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

 def        a__				(       self      :						List[Any]						)     ->			str:

       """simple docstring"""






       self.enable_attention_slicing(A_						)

 def        a__				(       self      :						Optional[int]						)     ->			Optional[int]:

       """simple docstring"""






       set_requires_grad(self.vae		,							A_						)

 def        a__				(       self      :						Union[str, Any]						)     ->			Optional[Any]:

       """simple docstring"""






       set_requires_grad(self.vae		,							A_						)

 def        a__				(       self      :						Any						)     ->			Tuple:

       """simple docstring"""






       set_requires_grad(self.unet		,							A_						)

 def        a__				(       self      :						Tuple						)     ->			List[Any]:

       """simple docstring"""






       set_requires_grad(self.unet		,							A_						)

 def        a__				(       self      :						List[Any]		,							A_      :						Any		,							A_      :						int		,							A_      :						Dict						)     ->			Dict:

       """simple docstring"""






       lowerCamelCase_	  =					min(int(num_inference_steps * strength						)		,							A_						)

       lowerCamelCase_	  =					max(num_inference_steps - init_timestep		,							0						)
       lowerCamelCase_	  =					self.scheduler.timesteps[t_start:]

       return timesteps, num_inference_steps - t_start

 def        a__				(       self      :						Optional[int]		,							A_      :						Tuple		,							A_      :						Tuple		,							A_      :						Dict		,							A_      :						List[Any]		,							A_      :						Dict		,							A_      :						Dict=None						)     ->			Optional[int]:

       """simple docstring"""






       if not isinstance(A_		,							torch.Tensor						):
             raise ValueError(f"""`image` has to be of type `torch.Tensor` but is {type(A_						)}"""						)

       lowerCamelCase_	  =					image.to(device=A_		,							dtype=A_						)

       if isinstance(A_		,							A_						):
             lowerCamelCase_	  =					[
                 self.vae.encode(image[i : i + 1]						).latent_dist.sample(generator[i]						) for i in range(A_						)
             ]
             lowerCamelCase_	  =					torch.cat(A_		,							dim=0						)
       else:
             lowerCamelCase_	  =					self.vae.encode(A_						).latent_dist.sample(A_						)

       # Hardcode 0.18215 because stable-diffusion-2-base has not self.vae.config.scaling_factor
       lowerCamelCase_	  =					0.18215 * init_latents
       lowerCamelCase_	  =					init_latents.repeat_interleave(A_		,							dim=0						)

       lowerCamelCase_	  =					randn_tensor(init_latents.shape		,							generator=A_		,							device=A_		,							dtype=A_						)

       # get latents
       lowerCamelCase_	  =					self.scheduler.add_noise(A_		,							A_		,							A_						)
       lowerCamelCase_	  =					init_latents

       return latents

 def        a__				(       self      :						List[str]		,							A_      :						Optional[int]						)     ->			List[Any]:

       """simple docstring"""






       lowerCamelCase_	  =					self.coca_transform(A_						).unsqueeze(0						)
       with torch.no_grad(), torch.cuda.amp.autocast():
             lowerCamelCase_	  =					self.coca_model.generate(transformed_image.to(device=self.device		,							dtype=self.coca_model.dtype						)						)
       lowerCamelCase_	  =					self.coca_tokenizer.decode(generated[0].cpu().numpy()						)
       return generated.split('<end_of_text>'						)[0].replace('<start_of_text>'		,							''						).rstrip(' .,'						)

 def        a__				(       self      :						Union[str, Any]		,							A_      :						str		,							A_      :						Union[str, Any]						)     ->			List[str]:

       """simple docstring"""






       lowerCamelCase_	  =					self.feature_extractor.preprocess(A_						)
       lowerCamelCase_	  =					torch.from_numpy(clip_image_input['pixel_values'][0]						).unsqueeze(0						).to(self.device						).half()
       lowerCamelCase_	  =					self.clip_model.get_image_features(A_						)
       lowerCamelCase_	  =					image_embeddings_clip / image_embeddings_clip.norm(p=2		,							dim=-1		,							keepdim=A_						)
       lowerCamelCase_	  =					image_embeddings_clip.repeat_interleave(A_		,							dim=0						)
       return image_embeddings_clip

 @torch.enable_grad()
 def        a__				(       self      :						Optional[Any]		,							A_      :						List[Any]		,							A_      :						Dict		,							A_      :						List[str]		,							A_      :						Tuple		,							A_      :						Dict		,							A_      :						int		,							A_      :						Tuple		,							)     ->			str:

       """simple docstring"""






       lowerCamelCase_	  =					latents.detach().requires_grad_()

       lowerCamelCase_	  =					self.scheduler.scale_model_input(A_		,							A_						)

       # predict the noise residual
       lowerCamelCase_	  =					self.unet(A_		,							A_		,							encoder_hidden_states=A_						).sample

       if isinstance(self.scheduler		,							(PNDMScheduler, DDIMScheduler, DPMSolverMultistepScheduler)						):
             lowerCamelCase_	  =					self.scheduler.alphas_cumprod[timestep]
             lowerCamelCase_	  =					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_	  =					(latents - beta_prod_t ** 0.5 * noise_pred) / alpha_prod_t ** 0.5

             lowerCamelCase_	  =					torch.sqrt(A_						)
             lowerCamelCase_	  =					pred_original_sample * (fac) + latents * (1 - fac)
       elif isinstance(self.scheduler		,							A_						):
             lowerCamelCase_	  =					self.scheduler.sigmas[index]
             lowerCamelCase_	  =					latents - sigma * noise_pred
       else:
             raise ValueError(f"""scheduler type {type(self.scheduler						)} not supported"""						)

       # Hardcode 0.18215 because stable-diffusion-2-base has not self.vae.config.scaling_factor
       lowerCamelCase_	  =					1 / 0.18215 * sample
       lowerCamelCase_	  =					self.vae.decode(A_						).sample
       lowerCamelCase_	  =					(image / 2 + 0.5).clamp(0		,							1						)

       lowerCamelCase_	  =					transforms.Resize(self.feature_extractor_size						)(A_						)
       lowerCamelCase_	  =					self.normalize(A_						).to(latents.dtype						)

       lowerCamelCase_	  =					self.clip_model.get_image_features(A_						)
       lowerCamelCase_	  =					image_embeddings_clip / image_embeddings_clip.norm(p=2		,							dim=-1		,							keepdim=A_						)

       lowerCamelCase_	  =					spherical_dist_loss(A_		,							A_						).mean() * clip_guidance_scale

       lowerCamelCase_	  =					-torch.autograd.grad(A_		,							A_						)[0]

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





 @torch.no_grad()
 def __call__(       self      :						int		,							A_      :						Union[torch.FloatTensor, PIL.Image.Image]		,							A_      :						Union[torch.FloatTensor, PIL.Image.Image]		,							A_      :						Optional[str] = None		,							A_      :						Optional[str] = None		,							A_      :						Optional[int] = 512		,							A_      :						Optional[int] = 512		,							A_      :						float = 0.6		,							A_      :						Optional[int] = 50		,							A_      :						Optional[float] = 7.5		,							A_      :						Optional[int] = 1		,							A_      :						float = 0.0		,							A_      :						Optional[float] = 100		,							A_      :						Optional[torch.Generator] = None		,							A_      :						Optional[str] = "pil"		,							A_      :						bool = True		,							A_      :						float = 0.8		,							A_      :						float = 0.1		,							A_      :						float = 0.1		,							)     ->			int:

       """simple docstring"""






       if isinstance(A_		,							A_						) and len(A_						) != batch_size:
             raise ValueError(f"""You have passed {batch_size} batch_size, but only {len(A_						)} generators."""						)

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

       if isinstance(A_		,							torch.Generator						) and batch_size > 1:
             lowerCamelCase_	  =					[generator] + [None] * (batch_size - 1)

       lowerCamelCase_	  =					[
           ('model', self.coca_model is None),
           ('tokenizer', self.coca_tokenizer is None),
           ('transform', self.coca_transform is None),
       ]
       lowerCamelCase_	  =					[x[0] for x in coca_is_none if x[1]]
       lowerCamelCase_	  =					', '.join(A_						)
       # generate prompts with coca model if prompt is None
       if content_prompt is None:
             if len(A_						):
                   raise ValueError(
                       f"""Content prompt is None and CoCa [{coca_is_none_str}] is None."""
                       f"""Set prompt or pass Coca [{coca_is_none_str}] to DiffusionPipeline."""						)
             lowerCamelCase_	  =					self.get_image_description(A_						)
       if style_prompt is None:
             if len(A_						):
                   raise ValueError(
                       f"""Style prompt is None and CoCa [{coca_is_none_str}] is None."""
                       f""" Set prompt or pass Coca [{coca_is_none_str}] to DiffusionPipeline."""						)
             lowerCamelCase_	  =					self.get_image_description(A_						)

       # get prompt text embeddings for content and style
       lowerCamelCase_	  =					self.tokenizer(
           A_		,							padding='max_length'		,							max_length=self.tokenizer.model_max_length		,							truncation=A_		,							return_tensors='pt'		,							)
       lowerCamelCase_	  =					self.text_encoder(content_text_input.input_ids.to(self.device						)						)[0]

       lowerCamelCase_	  =					self.tokenizer(
           A_		,							padding='max_length'		,							max_length=self.tokenizer.model_max_length		,							truncation=A_		,							return_tensors='pt'		,							)
       lowerCamelCase_	  =					self.text_encoder(style_text_input.input_ids.to(self.device						)						)[0]

       lowerCamelCase_	  =					slerp(A_		,							A_		,							A_						)

       # duplicate text embeddings for each generation per prompt
       lowerCamelCase_	  =					text_embeddings.repeat_interleave(A_		,							dim=0						)

       # set timesteps
       lowerCamelCase_	  =					'offset' in set(inspect.signature(self.scheduler.set_timesteps						).parameters.keys()						)
       lowerCamelCase_	  =					{}
       if accepts_offset:
             lowerCamelCase_	  =					1

       self.scheduler.set_timesteps(A_		,							**A_						)
       # Some schedulers like PNDM have timesteps as arrays
       # It's more optimized to move all timesteps to correct device beforehand
       self.scheduler.timesteps.to(self.device						)

       lowerCamelCase_  ,       lowerCamelCase_	  =					self.get_timesteps(A_		,							A_		,							self.device						)
       lowerCamelCase_	  =					timesteps[:1].repeat(A_						)

       # Preprocess image
       lowerCamelCase_	  =					preprocess(A_		,							A_		,							A_						)
       lowerCamelCase_	  =					self.prepare_latents(
           A_		,							A_		,							A_		,							text_embeddings.dtype		,							self.device		,							A_						)

       lowerCamelCase_	  =					preprocess(A_		,							A_		,							A_						)
       lowerCamelCase_	  =					self.prepare_latents(
           A_		,							A_		,							A_		,							text_embeddings.dtype		,							self.device		,							A_						)

       lowerCamelCase_	  =					slerp(A_		,							A_		,							A_						)

       if clip_guidance_scale > 0:
             lowerCamelCase_	  =					self.get_clip_image_embeddings(A_		,							A_						)
             lowerCamelCase_	  =					self.get_clip_image_embeddings(A_		,							A_						)
             lowerCamelCase_	  =					slerp(
                 A_		,							A_		,							A_						)

       # 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_	  =					guidance_scale > 1.0
       # get unconditional embeddings for classifier free guidance
       if do_classifier_free_guidance:
             lowerCamelCase_	  =					content_text_input.input_ids.shape[-1]
             lowerCamelCase_	  =					self.tokenizer(['']		,							padding='max_length'		,							max_length=A_		,							return_tensors='pt'						)
             lowerCamelCase_	  =					self.text_encoder(uncond_input.input_ids.to(self.device						)						)[0]
             # duplicate unconditional embeddings for each generation per prompt
             lowerCamelCase_	  =					uncond_embeddings.repeat_interleave(A_		,							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_	  =					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_	  =					(batch_size, self.unet.config.in_channels, height // 8, width // 8)
       lowerCamelCase_	  =					text_embeddings.dtype
       if latents is None:
             if self.device.type == "mps":
                   # randn does not work reproducibly on mps
                   lowerCamelCase_	  =					torch.randn(A_		,							generator=A_		,							device='cpu'		,							dtype=A_						).to(
                       self.device						)
             else:
                   lowerCamelCase_	  =					torch.randn(A_		,							generator=A_		,							device=self.device		,							dtype=A_						)
       else:
             if latents.shape != latents_shape:
                   raise ValueError(f"""Unexpected latents shape, got {latents.shape}, expected {latents_shape}"""						)
             lowerCamelCase_	  =					latents.to(self.device						)

       # scale the initial noise by the standard deviation required by the scheduler
       lowerCamelCase_	  =					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_	  =					'eta' in set(inspect.signature(self.scheduler.step						).parameters.keys()						)
       lowerCamelCase_	  =					{}
       if accepts_eta:
             lowerCamelCase_	  =					eta

       # check if the scheduler accepts generator
       lowerCamelCase_	  =					'generator' in set(inspect.signature(self.scheduler.step						).parameters.keys()						)
       if accepts_generator:
             lowerCamelCase_	  =					generator

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

                   # predict the noise residual
                   lowerCamelCase_	  =					self.unet(A_		,							A_		,							encoder_hidden_states=A_						).sample

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

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

                   # compute the previous noisy sample x_t -> x_t-1
                   lowerCamelCase_	  =					self.scheduler.step(A_		,							A_		,							A_		,							**A_						).prev_sample

        # Hardcode 0.18215 because stable-diffusion-2-base has not self.vae.config.scaling_factor
       lowerCamelCase_	  =					1 / 0.18215 * latents
       lowerCamelCase_	  =					self.vae.decode(A_						).sample

       lowerCamelCase_	  =					(image / 2 + 0.5).clamp(0		,							1						)
       lowerCamelCase_	  =					image.cpu().permute(0		,							2		,							3		,							1						).numpy()

       if output_type == "pil":
             lowerCamelCase_	  =					self.numpy_to_pil(A_						)

       if not return_dict:
             return (image, None)

       return StableDiffusionPipelineOutput(images=A_		,							nsfw_content_detected=A_						)









import platform
from argparse import ArgumentParser

import huggingface_hub

from .. import __version__ as version
from ..utils import is_accelerate_available, is_torch_available, is_transformers_available, is_xformers_available
from . import BaseDiffusersCLICommand

def     _SCREAMING_SNAKE_CASE    (							lowercase			: Optional[Any]  ):


      '''simple docstring'''



      return EnvironmentCommand()







class 						A(		UpperCamelCase     ):



 '''simple docstring'''


 @staticmethod
 def        a__				(       A_      :						ArgumentParser						)     ->			str:

       """simple docstring"""






       lowerCamelCase_	  =					parser.add_parser('env'						)
       download_parser.set_defaults(func=A_						)

 def        a__				(       self      :						Optional[Any]						)     ->			Any:

       """simple docstring"""






       lowerCamelCase_	  =					huggingface_hub.__version__

       lowerCamelCase_	  =					'not installed'
       lowerCamelCase_	  =					'NA'
       if is_torch_available():
             import torch

             lowerCamelCase_	  =					torch.__version__
             lowerCamelCase_	  =					torch.cuda.is_available()

       lowerCamelCase_	  =					'not installed'
       if is_transformers_available():
             import transformers

             lowerCamelCase_	  =					transformers.__version__

       lowerCamelCase_	  =					'not installed'
       if is_accelerate_available():
             import accelerate

             lowerCamelCase_	  =					accelerate.__version__

       lowerCamelCase_	  =					'not installed'
       if is_xformers_available():
             import xformers

             lowerCamelCase_	  =					xformers.__version__

       lowerCamelCase_	  =					{
           '`diffusers` version': version,
           'Platform': platform.platform(),
           'Python version': platform.python_version(),
           'PyTorch version (GPU?)': f"""{pt_version} ({pt_cuda_available})""",
           'Huggingface_hub version': hub_version,
           'Transformers version': transformers_version,
           'Accelerate version': accelerate_version,
           'xFormers version': xformers_version,
           'Using GPU in script?': '<fill in>',
           'Using distributed or parallel set-up in script?': '<fill in>',
       }

       print('\nCopy-and-paste the text below in your GitHub issue and FILL OUT the two last points.\n'						)
       print(self.format_dict(A_						)						)

       return info





 @staticmethod
 def        a__				(       A_      :						Dict						)     ->			Any:

       """simple docstring"""






       return "\n".join([f"""- {prop}: {val}""" for prop, val in d.items()]						) + "\n"



