




import unittest

import numpy as np
import timeout_decorator  # noqa

from transformers import OPTConfig, is_flax_available
from transformers.testing_utils import require_flax, require_sentencepiece, slow

from ...generation.test_flax_utils import FlaxGenerationTesterMixin
from ...test_modeling_flax_common import FlaxModelTesterMixin, ids_tensor


if is_flax_available():
 import os

 # The slow tests are often failing with OOM error on GPU
 # This makes JAX allocate exactly what is needed on demand, and deallocate memory that is no longer needed
 # but will be slower as stated here https://jax.readthedocs.io/en/latest/gpu_memory_allocation.html
 SCREAMING_SNAKE_CASE_						:					Dict			   =       "platform"

 import jax
 import jax.numpy as jnp

 from transformers import FlaxOPTForCausalLM, FlaxOPTModel, GPTaTokenizer
def        __lowerCAmelCase (UpperCAmelCase_   ,    UpperCAmelCase_   ,    UpperCAmelCase_=None   ,    UpperCAmelCase_=None		):

    '''simple docstring'''





    if attention_mask is None:
        a__      =       np.where(input_ids != config.pad_token_id   ,    1   ,    0		)
    return {
        "input_ids": input_ids,
        "attention_mask": attention_mask,
    }







@require_flax
class 		__UpperCAmelCase:





     """simple docstring"""



     def __init__(      self :       Any						,				A :       List[Any]						,				A :       Tuple=13						,				A :       Optional[int]=7						,				A :       str=True						,				A :       str=False						,				A :       List[Any]=99						,				A :       Optional[Any]=16						,				A :       Union[str, Any]=2						,				A :       List[Any]=4						,				A :       Tuple=4						,				A :       Any="gelu"						,				A :       Any=0.1						,				A :       Tuple=0.1						,				A :       List[str]=20						,				A :       int=2						,				A :       Optional[int]=1						,				A :       Dict=0						,				A :       Tuple=16						,				A :       int=16						,				A :       int=0.0_2						,				)			->			Optional[int]:






         """simple docstring"""
         a__      =       parent
         a__      =       batch_size
         a__      =       seq_length
         a__      =       is_training
         a__      =       use_labels
         a__      =       vocab_size
         a__      =       hidden_size
         a__      =       num_hidden_layers
         a__      =       num_attention_heads
         a__      =       intermediate_size
         a__      =       hidden_act
         a__      =       hidden_dropout_prob
         a__      =       attention_probs_dropout_prob
         a__      =       max_position_embeddings
         a__      =       eos_token_id
         a__      =       pad_token_id
         a__      =       bos_token_id
         a__      =       embed_dim
         a__      =       word_embed_proj_dim
         a__      =       initializer_range
         a__      =       False

     def 				__SCREAMING_SNAKE_CASE		(      self :       Tuple    )			->			Optional[int]:






         """simple docstring"""
         a__      =       np.clip(ids_tensor([self.batch_size, self.seq_length - 1]						,				self.vocab_size    )						,				3						,				self.vocab_size    )
         a__      =       np.concatenate((input_ids, 2 * np.ones((self.batch_size, 1)						,				dtype=np.intaa    ))						,				-1    )

         a__      =       OPTConfig(
             vocab_size=self.vocab_size						,				hidden_size=self.hidden_size						,				num_hidden_layers=self.num_hidden_layers						,				num_attention_heads=self.num_attention_heads						,				ffn_dim=self.intermediate_size						,				dropout=self.hidden_dropout_prob						,				attention_dropout=self.attention_probs_dropout_prob						,				max_position_embeddings=self.max_position_embeddings						,				eos_token_id=self.eos_token_id						,				bos_token_id=self.bos_token_id						,				pad_token_id=self.pad_token_id						,				embed_dim=self.embed_dim						,				is_encoder_decoder=a_						,				word_embed_proj_dim=self.word_embed_proj_dim						,				initializer_range=self.initializer_range						,				use_cache=a_						,				)
         a__      =       prepare_opt_inputs_dict(a_						,				a_    )
         return config, inputs_dict

     def 				__SCREAMING_SNAKE_CASE		(      self :       Any    )			->			Tuple:






         """simple docstring"""
         a__ ,		a__      =       self.prepare_config_and_inputs()
         return config, inputs_dict

     def 				__SCREAMING_SNAKE_CASE		(      self :       Any						,				A :       Optional[Any]						,				A :       List[str]						,				A :       Optional[Any]    )			->			int:






         """simple docstring"""
         a__      =       20
         a__      =       model_class_name(a_    )

         a__      =       inputs_dict["""input_ids"""]
         a__      =       inputs_dict["""attention_mask"""]

         a__      =       model.init_cache(input_ids.shape[0]						,				a_    )
         a__      =       jnp.ones((input_ids.shape[0], max_length)						,				dtype="""i4"""    )

         a__      =       jnp.broadcast_to(
             jnp.arange(input_ids.shape[-1] - 1    )[None, :]						,				(input_ids.shape[0], input_ids.shape[-1] - 1)						,				)
         a__      =       model(
             input_ids[:, :-1]						,				attention_mask=a_						,				past_key_values=a_						,				position_ids=a_						,				)

         a__      =       jnp.array(input_ids.shape[0] * [[input_ids.shape[-1] - 1]]						,				dtype="""i4"""    )
         a__      =       model(
             input_ids[:, -1:]						,				attention_mask=a_						,				past_key_values=outputs_cache.past_key_values						,				position_ids=a_						,				)

         a__      =       model(a_    )

         a__      =       np.max(np.abs((outputs_cache_next[0][:, -1, :5] - outputs[0][:, -1, :5])    )    )
         self.parent.assertTrue(diff < 1E-3						,				msg=f"""Max diff is {diff}"""    )






     def 				__SCREAMING_SNAKE_CASE		(      self :       Optional[Any]						,				A :       Optional[Any]						,				A :       str						,				A :       str    )			->			Tuple:






         """simple docstring"""
         a__      =       20
         a__      =       model_class_name(a_    )

         a__ ,		a__      =       (
             inputs_dict["""input_ids"""],
             inputs_dict["""attention_mask"""],
         )

         a__      =       jnp.concatenate(
             [
                 attention_mask,
                 jnp.zeros((attention_mask.shape[0], max_length - attention_mask.shape[1])    ),
             ]						,				axis=-1						,				)

         a__      =       model.init_cache(input_ids.shape[0]						,				a_    )
         a__      =       jnp.broadcast_to(
             jnp.arange(input_ids.shape[-1] - 1    )[None, :]						,				(input_ids.shape[0], input_ids.shape[-1] - 1)						,				)

         a__      =       model(
             input_ids[:, :-1]						,				attention_mask=a_						,				past_key_values=a_						,				position_ids=a_						,				)
         a__      =       jnp.array(input_ids.shape[0] * [[input_ids.shape[-1] - 1]]						,				dtype="""i4"""    )
         a__      =       model(
             input_ids[:, -1:]						,				past_key_values=outputs_cache.past_key_values						,				attention_mask=a_						,				position_ids=a_						,				)

         a__      =       model(a_						,				attention_mask=a_    )

         a__      =       np.max(np.abs((outputs_cache_next[0][:, -1, :5] - outputs[0][:, -1, :5])    )    )
         self.parent.assertTrue(diff < 1E-3						,				msg=f"""Max diff is {diff}"""    )







