


'''simple docstring'''






# Lint as: python3
import os
import re
import urllib.parse
from pathlib import Path
from typing import Callable, List, Optional, Union
from zipfile import ZipFile

from ..utils.file_utils import cached_path, hf_github_url
from ..utils.logging import get_logger
from ..utils.version import Version


lowercase							=get_logger(__name__)






class  __magic_name__							:
       UpperCAmelCase										="dummy_data"
       UpperCAmelCase										="datasets"
       UpperCAmelCase										=False






       def __init__(      self  ,							snake_case  ,							snake_case  ,							snake_case  ,							snake_case = None  ,							snake_case = False  ,							snake_case = True  ,							snake_case = None  ,							)						->  Dict:




          '''simple docstring'''
          _UpperCAmelCase       :       Optional[int]							   =0
          _UpperCAmelCase       :       str							   =dataset_name
          _UpperCAmelCase       :       Optional[Any]							   =cache_dir
          _UpperCAmelCase       :       str							   =use_local_dummy_data
          _UpperCAmelCase       :       Optional[int]							   =config
          # download_callbacks take a single url as input
          _UpperCAmelCase       :       List[Callable]							   =download_callbacks or []
          # if False, it doesn't load existing files and it returns the paths of the dummy files relative
          # to the dummy_data zip file root
          _UpperCAmelCase       :       List[str]							   =load_existing_dummy_data

          # TODO(PVP, QL) might need to make this more general
          _UpperCAmelCase       :       Tuple							   =str(snake_case)
          # to be downloaded
          _UpperCAmelCase       :       Dict							   =None
          _UpperCAmelCase       :       List[Any]							   =None






       @property
       def       lowerCAmelCase		(      self)						->  List[Any]:




          '''simple docstring'''
          if self._dummy_file is None:
             _UpperCAmelCase       :       Dict							   =self.download_dummy_data()
          return self._dummy_file






       @property
       def       lowerCAmelCase		(      self)						->  Tuple:




          '''simple docstring'''
          if self.config is not None:
             # structure is dummy / config_name / version_name
             return os.path.join('dummy'  ,							self.config.name  ,							self.version_name)
          # structure is dummy / version_name
          return os.path.join('dummy'  ,							self.version_name)






       @property
       def       lowerCAmelCase		(      self)						->  Dict:




          '''simple docstring'''
          return os.path.join(self.dummy_data_folder  ,							'dummy_data.zip')






       def       lowerCAmelCase		(      self)						->  Union[str, Any]:




          '''simple docstring'''
          _UpperCAmelCase       :       Optional[int]							   =(
              self.local_path_to_dummy_data if self.use_local_dummy_data is True else self.github_path_to_dummy_data
          )

          _UpperCAmelCase       :       List[Any]							   =cached_path(
              snake_case  ,							cache_dir=self.cache_dir  ,							extract_compressed_file=snake_case  ,							force_extract=snake_case)

          return os.path.join(snake_case  ,							self.dummy_file_name)






       @property
       def       lowerCAmelCase		(      self)						->  Dict:




          '''simple docstring'''
          return os.path.join(self.datasets_scripts_dir  ,							self.dataset_name  ,							self.dummy_zip_file)






       @property
       def       lowerCAmelCase		(      self)						->  List[Any]:




          '''simple docstring'''
          if self._bucket_url is None:
             _UpperCAmelCase       :       str							   =hf_github_url(self.dataset_name  ,							self.dummy_zip_file.replace(os.sep  ,							'/'))
          return self._bucket_url






       @property
       def       lowerCAmelCase		(      self)						->  Any:




          '''simple docstring'''
          # return full path if its a dir
          if os.path.isdir(self.dummy_file):
             return self.dummy_file
          # else cut off path to file -> example `xsum`.
          return "/".join(self.dummy_file.replace(os.sep  ,							'/').split('/')[:-1])






       def       lowerCAmelCase		(      self  ,							snake_case  ,							*snake_case)						->  Optional[int]:




          '''simple docstring'''
          if self.load_existing_dummy_data:
             # dummy data is downloaded and tested
             _UpperCAmelCase       :       Optional[int]							   =self.dummy_file
          else:
             # dummy data cannot be downloaded and only the path to dummy file is returned
             _UpperCAmelCase       :       List[Any]							   =self.dummy_file_name

          # special case when data_url is a dict
          if isinstance(snake_case  ,							snake_case):
             return self.create_dummy_data_dict(snake_case  ,							snake_case)
          elif isinstance(snake_case  ,							(list, tuple)):
             return self.create_dummy_data_list(snake_case  ,							snake_case)
          else:
             return self.create_dummy_data_single(snake_case  ,							snake_case)






       def       lowerCAmelCase		(      self  ,							snake_case  ,							*snake_case)						->  str:




          '''simple docstring'''
          return self.download_and_extract(snake_case)






       def       lowerCAmelCase		(      self  ,							snake_case  ,							snake_case)						->  Union[str, Any]:




          '''simple docstring'''
          return self.download_and_extract(snake_case)






       def       lowerCAmelCase		(      self  ,							snake_case  ,							*snake_case  ,							**snake_case)						->  Tuple:




          '''simple docstring'''
          return path






       def       lowerCAmelCase		(      self)						->  Tuple:




          '''simple docstring'''
          return {}






       def       lowerCAmelCase		(      self  ,							snake_case  ,							snake_case)						->  int:




          '''simple docstring'''
          _UpperCAmelCase       :       Tuple							   ={}
          for key, single_urls in data_url.items():
             for download_callback in self.download_callbacks:
                if isinstance(snake_case  ,							snake_case):
                   for single_url in single_urls:
                      download_callback(snake_case)
                else:
                   _UpperCAmelCase       :       Any							   =single_urls
                   download_callback(snake_case)
            # we force the name of each key to be the last file / folder name of the url path
            # if the url has arguments, we need to encode them with urllib.parse.quote_plus
             if isinstance(snake_case  ,							snake_case):
                _UpperCAmelCase       :       Union[str, Any]							   =[os.path.join(snake_case  ,							urllib.parse.quote_plus(Path(snake_case).name)) for x in single_urls]
             else:
                _UpperCAmelCase       :       Dict							   =single_urls
                _UpperCAmelCase       :       Optional[int]							   =os.path.join(snake_case  ,							urllib.parse.quote_plus(Path(snake_case).name))
             _UpperCAmelCase       :       str							   =value

          # make sure that values are unique
          if all(isinstance(snake_case  ,							snake_case) for i in dummy_data_dict.values()) and len(set(dummy_data_dict.values())) < len(
              dummy_data_dict.values()):
             # append key to value to make its name unique
             _UpperCAmelCase       :       List[str]							   ={key: value + key for key, value in dummy_data_dict.items()}

          return dummy_data_dict






       def       lowerCAmelCase		(      self  ,							snake_case  ,							snake_case)						->  Any:




          '''simple docstring'''
          _UpperCAmelCase       :       Tuple							   =[]
          # trick: if there are many shards named like `data.txt-000001-of-00300`, only use the first one
          _UpperCAmelCase       :       Any							   =all(bool(re.findall('[0-9]{3,}-of-[0-9]{3,}'  ,							snake_case)) for url in data_url)
          _UpperCAmelCase       :       Optional[int]							   =all(
              url.startswith('https://ftp.ncbi.nlm.nih.gov/pubmed/baseline/pubmed') for url in data_url)
          if data_url and (is_tf_records or is_pubmed_records):
             _UpperCAmelCase       :       List[Any]							   =[data_url[0]] * len(snake_case)
          for single_url in data_url:
             for download_callback in self.download_callbacks:
                download_callback(snake_case)
             # we force the name of each key to be the last file / folder name of the url path
             # if the url has arguments, we need to encode them with urllib.parse.quote_plus
             _UpperCAmelCase       :       Optional[Any]							   =os.path.join(snake_case  ,							urllib.parse.quote_plus(single_url.split('/')[-1]))
             dummy_data_list.append(snake_case)
          return dummy_data_list






       def       lowerCAmelCase		(      self  ,							snake_case  ,							snake_case)						->  List[str]:




          '''simple docstring'''
          for download_callback in self.download_callbacks:
             download_callback(snake_case)
          # we force the name of each key to be the last file / folder name of the url path
          # if the url has arguments, we need to encode them with urllib.parse.quote_plus
          _UpperCAmelCase       :       int							   =os.path.join(snake_case  ,							urllib.parse.quote_plus(data_url.split('/')[-1]))
          if os.path.exists(snake_case) or not self.load_existing_dummy_data:
             return value
          else:
             # Backward compatibility, maybe deprecate at one point.
             # For many datasets with single url calls to dl_manager.download_and_extract,
             # the dummy_data.zip file is actually the zipped downloaded file
             # while now we expected the dummy_data.zip file to be a directory containing
             # the downloaded file.
             return path_to_dummy_data






       def       lowerCAmelCase		(      self)						->  Tuple:




          '''simple docstring'''
          pass






       def       lowerCAmelCase		(      self)						->  Union[str, Any]:




          '''simple docstring'''
          pass






       def       lowerCAmelCase		(      self  ,							snake_case)						->  Union[str, Any]:




          '''simple docstring'''
          def _iter_archive_members(snake_case):
             # this preserves the order of the members inside the ZIP archive
             _UpperCAmelCase       :       Union[str, Any]							   =Path(self.dummy_file).parent
             _UpperCAmelCase       :       Union[str, Any]							   =path.relative_to(snake_case)
             with ZipFile(self.local_path_to_dummy_data) as zip_file:
                _UpperCAmelCase       :       List[str]							   =zip_file.namelist()
             for member in members:
                if member.startswith(relative_path.as_posix()):
                   yield dummy_parent_path.joinpath(snake_case)

          _UpperCAmelCase       :       List[Any]							   =Path(snake_case)
          _UpperCAmelCase       :       Union[str, Any]							   =_iter_archive_members(snake_case) if self.use_local_dummy_data else path.rglob('*')
          for file_path in file_paths:
             if file_path.is_file() and not file_path.name.startswith(('.', '__')):
                yield file_path.relative_to(snake_case).as_posix(), file_path.open('rb')






       def       lowerCAmelCase		(      self  ,							snake_case)						->  List[str]:




          '''simple docstring'''
          if not isinstance(snake_case  ,							snake_case):
             _UpperCAmelCase       :       Optional[int]							   =[paths]
          for path in paths:
             if os.path.isfile(snake_case):
                if os.path.basename(snake_case).startswith(('.', '__')):
                   return
                yield path
             else:
                for dirpath, dirnames, filenames in os.walk(snake_case):
                   if os.path.basename(snake_case).startswith(('.', '__')):
                      continue
                   dirnames.sort()
                   for filename in sorted(snake_case):
                      if filename.startswith(('.', '__')):
                         continue
                      yield os.path.join(snake_case  ,							snake_case)





