






"""simple docstring"""







import collections
import tempfile
import unittest

import numpy as np

from transformers.testing_utils import (
    is_pt_flax_cross_test,
    require_flax,
    require_torch,
    require_vision,
    slow,
    torch_device,
)
from transformers.utils import is_flax_available, is_torch_available, is_vision_available

from ...test_modeling_flax_common import floats_tensor, ids_tensor, random_attention_mask
from ..bert.test_modeling_flax_bert import FlaxBertModelTester
from ..clip.test_modeling_flax_clip import FlaxCLIPVisionModelTester
from ..vit.test_modeling_flax_vit import FlaxViTModelTester


if is_flax_available():
			from transformers import (
			    FlaxBertModel,
			    FlaxCLIPVisionModel,
			    FlaxVisionTextDualEncoderModel,
			    FlaxViTModel,
			    VisionTextDualEncoderConfig,
			    VisionTextDualEncoderProcessor,
			)
			from transformers.modeling_flax_pytorch_utils import (
			    convert_pytorch_state_dict_to_flax,
			    load_flax_weights_in_pytorch_model,
			)


if is_torch_available():
			import torch

			from transformers import VisionTextDualEncoderModel

if is_vision_available():
			from PIL import Image





def 		lowercase__(A		)							->Tuple:




						"""simple docstring"""



						if isinstance(A__   ,						collections.abc.Iterable		):
												return x
						return (x, x)


