code: 
  
  
  
  import copy
  import logging
  import random
  
  import joblib
  import numpy as np
  import torch
  import torch.nn as nn
  from torch.utils.data import DataLoader
  from tqdm import tqdm
  
  from transformers import AdamW, GPTaLMHeadModel, get_linear_schedule_with_warmup
  
  
  lowerCAmelCase__	:  int 		=logging.getLogger(__name__)
  
  
  
  
  
  
  def 			__lowercase     (				a__     )       ->				Union[str, Any]:
  							random.seed(a__     )
  							np.random.seed(a__     )
  							torch.manual_seed(a__     )
  							torch.cuda.manual_seed_all(a__     )
  
  
  
  
  
  
  def 			__lowercase     (				a__							,     a__							,     a__     )       ->				List[Any]:
  							model.eval()
  							__SCREAMING_SNAKE_CASE								=      next(model.parameters()     ).device
  							__SCREAMING_SNAKE_CASE								=      1
  							__SCREAMING_SNAKE_CASE								=      torch.zeros((eval_batch_size, context_len)							,     dtype=torch.long							,     device=a__     )
  							__SCREAMING_SNAKE_CASE								=      DataLoader(a__							,     shuffle=a__							,     batch_size=a__     )
  							__SCREAMING_SNAKE_CASE								=      torch.zeros(1							,     device=a__     )
  							__SCREAMING_SNAKE_CASE								=      0
  							for batch in eval_dataloader:
  														batch.to(a__     )
  														# pad
  														context.zero_()
  														for i in range(a__     ):
  																					__SCREAMING_SNAKE_CASE								=      batch[i]
  														__SCREAMING_SNAKE_CASE								=      model(a__							,     labels=a__     )
  														eval_loss += outputs[0].sum().item()
  														nb_eval_examples += batch.size(0     )
  							__SCREAMING_SNAKE_CASE								=      eval_loss / nb_eval_examples
  							__SCREAMING_SNAKE_CASE								=      torch.exp(a__     )
  							model.train()
  							return perplexity
  
  
  
  
  
  
  def 			__lowercase     (				a__="gpt2"     )       ->				str:
  							__SCREAMING_SNAKE_CASE								=      GPTaLMHeadModel.from_pretrained(a__							,     output_hidden_states=a__     )
  							torch.save(model.state_dict()							,     model_name + 'local.pt'     )
  							return model
  
  
  
  
  
  
  def 			__lowercase     (				a__							,     a__							,     a__     )       ->				Optional[Any]:
  							__SCREAMING_SNAKE_CASE								=      copy.deepcopy(a__     )
  							model.to(a__     )
  
  							__SCREAMING_SNAKE_CASE								=      ['bias', 'LayerNorm.weight']
  							__SCREAMING_SNAKE_CASE								=      [
  							    {
  							        'params': [p for n, p in model.named_parameters() if not any(nd in n for nd in no_decay     )],
  							        'weight_decay': 0.0,
  							    },
  							    {'params': [p for n, p in model.named_parameters() if any(nd in n for nd in no_decay     )], 'weight_decay': 0.0},
  							]
  							__SCREAMING_SNAKE_CASE								=      AdamW(a__							,     lr=5E-5							,     eps=1E-8     )
  							__SCREAMING_SNAKE_CASE								=      get_linear_schedule_with_warmup(a__							,     0							,     a__     )
  							torch.cuda.empty_cache()
  							return model, lm_optimizer, lm_scheduler
  
  
  
  
  
  
  def 			__lowercase     (				a__							,     a__							,     a__							,     a__     )       ->				Optional[int]:
  							__SCREAMING_SNAKE_CASE								=      np.array(a__     ).mean()
  							__SCREAMING_SNAKE_CASE								=      np.array(a__     ).std()
  							__SCREAMING_SNAKE_CASE								=      (real_perps - avg) / std
  							__SCREAMING_SNAKE_CASE								=      list(zip(a__							,     a__							,     a__     )     )
  							joblib.dump(a__							,     a__     )
  
  
  
  
  
  
  def 			__lowercase     (				a__							,     a__							,     a__							,     a__							,     a__							,     a__							,     a__							,     a__="dev.jbl"							,     a__=recopy_gpta							,     )       ->				Dict:
  							__SCREAMING_SNAKE_CASE								=      []
  							__SCREAMING_SNAKE_CASE								=      []
  							__SCREAMING_SNAKE_CASE								=      []
  
  							# Initialize the transformer model
  							__SCREAMING_SNAKE_CASE								=      copy.deepcopy(a__     )
  							orig_model.to(device='cpu'     )
  							torch.cuda.empty_cache()
  
  							# Compute perplexity of initial transformer model for comparison
  							model.train()
  							__SCREAMING_SNAKE_CASE							, __SCREAMING_SNAKE_CASE							, __SCREAMING_SNAKE_CASE								=      recopy_model(a__							,     a__							,     a__     )
  
  							for step in tqdm(range(a__     )     ):
  														__SCREAMING_SNAKE_CASE								=      torch.zeros((1, context_len)							,     dtype=torch.long							,     device=a__     )
  														__SCREAMING_SNAKE_CASE								=      random.choice(a__     )
  														__SCREAMING_SNAKE_CASE								=      random.randint(0							,     len(story[0]     ) - context_len - 1     )
  														__SCREAMING_SNAKE_CASE								=      story[0][start : start + context_len]
  														lm_optimizer.zero_grad()
  														__SCREAMING_SNAKE_CASE								=      model(a__							,     labels=a__     )
  														__SCREAMING_SNAKE_CASE								=      outputs[0]
  														__SCREAMING_SNAKE_CASE								=      compute_perplexity(a__							,     a__							,     a__     )
  														model.train()
  														lm_loss.backward()
  														# Do LM backprop
  														torch.nn.utils.clip_grad_norm_(model.parameters()							,     3.0     )
  														lm_optimizer.step()
  														lm_scheduler.step()  # Update learning rate schedule
  
  														# Compute perplexity after back-propagating on the selected context
  														__SCREAMING_SNAKE_CASE								=      compute_perplexity(a__							,     a__							,     a__     )
  
  														# Periodically save the stored (X, IG(X)) pairs
  														if step % 10_00 == 0 and step > 1:
  																					intermittent_save(a__							,     a__							,     a__							,     a__     )
  
  														# Reset the pretrained model to the original pretrained GPT-2 weights after each iteration
  														__SCREAMING_SNAKE_CASE							, __SCREAMING_SNAKE_CASE							, __SCREAMING_SNAKE_CASE								=      recopy_model(a__							,     a__							,     a__     )
  
  														past_perps.append(past_perp.item()     )
  														real_perps.append(orig_perp - real_perp.item()     )
  														contexts.append(np.array(context.cpu()     )     )
  
  							intermittent_save(a__							,     a__							,     a__							,     a__     )
  
  
  
  
  
  
  def 			__lowercase     (				a__							,     a__="data/tokenized_stories_train_wikitext103.jbl"							,     a__=1_00							,     a__=10_26							,     a__=True     )       ->				int:
  							__SCREAMING_SNAKE_CASE								=      joblib.load(a__     )
  							print('data loaded'     )
  							__SCREAMING_SNAKE_CASE								=      []
  							if trim:
  														for i, example in enumerate(a__     ):
  																					if len(example[0]     ) > min_len:
  																												__SCREAMING_SNAKE_CASE								=      random.randint(0							,     len(example[0]     ) - context_len - 1     )
  																												objective_set.append(example[0, start : start + context_len]     )
  																					if len(a__     ) >= number:
  																												break
  														__SCREAMING_SNAKE_CASE								=      []
  														for j in range(i + 1							,     len(a__     )     ):
  																					if len(data[j][0]     ) > min_len:
  																												train_data.append(data[j]     )
  							else:
  														__SCREAMING_SNAKE_CASE								=      data[0:number]
  														__SCREAMING_SNAKE_CASE								=      data[number:]
  
  							joblib.dump(a__							,     'objective_set.jbl'     )
  							print('objective set saved'     )
  							return train_data, objective_set
  
  
  
  
  
  
  def 			__lowercase     (				a__							,     a__							,     a__							,     a__							,     a__=50							,     a__="secondary_learner.pt"     )       ->				List[Any]:
  							__SCREAMING_SNAKE_CASE								=      torch.device('cuda:0' if torch.cuda.is_available() else 'cpu'     )
  							# We will use the first 512 pairs from our dataset as a test set for
  							# our secondary learner and the rest to train
  							__SCREAMING_SNAKE_CASE								=      train_dataset[:5_12]
  							__SCREAMING_SNAKE_CASE								=      train_dataset[5_12:]
  							__SCREAMING_SNAKE_CASE								=      DataLoader(a__							,     shuffle=a__							,     batch_size=a__     )
  							__SCREAMING_SNAKE_CASE								=      DataLoader(a__							,     shuffle=a__							,     batch_size=a__     )
  
  							# secondary learner model set up
  							__SCREAMING_SNAKE_CASE								=      nn.MSELoss()
  							__SCREAMING_SNAKE_CASE								=      nn.MSELoss(reduction='sum'     )
  							secondary_learner.to(a__     )
  							__SCREAMING_SNAKE_CASE								=      torch.optim.Adam(secondary_learner.parameters()							,     lr=0.00001     )
  							secondary_learner.train()
  
  							# TODO in original code this is written as number of actual batches seen
  							# not number of items seen but other places it is number of items instead.
  							# improve consistency! changed this to epochs for clarity
  							__SCREAMING_SNAKE_CASE								=      float('inf'     )
  							# Iterate through batches until we've used max_steps batches
  							for epoch in range(int(a__     )     ):
  														__SCREAMING_SNAKE_CASE								=      0.0
  														secondary_learner.train()
  														for step, batch in enumerate(a__     ):
  																					__SCREAMING_SNAKE_CASE								=      batch[0].to(a__     )
  																					__SCREAMING_SNAKE_CASE								=      batch[1].to(a__     )
  																					__SCREAMING_SNAKE_CASE								=      secondary_learner(a__     )
  																					q_optimizer.zero_grad()
  																					__SCREAMING_SNAKE_CASE								=      loss(a__							,     real_q.float()     )
  																					q_loss.backward()
  																					q_optimizer.step()
  																					tr_q_loss += q_loss.item()
  
  																					# model trains fairly quickly so we won't wait for a full epoch
  																					# eval is triggered at eval_freq and end of epochs
  																					if (step % eval_freq == 0 and step > 0) or ((step + 1) == len(a__     )):
  																												__SCREAMING_SNAKE_CASE								=      tr_q_loss / (step + 1)
  
  																												secondary_learner.eval()
  																												__SCREAMING_SNAKE_CASE								=      0.0
  																												__SCREAMING_SNAKE_CASE								=      0.0
  																												__SCREAMING_SNAKE_CASE								=      []
  																												__SCREAMING_SNAKE_CASE								=      []
  																												# Compute performance of the secondary learner after this batch
  																												for stepa, batcha in enumerate(a__     ):
  																																			__SCREAMING_SNAKE_CASE								=      batcha[0].to(a__     )
  																																			__SCREAMING_SNAKE_CASE								=      batcha[1].to(a__     )
  																																			__SCREAMING_SNAKE_CASE								=      secondary_learner(a__     )
  																																			q_lossa += test_loss(a__							,     a__     ).item()
  																																			sum_qa += torch.sum(a__     ).item()
  																																			for ei, i in enumerate(predicted_qa.cpu().detach().numpy()     ):
  																																										predicted.append(i.item()     )
  																																			for ei, i in enumerate(real_qa.cpu().detach().numpy()     ):
  																																										actual.append(i.item()     )
  
  																												q_lossa /= len(a__     )
  																												print(
  																												    'Epoch: '							,     a__							,     'step: '							,     a__							,     'Avg. q:'							,     sum_qa / len(a__     )							,     'Train Loss: '							,     a__							,     'Test Loss: '							,     a__							,     )
  																												if q_lossa < best_test_loss:
  																																			joblib.dump((predicted, actual)							,     'pred_vs_actual.jbl'     )
  																																			torch.save(secondary_learner.state_dict()							,     a__     )
  																																			__SCREAMING_SNAKE_CASE								=      q_lossa
  
  																					secondary_learner.train()
  							return secondary_learner
  
  class      UpperCAmelCase_						(     nn.Module ):
  
  						'''simple docstring'''
  
  
  
  
  						def __init__(					self   ,							_A	):
  
  
  
  
  
  
  
  													'''simple docstring'''
  
  													super(_A   ,							self	).__init__()
  													__SCREAMING_SNAKE_CASE								=      model.transformer.wte
  													__SCREAMING_SNAKE_CASE								=      copy.deepcopy(model.transformer.wte.weight	)
  													__SCREAMING_SNAKE_CASE								=      nn.Convad(self.embeddings.weight.size(1	)   ,							256   ,							3   ,							padding=1	)
  													__SCREAMING_SNAKE_CASE								=      nn.Sequential(nn.Linear(256   ,							32	)   ,							nn.Dropout(p=0.1	)   ,							nn.Linear(32   ,							32	)   ,							nn.Linear(32   ,							1	)	)
  
  
  
  
  
  
  						def  _A			(					self   ,							_A	):
  
  
  
  
  
  
  
  													'''simple docstring'''
  
  													__SCREAMING_SNAKE_CASE								=      torch.max(self.conv(self.embeddings(_A	).squeeze(1	).transpose(1   ,							2	)	)   ,							2	)[0]
  													__SCREAMING_SNAKE_CASE								=      self.fc(_A	)
  													return qs.squeeze(1	)
  
  
  
  
  
  
  						@classmethod
  						def  _A			(					cls   ,							_A   ,							_A	):
  
  
  
  
  
  
  
  													'''simple docstring'''
  
  													__SCREAMING_SNAKE_CASE								=      cls(_A	)  # this calls __init__
  													__SCREAMING_SNAKE_CASE								=      torch.load(_A	)
  													secondary_learner.load_state_dict(_A	)
  													__SCREAMING_SNAKE_CASE								=      model.transformer.wte
  													__SCREAMING_SNAKE_CASE								=      copy.deepcopy(model.transformer.wte.weight	)
  													return secondary_learner
  
  
  
  
  
  
  
