





'''simple docstring'''



from argparse import ArgumentParser

from .env import EnvironmentCommand




def    _lowerCAmelCase						(				):
			__lowercase										=							ArgumentParser('''Diffusers CLI tool'''		,       usage='''diffusers-cli <command> [<args>]'''				)
			__lowercase										=							parser.add_subparsers(help='''diffusers-cli command helpers'''				)

			# Register commands
			EnvironmentCommand.register_subcommand(lowerCamelCase_				)

			# Let's go
			__lowercase										=							parser.parse_args()

			if not hasattr(lowerCamelCase_		,       '''func'''				):
						parser.print_help()
						exit(1				)

			# Run
			__lowercase										=							args.func(lowerCamelCase_				)
			service.run()


if __name__ == "__main__":
	main()












'''simple docstring'''



import gc
import unittest

import numpy as np
import torch
from transformers import CLIPTextConfig, CLIPTextModel, CLIPTokenizer

from diffusers import TransformeraDModel, VQDiffusionPipeline, VQDiffusionScheduler, VQModel
from diffusers.pipelines.vq_diffusion.pipeline_vq_diffusion import LearnedClassifierFreeSamplingEmbeddings
from diffusers.utils import load_numpy, slow, torch_device
from diffusers.utils.testing_utils import require_torch_gpu


_SCREAMING_SNAKE_CASE					=    False






