revert other stages

examples/inference/basic/basic.py

@@ -1,18 +1,17 @@
 from fastvideo import VideoGenerator
 
-from fastvideo.configs.sample import SamplingParam
+# from fastvideo.configs.sample import SamplingParam
 
-OUTPUT_PATH = "video_samples_5b"
+OUTPUT_PATH = "video_samples"
 def main():
     # FastVideo will automatically use the optimal default arguments for the
     # model.
     # If a local path is provided, FastVideo will make a best effort
     # attempt to identify the optimal arguments.
-    model_id = "Wan-AI/Wan2.2-TI2V-5B-Diffusers"
     generator = VideoGenerator.from_pretrained(
-        model_id,
+        "Wan-AI/Wan2.1-T2V-1.3B-Diffusers",
         # FastVideo will automatically handle distributed setup
-        num_gpus=1,
+        num_gpus=2,
         use_fsdp_inference=True,
         dit_cpu_offload=False,
         vae_cpu_offload=False,
@@ -20,22 +19,27 @@ def main():
         # image_encoder_cpu_offload=False,
     )
 
-    sampling_param = SamplingParam.from_pretrained("Wan-AI/Wan2.2-TI2V-5B-Diffusers")
+    # sampling_param = SamplingParam.from_pretrained("Wan-AI/Wan2.1-T2V-1.3B-Diffusers")
     # sampling_param.num_frames = 45
     # sampling_param.image_path = "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/astronaut.jpg"
     # Generate videos with the same simple API, regardless of GPU count
     prompt = (
-        "A curious raccoon peers through a vibrant field of yellow sunflowers, its eyes wide with interest. The playful yet serene atmosphere is complemented by soft natural light filtering through the petals. Mid-shot, warm and cheerful tones."
+        "A curious raccoon peers through a vibrant field of yellow sunflowers, its eyes "
+        "wide with interest. The playful yet serene atmosphere is complemented by soft "
+        "natural light filtering through the petals. Mid-shot, warm and cheerful tones."
     )
-    video = generator.generate_video(prompt, output_path=OUTPUT_PATH, save_video=True, sampling_param=sampling_param)
+    video = generator.generate_video(prompt, output_path=OUTPUT_PATH, save_video=True)
     # video = generator.generate_video(prompt, sampling_param=sampling_param, output_path="wan_t2v_videos/")
 
     # Generate another video with a different prompt, without reloading the
     # model!
     prompt2 = (
-        "A majestic lion strides across the golden savanna, its powerful frame glistening under the warm afternoon sun. The tall grass ripples gently in the breeze, enhancing the lion's commanding presence. The tone is vibrant, embodying the raw energy of the wild. Low angle, steady tracking shot, cinematic."
-    )
-    video2 = generator.generate_video(prompt2, output_path=OUTPUT_PATH, save_video=True, sampling_param=sampling_param)
+        "A majestic lion strides across the golden savanna, its powerful frame "
+        "glistening under the warm afternoon sun. The tall grass ripples gently in "
+        "the breeze, enhancing the lion's commanding presence. The tone is vibrant, "
+        "embodying the raw energy of the wild. Low angle, steady tracking shot, "
+        "cinematic.")
+    video2 = generator.generate_video(prompt2, output_path=OUTPUT_PATH, save_video=True)
 
 
 if __name__ == "__main__":

examples/inference/basic/basic_i2v.py

@@ -1,42 +0,0 @@
-from fastvideo import VideoGenerator
-
-from fastvideo.configs.sample import SamplingParam
-
-OUTPUT_PATH = "video_samples_5b"
-def main():
-    # FastVideo will automatically use the optimal default arguments for the
-    # model.
-    # If a local path is provided, FastVideo will make a best effort
-    # attempt to identify the optimal arguments.
-    model_id = "Wan-AI/Wan2.2-TI2V-5B-Diffusers"
-    generator = VideoGenerator.from_pretrained(
-        model_id,
-        # FastVideo will automatically handle distributed setup
-        num_gpus=1,
-        use_fsdp_inference=True,
-        dit_cpu_offload=False,
-        vae_cpu_offload=False,
-        text_encoder_cpu_offload=False,
-        # image_encoder_cpu_offload=False,
-    )
-
-    sampling_param = SamplingParam.from_pretrained("Wan-AI/Wan2.2-TI2V-5B-Diffusers")
-    # sampling_param.num_frames = 45
-    sampling_param.image_path = "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/astronaut.jpg"
-    # Generate videos with the same simple API, regardless of GPU count
-    prompt = (
-        "An astronaut hatching from an egg, on the surface of the moon, the darkness and depth of space realised in the background. High quality, ultrarealistic detail and breath-taking movie-like camera shot."
-    )
-    video = generator.generate_video(prompt, output_path=OUTPUT_PATH, save_video=True, sampling_param=sampling_param)
-    # video = generator.generate_video(prompt, sampling_param=sampling_param, output_path="wan_t2v_videos/")
-
-    # Generate another video with a different prompt, without reloading the
-    # model!
-    # prompt2 = (
-    #     "A majestic lion strides across the golden savanna, its powerful frame glistening under the warm afternoon sun. The tall grass ripples gently in the breeze, enhancing the lion's commanding presence. The tone is vibrant, embodying the raw energy of the wild. Low angle, steady tracking shot, cinematic."
-    # )
-    # video2 = generator.generate_video(prompt2, output_path=OUTPUT_PATH, save_video=True, sampling_param=sampling_param)
-
-
-if __name__ == "__main__":
-    main()

fastvideo/configs/models/dits/wanvideo.py

@@ -1,5 +1,4 @@
 # SPDX-License-Identifier: Apache-2.0
-from typing import Optional
 from dataclasses import dataclass, field
 
 from fastvideo.configs.models.dits.base import DiTArchConfig, DiTConfig
@@ -90,7 +89,6 @@ class WanVideoArchConfig(DiTArchConfig):
     image_dim: int | None = None
     added_kv_proj_dim: int | None = None
     rope_max_seq_len: int = 1024
-    pos_embed_seq_len: Optional[int] = None
     exclude_lora_layers: list[str] = field(default_factory=lambda: ["embedder"])
 
     def __post_init__(self):

fastvideo/configs/models/vaes/wanvae.py

@@ -1,5 +1,4 @@
 # SPDX-License-Identifier: Apache-2.0
-from typing import Optional
 from dataclasses import dataclass, field
 
 import torch
@@ -10,7 +9,6 @@ from fastvideo.configs.models.vaes.base import VAEArchConfig, VAEConfig
 @dataclass
 class WanVAEArchConfig(VAEArchConfig):
     base_dim: int = 96
-    decoder_base_dim: Optional[int] = None
     z_dim: int = 16
     dim_mult: tuple[int, ...] = (1, 2, 4, 4)
     num_res_blocks: int = 2
@@ -53,21 +51,14 @@ class WanVAEArchConfig(VAEArchConfig):
         2.8251,
         1.9160,
     )
-    is_residual: bool = False
-    in_channels: int = 3
-    out_channels: int = 3
-    patch_size: Optional[int] = None
-    scale_factor_temporal: int = 4
-    scale_factor_spatial: int = 8
-    clip_output: bool = True
+    temporal_compression_ratio = 4
+    spatial_compression_ratio = 8
 
     def __post_init__(self):
         self.scaling_factor: torch.tensor = 1.0 / torch.tensor(
             self.latents_std).view(1, self.z_dim, 1, 1, 1)
         self.shift_factor: torch.tensor = torch.tensor(self.latents_mean).view(
             1, self.z_dim, 1, 1, 1)
-        self.temporal_compression_ratio = self.scale_factor_temporal
-        self.spatial_compression_ratio = self.scale_factor_spatial
 
 
 @dataclass

fastvideo/configs/sample/registry.py

@@ -11,9 +11,7 @@ from fastvideo.configs.sample.wan import (Wan2_1_Fun_1_3B_InP_SamplingParam,
                                           WanI2V_14B_480P_SamplingParam,
                                           WanI2V_14B_720P_SamplingParam,
                                           WanT2V_1_3B_SamplingParam,
-                                          WanT2V_14B_SamplingParam,
-                                          Wan2_2_TI2V_5B_SamplingParam,
-                                          Wan2_1_Fun_1_3B_InP_SamplingParam)
+                                          WanT2V_14B_SamplingParam)
 from fastvideo.logger import init_logger
 from fastvideo.utils import (maybe_download_model_index,
                              verify_model_config_and_directory)

fastvideo/configs/sample/wan.py

@@ -134,13 +134,10 @@ class Wan2_2_Base_SamplingParam(SamplingParam):
 @dataclass
 class Wan2_2_TI2V_5B_SamplingParam(Wan2_2_Base_SamplingParam):
     """Sampling parameters for Wan2.2 TI2V 5B model."""
-    # Video parameters
     height: int = 704
     width: int = 1280
     num_frames: int = 121
     fps: int = 24
