





from typing import List, Optional, Union

import torch
from transformers import (
    XLMRobertaTokenizer,
)

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


UpperCamelCase__   : List[str]					  = logging.get_logger(__name__)  # pylint: disable=invalid-name

UpperCamelCase__   : Tuple					  = """
    Examples:
        ```py
        >>> from diffusers import KandinskyPipeline, KandinskyPriorPipeline
        >>> import torch

        >>> pipe_prior = KandinskyPriorPipeline.from_pretrained(\"kandinsky-community/Kandinsky-2-1-prior\")
        >>> pipe_prior.to(\"cuda\")

        >>> prompt = \"red cat, 4k photo\"
        >>> out = pipe_prior(prompt)
        >>> image_emb = out.image_embeds
        >>> negative_image_emb = out.negative_image_embeds

        >>> pipe = KandinskyPipeline.from_pretrained(\"kandinsky-community/kandinsky-2-1\")
        >>> pipe.to(\"cuda\")

        >>> image = pipe(
        ...     prompt,
        ...     image_embeds=image_emb,
        ...     negative_image_embeds=negative_image_emb,
        ...     height=768,
        ...     width=768,
        ...     num_inference_steps=100,
        ... ).images

        >>> image[0].save(\"cat.png\")
        ```
"""




def    A_(			A				,      A				,      A=8   ):
	UpperCAmelCase_								=     h // scale_factor**2
	if h % scale_factor**2 != 0:
		new_h += 1
	UpperCAmelCase_								=     w // scale_factor**2
	if w % scale_factor**2 != 0:
		new_w += 1
	return new_h * scale_factor, new_w * scale_factor







