


"""simple docstring"""







import gc
import importlib.metadata
import tempfile
import unittest

from packaging import version

from transformers import (
    AutoModel,
    AutoModelForCausalLM,
    AutoModelForSeqaSeqLM,
    AutoModelForSequenceClassification,
    AutoTokenizer,
    BitsAndBytesConfig,
    pipeline,
)
from transformers.testing_utils import (
    is_torch_available,
    require_accelerate,
    require_bitsandbytes,
    require_torch,
    require_torch_gpu,
    require_torch_multi_gpu,
    slow,
)




def        UpperCAmelCase__	(lowerCAmelCase_       ):




					'''simple docstring'''


					if model.config.model_type == "gpt2":
										return model.transformer.h[0].mlp.c_fc
					return model.transformer.h[0].mlp.dense_ah_to_h


if is_torch_available():
		import torch
		import torch.nn as nn




		class 			UpperCamelCase_     (			nn.Module):

			"""simple docstring"""





			def __init__(			self      :    Tuple      ,      UpperCAmelCase__      :    nn.Module      ,      UpperCAmelCase__      :    int       )	->     Optional[Any]:
								super().__init__()
								__SCREAMING_SNAKE_CASE						       =    module
								__SCREAMING_SNAKE_CASE						       =    nn.Sequential(
								    nn.Linear(module.in_features      ,      UpperCAmelCase__      ,      bias=UpperCAmelCase__       )      ,      nn.Linear(UpperCAmelCase__      ,      module.out_features      ,      bias=UpperCAmelCase__       )      ,      )
								__SCREAMING_SNAKE_CASE						       =    (2.0 / (5 * min(module.in_features      ,      module.out_features       ))) ** 0.5
								nn.init.normal_(self.adapter[0].weight      ,      std=UpperCAmelCase__       )
								nn.init.zeros_(self.adapter[1].weight       )
								self.adapter.to(module.weight.device       )




			def   UpperCAmelCase_				(			self      :    Optional[Any]      ,      UpperCAmelCase__      :    Optional[int]      ,      *UpperCAmelCase__      :    str      ,      **UpperCAmelCase__      :    List[Any]       )	->     List[str]:
								return self.module(UpperCAmelCase__      ,      *UpperCAmelCase__      ,      **UpperCAmelCase__       ) + self.adapter(UpperCAmelCase__       )




@require_bitsandbytes
@require_accelerate
@require_torch
@require_torch_gpu
@slow
class 			UpperCamelCase_     (			unittest.TestCase):

	"""simple docstring"""





	snake_case__    :  Optional[int]    =      "bigscience/bloom-1b7"

	# Constant values
	snake_case__    :  Dict    =      2.1_09_65_95_52_69_25_74

	snake_case__    :  Union[str, Any]    =      "Hello my name is"
	snake_case__    :  Optional[int]    =      set()
	EXPECTED_OUTPUTS.add("Hello my name is John and I am a professional photographer. I")
	EXPECTED_OUTPUTS.add("Hello my name is John.\nI am a friend of your father.\n")
	EXPECTED_OUTPUTS.add("Hello my name is John Doe, I am a student at the University")
	snake_case__    :  str    =      10




	def   UpperCAmelCase_				(			self      :    Tuple       )	->     Dict:
						# Models and tokenizer
						__SCREAMING_SNAKE_CASE						       =    AutoTokenizer.from_pretrained(self.model_name       )




