85 lines
2.5 KiB
Python
85 lines
2.5 KiB
Python
# // Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
|
|
# //
|
|
# // 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.
|
|
|
|
"""
|
|
Utility functions.
|
|
"""
|
|
|
|
from typing import Callable
|
|
import torch
|
|
|
|
|
|
def expand_dims(tensor: torch.Tensor, ndim: int):
|
|
"""
|
|
Expand tensor to target ndim. New dims are added to the right.
|
|
For example, if the tensor shape was (8,), target ndim is 4, return (8, 1, 1, 1).
|
|
"""
|
|
shape = tensor.shape + (1,) * (ndim - tensor.ndim)
|
|
return tensor.reshape(shape)
|
|
|
|
|
|
def assert_schedule_timesteps_compatible(schedule, timesteps):
|
|
"""
|
|
Check if schedule and timesteps are compatible.
|
|
"""
|
|
if schedule.T != timesteps.T:
|
|
raise ValueError("Schedule and timesteps must have the same T.")
|
|
if schedule.is_continuous() != timesteps.is_continuous():
|
|
raise ValueError("Schedule and timesteps must have the same continuity.")
|
|
|
|
|
|
def classifier_free_guidance(
|
|
pos: torch.Tensor,
|
|
neg: torch.Tensor,
|
|
scale: float,
|
|
rescale: float = 0.0,
|
|
):
|
|
"""
|
|
Apply classifier-free guidance.
|
|
"""
|
|
# Classifier-free guidance (https://arxiv.org/abs/2207.12598)
|
|
cfg = neg + scale * (pos - neg)
|
|
|
|
# Classifier-free guidance rescale (https://arxiv.org/pdf/2305.08891.pdf)
|
|
if rescale != 0.0:
|
|
pos_std = pos.std(dim=list(range(1, pos.ndim)), keepdim=True)
|
|
cfg_std = cfg.std(dim=list(range(1, cfg.ndim)), keepdim=True)
|
|
factor = pos_std / cfg_std
|
|
factor = rescale * factor + (1 - rescale)
|
|
cfg *= factor
|
|
|
|
return cfg
|
|
|
|
|
|
def classifier_free_guidance_dispatcher(
|
|
pos: Callable,
|
|
neg: Callable,
|
|
scale: float,
|
|
rescale: float = 0.0,
|
|
):
|
|
"""
|
|
Optionally execute models depending on classifer-free guidance scale.
|
|
"""
|
|
# If scale is 1, no need to execute neg model.
|
|
if scale == 1.0:
|
|
return pos()
|
|
|
|
# Otherwise, execute both pos nad neg models and apply cfg.
|
|
return classifier_free_guidance(
|
|
pos=pos(),
|
|
neg=neg(),
|
|
scale=scale,
|
|
rescale=rescale,
|
|
)
|