class   _UpperCamelCase  (       A_  ):







					'''simple docstring'''

					def __init__(						self    :  List[str] ,					__lowercase    :  MultilingualCLIP ,					__lowercase    :  XLMRobertaTokenizer ,					__lowercase    :  UNetaDConditionModel ,					__lowercase    :  Union[DDIMScheduler, DDPMScheduler] ,					__lowercase    :  VQModel ,					):






						'''simple docstring'''







						super().__init__()

						self.register_modules(
						    text_encoder=__lowercase ,					tokenizer=__lowercase ,					unet=__lowercase ,					scheduler=__lowercase ,					movq=__lowercase ,					)
						UpperCAmelCase_								=     2 ** (len(self.movq.config.block_out_channels		) - 1)





					def        SCREAMING_SNAKE_CASE							(						self    :  List[Any] ,					__lowercase    :  Optional[int] ,					__lowercase    :  int ,					__lowercase    :  Any ,					__lowercase    :  Optional[Any] ,					__lowercase    :  Union[str, Any] ,					__lowercase    :  Tuple		):






						'''simple docstring'''







						if latents is None:
							UpperCAmelCase_								=     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}"""		)
							UpperCAmelCase_								=     latents.to(__lowercase		)

						UpperCAmelCase_								=     latents * scheduler.init_noise_sigma
						return latents





					def        SCREAMING_SNAKE_CASE							(						self    :  Any ,					__lowercase    :  Optional[int] ,					__lowercase    :  List[Any] ,					__lowercase    :  Optional[int] ,					__lowercase    :  Optional[Any] ,					__lowercase    :  Tuple=None ,					):






						'''simple docstring'''







						UpperCAmelCase_								=     len(__lowercase		) if isinstance(__lowercase ,					__lowercase		) else 1
						# get prompt text embeddings
						UpperCAmelCase_								=     self.tokenizer(
						    __lowercase ,					padding="""max_length""" ,					truncation=__lowercase ,					max_length=77 ,					return_attention_mask=__lowercase ,					add_special_tokens=__lowercase ,					return_tensors="""pt""" ,					)

						UpperCAmelCase_								=     text_inputs.input_ids
						UpperCAmelCase_								=     self.tokenizer(__lowercase ,					padding="""longest""" ,					return_tensors="""pt"""		).input_ids

						if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not torch.equal(__lowercase ,					__lowercase		):
							UpperCAmelCase_								=     self.tokenizer.batch_decode(untruncated_ids[:, self.tokenizer.model_max_length - 1 : -1]		)
							logger.warning(
							    """The following part of your input was truncated because CLIP can only handle sequences up to"""
							    F""" {self.tokenizer.model_max_length} tokens: {removed_text}"""		)

						UpperCAmelCase_								=     text_input_ids.to(__lowercase		)
						UpperCAmelCase_								=     text_inputs.attention_mask.to(__lowercase		)

						UpperCAmelCase_						,	UpperCAmelCase_								=     self.text_encoder(
						    input_ids=__lowercase ,					attention_mask=__lowercase		)

						UpperCAmelCase_								=     prompt_embeds.repeat_interleave(__lowercase ,					dim=0		)
						UpperCAmelCase_								=     text_encoder_hidden_states.repeat_interleave(__lowercase ,					dim=0		)
						UpperCAmelCase_								=     text_mask.repeat_interleave(__lowercase ,					dim=0		)

						if do_classifier_free_guidance:
							UpperCAmelCase_								=     42
							if negative_prompt is None:
								UpperCAmelCase_								=     [""""""] * batch_size
							elif type(__lowercase		) is not type(__lowercase		):
								raise TypeError(
								    F"""`negative_prompt` should be the same type to `prompt`, but got {type(__lowercase		)} !="""
								    F""" {type(__lowercase		)}."""		)
							elif isinstance(__lowercase ,					__lowercase		):
								UpperCAmelCase_								=     [negative_prompt]
							elif batch_size != len(__lowercase		):
								raise ValueError(
								    F"""`negative_prompt`: {negative_prompt} has batch size {len(__lowercase		)}, but `prompt`:"""
								    F""" {prompt} has batch size {batch_size}. Please make sure that passed `negative_prompt` matches"""
								    """ the batch size of `prompt`."""		)
							else:
								UpperCAmelCase_								=     negative_prompt

							UpperCAmelCase_								=     self.tokenizer(
							    __lowercase ,					padding="""max_length""" ,					max_length=77 ,					truncation=__lowercase ,					return_attention_mask=__lowercase ,					add_special_tokens=__lowercase ,					return_tensors="""pt""" ,					)
							UpperCAmelCase_								=     uncond_input.input_ids.to(__lowercase		)
							UpperCAmelCase_								=     uncond_input.attention_mask.to(__lowercase		)

							UpperCAmelCase_						,	UpperCAmelCase_								=     self.text_encoder(
							    input_ids=__lowercase ,					attention_mask=__lowercase		)

							# duplicate unconditional embeddings for each generation per prompt, using mps friendly method

							UpperCAmelCase_								=     negative_prompt_embeds.shape[1]
							UpperCAmelCase_								=     negative_prompt_embeds.repeat(1 ,					__lowercase		)
							UpperCAmelCase_								=     negative_prompt_embeds.view(batch_size * num_images_per_prompt ,					__lowercase		)

							UpperCAmelCase_								=     uncond_text_encoder_hidden_states.shape[1]
							UpperCAmelCase_								=     uncond_text_encoder_hidden_states.repeat(1 ,					__lowercase ,					1		)
							UpperCAmelCase_								=     uncond_text_encoder_hidden_states.view(
							    batch_size * num_images_per_prompt ,					__lowercase ,					-1		)
							UpperCAmelCase_								=     uncond_text_mask.repeat_interleave(__lowercase ,					dim=0		)

							# done duplicates

							# For classifier free guidance, we need to do two forward passes.
							# Here we concatenate the unconditional and text embeddings into a single batch
							# to avoid doing two forward passes
							UpperCAmelCase_								=     torch.cat([negative_prompt_embeds, prompt_embeds]		)
							UpperCAmelCase_								=     torch.cat([uncond_text_encoder_hidden_states, text_encoder_hidden_states]		)

							UpperCAmelCase_								=     torch.cat([uncond_text_mask, text_mask]		)

						return prompt_embeds, text_encoder_hidden_states, text_mask





					def        SCREAMING_SNAKE_CASE							(						self    :  str ,					__lowercase    :  Union[str, Any]=0		):






						'''simple docstring'''







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

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

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





					def        SCREAMING_SNAKE_CASE							(						self    :  Dict ,					__lowercase    :  Union[str, Any]=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."""		)

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

						UpperCAmelCase_								=     None
						for cpu_offloaded_model in [self.text_encoder, self.unet, self.movq]:
							UpperCAmelCase_						,	UpperCAmelCase_								=     cpu_offload_with_hook(__lowercase ,					__lowercase ,					prev_module_hook=__lowercase		)

						if self.safety_checker is not None:
							UpperCAmelCase_						,	UpperCAmelCase_								=     cpu_offload_with_hook(self.safety_checker ,					__lowercase ,					prev_module_hook=__lowercase		)

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





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






						'''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    :  List[str] ,					__lowercase    :  Union[str, List[str]] ,					__lowercase    :  Union[torch.FloatTensor, List[torch.FloatTensor]] ,					__lowercase    :  Union[torch.FloatTensor, List[torch.FloatTensor]] ,					__lowercase    :  Optional[Union[str, List[str]]] = None ,					__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'''







						if isinstance(__lowercase ,					__lowercase		):
							UpperCAmelCase_								=     1
						elif isinstance(__lowercase ,					__lowercase		):
							UpperCAmelCase_								=     len(__lowercase		)
						else:
							raise ValueError(F"""`prompt` has to be of type `str` or `list` but is {type(__lowercase		)}"""		)

						UpperCAmelCase_								=     self._execution_device

						UpperCAmelCase_								=     batch_size * num_images_per_prompt
						UpperCAmelCase_								=     guidance_scale > 1.0

						UpperCAmelCase_						,	UpperCAmelCase_						,	UpperCAmelCase_								=     self._encode_prompt(
						    __lowercase ,					__lowercase ,					__lowercase ,					__lowercase ,					__lowercase		)

						if isinstance(__lowercase ,					__lowercase		):
							UpperCAmelCase_								=     torch.cat(__lowercase ,					dim=0		)
						if isinstance(__lowercase ,					__lowercase		):
							UpperCAmelCase_								=     torch.cat(__lowercase ,					dim=0		)

						if do_classifier_free_guidance:
							UpperCAmelCase_								=     image_embeds.repeat_interleave(__lowercase ,					dim=0		)
							UpperCAmelCase_								=     negative_image_embeds.repeat_interleave(__lowercase ,					dim=0		)

						UpperCAmelCase_								=     torch.cat([negative_image_embeds, image_embeds] ,					dim=0		).to(
						    dtype=prompt_embeds.dtype ,					device=__lowercase		)

						self.scheduler.set_timesteps(__lowercase ,					device=__lowercase		)
						UpperCAmelCase_								=     self.scheduler.timesteps

						UpperCAmelCase_								=     self.unet.config.in_channels

						UpperCAmelCase_						,	UpperCAmelCase_								=     get_new_h_w(__lowercase ,					__lowercase ,					self.movq_scale_factor		)

						# create initial latent
						UpperCAmelCase_								=     self.prepare_latents(
						    (batch_size, num_channels_latents, height, width) ,					text_encoder_hidden_states.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
							UpperCAmelCase_								=     torch.cat([latents] * 2		) if do_classifier_free_guidance else latents

							UpperCAmelCase_								=     {"""text_embeds""": prompt_embeds, """image_embeds""": image_embeds}
							UpperCAmelCase_								=     self.unet(
							    sample=__lowercase ,					timestep=__lowercase ,					encoder_hidden_states=__lowercase ,					added_cond_kwargs=__lowercase ,					return_dict=__lowercase ,					)[0]

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

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

						if output_type == "pil":
							UpperCAmelCase_								=     self.numpy_to_pil(__lowercase		)

						if not return_dict:
							return (image,)

						return ImagePipelineOutput(images=__lowercase		)









import warnings
from pathlib import Path
from typing import List, Tuple, Union

import fire
from torch import nn

from transformers import AutoModelForSeqaSeqLM, AutoTokenizer, PreTrainedModel
from transformers.utils import logging


UpperCamelCase__   : Tuple					  = logging.get_logger(__name__)




def    A_(			A				,      A				,      A   ):
	UpperCAmelCase_								=     nn.ModuleList([src_layers[i] for i in layers_to_copy]   )
	assert len(A   ) == len(A   ), f"""{len(A   )} != {len(A   )}"""
	dest_layers.load_state_dict(layers_to_copy.state_dict()   )


UpperCamelCase__   : str					  = {
    # maps  num layers in teacher -> num_layers in student -> which teacher layers to copy.
    # 12: bart, 16: pegasus, 6: marian/Helsinki-NLP
    12: {
        1: [0],  # This says that if the teacher has 12 layers and the student has 1, copy layer 0 of the teacher
        2: [0, 6],
        3: [0, 6, 11],
        4: [0, 4, 8, 11],
        6: [0, 2, 4, 7, 9, 11],
        9: [0, 1, 2, 4, 5, 7, 9, 10, 11],
        12: list(range(12)),
    },
    16: {  # maps  num layers in student -> which teacher layers to copy
        1: [0],
        2: [0, 15],
        3: [0, 8, 15],
        4: [0, 5, 10, 15],
        6: [0, 3, 6, 9, 12, 15],
        8: [0, 2, 4, 6, 8, 10, 12, 15],
        9: [0, 1, 3, 5, 7, 9, 11, 13, 15],
        12: [0, 1, 2, 3, 4, 5, 6, 7, 9, 11, 13, 15],
        16: list(range(16)),
    },
    6: {1: [0], 2: [0, 5], 3: [0, 2, 5], 4: [0, 1, 3, 5], 6: list(range(6))},
}
UpperCamelCase__   : Any					  = {
    # maps  num layers in student -> which teacher layers to copy.
    6: {1: [5], 2: [3, 5], 3: [1, 4, 5], 4: [1, 2, 4, 5]},
    12: {1: [11], 2: [5, 11], 3: [3, 7, 11], 6: [1, 3, 5, 8, 10, 11]},
    16: {1: [15], 4: [4, 9, 12, 15], 8: [1, 3, 5, 7, 9, 11, 13, 15]},
}




def    A_(			A				,      A   ):
	try:
		UpperCAmelCase_								=     LAYERS_TO_COPY[n_teacher][n_student]
		return val
	except KeyError:
		if n_student != n_teacher:
			warnings.warn(
			    f"""no hardcoded layers to copy for teacher {n_teacher} -> student {n_student}, defaulting to first"""
			    f""" {n_student}"""   )
		return list(range(A   )   )




def    A_(			A				,      A   ):
	if n_student > n_teacher:
		raise ValueError(f"""Cannot perform intermediate supervision for student {n_student} > teacher {n_teacher}"""   )
	elif n_teacher == n_student:
		return list(range(A   )   )
	elif n_student == 1:
		return [n_teacher - 1]
	else:
		return LAYERS_TO_SUPERVISE[n_teacher][n_student]




def    A_(			A				,      A = "student"				,      A = None				,      A = None				,      A=False				,      A=None				,      A=None				,      **A				,      ):
	UpperCAmelCase_								=     """encoder_layers and decoder_layers cannot be both None-- you would just have an identical teacher."""
	assert (e is not None) or (d is not None), _msg
	if isinstance(A				,      A   ):
		AutoTokenizer.from_pretrained(A   ).save_pretrained(A   )  # purely for convenience
		UpperCAmelCase_								=     AutoModelForSeqaSeqLM.from_pretrained(A   ).eval()
	else:
		assert isinstance(A				,      A   ), f"""teacher must be a model or string got type {type(A   )}"""
	UpperCAmelCase_								=     teacher.config.to_diff_dict()

	try:
		UpperCAmelCase_						,	UpperCAmelCase_								=     teacher.config.encoder_layers, teacher.config.decoder_layers
		if e is None:
			UpperCAmelCase_								=     teacher_e
		if d is None:
			UpperCAmelCase_								=     teacher_d
		init_kwargs.update({"""encoder_layers""": e, """decoder_layers""": d}   )
	except AttributeError:  # T5
		if hasattr(teacher.config				,      """num_encoder_layers"""   ):
			UpperCAmelCase_						,	UpperCAmelCase_								=     teacher.config.num_encoder_layers, teacher.config.num_decoder_layers
		else:
			UpperCAmelCase_						,	UpperCAmelCase_								=     teacher.config.num_layers, teacher.config.num_decoder_layers
		if e is None:
			UpperCAmelCase_								=     teacher_e
		if d is None:
			UpperCAmelCase_								=     teacher_d
		if hasattr(teacher.config				,      """num_encoder_layers"""   ):
			init_kwargs.update({"""num_encoder_layers""": e, """num_decoder_layers""": d}   )
		else:
			init_kwargs.update({"""num_layers""": e, """num_decoder_layers""": d}   )

    # Kwargs to instantiate student: teacher kwargs with updated layer numbers + **extra_config_kwargs
	init_kwargs.update(A   )

	# Copy weights
	UpperCAmelCase_								=     teacher.config_class(**A   )
	UpperCAmelCase_								=     AutoModelForSeqaSeqLM.from_config(A   )
	# Start by copying the full teacher state dict this will copy the first N teacher layers to the student.
	UpperCAmelCase_								=     student.load_state_dict(teacher.state_dict()				,      strict=A   )
	assert info.missing_keys == [], info.missing_keys  # every student key should have a teacher keys.

	if copy_first_teacher_layers:  # Our copying is done. We just log and save
		UpperCAmelCase_						,	UpperCAmelCase_								=     list(range(A   )   ), list(range(A   )   )
		logger.info(
		    f"""Copied encoder layers {e_layers_to_copy} and decoder layers {d_layers_to_copy}. Saving them to"""
		    f""" {save_path}"""   )
		student.save_pretrained(A   )
		return student, e_layers_to_copy, d_layers_to_copy

	# Decide which layers of the teacher to copy. Not exactly alternating -- we try to keep first and last layer.
	if e_layers_to_copy is None:
		UpperCAmelCase_								=     pick_layers_to_copy(A				,      A   )
	if d_layers_to_copy is None:
		UpperCAmelCase_								=     pick_layers_to_copy(A				,      A   )

	try:
		if hasattr(
		    A				,      """prophetnet"""   ):  # For ProphetNet, student.model.encoder.layers is called student.prophetnet.encoder.layers
			copy_layers(teacher.prophetnet.encoder.layers				,      student.prophetnet.encoder.layers				,      A   )
			copy_layers(teacher.prophetnet.decoder.layers				,      student.prophetnet.decoder.layers				,      A   )
		else:
			copy_layers(teacher.model.encoder.layers				,      student.model.encoder.layers				,      A   )
			copy_layers(teacher.model.decoder.layers				,      student.model.decoder.layers				,      A   )
	except AttributeError:  # For t5, student.model.encoder.layers is called student.encoder.block
		copy_layers(teacher.encoder.block				,      student.encoder.block				,      A   )
		copy_layers(teacher.decoder.block				,      student.decoder.block				,      A   )
	logger.info(
	    f"""Copied encoder layers {e_layers_to_copy} and decoder layers {d_layers_to_copy}. Saving them to {save_path}"""   )
	UpperCAmelCase_								=     {
	    """teacher_type""": teacher.config.model_type,
	    """copied_encoder_layers""": e_layers_to_copy,
	    """copied_decoder_layers""": d_layers_to_copy,
	}
	student.save_pretrained(A   )
	# Save information about copying for easier reproducibility

	return student, e_layers_to_copy, d_layers_to_copy


if __name__ == "__main__":
				fire.Fire(create_student_by_copying_alternating_layers)


