





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,
)


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

_snake_case           =			"""
    Examples:
        ```py
        >>> import torch
        >>> import numpy as np

        >>> from diffusers import KandinskyV22PriorPipeline, KandinskyV22ControlnetPipeline
        >>> from transformers import pipeline
        >>> from diffusers.utils import load_image


        >>> def make_hint(image, depth_estimator):
        ...     image = depth_estimator(image)[\"depth\"]
        ...     image = np.array(image)
        ...     image = image[:, :, None]
        ...     image = np.concatenate([image, image, image], axis=2)
        ...     detected_map = torch.from_numpy(image).float() / 255.0
        ...     hint = detected_map.permute(2, 0, 1)
        ...     return hint


        >>> depth_estimator = pipeline(\"depth-estimation\")

        >>> pipe_prior = KandinskyV22PriorPipeline.from_pretrained(
        ...     \"kandinsky-community/kandinsky-2-2-prior\", torch_dtype=torch.float16
        ... )
        >>> pipe_prior = pipe_prior.to(\"cuda\")

        >>> pipe = KandinskyV22ControlnetPipeline.from_pretrained(
        ...     \"kandinsky-community/kandinsky-2-2-controlnet-depth\", torch_dtype=torch.float16
        ... )
        >>> pipe = pipe.to(\"cuda\")


        >>> img = load_image(
        ...     \"https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main\"
        ...     \"/kandinsky/cat.png\"
        ... ).resize((768, 768))

        >>> hint = make_hint(img, depth_estimator).unsqueeze(0).half().to(\"cuda\")

        >>> prompt = \"A robot, 4k photo\"
        >>> 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\"

        >>> generator = torch.Generator(device=\"cuda\").manual_seed(43)

        >>> image_emb, zero_image_emb = pipe_prior(
        ...     prompt=prompt, negative_prompt=negative_prior_prompt, generator=generator
        ... ).to_tuple()

        >>> images = pipe(
        ...     image_embeds=image_emb,
        ...     negative_image_embeds=zero_image_emb,
        ...     hint=hint,
        ...     num_inference_steps=50,
        ...     generator=generator,
        ...     height=768,
        ...     width=768,
        ... ).images

        >>> images[0].save(\"robot_cat.png\")
        ```
"""


def _A       (   __magic_name__  , __magic_name__  , __magic_name__=8	):
   lowercase__			       =  height // scale_factor**2
   if height % scale_factor**2 != 0:
      new_height += 1
   lowercase__			       =  width // scale_factor**2
   if width % scale_factor**2 != 0:
      new_width += 1
   return new_height * scale_factor, new_width * scale_factor