class 							__lowercase       (      unittest.TestCase						):







	'''simple docstring'''

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







				'''simple docstring'''





				super().tearDown()
				gc.collect()
				torch.cuda.empty_cache()


	@property
	def 						_UpperCAmelCase     (self )							->							Dict:







				'''simple docstring'''





				return 12


	@property
	def 						_UpperCAmelCase     (self )							->							List[Any]:







				'''simple docstring'''





				return 12


	@property
	def 						_UpperCAmelCase     (self )							->							List[Any]:







				'''simple docstring'''





				return 32


	@property
	def 						_UpperCAmelCase     (self )							->							List[str]:







				'''simple docstring'''





				torch.manual_seed(0 )
				__lowercase										=							VQModel(
				    block_out_channels=[32, 64]     ,in_channels=3     ,out_channels=3     ,down_block_types=['''DownEncoderBlock2D''', '''DownEncoderBlock2D''']     ,up_block_types=['''UpDecoderBlock2D''', '''UpDecoderBlock2D''']     ,latent_channels=3     ,num_vq_embeddings=self.num_embed     ,vq_embed_dim=3     ,)
				return model


	@property
	def 						_UpperCAmelCase     (self )							->							int:







				'''simple docstring'''





				__lowercase										=							CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip''' )
				return tokenizer


	@property
	def 						_UpperCAmelCase     (self )							->							Tuple:







				'''simple docstring'''





				torch.manual_seed(0 )
				__lowercase										=							CLIPTextConfig(
				    bos_token_id=0     ,eos_token_id=2     ,hidden_size=self.text_embedder_hidden_size     ,intermediate_size=37     ,layer_norm_eps=1E-0_5     ,num_attention_heads=4     ,num_hidden_layers=5     ,pad_token_id=1     ,vocab_size=1000     ,)
				return CLIPTextModel(_lowerCamelCase )


	@property
	def 						_UpperCAmelCase     (self )							->							Dict:







				'''simple docstring'''





				torch.manual_seed(0 )

				__lowercase										=							12
				__lowercase										=							12

				__lowercase										=							{
				    '''attention_bias''': True,
				    '''cross_attention_dim''': 32,
				    '''attention_head_dim''': height * width,
				    '''num_attention_heads''': 1,
				    '''num_vector_embeds''': self.num_embed,
				    '''num_embeds_ada_norm''': self.num_embeds_ada_norm,
				    '''norm_num_groups''': 32,
				    '''sample_size''': width,
				    '''activation_fn''': '''geglu-approximate''',
				}

				__lowercase										=							TransformeraDModel(**_lowerCamelCase )
				return model


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







				'''simple docstring'''





				__lowercase										=							'''cpu'''

				__lowercase										=							self.dummy_vqvae
				__lowercase										=							self.dummy_text_encoder
				__lowercase										=							self.dummy_tokenizer
				__lowercase										=							self.dummy_transformer
				__lowercase										=							VQDiffusionScheduler(self.num_embed )
				__lowercase										=							LearnedClassifierFreeSamplingEmbeddings(learnable=_lowerCamelCase )

				__lowercase										=							VQDiffusionPipeline(
				    vqvae=_lowerCamelCase     ,text_encoder=_lowerCamelCase     ,tokenizer=_lowerCamelCase     ,transformer=_lowerCamelCase     ,scheduler=_lowerCamelCase     ,learned_classifier_free_sampling_embeddings=_lowerCamelCase     ,)
				__lowercase										=							pipe.to(_lowerCamelCase )
				pipe.set_progress_bar_config(disable=_lowerCamelCase )

				__lowercase										=							'''teddy bear playing in the pool'''

				__lowercase										=							torch.Generator(device=_lowerCamelCase ).manual_seed(0 )
				__lowercase										=							pipe([prompt]     ,generator=_lowerCamelCase     ,num_inference_steps=2     ,output_type='''np''' )
				__lowercase										=							output.images

				__lowercase										=							torch.Generator(device=_lowerCamelCase ).manual_seed(0 )
				__lowercase										=							pipe(
				    [prompt]     ,generator=_lowerCamelCase     ,output_type='''np'''     ,return_dict=_lowerCamelCase     ,num_inference_steps=2 )[0]

				__lowercase										=							image[0, -3:, -3:, -1]
				__lowercase										=							image_from_tuple[0, -3:, -3:, -1]

				assert image.shape == (1, 24, 24, 3)

				__lowercase										=							np.array([0.6_5_5_1, 0.6_1_6_8, 0.5_0_0_8, 0.5_6_7_6, 0.5_6_5_9, 0.4_2_9_5, 0.6_0_7_3, 0.5_5_9_9, 0.4_9_9_2] )

				assert np.abs(image_slice.flatten() - expected_slice ).max() < 1E-2
				assert np.abs(image_from_tuple_slice.flatten() - expected_slice ).max() < 1E-2


	def 						_UpperCAmelCase     (self )							->							Any:







				'''simple docstring'''





				__lowercase										=							'''cpu'''

				__lowercase										=							self.dummy_vqvae
				__lowercase										=							self.dummy_text_encoder
				__lowercase										=							self.dummy_tokenizer
				__lowercase										=							self.dummy_transformer
				__lowercase										=							VQDiffusionScheduler(self.num_embed )
				__lowercase										=							LearnedClassifierFreeSamplingEmbeddings(
				    learnable=_lowerCamelCase     ,hidden_size=self.text_embedder_hidden_size     ,length=tokenizer.model_max_length )

				__lowercase										=							VQDiffusionPipeline(
				    vqvae=_lowerCamelCase     ,text_encoder=_lowerCamelCase     ,tokenizer=_lowerCamelCase     ,transformer=_lowerCamelCase     ,scheduler=_lowerCamelCase     ,learned_classifier_free_sampling_embeddings=_lowerCamelCase     ,)
				__lowercase										=							pipe.to(_lowerCamelCase )
				pipe.set_progress_bar_config(disable=_lowerCamelCase )

				__lowercase										=							'''teddy bear playing in the pool'''

				__lowercase										=							torch.Generator(device=_lowerCamelCase ).manual_seed(0 )
				__lowercase										=							pipe([prompt]     ,generator=_lowerCamelCase     ,num_inference_steps=2     ,output_type='''np''' )
				__lowercase										=							output.images

				__lowercase										=							torch.Generator(device=_lowerCamelCase ).manual_seed(0 )
				__lowercase										=							pipe(
				    [prompt]     ,generator=_lowerCamelCase     ,output_type='''np'''     ,return_dict=_lowerCamelCase     ,num_inference_steps=2 )[0]

				__lowercase										=							image[0, -3:, -3:, -1]
				__lowercase										=							image_from_tuple[0, -3:, -3:, -1]

				assert image.shape == (1, 24, 24, 3)

				__lowercase										=							np.array([0.6_6_9_3, 0.6_0_7_5, 0.4_9_5_9, 0.5_7_0_1, 0.5_5_8_3, 0.4_3_3_3, 0.6_1_7_1, 0.5_6_8_4, 0.4_9_8_8] )

				assert np.abs(image_slice.flatten() - expected_slice ).max() < 2.0
				assert np.abs(image_from_tuple_slice.flatten() - expected_slice ).max() < 1E-2






@slow
@require_torch_gpu
class 							__lowercase       (      unittest.TestCase						):







	'''simple docstring'''

	def 						_UpperCAmelCase     (self )							->							Dict:







				'''simple docstring'''





				super().tearDown()
				gc.collect()
				torch.cuda.empty_cache()


	def 						_UpperCAmelCase     (self )							->							str:







				'''simple docstring'''





				__lowercase										=							load_numpy(
				    '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'''
				    '''/vq_diffusion/teddy_bear_pool_classifier_free_sampling.npy''' )

				__lowercase										=							VQDiffusionPipeline.from_pretrained('''microsoft/vq-diffusion-ithq''' )
				__lowercase										=							pipeline.to(_lowerCamelCase )
				pipeline.set_progress_bar_config(disable=_lowerCamelCase )

				# requires GPU generator for gumbel softmax
				# don't use GPU generator in tests though
				__lowercase										=							torch.Generator(device=_lowerCamelCase ).manual_seed(0 )
				__lowercase										=							pipeline(
				    '''teddy bear playing in the pool'''     ,num_images_per_prompt=1     ,generator=_lowerCamelCase     ,output_type='''np'''     ,)

				__lowercase										=							output.images[0]

				assert image.shape == (256, 256, 3)
				assert np.abs(expected_image - image ).max() < 2.0





