
import gc
import random
import unittest

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

from diffusers import AutoencoderKL, DDIMScheduler, DDPMScheduler, StableDiffusionUpscalePipeline, UNetaDConditionModel
from diffusers.utils import floats_tensor, load_image, load_numpy, slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu


enable_full_determinism()




class 					_SCREAMING_SNAKE_CASE  (      unittest.TestCase ):
      def A_			(			self ):
       super().tearDown()
       gc.collect()
       torch.cuda.empty_cache()
      @property
      def A_			(			self ):
       snake_case__						=				1
       snake_case__						=				3
       snake_case__						=				(32, 32)

       snake_case__						=				floats_tensor((batch_size, num_channels) + sizes		,	rng=random.Random(0 ) ).to(lowerCamelCase )
       return image
      @property
      def A_			(			self ):
       torch.manual_seed(0 )
       snake_case__						=				UNetaDConditionModel(
           block_out_channels=(32, 32, 64)		,	layers_per_block=2		,	sample_size=32		,	in_channels=7		,	out_channels=4		,	down_block_types=("DownBlock2D", "CrossAttnDownBlock2D", "CrossAttnDownBlock2D")		,	up_block_types=("CrossAttnUpBlock2D", "CrossAttnUpBlock2D", "UpBlock2D")		,	cross_attention_dim=32		,	attention_head_dim=8		,	use_linear_projection=lowerCamelCase		,	only_cross_attention=(True, True, False)		,	num_class_embeds=1_00		,	)
       return model
      @property
      def A_			(			self ):
       torch.manual_seed(0 )
       snake_case__						=				AutoencoderKL(
           block_out_channels=[32, 32, 64]		,	in_channels=3		,	out_channels=3		,	down_block_types=["DownEncoderBlock2D", "DownEncoderBlock2D", "DownEncoderBlock2D"]		,	up_block_types=["UpDecoderBlock2D", "UpDecoderBlock2D", "UpDecoderBlock2D"]		,	latent_channels=4		,	)
       return model
      @property
      def A_			(			self ):
       torch.manual_seed(0 )
       snake_case__						=				CLIPTextConfig(
           bos_token_id=0		,	eos_token_id=2		,	hidden_size=32		,	intermediate_size=37		,	layer_norm_eps=1e-05		,	num_attention_heads=4		,	num_hidden_layers=5		,	pad_token_id=1		,	vocab_size=10_00		,	hidden_act="gelu"		,	projection_dim=5_12		,	)
       return CLIPTextModel(lowerCamelCase )
      def A_			(			self ):
       snake_case__						=				"cpu"  # ensure determinism for the device-dependent torch.Generator
       snake_case__						=				self.dummy_cond_unet_upscale
       snake_case__						=				DDPMScheduler()
       snake_case__						=				DDIMScheduler(prediction_type="v_prediction" )
       snake_case__						=				self.dummy_vae
       snake_case__						=				self.dummy_text_encoder
       snake_case__						=				CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip" )

       snake_case__						=				self.dummy_image.cpu().permute(0		,	2		,	3		,	1 )[0]
       snake_case__						=				Image.fromarray(np.uinta(lowerCamelCase ) ).convert("RGB" ).resize((64, 64) )

       # make sure here that pndm scheduler skips prk
       snake_case__						=				StableDiffusionUpscalePipeline(
           unet=lowerCamelCase		,	low_res_scheduler=lowerCamelCase		,	scheduler=lowerCamelCase		,	vae=lowerCamelCase		,	text_encoder=lowerCamelCase		,	tokenizer=lowerCamelCase		,	max_noise_level=3_50		,	)
       snake_case__						=				sd_pipe.to(lowerCamelCase )
       sd_pipe.set_progress_bar_config(disable=lowerCamelCase )

       snake_case__						=				"A painting of a squirrel eating a burger"
       snake_case__						=				torch.Generator(device=lowerCamelCase ).manual_seed(0 )
       snake_case__						=				sd_pipe(
           [prompt]		,	image=lowerCamelCase		,	generator=lowerCamelCase		,	guidance_scale=6.0		,	noise_level=20		,	num_inference_steps=2		,	output_type="np"		,	)

       snake_case__						=				output.images

       snake_case__						=				torch.Generator(device=lowerCamelCase ).manual_seed(0 )
       snake_case__						=				sd_pipe(
           [prompt]		,	image=lowerCamelCase		,	generator=lowerCamelCase		,	guidance_scale=6.0		,	noise_level=20		,	num_inference_steps=2		,	output_type="np"		,	return_dict=lowerCamelCase		,	)[0]

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

       snake_case__						=				low_res_image.size[0] * 4
       assert image.shape == (1, expected_height_width, expected_height_width, 3)
       snake_case__						=				np.array([0.3_1_1_3, 0.3_9_1_0, 0.4_2_7_2, 0.4_8_5_9, 0.5_0_6_1, 0.4_6_5_2, 0.5_3_6_2, 0.5_7_1_5, 0.5_6_6_1] )

       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 A_			(			self ):
       snake_case__						=				"cpu"  # ensure determinism for the device-dependent torch.Generator
       snake_case__						=				self.dummy_cond_unet_upscale
       snake_case__						=				DDPMScheduler()
       snake_case__						=				DDIMScheduler(prediction_type="v_prediction" )
       snake_case__						=				self.dummy_vae
       snake_case__						=				self.dummy_text_encoder
       snake_case__						=				CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip" )

       snake_case__						=				self.dummy_image.cpu().permute(0		,	2		,	3		,	1 )[0]
       snake_case__						=				Image.fromarray(np.uinta(lowerCamelCase ) ).convert("RGB" ).resize((64, 64) )

       # make sure here that pndm scheduler skips prk
       snake_case__						=				StableDiffusionUpscalePipeline(
           unet=lowerCamelCase		,	low_res_scheduler=lowerCamelCase		,	scheduler=lowerCamelCase		,	vae=lowerCamelCase		,	text_encoder=lowerCamelCase		,	tokenizer=lowerCamelCase		,	max_noise_level=3_50		,	)
       snake_case__						=				sd_pipe.to(lowerCamelCase )
       sd_pipe.set_progress_bar_config(disable=lowerCamelCase )

       snake_case__						=				"A painting of a squirrel eating a burger"
       snake_case__						=				sd_pipe(
           2 * [prompt]		,	image=2 * [low_res_image]		,	guidance_scale=6.0		,	noise_level=20		,	num_inference_steps=2		,	output_type="np"		,	)
       snake_case__						=				output.images
       assert image.shape[0] == 2

       snake_case__						=				torch.Generator(device=lowerCamelCase ).manual_seed(0 )
       snake_case__						=				sd_pipe(
           [prompt]		,	image=lowerCamelCase		,	generator=lowerCamelCase		,	num_images_per_prompt=2		,	guidance_scale=6.0		,	noise_level=20		,	num_inference_steps=2		,	output_type="np"		,	)
       snake_case__						=				output.images
       assert image.shape[0] == 2



      @unittest.skipIf(torch_device != "cuda"		,	"This test requires a GPU" )
      def A_			(			self ):
       snake_case__						=				self.dummy_cond_unet_upscale
       snake_case__						=				DDPMScheduler()
       snake_case__						=				DDIMScheduler(prediction_type="v_prediction" )
       snake_case__						=				self.dummy_vae
       snake_case__						=				self.dummy_text_encoder
       snake_case__						=				CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip" )

       snake_case__						=				self.dummy_image.cpu().permute(0		,	2		,	3		,	1 )[0]
       snake_case__						=				Image.fromarray(np.uinta(lowerCamelCase ) ).convert("RGB" ).resize((64, 64) )

       # put models in fp16, except vae as it overflows in fp16
       snake_case__						=				unet.half()
       snake_case__						=				text_encoder.half()

       # make sure here that pndm scheduler skips prk
       snake_case__						=				StableDiffusionUpscalePipeline(
           unet=lowerCamelCase		,	low_res_scheduler=lowerCamelCase		,	scheduler=lowerCamelCase		,	vae=lowerCamelCase		,	text_encoder=lowerCamelCase		,	tokenizer=lowerCamelCase		,	max_noise_level=3_50		,	)
       snake_case__						=				sd_pipe.to(lowerCamelCase )
       sd_pipe.set_progress_bar_config(disable=lowerCamelCase )

       snake_case__						=				"A painting of a squirrel eating a burger"
       snake_case__						=				torch.manual_seed(0 )
       snake_case__						=				sd_pipe(
           [prompt]		,	image=lowerCamelCase		,	generator=lowerCamelCase		,	num_inference_steps=2		,	output_type="np"		,	).images

       snake_case__						=				low_res_image.size[0] * 4
       assert image.shape == (1, expected_height_width, expected_height_width, 3)





