fix pretrained model download to work with xl

This commit is contained in:
Lincoln Stein 2023-07-18 21:10:33 -04:00
parent e943913f58
commit 39c14eb2ac

View File

@ -10,7 +10,7 @@ from tempfile import TemporaryDirectory
from typing import List, Dict, Callable, Union, Set from typing import List, Dict, Callable, Union, Set
import requests import requests
from diffusers import StableDiffusionPipeline from diffusers import DiffusionPipeline
from diffusers import logging as dlogging from diffusers import logging as dlogging
from huggingface_hub import hf_hub_url, HfFolder, HfApi from huggingface_hub import hf_hub_url, HfFolder, HfApi
from omegaconf import OmegaConf from omegaconf import OmegaConf
@ -310,6 +310,8 @@ class ModelInstall(object):
if key := self.reverse_paths.get(path_name): if key := self.reverse_paths.get(path_name):
(name, base, mtype) = ModelManager.parse_key(key) (name, base, mtype) = ModelManager.parse_key(key)
return name return name
elif location.is_dir():
return location.name
else: else:
return location.stem return location.stem
@ -365,7 +367,7 @@ class ModelInstall(object):
model = None model = None
for revision in revisions: for revision in revisions:
try: try:
model = StableDiffusionPipeline.from_pretrained(repo_id,revision=revision,safety_checker=None) model = DiffusionPipeline.from_pretrained(repo_id,revision=revision,safety_checker=None)
except: # most errors are due to fp16 not being present. Fix this to catch other errors except: # most errors are due to fp16 not being present. Fix this to catch other errors
pass pass
if model: if model: