
import gc
import unittest

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

from diffusers import TransformeraDModel, VQDiffusionPipeline, VQDiffusionScheduler, VQModel
from diffusers.pipelines.vq_diffusion.pipeline_vq_diffusion import LearnedClassifierFreeSamplingEmbeddings
from diffusers.utils import load_numpy, slow, torch_device
from diffusers.utils.testing_utils import require_torch_gpu


UpperCamelCase_   	=	False

class       a		(   unittest.TestCase					):







       def       UpperCAmelCase__					(						self  :  List[Any] ):
        """simple docstring"""







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







       @property
       def       UpperCAmelCase__					(						self  :  Optional[Any] ):
        """simple docstring"""







        return 12







       @property
       def       UpperCAmelCase__					(						self  :  List[Any] ):
        """simple docstring"""







        return 12







       @property
       def       UpperCAmelCase__					(						self  :  str ):
        """simple docstring"""







        return 32







       @property
       def       UpperCAmelCase__					(						self  :  Union[str, Any] ):
        """simple docstring"""







        torch.manual_seed(0 )
        __lowerCAmelCase         =							VQModel(
            block_out_channels=[32, 64]      ,				in_channels=3      ,				out_channels=3      ,				down_block_types=["DownEncoderBlock2D", "DownEncoderBlock2D"]      ,				up_block_types=["UpDecoderBlock2D", "UpDecoderBlock2D"]      ,				latent_channels=3      ,				num_vq_embeddings=self.num_embed      ,				vq_embed_dim=3      ,				)
        return model







       @property
       def       UpperCAmelCase__					(						self  :  Dict ):
        """simple docstring"""







        __lowerCAmelCase         =							CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip" )
        return tokenizer







       @property
       def       UpperCAmelCase__					(						self  :  Dict ):
        """simple docstring"""







        torch.manual_seed(0 )
        __lowerCAmelCase         =							CLIPTextConfig(
            bos_token_id=0      ,				eos_token_id=2      ,				hidden_size=self.text_embedder_hidden_size      ,				intermediate_size=37      ,				layer_norm_eps=1E-05      ,				num_attention_heads=4      ,				num_hidden_layers=5      ,				pad_token_id=1      ,				vocab_size=1_000      ,				)
        return CLIPTextModel(_UpperCAmelCase )







       @property
       def       UpperCAmelCase__					(						self  :  Tuple ):
        """simple docstring"""







        torch.manual_seed(0 )

        __lowerCAmelCase         =							12
        __lowerCAmelCase         =							12

        __lowerCAmelCase         =							{
            '''attention_bias''': True,
            '''cross_attention_dim''': 32,
            '''attention_head_dim''': height * width,
            '''num_attention_heads''': 1,
            '''num_vector_embeds''': self.num_embed,
            '''num_embeds_ada_norm''': self.num_embeds_ada_norm,
            '''norm_num_groups''': 32,
            '''sample_size''': width,
            '''activation_fn''': '''geglu-approximate''',
        }

        __lowerCAmelCase         =							TransformeraDModel(**_UpperCAmelCase )
        return model







       def       UpperCAmelCase__					(						self  :  List[Any] ):
        """simple docstring"""







        __lowerCAmelCase         =							'''cpu'''

        __lowerCAmelCase         =							self.dummy_vqvae
        __lowerCAmelCase         =							self.dummy_text_encoder
        __lowerCAmelCase         =							self.dummy_tokenizer
        __lowerCAmelCase         =							self.dummy_transformer
        __lowerCAmelCase         =							VQDiffusionScheduler(self.num_embed )
        __lowerCAmelCase         =							LearnedClassifierFreeSamplingEmbeddings(learnable=_UpperCAmelCase )

        __lowerCAmelCase         =							VQDiffusionPipeline(
            vqvae=_UpperCAmelCase      ,				text_encoder=_UpperCAmelCase      ,				tokenizer=_UpperCAmelCase      ,				transformer=_UpperCAmelCase      ,				scheduler=_UpperCAmelCase      ,				learned_classifier_free_sampling_embeddings=_UpperCAmelCase      ,				)
        __lowerCAmelCase         =							pipe.to(_UpperCAmelCase )
        pipe.set_progress_bar_config(disable=_UpperCAmelCase )

        __lowerCAmelCase         =							'''teddy bear playing in the pool'''

        __lowerCAmelCase         =							torch.Generator(device=_UpperCAmelCase ).manual_seed(0 )
        __lowerCAmelCase         =							pipe([prompt]      ,				generator=_UpperCAmelCase      ,				num_inference_steps=2      ,				output_type="np" )
        __lowerCAmelCase         =							output.images

        __lowerCAmelCase         =							torch.Generator(device=_UpperCAmelCase ).manual_seed(0 )
        __lowerCAmelCase         =							pipe(
            [prompt]      ,				generator=_UpperCAmelCase      ,				output_type="np"      ,				return_dict=_UpperCAmelCase      ,				num_inference_steps=2 )[0]

        __lowerCAmelCase         =							image[0, -3:, -3:, -1]
        __lowerCAmelCase         =							image_from_tuple[0, -3:, -3:, -1]

        assert image.shape == (1, 24, 24, 3)

        __lowerCAmelCase         =							np.array([0.6_5_5_1, 0.6_1_6_8, 0.5_0_0_8, 0.5_6_7_6, 0.5_6_5_9, 0.4_2_9_5, 0.6_0_7_3, 0.5_5_9_9, 0.4_9_9_2] )

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







       def       UpperCAmelCase__					(						self  :  Dict ):
        """simple docstring"""







        __lowerCAmelCase         =							'''cpu'''

        __lowerCAmelCase         =							self.dummy_vqvae
        __lowerCAmelCase         =							self.dummy_text_encoder
        __lowerCAmelCase         =							self.dummy_tokenizer
        __lowerCAmelCase         =							self.dummy_transformer
        __lowerCAmelCase         =							VQDiffusionScheduler(self.num_embed )
        __lowerCAmelCase         =							LearnedClassifierFreeSamplingEmbeddings(
            learnable=_UpperCAmelCase      ,				hidden_size=self.text_embedder_hidden_size      ,				length=tokenizer.model_max_length )

        __lowerCAmelCase         =							VQDiffusionPipeline(
            vqvae=_UpperCAmelCase      ,				text_encoder=_UpperCAmelCase      ,				tokenizer=_UpperCAmelCase      ,				transformer=_UpperCAmelCase      ,				scheduler=_UpperCAmelCase      ,				learned_classifier_free_sampling_embeddings=_UpperCAmelCase      ,				)
        __lowerCAmelCase         =							pipe.to(_UpperCAmelCase )
        pipe.set_progress_bar_config(disable=_UpperCAmelCase )

        __lowerCAmelCase         =							'''teddy bear playing in the pool'''

        __lowerCAmelCase         =							torch.Generator(device=_UpperCAmelCase ).manual_seed(0 )
        __lowerCAmelCase         =							pipe([prompt]      ,				generator=_UpperCAmelCase      ,				num_inference_steps=2      ,				output_type="np" )
        __lowerCAmelCase         =							output.images

        __lowerCAmelCase         =							torch.Generator(device=_UpperCAmelCase ).manual_seed(0 )
        __lowerCAmelCase         =							pipe(
            [prompt]      ,				generator=_UpperCAmelCase      ,				output_type="np"      ,				return_dict=_UpperCAmelCase      ,				num_inference_steps=2 )[0]

        __lowerCAmelCase         =							image[0, -3:, -3:, -1]
        __lowerCAmelCase         =							image_from_tuple[0, -3:, -3:, -1]

        assert image.shape == (1, 24, 24, 3)

        __lowerCAmelCase         =							np.array([0.6_6_9_3, 0.6_0_7_5, 0.4_9_5_9, 0.5_7_0_1, 0.5_5_8_3, 0.4_3_3_3, 0.6_1_7_1, 0.5_6_8_4, 0.4_9_8_8] )

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

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







       def       UpperCAmelCase__					(						self  :  Union[str, Any] ):
        """simple docstring"""







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







       def       UpperCAmelCase__					(						self  :  Tuple ):
        """simple docstring"""







        __lowerCAmelCase         =							load_numpy(
            "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"
            "/vq_diffusion/teddy_bear_pool_classifier_free_sampling.npy" )

        __lowerCAmelCase         =							VQDiffusionPipeline.from_pretrained("microsoft/vq-diffusion-ithq" )
        __lowerCAmelCase         =							pipeline.to(_UpperCAmelCase )
        pipeline.set_progress_bar_config(disable=_UpperCAmelCase )

        # requires GPU generator for gumbel softmax
        # don't use GPU generator in tests though
        __lowerCAmelCase         =							torch.Generator(device=_UpperCAmelCase ).manual_seed(0 )
        __lowerCAmelCase         =							pipeline(
            "teddy bear playing in the pool"      ,				num_images_per_prompt=1      ,				generator=_UpperCAmelCase      ,				output_type="np"      ,				)

        __lowerCAmelCase         =							output.images[0]

        assert image.shape == (256, 256, 3)
        assert np.abs(expected_image - image ).max() < 2.0








