[Model]: Helios support (#1604)
Signed-off-by: princepride <wangzhipeng628@gmail.com> Signed-off-by: 汪志鹏 <wangzhipeng628@gmail.com> Signed-off-by: gcanlin <canlinguosdu@gmail.com> Co-authored-by: gcanlin <canlinguosdu@gmail.com>
This commit is contained in:
汪志鹏andgcanlin authored and GitHub committed
1 parent
675c35db97
commit
f9ec07f0e2
8 files changed
+3363
No files matched your search
@@ -0,0 +1,19 @@
|
||||
from .helios_transformer import HeliosTransformer3DModel
|
||||
from .pipeline_helios import (
|
||||
HeliosPipeline,
|
||||
create_transformer_from_config,
|
||||
get_helios_post_process_func,
|
||||
get_helios_pre_process_func,
|
||||
load_transformer_config,
|
||||
)
|
||||
from .scheduling_helios import HeliosScheduler
|
||||
|
||||
__all__ = [
|
||||
"HeliosPipeline",
|
||||
"HeliosTransformer3DModel",
|
||||
"HeliosScheduler",
|
||||
"get_helios_post_process_func",
|
||||
"get_helios_pre_process_func",
|
||||
"load_transformer_config",
|
||||
"create_transformer_from_config",
|
||||
]
|
||||
@@ -0,0 +1,929 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
# Adapted from Helios (https://github.com/BestWishYsh/Helios)
|
||||
|
||||
import math
|
||||
from collections.abc import Iterable
|
||||
from functools import lru_cache
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
from diffusers.models.embeddings import PixArtAlphaTextProjection, TimestepEmbedding, Timesteps
|
||||
from diffusers.models.modeling_outputs import Transformer2DModelOutput
|
||||
from diffusers.models.normalization import FP32LayerNorm
|
||||
from vllm.distributed import (
|
||||
get_tensor_model_parallel_rank,
|
||||
get_tensor_model_parallel_world_size,
|
||||
tensor_model_parallel_all_reduce,
|
||||
)
|
||||
from vllm.logger import init_logger
|
||||
from vllm.model_executor.layers.conv import Conv3dLayer
|
||||
from vllm.model_executor.layers.linear import ColumnParallelLinear, QKVParallelLinear, RowParallelLinear
|
||||
from vllm.model_executor.model_loader.weight_utils import default_weight_loader
|
||||
|
||||
from vllm_omni.diffusion.attention.layer import Attention
|
||||
from vllm_omni.diffusion.distributed.sp_plan import (
|
||||
SequenceParallelInput,
|
||||
SequenceParallelOutput,
|
||||
)
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
|
||||
def pad_for_3d_conv(x, kernel_size):
|
||||
b, c, t, h, w = x.shape
|
||||
pt, ph, pw = kernel_size
|
||||
pad_t = (pt - (t % pt)) % pt
|
||||
pad_h = (ph - (h % ph)) % ph
|
||||
pad_w = (pw - (w % pw)) % pw
|
||||
return F.pad(x, (0, pad_w, 0, pad_h, 0, pad_t), mode="replicate")
|
||||
|
||||
|
||||
def center_down_sample_3d(x, kernel_size):
|
||||
return F.avg_pool3d(x, kernel_size, stride=kernel_size)
|
||||
|
||||
|
||||
def apply_rotary_emb_helios(
|
||||
hidden_states: torch.Tensor,
|
||||
freqs_cis: torch.Tensor,
|
||||
) -> torch.Tensor:
|
||||
"""Apply Helios-style rotary embeddings.
|
||||
|
||||
freqs_cis contains [cos_t, cos_y, cos_x, sin_t, sin_y, sin_x] concatenated
|
||||
along the last dimension, with shape [B, seq, D*2] where D = DT+DY+DX.
|
||||
hidden_states has shape [B, seq, H, head_dim].
|
||||
"""
|
||||
x_1, x_2 = hidden_states.unflatten(-1, (-1, 2)).unbind(-1)
|
||||
cos, sin = freqs_cis.unsqueeze(-2).chunk(2, dim=-1)
|
||||
out = torch.empty_like(hidden_states)
|
||||
out[..., 0::2] = x_1 * cos[..., 0::2] - x_2 * sin[..., 1::2]
|
||||
out[..., 1::2] = x_1 * sin[..., 1::2] + x_2 * cos[..., 0::2]
|
||||
return out.type_as(hidden_states)
|
||||
|
||||
|
||||
class DistributedRMSNorm(nn.Module):
|
||||
"""RMSNorm that computes global RMS across tensor parallel ranks."""
|
||||
|
||||
def __init__(self, hidden_size: int, eps: float = 1e-6):
|
||||
super().__init__()
|
||||
self.eps = eps
|
||||
self.weight = nn.Parameter(torch.ones(hidden_size))
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
tp_size = get_tensor_model_parallel_world_size()
|
||||
|
||||
input_dtype = x.dtype
|
||||
x_float = x.float()
|
||||
local_sum_sq = (x_float**2).sum(dim=-1, keepdim=True)
|
||||
local_count = x.shape[-1]
|
||||
|
||||
if tp_size > 1:
|
||||
global_sum_sq = local_sum_sq.clone()
|
||||
tensor_model_parallel_all_reduce(global_sum_sq)
|
||||
global_count = local_count * tp_size
|
||||
else:
|
||||
global_sum_sq = local_sum_sq
|
||||
global_count = local_count
|
||||
|
||||
rms = torch.sqrt(global_sum_sq / global_count + self.eps)
|
||||
|
||||
output = (x_float / rms) * self.weight.float()
|
||||
return output.to(input_dtype)
|
||||
|
||||
|
||||
class ColumnParallelGELU(nn.Module):
|
||||
"""Column parallel linear with GELU activation."""
|
||||
|
||||
def __init__(self, dim_in: int, dim_out: int, *, approximate: str = "tanh", bias: bool = True):
|
||||
super().__init__()
|
||||
self.proj = ColumnParallelLinear(
|
||||
dim_in,
|
||||
dim_out,
|
||||
bias=bias,
|
||||
gather_output=False,
|
||||
return_bias=False,
|
||||
)
|
||||
self.approximate = approximate
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
x = self.proj(x)
|
||||
return F.gelu(x, approximate=self.approximate)
|
||||
|
||||
|
||||
class HeliosFeedForward(nn.Module):
|
||||
"""TP-enabled FeedForward network for Helios."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
dim: int,
|
||||
inner_dim: int,
|
||||
dim_out: int | None = None,
|
||||
bias: bool = True,
|
||||
) -> None:
|
||||
super().__init__()
|
||||
dim_out = dim_out or dim
|
||||
|
||||
self.net_0 = ColumnParallelGELU(dim, inner_dim, approximate="tanh", bias=bias)
|
||||
self.net_1 = nn.Identity()
|
||||
self.net_2 = RowParallelLinear(
|
||||
inner_dim,
|
||||
dim_out,
|
||||
bias=bias,
|
||||
input_is_parallel=True,
|
||||
return_bias=False,
|
||||
)
|
||||
|
||||
def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
|
||||
hidden_states = self.net_0(hidden_states)
|
||||
hidden_states = self.net_1(hidden_states)
|
||||
hidden_states = self.net_2(hidden_states)
|
||||
return hidden_states
|
||||
|
||||
|
||||
class HeliosRotaryPosEmbed(nn.Module):
|
||||
"""Helios-style 3D rotary position embeddings using explicit frame indices."""
|
||||
|
||||
def __init__(self, rope_dim: tuple[int, int, int], theta: float):
|
||||
super().__init__()
|
||||
self.DT, self.DY, self.DX = rope_dim
|
||||
self.theta = theta
|
||||
self.register_buffer("freqs_base_t", self._get_freqs_base(self.DT), persistent=False)
|
||||
self.register_buffer("freqs_base_y", self._get_freqs_base(self.DY), persistent=False)
|
||||
self.register_buffer("freqs_base_x", self._get_freqs_base(self.DX), persistent=False)
|
||||
|
||||
def _get_freqs_base(self, dim):
|
||||
return 1.0 / (self.theta ** (torch.arange(0, dim, 2, dtype=torch.float32)[: (dim // 2)] / dim))
|
||||
|
||||
@torch.no_grad()
|
||||
def get_frequency_batched(self, freqs_base, pos):
|
||||
freqs = torch.einsum("d,bthw->dbthw", freqs_base, pos)
|
||||
freqs = freqs.repeat_interleave(2, dim=0)
|
||||
return freqs.cos(), freqs.sin()
|
||||
|
||||
@torch.no_grad()
|
||||
@lru_cache(maxsize=32)
|
||||
def _get_spatial_meshgrid(self, height, width, device_str):
|
||||
device = torch.device(device_str)
|
||||
grid_y_coords = torch.arange(height, device=device, dtype=torch.float32)
|
||||
grid_x_coords = torch.arange(width, device=device, dtype=torch.float32)
|
||||
grid_y, grid_x = torch.meshgrid(grid_y_coords, grid_x_coords, indexing="ij")
|
||||
return grid_y, grid_x
|
||||
|
||||
@torch.no_grad()
|
||||
def forward(self, frame_indices, height, width, device):
|
||||
batch_size = frame_indices.shape[0]
|
||||
num_frames = frame_indices.shape[1]
|
||||
|
||||
frame_indices = frame_indices.to(device=device, dtype=torch.float32)
|
||||
grid_y, grid_x = self._get_spatial_meshgrid(height, width, str(device))
|
||||
|
||||
grid_t = frame_indices[:, :, None, None].expand(batch_size, num_frames, height, width)
|
||||
grid_y_batch = grid_y[None, None, :, :].expand(batch_size, num_frames, -1, -1)
|
||||
grid_x_batch = grid_x[None, None, :, :].expand(batch_size, num_frames, -1, -1)
|
||||
|
||||
freqs_cos_t, freqs_sin_t = self.get_frequency_batched(self.freqs_base_t, grid_t)
|
||||
freqs_cos_y, freqs_sin_y = self.get_frequency_batched(self.freqs_base_y, grid_y_batch)
|
||||
freqs_cos_x, freqs_sin_x = self.get_frequency_batched(self.freqs_base_x, grid_x_batch)
|
||||
|
||||
result = torch.cat(
|
||||
[freqs_cos_t, freqs_cos_y, freqs_cos_x, freqs_sin_t, freqs_sin_y, freqs_sin_x],
|
||||
dim=0,
|
||||
)
|
||||
|
||||
return result.permute(1, 0, 2, 3, 4)
|
||||
|
||||
|
||||
class HeliosTimeTextEmbedding(nn.Module):
|
||||
"""Combined time and text condition embeddings for Helios."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
dim: int,
|
||||
time_freq_dim: int,
|
||||
time_proj_dim: int,
|
||||
text_embed_dim: int,
|
||||
):
|
||||
super().__init__()
|
||||
|
||||
self.timesteps_proj = Timesteps(num_channels=time_freq_dim, flip_sin_to_cos=True, downscale_freq_shift=0)
|
||||
self.time_embedder = TimestepEmbedding(in_channels=time_freq_dim, time_embed_dim=dim)
|
||||
self.act_fn = nn.SiLU()
|
||||
self.time_proj = nn.Linear(dim, time_proj_dim)
|
||||
self.text_embedder = PixArtAlphaTextProjection(text_embed_dim, dim, act_fn="gelu_tanh")
|
||||
|
||||
def forward(
|
||||
self,
|
||||
timestep: torch.Tensor,
|
||||
encoder_hidden_states: torch.Tensor | None = None,
|
||||
is_return_encoder_hidden_states: bool = True,
|
||||
):
|
||||
timestep = self.timesteps_proj(timestep)
|
||||
|
||||
time_embedder_dtype = next(iter(self.time_embedder.parameters())).dtype
|
||||
if timestep.dtype != time_embedder_dtype and time_embedder_dtype != torch.int8:
|
||||
timestep = timestep.to(time_embedder_dtype)
|
||||
temb = self.time_embedder(timestep).type_as(encoder_hidden_states)
|
||||
timestep_proj = self.time_proj(self.act_fn(temb))
|
||||
|
||||
if encoder_hidden_states is not None and is_return_encoder_hidden_states:
|
||||
encoder_hidden_states = self.text_embedder(encoder_hidden_states)
|
||||
|
||||
return temb, timestep_proj, encoder_hidden_states
|
||||
|
||||
|
||||
class HeliosOutputNorm(nn.Module):
|
||||
"""Output normalization that extracts only original_context_length tokens."""
|
||||
|
||||
def __init__(self, dim: int, eps: float = 1e-6):
|
||||
super().__init__()
|
||||
self.scale_shift_table = nn.Parameter(torch.randn(1, 2, dim) / dim**0.5)
|
||||
self.norm = FP32LayerNorm(dim, eps, elementwise_affine=False)
|
||||
|
||||
def forward(self, hidden_states: torch.Tensor, temb: torch.Tensor, original_context_length: int):
|
||||
temb = temb[:, -original_context_length:, :]
|
||||
shift, scale = (self.scale_shift_table.unsqueeze(0).to(temb.device) + temb.unsqueeze(2)).chunk(2, dim=2)
|
||||
shift, scale = shift.squeeze(2).to(hidden_states.device), scale.squeeze(2).to(hidden_states.device)
|
||||
hidden_states = hidden_states[:, -original_context_length:, :]
|
||||
hidden_states = (self.norm(hidden_states.float()) * (1 + scale) + shift).type_as(hidden_states)
|
||||
return hidden_states
|
||||
|
||||
|
||||
class HeliosSelfAttention(nn.Module):
|
||||
"""Optimized self-attention for Helios with history amplification support."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
dim: int,
|
||||
num_heads: int,
|
||||
head_dim: int,
|
||||
eps: float = 1e-5,
|
||||
dropout: float = 0.0,
|
||||
is_amplify_history: bool = False,
|
||||
history_scale_mode: str = "per_head",
|
||||
):
|
||||
super().__init__()
|
||||
|
||||
self.dim = dim
|
||||
self.num_heads = num_heads
|
||||
self.head_dim = head_dim
|
||||
self.inner_dim = num_heads * head_dim
|
||||
|
||||
self.to_qkv = QKVParallelLinear(
|
||||
hidden_size=dim,
|
||||
head_size=head_dim,
|
||||
total_num_heads=num_heads,
|
||||
bias=True,
|
||||
)
|
||||
|
||||
self.num_heads = self.to_qkv.num_heads
|
||||
self.num_kv_heads = self.to_qkv.num_kv_heads
|
||||
self.tp_inner_dim = self.num_heads * head_dim
|
||||
|
||||
self.norm_q = DistributedRMSNorm(self.tp_inner_dim, eps=eps)
|
||||
self.norm_k = DistributedRMSNorm(self.tp_inner_dim, eps=eps)
|
||||
|
||||
self.to_out = RowParallelLinear(
|
||||
self.inner_dim,
|
||||
dim,
|
||||
bias=True,
|
||||
input_is_parallel=True,
|
||||
return_bias=False,
|
||||
)
|
||||
self.dropout = nn.Dropout(dropout)
|
||||
|
||||
self.attn = Attention(
|
||||
num_heads=self.num_heads,
|
||||
head_size=head_dim,
|
||||
num_kv_heads=self.num_kv_heads,
|
||||
softmax_scale=1.0 / (head_dim**0.5),
|
||||
causal=False,
|
||||
)
|
||||
|
||||
self.is_amplify_history = is_amplify_history
|
||||
if is_amplify_history:
|
||||
tp_size = get_tensor_model_parallel_world_size()
|
||||
local_heads = num_heads // tp_size
|
||||
if history_scale_mode == "scalar":
|
||||
self.history_key_scale = nn.Parameter(torch.ones(1))
|
||||
elif history_scale_mode == "per_head":
|
||||
self.history_key_scale = nn.Parameter(torch.ones(local_heads))
|
||||
else:
|
||||
raise ValueError(f"Unknown history_scale_mode: {history_scale_mode}")
|
||||
self.history_scale_mode = history_scale_mode
|
||||
self.max_scale = 10.0
|
||||
|
||||
def forward(
|
||||
self,
|
||||
hidden_states: torch.Tensor,
|
||||
rotary_emb: torch.Tensor | None = None,
|
||||
original_context_length: int | None = None,
|
||||
) -> torch.Tensor:
|
||||
qkv, _ = self.to_qkv(hidden_states)
|
||||
|
||||
q_size = self.num_heads * self.head_dim
|
||||
kv_size = self.num_kv_heads * self.head_dim
|
||||
query, key, value = qkv.split([q_size, kv_size, kv_size], dim=-1)
|
||||
|
||||
query = self.norm_q(query)
|
||||
key = self.norm_k(key)
|
||||
|
||||
query = query.unflatten(2, (self.num_heads, self.head_dim))
|
||||
key = key.unflatten(2, (self.num_kv_heads, self.head_dim))
|
||||
value = value.unflatten(2, (self.num_kv_heads, self.head_dim))
|
||||
|
||||
if rotary_emb is not None:
|
||||
query = apply_rotary_emb_helios(query, rotary_emb)
|
||||
key = apply_rotary_emb_helios(key, rotary_emb)
|
||||
|
||||
if self.is_amplify_history and original_context_length is not None:
|
||||
history_seq_len = hidden_states.shape[1] - original_context_length
|
||||
if history_seq_len > 0:
|
||||
scale_key = 1.0 + torch.sigmoid(self.history_key_scale) * (self.max_scale - 1.0)
|
||||
if self.history_scale_mode == "per_head":
|
||||
scale_key = scale_key.view(1, 1, -1, 1)
|
||||
key = torch.cat(
|
||||
[key[:, :history_seq_len] * scale_key, key[:, history_seq_len:]],
|
||||
dim=1,
|
||||
)
|
||||
|
||||
hidden_states = self.attn(query, key, value)
|
||||
hidden_states = hidden_states.flatten(2, 3)
|
||||
hidden_states = hidden_states.type_as(query)
|
||||
|
||||
hidden_states = self.to_out(hidden_states)
|
||||
hidden_states = self.dropout(hidden_states)
|
||||
|
||||
return hidden_states
|
||||
|
||||
|
||||
class HeliosCrossAttention(nn.Module):
|
||||
"""Optimized cross-attention for Helios."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
dim: int,
|
||||
num_heads: int,
|
||||
head_dim: int,
|
||||
eps: float = 1e-5,
|
||||
dropout: float = 0.0,
|
||||
):
|
||||
super().__init__()
|
||||
|
||||
self.dim = dim
|
||||
self.num_heads = num_heads
|
||||
self.head_dim = head_dim
|
||||
self.inner_dim = num_heads * head_dim
|
||||
|
||||
self.to_q = ColumnParallelLinear(
|
||||
dim,
|
||||
self.inner_dim,
|
||||
bias=True,
|
||||
gather_output=False,
|
||||
return_bias=False,
|
||||
)
|
||||
self.to_k = ColumnParallelLinear(
|
||||
dim,
|
||||
self.inner_dim,
|
||||
bias=True,
|
||||
gather_output=False,
|
||||
return_bias=False,
|
||||
)
|
||||
self.to_v = ColumnParallelLinear(
|
||||
dim,
|
||||
self.inner_dim,
|
||||
bias=True,
|
||||
gather_output=False,
|
||||
return_bias=False,
|
||||
)
|
||||
|
||||
tp_size = get_tensor_model_parallel_world_size()
|
||||
self.num_heads = num_heads // tp_size
|
||||
self.tp_inner_dim = self.num_heads * head_dim
|
||||
|
||||
self.norm_q = DistributedRMSNorm(self.tp_inner_dim, eps=eps)
|
||||
self.norm_k = DistributedRMSNorm(self.tp_inner_dim, eps=eps)
|
||||
|
||||
self.to_out = RowParallelLinear(
|
||||
self.inner_dim,
|
||||
dim,
|
||||
bias=True,
|
||||
input_is_parallel=True,
|
||||
return_bias=False,
|
||||
)
|
||||
self.dropout = nn.Dropout(dropout)
|
||||
|
||||
self.attn = Attention(
|
||||
num_heads=self.num_heads,
|
||||
head_size=head_dim,
|
||||
num_kv_heads=self.num_heads,
|
||||
softmax_scale=1.0 / (head_dim**0.5),
|
||||
causal=False,
|
||||
)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
hidden_states: torch.Tensor,
|
||||
encoder_hidden_states: torch.Tensor,
|
||||
) -> torch.Tensor:
|
||||
query = self.to_q(hidden_states)
|
||||
query = self.norm_q(query)
|
||||
|
||||
key = self.to_k(encoder_hidden_states)
|
||||
value = self.to_v(encoder_hidden_states)
|
||||
key = self.norm_k(key)
|
||||
|
||||
query = query.unflatten(2, (self.num_heads, self.head_dim))
|
||||
key = key.unflatten(2, (self.num_heads, self.head_dim))
|
||||
value = value.unflatten(2, (self.num_heads, self.head_dim))
|
||||
|
||||
hidden_states = self.attn(query, key, value)
|
||||
hidden_states = hidden_states.flatten(2, 3)
|
||||
hidden_states = hidden_states.type_as(query)
|
||||
|
||||
hidden_states = self.to_out(hidden_states)
|
||||
hidden_states = self.dropout(hidden_states)
|
||||
|
||||
return hidden_states
|
||||
|
||||
|
||||
class HeliosTransformerBlock(nn.Module):
|
||||
"""Transformer block with guidance cross-attention and history support."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
dim: int,
|
||||
ffn_dim: int,
|
||||
num_heads: int,
|
||||
eps: float = 1e-6,
|
||||
cross_attn_norm: bool = False,
|
||||
guidance_cross_attn: bool = False,
|
||||
is_amplify_history: bool = False,
|
||||
history_scale_mode: str = "per_head",
|
||||
):
|
||||
super().__init__()
|
||||
|
||||
head_dim = dim // num_heads
|
||||
|
||||
# 1. Self-attention
|
||||
self.norm1 = FP32LayerNorm(dim, eps, elementwise_affine=False)
|
||||
self.attn1 = HeliosSelfAttention(
|
||||
dim=dim,
|
||||
num_heads=num_heads,
|
||||
head_dim=head_dim,
|
||||
eps=eps,
|
||||
is_amplify_history=is_amplify_history,
|
||||
history_scale_mode=history_scale_mode,
|
||||
)
|
||||
|
||||
# 2. Cross-attention
|
||||
self.attn2 = HeliosCrossAttention(
|
||||
dim=dim,
|
||||
num_heads=num_heads,
|
||||
head_dim=head_dim,
|
||||
eps=eps,
|
||||
)
|
||||
self.norm2 = FP32LayerNorm(dim, eps, elementwise_affine=True) if cross_attn_norm else nn.Identity()
|
||||
|
||||
# 3. Feed-forward
|
||||
self.ffn = HeliosFeedForward(dim=dim, inner_dim=ffn_dim, dim_out=dim)
|
||||
self.norm3 = FP32LayerNorm(dim, eps, elementwise_affine=False)
|
||||
|
||||
self.scale_shift_table = nn.Parameter(torch.randn(1, 6, dim) / dim**0.5)
|
||||
|
||||
# 4. Guidance cross-attention flag
|
||||
self.guidance_cross_attn = guidance_cross_attn
|
||||
|
||||
def forward(
|
||||
self,
|
||||
hidden_states: torch.Tensor,
|
||||
encoder_hidden_states: torch.Tensor,
|
||||
temb: torch.Tensor,
|
||||
rotary_emb: torch.Tensor,
|
||||
original_context_length: int | None = None,
|
||||
) -> torch.Tensor:
|
||||
if temb.ndim == 4:
|
||||
shift_msa, scale_msa, gate_msa, c_shift_msa, c_scale_msa, c_gate_msa = (
|
||||
self.scale_shift_table.unsqueeze(0) + temb.float()
|
||||
).chunk(6, dim=2)
|
||||
shift_msa = shift_msa.squeeze(2)
|
||||
scale_msa = scale_msa.squeeze(2)
|
||||
gate_msa = gate_msa.squeeze(2)
|
||||
c_shift_msa = c_shift_msa.squeeze(2)
|
||||
c_scale_msa = c_scale_msa.squeeze(2)
|
||||
c_gate_msa = c_gate_msa.squeeze(2)
|
||||
else:
|
||||
shift_msa, scale_msa, gate_msa, c_shift_msa, c_scale_msa, c_gate_msa = (
|
||||
self.scale_shift_table + temb.float()
|
||||
).chunk(6, dim=1)
|
||||
|
||||
# 1. Self-attention
|
||||
norm_hidden_states = (self.norm1(hidden_states.float()) * (1 + scale_msa) + shift_msa).type_as(hidden_states)
|
||||
attn_output = self.attn1(norm_hidden_states, rotary_emb, original_context_length)
|
||||
hidden_states = (hidden_states.float() + attn_output * gate_msa).type_as(hidden_states)
|
||||
|
||||
# 2. Cross-attention (with optional guidance: only current chunk attends to text)
|
||||
if self.guidance_cross_attn and original_context_length is not None:
|
||||
history_seq_len = hidden_states.shape[1] - original_context_length
|
||||
|
||||
history_hidden_states, current_hidden_states = (
|
||||
hidden_states[:, :history_seq_len],
|
||||
hidden_states[:, history_seq_len:],
|
||||
)
|
||||
norm_hidden_states = self.norm2(current_hidden_states.float()).type_as(current_hidden_states)
|
||||
attn_output = self.attn2(norm_hidden_states, encoder_hidden_states)
|
||||
current_hidden_states = current_hidden_states + attn_output
|
||||
hidden_states = torch.cat([history_hidden_states, current_hidden_states], dim=1)
|
||||
else:
|
||||
norm_hidden_states = self.norm2(hidden_states.float()).type_as(hidden_states)
|
||||
attn_output = self.attn2(norm_hidden_states, encoder_hidden_states)
|
||||
hidden_states = hidden_states + attn_output
|
||||
|
||||
# 3. Feed-forward
|
||||
norm_hidden_states = (self.norm3(hidden_states.float()) * (1 + c_scale_msa) + c_shift_msa).type_as(
|
||||
hidden_states
|
||||
)
|
||||
ff_output = self.ffn(norm_hidden_states)
|
||||
hidden_states = (hidden_states.float() + ff_output.float() * c_gate_msa).type_as(hidden_states)
|
||||
|
||||
return hidden_states
|
||||
|
||||
|
||||
class HeliosTransformer3DModel(nn.Module):
|
||||
"""Optimized Helios Transformer model for video generation using vLLM layers.
|
||||
|
||||
Helios extends the Wan2.2 architecture with multi-term memory patches,
|
||||
guidance cross-attention, and chunked video generation support.
|
||||
"""
|
||||
|
||||
_repeated_blocks = ["HeliosTransformerBlock"]
|
||||
_layerwise_offload_blocks_attr = "blocks"
|
||||
packed_modules_mapping = {
|
||||
"to_qkv": ["to_q", "to_k", "to_v"],
|
||||
}
|
||||
|
||||
@staticmethod
|
||||
def _is_transformer_block(name: str, module) -> bool:
|
||||
return "blocks" in name and name.split(".")[-1].isdigit()
|
||||
|
||||
_hsdp_shard_conditions = [_is_transformer_block]
|
||||
|
||||
_sp_plan = {
|
||||
"rope": {
|
||||
0: SequenceParallelInput(split_dim=1, expected_dims=4, split_output=True, auto_pad=True),
|
||||
1: SequenceParallelInput(split_dim=1, expected_dims=4, split_output=True, auto_pad=True),
|
||||
},
|
||||
"blocks.0": {
|
||||
"hidden_states": SequenceParallelInput(split_dim=1, expected_dims=3, auto_pad=True),
|
||||
},
|
||||
"proj_out": SequenceParallelOutput(gather_dim=1, expected_dims=3),
|
||||
}
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
patch_size: tuple[int, ...] = (1, 2, 2),
|
||||
num_attention_heads: int = 40,
|
||||
attention_head_dim: int = 128,
|
||||
in_channels: int = 16,
|
||||
out_channels: int = 16,
|
||||
text_dim: int = 4096,
|
||||
freq_dim: int = 256,
|
||||
ffn_dim: int = 13824,
|
||||
num_layers: int = 40,
|
||||
cross_attn_norm: bool = True,
|
||||
qk_norm: str | None = "rms_norm_across_heads",
|
||||
eps: float = 1e-6,
|
||||
added_kv_proj_dim: int | None = None,
|
||||
rope_dim: tuple[int, ...] = (44, 42, 42),
|
||||
rope_theta: float = 10000.0,
|
||||
guidance_cross_attn: bool = True,
|
||||
zero_history_timestep: bool = True,
|
||||
has_multi_term_memory_patch: bool = True,
|
||||
is_amplify_history: bool = False,
|
||||
history_scale_mode: str = "per_head",
|
||||
) -> None:
|
||||
super().__init__()
|
||||
|
||||
inner_dim = num_attention_heads * attention_head_dim
|
||||
out_channels = out_channels or in_channels
|
||||
|
||||
self.config = type(
|
||||
"Config",
|
||||
(),
|
||||
{
|
||||
"patch_size": patch_size,
|
||||
"num_attention_heads": num_attention_heads,
|
||||
"attention_head_dim": attention_head_dim,
|
||||
"in_channels": in_channels,
|
||||
"out_channels": out_channels,
|
||||
"text_dim": text_dim,
|
||||
"freq_dim": freq_dim,
|
||||
"ffn_dim": ffn_dim,
|
||||
"num_layers": num_layers,
|
||||
"cross_attn_norm": cross_attn_norm,
|
||||
"qk_norm": qk_norm,
|
||||
"eps": eps,
|
||||
"added_kv_proj_dim": added_kv_proj_dim,
|
||||
"rope_dim": rope_dim,
|
||||
"rope_theta": rope_theta,
|
||||
"guidance_cross_attn": guidance_cross_attn,
|
||||
"zero_history_timestep": zero_history_timestep,
|
||||
"has_multi_term_memory_patch": has_multi_term_memory_patch,
|
||||
"is_amplify_history": is_amplify_history,
|
||||
"history_scale_mode": history_scale_mode,
|
||||
},
|
||||
)()
|
||||
|
||||
self.inner_dim = inner_dim
|
||||
self.zero_history_timestep = zero_history_timestep
|
||||
self.has_multi_term_memory_patch = has_multi_term_memory_patch
|
||||
|
||||
# 1. Patch & position embedding
|
||||
self.rope = HeliosRotaryPosEmbed(rope_dim=rope_dim, theta=rope_theta)
|
||||
self.patch_embedding = Conv3dLayer(
|
||||
in_channels=in_channels,
|
||||
out_channels=inner_dim,
|
||||
kernel_size=patch_size,
|
||||
stride=patch_size,
|
||||
)
|
||||
|
||||
# 2. Multi-term memory patches
|
||||
if has_multi_term_memory_patch:
|
||||
self.patch_short = Conv3dLayer(
|
||||
in_channels=in_channels,
|
||||
out_channels=inner_dim,
|
||||
kernel_size=(1, 2, 2),
|
||||
stride=(1, 2, 2),
|
||||
)
|
||||
self.patch_mid = Conv3dLayer(
|
||||
in_channels=in_channels,
|
||||
out_channels=inner_dim,
|
||||
kernel_size=(2, 4, 4),
|
||||
stride=(2, 4, 4),
|
||||
)
|
||||
self.patch_long = Conv3dLayer(
|
||||
in_channels=in_channels,
|
||||
out_channels=inner_dim,
|
||||
kernel_size=(4, 8, 8),
|
||||
stride=(4, 8, 8),
|
||||
)
|
||||
|
||||
# 3. Condition embeddings
|
||||
self.condition_embedder = HeliosTimeTextEmbedding(
|
||||
dim=inner_dim,
|
||||
time_freq_dim=freq_dim,
|
||||
time_proj_dim=inner_dim * 6,
|
||||
text_embed_dim=text_dim,
|
||||
)
|
||||
|
||||
# 4. Transformer blocks
|
||||
self.blocks = nn.ModuleList(
|
||||
[
|
||||
HeliosTransformerBlock(
|
||||
inner_dim,
|
||||
ffn_dim,
|
||||
num_attention_heads,
|
||||
eps,
|
||||
cross_attn_norm,
|
||||
guidance_cross_attn=guidance_cross_attn,
|
||||
is_amplify_history=is_amplify_history,
|
||||
history_scale_mode=history_scale_mode,
|
||||
)
|
||||
for _ in range(num_layers)
|
||||
]
|
||||
)
|
||||
|
||||
# 5. Output norm & projection
|
||||
self.norm_out = HeliosOutputNorm(inner_dim, eps)
|
||||
self.proj_out = nn.Linear(inner_dim, out_channels * math.prod(patch_size))
|
||||
|
||||
@property
|
||||
def dtype(self) -> torch.dtype:
|
||||
return next(self.parameters()).dtype
|
||||
|
||||
def forward(
|
||||
self,
|
||||
hidden_states: torch.Tensor,
|
||||
timestep: torch.LongTensor,
|
||||
encoder_hidden_states: torch.Tensor,
|
||||
indices_hidden_states: torch.Tensor | None = None,
|
||||
indices_latents_history_short: torch.Tensor | None = None,
|
||||
indices_latents_history_mid: torch.Tensor | None = None,
|
||||
indices_latents_history_long: torch.Tensor | None = None,
|
||||
latents_history_short: torch.Tensor | None = None,
|
||||
latents_history_mid: torch.Tensor | None = None,
|
||||
latents_history_long: torch.Tensor | None = None,
|
||||
return_dict: bool = True,
|
||||
attention_kwargs: dict[str, Any] | None = None,
|
||||
) -> torch.Tensor | Transformer2DModelOutput:
|
||||
batch_size = hidden_states.shape[0]
|
||||
p_t, p_h, p_w = self.config.patch_size
|
||||
|
||||
# 1. Process noisy latents
|
||||
hidden_states = self.patch_embedding(hidden_states)
|
||||
_, _, post_patch_num_frames, post_patch_height, post_patch_width = hidden_states.shape
|
||||
|
||||
if indices_hidden_states is None:
|
||||
indices_hidden_states = torch.arange(0, post_patch_num_frames).unsqueeze(0).expand(batch_size, -1)
|
||||
|
||||
hidden_states = hidden_states.flatten(2).transpose(1, 2)
|
||||
rotary_emb = self.rope(
|
||||
frame_indices=indices_hidden_states,
|
||||
height=post_patch_height,
|
||||
width=post_patch_width,
|
||||
device=hidden_states.device,
|
||||
)
|
||||
rotary_emb = rotary_emb.flatten(2).transpose(1, 2)
|
||||
original_context_length = hidden_states.shape[1]
|
||||
|
||||
# 2. Process short history latents
|
||||
if latents_history_short is not None and indices_latents_history_short is not None:
|
||||
latents_history_short = latents_history_short.to(hidden_states)
|
||||
latents_history_short = self.patch_short(latents_history_short)
|
||||
_, _, _, H1, W1 = latents_history_short.shape
|
||||
latents_history_short = latents_history_short.flatten(2).transpose(1, 2)
|
||||
|
||||
rotary_emb_history_short = self.rope(
|
||||
frame_indices=indices_latents_history_short,
|
||||
height=H1,
|
||||
width=W1,
|
||||
device=latents_history_short.device,
|
||||
)
|
||||
rotary_emb_history_short = rotary_emb_history_short.flatten(2).transpose(1, 2)
|
||||
|
||||
hidden_states = torch.cat([latents_history_short, hidden_states], dim=1)
|
||||
rotary_emb = torch.cat([rotary_emb_history_short, rotary_emb], dim=1)
|
||||
|
||||
# 3. Process mid history latents
|
||||
if latents_history_mid is not None and indices_latents_history_mid is not None:
|
||||
latents_history_mid = latents_history_mid.to(hidden_states)
|
||||
latents_history_mid = pad_for_3d_conv(latents_history_mid, (2, 4, 4))
|
||||
latents_history_mid = self.patch_mid(latents_history_mid)
|
||||
latents_history_mid = latents_history_mid.flatten(2).transpose(1, 2)
|
||||
|
||||
rotary_emb_history_mid = self.rope(
|
||||
frame_indices=indices_latents_history_mid,
|
||||
height=H1,
|
||||
width=W1,
|
||||
device=latents_history_mid.device,
|
||||
)
|
||||
rotary_emb_history_mid = pad_for_3d_conv(rotary_emb_history_mid, (2, 2, 2))
|
||||
rotary_emb_history_mid = center_down_sample_3d(rotary_emb_history_mid, (2, 2, 2))
|
||||
rotary_emb_history_mid = rotary_emb_history_mid.flatten(2).transpose(1, 2)
|
||||
|
||||
hidden_states = torch.cat([latents_history_mid, hidden_states], dim=1)
|
||||
rotary_emb = torch.cat([rotary_emb_history_mid, rotary_emb], dim=1)
|
||||
|
||||
# 4. Process long history latents
|
||||
if latents_history_long is not None and indices_latents_history_long is not None:
|
||||
latents_history_long = latents_history_long.to(hidden_states)
|
||||
latents_history_long = pad_for_3d_conv(latents_history_long, (4, 8, 8))
|
||||
latents_history_long = self.patch_long(latents_history_long)
|
||||
latents_history_long = latents_history_long.flatten(2).transpose(1, 2)
|
||||
|
||||
rotary_emb_history_long = self.rope(
|
||||
frame_indices=indices_latents_history_long,
|
||||
height=H1,
|
||||
width=W1,
|
||||
device=latents_history_long.device,
|
||||
)
|
||||
rotary_emb_history_long = pad_for_3d_conv(rotary_emb_history_long, (4, 4, 4))
|
||||
rotary_emb_history_long = center_down_sample_3d(rotary_emb_history_long, (4, 4, 4))
|
||||
rotary_emb_history_long = rotary_emb_history_long.flatten(2).transpose(1, 2)
|
||||
|
||||
hidden_states = torch.cat([latents_history_long, hidden_states], dim=1)
|
||||
rotary_emb = torch.cat([rotary_emb_history_long, rotary_emb], dim=1)
|
||||
|
||||
history_context_length = hidden_states.shape[1] - original_context_length
|
||||
|
||||
# 5. Compute timestep embeddings
|
||||
if indices_hidden_states is not None and self.zero_history_timestep:
|
||||
timestep_t0 = torch.zeros((1), dtype=timestep.dtype, device=timestep.device)
|
||||
temb_t0, timestep_proj_t0, _ = self.condition_embedder(
|
||||
timestep_t0, encoder_hidden_states, is_return_encoder_hidden_states=False
|
||||
)
|
||||
temb_t0 = temb_t0.unsqueeze(1).expand(batch_size, history_context_length, -1)
|
||||
timestep_proj_t0 = (
|
||||
timestep_proj_t0.unflatten(-1, (6, -1))
|
||||
.view(1, 6, 1, -1)
|
||||
.expand(batch_size, -1, history_context_length, -1)
|
||||
)
|
||||
|
||||
temb, timestep_proj, encoder_hidden_states = self.condition_embedder(timestep, encoder_hidden_states)
|
||||
timestep_proj = timestep_proj.unflatten(-1, (6, -1))
|
||||
|
||||
if indices_hidden_states is not None and not self.zero_history_timestep:
|
||||
main_repeat_size = hidden_states.shape[1]
|
||||
else:
|
||||
main_repeat_size = original_context_length
|
||||
temb = temb.view(batch_size, 1, -1).expand(batch_size, main_repeat_size, -1)
|
||||
timestep_proj = timestep_proj.view(batch_size, 6, 1, -1).expand(batch_size, 6, main_repeat_size, -1)
|
||||
|
||||
if indices_hidden_states is not None and self.zero_history_timestep:
|
||||
temb = torch.cat([temb_t0, temb], dim=1)
|
||||
timestep_proj = torch.cat([timestep_proj_t0, timestep_proj], dim=2)
|
||||
|
||||
if timestep_proj.ndim == 4:
|
||||
timestep_proj = timestep_proj.permute(0, 2, 1, 3)
|
||||
|
||||
# 6. Transformer blocks
|
||||
hidden_states = hidden_states.contiguous()
|
||||
encoder_hidden_states = encoder_hidden_states.contiguous()
|
||||
rotary_emb = rotary_emb.contiguous()
|
||||
|
||||
for block in self.blocks:
|
||||
hidden_states = block(
|
||||
hidden_states,
|
||||
encoder_hidden_states,
|
||||
timestep_proj,
|
||||
rotary_emb,
|
||||
original_context_length,
|
||||
)
|
||||
|
||||
# 7. Output normalization
|
||||
hidden_states = self.norm_out(hidden_states, temb, original_context_length)
|
||||
hidden_states = self.proj_out(hidden_states)
|
||||
|
||||
# 8. Unpatchify
|
||||
hidden_states = hidden_states.reshape(
|
||||
batch_size, post_patch_num_frames, post_patch_height, post_patch_width, p_t, p_h, p_w, -1
|
||||
)
|
||||
hidden_states = hidden_states.permute(0, 7, 1, 4, 2, 5, 3, 6)
|
||||
output = hidden_states.flatten(6, 7).flatten(4, 5).flatten(2, 3)
|
||||
|
||||
if not return_dict:
|
||||
return (output,)
|
||||
|
||||
return Transformer2DModelOutput(sample=output)
|
||||
|
||||
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
|
||||
"""Load weights with QKV fusion, FFN remapping, and TP norm sharding."""
|
||||
tp_rank = get_tensor_model_parallel_rank()
|
||||
tp_size = get_tensor_model_parallel_world_size()
|
||||
|
||||
stacked_params_mapping = [
|
||||
(".attn1.to_qkv", ".attn1.to_q", "q"),
|
||||
(".attn1.to_qkv", ".attn1.to_k", "k"),
|
||||
(".attn1.to_qkv", ".attn1.to_v", "v"),
|
||||
]
|
||||
|
||||
params_dict = dict(self.named_parameters())
|
||||
loaded_params: set[str] = set()
|
||||
|
||||
for name, loaded_weight in weights:
|
||||
original_name = name
|
||||
lookup_name = name
|
||||
|
||||
# Handle QKV fusion for self-attention
|
||||
for param_name, weight_name, shard_id in stacked_params_mapping:
|
||||
if weight_name not in original_name:
|
||||
continue
|
||||
lookup_name = original_name.replace(weight_name, param_name)
|
||||
param = params_dict[lookup_name]
|
||||
weight_loader = param.weight_loader
|
||||
weight_loader(param, loaded_weight, shard_id)
|
||||
break
|
||||
else:
|
||||
# FFN remapping: net.0 -> net_0, net.2 -> net_2
|
||||
if ".ffn.net.0." in lookup_name:
|
||||
lookup_name = lookup_name.replace(".ffn.net.0.", ".ffn.net_0.")
|
||||
elif ".ffn.net.2." in lookup_name:
|
||||
lookup_name = lookup_name.replace(".ffn.net.2.", ".ffn.net_2.")
|
||||
|
||||
# Output projection: to_out.0 -> to_out
|
||||
if ".to_out.0." in lookup_name:
|
||||
lookup_name = lookup_name.replace(".to_out.0.", ".to_out.")
|
||||
|
||||
if lookup_name not in params_dict:
|
||||
logger.warning(f"Skipping weight {original_name} -> {lookup_name}")
|
||||
continue
|
||||
|
||||
param = params_dict[lookup_name]
|
||||
|
||||
# TP sharding for RMSNorm weights
|
||||
if tp_size > 1 and any(
|
||||
norm_name in lookup_name
|
||||
for norm_name in [
|
||||
".attn1.norm_q.",
|
||||
".attn1.norm_k.",
|
||||
".attn2.norm_q.",
|
||||
".attn2.norm_k.",
|
||||
]
|
||||
):
|
||||
shard_size = loaded_weight.shape[0] // tp_size
|
||||
loaded_weight = loaded_weight[tp_rank * shard_size : (tp_rank + 1) * shard_size]
|
||||
|
||||
# TP sharding for history_key_scale (per_head mode)
|
||||
if tp_size > 1 and ".history_key_scale" in lookup_name and loaded_weight.dim() == 1:
|
||||
shard_size = loaded_weight.shape[0] // tp_size
|
||||
if shard_size > 0:
|
||||
loaded_weight = loaded_weight[tp_rank * shard_size : (tp_rank + 1) * shard_size]
|
||||
|
||||
weight_loader = getattr(param, "weight_loader", default_weight_loader)
|
||||
weight_loader(param, loaded_weight)
|
||||
|
||||
loaded_params.add(original_name)
|
||||
loaded_params.add(lookup_name)
|
||||
|
||||
return loaded_params
|
||||
File diff suppressed because it is too large.
Load diff
@@ -0,0 +1,800 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
# Adapted from Helios (https://github.com/BestWishYsh/Helios)
|
||||
|
||||
import math
|
||||
from dataclasses import dataclass
|
||||
from typing import Literal
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
from diffusers.configuration_utils import ConfigMixin, register_to_config
|
||||
from diffusers.schedulers.scheduling_utils import SchedulerMixin
|
||||
from diffusers.utils import BaseOutput
|
||||
|
||||
|
||||
@dataclass
|
||||
class HeliosSchedulerOutput(BaseOutput):
|
||||
prev_sample: torch.FloatTensor
|
||||
model_outputs: torch.FloatTensor | None = None
|
||||
last_sample: torch.FloatTensor | None = None
|
||||
this_order: int | None = None
|
||||
|
||||
|
||||
class HeliosScheduler(SchedulerMixin, ConfigMixin):
|
||||
_compatibles = []
|
||||
order = 1
|
||||
|
||||
@register_to_config
|
||||
def __init__(
|
||||
self,
|
||||
num_train_timesteps: int = 1000,
|
||||
shift: float = 1.0,
|
||||
stages: int = 3,
|
||||
stage_range: list = [0, 1 / 3, 2 / 3, 1],
|
||||
gamma: float = 1 / 3,
|
||||
thresholding: bool = False,
|
||||
prediction_type: str = "flow_prediction",
|
||||
solver_order: int = 2,
|
||||
predict_x0: bool = True,
|
||||
solver_type: str = "bh2",
|
||||
lower_order_final: bool = True,
|
||||
disable_corrector: list[int] = [],
|
||||
solver_p: SchedulerMixin = None,
|
||||
use_flow_sigmas: bool = True,
|
||||
scheduler_type: str = "unipc",
|
||||
use_dynamic_shifting: bool = False,
|
||||
time_shift_type: Literal["exponential", "linear"] = "linear",
|
||||
):
|
||||
self.timestep_ratios = {}
|
||||
self.timesteps_per_stage = {}
|
||||
self.sigmas_per_stage = {}
|
||||
self.start_sigmas = {}
|
||||
self.end_sigmas = {}
|
||||
self.ori_start_sigmas = {}
|
||||
|
||||
self.init_sigmas_for_each_stage()
|
||||
self.sigma_min = self.sigmas[-1].item()
|
||||
self.sigma_max = self.sigmas[0].item()
|
||||
self.gamma = gamma
|
||||
|
||||
if solver_type not in ["bh1", "bh2"]:
|
||||
if solver_type in ["midpoint", "heun", "logrho"]:
|
||||
self.register_to_config(solver_type="bh2")
|
||||
else:
|
||||
raise NotImplementedError(f"{solver_type} is not implemented for {self.__class__}")
|
||||
|
||||
self.predict_x0 = predict_x0
|
||||
self.model_outputs = [None] * solver_order
|
||||
self.timestep_list = [None] * solver_order
|
||||
self.lower_order_nums = 0
|
||||
self.disable_corrector = disable_corrector
|
||||
self.solver_p = solver_p
|
||||
self.last_sample = None
|
||||
self._step_index = None
|
||||
self._begin_index = None
|
||||
|
||||
def init_sigmas(self):
|
||||
num_train_timesteps = self.config.num_train_timesteps
|
||||
shift = self.config.shift
|
||||
|
||||
alphas = np.linspace(1, 1 / num_train_timesteps, num_train_timesteps + 1)
|
||||
sigmas = 1.0 - alphas
|
||||
sigmas = np.flip(shift * sigmas / (1 + (shift - 1) * sigmas))[:-1].copy()
|
||||
sigmas = torch.from_numpy(sigmas)
|
||||
timesteps = (sigmas * num_train_timesteps).clone()
|
||||
|
||||
self._step_index = None
|
||||
self._begin_index = None
|
||||
self.timesteps = timesteps
|
||||
self.sigmas = sigmas.to("cpu")
|
||||
|
||||
def init_sigmas_for_each_stage(self):
|
||||
self.init_sigmas()
|
||||
|
||||
stage_distance = []
|
||||
stages = self.config.stages
|
||||
training_steps = self.config.num_train_timesteps
|
||||
stage_range = self.config.stage_range
|
||||
|
||||
for i_s in range(stages):
|
||||
start_indice = int(stage_range[i_s] * training_steps)
|
||||
start_indice = max(start_indice, 0)
|
||||
end_indice = int(stage_range[i_s + 1] * training_steps)
|
||||
end_indice = min(end_indice, training_steps)
|
||||
start_sigma = self.sigmas[start_indice].item()
|
||||
end_sigma = self.sigmas[end_indice].item() if end_indice < training_steps else 0.0
|
||||
self.ori_start_sigmas[i_s] = start_sigma
|
||||
|
||||
if i_s != 0:
|
||||
ori_sigma = 1 - start_sigma
|
||||
gamma = self.config.gamma
|
||||
corrected_sigma = (1 / (math.sqrt(1 + (1 / gamma)) * (1 - ori_sigma) + ori_sigma)) * ori_sigma
|
||||
start_sigma = 1 - corrected_sigma
|
||||
|
||||
stage_distance.append(start_sigma - end_sigma)
|
||||
self.start_sigmas[i_s] = start_sigma
|
||||
self.end_sigmas[i_s] = end_sigma
|
||||
|
||||
tot_distance = sum(stage_distance)
|
||||
for i_s in range(stages):
|
||||
if i_s == 0:
|
||||
start_ratio = 0.0
|
||||
else:
|
||||
start_ratio = sum(stage_distance[:i_s]) / tot_distance
|
||||
if i_s == stages - 1:
|
||||
end_ratio = 0.9999999999999999
|
||||
else:
|
||||
end_ratio = sum(stage_distance[: i_s + 1]) / tot_distance
|
||||
|
||||
self.timestep_ratios[i_s] = (start_ratio, end_ratio)
|
||||
|
||||
for i_s in range(stages):
|
||||
timestep_ratio = self.timestep_ratios[i_s]
|
||||
timestep_max = min(self.timesteps[int(timestep_ratio[0] * training_steps)], 999)
|
||||
timestep_min = self.timesteps[min(int(timestep_ratio[1] * training_steps), training_steps - 1)]
|
||||
timesteps = np.linspace(timestep_max, timestep_min, training_steps + 1)
|
||||
self.timesteps_per_stage[i_s] = (
|
||||
timesteps[:-1] if isinstance(timesteps, torch.Tensor) else torch.from_numpy(timesteps[:-1])
|
||||
)
|
||||
stage_sigmas = np.linspace(0.999, 0, training_steps + 1)
|
||||
self.sigmas_per_stage[i_s] = torch.from_numpy(stage_sigmas[:-1])
|
||||
|
||||
@property
|
||||
def step_index(self):
|
||||
return self._step_index
|
||||
|
||||
@property
|
||||
def begin_index(self):
|
||||
return self._begin_index
|
||||
|
||||
def set_begin_index(self, begin_index: int = 0):
|
||||
self._begin_index = begin_index
|
||||
|
||||
def _sigma_to_t(self, sigma):
|
||||
return sigma * self.config.num_train_timesteps
|
||||
|
||||
def set_timesteps(
|
||||
self,
|
||||
num_inference_steps: int,
|
||||
stage_index: int | None = None,
|
||||
device: str | torch.device = None,
|
||||
sigmas: bool | None = None,
|
||||
mu: bool | None = None,
|
||||
is_amplify_first_chunk: bool = False,
|
||||
):
|
||||
if self.config.scheduler_type == "dmd":
|
||||
if is_amplify_first_chunk:
|
||||
num_inference_steps = num_inference_steps * 2 + 1
|
||||
else:
|
||||
num_inference_steps = num_inference_steps + 1
|
||||
|
||||
self.num_inference_steps = num_inference_steps
|
||||
self.init_sigmas()
|
||||
|
||||
if self.config.stages == 1 or stage_index is None:
|
||||
if sigmas is None:
|
||||
sigmas = np.linspace(1, 1 / self.config.num_train_timesteps, num_inference_steps + 1)[:-1].astype(
|
||||
np.float32
|
||||
)
|
||||
if self.config.shift != 1.0:
|
||||
assert not self.config.use_dynamic_shifting
|
||||
sigmas = self.time_shift(self.config.shift, 1.0, sigmas)
|
||||
timesteps = (sigmas * self.config.num_train_timesteps).copy()
|
||||
sigmas = torch.from_numpy(sigmas)
|
||||
else:
|
||||
stage_timesteps = self.timesteps_per_stage[stage_index]
|
||||
timesteps = np.linspace(
|
||||
stage_timesteps[0].item(),
|
||||
stage_timesteps[-1].item(),
|
||||
num_inference_steps,
|
||||
)
|
||||
|
||||
stage_sigmas = self.sigmas_per_stage[stage_index]
|
||||
ratios = np.linspace(stage_sigmas[0].item(), stage_sigmas[-1].item(), num_inference_steps)
|
||||
sigmas = torch.from_numpy(ratios)
|
||||
|
||||
self.timesteps = torch.from_numpy(timesteps).to(device=device)
|
||||
self.sigmas = torch.cat([sigmas, torch.zeros(1)]).to(device=device)
|
||||
|
||||
self._step_index = None
|
||||
self.reset_scheduler_history()
|
||||
|
||||
if self.config.scheduler_type == "dmd":
|
||||
self.timesteps = self.timesteps[:-1]
|
||||
self.sigmas = torch.cat([self.sigmas[:-2], self.sigmas[-1:]])
|
||||
|
||||
if self.config.use_dynamic_shifting:
|
||||
assert self.config.shift == 1.0
|
||||
self.sigmas = self.time_shift(mu, 1.0, self.sigmas)
|
||||
if self.config.stages == 1 or stage_index is None:
|
||||
self.timesteps = self.sigmas[:-1] * self.config.num_train_timesteps
|
||||
else:
|
||||
self.timesteps = self.timesteps_per_stage[stage_index].min() + self.sigmas[:-1] * (
|
||||
self.timesteps_per_stage[stage_index].max() - self.timesteps_per_stage[stage_index].min()
|
||||
)
|
||||
|
||||
def time_shift(self, mu: float, sigma: float, t: torch.Tensor):
|
||||
if self.config.time_shift_type == "exponential":
|
||||
return self._time_shift_exponential(mu, sigma, t)
|
||||
elif self.config.time_shift_type == "linear":
|
||||
return self._time_shift_linear(mu, sigma, t)
|
||||
|
||||
def _time_shift_exponential(self, mu, sigma, t):
|
||||
return math.exp(mu) / (math.exp(mu) + (1 / t - 1) ** sigma)
|
||||
|
||||
def _time_shift_linear(self, mu, sigma, t):
|
||||
return mu / (mu + (1 / t - 1) ** sigma)
|
||||
|
||||
# ---------------------------------- Euler ----------------------------------
|
||||
def index_for_timestep(self, timestep, schedule_timesteps=None):
|
||||
if schedule_timesteps is None:
|
||||
schedule_timesteps = self.timesteps
|
||||
|
||||
indices = (schedule_timesteps == timestep).nonzero()
|
||||
pos = 1 if len(indices) > 1 else 0
|
||||
return indices[pos].item()
|
||||
|
||||
def _init_step_index(self, timestep):
|
||||
if self.begin_index is None:
|
||||
if isinstance(timestep, torch.Tensor):
|
||||
timestep = timestep.to(self.timesteps.device)
|
||||
self._step_index = self.index_for_timestep(timestep)
|
||||
else:
|
||||
self._step_index = self._begin_index
|
||||
|
||||
def step_euler(
|
||||
self,
|
||||
model_output: torch.FloatTensor,
|
||||
timestep: float | torch.FloatTensor = None,
|
||||
sample: torch.FloatTensor = None,
|
||||
generator: torch.Generator | None = None,
|
||||
sigma: torch.FloatTensor | None = None,
|
||||
sigma_next: torch.FloatTensor | None = None,
|
||||
return_dict: bool = True,
|
||||
) -> HeliosSchedulerOutput | tuple:
|
||||
assert (sigma is None) == (sigma_next is None)
|
||||
|
||||
if sigma is None and sigma_next is None:
|
||||
if (
|
||||
isinstance(timestep, int)
|
||||
or isinstance(timestep, torch.IntTensor)
|
||||
or isinstance(timestep, torch.LongTensor)
|
||||
):
|
||||
raise ValueError(
|
||||
"Passing integer indices as timesteps to EulerDiscreteScheduler.step() is not supported. "
|
||||
"Pass one of scheduler.timesteps as a timestep."
|
||||
)
|
||||
|
||||
if self.step_index is None:
|
||||
self._step_index = 0
|
||||
|
||||
sample = sample.to(torch.float32)
|
||||
|
||||
if sigma is None and sigma_next is None:
|
||||
sigma = self.sigmas[self.step_index]
|
||||
sigma_next = self.sigmas[self.step_index + 1]
|
||||
|
||||
prev_sample = sample + (sigma_next - sigma) * model_output
|
||||
prev_sample = prev_sample.to(model_output.dtype)
|
||||
|
||||
self._step_index += 1
|
||||
|
||||
if not return_dict:
|
||||
return (prev_sample,)
|
||||
|
||||
return HeliosSchedulerOutput(prev_sample=prev_sample)
|
||||
|
||||
# ---------------------------------- UniPC ----------------------------------
|
||||
def _sigma_to_alpha_sigma_t(self, sigma):
|
||||
if self.config.use_flow_sigmas:
|
||||
alpha_t = 1 - sigma
|
||||
sigma_t = torch.clamp(sigma, min=1e-8)
|
||||
else:
|
||||
alpha_t = 1 / ((sigma**2 + 1) ** 0.5)
|
||||
sigma_t = sigma * alpha_t
|
||||
|
||||
return alpha_t, sigma_t
|
||||
|
||||
def convert_model_output(
|
||||
self,
|
||||
model_output: torch.Tensor,
|
||||
*args,
|
||||
sample: torch.Tensor = None,
|
||||
sigma: torch.Tensor = None,
|
||||
**kwargs,
|
||||
) -> torch.Tensor:
|
||||
if sample is None:
|
||||
if len(args) > 1:
|
||||
sample = args[1]
|
||||
else:
|
||||
raise ValueError("missing `sample` as a required keyword argument")
|
||||
|
||||
flag = False
|
||||
if sigma is None:
|
||||
flag = True
|
||||
sigma = self.sigmas[self.step_index]
|
||||
alpha_t, sigma_t = self._sigma_to_alpha_sigma_t(sigma)
|
||||
|
||||
if self.predict_x0:
|
||||
if self.config.prediction_type == "epsilon":
|
||||
x0_pred = (sample - sigma_t * model_output) / alpha_t
|
||||
elif self.config.prediction_type == "sample":
|
||||
x0_pred = model_output
|
||||
elif self.config.prediction_type == "v_prediction":
|
||||
x0_pred = alpha_t * sample - sigma_t * model_output
|
||||
elif self.config.prediction_type == "flow_prediction":
|
||||
if flag:
|
||||
sigma_t = self.sigmas[self.step_index]
|
||||
else:
|
||||
sigma_t = sigma
|
||||
x0_pred = sample - sigma_t * model_output
|
||||
else:
|
||||
raise ValueError(
|
||||
f"prediction_type given as {self.config.prediction_type} must be one of `epsilon`, `sample`, "
|
||||
"`v_prediction`, or `flow_prediction`."
|
||||
)
|
||||
|
||||
if self.config.thresholding:
|
||||
x0_pred = self._threshold_sample(x0_pred)
|
||||
|
||||
return x0_pred
|
||||
else:
|
||||
if self.config.prediction_type == "epsilon":
|
||||
return model_output
|
||||
elif self.config.prediction_type == "sample":
|
||||
epsilon = (sample - alpha_t * model_output) / sigma_t
|
||||
return epsilon
|
||||
elif self.config.prediction_type == "v_prediction":
|
||||
epsilon = alpha_t * model_output + sigma_t * sample
|
||||
return epsilon
|
||||
else:
|
||||
raise ValueError(
|
||||
f"prediction_type given as {self.config.prediction_type} must be one of `epsilon`, `sample`, or"
|
||||
" `v_prediction`."
|
||||
)
|
||||
|
||||
def multistep_uni_p_bh_update(
|
||||
self,
|
||||
model_output: torch.Tensor,
|
||||
*args,
|
||||
sample: torch.Tensor = None,
|
||||
order: int = None,
|
||||
sigma: torch.Tensor = None,
|
||||
sigma_next: torch.Tensor = None,
|
||||
**kwargs,
|
||||
) -> torch.Tensor:
|
||||
if sample is None:
|
||||
if len(args) > 1:
|
||||
sample = args[1]
|
||||
else:
|
||||
raise ValueError("missing `sample` as a required keyword argument")
|
||||
if order is None:
|
||||
if len(args) > 2:
|
||||
order = args[2]
|
||||
else:
|
||||
raise ValueError("missing `order` as a required keyword argument")
|
||||
|
||||
model_output_list = self.model_outputs
|
||||
|
||||
s0 = self.timestep_list[-1]
|
||||
m0 = model_output_list[-1]
|
||||
x = sample
|
||||
|
||||
if self.solver_p:
|
||||
x_t = self.solver_p.step(model_output, s0, x).prev_sample
|
||||
return x_t
|
||||
|
||||
if sigma_next is None and sigma is None:
|
||||
sigma_t, sigma_s0 = self.sigmas[self.step_index + 1], self.sigmas[self.step_index]
|
||||
else:
|
||||
sigma_t, sigma_s0 = sigma_next, sigma
|
||||
alpha_t, sigma_t = self._sigma_to_alpha_sigma_t(sigma_t)
|
||||
alpha_s0, sigma_s0 = self._sigma_to_alpha_sigma_t(sigma_s0)
|
||||
|
||||
lambda_t = torch.log(alpha_t) - torch.log(sigma_t)
|
||||
lambda_s0 = torch.log(alpha_s0) - torch.log(sigma_s0)
|
||||
|
||||
h = lambda_t - lambda_s0
|
||||
device = sample.device
|
||||
|
||||
rks = []
|
||||
D1s = []
|
||||
for i in range(1, order):
|
||||
si = self.step_index - i
|
||||
mi = model_output_list[-(i + 1)]
|
||||
alpha_si, sigma_si = self._sigma_to_alpha_sigma_t(self.sigmas[si])
|
||||
lambda_si = torch.log(alpha_si) - torch.log(sigma_si)
|
||||
rk = (lambda_si - lambda_s0) / h
|
||||
rks.append(rk)
|
||||
D1s.append((mi - m0) / rk)
|
||||
|
||||
rks.append(1.0)
|
||||
rks = torch.tensor(rks, device=device, dtype=torch.float32)
|
||||
|
||||
R = []
|
||||
b = []
|
||||
|
||||
hh = -h if self.predict_x0 else h
|
||||
hh = hh.float()
|
||||
h_phi_1 = torch.expm1(hh)
|
||||
h_phi_k = h_phi_1 / hh - 1
|
||||
|
||||
factorial_i = 1
|
||||
|
||||
if self.config.solver_type == "bh1":
|
||||
B_h = hh
|
||||
elif self.config.solver_type == "bh2":
|
||||
B_h = torch.expm1(hh)
|
||||
else:
|
||||
raise NotImplementedError()
|
||||
|
||||
for i in range(1, order + 1):
|
||||
R.append(torch.pow(rks, i - 1))
|
||||
b.append(h_phi_k * factorial_i / B_h)
|
||||
factorial_i *= i + 1
|
||||
h_phi_k = h_phi_k / hh - 1 / factorial_i
|
||||
|
||||
R = torch.stack(R)
|
||||
b = torch.tensor(b, device=device, dtype=torch.float32)
|
||||
|
||||
if len(D1s) > 0:
|
||||
D1s = torch.stack(D1s, dim=1)
|
||||
if order == 2:
|
||||
rhos_p = torch.tensor([0.5], dtype=x.dtype, device=device)
|
||||
else:
|
||||
rhos_p = torch.linalg.solve(R[:-1, :-1], b[:-1]).to(device).to(x.dtype)
|
||||
else:
|
||||
D1s = None
|
||||
|
||||
if self.predict_x0:
|
||||
x_t_ = sigma_t / sigma_s0 * x - alpha_t * h_phi_1 * m0
|
||||
if D1s is not None:
|
||||
pred_res = torch.einsum("k,bkc...->bc...", rhos_p, D1s)
|
||||
else:
|
||||
pred_res = 0
|
||||
x_t = x_t_ - alpha_t * B_h * pred_res
|
||||
else:
|
||||
x_t_ = alpha_t / alpha_s0 * x - sigma_t * h_phi_1 * m0
|
||||
if D1s is not None:
|
||||
pred_res = torch.einsum("k,bkc...->bc...", rhos_p, D1s)
|
||||
else:
|
||||
pred_res = 0
|
||||
x_t = x_t_ - sigma_t * B_h * pred_res
|
||||
|
||||
x_t = x_t.to(x.dtype)
|
||||
return x_t
|
||||
|
||||
def multistep_uni_c_bh_update(
|
||||
self,
|
||||
this_model_output: torch.Tensor,
|
||||
*args,
|
||||
last_sample: torch.Tensor = None,
|
||||
this_sample: torch.Tensor = None,
|
||||
order: int = None,
|
||||
sigma_before: torch.Tensor = None,
|
||||
sigma: torch.Tensor = None,
|
||||
**kwargs,
|
||||
) -> torch.Tensor:
|
||||
if last_sample is None:
|
||||
if len(args) > 1:
|
||||
last_sample = args[1]
|
||||
else:
|
||||
raise ValueError("missing `last_sample` as a required keyword argument")
|
||||
if this_sample is None:
|
||||
if len(args) > 2:
|
||||
this_sample = args[2]
|
||||
else:
|
||||
raise ValueError("missing `this_sample` as a required keyword argument")
|
||||
if order is None:
|
||||
if len(args) > 3:
|
||||
order = args[3]
|
||||
else:
|
||||
raise ValueError("missing `order` as a required keyword argument")
|
||||
|
||||
model_output_list = self.model_outputs
|
||||
|
||||
m0 = model_output_list[-1]
|
||||
x = last_sample
|
||||
x_t = this_sample
|
||||
model_t = this_model_output
|
||||
|
||||
if sigma_before is None and sigma is None:
|
||||
sigma_t, sigma_s0 = self.sigmas[self.step_index], self.sigmas[self.step_index - 1]
|
||||
else:
|
||||
sigma_t, sigma_s0 = sigma, sigma_before
|
||||
alpha_t, sigma_t = self._sigma_to_alpha_sigma_t(sigma_t)
|
||||
alpha_s0, sigma_s0 = self._sigma_to_alpha_sigma_t(sigma_s0)
|
||||
|
||||
lambda_t = torch.log(alpha_t) - torch.log(sigma_t)
|
||||
lambda_s0 = torch.log(alpha_s0) - torch.log(sigma_s0)
|
||||
|
||||
h = lambda_t - lambda_s0
|
||||
device = this_sample.device
|
||||
|
||||
rks = []
|
||||
D1s = []
|
||||
for i in range(1, order):
|
||||
si = self.step_index - (i + 1)
|
||||
mi = model_output_list[-(i + 1)]
|
||||
alpha_si, sigma_si = self._sigma_to_alpha_sigma_t(self.sigmas[si])
|
||||
lambda_si = torch.log(alpha_si) - torch.log(sigma_si)
|
||||
rk = (lambda_si - lambda_s0) / h
|
||||
rks.append(rk)
|
||||
D1s.append((mi - m0) / rk)
|
||||
|
||||
rks.append(1.0)
|
||||
rks = torch.tensor(rks, device=device, dtype=torch.float32)
|
||||
|
||||
R = []
|
||||
b = []
|
||||
|
||||
hh = -h if self.predict_x0 else h
|
||||
hh = hh.float()
|
||||
h_phi_1 = torch.expm1(hh)
|
||||
h_phi_k = h_phi_1 / hh - 1
|
||||
|
||||
factorial_i = 1
|
||||
|
||||
if self.config.solver_type == "bh1":
|
||||
B_h = hh
|
||||
elif self.config.solver_type == "bh2":
|
||||
B_h = torch.expm1(hh)
|
||||
else:
|
||||
raise NotImplementedError()
|
||||
|
||||
for i in range(1, order + 1):
|
||||
R.append(torch.pow(rks, i - 1))
|
||||
b.append(h_phi_k * factorial_i / B_h)
|
||||
factorial_i *= i + 1
|
||||
h_phi_k = h_phi_k / hh - 1 / factorial_i
|
||||
|
||||
R = torch.stack(R)
|
||||
b = torch.tensor(b, device=device, dtype=torch.float32)
|
||||
|
||||
if len(D1s) > 0:
|
||||
D1s = torch.stack(D1s, dim=1)
|
||||
else:
|
||||
D1s = None
|
||||
|
||||
if order == 1:
|
||||
rhos_c = torch.tensor([0.5], dtype=x.dtype, device=device)
|
||||
else:
|
||||
rhos_c = torch.linalg.solve(R, b).to(device).to(x.dtype)
|
||||
|
||||
if self.predict_x0:
|
||||
x_t_ = sigma_t / sigma_s0 * x - alpha_t * h_phi_1 * m0
|
||||
if D1s is not None:
|
||||
corr_res = torch.einsum("k,bkc...->bc...", rhos_c[:-1], D1s)
|
||||
else:
|
||||
corr_res = 0
|
||||
D1_t = model_t - m0
|
||||
x_t = x_t_ - alpha_t * B_h * (corr_res + rhos_c[-1] * D1_t)
|
||||
else:
|
||||
x_t_ = alpha_t / alpha_s0 * x - sigma_t * h_phi_1 * m0
|
||||
if D1s is not None:
|
||||
corr_res = torch.einsum("k,bkc...->bc...", rhos_c[:-1], D1s)
|
||||
else:
|
||||
corr_res = 0
|
||||
D1_t = model_t - m0
|
||||
x_t = x_t_ - sigma_t * B_h * (corr_res + rhos_c[-1] * D1_t)
|
||||
x_t = x_t.to(x.dtype)
|
||||
return x_t
|
||||
|
||||
def step_unipc(
|
||||
self,
|
||||
model_output: torch.Tensor,
|
||||
timestep: int | torch.Tensor = None,
|
||||
sample: torch.Tensor = None,
|
||||
return_dict: bool = True,
|
||||
model_outputs: list = None,
|
||||
timestep_list: list = None,
|
||||
sigma_before: torch.Tensor = None,
|
||||
sigma: torch.Tensor = None,
|
||||
sigma_next: torch.Tensor = None,
|
||||
cus_step_index: int = None,
|
||||
cus_lower_order_num: int = None,
|
||||
cus_this_order: int = None,
|
||||
cus_last_sample: torch.Tensor = None,
|
||||
) -> HeliosSchedulerOutput | tuple:
|
||||
if self.num_inference_steps is None:
|
||||
raise ValueError(
|
||||
"Number of inference steps is 'None', you need to run 'set_timesteps' after creating the scheduler"
|
||||
)
|
||||
|
||||
if cus_step_index is None:
|
||||
if self.step_index is None:
|
||||
self._step_index = 0
|
||||
else:
|
||||
self._step_index = cus_step_index
|
||||
|
||||
if cus_lower_order_num is not None:
|
||||
self.lower_order_nums = cus_lower_order_num
|
||||
|
||||
if cus_this_order is not None:
|
||||
self.this_order = cus_this_order
|
||||
|
||||
if cus_last_sample is not None:
|
||||
self.last_sample = cus_last_sample
|
||||
|
||||
use_corrector = (
|
||||
self.step_index > 0 and self.step_index - 1 not in self.disable_corrector and self.last_sample is not None
|
||||
)
|
||||
|
||||
model_output_convert = self.convert_model_output(model_output, sample=sample, sigma=sigma)
|
||||
|
||||
if model_outputs is not None and timestep_list is not None:
|
||||
self.model_outputs = model_outputs[:-1]
|
||||
self.timestep_list = timestep_list[:-1]
|
||||
|
||||
if use_corrector:
|
||||
sample = self.multistep_uni_c_bh_update(
|
||||
this_model_output=model_output_convert,
|
||||
last_sample=self.last_sample,
|
||||
this_sample=sample,
|
||||
order=self.this_order,
|
||||
sigma_before=sigma_before,
|
||||
sigma=sigma,
|
||||
)
|
||||
|
||||
if model_outputs is not None and timestep_list is not None:
|
||||
model_outputs[-1] = model_output_convert
|
||||
self.model_outputs = model_outputs[1:]
|
||||
self.timestep_list = timestep_list[1:]
|
||||
else:
|
||||
for i in range(self.config.solver_order - 1):
|
||||
self.model_outputs[i] = self.model_outputs[i + 1]
|
||||
self.timestep_list[i] = self.timestep_list[i + 1]
|
||||
self.model_outputs[-1] = model_output_convert
|
||||
self.timestep_list[-1] = timestep
|
||||
|
||||
if self.config.lower_order_final:
|
||||
this_order = min(self.config.solver_order, len(self.timesteps) - self.step_index)
|
||||
else:
|
||||
this_order = self.config.solver_order
|
||||
self.this_order = min(this_order, self.lower_order_nums + 1)
|
||||
assert self.this_order > 0
|
||||
|
||||
self.last_sample = sample
|
||||
prev_sample = self.multistep_uni_p_bh_update(
|
||||
model_output=model_output,
|
||||
sample=sample,
|
||||
order=self.this_order,
|
||||
sigma=sigma,
|
||||
sigma_next=sigma_next,
|
||||
)
|
||||
|
||||
if cus_lower_order_num is None:
|
||||
if self.lower_order_nums < self.config.solver_order:
|
||||
self.lower_order_nums += 1
|
||||
|
||||
if cus_step_index is None:
|
||||
self._step_index += 1
|
||||
|
||||
if not return_dict:
|
||||
return (prev_sample, model_outputs, self.last_sample, self.this_order)
|
||||
|
||||
return HeliosSchedulerOutput(
|
||||
prev_sample=prev_sample,
|
||||
model_outputs=model_outputs,
|
||||
last_sample=self.last_sample,
|
||||
this_order=self.this_order,
|
||||
)
|
||||
|
||||
# ---------------------------------- DMD ----------------------------------
|
||||
def add_noise(self, original_samples, noise, timestep, sigmas, timesteps):
|
||||
sigmas = sigmas.to(noise.device)
|
||||
timesteps = timesteps.to(noise.device)
|
||||
timestep_id = torch.argmin((timesteps.unsqueeze(0) - timestep.unsqueeze(1)).abs(), dim=1)
|
||||
sigma = sigmas[timestep_id].reshape(-1, 1, 1, 1, 1)
|
||||
sample = (1 - sigma) * original_samples + sigma * noise
|
||||
return sample.type_as(noise)
|
||||
|
||||
def convert_flow_pred_to_x0(self, flow_pred, xt, timestep, sigmas, timesteps):
|
||||
original_dtype = flow_pred.dtype
|
||||
device = flow_pred.device
|
||||
flow_pred, xt, sigmas, timesteps = (x.double().to(device) for x in (flow_pred, xt, sigmas, timesteps))
|
||||
|
||||
timestep_id = torch.argmin((timesteps.unsqueeze(0) - timestep.unsqueeze(1)).abs(), dim=1)
|
||||
sigma_t = sigmas[timestep_id].reshape(-1, 1, 1, 1, 1)
|
||||
x0_pred = xt - sigma_t * flow_pred
|
||||
return x0_pred.to(original_dtype)
|
||||
|
||||
def step_dmd(
|
||||
self,
|
||||
model_output: torch.FloatTensor,
|
||||
timestep: float | torch.FloatTensor = None,
|
||||
sample: torch.FloatTensor = None,
|
||||
generator: torch.Generator | None = None,
|
||||
return_dict: bool = True,
|
||||
cur_sampling_step: int = 0,
|
||||
dmd_noisy_tensor: torch.FloatTensor | None = None,
|
||||
dmd_sigmas: torch.FloatTensor | None = None,
|
||||
dmd_timesteps: torch.FloatTensor | None = None,
|
||||
all_timesteps: torch.FloatTensor | None = None,
|
||||
):
|
||||
pred_image_or_video = self.convert_flow_pred_to_x0(
|
||||
flow_pred=model_output,
|
||||
xt=sample,
|
||||
timestep=torch.full((model_output.shape[0],), timestep, dtype=torch.long, device=model_output.device),
|
||||
sigmas=dmd_sigmas,
|
||||
timesteps=dmd_timesteps,
|
||||
)
|
||||
if cur_sampling_step < len(all_timesteps) - 1:
|
||||
prev_sample = self.add_noise(
|
||||
pred_image_or_video,
|
||||
dmd_noisy_tensor,
|
||||
torch.full(
|
||||
(model_output.shape[0],),
|
||||
all_timesteps[cur_sampling_step + 1],
|
||||
dtype=torch.long,
|
||||
device=model_output.device,
|
||||
),
|
||||
sigmas=dmd_sigmas,
|
||||
timesteps=dmd_timesteps,
|
||||
)
|
||||
else:
|
||||
prev_sample = pred_image_or_video
|
||||
|
||||
if not return_dict:
|
||||
return (prev_sample,)
|
||||
|
||||
return HeliosSchedulerOutput(prev_sample=prev_sample)
|
||||
|
||||
# ---------------------------------- Merge ----------------------------------
|
||||
def step(
|
||||
self,
|
||||
model_output: torch.FloatTensor,
|
||||
timestep: float | torch.FloatTensor = None,
|
||||
sample: torch.FloatTensor = None,
|
||||
generator: torch.Generator | None = None,
|
||||
return_dict: bool = True,
|
||||
cur_sampling_step: int = 0,
|
||||
dmd_noisy_tensor: torch.FloatTensor | None = None,
|
||||
dmd_sigmas: torch.FloatTensor | None = None,
|
||||
dmd_timesteps: torch.FloatTensor | None = None,
|
||||
all_timesteps: torch.FloatTensor | None = None,
|
||||
) -> HeliosSchedulerOutput | tuple:
|
||||
if self.config.scheduler_type == "euler":
|
||||
return self.step_euler(
|
||||
model_output=model_output,
|
||||
timestep=timestep,
|
||||
sample=sample,
|
||||
generator=generator,
|
||||
return_dict=return_dict,
|
||||
)
|
||||
elif self.config.scheduler_type == "unipc":
|
||||
return self.step_unipc(
|
||||
model_output=model_output,
|
||||
timestep=timestep,
|
||||
sample=sample,
|
||||
return_dict=return_dict,
|
||||
)
|
||||
elif self.config.scheduler_type == "dmd":
|
||||
return self.step_dmd(
|
||||
model_output=model_output,
|
||||
timestep=timestep,
|
||||
sample=sample,
|
||||
generator=generator,
|
||||
return_dict=return_dict,
|
||||
cur_sampling_step=cur_sampling_step,
|
||||
dmd_noisy_tensor=dmd_noisy_tensor,
|
||||
dmd_sigmas=dmd_sigmas,
|
||||
dmd_timesteps=dmd_timesteps,
|
||||
all_timesteps=all_timesteps,
|
||||
)
|
||||
else:
|
||||
raise NotImplementedError
|
||||
|
||||
def reset_scheduler_history(self):
|
||||
self.model_outputs = [None] * self.config.solver_order
|
||||
self.timestep_list = [None] * self.config.solver_order
|
||||
self.lower_order_nums = 0
|
||||
self.disable_corrector = self.config.disable_corrector
|
||||
self.solver_p = self.config.solver_p
|
||||
self.last_sample = None
|
||||
self._step_index = None
|
||||
self._begin_index = None
|
||||
|
||||
def __len__(self):
|
||||
return self.config.num_train_timesteps
|
||||
Reference in new issue
Block a user