


'''simple docstring'''






import random
import unittest

import numpy as np
import torch
from transformers import CLIPTextConfig, CLIPTextModel, CLIPTokenizer

from diffusers import (
    AutoencoderKL,
    DDIMScheduler,
    UNetaDConditionModel,
    VideoToVideoSDPipeline,
)
from diffusers.utils import floats_tensor, is_xformers_available, skip_mps
from diffusers.utils.testing_utils import enable_full_determinism, slow, torch_device

from ..pipeline_params import (
    TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS,
    TEXT_GUIDED_IMAGE_VARIATION_PARAMS,
)
from ..test_pipelines_common import PipelineTesterMixin


enable_full_determinism()
@skip_mps
class 						__UpperCamelCase	(						lowercase__						,  unittest.TestCase							):
							lowercase			:		Optional[int]   		=  VideoToVideoSDPipeline
							lowercase			:		List[Any]   		=  TEXT_GUIDED_IMAGE_VARIATION_PARAMS.union({'video'}							) - {'image', 'width', 'height'}
							lowercase			:		List[Any]   		=  TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS.union({'video'}							) - {'image'}
							lowercase			:		str   		=  PipelineTesterMixin.required_optional_params - {'latents'}
							lowercase			:		Optional[int]   		=  False

							# No `output_type`.
							lowercase			:		List[str]   		=  frozenset(
							    [
							        'num_inference_steps',
							        'generator',
							        'latents',
							        'return_dict',
							        'callback',
							        'callback_steps',
							    ]							)





							def 		a__   (  self		:Tuple			):
									torch.manual_seed(0			)
									snake_case_		:							Any			       =						UNetaDConditionModel(
									    block_out_channels=(3_2, 6_4, 6_4, 6_4)  ,layers_per_block=2  ,sample_size=3_2  ,in_channels=4  ,out_channels=4  ,down_block_types=("""CrossAttnDownBlock3D""", """CrossAttnDownBlock3D""", """CrossAttnDownBlock3D""", """DownBlock3D""")  ,up_block_types=("""UpBlock3D""", """CrossAttnUpBlock3D""", """CrossAttnUpBlock3D""", """CrossAttnUpBlock3D""")  ,cross_attention_dim=3_2  ,attention_head_dim=4  ,)
									snake_case_		:							List[Any]			       =						DDIMScheduler(
									    beta_start=0.0_00_85  ,beta_end=0.0_12  ,beta_schedule="""scaled_linear"""  ,clip_sample=__a  ,set_alpha_to_one=__a  ,)
									torch.manual_seed(0			)
									snake_case_		:							List[str]			       =						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  ,sample_size=1_2_8  ,)
									torch.manual_seed(0			)
									snake_case_		:							Tuple			       =						CLIPTextConfig(
									    bos_token_id=0  ,eos_token_id=2  ,hidden_size=3_2  ,intermediate_size=3_7  ,layer_norm_eps=1E-0_5  ,num_attention_heads=4  ,num_hidden_layers=5  ,pad_token_id=1  ,vocab_size=1_0_0_0  ,hidden_act="""gelu"""  ,projection_dim=5_1_2  ,)
									snake_case_		:							List[Any]			       =						CLIPTextModel(__a			)
									snake_case_		:							Optional[Any]			       =						CLIPTokenizer.from_pretrained("""hf-internal-testing/tiny-random-clip"""			)

									snake_case_		:							Dict			       =						{
									    """unet""": unet,
									    """scheduler""": scheduler,
									    """vae""": vae,
									    """text_encoder""": text_encoder,
									    """tokenizer""": tokenizer,
									}
									return components





							def 		a__   (  self		:List[Any]  ,_UpperCamelCase		:List[str]  ,_UpperCamelCase		:str=0			):
									snake_case_		:							Dict			       =						floats_tensor((1, 3, 3, 3_2, 3_2)  ,rng=random.Random(__a			)			).to(__a			)

									if str(__a			).startswith("""mps"""			):
											snake_case_		:							int			       =						torch.manual_seed(__a			)
									else:
											snake_case_		:							Any			       =						torch.Generator(device=__a			).manual_seed(__a			)
									snake_case_		:							int			       =						{
									    """prompt""": """A painting of a squirrel eating a burger""",
									    """video""": video,
									    """generator""": generator,
									    """num_inference_steps""": 2,
									    """guidance_scale""": 6.0,
									    """output_type""": """pt""",
									}
									return inputs





							def 		a__   (  self		:str			):
									snake_case_		:							Union[str, Any]			       =						"""cpu"""  # ensure determinism for the device-dependent torch.Generator
									snake_case_		:							Any			       =						self.get_dummy_components()
									snake_case_		:							List[str]			       =						VideoToVideoSDPipeline(**__a			)
									snake_case_		:							Union[str, Any]			       =						sd_pipe.to(__a			)
									sd_pipe.set_progress_bar_config(disable=__a			)

									snake_case_		:							Dict			       =						self.get_dummy_inputs(__a			)
									snake_case_		:							Optional[Any]			       =						"""np"""
									snake_case_		:							List[str]			       =						sd_pipe(**__a			).frames
									snake_case_		:							int			       =						frames[0][-3:, -3:, -1]

									assert frames[0].shape == (3_2, 3_2, 3)
									snake_case_		:							Optional[Any]			       =						np.array([1_0_6, 1_1_7, 1_1_3, 1_7_4, 1_3_7, 1_1_2, 1_4_8, 1_5_1, 1_3_1]			)

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





							@unittest.skipIf(
							    torch_device != """cuda""" or not is_xformers_available()  ,reason="""XFormers attention is only available with CUDA and `xformers` installed"""  ,)
							def 		a__   (  self		:Union[str, Any]			):
									self._test_xformers_attention_forwardGenerator_pass(test_mean_pixel_difference=__a  ,expected_max_diff=5E-3			)





							@unittest.skip(reason="""Batching needs to be properly figured out first for this pipeline."""			)
							def 		a__   (  self		:List[str]			):
									pass





							@unittest.skip(reason="""Batching needs to be properly figured out first for this pipeline."""			)
							def 		a__   (  self		:Optional[int]			):
									pass





							@unittest.skip(reason="""`num_images_per_prompt` argument is not supported for this pipeline."""			)
							def 		a__   (  self		:Tuple			):
									pass





							def 		a__   (  self		:Optional[Any]			):
									return super().test_progress_bar()