-
-    # Denoising stage
     guidance_scale: float = 5.0
     num_inference_steps: int = 50
 

fastvideo/models/schedulers/scheduling_unipc_multistep.py

@@ -1,4 +1,6 @@
-# Copyright 2025 TSAIL Team and The HuggingFace Team. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+# Copyright 2024 TSAIL Team and The HuggingFace Team. All rights reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -12,36 +14,34 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-# DISCLAIMER: check https://huggingface.co/papers/2302.04867 and https://github.com/wl-zhao/UniPC for more info
+# DISCLAIMER: check https://arxiv.org/abs/2302.04867 and https://github.com/wl-zhao/UniPC for more info
 # The codebase is modified based on https://github.com/huggingface/diffusers/blob/main/src/diffusers/schedulers/scheduling_dpmsolver_multistep.py
 # ==============================================================================
 #
-# Modified from diffusers==0.35.0.dev0
+# Modified from diffusers==0.33.0.dev0
 #
 # ==============================================================================
 
 import math
-from typing import List, Optional, Tuple, Union
 
 import numpy as np
+import scipy.stats
 import torch
-
 from diffusers.configuration_utils import ConfigMixin, register_to_config
-from diffusers.utils import deprecate, is_scipy_available
-from diffusers.schedulers.scheduling_utils import KarrasDiffusionSchedulers, SchedulerMixin, SchedulerOutput
+from diffusers.schedulers.scheduling_utils import (KarrasDiffusionSchedulers,
+                                                   SchedulerMixin,
+                                                   SchedulerOutput)
+from diffusers.utils import deprecate
 
 from fastvideo.models.schedulers.base import BaseScheduler
 
-if is_scipy_available():
-    import scipy.stats
-
 
 # Copied from diffusers.schedulers.scheduling_ddpm.betas_for_alpha_bar
 def betas_for_alpha_bar(
     num_diffusion_timesteps,
     max_beta=0.999,
     alpha_transform_type="cosine",
-):
+) -> torch.Tensor:
     """
     Create a beta schedule that discretizes the given alpha_t_bar function, which defines the cumulative product of
     (1-beta) over time from t = [0,1].
@@ -62,16 +62,17 @@ def betas_for_alpha_bar(
     """
     if alpha_transform_type == "cosine":
 
-        def alpha_bar_fn(t):
-            return math.cos((t + 0.008) / 1.008 * math.pi / 2) ** 2
+        def alpha_bar_fn(t: float) -> float:
+            return math.cos((t + 0.008) / 1.008 * math.pi / 2)**2
 
     elif alpha_transform_type == "exp":
 
-        def alpha_bar_fn(t):
+        def alpha_bar_fn(t: float) -> float:
             return math.exp(t * -12.0)
 
     else:
-        raise ValueError(f"Unsupported alpha_transform_type: {alpha_transform_type}")
+        raise ValueError(
+            f"Unsupported alpha_transform_type: {alpha_transform_type}")
 
     betas = []
     for i in range(num_diffusion_timesteps):
@@ -82,9 +83,9 @@ def betas_for_alpha_bar(
 
 
 # Copied from diffusers.schedulers.scheduling_ddim.rescale_zero_terminal_snr
-def rescale_zero_terminal_snr(betas):
+def rescale_zero_terminal_snr(betas: torch.Tensor) -> torch.Tensor:
     """
-    Rescales betas to have zero terminal SNR Based on https://huggingface.co/papers/2305.08891 (Algorithm 1)
+    Rescales betas to have zero terminal SNR Based on https://arxiv.org/pdf/2305.08891.pdf (Algorithm 1)
 
 
     Args:
@@ -107,7 +108,8 @@ def rescale_zero_terminal_snr(betas):
     alphas_bar_sqrt -= alphas_bar_sqrt_T
 
     # Scale so the first timestep is back to the old value.
-    alphas_bar_sqrt *= alphas_bar_sqrt_0 / (alphas_bar_sqrt_0 - alphas_bar_sqrt_T)
+    alphas_bar_sqrt *= alphas_bar_sqrt_0 / (alphas_bar_sqrt_0 -
+                                            alphas_bar_sqrt_T)
 
     # Convert alphas_bar_sqrt to betas
     alphas_bar = alphas_bar_sqrt**2  # Revert sqrt
@@ -174,8 +176,6 @@ class UniPCMultistepScheduler(SchedulerMixin, ConfigMixin, BaseScheduler):
         use_beta_sigmas (`bool`, *optional*, defaults to `False`):
             Whether to use beta sigmas for step sizes in the noise schedule during the sampling process. Refer to [Beta
             Sampling is All You Need](https://huggingface.co/papers/2407.12173) for more information.
-        use_flow_sigmas (`bool`, *optional*, defaults to `False`):
-            Whether to use flow sigmas for step sizes in the noise schedule during the sampling process.
         timestep_spacing (`str`, defaults to `"linspace"`):
             The way the timesteps should be scaled. Refer to Table 2 of the [Common Diffusion Noise Schedules and
             Sample Steps are Flawed](https://huggingface.co/papers/2305.08891) for more information.
@@ -200,7 +200,7 @@ class UniPCMultistepScheduler(SchedulerMixin, ConfigMixin, BaseScheduler):
         beta_start: float = 0.0001,
         beta_end: float = 0.02,
         beta_schedule: str = "linear",
-        trained_betas: Optional[Union[np.ndarray, List[float]]] = None,
+        trained_betas: np.ndarray | list[float] | None = None,
         solver_order: int = 2,
         prediction_type: str = "epsilon",
         thresholding: bool = False,
@@ -209,38 +209,44 @@ class UniPCMultistepScheduler(SchedulerMixin, ConfigMixin, BaseScheduler):
         predict_x0: bool = True,
         solver_type: str = "bh2",
         lower_order_final: bool = True,
-        disable_corrector: List[int] = [],
+        disable_corrector: tuple[int, ...] = (),
         solver_p: SchedulerMixin = None,
-        use_karras_sigmas: Optional[bool] = False,
-        use_exponential_sigmas: Optional[bool] = False,
-        use_beta_sigmas: Optional[bool] = False,
-        use_flow_sigmas: Optional[bool] = False,
-        flow_shift: Optional[float] = 1.0,
+        use_karras_sigmas: bool | None = False,
+        use_exponential_sigmas: bool | None = False,
+        use_beta_sigmas: bool | None = False,
+        use_flow_sigmas: bool | None = False,
+        flow_shift: float | None = 1.0,
         timestep_spacing: str = "linspace",
         steps_offset: int = 0,
-        final_sigmas_type: Optional[str] = "zero",  # "zero", "sigma_min"
+        final_sigmas_type: str | None = "zero",  # "zero", "sigma_min"
         rescale_betas_zero_snr: bool = False,
-        use_dynamic_shifting: bool = False,
-        time_shift_type: str = "exponential",
     ):
-        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:
+        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 trained_betas is not None:
             self.betas = torch.tensor(trained_betas, dtype=torch.float32)
         elif beta_schedule == "linear":
-            self.betas = torch.linspace(beta_start, beta_end, num_train_timesteps, dtype=torch.float32)
+            self.betas = torch.linspace(beta_start,
+                                        beta_end,
+                                        num_train_timesteps,
+                                        dtype=torch.float32)
         elif beta_schedule == "scaled_linear":
             # this schedule is very specific to the latent diffusion model.
-            self.betas = torch.linspace(beta_start**0.5, beta_end**0.5, num_train_timesteps, dtype=torch.float32) ** 2
+            self.betas = torch.linspace(beta_start**0.5,
+                                        beta_end**0.5,
+                                        num_train_timesteps,
+                                        dtype=torch.float32)**2
         elif beta_schedule == "squaredcos_cap_v2":
             # Glide cosine schedule
             self.betas = betas_for_alpha_bar(num_train_timesteps)
         else:
-            raise NotImplementedError(f"{beta_schedule} is not implemented for {self.__class__}")
+            raise NotImplementedError(
+                f"{beta_schedule} is not implemented for {self.__class__}")
 
         if rescale_betas_zero_snr:
             self.betas = rescale_zero_terminal_snr(self.betas)
@@ -257,7 +263,7 @@ class UniPCMultistepScheduler(SchedulerMixin, ConfigMixin, BaseScheduler):
         self.alpha_t = torch.sqrt(self.alphas_cumprod)
         self.sigma_t = torch.sqrt(1 - self.alphas_cumprod)
         self.lambda_t = torch.log(self.alpha_t) - torch.log(self.sigma_t)
-        self.sigmas = ((1 - self.alphas_cumprod) / self.alphas_cumprod) ** 0.5
+        self.sigmas = ((1 - self.alphas_cumprod) / self.alphas_cumprod)**0.5
 
         # standard deviation of the initial noise distribution
         self.init_noise_sigma = 1.0
@@ -266,22 +272,27 @@ class UniPCMultistepScheduler(SchedulerMixin, ConfigMixin, BaseScheduler):
             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__}")
+                raise NotImplementedError(
+                    f"{solver_type} is not implemented for {self.__class__}")
 
         self.predict_x0 = predict_x0
         # setable values
-        self.num_inference_steps = None
-        timesteps = np.linspace(0, num_train_timesteps - 1, num_train_timesteps, dtype=np.float32)[::-1].copy()
+        self.num_inference_steps: int | None = None
+        timesteps = np.linspace(0,
+                                num_train_timesteps - 1,
+                                num_train_timesteps,
+                                dtype=np.float32)[::-1].copy()
         self.timesteps = torch.from_numpy(timesteps)
         self.model_outputs = [None] * solver_order
-        self.timestep_list = [None] * solver_order
+        self.timestep_list: list[int | torch.Tensor] = [None] * solver_order
         self.lower_order_nums = 0
-        self.disable_corrector = disable_corrector
+        self.disable_corrector = list(disable_corrector)
         self.solver_p = solver_p
         self.last_sample = None
-        self._step_index = None
-        self._begin_index = None
-        self.sigmas = self.sigmas.to("cpu")  # to avoid too much CPU/GPU communication
+        self._step_index: int | None = None
+        self._begin_index: int | None = None
+        self.sigmas = self.sigmas.to(
+            "cpu")  # to avoid too much CPU/GPU communication
 
         BaseScheduler.__init__(self)
 
@@ -313,9 +324,9 @@ class UniPCMultistepScheduler(SchedulerMixin, ConfigMixin, BaseScheduler):
         """
         self._begin_index = begin_index
 
-    def set_timesteps(
-        self, num_inference_steps: int, device: Union[str, torch.device] = None, mu: Optional[float] = None
-    ):
+    def set_timesteps(self,
+                      num_inference_steps: int,
+                      device: str | torch.device = None):
         """
         Sets the discrete timesteps used for the diffusion chain (to be run before inference).
 
@@ -325,40 +336,41 @@ class UniPCMultistepScheduler(SchedulerMixin, ConfigMixin, BaseScheduler):
             device (`str` or `torch.device`, *optional*):
                 The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
         """
-        # "linspace", "leading", "trailing" corresponds to annotation of Table 2. of https://huggingface.co/papers/2305.08891
-        if mu is not None:
-            assert self.config.use_dynamic_shifting and self.config.time_shift_type == "exponential"
-            self.config.flow_shift = np.exp(mu)
+        # "linspace", "leading", "trailing" corresponds to annotation of Table 2. of https://arxiv.org/abs/2305.08891
         if self.config.timestep_spacing == "linspace":
-            timesteps = (
-                np.linspace(0, self.config.num_train_timesteps - 1, num_inference_steps + 1)
-                .round()[::-1][:-1]
-                .copy()
-                .astype(np.int64)
-            )
+            timesteps = (np.linspace(
+                0, self.config.num_train_timesteps - 1, num_inference_steps +
+                1).round()[::-1][:-1].copy().astype(np.int64))
         elif self.config.timestep_spacing == "leading":
-            step_ratio = self.config.num_train_timesteps // (num_inference_steps + 1)
+            step_ratio = self.config.num_train_timesteps // (
+                num_inference_steps + 1)
             # creates integer timesteps by multiplying by ratio
             # casting to int to avoid issues when num_inference_step is power of 3