@slow
@require_torch_gpu
class 					_SCREAMING_SNAKE_CASE  (      unittest.TestCase ):
      def A_			(			self ):
       super().tearDown()
       gc.collect()
       torch.cuda.empty_cache()
      def A_			(			self ):
       snake_case__						=				load_image(
           "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"
           "/sd2-upscale/low_res_cat.png" )
       snake_case__						=				load_numpy(
           "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/sd2-upscale"
           "/upsampled_cat.npy" )

       snake_case__						=				"stabilityai/stable-diffusion-x4-upscaler"
       snake_case__						=				StableDiffusionUpscalePipeline.from_pretrained(lowerCamelCase )
       pipe.to(lowerCamelCase )
       pipe.set_progress_bar_config(disable=lowerCamelCase )
       pipe.enable_attention_slicing()

       snake_case__						=				"a cat sitting on a park bench"

       snake_case__						=				torch.manual_seed(0 )
       snake_case__						=				pipe(
           prompt=lowerCamelCase		,	image=lowerCamelCase		,	generator=lowerCamelCase		,	output_type="np"		,	)
       snake_case__						=				output.images[0]

       assert image.shape == (5_12, 5_12, 3)
       assert np.abs(expected_image - image ).max() < 1e-3
      def A_			(			self ):
       snake_case__						=				load_image(
           "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"
           "/sd2-upscale/low_res_cat.png" )
       snake_case__						=				load_numpy(
           "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/sd2-upscale"
           "/upsampled_cat_fp16.npy" )

       snake_case__						=				"stabilityai/stable-diffusion-x4-upscaler"
       snake_case__						=				StableDiffusionUpscalePipeline.from_pretrained(
           lowerCamelCase		,	torch_dtype=torch.floataa		,	)
       pipe.to(lowerCamelCase )
       pipe.set_progress_bar_config(disable=lowerCamelCase )
       pipe.enable_attention_slicing()

       snake_case__						=				"a cat sitting on a park bench"

       snake_case__						=				torch.manual_seed(0 )
       snake_case__						=				pipe(
           prompt=lowerCamelCase		,	image=lowerCamelCase		,	generator=lowerCamelCase		,	output_type="np"		,	)
       snake_case__						=				output.images[0]

       assert image.shape == (5_12, 5_12, 3)
       assert np.abs(expected_image - image ).max() < 5e-1



      def A_			(			self ):
       torch.cuda.empty_cache()
       torch.cuda.reset_max_memory_allocated()
       torch.cuda.reset_peak_memory_stats()

       snake_case__						=				load_image(
           "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"
           "/sd2-upscale/low_res_cat.png" )

       snake_case__						=				"stabilityai/stable-diffusion-x4-upscaler"
       snake_case__						=				StableDiffusionUpscalePipeline.from_pretrained(
           lowerCamelCase		,	torch_dtype=torch.floataa		,	)
       pipe.to(lowerCamelCase )
       pipe.set_progress_bar_config(disable=lowerCamelCase )
       pipe.enable_attention_slicing(1 )
       pipe.enable_sequential_cpu_offload()

       snake_case__						=				"a cat sitting on a park bench"

       snake_case__						=				torch.manual_seed(0 )
       snake_case__						=				pipe(
           prompt=lowerCamelCase		,	image=lowerCamelCase		,	generator=lowerCamelCase		,	num_inference_steps=5		,	output_type="np"		,	)

       snake_case__						=				torch.cuda.max_memory_allocated()
       # make sure that less than 2.9 GB is allocated
       assert mem_bytes < 2.9 * 10**9





