'''simple docstring'''

import importlib.metadata
import warnings
from copy import deepcopy

from packaging import version

from ..utils import logging
from .import_utils import is_accelerate_available, is_bitsandbytes_available


if is_bitsandbytes_available():
       import bitsandbytes as bnb
       import torch
       import torch.nn as nn

       from ..pytorch_utils import ConvaD

if is_accelerate_available():
       from accelerate import init_empty_weights
       from accelerate.utils import find_tied_parameters

lowerCAmelCase_ : Optional[int]						  =      logging.get_logger(__name__)


def __a (	__lowerCamelCase			:						int ,		__lowerCamelCase			:						Optional[int] ,		__lowerCamelCase			:						Optional[int] ,		__lowerCamelCase			:						Optional[int]=None ,		__lowerCamelCase			:						List[str]=None				)	->							Any:
      '''simple docstring'''


      if "." in tensor_name:
            lowercase_     =      tensor_name.split("."				)
            for split in splits[:-1]:
                  lowercase_     =      getattr(__lowerCamelCase ,		__lowerCamelCase				)
                  if new_module is None:
                        raise ValueError(f'{module} has no attribute {split}.'				)
                  lowercase_     =      new_module
            lowercase_     =      splits[-1]

      if tensor_name not in module._parameters and tensor_name not in module._buffers:
            raise ValueError(f'{module} does not have a parameter or a buffer named {tensor_name}.'				)
      lowercase_     =      tensor_name in module._buffers
      lowercase_     =      getattr(__lowerCamelCase ,		__lowerCamelCase				)

      if old_value.device == torch.device("meta"				) and device not in ["meta", torch.device("meta"				)] and value is None:
            raise ValueError(f'{tensor_name} is on the meta device, we need a `value` to put in on {device}.'				)

      lowercase_     =      False
      lowercase_     =      False
      if is_buffer or not is_bitsandbytes_available():
            lowercase_     =      False
            lowercase_     =      False
      else:
            lowercase_     =      hasattr(bnb.nn ,		"Params4bit"				) and isinstance(module._parameters[tensor_name] ,		bnb.nn.Paramsabit				)
            lowercase_     =      isinstance(module._parameters[tensor_name] ,		bnb.nn.IntaParams				)

      if is_abit or is_abit:
            lowercase_     =      module._parameters[tensor_name]
            if param.device.type != "cuda":
                  if value is None:
                        lowercase_     =      old_value.to(__lowerCamelCase				)
                  elif isinstance(__lowerCamelCase ,		torch.Tensor				):
                        lowercase_     =      value.to("cpu"				)
                        if value.dtype == torch.inta:
                              lowercase_     =      version.parse(importlib.metadata.version("bitsandbytes"				)				) > version.parse(
                                  "0.37.2"				)
                              if not is_abit_serializable:
                                    raise ValueError(
                                        "Detected int8 weights but the version of bitsandbytes is not compatible with int8 serialization. "
                                        "Make sure to download the latest `bitsandbytes` version. `pip install --upgrade bitsandbytes`."				)
                  else:
                        lowercase_     =      torch.tensor(__lowerCamelCase ,		device="cpu"				)

                  # Support models using `Conv1D` in place of `nn.Linear` (e.g. gpt2) by transposing the weight matrix prior to quantization.
                  # Since weights are saved in the correct "orientation", we skip transposing when loading.
                  if issubclass(module.source_cls ,		__lowerCamelCase				) and fpaa_statistics is None:
                        lowercase_     =      new_value.T

                  lowercase_     =      old_value.__dict__
                  if is_abit:
                        lowercase_     =      bnb.nn.IntaParams(__lowerCamelCase ,		requires_grad=__lowerCamelCase ,		**__lowerCamelCase				).to(__lowerCamelCase				)
                  elif is_abit:
                        lowercase_     =      bnb.nn.Paramsabit(__lowerCamelCase ,		requires_grad=__lowerCamelCase ,		**__lowerCamelCase				).to(__lowerCamelCase				)

                  lowercase_     =      new_value
                  if fpaa_statistics is not None:
                        setattr(module.weight ,		"SCB" ,		fpaa_statistics.to(__lowerCamelCase				)				)

      else:
            if value is None:
                  lowercase_     =      old_value.to(__lowerCamelCase				)
            elif isinstance(__lowerCamelCase ,		torch.Tensor				):
                  lowercase_     =      value.to(__lowerCamelCase				)
            else:
                  lowercase_     =      torch.tensor(__lowerCamelCase ,		device=__lowerCamelCase				)

            if is_buffer:
                  lowercase_     =      new_value
            else:
                  lowercase_     =      nn.Parameter(__lowerCamelCase ,		requires_grad=old_value.requires_grad				)
                  lowercase_     =      new_value


