




"""simple docstring"""
import argparse
import os

import align
import numpy as np
import requests
import tensorflow as tf
import torch
from PIL import Image
from tokenizer import Tokenizer

from transformers import (
    AlignConfig,
    AlignModel,
    AlignProcessor,
    BertConfig,
    BertTokenizer,
    EfficientNetConfig,
    EfficientNetImageProcessor,
)
from transformers.utils import logging


logging.set_verbosity_info()
_SCREAMING_SNAKE_CASE						 =     logging.get_logger(__name__)






def 				__lowerCAmelCase					(					__lowerCAmelCase			:		int						)    ->       Optional[int]:
 _UpperCamelCase      :	Optional[int]					      =  tf.image.resize(__lowerCAmelCase  ,	(346, 346)						)
 _UpperCamelCase      :	List[str]					      =  tf.image.crop_to_bounding_box(__lowerCAmelCase  ,	(346 - 289) // 2  ,	(346 - 289) // 2  ,	289  ,	289						)
 return image






def 				__lowerCAmelCase					(					)    ->       Tuple:
 _UpperCamelCase      :	Optional[int]					      =  EfficientNetConfig.from_pretrained("google/efficientnet-b7"						)
 _UpperCamelCase      :	Any					      =  289
 _UpperCamelCase      :	List[Any]					      =  640
 _UpperCamelCase      :	Union[str, Any]					      =  {"0": "LABEL_0", "1": "LABEL_1"}
 _UpperCamelCase      :	int					      =  {"LABEL_0": 0, "LABEL_1": 1}
 _UpperCamelCase      :	Optional[Any]					      =  []

 _UpperCamelCase      :	int					      =  BertConfig()
 _UpperCamelCase      :	Tuple					      =  AlignConfig.from_text_vision_configs(
     text_config=__lowerCAmelCase  ,	vision_config=__lowerCAmelCase  ,	projection_dim=640						)
 return config






def 				__lowerCAmelCase					(					)    ->       Union[str, Any]:
 _UpperCamelCase      :	List[str]					      =  "http://images.cocodataset.org/val2017/000000039769.jpg"
 _UpperCamelCase      :	Any					      =  Image.open(requests.get(__lowerCAmelCase  ,	stream=__lowerCAmelCase						).raw						)
 return im






def 				__lowerCAmelCase					(					)    ->       List[str]:
 _UpperCamelCase      :	List[str]					      =  EfficientNetImageProcessor(
     do_center_crop=__lowerCAmelCase  ,	rescale_factor=1 / 127.5  ,	rescale_offset=__lowerCAmelCase  ,	do_normalize=__lowerCAmelCase  ,	include_top=__lowerCAmelCase  ,	resample=Image.BILINEAR  ,	)
 _UpperCamelCase      :	List[Any]					      =  BertTokenizer.from_pretrained("bert-base-uncased"						)
 _UpperCamelCase      :	Optional[int]					      =  64
 _UpperCamelCase      :	List[Any]					      =  AlignProcessor(image_processor=__lowerCAmelCase  ,	tokenizer=__lowerCAmelCase						)
 return processor






def 				__lowerCAmelCase					(					__lowerCAmelCase			:		Optional[int]						)    ->       Optional[int]:
 # EfficientNet image encoder
 _UpperCamelCase      :	Optional[int]					      =  [v.split("_"						)[0].split("block"						)[1] for v in original_param_names if v.startswith("block"						)]
 _UpperCamelCase      :	Optional[int]					      =  list(set(__lowerCAmelCase						)						)
 _UpperCamelCase      :	List[str]					      =  sorted(__lowerCAmelCase						)
 _UpperCamelCase      :	Optional[Any]					      =  len(__lowerCAmelCase						)
 _UpperCamelCase      :	Union[str, Any]					      =  {b: str(__lowerCAmelCase						) for b, i in zip(__lowerCAmelCase  ,	range(__lowerCAmelCase						)						)}

 _UpperCamelCase      :	Optional[int]					      =  []
 rename_keys.append(("stem_conv/kernel:0", "embeddings.convolution.weight")						)
 rename_keys.append(("stem_bn/gamma:0", "embeddings.batchnorm.weight")						)
 rename_keys.append(("stem_bn/beta:0", "embeddings.batchnorm.bias")						)
 rename_keys.append(("stem_bn/moving_mean:0", "embeddings.batchnorm.running_mean")						)
 rename_keys.append(("stem_bn/moving_variance:0", "embeddings.batchnorm.running_var")						)

 for b in block_names:
  _UpperCamelCase      :	List[Any]					      =  block_name_mapping[b]
  rename_keys.append((f"block{b}_expand_conv/kernel:0", f"encoder.blocks.{hf_b}.expansion.expand_conv.weight")						)
  rename_keys.append((f"block{b}_expand_bn/gamma:0", f"encoder.blocks.{hf_b}.expansion.expand_bn.weight")						)
  rename_keys.append((f"block{b}_expand_bn/beta:0", f"encoder.blocks.{hf_b}.expansion.expand_bn.bias")						)
  rename_keys.append(
      (f"block{b}_expand_bn/moving_mean:0", f"encoder.blocks.{hf_b}.expansion.expand_bn.running_mean")						)
  rename_keys.append(
      (f"block{b}_expand_bn/moving_variance:0", f"encoder.blocks.{hf_b}.expansion.expand_bn.running_var")						)
  rename_keys.append(
      (f"block{b}_dwconv/depthwise_kernel:0", f"encoder.blocks.{hf_b}.depthwise_conv.depthwise_conv.weight")						)
  rename_keys.append((f"block{b}_bn/gamma:0", f"encoder.blocks.{hf_b}.depthwise_conv.depthwise_norm.weight")						)
  rename_keys.append((f"block{b}_bn/beta:0", f"encoder.blocks.{hf_b}.depthwise_conv.depthwise_norm.bias")						)
  rename_keys.append(
      (f"block{b}_bn/moving_mean:0", f"encoder.blocks.{hf_b}.depthwise_conv.depthwise_norm.running_mean")						)
  rename_keys.append(
      (f"block{b}_bn/moving_variance:0", f"encoder.blocks.{hf_b}.depthwise_conv.depthwise_norm.running_var")						)

  rename_keys.append((f"block{b}_se_reduce/kernel:0", f"encoder.blocks.{hf_b}.squeeze_excite.reduce.weight")						)
  rename_keys.append((f"block{b}_se_reduce/bias:0", f"encoder.blocks.{hf_b}.squeeze_excite.reduce.bias")						)
  rename_keys.append((f"block{b}_se_expand/kernel:0", f"encoder.blocks.{hf_b}.squeeze_excite.expand.weight")						)
  rename_keys.append((f"block{b}_se_expand/bias:0", f"encoder.blocks.{hf_b}.squeeze_excite.expand.bias")						)
  rename_keys.append(
      (f"block{b}_project_conv/kernel:0", f"encoder.blocks.{hf_b}.projection.project_conv.weight")						)
  rename_keys.append((f"block{b}_project_bn/gamma:0", f"encoder.blocks.{hf_b}.projection.project_bn.weight")						)
  rename_keys.append((f"block{b}_project_bn/beta:0", f"encoder.blocks.{hf_b}.projection.project_bn.bias")						)
  rename_keys.append(
      (f"block{b}_project_bn/moving_mean:0", f"encoder.blocks.{hf_b}.projection.project_bn.running_mean")						)
  rename_keys.append(
      (f"block{b}_project_bn/moving_variance:0", f"encoder.blocks.{hf_b}.projection.project_bn.running_var")						)

 _UpperCamelCase      :	int					      =  {}
 for item in rename_keys:
  if item[0] in original_param_names:
   _UpperCamelCase      :	Tuple					      =  "vision_model." + item[1]

    # BERT text encoder
 _UpperCamelCase      :	Dict					      =  []
 _UpperCamelCase      :	Optional[int]					      =  "tf_bert_model/bert"
 _UpperCamelCase      :	Any					      =  "text_model"
 for i in range(12						):
  rename_keys.append(
      (
          f"{old}/encoder/layer_._{i}/attention/self/query/kernel:0",
          f"{new}.encoder.layer.{i}.attention.self.query.weight",
      )						)
  rename_keys.append(
      (
          f"{old}/encoder/layer_._{i}/attention/self/query/bias:0",
          f"{new}.encoder.layer.{i}.attention.self.query.bias",
      )						)
  rename_keys.append(
      (
          f"{old}/encoder/layer_._{i}/attention/self/key/kernel:0",
          f"{new}.encoder.layer.{i}.attention.self.key.weight",
      )						)
  rename_keys.append(
      (
          f"{old}/encoder/layer_._{i}/attention/self/key/bias:0",
          f"{new}.encoder.layer.{i}.attention.self.key.bias",
      )						)
  rename_keys.append(
      (
          f"{old}/encoder/layer_._{i}/attention/self/value/kernel:0",
          f"{new}.encoder.layer.{i}.attention.self.value.weight",
      )						)
  rename_keys.append(
      (
          f"{old}/encoder/layer_._{i}/attention/self/value/bias:0",
          f"{new}.encoder.layer.{i}.attention.self.value.bias",
      )						)
  rename_keys.append(
      (
          f"{old}/encoder/layer_._{i}/attention/output/dense/kernel:0",
          f"{new}.encoder.layer.{i}.attention.output.dense.weight",
      )						)
  rename_keys.append(
      (
          f"{old}/encoder/layer_._{i}/attention/output/dense/bias:0",
          f"{new}.encoder.layer.{i}.attention.output.dense.bias",
      )						)
  rename_keys.append(
      (
          f"{old}/encoder/layer_._{i}/attention/output/LayerNorm/gamma:0",
          f"{new}.encoder.layer.{i}.attention.output.LayerNorm.weight",
      )						)
  rename_keys.append(
      (
          f"{old}/encoder/layer_._{i}/attention/output/LayerNorm/beta:0",
          f"{new}.encoder.layer.{i}.attention.output.LayerNorm.bias",
      )						)
  rename_keys.append(
      (
          f"{old}/encoder/layer_._{i}/intermediate/dense/kernel:0",
          f"{new}.encoder.layer.{i}.intermediate.dense.weight",
      )						)
  rename_keys.append(
      (
          f"{old}/encoder/layer_._{i}/intermediate/dense/bias:0",
          f"{new}.encoder.layer.{i}.intermediate.dense.bias",
      )						)
  rename_keys.append(
      (f"{old}/encoder/layer_._{i}/output/dense/kernel:0", f"{new}.encoder.layer.{i}.output.dense.weight")						)
  rename_keys.append(
      (f"{old}/encoder/layer_._{i}/output/dense/bias:0", f"{new}.encoder.layer.{i}.output.dense.bias")						)
  rename_keys.append(
      (f"{old}/encoder/layer_._{i}/output/LayerNorm/gamma:0", f"{new}.encoder.layer.{i}.output.LayerNorm.weight")						)
  rename_keys.append(
      (f"{old}/encoder/layer_._{i}/output/LayerNorm/beta:0", f"{new}.encoder.layer.{i}.output.LayerNorm.bias")						)

 rename_keys.append((f"{old}/embeddings/word_embeddings/weight:0", f"{new}.embeddings.word_embeddings.weight")						)
 rename_keys.append(
     (f"{old}/embeddings/position_embeddings/embeddings:0", f"{new}.embeddings.position_embeddings.weight")						)
 rename_keys.append(
     (f"{old}/embeddings/token_type_embeddings/embeddings:0", f"{new}.embeddings.token_type_embeddings.weight")						)
 rename_keys.append((f"{old}/embeddings/LayerNorm/gamma:0", f"{new}.embeddings.LayerNorm.weight")						)
 rename_keys.append((f"{old}/embeddings/LayerNorm/beta:0", f"{new}.embeddings.LayerNorm.bias")						)

 rename_keys.append((f"{old}/pooler/dense/kernel:0", f"{new}.pooler.dense.weight")						)
 rename_keys.append((f"{old}/pooler/dense/bias:0", f"{new}.pooler.dense.bias")						)
 rename_keys.append(("dense/kernel:0", "text_projection.weight")						)
 rename_keys.append(("dense/bias:0", "text_projection.bias")						)
 rename_keys.append(("dense/bias:0", "text_projection.bias")						)
 rename_keys.append(("temperature:0", "temperature")						)

 for item in rename_keys:
  if item[0] in original_param_names:
   _UpperCamelCase      :	Optional[int]					      =  item[1]
 return key_mapping






def 				__lowerCAmelCase					(					__lowerCAmelCase			:		Optional[Any]  ,	__lowerCAmelCase			:		Tuple  ,	__lowerCAmelCase			:		Optional[int]						)    ->       Union[str, Any]:
 list(hf_params.keys()						)

 for key, value in tf_params.items():
  if key not in key_mapping:
   continue

  _UpperCamelCase      :	List[Any]					      =  key_mapping[key]
  if "_conv" in key and "kernel" in key:
   _UpperCamelCase      :	Dict					      =  torch.from_numpy(__lowerCAmelCase						).permute(3  ,	2  ,	0  ,	1						)
  elif "embeddings" in key:
   _UpperCamelCase      :	Tuple					      =  torch.from_numpy(__lowerCAmelCase						)
  elif "depthwise_kernel" in key:
   _UpperCamelCase      :	Tuple					      =  torch.from_numpy(__lowerCAmelCase						).permute(2  ,	3  ,	0  ,	1						)
  elif "kernel" in key:
   _UpperCamelCase      :	Dict					      =  torch.from_numpy(np.transpose(__lowerCAmelCase						)						)
  elif "temperature" in key:
   _UpperCamelCase      :	Union[str, Any]					      =  value
  elif "bn/gamma" or "bn/beta" in key:
   _UpperCamelCase      :	str					      =  torch.from_numpy(np.transpose(__lowerCAmelCase						)						).squeeze()
  else:
   _UpperCamelCase      :	str					      =  torch.from_numpy(__lowerCAmelCase						)

  # Replace HF parameters with original TF model parameters
  hf_params[hf_key].copy_(__lowerCAmelCase						)






@torch.no_grad()
def 				__lowerCAmelCase					(					__lowerCAmelCase			:		int  ,	__lowerCAmelCase			:		int  ,	__lowerCAmelCase			:		Optional[int]  ,	__lowerCAmelCase			:		str						)    ->       Tuple:
 _UpperCamelCase      :	Dict					      =  64
 _UpperCamelCase      :	Tuple					      =  Tokenizer(__lowerCAmelCase						)
 _UpperCamelCase      :	Tuple					      =  align.Align("efficientnet-b7"  ,	"bert-base"  ,	640  ,	__lowerCAmelCase  ,	tok.get_vocab_size()						)
 original_model.compile()
 original_model.load_weights(__lowerCAmelCase						)

 _UpperCamelCase      :	Optional[Any]					      =  original_model.trainable_variables
 _UpperCamelCase      :	Any					      =  original_model.non_trainable_variables
 _UpperCamelCase      :	List[Any]					      =  {param.name: param.numpy() for param in tf_params}
 for param in tf_non_train_params:
  _UpperCamelCase      :	Optional[Any]					      =  param.numpy()
 _UpperCamelCase      :	Tuple					      =  list(tf_params.keys()						)

 # Load HuggingFace model
 _UpperCamelCase      :	Union[str, Any]					      =  get_align_config()
 _UpperCamelCase      :	Optional[int]					      =  AlignModel(__lowerCAmelCase						).eval()
 _UpperCamelCase      :	Dict					      =  hf_model.state_dict()

 # Create src-to-dst parameter name mapping dictionary
 print("Converting parameters..."						)
 _UpperCamelCase      :	Optional[int]					      =  rename_keys(__lowerCAmelCase						)
 replace_params(__lowerCAmelCase  ,	__lowerCAmelCase  ,	__lowerCAmelCase						)

 # Initialize processor
 _UpperCamelCase      :	Union[str, Any]					      =  get_processor()
 _UpperCamelCase      :	Dict					      =  processor(
     images=prepare_img()  ,	text="A picture of a cat"  ,	padding="max_length"  ,	max_length=64  ,	return_tensors="pt"						)

 # HF model inference
 hf_model.eval()
 with torch.no_grad():
  _UpperCamelCase      :	Union[str, Any]					      =  hf_model(**__lowerCAmelCase						)

 _UpperCamelCase      :	Union[str, Any]					      =  outputs.image_embeds.detach().numpy()
 _UpperCamelCase      :	int					      =  outputs.text_embeds.detach().numpy()

 # Original model inference
 _UpperCamelCase      :	Tuple					      =  False
 _UpperCamelCase      :	Dict					      =  EfficientNetImageProcessor(
     do_center_crop=__lowerCAmelCase  ,	do_rescale=__lowerCAmelCase  ,	do_normalize=__lowerCAmelCase  ,	include_top=__lowerCAmelCase  ,	resample=Image.BILINEAR  ,	)
 _UpperCamelCase      :	Optional[int]					      =  tf_image_processor(images=prepare_img()  ,	return_tensors="tf"  ,	data_format="channels_last"						)["pixel_values"]
 _UpperCamelCase      :	Union[str, Any]					      =  tok(tf.constant(["A picture of a cat"]						)						)

 _UpperCamelCase      :	Any					      =  original_model.image_encoder(__lowerCAmelCase  ,	training=__lowerCAmelCase						)
 _UpperCamelCase      :	Dict					      =  original_model.text_encoder(__lowerCAmelCase  ,	training=__lowerCAmelCase						)

 _UpperCamelCase      :	Union[str, Any]					      =  tf.nn.la_normalize(__lowerCAmelCase  ,	axis=-1						)
 _UpperCamelCase      :	List[Any]					      =  tf.nn.la_normalize(__lowerCAmelCase  ,	axis=-1						)

 # Check whether original and HF model outputs match  -> np.allclose
 if not np.allclose(__lowerCAmelCase  ,	__lowerCAmelCase  ,	atol=1e-3						):
  raise ValueError("The predicted image features are not the same."						)
 if not np.allclose(__lowerCAmelCase  ,	__lowerCAmelCase  ,	atol=1e-3						):
  raise ValueError("The predicted text features are not the same."						)
 print("Model outputs match!"						)

 if save_model:
  # Create folder to save model
  if not os.path.isdir(__lowerCAmelCase						):
   os.mkdir(__lowerCAmelCase						)
  # Save converted model and image processor
  hf_model.save_pretrained(__lowerCAmelCase						)
  processor.save_pretrained(__lowerCAmelCase						)

 if push_to_hub:
  # Push model and image processor to hub
  print("Pushing converted ALIGN to the hub..."						)
  processor.push_to_hub("align-base"						)
  hf_model.push_to_hub("align-base"						)


if __name__ == "__main__":
    _SCREAMING_SNAKE_CASE						 =     argparse.ArgumentParser()
    # Required parameters
    parser.add_argument(
        """--checkpoint_path""",
        default="""./weights/model-weights""",
        type=str,
        help="""Path to the pretrained TF ALIGN checkpoint.""",
    )
    parser.add_argument(
        """--pytorch_dump_folder_path""",
        default="""hf_model""",
        type=str,
        help="""Path to the output PyTorch model directory.""",
    )
    parser.add_argument("""--save_model""", action="""store_true""", help="""Save model to local""")
    parser.add_argument("""--push_to_hub""", action="""store_true""", help="""Push model and image processor to the hub""")

    _SCREAMING_SNAKE_CASE						 =     parser.parse_args()
    convert_align_checkpoint(args.checkpoint_path, args.pytorch_dump_folder_path, args.save_model, args.push_to_hub)









"""simple docstring"""
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  _SCREAMING_SNAKE_CASE       (     nn.Module				):



 '''simple docstring'''





 def __init__(self					,     lowerCAmelCase__					,     lowerCAmelCase__=1.0   ):

  '''simple docstring'''







  super().__init__()

  _UpperCamelCase      :	Optional[Any]					      =  cut_size
  _UpperCamelCase      :	List[str]					      =  cut_power






 def lowercase_      (self					,     lowerCAmelCase__					,     lowerCAmelCase__   ):

  '''simple docstring'''







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






def 				__lowerCAmelCase					(					__lowerCAmelCase			:		str  ,	__lowerCAmelCase			:		Optional[int]						)    ->       Union[str, Any]:
 _UpperCamelCase      :	List[str]					      =  F.normalize(__lowerCAmelCase  ,	dim=-1						)
 _UpperCamelCase      :	Optional[Any]					      =  F.normalize(__lowerCAmelCase  ,	dim=-1						)
 return (x - y).norm(dim=-1						).div(2						).arcsin().pow(2						).mul(2						)






def 				__lowerCAmelCase					(					__lowerCAmelCase			:		Optional[Any]  ,	__lowerCAmelCase			:		str						)    ->       List[str]:
 for param in model.parameters():
  _UpperCamelCase      :	Any					      =  value







class  _SCREAMING_SNAKE_CASE       (     SCREAMING_SNAKE_CASE__				):



 '''simple docstring'''





 def __init__(self					,     lowerCAmelCase__					,     lowerCAmelCase__					,     lowerCAmelCase__					,     lowerCAmelCase__					,     lowerCAmelCase__					,     lowerCAmelCase__					,     lowerCAmelCase__					,     ):

  '''simple docstring'''







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

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

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






 def lowercase_      (self					,     lowerCAmelCase__ = "auto"   ):

  '''simple docstring'''







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






 def lowercase_      (self   ):

  '''simple docstring'''







  self.enable_attention_slicing(lowerCAmelCase__   )






 def lowercase_      (self   ):

  '''simple docstring'''







  set_requires_grad(self.vae					,     lowerCAmelCase__   )






 def lowercase_      (self   ):

  '''simple docstring'''







  set_requires_grad(self.vae					,     lowerCAmelCase__   )






 def lowercase_      (self   ):

  '''simple docstring'''







  set_requires_grad(self.unet					,     lowerCAmelCase__   )






 def lowercase_      (self   ):

  '''simple docstring'''







  set_requires_grad(self.unet					,     lowerCAmelCase__   )






 @torch.enable_grad()
 def lowercase_      (self					,     lowerCAmelCase__					,     lowerCAmelCase__					,     lowerCAmelCase__					,     lowerCAmelCase__					,     lowerCAmelCase__					,     lowerCAmelCase__					,     lowerCAmelCase__					,     lowerCAmelCase__					,     lowerCAmelCase__=True					,     ):

  '''simple docstring'''







  _UpperCamelCase      :	Dict					      =  latents.detach().requires_grad_()

  _UpperCamelCase      :	Optional[Any]					      =  self.scheduler.scale_model_input(lowerCAmelCase__					,     lowerCAmelCase__   )

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

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

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

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

  if use_cutouts:
   _UpperCamelCase      :	Dict					      =  self.make_cutouts(lowerCAmelCase__					,     lowerCAmelCase__   )
  else:
   _UpperCamelCase      :	List[str]					      =  transforms.Resize(self.cut_out_size   )(lowerCAmelCase__   )
  _UpperCamelCase      :	str					      =  self.normalize(lowerCAmelCase__   ).to(latents.dtype   )

  _UpperCamelCase      :	List[Any]					      =  self.clip_model.get_image_features(lowerCAmelCase__   )
  _UpperCamelCase      :	Dict					      =  image_embeddings_clip / image_embeddings_clip.norm(p=2					,     dim=-1					,     keepdim=lowerCAmelCase__   )

  if use_cutouts:
   _UpperCamelCase      :	str					      =  spherical_dist_loss(lowerCAmelCase__					,     lowerCAmelCase__   )
   _UpperCamelCase      :	str					      =  dists.view([num_cutouts, sample.shape[0], -1]   )
   _UpperCamelCase      :	Optional[Any]					      =  dists.sum(2   ).mean(0   ).sum() * clip_guidance_scale
  else:
   _UpperCamelCase      :	Any					      =  spherical_dist_loss(lowerCAmelCase__					,     lowerCAmelCase__   ).mean() * clip_guidance_scale

  _UpperCamelCase      :	List[str]					      =  -torch.autograd.grad(lowerCAmelCase__					,     lowerCAmelCase__   )[0]

  if isinstance(self.scheduler					,     lowerCAmelCase__   ):
   _UpperCamelCase      :	Tuple					      =  latents.detach() + grads * (sigma**2)
   _UpperCamelCase      :	int					      =  noise_pred_original
  else:
   _UpperCamelCase      :	Union[str, Any]					      =  noise_pred_original - torch.sqrt(lowerCAmelCase__   ) * grads
  return noise_pred, latents






 @torch.no_grad()
 def __call__(self					,     lowerCAmelCase__					,     lowerCAmelCase__ = 5_12					,     lowerCAmelCase__ = 5_12					,     lowerCAmelCase__ = 50					,     lowerCAmelCase__ = 7.5					,     lowerCAmelCase__ = 1					,     lowerCAmelCase__ = 0.0					,     lowerCAmelCase__ = 1_00					,     lowerCAmelCase__ = None					,     lowerCAmelCase__ = 4					,     lowerCAmelCase__ = True					,     lowerCAmelCase__ = None					,     lowerCAmelCase__ = None					,     lowerCAmelCase__ = "pil"					,     lowerCAmelCase__ = True					,     ):

  '''simple docstring'''







  if isinstance(lowerCAmelCase__					,     lowerCAmelCase__   ):
   _UpperCamelCase      :	str					      =  1
  elif isinstance(lowerCAmelCase__					,     lowerCAmelCase__   ):
   _UpperCamelCase      :	Dict					      =  len(lowerCAmelCase__   )
  else:
   raise ValueError(F"`prompt` has to be of type `str` or `list` but is {type(lowerCAmelCase__   )}"   )

  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
  _UpperCamelCase      :	Dict					      =  self.tokenizer(
      lowerCAmelCase__					,     padding="max_length"					,     max_length=self.tokenizer.model_max_length					,     truncation=lowerCAmelCase__					,     return_tensors="pt"					,     )
  _UpperCamelCase      :	List[str]					      =  self.text_encoder(text_input.input_ids.to(self.device   )   )[0]
  # duplicate text embeddings for each generation per prompt
  _UpperCamelCase      :	str					      =  text_embeddings.repeat_interleave(lowerCAmelCase__					,     dim=0   )

  if clip_guidance_scale > 0:
   if clip_prompt is not None:
    _UpperCamelCase      :	List[Any]					      =  self.tokenizer(
        lowerCAmelCase__					,     padding="max_length"					,     max_length=self.tokenizer.model_max_length					,     truncation=lowerCAmelCase__					,     return_tensors="pt"					,     ).input_ids.to(self.device   )
   else:
    _UpperCamelCase      :	str					      =  text_input.input_ids.to(self.device   )
   _UpperCamelCase      :	Dict					      =  self.clip_model.get_text_features(lowerCAmelCase__   )
   _UpperCamelCase      :	Optional[int]					      =  text_embeddings_clip / text_embeddings_clip.norm(p=2					,     dim=-1					,     keepdim=lowerCAmelCase__   )
   # duplicate text embeddings clip for each generation per prompt
   _UpperCamelCase      :	int					      =  text_embeddings_clip.repeat_interleave(lowerCAmelCase__					,     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.
  _UpperCamelCase      :	Optional[Any]					      =  guidance_scale > 1.0
  # get unconditional embeddings for classifier free guidance
  if do_classifier_free_guidance:
   _UpperCamelCase      :	Optional[int]					      =  text_input.input_ids.shape[-1]
   _UpperCamelCase      :	List[str]					      =  self.tokenizer([""]					,     padding="max_length"					,     max_length=lowerCAmelCase__					,     return_tensors="pt"   )
   _UpperCamelCase      :	List[str]					      =  self.text_encoder(uncond_input.input_ids.to(self.device   )   )[0]
   # duplicate unconditional embeddings for each generation per prompt
   _UpperCamelCase      :	Dict					      =  uncond_embeddings.repeat_interleave(lowerCAmelCase__					,     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
   _UpperCamelCase      :	Tuple					      =  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`.
  _UpperCamelCase      :	Optional[Any]					      =  (batch_size * num_images_per_prompt, self.unet.config.in_channels, height // 8, width // 8)
  _UpperCamelCase      :	Optional[Any]					      =  text_embeddings.dtype
  if latents is None:
   if self.device.type == "mps":
    # randn does not work reproducibly on mps
    _UpperCamelCase      :	Optional[Any]					      =  torch.randn(lowerCAmelCase__					,     generator=lowerCAmelCase__					,     device="cpu"					,     dtype=lowerCAmelCase__   ).to(
        self.device   )
   else:
    _UpperCamelCase      :	Union[str, Any]					      =  torch.randn(lowerCAmelCase__					,     generator=lowerCAmelCase__					,     device=self.device					,     dtype=lowerCAmelCase__   )
  else:
   if latents.shape != latents_shape:
    raise ValueError(F"Unexpected latents shape, got {latents.shape}, expected {latents_shape}"   )
   _UpperCamelCase      :	Optional[int]					      =  latents.to(self.device   )

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

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

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

  # scale the initial noise by the standard deviation required by the scheduler
  _UpperCamelCase      :	Tuple					      =  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]
  _UpperCamelCase      :	str					      =  "eta" in set(inspect.signature(self.scheduler.step   ).parameters.keys()   )
  _UpperCamelCase      :	Optional[int]					      =  {}
  if accepts_eta:
   _UpperCamelCase      :	Union[str, Any]					      =  eta

  # check if the scheduler accepts generator
  _UpperCamelCase      :	Optional[int]					      =  "generator" in set(inspect.signature(self.scheduler.step   ).parameters.keys()   )
  if accepts_generator:
   _UpperCamelCase      :	List[str]					      =  generator

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

   # predict the noise residual
   _UpperCamelCase      :	List[str]					      =  self.unet(lowerCAmelCase__					,     lowerCAmelCase__					,     encoder_hidden_states=lowerCAmelCase__   ).sample

   # perform classifier free guidance
   if do_classifier_free_guidance:
    _UpperCamelCase      ,	_UpperCamelCase      :	List[Any]					      =  noise_pred.chunk(2   )
    _UpperCamelCase      :	List[Any]					      =  noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)

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

   # compute the previous noisy sample x_t -> x_t-1
   _UpperCamelCase      :	Tuple					      =  self.scheduler.step(lowerCAmelCase__					,     lowerCAmelCase__					,     lowerCAmelCase__					,     **lowerCAmelCase__   ).prev_sample

  # scale and decode the image latents with vae
  _UpperCamelCase      :	Tuple					      =  1 / self.vae.config.scaling_factor * latents
  _UpperCamelCase      :	Union[str, Any]					      =  self.vae.decode(lowerCAmelCase__   ).sample

  _UpperCamelCase      :	List[Any]					      =  (image / 2 + 0.5).clamp(0					,     1   )
  _UpperCamelCase      :	Any					      =  image.cpu().permute(0					,     2					,     3					,     1   ).numpy()

  if output_type == "pil":
   _UpperCamelCase      :	Dict					      =  self.numpy_to_pil(lowerCAmelCase__   )

  if not return_dict:
   return (image, None)

  return StableDiffusionPipelineOutput(images=lowerCAmelCase__					,     nsfw_content_detected=lowerCAmelCase__   )



