
import math
import unittest

from transformers import BioGptConfig, is_torch_available
from transformers.testing_utils import require_torch, slow, torch_device

from ...generation.test_utils import GenerationTesterMixin
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 (
		    BioGptForCausalLM,
		    BioGptForSequenceClassification,
		    BioGptForTokenClassification,
		    BioGptModel,
		    BioGptTokenizer,
		)
		from transformers.models.biogpt.modeling_biogpt import BIOGPT_PRETRAINED_MODEL_ARCHIVE_LIST



class    a_ :






			def __init__(      self				,			__UpperCamelCase				,			__UpperCamelCase=13				,			__UpperCamelCase=7				,			__UpperCamelCase=True				,			__UpperCamelCase=True				,			__UpperCamelCase=False				,			__UpperCamelCase=True				,			__UpperCamelCase=99				,			__UpperCamelCase=32				,			__UpperCamelCase=5				,			__UpperCamelCase=4				,			__UpperCamelCase=37				,			__UpperCamelCase="gelu"				,			__UpperCamelCase=0.1				,			__UpperCamelCase=0.1				,			__UpperCamelCase=512				,			__UpperCamelCase=16				,			__UpperCamelCase=2				,			__UpperCamelCase=0.02				,			__UpperCamelCase=3				,			__UpperCamelCase=4				,			__UpperCamelCase=None				,			):
									_lowercase					       =					parent
									_lowercase					       =					batch_size
									_lowercase					       =					seq_length
									_lowercase					       =					is_training
									_lowercase					       =					use_input_mask
									_lowercase					       =					use_token_type_ids
									_lowercase					       =					use_labels
									_lowercase					       =					vocab_size
									_lowercase					       =					hidden_size
									_lowercase					       =					num_hidden_layers
									_lowercase					       =					num_attention_heads
									_lowercase					       =					intermediate_size
									_lowercase					       =					hidden_act
									_lowercase					       =					hidden_dropout_prob
									_lowercase					       =					attention_probs_dropout_prob
									_lowercase					       =					max_position_embeddings
									_lowercase					       =					type_vocab_size
									_lowercase					       =					type_sequence_label_size
									_lowercase					       =					initializer_range
									_lowercase					       =					num_labels
									_lowercase					       =					num_choices
									_lowercase					       =					scope






			def    UpperCamelCase_				(      self  ):
									_lowercase					       =					ids_tensor([self.batch_size, self.seq_length]				,			self.vocab_size  )

									_lowercase					       =					None
									if self.use_input_mask:
															_lowercase					       =					random_attention_mask([self.batch_size, self.seq_length]  )

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

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

									_lowercase					       =					self.get_config()

									return config, input_ids, token_type_ids, input_mask, sequence_labels, token_labels, choice_labels






			def    UpperCamelCase_				(      self  ):
									return BioGptConfig(
									    vocab_size=self.vocab_size				,			hidden_size=self.hidden_size				,			num_hidden_layers=self.num_hidden_layers				,			num_attention_heads=self.num_attention_heads				,			intermediate_size=self.intermediate_size				,			hidden_act=self.hidden_act				,			hidden_dropout_prob=self.hidden_dropout_prob				,			attention_probs_dropout_prob=self.attention_probs_dropout_prob				,			max_position_embeddings=self.max_position_embeddings				,			type_vocab_size=self.type_vocab_size				,			is_decoder=__UpperCamelCase				,			initializer_range=self.initializer_range				,			)






			def    UpperCamelCase_				(      self				,			__UpperCamelCase				,			__UpperCamelCase				,			__UpperCamelCase				,			__UpperCamelCase				,			__UpperCamelCase				,			__UpperCamelCase				,			__UpperCamelCase  ):
									_lowercase					       =					BioGptModel(config=__UpperCamelCase  )
									model.to(__UpperCamelCase  )
									model.eval()
									_lowercase					       =					model(__UpperCamelCase				,			attention_mask=__UpperCamelCase  )
									_lowercase					       =					model(__UpperCamelCase  )
									self.parent.assertEqual(result.last_hidden_state.shape				,			(self.batch_size, self.seq_length, self.hidden_size)  )






			def    UpperCamelCase_				(      self				,			__UpperCamelCase				,			__UpperCamelCase				,			__UpperCamelCase				,			__UpperCamelCase				,			__UpperCamelCase				,			__UpperCamelCase				,			__UpperCamelCase				,			__UpperCamelCase				,			__UpperCamelCase				,			):
									_lowercase					       =					BioGptForCausalLM(config=__UpperCamelCase  )
									model.to(__UpperCamelCase  )
									model.eval()
									_lowercase					       =					model(__UpperCamelCase				,			attention_mask=__UpperCamelCase				,			token_type_ids=__UpperCamelCase				,			labels=__UpperCamelCase  )
									self.parent.assertEqual(result.logits.shape				,			(self.batch_size, self.seq_length, self.vocab_size)  )






			def    UpperCamelCase_				(      self				,			__UpperCamelCase				,			__UpperCamelCase				,			__UpperCamelCase				,			__UpperCamelCase				,			__UpperCamelCase				,			*__UpperCamelCase  ):
									_lowercase					       =					BioGptModel(config=__UpperCamelCase  )
									model.to(__UpperCamelCase  )
									model.eval()

									# create attention mask
									_lowercase					       =					torch.ones(input_ids.shape				,			dtype=torch.long				,			device=__UpperCamelCase  )
									_lowercase					       =					self.seq_length // 2
									_lowercase					       =					0

									# first forward pass
									_lowercase       ,						_lowercase					       =					model(__UpperCamelCase				,			attention_mask=__UpperCamelCase  ).to_tuple()

									# create hypothetical next token and extent to next_input_ids
									_lowercase					       =					ids_tensor((self.batch_size, 1)				,			config.vocab_size  )

									# change a random masked slice from input_ids
									_lowercase					       =					ids_tensor((1,)				,			__UpperCamelCase  ).item() + 1
									_lowercase					       =					ids_tensor((self.batch_size, 1)				,			config.vocab_size  ).squeeze(-1  )
									_lowercase					       =					random_other_next_tokens

									# append to next input_ids and attn_mask
									_lowercase					       =					torch.cat([input_ids, next_tokens]				,			dim=-1  )
									_lowercase					       =					torch.cat(
									    [attn_mask, torch.ones((attn_mask.shape[0], 1)				,			dtype=torch.long				,			device=__UpperCamelCase  )]				,			dim=1				,			)

									# get two different outputs
									_lowercase					       =					model(__UpperCamelCase				,			attention_mask=__UpperCamelCase  )["""last_hidden_state"""]
									_lowercase					       =					model(__UpperCamelCase				,			past_key_values=__UpperCamelCase				,			attention_mask=__UpperCamelCase  )["""last_hidden_state"""]

									# select random slice
									_lowercase					       =					ids_tensor((1,)				,			output_from_past.shape[-1]  ).item()
									_lowercase					       =					output_from_no_past[:, -1, random_slice_idx].detach()
									_lowercase					       =					output_from_past[:, 0, random_slice_idx].detach()

									# test that outputs are equal for slice
									self.parent.assertTrue(torch.allclose(__UpperCamelCase				,			__UpperCamelCase				,			atol=1E-3  )  )






			def    UpperCamelCase_				(      self				,			__UpperCamelCase				,			__UpperCamelCase				,			__UpperCamelCase				,			__UpperCamelCase				,			__UpperCamelCase				,			*__UpperCamelCase  ):
									_lowercase					       =					BioGptModel(config=__UpperCamelCase  ).to(__UpperCamelCase  ).eval()

									_lowercase					       =					torch.ones(input_ids.shape				,			dtype=torch.long				,			device=__UpperCamelCase  )

									# first forward pass
									_lowercase					       =					model(__UpperCamelCase				,			attention_mask=__UpperCamelCase				,			use_cache=__UpperCamelCase  )

									_lowercase       ,						_lowercase					       =					outputs.to_tuple()

									# create hypothetical multiple next token and extent to next_input_ids
									_lowercase					       =					ids_tensor((self.batch_size, 3)				,			config.vocab_size  )
									_lowercase					       =					ids_tensor((self.batch_size, 3)				,			2  )

									# append to next input_ids and
									_lowercase					       =					torch.cat([input_ids, next_tokens]				,			dim=-1  )
									_lowercase					       =					torch.cat([attention_mask, next_attn_mask]				,			dim=-1  )

									_lowercase					       =					model(__UpperCamelCase				,			attention_mask=__UpperCamelCase  )["""last_hidden_state"""]
									_lowercase					       =					model(__UpperCamelCase				,			attention_mask=__UpperCamelCase				,			past_key_values=__UpperCamelCase  )[
									    """last_hidden_state"""
									]

									# select random slice
									_lowercase					       =					ids_tensor((1,)				,			output_from_past.shape[-1]  ).item()
									_lowercase					       =					output_from_no_past[:, -3:, random_slice_idx].detach()
									_lowercase					       =					output_from_past[:, :, random_slice_idx].detach()

									self.parent.assertTrue(output_from_past_slice.shape[1] == next_tokens.shape[1]  )

									# test that outputs are equal for slice
									self.parent.assertTrue(torch.allclose(__UpperCamelCase				,			__UpperCamelCase				,			atol=1E-3  )  )






			def    UpperCamelCase_				(      self				,			__UpperCamelCase				,			__UpperCamelCase				,			__UpperCamelCase				,			__UpperCamelCase				,			__UpperCamelCase				,			*__UpperCamelCase				,			__UpperCamelCase=False  ):
									_lowercase					       =					BioGptForCausalLM(__UpperCamelCase  )
									model.to(__UpperCamelCase  )
									if gradient_checkpointing:
															model.gradient_checkpointing_enable()

									_lowercase					       =					model(__UpperCamelCase				,			labels=__UpperCamelCase  )
									self.parent.assertEqual(result.loss.shape				,			()  )
									self.parent.assertEqual(result.logits.shape				,			(self.batch_size, self.seq_length, self.vocab_size)  )
									result.loss.backward()






			def    UpperCamelCase_				(      self				,			__UpperCamelCase				,			*__UpperCamelCase  ):
									_lowercase					       =					BioGptModel(__UpperCamelCase  )
									_lowercase					       =					model.config.initializer_range / math.sqrt(2 * model.config.num_hidden_layers  )
									for key in model.state_dict().keys():
															if "c_proj" in key and "weight" in key:
																					self.parent.assertLessEqual(abs(torch.std(model.state_dict()[key]  ) - model_std  )				,			0.0_01  )
																					self.parent.assertLessEqual(abs(torch.mean(model.state_dict()[key]  ) - 0.0  )				,			0.01  )






			def    UpperCamelCase_				(      self				,			__UpperCamelCase				,			__UpperCamelCase				,			__UpperCamelCase				,			__UpperCamelCase				,			__UpperCamelCase				,			*__UpperCamelCase  ):
									_lowercase					       =					self.num_labels
									_lowercase					       =					BioGptForTokenClassification(__UpperCamelCase  )
									model.to(__UpperCamelCase  )
									model.eval()
									_lowercase					       =					model(__UpperCamelCase				,			attention_mask=__UpperCamelCase				,			token_type_ids=__UpperCamelCase  )
									self.parent.assertEqual(result.logits.shape				,			(self.batch_size, self.seq_length, self.num_labels)  )






			def    UpperCamelCase_				(      self  ):
									_lowercase					       =					self.prepare_config_and_inputs()
									(
									    (
									    _lowercase
									)       ,						(
									    _lowercase
									)       ,						(
									    _lowercase
									)       ,						(
									    _lowercase
									)       ,						(
									    _lowercase
									)       ,						(
									    _lowercase
									)       ,						(
									    _lowercase
									)       ,						
									)					       =					config_and_inputs
									_lowercase					       =					{"""input_ids""": input_ids, """attention_mask""": input_mask}
									return config, inputs_dict







