






'''simple docstring'''

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        a_		(    lowerCamelCase			, lowerCamelCase			, lowerCamelCase			, unittest.TestCase       ):
  lowercase  	=   AltDiffusionPipeline
  lowercase  	=   TEXT_TO_IMAGE_PARAMS
  lowercase  	=   TEXT_TO_IMAGE_BATCH_PARAMS
  lowercase  	=   TEXT_TO_IMAGE_IMAGE_PARAMS
  lowercase  	=   TEXT_TO_IMAGE_IMAGE_PARAMS

  def 						A__		(   self   )						->   int:



        """simple docstring"""






        torch.manual_seed(0   )
        UpperCamelCase       						=      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    ,	)
        UpperCamelCase       						=      DDIMScheduler(
            beta_start=0.0_0_0_8_5    ,	beta_end=0.0_1_2    ,	beta_schedule="""scaled_linear"""    ,	clip_sample=_SCREAMING_SNAKE_CASE    ,	set_alpha_to_one=_SCREAMING_SNAKE_CASE    ,	)
        torch.manual_seed(0   )
        UpperCamelCase       						=      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   )
        UpperCamelCase       						=      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    ,	)
        UpperCamelCase       						=      CLIPTextModel(_SCREAMING_SNAKE_CASE   )

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

        UpperCamelCase       						=      {
            """unet""": unet,
            """scheduler""": scheduler,
            """vae""": vae,
            """text_encoder""": text_encoder,
            """tokenizer""": tokenizer,
            """safety_checker""": None,
            """feature_extractor""": None,
        }
        return components

  def 						A__		(   self    ,	_SCREAMING_SNAKE_CASE    ,	_SCREAMING_SNAKE_CASE=0   )						->   Union[str, Any]:



        """simple docstring"""






        if str(_SCREAMING_SNAKE_CASE   ).startswith("""mps"""   ):
              UpperCamelCase       						=      torch.manual_seed(_SCREAMING_SNAKE_CASE   )
        else:
              UpperCamelCase       						=      torch.Generator(device=_SCREAMING_SNAKE_CASE   ).manual_seed(_SCREAMING_SNAKE_CASE   )
        UpperCamelCase       						=      {
            """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   )						->   Any:



        """simple docstring"""






        super().test_attention_slicing_forward_pass(expected_max_diff=3e-3   )

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



        """simple docstring"""






        super().test_inference_batch_single_identical(expected_max_diff=3e-3   )

  def 						A__		(   self   )						->   Any:



        """simple docstring"""






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

        UpperCamelCase       						=      self.get_dummy_components()
        torch.manual_seed(0   )
        UpperCamelCase       						=      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
        UpperCamelCase       						=      RobertaSeriesModelWithTransformation(_SCREAMING_SNAKE_CASE   )
        UpperCamelCase       						=      text_encoder

        UpperCamelCase       						=      AltDiffusionPipeline(**_SCREAMING_SNAKE_CASE   )
        UpperCamelCase       						=      alt_pipe.to(_SCREAMING_SNAKE_CASE   )
        alt_pipe.set_progress_bar_config(disable=_SCREAMING_SNAKE_CASE   )

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

        assert image.shape == (1, 64, 64, 3)
        UpperCamelCase       						=      np.array(
            [0.5_7_4_8_1_6_2, 0.6_0_4_4_7_1_4_5, 0.4_8_8_2_1_2_1_7, 0.5_0_1_0_0_6_3_6, 0.5_4_3_1_1_8_5, 0.4_5_7_6_3_6_8_3, 0.4_9_6_5_7_6_9_6, 0.4_8_1_3_2_7_3_3, 0.4_7_5_7_3_0_9_3]   )

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


  def 						A__		(   self   )						->   int:



        """simple docstring"""






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

        UpperCamelCase       						=      self.get_dummy_components()
        UpperCamelCase       						=      PNDMScheduler(skip_prk_steps=_SCREAMING_SNAKE_CASE   )
        torch.manual_seed(0   )
        UpperCamelCase       						=      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
        UpperCamelCase       						=      RobertaSeriesModelWithTransformation(_SCREAMING_SNAKE_CASE   )
        UpperCamelCase       						=      text_encoder
        UpperCamelCase       						=      AltDiffusionPipeline(**_SCREAMING_SNAKE_CASE   )
        UpperCamelCase       						=      alt_pipe.to(_SCREAMING_SNAKE_CASE   )
        alt_pipe.set_progress_bar_config(disable=_SCREAMING_SNAKE_CASE   )

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

        assert image.shape == (1, 64, 64, 3)
        UpperCamelCase       						=      np.array(
            [0.5_1_6_0_5_0_9_3, 0.5_7_0_7_2_4_1, 0.4_7_3_6_5_5_0_7, 0.5_0_5_7_8_8_8_6, 0.5_6_3_3_8_7_7, 0.4_6_4_2_5_0_3, 0.5_1_8_2_0_8_1, 0.4_8_7_6_3_4_8_4, 0.4_9_0_8_4_2_3_7]   )

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






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

  def 						A__		(   self   )						->   Tuple:



        """simple docstring"""






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

  def 						A__		(   self   )						->   int:



        """simple docstring"""






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

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

        UpperCamelCase       						=      output.images

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

        assert image.shape == (1, 512, 512, 3)
        UpperCamelCase       						=      np.array([0.1_0_1_0, 0.0_8_0_0, 0.0_7_9_4, 0.0_8_8_5, 0.0_8_4_3, 0.0_7_6_2, 0.0_7_6_9, 0.0_7_2_9, 0.0_5_8_6]   )

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


  def 						A__		(   self   )						->   Optional[int]:



        """simple docstring"""






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

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

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

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

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

        assert image.shape == (1, 512, 512, 3)
        UpperCamelCase       						=      np.array([0.4_0_1_9, 0.4_0_5_2, 0.3_8_1_0, 0.4_1_1_9, 0.3_9_1_6, 0.3_9_8_2, 0.4_6_5_1, 0.4_1_9_5, 0.5_3_2_3]   )

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








'''simple docstring'''

# Copyright 2022 The HuggingFace Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import argparse
import os
import platform

import numpy as np
import psutil
import torch

from accelerate import __version__ as version
from accelerate.commands.config import default_config_file, load_config_from_file

from ..utils import is_npu_available, is_xpu_available







def 					lowercase__				(   __UpperCamelCase=None					)->				Union[str, Any]:
      if subparsers is not None:
            UpperCamelCase       						=      subparsers.add_parser("""env"""					)
      else:
            UpperCamelCase       						=      argparse.ArgumentParser("""Accelerate env command"""					)

      parser.add_argument(
          """--config_file"""	,		default=__UpperCamelCase	,		help="""The config file to use for the default values in the launching script."""					)

      if subparsers is not None:
            parser.set_defaults(func=__UpperCamelCase					)
      return parser







def 					lowercase__				(   __UpperCamelCase					)->				List[str]:
      UpperCamelCase       						=      torch.__version__
      UpperCamelCase       						=      torch.cuda.is_available()
      UpperCamelCase       						=      is_xpu_available()
      UpperCamelCase       						=      is_npu_available()

      UpperCamelCase       						=      """Not found"""
      # Get the default from the config file.
      if args.config_file is not None or os.path.isfile(__UpperCamelCase					):
            UpperCamelCase       						=      load_config_from_file(args.config_file					).to_dict()

      UpperCamelCase       						=      {
          """`Accelerate` version""": version,
          """Platform""": platform.platform(),
          """Python version""": platform.python_version(),
          """Numpy version""": np.__version__,
          """PyTorch version (GPU?)""": F"{pt_version} ({pt_cuda_available})",
          """PyTorch XPU available""": str(__UpperCamelCase					),
          """PyTorch NPU available""": str(__UpperCamelCase					),
          """System RAM""": F"{psutil.virtual_memory().total / 1024 ** 3:.2f} GB",
      }
      if pt_cuda_available:
            UpperCamelCase       						=      torch.cuda.get_device_name()

      print("""\nCopy-and-paste the text below in your GitHub issue\n"""					)
      print("""\n""".join([F"- {prop}: {val}" for prop, val in info.items()]					)					)

      print("""- `Accelerate` default config:""" if args.config_file is None else """- `Accelerate` config passed:"""					)
      UpperCamelCase       						=      (
          """\n""".join([F"\t- {prop}: {val}" for prop, val in accelerate_config.items()]					)
          if isinstance(__UpperCamelCase	,		__UpperCamelCase					)
          else F"\t{accelerate_config}"
      )
      print(__UpperCamelCase					)

      UpperCamelCase       						=      accelerate_config

      return info







def 					lowercase__				(   )->				int:
      UpperCamelCase       						=      env_command_parser()
      UpperCamelCase       						=      parser.parse_args()
      env_command(__UpperCamelCase					)
      return 0


if __name__ == "__main__":
     raise SystemExit(main())
