

"""simple docstring"""




import gc
import random
import unittest

import numpy as np
import torch

from diffusers import (
    DDIMScheduler,
    KandinskyVaaControlnetPipeline,
    KandinskyVaaPriorPipeline,
    UNetaDConditionModel,
    VQModel,
)
from diffusers.utils import floats_tensor, load_image, load_numpy, slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu

from ..test_pipelines_common import PipelineTesterMixin, assert_mean_pixel_difference


enable_full_determinism()


class 			A						(							_UpperCAmelCase   ,							unittest.TestCase			):







    """simple docstring"""





    lowerCamelCase		     =							KandinskyVaaControlnetPipeline
    lowerCamelCase		     =							['image_embeds', 'negative_image_embeds', 'hint']
    lowerCamelCase		     =							['image_embeds', 'negative_image_embeds', 'hint']
    lowerCamelCase		     =							[
        'generator',
        'height',
        'width',
        'latents',
        'guidance_scale',
        'num_inference_steps',
        'return_dict',
        'guidance_scale',
        'num_images_per_prompt',
        'output_type',
        'return_dict',
    ]
    lowerCamelCase		     =							False


    @property
    def 						snake_case__					(				self			:	Tuple		)->   Optional[int]:





        '''simple docstring'''







        return 3_2


    @property
    def 						snake_case__					(				self			:	Optional[Any]		)->   Tuple:





        '''simple docstring'''







        return 3_2


    @property
    def 						snake_case__					(				self			:	Tuple		)->   Union[str, Any]:





        '''simple docstring'''







        return self.time_input_dim


    @property
    def 						snake_case__					(				self			:	List[str]		)->   int:





        '''simple docstring'''







        return self.time_input_dim * 4


    @property
    def 						snake_case__					(				self			:	Union[str, Any]		)->   Optional[Any]:





        '''simple docstring'''







        return 1_0_0


    @property
    def 						snake_case__					(				self			:	Optional[Any]		)->   Union[str, Any]:





        '''simple docstring'''







        torch.manual_seed(0		)

        A__   				=       {
            'in_channels': 8,
            # Out channels is double in channels because predicts mean and variance
            'out_channels': 8,
            'addition_embed_type': 'image_hint',
            'down_block_types': ('ResnetDownsampleBlock2D', 'SimpleCrossAttnDownBlock2D'),
            'up_block_types': ('SimpleCrossAttnUpBlock2D', 'ResnetUpsampleBlock2D'),
            'mid_block_type': 'UNetMidBlock2DSimpleCrossAttn',
            'block_out_channels': (self.block_out_channels_a, self.block_out_channels_a * 2),
            'layers_per_block': 1,
            'encoder_hid_dim': self.text_embedder_hidden_size,
            'encoder_hid_dim_type': 'image_proj',
            'cross_attention_dim': self.cross_attention_dim,
            'attention_head_dim': 4,
            'resnet_time_scale_shift': 'scale_shift',
            'class_embed_type': None,
        }

        A__   				=       UNetaDConditionModel(**lowercase_		)
        return model


    @property
    def 						snake_case__					(				self			:	Union[str, Any]		)->   int:





        '''simple docstring'''







        return {
            "block_out_channels": [3_2, 3_2, 6_4, 6_4],
            "down_block_types": [
                "DownEncoderBlock2D",
                "DownEncoderBlock2D",
                "DownEncoderBlock2D",
                "AttnDownEncoderBlock2D",
            ],
            "in_channels": 3,
            "latent_channels": 4,
            "layers_per_block": 1,
            "norm_num_groups": 8,
            "norm_type": "spatial",
            "num_vq_embeddings": 1_2,
            "out_channels": 3,
            "up_block_types": ["AttnUpDecoderBlock2D", "UpDecoderBlock2D", "UpDecoderBlock2D", "UpDecoderBlock2D"],
            "vq_embed_dim": 4,
        }


    @property
    def 						snake_case__					(				self			:	str		)->   str:





        '''simple docstring'''







        torch.manual_seed(0		)
        A__   				=       VQModel(**self.dummy_movq_kwargs		)
        return model


    def 						snake_case__					(				self			:	Optional[Any]		)->   Tuple:





        '''simple docstring'''







        A__   				=       self.dummy_unet
        A__   				=       self.dummy_movq

        A__   				=       DDIMScheduler(
            num_train_timesteps=1_0_0_0,beta_schedule='linear',beta_start=0.00_085,beta_end=0.012,clip_sample=lowercase_,set_alpha_to_one=lowercase_,steps_offset=1,prediction_type='epsilon',thresholding=lowercase_,)

        A__   				=       {
            'unet': unet,
            'scheduler': scheduler,
            'movq': movq,
        }
        return components


    def 						snake_case__					(				self			:	Tuple,lowercase_			:	Any,lowercase_			:	List[Any]=0		)->   Optional[Any]:





        '''simple docstring'''







        A__   				=       floats_tensor((1, self.text_embedder_hidden_size),rng=random.Random(lowercase_		)		).to(lowercase_		)
        A__   				=       floats_tensor((1, self.text_embedder_hidden_size),rng=random.Random(seed + 1		)		).to(
            lowercase_		)

        # create hint
        A__   				=       floats_tensor((1, 3, 6_4, 6_4),rng=random.Random(lowercase_		)		).to(lowercase_		)

        if str(lowercase_		).startswith('mps'		):
            A__   				=       torch.manual_seed(lowercase_		)
        else:
            A__   				=       torch.Generator(device=lowercase_		).manual_seed(lowercase_		)
        A__   				=       {
            'image_embeds': image_embeds,
            'negative_image_embeds': negative_image_embeds,
            'hint': hint,
            'generator': generator,
            'height': 6_4,
            'width': 6_4,
            'guidance_scale': 4.0,
            'num_inference_steps': 2,
            'output_type': 'np',
        }
        return inputs


    def 						snake_case__					(				self			:	Optional[Any]		)->   Optional[Any]:





        '''simple docstring'''







        A__   				=       'cpu'

        A__   				=       self.get_dummy_components()

        A__   				=       self.pipeline_class(**lowercase_		)
        A__   				=       pipe.to(lowercase_		)

        pipe.set_progress_bar_config(disable=lowercase_		)

        A__   				=       pipe(**self.get_dummy_inputs(lowercase_		)		)
        A__   				=       output.images

        A__   				=       pipe(
            **self.get_dummy_inputs(lowercase_		),return_dict=lowercase_,)[0]

        A__   				=       image[0, -3:, -3:, -1]
        A__   				=       image_from_tuple[0, -3:, -3:, -1]

        assert image.shape == (1, 6_4, 6_4, 3)

        A__   				=       np.array(
            [0.6_959_826, 0.868_279, 0.7_558_092, 0.68_769_467, 0.85_805_804, 0.65_977_496, 0.44_885_302, 0.5_959_111, 0.4_251_595]		)

        assert (
            np.abs(image_slice.flatten() - expected_slice		).max() < 1E-2
        ), F' expected_slice {expected_slice}, but got {image_slice.flatten()}'

        assert (
            np.abs(image_from_tuple_slice.flatten() - expected_slice		).max() < 1E-2
        ), F' expected_slice {expected_slice}, but got {image_from_tuple_slice.flatten()}'




