

import gc
import unittest

import numpy as np
import torch
import torch.nn.functional as F
from transformers import (
    ClapTextConfig,
    ClapTextModelWithProjection,
    RobertaTokenizer,
    SpeechTaHifiGan,
    SpeechTaHifiGanConfig,
)

from diffusers import (
    AudioLDMPipeline,
    AutoencoderKL,
    DDIMScheduler,
    LMSDiscreteScheduler,
    PNDMScheduler,
    UNetaDConditionModel,
)
from diffusers.utils import is_xformers_available, slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism

from ..pipeline_params import TEXT_TO_AUDIO_BATCH_PARAMS, TEXT_TO_AUDIO_PARAMS
from ..test_pipelines_common import PipelineTesterMixin


enable_full_determinism()




class      UpperCamelCase (     SCREAMING_SNAKE_CASE	,			unittest.TestCase   ):
    __UpperCamelCase         =AudioLDMPipeline
    __UpperCamelCase         =TEXT_TO_AUDIO_PARAMS
    __UpperCamelCase         =TEXT_TO_AUDIO_BATCH_PARAMS
    __UpperCamelCase         =frozenset(
        [
            "num_inference_steps",
            "num_waveforms_per_prompt",
            "generator",
            "latents",
            "output_type",
            "return_dict",
            "callback",
            "callback_steps",
        ]   )

    def UpperCamelCase			(			self :      Optional[Any]    ):

     """simple docstring"""






     torch.manual_seed(0    )
     SCREAMING_SNAKE_CASE								=			UNetaDConditionModel(
         block_out_channels=(3_2, 6_4)							,    layers_per_block=2							,    sample_size=3_2							,    in_channels=4							,    out_channels=4							,    down_block_types=('DownBlock2D', 'CrossAttnDownBlock2D')							,    up_block_types=('CrossAttnUpBlock2D', 'UpBlock2D')							,    cross_attention_dim=(3_2, 6_4)							,    class_embed_type='simple_projection'							,    projection_class_embeddings_input_dim=3_2							,    class_embeddings_concat=snake_case__							,    )
     SCREAMING_SNAKE_CASE								=			DDIMScheduler(
         beta_start=0.00_085							,    beta_end=0.012							,    beta_schedule='scaled_linear'							,    clip_sample=snake_case__							,    set_alpha_to_one=snake_case__							,    )
     torch.manual_seed(0    )
     SCREAMING_SNAKE_CASE								=			AutoencoderKL(
         block_out_channels=[3_2, 6_4]							,    in_channels=1							,    out_channels=1							,    down_block_types=['DownEncoderBlock2D', 'DownEncoderBlock2D']							,    up_block_types=['UpDecoderBlock2D', 'UpDecoderBlock2D']							,    latent_channels=4							,    )
     torch.manual_seed(0    )
     SCREAMING_SNAKE_CASE								=			ClapTextConfig(
         bos_token_id=0							,    eos_token_id=2							,    hidden_size=3_2							,    intermediate_size=3_7							,    layer_norm_eps=1E-05							,    num_attention_heads=4							,    num_hidden_layers=5							,    pad_token_id=1							,    vocab_size=1_0_0_0							,    projection_dim=3_2							,    )
     SCREAMING_SNAKE_CASE								=			ClapTextModelWithProjection(snake_case__    )
     SCREAMING_SNAKE_CASE								=			RobertaTokenizer.from_pretrained('hf-internal-testing/tiny-random-roberta'							,    model_max_length=7_7    )

     SCREAMING_SNAKE_CASE								=			SpeechTaHifiGanConfig(
         model_in_dim=8							,    sampling_rate=1_6_0_0_0							,    upsample_initial_channel=1_6							,    upsample_rates=[2, 2]							,    upsample_kernel_sizes=[4, 4]							,    resblock_kernel_sizes=[3, 7]							,    resblock_dilation_sizes=[[1, 3, 5], [1, 3, 5]]							,    normalize_before=snake_case__							,    )

     SCREAMING_SNAKE_CASE								=			SpeechTaHifiGan(snake_case__    )

     SCREAMING_SNAKE_CASE								=			{
         'unet': unet,
         'scheduler': scheduler,
         'vae': vae,
         'text_encoder': text_encoder,
         'tokenizer': tokenizer,
         'vocoder': vocoder,
     }
     return components

    def UpperCamelCase			(			self :      Optional[int]							,    snake_case__ :      int							,    snake_case__ :      int=0    ):

     """simple docstring"""






     if str(snake_case__    ).startswith('mps'    ):
      SCREAMING_SNAKE_CASE								=			torch.manual_seed(snake_case__    )
     else:
      SCREAMING_SNAKE_CASE								=			torch.Generator(device=snake_case__    ).manual_seed(snake_case__    )
     SCREAMING_SNAKE_CASE								=			{
         'prompt': 'A hammer hitting a wooden surface',
         'generator': generator,
         'num_inference_steps': 2,
         'guidance_scale': 6.0,
     }
     return inputs

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

     """simple docstring"""






     SCREAMING_SNAKE_CASE								=			'cpu'  # ensure determinism for the device-dependent torch.Generator

     SCREAMING_SNAKE_CASE								=			self.get_dummy_components()
     SCREAMING_SNAKE_CASE								=			AudioLDMPipeline(**snake_case__    )
     SCREAMING_SNAKE_CASE								=			audioldm_pipe.to(snake_case__    )
     audioldm_pipe.set_progress_bar_config(disable=snake_case__    )

     SCREAMING_SNAKE_CASE								=			self.get_dummy_inputs(snake_case__    )
     SCREAMING_SNAKE_CASE								=			audioldm_pipe(**snake_case__    )
     SCREAMING_SNAKE_CASE								=			output.audios[0]

     assert audio.ndim == 1
     assert len(snake_case__    ) == 2_5_6

     SCREAMING_SNAKE_CASE								=			audio[:1_0]
     SCREAMING_SNAKE_CASE								=			np.array(
         [-0.0_050, 0.0_050, -0.0_060, 0.0_033, -0.0_026, 0.0_033, -0.0_027, 0.0_033, -0.0_028, 0.0_033]    )

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

    def UpperCamelCase			(			self :      Dict    ):

     """simple docstring"""






     SCREAMING_SNAKE_CASE								=			self.get_dummy_components()
     SCREAMING_SNAKE_CASE								=			AudioLDMPipeline(**snake_case__    )
     SCREAMING_SNAKE_CASE								=			audioldm_pipe.to(snake_case__    )
     SCREAMING_SNAKE_CASE								=			audioldm_pipe.to(snake_case__    )
     audioldm_pipe.set_progress_bar_config(disable=snake_case__    )

     SCREAMING_SNAKE_CASE								=			self.get_dummy_inputs(snake_case__    )
     SCREAMING_SNAKE_CASE								=			3 * [inputs['prompt']]

     # forward
     SCREAMING_SNAKE_CASE								=			audioldm_pipe(**snake_case__    )
     SCREAMING_SNAKE_CASE								=			output.audios[0]

     SCREAMING_SNAKE_CASE								=			self.get_dummy_inputs(snake_case__    )
     SCREAMING_SNAKE_CASE								=			3 * [inputs.pop('prompt'    )]

     SCREAMING_SNAKE_CASE								=			audioldm_pipe.tokenizer(
         snake_case__							,    padding='max_length'							,    max_length=audioldm_pipe.tokenizer.model_max_length							,    truncation=snake_case__							,    return_tensors='pt'							,    )
     SCREAMING_SNAKE_CASE								=			text_inputs['input_ids'].to(snake_case__    )

     SCREAMING_SNAKE_CASE								=			audioldm_pipe.text_encoder(
         snake_case__							,    )
     SCREAMING_SNAKE_CASE								=			prompt_embeds.text_embeds
     # additional L_2 normalization over each hidden-state
     SCREAMING_SNAKE_CASE								=			F.normalize(snake_case__							,    dim=-1    )

     SCREAMING_SNAKE_CASE								=			prompt_embeds

     # forward
     SCREAMING_SNAKE_CASE								=			audioldm_pipe(**snake_case__    )
     SCREAMING_SNAKE_CASE								=			output.audios[0]

     assert np.abs(audio_a - audio_a    ).max() < 1E-2

    def UpperCamelCase			(			self :      int    ):

     """simple docstring"""






     SCREAMING_SNAKE_CASE								=			self.get_dummy_components()
     SCREAMING_SNAKE_CASE								=			AudioLDMPipeline(**snake_case__    )
     SCREAMING_SNAKE_CASE								=			audioldm_pipe.to(snake_case__    )
     SCREAMING_SNAKE_CASE								=			audioldm_pipe.to(snake_case__    )
     audioldm_pipe.set_progress_bar_config(disable=snake_case__    )

     SCREAMING_SNAKE_CASE								=			self.get_dummy_inputs(snake_case__    )
     SCREAMING_SNAKE_CASE								=			3 * ['this is a negative prompt']
     SCREAMING_SNAKE_CASE								=			negative_prompt
     SCREAMING_SNAKE_CASE								=			3 * [inputs['prompt']]

     # forward
     SCREAMING_SNAKE_CASE								=			audioldm_pipe(**snake_case__    )
     SCREAMING_SNAKE_CASE								=			output.audios[0]

     SCREAMING_SNAKE_CASE								=			self.get_dummy_inputs(snake_case__    )
     SCREAMING_SNAKE_CASE								=			3 * [inputs.pop('prompt'    )]

     SCREAMING_SNAKE_CASE								=			[]
     for p in [prompt, negative_prompt]:
      SCREAMING_SNAKE_CASE								=			audioldm_pipe.tokenizer(
          snake_case__							,    padding='max_length'							,    max_length=audioldm_pipe.tokenizer.model_max_length							,    truncation=snake_case__							,    return_tensors='pt'							,    )
      SCREAMING_SNAKE_CASE								=			text_inputs['input_ids'].to(snake_case__    )

      SCREAMING_SNAKE_CASE								=			audioldm_pipe.text_encoder(
          snake_case__							,    )
      SCREAMING_SNAKE_CASE								=			text_embeds.text_embeds
      # additional L_2 normalization over each hidden-state
      SCREAMING_SNAKE_CASE								=			F.normalize(snake_case__							,    dim=-1    )

      embeds.append(snake_case__    )

     SCREAMING_SNAKE_CASE						,	SCREAMING_SNAKE_CASE								=			embeds

     # forward
     SCREAMING_SNAKE_CASE								=			audioldm_pipe(**snake_case__    )
     SCREAMING_SNAKE_CASE								=			output.audios[0]

     assert np.abs(audio_a - audio_a    ).max() < 1E-2

    def UpperCamelCase			(			self :      int    ):

     """simple docstring"""






     SCREAMING_SNAKE_CASE								=			'cpu'  # ensure determinism for the device-dependent torch.Generator
     SCREAMING_SNAKE_CASE								=			self.get_dummy_components()
     SCREAMING_SNAKE_CASE								=			PNDMScheduler(skip_prk_steps=snake_case__    )
     SCREAMING_SNAKE_CASE								=			AudioLDMPipeline(**snake_case__    )
     SCREAMING_SNAKE_CASE								=			audioldm_pipe.to(snake_case__    )
     audioldm_pipe.set_progress_bar_config(disable=snake_case__    )

     SCREAMING_SNAKE_CASE								=			self.get_dummy_inputs(snake_case__    )
     SCREAMING_SNAKE_CASE								=			'egg cracking'
     SCREAMING_SNAKE_CASE								=			audioldm_pipe(**snake_case__							,    negative_prompt=snake_case__    )
     SCREAMING_SNAKE_CASE								=			output.audios[0]

     assert audio.ndim == 1
     assert len(snake_case__    ) == 2_5_6

     SCREAMING_SNAKE_CASE								=			audio[:1_0]
     SCREAMING_SNAKE_CASE								=			np.array(
         [-0.0_051, 0.0_050, -0.0_060, 0.0_034, -0.0_026, 0.0_033, -0.0_027, 0.0_033, -0.0_028, 0.0_032]    )

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

    def UpperCamelCase			(			self :      str    ):

     """simple docstring"""






     SCREAMING_SNAKE_CASE								=			'cpu'  # ensure determinism for the device-dependent torch.Generator
     SCREAMING_SNAKE_CASE								=			self.get_dummy_components()
     SCREAMING_SNAKE_CASE								=			PNDMScheduler(skip_prk_steps=snake_case__    )
     SCREAMING_SNAKE_CASE								=			AudioLDMPipeline(**snake_case__    )
     SCREAMING_SNAKE_CASE								=			audioldm_pipe.to(snake_case__    )
     audioldm_pipe.set_progress_bar_config(disable=snake_case__    )

     SCREAMING_SNAKE_CASE								=			'A hammer hitting a wooden surface'

     # test num_waveforms_per_prompt=1 (default)
     SCREAMING_SNAKE_CASE								=			audioldm_pipe(snake_case__							,    num_inference_steps=2    ).audios

     assert audios.shape == (1, 2_5_6)

     # test num_waveforms_per_prompt=1 (default) for batch of prompts
     SCREAMING_SNAKE_CASE								=			2
     SCREAMING_SNAKE_CASE								=			audioldm_pipe([prompt] * batch_size							,    num_inference_steps=2    ).audios

     assert audios.shape == (batch_size, 2_5_6)

     # test num_waveforms_per_prompt for single prompt
     SCREAMING_SNAKE_CASE								=			2
     SCREAMING_SNAKE_CASE								=			audioldm_pipe(snake_case__							,    num_inference_steps=2							,    num_waveforms_per_prompt=snake_case__    ).audios

     assert audios.shape == (num_waveforms_per_prompt, 2_5_6)

     # test num_waveforms_per_prompt for batch of prompts
     SCREAMING_SNAKE_CASE								=			2
     SCREAMING_SNAKE_CASE								=			audioldm_pipe(
         [prompt] * batch_size							,    num_inference_steps=2							,    num_waveforms_per_prompt=snake_case__    ).audios

     assert audios.shape == (batch_size * num_waveforms_per_prompt, 2_5_6)

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

     """simple docstring"""






     SCREAMING_SNAKE_CASE								=			'cpu'  # ensure determinism for the device-dependent torch.Generator
     SCREAMING_SNAKE_CASE								=			self.get_dummy_components()
     SCREAMING_SNAKE_CASE								=			AudioLDMPipeline(**snake_case__    )
     SCREAMING_SNAKE_CASE								=			audioldm_pipe.to(snake_case__    )
     audioldm_pipe.set_progress_bar_config(disable=snake_case__    )
     SCREAMING_SNAKE_CASE								=			audioldm_pipe.vocoder.config.sampling_rate

     SCREAMING_SNAKE_CASE								=			self.get_dummy_inputs(snake_case__    )
     SCREAMING_SNAKE_CASE								=			audioldm_pipe(audio_length_in_s=0.016							,    **snake_case__    )
     SCREAMING_SNAKE_CASE								=			output.audios[0]

     assert audio.ndim == 1
     assert len(snake_case__    ) / vocoder_sampling_rate == 0.016

     SCREAMING_SNAKE_CASE								=			audioldm_pipe(audio_length_in_s=0.032							,    **snake_case__    )
     SCREAMING_SNAKE_CASE								=			output.audios[0]

     assert audio.ndim == 1
     assert len(snake_case__    ) / vocoder_sampling_rate == 0.032

    def UpperCamelCase			(			self :      Optional[Any]    ):

     """simple docstring"""






     SCREAMING_SNAKE_CASE								=			self.get_dummy_components()
     SCREAMING_SNAKE_CASE								=			AudioLDMPipeline(**snake_case__    )
     SCREAMING_SNAKE_CASE								=			audioldm_pipe.to(snake_case__    )
     audioldm_pipe.set_progress_bar_config(disable=snake_case__    )

     SCREAMING_SNAKE_CASE								=			['hey']

     SCREAMING_SNAKE_CASE								=			audioldm_pipe(snake_case__							,    num_inference_steps=1    )
     SCREAMING_SNAKE_CASE								=			output.audios.shape
     assert audio_shape == (1, 2_5_6)

     SCREAMING_SNAKE_CASE								=			audioldm_pipe.vocoder.config
     config.model_in_dim *= 2
     SCREAMING_SNAKE_CASE								=			SpeechTaHifiGan(snake_case__    ).to(snake_case__    )
     SCREAMING_SNAKE_CASE								=			audioldm_pipe(snake_case__							,    num_inference_steps=1    )
     SCREAMING_SNAKE_CASE								=			output.audios.shape
     # waveform shape is unchanged, we just have 2x the number of mel channels in the spectrogram
     assert audio_shape == (1, 2_5_6)

    def UpperCamelCase			(			self :      Tuple    ):

     """simple docstring"""






     self._test_attention_slicing_forward_pass(test_mean_pixel_difference=snake_case__    )

    def UpperCamelCase			(			self :      int    ):

     """simple docstring"""






     self._test_inference_batch_single_identical(test_mean_pixel_difference=snake_case__    )




    @unittest.skipIf(
        torch_device != 'cuda' or not is_xformers_available()							,    reason='XFormers attention is only available with CUDA and `xformers` installed'							,    )
    def UpperCamelCase			(			self :      Dict    ):

     """simple docstring"""






     self._test_xformers_attention_forwardGenerator_pass(test_mean_pixel_difference=snake_case__    )




