(args)
| 53 | |
| 54 | |
| 55 | def main(args): |
| 56 | tokenizer = T5Tokenizer.from_pretrained("google/t5-v1_1-xxl") |
| 57 | text_encoder = T5EncoderModel.from_pretrained("google/t5-v1_1-xxl") |
| 58 | |
| 59 | feature_extractor = CLIPImageProcessor.from_pretrained("openai/clip-vit-large-patch14") |
| 60 | safety_checker = convert_safety_checker(p_head_path=args.p_head_path, w_head_path=args.w_head_path) |
| 61 | |
| 62 | if args.unet_config is not None and args.unet_checkpoint_path is not None and args.dump_path is not None: |
| 63 | convert_stage_1_pipeline(tokenizer, text_encoder, feature_extractor, safety_checker, args) |
| 64 | |
| 65 | if args.unet_checkpoint_path_stage_2 is not None and args.dump_path_stage_2 is not None: |
| 66 | convert_super_res_pipeline(tokenizer, text_encoder, feature_extractor, safety_checker, args, stage=2) |
| 67 | |
| 68 | if args.unet_checkpoint_path_stage_3 is not None and args.dump_path_stage_3 is not None: |
| 69 | convert_super_res_pipeline(tokenizer, text_encoder, feature_extractor, safety_checker, args, stage=3) |
| 70 | |
| 71 | |
| 72 | def convert_stage_1_pipeline(tokenizer, text_encoder, feature_extractor, safety_checker, args): |
no test coverage detected