Upload run_wan_I2V_FB.py

run_wan_I2V_FB.py

@@ -0,0 +1,88 @@
+import torch
+import torch.distributed as dist
+import numpy as np
+import time
+from diffusers import AutoencoderKLWan, WanImageToVideoPipeline
+from diffusers.utils import export_to_video, load_image
+from transformers import CLIPVisionModel
+from PIL import Image
+from diffusers.schedulers.scheduling_unipc_multistep import UniPCMultistepScheduler
+import io
+
+start_loading = time.time()
+data_type = torch.bfloat16
+
+# Available models: Wan-AI/Wan2.1-I2V-14B-480P-Diffusers, Wan-AI/Wan2.1-I2V-14B-720P-Diffusers
+model_id = "Wan-AI/Wan2.1-I2V-14B-720P-Diffusers"
+
+image_encoder = CLIPVisionModel.from_pretrained(
+    model_id, subfolder="image_encoder", torch_dtype=data_type
+)
+vae = AutoencoderKLWan.from_pretrained(
+    model_id, subfolder="vae", torch_dtype=torch.float32
+)
+pipe = WanImageToVideoPipeline.from_pretrained(
+    model_id, vae=vae, image_encoder=image_encoder, torch_dtype=data_type
+)
+pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config, flow_shift=5.0)
+pipe.to("cuda")
+
+from para_attn.first_block_cache.diffusers_adapters import apply_cache_on_pipe
+
+apply_cache_on_pipe(pipe , residual_diff_threshold=0.1)
+
+print("Pipeline Loaded.....")
+end_loading = time.time()
+loading_time = end_loading - start_loading
+
+prompt = (
+    "Cars racing in slow motion"
+)
+negative_prompt = (
+    "bright colors, overexposed, static, blurred details, subtitles, style, artwork, painting, picture, still, overall gray, worst quality, low quality, JPEG compression residue, ugly, incomplete, extra fingers, poorly drawn hands, poorly drawn faces, deformed, disfigured, malformed limbs, fused fingers, still picture, cluttered background, three legs, many people in the background, walking backwards"
+)
+
+image = load_image(
+    "https://storage.googleapis.com/falserverless/gallery/car_720p.png"
+)
+
+# max_area = 1024 * 1024
+# aspect_ratio = image.height / image.width
+# mod_value = pipe.vae_scale_factor_spatial * pipe.transformer.config.patch_size[1]
+# print(f"MOD VALUE :{mod_value}")
+# height = round(np.sqrt(max_area * aspect_ratio)) // mod_value * mod_value
+# width = round(np.sqrt(max_area / aspect_ratio)) // mod_value * mod_value
+width , height = image.size
+# image = image.resize((width, height))
+
+# Start measuring inference time
+start_inference = time.time()
+
+# Run the pipeline
+output = pipe(
+    image=image,
+    prompt=prompt,
+    negative_prompt=negative_prompt,
+    height=height,
+    width=width,
+    num_frames=81,
+    guidance_scale=5.0,
+    num_inference_steps=30,
+    output_type="pil",
+).frames[0]
+
+# End of inference time measurement
+end_inference = time.time()
+inference_time = end_inference - start_inference
+
+# Save output and print timing info
+# if dist.get_rank() == 0:
+print(f"{'=' * 50}")
+print(f"Device: {torch.cuda.get_device_name()}")
+print(f"Number of GPUs: {dist.get_world_size()}")
+print(f"Pipeline Loading Time: {loading_time:.2f} seconds")
+print(f"Pipeline Inference Time: {inference_time:.2f} seconds")
+print(f"Resolution: {width}x{height}")
+print(f"{'=' * 50}")
+print("Saving video to wan-i2v.mp4")
+export_to_video(output, "wan-i2v_fb.mp4", fps=16)