import argparse
import collections
import json
import os
import re
import string
import sys

import numpy as np


__snake_case      :      Any             =       re.compile(R"""\b(a|an|the)\b""", re.UNICODE)

__snake_case      :      List[Any]             =       None







def 		_UpperCamelCase				(      )					->				Tuple:

						"""simple docstring"""



						lowerCAmelCase__        =  argparse.ArgumentParser('Official evaluation script for SQuAD version 2.0.'						)
						parser.add_argument('data_file'      ,    metavar='data.json'      ,    help='Input data JSON file.'						)
						parser.add_argument('pred_file'      ,    metavar='pred.json'      ,    help='Model predictions.'						)
						parser.add_argument(
						    '--out-file'      ,    '-o'      ,    metavar='eval.json'      ,    help='Write accuracy metrics to file (default is stdout).'						)
						parser.add_argument(
						    '--na-prob-file'      ,    '-n'      ,    metavar='na_prob.json'      ,    help='Model estimates of probability of no answer.'						)
						parser.add_argument(
						    '--na-prob-thresh'      ,    '-t'      ,    type=UpperCamelCase_      ,    default=1.0      ,    help='Predict "" if no-answer probability exceeds this (default = 1.0).'      ,    )
						parser.add_argument(
						    '--out-image-dir'      ,    '-p'      ,    metavar='out_images'      ,    default=UpperCamelCase_      ,    help='Save precision-recall curves to directory.'						)
						parser.add_argument('--verbose'      ,    '-v'      ,    action='store_true'						)
						if len(sys.argv						) == 1:
												parser.print_help()
												sys.exit(1						)
						return parser.parse_args()







def 		_UpperCamelCase				(      UpperCamelCase_    :   Dict						)					->				List[str]:

						"""simple docstring"""



						lowerCAmelCase__        =  {}
						for article in dataset:
												for p in article["paragraphs"]:
																		for qa in p["qas"]:
																								lowerCAmelCase__        =  bool(qa['answers']['text']						)
						return qid_to_has_ans







def 		_UpperCamelCase				(      UpperCamelCase_    :   List[Any]						)					->				Any:

						"""simple docstring"""



						def remove_articles(UpperCamelCase_    :   Optional[int]						):
												return ARTICLES_REGEX.sub(' '      ,    UpperCamelCase_						)

						def white_space_fix(UpperCamelCase_    :   Any						):
												return " ".join(text.split()						)

						def remove_punc(UpperCamelCase_    :   Union[str, Any]						):
												lowerCAmelCase__        =  set(string.punctuation						)
												return "".join(ch for ch in text if ch not in exclude						)

						def lower(UpperCamelCase_    :   str						):
												return text.lower()

						return white_space_fix(remove_articles(remove_punc(lower(UpperCamelCase_						)						)						)						)







def 		_UpperCamelCase				(      UpperCamelCase_    :   Union[str, Any]						)					->				Union[str, Any]:

						"""simple docstring"""



						if not s:
												return []
						return normalize_answer(UpperCamelCase_						).split()







def 		_UpperCamelCase				(      UpperCamelCase_    :   int      ,    UpperCamelCase_    :   List[str]						)					->				Tuple:

						"""simple docstring"""



						return int(normalize_answer(UpperCamelCase_						) == normalize_answer(UpperCamelCase_						)						)







def 		_UpperCamelCase				(      UpperCamelCase_    :   List[Any]      ,    UpperCamelCase_    :   List[Any]						)					->				Tuple:

						"""simple docstring"""



						lowerCAmelCase__        =  get_tokens(UpperCamelCase_						)
						lowerCAmelCase__        =  get_tokens(UpperCamelCase_						)
						lowerCAmelCase__        =  collections.Counter(UpperCamelCase_						) & collections.Counter(UpperCamelCase_						)
						lowerCAmelCase__        =  sum(common.values()						)
						if len(UpperCamelCase_						) == 0 or len(UpperCamelCase_						) == 0:
												# If either is no-answer, then F1 is 1 if they agree, 0 otherwise
												return int(gold_toks == pred_toks						)
						if num_same == 0:
												return 0
						lowerCAmelCase__        =  1.0 * num_same / len(UpperCamelCase_						)
						lowerCAmelCase__        =  1.0 * num_same / len(UpperCamelCase_						)
						lowerCAmelCase__        =  (2 * precision * recall) / (precision + recall)
						return fa