class  lowerCAmelCase							(       lowercase_ ):







  def __init__(      self  :str   ,							_lowercase  :UNetaDConditionModel   ,							_lowercase  :DDPMScheduler   ,							_lowercase  :VQModel   ,							):






     '''simple docstring'''





     super().__init__()

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







  def   UpperCAmelCase					(      self  :Any   ,							_lowercase  :Tuple   ,							_lowercase  :Dict   ,							_lowercase  :List[str]   ,							_lowercase  :List[Any]   ,							_lowercase  :str   ,							_lowercase  :Optional[Any] ):






     '''simple docstring'''





     if latents is None:
        lowercase__			       =  randn_tensor(_lowercase   ,							generator=_lowercase   ,							device=_lowercase   ,							dtype=_lowercase )
     else:
        if latents.shape != shape:
           raise ValueError(f'''Unexpected latents shape, got {latents.shape}, expected {shape}''' )
        lowercase__			       =  latents.to(_lowercase )

     lowercase__			       =  latents * scheduler.init_noise_sigma
     return latents







  def   UpperCAmelCase					(      self  :Tuple   ,							_lowercase  :List[str]=0 ):






     '''simple docstring'''





     if is_accelerate_available():
        from accelerate import cpu_offload
     else:
        raise ImportError("Please install accelerate via `pip install accelerate`" )

     lowercase__			       =  torch.device(f'''cuda:{gpu_id}''' )

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







  def   UpperCAmelCase					(      self  :str   ,							_lowercase  :Tuple=0 ):






     '''simple docstring'''





     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." )

     lowercase__			       =  torch.device(f'''cuda:{gpu_id}''' )

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

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

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







  @property
  # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline._execution_device
  def   UpperCAmelCase					(      self  :Any ):






     '''simple docstring'''





     if not hasattr(self.unet   ,							"_hf_hook" ):
        return self.device
     for module in self.unet.modules():
        if (
            hasattr(_lowercase   ,							"_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(_lowercase )
  def __call__(      self  :Any   ,							_lowercase  :Union[torch.FloatTensor, List[torch.FloatTensor]]   ,							_lowercase  :Union[torch.FloatTensor, List[torch.FloatTensor]]   ,							_lowercase  :torch.FloatTensor   ,							_lowercase  :int = 5_12   ,							_lowercase  :int = 5_12   ,							_lowercase  :int = 1_00   ,							_lowercase  :float = 4.0   ,							_lowercase  :int = 1   ,							_lowercase  :Optional[Union[torch.Generator, List[torch.Generator]]] = None   ,							_lowercase  :Optional[torch.FloatTensor] = None   ,							_lowercase  :Optional[str] = "pil"   ,							_lowercase  :bool = True   ,							):






     '''simple docstring'''





     lowercase__			       =  self._execution_device

     lowercase__			       =  guidance_scale > 1.0

     if isinstance(_lowercase   ,							_lowercase ):
        lowercase__			       =  torch.cat(_lowercase   ,							dim=0 )
     if isinstance(_lowercase   ,							_lowercase ):
        lowercase__			       =  torch.cat(_lowercase   ,							dim=0 )
     if isinstance(_lowercase   ,							_lowercase ):
        lowercase__			       =  torch.cat(_lowercase   ,							dim=0 )

     lowercase__			       =  image_embeds.shape[0] * num_images_per_prompt

     if do_classifier_free_guidance:
        lowercase__			       =  image_embeds.repeat_interleave(_lowercase   ,							dim=0 )
        lowercase__			       =  negative_image_embeds.repeat_interleave(_lowercase   ,							dim=0 )
        lowercase__			       =  hint.repeat_interleave(_lowercase   ,							dim=0 )

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

     lowercase__			       =  self.movq.config.latent_channels

     lowercase__       ,							lowercase__			       =  downscale_height_and_width(_lowercase   ,							_lowercase   ,							self.movq_scale_factor )

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

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

        lowercase__			       =  {"image_embeds": image_embeds, "hint": hint}
        lowercase__			       =  self.unet(
            sample=_lowercase   ,							timestep=_lowercase   ,							encoder_hidden_states=_lowercase   ,							added_cond_kwargs=_lowercase   ,							return_dict=_lowercase   ,							)[0]

        if do_classifier_free_guidance:
           lowercase__       ,							lowercase__			       =  noise_pred.split(latents.shape[1]   ,							dim=1 )
           lowercase__       ,							lowercase__			       =  noise_pred.chunk(2 )
           lowercase__       ,							lowercase__			       =  variance_pred.chunk(2 )
           lowercase__			       =  noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)
           lowercase__			       =  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"]
        ):
           lowercase__       ,							lowercase__			       =  noise_pred.split(latents.shape[1]   ,							dim=1 )

        # compute the previous noisy sample x_t -> x_t-1
        lowercase__			       =  self.scheduler.step(
            _lowercase   ,							_lowercase   ,							_lowercase   ,							generator=_lowercase   ,							)[0]
     # post-processing
     lowercase__			       =  self.movq.decode(_lowercase   ,							force_not_quantize=_lowercase )["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"]:
        lowercase__			       =  image * 0.5 + 0.5
        lowercase__			       =  image.clamp(0   ,							1 )
        lowercase__			       =  image.cpu().permute(0   ,							2   ,							3   ,							1 ).float().numpy()

     if output_type == "pil":
        lowercase__			       =  self.numpy_to_pil(_lowercase )

     if not return_dict:
        return (image,)

     return ImagePipelineOutput(images=_lowercase )













import time
from dataclasses import dataclass
from multiprocessing import Pool
from unittest import TestCase
from unittest.mock import patch

import multiprocess
import numpy as np
import pytest

from datasets.utils.py_utils import (
    NestedDataStructure,
    asdict,
    iflatmap_unordered,
    map_nested,
    temp_seed,
    temporary_assignment,
    zip_dict,
)

from .utils import require_tf, require_torch


def _A       (   __magic_name__	):  # picklable for multiprocessing
   return x.sum()


