


"""simple docstring"""




import gc
import random
import unittest

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

from diffusers import AutoencoderKL, CycleDiffusionPipeline, DDIMScheduler, 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, skip_mps

from ..pipeline_params import (
    IMAGE_TO_IMAGE_IMAGE_PARAMS,
    TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS,
    TEXT_GUIDED_IMAGE_VARIATION_PARAMS,
)
from ..test_pipelines_common import PipelineLatentTesterMixin, PipelineTesterMixin


enable_full_determinism()
class 					__lowerCAmelCase					(   _SCREAMING_SNAKE_CASE     ,			_SCREAMING_SNAKE_CASE     ,			unittest.TestCase       ):


		'''simple docstring'''



		__UpperCAmelCase		:			Any										=				CycleDiffusionPipeline
		__UpperCAmelCase		:			Tuple										=				TEXT_GUIDED_IMAGE_VARIATION_PARAMS - {
		    "negative_prompt",
		    "height",
		    "width",
		    "negative_prompt_embeds",
		}
		__UpperCAmelCase		:			Union[str, Any]										=				PipelineTesterMixin.required_optional_params - {"latents"}
		__UpperCAmelCase		:			List[Any]										=				TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS.union({'source_prompt'}       )
		__UpperCAmelCase		:			Any										=				IMAGE_TO_IMAGE_IMAGE_PARAMS
		__UpperCAmelCase		:			Optional[Any]										=				IMAGE_TO_IMAGE_IMAGE_PARAMS

		def 			__UpperCAmelCase       (  self      ):
				torch.manual_seed(0      )
				__a           =			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							, )
				__a           =			DDIMScheduler(
				    beta_start=0.0_0085							, beta_end=0.012							, beta_schedule='''scaled_linear'''							, num_train_timesteps=1_000							, clip_sample=A_							, set_alpha_to_one=A_							, )
				torch.manual_seed(0      )
				__a           =			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      )
				__a           =			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=1_000							, )
				__a           =			CLIPTextModel(A_      )
				__a           =			CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''      )

				__a           =			{
				    '''unet''': unet,
				    '''scheduler''': scheduler,
				    '''vae''': vae,
				    '''text_encoder''': text_encoder,
				    '''tokenizer''': tokenizer,
				    '''safety_checker''': None,
				    '''feature_extractor''': None,
				}
				return components

		def 			__UpperCAmelCase       (  self							, _a							, _a=0      ):
				__a           =			floats_tensor((1, 3, 32, 32)							, rng=random.Random(A_      )      ).to(A_      )
				__a           =			image / 2 + 0.5
				if str(A_      ).startswith('''mps'''      ):
						__a           =			torch.manual_seed(A_      )
				else:
						__a           =			torch.Generator(device=A_      ).manual_seed(A_      )
				__a           =			{
				    '''prompt''': '''An astronaut riding an elephant''',
				    '''source_prompt''': '''An astronaut riding a horse''',
				    '''image''': image,
				    '''generator''': generator,
				    '''num_inference_steps''': 2,
				    '''eta''': 0.1,
				    '''strength''': 0.8,
				    '''guidance_scale''': 3,
				    '''source_guidance_scale''': 1,
				    '''output_type''': '''numpy''',
				}
				return inputs

		def 			__UpperCAmelCase       (  self      ):
				__a           =			'''cpu'''  # ensure determinism for the device-dependent torch.Generator

				__a           =			self.get_dummy_components()
				__a           =			CycleDiffusionPipeline(**A_      )
				__a           =			pipe.to(A_      )
				pipe.set_progress_bar_config(disable=A_      )

				__a           =			self.get_dummy_inputs(A_      )
				__a           =			pipe(**A_      )
				__a           =			output.images

				__a           =			images[0, -3:, -3:, -1]

				assert images.shape == (1, 32, 32, 3)
				__a           =			np.array([0.4459, 0.4943, 0.4544, 0.6643, 0.5474, 0.4327, 0.5701, 0.5959, 0.5179]      )

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

		@unittest.skipIf(torch_device != '''cuda'''							, '''This test requires a GPU'''      )
		def 			__UpperCAmelCase       (  self      ):
				__a           =			self.get_dummy_components()
				for name, module in components.items():
						if hasattr(A_							, '''half'''      ):
								__a           =			module.half()
				__a           =			CycleDiffusionPipeline(**A_      )
				__a           =			pipe.to(A_      )
				pipe.set_progress_bar_config(disable=A_      )

				__a           =			self.get_dummy_inputs(A_      )
				__a           =			pipe(**A_      )
				__a           =			output.images

				__a           =			images[0, -3:, -3:, -1]

				assert images.shape == (1, 32, 32, 3)
				__a           =			np.array([0.3506, 0.4543, 0.446, 0.4575, 0.5195, 0.4155, 0.5273, 0.518, 0.4116]      )

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

		@skip_mps
		def 			__UpperCAmelCase       (  self      ):
				return super().test_save_load_local()

		@unittest.skip('''non-deterministic pipeline'''      )
		def 			__UpperCAmelCase       (  self      ):
				return super().test_inference_batch_single_identical()

		@skip_mps
		def 			__UpperCAmelCase       (  self      ):
				return super().test_dict_tuple_outputs_equivalent()

		@skip_mps
		def 			__UpperCAmelCase       (  self      ):
				return super().test_save_load_optional_components()





		@skip_mps
		def 			__UpperCAmelCase       (  self      ):
				return super().test_attention_slicing_forward_pass()







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


		'''simple docstring'''



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

		def 			__UpperCAmelCase       (  self      ):
				__a           =			load_image(
				    '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'''
				    '''/cycle-diffusion/black_colored_car.png'''      )
				__a           =			load_numpy(
				    '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/cycle-diffusion/blue_colored_car_fp16.npy'''      )
				__a           =			init_image.resize((512, 512)      )

				__a           =			'''CompVis/stable-diffusion-v1-4'''
				__a           =			DDIMScheduler.from_pretrained(A_							, subfolder='''scheduler'''      )
				__a           =			CycleDiffusionPipeline.from_pretrained(
				    A_							, scheduler=A_							, safety_checker=A_							, torch_dtype=torch.floataa							, revision='''fp16'''      )

				pipe.to(A_      )
				pipe.set_progress_bar_config(disable=A_      )
				pipe.enable_attention_slicing()

				__a           =			'''A black colored car'''
				__a           =			'''A blue colored car'''

				__a           =			torch.manual_seed(0      )
				__a           =			pipe(
				    prompt=A_							, source_prompt=A_							, image=A_							, num_inference_steps=100							, eta=0.1							, strength=0.85							, guidance_scale=3							, source_guidance_scale=1							, generator=A_							, output_type='''np'''							, )
				__a           =			output.images

				# the values aren't exactly equal, but the images look the same visually
				assert np.abs(image - expected_image      ).max() < 5E-1





		def 			__UpperCAmelCase       (  self      ):
				__a           =			load_image(
				    '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'''
				    '''/cycle-diffusion/black_colored_car.png'''      )
				__a           =			load_numpy(
				    '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/cycle-diffusion/blue_colored_car.npy'''      )
				__a           =			init_image.resize((512, 512)      )

				__a           =			'''CompVis/stable-diffusion-v1-4'''
				__a           =			DDIMScheduler.from_pretrained(A_							, subfolder='''scheduler'''      )
				__a           =			CycleDiffusionPipeline.from_pretrained(A_							, scheduler=A_							, safety_checker=A_      )

				pipe.to(A_      )
				pipe.set_progress_bar_config(disable=A_      )
				pipe.enable_attention_slicing()

				__a           =			'''A black colored car'''
				__a           =			'''A blue colored car'''

				__a           =			torch.manual_seed(0      )
				__a           =			pipe(
				    prompt=A_							, source_prompt=A_							, image=A_							, num_inference_steps=100							, eta=0.1							, strength=0.85							, guidance_scale=3							, source_guidance_scale=1							, generator=A_							, output_type='''np'''							, )
				__a           =			output.images

				assert np.abs(image - expected_image      ).max() < 2E-2








