






import json
import os
import pickle
import shutil
import tempfile
from unittest import TestCase
from unittest.mock import patch

import numpy as np
from datasets import Dataset

from transformers import is_faiss_available
from transformers.models.bart.configuration_bart import BartConfig
from transformers.models.bart.tokenization_bart import BartTokenizer
from transformers.models.bert.tokenization_bert import VOCAB_FILES_NAMES as DPR_VOCAB_FILES_NAMES
from transformers.models.dpr.configuration_dpr import DPRConfig
from transformers.models.dpr.tokenization_dpr import DPRContextEncoderTokenizer, DPRQuestionEncoderTokenizer
from transformers.models.rag.configuration_rag import RagConfig
from transformers.models.rag.retrieval_rag import CustomHFIndex, RagRetriever
from transformers.models.roberta.tokenization_roberta import VOCAB_FILES_NAMES as BART_VOCAB_FILES_NAMES
from transformers.testing_utils import require_faiss, require_sentencepiece, require_tokenizers, require_torch


if is_faiss_available():
   import faiss


@require_faiss
class  SCREAMING_SNAKE_CASE			(			lowerCAmelCase					):
     '''simple docstring'''


     def 				_A       (				self       :	str      ):
       SCREAMING_SNAKE_CASE       :      Optional[int]							=    tempfile.mkdtemp()
       SCREAMING_SNAKE_CASE       :      Union[str, Any]							=    8

       # DPR tok
       SCREAMING_SNAKE_CASE       :      Optional[Any]							=    [
           "[UNK]",
           "[CLS]",
           "[SEP]",
           "[PAD]",
           "[MASK]",
           "want",
           "##want",
           "##ed",
           "wa",
           "un",
           "runn",
           "##ing",
           ",",
           "low",
           "lowest",
       ]
       SCREAMING_SNAKE_CASE       :      List[Any]							=    os.path.join(self.tmpdirname					, "dpr_tokenizer"      )
       os.makedirs(UpperCAmelCase_					, exist_ok=UpperCAmelCase_      )
       SCREAMING_SNAKE_CASE       :      Tuple							=    os.path.join(UpperCAmelCase_					, DPR_VOCAB_FILES_NAMES["vocab_file"]      )
       with open(self.vocab_file					, "w"					, encoding="utf-8"      ) as vocab_writer:
         vocab_writer.write("".join([x + "\n" for x in vocab_tokens]      )      )

       # BART tok
       SCREAMING_SNAKE_CASE       :      str							=    [
           "l",
           "o",
           "w",
           "e",
           "r",
           "s",
           "t",
           "i",
           "d",
           "n",
           "\u0120",
           "\u0120l",
           "\u0120n",
           "\u0120lo",
           "\u0120low",
           "er",
           "\u0120lowest",
           "\u0120newer",
           "\u0120wider",
           "<unk>",
       ]
       SCREAMING_SNAKE_CASE       :      int							=    dict(zip(UpperCAmelCase_					, range(len(UpperCAmelCase_      )      )      )      )
       SCREAMING_SNAKE_CASE       :      Dict							=    ["#version: 0.2", "\u0120 l", "\u0120l o", "\u0120lo w", "e r", ""]
       SCREAMING_SNAKE_CASE       :      Optional[Any]							=    {"unk_token": "<unk>"}

       SCREAMING_SNAKE_CASE       :      Any							=    os.path.join(self.tmpdirname					, "bart_tokenizer"      )
       os.makedirs(UpperCAmelCase_					, exist_ok=UpperCAmelCase_      )
       SCREAMING_SNAKE_CASE       :      str							=    os.path.join(UpperCAmelCase_					, BART_VOCAB_FILES_NAMES["vocab_file"]      )
       SCREAMING_SNAKE_CASE       :      List[Any]							=    os.path.join(UpperCAmelCase_					, BART_VOCAB_FILES_NAMES["merges_file"]      )
       with open(self.vocab_file					, "w"					, encoding="utf-8"      ) as fp:
         fp.write(json.dumps(UpperCAmelCase_      ) + "\n"      )
       with open(self.merges_file					, "w"					, encoding="utf-8"      ) as fp:
         fp.write("\n".join(UpperCAmelCase_      )      )


     def 				_A       (				self       :	Tuple      ):
       return DPRQuestionEncoderTokenizer.from_pretrained(os.path.join(self.tmpdirname					, "dpr_tokenizer"      )      )


     def 				_A       (				self       :	str      ):
       return DPRContextEncoderTokenizer.from_pretrained(os.path.join(self.tmpdirname					, "dpr_tokenizer"      )      )


     def 				_A       (				self       :	Optional[int]      ):
       return BartTokenizer.from_pretrained(os.path.join(self.tmpdirname					, "bart_tokenizer"      )      )


     def 				_A       (				self       :	Optional[Any]      ):
       shutil.rmtree(self.tmpdirname      )


     def 				_A       (				self       :	List[str]      ):
       SCREAMING_SNAKE_CASE       :      Any							=    Dataset.from_dict(
           {
               "id": ["0", "1"],
               "text": ["foo", "bar"],
               "title": ["Foo", "Bar"],
               "embeddings": [np.ones(self.retrieval_vector_size      ), 2 * np.ones(self.retrieval_vector_size      )],
           }      )
       dataset.add_faiss_index("embeddings"					, string_factory="Flat"					, metric_type=faiss.METRIC_INNER_PRODUCT      )
       return dataset


     def 				_A       (				self       :	List[str]      ):
       SCREAMING_SNAKE_CASE       :      List[str]							=    self.get_dummy_dataset()
       SCREAMING_SNAKE_CASE       :      int							=    RagConfig(
           retrieval_vector_size=self.retrieval_vector_size					, question_encoder=DPRConfig().to_dict()					, generator=BartConfig().to_dict()					, )
       with patch("transformers.models.rag.retrieval_rag.load_dataset"      ) as mock_load_dataset:
         SCREAMING_SNAKE_CASE       :      Optional[int]							=    dataset
         SCREAMING_SNAKE_CASE       :      str							=    RagRetriever(
             UpperCAmelCase_					, question_encoder_tokenizer=self.get_dpr_tokenizer()					, generator_tokenizer=self.get_bart_tokenizer()					, )
       return retriever


     def 				_A       (				self       :	str					, UpperCAmelCase_       :	bool      ):
       SCREAMING_SNAKE_CASE       :      int							=    self.get_dummy_dataset()
       SCREAMING_SNAKE_CASE       :      Any							=    RagConfig(
           retrieval_vector_size=self.retrieval_vector_size					, question_encoder=DPRConfig().to_dict()					, generator=BartConfig().to_dict()					, index_name="custom"					, )
       if from_disk:
         SCREAMING_SNAKE_CASE       :      Optional[Any]							=    os.path.join(self.tmpdirname					, "dataset"      )
         SCREAMING_SNAKE_CASE       :      List[str]							=    os.path.join(self.tmpdirname					, "index.faiss"      )
         dataset.get_index("embeddings"      ).save(os.path.join(self.tmpdirname					, "index.faiss"      )      )
         dataset.drop_index("embeddings"      )
         dataset.save_to_disk(os.path.join(self.tmpdirname					, "dataset"      )      )
         del dataset
         SCREAMING_SNAKE_CASE       :      Any							=    RagRetriever(
             UpperCAmelCase_					, question_encoder_tokenizer=self.get_dpr_tokenizer()					, generator_tokenizer=self.get_bart_tokenizer()					, )
       else:
         SCREAMING_SNAKE_CASE       :      str							=    RagRetriever(
             UpperCAmelCase_					, question_encoder_tokenizer=self.get_dpr_tokenizer()					, generator_tokenizer=self.get_bart_tokenizer()					, index=CustomHFIndex(config.retrieval_vector_size					, UpperCAmelCase_      )					, )
       return retriever


     def 				_A       (				self       :	List[Any]      ):
       SCREAMING_SNAKE_CASE       :      Union[str, Any]							=    Dataset.from_dict(
           {
               "id": ["0", "1"],
               "text": ["foo", "bar"],
               "title": ["Foo", "Bar"],
               "embeddings": [np.ones(self.retrieval_vector_size + 1      ), 2 * np.ones(self.retrieval_vector_size + 1      )],
           }      )
       dataset.add_faiss_index("embeddings"					, string_factory="Flat"					, metric_type=faiss.METRIC_INNER_PRODUCT      )

       SCREAMING_SNAKE_CASE       :      List[Any]							=    os.path.join(self.tmpdirname					, "hf_bert_base.hnswSQ8_correct_phi_128.c_index"      )
       dataset.save_faiss_index("embeddings"					, index_file_name + ".index.dpr"      )
       pickle.dump(dataset["id"]					, open(index_file_name + ".index_meta.dpr"					, "wb"      )      )

       SCREAMING_SNAKE_CASE       :      List[Any]							=    os.path.join(self.tmpdirname					, "psgs_w100.tsv.pkl"      )
       SCREAMING_SNAKE_CASE       :      List[str]							=    {sample["id"]: [sample["text"], sample["title"]] for sample in dataset}
       pickle.dump(UpperCAmelCase_					, open(UpperCAmelCase_					, "wb"      )      )

       SCREAMING_SNAKE_CASE       :      Any							=    RagConfig(
           retrieval_vector_size=self.retrieval_vector_size					, question_encoder=DPRConfig().to_dict()					, generator=BartConfig().to_dict()					, index_name="legacy"					, index_path=self.tmpdirname					, )
       SCREAMING_SNAKE_CASE       :      str							=    RagRetriever(
           UpperCAmelCase_					, question_encoder_tokenizer=self.get_dpr_tokenizer()					, generator_tokenizer=self.get_bart_tokenizer()      )
       return retriever


     def 				_A       (				self       :	int      ):
       SCREAMING_SNAKE_CASE       :      Tuple							=    1
       SCREAMING_SNAKE_CASE       :      str							=    self.get_dummy_canonical_hf_index_retriever()
       SCREAMING_SNAKE_CASE       :      Union[str, Any]							=    np.array(
           [np.ones(self.retrieval_vector_size      ), -np.ones(self.retrieval_vector_size      )]					, dtype=np.floataa      )
       SCREAMING_SNAKE_CASE       :      Union[str, Any]							=    retriever.retrieve(UpperCAmelCase_					, n_docs=UpperCAmelCase_      )
       self.assertEqual(retrieved_doc_embeds.shape					, (2, n_docs, self.retrieval_vector_size)      )
       self.assertEqual(len(UpperCAmelCase_      )					, 2      )
       self.assertEqual(sorted(doc_dicts[0]      )					, ["embeddings", "id", "text", "title"]      )
       self.assertEqual(len(doc_dicts[0]["id"]      )					, UpperCAmelCase_      )
       self.assertEqual(doc_dicts[0]["id"][0]					, "1"      )  # max inner product is reached with second doc
       self.assertEqual(doc_dicts[1]["id"][0]					, "0"      )  # max inner product is reached with first doc
       self.assertListEqual(doc_ids.tolist()					, [[1], [0]]      )


     def 				_A       (				self       :	List[Any]      ):
       SCREAMING_SNAKE_CASE       :      Optional[int]							=    self.get_dummy_canonical_hf_index_retriever()
       with tempfile.TemporaryDirectory() as tmp_dirname:
         with patch("transformers.models.rag.retrieval_rag.load_dataset"      ) as mock_load_dataset:
           SCREAMING_SNAKE_CASE       :      Tuple							=    self.get_dummy_dataset()
           retriever.save_pretrained(UpperCAmelCase_      )
           SCREAMING_SNAKE_CASE       :      int							=    RagRetriever.from_pretrained(UpperCAmelCase_      )
           self.assertIsInstance(UpperCAmelCase_					, UpperCAmelCase_      )
           SCREAMING_SNAKE_CASE       :      List[Any]							=    np.array(
               [np.ones(self.retrieval_vector_size      ), -np.ones(self.retrieval_vector_size      )]					, dtype=np.floataa      )
           SCREAMING_SNAKE_CASE       :      List[str]							=    retriever.retrieve(UpperCAmelCase_					, n_docs=1      )
           self.assertTrue(out is not None      )


     def 				_A       (				self       :	Any      ):
       SCREAMING_SNAKE_CASE       :      Any							=    1
       SCREAMING_SNAKE_CASE       :      List[Any]							=    self.get_dummy_custom_hf_index_retriever(from_disk=UpperCAmelCase_      )
       SCREAMING_SNAKE_CASE       :      Any							=    np.array(
           [np.ones(self.retrieval_vector_size      ), -np.ones(self.retrieval_vector_size      )]					, dtype=np.floataa      )
       SCREAMING_SNAKE_CASE       :      Optional[int]							=    retriever.retrieve(UpperCAmelCase_					, n_docs=UpperCAmelCase_      )
       self.assertEqual(retrieved_doc_embeds.shape					, (2, n_docs, self.retrieval_vector_size)      )
       self.assertEqual(len(UpperCAmelCase_      )					, 2      )
       self.assertEqual(sorted(doc_dicts[0]      )					, ["embeddings", "id", "text", "title"]      )
       self.assertEqual(len(doc_dicts[0]["id"]      )					, UpperCAmelCase_      )
       self.assertEqual(doc_dicts[0]["id"][0]					, "1"      )  # max inner product is reached with second doc
       self.assertEqual(doc_dicts[1]["id"][0]					, "0"      )  # max inner product is reached with first doc
       self.assertListEqual(doc_ids.tolist()					, [[1], [0]]      )


     def 				_A       (				self       :	List[Any]      ):
       SCREAMING_SNAKE_CASE       :      Optional[Any]							=    self.get_dummy_custom_hf_index_retriever(from_disk=UpperCAmelCase_      )
       with tempfile.TemporaryDirectory() as tmp_dirname:
         retriever.save_pretrained(UpperCAmelCase_      )
         SCREAMING_SNAKE_CASE       :      Any							=    RagRetriever.from_pretrained(UpperCAmelCase_      )
         self.assertIsInstance(UpperCAmelCase_					, UpperCAmelCase_      )
         SCREAMING_SNAKE_CASE       :      Tuple							=    np.array(
             [np.ones(self.retrieval_vector_size      ), -np.ones(self.retrieval_vector_size      )]					, dtype=np.floataa      )
         SCREAMING_SNAKE_CASE       :      Any							=    retriever.retrieve(UpperCAmelCase_					, n_docs=1      )
         self.assertTrue(out is not None      )


     def 				_A       (				self       :	str      ):
       SCREAMING_SNAKE_CASE       :      Optional[Any]							=    1
       SCREAMING_SNAKE_CASE       :      Any							=    self.get_dummy_custom_hf_index_retriever(from_disk=UpperCAmelCase_      )
       SCREAMING_SNAKE_CASE       :      List[str]							=    np.array(
           [np.ones(self.retrieval_vector_size      ), -np.ones(self.retrieval_vector_size      )]					, dtype=np.floataa      )
       SCREAMING_SNAKE_CASE       :      Tuple							=    retriever.retrieve(UpperCAmelCase_					, n_docs=UpperCAmelCase_      )
       self.assertEqual(retrieved_doc_embeds.shape					, (2, n_docs, self.retrieval_vector_size)      )
       self.assertEqual(len(UpperCAmelCase_      )					, 2      )
       self.assertEqual(sorted(doc_dicts[0]      )					, ["embeddings", "id", "text", "title"]      )
       self.assertEqual(len(doc_dicts[0]["id"]      )					, UpperCAmelCase_      )
       self.assertEqual(doc_dicts[0]["id"][0]					, "1"      )  # max inner product is reached with second doc
       self.assertEqual(doc_dicts[1]["id"][0]					, "0"      )  # max inner product is reached with first doc
       self.assertListEqual(doc_ids.tolist()					, [[1], [0]]      )


     def 				_A       (				self       :	int      ):
       SCREAMING_SNAKE_CASE       :      List[Any]							=    self.get_dummy_custom_hf_index_retriever(from_disk=UpperCAmelCase_      )
       with tempfile.TemporaryDirectory() as tmp_dirname:
         retriever.save_pretrained(UpperCAmelCase_      )
         SCREAMING_SNAKE_CASE       :      str							=    RagRetriever.from_pretrained(UpperCAmelCase_      )
         self.assertIsInstance(UpperCAmelCase_					, UpperCAmelCase_      )
         SCREAMING_SNAKE_CASE       :      str							=    np.array(
             [np.ones(self.retrieval_vector_size      ), -np.ones(self.retrieval_vector_size      )]					, dtype=np.floataa      )
         SCREAMING_SNAKE_CASE       :      Dict							=    retriever.retrieve(UpperCAmelCase_					, n_docs=1      )
         self.assertTrue(out is not None      )


     def 				_A       (				self       :	Any      ):
       SCREAMING_SNAKE_CASE       :      int							=    1
       SCREAMING_SNAKE_CASE       :      List[str]							=    self.get_dummy_legacy_index_retriever()
       SCREAMING_SNAKE_CASE       :      Tuple							=    np.array(
           [np.ones(self.retrieval_vector_size      ), -np.ones(self.retrieval_vector_size      )]					, dtype=np.floataa      )
       SCREAMING_SNAKE_CASE       :      Any							=    retriever.retrieve(UpperCAmelCase_					, n_docs=UpperCAmelCase_      )
       self.assertEqual(retrieved_doc_embeds.shape					, (2, n_docs, self.retrieval_vector_size)      )
       self.assertEqual(len(UpperCAmelCase_      )					, 2      )
       self.assertEqual(sorted(doc_dicts[0]      )					, ["text", "title"]      )
       self.assertEqual(len(doc_dicts[0]["text"]      )					, UpperCAmelCase_      )
       self.assertEqual(doc_dicts[0]["text"][0]					, "bar"      )  # max inner product is reached with second doc
       self.assertEqual(doc_dicts[1]["text"][0]					, "foo"      )  # max inner product is reached with first doc
       self.assertListEqual(doc_ids.tolist()					, [[1], [0]]      )


     def 				_A       (				self       :	List[Any]      ):
       SCREAMING_SNAKE_CASE       :      Any							=    self.get_dummy_legacy_index_retriever()
       with tempfile.TemporaryDirectory() as tmp_dirname:
         retriever.save_pretrained(UpperCAmelCase_      )
         SCREAMING_SNAKE_CASE       :      Optional[Any]							=    RagRetriever.from_pretrained(UpperCAmelCase_      )
         self.assertIsInstance(UpperCAmelCase_					, UpperCAmelCase_      )
         SCREAMING_SNAKE_CASE       :      Any							=    np.array(
             [np.ones(self.retrieval_vector_size      ), -np.ones(self.retrieval_vector_size      )]					, dtype=np.floataa      )
         SCREAMING_SNAKE_CASE       :      Optional[int]							=    retriever.retrieve(UpperCAmelCase_					, n_docs=1      )
         self.assertTrue(out is not None      )


     @require_torch
     @require_tokenizers
     @require_sentencepiece
     def 				_A       (				self       :	Optional[Any]      ):
       import torch

       SCREAMING_SNAKE_CASE       :      str							=    1
       SCREAMING_SNAKE_CASE       :      List[str]							=    self.get_dummy_canonical_hf_index_retriever()
       SCREAMING_SNAKE_CASE       :      str							=    [[5, 7], [10, 11]]
       SCREAMING_SNAKE_CASE       :      Tuple							=    np.array(
           [np.ones(self.retrieval_vector_size      ), -np.ones(self.retrieval_vector_size      )]					, dtype=np.floataa      )
       SCREAMING_SNAKE_CASE       :      str							=    retriever(UpperCAmelCase_					, UpperCAmelCase_					, prefix=retriever.config.generator.prefix					, n_docs=UpperCAmelCase_      )
       SCREAMING_SNAKE_CASE       :      str							=    (
           out["context_input_ids"],
           out["context_attention_mask"],
           out["retrieved_doc_embeds"],
       )
       self.assertEqual(retrieved_doc_embeds.shape					, (2, n_docs, self.retrieval_vector_size)      )
       self.assertIsInstance(UpperCAmelCase_					, UpperCAmelCase_      )
       self.assertIsInstance(UpperCAmelCase_					, UpperCAmelCase_      )
       self.assertIsInstance(UpperCAmelCase_					, np.ndarray      )

       SCREAMING_SNAKE_CASE       :      Union[str, Any]							=    retriever(
           UpperCAmelCase_					, UpperCAmelCase_					, prefix=retriever.config.generator.prefix					, n_docs=UpperCAmelCase_					, return_tensors="pt"					, )
       SCREAMING_SNAKE_CASE       :      Union[str, Any]							=    (  # noqa: F841
           out["context_input_ids"],
           out["context_attention_mask"],
           out["retrieved_doc_embeds"],
           out["doc_ids"],
       )
       self.assertEqual(retrieved_doc_embeds.shape					, (2, n_docs, self.retrieval_vector_size)      )
       self.assertIsInstance(UpperCAmelCase_					, torch.Tensor      )
       self.assertIsInstance(UpperCAmelCase_					, torch.Tensor      )
       self.assertIsInstance(UpperCAmelCase_					, torch.Tensor      )






     @require_torch
     @require_tokenizers
     @require_sentencepiece
     def 				_A       (				self       :	List[Any]      ):
       SCREAMING_SNAKE_CASE       :      Optional[Any]							=    self.get_dpr_ctx_encoder_tokenizer()
       SCREAMING_SNAKE_CASE       :      Optional[int]							=    1
       SCREAMING_SNAKE_CASE       :      str							=    self.get_dummy_custom_hf_index_retriever(from_disk=UpperCAmelCase_      )
       retriever.set_ctx_encoder_tokenizer(UpperCAmelCase_      )

       SCREAMING_SNAKE_CASE       :      List[Any]							=    [[5, 7], [10, 11]]
       SCREAMING_SNAKE_CASE       :      str							=    np.array(
           [np.ones(self.retrieval_vector_size      ), -np.ones(self.retrieval_vector_size      )]					, dtype=np.floataa      )
       SCREAMING_SNAKE_CASE       :      Optional[int]							=    retriever(UpperCAmelCase_					, UpperCAmelCase_					, prefix=retriever.config.generator.prefix					, n_docs=UpperCAmelCase_      )

       self.assertEqual(
           len(UpperCAmelCase_      )					, 6      )  # check whether the retriever output consist of 6 attributes including tokenized docs
       self.assertEqual(
           all(k in out for k in ("tokenized_doc_ids", "tokenized_doc_attention_mask")      )					, UpperCAmelCase_      )  # check for doc token related keys in dictionary.














