






'''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 	SCREAMING_SNAKE_CASE	(      nn.Module  ):




       def __init__(					self      : str					):


            """simple docstring"""







            super().__init__()
            __lowerCamelCase						:       Tuple				   =				nn.Linear(3	,    4					)
            __lowerCamelCase						:       Optional[int]				   =				nn.BatchNormad(4					)
            __lowerCamelCase						:       int				   =				nn.Linear(4	,    5					)




       def 					a_					(					self      : Tuple	,    A__      : List[Any]					):


            """simple docstring"""







            return self.lineara(self.batchnorm(self.lineara(A__					)					)					)







class 	SCREAMING_SNAKE_CASE	(      lowerCAmelCase_  ):




       def 					a_					(					self      : Optional[Any]	,    A__      : List[Any]	,    *A__      : int	,    **A__      : List[str]					):


            """simple docstring"""







            return (args[0] + 1,) + args[1:], kwargs







class 	SCREAMING_SNAKE_CASE	(      lowerCAmelCase_  ):




       def 					a_					(					self      : Union[str, Any]	,    A__      : Union[str, Any]	,    A__      : Any					):


            """simple docstring"""







            return output + 1







class 	SCREAMING_SNAKE_CASE	(      unittest.TestCase  ):




       def 					a_					(					self      : Dict					):


            """simple docstring"""







            __lowerCamelCase						:       Optional[int]				   =				ModelForTest()
            __lowerCamelCase						:       Optional[int]				   =				ModelHook()

            add_hook_to_module(A__	,    A__					)
            self.assertEqual(test_model._hf_hook	,    A__					)
            self.assertTrue(hasattr(A__	,    """_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(A__					)
            self.assertFalse(hasattr(A__	,    """_hf_hook"""					)					)
            self.assertFalse(hasattr(A__	,    """_old_forward"""					)					)




       def 					a_					(					self      : str					):


            """simple docstring"""







            __lowerCamelCase						:       Any				   =				ModelForTest()
            __lowerCamelCase						:       Dict				   =				ModelHook()

            add_hook_to_module(A__	,    A__					)
            add_hook_to_module(A__	,    A__	,    append=A__					)

            self.assertEqual(isinstance(test_model._hf_hook	,    A__					)	,    A__					)
            self.assertEqual(len(test_model._hf_hook.hooks					)	,    2					)
            self.assertTrue(hasattr(A__	,    """_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(A__					)
            self.assertFalse(hasattr(A__	,    """_hf_hook"""					)					)
            self.assertFalse(hasattr(A__	,    """_old_forward"""					)					)




       def 					a_					(					self      : List[Any]					):


            """simple docstring"""







            __lowerCamelCase						:       Optional[Any]				   =				ModelForTest()
            __lowerCamelCase						:       Tuple				   =				torch.randn(2	,    3					)
            __lowerCamelCase						:       Dict				   =				test_model(x + 1					)
            __lowerCamelCase						:       str				   =				test_model(x + 2					)

            __lowerCamelCase						:       Tuple				   =				PreForwardHook()
            add_hook_to_module(A__	,    A__					)
            __lowerCamelCase						:       Any				   =				test_model(A__					)
            self.assertTrue(torch.allclose(A__	,    A__	,    atol=1e-5					)					)

            # Attaching a hook to a model when it already has one replaces, does not chain
            __lowerCamelCase						:       Union[str, Any]				   =				PreForwardHook()
            add_hook_to_module(A__	,    A__					)
            __lowerCamelCase						:       Tuple				   =				test_model(A__					)
            self.assertTrue(torch.allclose(A__	,    A__	,    atol=1e-5					)					)

            # You need to use the sequential hook to chain two or more hooks
            __lowerCamelCase						:       List[Any]				   =				SequentialHook(PreForwardHook()	,    PreForwardHook()					)
            add_hook_to_module(A__	,    A__					)

            __lowerCamelCase						:       Optional[Any]				   =				test_model(A__					)
            assert torch.allclose(A__	,    A__	,    atol=1e-5					)




       def 					a_					(					self      : Any					):


            """simple docstring"""







            __lowerCamelCase						:       Optional[int]				   =				ModelForTest()
            __lowerCamelCase						:       int				   =				torch.randn(2	,    3					)
            __lowerCamelCase						:       Optional[Any]				   =				test_model(A__					)

            __lowerCamelCase						:       Any				   =				PostForwardHook()
            add_hook_to_module(A__	,    A__					)
            __lowerCamelCase						:       Optional[Any]				   =				test_model(A__					)
            self.assertTrue(torch.allclose(A__	,    output + 1	,    atol=1e-5					)					)

            # Attaching a hook to a model when it already has one replaces, does not chain
            __lowerCamelCase						:       Tuple				   =				PostForwardHook()
            add_hook_to_module(A__	,    A__					)
            __lowerCamelCase						:       Dict				   =				test_model(A__					)
            self.assertTrue(torch.allclose(A__	,    output + 1	,    atol=1e-5					)					)

            # You need to use the sequential hook to chain two or more hooks
            __lowerCamelCase						:       Tuple				   =				SequentialHook(PostForwardHook()	,    PostForwardHook()					)
            add_hook_to_module(A__	,    A__					)

            __lowerCamelCase						:       Optional[Any]				   =				test_model(A__					)
            assert torch.allclose(A__	,    output + 2	,    atol=1e-5					)




       def 					a_					(					self      : Dict					):


            """simple docstring"""







            __lowerCamelCase						:       Union[str, Any]				   =				ModelForTest()
            __lowerCamelCase						:       List[str]				   =				torch.randn(2	,    3					)
            __lowerCamelCase						:       Union[str, Any]				   =				test_model(A__					)

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

            __lowerCamelCase						:       List[Any]				   =				True
            __lowerCamelCase						:       Tuple				   =				test_model(A__					)
            self.assertFalse(outputa.requires_grad					)




       @require_multi_gpu
       def 					a_					(					self      : Any					):


            """simple docstring"""







            __lowerCamelCase						:       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
            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
            __lowerCamelCase						:       List[Any]				   =				torch.randn(2	,    3					)
            __lowerCamelCase						:       int				   =				model(A__					)
            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(A__	,    AlignDevicesHook(io_same_device=A__					)					)
            __lowerCamelCase						:       str				   =				torch.randn(2	,    3					).to(0					)
            __lowerCamelCase						:       Dict				   =				model(A__					)
            self.assertEqual(output.device	,    torch.device(0					)					)




       def 					a_					(					self      : Union[str, Any]					):


            """simple docstring"""







            __lowerCamelCase						:       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
            __lowerCamelCase						:       str				   =				{"""execution_device""": 0 if torch.cuda.is_available() else """cpu""", """offload""": True}

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

            # 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
            __lowerCamelCase						:       Union[str, Any]				   =				torch.device(hook_kwargs["""execution_device"""]					)
            self.assertEqual(model.batchnorm.running_mean.device	,    A__					)

            __lowerCamelCase						:       int				   =				torch.randn(2	,    3					)
            __lowerCamelCase						:       Optional[Any]				   =				model(A__					)
            self.assertEqual(output.device	,    A__					)

            # 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
            __lowerCamelCase						:       int				   =				{
                """execution_device""": 0 if torch.cuda.is_available() else """cpu""",
                """offload""": True,
                """offload_buffers""": True,
            }

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

            # 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"""					)					)

            __lowerCamelCase						:       Optional[Any]				   =				torch.randn(2	,    3					)
            __lowerCamelCase						:       Optional[Any]				   =				model(A__					)
            self.assertEqual(output.device	,    A__					)

            # 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 					a_					(					self      : Dict					):


            """simple docstring"""







            __lowerCamelCase						:       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
            __lowerCamelCase						:       Optional[int]				   =				0 if torch.cuda.is_available() else """cpu"""
            attach_align_device_hook(A__	,    execution_device=A__	,    offload=A__					)

            # 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
            __lowerCamelCase						:       Dict				   =				torch.device(A__					)
            self.assertEqual(model.batchnorm.running_mean.device	,    A__					)

            __lowerCamelCase						:       Tuple				   =				torch.randn(2	,    3					)
            __lowerCamelCase						:       str				   =				model(A__					)
            self.assertEqual(output.device	,    A__					)

            # Removing hooks loads back the weights in the model.
            remove_hook_from_submodules(A__					)
            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(A__	,    execution_device=A__	,    offload=A__	,    offload_buffers=A__					)

            # 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"""					)					)

            __lowerCamelCase						:       Optional[Any]				   =				torch.randn(2	,    3					)
            __lowerCamelCase						:       List[Any]				   =				model(A__					)
            self.assertEqual(output.device	,    A__					)

            # Removing hooks loads back the weights in the model.
            remove_hook_from_submodules(A__					)
            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 					a_					(					self      : str					):


            """simple docstring"""







            __lowerCamelCase						:       Union[str, 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
            __lowerCamelCase						:       Optional[int]				   =				0 if torch.cuda.is_available() else """cpu"""
            attach_align_device_hook(
                A__	,    execution_device=A__	,    offload=A__	,    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
            __lowerCamelCase						:       List[Any]				   =				torch.device(A__					)
            self.assertEqual(model.batchnorm.running_mean.device	,    A__					)

            __lowerCamelCase						:       int				   =				torch.randn(2	,    3					)
            __lowerCamelCase						:       int				   =				model(A__					)
            self.assertEqual(output.device	,    A__					)

            # Removing hooks loads back the weights in the model.
            remove_hook_from_submodules(A__					)
            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(
                A__	,    execution_device=A__	,    offload=A__	,    weights_map=model.state_dict()	,    offload_buffers=A__	,    )

            # 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"""					)					)

            __lowerCamelCase						:       Dict				   =				torch.randn(2	,    3					)
            __lowerCamelCase						:       Optional[int]				   =				model(A__					)
            self.assertEqual(output.device	,    A__					)

            # Removing hooks loads back the weights in the model.
            remove_hook_from_submodules(A__					)
            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"""					)					)











'''simple docstring'''

from typing import List, Optional, Tuple, Union

import torch

from ...utils import logging, randn_tensor
from ..pipeline_utils import AudioPipelineOutput, DiffusionPipeline


UpperCAmelCase__						:Dict          =			logging.get_logger(__name__)  # pylint: disable=invalid-name







class 	SCREAMING_SNAKE_CASE	(      lowerCAmelCase_  ):




       def __init__(					self      : int	,    A__      : List[Any]	,    A__      : Tuple					):


            """simple docstring"""







            super().__init__()
            self.register_modules(unet=A__	,    scheduler=A__					)




       @torch.no_grad()
       def __call__(					self      : List[str]	,    A__      : int = 1	,    A__      : int = 100	,    A__      : Optional[Union[torch.Generator, List[torch.Generator]]] = None	,    A__      : Optional[float] = None	,    A__      : bool = True	,    ):


            """simple docstring"""







            if audio_length_in_s is None:
                 __lowerCamelCase						:       Union[str, Any]				   =				self.unet.config.sample_size / self.unet.config.sample_rate

            __lowerCamelCase						:       int				   =				audio_length_in_s * self.unet.config.sample_rate

            __lowerCamelCase						:       List[Any]				   =				2 ** len(self.unet.up_blocks					)
            if sample_size < 3 * down_scale_factor:
                 raise ValueError(
                     f"{audio_length_in_s} is too small. Make sure it's bigger or equal to"
                     f" {3 * down_scale_factor / self.unet.config.sample_rate}."					)

            __lowerCamelCase						:       Any				   =				int(A__					)
            if sample_size % down_scale_factor != 0:
                 __lowerCamelCase						:       Optional[Any]				   =				(
                     (audio_length_in_s * self.unet.config.sample_rate) // down_scale_factor + 1
                 ) * down_scale_factor
                 logger.info(
                     f"{audio_length_in_s} is increased to {sample_size / self.unet.config.sample_rate} so that it can be handled"
                     f" by the model. It will be cut to {original_sample_size / self.unet.config.sample_rate} after the denoising"
                     """ process."""					)
            __lowerCamelCase						:       Dict				   =				int(A__					)

            __lowerCamelCase						:       Optional[int]				   =				next(iter(self.unet.parameters()					)					).dtype
            __lowerCamelCase						:       Optional[Any]				   =				(batch_size, self.unet.config.in_channels, sample_size)
            if isinstance(A__	,    A__					) and len(A__					) != batch_size:
                 raise ValueError(
                     f"You have passed a list of generators of length {len(A__					)}, but requested an effective batch"
                     f" size of {batch_size}. Make sure the batch size matches the length of the generators."					)

            __lowerCamelCase						:       int				   =				randn_tensor(A__	,    generator=A__	,    device=self.device	,    dtype=A__					)

            # set step values
            self.scheduler.set_timesteps(A__	,    device=audio.device					)
            __lowerCamelCase						:       List[Any]				   =				self.scheduler.timesteps.to(A__					)

            for t in self.progress_bar(self.scheduler.timesteps					):
                 # 1. predict noise model_output
                 __lowerCamelCase						:       List[str]				   =				self.unet(A__	,    A__					).sample

                 # 2. compute previous image: x_t -> t_t-1
                 __lowerCamelCase						:       Dict				   =				self.scheduler.step(A__	,    A__	,    A__					).prev_sample

            __lowerCamelCase						:       Optional[int]				   =				audio.clamp(-1	,    1					).float().cpu().numpy()

            __lowerCamelCase						:       int				   =				audio[:, :, :original_sample_size]

            if not return_dict:
                 return (audio,)

            return AudioPipelineOutput(audios=A__					)



