From e6faf607f71b86bf210cc4eca06555304f445611 Mon Sep 17 00:00:00 2001 From: Sayak Paul Date: Thu, 5 Oct 2023 14:29:00 +0200 Subject: [PATCH 01/25] add: entry for DDPO support. (#5250) * add: entry for DDPO support. * move to training * address steven's comments./ --- docs/source/en/_toctree.yml | 2 ++ docs/source/en/training/ddpo.md | 17 +++++++++++++++++ 2 files changed, 19 insertions(+) create mode 100644 docs/source/en/training/ddpo.md diff --git a/docs/source/en/_toctree.yml b/docs/source/en/_toctree.yml index d95e553bd39a..b8aa71dacbe2 100644 --- a/docs/source/en/_toctree.yml +++ b/docs/source/en/_toctree.yml @@ -106,6 +106,8 @@ title: Custom Diffusion - local: training/t2i_adapters title: T2I-Adapters + - local: training/ddpo + title: Reinforcement learning training with DDPO title: Training - sections: - local: using-diffusers/other-modalities diff --git a/docs/source/en/training/ddpo.md b/docs/source/en/training/ddpo.md new file mode 100644 index 000000000000..1ec961dfdd04 --- /dev/null +++ b/docs/source/en/training/ddpo.md @@ -0,0 +1,17 @@ + + +# Reinforcement learning training with DDPO + +You can fine-tune Stable Diffusion on a reward function via reinforcement learning with the 🤗 TRL library and 🤗 Diffusers. This is done with the Denoising Diffusion Policy Optimization (DDPO) algorithm introduced by Black et al. in [Training Diffusion Models with Reinforcement Learning](https://arxiv.org/abs/2305.13301), which is implemented in 🤗 TRL with the [`~trl.DDPOTrainer`]. + +For more information, check out the [`~trl.DDPOTrainer`] API reference and the [Finetune Stable Diffusion Models with DDPO via TRL](https://huggingface.co/blog/trl-ddpo) blog post. \ No newline at end of file From 02a8d664a2a49289c4058d584b3694f3e4a1b0d2 Mon Sep 17 00:00:00 2001 From: Bagheera <59658056+bghira@users.noreply.github.com> Date: Thu, 5 Oct 2023 11:52:27 -0700 Subject: [PATCH 02/25] =?UTF-8?q?Min-SNR=20Gamma:=20correct=20the=20fix=20?= =?UTF-8?q?for=20SNR=20weighted=20loss=20in=20v-prediction=20=E2=80=A6=20(?= =?UTF-8?q?#5238)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Min-SNR Gamma: correct the fix for SNR weighted loss in v-prediction by adding 1 to SNR rather than the resulting loss weights Co-authored-by: bghira Co-authored-by: Sayak Paul --- examples/controlnet/train_controlnet_flax.py | 13 +++--------- .../train_text_to_image_decoder.py | 20 ++++-------------- .../train_text_to_image_lora_decoder.py | 20 ++++-------------- .../train_text_to_image_lora_prior.py | 20 ++++-------------- .../train_text_to_image_prior.py | 20 ++++-------------- .../text_to_image/train_text_to_image.py | 21 +++++-------------- examples/text_to_image/train_text_to_image.py | 20 ++++-------------- .../text_to_image/train_text_to_image_lora.py | 20 ++++-------------- .../train_text_to_image_lora_sdxl.py | 20 ++++-------------- .../text_to_image/train_text_to_image_sdxl.py | 20 ++++-------------- 10 files changed, 40 insertions(+), 154 deletions(-) diff --git a/examples/controlnet/train_controlnet_flax.py b/examples/controlnet/train_controlnet_flax.py index 34e8c69ff64b..68162d7824ab 100644 --- a/examples/controlnet/train_controlnet_flax.py +++ b/examples/controlnet/train_controlnet_flax.py @@ -907,17 +907,10 @@ def compute_loss(params, minibatch, sample_rng): if args.snr_gamma is not None: snr = jnp.array(compute_snr(timesteps)) - base_weights = jnp.where(snr < args.snr_gamma, snr, jnp.ones_like(snr) * args.snr_gamma) / snr if noise_scheduler.config.prediction_type == "v_prediction": - snr_loss_weights = base_weights + 1 - else: - # Epsilon and sample prediction use the base weights. - snr_loss_weights = base_weights - # For zero-terminal SNR, we have to handle the case where a sigma of Zero results in a Inf value. - # When we run this, the MSE loss weights for this timestep is set unconditionally to 1. - # If we do not run this, the loss value will go to NaN almost immediately, usually within one step. - snr_loss_weights[snr == 0] = 1.0 - + # Velocity objective requires that we add one to SNR values before we divide by them. + snr = snr + 1 + snr_loss_weights = jnp.where(snr < args.snr_gamma, snr, jnp.ones_like(snr) * args.snr_gamma) / snr loss = loss * snr_loss_weights loss = loss.mean() diff --git a/examples/kandinsky2_2/text_to_image/train_text_to_image_decoder.py b/examples/kandinsky2_2/text_to_image/train_text_to_image_decoder.py index 0938d38ba487..4ca95ecebea9 100644 --- a/examples/kandinsky2_2/text_to_image/train_text_to_image_decoder.py +++ b/examples/kandinsky2_2/text_to_image/train_text_to_image_decoder.py @@ -781,25 +781,13 @@ def collate_fn(examples): # Since we predict the noise instead of x_0, the original formulation is slightly changed. # This is discussed in Section 4.2 of the same paper. snr = compute_snr(noise_scheduler, timesteps) - base_weight = ( + if noise_scheduler.config.prediction_type == "v_prediction": + # Velocity objective requires that we add one to SNR values before we divide by them. + snr = snr + 1 + mse_loss_weights = ( torch.stack([snr, args.snr_gamma * torch.ones_like(timesteps)], dim=1).min(dim=1)[0] / snr ) - if noise_scheduler.config.prediction_type == "v_prediction": - # Velocity objective needs to be floored to an SNR weight of one. - mse_loss_weights = base_weight + 1 - else: - # Epsilon and sample both use the same loss weights. - mse_loss_weights = base_weight - - # For zero-terminal SNR, we have to handle the case where a sigma of Zero results in a Inf value. - # When we run this, the MSE loss weights for this timestep is set unconditionally to 1. - # If we do not run this, the loss value will go to NaN almost immediately, usually within one step. - mse_loss_weights[snr == 0] = 1.0 - - # We first calculate the original loss. Then we mean over the non-batch dimensions and - # rebalance the sample-wise losses with their respective loss weights. - # Finally, we take the mean of the rebalanced loss. loss = F.mse_loss(model_pred.float(), target.float(), reduction="none") loss = loss.mean(dim=list(range(1, len(loss.shape)))) * mse_loss_weights loss = loss.mean() diff --git a/examples/kandinsky2_2/text_to_image/train_text_to_image_lora_decoder.py b/examples/kandinsky2_2/text_to_image/train_text_to_image_lora_decoder.py index 91091f4d80fb..19245724ecf5 100644 --- a/examples/kandinsky2_2/text_to_image/train_text_to_image_lora_decoder.py +++ b/examples/kandinsky2_2/text_to_image/train_text_to_image_lora_decoder.py @@ -631,25 +631,13 @@ def collate_fn(examples): # Since we predict the noise instead of x_0, the original formulation is slightly changed. # This is discussed in Section 4.2 of the same paper. snr = compute_snr(noise_scheduler, timesteps) - base_weight = ( + if noise_scheduler.config.prediction_type == "v_prediction": + # Velocity objective requires that we add one to SNR values before we divide by them. + snr = snr + 1 + mse_loss_weights = ( torch.stack([snr, args.snr_gamma * torch.ones_like(timesteps)], dim=1).min(dim=1)[0] / snr ) - if noise_scheduler.config.prediction_type == "v_prediction": - # Velocity objective needs to be floored to an SNR weight of one. - mse_loss_weights = base_weight + 1 - else: - # Epsilon and sample both use the same loss weights. - mse_loss_weights = base_weight - - # For zero-terminal SNR, we have to handle the case where a sigma of Zero results in a Inf value. - # When we run this, the MSE loss weights for this timestep is set unconditionally to 1. - # If we do not run this, the loss value will go to NaN almost immediately, usually within one step. - mse_loss_weights[snr == 0] = 1.0 - - # We first calculate the original loss. Then we mean over the non-batch dimensions and - # rebalance the sample-wise losses with their respective loss weights. - # Finally, we take the mean of the rebalanced loss. loss = F.mse_loss(model_pred.float(), target.float(), reduction="none") loss = loss.mean(dim=list(range(1, len(loss.shape)))) * mse_loss_weights loss = loss.mean() diff --git a/examples/kandinsky2_2/text_to_image/train_text_to_image_lora_prior.py b/examples/kandinsky2_2/text_to_image/train_text_to_image_lora_prior.py index 3099c613df73..7305137218ef 100644 --- a/examples/kandinsky2_2/text_to_image/train_text_to_image_lora_prior.py +++ b/examples/kandinsky2_2/text_to_image/train_text_to_image_lora_prior.py @@ -664,25 +664,13 @@ def collate_fn(examples): # Since we predict the noise instead of x_0, the original formulation is slightly changed. # This is discussed in Section 4.2 of the same paper. snr = compute_snr(noise_scheduler, timesteps) - base_weight = ( + if noise_scheduler.config.prediction_type == "v_prediction": + # Velocity objective requires that we add one to SNR values before we divide by them. + snr = snr + 1 + mse_loss_weights = ( torch.stack([snr, args.snr_gamma * torch.ones_like(timesteps)], dim=1).min(dim=1)[0] / snr ) - if noise_scheduler.config.prediction_type == "v_prediction": - # Velocity objective needs to be floored to an SNR weight of one. - mse_loss_weights = base_weight + 1 - else: - # Epsilon and sample both use the same loss weights. - mse_loss_weights = base_weight - - # For zero-terminal SNR, we have to handle the case where a sigma of Zero results in a Inf value. - # When we run this, the MSE loss weights for this timestep is set unconditionally to 1. - # If we do not run this, the loss value will go to NaN almost immediately, usually within one step. - mse_loss_weights[snr == 0] = 1.0 - - # We first calculate the original loss. Then we mean over the non-batch dimensions and - # rebalance the sample-wise losses with their respective loss weights. - # Finally, we take the mean of the rebalanced loss. loss = F.mse_loss(model_pred.float(), target.float(), reduction="none") loss = loss.mean(dim=list(range(1, len(loss.shape)))) * mse_loss_weights loss = loss.mean() diff --git a/examples/kandinsky2_2/text_to_image/train_text_to_image_prior.py b/examples/kandinsky2_2/text_to_image/train_text_to_image_prior.py index 74fa504345fe..d21eaf3dd0b0 100644 --- a/examples/kandinsky2_2/text_to_image/train_text_to_image_prior.py +++ b/examples/kandinsky2_2/text_to_image/train_text_to_image_prior.py @@ -811,25 +811,13 @@ def collate_fn(examples): # Since we predict the noise instead of x_0, the original formulation is slightly changed. # This is discussed in Section 4.2 of the same paper. snr = compute_snr(noise_scheduler, timesteps) - base_weight = ( + if noise_scheduler.config.prediction_type == "v_prediction": + # Velocity objective requires that we add one to SNR values before we divide by them. + snr = snr + 1 + mse_loss_weights = ( torch.stack([snr, args.snr_gamma * torch.ones_like(timesteps)], dim=1).min(dim=1)[0] / snr ) - if noise_scheduler.config.prediction_type == "v_prediction": - # Velocity objective needs to be floored to an SNR weight of one. - mse_loss_weights = base_weight + 1 - else: - # Epsilon and sample both use the same loss weights. - mse_loss_weights = base_weight - - # For zero-terminal SNR, we have to handle the case where a sigma of Zero results in a Inf value. - # When we run this, the MSE loss weights for this timestep is set unconditionally to 1. - # If we do not run this, the loss value will go to NaN almost immediately, usually within one step. - mse_loss_weights[snr == 0] = 1.0 - - # We first calculate the original loss. Then we mean over the non-batch dimensions and - # rebalance the sample-wise losses with their respective loss weights. - # Finally, we take the mean of the rebalanced loss. loss = F.mse_loss(model_pred.float(), target.float(), reduction="none") loss = loss.mean(dim=list(range(1, len(loss.shape)))) * mse_loss_weights loss = loss.mean() diff --git a/examples/research_projects/onnxruntime/text_to_image/train_text_to_image.py b/examples/research_projects/onnxruntime/text_to_image/train_text_to_image.py index 15c17063bd68..f7100788cde2 100644 --- a/examples/research_projects/onnxruntime/text_to_image/train_text_to_image.py +++ b/examples/research_projects/onnxruntime/text_to_image/train_text_to_image.py @@ -848,24 +848,13 @@ def collate_fn(examples): # Since we predict the noise instead of x_0, the original formulation is slightly changed. # This is discussed in Section 4.2 of the same paper. snr = compute_snr(noise_scheduler, timesteps) - base_weight = ( + if noise_scheduler.config.prediction_type == "v_prediction": + # Velocity objective requires that we add one to SNR values before we divide by them. + snr = snr + 1 + mse_loss_weights = ( torch.stack([snr, args.snr_gamma * torch.ones_like(timesteps)], dim=1).min(dim=1)[0] / snr ) - if noise_scheduler.config.prediction_type == "v_prediction": - # velocity objective prediction requires SNR weights to be floored to a min value of 1. - mse_loss_weights = base_weight + 1 - else: - # Epsilon and sample prediction use the base weights. - mse_loss_weights = base_weight - - # For zero-terminal SNR, we have to handle the case where a sigma of Zero results in a Inf value. - # When we run this, the MSE loss weights for this timestep is set unconditionally to 1. - # If we do not run this, the loss value will go to NaN almost immediately, usually within one step. - mse_loss_weights[snr == 0] = 1.0 - - # We first calculate the original loss. Then we mean over the non-batch dimensions and - # rebalance the sample-wise losses with their respective loss weights. - # Finally, we take the mean of the rebalanced loss. + loss = F.mse_loss(model_pred.float(), target.float(), reduction="none") loss = loss.mean(dim=list(range(1, len(loss.shape)))) * mse_loss_weights loss = loss.mean() diff --git a/examples/text_to_image/train_text_to_image.py b/examples/text_to_image/train_text_to_image.py index 03efc3fa13c5..e216529b2f54 100644 --- a/examples/text_to_image/train_text_to_image.py +++ b/examples/text_to_image/train_text_to_image.py @@ -929,25 +929,13 @@ def collate_fn(examples): # Since we predict the noise instead of x_0, the original formulation is slightly changed. # This is discussed in Section 4.2 of the same paper. snr = compute_snr(noise_scheduler, timesteps) - base_weight = ( + if noise_scheduler.config.prediction_type == "v_prediction": + # Velocity objective requires that we add one to SNR values before we divide by them. + snr = snr + 1 + mse_loss_weights = ( torch.stack([snr, args.snr_gamma * torch.ones_like(timesteps)], dim=1).min(dim=1)[0] / snr ) - if noise_scheduler.config.prediction_type == "v_prediction": - # Velocity objective needs to be floored to an SNR weight of one. - mse_loss_weights = base_weight + 1 - else: - # Epsilon and sample both use the same loss weights. - mse_loss_weights = base_weight - - # For zero-terminal SNR, we have to handle the case where a sigma of Zero results in a Inf value. - # When we run this, the MSE loss weights for this timestep is set unconditionally to 1. - # If we do not run this, the loss value will go to NaN almost immediately, usually within one step. - mse_loss_weights[snr == 0] = 1.0 - - # We first calculate the original loss. Then we mean over the non-batch dimensions and - # rebalance the sample-wise losses with their respective loss weights. - # Finally, we take the mean of the rebalanced loss. loss = F.mse_loss(model_pred.float(), target.float(), reduction="none") loss = loss.mean(dim=list(range(1, len(loss.shape)))) * mse_loss_weights loss = loss.mean() diff --git a/examples/text_to_image/train_text_to_image_lora.py b/examples/text_to_image/train_text_to_image_lora.py index 6ae157bce9b3..eac0f18f49f4 100644 --- a/examples/text_to_image/train_text_to_image_lora.py +++ b/examples/text_to_image/train_text_to_image_lora.py @@ -759,25 +759,13 @@ def collate_fn(examples): # Since we predict the noise instead of x_0, the original formulation is slightly changed. # This is discussed in Section 4.2 of the same paper. snr = compute_snr(noise_scheduler, timesteps) - base_weight = ( + if noise_scheduler.config.prediction_type == "v_prediction": + # Velocity objective requires that we add one to SNR values before we divide by them. + snr = snr + 1 + mse_loss_weights = ( torch.stack([snr, args.snr_gamma * torch.ones_like(timesteps)], dim=1).min(dim=1)[0] / snr ) - if noise_scheduler.config.prediction_type == "v_prediction": - # Velocity objective needs to be floored to an SNR weight of one. - mse_loss_weights = base_weight + 1 - else: - # Epsilon and sample both use the same loss weights. - mse_loss_weights = base_weight - - # For zero-terminal SNR, we have to handle the case where a sigma of Zero results in a Inf value. - # When we run this, the MSE loss weights for this timestep is set unconditionally to 1. - # If we do not run this, the loss value will go to NaN almost immediately, usually within one step. - mse_loss_weights[snr == 0] = 1.0 - - # We first calculate the original loss. Then we mean over the non-batch dimensions and - # rebalance the sample-wise losses with their respective loss weights. - # Finally, we take the mean of the rebalanced loss. loss = F.mse_loss(model_pred.float(), target.float(), reduction="none") loss = loss.mean(dim=list(range(1, len(loss.shape)))) * mse_loss_weights loss = loss.mean() diff --git a/examples/text_to_image/train_text_to_image_lora_sdxl.py b/examples/text_to_image/train_text_to_image_lora_sdxl.py index d685d468db4d..ed7a15cd95fe 100644 --- a/examples/text_to_image/train_text_to_image_lora_sdxl.py +++ b/examples/text_to_image/train_text_to_image_lora_sdxl.py @@ -1050,25 +1050,13 @@ def compute_time_ids(original_size, crops_coords_top_left): # Since we predict the noise instead of x_0, the original formulation is slightly changed. # This is discussed in Section 4.2 of the same paper. snr = compute_snr(noise_scheduler, timesteps) - base_weight = ( + if noise_scheduler.config.prediction_type == "v_prediction": + # Velocity objective requires that we add one to SNR values before we divide by them. + snr = snr + 1 + mse_loss_weights = ( torch.stack([snr, args.snr_gamma * torch.ones_like(timesteps)], dim=1).min(dim=1)[0] / snr ) - if noise_scheduler.config.prediction_type == "v_prediction": - # Velocity objective needs to be floored to an SNR weight of one. - mse_loss_weights = base_weight + 1 - else: - # Epsilon and sample both use the same loss weights. - mse_loss_weights = base_weight - - # For zero-terminal SNR, we have to handle the case where a sigma of Zero results in a Inf value. - # When we run this, the MSE loss weights for this timestep is set unconditionally to 1. - # If we do not run this, the loss value will go to NaN almost immediately, usually within one step. - mse_loss_weights[snr == 0] = 1.0 - - # We first calculate the original loss. Then we mean over the non-batch dimensions and - # rebalance the sample-wise losses with their respective loss weights. - # Finally, we take the mean of the rebalanced loss. loss = F.mse_loss(model_pred.float(), target.float(), reduction="none") loss = loss.mean(dim=list(range(1, len(loss.shape)))) * mse_loss_weights loss = loss.mean() diff --git a/examples/text_to_image/train_text_to_image_sdxl.py b/examples/text_to_image/train_text_to_image_sdxl.py index 4cf966af77cd..c681943f2e94 100644 --- a/examples/text_to_image/train_text_to_image_sdxl.py +++ b/examples/text_to_image/train_text_to_image_sdxl.py @@ -1067,25 +1067,13 @@ def compute_time_ids(original_size, crops_coords_top_left): # Since we predict the noise instead of x_0, the original formulation is slightly changed. # This is discussed in Section 4.2 of the same paper. snr = compute_snr(noise_scheduler, timesteps) - base_weight = ( + if noise_scheduler.config.prediction_type == "v_prediction": + # Velocity objective requires that we add one to SNR values before we divide by them. + snr = snr + 1 + mse_loss_weights = ( torch.stack([snr, args.snr_gamma * torch.ones_like(timesteps)], dim=1).min(dim=1)[0] / snr ) - if noise_scheduler.config.prediction_type == "v_prediction": - # Velocity objective needs to be floored to an SNR weight of one. - mse_loss_weights = base_weight + 1 - else: - # Epsilon and sample both use the same loss weights. - mse_loss_weights = base_weight - - # For zero-terminal SNR, we have to handle the case where a sigma of Zero results in a Inf value. - # When we run this, the MSE loss weights for this timestep is set unconditionally to 1. - # If we do not run this, the loss value will go to NaN almost immediately, usually within one step. - mse_loss_weights[snr == 0] = 1.0 - - # We first calculate the original loss. Then we mean over the non-batch dimensions and - # rebalance the sample-wise losses with their respective loss weights. - # Finally, we take the mean of the rebalanced loss. loss = F.mse_loss(model_pred.float(), target.float(), reduction="none") loss = loss.mean(dim=list(range(1, len(loss.shape)))) * mse_loss_weights loss = loss.mean() From 0922210c5cf2ca0f93fa5f924c6ed195ec7da7e2 Mon Sep 17 00:00:00 2001 From: Patrick von Platen Date: Fri, 6 Oct 2023 09:42:20 +0200 Subject: [PATCH 03/25] Update bug-report.yml --- .github/ISSUE_TEMPLATE/bug-report.yml | 45 +++++++++++++++++++++------ 1 file changed, 35 insertions(+), 10 deletions(-) diff --git a/.github/ISSUE_TEMPLATE/bug-report.yml b/.github/ISSUE_TEMPLATE/bug-report.yml index 483bbd1650da..e9e672e72dd5 100644 --- a/.github/ISSUE_TEMPLATE/bug-report.yml +++ b/.github/ISSUE_TEMPLATE/bug-report.yml @@ -13,7 +13,7 @@ body: *Give your issue a fitting title. Assume that someone which very limited knowledge of diffusers can understand your issue. Add links to the source code, documentation other issues, pull requests etc...* - 2. If your issue is about something not working, **always** provide a reproducible code snippet. The reader should be able to reproduce your issue by **only copy-pasting your code snippet into a Python shell**. *The community cannot solve your issue if it cannot reproduce it. If your bug is related to training, add your training script and make everything needed to train public. Otherwise, just add a simple Python code snippet.* - - 3. Add the **minimum amount of code / context that is needed to understand, reproduce your issue**. + - 3. Add the **minimum** amount of code / context that is needed to understand, reproduce your issue. *Make the life of maintainers easy. `diffusers` is getting many issues every day. Make sure your issue is about one bug and one bug only. Make sure you add only the context, code needed to understand your issues - nothing more. Generally, every issue is a way of documenting this library, try to make it a good documentation entry.* - 4. For issues related to community pipelines (i.e., the pipelines located in the `examples/community` folder), please tag the author of the pipeline in your issue thread as those pipelines are not maintained. - type: markdown @@ -61,21 +61,46 @@ body: All issues are read by one of the core maintainers, so if you don't know who to tag, just leave this blank and a core maintainer will ping the right person. - Please tag fewer than 3 people. - - General library related questions: @patrickvonplaten and @sayakpaul + Please tag a maximum of 2 people. + + Questions on DiffusionPipeline (Saving, Loading, From pretrained, ...): + + Questions on pipelines: + - Stable Diffusion @yiyixuxu @DN6 @patrickvonplaten @sayakpaul @patrickvonplaten + - Stable Diffusion XL @yiyixuxu @sayakpaul @DN6 @patrickvonplaten + - Kandinsky @yiyixuxu @patrickvonplaten + - ControlNet @sayakpaul @yiyixuxu @DN6 @patrickvonplaten + - T2I Adapter @sayakpaul @yiyixuxu @DN6 @patrickvonplaten + - IF @DN6 @patrickvonplaten + - Text-to-Video / Video-to-Video @DN6 @sayakpaul @patrickvonplaten + - Wuerstchen @DN6 @patrickvonplaten + - Other: @yiyixuxu @DN6 + + Questions on models: + - UNet @DN6 @yiyixuxu @sayakpaul @patrickvonplaten + - VAE @sayakpaul @DN6 @yiyixuxu @patrickvonplaten + - Transformers/Attention @DN6 @yiyixuxu @sayakpaul @DN6 @patrickvonplaten - Questions on the training examples: @williamberman, @sayakpaul, @yiyixuxu + Questions on Schedulers: @yiyixuxu @patrickvonplaten - Questions on memory optimizations, LoRA, float16, etc.: @williamberman, @patrickvonplaten, and @sayakpaul + Questions on LoRA: @sayakpaul @patrickvonplaten - Questions on schedulers: @patrickvonplaten and @williamberman + Questions on Textual Inversion: @sayakpaul @patrickvonplaten - Questions on models and pipelines: @patrickvonplaten, @sayakpaul, and @williamberman (for community pipelines, please tag the original author of the pipeline) + Questions on Training: + - DreamBooth @sayakpaul @patrickvonplaten + - Text-to-Image Fine-tuning @sayakpaul @patrickvonplaten + - Textual Inversion @sayakpaul @patrickvonplaten + - ControlNet @sayakpaul @patrickvonplaten + + Questions on Tests: @DN6 @sayakpaul @yiyixuxu + + Questions on Documentation: @stevhliu Questions on JAX- and MPS-related things: @pcuenca - Questions on audio pipelines: @patrickvonplaten, @kashif, and @sanchit-gandhi + Questions on audio pipelines: @DN6 @patrickvonplaten + + - Documentation: @stevhliu and @yiyixuxu placeholder: "@Username ..." From 6ce01bd647c90edb520aaf2c2aeaf9be144bee75 Mon Sep 17 00:00:00 2001 From: Dhruv Nair Date: Fri, 6 Oct 2023 10:25:18 +0200 Subject: [PATCH 04/25] Bump tolerance on shape test (#5289) bump tolerance on shape test --- tests/pipelines/shap_e/test_shap_e_img2img.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/pipelines/shap_e/test_shap_e_img2img.py b/tests/pipelines/shap_e/test_shap_e_img2img.py index 55c0ae6bd02e..055dbe7a97d4 100644 --- a/tests/pipelines/shap_e/test_shap_e_img2img.py +++ b/tests/pipelines/shap_e/test_shap_e_img2img.py @@ -250,7 +250,7 @@ def test_float16_inference(self): super().test_float16_inference(expected_max_diff=1e-1) def test_save_load_local(self): - super().test_save_load_local(expected_max_difference=1e-3) + super().test_save_load_local(expected_max_difference=5e-3) @unittest.skip("Key error is raised with accelerate") def test_sequential_cpu_offload_forward_pass(self): From 872ae1dd12b0f7683834fdf9037810f70a661901 Mon Sep 17 00:00:00 2001 From: Dhruv Nair Date: Fri, 6 Oct 2023 13:18:13 +0200 Subject: [PATCH 05/25] Add from single file to StableDiffusionUpscalePipeline and StableDiffusionLatentUpscalePipeline (#5194) * add from single file * clean up * make style * add single file loading for upscaling --- .../stable_diffusion/convert_from_ckpt.py | 53 +++++++++++++++++-- ...ipeline_stable_diffusion_latent_upscale.py | 3 +- .../pipeline_stable_diffusion_upscale.py | 6 ++- .../test_stable_diffusion_upscale.py | 34 ++++++++++++ 4 files changed, 89 insertions(+), 7 deletions(-) diff --git a/src/diffusers/pipelines/stable_diffusion/convert_from_ckpt.py b/src/diffusers/pipelines/stable_diffusion/convert_from_ckpt.py index 618ee1942224..e97f66bbcb24 100644 --- a/src/diffusers/pipelines/stable_diffusion/convert_from_ckpt.py +++ b/src/diffusers/pipelines/stable_diffusion/convert_from_ckpt.py @@ -304,8 +304,6 @@ def create_unet_diffusers_config(original_config, image_size: int, controlnet=Fa class_embed_type = "projection" assert "adm_in_channels" in unet_params projection_class_embeddings_input_dim = unet_params.adm_in_channels - else: - raise NotImplementedError(f"Unknown conditional unet num_classes config: {unet_params.num_classes}") config = { "sample_size": image_size // vae_scale_factor, @@ -323,6 +321,12 @@ def create_unet_diffusers_config(original_config, image_size: int, controlnet=Fa "transformer_layers_per_block": transformer_layers_per_block, } + if "disable_self_attentions" in unet_params: + config["only_cross_attention"] = unet_params.disable_self_attentions + + if "num_classes" in unet_params and type(unet_params.num_classes) == int: + config["num_class_embeds"] = unet_params.num_classes + if controlnet: config["conditioning_channels"] = unet_params.hint_channels else: @@ -441,6 +445,10 @@ def convert_ldm_unet_checkpoint( new_checkpoint["add_embedding.linear_2.weight"] = unet_state_dict["label_emb.0.2.weight"] new_checkpoint["add_embedding.linear_2.bias"] = unet_state_dict["label_emb.0.2.bias"] + # Relevant to StableDiffusionUpscalePipeline + if "num_class_embeds" in config: + new_checkpoint["class_embedding.weight"] = unet_state_dict["label_emb.weight"] + new_checkpoint["conv_in.weight"] = unet_state_dict["input_blocks.0.0.weight"] new_checkpoint["conv_in.bias"] = unet_state_dict["input_blocks.0.0.bias"] @@ -496,6 +504,7 @@ def convert_ldm_unet_checkpoint( if len(attentions): paths = renew_attention_paths(attentions) + meta_path = {"old": f"input_blocks.{i}.1", "new": f"down_blocks.{block_id}.attentions.{layer_in_block_id}"} assign_to_checkpoint( paths, new_checkpoint, unet_state_dict, additional_replacements=[meta_path], config=config @@ -1210,6 +1219,7 @@ def download_from_original_stable_diffusion_ckpt( StableDiffusionControlNetPipeline, StableDiffusionInpaintPipeline, StableDiffusionPipeline, + StableDiffusionUpscalePipeline, StableDiffusionXLImg2ImgPipeline, StableUnCLIPImg2ImgPipeline, StableUnCLIPPipeline, @@ -1256,6 +1266,8 @@ def download_from_original_stable_diffusion_ckpt( key_name_v2_1 = "model.diffusion_model.input_blocks.2.1.transformer_blocks.0.attn2.to_k.weight" key_name_sd_xl_base = "conditioner.embedders.1.model.transformer.resblocks.9.mlp.c_proj.bias" key_name_sd_xl_refiner = "conditioner.embedders.0.model.transformer.resblocks.9.mlp.c_proj.bias" + is_upscale = pipeline_class == StableDiffusionUpscalePipeline + config_url = None # model_type = "v1" @@ -1285,6 +1297,10 @@ def download_from_original_stable_diffusion_ckpt( original_config_file = config_files["xl_refiner"] else: config_url = "https://raw.githubusercontent.com/Stability-AI/generative-models/main/configs/inference/sd_xl_refiner.yaml" + + if is_upscale: + config_url = "https://raw.githubusercontent.com/Stability-AI/stablediffusion/main/configs/stable-diffusion/x4-upscaling.yaml" + if config_url is not None: original_config_file = BytesIO(requests.get(config_url).content) @@ -1308,6 +1324,8 @@ def download_from_original_stable_diffusion_ckpt( if num_in_channels is None and pipeline_class == StableDiffusionInpaintPipeline: num_in_channels = 9 + if num_in_channels is None and pipeline_class == StableDiffusionUpscalePipeline: + num_in_channels = 7 elif num_in_channels is None: num_in_channels = 4 @@ -1391,9 +1409,13 @@ def download_from_original_stable_diffusion_ckpt( else: raise ValueError(f"Scheduler of type {scheduler_type} doesn't exist!") + if pipeline_class == StableDiffusionUpscalePipeline: + image_size = original_config.model.params.unet_config.params.image_size + # Convert the UNet2DConditionModel model. unet_config = create_unet_diffusers_config(original_config, image_size=image_size) unet_config["upcast_attention"] = upcast_attention + path = checkpoint_path_or_dict if isinstance(checkpoint_path_or_dict, str) else "" converted_unet_checkpoint = convert_ldm_unet_checkpoint( checkpoint, unet_config, path=path, extract_ema=extract_ema @@ -1458,8 +1480,29 @@ def download_from_original_stable_diffusion_ckpt( controlnet=controlnet, safety_checker=None, feature_extractor=None, - requires_safety_checker=False, ) + if hasattr(pipe, "requires_safety_checker"): + pipe.requires_safety_checker = False + + elif pipeline_class == StableDiffusionUpscalePipeline: + scheduler = DDIMScheduler.from_pretrained( + "stabilityai/stable-diffusion-x4-upscaler", subfolder="scheduler" + ) + low_res_scheduler = DDPMScheduler.from_pretrained( + "stabilityai/stable-diffusion-x4-upscaler", subfolder="low_res_scheduler" + ) + + pipe = pipeline_class( + vae=vae, + text_encoder=text_model, + tokenizer=tokenizer, + unet=unet, + scheduler=scheduler, + low_res_scheduler=low_res_scheduler, + safety_checker=None, + feature_extractor=None, + ) + else: pipe = pipeline_class( vae=vae, @@ -1469,8 +1512,10 @@ def download_from_original_stable_diffusion_ckpt( scheduler=scheduler, safety_checker=None, feature_extractor=None, - requires_safety_checker=False, ) + if hasattr(pipe, "requires_safety_checker"): + pipe.requires_safety_checker = False + else: image_normalizer, image_noising_scheduler = stable_unclip_image_noising_components( original_config, clip_stats_path=clip_stats_path, device=device diff --git a/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_latent_upscale.py b/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_latent_upscale.py index 57d23e89c639..1e38142b9c66 100644 --- a/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_latent_upscale.py +++ b/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_latent_upscale.py @@ -22,6 +22,7 @@ from transformers import CLIPTextModel, CLIPTokenizer from ...image_processor import PipelineImageInput, VaeImageProcessor +from ...loaders import FromSingleFileMixin from ...models import AutoencoderKL, UNet2DConditionModel from ...schedulers import EulerDiscreteScheduler from ...utils import deprecate, logging @@ -59,7 +60,7 @@ def preprocess(image): return image -class StableDiffusionLatentUpscalePipeline(DiffusionPipeline): +class StableDiffusionLatentUpscalePipeline(DiffusionPipeline, FromSingleFileMixin): r""" Pipeline for upscaling Stable Diffusion output image resolution by a factor of 2. diff --git a/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_upscale.py b/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_upscale.py index adf950a1df29..8d01e0a0d086 100644 --- a/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_upscale.py +++ b/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_upscale.py @@ -22,7 +22,7 @@ from transformers import CLIPImageProcessor, CLIPTextModel, CLIPTokenizer from ...image_processor import PipelineImageInput, VaeImageProcessor -from ...loaders import LoraLoaderMixin, TextualInversionLoaderMixin +from ...loaders import FromSingleFileMixin, LoraLoaderMixin, TextualInversionLoaderMixin from ...models import AutoencoderKL, UNet2DConditionModel from ...models.attention_processor import ( AttnProcessor2_0, @@ -67,7 +67,9 @@ def preprocess(image): return image -class StableDiffusionUpscalePipeline(DiffusionPipeline, TextualInversionLoaderMixin, LoraLoaderMixin): +class StableDiffusionUpscalePipeline( + DiffusionPipeline, TextualInversionLoaderMixin, LoraLoaderMixin, FromSingleFileMixin +): r""" Pipeline for text-guided image super-resolution using Stable Diffusion 2. diff --git a/tests/pipelines/stable_diffusion_2/test_stable_diffusion_upscale.py b/tests/pipelines/stable_diffusion_2/test_stable_diffusion_upscale.py index 2c0f37519ad8..aa5b3e38b0c1 100644 --- a/tests/pipelines/stable_diffusion_2/test_stable_diffusion_upscale.py +++ b/tests/pipelines/stable_diffusion_2/test_stable_diffusion_upscale.py @@ -29,6 +29,7 @@ floats_tensor, load_image, load_numpy, + numpy_cosine_similarity_distance, require_torch_gpu, slow, torch_device, @@ -479,3 +480,36 @@ def test_stable_diffusion_pipeline_with_sequential_cpu_offloading(self): mem_bytes = torch.cuda.max_memory_allocated() # make sure that less than 2.9 GB is allocated assert mem_bytes < 2.9 * 10**9 + + def test_download_ckpt_diff_format_is_same(self): + image = load_image( + "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main" + "/sd2-upscale/low_res_cat.png" + ) + + prompt = "a cat sitting on a park bench" + model_id = "stabilityai/stable-diffusion-x4-upscaler" + pipe = StableDiffusionUpscalePipeline.from_pretrained(model_id) + pipe.enable_model_cpu_offload() + + generator = torch.Generator("cpu").manual_seed(0) + output = pipe(prompt=prompt, image=image, generator=generator, output_type="np", num_inference_steps=3) + image_from_pretrained = output.images[0] + + single_file_path = ( + "https://huggingface.co/stabilityai/stable-diffusion-x4-upscaler/blob/main/x4-upscaler-ema.safetensors" + ) + pipe_from_single_file = StableDiffusionUpscalePipeline.from_single_file(single_file_path) + pipe_from_single_file.enable_model_cpu_offload() + + generator = torch.Generator("cpu").manual_seed(0) + output_from_single_file = pipe_from_single_file( + prompt=prompt, image=image, generator=generator, output_type="np", num_inference_steps=3 + ) + image_from_single_file = output_from_single_file.images[0] + + assert image_from_pretrained.shape == (512, 512, 3) + assert image_from_single_file.shape == (512, 512, 3) + assert ( + numpy_cosine_similarity_distance(image_from_pretrained.flatten(), image_from_single_file.flatten()) < 1e-3 + ) From 7eaae83f168e0d7ebca1ef238cbc58c5c46d39d9 Mon Sep 17 00:00:00 2001 From: Sayak Paul Date: Fri, 6 Oct 2023 17:14:47 +0200 Subject: [PATCH 06/25] [LoRA] fix: torch.compile() for lora conv (#5298) fix: torch.compile() for lora conv --- src/diffusers/models/lora.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/src/diffusers/models/lora.py b/src/diffusers/models/lora.py index cc8e3e231e2b..a777bb93e1c8 100644 --- a/src/diffusers/models/lora.py +++ b/src/diffusers/models/lora.py @@ -164,7 +164,10 @@ def forward(self, hidden_states, scale: float = 1.0): hidden_states, self.weight, self.bias, self.stride, self.padding, self.dilation, self.groups ) else: - return super().forward(hidden_states) + (scale * self.lora_layer(hidden_states)) + original_outputs = F.conv2d( + hidden_states, self.weight, self.bias, self.stride, self.padding, self.dilation, self.groups + ) + return original_outputs + (scale * self.lora_layer(hidden_states)) class LoRACompatibleLinear(nn.Linear): From f0a2c637538458bea4b751903522766ff7d7f9c8 Mon Sep 17 00:00:00 2001 From: Steven Liu <59462357+stevhliu@users.noreply.github.com> Date: Fri, 6 Oct 2023 09:44:24 -0700 Subject: [PATCH 07/25] [docs] Improved inpaint docs (#5210) * start * finish draft * add section * edits * feedback * make fix-copies * rebase --- docs/source/en/using-diffusers/img2img.md | 8 +- docs/source/en/using-diffusers/inpaint.md | 549 ++++++++++++++++++++-- 2 files changed, 506 insertions(+), 51 deletions(-) diff --git a/docs/source/en/using-diffusers/img2img.md b/docs/source/en/using-diffusers/img2img.md index 82aa328d2b9c..c0bf4dc52153 100644 --- a/docs/source/en/using-diffusers/img2img.md +++ b/docs/source/en/using-diffusers/img2img.md @@ -33,7 +33,7 @@ pipeline.enable_xformers_memory_efficient_attention() -You'll notice throughout the guide, we use [`~DiffusionPipeline.enable_model_cpu_offload`] and [`~DiffusionPipeline.enable_xformers_memory_efficient_attention`], to save memory and increase inference speed. If you're using PyTorch 2.0, then you don't need to call [`~DiffusionPipeline.enable_xformers_memory_efficient_attention`] on your pipeline because it'll already be using PyTorch 2.0's native [scaled-dot product attention](/optimization/torch2.0#scaled-dot-product-attention). +You'll notice throughout the guide, we use [`~DiffusionPipeline.enable_model_cpu_offload`] and [`~DiffusionPipeline.enable_xformers_memory_efficient_attention`], to save memory and increase inference speed. If you're using PyTorch 2.0, then you don't need to call [`~DiffusionPipeline.enable_xformers_memory_efficient_attention`] on your pipeline because it'll already be using PyTorch 2.0's native [scaled-dot product attention](../optimization/torch2.0#scaled-dot-product-attention). @@ -590,17 +590,17 @@ image ## Optimize -Running diffusion models is computationally expensive and intensive, but with a few optimization tricks, it is entirely possible to run them on consumer and free-tier GPUs. For example, you can use a more memory-efficient form of attention such as PyTorch 2.0's [scaled-dot product attention](optimization/torch2.0#scaled-dot-product-attention) or [xFormers](optimization/xformers) (you can use one or the other, but there's no need to use both). You can also offload the model to the GPU while the other pipeline components wait on the CPU. +Running diffusion models is computationally expensive and intensive, but with a few optimization tricks, it is entirely possible to run them on consumer and free-tier GPUs. For example, you can use a more memory-efficient form of attention such as PyTorch 2.0's [scaled-dot product attention](../optimization/torch2.0#scaled-dot-product-attention) or [xFormers](../optimization/xformers) (you can use one or the other, but there's no need to use both). You can also offload the model to the GPU while the other pipeline components wait on the CPU. ```diff + pipeline.enable_model_cpu_offload() + pipeline.enable_xformers_memory_efficient_attention() ``` -With [`torch.compile`](optimization/torch2.0#torch.compile), you can boost your inference speed even more by wrapping your UNet with it: +With [`torch.compile`](../optimization/torch2.0#torch.compile), you can boost your inference speed even more by wrapping your UNet with it: ```py pipe.unet = torch.compile(pipe.unet, mode="reduce-overhead", fullgraph=True) ``` -To learn more, take a look at the [Reduce memory usage](optimization/memory) and [Torch 2.0](optimization/torch2.0) guides. +To learn more, take a look at the [Reduce memory usage](../optimization/memory) and [Torch 2.0](../optimization/torch2.0) guides. diff --git a/docs/source/en/using-diffusers/inpaint.md b/docs/source/en/using-diffusers/inpaint.md index 7f10e43243a3..4d99fca26eb6 100644 --- a/docs/source/en/using-diffusers/inpaint.md +++ b/docs/source/en/using-diffusers/inpaint.md @@ -10,87 +10,289 @@ an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express o specific language governing permissions and limitations under the License. --> -# Text-guided image-inpainting +# Inpainting [[open-in-colab]] -The [`StableDiffusionInpaintPipeline`] allows you to edit specific parts of an image by providing a mask and a text prompt. It uses a version of Stable Diffusion, like [`runwayml/stable-diffusion-inpainting`](https://huggingface.co/runwayml/stable-diffusion-inpainting) specifically trained for inpainting tasks. +Inpainting replaces or edits specific areas of an image. This makes it a useful tool for image restoration like removing defects and artifacts, or even replacing an image area with something entirely new. Inpainting relies on a mask to determine which regions of an image to fill in; the area to inpaint is represented by white pixels and the area to keep is represented by black pixels. The white pixels are filled in by the prompt. -Get started by loading an instance of the [`StableDiffusionInpaintPipeline`]: +With 🤗 Diffusers, here is how you can do inpainting: -```python -import PIL -import requests +1. Load an inpainting checkpoint with the [`AutoPipelineForInpainting`] class. This'll automatically detect the appropriate pipeline class to load based on the checkpoint: + +```py import torch -from io import BytesIO +from diffusers import AutoPipelineForInpainting +from diffusers.utils import load_image -from diffusers import StableDiffusionInpaintPipeline +pipeline = AutoPipelineForInpainting.from_pretrained( + "kandinsky-community/kandinsky-2-2-decoder-inpaint", torch_dtype=torch.float16 +).to("cuda") +pipeline.enable_model_cpu_offload() +pipeline.enable_xformers_memory_efficient_attention() +``` -pipeline = StableDiffusionInpaintPipeline.from_pretrained( - "runwayml/stable-diffusion-inpainting", - torch_dtype=torch.float16, - use_safetensors=True, - variant="fp16", -) -pipeline = pipeline.to("cuda") + + +You'll notice throughout the guide, we use [`~DiffusionPipeline.enable_model_cpu_offload`] and [`~DiffusionPipeline.enable_xformers_memory_efficient_attention`], to save memory and increase inference speed. If you're using PyTorch 2.0, it's not necessary to call [`~DiffusionPipeline.enable_xformers_memory_efficient_attention`] on your pipeline because it'll already be using PyTorch 2.0's native [scaled-dot product attention](../optimization/torch2.0#scaled-dot-product-attention). + + + +2. Load the base and mask images: + +```py +init_image = load_image("https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/inpaint.png").convert("RGB") +mask_image = load_image("https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/inpaint_mask.png").convert("RGB") ``` -Download an image and a mask of a dog which you'll eventually replace: +3. Create a prompt to inpaint the image with and pass it to the pipeline with the base and mask images: -```python -def download_image(url): - response = requests.get(url) - return PIL.Image.open(BytesIO(response.content)).convert("RGB") +```py +prompt = "a black cat with glowing eyes, cute, adorable, disney, pixar, highly detailed, 8k" +negative_prompt = "bad anatomy, deformed, ugly, disfigured" +image = pipeline(prompt=prompt, negative_prompt=negative_prompt, image=init_image, mask_image=mask_image).images[0] +``` +
+
+ +
base image
+
+
+ +
generated image
+
+
-img_url = "https://raw.githubusercontent.com/CompVis/latent-diffusion/main/data/inpainting_examples/overture-creations-5sI6fQgYIuo.png" -mask_url = "https://raw.githubusercontent.com/CompVis/latent-diffusion/main/data/inpainting_examples/overture-creations-5sI6fQgYIuo_mask.png" +## Popular models + +[Stable Diffusion Inpainting](https://huggingface.co/runwayml/stable-diffusion-inpainting), [Stable Diffusion XL (SDXL) Inpainting](https://huggingface.co/diffusers/stable-diffusion-xl-1.0-inpainting-0.1), and [Kandinsky 2.2](https://huggingface.co/kandinsky-community/kandinsky-2-2-decoder-inpaint) are among the most popular models for inpainting. SDXL typically produces higher resolution images than Stable Diffusion v1.5, and Kandinsky 2.2 is also capable of generating high-quality images. + +### Stable Diffusion Inpainting + +Stable Diffusion Inpainting is a latent diffusion model finetuned on 512x512 images on inpainting. It is a good starting point because it is relatively fast and generates good quality images. To use this model for inpainting, you'll need to pass a prompt, base and mask image to the pipeline: -init_image = download_image(img_url).resize((512, 512)) -mask_image = download_image(mask_url).resize((512, 512)) +```py +import torch +from diffusers import AutoPipelineForInpainting +from diffusers.utils import load_image + +pipeline = AutoPipelineForInpainting.from_pretrained( + "runwayml/stable-diffusion-inpainting", torch_dtype=torch.float16, variant="fp16" +).to("cuda") +pipeline.enable_model_cpu_offload() +pipeline.enable_xformers_memory_efficient_attention() + +# load base and mask image +init_image = load_image("https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/inpaint.png").convert("RGB") +mask_image = load_image("https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/inpaint_mask.png").convert("RGB") + +generator = torch.Generator("cuda").manual_seed(92) +prompt = "concept art digital painting of an elven castle, inspired by lord of the rings, highly detailed, 8k" +image = pipeline(prompt=prompt, image=init_image, mask_image=mask_image, generator=generator).images[0] ``` -Now you can create a prompt to replace the mask with something else: +### Stable Diffusion XL (SDXL) Inpainting -```python -prompt = "Face of a yellow cat, high resolution, sitting on a park bench" -image = pipeline(prompt=prompt, image=init_image, mask_image=mask_image).images[0] +SDXL is a larger and more powerful version of Stable Diffusion v1.5. This model can follow a two-stage model process (though each model can also be used alone); the base model generates an image, and a refiner model takes that image and further enhances its details and quality. Take a look at the [SDXL](sdxl) guide for a more comprehensive guide on how to use SDXL and configure it's parameters. + +```py +import torch +from diffusers import AutoPipelineForInpainting +from diffusers.utils import load_image + +pipeline = AutoPipelineForInpainting.from_pretrained( + "diffusers/stable-diffusion-xl-1.0-inpainting-0.1", torch_dtype=torch.float16, variant="fp16" +).to("cuda") +pipeline.enable_model_cpu_offload() +pipeline.enable_xformers_memory_efficient_attention() + +# load base and mask image +init_image = load_image("https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/inpaint.png").convert("RGB") +mask_image = load_image("https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/inpaint_mask.png").convert("RGB") + +generator = torch.Generator("cuda").manual_seed(92) +prompt = "concept art digital painting of an elven castle, inspired by lord of the rings, highly detailed, 8k" +image = pipeline(prompt=prompt, image=init_image, mask_image=mask_image, generator=generator).images[0] ``` -`image` | `mask_image` | `prompt` | output | -:-------------------------:|:-------------------------:|:-------------------------:|-------------------------:| -drawing | drawing | ***Face of a yellow cat, high resolution, sitting on a park bench*** | drawing | +### Kandinsky 2.2 Inpainting +The Kandinsky model family is similar to SDXL because it uses two models as well; the image prior model creates image embeddings, and the diffusion model generates images from them. You can load the image prior and diffusion model separately, but the easiest way to use Kandinsky 2.2 is to load it into the [`AutoPipelineForInpainting`] class which uses the [`KandinskyV22InpaintCombinedPipeline`] under the hood. - +```py +import torch +from diffusers import AutoPipelineForInpainting +from diffusers.utils import load_image -A previous experimental implementation of inpainting used a different, lower-quality process. To ensure backwards compatibility, loading a pretrained pipeline that doesn't contain the new model will still apply the old inpainting method. +pipeline = AutoPipelineForInpainting.from_pretrained( + "kandinsky-community/kandinsky-2-2-decoder-inpaint", torch_dtype=torch.float16 +).to("cuda") +pipeline.enable_model_cpu_offload() +pipeline.enable_xformers_memory_efficient_attention() - +# load base and mask image +init_image = load_image("https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/inpaint.png").convert("RGB") +mask_image = load_image("https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/inpaint_mask.png").convert("RGB") + +generator = torch.Generator("cuda").manual_seed(92) +prompt = "concept art digital painting of an elven castle, inspired by lord of the rings, highly detailed, 8k" +image = pipeline(prompt=prompt, image=init_image, mask_image=mask_image, generator=generator).images[0] +``` -Check out the Spaces below to try out image inpainting yourself! +
+
+ +
base image
+
+
+ +
Stable Diffusion Inpainting
+
+
+ +
Stable Diffusion XL Inpainting
+
+
+ +
Kandinsky 2.2 Inpainting
+
+
+ +## Configure pipeline parameters + +Image features - like quality and "creativity" - are dependent on pipeline parameters. Knowing what these parameters do is important for getting the results you want. Let's take a look at the most important parameters and see how changing them affects the output. + +## Strength + +`strength` is a measure of how much noise is added to the base image, which influences how similar the output is to the base image. + +* 📈 a high `strength` value means more noise is added to an image and the denoising process takes longer, but you'll get higher quality images that are more different from the base image +* 📉 a low `strength` value means less noise is added to an image and the denoising process is faster, but the image quality may not be as great and the generated image resembles the base image more + +```py +import torch +from diffusers import AutoPipelineForInpainting +from diffusers.utils import load_image + +pipeline = AutoPipelineForInpainting.from_pretrained( + "runwayml/stable-diffusion-inpainting", torch_dtype=torch.float16, variant="fp16" +).to("cuda") +pipeline.enable_model_cpu_offload() +pipeline.enable_xformers_memory_efficient_attention() + +# load base and mask image +init_image = load_image("https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/inpaint.png").convert("RGB") +mask_image = load_image("https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/inpaint_mask.png").convert("RGB") + +prompt = "concept art digital painting of an elven castle, inspired by lord of the rings, highly detailed, 8k" +image = pipeline(prompt=prompt, image=init_image, mask_image=mask_image, strength=0.6).images[0] +``` + +
+
+ +
strength = 0.6
+
+
+ +
strength = 0.8
+
+
+ +
strength = 1.0
+
+
+ +## Guidance scale + +`guidance_scale` affects how aligned the text prompt and generated image are. + +* 📈 a high `guidance_scale` value means the prompt and generated image are closely aligned, so the output is a stricter interpretation of the prompt +* 📉 a low `guidance_scale` value means the prompt and generated image are more loosely aligned, so the output may be more varied from the prompt + +You can use `strength` and `guidance_scale` together for more control over how expressive the model is. For example, a combination high `strength` and `guidance_scale` values gives the model the most creative freedom. + +```py +import torch +from diffusers import AutoPipelineForInpainting +from diffusers.utils import load_image + +pipeline = AutoPipelineForInpainting.from_pretrained( + "runwayml/stable-diffusion-inpainting", torch_dtype=torch.float16, variant="fp16" +).to("cuda") +pipeline.enable_model_cpu_offload() +pipeline.enable_xformers_memory_efficient_attention() + +# load base and mask image +init_image = load_image("https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/inpaint.png").convert("RGB") +mask_image = load_image("https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/inpaint_mask.png").convert("RGB") + +prompt = "concept art digital painting of an elven castle, inspired by lord of the rings, highly detailed, 8k" +image = pipeline(prompt=prompt, image=init_image, mask_image=mask_image, guidance_scale=2.5).images[0] +``` - +
+
+ +
guidance_scale = 2.5
+
+
+ +
guidance_scale = 7.5
+
+
+ +
guidance_scale = 12.5
+
+
+ +### Negative prompt + +A negative prompt assumes the opposite role of a prompt; it guides the model away from generating certain things in an image. This is useful for quickly improving image quality and preventing the model from generating things you don't want. + +```py +import torch +from diffusers import AutoPipelineForInpainting +from diffusers.utils import load_image + +pipeline = AutoPipelineForInpainting.from_pretrained( + "runwayml/stable-diffusion-inpainting", torch_dtype=torch.float16, variant="fp16" +).to("cuda") +pipeline.enable_model_cpu_offload() +pipeline.enable_xformers_memory_efficient_attention() -## Preserving the Unmasked Area of the Image +# load base and mask image +init_image = load_image("https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/inpaint.png").convert("RGB") +mask_image = load_image("https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/inpaint_mask.png").convert("RGB") + +prompt = "concept art digital painting of an elven castle, inspired by lord of the rings, highly detailed, 8k" +negative_prompt = "bad architecture, unstable, poor details, blurry" +image = pipeline(prompt=prompt, negative_prompt=negative_prompt, image=init_image, mask_image=mask_image).images[0] +image +``` -Generally speaking, [`StableDiffusionInpaintPipeline`] (and other inpainting pipelines) will change the unmasked part of the image as well. If this behavior is undesirable, you can force the unmasked area to remain the same as follows: +
+
+ +
negative_prompt = "bad architecture, unstable, poor details, blurry"
+
+
-```python +## Preserve unmasked areas + +The [`AutoPipelineForInpainting`] (and other inpainting pipelines) generally changes the unmasked parts of an image to create a more natural transition between the masked and unmasked region. If this behavior is undesirable, you can force the unmasked area to remain the same. However, forcing the unmasked portion of the image to remain the same may result in some unusual transitions between the unmasked and masked areas. + +```py import PIL import numpy as np import torch -from diffusers import StableDiffusionInpaintPipeline +from diffusers import AutoPipelineForInpainting from diffusers.utils import load_image device = "cuda" -pipeline = StableDiffusionInpaintPipeline.from_pretrained( +pipeline = AutoPipelineForInpainting.from_pretrained( "runwayml/stable-diffusion-inpainting", torch_dtype=torch.float16, ) @@ -121,4 +323,257 @@ unmasked_unchanged_image = PIL.Image.fromarray(unmasked_unchanged_image_arr.roun unmasked_unchanged_image.save("force_unmasked_unchanged.png") ``` -Forcing the unmasked portion of the image to remain the same might result in some weird transitions between the unmasked and masked areas, since the model will typically change the masked and unmasked areas to make the transition more natural. +## Chained inpainting pipelines + +[`AutoPipelineForInpainting`] can be chained with other 🤗 Diffusers pipelines to edit their outputs. This is often useful for improving the output quality from your other diffusion pipelines, and if you're using multiple pipelines, it can be more memory-efficient to chain them together to keep the outputs in latent space and reuse the same pipeline components. + +### Text-to-image-to-inpaint + +Chaining a text-to-image and inpainting pipeline allows you to inpaint the generated image, and you don't have to provide a base image to begin with. This makes it convenient to edit your favorite text-to-image outputs without having to generate an entirely new image. + +Start with the text-to-image pipeline to create a castle: + +```py +import torch +from diffusers import AutoPipelineForText2Image, AutoPipelineForInpainting +from diffusers.utils import load_image + +pipeline = AutoPipelineForText2Image.from_pretrained( + "runwayml/stable-diffusion-v1-5", torch_dtype=torch.float16, variant="fp16", use_safetensors=True +).to("cuda") +pipeline.enable_model_cpu_offload() +pipeline.enable_xformers_memory_efficient_attention() + +image = pipeline("concept art digital painting of an elven castle, inspired by lord of the rings, highly detailed, 8k").images[0] +``` + +Load the mask image of the output from above: + +```py +mask_image = load_image("https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/inpaint_text-chain-mask.png").convert("RGB") +``` + +And let's inpaint the masked area with a waterfall: + +```py +pipeline = AutoPipelineForInpainting.from_pretrained( + "kandinsky-community/kandinsky-2-2-decoder-inpaint", torch_dtype=torch.float16, variant="fp16" +).to("cuda") +pipeline.enable_model_cpu_offload() +pipeline.enable_xformers_memory_efficient_attention() + +prompt = "digital painting of a fantasy waterfall, cloudy" +image = pipeline(prompt=prompt, image=image, mask_image=mask_image).images[0] +image +``` + +
+
+ +
text-to-image
+
+
+ +
inpaint
+
+
+ + +### Inpaint-to-image-to-image + +You can also chain an inpainting pipeline before another pipeline like image-to-image or an upscaler to improve the quality. + +Begin by inpainting an image: + +```py +import torch +from diffusers import AutoPipelineForInpainting, AutoPipelineForImage2Image +from diffusers.utils import load_image + +pipeline = AutoPipelineForInpainting.from_pretrained( + "runwayml/stable-diffusion-inpainting", torch_dtype=torch.float16, variant="fp16" +).to("cuda") +pipeline.enable_model_cpu_offload() +pipeline.enable_xformers_memory_efficient_attention() + +# load base and mask image +init_image = load_image("https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/inpaint.png").convert("RGB") +mask_image = load_image("https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/inpaint_mask.png").convert("RGB") + +prompt = "concept art digital painting of an elven castle, inspired by lord of the rings, highly detailed, 8k" +image = pipeline(prompt=prompt, image=init_image, mask_image=mask_image).images[0] + +# resize image to 1024x1024 for SDXL +image = image.resize((1024, 1024)) +``` + +Now let's pass the image to another inpainting pipeline with SDXL's refiner model to enhance the image details and quality: + +```py +pipeline = AutoPipelineForInpainting.from_pretrained( + "stabilityai/stable-diffusion-xl-refiner-1.0", torch_dtype=torch.float16, variant="fp16" +).to("cuda") +pipeline.enable_model_cpu_offload() +pipeline.enable_xformers_memory_efficient_attention() + +image = pipeline(prompt=prompt, image=image, mask_image=mask_image, output_type="latent").images[0] +``` + + + +It is important to specify `output_type="latent"` in the pipeline to keep all the outputs in latent space to avoid an unnecessary decode-encode step. This only works if the chained pipelines are using the same VAE. For example, in the [Text-to-image-to-inpaint](#text-to-image-to-inpaint) section, Kandinsky 2.2 uses a different VAE class than the Stable Diffusion model so it won't work. But if you use Stable Diffusion v1.5 for both pipelines, then you can keep everything in latent space because they both use [`AutoencoderKL`]. + + + +Finally, you can pass this image to an image-to-image pipeline to put the finishing touches on it. It is more efficient to use the [`~AutoPipelineForImage2Image.from_pipe`] method to reuse the existing pipeline components, and avoid unnecessarily loading all the pipeline components into memory again. + +```py +pipeline = AutoPipelineForImage2Image.from_pipe(pipeline) +pipeline.enable_xformers_memory_efficient_attention() + +image = pipeline(prompt=prompt, image=image).images[0] +``` + +
+
+ +
initial image
+
+
+ +
inpaint
+
+
+ +
image-to-image
+
+
+ +Image-to-image and inpainting are actually very similar tasks. Image-to-image generates a new image that resembles the existing provided image. Inpainting does the same thing, but it only transforms the image area defined by the mask and the rest of the image is unchanged. You can think of inpainting as a more precise tool for making specific changes and image-to-image has a broader scope for making more sweeping changes. + +## Control image generation + +Getting an image to look exactly the way you want is challenging because the denoising process is random. While you can control certain aspects of generation by configuring parameters like `negative_prompt`, there are better and more efficient methods for controlling image generation. + +### Prompt weighting + +Prompt weighting provides a quantifiable way to scale the representation of concepts in a prompt. You can use it to increase or decrease the magnitude of the text embedding vector for each concept in the prompt, which subsequently determines how much of each concept is generated. The [Compel](https://github.com/damian0815/compel) library offers an intuitive syntax for scaling the prompt weights and generating the embeddings. Learn how to create the embeddings in the [Prompt weighting](../using-diffusers/weighted_prompts) guide. + +Once you've generated the embeddings, pass them to the `prompt_embeds` (and `negative_prompt_embeds` if you're using a negative prompt) parameter in the [`AutoPipelineForInpainting`]. The embeddings replace the `prompt` parameter: + +```py +import torch +from diffusers import AutoPipelineForInpainting + +pipeline = AutoPipelineForInpainting.from_pretrained( + "runwayml/stable-diffusion-inpainting", torch_dtype=torch.float16, +).to("cuda") +pipeline.enable_model_cpu_offload() +pipeline.enable_xformers_memory_efficient_attention() + +image = pipeline(prompt_emebds=prompt_embeds, # generated from Compel + negative_prompt_embeds, # generated from Compel + image=init_image, + mask_image=mask_image +).images[0] +``` + +### ControlNet + +ControlNet models are used with other diffusion models like Stable Diffusion, and they provide an even more flexible and accurate way to control how an image is generated. A ControlNet accepts an additional conditioning image input that guides the diffusion model to preserve the features in it. + +For example, let's condition an image with a ControlNet pretrained on inpaint images: + +```py +import torch +import numpy as np +from diffusers import ControlNetModel, StableDiffusionControlNetInpaintPipeline +from diffusers.utils import load_image + +# load ControlNet +controlnet = ControlNetModel.from_pretrained("lllyasviel/control_v11p_sd15_inpaint", torch_dtype=torch.float16, variant="fp16") + +# pass ControlNet to the pipeline +pipeline = StableDiffusionControlNetInpaintPipeline.from_pretrained( + "runwayml/stable-diffusion-inpainting", controlnet=controlnet, torch_dtype=torch.float16, variant="fp16" +).to("cuda") +pipeline.enable_model_cpu_offload() +pipeline.enable_xformers_memory_efficient_attention() + +# load base and mask image +init_image = load_image("https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/inpaint.png").convert("RGB") +mask_image = load_image("https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/inpaint_mask.png").convert("RGB") + +# prepare control image +def make_inpaint_condition(init_image, mask_image): + init_image = np.array(init_image.convert("RGB")).astype(np.float32) / 255.0 + mask_image = np.array(mask_image.convert("L")).astype(np.float32) / 255.0 + + assert init_image.shape[0:1] == mask_image.shape[0:1], "image and image_mask must have the same image size" + init_image[mask_image > 0.5] = -1.0 # set as masked pixel + init_image = np.expand_dims(init_image, 0).transpose(0, 3, 1, 2) + init_image = torch.from_numpy(init_image) + return init_image + +control_image = make_inpaint_condition(init_image, mask_image) +``` + +Now generate an image from the base, mask and control images. You'll notice features of the base image are strongly preserved in the generated image. + +```py +prompt = "concept art digital painting of an elven castle, inspired by lord of the rings, highly detailed, 8k" +image = pipeline(prompt=prompt, image=init_image, mask_image=mask_image, control_image=control_image).images[0] +image +``` + +You can take this a step further and chain it with an image-to-image pipeline to apply a new [style](https://huggingface.co/nitrosocke/elden-ring-diffusion): + +```py +from diffusers import AutoPipelineForImage2Image + +pipeline = AutoPipelineForImage2Image.from_pretrained( + "nitrosocke/elden-ring-diffusion", torch_dtype=torch.float16, +).to("cuda") +pipeline.enable_model_cpu_offload() +pipeline.enable_xformers_memory_efficient_attention() + +prompt = "elden ring style castle" # include the token "elden ring style" in the prompt +negative_prompt = "bad architecture, deformed, disfigured, poor details" + +image = pipeline(prompt, negative_prompt=negative_prompt, image=image).images[0] +image +``` + +
+
+ +
initial image
+
+
+ +
ControlNet inpaint
+
+
+ +
image-to-image
+
+
+ +## Optimize + +It can be difficult and slow to run diffusion models if you're resource constrained, but it dosen't have to be with a few optimization tricks. One of the biggest (and easiest) optimizations you can enable is switching to memory-efficient attention. If you're using PyTorch 2.0, [scaled-dot product attention](../optimization/torch2.0#scaled-dot-product-attention) is automatically enabled and you don't need to do anything else. For non-PyTorch 2.0 users, you can install and use [xFormers](../optimization/xformers)'s implementation of memory-efficient attention. Both options reduce memory usage and accelerate inference. + +You can also offload the model to the GPU to save even more memory: + +```diff ++ pipeline.enable_xformers_memory_efficient_attention() ++ pipeline.enable_model_cpu_offload() +``` + +To speed-up your inference code even more, use [`torch_compile`](../optimization/torch2.0#torch.compile). You should wrap `torch.compile` around the most intensive component in the pipeline which is typically the UNet: + +```py +pipe.unet = torch.compile(pipe.unet, mode="reduce-overhead", fullgraph=True) +``` + +Learn more in the [Reduce memory usage](../optimization/memory) and [Torch 2.0](../optimization/torch2.0) guides. \ No newline at end of file From 016866792de772dadf07a32ee6ad6bae09791f1a Mon Sep 17 00:00:00 2001 From: TimothyAlexisVass <55708319+TimothyAlexisVass@users.noreply.github.com> Date: Fri, 6 Oct 2023 20:20:06 +0200 Subject: [PATCH 08/25] Minor fixes (#5309) tiny fixes --- docs/source/en/using-diffusers/write_own_pipeline.md | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/docs/source/en/using-diffusers/write_own_pipeline.md b/docs/source/en/using-diffusers/write_own_pipeline.md index 42b3e4d6761d..a9243a7b9adc 100644 --- a/docs/source/en/using-diffusers/write_own_pipeline.md +++ b/docs/source/en/using-diffusers/write_own_pipeline.md @@ -112,7 +112,7 @@ As you can see, this is already more complex than the DDPM pipeline which only c -💡 Read the [How does Stable Diffusion work?](https://huggingface.co/blog/stable_diffusion#how-does-stable-diffusion-work) blog for more details about how the VAE, UNet, and text encoder models. +💡 Read the [How does Stable Diffusion work?](https://huggingface.co/blog/stable_diffusion#how-does-stable-diffusion-work) blog for more details about how the VAE, UNet, and text encoder models work. @@ -214,7 +214,7 @@ Next, generate some initial random noise as a starting point for the diffusion p ```py >>> latents = torch.randn( -... (batch_size, unet.in_channels, height // 8, width // 8), +... (batch_size, unet.config.in_channels, height // 8, width // 8), ... generator=generator, ... ) >>> latents = latents.to(torch_device) From dd25ef567927bdd7d29fb1862eaf447bc4716473 Mon Sep 17 00:00:00 2001 From: vedant2003 <97391064+jgyfutub@users.noreply.github.com> Date: Sat, 7 Oct 2023 02:49:38 +0530 Subject: [PATCH 09/25] [Hacktoberfest]Fixing issues #5241 (#5255) * Update pipeline_wuerstchen_prior.py * prior_num_inference_steps updated * height, width, num_inference_steps, and guidance_scale synced * parameters synced * latent_mean, latent_std, and resolution_multiple synced * prior_num_inference_steps changed * Formatted pipeline_wuerstchen_prior.py * Update src/diffusers/pipelines/wuerstchen/pipeline_wuerstchen_prior.py --------- Co-authored-by: Kashif Rasul --- .../pipelines/wuerstchen/pipeline_wuerstchen.py | 4 ++-- .../wuerstchen/pipeline_wuerstchen_combined.py | 2 +- .../wuerstchen/pipeline_wuerstchen_prior.py | 14 ++++++++++---- 3 files changed, 13 insertions(+), 7 deletions(-) diff --git a/src/diffusers/pipelines/wuerstchen/pipeline_wuerstchen.py b/src/diffusers/pipelines/wuerstchen/pipeline_wuerstchen.py index 9189691bd0d5..6caa09a46ce0 100644 --- a/src/diffusers/pipelines/wuerstchen/pipeline_wuerstchen.py +++ b/src/diffusers/pipelines/wuerstchen/pipeline_wuerstchen.py @@ -213,13 +213,13 @@ def __call__( Image Embeddings either extracted from an image or generated by a Prior Model. prompt (`str` or `List[str]`): The prompt or prompts to guide the image generation. - num_inference_steps (`int`, *optional*, defaults to 30): + num_inference_steps (`int`, *optional*, defaults to 12): The number of denoising steps. More denoising steps usually lead to a higher quality image at the expense of slower inference. timesteps (`List[int]`, *optional*): Custom timesteps to use for the denoising process. If not defined, equal spaced `num_inference_steps` timesteps are used. Must be in descending order. - guidance_scale (`float`, *optional*, defaults to 4.0): + guidance_scale (`float`, *optional*, defaults to 0.0): Guidance scale as defined in [Classifier-Free Diffusion Guidance](https://arxiv.org/abs/2207.12598). `decoder_guidance_scale` is defined as `w` of equation 2. of [Imagen Paper](https://arxiv.org/pdf/2205.11487.pdf). Guidance scale is enabled by setting diff --git a/src/diffusers/pipelines/wuerstchen/pipeline_wuerstchen_combined.py b/src/diffusers/pipelines/wuerstchen/pipeline_wuerstchen_combined.py index 6b5ce9530d4c..888d3c0dd74b 100644 --- a/src/diffusers/pipelines/wuerstchen/pipeline_wuerstchen_combined.py +++ b/src/diffusers/pipelines/wuerstchen/pipeline_wuerstchen_combined.py @@ -194,7 +194,7 @@ def __call__( Paper](https://arxiv.org/pdf/2205.11487.pdf). Guidance scale is enabled by setting `prior_guidance_scale > 1`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`, usually at the expense of lower image quality. - prior_num_inference_steps (`Union[int, Dict[float, int]]`, *optional*, defaults to 30): + prior_num_inference_steps (`Union[int, Dict[float, int]]`, *optional*, defaults to 60): The number of prior denoising steps. More denoising steps usually lead to a higher quality image at the expense of slower inference. For more specific timestep spacing, you can pass customized `prior_timesteps` diff --git a/src/diffusers/pipelines/wuerstchen/pipeline_wuerstchen_prior.py b/src/diffusers/pipelines/wuerstchen/pipeline_wuerstchen_prior.py index ca9568f9d39f..dba6d7bb06db 100644 --- a/src/diffusers/pipelines/wuerstchen/pipeline_wuerstchen_prior.py +++ b/src/diffusers/pipelines/wuerstchen/pipeline_wuerstchen_prior.py @@ -82,6 +82,12 @@ class WuerstchenPriorPipeline(DiffusionPipeline): [CLIPTokenizer](https://huggingface.co/docs/transformers/v4.21.0/en/model_doc/clip#transformers.CLIPTokenizer). scheduler ([`DDPMWuerstchenScheduler`]): A scheduler to be used in combination with `prior` to generate image embedding. + latent_mean ('float', *optional*, defaults to 42.0): + Mean value for latent diffusers. + latent_std ('float', *optional*, defaults to 1.0): + Standard value for latent diffusers. + resolution_multiple ('float', *optional*, defaults to 42.67): + Default resolution for multiple images generated. """ model_cpu_offload_seq = "text_encoder->prior" @@ -282,17 +288,17 @@ def __call__( Args: prompt (`str` or `List[str]`): The prompt or prompts to guide the image generation. - height (`int`, *optional*, defaults to 512): + height (`int`, *optional*, defaults to 1024): The height in pixels of the generated image. - width (`int`, *optional*, defaults to 512): + width (`int`, *optional*, defaults to 1024): The width in pixels of the generated image. - num_inference_steps (`int`, *optional*, defaults to 30): + num_inference_steps (`int`, *optional*, defaults to 60): The number of denoising steps. More denoising steps usually lead to a higher quality image at the expense of slower inference. timesteps (`List[int]`, *optional*): Custom timesteps to use for the denoising process. If not defined, equal spaced `num_inference_steps` timesteps are used. Must be in descending order. - guidance_scale (`float`, *optional*, defaults to 4.0): + guidance_scale (`float`, *optional*, defaults to 8.0): Guidance scale as defined in [Classifier-Free Diffusion Guidance](https://arxiv.org/abs/2207.12598). `decoder_guidance_scale` is defined as `w` of equation 2. of [Imagen Paper](https://arxiv.org/pdf/2205.11487.pdf). Guidance scale is enabled by setting From 306dc6e751a45b6b5357fbf512f21e63505c55b7 Mon Sep 17 00:00:00 2001 From: Shubham S Jagtap <63872951+ShubhamJagtap2000@users.noreply.github.com> Date: Sat, 7 Oct 2023 03:20:18 +0530 Subject: [PATCH 10/25] Update README.md (#5267) Co-authored-by: YiYi Xu --- docs/README.md | 16 ++++++++-------- 1 file changed, 8 insertions(+), 8 deletions(-) diff --git a/docs/README.md b/docs/README.md index e6408dc976fd..fd0a3a58b0aa 100644 --- a/docs/README.md +++ b/docs/README.md @@ -128,7 +128,7 @@ When adding a new pipeline: - Possible an end-to-end example of how to use it - Add all the pipeline classes that should be linked in the diffusion model. These classes should be added using our Markdown syntax. By default as follows: -``` +```py ## XXXPipeline [[autodoc]] XXXPipeline @@ -138,7 +138,7 @@ When adding a new pipeline: This will include every public method of the pipeline that is documented, as well as the `__call__` method that is not documented by default. If you just want to add additional methods that are not documented, you can put the list of all methods to add in a list that contains `all`. -``` +```py [[autodoc]] XXXPipeline - all - __call__ @@ -172,7 +172,7 @@ Arguments should be defined with the `Args:` (or `Arguments:` or `Parameters:`) an indentation. The argument should be followed by its type, with its shape if it is a tensor, a colon, and its description: -``` +```py Args: n_layers (`int`): The number of layers of the model. ``` @@ -182,7 +182,7 @@ after the argument. Here's an example showcasing everything so far: -``` +```py Args: input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`): Indices of input sequence tokens in the vocabulary. @@ -196,13 +196,13 @@ Here's an example showcasing everything so far: For optional arguments or arguments with defaults we follow the following syntax: imagine we have a function with the following signature: -``` +```py def my_function(x: str = None, a: float = 1): ``` then its documentation should look like this: -``` +```py Args: x (`str`, *optional*): This argument controls ... @@ -235,14 +235,14 @@ building the return. Here's an example of a single value return: -``` +```py Returns: `List[int]`: A list of integers in the range [0, 1] --- 1 for a special token, 0 for a sequence token. ``` Here's an example of a tuple return, comprising several objects: -``` +```py Returns: `tuple(torch.FloatTensor)` comprising various elements depending on the configuration ([`BertConfig`]) and inputs: - ** loss** (*optional*, returned when `masked_lm_labels` is provided) `torch.FloatTensor` of shape `(1,)` -- From 0513a8cfd8c6a626da0e0bdafc469305a57a723e Mon Sep 17 00:00:00 2001 From: Zeng Xian Date: Sun, 8 Oct 2023 20:54:33 +0800 Subject: [PATCH 11/25] fix typo in train dreambooth lora description (#5332) --- examples/dreambooth/train_dreambooth_lora.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/examples/dreambooth/train_dreambooth_lora.py b/examples/dreambooth/train_dreambooth_lora.py index 5bb6c78b7b74..47de88f338d1 100644 --- a/examples/dreambooth/train_dreambooth_lora.py +++ b/examples/dreambooth/train_dreambooth_lora.py @@ -854,7 +854,7 @@ def main(args): # For Stable Diffusion, it should be equal to: # - down blocks (2x attention layers) * (2x transformer layers) * (3x down blocks) = 12 # - mid blocks (2x attention layers) * (1x transformer layers) * (1x mid blocks) = 2 - # - up blocks (2x attention layers) * (3x transformer layers) * (3x down blocks) = 18 + # - up blocks (2x attention layers) * (3x transformer layers) * (3x up blocks) = 18 # => 32 layers # Set correct lora layers From 6bd55b54bc63fc3c0ae3996f053ec5670524b30e Mon Sep 17 00:00:00 2001 From: chuzh <50682966+rchuzh99@users.noreply.github.com> Date: Mon, 9 Oct 2023 15:54:01 +0800 Subject: [PATCH 12/25] Fix [core/GLIGEN]: TypeError when iterating over 0-d tensor with In-painting mode when EulerAncestralDiscreteScheduler is used (#5305) * fix(gligen_inpaint_pipeline): :bug: Wrap the timestep() 0-d tensor in a list to convert to 1-d tensor. This avoids the TypeError caused by trying to directly iterate over a 0-dimensional tensor in the denoising stage * test(gligen/gligen_text_image): unit test using the EulerAncestralDiscreteScheduler --------- Co-authored-by: zhen-hao.chu Co-authored-by: Sayak Paul --- .../train_custom_diffusion.py | 2 +- .../pipeline_stable_diffusion_gligen.py | 4 +++- ...line_stable_diffusion_gligen_text_image.py | 4 +++- .../test_stable_diffusion_gligen.py | 21 ++++++++++++++++++- ...test_stable_diffusion_gligen_text_image.py | 21 ++++++++++++++++++- 5 files changed, 47 insertions(+), 5 deletions(-) diff --git a/examples/custom_diffusion/train_custom_diffusion.py b/examples/custom_diffusion/train_custom_diffusion.py index 3288fe3258ac..8d90998700f4 100644 --- a/examples/custom_diffusion/train_custom_diffusion.py +++ b/examples/custom_diffusion/train_custom_diffusion.py @@ -207,7 +207,7 @@ def __init__( with open(concept["class_prompt"], "r") as f: class_prompt = f.read().splitlines() - class_img_path = [(x, y) for (x, y) in zip(class_images_path, class_prompt)] + class_img_path = list(zip(class_images_path, class_prompt)) self.class_images_path.extend(class_img_path[:num_class_images]) random.shuffle(self.instance_images_path) diff --git a/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_gligen.py b/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_gligen.py index e39b7cde3d36..f176f08d5d8c 100644 --- a/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_gligen.py +++ b/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_gligen.py @@ -803,7 +803,9 @@ def __call__( if gligen_inpaint_image is not None: gligen_inpaint_latent_with_noise = ( - self.scheduler.add_noise(gligen_inpaint_latent, torch.randn_like(gligen_inpaint_latent), t) + self.scheduler.add_noise( + gligen_inpaint_latent, torch.randn_like(gligen_inpaint_latent), torch.tensor([t]) + ) .expand(latents.shape[0], -1, -1, -1) .clone() ) diff --git a/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_gligen_text_image.py b/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_gligen_text_image.py index 59ba0ab666b3..ba418b4cb3c3 100644 --- a/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_gligen_text_image.py +++ b/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_gligen_text_image.py @@ -965,7 +965,9 @@ def __call__( if gligen_inpaint_image is not None: gligen_inpaint_latent_with_noise = ( - self.scheduler.add_noise(gligen_inpaint_latent, torch.randn_like(gligen_inpaint_latent), t) + self.scheduler.add_noise( + gligen_inpaint_latent, torch.randn_like(gligen_inpaint_latent), torch.tensor([t]) + ) .expand(latents.shape[0], -1, -1, -1) .clone() ) diff --git a/tests/pipelines/stable_diffusion/test_stable_diffusion_gligen.py b/tests/pipelines/stable_diffusion/test_stable_diffusion_gligen.py index 19d44e0cd1d9..388ad9672e15 100644 --- a/tests/pipelines/stable_diffusion/test_stable_diffusion_gligen.py +++ b/tests/pipelines/stable_diffusion/test_stable_diffusion_gligen.py @@ -22,6 +22,7 @@ from diffusers import ( AutoencoderKL, DDIMScheduler, + EulerAncestralDiscreteScheduler, StableDiffusionGLIGENPipeline, UNet2DConditionModel, ) @@ -120,7 +121,7 @@ def get_dummy_inputs(self, device, seed=0): } return inputs - def test_gligen(self): + def test_stable_diffusion_gligen_default_case(self): device = "cpu" # ensure determinism for the device-dependent torch.Generator components = self.get_dummy_components() sd_pipe = StableDiffusionGLIGENPipeline(**components) @@ -136,6 +137,24 @@ def test_gligen(self): assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-2 + def test_stable_diffusion_gligen_k_euler_ancestral(self): + device = "cpu" # ensure determinism for the device-dependent torch.Generator + components = self.get_dummy_components() + sd_pipe = StableDiffusionGLIGENPipeline(**components) + sd_pipe.scheduler = EulerAncestralDiscreteScheduler.from_config(sd_pipe.scheduler.config) + sd_pipe = sd_pipe.to(device) + sd_pipe.set_progress_bar_config(disable=None) + + inputs = self.get_dummy_inputs(device) + output = sd_pipe(**inputs) + image = output.images + image_slice = image[0, -3:, -3:, -1] + + assert image.shape == (1, 64, 64, 3) + expected_slice = np.array([0.425, 0.494, 0.429, 0.469, 0.525, 0.417, 0.533, 0.5, 0.47]) + + assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-2 + def test_attention_slicing_forward_pass(self): super().test_attention_slicing_forward_pass(expected_max_diff=3e-3) diff --git a/tests/pipelines/stable_diffusion/test_stable_diffusion_gligen_text_image.py b/tests/pipelines/stable_diffusion/test_stable_diffusion_gligen_text_image.py index 4e14adc81f42..f8f32643aec1 100644 --- a/tests/pipelines/stable_diffusion/test_stable_diffusion_gligen_text_image.py +++ b/tests/pipelines/stable_diffusion/test_stable_diffusion_gligen_text_image.py @@ -29,6 +29,7 @@ from diffusers import ( AutoencoderKL, DDIMScheduler, + EulerAncestralDiscreteScheduler, StableDiffusionGLIGENTextImagePipeline, UNet2DConditionModel, ) @@ -150,7 +151,7 @@ def get_dummy_inputs(self, device, seed=0): } return inputs - def test_gligen(self): + def test_stable_diffusion_gligen_text_image_default_case(self): device = "cpu" # ensure determinism for the device-dependent torch.Generator components = self.get_dummy_components() sd_pipe = StableDiffusionGLIGENTextImagePipeline(**components) @@ -166,6 +167,24 @@ def test_gligen(self): assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-2 + def test_stable_diffusion_gligen_k_euler_ancestral(self): + device = "cpu" # ensure determinism for the device-dependent torch.Generator + components = self.get_dummy_components() + sd_pipe = StableDiffusionGLIGENTextImagePipeline(**components) + sd_pipe.scheduler = EulerAncestralDiscreteScheduler.from_config(sd_pipe.scheduler.config) + sd_pipe = sd_pipe.to(device) + sd_pipe.set_progress_bar_config(disable=None) + + inputs = self.get_dummy_inputs(device) + image = sd_pipe(**inputs).images + image_slice = image[0, -3:, -3:, -1] + + assert image.shape == (1, 64, 64, 3) + + expected_slice = np.array([0.425, 0.494, 0.429, 0.469, 0.525, 0.417, 0.533, 0.5, 0.47]) + + assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-2 + def test_attention_slicing_forward_pass(self): super().test_attention_slicing_forward_pass(expected_max_diff=3e-3) From cc2c4ae759c915e76f020a5ddd1764b8063dc79d Mon Sep 17 00:00:00 2001 From: Pu Cao <48318302+caopulan@users.noreply.github.com> Date: Mon, 9 Oct 2023 16:08:01 +0800 Subject: [PATCH 13/25] fix inference in custom diffusion (#5329) * Update train_custom_diffusion.py * make style * Empty-Commit --------- Co-authored-by: Sayak Paul --- .../train_custom_diffusion.py | 86 ++++++++++--------- 1 file changed, 44 insertions(+), 42 deletions(-) diff --git a/examples/custom_diffusion/train_custom_diffusion.py b/examples/custom_diffusion/train_custom_diffusion.py index 8d90998700f4..4773446a615b 100644 --- a/examples/custom_diffusion/train_custom_diffusion.py +++ b/examples/custom_diffusion/train_custom_diffusion.py @@ -1214,50 +1214,52 @@ def main(args): if global_step >= args.max_train_steps: break - if accelerator.is_main_process: - images = [] + if accelerator.is_main_process: + images = [] - if args.validation_prompt is not None and global_step % args.validation_steps == 0: - logger.info( - f"Running validation... \n Generating {args.num_validation_images} images with prompt:" - f" {args.validation_prompt}." - ) - # create pipeline - pipeline = DiffusionPipeline.from_pretrained( - args.pretrained_model_name_or_path, - unet=accelerator.unwrap_model(unet), - text_encoder=accelerator.unwrap_model(text_encoder), - tokenizer=tokenizer, - revision=args.revision, - torch_dtype=weight_dtype, - ) - pipeline.scheduler = DPMSolverMultistepScheduler.from_config(pipeline.scheduler.config) - pipeline = pipeline.to(accelerator.device) - pipeline.set_progress_bar_config(disable=True) - - # run inference - generator = torch.Generator(device=accelerator.device).manual_seed(args.seed) - images = [ - pipeline(args.validation_prompt, num_inference_steps=25, generator=generator, eta=1.0).images[0] - for _ in range(args.num_validation_images) - ] - - for tracker in accelerator.trackers: - if tracker.name == "tensorboard": - np_images = np.stack([np.asarray(img) for img in images]) - tracker.writer.add_images("validation", np_images, epoch, dataformats="NHWC") - if tracker.name == "wandb": - tracker.log( - { - "validation": [ - wandb.Image(image, caption=f"{i}: {args.validation_prompt}") - for i, image in enumerate(images) - ] - } - ) + if args.validation_prompt is not None and global_step % args.validation_steps == 0: + logger.info( + f"Running validation... \n Generating {args.num_validation_images} images with prompt:" + f" {args.validation_prompt}." + ) + # create pipeline + pipeline = DiffusionPipeline.from_pretrained( + args.pretrained_model_name_or_path, + unet=accelerator.unwrap_model(unet), + text_encoder=accelerator.unwrap_model(text_encoder), + tokenizer=tokenizer, + revision=args.revision, + torch_dtype=weight_dtype, + ) + pipeline.scheduler = DPMSolverMultistepScheduler.from_config(pipeline.scheduler.config) + pipeline = pipeline.to(accelerator.device) + pipeline.set_progress_bar_config(disable=True) - del pipeline - torch.cuda.empty_cache() + # run inference + generator = torch.Generator(device=accelerator.device).manual_seed(args.seed) + images = [ + pipeline(args.validation_prompt, num_inference_steps=25, generator=generator, eta=1.0).images[ + 0 + ] + for _ in range(args.num_validation_images) + ] + + for tracker in accelerator.trackers: + if tracker.name == "tensorboard": + np_images = np.stack([np.asarray(img) for img in images]) + tracker.writer.add_images("validation", np_images, epoch, dataformats="NHWC") + if tracker.name == "wandb": + tracker.log( + { + "validation": [ + wandb.Image(image, caption=f"{i}: {args.validation_prompt}") + for i, image in enumerate(images) + ] + } + ) + + del pipeline + torch.cuda.empty_cache() # Save the custom diffusion layers accelerator.wait_for_everyone() From 2ed7e05fc26817e5dc04f13a39c62a61c372608b Mon Sep 17 00:00:00 2001 From: Sebastian Date: Mon, 9 Oct 2023 11:19:56 +0200 Subject: [PATCH 14/25] Improve performance of fast test by reducing down blocks (#5290) * Reduce number of down block channels * Remove debug code * Set new excepted image slice values for sdxl euler test --- .../stable_diffusion_xl/test_stable_diffusion_xl.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/tests/pipelines/stable_diffusion_xl/test_stable_diffusion_xl.py b/tests/pipelines/stable_diffusion_xl/test_stable_diffusion_xl.py index 65c7526e3aa2..cebd860a4379 100644 --- a/tests/pipelines/stable_diffusion_xl/test_stable_diffusion_xl.py +++ b/tests/pipelines/stable_diffusion_xl/test_stable_diffusion_xl.py @@ -51,7 +51,7 @@ class StableDiffusionXLPipelineFastTests(PipelineLatentTesterMixin, PipelineTest def get_dummy_components(self): torch.manual_seed(0) unet = UNet2DConditionModel( - block_out_channels=(32, 64), + block_out_channels=(2, 4), layers_per_block=2, sample_size=32, in_channels=4, @@ -66,6 +66,7 @@ def get_dummy_components(self): transformer_layers_per_block=(1, 2), projection_class_embeddings_input_dim=80, # 6 * 8 + 32 cross_attention_dim=64, + norm_num_groups=1, ) scheduler = EulerDiscreteScheduler( beta_start=0.00085, @@ -144,7 +145,7 @@ def test_stable_diffusion_xl_euler(self): image_slice = image[0, -3:, -3:, -1] assert image.shape == (1, 64, 64, 3) - expected_slice = np.array([0.5873, 0.6128, 0.4797, 0.5122, 0.5674, 0.4639, 0.5227, 0.5149, 0.4747]) + expected_slice = np.array([0.5552, 0.5569, 0.4725, 0.4348, 0.4994, 0.4632, 0.5142, 0.5012, 0.47]) assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-2 From c4d66200b7a747a3657e81c188a5d833f23a7d47 Mon Sep 17 00:00:00 2001 From: __mo_san__ <50895527+m0saan@users.noreply.github.com> Date: Mon, 9 Oct 2023 10:26:27 +0100 Subject: [PATCH 15/25] make-fast-test-for-StableDiffusionControlNetPipeline-faster (#5292) * decrease UNet2DConditionModel & ControlNetModel blocks * decrease UNet2DConditionModel & ControlNetModel blocks * decrease even more blocks & number of norm groups * decrease vae block out channels and n of norm goups * fix code style --------- Co-authored-by: Sayak Paul --- tests/pipelines/controlnet/test_controlnet.py | 30 ++++++++++++------- 1 file changed, 20 insertions(+), 10 deletions(-) diff --git a/tests/pipelines/controlnet/test_controlnet.py b/tests/pipelines/controlnet/test_controlnet.py index d11b59cc1510..64baeea910b8 100644 --- a/tests/pipelines/controlnet/test_controlnet.py +++ b/tests/pipelines/controlnet/test_controlnet.py @@ -119,7 +119,7 @@ class ControlNetPipelineFastTests( def get_dummy_components(self): torch.manual_seed(0) unet = UNet2DConditionModel( - block_out_channels=(32, 64), + block_out_channels=(4, 8), layers_per_block=2, sample_size=32, in_channels=4, @@ -127,15 +127,17 @@ def get_dummy_components(self): down_block_types=("DownBlock2D", "CrossAttnDownBlock2D"), up_block_types=("CrossAttnUpBlock2D", "UpBlock2D"), cross_attention_dim=32, + norm_num_groups=1, ) torch.manual_seed(0) controlnet = ControlNetModel( - block_out_channels=(32, 64), + block_out_channels=(4, 8), layers_per_block=2, in_channels=4, down_block_types=("DownBlock2D", "CrossAttnDownBlock2D"), cross_attention_dim=32, conditioning_embedding_out_channels=(16, 32), + norm_num_groups=1, ) torch.manual_seed(0) scheduler = DDIMScheduler( @@ -147,12 +149,13 @@ def get_dummy_components(self): ) torch.manual_seed(0) vae = AutoencoderKL( - block_out_channels=[32, 64], + block_out_channels=[4, 8], in_channels=3, out_channels=3, down_block_types=["DownEncoderBlock2D", "DownEncoderBlock2D"], up_block_types=["UpDecoderBlock2D", "UpDecoderBlock2D"], latent_channels=4, + norm_num_groups=2, ) torch.manual_seed(0) text_encoder_config = CLIPTextConfig( @@ -230,7 +233,7 @@ class StableDiffusionMultiControlNetPipelineFastTests( def get_dummy_components(self): torch.manual_seed(0) unet = UNet2DConditionModel( - block_out_channels=(32, 64), + block_out_channels=(4, 8), layers_per_block=2, sample_size=32, in_channels=4, @@ -238,6 +241,7 @@ def get_dummy_components(self): down_block_types=("DownBlock2D", "CrossAttnDownBlock2D"), up_block_types=("CrossAttnUpBlock2D", "UpBlock2D"), cross_attention_dim=32, + norm_num_groups=1, ) torch.manual_seed(0) @@ -247,23 +251,25 @@ def init_weights(m): m.bias.data.fill_(1.0) controlnet1 = ControlNetModel( - block_out_channels=(32, 64), + block_out_channels=(4, 8), layers_per_block=2, in_channels=4, down_block_types=("DownBlock2D", "CrossAttnDownBlock2D"), cross_attention_dim=32, conditioning_embedding_out_channels=(16, 32), + norm_num_groups=1, ) controlnet1.controlnet_down_blocks.apply(init_weights) torch.manual_seed(0) controlnet2 = ControlNetModel( - block_out_channels=(32, 64), + block_out_channels=(4, 8), layers_per_block=2, in_channels=4, down_block_types=("DownBlock2D", "CrossAttnDownBlock2D"), cross_attention_dim=32, conditioning_embedding_out_channels=(16, 32), + norm_num_groups=1, ) controlnet2.controlnet_down_blocks.apply(init_weights) @@ -277,12 +283,13 @@ def init_weights(m): ) torch.manual_seed(0) vae = AutoencoderKL( - block_out_channels=[32, 64], + block_out_channels=[4, 8], in_channels=3, out_channels=3, down_block_types=["DownEncoderBlock2D", "DownEncoderBlock2D"], up_block_types=["UpDecoderBlock2D", "UpDecoderBlock2D"], latent_channels=4, + norm_num_groups=2, ) torch.manual_seed(0) text_encoder_config = CLIPTextConfig( @@ -415,7 +422,7 @@ class StableDiffusionMultiControlNetOneModelPipelineFastTests( def get_dummy_components(self): torch.manual_seed(0) unet = UNet2DConditionModel( - block_out_channels=(32, 64), + block_out_channels=(4, 8), layers_per_block=2, sample_size=32, in_channels=4, @@ -423,6 +430,7 @@ def get_dummy_components(self): down_block_types=("DownBlock2D", "CrossAttnDownBlock2D"), up_block_types=("CrossAttnUpBlock2D", "UpBlock2D"), cross_attention_dim=32, + norm_num_groups=1, ) torch.manual_seed(0) @@ -432,12 +440,13 @@ def init_weights(m): m.bias.data.fill_(1.0) controlnet = ControlNetModel( - block_out_channels=(32, 64), + block_out_channels=(4, 8), layers_per_block=2, in_channels=4, down_block_types=("DownBlock2D", "CrossAttnDownBlock2D"), cross_attention_dim=32, conditioning_embedding_out_channels=(16, 32), + norm_num_groups=1, ) controlnet.controlnet_down_blocks.apply(init_weights) @@ -451,12 +460,13 @@ def init_weights(m): ) torch.manual_seed(0) vae = AutoencoderKL( - block_out_channels=[32, 64], + block_out_channels=[4, 8], in_channels=3, out_channels=3, down_block_types=["DownEncoderBlock2D", "DownEncoderBlock2D"], up_block_types=["UpDecoderBlock2D", "UpDecoderBlock2D"], latent_channels=4, + norm_num_groups=2, ) torch.manual_seed(0) text_encoder_config = CLIPTextConfig( From bd72927c63447e46ebc595c91adb20a49bec4718 Mon Sep 17 00:00:00 2001 From: Aryan V S Date: Mon, 9 Oct 2023 19:59:23 +0530 Subject: [PATCH 16/25] Improve typehints and docs in `diffusers/models` (#5299) * improvement: add typehints and docs to diffusers/models/activations.py * improvement: add typehints and docs to diffusers/models/resnet.py --- src/diffusers/models/activations.py | 10 +- src/diffusers/models/resnet.py | 195 +++++++++++++++++++++------- 2 files changed, 159 insertions(+), 46 deletions(-) diff --git a/src/diffusers/models/activations.py b/src/diffusers/models/activations.py index 04c978403f41..46da899096c2 100644 --- a/src/diffusers/models/activations.py +++ b/src/diffusers/models/activations.py @@ -1,7 +1,15 @@ from torch import nn -def get_activation(act_fn): +def get_activation(act_fn: str) -> nn.Module: + """Helper function to get activation function from string. + + Args: + act_fn (str): Name of activation function. + + Returns: + nn.Module: Activation function. + """ if act_fn in ["swish", "silu"]: return nn.SiLU() elif act_fn == "mish": diff --git a/src/diffusers/models/resnet.py b/src/diffusers/models/resnet.py index ac66e2271c61..3972b438b076 100644 --- a/src/diffusers/models/resnet.py +++ b/src/diffusers/models/resnet.py @@ -14,7 +14,7 @@ # limitations under the License. from functools import partial -from typing import Optional +from typing import Optional, Tuple, Union import torch import torch.nn as nn @@ -38,9 +38,18 @@ class Upsample1D(nn.Module): option to use a convolution transpose. out_channels (`int`, optional): number of output channels. Defaults to `channels`. + name (`str`, default `conv`): + name of the upsampling 1D layer. """ - def __init__(self, channels, use_conv=False, use_conv_transpose=False, out_channels=None, name="conv"): + def __init__( + self, + channels: int, + use_conv: bool = False, + use_conv_transpose: bool = False, + out_channels: Optional[int] = None, + name: str = "conv", + ): super().__init__() self.channels = channels self.out_channels = out_channels or channels @@ -54,7 +63,7 @@ def __init__(self, channels, use_conv=False, use_conv_transpose=False, out_chann elif use_conv: self.conv = nn.Conv1d(self.channels, self.out_channels, 3, padding=1) - def forward(self, inputs): + def forward(self, inputs: torch.Tensor) -> torch.Tensor: assert inputs.shape[1] == self.channels if self.use_conv_transpose: return self.conv(inputs) @@ -79,9 +88,18 @@ class Downsample1D(nn.Module): number of output channels. Defaults to `channels`. padding (`int`, default `1`): padding for the convolution. + name (`str`, default `conv`): + name of the downsampling 1D layer. """ - def __init__(self, channels, use_conv=False, out_channels=None, padding=1, name="conv"): + def __init__( + self, + channels: int, + use_conv: bool = False, + out_channels: Optional[int] = None, + padding: int = 1, + name: str = "conv", + ): super().__init__() self.channels = channels self.out_channels = out_channels or channels @@ -96,7 +114,7 @@ def __init__(self, channels, use_conv=False, out_channels=None, padding=1, name= assert self.channels == self.out_channels self.conv = nn.AvgPool1d(kernel_size=stride, stride=stride) - def forward(self, inputs): + def forward(self, inputs: torch.Tensor) -> torch.Tensor: assert inputs.shape[1] == self.channels return self.conv(inputs) @@ -113,9 +131,18 @@ class Upsample2D(nn.Module): option to use a convolution transpose. out_channels (`int`, optional): number of output channels. Defaults to `channels`. + name (`str`, default `conv`): + name of the upsampling 2D layer. """ - def __init__(self, channels, use_conv=False, use_conv_transpose=False, out_channels=None, name="conv"): + def __init__( + self, + channels: int, + use_conv: bool = False, + use_conv_transpose: bool = False, + out_channels: Optional[int] = None, + name: str = "conv", + ): super().__init__() self.channels = channels self.out_channels = out_channels or channels @@ -135,7 +162,7 @@ def __init__(self, channels, use_conv=False, use_conv_transpose=False, out_chann else: self.Conv2d_0 = conv - def forward(self, hidden_states, output_size=None, scale: float = 1.0): + def forward(self, hidden_states: torch.Tensor, output_size: Optional[int] = None, scale: float = 1.0): assert hidden_states.shape[1] == self.channels if self.use_conv_transpose: @@ -191,9 +218,18 @@ class Downsample2D(nn.Module): number of output channels. Defaults to `channels`. padding (`int`, default `1`): padding for the convolution. + name (`str`, default `conv`): + name of the downsampling 2D layer. """ - def __init__(self, channels, use_conv=False, out_channels=None, padding=1, name="conv"): + def __init__( + self, + channels: int, + use_conv: bool = False, + out_channels: Optional[int] = None, + padding: int = 1, + name: str = "conv", + ): super().__init__() self.channels = channels self.out_channels = out_channels or channels @@ -246,7 +282,13 @@ class FirUpsample2D(nn.Module): kernel for the FIR filter. """ - def __init__(self, channels=None, out_channels=None, use_conv=False, fir_kernel=(1, 3, 3, 1)): + def __init__( + self, + channels: int = None, + out_channels: Optional[int] = None, + use_conv: bool = False, + fir_kernel: Tuple[int, int, int, int] = (1, 3, 3, 1), + ): super().__init__() out_channels = out_channels if out_channels else channels if use_conv: @@ -255,7 +297,14 @@ def __init__(self, channels=None, out_channels=None, use_conv=False, fir_kernel= self.fir_kernel = fir_kernel self.out_channels = out_channels - def _upsample_2d(self, hidden_states, weight=None, kernel=None, factor=2, gain=1): + def _upsample_2d( + self, + hidden_states: torch.Tensor, + weight: Optional[torch.Tensor] = None, + kernel: Optional[torch.FloatTensor] = None, + factor: int = 2, + gain: float = 1, + ) -> torch.Tensor: """Fused `upsample_2d()` followed by `Conv2d()`. Padding is performed only once at the beginning, not between the operations. The fused op is considerably more @@ -335,7 +384,7 @@ def _upsample_2d(self, hidden_states, weight=None, kernel=None, factor=2, gain=1 return output - def forward(self, hidden_states): + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: if self.use_conv: height = self._upsample_2d(hidden_states, self.Conv2d_0.weight, kernel=self.fir_kernel) height = height + self.Conv2d_0.bias.reshape(1, -1, 1, 1) @@ -359,7 +408,13 @@ class FirDownsample2D(nn.Module): kernel for the FIR filter. """ - def __init__(self, channels=None, out_channels=None, use_conv=False, fir_kernel=(1, 3, 3, 1)): + def __init__( + self, + channels: int = None, + out_channels: Optional[int] = None, + use_conv: bool = False, + fir_kernel: Tuple[int, int, int, int] = (1, 3, 3, 1), + ): super().__init__() out_channels = out_channels if out_channels else channels if use_conv: @@ -368,7 +423,14 @@ def __init__(self, channels=None, out_channels=None, use_conv=False, fir_kernel= self.use_conv = use_conv self.out_channels = out_channels - def _downsample_2d(self, hidden_states, weight=None, kernel=None, factor=2, gain=1): + def _downsample_2d( + self, + hidden_states: torch.Tensor, + weight: Optional[torch.Tensor] = None, + kernel: Optional[torch.FloatTensor] = None, + factor: int = 2, + gain: float = 1, + ) -> torch.Tensor: """Fused `Conv2d()` followed by `downsample_2d()`. Padding is performed only once at the beginning, not between the operations. The fused op is considerably more efficient than performing the same calculation using standard TensorFlow ops. It supports gradients of @@ -422,7 +484,7 @@ def _downsample_2d(self, hidden_states, weight=None, kernel=None, factor=2, gain return output - def forward(self, hidden_states): + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: if self.use_conv: downsample_input = self._downsample_2d(hidden_states, weight=self.Conv2d_0.weight, kernel=self.fir_kernel) hidden_states = downsample_input + self.Conv2d_0.bias.reshape(1, -1, 1, 1) @@ -434,14 +496,20 @@ def forward(self, hidden_states): # downsample/upsample layer used in k-upscaler, might be able to use FirDownsample2D/DirUpsample2D instead class KDownsample2D(nn.Module): - def __init__(self, pad_mode="reflect"): + r"""A 2D K-downsampling layer. + + Parameters: + pad_mode (`str`, *optional*, default to `"reflect"`): the padding mode to use. + """ + + def __init__(self, pad_mode: str = "reflect"): super().__init__() self.pad_mode = pad_mode kernel_1d = torch.tensor([[1 / 8, 3 / 8, 3 / 8, 1 / 8]]) self.pad = kernel_1d.shape[1] // 2 - 1 self.register_buffer("kernel", kernel_1d.T @ kernel_1d, persistent=False) - def forward(self, inputs): + def forward(self, inputs: torch.Tensor) -> torch.Tensor: inputs = F.pad(inputs, (self.pad,) * 4, self.pad_mode) weight = inputs.new_zeros([inputs.shape[1], inputs.shape[1], self.kernel.shape[0], self.kernel.shape[1]]) indices = torch.arange(inputs.shape[1], device=inputs.device) @@ -451,14 +519,20 @@ def forward(self, inputs): class KUpsample2D(nn.Module): - def __init__(self, pad_mode="reflect"): + r"""A 2D K-upsampling layer. + + Parameters: + pad_mode (`str`, *optional*, default to `"reflect"`): the padding mode to use. + """ + + def __init__(self, pad_mode: str = "reflect"): super().__init__() self.pad_mode = pad_mode kernel_1d = torch.tensor([[1 / 8, 3 / 8, 3 / 8, 1 / 8]]) * 2 self.pad = kernel_1d.shape[1] // 2 - 1 self.register_buffer("kernel", kernel_1d.T @ kernel_1d, persistent=False) - def forward(self, inputs): + def forward(self, inputs: torch.Tensor) -> torch.Tensor: inputs = F.pad(inputs, ((self.pad + 1) // 2,) * 4, self.pad_mode) weight = inputs.new_zeros([inputs.shape[1], inputs.shape[1], self.kernel.shape[0], self.kernel.shape[1]]) indices = torch.arange(inputs.shape[1], device=inputs.device) @@ -501,23 +575,23 @@ class ResnetBlock2D(nn.Module): def __init__( self, *, - in_channels, - out_channels=None, - conv_shortcut=False, - dropout=0.0, - temb_channels=512, - groups=32, - groups_out=None, - pre_norm=True, - eps=1e-6, - non_linearity="swish", - skip_time_act=False, - time_embedding_norm="default", # default, scale_shift, ada_group, spatial - kernel=None, - output_scale_factor=1.0, - use_in_shortcut=None, - up=False, - down=False, + in_channels: int, + out_channels: Optional[int] = None, + conv_shortcut: bool = False, + dropout: float = 0.0, + temb_channels: int = 512, + groups: int = 32, + groups_out: Optional[int] = None, + pre_norm: bool = True, + eps: float = 1e-6, + non_linearity: str = "swish", + skip_time_act: bool = False, + time_embedding_norm: str = "default", # default, scale_shift, ada_group, spatial + kernel: Optional[torch.FloatTensor] = None, + output_scale_factor: float = 1.0, + use_in_shortcut: Optional[bool] = None, + up: bool = False, + down: bool = False, conv_shortcut_bias: bool = True, conv_2d_out_channels: Optional[int] = None, ): @@ -667,7 +741,7 @@ def forward(self, input_tensor, temb, scale: float = 1.0): # unet_rl.py -def rearrange_dims(tensor): +def rearrange_dims(tensor: torch.Tensor) -> torch.Tensor: if len(tensor.shape) == 2: return tensor[:, :, None] if len(tensor.shape) == 3: @@ -681,16 +755,24 @@ def rearrange_dims(tensor): class Conv1dBlock(nn.Module): """ Conv1d --> GroupNorm --> Mish + + Parameters: + inp_channels (`int`): Number of input channels. + out_channels (`int`): Number of output channels. + kernel_size (`int` or `tuple`): Size of the convolving kernel. + n_groups (`int`, default `8`): Number of groups to separate the channels into. """ - def __init__(self, inp_channels, out_channels, kernel_size, n_groups=8): + def __init__( + self, inp_channels: int, out_channels: int, kernel_size: Union[int, Tuple[int, int]], n_groups: int = 8 + ): super().__init__() self.conv1d = nn.Conv1d(inp_channels, out_channels, kernel_size, padding=kernel_size // 2) self.group_norm = nn.GroupNorm(n_groups, out_channels) self.mish = nn.Mish() - def forward(self, inputs): + def forward(self, inputs: torch.Tensor) -> torch.Tensor: intermediate_repr = self.conv1d(inputs) intermediate_repr = rearrange_dims(intermediate_repr) intermediate_repr = self.group_norm(intermediate_repr) @@ -701,7 +783,19 @@ def forward(self, inputs): # unet_rl.py class ResidualTemporalBlock1D(nn.Module): - def __init__(self, inp_channels, out_channels, embed_dim, kernel_size=5): + """ + Residual 1D block with temporal convolutions. + + Parameters: + inp_channels (`int`): Number of input channels. + out_channels (`int`): Number of output channels. + embed_dim (`int`): Embedding dimension. + kernel_size (`int` or `tuple`): Size of the convolving kernel. + """ + + def __init__( + self, inp_channels: int, out_channels: int, embed_dim: int, kernel_size: Union[int, Tuple[int, int]] = 5 + ): super().__init__() self.conv_in = Conv1dBlock(inp_channels, out_channels, kernel_size) self.conv_out = Conv1dBlock(out_channels, out_channels, kernel_size) @@ -713,7 +807,7 @@ def __init__(self, inp_channels, out_channels, embed_dim, kernel_size=5): nn.Conv1d(inp_channels, out_channels, 1) if inp_channels != out_channels else nn.Identity() ) - def forward(self, inputs, t): + def forward(self, inputs: torch.Tensor, t: torch.Tensor) -> torch.Tensor: """ Args: inputs : [ batch_size x inp_channels x horizon ] @@ -729,7 +823,9 @@ def forward(self, inputs, t): return out + self.residual_conv(inputs) -def upsample_2d(hidden_states, kernel=None, factor=2, gain=1): +def upsample_2d( + hidden_states: torch.Tensor, kernel: Optional[torch.FloatTensor] = None, factor: int = 2, gain: float = 1 +) -> torch.Tensor: r"""Upsample2D a batch of 2D images with the given filter. Accepts a batch of 2D images of the shape `[N, C, H, W]` or `[N, H, W, C]` and upsamples each image with the given filter. The filter is normalized so that if the input pixels are constant, they will be scaled by the specified @@ -766,7 +862,9 @@ def upsample_2d(hidden_states, kernel=None, factor=2, gain=1): return output -def downsample_2d(hidden_states, kernel=None, factor=2, gain=1): +def downsample_2d( + hidden_states: torch.Tensor, kernel: Optional[torch.FloatTensor] = None, factor: int = 2, gain: float = 1 +) -> torch.Tensor: r"""Downsample2D a batch of 2D images with the given filter. Accepts a batch of 2D images of the shape `[N, C, H, W]` or `[N, H, W, C]` and downsamples each image with the given filter. The filter is normalized so that if the input pixels are constant, they will be scaled by the @@ -801,7 +899,9 @@ def downsample_2d(hidden_states, kernel=None, factor=2, gain=1): return output -def upfirdn2d_native(tensor, kernel, up=1, down=1, pad=(0, 0)): +def upfirdn2d_native( + tensor: torch.Tensor, kernel: torch.Tensor, up: int = 1, down: int = 1, pad: Tuple[int, int] = (0, 0) +) -> torch.Tensor: up_x = up_y = up down_x = down_y = down pad_x0 = pad_y0 = pad[0] @@ -849,9 +949,14 @@ class TemporalConvLayer(nn.Module): """ Temporal convolutional layer that can be used for video (sequence of images) input Code mostly copied from: https://github.com/modelscope/modelscope/blob/1509fdb973e5871f37148a4b5e5964cafd43e64d/modelscope/models/multi_modal/video_synthesis/unet_sd.py#L1016 + + Parameters: + in_dim (`int`): Number of input channels. + out_dim (`int`): Number of output channels. + dropout (`float`, *optional*, defaults to `0.0`): The dropout probability to use. """ - def __init__(self, in_dim, out_dim=None, dropout=0.0): + def __init__(self, in_dim: int, out_dim: Optional[int] = None, dropout: float = 0.0): super().__init__() out_dim = out_dim or in_dim self.in_dim = in_dim @@ -884,7 +989,7 @@ def __init__(self, in_dim, out_dim=None, dropout=0.0): nn.init.zeros_(self.conv4[-1].weight) nn.init.zeros_(self.conv4[-1].bias) - def forward(self, hidden_states, num_frames=1): + def forward(self, hidden_states: torch.Tensor, num_frames: int = 1) -> torch.Tensor: hidden_states = ( hidden_states[None, :].reshape((-1, num_frames) + hidden_states.shape[1:]).permute(0, 2, 1, 3, 4) ) From e2c0208c86fa04040900be74b7b880085b69c5a7 Mon Sep 17 00:00:00 2001 From: Brian Yarbrough <6315292+byarbrough@users.noreply.github.com> Date: Mon, 9 Oct 2023 09:08:55 -0600 Subject: [PATCH 17/25] Add py.typed for PEP 561 compliance (#5326) See #5325 --- setup.py | 1 + src/diffusers/py.typed | 0 2 files changed, 1 insertion(+) create mode 100644 src/diffusers/py.typed diff --git a/setup.py b/setup.py index a2201ac5b3b1..16a77c1fe313 100644 --- a/setup.py +++ b/setup.py @@ -255,6 +255,7 @@ def run(self): url="https://github.com/huggingface/diffusers", package_dir={"": "src"}, packages=find_packages("src"), + package_data={"diffusers": ["py.typed"]}, include_package_data=True, python_requires=">=3.8.0", install_requires=list(install_requires), diff --git a/src/diffusers/py.typed b/src/diffusers/py.typed new file mode 100644 index 000000000000..e69de29bb2d1 From 8d314c96eeb2c953415c5630f3df5620bab11c58 Mon Sep 17 00:00:00 2001 From: Aryan V S Date: Mon, 9 Oct 2023 20:39:16 +0530 Subject: [PATCH 18/25] [HacktoberFest] Add missing docstrings to diffusers/models (#5248) * add missing docstrings * chore: run make quality * improvement: include docs suggestion by @yiyixuxu --------- Co-authored-by: Patrick von Platen --- src/diffusers/models/controlnet.py | 6 +++++ src/diffusers/models/dual_transformer_2d.py | 8 ++++-- src/diffusers/models/transformer_2d.py | 8 ++++++ src/diffusers/models/transformer_temporal.py | 6 +++++ src/diffusers/models/unet_2d_condition.py | 20 +++++++++++++++ .../models/unet_2d_condition_flax.py | 7 ++++++ src/diffusers/models/unet_3d_condition.py | 17 +++++++++++++ .../versatile_diffusion/modeling_text_unet.py | 25 +++++++++++++++---- 8 files changed, 90 insertions(+), 7 deletions(-) diff --git a/src/diffusers/models/controlnet.py b/src/diffusers/models/controlnet.py index 1a82b0421f88..c0d2da9b8c5f 100644 --- a/src/diffusers/models/controlnet.py +++ b/src/diffusers/models/controlnet.py @@ -671,7 +671,13 @@ def forward( class_labels (`torch.Tensor`, *optional*, defaults to `None`): Optional class labels for conditioning. Their embeddings will be summed with the timestep embeddings. timestep_cond (`torch.Tensor`, *optional*, defaults to `None`): + Additional conditional embeddings for timestep. If provided, the embeddings will be summed with the + timestep_embedding passed through the `self.time_embedding` layer to obtain the final timestep + embeddings. attention_mask (`torch.Tensor`, *optional*, defaults to `None`): + An attention mask of shape `(batch, key_tokens)` is applied to `encoder_hidden_states`. If `1` the mask + is kept, otherwise if `0` it is discarded. Mask will be converted into a bias, which adds large + negative values to the attention scores corresponding to "discard" tokens. added_cond_kwargs (`dict`): Additional conditions for the Stable Diffusion XL UNet. cross_attention_kwargs (`dict[str]`, *optional*, defaults to `None`): diff --git a/src/diffusers/models/dual_transformer_2d.py b/src/diffusers/models/dual_transformer_2d.py index 3db7e73ca6af..02568298409c 100644 --- a/src/diffusers/models/dual_transformer_2d.py +++ b/src/diffusers/models/dual_transformer_2d.py @@ -107,14 +107,18 @@ def forward( Args: hidden_states ( When discrete, `torch.LongTensor` of shape `(batch size, num latent pixels)`. When continuous, `torch.FloatTensor` of shape `(batch size, channel, height, width)`): Input - hidden_states + hidden_states. encoder_hidden_states ( `torch.LongTensor` of shape `(batch size, encoder_hidden_states dim)`, *optional*): Conditional embeddings for cross attention layer. If not given, cross-attention defaults to self-attention. timestep ( `torch.long`, *optional*): Optional timestep to be applied as an embedding in AdaLayerNorm's. Used to indicate denoising step. attention_mask (`torch.FloatTensor`, *optional*): - Optional attention mask to be applied in Attention + Optional attention mask to be applied in Attention. + cross_attention_kwargs (`dict`, *optional*): + A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under + `self.processor` in + [diffusers.models.attention_processor](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py). return_dict (`bool`, *optional*, defaults to `True`): Whether or not to return a [`models.unet_2d_condition.UNet2DConditionOutput`] instead of a plain tuple. diff --git a/src/diffusers/models/transformer_2d.py b/src/diffusers/models/transformer_2d.py index c96aef65f339..e7780a7bca3d 100644 --- a/src/diffusers/models/transformer_2d.py +++ b/src/diffusers/models/transformer_2d.py @@ -235,6 +235,14 @@ def forward( class_labels ( `torch.LongTensor` of shape `(batch size, num classes)`, *optional*): Used to indicate class labels conditioning. Optional class labels to be applied as an embedding in `AdaLayerZeroNorm`. + cross_attention_kwargs ( `Dict[str, Any]`, *optional*): + A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under + `self.processor` in + [diffusers.models.attention_processor](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py). + attention_mask ( `torch.Tensor`, *optional*): + An attention mask of shape `(batch, key_tokens)` is applied to `encoder_hidden_states`. If `1` the mask + is kept, otherwise if `0` it is discarded. Mask will be converted into a bias, which adds large + negative values to the attention scores corresponding to "discard" tokens. encoder_attention_mask ( `torch.Tensor`, *optional*): Cross-attention mask applied to `encoder_hidden_states`. Two formats supported: diff --git a/src/diffusers/models/transformer_temporal.py b/src/diffusers/models/transformer_temporal.py index cfafdb055bcf..d002cb3315fa 100644 --- a/src/diffusers/models/transformer_temporal.py +++ b/src/diffusers/models/transformer_temporal.py @@ -128,6 +128,12 @@ def forward( class_labels ( `torch.LongTensor` of shape `(batch size, num classes)`, *optional*): Used to indicate class labels conditioning. Optional class labels to be applied as an embedding in `AdaLayerZeroNorm`. + num_frames (`int`, *optional*, defaults to 1): + The number of frames to be processed per batch. This is used to reshape the hidden states. + cross_attention_kwargs (`dict`, *optional*): + A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under + `self.processor` in + [diffusers.models.attention_processor](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py). return_dict (`bool`, *optional*, defaults to `True`): Whether or not to return a [`~models.unet_2d_condition.UNet2DConditionOutput`] instead of a plain tuple. diff --git a/src/diffusers/models/unet_2d_condition.py b/src/diffusers/models/unet_2d_condition.py index 52c3fc141e59..f858a7685360 100644 --- a/src/diffusers/models/unet_2d_condition.py +++ b/src/diffusers/models/unet_2d_condition.py @@ -790,6 +790,26 @@ def forward( timestep (`torch.FloatTensor` or `float` or `int`): The number of timesteps to denoise an input. encoder_hidden_states (`torch.FloatTensor`): The encoder hidden states with shape `(batch, sequence_length, feature_dim)`. + class_labels (`torch.Tensor`, *optional*, defaults to `None`): + Optional class labels for conditioning. Their embeddings will be summed with the timestep embeddings. + timestep_cond: (`torch.Tensor`, *optional*, defaults to `None`): + Conditional embeddings for timestep. If provided, the embeddings will be summed with the samples passed + through the `self.time_embedding` layer to obtain the timestep embeddings. + attention_mask (`torch.Tensor`, *optional*, defaults to `None`): + An attention mask of shape `(batch, key_tokens)` is applied to `encoder_hidden_states`. If `1` the mask + is kept, otherwise if `0` it is discarded. Mask will be converted into a bias, which adds large + negative values to the attention scores corresponding to "discard" tokens. + cross_attention_kwargs (`dict`, *optional*): + A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under + `self.processor` in + [diffusers.models.attention_processor](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py). + added_cond_kwargs: (`dict`, *optional*): + A kwargs dictionary containing additional embeddings that if specified are added to the embeddings that + are passed along to the UNet blocks. + down_block_additional_residuals: (`tuple` of `torch.Tensor`, *optional*): + A tuple of tensors that if specified are added to the residuals of down unet blocks. + mid_block_additional_residual: (`torch.Tensor`, *optional*): + A tensor that if specified is added to the residual of the middle unet block. encoder_attention_mask (`torch.Tensor`): A cross-attention mask of shape `(batch, sequence_length)` is applied to `encoder_hidden_states`. If `True` the mask is kept, otherwise if `False` it is discarded. Mask will be converted into a bias, diff --git a/src/diffusers/models/unet_2d_condition_flax.py b/src/diffusers/models/unet_2d_condition_flax.py index a3aebde7bf16..a56db67b6a4e 100644 --- a/src/diffusers/models/unet_2d_condition_flax.py +++ b/src/diffusers/models/unet_2d_condition_flax.py @@ -334,6 +334,13 @@ def __call__( sample (`jnp.ndarray`): (batch, channel, height, width) noisy inputs tensor timestep (`jnp.ndarray` or `float` or `int`): timesteps encoder_hidden_states (`jnp.ndarray`): (batch_size, sequence_length, hidden_size) encoder hidden states + added_cond_kwargs: (`dict`, *optional*): + A kwargs dictionary containing additional embeddings that if specified are added to the embeddings that + are passed along to the UNet blocks. + down_block_additional_residuals: (`tuple` of `torch.Tensor`, *optional*): + A tuple of tensors that if specified are added to the residuals of down unet blocks. + mid_block_additional_residual: (`torch.Tensor`, *optional*): + A tensor that if specified is added to the residual of the middle unet block. return_dict (`bool`, *optional*, defaults to `True`): Whether or not to return a [`models.unet_2d_condition_flax.FlaxUNet2DConditionOutput`] instead of a plain tuple. diff --git a/src/diffusers/models/unet_3d_condition.py b/src/diffusers/models/unet_3d_condition.py index 4e6de97390e7..2ab1d4060e17 100644 --- a/src/diffusers/models/unet_3d_condition.py +++ b/src/diffusers/models/unet_3d_condition.py @@ -519,6 +519,23 @@ def forward( timestep (`torch.FloatTensor` or `float` or `int`): The number of timesteps to denoise an input. encoder_hidden_states (`torch.FloatTensor`): The encoder hidden states with shape `(batch, sequence_length, feature_dim)`. + class_labels (`torch.Tensor`, *optional*, defaults to `None`): + Optional class labels for conditioning. Their embeddings will be summed with the timestep embeddings. + timestep_cond: (`torch.Tensor`, *optional*, defaults to `None`): + Conditional embeddings for timestep. If provided, the embeddings will be summed with the samples passed + through the `self.time_embedding` layer to obtain the timestep embeddings. + attention_mask (`torch.Tensor`, *optional*, defaults to `None`): + An attention mask of shape `(batch, key_tokens)` is applied to `encoder_hidden_states`. If `1` the mask + is kept, otherwise if `0` it is discarded. Mask will be converted into a bias, which adds large + negative values to the attention scores corresponding to "discard" tokens. + cross_attention_kwargs (`dict`, *optional*): + A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under + `self.processor` in + [diffusers.models.attention_processor](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py). + down_block_additional_residuals: (`tuple` of `torch.Tensor`, *optional*): + A tuple of tensors that if specified are added to the residuals of down unet blocks. + mid_block_additional_residual: (`torch.Tensor`, *optional*): + A tensor that if specified is added to the residual of the middle unet block. return_dict (`bool`, *optional*, defaults to `True`): Whether or not to return a [`~models.unet_3d_condition.UNet3DConditionOutput`] instead of a plain tuple. diff --git a/src/diffusers/pipelines/versatile_diffusion/modeling_text_unet.py b/src/diffusers/pipelines/versatile_diffusion/modeling_text_unet.py index 7d1ca1d934d8..ad8c326108b5 100644 --- a/src/diffusers/pipelines/versatile_diffusion/modeling_text_unet.py +++ b/src/diffusers/pipelines/versatile_diffusion/modeling_text_unet.py @@ -999,6 +999,26 @@ def forward( timestep (`torch.FloatTensor` or `float` or `int`): The number of timesteps to denoise an input. encoder_hidden_states (`torch.FloatTensor`): The encoder hidden states with shape `(batch, sequence_length, feature_dim)`. + class_labels (`torch.Tensor`, *optional*, defaults to `None`): + Optional class labels for conditioning. Their embeddings will be summed with the timestep embeddings. + timestep_cond: (`torch.Tensor`, *optional*, defaults to `None`): + Conditional embeddings for timestep. If provided, the embeddings will be summed with the samples passed + through the `self.time_embedding` layer to obtain the timestep embeddings. + attention_mask (`torch.Tensor`, *optional*, defaults to `None`): + An attention mask of shape `(batch, key_tokens)` is applied to `encoder_hidden_states`. If `1` the mask + is kept, otherwise if `0` it is discarded. Mask will be converted into a bias, which adds large + negative values to the attention scores corresponding to "discard" tokens. + cross_attention_kwargs (`dict`, *optional*): + A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under + `self.processor` in + [diffusers.models.attention_processor](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py). + added_cond_kwargs: (`dict`, *optional*): + A kwargs dictionary containing additional embeddings that if specified are added to the embeddings that + are passed along to the UNet blocks. + down_block_additional_residuals: (`tuple` of `torch.Tensor`, *optional*): + A tuple of tensors that if specified are added to the residuals of down blocks. + mid_block_additional_residual: (`torch.Tensor`, *optional*): + A tensor that if specified is added to the residual of the middle block. encoder_attention_mask (`torch.Tensor`): A cross-attention mask of shape `(batch, sequence_length)` is applied to `encoder_hidden_states`. If `True` the mask is kept, otherwise if `False` it is discarded. Mask will be converted into a bias, @@ -1006,11 +1026,6 @@ def forward( return_dict (`bool`, *optional*, defaults to `True`): Whether or not to return a [`~models.unet_2d_condition.UNet2DConditionOutput`] instead of a plain tuple. - cross_attention_kwargs (`dict`, *optional*): - A kwargs dictionary that if specified is passed along to the [`AttnProcessor`]. - added_cond_kwargs: (`dict`, *optional*): - A kwargs dictionary containin additional embeddings that if specified are added to the embeddings that - are passed along to the UNet blocks. Returns: [`~models.unet_2d_condition.UNet2DConditionOutput`] or `tuple`: From d199bc62ecd0d9d94435ac1146c0fadda86b9459 Mon Sep 17 00:00:00 2001 From: Patrick von Platen Date: Mon, 9 Oct 2023 17:12:12 +0200 Subject: [PATCH 19/25] make style --- .../pipelines/versatile_diffusion/modeling_text_unet.py | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/src/diffusers/pipelines/versatile_diffusion/modeling_text_unet.py b/src/diffusers/pipelines/versatile_diffusion/modeling_text_unet.py index ad8c326108b5..4e50bbefe933 100644 --- a/src/diffusers/pipelines/versatile_diffusion/modeling_text_unet.py +++ b/src/diffusers/pipelines/versatile_diffusion/modeling_text_unet.py @@ -1016,9 +1016,9 @@ def forward( A kwargs dictionary containing additional embeddings that if specified are added to the embeddings that are passed along to the UNet blocks. down_block_additional_residuals: (`tuple` of `torch.Tensor`, *optional*): - A tuple of tensors that if specified are added to the residuals of down blocks. + A tuple of tensors that if specified are added to the residuals of down unet blocks. mid_block_additional_residual: (`torch.Tensor`, *optional*): - A tensor that if specified is added to the residual of the middle block. + A tensor that if specified is added to the residual of the middle unet block. encoder_attention_mask (`torch.Tensor`): A cross-attention mask of shape `(batch, sequence_length)` is applied to `encoder_hidden_states`. If `True` the mask is kept, otherwise if `False` it is discarded. Mask will be converted into a bias, @@ -1026,6 +1026,11 @@ def forward( return_dict (`bool`, *optional*, defaults to `True`): Whether or not to return a [`~models.unet_2d_condition.UNet2DConditionOutput`] instead of a plain tuple. + cross_attention_kwargs (`dict`, *optional*): + A kwargs dictionary that if specified is passed along to the [`AttnProcessor`]. + added_cond_kwargs: (`dict`, *optional*): + A kwargs dictionary containin additional embeddings that if specified are added to the embeddings that + are passed along to the UNet blocks. Returns: [`~models.unet_2d_condition.UNet2DConditionOutput`] or `tuple`: From 35952e61c18dfcabb9e60b54c6205ca7f0467962 Mon Sep 17 00:00:00 2001 From: Jonathan Whitaker Date: Mon, 9 Oct 2023 08:20:12 -0700 Subject: [PATCH 20/25] Fix links in docs to adapter code (#5323) Update adapter.md to fix links to adapter pipelines --- docs/source/en/api/pipelines/stable_diffusion/adapter.md | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/docs/source/en/api/pipelines/stable_diffusion/adapter.md b/docs/source/en/api/pipelines/stable_diffusion/adapter.md index 4c7415ddb02b..cf3aca4bfa52 100644 --- a/docs/source/en/api/pipelines/stable_diffusion/adapter.md +++ b/docs/source/en/api/pipelines/stable_diffusion/adapter.md @@ -28,8 +28,8 @@ This model was contributed by the community contributor [HimariO](https://github | Pipeline | Tasks | Demo |---|---|:---:| -| [StableDiffusionAdapterPipeline](https://github.com/huggingface/diffusers/blob/main/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_adapter.py) | *Text-to-Image Generation with T2I-Adapter Conditioning* | - -| [StableDiffusionXLAdapterPipeline](https://github.com/huggingface/diffusers/blob/main/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_xl_adapter.py) | *Text-to-Image Generation with T2I-Adapter Conditioning on StableDiffusion-XL* | - +| [StableDiffusionAdapterPipeline](https://github.com/huggingface/diffusers/blob/main/src/diffusers/pipelines/t2i_adapter/pipeline_stable_diffusion_adapter.py) | *Text-to-Image Generation with T2I-Adapter Conditioning* | - +| [StableDiffusionXLAdapterPipeline](https://github.com/huggingface/diffusers/blob/main/src/diffusers/pipelines/t2i_adapter/pipeline_stable_diffusion_xl_adapter.py) | *Text-to-Image Generation with T2I-Adapter Conditioning on StableDiffusion-XL* | - ## Usage example with the base model of StableDiffusion-1.4/1.5 From a844065384abf4f788cbcb65dfac80a2612ecdbe Mon Sep 17 00:00:00 2001 From: Jake Vanderplas Date: Mon, 9 Oct 2023 08:31:50 -0700 Subject: [PATCH 21/25] replace references to deprecated KeyArray & PRNGKeyArray (#5324) --- setup.py | 4 ++-- src/diffusers/dependency_versions_table.py | 4 ++-- src/diffusers/models/controlnet_flax.py | 2 +- src/diffusers/models/modeling_flax_utils.py | 2 +- src/diffusers/models/unet_2d_condition_flax.py | 2 +- src/diffusers/models/vae_flax.py | 2 +- .../pipelines/controlnet/pipeline_flax_controlnet.py | 6 +++--- .../stable_diffusion/pipeline_flax_stable_diffusion.py | 4 ++-- .../pipeline_flax_stable_diffusion_img2img.py | 6 +++--- .../pipeline_flax_stable_diffusion_inpaint.py | 4 ++-- .../pipelines/stable_diffusion/safety_checker_flax.py | 2 +- .../pipeline_flax_stable_diffusion_xl.py | 4 ++-- src/diffusers/schedulers/scheduling_ddpm_flax.py | 4 ++-- src/diffusers/schedulers/scheduling_karras_ve_flax.py | 3 ++- src/diffusers/schedulers/scheduling_sde_ve_flax.py | 5 +++-- 15 files changed, 28 insertions(+), 26 deletions(-) diff --git a/setup.py b/setup.py index 16a77c1fe313..7ad5646d4fca 100644 --- a/setup.py +++ b/setup.py @@ -102,8 +102,8 @@ "importlib_metadata", "invisible-watermark>=0.2.0", "isort>=5.5.4", - "jax>=0.2.8,!=0.3.2", - "jaxlib>=0.1.65", + "jax>=0.4.1", + "jaxlib>=0.4.1", "Jinja2", "k-diffusion>=0.0.12", "torchsde", diff --git a/src/diffusers/dependency_versions_table.py b/src/diffusers/dependency_versions_table.py index d4b94ba6d4ed..970013c31a20 100644 --- a/src/diffusers/dependency_versions_table.py +++ b/src/diffusers/dependency_versions_table.py @@ -15,8 +15,8 @@ "importlib_metadata": "importlib_metadata", "invisible-watermark": "invisible-watermark>=0.2.0", "isort": "isort>=5.5.4", - "jax": "jax>=0.2.8,!=0.3.2", - "jaxlib": "jaxlib>=0.1.65", + "jax": "jax>=0.4.1", + "jaxlib": "jaxlib>=0.4.1", "Jinja2": "Jinja2", "k-diffusion": "k-diffusion>=0.0.12", "torchsde": "torchsde", diff --git a/src/diffusers/models/controlnet_flax.py b/src/diffusers/models/controlnet_flax.py index a826df48e41a..076e6183211b 100644 --- a/src/diffusers/models/controlnet_flax.py +++ b/src/diffusers/models/controlnet_flax.py @@ -168,7 +168,7 @@ class FlaxControlNetModel(nn.Module, FlaxModelMixin, ConfigMixin): controlnet_conditioning_channel_order: str = "rgb" conditioning_embedding_out_channels: Tuple[int] = (16, 32, 96, 256) - def init_weights(self, rng: jax.random.KeyArray) -> FrozenDict: + def init_weights(self, rng: jax.Array) -> FrozenDict: # init input tensors sample_shape = (1, self.in_channels, self.sample_size, self.sample_size) sample = jnp.zeros(sample_shape, dtype=jnp.float32) diff --git a/src/diffusers/models/modeling_flax_utils.py b/src/diffusers/models/modeling_flax_utils.py index 97f7b43bc64e..ea4d1bfea548 100644 --- a/src/diffusers/models/modeling_flax_utils.py +++ b/src/diffusers/models/modeling_flax_utils.py @@ -192,7 +192,7 @@ def to_fp16(self, params: Union[Dict, FrozenDict], mask: Any = None): ```""" return self._cast_floating_to(params, jnp.float16, mask) - def init_weights(self, rng: jax.random.KeyArray) -> Dict: + def init_weights(self, rng: jax.Array) -> Dict: raise NotImplementedError(f"init_weights method has to be implemented for {self}") @classmethod diff --git a/src/diffusers/models/unet_2d_condition_flax.py b/src/diffusers/models/unet_2d_condition_flax.py index a56db67b6a4e..770cbf09ccac 100644 --- a/src/diffusers/models/unet_2d_condition_flax.py +++ b/src/diffusers/models/unet_2d_condition_flax.py @@ -126,7 +126,7 @@ class FlaxUNet2DConditionModel(nn.Module, FlaxModelMixin, ConfigMixin): addition_embed_type_num_heads: int = 64 projection_class_embeddings_input_dim: Optional[int] = None - def init_weights(self, rng: jax.random.KeyArray) -> FrozenDict: + def init_weights(self, rng: jax.Array) -> FrozenDict: # init input tensors sample_shape = (1, self.in_channels, self.sample_size, self.sample_size) sample = jnp.zeros(sample_shape, dtype=jnp.float32) diff --git a/src/diffusers/models/vae_flax.py b/src/diffusers/models/vae_flax.py index b8f5b1d0e399..d2dde2ba197b 100644 --- a/src/diffusers/models/vae_flax.py +++ b/src/diffusers/models/vae_flax.py @@ -817,7 +817,7 @@ def setup(self): dtype=self.dtype, ) - def init_weights(self, rng: jax.random.KeyArray) -> FrozenDict: + def init_weights(self, rng: jax.Array) -> FrozenDict: # init input tensors sample_shape = (1, self.in_channels, self.sample_size, self.sample_size) sample = jnp.zeros(sample_shape, dtype=jnp.float32) diff --git a/src/diffusers/pipelines/controlnet/pipeline_flax_controlnet.py b/src/diffusers/pipelines/controlnet/pipeline_flax_controlnet.py index b2c8871aa0d6..b57e776e49eb 100644 --- a/src/diffusers/pipelines/controlnet/pipeline_flax_controlnet.py +++ b/src/diffusers/pipelines/controlnet/pipeline_flax_controlnet.py @@ -241,7 +241,7 @@ def _generate( prompt_ids: jnp.array, image: jnp.array, params: Union[Dict, FrozenDict], - prng_seed: jax.random.KeyArray, + prng_seed: jax.Array, num_inference_steps: int, guidance_scale: float, latents: Optional[jnp.array] = None, @@ -351,7 +351,7 @@ def __call__( prompt_ids: jnp.array, image: jnp.array, params: Union[Dict, FrozenDict], - prng_seed: jax.random.KeyArray, + prng_seed: jax.Array, num_inference_steps: int = 50, guidance_scale: Union[float, jnp.array] = 7.5, latents: jnp.array = None, @@ -370,7 +370,7 @@ def __call__( Array representing the ControlNet input condition to provide guidance to the `unet` for generation. params (`Dict` or `FrozenDict`): Dictionary containing the model parameters/weights. - prng_seed (`jax.random.KeyArray` or `jax.Array`): + prng_seed (`jax.Array` or `jax.Array`): Array containing random number generator key. num_inference_steps (`int`, *optional*, defaults to 50): The number of denoising steps. More denoising steps usually lead to a higher quality image at the diff --git a/src/diffusers/pipelines/stable_diffusion/pipeline_flax_stable_diffusion.py b/src/diffusers/pipelines/stable_diffusion/pipeline_flax_stable_diffusion.py index 131a7c7bc2bd..a847cd15c6ce 100644 --- a/src/diffusers/pipelines/stable_diffusion/pipeline_flax_stable_diffusion.py +++ b/src/diffusers/pipelines/stable_diffusion/pipeline_flax_stable_diffusion.py @@ -215,7 +215,7 @@ def _generate( self, prompt_ids: jnp.array, params: Union[Dict, FrozenDict], - prng_seed: jax.random.KeyArray, + prng_seed: jax.Array, num_inference_steps: int, height: int, width: int, @@ -312,7 +312,7 @@ def __call__( self, prompt_ids: jnp.array, params: Union[Dict, FrozenDict], - prng_seed: jax.random.KeyArray, + prng_seed: jax.Array, num_inference_steps: int = 50, height: Optional[int] = None, width: Optional[int] = None, diff --git a/src/diffusers/pipelines/stable_diffusion/pipeline_flax_stable_diffusion_img2img.py b/src/diffusers/pipelines/stable_diffusion/pipeline_flax_stable_diffusion_img2img.py index a9717533fa93..42a79db6b2b2 100644 --- a/src/diffusers/pipelines/stable_diffusion/pipeline_flax_stable_diffusion_img2img.py +++ b/src/diffusers/pipelines/stable_diffusion/pipeline_flax_stable_diffusion_img2img.py @@ -235,7 +235,7 @@ def _generate( prompt_ids: jnp.array, image: jnp.array, params: Union[Dict, FrozenDict], - prng_seed: jax.random.KeyArray, + prng_seed: jax.Array, start_timestep: int, num_inference_steps: int, height: int, @@ -340,7 +340,7 @@ def __call__( prompt_ids: jnp.array, image: jnp.array, params: Union[Dict, FrozenDict], - prng_seed: jax.random.KeyArray, + prng_seed: jax.Array, strength: float = 0.8, num_inference_steps: int = 50, height: Optional[int] = None, @@ -361,7 +361,7 @@ def __call__( Array representing an image batch to be used as the starting point. params (`Dict` or `FrozenDict`): Dictionary containing the model parameters/weights. - prng_seed (`jax.random.KeyArray` or `jax.Array`): + prng_seed (`jax.Array` or `jax.Array`): Array containing random number generator key. strength (`float`, *optional*, defaults to 0.8): Indicates extent to transform the reference `image`. Must be between 0 and 1. `image` is used as a diff --git a/src/diffusers/pipelines/stable_diffusion/pipeline_flax_stable_diffusion_inpaint.py b/src/diffusers/pipelines/stable_diffusion/pipeline_flax_stable_diffusion_inpaint.py index b43fa3837062..153267da1067 100644 --- a/src/diffusers/pipelines/stable_diffusion/pipeline_flax_stable_diffusion_inpaint.py +++ b/src/diffusers/pipelines/stable_diffusion/pipeline_flax_stable_diffusion_inpaint.py @@ -270,7 +270,7 @@ def _generate( mask: jnp.array, masked_image: jnp.array, params: Union[Dict, FrozenDict], - prng_seed: jax.random.KeyArray, + prng_seed: jax.Array, num_inference_steps: int, height: int, width: int, @@ -398,7 +398,7 @@ def __call__( mask: jnp.array, masked_image: jnp.array, params: Union[Dict, FrozenDict], - prng_seed: jax.random.KeyArray, + prng_seed: jax.Array, num_inference_steps: int = 50, height: Optional[int] = None, width: Optional[int] = None, diff --git a/src/diffusers/pipelines/stable_diffusion/safety_checker_flax.py b/src/diffusers/pipelines/stable_diffusion/safety_checker_flax.py index 3a8c31679540..5966600462bf 100644 --- a/src/diffusers/pipelines/stable_diffusion/safety_checker_flax.py +++ b/src/diffusers/pipelines/stable_diffusion/safety_checker_flax.py @@ -87,7 +87,7 @@ def __init__( module = self.module_class(config=config, dtype=dtype, **kwargs) super().__init__(config, module, input_shape=input_shape, seed=seed, dtype=dtype, _do_init=_do_init) - def init_weights(self, rng: jax.random.KeyArray, input_shape: Tuple, params: FrozenDict = None) -> FrozenDict: + def init_weights(self, rng: jax.Array, input_shape: Tuple, params: FrozenDict = None) -> FrozenDict: # init input tensor clip_input = jax.random.normal(rng, input_shape) diff --git a/src/diffusers/pipelines/stable_diffusion_xl/pipeline_flax_stable_diffusion_xl.py b/src/diffusers/pipelines/stable_diffusion_xl/pipeline_flax_stable_diffusion_xl.py index 3acb5ae538a4..8f043c7c6657 100644 --- a/src/diffusers/pipelines/stable_diffusion_xl/pipeline_flax_stable_diffusion_xl.py +++ b/src/diffusers/pipelines/stable_diffusion_xl/pipeline_flax_stable_diffusion_xl.py @@ -89,7 +89,7 @@ def __call__( self, prompt_ids: jax.Array, params: Union[Dict, FrozenDict], - prng_seed: jax.random.KeyArray, + prng_seed: jax.Array, num_inference_steps: int = 50, guidance_scale: Union[float, jax.Array] = 7.5, height: Optional[int] = None, @@ -170,7 +170,7 @@ def _generate( self, prompt_ids: jnp.array, params: Union[Dict, FrozenDict], - prng_seed: jax.random.KeyArray, + prng_seed: jax.Array, num_inference_steps: int, height: int, width: int, diff --git a/src/diffusers/schedulers/scheduling_ddpm_flax.py b/src/diffusers/schedulers/scheduling_ddpm_flax.py index 529d2bd03a75..ab7d70f466e6 100644 --- a/src/diffusers/schedulers/scheduling_ddpm_flax.py +++ b/src/diffusers/schedulers/scheduling_ddpm_flax.py @@ -198,7 +198,7 @@ def step( model_output: jnp.ndarray, timestep: int, sample: jnp.ndarray, - key: Optional[jax.random.KeyArray] = None, + key: Optional[jax.Array] = None, return_dict: bool = True, ) -> Union[FlaxDDPMSchedulerOutput, Tuple]: """ @@ -211,7 +211,7 @@ def step( timestep (`int`): current discrete timestep in the diffusion chain. sample (`jnp.ndarray`): current instance of sample being created by diffusion process. - key (`jax.random.KeyArray`): a PRNG key. + key (`jax.Array`): a PRNG key. return_dict (`bool`): option for returning tuple rather than FlaxDDPMSchedulerOutput class Returns: diff --git a/src/diffusers/schedulers/scheduling_karras_ve_flax.py b/src/diffusers/schedulers/scheduling_karras_ve_flax.py index 45c0dbddf7ef..4a8606007d5f 100644 --- a/src/diffusers/schedulers/scheduling_karras_ve_flax.py +++ b/src/diffusers/schedulers/scheduling_karras_ve_flax.py @@ -17,6 +17,7 @@ from typing import Optional, Tuple, Union import flax +import jax import jax.numpy as jnp from jax import random @@ -139,7 +140,7 @@ def add_noise_to_input( state: KarrasVeSchedulerState, sample: jnp.ndarray, sigma: float, - key: random.KeyArray, + key: jax.Array, ) -> Tuple[jnp.ndarray, float]: """ Explicit Langevin-like "churn" step of adding noise to the sample according to a factor gamma_i ≥ 0 to reach a diff --git a/src/diffusers/schedulers/scheduling_sde_ve_flax.py b/src/diffusers/schedulers/scheduling_sde_ve_flax.py index b6240559fc88..935f972a9bdb 100644 --- a/src/diffusers/schedulers/scheduling_sde_ve_flax.py +++ b/src/diffusers/schedulers/scheduling_sde_ve_flax.py @@ -18,6 +18,7 @@ from typing import Optional, Tuple, Union import flax +import jax import jax.numpy as jnp from jax import random @@ -169,7 +170,7 @@ def step_pred( model_output: jnp.ndarray, timestep: int, sample: jnp.ndarray, - key: random.KeyArray, + key: jax.Array, return_dict: bool = True, ) -> Union[FlaxSdeVeOutput, Tuple]: """ @@ -228,7 +229,7 @@ def step_correct( state: ScoreSdeVeSchedulerState, model_output: jnp.ndarray, sample: jnp.ndarray, - key: random.KeyArray, + key: jax.Array, return_dict: bool = True, ) -> Union[FlaxSdeVeOutput, Tuple]: """ From ed2f956072a3b446d984f359ba6c427c259ab4ee Mon Sep 17 00:00:00 2001 From: Patrick von Platen Date: Mon, 9 Oct 2023 18:01:55 +0200 Subject: [PATCH 22/25] Fix loading broken LoRAs that could give NaN (#5316) * Fix fuse Lora * improve a bit * make style * Update src/diffusers/models/lora.py Co-authored-by: Benjamin Bossan * ciao C file * ciao C file * test & make style --------- Co-authored-by: Benjamin Bossan --- src/diffusers/loaders.py | 48 +++++++++++++++------- src/diffusers/models/lora.py | 20 ++++++++- tests/lora/test_lora_layers_old_backend.py | 41 ++++++++++++++++++ 3 files changed, 92 insertions(+), 17 deletions(-) diff --git a/src/diffusers/loaders.py b/src/diffusers/loaders.py index 11858ac3bb90..2cc547be0178 100644 --- a/src/diffusers/loaders.py +++ b/src/diffusers/loaders.py @@ -121,7 +121,7 @@ def state_dict(self, *args, destination=None, prefix="", keep_vars=False): return super().state_dict(*args, destination=destination, prefix=prefix, keep_vars=keep_vars) - def _fuse_lora(self, lora_scale=1.0): + def _fuse_lora(self, lora_scale=1.0, safe_fusing=False): if self.lora_linear_layer is None: return @@ -135,6 +135,14 @@ def _fuse_lora(self, lora_scale=1.0): w_up = w_up * self.lora_linear_layer.network_alpha / self.lora_linear_layer.rank fused_weight = w_orig + (lora_scale * torch.bmm(w_up[None, :], w_down[None, :])[0]) + + if safe_fusing and torch.isnan(fused_weight).any().item(): + raise ValueError( + "This LoRA weight seems to be broken. " + f"Encountered NaN values when trying to fuse LoRA weights for {self}." + "LoRA weights will not be fused." + ) + self.regular_linear_layer.weight.data = fused_weight.to(device=device, dtype=dtype) # we can drop the lora layer now @@ -672,13 +680,14 @@ def save_function(weights, filename): save_function(state_dict, os.path.join(save_directory, weight_name)) logger.info(f"Model weights saved in {os.path.join(save_directory, weight_name)}") - def fuse_lora(self, lora_scale=1.0): + def fuse_lora(self, lora_scale=1.0, safe_fusing=False): self.lora_scale = lora_scale + self._safe_fusing = safe_fusing self.apply(self._fuse_lora_apply) def _fuse_lora_apply(self, module): if hasattr(module, "_fuse_lora"): - module._fuse_lora(self.lora_scale) + module._fuse_lora(self.lora_scale, self._safe_fusing) def unfuse_lora(self): self.apply(self._unfuse_lora_apply) @@ -2086,7 +2095,13 @@ def unload_lora_weights(self): # Safe to call the following regardless of LoRA. self._remove_text_encoder_monkey_patch() - def fuse_lora(self, fuse_unet: bool = True, fuse_text_encoder: bool = True, lora_scale: float = 1.0): + def fuse_lora( + self, + fuse_unet: bool = True, + fuse_text_encoder: bool = True, + lora_scale: float = 1.0, + safe_fusing: bool = False, + ): r""" Fuses the LoRA parameters into the original parameters of the corresponding blocks. @@ -2103,6 +2118,8 @@ def fuse_lora(self, fuse_unet: bool = True, fuse_text_encoder: bool = True, lora LoRA parameters then it won't have any effect. lora_scale (`float`, defaults to 1.0): Controls how much to influence the outputs with the LoRA parameters. + safe_fusing (`bool`, defaults to `False`): + Whether to check fused weights for NaN values before fusing and if values are NaN not fusing them. """ if fuse_unet or fuse_text_encoder: self.num_fused_loras += 1 @@ -2112,12 +2129,13 @@ def fuse_lora(self, fuse_unet: bool = True, fuse_text_encoder: bool = True, lora ) if fuse_unet: - self.unet.fuse_lora(lora_scale) + self.unet.fuse_lora(lora_scale, safe_fusing=safe_fusing) if self.use_peft_backend: from peft.tuners.tuners_utils import BaseTunerLayer - def fuse_text_encoder_lora(text_encoder, lora_scale=1.0): + def fuse_text_encoder_lora(text_encoder, lora_scale=1.0, safe_fusing=False): + # TODO(Patrick, Younes): enable "safe" fusing for module in text_encoder.modules(): if isinstance(module, BaseTunerLayer): if lora_scale != 1.0: @@ -2129,24 +2147,24 @@ def fuse_text_encoder_lora(text_encoder, lora_scale=1.0): if version.parse(__version__) > version.parse("0.23"): deprecate("fuse_text_encoder_lora", "0.25", LORA_DEPRECATION_MESSAGE) - def fuse_text_encoder_lora(text_encoder, lora_scale=1.0): + def fuse_text_encoder_lora(text_encoder, lora_scale=1.0, safe_fusing=False): for _, attn_module in text_encoder_attn_modules(text_encoder): if isinstance(attn_module.q_proj, PatchedLoraProjection): - attn_module.q_proj._fuse_lora(lora_scale) - attn_module.k_proj._fuse_lora(lora_scale) - attn_module.v_proj._fuse_lora(lora_scale) - attn_module.out_proj._fuse_lora(lora_scale) + attn_module.q_proj._fuse_lora(lora_scale, safe_fusing) + attn_module.k_proj._fuse_lora(lora_scale, safe_fusing) + attn_module.v_proj._fuse_lora(lora_scale, safe_fusing) + attn_module.out_proj._fuse_lora(lora_scale, safe_fusing) for _, mlp_module in text_encoder_mlp_modules(text_encoder): if isinstance(mlp_module.fc1, PatchedLoraProjection): - mlp_module.fc1._fuse_lora(lora_scale) - mlp_module.fc2._fuse_lora(lora_scale) + mlp_module.fc1._fuse_lora(lora_scale, safe_fusing) + mlp_module.fc2._fuse_lora(lora_scale, safe_fusing) if fuse_text_encoder: if hasattr(self, "text_encoder"): - fuse_text_encoder_lora(self.text_encoder, lora_scale) + fuse_text_encoder_lora(self.text_encoder, lora_scale, safe_fusing) if hasattr(self, "text_encoder_2"): - fuse_text_encoder_lora(self.text_encoder_2, lora_scale) + fuse_text_encoder_lora(self.text_encoder_2, lora_scale, safe_fusing) def unfuse_lora(self, unfuse_unet: bool = True, unfuse_text_encoder: bool = True): r""" diff --git a/src/diffusers/models/lora.py b/src/diffusers/models/lora.py index a777bb93e1c8..aec7200afdfe 100644 --- a/src/diffusers/models/lora.py +++ b/src/diffusers/models/lora.py @@ -112,7 +112,7 @@ def __init__(self, *args, lora_layer: Optional[LoRAConv2dLayer] = None, **kwargs def set_lora_layer(self, lora_layer: Optional[LoRAConv2dLayer]): self.lora_layer = lora_layer - def _fuse_lora(self, lora_scale=1.0): + def _fuse_lora(self, lora_scale=1.0, safe_fusing=False): if self.lora_layer is None: return @@ -128,6 +128,14 @@ def _fuse_lora(self, lora_scale=1.0): fusion = torch.mm(w_up.flatten(start_dim=1), w_down.flatten(start_dim=1)) fusion = fusion.reshape((w_orig.shape)) fused_weight = w_orig + (lora_scale * fusion) + + if safe_fusing and torch.isnan(fused_weight).any().item(): + raise ValueError( + "This LoRA weight seems to be broken. " + f"Encountered NaN values when trying to fuse LoRA weights for {self}." + "LoRA weights will not be fused." + ) + self.weight.data = fused_weight.to(device=device, dtype=dtype) # we can drop the lora layer now @@ -182,7 +190,7 @@ def __init__(self, *args, lora_layer: Optional[LoRALinearLayer] = None, **kwargs def set_lora_layer(self, lora_layer: Optional[LoRALinearLayer]): self.lora_layer = lora_layer - def _fuse_lora(self, lora_scale=1.0): + def _fuse_lora(self, lora_scale=1.0, safe_fusing=False): if self.lora_layer is None: return @@ -196,6 +204,14 @@ def _fuse_lora(self, lora_scale=1.0): w_up = w_up * self.lora_layer.network_alpha / self.lora_layer.rank fused_weight = w_orig + (lora_scale * torch.bmm(w_up[None, :], w_down[None, :])[0]) + + if safe_fusing and torch.isnan(fused_weight).any().item(): + raise ValueError( + "This LoRA weight seems to be broken. " + f"Encountered NaN values when trying to fuse LoRA weights for {self}." + "LoRA weights will not be fused." + ) + self.weight.data = fused_weight.to(device=device, dtype=dtype) # we can drop the lora layer now diff --git a/tests/lora/test_lora_layers_old_backend.py b/tests/lora/test_lora_layers_old_backend.py index d616ef8c78b8..8c1fb4877653 100644 --- a/tests/lora/test_lora_layers_old_backend.py +++ b/tests/lora/test_lora_layers_old_backend.py @@ -1028,6 +1028,47 @@ def test_load_lora_locally_safetensors(self): sd_pipe.unload_lora_weights() + def test_lora_fuse_nan(self): + pipeline_components, lora_components = self.get_dummy_components() + sd_pipe = StableDiffusionXLPipeline(**pipeline_components) + sd_pipe = sd_pipe.to(torch_device) + sd_pipe.set_progress_bar_config(disable=None) + + _, _, pipeline_inputs = self.get_dummy_inputs(with_generator=False) + + # Emulate training. + set_lora_weights(lora_components["unet_lora_layers"].parameters(), randn_weight=True) + set_lora_weights(lora_components["text_encoder_one_lora_layers"].parameters(), randn_weight=True) + set_lora_weights(lora_components["text_encoder_two_lora_layers"].parameters(), randn_weight=True) + + with tempfile.TemporaryDirectory() as tmpdirname: + StableDiffusionXLPipeline.save_lora_weights( + save_directory=tmpdirname, + unet_lora_layers=lora_components["unet_lora_layers"], + text_encoder_lora_layers=lora_components["text_encoder_one_lora_layers"], + text_encoder_2_lora_layers=lora_components["text_encoder_two_lora_layers"], + safe_serialization=True, + ) + self.assertTrue(os.path.isfile(os.path.join(tmpdirname, "pytorch_lora_weights.safetensors"))) + sd_pipe.load_lora_weights(os.path.join(tmpdirname, "pytorch_lora_weights.safetensors")) + + # corrupt one LoRA weight with `inf` values + with torch.no_grad(): + sd_pipe.unet.mid_block.attentions[0].transformer_blocks[0].attn1.to_q.lora_layer.down.weight += float( + "inf" + ) + + # with `safe_fusing=True` we should see an Error + with self.assertRaises(ValueError): + sd_pipe.fuse_lora(safe_fusing=True) + + # without we should not see an error, but every image will be black + sd_pipe.fuse_lora(safe_fusing=False) + + out = sd_pipe("test", num_inference_steps=2, output_type="np").images + + assert np.isnan(out).all() + def test_lora_fusion(self): pipeline_components, lora_components = self.get_dummy_components() sd_pipe = StableDiffusionXLPipeline(**pipeline_components) From 4ac205e32fc44b6d92488304d74781a83945b9e6 Mon Sep 17 00:00:00 2001 From: Roy Hvaara Date: Mon, 9 Oct 2023 13:04:40 -0700 Subject: [PATCH 23/25] [JAX] Replace uses of `jnp.array` in types with `jnp.ndarray`. (#4719) `jnp.array` is a function, not a type: https://jax.readthedocs.io/en/latest/_autosummary/jax.numpy.array.html so it never makes sense to use `jnp.array` in a type annotation. Presumably the intent was to write `jnp.ndarray` aka `jax.Array`. Change uses of `jnp.array` to `jnp.ndarray`. Co-authored-by: Peter Hawkins --- .../controlnet/pipeline_flax_controlnet.py | 30 +++++++++---------- .../pipeline_flax_stable_diffusion.py | 12 ++++---- .../pipeline_flax_stable_diffusion_img2img.py | 24 +++++++-------- .../pipeline_flax_stable_diffusion_inpaint.py | 24 +++++++-------- 4 files changed, 45 insertions(+), 45 deletions(-) diff --git a/src/diffusers/pipelines/controlnet/pipeline_flax_controlnet.py b/src/diffusers/pipelines/controlnet/pipeline_flax_controlnet.py index b57e776e49eb..e1f508dc1e36 100644 --- a/src/diffusers/pipelines/controlnet/pipeline_flax_controlnet.py +++ b/src/diffusers/pipelines/controlnet/pipeline_flax_controlnet.py @@ -238,14 +238,14 @@ def _run_safety_checker(self, images, safety_model_params, jit=False): def _generate( self, - prompt_ids: jnp.array, - image: jnp.array, + prompt_ids: jnp.ndarray, + image: jnp.ndarray, params: Union[Dict, FrozenDict], prng_seed: jax.Array, num_inference_steps: int, guidance_scale: float, - latents: Optional[jnp.array] = None, - neg_prompt_ids: Optional[jnp.array] = None, + latents: Optional[jnp.ndarray] = None, + neg_prompt_ids: Optional[jnp.ndarray] = None, controlnet_conditioning_scale: float = 1.0, ): height, width = image.shape[-2:] @@ -348,15 +348,15 @@ def loop_body(step, args): @replace_example_docstring(EXAMPLE_DOC_STRING) def __call__( self, - prompt_ids: jnp.array, - image: jnp.array, + prompt_ids: jnp.ndarray, + image: jnp.ndarray, params: Union[Dict, FrozenDict], prng_seed: jax.Array, num_inference_steps: int = 50, - guidance_scale: Union[float, jnp.array] = 7.5, - latents: jnp.array = None, - neg_prompt_ids: jnp.array = None, - controlnet_conditioning_scale: Union[float, jnp.array] = 1.0, + guidance_scale: Union[float, jnp.ndarray] = 7.5, + latents: jnp.ndarray = None, + neg_prompt_ids: jnp.ndarray = None, + controlnet_conditioning_scale: Union[float, jnp.ndarray] = 1.0, return_dict: bool = True, jit: bool = False, ): @@ -364,13 +364,13 @@ def __call__( The call function to the pipeline for generation. Args: - prompt_ids (`jnp.array`): + prompt_ids (`jnp.ndarray`): The prompt or prompts to guide the image generation. - image (`jnp.array`): + image (`jnp.ndarray`): Array representing the ControlNet input condition to provide guidance to the `unet` for generation. params (`Dict` or `FrozenDict`): Dictionary containing the model parameters/weights. - prng_seed (`jax.Array` or `jax.Array`): + prng_seed (`jax.Array`): Array containing random number generator key. num_inference_steps (`int`, *optional*, defaults to 50): The number of denoising steps. More denoising steps usually lead to a higher quality image at the @@ -378,11 +378,11 @@ def __call__( guidance_scale (`float`, *optional*, defaults to 7.5): A higher guidance scale value encourages the model to generate images closely linked to the text `prompt` at the expense of lower image quality. Guidance scale is enabled when `guidance_scale > 1`. - latents (`jnp.array`, *optional*): + latents (`jnp.ndarray`, *optional*): Pre-generated noisy latents sampled from a Gaussian distribution, to be used as inputs for image generation. Can be used to tweak the same generation with different prompts. If not provided, a latents array is generated by sampling using the supplied random `generator`. - controlnet_conditioning_scale (`float` or `jnp.array`, *optional*, defaults to 1.0): + controlnet_conditioning_scale (`float` or `jnp.ndarray`, *optional*, defaults to 1.0): The outputs of the ControlNet are multiplied by `controlnet_conditioning_scale` before they are added to the residual in the original `unet`. return_dict (`bool`, *optional*, defaults to `True`): diff --git a/src/diffusers/pipelines/stable_diffusion/pipeline_flax_stable_diffusion.py b/src/diffusers/pipelines/stable_diffusion/pipeline_flax_stable_diffusion.py index a847cd15c6ce..bcf2a6217772 100644 --- a/src/diffusers/pipelines/stable_diffusion/pipeline_flax_stable_diffusion.py +++ b/src/diffusers/pipelines/stable_diffusion/pipeline_flax_stable_diffusion.py @@ -220,8 +220,8 @@ def _generate( height: int, width: int, guidance_scale: float, - latents: Optional[jnp.array] = None, - neg_prompt_ids: Optional[jnp.array] = None, + latents: Optional[jnp.ndarray] = None, + neg_prompt_ids: Optional[jnp.ndarray] = None, ): if height % 8 != 0 or width % 8 != 0: raise ValueError(f"`height` and `width` have to be divisible by 8 but are {height} and {width}.") @@ -316,9 +316,9 @@ def __call__( num_inference_steps: int = 50, height: Optional[int] = None, width: Optional[int] = None, - guidance_scale: Union[float, jnp.array] = 7.5, - latents: jnp.array = None, - neg_prompt_ids: jnp.array = None, + guidance_scale: Union[float, jnp.ndarray] = 7.5, + latents: jnp.ndarray = None, + neg_prompt_ids: jnp.ndarray = None, return_dict: bool = True, jit: bool = False, ): @@ -338,7 +338,7 @@ def __call__( guidance_scale (`float`, *optional*, defaults to 7.5): A higher guidance scale value encourages the model to generate images closely linked to the text `prompt` at the expense of lower image quality. Guidance scale is enabled when `guidance_scale > 1`. - latents (`jnp.array`, *optional*): + latents (`jnp.ndarray`, *optional*): Pre-generated noisy latents sampled from a Gaussian distribution, to be used as inputs for image generation. Can be used to tweak the same generation with different prompts. If not provided, a latents array is generated by sampling using the supplied random `generator`. diff --git a/src/diffusers/pipelines/stable_diffusion/pipeline_flax_stable_diffusion_img2img.py b/src/diffusers/pipelines/stable_diffusion/pipeline_flax_stable_diffusion_img2img.py index 42a79db6b2b2..c1fd310ea582 100644 --- a/src/diffusers/pipelines/stable_diffusion/pipeline_flax_stable_diffusion_img2img.py +++ b/src/diffusers/pipelines/stable_diffusion/pipeline_flax_stable_diffusion_img2img.py @@ -232,8 +232,8 @@ def get_timestep_start(self, num_inference_steps, strength): def _generate( self, - prompt_ids: jnp.array, - image: jnp.array, + prompt_ids: jnp.ndarray, + image: jnp.ndarray, params: Union[Dict, FrozenDict], prng_seed: jax.Array, start_timestep: int, @@ -241,8 +241,8 @@ def _generate( height: int, width: int, guidance_scale: float, - noise: Optional[jnp.array] = None, - neg_prompt_ids: Optional[jnp.array] = None, + noise: Optional[jnp.ndarray] = None, + neg_prompt_ids: Optional[jnp.ndarray] = None, ): if height % 8 != 0 or width % 8 != 0: raise ValueError(f"`height` and `width` have to be divisible by 8 but are {height} and {width}.") @@ -337,17 +337,17 @@ def loop_body(step, args): @replace_example_docstring(EXAMPLE_DOC_STRING) def __call__( self, - prompt_ids: jnp.array, - image: jnp.array, + prompt_ids: jnp.ndarray, + image: jnp.ndarray, params: Union[Dict, FrozenDict], prng_seed: jax.Array, strength: float = 0.8, num_inference_steps: int = 50, height: Optional[int] = None, width: Optional[int] = None, - guidance_scale: Union[float, jnp.array] = 7.5, - noise: jnp.array = None, - neg_prompt_ids: jnp.array = None, + guidance_scale: Union[float, jnp.ndarray] = 7.5, + noise: jnp.ndarray = None, + neg_prompt_ids: jnp.ndarray = None, return_dict: bool = True, jit: bool = False, ): @@ -355,9 +355,9 @@ def __call__( The call function to the pipeline for generation. Args: - prompt_ids (`jnp.array`): + prompt_ids (`jnp.ndarray`): The prompt or prompts to guide image generation. - image (`jnp.array`): + image (`jnp.ndarray`): Array representing an image batch to be used as the starting point. params (`Dict` or `FrozenDict`): Dictionary containing the model parameters/weights. @@ -379,7 +379,7 @@ def __call__( guidance_scale (`float`, *optional*, defaults to 7.5): A higher guidance scale value encourages the model to generate images closely linked to the text `prompt` at the expense of lower image quality. Guidance scale is enabled when `guidance_scale > 1`. - noise (`jnp.array`, *optional*): + noise (`jnp.ndarray`, *optional*): Pre-generated noisy latents sampled from a Gaussian distribution to be used as inputs for image generation. Can be used to tweak the same generation with different prompts. The array is generated by sampling using the supplied random `generator`. diff --git a/src/diffusers/pipelines/stable_diffusion/pipeline_flax_stable_diffusion_inpaint.py b/src/diffusers/pipelines/stable_diffusion/pipeline_flax_stable_diffusion_inpaint.py index 153267da1067..b9a2331a061c 100644 --- a/src/diffusers/pipelines/stable_diffusion/pipeline_flax_stable_diffusion_inpaint.py +++ b/src/diffusers/pipelines/stable_diffusion/pipeline_flax_stable_diffusion_inpaint.py @@ -266,17 +266,17 @@ def _run_safety_checker(self, images, safety_model_params, jit=False): def _generate( self, - prompt_ids: jnp.array, - mask: jnp.array, - masked_image: jnp.array, + prompt_ids: jnp.ndarray, + mask: jnp.ndarray, + masked_image: jnp.ndarray, params: Union[Dict, FrozenDict], prng_seed: jax.Array, num_inference_steps: int, height: int, width: int, guidance_scale: float, - latents: Optional[jnp.array] = None, - neg_prompt_ids: Optional[jnp.array] = None, + latents: Optional[jnp.ndarray] = None, + neg_prompt_ids: Optional[jnp.ndarray] = None, ): if height % 8 != 0 or width % 8 != 0: raise ValueError(f"`height` and `width` have to be divisible by 8 but are {height} and {width}.") @@ -394,17 +394,17 @@ def loop_body(step, args): @replace_example_docstring(EXAMPLE_DOC_STRING) def __call__( self, - prompt_ids: jnp.array, - mask: jnp.array, - masked_image: jnp.array, + prompt_ids: jnp.ndarray, + mask: jnp.ndarray, + masked_image: jnp.ndarray, params: Union[Dict, FrozenDict], prng_seed: jax.Array, num_inference_steps: int = 50, height: Optional[int] = None, width: Optional[int] = None, - guidance_scale: Union[float, jnp.array] = 7.5, - latents: jnp.array = None, - neg_prompt_ids: jnp.array = None, + guidance_scale: Union[float, jnp.ndarray] = 7.5, + latents: jnp.ndarray = None, + neg_prompt_ids: jnp.ndarray = None, return_dict: bool = True, jit: bool = False, ): @@ -424,7 +424,7 @@ def __call__( guidance_scale (`float`, *optional*, defaults to 7.5): A higher guidance scale value encourages the model to generate images closely linked to the text `prompt` at the expense of lower image quality. Guidance scale is enabled when `guidance_scale > 1`. - latents (`jnp.array`, *optional*): + latents (`jnp.ndarray`, *optional*): Pre-generated noisy latents sampled from a Gaussian distribution, to be used as inputs for image generation. Can be used to tweak the same generation with different prompts. If not provided, a latents array is generated by sampling using the supplied random `generator`. From d3e0750d5d0089d84ebaffa2fb7adbd6619abc57 Mon Sep 17 00:00:00 2001 From: Julien Simon <3436143+juliensimon@users.noreply.github.com> Date: Tue, 10 Oct 2023 12:07:58 +0200 Subject: [PATCH 24/25] Add missing dependency in requirements file (#5345) Update requirements_sdxl.txt Add missing 'datasets' Co-authored-by: Sayak Paul --- examples/text_to_image/requirements_sdxl.txt | 1 + 1 file changed, 1 insertion(+) diff --git a/examples/text_to_image/requirements_sdxl.txt b/examples/text_to_image/requirements_sdxl.txt index 5d67662fadbe..cdd3336e3617 100644 --- a/examples/text_to_image/requirements_sdxl.txt +++ b/examples/text_to_image/requirements_sdxl.txt @@ -4,3 +4,4 @@ transformers>=4.25.1 ftfy tensorboard Jinja2 +datasets From 9c82b68f074df116f0c5044d20a0cf9c0086b5ac Mon Sep 17 00:00:00 2001 From: Humphrey009 <60021713+Humphrey009@users.noreply.github.com> Date: Tue, 10 Oct 2023 18:09:22 +0800 Subject: [PATCH 25/25] fix problem of 'accelerator.is_main_process' to run in mutiple GPUs (#5340) fix problem of 'accelerator.is_main_process' to run in mutiple GPUs or NPUs Co-authored-by: jiaqiw --- .../train_unconditional.py | 44 +++++++++---------- 1 file changed, 22 insertions(+), 22 deletions(-) diff --git a/examples/unconditional_image_generation/train_unconditional.py b/examples/unconditional_image_generation/train_unconditional.py index 4925c74c8ccf..a3baa3b85b36 100644 --- a/examples/unconditional_image_generation/train_unconditional.py +++ b/examples/unconditional_image_generation/train_unconditional.py @@ -607,28 +607,28 @@ def transform_images(examples): progress_bar.update(1) global_step += 1 - if global_step % args.checkpointing_steps == 0: - # _before_ saving state, check if this save would set us over the `checkpoints_total_limit` - if args.checkpoints_total_limit is not None: - checkpoints = os.listdir(args.output_dir) - checkpoints = [d for d in checkpoints if d.startswith("checkpoint")] - checkpoints = sorted(checkpoints, key=lambda x: int(x.split("-")[1])) - - # before we save the new checkpoint, we need to have at _most_ `checkpoints_total_limit - 1` checkpoints - if len(checkpoints) >= args.checkpoints_total_limit: - num_to_remove = len(checkpoints) - args.checkpoints_total_limit + 1 - removing_checkpoints = checkpoints[0:num_to_remove] - - logger.info( - f"{len(checkpoints)} checkpoints already exist, removing {len(removing_checkpoints)} checkpoints" - ) - logger.info(f"removing checkpoints: {', '.join(removing_checkpoints)}") - - for removing_checkpoint in removing_checkpoints: - removing_checkpoint = os.path.join(args.output_dir, removing_checkpoint) - shutil.rmtree(removing_checkpoint) - - if accelerator.is_main_process: + if accelerator.is_main_process: + if global_step % args.checkpointing_steps == 0: + # _before_ saving state, check if this save would set us over the `checkpoints_total_limit` + if args.checkpoints_total_limit is not None: + checkpoints = os.listdir(args.output_dir) + checkpoints = [d for d in checkpoints if d.startswith("checkpoint")] + checkpoints = sorted(checkpoints, key=lambda x: int(x.split("-")[1])) + + # before we save the new checkpoint, we need to have at _most_ `checkpoints_total_limit - 1` checkpoints + if len(checkpoints) >= args.checkpoints_total_limit: + num_to_remove = len(checkpoints) - args.checkpoints_total_limit + 1 + removing_checkpoints = checkpoints[0:num_to_remove] + + logger.info( + f"{len(checkpoints)} checkpoints already exist, removing {len(removing_checkpoints)} checkpoints" + ) + logger.info(f"removing checkpoints: {', '.join(removing_checkpoints)}") + + for removing_checkpoint in removing_checkpoints: + removing_checkpoint = os.path.join(args.output_dir, removing_checkpoint) + shutil.rmtree(removing_checkpoint) + save_path = os.path.join(args.output_dir, f"checkpoint-{global_step}") accelerator.save_state(save_path) logger.info(f"Saved state to {save_path}")