




'''simple docstring'''


from typing import List, Optional, Union

import torch

from ...models import UNetaDConditionModel, VQModel
from ...pipelines import DiffusionPipeline
from ...pipelines.pipeline_utils import ImagePipelineOutput
from ...schedulers import DDPMScheduler
from ...utils import (
    is_accelerate_available,
    is_accelerate_version,
    logging,
    randn_tensor,
    replace_example_docstring,
)


UpperCAmelCase_					=   logging.get_logger(__name__)  # pylint: disable=invalid-name

UpperCAmelCase_					=   "\n    Examples:\n        ```py\n        >>> import torch\n        >>> import numpy as np\n\n        >>> from diffusers import KandinskyV22PriorPipeline, KandinskyV22ControlnetPipeline\n        >>> from transformers import pipeline\n        >>> from diffusers.utils import load_image\n\n\n        >>> def make_hint(image, depth_estimator):\n        ...     image = depth_estimator(image)[\"depth\"]\n        ...     image = np.array(image)\n        ...     image = image[:, :, None]\n        ...     image = np.concatenate([image, image, image], axis=2)\n        ...     detected_map = torch.from_numpy(image).float() / 255.0\n        ...     hint = detected_map.permute(2, 0, 1)\n        ...     return hint\n\n\n        >>> depth_estimator = pipeline(\"depth-estimation\")\n\n        >>> pipe_prior = KandinskyV22PriorPipeline.from_pretrained(\n        ...     \"kandinsky-community/kandinsky-2-2-prior\", torch_dtype=torch.float16\n        ... )\n        >>> pipe_prior = pipe_prior.to(\"cuda\")\n\n        >>> pipe = KandinskyV22ControlnetPipeline.from_pretrained(\n        ...     \"kandinsky-community/kandinsky-2-2-controlnet-depth\", torch_dtype=torch.float16\n        ... )\n        >>> pipe = pipe.to(\"cuda\")\n\n\n        >>> img = load_image(\n        ...     \"https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main\"\n        ...     \"/kandinsky/cat.png\"\n        ... ).resize((768, 768))\n\n        >>> hint = make_hint(img, depth_estimator).unsqueeze(0).half().to(\"cuda\")\n\n        >>> prompt = \"A robot, 4k photo\"\n        >>> negative_prior_prompt = \"lowres, text, error, cropped, worst quality, low quality, jpeg artifacts, ugly, duplicate, morbid, mutilated, out of frame, extra fingers, mutated hands, poorly drawn hands, poorly drawn face, mutation, deformed, blurry, dehydrated, bad anatomy, bad proportions, extra limbs, cloned face, disfigured, gross proportions, malformed limbs, missing arms, missing legs, extra arms, extra legs, fused fingers, too many fingers, long neck, username, watermark, signature\"\n\n        >>> generator = torch.Generator(device=\"cuda\").manual_seed(43)\n\n        >>> image_emb, zero_image_emb = pipe_prior(\n        ...     prompt=prompt, negative_prompt=negative_prior_prompt, generator=generator\n        ... ).to_tuple()\n\n        >>> images = pipe(\n        ...     image_embeds=image_emb,\n        ...     negative_image_embeds=zero_image_emb,\n        ...     hint=hint,\n        ...     num_inference_steps=50,\n        ...     generator=generator,\n        ...     height=768,\n        ...     width=768,\n        ... ).images\n\n        >>> images[0].save(\"robot_cat.png\")\n        ```\n"


def     SCREAMING_SNAKE_CASE							(   a_						:				Any					,       a_						:				Any					,       a_						:				List[str]=8      ):
  __a								=       height // scale_factor**2
  if height % scale_factor**2 != 0:
    new_height += 1
  __a								=       width // scale_factor**2
  if width % scale_factor**2 != 0:
    new_width += 1
  return new_height * scale_factor, new_width * scale_factor