class 			UpperCamelCase_     (			UpperCamelCase):

	"""simple docstring"""





	def   UpperCAmelCase_				(			self      :    Optional[Any]       )	->     Optional[Any]:
						super().setUp()

						# Models and tokenizer
						__SCREAMING_SNAKE_CASE						       =    AutoModelForCausalLM.from_pretrained(
						    self.model_name      ,      torch_dtype=torch.floataa      ,      device_map="auto"       )
						__SCREAMING_SNAKE_CASE						       =    AutoModelForCausalLM.from_pretrained(self.model_name      ,      load_in_abit=UpperCAmelCase__      ,      device_map="auto"       )




	def   UpperCAmelCase_				(			self      :    Tuple       )	->     str:
						del self.model_fpaa
						del self.model_abit

						gc.collect()
						torch.cuda.empty_cache()




	def   UpperCAmelCase_				(			self      :    Dict       )	->     str:
						__SCREAMING_SNAKE_CASE						       =    self.model_abit.config

						self.assertTrue(hasattr(UpperCAmelCase__      ,      "quantization_config"       )       )

						__SCREAMING_SNAKE_CASE						       =    config.to_dict()
						__SCREAMING_SNAKE_CASE						       =    config.to_diff_dict()

						__SCREAMING_SNAKE_CASE						       =    config.to_json_string()




	def   UpperCAmelCase_				(			self      :    Union[str, Any]       )	->     Optional[Any]:
						from bitsandbytes.nn import Paramsabit

						__SCREAMING_SNAKE_CASE						       =    self.model_fpaa.get_memory_footprint()
						__SCREAMING_SNAKE_CASE						       =    self.model_abit.get_memory_footprint()

						self.assertAlmostEqual(mem_fpaa / mem_abit      ,      self.EXPECTED_RELATIVE_DIFFERENCE       )
						__SCREAMING_SNAKE_CASE						       =    get_some_linear_layer(self.model_abit       )
						self.assertTrue(linear.weight.__class__ == Paramsabit       )




	def   UpperCAmelCase_				(			self      :    List[Any]       )	->     Dict:
						from transformers import TaPreTrainedModel

						self.model_fpaa.get_memory_footprint()
						self.model_abit.get_memory_footprint()

						for name, module in self.model_abit.named_modules():
											if isinstance(UpperCAmelCase__      ,      torch.nn.Linear       ):
																if name not in ["lm_head"] + TaPreTrainedModel._keep_in_fpaa_modules:
																					# 4-bit parameters are packed in uint8 variables
																					self.assertTrue(module.weight.dtype == torch.uinta       )




	def   UpperCAmelCase_				(			self      :    Optional[Any]       )	->     Union[str, Any]:
						__SCREAMING_SNAKE_CASE						       =    self.tokenizer(self.input_text      ,      return_tensors="pt"       )
						__SCREAMING_SNAKE_CASE						       =    self.model_abit.generate(input_ids=encoded_input["input_ids"].to(0       )      ,      max_new_tokens=1_0       )

						self.assertIn(self.tokenizer.decode(output_sequences[0]      ,      skip_special_tokens=UpperCAmelCase__       )      ,      self.EXPECTED_OUTPUTS       )




	def   UpperCAmelCase_				(			self      :    Dict       )	->     List[str]:
						__SCREAMING_SNAKE_CASE						       =    BitsAndBytesConfig()
						__SCREAMING_SNAKE_CASE						       =    True

						__SCREAMING_SNAKE_CASE						       =    AutoModelForCausalLM.from_pretrained(
						    self.model_name      ,      quantization_config=UpperCAmelCase__      ,      device_map="auto"       )

						__SCREAMING_SNAKE_CASE						       =    self.tokenizer(self.input_text      ,      return_tensors="pt"       )
						__SCREAMING_SNAKE_CASE						       =    model_abit_from_config.generate(
						    input_ids=encoded_input["input_ids"].to(0       )      ,      max_new_tokens=1_0       )

						self.assertIn(self.tokenizer.decode(output_sequences[0]      ,      skip_special_tokens=UpperCAmelCase__       )      ,      self.EXPECTED_OUTPUTS       )




	def   UpperCAmelCase_				(			self      :    List[Any]       )	->     Optional[Any]:
						with self.assertRaises(UpperCAmelCase__       ), tempfile.TemporaryDirectory() as tmpdirname:
											self.model_abit.save_pretrained(UpperCAmelCase__       )




	def   UpperCAmelCase_				(			self      :    Tuple       )	->     Dict:
						__SCREAMING_SNAKE_CASE						       =    BitsAndBytesConfig()

						with self.assertRaises(UpperCAmelCase__       ):
											__SCREAMING_SNAKE_CASE						       =    AutoModelForCausalLM.from_pretrained(
											    self.model_name      ,      quantization_config=UpperCAmelCase__      ,      load_in_abit=UpperCAmelCase__      ,      device_map="auto"      ,      bnb_abit_quant_type="nf4"      ,      )




	def   UpperCAmelCase_				(			self      :    int       )	->     List[str]:
						with self.assertRaises(UpperCAmelCase__       ):
											# Tries with `str`
											self.model_abit.to("cpu"       )

						with self.assertRaises(UpperCAmelCase__       ):
											# Tries with a `dtype``
											self.model_abit.to(torch.floataa       )

						with self.assertRaises(UpperCAmelCase__       ):
											# Tries with a `device`
											self.model_abit.to(torch.device("cuda:0"       )       )

						with self.assertRaises(UpperCAmelCase__       ):
											# Tries with a `device`
											self.model_abit.float()

						with self.assertRaises(UpperCAmelCase__       ):
											# Tries with a `device`
											self.model_abit.half()

						# Test if we did not break anything
						__SCREAMING_SNAKE_CASE						       =    self.tokenizer(self.input_text      ,      return_tensors="pt"       )

						__SCREAMING_SNAKE_CASE						       =    self.model_fpaa.to(torch.floataa       )
						__SCREAMING_SNAKE_CASE						       =    self.model_fpaa.generate(input_ids=encoded_input["input_ids"].to(0       )      ,      max_new_tokens=1_0       )

						# Check this does not throw an error
						__SCREAMING_SNAKE_CASE						       =    self.model_fpaa.to("cpu"       )

						# Check this does not throw an error
						__SCREAMING_SNAKE_CASE						       =    self.model_fpaa.half()

						# Check this does not throw an error
						__SCREAMING_SNAKE_CASE						       =    self.model_fpaa.float()




	def   UpperCAmelCase_				(			self      :    List[str]       )	->     Dict:
						__SCREAMING_SNAKE_CASE						       =    AutoModelForSeqaSeqLM.from_pretrained("t5-small"      ,      load_in_abit=UpperCAmelCase__      ,      device_map="auto"       )
						self.assertTrue(model.decoder.block[0].layer[2].DenseReluDense.wo.weight.dtype == torch.floataa       )