def 		_UpperCamelCase				(      UpperCamelCase_    :   str      ,    UpperCamelCase_    :   Optional[int]						)					->				List[Any]:

						"""simple docstring"""



						lowerCAmelCase__        =  {}
						lowerCAmelCase__        =  {}
						for article in dataset:
												for p in article["paragraphs"]:
																		for qa in p["qas"]:
																								lowerCAmelCase__        =  qa['id']
																								lowerCAmelCase__        =  [t for t in qa['answers']['text'] if normalize_answer(UpperCamelCase_						)]
																								if not gold_answers:
																														# For unanswerable questions, only correct answer is empty string
																														lowerCAmelCase__        =  ['']
																								if qid not in preds:
																														print(F"Missing prediction for {qid}"						)
																														continue
																								lowerCAmelCase__        =  preds[qid]
																								# Take max over all gold answers
																								lowerCAmelCase__        =  max(compute_exact(UpperCamelCase_      ,    UpperCamelCase_						) for a in gold_answers						)
																								lowerCAmelCase__        =  max(compute_fa(UpperCamelCase_      ,    UpperCamelCase_						) for a in gold_answers						)
						return exact_scores, fa_scores







def 		_UpperCamelCase				(      UpperCamelCase_    :   Dict      ,    UpperCamelCase_    :   List[Any]      ,    UpperCamelCase_    :   str      ,    UpperCamelCase_    :   List[Any]						)					->				str:

						"""simple docstring"""



						lowerCAmelCase__        =  {}
						for qid, s in scores.items():
												lowerCAmelCase__        =  na_probs[qid] > na_prob_thresh
												if pred_na:
																		lowerCAmelCase__        =  float(not qid_to_has_ans[qid]						)
												else:
																		lowerCAmelCase__        =  s
						return new_scores







def 		_UpperCamelCase				(      UpperCamelCase_    :   List[str]      ,    UpperCamelCase_    :   Any      ,    UpperCamelCase_    :   str=None						)					->				Union[str, Any]:

						"""simple docstring"""



						if not qid_list:
												lowerCAmelCase__        =  len(UpperCamelCase_						)
												return collections.OrderedDict(
												    [
												        ('exact', 100.0 * sum(exact_scores.values()						) / total),
												        ('f1', 100.0 * sum(fa_scores.values()						) / total),
												        ('total', total),
												    ]						)
						else:
												lowerCAmelCase__        =  len(UpperCamelCase_						)
												return collections.OrderedDict(
												    [
												        ('exact', 100.0 * sum(exact_scores[k] for k in qid_list						) / total),
												        ('f1', 100.0 * sum(fa_scores[k] for k in qid_list						) / total),
												        ('total', total),
												    ]						)







def 		_UpperCamelCase				(      UpperCamelCase_    :   Dict      ,    UpperCamelCase_    :   str      ,    UpperCamelCase_    :   Tuple						)					->				Tuple:

						"""simple docstring"""



						for k in new_eval:
												lowerCAmelCase__        =  new_eval[k]







def 		_UpperCamelCase				(      UpperCamelCase_    :   Tuple      ,    UpperCamelCase_    :   Optional[Any]      ,    UpperCamelCase_    :   Optional[int]      ,    UpperCamelCase_    :   Optional[int]						)					->				int:

						"""simple docstring"""



						plt.step(UpperCamelCase_      ,    UpperCamelCase_      ,    color='b'      ,    alpha=0.2      ,    where='post'						)
						plt.fill_between(UpperCamelCase_      ,    UpperCamelCase_      ,    step='post'      ,    alpha=0.2      ,    color='b'						)
						plt.xlabel('Recall'						)
						plt.ylabel('Precision'						)
						plt.xlim([0.0, 1.05]						)
						plt.ylim([0.0, 1.05]						)
						plt.title(UpperCamelCase_						)
						plt.savefig(UpperCamelCase_						)
						plt.clf()







