


"""simple docstring"""






import gc
import random
import unittest

import numpy as np
import torch
from PIL import Image
from transformers import CLIPImageProcessor, CLIPVisionConfig, CLIPVisionModelWithProjection

from diffusers import (
    AutoencoderKL,
    DPMSolverMultistepScheduler,
    PNDMScheduler,
    StableDiffusionImageVariationPipeline,
    UNetaDConditionModel,
)
from diffusers.utils import floats_tensor, load_image, load_numpy, nightly, slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu

from ..pipeline_params import IMAGE_VARIATION_BATCH_PARAMS, IMAGE_VARIATION_PARAMS
from ..test_pipelines_common import PipelineKarrasSchedulerTesterMixin, PipelineLatentTesterMixin, PipelineTesterMixin


enable_full_determinism()


class       lowerCamelCase_		(      lowercase							,	lowercase							,	lowercase							,	unittest.TestCase ):





				"""simple docstring"""







				_lowerCAmelCase   :   Tuple   				=       StableDiffusionImageVariationPipeline
				_lowerCAmelCase   :   Dict   				=       IMAGE_VARIATION_PARAMS
				_lowerCAmelCase   :   Optional[int]   				=       IMAGE_VARIATION_BATCH_PARAMS
				_lowerCAmelCase   :   str   				=       frozenset([] )
				# TO-DO: update image_params once pipeline is refactored with VaeImageProcessor.preprocess
				_lowerCAmelCase   :   Any   				=       frozenset([] )



				def 	lowerCAmelCase__		(    self  ):
								torch.manual_seed(0  )
								SCREAMING_SNAKE_CASE__  				=					UNetaDConditionModel(
								    block_out_channels=(32, 64)							,		layers_per_block=2							,		sample_size=32							,		in_channels=4							,		out_channels=4							,		down_block_types=("DownBlock2D", "CrossAttnDownBlock2D")							,		up_block_types=("CrossAttnUpBlock2D", "UpBlock2D")							,		cross_attention_dim=32							,		)
								SCREAMING_SNAKE_CASE__  				=					PNDMScheduler(skip_prk_steps=UpperCAmelCase__  )
								torch.manual_seed(0  )
								SCREAMING_SNAKE_CASE__  				=					AutoencoderKL(
								    block_out_channels=[32, 64]							,		in_channels=3							,		out_channels=3							,		down_block_types=["DownEncoderBlock2D", "DownEncoderBlock2D"]							,		up_block_types=["UpDecoderBlock2D", "UpDecoderBlock2D"]							,		latent_channels=4							,		)
								torch.manual_seed(0  )
								SCREAMING_SNAKE_CASE__  				=					CLIPVisionConfig(
								    hidden_size=32							,		projection_dim=32							,		intermediate_size=37							,		layer_norm_eps=1e-05							,		num_attention_heads=4							,		num_hidden_layers=5							,		image_size=32							,		patch_size=4							,		)
								SCREAMING_SNAKE_CASE__  				=					CLIPVisionModelWithProjection(UpperCAmelCase__  )
								SCREAMING_SNAKE_CASE__  				=					CLIPImageProcessor(crop_size=32							,		size=32  )

								SCREAMING_SNAKE_CASE__  				=					{
								    "unet": unet,
								    "scheduler": scheduler,
								    "vae": vae,
								    "image_encoder": image_encoder,
								    "feature_extractor": feature_extractor,
								    "safety_checker": None,
								}
								return components



				def 	lowerCAmelCase__		(    self							,		UpperCAmelCase__							,		UpperCAmelCase__=0  ):
								SCREAMING_SNAKE_CASE__  				=					floats_tensor((1, 3, 32, 32)							,		rng=random.Random(UpperCAmelCase__  )  )
								SCREAMING_SNAKE_CASE__  				=					image.cpu().permute(0							,		2							,		3							,		1  )[0]
								SCREAMING_SNAKE_CASE__  				=					Image.fromarray(np.uinta(UpperCAmelCase__  )  ).convert("RGB"  ).resize((32, 32)  )
								if str(UpperCAmelCase__  ).startswith("mps"  ):
												SCREAMING_SNAKE_CASE__  				=					torch.manual_seed(UpperCAmelCase__  )
								else:
												SCREAMING_SNAKE_CASE__  				=					torch.Generator(device=UpperCAmelCase__  ).manual_seed(UpperCAmelCase__  )
								SCREAMING_SNAKE_CASE__  				=					{
								    "image": image,
								    "generator": generator,
								    "num_inference_steps": 2,
								    "guidance_scale": 6.0,
								    "output_type": "numpy",
								}
								return inputs



				def 	lowerCAmelCase__		(    self  ):
								SCREAMING_SNAKE_CASE__  				=					"cpu"  # ensure determinism for the device-dependent torch.Generator
								SCREAMING_SNAKE_CASE__  				=					self.get_dummy_components()
								SCREAMING_SNAKE_CASE__  				=					StableDiffusionImageVariationPipeline(**UpperCAmelCase__  )
								SCREAMING_SNAKE_CASE__  				=					sd_pipe.to(UpperCAmelCase__  )
								sd_pipe.set_progress_bar_config(disable=UpperCAmelCase__  )

								SCREAMING_SNAKE_CASE__  				=					self.get_dummy_inputs(UpperCAmelCase__  )
								SCREAMING_SNAKE_CASE__  				=					sd_pipe(**UpperCAmelCase__  ).images
								SCREAMING_SNAKE_CASE__  				=					image[0, -3:, -3:, -1]

								assert image.shape == (1, 64, 64, 3)
								SCREAMING_SNAKE_CASE__  				=					np.array([0.5_239, 0.5_723, 0.4_796, 0.5_049, 0.5_550, 0.4_685, 0.5_329, 0.4_891, 0.4_921]  )

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



				def 	lowerCAmelCase__		(    self  ):
								SCREAMING_SNAKE_CASE__  				=					"cpu"  # ensure determinism for the device-dependent torch.Generator
								SCREAMING_SNAKE_CASE__  				=					self.get_dummy_components()
								SCREAMING_SNAKE_CASE__  				=					StableDiffusionImageVariationPipeline(**UpperCAmelCase__  )
								SCREAMING_SNAKE_CASE__  				=					sd_pipe.to(UpperCAmelCase__  )
								sd_pipe.set_progress_bar_config(disable=UpperCAmelCase__  )

								SCREAMING_SNAKE_CASE__  				=					self.get_dummy_inputs(UpperCAmelCase__  )
								SCREAMING_SNAKE_CASE__  				=					2 * [inputs["image"]]
								SCREAMING_SNAKE_CASE__  				=					sd_pipe(**UpperCAmelCase__  )

								SCREAMING_SNAKE_CASE__  				=					output.images

								SCREAMING_SNAKE_CASE__  				=					image[-1, -3:, -3:, -1]

								assert image.shape == (2, 64, 64, 3)
								SCREAMING_SNAKE_CASE__  				=					np.array([0.6_892, 0.5_637, 0.5_836, 0.5_771, 0.6_254, 0.6_409, 0.5_580, 0.5_569, 0.5_289]  )

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






				def 	lowerCAmelCase__		(    self  ):
								super().test_inference_batch_single_identical(expected_max_diff=3e-3  )


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





				"""simple docstring"""







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



				def 	lowerCAmelCase__		(    self							,		UpperCAmelCase__							,		UpperCAmelCase__="cpu"							,		UpperCAmelCase__=torch.floataa							,		UpperCAmelCase__=0  ):
								SCREAMING_SNAKE_CASE__  				=					torch.Generator(device=UpperCAmelCase__  ).manual_seed(UpperCAmelCase__  )
								SCREAMING_SNAKE_CASE__  				=					load_image(
								    "https://huggingface.co/datasets/diffusers/test-arrays/resolve/main"
								    "/stable_diffusion_imgvar/input_image_vermeer.png"  )
								SCREAMING_SNAKE_CASE__  				=					np.random.RandomState(UpperCAmelCase__  ).standard_normal((1, 4, 64, 64)  )
								SCREAMING_SNAKE_CASE__  				=					torch.from_numpy(UpperCAmelCase__  ).to(device=UpperCAmelCase__							,		dtype=UpperCAmelCase__  )
								SCREAMING_SNAKE_CASE__  				=					{
								    "image": init_image,
								    "latents": latents,
								    "generator": generator,
								    "num_inference_steps": 3,
								    "guidance_scale": 7.5,
								    "output_type": "numpy",
								}
								return inputs



				def 	lowerCAmelCase__		(    self  ):
								SCREAMING_SNAKE_CASE__  				=					StableDiffusionImageVariationPipeline.from_pretrained(
								    "lambdalabs/sd-image-variations-diffusers"							,		safety_checker=UpperCAmelCase__  )
								SCREAMING_SNAKE_CASE__  				=					sd_pipe.to(UpperCAmelCase__  )
								sd_pipe.set_progress_bar_config(disable=UpperCAmelCase__  )

								SCREAMING_SNAKE_CASE__  				=					self.get_inputs(UpperCAmelCase__  )
								SCREAMING_SNAKE_CASE__  				=					sd_pipe(**UpperCAmelCase__  ).images
								SCREAMING_SNAKE_CASE__  				=					image[0, -3:, -3:, -1].flatten()

								assert image.shape == (1, 512, 512, 3)
								SCREAMING_SNAKE_CASE__  				=					np.array([0.84_491, 0.90_789, 0.75_708, 0.78_734, 0.83_485, 0.70_099, 0.66_938, 0.68_727, 0.61_379]  )
								assert np.abs(image_slice - expected_slice  ).max() < 6e-3



				def 	lowerCAmelCase__		(    self  ):
								SCREAMING_SNAKE_CASE__  				=					0

								def callback_fn(UpperCAmelCase__							,		UpperCAmelCase__							,		UpperCAmelCase__  ) -> None:
												SCREAMING_SNAKE_CASE__  				=					True
												nonlocal number_of_steps
												number_of_steps += 1
												if step == 1:
																SCREAMING_SNAKE_CASE__  				=					latents.detach().cpu().numpy()
																assert latents.shape == (1, 4, 64, 64)
																SCREAMING_SNAKE_CASE__  				=					latents[0, -3:, -3:, -1]
																SCREAMING_SNAKE_CASE__  				=					np.array(
																    [-0.1_621, 0.2_837, -0.7_979, -0.1_221, -1.3_057, 0.7_681, -2.1_191, 0.0_464, 1.6_309]  )

																assert np.abs(latents_slice.flatten() - expected_slice  ).max() < 5e-2
												elif step == 2:
																SCREAMING_SNAKE_CASE__  				=					latents.detach().cpu().numpy()
																assert latents.shape == (1, 4, 64, 64)
																SCREAMING_SNAKE_CASE__  				=					latents[0, -3:, -3:, -1]
																SCREAMING_SNAKE_CASE__  				=					np.array([0.6_299, 1.7_500, 1.1_992, -2.1_582, -1.8_994, 0.7_334, -0.7_090, 1.0_137, 1.5_273]  )

																assert np.abs(latents_slice.flatten() - expected_slice  ).max() < 5e-2

								SCREAMING_SNAKE_CASE__  				=					False

								SCREAMING_SNAKE_CASE__  				=					StableDiffusionImageVariationPipeline.from_pretrained(
								    "fusing/sd-image-variations-diffusers"							,		safety_checker=UpperCAmelCase__							,		torch_dtype=torch.floataa							,		)
								pipe.to(UpperCAmelCase__  )
								pipe.set_progress_bar_config(disable=UpperCAmelCase__  )
								pipe.enable_attention_slicing()

								SCREAMING_SNAKE_CASE__  				=					self.get_inputs(UpperCAmelCase__							,		dtype=torch.floataa  )
								pipe(**UpperCAmelCase__							,		callback=UpperCAmelCase__							,		callback_steps=1  )
								assert callback_fn.has_been_called
								assert number_of_steps == inputs["num_inference_steps"]






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

								SCREAMING_SNAKE_CASE__  				=					"fusing/sd-image-variations-diffusers"
								SCREAMING_SNAKE_CASE__  				=					StableDiffusionImageVariationPipeline.from_pretrained(
								    UpperCAmelCase__							,		safety_checker=UpperCAmelCase__							,		torch_dtype=torch.floataa  )
								SCREAMING_SNAKE_CASE__  				=					pipe.to(UpperCAmelCase__  )
								pipe.set_progress_bar_config(disable=UpperCAmelCase__  )
								pipe.enable_attention_slicing(1  )
								pipe.enable_sequential_cpu_offload()

								SCREAMING_SNAKE_CASE__  				=					self.get_inputs(UpperCAmelCase__							,		dtype=torch.floataa  )
								SCREAMING_SNAKE_CASE__  				=					pipe(**UpperCAmelCase__  )

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