'''simple docstring'''






from typing import Optional, Tuple, Union

import flax
import flax.linen as nn
import jax
import jax.numpy as jnp
from flax.core.frozen_dict import FrozenDict

from ..configuration_utils import ConfigMixin, flax_register_to_config
from ..utils import BaseOutput
from .embeddings_flax import FlaxTimestepEmbedding, FlaxTimesteps
from .modeling_flax_utils import FlaxModelMixin
from .unet_ad_blocks_flax import (
    FlaxCrossAttnDownBlockaD,
    FlaxDownBlockaD,
    FlaxUNetMidBlockaDCrossAttn,
)






@flax.struct.dataclass
class  __magic_name__							(		lowerCAmelCase    ):
       UpperCAmelCase										=42
       UpperCAmelCase										=42






class  __magic_name__							(		nn.Module    ):
       UpperCAmelCase										=42
       UpperCAmelCase										=(1_6, 3_2, 9_6, 2_5_6)
       UpperCAmelCase										=jnp.floataa






       def       lowerCAmelCase		(      self)						->  Optional[int]:




          '''simple docstring'''
          _UpperCAmelCase       :       str							   =nn.Conv(
              self.block_out_channels[0]  ,							kernel_size=(3, 3)  ,							padding=((1, 1), (1, 1))  ,							dtype=self.dtype  ,							)

          _UpperCAmelCase       :       Tuple							   =[]
          for i in range(len(self.block_out_channels) - 1):
             _UpperCAmelCase       :       Optional[int]							   =self.block_out_channels[i]
             _UpperCAmelCase       :       List[Any]							   =self.block_out_channels[i + 1]
             _UpperCAmelCase       :       Tuple							   =nn.Conv(
                 snake_case  ,							kernel_size=(3, 3)  ,							padding=((1, 1), (1, 1))  ,							dtype=self.dtype  ,							)
             blocks.append(snake_case)
             _UpperCAmelCase       :       Optional[int]							   =nn.Conv(
                 snake_case  ,							kernel_size=(3, 3)  ,							strides=(2, 2)  ,							padding=((1, 1), (1, 1))  ,							dtype=self.dtype  ,							)
             blocks.append(snake_case)
          _UpperCAmelCase       :       Dict							   =blocks

          _UpperCAmelCase       :       Tuple							   =nn.Conv(
              self.conditioning_embedding_channels  ,							kernel_size=(3, 3)  ,							padding=((1, 1), (1, 1))  ,							kernel_init=nn.initializers.zeros_init()  ,							bias_init=nn.initializers.zeros_init()  ,							dtype=self.dtype  ,							)






       def __call__(      self  ,							snake_case)						->  List[str]:




          '''simple docstring'''
          _UpperCAmelCase       :       int							   =self.conv_in(snake_case)
          _UpperCAmelCase       :       Any							   =nn.silu(snake_case)

          for block in self.blocks:
             _UpperCAmelCase       :       Optional[Any]							   =block(snake_case)
             _UpperCAmelCase       :       Union[str, Any]							   =nn.silu(snake_case)

          _UpperCAmelCase       :       str							   =self.conv_out(snake_case)

          return embedding