from typing import Optional, Tuple, Union

import torch
from einops import rearrange, reduce

from diffusers import DDIMScheduler, DDPMScheduler, DiffusionPipeline, ImagePipelineOutput, UNetaDConditionModel
from diffusers.schedulers.scheduling_ddim import DDIMSchedulerOutput
from diffusers.schedulers.scheduling_ddpm import DDPMSchedulerOutput


__magic_name__            =  8





def 	SCREAMING_SNAKE_CASE__    (      __lowerCAmelCase    ,		__lowerCAmelCase=BITS							):
    snake_case__						=				x.device

    snake_case__						=				(x * 255).int().clamp(0    ,		255							)

    snake_case__						=				2 ** torch.arange(bits - 1    ,		-1    ,		-1    ,		device=__lowerCAmelCase							)
    snake_case__						=				rearrange(__lowerCAmelCase    ,		"d -> d 1 1"							)
    snake_case__						=				rearrange(__lowerCAmelCase    ,		"b c h w -> b c 1 h w"							)

    snake_case__						=				((x & mask) != 0).float()
    snake_case__						=				rearrange(__lowerCAmelCase    ,		"b c d h w -> b (c d) h w"							)
    snake_case__						=				bits * 2 - 1
    return bits





def 	SCREAMING_SNAKE_CASE__    (      __lowerCAmelCase    ,		__lowerCAmelCase=BITS							):
    snake_case__						=				x.device

    snake_case__						=				(x > 0).int()
    snake_case__						=				2 ** torch.arange(bits - 1    ,		-1    ,		-1    ,		device=__lowerCAmelCase    ,		dtype=torch.intaa							)

    snake_case__						=				rearrange(__lowerCAmelCase    ,		"d -> d 1 1"							)
    snake_case__						=				rearrange(__lowerCAmelCase    ,		"b (c d) h w -> b c d h w"    ,		d=8							)
    snake_case__						=				reduce(x * mask    ,		"b c d h w -> b c h w"    ,		"sum"							)
    return (dec / 255).clamp(0.0    ,		1.0							)