def _A       (   __magic_name__	):  # picklable for multiprocessing
   return i + 1





@dataclass
class  lowerCAmelCase							:
  __lowerCamelCase		=      42
  __lowerCamelCase		=      42





class  lowerCAmelCase							(       lowercase_ ):







  def   UpperCAmelCase					(      self  :Optional[Any] ):






     '''simple docstring'''





     lowercase__			       =  {}
     lowercase__			       =  []
     lowercase__			       =  1
     lowercase__			       =  [1, 2]
     lowercase__			       =  {"a": 1, "b": 2}
     lowercase__			       =  {"a": [1, 2], "b": [3, 4]}
     lowercase__			       =  {"a": {"1": 1}, "b": 2}
     lowercase__			       =  {"a": 1, "b": 2, "c": 3, "d": 4}
     lowercase__			       =  {}
     lowercase__			       =  []
     lowercase__			       =  2
     lowercase__			       =  [2, 3]
     lowercase__			       =  {"a": 2, "b": 3}
     lowercase__			       =  {"a": [2, 3], "b": [4, 5]}
     lowercase__			       =  {"a": {"1": 2}, "b": 3}
     lowercase__			       =  {"a": 2, "b": 3, "c": 4, "d": 5}
     self.assertEqual(map_nested(_lowercase   ,							_lowercase )   ,							_lowercase )
     self.assertEqual(map_nested(_lowercase   ,							_lowercase )   ,							_lowercase )
     self.assertEqual(map_nested(_lowercase   ,							_lowercase )   ,							_lowercase )
     self.assertEqual(map_nested(_lowercase   ,							_lowercase )   ,							_lowercase )
     self.assertEqual(map_nested(_lowercase   ,							_lowercase )   ,							_lowercase )
     self.assertEqual(map_nested(_lowercase   ,							_lowercase )   ,							_lowercase )
     self.assertEqual(map_nested(_lowercase   ,							_lowercase )   ,							_lowercase )
     self.assertEqual(map_nested(_lowercase   ,							_lowercase )   ,							_lowercase )

     lowercase__			       =  2
     self.assertEqual(map_nested(_lowercase   ,							_lowercase   ,							num_proc=_lowercase )   ,							_lowercase )
     self.assertEqual(map_nested(_lowercase   ,							_lowercase   ,							num_proc=_lowercase )   ,							_lowercase )
     self.assertEqual(map_nested(_lowercase   ,							_lowercase   ,							num_proc=_lowercase )   ,							_lowercase )
     self.assertEqual(map_nested(_lowercase   ,							_lowercase   ,							num_proc=_lowercase )   ,							_lowercase )
     self.assertEqual(map_nested(_lowercase   ,							_lowercase   ,							num_proc=_lowercase )   ,							_lowercase )
     self.assertEqual(map_nested(_lowercase   ,							_lowercase   ,							num_proc=_lowercase )   ,							_lowercase )
     self.assertEqual(map_nested(_lowercase   ,							_lowercase   ,							num_proc=_lowercase )   ,							_lowercase )
     self.assertEqual(map_nested(_lowercase   ,							_lowercase   ,							num_proc=_lowercase )   ,							_lowercase )

     lowercase__			       =  {"a": np.eye(2 ), "b": np.zeros(3 ), "c": np.ones(2 )}
     lowercase__			       =  {"a": 2, "b": 0, "c": 2}
     lowercase__			       =  {
         "a": np.eye(2 ).astype(_lowercase ),
         "b": np.zeros(3 ).astype(_lowercase ),
         "c": np.ones(2 ).astype(_lowercase ),
     }
     self.assertEqual(map_nested(_lowercase   ,							_lowercase   ,							map_numpy=_lowercase )   ,							_lowercase )
     self.assertEqual(
         {k: v.tolist() for k, v in map_nested(_lowercase   ,							_lowercase   ,							map_numpy=_lowercase ).items()}   ,							{k: v.tolist() for k, v in expected_map_nested_sna_int.items()}   ,							)
     self.assertEqual(map_nested(_lowercase   ,							_lowercase   ,							map_numpy=_lowercase   ,							num_proc=_lowercase )   ,							_lowercase )
     self.assertEqual(
         {k: v.tolist() for k, v in map_nested(_lowercase   ,							_lowercase   ,							map_numpy=_lowercase   ,							num_proc=_lowercase ).items()}   ,							{k: v.tolist() for k, v in expected_map_nested_sna_int.items()}   ,							)
     with self.assertRaises(_lowercase ):  # can't pickle a local lambda
        map_nested(lambda _lowercase : x + 1   ,							_lowercase   ,							num_proc=_lowercase )







  def   UpperCAmelCase					(      self  :List[Any] ):






     '''simple docstring'''





     lowercase__			       =  {"a": 1, "b": 2}
     lowercase__			       =  {"a": 3, "b": 4}
     lowercase__			       =  {"a": 5, "b": 6}
     lowercase__			       =  sorted([("a", (1, 3, 5)), ("b", (2, 4, 6))] )
     self.assertEqual(sorted(zip_dict(_lowercase   ,							_lowercase   ,							_lowercase ) )   ,							_lowercase )







  def   UpperCAmelCase					(      self  :Union[str, Any] ):






     '''simple docstring'''





     class  lowerCAmelCase							:
       __lowerCamelCase		=      'bar'
     lowercase__			       =  Foo()
     self.assertEqual(foo.my_attr   ,							"bar" )
     with temporary_assignment(_lowercase   ,							"my_attr"   ,							"BAR" ):
        self.assertEqual(foo.my_attr   ,							"BAR" )
     self.assertEqual(foo.my_attr   ,							"bar" )


