





from pickle import UnpicklingError

import jax
import jax.numpy as jnp
import numpy as np
from flax.serialization import from_bytes
from flax.traverse_util import flatten_dict

from ..utils import logging


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



def 	A_       (  _lowerCAmelCase							,     _lowerCAmelCase  )		->				Tuple:





 try:
  with open(UpperCamelCase__							,     "rb"  ) as flax_state_f:
   UpperCamelCase			:      Union[str, Any] 			=     from_bytes(UpperCamelCase__							,     flax_state_f.read()  )
 except UnpicklingError as e:
  try:
   with open(UpperCamelCase__  ) as f:
    if f.read().startswith("version"  ):
     raise OSError(
         "You seem to have cloned a repository without having git-lfs installed. Please"
         " install git-lfs and run `git lfs install` followed by `git lfs pull` in the"
         " folder you cloned."  )
    else:
     raise ValueError from e
  except (UnicodeDecodeError, ValueError):
   raise EnvironmentError(F"""Unable to convert {model_file} to Flax deserializable object. """  )

 return load_flax_weights_in_pytorch_model(UpperCamelCase__							,     UpperCamelCase__  )



def 	A_       (  _lowerCAmelCase							,     _lowerCAmelCase  )		->				Union[str, Any]:





 try:
  import torch  # noqa: F401
 except ImportError:
  logger.error(
      "Loading Flax weights in PyTorch requires both PyTorch and Flax to be installed. Please see"
      " https://pytorch.org/ and https://flax.readthedocs.io/en/latest/installation.html for installation"
      " instructions."  )
  raise

 # check if we have bf16 weights
 UpperCamelCase			:      Optional[int] 			=     flatten_dict(jax.tree_util.tree_map(lambda _lowerCAmelCase  : x.dtype == jnp.bfloataa							,     UpperCamelCase__  )  ).values()
 if any(UpperCamelCase__  ):
  # convert all weights to fp32 if they are bf16 since torch.from_numpy can-not handle bf16

  # and bf16 is not fully supported in PT yet.
  logger.warning(
      "Found ``bfloat16`` weights in Flax model. Casting all ``bfloat16`` weights to ``float32`` "
      "before loading those in PyTorch model."  )
  UpperCamelCase			:      List[Any] 			=     jax.tree_util.tree_map(
      lambda _lowerCAmelCase  : params.astype(np.floataa  ) if params.dtype == jnp.bfloataa else params							,     UpperCamelCase__  )

 UpperCamelCase			:      List[str] 			=     ""

 UpperCamelCase			:      Union[str, Any] 			=     flatten_dict(UpperCamelCase__							,     sep="."  )
 UpperCamelCase			:      Union[str, Any] 			=     pt_model.state_dict()

 # keep track of unexpected & missing keys
 UpperCamelCase			:      List[Any] 			=     []
 UpperCamelCase			:      Optional[int] 			=     set(pt_model_dict.keys()  )

 for flax_key_tuple, flax_tensor in flax_state_dict.items():
  UpperCamelCase			:      Any 			=     flax_key_tuple.split("."  )

  if flax_key_tuple_array[-1] == "kernel" and flax_tensor.ndim == 4:
   UpperCamelCase			:      Optional[Any] 			=     flax_key_tuple_array[:-1] + ["weight"]
   UpperCamelCase			:      int 			=     jnp.transpose(UpperCamelCase__							,     (3, 2, 0, 1)  )
  elif flax_key_tuple_array[-1] == "kernel":
   UpperCamelCase			:      Tuple 			=     flax_key_tuple_array[:-1] + ["weight"]
   UpperCamelCase			:      List[Any] 			=     flax_tensor.T
  elif flax_key_tuple_array[-1] == "scale":
   UpperCamelCase			:      List[str] 			=     flax_key_tuple_array[:-1] + ["weight"]

  if "time_embedding" not in flax_key_tuple_array:
   for i, flax_key_tuple_string in enumerate(UpperCamelCase__  ):
    UpperCamelCase			:      List[Any] 			=     (
        flax_key_tuple_string.replace("_0"							,     ".0"  )
        .replace("_1"							,     ".1"  )
        .replace("_2"							,     ".2"  )
        .replace("_3"							,     ".3"  )
        .replace("_4"							,     ".4"  )
        .replace("_5"							,     ".5"  )
        .replace("_6"							,     ".6"  )
        .replace("_7"							,     ".7"  )
        .replace("_8"							,     ".8"  )
        .replace("_9"							,     ".9"  )
    )

  UpperCamelCase			:      Any 			=     ".".join(UpperCamelCase__  )

  if flax_key in pt_model_dict:
   if flax_tensor.shape != pt_model_dict[flax_key].shape:
    raise ValueError(
        F"""Flax checkpoint seems to be incorrect. Weight {flax_key_tuple} was expected """
        F"""to be of shape {pt_model_dict[flax_key].shape}, but is {flax_tensor.shape}."""  )
   else:
    # add weight to pytorch dict
    UpperCamelCase			:      Union[str, Any] 			=     np.asarray(UpperCamelCase__  ) if not isinstance(UpperCamelCase__							,     np.ndarray  ) else flax_tensor
    UpperCamelCase			:      str 			=     torch.from_numpy(UpperCamelCase__  )
    # remove from missing keys
    missing_keys.remove(UpperCamelCase__  )
  else:
   # weight is not expected by PyTorch model
   unexpected_keys.append(UpperCamelCase__  )

 pt_model.load_state_dict(UpperCamelCase__  )

 # re-transform missing_keys to list
 UpperCamelCase			:      List[Any] 			=     list(UpperCamelCase__  )

 if len(UpperCamelCase__  ) > 0:
  logger.warning(
      "Some weights of the Flax model were not used when initializing the PyTorch model"
      F""" {pt_model.__class__.__name__}: {unexpected_keys}\n- This IS expected if you are initializing"""
      F""" {pt_model.__class__.__name__} from a Flax model trained on another task or with another architecture"""
      " (e.g. initializing a BertForSequenceClassification model from a FlaxBertForPreTraining model).\n- This"
      F""" IS NOT expected if you are initializing {pt_model.__class__.__name__} from a Flax model that you expect"""
      " to be exactly identical (e.g. initializing a BertForSequenceClassification model from a"
      " FlaxBertForSequenceClassification model)."  )
 if len(UpperCamelCase__  ) > 0:
  logger.warning(
      F"""Some weights of {pt_model.__class__.__name__} were not initialized from the Flax model and are newly"""
      F""" initialized: {missing_keys}\nYou should probably TRAIN this model on a down-stream task to be able to"""
      " use it for predictions and inference."  )

 return pt_model







