Skip to content

Commit 47649f3

Browse files
committed
Add support for DC encoder and fix latent shape errors
1 parent 7a04604 commit 47649f3

File tree

1 file changed

+5
-1
lines changed

1 file changed

+5
-1
lines changed

src/diffusers/pipelines/pixart_alpha/pipeline_pixart_sigma.py

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -220,7 +220,10 @@ def __init__(
220220
tokenizer=tokenizer, text_encoder=text_encoder, vae=vae, transformer=transformer, scheduler=scheduler
221221
)
222222

223-
self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1) if getattr(self, "vae", None) else 8
223+
if hasattr(self.vae.config, 'decoder_block_out_channels'):
224+
self.vae_scale_factor = 2 ** (len(self.vae.config.decoder_block_out_channels) - 1) if getattr(self, "vae", None) else 8
225+
else:
226+
self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1) if getattr(self, "vae", None) else 8
224227
self.image_processor = PixArtImageProcessor(vae_scale_factor=self.vae_scale_factor)
225228

226229
# Copied from diffusers.pipelines.pixart_alpha.pipeline_pixart_alpha.PixArtAlphaPipeline.encode_prompt with 120->300
@@ -855,6 +858,7 @@ def __call__(
855858
noise_pred = noise_pred
856859

857860
# compute previous image: x_t -> x_t-1
861+
latents = latents[:noise_pred.shape[0], :noise_pred.shape[1], :noise_pred.shape[2], :noise_pred.shape[3]]
858862
latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs, return_dict=False)[0]
859863

860864
# call the callback, if provided

0 commit comments

Comments
 (0)