@require_bitsandbytes
@require_accelerate
@require_torch
@require_torch_gpu
@slow
class 			UpperCamelCase_     (			unittest.TestCase):

	"""simple docstring"""





	@classmethod
	def   UpperCAmelCase_				(			cls      :    Tuple       )	->     List[Any]:
						__SCREAMING_SNAKE_CASE						       =    "t5-small"
						__SCREAMING_SNAKE_CASE						       =    "google/flan-t5-small"  # flan-t5 uses dense-act instead of dense-relu-dense
						__SCREAMING_SNAKE_CASE						       =    AutoTokenizer.from_pretrained(cls.model_name       )
						__SCREAMING_SNAKE_CASE						       =    "Translate in German: Hello, my dog is cute"




	def   UpperCAmelCase_				(			self      :    List[str]       )	->     Tuple:
						gc.collect()
						torch.cuda.empty_cache()




	def   UpperCAmelCase_				(			self      :    List[Any]       )	->     int:
						from transformers import TaForConditionalGeneration

						__SCREAMING_SNAKE_CASE						       =    TaForConditionalGeneration._keep_in_fpaa_modules
						__SCREAMING_SNAKE_CASE						       =    None

						# test with `t5-small`
						__SCREAMING_SNAKE_CASE						       =    TaForConditionalGeneration.from_pretrained(self.model_name      ,      load_in_abit=UpperCAmelCase__      ,      device_map="auto"       )
						__SCREAMING_SNAKE_CASE						       =    self.tokenizer(self.input_text      ,      return_tensors="pt"       ).to(0       )
						__SCREAMING_SNAKE_CASE						       =    model.generate(**UpperCAmelCase__       )

						# test with `flan-t5-small`
						__SCREAMING_SNAKE_CASE						       =    TaForConditionalGeneration.from_pretrained(
						    self.dense_act_model_name      ,      load_in_abit=UpperCAmelCase__      ,      device_map="auto"       )
						__SCREAMING_SNAKE_CASE						       =    self.tokenizer(self.input_text      ,      return_tensors="pt"       ).to(0       )
						__SCREAMING_SNAKE_CASE						       =    model.generate(**UpperCAmelCase__       )
						__SCREAMING_SNAKE_CASE						       =    modules




	def   UpperCAmelCase_				(			self      :    Optional[Any]       )	->     str:
						import bitsandbytes as bnb

						from transformers import TaForConditionalGeneration

						# test with `t5-small`
						__SCREAMING_SNAKE_CASE						       =    TaForConditionalGeneration.from_pretrained(self.model_name      ,      load_in_abit=UpperCAmelCase__      ,      device_map="auto"       )

						# there was a bug with decoders - this test checks that it is fixed
						self.assertTrue(isinstance(model.decoder.block[0].layer[0].SelfAttention.q      ,      bnb.nn.Linearabit       )       )

						__SCREAMING_SNAKE_CASE						       =    self.tokenizer(self.input_text      ,      return_tensors="pt"       ).to(0       )
						__SCREAMING_SNAKE_CASE						       =    model.generate(**UpperCAmelCase__       )

						# test with `flan-t5-small`
						__SCREAMING_SNAKE_CASE						       =    TaForConditionalGeneration.from_pretrained(
						    self.dense_act_model_name      ,      load_in_abit=UpperCAmelCase__      ,      device_map="auto"       )
						__SCREAMING_SNAKE_CASE						       =    self.tokenizer(self.input_text      ,      return_tensors="pt"       ).to(0       )
						__SCREAMING_SNAKE_CASE						       =    model.generate(**UpperCAmelCase__       )




