diff --git a/examples/inference/basic/basic.py b/examples/inference/basic/basic.py index 337b89550..9174c6abf 100644 --- a/examples/inference/basic/basic.py +++ b/examples/inference/basic/basic.py @@ -1,6 +1,6 @@ from fastvideo import VideoGenerator -# from fastvideo.configs.sample import SamplingParam +from fastvideo.configs.sample import SamplingParam OUTPUT_PATH = "video_samples" def main(): @@ -8,30 +8,35 @@ def main(): # model. # If a local path is provided, FastVideo will make a best effort # attempt to identify the optimal arguments. + model_name = "Wan-AI/Wan2.2-TI2V-5B-Diffusers" generator = VideoGenerator.from_pretrained( - "Wan-AI/Wan2.1-T2V-1.3B-Diffusers", + model_name, # 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=True, + text_encoder_cpu_offload=False, # Set pin_cpu_memory to false if CPU RAM is limited and there're no frequent CPU-GPU transfer - pin_cpu_memory=False, + pin_cpu_memory=True, # image_encoder_cpu_offload=False, ) - # 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" + sampling_param = SamplingParam.from_pretrained(model_name) + sampling_param.image_path = "test.jpg" + # sampling_param.num_inference_steps = 0 # 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." - ) - video = generator.generate_video(prompt, output_path=OUTPUT_PATH, save_video=True) + i2v_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." + i2v_prompt = "A little girl is packing a suitcase and the contents starts flying out of the suitcase everywhere." + prompt = i2v_prompt + # 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." + # ) + 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/") + return # Generate another video with a different prompt, without reloading the # model! diff --git a/fastvideo/entrypoints/video_generator.py b/fastvideo/entrypoints/video_generator.py index 70270f9e0..d5668cf8e 100644 --- a/fastvideo/entrypoints/video_generator.py +++ b/fastvideo/entrypoints/video_generator.py @@ -278,6 +278,7 @@ class VideoGenerator: width: {target_width} video_length: {sampling_param.num_frames} prompt: {prompt} + image_path: {sampling_param.image_path} neg_prompt: {sampling_param.negative_prompt} seed: {sampling_param.seed} infer_steps: {sampling_param.num_inference_steps} diff --git a/fastvideo/models/dits/wanvideo.py b/fastvideo/models/dits/wanvideo.py index 2ca840622..70d83f190 100644 --- a/fastvideo/models/dits/wanvideo.py +++ b/fastvideo/models/dits/wanvideo.py @@ -82,6 +82,7 @@ class WanTimeTextImageEmbedding(nn.Module): encoder_hidden_states: torch.Tensor, encoder_hidden_states_image: torch.Tensor | None = None, ): + logger.info(f"WTF timestep shape: {timestep.shape}") temb = self.time_embedder(timestep) timestep_proj = self.time_modulation(temb) 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 c95c68cca..4141ca718 100644 --- a/fastvideo/pipelines/stages/denoising.py +++ b/fastvideo/pipelines/stages/denoising.py @@ -10,6 +10,7 @@ from collections.abc import Iterable from typing import Any import torch +import torchvision.transforms.functional as TF from einops import rearrange from tqdm.auto import tqdm @@ -30,7 +31,8 @@ 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 dict_to_3d_list, masks_like, best_output_size +from PIL import Image 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=None, 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( @@ -184,8 +187,43 @@ class DenoisingStage(PipelineStage): assert neg_prompt_embeds is not None assert torch.isnan(neg_prompt_embeds[0]).sum() == 0 + latent_model_input = latents.to(target_dtype) + assert latent_model_input.shape[0] == 1, "only support batch size 1" + + if fastvideo_args.pipeline_config.ti2v_task and batch.pil_image is not None: + logger.info("===========Using TI2V task===========") + # 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" + assert self.vae is not None, "VAE is not provided for TI2V task" + z = self.vae.encode(batch.pil_image).mean.float() + logger.info(f"z shape: {z.shape}") + logger.info(f"latent_model_input shape: {latent_model_input.shape}") + latent_model_input = latent_model_input.squeeze(0) + mask1, mask2 = masks_like([latent_model_input], + zero=True) + # logger.info(f"mask1 shape: {mask1.shape}") + # logger.info(f"mask2 shape: {mask2.shape}") + latent_model_input = ( + 1. - mask2[0]) * z + mask2[0] * latent_model_input + # latent_model_input = latent_model_input.unsqueeze(0) + latent_model_input = latent_model_input.to(get_local_torch_device()) + latents = latent_model_input + F = batch.num_frames + temporal_scale = fastvideo_args.pipeline_config.vae_config.arch_config.scale_factor_temporal + spatial_scale = fastvideo_args.pipeline_config.vae_config.arch_config.scale_factor_spatial + patch_size = fastvideo_args.pipeline_config.dit_config.arch_config.patch_size + seq_len = ((F - 1) // temporal_scale + 1) * ( + batch.height // spatial_scale) * (batch.width // spatial_scale) // ( + patch_size[1] * patch_size[2]) + import math + seq_len = int(math.ceil(seq_len / sp_world_size)) * sp_world_size + logger.info("latents shape: %s", latents.shape) + # Run denoising loop with self.progress_bar(total=num_inference_steps) as progress_bar: + # logger.info(f"seq_len: {seq_len}") + logger.info(f"init timesteps: {timesteps}") for i, t in enumerate(timesteps): # Skip if interrupted if hasattr(self, 'interrupt') and self.interrupt: @@ -198,22 +236,35 @@ class DenoisingStage(PipelineStage): latent_model_input = torch.cat( [latent_model_input, batch.image_latent], dim=1).to(target_dtype) - elif fastvideo_args.pipeline_config.ti2v_task and batch.pil_image is not None: - # 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]) - mask1, mask2 = masks_like([latent_model_input], - zero=True, - generator=batch.generator) - latent_model_input = ( - 1. - mask2[0]) * z[0] + mask2[0] * latent_model_input + if fastvideo_args.pipeline_config.ti2v_task and batch.pil_image is not None: + logger.info(f"before ti2v timestep: {t}") + timestep = [t] + timestep = torch.stack(timestep).to(get_local_torch_device()) + + logger.info(f"mask2 shape: {mask2[0].shape}") + logger.info(f"mask[0][0] shape: {mask2[0][0].shape}") + logger.info(f"mask[0][0][:, ::2, ::2] shape: {mask2[0][0][:, ::2, ::2].shape}") + temp_ts = (mask2[0][0][:, ::2, ::2] * timestep).flatten() + logger.info(f"temp_ts: {temp_ts}") + logger.info(f"temp_ts shape: {temp_ts.shape}") + temp_ts = torch.cat([ + temp_ts, + temp_ts.new_ones(seq_len - temp_ts.size(0)) * timestep + ]) + # timestep = temp_ts.unsqueeze(0) + timestep = temp_ts + logger.info(f"after ti2v timestep: {timestep}") + t = timestep + # else: + t_expand = t.repeat(latent_model_input.shape[0]) + logger.info(f"t_expand shape: {t_expand.shape}") + logger.info(f"t_expand: {t_expand}") + assert torch.isnan(latent_model_input).sum() == 0 latent_model_input = self.scheduler.scale_model_input( latent_model_input, t) # Prepare inputs for transformer - t_expand = t.repeat(latent_model_input.shape[0]) guidance_expand = ( torch.tensor( [fastvideo_args.pipeline_config.embedded_cfg_scale] * @@ -314,6 +365,10 @@ class DenoisingStage(PipelineStage): latents, **extra_step_kwargs, return_dict=False)[0] + if fastvideo_args.pipeline_config.ti2v_task and batch.pil_image is not None: + latents = latents.squeeze(0) + latents = (1. - mask2[0]) * z + mask2[0] * latents + # latents = latents.unsqueeze(0) # Update progress bar if i == len(timesteps) - 1 or ( diff --git a/fastvideo/pipelines/stages/input_validation.py b/fastvideo/pipelines/stages/input_validation.py index 466ece75c..d69b3523c 100644 --- a/fastvideo/pipelines/stages/input_validation.py +++ b/fastvideo/pipelines/stages/input_validation.py @@ -12,6 +12,9 @@ from fastvideo.pipelines.pipeline_batch_info import ForwardBatch from fastvideo.pipelines.stages.base import PipelineStage from fastvideo.pipelines.stages.validators import (StageValidators, VerificationResult) +from fastvideo.utils import best_output_size +from PIL import Image +import torchvision.transforms.functional as TF logger = init_logger(__name__) @@ -100,6 +103,35 @@ class InputValidationStage(PipelineStage): image = load_image(batch.image_path) batch.pil_image = image + img = batch.pil_image + ih, iw = img.height, img.width + logger.info(f"img height: {ih}, img width: {iw}") + patch_size = fastvideo_args.pipeline_config.dit_config.arch_config.patch_size + vae_stride = fastvideo_args.pipeline_config.vae_config.arch_config.scale_factor_spatial + logger.info(f"patch_size: {patch_size}, vae_stride: {vae_stride}") + dh, dw = patch_size[1] * vae_stride, patch_size[2] * vae_stride + 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) + logger.info(f"resized img height: {img.height}, img width: {img.width}") + + # 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 + # logger.info(f"img shape: {img.shape}") + + # to tensor + img = TF.to_tensor(img).sub_(0.5).div_(0.5).to(self.device).unsqueeze(1) + logger.info(f"img shape: {img.shape}") + img = img.unsqueeze(0) + batch.height = oh + batch.width = ow + batch.pil_image = img + return batch def verify_input(self, batch: ForwardBatch,