


"""simple docstring"""






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,
)


UpperCamelCase__				:str			   =	logging.get_logger(__name__)  # pylint: disable=invalid-name

UpperCamelCase__				:List[Any]			   =	"""
    Examples:
        ```py
        >>> from diffusers import KandinskyV22Pipeline, KandinskyV22PriorPipeline
        >>> import torch

        >>> pipe_prior = KandinskyV22PriorPipeline.from_pretrained(\"kandinsky-community/kandinsky-2-2-prior\")
        >>> pipe_prior.to(\"cuda\")
        >>> prompt = \"red cat, 4k photo\"
        >>> out = pipe_prior(prompt)
        >>> image_emb = out.image_embeds
        >>> zero_image_emb = out.negative_image_embeds
        >>> pipe = KandinskyV22Pipeline.from_pretrained(\"kandinsky-community/kandinsky-2-2-decoder\")
        >>> pipe.to(\"cuda\")
        >>> image = pipe(
        ...     image_embeds=image_emb,
        ...     negative_image_embeds=zero_image_emb,
        ...     height=768,
        ...     width=768,
        ...     num_inference_steps=50,
        ... ).images
        >>> image[0].save(\"cat.png\")
        ```
"""





def 				A_       ( snake_case__       ,					snake_case__       ,					snake_case__=8		)					->  Any:
			_UpperCamelCase					:int					     =	height // scale_factor**2
			if height % scale_factor**2 != 0:
						new_height += 1
			_UpperCamelCase					:int					     =	width // scale_factor**2
			if width % scale_factor**2 != 0:
						new_width += 1
			return new_height * scale_factor, new_width * scale_factor