@pytest.mark.parametrize(
    "iterable_length, num_proc, expected_num_proc"  , [
        (1, None, 1),
        (1, 1, 1),
        (2, None, 1),
        (2, 1, 1),
        (2, 2, 1),
        (2, 3, 1),
        (3, 2, 1),
        (16, 16, 16),
        (16, 17, 16),
        (17, 16, 16),
    ]  , )
def _A       (   __magic_name__  , __magic_name__  , __magic_name__	):
   with patch("datasets.utils.py_utils._single_map_nested"	) as mock_single_map_nested, patch(
       "datasets.parallel.parallel.Pool"	) as mock_multiprocessing_pool:
      lowercase__			       =  {f'''{i}''': i for i in range(__magic_name__	)}
      lowercase__			       =  map_nested(lambda __magic_name__	: x + 10  , __magic_name__  , num_proc=__magic_name__  , parallel_min_length=16	)
      if expected_num_proc == 1:
         assert mock_single_map_nested.called
         assert not mock_multiprocessing_pool.called
      else:
         assert not mock_single_map_nested.called
         assert mock_multiprocessing_pool.called
         assert mock_multiprocessing_pool.call_args[0][0] == expected_num_proc





class  lowerCAmelCase							(       lowercase_ ):







  @require_tf
  def   UpperCAmelCase					(      self  :List[Any] ):






     '''simple docstring'''





     import tensorflow as tf
     from tensorflow.keras import layers

     lowercase__			       =  layers.Dense(2 )

     def gen_random_output():
        lowercase__			       =  tf.random.uniform((1, 3) )
        return model(_lowercase ).numpy()

     with temp_seed(42   ,							set_tensorflow=_lowercase ):
        lowercase__			       =  gen_random_output()
     with temp_seed(42   ,							set_tensorflow=_lowercase ):
        lowercase__			       =  gen_random_output()
     lowercase__			       =  gen_random_output()

     np.testing.assert_equal(_lowercase   ,							_lowercase )
     self.assertGreater(np.abs(outa - outa ).sum()   ,							0 )







  @require_torch
  def   UpperCAmelCase					(      self  :Optional[int] ):






     '''simple docstring'''





     import torch

     def gen_random_output():
        lowercase__			       =  torch.nn.Linear(3   ,							2 )
        lowercase__			       =  torch.rand(1   ,							3 )
        return model(_lowercase ).detach().numpy()

     with temp_seed(42   ,							set_pytorch=_lowercase ):
        lowercase__			       =  gen_random_output()
     with temp_seed(42   ,							set_pytorch=_lowercase ):
        lowercase__			       =  gen_random_output()
     lowercase__			       =  gen_random_output()

     np.testing.assert_equal(_lowercase   ,							_lowercase )
     self.assertGreater(np.abs(outa - outa ).sum()   ,							0 )







  def   UpperCAmelCase					(      self  :str ):






     '''simple docstring'''





     def gen_random_output():
        return np.random.rand(1   ,							3 )

     with temp_seed(42 ):
        lowercase__			       =  gen_random_output()
     with temp_seed(42 ):
        lowercase__			       =  gen_random_output()
     lowercase__			       =  gen_random_output()

     np.testing.assert_equal(_lowercase   ,							_lowercase )
     self.assertGreater(np.abs(outa - outa ).sum()   ,							0 )