def 		_UpperCamelCase				(      UpperCamelCase_    :   Optional[int]      ,    UpperCamelCase_    :   Optional[int]      ,    UpperCamelCase_    :   int      ,    UpperCamelCase_    :   str      ,    UpperCamelCase_    :   List[str]=None      ,    UpperCamelCase_    :   Any=None						)					->				List[str]:

						"""simple docstring"""



						lowerCAmelCase__        =  sorted(UpperCamelCase_      ,    key=lambda UpperCamelCase_						: na_probs[k]						)
						lowerCAmelCase__        =  0.0
						lowerCAmelCase__        =  1.0
						lowerCAmelCase__        =  0.0
						lowerCAmelCase__        =  [1.0]
						lowerCAmelCase__        =  [0.0]
						lowerCAmelCase__        =  0.0
						for i, qid in enumerate(UpperCamelCase_						):
												if qid_to_has_ans[qid]:
																		true_pos += scores[qid]
												lowerCAmelCase__        =  true_pos / float(i + 1						)
												lowerCAmelCase__        =  true_pos / float(UpperCamelCase_						)
												if i == len(UpperCamelCase_						) - 1 or na_probs[qid] != na_probs[qid_list[i + 1]]:
																		# i.e., if we can put a threshold after this point
																		avg_prec += cur_p * (cur_r - recalls[-1])
																		precisions.append(UpperCamelCase_						)
																		recalls.append(UpperCamelCase_						)
						if out_image:
												plot_pr_curve(UpperCamelCase_      ,    UpperCamelCase_      ,    UpperCamelCase_      ,    UpperCamelCase_						)
						return {"ap": 100.0 * avg_prec}







def 		_UpperCamelCase				(      UpperCamelCase_    :   str      ,    UpperCamelCase_    :   Optional[Any]      ,    UpperCamelCase_    :   Optional[int]      ,    UpperCamelCase_    :   List[str]      ,    UpperCamelCase_    :   str      ,    UpperCamelCase_    :   Tuple						)					->				Tuple:

						"""simple docstring"""



						if out_image_dir and not os.path.exists(UpperCamelCase_						):
												os.makedirs(UpperCamelCase_						)
						lowerCAmelCase__        =  sum(1 for v in qid_to_has_ans.values() if v						)
						if num_true_pos == 0:
												return
						lowerCAmelCase__        =  make_precision_recall_eval(
						    UpperCamelCase_      ,    UpperCamelCase_      ,    UpperCamelCase_      ,    UpperCamelCase_      ,    out_image=os.path.join(UpperCamelCase_      ,    'pr_exact.png'						)      ,    title='Precision-Recall curve for Exact Match score'      ,    )
						lowerCAmelCase__        =  make_precision_recall_eval(
						    UpperCamelCase_      ,    UpperCamelCase_      ,    UpperCamelCase_      ,    UpperCamelCase_      ,    out_image=os.path.join(UpperCamelCase_      ,    'pr_f1.png'						)      ,    title='Precision-Recall curve for F1 score'      ,    )
						lowerCAmelCase__        =  {k: float(UpperCamelCase_						) for k, v in qid_to_has_ans.items()}
						lowerCAmelCase__        =  make_precision_recall_eval(
						    UpperCamelCase_      ,    UpperCamelCase_      ,    UpperCamelCase_      ,    UpperCamelCase_      ,    out_image=os.path.join(UpperCamelCase_      ,    'pr_oracle.png'						)      ,    title='Oracle Precision-Recall curve (binary task of HasAns vs. NoAns)'      ,    )
						merge_eval(UpperCamelCase_      ,    UpperCamelCase_      ,    'pr_exact'						)
						merge_eval(UpperCamelCase_      ,    UpperCamelCase_      ,    'pr_f1'						)
						merge_eval(UpperCamelCase_      ,    UpperCamelCase_      ,    'pr_oracle'						)