@slow
@require_torch_gpu
class 			A						(							unittest.TestCase			):







    """simple docstring"""





    def 						snake_case__					(				self			:	int		)->   List[Any]:





        '''simple docstring'''







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


    def 						snake_case__					(				self			:	Any		)->   int:





        '''simple docstring'''







        A__   				=       load_numpy(
            'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'
            '/kandinskyv22/kandinskyv22_controlnet_robotcat_fp16.npy'		)

        A__   				=       load_image(
            'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'
            '/kandinskyv22/hint_image_cat.png'		)
        A__   				=       torch.from_numpy(np.array(lowercase_		)		).float() / 255.0
        A__   				=       hint.permute(2,0,1		).unsqueeze(0		)

        A__   				=       KandinskyVaaPriorPipeline.from_pretrained(
            'kandinsky-community/kandinsky-2-2-prior',torch_dtype=torch.floataa		)
        pipe_prior.to(lowercase_		)

        A__   				=       KandinskyVaaControlnetPipeline.from_pretrained(
            'kandinsky-community/kandinsky-2-2-controlnet-depth',torch_dtype=torch.floataa		)
        A__   				=       pipeline.to(lowercase_		)
        pipeline.set_progress_bar_config(disable=lowercase_		)

        A__   				=       'A robot, 4k photo'

        A__   				=       torch.Generator(device='cuda'		).manual_seed(0		)
        A__				,						A__   				=       pipe_prior(
            lowercase_,generator=lowercase_,num_inference_steps=5,negative_prompt='',).to_tuple()

        A__   				=       torch.Generator(device='cuda'		).manual_seed(0		)
        A__   				=       pipeline(
            image_embeds=lowercase_,negative_image_embeds=lowercase_,hint=lowercase_,generator=lowercase_,num_inference_steps=1_0_0,output_type='np',)

        A__   				=       output.images[0]

        assert image.shape == (5_1_2, 5_1_2, 3)

        assert_mean_pixel_difference(lowercase_,lowercase_		)




