


'''simple docstring'''






import json
import os
from dataclasses import dataclass
from functools import partial
from typing import Callable

import flax.linen as nn
import jax
import jax.numpy as jnp
import joblib
import optax
import wandb
from flax import jax_utils, struct, traverse_util
from flax.serialization import from_bytes, to_bytes
from flax.training import train_state
from flax.training.common_utils import shard
from tqdm.auto import tqdm

from transformers import BigBirdConfig, FlaxBigBirdForQuestionAnswering
from transformers.models.big_bird.modeling_flax_big_bird import FlaxBigBirdForQuestionAnsweringModule






class    __magic_name__				(						lowerCAmelCase							):
       UpperCAmelCase		 =42
       UpperCAmelCase		 =jnp.floataa
       UpperCAmelCase		 =True






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




          '''simple docstring'''
          super().setup()
          _UpperCAmelCase  :  Any				       =nn.Dense(5    ,				dtype=self.dtype)






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




          '''simple docstring'''
          _UpperCAmelCase  :  Any				       =super().__call__(*snake_case    ,				**snake_case)
          _UpperCAmelCase  :  List[str]				       =self.cls(outputs[2])
          return outputs[:2] + (cls_out,)






class    __magic_name__				(						lowerCAmelCase							):
       UpperCAmelCase		 =FlaxBigBirdForNaturalQuestionsModule






def lowerCamelCase__     (		__lowerCamelCase			:				Dict			,       __lowerCamelCase			:				Tuple			,       __lowerCamelCase			:				Tuple			,       __lowerCamelCase			:				Tuple			,       __lowerCamelCase			:				Tuple			,       __lowerCamelCase			:				str				):







   '''simple docstring'''




   def cross_entropy(__lowerCamelCase			:				List[Any]			,       __lowerCamelCase			:				List[str]			,       __lowerCamelCase			:				Union[str, Any]=None				):
      _UpperCAmelCase  :  int				       =logits.shape[-1]
      _UpperCAmelCase  :  List[Any]				       =(labels[..., None] == jnp.arange(__lowerCamelCase				)[None]).astype('f4'				)
      _UpperCAmelCase  :  Any				       =jax.nn.log_softmax(__lowerCamelCase			,       axis=-1				)
      _UpperCAmelCase  :  Dict				       =-jnp.sum(labels * logits			,       axis=-1				)
      if reduction is not None:
         _UpperCAmelCase  :  List[Any]				       =reduction(__lowerCamelCase				)
      return loss

   _UpperCAmelCase  :  Tuple				       =partial(__lowerCamelCase			,       reduction=jnp.mean				)
   _UpperCAmelCase  :  List[Any]				       =cross_entropy(__lowerCamelCase			,       __lowerCamelCase				)
   _UpperCAmelCase  :  Any				       =cross_entropy(__lowerCamelCase			,       __lowerCamelCase				)
   _UpperCAmelCase  :  int				       =cross_entropy(__lowerCamelCase			,       __lowerCamelCase				)
   return (start_loss + end_loss + pooled_loss) / 3






@dataclass
class    __magic_name__				:
       UpperCAmelCase		 ="google/bigbird-roberta-base"
       UpperCAmelCase		 =3_0_0_0
       UpperCAmelCase		 =1_0_5_0_0

       UpperCAmelCase		 =1_2_8
       UpperCAmelCase		 =3

       UpperCAmelCase		 =1
       UpperCAmelCase		 =5

       # tx_args
       UpperCAmelCase		 =3e-5
       UpperCAmelCase		 =0.0
       UpperCAmelCase		 =2_0_0_0_0
       UpperCAmelCase		 =0.0095

       UpperCAmelCase		 ="bigbird-roberta-natural-questions"
       UpperCAmelCase		 ="training-expt"
       UpperCAmelCase		 ="data/nq-training.jsonl"
       UpperCAmelCase		 ="data/nq-validation.jsonl"






       def      lowerCAmelCase						(     self)		->    Tuple:




          '''simple docstring'''
          os.makedirs(self.base_dir    ,				exist_ok=snake_case)
          _UpperCAmelCase  :  int				       =os.path.join(self.base_dir    ,				self.save_dir)
          _UpperCAmelCase  :  str				       =self.batch_size_per_device * jax.device_count()






