code: 
  
  
  """simple docstring"""
  import random
  import unittest
  
  import numpy as np
  import torch
  
  from diffusers import DDIMScheduler, LDMSuperResolutionPipeline, UNetaDModel, VQModel
  from diffusers.utils import PIL_INTERPOLATION, floats_tensor, load_image, slow, torch_device
  from diffusers.utils.testing_utils import enable_full_determinism, require_torch
  
  
  enable_full_determinism()
  
  
  
  
  class   lowercase__							(				unittest.TestCase							):
  
  
  
  
  
     @property
     def 		UpperCamelCase_     (       self)     ->     List[str]:
        _lowerCamelCase  :				Union[str, Any]      =		1
        _lowerCamelCase  :				Optional[int]      =		3
        _lowerCamelCase  :				List[str]      =		(32, 32)
  
        _lowerCamelCase  :				str      =		floats_tensor((batch_size, num_channels) + sizes			,      rng=random.Random(0)).to(A_)
        return image
  
  
  
  
  
     @property
     def 		UpperCamelCase_     (       self)     ->     Tuple:
        torch.manual_seed(0)
        _lowerCamelCase  :				List[str]      =		UNetaDModel(
            block_out_channels=(32, 64)			,      layers_per_block=2			,      sample_size=32			,      in_channels=6			,      out_channels=3			,      down_block_types=("""DownBlock2D""", """AttnDownBlock2D""")			,      up_block_types=("""AttnUpBlock2D""", """UpBlock2D""")			,      )
        return model
  
  
  
  
  
     @property
     def 		UpperCamelCase_     (       self)     ->     List[Any]:
        torch.manual_seed(0)
        _lowerCamelCase  :				List[Any]      =		VQModel(
            block_out_channels=[32, 64]			,      in_channels=3			,      out_channels=3			,      down_block_types=["""DownEncoderBlock2D""", """DownEncoderBlock2D"""]			,      up_block_types=["""UpDecoderBlock2D""", """UpDecoderBlock2D"""]			,      latent_channels=3			,      )
        return model
  
  
  
  
  
     def 		UpperCamelCase_     (       self)     ->     Union[str, Any]:
        _lowerCamelCase  :				List[str]      =		"""cpu"""
        _lowerCamelCase  :				Union[str, Any]      =		self.dummy_uncond_unet
        _lowerCamelCase  :				int      =		DDIMScheduler()
        _lowerCamelCase  :				List[Any]      =		self.dummy_vq_model
  
        _lowerCamelCase  :				Optional[int]      =		LDMSuperResolutionPipeline(unet=A_			,      vqvae=A_			,      scheduler=A_)
        ldm.to(A_)
        ldm.set_progress_bar_config(disable=A_)
  
        _lowerCamelCase  :				Dict      =		self.dummy_image.to(A_)
  
        _lowerCamelCase  :				Optional[Any]      =		torch.Generator(device=A_).manual_seed(0)
        _lowerCamelCase  :				Dict      =		ldm(image=A_			,      generator=A_			,      num_inference_steps=2			,      output_type="""numpy""").images
  
        _lowerCamelCase  :				Union[str, Any]      =		image[0, -3:, -3:, -1]
  
        assert image.shape == (1, 64, 64, 3)
        _lowerCamelCase  :				Optional[Any]      =		np.array([0.86_78, 0.82_45, 0.63_81, 0.68_30, 0.43_85, 0.55_99, 0.46_41, 0.62_01, 0.51_50])
  
        assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-2
  
  
  
  
  
     @unittest.skipIf(torch_device != """cuda"""			,      """This test requires a GPU""")
     def 		UpperCamelCase_     (       self)     ->     Tuple:
        _lowerCamelCase  :				Any      =		self.dummy_uncond_unet
        _lowerCamelCase  :				List[str]      =		DDIMScheduler()
        _lowerCamelCase  :				Union[str, Any]      =		self.dummy_vq_model
  
        # put models in fp16
        _lowerCamelCase  :				Tuple      =		unet.half()
        _lowerCamelCase  :				str      =		vqvae.half()
  
        _lowerCamelCase  :				int      =		LDMSuperResolutionPipeline(unet=A_			,      vqvae=A_			,      scheduler=A_)
        ldm.to(A_)
        ldm.set_progress_bar_config(disable=A_)
  
        _lowerCamelCase  :				Optional[int]      =		self.dummy_image.to(A_)
  
        _lowerCamelCase  :				int      =		ldm(A_			,      num_inference_steps=2			,      output_type="""numpy""").images
  
        assert image.shape == (1, 64, 64, 3)
  
  
  
  
  
  
  
  @slow
  @require_torch
  class   lowercase__							(				unittest.TestCase							):
  
  
  
  
  
     def 		UpperCamelCase_     (       self)     ->     Optional[int]:
        _lowerCamelCase  :				List[str]      =		load_image(
            """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"""
            """/vq_diffusion/teddy_bear_pool.png""")
        _lowerCamelCase  :				Union[str, Any]      =		init_image.resize((64, 64)			,      resample=PIL_INTERPOLATION["""lanczos"""])
  
        _lowerCamelCase  :				Union[str, Any]      =		LDMSuperResolutionPipeline.from_pretrained("""duongna/ldm-super-resolution"""			,      device_map="""auto""")
        ldm.set_progress_bar_config(disable=A_)
  
        _lowerCamelCase  :				str      =		torch.manual_seed(0)
        _lowerCamelCase  :				Optional[int]      =		ldm(image=A_			,      generator=A_			,      num_inference_steps=20			,      output_type="""numpy""").images
  
        _lowerCamelCase  :				Union[str, Any]      =		image[0, -3:, -3:, -1]
  
        assert image.shape == (1, 256, 256, 3)
        _lowerCamelCase  :				str      =		np.array([0.76_44, 0.76_79, 0.76_42, 0.76_33, 0.76_66, 0.75_60, 0.74_25, 0.72_57, 0.69_07])
  
        assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-2
  
  