class  __lowercase    (	__magic_name__			):







   def __init__(						self   ,				UpperCamelCase   ,				UpperCamelCase   ,				UpperCamelCase   ,				)					->     Optional[int]:
     super().__init__()

     self.register_modules(
         unet=UpperCamelCase   ,				scheduler=UpperCamelCase   ,				movq=UpperCamelCase   ,				)
     __a								=       2 ** (len(self.movq.config.block_out_channels   ) - 1)







   def 		UpperCamelCase__      (						self   ,				UpperCamelCase   ,				UpperCamelCase   ,				UpperCamelCase   ,				UpperCamelCase   ,				UpperCamelCase   ,				UpperCamelCase   )					->     Optional[Any]:
     if latents is None:
       __a								=       randn_tensor(UpperCamelCase   ,				generator=UpperCamelCase   ,				device=UpperCamelCase   ,				dtype=UpperCamelCase   )
     else:
       if latents.shape != shape:
         raise ValueError(f"Unexpected latents shape, got {latents.shape}, expected {shape}"   )
       __a								=       latents.to(UpperCamelCase   )

     __a								=       latents * scheduler.init_noise_sigma
     return latents







   def 		UpperCamelCase__      (						self   ,				UpperCamelCase=0   )					->     str:
     if is_accelerate_available():
       from accelerate import cpu_offload
     else:
       raise ImportError('Please install accelerate via `pip install accelerate`'   )

     __a								=       torch.device(f"cuda:{gpu_id}"   )

     __a								=       [
         self.unet,
         self.movq,
     ]
     for cpu_offloaded_model in models:
       if cpu_offloaded_model is not None:
         cpu_offload(UpperCamelCase   ,				UpperCamelCase   )







   def 		UpperCamelCase__      (						self   ,				UpperCamelCase=0   )					->     Union[str, Any]:
     if is_accelerate_available() and is_accelerate_version('>='   ,				'0.17.0.dev0'   ):
       from accelerate import cpu_offload_with_hook
     else:
       raise ImportError('`enable_model_cpu_offload` requires `accelerate v0.17.0` or higher.'   )

     __a								=       torch.device(f"cuda:{gpu_id}"   )

     if self.device.type != "cpu":
       self.to('cpu'   ,				silence_dtype_warnings=UpperCamelCase   )
       torch.cuda.empty_cache()  # otherwise we don't see the memory savings (but they probably exist)

     __a								=       None
     for cpu_offloaded_model in [self.unet, self.movq]:
       __a						,							__a								=       cpu_offload_with_hook(UpperCamelCase   ,				UpperCamelCase   ,				prev_module_hook=UpperCamelCase   )

     # We'll offload the last model manually.
     __a								=       hook







   @property
   # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline._execution_device
   def 		UpperCamelCase__      (						self   )					->     Optional[int]:
     if not hasattr(self.unet   ,				'_hf_hook'   ):
       return self.device
     for module in self.unet.modules():
       if (
           hasattr(UpperCamelCase   ,				'_hf_hook'   )
           and hasattr(module._hf_hook   ,				'execution_device'   )
           and module._hf_hook.execution_device is not None
       ):
         return torch.device(module._hf_hook.execution_device   )
     return self.device







   @torch.no_grad()
   @replace_example_docstring(UpperCamelCase   )
   def __call__(						self   ,				UpperCamelCase   ,				UpperCamelCase   ,				UpperCamelCase   ,				UpperCamelCase = 512   ,				UpperCamelCase = 512   ,				UpperCamelCase = 100   ,				UpperCamelCase = 4.0   ,				UpperCamelCase = 1   ,				UpperCamelCase = None   ,				UpperCamelCase = None   ,				UpperCamelCase = "pil"   ,				UpperCamelCase = True   ,				)					->     Any:
     __a								=       self._execution_device

     __a								=       guidance_scale > 1.0

     if isinstance(UpperCamelCase   ,				UpperCamelCase   ):
       __a								=       torch.cat(UpperCamelCase   ,				dim=0   )
     if isinstance(UpperCamelCase   ,				UpperCamelCase   ):
       __a								=       torch.cat(UpperCamelCase   ,				dim=0   )
     if isinstance(UpperCamelCase   ,				UpperCamelCase   ):
       __a								=       torch.cat(UpperCamelCase   ,				dim=0   )

     __a								=       image_embeds.shape[0] * num_images_per_prompt

     if do_classifier_free_guidance:
       __a								=       image_embeds.repeat_interleave(UpperCamelCase   ,				dim=0   )
       __a								=       negative_image_embeds.repeat_interleave(UpperCamelCase   ,				dim=0   )
       __a								=       hint.repeat_interleave(UpperCamelCase   ,				dim=0   )

     __a								=       torch.cat([negative_image_embeds, image_embeds]   ,				dim=0   ).to(dtype=self.unet.dtype   ,				device=UpperCamelCase   )
     __a								=       torch.cat([hint, hint]   ,				dim=0   ).to(dtype=self.unet.dtype   ,				device=UpperCamelCase   )
     self.scheduler.set_timesteps(UpperCamelCase   ,				device=UpperCamelCase   )
     __a								=       self.scheduler.timesteps

     __a								=       self.movq.config.latent_channels

     __a						,							__a								=       downscale_height_and_width(UpperCamelCase   ,				UpperCamelCase   ,				self.movq_scale_factor   )

     # create initial latent
     __a								=       self.prepare_latents(
         (batch_size, num_channels_latents, height, width)   ,				image_embeds.dtype   ,				UpperCamelCase   ,				UpperCamelCase   ,				UpperCamelCase   ,				self.scheduler   ,				)

     for i, t in enumerate(self.progress_bar(UpperCamelCase   )   ):
       # expand the latents if we are doing classifier free guidance
       __a								=       torch.cat([latents] * 2   ) if do_classifier_free_guidance else latents

       __a								=       {'image_embeds': image_embeds, 'hint': hint}
       __a								=       self.unet(
           sample=UpperCamelCase   ,				timestep=UpperCamelCase   ,				encoder_hidden_states=UpperCamelCase   ,				added_cond_kwargs=UpperCamelCase   ,				return_dict=UpperCamelCase   ,				)[0]

       if do_classifier_free_guidance:
         __a						,							__a								=       noise_pred.split(latents.shape[1]   ,				dim=1   )
         __a						,							__a								=       noise_pred.chunk(2   )
         __a						,							__a								=       variance_pred.chunk(2   )
         __a								=       noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)
         __a								=       torch.cat([noise_pred, variance_pred_text]   ,				dim=1   )

       if not (
           hasattr(self.scheduler.config   ,				'variance_type'   )
           and self.scheduler.config.variance_type in ["learned", "learned_range"]
       ):
         __a						,							__a								=       noise_pred.split(latents.shape[1]   ,				dim=1   )

       # compute the previous noisy sample x_t -> x_t-1
       __a								=       self.scheduler.step(
           UpperCamelCase   ,				UpperCamelCase   ,				UpperCamelCase   ,				generator=UpperCamelCase   ,				)[0]
     # post-processing
     __a								=       self.movq.decode(UpperCamelCase   ,				force_not_quantize=UpperCamelCase   )['sample']

     if output_type not in ["pt", "np", "pil"]:
       raise ValueError(f"Only the output types `pt`, `pil` and `np` are supported not output_type={output_type}"   )

     if output_type in ["np", "pil"]:
       __a								=       image * 0.5 + 0.5
       __a								=       image.clamp(0   ,				1   )
       __a								=       image.cpu().permute(0   ,				2   ,				3   ,				1   ).float().numpy()

     if output_type == "pil":
       __a								=       self.numpy_to_pil(UpperCamelCase   )

     if not return_dict:
       return (image,)

     return ImagePipelineOutput(images=UpperCamelCase   )