@dataclass
class    __magic_name__				:
       UpperCAmelCase		 =42
       UpperCAmelCase		 =4_0_9_6  # no dynamic padding on TPUs






       def __call__(     self    ,				snake_case)		->    str:




          '''simple docstring'''
          _UpperCAmelCase  :  List[Any]				       =self.collate_fn(snake_case)
          _UpperCAmelCase  :  str				       =jax.tree_util.tree_map(snake_case    ,				snake_case)
          return batch






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




          '''simple docstring'''
          _UpperCAmelCase       ,			_UpperCAmelCase  :  Tuple				       =self.fetch_inputs(features['input_ids'])
          _UpperCAmelCase  :  Any				       ={
              'input_ids': jnp.array(snake_case    ,				dtype=jnp.intaa),
              'attention_mask': jnp.array(snake_case    ,				dtype=jnp.intaa),
              'start_labels': jnp.array(features['start_token']    ,				dtype=jnp.intaa),
              'end_labels': jnp.array(features['end_token']    ,				dtype=jnp.intaa),
              'pooled_labels': jnp.array(features['category']    ,				dtype=jnp.intaa),
          }
          return batch






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




          '''simple docstring'''
          _UpperCAmelCase  :  Optional[Any]				       =[self._fetch_inputs(snake_case) for ids in input_ids]
          return zip(*snake_case)






       def      lowerCAmelCase						(     self    ,				snake_case)		->    int:




          '''simple docstring'''
          _UpperCAmelCase  :  int				       =[1 for _ in range(len(snake_case))]
          while len(snake_case) < self.max_length:
             input_ids.append(self.pad_id)
             attention_mask.append(0)
          return input_ids, attention_mask