@require_flax
class 		__UpperCAmelCase(		SCREAMING_SNAKE_CASE__       ,	unittest.TestCase       ,	SCREAMING_SNAKE_CASE__ ):





     """simple docstring"""



     __UpperCAmelCase	: str						=    (FlaxOPTModel, FlaxOPTForCausalLM) if is_flax_available() else ()
     __UpperCAmelCase	: Optional[int]						=    () if is_flax_available() else ()

     def 				__SCREAMING_SNAKE_CASE		(      self :       str    )			->			List[str]:






         """simple docstring"""
         a__      =       FlaxOPTModelTester(self    )

     def 				__SCREAMING_SNAKE_CASE		(      self :       Any    )			->			List[Any]:






         """simple docstring"""
         a__ ,		a__      =       self.model_tester.prepare_config_and_inputs()
         for model_class in self.all_model_classes:
             self.model_tester.check_use_cache_forward(a_						,				a_						,				a_    )

     def 				__SCREAMING_SNAKE_CASE		(      self :       Optional[int]    )			->			List[Any]:






         """simple docstring"""
         a__ ,		a__      =       self.model_tester.prepare_config_and_inputs()
         for model_class in self.all_model_classes:
             self.model_tester.check_use_cache_forward_with_attn_mask(a_						,				a_						,				a_    )






     @slow
     def 				__SCREAMING_SNAKE_CASE		(      self :       List[Any]    )			->			List[str]:






         """simple docstring"""



         for model_class_name in self.all_model_classes:
             a__      =       model_class_name.from_pretrained("""facebook/opt-125m"""    )
             a__      =       np.ones((1, 1)    ) * model.config.eos_token_id
             a__      =       model(a_    )
             self.assertIsNotNone(a_    )







