checkpoint

examples/inference/basic/basic.py

@@ -1,17 +1,18 @@
 from fastvideo import VideoGenerator
 
-# from fastvideo.configs.sample import SamplingParam
+from fastvideo.configs.sample import SamplingParam
 
-OUTPUT_PATH = "video_samples"
+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(
-        "Wan-AI/Wan2.1-T2V-1.3B-Diffusers",
+        model_id,
         # FastVideo will automatically handle distributed setup
-        num_gpus=2,
+        num_gpus=1,
         use_fsdp_inference=True,
         dit_cpu_offload=False,
         vae_cpu_offload=False,
@@ -19,27 +20,22 @@ def main():
         # image_encoder_cpu_offload=False,
     )
 
-    # sampling_param = SamplingParam.from_pretrained("Wan-AI/Wan2.1-T2V-1.3B-Diffusers")
+    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 = (
-        "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)
+    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)
+        "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__":

examples/inference/basic/basic_i2v.py

@@ -0,0 +1,42 @@
+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/pipelines/basic/wan/wan_pipeline.py

@@ -62,6 +62,7 @@ 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/stages/denoising.py

@@ -10,7 +10,9 @@ 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
@@ -30,7 +32,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 dict_to_3d_list, masks_like
+from fastvideo.utils import best_output_size, dict_to_3d_list, masks_like
 
 try:
     from fastvideo.attention.backends.sliding_tile_attn import (
@@ -57,10 +59,11 @@ class DenoisingStage(PipelineStage):
     the initial noise into the final output.
     """
 
-    def __init__(self, transformer, scheduler, pipeline=None) -> None:
+    def __init__(self, transformer, scheduler, vae, 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(
@@ -202,7 +205,35 @@ 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"
-                    z = self.get_module("vae").encode([batch.pil_image])
+                    # 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))
                     mask1, mask2 = masks_like([latent_model_input],
                                               zero=True,
                                               generator=batch.generator)

fastvideo/utils.py

@@ -846,3 +846,29 @@ 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