@require_torch
class    a_ (			_a					,			_a					,			_a					,			unittest.TestCase    ):
			a	:			List[str]					   =							(
			    (BioGptModel, BioGptForCausalLM, BioGptForSequenceClassification, BioGptForTokenClassification)
			    if is_torch_available()
			    else ()
			)
			a	:			int					   =							(BioGptForCausalLM,) if is_torch_available() else ()
			a	:			int					   =							(
			    {
			        '''feature-extraction''': BioGptModel,
			        '''text-classification''': BioGptForSequenceClassification,
			        '''text-generation''': BioGptForCausalLM,
			        '''token-classification''': BioGptForTokenClassification,
			        '''zero-shot''': BioGptForSequenceClassification,
			    }
			    if is_torch_available()
			    else {}
			)
			a	:			Any					   =							False






			def    UpperCamelCase_				(      self  ):
									_lowercase					       =					BioGptModelTester(self  )
									_lowercase					       =					ConfigTester(self				,			config_class=__UpperCamelCase				,			hidden_size=37  )






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






			def    UpperCamelCase_				(      self  ):
									_lowercase					       =					self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_model(*__UpperCamelCase  )






			def    UpperCamelCase_				(      self  ):
									_lowercase					       =					self.model_tester.prepare_config_and_inputs()
									for type in ["absolute", "relative_key", "relative_key_query"]:
															_lowercase					       =					type
															self.model_tester.create_and_check_model(*__UpperCamelCase  )






			def    UpperCamelCase_				(      self  ):
									_lowercase					       =					self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_biogpt_model_attention_mask_past(*__UpperCamelCase  )






			def    UpperCamelCase_				(      self  ):
									_lowercase					       =					self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_forward_and_backwards(*__UpperCamelCase				,			gradient_checkpointing=__UpperCamelCase  )






			def    UpperCamelCase_				(      self  ):
									_lowercase					       =					self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_biogpt_model_past_large_inputs(*__UpperCamelCase  )






			def    UpperCamelCase_				(      self  ):
									_lowercase					       =					self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_biogpt_weight_initialization(*__UpperCamelCase  )






			def    UpperCamelCase_				(      self  ):
									_lowercase					       =					self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_biogpt_for_token_classification(*__UpperCamelCase  )






			@slow
			def    UpperCamelCase_				(      self  ):
									_lowercase					       =					BioGptForCausalLM.from_pretrained("""microsoft/biogpt"""  )
									model.to(__UpperCamelCase  )
									_lowercase					       =					BioGptTokenizer.from_pretrained("""microsoft/biogpt"""  )

									_lowercase					       =					"""left"""

									# Define PAD Token = EOS Token = 50256
									_lowercase					       =					tokenizer.eos_token
									_lowercase					       =					model.config.eos_token_id

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

									_lowercase					       =					tokenizer(__UpperCamelCase				,			return_tensors="""pt"""				,			padding=__UpperCamelCase  )
									_lowercase					       =					inputs["""input_ids"""].to(__UpperCamelCase  )

									_lowercase					       =					model.generate(
									    input_ids=__UpperCamelCase				,			attention_mask=inputs["""attention_mask"""].to(__UpperCamelCase  )				,			)

									_lowercase					       =					tokenizer(sentences[0]				,			return_tensors="""pt"""  ).input_ids.to(__UpperCamelCase  )
									_lowercase					       =					model.generate(input_ids=__UpperCamelCase  )

									_lowercase					       =					inputs_non_padded.shape[-1] - inputs["""attention_mask"""][-1].long().sum().cpu().item()
									_lowercase					       =					tokenizer(sentences[1]				,			return_tensors="""pt"""  ).input_ids.to(__UpperCamelCase  )
									_lowercase					       =					model.generate(input_ids=__UpperCamelCase				,			max_length=model.config.max_length - num_paddings  )

									_lowercase					       =					tokenizer.batch_decode(__UpperCamelCase				,			skip_special_tokens=__UpperCamelCase  )
									_lowercase					       =					tokenizer.decode(output_non_padded[0]				,			skip_special_tokens=__UpperCamelCase  )
									_lowercase					       =					tokenizer.decode(output_padded[0]				,			skip_special_tokens=__UpperCamelCase  )

									_lowercase					       =					[
									    """Hello, my dog is a little bit bigger than a little bit.""",
									    """Today, I have a good idea of how to use the information""",
									]
									self.assertListEqual(__UpperCamelCase				,			__UpperCamelCase  )
									self.assertListEqual(__UpperCamelCase				,			[non_padded_sentence, padded_sentence]  )






			@slow
			def    UpperCamelCase_				(      self  ):
									for model_name in BIOGPT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
															_lowercase					       =					BioGptModel.from_pretrained(__UpperCamelCase  )
															self.assertIsNotNone(__UpperCamelCase  )






			def    UpperCamelCase_				(      self  ):
									_lowercase       ,						_lowercase					       =					self.model_tester.prepare_config_and_inputs_for_common()
									_lowercase					       =					3
									_lowercase					       =					input_dict["""input_ids"""]
									_lowercase					       =					input_ids.ne(1  ).to(__UpperCamelCase  )
									_lowercase					       =					ids_tensor([self.model_tester.batch_size]				,			self.model_tester.type_sequence_label_size  )
									_lowercase					       =					BioGptForSequenceClassification(__UpperCamelCase  )
									model.to(__UpperCamelCase  )
									model.eval()
									_lowercase					       =					model(__UpperCamelCase				,			attention_mask=__UpperCamelCase				,			labels=__UpperCamelCase  )
									self.assertEqual(result.logits.shape				,			(self.model_tester.batch_size, self.model_tester.num_labels)  )






			def    UpperCamelCase_				(      self  ):
									_lowercase       ,						_lowercase					       =					self.model_tester.prepare_config_and_inputs_for_common()
									_lowercase					       =					3
									_lowercase					       =					"""multi_label_classification"""
									_lowercase					       =					input_dict["""input_ids"""]
									_lowercase					       =					input_ids.ne(1  ).to(__UpperCamelCase  )
									_lowercase					       =					ids_tensor(
									    [self.model_tester.batch_size, config.num_labels]				,			self.model_tester.type_sequence_label_size  ).to(torch.float  )
									_lowercase					       =					BioGptForSequenceClassification(__UpperCamelCase  )
									model.to(__UpperCamelCase  )
									model.eval()
									_lowercase					       =					model(__UpperCamelCase				,			attention_mask=__UpperCamelCase				,			labels=__UpperCamelCase  )
									self.assertEqual(result.logits.shape				,			(self.model_tester.batch_size, self.model_tester.num_labels)  )