import unittest

import numpy as np

from transformers.testing_utils import require_torch, require_vision
from transformers.utils import is_torch_available, is_vision_available

from ...test_image_processing_common import ImageProcessingSavingTestMixin, prepare_video_inputs


if is_torch_available():
    import torch

if is_vision_available():
    from PIL import Image

    from transformers import VivitImageProcessor



class       A__     ( unittest.TestCase    ):





      def __init__(				self ,							A_ ,							A_=7 ,							A_=3 ,							A_=10 ,							A_=18 ,							A_=30 ,							A_=400 ,							A_=True ,							A_=None ,							A_=True ,							A_=[0.5, 0.5, 0.5] ,							A_=[0.5, 0.5, 0.5] ,							A_=None ,							):
       '''simple docstring'''




       UpperCamelCase			:      Optional[int] 			=     size if size is not None else {"shortest_edge": 18}
       UpperCamelCase			:      Tuple 			=     crop_size if crop_size is not None else {"height": 18, "width": 18}

       UpperCamelCase			:      Optional[Any] 			=     parent
       UpperCamelCase			:      Optional[int] 			=     batch_size
       UpperCamelCase			:      List[Any] 			=     num_channels
       UpperCamelCase			:      Union[str, Any] 			=     num_frames
       UpperCamelCase			:      Any 			=     image_size
       UpperCamelCase			:      Tuple 			=     min_resolution
       UpperCamelCase			:      Optional[Any] 			=     max_resolution
       UpperCamelCase			:      Any 			=     do_resize
       UpperCamelCase			:      Tuple 			=     size
       UpperCamelCase			:      List[Any] 			=     do_normalize
       UpperCamelCase			:      Optional[int] 			=     image_mean
       UpperCamelCase			:      Any 			=     image_std
       UpperCamelCase			:      str 			=     crop_size





      def  __UpperCamelCase(				self     ):
       '''simple docstring'''




       return {
           "image_mean": self.image_mean,
           "image_std": self.image_std,
           "do_normalize": self.do_normalize,
           "do_resize": self.do_resize,
           "size": self.size,
           "crop_size": self.crop_size,
       }