style_context: 
  
  
  
  import gc
  import random
  import tempfile
  import unittest
  
  import numpy as np
  import torch
  from transformers import CLIPTextConfig, CLIPTextModel, CLIPTokenizer
  
  from diffusers import AutoencoderKL, DDIMScheduler, LMSDiscreteScheduler, PNDMScheduler, UNetaDConditionModel
  from diffusers.pipelines.semantic_stable_diffusion import SemanticStableDiffusionPipeline as StableDiffusionPipeline
  from diffusers.utils import floats_tensor, nightly, torch_device
  from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu
  
  
  enable_full_determinism()
  
  class      UpperCAmelCase_						(     unittest.TestCase ):
  
  						'''simple docstring'''
  
  
  
  
  						def  _A			(					self	):
  
  
  
  
  
  
  
  													'''simple docstring'''
  
  													super().tearDown()
  													gc.collect()
  													torch.cuda.empty_cache()
  
  
  
  
  
  
  						@property
  						def  _A			(					self	):
  
  
  
  
  
  
  
  													'''simple docstring'''
  
  													__SCREAMING_SNAKE_CASE								=      1
  													__SCREAMING_SNAKE_CASE								=      3
  													__SCREAMING_SNAKE_CASE								=      (32, 32)
  
  													__SCREAMING_SNAKE_CASE								=      floats_tensor((batch_size, num_channels) + sizes   ,							rng=random.Random(0	)	).to(_A	)
  													return image
  
  
  
  
  
  
  						@property
  						def  _A			(					self	):
  
  
  
  
  
  
  
  													'''simple docstring'''
  
  													torch.manual_seed(0	)
  													__SCREAMING_SNAKE_CASE								=      UNetaDConditionModel(
  													    block_out_channels=(32, 64)   ,							layers_per_block=2   ,							sample_size=32   ,							in_channels=4   ,							out_channels=4   ,							down_block_types=('DownBlock2D', 'CrossAttnDownBlock2D')   ,							up_block_types=('CrossAttnUpBlock2D', 'UpBlock2D')   ,							cross_attention_dim=32   ,							)
  													return model
  
  
  
  
  
  
  						@property
  						def  _A			(					self	):
  
  
  
  
  
  
  
  													'''simple docstring'''
  
  													torch.manual_seed(0	)
  													__SCREAMING_SNAKE_CASE								=      AutoencoderKL(
  													    block_out_channels=[32, 64]   ,							in_channels=3   ,							out_channels=3   ,							down_block_types=['DownEncoderBlock2D', 'DownEncoderBlock2D']   ,							up_block_types=['UpDecoderBlock2D', 'UpDecoderBlock2D']   ,							latent_channels=4   ,							)
  													return model
  
  
  
  
  
  
  						@property
  						def  _A			(					self	):
  
  
  
  
  
  
  
  													'''simple docstring'''
  
  													torch.manual_seed(0	)
  													__SCREAMING_SNAKE_CASE								=      CLIPTextConfig(
  													    bos_token_id=0   ,							eos_token_id=2   ,							hidden_size=32   ,							intermediate_size=37   ,							layer_norm_eps=1e-05   ,							num_attention_heads=4   ,							num_hidden_layers=5   ,							pad_token_id=1   ,							vocab_size=1_000   ,							)
  													return CLIPTextModel(_A	)
  
  
  
  
  
  
  						@property
  						def  _A			(					self	):
  
  
  
  
  
  
  
  													'''simple docstring'''
  
  													def extract(*_A   ,							**_A	):
  
  																				class      UpperCAmelCase_						:
  
  																										'''simple docstring'''
  
  
  
  
  																										def __init__(					self	):
  
  
  
  
  
  
  
  																																	'''simple docstring'''
  
  																																	__SCREAMING_SNAKE_CASE								=      torch.ones([0]	)
  
  
  
  
  
  
  																										def  _A			(					self   ,							_A	):
  
  
  
  
  
  
  
  																																	'''simple docstring'''
  
  																																	self.pixel_values.to(_A	)
  																																	return self
  
  
  
  
  
  
  
  																				return Out()
  
  													return extract
  
  
  
  
  
  
  						def  _A			(					self	):
  
  
  
  
  
  
  
  													'''simple docstring'''
  
  													__SCREAMING_SNAKE_CASE								=      'cpu'  # ensure determinism for the device-dependent torch.Generator
  													__SCREAMING_SNAKE_CASE								=      self.dummy_cond_unet
  													__SCREAMING_SNAKE_CASE								=      DDIMScheduler(
  													    beta_start=0.0_0_0_8_5   ,							beta_end=0.0_1_2   ,							beta_schedule='scaled_linear'   ,							clip_sample=_A   ,							set_alpha_to_one=_A   ,							)
  
  													__SCREAMING_SNAKE_CASE								=      self.dummy_vae
  													__SCREAMING_SNAKE_CASE								=      self.dummy_text_encoder
  													__SCREAMING_SNAKE_CASE								=      CLIPTokenizer.from_pretrained('hf-internal-testing/tiny-random-clip'	)
  
  													# make sure here that pndm scheduler skips prk
  													__SCREAMING_SNAKE_CASE								=      StableDiffusionPipeline(
  													    unet=_A   ,							scheduler=_A   ,							vae=_A   ,							text_encoder=_A   ,							tokenizer=_A   ,							safety_checker=_A   ,							feature_extractor=self.dummy_extractor   ,							)
  													__SCREAMING_SNAKE_CASE								=      sd_pipe.to(_A	)
  													sd_pipe.set_progress_bar_config(disable=_A	)
  
  													__SCREAMING_SNAKE_CASE								=      'A painting of a squirrel eating a burger'
  
  													__SCREAMING_SNAKE_CASE								=      torch.Generator(device=_A	).manual_seed(0	)
  													__SCREAMING_SNAKE_CASE								=      sd_pipe([prompt]   ,							generator=_A   ,							guidance_scale=6.0   ,							num_inference_steps=2   ,							output_type='np'	)
  													__SCREAMING_SNAKE_CASE								=      output.images
  
  													__SCREAMING_SNAKE_CASE								=      torch.Generator(device=_A	).manual_seed(0	)
  													__SCREAMING_SNAKE_CASE								=      sd_pipe(
  													    [prompt]   ,							generator=_A   ,							guidance_scale=6.0   ,							num_inference_steps=2   ,							output_type='np'   ,							return_dict=_A   ,							)[0]
  
  													__SCREAMING_SNAKE_CASE								=      image[0, -3:, -3:, -1]
  													__SCREAMING_SNAKE_CASE								=      image_from_tuple[0, -3:, -3:, -1]
  
  													assert image.shape == (1, 64, 64, 3)
  													__SCREAMING_SNAKE_CASE								=      np.array([0.5_7_5_3, 0.6_1_1_4, 0.5_0_0_1, 0.5_0_3_4, 0.5_4_7_0, 0.4_7_2_9, 0.4_9_7_1, 0.4_8_6_7, 0.4_8_6_7]	)
  
  													assert np.abs(image_slice.flatten() - expected_slice	).max() < 1e-2
  													assert np.abs(image_from_tuple_slice.flatten() - expected_slice	).max() < 1e-2
  
  
  
  
  
  
  						def  _A			(					self	):
  
  
  
  
  
  
  
  													'''simple docstring'''
  
  													__SCREAMING_SNAKE_CASE								=      'cpu'  # ensure determinism for the device-dependent torch.Generator
  													__SCREAMING_SNAKE_CASE								=      self.dummy_cond_unet
  													__SCREAMING_SNAKE_CASE								=      PNDMScheduler(skip_prk_steps=_A	)
  													__SCREAMING_SNAKE_CASE								=      self.dummy_vae
  													__SCREAMING_SNAKE_CASE								=      self.dummy_text_encoder
  													__SCREAMING_SNAKE_CASE								=      CLIPTokenizer.from_pretrained('hf-internal-testing/tiny-random-clip'	)
  
  													# make sure here that pndm scheduler skips prk
  													__SCREAMING_SNAKE_CASE								=      StableDiffusionPipeline(
  													    unet=_A   ,							scheduler=_A   ,							vae=_A   ,							text_encoder=_A   ,							tokenizer=_A   ,							safety_checker=_A   ,							feature_extractor=self.dummy_extractor   ,							)
  													__SCREAMING_SNAKE_CASE								=      sd_pipe.to(_A	)
  													sd_pipe.set_progress_bar_config(disable=_A	)
  
  													__SCREAMING_SNAKE_CASE								=      'A painting of a squirrel eating a burger'
  													__SCREAMING_SNAKE_CASE								=      torch.Generator(device=_A	).manual_seed(0	)
  													__SCREAMING_SNAKE_CASE								=      sd_pipe([prompt]   ,							generator=_A   ,							guidance_scale=6.0   ,							num_inference_steps=2   ,							output_type='np'	)
  
  													__SCREAMING_SNAKE_CASE								=      output.images
  
  													__SCREAMING_SNAKE_CASE								=      torch.Generator(device=_A	).manual_seed(0	)
  													__SCREAMING_SNAKE_CASE								=      sd_pipe(
  													    [prompt]   ,							generator=_A   ,							guidance_scale=6.0   ,							num_inference_steps=2   ,							output_type='np'   ,							return_dict=_A   ,							)[0]
  
  													__SCREAMING_SNAKE_CASE								=      image[0, -3:, -3:, -1]
  													__SCREAMING_SNAKE_CASE								=      image_from_tuple[0, -3:, -3:, -1]
  
  													assert image.shape == (1, 64, 64, 3)
  													__SCREAMING_SNAKE_CASE								=      np.array([0.5_1_2_2, 0.5_7_1_2, 0.4_8_2_5, 0.5_0_5_3, 0.5_6_4_6, 0.4_7_6_9, 0.5_1_7_9, 0.4_8_9_4, 0.4_9_9_4]	)
  
  													assert np.abs(image_slice.flatten() - expected_slice	).max() < 1e-2
  													assert np.abs(image_from_tuple_slice.flatten() - expected_slice	).max() < 1e-2
  
  
  
  
  
  
  						def  _A			(					self	):
  
  
  
  
  
  
  
  													'''simple docstring'''
  
  													__SCREAMING_SNAKE_CASE								=      StableDiffusionPipeline.from_pretrained(
  													    'hf-internal-testing/tiny-stable-diffusion-lms-pipe'   ,							safety_checker=_A	)
  													assert isinstance(_A   ,							_A	)
  													assert isinstance(pipe.scheduler   ,							_A	)
  													assert pipe.safety_checker is None
  
  													__SCREAMING_SNAKE_CASE								=      pipe('example prompt'   ,							num_inference_steps=2	).images[0]
  													assert image is not None
  
  													# check that there's no error when saving a pipeline with one of the models being None
  													with tempfile.TemporaryDirectory() as tmpdirname:
  																				pipe.save_pretrained(_A	)
  																				__SCREAMING_SNAKE_CASE								=      StableDiffusionPipeline.from_pretrained(_A	)
  
  													# sanity check that the pipeline still works
  													assert pipe.safety_checker is None
  													__SCREAMING_SNAKE_CASE								=      pipe('example prompt'   ,							num_inference_steps=2	).images[0]
  													assert image is not None
  
  
  
  
  
  
  						@unittest.skipIf(torch_device != 'cuda'   ,							'This test requires a GPU'	)
  						def  _A			(					self	):
  
  
  
  
  
  
  
  													'''simple docstring'''
  
  													__SCREAMING_SNAKE_CASE								=      self.dummy_cond_unet
  													__SCREAMING_SNAKE_CASE								=      PNDMScheduler(skip_prk_steps=_A	)
  													__SCREAMING_SNAKE_CASE								=      self.dummy_vae
  													__SCREAMING_SNAKE_CASE								=      self.dummy_text_encoder
  													__SCREAMING_SNAKE_CASE								=      CLIPTokenizer.from_pretrained('hf-internal-testing/tiny-random-clip'	)
  
  													# put models in fp16
  													__SCREAMING_SNAKE_CASE								=      unet.half()
  													__SCREAMING_SNAKE_CASE								=      vae.half()
  													__SCREAMING_SNAKE_CASE								=      bert.half()
  
  													# make sure here that pndm scheduler skips prk
  													__SCREAMING_SNAKE_CASE								=      StableDiffusionPipeline(
  													    unet=_A   ,							scheduler=_A   ,							vae=_A   ,							text_encoder=_A   ,							tokenizer=_A   ,							safety_checker=_A   ,							feature_extractor=self.dummy_extractor   ,							)
  													__SCREAMING_SNAKE_CASE								=      sd_pipe.to(_A	)
  													sd_pipe.set_progress_bar_config(disable=_A	)
  
  													__SCREAMING_SNAKE_CASE								=      'A painting of a squirrel eating a burger'
  													__SCREAMING_SNAKE_CASE								=      sd_pipe([prompt]   ,							num_inference_steps=2   ,							output_type='np'	).images
  
  													assert image.shape == (1, 64, 64, 3)
  
  
  
  
  
  
  
  @nightly
  @require_torch_gpu
  class      UpperCAmelCase_						(     unittest.TestCase ):
  
  						'''simple docstring'''
  
  
  
  
  						def  _A			(					self	):
  
  
  
  
  
  
  
  													'''simple docstring'''
  
  													super().tearDown()
  													gc.collect()
  													torch.cuda.empty_cache()
  
  
  
  
  
  
  						def  _A			(					self	):
  
  
  
  
  
  
  
  													'''simple docstring'''
  
  													__SCREAMING_SNAKE_CASE								=      'cuda'
  													__SCREAMING_SNAKE_CASE								=      StableDiffusionPipeline.from_pretrained('runwayml/stable-diffusion-v1-5'	)
  													__SCREAMING_SNAKE_CASE								=      pipe.to(_A	)
  													pipe.set_progress_bar_config(disable=_A	)
  
  													__SCREAMING_SNAKE_CASE								=      'a photo of a cat'
  													__SCREAMING_SNAKE_CASE								=      {
  													    'editing_prompt': ['sunglasses'],
  													    'reverse_editing_direction': [False],
  													    'edit_warmup_steps': 10,
  													    'edit_guidance_scale': 6,
  													    'edit_threshold': 0.9_5,
  													    'edit_momentum_scale': 0.5,
  													    'edit_mom_beta': 0.6,
  													}
  
  													__SCREAMING_SNAKE_CASE								=      3
  													__SCREAMING_SNAKE_CASE								=      7
  
  													# no sega enabled
  													__SCREAMING_SNAKE_CASE								=      torch.Generator(_A	)
  													generator.manual_seed(_A	)
  													__SCREAMING_SNAKE_CASE								=      pipe(
  													    [prompt]   ,							generator=_A   ,							guidance_scale=_A   ,							num_inference_steps=50   ,							output_type='np'   ,							width=512   ,							height=512   ,							)
  
  													__SCREAMING_SNAKE_CASE								=      output.images
  													__SCREAMING_SNAKE_CASE								=      image[0, -3:, -3:, -1]
  													__SCREAMING_SNAKE_CASE								=      [
  													    0.3_4_6_7_3_1_1_3,
  													    0.3_8_4_9_2_7_3_3,
  													    0.3_7_5_9_7_3_5_2,
  													    0.3_4_0_8_6_3_3_5,
  													    0.3_5_6_5_0_7_4_8,
  													    0.3_5_5_7_9_2_0_5,
  													    0.3_3_8_4_7_6_3,
  													    0.3_4_3_4_0_2_3_6,
  													    0.3_5_7_3_2_7_1,
  													]
  
  													assert image.shape == (1, 512, 512, 3)
  
  													assert np.abs(image_slice.flatten() - expected_slice	).max() < 1e-2
  
  													# with sega enabled
  													# generator = torch.manual_seed(seed)
  													generator.manual_seed(_A	)
  													__SCREAMING_SNAKE_CASE								=      pipe(
  													    [prompt]   ,							generator=_A   ,							guidance_scale=_A   ,							num_inference_steps=50   ,							output_type='np'   ,							width=512   ,							height=512   ,							**_A   ,							)
  
  													__SCREAMING_SNAKE_CASE								=      output.images
  													__SCREAMING_SNAKE_CASE								=      image[0, -3:, -3:, -1]
  													__SCREAMING_SNAKE_CASE								=      [
  													    0.4_1_8_8_7_8_2_6,
  													    0.3_7_7_2_8_7_6_6,
  													    0.3_0_1_3_8_2_7_2,
  													    0.4_1_4_1_6_3_3_5,
  													    0.4_1_6_6_4_9_8_5,
  													    0.3_6_2_8_3_3_9_2,
  													    0.3_6_1_9_1_2_4_6,
  													    0.4_3_3_6_4_4_6_5,
  													    0.4_3_0_0_1_7_3_2,
  													]
  
  													assert image.shape == (1, 512, 512, 3)
  
  													assert np.abs(image_slice.flatten() - expected_slice	).max() < 1e-2
  
  
  
  
  
  
  						def  _A			(					self	):
  
  
  
  
  
  
  
  													'''simple docstring'''
  
  													__SCREAMING_SNAKE_CASE								=      'cuda'
  													__SCREAMING_SNAKE_CASE								=      StableDiffusionPipeline.from_pretrained('runwayml/stable-diffusion-v1-5'	)
  													__SCREAMING_SNAKE_CASE								=      pipe.to(_A	)
  													pipe.set_progress_bar_config(disable=_A	)
  
  													__SCREAMING_SNAKE_CASE								=      'an image of a crowded boulevard, realistic, 4k'
  													__SCREAMING_SNAKE_CASE								=      {
  													    'editing_prompt': 'crowd, crowded, people',
  													    'reverse_editing_direction': True,
  													    'edit_warmup_steps': 10,
  													    'edit_guidance_scale': 8.3,
  													    'edit_threshold': 0.9,
  													    'edit_momentum_scale': 0.5,
  													    'edit_mom_beta': 0.6,
  													}
  
  													__SCREAMING_SNAKE_CASE								=      9
  													__SCREAMING_SNAKE_CASE								=      7
  
  													# no sega enabled
  													__SCREAMING_SNAKE_CASE								=      torch.Generator(_A	)
  													generator.manual_seed(_A	)
  													__SCREAMING_SNAKE_CASE								=      pipe(
  													    [prompt]   ,							generator=_A   ,							guidance_scale=_A   ,							num_inference_steps=50   ,							output_type='np'   ,							width=512   ,							height=512   ,							)
  
  													__SCREAMING_SNAKE_CASE								=      output.images
  													__SCREAMING_SNAKE_CASE								=      image[0, -3:, -3:, -1]
  													__SCREAMING_SNAKE_CASE								=      [
  													    0.4_3_4_9_7_9_9_8,
  													    0.9_1_8_1_4_0_6_5,
  													    0.7_5_4_0_7_3_9,
  													    0.5_5_5_8_0_2_0_5,
  													    0.8_4_6_7_2_6_5,
  													    0.5_3_8_9_6_9_1,
  													    0.6_2_5_7_4_5_0_6,
  													    0.5_8_8_9_7_7_6_3,
  													    0.5_0_9_2_6_7_5_7,
  													]
  
  													assert image.shape == (1, 512, 512, 3)
  
  													assert np.abs(image_slice.flatten() - expected_slice	).max() < 1e-2
  
  													# with sega enabled
  													# generator = torch.manual_seed(seed)
  													generator.manual_seed(_A	)
  													__SCREAMING_SNAKE_CASE								=      pipe(
  													    [prompt]   ,							generator=_A   ,							guidance_scale=_A   ,							num_inference_steps=50   ,							output_type='np'   ,							width=512   ,							height=512   ,							**_A   ,							)
  
  													__SCREAMING_SNAKE_CASE								=      output.images
  													__SCREAMING_SNAKE_CASE								=      image[0, -3:, -3:, -1]
  													__SCREAMING_SNAKE_CASE								=      [
  													    0.3_0_8_9_7_1_9,
  													    0.3_0_5_0_0_1_4_4,
  													    0.2_9_0_1_6_0_4_2,
  													    0.3_0_6_3_0_9_6_4,
  													    0.3_2_5_6_8_7,
  													    0.2_9_4_1_9_2_2_5,
  													    0.2_9_0_8_0_9_1,
  													    0.2_8_7_2_3_5_9_8,
  													    0.2_7_6_9_6_2_9_4,
  													]
  
  													assert image.shape == (1, 512, 512, 3)
  
  													assert np.abs(image_slice.flatten() - expected_slice	).max() < 1e-2
  
  
  
  
  
  
  						def  _A			(					self	):
  
  
  
  
  
  
  
  													'''simple docstring'''
  
  													__SCREAMING_SNAKE_CASE								=      'cuda'
  													__SCREAMING_SNAKE_CASE								=      StableDiffusionPipeline.from_pretrained('runwayml/stable-diffusion-v1-5'	)
  													__SCREAMING_SNAKE_CASE								=      pipe.to(_A	)
  													pipe.set_progress_bar_config(disable=_A	)
  
  													__SCREAMING_SNAKE_CASE								=      'a castle next to a river'
  													__SCREAMING_SNAKE_CASE								=      {
  													    'editing_prompt': ['boat on a river, boat', 'monet, impression, sunrise'],
  													    'reverse_editing_direction': False,
  													    'edit_warmup_steps': [15, 18],
  													    'edit_guidance_scale': 6,
  													    'edit_threshold': [0.9, 0.8],
  													    'edit_momentum_scale': 0.5,
  													    'edit_mom_beta': 0.6,
  													}
  
  													__SCREAMING_SNAKE_CASE								=      48
  													__SCREAMING_SNAKE_CASE								=      7
  
  													# no sega enabled
  													__SCREAMING_SNAKE_CASE								=      torch.Generator(_A	)
  													generator.manual_seed(_A	)
  													__SCREAMING_SNAKE_CASE								=      pipe(
  													    [prompt]   ,							generator=_A   ,							guidance_scale=_A   ,							num_inference_steps=50   ,							output_type='np'   ,							width=512   ,							height=512   ,							)
  
  													__SCREAMING_SNAKE_CASE								=      output.images
  													__SCREAMING_SNAKE_CASE								=      image[0, -3:, -3:, -1]
  													__SCREAMING_SNAKE_CASE								=      [
  													    0.7_5_1_6_3_5_5_5,
  													    0.7_6_0_3_7_1_4_5,
  													    0.6_1_7_8_5,
  													    0.9_1_8_9_6_7_3,
  													    0.8_6_2_7_7_0_1,
  													    0.8_5_1_8_9_6_9_4,
  													    0.8_5_1_2_8_1_3,
  													    0.8_7_0_1_2_0_7_6,
  													    0.8_3_1_2_8_5_7,
  													]
  
  													assert image.shape == (1, 512, 512, 3)
  
  													assert np.abs(image_slice.flatten() - expected_slice	).max() < 1e-2
  
  													# with sega enabled
  													# generator = torch.manual_seed(seed)
  													generator.manual_seed(_A	)
  													__SCREAMING_SNAKE_CASE								=      pipe(
  													    [prompt]   ,							generator=_A   ,							guidance_scale=_A   ,							num_inference_steps=50   ,							output_type='np'   ,							width=512   ,							height=512   ,							**_A   ,							)
  
  													__SCREAMING_SNAKE_CASE								=      output.images
  													__SCREAMING_SNAKE_CASE								=      image[0, -3:, -3:, -1]
  													__SCREAMING_SNAKE_CASE								=      [
  													    0.7_3_5_5_3_3_6_5,
  													    0.7_5_3_7_2_7_1,
  													    0.7_4_3_4_1_9_0_5,
  													    0.6_6_4_8_0_3_5_6,
  													    0.6_4_7_2_9_2_5,
  													    0.6_3_0_3_9_4_1_6,
  													    0.6_4_8_1_2_9_0_5,
  													    0.6_7_4_9_7_1_7,
  													    0.6_5_1_7_1_0_2,
  													]
  
  													assert image.shape == (1, 512, 512, 3)
  
  													assert np.abs(image_slice.flatten() - expected_slice	).max() < 1e-2
  
  
  
  
  
  
  						def  _A			(					self	):
  
  
  
  
  
  
  
  													'''simple docstring'''
  
  													__SCREAMING_SNAKE_CASE								=      'cuda'
  													__SCREAMING_SNAKE_CASE								=      StableDiffusionPipeline.from_pretrained('runwayml/stable-diffusion-v1-5'   ,							torch_dtype=torch.floataa	)
  													__SCREAMING_SNAKE_CASE								=      pipe.to(_A	)
  													pipe.set_progress_bar_config(disable=_A	)
  
  													__SCREAMING_SNAKE_CASE								=      'a photo of a cat'
  													__SCREAMING_SNAKE_CASE								=      {
  													    'editing_prompt': ['sunglasses'],
  													    'reverse_editing_direction': [False],
  													    'edit_warmup_steps': 10,
  													    'edit_guidance_scale': 6,
  													    'edit_threshold': 0.9_5,
  													    'edit_momentum_scale': 0.5,
  													    'edit_mom_beta': 0.6,
  													}
  
  													__SCREAMING_SNAKE_CASE								=      3
  													__SCREAMING_SNAKE_CASE								=      7
  
  													# no sega enabled
  													__SCREAMING_SNAKE_CASE								=      torch.Generator(_A	)
  													generator.manual_seed(_A	)
  													__SCREAMING_SNAKE_CASE								=      pipe(
  													    [prompt]   ,							generator=_A   ,							guidance_scale=_A   ,							num_inference_steps=50   ,							output_type='np'   ,							width=512   ,							height=512   ,							)
  
  													__SCREAMING_SNAKE_CASE								=      output.images
  													__SCREAMING_SNAKE_CASE								=      image[0, -3:, -3:, -1]
  													__SCREAMING_SNAKE_CASE								=      [
  													    0.3_4_8_8_7_6_9_5,
  													    0.3_8_7_6_9_5_3,
  													    0.3_7_5,
  													    0.3_4_4_2_3_8_2_8,
  													    0.3_5_8_1_5_4_3,
  													    0.3_5_7_1_7_7_7_3,
  													    0.3_3_8_3_7_8_9,
  													    0.3_4_5_7_0_3_1_2,
  													    0.3_5_9_3_7_5,
  													]
  
  													assert image.shape == (1, 512, 512, 3)
  
  													assert np.abs(image_slice.flatten() - expected_slice	).max() < 1e-2
  
  													# with sega enabled
  													# generator = torch.manual_seed(seed)
  													generator.manual_seed(_A	)
  													__SCREAMING_SNAKE_CASE								=      pipe(
  													    [prompt]   ,							generator=_A   ,							guidance_scale=_A   ,							num_inference_steps=50   ,							output_type='np'   ,							width=512   ,							height=512   ,							**_A   ,							)
  
  													__SCREAMING_SNAKE_CASE								=      output.images
  													__SCREAMING_SNAKE_CASE								=      image[0, -3:, -3:, -1]
  													__SCREAMING_SNAKE_CASE								=      [
  													    0.4_2_2_8_5_1_5_6,
  													    0.3_6_9_1_4_0_6_2,
  													    0.2_9_0_7_7_1_4_8,
  													    0.4_2_0_4_1_0_1_6,
  													    0.4_1_9_1_8_9_4_5,
  													    0.3_5_4_9_8_0_4_7,
  													    0.3_6_1_8_1_6_4,
  													    0.4_4_2_3_8_2_8,
  													    0.4_3_1_1_5_2_3_4,
  													]
  
  													assert image.shape == (1, 512, 512, 3)
  
  													assert np.abs(image_slice.flatten() - expected_slice	).max() < 1e-2
  
  
  
  
  
  
  