@require_torch
class    a_ (			unittest.TestCase    ):






			@slow
			def    UpperCamelCase_				(      self  ):
									_lowercase					       =					BioGptForCausalLM.from_pretrained("""microsoft/biogpt"""  )
									_lowercase					       =					torch.tensor([[2, 4_805, 9, 656, 21]]  )
									_lowercase					       =					model(__UpperCamelCase  )[0]

									_lowercase					       =					42_384

									_lowercase					       =					torch.Size((1, 5, vocab_size)  )
									self.assertEqual(output.shape				,			__UpperCamelCase  )

									_lowercase					       =					torch.tensor(
									    [[[-9.52_36, -9.89_18, 10.45_57], [-11.04_69, -9.64_23, 8.10_22], [-8.86_64, -7.88_26, 5.53_25]]]  )

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






			@slow
			def    UpperCamelCase_				(      self  ):
									_lowercase					       =					BioGptTokenizer.from_pretrained("""microsoft/biogpt"""  )
									_lowercase					       =					BioGptForCausalLM.from_pretrained("""microsoft/biogpt"""  )
									model.to(__UpperCamelCase  )

									torch.manual_seed(0  )
									_lowercase					       =					tokenizer("""COVID-19 is"""				,			return_tensors="""pt"""  ).to(__UpperCamelCase  )
									_lowercase					       =					model.generate(
									    **__UpperCamelCase				,			min_length=100				,			max_length=1_024				,			num_beams=5				,			early_stopping=__UpperCamelCase				,			)
									_lowercase					       =					tokenizer.decode(output_ids[0]				,			skip_special_tokens=__UpperCamelCase  )

									_lowercase					       =					(
									    """COVID-19 is a global pandemic caused by severe acute respiratory syndrome coronavirus 2 (SARS-CoV-2), the"""
									    """ causative agent of coronavirus disease 2019 (COVID-19), which has spread to more than 200 countries and"""
									    """ territories, including the United States (US), Canada, Australia, New Zealand, the United Kingdom (UK),"""
									    """ and the United States of America (USA), as of March 11, 2020, with more than 800,000 confirmed cases and"""
									    """ more than 800,000 deaths."""
									)
									self.assertEqual(__UpperCamelCase				,			__UpperCamelCase  )