def lowerCamelCase__     (		__lowerCamelCase			:				Tuple			,       __lowerCamelCase			:				Any			,       __lowerCamelCase			:				List[Any]=None				):







   '''simple docstring'''




   if seed is not None:
      _UpperCAmelCase  :  List[Any]				       =dataset.shuffle(seed=__lowerCamelCase				)
   for i in range(len(__lowerCamelCase				) // batch_size				):
      _UpperCAmelCase  :  int				       =dataset[i * batch_size : (i + 1) * batch_size]
      yield dict(__lowerCamelCase				)





@partial(jax.pmap			,       axis_name='batch'				)
def lowerCamelCase__     (		__lowerCamelCase			:				Any			,       __lowerCamelCase			:				Optional[Any]			,       **__lowerCamelCase			:				List[str]				):







   '''simple docstring'''




   def loss_fn(__lowerCamelCase			:				List[str]				):
      _UpperCAmelCase  :  Dict				       =model_inputs.pop('start_labels'				)
      _UpperCAmelCase  :  str				       =model_inputs.pop('end_labels'				)
      _UpperCAmelCase  :  Optional[Any]				       =model_inputs.pop('pooled_labels'				)

      _UpperCAmelCase  :  Tuple				       =state.apply_fn(**__lowerCamelCase			,       params=__lowerCamelCase			,       dropout_rng=__lowerCamelCase			,       train=__lowerCamelCase				)
      _UpperCAmelCase       ,			_UpperCAmelCase       ,			_UpperCAmelCase  :  Optional[int]				       =outputs

      return state.loss_fn(
          __lowerCamelCase			,       __lowerCamelCase			,       __lowerCamelCase			,       __lowerCamelCase			,       __lowerCamelCase			,       __lowerCamelCase			,       )

   _UpperCAmelCase       ,			_UpperCAmelCase  :  Tuple				       =jax.random.split(__lowerCamelCase				)
   _UpperCAmelCase  :  Dict				       =jax.value_and_grad(__lowerCamelCase				)
   _UpperCAmelCase       ,			_UpperCAmelCase  :  str				       =grad_fn(state.params				)
   _UpperCAmelCase  :  Any				       =jax.lax.pmean({'loss': loss}			,       axis_name='batch'				)
   _UpperCAmelCase  :  List[Any]				       =jax.lax.pmean(__lowerCamelCase			,       'batch'				)

   _UpperCAmelCase  :  int				       =state.apply_gradients(grads=__lowerCamelCase				)
   return state, metrics, new_drp_rng





@partial(jax.pmap			,       axis_name='batch'				)
def lowerCamelCase__     (		__lowerCamelCase			:				Any			,       **__lowerCamelCase			:				List[Any]				):







   '''simple docstring'''




   _UpperCAmelCase  :  List[str]				       =model_inputs.pop('start_labels'				)
   _UpperCAmelCase  :  Any				       =model_inputs.pop('end_labels'				)
   _UpperCAmelCase  :  Dict				       =model_inputs.pop('pooled_labels'				)

   _UpperCAmelCase  :  Dict				       =state.apply_fn(**__lowerCamelCase			,       params=state.params			,       train=__lowerCamelCase				)
   _UpperCAmelCase       ,			_UpperCAmelCase       ,			_UpperCAmelCase  :  Optional[int]				       =outputs

   _UpperCAmelCase  :  Tuple				       =state.loss_fn(__lowerCamelCase			,       __lowerCamelCase			,       __lowerCamelCase			,       __lowerCamelCase			,       __lowerCamelCase			,       __lowerCamelCase				)
   _UpperCAmelCase  :  List[str]				       =jax.lax.pmean({'loss': loss}			,       axis_name='batch'				)
   return metrics






class    __magic_name__				(						train_state.TrainState							):
       UpperCAmelCase		 =struct.field(pytree_node=lowerCAmelCase							)






@dataclass
class    __magic_name__				:
       UpperCAmelCase		 =42
       UpperCAmelCase		 =42
       UpperCAmelCase		 =42
       UpperCAmelCase		 =42
       UpperCAmelCase		 =42
       UpperCAmelCase		 =42
       UpperCAmelCase		 =None






       def      lowerCAmelCase						(     self    ,				snake_case    ,				snake_case    ,				snake_case    ,				snake_case=None)		->    Tuple:




          '''simple docstring'''
          _UpperCAmelCase  :  Dict				       =model.params
          _UpperCAmelCase  :  List[str]				       =TrainState.create(
              apply_fn=model.__call__    ,				params=snake_case    ,				tx=snake_case    ,				loss_fn=snake_case    ,				)
          if ckpt_dir is not None:
             _UpperCAmelCase       ,			_UpperCAmelCase       ,			_UpperCAmelCase       ,			_UpperCAmelCase       ,			_UpperCAmelCase  :  List[str]				       =restore_checkpoint(snake_case    ,				snake_case)
             _UpperCAmelCase  :  int				       ={
                 'lr': args.lr,
                 'init_lr': args.init_lr,
                 'warmup_steps': args.warmup_steps,
                 'num_train_steps': num_train_steps,
                 'weight_decay': args.weight_decay,
             }
             _UpperCAmelCase       ,			_UpperCAmelCase  :  Any				       =build_tx(**snake_case)
             _UpperCAmelCase  :  Dict				       =train_state.TrainState(
                 step=snake_case    ,				apply_fn=model.__call__    ,				params=snake_case    ,				tx=snake_case    ,				opt_state=snake_case    ,				)
             _UpperCAmelCase  :  List[str]				       =args
             _UpperCAmelCase  :  int				       =data_collator
             _UpperCAmelCase  :  List[str]				       =lr
             _UpperCAmelCase  :  str				       =params
          _UpperCAmelCase  :  Tuple				       =jax_utils.replicate(snake_case)
          return state






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




          '''simple docstring'''
          _UpperCAmelCase  :  int				       =self.args
          _UpperCAmelCase  :  Optional[int]				       =len(snake_case) // args.batch_size

          _UpperCAmelCase  :  List[str]				       =jax.random.PRNGKey(0)
          _UpperCAmelCase  :  Dict				       =jax.random.split(snake_case    ,				jax.device_count())
          for epoch in range(args.max_epochs):
             _UpperCAmelCase  :  List[str]				       =jnp.array(0    ,				dtype=jnp.floataa)
             _UpperCAmelCase  :  Dict				       =get_batched_dataset(snake_case    ,				args.batch_size    ,				seed=snake_case)
             _UpperCAmelCase  :  Optional[int]				       =0
             for batch in tqdm(snake_case    ,				total=snake_case    ,				desc=f"Running EPOCH-{epoch}"):
                _UpperCAmelCase  :  List[Any]				       =self.data_collator(snake_case)
                _UpperCAmelCase       ,			_UpperCAmelCase       ,			_UpperCAmelCase  :  List[Any]				       =self.train_step_fn(snake_case    ,				snake_case    ,				**snake_case)
                running_loss += jax_utils.unreplicate(metrics['loss'])
                i += 1
                if i % args.logging_steps == 0:
                   _UpperCAmelCase  :  List[Any]				       =jax_utils.unreplicate(state.step)
                   _UpperCAmelCase  :  Tuple				       =running_loss.item() / i
                   _UpperCAmelCase  :  List[Any]				       =self.scheduler_fn(state_step - 1)

                   _UpperCAmelCase  :  Dict				       =self.evaluate(snake_case    ,				snake_case)
                   _UpperCAmelCase  :  Optional[int]				       ={
                       'step': state_step.item(),
                       'eval_loss': eval_loss.item(),
                       'tr_loss': tr_loss,
                       'lr': lr.item(),
                   }
                   tqdm.write(str(snake_case))
                   self.logger.log(snake_case    ,				commit=snake_case)

                if i % args.save_steps == 0:
                   self.save_checkpoint(args.save_dir + f"-e{epoch}-s{i}"    ,				state=snake_case)






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




          '''simple docstring'''
          _UpperCAmelCase  :  Optional[Any]				       =get_batched_dataset(snake_case    ,				self.args.batch_size)
          _UpperCAmelCase  :  Tuple				       =len(snake_case) // self.args.batch_size
          _UpperCAmelCase  :  Tuple				       =jnp.array(0    ,				dtype=jnp.floataa)
          _UpperCAmelCase  :  List[str]				       =0
          for batch in tqdm(snake_case    ,				total=snake_case    ,				desc='Evaluating ... '):
             _UpperCAmelCase  :  Any				       =self.data_collator(snake_case)
             _UpperCAmelCase  :  Dict				       =self.val_step_fn(snake_case    ,				**snake_case)
             running_loss += jax_utils.unreplicate(metrics['loss'])
             i += 1
          return running_loss / i






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




          '''simple docstring'''
          _UpperCAmelCase  :  int				       =jax_utils.unreplicate(snake_case)
          print(f"SAVING CHECKPOINT IN {save_dir}"    ,				end=' ... ')
          self.model_save_fn(snake_case    ,				params=state.params)
          with open(os.path.join(snake_case    ,				'opt_state.msgpack')    ,				'wb') as f:
             f.write(to_bytes(state.opt_state))
          joblib.dump(self.args    ,				os.path.join(snake_case    ,				'args.joblib'))
          joblib.dump(self.data_collator    ,				os.path.join(snake_case    ,				'data_collator.joblib'))
          with open(os.path.join(snake_case    ,				'training_state.json')    ,				'w') as f:
             json.dump({'step': state.step.item()}    ,				snake_case)
          print('DONE')






def lowerCamelCase__     (		__lowerCamelCase			:				str			,       __lowerCamelCase			:				List[Any]				):







   '''simple docstring'''




   print(f"RESTORING CHECKPOINT FROM {save_dir}"			,       end=' ... '				)
   with open(os.path.join(__lowerCamelCase			,       'flax_model.msgpack'				)			,       'rb'				) as f:
      _UpperCAmelCase  :  Union[str, Any]				       =from_bytes(state.params			,       f.read()				)

   with open(os.path.join(__lowerCamelCase			,       'opt_state.msgpack'				)			,       'rb'				) as f:
      _UpperCAmelCase  :  Dict				       =from_bytes(state.opt_state			,       f.read()				)

   _UpperCAmelCase  :  Tuple				       =joblib.load(os.path.join(__lowerCamelCase			,       'args.joblib'				)				)
   _UpperCAmelCase  :  List[Any]				       =joblib.load(os.path.join(__lowerCamelCase			,       'data_collator.joblib'				)				)

   with open(os.path.join(__lowerCamelCase			,       'training_state.json'				)			,       'r'				) as f:
      _UpperCAmelCase  :  Any				       =json.load(__lowerCamelCase				)
   _UpperCAmelCase  :  int				       =training_state['step']

   print('DONE'				)
   return params, opt_state, step, args, data_collator





def lowerCamelCase__     (		__lowerCamelCase			:				List[Any]			,       __lowerCamelCase			:				Optional[Any]			,       __lowerCamelCase			:				Tuple			,       __lowerCamelCase			:				Union[str, Any]				):







   '''simple docstring'''




   _UpperCAmelCase  :  Dict				       =num_train_steps - warmup_steps
   _UpperCAmelCase  :  Union[str, Any]				       =optax.linear_schedule(init_value=__lowerCamelCase			,       end_value=__lowerCamelCase			,       transition_steps=__lowerCamelCase				)
   _UpperCAmelCase  :  Union[str, Any]				       =optax.linear_schedule(init_value=__lowerCamelCase			,       end_value=1e-7			,       transition_steps=__lowerCamelCase				)
   _UpperCAmelCase  :  Dict				       =optax.join_schedules(schedules=[warmup_fn, decay_fn]			,       boundaries=[warmup_steps]				)
   return lr





def lowerCamelCase__     (		__lowerCamelCase			:				List[str]			,       __lowerCamelCase			:				Optional[Any]			,       __lowerCamelCase			:				Optional[Any]			,       __lowerCamelCase			:				List[Any]			,       __lowerCamelCase			:				Tuple				):







   '''simple docstring'''




   def weight_decay_mask(__lowerCamelCase			:				Optional[Any]				):
      _UpperCAmelCase  :  Tuple				       =traverse_util.flatten_dict(__lowerCamelCase				)
      _UpperCAmelCase  :  Optional[Any]				       ={k: (v[-1] != 'bias' and v[-2:] != ('LayerNorm', 'scale')) for k, v in params.items()}
      return traverse_util.unflatten_dict(__lowerCamelCase				)

   _UpperCAmelCase  :  Tuple				       =scheduler_fn(__lowerCamelCase			,       __lowerCamelCase			,       __lowerCamelCase			,       __lowerCamelCase				)

   _UpperCAmelCase  :  Optional[Any]				       =optax.adamw(learning_rate=__lowerCamelCase			,       weight_decay=__lowerCamelCase			,       mask=__lowerCamelCase				)
   return tx, lr





'''simple docstring'''






import gc
import random
import unittest

import numpy as np
import torch
from PIL import Image
from transformers import XLMRobertaTokenizerFast

from diffusers import DDIMScheduler, KandinskyImgaImgPipeline, KandinskyPriorPipeline, UNetaDConditionModel, VQModel
from diffusers.pipelines.kandinsky.text_encoder import MCLIPConfig, MultilingualCLIP
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    __magic_name__				(						lowerCAmelCase   ,unittest.TestCase							):
       UpperCAmelCase		 =KandinskyImgaImgPipeline
       UpperCAmelCase		 =["prompt", "image_embeds", "negative_image_embeds", "image"]
       UpperCAmelCase		 =[
           "prompt",
           "negative_prompt",
           "image_embeds",
           "negative_image_embeds",
           "image",
       ]
       UpperCAmelCase		 =[
           "generator",
           "height",
           "width",
           "strength",
           "guidance_scale",
           "negative_prompt",
           "num_inference_steps",
           "return_dict",
           "guidance_scale",
           "num_images_per_prompt",
           "output_type",
           "return_dict",
       ]
       UpperCAmelCase		 =False






       @property
       def      lowerCAmelCase						(     self)		->    str:




          '''simple docstring'''
          return 3_2






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




          '''simple docstring'''
          return 3_2






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




          '''simple docstring'''
          return self.time_input_dim






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




          '''simple docstring'''
          return self.time_input_dim * 4






       @property
       def      lowerCAmelCase						(     self)		->    Tuple:




          '''simple docstring'''
          return 1_0_0






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




          '''simple docstring'''
          _UpperCAmelCase  :  int				       =XLMRobertaTokenizerFast.from_pretrained('YiYiXu/tiny-random-mclip-base')
          return tokenizer






       @property
       def      lowerCAmelCase						(     self)		->    Dict:




          '''simple docstring'''
          torch.manual_seed(0)
          _UpperCAmelCase  :  Optional[Any]				       =MCLIPConfig(
              numDims=self.cross_attention_dim    ,				transformerDimensions=self.text_embedder_hidden_size    ,				hidden_size=self.text_embedder_hidden_size    ,				intermediate_size=3_7    ,				num_attention_heads=4    ,				num_hidden_layers=5    ,				vocab_size=1_0_0_5    ,				)

          _UpperCAmelCase  :  List[Any]				       =MultilingualCLIP(snake_case)
          _UpperCAmelCase  :  List[Any]				       =text_encoder.eval()

          return text_encoder






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




          '''simple docstring'''
          torch.manual_seed(0)

          _UpperCAmelCase  :  Union[str, Any]				       ={
              'in_channels': 4,
              # Out channels is double in channels because predicts mean and variance
              'out_channels': 8,
              'addition_embed_type': 'text_image',
              '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': 'text_image_proj',
              'cross_attention_dim': self.cross_attention_dim,
              'attention_head_dim': 4,
              'resnet_time_scale_shift': 'scale_shift',
              'class_embed_type': None,
          }

          _UpperCAmelCase  :  Tuple				       =UNetaDConditionModel(**snake_case)
          return model






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




          '''simple docstring'''
          return {
              "block_out_channels": [3_2, 6_4],
              "down_block_types": ["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",
              ],
              "vq_embed_dim": 4,
          }






       @property
       def      lowerCAmelCase						(     self)		->    int:




          '''simple docstring'''
          torch.manual_seed(0)
          _UpperCAmelCase  :  List[Any]				       =VQModel(**self.dummy_movq_kwargs)
          return model






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




          '''simple docstring'''
          _UpperCAmelCase  :  Tuple				       =self.dummy_text_encoder
          _UpperCAmelCase  :  Any				       =self.dummy_tokenizer
          _UpperCAmelCase  :  List[str]				       =self.dummy_unet
          _UpperCAmelCase  :  Union[str, Any]				       =self.dummy_movq

          _UpperCAmelCase  :  List[str]				       ={
              'num_train_timesteps': 1_0_0_0,
              'beta_schedule': 'linear',
              'beta_start': 0.0_00_85,
              'beta_end': 0.0_12,
              'clip_sample': False,
              'set_alpha_to_one': False,
              'steps_offset': 0,
              'prediction_type': 'epsilon',
              'thresholding': False,
          }

          _UpperCAmelCase  :  Dict				       =DDIMScheduler(**snake_case)

          _UpperCAmelCase  :  Optional[int]				       ={
              'text_encoder': text_encoder,
              'tokenizer': tokenizer,
              'unet': unet,
              'scheduler': scheduler,
              'movq': movq,
          }

          return components






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




          '''simple docstring'''
          _UpperCAmelCase  :  List[Any]				       =floats_tensor((1, self.cross_attention_dim)    ,				rng=random.Random(snake_case)).to(snake_case)
          _UpperCAmelCase  :  Tuple				       =floats_tensor((1, self.cross_attention_dim)    ,				rng=random.Random(seed + 1)).to(snake_case)
          # create init_image
          _UpperCAmelCase  :  Dict				       =floats_tensor((1, 3, 6_4, 6_4)    ,				rng=random.Random(snake_case)).to(snake_case)
          _UpperCAmelCase  :  Optional[int]				       =image.cpu().permute(0    ,				2    ,				3    ,				1)[0]
          _UpperCAmelCase  :  List[Any]				       =Image.fromarray(np.uinta(snake_case)).convert('RGB').resize((2_5_6, 2_5_6))

          if str(snake_case).startswith('mps'):
             _UpperCAmelCase  :  str				       =torch.manual_seed(snake_case)
          else:
             _UpperCAmelCase  :  Optional[Any]				       =torch.Generator(device=snake_case).manual_seed(snake_case)
          _UpperCAmelCase  :  str				       ={
              'prompt': 'horse',
              'image': init_image,
              'image_embeds': image_embeds,
              'negative_image_embeds': negative_image_embeds,
              'generator': generator,
              'height': 6_4,
              'width': 6_4,
              'num_inference_steps': 1_0,
              'guidance_scale': 7.0,
              'strength': 0.2,
              'output_type': 'np',
          }
          return inputs






       def      lowerCAmelCase						(     self)		->    str:




          '''simple docstring'''
          _UpperCAmelCase  :  List[Any]				       ='cpu'

          _UpperCAmelCase  :  Dict				       =self.get_dummy_components()

          _UpperCAmelCase  :  Optional[Any]				       =self.pipeline_class(**snake_case)
          _UpperCAmelCase  :  List[str]				       =pipe.to(snake_case)

          pipe.set_progress_bar_config(disable=snake_case)

          _UpperCAmelCase  :  List[Any]				       =pipe(**self.get_dummy_inputs(snake_case))
          _UpperCAmelCase  :  List[str]				       =output.images

          _UpperCAmelCase  :  Any				       =pipe(
              **self.get_dummy_inputs(snake_case)    ,				return_dict=snake_case    ,				)[0]

          _UpperCAmelCase  :  Any				       =image[0, -3:, -3:, -1]
          _UpperCAmelCase  :  Tuple				       =image_from_tuple[0, -3:, -3:, -1]

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

          _UpperCAmelCase  :  List[Any]				       =np.array(
              [0.61_47_49_43, 0.6_07_35_39, 0.43_30_85_44, 0.5_92_82_69, 0.47_49_35_95, 0.46_75_59_73, 0.4_61_38_38, 0.45_36_87_97, 0.50_11_92_33])
          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    __magic_name__				(						unittest.TestCase							):






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




          '''simple docstring'''
          # clean up the VRAM after each test
          super().tearDown()
          gc.collect()
          torch.cuda.empty_cache()






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




          '''simple docstring'''
          _UpperCAmelCase  :  Tuple				       =load_numpy(
              'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'
              '/kandinsky/kandinsky_img2img_frog.npy')

          _UpperCAmelCase  :  List[Any]				       =load_image(
              'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main' '/kandinsky/cat.png')
          _UpperCAmelCase  :  Union[str, Any]				       ='A red cartoon frog, 4k'

          _UpperCAmelCase  :  Union[str, Any]				       =KandinskyPriorPipeline.from_pretrained(
              'kandinsky-community/kandinsky-2-1-prior'    ,				torch_dtype=torch.floataa)
          pipe_prior.to(snake_case)

          _UpperCAmelCase  :  Optional[Any]				       =KandinskyImgaImgPipeline.from_pretrained(
              'kandinsky-community/kandinsky-2-1'    ,				torch_dtype=torch.floataa)
          _UpperCAmelCase  :  Tuple				       =pipeline.to(snake_case)

          pipeline.set_progress_bar_config(disable=snake_case)

          _UpperCAmelCase  :  List[str]				       =torch.Generator(device='cpu').manual_seed(0)
          _UpperCAmelCase       ,			_UpperCAmelCase  :  List[Any]				       =pipe_prior(
              snake_case    ,				generator=snake_case    ,				num_inference_steps=5    ,				negative_prompt=''    ,				).to_tuple()

          _UpperCAmelCase  :  str				       =pipeline(
              snake_case    ,				image=snake_case    ,				image_embeds=snake_case    ,				negative_image_embeds=snake_case    ,				generator=snake_case    ,				num_inference_steps=1_0_0    ,				height=7_6_8    ,				width=7_6_8    ,				strength=0.2    ,				output_type='np'    ,				)

          _UpperCAmelCase  :  Any				       =output.images[0]

          assert image.shape == (7_6_8, 7_6_8, 3)

          assert_mean_pixel_difference(snake_case    ,				snake_case)

