
'''simple docstring'''



from __future__ import annotations

import inspect
import unittest

import numpy as np

from transformers import SwinConfig
from transformers.testing_utils import require_tf, require_vision, slow, to_atuple
from transformers.utils import cached_property, is_tf_available, is_vision_available

from ...test_configuration_common import ConfigTester
from ...test_modeling_tf_common import TFModelTesterMixin, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_tf_available():
     import tensorflow as tf

     from transformers.models.swin.modeling_tf_swin import (
         TF_SWIN_PRETRAINED_MODEL_ARCHIVE_LIST,
         TFSwinForImageClassification,
         TFSwinForMaskedImageModeling,
         TFSwinModel,
     )


if is_vision_available():
     from PIL import Image

     from transformers import AutoImageProcessor

class       _a :





       """simple docstring"""




       def __init__(							self				,							A__				,							A__=13				,							A__=32				,							A__=2				,							A__=3				,							A__=16				,							A__=[1, 2, 1]				,							A__=[2, 2, 4]				,							A__=2				,							A__=2.0				,							A__=True				,							A__=0.0				,							A__=0.0				,							A__=0.1				,							A__="gelu"				,							A__=False				,							A__=True				,							A__=0.02				,							A__=1E-5				,							A__=True				,							A__=None				,							A__=True				,							A__=10				,							A__=8				,							)			->	None:
          _SCREAMING_SNAKE_CASE          = parent
          _SCREAMING_SNAKE_CASE          = batch_size
          _SCREAMING_SNAKE_CASE          = image_size
          _SCREAMING_SNAKE_CASE          = patch_size
          _SCREAMING_SNAKE_CASE          = num_channels
          _SCREAMING_SNAKE_CASE          = embed_dim
          _SCREAMING_SNAKE_CASE          = depths
          _SCREAMING_SNAKE_CASE          = num_heads
          _SCREAMING_SNAKE_CASE          = window_size
          _SCREAMING_SNAKE_CASE          = mlp_ratio
          _SCREAMING_SNAKE_CASE          = qkv_bias
          _SCREAMING_SNAKE_CASE          = hidden_dropout_prob
          _SCREAMING_SNAKE_CASE          = attention_probs_dropout_prob
          _SCREAMING_SNAKE_CASE          = drop_path_rate
          _SCREAMING_SNAKE_CASE          = hidden_act
          _SCREAMING_SNAKE_CASE          = use_absolute_embeddings
          _SCREAMING_SNAKE_CASE          = patch_norm
          _SCREAMING_SNAKE_CASE          = layer_norm_eps
          _SCREAMING_SNAKE_CASE          = initializer_range
          _SCREAMING_SNAKE_CASE          = is_training
          _SCREAMING_SNAKE_CASE          = scope
          _SCREAMING_SNAKE_CASE          = use_labels
          _SCREAMING_SNAKE_CASE          = type_sequence_label_size
          _SCREAMING_SNAKE_CASE          = encoder_stride






       def 				UpperCamelCase (							self					)			->	List[Any]:
          _SCREAMING_SNAKE_CASE          = floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]					)

          _SCREAMING_SNAKE_CASE          = None
          if self.use_labels:
             _SCREAMING_SNAKE_CASE          = ids_tensor([self.batch_size]				,							self.type_sequence_label_size					)

          _SCREAMING_SNAKE_CASE          = self.get_config()

          return config, pixel_values, labels






       def 				UpperCamelCase (							self					)			->	Dict:
          return SwinConfig(
              image_size=self.image_size				,							patch_size=self.patch_size				,							num_channels=self.num_channels				,							embed_dim=self.embed_dim				,							depths=self.depths				,							num_heads=self.num_heads				,							window_size=self.window_size				,							mlp_ratio=self.mlp_ratio				,							qkv_bias=self.qkv_bias				,							hidden_dropout_prob=self.hidden_dropout_prob				,							attention_probs_dropout_prob=self.attention_probs_dropout_prob				,							drop_path_rate=self.drop_path_rate				,							hidden_act=self.hidden_act				,							use_absolute_embeddings=self.use_absolute_embeddings				,							path_norm=self.patch_norm				,							layer_norm_eps=self.layer_norm_eps				,							initializer_range=self.initializer_range				,							encoder_stride=self.encoder_stride				,							)






       def 				UpperCamelCase (							self				,							A__				,							A__				,							A__					)			->	Optional[Any]:
          _SCREAMING_SNAKE_CASE          = TFSwinModel(config=A__					)
          _SCREAMING_SNAKE_CASE          = model(A__					)

          _SCREAMING_SNAKE_CASE          = ((config.image_size // config.patch_size) ** 2) // (4 ** (len(config.depths					) - 1))
          _SCREAMING_SNAKE_CASE          = int(config.embed_dim * 2 ** (len(config.depths					) - 1)					)

          self.parent.assertEqual(result.last_hidden_state.shape				,							(self.batch_size, expected_seq_len, expected_dim)					)






       def 				UpperCamelCase (							self				,							A__				,							A__				,							A__					)			->	Optional[int]:
          _SCREAMING_SNAKE_CASE          = TFSwinForMaskedImageModeling(config=A__					)
          _SCREAMING_SNAKE_CASE          = model(A__					)
          self.parent.assertEqual(
              result.logits.shape				,							(self.batch_size, self.num_channels, self.image_size, self.image_size)					)

          # test greyscale images
          _SCREAMING_SNAKE_CASE          = 1
          _SCREAMING_SNAKE_CASE          = TFSwinForMaskedImageModeling(A__					)

          _SCREAMING_SNAKE_CASE          = floats_tensor([self.batch_size, 1, self.image_size, self.image_size]					)
          _SCREAMING_SNAKE_CASE          = model(A__					)
          self.parent.assertEqual(result.logits.shape				,							(self.batch_size, 1, self.image_size, self.image_size)					)






       def 				UpperCamelCase (							self				,							A__				,							A__				,							A__					)			->	Union[str, Any]:
          _SCREAMING_SNAKE_CASE          = self.type_sequence_label_size
          _SCREAMING_SNAKE_CASE          = TFSwinForImageClassification(A__					)
          _SCREAMING_SNAKE_CASE          = model(A__				,							labels=A__					)
          self.parent.assertEqual(result.logits.shape				,							(self.batch_size, self.type_sequence_label_size)					)

          # test greyscale images
          _SCREAMING_SNAKE_CASE          = 1
          _SCREAMING_SNAKE_CASE          = TFSwinForImageClassification(A__					)
          _SCREAMING_SNAKE_CASE          = floats_tensor([self.batch_size, 1, self.image_size, self.image_size]					)
          _SCREAMING_SNAKE_CASE          = model(A__					)
          self.parent.assertEqual(result.logits.shape				,							(self.batch_size, self.type_sequence_label_size)					)






       def 				UpperCamelCase (							self					)			->	Tuple:
          _SCREAMING_SNAKE_CASE          = self.prepare_config_and_inputs()
          _SCREAMING_SNAKE_CASE       ,							_SCREAMING_SNAKE_CASE       ,							_SCREAMING_SNAKE_CASE          = config_and_inputs
          _SCREAMING_SNAKE_CASE          = {"""pixel_values""": pixel_values}
          return config, inputs_dict






@require_tf
class       _a (_lowerCamelCase     ,    _lowerCamelCase     ,    unittest.TestCase):





       """simple docstring"""




       SCREAMING_SNAKE_CASE					=     (
           (
               TFSwinModel,
               TFSwinForImageClassification,
               TFSwinForMaskedImageModeling,
           )
           if is_tf_available()
           else ()
       )
       SCREAMING_SNAKE_CASE					=     (
           {'feature-extraction': TFSwinModel, 'image-classification': TFSwinForImageClassification}
           if is_tf_available()
           else {}
       )

       SCREAMING_SNAKE_CASE					=     False
       SCREAMING_SNAKE_CASE					=     False
       SCREAMING_SNAKE_CASE					=     False
       SCREAMING_SNAKE_CASE					=     False






       def 				UpperCamelCase (							self					)			->	List[str]:
          _SCREAMING_SNAKE_CASE          = TFSwinModelTester(self					)
          _SCREAMING_SNAKE_CASE          = ConfigTester(self				,							config_class=A__				,							embed_dim=37					)






       def 				UpperCamelCase (							self					)			->	Dict:
          self.create_and_test_config_common_properties()
          self.config_tester.create_and_test_config_to_json_string()
          self.config_tester.create_and_test_config_to_json_file()
          self.config_tester.create_and_test_config_from_and_save_pretrained()
          self.config_tester.create_and_test_config_with_num_labels()
          self.config_tester.check_config_can_be_init_without_params()
          self.config_tester.check_config_arguments_init()






       def 				UpperCamelCase (							self					)			->	Optional[Any]:
          return






       def 				UpperCamelCase (							self					)			->	List[Any]:
          _SCREAMING_SNAKE_CASE          = self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_model(*A__					)






       def 				UpperCamelCase (							self					)			->	Union[str, Any]:
          _SCREAMING_SNAKE_CASE          = self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_for_masked_image_modeling(*A__					)






       def 				UpperCamelCase (							self					)			->	List[Any]:
          _SCREAMING_SNAKE_CASE          = self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_for_image_classification(*A__					)






       @unittest.skip(reason="""Swin does not use inputs_embeds"""					)
       def 				UpperCamelCase (							self					)			->	int:
          pass






       def 				UpperCamelCase (							self					)			->	Any:
          _SCREAMING_SNAKE_CASE       ,							_SCREAMING_SNAKE_CASE          = self.model_tester.prepare_config_and_inputs_for_common()

          for model_class in self.all_model_classes:
             _SCREAMING_SNAKE_CASE          = model_class(A__					)
             self.assertIsInstance(model.get_input_embeddings()				,							tf.keras.layers.Layer					)
             _SCREAMING_SNAKE_CASE          = model.get_output_embeddings()
             self.assertTrue(x is None or isinstance(A__				,							tf.keras.layers.Dense					)					)






       def 				UpperCamelCase (							self					)			->	List[Any]:
          _SCREAMING_SNAKE_CASE       ,							_SCREAMING_SNAKE_CASE          = self.model_tester.prepare_config_and_inputs_for_common()

          for model_class in self.all_model_classes:
             _SCREAMING_SNAKE_CASE          = model_class(A__					)
             _SCREAMING_SNAKE_CASE          = inspect.signature(model.call					)
             # signature.parameters is an OrderedDict => so arg_names order is deterministic
             _SCREAMING_SNAKE_CASE          = [*signature.parameters.keys()]

             _SCREAMING_SNAKE_CASE          = ["""pixel_values"""]
             self.assertListEqual(arg_names[:1]				,							A__					)






       def 				UpperCamelCase (							self					)			->	Optional[int]:
          _SCREAMING_SNAKE_CASE       ,							_SCREAMING_SNAKE_CASE          = self.model_tester.prepare_config_and_inputs_for_common()
          _SCREAMING_SNAKE_CASE          = True

          for model_class in self.all_model_classes:
             _SCREAMING_SNAKE_CASE          = True
             _SCREAMING_SNAKE_CASE          = False
             _SCREAMING_SNAKE_CASE          = True
             _SCREAMING_SNAKE_CASE          = model_class(A__					)
             _SCREAMING_SNAKE_CASE          = model(**self._prepare_for_class(A__				,							A__					)					)
             _SCREAMING_SNAKE_CASE          = outputs.attentions
             _SCREAMING_SNAKE_CASE          = len(self.model_tester.depths					)
             self.assertEqual(len(A__					)				,							A__					)

             # check that output_attentions also work using config
             del inputs_dict["output_attentions"]
             _SCREAMING_SNAKE_CASE          = True
             _SCREAMING_SNAKE_CASE          = config.window_size**2
             _SCREAMING_SNAKE_CASE          = model_class(A__					)
             _SCREAMING_SNAKE_CASE          = model(**self._prepare_for_class(A__				,							A__					)					)
             _SCREAMING_SNAKE_CASE          = outputs.attentions
             self.assertEqual(len(A__					)				,							A__					)

             self.assertListEqual(
                 list(attentions[0].shape[-3:]					)				,							[self.model_tester.num_heads[0], window_size_squared, window_size_squared]				,							)
             _SCREAMING_SNAKE_CASE          = len(A__					)

             # Check attention is always last and order is fine
             _SCREAMING_SNAKE_CASE          = True
             _SCREAMING_SNAKE_CASE          = True
             _SCREAMING_SNAKE_CASE          = model_class(A__					)
             _SCREAMING_SNAKE_CASE          = model(**self._prepare_for_class(A__				,							A__					)					)

             if hasattr(self.model_tester				,							"""num_hidden_states_types"""					):
                _SCREAMING_SNAKE_CASE          = self.model_tester.num_hidden_states_types
             else:
                # also another +1 for reshaped_hidden_states
                _SCREAMING_SNAKE_CASE          = 2
             self.assertEqual(out_len + added_hidden_states				,							len(A__					)					)

             _SCREAMING_SNAKE_CASE          = outputs.attentions

             self.assertEqual(len(A__					)				,							A__					)

             self.assertListEqual(
                 list(self_attentions[0].shape[-3:]					)				,							[self.model_tester.num_heads[0], window_size_squared, window_size_squared]				,							)






       def 				UpperCamelCase (							self				,							A__				,							A__				,							A__				,							A__					)			->	List[Any]:
          _SCREAMING_SNAKE_CASE          = model_class(A__					)
          _SCREAMING_SNAKE_CASE          = model(**self._prepare_for_class(A__				,							A__					)					)
          _SCREAMING_SNAKE_CASE          = outputs.hidden_states

          _SCREAMING_SNAKE_CASE          = getattr(
              self.model_tester				,							"""expected_num_hidden_layers"""				,							len(self.model_tester.depths					) + 1					)
          self.assertEqual(len(A__					)				,							A__					)

          # Swin has a different seq_length
          _SCREAMING_SNAKE_CASE          = to_atuple(config.patch_size					)

          _SCREAMING_SNAKE_CASE          = (image_size[1] // patch_size[1]) * (image_size[0] // patch_size[0])

          self.assertListEqual(
              list(hidden_states[0].shape[-2:]					)				,							[num_patches, self.model_tester.embed_dim]				,							)

          _SCREAMING_SNAKE_CASE          = outputs.reshaped_hidden_states
          self.assertEqual(len(A__					)				,							A__					)

          _SCREAMING_SNAKE_CASE       ,							_SCREAMING_SNAKE_CASE       ,							_SCREAMING_SNAKE_CASE       ,							_SCREAMING_SNAKE_CASE          = reshaped_hidden_states[0].shape

          _SCREAMING_SNAKE_CASE          = tf.reshape(reshaped_hidden_states[0]				,							(batch_size, num_channels, height * width)					)
          _SCREAMING_SNAKE_CASE          = tf.transpose(A__				,							(0, 2, 1)					)

          self.assertListEqual(
              list(reshaped_hidden_states.shape[-2:]					)				,							[num_patches, self.model_tester.embed_dim]				,							)






       def 				UpperCamelCase (							self					)			->	Optional[int]:
          _SCREAMING_SNAKE_CASE       ,							_SCREAMING_SNAKE_CASE          = self.model_tester.prepare_config_and_inputs_for_common()

          _SCREAMING_SNAKE_CASE          = to_atuple(self.model_tester.image_size					)

          for model_class in self.all_model_classes:
             _SCREAMING_SNAKE_CASE          = True
             self.check_hidden_states_output(A__				,							A__				,							A__				,							A__					)

             # check that output_hidden_states also work using config
             del inputs_dict["output_hidden_states"]
             _SCREAMING_SNAKE_CASE          = True

             self.check_hidden_states_output(A__				,							A__				,							A__				,							A__					)






       def 				UpperCamelCase (							self					)			->	Union[str, Any]:
          _SCREAMING_SNAKE_CASE       ,							_SCREAMING_SNAKE_CASE          = self.model_tester.prepare_config_and_inputs_for_common()
          _SCREAMING_SNAKE_CASE          = 3

          _SCREAMING_SNAKE_CASE          = to_atuple(self.model_tester.image_size					)
          _SCREAMING_SNAKE_CASE          = to_atuple(config.patch_size					)

          _SCREAMING_SNAKE_CASE          = image_size[0] + patch_size[0] - (image_size[0] % patch_size[0])
          _SCREAMING_SNAKE_CASE          = image_size[1] + patch_size[1] - (image_size[1] % patch_size[1])

          for model_class in self.all_model_classes:
             _SCREAMING_SNAKE_CASE          = True
             self.check_hidden_states_output(A__				,							A__				,							A__				,							(padded_height, padded_width)					)

             # check that output_hidden_states also work using config
             del inputs_dict["output_hidden_states"]
             _SCREAMING_SNAKE_CASE          = True
             self.check_hidden_states_output(A__				,							A__				,							A__				,							(padded_height, padded_width)					)






       @slow
       def 				UpperCamelCase (							self					)			->	List[str]:
          for model_name in TF_SWIN_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
             _SCREAMING_SNAKE_CASE          = TFSwinModel.from_pretrained(A__					)
             self.assertIsNotNone(A__					)






@require_vision
@require_tf
class       _a (unittest.TestCase):





       """simple docstring"""




       @cached_property
       def 				UpperCamelCase (							self					)			->	Union[str, Any]:
          return (
              AutoImageProcessor.from_pretrained("""microsoft/swin-tiny-patch4-window7-224"""					)
              if is_vision_available()
              else None
          )






       @slow
       def 				UpperCamelCase (							self					)			->	Union[str, Any]:
          _SCREAMING_SNAKE_CASE          = TFSwinForImageClassification.from_pretrained("""microsoft/swin-tiny-patch4-window7-224"""					)
          _SCREAMING_SNAKE_CASE          = self.default_image_processor

          _SCREAMING_SNAKE_CASE          = Image.open("""./tests/fixtures/tests_samples/COCO/000000039769.png"""					)
          _SCREAMING_SNAKE_CASE          = image_processor(images=A__				,							return_tensors="""tf"""					)

          # forward pass
          _SCREAMING_SNAKE_CASE          = model(A__					)

          # verify the logits
          _SCREAMING_SNAKE_CASE          = tf.TensorShape((1, 10_00)					)
          self.assertEqual(outputs.logits.shape				,							A__					)
          _SCREAMING_SNAKE_CASE          = tf.constant([-0.0948, -0.6454, -0.0921]					)
          self.assertTrue(np.allclose(outputs.logits[0, :3]				,							A__				,							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    UpperCamelCase     (			nn.Module    ):




	"""simple docstring"""






	def __init__(			self							,UpperCAmelCase_							,UpperCAmelCase_=1.0					):
			super().__init__()

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






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







def       __SCREAMING_SNAKE_CASE      (   __UpperCAmelCase					,							__UpperCAmelCase	):
		_lowercase							:      Dict			=					F.normalize(__UpperCAmelCase					,							dim=-1	)
		_lowercase							:      Optional[int]			=					F.normalize(__UpperCAmelCase					,							dim=-1	)
		return (x - y).norm(dim=-1	).div(2	).arcsin().pow(2	).mul(2	)


def       __SCREAMING_SNAKE_CASE      (   __UpperCAmelCase					,							__UpperCAmelCase	):
		for param in model.parameters():
				_lowercase							:      Union[str, Any]			=					value





class    UpperCamelCase     (			snake_case    ):




	"""simple docstring"""






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

			_lowercase							:      Tuple			=					transforms.Normalize(mean=feature_extractor.image_mean							,std=feature_extractor.image_std					)
			_lowercase							:      Tuple			=					(
			    feature_extractor.size
			    if isinstance(feature_extractor.size							,UpperCAmelCase_					)
			    else feature_extractor.size["""shortest_edge"""]
			)
			_lowercase							:      List[str]			=					MakeCutouts(self.cut_out_size					)

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






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






	def 					lowerCamelCase__ (			self					):
			self.enable_attention_slicing(UpperCAmelCase_					)






	def 					lowerCamelCase__ (			self					):
			set_requires_grad(self.vae							,UpperCAmelCase_					)






	def 					lowerCamelCase__ (			self					):
			set_requires_grad(self.vae							,UpperCAmelCase_					)






	def 					lowerCamelCase__ (			self					):
			set_requires_grad(self.unet							,UpperCAmelCase_					)






	def 					lowerCamelCase__ (			self					):
			set_requires_grad(self.unet							,UpperCAmelCase_					)






	@torch.enable_grad()
	def 					lowerCamelCase__ (			self							,UpperCAmelCase_							,UpperCAmelCase_							,UpperCAmelCase_							,UpperCAmelCase_							,UpperCAmelCase_							,UpperCAmelCase_							,UpperCAmelCase_							,UpperCAmelCase_							,UpperCAmelCase_=True							,):
			_lowercase							:      Union[str, Any]			=					latents.detach().requires_grad_()

			_lowercase							:      Optional[Any]			=					self.scheduler.scale_model_input(UpperCAmelCase_							,UpperCAmelCase_					)

			# predict the noise residual
			_lowercase							:      str			=					self.unet(UpperCAmelCase_							,UpperCAmelCase_							,encoder_hidden_states=UpperCAmelCase_					).sample

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

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

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

			if use_cutouts:
					_lowercase							:      Any			=					self.make_cutouts(UpperCAmelCase_							,UpperCAmelCase_					)
			else:
					_lowercase							:      Dict			=					transforms.Resize(self.cut_out_size					)(UpperCAmelCase_					)
			_lowercase							:      List[Any]			=					self.normalize(UpperCAmelCase_					).to(latents.dtype					)

			_lowercase							:      Optional[Any]			=					self.clip_model.get_image_features(UpperCAmelCase_					)
			_lowercase							:      List[Any]			=					image_embeddings_clip / image_embeddings_clip.norm(p=2							,dim=-1							,keepdim=UpperCAmelCase_					)

			if use_cutouts:
					_lowercase							:      Dict			=					spherical_dist_loss(UpperCAmelCase_							,UpperCAmelCase_					)
					_lowercase							:      List[str]			=					dists.view([num_cutouts, sample.shape[0], -1]					)
					_lowercase							:      Dict			=					dists.sum(2					).mean(0					).sum() * clip_guidance_scale
			else:
					_lowercase							:      int			=					spherical_dist_loss(UpperCAmelCase_							,UpperCAmelCase_					).mean() * clip_guidance_scale

			_lowercase							:      Dict			=					-torch.autograd.grad(UpperCAmelCase_							,UpperCAmelCase_					)[0]

			if isinstance(self.scheduler							,UpperCAmelCase_					):
					_lowercase							:      List[Any]			=					latents.detach() + grads * (sigma**2)
					_lowercase							:      Optional[int]			=					noise_pred_original
			else:
					_lowercase							:      str			=					noise_pred_original - torch.sqrt(UpperCAmelCase_					) * grads
			return noise_pred, latents






	@torch.no_grad()
	def __call__(			self							,UpperCAmelCase_							,UpperCAmelCase_ = 5_12							,UpperCAmelCase_ = 5_12							,UpperCAmelCase_ = 50							,UpperCAmelCase_ = 7.5							,UpperCAmelCase_ = 1							,UpperCAmelCase_ = 0.0							,UpperCAmelCase_ = 1_00							,UpperCAmelCase_ = None							,UpperCAmelCase_ = 4							,UpperCAmelCase_ = True							,UpperCAmelCase_ = None							,UpperCAmelCase_ = None							,UpperCAmelCase_ = "pil"							,UpperCAmelCase_ = True							,):
			if isinstance(UpperCAmelCase_							,UpperCAmelCase_					):
					_lowercase							:      Optional[int]			=					1
			elif isinstance(UpperCAmelCase_							,UpperCAmelCase_					):
					_lowercase							:      str			=					len(UpperCAmelCase_					)
			else:
					raise ValueError(f"""`prompt` has to be of type `str` or `list` but is {type(UpperCAmelCase_					)}"""					)

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

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

			# set timesteps
			_lowercase							:      Dict			=					"""offset""" in set(inspect.signature(self.scheduler.set_timesteps					).parameters.keys()					)
			_lowercase							:      Optional[Any]			=					{}
			if accepts_offset:
					_lowercase							:      Tuple			=					1

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

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

			# scale the initial noise by the standard deviation required by the scheduler
			_lowercase							:      Any			=					latents * self.scheduler.init_noise_sigma

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

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

			for i, t in enumerate(self.progress_bar(UpperCAmelCase_					)					):
					# expand the latents if we are doing classifier free guidance
					_lowercase							:      Any			=					torch.cat([latents] * 2					) if do_classifier_free_guidance else latents
					_lowercase							:      List[str]			=					self.scheduler.scale_model_input(UpperCAmelCase_							,UpperCAmelCase_					)

					# predict the noise residual
					_lowercase							:      Optional[int]			=					self.unet(UpperCAmelCase_							,UpperCAmelCase_							,encoder_hidden_states=UpperCAmelCase_					).sample

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

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

					# compute the previous noisy sample x_t -> x_t-1
					_lowercase							:      Optional[int]			=					self.scheduler.step(UpperCAmelCase_							,UpperCAmelCase_							,UpperCAmelCase_							,**UpperCAmelCase_					).prev_sample

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

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

			if output_type == "pil":
					_lowercase							:      Any			=					self.numpy_to_pil(UpperCAmelCase_					)

			if not return_dict:
					return (image, None)

			return StableDiffusionPipelineOutput(images=UpperCAmelCase_							,nsfw_content_detected=UpperCAmelCase_					)