def __a (	__lowerCamelCase			:						int ,		__lowerCamelCase			:						List[Any]=None ,		__lowerCamelCase			:						Optional[Any]=None ,		__lowerCamelCase			:						Optional[Any]=None ,		__lowerCamelCase			:						List[str]=False				)	->							Tuple:
      '''simple docstring'''


      for name, module in model.named_children():
            if current_key_name is None:
                  lowercase_     =      []
            current_key_name.append(__lowerCamelCase				)

            if (isinstance(__lowerCamelCase ,		nn.Linear				) or isinstance(__lowerCamelCase ,		__lowerCamelCase				)) and name not in modules_to_not_convert:
                  # Check if the current key is not in the `modules_to_not_convert`
                  if not any(key in ".".join(__lowerCamelCase				) for key in modules_to_not_convert				):
                        with init_empty_weights():
                              if isinstance(__lowerCamelCase ,		__lowerCamelCase				):
                                    lowercase_ ,							lowercase_     =      module.weight.shape
                              else:
                                    lowercase_     =      module.in_features
                                    lowercase_     =      module.out_features

                              if quantization_config.quantization_method() == "llm_int8":
                                    lowercase_     =      bnb.nn.LinearabitLt(
                                        __lowerCamelCase ,		__lowerCamelCase ,		module.bias is not None ,		has_fpaa_weights=quantization_config.llm_inta_has_fpaa_weight ,		threshold=quantization_config.llm_inta_threshold ,		)
                                    lowercase_     =      True
                              else:
                                    if (
                                        quantization_config.llm_inta_skip_modules is not None
                                        and name in quantization_config.llm_inta_skip_modules
                                    ):
                                          pass
                                    else:
                                          lowercase_     =      bnb.nn.Linearabit(
                                              __lowerCamelCase ,		__lowerCamelCase ,		module.bias is not None ,		quantization_config.bnb_abit_compute_dtype ,		compress_statistics=quantization_config.bnb_abit_use_double_quant ,		quant_type=quantization_config.bnb_abit_quant_type ,		)
                                          lowercase_     =      True
                    # Store the module class in case we need to transpose the weight later
                              lowercase_     =      type(__lowerCamelCase				)
                              # Force requires grad to False to avoid unexpected errors
                              model._modules[name].requires_grad_(__lowerCamelCase				)
            if len(list(module.children()				)				) > 0:
                  lowercase_ ,							lowercase_     =      _replace_with_bnb_linear(
                      __lowerCamelCase ,		__lowerCamelCase ,		__lowerCamelCase ,		__lowerCamelCase ,		has_been_replaced=__lowerCamelCase ,		)
            # Remove the last key for recursion
            current_key_name.pop(-1				)
      return model, has_been_replaced


def __a (	__lowerCamelCase			:						int ,		__lowerCamelCase			:						List[Any]=None ,		__lowerCamelCase			:						int=None ,		__lowerCamelCase			:						Tuple=None				)	->							Tuple:
      '''simple docstring'''


      lowercase_     =      ["lm_head"] if modules_to_not_convert is None else modules_to_not_convert
      lowercase_ ,							lowercase_     =      _replace_with_bnb_linear(
          __lowerCamelCase ,		__lowerCamelCase ,		__lowerCamelCase ,		__lowerCamelCase				)

      if not has_been_replaced:
            logger.warning(
                "You are loading your model in 8bit or 4bit but no linear modules were found in your model."
                " Please double check your model architecture, or submit an issue on github if you think this is"
                " a bug."				)

      return model