@require_flax
class        __UpperCAmelCase:




						"""simple docstring"""




						def 						UpperCAmelCase_			(  self			,    snake_case__			,    snake_case__     ):
												'''simple docstring'''
												pass




						def 						UpperCAmelCase_			(  self     ):
												'''simple docstring'''
												pass




						def 						UpperCAmelCase_			(  self     ):
												'''simple docstring'''
												pass




						def 						UpperCAmelCase_			(  self			,    snake_case__			,    snake_case__			,    snake_case__     ):
												'''simple docstring'''
												lowercase__							: Tuple=			np.abs((a - b)     ).max()
												self.assertLessEqual(lowerCAmelCase_			,    lowerCAmelCase_			,    F'''Difference between torch and flax is {diff} (>= {tol}).'''     )




						def 						UpperCAmelCase_			(  self			,    snake_case__			,    snake_case__			,    snake_case__			,    snake_case__			,    snake_case__=None			,    **snake_case__     ):
												'''simple docstring'''
												lowercase__							: Tuple=			VisionTextDualEncoderConfig.from_vision_text_configs(lowerCAmelCase_			,    lowerCAmelCase_     )

												lowercase__							: List[Any]=			FlaxVisionTextDualEncoderModel(lowerCAmelCase_     )

												lowercase__							: List[Any]=			model(input_ids=lowerCAmelCase_			,    pixel_values=lowerCAmelCase_			,    attention_mask=lowerCAmelCase_     )

												self.assertEqual(output["text_embeds"].shape			,    (input_ids.shape[0], config.projection_dim)     )
												self.assertEqual(output["image_embeds"].shape			,    (pixel_values.shape[0], config.projection_dim)     )




						def 						UpperCAmelCase_			(  self			,    snake_case__			,    snake_case__			,    snake_case__			,    snake_case__			,    snake_case__=None			,    **snake_case__     ):
												'''simple docstring'''
												lowercase__							: Dict=			self.get_vision_text_model(lowerCAmelCase_			,    lowerCAmelCase_     )
												lowercase__							: Optional[Any]=			{"vision_model": vision_model, "text_model": text_model}
												lowercase__							: str=			FlaxVisionTextDualEncoderModel.from_vision_text_pretrained(**lowerCAmelCase_     )

												lowercase__							: Tuple=			model(input_ids=lowerCAmelCase_			,    pixel_values=lowerCAmelCase_			,    attention_mask=lowerCAmelCase_     )

												self.assertEqual(output["text_embeds"].shape			,    (input_ids.shape[0], model.config.projection_dim)     )
												self.assertEqual(output["image_embeds"].shape			,    (pixel_values.shape[0], model.config.projection_dim)     )




						def 						UpperCAmelCase_			(  self			,    snake_case__			,    snake_case__			,    snake_case__			,    snake_case__			,    snake_case__=None			,    **snake_case__     ):
												'''simple docstring'''
												lowercase__							: str=			self.get_vision_text_model(lowerCAmelCase_			,    lowerCAmelCase_     )
												lowercase__							: Tuple=			{"vision_model": vision_model, "text_model": text_model}
												lowercase__							: Optional[int]=			FlaxVisionTextDualEncoderModel.from_vision_text_pretrained(**lowerCAmelCase_     )

												lowercase__							: List[str]=			model(input_ids=lowerCAmelCase_			,    pixel_values=lowerCAmelCase_			,    attention_mask=lowerCAmelCase_     )
												lowercase__							: List[str]=			output[0]

												with tempfile.TemporaryDirectory() as tmpdirname:
																		model.save_pretrained(lowerCAmelCase_     )
																		lowercase__							: Optional[int]=			FlaxVisionTextDualEncoderModel.from_pretrained(lowerCAmelCase_     )

																		lowercase__							: Dict=			model(input_ids=lowerCAmelCase_			,    pixel_values=lowerCAmelCase_			,    attention_mask=lowerCAmelCase_     )
																		lowercase__							: Union[str, Any]=			after_output[0]
																		lowercase__							: Any=			np.amax(np.abs(out_a - out_a     )     )
																		self.assertLessEqual(lowerCAmelCase_			,    1e-3     )




						def 						UpperCAmelCase_			(  self			,    snake_case__			,    snake_case__			,    snake_case__			,    snake_case__			,    snake_case__=None			,    **snake_case__     ):
												'''simple docstring'''
												lowercase__							: Optional[int]=			self.get_vision_text_model(lowerCAmelCase_			,    lowerCAmelCase_     )
												lowercase__							: List[str]=			{"vision_model": vision_model, "text_model": text_model}
												lowercase__							: Any=			FlaxVisionTextDualEncoderModel.from_vision_text_pretrained(**lowerCAmelCase_     )

												lowercase__							: List[Any]=			model(
												    input_ids=lowerCAmelCase_			,    pixel_values=lowerCAmelCase_			,    attention_mask=lowerCAmelCase_			,    output_attentions=lowerCAmelCase_     )

												lowercase__							: List[str]=			output.vision_model_output.attentions
												self.assertEqual(len(lowerCAmelCase_     )			,    vision_config.num_hidden_layers     )

												# in ViT, the seq_len equals the number of patches + 1 (we add 1 for the [CLS] token)
												lowercase__							: str=			to_atuple(vision_model.config.image_size     )
												lowercase__							: Optional[Any]=			to_atuple(vision_model.config.patch_size     )
												lowercase__							: Optional[Any]=			(image_size[1] // patch_size[1]) * (image_size[0] // patch_size[0])
												lowercase__							: Dict=			num_patches + 1
												self.assertEqual(vision_attentions[0].shape[-3:]			,    (vision_config.num_attention_heads, seq_len, seq_len)     )

												lowercase__							: List[str]=			output.text_model_output.attentions
												self.assertEqual(len(lowerCAmelCase_     )			,    text_config.num_hidden_layers     )

												self.assertEqual(
												    text_attentions[0].shape[-3:]			,    (text_config.num_attention_heads, input_ids.shape[-1], input_ids.shape[-1])			,    )




						def 						UpperCAmelCase_			(  self			,    snake_case__			,    snake_case__			,    snake_case__     ):
												'''simple docstring'''
												pt_model.to(lowerCAmelCase_     )
												pt_model.eval()

												# prepare inputs
												lowercase__							: Dict=			inputs_dict
												lowercase__							: List[str]=			{k: torch.tensor(v.tolist()     ) for k, v in flax_inputs.items()}

												with torch.no_grad():
																		lowercase__							: int=			pt_model(**lowerCAmelCase_     ).to_tuple()

												lowercase__							: int=			fx_model(**lowerCAmelCase_     ).to_tuple()
												self.assertEqual(len(lowerCAmelCase_     )			,    len(lowerCAmelCase_     )			,    "Output lengths differ between Flax and PyTorch"     )
												for fx_output, pt_output in zip(fx_outputs[:4]			,    pt_outputs[:4]     ):
																		self.assert_almost_equals(lowerCAmelCase_			,    pt_output.numpy()			,    4e-2     )

												# PT -> Flax
												with tempfile.TemporaryDirectory() as tmpdirname:
																		pt_model.save_pretrained(lowerCAmelCase_     )
																		lowercase__							: List[Any]=			FlaxVisionTextDualEncoderModel.from_pretrained(lowerCAmelCase_			,    from_pt=lowerCAmelCase_     )

												lowercase__							: Optional[Any]=			fx_model_loaded(**lowerCAmelCase_     ).to_tuple()
												self.assertEqual(len(lowerCAmelCase_     )			,    len(lowerCAmelCase_     )			,    "Output lengths differ between Flax and PyTorch"     )
												for fx_output_loaded, pt_output in zip(fx_outputs_loaded[:4]			,    pt_outputs[:4]     ):
																		self.assert_almost_equals(lowerCAmelCase_			,    pt_output.numpy()			,    4e-2     )

												# Flax -> PT
												with tempfile.TemporaryDirectory() as tmpdirname:
																		fx_model.save_pretrained(lowerCAmelCase_     )
																		lowercase__							: Optional[Any]=			VisionTextDualEncoderModel.from_pretrained(lowerCAmelCase_			,    from_flax=lowerCAmelCase_     )

												pt_model_loaded.to(lowerCAmelCase_     )
												pt_model_loaded.eval()

												with torch.no_grad():
																		lowercase__							: Tuple=			pt_model_loaded(**lowerCAmelCase_     ).to_tuple()

												self.assertEqual(len(lowerCAmelCase_     )			,    len(lowerCAmelCase_     )			,    "Output lengths differ between Flax and PyTorch"     )
												for fx_output, pt_output_loaded in zip(fx_outputs[:4]			,    pt_outputs_loaded[:4]     ):
																		self.assert_almost_equals(lowerCAmelCase_			,    pt_output_loaded.numpy()			,    4e-2     )




						def 						UpperCAmelCase_			(  self			,    snake_case__			,    snake_case__			,    snake_case__     ):
												'''simple docstring'''
												lowercase__							: Union[str, Any]=			VisionTextDualEncoderConfig.from_vision_text_configs(lowerCAmelCase_			,    lowerCAmelCase_     )

												lowercase__							: List[Any]=			VisionTextDualEncoderModel(lowerCAmelCase_     )
												lowercase__							: Dict=			FlaxVisionTextDualEncoderModel(lowerCAmelCase_     )

												lowercase__							: Any=			convert_pytorch_state_dict_to_flax(pt_model.state_dict()			,    lowerCAmelCase_     )
												lowercase__							: Dict=			fx_state

												self.check_pt_flax_equivalence(lowerCAmelCase_			,    lowerCAmelCase_			,    lowerCAmelCase_     )




						def 						UpperCAmelCase_			(  self			,    snake_case__			,    snake_case__			,    snake_case__     ):
												'''simple docstring'''
												lowercase__							: str=			VisionTextDualEncoderConfig.from_vision_text_configs(lowerCAmelCase_			,    lowerCAmelCase_     )

												lowercase__							: List[Any]=			VisionTextDualEncoderModel(lowerCAmelCase_     )
												lowercase__							: Dict=			FlaxVisionTextDualEncoderModel(lowerCAmelCase_     )

												lowercase__							: List[Any]=			load_flax_weights_in_pytorch_model(lowerCAmelCase_			,    fx_model.params     )

												self.check_pt_flax_equivalence(lowerCAmelCase_			,    lowerCAmelCase_			,    lowerCAmelCase_     )




						def 						UpperCAmelCase_			(  self     ):
												'''simple docstring'''
												lowercase__							: str=			self.prepare_config_and_inputs()
												self.check_model_from_pretrained_configs(**lowerCAmelCase_     )




						def 						UpperCAmelCase_			(  self     ):
												'''simple docstring'''
												lowercase__							: Union[str, Any]=			self.prepare_config_and_inputs()
												self.check_vision_text_dual_encoder_from_pretrained(**lowerCAmelCase_     )




						def 						UpperCAmelCase_			(  self     ):
												'''simple docstring'''
												lowercase__							: Optional[Any]=			self.prepare_config_and_inputs()
												self.check_save_load(**lowerCAmelCase_     )




						def 						UpperCAmelCase_			(  self     ):
												'''simple docstring'''
												lowercase__							: str=			self.prepare_config_and_inputs()
												self.check_vision_text_output_attention(**lowerCAmelCase_     )




						@is_pt_flax_cross_test
						def 						UpperCAmelCase_			(  self     ):
												'''simple docstring'''
												lowercase__							: int=			self.prepare_config_and_inputs()
												lowercase__							: int=			config_inputs_dict.pop("vision_config"     )
												lowercase__							: int=			config_inputs_dict.pop("text_config"     )

												lowercase__							: Optional[Any]=			config_inputs_dict

												self.check_equivalence_pt_to_flax(lowerCAmelCase_			,    lowerCAmelCase_			,    lowerCAmelCase_     )
												self.check_equivalence_flax_to_pt(lowerCAmelCase_			,    lowerCAmelCase_			,    lowerCAmelCase_     )




						@slow
						def 						UpperCAmelCase_			(  self     ):
												'''simple docstring'''
												lowercase__							: int=			self.get_pretrained_model_and_inputs()

												lowercase__							: List[Any]=			model_a(**lowerCAmelCase_     )
												lowercase__							: int=			outputs[0]

												with tempfile.TemporaryDirectory() as tmp_dirname:
																		model_a.save_pretrained(lowerCAmelCase_     )
																		lowercase__							: List[str]=			FlaxVisionTextDualEncoderModel.from_pretrained(lowerCAmelCase_     )

																		lowercase__							: Any=			model_a(**lowerCAmelCase_     )
																		lowercase__							: Optional[int]=			after_outputs[0]
																		lowercase__							: int=			np.amax(np.abs(out_a - out_a     )     )
																		self.assertLessEqual(lowerCAmelCase_			,    1e-5     )


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




						"""simple docstring"""




						def 						UpperCAmelCase_			(  self     ):
												'''simple docstring'''
												lowercase__							: Optional[Any]=			FlaxVisionTextDualEncoderModel.from_vision_text_pretrained(
												    "hf-internal-testing/tiny-random-vit"			,    "hf-internal-testing/tiny-bert"			,    vision_from_pt=lowerCAmelCase_			,    text_from_pt=lowerCAmelCase_			,    )
												lowercase__							: List[Any]=			13
												lowercase__							: Optional[Any]=			floats_tensor(
												    [
												        batch_size,
												        model.config.vision_config.num_channels,
												        model.config.vision_config.image_size,
												        model.config.vision_config.image_size,
												    ]     )
												lowercase__							: Tuple=			ids_tensor([batch_size, 4]			,    model.config.text_config.vocab_size     )
												lowercase__							: Optional[Any]=			random_attention_mask([batch_size, 4]     )
												lowercase__							: str=			{"pixel_values": pixel_values, "input_ids": input_ids, "attention_mask": attention_mask}

												return model, inputs




						def 						UpperCAmelCase_			(  self			,    snake_case__			,    snake_case__     ):
												'''simple docstring'''
												lowercase__							: List[str]=			FlaxViTModel(lowerCAmelCase_     )
												lowercase__							: Any=			FlaxBertModel(lowerCAmelCase_     )
												return vision_model, text_model




						def 						UpperCAmelCase_			(  self     ):
												'''simple docstring'''
												lowercase__							: Union[str, Any]=			FlaxViTModelTester(self     )
												lowercase__							: Optional[int]=			FlaxBertModelTester(self     )
												lowercase__							: Any=			vit_model_tester.prepare_config_and_inputs()
												lowercase__							: Optional[Any]=			bert_model_tester.prepare_config_and_inputs()

												lowercase__							: Any=			vision_config_and_inputs

												lowercase__							: int=			text_config_and_inputs

												# make sure that cross attention layers are added
												return {
												    "text_config": text_config,
												    "vision_config": vision_config,
												    "pixel_values": pixel_values,
												    "attention_mask": attention_mask,
												    "input_ids": input_ids,
												    "token_type_ids": token_type_ids,
												}


@require_torch
class        __UpperCAmelCase(  __A    ,    unittest.TestCase				):




						"""simple docstring"""




						def 						UpperCAmelCase_			(  self     ):
												'''simple docstring'''
												lowercase__							: Union[str, Any]=			FlaxVisionTextDualEncoderModel.from_vision_text_pretrained(
												    "hf-internal-testing/tiny-random-clip"			,    "hf-internal-testing/tiny-bert"			,    vision_from_pt=lowerCAmelCase_			,    text_from_pt=lowerCAmelCase_			,    )
												lowercase__							: Any=			13
												lowercase__							: Union[str, Any]=			floats_tensor(
												    [
												        batch_size,
												        model.config.vision_config.num_channels,
												        model.config.vision_config.image_size,
												        model.config.vision_config.image_size,
												    ]     )
												lowercase__							: Optional[int]=			ids_tensor([batch_size, 4]			,    model.config.text_config.vocab_size     )
												lowercase__							: int=			random_attention_mask([batch_size, 4]     )
												lowercase__							: int=			{"pixel_values": pixel_values, "input_ids": input_ids, "attention_mask": attention_mask}

												return model, inputs




						def 						UpperCAmelCase_			(  self			,    snake_case__			,    snake_case__     ):
												'''simple docstring'''
												lowercase__							: Tuple=			FlaxCLIPVisionModel(lowerCAmelCase_     )
												lowercase__							: Optional[Any]=			FlaxBertModel(lowerCAmelCase_     )
												return vision_model, text_model




						def 						UpperCAmelCase_			(  self     ):
												'''simple docstring'''
												lowercase__							: List[Any]=			FlaxCLIPVisionModelTester(self     )
												lowercase__							: Union[str, Any]=			FlaxBertModelTester(self     )
												lowercase__							: List[str]=			clip_model_tester.prepare_config_and_inputs()
												lowercase__							: Union[str, Any]=			bert_model_tester.prepare_config_and_inputs()

												lowercase__							: Any=			vision_config_and_inputs

												lowercase__							: Dict=			text_config_and_inputs

												# make sure that cross attention layers are added
												return {
												    "text_config": text_config,
												    "vision_config": vision_config,
												    "pixel_values": pixel_values,
												    "attention_mask": attention_mask,
												    "input_ids": input_ids,
												    "token_type_ids": token_type_ids,
												}


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




						"""simple docstring"""




						@slow
						def 						UpperCAmelCase_			(  self     ):
												'''simple docstring'''
												lowercase__							: Any=			FlaxVisionTextDualEncoderModel.from_pretrained("clip-italian/clip-italian"			,    logit_scale_init_value=1.0     )
												lowercase__							: Tuple=			VisionTextDualEncoderProcessor.from_pretrained("clip-italian/clip-italian"     )

												lowercase__							: Optional[Any]=			Image.open("./tests/fixtures/tests_samples/COCO/000000039769.png"     )
												lowercase__							: int=			processor(
												    text=["una foto di un gatto", "una foto di un cane"]			,    images=lowerCAmelCase_			,    padding=lowerCAmelCase_			,    return_tensors="np"     )

												lowercase__							: Optional[int]=			model(**lowerCAmelCase_     )

												# verify the logits
												self.assertEqual(outputs.logits_per_image.shape			,    (inputs.pixel_values.shape[0], inputs.input_ids.shape[0])     )
												self.assertEqual(
												    outputs.logits_per_text.shape			,    (inputs.input_ids.shape[0], inputs.pixel_values.shape[0])			,    )

												lowercase__							: List[str]=			np.array([[1.2_28_47_27, 0.3_10_41_22]]     )

												self.assertTrue(np.allclose(outputs.logits_per_image			,    lowerCAmelCase_			,    atol=1e-3     )     )














"""simple docstring"""
from typing import Dict, List, Optional, Tuple, Union

import torch

from ...models import AutoencoderKL, TransformeraDModel
from ...schedulers import KarrasDiffusionSchedulers
from ...utils import randn_tensor
from ..pipeline_utils import DiffusionPipeline, ImagePipelineOutput







class      UpperCamelCase_			(__A	):





	def __init__( self   :					List[Any]      ,      lowerCAmelCase_   :					TransformeraDModel      ,      lowerCAmelCase_   :					AutoencoderKL      ,      lowerCAmelCase_   :					KarrasDiffusionSchedulers      ,      lowerCAmelCase_   :					Optional[Dict[int, str]] = None      ,      )		->     List[Any]:
				super().__init__()
				self.register_modules(transformer=lowerCAmelCase_      ,      vae=lowerCAmelCase_      ,      scheduler=lowerCAmelCase_		)

				# create a imagenet -> id dictionary for easier use
				UpperCAmelCase_	:					str     	=   {}
				if idalabel is not None:
							for key, value in idalabel.items():
										for label in value.split(","		):
													UpperCAmelCase_	:					Union[str, Any]     	=   int(lowerCAmelCase_		)
							UpperCAmelCase_	:					Optional[int]     	=   dict(sorted(self.labels.items()		)		)





	def 	_SCREAMING_SNAKE_CASE ( self   :					str      ,      lowerCAmelCase_   :					Union[str, List[str]]		)		->     List[int]:

				if not isinstance(lowerCAmelCase_      ,      lowerCAmelCase_		):
							UpperCAmelCase_	:					str     	=   list(lowerCAmelCase_		)

				for l in label:
							if l not in self.labels:
										raise ValueError(
										    f"""{l} does not exist. Please make sure to select one of the following labels: \n {self.labels}."""		)

				return [self.labels[l] for l in label]





	@torch.no_grad()
	def __call__( self   :					List[str]      ,      lowerCAmelCase_   :					List[int]      ,      lowerCAmelCase_   :					float = 4.0      ,      lowerCAmelCase_   :					Optional[Union[torch.Generator, List[torch.Generator]]] = None      ,      lowerCAmelCase_   :					int = 50      ,      lowerCAmelCase_   :					Optional[str] = "pil"      ,      lowerCAmelCase_   :					bool = True      ,      )		->     Union[ImagePipelineOutput, Tuple]:
				UpperCAmelCase_	:					Optional[int]     	=   len(lowerCAmelCase_		)
				UpperCAmelCase_	:					int     	=   self.transformer.config.sample_size
				UpperCAmelCase_	:					Optional[Any]     	=   self.transformer.config.in_channels

				UpperCAmelCase_	:					Tuple     	=   randn_tensor(
				    shape=(batch_size, latent_channels, latent_size, latent_size)      ,      generator=lowerCAmelCase_      ,      device=self.device      ,      dtype=self.transformer.dtype      ,      )
				UpperCAmelCase_	:					str     	=   torch.cat([latents] * 2		) if guidance_scale > 1 else latents

				UpperCAmelCase_	:					List[Any]     	=   torch.tensor(lowerCAmelCase_      ,      device=self.device		).reshape(-1		)
				UpperCAmelCase_	:					Optional[Any]     	=   torch.tensor([1_000] * batch_size      ,      device=self.device		)
				UpperCAmelCase_	:					int     	=   torch.cat([class_labels, class_null]      ,      0		) if guidance_scale > 1 else class_labels

				# set step values
				self.scheduler.set_timesteps(lowerCAmelCase_		)

				for t in self.progress_bar(self.scheduler.timesteps		):
							if guidance_scale > 1:
										UpperCAmelCase_	:					Optional[Any]     	=   latent_model_input[: len(lowerCAmelCase_		) // 2]
										UpperCAmelCase_	:					int     	=   torch.cat([half, half]      ,      dim=0		)
							UpperCAmelCase_	:					Optional[int]     	=   self.scheduler.scale_model_input(lowerCAmelCase_      ,      lowerCAmelCase_		)

							UpperCAmelCase_	:					Union[str, Any]     	=   t
							if not torch.is_tensor(lowerCAmelCase_		):
										# TODO: this requires sync between CPU and GPU. So try to pass timesteps as tensors if you can
										# This would be a good case for the `match` statement (Python 3.10+)
										UpperCAmelCase_	:					List[Any]     	=   latent_model_input.device.type == "mps"
										if isinstance(lowerCAmelCase_      ,      lowerCAmelCase_		):
													UpperCAmelCase_	:					Optional[int]     	=   torch.floataa if is_mps else torch.floataa
										else:
													UpperCAmelCase_	:					int     	=   torch.intaa if is_mps else torch.intaa
										UpperCAmelCase_	:					Optional[int]     	=   torch.tensor([timesteps]      ,      dtype=lowerCAmelCase_      ,      device=latent_model_input.device		)
							elif len(timesteps.shape		) == 0:
										UpperCAmelCase_	:					List[Any]     	=   timesteps[None].to(latent_model_input.device		)
							# broadcast to batch dimension in a way that's compatible with ONNX/Core ML
							UpperCAmelCase_	:					List[Any]     	=   timesteps.expand(latent_model_input.shape[0]		)
							# predict noise model_output
							UpperCAmelCase_	:					List[Any]     	=   self.transformer(
							    lowerCAmelCase_      ,      timestep=lowerCAmelCase_      ,      class_labels=lowerCAmelCase_		).sample

							# perform guidance
							if guidance_scale > 1:
										UpperCAmelCase_	,      UpperCAmelCase_	:					Optional[int]     	=   noise_pred[:, :latent_channels], noise_pred[:, latent_channels:]
										UpperCAmelCase_	,      UpperCAmelCase_	:					Union[str, Any]     	=   torch.split(lowerCAmelCase_      ,      len(lowerCAmelCase_		) // 2      ,      dim=0		)

										UpperCAmelCase_	:					Optional[Any]     	=   uncond_eps + guidance_scale * (cond_eps - uncond_eps)
										UpperCAmelCase_	:					Tuple     	=   torch.cat([half_eps, half_eps]      ,      dim=0		)

										UpperCAmelCase_	:					Optional[Any]     	=   torch.cat([eps, rest]      ,      dim=1		)

							# learned sigma
							if self.transformer.config.out_channels // 2 == latent_channels:
										UpperCAmelCase_	,      UpperCAmelCase_	:					List[str]     	=   torch.split(lowerCAmelCase_      ,      lowerCAmelCase_      ,      dim=1		)
							else:
										UpperCAmelCase_	:					Optional[Any]     	=   noise_pred

							# compute previous image: x_t -> x_t-1
							UpperCAmelCase_	:					Union[str, Any]     	=   self.scheduler.step(lowerCAmelCase_      ,      lowerCAmelCase_      ,      lowerCAmelCase_		).prev_sample

				if guidance_scale > 1:
							UpperCAmelCase_	,      UpperCAmelCase_	:					Optional[Any]     	=   latent_model_input.chunk(2      ,      dim=0		)
				else:
							UpperCAmelCase_	:					Optional[int]     	=   latent_model_input

				UpperCAmelCase_	:					Dict     	=   1 / self.vae.config.scaling_factor * latents
				UpperCAmelCase_	:					Dict     	=   self.vae.decode(lowerCAmelCase_		).sample

				UpperCAmelCase_	:					Tuple     	=   (samples / 2 + 0.5).clamp(0      ,      1		)

				# we always cast to float32 as this does not cause significant overhead and is compatible with bfloat16
				UpperCAmelCase_	:					Optional[Any]     	=   samples.cpu().permute(0      ,      2      ,      3      ,      1		).float().numpy()

				if output_type == "pil":
							UpperCAmelCase_	:					Optional[int]     	=   self.numpy_to_pil(lowerCAmelCase_		)

				if not return_dict:
							return (samples,)

				return ImagePipelineOutput(images=lowerCAmelCase_		)