@slow
@skip_mps
class 						__UpperCamelCase	(						unittest.TestCase							):





							def 		a__   (  self		:List[Any]			):
									snake_case_		:							Union[str, Any]			       =						VideoToVideoSDPipeline.from_pretrained("""cerspense/zeroscope_v2_XL"""  ,torch_dtype=torch.floataa			)
									pipe.enable_model_cpu_offload()

									# 10 frames
									snake_case_		:							str			       =						torch.Generator(device="""cpu"""			).manual_seed(0			)
									snake_case_		:							str			       =						torch.randn((1, 1_0, 3, 1_0_2_4, 5_7_6)  ,generator=__a			)
									snake_case_		:							int			       =						video.to("""cuda"""			)

									snake_case_		:							List[Any]			       =						"""Spiderman is surfing"""

									snake_case_		:							Union[str, Any]			       =						pipe(__a  ,video=__a  ,generator=__a  ,num_inference_steps=3  ,output_type="""pt"""			).frames

									snake_case_		:							int			       =						np.array([-1.0_45_89_84, -1.1_27_92_97, -0.9_66_30_86, -0.91_50_39_06, -0.75_09_76_56]			)
									assert np.abs(video_frames.cpu().numpy()[0, 0, 0, 0, -5:] - expected_array			).sum() < 1E-2



'''simple docstring'''






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


__A	:	List[Any]     =     logging.get_logger(__name__)





def      UpperCAmelCase      (		lowerCamelCase_  :nn.ModuleList				,      lowerCamelCase_  :nn.ModuleList				,      lowerCamelCase_  :List[int] ):







		'''simple docstring'''
		snake_case_		:							Optional[Any]			       =						nn.ModuleList([src_layers[i] for i in layers_to_copy] )
		assert len(lowerCamelCase_ ) == len(lowerCamelCase_ ), F'''{len(lowerCamelCase_ )} != {len(lowerCamelCase_ )}'''
		dest_layers.load_state_dict(layers_to_copy.state_dict() )


__A	:	Optional[int]     =     {
    # 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))},
}
__A	:	Dict     =     {
    # 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      UpperCAmelCase      (		lowerCamelCase_  :Union[str, Any]				,      lowerCamelCase_  :int ):







		'''simple docstring'''
		try:
				snake_case_		:							int			       =						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(lowerCamelCase_ ) )





def      UpperCAmelCase      (		lowerCamelCase_  :List[Any]				,      lowerCamelCase_  :List[Any] ):







		'''simple docstring'''
		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(lowerCamelCase_ ) )
		elif n_student == 1:
				return [n_teacher - 1]
		else:
				return LAYERS_TO_SUPERVISE[n_teacher][n_student]





