Spaces:
Running on Zero
Running on Zero
| # Copyright 2024 Stability AI, Katherine Crowson and The HuggingFace Team. All rights reserved. | |
| # | |
| # Licensed under the Apache License, Version 2.0 (the "License"); | |
| # you may not use this file except in compliance with the License. | |
| # You may obtain a copy of the License at | |
| # | |
| # http://www.apache.org/licenses/LICENSE-2.0 | |
| # | |
| # Unless required by applicable law or agreed to in writing, software | |
| # distributed under the License is distributed on an "AS IS" BASIS, | |
| # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. | |
| # See the License for the specific language governing permissions and | |
| # limitations under the License. | |
| import math | |
| import numpy as np | |
| import torch | |
| def compute_density_for_timestep_sampling( | |
| weighting_scheme: str, batch_size: int, logit_mean: float = None, logit_std: float = None, mode_scale: float = None | |
| ): | |
| """ | |
| Compute the density for sampling the timesteps when doing SD3 training. | |
| Courtesy: This was contributed by Rafie Walker in https://github.com/huggingface/diffusers/pull/8528. | |
| SD3 paper reference: https://arxiv.org/abs/2403.03206v1. | |
| """ | |
| if weighting_scheme == "logit_normal": | |
| # See 3.1 in the SD3 paper ($rf/lognorm(0.00,1.00)$). | |
| u = torch.normal(mean=logit_mean, std=logit_std, size=(batch_size,), device="cpu") | |
| u = torch.nn.functional.sigmoid(u) | |
| elif weighting_scheme == "mode": | |
| u = torch.rand(size=(batch_size,), device="cpu") | |
| u = 1 - u - mode_scale * (torch.cos(math.pi * u / 2) ** 2 - 1 + u) | |
| else: | |
| u = torch.rand(size=(batch_size,), device="cpu") | |
| return u | |
| def compute_loss_weighting_for_sd3(weighting_scheme: str, sigmas=None): | |
| """ | |
| Computes loss weighting scheme for SD3 training. | |
| Courtesy: This was contributed by Rafie Walker in https://github.com/huggingface/diffusers/pull/8528. | |
| SD3 paper reference: https://arxiv.org/abs/2403.03206v1. | |
| """ | |
| if weighting_scheme == "sigma_sqrt": | |
| weighting = (sigmas**-2.0).float() | |
| elif weighting_scheme == "cosmap": | |
| bot = 1 - 2 * sigmas + 2 * sigmas**2 | |
| weighting = 2 / (math.pi * bot) | |
| else: | |
| weighting = torch.ones_like(sigmas) | |
| return weighting | |