import numpy as np
import torch
import tqdm

from ...models.unet_ad import UNetaDModel
from ...pipelines import DiffusionPipeline
from ...utils import randn_tensor
from ...utils.dummy_pt_objects import DDPMScheduler


class 				lowercase					(		_SCREAMING_SNAKE_CASE   ):



		def __init__( self				,      A_				,      A_				,      A_				,      A_				,      )  ->					Optional[int]:
				"""simple docstring"""
				super().__init__()
				UpperCamelCase		=     value_function
				UpperCamelCase		=     unet
				UpperCamelCase		=     scheduler
				UpperCamelCase		=     env
				UpperCamelCase		=     env.get_dataset()
				UpperCamelCase		=     {}
				for key in self.data.keys():
						try:
								UpperCamelCase		=     self.data[key].mean()
						except:  # noqa: E722
								pass
				UpperCamelCase		=     {}
				for key in self.data.keys():
						try:
								UpperCamelCase		=     self.data[key].std()
						except:  # noqa: E722
								pass
				UpperCamelCase		=     env.observation_space.shape[0]
				UpperCamelCase		=     env.action_space.shape[0]



		def __UpperCamelCase ( self				,      A_				,      A_       )  ->					Optional[Any]:
				"""simple docstring"""
				return (x_in - self.means[key]) / self.stds[key]



		def __UpperCamelCase ( self				,      A_				,      A_       )  ->					Union[str, Any]:
				"""simple docstring"""
				return x_in * self.stds[key] + self.means[key]



		def __UpperCamelCase ( self				,      A_       )  ->					Any:
				"""simple docstring"""
				if type(A_       ) is dict:
						return {k: self.to_torch(A_       ) for k, v in x_in.items()}
				elif torch.is_tensor(A_       ):
						return x_in.to(self.unet.device       )
				return torch.tensor(A_				,      device=self.unet.device       )



		def __UpperCamelCase ( self				,      A_				,      A_				,      A_       )  ->					Optional[int]:
				"""simple docstring"""
				for key, val in cond.items():
						UpperCamelCase		=     val.clone()
				return x_in



		def __UpperCamelCase ( self				,      A_				,      A_				,      A_				,      A_       )  ->					Union[str, Any]:
				"""simple docstring"""
				UpperCamelCase		=     x.shape[0]
				UpperCamelCase		=     None
				for i in tqdm.tqdm(self.scheduler.timesteps       ):
						# create batch of timesteps to pass into model
						UpperCamelCase		=     torch.full((batch_size,)				,      A_				,      device=self.unet.device				,      dtype=torch.long       )
						for _ in range(A_       ):
								with torch.enable_grad():
										x.requires_grad_()

										# permute to match dimension for pre-trained models
										UpperCamelCase		=     self.value_function(x.permute(0				,      2				,      1       )				,      A_       ).sample
										UpperCamelCase		=     torch.autograd.grad([y.sum()]				,      [x]       )[0]

										UpperCamelCase		=     self.scheduler._get_variance(A_       )
										UpperCamelCase		=     torch.exp(0.5 * posterior_variance       )
										UpperCamelCase		=     model_std * grad

								UpperCamelCase		=     0
								UpperCamelCase		=     x.detach()
								UpperCamelCase		=     x + scale * grad
								UpperCamelCase		=     self.reset_xa(A_				,      A_				,      self.action_dim       )

						UpperCamelCase		=     self.unet(x.permute(0				,      2				,      1       )				,      A_       ).sample.permute(0				,      2				,      1       )

						# TODO: verify deprecation of this kwarg
						UpperCamelCase		=     self.scheduler.step(A_				,      A_				,      A_				,      predict_epsilon=A_       )['prev_sample']

						# apply conditions to the trajectory (set the initial state)
						UpperCamelCase		=     self.reset_xa(A_				,      A_				,      self.action_dim       )
						UpperCamelCase		=     self.to_torch(A_       )
				return x, y



		def __call__( self				,      A_				,      A_=64				,      A_=32				,      A_=2				,      A_=0.1       )  ->					List[str]:
				"""simple docstring"""
				# normalize the observations and create  batch dimension
				UpperCamelCase		=     self.normalize(A_				,      'observations'       )
				UpperCamelCase		=     obs[None].repeat(A_				,      axis=0       )

				UpperCamelCase		=     {0: self.to_torch(A_       )}
				UpperCamelCase		=     (batch_size, planning_horizon, self.state_dim + self.action_dim)

				# generate initial noise and apply our conditions (to make the trajectories start at current state)
				UpperCamelCase		=     randn_tensor(A_				,      device=self.unet.device       )
				UpperCamelCase		=     self.reset_xa(A_				,      A_				,      self.action_dim       )
				UpperCamelCase		=     self.to_torch(A_       )

				# run the diffusion process
				UpperCamelCase	,  UpperCamelCase		=     self.run_diffusion(A_				,      A_				,      A_				,      A_       )

				# sort output trajectories by value
				UpperCamelCase		=     y.argsort(0				,      descending=A_       ).squeeze()
				UpperCamelCase		=     x[sorted_idx]
				UpperCamelCase		=     sorted_values[:, :, : self.action_dim]
				UpperCamelCase		=     actions.detach().cpu().numpy()
				UpperCamelCase		=     self.de_normalize(A_				,      key='actions'       )

				# select the action with the highest value
				if y is not None:
						UpperCamelCase		=     0
				else:
						# if we didn't run value guiding, select a random action
						UpperCamelCase		=     np.random.randint(0				,      A_       )

				UpperCamelCase		=     denorm_actions[selected_index, 0]
				return denorm_actions