@flax_register_to_config
class  __magic_name__							(		nn.Module ,lowerCAmelCase ,lowerCAmelCase    ):
       UpperCAmelCase										=3_2
       UpperCAmelCase										=4
       UpperCAmelCase										=(
           "CrossAttnDownBlock2D",
           "CrossAttnDownBlock2D",
           "CrossAttnDownBlock2D",
           "DownBlock2D",
       )
       UpperCAmelCase										=False
       UpperCAmelCase										=(3_2_0, 6_4_0, 1_2_8_0, 1_2_8_0)
       UpperCAmelCase										=2
       UpperCAmelCase										=8
       UpperCAmelCase										=None
       UpperCAmelCase										=1_2_8_0
       UpperCAmelCase										=0.0
       UpperCAmelCase										=False
       UpperCAmelCase										=jnp.floataa
       UpperCAmelCase										=True
       UpperCAmelCase										=0
       UpperCAmelCase										="rgb"
       UpperCAmelCase										=(1_6, 3_2, 9_6, 2_5_6)






       def       lowerCAmelCase		(      self  ,							snake_case)						->  FrozenDict:




          '''simple docstring'''
          # init input tensors
          _UpperCAmelCase       :       Any							   =(1, self.in_channels, self.sample_size, self.sample_size)
          _UpperCAmelCase       :       Optional[Any]							   =jnp.zeros(snake_case  ,							dtype=jnp.floataa)
          _UpperCAmelCase       :       Optional[int]							   =jnp.ones((1,)  ,							dtype=jnp.intaa)
          _UpperCAmelCase       :       str							   =jnp.zeros((1, 1, self.cross_attention_dim)  ,							dtype=jnp.floataa)
          _UpperCAmelCase       :       Optional[Any]							   =(1, 3, self.sample_size * 8, self.sample_size * 8)
          _UpperCAmelCase       :       int							   =jnp.zeros(snake_case  ,							dtype=jnp.floataa)

          _UpperCAmelCase   ,	_UpperCAmelCase       :       List[Any]							   =jax.random.split(snake_case)
          _UpperCAmelCase       :       str							   ={'params': params_rng, 'dropout': dropout_rng}

          return self.init(snake_case  ,							snake_case  ,							snake_case  ,							snake_case  ,							snake_case)["params"]






       def       lowerCAmelCase		(      self)						->  Tuple:




          '''simple docstring'''
          _UpperCAmelCase       :       Optional[Any]							   =self.block_out_channels
          _UpperCAmelCase       :       Tuple							   =block_out_channels[0] * 4

          # If `num_attention_heads` is not defined (which is the case for most models)
          # it will default to `attention_head_dim`. This looks weird upon first reading it and it is.
          # The reason for this behavior is to correct for incorrectly named variables that were introduced
          # when this library was created. The incorrect naming was only discovered much later in https://github.com/huggingface/diffusers/issues/2011#issuecomment-1547958131
          # Changing `attention_head_dim` to `num_attention_heads` for 40,000+ configurations is too backwards breaking
          # which is why we correct for the naming here.
          _UpperCAmelCase       :       Optional[Any]							   =self.num_attention_heads or self.attention_head_dim

          # input
          _UpperCAmelCase       :       Tuple							   =nn.Conv(
              block_out_channels[0]  ,							kernel_size=(3, 3)  ,							strides=(1, 1)  ,							padding=((1, 1), (1, 1))  ,							dtype=self.dtype  ,							)

          # time
          _UpperCAmelCase       :       Union[str, Any]							   =FlaxTimesteps(
              block_out_channels[0]  ,							flip_sin_to_cos=self.flip_sin_to_cos  ,							freq_shift=self.config.freq_shift)
          _UpperCAmelCase       :       str							   =FlaxTimestepEmbedding(snake_case  ,							dtype=self.dtype)

          _UpperCAmelCase       :       Optional[Any]							   =FlaxControlNetConditioningEmbedding(
              conditioning_embedding_channels=block_out_channels[0]  ,							block_out_channels=self.conditioning_embedding_out_channels  ,							)

          _UpperCAmelCase       :       Optional[int]							   =self.only_cross_attention
          if isinstance(snake_case  ,							snake_case):
             _UpperCAmelCase       :       Dict							   =(only_cross_attention,) * len(self.down_block_types)

          if isinstance(snake_case  ,							snake_case):
             _UpperCAmelCase       :       Optional[Any]							   =(num_attention_heads,) * len(self.down_block_types)

          # down
          _UpperCAmelCase       :       int							   =[]
          _UpperCAmelCase       :       Optional[int]							   =[]

          _UpperCAmelCase       :       List[str]							   =block_out_channels[0]

          _UpperCAmelCase       :       int							   =nn.Conv(
              snake_case  ,							kernel_size=(1, 1)  ,							padding='VALID'  ,							kernel_init=nn.initializers.zeros_init()  ,							bias_init=nn.initializers.zeros_init()  ,							dtype=self.dtype  ,							)
          controlnet_down_blocks.append(snake_case)

          for i, down_block_type in enumerate(self.down_block_types):
             _UpperCAmelCase       :       Tuple							   =output_channel
             _UpperCAmelCase       :       Dict							   =block_out_channels[i]
             _UpperCAmelCase       :       str							   =i == len(snake_case) - 1

             if down_block_type == "CrossAttnDownBlock2D":
                _UpperCAmelCase       :       Tuple							   =FlaxCrossAttnDownBlockaD(
                    in_channels=snake_case  ,							out_channels=snake_case  ,							dropout=self.dropout  ,							num_layers=self.layers_per_block  ,							num_attention_heads=num_attention_heads[i]  ,							add_downsample=not is_final_block  ,							use_linear_projection=self.use_linear_projection  ,							only_cross_attention=only_cross_attention[i]  ,							dtype=self.dtype  ,							)
             else:
                _UpperCAmelCase       :       Optional[Any]							   =FlaxDownBlockaD(
                    in_channels=snake_case  ,							out_channels=snake_case  ,							dropout=self.dropout  ,							num_layers=self.layers_per_block  ,							add_downsample=not is_final_block  ,							dtype=self.dtype  ,							)

             down_blocks.append(snake_case)

             for _ in range(self.layers_per_block):
                _UpperCAmelCase       :       Tuple							   =nn.Conv(
                    snake_case  ,							kernel_size=(1, 1)  ,							padding='VALID'  ,							kernel_init=nn.initializers.zeros_init()  ,							bias_init=nn.initializers.zeros_init()  ,							dtype=self.dtype  ,							)
                controlnet_down_blocks.append(snake_case)

             if not is_final_block:
                _UpperCAmelCase       :       List[str]							   =nn.Conv(
                    snake_case  ,							kernel_size=(1, 1)  ,							padding='VALID'  ,							kernel_init=nn.initializers.zeros_init()  ,							bias_init=nn.initializers.zeros_init()  ,							dtype=self.dtype  ,							)
                controlnet_down_blocks.append(snake_case)

          _UpperCAmelCase       :       List[Any]							   =down_blocks
          _UpperCAmelCase       :       Optional[Any]							   =controlnet_down_blocks

          # mid
          _UpperCAmelCase       :       int							   =block_out_channels[-1]
          _UpperCAmelCase       :       Optional[Any]							   =FlaxUNetMidBlockaDCrossAttn(
              in_channels=snake_case  ,							dropout=self.dropout  ,							num_attention_heads=num_attention_heads[-1]  ,							use_linear_projection=self.use_linear_projection  ,							dtype=self.dtype  ,							)

          _UpperCAmelCase       :       Optional[int]							   =nn.Conv(
              snake_case  ,							kernel_size=(1, 1)  ,							padding='VALID'  ,							kernel_init=nn.initializers.zeros_init()  ,							bias_init=nn.initializers.zeros_init()  ,							dtype=self.dtype  ,							)






       def __call__(      self  ,							snake_case  ,							snake_case  ,							snake_case  ,							snake_case  ,							snake_case = 1.0  ,							snake_case = True  ,							snake_case = False  ,							)						->  Union[FlaxControlNetOutput, Tuple]:




          '''simple docstring'''
          _UpperCAmelCase       :       Union[str, Any]							   =self.controlnet_conditioning_channel_order
          if channel_order == "bgr":
             _UpperCAmelCase       :       Optional[int]							   =jnp.flip(snake_case  ,							axis=1)

          # 1. time
          if not isinstance(snake_case  ,							jnp.ndarray):
             _UpperCAmelCase       :       Optional[int]							   =jnp.array([timesteps]  ,							dtype=jnp.intaa)
          elif isinstance(snake_case  ,							jnp.ndarray) and len(timesteps.shape) == 0:
             _UpperCAmelCase       :       str							   =timesteps.astype(dtype=jnp.floataa)
             _UpperCAmelCase       :       Dict							   =jnp.expand_dims(snake_case  ,							0)

          _UpperCAmelCase       :       int							   =self.time_proj(snake_case)
          _UpperCAmelCase       :       Any							   =self.time_embedding(snake_case)

          # 2. pre-process
          _UpperCAmelCase       :       str							   =jnp.transpose(snake_case  ,							(0, 2, 3, 1))
          _UpperCAmelCase       :       Any							   =self.conv_in(snake_case)

          _UpperCAmelCase       :       List[str]							   =jnp.transpose(snake_case  ,							(0, 2, 3, 1))
          _UpperCAmelCase       :       Optional[int]							   =self.controlnet_cond_embedding(snake_case)
          sample += controlnet_cond

          # 3. down
          _UpperCAmelCase       :       Tuple							   =(sample,)
          for down_block in self.down_blocks:
             if isinstance(snake_case  ,							snake_case):
                _UpperCAmelCase   ,	_UpperCAmelCase       :       Dict							   =down_block(snake_case  ,							snake_case  ,							snake_case  ,							deterministic=not train)
             else:
                _UpperCAmelCase   ,	_UpperCAmelCase       :       Dict							   =down_block(snake_case  ,							snake_case  ,							deterministic=not train)
             down_block_res_samples += res_samples

          # 4. mid
          _UpperCAmelCase       :       List[Any]							   =self.mid_block(snake_case  ,							snake_case  ,							snake_case  ,							deterministic=not train)

          # 5. contronet blocks
          _UpperCAmelCase       :       Union[str, Any]							   =()
          for down_block_res_sample, controlnet_block in zip(snake_case  ,							self.controlnet_down_blocks):
             _UpperCAmelCase       :       List[str]							   =controlnet_block(snake_case)
             controlnet_down_block_res_samples += (down_block_res_sample,)

          _UpperCAmelCase       :       Optional[int]							   =controlnet_down_block_res_samples

          _UpperCAmelCase       :       List[str]							   =self.controlnet_mid_block(snake_case)

          # 6. scaling
          _UpperCAmelCase       :       Tuple							   =[sample * conditioning_scale for sample in down_block_res_samples]
          mid_block_res_sample *= conditioning_scale

          if not return_dict:
             return (down_block_res_samples, mid_block_res_sample)

          return FlaxControlNetOutput(
              down_block_res_samples=snake_case  ,							mid_block_res_sample=snake_case)

