






"""simple docstring"""




import argparse
import math
import os
from copy import deepcopy

import torch
from audio_diffusion.models import DiffusionAttnUnetaD
from diffusion import sampling
from torch import nn

from diffusers import DanceDiffusionPipeline, IPNDMScheduler, UNetaDModel


_snake_case          =   {
    'gwf-440k': {
        'url': 'https://model-server.zqevans2.workers.dev/gwf-440k.ckpt',
        'sample_rate': 4_8000,
        'sample_size': 6_5536,
    },
    'jmann-small-190k': {
        'url': 'https://model-server.zqevans2.workers.dev/jmann-small-190k.ckpt',
        'sample_rate': 4_8000,
        'sample_size': 6_5536,
    },
    'jmann-large-580k': {
        'url': 'https://model-server.zqevans2.workers.dev/jmann-large-580k.ckpt',
        'sample_rate': 4_8000,
        'sample_size': 13_1072,
    },
    'maestro-uncond-150k': {
        'url': 'https://model-server.zqevans2.workers.dev/maestro-uncond-150k.ckpt',
        'sample_rate': 1_6000,
        'sample_size': 6_5536,
    },
    'unlocked-uncond-250k': {
        'url': 'https://model-server.zqevans2.workers.dev/unlocked-uncond-250k.ckpt',
        'sample_rate': 1_6000,
        'sample_size': 6_5536,
    },
    'honk-140k': {
        'url': 'https://model-server.zqevans2.workers.dev/honk-140k.ckpt',
        'sample_rate': 1_6000,
        'sample_size': 6_5536,
    },
}

def 			lowerCAmelCase__     (     UpperCamelCase__     ,   UpperCamelCase__				):



  '''simple docstring'''
  return torch.atana(UpperCamelCase__     ,   UpperCamelCase__				) / math.pi * 2

def 			lowerCAmelCase__     (     UpperCamelCase__				):



  '''simple docstring'''
  _a  :       Dict             =	torch.sin(t * math.pi / 2				) ** 2
  _a  :       Union[str, Any]             =	(1 - sigma**2) ** 0.5
  return alpha_sigma_to_t(UpperCamelCase__     ,   UpperCamelCase__				)






class        UpperCamelCase			(	snake_case_      ):
 pass







class        UpperCamelCase			(	nn.Module      ):


 def __init__(     self		:						Any							,						UpperCAmelCase__		:						Optional[Any]							)					->       Union[str, Any]:
   super().__init__()

   _a  :       int             =	DiffusionAttnUnetaD(UpperCAmelCase__							,						n_attn_layers=4							)
   _a  :       Union[str, Any]             =	deepcopy(self.diffusion							)
   _a  :       Dict             =	torch.quasirandom.SobolEngine(1							,						scramble=UpperCAmelCase__							)







def 			lowerCAmelCase__     (     UpperCamelCase__				):



  '''simple docstring'''
  _a  :       str             =	MODELS_MAP[model_name]["""url"""]
  os.system(F"""wget {url} ./"""				)

  return F"""./{model_name}.ckpt"""


_snake_case          =   {
    '1': 'resnets.0',
    '2': 'attentions.0',
    '3': 'resnets.1',
    '4': 'attentions.1',
    '5': 'resnets.2',
    '6': 'attentions.2',
}
_snake_case          =   {
    '8': 'resnets.0',
    '9': 'attentions.0',
    '10': 'resnets.1',
    '11': 'attentions.1',
    '12': 'resnets.2',
    '13': 'attentions.2',
}
_snake_case          =   {
    '1': 'resnets.0',
    '2': 'attentions.0',
    '3': 'resnets.1',
    '4': 'attentions.1',
    '5': 'resnets.2',
    '6': 'attentions.2',
    '8': 'resnets.3',
    '9': 'attentions.3',
    '10': 'resnets.4',
    '11': 'attentions.4',
    '12': 'resnets.5',
    '13': 'attentions.5',
}
_snake_case          =   {
    '0': 'resnets.0',
    '1': 'resnets.1',
    '2': 'resnets.2',
    '4': 'resnets.0',
    '5': 'resnets.1',
    '6': 'resnets.2',
}

_snake_case          =   {
    'skip': 'conv_skip',
    'main.0': 'conv_1',
    'main.1': 'group_norm_1',
    'main.3': 'conv_2',
    'main.4': 'group_norm_2',
}

