





import unittest

import torch

from diffusers import DDIMScheduler, DDPMScheduler, UNetaDModel
from diffusers.training_utils import set_seed
from diffusers.utils.testing_utils import slow


UpperCamelCase__           =			False

class 							__SCREAMING_SNAKE_CASE       (  unittest.TestCase							):






							def 					_lowerCamelCase		(    self	,			__lowerCAmelCase=32	):
														set_seed(0	)
														UpperCamelCase__					  =      UNetaDModel(sample_size=__lowerCAmelCase	,			in_channels=3	,			out_channels=3	)
														UpperCamelCase__					  =      torch.optim.SGD(model.parameters()	,			lr=0.0001	)
														return model, optimizer







							@slow
							def 					_lowerCamelCase		(    self	):
														UpperCamelCase__					  =      """cpu"""  # ensure full determinism without setting the CUBLAS_WORKSPACE_CONFIG env variable
														UpperCamelCase__					  =      DDPMScheduler(
														    num_train_timesteps=1000	,			beta_start=0.0001	,			beta_end=0.02	,			beta_schedule="""linear"""	,			clip_sample=__lowerCAmelCase	,			)
														UpperCamelCase__					  =      DDIMScheduler(
														    num_train_timesteps=1000	,			beta_start=0.0001	,			beta_end=0.02	,			beta_schedule="""linear"""	,			clip_sample=__lowerCAmelCase	,			)

														assert ddpm_scheduler.config.num_train_timesteps == ddim_scheduler.config.num_train_timesteps

														# shared batches for DDPM and DDIM
														set_seed(0	)
														UpperCamelCase__					  =      [torch.randn((4, 3, 32, 32)	).clip(-1	,			1	).to(__lowerCAmelCase	) for _ in range(4	)]
														UpperCamelCase__					  =      [torch.randn((4, 3, 32, 32)	).to(__lowerCAmelCase	) for _ in range(4	)]
														UpperCamelCase__					  =      [torch.randint(0	,			1000	,			(4,)	).long().to(__lowerCAmelCase	) for _ in range(4	)]

														# train with a DDPM scheduler
														UpperCamelCase__  ,  UpperCamelCase__					  =      self.get_model_optimizer(resolution=32	)
														model.train().to(__lowerCAmelCase	)
														for i in range(4	):
																					optimizer.zero_grad()
																					UpperCamelCase__					  =      ddpm_scheduler.add_noise(clean_images[i]	,			noise[i]	,			timesteps[i]	)
																					UpperCamelCase__					  =      model(__lowerCAmelCase	,			timesteps[i]	).sample
																					UpperCamelCase__					  =      torch.nn.functional.mse_loss(__lowerCAmelCase	,			noise[i]	)
																					loss.backward()
																					optimizer.step()
														del model, optimizer

														# recreate the model and optimizer, and retry with DDIM
														UpperCamelCase__  ,  UpperCamelCase__					  =      self.get_model_optimizer(resolution=32	)
														model.train().to(__lowerCAmelCase	)
														for i in range(4	):
																					optimizer.zero_grad()
																					UpperCamelCase__					  =      ddim_scheduler.add_noise(clean_images[i]	,			noise[i]	,			timesteps[i]	)
																					UpperCamelCase__					  =      model(__lowerCAmelCase	,			timesteps[i]	).sample
																					UpperCamelCase__					  =      torch.nn.functional.mse_loss(__lowerCAmelCase	,			noise[i]	)
																					loss.backward()
																					optimizer.step()
														del model, optimizer

														self.assertTrue(torch.allclose(__lowerCAmelCase	,			__lowerCAmelCase	,			atol=1E-5	)	)
														self.assertTrue(torch.allclose(__lowerCAmelCase	,			__lowerCAmelCase	,			atol=1E-5	)	)











import unittest

import torch

from diffusers import VQModel
from diffusers.utils import floats_tensor, torch_device
from diffusers.utils.testing_utils import enable_full_determinism

from .test_modeling_common import ModelTesterMixin, UNetTesterMixin


enable_full_determinism()

class 							__SCREAMING_SNAKE_CASE       (  _a ,		_a ,		unittest.TestCase							):
							snake_case  :   List[str]      =							VQModel
							snake_case  :   Dict      =							"""sample"""






							@property
							def 					_lowerCamelCase		(    self	,			__lowerCAmelCase=(32, 32)	):
														UpperCamelCase__					  =      4
														UpperCamelCase__					  =      3

														UpperCamelCase__					  =      floats_tensor((batch_size, num_channels) + sizes	).to(__lowerCAmelCase	)

														return {"sample": image}






							@property
							def 					_lowerCamelCase		(    self	):
														return (3, 32, 32)






							@property
							def 					_lowerCamelCase		(    self	):
														return (3, 32, 32)






							def 					_lowerCamelCase		(    self	):
														UpperCamelCase__					  =      {
														    """block_out_channels""": [32, 64],
														    """in_channels""": 3,
														    """out_channels""": 3,
														    """down_block_types""": ["""DownEncoderBlock2D""", """DownEncoderBlock2D"""],
														    """up_block_types""": ["""UpDecoderBlock2D""", """UpDecoderBlock2D"""],
														    """latent_channels""": 3,
														}
														UpperCamelCase__					  =      self.dummy_input
														return init_dict, inputs_dict






							def 					_lowerCamelCase		(    self	):
														pass






							def 					_lowerCamelCase		(    self	):
														pass






							def 					_lowerCamelCase		(    self	):
														UpperCamelCase__  ,  UpperCamelCase__					  =      VQModel.from_pretrained("""fusing/vqgan-dummy"""	,			output_loading_info=__lowerCAmelCase	)
														self.assertIsNotNone(__lowerCAmelCase	)
														self.assertEqual(len(loading_info["""missing_keys"""]	)	,			0	)

														model.to(__lowerCAmelCase	)
														UpperCamelCase__					  =      model(**self.dummy_input	)

														assert image is not None, "Make sure output is not None"







							def 					_lowerCamelCase		(    self	):
														UpperCamelCase__					  =      VQModel.from_pretrained("""fusing/vqgan-dummy"""	)
														model.to(__lowerCAmelCase	).eval()

														torch.manual_seed(0	)
														if torch.cuda.is_available():
																					torch.cuda.manual_seed_all(0	)

														UpperCamelCase__					  =      torch.randn(1	,			model.config.in_channels	,			model.config.sample_size	,			model.config.sample_size	)
														UpperCamelCase__					  =      image.to(__lowerCAmelCase	)
														with torch.no_grad():
																					UpperCamelCase__					  =      model(__lowerCAmelCase	).sample

														UpperCamelCase__					  =      output[0, -1, -3:, -3:].flatten().cpu()
														# fmt: off
														UpperCamelCase__					  =      torch.tensor([-0.0153, -0.4044, -0.1880, -0.5161, -0.2418, -0.4072, -0.1612, -0.0633, -0.0143]	)
														# fmt: on
														self.assertTrue(torch.allclose(__lowerCAmelCase	,			__lowerCAmelCase	,			atol=1E-3	)	)




