nodes/hao-ai-lab-FastVideo
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

lint

Browse Source
This commit is contained in:
Will Lin
2025-05-08 16:30:53 +08:00
committed by Jinzhe Pan
parent 93c30ce848
commit fd80ccbf88
16 changed files with 280 additions and 224 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.pre-commit-config.yaml
+1
View File
@@ -20,6 +20,7 @@ exclude: |
fastvideo/train\.py|
fastvideo/utils/.*|
examples/.*|
fastvideo/v1/models/schedulers/scheduling_flow_match_euler_discrete.py|
.github/workflows/fastvideo-publish.yml|
.github/workflows/sta-publish.yml|
.github/workflows/build-image-template.yml
fastvideo/v1/configs/models/base.py
+4 -1
View File
@@ -43,7 +43,10 @@ class ModelConfig:
# This should be used only when loading from transformers/diffusers
def update_model_arch(self, source_model_dict: Dict[str, Any]) -> None:
# Remove all keys that start with "_"
keys_to_remove = [key for key in list(source_model_dict.keys()) if str(key).startswith("_")]
keys_to_remove = [
key for key in list(source_model_dict.keys())
if str(key).startswith("_")
]
for key in keys_to_remove:
source_model_dict.pop(key)
fastvideo/v1/configs/models/dits/__init__.py
+1 -1
View File
@@ -1,5 +1,5 @@
from fastvideo.v1.configs.models.dits.flux import FluxImageConfig
from fastvideo.v1.configs.models.dits.hunyuanvideo import HunyuanVideoConfig
from fastvideo.v1.configs.models.dits.wanvideo import WanVideoConfig
from fastvideo.v1.configs.models.dits.flux import FluxImageConfig
__all__ = ["HunyuanVideoConfig", "WanVideoConfig", "FluxImageConfig"]
fastvideo/v1/configs/models/dits/flux.py
+2 -1
View File
@@ -2,6 +2,7 @@ from dataclasses import dataclass, field
from typing import Optional, Tuple
import torch
from fastvideo.v1.configs.models.dits.base import DiTArchConfig, DiTConfig
@@ -112,7 +113,7 @@ class FluxImageArchConfig(DiTArchConfig):
joint_attention_dim: int = 4096
pooled_projection_dim: int = 768
guidance_embeds: bool = False
axes_dims_rope: Tuple[int] = (16, 56, 56)
axes_dims_rope: Tuple[int, ...] = (16, 56, 56)
rope_theta: int = 10000
dtype: Optional[torch.dtype] = torch.bfloat16
fastvideo/v1/configs/models/vaes/__init__.py
+1 -1
View File
@@ -1,6 +1,6 @@
from fastvideo.v1.configs.models.vaes.hunyuanvae import HunyuanVAEConfig
from fastvideo.v1.configs.models.vaes.wanvae import WanVAEConfig
from fastvideo.v1.configs.models.vaes.image_vae import ImageVAEConfig
from fastvideo.v1.configs.models.vaes.wanvae import WanVAEConfig
__all__ = [
"HunyuanVAEConfig",
fastvideo/v1/configs/models/vaes/image_vae.py
+21 -20
View File
@@ -1,35 +1,36 @@
from dataclasses import dataclass, field
from typing import Tuple, Optional
from typing import Optional, Tuple
from fastvideo.v1.configs.models.vaes.base import VAEArchConfig, VAEConfig
@dataclass
class ImageVAEArchConfig(VAEArchConfig):
in_channels: int = 3,
out_channels: int = 3,
down_block_types: Tuple[str] = ("DownEncoderBlock2D",),
up_block_types: Tuple[str] = ("UpDecoderBlock2D",),
block_out_channels: Tuple[int] = (64,),
layers_per_block: int = 1,
act_fn: str = "silu",
latent_channels: int = 4,
norm_num_groups: int = 32,
sample_size: int = 32,
scaling_factor: float = 0.18215,
shift_factor: Optional[float] = None,
latents_mean: Optional[Tuple[float]] = None,
latents_std: Optional[Tuple[float]] = None,
force_upcast: float = True,
use_quant_conv: bool = True,
use_post_quant_conv: bool = True,
mid_block_add_attention: bool = True,
in_channels: int = 3
out_channels: int = 3
down_block_types: Tuple[str] = ("DownEncoderBlock2D", )
up_block_types: Tuple[str] = ("UpDecoderBlock2D", )
block_out_channels: Tuple[int] = (64, )
layers_per_block: int = 1
act_fn: str = "silu"
latent_channels: int = 4
norm_num_groups: int = 32
sample_size: int = 32
scaling_factor: float = 0.18215
shift_factor: Optional[float] = None
latents_mean: Optional[Tuple[float]] = None
latents_std: Optional[Tuple[float]] = None
force_upcast: float = True
use_quant_conv: bool = True
use_post_quant_conv: bool = True
mid_block_add_attention: bool = True
def __post_init__(self):
self.spatial_compression_ratio: int = 2**(len(self.block_out_channels) -
1)
self.temporal_compression_ratio = 1
@dataclass
class ImageVAEConfig(VAEConfig):
arch_config: VAEArchConfig = field(default_factory=ImageVAEArchConfig)
@@ -37,4 +38,4 @@ class ImageVAEConfig(VAEConfig):
# overrides VAEConfig
use_tiling: bool = False
use_temporal_tiling: bool = False
use_parallel_tiling: bool = False
use_parallel_tiling: bool = False
fastvideo/v1/configs/pipelines/flux.py
+11 -9
View File
@@ -10,24 +10,26 @@ from fastvideo.v1.configs.models.encoders import (BaseEncoderOutput,
from fastvideo.v1.configs.models.vaes import ImageVAEConfig
from fastvideo.v1.configs.pipelines.base import PipelineConfig
def t5_preprocess_text(prompt: str) -> str:
return prompt
def t5_postprocess_text(outputs: BaseEncoderOutput) -> torch.tensor:
hidden_state: torch.tensor = outputs.last_hidden_state
assert torch.isnan(hidden_state).sum() == 0
prompt_embeds_tensor: torch.tensor = torch.stack(
[
torch.cat([u, u.new_zeros(512 - u.size(0), u.size(1))])
for u in hidden_state
],
dim=0
)
prompt_embeds_tensor: torch.tensor = torch.stack([
torch.cat([u, u.new_zeros(512 - u.size(0), u.size(1))])
for u in hidden_state
],
dim=0)
return prompt_embeds_tensor
def clip_preprocess_text(prompt: str) -> str:
return prompt
def clip_postprocess_text(outputs: BaseEncoderOutput) -> torch.tensor:
pooler_output: torch.tensor = outputs.pooler_output
return pooler_output
@@ -43,7 +45,7 @@ class FluxConfig(PipelineConfig):
# VAE
vae_config: VAEConfig = field(default_factory=ImageVAEConfig)
# Denoising stage
embedded_cfg_scale: int = 3.5
embedded_cfg_scale: float = 3.5
# Text encoding stage
text_encoder_configs: Tuple[EncoderConfig, ...] = field(
@@ -63,4 +65,4 @@ class FluxConfig(PipelineConfig):
def __post_init__(self):
self.vae_config.load_encoder = False
self.vae_config.load_decoder = True
self.vae_config.load_decoder = True
fastvideo/v1/configs/pipelines/registry.py
+1 -1
View File
@@ -4,13 +4,13 @@ import os
from typing import Callable, Dict, Optional, Type
from fastvideo.v1.configs.pipelines.base import PipelineConfig
from fastvideo.v1.configs.pipelines.flux import FluxConfig
from fastvideo.v1.configs.pipelines.hunyuan import (FastHunyuanConfig,
HunyuanConfig)
from fastvideo.v1.configs.pipelines.wan import (WanI2V480PConfig,
WanI2V720PConfig,
WanT2V480PConfig,
WanT2V720PConfig)
from fastvideo.v1.configs.pipelines.flux import FluxConfig
from fastvideo.v1.logger import init_logger
from fastvideo.v1.utils import (maybe_download_model_index,
verify_model_config_and_directory)
fastvideo/v1/models/dits/flux.py
+32 -34
View File
@@ -1,24 +1,23 @@
from typing import Optional, Tuple, List, Union
from typing import List, Optional, Tuple, Union
import torch
import torch.nn as nn
from diffusers.models.normalization import RMSNorm
from fastvideo.v1.attention import DistributedAttention
from fastvideo.v1.configs.models.dits import FluxImageConfig
from fastvideo.v1.models.dits.base import BaseDiT
from fastvideo.v1.attention import DistributedAttention, LocalAttention
from fastvideo.v1.layers.layernorm import (LayerNormScaleShift, ScaleResidual,
ScaleResidualLayerNormScaleShift)
from fastvideo.v1.layers.linear import ReplicatedLinear
from fastvideo.v1.layers.visual_embedding import (ModulateProjection,
PatchEmbed, TimestepEmbedder,
unpatchify)
from fastvideo.v1.layers.mlp import MLP
from fastvideo.v1.layers.rotary_embedding import (_apply_rotary_emb,
get_rotary_pos_embed)
from fastvideo.v1.layers.layernorm import (LayerNormScaleShift, ScaleResidual,
ScaleResidualLayerNormScaleShift)
from fastvideo.v1.layers.visual_embedding import (ModulateProjection,
TimestepEmbedder)
from fastvideo.v1.models.dits.base import BaseDiT
from fastvideo.v1.platforms import _Backend
from fastvideo.v1.distributed.parallel_state import get_sequence_model_parallel_world_size
class MMDoubleStreamBlock(nn.Module):
"""
A multimodal DiT block with separate modulation for text and image/video,
@@ -225,6 +224,7 @@ class MMDoubleStreamBlock(nn.Module):
return img, txt
class MMSingleStreamBlock(nn.Module):
"""
A DiT block with parallel linear layers using distributed attention
@@ -354,6 +354,7 @@ class MMSingleStreamBlock(nn.Module):
# Apply residual connection with gating using fused operation
return self.output_residual(x, output, mod_gate)
class FinalLayer(nn.Module):
"""
The final layer of DiT that projects features to pixel space.
@@ -373,7 +374,7 @@ class FinalLayer(nn.Module):
elementwise_affine=False,
dtype=dtype)
output_dim = patch_size ** 3 * out_channels
output_dim = patch_size**3 * out_channels
self.linear = ReplicatedLinear(hidden_size,
output_dim,
@@ -391,15 +392,18 @@ class FinalLayer(nn.Module):
def forward(self, img, vec):
scale, shift = self.adaLN_modulation(vec).chunk(2, dim=-1)
img = self.norm_final(img) * (1.0 + scale.unsqueeze(1)) + shift.unsqueeze(1)
img = self.norm_final(img) * (1.0 +
scale.unsqueeze(1)) + shift.unsqueeze(1)
img, _ = self.linear(img)
return img
class FluxTransformer2DModel(BaseDiT):
_fsdp_shard_conditions = FluxImageConfig()._fsdp_shard_conditions
_supported_attention_backends = FluxImageConfig()._supported_attention_backends
_supported_attention_backends = FluxImageConfig(
)._supported_attention_backends
_param_names_mapping = FluxImageConfig()._param_names_mapping
def __init__(self, config: FluxImageConfig) -> None:
super().__init__(config=config)
@@ -426,21 +430,18 @@ class FluxTransformer2DModel(BaseDiT):
dtype=config.dtype,
prefix=f"{config.prefix}.time_in")
self.txt2_in = MLP(self.text_states_dim_2,
self.hidden_size,
self.hidden_size,
act_type="silu",
dtype=config.dtype,
prefix=f"{config.prefix}.txt2_in")
self.guidance_in = (
TimestepEmbedder(
self.hidden_size,
act_layer="silu",
dtype=config.dtype,
prefix=f"{config.prefix}.guidance_in"
)
if config.guidance_embeds else None
)
self.hidden_size,
self.hidden_size,
act_type="silu",
dtype=config.dtype,
prefix=f"{config.prefix}.txt2_in")
self.guidance_in = (TimestepEmbedder(
self.hidden_size,
act_layer="silu",
dtype=config.dtype,
prefix=f"{config.prefix}.guidance_in")
if config.guidance_embeds else None)
# Double blocks
self.double_blocks = nn.ModuleList([
MMDoubleStreamBlock(
@@ -506,10 +507,7 @@ class FluxTransformer2DModel(BaseDiT):
# Get spatial dimensions
# _, _, oh, ow = img.shape
th, tw = (
h // self.patch_size // 2,
w // self.patch_size // 2
)
th, tw = (h // self.patch_size // 2, w // self.patch_size // 2)
# Get rotary embeddings
freqs_cos_img, freqs_sin_img = get_rotary_pos_embed(
@@ -555,7 +553,7 @@ class FluxTransformer2DModel(BaseDiT):
# Extract image features
img = x[:, :img_seq_len, ...]
# Final layer
img = self.final_layer(img, vec)
fastvideo/v1/models/schedulers/scheduling_flow_match_euler_discrete.py
+78 -63
View File
@@ -19,16 +19,16 @@
#
# ==============================================================================
from dataclasses import dataclass
import math
from typing import Optional, Tuple, Union, List
from dataclasses import dataclass
from typing import List, Optional, Tuple, Union
import numpy as np
import scipy
import torch
from diffusers.configuration_utils import ConfigMixin, register_to_config
from diffusers.schedulers.scheduling_utils import SchedulerMixin
from diffusers.utils import BaseOutput, logging, is_scipy_available
from diffusers.utils import BaseOutput, is_scipy_available, logging
from fastvideo.v1.models.schedulers.base import BaseScheduler
@@ -49,7 +49,8 @@ class FlowMatchEulerDiscreteSchedulerOutput(BaseOutput):
prev_sample: torch.FloatTensor
class FlowMatchEulerDiscreteScheduler(SchedulerMixin, ConfigMixin, BaseScheduler):
class FlowMatchEulerDiscreteScheduler(SchedulerMixin, ConfigMixin,
BaseScheduler):
"""
Euler scheduler.
@@ -111,15 +112,23 @@ class FlowMatchEulerDiscreteScheduler(SchedulerMixin, ConfigMixin, BaseScheduler
stochastic_sampling: bool = False,
):
if self.config.use_beta_sigmas and not is_scipy_available():
raise ImportError("Make sure to install scipy if you want to use beta sigmas.")
if sum([self.config.use_beta_sigmas, self.config.use_exponential_sigmas, self.config.use_karras_sigmas]) > 1:
raise ImportError(
"Make sure to install scipy if you want to use beta sigmas.")
if sum([
self.config.use_beta_sigmas, self.config.use_exponential_sigmas,
self.config.use_karras_sigmas
]) > 1:
raise ValueError(
"Only one of `config.use_beta_sigmas`, `config.use_exponential_sigmas`, `config.use_karras_sigmas` can be used."
)
if time_shift_type not in {"exponential", "linear"}:
raise ValueError("`time_shift_type` must either be 'exponential' or 'linear'.")
raise ValueError(
"`time_shift_type` must either be 'exponential' or 'linear'.")
timesteps = np.linspace(1, num_train_timesteps, num_train_timesteps, dtype=np.float32)[::-1].copy()
timesteps = np.linspace(1,
num_train_timesteps,
num_train_timesteps,
dtype=np.float32)[::-1].copy()
timesteps = torch.from_numpy(timesteps).to(dtype=torch.float32)
sigmas = timesteps / num_train_timesteps
@@ -134,7 +143,8 @@ class FlowMatchEulerDiscreteScheduler(SchedulerMixin, ConfigMixin, BaseScheduler
self._shift = shift
self.sigmas = sigmas.to("cpu") # to avoid too much CPU/GPU communication
self.sigmas = sigmas.to(
"cpu") # to avoid too much CPU/GPU communication
self.sigma_min = self.sigmas[-1].item()
self.sigma_max = self.sigmas[0].item()
@@ -197,7 +207,8 @@ class FlowMatchEulerDiscreteScheduler(SchedulerMixin, ConfigMixin, BaseScheduler
if sample.device.type == "mps" and torch.is_floating_point(timestep):
# mps does not support float64
schedule_timesteps = self.timesteps.to(sample.device, dtype=torch.float32)
schedule_timesteps = self.timesteps.to(sample.device,
dtype=torch.float32)
timestep = timestep.to(sample.device, dtype=torch.float32)
else:
schedule_timesteps = self.timesteps.to(sample.device)
@@ -205,7 +216,9 @@ class FlowMatchEulerDiscreteScheduler(SchedulerMixin, ConfigMixin, BaseScheduler
# self.begin_index is None when scheduler is used for training, or pipeline does not implement set_begin_index
if self.begin_index is None:
step_indices = [self.index_for_timestep(t, schedule_timesteps) for t in timestep]
step_indices = [
self.index_for_timestep(t, schedule_timesteps) for t in timestep
]
elif self.step_index is not None:
# add_noise is called after first denoising step (for inpainting)
step_indices = [self.step_index] * timestep.shape[0]
@@ -279,21 +292,25 @@ class FlowMatchEulerDiscreteScheduler(SchedulerMixin, ConfigMixin, BaseScheduler
automatically.
"""
if self.config.use_dynamic_shifting and mu is None:
raise ValueError("`mu` must be passed when `use_dynamic_shifting` is set to be `True`")
raise ValueError(
"`mu` must be passed when `use_dynamic_shifting` is set to be `True`"
)
if sigmas is not None and timesteps is not None:
if len(sigmas) != len(timesteps):
raise ValueError("`sigmas` and `timesteps` should have the same length")
if sigmas is not None and timesteps is not None and len(sigmas) != len(
timesteps):
raise ValueError(
"`sigmas` and `timesteps` should have the same length")
if num_inference_steps is not None:
if (sigmas is not None and len(sigmas) != num_inference_steps) or (
timesteps is not None and len(timesteps) != num_inference_steps
):
timesteps is not None
and len(timesteps) != num_inference_steps):
raise ValueError(
"`sigmas` and `timesteps` should have the same length as num_inference_steps, if `num_inference_steps` is provided"
)
else:
num_inference_steps = len(sigmas) if sigmas is not None else len(timesteps)
num_inference_steps = len(sigmas) if sigmas is not None else len(
timesteps)
self.num_inference_steps = num_inference_steps
@@ -305,9 +322,9 @@ class FlowMatchEulerDiscreteScheduler(SchedulerMixin, ConfigMixin, BaseScheduler
if sigmas is None:
if timesteps is None:
timesteps = np.linspace(
self._sigma_to_t(self.sigma_max), self._sigma_to_t(self.sigma_min), num_inference_steps
)
timesteps = np.linspace(self._sigma_to_t(self.sigma_max),
self._sigma_to_t(self.sigma_min),
num_inference_steps)
sigmas = timesteps / self.config.num_train_timesteps
else:
sigmas = np.array(sigmas).astype(np.float32)
@@ -326,18 +343,22 @@ class FlowMatchEulerDiscreteScheduler(SchedulerMixin, ConfigMixin, BaseScheduler
# 4. If required, convert sigmas to one of karras, exponential, or beta sigma schedules
if self.config.use_karras_sigmas:
sigmas = self._convert_to_karras(in_sigmas=sigmas, num_inference_steps=num_inference_steps)
sigmas = self._convert_to_karras(
in_sigmas=sigmas, num_inference_steps=num_inference_steps)
elif self.config.use_exponential_sigmas:
sigmas = self._convert_to_exponential(in_sigmas=sigmas, num_inference_steps=num_inference_steps)
sigmas = self._convert_to_exponential(
in_sigmas=sigmas, num_inference_steps=num_inference_steps)
elif self.config.use_beta_sigmas:
sigmas = self._convert_to_beta(in_sigmas=sigmas, num_inference_steps=num_inference_steps)
sigmas = self._convert_to_beta(
in_sigmas=sigmas, num_inference_steps=num_inference_steps)
# 5. Convert sigmas and timesteps to tensors and move to specified device
sigmas = torch.from_numpy(sigmas).to(dtype=torch.float32, device=device)
if not is_timesteps_provided:
timesteps = sigmas * self.config.num_train_timesteps
else:
timesteps = torch.from_numpy(timesteps).to(dtype=torch.float32, device=device)
timesteps = torch.from_numpy(timesteps).to(dtype=torch.float32,
device=device)
# 6. Append the terminal sigma value.
# If a model requires inverted sigma schedule for denoising but timesteps without inversion, the
@@ -354,11 +375,9 @@ class FlowMatchEulerDiscreteScheduler(SchedulerMixin, ConfigMixin, BaseScheduler
self._step_index = None
self._begin_index = None
def scale_model_input(
self,
sample: torch.Tensor,
timestep: Optional[int] = None
) -> torch.Tensor:
def scale_model_input(self,
sample: torch.Tensor,
timestep: Optional[int] = None) -> torch.Tensor:
return sample
def index_for_timestep(self, timestep, schedule_timesteps=None):
@@ -427,18 +446,11 @@ class FlowMatchEulerDiscreteScheduler(SchedulerMixin, ConfigMixin, BaseScheduler
otherwise a tuple is returned where the first element is the sample tensor.
"""
if (
isinstance(timestep, int)
or isinstance(timestep, torch.IntTensor)
or isinstance(timestep, torch.LongTensor)
):
raise ValueError(
(
"Passing integer indices (e.g. from `enumerate(timesteps)`) as timesteps to"
" `FlowMatchEulerDiscreteScheduler.step()` is not supported. Make sure to pass"
" one of the `scheduler.timesteps` as a timestep."
),
)
if isinstance(timestep, (int, torch.IntTensor, torch.LongTensor)):
raise ValueError((
"Passing integer indices (e.g. from `enumerate(timesteps)`) as timesteps to"
" `FlowMatchEulerDiscreteScheduler.step()` is not supported. Make sure to pass"
" one of the `scheduler.timesteps` as a timestep."), )
if self.step_index is None:
self._init_step_index(timestep)
@@ -480,12 +492,13 @@ class FlowMatchEulerDiscreteScheduler(SchedulerMixin, ConfigMixin, BaseScheduler
prev_sample = prev_sample.to(model_output.dtype)
if not return_dict:
return (prev_sample,)
return (prev_sample, )
return FlowMatchEulerDiscreteSchedulerOutput(prev_sample=prev_sample)
# Copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler._convert_to_karras
def _convert_to_karras(self, in_sigmas: torch.Tensor, num_inference_steps) -> torch.Tensor:
def _convert_to_karras(self, in_sigmas: torch.Tensor,
num_inference_steps) -> torch.Tensor:
"""Constructs the noise schedule of Karras et al. (2022)."""
# Hack to make sure that other schedulers which copy this function don't break
@@ -505,13 +518,14 @@ class FlowMatchEulerDiscreteScheduler(SchedulerMixin, ConfigMixin, BaseScheduler
rho = 7.0 # 7.0 is the value used in the paper
ramp = np.linspace(0, 1, num_inference_steps)
min_inv_rho = sigma_min ** (1 / rho)
max_inv_rho = sigma_max ** (1 / rho)
sigmas = (max_inv_rho + ramp * (min_inv_rho - max_inv_rho)) ** rho
min_inv_rho = sigma_min**(1 / rho)
max_inv_rho = sigma_max**(1 / rho)
sigmas = (max_inv_rho + ramp * (min_inv_rho - max_inv_rho))**rho
return sigmas
# Copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler._convert_to_exponential
def _convert_to_exponential(self, in_sigmas: torch.Tensor, num_inference_steps: int) -> torch.Tensor:
def _convert_to_exponential(self, in_sigmas: torch.Tensor,
num_inference_steps: int) -> torch.Tensor:
"""Constructs an exponential noise schedule."""
# Hack to make sure that other schedulers which copy this function don't break
@@ -529,13 +543,17 @@ class FlowMatchEulerDiscreteScheduler(SchedulerMixin, ConfigMixin, BaseScheduler
sigma_min = sigma_min if sigma_min is not None else in_sigmas[-1].item()
sigma_max = sigma_max if sigma_max is not None else in_sigmas[0].item()
sigmas = np.exp(np.linspace(math.log(sigma_max), math.log(sigma_min), num_inference_steps))
sigmas = np.exp(
np.linspace(math.log(sigma_max), math.log(sigma_min),
num_inference_steps))
return sigmas
# Copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler._convert_to_beta
def _convert_to_beta(
self, in_sigmas: torch.Tensor, num_inference_steps: int, alpha: float = 0.6, beta: float = 0.6
) -> torch.Tensor:
def _convert_to_beta(self,
in_sigmas: torch.Tensor,
num_inference_steps: int,
alpha: float = 0.6,
beta: float = 0.6) -> torch.Tensor:
"""From "Beta Sampling is All You Need" [arXiv:2407.12173] (Lee et. al, 2024)"""
# Hack to make sure that other schedulers which copy this function don't break
@@ -553,22 +571,19 @@ class FlowMatchEulerDiscreteScheduler(SchedulerMixin, ConfigMixin, BaseScheduler
sigma_min = sigma_min if sigma_min is not None else in_sigmas[-1].item()
sigma_max = sigma_max if sigma_max is not None else in_sigmas[0].item()
sigmas = np.array(
[
sigma_min + (ppf * (sigma_max - sigma_min))
for ppf in [
scipy.stats.beta.ppf(timestep, alpha, beta)
for timestep in 1 - np.linspace(0, 1, num_inference_steps)
]
sigmas = np.array([
sigma_min + (ppf * (sigma_max - sigma_min)) for ppf in [
scipy.stats.beta.ppf(timestep, alpha, beta)
for timestep in 1 - np.linspace(0, 1, num_inference_steps)
]
)
])
return sigmas
def _time_shift_exponential(self, mu, sigma, t):
return math.exp(mu) / (math.exp(mu) + (1 / t - 1) ** sigma)
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)
return mu / (mu + (1 / t - 1)**sigma)
def __len__(self):
return self.config.num_train_timesteps
return self.config.num_train_timesteps
fastvideo/v1/models/vaes/image_vae.py
+68 -44
View File
@@ -32,16 +32,13 @@ from typing import Dict, Optional, Union
import torch
import torch.nn as nn
from diffusers.models.autoencoders.vae import Encoder, Decoder
from diffusers.models.attention_processor import (
AttentionProcessor,
AttnAddedKVProcessor,
AttnProcessor,
Attention,
FusedAttnProcessor2_0,
ADDED_KV_ATTENTION_PROCESSORS,
CROSS_ATTENTION_PROCESSORS,
)
from diffusers.models.attention_processor import (ADDED_KV_ATTENTION_PROCESSORS,
CROSS_ATTENTION_PROCESSORS,
Attention, AttentionProcessor,
AttnAddedKVProcessor,
AttnProcessor,
FusedAttnProcessor2_0)
from diffusers.models.autoencoders.vae import Decoder, Encoder
from fastvideo.v1.configs.models.vaes import ImageVAEConfig
from fastvideo.v1.models.vaes.common import DiagonalGaussianDistribution
@@ -75,7 +72,7 @@ class AutoencoderKL(nn.Module):
Synthesis with Latent Diffusion Models](https://arxiv.org/abs/2112.10752) paper.
force_upcast (`bool`, *optional*, default to `True`):
If enabled it will force the VAE to run in float32 for high image resolution pipelines, such as SD-XL. VAE
can be fine-tuned / trained to a lower range without loosing too much precision in which case
can be fine-tuned / trained to a lower range without losing too much precision in which case
`force_upcast` can be set to `False` - see: https://huggingface.co/madebyollin/sdxl-vae-fp16-fix
mid_block_add_attention (`bool`, *optional*, default to `True`):
If enabled, the mid_block of the Encoder and Decoder will have attention blocks. If set to false, the
@@ -106,7 +103,9 @@ class AutoencoderKL(nn.Module):
double_z=True,
mid_block_add_attention=config.mid_block_add_attention,
)
self.quant_conv = nn.Conv2d(2 * config.latent_channels, 2 * config.latent_channels, 1) if config.use_quant_conv else None
self.quant_conv = nn.Conv2d(2 * config.latent_channels, 2 *
config.latent_channels,
1) if config.use_quant_conv else None
if config.load_decoder:
# pass init params to Decoder
@@ -120,19 +119,19 @@ class AutoencoderKL(nn.Module):
act_fn=config.act_fn,
mid_block_add_attention=config.mid_block_add_attention,
)
self.post_quant_conv = nn.Conv2d(config.latent_channels, config.latent_channels, 1) if config.use_post_quant_conv else None
self.post_quant_conv = nn.Conv2d(
config.latent_channels, config.latent_channels,
1) if config.use_post_quant_conv else None
self.use_slicing = False
self.use_tiling = False
# only relevant if vae tiling is enabled
self.tile_sample_min_size = config.sample_size
sample_size = (
config.sample_size[0]
if isinstance(config.sample_size, (list, tuple))
else config.sample_size
)
self.tile_latent_min_size = int(sample_size / (2 ** (len(config.block_out_channels) - 1)))
sample_size = (config.sample_size[0] if isinstance(
config.sample_size, (list, tuple)) else config.sample_size)
self.tile_latent_min_size = int(
sample_size / (2**(len(config.block_out_channels) - 1)))
self.tile_overlap_factor = 0.25
def enable_tiling(self, use_tiling: bool = True):
@@ -173,14 +172,17 @@ class AutoencoderKL(nn.Module):
indexed by its weight name.
"""
# set recursively
processors = {}
processors: dict[str, AttentionProcessor] = {}
def fn_recursive_add_processors(name: str, module: torch.nn.Module, processors: Dict[str, AttentionProcessor]):
def fn_recursive_add_processors(name: str, module: torch.nn.Module,
processors: Dict[str,
AttentionProcessor]):
if hasattr(module, "get_processor"):
processors[f"{name}.processor"] = module.get_processor()
for sub_name, child in module.named_children():
fn_recursive_add_processors(f"{name}.{sub_name}", child, processors)
fn_recursive_add_processors(f"{name}.{sub_name}", child,
processors)
return processors
@@ -190,7 +192,9 @@ class AutoencoderKL(nn.Module):
return processors
# Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.set_attn_processor
def set_attn_processor(self, processor: Union[AttentionProcessor, Dict[str, AttentionProcessor]]):
def set_attn_processor(self, processor: Union[AttentionProcessor,
Dict[str,
AttentionProcessor]]):
r"""
Sets the attention processor to use to compute attention.
@@ -211,7 +215,8 @@ class AutoencoderKL(nn.Module):
f" number of attention layers: {count}. Please make sure to pass {count} processor classes."
)
def fn_recursive_attn_processor(name: str, module: torch.nn.Module, processor):
def fn_recursive_attn_processor(name: str, module: torch.nn.Module,
processor):
if hasattr(module, "set_processor"):
if not isinstance(processor, dict):
module.set_processor(processor)
@@ -219,7 +224,8 @@ class AutoencoderKL(nn.Module):
module.set_processor(processor.pop(f"{name}.processor"))
for sub_name, child in module.named_children():
fn_recursive_attn_processor(f"{name}.{sub_name}", child, processor)
fn_recursive_attn_processor(f"{name}.{sub_name}", child,
processor)
for name, module in self.named_children():
fn_recursive_attn_processor(name, module, processor)
@@ -229,9 +235,11 @@ class AutoencoderKL(nn.Module):
"""
Disables custom attention processors and sets the default attention implementation.
"""
if all(proc.__class__ in ADDED_KV_ATTENTION_PROCESSORS for proc in self.attn_processors.values()):
if all(proc.__class__ in ADDED_KV_ATTENTION_PROCESSORS
for proc in self.attn_processors.values()):
processor = AttnAddedKVProcessor()
elif all(proc.__class__ in CROSS_ATTENTION_PROCESSORS for proc in self.attn_processors.values()):
elif all(proc.__class__ in CROSS_ATTENTION_PROCESSORS
for proc in self.attn_processors.values()):
processor = AttnProcessor()
else:
raise ValueError(
@@ -243,7 +251,8 @@ class AutoencoderKL(nn.Module):
def _encode(self, x: torch.Tensor) -> torch.Tensor:
batch_size, num_channels, height, width = x.shape
if self.use_tiling and (width > self.tile_sample_min_size or height > self.tile_sample_min_size):
if self.use_tiling and (width > self.tile_sample_min_size
or height > self.tile_sample_min_size):
return self._tiled_encode(x)
enc = self.encoder(x)
@@ -276,7 +285,8 @@ class AutoencoderKL(nn.Module):
return enc
def _decode(self, z: torch.Tensor) -> torch.Tensor:
if self.use_tiling and (z.shape[-1] > self.tile_latent_min_size or z.shape[-2] > self.tile_latent_min_size):
if self.use_tiling and (z.shape[-1] > self.tile_latent_min_size
or z.shape[-2] > self.tile_latent_min_size):
return self.tiled_decode(z)
if self.post_quant_conv is not None:
@@ -302,23 +312,29 @@ class AutoencoderKL(nn.Module):
"""
if self.use_slicing and z.shape[0] > 1:
decoded_slices = [self._decode(z_slice).sample for z_slice in z.split(1)]
decoded_slices = [
self._decode(z_slice).sample for z_slice in z.split(1)
]
decoded = torch.cat(decoded_slices)
else:
decoded = self._decode(z)
return decoded
def blend_v(self, a: torch.Tensor, b: torch.Tensor, blend_extent: int) -> torch.Tensor:
def blend_v(self, a: torch.Tensor, b: torch.Tensor,
blend_extent: int) -> torch.Tensor:
blend_extent = min(a.shape[2], b.shape[2], blend_extent)
for y in range(blend_extent):
b[:, :, y, :] = a[:, :, -blend_extent + y, :] * (1 - y / blend_extent) + b[:, :, y, :] * (y / blend_extent)
b[:, :, y, :] = a[:, :, -blend_extent + y, :] * (
1 - y / blend_extent) + b[:, :, y, :] * (y / blend_extent)
return b
def blend_h(self, a: torch.Tensor, b: torch.Tensor, blend_extent: int) -> torch.Tensor:
def blend_h(self, a: torch.Tensor, b: torch.Tensor,
blend_extent: int) -> torch.Tensor:
blend_extent = min(a.shape[3], b.shape[3], blend_extent)
for x in range(blend_extent):
b[:, :, :, x] = a[:, :, :, -blend_extent + x] * (1 - x / blend_extent) + b[:, :, :, x] * (x / blend_extent)
b[:, :, :, x] = a[:, :, :, -blend_extent + x] * (
1 - x / blend_extent) + b[:, :, :, x] * (x / blend_extent)
return b
def _tiled_encode(self, x: torch.Tensor) -> torch.Tensor:
@@ -338,7 +354,8 @@ class AutoencoderKL(nn.Module):
The latent representation of the encoded videos.
"""
overlap_size = int(self.tile_sample_min_size * (1 - self.tile_overlap_factor))
overlap_size = int(self.tile_sample_min_size *
(1 - self.tile_overlap_factor))
blend_extent = int(self.tile_latent_min_size * self.tile_overlap_factor)
row_limit = self.tile_latent_min_size - blend_extent
@@ -347,9 +364,10 @@ class AutoencoderKL(nn.Module):
for i in range(0, x.shape[2], overlap_size):
row = []
for j in range(0, x.shape[3], overlap_size):
tile = x[:, :, i : i + self.tile_sample_min_size, j : j + self.tile_sample_min_size]
tile = x[:, :, i:i + self.tile_sample_min_size,
j:j + self.tile_sample_min_size]
tile = self.encoder(tile)
if self.config.use_quant_conv:
if self.quant_conv:
tile = self.quant_conv(tile)
row.append(tile)
rows.append(row)
@@ -389,7 +407,8 @@ class AutoencoderKL(nn.Module):
`tuple` is returned.
"""
overlap_size = int(self.tile_sample_min_size * (1 - self.tile_overlap_factor))
overlap_size = int(self.tile_sample_min_size *
(1 - self.tile_overlap_factor))
blend_extent = int(self.tile_latent_min_size * self.tile_overlap_factor)
row_limit = self.tile_latent_min_size - blend_extent
@@ -398,9 +417,10 @@ class AutoencoderKL(nn.Module):
for i in range(0, x.shape[2], overlap_size):
row = []
for j in range(0, x.shape[3], overlap_size):
tile = x[:, :, i : i + self.tile_sample_min_size, j : j + self.tile_sample_min_size]
tile = x[:, :, i:i + self.tile_sample_min_size,
j:j + self.tile_sample_min_size]
tile = self.encoder(tile)
if self.config.use_quant_conv:
if self.quant_conv:
tile = self.quant_conv(tile)
row.append(tile)
rows.append(row)
@@ -436,7 +456,8 @@ class AutoencoderKL(nn.Module):
If return_dict is True, a [`~models.vae.DecoderOutput`] is returned, otherwise a plain `tuple` is
returned.
"""
overlap_size = int(self.tile_latent_min_size * (1 - self.tile_overlap_factor))
overlap_size = int(self.tile_latent_min_size *
(1 - self.tile_overlap_factor))
blend_extent = int(self.tile_sample_min_size * self.tile_overlap_factor)
row_limit = self.tile_sample_min_size - blend_extent
@@ -446,8 +467,9 @@ class AutoencoderKL(nn.Module):
for i in range(0, z.shape[2], overlap_size):
row = []
for j in range(0, z.shape[3], overlap_size):
tile = z[:, :, i : i + self.tile_latent_min_size, j : j + self.tile_latent_min_size]
if self.config.use_post_quant_conv:
tile = z[:, :, i:i + self.tile_latent_min_size,
j:j + self.tile_latent_min_size]
if self.post_quant_conv:
tile = self.post_quant_conv(tile)
decoded = self.decoder(tile)
row.append(decoded)
@@ -507,7 +529,9 @@ class AutoencoderKL(nn.Module):
for _, attn_processor in self.attn_processors.items():
if "Added" in str(attn_processor.__class__.__name__):
raise ValueError("`fuse_qkv_projections()` is not supported for models having added KV projections.")
raise ValueError(
"`fuse_qkv_projections()` is not supported for models having added KV projections."
)
self.original_attn_processors = self.attn_processors
fastvideo/v1/pipelines/composed_pipeline_base.py
+4 -1
View File
@@ -128,7 +128,10 @@ class ComposedPipelineBase(ABC):
modules_config = deepcopy(self.config)
# remove keys that are not pipeline modules
keys_to_remove = [key for key in list(modules_config.keys()) if str(key).startswith("_")]
keys_to_remove = [
key for key in list(modules_config.keys())
if str(key).startswith("_")
]
for key in keys_to_remove:
modules_config.pop(key)
fastvideo/v1/pipelines/flux/custom_stages.py
+45 -34
View File
@@ -1,18 +1,12 @@
import importlib
from einops import rearrange
import numpy as np
import torch
from fastvideo.v1.fastvideo_args import FastVideoArgs
from fastvideo.v1.pipelines.pipeline_batch_info import ForwardBatch
from fastvideo.v1.pipelines.stages.base import PipelineStage
from fastvideo.v1.pipelines.stages.denoising import DenoisingStage
from fastvideo.v1.forward_context import set_forward_context
from einops import rearrange
from fastvideo.v1.attention import get_attn_backend
from fastvideo.v1.distributed import (get_sequence_model_parallel_rank,
get_sequence_model_parallel_world_size,)
get_sequence_model_parallel_world_size)
from fastvideo.v1.distributed.communication_op import (
sequence_model_parallel_all_gather)
from fastvideo.v1.fastvideo_args import FastVideoArgs
@@ -20,54 +14,61 @@ from fastvideo.v1.forward_context import set_forward_context
from fastvideo.v1.logger import init_logger
from fastvideo.v1.pipelines.pipeline_batch_info import ForwardBatch
from fastvideo.v1.pipelines.stages.base import PipelineStage
from fastvideo.v1.pipelines.stages.denoising import DenoisingStage
from fastvideo.v1.platforms import _Backend
from fastvideo.v1.utils import PRECISION_TO_TYPE
logger = init_logger(__name__)
class TimestepsPreparationPreStage(PipelineStage):
def __init__(self, scheduler):
def __init__(self, scheduler) -> None:
super().__init__()
self.scheduler = scheduler
def calculate_shift(
self,
image_seq_len,
base_seq_len: int = 256,
max_seq_len: int = 4096,
base_shift: float = 0.5,
max_shift: float = 1.15
):
def calculate_shift(self,
image_seq_len,
base_seq_len: int = 256,
max_seq_len: int = 4096,
base_shift: float = 0.5,
max_shift: float = 1.15):
m = (max_shift - base_shift) / (max_seq_len - base_seq_len)
b = base_shift - m * base_seq_len
mu = image_seq_len * m + b
return mu
def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch:
def forward(self, batch: ForwardBatch,
fastvideo_args: FastVideoArgs) -> ForwardBatch:
assert batch.height is not None
assert batch.width is not None
batch.sigmas = np.linspace(
1.0, 1 / batch.num_inference_steps, batch.num_inference_steps
) if batch.sigmas is None else batch.sigmas
1.0, 1 / batch.num_inference_steps,
batch.num_inference_steps) if batch.sigmas is None else batch.sigmas
spatial_compression_ratio = fastvideo_args.vae_config.arch_config.spatial_compression_ratio
batch.extra_set_timesteps_kwargs["mu"] = (
batch.extra_set_timesteps_kwargs.get("mu", None) or
self.calculate_shift(
(batch.height // spatial_compression_ratio // 2) *
(batch.width // spatial_compression_ratio // 2),
batch.extra_set_timesteps_kwargs.get("mu", None)
or self.calculate_shift(
(batch.height // spatial_compression_ratio // 2) *
(batch.width // spatial_compression_ratio // 2),
self.scheduler.config.get("base_image_seq_len", 256),
self.scheduler.config.get("max_image_seq_len", 4096),
self.scheduler.config.get("base_shift", 0.5),
self.scheduler.config.get("max_shift", 1.15),
)
)
))
return batch
class DenoisingPreprocessingStage(PipelineStage):
def __init__(self):
def __init__(self) -> None:
super().__init__()
def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch:
def forward(self, batch: ForwardBatch,
fastvideo_args: FastVideoArgs) -> ForwardBatch:
# [B, in_channels // 4, 1, H, W] -> [B, H // 2, W // 2, in_channels]
assert batch.latents is not None
b, c, _, h, w = batch.latents.shape
latents = batch.latents.view(b, c, h // 2, 2, w // 2, 2)
latents = latents.permute(0, 2, 4, 1, 3, 5)
@@ -75,11 +76,18 @@ class DenoisingPreprocessingStage(PipelineStage):
batch.latents = latents
return batch
class DenoisingPostprocessingStage(PipelineStage):
def __init__(self):
def __init__(self) -> None:
super().__init__()
def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch:
def forward(self, batch: ForwardBatch,
fastvideo_args: FastVideoArgs) -> ForwardBatch:
assert batch.latents is not None
assert batch.height_latents is not None
assert batch.width_latents is not None
latents = batch.latents
# Skip decoding if output type is latent
if fastvideo_args.output_type == "latent":
@@ -105,6 +113,7 @@ if spec is not None:
from fastvideo.v1.attention.backends.sliding_tile_attn import (
SlidingTileAttentionBackend)
class FluxDenoisingStage(DenoisingStage):
"""
Stage for running the denoising loop in diffusion pipelines.
@@ -338,11 +347,13 @@ class FluxDenoisingStage(DenoisingStage):
class ImageOutputStage(PipelineStage):
def __init__(self):
def __init__(self) -> None:
super().__init__()
def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch:
def forward(self, batch: ForwardBatch,
fastvideo_args: FastVideoArgs) -> ForwardBatch:
output = batch.output
output = output.unsqueeze(2)
batch.output = output
return batch
return batch
fastvideo/v1/pipelines/flux/flux_pipeline.py
+7 -10
View File
@@ -9,17 +9,14 @@ using the modular pipeline architecture.
from fastvideo.v1.fastvideo_args import FastVideoArgs
from fastvideo.v1.logger import init_logger
from fastvideo.v1.pipelines.composed_pipeline_base import ComposedPipelineBase
from fastvideo.v1.pipelines.stages import (TextEncodingStage,
ConditioningStage,
DecodingStage,
from fastvideo.v1.pipelines.flux.custom_stages import (
DenoisingPostprocessingStage, DenoisingPreprocessingStage,
FluxDenoisingStage, ImageOutputStage, TimestepsPreparationPreStage)
from fastvideo.v1.pipelines.stages import (ConditioningStage, DecodingStage,
InputValidationStage,
LatentPreparationStage,
TextEncodingStage,
TimestepPreparationStage)
from fastvideo.v1.pipelines.flux.custom_stages import (DenoisingPreprocessingStage,
DenoisingPostprocessingStage,
ImageOutputStage,
FluxDenoisingStage,
TimestepsPreparationPreStage)
logger = init_logger(__name__)
@@ -79,7 +76,7 @@ class FluxPipeline(ComposedPipelineBase):
self.add_stage(stage_name="decoding_stage",
stage=DecodingStage(vae=self.get_module("vae")))
self.add_stage(stage_name="output_stage",
stage=ImageOutputStage())
self.add_stage(stage_name="output_stage", stage=ImageOutputStage())
EntryClass = FluxPipeline
fastvideo/v1/pipelines/stages/denoising.py
+4 -2
View File
@@ -64,6 +64,7 @@ class DenoisingStage(PipelineStage):
Returns:
The batch with denoised latents.
"""
assert batch.latents is not None
# Prepare extra step kwargs for scheduler
extra_step_kwargs = self.prepare_extra_func_kwargs(
self.scheduler.step,
@@ -301,8 +302,9 @@ class DenoisingStage(PipelineStage):
"""
extra_step_kwargs = {}
for k, v in kwargs.items():
accepts = (k in set(inspect.signature(func).parameters.keys()) or
"kwargs" in set(inspect.signature(func).parameters.keys()))
accepts = (k in set(inspect.signature(func).parameters.keys())
or "kwargs" in set(
inspect.signature(func).parameters.keys()))
if accepts:
extra_step_kwargs[k] = v
return extra_step_kwargs
fastvideo/v1/pipelines/stages/timestep_preparation.py
-2
View File
@@ -7,8 +7,6 @@ This module contains implementations of timestep preparation stages for diffusio
import inspect
import numpy as np
from fastvideo.v1.fastvideo_args import FastVideoArgs
from fastvideo.v1.logger import init_logger
from fastvideo.v1.pipelines.pipeline_batch_info import ForwardBatch