@slow
class      UpperCamelCase (     unittest.TestCase   ):

    def UpperCamelCase			(			self :      Any    ):

     """simple docstring"""






     super().tearDown()
     gc.collect()
     torch.cuda.empty_cache()

    def UpperCamelCase			(			self :      int							,    snake_case__ :      int							,    snake_case__ :      Tuple="cpu"							,    snake_case__ :      List[str]=torch.floataa							,    snake_case__ :      Optional[Any]=0    ):

     """simple docstring"""






     SCREAMING_SNAKE_CASE								=			torch.Generator(device=snake_case__    ).manual_seed(snake_case__    )
     SCREAMING_SNAKE_CASE								=			np.random.RandomState(snake_case__    ).standard_normal((1, 8, 1_2_8, 1_6)    )
     SCREAMING_SNAKE_CASE								=			torch.from_numpy(snake_case__    ).to(device=snake_case__							,    dtype=snake_case__    )
     SCREAMING_SNAKE_CASE								=			{
         'prompt': 'A hammer hitting a wooden surface',
         'latents': latents,
         'generator': generator,
         'num_inference_steps': 3,
         'guidance_scale': 2.5,
     }
     return inputs

    def UpperCamelCase			(			self :      int    ):

     """simple docstring"""






     SCREAMING_SNAKE_CASE								=			AudioLDMPipeline.from_pretrained('cvssp/audioldm'    )
     SCREAMING_SNAKE_CASE								=			audioldm_pipe.to(snake_case__    )
     audioldm_pipe.set_progress_bar_config(disable=snake_case__    )

     SCREAMING_SNAKE_CASE								=			self.get_inputs(snake_case__    )
     SCREAMING_SNAKE_CASE								=			2_5
     SCREAMING_SNAKE_CASE								=			audioldm_pipe(**snake_case__    ).audios[0]

     assert audio.ndim == 1
     assert len(snake_case__    ) == 8_1_9_2_0

     SCREAMING_SNAKE_CASE								=			audio[7_7_2_3_0:7_7_2_4_0]
     SCREAMING_SNAKE_CASE								=			np.array(
         [-0.4_884, -0.4_607, 0.0_023, 0.5_007, 0.5_896, 0.5_151, 0.3_813, -0.0_208, -0.3_687, -0.4_315]    )
     SCREAMING_SNAKE_CASE								=			np.abs(expected_slice - audio_slice    ).max()
     assert max_diff < 1E-2




    def UpperCamelCase			(			self :      Optional[Any]    ):

     """simple docstring"""






     SCREAMING_SNAKE_CASE								=			AudioLDMPipeline.from_pretrained('cvssp/audioldm'    )
     SCREAMING_SNAKE_CASE								=			LMSDiscreteScheduler.from_config(audioldm_pipe.scheduler.config    )
     SCREAMING_SNAKE_CASE								=			audioldm_pipe.to(snake_case__    )
     audioldm_pipe.set_progress_bar_config(disable=snake_case__    )

     SCREAMING_SNAKE_CASE								=			self.get_inputs(snake_case__    )
     SCREAMING_SNAKE_CASE								=			audioldm_pipe(**snake_case__    ).audios[0]

     assert audio.ndim == 1
     assert len(snake_case__    ) == 8_1_9_2_0

     SCREAMING_SNAKE_CASE								=			audio[2_7_7_8_0:2_7_7_9_0]
     SCREAMING_SNAKE_CASE								=			np.array([-0.2_131, -0.0_873, -0.0_124, -0.0_189, 0.0_569, 0.1_373, 0.1_883, 0.2_886, 0.3_297, 0.2_212]    )
     SCREAMING_SNAKE_CASE								=			np.abs(expected_slice - audio_slice    ).max()
     assert max_diff < 3E-2




