sd-scripts
sd-scripts copied to clipboard
load_target_model support symbolic link file
in train_util.py # determine SD or Diffusers. Maybe because my ckpt is a symbolic link to save disk usage. So that path.isfile returns False, then causing from_pretrained("XXX.ckpt") throw errors.
Maybe cope with os.readlink
def load_target_model(args: argparse.Namespace, weight_dtype):
pth = args.pretrained_model_name_or_path
pth = os.readlink(pth) if os.path.islink(pth) else pth
load_stable_diffusion_format = os.path.isfile(pth) # determine SD or Diffusers
if load_stable_diffusion_format:
print("load StableDiffusion checkpoint", pth)
text_encoder, vae, unet = model_util.load_models_from_stable_diffusion_checkpoint(args.v2, pth)
else:
print("load Diffusers pretrained models", pth)
pipe = StableDiffusionPipeline.from_pretrained(pth, tokenizer=None, safety_checker=None)
text_encoder = pipe.text_encoder
vae = pipe.vae
unet = pipe.unet
del pipe
# VAEを読み込む
if args.vae is not None:
vae = model_util.load_vae(args.vae, weight_dtype)
print("additional VAE loaded")
return text_encoder, vae, unet, load_stable_diffusion_format