




"""simple docstring"""
import inspect
import unittest

from transformers import is_torch_available, is_vision_available
from transformers.models.auto import get_values
from transformers.testing_utils import require_torch, require_vision, slow, torch_device

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
    import torch

    from transformers import MODEL_MAPPING, GLPNConfig, GLPNForDepthEstimation, GLPNModel
    from transformers.models.glpn.modeling_glpn import GLPN_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
    from PIL import Image

    from transformers import GLPNImageProcessor







class  _SCREAMING_SNAKE_CASE       (     SCREAMING_SNAKE_CASE__				):



 '''simple docstring'''





 def lowercase_      (self   ):

  '''simple docstring'''







  _UpperCamelCase      :	Dict					      =  self.config_class(**self.inputs_dict   )
  self.parent.assertTrue(hasattr(lowerCAmelCase__					,     "hidden_sizes"   )   )
  self.parent.assertTrue(hasattr(lowerCAmelCase__					,     "num_attention_heads"   )   )
  self.parent.assertTrue(hasattr(lowerCAmelCase__					,     "num_encoder_blocks"   )   )







class  _SCREAMING_SNAKE_CASE       :



 '''simple docstring'''





 def __init__(self					,     lowerCAmelCase__					,     lowerCAmelCase__=13					,     lowerCAmelCase__=64					,     lowerCAmelCase__=3					,     lowerCAmelCase__=4					,     lowerCAmelCase__=[2, 2, 2, 2]					,     lowerCAmelCase__=[8, 4, 2, 1]					,     lowerCAmelCase__=[16, 32, 64, 1_28]					,     lowerCAmelCase__=[1, 4, 8, 16]					,     lowerCAmelCase__=[1, 2, 4, 8]					,     lowerCAmelCase__=True					,     lowerCAmelCase__=True					,     lowerCAmelCase__="gelu"					,     lowerCAmelCase__=0.1					,     lowerCAmelCase__=0.1					,     lowerCAmelCase__=0.02					,     lowerCAmelCase__=16					,     lowerCAmelCase__=3					,     lowerCAmelCase__=None					,     ):

  '''simple docstring'''







  _UpperCamelCase      :	Union[str, Any]					      =  parent
  _UpperCamelCase      :	Any					      =  batch_size
  _UpperCamelCase      :	Any					      =  image_size
  _UpperCamelCase      :	Optional[Any]					      =  num_channels
  _UpperCamelCase      :	List[str]					      =  num_encoder_blocks
  _UpperCamelCase      :	Union[str, Any]					      =  sr_ratios
  _UpperCamelCase      :	List[str]					      =  depths
  _UpperCamelCase      :	Any					      =  hidden_sizes
  _UpperCamelCase      :	Any					      =  downsampling_rates
  _UpperCamelCase      :	Any					      =  num_attention_heads
  _UpperCamelCase      :	Optional[Any]					      =  is_training
  _UpperCamelCase      :	Optional[Any]					      =  use_labels
  _UpperCamelCase      :	Dict					      =  hidden_act
  _UpperCamelCase      :	Optional[Any]					      =  hidden_dropout_prob
  _UpperCamelCase      :	Union[str, Any]					      =  attention_probs_dropout_prob
  _UpperCamelCase      :	str					      =  initializer_range
  _UpperCamelCase      :	Union[str, Any]					      =  decoder_hidden_size
  _UpperCamelCase      :	List[str]					      =  num_labels
  _UpperCamelCase      :	Dict					      =  scope






 def lowercase_      (self   ):

  '''simple docstring'''







  _UpperCamelCase      :	int					      =  floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]   )

  _UpperCamelCase      :	List[Any]					      =  None
  if self.use_labels:
   _UpperCamelCase      :	Optional[Any]					      =  ids_tensor([self.batch_size, self.image_size, self.image_size]					,     self.num_labels   )

  _UpperCamelCase      :	Any					      =  self.get_config()
  return config, pixel_values, labels






 def lowercase_      (self   ):

  '''simple docstring'''







  return GLPNConfig(
      image_size=self.image_size					,     num_channels=self.num_channels					,     num_encoder_blocks=self.num_encoder_blocks					,     depths=self.depths					,     hidden_sizes=self.hidden_sizes					,     num_attention_heads=self.num_attention_heads					,     hidden_act=self.hidden_act					,     hidden_dropout_prob=self.hidden_dropout_prob					,     attention_probs_dropout_prob=self.attention_probs_dropout_prob					,     initializer_range=self.initializer_range					,     decoder_hidden_size=self.decoder_hidden_size					,     )






 def lowercase_      (self					,     lowerCAmelCase__					,     lowerCAmelCase__					,     lowerCAmelCase__   ):

  '''simple docstring'''







  _UpperCamelCase      :	int					      =  GLPNModel(config=lowerCAmelCase__   )
  model.to(lowerCAmelCase__   )
  model.eval()
  _UpperCamelCase      :	Optional[int]					      =  model(lowerCAmelCase__   )
  _UpperCamelCase      :	Optional[Any]					      =  self.image_size // (self.downsampling_rates[-1] * 2)
  self.parent.assertEqual(
      result.last_hidden_state.shape					,     (self.batch_size, self.hidden_sizes[-1], expected_height, expected_width)   )






 def lowercase_      (self					,     lowerCAmelCase__					,     lowerCAmelCase__					,     lowerCAmelCase__   ):

  '''simple docstring'''







  _UpperCamelCase      :	List[str]					      =  self.num_labels
  _UpperCamelCase      :	Optional[int]					      =  GLPNForDepthEstimation(lowerCAmelCase__   )
  model.to(lowerCAmelCase__   )
  model.eval()
  _UpperCamelCase      :	Tuple					      =  model(lowerCAmelCase__   )
  self.parent.assertEqual(result.predicted_depth.shape					,     (self.batch_size, self.image_size, self.image_size)   )
  _UpperCamelCase      :	Any					      =  model(lowerCAmelCase__					,     labels=lowerCAmelCase__   )
  self.parent.assertEqual(result.predicted_depth.shape					,     (self.batch_size, self.image_size, self.image_size)   )






 def lowercase_      (self   ):

  '''simple docstring'''







  _UpperCamelCase      :	int					      =  self.prepare_config_and_inputs()
  _UpperCamelCase      ,	_UpperCamelCase      ,	_UpperCamelCase      :	Any					      =  config_and_inputs
  _UpperCamelCase      :	int					      =  {"pixel_values": pixel_values}
  return config, inputs_dict