from typing import List, Optional, Union

import numpy as np
import PIL
import torch
from PIL import Image

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


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

a_    :						str    		=      "\n    Examples:\n        ```py\n        >>> from diffusers import KandinskyV22Img2ImgPipeline, KandinskyV22PriorPipeline\n        >>> from diffusers.utils import load_image\n        >>> import torch\n\n        >>> pipe_prior = KandinskyV22PriorPipeline.from_pretrained(\n        ...     \"kandinsky-community/kandinsky-2-2-prior\", torch_dtype=torch.float16\n        ... )\n        >>> pipe_prior.to(\"cuda\")\n\n        >>> prompt = \"A red cartoon frog, 4k\"\n        >>> image_emb, zero_image_emb = pipe_prior(prompt, return_dict=False)\n\n        >>> pipe = KandinskyV22Img2ImgPipeline.from_pretrained(\n        ...     \"kandinsky-community/kandinsky-2-2-decoder\", torch_dtype=torch.float16\n        ... )\n        >>> pipe.to(\"cuda\")\n\n        >>> init_image = load_image(\n        ...     \"https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main\"\n        ...     \"/kandinsky/frog.png\"\n        ... )\n\n        >>> image = pipe(\n        ...     image=init_image,\n        ...     image_embeds=image_emb,\n        ...     negative_image_embeds=zero_image_emb,\n        ...     height=768,\n        ...     width=768,\n        ...     num_inference_steps=100,\n        ...     strength=0.2,\n        ... ).images\n\n        >>> image[0].save(\"red_frog.png\")\n        ```\n"


