



"""simple docstring"""
import unittest

from transformers import FunnelConfig, FunnelTokenizer, is_torch_available
from transformers.models.auto import get_values
from transformers.testing_utils import require_sentencepiece, require_tokenizers, require_torch, slow, torch_device

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


if is_torch_available():
			import torch

			from transformers import (
			    MODEL_FOR_PRETRAINING_MAPPING,
			    FunnelBaseModel,
			    FunnelForMaskedLM,
			    FunnelForMultipleChoice,
			    FunnelForPreTraining,
			    FunnelForQuestionAnswering,
			    FunnelForSequenceClassification,
			    FunnelForTokenClassification,
			    FunnelModel,
			)


class 							lowerCamelCase_:





					'''simple docstring'''



					def __init__(			self ,				lowerCamelCase__ ,				lowerCamelCase__=1_3 ,				lowerCamelCase__=7 ,				lowerCamelCase__=True ,				lowerCamelCase__=True ,				lowerCamelCase__=True ,				lowerCamelCase__=True ,				lowerCamelCase__=9_9 ,				lowerCamelCase__=[1, 1, 2] ,				lowerCamelCase__=1 ,				lowerCamelCase__=3_2 ,				lowerCamelCase__=4 ,				lowerCamelCase__=8 ,				lowerCamelCase__=3_7 ,				lowerCamelCase__="gelu_new" ,				lowerCamelCase__=0.1 ,				lowerCamelCase__=0.1 ,				lowerCamelCase__=0.0 ,				lowerCamelCase__=5_1_2 ,				lowerCamelCase__=3 ,				lowerCamelCase__=0.0_2 ,				lowerCamelCase__=3 ,				lowerCamelCase__=4 ,				lowerCamelCase__=None ,				lowerCamelCase__=False ,				):
									_lowerCamelCase					     = parent
									_lowerCamelCase					     = batch_size
									_lowerCamelCase					     = seq_length
									_lowerCamelCase					     = is_training
									_lowerCamelCase					     = use_input_mask
									_lowerCamelCase					     = use_token_type_ids
									_lowerCamelCase					     = use_labels
									_lowerCamelCase					     = vocab_size
									_lowerCamelCase					     = block_sizes
									_lowerCamelCase					     = num_decoder_layers
									_lowerCamelCase					     = d_model
									_lowerCamelCase					     = n_head
									_lowerCamelCase					     = d_head
									_lowerCamelCase					     = d_inner
									_lowerCamelCase					     = hidden_act
									_lowerCamelCase					     = hidden_dropout
									_lowerCamelCase					     = attention_dropout
									_lowerCamelCase					     = activation_dropout
									_lowerCamelCase					     = max_position_embeddings
									_lowerCamelCase					     = type_vocab_size
									_lowerCamelCase					     = 2
									_lowerCamelCase					     = num_labels
									_lowerCamelCase					     = num_choices
									_lowerCamelCase					     = scope
									_lowerCamelCase					     = initializer_std

									# Used in the tests to check the size of the first attention layer
									_lowerCamelCase					     = n_head
									# Used in the tests to check the size of the first hidden state
									_lowerCamelCase					     = self.d_model
									# Used in the tests to check the number of output hidden states/attentions
									_lowerCamelCase					     = sum(self.block_sizes      ) + (0 if base else self.num_decoder_layers)
									# FunnelModel adds two hidden layers: input embeddings and the sum of the upsampled encoder hidden state with
									# the last hidden state of the first block (which is the first hidden state of the decoder).
									if not base:
													_lowerCamelCase					     = self.num_hidden_layers + 2

					def 	snake_case__					(			self      ):
									_lowerCamelCase					     = ids_tensor([self.batch_size, self.seq_length] ,				self.vocab_size      )

									_lowerCamelCase					     = None
									if self.use_input_mask:
													_lowerCamelCase					     = ids_tensor([self.batch_size, self.seq_length] ,				vocab_size=2      )

									_lowerCamelCase					     = None
									if self.use_token_type_ids:
													_lowerCamelCase					     = ids_tensor([self.batch_size, self.seq_length] ,				self.type_vocab_size      )

									_lowerCamelCase					     = None
									_lowerCamelCase					     = None
									_lowerCamelCase					     = None
									if self.use_labels:
													_lowerCamelCase					     = ids_tensor([self.batch_size] ,				self.type_sequence_label_size      )
													_lowerCamelCase					     = ids_tensor([self.batch_size, self.seq_length] ,				self.num_labels      )
													_lowerCamelCase					     = ids_tensor([self.batch_size] ,				self.num_choices      )
													_lowerCamelCase					     = ids_tensor([self.batch_size, self.seq_length] ,				1      )

									_lowerCamelCase					     = self.get_config()

									return (
									    config,
									    input_ids,
									    token_type_ids,
									    input_mask,
									    sequence_labels,
									    token_labels,
									    choice_labels,
									    fake_token_labels,
									)

					def 	snake_case__					(			self      ):
									return FunnelConfig(
									    vocab_size=self.vocab_size ,				block_sizes=self.block_sizes ,				num_decoder_layers=self.num_decoder_layers ,				d_model=self.d_model ,				n_head=self.n_head ,				d_head=self.d_head ,				d_inner=self.d_inner ,				hidden_act=self.hidden_act ,				hidden_dropout=self.hidden_dropout ,				attention_dropout=self.attention_dropout ,				activation_dropout=self.activation_dropout ,				max_position_embeddings=self.max_position_embeddings ,				type_vocab_size=self.type_vocab_size ,				initializer_std=self.initializer_std ,				)

					def 	snake_case__					(			self ,				lowerCamelCase__ ,				lowerCamelCase__ ,				lowerCamelCase__ ,				lowerCamelCase__ ,				lowerCamelCase__ ,				lowerCamelCase__ ,				lowerCamelCase__ ,				lowerCamelCase__ ,				):
									_lowerCamelCase					     = FunnelModel(config=lowerCamelCase__      )
									model.to(lowerCamelCase__      )
									model.eval()
									_lowerCamelCase					     = model(lowerCamelCase__ ,				attention_mask=lowerCamelCase__ ,				token_type_ids=lowerCamelCase__      )
									_lowerCamelCase					     = model(lowerCamelCase__ ,				token_type_ids=lowerCamelCase__      )
									_lowerCamelCase					     = model(lowerCamelCase__      )
									self.parent.assertEqual(result.last_hidden_state.shape ,				(self.batch_size, self.seq_length, self.d_model)      )

									_lowerCamelCase					     = False
									_lowerCamelCase					     = model(lowerCamelCase__      )
									self.parent.assertEqual(result.last_hidden_state.shape ,				(self.batch_size, self.seq_length, self.d_model)      )

									_lowerCamelCase					     = False
									_lowerCamelCase					     = model(lowerCamelCase__      )
									self.parent.assertEqual(result.last_hidden_state.shape ,				(self.batch_size, self.seq_length, self.d_model)      )

					def 	snake_case__					(			self ,				lowerCamelCase__ ,				lowerCamelCase__ ,				lowerCamelCase__ ,				lowerCamelCase__ ,				lowerCamelCase__ ,				lowerCamelCase__ ,				lowerCamelCase__ ,				lowerCamelCase__ ,				):
									_lowerCamelCase					     = FunnelBaseModel(config=lowerCamelCase__      )
									model.to(lowerCamelCase__      )
									model.eval()
									_lowerCamelCase					     = model(lowerCamelCase__ ,				attention_mask=lowerCamelCase__ ,				token_type_ids=lowerCamelCase__      )
									_lowerCamelCase					     = model(lowerCamelCase__ ,				token_type_ids=lowerCamelCase__      )
									_lowerCamelCase					     = model(lowerCamelCase__      )
									self.parent.assertEqual(result.last_hidden_state.shape ,				(self.batch_size, 2, self.d_model)      )

									_lowerCamelCase					     = False
									_lowerCamelCase					     = model(lowerCamelCase__      )
									self.parent.assertEqual(result.last_hidden_state.shape ,				(self.batch_size, 3, self.d_model)      )

									_lowerCamelCase					     = False
									_lowerCamelCase					     = model(lowerCamelCase__      )
									self.parent.assertEqual(result.last_hidden_state.shape ,				(self.batch_size, 2, self.d_model)      )

					def 	snake_case__					(			self ,				lowerCamelCase__ ,				lowerCamelCase__ ,				lowerCamelCase__ ,				lowerCamelCase__ ,				lowerCamelCase__ ,				lowerCamelCase__ ,				lowerCamelCase__ ,				lowerCamelCase__ ,				):
									_lowerCamelCase					     = self.num_labels
									_lowerCamelCase					     = FunnelForPreTraining(config=lowerCamelCase__      )
									model.to(lowerCamelCase__      )
									model.eval()
									_lowerCamelCase					     = model(lowerCamelCase__ ,				attention_mask=lowerCamelCase__ ,				token_type_ids=lowerCamelCase__ ,				labels=lowerCamelCase__      )
									self.parent.assertEqual(result.logits.shape ,				(self.batch_size, self.seq_length)      )

					def 	snake_case__					(			self ,				lowerCamelCase__ ,				lowerCamelCase__ ,				lowerCamelCase__ ,				lowerCamelCase__ ,				lowerCamelCase__ ,				lowerCamelCase__ ,				lowerCamelCase__ ,				lowerCamelCase__ ,				):
									_lowerCamelCase					     = FunnelForMaskedLM(config=lowerCamelCase__      )
									model.to(lowerCamelCase__      )
									model.eval()
									_lowerCamelCase					     = model(lowerCamelCase__ ,				attention_mask=lowerCamelCase__ ,				token_type_ids=lowerCamelCase__ ,				labels=lowerCamelCase__      )
									self.parent.assertEqual(result.logits.shape ,				(self.batch_size, self.seq_length, self.vocab_size)      )

					def 	snake_case__					(			self ,				lowerCamelCase__ ,				lowerCamelCase__ ,				lowerCamelCase__ ,				lowerCamelCase__ ,				lowerCamelCase__ ,				lowerCamelCase__ ,				lowerCamelCase__ ,				lowerCamelCase__ ,				):
									_lowerCamelCase					     = self.num_labels
									_lowerCamelCase					     = FunnelForSequenceClassification(lowerCamelCase__      )
									model.to(lowerCamelCase__      )
									model.eval()
									_lowerCamelCase					     = model(lowerCamelCase__ ,				attention_mask=lowerCamelCase__ ,				token_type_ids=lowerCamelCase__ ,				labels=lowerCamelCase__      )
									self.parent.assertEqual(result.logits.shape ,				(self.batch_size, self.num_labels)      )

					def 	snake_case__					(			self ,				lowerCamelCase__ ,				lowerCamelCase__ ,				lowerCamelCase__ ,				lowerCamelCase__ ,				lowerCamelCase__ ,				lowerCamelCase__ ,				lowerCamelCase__ ,				lowerCamelCase__ ,				):
									_lowerCamelCase					     = self.num_choices
									_lowerCamelCase					     = FunnelForMultipleChoice(config=lowerCamelCase__      )
									model.to(lowerCamelCase__      )
									model.eval()
									_lowerCamelCase					     = input_ids.unsqueeze(1      ).expand(-1 ,				self.num_choices ,				-1      ).contiguous()
									_lowerCamelCase					     = token_type_ids.unsqueeze(1      ).expand(-1 ,				self.num_choices ,				-1      ).contiguous()
									_lowerCamelCase					     = input_mask.unsqueeze(1      ).expand(-1 ,				self.num_choices ,				-1      ).contiguous()
									_lowerCamelCase					     = model(
									    lowerCamelCase__ ,				attention_mask=lowerCamelCase__ ,				token_type_ids=lowerCamelCase__ ,				labels=lowerCamelCase__ ,				)
									self.parent.assertEqual(result.logits.shape ,				(self.batch_size, self.num_choices)      )

					def 	snake_case__					(			self ,				lowerCamelCase__ ,				lowerCamelCase__ ,				lowerCamelCase__ ,				lowerCamelCase__ ,				lowerCamelCase__ ,				lowerCamelCase__ ,				lowerCamelCase__ ,				lowerCamelCase__ ,				):
									_lowerCamelCase					     = self.num_labels
									_lowerCamelCase					     = FunnelForTokenClassification(config=lowerCamelCase__      )
									model.to(lowerCamelCase__      )
									model.eval()
									_lowerCamelCase					     = model(lowerCamelCase__ ,				attention_mask=lowerCamelCase__ ,				token_type_ids=lowerCamelCase__ ,				labels=lowerCamelCase__      )
									self.parent.assertEqual(result.logits.shape ,				(self.batch_size, self.seq_length, self.num_labels)      )

					def 	snake_case__					(			self ,				lowerCamelCase__ ,				lowerCamelCase__ ,				lowerCamelCase__ ,				lowerCamelCase__ ,				lowerCamelCase__ ,				lowerCamelCase__ ,				lowerCamelCase__ ,				lowerCamelCase__ ,				):
									_lowerCamelCase					     = FunnelForQuestionAnswering(config=lowerCamelCase__      )
									model.to(lowerCamelCase__      )
									model.eval()
									_lowerCamelCase					     = model(
									    lowerCamelCase__ ,				attention_mask=lowerCamelCase__ ,				token_type_ids=lowerCamelCase__ ,				start_positions=lowerCamelCase__ ,				end_positions=lowerCamelCase__ ,				)
									self.parent.assertEqual(result.start_logits.shape ,				(self.batch_size, self.seq_length)      )
									self.parent.assertEqual(result.end_logits.shape ,				(self.batch_size, self.seq_length)      )





					def 	snake_case__					(			self      ):
									_lowerCamelCase					     = self.prepare_config_and_inputs()
									(
									    (
									    _lowerCamelCase
									)     ,   (
									    _lowerCamelCase
									)     ,   (
									    _lowerCamelCase
									)     ,   (
									    _lowerCamelCase
									)     ,   (
									    _lowerCamelCase
									)     ,   (
									    _lowerCamelCase
									)     ,   (
									    _lowerCamelCase
									)     ,   (
									    _lowerCamelCase
									)     ,   
									)					     = config_and_inputs
									_lowerCamelCase					     = {'''input_ids''': input_ids, '''token_type_ids''': token_type_ids, '''attention_mask''': input_mask}
									return config, inputs_dict


