




import gc
import unittest

import numpy as np
import torch
from torch.backends.cuda import sdp_kernel

from diffusers import (
    CMStochasticIterativeScheduler,
    ConsistencyModelPipeline,
    UNetaDModel,
)
from diffusers.utils import randn_tensor, slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_a, require_torch_gpu

from ..pipeline_params import UNCONDITIONAL_IMAGE_GENERATION_BATCH_PARAMS, UNCONDITIONAL_IMAGE_GENERATION_PARAMS
from ..test_pipelines_common import PipelineTesterMixin


enable_full_determinism()







class 						lowercase							(   UpperCamelCase__,unittest.TestCase		):
     _a							 =						ConsistencyModelPipeline
     _a							 =						UNCONDITIONAL_IMAGE_GENERATION_PARAMS
     _a							 =						UNCONDITIONAL_IMAGE_GENERATION_BATCH_PARAMS

     # Override required_optional_params to remove num_images_per_prompt
     _a							 =						frozenset(
         [
             "num_inference_steps",
             "generator",
             "latents",
             "output_type",
             "return_dict",
             "callback",
             "callback_steps",
         ]		)

     @property
     def 					a__	(  self				)						->							Dict:
           _A							:     List[str]			=						UNetaDModel.from_pretrained(
               """diffusers/consistency-models-test"""		,     subfolder="""test_unet"""		,     )
           return unet

     @property
     def 					a__	(  self				)						->							List[str]:
           _A							:     Optional[Any]			=						UNetaDModel.from_pretrained(
               """diffusers/consistency-models-test"""		,     subfolder="""test_unet_class_cond"""		,     )
           return unet

     def 					a__	(  self		,     _a=False				)						->							List[Any]:
           if class_cond:
                 _A							:     List[str]			=						self.dummy_cond_unet
           else:
                 _A							:     List[str]			=						self.dummy_uncond_unet

           # Default to CM multistep sampler
           _A							:     int			=						CMStochasticIterativeScheduler(
               num_train_timesteps=40		,     sigma_min=0.002		,     sigma_max=80.0		,     )

           _A							:     List[Any]			=						{
               """unet""": unet,
               """scheduler""": scheduler,
           }

           return components

     def 					a__	(  self		,     _a		,     _a=0				)						->							Any:
           if str(_a				).startswith("""mps"""				):
                 _A							:     int			=						torch.manual_seed(_a				)
           else:
                 _A							:     Union[str, Any]			=						torch.Generator(device=_a				).manual_seed(_a				)

           _A							:     List[str]			=						{
               """batch_size""": 1,
               """num_inference_steps""": None,
               """timesteps""": [22, 0],
               """generator""": generator,
               """output_type""": """np""",
           }

           return inputs

     def 					a__	(  self				)						->							Any:
           _A							:     Any			=						"""cpu"""  # ensure determinism for the device-dependent torch.Generator
           _A							:     Union[str, Any]			=						self.get_dummy_components()
           _A							:     Any			=						ConsistencyModelPipeline(**_a				)
           _A							:     Any			=						pipe.to(_a				)
           pipe.set_progress_bar_config(disable=_a				)

           _A							:     Optional[int]			=						self.get_dummy_inputs(_a				)
           _A							:     Any			=						pipe(**_a				).images
           assert image.shape == (1, 32, 32, 3)

           _A							:     Union[str, Any]			=						image[0, -3:, -3:, -1]
           _A							:     Dict			=						np.array([0.3572, 0.6273, 0.4031, 0.3961, 0.4321, 0.5730, 0.5266, 0.4780, 0.5004]				)

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

     def 					a__	(  self				)						->							Dict:
           _A							:     List[Any]			=						"""cpu"""  # ensure determinism for the device-dependent torch.Generator
           _A							:     Any			=						self.get_dummy_components(class_cond=_a				)
           _A							:     str			=						ConsistencyModelPipeline(**_a				)
           _A							:     List[str]			=						pipe.to(_a				)
           pipe.set_progress_bar_config(disable=_a				)

           _A							:     List[str]			=						self.get_dummy_inputs(_a				)
           _A							:     Tuple			=						0
           _A							:     int			=						pipe(**_a				).images
           assert image.shape == (1, 32, 32, 3)

           _A							:     Union[str, Any]			=						image[0, -3:, -3:, -1]
           _A							:     List[Any]			=						np.array([0.3572, 0.6273, 0.4031, 0.3961, 0.4321, 0.5730, 0.5266, 0.4780, 0.5004]				)

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

     def 					a__	(  self				)						->							Any:
           _A							:     Optional[int]			=						"""cpu"""  # ensure determinism for the device-dependent torch.Generator
           _A							:     Tuple			=						self.get_dummy_components()
           _A							:     Dict			=						ConsistencyModelPipeline(**_a				)
           _A							:     Optional[Any]			=						pipe.to(_a				)
           pipe.set_progress_bar_config(disable=_a				)

           _A							:     Union[str, Any]			=						self.get_dummy_inputs(_a				)
           _A							:     Optional[Any]			=						1
           _A							:     Optional[Any]			=						None
           _A							:     Dict			=						pipe(**_a				).images
           assert image.shape == (1, 32, 32, 3)

           _A							:     Optional[int]			=						image[0, -3:, -3:, -1]
           _A							:     Optional[Any]			=						np.array([0.5004, 0.5004, 0.4994, 0.5008, 0.4976, 0.5018, 0.4990, 0.4982, 0.4987]				)

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



     def 					a__	(  self				)						->							List[str]:
           _A							:     Tuple			=						"""cpu"""  # ensure determinism for the device-dependent torch.Generator
           _A							:     Any			=						self.get_dummy_components(class_cond=_a				)
           _A							:     str			=						ConsistencyModelPipeline(**_a				)
           _A							:     List[str]			=						pipe.to(_a				)
           pipe.set_progress_bar_config(disable=_a				)

           _A							:     Optional[int]			=						self.get_dummy_inputs(_a				)
           _A							:     Optional[int]			=						1
           _A							:     Any			=						None
           _A							:     str			=						0
           _A							:     Dict			=						pipe(**_a				).images
           assert image.shape == (1, 32, 32, 3)

           _A							:     str			=						image[0, -3:, -3:, -1]
           _A							:     str			=						np.array([0.5004, 0.5004, 0.4994, 0.5008, 0.4976, 0.5018, 0.4990, 0.4982, 0.4987]				)

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







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

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

     def 					a__	(  self		,     _a=0		,     _a=False		,     _a="cpu"		,     _a=torch.floataa		,     _a=(1, 3, 64, 64)				)						->							Optional[int]:
           _A							:     Optional[int]			=						torch.manual_seed(_a				)

           _A							:     List[Any]			=						{
               """num_inference_steps""": None,
               """timesteps""": [22, 0],
               """class_labels""": 0,
               """generator""": generator,
               """output_type""": """np""",
           }

           if get_fixed_latents:
                 _A							:     Tuple			=						self.get_fixed_latents(seed=_a		,     device=_a		,     dtype=_a		,     shape=_a				)
                 _A							:     Any			=						latents

           return inputs

     def 					a__	(  self		,     _a=0		,     _a="cpu"		,     _a=torch.floataa		,     _a=(1, 3, 64, 64)				)						->							Optional[Any]:
           if type(_a				) == str:
                 _A							:     int			=						torch.device(_a				)
           _A							:     Optional[int]			=						torch.Generator(device=_a				).manual_seed(_a				)
           _A							:     Tuple			=						randn_tensor(_a		,     generator=_a		,     device=_a		,     dtype=_a				)
           return latents

     def 					a__	(  self				)						->							Any:
           _A							:     List[str]			=						UNetaDModel.from_pretrained("""diffusers/consistency_models"""		,     subfolder="""diffusers_cd_imagenet64_l2"""				)
           _A							:     Optional[Any]			=						CMStochasticIterativeScheduler(
               num_train_timesteps=40		,     sigma_min=0.002		,     sigma_max=80.0		,     )
           _A							:     Dict			=						ConsistencyModelPipeline(unet=_a		,     scheduler=_a				)
           pipe.to(torch_device=_a				)
           pipe.set_progress_bar_config(disable=_a				)

           _A							:     List[Any]			=						self.get_inputs()
           _A							:     Optional[Any]			=						pipe(**_a				).images
           assert image.shape == (1, 64, 64, 3)

           _A							:     int			=						image[0, -3:, -3:, -1]

           _A							:     Dict			=						np.array([0.0888, 0.0881, 0.0666, 0.0479, 0.0292, 0.0195, 0.0201, 0.0163, 0.0254]				)

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

     def 					a__	(  self				)						->							Optional[int]:
           _A							:     Any			=						UNetaDModel.from_pretrained("""diffusers/consistency_models"""		,     subfolder="""diffusers_cd_imagenet64_l2"""				)
           _A							:     Tuple			=						CMStochasticIterativeScheduler(
               num_train_timesteps=40		,     sigma_min=0.002		,     sigma_max=80.0		,     )
           _A							:     Dict			=						ConsistencyModelPipeline(unet=_a		,     scheduler=_a				)
           pipe.to(torch_device=_a				)
           pipe.set_progress_bar_config(disable=_a				)

           _A							:     Union[str, Any]			=						self.get_inputs()
           _A							:     Optional[Any]			=						1
           _A							:     str			=						None
           _A							:     int			=						pipe(**_a				).images
           assert image.shape == (1, 64, 64, 3)

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

           _A							:     Tuple			=						np.array([0.0340, 0.0152, 0.0063, 0.0267, 0.0221, 0.0107, 0.0416, 0.0186, 0.0217]				)

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

     @require_torch_a
     def 					a__	(  self				)						->							Optional[Any]:
           _A							:     Any			=						UNetaDModel.from_pretrained("""diffusers/consistency_models"""		,     subfolder="""diffusers_cd_imagenet64_l2"""				)
           _A							:     Optional[int]			=						CMStochasticIterativeScheduler(
               num_train_timesteps=40		,     sigma_min=0.002		,     sigma_max=80.0		,     )
           _A							:     Dict			=						ConsistencyModelPipeline(unet=_a		,     scheduler=_a				)
           pipe.to(torch_device=_a		,     torch_dtype=torch.floataa				)
           pipe.set_progress_bar_config(disable=_a				)

           _A							:     int			=						self.get_inputs(get_fixed_latents=_a		,     device=_a				)
           # Ensure usage of flash attention in torch 2.0
           with sdp_kernel(enable_flash=_a		,     enable_math=_a		,     enable_mem_efficient=_a				):
                 _A							:     Union[str, Any]			=						pipe(**_a				).images
           assert image.shape == (1, 64, 64, 3)

           _A							:     Optional[Any]			=						image[0, -3:, -3:, -1]

           _A							:     Dict			=						np.array([0.1875, 0.1428, 0.1289, 0.2151, 0.2092, 0.1477, 0.1877, 0.1641, 0.1353]				)

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



     @require_torch_a
     def 					a__	(  self				)						->							Optional[int]:
           _A							:     int			=						UNetaDModel.from_pretrained("""diffusers/consistency_models"""		,     subfolder="""diffusers_cd_imagenet64_l2"""				)
           _A							:     Any			=						CMStochasticIterativeScheduler(
               num_train_timesteps=40		,     sigma_min=0.002		,     sigma_max=80.0		,     )
           _A							:     str			=						ConsistencyModelPipeline(unet=_a		,     scheduler=_a				)
           pipe.to(torch_device=_a		,     torch_dtype=torch.floataa				)
           pipe.set_progress_bar_config(disable=_a				)

           _A							:     Any			=						self.get_inputs(get_fixed_latents=_a		,     device=_a				)
           _A							:     Dict			=						1
           _A							:     Union[str, Any]			=						None
           # Ensure usage of flash attention in torch 2.0
           with sdp_kernel(enable_flash=_a		,     enable_math=_a		,     enable_mem_efficient=_a				):
                 _A							:     int			=						pipe(**_a				).images
           assert image.shape == (1, 64, 64, 3)

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

           _A							:     Dict			=						np.array([0.1663, 0.1948, 0.2275, 0.1680, 0.1204, 0.1245, 0.1858, 0.1338, 0.2095]				)

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