class 			UpperCamelCase_     (			UpperCamelCase):

	"""simple docstring"""





	def   UpperCAmelCase_				(			self      :    Union[str, Any]       )	->     Optional[Any]:
						super().setUp()
						# model_name
						__SCREAMING_SNAKE_CASE						       =    "bigscience/bloom-560m"
						__SCREAMING_SNAKE_CASE						       =    "t5-small"

						# Different types of model

						__SCREAMING_SNAKE_CASE						       =    AutoModel.from_pretrained(self.model_name      ,      load_in_abit=UpperCAmelCase__      ,      device_map="auto"       )
						# Sequence classification model
						__SCREAMING_SNAKE_CASE						       =    AutoModelForSequenceClassification.from_pretrained(
						    self.model_name      ,      load_in_abit=UpperCAmelCase__      ,      device_map="auto"       )
						# CausalLM model
						__SCREAMING_SNAKE_CASE						       =    AutoModelForCausalLM.from_pretrained(self.model_name      ,      load_in_abit=UpperCAmelCase__      ,      device_map="auto"       )
						# Seq2seq model
						__SCREAMING_SNAKE_CASE						       =    AutoModelForSeqaSeqLM.from_pretrained(
						    self.seq_to_seq_name      ,      load_in_abit=UpperCAmelCase__      ,      device_map="auto"       )




	def   UpperCAmelCase_				(			self      :    int       )	->     Optional[int]:
						del self.base_model
						del self.sequence_model
						del self.model_abit
						del self.seq_to_seq_model

						gc.collect()
						torch.cuda.empty_cache()




	def   UpperCAmelCase_				(			self      :    Any       )	->     Optional[int]:
						from bitsandbytes.nn import Paramsabit

						self.assertTrue(self.base_model.h[-1].mlp.dense_ah_to_h.weight.__class__ == Paramsabit       )

						# Other heads should be nn.Parameter
						self.assertTrue(self.model_abit.lm_head.weight.__class__ == torch.nn.Parameter       )
						self.assertTrue(self.sequence_model.score.weight.__class__ == torch.nn.Parameter       )
						self.assertTrue(self.seq_to_seq_model.lm_head.weight.__class__ == torch.nn.Parameter       )




class 			UpperCamelCase_     (			UpperCamelCase):

	"""simple docstring"""





	def   UpperCAmelCase_				(			self      :    List[str]       )	->     Tuple:
						super().setUp()




	def   UpperCAmelCase_				(			self      :    Dict       )	->     Optional[Any]:
						del self.pipe

						gc.collect()
						torch.cuda.empty_cache()




	def   UpperCAmelCase_				(			self      :    Union[str, Any]       )	->     Dict:
						__SCREAMING_SNAKE_CASE						       =    pipeline(
						    "text-generation"      ,      model=self.model_name      ,      model_kwargs={"device_map": "auto", "load_in_4bit": True, "torch_dtype": torch.floataa}      ,      max_new_tokens=self.MAX_NEW_TOKENS      ,      )

						# Real second forward pass
						__SCREAMING_SNAKE_CASE						       =    self.pipe(self.input_text       )
						self.assertIn(pipeline_output[0]["generated_text"]      ,      self.EXPECTED_OUTPUTS       )