_snake_case          =   {
    'norm': 'group_norm',
    'qkv_proj': ['query', 'key', 'value'],
    'out_proj': ['proj_attn'],
}

def 			lowerCAmelCase__     (     UpperCamelCase__				):



  '''simple docstring'''
  if name.startswith("""skip"""				):
    return name.replace("""skip"""     ,   RES_CONV_MAP["""skip"""]				)

  # name has to be of format main.{digit}
  if not name.startswith("""main."""				):
    raise ValueError(F"""ResConvBlock error with {name}"""				)

  return name.replace(name[:6]     ,   RES_CONV_MAP[name[:6]]				)

def 			lowerCAmelCase__     (     UpperCamelCase__				):



  '''simple docstring'''
  for key, value in ATTN_MAP.items():
    if name.startswith(UpperCamelCase__				) and not isinstance(UpperCamelCase__     ,   UpperCamelCase__				):
      return name.replace(UpperCamelCase__     ,   UpperCamelCase__				)
    elif name.startswith(UpperCamelCase__				):
      return [name.replace(UpperCamelCase__     ,   UpperCamelCase__				) for v in value]
  raise ValueError(F"""Attn error with {name}"""				)

def 			lowerCAmelCase__     (     UpperCamelCase__     ,   UpperCamelCase__=1_3				):



  '''simple docstring'''
  _a  :       Dict             =	input_string

  if string.split("""."""				)[0] == "timestep_embed":
    return string.replace("""timestep_embed"""     ,   """time_proj"""				)

  _a  :       Any             =	0
  if string.startswith("""net.3."""				):
    depth += 1
    _a  :       Union[str, Any]             =	string[6:]
  elif string.startswith("""net."""				):
    _a  :       List[Any]             =	string[4:]

  while string.startswith("""main.7."""				):
    depth += 1
    _a  :       Tuple             =	string[7:]

  if string.startswith("""main."""				):
    _a  :       Tuple             =	string[5:]

  # mid block
  if string[:2].isdigit():
    _a  :       Dict             =	string[:2]
    _a  :       List[Any]             =	string[2:]
  else:
    _a  :       Union[str, Any]             =	string[0]
    _a  :       List[Any]             =	string[1:]

  if depth == max_depth:
    _a  :       Tuple             =	MID_NUM_TO_LAYER[layer_num]
    _a  :       Dict             =	"""mid_block"""
  elif depth > 0 and int(UpperCamelCase__				) < 7:
    _a  :       Union[str, Any]             =	DOWN_NUM_TO_LAYER[layer_num]
    _a  :       Tuple             =	F"""down_blocks.{depth}"""
  elif depth > 0 and int(UpperCamelCase__				) > 7:
    _a  :       Any             =	UP_NUM_TO_LAYER[layer_num]
    _a  :       Dict             =	F"""up_blocks.{max_depth - depth - 1}"""
  elif depth == 0:
    _a  :       Dict             =	DEPTH_0_TO_LAYER[layer_num]
    _a  :       str             =	F"""up_blocks.{max_depth - 1}""" if int(UpperCamelCase__				) > 3 else """down_blocks.0"""

  if not string_left.startswith("""."""				):
    raise ValueError(F"""Naming error with {input_string} and string_left: {string_left}."""				)

  _a  :       int             =	string_left[1:]

  if "resnets" in new_layer:
    _a  :       Optional[Any]             =	convert_resconv_naming(UpperCamelCase__				)
  elif "attentions" in new_layer:
    _a  :       Tuple             =	convert_attn_naming(UpperCamelCase__				)
    _a  :       Tuple             =	new_string_left

  if not isinstance(UpperCamelCase__     ,   UpperCamelCase__				):
    _a  :       Union[str, Any]             =	prefix + """.""" + new_layer + """.""" + string_left
  else:
    _a  :       Tuple             =	[prefix + """.""" + new_layer + """.""" + s for s in string_left]
  return new_string