import argparse
import collections

import torch
from flax import traverse_util
from tax import checkpoints

from transformers import TaConfig, TaEncoderModel, TaForConditionalGeneration
from transformers.utils import logging


logging.set_verbosity_info()
def 						_snake_case(   SCREAMING_SNAKE_CASE__      :     List[Any]   ,      SCREAMING_SNAKE_CASE__      :     List[str]   ,      SCREAMING_SNAKE_CASE__      :     Optional[int]   ,      SCREAMING_SNAKE_CASE__      :     Optional[int]="attention"				)							->				Union[str, Any]:





    '''simple docstring'''






    A__   				=       params[f'{prefix}/layers_{i}/{layer_name}/key/kernel']
    A__   				=       params[f'{prefix}/layers_{i}/{layer_name}/out/kernel']
    A__   				=       params[f'{prefix}/layers_{i}/{layer_name}/query/kernel']
    A__   				=       params[f'{prefix}/layers_{i}/{layer_name}/value/kernel']
    return k, o, q, v
def 						_snake_case(   SCREAMING_SNAKE_CASE__      :     Any   ,      SCREAMING_SNAKE_CASE__      :     Tuple   ,      SCREAMING_SNAKE_CASE__      :     Optional[Any]   ,      SCREAMING_SNAKE_CASE__      :     Dict=False				)							->				str:





    '''simple docstring'''






    if split_mlp_wi:
        A__   				=       params[f'{prefix}/layers_{i}/mlp/wi_0/kernel']
        A__   				=       params[f'{prefix}/layers_{i}/mlp/wi_1/kernel']
        A__   				=       (wi_a, wi_a)
    else:
        A__   				=       params[f'{prefix}/layers_{i}/mlp/wi/kernel']

    A__   				=       params[f'{prefix}/layers_{i}/mlp/wo/kernel']
    return wi, wo
def 						_snake_case(   SCREAMING_SNAKE_CASE__      :     Optional[Any]   ,      SCREAMING_SNAKE_CASE__      :     Tuple   ,      SCREAMING_SNAKE_CASE__      :     Dict   ,      SCREAMING_SNAKE_CASE__      :     int				)							->				str:





    '''simple docstring'''






    return params[f'{prefix}/layers_{i}/{layer_name}/scale']