@require_torch
@require_vision
class       A__     ( __snake_case ,     unittest.TestCase    ):
      _UpperCAmelCase			:List[str]							=     VivitImageProcessor if is_vision_available() else None





      def  __UpperCamelCase(				self     ):
       '''simple docstring'''




       UpperCamelCase			:      List[Any] 			=     VivitImageProcessingTester(self     )





      @property
      def  __UpperCamelCase(				self     ):
       '''simple docstring'''




       return self.image_processor_tester.prepare_image_processor_dict()





      def  __UpperCamelCase(				self     ):
       '''simple docstring'''




       UpperCamelCase			:      List[Any] 			=     self.image_processing_class(**self.image_processor_dict     )
       self.assertTrue(hasattr(A_ ,							"image_mean"     )     )
       self.assertTrue(hasattr(A_ ,							"image_std"     )     )
       self.assertTrue(hasattr(A_ ,							"do_normalize"     )     )
       self.assertTrue(hasattr(A_ ,							"do_resize"     )     )
       self.assertTrue(hasattr(A_ ,							"do_center_crop"     )     )
       self.assertTrue(hasattr(A_ ,							"size"     )     )





      def  __UpperCamelCase(				self     ):
       '''simple docstring'''




       UpperCamelCase			:      Dict 			=     self.image_processing_class.from_dict(self.image_processor_dict     )
       self.assertEqual(image_processor.size ,							{"shortest_edge": 18}     )
       self.assertEqual(image_processor.crop_size ,							{"height": 18, "width": 18}     )

       UpperCamelCase			:      Dict 			=     self.image_processing_class.from_dict(self.image_processor_dict ,							size=42 ,							crop_size=84     )
       self.assertEqual(image_processor.size ,							{"shortest_edge": 42}     )
       self.assertEqual(image_processor.crop_size ,							{"height": 84, "width": 84}     )





      def  __UpperCamelCase(				self     ):
       '''simple docstring'''




       UpperCamelCase			:      Optional[Any] 			=     self.image_processing_class(**self.image_processor_dict     )
       # create random PIL videos
       UpperCamelCase			:      Union[str, Any] 			=     prepare_video_inputs(self.image_processor_tester ,							equal_resolution=A_     )
       for video in video_inputs:
        self.assertIsInstance(A_ ,							A_     )
        self.assertIsInstance(video[0] ,							Image.Image     )

       # Test not batched input
       UpperCamelCase			:      Any 			=     image_processing(video_inputs[0] ,							return_tensors="pt"     ).pixel_values
       self.assertEqual(
           encoded_videos.shape ,							(
               1,
               self.image_processor_tester.num_frames,
               self.image_processor_tester.num_channels,
               self.image_processor_tester.crop_size["height"],
               self.image_processor_tester.crop_size["width"],
           ) ,							)

       # Test batched
       UpperCamelCase			:      str 			=     image_processing(A_ ,							return_tensors="pt"     ).pixel_values
       self.assertEqual(
           encoded_videos.shape ,							(
               self.image_processor_tester.batch_size,
               self.image_processor_tester.num_frames,
               self.image_processor_tester.num_channels,
               self.image_processor_tester.crop_size["height"],
               self.image_processor_tester.crop_size["width"],
           ) ,							)





      def  __UpperCamelCase(				self     ):
       '''simple docstring'''




       UpperCamelCase			:      List[Any] 			=     self.image_processing_class(**self.image_processor_dict     )
       # create random numpy tensors
       UpperCamelCase			:      str 			=     prepare_video_inputs(self.image_processor_tester ,							equal_resolution=A_ ,							numpify=A_     )
       for video in video_inputs:
        self.assertIsInstance(A_ ,							A_     )
        self.assertIsInstance(video[0] ,							np.ndarray     )

       # Test not batched input
       UpperCamelCase			:      Tuple 			=     image_processing(video_inputs[0] ,							return_tensors="pt"     ).pixel_values
       self.assertEqual(
           encoded_videos.shape ,							(
               1,
               self.image_processor_tester.num_frames,
               self.image_processor_tester.num_channels,
               self.image_processor_tester.crop_size["height"],
               self.image_processor_tester.crop_size["width"],
           ) ,							)

       # Test batched
       UpperCamelCase			:      Any 			=     image_processing(A_ ,							return_tensors="pt"     ).pixel_values
       self.assertEqual(
           encoded_videos.shape ,							(
               self.image_processor_tester.batch_size,
               self.image_processor_tester.num_frames,
               self.image_processor_tester.num_channels,
               self.image_processor_tester.crop_size["height"],
               self.image_processor_tester.crop_size["width"],
           ) ,							)





      def  __UpperCamelCase(				self     ):
       '''simple docstring'''




       UpperCamelCase			:      str 			=     self.image_processing_class(**self.image_processor_dict     )
       # create random PyTorch tensors
       UpperCamelCase			:      Union[str, Any] 			=     prepare_video_inputs(self.image_processor_tester ,							equal_resolution=A_ ,							torchify=A_     )
       for video in video_inputs:
        self.assertIsInstance(A_ ,							A_     )
        self.assertIsInstance(video[0] ,							torch.Tensor     )

       # Test not batched input
       UpperCamelCase			:      Tuple 			=     image_processing(video_inputs[0] ,							return_tensors="pt"     ).pixel_values
       self.assertEqual(
           encoded_videos.shape ,							(
               1,
               self.image_processor_tester.num_frames,
               self.image_processor_tester.num_channels,
               self.image_processor_tester.crop_size["height"],
               self.image_processor_tester.crop_size["width"],
           ) ,							)

       # Test batched
       UpperCamelCase			:      List[Any] 			=     image_processing(A_ ,							return_tensors="pt"     ).pixel_values
       self.assertEqual(
           encoded_videos.shape ,							(
               self.image_processor_tester.batch_size,
               self.image_processor_tester.num_frames,
               self.image_processor_tester.num_channels,
               self.image_processor_tester.crop_size["height"],
               self.image_processor_tester.crop_size["width"],
           ) ,							)