@require_sentencepiece
@require_flax
class 		__UpperCAmelCase(		unittest.TestCase ):





     """simple docstring"""



     @slow
     def 				__SCREAMING_SNAKE_CASE		(      self :       Dict    )			->			Union[str, Any]:






         """simple docstring"""
         a__      =       FlaxOPTModel.from_pretrained("""facebook/opt-350m"""    )
         a__      =       jnp.array([[0, 3_14_14, 2_32, 3_28, 7_40, 11_40, 1_26_95, 69, 4_60_78, 15_88, 2]]    )
         a__      =       model(input_ids=a_    ).last_hidden_state
         a__      =       (1, 11, 5_12)
         self.assertEqual(output.shape						,				a_    )
         a__      =       jnp.array(
             [[-0.2_8_6_7, -1.9_2_5_6, -0.3_0_6_2], [-1.2_7_1_1, -0.1_3_3_7, -0.1_8_9_7], [0.4_1_0_9, 0.1_1_8_7, -1.3_1_4_2]]    )
         self.assertTrue(jnp.allclose(output[:, :3, :3]						,				a_						,				atol=4E-2    )    )







@require_flax
@slow
class 		__UpperCAmelCase(		unittest.TestCase ):





     """simple docstring"""



     def 				__SCREAMING_SNAKE_CASE		(      self :       Optional[Any]    )			->			Optional[Any]:






         """simple docstring"""
         super().setUp()
         a__      =       """facebook/opt-350m"""






     def 				__SCREAMING_SNAKE_CASE		(      self :       str    )			->			List[str]:






         """simple docstring"""
         a__      =       FlaxOPTForCausalLM.from_pretrained(self.path_model    )
         a__      =       GPTaTokenizer.from_pretrained(self.path_model    )

         a__      =       [
             """Today is a beautiful day and I want to""",
             """In the city of""",
             """Paris is the capital of France and""",
             """Computers and mobile phones have taken""",
         ]
         # verify that prompt without BOS token is identical to Metaseq -> add_special_tokens=False
         a__      =       tokenizer(a_						,				return_tensors="""jax"""						,				padding=a_						,				add_special_tokens=a_    )
         a__      =       model(inputs.input_ids						,				attention_mask=inputs.attention_mask    )[0].mean(axis=-1    )
         a__      =       jnp.array(
             [
                 [1.3_8_5_1, -1_3.8_9_2_3, -1_0.5_2_2_9, -1_0.7_5_3_3, -0.2_3_0_9, -1_0.2_3_8_4, -0.5_3_6_5, -9.0_9_4_7, -5.1_6_7_0],
                 [-4.7_0_7_3, -1_0.6_2_7_6, -3.9_4_1_5, -2_1.5_2_4_2, -0.2_8_2_2, -0.2_8_2_2, -0.2_8_2_2, -0.2_8_2_2, -0.2_8_2_2],
                 [0.6_2_4_7, -3.4_2_2_9, -8.9_1_7_9, -1.4_2_9_7, -1_4.1_6_5_0, 1.4_1_4_6, -9.0_2_1_8, -0.2_7_0_3, -0.2_7_0_3],
                 [6.4_7_8_3, -1.9_9_1_3, -1_0.7_9_2_6, -2.3_3_3_6, 1.5_0_9_2, -0.9_9_7_4, -6.8_2_1_3, 1.3_4_7_7, 1.3_4_7_7],
             ]    )
         self.assertTrue(jnp.allclose(a_						,				a_						,				atol=4E-2    )    )

         a__      =       jax.jit(a_    )
         a__      =       model(inputs.input_ids						,				attention_mask=inputs.attention_mask    )[0].mean(axis=-1    )
         self.assertTrue(jnp.allclose(a_						,				a_						,				atol=4E-2    )    )







