Skip to content

Commit 1e216be

Browse files
patil-surajanton-lpcuencapatrickvonplaten
authored
make scaling factor a config arg of vae/vqvae (huggingface#1860)
* make scaling factor cnfig arg of vae * fix * make flake happy * fix ldm * fix upscaler * qualirty * Apply suggestions from code review Co-authored-by: Anton Lozhkov <anton@huggingface.co> Co-authored-by: Pedro Cuenca <pedro@huggingface.co> Co-authored-by: Patrick von Platen <patrick.v.platen@gmail.com> * solve conflicts, addres some comments * examples * examples min version * doc * fix type * typo * Update src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_upscale.py Co-authored-by: Pedro Cuenca <pedro@huggingface.co> * remove duplicate line * Apply suggestions from code review Co-authored-by: Patrick von Platen <patrick.v.platen@gmail.com> Co-authored-by: Anton Lozhkov <anton@huggingface.co> Co-authored-by: Pedro Cuenca <pedro@huggingface.co> Co-authored-by: Patrick von Platen <patrick.v.platen@gmail.com>
1 parent 915a563 commit 1e216be

39 files changed

Lines changed: 95 additions & 55 deletions

File tree

examples/community/clip_guided_stable_diffusion.py

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -150,7 +150,7 @@ def cond_fn(
150150
else:
151151
raise ValueError(f"scheduler type {type(self.scheduler)} not supported")
152152

153-
sample = 1 / 0.18215 * sample
153+
sample = 1 / self.vae.config.scaling_factor * sample
154154
image = self.vae.decode(sample).sample
155155
image = (image / 2 + 0.5).clamp(0, 1)
156156

@@ -336,7 +336,7 @@ def __call__(
336336
latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs).prev_sample
337337

338338
# scale and decode the image latents with vae
339-
latents = 1 / 0.18215 * latents
339+
latents = 1 / self.vae.config.scaling_factor * latents
340340
image = self.vae.decode(latents).sample
341341

342342
image = (image / 2 + 0.5).clamp(0, 1)

examples/dreambooth/train_dreambooth.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -803,7 +803,7 @@ def main(args):
803803
with accelerator.accumulate(unet):
804804
# Convert images to latent space
805805
latents = vae.encode(batch["pixel_values"].to(dtype=weight_dtype)).latent_dist.sample()
806-
latents = latents * 0.18215
806+
latents = latents * vae.config.scaling_factor
807807

808808
# Sample noise that we'll add to the latents
809809
noise = torch.randn_like(latents)

examples/dreambooth/train_dreambooth_flax.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -533,7 +533,7 @@ def compute_loss(params):
533533
latents = vae_outputs.latent_dist.sample(sample_rng)
534534
# (NHWC) -> (NCHW)
535535
latents = jnp.transpose(latents, (0, 3, 1, 2))
536-
latents = latents * 0.18215
536+
latents = latents * vae.config.scaling_factor
537537

538538
# Sample noise that we'll add to the latents
539539
noise_rng, timestep_rng = jax.random.split(sample_rng)

examples/dreambooth/train_dreambooth_lora.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -853,7 +853,7 @@ def main(args):
853853
with accelerator.accumulate(unet):
854854
# Convert images to latent space
855855
latents = vae.encode(batch["pixel_values"].to(dtype=weight_dtype)).latent_dist.sample()
856-
latents = latents * 0.18215
856+
latents = latents * vae.config.scaling_factor
857857

858858
# Sample noise that we'll add to the latents
859859
noise = torch.randn_like(latents)

examples/research_projects/colossalai/train_dreambooth_colossalai.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -607,7 +607,7 @@ def collate_fn(examples):
607607
optimizer.zero_grad()
608608

609609
latents = vae.encode(batch["pixel_values"].to(dtype=weight_dtype)).latent_dist.sample()
610-
latents = latents * 0.18215
610+
latents = latents * vae.config.scaling_factor
611611

612612
# Sample noise that we'll add to the latents
613613
noise = torch.randn_like(latents)

examples/research_projects/dreambooth_inpaint/train_dreambooth_inpaint.py

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -33,7 +33,7 @@
3333

3434

3535
# Will error if the minimal version of diffusers is not installed. Remove at your own risks.
36-
check_min_version("0.10.0.dev0")
36+
check_min_version("0.13.0.dev0")
3737

3838
logger = get_logger(__name__)
3939

@@ -699,13 +699,13 @@ def collate_fn(examples):
699699
# Convert images to latent space
700700

701701
latents = vae.encode(batch["pixel_values"].to(dtype=weight_dtype)).latent_dist.sample()
702-
latents = latents * 0.18215
702+
latents = latents * vae.config.scaling_factor
703703

704704
# Convert masked images to latent space
705705
masked_latents = vae.encode(
706706
batch["masked_images"].reshape(batch["pixel_values"].shape).to(dtype=weight_dtype)
707707
).latent_dist.sample()
708-
masked_latents = masked_latents * 0.18215
708+
masked_latents = masked_latents * vae.config.scaling_factor
709709

710710
masks = batch["masks"]
711711
# resize the mask to latents shape as we concatenate the mask to the latents

examples/research_projects/intel_opts/textual_inversion/textual_inversion_bf16.py

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -51,7 +51,7 @@
5151

5252

5353
# Will error if the minimal version of diffusers is not installed. Remove at your own risks.
54-
check_min_version("0.10.0.dev0")
54+
check_min_version("0.13.0.dev0")
5555

5656

5757
logger = get_logger(__name__)
@@ -555,7 +555,7 @@ def main():
555555
with accelerator.accumulate(text_encoder):
556556
# Convert images to latent space
557557
latents = vae.encode(batch["pixel_values"]).latent_dist.sample().detach()
558-
latents = latents * 0.18215
558+
latents = latents * vae.config.scaling_factor
559559

560560
# Sample noise that we'll add to the latents
561561
noise = torch.randn(latents.shape).to(latents.device)

examples/research_projects/multi_subject_dreambooth/train_multi_subject_dreambooth.py

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -31,7 +31,7 @@
3131

3232

3333
# Will error if the minimal version of diffusers is not installed. Remove at your own risks.
34-
check_min_version("0.10.0.dev0")
34+
check_min_version("0.13.0.dev0")
3535

3636
logger = get_logger(__name__)
3737

@@ -788,7 +788,7 @@ def main(args):
788788
with accelerator.accumulate(unet):
789789
# Convert images to latent space
790790
latents = vae.encode(batch["pixel_values"].to(dtype=weight_dtype)).latent_dist.sample()
791-
latents = latents * 0.18215
791+
latents = latents * vae.config.scaling_factor
792792

793793
# Sample noise that we'll add to the latents
794794
noise = torch.randn_like(latents)

examples/text_to_image/train_text_to_image.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -636,7 +636,7 @@ def collate_fn(examples):
636636
with accelerator.accumulate(unet):
637637
# Convert images to latent space
638638
latents = vae.encode(batch["pixel_values"].to(weight_dtype)).latent_dist.sample()
639-
latents = latents * 0.18215
639+
latents = latents * vae.config.scaling_factor
640640

641641
# Sample noise that we'll add to the latents
642642
noise = torch.randn_like(latents)

examples/text_to_image/train_text_to_image_flax.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -438,7 +438,7 @@ def compute_loss(params):
438438
latents = vae_outputs.latent_dist.sample(sample_rng)
439439
# (NHWC) -> (NCHW)
440440
latents = jnp.transpose(latents, (0, 3, 1, 2))
441-
latents = latents * 0.18215
441+
latents = latents * vae.config.scaling_factor
442442

443443
# Sample noise that we'll add to the latents
444444
noise_rng, timestep_rng = jax.random.split(sample_rng)

0 commit comments

Comments
 (0)