import tempfile

import torch

from diffusers import (
    DEISMultistepScheduler,
    DPMSolverMultistepScheduler,
    DPMSolverSinglestepScheduler,
    UniPCMultistepScheduler,
)

from .test_schedulers import SchedulerCommonTest

class       a		(   __UpperCAmelCase					):
							lowercase_		:						Dict   							=  (DPMSolverSinglestepScheduler,)
							lowercase_		:						Optional[Any]   							=  (('num_inference_steps', 25),)







							def       UpperCAmelCase__					(						self  :  Optional[Any]      ,				**snake_case__  :  Union[str, Any] ):
								"""simple docstring"""







								__lowerCAmelCase         =							{
								    "num_train_timesteps": 1_000,
								    "beta_start": 0.0_0_0_1,
								    "beta_end": 0.0_2,
								    "beta_schedule": "linear",
								    "solver_order": 2,
								    "prediction_type": "epsilon",
								    "thresholding": False,
								    "sample_max_value": 1.0,
								    "algorithm_type": "dpmsolver++",
								    "solver_type": "midpoint",
								    "lambda_min_clipped": -float("inf" ),
								    "variance_type": None,
								}

								config.update(**snake_case__ )
								return config







							def       UpperCAmelCase__					(						self  :  Any      ,				snake_case__  :  Tuple=0      ,				**snake_case__  :  Optional[Any] ):
								"""simple docstring"""







								__lowerCAmelCase         =							dict(self.forward_default_kwargs )
								__lowerCAmelCase         =							kwargs.pop("num_inference_steps"      ,				snake_case__ )
								__lowerCAmelCase         =							self.dummy_sample
								__lowerCAmelCase         =							0.1 * sample
								__lowerCAmelCase         =							[residual + 0.2, residual + 0.1_5, residual + 0.1_0]

								for scheduler_class in self.scheduler_classes:
									__lowerCAmelCase         =							self.get_scheduler_config(**snake_case__ )
									__lowerCAmelCase         =							scheduler_class(**snake_case__ )
									scheduler.set_timesteps(snake_case__ )
									# copy over dummy past residuals
									__lowerCAmelCase         =							dummy_past_residuals[: scheduler.config.solver_order]

									with tempfile.TemporaryDirectory() as tmpdirname:
										scheduler.save_config(snake_case__ )
										__lowerCAmelCase         =							scheduler_class.from_pretrained(snake_case__ )
										new_scheduler.set_timesteps(snake_case__ )
										# copy over dummy past residuals
										__lowerCAmelCase         =							dummy_past_residuals[: new_scheduler.config.solver_order]

									__lowerCAmelCase   ,						__lowerCAmelCase         =							sample, sample
									for t in range(snake_case__      ,				time_step + scheduler.config.solver_order + 1 ):
										__lowerCAmelCase         =							scheduler.step(snake_case__      ,				snake_case__      ,				snake_case__      ,				**snake_case__ ).prev_sample
										__lowerCAmelCase         =							new_scheduler.step(snake_case__      ,				snake_case__      ,				snake_case__      ,				**snake_case__ ).prev_sample

										assert torch.sum(torch.abs(output - new_output ) ) < 1E-5, "Scheduler outputs are not identical"







							def       UpperCAmelCase__					(						self  :  List[str] ):
								"""simple docstring"""







								pass







							def       UpperCAmelCase__					(						self  :  str      ,				snake_case__  :  Tuple=0      ,				**snake_case__  :  Tuple ):
								"""simple docstring"""







								__lowerCAmelCase         =							dict(self.forward_default_kwargs )
								__lowerCAmelCase         =							kwargs.pop("num_inference_steps"      ,				snake_case__ )
								__lowerCAmelCase         =							self.dummy_sample
								__lowerCAmelCase         =							0.1 * sample
								__lowerCAmelCase         =							[residual + 0.2, residual + 0.1_5, residual + 0.1_0]

								for scheduler_class in self.scheduler_classes:
									__lowerCAmelCase         =							self.get_scheduler_config()
									__lowerCAmelCase         =							scheduler_class(**snake_case__ )
									scheduler.set_timesteps(snake_case__ )

									# copy over dummy past residuals (must be after setting timesteps)
									__lowerCAmelCase         =							dummy_past_residuals[: scheduler.config.solver_order]

									with tempfile.TemporaryDirectory() as tmpdirname:
										scheduler.save_config(snake_case__ )
										__lowerCAmelCase         =							scheduler_class.from_pretrained(snake_case__ )
										# copy over dummy past residuals
										new_scheduler.set_timesteps(snake_case__ )

										# copy over dummy past residual (must be after setting timesteps)
										__lowerCAmelCase         =							dummy_past_residuals[: new_scheduler.config.solver_order]

									__lowerCAmelCase         =							scheduler.step(snake_case__      ,				snake_case__      ,				snake_case__      ,				**snake_case__ ).prev_sample
									__lowerCAmelCase         =							new_scheduler.step(snake_case__      ,				snake_case__      ,				snake_case__      ,				**snake_case__ ).prev_sample

									assert torch.sum(torch.abs(output - new_output ) ) < 1E-5, "Scheduler outputs are not identical"







							def       UpperCAmelCase__					(						self  :  Optional[Any]      ,				snake_case__  :  Optional[Any]=None      ,				**snake_case__  :  Optional[int] ):
								"""simple docstring"""







								if scheduler is None:
									__lowerCAmelCase         =							self.scheduler_classes[0]
									__lowerCAmelCase         =							self.get_scheduler_config(**snake_case__ )
									__lowerCAmelCase         =							scheduler_class(**snake_case__ )

								__lowerCAmelCase         =							self.scheduler_classes[0]
								__lowerCAmelCase         =							self.get_scheduler_config(**snake_case__ )
								__lowerCAmelCase         =							scheduler_class(**snake_case__ )

								__lowerCAmelCase         =							10
								__lowerCAmelCase         =							self.dummy_model()
								__lowerCAmelCase         =							self.dummy_sample_deter
								scheduler.set_timesteps(snake_case__ )

								for i, t in enumerate(scheduler.timesteps ):
									__lowerCAmelCase         =							model(snake_case__      ,				snake_case__ )
									__lowerCAmelCase         =							scheduler.step(snake_case__      ,				snake_case__      ,				snake_case__ ).prev_sample

								return sample







							def       UpperCAmelCase__					(						self  :  Union[str, Any] ):
								"""simple docstring"""







								__lowerCAmelCase         =							DPMSolverSinglestepScheduler(**self.get_scheduler_config() )
								__lowerCAmelCase         =							50
								__lowerCAmelCase         =							self.dummy_model()
								__lowerCAmelCase         =							self.dummy_sample_deter
								scheduler.set_timesteps(snake_case__ )

								# make sure that the first t is uneven
								for i, t in enumerate(scheduler.timesteps[3:] ):
									__lowerCAmelCase         =							model(snake_case__      ,				snake_case__ )
									__lowerCAmelCase         =							scheduler.step(snake_case__      ,				snake_case__      ,				snake_case__ ).prev_sample

								__lowerCAmelCase         =							torch.mean(torch.abs(snake_case__ ) )

								assert abs(result_mean.item() - 0.2_5_7_4 ) < 1E-3







							def       UpperCAmelCase__					(						self  :  List[str] ):
								"""simple docstring"""







								for timesteps in [25, 50, 100, 999, 1_000]:
									self.check_over_configs(num_train_timesteps=snake_case__ )







							def       UpperCAmelCase__					(						self  :  Optional[int] ):
								"""simple docstring"""







								__lowerCAmelCase         =							DPMSolverSinglestepScheduler(**self.get_scheduler_config() )
								__lowerCAmelCase         =							self.full_loop(scheduler=snake_case__ )
								__lowerCAmelCase         =							torch.mean(torch.abs(snake_case__ ) )

								assert abs(result_mean.item() - 0.2_7_9_1 ) < 1E-3

								__lowerCAmelCase         =							DEISMultistepScheduler.from_config(scheduler.config )
								__lowerCAmelCase         =							DPMSolverMultistepScheduler.from_config(scheduler.config )
								__lowerCAmelCase         =							UniPCMultistepScheduler.from_config(scheduler.config )
								__lowerCAmelCase         =							DPMSolverSinglestepScheduler.from_config(scheduler.config )

								__lowerCAmelCase         =							self.full_loop(scheduler=snake_case__ )
								__lowerCAmelCase         =							torch.mean(torch.abs(snake_case__ ) )

								assert abs(result_mean.item() - 0.2_7_9_1 ) < 1E-3







							def       UpperCAmelCase__					(						self  :  str ):
								"""simple docstring"""







								self.check_over_configs(thresholding=snake_case__ )
								for order in [1, 2, 3]:
									for solver_type in ["midpoint", "heun"]:
										for threshold in [0.5, 1.0, 2.0]:
											for prediction_type in ["epsilon", "sample"]:
												self.check_over_configs(
												    thresholding=snake_case__      ,				prediction_type=snake_case__      ,				sample_max_value=snake_case__      ,				algorithm_type="dpmsolver++"      ,				solver_order=snake_case__      ,				solver_type=snake_case__      ,				)







							def       UpperCAmelCase__					(						self  :  Union[str, Any] ):
								"""simple docstring"""







								for prediction_type in ["epsilon", "v_prediction"]:
									self.check_over_configs(prediction_type=snake_case__ )







							def       UpperCAmelCase__					(						self  :  str ):
								"""simple docstring"""







								for algorithm_type in ["dpmsolver", "dpmsolver++"]:
									for solver_type in ["midpoint", "heun"]:
										for order in [1, 2, 3]:
											for prediction_type in ["epsilon", "sample"]:
												self.check_over_configs(
												    solver_order=snake_case__      ,				solver_type=snake_case__      ,				prediction_type=snake_case__      ,				algorithm_type=snake_case__      ,				)
												__lowerCAmelCase         =							self.full_loop(
												    solver_order=snake_case__      ,				solver_type=snake_case__      ,				prediction_type=snake_case__      ,				algorithm_type=snake_case__      ,				)
												assert not torch.isnan(snake_case__ ).any(), "Samples have nan numbers"







							def       UpperCAmelCase__					(						self  :  List[str] ):
								"""simple docstring"""







								self.check_over_configs(lower_order_final=snake_case__ )
								self.check_over_configs(lower_order_final=snake_case__ )







							def       UpperCAmelCase__					(						self  :  Dict ):
								"""simple docstring"""







								self.check_over_configs(lambda_min_clipped=-float("inf" ) )
								self.check_over_configs(lambda_min_clipped=-5.1 )







							def       UpperCAmelCase__					(						self  :  Any ):
								"""simple docstring"""







								self.check_over_configs(variance_type=snake_case__ )
								self.check_over_configs(variance_type="learned_range" )







							def       UpperCAmelCase__					(						self  :  List[Any] ):
								"""simple docstring"""







								for num_inference_steps in [1, 2, 3, 5, 10, 50, 100, 999, 1_000]:
									self.check_over_forward(num_inference_steps=snake_case__      ,				time_step=0 )







							def       UpperCAmelCase__					(						self  :  List[Any] ):
								"""simple docstring"""







								__lowerCAmelCase         =							self.full_loop()
								__lowerCAmelCase         =							torch.mean(torch.abs(snake_case__ ) )

								assert abs(result_mean.item() - 0.2_7_9_1 ) < 1E-3







							def       UpperCAmelCase__					(						self  :  List[str] ):
								"""simple docstring"""







								__lowerCAmelCase         =							self.full_loop(use_karras_sigmas=snake_case__ )
								__lowerCAmelCase         =							torch.mean(torch.abs(snake_case__ ) )

								assert abs(result_mean.item() - 0.2_2_4_8 ) < 1E-3







							def       UpperCAmelCase__					(						self  :  List[str] ):
								"""simple docstring"""







								__lowerCAmelCase         =							self.full_loop(prediction_type="v_prediction" )
								__lowerCAmelCase         =							torch.mean(torch.abs(snake_case__ ) )

								assert abs(result_mean.item() - 0.1_4_5_3 ) < 1E-3







							def       UpperCAmelCase__					(						self  :  Optional[Any] ):
								"""simple docstring"""







								__lowerCAmelCase         =							self.full_loop(prediction_type="v_prediction"      ,				use_karras_sigmas=snake_case__ )
								__lowerCAmelCase         =							torch.mean(torch.abs(snake_case__ ) )

								assert abs(result_mean.item() - 0.0_6_4_9 ) < 1E-3







							def       UpperCAmelCase__					(						self  :  int ):
								"""simple docstring"""







								__lowerCAmelCase         =							self.scheduler_classes[0]
								__lowerCAmelCase         =							self.get_scheduler_config(thresholding=snake_case__      ,				dynamic_thresholding_ratio=0 )
								__lowerCAmelCase         =							scheduler_class(**snake_case__ )

								__lowerCAmelCase         =							10
								__lowerCAmelCase         =							self.dummy_model()
								__lowerCAmelCase         =							self.dummy_sample_deter.half()
								scheduler.set_timesteps(snake_case__ )

								for i, t in enumerate(scheduler.timesteps ):
									__lowerCAmelCase         =							model(snake_case__      ,				snake_case__ )
									__lowerCAmelCase         =							scheduler.step(snake_case__      ,				snake_case__      ,				snake_case__ ).prev_sample

								assert sample.dtype == torch.floataa