@require_flax
@slow
class 		__UpperCAmelCase(		unittest.TestCase ):





     """simple docstring"""



     @property
     def 				__SCREAMING_SNAKE_CASE		(      self :       Any    )			->			Tuple:






         """simple docstring"""
         return [
             "Today is a beautiful day and I want",
             "In the city of",
             "Paris is the capital of France and",
             "Computers and mobile phones have taken",
         ]

     def 				__SCREAMING_SNAKE_CASE		(      self :       Union[str, Any]    )			->			Any:






         """simple docstring"""
         a__      =       """facebook/opt-125m"""

         a__      =       [
             """Today is a beautiful day and I want to""",
             """In the city of New York, the city""",
             """Paris is the capital of France and the capital""",
             """Computers and mobile phones have taken over the""",
         ]

         a__      =       []

         a__      =       FlaxOPTForCausalLM.from_pretrained(a_    )
         a__      =       GPTaTokenizer.from_pretrained(a_    )

         for prompt in self.prompts:
             a__      =       tokenizer(a_						,				return_tensors="""jax"""    ).input_ids

             a__      =       model.generate(a_						,				max_length=10    )
             a__      =       generated_ids[0]

             a__      =       tokenizer.batch_decode(a_						,				skip_special_tokens=a_    )
             predicted_outputs += generated_string

         self.assertListEqual(a_						,				a_    )

     def 				__SCREAMING_SNAKE_CASE		(      self :       Any    )			->			int:






         """simple docstring"""
         a__      =       """facebook/opt-350m"""

         a__      =       [
             """Today is a beautiful day and I want to""",
             """In the city of San Francisco, the city""",
             """Paris is the capital of France and the capital""",
             """Computers and mobile phones have taken over the""",
         ]

         a__      =       []
         a__      =       FlaxOPTForCausalLM.from_pretrained(a_    )
         a__      =       GPTaTokenizer.from_pretrained(a_    )

         for prompt in self.prompts:
             a__      =       tokenizer(a_						,				return_tensors="""jax"""    ).input_ids

             a__      =       model.generate(a_						,				max_length=10    )
             a__      =       generated_ids[0]

             a__      =       tokenizer.batch_decode(a_						,				skip_special_tokens=a_    )
             predicted_outputs += generated_string

         self.assertListEqual(a_						,				a_    )

     def 				__SCREAMING_SNAKE_CASE		(      self :       Optional[int]    )			->			Optional[int]:






         """simple docstring"""
         a__      =       """facebook/opt-125m"""
         a__      =       [
             """Today is a beautiful day and I want to thank""",
             """In the city of Rome Canaver Canaver Canaver Canaver""",
         ]
         a__      =       FlaxOPTForCausalLM.from_pretrained(a_    )
         a__      =       GPTaTokenizer.from_pretrained(a_    )
         a__      =       tokenizer(
             [
                 """Today is a beautiful day and I want to""",
                 """In the city of""",
             ]						,				return_tensors="""jax"""						,				padding=a_						,				)

         a__      =       jax.jit(model.generate    )

         a__      =       jit_generate(inputs["""input_ids"""]						,				attention_mask=inputs["""attention_mask"""]    ).sequences

         a__      =       tokenizer.batch_decode(a_						,				skip_special_tokens=a_    )

         self.assertIsNotNone(a_						,				a_    )






     def 				__SCREAMING_SNAKE_CASE		(      self :       Any    )			->			Optional[Any]:






         """simple docstring"""
         a__      =       """facebook/opt-350m"""

         a__      =       GPTaTokenizer.from_pretrained(a_    )
         a__      =       FlaxOPTForCausalLM.from_pretrained(a_    )

         a__      =       """left"""

         # use different length sentences to test batching
         a__      =       [
             """Hello, my dog is a little""",
             """Today, I""",
         ]

         a__      =       tokenizer(a_						,				return_tensors="""jax"""						,				padding=a_    )
         a__      =       inputs["""input_ids"""]

         a__      =       model.generate(input_ids=a_						,				attention_mask=inputs["""attention_mask"""]						,				trace=a_    )

         a__      =       tokenizer(sentences[0]						,				return_tensors="""jax"""    ).input_ids
         a__      =       model.generate(input_ids=a_    )

         a__      =       inputs_non_padded.shape[-1] - inputs["""attention_mask"""][-1].sum()
         a__      =       tokenizer(sentences[1]						,				return_tensors="""jax"""    ).input_ids
         a__      =       model.generate(input_ids=a_						,				max_length=model.config.max_length - num_paddings    )

         a__      =       tokenizer.batch_decode(outputs[0]						,				skip_special_tokens=a_    )
         a__      =       tokenizer.decode(output_non_padded[0][0]						,				skip_special_tokens=a_    )
         a__      =       tokenizer.decode(output_padded[0][0]						,				skip_special_tokens=a_    )

         a__      =       [
             """Hello, my dog is a little bit of a dork.\nI'm a little bit""",
             """Today, I was in the middle of a conversation with a friend about the""",
         ]
         self.assertListEqual(a_						,				a_    )
         self.assertListEqual(a_						,				[non_padded_sentence, padded_sentence]    )