from typing import List, Optional, Union

import torch

from ...models import UNetaDConditionModel, VQModel
from ...pipelines import DiffusionPipeline
from ...pipelines.pipeline_utils import ImagePipelineOutput
from ...schedulers import DDPMScheduler
from ...utils import (
    is_accelerate_available,
    is_accelerate_version,
    logging,
    randn_tensor,
    replace_example_docstring,
)


A		: Tuple								=     logging.get_logger(__name__)  # pylint: disable=invalid-name

A		: Optional[Any]								=     '''
    Examples:
        ```py
        >>> import torch
        >>> import numpy as np

        >>> from diffusers import KandinskyV22PriorPipeline, KandinskyV22ControlnetPipeline
        >>> from transformers import pipeline
        >>> from diffusers.utils import load_image


        >>> def make_hint(image, depth_estimator):
        ...     image = depth_estimator(image)["depth"]
        ...     image = np.array(image)
        ...     image = image[:, :, None]
        ...     image = np.concatenate([image, image, image], axis=2)
        ...     detected_map = torch.from_numpy(image).float() / 255.0
        ...     hint = detected_map.permute(2, 0, 1)
        ...     return hint


        >>> depth_estimator = pipeline("depth-estimation")

        >>> pipe_prior = KandinskyV22PriorPipeline.from_pretrained(
        ...     "kandinsky-community/kandinsky-2-2-prior", torch_dtype=torch.float16
        ... )
        >>> pipe_prior = pipe_prior.to("cuda")

        >>> pipe = KandinskyV22ControlnetPipeline.from_pretrained(
        ...     "kandinsky-community/kandinsky-2-2-controlnet-depth", torch_dtype=torch.float16
        ... )
        >>> pipe = pipe.to("cuda")


        >>> img = load_image(
        ...     "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"
        ...     "/kandinsky/cat.png"
        ... ).resize((768, 768))

        >>> hint = make_hint(img, depth_estimator).unsqueeze(0).half().to("cuda")

        >>> prompt = "A robot, 4k photo"
        >>> negative_prior_prompt = "lowres, text, error, cropped, worst quality, low quality, jpeg artifacts, ugly, duplicate, morbid, mutilated, out of frame, extra fingers, mutated hands, poorly drawn hands, poorly drawn face, mutation, deformed, blurry, dehydrated, bad anatomy, bad proportions, extra limbs, cloned face, disfigured, gross proportions, malformed limbs, missing arms, missing legs, extra arms, extra legs, fused fingers, too many fingers, long neck, username, watermark, signature"

        >>> generator = torch.Generator(device="cuda").manual_seed(43)

        >>> image_emb, zero_image_emb = pipe_prior(
        ...     prompt=prompt, negative_prompt=negative_prior_prompt, generator=generator
        ... ).to_tuple()

        >>> images = pipe(
        ...     image_embeds=image_emb,
        ...     negative_image_embeds=zero_image_emb,
        ...     hint=hint,
        ...     num_inference_steps=50,
        ...     generator=generator,
        ...     height=768,
        ...     width=768,
        ... ).images

        >>> images[0].save("robot_cat.png")
        ```
'''