@require_torch
class 							lowerCamelCase_(      A__, A__, unittest.TestCase			):





					'''simple docstring'''



					lowercase__      :				int        =       False
					lowercase__      :				List[Any]        =       False
					lowercase__      :				List[str]        =       (
					    (
					        FunnelModel,
					        FunnelForMaskedLM,
					        FunnelForPreTraining,
					        FunnelForQuestionAnswering,
					        FunnelForTokenClassification,
					    )
					    if is_torch_available()
					    else ()
					)
					lowercase__      :				Tuple        =       (
					    {
					        'feature-extraction': (FunnelBaseModel, FunnelModel),
					        'fill-mask': FunnelForMaskedLM,
					        'question-answering': FunnelForQuestionAnswering,
					        'text-classification': FunnelForSequenceClassification,
					        'token-classification': FunnelForTokenClassification,
					        'zero-shot': FunnelForSequenceClassification,
					    }
					    if is_torch_available()
					    else {}
					)

					def 	snake_case__					(			self ,				lowerCamelCase__ ,				lowerCamelCase__ ,				lowerCamelCase__=False      ):
									_lowerCamelCase					     = super()._prepare_for_class(lowerCamelCase__ ,				lowerCamelCase__ ,				return_labels=lowerCamelCase__      )

									if return_labels:
													if model_class in get_values(lowerCamelCase__      ):
																	_lowerCamelCase					     = torch.zeros(
																	    (self.model_tester.batch_size, self.model_tester.seq_length) ,				dtype=torch.long ,				device=lowerCamelCase__      )
									return inputs_dict

					def 	snake_case__					(			self      ):
									_lowerCamelCase					     = FunnelModelTester(self      )
									_lowerCamelCase					     = ConfigTester(self ,				config_class=lowerCamelCase__      )

					def 	snake_case__					(			self      ):
									self.config_tester.run_common_tests()

					def 	snake_case__					(			self      ):
									_lowerCamelCase					     = self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_model(*lowerCamelCase__      )

					def 	snake_case__					(			self      ):
									_lowerCamelCase					     = self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_for_pretraining(*lowerCamelCase__      )

					def 	snake_case__					(			self      ):
									_lowerCamelCase					     = self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_for_masked_lm(*lowerCamelCase__      )

					def 	snake_case__					(			self      ):
									_lowerCamelCase					     = self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_for_token_classification(*lowerCamelCase__      )

					def 	snake_case__					(			self      ):
									_lowerCamelCase					     = self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_for_question_answering(*lowerCamelCase__      )





					def 	snake_case__					(			self ,				lowerCamelCase__      ):
									if hasattr(lowerCamelCase__ ,				'''weight'''      ) and module.weight is not None:
													module.weight.data.fill_(3      )
									if hasattr(lowerCamelCase__ ,				'''bias'''      ) and module.bias is not None:
													module.bias.data.fill_(3      )

									for param in ["r_w_bias", "r_r_bias", "r_kernel", "r_s_bias", "seg_embed"]:
													if hasattr(lowerCamelCase__ ,				lowerCamelCase__      ) and getattr(lowerCamelCase__ ,				lowerCamelCase__      ) is not None:
																	_lowerCamelCase					     = getattr(lowerCamelCase__ ,				lowerCamelCase__      )
																	weight.data.fill_(3      )