'''simple docstring'''


import itertools
import os
import random
import tempfile
import unittest

import numpy as np
from datasets import load_dataset

from transformers import is_speech_available
from transformers.testing_utils import check_json_file_has_correct_format, require_torch, require_torchaudio
from transformers.utils.import_utils import is_torch_available

from ...test_sequence_feature_extraction_common import SequenceFeatureExtractionTestMixin


if is_speech_available():
  from transformers import WhisperFeatureExtractor

if is_torch_available():
  import torch

UpperCAmelCase_					=   random.Random()


def     SCREAMING_SNAKE_CASE							(   a_						:				Optional[Any]					,       a_						:				Any=1.0					,       a_						:				Optional[Any]=None					,       a_						:				Union[str, Any]=None      ):
  if rng is None:
    __a								=       global_rng

  __a								=       []
  for batch_idx in range(shape[0]      ):
    values.append([]      )
    for _ in range(shape[1]      ):
      values[-1].append(rng.random() * scale      )

  return values





@require_torch
@require_torchaudio
class  __lowercase    (	unittest.TestCase			):







   def __init__(						self   ,				UpperCamelCase   ,				UpperCamelCase=7   ,				UpperCamelCase=400   ,				UpperCamelCase=2000   ,				UpperCamelCase=10   ,				UpperCamelCase=160   ,				UpperCamelCase=8   ,				UpperCamelCase=0.0   ,				UpperCamelCase=4000   ,				UpperCamelCase=False   ,				UpperCamelCase=True   ,				)					->     Optional[Any]:
     __a								=       parent
     __a								=       batch_size
     __a								=       min_seq_length
     __a								=       max_seq_length
     __a								=       (self.max_seq_length - self.min_seq_length) // (self.batch_size - 1)
     __a								=       padding_value
     __a								=       sampling_rate
     __a								=       return_attention_mask
     __a								=       do_normalize
     __a								=       feature_size
     __a								=       chunk_length
     __a								=       hop_length







   def 		UpperCamelCase__      (						self   )					->     Tuple:
     return {
         "feature_size": self.feature_size,
         "hop_length": self.hop_length,
         "chunk_length": self.chunk_length,
         "padding_value": self.padding_value,
         "sampling_rate": self.sampling_rate,
         "return_attention_mask": self.return_attention_mask,
         "do_normalize": self.do_normalize,
     }







   def 		UpperCamelCase__      (						self   ,				UpperCamelCase=False   ,				UpperCamelCase=False   )					->     int:
     def _flatten(UpperCamelCase   ):
       return list(itertools.chain(*UpperCamelCase   )   )

     if equal_length:
       __a								=       [floats_list((self.max_seq_length, self.feature_size)   ) for _ in range(self.batch_size   )]
     else:
       # make sure that inputs increase in size
       __a								=       [
           floats_list((x, self.feature_size)   )
           for x in range(self.min_seq_length   ,				self.max_seq_length   ,				self.seq_length_diff   )
       ]
     if numpify:
       __a								=       [np.asarray(UpperCamelCase   ) for x in speech_inputs]
     return speech_inputs