def 						_snake_case(   SCREAMING_SNAKE_CASE__      :     dict   ,      *, SCREAMING_SNAKE_CASE__      :     int   ,      SCREAMING_SNAKE_CASE__      :     bool				)							->				int:





    '''simple docstring'''






    A__   				=       traverse_util.flatten_dict(variables['target']				)
    A__   				=       {'/'.join(SCREAMING_SNAKE_CASE__				): v for k, v in old.items()}

    # v1.1 models have a gated GeLU with wi_0 and wi_1 instead of wi
    A__   				=       'encoder/layers_0/mlp/wi_0/kernel' in old
    print('Split MLP:'   ,      SCREAMING_SNAKE_CASE__				)

    A__   				=       collections.OrderedDict()

    # Shared embeddings.
    A__   				=       old['token_embedder/embedding']

    # Encoder.
    for i in range(SCREAMING_SNAKE_CASE__				):
        # Block i, layer 0 (Self Attention).
        A__   				=       tax_layer_norm_lookup(SCREAMING_SNAKE_CASE__   ,      SCREAMING_SNAKE_CASE__   ,      'encoder'   ,      'pre_attention_layer_norm'				)
        A__				,						A__				,						A__				,						A__   				=       tax_attention_lookup(SCREAMING_SNAKE_CASE__   ,      SCREAMING_SNAKE_CASE__   ,      'encoder'   ,      'attention'				)
        A__   				=       layer_norm
        A__   				=       k.T
        A__   				=       o.T
        A__   				=       q.T
        A__   				=       v.T

        # Block i, layer 1 (MLP).
        A__   				=       tax_layer_norm_lookup(SCREAMING_SNAKE_CASE__   ,      SCREAMING_SNAKE_CASE__   ,      'encoder'   ,      'pre_mlp_layer_norm'				)
        A__				,						A__   				=       tax_mlp_lookup(SCREAMING_SNAKE_CASE__   ,      SCREAMING_SNAKE_CASE__   ,      'encoder'   ,      SCREAMING_SNAKE_CASE__				)
        A__   				=       layer_norm
        if split_mlp_wi:
            A__   				=       wi[0].T
            A__   				=       wi[1].T
        else:
            A__   				=       wi.T
        A__   				=       wo.T

    A__   				=       old[
        'encoder/relpos_bias/rel_embedding'
    ].T
    A__   				=       old['encoder/encoder_norm/scale']

    if not is_encoder_only:
        # Decoder.
        for i in range(SCREAMING_SNAKE_CASE__				):
            # Block i, layer 0 (Self Attention).
            A__   				=       tax_layer_norm_lookup(SCREAMING_SNAKE_CASE__   ,      SCREAMING_SNAKE_CASE__   ,      'decoder'   ,      'pre_self_attention_layer_norm'				)
            A__				,						A__				,						A__				,						A__   				=       tax_attention_lookup(SCREAMING_SNAKE_CASE__   ,      SCREAMING_SNAKE_CASE__   ,      'decoder'   ,      'self_attention'				)
            A__   				=       layer_norm
            A__   				=       k.T
            A__   				=       o.T
            A__   				=       q.T
            A__   				=       v.T

            # Block i, layer 1 (Cross Attention).
            A__   				=       tax_layer_norm_lookup(SCREAMING_SNAKE_CASE__   ,      SCREAMING_SNAKE_CASE__   ,      'decoder'   ,      'pre_cross_attention_layer_norm'				)
            A__				,						A__				,						A__				,						A__   				=       tax_attention_lookup(SCREAMING_SNAKE_CASE__   ,      SCREAMING_SNAKE_CASE__   ,      'decoder'   ,      'encoder_decoder_attention'				)
            A__   				=       layer_norm
            A__   				=       k.T
            A__   				=       o.T
            A__   				=       q.T
            A__   				=       v.T

            # Block i, layer 2 (MLP).
            A__   				=       tax_layer_norm_lookup(SCREAMING_SNAKE_CASE__   ,      SCREAMING_SNAKE_CASE__   ,      'decoder'   ,      'pre_mlp_layer_norm'				)
            A__				,						A__   				=       tax_mlp_lookup(SCREAMING_SNAKE_CASE__   ,      SCREAMING_SNAKE_CASE__   ,      'decoder'   ,      SCREAMING_SNAKE_CASE__				)
            A__   				=       layer_norm
            if split_mlp_wi:
                A__   				=       wi[0].T
                A__   				=       wi[1].T
            else:
                A__   				=       wi.T
            A__   				=       wo.T

        A__   				=       old['decoder/decoder_norm/scale']
        A__   				=       old[
            'decoder/relpos_bias/rel_embedding'
        ].T

        # LM Head (only in v1.1 checkpoints, in v1.0 embeddings are used instead)
        if "decoder/logits_dense/kernel" in old:
            A__   				=       old['decoder/logits_dense/kernel'].T

    return new