-            timesteps = (np.arange(0, num_inference_steps + 1) * step_ratio).round()[::-1][:-1].copy().astype(np.int64)
+            timesteps = (np.arange(0, num_inference_steps + 1) *
+                         step_ratio).round()[::-1][:-1].copy().astype(np.int64)
             timesteps += self.config.steps_offset
         elif self.config.timestep_spacing == "trailing":
             step_ratio = self.config.num_train_timesteps / num_inference_steps
             # creates integer timesteps by multiplying by ratio
             # casting to int to avoid issues when num_inference_step is power of 3
-            timesteps = np.arange(self.config.num_train_timesteps, 0, -step_ratio).round().copy().astype(np.int64)
+            timesteps = np.arange(self.config.num_train_timesteps, 0,
+                                  -step_ratio).round().copy().astype(np.int64)
             timesteps -= 1
         else:
             raise ValueError(
                 f"{self.config.timestep_spacing} is not supported. Please make sure to choose one of 'linspace', 'leading' or 'trailing'."
             )
 
-        sigmas = np.array(((1 - self.alphas_cumprod) / self.alphas_cumprod) ** 0.5)
+        sigmas = np.array(
+            ((1 - self.alphas_cumprod) / self.alphas_cumprod)**0.5)
         if self.config.use_karras_sigmas:
             log_sigmas = np.log(sigmas)
             sigmas = np.flip(sigmas).copy()
-            sigmas = self._convert_to_karras(in_sigmas=sigmas, num_inference_steps=num_inference_steps)
-            timesteps = np.array([self._sigma_to_t(sigma, log_sigmas) for sigma in sigmas]).round()
+            sigmas = self._convert_to_karras(
+                in_sigmas=sigmas, num_inference_steps=num_inference_steps)
+            timesteps = np.array([
+                self._sigma_to_t(sigma, log_sigmas) for sigma in sigmas
+            ]).round()
             if self.config.final_sigmas_type == "sigma_min":
                 sigma_last = sigmas[-1]
             elif self.config.final_sigmas_type == "zero":
@@ -371,8 +383,10 @@ class UniPCMultistepScheduler(SchedulerMixin, ConfigMixin, BaseScheduler):
         elif self.config.use_exponential_sigmas:
             log_sigmas = np.log(sigmas)
             sigmas = np.flip(sigmas).copy()
-            sigmas = self._convert_to_exponential(in_sigmas=sigmas, num_inference_steps=num_inference_steps)
-            timesteps = np.array([self._sigma_to_t(sigma, log_sigmas) for sigma in sigmas])
+            sigmas = self._convert_to_exponential(
+                in_sigmas=sigmas, num_inference_steps=num_inference_steps)
+            timesteps = np.array(
+                [self._sigma_to_t(sigma, log_sigmas) for sigma in sigmas])
             if self.config.final_sigmas_type == "sigma_min":
                 sigma_last = sigmas[-1]
             elif self.config.final_sigmas_type == "zero":
@@ -385,8 +399,10 @@ class UniPCMultistepScheduler(SchedulerMixin, ConfigMixin, BaseScheduler):
         elif self.config.use_beta_sigmas:
             log_sigmas = np.log(sigmas)
             sigmas = np.flip(sigmas).copy()
-            sigmas = self._convert_to_beta(in_sigmas=sigmas, num_inference_steps=num_inference_steps)
-            timesteps = np.array([self._sigma_to_t(sigma, log_sigmas) for sigma in sigmas])
+            sigmas = self._convert_to_beta(
+                in_sigmas=sigmas, num_inference_steps=num_inference_steps)
+            timesteps = np.array(
+                [self._sigma_to_t(sigma, log_sigmas) for sigma in sigmas])
             if self.config.final_sigmas_type == "sigma_min":
                 sigma_last = sigmas[-1]
             elif self.config.final_sigmas_type == "zero":
@@ -397,9 +413,12 @@ class UniPCMultistepScheduler(SchedulerMixin, ConfigMixin, BaseScheduler):
                 )
             sigmas = np.concatenate([sigmas, [sigma_last]]).astype(np.float32)
         elif self.config.use_flow_sigmas:
-            alphas = np.linspace(1, 1 / self.config.num_train_timesteps, num_inference_steps + 1)
+            alphas = np.linspace(1, 1 / self.config.num_train_timesteps,
+                                 num_inference_steps + 1)
             sigmas = 1.0 - alphas
-            sigmas = np.flip(self.config.flow_shift * sigmas / (1 + (self.config.flow_shift - 1) * sigmas))[:-1].copy()
+            sigmas = np.flip(
+                self.config.flow_shift * sigmas /
+                (1 + (self.config.flow_shift - 1) * sigmas))[:-1].copy()
             timesteps = (sigmas * self.config.num_train_timesteps).copy()
             if self.config.final_sigmas_type == "sigma_min":
                 sigma_last = sigmas[-1]
@@ -413,7 +432,8 @@ class UniPCMultistepScheduler(SchedulerMixin, ConfigMixin, BaseScheduler):
         else:
             sigmas = np.interp(timesteps, np.arange(0, len(sigmas)), sigmas)
             if self.config.final_sigmas_type == "sigma_min":
-                sigma_last = ((1 - self.alphas_cumprod[0]) / self.alphas_cumprod[0]) ** 0.5
+                sigma_last = ((1 - self.alphas_cumprod[0]) /
+                              self.alphas_cumprod[0])**0.5
             elif self.config.final_sigmas_type == "zero":
                 sigma_last = 0
             else:
@@ -423,7 +443,8 @@ class UniPCMultistepScheduler(SchedulerMixin, ConfigMixin, BaseScheduler):
             sigmas = np.concatenate([sigmas, [sigma_last]]).astype(np.float32)
 
         self.sigmas = torch.from_numpy(sigmas)
