
import os
import tempfile
import unittest

from transformers import FlaubertConfig, is_torch_available
from transformers.testing_utils import require_torch, require_torch_gpu, slow, torch_device

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


if is_torch_available():
   import torch

   from transformers import (
       FlaubertForMultipleChoice,
       FlaubertForQuestionAnswering,
       FlaubertForQuestionAnsweringSimple,
       FlaubertForSequenceClassification,
       FlaubertForTokenClassification,
       FlaubertModel,
       FlaubertWithLMHeadModel,
   )
   from transformers.models.flaubert.modeling_flaubert import FLAUBERT_PRETRAINED_MODEL_ARCHIVE_LIST







class   lowerCamelCase	(						A_						):





     def __init__(self   :					List[str]       ,					_A   :					str       ,					_A   :					List[Any]=1_3       ,					_A   :					Any=7       ,					_A   :					str=True       ,					_A   :					int=True       ,					_A   :					List[Any]=True       ,					_A   :					Optional[Any]=True       ,					_A   :					List[Any]=True       ,					_A   :					int=False       ,					_A   :					List[Any]=False       ,					_A   :					Tuple=False       ,					_A   :					int=2       ,					_A   :					Optional[int]=9_9       ,					_A   :					Union[str, Any]=0       ,					_A   :					Optional[Any]=3_2       ,					_A   :					List[str]=5       ,					_A   :					Optional[Any]=4       ,					_A   :					int=0.1       ,					_A   :					Optional[int]=0.1       ,					_A   :					List[Any]=5_1_2       ,					_A   :					str=1_2       ,					_A   :					str=2       ,					_A   :					Optional[Any]=0.02       ,					_A   :					int=3       ,					_A   :					Optional[int]=4       ,					_A   :					List[Any]="last"       ,					_A   :					Dict=None       ,					_A   :					List[Any]=None       ,					)		->		Tuple:
         snake_case				  =      parent
         snake_case				  =      batch_size
         snake_case				  =      seq_length
         snake_case				  =      is_training
         snake_case				  =      use_input_lengths
         snake_case				  =      use_token_type_ids
         snake_case				  =      use_labels
         snake_case				  =      gelu_activation
         snake_case				  =      sinusoidal_embeddings
         snake_case				  =      causal
         snake_case				  =      asm
         snake_case				  =      n_langs
         snake_case				  =      vocab_size
         snake_case				  =      n_special
         snake_case				  =      hidden_size
         snake_case				  =      num_hidden_layers
         snake_case				  =      num_attention_heads
         snake_case				  =      hidden_dropout_prob
         snake_case				  =      attention_probs_dropout_prob
         snake_case				  =      max_position_embeddings
         snake_case				  =      type_vocab_size
         snake_case				  =      type_sequence_label_size
         snake_case				  =      initializer_range
         snake_case				  =      num_labels
         snake_case				  =      num_choices
         snake_case				  =      summary_type
         snake_case				  =      use_proj
         snake_case				  =      scope





     def 	UpperCAmelCase(self   :					Optional[Any]     )		->		Dict:
         snake_case				  =      ids_tensor([self.batch_size, self.seq_length]       ,					self.vocab_size     )
         snake_case				  =      random_attention_mask([self.batch_size, self.seq_length]     )

         snake_case				  =      None
         if self.use_input_lengths:
             snake_case				  =      (
                 ids_tensor([self.batch_size]       ,					vocab_size=2     ) + self.seq_length - 2
             )  # small variation of seq_length

         snake_case				  =      None
         if self.use_token_type_ids:
             snake_case				  =      ids_tensor([self.batch_size, self.seq_length]       ,					self.n_langs     )

         snake_case				  =      None
         snake_case				  =      None
         snake_case				  =      None
         if self.use_labels:
             snake_case				  =      ids_tensor([self.batch_size]       ,					self.type_sequence_label_size     )
             snake_case				  =      ids_tensor([self.batch_size, self.seq_length]       ,					self.num_labels     )
             snake_case				  =      ids_tensor([self.batch_size]       ,					2     ).float()
             snake_case				  =      ids_tensor([self.batch_size]       ,					self.num_choices     )

         snake_case				  =      self.get_config()

         return (
             config,
             input_ids,
             token_type_ids,
             input_lengths,
             sequence_labels,
             token_labels,
             is_impossible_labels,
             choice_labels,
             input_mask,
         )





     def 	UpperCAmelCase(self   :					List[str]     )		->		int:
         return FlaubertConfig(
             vocab_size=self.vocab_size       ,					n_special=self.n_special       ,					emb_dim=self.hidden_size       ,					n_layers=self.num_hidden_layers       ,					n_heads=self.num_attention_heads       ,					dropout=self.hidden_dropout_prob       ,					attention_dropout=self.attention_probs_dropout_prob       ,					gelu_activation=self.gelu_activation       ,					sinusoidal_embeddings=self.sinusoidal_embeddings       ,					asm=self.asm       ,					causal=self.causal       ,					n_langs=self.n_langs       ,					max_position_embeddings=self.max_position_embeddings       ,					initializer_range=self.initializer_range       ,					summary_type=self.summary_type       ,					use_proj=self.use_proj       ,					)





     def 	UpperCAmelCase(self   :					Dict       ,					_A   :					int       ,					_A   :					List[Any]       ,					_A   :					Dict       ,					_A   :					Any       ,					_A   :					Optional[int]       ,					_A   :					List[Any]       ,					_A   :					List[str]       ,					_A   :					Optional[int]       ,					_A   :					int       ,					)		->		Dict:
         snake_case				  =      FlaubertModel(config=_A     )
         model.to(_A     )
         model.eval()
         snake_case				  =      model(_A       ,					lengths=_A       ,					langs=_A     )
         snake_case				  =      model(_A       ,					langs=_A     )
         snake_case				  =      model(_A     )
         self.parent.assertEqual(result.last_hidden_state.shape       ,					(self.batch_size, self.seq_length, self.hidden_size)     )





     def 	UpperCAmelCase(self   :					Optional[Any]       ,					_A   :					Dict       ,					_A   :					Dict       ,					_A   :					Union[str, Any]       ,					_A   :					int       ,					_A   :					Tuple       ,					_A   :					Optional[Any]       ,					_A   :					Dict       ,					_A   :					int       ,					_A   :					str       ,					)		->		int:
         snake_case				  =      FlaubertWithLMHeadModel(_A     )
         model.to(_A     )
         model.eval()

         snake_case				  =      model(_A       ,					token_type_ids=_A       ,					labels=_A     )
         self.parent.assertEqual(result.loss.shape       ,					()     )
         self.parent.assertEqual(result.logits.shape       ,					(self.batch_size, self.seq_length, self.vocab_size)     )





     def 	UpperCAmelCase(self   :					Any       ,					_A   :					Tuple       ,					_A   :					List[str]       ,					_A   :					Union[str, Any]       ,					_A   :					Optional[Any]       ,					_A   :					Optional[int]       ,					_A   :					Optional[int]       ,					_A   :					Union[str, Any]       ,					_A   :					Union[str, Any]       ,					_A   :					List[Any]       ,					)		->		Union[str, Any]:
         snake_case				  =      FlaubertForQuestionAnsweringSimple(_A     )
         model.to(_A     )
         model.eval()

         snake_case				  =      model(_A     )

         snake_case				  =      model(_A       ,					start_positions=_A       ,					end_positions=_A     )
         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 	UpperCAmelCase(self   :					Dict       ,					_A   :					List[Any]       ,					_A   :					Tuple       ,					_A   :					List[str]       ,					_A   :					str       ,					_A   :					Tuple       ,					_A   :					int       ,					_A   :					Tuple       ,					_A   :					Tuple       ,					_A   :					Dict       ,					)		->		str:
         snake_case				  =      FlaubertForQuestionAnswering(_A     )
         model.to(_A     )
         model.eval()

         snake_case				  =      model(_A     )

         snake_case				  =      model(
             _A       ,					start_positions=_A       ,					end_positions=_A       ,					cls_index=_A       ,					is_impossible=_A       ,					p_mask=_A       ,					)

         snake_case				  =      model(
             _A       ,					start_positions=_A       ,					end_positions=_A       ,					cls_index=_A       ,					is_impossible=_A       ,					)

         ((snake_case)  ,							)				  =      result_with_labels.to_tuple()

         snake_case				  =      model(_A       ,					start_positions=_A       ,					end_positions=_A     )

         ((snake_case)  ,							)				  =      result_with_labels.to_tuple()

         self.parent.assertEqual(result_with_labels.loss.shape       ,					()     )
         self.parent.assertEqual(result.start_top_log_probs.shape       ,					(self.batch_size, model.config.start_n_top)     )
         self.parent.assertEqual(result.start_top_index.shape       ,					(self.batch_size, model.config.start_n_top)     )
         self.parent.assertEqual(
             result.end_top_log_probs.shape       ,					(self.batch_size, model.config.start_n_top * model.config.end_n_top)     )
         self.parent.assertEqual(
             result.end_top_index.shape       ,					(self.batch_size, model.config.start_n_top * model.config.end_n_top)     )
         self.parent.assertEqual(result.cls_logits.shape       ,					(self.batch_size,)     )





     def 	UpperCAmelCase(self   :					Optional[Any]       ,					_A   :					List[str]       ,					_A   :					Union[str, Any]       ,					_A   :					str       ,					_A   :					Optional[Any]       ,					_A   :					Tuple       ,					_A   :					Optional[Any]       ,					_A   :					Union[str, Any]       ,					_A   :					Tuple       ,					_A   :					str       ,					)		->		str:
         snake_case				  =      FlaubertForSequenceClassification(_A     )
         model.to(_A     )
         model.eval()

         snake_case				  =      model(_A     )
         snake_case				  =      model(_A       ,					labels=_A     )

         self.parent.assertEqual(result.loss.shape       ,					()     )
         self.parent.assertEqual(result.logits.shape       ,					(self.batch_size, self.type_sequence_label_size)     )





     def 	UpperCAmelCase(self   :					Union[str, Any]       ,					_A   :					Optional[int]       ,					_A   :					str       ,					_A   :					List[str]       ,					_A   :					Optional[int]       ,					_A   :					Union[str, Any]       ,					_A   :					List[str]       ,					_A   :					List[Any]       ,					_A   :					Tuple       ,					_A   :					Any       ,					)		->		List[str]:
         snake_case				  =      self.num_labels
         snake_case				  =      FlaubertForTokenClassification(_A     )
         model.to(_A     )
         model.eval()

         snake_case				  =      model(_A       ,					attention_mask=_A       ,					labels=_A     )
         self.parent.assertEqual(result.logits.shape       ,					(self.batch_size, self.seq_length, self.num_labels)     )





     def 	UpperCAmelCase(self   :					Optional[Any]       ,					_A   :					List[str]       ,					_A   :					Optional[int]       ,					_A   :					Optional[Any]       ,					_A   :					Tuple       ,					_A   :					int       ,					_A   :					str       ,					_A   :					int       ,					_A   :					Union[str, Any]       ,					_A   :					Tuple       ,					)		->		List[str]:
         snake_case				  =      self.num_choices
         snake_case				  =      FlaubertForMultipleChoice(config=_A     )
         model.to(_A     )
         model.eval()
         snake_case				  =      input_ids.unsqueeze(1     ).expand(-1       ,					self.num_choices       ,					-1     ).contiguous()
         snake_case				  =      token_type_ids.unsqueeze(1     ).expand(-1       ,					self.num_choices       ,					-1     ).contiguous()
         snake_case				  =      input_mask.unsqueeze(1     ).expand(-1       ,					self.num_choices       ,					-1     ).contiguous()
         snake_case				  =      model(
             _A       ,					attention_mask=_A       ,					token_type_ids=_A       ,					labels=_A       ,					)
         self.parent.assertEqual(result.logits.shape       ,					(self.batch_size, self.num_choices)     )





     def 	UpperCAmelCase(self   :					Tuple     )		->		Optional[int]:
         snake_case				  =      self.prepare_config_and_inputs()
         (
             (
             snake_case
         )  ,							(
             snake_case
         )  ,							(
             snake_case
         )  ,							(
             snake_case
         )  ,							(
             snake_case
         )  ,							(
             snake_case
         )  ,							(
             snake_case
         )  ,							(
             snake_case
         )  ,							(
             snake_case
         )  ,							
         )				  =      config_and_inputs
         snake_case				  =      {
             "input_ids": input_ids,
             "token_type_ids": token_type_ids,
             "lengths": input_lengths,
             "attention_mask": input_mask,
         }
         return config, inputs_dict