@pytest.mark.parametrize("input_data"  , [{}]	)
def _A       (   __magic_name__	):
   lowercase__			       =  NestedDataStructure(__magic_name__	).data
   assert output_data == input_data


@pytest.mark.parametrize(
    "data, expected_output"  , [
        ({}, []),
        ([], []),
        ("foo", ["foo"]),
        (["foo", "bar"], ["foo", "bar"]),
        ([["foo", "bar"]], ["foo", "bar"]),
        ([[["foo"], ["bar"]]], ["foo", "bar"]),
        ([[["foo"], "bar"]], ["foo", "bar"]),
        ({"a": 1, "b": 2}, [1, 2]),
        ({"a": [1, 2], "b": [3, 4]}, [1, 2, 3, 4]),
        ({"a": [[1, 2]], "b": [[3, 4]]}, [1, 2, 3, 4]),
        ({"a": [[1, 2]], "b": [3, 4]}, [1, 2, 3, 4]),
        ({"a": [[[1], [2]]], "b": [[[3], [4]]]}, [1, 2, 3, 4]),
        ({"a": [[[1], [2]]], "b": [[3, 4]]}, [1, 2, 3, 4]),
        ({"a": [[[1], [2]]], "b": [3, 4]}, [1, 2, 3, 4]),
        ({"a": [[[1], [2]]], "b": [3, [4]]}, [1, 2, 3, 4]),
        ({"a": {"1": 1}, "b": 2}, [1, 2]),
        ({"a": {"1": [1]}, "b": 2}, [1, 2]),
        ({"a": {"1": [1]}, "b": [2]}, [1, 2]),
    ]  , )
def _A       (   __magic_name__  , __magic_name__	):
   lowercase__			       =  NestedDataStructure(__magic_name__	).flatten()
   assert output == expected_output


def _A       (   ):
   lowercase__			       =  A(x=1  , y="foobar"	)
   lowercase__			       =  {"x": 1, "y": "foobar"}
   assert asdict(__magic_name__	) == expected_output

   lowercase__			       =  {"a": {"b": A(x=10  , y="foo"	)}, "c": [A(x=20  , y="bar"	)]}
   lowercase__			       =  {"a": {"b": {"x": 10, "y": "foo"}}, "c": [{"x": 20, "y": "bar"}]}
   assert asdict(__magic_name__	) == expected_output

   with pytest.raises(__magic_name__	):
      asdict([1, A(x=10  , y="foo"	)]	)


def _A       (   __magic_name__	):
   return text.split()


def _A       (   __magic_name__	):
   yield (time.time(), content)
   time.sleep(2	)
   yield (time.time(), content)







def _A       (   ):
   with Pool(2	) as pool:
      lowercase__			       =  list(iflatmap_unordered(__magic_name__  , _split_text  , kwargs_iterable=[{"text": "hello there"}] * 10	)	)
      assert out.count("hello"	) == 10
      assert out.count("there"	) == 10
      assert len(__magic_name__	) == 20

   # check multiprocess from pathos (uses dill for pickling)
   with multiprocess.Pool(2	) as pool:
      lowercase__			       =  list(iflatmap_unordered(__magic_name__  , _split_text  , kwargs_iterable=[{"text": "hello there"}] * 10	)	)
      assert out.count("hello"	) == 10
      assert out.count("there"	) == 10
      assert len(__magic_name__	) == 20

   # check that we get items as fast as possible
   with Pool(2	) as pool:
      lowercase__			       =  []
      for yield_time, content in iflatmap_unordered(
          __magic_name__  , _aseconds_generator_of_aitems_with_timing  , kwargs_iterable=[{"content": "a"}, {"content": "b"}]	):
         assert yield_time < time.time() + 0.1, "we should each item directly after it was yielded"
         out.append(__magic_name__	)
      assert out.count("a"	) == 2
      assert out.count("b"	) == 2
      assert len(__magic_name__	) == 4