def 		_UpperCamelCase				(      UpperCamelCase_    :   Optional[Any]      ,    UpperCamelCase_    :   Optional[int]      ,    UpperCamelCase_    :   Dict      ,    UpperCamelCase_    :   Dict						)					->				int:

						"""simple docstring"""



						if not qid_list:
												return
						lowerCAmelCase__        =  [na_probs[k] for k in qid_list]
						lowerCAmelCase__        =  np.ones_like(UpperCamelCase_						) / float(len(UpperCamelCase_						)						)
						plt.hist(UpperCamelCase_      ,    weights=UpperCamelCase_      ,    bins=20      ,    range=(0.0, 1.0)						)
						plt.xlabel('Model probability of no-answer'						)
						plt.ylabel('Proportion of dataset'						)
						plt.title(F"Histogram of no-answer probability: {name}"						)
						plt.savefig(os.path.join(UpperCamelCase_      ,    F"na_prob_hist_{name}.png"						)						)
						plt.clf()







def 		_UpperCamelCase				(      UpperCamelCase_    :   Any      ,    UpperCamelCase_    :   Tuple      ,    UpperCamelCase_    :   Optional[int]      ,    UpperCamelCase_    :   str						)					->				int:

						"""simple docstring"""



						lowerCAmelCase__        =  sum(1 for k in qid_to_has_ans if not qid_to_has_ans[k]						)
						lowerCAmelCase__        =  num_no_ans
						lowerCAmelCase__        =  cur_score
						lowerCAmelCase__        =  0.0
						lowerCAmelCase__        =  sorted(UpperCamelCase_      ,    key=lambda UpperCamelCase_						: na_probs[k]						)
						for i, qid in enumerate(UpperCamelCase_						):
												if qid not in scores:
																		continue
												if qid_to_has_ans[qid]:
																		lowerCAmelCase__        =  scores[qid]
												else:
																		if preds[qid]:
																								lowerCAmelCase__        =  -1
																		else:
																								lowerCAmelCase__        =  0
												cur_score += diff
												if cur_score > best_score:
																		lowerCAmelCase__        =  cur_score
																		lowerCAmelCase__        =  na_probs[qid]
						return 100.0 * best_score / len(UpperCamelCase_						), best_thresh







def 		_UpperCamelCase				(      UpperCamelCase_    :   List[str]      ,    UpperCamelCase_    :   Optional[Any]      ,    UpperCamelCase_    :   str      ,    UpperCamelCase_    :   int      ,    UpperCamelCase_    :   Optional[int]      ,    UpperCamelCase_    :   List[str]						)					->				str:

						"""simple docstring"""



						lowerCAmelCase__  ,    lowerCAmelCase__        =  find_best_thresh(UpperCamelCase_      ,    UpperCamelCase_      ,    UpperCamelCase_      ,    UpperCamelCase_						)
						lowerCAmelCase__  ,    lowerCAmelCase__        =  find_best_thresh(UpperCamelCase_      ,    UpperCamelCase_      ,    UpperCamelCase_      ,    UpperCamelCase_						)
						lowerCAmelCase__        =  best_exact
						lowerCAmelCase__        =  exact_thresh
						lowerCAmelCase__        =  best_fa
						lowerCAmelCase__        =  fa_thresh