@nightly
@require_torch_gpu
class       lowerCamelCase_		(      unittest.TestCase ):





				"""simple docstring"""







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



				def 	lowerCAmelCase__		(    self							,		UpperCAmelCase__							,		UpperCAmelCase__="cpu"							,		UpperCAmelCase__=torch.floataa							,		UpperCAmelCase__=0  ):
								SCREAMING_SNAKE_CASE__  				=					torch.Generator(device=UpperCAmelCase__  ).manual_seed(UpperCAmelCase__  )
								SCREAMING_SNAKE_CASE__  				=					load_image(
								    "https://huggingface.co/datasets/diffusers/test-arrays/resolve/main"
								    "/stable_diffusion_imgvar/input_image_vermeer.png"  )
								SCREAMING_SNAKE_CASE__  				=					np.random.RandomState(UpperCAmelCase__  ).standard_normal((1, 4, 64, 64)  )
								SCREAMING_SNAKE_CASE__  				=					torch.from_numpy(UpperCAmelCase__  ).to(device=UpperCAmelCase__							,		dtype=UpperCAmelCase__  )
								SCREAMING_SNAKE_CASE__  				=					{
								    "image": init_image,
								    "latents": latents,
								    "generator": generator,
								    "num_inference_steps": 50,
								    "guidance_scale": 7.5,
								    "output_type": "numpy",
								}
								return inputs



				def 	lowerCAmelCase__		(    self  ):
								SCREAMING_SNAKE_CASE__  				=					StableDiffusionImageVariationPipeline.from_pretrained("fusing/sd-image-variations-diffusers"  )
								sd_pipe.to(UpperCAmelCase__  )
								sd_pipe.set_progress_bar_config(disable=UpperCAmelCase__  )

								SCREAMING_SNAKE_CASE__  				=					self.get_inputs(UpperCAmelCase__  )
								SCREAMING_SNAKE_CASE__  				=					sd_pipe(**UpperCAmelCase__  ).images[0]

								SCREAMING_SNAKE_CASE__  				=					load_numpy(
								    "https://huggingface.co/datasets/diffusers/test-arrays/resolve/main"
								    "/stable_diffusion_imgvar/lambdalabs_variations_pndm.npy"  )
								SCREAMING_SNAKE_CASE__  				=					np.abs(expected_image - image  ).max()
								assert max_diff < 1e-3






				def 	lowerCAmelCase__		(    self  ):
								SCREAMING_SNAKE_CASE__  				=					StableDiffusionImageVariationPipeline.from_pretrained("fusing/sd-image-variations-diffusers"  )
								SCREAMING_SNAKE_CASE__  				=					DPMSolverMultistepScheduler.from_config(sd_pipe.scheduler.config  )
								sd_pipe.to(UpperCAmelCase__  )
								sd_pipe.set_progress_bar_config(disable=UpperCAmelCase__  )

								SCREAMING_SNAKE_CASE__  				=					self.get_inputs(UpperCAmelCase__  )
								SCREAMING_SNAKE_CASE__  				=					25
								SCREAMING_SNAKE_CASE__  				=					sd_pipe(**UpperCAmelCase__  ).images[0]

								SCREAMING_SNAKE_CASE__  				=					load_numpy(
								    "https://huggingface.co/datasets/diffusers/test-arrays/resolve/main"
								    "/stable_diffusion_imgvar/lambdalabs_variations_dpm_multi.npy"  )
								SCREAMING_SNAKE_CASE__  				=					np.abs(expected_image - image  ).max()
								assert max_diff < 1e-3




"""simple docstring"""






from __future__ import annotations

from collections.abc import Generator


def __lowercase  (   lowerCamelCase_	:   int	):
				if not isinstance(lowerCamelCase_    ,       lowerCamelCase_	) or n < 1:
								raise Exception("Sequence only defined for positive integers"	)

				yield n
				while n != 1:
								if n % 2 == 0:
												n //= 2
								else:
												SCREAMING_SNAKE_CASE__  				=					3 * n + 1
								yield n


def __lowercase  (   ):
				SCREAMING_SNAKE_CASE__  				=					43
				SCREAMING_SNAKE_CASE__  				=					tuple(collatz_sequence(lowerCamelCase_	)	)
				print(lowerCamelCase_	)
				print(F'''Collatz sequence from {n} took {len(lowerCamelCase_	)} steps.'''	)


if __name__ == "__main__":
	main()
