'''simple docstring'''
import enum
import warnings

from ..tokenization_utils import TruncationStrategy
from ..utils import add_end_docstrings, is_tf_available, is_torch_available, logging
from .base import PIPELINE_INIT_ARGS, Pipeline


if is_tf_available():
      import tensorflow as tf

      from ..models.auto.modeling_tf_auto import TF_MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING

if is_torch_available():
      from ..models.auto.modeling_auto import MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING

lowercase__  :      int        =					logging.get_logger(__name__)







class lowerCamelCase							(	enum.Enum					):




  '''simple docstring'''



  lowerCAmelCase__										=	0
  lowerCAmelCase__										=	1







@add_end_docstrings(lowerCamelCase					)
class lowerCamelCase							(	lowerCamelCase					):




  '''simple docstring'''



  lowerCAmelCase__										=	'''generated'''






  def __init__(      self			:       str  ,							*UpperCAmelCase__			:       Union[str, Any]  ,							**UpperCAmelCase__			:       Optional[int]  )				->Any:
        super().__init__(*UpperCAmelCase__  ,							**UpperCAmelCase__  )

        self.check_model_type(
            TF_MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING
            if self.framework == '''tf'''
            else MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING  )






  def  lowerCAmelCase__ (      self			:       Optional[Any]  ,							UpperCAmelCase__			:       List[Any]=None  ,							UpperCAmelCase__			:       Tuple=None  ,							UpperCAmelCase__			:       int=None  ,							UpperCAmelCase__			:       Dict=None  ,							UpperCAmelCase__			:       List[str]=None  ,							UpperCAmelCase__			:       str=None  ,							**UpperCAmelCase__			:       List[str]  ,							)				->int:
        UpperCAmelCase_            =  {}
        if truncation is not None:
              UpperCAmelCase_            =  truncation

        UpperCAmelCase_            =  generate_kwargs

        UpperCAmelCase_            =  {}
        if return_tensors is not None and return_type is None:
              UpperCAmelCase_            =  ReturnType.TENSORS if return_tensors else ReturnType.TEXT
        if return_type is not None:
              UpperCAmelCase_            =  return_type

        if clean_up_tokenization_spaces is not None:
              UpperCAmelCase_            =  clean_up_tokenization_spaces

        if stop_sequence is not None:
              UpperCAmelCase_            =  self.tokenizer.encode(UpperCAmelCase__  ,							add_special_tokens=UpperCAmelCase__  )
              if len(UpperCAmelCase__  ) > 1:
                    warnings.warn(
                        '''Stopping on a multiple token sequence is not yet supported on transformers. The first token of'''
                        ''' the stop sequence will be used as the stop sequence string in the interim.'''  )
              UpperCAmelCase_            =  stop_sequence_ids[0]

        return preprocess_params, forward_params, postprocess_params






  def  lowerCAmelCase__ (      self			:       List[Any]  ,							UpperCAmelCase__			:       int  ,							UpperCAmelCase__			:       int  ,							UpperCAmelCase__			:       int  )				->Union[str, Any]:
        return True






  def  lowerCAmelCase__ (      self			:       Optional[Any]  ,							*UpperCAmelCase__			:       Dict  ,							UpperCAmelCase__			:       Tuple  )				->List[Any]:
        UpperCAmelCase_            =  self.model.config.prefix if self.model.config.prefix is not None else ''''''
        if isinstance(args[0]  ,							UpperCAmelCase__  ):
              if self.tokenizer.pad_token_id is None:
                    raise ValueError('''Please make sure that the tokenizer has a pad_token_id when using a batch input'''  )
              UpperCAmelCase_            =  ([prefix + arg for arg in args[0]],)
              UpperCAmelCase_            =  True

        elif isinstance(args[0]  ,							UpperCAmelCase__  ):
              UpperCAmelCase_            =  (prefix + args[0],)
              UpperCAmelCase_            =  False
        else:
              raise ValueError(
                  f""" `args[0]`: {args[0]} have the wrong format. The should be either of type `str` or type `list`"""  )
        UpperCAmelCase_            =  self.tokenizer(*UpperCAmelCase__  ,							padding=UpperCAmelCase__  ,							truncation=UpperCAmelCase__  ,							return_tensors=self.framework  )
        # This is produced by tokenizers but is an invalid generate kwargs
        if "token_type_ids" in inputs:
              del inputs["token_type_ids"]
        return inputs






  def __call__(      self			:       Optional[int]  ,							*UpperCAmelCase__			:       Dict  ,							**UpperCAmelCase__			:       Optional[Any]  )				->Dict:
        UpperCAmelCase_            =  super().__call__(*UpperCAmelCase__  ,							**UpperCAmelCase__  )
        if (
            isinstance(args[0]  ,							UpperCAmelCase__  )
            and all(isinstance(UpperCAmelCase__  ,							UpperCAmelCase__  ) for el in args[0]  )
            and all(len(UpperCAmelCase__  ) == 1 for res in result  )
        ):
              return [res[0] for res in result]
        return result






  def  lowerCAmelCase__ (      self			:       int  ,							UpperCAmelCase__			:       Optional[Any]  ,							UpperCAmelCase__			:       Union[str, Any]=TruncationStrategy.DO_NOT_TRUNCATE  ,							**UpperCAmelCase__			:       List[str]  )				->Optional[int]:
        UpperCAmelCase_            =  self._parse_and_tokenize(UpperCAmelCase__  ,							truncation=UpperCAmelCase__  ,							**UpperCAmelCase__  )
        return inputs






  def  lowerCAmelCase__ (      self			:       Union[str, Any]  ,							UpperCAmelCase__			:       Dict  ,							**UpperCAmelCase__			:       Any  )				->Union[str, Any]:
        if self.framework == "pt":
              UpperCAmelCase_      ,		UpperCAmelCase_            =  model_inputs['''input_ids'''].shape
        elif self.framework == "tf":
              UpperCAmelCase_      ,		UpperCAmelCase_            =  tf.shape(model_inputs['''input_ids''']  ).numpy()

        UpperCAmelCase_            =  generate_kwargs.get('''min_length'''  ,							self.model.config.min_length  )
        UpperCAmelCase_            =  generate_kwargs.get('''max_length'''  ,							self.model.config.max_length  )
        self.check_inputs(UpperCAmelCase__  ,							generate_kwargs['''min_length''']  ,							generate_kwargs['''max_length''']  )
        UpperCAmelCase_            =  self.model.generate(**UpperCAmelCase__  ,							**UpperCAmelCase__  )
        UpperCAmelCase_            =  output_ids.shape[0]
        if self.framework == "pt":
              UpperCAmelCase_            =  output_ids.reshape(UpperCAmelCase__  ,							out_b // in_b  ,							*output_ids.shape[1:]  )
        elif self.framework == "tf":
              UpperCAmelCase_            =  tf.reshape(UpperCAmelCase__  ,							(in_b, out_b // in_b, *output_ids.shape[1:])  )
        return {"output_ids": output_ids}






  def  lowerCAmelCase__ (      self			:       int  ,							UpperCAmelCase__			:       Optional[Any]  ,							UpperCAmelCase__			:       Optional[int]=ReturnType.TEXT  ,							UpperCAmelCase__			:       Dict=False  )				->Any:
        UpperCAmelCase_            =  []
        for output_ids in model_outputs["output_ids"][0]:
              if return_type == ReturnType.TENSORS:
                    UpperCAmelCase_            =  {f"""{self.return_name}_token_ids""": output_ids}
              elif return_type == ReturnType.TEXT:
                    UpperCAmelCase_            =  {
                        f"""{self.return_name}_text""": self.tokenizer.decode(
                            UpperCAmelCase__  ,							skip_special_tokens=UpperCAmelCase__  ,							clean_up_tokenization_spaces=UpperCAmelCase__  ,							)
                    }
              records.append(UpperCAmelCase__  )
        return records







@add_end_docstrings(lowerCamelCase					)
class lowerCamelCase							(	lowerCamelCase					):




  '''simple docstring'''



  lowerCAmelCase__										=	'''summary'''






  def __call__(      self			:       Union[str, Any]  ,							*UpperCAmelCase__			:       Optional[int]  ,							**UpperCAmelCase__			:       Union[str, Any]  )				->Any:
        return super().__call__(*UpperCAmelCase__  ,							**UpperCAmelCase__  )






  def  lowerCAmelCase__ (      self			:       str  ,							UpperCAmelCase__			:       int  ,							UpperCAmelCase__			:       int  ,							UpperCAmelCase__			:       int  )				->bool:
        if max_length < min_length:
              logger.warning(f"""Your min_length={min_length} must be inferior than your max_length={max_length}."""  )

        if input_length < max_length:
              logger.warning(
                  f"""Your max_length is set to {max_length}, but your input_length is only {input_length}. Since this is """
                  '''a summarization task, where outputs shorter than the input are typically wanted, you might '''
                  f"""consider decreasing max_length manually, e.g. summarizer('...', max_length={input_length//2})"""  )







@add_end_docstrings(lowerCamelCase					)
class lowerCamelCase							(	lowerCamelCase					):




  '''simple docstring'''



  lowerCAmelCase__										=	'''translation'''






  def  lowerCAmelCase__ (      self			:       List[Any]  ,							UpperCAmelCase__			:       int  ,							UpperCAmelCase__			:       int  ,							UpperCAmelCase__			:       int  )				->List[str]:
        if input_length > 0.9 * max_length:
              logger.warning(
                  f"""Your input_length: {input_length} is bigger than 0.9 * max_length: {max_length}. You might consider """
                  '''increasing your max_length manually, e.g. translator(\'...\', max_length=400)'''  )
        return True






  def  lowerCAmelCase__ (      self			:       int  ,							*UpperCAmelCase__			:       Optional[int]  ,							UpperCAmelCase__			:       Any=TruncationStrategy.DO_NOT_TRUNCATE  ,							UpperCAmelCase__			:       Optional[Any]=None  ,							UpperCAmelCase__			:       List[Any]=None  )				->Optional[Any]:
        if getattr(self.tokenizer  ,							'''_build_translation_inputs'''  ,							UpperCAmelCase__  ):
              return self.tokenizer._build_translation_inputs(
                  *UpperCAmelCase__  ,							return_tensors=self.framework  ,							truncation=UpperCAmelCase__  ,							src_lang=UpperCAmelCase__  ,							tgt_lang=UpperCAmelCase__  )
        else:
              return super()._parse_and_tokenize(*UpperCAmelCase__  ,							truncation=UpperCAmelCase__  )






  def  lowerCAmelCase__ (      self			:       Tuple  ,							UpperCAmelCase__			:       Dict=None  ,							UpperCAmelCase__			:       Optional[int]=None  ,							**UpperCAmelCase__			:       int  )				->Dict:
        UpperCAmelCase_      ,		UpperCAmelCase_      ,		UpperCAmelCase_            =  super()._sanitize_parameters(**UpperCAmelCase__  )
        if src_lang is not None:
              UpperCAmelCase_            =  src_lang
        if tgt_lang is not None:
              UpperCAmelCase_            =  tgt_lang
        if src_lang is None and tgt_lang is None:
              # Backward compatibility, direct arguments use is preferred.
              UpperCAmelCase_            =  kwargs.get('''task'''  ,							self.task  )
              UpperCAmelCase_            =  task.split('''_'''  )
              if task and len(UpperCAmelCase__  ) == 4:
                    # translation, XX, to YY
                    UpperCAmelCase_            =  items[1]
                    UpperCAmelCase_            =  items[3]
        return preprocess_params, forward_params, postprocess_params






  def __call__(      self			:       List[str]  ,							*UpperCAmelCase__			:       List[str]  ,							**UpperCAmelCase__			:       Any  )				->Tuple:
        return super().__call__(*UpperCAmelCase__  ,							**UpperCAmelCase__  )






'''simple docstring'''
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 lowerCamelCase							(	unittest.TestCase					):




  '''simple docstring'''



  def  lowerCAmelCase__ (      self			:       Any  )				->Optional[int]:
        # clean up the VRAM after each test
        super().tearDown()
        gc.collect()
        torch.cuda.empty_cache()






  def  lowerCAmelCase__ (      self			:       str  )				->List[str]:
        UpperCAmelCase_            =  StableDiffusionKDiffusionPipeline.from_pretrained('''CompVis/stable-diffusion-v1-4'''  )
        UpperCAmelCase_            =  sd_pipe.to(UpperCAmelCase__  )
        sd_pipe.set_progress_bar_config(disable=UpperCAmelCase__  )

        sd_pipe.set_scheduler('''sample_euler'''  )

        UpperCAmelCase_            =  '''A painting of a squirrel eating a burger'''
        UpperCAmelCase_            =  torch.manual_seed(0  )
        UpperCAmelCase_            =  sd_pipe([prompt]  ,							generator=UpperCAmelCase__  ,							guidance_scale=9.0  ,							num_inference_steps=20  ,							output_type='''np'''  )

        UpperCAmelCase_            =  output.images

        UpperCAmelCase_            =  image[0, -3:, -3:, -1]

        assert image.shape == (1, 512, 512, 3)
        UpperCAmelCase_            =  np.array([0.0447, 0.0492, 0.0468, 0.0408, 0.0383, 0.0408, 0.0354, 0.0380, 0.0339]  )

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






  def  lowerCAmelCase__ (      self			:       List[str]  )				->int:
        UpperCAmelCase_            =  StableDiffusionKDiffusionPipeline.from_pretrained('''stabilityai/stable-diffusion-2-1-base'''  )
        UpperCAmelCase_            =  sd_pipe.to(UpperCAmelCase__  )
        sd_pipe.set_progress_bar_config(disable=UpperCAmelCase__  )

        sd_pipe.set_scheduler('''sample_euler'''  )

        UpperCAmelCase_            =  '''A painting of a squirrel eating a burger'''
        UpperCAmelCase_            =  torch.manual_seed(0  )
        UpperCAmelCase_            =  sd_pipe([prompt]  ,							generator=UpperCAmelCase__  ,							guidance_scale=9.0  ,							num_inference_steps=20  ,							output_type='''np'''  )

        UpperCAmelCase_            =  output.images

        UpperCAmelCase_            =  image[0, -3:, -3:, -1]

        assert image.shape == (1, 512, 512, 3)
        UpperCAmelCase_            =  np.array([0.1237, 0.1320, 0.1438, 0.1359, 0.1390, 0.1132, 0.1277, 0.1175, 0.1112]  )

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






  def  lowerCAmelCase__ (      self			:       List[Any]  )				->Optional[int]:
        UpperCAmelCase_            =  StableDiffusionKDiffusionPipeline.from_pretrained('''stabilityai/stable-diffusion-2-1-base'''  )
        UpperCAmelCase_            =  sd_pipe.to(UpperCAmelCase__  )
        sd_pipe.set_progress_bar_config(disable=UpperCAmelCase__  )

        sd_pipe.set_scheduler('''sample_dpmpp_2m'''  )

        UpperCAmelCase_            =  '''A painting of a squirrel eating a burger'''
        UpperCAmelCase_            =  torch.manual_seed(0  )
        UpperCAmelCase_            =  sd_pipe(
            [prompt]  ,							generator=UpperCAmelCase__  ,							guidance_scale=7.5  ,							num_inference_steps=15  ,							output_type='''np'''  ,							use_karras_sigmas=UpperCAmelCase__  ,							)

        UpperCAmelCase_            =  output.images

        UpperCAmelCase_            =  image[0, -3:, -3:, -1]

        assert image.shape == (1, 512, 512, 3)
        UpperCAmelCase_            =  np.array(
            [0.1138_1689, 0.1211_2921, 0.138_9457, 0.1254_9606, 0.124_4964, 0.1083_1517, 0.1156_2866, 0.1086_7816, 0.1049_9048]  )

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