-        self.timesteps = torch.from_numpy(timesteps).to(device=device, dtype=torch.int64)
+        self.timesteps = torch.from_numpy(timesteps).to(device=device,
+                                                        dtype=torch.int64)
 
         self.num_inference_steps = len(timesteps)
 
@@ -438,7 +459,8 @@ class UniPCMultistepScheduler(SchedulerMixin, ConfigMixin, BaseScheduler):
         # add an index counter for schedulers that allow duplicated timesteps
         self._step_index = None
         self._begin_index = None
-        self.sigmas = self.sigmas.to("cpu")  # to avoid too much CPU/GPU communication
+        self.sigmas = self.sigmas.to(
+            "cpu")  # to avoid too much CPU/GPU communication
 
     # Copied from diffusers.schedulers.scheduling_ddpm.DDPMScheduler._threshold_sample
     def _threshold_sample(self, sample: torch.Tensor) -> torch.Tensor:
@@ -449,25 +471,31 @@ class UniPCMultistepScheduler(SchedulerMixin, ConfigMixin, BaseScheduler):
         pixels from saturation at each step. We find that dynamic thresholding results in significantly better
         photorealism as well as better image-text alignment, especially when using very large guidance weights."
 
-        https://huggingface.co/papers/2205.11487
+        https://arxiv.org/abs/2205.11487
         """
         dtype = sample.dtype
         batch_size, channels, *remaining_dims = sample.shape
 
         if dtype not in (torch.float32, torch.float64):
-            sample = sample.float()  # upcast for quantile calculation, and clamp not implemented for cpu half
+            sample = sample.float(
+            )  # upcast for quantile calculation, and clamp not implemented for cpu half
 
         # Flatten sample for doing quantile calculation along each image
         sample = sample.reshape(batch_size, channels * np.prod(remaining_dims))
 
         abs_sample = sample.abs()  # "a certain percentile absolute pixel value"
 
-        s = torch.quantile(abs_sample, self.config.dynamic_thresholding_ratio, dim=1)
+        s = torch.quantile(abs_sample,
+                           self.config.dynamic_thresholding_ratio,
+                           dim=1)
         s = torch.clamp(
             s, min=1, max=self.config.sample_max_value
         )  # When clamped to min=1, equivalent to standard clipping to [-1, 1]
-        s = s.unsqueeze(1)  # (batch_size, 1) because clamp will broadcast along dim=0
-        sample = torch.clamp(sample, -s, s) / s  # "we threshold xt0 to the range [-s, s] and then divide by s"
+        s = s.unsqueeze(
+            1)  # (batch_size, 1) because clamp will broadcast along dim=0
+        sample = torch.clamp(
+            sample, -s, s
+        ) / s  # "we threshold xt0 to the range [-s, s] and then divide by s"
 
         sample = sample.reshape(batch_size, channels, *remaining_dims)
         sample = sample.to(dtype)
@@ -475,7 +503,7 @@ class UniPCMultistepScheduler(SchedulerMixin, ConfigMixin, BaseScheduler):
         return sample
 
     # Copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler._sigma_to_t
-    def _sigma_to_t(self, sigma, log_sigmas):
+    def _sigma_to_t(self, sigma, log_sigmas) -> torch.Tensor:
         # get log sigma
         log_sigma = np.log(np.maximum(sigma, 1e-10))
 
@@ -483,7 +511,9 @@ class UniPCMultistepScheduler(SchedulerMixin, ConfigMixin, BaseScheduler):
         dists = log_sigma - log_sigmas[:, np.newaxis]
 
         # get sigmas range
-        low_idx = np.cumsum((dists >= 0), axis=0).argmax(axis=0).clip(max=log_sigmas.shape[0] - 2)
+        low_idx = np.cumsum(
+            (dists >= 0),
+            axis=0).argmax(axis=0).clip(max=log_sigmas.shape[0] - 2)
         high_idx = low_idx + 1
 
         low = log_sigmas[low_idx]
@@ -499,18 +529,20 @@ class UniPCMultistepScheduler(SchedulerMixin, ConfigMixin, BaseScheduler):
         return t
 
     # Copied from diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler._sigma_to_alpha_sigma_t
-    def _sigma_to_alpha_sigma_t(self, sigma):
+    def _sigma_to_alpha_sigma_t(
+            self, sigma: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
         if self.config.use_flow_sigmas:
             alpha_t = 1 - sigma
             sigma_t = sigma
         else:
-            alpha_t = 1 / ((sigma**2 + 1) ** 0.5)
+            alpha_t = 1 / ((sigma**2 + 1)**0.5)
             sigma_t = sigma * alpha_t
 
         return alpha_t, sigma_t
 
     # 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
@@ -530,13 +562,14 @@ class UniPCMultistepScheduler(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
@@ -554,13 +587,17 @@ class UniPCMultistepScheduler(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
@@ -578,15 +615,12 @@ class UniPCMultistepScheduler(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 convert_model_output(
@@ -616,7 +650,8 @@ class UniPCMultistepScheduler(SchedulerMixin, ConfigMixin, BaseScheduler):
             if len(args) > 1:
                 sample = args[1]
             else:
-                raise ValueError("missing `sample` as a required keyword argument")
+                raise ValueError(
+                    "missing `sample` as a required keyword argument")
         if timestep is not None:
             deprecate(
                 "timesteps",
@@ -659,15 +694,14 @@ class UniPCMultistepScheduler(SchedulerMixin, ConfigMixin, BaseScheduler):
             else:
                 raise ValueError(
                     f"prediction_type given as {self.config.prediction_type} must be one of `epsilon`, `sample`, or"
-                    " `v_prediction` for the UniPCMultistepScheduler."
-                )
+                    " `v_prediction` for the UniPCMultistepScheduler.")
 
     def multistep_uni_p_bh_update(
         self,
         model_output: torch.Tensor,
         *args,
         sample: torch.Tensor = None,
-        order: int = None,
+        order: int | None = None,
         **kwargs,
     ) -> torch.Tensor:
         """
@@ -687,17 +721,20 @@ class UniPCMultistepScheduler(SchedulerMixin, ConfigMixin, BaseScheduler):
             `torch.Tensor`:
                 The sample tensor at the previous timestep.
         """
-        prev_timestep = args[0] if len(args) > 0 else kwargs.pop("prev_timestep", None)
+        prev_timestep = args[0] if len(args) > 0 else kwargs.pop(
+            "prev_timestep", None)
         if sample is None:
             if len(args) > 1:
                 sample = args[1]
             else:
-                raise ValueError("missing `sample` as a required keyword argument")
+                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")
+                raise ValueError(
+                    " missing `order` as a required keyword argument")
         if prev_timestep is not None:
             deprecate(
                 "prev_timestep",
@@ -714,7 +751,8 @@ class UniPCMultistepScheduler(SchedulerMixin, ConfigMixin, BaseScheduler):
             x_t = self.solver_p.step(model_output, s0, x).prev_sample
             return x_t
 
-        sigma_t, sigma_s0 = self.sigmas[self.step_index + 1], self.sigmas[self.step_index]
+        sigma_t, sigma_s0 = self.sigmas[self.step_index +
+                                        1], self.sigmas[self.step_index]
         alpha_t, sigma_t = self._sigma_to_alpha_sigma_t(sigma_t)
         alpha_s0, sigma_s0 = self._sigma_to_alpha_sigma_t(sigma_s0)
 
@@ -728,7 +766,7 @@ class UniPCMultistepScheduler(SchedulerMixin, ConfigMixin, BaseScheduler):
         D1s = []
         for i in range(1, order):
             si = self.step_index - i
-            mi = model_output_list[-(i + 1)]
+            mi: torch.Tensor = 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
@@ -760,30 +798,32 @@ class UniPCMultistepScheduler(SchedulerMixin, ConfigMixin, BaseScheduler):
             factorial_i *= i + 1
             h_phi_k = h_phi_k / hh - 1 / factorial_i
 
-        R = torch.stack(R)
+        R_tensor: torch.Tensor = torch.stack(R)
         b = torch.tensor(b, device=device)
 
+        D1s_tensor: torch.Tensor | None = None
         if len(D1s) > 0:
-            D1s = torch.stack(D1s, dim=1)  # (B, K)
+            D1s_tensor = torch.stack(D1s, dim=1)  # (B, K)
             # for order 2, we use a simplified version
             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)
+                rhos_p = torch.linalg.solve(R_tensor[:-1, :-1],
+                                            b[:-1]).to(device).to(x.dtype)
         else:
-            D1s = None
+            D1s_tensor = 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)
+            if D1s_tensor is not None:
+                pred_res = torch.einsum("k,bkc...->bc...", rhos_p, D1s_tensor)
             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)
+            if D1s_tensor is not None:
+                pred_res = torch.einsum("k,bkc...->bc...", rhos_p, D1s_tensor)
             else:
                 pred_res = 0
             x_t = x_t_ - sigma_t * B_h * pred_res
@@ -795,9 +835,9 @@ class UniPCMultistepScheduler(SchedulerMixin, ConfigMixin, BaseScheduler):
         self,
         this_model_output: torch.Tensor,
         *args,
-        last_sample: torch.Tensor = None,
-        this_sample: torch.Tensor = None,
-        order: int = None,
+        last_sample: torch.Tensor | None = None,
+        this_sample: torch.Tensor | None = None,
+        order: int | None = None,
         **kwargs,
     ) -> torch.Tensor:
         """
