
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					)							->			Tuple:
					super().__init__()
					lowerCamelCase_         =     nn.Linear(3		,    4					)
					lowerCamelCase_         =     nn.BatchNormad(4					)
					lowerCamelCase_         =     nn.Linear(4		,    5					)



	def        SCREAMING_SNAKE_CASE_(       self		,    lowercase					)							->			List[str]:
					return self.lineara(self.batchnorm(self.lineara(lowercase					)					)					)




class        _SCREAMING_SNAKE_CASE	( snake_case_		):



	def        SCREAMING_SNAKE_CASE_(       self		,    lowercase		,    *lowercase		,    **lowercase					)							->			str:
					return (args[0] + 1,) + args[1:], kwargs




class        _SCREAMING_SNAKE_CASE	( snake_case_		):



	def        SCREAMING_SNAKE_CASE_(       self		,    lowercase		,    lowercase					)							->			Union[str, Any]:
					return output + 1




class        _SCREAMING_SNAKE_CASE	( unittest.TestCase		):



	def        SCREAMING_SNAKE_CASE_(       self					)							->			Optional[int]:
					lowerCamelCase_         =     ModelForTest()
					lowerCamelCase_         =     ModelHook()

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



	def        SCREAMING_SNAKE_CASE_(       self					)							->			int:
					lowerCamelCase_         =     ModelForTest()
					lowerCamelCase_         =     ModelHook()

					add_hook_to_module(lowercase		,    lowercase					)
					add_hook_to_module(lowercase		,    lowercase		,    append=lowercase					)

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



	def        SCREAMING_SNAKE_CASE_(       self					)							->			Tuple:
					lowerCamelCase_         =     ModelForTest()
					lowerCamelCase_         =     torch.randn(2		,    3					)
					lowerCamelCase_         =     test_model(x + 1					)
					lowerCamelCase_         =     test_model(x + 2					)

					lowerCamelCase_         =     PreForwardHook()
					add_hook_to_module(lowercase		,    lowercase					)
					lowerCamelCase_         =     test_model(lowercase					)
					self.assertTrue(torch.allclose(lowercase		,    lowercase		,    atol=1e-5					)					)

					# Attaching a hook to a model when it already has one replaces, does not chain
					lowerCamelCase_         =     PreForwardHook()
					add_hook_to_module(lowercase		,    lowercase					)
					lowerCamelCase_         =     test_model(lowercase					)
					self.assertTrue(torch.allclose(lowercase		,    lowercase		,    atol=1e-5					)					)

					# You need to use the sequential hook to chain two or more hooks
					lowerCamelCase_         =     SequentialHook(PreForwardHook()		,    PreForwardHook()					)
					add_hook_to_module(lowercase		,    lowercase					)

					lowerCamelCase_         =     test_model(lowercase					)
					assert torch.allclose(lowercase		,    lowercase		,    atol=1e-5					)



	def        SCREAMING_SNAKE_CASE_(       self					)							->			Optional[Any]:
					lowerCamelCase_         =     ModelForTest()
					lowerCamelCase_         =     torch.randn(2		,    3					)
					lowerCamelCase_         =     test_model(lowercase					)

					lowerCamelCase_         =     PostForwardHook()
					add_hook_to_module(lowercase		,    lowercase					)
					lowerCamelCase_         =     test_model(lowercase					)
					self.assertTrue(torch.allclose(lowercase		,    output + 1		,    atol=1e-5					)					)

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

					# You need to use the sequential hook to chain two or more hooks
					lowerCamelCase_         =     SequentialHook(PostForwardHook()		,    PostForwardHook()					)
					add_hook_to_module(lowercase		,    lowercase					)

					lowerCamelCase_         =     test_model(lowercase					)
					assert torch.allclose(lowercase		,    output + 2		,    atol=1e-5					)



	def        SCREAMING_SNAKE_CASE_(       self					)							->			List[str]:
					lowerCamelCase_         =     ModelForTest()
					lowerCamelCase_         =     torch.randn(2		,    3					)
					lowerCamelCase_         =     test_model(lowercase					)

					lowerCamelCase_         =     PostForwardHook()
					add_hook_to_module(lowercase		,    lowercase					)
					lowerCamelCase_         =     test_model(lowercase					)
					self.assertTrue(torch.allclose(lowercase		,    output + 1					)					)
					self.assertTrue(outputa.requires_grad					)

					lowerCamelCase_         =     True
					lowerCamelCase_         =     test_model(lowercase					)
					self.assertFalse(outputa.requires_grad					)



	@require_multi_gpu
	def        SCREAMING_SNAKE_CASE_(       self					)							->			str:
					lowerCamelCase_         =     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_         =     torch.randn(2		,    3					)
					lowerCamelCase_         =     model(lowercase					)
					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(lowercase		,    AlignDevicesHook(io_same_device=lowercase					)					)
					lowerCamelCase_         =     torch.randn(2		,    3					).to(0					)
					lowerCamelCase_         =     model(lowercase					)
					self.assertEqual(output.device		,    torch.device(0					)					)



	def        SCREAMING_SNAKE_CASE_(       self					)							->			Optional[Any]:
					lowerCamelCase_         =     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_         =     {"execution_device": 0 if torch.cuda.is_available() else "cpu", "offload": True}

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

					# 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_         =     torch.device(hook_kwargs["execution_device"]					)
					self.assertEqual(model.batchnorm.running_mean.device		,    lowercase					)

					lowerCamelCase_         =     torch.randn(2		,    3					)
					lowerCamelCase_         =     model(lowercase					)
					self.assertEqual(output.device		,    lowercase					)

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

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

					# 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_         =     torch.randn(2		,    3					)
					lowerCamelCase_         =     model(lowercase					)
					self.assertEqual(output.device		,    lowercase					)

					# 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        SCREAMING_SNAKE_CASE_(       self					)							->			Optional[int]:
					lowerCamelCase_         =     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_         =     0 if torch.cuda.is_available() else "cpu"
					attach_align_device_hook(lowercase		,    execution_device=lowercase		,    offload=lowercase					)

					# 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_         =     torch.device(lowercase					)
					self.assertEqual(model.batchnorm.running_mean.device		,    lowercase					)

					lowerCamelCase_         =     torch.randn(2		,    3					)
					lowerCamelCase_         =     model(lowercase					)
					self.assertEqual(output.device		,    lowercase					)

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

					# 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_         =     torch.randn(2		,    3					)
					lowerCamelCase_         =     model(lowercase					)
					self.assertEqual(output.device		,    lowercase					)

					# Removing hooks loads back the weights in the model.
					remove_hook_from_submodules(lowercase					)
					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        SCREAMING_SNAKE_CASE_(       self					)							->			Tuple:
					lowerCamelCase_         =     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_         =     0 if torch.cuda.is_available() else "cpu"
					attach_align_device_hook(
					    lowercase		,    execution_device=lowercase		,    offload=lowercase		,    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_         =     torch.device(lowercase					)
					self.assertEqual(model.batchnorm.running_mean.device		,    lowercase					)

					lowerCamelCase_         =     torch.randn(2		,    3					)
					lowerCamelCase_         =     model(lowercase					)
					self.assertEqual(output.device		,    lowercase					)

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

					# 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_         =     torch.randn(2		,    3					)
					lowerCamelCase_         =     model(lowercase					)
					self.assertEqual(output.device		,    lowercase					)

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




import numpy as np
import torch
import torch.nn as nn
from transformers import CLIPConfig, CLIPVisionModelWithProjection, PreTrainedModel

from ...utils import logging


__A							=logging.get_logger(__name__)




class        _SCREAMING_SNAKE_CASE	( snake_case_		):
	lowerCAmelCase__   =					CLIPConfig

	lowerCAmelCase__   =					['CLIPEncoderLayer']



	def __init__(       self		,    lowercase					)							->			int:
					super().__init__(lowercase					)

					lowerCamelCase_         =     CLIPVisionModelWithProjection(config.vision_config					)

					lowerCamelCase_         =     nn.Linear(config.vision_config.projection_dim		,    1					)
					lowerCamelCase_         =     nn.Linear(config.vision_config.projection_dim		,    1					)



	@torch.no_grad()
	def        SCREAMING_SNAKE_CASE_(       self		,    lowercase		,    lowercase		,    lowercase=0.5		,    lowercase=0.5					)							->			int:
					lowerCamelCase_         =     self.vision_model(lowercase					)[0]

					lowerCamelCase_         =     self.p_head(lowercase					)
					lowerCamelCase_         =     nsfw_detected.flatten()
					lowerCamelCase_         =     nsfw_detected > p_threshold
					lowerCamelCase_         =     nsfw_detected.tolist()

					if any(lowercase					):
									logger.warning(
									    "Potential NSFW content was detected in one or more images. A black image will be returned instead."
									    " Try again with a different prompt and/or seed."					)

					for idx, nsfw_detected_ in enumerate(lowercase					):
									if nsfw_detected_:
													lowerCamelCase_         =     np.zeros(images[idx].shape					)

					lowerCamelCase_         =     self.w_head(lowercase					)
					lowerCamelCase_         =     watermark_detected.flatten()
					lowerCamelCase_         =     watermark_detected > w_threshold
					lowerCamelCase_         =     watermark_detected.tolist()

					if any(lowercase					):
									logger.warning(
									    "Potential watermarked content was detected in one or more images. A black image will be returned instead."
									    " Try again with a different prompt and/or seed."					)

					for idx, watermark_detected_ in enumerate(lowercase					):
									if watermark_detected_:
													lowerCamelCase_         =     np.zeros(images[idx].shape					)

					return images, nsfw_detected, watermark_detected


