

'''simple docstring'''
import gc
import unittest

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

from diffusers import AutoencoderKL, DDIMScheduler, LDMTextToImagePipeline, UNetaDConditionModel
from diffusers.utils.testing_utils import (
    enable_full_determinism,
    load_numpy,
    nightly,
    require_torch_gpu,
    slow,
    torch_device,
)

from ..pipeline_params import TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_PARAMS
from ..test_pipelines_common import PipelineTesterMixin


enable_full_determinism()







class 							A__    (       A__     ,   unittest.TestCase  ):
  A__		 =				LDMTextToImagePipeline
  A__		 =				TEXT_TO_IMAGE_PARAMS - {
      'negative_prompt',
      'negative_prompt_embeds',
      'cross_attention_kwargs',
      'prompt_embeds',
  }
  A__		 =				PipelineTesterMixin.required_optional_params - {
      'num_images_per_prompt',
      'callback',
      'callback_steps',
  }
  A__		 =				TEXT_TO_IMAGE_BATCH_PARAMS
  A__		 =				False





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





       '''simple docstring'''






       torch.manual_seed(0 )
       _SCREAMING_SNAKE_CASE							    =UNetaDConditionModel(
           block_out_channels=(32, 64)							,					layers_per_block=2							,					sample_size=32							,					in_channels=4							,					out_channels=4							,					down_block_types=('DownBlock2D', 'CrossAttnDownBlock2D')							,					up_block_types=('CrossAttnUpBlock2D', 'UpBlock2D')							,					cross_attention_dim=32							,					)
       _SCREAMING_SNAKE_CASE							    =DDIMScheduler(
           beta_start=0.0_00_85							,					beta_end=0.0_12							,					beta_schedule='scaled_linear'							,					clip_sample=_a							,					set_alpha_to_one=_a							,					)
       torch.manual_seed(0 )
       _SCREAMING_SNAKE_CASE							    =AutoencoderKL(
           block_out_channels=(32, 64)							,					in_channels=3							,					out_channels=3							,					down_block_types=('DownEncoderBlock2D', 'DownEncoderBlock2D')							,					up_block_types=('UpDecoderBlock2D', 'UpDecoderBlock2D')							,					latent_channels=4							,					)
       torch.manual_seed(0 )
       _SCREAMING_SNAKE_CASE							    =CLIPTextConfig(
           bos_token_id=0							,					eos_token_id=2							,					hidden_size=32							,					intermediate_size=37							,					layer_norm_eps=1e-05							,					num_attention_heads=4							,					num_hidden_layers=5							,					pad_token_id=1							,					vocab_size=1000							,					)
       _SCREAMING_SNAKE_CASE							    =CLIPTextModel(_a )
       _SCREAMING_SNAKE_CASE							    =CLIPTokenizer.from_pretrained('hf-internal-testing/tiny-random-clip' )

       _SCREAMING_SNAKE_CASE							    ={
           'unet': unet,
           'scheduler': scheduler,
           'vqvae': vae,
           'bert': text_encoder,
           'tokenizer': tokenizer,
       }
       return components





  def 						A  (		self      :	Union[str, Any]							,					_a      :	List[str]							,					_a      :	Dict=0 )	->    Optional[int]:





       '''simple docstring'''






       if str(_a ).startswith('mps' ):
            _SCREAMING_SNAKE_CASE							    =torch.manual_seed(_a )
       else:
            _SCREAMING_SNAKE_CASE							    =torch.Generator(device=_a ).manual_seed(_a )
       _SCREAMING_SNAKE_CASE							    ={
           'prompt': 'A painting of a squirrel eating a burger',
           'generator': generator,
           'num_inference_steps': 2,
           'guidance_scale': 6.0,
           'output_type': 'numpy',
       }
       return inputs





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





       '''simple docstring'''






       _SCREAMING_SNAKE_CASE							    ='cpu'  # ensure determinism for the device-dependent torch.Generator

       _SCREAMING_SNAKE_CASE							    =self.get_dummy_components()
       _SCREAMING_SNAKE_CASE							    =LDMTextToImagePipeline(**_a )
       pipe.to(_a )
       pipe.set_progress_bar_config(disable=_a )

       _SCREAMING_SNAKE_CASE							    =self.get_dummy_inputs(_a )
       _SCREAMING_SNAKE_CASE							    =pipe(**_a ).images
       _SCREAMING_SNAKE_CASE							    =image[0, -3:, -3:, -1]

       assert image.shape == (1, 16, 16, 3)
       _SCREAMING_SNAKE_CASE							    =np.array([0.61_01, 0.61_56, 0.56_22, 0.48_95, 0.66_61, 0.38_04, 0.57_48, 0.61_36, 0.50_14] )

       assert np.abs(image_slice.flatten() - expected_slice ).max() < 1e-3







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





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





       '''simple docstring'''






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





  def 						A  (		self      :	int							,					_a      :	Optional[Any]							,					_a      :	Any=torch.floataa							,					_a      :	int=0 )	->    Optional[Any]:





       '''simple docstring'''






       _SCREAMING_SNAKE_CASE							    =torch.manual_seed(_a )
       _SCREAMING_SNAKE_CASE							    =np.random.RandomState(_a ).standard_normal((1, 4, 32, 32) )
       _SCREAMING_SNAKE_CASE							    =torch.from_numpy(_a ).to(device=_a							,					dtype=_a )
       _SCREAMING_SNAKE_CASE							    ={
           'prompt': 'A painting of a squirrel eating a burger',
           'latents': latents,
           'generator': generator,
           'num_inference_steps': 3,
           'guidance_scale': 6.0,
           'output_type': 'numpy',
       }
       return inputs





  def 						A  (		self      :	Dict )	->    Dict:





       '''simple docstring'''






       _SCREAMING_SNAKE_CASE							    =LDMTextToImagePipeline.from_pretrained('CompVis/ldm-text2im-large-256' ).to(_a )
       pipe.set_progress_bar_config(disable=_a )

       _SCREAMING_SNAKE_CASE							    =self.get_inputs(_a )
       _SCREAMING_SNAKE_CASE							    =pipe(**_a ).images
       _SCREAMING_SNAKE_CASE							    =image[0, -3:, -3:, -1].flatten()

       assert image.shape == (1, 256, 256, 3)
       _SCREAMING_SNAKE_CASE							    =np.array([0.5_18_25, 0.5_28_50, 0.5_25_43, 0.5_42_58, 0.5_23_04, 0.5_25_69, 0.5_43_63, 0.5_52_76, 0.5_68_78] )
       _SCREAMING_SNAKE_CASE							    =np.abs(expected_slice - image_slice ).max()
       assert max_diff < 1e-3