@require_torch
@require_torchaudio
class  __lowercase    (	__magic_name__							,				unittest.TestCase			):
   _a    						=   WhisperFeatureExtractor if is_speech_available() else None







   def 		UpperCamelCase__      (						self   )					->     str:
     __a								=       WhisperFeatureExtractionTester(self   )







   def 		UpperCamelCase__      (						self   )					->     Union[str, Any]:
     __a								=       self.feature_extraction_class(**self.feat_extract_dict   )

     with tempfile.TemporaryDirectory() as tmpdirname:
       __a								=       feat_extract_first.save_pretrained(UpperCamelCase   )[0]
       check_json_file_has_correct_format(UpperCamelCase   )
       __a								=       self.feature_extraction_class.from_pretrained(UpperCamelCase   )

     __a								=       feat_extract_first.to_dict()
     __a								=       feat_extract_second.to_dict()
     __a								=       feat_extract_first.mel_filters
     __a								=       feat_extract_second.mel_filters
     self.assertTrue(np.allclose(UpperCamelCase   ,				UpperCamelCase   )   )
     self.assertEqual(UpperCamelCase   ,				UpperCamelCase   )







   def 		UpperCamelCase__      (						self   )					->     Dict:
     __a								=       self.feature_extraction_class(**self.feat_extract_dict   )

     with tempfile.TemporaryDirectory() as tmpdirname:
       __a								=       os.path.join(UpperCamelCase   ,				'feat_extract.json'   )
       feat_extract_first.to_json_file(UpperCamelCase   )
       __a								=       self.feature_extraction_class.from_json_file(UpperCamelCase   )

     __a								=       feat_extract_first.to_dict()
     __a								=       feat_extract_second.to_dict()
     __a								=       feat_extract_first.mel_filters
     __a								=       feat_extract_second.mel_filters
     self.assertTrue(np.allclose(UpperCamelCase   ,				UpperCamelCase   )   )
     self.assertEqual(UpperCamelCase   ,				UpperCamelCase   )







   def 		UpperCamelCase__      (						self   )					->     Optional[Any]:
     # Tests that all call wrap to encode_plus and batch_encode_plus
     __a								=       self.feature_extraction_class(**self.feat_extract_tester.prepare_feat_extract_dict()   )
     # create three inputs of length 800, 1000, and 1200
     __a								=       [floats_list((1, x)   )[0] for x in range(800   ,				1400   ,				200   )]
     __a								=       [np.asarray(UpperCamelCase   ) for speech_input in speech_inputs]

     # Test feature size
     __a								=       feature_extractor(UpperCamelCase   ,				padding='max_length'   ,				return_tensors='np'   ).input_features
     self.assertTrue(input_features.ndim == 3   )
     self.assertTrue(input_features.shape[-1] == feature_extractor.nb_max_frames   )
     self.assertTrue(input_features.shape[-2] == feature_extractor.feature_size   )

     # Test not batched input
     __a								=       feature_extractor(speech_inputs[0]   ,				return_tensors='np'   ).input_features
     __a								=       feature_extractor(np_speech_inputs[0]   ,				return_tensors='np'   ).input_features
     self.assertTrue(np.allclose(UpperCamelCase   ,				UpperCamelCase   ,				atol=1e-3   )   )

     # Test batched
     __a								=       feature_extractor(UpperCamelCase   ,				return_tensors='np'   ).input_features
     __a								=       feature_extractor(UpperCamelCase   ,				return_tensors='np'   ).input_features
     for enc_seq_a, enc_seq_a in zip(UpperCamelCase   ,				UpperCamelCase   ):
       self.assertTrue(np.allclose(UpperCamelCase   ,				UpperCamelCase   ,				atol=1e-3   )   )

     # Test 2-D numpy arrays are batched.
     __a								=       [floats_list((1, x)   )[0] for x in (800, 800, 800)]
     __a								=       np.asarray(UpperCamelCase   )
     __a								=       feature_extractor(UpperCamelCase   ,				return_tensors='np'   ).input_features
     __a								=       feature_extractor(UpperCamelCase   ,				return_tensors='np'   ).input_features
     for enc_seq_a, enc_seq_a in zip(UpperCamelCase   ,				UpperCamelCase   ):
       self.assertTrue(np.allclose(UpperCamelCase   ,				UpperCamelCase   ,				atol=1e-3   )   )

     # Test truncation required
     __a								=       [floats_list((1, x)   )[0] for x in range(200   ,				(feature_extractor.n_samples + 500)   ,				200   )]
     __a								=       [np.asarray(UpperCamelCase   ) for speech_input in speech_inputs]

     __a								=       [x[: feature_extractor.n_samples] for x in speech_inputs]
     __a								=       [np.asarray(UpperCamelCase   ) for speech_input in speech_inputs_truncated]

     __a								=       feature_extractor(UpperCamelCase   ,				return_tensors='np'   ).input_features
     __a								=       feature_extractor(UpperCamelCase   ,				return_tensors='np'   ).input_features
     for enc_seq_a, enc_seq_a in zip(UpperCamelCase   ,				UpperCamelCase   ):
       self.assertTrue(np.allclose(UpperCamelCase   ,				UpperCamelCase   ,				atol=1e-3   )   )







   def 		UpperCamelCase__      (						self   )					->     int:
     import torch

     __a								=       self.feature_extraction_class(**self.feat_extract_tester.prepare_feat_extract_dict()   )
     __a								=       np.random.rand(100   ,				32   ).astype(np.floataa   )
     __a								=       np_speech_inputs.tolist()

     for inputs in [py_speech_inputs, np_speech_inputs]:
       __a								=       feature_extractor.pad([{'input_features': inputs}]   ,				return_tensors='np'   )
       self.assertTrue(np_processed.input_features.dtype == np.floataa   )
       __a								=       feature_extractor.pad([{'input_features': inputs}]   ,				return_tensors='pt'   )
       self.assertTrue(pt_processed.input_features.dtype == torch.floataa   )







   def 		UpperCamelCase__      (						self   ,				UpperCamelCase   )					->     Tuple:
     __a								=       load_dataset('hf-internal-testing/librispeech_asr_dummy'   ,				'clean'   ,				split='validation'   )
     # automatic decoding with librispeech
     __a								=       ds.sort('id'   ).select(range(UpperCamelCase   )   )[:num_samples]['audio']

     return [x["array"] for x in speech_samples]







   def 		UpperCamelCase__      (						self   )					->     int:
     # fmt: off
     __a								=       torch.tensor(
         [
             0.1_193, -0.0_946, -0.1_098, -0.0_196, 0.0_225, -0.0_690, -0.1_736, 0.0_951,
             0.0_971, -0.0_817, -0.0_702, 0.0_162, 0.0_260, 0.0_017, -0.0_192, -0.1_678,
             0.0_709, -0.1_867, -0.0_655, -0.0_274, -0.0_234, -0.1_884, -0.0_516, -0.0_554,
             -0.0_274, -0.1_425, -0.1_423, 0.0_837, 0.0_377, -0.0_854
         ]   )
     # fmt: on

     __a								=       self._load_datasamples(1   )
     __a								=       WhisperFeatureExtractor()
     __a								=       feature_extractor(UpperCamelCase   ,				return_tensors='pt'   ).input_features
     self.assertEqual(input_features.shape   ,				(1, 80, 3000)   )
     self.assertTrue(torch.allclose(input_features[0, 0, :30]   ,				UpperCamelCase   ,				atol=1e-4   )   )







   def 		UpperCamelCase__      (						self   )					->     List[Any]:
     __a								=       self.feature_extraction_class(**self.feat_extract_tester.prepare_feat_extract_dict()   )
     __a								=       self._load_datasamples(1   )[0]
     __a								=       ((audio - audio.min()) / (audio.max() - audio.min())) * 6_5535  # Rescale to [0, 65535] to show issue
     __a								=       feat_extract.zero_mean_unit_var_norm([audio]   ,				attention_mask=UpperCamelCase   )[0]

     self.assertTrue(np.all(np.mean(UpperCamelCase   ) < 1e-3   )   )
     self.assertTrue(np.all(np.abs(np.var(UpperCamelCase   ) - 1   ) < 1e-3   )   )