def 			lowerCAmelCase__     (     UpperCamelCase__				):



  '''simple docstring'''
  _a  :       List[str]             =	{}
  for k, v in state_dict.items():
    if k.endswith("""kernel"""				):
      # up- and downsample layers, don't have trainable weights
      continue

    _a  :       Optional[int]             =	rename(UpperCamelCase__				)

    # check if we need to transform from Conv => Linear for attention
    if isinstance(UpperCamelCase__     ,   UpperCamelCase__				):
      _a  :       int             =	transform_conv_attns(UpperCamelCase__     ,   UpperCamelCase__     ,   UpperCamelCase__				)
    else:
      _a  :       int             =	v

  return new_state_dict

def 			lowerCAmelCase__     (     UpperCamelCase__     ,   UpperCamelCase__     ,   UpperCamelCase__				):



  '''simple docstring'''
  if len(UpperCamelCase__				) == 1:
    if len(v.shape				) == 3:
      # weight
      _a  :       List[Any]             =	v[:, :, 0]
    else:
      # bias
      _a  :       str             =	v
  else:
    # qkv matrices
    _a  :       Any             =	v.shape[0]
    _a  :       int             =	trippled_shape // 3
    for i in range(3				):
      if len(v.shape				) == 3:
        _a  :       List[Any]             =	v[i * single_shape : (i + 1) * single_shape, :, 0]
      else:
        _a  :       int             =	v[i * single_shape : (i + 1) * single_shape]
  return new_state_dict

def 			lowerCAmelCase__     (     UpperCamelCase__				):



  '''simple docstring'''
  _a  :       Tuple             =	torch.device("""cuda""" if torch.cuda.is_available() else """cpu"""				)

  _a  :       Dict             =	args.model_path.split("""/"""				)[-1].split("""."""				)[0]
  if not os.path.isfile(args.model_path				):
    assert (
        model_name == args.model_path
    ), F"""Make sure to provide one of the official model names {MODELS_MAP.keys()}"""
    _a  :       Any             =	download(UpperCamelCase__				)

  _a  :       Dict             =	MODELS_MAP[model_name]["""sample_rate"""]
  _a  :       Optional[int]             =	MODELS_MAP[model_name]["""sample_size"""]

  _a  :       Tuple             =	Object()
  _a  :       List[Any]             =	sample_size
  _a  :       Dict             =	sample_rate
  _a  :       str             =	0

  _a  :       str             =	UNetaDModel(sample_size=UpperCamelCase__     ,   sample_rate=UpperCamelCase__				)
  _a  :       int             =	diffusers_model.state_dict()

  _a  :       int             =	DiffusionUncond(UpperCamelCase__				)
  orig_model.load_state_dict(torch.load(args.model_path     ,   map_location=UpperCamelCase__				)["""state_dict"""]				)
  _a  :       Union[str, Any]             =	orig_model.diffusion_ema.eval()
  _a  :       int             =	orig_model.state_dict()
  _a  :       int             =	rename_orig_weights(UpperCamelCase__				)

  _a  :       str             =	set(renamed_state_dict.keys()				) - set(diffusers_state_dict.keys()				)
  _a  :       Optional[int]             =	set(diffusers_state_dict.keys()				) - set(renamed_state_dict.keys()				)

  assert len(UpperCamelCase__				) == 0, F"""Problem with {renamed_minus_diffusers}"""
  assert all(k.endswith("""kernel"""				) for k in list(UpperCamelCase__				)				), F"""Problem with {diffusers_minus_renamed}"""

  for key, value in renamed_state_dict.items():
    assert (
        diffusers_state_dict[key].squeeze().shape == value.squeeze().shape
    ), F"""Shape for {key} doesn't match. Diffusers: {diffusers_state_dict[key].shape} vs. {value.shape}"""
    if key == "time_proj.weight":
      _a  :       int             =	value.squeeze()

    _a  :       List[Any]             =	value

  diffusers_model.load_state_dict(UpperCamelCase__				)

  _a  :       str             =	1_0_0
  _a  :       Optional[int]             =	3_3

  _a  :       Tuple             =	IPNDMScheduler(num_train_timesteps=UpperCamelCase__				)

  _a  :       List[str]             =	torch.manual_seed(UpperCamelCase__				)
  _a  :       Tuple             =	torch.randn([1, 2, config.sample_size]     ,   generator=UpperCamelCase__				).to(UpperCamelCase__				)

  _a  :       Optional[int]             =	torch.linspace(1     ,   0     ,   steps + 1     ,   device=UpperCamelCase__				)[:-1]
  _a  :       Tuple             =	get_crash_schedule(UpperCamelCase__				)

  _a  :       Tuple             =	DanceDiffusionPipeline(unet=UpperCamelCase__     ,   scheduler=UpperCamelCase__				)

  _a  :       Union[str, Any]             =	torch.manual_seed(3_3				)
  _a  :       Tuple             =	pipe(num_inference_steps=UpperCamelCase__     ,   generator=UpperCamelCase__				).audios

  _a  :       int             =	sampling.iplms_sample(UpperCamelCase__     ,   UpperCamelCase__     ,   UpperCamelCase__     ,   {}				)
  _a  :       List[Any]             =	generated.clamp(-1     ,   1				)

  _a  :       int             =	(generated - audio).abs().sum()
  _a  :       Any             =	(generated - audio).abs().max()

  if args.save:
    pipe.save_pretrained(args.checkpoint_path				)

  print("""Diff sum"""     ,   UpperCamelCase__				)
  print("""Diff max"""     ,   UpperCamelCase__				)

  assert diff_max < 1e-3, F"""Diff max: {diff_max} is too much :-/"""

  print(F"""Conversion for {model_name} successful!"""				)