@require_torch
class 							lowerCamelCase_(      A__, unittest.TestCase			):





					'''simple docstring'''



					lowercase__      :				List[Any]        =       False
					lowercase__      :				List[str]        =       False
					lowercase__      :				int        =       (
					    (FunnelBaseModel, FunnelForMultipleChoice, FunnelForSequenceClassification) if is_torch_available() else ()
					)

					def 	snake_case__					(			self      ):
									_lowerCamelCase					     = FunnelModelTester(self ,				base=lowerCamelCase__      )
									_lowerCamelCase					     = ConfigTester(self ,				config_class=lowerCamelCase__      )

					def 	snake_case__					(			self      ):
									self.config_tester.run_common_tests()

					def 	snake_case__					(			self      ):
									_lowerCamelCase					     = self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_base_model(*lowerCamelCase__      )

					def 	snake_case__					(			self      ):
									_lowerCamelCase					     = self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_for_sequence_classification(*lowerCamelCase__      )

					def 	snake_case__					(			self      ):
									_lowerCamelCase					     = self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_for_multiple_choice(*lowerCamelCase__      )

					def 	snake_case__					(			self      ):
									_lowerCamelCase     ,   _lowerCamelCase					     = self.model_tester.prepare_config_and_inputs_for_common()
									_lowerCamelCase					     = True

									for model_class in self.all_model_classes:
													if model_class.__name__ == "FunnelBaseModel":
																	continue
													_lowerCamelCase					     = model_class(lowerCamelCase__      )
													model.to(lowerCamelCase__      )
													model.train()
													_lowerCamelCase					     = self._prepare_for_class(lowerCamelCase__ ,				lowerCamelCase__ ,				return_labels=lowerCamelCase__      )
													_lowerCamelCase					     = model(**lowerCamelCase__      ).loss
													loss.backward()





					def 	snake_case__					(			self ,				lowerCamelCase__      ):
									if hasattr(lowerCamelCase__ ,				'''weight'''      ) and module.weight is not None:
													module.weight.data.fill_(3      )
									if hasattr(lowerCamelCase__ ,				'''bias'''      ) and module.bias is not None:
													module.bias.data.fill_(3      )

									for param in ["r_w_bias", "r_r_bias", "r_kernel", "r_s_bias", "seg_embed"]:
													if hasattr(lowerCamelCase__ ,				lowerCamelCase__      ) and getattr(lowerCamelCase__ ,				lowerCamelCase__      ) is not None:
																	_lowerCamelCase					     = getattr(lowerCamelCase__ ,				lowerCamelCase__      )
																	weight.data.fill_(3      )