def 	UpperCamelCase__  (     SCREAMING_SNAKE_CASE_	:	str		,   SCREAMING_SNAKE_CASE_	:	Dict		,   SCREAMING_SNAKE_CASE_	:	int=8       )  ->     Optional[int]:
						_lowercase					       =					height // scale_factor**2
						if height % scale_factor**2 != 0:
												new_height += 1
						_lowercase					       =					width // scale_factor**2
						if width % scale_factor**2 != 0:
												new_width += 1
						return new_height * scale_factor, new_width * scale_factor



class    a_ (			_a    ):






			def __init__(      self				,			__UpperCamelCase				,			__UpperCamelCase				,			__UpperCamelCase				,			):
									super().__init__()

									self.register_modules(
									    unet=__UpperCamelCase				,			scheduler=__UpperCamelCase				,			movq=__UpperCamelCase				,			)
									_lowercase					       =					2 ** (len(self.movq.config.block_out_channels  ) - 1)






			def    UpperCamelCase_				(      self				,			__UpperCamelCase				,			__UpperCamelCase				,			__UpperCamelCase				,			__UpperCamelCase				,			__UpperCamelCase				,			__UpperCamelCase  ):
									if latents is None:
															_lowercase					       =					randn_tensor(__UpperCamelCase				,			generator=__UpperCamelCase				,			device=__UpperCamelCase				,			dtype=__UpperCamelCase  )
									else:
															if latents.shape != shape:
																					raise ValueError(f"""Unexpected latents shape, got {latents.shape}, expected {shape}"""  )
															_lowercase					       =					latents.to(__UpperCamelCase  )

									_lowercase					       =					latents * scheduler.init_noise_sigma
									return latents






			def    UpperCamelCase_				(      self				,			__UpperCamelCase=0  ):
									if is_accelerate_available():
															from accelerate import cpu_offload
									else:
															raise ImportError("""Please install accelerate via `pip install accelerate`"""  )

									_lowercase					       =					torch.device(f"""cuda:{gpu_id}"""  )

									_lowercase					       =					[
									    self.unet,
									    self.movq,
									]
									for cpu_offloaded_model in models:
															if cpu_offloaded_model is not None:
																					cpu_offload(__UpperCamelCase				,			__UpperCamelCase  )






			def    UpperCamelCase_				(      self				,			__UpperCamelCase=0  ):
									if is_accelerate_available() and is_accelerate_version(""">="""				,			"""0.17.0.dev0"""  ):
															from accelerate import cpu_offload_with_hook
									else:
															raise ImportError("""`enable_model_cpu_offload` requires `accelerate v0.17.0` or higher."""  )

									_lowercase					       =					torch.device(f"""cuda:{gpu_id}"""  )

									if self.device.type != "cpu":
															self.to("""cpu"""				,			silence_dtype_warnings=__UpperCamelCase  )
															torch.cuda.empty_cache()  # otherwise we don't see the memory savings (but they probably exist)

									_lowercase					       =					None
									for cpu_offloaded_model in [self.unet, self.movq]:
															_lowercase       ,						_lowercase					       =					cpu_offload_with_hook(__UpperCamelCase				,			__UpperCamelCase				,			prev_module_hook=__UpperCamelCase  )

									# We'll offload the last model manually.
									_lowercase					       =					hook






			@property
			# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline._execution_device
			def    UpperCamelCase_				(      self  ):
									if not hasattr(self.unet				,			"""_hf_hook"""  ):
															return self.device
									for module in self.unet.modules():
															if (
															    hasattr(__UpperCamelCase				,			"""_hf_hook"""  )
															    and hasattr(module._hf_hook				,			"""execution_device"""  )
															    and module._hf_hook.execution_device is not None
															):
																					return torch.device(module._hf_hook.execution_device  )
									return self.device






			@torch.no_grad()
			@replace_example_docstring(__UpperCamelCase  )
			def __call__(      self				,			__UpperCamelCase				,			__UpperCamelCase				,			__UpperCamelCase				,			__UpperCamelCase = 512				,			__UpperCamelCase = 512				,			__UpperCamelCase = 100				,			__UpperCamelCase = 4.0				,			__UpperCamelCase = 1				,			__UpperCamelCase = None				,			__UpperCamelCase = None				,			__UpperCamelCase = "pil"				,			__UpperCamelCase = True				,			):
									_lowercase					       =					self._execution_device

									_lowercase					       =					guidance_scale > 1.0

									if isinstance(__UpperCamelCase				,			__UpperCamelCase  ):
															_lowercase					       =					torch.cat(__UpperCamelCase				,			dim=0  )
									if isinstance(__UpperCamelCase				,			__UpperCamelCase  ):
															_lowercase					       =					torch.cat(__UpperCamelCase				,			dim=0  )
									if isinstance(__UpperCamelCase				,			__UpperCamelCase  ):
															_lowercase					       =					torch.cat(__UpperCamelCase				,			dim=0  )

									_lowercase					       =					image_embeds.shape[0] * num_images_per_prompt

									if do_classifier_free_guidance:
															_lowercase					       =					image_embeds.repeat_interleave(__UpperCamelCase				,			dim=0  )
															_lowercase					       =					negative_image_embeds.repeat_interleave(__UpperCamelCase				,			dim=0  )
															_lowercase					       =					hint.repeat_interleave(__UpperCamelCase				,			dim=0  )

									_lowercase					       =					torch.cat([negative_image_embeds, image_embeds]				,			dim=0  ).to(dtype=self.unet.dtype				,			device=__UpperCamelCase  )
									_lowercase					       =					torch.cat([hint, hint]				,			dim=0  ).to(dtype=self.unet.dtype				,			device=__UpperCamelCase  )
									self.scheduler.set_timesteps(__UpperCamelCase				,			device=__UpperCamelCase  )
									_lowercase					       =					self.scheduler.timesteps

									_lowercase					       =					self.movq.config.latent_channels

									_lowercase       ,						_lowercase					       =					downscale_height_and_width(__UpperCamelCase				,			__UpperCamelCase				,			self.movq_scale_factor  )

									# create initial latent
									_lowercase					       =					self.prepare_latents(
									    (batch_size, num_channels_latents, height, width)				,			image_embeds.dtype				,			__UpperCamelCase				,			__UpperCamelCase				,			__UpperCamelCase				,			self.scheduler				,			)

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

															_lowercase					       =					{"""image_embeds""": image_embeds, """hint""": hint}
															_lowercase					       =					self.unet(
															    sample=__UpperCamelCase				,			timestep=__UpperCamelCase				,			encoder_hidden_states=__UpperCamelCase				,			added_cond_kwargs=__UpperCamelCase				,			return_dict=__UpperCamelCase				,			)[0]

															if do_classifier_free_guidance:
																					_lowercase       ,						_lowercase					       =					noise_pred.split(latents.shape[1]				,			dim=1  )
																					_lowercase       ,						_lowercase					       =					noise_pred.chunk(2  )
																					_lowercase       ,						_lowercase					       =					variance_pred.chunk(2  )
																					_lowercase					       =					noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)
																					_lowercase					       =					torch.cat([noise_pred, variance_pred_text]				,			dim=1  )

															if not (
															    hasattr(self.scheduler.config				,			"""variance_type"""  )
															    and self.scheduler.config.variance_type in ["learned", "learned_range"]
															):
																					_lowercase       ,						_lowercase					       =					noise_pred.split(latents.shape[1]				,			dim=1  )

															# compute the previous noisy sample x_t -> x_t-1
															_lowercase					       =					self.scheduler.step(
															    __UpperCamelCase				,			__UpperCamelCase				,			__UpperCamelCase				,			generator=__UpperCamelCase				,			)[0]
									# post-processing
									_lowercase					       =					self.movq.decode(__UpperCamelCase				,			force_not_quantize=__UpperCamelCase  )["""sample"""]

									if output_type not in ["pt", "np", "pil"]:
															raise ValueError(f"""Only the output types `pt`, `pil` and `np` are supported not output_type={output_type}"""  )

									if output_type in ["np", "pil"]:
															_lowercase					       =					image * 0.5 + 0.5
															_lowercase					       =					image.clamp(0				,			1  )
															_lowercase					       =					image.cpu().permute(0				,			2				,			3				,			1  ).float().numpy()

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

									if not return_dict:
															return (image,)

									return ImagePipelineOutput(images=__UpperCamelCase  )