@require_torch_multi_gpu
class 			UpperCamelCase_     (			UpperCamelCase):

	"""simple docstring"""





	def   UpperCAmelCase_				(			self      :    Dict       )	->     Dict:
						super().setUp()




	def   UpperCAmelCase_				(			self      :    Any       )	->     Union[str, Any]:
						__SCREAMING_SNAKE_CASE						       =    AutoModelForCausalLM.from_pretrained(
						    self.model_name      ,      load_in_abit=UpperCAmelCase__      ,      device_map="balanced"       )

						# Check correct device map
						self.assertEqual(set(model_parallel.hf_device_map.values()       )      ,      {0, 1}       )

						# Check that inference pass works on the model
						__SCREAMING_SNAKE_CASE						       =    self.tokenizer(self.input_text      ,      return_tensors="pt"       )

						# Second real batch
						__SCREAMING_SNAKE_CASE						       =    model_parallel.generate(input_ids=encoded_input["input_ids"].to(0       )      ,      max_new_tokens=1_0       )
						self.assertIn(self.tokenizer.decode(output_parallel[0]      ,      skip_special_tokens=UpperCAmelCase__       )      ,      self.EXPECTED_OUTPUTS       )




class 			UpperCamelCase_     (			UpperCamelCase):

	"""simple docstring"""





	def   UpperCAmelCase_				(			self      :    Optional[int]       )	->     Optional[int]:
						__SCREAMING_SNAKE_CASE						       =    "facebook/opt-350m"
						super().setUp()




	def   UpperCAmelCase_				(			self      :    int       )	->     Optional[Any]:
						if version.parse(importlib.metadata.version("bitsandbytes"       )       ) < version.parse("0.37.0"       ):
											return

						# Step 1: freeze all parameters
						__SCREAMING_SNAKE_CASE						       =    AutoModelForCausalLM.from_pretrained(self.model_name      ,      load_in_abit=UpperCAmelCase__       )

						self.assertEqual(set(model.hf_device_map.values()       )      ,      {torch.cuda.current_device()}       )

						for param in model.parameters():
											__SCREAMING_SNAKE_CASE						       =    False  # freeze the model - train adapters later
											if param.ndim == 1:
																# cast the small parameters (e.g. layernorm) to fp32 for stability
																__SCREAMING_SNAKE_CASE						       =    param.data.to(torch.floataa       )

        # Step 2: add adapters
						for _, module in model.named_modules():
											if "OPTAttention" in repr(type(UpperCAmelCase__       )       ):
																__SCREAMING_SNAKE_CASE						       =    LoRALayer(module.q_proj      ,      rank=1_6       )
																__SCREAMING_SNAKE_CASE						       =    LoRALayer(module.k_proj      ,      rank=1_6       )
																__SCREAMING_SNAKE_CASE						       =    LoRALayer(module.v_proj      ,      rank=1_6       )

        # Step 3: dummy batch
						__SCREAMING_SNAKE_CASE						       =    self.tokenizer("Test batch "      ,      return_tensors="pt"       ).to(0       )

						# Step 4: Check if the gradient is not None
						with torch.cuda.amp.autocast():
											__SCREAMING_SNAKE_CASE						       =    model.forward(**UpperCAmelCase__       )
											out.logits.norm().backward()

						for module in model.modules():
											if isinstance(UpperCAmelCase__      ,      UpperCAmelCase__       ):
																self.assertTrue(module.adapter[1].weight.grad is not None       )
																self.assertTrue(module.adapter[1].weight.grad.norm().item() > 0       )
											elif isinstance(UpperCAmelCase__      ,      nn.Embedding       ):
																self.assertTrue(module.weight.grad is None       )




class 			UpperCamelCase_     (			UpperCamelCase):

	"""simple docstring"""





	snake_case__    :  List[Any]    =      "gpt2-xl"
	snake_case__    :  Optional[int]    =      3.31_91_85_48_54_15_21_87




"""simple docstring"""