@require_torch
@require_sentencepiece
@require_tokenizers
class 							lowerCamelCase_(      unittest.TestCase			):





					'''simple docstring'''



					def 	snake_case__					(			self      ):
									_lowerCamelCase					     = 1_3
									_lowerCamelCase					     = 7
									_lowerCamelCase					     = torch.arange(0 ,				batch_size * sequence_length      ).long().reshape(lowerCamelCase__ ,				lowerCamelCase__      )
									_lowerCamelCase					     = [0, 1, 2, 3, 4, 5, 6, 4, 1, 3, 5, 0, 1]
									_lowerCamelCase					     = torch.tensor([[2] + [0] * a + [1] * (sequence_length - a - 1) for a in lengths]      )

									_lowerCamelCase					     = FunnelModel.from_pretrained('''sgugger/funnel-random-tiny'''      )
									_lowerCamelCase					     = model(lowerCamelCase__ ,				token_type_ids=lowerCamelCase__      )[0].abs()

									_lowerCamelCase					     = torch.tensor(2_3_4_4.8_3_5_2      )
									_lowerCamelCase					     = torch.tensor(0.8_0_5_2      )
									self.assertTrue(torch.allclose(output.sum() ,				lowerCamelCase__ ,				atol=1e-4      )      )
									self.assertTrue(torch.allclose(output.mean() ,				lowerCamelCase__ ,				atol=1e-4      )      )

									_lowerCamelCase					     = torch.tensor([[1] * 7, [1] * 4 + [0] * 3] * 6 + [[0, 1, 1, 0, 0, 1, 1]]      )
									_lowerCamelCase					     = model(lowerCamelCase__ ,				attention_mask=lowerCamelCase__ ,				token_type_ids=lowerCamelCase__      )[0].abs()

									_lowerCamelCase					     = torch.tensor(2_3_4_3.8_4_2_5      )
									_lowerCamelCase					     = torch.tensor(0.8_0_4_9      )
									self.assertTrue(torch.allclose(output.sum() ,				lowerCamelCase__ ,				atol=1e-4      )      )
									self.assertTrue(torch.allclose(output.mean() ,				lowerCamelCase__ ,				atol=1e-4      )      )





					@slow
					def 	snake_case__					(			self      ):
									_lowerCamelCase					     = FunnelTokenizer.from_pretrained('''huggingface/funnel-small'''      )
									_lowerCamelCase					     = FunnelModel.from_pretrained('''huggingface/funnel-small'''      )
									_lowerCamelCase					     = tokenizer('''Hello! I am the Funnel Transformer model.''' ,				return_tensors='''pt'''      )
									_lowerCamelCase					     = model(**lowerCamelCase__      )[0]

									_lowerCamelCase					     = torch.tensor(2_3_5.7_2_4_6      )
									_lowerCamelCase					     = torch.tensor(0.0_2_5_6      )
									self.assertTrue(torch.allclose(output.sum() ,				lowerCamelCase__ ,				atol=1e-4      )      )
									self.assertTrue(torch.allclose(output.mean() ,				lowerCamelCase__ ,				atol=1e-4      )      )








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 							a_ (						nn.Module		):


							'''simple docstring'''






							def __init__(		self			,			lowercase_			,			lowercase_=1.0   )					->					Any:
														'''simple docstring'''

														super().__init__()

														lowerCAmelCase_										=     cut_size
														lowerCAmelCase_										=     cut_power






							def      _lowercase       (		self			,			lowercase_			,			lowercase_   )					->					int:
														'''simple docstring'''

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


