From fd80ccbf8875c30aefd24979fd1cbe617ee4b740 Mon Sep 17 00:00:00 2001 From: Will Lin Date: Tue, 6 May 2025 14:33:35 -0700 Subject: [PATCH] lint --- .pre-commit-config.yaml | 1 + fastvideo/v1/configs/models/base.py | 5 +- fastvideo/v1/configs/models/dits/__init__.py | 2 +- fastvideo/v1/configs/models/dits/flux.py | 3 +- fastvideo/v1/configs/models/vaes/__init__.py | 2 +- fastvideo/v1/configs/models/vaes/image_vae.py | 41 ++--- fastvideo/v1/configs/pipelines/flux.py | 20 +-- fastvideo/v1/configs/pipelines/registry.py | 2 +- fastvideo/v1/models/dits/flux.py | 66 ++++---- .../scheduling_flow_match_euler_discrete.py | 141 ++++++++++-------- fastvideo/v1/models/vaes/image_vae.py | 112 ++++++++------ .../v1/pipelines/composed_pipeline_base.py | 5 +- fastvideo/v1/pipelines/flux/custom_stages.py | 79 +++++----- fastvideo/v1/pipelines/flux/flux_pipeline.py | 17 +-- fastvideo/v1/pipelines/stages/denoising.py | 6 +- .../pipelines/stages/timestep_preparation.py | 2 - 16 files changed, 280 insertions(+), 224 deletions(-) diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index f50aba44b..03464f48f 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -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 diff --git a/fastvideo/v1/configs/models/base.py b/fastvideo/v1/configs/models/base.py index 7142c5610..9fb8fa88a 100644 --- a/fastvideo/v1/configs/models/base.py +++ b/fastvideo/v1/configs/models/base.py @@ -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) diff --git a/fastvideo/v1/configs/models/dits/__init__.py b/fastvideo/v1/configs/models/dits/__init__.py index d402302b8..c1be86a48 100644 --- a/fastvideo/v1/configs/models/dits/__init__.py +++ b/fastvideo/v1/configs/models/dits/__init__.py @@ -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"] diff --git a/fastvideo/v1/configs/models/dits/flux.py b/fastvideo/v1/configs/models/dits/flux.py index aea2215c4..afd89d0a4 100644 --- a/fastvideo/v1/configs/models/dits/flux.py +++ b/fastvideo/v1/configs/models/dits/flux.py @@ -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 diff --git a/fastvideo/v1/configs/models/vaes/__init__.py b/fastvideo/v1/configs/models/vaes/__init__.py index 9d3b69b36..6909550a6 100644 --- a/fastvideo/v1/configs/models/vaes/__init__.py +++ b/fastvideo/v1/configs/models/vaes/__init__.py @@ -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", diff --git a/fastvideo/v1/configs/models/vaes/image_vae.py b/fastvideo/v1/configs/models/vaes/image_vae.py index e38f0f3fd..243d7be52 100644 --- a/fastvideo/v1/configs/models/vaes/image_vae.py +++ b/fastvideo/v1/configs/models/vaes/image_vae.py @@ -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 \ No newline at end of file + use_parallel_tiling: bool = False diff --git a/fastvideo/v1/configs/pipelines/flux.py b/fastvideo/v1/configs/pipelines/flux.py index d8a16d7d8..408fdb5d2 100644 --- a/fastvideo/v1/configs/pipelines/flux.py +++ b/fastvideo/v1/configs/pipelines/flux.py @@ -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 \ No newline at end of file + self.vae_config.load_decoder = True diff --git a/fastvideo/v1/configs/pipelines/registry.py b/fastvideo/v1/configs/pipelines/registry.py index 668e60e17..4ae3f7f95 100644 --- a/fastvideo/v1/configs/pipelines/registry.py +++ b/fastvideo/v1/configs/pipelines/registry.py @@ -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) diff --git a/fastvideo/v1/models/dits/flux.py b/fastvideo/v1/models/dits/flux.py index 876d0e009..0c218f1b7 100644 --- a/fastvideo/v1/models/dits/flux.py +++ b/fastvideo/v1/models/dits/flux.py @@ -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) diff --git a/fastvideo/v1/models/schedulers/scheduling_flow_match_euler_discrete.py b/fastvideo/v1/models/schedulers/scheduling_flow_match_euler_discrete.py index e714a2a25..d65900f22 100644 --- a/fastvideo/v1/models/schedulers/scheduling_flow_match_euler_discrete.py +++ b/fastvideo/v1/models/schedulers/scheduling_flow_match_euler_discrete.py @@ -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 \ No newline at end of file + return self.config.num_train_timesteps diff --git a/fastvideo/v1/models/vaes/image_vae.py b/fastvideo/v1/models/vaes/image_vae.py index a454c4329..580bd8700 100644 --- a/fastvideo/v1/models/vaes/image_vae.py +++ b/fastvideo/v1/models/vaes/image_vae.py @@ -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 diff --git a/fastvideo/v1/pipelines/composed_pipeline_base.py b/fastvideo/v1/pipelines/composed_pipeline_base.py index 9ede44e10..858456493 100644 --- a/fastvideo/v1/pipelines/composed_pipeline_base.py +++ b/fastvideo/v1/pipelines/composed_pipeline_base.py @@ -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) diff --git a/fastvideo/v1/pipelines/flux/custom_stages.py b/fastvideo/v1/pipelines/flux/custom_stages.py index 6149822da..973db0755 100644 --- a/fastvideo/v1/pipelines/flux/custom_stages.py +++ b/fastvideo/v1/pipelines/flux/custom_stages.py @@ -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 \ No newline at end of file + return batch diff --git a/fastvideo/v1/pipelines/flux/flux_pipeline.py b/fastvideo/v1/pipelines/flux/flux_pipeline.py index 09fcdf8ac..12ea1f7a8 100644 --- a/fastvideo/v1/pipelines/flux/flux_pipeline.py +++ b/fastvideo/v1/pipelines/flux/flux_pipeline.py @@ -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 diff --git a/fastvideo/v1/pipelines/stages/denoising.py b/fastvideo/v1/pipelines/stages/denoising.py index e2164f3fc..d173d8109 100644 --- a/fastvideo/v1/pipelines/stages/denoising.py +++ b/fastvideo/v1/pipelines/stages/denoising.py @@ -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 diff --git a/fastvideo/v1/pipelines/stages/timestep_preparation.py b/fastvideo/v1/pipelines/stages/timestep_preparation.py index 0f5cdfba8..4fb58caf3 100644 --- a/fastvideo/v1/pipelines/stages/timestep_preparation.py +++ b/fastvideo/v1/pipelines/stages/timestep_preparation.py @@ -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