from logging import getLogger
from typing import Any, Callable, List, Optional, Union

import numpy as np
import PIL
import torch

from ...schedulers import DDPMScheduler
from ..onnx_utils import ORT_TO_NP_TYPE, OnnxRuntimeModel
from ..pipeline_utils import ImagePipelineOutput
from . import StableDiffusionUpscalePipeline


SCREAMING_SNAKE_CASE_						:					Optional[int]			   =       getLogger(__name__)


SCREAMING_SNAKE_CASE_						:					Union[str, Any]			   =       4
SCREAMING_SNAKE_CASE_						:					List[str]			   =       7

SCREAMING_SNAKE_CASE_						:					List[str]			   =       {
    '''float16''': torch.floataa,
    '''float32''': torch.floataa,
}
def        __lowerCAmelCase (UpperCAmelCase_		):
    if isinstance(UpperCAmelCase_   ,    torch.Tensor		):
        return image
    elif isinstance(UpperCAmelCase_   ,    PIL.Image.Image		):
        a__      =       [image]

    if isinstance(image[0]   ,    PIL.Image.Image		):
        a__ ,		a__      =       image[0].size
        a__ ,		a__      =       (x - x % 64 for x in (w, h))  # resize to integer multiple of 32

        a__      =       [np.array(i.resize((w, h)		)		)[None, :] for i in image]
        a__      =       np.concatenate(UpperCAmelCase_   ,    axis=0		)
        a__      =       np.array(UpperCAmelCase_		).astype(np.floataa		) / 2_5_5.0
        a__      =       image.transpose(0   ,    3   ,    1   ,    2		)
        a__      =       2.0 * image - 1.0
        a__      =       torch.from_numpy(UpperCAmelCase_		)
    elif isinstance(image[0]   ,    torch.Tensor		):
        a__      =       torch.cat(UpperCAmelCase_   ,    dim=0		)

    return image