def __a (	*__lowerCamelCase			:						Optional[int] ,		**__lowerCamelCase			:						List[str]				)	->							Union[str, Any]:
      '''simple docstring'''


      warnings.warn(
          "`replace_8bit_linear` will be deprecated in a future version, please use `replace_with_bnb_linear` instead" ,		__lowerCamelCase ,		)
      return replace_with_bnb_linear(*__lowerCamelCase ,		**__lowerCamelCase				)


def __a (	*__lowerCamelCase			:						int ,		**__lowerCamelCase			:						Any				)	->							int:
      '''simple docstring'''


      warnings.warn(
          "`set_module_8bit_tensor_to_device` will be deprecated in a future version, please use `set_module_quantized_tensor_to_device` instead" ,		__lowerCamelCase ,		)
      return set_module_quantized_tensor_to_device(*__lowerCamelCase ,		**__lowerCamelCase				)






def __a (	__lowerCamelCase			:						Dict				)	->							str:
      '''simple docstring'''


      lowercase_     =      deepcopy(__lowerCamelCase				)  # this has 0 cost since it is done inside `init_empty_weights` context manager`
      tied_model.tie_weights()

      lowercase_     =      find_tied_parameters(__lowerCamelCase				)
      # For compatibility with Accelerate < 0.18
      if isinstance(__lowerCamelCase ,		__lowerCamelCase				):
            lowercase_     =      sum(list(tied_params.values()				) ,		[]				) + list(tied_params.keys()				)
      else:
            lowercase_     =      sum(__lowerCamelCase ,		[]				)
      lowercase_     =      len(__lowerCamelCase				) > 0

      # Check if it is a base model
      lowercase_     =      not hasattr(__lowerCamelCase ,		model.base_model_prefix				)

      # Ignore this for base models (BertModel, GPT2Model, etc.)
      if (not has_tied_params) and is_base_model:
            return []

      # otherwise they have an attached head
      lowercase_     =      list(model.named_children()				)
      lowercase_     =      [list_modules[-1][0]]

      # add last module together with tied weights
      lowercase_     =      set(__lowerCamelCase				) - set(__lowerCamelCase				)
      lowercase_     =      list(set(__lowerCamelCase				)				) + list(__lowerCamelCase				)

      # remove ".weight" from the keys
      lowercase_     =      [".weight", ".bias"]
      lowercase_     =      []
      for name in list_untouched:
            for name_to_remove in names_to_remove:
                  if name_to_remove in name:
                        lowercase_     =      name.replace(__lowerCamelCase ,		""				)
            filtered_module_names.append(__lowerCamelCase				)

      return filtered_module_names







'''simple docstring'''

import gc
import unittest

from diffusers import FlaxControlNetModel, FlaxStableDiffusionControlNetPipeline
from diffusers.utils import is_flax_available, load_image, slow
from diffusers.utils.testing_utils import require_flax


if is_flax_available():
       import jax
       import jax.numpy as jnp
       from flax.jax_utils import replicate
       from flax.training.common_utils import shard