@require_torch
class   lowerCamelCase	(						A_							,				A_							,				unittest.TestCase						):
     UpperCAmelCase__       :							List[str]     =	(
         (
             FlaubertModel,
             FlaubertWithLMHeadModel,
             FlaubertForQuestionAnswering,
             FlaubertForQuestionAnsweringSimple,
             FlaubertForSequenceClassification,
             FlaubertForTokenClassification,
             FlaubertForMultipleChoice,
         )
         if is_torch_available()
         else ()
     )
     UpperCAmelCase__       :							str     =	(
         {
             "feature-extraction": FlaubertModel,
             "fill-mask": FlaubertWithLMHeadModel,
             "question-answering": FlaubertForQuestionAnsweringSimple,
             "text-classification": FlaubertForSequenceClassification,
             "token-classification": FlaubertForTokenClassification,
             "zero-shot": FlaubertForSequenceClassification,
         }
         if is_torch_available()
         else {}
     )





     def 	UpperCAmelCase(self   :					Union[str, Any]       ,					_A   :					List[Any]       ,					_A   :					Tuple       ,					_A   :					Tuple       ,					_A   :					Optional[int]       ,					_A   :					Optional[Any]     )		->		Dict:
         if (
             pipeline_test_casse_name == "QAPipelineTests"
             and tokenizer_name is not None
             and not tokenizer_name.endswith("Fast"     )
         ):
             # `QAPipelineTests` fails for a few models when the slower tokenizer are used.
             # (The slower tokenizers were never used for pipeline tests before the pipeline testing rework)
             # TODO: check (and possibly fix) the `QAPipelineTests` with slower tokenizer
             return True

         return False





     def 	UpperCAmelCase(self   :					str       ,					_A   :					List[Any]       ,					_A   :					Union[str, Any]       ,					_A   :					Optional[Any]=False     )		->		int:
         snake_case				  =      super()._prepare_for_class(_A       ,					_A       ,					return_labels=_A     )

         if return_labels:
             if model_class.__name__ == "FlaubertForQuestionAnswering":
                 snake_case				  =      torch.zeros(
                     self.model_tester.batch_size       ,					dtype=torch.long       ,					device=_A     )
                 snake_case				  =      torch.zeros(
                     self.model_tester.batch_size       ,					dtype=torch.long       ,					device=_A     )

         return inputs_dict





     def 	UpperCAmelCase(self   :					Optional[Any]     )		->		Optional[Any]:
         snake_case				  =      FlaubertModelTester(self     )
         snake_case				  =      ConfigTester(self       ,					config_class=_A       ,					emb_dim=3_7     )





     def 	UpperCAmelCase(self   :					Optional[Any]     )		->		int:
         self.config_tester.run_common_tests()





     def 	UpperCAmelCase(self   :					str     )		->		Tuple:
         snake_case				  =      self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_flaubert_model(*_A     )





     def 	UpperCAmelCase(self   :					Optional[int]     )		->		Optional[Any]:
         snake_case				  =      self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_flaubert_lm_head(*_A     )





     def 	UpperCAmelCase(self   :					Optional[Any]     )		->		int:
         snake_case				  =      self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_flaubert_simple_qa(*_A     )





     def 	UpperCAmelCase(self   :					List[Any]     )		->		Tuple:
         snake_case				  =      self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_flaubert_qa(*_A     )





     def 	UpperCAmelCase(self   :					str     )		->		Any:
         snake_case				  =      self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_flaubert_sequence_classif(*_A     )





     def 	UpperCAmelCase(self   :					str     )		->		Dict:
         snake_case				  =      self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_flaubert_token_classif(*_A     )





     def 	UpperCAmelCase(self   :					Union[str, Any]     )		->		List[str]:
         snake_case				  =      self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_flaubert_multiple_choice(*_A     )





     @slow
     def 	UpperCAmelCase(self   :					List[str]     )		->		Union[str, Any]:
         for model_name in FLAUBERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
             snake_case				  =      FlaubertModel.from_pretrained(_A     )
             self.assertIsNotNone(_A     )





     @slow
     @require_torch_gpu
     def 	UpperCAmelCase(self   :					Tuple     )		->		Optional[Any]:
         snake_case  ,							snake_case				  =      self.model_tester.prepare_config_and_inputs_for_common()
         for model_class in self.all_model_classes:
             # FlauBertForMultipleChoice behaves incorrectly in JIT environments.
             if model_class == FlaubertForMultipleChoice:
                 return

             snake_case				  =      True
             snake_case				  =      model_class(config=_A     )

             snake_case				  =      self._prepare_for_class(_A       ,					_A     )
             snake_case				  =      torch.jit.trace(
                 _A       ,					(inputs_dict["input_ids"].to("cpu"     ), inputs_dict["attention_mask"].to("cpu"     ))     )

             with tempfile.TemporaryDirectory() as tmp:
                 torch.jit.save(_A       ,					os.path.join(_A       ,					"traced_model.pt"     )     )
                 snake_case				  =      torch.jit.load(os.path.join(_A       ,					"traced_model.pt"     )       ,					map_location=_A     )
                 loaded(inputs_dict["input_ids"].to(_A     )       ,					inputs_dict["attention_mask"].to(_A     )     )