class 		__UpperCAmelCase(		lowercase__ ):





     """simple docstring"""



     def __init__(      self :       int						,				A :       OnnxRuntimeModel						,				A :       OnnxRuntimeModel						,				A :       Any						,				A :       OnnxRuntimeModel						,				A :       DDPMScheduler						,				A :       Any						,				A :       int = 3_50						,				)			->			List[Any]:






         """simple docstring"""
         super().__init__(
             vae=A						,				text_encoder=A						,				tokenizer=A						,				unet=A						,				low_res_scheduler=A						,				scheduler=A						,				safety_checker=A						,				feature_extractor=A						,				watermarker=A						,				max_noise_level=A						,				)

     def __call__(      self :       Optional[int]						,				A :       Union[str, List[str]]						,				A :       Union[torch.FloatTensor, PIL.Image.Image, List[PIL.Image.Image]]						,				A :       int = 75						,				A :       float = 9.0						,				A :       int = 20						,				A :       Optional[Union[str, List[str]]] = None						,				A :       Optional[int] = 1						,				A :       float = 0.0						,				A :       Optional[Union[torch.Generator, List[torch.Generator]]] = None						,				A :       Optional[torch.FloatTensor] = None						,				A :       Optional[np.ndarray] = None						,				A :       Optional[np.ndarray] = None						,				A :       Optional[str] = "pil"						,				A :       bool = True						,				A :       Optional[Callable[[int, int, torch.FloatTensor], None]] = None						,				A :       Optional[int] = 1						,				)			->			Tuple:






         """simple docstring"""
         self.check_inputs(A						,				A						,				A						,				A    )

         # 2. Define call parameters
         if prompt is not None and isinstance(A						,				A    ):
             a__      =       1
         elif prompt is not None and isinstance(A						,				A    ):
             a__      =       len(A    )
         else:
             a__      =       prompt_embeds.shape[0]

         a__      =       self._execution_device
         # 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.
         a__      =       guidance_scale > 1.0

         # 3. Encode input prompt
         a__      =       self._encode_prompt(
             A						,				A						,				A						,				A						,				A						,				prompt_embeds=A						,				negative_prompt_embeds=A						,				)

         a__      =       ORT_TO_PT_TYPE[str(text_embeddings.dtype    )]

         # 4. Preprocess image
         a__      =       preprocess(A    )
         a__      =       image.cpu()

         # 5. set timesteps
         self.scheduler.set_timesteps(A						,				device=A    )
         a__      =       self.scheduler.timesteps

         # 5. Add noise to image
         a__      =       torch.tensor([noise_level]						,				dtype=torch.long						,				device=A    )
         a__      =       torch.randn(image.shape						,				generator=A						,				device=A						,				dtype=A    )
         a__      =       self.low_res_scheduler.add_noise(A						,				A						,				A    )

         a__      =       2 if do_classifier_free_guidance else 1
         a__      =       np.concatenate([image] * batch_multiplier * num_images_per_prompt    )
         a__      =       np.concatenate([noise_level] * image.shape[0]    )

         # 6. Prepare latent variables
         a__ ,		a__      =       image.shape[2:]
         a__      =       self.prepare_latents(
             batch_size * num_images_per_prompt						,				A						,				A						,				A						,				A						,				A						,				A						,				A						,				)

         # 7. Check that sizes of image and latents match
         a__      =       image.shape[1]
         if NUM_LATENT_CHANNELS + num_channels_image != NUM_UNET_INPUT_CHANNELS:
             raise ValueError(
                 """Incorrect configuration settings! The config of `pipeline.unet` expects"""
                 f""" {NUM_UNET_INPUT_CHANNELS} but received `num_channels_latents`: {NUM_LATENT_CHANNELS} +"""
                 f""" `num_channels_image`: {num_channels_image} """
                 f""" = {NUM_LATENT_CHANNELS+num_channels_image}. Please verify the config of"""
                 """ `pipeline.unet` or your `image` input."""    )

         # 8. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline
         a__      =       self.prepare_extra_step_kwargs(A						,				A    )

         a__      =       next(
             (input.type for input in self.unet.model.get_inputs() if input.name == """timestep""")						,				"""tensor(float)"""    )
         a__      =       ORT_TO_NP_TYPE[timestep_dtype]

         # 9. Denoising loop
         a__      =       len(A    ) - num_inference_steps * self.scheduler.order
         with self.progress_bar(total=A    ) as progress_bar:
             for i, t in enumerate(A    ):
                 # expand the latents if we are doing classifier free guidance
                 a__      =       np.concatenate([latents] * 2    ) if do_classifier_free_guidance else latents

                 # concat latents, mask, masked_image_latents in the channel dimension
                 a__      =       self.scheduler.scale_model_input(A						,				A    )
                 a__      =       np.concatenate([latent_model_input, image]						,				axis=1    )

                 # timestep to tensor
                 a__      =       np.array([t]						,				dtype=A    )

                 # predict the noise residual
                 a__      =       self.unet(
                     sample=A						,				timestep=A						,				encoder_hidden_states=A						,				class_labels=noise_level.astype(np.intaa    )						,				)[0]

                 # perform guidance
                 if do_classifier_free_guidance:
                     a__ ,		a__      =       np.split(A						,				2    )
                 a__      =       noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)

                 # compute the previous noisy sample x_t -> x_t-1
                 a__      =       self.scheduler.step(
                     torch.from_numpy(A    )						,				A						,				A						,				**A    ).prev_sample

                 # call the callback, if provided
                 if i == len(A    ) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0):
                     progress_bar.update()
                     if callback is not None and i % callback_steps == 0:
                         callback(A						,				A						,				A    )

        # 10. Post-processing
         a__      =       self.decode_latents(latents.float()    )

         # 11. Convert to PIL
         if output_type == "pil":
             a__      =       self.numpy_to_pil(A    )

         if not return_dict:
             return (image,)

         return ImagePipelineOutput(images=A    )

     def 				__SCREAMING_SNAKE_CASE		(      self :       List[Any]						,				A :       Dict    )			->			List[str]:






         """simple docstring"""
         a__      =       1 / 0.0_8_3_3_3 * latents
         a__      =       self.vae(latent_sample=A    )[0]
         a__      =       np.clip(image / 2 + 0.5						,				0						,				1    )
         a__      =       image.transpose((0, 2, 3, 1)    )
         return image






     def 				__SCREAMING_SNAKE_CASE		(      self :       List[str]						,				A :       Union[str, List[str]]						,				A :       str						,				A :       Optional[int]						,				A :       bool						,				A :       Optional[str]						,				A :       Optional[torch.FloatTensor] = None						,				A :       Optional[torch.FloatTensor] = None						,				)			->			Tuple:






         """simple docstring"""
         if prompt is not None and isinstance(A						,				A    ):
             a__      =       1
         elif prompt is not None and isinstance(A						,				A    ):
             a__      =       len(A    )
         else:
             a__      =       prompt_embeds.shape[0]

         if prompt_embeds is None:
             a__      =       self.tokenizer(
                 A						,				padding="""max_length"""						,				max_length=self.tokenizer.model_max_length						,				truncation=A						,				return_tensors="""pt"""						,				)
             a__      =       text_inputs.input_ids
             a__      =       self.tokenizer(A						,				padding="""longest"""						,				return_tensors="""pt"""    ).input_ids

             if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not torch.equal(
                 A						,				A    ):
                 a__      =       self.tokenizer.batch_decode(
                     untruncated_ids[:, self.tokenizer.model_max_length - 1 : -1]    )
                 logger.warning(
                     """The following part of your input was truncated because CLIP can only handle sequences up to"""
                     f""" {self.tokenizer.model_max_length} tokens: {removed_text}"""    )

             # no positional arguments to text_encoder
             a__      =       self.text_encoder(
                 input_ids=text_input_ids.int().to(A    )						,				)
             a__      =       prompt_embeds[0]

         a__ ,		a__ ,		a__      =       prompt_embeds.shape
         # duplicate text embeddings for each generation per prompt, using mps friendly method
         a__      =       prompt_embeds.repeat(1						,				A    )
         a__      =       prompt_embeds.reshape(bs_embed * num_images_per_prompt						,				A						,				-1    )

         # get unconditional embeddings for classifier free guidance
         if do_classifier_free_guidance and negative_prompt_embeds is None:
             a__      =       42
             if negative_prompt is None:
                 a__      =       [""""""] * batch_size
             elif type(A    ) is not type(A    ):
                 raise TypeError(
                     f"""`negative_prompt` should be the same type to `prompt`, but got {type(A    )} !="""
                     f""" {type(A    )}."""    )
             elif isinstance(A						,				A    ):
                 a__      =       [negative_prompt]
             elif batch_size != len(A    ):
                 raise ValueError(
                     f"""`negative_prompt`: {negative_prompt} has batch size {len(A    )}, but `prompt`:"""
                     f""" {prompt} has batch size {batch_size}. Please make sure that passed `negative_prompt` matches"""
                     """ the batch size of `prompt`."""    )
             else:
                 a__      =       negative_prompt

             a__      =       text_input_ids.shape[-1]
             a__      =       self.tokenizer(
                 A						,				padding="""max_length"""						,				max_length=A						,				truncation=A						,				return_tensors="""pt"""						,				)

             # if hasattr(uncond_input, "attention_mask"):
             #     attention_mask = uncond_input.attention_mask.to(device)
             # else:
             #     attention_mask = None

             a__      =       self.text_encoder(
                 input_ids=uncond_input.input_ids.int().to(A    )						,				)
             a__      =       uncond_embeddings[0]

         if do_classifier_free_guidance:
             a__      =       uncond_embeddings.shape[1]
             # duplicate unconditional embeddings for each generation per prompt, using mps friendly method
             a__      =       uncond_embeddings.repeat(1						,				A    )
             a__      =       uncond_embeddings.reshape(batch_size * num_images_per_prompt						,				A						,				-1    )

             # 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
             a__      =       np.concatenate([uncond_embeddings, prompt_embeds]    )

         return prompt_embeds