@slow
@require_flax
class 							lowercase						(						unittest.TestCase ):


   def __UpperCAmelCase       (   self				:	int)      ->					List[str]:
         # clean up the VRAM after each test
         super().tearDown()
         gc.collect()


   def __UpperCAmelCase       (   self				:	Dict)      ->					Any:
         lowercase_ ,							lowercase_     =      FlaxControlNetModel.from_pretrained(
             "lllyasviel/sd-controlnet-canny"   ,		from_pt=__lowerCAmelCase   ,		dtype=jnp.bfloataa)
         lowercase_ ,							lowercase_     =      FlaxStableDiffusionControlNetPipeline.from_pretrained(
             "runwayml/stable-diffusion-v1-5"   ,		controlnet=__lowerCAmelCase   ,		from_pt=__lowerCAmelCase   ,		dtype=jnp.bfloataa)
         lowercase_     =      controlnet_params

         lowercase_     =      "bird"
         lowercase_     =      jax.device_count()
         lowercase_     =      pipe.prepare_text_inputs([prompts] * num_samples)

         lowercase_     =      load_image(
             "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/sd_controlnet/bird_canny.png")
         lowercase_     =      pipe.prepare_image_inputs([canny_image] * num_samples)

         lowercase_     =      jax.random.PRNGKey(0)
         lowercase_     =      jax.random.split(__lowerCAmelCase   ,		jax.device_count())

         lowercase_     =      replicate(__lowerCAmelCase)
         lowercase_     =      shard(__lowerCAmelCase)
         lowercase_     =      shard(__lowerCAmelCase)

         lowercase_     =      pipe(
             prompt_ids=__lowerCAmelCase   ,		image=__lowerCAmelCase   ,		params=__lowerCAmelCase   ,		prng_seed=__lowerCAmelCase   ,		num_inference_steps=50   ,		jit=__lowerCAmelCase   ,		).images
         assert images.shape == (jax.device_count(), 1, 768, 512, 3)

         lowercase_     =      images.reshape((images.shape[0] * images.shape[1],) + images.shape[-3:])
         lowercase_     =      images[0, 253:256, 253:256, -1]

         lowercase_     =      jnp.asarray(jax.device_get(image_slice.flatten()))
         lowercase_     =      jnp.array(
             [0.16_7969, 0.11_6699, 0.08_1543, 0.15_4297, 0.13_2812, 0.10_8887, 0.16_9922, 0.16_9922, 0.20_5078])
         print(F'output_slice: {output_slice}')
         assert jnp.abs(output_slice - expected_slice).max() < 1e-2


   def __UpperCAmelCase       (   self				:	Tuple)      ->					List[str]:
         lowercase_ ,							lowercase_     =      FlaxControlNetModel.from_pretrained(
             "lllyasviel/sd-controlnet-openpose"   ,		from_pt=__lowerCAmelCase   ,		dtype=jnp.bfloataa)
         lowercase_ ,							lowercase_     =      FlaxStableDiffusionControlNetPipeline.from_pretrained(
             "runwayml/stable-diffusion-v1-5"   ,		controlnet=__lowerCAmelCase   ,		from_pt=__lowerCAmelCase   ,		dtype=jnp.bfloataa)
         lowercase_     =      controlnet_params

         lowercase_     =      "Chef in the kitchen"
         lowercase_     =      jax.device_count()
         lowercase_     =      pipe.prepare_text_inputs([prompts] * num_samples)

         lowercase_     =      load_image(
             "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/sd_controlnet/pose.png")
         lowercase_     =      pipe.prepare_image_inputs([pose_image] * num_samples)

         lowercase_     =      jax.random.PRNGKey(0)
         lowercase_     =      jax.random.split(__lowerCAmelCase   ,		jax.device_count())

         lowercase_     =      replicate(__lowerCAmelCase)
         lowercase_     =      shard(__lowerCAmelCase)
         lowercase_     =      shard(__lowerCAmelCase)

         lowercase_     =      pipe(
             prompt_ids=__lowerCAmelCase   ,		image=__lowerCAmelCase   ,		params=__lowerCAmelCase   ,		prng_seed=__lowerCAmelCase   ,		num_inference_steps=50   ,		jit=__lowerCAmelCase   ,		).images
         assert images.shape == (jax.device_count(), 1, 768, 512, 3)

         lowercase_     =      images.reshape((images.shape[0] * images.shape[1],) + images.shape[-3:])
         lowercase_     =      images[0, 253:256, 253:256, -1]

         lowercase_     =      jnp.asarray(jax.device_get(image_slice.flatten()))
         lowercase_     =      jnp.array(
             [[0.27_1484, 0.26_1719, 0.27_5391, 0.27_7344, 0.27_9297, 0.29_1016, 0.29_4922, 0.30_2734, 0.30_2734]])
         print(F'output_slice: {output_slice}')
         assert jnp.abs(output_slice - expected_slice).max() < 1e-2