def        __lowerCAmelCase    (       _UpperCamelCase		:  List[str]		,							_UpperCamelCase		:  Any		,							_UpperCamelCase		:  Any=8	)		->							List[str]:
 '''simple docstring'''


 SCREAMING_SNAKE_CASE								=			height // scale_factor**2
 if height % scale_factor**2 != 0:
  new_height += 1
 SCREAMING_SNAKE_CASE								=			width // scale_factor**2
 if width % scale_factor**2 != 0:
  new_width += 1
 return new_height * scale_factor, new_width * scale_factor


def        __lowerCAmelCase    (       _UpperCamelCase		:  Optional[int]		,							_UpperCamelCase		:  List[Any]=5_12		,							_UpperCamelCase		:  Union[str, Any]=5_12	)		->							Union[str, Any]:
 '''simple docstring'''


 SCREAMING_SNAKE_CASE								=			pil_image.resize((w, h)		,							resample=Image.BICUBIC		,							reducing_gap=1	)
 SCREAMING_SNAKE_CASE								=			np.array(pil_image.convert('RGB'	)	)
 SCREAMING_SNAKE_CASE								=			arr.astype(np.floataa	) / 1_27.5 - 1
 SCREAMING_SNAKE_CASE								=			np.transpose(_UpperCamelCase		,							[2, 0, 1]	)
 SCREAMING_SNAKE_CASE								=			torch.from_numpy(_UpperCamelCase	).unsqueeze(0	)
 return image




