


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  (						_UpperCAmelCase    ,       _UpperCAmelCase    ,       _UpperCAmelCase    ,       unittest.TestCase			):
     lowerCamelCase			:			Optional[int]      = AltDiffusionPipeline
     lowerCamelCase			:			List[Any]      = TEXT_TO_IMAGE_PARAMS
     lowerCamelCase			:			Tuple      = TEXT_TO_IMAGE_BATCH_PARAMS
     lowerCamelCase			:			Any      = TEXT_TO_IMAGE_IMAGE_PARAMS
     lowerCamelCase			:			Dict      = TEXT_TO_IMAGE_IMAGE_PARAMS







     def 							lowercase__       (						self :						Optional[Any]      ):
           torch.manual_seed(0      )
           SCREAMING_SNAKE_CASE__       :  Dict						     =		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__       :  List[str]						     =		DDIMScheduler(
               beta_start=0.00085   ,   beta_end=0.012   ,   beta_schedule='''scaled_linear'''   ,   clip_sample=_lowercase   ,   set_alpha_to_one=_lowercase   ,   )
           torch.manual_seed(0      )
           SCREAMING_SNAKE_CASE__       :  Optional[int]						     =		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      )
           SCREAMING_SNAKE_CASE__       :  int						     =		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=50_02   ,   )
           SCREAMING_SNAKE_CASE__       :  Dict						     =		CLIPTextModel(_lowercase      )

           SCREAMING_SNAKE_CASE__       :  Optional[int]						     =		XLMRobertaTokenizer.from_pretrained('''hf-internal-testing/tiny-xlm-roberta'''      )
           SCREAMING_SNAKE_CASE__       :  List[str]						     =		77

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







     def 							lowercase__       (						self :						Union[str, Any]   ,   _lowercase :						int   ,   _lowercase :						List[Any]=0      ):
           if str(_lowercase      ).startswith('''mps'''      ):
                 SCREAMING_SNAKE_CASE__       :  Tuple						     =		torch.manual_seed(_lowercase      )
           else:
                 SCREAMING_SNAKE_CASE__       :  str						     =		torch.Generator(device=_lowercase      ).manual_seed(_lowercase      )
           SCREAMING_SNAKE_CASE__       :  List[Any]						     =		{
               '''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 :						Optional[Any]      ):
           super().test_attention_slicing_forward_pass(expected_max_diff=3E-3      )







     def 							lowercase__       (						self :						Dict      ):
           super().test_inference_batch_single_identical(expected_max_diff=3E-3      )







     def 							lowercase__       (						self :						Optional[Any]      ):
           SCREAMING_SNAKE_CASE__       :  List[str]						     =		'''cpu'''  # ensure determinism for the device-dependent torch.Generator

           SCREAMING_SNAKE_CASE__       :  str						     =		self.get_dummy_components()
           torch.manual_seed(0      )
           SCREAMING_SNAKE_CASE__       :  Dict						     =		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=50_02   ,   )
           # TODO: remove after fixing the non-deterministic text encoder
           SCREAMING_SNAKE_CASE__       :  Optional[Any]						     =		RobertaSeriesModelWithTransformation(_lowercase      )
           SCREAMING_SNAKE_CASE__       :  Union[str, Any]						     =		text_encoder

           SCREAMING_SNAKE_CASE__       :  Optional[Any]						     =		AltDiffusionPipeline(**_lowercase      )
           SCREAMING_SNAKE_CASE__       :  Any						     =		alt_pipe.to(_lowercase      )
           alt_pipe.set_progress_bar_config(disable=_lowercase      )

           SCREAMING_SNAKE_CASE__       :  List[str]						     =		self.get_dummy_inputs(_lowercase      )
           SCREAMING_SNAKE_CASE__       :  List[str]						     =		'''A photo of an astronaut'''
           SCREAMING_SNAKE_CASE__       :  Optional[Any]						     =		alt_pipe(**_lowercase      )
           SCREAMING_SNAKE_CASE__       :  List[Any]						     =		output.images
           SCREAMING_SNAKE_CASE__       :  Dict						     =		image[0, -3:, -3:, -1]

           assert image.shape == (1, 64, 64, 3)
           SCREAMING_SNAKE_CASE__       :  List[Any]						     =		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 :						List[Any]      ):
           SCREAMING_SNAKE_CASE__       :  Tuple						     =		'''cpu'''  # ensure determinism for the device-dependent torch.Generator

           SCREAMING_SNAKE_CASE__       :  Optional[int]						     =		self.get_dummy_components()
           SCREAMING_SNAKE_CASE__       :  Any						     =		PNDMScheduler(skip_prk_steps=_lowercase      )
           torch.manual_seed(0      )
           SCREAMING_SNAKE_CASE__       :  Optional[Any]						     =		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=50_02   ,   )
           # TODO: remove after fixing the non-deterministic text encoder
           SCREAMING_SNAKE_CASE__       :  Optional[int]						     =		RobertaSeriesModelWithTransformation(_lowercase      )
           SCREAMING_SNAKE_CASE__       :  str						     =		text_encoder
           SCREAMING_SNAKE_CASE__       :  Any						     =		AltDiffusionPipeline(**_lowercase      )
           SCREAMING_SNAKE_CASE__       :  Optional[int]						     =		alt_pipe.to(_lowercase      )
           alt_pipe.set_progress_bar_config(disable=_lowercase      )

           SCREAMING_SNAKE_CASE__       :  List[str]						     =		self.get_dummy_inputs(_lowercase      )
           SCREAMING_SNAKE_CASE__       :  int						     =		alt_pipe(**_lowercase      )
           SCREAMING_SNAKE_CASE__       :  Dict						     =		output.images
           SCREAMING_SNAKE_CASE__       :  int						     =		image[0, -3:, -3:, -1]

           assert image.shape == (1, 64, 64, 3)
           SCREAMING_SNAKE_CASE__       :  Union[str, Any]						     =		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 							lowercase  (						unittest.TestCase			):







     def 							lowercase__       (						self :						int      ):
           # clean up the VRAM after each test
           super().tearDown()
           gc.collect()
           torch.cuda.empty_cache()







     def 							lowercase__       (						self :						Dict      ):
           # make sure here that pndm scheduler skips prk
           SCREAMING_SNAKE_CASE__       :  Any						     =		AltDiffusionPipeline.from_pretrained('''BAAI/AltDiffusion'''   ,   safety_checker=_lowercase      )
           SCREAMING_SNAKE_CASE__       :  Any						     =		alt_pipe.to(_lowercase      )
           alt_pipe.set_progress_bar_config(disable=_lowercase      )

           SCREAMING_SNAKE_CASE__       :  str						     =		'''A painting of a squirrel eating a burger'''
           SCREAMING_SNAKE_CASE__       :  str						     =		torch.manual_seed(0      )
           SCREAMING_SNAKE_CASE__       :  Any						     =		alt_pipe([prompt]   ,   generator=_lowercase   ,   guidance_scale=6.0   ,   num_inference_steps=20   ,   output_type='''np'''      )

           SCREAMING_SNAKE_CASE__       :  List[Any]						     =		output.images

           SCREAMING_SNAKE_CASE__       :  List[str]						     =		image[0, -3:, -3:, -1]

           assert image.shape == (1, 5_12, 5_12, 3)
           SCREAMING_SNAKE_CASE__       :  int						     =		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 :						Optional[int]      ):
           SCREAMING_SNAKE_CASE__       :  Tuple						     =		DDIMScheduler.from_pretrained('''BAAI/AltDiffusion'''   ,   subfolder='''scheduler'''      )

           SCREAMING_SNAKE_CASE__       :  Optional[Any]						     =		AltDiffusionPipeline.from_pretrained('''BAAI/AltDiffusion'''   ,   scheduler=_lowercase   ,   safety_checker=_lowercase      )
           SCREAMING_SNAKE_CASE__       :  Optional[Any]						     =		alt_pipe.to(_lowercase      )
           alt_pipe.set_progress_bar_config(disable=_lowercase      )

           SCREAMING_SNAKE_CASE__       :  List[Any]						     =		'''A painting of a squirrel eating a burger'''
           SCREAMING_SNAKE_CASE__       :  Optional[Any]						     =		torch.manual_seed(0      )

           SCREAMING_SNAKE_CASE__       :  Tuple						     =		alt_pipe([prompt]   ,   generator=_lowercase   ,   num_inference_steps=2   ,   output_type='''numpy'''      )
           SCREAMING_SNAKE_CASE__       :  Any						     =		output.images

           SCREAMING_SNAKE_CASE__       :  Optional[int]						     =		image[0, -3:, -3:, -1]

           assert image.shape == (1, 5_12, 5_12, 3)
           SCREAMING_SNAKE_CASE__       :  Optional[Any]						     =		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













'''simple docstring'''




import os
from collections.abc import Iterator





def 	__lowerCamelCase			(      __lowerCAmelCase       :							str = "."		)		->  Iterator[str]:
   for dir_path, dir_names, filenames in os.walk(__lowerCAmelCase		):
      snake_case 						=   [d for d in dir_names if d != """scripts""" and d[0] not in """._"""]
      for filename in filenames:
         if filename == "__init__.py":
            continue
         if os.path.splitext(__lowerCAmelCase		)[1] in (".py", ".ipynb"):
            yield os.path.join(__lowerCAmelCase      ,    __lowerCAmelCase		).lstrip("""./"""		)





def 	__lowerCamelCase			(      __lowerCAmelCase       :							List[str]		)		->  Any:
   return F'''{i * "  "}*''' if i else "\n##"





def 	__lowerCamelCase			(      __lowerCAmelCase       :							str      ,    __lowerCAmelCase       :							str		)		->  str:
   snake_case 						=   old_path.split(os.sep		)
   for i, new_part in enumerate(new_path.split(os.sep		)		):
      if (i + 1 > len(__lowerCAmelCase		) or old_parts[i] != new_part) and new_part:
         print(F'''{md_prefix(__lowerCAmelCase		)} {new_part.replace("_"      ,    " "		).title()}'''		)
   return new_path





def 	__lowerCamelCase			(      __lowerCAmelCase       :							str = "."		)		->  None:
   snake_case 						=   """"""
   for filepath in sorted(good_file_paths(__lowerCAmelCase		)		):
      snake_case						,    snake_case 						=   os.path.split(__lowerCAmelCase		)
      if filepath != old_path:
         snake_case 						=   print_path(__lowerCAmelCase      ,    __lowerCAmelCase		)
      snake_case 						=   (filepath.count(os.sep		) + 1) if filepath else 0
      snake_case 						=   F'''{filepath}/{filename}'''.replace(""" """      ,    """%20"""		)
      snake_case 						=   os.path.splitext(filename.replace("""_"""      ,    """ """		).title()		)[0]
      print(F'''{md_prefix(__lowerCAmelCase		)} [{filename}]({url})'''		)


if __name__ == "__main__":
 print_directory_md(".")