@@ -819,22 +859,26 @@ class UniPCMultistepScheduler(SchedulerMixin, ConfigMixin, BaseScheduler):
             `torch.Tensor`:
                 The corrected sample tensor at the current timestep.
         """
-        this_timestep = args[0] if len(args) > 0 else kwargs.pop("this_timestep", None)
+        this_timestep = args[0] if len(args) > 0 else kwargs.pop(
+            "this_timestep", None)
         if last_sample is None:
             if len(args) > 1:
                 last_sample = args[1]
             else:
-                raise ValueError("missing `last_sample` as a required keyword argument")
+                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")
+                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")
+                raise ValueError(
+                    " missing`order` as a required keyword argument")
         if this_timestep is not None:
             deprecate(
                 "this_timestep",
@@ -849,7 +893,8 @@ class UniPCMultistepScheduler(SchedulerMixin, ConfigMixin, BaseScheduler):
         x_t = this_sample
         model_t = this_model_output
 
-        sigma_t, sigma_s0 = self.sigmas[self.step_index], self.sigmas[self.step_index - 1]
+        sigma_t, sigma_s0 = self.sigmas[self.step_index], self.sigmas[
+            self.step_index - 1]
         alpha_t, sigma_t = self._sigma_to_alpha_sigma_t(sigma_t)
         alpha_s0, sigma_s0 = self._sigma_to_alpha_sigma_t(sigma_s0)
 
@@ -863,7 +908,7 @@ class UniPCMultistepScheduler(SchedulerMixin, ConfigMixin, BaseScheduler):
         D1s = []
         for i in range(1, order):
             si = self.step_index - (i + 1)
-            mi = model_output_list[-(i + 1)]
+            mi: torch.Tensor = 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
@@ -898,10 +943,8 @@ class UniPCMultistepScheduler(SchedulerMixin, ConfigMixin, BaseScheduler):
         R = torch.stack(R)
         b = torch.tensor(b, device=device)
 
-        if len(D1s) > 0:
-            D1s = torch.stack(D1s, dim=1)
-        else:
-            D1s = None
+        D1s_tensor: torch.Tensor | None = torch.stack(
+            D1s, dim=1) if len(D1s) > 0 else None
 
         # for order 1, we use a simplified version
         if order == 1:
@@ -911,16 +954,18 @@ class UniPCMultistepScheduler(SchedulerMixin, ConfigMixin, BaseScheduler):
 
         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)
+            if D1s_tensor is not None:
+                corr_res = torch.einsum("k,bkc...->bc...", rhos_c[:-1],
+                                        D1s_tensor)
             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)
+            if D1s_tensor is not None:
+                corr_res = torch.einsum("k,bkc...->bc...", rhos_c[:-1],
+                                        D1s_tensor)
             else:
                 corr_res = 0
             D1_t = model_t - m0
@@ -929,7 +974,7 @@ class UniPCMultistepScheduler(SchedulerMixin, ConfigMixin, BaseScheduler):
         return x_t
 
     # Copied from diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler.index_for_timestep
-    def index_for_timestep(self, timestep, schedule_timesteps=None):
+    def index_for_timestep(self, timestep, schedule_timesteps=None) -> int:
         if schedule_timesteps is None:
             schedule_timesteps = self.timesteps
 
@@ -949,7 +994,7 @@ class UniPCMultistepScheduler(SchedulerMixin, ConfigMixin, BaseScheduler):
         return step_index
 
     # Copied from diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler._init_step_index
-    def _init_step_index(self, timestep):
+    def _init_step_index(self, timestep) -> None:
         """
         Initialize the step_index counter for the scheduler.
         """
@@ -964,10 +1009,10 @@ class UniPCMultistepScheduler(SchedulerMixin, ConfigMixin, BaseScheduler):
     def step(
         self,
         model_output: torch.Tensor,
-        timestep: Union[int, torch.Tensor],
+        timestep: int | torch.Tensor,
         sample: torch.Tensor,
         return_dict: bool = True,
-    ) -> Union[SchedulerOutput, Tuple]:
+    ) -> SchedulerOutput | tuple:
         """
         Predict the sample from the previous timestep by reversing the SDE. This function propagates the sample with
         the multistep UniPC.
@@ -996,11 +1041,13 @@ class UniPCMultistepScheduler(SchedulerMixin, ConfigMixin, BaseScheduler):
         if self.step_index is None:
             self._init_step_index(timestep)
 
-        use_corrector = (
-            self.step_index > 0 and self.step_index - 1 not in self.disable_corrector and self.last_sample is not None
-        )
+        assert self.step_index is not None
+        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)
+        model_output_convert = self.convert_model_output(model_output,
+                                                         sample=sample)
         if use_corrector:
             sample = self.multistep_uni_c_bh_update(
                 this_model_output=model_output_convert,
@@ -1017,16 +1064,19 @@ class UniPCMultistepScheduler(SchedulerMixin, ConfigMixin, BaseScheduler):
         self.timestep_list[-1] = timestep
 
         if self.config.lower_order_final:
-            this_order = min(self.config.solver_order, len(self.timesteps) - self.step_index)
+            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)  # warmup for multistep
+        self.this_order: int = min(this_order, self.lower_order_nums +
+                                   1)  # warmup for multistep
         assert self.this_order > 0
 
         self.last_sample = sample
         prev_sample = self.multistep_uni_p_bh_update(
-            model_output=model_output,  # pass the original non-converted model output, in case solver-p is used
+            model_output=
+            model_output,  # pass the original non-converted model output, in case solver-p is used
             sample=sample,
             order=self.this_order,
         )
@@ -1035,14 +1085,16 @@ class UniPCMultistepScheduler(SchedulerMixin, ConfigMixin, BaseScheduler):
             self.lower_order_nums += 1
 
         # upon completion increase step index by one
+        assert self._step_index is not None
         self._step_index += 1
 
         if not return_dict:
-            return (prev_sample,)
+            return (prev_sample, )
 
         return SchedulerOutput(prev_sample=prev_sample)
 
