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


lowercase__		:Optional[Any]   =    logging.get_logger(__name__)  # pylint: disable=invalid-name

lowercase__		:Optional[Any]   =    "\n    Examples:\n        ```py\n        >>> from diffusers import KandinskyV22Pipeline, KandinskyV22PriorPipeline\n        >>> import torch\n\n        >>> pipe_prior = KandinskyV22PriorPipeline.from_pretrained(\"kandinsky-community/kandinsky-2-2-prior\")\n        >>> pipe_prior.to(\"cuda\")\n        >>> prompt = \"red cat, 4k photo\"\n        >>> out = pipe_prior(prompt)\n        >>> image_emb = out.image_embeds\n        >>> zero_image_emb = out.negative_image_embeds\n        >>> pipe = KandinskyV22Pipeline.from_pretrained(\"kandinsky-community/kandinsky-2-2-decoder\")\n        >>> pipe.to(\"cuda\")\n        >>> image = pipe(\n        ...     image_embeds=image_emb,\n        ...     negative_image_embeds=zero_image_emb,\n        ...     height=768,\n        ...     width=768,\n        ...     num_inference_steps=50,\n        ... ).images\n        >>> image[0].save(\"cat.png\")\n        ```\n"



def        UpperCamelCase    (							lowerCAmelCase__ ,			lowerCAmelCase__ ,			lowerCAmelCase__=8		):
					'''simple docstring'''
					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 				lowercase		(					SCREAMING_SNAKE_CASE__		):




							def __init__(				self   ,A__   ,A__   ,A__   ,):
												super().__init__()

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




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

												lowercase							=						latents * scheduler.init_noise_sigma
												return latents




							def  A__      (				self   ,A__=0):
												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(A__   ,A__)




							def  A__      (				self   ,A__=0):
												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=A__)
																	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(A__   ,A__   ,prev_module_hook=A__)

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




							@property
							# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline._execution_device
							def  A__      (				self):
												if not hasattr(self.unet   ,'''_hf_hook'''):
																	return self.device
												for module in self.unet.modules():
																	if (
																	    hasattr(A__   ,'''_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(A__)
							def __call__(				self   ,A__   ,A__   ,A__ = 5_1_2   ,A__ = 5_1_2   ,A__ = 1_0_0   ,A__ = 4.0   ,A__ = 1   ,A__ = None   ,A__ = None   ,A__ = "pil"   ,A__ = True   ,):
												lowercase							=						self._execution_device

												lowercase							=						guidance_scale > 1.0

												if isinstance(A__   ,A__):
																	lowercase							=						torch.cat(A__   ,dim=0)
												lowercase							=						image_embeds.shape[0] * num_images_per_prompt
												if isinstance(A__   ,A__):
																	lowercase							=						torch.cat(A__   ,dim=0)

												if do_classifier_free_guidance:
																	lowercase							=						image_embeds.repeat_interleave(A__   ,dim=0)
																	lowercase							=						negative_image_embeds.repeat_interleave(A__   ,dim=0)

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

												self.scheduler.set_timesteps(A__   ,device=A__)
												lowercase							=						self.scheduler.timesteps

												lowercase							=						self.unet.config.in_channels

												lowercase	,						lowercase							=						downscale_height_and_width(A__   ,A__   ,self.movq_scale_factor)

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

												for i, t in enumerate(self.progress_bar(A__)):
																	# 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}
																	lowercase							=						self.unet(
																	    sample=A__   ,timestep=A__   ,encoder_hidden_states=A__   ,added_cond_kwargs=A__   ,return_dict=A__   ,)[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(
																	    A__   ,A__   ,A__   ,generator=A__   ,)[0]
												# post-processing
												lowercase							=						self.movq.decode(A__   ,force_not_quantize=A__)['''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(A__)

												if not return_dict:
																	return (image,)

												return ImagePipelineOutput(images=A__)






import gc
import random
import unittest

import numpy as np
import torch
from transformers import XLMRobertaTokenizer

from diffusers import (
    AltDiffusionImgaImgPipeline,
    AutoencoderKL,
    PNDMScheduler,
    UNetaDConditionModel,
)
from diffusers.image_processor import VaeImageProcessor
from diffusers.pipelines.alt_diffusion.modeling_roberta_series import (
    RobertaSeriesConfig,
    RobertaSeriesModelWithTransformation,
)
from diffusers.utils import floats_tensor, load_image, load_numpy, slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu


enable_full_determinism()







class 				lowercase		(					unittest.TestCase		):




							def  A__      (				self):
												# clean up the VRAM after each test
												super().tearDown()
												gc.collect()
												torch.cuda.empty_cache()




							@property
							def  A__      (				self):
												lowercase							=						1
												lowercase							=						3
												lowercase							=						(3_2, 3_2)

												lowercase							=						floats_tensor((batch_size, num_channels) + sizes   ,rng=random.Random(0)).to(A__)
												return image




							@property
							def  A__      (				self):
												torch.manual_seed(0)
												lowercase							=						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   ,)
												return model




							@property
							def  A__      (				self):
												torch.manual_seed(0)
												lowercase							=						AutoencoderKL(
												    block_out_channels=[3_2, 6_4]   ,in_channels=3   ,out_channels=3   ,down_block_types=['''DownEncoderBlock2D''', '''DownEncoderBlock2D''']   ,up_block_types=['''UpDecoderBlock2D''', '''UpDecoderBlock2D''']   ,latent_channels=4   ,)
												return model




							@property
							def  A__      (				self):
												torch.manual_seed(0)
												lowercase							=						RobertaSeriesConfig(
												    hidden_size=3_2   ,project_dim=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=5_0_0_6   ,)
												return RobertaSeriesModelWithTransformation(A__)




							@property
							def  A__      (				self):
												def extract(*A__   ,**A__):







																	class 				lowercase		:




																								def __init__(				self):
																													lowercase							=						torch.ones([0])





																								def  A__      (				self   ,A__):
																													self.pixel_values.to(A__)
																													return self


																	return Out()

												return extract




							def  A__      (				self):
												lowercase							=						'''cpu'''  # ensure determinism for the device-dependent torch.Generator
												lowercase							=						self.dummy_cond_unet
												lowercase							=						PNDMScheduler(skip_prk_steps=A__)
												lowercase							=						self.dummy_vae
												lowercase							=						self.dummy_text_encoder
												lowercase							=						XLMRobertaTokenizer.from_pretrained('''hf-internal-testing/tiny-xlm-roberta''')
												lowercase							=						7_7

												lowercase							=						self.dummy_image.to(A__)
												lowercase							=						init_image / 2 + 0.5

												# make sure here that pndm scheduler skips prk
												lowercase							=						AltDiffusionImgaImgPipeline(
												    unet=A__   ,scheduler=A__   ,vae=A__   ,text_encoder=A__   ,tokenizer=A__   ,safety_checker=A__   ,feature_extractor=self.dummy_extractor   ,)
												lowercase							=						VaeImageProcessor(vae_scale_factor=alt_pipe.vae_scale_factor   ,do_normalize=A__)
												lowercase							=						alt_pipe.to(A__)
												alt_pipe.set_progress_bar_config(disable=A__)

												lowercase							=						'''A painting of a squirrel eating a burger'''
												lowercase							=						torch.Generator(device=A__).manual_seed(0)
												lowercase							=						alt_pipe(
												    [prompt]   ,generator=A__   ,guidance_scale=6.0   ,num_inference_steps=2   ,output_type='''np'''   ,image=A__   ,)

												lowercase							=						output.images

												lowercase							=						torch.Generator(device=A__).manual_seed(0)
												lowercase							=						alt_pipe(
												    [prompt]   ,generator=A__   ,guidance_scale=6.0   ,num_inference_steps=2   ,output_type='''np'''   ,image=A__   ,return_dict=A__   ,)[0]

												lowercase							=						image[0, -3:, -3:, -1]
												lowercase							=						image_from_tuple[0, -3:, -3:, -1]

												assert image.shape == (1, 3_2, 3_2, 3)
												lowercase							=						np.array([0.4427, 0.3731, 0.4249, 0.4941, 0.4546, 0.4148, 0.4193, 0.4666, 0.4499])

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




							@unittest.skipIf(torch_device != '''cuda'''   ,'''This test requires a GPU''')
							def  A__      (				self):
												lowercase							=						self.dummy_cond_unet
												lowercase							=						PNDMScheduler(skip_prk_steps=A__)
												lowercase							=						self.dummy_vae
												lowercase							=						self.dummy_text_encoder
												lowercase							=						XLMRobertaTokenizer.from_pretrained('''hf-internal-testing/tiny-xlm-roberta''')
												lowercase							=						7_7

												lowercase							=						self.dummy_image.to(A__)

												# put models in fp16
												lowercase							=						unet.half()
												lowercase							=						vae.half()
												lowercase							=						bert.half()

												# make sure here that pndm scheduler skips prk
												lowercase							=						AltDiffusionImgaImgPipeline(
												    unet=A__   ,scheduler=A__   ,vae=A__   ,text_encoder=A__   ,tokenizer=A__   ,safety_checker=A__   ,feature_extractor=self.dummy_extractor   ,)
												lowercase							=						VaeImageProcessor(vae_scale_factor=alt_pipe.vae_scale_factor   ,do_normalize=A__)
												lowercase							=						alt_pipe.to(A__)
												alt_pipe.set_progress_bar_config(disable=A__)

												lowercase							=						'''A painting of a squirrel eating a burger'''
												lowercase							=						torch.manual_seed(0)
												lowercase							=						alt_pipe(
												    [prompt]   ,generator=A__   ,num_inference_steps=2   ,output_type='''np'''   ,image=A__   ,).images

												assert image.shape == (1, 3_2, 3_2, 3)





							@unittest.skipIf(torch_device != '''cuda'''   ,'''This test requires a GPU''')
							def  A__      (				self):
												lowercase							=						load_image(
												    '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'''
												    '''/img2img/sketch-mountains-input.jpg''')
												# resize to resolution that is divisible by 8 but not 16 or 32
												lowercase							=						init_image.resize((7_6_0, 5_0_4))

												lowercase							=						'''BAAI/AltDiffusion'''
												lowercase							=						AltDiffusionImgaImgPipeline.from_pretrained(
												    A__   ,safety_checker=A__   ,)
												pipe.to(A__)
												pipe.set_progress_bar_config(disable=A__)
												pipe.enable_attention_slicing()

												lowercase							=						'''A fantasy landscape, trending on artstation'''

												lowercase							=						torch.manual_seed(0)
												lowercase							=						pipe(
												    prompt=A__   ,image=A__   ,strength=0.75   ,guidance_scale=7.5   ,generator=A__   ,output_type='''np'''   ,)
												lowercase							=						output.images[0]

												lowercase							=						image[2_5_5:2_5_8, 3_8_3:3_8_6, -1]

												assert image.shape == (5_0_4, 7_6_0, 3)
												lowercase							=						np.array([0.9358, 0.9397, 0.9599, 0.9901, 1.0000, 1.0000, 0.9882, 1.0000, 1.0000])

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







@slow
@require_torch_gpu
class 				lowercase		(					unittest.TestCase		):




							def  A__      (				self):
												# clean up the VRAM after each test
												super().tearDown()
												gc.collect()
												torch.cuda.empty_cache()





							def  A__      (				self):
												lowercase							=						load_image(
												    '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'''
												    '''/img2img/sketch-mountains-input.jpg''')
												lowercase							=						init_image.resize((7_6_8, 5_1_2))
												lowercase							=						load_numpy(
												    '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/img2img/fantasy_landscape_alt.npy''')

												lowercase							=						'''BAAI/AltDiffusion'''
												lowercase							=						AltDiffusionImgaImgPipeline.from_pretrained(
												    A__   ,safety_checker=A__   ,)
												pipe.to(A__)
												pipe.set_progress_bar_config(disable=A__)
												pipe.enable_attention_slicing()

												lowercase							=						'''A fantasy landscape, trending on artstation'''

												lowercase							=						torch.manual_seed(0)
												lowercase							=						pipe(
												    prompt=A__   ,image=A__   ,strength=0.75   ,guidance_scale=7.5   ,generator=A__   ,output_type='''np'''   ,)
												lowercase							=						output.images[0]

												assert image.shape == (5_1_2, 7_6_8, 3)
												# img2img is flaky across GPUs even in fp32, so using MAE here
												assert np.abs(expected_image - image).max() < 1E-2




