@@ -169,6 +169,8 @@ class StableDiffusionXLInstructPix2PixPipeline(
169
169
Whether to use the [invisible_watermark library](https://github.com/ShieldMnt/invisible-watermark/) to
170
170
watermark output images. If not defined, it will default to True if the package is installed, otherwise no
171
171
watermarker will be used.
172
+ is_cosxl_edit (`bool`, *optional*):
173
+ When set the image latents are scaled.
172
174
"""
173
175
174
176
model_cpu_offload_seq = "text_encoder->text_encoder_2->unet->vae"
@@ -185,6 +187,7 @@ def __init__(
185
187
scheduler : KarrasDiffusionSchedulers ,
186
188
force_zeros_for_empty_prompt : bool = True ,
187
189
add_watermarker : Optional [bool ] = None ,
190
+ is_cosxl_edit : Optional [bool ] = False ,
188
191
):
189
192
super ().__init__ ()
190
193
@@ -201,6 +204,7 @@ def __init__(
201
204
self .vae_scale_factor = 2 ** (len (self .vae .config .block_out_channels ) - 1 )
202
205
self .image_processor = VaeImageProcessor (vae_scale_factor = self .vae_scale_factor )
203
206
self .default_sample_size = self .unet .config .sample_size
207
+ self .is_cosxl_edit = is_cosxl_edit
204
208
205
209
add_watermarker = add_watermarker if add_watermarker is not None else is_invisible_watermark_available ()
206
210
@@ -551,6 +555,9 @@ def prepare_image_latents(
551
555
if image_latents .dtype != self .vae .dtype :
552
556
image_latents = image_latents .to (dtype = self .vae .dtype )
553
557
558
+ if self .is_cosxl_edit :
559
+ image_latents = image_latents * self .vae .config .scaling_factor
560
+
554
561
return image_latents
555
562
556
563
# Copied from diffusers.pipelines.stable_diffusion_xl.pipeline_stable_diffusion_xl.StableDiffusionXLPipeline._get_add_time_ids
0 commit comments