fix sdxl refiner export (#2133)

huggingface · Dec 20, 2024 · 9840121 · 9840121
1 parent 34b3d8b
commit 9840121
Showing 1 changed file with 5 additions and 1 deletion.
diff --git a/optimum/exporters/utils.py b/optimum/exporters/utils.py
@@ -139,7 +139,11 @@ def _get_submodels_for_export_diffusion(
         # https://github.com/huggingface/diffusers/blob/v0.18.2/src/diffusers/pipelines/stable_diffusion_xl/pipeline_stable_diffusion_xl_img2img.py#L571
         unet.config.requires_aesthetics_score = getattr(pipeline.config, "requires_aesthetics_score", False)
         unet.config.time_cond_proj_dim = getattr(pipeline.unet.config, "time_cond_proj_dim", None)
-        unet.config.text_encoder_projection_dim = pipeline.text_encoder.config.projection_dim
+        unet.config.text_encoder_projection_dim = (
+            pipeline.text_encoder.config.projection_dim
+            if not is_sdxl
+            else pipeline.text_encoder_2.config.projection_dim
+        )
         unet.config.export_model_type = _get_diffusers_submodel_type(unet)
         models_for_export["unet"] = unet