def      UpperCAmelCase      (		lowerCamelCase_  :Union[str, PreTrainedModel]				,      lowerCamelCase_  :Union[str, Path] = "student"				,      lowerCamelCase_  :Union[int, None] = None				,      lowerCamelCase_  :Union[int, None] = None				,      lowerCamelCase_  :List[str]=False				,      lowerCamelCase_  :Optional[int]=None				,      lowerCamelCase_  :Optional[Any]=None				,      **lowerCamelCase_  :Dict				,      ):







		'''simple docstring'''
		snake_case_		:							Optional[Any]			       =						"""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(lowerCamelCase_				,      lowerCamelCase_ ):
				AutoTokenizer.from_pretrained(lowerCamelCase_ ).save_pretrained(lowerCamelCase_ )  # purely for convenience
				snake_case_		:							List[Any]			       =						AutoModelForSeqaSeqLM.from_pretrained(lowerCamelCase_ ).eval()
		else:
				assert isinstance(lowerCamelCase_				,      lowerCamelCase_ ), F'''teacher must be a model or string got type {type(lowerCamelCase_ )}'''
		snake_case_		:							Any			       =						teacher.config.to_diff_dict()

		try:
				snake_case_       ,      snake_case_		:							List[str]			       =						teacher.config.encoder_layers, teacher.config.decoder_layers
				if e is None:
						snake_case_		:							Dict			       =						teacher_e
				if d is None:
						snake_case_		:							List[Any]			       =						teacher_d
				init_kwargs.update({"""encoder_layers""": e, """decoder_layers""": d} )
		except AttributeError:  # T5
				if hasattr(teacher.config				,      """num_encoder_layers""" ):
						snake_case_       ,      snake_case_		:							Tuple			       =						teacher.config.num_encoder_layers, teacher.config.num_decoder_layers
				else:
						snake_case_       ,      snake_case_		:							Union[str, Any]			       =						teacher.config.num_layers, teacher.config.num_decoder_layers
				if e is None:
						snake_case_		:							Optional[int]			       =						teacher_e
				if d is None:
						snake_case_		:							List[str]			       =						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(lowerCamelCase_ )

		# Copy weights
		snake_case_		:							List[Any]			       =						teacher.config_class(**lowerCamelCase_ )
		snake_case_		:							Tuple			       =						AutoModelForSeqaSeqLM.from_config(lowerCamelCase_ )
		# Start by copying the full teacher state dict this will copy the first N teacher layers to the student.
		snake_case_		:							Tuple			       =						student.load_state_dict(teacher.state_dict()				,      strict=lowerCamelCase_ )
		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
				snake_case_       ,      snake_case_		:							List[str]			       =						list(range(lowerCamelCase_ ) ), list(range(lowerCamelCase_ ) )
				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(lowerCamelCase_ )
				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:
				snake_case_		:							List[int]			       =						pick_layers_to_copy(lowerCamelCase_				,      lowerCamelCase_ )
		if d_layers_to_copy is None:
				snake_case_		:							List[int]			       =						pick_layers_to_copy(lowerCamelCase_				,      lowerCamelCase_ )

		try:
				if hasattr(
				    lowerCamelCase_				,      """prophetnet""" ):  # For ProphetNet, student.model.encoder.layers is called student.prophetnet.encoder.layers
						copy_layers(teacher.prophetnet.encoder.layers				,      student.prophetnet.encoder.layers				,      lowerCamelCase_ )
						copy_layers(teacher.prophetnet.decoder.layers				,      student.prophetnet.decoder.layers				,      lowerCamelCase_ )
				else:
						copy_layers(teacher.model.encoder.layers				,      student.model.encoder.layers				,      lowerCamelCase_ )
						copy_layers(teacher.model.decoder.layers				,      student.model.decoder.layers				,      lowerCamelCase_ )
		except AttributeError:  # For t5, student.model.encoder.layers is called student.encoder.block
				copy_layers(teacher.encoder.block				,      student.encoder.block				,      lowerCamelCase_ )
				copy_layers(teacher.decoder.block				,      student.decoder.block				,      lowerCamelCase_ )
		logger.info(
		    F'''Copied encoder layers {e_layers_to_copy} and decoder layers {d_layers_to_copy}. Saving them to {save_path}''' )
		snake_case_		:							Any			       =						{
		    """teacher_type""": teacher.config.model_type,
		    """copied_encoder_layers""": e_layers_to_copy,
		    """copied_decoder_layers""": d_layers_to_copy,
		}
		student.save_pretrained(lowerCamelCase_ )
		# 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)