if __name__ == "__main__":
 _snake_case          =   argparse.ArgumentParser()

 parser.add_argument('--model_path', default=None, type=str, required=True, help='Path to the model to convert.')
 parser.add_argument(
     '--save', default=True, type=bool, required=False, help='Whether to save the converted model or not.'
 )
 parser.add_argument('--checkpoint_path', default=None, type=str, required=True, help='Path to the output model.')
 _snake_case          =   parser.parse_args()

 main(args)









"""simple docstring"""




import inspect
import unittest

import torch
import torch.nn as nn

from accelerate.hooks import (
    AlignDevicesHook,
    ModelHook,
    SequentialHook,
    add_hook_to_module,
    attach_align_device_hook,
    remove_hook_from_module,
    remove_hook_from_submodules,
)
from accelerate.test_utils import require_multi_gpu






class        UpperCamelCase			(	nn.Module      ):


 def __init__(     self		:						Union[str, Any]							)					->       int:
   super().__init__()
   _a  :       Optional[Any]             =	nn.Linear(3							,						4							)
   _a  :       Tuple             =	nn.BatchNormad(4							)
   _a  :       Dict             =	nn.Linear(4							,						5							)







 def _lowercase	(     self		:						Optional[int]							,						UpperCAmelCase__		:						List[str]							)					->       int:
   return self.lineara(self.batchnorm(self.lineara(UpperCAmelCase__							)							)							)







class        UpperCamelCase			(	snake_case_      ):


 def _lowercase	(     self		:						Any							,						UpperCAmelCase__		:						Any							,						*UpperCAmelCase__		:						List[str]							,						**UpperCAmelCase__		:						Optional[int]							)					->       Optional[int]:
   return (args[0] + 1,) + args[1:], kwargs







class        UpperCamelCase			(	snake_case_      ):


 def _lowercase	(     self		:						Optional[Any]							,						UpperCAmelCase__		:						Tuple							,						UpperCAmelCase__		:						Optional[Any]							)					->       List[str]:
   return output + 1