def      lowerCamelCase					(						a_ ,						a_	)	->			Optional[int]:
							lowerCAmelCase_										=     F.normalize(a_ ,						dim=-1	)
							lowerCAmelCase_										=     F.normalize(a_ ,						dim=-1	)
							return (x - y).norm(dim=-1	).div(2	).arcsin().pow(2	).mul(2	)
def      lowerCamelCase					(						a_ ,						a_	)	->			int:
							for param in model.parameters():
														lowerCAmelCase_										=     value




class 							a_ (						a_		):


							'''simple docstring'''






							def __init__(		self			,			lowercase_			,			lowercase_			,			lowercase_			,			lowercase_			,			lowercase_			,			lowercase_			,			lowercase_			,			)					->					str:
														'''simple docstring'''

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

														lowerCAmelCase_										=     transforms.Normalize(mean=feature_extractor.image_mean			,			std=feature_extractor.image_std   )
														lowerCAmelCase_										=     (
														    feature_extractor.size
														    if isinstance(feature_extractor.size			,			lowercase_   )
														    else feature_extractor.size['shortest_edge']
														)
														lowerCAmelCase_										=     MakeCutouts(self.cut_out_size   )

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

							def      _lowercase       (		self			,			lowercase_ = "auto"   )					->					Union[str, Any]:
														'''simple docstring'''

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

							def      _lowercase       (		self   )					->					List[Any]:
														'''simple docstring'''

														self.enable_attention_slicing(lowercase_   )

							def      _lowercase       (		self   )					->					str:
														'''simple docstring'''

														set_requires_grad(self.vae			,			lowercase_   )

							def      _lowercase       (		self   )					->					List[str]:
														'''simple docstring'''

														set_requires_grad(self.vae			,			lowercase_   )

							def      _lowercase       (		self   )					->					List[str]:
														'''simple docstring'''

														set_requires_grad(self.unet			,			lowercase_   )

							def      _lowercase       (		self   )					->					Tuple:
														'''simple docstring'''

														set_requires_grad(self.unet			,			lowercase_   )

							@torch.enable_grad()
							def      _lowercase       (		self			,			lowercase_			,			lowercase_			,			lowercase_			,			lowercase_			,			lowercase_			,			lowercase_			,			lowercase_			,			lowercase_			,			lowercase_=True			,			)					->					Union[str, Any]:
														'''simple docstring'''

														lowerCAmelCase_										=     latents.detach().requires_grad_()

														lowerCAmelCase_										=     self.scheduler.scale_model_input(lowercase_			,			lowercase_   )

														# predict the noise residual
														lowerCAmelCase_										=     self.unet(lowercase_			,			lowercase_			,			encoder_hidden_states=lowercase_   ).sample

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

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

														lowerCAmelCase_										=     1 / self.vae.config.scaling_factor * sample
														lowerCAmelCase_										=     self.vae.decode(lowercase_   ).sample
														lowerCAmelCase_										=     (image / 2 + 0.5).clamp(0			,			1   )

														if use_cutouts:
																					lowerCAmelCase_										=     self.make_cutouts(lowercase_			,			lowercase_   )
														else:
																					lowerCAmelCase_										=     transforms.Resize(self.cut_out_size   )(lowercase_   )
														lowerCAmelCase_										=     self.normalize(lowercase_   ).to(latents.dtype   )

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

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

														lowerCAmelCase_										=     -torch.autograd.grad(lowercase_			,			lowercase_   )[0]

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






							@torch.no_grad()
							def __call__(		self			,			lowercase_			,			lowercase_ = 5_1_2			,			lowercase_ = 5_1_2			,			lowercase_ = 5_0			,			lowercase_ = 7.5			,			lowercase_ = 1			,			lowercase_ = 0.0			,			lowercase_ = 1_0_0			,			lowercase_ = None			,			lowercase_ = 4			,			lowercase_ = True			,			lowercase_ = None			,			lowercase_ = None			,			lowercase_ = "pil"			,			lowercase_ = True			,			)					->					str:
														'''simple docstring'''

														if isinstance(lowercase_			,			lowercase_   ):
																					lowerCAmelCase_										=     1
														elif isinstance(lowercase_			,			lowercase_   ):
																					lowerCAmelCase_										=     len(lowercase_   )
														else:
																					raise ValueError(f'''`prompt` has to be of type `str` or `list` but is {type(lowercase_   )}'''   )

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

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

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

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

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

														# get the initial random noise unless the user supplied it

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

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

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

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

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

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

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

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

																					# predict the noise residual
																					lowerCAmelCase_										=     self.unet(lowercase_			,			lowercase_			,			encoder_hidden_states=lowercase_   ).sample

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

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

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

														# scale and decode the image latents with vae
														lowerCAmelCase_										=     1 / self.vae.config.scaling_factor * latents
														lowerCAmelCase_										=     self.vae.decode(lowercase_   ).sample

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

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

														if not return_dict:
																					return (image, None)

														return StableDiffusionPipelineOutput(images=lowercase_			,			nsfw_content_detected=lowercase_   )