import random
import unittest

import torch

from diffusers import IFInpaintingSuperResolutionPipeline
from diffusers.utils import floats_tensor
from diffusers.utils.import_utils import is_xformers_available
from diffusers.utils.testing_utils import skip_mps, torch_device

from ..pipeline_params import (
    TEXT_GUIDED_IMAGE_INPAINTING_BATCH_PARAMS,
    TEXT_GUIDED_IMAGE_INPAINTING_PARAMS,
)
from ..test_pipelines_common import PipelineTesterMixin
from . import IFPipelineTesterMixin




@skip_mps
class 			UpperCamelCase_     (			UpperCamelCase      ,  UpperCamelCase      ,  unittest.TestCase):

	"""simple docstring"""





	snake_case__    :  str    =      IFInpaintingSuperResolutionPipeline
	snake_case__    :  Optional[Any]    =      TEXT_GUIDED_IMAGE_INPAINTING_PARAMS - {"width", "height"}
	snake_case__    :  Union[str, Any]    =      TEXT_GUIDED_IMAGE_INPAINTING_BATCH_PARAMS.union({"original_image"})
	snake_case__    :  Dict    =      PipelineTesterMixin.required_optional_params - {"latents"}




	def   UpperCAmelCase_				(			self      :    List[Any]       )	->     Optional[int]:
						return self._get_superresolution_dummy_components()




	def   UpperCAmelCase_				(			self      :    List[str]      ,      UpperCAmelCase__      :    List[Any]      ,      UpperCAmelCase__      :    List[str]=0       )	->     List[str]:
						if str(UpperCAmelCase__       ).startswith("mps"       ):
											__SCREAMING_SNAKE_CASE						       =    torch.manual_seed(UpperCAmelCase__       )
						else:
											__SCREAMING_SNAKE_CASE						       =    torch.Generator(device=UpperCAmelCase__       ).manual_seed(UpperCAmelCase__       )

						__SCREAMING_SNAKE_CASE						       =    floats_tensor((1, 3, 1_6, 1_6)      ,      rng=random.Random(UpperCAmelCase__       )       ).to(UpperCAmelCase__       )
						__SCREAMING_SNAKE_CASE						       =    floats_tensor((1, 3, 3_2, 3_2)      ,      rng=random.Random(UpperCAmelCase__       )       ).to(UpperCAmelCase__       )
						__SCREAMING_SNAKE_CASE						       =    floats_tensor((1, 3, 3_2, 3_2)      ,      rng=random.Random(UpperCAmelCase__       )       ).to(UpperCAmelCase__       )

						__SCREAMING_SNAKE_CASE						       =    {
						    "prompt": "A painting of a squirrel eating a burger",
						    "image": image,
						    "original_image": original_image,
						    "mask_image": mask_image,
						    "generator": generator,
						    "num_inference_steps": 2,
						    "output_type": "numpy",
						}

						return inputs




	@unittest.skipIf(
	    torch_device != "cuda" or not is_xformers_available()      ,      reason="XFormers attention is only available with CUDA and `xformers` installed"      ,      )
	def   UpperCAmelCase_				(			self      :    str       )	->     Dict:
						self._test_xformers_attention_forwardGenerator_pass(expected_max_diff=1E-3       )




	def   UpperCAmelCase_				(			self      :    str       )	->     Tuple:
						self._test_save_load_optional_components()




	@unittest.skipIf(torch_device != "cuda"      ,      reason="float16 requires CUDA"       )
	def   UpperCAmelCase_				(			self      :    Tuple       )	->     Tuple:
						# Due to non-determinism in save load of the hf-internal-testing/tiny-random-t5 text encoder
						super().test_save_load_floataa(expected_max_diff=1E-1       )




	def   UpperCAmelCase_				(			self      :    Optional[Any]       )	->     List[Any]:
						self._test_attention_slicing_forward_pass(expected_max_diff=1E-2       )




	def   UpperCAmelCase_				(			self      :    Any       )	->     Union[str, Any]:
						self._test_save_load_local()




	def   UpperCAmelCase_				(			self      :    Union[str, Any]       )	->     Optional[Any]:
						self._test_inference_batch_single_identical(
						    expected_max_diff=1E-2      ,      )