@nightly
@require_torch_gpu
class 							A__    (       unittest.TestCase  ):





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





       '''simple docstring'''






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





  def 						A  (		self      :	Any							,					_a      :	Any							,					_a      :	Optional[int]=torch.floataa							,					_a      :	str=0 )	->    List[Any]:





       '''simple docstring'''






       _SCREAMING_SNAKE_CASE							    =torch.manual_seed(_a )
       _SCREAMING_SNAKE_CASE							    =np.random.RandomState(_a ).standard_normal((1, 4, 32, 32) )
       _SCREAMING_SNAKE_CASE							    =torch.from_numpy(_a ).to(device=_a							,					dtype=_a )
       _SCREAMING_SNAKE_CASE							    ={
           'prompt': 'A painting of a squirrel eating a burger',
           'latents': latents,
           'generator': generator,
           'num_inference_steps': 50,
           'guidance_scale': 6.0,
           'output_type': 'numpy',
       }
       return inputs





  def 						A  (		self      :	str )	->    Optional[int]:





       '''simple docstring'''






       _SCREAMING_SNAKE_CASE							    =LDMTextToImagePipeline.from_pretrained('CompVis/ldm-text2im-large-256' ).to(_a )
       pipe.set_progress_bar_config(disable=_a )

       _SCREAMING_SNAKE_CASE							    =self.get_inputs(_a )
       _SCREAMING_SNAKE_CASE							    =pipe(**_a ).images[0]

       _SCREAMING_SNAKE_CASE							    =load_numpy(
           'https://huggingface.co/datasets/diffusers/test-arrays/resolve/main/ldm_text2img/ldm_large_256_ddim.npy' )
       _SCREAMING_SNAKE_CASE							    =np.abs(expected_image - image ).max()
       assert max_diff < 1e-3



'''simple docstring'''
import importlib
import sys
from argparse import REMAINDER, ArgumentParser
from pathlib import Path

import torch_xla.distributed.xla_multiprocessing as xmp


def        _lowerCAmelCase					(						)      ->			Any:
     """simple docstring"""




     _SCREAMING_SNAKE_CASE							    =ArgumentParser(
         description=(
             'PyTorch TPU distributed training launch helper utility that will spawn up multiple distributed processes'
         )				)

     # Optional arguments for the launch helper
     parser.add_argument('--num_cores'							, type=_UpperCamelCase							, default=1							, help='Number of TPU cores to use (1 or 8).'				)

     # positional
     parser.add_argument(
         'training_script'							, type=_UpperCamelCase							, help=(
             'The full path to the single TPU training '
             'program/script to be launched in parallel, '
             'followed by all the arguments for the '
             'training script'
         )							, )

     # rest from the training program
     parser.add_argument('training_script_args'							, nargs=_UpperCamelCase				)

     return parser.parse_args()


def        _lowerCAmelCase					(						)      ->			Optional[int]:
     """simple docstring"""




     _SCREAMING_SNAKE_CASE							    =parse_args()

     # Import training_script as a module.
     _SCREAMING_SNAKE_CASE							    =Path(args.training_script				)
     sys.path.append(str(script_fpath.parent.resolve()				)				)
     _SCREAMING_SNAKE_CASE							    =script_fpath.stem
     _SCREAMING_SNAKE_CASE							    =importlib.import_module(_UpperCamelCase				)

     # Patch sys.argv
     _SCREAMING_SNAKE_CASE							    =[args.training_script] + args.training_script_args + ['--tpu_num_cores', str(args.num_cores				)]

     xmp.spawn(mod._mp_fn							, args=()							, nprocs=args.num_cores				)


if __name__ == "__main__":
     main()