@require_torch
class  _SCREAMING_SNAKE_CASE       (     SCREAMING_SNAKE_CASE__			,				SCREAMING_SNAKE_CASE__			,				unittest.TestCase				):



 '''simple docstring'''





 __UpperCAmelCase		   =							(GLPNModel, GLPNForDepthEstimation) if is_torch_available() else ()
 __UpperCAmelCase		   =							(
     {"""depth-estimation""": GLPNForDepthEstimation, """feature-extraction""": GLPNModel} if is_torch_available() else {}
 )

 __UpperCAmelCase		   =							False
 __UpperCAmelCase		   =							False
 __UpperCAmelCase		   =							False






 def lowercase_      (self   ):

  '''simple docstring'''







  _UpperCamelCase      :	List[Any]					      =  GLPNModelTester(self   )
  _UpperCamelCase      :	Any					      =  GLPNConfigTester(self					,     config_class=lowerCAmelCase__   )






 def lowercase_      (self   ):

  '''simple docstring'''







  self.config_tester.run_common_tests()






 def lowercase_      (self   ):

  '''simple docstring'''







  _UpperCamelCase      :	str					      =  self.model_tester.prepare_config_and_inputs()
  self.model_tester.create_and_check_model(*lowerCAmelCase__   )






 def lowercase_      (self   ):

  '''simple docstring'''







  _UpperCamelCase      :	List[str]					      =  self.model_tester.prepare_config_and_inputs()
  self.model_tester.create_and_check_for_depth_estimation(*lowerCAmelCase__   )






 @unittest.skip("GLPN does not use inputs_embeds"   )
 def lowercase_      (self   ):

  '''simple docstring'''







  pass






 @unittest.skip("GLPN does not have get_input_embeddings method and get_output_embeddings methods"   )
 def lowercase_      (self   ):

  '''simple docstring'''







  pass






 def lowercase_      (self   ):

  '''simple docstring'''







  _UpperCamelCase      ,	_UpperCamelCase      :	List[str]					      =  self.model_tester.prepare_config_and_inputs_for_common()

  for model_class in self.all_model_classes:
   _UpperCamelCase      :	Any					      =  model_class(lowerCAmelCase__   )
   _UpperCamelCase      :	Tuple					      =  inspect.signature(model.forward   )
   # signature.parameters is an OrderedDict => so arg_names order is deterministic
   _UpperCamelCase      :	Any					      =  [*signature.parameters.keys()]

   _UpperCamelCase      :	Dict					      =  ["pixel_values"]
   self.assertListEqual(arg_names[:1]					,     lowerCAmelCase__   )






 def lowercase_      (self   ):

  '''simple docstring'''







  _UpperCamelCase      ,	_UpperCamelCase      :	List[Any]					      =  self.model_tester.prepare_config_and_inputs_for_common()
  _UpperCamelCase      :	Any					      =  True

  for model_class in self.all_model_classes:
   _UpperCamelCase      :	Tuple					      =  True
   _UpperCamelCase      :	List[str]					      =  False
   _UpperCamelCase      :	Union[str, Any]					      =  True
   _UpperCamelCase      :	List[Any]					      =  model_class(lowerCAmelCase__   )
   model.to(lowerCAmelCase__   )
   model.eval()
   with torch.no_grad():
    _UpperCamelCase      :	Any					      =  model(**self._prepare_for_class(lowerCAmelCase__					,     lowerCAmelCase__   )   )
   _UpperCamelCase      :	Optional[Any]					      =  outputs.attentions

   _UpperCamelCase      :	Tuple					      =  sum(self.model_tester.depths   )
   self.assertEqual(len(lowerCAmelCase__   )					,     lowerCAmelCase__   )

   # check that output_attentions also work using config
   del inputs_dict["output_attentions"]
   _UpperCamelCase      :	int					      =  True
   _UpperCamelCase      :	Dict					      =  model_class(lowerCAmelCase__   )
   model.to(lowerCAmelCase__   )
   model.eval()
   with torch.no_grad():
    _UpperCamelCase      :	Any					      =  model(**self._prepare_for_class(lowerCAmelCase__					,     lowerCAmelCase__   )   )
   _UpperCamelCase      :	List[Any]					      =  outputs.attentions

   self.assertEqual(len(lowerCAmelCase__   )					,     lowerCAmelCase__   )

   # verify the first attentions (first block, first layer)
   _UpperCamelCase      :	str					      =  (self.model_tester.image_size // 4) ** 2
   _UpperCamelCase      :	Union[str, Any]					      =  (self.model_tester.image_size // (4 * self.model_tester.sr_ratios[0])) ** 2
   self.assertListEqual(
       list(attentions[0].shape[-3:]   )					,     [self.model_tester.num_attention_heads[0], expected_seq_len, expected_reduced_seq_len]					,     )

   # verify the last attentions (last block, last layer)
   _UpperCamelCase      :	str					      =  (self.model_tester.image_size // 32) ** 2
   _UpperCamelCase      :	Optional[int]					      =  (self.model_tester.image_size // (32 * self.model_tester.sr_ratios[-1])) ** 2
   self.assertListEqual(
       list(attentions[-1].shape[-3:]   )					,     [self.model_tester.num_attention_heads[-1], expected_seq_len, expected_reduced_seq_len]					,     )
   _UpperCamelCase      :	Dict					      =  len(lowerCAmelCase__   )

   # Check attention is always last and order is fine
   _UpperCamelCase      :	List[str]					      =  True
   _UpperCamelCase      :	Optional[Any]					      =  True
   _UpperCamelCase      :	Dict					      =  model_class(lowerCAmelCase__   )
   model.to(lowerCAmelCase__   )
   model.eval()
   with torch.no_grad():
    _UpperCamelCase      :	Optional[int]					      =  model(**self._prepare_for_class(lowerCAmelCase__					,     lowerCAmelCase__   )   )

   self.assertEqual(out_len + 1					,     len(lowerCAmelCase__   )   )

   _UpperCamelCase      :	List[str]					      =  outputs.attentions

   self.assertEqual(len(lowerCAmelCase__   )					,     lowerCAmelCase__   )
   # verify the first attentions (first block, first layer)
   _UpperCamelCase      :	Optional[Any]					      =  (self.model_tester.image_size // 4) ** 2
   _UpperCamelCase      :	List[Any]					      =  (self.model_tester.image_size // (4 * self.model_tester.sr_ratios[0])) ** 2
   self.assertListEqual(
       list(self_attentions[0].shape[-3:]   )					,     [self.model_tester.num_attention_heads[0], expected_seq_len, expected_reduced_seq_len]					,     )






 def lowercase_      (self   ):

  '''simple docstring'''







  def check_hidden_states_output(lowerCAmelCase__					,     lowerCAmelCase__					,     lowerCAmelCase__   ):
   _UpperCamelCase      :	Any					      =  model_class(lowerCAmelCase__   )
   model.to(lowerCAmelCase__   )
   model.eval()

   with torch.no_grad():
    _UpperCamelCase      :	int					      =  model(**self._prepare_for_class(lowerCAmelCase__					,     lowerCAmelCase__   )   )

   _UpperCamelCase      :	Union[str, Any]					      =  outputs.hidden_states

   _UpperCamelCase      :	Dict					      =  self.model_tester.num_encoder_blocks
   self.assertEqual(len(lowerCAmelCase__   )					,     lowerCAmelCase__   )

   # verify the first hidden states (first block)
   self.assertListEqual(
       list(hidden_states[0].shape[-3:]   )					,     [
           self.model_tester.hidden_sizes[0],
           self.model_tester.image_size // 4,
           self.model_tester.image_size // 4,
       ]					,     )

  _UpperCamelCase      ,	_UpperCamelCase      :	Optional[int]					      =  self.model_tester.prepare_config_and_inputs_for_common()

  for model_class in self.all_model_classes:
   _UpperCamelCase      :	int					      =  True
   check_hidden_states_output(lowerCAmelCase__					,     lowerCAmelCase__					,     lowerCAmelCase__   )

   # check that output_hidden_states also work using config
   del inputs_dict["output_hidden_states"]
   _UpperCamelCase      :	Optional[Any]					      =  True

   check_hidden_states_output(lowerCAmelCase__					,     lowerCAmelCase__					,     lowerCAmelCase__   )






 def lowercase_      (self   ):

  '''simple docstring'''







  if not self.model_tester.is_training:
   return

  _UpperCamelCase      ,	_UpperCamelCase      :	int					      =  self.model_tester.prepare_config_and_inputs_for_common()
  _UpperCamelCase      :	List[str]					      =  True

  for model_class in self.all_model_classes:
   if model_class in get_values(lowerCAmelCase__   ):
    continue
   # TODO: remove the following 3 lines once we have a MODEL_FOR_DEPTH_ESTIMATION_MAPPING
   # this can then be incorporated into _prepare_for_class in test_modeling_common.py
   if model_class.__name__ == "GLPNForDepthEstimation":
    _UpperCamelCase      ,	_UpperCamelCase      ,	_UpperCamelCase      ,	_UpperCamelCase      :	Optional[int]					      =  inputs_dict["pixel_values"].shape
    _UpperCamelCase      :	Optional[Any]					      =  torch.zeros(
        [self.model_tester.batch_size, height, width]					,     device=lowerCAmelCase__   ).long()
   _UpperCamelCase      :	List[Any]					      =  model_class(lowerCAmelCase__   )
   model.to(lowerCAmelCase__   )
   model.train()
   _UpperCamelCase      :	Optional[Any]					      =  self._prepare_for_class(lowerCAmelCase__					,     lowerCAmelCase__					,     return_labels=lowerCAmelCase__   )
   _UpperCamelCase      :	Tuple					      =  model(**lowerCAmelCase__   ).loss
   loss.backward()






 @slow
 def lowercase_      (self   ):

  '''simple docstring'''







  for model_name in GLPN_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
   _UpperCamelCase      :	int					      =  GLPNModel.from_pretrained(lowerCAmelCase__   )
   self.assertIsNotNone(lowerCAmelCase__   )






def 				__lowerCAmelCase					(					)    ->       Dict:
 _UpperCamelCase      :	Optional[Any]					      =  Image.open("./tests/fixtures/tests_samples/COCO/000000039769.png"						)
 return image







@require_torch
@require_vision
@slow
class  _SCREAMING_SNAKE_CASE       (     unittest.TestCase				):



 '''simple docstring'''





 @slow
 def lowercase_      (self   ):

  '''simple docstring'''







  _UpperCamelCase      :	Optional[int]					      =  GLPNImageProcessor.from_pretrained(GLPN_PRETRAINED_MODEL_ARCHIVE_LIST[0]   )
  _UpperCamelCase      :	Optional[Any]					      =  GLPNForDepthEstimation.from_pretrained(GLPN_PRETRAINED_MODEL_ARCHIVE_LIST[0]   ).to(lowerCAmelCase__   )

  _UpperCamelCase      :	Any					      =  prepare_img()
  _UpperCamelCase      :	List[Any]					      =  image_processor(images=lowerCAmelCase__					,     return_tensors="pt"   ).to(lowerCAmelCase__   )

  # forward pass
  with torch.no_grad():
   _UpperCamelCase      :	str					      =  model(**lowerCAmelCase__   )

  # verify the predicted depth
  _UpperCamelCase      :	List[Any]					      =  torch.Size([1, 4_80, 6_40]   )
  self.assertEqual(outputs.predicted_depth.shape					,     lowerCAmelCase__   )

  _UpperCamelCase      :	Optional[Any]					      =  torch.tensor(
      [[3.4291, 2.7865, 2.5151], [3.2841, 2.7021, 2.3502], [3.1147, 2.4625, 2.2481]]   ).to(lowerCAmelCase__   )

  self.assertTrue(torch.allclose(outputs.predicted_depth[0, :3, :3]					,     lowerCAmelCase__					,     atol=1E-4   )   )









"""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__   )



