From cfb059c5e85aec171c79d4a8499a035f4754ed76 Mon Sep 17 00:00:00 2001 From: Seunghoon Lee Date: Sat, 10 Feb 2024 06:03:39 +0900 Subject: [PATCH] fix onnx sdxl --- modules/onnx_impl/pipelines/__init__.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/modules/onnx_impl/pipelines/__init__.py b/modules/onnx_impl/pipelines/__init__.py index 919284b2d..5dfca6185 100644 --- a/modules/onnx_impl/pipelines/__init__.py +++ b/modules/onnx_impl/pipelines/__init__.py @@ -25,6 +25,8 @@ SUBMODELS_SD = ("text_encoder", "unet", "vae_encoder", "vae_decoder",) SUBMODELS_SDXL = ("text_encoder", "text_encoder_2", "unet", "vae_encoder", "vae_decoder",) SUBMODELS_SDXL_REFINER = ("text_encoder_2", "unet", "vae_encoder", "vae_decoder",) +SUBMODELS_LARGE = ("text_encoder_2", "unet",) + class PipelineBase(TorchCompatibleModule, diffusers.DiffusionPipeline, metaclass=ABCMeta): model_type: str @@ -177,7 +179,7 @@ class OnnxRawPipeline(PipelineBase): onnx.save_model( model, os.path.join(destination, "model.onnx"), - save_as_external_data=submodel == "unet", + save_as_external_data=submodel in SUBMODELS_LARGE, all_tensors_to_one_file=True, location="weights.pb", )