



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


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

lowerCAmelCase	:		Tuple				=   '''
    Examples:
        ```py
        >>> from diffusers import KandinskyV22Pipeline, KandinskyV22PriorPipeline
        >>> import torch

        >>> pipe_prior = KandinskyV22PriorPipeline.from_pretrained("kandinsky-community/kandinsky-2-2-prior")
        >>> pipe_prior.to("cuda")
        >>> prompt = "red cat, 4k photo"
        >>> out = pipe_prior(prompt)
        >>> image_emb = out.image_embeds
        >>> zero_image_emb = out.negative_image_embeds
        >>> pipe = KandinskyV22Pipeline.from_pretrained("kandinsky-community/kandinsky-2-2-decoder")
        >>> pipe.to("cuda")
        >>> image = pipe(
        ...     image_embeds=image_emb,
        ...     negative_image_embeds=zero_image_emb,
        ...     height=768,
        ...     width=768,
        ...     num_inference_steps=50,
        ... ).images
        >>> image[0].save("cat.png")
        ```
'''



def   __lowerCAmelCase					(							lowerCamelCase   :						int				,		lowerCamelCase   :						str				,		lowerCamelCase   :						Union[str, Any]=8					):






	'''simple docstring'''




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






class 					UpperCAmelCase__	(		UpperCamelCase__ ):



					def __init__(    self			,				UpperCamelCase			,				UpperCamelCase			,				UpperCamelCase			,				)     ->			Union[str, Any]:
						super().__init__()

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



					def        UpperCAmelCase_ (    self			,				UpperCamelCase			,				UpperCamelCase			,				UpperCamelCase			,				UpperCamelCase			,				UpperCamelCase			,				UpperCamelCase     )     ->			Tuple:
						if latents is None:
							__lowerCAmelCase												=  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}'''     )
							__lowerCAmelCase												=  latents.to(UpperCamelCase     )

						__lowerCAmelCase												=  latents * scheduler.init_noise_sigma
						return latents



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

						__lowerCAmelCase												=  torch.device(F'''cuda:{gpu_id}'''     )

						__lowerCAmelCase												=  [
						    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."     )

						__lowerCAmelCase												=  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)

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

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



					@property
					# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline._execution_device
					def        UpperCAmelCase_ (    self     )     ->			Union[str, Any]:
						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 = 512			,				UpperCamelCase = 512			,				UpperCamelCase = 100			,				UpperCamelCase = 4.0			,				UpperCamelCase = 1			,				UpperCamelCase = None			,				UpperCamelCase = None			,				UpperCamelCase = "pil"			,				UpperCamelCase = True			,				)     ->			int:
						__lowerCAmelCase												=  self._execution_device

						__lowerCAmelCase												=  guidance_scale > 1.0

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

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

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

						self.scheduler.set_timesteps(UpperCamelCase			,				device=UpperCamelCase     )
						__lowerCAmelCase												=  self.scheduler.timesteps

						__lowerCAmelCase												=  self.unet.config.in_channels

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

						# create initial latent
						__lowerCAmelCase												=  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
							__lowerCAmelCase												=  torch.cat([latents] * 2     ) if do_classifier_free_guidance else latents

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

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

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

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

						if not return_dict:
							return (image,)

						return ImagePipelineOutput(images=UpperCamelCase     )




'''simple docstring'''



from __future__ import annotations

import math



def   __lowerCAmelCase					(							lowerCamelCase   :						list				,		lowerCamelCase   :						list					):






	'''simple docstring'''




	if len(lowerCamelCase					) != 2 or len(a[0]					) != 2 or len(lowerCamelCase					) != 2 or len(b[0]					) != 2:
		raise Exception("Matrices are not 2x2"					)
	__lowerCAmelCase												=  [
	    [a[0][0] * b[0][0] + a[0][1] * b[1][0], a[0][0] * b[0][1] + a[0][1] * b[1][1]],
	    [a[1][0] * b[0][0] + a[1][1] * b[1][0], a[1][0] * b[0][1] + a[1][1] * b[1][1]],
	]
	return new_matrix



def   __lowerCAmelCase					(							lowerCamelCase   :						list				,		lowerCamelCase   :						list					):






	'''simple docstring'''




	return [
	    [matrix_a[row][col] + matrix_b[row][col] for col in range(len(matrix_a[row]					)					)]
	    for row in range(len(lowerCamelCase					)					)
	]



def   __lowerCAmelCase					(							lowerCamelCase   :						list				,		lowerCamelCase   :						list					):






	'''simple docstring'''




	return [
	    [matrix_a[row][col] - matrix_b[row][col] for col in range(len(matrix_a[row]					)					)]
	    for row in range(len(lowerCamelCase					)					)
	]



def   __lowerCAmelCase					(							lowerCamelCase   :						list					):






	'''simple docstring'''




	if len(lowerCamelCase					) % 2 != 0 or len(a[0]					) % 2 != 0:
		raise Exception("Odd matrices are not supported!"					)

	__lowerCAmelCase												=  len(lowerCamelCase					)
	__lowerCAmelCase												=  matrix_length // 2

	__lowerCAmelCase												=  [[a[i][j] for j in range(lowerCamelCase				,		lowerCamelCase					)] for i in range(lowerCamelCase					)]
	__lowerCAmelCase												=  [
	    [a[i][j] for j in range(lowerCamelCase				,		lowerCamelCase					)] for i in range(lowerCamelCase				,		lowerCamelCase					)
	]

	__lowerCAmelCase												=  [[a[i][j] for j in range(lowerCamelCase					)] for i in range(lowerCamelCase					)]
	__lowerCAmelCase												=  [[a[i][j] for j in range(lowerCamelCase					)] for i in range(lowerCamelCase				,		lowerCamelCase					)]

	return top_left, top_right, bot_left, bot_right



def   __lowerCAmelCase					(							lowerCamelCase   :						list					):






	'''simple docstring'''




	return len(lowerCamelCase					), len(matrix[0]					)



def   __lowerCAmelCase					(							lowerCamelCase   :						list					):






	'''simple docstring'''




	print("\n".join(str(lowerCamelCase					) for line in matrix					)					)



def   __lowerCAmelCase					(							lowerCamelCase   :						list				,		lowerCamelCase   :						list					):






	'''simple docstring'''




	if matrix_dimensions(lowerCamelCase					) == (2, 2):
		return default_matrix_multiplication(lowerCamelCase				,		lowerCamelCase					)

	__lowerCAmelCase							,       __lowerCAmelCase							,       __lowerCAmelCase							,       __lowerCAmelCase												=  split_matrix(lowerCamelCase					)
	__lowerCAmelCase							,       __lowerCAmelCase							,       __lowerCAmelCase							,       __lowerCAmelCase												=  split_matrix(lowerCamelCase					)

	__lowerCAmelCase												=  actual_strassen(lowerCamelCase				,		matrix_subtraction(lowerCamelCase				,		lowerCamelCase					)					)
	__lowerCAmelCase												=  actual_strassen(matrix_addition(lowerCamelCase				,		lowerCamelCase					)				,		lowerCamelCase					)
	__lowerCAmelCase												=  actual_strassen(matrix_addition(lowerCamelCase				,		lowerCamelCase					)				,		lowerCamelCase					)
	__lowerCAmelCase												=  actual_strassen(lowerCamelCase				,		matrix_subtraction(lowerCamelCase				,		lowerCamelCase					)					)
	__lowerCAmelCase												=  actual_strassen(matrix_addition(lowerCamelCase				,		lowerCamelCase					)				,		matrix_addition(lowerCamelCase				,		lowerCamelCase					)					)
	__lowerCAmelCase												=  actual_strassen(matrix_subtraction(lowerCamelCase				,		lowerCamelCase					)				,		matrix_addition(lowerCamelCase				,		lowerCamelCase					)					)
	__lowerCAmelCase												=  actual_strassen(matrix_subtraction(lowerCamelCase				,		lowerCamelCase					)				,		matrix_addition(lowerCamelCase				,		lowerCamelCase					)					)

	__lowerCAmelCase												=  matrix_addition(matrix_subtraction(matrix_addition(lowerCamelCase				,		lowerCamelCase					)				,		lowerCamelCase					)				,		lowerCamelCase					)
	__lowerCAmelCase												=  matrix_addition(lowerCamelCase				,		lowerCamelCase					)
	__lowerCAmelCase												=  matrix_addition(lowerCamelCase				,		lowerCamelCase					)
	__lowerCAmelCase												=  matrix_subtraction(matrix_subtraction(matrix_addition(lowerCamelCase				,		lowerCamelCase					)				,		lowerCamelCase					)				,		lowerCamelCase					)

	# construct the new matrix from our 4 quadrants
	__lowerCAmelCase												=  []
	for i in range(len(lowerCamelCase					)					):
		new_matrix.append(top_left[i] + top_right[i]					)
	for i in range(len(lowerCamelCase					)					):
		new_matrix.append(bot_left[i] + bot_right[i]					)
	return new_matrix



def   __lowerCAmelCase					(							lowerCamelCase   :						list				,		lowerCamelCase   :						list					):






	'''simple docstring'''




	if matrix_dimensions(lowerCamelCase					)[1] != matrix_dimensions(lowerCamelCase					)[0]:
		__lowerCAmelCase												=  (
		    "Unable to multiply these matrices, please check the dimensions.\n"
		    f'''Matrix A: {matrixa}\n'''
		    f'''Matrix B: {matrixa}'''
		)
		raise Exception(lowerCamelCase					)
	__lowerCAmelCase												=  matrix_dimensions(lowerCamelCase					)
	__lowerCAmelCase												=  matrix_dimensions(lowerCamelCase					)

	if dimensiona[0] == dimensiona[1] and dimensiona[0] == dimensiona[1]:
		return [matrixa, matrixa]

	__lowerCAmelCase												=  max(*lowerCamelCase				,		*lowerCamelCase					)
	__lowerCAmelCase												=  int(math.pow(2				,		math.ceil(math.loga(lowerCamelCase					)					)					)					)
	__lowerCAmelCase												=  matrixa
	__lowerCAmelCase												=  matrixa

	# Adding zeros to the matrices so that the arrays dimensions are the same and also
	# power of 2
	for i in range(0				,		lowerCamelCase					):
		if i < dimensiona[0]:
			for _ in range(dimensiona[1]				,		lowerCamelCase					):
				new_matrixa[i].append(0					)
		else:
			new_matrixa.append([0] * maxim					)
		if i < dimensiona[0]:
			for _ in range(dimensiona[1]				,		lowerCamelCase					):
				new_matrixa[i].append(0					)
		else:
			new_matrixa.append([0] * maxim					)

	__lowerCAmelCase												=  actual_strassen(lowerCamelCase				,		lowerCamelCase					)

	# Removing the additional zeros
	for i in range(0				,		lowerCamelCase					):
		if i < dimensiona[0]:
			for _ in range(dimensiona[1]				,		lowerCamelCase					):
				final_matrix[i].pop()
		else:
			final_matrix.pop()
	return final_matrix


if __name__ == "__main__":
	lowerCAmelCase	:		Tuple				=   [
	    [2, 3, 4, 5],
	    [6, 4, 3, 1],
	    [2, 3, 6, 7],
	    [3, 1, 2, 4],
	    [2, 3, 4, 5],
	    [6, 4, 3, 1],
	    [2, 3, 6, 7],
	    [3, 1, 2, 4],
	    [2, 3, 4, 5],
	    [6, 2, 3, 1],
	]
	lowerCAmelCase	:		Any				=   [[0, 2, 1, 1], [1_6, 2, 3, 3], [2, 2, 7, 7], [1_3, 1_1, 2_2, 4]]
	print(strassen(matrixa, matrixa))