@require_torch
class   lowerCamelCase	(						unittest.TestCase						):





     @slow
     def 	UpperCAmelCase(self   :					Tuple     )		->		str:
         snake_case				  =      FlaubertModel.from_pretrained("flaubert/flaubert_base_cased"     )
         snake_case				  =      torch.tensor([[0, 3_4_5, 2_3_2, 3_2_8, 7_4_0, 1_4_0, 1_6_9_5, 6_9, 6_0_7_8, 1_5_8_8, 2]]     )
         with torch.no_grad():
             snake_case				  =      model(_A     )[0]
         snake_case				  =      torch.Size((1, 1_1, 7_6_8)     )
         self.assertEqual(output.shape       ,					_A     )
         snake_case				  =      torch.tensor(
             [[[-2.62_51, -1.42_98, -0.02_27], [-2.85_10, -1.63_87, 0.22_58], [-2.81_14, -1.18_32, -0.30_66]]]     )

         self.assertTrue(torch.allclose(output[:, :3, :3]       ,					_A       ,					atol=1E-4     )     )



import inspect
from typing import Optional, Union

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

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



def 		lowercase_    (     A__				,	A__				,	A__	)							->							str:
    """simple docstring"""

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

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



def 		lowercase_    (     A__				,	A__				,	A__				,	A__=0.9995	)							->							List[Any]:
    """simple docstring"""

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

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

    if inputs_are_torch:
        snake_case				  =      torch.from_numpy(A__	).to(A__	)

    return va



