




'''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  __lowercase    (	__magic_name__							,				__magic_name__							,				__magic_name__							,				unittest.TestCase			):
   _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 		UpperCamelCase__      (						self   )					->     Optional[int]:
     torch.manual_seed(0   )
     __a								=       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   ,				)
     __a								=       DDIMScheduler(
         beta_start=0.00_085   ,				beta_end=0.012   ,				beta_schedule='scaled_linear'   ,				clip_sample=UpperCamelCase   ,				set_alpha_to_one=UpperCamelCase   ,				)
     torch.manual_seed(0   )
     __a								=       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   )
     __a								=       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   ,				)
     __a								=       CLIPTextModel(UpperCamelCase   )

     __a								=       XLMRobertaTokenizer.from_pretrained('hf-internal-testing/tiny-xlm-roberta'   )
     __a								=       77

     __a								=       {
         'unet': unet,
         'scheduler': scheduler,
         'vae': vae,
         'text_encoder': text_encoder,
         'tokenizer': tokenizer,
         'safety_checker': None,
         'feature_extractor': None,
     }
     return components







   def 		UpperCamelCase__      (						self   ,				UpperCamelCase   ,				UpperCamelCase=0   )					->     int:
     if str(UpperCamelCase   ).startswith('mps'   ):
       __a								=       torch.manual_seed(UpperCamelCase   )
     else:
       __a								=       torch.Generator(device=UpperCamelCase   ).manual_seed(UpperCamelCase   )
     __a								=       {
         '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 		UpperCamelCase__      (						self   )					->     Dict:
     super().test_attention_slicing_forward_pass(expected_max_diff=3e-3   )







   def 		UpperCamelCase__      (						self   )					->     Dict:
     super().test_inference_batch_single_identical(expected_max_diff=3e-3   )







   def 		UpperCamelCase__      (						self   )					->     Dict:
     __a								=       'cpu'  # ensure determinism for the device-dependent torch.Generator

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

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

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

     assert image.shape == (1, 64, 64, 3)
     __a								=       np.array(
         [0.5_748_162, 0.60_447_145, 0.48_821_217, 0.50_100_636, 0.5_431_185, 0.45_763_683, 0.49_657_696, 0.48_132_733, 0.47_573_093]   )

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







   def 		UpperCamelCase__      (						self   )					->     Any:
     __a								=       'cpu'  # ensure determinism for the device-dependent torch.Generator

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

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

     assert image.shape == (1, 64, 64, 3)
     __a								=       np.array(
         [0.51_605_093, 0.5_707_241, 0.47_365_507, 0.50_578_886, 0.5_633_877, 0.4_642_503, 0.5_182_081, 0.48_763_484, 0.49_084_237]   )

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





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







   def 		UpperCamelCase__      (						self   )					->     Optional[int]:
     # clean up the VRAM after each test
     super().tearDown()
     gc.collect()
     torch.cuda.empty_cache()







   def 		UpperCamelCase__      (						self   )					->     Optional[int]:
     # make sure here that pndm scheduler skips prk
     __a								=       AltDiffusionPipeline.from_pretrained('BAAI/AltDiffusion'   ,				safety_checker=UpperCamelCase   )
     __a								=       alt_pipe.to(UpperCamelCase   )
     alt_pipe.set_progress_bar_config(disable=UpperCamelCase   )

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

     __a								=       output.images

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

     assert image.shape == (1, 512, 512, 3)
     __a								=       np.array([0.1_010, 0.0_800, 0.0_794, 0.0_885, 0.0_843, 0.0_762, 0.0_769, 0.0_729, 0.0_586]   )

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







   def 		UpperCamelCase__      (						self   )					->     Optional[Any]:
     __a								=       DDIMScheduler.from_pretrained('BAAI/AltDiffusion'   ,				subfolder='scheduler'   )

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

     __a								=       'A painting of a squirrel eating a burger'
     __a								=       torch.manual_seed(0   )

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

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

     assert image.shape == (1, 512, 512, 3)
     __a								=       np.array([0.4_019, 0.4_052, 0.3_810, 0.4_119, 0.3_916, 0.3_982, 0.4_651, 0.4_195, 0.5_323]   )

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









'''simple docstring'''


from typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_tf_available, is_torch_available


UpperCAmelCase_					=   {"configuration_swin": ["SWIN_PRETRAINED_CONFIG_ARCHIVE_MAP", "SwinConfig", "SwinOnnxConfig"]}


try:
  if not is_torch_available():
    raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
  pass
else:
  UpperCAmelCase_					=   [
      "SWIN_PRETRAINED_MODEL_ARCHIVE_LIST",
      "SwinForImageClassification",
      "SwinForMaskedImageModeling",
      "SwinModel",
      "SwinPreTrainedModel",
      "SwinBackbone",
  ]

try:
  if not is_tf_available():
    raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
  pass
else:
  UpperCAmelCase_					=   [
      "TF_SWIN_PRETRAINED_MODEL_ARCHIVE_LIST",
      "TFSwinForImageClassification",
      "TFSwinForMaskedImageModeling",
      "TFSwinModel",
      "TFSwinPreTrainedModel",
  ]

if TYPE_CHECKING:
  from .configuration_swin import SWIN_PRETRAINED_CONFIG_ARCHIVE_MAP, SwinConfig, SwinOnnxConfig

  try:
    if not is_torch_available():
      raise OptionalDependencyNotAvailable()
  except OptionalDependencyNotAvailable:
    pass
  else:
    from .modeling_swin import (
        SWIN_PRETRAINED_MODEL_ARCHIVE_LIST,
        SwinBackbone,
        SwinForImageClassification,
        SwinForMaskedImageModeling,
        SwinModel,
        SwinPreTrainedModel,
    )

  try:
    if not is_tf_available():
      raise OptionalDependencyNotAvailable()
  except OptionalDependencyNotAvailable:
    pass
  else:
    from .modeling_tf_swin import (
        TF_SWIN_PRETRAINED_MODEL_ARCHIVE_LIST,
        TFSwinForImageClassification,
        TFSwinForMaskedImageModeling,
        TFSwinModel,
        TFSwinPreTrainedModel,
    )

else:
  import sys

  UpperCAmelCase_					=   _LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__)



