Spaces:
Running
Running
| import argparse | |
| from transformers import CLIPImageProcessor, CLIPVisionModelWithProjection | |
| from diffusers import UnCLIPImageVariationPipeline, UnCLIPPipeline | |
| if __name__ == "__main__": | |
| parser = 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.", | |
| ) | |
| args = parser.parse_args() | |
| txt2img = UnCLIPPipeline.from_pretrained(args.txt2img_unclip) | |
| feature_extractor = CLIPImageProcessor() | |
| image_encoder = CLIPVisionModelWithProjection.from_pretrained("openai/clip-vit-large-patch14") | |
| img2img = UnCLIPImageVariationPipeline( | |
| decoder=txt2img.decoder, | |
| text_encoder=txt2img.text_encoder, | |
| tokenizer=txt2img.tokenizer, | |
| text_proj=txt2img.text_proj, | |
| feature_extractor=feature_extractor, | |
| image_encoder=image_encoder, | |
| super_res_first=txt2img.super_res_first, | |
| super_res_last=txt2img.super_res_last, | |
| decoder_scheduler=txt2img.decoder_scheduler, | |
| super_res_scheduler=txt2img.super_res_scheduler, | |
| ) | |
| img2img.save_pretrained(args.dump_path) | |