class      UpperCamelCase (     SCREAMING_SNAKE_CASE   ):

    def __init__(			self :      int							,    snake_case__ :      UNetaDConditionModel							,    snake_case__ :      DDPMScheduler							,    snake_case__ :      VQModel							,    ):

     """simple docstring"""






     super().__init__()

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

    def UpperCamelCase			(			self :      Any							,    snake_case__ :      int							,    snake_case__ :      List[Any]							,    snake_case__ :      List[str]    ):

     """simple docstring"""






     SCREAMING_SNAKE_CASE								=			min(int(num_inference_steps * strength    )							,    snake_case__    )

     SCREAMING_SNAKE_CASE								=			max(num_inference_steps - init_timestep							,    0    )
     SCREAMING_SNAKE_CASE								=			self.scheduler.timesteps[t_start:]

     return timesteps, num_inference_steps - t_start

    def UpperCamelCase			(			self :      List[str]							,    snake_case__ :      Dict							,    snake_case__ :      int							,    snake_case__ :      Optional[int]							,    snake_case__ :      Optional[Any]							,    snake_case__ :      Union[str, Any]							,    snake_case__ :      Optional[int]							,    snake_case__ :      str=None    ):

     """simple docstring"""






     if not isinstance(snake_case__							,    (torch.Tensor, PIL.Image.Image, list)    ):
      raise ValueError(
          F"""`image` has to be of type `torch.Tensor`, `PIL.Image.Image` or list but is {type(snake_case__    )}"""    )

     SCREAMING_SNAKE_CASE								=			image.to(device=snake_case__							,    dtype=snake_case__    )

     SCREAMING_SNAKE_CASE								=			batch_size * num_images_per_prompt

     if image.shape[1] == 4:
      SCREAMING_SNAKE_CASE								=			image

     else:
      if isinstance(snake_case__							,    snake_case__    ) and len(snake_case__    ) != batch_size:
       raise ValueError(
           F"""You have passed a list of generators of length {len(snake_case__    )}, but requested an effective batch"""
           F""" size of {batch_size}. Make sure the batch size matches the length of the generators."""    )

      elif isinstance(snake_case__							,    snake_case__    ):
       SCREAMING_SNAKE_CASE								=			[
           self.movq.encode(image[i : i + 1]    ).latent_dist.sample(generator[i]    ) for i in range(snake_case__    )
       ]
       SCREAMING_SNAKE_CASE								=			torch.cat(snake_case__							,    dim=0    )
      else:
       SCREAMING_SNAKE_CASE								=			self.movq.encode(snake_case__    ).latent_dist.sample(snake_case__    )

      SCREAMING_SNAKE_CASE								=			self.movq.config.scaling_factor * init_latents

     SCREAMING_SNAKE_CASE								=			torch.cat([init_latents]							,    dim=0    )

     SCREAMING_SNAKE_CASE								=			init_latents.shape
     SCREAMING_SNAKE_CASE								=			randn_tensor(snake_case__							,    generator=snake_case__							,    device=snake_case__							,    dtype=snake_case__    )

     # get latents
     SCREAMING_SNAKE_CASE								=			self.scheduler.add_noise(snake_case__							,    snake_case__							,    snake_case__    )

     SCREAMING_SNAKE_CASE								=			init_latents

     return latents

    def UpperCamelCase			(			self :      int							,    snake_case__ :      List[str]=0    ):

     """simple docstring"""






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

     SCREAMING_SNAKE_CASE								=			torch.device(F"""cuda:{gpu_id}"""    )

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

    def UpperCamelCase			(			self :      Union[str, Any]							,    snake_case__ :      Optional[int]=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.'    )

     SCREAMING_SNAKE_CASE								=			torch.device(F"""cuda:{gpu_id}"""    )

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

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

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

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

     """simple docstring"""






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

     """simple docstring"""






     SCREAMING_SNAKE_CASE								=			self._execution_device

     SCREAMING_SNAKE_CASE								=			guidance_scale > 1.0

     if isinstance(snake_case__							,    snake_case__    ):
      SCREAMING_SNAKE_CASE								=			torch.cat(snake_case__							,    dim=0    )
     SCREAMING_SNAKE_CASE								=			image_embeds.shape[0]
     if isinstance(snake_case__							,    snake_case__    ):
      SCREAMING_SNAKE_CASE								=			torch.cat(snake_case__							,    dim=0    )

     if do_classifier_free_guidance:
      SCREAMING_SNAKE_CASE								=			image_embeds.repeat_interleave(snake_case__							,    dim=0    )
      SCREAMING_SNAKE_CASE								=			negative_image_embeds.repeat_interleave(snake_case__							,    dim=0    )

     SCREAMING_SNAKE_CASE								=			torch.cat([negative_image_embeds, image_embeds]							,    dim=0    ).to(dtype=self.unet.dtype							,    device=snake_case__    )

     if not isinstance(snake_case__							,    snake_case__    ):
      SCREAMING_SNAKE_CASE								=			[image]
     if not all(isinstance(snake_case__							,    (PIL.Image.Image, torch.Tensor)    ) for i in image    ):
      raise ValueError(
          F"""Input is in incorrect format: {[type(snake_case__    ) for i in image]}. Currently, we only support  PIL image and pytorch tensor"""    )

     SCREAMING_SNAKE_CASE								=			torch.cat([prepare_image(snake_case__							,    snake_case__							,    snake_case__    ) for i in image]							,    dim=0    )
     SCREAMING_SNAKE_CASE								=			image.to(dtype=image_embeds.dtype							,    device=snake_case__    )

     SCREAMING_SNAKE_CASE								=			self.movq.encode(snake_case__    )['latents']
     SCREAMING_SNAKE_CASE								=			latents.repeat_interleave(snake_case__							,    dim=0    )
     self.scheduler.set_timesteps(snake_case__							,    device=snake_case__    )
     SCREAMING_SNAKE_CASE						,	SCREAMING_SNAKE_CASE								=			self.get_timesteps(snake_case__							,    snake_case__							,    snake_case__    )
     SCREAMING_SNAKE_CASE								=			timesteps[:1].repeat(batch_size * num_images_per_prompt    )
     SCREAMING_SNAKE_CASE						,	SCREAMING_SNAKE_CASE								=			downscale_height_and_width(snake_case__							,    snake_case__							,    self.movq_scale_factor    )
     SCREAMING_SNAKE_CASE								=			self.prepare_latents(
         snake_case__							,    snake_case__							,    snake_case__							,    snake_case__							,    image_embeds.dtype							,    snake_case__							,    snake_case__    )
     for i, t in enumerate(self.progress_bar(snake_case__    )    ):
      # expand the latents if we are doing classifier free guidance
      SCREAMING_SNAKE_CASE								=			torch.cat([latents] * 2    ) if do_classifier_free_guidance else latents

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

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

      # compute the previous noisy sample x_t -> x_t-1
      SCREAMING_SNAKE_CASE								=			self.scheduler.step(
          snake_case__							,    snake_case__							,    snake_case__							,    generator=snake_case__							,    )[0]

     # post-processing
     SCREAMING_SNAKE_CASE								=			self.movq.decode(snake_case__							,    force_not_quantize=snake_case__    )['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"]:
      SCREAMING_SNAKE_CASE								=			image * 0.5 + 0.5
      SCREAMING_SNAKE_CASE								=			image.clamp(0							,    1    )
      SCREAMING_SNAKE_CASE								=			image.cpu().permute(0							,    2							,    3							,    1    ).float().numpy()

     if output_type == "pil":
      SCREAMING_SNAKE_CASE								=			self.numpy_to_pil(snake_case__    )

     if not return_dict:
      return (image,)

     return ImagePipelineOutput(images=snake_case__    )

