





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 				__lowerCAmelCase       (    UpperCamelCase__						,				unittest.TestCase):
      _lowercase				:					str						=			ConsistencyModelPipeline
      _lowercase				:					Optional[Any]						=			UNCONDITIONAL_IMAGE_GENERATION_PARAMS
      _lowercase				:					Dict						=			UNCONDITIONAL_IMAGE_GENERATION_BATCH_PARAMS

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







      @property
      def    _lowercase  (    self							)   ->						Optional[int]:



        '''simple docstring'''




        a__	:					str										=UNetaDModel.from_pretrained(
            "diffusers/consistency-models-test"		,      subfolder="test_unet"		,      )
        return unet







      @property
      def    _lowercase  (    self							)   ->						Dict:



        '''simple docstring'''




        a__	:					List[str]										=UNetaDModel.from_pretrained(
            "diffusers/consistency-models-test"		,      subfolder="test_unet_class_cond"		,      )
        return unet







      def    _lowercase  (    self		,      lowerCAmelCase__=False							)   ->						List[str]:



        '''simple docstring'''




        if class_cond:
          a__	:					int										=self.dummy_cond_unet
        else:
          a__	:					Union[str, Any]										=self.dummy_uncond_unet

        # Default to CM multistep sampler
        a__	:					int										=CMStochasticIterativeScheduler(
            num_train_timesteps=4_0		,      sigma_min=0.0_02		,      sigma_max=80.0		,      )

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

        return components







      def    _lowercase  (    self		,      lowerCAmelCase__		,      lowerCAmelCase__=0							)   ->						Any:



        '''simple docstring'''




        if str(lowerCAmelCase__							).startswith("mps"							):
          a__	:					Optional[Any]										=torch.manual_seed(lowerCAmelCase__							)
        else:
          a__	:					Dict										=torch.Generator(device=lowerCAmelCase__							).manual_seed(lowerCAmelCase__							)

        a__	:					Optional[int]										={
            "batch_size": 1,
            "num_inference_steps": None,
            "timesteps": [2_2, 0],
            "generator": generator,
            "output_type": "np",
        }

        return inputs







      def    _lowercase  (    self							)   ->						int:



        '''simple docstring'''




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

        a__	:					Any										=self.get_dummy_inputs(lowerCAmelCase__							)
        a__	:					List[Any]										=pipe(**lowerCAmelCase__							).images
        assert image.shape == (1, 3_2, 3_2, 3)

        a__	:					Union[str, Any]										=image[0, -3:, -3:, -1]
        a__	:					Any										=np.array([0.35_72, 0.62_73, 0.40_31, 0.39_61, 0.43_21, 0.57_30, 0.52_66, 0.47_80, 0.50_04]							)

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







      def    _lowercase  (    self							)   ->						List[Any]:



        '''simple docstring'''




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

        a__	:					List[Any]										=self.get_dummy_inputs(lowerCAmelCase__							)
        a__	:					Any										=0
        a__	:					Optional[int]										=pipe(**lowerCAmelCase__							).images
        assert image.shape == (1, 3_2, 3_2, 3)

        a__	:					int										=image[0, -3:, -3:, -1]
        a__	:					Any										=np.array([0.35_72, 0.62_73, 0.40_31, 0.39_61, 0.43_21, 0.57_30, 0.52_66, 0.47_80, 0.50_04]							)

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







      def    _lowercase  (    self							)   ->						Optional[int]:



        '''simple docstring'''




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

        a__	:					str										=self.get_dummy_inputs(lowerCAmelCase__							)
        a__	:					Union[str, Any]										=1
        a__	:					List[Any]										=None
        a__	:					List[str]										=pipe(**lowerCAmelCase__							).images
        assert image.shape == (1, 3_2, 3_2, 3)

        a__	:					Dict										=image[0, -3:, -3:, -1]
        a__	:					Any										=np.array([0.50_04, 0.50_04, 0.49_94, 0.50_08, 0.49_76, 0.50_18, 0.49_90, 0.49_82, 0.49_87]							)

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







      def    _lowercase  (    self							)   ->						Any:



        '''simple docstring'''




        a__	:					List[Any]										="cpu"  # ensure determinism for the device-dependent torch.Generator
        a__	:					Optional[int]										=self.get_dummy_components(class_cond=lowerCAmelCase__							)
        a__	:					Tuple										=ConsistencyModelPipeline(**lowerCAmelCase__							)
        a__	:					Union[str, Any]										=pipe.to(lowerCAmelCase__							)
        pipe.set_progress_bar_config(disable=lowerCAmelCase__							)

        a__	:					Union[str, Any]										=self.get_dummy_inputs(lowerCAmelCase__							)
        a__	:					Dict										=1
        a__	:					List[Any]										=None
        a__	:					int										=0
        a__	:					List[str]										=pipe(**lowerCAmelCase__							).images
        assert image.shape == (1, 3_2, 3_2, 3)

        a__	:					int										=image[0, -3:, -3:, -1]
        a__	:					Optional[int]										=np.array([0.50_04, 0.50_04, 0.49_94, 0.50_08, 0.49_76, 0.50_18, 0.49_90, 0.49_82, 0.49_87]							)

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







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







      def    _lowercase  (    self							)   ->						int:



        '''simple docstring'''




        super().tearDown()
        gc.collect()
        torch.cuda.empty_cache()







      def    _lowercase  (    self		,      lowerCAmelCase__=0		,      lowerCAmelCase__=False		,      lowerCAmelCase__="cpu"		,      lowerCAmelCase__=torch.floataa		,      lowerCAmelCase__=(1, 3, 6_4, 6_4)							)   ->						Union[str, Any]:



        '''simple docstring'''




        a__	:					Any										=torch.manual_seed(lowerCAmelCase__							)

        a__	:					Dict										={
            "num_inference_steps": None,
            "timesteps": [2_2, 0],
            "class_labels": 0,
            "generator": generator,
            "output_type": "np",
        }

        if get_fixed_latents:
          a__	:					Union[str, Any]										=self.get_fixed_latents(seed=lowerCAmelCase__		,      device=lowerCAmelCase__		,      dtype=lowerCAmelCase__		,      shape=lowerCAmelCase__							)
          a__	:					Any										=latents

        return inputs







      def    _lowercase  (    self		,      lowerCAmelCase__=0		,      lowerCAmelCase__="cpu"		,      lowerCAmelCase__=torch.floataa		,      lowerCAmelCase__=(1, 3, 6_4, 6_4)							)   ->						List[str]:



        '''simple docstring'''




        if type(lowerCAmelCase__							) == str:
          a__	:					List[str]										=torch.device(lowerCAmelCase__							)
        a__	:					str										=torch.Generator(device=lowerCAmelCase__							).manual_seed(lowerCAmelCase__							)
        a__	:					Any										=randn_tensor(lowerCAmelCase__		,      generator=lowerCAmelCase__		,      device=lowerCAmelCase__		,      dtype=lowerCAmelCase__							)
        return latents







      def    _lowercase  (    self							)   ->						Any:



        '''simple docstring'''




        a__	:					List[Any]										=UNetaDModel.from_pretrained("diffusers/consistency_models"		,      subfolder="diffusers_cd_imagenet64_l2"							)
        a__	:					Union[str, Any]										=CMStochasticIterativeScheduler(
            num_train_timesteps=4_0		,      sigma_min=0.0_02		,      sigma_max=80.0		,      )
        a__	:					int										=ConsistencyModelPipeline(unet=lowerCAmelCase__		,      scheduler=lowerCAmelCase__							)
        pipe.to(torch_device=lowerCAmelCase__							)
        pipe.set_progress_bar_config(disable=lowerCAmelCase__							)

        a__	:					Dict										=self.get_inputs()
        a__	:					List[str]										=pipe(**lowerCAmelCase__							).images
        assert image.shape == (1, 6_4, 6_4, 3)

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

        a__	:					Optional[Any]										=np.array([0.08_88, 0.08_81, 0.06_66, 0.04_79, 0.02_92, 0.01_95, 0.02_01, 0.01_63, 0.02_54]							)

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







      def    _lowercase  (    self							)   ->						Any:



        '''simple docstring'''




        a__	:					List[str]										=UNetaDModel.from_pretrained("diffusers/consistency_models"		,      subfolder="diffusers_cd_imagenet64_l2"							)
        a__	:					str										=CMStochasticIterativeScheduler(
            num_train_timesteps=4_0		,      sigma_min=0.0_02		,      sigma_max=80.0		,      )
        a__	:					List[Any]										=ConsistencyModelPipeline(unet=lowerCAmelCase__		,      scheduler=lowerCAmelCase__							)
        pipe.to(torch_device=lowerCAmelCase__							)
        pipe.set_progress_bar_config(disable=lowerCAmelCase__							)

        a__	:					int										=self.get_inputs()
        a__	:					Optional[Any]										=1
        a__	:					Tuple										=None
        a__	:					List[str]										=pipe(**lowerCAmelCase__							).images
        assert image.shape == (1, 6_4, 6_4, 3)

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

        a__	:					List[Any]										=np.array([0.03_40, 0.01_52, 0.00_63, 0.02_67, 0.02_21, 0.01_07, 0.04_16, 0.01_86, 0.02_17]							)

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







      @require_torch_a
      def    _lowercase  (    self							)   ->						List[Any]:



        '''simple docstring'''




        a__	:					str										=UNetaDModel.from_pretrained("diffusers/consistency_models"		,      subfolder="diffusers_cd_imagenet64_l2"							)
        a__	:					List[Any]										=CMStochasticIterativeScheduler(
            num_train_timesteps=4_0		,      sigma_min=0.0_02		,      sigma_max=80.0		,      )
        a__	:					Union[str, Any]										=ConsistencyModelPipeline(unet=lowerCAmelCase__		,      scheduler=lowerCAmelCase__							)
        pipe.to(torch_device=lowerCAmelCase__		,      torch_dtype=torch.floataa							)
        pipe.set_progress_bar_config(disable=lowerCAmelCase__							)

        a__	:					List[str]										=self.get_inputs(get_fixed_latents=lowerCAmelCase__		,      device=lowerCAmelCase__							)
        # Ensure usage of flash attention in torch 2.0
        with sdp_kernel(enable_flash=lowerCAmelCase__		,      enable_math=lowerCAmelCase__		,      enable_mem_efficient=lowerCAmelCase__							):
          a__	:					Optional[int]										=pipe(**lowerCAmelCase__							).images
        assert image.shape == (1, 6_4, 6_4, 3)

        a__	:					str										=image[0, -3:, -3:, -1]

        a__	:					Any										=np.array([0.18_75, 0.14_28, 0.12_89, 0.21_51, 0.20_92, 0.14_77, 0.18_77, 0.16_41, 0.13_53]							)

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







      @require_torch_a
      def    _lowercase  (    self							)   ->						Any:



        '''simple docstring'''




        a__	:					str										=UNetaDModel.from_pretrained("diffusers/consistency_models"		,      subfolder="diffusers_cd_imagenet64_l2"							)
        a__	:					Dict										=CMStochasticIterativeScheduler(
            num_train_timesteps=4_0		,      sigma_min=0.0_02		,      sigma_max=80.0		,      )
        a__	:					Any										=ConsistencyModelPipeline(unet=lowerCAmelCase__		,      scheduler=lowerCAmelCase__							)
        pipe.to(torch_device=lowerCAmelCase__		,      torch_dtype=torch.floataa							)
        pipe.set_progress_bar_config(disable=lowerCAmelCase__							)

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

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

        a__	:					int										=np.array([0.16_63, 0.19_48, 0.22_75, 0.16_80, 0.12_04, 0.12_45, 0.18_58, 0.13_38, 0.20_95]							)

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