-    def scale_model_input(self, sample: torch.Tensor, *args, **kwargs) -> torch.Tensor:
+    def scale_model_input(self, sample: torch.Tensor, *args,
+                          **kwargs) -> torch.Tensor:
         """
         Ensures interchangeability with schedulers that need to scale the denoising model input depending on the
         current timestep.
@@ -1065,18 +1117,25 @@ class UniPCMultistepScheduler(SchedulerMixin, ConfigMixin, BaseScheduler):
         timesteps: torch.IntTensor,
     ) -> torch.Tensor:
         # Make sure sigmas and timesteps have the same device and dtype as original_samples
-        sigmas = self.sigmas.to(device=original_samples.device, dtype=original_samples.dtype)
-        if original_samples.device.type == "mps" and torch.is_floating_point(timesteps):
+        sigmas = self.sigmas.to(device=original_samples.device,
+                                dtype=original_samples.dtype)
+        if original_samples.device.type == "mps" and torch.is_floating_point(
+                timesteps):
             # mps does not support float64
-            schedule_timesteps = self.timesteps.to(original_samples.device, dtype=torch.float32)
-            timesteps = timesteps.to(original_samples.device, dtype=torch.float32)
+            schedule_timesteps = self.timesteps.to(original_samples.device,
+                                                   dtype=torch.float32)
+            timesteps = timesteps.to(original_samples.device,
+                                     dtype=torch.float32)
         else:
             schedule_timesteps = self.timesteps.to(original_samples.device)
             timesteps = timesteps.to(original_samples.device)
 
         # begin_index is None when the 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 timesteps]
+            step_indices = [
+                self.index_for_timestep(t, schedule_timesteps)
+                for t in timesteps
+            ]
         elif self.step_index is not None:
             # add_noise is called after first denoising step (for inpainting)
             step_indices = [self.step_index] * timesteps.shape[0]
@@ -1093,4 +1152,4 @@ class UniPCMultistepScheduler(SchedulerMixin, ConfigMixin, BaseScheduler):
         return noisy_samples
 
     def __len__(self):
-        return self.config.num_train_timesteps
+        return self.config.num_train_timesteps

fastvideo/models/vaes/wanvae.py

@@ -32,123 +32,22 @@ CACHE_T = 2
 is_first_frame = contextvars.ContextVar("is_first_frame", default=False)
 feat_cache = contextvars.ContextVar("feat_cache", default=None)
 feat_idx = contextvars.ContextVar("feat_idx", default=0)
-first_chunk = contextvars.ContextVar("first_chunk", default=None)
 
 
 @contextmanager
 def forward_context(first_frame_arg=False,
                     feat_cache_arg=None,
-                    feat_idx_arg=None,
-                    first_chunk_arg=None):
+                    feat_idx_arg=None):
     is_first_frame_token = is_first_frame.set(first_frame_arg)
     feat_cache_token = feat_cache.set(feat_cache_arg)
     feat_idx_token = feat_idx.set(feat_idx_arg)
-    first_chunk_token = first_chunk.set(first_chunk_arg)
     try:
         yield
     finally:
         is_first_frame.reset(is_first_frame_token)
         feat_cache.reset(feat_cache_token)
         feat_idx.reset(feat_idx_token)
-        first_chunk.reset(first_chunk_token)
 
-class AvgDown3D(nn.Module):
-    def __init__(
-        self,
-        in_channels,
-        out_channels,
-        factor_t,
-        factor_s=1,
-    ):
-        super().__init__()
-        self.in_channels = in_channels
-        self.out_channels = out_channels
-        self.factor_t = factor_t
-        self.factor_s = factor_s
-        self.factor = self.factor_t * self.factor_s * self.factor_s
-
-        assert in_channels * self.factor % out_channels == 0
-        self.group_size = in_channels * self.factor // out_channels
-
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        pad_t = (self.factor_t - x.shape[2] % self.factor_t) % self.factor_t
-        pad = (0, 0, 0, 0, pad_t, 0)
-        x = F.pad(x, pad)
-        B, C, T, H, W = x.shape
-        x = x.view(
-            B,
-            C,
-            T // self.factor_t,
-            self.factor_t,
-            H // self.factor_s,
-            self.factor_s,
-            W // self.factor_s,
-            self.factor_s,
-        )
-        x = x.permute(0, 1, 3, 5, 7, 2, 4, 6).contiguous()
-        x = x.view(
-            B,
-            C * self.factor,
-            T // self.factor_t,
-            H // self.factor_s,
-            W // self.factor_s,
-        )
-        x = x.view(
-            B,
-            self.out_channels,
-            self.group_size,
-            T // self.factor_t,
-            H // self.factor_s,
-            W // self.factor_s,
-        )
-        x = x.mean(dim=2)
-        return x
-
-
-class DupUp3D(nn.Module):
-    def __init__(
-        self,
-        in_channels: int,
-        out_channels: int,
-        factor_t,
-        factor_s=1,
-    ):
-        super().__init__()
-        self.in_channels = in_channels
-        self.out_channels = out_channels
-
-        self.factor_t = factor_t
-        self.factor_s = factor_s
-        self.factor = self.factor_t * self.factor_s * self.factor_s
-
-        assert out_channels * self.factor % in_channels == 0
-        self.repeats = out_channels * self.factor // in_channels
-
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        x = x.repeat_interleave(self.repeats, dim=1)
-        x = x.view(
-            x.size(0),
-            self.out_channels,
-            self.factor_t,
-            self.factor_s,
-            self.factor_s,
-            x.size(2),
-            x.size(3),
-            x.size(4),
-        )
-        x = x.permute(0, 1, 5, 2, 6, 3, 7, 4).contiguous()
-        x = x.view(
-            x.size(0),
-            self.out_channels,
-            x.size(2) * self.factor_t,
-            x.size(4) * self.factor_s,
-            x.size(6) * self.factor_s,
-        )
-
-        _first_chunk = first_chunk.get()
-        if _first_chunk:
-            x = x[:, :, self.factor_t - 1 :, :, :]
-        return x
 
 class WanCausalConv3d(nn.Conv3d):
     r"""
@@ -259,26 +158,20 @@ class WanResample(nn.Module):
             - 'downsample3d': 3D downsampling with zero-padding, convolution, and causal 3D convolution.
     """
 
-    def __init__(self, dim: int, mode: str, upsample_out_dim: int = None) -> None:
+    def __init__(self, dim: int, mode: str) -> None:
         super().__init__()
         self.dim = dim
         self.mode = mode
 
-        # default to dim //2
-        if upsample_out_dim is None:
-            upsample_out_dim = dim // 2
-
         # layers
         if mode == "upsample2d":
             self.resample = nn.Sequential(
                 WanUpsample(scale_factor=(2.0, 2.0), mode="nearest-exact"),
-                nn.Conv2d(dim, upsample_out_dim, 3, padding=1),
-            )
+                nn.Conv2d(dim, dim // 2, 3, padding=1))
         elif mode == "upsample3d":
             self.resample = nn.Sequential(
                 WanUpsample(scale_factor=(2.0, 2.0), mode="nearest-exact"),
-                nn.Conv2d(dim, upsample_out_dim, 3, padding=1),
-            )
+                nn.Conv2d(dim, dim // 2, 3, padding=1))
             self.time_conv = WanCausalConv3d(dim,
                                              dim * 2, (3, 1, 1),
                                              padding=(1, 0, 0))
@@ -551,40 +444,6 @@ class WanMidBlock(nn.Module):
 
         return x
 
-class WanResidualDownBlock(nn.Module):
-    def __init__(self, in_dim, out_dim, dropout, num_res_blocks, temperal_downsample=False, down_flag=False):
-        super().__init__()
-
-        # Shortcut path with downsample
-        self.avg_shortcut = AvgDown3D(
-            in_dim,
-            out_dim,
-            factor_t=2 if temperal_downsample else 1,
-            factor_s=2 if down_flag else 1,
-        )
-
-        # Main path with residual blocks and downsample
-        resnets = []
-        for _ in range(num_res_blocks):
-            resnets.append(WanResidualBlock(in_dim, out_dim, dropout))
-            in_dim = out_dim
-        self.resnets = nn.ModuleList(resnets)
-
-        # Add the final downsample block
-        if down_flag:
-            mode = "downsample3d" if temperal_downsample else "downsample2d"
-            self.downsampler = WanResample(out_dim, mode=mode)
-        else:
-            self.downsampler = None
-
-    def forward(self, x):
-        x_copy = x.clone()
-        for resnet in self.resnets:
-            x = resnet(x)
-        if self.downsampler is not None:
-            x = self.downsampler(x)
-
-        return x + self.avg_shortcut(x_copy)
 
 class WanEncoder3d(nn.Module):
     r"""
@@ -603,7 +462,6 @@ class WanEncoder3d(nn.Module):
 
     def __init__(
         self,
-        in_channels: int = 3,
         dim=128,
         z_dim=4,
         dim_mult=(1, 2, 4, 4),
@@ -612,7 +470,6 @@ class WanEncoder3d(nn.Module):
         temperal_downsample=(True, True, False),
         dropout=0.0,
         non_linearity: str = "silu",
-        is_residual: bool = False,  # wan 2.2 vae use a residual downblock
     ):
         super().__init__()
         self.dim = dim
@@ -629,36 +486,26 @@ class WanEncoder3d(nn.Module):
         scale = 1.0
 
         # init block
-        self.conv_in = WanCausalConv3d(in_channels, dims[0], 3, padding=1)
+        self.conv_in = WanCausalConv3d(3, dims[0], 3, padding=1)
 
         # downsample blocks
         self.down_blocks = nn.ModuleList([])
         for i, (in_dim,
                 out_dim) in enumerate(zip(dims[:-1], dims[1:], strict=True)):
             # residual (+attention) blocks
-            if is_residual:
+            for _ in range(num_res_blocks):
                 self.down_blocks.append(
-                    WanResidualDownBlock(
-                        in_dim,
-                        out_dim,
-                        dropout,
-                        num_res_blocks,
-                        temperal_downsample=temperal_downsample[i] if i != len(dim_mult) - 1 else False,
-                        down_flag=i != len(dim_mult) - 1,
-                    )
-                )
-            else:
-                for _ in range(num_res_blocks):
-                    self.down_blocks.append(WanResidualBlock(in_dim, out_dim, dropout))
-                    if scale in attn_scales:
-                        self.down_blocks.append(WanAttentionBlock(out_dim))
-                    in_dim = out_dim
+                    WanResidualBlock(in_dim, out_dim, dropout))
+                if scale in attn_scales:
+                    self.down_blocks.append(WanAttentionBlock(out_dim))
+                in_dim = out_dim
 
