





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    __snake_case	(							__lowerCamelCase       ,      unittest.TestCase      ):





       '''simple docstring'''







       lowerCAmelCase__						=					ConsistencyModelPipeline
       lowerCAmelCase__						=					UNCONDITIONAL_IMAGE_GENERATION_PARAMS
       lowerCAmelCase__						=					UNCONDITIONAL_IMAGE_GENERATION_BATCH_PARAMS

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




       @property
       def 							UpperCAmelCase__			(  self			:   Union[str, Any]		):
         __snake_case:						Union[str, Any]			  =						UNetaDModel.from_pretrained(
             """diffusers/consistency-models-test"""							,    subfolder="""test_unet"""							,    )
         return unet




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




       def 							UpperCAmelCase__			(  self			:   List[str]							,    A			:   Optional[Any]=False		):
         if class_cond:
           __snake_case:						Dict			  =						self.dummy_cond_unet
         else:
           __snake_case:						int			  =						self.dummy_uncond_unet

         # Default to CM multistep sampler
         __snake_case:						Optional[int]			  =						CMStochasticIterativeScheduler(
             num_train_timesteps=40							,    sigma_min=0.002							,    sigma_max=80.0							,    )

         __snake_case:						Tuple			  =						{
             """unet""": unet,
             """scheduler""": scheduler,
         }

         return components




       def 							UpperCAmelCase__			(  self			:   Union[str, Any]							,    A			:   Union[str, Any]							,    A			:   Tuple=0		):
         if str(A		).startswith("""mps"""		):
           __snake_case:						Dict			  =						torch.manual_seed(A		)
         else:
           __snake_case:						Optional[Any]			  =						torch.Generator(device=A		).manual_seed(A		)

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

         return inputs




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

         __snake_case:						Any			  =						self.get_dummy_inputs(A		)
         __snake_case:						List[Any]			  =						pipe(**A		).images
         assert image.shape == (1, 32, 32, 3)

         __snake_case:						Optional[Any]			  =						image[0, -3:, -3:, -1]
         __snake_case:						List[str]			  =						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 							UpperCAmelCase__			(  self			:   Dict		):
         __snake_case:						Optional[int]			  =						"""cpu"""  # ensure determinism for the device-dependent torch.Generator
         __snake_case:						int			  =						self.get_dummy_components(class_cond=A		)
         __snake_case:						Optional[Any]			  =						ConsistencyModelPipeline(**A		)
         __snake_case:						Tuple			  =						pipe.to(A		)
         pipe.set_progress_bar_config(disable=A		)

         __snake_case:						Union[str, Any]			  =						self.get_dummy_inputs(A		)
         __snake_case:						Any			  =						0
         __snake_case:						Tuple			  =						pipe(**A		).images
         assert image.shape == (1, 32, 32, 3)

         __snake_case:						Any			  =						image[0, -3:, -3:, -1]
         __snake_case:						Optional[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 							UpperCAmelCase__			(  self			:   Any		):
         __snake_case:						Union[str, Any]			  =						"""cpu"""  # ensure determinism for the device-dependent torch.Generator
         __snake_case:						str			  =						self.get_dummy_components()
         __snake_case:						List[str]			  =						ConsistencyModelPipeline(**A		)
         __snake_case:						Dict			  =						pipe.to(A		)
         pipe.set_progress_bar_config(disable=A		)

         __snake_case:						Any			  =						self.get_dummy_inputs(A		)
         __snake_case:						Optional[int]			  =						1
         __snake_case:						List[Any]			  =						None
         __snake_case:						List[str]			  =						pipe(**A		).images
         assert image.shape == (1, 32, 32, 3)

         __snake_case:						Optional[Any]			  =						image[0, -3:, -3:, -1]
         __snake_case:						int			  =						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 							UpperCAmelCase__			(  self			:   List[str]		):
         __snake_case:						Tuple			  =						"""cpu"""  # ensure determinism for the device-dependent torch.Generator
         __snake_case:						Any			  =						self.get_dummy_components(class_cond=A		)
         __snake_case:						Optional[int]			  =						ConsistencyModelPipeline(**A		)
         __snake_case:						int			  =						pipe.to(A		)
         pipe.set_progress_bar_config(disable=A		)

         __snake_case:						Tuple			  =						self.get_dummy_inputs(A		)
         __snake_case:						Any			  =						1
         __snake_case:						Tuple			  =						None
         __snake_case:						Tuple			  =						0
         __snake_case:						str			  =						pipe(**A		).images
         assert image.shape == (1, 32, 32, 3)

         __snake_case:						str			  =						image[0, -3:, -3:, -1]
         __snake_case:						Tuple			  =						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    __snake_case	(							unittest.TestCase      ):





       '''simple docstring'''







       def 							UpperCAmelCase__			(  self			:   Dict		):
         super().tearDown()
         gc.collect()
         torch.cuda.empty_cache()




       def 							UpperCAmelCase__			(  self			:   str							,    A			:   Union[str, Any]=0							,    A			:   Any=False							,    A			:   List[Any]="cpu"							,    A			:   Any=torch.floataa							,    A			:   str=(1, 3, 64, 64)		):
         __snake_case:						Dict			  =						torch.manual_seed(A		)

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

         if get_fixed_latents:
           __snake_case:						Optional[int]			  =						self.get_fixed_latents(seed=A							,    device=A							,    dtype=A							,    shape=A		)
           __snake_case:						Optional[int]			  =						latents

         return inputs




       def 							UpperCAmelCase__			(  self			:   Dict							,    A			:   Dict=0							,    A			:   List[Any]="cpu"							,    A			:   Any=torch.floataa							,    A			:   int=(1, 3, 64, 64)		):
         if type(A		) == str:
           __snake_case:						List[Any]			  =						torch.device(A		)
         __snake_case:						Optional[int]			  =						torch.Generator(device=A		).manual_seed(A		)
         __snake_case:						List[str]			  =						randn_tensor(A							,    generator=A							,    device=A							,    dtype=A		)
         return latents




       def 							UpperCAmelCase__			(  self			:   Optional[int]		):
         __snake_case:						Optional[Any]			  =						UNetaDModel.from_pretrained("""diffusers/consistency_models"""							,    subfolder="""diffusers_cd_imagenet64_l2"""		)
         __snake_case:						str			  =						CMStochasticIterativeScheduler(
             num_train_timesteps=40							,    sigma_min=0.002							,    sigma_max=80.0							,    )
         __snake_case:						List[str]			  =						ConsistencyModelPipeline(unet=A							,    scheduler=A		)
         pipe.to(torch_device=A		)
         pipe.set_progress_bar_config(disable=A		)

         __snake_case:						List[str]			  =						self.get_inputs()
         __snake_case:						Dict			  =						pipe(**A		).images
         assert image.shape == (1, 64, 64, 3)

         __snake_case:						List[Any]			  =						image[0, -3:, -3:, -1]

         __snake_case:						Union[str, Any]			  =						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 							UpperCAmelCase__			(  self			:   Tuple		):
         __snake_case:						Dict			  =						UNetaDModel.from_pretrained("""diffusers/consistency_models"""							,    subfolder="""diffusers_cd_imagenet64_l2"""		)
         __snake_case:						int			  =						CMStochasticIterativeScheduler(
             num_train_timesteps=40							,    sigma_min=0.002							,    sigma_max=80.0							,    )
         __snake_case:						Dict			  =						ConsistencyModelPipeline(unet=A							,    scheduler=A		)
         pipe.to(torch_device=A		)
         pipe.set_progress_bar_config(disable=A		)

         __snake_case:						Dict			  =						self.get_inputs()
         __snake_case:						Optional[Any]			  =						1
         __snake_case:						Any			  =						None
         __snake_case:						List[str]			  =						pipe(**A		).images
         assert image.shape == (1, 64, 64, 3)

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

         __snake_case:						Union[str, Any]			  =						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 							UpperCAmelCase__			(  self			:   Dict		):
         __snake_case:						List[Any]			  =						UNetaDModel.from_pretrained("""diffusers/consistency_models"""							,    subfolder="""diffusers_cd_imagenet64_l2"""		)
         __snake_case:						str			  =						CMStochasticIterativeScheduler(
             num_train_timesteps=40							,    sigma_min=0.002							,    sigma_max=80.0							,    )
         __snake_case:						Optional[Any]			  =						ConsistencyModelPipeline(unet=A							,    scheduler=A		)
         pipe.to(torch_device=A							,    torch_dtype=torch.floataa		)
         pipe.set_progress_bar_config(disable=A		)

         __snake_case:						Optional[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		):
           __snake_case:						Dict			  =						pipe(**A		).images
         assert image.shape == (1, 64, 64, 3)

         __snake_case:						List[Any]			  =						image[0, -3:, -3:, -1]

         __snake_case:						int			  =						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 							UpperCAmelCase__			(  self			:   Union[str, Any]		):
         __snake_case:						List[str]			  =						UNetaDModel.from_pretrained("""diffusers/consistency_models"""							,    subfolder="""diffusers_cd_imagenet64_l2"""		)
         __snake_case:						Any			  =						CMStochasticIterativeScheduler(
             num_train_timesteps=40							,    sigma_min=0.002							,    sigma_max=80.0							,    )
         __snake_case:						Tuple			  =						ConsistencyModelPipeline(unet=A							,    scheduler=A		)
         pipe.to(torch_device=A							,    torch_dtype=torch.floataa		)
         pipe.set_progress_bar_config(disable=A		)

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

         __snake_case:						Any			  =						image[0, -3:, -3:, -1]

         __snake_case:						str			  =						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













import math
import time

from transformers import Trainer, is_torch_tpu_available
from transformers.trainer_utils import PredictionOutput, speed_metrics


if is_torch_tpu_available(check_device=False):
      import torch_xla.core.xla_model as xm
      import torch_xla.debug.metrics as met







class    __snake_case	(							__lowerCamelCase      ):





       '''simple docstring'''







       def __init__(  self			:   str							,    *A			:   Dict							,    A			:   Optional[int]=None							,    A			:   Tuple=None							,    **A			:   Optional[int]		):
         super().__init__(*A							,    **A		)
         __snake_case:						List[Any]			  =						eval_examples
         __snake_case:						str			  =						post_process_function




       def 							UpperCAmelCase__			(  self			:   List[Any]							,    A			:   Dict=None							,    A			:   int=None							,    A			:   List[Any]=None							,    A			:   str = "eval"		):
         __snake_case:						int			  =						self.eval_dataset if eval_dataset is None else eval_dataset
         __snake_case:						Any			  =						self.get_eval_dataloader(A		)
         __snake_case:						Optional[Any]			  =						self.eval_examples if eval_examples is None else eval_examples

         # Temporarily disable metric computation, we will do it in the loop here.
         __snake_case:						Union[str, Any]			  =						self.compute_metrics
         __snake_case:						List[str]			  =						None
         __snake_case:						Tuple			  =						self.prediction_loop if self.args.use_legacy_prediction_loop else self.evaluation_loop
         __snake_case:						Tuple			  =						time.time()
         try:
           __snake_case:						Any			  =						eval_loop(
               A							,    description="""Evaluation"""							,    prediction_loss_only=True if compute_metrics is None else None							,    ignore_keys=A							,    metric_key_prefix=A							,    )
         finally:
           __snake_case:						Optional[int]			  =						compute_metrics
         __snake_case:						Union[str, Any]			  =						self.args.eval_batch_size * self.args.world_size
         if f'''{metric_key_prefix}_jit_compilation_time''' in output.metrics:
           start_time += output.metrics[f'''{metric_key_prefix}_jit_compilation_time''']
         output.metrics.update(
             speed_metrics(
                 A							,    A							,    num_samples=output.num_samples							,    num_steps=math.ceil(output.num_samples / total_batch_size		)							,    )		)
         if self.post_process_function is not None and self.compute_metrics is not None and self.args.should_save:
           # Only the main node write the results by default
           __snake_case:						List[str]			  =						self.post_process_function(A							,    A							,    output.predictions		)
           __snake_case:						List[Any]			  =						self.compute_metrics(A		)

           # Prefix all keys with metric_key_prefix + '_'
           for key in list(metrics.keys()		):
             if not key.startswith(f'''{metric_key_prefix}_'''		):
               __snake_case:						str			  =						metrics.pop(A		)
           metrics.update(output.metrics		)
         else:
           __snake_case:						List[Any]			  =						output.metrics

         if self.args.should_log:
           # Only the main node log the results by default
           self.log(A		)

         if self.args.tpu_metrics_debug or self.args.debug:
           # tpu-comment: Logging debug metrics for PyTorch/XLA (compile, execute times, ops, etc.)
           xm.master_print(met.metrics_report()		)

         __snake_case:						str			  =						self.callback_handler.on_evaluate(self.args							,    self.state							,    self.control							,    A		)
         return metrics






       def 							UpperCAmelCase__			(  self			:   Optional[Any]							,    A			:   List[Any]							,    A			:   List[str]							,    A			:   str=None							,    A			:   str = "test"		):
         __snake_case:						Optional[Any]			  =						self.get_test_dataloader(A		)

         # Temporarily disable metric computation, we will do it in the loop here.
         __snake_case:						Optional[int]			  =						self.compute_metrics
         __snake_case:						List[Any]			  =						None
         __snake_case:						str			  =						self.prediction_loop if self.args.use_legacy_prediction_loop else self.evaluation_loop
         __snake_case:						Dict			  =						time.time()
         try:
           __snake_case:						str			  =						eval_loop(
               A							,    description="""Prediction"""							,    prediction_loss_only=True if compute_metrics is None else None							,    ignore_keys=A							,    metric_key_prefix=A							,    )
         finally:
           __snake_case:						List[Any]			  =						compute_metrics
         __snake_case:						Dict			  =						self.args.eval_batch_size * self.args.world_size
         if f'''{metric_key_prefix}_jit_compilation_time''' in output.metrics:
           start_time += output.metrics[f'''{metric_key_prefix}_jit_compilation_time''']
         output.metrics.update(
             speed_metrics(
                 A							,    A							,    num_samples=output.num_samples							,    num_steps=math.ceil(output.num_samples / total_batch_size		)							,    )		)

         if self.post_process_function is None or self.compute_metrics is None:
           return output

         __snake_case:						Union[str, Any]			  =						self.post_process_function(A							,    A							,    output.predictions							,    """predict"""		)
         __snake_case:						str			  =						self.compute_metrics(A		)

         # Prefix all keys with metric_key_prefix + '_'
         for key in list(metrics.keys()		):
           if not key.startswith(f'''{metric_key_prefix}_'''		):
             __snake_case:						List[str]			  =						metrics.pop(A		)
         metrics.update(output.metrics		)
         return PredictionOutput(predictions=predictions.predictions							,    label_ids=predictions.label_ids							,    metrics=A		)






