diff --git a/examples/inference/basic/basic.py b/examples/inference/basic/basic.py index 37f15b516..f276b6552 100644 --- a/examples/inference/basic/basic.py +++ b/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__": diff --git a/examples/inference/basic/basic_i2v.py b/examples/inference/basic/basic_i2v.py new file mode 100644 index 000000000..b61748c8f --- /dev/null +++ b/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() diff --git a/fastvideo/pipelines/basic/wan/wan_pipeline.py b/fastvideo/pipelines/basic/wan/wan_pipeline.py index 8574c2d29..3f3b9197e 100644 --- a/fastvideo/pipelines/basic/wan/wan_pipeline.py +++ b/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", diff --git a/fastvideo/pipelines/stages/denoising.py b/fastvideo/pipelines/stages/denoising.py index d1e228cee..2013d6344 100644 --- a/fastvideo/pipelines/stages/denoising.py +++ b/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) diff --git a/fastvideo/utils.py b/fastvideo/utils.py index e6fc70552..0117a7f9b 100644 --- a/fastvideo/utils.py +++ b/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