Upload run_wan_I2V_CP+FB_GroupOffload.py

run_wan_I2V_CP+FB_GroupOffload.py

@@ -0,0 +1,140 @@
+import torch
+import torch.distributed as dist
+import numpy as np
+import time
+from diffusers import AutoencoderKLWan, WanImageToVideoPipeline, UMT5EncoderModel, WanTransformer3DModel
+from diffusers.utils import export_to_video, load_image
+from transformers import CLIPVisionModel, UMT5EncoderModel
+from PIL import Image
+from diffusers.schedulers.scheduling_unipc_multistep import UniPCMultistepScheduler
+import io
+from diffusers.hooks.group_offloading import apply_group_offloading
+
+dist.init_process_group()
+
+torch.cuda.set_device(dist.get_rank())
+
+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
+# )
+
+image_encoder = CLIPVisionModel.from_pretrained(
+    model_id, subfolder="image_encoder", torch_dtype=torch.float32
+)
+
+text_encoder = UMT5EncoderModel.from_pretrained(model_id, subfolder="text_encoder", torch_dtype=torch.bfloat16)
+vae = AutoencoderKLWan.from_pretrained(model_id, subfolder="vae", torch_dtype=torch.float32)
+transformer = WanTransformer3DModel.from_pretrained(model_id, subfolder="transformer", torch_dtype=torch.bfloat16)
+
+
+onload_device = torch.device("cuda")
+offload_device = torch.device("cpu")
+
+apply_group_offloading(text_encoder,
+    onload_device=onload_device,
+    offload_device=offload_device,
+    offload_type="block_level",
+    num_blocks_per_group=4
+)
+# pipe = WanImageToVideoPipeline.from_pretrained(
+#     model_id, vae=vae, image_encoder=image_encoder, torch_dtype=data_type
+# )
+
+pipe = WanImageToVideoPipeline.from_pretrained(
+    model_id,
+    vae=vae,
+    transformer=transformer,
+    text_encoder=text_encoder,
+    image_encoder=image_encoder,
+    torch_dtype=torch.bfloat16
+)
+
+pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config, flow_shift=5.0)
+pipe.to("cuda")
+
+# Import and apply parallel attention
+from para_attn.context_parallel import init_context_parallel_mesh
+from para_attn.context_parallel.diffusers_adapters import parallelize_pipe
+
+parallelize_pipe(
+    pipe,
+    mesh=init_context_parallel_mesh(
+        pipe.device.type,
+    ),
+)
+end_loading = time.time()
+
+from para_attn.first_block_cache.diffusers_adapters import apply_cache_on_pipe
+apply_cache_on_pipe(pipe , residual_diff_threshold=0.05)
+
+# pipe.enable_model_cpu_offload(gpu_id=dist.get_rank())
+# pipe.enable_vae_tiling()
+
+# torch._inductor.config.reorder_for_compute_comm_overlap = True
+# pipe.transformer = torch.compile(pipe.transformer, mode="max-autotune-no-cudagraphs")
+print("Pipeline Loaded.....")
+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" if dist.get_rank() == 0 else "pt",
+).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.mp4", fps=16)
+
+dist.destroy_process_group()