


import gc
import unittest

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

from diffusers import (
    AutoencoderKL,
    DDIMParallelScheduler,
    DDPMParallelScheduler,
    StableDiffusionParadigmsPipeline,
    UNetaDConditionModel,
)
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 PipelineLatentTesterMixin, PipelineTesterMixin


enable_full_determinism()







class    lowerCamelCase__     ( lowerCAmelCase     ,							lowerCAmelCase     ,							unittest.TestCase):
     SCREAMING_SNAKE_CASE__       =    StableDiffusionParadigmsPipeline
     SCREAMING_SNAKE_CASE__       =    TEXT_TO_IMAGE_PARAMS
     SCREAMING_SNAKE_CASE__       =    TEXT_TO_IMAGE_BATCH_PARAMS
     SCREAMING_SNAKE_CASE__       =    TEXT_TO_IMAGE_IMAGE_PARAMS
     SCREAMING_SNAKE_CASE__       =    TEXT_TO_IMAGE_IMAGE_PARAMS





     def   __A      (self )  ->     Optional[int]:
      torch.manual_seed(0 )
      _lowercase							   =UNetaDConditionModel(
          block_out_channels=(3_2, 6_4)	,  layers_per_block=2	,  sample_size=3_2	,  in_channels=4	,  out_channels=4	,  down_block_types=('''DownBlock2D''', '''CrossAttnDownBlock2D''')	,  up_block_types=('''CrossAttnUpBlock2D''', '''UpBlock2D''')	,  cross_attention_dim=3_2	,  attention_head_dim=(2, 4)	,  use_linear_projection=UpperCAmelCase	,  )
      _lowercase							   =DDIMParallelScheduler(
          beta_start=0.0_0085	,  beta_end=0.012	,  beta_schedule='''scaled_linear'''	,  clip_sample=UpperCAmelCase	,  set_alpha_to_one=UpperCAmelCase	,  )
      torch.manual_seed(0 )
      _lowercase							   =AutoencoderKL(
          block_out_channels=[3_2, 6_4]	,  in_channels=3	,  out_channels=3	,  down_block_types=['''DownEncoderBlock2D''', '''DownEncoderBlock2D''']	,  up_block_types=['''UpDecoderBlock2D''', '''UpDecoderBlock2D''']	,  latent_channels=4	,  sample_size=1_2_8	,  )
      torch.manual_seed(0 )
      _lowercase							   =CLIPTextConfig(
          bos_token_id=0	,  eos_token_id=2	,  hidden_size=3_2	,  intermediate_size=3_7	,  layer_norm_eps=1e-05	,  num_attention_heads=4	,  num_hidden_layers=5	,  pad_token_id=1	,  vocab_size=1_0_0_0	,  hidden_act='''gelu'''	,  projection_dim=5_1_2	,  )
      _lowercase							   =CLIPTextModel(UpperCAmelCase )
      _lowercase							   =CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip''' )

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





     def   __A      (self	,  UpperCAmelCase	,  UpperCAmelCase=0 )  ->     List[Any]:
      if str(UpperCAmelCase ).startswith('''mps''' ):
       _lowercase							   =torch.manual_seed(UpperCAmelCase )
      else:
       _lowercase							   =torch.Generator(device=UpperCAmelCase ).manual_seed(UpperCAmelCase )
      _lowercase							   ={
          '''prompt''': '''a photograph of an astronaut riding a horse''',
          '''generator''': generator,
          '''num_inference_steps''': 1_0,
          '''guidance_scale''': 6.0,
          '''output_type''': '''numpy''',
          '''parallel''': 3,
          '''debug''': True,
      }
      return inputs





     def   __A      (self )  ->     Dict:
      _lowercase							   ='''cpu'''  # ensure determinism for the device-dependent torch.Generator
      _lowercase							   =self.get_dummy_components()
      _lowercase							   =StableDiffusionParadigmsPipeline(**UpperCAmelCase )
      _lowercase							   =sd_pipe.to(UpperCAmelCase )
      sd_pipe.set_progress_bar_config(disable=UpperCAmelCase )

      _lowercase							   =self.get_dummy_inputs(UpperCAmelCase )
      _lowercase							   =sd_pipe(**UpperCAmelCase ).images
      _lowercase							   =image[0, -3:, -3:, -1]
      assert image.shape == (1, 6_4, 6_4, 3)

      _lowercase							   =np.array([0.4773, 0.5417, 0.4723, 0.4925, 0.5631, 0.4752, 0.5240, 0.4935, 0.5023] )

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





     def   __A      (self )  ->     List[str]:
      _lowercase							   ='''cpu'''  # ensure determinism for the device-dependent torch.Generator
      _lowercase							   =self.get_dummy_components()
      torch.manual_seed(0 )
      _lowercase							   =DDPMParallelScheduler()
      torch.manual_seed(0 )
      _lowercase							   =StableDiffusionParadigmsPipeline(**UpperCAmelCase )
      _lowercase							   =sd_pipe.to(UpperCAmelCase )
      sd_pipe.set_progress_bar_config(disable=UpperCAmelCase )

      _lowercase							   =self.get_dummy_inputs(UpperCAmelCase )
      _lowercase							   =sd_pipe(**UpperCAmelCase ).images
      _lowercase							   =image[0, -3:, -3:, -1]
      assert image.shape == (1, 6_4, 6_4, 3)

      _lowercase							   =np.array([0.3573, 0.4420, 0.4960, 0.4799, 0.3796, 0.3879, 0.4819, 0.4365, 0.4468] )

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





     def   __A      (self )  ->     Optional[Any]:
      super().test_inference_batch_consistent(batch_sizes=[1, 2] )





     def   __A      (self )  ->     Tuple:
      super().test_inference_batch_single_identical(batch_size=2	,  expected_max_diff=3e-3 )





     def   __A      (self )  ->     List[Any]:
      _lowercase							   ='''cpu'''  # ensure determinism for the device-dependent torch.Generator
      _lowercase							   =self.get_dummy_components()
      _lowercase							   =StableDiffusionParadigmsPipeline(**UpperCAmelCase )
      _lowercase							   =sd_pipe.to(UpperCAmelCase )
      sd_pipe.set_progress_bar_config(disable=UpperCAmelCase )

      _lowercase							   =self.get_dummy_inputs(UpperCAmelCase )
      _lowercase							   ='''french fries'''
      _lowercase							   =sd_pipe(**UpperCAmelCase	,  negative_prompt=UpperCAmelCase )
      _lowercase							   =output.images
      _lowercase							   =image[0, -3:, -3:, -1]

      assert image.shape == (1, 6_4, 6_4, 3)

      _lowercase							   =np.array([0.4771, 0.5420, 0.4683, 0.4918, 0.5636, 0.4725, 0.5230, 0.4923, 0.5015] )

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







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





     def   __A      (self )  ->     List[str]:
      super().tearDown()
      gc.collect()
      torch.cuda.empty_cache()





     def   __A      (self	,  UpperCAmelCase=0 )  ->     Optional[Any]:
      _lowercase							   =torch.Generator(device=UpperCAmelCase ).manual_seed(UpperCAmelCase )
      _lowercase							   ={
          '''prompt''': '''a photograph of an astronaut riding a horse''',
          '''generator''': generator,
          '''num_inference_steps''': 1_0,
          '''guidance_scale''': 7.5,
          '''output_type''': '''numpy''',
          '''parallel''': 3,
          '''debug''': True,
      }
      return inputs





     def   __A      (self )  ->     List[str]:
      _lowercase							   ='''stabilityai/stable-diffusion-2-base'''
      _lowercase							   =DDIMParallelScheduler.from_pretrained(UpperCAmelCase	,  subfolder='''scheduler''' )
      _lowercase							   =StableDiffusionParadigmsPipeline.from_pretrained(UpperCAmelCase	,  scheduler=UpperCAmelCase	,  safety_checker=UpperCAmelCase )
      pipe.to(UpperCAmelCase )
      pipe.set_progress_bar_config(disable=UpperCAmelCase )
      pipe.enable_attention_slicing()

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

      assert image.shape == (1, 5_1_2, 5_1_2, 3)

      _lowercase							   =np.array([0.9622, 0.9602, 0.9748, 0.9591, 0.9630, 0.9691, 0.9661, 0.9631, 0.9741] )

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





import unittest
from functools import partial

import torch

from accelerate import Accelerator, debug_launcher
from accelerate.state import AcceleratorState, GradientState
from accelerate.test_utils import require_cpu, require_huggingface_suite
from accelerate.utils import GradientAccumulationPlugin
def        UpperCAmelCase_      (							__snake_case=2  ,		__snake_case=True  ,		__snake_case=False		)  ->    Dict:







 """simple docstring"""


 _lowercase							   =Accelerator(step_scheduler_with_optimizer=__snake_case  ,		split_batches=__snake_case		)
 _lowercase							   =torch.nn.Linear(2  ,		4		)
 _lowercase							   =torch.optim.AdamW(model.parameters()  ,		lr=1.0		)
 _lowercase							   =torch.optim.lr_scheduler.OneCycleLR(__snake_case  ,		max_lr=0.01  ,		steps_per_epoch=2  ,		epochs=1		)
 _lowercase   ,					_lowercase   ,					_lowercase							   =accelerator.prepare(__snake_case  ,		__snake_case  ,		__snake_case		)

 # Optimizer has stepped
 scheduler.step()
 if step_scheduler_with_optimizer or (num_processes == 1):
  assert (
      scheduler.scheduler.last_epoch == num_processes
  ), F"Last Epoch ({scheduler.scheduler.last_epoch}) != Num Processes ({num_processes})"
 else:
  assert (
      scheduler.scheduler.last_epoch != num_processes
  ), F"Last Epoch ({scheduler.scheduler.last_epoch}) == Num Processes ({num_processes})"
def        UpperCAmelCase_      (							__snake_case=2  ,		__snake_case=True  ,		__snake_case=False		)  ->    Optional[int]:







 """simple docstring"""


 _lowercase							   =Accelerator(step_scheduler_with_optimizer=__snake_case  ,		split_batches=__snake_case		)
 _lowercase							   =torch.nn.Linear(2  ,		4		)
 _lowercase							   =torch.optim.AdamW(model.parameters()  ,		lr=1.0		)
 _lowercase							   =torch.optim.lr_scheduler.LambdaLR(__snake_case  ,		lr_lambda=lambda __snake_case		: 1 - n / 10		)
 _lowercase   ,					_lowercase   ,					_lowercase							   =accelerator.prepare(__snake_case  ,		__snake_case  ,		__snake_case		)

 # Optimizer has stepped
 _lowercase							   =False
 scheduler.step()
 _lowercase							   =1 - (num_processes if (step_scheduler_with_optimizer and not split_batches) else 1) / 10
 assert (
     scheduler.get_last_lr()[0] == expected_lr
 ), F"Wrong lr found at first step, expected {expected_lr}, got {scheduler.get_last_lr()[0]}"

 # Optimizer has not stepped
 _lowercase							   =True
 scheduler.step()
 if not step_scheduler_with_optimizer:
  _lowercase							   =1 - 2 / 10
 assert (
     scheduler.get_last_lr()[0] == expected_lr
 ), F"Wrong lr found at second step, expected {expected_lr}, got {scheduler.get_last_lr()[0]}"
def        UpperCAmelCase_      (							__snake_case = 2		)  ->    Optional[int]:







 """simple docstring"""


 from transformers import get_linear_schedule_with_warmup

 _lowercase							   =[1, 2, 4]
 for num_steps in steps:
  _lowercase							   =GradientAccumulationPlugin(num_steps=__snake_case  ,		adjust_scheduler=num_steps > 1		)
  _lowercase							   =Accelerator(gradient_accumulation_plugin=__snake_case		)
  _lowercase							   =torch.nn.Linear(2  ,		4		)
  _lowercase							   =torch.optim.AdamW(model.parameters()  ,		lr=10.0		)
  _lowercase							   =get_linear_schedule_with_warmup(optimizer=__snake_case  ,		num_warmup_steps=0  ,		num_training_steps=20		)

  _lowercase   ,					_lowercase   ,					_lowercase							   =accelerator.prepare(__snake_case  ,		__snake_case  ,		__snake_case		)

  for i in range(10 * num_steps		):
   with accelerator.accumulate(__snake_case		):
    optimizer.step()
    scheduler.step()

   if i == (10 * num_steps - 2):
    assert (
        scheduler.get_last_lr()[0] != 0
    ), F"Wrong lr found at second-to-last step, expected non-zero, got {scheduler.get_last_lr()[0]}. num_steps: {num_steps}"
  assert (
      scheduler.get_last_lr()[0] == 0
  ), F"Wrong lr found at last step, expected 0, got {scheduler.get_last_lr()[0]}"
  GradientState._reset_state()







@require_cpu
class    lowerCamelCase__     ( unittest.TestCase):





     def   __A      (self )  ->     Optional[Any]:
      debug_launcher(partial(UpperCAmelCase	,  num_processes=1 )	,  num_processes=1 )
      debug_launcher(partial(UpperCAmelCase	,  num_processes=1	,  split_batches=UpperCAmelCase )	,  num_processes=1 )





     def   __A      (self )  ->     Any:
      debug_launcher(partial(UpperCAmelCase	,  num_processes=1 )	,  num_processes=1 )
      debug_launcher(partial(UpperCAmelCase	,  num_processes=1	,  split_batches=UpperCAmelCase )	,  num_processes=1 )





     def   __A      (self )  ->     Union[str, Any]:
      debug_launcher(partial(UpperCAmelCase	,  num_processes=1	,  step_scheduler_with_optimizer=UpperCAmelCase )	,  num_processes=1 )





     def   __A      (self )  ->     Dict:
      debug_launcher(partial(UpperCAmelCase	,  num_processes=1	,  step_scheduler_with_optimizer=UpperCAmelCase )	,  num_processes=1 )





     def   __A      (self )  ->     int:
      AcceleratorState._reset_state(UpperCAmelCase )
      debug_launcher(UpperCAmelCase )
      debug_launcher(partial(UpperCAmelCase	,  num_processes=1	,  split_batches=UpperCAmelCase )	,  num_processes=1 )





     def   __A      (self )  ->     Optional[Any]:
      AcceleratorState._reset_state(UpperCAmelCase )
      debug_launcher(UpperCAmelCase )
      debug_launcher(partial(UpperCAmelCase	,  num_processes=1	,  split_batches=UpperCAmelCase )	,  num_processes=1 )





     def   __A      (self )  ->     List[str]:
      AcceleratorState._reset_state(UpperCAmelCase )
      debug_launcher(partial(UpperCAmelCase	,  step_scheduler_with_optimizer=UpperCAmelCase ) )





     def   __A      (self )  ->     Union[str, Any]:
      AcceleratorState._reset_state(UpperCAmelCase )
      debug_launcher(partial(UpperCAmelCase	,  step_scheduler_with_optimizer=UpperCAmelCase ) )





     @require_huggingface_suite
     def   __A      (self )  ->     int:
      AcceleratorState._reset_state(UpperCAmelCase )
      debug_launcher(partial(UpperCAmelCase	,  num_processes=1 ) )
      debug_launcher(UpperCAmelCase )