def 		_UpperCamelCase				(      )					->				Dict:

						"""simple docstring"""



						with open(OPTS.data_file						) as f:
												lowerCAmelCase__        =  json.load(UpperCamelCase_						)
												lowerCAmelCase__        =  dataset_json['data']
						with open(OPTS.pred_file						) as f:
												lowerCAmelCase__        =  json.load(UpperCamelCase_						)
						if OPTS.na_prob_file:
												with open(OPTS.na_prob_file						) as f:
																		lowerCAmelCase__        =  json.load(UpperCamelCase_						)
						else:
												lowerCAmelCase__        =  {k: 0.0 for k in preds}
						lowerCAmelCase__        =  make_qid_to_has_ans(UpperCamelCase_						)  # maps qid to True/False
						lowerCAmelCase__        =  [k for k, v in qid_to_has_ans.items() if v]
						lowerCAmelCase__        =  [k for k, v in qid_to_has_ans.items() if not v]
						lowerCAmelCase__  ,    lowerCAmelCase__        =  get_raw_scores(UpperCamelCase_      ,    UpperCamelCase_						)
						lowerCAmelCase__        =  apply_no_ans_threshold(UpperCamelCase_      ,    UpperCamelCase_      ,    UpperCamelCase_      ,    OPTS.na_prob_thresh						)
						lowerCAmelCase__        =  apply_no_ans_threshold(UpperCamelCase_      ,    UpperCamelCase_      ,    UpperCamelCase_      ,    OPTS.na_prob_thresh						)
						lowerCAmelCase__        =  make_eval_dict(UpperCamelCase_      ,    UpperCamelCase_						)
						if has_ans_qids:
												lowerCAmelCase__        =  make_eval_dict(UpperCamelCase_      ,    UpperCamelCase_      ,    qid_list=UpperCamelCase_						)
												merge_eval(UpperCamelCase_      ,    UpperCamelCase_      ,    'HasAns'						)
						if no_ans_qids:
												lowerCAmelCase__        =  make_eval_dict(UpperCamelCase_      ,    UpperCamelCase_      ,    qid_list=UpperCamelCase_						)
												merge_eval(UpperCamelCase_      ,    UpperCamelCase_      ,    'NoAns'						)
						if OPTS.na_prob_file:
												find_all_best_thresh(UpperCamelCase_      ,    UpperCamelCase_      ,    UpperCamelCase_      ,    UpperCamelCase_      ,    UpperCamelCase_      ,    UpperCamelCase_						)
						if OPTS.na_prob_file and OPTS.out_image_dir:
												run_precision_recall_analysis(UpperCamelCase_      ,    UpperCamelCase_      ,    UpperCamelCase_      ,    UpperCamelCase_      ,    UpperCamelCase_      ,    OPTS.out_image_dir						)
												histogram_na_prob(UpperCamelCase_      ,    UpperCamelCase_      ,    OPTS.out_image_dir      ,    'hasAns'						)
												histogram_na_prob(UpperCamelCase_      ,    UpperCamelCase_      ,    OPTS.out_image_dir      ,    'noAns'						)
						if OPTS.out_file:
												with open(OPTS.out_file      ,    'w'						) as f:
																		json.dump(UpperCamelCase_      ,    UpperCamelCase_						)
						else:
												print(json.dumps(UpperCamelCase_      ,    indent=2						)						)


if __name__ == "__main__":
					__snake_case      :      int             =       parse_args()
					if OPTS.out_image_dir:
										import matplotlib

										matplotlib.use("""Agg""")
										import matplotlib.pyplot as plt
					main()













import gc
import random
import unittest

import numpy as np
import torch
from PIL import Image
from transformers import CLIPTextConfig, CLIPTextModel, CLIPTokenizer

from diffusers import AutoencoderKL, DDIMScheduler, DDPMScheduler, StableDiffusionUpscalePipeline, UNetaDConditionModel
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


enable_full_determinism()