import gc
import unittest

import numpy as np
import torch

from diffusers import StableDiffusionKDiffusionPipeline
from diffusers.utils import slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu


enable_full_determinism()


@slow
@require_torch_gpu
class  SCREAMING_SNAKE_CASE			(			unittest.TestCase					):
     '''simple docstring'''


     def 				_A       (				self       :	List[Any]      ):
       # clean up the VRAM after each test
       super().tearDown()
       gc.collect()
       torch.cuda.empty_cache()


     def 				_A       (				self       :	str      ):
       SCREAMING_SNAKE_CASE       :      List[str]							=    StableDiffusionKDiffusionPipeline.from_pretrained("CompVis/stable-diffusion-v1-4"      )
       SCREAMING_SNAKE_CASE       :      int							=    sd_pipe.to(UpperCAmelCase_      )
       sd_pipe.set_progress_bar_config(disable=UpperCAmelCase_      )

       sd_pipe.set_scheduler("sample_euler"      )

       SCREAMING_SNAKE_CASE       :      str							=    "A painting of a squirrel eating a burger"
       SCREAMING_SNAKE_CASE       :      Union[str, Any]							=    torch.manual_seed(0      )
       SCREAMING_SNAKE_CASE       :      Any							=    sd_pipe([prompt]					, generator=UpperCAmelCase_					, guidance_scale=9.0					, num_inference_steps=20					, output_type="np"      )

       SCREAMING_SNAKE_CASE       :      List[str]							=    output.images

       SCREAMING_SNAKE_CASE       :      str							=    image[0, -3:, -3:, -1]

       assert image.shape == (1, 512, 512, 3)
       SCREAMING_SNAKE_CASE       :      int							=    np.array([0.0_447, 0.0_492, 0.0_468, 0.0_408, 0.0_383, 0.0_408, 0.0_354, 0.0_380, 0.0_339]      )

       assert np.abs(image_slice.flatten() - expected_slice      ).max() < 1E-2


     def 				_A       (				self       :	List[Any]      ):
       SCREAMING_SNAKE_CASE       :      str							=    StableDiffusionKDiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-2-1-base"      )
       SCREAMING_SNAKE_CASE       :      str							=    sd_pipe.to(UpperCAmelCase_      )
       sd_pipe.set_progress_bar_config(disable=UpperCAmelCase_      )

       sd_pipe.set_scheduler("sample_euler"      )

       SCREAMING_SNAKE_CASE       :      List[str]							=    "A painting of a squirrel eating a burger"
       SCREAMING_SNAKE_CASE       :      Dict							=    torch.manual_seed(0      )
       SCREAMING_SNAKE_CASE       :      Optional[int]							=    sd_pipe([prompt]					, generator=UpperCAmelCase_					, guidance_scale=9.0					, num_inference_steps=20					, output_type="np"      )

       SCREAMING_SNAKE_CASE       :      Optional[int]							=    output.images

       SCREAMING_SNAKE_CASE       :      Dict							=    image[0, -3:, -3:, -1]

       assert image.shape == (1, 512, 512, 3)
       SCREAMING_SNAKE_CASE       :      int							=    np.array([0.1_237, 0.1_320, 0.1_438, 0.1_359, 0.1_390, 0.1_132, 0.1_277, 0.1_175, 0.1_112]      )

       assert np.abs(image_slice.flatten() - expected_slice      ).max() < 5E-1






     def 				_A       (				self       :	Optional[int]      ):
       SCREAMING_SNAKE_CASE       :      List[str]							=    StableDiffusionKDiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-2-1-base"      )
       SCREAMING_SNAKE_CASE       :      Any							=    sd_pipe.to(UpperCAmelCase_      )
       sd_pipe.set_progress_bar_config(disable=UpperCAmelCase_      )

       sd_pipe.set_scheduler("sample_dpmpp_2m"      )

       SCREAMING_SNAKE_CASE       :      List[str]							=    "A painting of a squirrel eating a burger"
       SCREAMING_SNAKE_CASE       :      Any							=    torch.manual_seed(0      )
       SCREAMING_SNAKE_CASE       :      Union[str, Any]							=    sd_pipe(
           [prompt]					, generator=UpperCAmelCase_					, guidance_scale=7.5					, num_inference_steps=15					, output_type="np"					, use_karras_sigmas=UpperCAmelCase_					, )

       SCREAMING_SNAKE_CASE       :      Optional[Any]							=    output.images

       SCREAMING_SNAKE_CASE       :      List[Any]							=    image[0, -3:, -3:, -1]

       assert image.shape == (1, 512, 512, 3)
       SCREAMING_SNAKE_CASE       :      Union[str, Any]							=    np.array(
           [0.11_381_689, 0.12_112_921, 0.1_389_457, 0.12_549_606, 0.1_244_964, 0.10_831_517, 0.11_562_866, 0.10_867_816, 0.10_499_048]      )

       assert np.abs(image_slice.flatten() - expected_slice      ).max() < 1E-2