def 		lowercase_    (     A__				,	A__	)							->							Any:
    """simple docstring"""

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



def 		lowercase_    (     A__				,	A__	)							->							Optional[int]:
    """simple docstring"""

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







class   lowerCamelCase	(						A_						):





     def __init__(self   :					str       ,					_A   :					AutoencoderKL       ,					_A   :					CLIPTextModel       ,					_A   :					CLIPModel       ,					_A   :					CLIPTokenizer       ,					_A   :					UNetaDConditionModel       ,					_A   :					Union[PNDMScheduler, LMSDiscreteScheduler, DDIMScheduler, DPMSolverMultistepScheduler]       ,					_A   :					CLIPFeatureExtractor       ,					_A   :					List[Any]=None       ,					_A   :					Tuple=None       ,					_A   :					Union[str, Any]=None       ,					)		->		Dict:
         super().__init__()
         self.register_modules(
             vae=_A       ,					text_encoder=_A       ,					clip_model=_A       ,					tokenizer=_A       ,					unet=_A       ,					scheduler=_A       ,					feature_extractor=_A       ,					coca_model=_A       ,					coca_tokenizer=_A       ,					coca_transform=_A       ,					)
         snake_case				  =      (
             feature_extractor.size
             if isinstance(feature_extractor.size       ,					_A     )
             else feature_extractor.size["shortest_edge"]
         )
         snake_case				  =      transforms.Normalize(mean=feature_extractor.image_mean       ,					std=feature_extractor.image_std     )
         set_requires_grad(self.text_encoder       ,					_A     )
         set_requires_grad(self.clip_model       ,					_A     )





     def 	UpperCAmelCase(self   :					List[Any]       ,					_A   :					Optional[Union[str, int]] = "auto"     )		->		Dict:
         if slice_size == "auto":
             # half the attention head size is usually a good trade-off between
             # speed and memory
             snake_case				  =      self.unet.config.attention_head_dim // 2
         self.unet.set_attention_slice(_A     )





     def 	UpperCAmelCase(self   :					str     )		->		Any:
         self.enable_attention_slicing(_A     )





     def 	UpperCAmelCase(self   :					List[Any]     )		->		int:
         set_requires_grad(self.vae       ,					_A     )





     def 	UpperCAmelCase(self   :					Any     )		->		Union[str, Any]:
         set_requires_grad(self.vae       ,					_A     )





     def 	UpperCAmelCase(self   :					Optional[Any]     )		->		int:
         set_requires_grad(self.unet       ,					_A     )





     def 	UpperCAmelCase(self   :					str     )		->		Tuple:
         set_requires_grad(self.unet       ,					_A     )





     def 	UpperCAmelCase(self   :					Tuple       ,					_A   :					str       ,					_A   :					str       ,					_A   :					List[str]     )		->		Dict:
         # get the original timestep using init_timestep
         snake_case				  =      min(int(num_inference_steps * strength     )       ,					_A     )

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

         return timesteps, num_inference_steps - t_start





     def 	UpperCAmelCase(self   :					List[Any]       ,					_A   :					Dict       ,					_A   :					List[Any]       ,					_A   :					Dict       ,					_A   :					Optional[Any]       ,					_A   :					List[Any]       ,					_A   :					List[Any]=None     )		->		str:
         if not isinstance(_A       ,					torch.Tensor     ):
             raise ValueError(f'`image` has to be of type `torch.Tensor` but is {type(_A     )}'     )

         snake_case				  =      image.to(device=_A       ,					dtype=_A     )

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

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

         snake_case				  =      randn_tensor(init_latents.shape       ,					generator=_A       ,					device=_A       ,					dtype=_A     )

         # get latents
         snake_case				  =      self.scheduler.add_noise(_A       ,					_A       ,					_A     )
         snake_case				  =      init_latents

         return latents





     def 	UpperCAmelCase(self   :					Dict       ,					_A   :					Optional[int]     )		->		Union[str, Any]:
         snake_case				  =      self.coca_transform(_A     ).unsqueeze(0     )
         with torch.no_grad(), torch.cuda.amp.autocast():
             snake_case				  =      self.coca_model.generate(transformed_image.to(device=self.device       ,					dtype=self.coca_model.dtype     )     )
         snake_case				  =      self.coca_tokenizer.decode(generated[0].cpu().numpy()     )
         return generated.split("<end_of_text>"     )[0].replace("<start_of_text>"       ,					""     ).rstrip(" .,"     )





     def 	UpperCAmelCase(self   :					int       ,					_A   :					Union[str, Any]       ,					_A   :					Optional[Any]     )		->		Tuple:
         snake_case				  =      self.feature_extractor.preprocess(_A     )
         snake_case				  =      torch.from_numpy(clip_image_input["pixel_values"][0]     ).unsqueeze(0     ).to(self.device     ).half()
         snake_case				  =      self.clip_model.get_image_features(_A     )
         snake_case				  =      image_embeddings_clip / image_embeddings_clip.norm(p=2       ,					dim=-1       ,					keepdim=_A     )
         snake_case				  =      image_embeddings_clip.repeat_interleave(_A       ,					dim=0     )
         return image_embeddings_clip





     @torch.enable_grad()
     def 	UpperCAmelCase(self   :					Union[str, Any]       ,					_A   :					Tuple       ,					_A   :					str       ,					_A   :					Tuple       ,					_A   :					Optional[Any]       ,					_A   :					List[str]       ,					_A   :					Any       ,					_A   :					str       ,					)		->		List[Any]:
         snake_case				  =      latents.detach().requires_grad_()

         snake_case				  =      self.scheduler.scale_model_input(_A       ,					_A     )

         # predict the noise residual
         snake_case				  =      self.unet(_A       ,					_A       ,					encoder_hidden_states=_A     ).sample

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

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

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

         snake_case				  =      transforms.Resize(self.feature_extractor_size     )(_A     )
         snake_case				  =      self.normalize(_A     ).to(latents.dtype     )

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

         snake_case				  =      spherical_dist_loss(_A       ,					_A     ).mean() * clip_guidance_scale

         snake_case				  =      -torch.autograd.grad(_A       ,					_A     )[0]

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





     @torch.no_grad()
     def __call__(self   :					str       ,					_A   :					Union[torch.FloatTensor, PIL.Image.Image]       ,					_A   :					Union[torch.FloatTensor, PIL.Image.Image]       ,					_A   :					Optional[str] = None       ,					_A   :					Optional[str] = None       ,					_A   :					Optional[int] = 5_1_2       ,					_A   :					Optional[int] = 5_1_2       ,					_A   :					float = 0.6       ,					_A   :					Optional[int] = 5_0       ,					_A   :					Optional[float] = 7.5       ,					_A   :					Optional[int] = 1       ,					_A   :					float = 0.0       ,					_A   :					Optional[float] = 1_0_0       ,					_A   :					Optional[torch.Generator] = None       ,					_A   :					Optional[str] = "pil"       ,					_A   :					bool = True       ,					_A   :					float = 0.8       ,					_A   :					float = 0.1       ,					_A   :					float = 0.1       ,					)		->		Any:
         if isinstance(_A       ,					_A     ) and len(_A     ) != batch_size:
             raise ValueError(f'You have passed {batch_size} batch_size, but only {len(_A     )} generators.'     )

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

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

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

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

         snake_case				  =      self.tokenizer(
             _A       ,					padding="max_length"       ,					max_length=self.tokenizer.model_max_length       ,					truncation=_A       ,					return_tensors="pt"       ,					)
         snake_case				  =      self.text_encoder(style_text_input.input_ids.to(self.device     )     )[0]

         snake_case				  =      slerp(_A       ,					_A       ,					_A     )

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

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

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

         snake_case  ,							snake_case				  =      self.get_timesteps(_A       ,					_A       ,					self.device     )
         snake_case				  =      timesteps[:1].repeat(_A     )

         # Preprocess image
         snake_case				  =      preprocess(_A       ,					_A       ,					_A     )
         snake_case				  =      self.prepare_latents(
             _A       ,					_A       ,					_A       ,					text_embeddings.dtype       ,					self.device       ,					_A     )

         snake_case				  =      preprocess(_A       ,					_A       ,					_A     )
         snake_case				  =      self.prepare_latents(
             _A       ,					_A       ,					_A       ,					text_embeddings.dtype       ,					self.device       ,					_A     )

         snake_case				  =      slerp(_A       ,					_A       ,					_A     )

         if clip_guidance_scale > 0:
             snake_case				  =      self.get_clip_image_embeddings(_A       ,					_A     )
             snake_case				  =      self.get_clip_image_embeddings(_A       ,					_A     )
             snake_case				  =      slerp(
                 _A       ,					_A       ,					_A     )

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

             # For classifier free guidance, we need to do two forward passes.
             # Here we concatenate the unconditional and text embeddings into a single batch
             # to avoid doing two forward passes
             snake_case				  =      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`.
         snake_case				  =      (batch_size, self.unet.config.in_channels, height // 8, width // 8)
         snake_case				  =      text_embeddings.dtype
         if latents is None:
             if self.device.type == "mps":
                 # randn does not work reproducibly on mps
                 snake_case				  =      torch.randn(_A       ,					generator=_A       ,					device="cpu"       ,					dtype=_A     ).to(
                     self.device     )
             else:
                 snake_case				  =      torch.randn(_A       ,					generator=_A       ,					device=self.device       ,					dtype=_A     )
         else:
             if latents.shape != latents_shape:
                 raise ValueError(f'Unexpected latents shape, got {latents.shape}, expected {latents_shape}'     )
             snake_case				  =      latents.to(self.device     )

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

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

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

                 # predict the noise residual
                 snake_case				  =      self.unet(_A       ,					_A       ,					encoder_hidden_states=_A     ).sample

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

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

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

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

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

         if output_type == "pil":
             snake_case				  =      self.numpy_to_pil(_A     )

         if not return_dict:
             return (image, None)

         return StableDiffusionPipelineOutput(images=_A       ,					nsfw_content_detected=_A     )