-                # downsample block
-                if i != len(dim_mult) - 1:
-                    mode = "downsample3d" if temperal_downsample[i] else "downsample2d"
-                    self.down_blocks.append(WanResample(out_dim, mode=mode))
-                    scale /= 2.0
+            # downsample block
+            if i != len(dim_mult) - 1:
+                mode = "downsample3d" if temperal_downsample[
+                    i] else "downsample2d"
+                self.down_blocks.append(WanResample(out_dim, mode=mode))
+                scale /= 2.0
 
         # middle blocks
         self.mid_block = WanMidBlock(out_dim,
@@ -725,83 +572,6 @@ class WanEncoder3d(nn.Module):
             x = self.conv_out(x)
         return x
 
-class WanResidualUpBlock(nn.Module):
-    """
-    A block that handles upsampling for the WanVAE decoder.
-    Args:
-        in_dim (int): Input dimension
-        out_dim (int): Output dimension
-        num_res_blocks (int): Number of residual blocks
-        dropout (float): Dropout rate
-        temperal_upsample (bool): Whether to upsample on temporal dimension
-        up_flag (bool): Whether to upsample or not
-        non_linearity (str): Type of non-linearity to use
-    """
-
-    def __init__(
-        self,
-        in_dim: int,
-        out_dim: int,
-        num_res_blocks: int,
-        dropout: float = 0.0,
-        temperal_upsample: bool = False,
-        up_flag: bool = False,
-        non_linearity: str = "silu",
-    ):
-        super().__init__()
-        self.in_dim = in_dim
-        self.out_dim = out_dim
-
-        if up_flag:
-            self.avg_shortcut = DupUp3D(
-                in_dim,
-                out_dim,
-                factor_t=2 if temperal_upsample else 1,
-                factor_s=2,
-            )
-        else:
-            self.avg_shortcut = None
-
-        # create residual blocks
-        resnets = []
-        current_dim = in_dim
-        for _ in range(num_res_blocks + 1):
-            resnets.append(WanResidualBlock(current_dim, out_dim, dropout, non_linearity))
-            current_dim = out_dim
-
-        self.resnets = nn.ModuleList(resnets)
-
-        # Add upsampling layer if needed
-        if up_flag:
-            upsample_mode = "upsample3d" if temperal_upsample else "upsample2d"
-            self.upsampler = WanResample(out_dim, mode=upsample_mode, upsample_out_dim=out_dim)
-        else:
-            self.upsampler = None
-
-        self.gradient_checkpointing = False
-
-    def forward(self, x):
-        """
-        Forward pass through the upsampling block.
-        Args:
-            x (torch.Tensor): Input tensor
-            feat_cache (list, optional): Feature cache for causal convolutions
-            feat_idx (list, optional): Feature index for cache management
-        Returns:
-            torch.Tensor: Output tensor
-        """
-        x_copy = x.clone()
-
-        for resnet in self.resnets:
-            x = resnet(x)
-
-        if self.upsampler is not None:
-            x = self.upsampler(x)
-
-        if self.avg_shortcut is not None:
-            x = x + self.avg_shortcut(x_copy)
-
-        return x
 
 class WanUpBlock(nn.Module):
     """
@@ -893,8 +663,6 @@ class WanDecoder3d(nn.Module):
         temperal_upsample=(False, True, True),
         dropout=0.0,
         non_linearity: str = "silu",
-        out_channels: int = 3,
-        is_residual: bool = False,
     ):
         super().__init__()
         self.dim = dim
@@ -909,6 +677,7 @@ class WanDecoder3d(nn.Module):
 
         # dimensions
         dims = [dim * u for u in [dim_mult[-1]] + dim_mult[::-1]]
+        scale = 1.0 / 2**(len(dim_mult) - 2)
 
         # init block
         self.conv_in = WanCausalConv3d(z_dim, dims[0], 3, padding=1)
@@ -924,44 +693,33 @@ class WanDecoder3d(nn.Module):
         for i, (in_dim,
                 out_dim) in enumerate(zip(dims[:-1], dims[1:], strict=True)):
             # residual (+attention) blocks
-            if i > 0 and not is_residual:
-                # wan vae 2.1
+            if i > 0:
                 in_dim = in_dim // 2
 
-            # determine if we need upsampling
-            up_flag = i != len(dim_mult) - 1
-            # determine upsampling mode, if not upsampling, set to None
+            # Determine if we need upsampling
             upsample_mode = None
-            if up_flag and temperal_upsample[i]:
-                upsample_mode = "upsample3d"
-            elif up_flag:
-                upsample_mode = "upsample2d"
+            if i != len(dim_mult) - 1:
+                upsample_mode = "upsample3d" if temperal_upsample[
+                    i] else "upsample2d"
 
             # Create and add the upsampling block
-            if is_residual:
-                up_block = WanResidualUpBlock(
-                    in_dim=in_dim,
-                    out_dim=out_dim,
-                    num_res_blocks=num_res_blocks,
-                    dropout=dropout,
-                    temperal_upsample=temperal_upsample[i] if up_flag else False,
-                    up_flag=up_flag,
-                    non_linearity=non_linearity,
-                )
-            else:
-                up_block = WanUpBlock(
-                    in_dim=in_dim,
-                    out_dim=out_dim,
-                    num_res_blocks=num_res_blocks,
-                    dropout=dropout,
-                    upsample_mode=upsample_mode,
-                    non_linearity=non_linearity,
-                )
+            up_block = WanUpBlock(
+                in_dim=in_dim,
+                out_dim=out_dim,
+                num_res_blocks=num_res_blocks,
+                dropout=dropout,
+                upsample_mode=upsample_mode,
+                non_linearity=non_linearity,
+            )
             self.up_blocks.append(up_block)
 
+            # Update scale for next iteration
+            if upsample_mode is not None:
+                scale *= 2.0
+
         # output blocks
         self.norm_out = WanRMS_norm(out_dim, images=False)
-        self.conv_out = WanCausalConv3d(out_dim, out_channels, 3, padding=1)
+        self.conv_out = WanCausalConv3d(out_dim, 3, 3, padding=1)
 
         self.gradient_checkpointing = False
 
@@ -1018,74 +776,6 @@ class WanDecoder3d(nn.Module):
             x = self.conv_out(x)
         return x
 
-def patchify(x, patch_size):
-    if patch_size == 1:
-        return x
-
-    if x.dim() == 4:
-        # x shape: [batch_size, channels, height, width]
-        batch_size, channels, height, width = x.shape
-
-        # Ensure height and width are divisible by patch_size
-        if height % patch_size != 0 or width % patch_size != 0:
-            raise ValueError(f"Height ({height}) and width ({width}) must be divisible by patch_size ({patch_size})")
-
-        # Reshape to [batch_size, channels, height//patch_size, patch_size, width//patch_size, patch_size]
-        x = x.view(batch_size, channels, height // patch_size, patch_size, width // patch_size, patch_size)
-
-        # Rearrange to [batch_size, channels * patch_size * patch_size, height//patch_size, width//patch_size]
-        x = x.permute(0, 1, 3, 5, 2, 4).contiguous()
-        x = x.view(batch_size, channels * patch_size * patch_size, height // patch_size, width // patch_size)
-
-    elif x.dim() == 5:
-        # x shape: [batch_size, channels, frames, height, width]
-        batch_size, channels, frames, height, width = x.shape
-
-        # Ensure height and width are divisible by patch_size
-        if height % patch_size != 0 or width % patch_size != 0:
-            raise ValueError(f"Height ({height}) and width ({width}) must be divisible by patch_size ({patch_size})")
-
-        # Reshape to [batch_size, channels, frames, height//patch_size, patch_size, width//patch_size, patch_size]
-        x = x.view(batch_size, channels, frames, height // patch_size, patch_size, width // patch_size, patch_size)
-
-        # Rearrange to [batch_size, channels * patch_size * patch_size, frames, height//patch_size, width//patch_size]
-        x = x.permute(0, 1, 4, 6, 2, 3, 5).contiguous()
-        x = x.view(batch_size, channels * patch_size * patch_size, frames, height // patch_size, width // patch_size)
-
-    else:
-        raise ValueError(f"Invalid input shape: {x.shape}")
-
-    return x
-
-def unpatchify(x, patch_size):
-    if patch_size == 1:
-        return x
-
-    if x.dim() == 4:
-        # x shape: [b, (c * patch_size * patch_size), h, w]
-        batch_size, c_patches, height, width = x.shape
-        channels = c_patches // (patch_size * patch_size)
-
-        # Reshape to [b, c, patch_size, patch_size, h, w]
-        x = x.view(batch_size, channels, patch_size, patch_size, height, width)
-
-        # Rearrange to [b, c, h * patch_size, w * patch_size]
-        x = x.permute(0, 1, 4, 2, 5, 3).contiguous()
-        x = x.view(batch_size, channels, height * patch_size, width * patch_size)
-
-    elif x.dim() == 5:
-        # x shape: [batch_size, (channels * patch_size * patch_size), frame, height, width]
-        batch_size, c_patches, frames, height, width = x.shape
-        channels = c_patches // (patch_size * patch_size)
-
-        # Reshape to [b, c, patch_size, patch_size, f, h, w]
-        x = x.view(batch_size, channels, patch_size, patch_size, frames, height, width)
-
-        # Rearrange to [b, c, f, h * patch_size, w * patch_size]
-        x = x.permute(0, 1, 4, 5, 2, 6, 3).contiguous()
-        x = x.view(batch_size, channels, frames, height * patch_size, width * patch_size)
-
-    return x
 
 class AutoencoderKLWan(nn.Module, ParallelTiledVAE):
     r"""
@@ -1105,43 +795,24 @@ class AutoencoderKLWan(nn.Module, ParallelTiledVAE):
         self.z_dim = config.z_dim
         self.temperal_downsample = list(config.temperal_downsample)
         self.temperal_upsample = list(config.temperal_downsample)[::-1]
-
-        if config.decoder_base_dim is None:
-            decoder_base_dim = config.base_dim
-        else:
-            decoder_base_dim = config.decoder_base_dim
-
         self.latents_mean = list(config.latents_mean)
         self.latents_std = list(config.latents_std)
         self.shift_factor = config.shift_factor
 
         if config.load_encoder:
-            self.encoder = WanEncoder3d(
-                in_channels=config.in_channels,
-                dim=config.base_dim, 
-                z_dim=self.z_dim * 2,
-                dim_mult=config.dim_mult, 
-                num_res_blocks=config.num_res_blocks,
-                attn_scales=config.attn_scales,
-                temperal_downsample=self.temperal_downsample,
-                dropout=config.dropout,
-                is_residual=config.is_residual,
-            )
+            self.encoder = WanEncoder3d(config.base_dim, self.z_dim * 2,
+                                        config.dim_mult, config.num_res_blocks,
+                                        config.attn_scales,
+                                        self.temperal_downsample,
+                                        config.dropout)
         self.quant_conv = WanCausalConv3d(self.z_dim * 2, self.z_dim * 2, 1)
         self.post_quant_conv = WanCausalConv3d(self.z_dim, self.z_dim, 1)
 
         if config.load_decoder:
-            self.decoder = WanDecoder3d(
-                dim=decoder_base_dim, 
-                z_dim=self.z_dim,
-                dim_mult=config.dim_mult, 
-                num_res_blocks=config.num_res_blocks,
-                attn_scales=config.attn_scales,
-                temperal_upsample=self.temperal_upsample, 
-                dropout=config.dropout,
-                out_channels=config.out_channels,
-                is_residual=config.is_residual,
-            )
+            self.decoder = WanDecoder3d(config.base_dim, self.z_dim,
+                                        config.dim_mult, config.num_res_blocks,
+                                        config.attn_scales,
+                                        self.temperal_upsample, config.dropout)
 
         self.use_feature_cache = config.use_feature_cache
 
@@ -1167,8 +838,6 @@ class AutoencoderKLWan(nn.Module, ParallelTiledVAE):
     def encode(self, x: torch.Tensor) -> torch.Tensor:
         if self.use_feature_cache:
             self.clear_cache()
-            if self.config.patch_size is not None:
-                x = patchify(x, patch_size=self.config.patch_size)
             with forward_context(feat_cache_arg=self._enc_feat_map,
                                  feat_idx_arg=self._enc_conv_idx):
                 t = x.shape[2]
@@ -1234,17 +903,12 @@ class AutoencoderKLWan(nn.Module, ParallelTiledVAE):
                 for i in range(iter_):
                     feat_idx.set(0)
                     if i == 0:
-                        first_chunk.set(True)
                         out = self.decoder(x[:, :, i:i + 1, :, :])
                     else:
-                        first_chunk.set(False)
                         out_ = self.decoder(x[:, :, i:i + 1, :, :])
                         out = torch.cat([out, out_], 2)
 
-            if self.config.clip_output:
-                out = torch.clamp(out, min=-1.0, max=1.0)
-            if self.config.patch_size is not None:
-                out = unpatchify(out, patch_size=self.config.patch_size)
+            out = torch.clamp(out, min=-1.0, max=1.0)
             self.clear_cache()
         else:
             out = ParallelTiledVAE.decode(self, z)

fastvideo/pipelines/basic/wan/wan_pipeline.py

@@ -62,7 +62,6 @@ class WanPipeline(LoRAPipeline, ComposedPipelineBase):
                        stage=DenoisingStage(
                            transformer=self.get_module("transformer"),
                            scheduler=self.get_module("scheduler"),
-                           vae=self.get_module("vae"),
                            pipeline=self))
 
         self.add_stage(stage_name="decoding_stage",

fastvideo/pipelines/composed_pipeline_base.py

@@ -234,9 +234,6 @@ class ComposedPipelineBase(ABC):
         # remove keys that are not pipeline modules
         model_index.pop("_class_name")
         model_index.pop("_diffusers_version")
-        # @TODO(Wei): Temporary hack
-        model_index.pop("boundary_ratio")
-        model_index.pop("expand_timesteps")
 
         # some sanity checks
         assert len(
@@ -258,8 +255,6 @@ class ComposedPipelineBase(ABC):
         modules = {}
         for module_name, (transformers_or_diffusers,
                           architecture) in model_index.items():
-            if transformers_or_diffusers is None:
-                continue
             if module_name not in required_modules:
                 logger.info("Skipping module %s", module_name)
                 continue

fastvideo/pipelines/stages/denoising.py

@@ -10,9 +10,7 @@ from collections.abc import Iterable
 from typing import Any
 
 import torch
-import torchvision.transforms.functional as TF
 from einops import rearrange
-from PIL import Image
 from tqdm.auto import tqdm
 
 from fastvideo.attention import get_attn_backend
@@ -32,7 +30,7 @@ from fastvideo.pipelines.stages.base import PipelineStage
 from fastvideo.pipelines.stages.validators import StageValidators as V
 from fastvideo.pipelines.stages.validators import VerificationResult
 from fastvideo.platforms import AttentionBackendEnum
-from fastvideo.utils import best_output_size, dict_to_3d_list, masks_like
+from fastvideo.utils import dict_to_3d_list, masks_like
 
 try:
     from fastvideo.attention.backends.sliding_tile_attn import (
@@ -59,11 +57,10 @@ class DenoisingStage(PipelineStage):
     the initial noise into the final output.
     """
 
-    def __init__(self, transformer, scheduler, vae, pipeline=None) -> None:
+    def __init__(self, transformer, scheduler, pipeline=None) -> None:
         super().__init__()
         self.transformer = transformer
         self.scheduler = scheduler
-        self.vae = vae
         self.pipeline = weakref.ref(pipeline) if pipeline else None
         attn_head_size = self.transformer.hidden_size // self.transformer.num_attention_heads
         self.attn_backend = get_attn_backend(
@@ -205,35 +202,7 @@ class DenoisingStage(PipelineStage):
                     # TI2V directly replaces the first frame of the latent with
                     # the image latent instead of appending along the channel dim
                     assert batch.image_latent is None, "TI2V task should not have image latents"
-                    # preprocess
-                    img = batch.pil_image
-                    assert img is not None
-                    ih, iw = img.height, img.width
-                    dh, dw = self.patch_size[1] * self.vae_stride[
-                        1], self.patch_size[2] * self.vae_stride[2]
-                    max_area = 704 * 1280
-                    ow, oh = best_output_size(iw, ih, dw, dh, max_area)
-
-                    scale = max(ow / iw, oh / ih)
-                    img = img.resize((round(iw * scale), round(ih * scale)),
-                                     Image.LANCZOS)
-
-                    # center-crop
-                    x1 = (img.width - ow) // 2
-                    y1 = (img.height - oh) // 2
-                    img = img.crop((x1, y1, x1 + ow, y1 + oh))
-                    assert img.width == ow and img.height == oh
-
-                    # to tensor
-                    img = TF.to_tensor(img).sub_(0.5).div_(0.5).to(
-                        self.device).unsqueeze(1)
-
-                    #                     return [
-                    #     self.model.encode(u.unsqueeze(0),
-                    #                       self.scale).float().squeeze(0)
-                    #     for u in videos
-                    # ]
-                    z = self.vae.encode(img.unsqueeze(0))
+                    z = self.get_module("vae").encode([batch.pil_image])
                     mask1, mask2 = masks_like([latent_model_input],
                                               zero=True,
                                               generator=batch.generator)

fastvideo/utils.py

@@ -846,29 +846,3 @@ def masks_like(tensor,
                 v[:, 0] = torch.zeros_like(v[:, 0])
 
     return out1, out2
-
-
-def best_output_size(w, h, dw, dh, expected_area):
-    # float output size
-    ratio = w / h
-    ow = (expected_area * ratio)**0.5
-    oh = expected_area / ow
-
-    # process width first
-    ow1 = int(ow // dw * dw)
-    oh1 = int(expected_area / ow1 // dh * dh)
-    assert ow1 % dw == 0 and oh1 % dh == 0 and ow1 * oh1 <= expected_area
-    ratio1 = ow1 / oh1
-
-    # process height first
-    oh2 = int(oh // dh * dh)
-    ow2 = int(expected_area / oh2 // dw * dw)
-    assert oh2 % dh == 0 and ow2 % dw == 0 and ow2 * oh2 <= expected_area
-    ratio2 = ow2 / oh2
-
-    # compare ratios
-    if max(ratio / ratio1, ratio1 / ratio) < max(ratio / ratio2,
-                                                 ratio2 / ratio):
-        return ow1, oh1
-    else:
-        return ow2, oh2