def 						_snake_case(   SCREAMING_SNAKE_CASE__      :     Optional[Any]   ,      SCREAMING_SNAKE_CASE__      :     bool				)							->				Dict:





    '''simple docstring'''






    A__   				=       collections.OrderedDict([(k, torch.from_numpy(v.copy()				)) for (k, v) in converted_params.items()]				)

    # Add what is missing.
    if "encoder.embed_tokens.weight" not in state_dict:
        A__   				=       state_dict['shared.weight']

    if not is_encoder_only:
        if "decoder.embed_tokens.weight" not in state_dict:
            A__   				=       state_dict['shared.weight']

        if "lm_head.weight" not in state_dict:  # For old 1.0 models.
            print('Using shared word embeddings as lm_head.'				)
            A__   				=       state_dict['shared.weight']

    return state_dict
def 						_snake_case(   SCREAMING_SNAKE_CASE__      :     Union[str, Any]   ,      SCREAMING_SNAKE_CASE__      :     Any   ,      SCREAMING_SNAKE_CASE__      :     Any   ,      SCREAMING_SNAKE_CASE__      :     Tuple				)							->				int:





    '''simple docstring'''






    A__   				=       checkpoints.load_tax_checkpoint(SCREAMING_SNAKE_CASE__				)
    A__   				=       convert_tax_to_pytorch(SCREAMING_SNAKE_CASE__   ,      num_layers=config.num_layers   ,      is_encoder_only=SCREAMING_SNAKE_CASE__				)
    A__   				=       make_state_dict(SCREAMING_SNAKE_CASE__   ,      SCREAMING_SNAKE_CASE__				)
    model.load_state_dict(SCREAMING_SNAKE_CASE__   ,      strict=SCREAMING_SNAKE_CASE__				)
def 						_snake_case(   SCREAMING_SNAKE_CASE__      :     Optional[int]   ,      SCREAMING_SNAKE_CASE__      :     List[Any]   ,      SCREAMING_SNAKE_CASE__      :     Tuple   ,      SCREAMING_SNAKE_CASE__      :     bool = False				)							->				Any:





    '''simple docstring'''






    A__   				=       TaConfig.from_json_file(SCREAMING_SNAKE_CASE__				)
    print(f'Building PyTorch model from configuration: {config}'				)
    # Non-v1.1 checkpoints could also use T5Model, but this works for all.
    # The v1.0 checkpoints will simply have an LM head that is the word embeddings.
    if is_encoder_only:
        A__   				=       TaEncoderModel(SCREAMING_SNAKE_CASE__				)
    else:
        A__   				=       TaForConditionalGeneration(SCREAMING_SNAKE_CASE__				)

    # Load weights from tf checkpoint
    load_tax_weights_in_ta(SCREAMING_SNAKE_CASE__   ,      SCREAMING_SNAKE_CASE__   ,      SCREAMING_SNAKE_CASE__   ,      SCREAMING_SNAKE_CASE__				)

    # Save pytorch-model
    print(f'Save PyTorch model to {pytorch_dump_path}'				)
    model.save_pretrained(SCREAMING_SNAKE_CASE__				)

    # Verify that we can load the checkpoint.
    model.from_pretrained(SCREAMING_SNAKE_CASE__				)
    print('Done'				)


if __name__ == "__main__":
     lowercase_			=  argparse.ArgumentParser(description="Converts a native T5X checkpoint into a PyTorch checkpoint.")
     # Required parameters
     parser.add_argument(
         "--t5x_checkpoint_path", default=None, type=str, required=True, help="Path to the T5X checkpoint."
     )
     parser.add_argument(
         "--config_file",
         default=None,
         type=str,
         required=True,
         help="The config json file corresponding to the pre-trained T5 model.\nThis specifies the model architecture.",
     )
     parser.add_argument(
         "--pytorch_dump_path", default=None, type=str, required=True, help="Path to the output PyTorch model."
     )
     parser.add_argument(
         "--is_encoder_only", action="store_true", help="Check if the model is encoder-decoder model", default=False
     )
     lowercase_			=  parser.parse_args()
     convert_tax_checkpoint_to_pytorch(
         args.tax_checkpoint_path, args.config_file, args.pytorch_dump_path, args.is_encoder_only
     )