from ...configuration_utils import PretrainedConfig
from ...utils import logging


UpperCAmelCase	:	Dict		     =       logging.get_logger(__name__)

UpperCAmelCase	:	Tuple		     =       {
    """caidas/swin2sr-classicalsr-x2-64""": (
        """https://huggingface.co/caidas/swin2sr-classicalsr-x2-64/resolve/main/config.json"""
    ),
}







class 				__lowerCAmelCase       (    UpperCamelCase__):
      _lowercase				:					Any						=			"""swin2sr"""

      _lowercase				:					Tuple						=			{
          """hidden_size""": """embed_dim""",
          """num_attention_heads""": """num_heads""",
          """num_hidden_layers""": """num_layers""",
      }







      def __init__(    self		,      lowerCAmelCase__=6_4		,      lowerCAmelCase__=1		,      lowerCAmelCase__=3		,      lowerCAmelCase__=1_8_0		,      lowerCAmelCase__=[6, 6, 6, 6, 6, 6]		,      lowerCAmelCase__=[6, 6, 6, 6, 6, 6]		,      lowerCAmelCase__=8		,      lowerCAmelCase__=2.0		,      lowerCAmelCase__=True		,      lowerCAmelCase__=0.0		,      lowerCAmelCase__=0.0		,      lowerCAmelCase__=0.1		,      lowerCAmelCase__="gelu"		,      lowerCAmelCase__=False		,      lowerCAmelCase__=0.02		,      lowerCAmelCase__=1E-5		,      lowerCAmelCase__=2		,      lowerCAmelCase__=1.0		,      lowerCAmelCase__="1conv"		,      lowerCAmelCase__="pixelshuffle"		,      **lowerCAmelCase__		,      )   ->						int:



        '''simple docstring'''




        super().__init__(**lowerCAmelCase__							)

        a__	:					Optional[Any]										=image_size
        a__	:					Dict										=patch_size
        a__	:					Tuple										=num_channels
        a__	:					Union[str, Any]										=embed_dim
        a__	:					Optional[Any]										=depths
        a__	:					List[str]										=len(lowerCAmelCase__							)
        a__	:					Any										=num_heads
        a__	:					Any										=window_size
        a__	:					str										=mlp_ratio
        a__	:					List[str]										=qkv_bias
        a__	:					Dict										=hidden_dropout_prob
        a__	:					List[str]										=attention_probs_dropout_prob
        a__	:					Dict										=drop_path_rate
        a__	:					Optional[Any]										=hidden_act
        a__	:					Union[str, Any]										=use_absolute_embeddings
        a__	:					Optional[Any]										=layer_norm_eps
        a__	:					List[Any]										=initializer_range
        a__	:					int										=upscale
        a__	:					Optional[int]										=img_range
        a__	:					Any										=resi_connection
        a__	:					Optional[Any]										=upsampler






