Skip to content

Commit

Permalink
fix(api): make sure diffusion models have a valid prefix
Browse files Browse the repository at this point in the history
  • Loading branch information
ssube committed Dec 9, 2023
1 parent 6e614aa commit 46d9fc0
Show file tree
Hide file tree
Showing 2 changed files with 21 additions and 1 deletion.
20 changes: 20 additions & 0 deletions api/onnx_web/convert/__main__.py
Original file line number Diff line number Diff line change
Expand Up @@ -265,6 +265,20 @@ def fetch_model(
return source, False


DIFFUSION_PREFIX = ["diffusion-", "diffusion/", "diffusion\\", "stable-diffusion-"]


def fix_diffusion_name(name: str):
if not any([name.startswith(prefix) for prefix in DIFFUSION_PREFIX]):
logger.warning(
"diffusion models must have names starting with diffusion- to be recognized by the server: %s does not match",
name,
)
return f"diffusion-{name}"

return name


def convert_models(conversion: ConversionContext, args, models: Models):
model_errors = []

Expand Down Expand Up @@ -351,6 +365,12 @@ def convert_models(conversion: ConversionContext, args, models: Models):
if name in args.skip:
logger.info("skipping model: %s", name)
else:
# fix up entries with missing prefixes
name = fix_diffusion_name(name)
if name != model["name"]:
# update the model in-memory if the name changed
model["name"] = name

model_format = source_format(model)

try:
Expand Down
2 changes: 1 addition & 1 deletion api/onnx_web/convert/diffusion/diffusion_xl.py
Original file line number Diff line number Diff line change
Expand Up @@ -64,7 +64,7 @@ def convert_diffusion_diffusers_xl(

if replace_vae is not None:
vae_path = path.join(conversion.model_path, replace_vae)
if check_ext(replace_vae, RESOLVE_FORMATS):
if check_ext(vae_path, RESOLVE_FORMATS):
logger.debug("loading VAE from single tensor file: %s", vae_path)
pipeline.vae = AutoencoderKL.from_single_file(vae_path)
else:
Expand Down

0 comments on commit 46d9fc0

Please sign in to comment.