123 lines
4.0 KiB
Python
123 lines
4.0 KiB
Python
from logging import getLogger
|
|
from os import path
|
|
from typing import Dict, Optional, Tuple
|
|
|
|
import onnx
|
|
import torch
|
|
from diffusers import AutoencoderKL, StableDiffusionXLPipeline
|
|
from onnx.shape_inference import infer_shapes_path
|
|
from onnxruntime.transformers.float16 import convert_float_to_float16
|
|
from optimum.exporters.onnx import main_export
|
|
|
|
from ...constants import ONNX_MODEL, ONNX_WEIGHTS
|
|
from ...utils import run_gc
|
|
from ..client import fetch_model
|
|
from ..utils import RESOLVE_FORMATS, ConversionContext, check_ext
|
|
|
|
logger = getLogger(__name__)
|
|
|
|
|
|
@torch.no_grad()
|
|
def convert_diffusion_diffusers_xl(
|
|
conversion: ConversionContext,
|
|
model: Dict,
|
|
format: Optional[str],
|
|
) -> Tuple[bool, str]:
|
|
"""
|
|
From https://github.com/huggingface/diffusers/blob/main/scripts/convert_stable_diffusion_checkpoint_to_onnx.py
|
|
"""
|
|
name = str(model.get("name")).strip()
|
|
source = model.get("source")
|
|
replace_vae = model.get("vae", None)
|
|
|
|
device = conversion.training_device
|
|
dtype = conversion.torch_dtype()
|
|
logger.debug("using Torch dtype %s for pipeline", dtype)
|
|
|
|
dest_path = path.join(conversion.model_path, name)
|
|
model_index = path.join(dest_path, "model_index.json")
|
|
model_hash = path.join(dest_path, "hash.txt")
|
|
|
|
# diffusers go into a directory rather than .onnx file
|
|
logger.info(
|
|
"converting Stable Diffusion XL model %s: %s -> %s/", name, source, dest_path
|
|
)
|
|
|
|
if path.exists(dest_path) and path.exists(model_index):
|
|
logger.info("ONNX model already exists, skipping conversion")
|
|
|
|
if "hash" in model and not path.exists(model_hash):
|
|
logger.info("ONNX model does not have hash file, adding one")
|
|
with open(model_hash, "w") as f:
|
|
f.write(model["hash"])
|
|
|
|
return (False, dest_path)
|
|
|
|
cache_path = fetch_model(conversion, name, model["source"], format=format)
|
|
# safetensors -> diffusers directory with torch models
|
|
temp_path = path.join(conversion.cache_path, f"{name}-torch")
|
|
|
|
if format == "safetensors":
|
|
pipeline = StableDiffusionXLPipeline.from_single_file(
|
|
cache_path, use_safetensors=True
|
|
)
|
|
else:
|
|
pipeline = StableDiffusionXLPipeline.from_pretrained(cache_path)
|
|
|
|
if replace_vae is not None:
|
|
vae_path = path.join(conversion.model_path, replace_vae)
|
|
vae_file = check_ext(vae_path, RESOLVE_FORMATS)
|
|
if vae_file[0]:
|
|
logger.debug("loading VAE from single tensor file: %s", vae_path)
|
|
pipeline.vae = AutoencoderKL.from_single_file(vae_path)
|
|
else:
|
|
logger.debug("loading pretrained VAE from path: %s", replace_vae)
|
|
pipeline.vae = AutoencoderKL.from_pretrained(replace_vae)
|
|
|
|
if path.exists(temp_path):
|
|
logger.debug("torch model already exists for %s: %s", source, temp_path)
|
|
else:
|
|
logger.debug("exporting torch model for %s: %s", source, temp_path)
|
|
pipeline.save_pretrained(temp_path)
|
|
|
|
# GC temporary pipeline
|
|
del pipeline
|
|
run_gc()
|
|
|
|
# directory -> onnx using optimum exporters
|
|
main_export(
|
|
temp_path,
|
|
output=dest_path,
|
|
task="stable-diffusion-xl",
|
|
device=device,
|
|
fp16=conversion.has_optimization(
|
|
"torch-fp16"
|
|
), # optimum's fp16 mode only works on CUDA or ROCm
|
|
framework="pt",
|
|
)
|
|
|
|
if "hash" in model:
|
|
logger.debug("adding hash file to ONNX model")
|
|
with open(model_hash, "w") as f:
|
|
f.write(model["hash"])
|
|
|
|
if conversion.half:
|
|
unet_path = path.join(dest_path, "unet", ONNX_MODEL)
|
|
infer_shapes_path(unet_path)
|
|
unet = onnx.load(unet_path)
|
|
opt_model = convert_float_to_float16(
|
|
unet,
|
|
disable_shape_infer=True,
|
|
force_fp16_initializers=True,
|
|
keep_io_types=True,
|
|
)
|
|
onnx.save_model(
|
|
opt_model,
|
|
unet_path,
|
|
save_as_external_data=True,
|
|
all_tensors_to_one_file=True,
|
|
location=ONNX_WEIGHTS,
|
|
)
|
|
|
|
return (True, dest_path)
|