
import gc
import unittest

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

from diffusers import AltDiffusionPipeline, AutoencoderKL, DDIMScheduler, PNDMScheduler, UNetaDConditionModel
from diffusers.pipelines.alt_diffusion.modeling_roberta_series import (
    RobertaSeriesConfig,
    RobertaSeriesModelWithTransformation,
)
from diffusers.utils import slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu

from ..pipeline_params import TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_IMAGE_PARAMS, TEXT_TO_IMAGE_PARAMS
from ..test_pipelines_common import PipelineKarrasSchedulerTesterMixin, PipelineLatentTesterMixin, PipelineTesterMixin


enable_full_determinism()




class        UpperCAmelCase_			( _A  ,       _A  ,       _A  ,       unittest.TestCase     ):


  '''simple docstring'''







  a__           =		AltDiffusionPipeline
  a__           =		TEXT_TO_IMAGE_PARAMS
  a__           =		TEXT_TO_IMAGE_BATCH_PARAMS
  a__           =		TEXT_TO_IMAGE_IMAGE_PARAMS
  a__           =		TEXT_TO_IMAGE_IMAGE_PARAMS
  def 							_lowercase							( self			: Dict    )			->	Optional[Any]:







    """simple docstring"""






    torch.manual_seed(0    )
    __magic_name__				     =       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			,							)
    __magic_name__				     =       DDIMScheduler(
        beta_start=0.00085			,							beta_end=0.012			,							beta_schedule="""scaled_linear"""			,							clip_sample=UpperCamelCase__			,							set_alpha_to_one=UpperCamelCase__			,							)
    torch.manual_seed(0    )
    __magic_name__				     =       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			,							)

    # TODO: address the non-deterministic text encoder (fails for save-load tests)
    # torch.manual_seed(0)
    # text_encoder_config = RobertaSeriesConfig(
    #     hidden_size=32,
    #     project_dim=32,
    #     intermediate_size=37,
    #     layer_norm_eps=1e-05,
    #     num_attention_heads=4,
    #     num_hidden_layers=5,
    #     vocab_size=5002,
    # )
    # text_encoder = RobertaSeriesModelWithTransformation(text_encoder_config)

    torch.manual_seed(0    )
    __magic_name__				     =       CLIPTextConfig(
        bos_token_id=0			,							eos_token_id=2			,							hidden_size=32			,							projection_dim=32			,							intermediate_size=37			,							layer_norm_eps=1E-05			,							num_attention_heads=4			,							num_hidden_layers=5			,							pad_token_id=1			,							vocab_size=5002			,							)
    __magic_name__				     =       CLIPTextModel(UpperCamelCase__    )

    __magic_name__				     =       XLMRobertaTokenizer.from_pretrained("""hf-internal-testing/tiny-xlm-roberta"""    )
    __magic_name__				     =       77

    __magic_name__				     =       {
        """unet""": unet,
        """scheduler""": scheduler,
        """vae""": vae,
        """text_encoder""": text_encoder,
        """tokenizer""": tokenizer,
        """safety_checker""": None,
        """feature_extractor""": None,
    }
    return components
  def 							_lowercase							( self			: Optional[int]			,							UpperCamelCase__			: int			,							UpperCamelCase__			: Optional[Any]=0    )			->	Any:







    """simple docstring"""






    if str(UpperCamelCase__    ).startswith("""mps"""    ):
      __magic_name__				     =       torch.manual_seed(UpperCamelCase__    )
    else:
      __magic_name__				     =       torch.Generator(device=UpperCamelCase__    ).manual_seed(UpperCamelCase__    )
    __magic_name__				     =       {
        """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 							_lowercase							( self			: Union[str, Any]    )			->	int:







    """simple docstring"""






    super().test_attention_slicing_forward_pass(expected_max_diff=3E-3    )
  def 							_lowercase							( self			: List[Any]    )			->	Union[str, Any]:







    """simple docstring"""






    super().test_inference_batch_single_identical(expected_max_diff=3E-3    )
  def 							_lowercase							( self			: Optional[Any]    )			->	str:







    """simple docstring"""






    __magic_name__				     =       """cpu"""  # ensure determinism for the device-dependent torch.Generator

    __magic_name__				     =       self.get_dummy_components()
    torch.manual_seed(0    )
    __magic_name__				     =       RobertaSeriesConfig(
        hidden_size=32			,							project_dim=32			,							intermediate_size=37			,							layer_norm_eps=1E-05			,							num_attention_heads=4			,							num_hidden_layers=5			,							vocab_size=5002			,							)
    # TODO: remove after fixing the non-deterministic text encoder
    __magic_name__				     =       RobertaSeriesModelWithTransformation(UpperCamelCase__    )
    __magic_name__				     =       text_encoder

    __magic_name__				     =       AltDiffusionPipeline(**UpperCamelCase__    )
    __magic_name__				     =       alt_pipe.to(UpperCamelCase__    )
    alt_pipe.set_progress_bar_config(disable=UpperCamelCase__    )

    __magic_name__				     =       self.get_dummy_inputs(UpperCamelCase__    )
    __magic_name__				     =       """A photo of an astronaut"""
    __magic_name__				     =       alt_pipe(**UpperCamelCase__    )
    __magic_name__				     =       output.images
    __magic_name__				     =       image[0, -3:, -3:, -1]

    assert image.shape == (1, 64, 64, 3)
    __magic_name__				     =       np.array(
        [0.5748162, 0.60447145, 0.48821217, 0.50100636, 0.5431185, 0.45763683, 0.49657696, 0.48132733, 0.47573093]    )

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





  def 							_lowercase							( self			: Any    )			->	int:







    """simple docstring"""






    __magic_name__				     =       """cpu"""  # ensure determinism for the device-dependent torch.Generator

    __magic_name__				     =       self.get_dummy_components()
    __magic_name__				     =       PNDMScheduler(skip_prk_steps=UpperCamelCase__    )
    torch.manual_seed(0    )
    __magic_name__				     =       RobertaSeriesConfig(
        hidden_size=32			,							project_dim=32			,							intermediate_size=37			,							layer_norm_eps=1E-05			,							num_attention_heads=4			,							num_hidden_layers=5			,							vocab_size=5002			,							)
    # TODO: remove after fixing the non-deterministic text encoder
    __magic_name__				     =       RobertaSeriesModelWithTransformation(UpperCamelCase__    )
    __magic_name__				     =       text_encoder
    __magic_name__				     =       AltDiffusionPipeline(**UpperCamelCase__    )
    __magic_name__				     =       alt_pipe.to(UpperCamelCase__    )
    alt_pipe.set_progress_bar_config(disable=UpperCamelCase__    )

    __magic_name__				     =       self.get_dummy_inputs(UpperCamelCase__    )
    __magic_name__				     =       alt_pipe(**UpperCamelCase__    )
    __magic_name__				     =       output.images
    __magic_name__				     =       image[0, -3:, -3:, -1]

    assert image.shape == (1, 64, 64, 3)
    __magic_name__				     =       np.array(
        [0.51605093, 0.5707241, 0.47365507, 0.50578886, 0.5633877, 0.4642503, 0.5182081, 0.48763484, 0.49084237]    )

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







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


  '''simple docstring'''







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







    """simple docstring"""






    super().tearDown()
    gc.collect()
    torch.cuda.empty_cache()
  def 							_lowercase							( self			: Tuple    )			->	Tuple:







    """simple docstring"""






    __magic_name__				     =       AltDiffusionPipeline.from_pretrained("""BAAI/AltDiffusion"""			,							safety_checker=UpperCamelCase__    )
    __magic_name__				     =       alt_pipe.to(UpperCamelCase__    )
    alt_pipe.set_progress_bar_config(disable=UpperCamelCase__    )

    __magic_name__				     =       """A painting of a squirrel eating a burger"""
    __magic_name__				     =       torch.manual_seed(0    )
    __magic_name__				     =       alt_pipe([prompt]			,							generator=UpperCamelCase__			,							guidance_scale=6.0			,							num_inference_steps=20			,							output_type="""np"""    )

    __magic_name__				     =       output.images

    __magic_name__				     =       image[0, -3:, -3:, -1]

    assert image.shape == (1, 512, 512, 3)
    __magic_name__				     =       np.array([0.1010, 0.0800, 0.0794, 0.0885, 0.0843, 0.0762, 0.0769, 0.0729, 0.0586]    )

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





  def 							_lowercase							( self			: str    )			->	Dict:







    """simple docstring"""






    __magic_name__				     =       DDIMScheduler.from_pretrained("""BAAI/AltDiffusion"""			,							subfolder="""scheduler"""    )

    __magic_name__				     =       AltDiffusionPipeline.from_pretrained("""BAAI/AltDiffusion"""			,							scheduler=UpperCamelCase__			,							safety_checker=UpperCamelCase__    )
    __magic_name__				     =       alt_pipe.to(UpperCamelCase__    )
    alt_pipe.set_progress_bar_config(disable=UpperCamelCase__    )

    __magic_name__				     =       """A painting of a squirrel eating a burger"""
    __magic_name__				     =       torch.manual_seed(0    )

    __magic_name__				     =       alt_pipe([prompt]			,							generator=UpperCamelCase__			,							num_inference_steps=2			,							output_type="""numpy"""    )
    __magic_name__				     =       output.images

    __magic_name__				     =       image[0, -3:, -3:, -1]

    assert image.shape == (1, 512, 512, 3)
    __magic_name__				     =       np.array([0.4019, 0.4052, 0.3810, 0.4119, 0.3916, 0.3982, 0.4651, 0.4195, 0.5323]    )

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



import importlib
import sys
from argparse import REMAINDER, ArgumentParser
from pathlib import Path

import torch_xla.distributed.xla_multiprocessing as xmp

def a__    (       ):






  '''simple docstring'''





  __magic_name__				     =       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=A_,     default=1,     help="""Number of TPU cores to use (1 or 8)."""					)

  # positional
  parser.add_argument(
      """training_script""",     type=A_,     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=A_					)

  return parser.parse_args()

def a__    (       ):






  '''simple docstring'''





  __magic_name__				     =       parse_args()

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

  # Patch sys.argv
  __magic_name__				     =       [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()

