




from typing import TYPE_CHECKING

from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_torch_available,
)


__lowerCAmelCase		:				List[Any]       			=		{
    '''configuration_mega''': ['''MEGA_PRETRAINED_CONFIG_ARCHIVE_MAP''', '''MegaConfig''', '''MegaOnnxConfig'''],
}

try:
		if not is_torch_available():
				raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
		pass
else:
		__lowerCAmelCase		:				List[str]       			=		[
		    '''MEGA_PRETRAINED_MODEL_ARCHIVE_LIST''',
		    '''MegaForCausalLM''',
		    '''MegaForMaskedLM''',
		    '''MegaForMultipleChoice''',
		    '''MegaForQuestionAnswering''',
		    '''MegaForSequenceClassification''',
		    '''MegaForTokenClassification''',
		    '''MegaModel''',
		    '''MegaPreTrainedModel''',
		]

if TYPE_CHECKING:
		from .configuration_mega import MEGA_PRETRAINED_CONFIG_ARCHIVE_MAP, MegaConfig, MegaOnnxConfig

		try:
				if not is_torch_available():
						raise OptionalDependencyNotAvailable()
		except OptionalDependencyNotAvailable:
				pass
		else:
				from .modeling_mega import (
				    MEGA_PRETRAINED_MODEL_ARCHIVE_LIST,
				    MegaForCausalLM,
				    MegaForMaskedLM,
				    MegaForMultipleChoice,
				    MegaForQuestionAnswering,
				    MegaForSequenceClassification,
				    MegaForTokenClassification,
				    MegaModel,
				    MegaPreTrainedModel,
				)

else:
		import sys

		__lowerCAmelCase		:				Any       			=		_LazyModule(__name__, globals()['__file__'], _import_structure, module_spec=__spec__)














"""simple docstring"""



import argparse

from transformers import CLIPImageProcessor, CLIPVisionModelWithProjection

from diffusers import UnCLIPImageVariationPipeline, UnCLIPPipeline


if __name__ == "__main__":
						__A		:       str  				=		argparse.ArgumentParser()

						parser.add_argument('''--dump_path''', default=None, type=str, required=True, help='''Path to the output model.''')

						parser.add_argument(
						    '''--txt2img_unclip''',
						    default='''kakaobrain/karlo-v1-alpha''',
						    type=str,
						    required=False,
						    help='''The pretrained txt2img unclip.''',
						)

						__A		:       str  				=		parser.parse_args()

						__A		:       List[Any]  				=		UnCLIPPipeline.from_pretrained(args.txtaimg_unclip)

						__A		:       Dict  				=		CLIPImageProcessor()
						__A		:       Union[str, Any]  				=		CLIPVisionModelWithProjection.from_pretrained('''openai/clip-vit-large-patch14''')

						__A		:       List[str]  				=		UnCLIPImageVariationPipeline(
						    decoder=txtaimg.decoder,
						    text_encoder=txtaimg.text_encoder,
						    tokenizer=txtaimg.tokenizer,
						    text_proj=txtaimg.text_proj,
						    feature_extractor=feature_extractor,
						    image_encoder=image_encoder,
						    super_res_first=txtaimg.super_res_first,
						    super_res_last=txtaimg.super_res_last,
						    decoder_scheduler=txtaimg.decoder_scheduler,
						    super_res_scheduler=txtaimg.super_res_scheduler,
						)

						imgaimg.save_pretrained(args.dump_path)



