| 888 | |
| 889 | |
| 890 | def text2img(*, args, checkpoint_map_location): |
| 891 | print("loading text2img") |
| 892 | |
| 893 | text2img_checkpoint = torch.load(args.text2img_checkpoint_path, map_location=checkpoint_map_location) |
| 894 | |
| 895 | unet_model = unet_model_from_original_config() |
| 896 | |
| 897 | unet_diffusers_checkpoint = unet_original_checkpoint_to_diffusers_checkpoint(unet_model, text2img_checkpoint) |
| 898 | |
| 899 | del text2img_checkpoint |
| 900 | |
| 901 | load_checkpoint_to_model(unet_diffusers_checkpoint, unet_model, strict=True) |
| 902 | |
| 903 | print("done loading text2img") |
| 904 | |
| 905 | return unet_model |
| 906 | |
| 907 | |
| 908 | def inpaint_text2img(*, args, checkpoint_map_location): |