from __future__ import annotations

import unittest

from transformers import is_tf_available
from transformers.testing_utils import require_sentencepiece, require_tf, require_tokenizers, slow


if is_tf_available():
  import numpy as np
  import tensorflow as tf

  from transformers import TFCamembertModel







@require_tf
@require_sentencepiece
@require_tokenizers
class 						lowercase							(   unittest.TestCase		):

     @slow
     def 					a__	(  self				)						->							Any:
           _A							:     Tuple			=						TFCamembertModel.from_pretrained("""jplu/tf-camembert-base"""				)

           _A							:     List[Any]			=						tf.convert_to_tensor(
               [[5, 121, 11, 660, 16, 730, 2_5543, 110, 83, 6]]		,     dtype=tf.intaa		,     )  # J'aime le camembert !"

           _A							:     List[str]			=						model(_a				)["""last_hidden_state"""]
           _A							:     Union[str, Any]			=						tf.TensorShape((1, 10, 768)				)
           self.assertEqual(output.shape		,     _a				)
           # compare the actual values for a slice.
           _A							:     List[Any]			=						tf.convert_to_tensor(
               [[[-0.0254, 0.0235, 0.1027], [0.0606, -0.1811, -0.0418], [-0.1561, -0.1127, 0.2687]]]		,     dtype=tf.floataa		,     )
           # camembert = torch.hub.load('pytorch/fairseq', 'camembert.v0')
           # camembert.eval()
           # expected_slice = roberta.model.forward(input_ids)[0][:, :3, :3].detach()

           self.assertTrue(np.allclose(output[:, :3, :3].numpy()		,     expected_slice.numpy()		,     atol=1e-4				)				)