class       __SCREAMING_SNAKE_CASE   (				unittest.TestCase):






						def  UpperCamelCase__ (		self				):






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






						@property
						def  UpperCamelCase__ (		self				):






												"""simple docstring"""
												lowerCAmelCase__        =  1
												lowerCAmelCase__        =  3
												lowerCAmelCase__        =  (32, 32)

												lowerCAmelCase__        =  floats_tensor((batch_size, num_channels) + sizes       ,  rng=random.Random(0				)				).to(_UpperCamelCase				)
												return image






						@property
						def  UpperCamelCase__ (		self				):






												"""simple docstring"""
												torch.manual_seed(0				)
												lowerCAmelCase__        =  UNetaDConditionModel(
												    block_out_channels=(32, 32, 64)       ,  layers_per_block=2       ,  sample_size=32       ,  in_channels=7       ,  out_channels=4       ,  down_block_types=('DownBlock2D', 'CrossAttnDownBlock2D', 'CrossAttnDownBlock2D')       ,  up_block_types=('CrossAttnUpBlock2D', 'CrossAttnUpBlock2D', 'UpBlock2D')       ,  cross_attention_dim=32       ,  attention_head_dim=8       ,  use_linear_projection=_UpperCamelCase       ,  only_cross_attention=(True, True, False)       ,  num_class_embeds=1_00       ,  )
												return model






						@property
						def  UpperCamelCase__ (		self				):






												"""simple docstring"""
												torch.manual_seed(0				)
												lowerCAmelCase__        =  AutoencoderKL(
												    block_out_channels=[32, 32, 64]       ,  in_channels=3       ,  out_channels=3       ,  down_block_types=['DownEncoderBlock2D', 'DownEncoderBlock2D', 'DownEncoderBlock2D']       ,  up_block_types=['UpDecoderBlock2D', 'UpDecoderBlock2D', 'UpDecoderBlock2D']       ,  latent_channels=4       ,  )
												return model






						@property
						def  UpperCamelCase__ (		self				):






												"""simple docstring"""
												torch.manual_seed(0				)
												lowerCAmelCase__        =  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=10_00       ,  hidden_act='gelu'       ,  projection_dim=5_12       ,  )
												return CLIPTextModel(_UpperCamelCase				)






						def  UpperCamelCase__ (		self				):






												"""simple docstring"""
												lowerCAmelCase__        =  'cpu'  # ensure determinism for the device-dependent torch.Generator
												lowerCAmelCase__        =  self.dummy_cond_unet_upscale
												lowerCAmelCase__        =  DDPMScheduler()
												lowerCAmelCase__        =  DDIMScheduler(prediction_type='v_prediction'				)
												lowerCAmelCase__        =  self.dummy_vae
												lowerCAmelCase__        =  self.dummy_text_encoder
												lowerCAmelCase__        =  CLIPTokenizer.from_pretrained('hf-internal-testing/tiny-random-clip'				)

												lowerCAmelCase__        =  self.dummy_image.cpu().permute(0       ,  2       ,  3       ,  1				)[0]
												lowerCAmelCase__        =  Image.fromarray(np.uinta(_UpperCamelCase				)				).convert('RGB'				).resize((64, 64)				)

												# make sure here that pndm scheduler skips prk
												lowerCAmelCase__        =  StableDiffusionUpscalePipeline(
												    unet=_UpperCamelCase       ,  low_res_scheduler=_UpperCamelCase       ,  scheduler=_UpperCamelCase       ,  vae=_UpperCamelCase       ,  text_encoder=_UpperCamelCase       ,  tokenizer=_UpperCamelCase       ,  max_noise_level=3_50       ,  )
												lowerCAmelCase__        =  sd_pipe.to(_UpperCamelCase				)
												sd_pipe.set_progress_bar_config(disable=_UpperCamelCase				)

												lowerCAmelCase__        =  'A painting of a squirrel eating a burger'
												lowerCAmelCase__        =  torch.Generator(device=_UpperCamelCase				).manual_seed(0				)
												lowerCAmelCase__        =  sd_pipe(
												    [prompt]       ,  image=_UpperCamelCase       ,  generator=_UpperCamelCase       ,  guidance_scale=6.0       ,  noise_level=20       ,  num_inference_steps=2       ,  output_type='np'       ,  )

												lowerCAmelCase__        =  output.images

												lowerCAmelCase__        =  torch.Generator(device=_UpperCamelCase				).manual_seed(0				)
												lowerCAmelCase__        =  sd_pipe(
												    [prompt]       ,  image=_UpperCamelCase       ,  generator=_UpperCamelCase       ,  guidance_scale=6.0       ,  noise_level=20       ,  num_inference_steps=2       ,  output_type='np'       ,  return_dict=_UpperCamelCase       ,  )[0]

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

												lowerCAmelCase__        =  low_res_image.size[0] * 4
												assert image.shape == (1, expected_height_width, expected_height_width, 3)
												lowerCAmelCase__        =  np.array([0.31_13, 0.39_10, 0.42_72, 0.48_59, 0.50_61, 0.46_52, 0.53_62, 0.57_15, 0.56_61]				)

												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  UpperCamelCase__ (		self				):






												"""simple docstring"""
												lowerCAmelCase__        =  'cpu'  # ensure determinism for the device-dependent torch.Generator
												lowerCAmelCase__        =  self.dummy_cond_unet_upscale
												lowerCAmelCase__        =  DDPMScheduler()
												lowerCAmelCase__        =  DDIMScheduler(prediction_type='v_prediction'				)
												lowerCAmelCase__        =  self.dummy_vae
												lowerCAmelCase__        =  self.dummy_text_encoder
												lowerCAmelCase__        =  CLIPTokenizer.from_pretrained('hf-internal-testing/tiny-random-clip'				)

												lowerCAmelCase__        =  self.dummy_image.cpu().permute(0       ,  2       ,  3       ,  1				)[0]
												lowerCAmelCase__        =  Image.fromarray(np.uinta(_UpperCamelCase				)				).convert('RGB'				).resize((64, 64)				)

												# make sure here that pndm scheduler skips prk
												lowerCAmelCase__        =  StableDiffusionUpscalePipeline(
												    unet=_UpperCamelCase       ,  low_res_scheduler=_UpperCamelCase       ,  scheduler=_UpperCamelCase       ,  vae=_UpperCamelCase       ,  text_encoder=_UpperCamelCase       ,  tokenizer=_UpperCamelCase       ,  max_noise_level=3_50       ,  )
												lowerCAmelCase__        =  sd_pipe.to(_UpperCamelCase				)
												sd_pipe.set_progress_bar_config(disable=_UpperCamelCase				)

												lowerCAmelCase__        =  'A painting of a squirrel eating a burger'
												lowerCAmelCase__        =  sd_pipe(
												    2 * [prompt]       ,  image=2 * [low_res_image]       ,  guidance_scale=6.0       ,  noise_level=20       ,  num_inference_steps=2       ,  output_type='np'       ,  )
												lowerCAmelCase__        =  output.images
												assert image.shape[0] == 2

												lowerCAmelCase__        =  torch.Generator(device=_UpperCamelCase				).manual_seed(0				)
												lowerCAmelCase__        =  sd_pipe(
												    [prompt]       ,  image=_UpperCamelCase       ,  generator=_UpperCamelCase       ,  num_images_per_prompt=2       ,  guidance_scale=6.0       ,  noise_level=20       ,  num_inference_steps=2       ,  output_type='np'       ,  )
												lowerCAmelCase__        =  output.images
												assert image.shape[0] == 2







						@unittest.skipIf(torch_device != 'cuda'       ,  'This test requires a GPU'				)
						def  UpperCamelCase__ (		self				):






												"""simple docstring"""
												lowerCAmelCase__        =  self.dummy_cond_unet_upscale
												lowerCAmelCase__        =  DDPMScheduler()
												lowerCAmelCase__        =  DDIMScheduler(prediction_type='v_prediction'				)
												lowerCAmelCase__        =  self.dummy_vae
												lowerCAmelCase__        =  self.dummy_text_encoder
												lowerCAmelCase__        =  CLIPTokenizer.from_pretrained('hf-internal-testing/tiny-random-clip'				)

												lowerCAmelCase__        =  self.dummy_image.cpu().permute(0       ,  2       ,  3       ,  1				)[0]
												lowerCAmelCase__        =  Image.fromarray(np.uinta(_UpperCamelCase				)				).convert('RGB'				).resize((64, 64)				)

												# put models in fp16, except vae as it overflows in fp16
												lowerCAmelCase__        =  unet.half()
												lowerCAmelCase__        =  text_encoder.half()

												# make sure here that pndm scheduler skips prk
												lowerCAmelCase__        =  StableDiffusionUpscalePipeline(
												    unet=_UpperCamelCase       ,  low_res_scheduler=_UpperCamelCase       ,  scheduler=_UpperCamelCase       ,  vae=_UpperCamelCase       ,  text_encoder=_UpperCamelCase       ,  tokenizer=_UpperCamelCase       ,  max_noise_level=3_50       ,  )
												lowerCAmelCase__        =  sd_pipe.to(_UpperCamelCase				)
												sd_pipe.set_progress_bar_config(disable=_UpperCamelCase				)

												lowerCAmelCase__        =  'A painting of a squirrel eating a burger'
												lowerCAmelCase__        =  torch.manual_seed(0				)
												lowerCAmelCase__        =  sd_pipe(
												    [prompt]       ,  image=_UpperCamelCase       ,  generator=_UpperCamelCase       ,  num_inference_steps=2       ,  output_type='np'       ,  ).images

												lowerCAmelCase__        =  low_res_image.size[0] * 4
												assert image.shape == (1, expected_height_width, expected_height_width, 3)




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






						def  UpperCamelCase__ (		self				):






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






						def  UpperCamelCase__ (		self				):






												"""simple docstring"""
												lowerCAmelCase__        =  load_image(
												    'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'
												    '/sd2-upscale/low_res_cat.png'				)
												lowerCAmelCase__        =  load_numpy(
												    'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/sd2-upscale'
												    '/upsampled_cat.npy'				)

												lowerCAmelCase__        =  'stabilityai/stable-diffusion-x4-upscaler'
												lowerCAmelCase__        =  StableDiffusionUpscalePipeline.from_pretrained(_UpperCamelCase				)
												pipe.to(_UpperCamelCase				)
												pipe.set_progress_bar_config(disable=_UpperCamelCase				)
												pipe.enable_attention_slicing()

												lowerCAmelCase__        =  'a cat sitting on a park bench'

												lowerCAmelCase__        =  torch.manual_seed(0				)
												lowerCAmelCase__        =  pipe(
												    prompt=_UpperCamelCase       ,  image=_UpperCamelCase       ,  generator=_UpperCamelCase       ,  output_type='np'       ,  )
												lowerCAmelCase__        =  output.images[0]

												assert image.shape == (5_12, 5_12, 3)
												assert np.abs(expected_image - image				).max() < 1E-3






						def  UpperCamelCase__ (		self				):






												"""simple docstring"""
												lowerCAmelCase__        =  load_image(
												    'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'
												    '/sd2-upscale/low_res_cat.png'				)
												lowerCAmelCase__        =  load_numpy(
												    'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/sd2-upscale'
												    '/upsampled_cat_fp16.npy'				)

												lowerCAmelCase__        =  'stabilityai/stable-diffusion-x4-upscaler'
												lowerCAmelCase__        =  StableDiffusionUpscalePipeline.from_pretrained(
												    _UpperCamelCase       ,  torch_dtype=torch.floataa       ,  )
												pipe.to(_UpperCamelCase				)
												pipe.set_progress_bar_config(disable=_UpperCamelCase				)
												pipe.enable_attention_slicing()

												lowerCAmelCase__        =  'a cat sitting on a park bench'

												lowerCAmelCase__        =  torch.manual_seed(0				)
												lowerCAmelCase__        =  pipe(
												    prompt=_UpperCamelCase       ,  image=_UpperCamelCase       ,  generator=_UpperCamelCase       ,  output_type='np'       ,  )
												lowerCAmelCase__        =  output.images[0]

												assert image.shape == (5_12, 5_12, 3)
												assert np.abs(expected_image - image				).max() < 5E-1







						def  UpperCamelCase__ (		self				):






												"""simple docstring"""
												torch.cuda.empty_cache()
												torch.cuda.reset_max_memory_allocated()
												torch.cuda.reset_peak_memory_stats()

												lowerCAmelCase__        =  load_image(
												    'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'
												    '/sd2-upscale/low_res_cat.png'				)

												lowerCAmelCase__        =  'stabilityai/stable-diffusion-x4-upscaler'
												lowerCAmelCase__        =  StableDiffusionUpscalePipeline.from_pretrained(
												    _UpperCamelCase       ,  torch_dtype=torch.floataa       ,  )
												pipe.to(_UpperCamelCase				)
												pipe.set_progress_bar_config(disable=_UpperCamelCase				)
												pipe.enable_attention_slicing(1				)
												pipe.enable_sequential_cpu_offload()

												lowerCAmelCase__        =  'a cat sitting on a park bench'

												lowerCAmelCase__        =  torch.manual_seed(0				)
												lowerCAmelCase__        =  pipe(
												    prompt=_UpperCamelCase       ,  image=_UpperCamelCase       ,  generator=_UpperCamelCase       ,  num_inference_steps=5       ,  output_type='np'       ,  )

												lowerCAmelCase__        =  torch.cuda.max_memory_allocated()
												# make sure that less than 2.9 GB is allocated
												assert mem_bytes < 2.9 * 10**9