def 	SCREAMING_SNAKE_CASE__    (      self    ,		__lowerCAmelCase    ,		__lowerCAmelCase    ,		__lowerCAmelCase    ,		__lowerCAmelCase = 0.0    ,		__lowerCAmelCase = True    ,		__lowerCAmelCase=None    ,		__lowerCAmelCase = True    ,		):
    if self.num_inference_steps is None:
        raise ValueError(
            "Number of inference steps is 'None', you need to run 'set_timesteps' after creating the scheduler"							)

    # See formulas (12) and (16) of DDIM paper https://arxiv.org/pdf/2010.02502.pdf
    # Ideally, read DDIM paper in-detail understanding

    # Notation (<variable name> -> <name in paper>
    # - pred_noise_t -> e_theta(x_t, t)
    # - pred_original_sample -> f_theta(x_t, t) or x_0
    # - std_dev_t -> sigma_t
    # - eta -> η
    # - pred_sample_direction -> "direction pointing to x_t"
    # - pred_prev_sample -> "x_t-1"

    # 1. get previous step value (=t-1)
    snake_case__						=				timestep - self.config.num_train_timesteps // self.num_inference_steps

    # 2. compute alphas, betas
    snake_case__						=				self.alphas_cumprod[timestep]
    snake_case__						=				self.alphas_cumprod[prev_timestep] if prev_timestep >= 0 else self.final_alpha_cumprod

    snake_case__						=				1 - alpha_prod_t

    # 3. compute predicted original sample from predicted noise also called
    # "predicted x_0" of formula (12) from https://arxiv.org/pdf/2010.02502.pdf
    snake_case__						=				(sample - beta_prod_t ** 0.5 * model_output) / alpha_prod_t ** 0.5

    # 4. Clip "predicted x_0"
    snake_case__						=				self.bit_scale
    if self.config.clip_sample:
        snake_case__						=				torch.clamp(__lowerCAmelCase    ,		-scale    ,		__lowerCAmelCase							)

    # 5. compute variance: "sigma_t(η)" -> see formula (16)
    # σ_t = sqrt((1 − α_t−1)/(1 − α_t)) * sqrt(1 − α_t/α_t−1)
    snake_case__						=				self._get_variance(__lowerCAmelCase    ,		__lowerCAmelCase							)
    snake_case__						=				eta * variance ** 0.5

    if use_clipped_model_output:
        # the model_output is always re-derived from the clipped x_0 in Glide
        snake_case__						=				(sample - alpha_prod_t ** 0.5 * pred_original_sample) / beta_prod_t ** 0.5

    # 6. compute "direction pointing to x_t" of formula (12) from https://arxiv.org/pdf/2010.02502.pdf
    snake_case__						=				(1 - alpha_prod_t_prev - std_dev_t**2) ** 0.5 * model_output

    # 7. compute x_t without "random noise" of formula (12) from https://arxiv.org/pdf/2010.02502.pdf
    snake_case__						=				alpha_prod_t_prev ** 0.5 * pred_original_sample + pred_sample_direction

    if eta > 0:
        # randn_like does not support generator https://github.com/pytorch/pytorch/issues/27072
        snake_case__						=				model_output.device if torch.is_tensor(__lowerCAmelCase							) else "cpu"
        snake_case__						=				torch.randn(model_output.shape    ,		dtype=model_output.dtype    ,		generator=__lowerCAmelCase							).to(__lowerCAmelCase							)
        snake_case__						=				self._get_variance(__lowerCAmelCase    ,		__lowerCAmelCase							) ** 0.5 * eta * noise

        snake_case__						=				prev_sample + variance

    if not return_dict:
        return (prev_sample,)

    return DDIMSchedulerOutput(prev_sample=__lowerCAmelCase    ,		pred_original_sample=__lowerCAmelCase							)





