diff --git a/src/diffusers/pipelines/z_image/pipeline_z_image_controlnet.py b/src/diffusers/pipelines/z_image/pipeline_z_image_controlnet.py index 288aee039ed8..4774ef5ea630 100644 --- a/src/diffusers/pipelines/z_image/pipeline_z_image_controlnet.py +++ b/src/diffusers/pipelines/z_image/pipeline_z_image_controlnet.py @@ -19,7 +19,7 @@ from transformers import AutoTokenizer, PreTrainedModel from ...image_processor import PipelineImageInput, VaeImageProcessor -from ...loaders import FromSingleFileMixin +from ...loaders import FromSingleFileMixin, ZImageLoraLoaderMixin from ...models.autoencoders import AutoencoderKL from ...models.controlnets import ZImageControlNetModel from ...models.transformers import ZImageTransformer2DModel @@ -185,7 +185,7 @@ def retrieve_timesteps( return timesteps, num_inference_steps -class ZImageControlNetPipeline(DiffusionPipeline, FromSingleFileMixin): +class ZImageControlNetPipeline(DiffusionPipeline, ZImageLoraLoaderMixin, FromSingleFileMixin): model_cpu_offload_seq = "text_encoder->transformer->vae" _optional_components = [] _callback_tensor_inputs = ["latents", "prompt_embeds"] diff --git a/src/diffusers/pipelines/z_image/pipeline_z_image_controlnet_inpaint.py b/src/diffusers/pipelines/z_image/pipeline_z_image_controlnet_inpaint.py index d278da912b5a..c311e1e4e546 100644 --- a/src/diffusers/pipelines/z_image/pipeline_z_image_controlnet_inpaint.py +++ b/src/diffusers/pipelines/z_image/pipeline_z_image_controlnet_inpaint.py @@ -20,7 +20,7 @@ from transformers import AutoTokenizer, PreTrainedModel from ...image_processor import PipelineImageInput, VaeImageProcessor -from ...loaders import FromSingleFileMixin +from ...loaders import FromSingleFileMixin, ZImageLoraLoaderMixin from ...models.autoencoders import AutoencoderKL from ...models.controlnets import ZImageControlNetModel from ...models.transformers import ZImageTransformer2DModel @@ -185,7 +185,7 @@ def retrieve_timesteps( return timesteps, num_inference_steps -class ZImageControlNetInpaintPipeline(DiffusionPipeline, FromSingleFileMixin): +class ZImageControlNetInpaintPipeline(DiffusionPipeline, ZImageLoraLoaderMixin, FromSingleFileMixin): model_cpu_offload_seq = "text_encoder->transformer->vae" _optional_components = [] _callback_tensor_inputs = ["latents", "prompt_embeds"]