style_context: 
  
  
  """simple docstring"""
  import torch
  import trajectory.utils as utils
  
  from transformers import TrajectoryTransformerModel
  
  
  
  
  class   lowercase__							(				utils.Parser							):
     __UpperCAmelCase       					=				"halfcheetah-medium-expert-v2"
     __UpperCAmelCase       					=				"config.offline"
  
  
  
  
  
  
  
  def 							_snake_case			(      __snake_case       :			Any	,   __snake_case       :			Optional[Any]	,   __snake_case       :			Optional[int]	,   __snake_case       :			List[str]	,   __snake_case       :			Union[str, Any]      ):
  
  
  
  
     """simple docstring"""
  
  
  
  
  
  
  
     _lowerCamelCase      ,						_lowerCamelCase  :				int      =		utils.load_model(__snake_case	,   __snake_case	,   __snake_case	,   epoch=__snake_case	,   device=__snake_case      )
     _lowerCamelCase  :				List[str]      =		TrajectoryTransformerModel(gpt.config      )
  
     trajectory_transformer.tok_emb.load_state_dict(gpt.tok_emb.state_dict()      )
     _lowerCamelCase  :				str      =		gpt.pos_emb
     trajectory_transformer.drop.load_state_dict(gpt.drop.state_dict()      )
     trajectory_transformer.ln_f.load_state_dict(gpt.ln_f.state_dict()      )
     trajectory_transformer.head.load_state_dict(gpt.head.state_dict()      )
  
     for i, block in enumerate(gpt.blocks      ):
        trajectory_transformer.blocks[i].lna.load_state_dict(gpt.blocks[i].lna.state_dict()      )
        trajectory_transformer.blocks[i].lna.load_state_dict(gpt.blocks[i].lna.state_dict()      )
        trajectory_transformer.blocks[i].attn.load_state_dict(gpt.blocks[i].attn.state_dict()      )
  
        trajectory_transformer.blocks[i].la.load_state_dict(gpt.blocks[i].mlp[0].state_dict()      )
        trajectory_transformer.blocks[i].act.load_state_dict(gpt.blocks[i].mlp[1].state_dict()      )
        trajectory_transformer.blocks[i].la.load_state_dict(gpt.blocks[i].mlp[2].state_dict()      )
        trajectory_transformer.blocks[i].drop.load_state_dict(gpt.blocks[i].mlp[3].state_dict()      )
  
     torch.save(trajectory_transformer.state_dict()	,   """pytorch_model.bin"""      )
  
  
  if __name__ == "__main__":
   UpperCAmelCase   =	Parser().parse_args("""plan""")
   convert_trajectory_transformer_original_pytorch_checkpoint_to_pytorch(
       args.logbase, args.dataset, args.gpt_loadpath, args.gpt_epoch, args.device
   )
  
  