def 	SCREAMING_SNAKE_CASE__    (      self    ,		__lowerCAmelCase    ,		__lowerCAmelCase    ,		__lowerCAmelCase    ,		__lowerCAmelCase="epsilon"    ,		__lowerCAmelCase=None    ,		__lowerCAmelCase = True    ,		):
    snake_case__						=				timestep

    if model_output.shape[1] == sample.shape[1] * 2 and self.variance_type in ["learned", "learned_range"]:
        snake_case__						,					snake_case__						=				torch.split(__lowerCAmelCase    ,		sample.shape[1]    ,		dim=1							)
    else:
        snake_case__						=				None

    # 1. compute alphas, betas
    snake_case__						=				self.alphas_cumprod[t]
    snake_case__						=				self.alphas_cumprod[t - 1] if t > 0 else self.one
    snake_case__						=				1 - alpha_prod_t
    snake_case__						=				1 - alpha_prod_t_prev

    # 2. compute predicted original sample from predicted noise also called
    # "predicted x_0" of formula (15) from https://arxiv.org/pdf/2006.11239.pdf
    if prediction_type == "epsilon":
        snake_case__						=				(sample - beta_prod_t ** 0.5 * model_output) / alpha_prod_t ** 0.5
    elif prediction_type == "sample":
        snake_case__						=				model_output
    else:
        raise ValueError(F"""Unsupported prediction_type {prediction_type}."""							)

    # 3. Clip "predicted x_0"
    snake_case__						=				self.bit_scale
    if self.config.clip_sample:
        snake_case__						=				torch.clamp(__lowerCAmelCase    ,		-scale    ,		__lowerCAmelCase							)

    # 4. Compute coefficients for pred_original_sample x_0 and current sample x_t
    # See formula (7) from https://arxiv.org/pdf/2006.11239.pdf
    snake_case__						=				(alpha_prod_t_prev ** 0.5 * self.betas[t]) / beta_prod_t
    snake_case__						=				self.alphas[t] ** 0.5 * beta_prod_t_prev / beta_prod_t

    # 5. Compute predicted previous sample µ_t
    # See formula (7) from https://arxiv.org/pdf/2006.11239.pdf
    snake_case__						=				pred_original_sample_coeff * pred_original_sample + current_sample_coeff * sample

    # 6. Add noise
    snake_case__						=				0
    if t > 0:
        snake_case__						=				torch.randn(
            model_output.size()    ,		dtype=model_output.dtype    ,		layout=model_output.layout    ,		generator=__lowerCAmelCase							).to(model_output.device							)
        snake_case__						=				(self._get_variance(__lowerCAmelCase    ,		predicted_variance=__lowerCAmelCase							) ** 0.5) * noise

    snake_case__						=				pred_prev_sample + variance

    if not return_dict:
        return (pred_prev_sample,)

    return DDPMSchedulerOutput(prev_sample=__lowerCAmelCase    ,		pred_original_sample=__lowerCAmelCase							)





class 					_SCREAMING_SNAKE_CASE  (      __UpperCamelCase ):
      def __init__(			self		,	lowerCamelCase		,	lowerCamelCase		,	lowerCamelCase = 1.0		,	):
          super().__init__()
          snake_case__						=				bit_scale
          snake_case__						=				(
              ddim_bit_scheduler_step if isinstance(lowerCamelCase		,	lowerCamelCase ) else ddpm_bit_scheduler_step
          )

          self.register_modules(unet=lowerCamelCase		,	scheduler=lowerCamelCase )



      @torch.no_grad()
      def __call__(			self		,	lowerCamelCase = 2_56		,	lowerCamelCase = 2_56		,	lowerCamelCase = 50		,	lowerCamelCase = None		,	lowerCamelCase = 1		,	lowerCamelCase = "pil"		,	lowerCamelCase = True		,	**lowerCamelCase		,	):
          snake_case__						=				torch.randn(
              (batch_size, self.unet.config.in_channels, height, width)		,	generator=lowerCamelCase		,	)
          snake_case__						=				decimal_to_bits(lowerCamelCase ) * self.bit_scale
          snake_case__						=				latents.to(self.device )

          self.scheduler.set_timesteps(lowerCamelCase )

          for t in self.progress_bar(self.scheduler.timesteps ):
              # predict the noise residual
              snake_case__						=				self.unet(lowerCamelCase		,	lowerCamelCase ).sample

              # compute the previous noisy sample x_t -> x_t-1
              snake_case__						=				self.scheduler.step(lowerCamelCase		,	lowerCamelCase		,	lowerCamelCase ).prev_sample

          snake_case__						=				bits_to_decimal(lowerCamelCase )

          if output_type == "pil":
              snake_case__						=				self.numpy_to_pil(lowerCamelCase )

          if not return_dict:
              return (image,)

          return ImagePipelineOutput(images=lowerCamelCase )



