sd-scripts icon indicating copy to clipboard operation
sd-scripts copied to clipboard

load_target_model support symbolic link file

Open TkskKurumi opened this issue 2 years ago • 0 comments

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

TkskKurumi avatar Feb 13 '23 05:02 TkskKurumi