class      A(  lowerCamelCase__		):

		"""simple docstring"""





		def __init__( self					,				SCREAMING_SNAKE_CASE__					,				SCREAMING_SNAKE_CASE__					,				SCREAMING_SNAKE_CASE__					,				)      ->	Tuple:

					"""simple docstring"""



					super().__init__()

					self.register_modules(
					    unet=SCREAMING_SNAKE_CASE__					,				scheduler=SCREAMING_SNAKE_CASE__					,				movq=SCREAMING_SNAKE_CASE__					,				)
					_UpperCamelCase					:Dict					     =	2 ** (len(self.movq.config.block_out_channels   ) - 1)
		def   _UpperCamelCase( self					,				SCREAMING_SNAKE_CASE__					,				SCREAMING_SNAKE_CASE__					,				SCREAMING_SNAKE_CASE__					,				SCREAMING_SNAKE_CASE__					,				SCREAMING_SNAKE_CASE__					,				SCREAMING_SNAKE_CASE__   )      ->	int:

					"""simple docstring"""



					if latents is None:
								_UpperCamelCase					:str					     =	randn_tensor(SCREAMING_SNAKE_CASE__					,				generator=SCREAMING_SNAKE_CASE__					,				device=SCREAMING_SNAKE_CASE__					,				dtype=SCREAMING_SNAKE_CASE__   )
					else:
								if latents.shape != shape:
											raise ValueError(f"Unexpected latents shape, got {latents.shape}, expected {shape}"   )
								_UpperCamelCase					:int					     =	latents.to(SCREAMING_SNAKE_CASE__   )

					_UpperCamelCase					:str					     =	latents * scheduler.init_noise_sigma
					return latents
		def   _UpperCamelCase( self					,				SCREAMING_SNAKE_CASE__=0   )      ->	Tuple:

					"""simple docstring"""



					if is_accelerate_available():
								from accelerate import cpu_offload
					else:
								raise ImportError('''Please install accelerate via `pip install accelerate`'''   )

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

					_UpperCamelCase					:Tuple					     =	[
					    self.unet,
					    self.movq,
					]
					for cpu_offloaded_model in models:
								if cpu_offloaded_model is not None:
											cpu_offload(SCREAMING_SNAKE_CASE__					,				SCREAMING_SNAKE_CASE__   )
		def   _UpperCamelCase( self					,				SCREAMING_SNAKE_CASE__=0   )      ->	Optional[Any]:

					"""simple docstring"""



					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.'''   )

					_UpperCamelCase					:Optional[int]					     =	torch.device(f"cuda:{gpu_id}"   )

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

					_UpperCamelCase					:Optional[Any]					     =	None
					for cpu_offloaded_model in [self.unet, self.movq]:
								_UpperCamelCase     ,      _UpperCamelCase					:str					     =	cpu_offload_with_hook(SCREAMING_SNAKE_CASE__					,				SCREAMING_SNAKE_CASE__					,				prev_module_hook=SCREAMING_SNAKE_CASE__   )

					# We'll offload the last model manually.
					_UpperCamelCase					:Tuple					     =	hook
		@property
		# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline._execution_device
		def   _UpperCamelCase( self   )      ->	Optional[int]:

					"""simple docstring"""



					if not hasattr(self.unet					,				'''_hf_hook'''   ):
								return self.device
					for module in self.unet.modules():
								if (
								    hasattr(SCREAMING_SNAKE_CASE__					,				'''_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(SCREAMING_SNAKE_CASE__   )
		def __call__( self					,				SCREAMING_SNAKE_CASE__					,				SCREAMING_SNAKE_CASE__					,				SCREAMING_SNAKE_CASE__ = 5_12					,				SCREAMING_SNAKE_CASE__ = 5_12					,				SCREAMING_SNAKE_CASE__ = 1_00					,				SCREAMING_SNAKE_CASE__ = 4.0					,				SCREAMING_SNAKE_CASE__ = 1					,				SCREAMING_SNAKE_CASE__ = None					,				SCREAMING_SNAKE_CASE__ = None					,				SCREAMING_SNAKE_CASE__ = "pil"					,				SCREAMING_SNAKE_CASE__ = True					,				)      ->	List[str]:

					"""simple docstring"""



					_UpperCamelCase					:List[Any]					     =	self._execution_device

					_UpperCamelCase					:Any					     =	guidance_scale > 1.0

					if isinstance(SCREAMING_SNAKE_CASE__					,				SCREAMING_SNAKE_CASE__   ):
								_UpperCamelCase					:Optional[Any]					     =	torch.cat(SCREAMING_SNAKE_CASE__					,				dim=0   )
					_UpperCamelCase					:List[Any]					     =	image_embeds.shape[0] * num_images_per_prompt
					if isinstance(SCREAMING_SNAKE_CASE__					,				SCREAMING_SNAKE_CASE__   ):
								_UpperCamelCase					:Tuple					     =	torch.cat(SCREAMING_SNAKE_CASE__					,				dim=0   )

					if do_classifier_free_guidance:
								_UpperCamelCase					:Union[str, Any]					     =	image_embeds.repeat_interleave(SCREAMING_SNAKE_CASE__					,				dim=0   )
								_UpperCamelCase					:Optional[Any]					     =	negative_image_embeds.repeat_interleave(SCREAMING_SNAKE_CASE__					,				dim=0   )

					_UpperCamelCase					:Optional[Any]					     =	torch.cat([negative_image_embeds, image_embeds]					,				dim=0   ).to(dtype=self.unet.dtype					,				device=SCREAMING_SNAKE_CASE__   )

					self.scheduler.set_timesteps(SCREAMING_SNAKE_CASE__					,				device=SCREAMING_SNAKE_CASE__   )
					_UpperCamelCase					:List[str]					     =	self.scheduler.timesteps

					_UpperCamelCase					:Optional[int]					     =	self.unet.config.in_channels

					_UpperCamelCase     ,      _UpperCamelCase					:int					     =	downscale_height_and_width(SCREAMING_SNAKE_CASE__					,				SCREAMING_SNAKE_CASE__					,				self.movq_scale_factor   )

					# create initial latent
					_UpperCamelCase					:Tuple					     =	self.prepare_latents(
					    (batch_size, num_channels_latents, height, width)					,				image_embeds.dtype					,				SCREAMING_SNAKE_CASE__					,				SCREAMING_SNAKE_CASE__					,				SCREAMING_SNAKE_CASE__					,				self.scheduler					,				)

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

								_UpperCamelCase					:List[Any]					     =	{'''image_embeds''': image_embeds}
								_UpperCamelCase					:List[Any]					     =	self.unet(
								    sample=SCREAMING_SNAKE_CASE__					,				timestep=SCREAMING_SNAKE_CASE__					,				encoder_hidden_states=SCREAMING_SNAKE_CASE__					,				added_cond_kwargs=SCREAMING_SNAKE_CASE__					,				return_dict=SCREAMING_SNAKE_CASE__					,				)[0]

								if do_classifier_free_guidance:
											_UpperCamelCase     ,      _UpperCamelCase					:Optional[Any]					     =	noise_pred.split(latents.shape[1]					,				dim=1   )
											_UpperCamelCase     ,      _UpperCamelCase					:List[str]					     =	noise_pred.chunk(2   )
											_UpperCamelCase     ,      _UpperCamelCase					:Union[str, Any]					     =	variance_pred.chunk(2   )
											_UpperCamelCase					:str					     =	noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)
											_UpperCamelCase					:Any					     =	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"]
								):
											_UpperCamelCase     ,      _UpperCamelCase					:int					     =	noise_pred.split(latents.shape[1]					,				dim=1   )

								# compute the previous noisy sample x_t -> x_t-1
								_UpperCamelCase					:List[Any]					     =	self.scheduler.step(
								    SCREAMING_SNAKE_CASE__					,				SCREAMING_SNAKE_CASE__					,				SCREAMING_SNAKE_CASE__					,				generator=SCREAMING_SNAKE_CASE__					,				)[0]
					# post-processing
					_UpperCamelCase					:str					     =	self.movq.decode(SCREAMING_SNAKE_CASE__					,				force_not_quantize=SCREAMING_SNAKE_CASE__   )['''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"]:
								_UpperCamelCase					:str					     =	image * 0.5 + 0.5
								_UpperCamelCase					:Dict					     =	image.clamp(0					,				1   )
								_UpperCamelCase					:Tuple					     =	image.cpu().permute(0					,				2					,				3					,				1   ).float().numpy()

					if output_type == "pil":
								_UpperCamelCase					:Optional[int]					     =	self.numpy_to_pil(SCREAMING_SNAKE_CASE__   )

					if not return_dict:
								return (image,)

					return ImagePipelineOutput(images=SCREAMING_SNAKE_CASE__   )





"""simple docstring"""






import os
import unittest

from transformers.models.cpmant.tokenization_cpmant import VOCAB_FILES_NAMES, CpmAntTokenizer
from transformers.testing_utils import require_jieba, tooslow

from ...test_tokenization_common import TokenizerTesterMixin






@require_jieba
class      A(  lowerCamelCase__				,							unittest.TestCase		):

		"""simple docstring"""





		A    =							CpmAntTokenizer
		A    =							False
		def   _UpperCamelCase( self   )      ->	Union[str, Any]:

					"""simple docstring"""



					super().setUp()

					_UpperCamelCase					:str					     =	[
					    '''<d>''',
					    '''</d>''',
					    '''<s>''',
					    '''</s>''',
					    '''</_>''',
					    '''<unk>''',
					    '''<pad>''',
					    '''</n>''',
					    '''我''',
					    '''是''',
					    '''C''',
					    '''P''',
					    '''M''',
					    '''A''',
					    '''n''',
					    '''t''',
					]
					_UpperCamelCase					:List[str]					     =	os.path.join(self.tmpdirname					,				VOCAB_FILES_NAMES['''vocab_file''']   )
					with open(self.vocab_file					,				'''w'''					,				encoding='''utf-8'''   ) as vocab_writer:
								vocab_writer.write(''''''.join([x + '''\n''' for x in vocab_tokens]   )   )





		@tooslow
		def   _UpperCamelCase( self   )      ->	Tuple:

					"""simple docstring"""



					_UpperCamelCase					:Any					     =	CpmAntTokenizer.from_pretrained('''openbmb/cpm-ant-10b'''   )
					_UpperCamelCase					:List[str]					     =	'''今天天气真好！'''
					_UpperCamelCase					:Tuple					     =	['''今天''', '''天气''', '''真''', '''好''', '''！''']
					_UpperCamelCase					:Tuple					     =	tokenizer.tokenize(SCREAMING_SNAKE_CASE__   )
					self.assertListEqual(SCREAMING_SNAKE_CASE__					,				SCREAMING_SNAKE_CASE__   )
					_UpperCamelCase					:Dict					     =	'''今天天气真好！'''
					_UpperCamelCase					:Optional[int]					     =	[tokenizer.bos_token] + tokens

					_UpperCamelCase					:List[str]					     =	[6, 98_02, 1_49_62, 20_82, 8_31, 2_44]
					self.assertListEqual(tokenizer.convert_tokens_to_ids(SCREAMING_SNAKE_CASE__   )					,				SCREAMING_SNAKE_CASE__   )

					_UpperCamelCase					:str					     =	tokenizer.decode(SCREAMING_SNAKE_CASE__   )
					self.assertEqual(SCREAMING_SNAKE_CASE__					,				SCREAMING_SNAKE_CASE__   )