class        UpperCamelCase			(	unittest.TestCase      ):


 def _lowercase	(     self		:						Dict							)					->       str:
   _a  :       List[Any]             =	ModelForTest()
   _a  :       str             =	ModelHook()

   add_hook_to_module(UpperCAmelCase__							,						UpperCAmelCase__							)
   self.assertEqual(test_model._hf_hook							,						UpperCAmelCase__							)
   self.assertTrue(hasattr(UpperCAmelCase__							,						"""_old_forward"""							)							)

   # Check adding the hook did not change the name or the signature
   self.assertEqual(test_model.forward.__name__							,						"""forward"""							)
   self.assertListEqual(list(inspect.signature(test_model.forward							).parameters							)							,						["""x"""]							)

   remove_hook_from_module(UpperCAmelCase__							)
   self.assertFalse(hasattr(UpperCAmelCase__							,						"""_hf_hook"""							)							)
   self.assertFalse(hasattr(UpperCAmelCase__							,						"""_old_forward"""							)							)


 def _lowercase	(     self		:						Optional[int]							)					->       Optional[int]:
   _a  :       Dict             =	ModelForTest()
   _a  :       Dict             =	ModelHook()

   add_hook_to_module(UpperCAmelCase__							,						UpperCAmelCase__							)
   add_hook_to_module(UpperCAmelCase__							,						UpperCAmelCase__							,						append=UpperCAmelCase__							)

   self.assertEqual(isinstance(test_model._hf_hook							,						UpperCAmelCase__							)							,						UpperCAmelCase__							)
   self.assertEqual(len(test_model._hf_hook.hooks							)							,						2							)
   self.assertTrue(hasattr(UpperCAmelCase__							,						"""_old_forward"""							)							)

   # Check adding the hook did not change the name or the signature
   self.assertEqual(test_model.forward.__name__							,						"""forward"""							)
   self.assertListEqual(list(inspect.signature(test_model.forward							).parameters							)							,						["""x"""]							)

   remove_hook_from_module(UpperCAmelCase__							)
   self.assertFalse(hasattr(UpperCAmelCase__							,						"""_hf_hook"""							)							)
   self.assertFalse(hasattr(UpperCAmelCase__							,						"""_old_forward"""							)							)


 def _lowercase	(     self		:						Dict							)					->       int:
   _a  :       str             =	ModelForTest()
   _a  :       List[Any]             =	torch.randn(2							,						3							)
   _a  :       Optional[Any]             =	test_model(x + 1							)
   _a  :       str             =	test_model(x + 2							)

   _a  :       Union[str, Any]             =	PreForwardHook()
   add_hook_to_module(UpperCAmelCase__							,						UpperCAmelCase__							)
   _a  :       Tuple             =	test_model(UpperCAmelCase__							)
   self.assertTrue(torch.allclose(UpperCAmelCase__							,						UpperCAmelCase__							,						atol=1E-5							)							)

   # Attaching a hook to a model when it already has one replaces, does not chain
   _a  :       int             =	PreForwardHook()
   add_hook_to_module(UpperCAmelCase__							,						UpperCAmelCase__							)
   _a  :       str             =	test_model(UpperCAmelCase__							)
   self.assertTrue(torch.allclose(UpperCAmelCase__							,						UpperCAmelCase__							,						atol=1E-5							)							)

   # You need to use the sequential hook to chain two or more hooks
   _a  :       int             =	SequentialHook(PreForwardHook()							,						PreForwardHook()							)
   add_hook_to_module(UpperCAmelCase__							,						UpperCAmelCase__							)

   _a  :       Tuple             =	test_model(UpperCAmelCase__							)
   assert torch.allclose(UpperCAmelCase__							,						UpperCAmelCase__							,						atol=1E-5							)


 def _lowercase	(     self		:						Tuple							)					->       int:
   _a  :       Tuple             =	ModelForTest()
   _a  :       Union[str, Any]             =	torch.randn(2							,						3							)
   _a  :       Optional[int]             =	test_model(UpperCAmelCase__							)

   _a  :       int             =	PostForwardHook()
   add_hook_to_module(UpperCAmelCase__							,						UpperCAmelCase__							)
   _a  :       List[str]             =	test_model(UpperCAmelCase__							)
   self.assertTrue(torch.allclose(UpperCAmelCase__							,						output + 1							,						atol=1E-5							)							)

   # Attaching a hook to a model when it already has one replaces, does not chain
   _a  :       List[Any]             =	PostForwardHook()
   add_hook_to_module(UpperCAmelCase__							,						UpperCAmelCase__							)
   _a  :       Dict             =	test_model(UpperCAmelCase__							)
   self.assertTrue(torch.allclose(UpperCAmelCase__							,						output + 1							,						atol=1E-5							)							)

   # You need to use the sequential hook to chain two or more hooks
   _a  :       Any             =	SequentialHook(PostForwardHook()							,						PostForwardHook()							)
   add_hook_to_module(UpperCAmelCase__							,						UpperCAmelCase__							)

   _a  :       Optional[int]             =	test_model(UpperCAmelCase__							)
   assert torch.allclose(UpperCAmelCase__							,						output + 2							,						atol=1E-5							)


 def _lowercase	(     self		:						Dict							)					->       Optional[Any]:
   _a  :       Any             =	ModelForTest()
   _a  :       List[Any]             =	torch.randn(2							,						3							)
   _a  :       Dict             =	test_model(UpperCAmelCase__							)

   _a  :       Any             =	PostForwardHook()
   add_hook_to_module(UpperCAmelCase__							,						UpperCAmelCase__							)
   _a  :       List[str]             =	test_model(UpperCAmelCase__							)
   self.assertTrue(torch.allclose(UpperCAmelCase__							,						output + 1							)							)
   self.assertTrue(outputa.requires_grad							)

   _a  :       Any             =	True
   _a  :       Union[str, Any]             =	test_model(UpperCAmelCase__							)
   self.assertFalse(outputa.requires_grad							)


 @require_multi_gpu
 def _lowercase	(     self		:						Optional[Any]							)					->       str:
   _a  :       List[Any]             =	ModelForTest()
   # Everything is on CPU
   self.assertEqual(model.lineara.weight.device							,						torch.device("""cpu"""							)							)
   self.assertEqual(model.batchnorm.weight.device							,						torch.device("""cpu"""							)							)
   self.assertEqual(model.lineara.weight.device							,						torch.device("""cpu"""							)							)

   # This will move each submodule on different devices
   add_hook_to_module(model.lineara							,						AlignDevicesHook(execution_device=0							)							)
   add_hook_to_module(model.batchnorm							,						AlignDevicesHook(execution_device=0							)							)
   add_hook_to_module(model.lineara							,						AlignDevicesHook(execution_device=1							)							)

   self.assertEqual(model.lineara.weight.device							,						torch.device(0							)							)
   self.assertEqual(model.batchnorm.weight.device							,						torch.device(0							)							)
   self.assertEqual(model.batchnorm.running_mean.device							,						torch.device(0							)							)
   self.assertEqual(model.lineara.weight.device							,						torch.device(1							)							)

   # We can still make a forward pass. The input does not need to be on any particular device
   _a  :       Optional[int]             =	torch.randn(2							,						3							)
   _a  :       Any             =	model(UpperCAmelCase__							)
   self.assertEqual(output.device							,						torch.device(1							)							)

   # We can add a general hook to put back output on same device as input.
   add_hook_to_module(UpperCAmelCase__							,						AlignDevicesHook(io_same_device=UpperCAmelCase__							)							)
   _a  :       str             =	torch.randn(2							,						3							).to(0							)
   _a  :       Union[str, Any]             =	model(UpperCAmelCase__							)
   self.assertEqual(output.device							,						torch.device(0							)							)


 def _lowercase	(     self		:						str							)					->       Union[str, Any]:
   _a  :       int             =	ModelForTest()

   # Everything is on CPU
   self.assertEqual(model.lineara.weight.device							,						torch.device("""cpu"""							)							)
   self.assertEqual(model.batchnorm.weight.device							,						torch.device("""cpu"""							)							)
   self.assertEqual(model.lineara.weight.device							,						torch.device("""cpu"""							)							)

   # This will move each submodule on different devices
   _a  :       List[Any]             =	{"""execution_device""": 0 if torch.cuda.is_available() else """cpu""", """offload""": True}

   add_hook_to_module(model.lineara							,						AlignDevicesHook(**UpperCAmelCase__							)							)
   add_hook_to_module(model.batchnorm							,						AlignDevicesHook(**UpperCAmelCase__							)							)
   add_hook_to_module(model.lineara							,						AlignDevicesHook(**UpperCAmelCase__							)							)

   # Parameters have been offloaded, so on the meta device
   self.assertEqual(model.lineara.weight.device							,						torch.device("""meta"""							)							)
   self.assertEqual(model.batchnorm.weight.device							,						torch.device("""meta"""							)							)
   self.assertEqual(model.lineara.weight.device							,						torch.device("""meta"""							)							)
   # Buffers are not included in the offload by default, so are on the execution device
   _a  :       Dict             =	torch.device(hook_kwargs["""execution_device"""]							)
   self.assertEqual(model.batchnorm.running_mean.device							,						UpperCAmelCase__							)

   _a  :       int             =	torch.randn(2							,						3							)
   _a  :       str             =	model(UpperCAmelCase__							)
   self.assertEqual(output.device							,						UpperCAmelCase__							)

   # Removing hooks loads back the weights in the model.
   remove_hook_from_module(model.lineara							)
   remove_hook_from_module(model.batchnorm							)
   remove_hook_from_module(model.lineara							)
   self.assertEqual(model.lineara.weight.device							,						torch.device("""cpu"""							)							)
   self.assertEqual(model.batchnorm.weight.device							,						torch.device("""cpu"""							)							)
   self.assertEqual(model.lineara.weight.device							,						torch.device("""cpu"""							)							)

   # Now test with buffers included in the offload
   _a  :       List[str]             =	{
       """execution_device""": 0 if torch.cuda.is_available() else """cpu""",
       """offload""": True,
       """offload_buffers""": True,
   }

   add_hook_to_module(model.lineara							,						AlignDevicesHook(**UpperCAmelCase__							)							)
   add_hook_to_module(model.batchnorm							,						AlignDevicesHook(**UpperCAmelCase__							)							)
   add_hook_to_module(model.lineara							,						AlignDevicesHook(**UpperCAmelCase__							)							)

   # Parameters have been offloaded, so on the meta device, buffers included
   self.assertEqual(model.lineara.weight.device							,						torch.device("""meta"""							)							)
   self.assertEqual(model.batchnorm.weight.device							,						torch.device("""meta"""							)							)
   self.assertEqual(model.lineara.weight.device							,						torch.device("""meta"""							)							)
   self.assertEqual(model.batchnorm.running_mean.device							,						torch.device("""meta"""							)							)

   _a  :       Tuple             =	torch.randn(2							,						3							)
   _a  :       Union[str, Any]             =	model(UpperCAmelCase__							)
   self.assertEqual(output.device							,						UpperCAmelCase__							)

   # Removing hooks loads back the weights in the model.
   remove_hook_from_module(model.lineara							)
   remove_hook_from_module(model.batchnorm							)
   remove_hook_from_module(model.lineara							)
   self.assertEqual(model.lineara.weight.device							,						torch.device("""cpu"""							)							)
   self.assertEqual(model.batchnorm.weight.device							,						torch.device("""cpu"""							)							)
   self.assertEqual(model.lineara.weight.device							,						torch.device("""cpu"""							)							)


 def _lowercase	(     self		:						Tuple							)					->       List[str]:
   _a  :       str             =	ModelForTest()

   # Everything is on CPU
   self.assertEqual(model.lineara.weight.device							,						torch.device("""cpu"""							)							)
   self.assertEqual(model.batchnorm.weight.device							,						torch.device("""cpu"""							)							)
   self.assertEqual(model.lineara.weight.device							,						torch.device("""cpu"""							)							)

   # This will move each submodule on different devices
   _a  :       Union[str, Any]             =	0 if torch.cuda.is_available() else """cpu"""
   attach_align_device_hook(UpperCAmelCase__							,						execution_device=UpperCAmelCase__							,						offload=UpperCAmelCase__							)

   # Parameters have been offloaded, so on the meta device
   self.assertEqual(model.lineara.weight.device							,						torch.device("""meta"""							)							)
   self.assertEqual(model.batchnorm.weight.device							,						torch.device("""meta"""							)							)
   self.assertEqual(model.lineara.weight.device							,						torch.device("""meta"""							)							)
   # Buffers are not included in the offload by default, so are on the execution device
   _a  :       Dict             =	torch.device(UpperCAmelCase__							)
   self.assertEqual(model.batchnorm.running_mean.device							,						UpperCAmelCase__							)

   _a  :       Union[str, Any]             =	torch.randn(2							,						3							)
   _a  :       List[Any]             =	model(UpperCAmelCase__							)
   self.assertEqual(output.device							,						UpperCAmelCase__							)

   # Removing hooks loads back the weights in the model.
   remove_hook_from_submodules(UpperCAmelCase__							)
   self.assertEqual(model.lineara.weight.device							,						torch.device("""cpu"""							)							)
   self.assertEqual(model.batchnorm.weight.device							,						torch.device("""cpu"""							)							)
   self.assertEqual(model.lineara.weight.device							,						torch.device("""cpu"""							)							)

   # Now test with buffers included in the offload
   attach_align_device_hook(UpperCAmelCase__							,						execution_device=UpperCAmelCase__							,						offload=UpperCAmelCase__							,						offload_buffers=UpperCAmelCase__							)

   # Parameters have been offloaded, so on the meta device, buffers included
   self.assertEqual(model.lineara.weight.device							,						torch.device("""meta"""							)							)
   self.assertEqual(model.batchnorm.weight.device							,						torch.device("""meta"""							)							)
   self.assertEqual(model.lineara.weight.device							,						torch.device("""meta"""							)							)
   self.assertEqual(model.batchnorm.running_mean.device							,						torch.device("""meta"""							)							)

   _a  :       List[str]             =	torch.randn(2							,						3							)
   _a  :       Union[str, Any]             =	model(UpperCAmelCase__							)
   self.assertEqual(output.device							,						UpperCAmelCase__							)

   # Removing hooks loads back the weights in the model.
   remove_hook_from_submodules(UpperCAmelCase__							)
   self.assertEqual(model.lineara.weight.device							,						torch.device("""cpu"""							)							)
   self.assertEqual(model.batchnorm.weight.device							,						torch.device("""cpu"""							)							)
   self.assertEqual(model.lineara.weight.device							,						torch.device("""cpu"""							)							)







 def _lowercase	(     self		:						Dict							)					->       str:
   _a  :       Optional[Any]             =	ModelForTest()

   # Everything is on CPU
   self.assertEqual(model.lineara.weight.device							,						torch.device("""cpu"""							)							)
   self.assertEqual(model.batchnorm.weight.device							,						torch.device("""cpu"""							)							)
   self.assertEqual(model.lineara.weight.device							,						torch.device("""cpu"""							)							)

   # This will move each submodule on different devices
   _a  :       str             =	0 if torch.cuda.is_available() else """cpu"""
   attach_align_device_hook(
       UpperCAmelCase__							,						execution_device=UpperCAmelCase__							,						offload=UpperCAmelCase__							,						weights_map=model.state_dict()							)

   # Parameters have been offloaded, so on the meta device
   self.assertEqual(model.lineara.weight.device							,						torch.device("""meta"""							)							)
   self.assertEqual(model.batchnorm.weight.device							,						torch.device("""meta"""							)							)
   self.assertEqual(model.lineara.weight.device							,						torch.device("""meta"""							)							)
   # Buffers are not included in the offload by default, so are on the execution device
   _a  :       Union[str, Any]             =	torch.device(UpperCAmelCase__							)
   self.assertEqual(model.batchnorm.running_mean.device							,						UpperCAmelCase__							)

   _a  :       Union[str, Any]             =	torch.randn(2							,						3							)
   _a  :       int             =	model(UpperCAmelCase__							)
   self.assertEqual(output.device							,						UpperCAmelCase__							)

   # Removing hooks loads back the weights in the model.
   remove_hook_from_submodules(UpperCAmelCase__							)
   self.assertEqual(model.lineara.weight.device							,						torch.device("""cpu"""							)							)
   self.assertEqual(model.batchnorm.weight.device							,						torch.device("""cpu"""							)							)
   self.assertEqual(model.lineara.weight.device							,						torch.device("""cpu"""							)							)

   # Now test with buffers included in the offload
   attach_align_device_hook(
       UpperCAmelCase__							,						execution_device=UpperCAmelCase__							,						offload=UpperCAmelCase__							,						weights_map=model.state_dict()							,						offload_buffers=UpperCAmelCase__							,						)

   # Parameters have been offloaded, so on the meta device, buffers included
   self.assertEqual(model.lineara.weight.device							,						torch.device("""meta"""							)							)
   self.assertEqual(model.batchnorm.weight.device							,						torch.device("""meta"""							)							)
   self.assertEqual(model.lineara.weight.device							,						torch.device("""meta"""							)							)
   self.assertEqual(model.batchnorm.running_mean.device							,						torch.device("""meta"""							)							)

   _a  :       Any             =	torch.randn(2							,						3							)
   _a  :       int             =	model(UpperCAmelCase__							)
   self.assertEqual(output.device							,						UpperCAmelCase__							)

   # Removing hooks loads back the weights in the model.
   remove_hook_from_submodules(UpperCAmelCase__							)
   self.assertEqual(model.lineara.weight.device							,						torch.device("""cpu"""							)							)
   self.assertEqual(model.batchnorm.weight.device							,						torch.device("""cpu"""							)							)
   self.assertEqual(model.lineara.weight.device							,						torch.device("""cpu"""							)							)

