"""simple docstring"""






import tempfile

import torch

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

from .test_schedulers import SchedulerCommonTest





class  _A						(       lowerCAmelCase	):
      snake_case__			:   Any							  =	(DEISMultistepScheduler,)
      snake_case__			:   int							  =	(('num_inference_steps', 25),)

      def 				A__       (						self		,       **__lowerCAmelCase						):







            """simple docstring"""



            lowercase														=	{
                """num_train_timesteps""": 1000,
                """beta_start""": 0.0_0_0_1,
                """beta_end""": 0.0_2,
                """beta_schedule""": """linear""",
                """solver_order""": 2,
            }

            config.update(**__lowerCAmelCase						)
            return config

      def 				A__       (						self		,       __lowerCAmelCase=0		,       **__lowerCAmelCase						):







            """simple docstring"""



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

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

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

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

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

      def 				A__       (						self						):







            """simple docstring"""



            pass

      def 				A__       (						self		,       __lowerCAmelCase=0		,       **__lowerCAmelCase						):







            """simple docstring"""



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

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

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

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

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

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

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

      def 				A__       (						self		,       __lowerCAmelCase=None		,       **__lowerCAmelCase						):







            """simple docstring"""



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

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

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

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

            return sample

      def 				A__       (						self						):







            """simple docstring"""



            lowercase														=	dict(self.forward_default_kwargs						)

            lowercase														=	kwargs.pop("""num_inference_steps"""		,       __lowerCAmelCase						)

            for scheduler_class in self.scheduler_classes:
                  lowercase														=	self.get_scheduler_config()
                  lowercase														=	scheduler_class(**__lowerCAmelCase						)

                  lowercase														=	self.dummy_sample
                  lowercase														=	0.1 * sample

                  if num_inference_steps is not None and hasattr(__lowerCAmelCase		,       """set_timesteps"""						):
                        scheduler.set_timesteps(__lowerCAmelCase						)
                  elif num_inference_steps is not None and not hasattr(__lowerCAmelCase		,       """set_timesteps"""						):
                        lowercase														=	num_inference_steps

                  # copy over dummy past residuals (must be done after set_timesteps)
                  lowercase														=	[residual + 0.2, residual + 0.1_5, residual + 0.1_0]
                  lowercase														=	dummy_past_residuals[: scheduler.config.solver_order]

                  lowercase														=	scheduler.timesteps[5]
                  lowercase														=	scheduler.timesteps[6]

                  lowercase														=	scheduler.step(__lowerCAmelCase		,       __lowerCAmelCase		,       __lowerCAmelCase		,       **__lowerCAmelCase						).prev_sample
                  lowercase														=	scheduler.step(__lowerCAmelCase		,       __lowerCAmelCase		,       __lowerCAmelCase		,       **__lowerCAmelCase						).prev_sample

                  self.assertEqual(output_a.shape		,       sample.shape						)
                  self.assertEqual(output_a.shape		,       output_a.shape						)

      def 				A__       (						self						):







            """simple docstring"""



            lowercase														=	DEISMultistepScheduler(**self.get_scheduler_config()						)
            lowercase														=	self.full_loop(scheduler=__lowerCAmelCase						)
            lowercase														=	torch.mean(torch.abs(__lowerCAmelCase						)						)

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

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

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

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

      def 				A__       (						self						):







            """simple docstring"""



            for timesteps in [25, 50, 100, 999, 1000]:
                  self.check_over_configs(num_train_timesteps=__lowerCAmelCase						)

      def 				A__       (						self						):







            """simple docstring"""



            self.check_over_configs(thresholding=__lowerCAmelCase						)
            for order in [1, 2, 3]:
                  for solver_type in ["logrho"]:
                        for threshold in [0.5, 1.0, 2.0]:
                              for prediction_type in ["epsilon", "sample"]:
                                    self.check_over_configs(
                                        thresholding=__lowerCAmelCase		,       prediction_type=__lowerCAmelCase		,       sample_max_value=__lowerCAmelCase		,       algorithm_type="""deis"""		,       solver_order=__lowerCAmelCase		,       solver_type=__lowerCAmelCase		,       )

      def 				A__       (						self						):







            """simple docstring"""



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

      def 				A__       (						self						):







            """simple docstring"""



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

      def 				A__       (						self						):







            """simple docstring"""



            self.check_over_configs(lower_order_final=__lowerCAmelCase						)
            self.check_over_configs(lower_order_final=__lowerCAmelCase						)

      def 				A__       (						self						):







            """simple docstring"""



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

      def 				A__       (						self						):







            """simple docstring"""



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

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

      def 				A__       (						self						):







            """simple docstring"""



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

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







      def 				A__       (						self						):







            """simple docstring"""



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

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

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

            assert sample.dtype == torch.floataa




"""simple docstring"""






import numpy as np
import torch
from imwatermark import WatermarkEncoder


# Copied from https://github.com/Stability-AI/generative-models/blob/613af104c6b85184091d42d374fef420eddb356d/scripts/demo/streamlit_helpers.py#L66
__lowerCAmelCase    :       Union[str, Any]					     =0B1011_0011_1110_1100_1001_0000_0111_1011_1011_0001_1001_1110
# bin(x)[2:] gives bits of x as str, use int to convert them to 0/1
__lowerCAmelCase    :       Optional[int]					     =[int(bit) for bit in bin(WATERMARK_MESSAGE)[2:]]





class  _A						:

      def __init__(						self						):







            """simple docstring"""



            lowercase														=	WATERMARK_BITS
            lowercase														=	WatermarkEncoder()

            self.encoder.set_watermark("""bits"""		,       self.watermark						)







      def 				A__       (						self		,       __lowerCAmelCase						):







            """simple docstring"""



            if images.shape[-1] < 256:
                  return images

            lowercase														=	(255 * (images / 2 + 0.5)).cpu().permute(0		,       2		,       3		,       1						).float().numpy()

            lowercase														=	[self.encoder.encode(__lowerCAmelCase		,       """dwtDct"""						) for image in images]

            lowercase														=	torch.from_numpy(np.array(__lowerCAmelCase						)						).permute(0		,       3		,       1		,       2						)

            lowercase														=	torch.clamp(2 * (images / 255 - 0.5)		,       min=-1.0		,       max=1.0						)
            return images



