Update Readme

config/qwenimage21/qwenimage21_control.yaml

@@ -1,7 +0,0 @@
-format: diffusers
-pipeline: qwenimage21
-transformer_additional_kwargs:
-    # Dense control: inject a skip at every 2nd block (16 of 32 layers), matching the model's built-in
-    # `control_layers=None` default. Must contain 0. Dial back (e.g. [0, 4, 8, ...]) to shrink the trainable adapter.
-    control_layers: [0, 2, 4, 6, 8, 10, 12, 14, 16, 18, 20, 22, 24, 26, 28, 30]
-    control_in_dim: 129

examples/cogvideox_fun/app.py

@@ -10,9 +10,9 @@ for project_root in project_roots:
     sys.path.insert(0, project_root) if project_root not in sys.path else None
 
 from videox_fun.api.api import (infer_forward_api,
-                                update_diffusion_transformer_api)
-from videox_fun.ui.cogvideox_fun_ui import ui, ui_client, ui_host
+                               update_diffusion_transformer_api)
 from videox_fun.ui.controller import ddpm_scheduler_dict
+from videox_fun.ui.cogvideox_fun_ui import ui, ui_client, ui_host
 
 if __name__ == "__main__":
     # Choose the ui mode  

examples/cogvideox_fun/launch_api.py

@@ -4,6 +4,7 @@ import sys
 import time
 
 import gradio as gr
+import ray
 import torch
 
 current_file_path = os.path.abspath(__file__)
@@ -12,10 +13,9 @@ for project_root in project_roots:
     sys.path.insert(0, project_root) if project_root not in sys.path else None
 
 from videox_fun.api.api_multi_nodes import (MultiNodesEngine,
-                                            multi_nodes_infer_forward_api)
-from videox_fun.ui.cogvideox_fun_ui import CogVideoXFunController
+                                           multi_nodes_infer_forward_api)
 from videox_fun.ui.controller import flow_scheduler_dict
-
+from videox_fun.ui.cogvideox_fun_ui import CogVideoXFunController
 
 def main():
     parser = argparse.ArgumentParser(description='xDiT HTTP Service')

examples/cogvideox_fun/predict_i2v.py

@@ -16,14 +16,18 @@ for project_root in project_roots:
 
 from videox_fun.dist import set_multi_gpus_devices, shard_model
 from videox_fun.models import (AutoencoderKLCogVideoX,
-                               CogVideoXTransformer3DModel, T5EncoderModel,
-                               T5Tokenizer)
+                              CogVideoXTransformer3DModel, T5EncoderModel,
+                              T5Tokenizer)
 from videox_fun.pipeline import (CogVideoXFunInpaintPipeline,
-                                 CogVideoXFunPipeline)
-from videox_fun.utils import (apply_gpu_memory_mode, get_image_to_video_latent,
-                              merge_lora, save_videos_grid, unmerge_lora)
+                                CogVideoXFunPipeline)
+from videox_fun.utils import (register_auto_device_hook,
+                              safe_enable_group_offload)
+from videox_fun.utils.fp8_optimization import (convert_model_weight_to_float8, replace_parameters_by_name,
+                                              convert_weight_dtype_wrapper)
+from videox_fun.utils.lora_utils import merge_lora, unmerge_lora
+from videox_fun.utils.utils import get_image_to_video_latent, save_videos_grid
 
-# GPU memory mode, which can be chosen in [model_full_load, model_full_load_and_qfloat8, model_cpu_offload, model_cpu_offload_and_qfloat8, model_group_offload, sequential_cpu_offload].
+# GPU memory mode, which can be chosen in [model_full_load, model_full_load_and_qfloat8, model_cpu_offload, model_cpu_offload_and_qfloat8, sequential_cpu_offload].
 # model_full_load means that the entire model will be moved to the GPU.
 # 
 # model_full_load_and_qfloat8 means that the entire model will be moved to the GPU,
@@ -76,7 +80,7 @@ partial_video_length = None
 overlap_video_length = 4
 
 # Use torch.float16 if GPU does not support torch.bfloat16
-# Some graphics cards, such as v100, 2080ti, do not support torch.bfloat16
+# ome graphics cards, such as v100, 2080ti, do not support torch.bfloat16
 weight_dtype            = torch.bfloat16
 # If you want to generate from text, please set the validation_image_start = None and validation_image_end = None
 validation_image_start  = "asset/1.png"
@@ -103,7 +107,7 @@ transformer = CogVideoXTransformer3DModel.from_pretrained(
 if transformer_path is not None:
     print(f"From checkpoint: {transformer_path}")
     if transformer_path.endswith("safetensors"):
-        from safetensors.torch import load_file
+        from safetensors.torch import load_file, safe_open
         state_dict = load_file(transformer_path)
     else:
         state_dict = torch.load(transformer_path, map_location="cpu")
@@ -121,7 +125,7 @@ vae = AutoencoderKLCogVideoX.from_pretrained(
 if vae_path is not None:
     print(f"From checkpoint: {vae_path}")
     if vae_path.endswith("safetensors"):
-        from safetensors.torch import load_file
+        from safetensors.torch import load_file, safe_open
         state_dict = load_file(vae_path)
     else:
         state_dict = torch.load(vae_path, map_location="cpu")
@@ -139,7 +143,7 @@ text_encoder = T5EncoderModel.from_pretrained(
 )
 
 # Get Scheduler
-Chosen_Scheduler = {
+Chosen_Scheduler = scheduler_dict = {
     "Euler": EulerDiscreteScheduler,
     "Euler A": EulerAncestralDiscreteScheduler,
     "DPM++": DPMSolverMultistepScheduler, 
@@ -185,10 +189,23 @@ if compile_dit:
         pipeline.transformer.transformer_blocks[i] = torch.compile(pipeline.transformer.transformer_blocks[i])
     print("Add Compile")
 
-# Quantize (when the mode carries an "_and_<quant>" suffix) and then place the pipeline.
-# The order lives inside the helper: quantization has to happen before the offload hooks
-# are installed.
-apply_gpu_memory_mode(pipeline, GPU_memory_mode, device, weight_dtype, exclude_module_name=[])
+if GPU_memory_mode == "sequential_cpu_offload":
+    pipeline.enable_sequential_cpu_offload(device=device)
+elif GPU_memory_mode == "model_group_offload":
+    register_auto_device_hook(pipeline.transformer)
+    safe_enable_group_offload(pipeline, onload_device=device, offload_device="cpu", offload_type="leaf_level", use_stream=True)
+elif GPU_memory_mode == "model_cpu_offload_and_qfloat8":
+    convert_model_weight_to_float8(transformer, exclude_module_name=[], device=device)
+    convert_weight_dtype_wrapper(transformer, weight_dtype)
+    pipeline.enable_model_cpu_offload(device=device)
+elif GPU_memory_mode == "model_cpu_offload":
+    pipeline.enable_model_cpu_offload(device=device)
+elif GPU_memory_mode == "model_full_load_and_qfloat8":
+    convert_model_weight_to_float8(transformer, exclude_module_name=[], device=device)
+    convert_weight_dtype_wrapper(transformer, weight_dtype)
+    pipeline.to(device=device)
+else:
+    pipeline.to(device=device)
 
 generator = torch.Generator(device=device).manual_seed(seed)
 

examples/cogvideox_fun/predict_t2v.py

@@ -15,16 +15,20 @@ project_roots = [os.path.dirname(current_file_path), os.path.dirname(os.path.dir
 for project_root in project_roots:
     sys.path.insert(0, project_root) if project_root not in sys.path else None
 
-from videox_fun.dist import set_multi_gpus_devices, shard_model
 from videox_fun.models import (AutoencoderKLCogVideoX,
-                               CogVideoXTransformer3DModel, T5EncoderModel,
-                               T5Tokenizer)
-from videox_fun.pipeline import (CogVideoXFunInpaintPipeline,
-                                 CogVideoXFunPipeline)
-from videox_fun.utils import (apply_gpu_memory_mode, get_image_to_video_latent,
-                              merge_lora, save_videos_grid, unmerge_lora)
+                              CogVideoXTransformer3DModel, T5EncoderModel,
+                              T5Tokenizer)
+from videox_fun.pipeline import (CogVideoXFunPipeline,
+                                CogVideoXFunInpaintPipeline)
+from videox_fun.utils import (register_auto_device_hook,
+                              safe_enable_group_offload)
+from videox_fun.utils.fp8_optimization import (convert_model_weight_to_float8, replace_parameters_by_name,
+                                              convert_weight_dtype_wrapper)
+from videox_fun.utils.lora_utils import merge_lora, unmerge_lora
+from videox_fun.utils.utils import get_image_to_video_latent, save_videos_grid
+from videox_fun.dist import set_multi_gpus_devices, shard_model
 
-# GPU memory mode, which can be chosen in [model_full_load, model_full_load_and_qfloat8, model_cpu_offload, model_cpu_offload_and_qfloat8, model_group_offload, sequential_cpu_offload].
+# GPU memory mode, which can be chosen in [model_full_load, model_full_load_and_qfloat8, model_cpu_offload, model_cpu_offload_and_qfloat8, sequential_cpu_offload].
 # model_full_load means that the entire model will be moved to the GPU.
 # 
 # model_full_load_and_qfloat8 means that the entire model will be moved to the GPU,
@@ -73,7 +77,7 @@ video_length        = 49
 fps                 = 8
 
 # Use torch.float16 if GPU does not support torch.bfloat16
-# Some graphics cards, such as v100, 2080ti, do not support torch.bfloat16
+# ome graphics cards, such as v100, 2080ti, do not support torch.bfloat16
 weight_dtype        = torch.bfloat16
 prompt              = "A young woman with beautiful and clear eyes and blonde hair standing and white dress in a forest wearing a crown. She seems to be lost in thought, and the camera focuses on her face. The video is of high quality, and the view is very clear. High quality, masterpiece, best quality, highres, ultra-detailed, fantastic."
 negative_prompt     = "The video is not of a high quality, it has a low resolution. Watermark present in each frame. The background is solid. Strange body and strange trajectory. Distortion. "
@@ -95,7 +99,7 @@ transformer = CogVideoXTransformer3DModel.from_pretrained(
 if transformer_path is not None:
     print(f"From checkpoint: {transformer_path}")
     if transformer_path.endswith("safetensors"):
-        from safetensors.torch import load_file
+        from safetensors.torch import load_file, safe_open
         state_dict = load_file(transformer_path)
     else:
         state_dict = torch.load(transformer_path, map_location="cpu")
@@ -113,7 +117,7 @@ vae = AutoencoderKLCogVideoX.from_pretrained(
 if vae_path is not None:
     print(f"From checkpoint: {vae_path}")
     if vae_path.endswith("safetensors"):
-        from safetensors.torch import load_file
+        from safetensors.torch import load_file, safe_open
         state_dict = load_file(vae_path)
     else:
         state_dict = torch.load(vae_path, map_location="cpu")
@@ -131,7 +135,7 @@ text_encoder = T5EncoderModel.from_pretrained(
 )
 
 # Get Scheduler
-Chosen_Scheduler = {
+Chosen_Scheduler = scheduler_dict = {
     "Euler": EulerDiscreteScheduler,
     "Euler A": EulerAncestralDiscreteScheduler,
     "DPM++": DPMSolverMultistepScheduler, 
@@ -177,10 +181,23 @@ if compile_dit:
         pipeline.transformer.transformer_blocks[i] = torch.compile(pipeline.transformer.transformer_blocks[i])
     print("Add Compile")
 
-# Quantize (when the mode carries an "_and_<quant>" suffix) and then place the pipeline.
-# The order lives inside the helper: quantization has to happen before the offload hooks
-# are installed.
-apply_gpu_memory_mode(pipeline, GPU_memory_mode, device, weight_dtype, exclude_module_name=[])
+if GPU_memory_mode == "sequential_cpu_offload":
+    pipeline.enable_sequential_cpu_offload(device=device)
+elif GPU_memory_mode == "model_group_offload":
+    register_auto_device_hook(pipeline.transformer)
+    safe_enable_group_offload(pipeline, onload_device=device, offload_device="cpu", offload_type="leaf_level", use_stream=True)
+elif GPU_memory_mode == "model_cpu_offload_and_qfloat8":
+    convert_model_weight_to_float8(transformer, exclude_module_name=[], device=device)
+    convert_weight_dtype_wrapper(transformer, weight_dtype)
+    pipeline.enable_model_cpu_offload(device=device)
+elif GPU_memory_mode == "model_cpu_offload":
+    pipeline.enable_model_cpu_offload(device=device)
+elif GPU_memory_mode == "model_full_load_and_qfloat8":
+    convert_model_weight_to_float8(transformer, exclude_module_name=[], device=device)
+    convert_weight_dtype_wrapper(transformer, weight_dtype)
+    pipeline.to(device=device)
+else:
+    pipeline.to(device=device)
 
 generator = torch.Generator(device=device).manual_seed(seed)
 

examples/cogvideox_fun/predict_v2v.py

@@ -14,16 +14,20 @@ project_roots = [os.path.dirname(current_file_path), os.path.dirname(os.path.dir
 for project_root in project_roots:
     sys.path.insert(0, project_root) if project_root not in sys.path else None
 
-from videox_fun.dist import set_multi_gpus_devices, shard_model
 from videox_fun.models import (AutoencoderKLCogVideoX,
-                               CogVideoXTransformer3DModel, T5EncoderModel,
-                               T5Tokenizer)
-from videox_fun.pipeline import (CogVideoXFunInpaintPipeline,
-                                 CogVideoXFunPipeline)
-from videox_fun.utils import (apply_gpu_memory_mode, get_video_to_video_latent,
-                              merge_lora, save_videos_grid, unmerge_lora)
+                              CogVideoXTransformer3DModel, T5EncoderModel,
+                              T5Tokenizer)
+from videox_fun.pipeline import (CogVideoXFunPipeline,
+                                CogVideoXFunInpaintPipeline)
+from videox_fun.utils import (register_auto_device_hook,
+                              safe_enable_group_offload)
+from videox_fun.utils.lora_utils import merge_lora, unmerge_lora
+from videox_fun.utils.fp8_optimization import (convert_model_weight_to_float8, replace_parameters_by_name,
+                                              convert_weight_dtype_wrapper)
+from videox_fun.utils.utils import get_video_to_video_latent, save_videos_grid
+from videox_fun.dist import set_multi_gpus_devices, shard_model
 
-# GPU memory mode, which can be chosen in [model_full_load, model_full_load_and_qfloat8, model_cpu_offload, model_cpu_offload_and_qfloat8, model_group_offload, sequential_cpu_offload].
+# GPU memory mode, which can be chosen in [model_full_load, model_full_load_and_qfloat8, model_cpu_offload, model_cpu_offload_and_qfloat8, sequential_cpu_offload].
 # model_full_load means that the entire model will be moved to the GPU.
 # 
 # model_full_load_and_qfloat8 means that the entire model will be moved to the GPU,
@@ -71,7 +75,7 @@ video_length        = 49
 fps                 = 8
 
 # Use torch.float16 if GPU does not support torch.bfloat16
-# Some graphics cards, such as v100, 2080ti, do not support torch.bfloat16
+# ome graphics cards, such as v100, 2080ti, do not support torch.bfloat16
 weight_dtype            = torch.bfloat16
 # If you are preparing to redraw the reference video, set validation_video and validation_video_mask. 
 # If you do not use validation_video_mask, the entire video will be redrawn; 
@@ -102,7 +106,7 @@ transformer = CogVideoXTransformer3DModel.from_pretrained(
 if transformer_path is not None:
     print(f"From checkpoint: {transformer_path}")
     if transformer_path.endswith("safetensors"):
-        from safetensors.torch import load_file
+        from safetensors.torch import load_file, safe_open
         state_dict = load_file(transformer_path)
     else:
         state_dict = torch.load(transformer_path, map_location="cpu")
@@ -120,7 +124,7 @@ vae = AutoencoderKLCogVideoX.from_pretrained(
 if vae_path is not None:
     print(f"From checkpoint: {vae_path}")
     if vae_path.endswith("safetensors"):
-        from safetensors.torch import load_file
+        from safetensors.torch import load_file, safe_open
         state_dict = load_file(vae_path)
     else:
         state_dict = torch.load(vae_path, map_location="cpu")
@@ -138,7 +142,7 @@ text_encoder = T5EncoderModel.from_pretrained(
 )
 
 # Get Scheduler
-Chosen_Scheduler = {
+Chosen_Scheduler = scheduler_dict = {
     "Euler": EulerDiscreteScheduler,
     "Euler A": EulerAncestralDiscreteScheduler,
     "DPM++": DPMSolverMultistepScheduler, 
@@ -184,10 +188,23 @@ if compile_dit:
         pipeline.transformer.transformer_blocks[i] = torch.compile(pipeline.transformer.transformer_blocks[i])
     print("Add Compile")
 
-# Quantize (when the mode carries an "_and_<quant>" suffix) and then place the pipeline.
-# The order lives inside the helper: quantization has to happen before the offload hooks
-# are installed.
-apply_gpu_memory_mode(pipeline, GPU_memory_mode, device, weight_dtype, exclude_module_name=[])
+if GPU_memory_mode == "sequential_cpu_offload":
+    pipeline.enable_sequential_cpu_offload(device=device)
+elif GPU_memory_mode == "model_group_offload":
+    register_auto_device_hook(pipeline.transformer)
+    safe_enable_group_offload(pipeline, onload_device=device, offload_device="cpu", offload_type="leaf_level", use_stream=True)
+elif GPU_memory_mode == "model_cpu_offload_and_qfloat8":
+    convert_model_weight_to_float8(transformer, exclude_module_name=[], device=device)
+    convert_weight_dtype_wrapper(transformer, weight_dtype)
+    pipeline.enable_model_cpu_offload(device=device)
+elif GPU_memory_mode == "model_cpu_offload":
+    pipeline.enable_model_cpu_offload(device=device)
+elif GPU_memory_mode == "model_full_load_and_qfloat8":
+    convert_model_weight_to_float8(transformer, exclude_module_name=[], device=device)
+    convert_weight_dtype_wrapper(transformer, weight_dtype)
+    pipeline.to(device=device)
+else:
+    pipeline.to(device=device)
 
 generator = torch.Generator(device=device).manual_seed(seed)
 

examples/cogvideox_fun/predict_v2v_control.py

@@ -1,6 +1,7 @@
 import os
 import sys
 
+import cv2
 import numpy as np
 import torch
 from diffusers import (CogVideoXDDIMScheduler, DDIMScheduler,
@@ -15,15 +16,20 @@ project_roots = [os.path.dirname(current_file_path), os.path.dirname(os.path.dir
 for project_root in project_roots:
     sys.path.insert(0, project_root) if project_root not in sys.path else None
 
-from videox_fun.dist import set_multi_gpus_devices, shard_model
 from videox_fun.models import (AutoencoderKLCogVideoX,
-                               CogVideoXTransformer3DModel, T5EncoderModel,
-                               T5Tokenizer)
-from videox_fun.pipeline import CogVideoXFunControlPipeline
-from videox_fun.utils import (apply_gpu_memory_mode, get_video_to_video_latent,
-                              merge_lora, save_videos_grid, unmerge_lora)
+                              CogVideoXTransformer3DModel, T5EncoderModel,
+                              T5Tokenizer)
+from videox_fun.pipeline import (CogVideoXFunControlPipeline,
+                                CogVideoXFunInpaintPipeline)
+from videox_fun.utils import (register_auto_device_hook,
+                              safe_enable_group_offload)
+from videox_fun.utils.fp8_optimization import (convert_model_weight_to_float8, replace_parameters_by_name,
+                                              convert_weight_dtype_wrapper)
+from videox_fun.utils.lora_utils import merge_lora, unmerge_lora
+from videox_fun.utils.utils import get_video_to_video_latent, save_videos_grid
+from videox_fun.dist import set_multi_gpus_devices, shard_model
 
-# GPU memory mode, which can be chosen in [model_full_load, model_full_load_and_qfloat8, model_cpu_offload, model_cpu_offload_and_qfloat8, model_group_offload, sequential_cpu_offload].
+# GPU memory mode, which can be chosen in [model_full_load, model_full_load_and_qfloat8, model_cpu_offload, model_cpu_offload_and_qfloat8, sequential_cpu_offload].
 # model_full_load means that the entire model will be moved to the GPU.
 # 
 # model_full_load_and_qfloat8 means that the entire model will be moved to the GPU,
@@ -71,7 +77,7 @@ video_length        = 49
 fps                 = 8
 
 # Use torch.float16 if GPU does not support torch.bfloat16
-# Some graphics cards, such as v100, 2080ti, do not support torch.bfloat16
+# ome graphics cards, such as v100, 2080ti, do not support torch.bfloat16
 weight_dtype            = torch.bfloat16
 control_video           = "asset/pose.mp4"
 
@@ -96,7 +102,7 @@ transformer = CogVideoXTransformer3DModel.from_pretrained(
 if transformer_path is not None:
     print(f"From checkpoint: {transformer_path}")
     if transformer_path.endswith("safetensors"):
-        from safetensors.torch import load_file
+        from safetensors.torch import load_file, safe_open
         state_dict = load_file(transformer_path)
     else:
         state_dict = torch.load(transformer_path, map_location="cpu")
@@ -114,7 +120,7 @@ vae = AutoencoderKLCogVideoX.from_pretrained(
 if vae_path is not None:
     print(f"From checkpoint: {vae_path}")
     if vae_path.endswith("safetensors"):
-        from safetensors.torch import load_file
+        from safetensors.torch import load_file, safe_open
         state_dict = load_file(vae_path)
     else:
         state_dict = torch.load(vae_path, map_location="cpu")
@@ -132,7 +138,7 @@ text_encoder = T5EncoderModel.from_pretrained(
 )
 
 # Get Scheduler
-Chosen_Scheduler = {
+Chosen_Scheduler = scheduler_dict = {
     "Euler": EulerDiscreteScheduler,
     "Euler A": EulerAncestralDiscreteScheduler,
     "DPM++": DPMSolverMultistepScheduler, 
@@ -169,10 +175,23 @@ if compile_dit:
         pipeline.transformer.transformer_blocks[i] = torch.compile(pipeline.transformer.transformer_blocks[i])
     print("Add Compile")
 
-# Quantize (when the mode carries an "_and_<quant>" suffix) and then place the pipeline.
-# The order lives inside the helper: quantization has to happen before the offload hooks
-# are installed.
-apply_gpu_memory_mode(pipeline, GPU_memory_mode, device, weight_dtype, exclude_module_name=[])
+if GPU_memory_mode == "sequential_cpu_offload":
+    pipeline.enable_sequential_cpu_offload(device=device)
+elif GPU_memory_mode == "model_group_offload":
+    register_auto_device_hook(pipeline.transformer)
+    safe_enable_group_offload(pipeline, onload_device=device, offload_device="cpu", offload_type="leaf_level", use_stream=True)
+elif GPU_memory_mode == "model_cpu_offload_and_qfloat8":
+    convert_model_weight_to_float8(transformer, exclude_module_name=[], device=device)
+    convert_weight_dtype_wrapper(transformer, weight_dtype)
+    pipeline.enable_model_cpu_offload(device=device)
+elif GPU_memory_mode == "model_cpu_offload":
+    pipeline.enable_model_cpu_offload(device=device)
+elif GPU_memory_mode == "model_full_load_and_qfloat8":
+    convert_model_weight_to_float8(transformer, exclude_module_name=[], device=device)
+    convert_weight_dtype_wrapper(transformer, weight_dtype)
+    pipeline.to(device=device)
+else:
+    pipeline.to(device=device)
 
 generator = torch.Generator(device=device).manual_seed(seed)
 

examples/ernie_image/predict_t2i.py

@@ -13,11 +13,15 @@ from videox_fun.dist import set_multi_gpus_devices, shard_model
 from videox_fun.models import (AutoencoderKLFlux2, AutoTokenizer,
                                ErnieImageTransformer2DModel, Mistral3Model)
 from videox_fun.pipeline import ErnieImagePipeline
-from videox_fun.utils import (FlowDPMSolverMultistepScheduler,
-                              FlowUniPCMultistepScheduler,
-                              apply_gpu_memory_mode, merge_lora, unmerge_lora)
+from videox_fun.utils import (register_auto_device_hook,
+                              safe_enable_group_offload)
+from videox_fun.utils.fm_solvers import FlowDPMSolverMultistepScheduler
+from videox_fun.utils.fm_solvers_unipc import FlowUniPCMultistepScheduler
+from videox_fun.utils.fp8_optimization import (convert_model_weight_to_float8,
+                                               convert_weight_dtype_wrapper)
+from videox_fun.utils.lora_utils import merge_lora, unmerge_lora
 
-# GPU memory mode, which can be chosen in [model_full_load, model_full_load_and_qfloat8, model_cpu_offload, model_cpu_offload_and_qfloat8, model_group_offload, sequential_cpu_offload].
+# GPU memory mode, which can be chosen in [model_full_load, model_full_load_and_qfloat8, model_cpu_offload, model_cpu_offload_and_qfloat8, sequential_cpu_offload].
 # model_full_load means that the entire model will be moved to the GPU.
 # 
 # model_full_load_and_qfloat8 means that the entire model will be moved to the GPU,
@@ -62,7 +66,7 @@ lora_path           = None
 sample_size         = [1728, 992]
 
 # Use torch.float16 if GPU does not support torch.bfloat16
-# Some graphics cards, such as v100, 2080ti, do not support torch.bfloat16
+# ome graphics cards, such as v100, 2080ti, do not support torch.bfloat16
 weight_dtype        = torch.bfloat16
 prompt              = "1girl, black_hair, brown_eyes, earrings, freckles, grey_background, jewelry, lips, long_hair, looking_at_viewer, nose, piercing, realistic, red_lips, solo, upper_body"
 negative_prompt     = "低分辨率，低画质，肢体畸形，手指畸形，画面过饱和，蜡像感，人脸无细节，过度光滑，画面具有AI感。构图混乱。文字模糊，扭曲。"
@@ -84,7 +88,7 @@ transformer = ErnieImageTransformer2DModel.from_pretrained(
 if transformer_path is not None:
     print(f"From checkpoint: {transformer_path}")
     if transformer_path.endswith("safetensors"):
-        from safetensors.torch import load_file
+        from safetensors.torch import load_file, safe_open
         state_dict = load_file(transformer_path)
     else:
         state_dict = torch.load(transformer_path, map_location="cpu")
@@ -102,7 +106,7 @@ vae = AutoencoderKLFlux2.from_pretrained(
 if vae_path is not None:
     print(f"From checkpoint: {vae_path}")
     if vae_path.endswith("safetensors"):
-        from safetensors.torch import load_file
+        from safetensors.torch import load_file, safe_open
         state_dict = load_file(vae_path)
     else:
         state_dict = torch.load(vae_path, map_location="cpu")
@@ -120,7 +124,7 @@ text_encoder = Mistral3Model.from_pretrained(
 )
 
 # Get Scheduler
-Chosen_Scheduler = {
+Chosen_Scheduler = scheduler_dict = {
     "Flow": FlowMatchEulerDiscreteScheduler,
     "Flow_Unipc": FlowUniPCMultistepScheduler,
     "Flow_DPM++": FlowDPMSolverMultistepScheduler,
@@ -151,10 +155,23 @@ if compile_dit:
         pipeline.transformer.layers[i] = torch.compile(pipeline.transformer.layers[i])
     print("Add Compile")
 
-# Quantize (when the mode carries an "_and_<quant>" suffix) and then place the pipeline.
-# The order lives inside the helper: quantization has to happen before the offload hooks
-# are installed.
-apply_gpu_memory_mode(pipeline, GPU_memory_mode, device, weight_dtype, exclude_module_name=['img_in', 'txt_in', 'timestep'])
+if GPU_memory_mode == "sequential_cpu_offload":
+    pipeline.enable_sequential_cpu_offload(device=device)
+elif GPU_memory_mode == "model_group_offload":
+    register_auto_device_hook(pipeline.transformer)
+    safe_enable_group_offload(pipeline, onload_device=device, offload_device="cpu", offload_type="leaf_level", use_stream=True)
+elif GPU_memory_mode == "model_cpu_offload_and_qfloat8":
+    convert_model_weight_to_float8(transformer, exclude_module_name=["img_in", "txt_in", "timestep"], device=device)
+    convert_weight_dtype_wrapper(transformer, weight_dtype)
+    pipeline.enable_model_cpu_offload(device=device)
+elif GPU_memory_mode == "model_cpu_offload":
+    pipeline.enable_model_cpu_offload(device=device)
+elif GPU_memory_mode == "model_full_load_and_qfloat8":
+    convert_model_weight_to_float8(transformer, exclude_module_name=["img_in", "txt_in", "timestep"], device=device)
+    convert_weight_dtype_wrapper(transformer, weight_dtype)
+    pipeline.to(device=device)
+else:
+    pipeline.to(device=device)
 
 generator = torch.Generator(device=device).manual_seed(seed)
 

examples/fantasytalking/predict_s2v.py

@@ -15,19 +15,24 @@ for project_root in project_roots:
 from videox_fun.dist import set_multi_gpus_devices, shard_model
 from videox_fun.models import (AutoencoderKLWan, AutoencoderKLWan3_8,
                                AutoTokenizer, CLIPModel,
-                               FantasyTalkingAudioEncoder,
-                               FantasyTalkingTransformer3DModel,
+                               FantasyTalkingTransformer3DModel, FantasyTalkingAudioEncoder,
                                WanT5EncoderModel)
 from videox_fun.models.cache_utils import get_teacache_coefficients
 from videox_fun.pipeline import FantasyTalkingPipeline
-from videox_fun.utils import (FlowDPMSolverMultistepScheduler,
-                              FlowUniPCMultistepScheduler,
-                              apply_gpu_memory_mode, filter_kwargs,
-                              get_image_to_video_latent, merge_lora,
-                              merge_video_audio, save_videos_grid,
-                              unmerge_lora)
+from videox_fun.utils.fm_solvers import FlowDPMSolverMultistepScheduler
+from videox_fun.utils.fm_solvers_unipc import FlowUniPCMultistepScheduler
+from videox_fun.utils import (register_auto_device_hook,
+                              safe_enable_group_offload)
+from videox_fun.utils.fp8_optimization import (convert_model_weight_to_float8,
+                                               convert_weight_dtype_wrapper,
+                                               replace_parameters_by_name)
+from videox_fun.utils.lora_utils import merge_lora, unmerge_lora
+from videox_fun.utils.utils import (filter_kwargs, get_image_latent,
+                                    get_image_to_video_latent,
+                                    get_video_to_video_latent,
+                                    merge_video_audio, save_videos_grid)
 
-# GPU memory mode, which can be chosen in [model_full_load, model_full_load_and_qfloat8, model_cpu_offload, model_cpu_offload_and_qfloat8, model_group_offload, sequential_cpu_offload].
+# GPU memory mode, which can be chosen in [model_full_load, model_full_load_and_qfloat8, model_cpu_offload, model_cpu_offload_and_qfloat8, sequential_cpu_offload].
 # model_full_load means that the entire model will be moved to the GPU.
 # 
 # model_full_load_and_qfloat8 means that the entire model will be moved to the GPU,
@@ -106,7 +111,7 @@ video_length        = 81
 fps                 = 23
 
 # Use torch.float16 if GPU does not support torch.bfloat16
-# Some graphics cards, such as v100, 2080ti, do not support torch.bfloat16
+# ome graphics cards, such as v100, 2080ti, do not support torch.bfloat16
 weight_dtype            = torch.bfloat16
 # If you want to generate from text, please set the validation_image_start = None
 validation_image_start  = "asset/8.png"
@@ -135,7 +140,7 @@ transformer = FantasyTalkingTransformer3DModel.from_pretrained(
 if transformer_path is not None:
     print(f"From checkpoint: {transformer_path}")
     if transformer_path.endswith("safetensors"):
-        from safetensors.torch import load_file
+        from safetensors.torch import load_file, safe_open
         state_dict = load_file(transformer_path)
     else:
         state_dict = torch.load(transformer_path, map_location="cpu")
@@ -168,7 +173,7 @@ vae = Chosen_AutoencoderKL.from_pretrained(
 if vae_path is not None:
     print(f"From checkpoint: {vae_path}")
     if vae_path.endswith("safetensors"):
-        from safetensors.torch import load_file
+        from safetensors.torch import load_file, safe_open
         state_dict = load_file(vae_path)
     else:
         state_dict = torch.load(vae_path, map_location="cpu")
@@ -201,7 +206,7 @@ audio_encoder_path = model_name_audio if model_name_audio is not None else os.pa
 audio_encoder = FantasyTalkingAudioEncoder(audio_encoder_path)
 
 # Get Scheduler
-Chosen_Scheduler = {
+Chosen_Scheduler = scheduler_dict = {
     "Flow": FlowMatchEulerDiscreteScheduler,
     "Flow_Unipc": FlowUniPCMultistepScheduler,
     "Flow_DPM++": FlowDPMSolverMultistepScheduler,
@@ -239,10 +244,25 @@ if compile_dit:
         pipeline.transformer.blocks[i] = torch.compile(pipeline.transformer.blocks[i])
     print("Add Compile")
 
-# Quantize (when the mode carries an "_and_<quant>" suffix) and then place the pipeline.
-# The order lives inside the helper: quantization has to happen before the offload hooks
-# are installed.
-apply_gpu_memory_mode(pipeline, GPU_memory_mode, device, weight_dtype)
+if GPU_memory_mode == "sequential_cpu_offload":
+    replace_parameters_by_name(transformer, ["modulation",], device=device)
+    transformer.freqs = transformer.freqs.to(device=device)
+    pipeline.enable_sequential_cpu_offload(device=device)
+elif GPU_memory_mode == "model_group_offload":
+    register_auto_device_hook(pipeline.transformer)
+    safe_enable_group_offload(pipeline, onload_device=device, offload_device="cpu", offload_type="leaf_level", use_stream=True)
+elif GPU_memory_mode == "model_cpu_offload_and_qfloat8":
+    convert_model_weight_to_float8(transformer, exclude_module_name=["modulation",], device=device)
+    convert_weight_dtype_wrapper(transformer, weight_dtype)
+    pipeline.enable_model_cpu_offload(device=device)
+elif GPU_memory_mode == "model_cpu_offload":
+    pipeline.enable_model_cpu_offload(device=device)
+elif GPU_memory_mode == "model_full_load_and_qfloat8":
+    convert_model_weight_to_float8(transformer, exclude_module_name=["modulation",], device=device)
+    convert_weight_dtype_wrapper(transformer, weight_dtype)
+    pipeline.to(device=device)
+else:
+    pipeline.to(device=device)
 
 coefficients = get_teacache_coefficients(model_name) if enable_teacache else None
 if coefficients is not None:

examples/flashhead/predict_s2v.py

@@ -13,16 +13,22 @@ for project_root in project_roots:
     sys.path.insert(0, project_root) if project_root not in sys.path else None
 
 from videox_fun.dist import set_multi_gpus_devices, shard_model
-from videox_fun.models import (AutoencoderKLWan, FlashHeadAudioEncoder,
-                               FlashHeadTransformer3DModel)
+from videox_fun.models import (AutoencoderKLWan, AutoencoderKLWan3_8,
+                               FlashHeadTransformer3DModel, FlashHeadAudioEncoder)
 from videox_fun.pipeline import FlashHeadPipeline
-from videox_fun.utils import (FlowDPMSolverMultistepScheduler,
-                              FlowUniPCMultistepScheduler,
-                              apply_gpu_memory_mode, filter_kwargs,
-                              get_image_latent, merge_lora, merge_video_audio,
-                              save_videos_grid, unmerge_lora)
+from videox_fun.utils.fm_solvers import FlowDPMSolverMultistepScheduler
+from videox_fun.utils.fm_solvers_unipc import FlowUniPCMultistepScheduler
+from videox_fun.utils import (register_auto_device_hook,
+                              safe_enable_group_offload)
+from videox_fun.utils.fp8_optimization import (convert_model_weight_to_float8,
+                                               convert_weight_dtype_wrapper,
+                                               replace_parameters_by_name)
+from videox_fun.utils.lora_utils import merge_lora, unmerge_lora
+from videox_fun.utils.utils import (filter_kwargs, get_image_latent, get_image, 
+                                    get_video_to_video_latent,
+                                    merge_video_audio, save_videos_grid)
 
-# GPU memory mode, which can be chosen in [model_full_load, model_full_load_and_qfloat8, model_cpu_offload, model_cpu_offload_and_qfloat8, model_group_offload, sequential_cpu_offload].
+# GPU memory mode, which can be chosen in [model_full_load, model_full_load_and_qfloat8, model_cpu_offload, model_cpu_offload_and_qfloat8, sequential_cpu_offload].
 # model_full_load means that the entire model will be moved to the GPU.
 # 
 # model_full_load_and_qfloat8 means that the entire model will be moved to the GPU,
@@ -71,7 +77,7 @@ segment_frame_length    = 33
 fps                     = 25 
 
 # Use torch.float16 if GPU does not support torch.bfloat16
-# Some graphics cards, such as v100, 2080ti, do not support torch.bfloat16
+# ome graphics cards, such as v100, 2080ti, do not support torch.bfloat16
 weight_dtype            = torch.bfloat16
 # The path of the reference image
 ref_image               = "asset/9.png"
@@ -124,7 +130,7 @@ vae = AutoencoderKLWan.from_pretrained(
 if vae_path is not None:
     print(f"From checkpoint: {vae_path}")
     if vae_path.endswith("safetensors"):
-        from safetensors.torch import load_file
+        from safetensors.torch import load_file, safe_open
         state_dict = load_file(vae_path)
     else:
         state_dict = torch.load(vae_path, map_location="cpu")
@@ -140,7 +146,7 @@ audio_encoder = FlashHeadAudioEncoder(
 )
 
 # Get Scheduler
-Chosen_Scheduler = {
+Chosen_Scheduler = scheduler_dict = {
     "Flow": FlowMatchEulerDiscreteScheduler,
     "Flow_Unipc": FlowUniPCMultistepScheduler,
     "Flow_DPM++": FlowDPMSolverMultistepScheduler,
@@ -171,10 +177,25 @@ if compile_dit:
         pipeline.transformer.blocks[i] = torch.compile(pipeline.transformer.blocks[i])
     print("Add Compile")
 
-# Quantize (when the mode carries an "_and_<quant>" suffix) and then place the pipeline.
-# The order lives inside the helper: quantization has to happen before the offload hooks
-# are installed.
-apply_gpu_memory_mode(pipeline, GPU_memory_mode, device, weight_dtype)
+if GPU_memory_mode == "sequential_cpu_offload":
+    replace_parameters_by_name(transformer, ["modulation",], device=device)
+    transformer.freqs = transformer.freqs.to(device=device)
+    pipeline.enable_sequential_cpu_offload(device=device)
+elif GPU_memory_mode == "model_group_offload":
+    register_auto_device_hook(pipeline.transformer)
+    safe_enable_group_offload(pipeline, onload_device=device, offload_device="cpu", offload_type="leaf_level", use_stream=True)
+elif GPU_memory_mode == "model_cpu_offload_and_qfloat8":
+    convert_model_weight_to_float8(transformer, exclude_module_name=["modulation",], device=device)
+    convert_weight_dtype_wrapper(transformer, weight_dtype)
+    pipeline.enable_model_cpu_offload(device=device)
+elif GPU_memory_mode == "model_cpu_offload":
+    pipeline.enable_model_cpu_offload(device=device)
+elif GPU_memory_mode == "model_full_load_and_qfloat8":
+    convert_model_weight_to_float8(transformer, exclude_module_name=["modulation",], device=device)
+    convert_weight_dtype_wrapper(transformer, weight_dtype)
+    pipeline.to(device=device)
+else:
+    pipeline.to(device=device)
 
 generator = torch.Generator(device=device).manual_seed(seed)
 

examples/flux/predict_t2i.py

@@ -14,11 +14,15 @@ from videox_fun.models import (AutoencoderKL, CLIPTextModel, CLIPTokenizer,
                                FluxTransformer2DModel, T5EncoderModel,
                                T5TokenizerFast)
 from videox_fun.pipeline import FluxPipeline
-from videox_fun.utils import (FlowDPMSolverMultistepScheduler,
-                              FlowUniPCMultistepScheduler,
-                              apply_gpu_memory_mode, merge_lora, unmerge_lora)
+from videox_fun.utils import (register_auto_device_hook,
+                              safe_enable_group_offload)
+from videox_fun.utils.fm_solvers import FlowDPMSolverMultistepScheduler
+from videox_fun.utils.fm_solvers_unipc import FlowUniPCMultistepScheduler
+from videox_fun.utils.fp8_optimization import (convert_model_weight_to_float8,
+                                               convert_weight_dtype_wrapper)
+from videox_fun.utils.lora_utils import merge_lora, unmerge_lora
 
-# GPU memory mode, which can be chosen in [model_full_load, model_full_load_and_qfloat8, model_cpu_offload, model_cpu_offload_and_qfloat8, model_group_offload, sequential_cpu_offload].
+# GPU memory mode, which can be chosen in [model_full_load, model_full_load_and_qfloat8, model_cpu_offload, model_cpu_offload_and_qfloat8, sequential_cpu_offload].
 # model_full_load means that the entire model will be moved to the GPU.
 # 
 # model_full_load_and_qfloat8 means that the entire model will be moved to the GPU,
@@ -63,7 +67,7 @@ lora_path           = None
 sample_size         = [1344, 768]
 
 # Use torch.float16 if GPU does not support torch.bfloat16
-# Some graphics cards, such as v100, 2080ti, do not support torch.bfloat16
+# ome graphics cards, such as v100, 2080ti, do not support torch.bfloat16
 weight_dtype        = torch.bfloat16
 prompt              = "1girl, black_hair, brown_eyes, earrings, freckles, grey_background, jewelry, lips, long_hair, looking_at_viewer, nose, piercing, realistic, red_lips, solo, upper_body"
 negative_prompt     = " "
@@ -85,7 +89,7 @@ transformer = FluxTransformer2DModel.from_pretrained(
 if transformer_path is not None:
     print(f"From checkpoint: {transformer_path}")
     if transformer_path.endswith("safetensors"):
-        from safetensors.torch import load_file
+        from safetensors.torch import load_file, safe_open
         state_dict = load_file(transformer_path)
     else:
         state_dict = torch.load(transformer_path, map_location="cpu")
@@ -103,7 +107,7 @@ vae = AutoencoderKL.from_pretrained(
 if vae_path is not None:
     print(f"From checkpoint: {vae_path}")
     if vae_path.endswith("safetensors"):
-        from safetensors.torch import load_file
+        from safetensors.torch import load_file, safe_open
         state_dict = load_file(vae_path)
     else:
         state_dict = torch.load(vae_path, map_location="cpu")
@@ -128,7 +132,7 @@ text_encoder_2 = T5EncoderModel.from_pretrained(
 )
 
 # Get Scheduler
-Chosen_Scheduler = {
+Chosen_Scheduler = scheduler_dict = {
     "Flow": FlowMatchEulerDiscreteScheduler,
     "Flow_Unipc": FlowUniPCMultistepScheduler,
     "Flow_DPM++": FlowDPMSolverMultistepScheduler,
@@ -165,10 +169,23 @@ if compile_dit:
         pipeline.transformer.transformer_blocks[i] = torch.compile(pipeline.transformer.transformer_blocks[i])
     print("Add Compile")
 
-# Quantize (when the mode carries an "_and_<quant>" suffix) and then place the pipeline.
-# The order lives inside the helper: quantization has to happen before the offload hooks
-# are installed.
-apply_gpu_memory_mode(pipeline, GPU_memory_mode, device, weight_dtype, exclude_module_name=['img_in', 'txt_in', 'timestep'])
+if GPU_memory_mode == "sequential_cpu_offload":
+    pipeline.enable_sequential_cpu_offload(device=device)
+elif GPU_memory_mode == "model_group_offload":
+    register_auto_device_hook(pipeline.transformer)
+    safe_enable_group_offload(pipeline, onload_device=device, offload_device="cpu", offload_type="leaf_level", use_stream=True)
+elif GPU_memory_mode == "model_cpu_offload_and_qfloat8":
+    convert_model_weight_to_float8(transformer, exclude_module_name=["img_in", "txt_in", "timestep"], device=device)
+    convert_weight_dtype_wrapper(transformer, weight_dtype)
+    pipeline.enable_model_cpu_offload(device=device)
+elif GPU_memory_mode == "model_cpu_offload":
+    pipeline.enable_model_cpu_offload(device=device)
+elif GPU_memory_mode == "model_full_load_and_qfloat8":
+    convert_model_weight_to_float8(transformer, exclude_module_name=["img_in", "txt_in", "timestep"], device=device)
+    convert_weight_dtype_wrapper(transformer, weight_dtype)
+    pipeline.to(device=device)
+else:
+    pipeline.to(device=device)
 
 generator = torch.Generator(device=device).manual_seed(seed)
 

examples/flux2/predict_t2i.py

@@ -2,7 +2,8 @@ import os
 import sys
 
 import torch
-from diffusers import FlowMatchEulerDiscreteScheduler
+
+from diffusers import (FlowMatchEulerDiscreteScheduler)
 
 current_file_path = os.path.abspath(__file__)
 project_roots = [os.path.dirname(current_file_path), os.path.dirname(os.path.dirname(current_file_path)), os.path.dirname(os.path.dirname(os.path.dirname(current_file_path)))]
@@ -10,15 +11,20 @@ for project_root in project_roots:
     sys.path.insert(0, project_root) if project_root not in sys.path else None
 
 from videox_fun.dist import set_multi_gpus_devices, shard_model
-from videox_fun.models import (AutoencoderKLFlux2, Flux2Transformer2DModel,
+from videox_fun.models import (AutoencoderKLFlux2,
                                Mistral3ForConditionalGeneration,
-                               PixtralProcessor)
+                               PixtralProcessor, Flux2Transformer2DModel)
+from videox_fun.models.cache_utils import get_teacache_coefficients
 from videox_fun.pipeline import Flux2Pipeline
-from videox_fun.utils import (FlowDPMSolverMultistepScheduler,
-                              FlowUniPCMultistepScheduler,
-                              apply_gpu_memory_mode, merge_lora, unmerge_lora)
+from videox_fun.utils import (register_auto_device_hook,
+                              safe_enable_group_offload)
+from videox_fun.utils.fm_solvers import FlowDPMSolverMultistepScheduler
+from videox_fun.utils.fm_solvers_unipc import FlowUniPCMultistepScheduler
+from videox_fun.utils.fp8_optimization import (convert_model_weight_to_float8,
+                                               convert_weight_dtype_wrapper)
+from videox_fun.utils.lora_utils import merge_lora, unmerge_lora
 
-# GPU memory mode, which can be chosen in [model_full_load, model_full_load_and_qfloat8, model_cpu_offload, model_cpu_offload_and_qfloat8, model_group_offload, sequential_cpu_offload].
+# GPU memory mode, which can be chosen in [model_full_load, model_full_load_and_qfloat8, model_cpu_offload, model_cpu_offload_and_qfloat8, sequential_cpu_offload].
 # model_full_load means that the entire model will be moved to the GPU.
 # 
 # model_full_load_and_qfloat8 means that the entire model will be moved to the GPU,
@@ -63,7 +69,7 @@ lora_path           = None
 sample_size         = [1344, 768]
 
 # Use torch.float16 if GPU does not support torch.bfloat16
-# Some graphics cards, such as v100, 2080ti, do not support torch.bfloat16
+# ome graphics cards, such as v100, 2080ti, do not support torch.bfloat16
 weight_dtype        = torch.bfloat16
 # Please use as detailed a prompt as possible to describe the object that needs to be generated.
 prompt              = "1girl, black_hair, brown_eyes, earrings, freckles, grey_background, jewelry, lips, long_hair, looking_at_viewer, nose, piercing, realistic, red_lips, solo, upper_body"
@@ -86,7 +92,7 @@ transformer = Flux2Transformer2DModel.from_pretrained(
 if transformer_path is not None:
     print(f"From checkpoint: {transformer_path}")
     if transformer_path.endswith("safetensors"):
-        from safetensors.torch import load_file
+        from safetensors.torch import load_file, safe_open
         state_dict = load_file(transformer_path)
     else:
         state_dict = torch.load(transformer_path, map_location="cpu")
@@ -104,7 +110,7 @@ vae = AutoencoderKLFlux2.from_pretrained(
 if vae_path is not None:
     print(f"From checkpoint: {vae_path}")
     if vae_path.endswith("safetensors"):
-        from safetensors.torch import load_file
+        from safetensors.torch import load_file, safe_open
         state_dict = load_file(vae_path)
     else:
         state_dict = torch.load(vae_path, map_location="cpu")
@@ -123,7 +129,7 @@ text_encoder = Mistral3ForConditionalGeneration.from_pretrained(
 )
 
 # Get Scheduler
-Chosen_Scheduler = {
+Chosen_Scheduler = scheduler_dict = {
     "Flow": FlowMatchEulerDiscreteScheduler,
     "Flow_Unipc": FlowUniPCMultistepScheduler,
     "Flow_DPM++": FlowDPMSolverMultistepScheduler,
@@ -158,10 +164,23 @@ if compile_dit:
         pipeline.transformer.transformer_blocks[i] = torch.compile(pipeline.transformer.transformer_blocks[i])
     print("Add Compile")
 
-# Quantize (when the mode carries an "_and_<quant>" suffix) and then place the pipeline.
-# The order lives inside the helper: quantization has to happen before the offload hooks
-# are installed.
-apply_gpu_memory_mode(pipeline, GPU_memory_mode, device, weight_dtype, exclude_module_name=['img_in', 'txt_in', 'timestep'])
+if GPU_memory_mode == "sequential_cpu_offload":
+    pipeline.enable_sequential_cpu_offload(device=device)
+elif GPU_memory_mode == "model_group_offload":
+    register_auto_device_hook(pipeline.transformer)
+    safe_enable_group_offload(pipeline, onload_device=device, offload_device="cpu", offload_type="leaf_level", use_stream=True)
+elif GPU_memory_mode == "model_cpu_offload_and_qfloat8":
+    convert_model_weight_to_float8(transformer, exclude_module_name=["img_in", "txt_in", "timestep"], device=device)
+    convert_weight_dtype_wrapper(transformer, weight_dtype)
+    pipeline.enable_model_cpu_offload(device=device)
+elif GPU_memory_mode == "model_cpu_offload":
+    pipeline.enable_model_cpu_offload(device=device)
+elif GPU_memory_mode == "model_full_load_and_qfloat8":
+    convert_model_weight_to_float8(transformer, exclude_module_name=["img_in", "txt_in", "timestep"], device=device)
+    convert_weight_dtype_wrapper(transformer, weight_dtype)
+    pipeline.to(device=device)
+else:
+    pipeline.to(device=device)
 
 generator = torch.Generator(device=device).manual_seed(seed)
 

examples/flux2_fun/predict_i2i_inpaint.py

@@ -1,9 +1,11 @@
 import os
 import sys
 
+import numpy as np
 import torch
 from diffusers import FlowMatchEulerDiscreteScheduler
 from omegaconf import OmegaConf
+from PIL import Image
 
 current_file_path = os.path.abspath(__file__)
 project_roots = [os.path.dirname(current_file_path), os.path.dirname(os.path.dirname(current_file_path)), os.path.dirname(os.path.dirname(os.path.dirname(current_file_path)))]
@@ -12,16 +14,23 @@ for project_root in project_roots:
 
 from videox_fun.dist import set_multi_gpus_devices, shard_model
 from videox_fun.models import (AutoencoderKLFlux2,
-                               Flux2ControlTransformer2DModel,
                                Mistral3ForConditionalGeneration,
-                               PixtralProcessor)
+                               PixtralProcessor, Flux2ControlTransformer2DModel)
+from videox_fun.models.cache_utils import get_teacache_coefficients
 from videox_fun.pipeline import Flux2ControlPipeline
-from videox_fun.utils import (FlowDPMSolverMultistepScheduler,
-                              FlowUniPCMultistepScheduler,
-                              apply_gpu_memory_mode, get_image,
-                              get_image_latent, merge_lora, unmerge_lora)
+from videox_fun.utils import (register_auto_device_hook,
+                              safe_enable_group_offload)
+from videox_fun.utils.fm_solvers import FlowDPMSolverMultistepScheduler
+from videox_fun.utils.fm_solvers_unipc import FlowUniPCMultistepScheduler
+from videox_fun.utils.fp8_optimization import (convert_model_weight_to_float8,
+                                               convert_weight_dtype_wrapper)
+from videox_fun.utils.lora_utils import merge_lora, unmerge_lora
+from videox_fun.utils.utils import (filter_kwargs, get_image, get_image_latent,
+                                    get_image_to_video_latent,
+                                    get_video_to_video_latent,
+                                    save_videos_grid)
 
-# GPU memory mode, which can be chosen in [model_full_load, model_full_load_and_qfloat8, model_cpu_offload, model_cpu_offload_and_qfloat8, model_group_offload, sequential_cpu_offload].
+# GPU memory mode, which can be chosen in [model_full_load, model_full_load_and_qfloat8, model_cpu_offload, model_cpu_offload_and_qfloat8, sequential_cpu_offload].
 # model_full_load means that the entire model will be moved to the GPU.
 # 
 # model_full_load_and_qfloat8 means that the entire model will be moved to the GPU,
@@ -68,7 +77,7 @@ lora_path           = None
 sample_size         = [1728, 992]
 
 # Use torch.float16 if GPU does not support torch.bfloat16
-# Some graphics cards, such as v100, 2080ti, do not support torch.bfloat16
+# ome graphics cards, such as v100, 2080ti, do not support torch.bfloat16
 weight_dtype        = torch.bfloat16
 image               = None 
 control_image       = None
@@ -99,7 +108,7 @@ transformer = Flux2ControlTransformer2DModel.from_pretrained(
 if transformer_path is not None:
     print(f"From checkpoint: {transformer_path}")
     if transformer_path.endswith("safetensors"):
-        from safetensors.torch import load_file
+        from safetensors.torch import load_file, safe_open
         state_dict = load_file(transformer_path)
     else:
         state_dict = torch.load(transformer_path, map_location="cpu")
@@ -117,7 +126,7 @@ vae = AutoencoderKLFlux2.from_pretrained(
 if vae_path is not None:
     print(f"From checkpoint: {vae_path}")
     if vae_path.endswith("safetensors"):
-        from safetensors.torch import load_file
+        from safetensors.torch import load_file, safe_open
         state_dict = load_file(vae_path)
     else:
         state_dict = torch.load(vae_path, map_location="cpu")
@@ -136,7 +145,7 @@ text_encoder = Mistral3ForConditionalGeneration.from_pretrained(
 )
 
 # Get Scheduler
-Chosen_Scheduler = {
+Chosen_Scheduler = scheduler_dict = {
     "Flow": FlowMatchEulerDiscreteScheduler,
     "Flow_Unipc": FlowUniPCMultistepScheduler,
     "Flow_DPM++": FlowDPMSolverMultistepScheduler,
@@ -171,10 +180,23 @@ if compile_dit:
         pipeline.transformer.transformer_blocks[i] = torch.compile(pipeline.transformer.transformer_blocks[i])
     print("Add Compile")
 
-# Quantize (when the mode carries an "_and_<quant>" suffix) and then place the pipeline.
-# The order lives inside the helper: quantization has to happen before the offload hooks
-# are installed.
-apply_gpu_memory_mode(pipeline, GPU_memory_mode, device, weight_dtype, exclude_module_name=['img_in', 'txt_in', 'timestep'])
+if GPU_memory_mode == "sequential_cpu_offload":
+    pipeline.enable_sequential_cpu_offload(device=device)
+elif GPU_memory_mode == "model_group_offload":
+    register_auto_device_hook(pipeline.transformer)
+    safe_enable_group_offload(pipeline, onload_device=device, offload_device="cpu", offload_type="leaf_level", use_stream=True)
+elif GPU_memory_mode == "model_cpu_offload_and_qfloat8":
+    convert_model_weight_to_float8(transformer, exclude_module_name=["img_in", "txt_in", "timestep"], device=device)
+    convert_weight_dtype_wrapper(transformer, weight_dtype)
+    pipeline.enable_model_cpu_offload(device=device)
+elif GPU_memory_mode == "model_cpu_offload":
+    pipeline.enable_model_cpu_offload(device=device)
+elif GPU_memory_mode == "model_full_load_and_qfloat8":
+    convert_model_weight_to_float8(transformer, exclude_module_name=["img_in", "txt_in", "timestep"], device=device)
+    convert_weight_dtype_wrapper(transformer, weight_dtype)
+    pipeline.to(device=device)
+else:
+    pipeline.to(device=device)
 
 generator = torch.Generator(device=device).manual_seed(seed)
 

examples/flux2_fun/predict_t2i_control.py

@@ -1,9 +1,11 @@
 import os
 import sys
 
+import numpy as np
 import torch
 from diffusers import FlowMatchEulerDiscreteScheduler
 from omegaconf import OmegaConf
+from PIL import Image
 
 current_file_path = os.path.abspath(__file__)
 project_roots = [os.path.dirname(current_file_path), os.path.dirname(os.path.dirname(current_file_path)), os.path.dirname(os.path.dirname(os.path.dirname(current_file_path)))]
@@ -12,16 +14,23 @@ for project_root in project_roots:
 
 from videox_fun.dist import set_multi_gpus_devices, shard_model
 from videox_fun.models import (AutoencoderKLFlux2,
-                               Flux2ControlTransformer2DModel,
                                Mistral3ForConditionalGeneration,
-                               PixtralProcessor)
+                               PixtralProcessor, Flux2ControlTransformer2DModel)
+from videox_fun.models.cache_utils import get_teacache_coefficients
 from videox_fun.pipeline import Flux2ControlPipeline
-from videox_fun.utils import (FlowDPMSolverMultistepScheduler,
-                              FlowUniPCMultistepScheduler,
-                              apply_gpu_memory_mode, get_image,
-                              get_image_latent, merge_lora, unmerge_lora)
+from videox_fun.utils import (register_auto_device_hook,
+                              safe_enable_group_offload)
+from videox_fun.utils.fm_solvers import FlowDPMSolverMultistepScheduler
+from videox_fun.utils.fm_solvers_unipc import FlowUniPCMultistepScheduler
+from videox_fun.utils.fp8_optimization import (convert_model_weight_to_float8,
+                                               convert_weight_dtype_wrapper)
+from videox_fun.utils.lora_utils import merge_lora, unmerge_lora
+from videox_fun.utils.utils import (filter_kwargs, get_image, get_image_latent,
+                                    get_image_to_video_latent,
+                                    get_video_to_video_latent,
+                                    save_videos_grid)
 
-# GPU memory mode, which can be chosen in [model_full_load, model_full_load_and_qfloat8, model_cpu_offload, model_cpu_offload_and_qfloat8, model_group_offload, sequential_cpu_offload].
+# GPU memory mode, which can be chosen in [model_full_load, model_full_load_and_qfloat8, model_cpu_offload, model_cpu_offload_and_qfloat8, sequential_cpu_offload].
 # model_full_load means that the entire model will be moved to the GPU.
 # 
 # model_full_load_and_qfloat8 means that the entire model will be moved to the GPU,
@@ -68,7 +77,7 @@ lora_path           = None
 sample_size         = [1728, 992]
 
 # Use torch.float16 if GPU does not support torch.bfloat16
-# Some graphics cards, such as v100, 2080ti, do not support torch.bfloat16
+# ome graphics cards, such as v100, 2080ti, do not support torch.bfloat16
 weight_dtype        = torch.bfloat16
 image               = None
 control_image       = "asset/pose.jpg"
@@ -99,7 +108,7 @@ transformer = Flux2ControlTransformer2DModel.from_pretrained(
 if transformer_path is not None:
     print(f"From checkpoint: {transformer_path}")
     if transformer_path.endswith("safetensors"):
-        from safetensors.torch import load_file
+        from safetensors.torch import load_file, safe_open
         state_dict = load_file(transformer_path)
     else:
         state_dict = torch.load(transformer_path, map_location="cpu")
@@ -117,7 +126,7 @@ vae = AutoencoderKLFlux2.from_pretrained(
 if vae_path is not None:
     print(f"From checkpoint: {vae_path}")
     if vae_path.endswith("safetensors"):
-        from safetensors.torch import load_file
+        from safetensors.torch import load_file, safe_open
         state_dict = load_file(vae_path)
     else:
         state_dict = torch.load(vae_path, map_location="cpu")
@@ -136,7 +145,7 @@ text_encoder = Mistral3ForConditionalGeneration.from_pretrained(
 )
 
 # Get Scheduler
-Chosen_Scheduler = {
+Chosen_Scheduler = scheduler_dict = {
     "Flow": FlowMatchEulerDiscreteScheduler,
     "Flow_Unipc": FlowUniPCMultistepScheduler,
     "Flow_DPM++": FlowDPMSolverMultistepScheduler,
@@ -171,10 +180,23 @@ if compile_dit:
         pipeline.transformer.transformer_blocks[i] = torch.compile(pipeline.transformer.transformer_blocks[i])
     print("Add Compile")
 
-# Quantize (when the mode carries an "_and_<quant>" suffix) and then place the pipeline.
-# The order lives inside the helper: quantization has to happen before the offload hooks
-# are installed.
-apply_gpu_memory_mode(pipeline, GPU_memory_mode, device, weight_dtype, exclude_module_name=['img_in', 'txt_in', 'timestep'])
+if GPU_memory_mode == "sequential_cpu_offload":
+    pipeline.enable_sequential_cpu_offload(device=device)
+elif GPU_memory_mode == "model_group_offload":
+    register_auto_device_hook(pipeline.transformer)
+    safe_enable_group_offload(pipeline, onload_device=device, offload_device="cpu", offload_type="leaf_level", use_stream=True)
+elif GPU_memory_mode == "model_cpu_offload_and_qfloat8":
+    convert_model_weight_to_float8(transformer, exclude_module_name=["img_in", "txt_in", "timestep"], device=device)
+    convert_weight_dtype_wrapper(transformer, weight_dtype)
+    pipeline.enable_model_cpu_offload(device=device)
+elif GPU_memory_mode == "model_cpu_offload":
+    pipeline.enable_model_cpu_offload(device=device)
+elif GPU_memory_mode == "model_full_load_and_qfloat8":
+    convert_model_weight_to_float8(transformer, exclude_module_name=["img_in", "txt_in", "timestep"], device=device)
+    convert_weight_dtype_wrapper(transformer, weight_dtype)
+    pipeline.to(device=device)
+else:
+    pipeline.to(device=device)
 
 generator = torch.Generator(device=device).manual_seed(seed)
 

examples/flux2_fun/predict_t2i_control_ref.py

@@ -1,9 +1,11 @@
 import os
 import sys
 
+import numpy as np
 import torch
 from diffusers import FlowMatchEulerDiscreteScheduler
 from omegaconf import OmegaConf
+from PIL import Image
 
 current_file_path = os.path.abspath(__file__)
 project_roots = [os.path.dirname(current_file_path), os.path.dirname(os.path.dirname(current_file_path)), os.path.dirname(os.path.dirname(os.path.dirname(current_file_path)))]
@@ -12,16 +14,23 @@ for project_root in project_roots:
 
 from videox_fun.dist import set_multi_gpus_devices, shard_model
 from videox_fun.models import (AutoencoderKLFlux2,
-                               Flux2ControlTransformer2DModel,
                                Mistral3ForConditionalGeneration,
-                               PixtralProcessor)
+                               PixtralProcessor, Flux2ControlTransformer2DModel)
+from videox_fun.models.cache_utils import get_teacache_coefficients
 from videox_fun.pipeline import Flux2ControlPipeline
-from videox_fun.utils import (FlowDPMSolverMultistepScheduler,
-                              FlowUniPCMultistepScheduler,
-                              apply_gpu_memory_mode, get_image,
-                              get_image_latent, merge_lora, unmerge_lora)
+from videox_fun.utils import (register_auto_device_hook,
+                              safe_enable_group_offload)
+from videox_fun.utils.fm_solvers import FlowDPMSolverMultistepScheduler
+from videox_fun.utils.fm_solvers_unipc import FlowUniPCMultistepScheduler
+from videox_fun.utils.fp8_optimization import (convert_model_weight_to_float8,
+                                               convert_weight_dtype_wrapper)
+from videox_fun.utils.lora_utils import merge_lora, unmerge_lora
+from videox_fun.utils.utils import (filter_kwargs, get_image, get_image_latent,
+                                    get_image_to_video_latent,
+                                    get_video_to_video_latent,
+                                    save_videos_grid)
 
-# GPU memory mode, which can be chosen in [model_full_load, model_full_load_and_qfloat8, model_cpu_offload, model_cpu_offload_and_qfloat8, model_group_offload, sequential_cpu_offload].
+# GPU memory mode, which can be chosen in [model_full_load, model_full_load_and_qfloat8, model_cpu_offload, model_cpu_offload_and_qfloat8, sequential_cpu_offload].
 # model_full_load means that the entire model will be moved to the GPU.
 # 
 # model_full_load_and_qfloat8 means that the entire model will be moved to the GPU,
@@ -68,7 +77,7 @@ lora_path           = None
 sample_size         = [1728, 992]
 
 # Use torch.float16 if GPU does not support torch.bfloat16
-# Some graphics cards, such as v100, 2080ti, do not support torch.bfloat16
+# ome graphics cards, such as v100, 2080ti, do not support torch.bfloat16
 weight_dtype        = torch.bfloat16
 image               = "asset/8.png"
 control_image       = "asset/pose.jpg"
@@ -99,7 +108,7 @@ transformer = Flux2ControlTransformer2DModel.from_pretrained(
 if transformer_path is not None:
     print(f"From checkpoint: {transformer_path}")
     if transformer_path.endswith("safetensors"):
-        from safetensors.torch import load_file
+        from safetensors.torch import load_file, safe_open
         state_dict = load_file(transformer_path)
     else:
         state_dict = torch.load(transformer_path, map_location="cpu")
@@ -117,7 +126,7 @@ vae = AutoencoderKLFlux2.from_pretrained(
 if vae_path is not None:
     print(f"From checkpoint: {vae_path}")
     if vae_path.endswith("safetensors"):
-        from safetensors.torch import load_file
+        from safetensors.torch import load_file, safe_open
         state_dict = load_file(vae_path)
     else:
         state_dict = torch.load(vae_path, map_location="cpu")
@@ -136,7 +145,7 @@ text_encoder = Mistral3ForConditionalGeneration.from_pretrained(
 )
 
 # Get Scheduler
-Chosen_Scheduler = {
+Chosen_Scheduler = scheduler_dict = {
     "Flow": FlowMatchEulerDiscreteScheduler,
     "Flow_Unipc": FlowUniPCMultistepScheduler,
     "Flow_DPM++": FlowDPMSolverMultistepScheduler,
@@ -171,10 +180,23 @@ if compile_dit:
         pipeline.transformer.transformer_blocks[i] = torch.compile(pipeline.transformer.transformer_blocks[i])
     print("Add Compile")
 
-# Quantize (when the mode carries an "_and_<quant>" suffix) and then place the pipeline.
-# The order lives inside the helper: quantization has to happen before the offload hooks
-# are installed.
-apply_gpu_memory_mode(pipeline, GPU_memory_mode, device, weight_dtype, exclude_module_name=['img_in', 'txt_in', 'timestep'])
+if GPU_memory_mode == "sequential_cpu_offload":
+    pipeline.enable_sequential_cpu_offload(device=device)
+elif GPU_memory_mode == "model_group_offload":
+    register_auto_device_hook(pipeline.transformer)
+    safe_enable_group_offload(pipeline, onload_device=device, offload_device="cpu", offload_type="leaf_level", use_stream=True)
+elif GPU_memory_mode == "model_cpu_offload_and_qfloat8":
+    convert_model_weight_to_float8(transformer, exclude_module_name=["img_in", "txt_in", "timestep"], device=device)
+    convert_weight_dtype_wrapper(transformer, weight_dtype)
+    pipeline.enable_model_cpu_offload(device=device)
+elif GPU_memory_mode == "model_cpu_offload":
+    pipeline.enable_model_cpu_offload(device=device)
+elif GPU_memory_mode == "model_full_load_and_qfloat8":
+    convert_model_weight_to_float8(transformer, exclude_module_name=["img_in", "txt_in", "timestep"], device=device)
+    convert_weight_dtype_wrapper(transformer, weight_dtype)
+    pipeline.to(device=device)
+else:
+    pipeline.to(device=device)
 
 generator = torch.Generator(device=device).manual_seed(seed)
 

examples/hunyuanvideo/predict_i2v.py

@@ -4,6 +4,8 @@ import sys
 import numpy as np
 import torch
 from diffusers import FlowMatchEulerDiscreteScheduler
+from diffusers.utils import export_to_video
+from omegaconf import OmegaConf
 from PIL import Image
 
 current_file_path = os.path.abspath(__file__)
@@ -11,20 +13,28 @@ project_roots = [os.path.dirname(current_file_path), os.path.dirname(os.path.dir
 for project_root in project_roots:
     sys.path.insert(0, project_root) if project_root not in sys.path else None
 
+from diffusers.schedulers.scheduling_unipc_multistep import \
+    UniPCMultistepScheduler
 
 from videox_fun.dist import set_multi_gpus_devices, shard_model
-from videox_fun.models import (AutoencoderKLHunyuanVideo, CLIPImageProcessor,
-                               CLIPTextModel, CLIPTokenizer,
-                               HunyuanVideoTransformer3DModel,
-                               LlamaTokenizerFast,
-                               LlavaForConditionalGeneration)
-from videox_fun.pipeline import HunyuanVideoI2VPipeline
-from videox_fun.utils import (FlowDPMSolverMultistepScheduler,
-                              FlowUniPCMultistepScheduler,
-                              apply_gpu_memory_mode, get_image, merge_lora,
-                              save_videos_grid, unmerge_lora)
+from videox_fun.models import (AutoencoderKLHunyuanVideo, CLIPTextModel, CLIPImageProcessor,
+                               CLIPTokenizer, HunyuanVideoTransformer3DModel,
+                               LlavaForConditionalGeneration, LlamaTokenizerFast)
+from videox_fun.models.cache_utils import get_teacache_coefficients
+from videox_fun.pipeline import HunyuanVideoPipeline, HunyuanVideoI2VPipeline
+from videox_fun.utils import (register_auto_device_hook,
+                              safe_enable_group_offload)
+from videox_fun.utils.fm_solvers import FlowDPMSolverMultistepScheduler
+from videox_fun.utils.fm_solvers_unipc import FlowUniPCMultistepScheduler
+from videox_fun.utils.fp8_optimization import (convert_model_weight_to_float8,
+                                               convert_weight_dtype_wrapper,
+                                               replace_parameters_by_name)
+from videox_fun.utils.lora_utils import merge_lora, unmerge_lora
+from videox_fun.utils.utils import (filter_kwargs, get_image_to_video_latent,
+                                    save_videos_grid)
+from videox_fun.utils.utils import get_image
 
-# GPU memory mode, which can be chosen in [model_full_load, model_full_load_and_qfloat8, model_cpu_offload, model_cpu_offload_and_qfloat8, model_group_offload, sequential_cpu_offload].
+# GPU memory mode, which can be chosen in [model_full_load, model_full_load_and_qfloat8, model_cpu_offload, model_cpu_offload_and_qfloat8, sequential_cpu_offload].
 # model_full_load means that the entire model will be moved to the GPU.
 # 
 # model_full_load_and_qfloat8 means that the entire model will be moved to the GPU,
@@ -71,7 +81,7 @@ video_length        = 81
 fps                 = 16
 
 # Use torch.float16 if GPU does not support torch.bfloat16
-# Some graphics cards, such as v100, 2080ti, do not support torch.bfloat16
+# ome graphics cards, such as v100, 2080ti, do not support torch.bfloat16
 weight_dtype        = torch.bfloat16
 # If you want to generate from text, please set the validation_image_start = None and validation_image_end = None
 validation_image_start  = "asset/1.png"
@@ -96,7 +106,7 @@ transformer = HunyuanVideoTransformer3DModel.from_pretrained(
 if transformer_path is not None:
     print(f"From checkpoint: {transformer_path}")
     if transformer_path.endswith("safetensors"):
-        from safetensors.torch import load_file
+        from safetensors.torch import load_file, safe_open
         state_dict = load_file(transformer_path)
     else:
         state_dict = torch.load(transformer_path, map_location="cpu")
@@ -113,7 +123,7 @@ vae = AutoencoderKLHunyuanVideo.from_pretrained(
 if vae_path is not None:
     print(f"From checkpoint: {vae_path}")
     if vae_path.endswith("safetensors"):
-        from safetensors.torch import load_file
+        from safetensors.torch import load_file, safe_open
         state_dict = load_file(vae_path)
     else:
         state_dict = torch.load(vae_path, map_location="cpu")
@@ -152,7 +162,7 @@ image_processor = CLIPImageProcessor.from_pretrained(
 )
 
 # Get Scheduler
-Chosen_Scheduler = {
+Chosen_Scheduler = scheduler_dict = {
     "Flow": FlowMatchEulerDiscreteScheduler,
     "Flow_Unipc": FlowUniPCMultistepScheduler,
     "Flow_DPM++": FlowDPMSolverMultistepScheduler,
@@ -189,10 +199,23 @@ if compile_dit:
         pipeline.transformer.blocks[i] = torch.compile(pipeline.transformer.blocks[i])
     print("Add Compile")
 
-# Quantize (when the mode carries an "_and_<quant>" suffix) and then place the pipeline.
-# The order lives inside the helper: quantization has to happen before the offload hooks
-# are installed.
-apply_gpu_memory_mode(pipeline, GPU_memory_mode, device, weight_dtype, exclude_module_name=['x_embedder', 'context_embedder', 'time_text_embed', 'rope', 'proj_out'])
+if GPU_memory_mode == "sequential_cpu_offload":
+    pipeline.enable_sequential_cpu_offload(device=device)
+elif GPU_memory_mode == "model_group_offload":
+    register_auto_device_hook(pipeline.transformer)
+    safe_enable_group_offload(pipeline, onload_device=device, offload_device="cpu", offload_type="leaf_level", use_stream=True)
+elif GPU_memory_mode == "model_cpu_offload_and_qfloat8":
+    convert_model_weight_to_float8(transformer, exclude_module_name=["x_embedder", "context_embedder", "time_text_embed", "rope", "proj_out"], device=device)
+    convert_weight_dtype_wrapper(transformer, weight_dtype)
+    pipeline.enable_model_cpu_offload(device=device)
+elif GPU_memory_mode == "model_cpu_offload":
+    pipeline.enable_model_cpu_offload(device=device)
+elif GPU_memory_mode == "model_full_load_and_qfloat8":
+    convert_model_weight_to_float8(transformer, exclude_module_name=["x_embedder", "context_embedder", "time_text_embed", "rope", "proj_out"], device=device)
+    convert_weight_dtype_wrapper(transformer, weight_dtype)
+    pipeline.to(device=device)
+else:
+    pipeline.to(device=device)
 
 generator = torch.Generator(device=device).manual_seed(seed)
 

examples/hunyuanvideo/predict_t2v.py

@@ -4,6 +4,8 @@ import sys
 import numpy as np
 import torch
 from diffusers import FlowMatchEulerDiscreteScheduler
+from diffusers.utils import export_to_video
+from omegaconf import OmegaConf
 from PIL import Image
 
 current_file_path = os.path.abspath(__file__)
@@ -11,18 +13,27 @@ project_roots = [os.path.dirname(current_file_path), os.path.dirname(os.path.dir
 for project_root in project_roots:
     sys.path.insert(0, project_root) if project_root not in sys.path else None
 
+from diffusers.schedulers.scheduling_unipc_multistep import \
+    UniPCMultistepScheduler
 
 from videox_fun.dist import set_multi_gpus_devices, shard_model
 from videox_fun.models import (AutoencoderKLHunyuanVideo, CLIPTextModel,
                                CLIPTokenizer, HunyuanVideoTransformer3DModel,
                                LlamaModel, LlamaTokenizerFast)
+from videox_fun.models.cache_utils import get_teacache_coefficients
 from videox_fun.pipeline import HunyuanVideoPipeline
-from videox_fun.utils import (FlowDPMSolverMultistepScheduler,
-                              FlowUniPCMultistepScheduler,
-                              apply_gpu_memory_mode, merge_lora,
-                              save_videos_grid, unmerge_lora)
+from videox_fun.utils import (register_auto_device_hook,
+                              safe_enable_group_offload)
+from videox_fun.utils.fm_solvers import FlowDPMSolverMultistepScheduler
+from videox_fun.utils.fm_solvers_unipc import FlowUniPCMultistepScheduler
+from videox_fun.utils.fp8_optimization import (convert_model_weight_to_float8,
+                                               convert_weight_dtype_wrapper,
+                                               replace_parameters_by_name)
+from videox_fun.utils.lora_utils import merge_lora, unmerge_lora
+from videox_fun.utils.utils import (filter_kwargs, get_image_to_video_latent,
+                                    save_videos_grid)
 
-# GPU memory mode, which can be chosen in [model_full_load, model_full_load_and_qfloat8, model_cpu_offload, model_cpu_offload_and_qfloat8, model_group_offload, sequential_cpu_offload].
+# GPU memory mode, which can be chosen in [model_full_load, model_full_load_and_qfloat8, model_cpu_offload, model_cpu_offload_and_qfloat8, sequential_cpu_offload].
 # model_full_load means that the entire model will be moved to the GPU.
 # 
 # model_full_load_and_qfloat8 means that the entire model will be moved to the GPU,
@@ -69,7 +80,7 @@ video_length        = 81
 fps                 = 16
 
 # Use torch.float16 if GPU does not support torch.bfloat16
-# Some graphics cards, such as v100, 2080ti, do not support torch.bfloat16
+# ome graphics cards, such as v100, 2080ti, do not support torch.bfloat16
 weight_dtype        = torch.bfloat16
 prompt              = "1girl, black_hair, brown_eyes, earrings, freckles, grey_background, jewelry, lips, long_hair, looking_at_viewer, nose, piercing, realistic, red_lips, solo, upper_body"
 negative_prompt     = "The video is not of a high quality, it has a low resolution. Watermark present in each frame. The background is solid. Strange body and strange trajectory. Distortion. "
@@ -90,7 +101,7 @@ transformer = HunyuanVideoTransformer3DModel.from_pretrained(
 if transformer_path is not None:
     print(f"From checkpoint: {transformer_path}")
     if transformer_path.endswith("safetensors"):
-        from safetensors.torch import load_file
+        from safetensors.torch import load_file, safe_open
         state_dict = load_file(transformer_path)
     else:
         state_dict = torch.load(transformer_path, map_location="cpu")
@@ -107,7 +118,7 @@ vae = AutoencoderKLHunyuanVideo.from_pretrained(
 if vae_path is not None:
     print(f"From checkpoint: {vae_path}")
     if vae_path.endswith("safetensors"):
-        from safetensors.torch import load_file
+        from safetensors.torch import load_file, safe_open
         state_dict = load_file(vae_path)
     else:
         state_dict = torch.load(vae_path, map_location="cpu")
@@ -141,7 +152,7 @@ text_encoder_2 = CLIPTextModel.from_pretrained(
 )
 
 # Get Scheduler
-Chosen_Scheduler = {
+Chosen_Scheduler = scheduler_dict = {
     "Flow": FlowMatchEulerDiscreteScheduler,
     "Flow_Unipc": FlowUniPCMultistepScheduler,
     "Flow_DPM++": FlowDPMSolverMultistepScheduler,
@@ -177,10 +188,23 @@ if compile_dit:
         pipeline.transformer.blocks[i] = torch.compile(pipeline.transformer.blocks[i])
     print("Add Compile")
 
-# Quantize (when the mode carries an "_and_<quant>" suffix) and then place the pipeline.
-# The order lives inside the helper: quantization has to happen before the offload hooks
-# are installed.
-apply_gpu_memory_mode(pipeline, GPU_memory_mode, device, weight_dtype, exclude_module_name=['x_embedder', 'context_embedder', 'time_text_embed', 'rope', 'proj_out'])
+if GPU_memory_mode == "sequential_cpu_offload":
+    pipeline.enable_sequential_cpu_offload(device=device)
+elif GPU_memory_mode == "model_group_offload":
+    register_auto_device_hook(pipeline.transformer)
+    safe_enable_group_offload(pipeline, onload_device=device, offload_device="cpu", offload_type="leaf_level", use_stream=True)
+elif GPU_memory_mode == "model_cpu_offload_and_qfloat8":
+    convert_model_weight_to_float8(transformer, exclude_module_name=["x_embedder", "context_embedder", "time_text_embed", "rope", "proj_out"], device=device)
+    convert_weight_dtype_wrapper(transformer, weight_dtype)
+    pipeline.enable_model_cpu_offload(device=device)
+elif GPU_memory_mode == "model_cpu_offload":
+    pipeline.enable_model_cpu_offload(device=device)
+elif GPU_memory_mode == "model_full_load_and_qfloat8":
+    convert_model_weight_to_float8(transformer, exclude_module_name=["x_embedder", "context_embedder", "time_text_embed", "rope", "proj_out"], device=device)
+    convert_weight_dtype_wrapper(transformer, weight_dtype)
+    pipeline.to(device=device)
+else:
+    pipeline.to(device=device)
 
 generator = torch.Generator(device=device).manual_seed(seed)
 

examples/infinitetalk/predict_s2v.py

@@ -13,19 +13,25 @@ for project_root in project_roots:
     sys.path.insert(0, project_root) if project_root not in sys.path else None
 
 from videox_fun.dist import set_multi_gpus_devices, shard_model
-from videox_fun.models import (AutoencoderKLWan, AutoTokenizer, CLIPModel,
-                               InfiniteTalkAudioEncoder,
-                               InfiniteTalkTransformer3DModel,
-                               WanT5EncoderModel)
 from videox_fun.models.cache_utils import get_teacache_coefficients
+from videox_fun.models import (AutoencoderKLWan, AutoencoderKLWan3_8,
+                               AutoTokenizer, CLIPModel,
+                               InfiniteTalkTransformer3DModel, InfiniteTalkAudioEncoder,
+                               WanT5EncoderModel)
 from videox_fun.pipeline import InfiniteTalkPipeline
-from videox_fun.utils import (FlowDPMSolverMultistepScheduler,
-                              FlowUniPCMultistepScheduler,
-                              apply_gpu_memory_mode, filter_kwargs, get_image,
-                              get_image_latent, merge_lora, merge_video_audio,
-                              save_videos_grid, unmerge_lora)
+from videox_fun.utils.fm_solvers import FlowDPMSolverMultistepScheduler
+from videox_fun.utils.fm_solvers_unipc import FlowUniPCMultistepScheduler
+from videox_fun.utils import (register_auto_device_hook,
+                              safe_enable_group_offload)
+from videox_fun.utils.fp8_optimization import (convert_model_weight_to_float8,
+                                               convert_weight_dtype_wrapper,
+                                               replace_parameters_by_name)
+from videox_fun.utils.lora_utils import merge_lora, unmerge_lora
+from videox_fun.utils.utils import (filter_kwargs, get_image_latent, get_image, 
+                                    get_video_to_video_latent,
+                                    merge_video_audio, save_videos_grid)
 
-# GPU memory mode, which can be chosen in [model_full_load, model_full_load_and_qfloat8, model_cpu_offload, model_cpu_offload_and_qfloat8, model_group_offload, sequential_cpu_offload].
+# GPU memory mode, which can be chosen in [model_full_load, model_full_load_and_qfloat8, model_cpu_offload, model_cpu_offload_and_qfloat8, sequential_cpu_offload].
 # model_full_load means that the entire model will be moved to the GPU.
 # 
 # model_full_load_and_qfloat8 means that the entire model will be moved to the GPU,
@@ -89,7 +95,7 @@ segment_frame_length    = 81
 fps                     = 25
 
 # Use torch.float16 if GPU does not support torch.bfloat16
-# Some graphics cards, such as v100, 2080ti, do not support torch.bfloat16
+# ome graphics cards, such as v100, 2080ti, do not support torch.bfloat16
 weight_dtype            = torch.bfloat16
 # The path of the reference image
 ref_image               = "asset/8.png"
@@ -144,7 +150,7 @@ vae = AutoencoderKLWan.from_pretrained(
 if vae_path is not None:
     print(f"From checkpoint: {vae_path}")
     if vae_path.endswith("safetensors"):
-        from safetensors.torch import load_file
+        from safetensors.torch import load_file, safe_open
         state_dict = load_file(vae_path)
     else:
         state_dict = torch.load(vae_path, map_location="cpu")
@@ -180,7 +186,7 @@ clip_image_encoder = CLIPModel.from_pretrained(
 clip_image_encoder = clip_image_encoder.eval()
 
 # Get Scheduler
-Chosen_Scheduler = {
+Chosen_Scheduler = scheduler_dict = {
     "Flow": FlowMatchEulerDiscreteScheduler,
     "Flow_Unipc": FlowUniPCMultistepScheduler,
     "Flow_DPM++": FlowDPMSolverMultistepScheduler,
@@ -218,10 +224,25 @@ if compile_dit:
         pipeline.transformer.blocks[i] = torch.compile(pipeline.transformer.blocks[i])
     print("Add Compile")
 
-# Quantize (when the mode carries an "_and_<quant>" suffix) and then place the pipeline.
-# The order lives inside the helper: quantization has to happen before the offload hooks
-# are installed.
-apply_gpu_memory_mode(pipeline, GPU_memory_mode, device, weight_dtype)
+if GPU_memory_mode == "sequential_cpu_offload":
+    replace_parameters_by_name(transformer, ["modulation",], device=device)
+    transformer.freqs = transformer.freqs.to(device=device)
+    pipeline.enable_sequential_cpu_offload(device=device)
+elif GPU_memory_mode == "model_group_offload":
+    register_auto_device_hook(pipeline.transformer)
+    safe_enable_group_offload(pipeline, onload_device=device, offload_device="cpu", offload_type="leaf_level", use_stream=True)
+elif GPU_memory_mode == "model_cpu_offload_and_qfloat8":
+    convert_model_weight_to_float8(transformer, exclude_module_name=["modulation",], device=device)
+    convert_weight_dtype_wrapper(transformer, weight_dtype)
+    pipeline.enable_model_cpu_offload(device=device)
+elif GPU_memory_mode == "model_cpu_offload":
+    pipeline.enable_model_cpu_offload(device=device)
+elif GPU_memory_mode == "model_full_load_and_qfloat8":
+    convert_model_weight_to_float8(transformer, exclude_module_name=["modulation",], device=device)
+    convert_weight_dtype_wrapper(transformer, weight_dtype)
+    pipeline.to(device=device)
+else:
+    pipeline.to(device=device)
 
 coefficients = get_teacache_coefficients(model_name) if enable_teacache else None
 if coefficients is not None:

examples/lens/predict_t2i.py

@@ -13,11 +13,15 @@ from videox_fun.dist import set_multi_gpus_devices, shard_model
 from videox_fun.models import (AutoencoderKLFlux2, AutoTokenizer,
                                LensGptOssEncoder, LensTransformer2DModel)
 from videox_fun.pipeline import LensPipeline
-from videox_fun.utils import (FlowDPMSolverMultistepScheduler,
-                              FlowUniPCMultistepScheduler,
-                              apply_gpu_memory_mode, merge_lora, unmerge_lora)
+from videox_fun.utils import (register_auto_device_hook,
+                              safe_enable_group_offload)
+from videox_fun.utils.fm_solvers import FlowDPMSolverMultistepScheduler
+from videox_fun.utils.fm_solvers_unipc import FlowUniPCMultistepScheduler
+from videox_fun.utils.fp8_optimization import (convert_model_weight_to_float8,
+                                               convert_weight_dtype_wrapper)
+from videox_fun.utils.lora_utils import merge_lora, unmerge_lora
 
-# GPU memory mode, which can be chosen in [model_full_load, model_full_load_and_qfloat8, model_cpu_offload, model_cpu_offload_and_qfloat8, model_group_offload, sequential_cpu_offload].
+# GPU memory mode, which can be chosen in [model_full_load, model_full_load_and_qfloat8, model_cpu_offload, model_cpu_offload_and_qfloat8, sequential_cpu_offload].
 # model_full_load means that the entire model will be moved to the GPU.
 # 
 # model_full_load_and_qfloat8 means that the entire model will be moved to the GPU,
@@ -62,7 +66,7 @@ lora_path           = None
 sample_size         = [1728, 992]
 
 # Use torch.float16 if GPU does not support torch.bfloat16
-# Some graphics cards, such as v100, 2080ti, do not support torch.bfloat16
+# ome graphics cards, such as v100, 2080ti, do not support torch.bfloat16
 weight_dtype        = torch.bfloat16
 # Set to True on A100/V100 to dequantize MXFP4 GPT-OSS weights.
 dequantize_mxfp4    = False
@@ -87,7 +91,7 @@ transformer = LensTransformer2DModel.from_pretrained(
 if transformer_path is not None:
     print(f"From checkpoint: {transformer_path}")
     if transformer_path.endswith("safetensors"):
-        from safetensors.torch import load_file
+        from safetensors.torch import load_file, safe_open
         state_dict = load_file(transformer_path)
     else:
         state_dict = torch.load(transformer_path, map_location="cpu")
@@ -105,7 +109,7 @@ vae = AutoencoderKLFlux2.from_pretrained(
 if vae_path is not None:
     print(f"From checkpoint: {vae_path}")
     if vae_path.endswith("safetensors"):
-        from safetensors.torch import load_file
+        from safetensors.torch import load_file, safe_open
         state_dict = load_file(vae_path)
     else:
         state_dict = torch.load(vae_path, map_location="cpu")
@@ -132,7 +136,7 @@ text_encoder = LensGptOssEncoder.from_pretrained(
 )
 
 # Get Scheduler
-Chosen_Scheduler = {
+Chosen_Scheduler = scheduler_dict = {
     "Flow": FlowMatchEulerDiscreteScheduler,
     "Flow_Unipc": FlowUniPCMultistepScheduler,
     "Flow_DPM++": FlowDPMSolverMultistepScheduler,
@@ -167,10 +171,23 @@ if compile_dit:
         pipeline.transformer.transformer_blocks[i] = torch.compile(pipeline.transformer.transformer_blocks[i])
     print("Add Compile")
 
-# Quantize (when the mode carries an "_and_<quant>" suffix) and then place the pipeline.
-# The order lives inside the helper: quantization has to happen before the offload hooks
-# are installed.
-apply_gpu_memory_mode(pipeline, GPU_memory_mode, device, weight_dtype, exclude_module_name=['img_in', 'txt_in', 'timestep'])
+if GPU_memory_mode == "sequential_cpu_offload":
+    pipeline.enable_sequential_cpu_offload(device=device)
+elif GPU_memory_mode == "model_group_offload":
+    register_auto_device_hook(pipeline.transformer)
+    safe_enable_group_offload(pipeline, onload_device=device, offload_device="cpu", offload_type="leaf_level", use_stream=True)
+elif GPU_memory_mode == "model_cpu_offload_and_qfloat8":
+    convert_model_weight_to_float8(transformer, exclude_module_name=["img_in", "txt_in", "timestep"], device=device)
+    convert_weight_dtype_wrapper(transformer, weight_dtype)
+    pipeline.enable_model_cpu_offload(device=device)
+elif GPU_memory_mode == "model_cpu_offload":
+    pipeline.enable_model_cpu_offload(device=device)
+elif GPU_memory_mode == "model_full_load_and_qfloat8":
+    convert_model_weight_to_float8(transformer, exclude_module_name=["img_in", "txt_in", "timestep"], device=device)
+    convert_weight_dtype_wrapper(transformer, weight_dtype)
+    pipeline.to(device=device)
+else:
+    pipeline.to(device=device)
 
 generator = torch.Generator(device=device).manual_seed(seed)
 

examples/lingbot_video/predict_i2v.py

@@ -18,9 +18,13 @@ from videox_fun.models import (AutoencoderKLQwenImage,
 from videox_fun.models.lingbot_video_rewriter import ensure_json_caption
 from videox_fun.pipeline import LingBotVideoI2VPipeline
 from videox_fun.pipeline.pipeline_lingbot_video import DEFAULT_NEGATIVE_PROMPT
-from videox_fun.utils import (FlowUniPCMultistepScheduler,
-                              apply_gpu_memory_mode, merge_lora,
-                              save_videos_grid, unmerge_lora)
+from videox_fun.utils import (register_auto_device_hook,
+                              safe_enable_group_offload)
+from videox_fun.utils.fm_solvers_unipc import FlowUniPCMultistepScheduler
+from videox_fun.utils.fp8_optimization import (convert_model_weight_to_float8,
+                                               convert_weight_dtype_wrapper)
+from videox_fun.utils.lora_utils import merge_lora, unmerge_lora
+from videox_fun.utils.utils import save_videos_grid
 
 # GPU memory mode, which can be chosen in [model_full_load, model_full_load_and_qfloat8, model_cpu_offload, model_cpu_offload_and_qfloat8, model_group_offload].
 # model_full_load means that the entire model will be moved to the GPU.
@@ -72,7 +76,7 @@ video_length        = 81
 fps                 = 24
 
 # Use torch.float16 if GPU does not support torch.bfloat16
-# Some graphics cards, such as v100, 2080ti, do not support torch.bfloat16
+# some graphics cards, such as v100, 2080ti, do not support torch.bfloat16
 weight_dtype        = torch.bfloat16
 # The condition image is used twice: as Qwen3-VL visual input and as a clean
 # first-frame latent injected into the diffusion latent (ti2v).
@@ -104,6 +108,7 @@ prompt = ensure_json_caption(
 
 device = set_multi_gpus_devices(ulysses_degree, ring_degree)
 
+
 transformer = LingBotVideoTransformer3DModel.from_pretrained(
     os.path.join(model_name, "transformer"),
     low_cpu_mem_usage=True,
@@ -115,7 +120,7 @@ transformer = transformer.to(weight_dtype)
 if transformer_path is not None:
     print(f"From checkpoint: {transformer_path}")
     if transformer_path.endswith("safetensors"):
-        from safetensors.torch import load_file
+        from safetensors.torch import load_file, safe_open
         state_dict = load_file(transformer_path)
     else:
         state_dict = torch.load(transformer_path, map_location="cpu")
@@ -133,7 +138,7 @@ vae = AutoencoderKLQwenImage.from_pretrained(
 if vae_path is not None:
     print(f"From checkpoint: {vae_path}")
     if vae_path.endswith("safetensors"):
-        from safetensors.torch import load_file
+        from safetensors.torch import load_file, safe_open
         state_dict = load_file(vae_path)
     else:
         state_dict = torch.load(vae_path, map_location="cpu")
@@ -155,7 +160,7 @@ text_encoder = Qwen3VLForConditionalGeneration.from_pretrained(
 )
 
 # Get Scheduler
-Chosen_Scheduler = {
+Chosen_Scheduler = scheduler_dict = {
     "Flow_Unipc": FlowUniPCMultistepScheduler,
 }[sampler_name]
 scheduler = Chosen_Scheduler.from_pretrained(
@@ -172,10 +177,21 @@ pipeline = LingBotVideoI2VPipeline(
     scheduler=scheduler,
 )
 
-# Quantize (when the mode carries an "_and_<quant>" suffix) and then place the pipeline.
-# The order lives inside the helper: quantization has to happen before the offload hooks
-# are installed.
-apply_gpu_memory_mode(pipeline, GPU_memory_mode, device, weight_dtype, exclude_module_name=['time_embedder', 'time_modulation', 'text_embedder', 'norm', 'router', 'scale_shift_table', 'proj_out'])
+if GPU_memory_mode == "model_group_offload":
+    register_auto_device_hook(pipeline.transformer)
+    safe_enable_group_offload(pipeline, onload_device=device, offload_device="cpu", offload_type="leaf_level", use_stream=True)
+elif GPU_memory_mode == "model_cpu_offload_and_qfloat8":
+    convert_model_weight_to_float8(transformer, exclude_module_name=["time_embedder", "time_modulation", "text_embedder", "norm", "router", "scale_shift_table", "proj_out"], device=device)
+    convert_weight_dtype_wrapper(transformer, weight_dtype)
+    pipeline.enable_model_cpu_offload(device=device)
+elif GPU_memory_mode == "model_cpu_offload":
+    pipeline.enable_model_cpu_offload(device=device)
+elif GPU_memory_mode == "model_full_load_and_qfloat8":
+    convert_model_weight_to_float8(transformer, exclude_module_name=["time_embedder", "time_modulation", "text_embedder", "norm", "router", "scale_shift_table", "proj_out"], device=device)
+    convert_weight_dtype_wrapper(transformer, weight_dtype)
+    pipeline.to(device=device)
+else:
+    pipeline.to(device=device)
 
 if ulysses_degree > 1 or ring_degree > 1:
     from functools import partial

examples/lingbot_video/predict_t2v.py

@@ -15,12 +15,17 @@ from videox_fun.dist import set_multi_gpus_devices, shard_model
 from videox_fun.models import (AutoencoderKLQwenImage,
                                LingBotVideoTransformer3DModel,
                                Qwen3VLForConditionalGeneration)
-from videox_fun.models.lingbot_video_rewriter import ensure_json_caption
 from videox_fun.pipeline import LingBotVideoPipeline
 from videox_fun.pipeline.pipeline_lingbot_video import DEFAULT_NEGATIVE_PROMPT
-from videox_fun.utils import (FlowUniPCMultistepScheduler,
-                              apply_gpu_memory_mode, merge_lora,
-                              save_videos_grid, unmerge_lora)
+from videox_fun.utils import (register_auto_device_hook,
+                              safe_enable_group_offload)
+from videox_fun.utils.fp8_optimization import (convert_model_weight_to_float8,
+                                               convert_weight_dtype_wrapper)
+from videox_fun.utils.fm_solvers_unipc import FlowUniPCMultistepScheduler
+from videox_fun.utils.lora_utils import merge_lora, unmerge_lora
+from videox_fun.utils.utils import save_videos_grid
+
+from videox_fun.models.lingbot_video_rewriter import ensure_json_caption
 
 # GPU memory mode, which can be chosen in [model_full_load, model_full_load_and_qfloat8, model_cpu_offload, model_cpu_offload_and_qfloat8, model_group_offload].
 # model_full_load means that the entire model will be moved to the GPU.
@@ -73,7 +78,7 @@ video_length        = 81
 fps                 = 24
 
 # Use torch.float16 if GPU does not support torch.bfloat16
-# Some graphics cards, such as v100, 2080ti, do not support torch.bfloat16
+# some graphics cards, such as v100, 2080ti, do not support torch.bfloat16
 weight_dtype        = torch.bfloat16
 # prompts
 # Write a plain natural-language prompt: it is ALWAYS rewritten into the
@@ -106,6 +111,7 @@ prompt = ensure_json_caption(
 
 device = set_multi_gpus_devices(ulysses_degree, ring_degree)
 
+
 transformer = LingBotVideoTransformer3DModel.from_pretrained(
     os.path.join(model_name, "transformer"),
     low_cpu_mem_usage=True,
@@ -117,7 +123,7 @@ transformer = transformer.to(weight_dtype)
 if transformer_path is not None:
     print(f"From checkpoint: {transformer_path}")
     if transformer_path.endswith("safetensors"):
-        from safetensors.torch import load_file
+        from safetensors.torch import load_file, safe_open
         state_dict = load_file(transformer_path)
     else:
         state_dict = torch.load(transformer_path, map_location="cpu")
@@ -135,7 +141,7 @@ vae = AutoencoderKLQwenImage.from_pretrained(
 if vae_path is not None:
     print(f"From checkpoint: {vae_path}")
     if vae_path.endswith("safetensors"):
-        from safetensors.torch import load_file
+        from safetensors.torch import load_file, safe_open
         state_dict = load_file(vae_path)
     else:
         state_dict = torch.load(vae_path, map_location="cpu")
@@ -157,7 +163,7 @@ text_encoder = Qwen3VLForConditionalGeneration.from_pretrained(
 )
 
 # Get Scheduler
-Chosen_Scheduler = {
+Chosen_Scheduler = scheduler_dict = {
     "Flow_Unipc": FlowUniPCMultistepScheduler,
 }[sampler_name]
 scheduler = Chosen_Scheduler.from_pretrained(
@@ -174,10 +180,21 @@ pipeline = LingBotVideoPipeline(
     scheduler=scheduler,
 )
 
-# Quantize (when the mode carries an "_and_<quant>" suffix) and then place the pipeline.
-# The order lives inside the helper: quantization has to happen before the offload hooks
-# are installed.
-apply_gpu_memory_mode(pipeline, GPU_memory_mode, device, weight_dtype, exclude_module_name=['time_embedder', 'time_modulation', 'text_embedder', 'norm', 'router', 'scale_shift_table', 'proj_out'])
+if GPU_memory_mode == "model_group_offload":
+    register_auto_device_hook(pipeline.transformer)
+    safe_enable_group_offload(pipeline, onload_device=device, offload_device="cpu", offload_type="leaf_level", use_stream=True)
+elif GPU_memory_mode == "model_cpu_offload_and_qfloat8":
+    convert_model_weight_to_float8(transformer, exclude_module_name=["time_embedder", "time_modulation", "text_embedder", "norm", "router", "scale_shift_table", "proj_out"], device=device)
+    convert_weight_dtype_wrapper(transformer, weight_dtype)
+    pipeline.enable_model_cpu_offload(device=device)
+elif GPU_memory_mode == "model_cpu_offload":
+    pipeline.enable_model_cpu_offload(device=device)
+elif GPU_memory_mode == "model_full_load_and_qfloat8":
+    convert_model_weight_to_float8(transformer, exclude_module_name=["time_embedder", "time_modulation", "text_embedder", "norm", "router", "scale_shift_table", "proj_out"], device=device)
+    convert_weight_dtype_wrapper(transformer, weight_dtype)
+    pipeline.to(device=device)
+else:
+    pipeline.to(device=device)
 
 if ulysses_degree > 1 or ring_degree > 1:
     from functools import partial

examples/lingbot_video/predict_t2v_refine.py

@@ -15,12 +15,17 @@ from videox_fun.dist import set_multi_gpus_devices, shard_model
 from videox_fun.models import (AutoencoderKLQwenImage,
                                LingBotVideoTransformer3DModel,
                                Qwen3VLForConditionalGeneration)
-from videox_fun.models.lingbot_video_rewriter import ensure_json_caption
 from videox_fun.pipeline import LingBotVideoPipeline
-from videox_fun.pipeline.pipeline_lingbot_video import (
-    DEFAULT_NEGATIVE_PROMPT, prepare_refiner_latent)
-from videox_fun.utils import (FlowUniPCMultistepScheduler,
-                              apply_gpu_memory_mode, save_videos_grid)
+from videox_fun.pipeline.pipeline_lingbot_video import (DEFAULT_NEGATIVE_PROMPT,
+                                                       prepare_refiner_latent)
+from videox_fun.utils import (register_auto_device_hook,
+                              safe_enable_group_offload)
+from videox_fun.utils.fp8_optimization import (convert_model_weight_to_float8,
+                                               convert_weight_dtype_wrapper)
+from videox_fun.utils.fm_solvers_unipc import FlowUniPCMultistepScheduler
+from videox_fun.utils.utils import save_videos_grid
+
+from videox_fun.models.lingbot_video_rewriter import ensure_json_caption
 
 # Two-stage LingBot-Video t2v: the base DiT samples at a low resolution, then the
 # "refiner" DiT re-noises the upsampled latent to sigma = refiner_t_thresh and
@@ -84,7 +89,7 @@ video_length        = 81
 fps                 = 24
 
 # Use torch.float16 if GPU does not support torch.bfloat16
-# Some graphics cards, such as v100, 2080ti, do not support torch.bfloat16
+# some graphics cards, such as v100, 2080ti, do not support torch.bfloat16
 weight_dtype        = torch.bfloat16
 # prompts
 # Write a plain natural-language prompt: it is ALWAYS rewritten into the
@@ -139,7 +144,7 @@ def load_transformer(root, subpath, checkpoint_path):
     if checkpoint_path is not None:
         print(f"From checkpoint: {checkpoint_path}")
         if checkpoint_path.endswith("safetensors"):
-            from safetensors.torch import load_file
+            from safetensors.torch import load_file, safe_open
             state_dict = load_file(checkpoint_path)
         else:
             state_dict = torch.load(checkpoint_path, map_location="cpu")
@@ -172,7 +177,7 @@ vae = AutoencoderKLQwenImage.from_pretrained(
 if vae_path is not None:
     print(f"From checkpoint: {vae_path}")
     if vae_path.endswith("safetensors"):
-        from safetensors.torch import load_file
+        from safetensors.torch import load_file, safe_open
         state_dict = load_file(vae_path)
     else:
         state_dict = torch.load(vae_path, map_location="cpu")
@@ -187,7 +192,7 @@ processor = AutoProcessor.from_pretrained(
 )
 
 # Get Scheduler
-Chosen_Scheduler = {
+Chosen_Scheduler = scheduler_dict = {
     "Flow_Unipc": FlowUniPCMultistepScheduler,
 }[sampler_name]
 scheduler = Chosen_Scheduler.from_pretrained(
@@ -226,10 +231,21 @@ pipeline = LingBotVideoPipeline(
     processor=processor,
     scheduler=scheduler,
 )
-# Quantize (when the mode carries an "_and_<quant>" suffix) and then place the pipeline.
-# The order lives inside the helper: quantization has to happen before the offload hooks
-# are installed.
-apply_gpu_memory_mode(pipeline, GPU_memory_mode, device, weight_dtype, exclude_module_name=['time_embedder', 'time_modulation', 'text_embedder', 'norm', 'router', 'scale_shift_table', 'proj_out'])
+if GPU_memory_mode == "model_group_offload":
+    register_auto_device_hook(pipeline.transformer)
+    safe_enable_group_offload(pipeline, onload_device=device, offload_device="cpu", offload_type="leaf_level", use_stream=True)
+elif GPU_memory_mode == "model_cpu_offload_and_qfloat8":
+    convert_model_weight_to_float8(transformer, exclude_module_name=["time_embedder", "time_modulation", "text_embedder", "norm", "router", "scale_shift_table", "proj_out"], device=device)
+    convert_weight_dtype_wrapper(transformer, weight_dtype)
+    pipeline.enable_model_cpu_offload(device=device)
+elif GPU_memory_mode == "model_cpu_offload":
+    pipeline.enable_model_cpu_offload(device=device)
+elif GPU_memory_mode == "model_full_load_and_qfloat8":
+    convert_model_weight_to_float8(transformer, exclude_module_name=["time_embedder", "time_modulation", "text_embedder", "norm", "router", "scale_shift_table", "proj_out"], device=device)
+    convert_weight_dtype_wrapper(transformer, weight_dtype)
+    pipeline.to(device=device)
+else:
+    pipeline.to(device=device)
 
 if ulysses_degree > 1 or ring_degree > 1:
     from functools import partial
@@ -281,12 +297,24 @@ refiner_pipeline = LingBotVideoPipeline(
     processor=processor,
     scheduler=scheduler,
 )
-# Quantize (when the mode carries an "_and_<quant>" suffix) and then place the pipeline.
-# The order lives inside the helper: quantization has to happen before the offload hooks
-# are installed.
-apply_gpu_memory_mode(refiner_pipeline, GPU_memory_mode, device, weight_dtype, exclude_module_name=["time_embedder", "time_modulation", "text_embedder", "norm", "router", "scale_shift_table", "proj_out"])
+if GPU_memory_mode == "model_group_offload":
+    register_auto_device_hook(refiner_pipeline.transformer)
+    safe_enable_group_offload(refiner_pipeline, onload_device=device, offload_device="cpu", offload_type="leaf_level", use_stream=True)
+elif GPU_memory_mode == "model_cpu_offload_and_qfloat8":
+    convert_model_weight_to_float8(refiner, exclude_module_name=["time_embedder", "time_modulation", "text_embedder", "norm", "router", "scale_shift_table", "proj_out"], device=device)
+    convert_weight_dtype_wrapper(refiner, weight_dtype)
+    refiner_pipeline.enable_model_cpu_offload(device=device)
+elif GPU_memory_mode == "model_cpu_offload":
+    refiner_pipeline.enable_model_cpu_offload(device=device)
+elif GPU_memory_mode == "model_full_load_and_qfloat8":
+    convert_model_weight_to_float8(refiner, exclude_module_name=["time_embedder", "time_modulation", "text_embedder", "norm", "router", "scale_shift_table", "proj_out"], device=device)
+    convert_weight_dtype_wrapper(refiner, weight_dtype)
+    refiner_pipeline.to(device=device)
+else:
+    refiner_pipeline.to(device=device)
 
 if ulysses_degree > 1 or ring_degree > 1:
+    from functools import partial
     refiner.enable_multi_gpus_inference()
     if fsdp_dit:
         shard_fn = partial(shard_model, device_id=device, param_dtype=weight_dtype)

examples/lingbot_world/predict_i2v.py

@@ -12,17 +12,22 @@ project_roots = [os.path.dirname(current_file_path), os.path.dirname(os.path.dir
 for project_root in project_roots:
     sys.path.insert(0, project_root) if project_root not in sys.path else None
 
-from videox_fun.data.utils import prepare_lingbot_dit_cond_dict
 from videox_fun.dist import set_multi_gpus_devices, shard_model
 from videox_fun.models import (AutoencoderKLWan, AutoencoderKLWan3_8,
                                AutoTokenizer, WanT5EncoderModel,
                                WanTransformer3DModel_LingbotWorld)
 from videox_fun.models.cache_utils import get_teacache_coefficients
 from videox_fun.pipeline import Wan2_2I2VPipeline
-from videox_fun.utils import (FlowDPMSolverMultistepScheduler,
-                              FlowUniPCMultistepScheduler,
-                              apply_gpu_memory_mode, filter_kwargs,
-                              get_image_to_video_latent, save_videos_grid)
+from videox_fun.utils import (register_auto_device_hook,
+                              safe_enable_group_offload)
+from videox_fun.utils.fm_solvers import FlowDPMSolverMultistepScheduler
+from videox_fun.utils.fm_solvers_unipc import FlowUniPCMultistepScheduler
+from videox_fun.utils.fp8_optimization import (convert_model_weight_to_float8,
+                                               convert_weight_dtype_wrapper,
+                                               replace_parameters_by_name)
+from videox_fun.data.utils import prepare_lingbot_dit_cond_dict
+from videox_fun.utils.utils import (filter_kwargs, get_image_to_video_latent,
+                                    save_videos_grid)
 
 # GPU memory mode, which can be chosen in [model_full_load, model_full_load_and_qfloat8, model_cpu_offload, model_cpu_offload_and_qfloat8, model_group_offload, sequential_cpu_offload].
 # model_full_load means that the entire model will be moved to the GPU.
@@ -92,7 +97,7 @@ video_length        = 81
 fps                 = 16
 
 # Use torch.float16 if GPU does not support torch.bfloat16
-# Some graphics cards, such as v100, 2080ti, do not support torch.bfloat16
+# some graphics cards, such as v100, 2080ti, do not support torch.bfloat16
 weight_dtype            = torch.bfloat16
 
 # Camera trajectory (poses.npy / intrinsics.npy) + reference image + prompt.
@@ -132,7 +137,7 @@ else:
 if transformer_path is not None:
     print(f"From checkpoint: {transformer_path}")
     if transformer_path.endswith("safetensors"):
-        from safetensors.torch import load_file
+        from safetensors.torch import load_file, safe_open
         state_dict = load_file(transformer_path)
     else:
         state_dict = torch.load(transformer_path, map_location="cpu")
@@ -145,7 +150,7 @@ if transformer_2 is not None:
     if transformer_high_path is not None:
         print(f"From checkpoint: {transformer_high_path}")
         if transformer_high_path.endswith("safetensors"):
-            from safetensors.torch import load_file
+            from safetensors.torch import load_file, safe_open
             state_dict = load_file(transformer_high_path)
         else:
             state_dict = torch.load(transformer_high_path, map_location="cpu")
@@ -167,7 +172,7 @@ vae = Chosen_AutoencoderKL.from_pretrained(
 if vae_path is not None:
     print(f"From checkpoint: {vae_path}")
     if vae_path.endswith("safetensors"):
-        from safetensors.torch import load_file
+        from safetensors.torch import load_file, safe_open
         state_dict = load_file(vae_path)
     else:
         state_dict = torch.load(vae_path, map_location="cpu")
@@ -191,7 +196,7 @@ text_encoder = WanT5EncoderModel.from_pretrained(
 text_encoder = text_encoder.eval()
 
 # Get Scheduler
-Chosen_Scheduler = {
+Chosen_Scheduler = scheduler_dict = {
     "Flow": FlowMatchEulerDiscreteScheduler,
     "Flow_Unipc": FlowUniPCMultistepScheduler,
     "Flow_DPM++": FlowDPMSolverMultistepScheduler,
@@ -235,11 +240,36 @@ if compile_dit:
             pipeline.transformer_2.blocks[i] = torch.compile(pipeline.transformer_2.blocks[i])
     print("Add Compile")
 
-# Quantize (when the mode carries an "_and_<quant>" suffix) and then place the pipeline.
-# The order lives inside the helper: quantization has to happen before the offload hooks
-# are installed, and both transformers of this MoE setup are handled in one call, which
-# is exactly the bookkeeping the old 30-line if/elif chain repeated per script.
-apply_gpu_memory_mode(pipeline, GPU_memory_mode, device, weight_dtype)
+if GPU_memory_mode == "sequential_cpu_offload":
+    replace_parameters_by_name(transformer, ["modulation",], device=device)
+    transformer.freqs = transformer.freqs.to(device=device)
+    if transformer_2 is not None:
+        replace_parameters_by_name(transformer_2, ["modulation",], device=device)
+        transformer_2.freqs = transformer_2.freqs.to(device=device)
+    pipeline.enable_sequential_cpu_offload(device=device)
+elif GPU_memory_mode == "model_group_offload":
+    register_auto_device_hook(pipeline.transformer)
+    if transformer_2 is not None:
+        register_auto_device_hook(pipeline.transformer_2)
+    safe_enable_group_offload(pipeline, onload_device=device, offload_device="cpu", offload_type="leaf_level", use_stream=True)
+elif GPU_memory_mode == "model_cpu_offload_and_qfloat8":
+    convert_model_weight_to_float8(transformer, exclude_module_name=["modulation",], device=device)
+    convert_weight_dtype_wrapper(transformer, weight_dtype)
+    if transformer_2 is not None:
+        convert_model_weight_to_float8(transformer_2, exclude_module_name=["modulation",], device=device)
+        convert_weight_dtype_wrapper(transformer_2, weight_dtype)
+    pipeline.enable_model_cpu_offload(device=device)
+elif GPU_memory_mode == "model_cpu_offload":
+    pipeline.enable_model_cpu_offload(device=device)
+elif GPU_memory_mode == "model_full_load_and_qfloat8":
+    convert_model_weight_to_float8(transformer, exclude_module_name=["modulation",], device=device)
+    convert_weight_dtype_wrapper(transformer, weight_dtype)
+    if transformer_2 is not None:
+        convert_model_weight_to_float8(transformer_2, exclude_module_name=["modulation",], device=device)
+        convert_weight_dtype_wrapper(transformer_2, weight_dtype)
+    pipeline.to(device=device)
+else:
+    pipeline.to(device=device)
 
 coefficients = get_teacache_coefficients(model_name) if enable_teacache else None
 if coefficients is not None:
@@ -259,10 +289,10 @@ if cfg_skip_ratio is not None:
 generator = torch.Generator(device=device).manual_seed(seed)
 
 if lora_path is not None:
-    from videox_fun.utils import merge_lora, unmerge_lora
+    from videox_fun.utils.lora_utils import merge_lora, unmerge_lora
     pipeline = merge_lora(pipeline, lora_path, lora_weight, device=device, dtype=weight_dtype)
 if lora_high_path is not None and transformer_2 is not None:
-    from videox_fun.utils import merge_lora, unmerge_lora
+    from videox_fun.utils.lora_utils import merge_lora, unmerge_lora
     pipeline = merge_lora(pipeline, lora_high_path, lora_high_weight, device=device, dtype=weight_dtype, sub_transformer_name="transformer_2")
 
 with torch.no_grad():

examples/lingbot_world/predict_i2v_fast.py

@@ -17,10 +17,15 @@ from videox_fun.models import (AutoencoderKLWan, AutoTokenizer,
                                WanT5EncoderModel,
                                WanTransformer3DModel_LingbotWorldFast)
 from videox_fun.pipeline import WanFunLingbotWorldFastPipeline
-from videox_fun.utils import (FlowDPMSolverMultistepScheduler,
-                              FlowUniPCMultistepScheduler,
-                              apply_gpu_memory_mode, filter_kwargs, merge_lora,
-                              save_videos_grid, unmerge_lora)
+from videox_fun.utils import (register_auto_device_hook,
+                              safe_enable_group_offload)
+from videox_fun.utils.fm_solvers import FlowDPMSolverMultistepScheduler
+from videox_fun.utils.fm_solvers_unipc import FlowUniPCMultistepScheduler
+from videox_fun.utils.fp8_optimization import (convert_model_weight_to_float8,
+                                               convert_weight_dtype_wrapper,
+                                               replace_parameters_by_name)
+from videox_fun.utils.utils import filter_kwargs, save_videos_grid
+from videox_fun.utils.lora_utils import merge_lora, unmerge_lora
 
 # GPU memory mode, which can be chosen in [model_full_load, model_full_load_and_qfloat8, model_cpu_offload, model_cpu_offload_and_qfloat8, model_group_offload, sequential_cpu_offload].
 # model_full_load means that the entire model will be moved to the GPU.
@@ -108,7 +113,7 @@ video_length        = 81
 fps                 = 16
 
 # Use torch.float16 if GPU does not support torch.bfloat16
-# Some graphics cards, such as v100, 2080ti, do not support torch.bfloat16
+# some graphics cards, such as v100, 2080ti, do not support torch.bfloat16
 weight_dtype        = torch.bfloat16
 
 # Camera trajectory (poses.npy / intrinsics.npy) + reference image + prompt.
@@ -148,7 +153,7 @@ transformer = WanTransformer3DModel_LingbotWorldFast.from_pretrained(
 if transformer_path is not None:
     print(f"From checkpoint: {transformer_path}")
     if transformer_path.endswith("safetensors"):
-        from safetensors.torch import load_file
+        from safetensors.torch import load_file, safe_open
         state_dict = load_file(transformer_path)
     else:
         state_dict = torch.load(transformer_path, map_location="cpu")
@@ -166,7 +171,7 @@ vae = AutoencoderKLWan.from_pretrained(
 if vae_path is not None:
     print(f"From checkpoint: {vae_path}")
     if vae_path.endswith("safetensors"):
-        from safetensors.torch import load_file
+        from safetensors.torch import load_file, safe_open
         state_dict = load_file(vae_path)
     else:
         state_dict = torch.load(vae_path, map_location="cpu")
@@ -190,7 +195,7 @@ text_encoder = WanT5EncoderModel.from_pretrained(
 text_encoder = text_encoder.eval()
 
 # Get Scheduler
-Chosen_Scheduler = {
+Chosen_Scheduler = scheduler_dict = {
     "Flow": FlowMatchEulerDiscreteScheduler,
     "Flow_Unipc": FlowUniPCMultistepScheduler,
     "Flow_DPM++": FlowDPMSolverMultistepScheduler,
@@ -233,10 +238,25 @@ if compile_dit:
         pipeline.transformer.blocks[i] = torch.compile(pipeline.transformer.blocks[i])
     print("Add Compile")
 
-# Quantize (when the mode carries an "_and_<quant>" suffix) and then place the pipeline.
-# The order lives inside the helper: quantization has to happen before the offload hooks
-# are installed.
-apply_gpu_memory_mode(pipeline, GPU_memory_mode, device, weight_dtype)
+if GPU_memory_mode == "sequential_cpu_offload":
+    replace_parameters_by_name(transformer, ["modulation",], device=device)
+    transformer.freqs = transformer.freqs.to(device=device)
+    pipeline.enable_sequential_cpu_offload(device=device)
+elif GPU_memory_mode == "model_group_offload":
+    register_auto_device_hook(pipeline.transformer)
+    safe_enable_group_offload(pipeline, onload_device=device, offload_device="cpu", offload_type="leaf_level", use_stream=True)
+elif GPU_memory_mode == "model_cpu_offload_and_qfloat8":
+    convert_model_weight_to_float8(transformer, exclude_module_name=["modulation",], device=device)
+    convert_weight_dtype_wrapper(transformer, weight_dtype)
+    pipeline.enable_model_cpu_offload(device=device)
+elif GPU_memory_mode == "model_cpu_offload":
+    pipeline.enable_model_cpu_offload(device=device)
+elif GPU_memory_mode == "model_full_load_and_qfloat8":
+    convert_model_weight_to_float8(transformer, exclude_module_name=["modulation",], device=device)
+    convert_weight_dtype_wrapper(transformer, weight_dtype)
+    pipeline.to(device=device)
+else:
+    pipeline.to(device=device)
 
 generator = torch.Generator(device=device).manual_seed(seed)
 

examples/longcatvideo/predict_i2v.py

@@ -4,6 +4,7 @@ import sys
 import numpy as np
 import torch
 from diffusers import FlowMatchEulerDiscreteScheduler
+from omegaconf import OmegaConf
 from PIL import Image
 
 current_file_path = os.path.abspath(__file__)
@@ -12,16 +13,21 @@ for project_root in project_roots:
     sys.path.insert(0, project_root) if project_root not in sys.path else None
 
 from videox_fun.dist import set_multi_gpus_devices, shard_model
-from videox_fun.models import (AutoencoderKLLongCatVideo, AutoTokenizer,
-                               LongCatVideoTransformer3DModel,
-                               UMT5EncoderModel)
+from videox_fun.models import (AutoencoderKLLongCatVideo, UMT5EncoderModel, AutoTokenizer,
+                              LongCatVideoTransformer3DModel)
+from videox_fun.models.cache_utils import get_teacache_coefficients
 from videox_fun.pipeline import LongCatVideoPipeline
-from videox_fun.utils import (FlowDPMSolverMultistepScheduler,
-                              FlowUniPCMultistepScheduler,
-                              apply_gpu_memory_mode, get_image_to_video_latent,
-                              merge_lora, save_videos_grid, unmerge_lora)
+from videox_fun.utils import (register_auto_device_hook,
+                              safe_enable_group_offload)
+from videox_fun.utils.fp8_optimization import (convert_model_weight_to_float8, replace_parameters_by_name,
+                                              convert_weight_dtype_wrapper)
+from videox_fun.utils.lora_utils import merge_lora, unmerge_lora
+from videox_fun.utils.utils import (filter_kwargs, get_image_to_video_latent,
+                                   save_videos_grid)
+from videox_fun.utils.fm_solvers import FlowDPMSolverMultistepScheduler
+from videox_fun.utils.fm_solvers_unipc import FlowUniPCMultistepScheduler
 
-# GPU memory mode, which can be chosen in [model_full_load, model_full_load_and_qfloat8, model_cpu_offload, model_cpu_offload_and_qfloat8, model_group_offload, sequential_cpu_offload].
+# GPU memory mode, which can be chosen in [model_full_load, model_full_load_and_qfloat8, model_cpu_offload, model_cpu_offload_and_qfloat8, sequential_cpu_offload].
 # model_full_load means that the entire model will be moved to the GPU.
 # 
 # model_full_load_and_qfloat8 means that the entire model will be moved to the GPU,
@@ -68,7 +74,7 @@ video_length        = 81
 fps                 = 16
 
 # Use torch.float16 if GPU does not support torch.bfloat16
-# Some graphics cards, such as v100, 2080ti, do not support torch.bfloat16
+# ome graphics cards, such as v100, 2080ti, do not support torch.bfloat16
 weight_dtype            = torch.bfloat16
 validation_image_start  = "asset/1.png"
 
@@ -92,7 +98,7 @@ transformer = LongCatVideoTransformer3DModel.from_pretrained(
 if transformer_path is not None:
     print(f"From checkpoint: {transformer_path}")
     if transformer_path.endswith("safetensors"):
-        from safetensors.torch import load_file
+        from safetensors.torch import load_file, safe_open
         state_dict = load_file(transformer_path)
     else:
         state_dict = torch.load(transformer_path, map_location="cpu")
@@ -109,7 +115,7 @@ vae = AutoencoderKLLongCatVideo.from_pretrained(
 if vae_path is not None:
     print(f"From checkpoint: {vae_path}")
     if vae_path.endswith("safetensors"):
-        from safetensors.torch import load_file
+        from safetensors.torch import load_file, safe_open
         state_dict = load_file(vae_path)
     else:
         state_dict = torch.load(vae_path, map_location="cpu")
@@ -131,7 +137,7 @@ text_encoder = UMT5EncoderModel.from_pretrained(
 )
 
 # Get Scheduler
-Chosen_Scheduler = {
+Chosen_Scheduler = scheduler_dict = {
     "Flow": FlowMatchEulerDiscreteScheduler,
     "Flow_Unipc": FlowUniPCMultistepScheduler,
     "Flow_DPM++": FlowDPMSolverMultistepScheduler,
@@ -166,10 +172,24 @@ if compile_dit:
         pipeline.transformer.blocks[i] = torch.compile(pipeline.transformer.blocks[i])
     print("Add Compile")
 
-# Quantize (when the mode carries an "_and_<quant>" suffix) and then place the pipeline.
-# The order lives inside the helper: quantization has to happen before the offload hooks
-# are installed.
-apply_gpu_memory_mode(pipeline, GPU_memory_mode, device, weight_dtype)
+if GPU_memory_mode == "sequential_cpu_offload":
+    replace_parameters_by_name(transformer, ["modulation",], device=device)
+    pipeline.enable_sequential_cpu_offload(device=device)
+elif GPU_memory_mode == "model_group_offload":
+    register_auto_device_hook(pipeline.transformer)
+    safe_enable_group_offload(pipeline, onload_device=device, offload_device="cpu", offload_type="leaf_level", use_stream=True)
+elif GPU_memory_mode == "model_cpu_offload_and_qfloat8":
+    convert_model_weight_to_float8(transformer, exclude_module_name=["modulation",], device=device)
+    convert_weight_dtype_wrapper(transformer, weight_dtype)
+    pipeline.enable_model_cpu_offload(device=device)
+elif GPU_memory_mode == "model_cpu_offload":
+    pipeline.enable_model_cpu_offload(device=device)
+elif GPU_memory_mode == "model_full_load_and_qfloat8":
+    convert_model_weight_to_float8(transformer, exclude_module_name=["modulation",], device=device)
+    convert_weight_dtype_wrapper(transformer, weight_dtype)
+    pipeline.to(device=device)
+else:
+    pipeline.to(device=device)
 
 generator = torch.Generator(device=device).manual_seed(seed)
 

examples/longcatvideo/predict_s2v_avatar.py

@@ -6,6 +6,7 @@ import numpy as np
 import torch
 from audio_separator.separator import Separator
 from diffusers import FlowMatchEulerDiscreteScheduler
+from omegaconf import OmegaConf
 from PIL import Image
 
 current_file_path = os.path.abspath(__file__)
@@ -18,14 +19,20 @@ from videox_fun.models import (AutoencoderKLLongCatVideo, AutoTokenizer,
                                LongCatVideoAudioEncoder,
                                LongCatVideoAvatarTransformer3DModel,
                                UMT5EncoderModel)
+from videox_fun.models.cache_utils import get_teacache_coefficients
 from videox_fun.pipeline import LongCatVideoAvatarPipeline
-from videox_fun.utils import (FlowDPMSolverMultistepScheduler,
-                              FlowUniPCMultistepScheduler,
-                              apply_gpu_memory_mode, get_image_to_video_latent,
-                              merge_lora, merge_video_audio, save_videos_grid,
-                              unmerge_lora)
+from videox_fun.utils.fm_solvers import FlowDPMSolverMultistepScheduler
+from videox_fun.utils.fm_solvers_unipc import FlowUniPCMultistepScheduler
+from videox_fun.utils import (register_auto_device_hook,
+                              safe_enable_group_offload)
+from videox_fun.utils.fp8_optimization import (convert_model_weight_to_float8,
+                                               convert_weight_dtype_wrapper,
+                                               replace_parameters_by_name)
+from videox_fun.utils.lora_utils import merge_lora, unmerge_lora
+from videox_fun.utils.utils import (filter_kwargs, get_image_to_video_latent,
+                                    merge_video_audio, save_videos_grid)
 
-# GPU memory mode, which can be chosen in [model_full_load, model_full_load_and_qfloat8, model_cpu_offload, model_cpu_offload_and_qfloat8, model_group_offload, sequential_cpu_offload].
+# GPU memory mode, which can be chosen in [model_full_load, model_full_load_and_qfloat8, model_cpu_offload, model_cpu_offload_and_qfloat8, sequential_cpu_offload].
 # model_full_load means that the entire model will be moved to the GPU.
 # 
 # model_full_load_and_qfloat8 means that the entire model will be moved to the GPU,
@@ -80,7 +87,7 @@ audio_path = "asset/talk.wav"
 use_audio_vocal_separator = False
 
 # Use torch.float16 if GPU does not support torch.bfloat16
-# Some graphics cards, such as v100, 2080ti, do not support torch.bfloat16
+# ome graphics cards, such as v100, 2080ti, do not support torch.bfloat16
 weight_dtype        = torch.bfloat16
 # Prompt
 prompt              = "A young woman with long flowing purple hair stands by the seaside on a sunny day, singing. Wearing a white sleeveless dress with a navy blue bow at the collar, her hair gently sways in the ocean breeze. The sparkling sea, blue sky with white clouds, and pink wildflowers along the shore create a beautiful and vibrant scene."
@@ -102,7 +109,7 @@ transformer = LongCatVideoAvatarTransformer3DModel.from_pretrained(
 if transformer_path is not None:
     print(f"From checkpoint: {transformer_path}")
     if transformer_path.endswith("safetensors"):
-        from safetensors.torch import load_file
+        from safetensors.torch import load_file, safe_open
         state_dict = load_file(transformer_path)
     else:
         state_dict = torch.load(transformer_path, map_location="cpu")
@@ -119,7 +126,7 @@ vae = AutoencoderKLLongCatVideo.from_pretrained(
 if vae_path is not None:
     print(f"From checkpoint: {vae_path}")
     if vae_path.endswith("safetensors"):
-        from safetensors.torch import load_file
+        from safetensors.torch import load_file, safe_open
         state_dict = load_file(vae_path)
     else:
         state_dict = torch.load(vae_path, map_location="cpu")
@@ -147,7 +154,7 @@ audio_encoder = LongCatVideoAudioEncoder(
 audio_encoder.audio_encoder.feature_extractor._freeze_parameters()
 
 # Get Scheduler
-Chosen_Scheduler = {
+Chosen_Scheduler = scheduler_dict = {
     "Flow": FlowMatchEulerDiscreteScheduler,
     "Flow_Unipc": FlowUniPCMultistepScheduler,
     "Flow_DPM++": FlowDPMSolverMultistepScheduler,
@@ -183,10 +190,24 @@ if compile_dit:
         pipeline.transformer.blocks[i] = torch.compile(pipeline.transformer.blocks[i])
     print("Add Compile")
 
-# Quantize (when the mode carries an "_and_<quant>" suffix) and then place the pipeline.
-# The order lives inside the helper: quantization has to happen before the offload hooks
-# are installed.
-apply_gpu_memory_mode(pipeline, GPU_memory_mode, device, weight_dtype)
+if GPU_memory_mode == "sequential_cpu_offload":
+    replace_parameters_by_name(transformer, ["modulation",], device=device)
+    pipeline.enable_sequential_cpu_offload(device=device)
+elif GPU_memory_mode == "model_group_offload":
+    register_auto_device_hook(pipeline.transformer)
+    safe_enable_group_offload(pipeline, onload_device=device, offload_device="cpu", offload_type="leaf_level", use_stream=True)
+elif GPU_memory_mode == "model_cpu_offload_and_qfloat8":
+    convert_model_weight_to_float8(transformer, exclude_module_name=["modulation",], device=device)
+    convert_weight_dtype_wrapper(transformer, weight_dtype)
+    pipeline.enable_model_cpu_offload(device=device)
+elif GPU_memory_mode == "model_cpu_offload":
+    pipeline.enable_model_cpu_offload(device=device)
+elif GPU_memory_mode == "model_full_load_and_qfloat8":
+    convert_model_weight_to_float8(transformer, exclude_module_name=["modulation",], device=device)
+    convert_weight_dtype_wrapper(transformer, weight_dtype)
+    pipeline.to(device=device)
+else:
+    pipeline.to(device=device)
 
 generator = torch.Generator(device=device).manual_seed(seed)
 

examples/longcatvideo/predict_t2v.py

@@ -4,6 +4,7 @@ import sys
 import numpy as np
 import torch
 from diffusers import FlowMatchEulerDiscreteScheduler
+from omegaconf import OmegaConf
 from PIL import Image
 
 current_file_path = os.path.abspath(__file__)
@@ -12,16 +13,21 @@ for project_root in project_roots:
     sys.path.insert(0, project_root) if project_root not in sys.path else None
 
 from videox_fun.dist import set_multi_gpus_devices, shard_model
-from videox_fun.models import (AutoencoderKLLongCatVideo, AutoTokenizer,
-                               LongCatVideoTransformer3DModel,
-                               UMT5EncoderModel)
+from videox_fun.models import (AutoencoderKLLongCatVideo, UMT5EncoderModel, AutoTokenizer,
+                              LongCatVideoTransformer3DModel)
+from videox_fun.models.cache_utils import get_teacache_coefficients
 from videox_fun.pipeline import LongCatVideoPipeline
-from videox_fun.utils import (FlowDPMSolverMultistepScheduler,
-                              FlowUniPCMultistepScheduler,
-                              apply_gpu_memory_mode, merge_lora,
-                              save_videos_grid, unmerge_lora)
+from videox_fun.utils import (register_auto_device_hook,
+                              safe_enable_group_offload)
+from videox_fun.utils.fp8_optimization import (convert_model_weight_to_float8, replace_parameters_by_name,
+                                              convert_weight_dtype_wrapper)
+from videox_fun.utils.lora_utils import merge_lora, unmerge_lora
+from videox_fun.utils.utils import (filter_kwargs, get_image_to_video_latent,
+                                   save_videos_grid)
+from videox_fun.utils.fm_solvers import FlowDPMSolverMultistepScheduler
+from videox_fun.utils.fm_solvers_unipc import FlowUniPCMultistepScheduler
 
-# GPU memory mode, which can be chosen in [model_full_load, model_full_load_and_qfloat8, model_cpu_offload, model_cpu_offload_and_qfloat8, model_group_offload, sequential_cpu_offload].
+# GPU memory mode, which can be chosen in [model_full_load, model_full_load_and_qfloat8, model_cpu_offload, model_cpu_offload_and_qfloat8, sequential_cpu_offload].
 # model_full_load means that the entire model will be moved to the GPU.
 # 
 # model_full_load_and_qfloat8 means that the entire model will be moved to the GPU,
@@ -68,7 +74,7 @@ video_length        = 81
 fps                 = 16
 
 # Use torch.float16 if GPU does not support torch.bfloat16
-# Some graphics cards, such as v100, 2080ti, do not support torch.bfloat16
+# ome graphics cards, such as v100, 2080ti, do not support torch.bfloat16
 weight_dtype        = torch.bfloat16
 # Prompt
 prompt              = "1girl, black_hair, brown_eyes, earrings, freckles, grey_background, jewelry, lips, long_hair, looking_at_viewer, nose, piercing, realistic, red_lips, solo, upper_body"
@@ -90,7 +96,7 @@ transformer = LongCatVideoTransformer3DModel.from_pretrained(
 if transformer_path is not None:
     print(f"From checkpoint: {transformer_path}")
     if transformer_path.endswith("safetensors"):
-        from safetensors.torch import load_file
+        from safetensors.torch import load_file, safe_open
         state_dict = load_file(transformer_path)
     else:
         state_dict = torch.load(transformer_path, map_location="cpu")
@@ -107,7 +113,7 @@ vae = AutoencoderKLLongCatVideo.from_pretrained(
 if vae_path is not None:
     print(f"From checkpoint: {vae_path}")
     if vae_path.endswith("safetensors"):
-        from safetensors.torch import load_file
+        from safetensors.torch import load_file, safe_open
         state_dict = load_file(vae_path)
     else:
         state_dict = torch.load(vae_path, map_location="cpu")
@@ -129,7 +135,7 @@ text_encoder = UMT5EncoderModel.from_pretrained(
 )
 
 # Get Scheduler
-Chosen_Scheduler = {
+Chosen_Scheduler = scheduler_dict = {
     "Flow": FlowMatchEulerDiscreteScheduler,
     "Flow_Unipc": FlowUniPCMultistepScheduler,
     "Flow_DPM++": FlowDPMSolverMultistepScheduler,
@@ -165,10 +171,24 @@ if compile_dit:
         pipeline.transformer.blocks[i] = torch.compile(pipeline.transformer.blocks[i])
     print("Add Compile")
 
-# Quantize (when the mode carries an "_and_<quant>" suffix) and then place the pipeline.
-# The order lives inside the helper: quantization has to happen before the offload hooks
-# are installed.
-apply_gpu_memory_mode(pipeline, GPU_memory_mode, device, weight_dtype)
+if GPU_memory_mode == "sequential_cpu_offload":
+    replace_parameters_by_name(transformer, ["modulation",], device=device)
+    pipeline.enable_sequential_cpu_offload(device=device)
+elif GPU_memory_mode == "model_group_offload":
+    register_auto_device_hook(pipeline.transformer)
+    safe_enable_group_offload(pipeline, onload_device=device, offload_device="cpu", offload_type="leaf_level", use_stream=True)
+elif GPU_memory_mode == "model_cpu_offload_and_qfloat8":
+    convert_model_weight_to_float8(transformer, exclude_module_name=["modulation",], device=device)
+    convert_weight_dtype_wrapper(transformer, weight_dtype)
+    pipeline.enable_model_cpu_offload(device=device)
+elif GPU_memory_mode == "model_cpu_offload":
+    pipeline.enable_model_cpu_offload(device=device)
+elif GPU_memory_mode == "model_full_load_and_qfloat8":
+    convert_model_weight_to_float8(transformer, exclude_module_name=["modulation",], device=device)
+    convert_weight_dtype_wrapper(transformer, weight_dtype)
+    pipeline.to(device=device)
+else:
+    pipeline.to(device=device)
 
 generator = torch.Generator(device=device).manual_seed(seed)
 

examples/ltx2.3/predict_i2v.py

@@ -11,18 +11,25 @@ project_roots = [os.path.dirname(current_file_path), os.path.dirname(os.path.dir
 for project_root in project_roots:
     sys.path.insert(0, project_root) if project_root not in sys.path else None
 
-from videox_fun.dist import set_multi_gpus_devices, shard_model
 from videox_fun.models import (AutoencoderKLLTX2Audio, AutoencoderKLLTX2Video,
-                               Gemma3ForConditionalGeneration, Gemma3Processor,
-                               LTX2TextConnectors, LTX2VideoTransformer3DModel,
-                               LTX2VocoderWithBWE)
+                               Gemma3ForConditionalGeneration,
+                               GemmaTokenizerFast, LTX2TextConnectors, Gemma3Processor,
+                               LTX2VideoTransformer3DModel, LTX2VocoderWithBWE)
 from videox_fun.pipeline import LTX2I2VPipeline
-from videox_fun.utils import (FlowDPMSolverMultistepScheduler,
-                              FlowUniPCMultistepScheduler,
-                              apply_gpu_memory_mode, merge_lora,
-                              save_videos_with_audio_grid, unmerge_lora)
+from videox_fun.utils import (register_auto_device_hook,
+                              safe_enable_group_offload)
+from videox_fun.dist import set_multi_gpus_devices, shard_model
+from videox_fun.utils.fm_solvers import FlowDPMSolverMultistepScheduler
+from videox_fun.utils.fm_solvers_unipc import FlowUniPCMultistepScheduler
+from videox_fun.utils.fp8_optimization import (convert_model_weight_to_float8,
+                                               convert_weight_dtype_wrapper,
+                                               replace_parameters_by_name)
+from videox_fun.utils.lora_utils import merge_lora, unmerge_lora
+from videox_fun.utils.utils import (filter_kwargs, get_image_to_video_latent,
+                                    save_videos_grid,
+                                    save_videos_with_audio_grid)
 
-# GPU memory mode, which can be chosen in [model_full_load, model_full_load_and_qfloat8, model_cpu_offload, model_cpu_offload_and_qfloat8, model_group_offload, sequential_cpu_offload].
+# GPU memory mode, which can be chosen in [model_full_load, model_full_load_and_qfloat8, model_cpu_offload, model_cpu_offload_and_qfloat8, sequential_cpu_offload].
 # model_full_load means that the entire model will be moved to the GPU.
 # 
 # model_full_load_and_qfloat8 means that the entire model will be moved to the GPU,
@@ -68,7 +75,7 @@ video_length        = 121
 fps                 = 24
 
 # Use torch.float16 if GPU does not support torch.bfloat16
-# Some graphics cards, such as v100, 2080ti, do not support torch.bfloat16
+# ome graphics cards, such as v100, 2080ti, do not support torch.bfloat16
 weight_dtype        = torch.bfloat16
 # If you want to generate from text, please set the validation_image_start = None and validation_image_end = None
 validation_image_start  = "asset/1.png"
@@ -110,7 +117,7 @@ transformer = LTX2VideoTransformer3DModel.from_pretrained(
 if transformer_path is not None:
     print(f"From checkpoint: {transformer_path}")
     if transformer_path.endswith("safetensors"):
-        from safetensors.torch import load_file
+        from safetensors.torch import load_file, safe_open
         state_dict = load_file(transformer_path)
     else:
         state_dict = torch.load(transformer_path, map_location="cpu")
@@ -129,7 +136,7 @@ vae = AutoencoderKLLTX2Video.from_pretrained(
 if vae_path is not None:
     print(f"From checkpoint: {vae_path}")
     if vae_path.endswith("safetensors"):
-        from safetensors.torch import load_file
+        from safetensors.torch import load_file, safe_open
         state_dict = load_file(vae_path)
     else:
         state_dict = torch.load(vae_path, map_location="cpu")
@@ -219,10 +226,23 @@ if compile_dit:
         pipeline.transformer.transformer_blocks[i] = torch.compile(pipeline.transformer.transformer_blocks[i])
     print("Add Compile")
 
-# Quantize (when the mode carries an "_and_<quant>" suffix) and then place the pipeline.
-# The order lives inside the helper: quantization has to happen before the offload hooks
-# are installed.
-apply_gpu_memory_mode(pipeline, GPU_memory_mode, device, weight_dtype, exclude_module_name=['scale_shift_table', 'audio_scale_shift_table', 'video_a2v_cross_attn_scale_shift_table', 'audio_a2v_cross_attn_scale_shift_table', ''])
+if GPU_memory_mode == "sequential_cpu_offload":
+    pipeline.enable_sequential_cpu_offload(device=device)
+elif GPU_memory_mode == "model_group_offload":
+    register_auto_device_hook(pipeline.transformer)
+    safe_enable_group_offload(pipeline, onload_device=device, offload_device="cpu", offload_type="leaf_level", use_stream=True)
+elif GPU_memory_mode == "model_cpu_offload_and_qfloat8":
+    convert_model_weight_to_float8(transformer, exclude_module_name=["scale_shift_table", "audio_scale_shift_table", "video_a2v_cross_attn_scale_shift_table", "audio_a2v_cross_attn_scale_shift_table", ""], device=device)
+    convert_weight_dtype_wrapper(transformer, weight_dtype)
+    pipeline.enable_model_cpu_offload(device=device)
+elif GPU_memory_mode == "model_cpu_offload":
+    pipeline.enable_model_cpu_offload(device=device)
+elif GPU_memory_mode == "model_full_load_and_qfloat8":
+    convert_model_weight_to_float8(transformer, exclude_module_name=["scale_shift_table", "audio_scale_shift_table", "video_a2v_cross_attn_scale_shift_table", "audio_a2v_cross_attn_scale_shift_table", ""], device=device)
+    convert_weight_dtype_wrapper(transformer, weight_dtype)
+    pipeline.to(device=device)
+else:
+    pipeline.to(device=device)
 
 generator = torch.Generator(device=device).manual_seed(seed)
 

examples/ltx2.3/predict_t2v.py

@@ -11,18 +11,25 @@ project_roots = [os.path.dirname(current_file_path), os.path.dirname(os.path.dir
 for project_root in project_roots:
     sys.path.insert(0, project_root) if project_root not in sys.path else None
 
-from videox_fun.dist import set_multi_gpus_devices, shard_model
 from videox_fun.models import (AutoencoderKLLTX2Audio, AutoencoderKLLTX2Video,
-                               Gemma3ForConditionalGeneration, Gemma3Processor,
-                               LTX2TextConnectors, LTX2VideoTransformer3DModel,
-                               LTX2VocoderWithBWE)
+                               Gemma3ForConditionalGeneration,
+                               GemmaTokenizerFast, LTX2TextConnectors, Gemma3Processor,
+                               LTX2VideoTransformer3DModel, LTX2VocoderWithBWE)
 from videox_fun.pipeline import LTX2Pipeline
-from videox_fun.utils import (FlowDPMSolverMultistepScheduler,
-                              FlowUniPCMultistepScheduler,
-                              apply_gpu_memory_mode, merge_lora,
-                              save_videos_with_audio_grid, unmerge_lora)
+from videox_fun.utils import (register_auto_device_hook,
+                              safe_enable_group_offload)
+from videox_fun.dist import set_multi_gpus_devices, shard_model
+from videox_fun.utils.fm_solvers import FlowDPMSolverMultistepScheduler
+from videox_fun.utils.fm_solvers_unipc import FlowUniPCMultistepScheduler
+from videox_fun.utils.fp8_optimization import (convert_model_weight_to_float8,
+                                               convert_weight_dtype_wrapper,
+                                               replace_parameters_by_name)
+from videox_fun.utils.lora_utils import merge_lora, unmerge_lora
+from videox_fun.utils.utils import (filter_kwargs, get_image_to_video_latent,
+                                    save_videos_grid,
+                                    save_videos_with_audio_grid)
 
-# GPU memory mode, which can be chosen in [model_full_load, model_full_load_and_qfloat8, model_cpu_offload, model_cpu_offload_and_qfloat8, model_group_offload, sequential_cpu_offload].
+# GPU memory mode, which can be chosen in [model_full_load, model_full_load_and_qfloat8, model_cpu_offload, model_cpu_offload_and_qfloat8, sequential_cpu_offload].
 # model_full_load means that the entire model will be moved to the GPU.
 # 
 # model_full_load_and_qfloat8 means that the entire model will be moved to the GPU,
@@ -68,7 +75,7 @@ video_length        = 121
 fps                 = 24
 
 # Use torch.float16 if GPU does not support torch.bfloat16
-# Some graphics cards, such as v100, 2080ti, do not support torch.bfloat16
+# ome graphics cards, such as v100, 2080ti, do not support torch.bfloat16
 weight_dtype        = torch.bfloat16
 prompt              = "A brown dog barks on a sofa, sitting on a light-colored couch in a cozy room. Behind the dog, there is a framed painting on a shelf, surrounded by pink flowers. "
 negative_prompt     = "worst quality, inconsistent motion, blurry, jittery, distorted, static, low quality, artifacts"
@@ -106,7 +113,7 @@ transformer = LTX2VideoTransformer3DModel.from_pretrained(
 if transformer_path is not None:
     print(f"From checkpoint: {transformer_path}")
     if transformer_path.endswith("safetensors"):
-        from safetensors.torch import load_file
+        from safetensors.torch import load_file, safe_open
         state_dict = load_file(transformer_path)
     else:
         state_dict = torch.load(transformer_path, map_location="cpu")
@@ -125,7 +132,7 @@ vae = AutoencoderKLLTX2Video.from_pretrained(
 if vae_path is not None:
     print(f"From checkpoint: {vae_path}")
     if vae_path.endswith("safetensors"):
-        from safetensors.torch import load_file
+        from safetensors.torch import load_file, safe_open
         state_dict = load_file(vae_path)
     else:
         state_dict = torch.load(vae_path, map_location="cpu")
@@ -215,10 +222,23 @@ if compile_dit:
         pipeline.transformer.transformer_blocks[i] = torch.compile(pipeline.transformer.transformer_blocks[i])
     print("Add Compile")
 
-# Quantize (when the mode carries an "_and_<quant>" suffix) and then place the pipeline.
-# The order lives inside the helper: quantization has to happen before the offload hooks
-# are installed.
-apply_gpu_memory_mode(pipeline, GPU_memory_mode, device, weight_dtype, exclude_module_name=['scale_shift_table', 'audio_scale_shift_table', 'video_a2v_cross_attn_scale_shift_table', 'audio_a2v_cross_attn_scale_shift_table', ''])
+if GPU_memory_mode == "sequential_cpu_offload":
+    pipeline.enable_sequential_cpu_offload(device=device)
+elif GPU_memory_mode == "model_group_offload":
+    register_auto_device_hook(pipeline.transformer)
+    safe_enable_group_offload(pipeline, onload_device=device, offload_device="cpu", offload_type="leaf_level", use_stream=True)
+elif GPU_memory_mode == "model_cpu_offload_and_qfloat8":
+    convert_model_weight_to_float8(transformer, exclude_module_name=["scale_shift_table", "audio_scale_shift_table", "video_a2v_cross_attn_scale_shift_table", "audio_a2v_cross_attn_scale_shift_table", ""], device=device)
+    convert_weight_dtype_wrapper(transformer, weight_dtype)
+    pipeline.enable_model_cpu_offload(device=device)
+elif GPU_memory_mode == "model_cpu_offload":
+    pipeline.enable_model_cpu_offload(device=device)
+elif GPU_memory_mode == "model_full_load_and_qfloat8":
+    convert_model_weight_to_float8(transformer, exclude_module_name=["scale_shift_table", "audio_scale_shift_table", "video_a2v_cross_attn_scale_shift_table", "audio_a2v_cross_attn_scale_shift_table", ""], device=device)
+    convert_weight_dtype_wrapper(transformer, weight_dtype)
+    pipeline.to(device=device)
+else:
+    pipeline.to(device=device)
 
 generator = torch.Generator(device=device).manual_seed(seed)
 

examples/ltx2/predict_i2v.py

@@ -11,18 +11,25 @@ project_roots = [os.path.dirname(current_file_path), os.path.dirname(os.path.dir
 for project_root in project_roots:
     sys.path.insert(0, project_root) if project_root not in sys.path else None
 
-from videox_fun.dist import set_multi_gpus_devices, shard_model
 from videox_fun.models import (AutoencoderKLLTX2Audio, AutoencoderKLLTX2Video,
                                Gemma3ForConditionalGeneration,
                                GemmaTokenizerFast, LTX2TextConnectors,
                                LTX2VideoTransformer3DModel, LTX2Vocoder)
 from videox_fun.pipeline import LTX2I2VPipeline
-from videox_fun.utils import (FlowDPMSolverMultistepScheduler,
-                              FlowUniPCMultistepScheduler,
-                              apply_gpu_memory_mode, merge_lora,
-                              save_videos_with_audio_grid, unmerge_lora)
+from videox_fun.utils import (register_auto_device_hook,
+                              safe_enable_group_offload)
+from videox_fun.dist import set_multi_gpus_devices, shard_model
+from videox_fun.utils.fm_solvers import FlowDPMSolverMultistepScheduler
+from videox_fun.utils.fm_solvers_unipc import FlowUniPCMultistepScheduler
+from videox_fun.utils.fp8_optimization import (convert_model_weight_to_float8,
+                                               convert_weight_dtype_wrapper,
+                                               replace_parameters_by_name)
+from videox_fun.utils.lora_utils import merge_lora, unmerge_lora
+from videox_fun.utils.utils import (filter_kwargs, get_image_to_video_latent,
+                                    save_videos_grid,
+                                    save_videos_with_audio_grid)
 
-# GPU memory mode, which can be chosen in [model_full_load, model_full_load_and_qfloat8, model_cpu_offload, model_cpu_offload_and_qfloat8, model_group_offload, sequential_cpu_offload].
+# GPU memory mode, which can be chosen in [model_full_load, model_full_load_and_qfloat8, model_cpu_offload, model_cpu_offload_and_qfloat8, sequential_cpu_offload].
 # model_full_load means that the entire model will be moved to the GPU.
 # 
 # model_full_load_and_qfloat8 means that the entire model will be moved to the GPU,
@@ -68,7 +75,7 @@ video_length        = 121
 fps                 = 24
 
 # Use torch.float16 if GPU does not support torch.bfloat16
-# Some graphics cards, such as v100, 2080ti, do not support torch.bfloat16
+# ome graphics cards, such as v100, 2080ti, do not support torch.bfloat16
 weight_dtype        = torch.bfloat16
 # If you want to generate from text, please set the validation_image_start = None and validation_image_end = None
 validation_image_start  = "asset/1.png"
@@ -98,7 +105,7 @@ transformer = LTX2VideoTransformer3DModel.from_pretrained(
 if transformer_path is not None:
     print(f"From checkpoint: {transformer_path}")
     if transformer_path.endswith("safetensors"):
-        from safetensors.torch import load_file
+        from safetensors.torch import load_file, safe_open
         state_dict = load_file(transformer_path)
     else:
         state_dict = torch.load(transformer_path, map_location="cpu")
@@ -117,7 +124,7 @@ vae = AutoencoderKLLTX2Video.from_pretrained(
 if vae_path is not None:
     print(f"From checkpoint: {vae_path}")
     if vae_path.endswith("safetensors"):
-        from safetensors.torch import load_file
+        from safetensors.torch import load_file, safe_open
         state_dict = load_file(vae_path)
     else:
         state_dict = torch.load(vae_path, map_location="cpu")
@@ -203,10 +210,23 @@ if compile_dit:
         pipeline.transformer.transformer_blocks[i] = torch.compile(pipeline.transformer.transformer_blocks[i])
     print("Add Compile")
 
-# Quantize (when the mode carries an "_and_<quant>" suffix) and then place the pipeline.
-# The order lives inside the helper: quantization has to happen before the offload hooks
-# are installed.
-apply_gpu_memory_mode(pipeline, GPU_memory_mode, device, weight_dtype, exclude_module_name=['scale_shift_table', 'audio_scale_shift_table', 'video_a2v_cross_attn_scale_shift_table', 'audio_a2v_cross_attn_scale_shift_table', ''])
+if GPU_memory_mode == "sequential_cpu_offload":
+    pipeline.enable_sequential_cpu_offload(device=device)
+elif GPU_memory_mode == "model_group_offload":
+    register_auto_device_hook(pipeline.transformer)
+    safe_enable_group_offload(pipeline, onload_device=device, offload_device="cpu", offload_type="leaf_level", use_stream=True)
+elif GPU_memory_mode == "model_cpu_offload_and_qfloat8":
+    convert_model_weight_to_float8(transformer, exclude_module_name=["scale_shift_table", "audio_scale_shift_table", "video_a2v_cross_attn_scale_shift_table", "audio_a2v_cross_attn_scale_shift_table", ""], device=device)
+    convert_weight_dtype_wrapper(transformer, weight_dtype)
+    pipeline.enable_model_cpu_offload(device=device)
+elif GPU_memory_mode == "model_cpu_offload":
+    pipeline.enable_model_cpu_offload(device=device)
+elif GPU_memory_mode == "model_full_load_and_qfloat8":
+    convert_model_weight_to_float8(transformer, exclude_module_name=["scale_shift_table", "audio_scale_shift_table", "video_a2v_cross_attn_scale_shift_table", "audio_a2v_cross_attn_scale_shift_table", ""], device=device)
+    convert_weight_dtype_wrapper(transformer, weight_dtype)
+    pipeline.to(device=device)
+else:
+    pipeline.to(device=device)
 
 generator = torch.Generator(device=device).manual_seed(seed)
 

examples/ltx2/predict_i2v_upsample.py

@@ -11,19 +11,26 @@ project_roots = [os.path.dirname(current_file_path), os.path.dirname(os.path.dir
 for project_root in project_roots:
     sys.path.insert(0, project_root) if project_root not in sys.path else None
 
-from videox_fun.dist import set_multi_gpus_devices, shard_model
 from videox_fun.models import (AutoencoderKLLTX2Audio, AutoencoderKLLTX2Video,
                                Gemma3ForConditionalGeneration,
                                GemmaTokenizerFast, LTX2LatentUpsamplerModel,
-                               LTX2TextConnectors, LTX2VideoTransformer3DModel,
-                               LTX2Vocoder)
+                               LTX2TextConnectors,
+                               LTX2VideoTransformer3DModel, LTX2Vocoder)
 from videox_fun.pipeline import LTX2I2VPipeline, LTX2LatentUpsamplePipeline
-from videox_fun.utils import (FlowDPMSolverMultistepScheduler,
-                              FlowUniPCMultistepScheduler,
-                              apply_gpu_memory_mode, merge_lora,
-                              save_videos_with_audio_grid, unmerge_lora)
+from videox_fun.utils import (register_auto_device_hook,
+                              safe_enable_group_offload)
+from videox_fun.dist import set_multi_gpus_devices, shard_model
+from videox_fun.utils.fm_solvers import FlowDPMSolverMultistepScheduler
+from videox_fun.utils.fm_solvers_unipc import FlowUniPCMultistepScheduler
+from videox_fun.utils.fp8_optimization import (convert_model_weight_to_float8,
+                                               convert_weight_dtype_wrapper,
+                                               replace_parameters_by_name)
+from videox_fun.utils.lora_utils import merge_lora, unmerge_lora
+from videox_fun.utils.utils import (filter_kwargs, get_image_to_video_latent,
+                                    save_videos_grid,
+                                    save_videos_with_audio_grid)
 
-# GPU memory mode, which can be chosen in [model_full_load, model_full_load_and_qfloat8, model_cpu_offload, model_cpu_offload_and_qfloat8, model_group_offload, sequential_cpu_offload].
+# GPU memory mode, which can be chosen in [model_full_load, model_full_load_and_qfloat8, model_cpu_offload, model_cpu_offload_and_qfloat8, sequential_cpu_offload].
 # model_full_load means that the entire model will be moved to the GPU.
 # 
 # model_full_load_and_qfloat8 means that the entire model will be moved to the GPU,
@@ -72,7 +79,7 @@ fps                 = 24
 enable_latent_upsample = True
 
 # Use torch.float16 if GPU does not support torch.bfloat16
-# Some graphics cards, such as v100, 2080ti, do not support torch.bfloat16
+# ome graphics cards, such as v100, 2080ti, do not support torch.bfloat16
 weight_dtype        = torch.bfloat16
 # If you want to generate from text, please set the validation_image_start = None and validation_image_end = None
 validation_image_start  = "asset/1.png"
@@ -102,7 +109,7 @@ transformer = LTX2VideoTransformer3DModel.from_pretrained(
 if transformer_path is not None:
     print(f"From checkpoint: {transformer_path}")
     if transformer_path.endswith("safetensors"):
-        from safetensors.torch import load_file
+        from safetensors.torch import load_file, safe_open
         state_dict = load_file(transformer_path)
     else:
         state_dict = torch.load(transformer_path, map_location="cpu")
@@ -121,7 +128,7 @@ vae = AutoencoderKLLTX2Video.from_pretrained(
 if vae_path is not None:
     print(f"From checkpoint: {vae_path}")
     if vae_path.endswith("safetensors"):
-        from safetensors.torch import load_file
+        from safetensors.torch import load_file, safe_open
         state_dict = load_file(vae_path)
     else:
         state_dict = torch.load(vae_path, map_location="cpu")
@@ -207,10 +214,23 @@ if compile_dit:
         pipeline.transformer.transformer_blocks[i] = torch.compile(pipeline.transformer.transformer_blocks[i])
     print("Add Compile")
 
-# Quantize (when the mode carries an "_and_<quant>" suffix) and then place the pipeline.
-# The order lives inside the helper: quantization has to happen before the offload hooks
-# are installed.
-apply_gpu_memory_mode(pipeline, GPU_memory_mode, device, weight_dtype, exclude_module_name=['scale_shift_table', 'audio_scale_shift_table', 'video_a2v_cross_attn_scale_shift_table', 'audio_a2v_cross_attn_scale_shift_table', ''])
+if GPU_memory_mode == "sequential_cpu_offload":
+    pipeline.enable_sequential_cpu_offload(device=device)
+elif GPU_memory_mode == "model_group_offload":
+    register_auto_device_hook(pipeline.transformer)
+    safe_enable_group_offload(pipeline, onload_device=device, offload_device="cpu", offload_type="leaf_level", use_stream=True)
+elif GPU_memory_mode == "model_cpu_offload_and_qfloat8":
+    convert_model_weight_to_float8(transformer, exclude_module_name=["scale_shift_table", "audio_scale_shift_table", "video_a2v_cross_attn_scale_shift_table", "audio_a2v_cross_attn_scale_shift_table", ""], device=device)
+    convert_weight_dtype_wrapper(transformer, weight_dtype)
+    pipeline.enable_model_cpu_offload(device=device)
+elif GPU_memory_mode == "model_cpu_offload":
+    pipeline.enable_model_cpu_offload(device=device)
+elif GPU_memory_mode == "model_full_load_and_qfloat8":
+    convert_model_weight_to_float8(transformer, exclude_module_name=["scale_shift_table", "audio_scale_shift_table", "video_a2v_cross_attn_scale_shift_table", "audio_a2v_cross_attn_scale_shift_table", ""], device=device)
+    convert_weight_dtype_wrapper(transformer, weight_dtype)
+    pipeline.to(device=device)
+else:
+    pipeline.to(device=device)
 
 generator = torch.Generator(device=device).manual_seed(seed)
 

examples/ltx2/predict_t2v.py

@@ -11,18 +11,25 @@ project_roots = [os.path.dirname(current_file_path), os.path.dirname(os.path.dir
 for project_root in project_roots:
     sys.path.insert(0, project_root) if project_root not in sys.path else None
 
-from videox_fun.dist import set_multi_gpus_devices, shard_model
 from videox_fun.models import (AutoencoderKLLTX2Audio, AutoencoderKLLTX2Video,
                                Gemma3ForConditionalGeneration,
                                GemmaTokenizerFast, LTX2TextConnectors,
                                LTX2VideoTransformer3DModel, LTX2Vocoder)
 from videox_fun.pipeline import LTX2Pipeline
-from videox_fun.utils import (FlowDPMSolverMultistepScheduler,
-                              FlowUniPCMultistepScheduler,
-                              apply_gpu_memory_mode, merge_lora,
-                              save_videos_with_audio_grid, unmerge_lora)
+from videox_fun.utils import (register_auto_device_hook,
+                              safe_enable_group_offload)
+from videox_fun.dist import set_multi_gpus_devices, shard_model
+from videox_fun.utils.fm_solvers import FlowDPMSolverMultistepScheduler
+from videox_fun.utils.fm_solvers_unipc import FlowUniPCMultistepScheduler
+from videox_fun.utils.fp8_optimization import (convert_model_weight_to_float8,
+                                               convert_weight_dtype_wrapper,
+                                               replace_parameters_by_name)
+from videox_fun.utils.lora_utils import merge_lora, unmerge_lora
+from videox_fun.utils.utils import (filter_kwargs, get_image_to_video_latent,
+                                    save_videos_grid,
+                                    save_videos_with_audio_grid)
 
-# GPU memory mode, which can be chosen in [model_full_load, model_full_load_and_qfloat8, model_cpu_offload, model_cpu_offload_and_qfloat8, model_group_offload, sequential_cpu_offload].
+# GPU memory mode, which can be chosen in [model_full_load, model_full_load_and_qfloat8, model_cpu_offload, model_cpu_offload_and_qfloat8, sequential_cpu_offload].
 # model_full_load means that the entire model will be moved to the GPU.
 # 
 # model_full_load_and_qfloat8 means that the entire model will be moved to the GPU,
@@ -68,7 +75,7 @@ video_length        = 121
 fps                 = 24
 
 # Use torch.float16 if GPU does not support torch.bfloat16
-# Some graphics cards, such as v100, 2080ti, do not support torch.bfloat16
+# ome graphics cards, such as v100, 2080ti, do not support torch.bfloat16
 weight_dtype        = torch.bfloat16
 prompt              = "A brown dog barks on a sofa, sitting on a light-colored couch in a cozy room. Behind the dog, there is a framed painting on a shelf, surrounded by pink flowers. "
 negative_prompt     = "worst quality, inconsistent motion, blurry, jittery, distorted, static, low quality, artifacts"
@@ -94,7 +101,7 @@ transformer = LTX2VideoTransformer3DModel.from_pretrained(
 if transformer_path is not None:
     print(f"From checkpoint: {transformer_path}")
     if transformer_path.endswith("safetensors"):
-        from safetensors.torch import load_file
+        from safetensors.torch import load_file, safe_open
         state_dict = load_file(transformer_path)
     else:
         state_dict = torch.load(transformer_path, map_location="cpu")
@@ -113,7 +120,7 @@ vae = AutoencoderKLLTX2Video.from_pretrained(
 if vae_path is not None:
     print(f"From checkpoint: {vae_path}")
     if vae_path.endswith("safetensors"):
-        from safetensors.torch import load_file
+        from safetensors.torch import load_file, safe_open
         state_dict = load_file(vae_path)
     else:
         state_dict = torch.load(vae_path, map_location="cpu")
@@ -199,10 +206,23 @@ if compile_dit:
         pipeline.transformer.transformer_blocks[i] = torch.compile(pipeline.transformer.transformer_blocks[i])
     print("Add Compile")
 
-# Quantize (when the mode carries an "_and_<quant>" suffix) and then place the pipeline.
-# The order lives inside the helper: quantization has to happen before the offload hooks
-# are installed.
-apply_gpu_memory_mode(pipeline, GPU_memory_mode, device, weight_dtype, exclude_module_name=['scale_shift_table', 'audio_scale_shift_table', 'video_a2v_cross_attn_scale_shift_table', 'audio_a2v_cross_attn_scale_shift_table', ''])
+if GPU_memory_mode == "sequential_cpu_offload":
+    pipeline.enable_sequential_cpu_offload(device=device)
+elif GPU_memory_mode == "model_group_offload":
+    register_auto_device_hook(pipeline.transformer)
+    safe_enable_group_offload(pipeline, onload_device=device, offload_device="cpu", offload_type="leaf_level", use_stream=True)
+elif GPU_memory_mode == "model_cpu_offload_and_qfloat8":
+    convert_model_weight_to_float8(transformer, exclude_module_name=["scale_shift_table", "audio_scale_shift_table", "video_a2v_cross_attn_scale_shift_table", "audio_a2v_cross_attn_scale_shift_table", ""], device=device)
+    convert_weight_dtype_wrapper(transformer, weight_dtype)
+    pipeline.enable_model_cpu_offload(device=device)
+elif GPU_memory_mode == "model_cpu_offload":
+    pipeline.enable_model_cpu_offload(device=device)
+elif GPU_memory_mode == "model_full_load_and_qfloat8":
+    convert_model_weight_to_float8(transformer, exclude_module_name=["scale_shift_table", "audio_scale_shift_table", "video_a2v_cross_attn_scale_shift_table", "audio_a2v_cross_attn_scale_shift_table", ""], device=device)
+    convert_weight_dtype_wrapper(transformer, weight_dtype)
+    pipeline.to(device=device)
+else:
+    pipeline.to(device=device)
 
 generator = torch.Generator(device=device).manual_seed(seed)
 

examples/minimax_h3/predict_i2v.py

@@ -1,6 +1,7 @@
 import os
 import sys
 
+import numpy as np
 import torch
 from PIL import Image
 
@@ -16,9 +17,12 @@ from videox_fun.models import (AutoencoderKLMiniMaxH3,
                                Qwen3VLForConditionalGeneration,
                                Qwen3VLProcessor)
 from videox_fun.pipeline import MiniMaxH3Pipeline
-from videox_fun.utils import (MiniMaxH3Scheduler, apply_gpu_memory_mode,
-                              convert_model_weight_to_float8, merge_lora,
-                              save_videos_with_audio_grid, unmerge_lora)
+from videox_fun.utils import (MiniMaxH3Scheduler, register_auto_device_hook,
+                              safe_enable_group_offload)
+from videox_fun.utils.fp8_optimization import (convert_model_weight_to_float8,
+                                               convert_weight_dtype_wrapper)
+from videox_fun.utils.lora_utils import merge_lora, unmerge_lora
+from videox_fun.utils.utils import save_videos_with_audio_grid
 
 # GPU memory mode, which can be chosen in [model_full_load, model_full_load_and_qfloat8, model_cpu_offload, model_cpu_offload_and_qfloat8, model_group_offload, sequential_cpu_offload].
 # model_full_load means that the entire model will be moved to the GPU.
@@ -80,7 +84,7 @@ validation_image_start  = "asset/1.png"
 validation_image_end    = None
 
 # Use torch.float16 if GPU does not support torch.bfloat16
-# Some graphics cards, such as v100, 2080ti, do not support torch.bfloat16
+# ome graphics cards, such as v100, 2080ti, do not support torch.bfloat16
 weight_dtype        = torch.bfloat16
 prompt              = "一只棕色的狗摇着头，坐在舒适房间里的浅色沙发上。在狗的后面，架子上有一幅镶框的画，周围是粉红色的花朵。房间里柔和温暖的灯光营造出舒适的氛围。"
 seed                = 43
@@ -119,7 +123,7 @@ if transformer_path is not None:
         )
     else:
         if transformer_path.endswith("safetensors"):
-            from safetensors.torch import load_file
+            from safetensors.torch import load_file, safe_open
             state_dict = load_file(transformer_path)
         else:
             state_dict = torch.load(transformer_path, map_location="cpu")
@@ -155,7 +159,7 @@ vae = AutoencoderKLMiniMaxH3.from_pretrained(
 if vae_path is not None:
     print(f"From checkpoint: {vae_path}")
     if vae_path.endswith("safetensors"):
-        from safetensors.torch import load_file
+        from safetensors.torch import load_file, safe_open
         state_dict = load_file(vae_path)
     else:
         state_dict = torch.load(vae_path, map_location="cpu")
@@ -229,14 +233,21 @@ if compile_dit:
         pipeline.transformer.transformer_blocks[i] = torch.compile(pipeline.transformer.transformer_blocks[i])
     print("Add Compile")
 
-# Quantize (when the mode carries an "_and_<quant>" suffix) and then place the pipeline.
-# The order lives inside the helper: quantization has to happen before the offload hooks
-# are installed.
-# The FP8 conversion above has already run (before the FSDP sharding, on purpose); only the dequant
-# wrapper and the memory placement are left, which is what the preconverted tag installs.
-apply_gpu_memory_mode(pipeline, GPU_memory_mode, device, weight_dtype,
-                      quant_tag="qfloat8_preconverted" if GPU_memory_mode.endswith("_and_qfloat8") else None,
-                      exclude_module_name=[])
+if GPU_memory_mode == "sequential_cpu_offload":
+    pipeline.enable_sequential_cpu_offload(device=device)
+elif GPU_memory_mode == "model_group_offload":
+    register_auto_device_hook(pipeline.transformer)
+    safe_enable_group_offload(pipeline, onload_device=device, offload_device="cpu", offload_type="leaf_level", use_stream=True)
+elif GPU_memory_mode == "model_cpu_offload_and_qfloat8":
+    convert_weight_dtype_wrapper(pipeline.transformer, weight_dtype)
+    pipeline.enable_model_cpu_offload(device=device)
+elif GPU_memory_mode == "model_cpu_offload":
+    pipeline.enable_model_cpu_offload(device=device)
+elif GPU_memory_mode == "model_full_load_and_qfloat8":
+    convert_weight_dtype_wrapper(pipeline.transformer, weight_dtype)
+    pipeline.to(device=device)
+else:
+    pipeline.to(device=device)
 
 generator = torch.Generator(device=device).manual_seed(seed)
 

examples/minimax_h3/predict_ref2va.py

@@ -11,15 +11,20 @@ for project_root in project_roots:
 from videox_fun.dist import set_multi_gpus_devices, shard_model
 from videox_fun.models import (AutoencoderKLMiniMaxH3,
                                AutoencoderKLMiniMaxH3Audio,
-                               MiniMaxH3Transformer3DModel, Qwen2TokenizerFast,
+                               MiniMaxH3Transformer3DModel,
+                               Qwen2TokenizerFast,
                                Qwen3VLForConditionalGeneration,
                                Qwen3VLProcessor)
 from videox_fun.pipeline import (MiniMaxH3AudioReference,
-                                 MiniMaxH3ImageReference, MiniMaxH3Pipeline,
+                                 MiniMaxH3ImageReference,
+                                 MiniMaxH3Pipeline,
                                  MiniMaxH3VideoReference)
-from videox_fun.utils import (MiniMaxH3Scheduler, apply_gpu_memory_mode,
-                              convert_model_weight_to_float8, merge_lora,
-                              save_videos_with_audio_grid, unmerge_lora)
+from videox_fun.utils import (MiniMaxH3Scheduler, register_auto_device_hook,
+                              safe_enable_group_offload)
+from videox_fun.utils.fp8_optimization import (convert_model_weight_to_float8,
+                                               convert_weight_dtype_wrapper)
+from videox_fun.utils.lora_utils import merge_lora, unmerge_lora
+from videox_fun.utils.utils import save_videos_with_audio_grid
 
 # GPU memory mode, which can be chosen in [model_full_load, model_full_load_and_qfloat8, model_cpu_offload, model_cpu_offload_and_qfloat8, model_group_offload, sequential_cpu_offload].
 # model_full_load means that the entire model will be moved to the GPU.
@@ -125,7 +130,7 @@ if transformer_path is not None:
         )
     else:
         if transformer_path.endswith("safetensors"):
-            from safetensors.torch import load_file
+            from safetensors.torch import load_file, safe_open
             state_dict = load_file(transformer_path)
         else:
             state_dict = torch.load(transformer_path, map_location="cpu")
@@ -161,7 +166,7 @@ vae = AutoencoderKLMiniMaxH3.from_pretrained(
 if vae_path is not None:
     print(f"From checkpoint: {vae_path}")
     if vae_path.endswith("safetensors"):
-        from safetensors.torch import load_file
+        from safetensors.torch import load_file, safe_open
         state_dict = load_file(vae_path)
     else:
         state_dict = torch.load(vae_path, map_location="cpu")
@@ -235,14 +240,21 @@ if compile_dit:
         pipeline.transformer.transformer_blocks[i] = torch.compile(pipeline.transformer.transformer_blocks[i])
     print("Add Compile")
 
-# Quantize (when the mode carries an "_and_<quant>" suffix) and then place the pipeline.
-# The order lives inside the helper: quantization has to happen before the offload hooks
-# are installed.
-# The FP8 conversion above has already run (before the FSDP sharding, on purpose); only the dequant
-# wrapper and the memory placement are left, which is what the preconverted tag installs.
-apply_gpu_memory_mode(pipeline, GPU_memory_mode, device, weight_dtype,
-                      quant_tag="qfloat8_preconverted" if GPU_memory_mode.endswith("_and_qfloat8") else None,
-                      exclude_module_name=[])
+if GPU_memory_mode == "sequential_cpu_offload":
+    pipeline.enable_sequential_cpu_offload(device=device)
+elif GPU_memory_mode == "model_group_offload":
+    register_auto_device_hook(pipeline.transformer)
+    safe_enable_group_offload(pipeline, onload_device=device, offload_device="cpu", offload_type="leaf_level", use_stream=True)
+elif GPU_memory_mode == "model_cpu_offload_and_qfloat8":
+    convert_weight_dtype_wrapper(pipeline.transformer, weight_dtype)
+    pipeline.enable_model_cpu_offload(device=device)
+elif GPU_memory_mode == "model_cpu_offload":
+    pipeline.enable_model_cpu_offload(device=device)
+elif GPU_memory_mode == "model_full_load_and_qfloat8":
+    convert_weight_dtype_wrapper(pipeline.transformer, weight_dtype)
+    pipeline.to(device=device)
+else:
+    pipeline.to(device=device)
 
 generator = torch.Generator(device=device).manual_seed(seed)
 
@@ -263,6 +275,7 @@ def parse_reference(entry: str):
         return MiniMaxH3AudioReference.from_file(media)
     raise ValueError(f"A reference entry must start with `image=`, `video=` or `audio=`, got {entry!r}.")
 
+
 # Decode every reference at the rate its container carries, which the pipeline's setup resamples onto MiniMax-H3's
 # own 24 fps and the audio VAE's sample rate.
 parsed_references = [parse_reference(entry) for entry in references]

examples/minimax_h3/predict_t2v.py

@@ -1,7 +1,9 @@
 import os
 import sys
 
+import numpy as np
 import torch
+from PIL import Image
 
 current_file_path = os.path.abspath(__file__)
 project_roots = [os.path.dirname(current_file_path), os.path.dirname(os.path.dirname(current_file_path)), os.path.dirname(os.path.dirname(os.path.dirname(current_file_path)))]
@@ -15,9 +17,12 @@ from videox_fun.models import (AutoencoderKLMiniMaxH3,
                                Qwen3VLForConditionalGeneration,
                                Qwen3VLProcessor)
 from videox_fun.pipeline import MiniMaxH3Pipeline
-from videox_fun.utils import (MiniMaxH3Scheduler, apply_gpu_memory_mode,
-                              convert_model_weight_to_float8, merge_lora,
-                              save_videos_with_audio_grid, unmerge_lora)
+from videox_fun.utils import (MiniMaxH3Scheduler, register_auto_device_hook,
+                              safe_enable_group_offload)
+from videox_fun.utils.fp8_optimization import (convert_model_weight_to_float8,
+                                               convert_weight_dtype_wrapper)
+from videox_fun.utils.lora_utils import merge_lora, unmerge_lora
+from videox_fun.utils.utils import save_videos_with_audio_grid
 
 # GPU memory mode, which can be chosen in [model_full_load, model_full_load_and_qfloat8, model_cpu_offload, model_cpu_offload_and_qfloat8, model_group_offload, sequential_cpu_offload].
 # model_full_load means that the entire model will be moved to the GPU.
@@ -111,7 +116,7 @@ if transformer_path is not None:
         )
     else:
         if transformer_path.endswith("safetensors"):
-            from safetensors.torch import load_file
+            from safetensors.torch import load_file, safe_open
             state_dict = load_file(transformer_path)
         else:
             state_dict = torch.load(transformer_path, map_location="cpu")
@@ -147,7 +152,7 @@ vae = AutoencoderKLMiniMaxH3.from_pretrained(
 if vae_path is not None:
     print(f"From checkpoint: {vae_path}")
     if vae_path.endswith("safetensors"):
-        from safetensors.torch import load_file
+        from safetensors.torch import load_file, safe_open
         state_dict = load_file(vae_path)
     else:
         state_dict = torch.load(vae_path, map_location="cpu")
@@ -221,14 +226,21 @@ if compile_dit:
         pipeline.transformer.transformer_blocks[i] = torch.compile(pipeline.transformer.transformer_blocks[i])
     print("Add Compile")
 
-# Quantize (when the mode carries an "_and_<quant>" suffix) and then place the pipeline.
-# The order lives inside the helper: quantization has to happen before the offload hooks
-# are installed.
-# The FP8 conversion above has already run (before the FSDP sharding, on purpose); only the dequant
-# wrapper and the memory placement are left, which is what the preconverted tag installs.
-apply_gpu_memory_mode(pipeline, GPU_memory_mode, device, weight_dtype,
-                      quant_tag="qfloat8_preconverted" if GPU_memory_mode.endswith("_and_qfloat8") else None,
-                      exclude_module_name=[])
+if GPU_memory_mode == "sequential_cpu_offload":
+    pipeline.enable_sequential_cpu_offload(device=device)
+elif GPU_memory_mode == "model_group_offload":
+    register_auto_device_hook(pipeline.transformer)
+    safe_enable_group_offload(pipeline, onload_device=device, offload_device="cpu", offload_type="leaf_level", use_stream=True)
+elif GPU_memory_mode == "model_cpu_offload_and_qfloat8":
+    convert_weight_dtype_wrapper(pipeline.transformer, weight_dtype)
+    pipeline.enable_model_cpu_offload(device=device)
+elif GPU_memory_mode == "model_cpu_offload":
+    pipeline.enable_model_cpu_offload(device=device)
+elif GPU_memory_mode == "model_full_load_and_qfloat8":
+    convert_weight_dtype_wrapper(pipeline.transformer, weight_dtype)
+    pipeline.to(device=device)
+else:
+    pipeline.to(device=device)
 
 generator = torch.Generator(device=device).manual_seed(seed)
 

examples/minimax_h3_fun/predict_v2v_control.py

@@ -1,7 +1,9 @@
 import os
 import sys
 
+import numpy as np
 import torch
+from PIL import Image
 
 current_file_path = os.path.abspath(__file__)
 project_roots = [os.path.dirname(current_file_path), os.path.dirname(os.path.dirname(current_file_path)), os.path.dirname(os.path.dirname(os.path.dirname(current_file_path)))]
@@ -16,10 +18,14 @@ from videox_fun.models import (AutoencoderKLMiniMaxH3,
                                Qwen3VLForConditionalGeneration,
                                Qwen3VLProcessor)
 from videox_fun.pipeline import MiniMaxH3ControlPipeline
-from videox_fun.utils import (MiniMaxH3Scheduler, apply_gpu_memory_mode,
-                              convert_model_weight_to_float8,
-                              get_video_to_video_latent, merge_lora,
-                              save_videos_with_audio_grid, unmerge_lora)
+from videox_fun.utils import (MiniMaxH3Scheduler, register_auto_device_hook,
+                              safe_enable_group_offload)
+from videox_fun.utils.fp8_optimization import (convert_model_weight_to_float8,
+                                               convert_weight_dtype_wrapper)
+from videox_fun.utils.lora_utils import merge_lora, unmerge_lora
+from videox_fun.utils.utils import (get_video_to_video_latent,
+                                    save_videos_with_audio_grid)
+
 
 # GPU memory mode, which can be chosen in [model_full_load, model_full_load_and_qfloat8, model_cpu_offload, model_cpu_offload_and_qfloat8, model_group_offload, sequential_cpu_offload].
 # model_full_load means that the entire model will be moved to the GPU.
@@ -86,7 +92,7 @@ fps                 = 24
 control_context_scale = 1.00
 
 # Use torch.float16 if GPU does not support torch.bfloat16
-# Some graphics cards, such as v100, 2080ti, do not support torch.bfloat16
+# ome graphics cards, such as v100, 2080ti, do not support torch.bfloat16
 weight_dtype        = torch.bfloat16
 control_video       = "asset/pose.mp4"
 # Inpaint inputs, only read by checkpoints trained with `--enable_inpaint` (control_in_dim widened, e.g. 49):
@@ -140,7 +146,7 @@ transformer = MiniMaxH3ControlTransformer3DModel.from_pretrained(
 if transformer_path is not None:
     print(f"From checkpoint: {transformer_path}")
     if transformer_path.endswith("safetensors"):
-        from safetensors.torch import load_file
+        from safetensors.torch import load_file, safe_open
         state_dict = load_file(transformer_path)
     else:
         state_dict = torch.load(transformer_path, map_location="cpu")
@@ -161,7 +167,7 @@ vae = AutoencoderKLMiniMaxH3.from_pretrained(
 if vae_path is not None:
     print(f"From checkpoint: {vae_path}")
     if vae_path.endswith("safetensors"):
-        from safetensors.torch import load_file
+        from safetensors.torch import load_file, safe_open
         state_dict = load_file(vae_path)
     else:
         state_dict = torch.load(vae_path, map_location="cpu")
@@ -237,14 +243,21 @@ if compile_dit:
         pipeline.transformer.transformer_blocks[i] = torch.compile(pipeline.transformer.transformer_blocks[i])
     print("Add Compile")
 
-# Quantize (when the mode carries an "_and_<quant>" suffix) and then place the pipeline.
-# The order lives inside the helper: quantization has to happen before the offload hooks
-# are installed.
-# The FP8 conversion above has already run (before the FSDP sharding, on purpose); only the dequant
-# wrapper and the memory placement are left, which is what the preconverted tag installs.
-apply_gpu_memory_mode(pipeline, GPU_memory_mode, device, weight_dtype,
-                      quant_tag="qfloat8_preconverted" if GPU_memory_mode.endswith("_and_qfloat8") else None,
-                      exclude_module_name=[])
+if GPU_memory_mode == "sequential_cpu_offload":
+    pipeline.enable_sequential_cpu_offload(device=device)
+elif GPU_memory_mode == "model_group_offload":
+    register_auto_device_hook(pipeline.transformer)
+    safe_enable_group_offload(pipeline, onload_device=device, offload_device="cpu", offload_type="leaf_level", use_stream=True)
+elif GPU_memory_mode == "model_cpu_offload_and_qfloat8":
+    convert_weight_dtype_wrapper(pipeline.transformer, weight_dtype)
+    pipeline.enable_model_cpu_offload(device=device)
+elif GPU_memory_mode == "model_cpu_offload":
+    pipeline.enable_model_cpu_offload(device=device)
+elif GPU_memory_mode == "model_full_load_and_qfloat8":
+    convert_weight_dtype_wrapper(pipeline.transformer, weight_dtype)
+    pipeline.to(device=device)
+else:
+    pipeline.to(device=device)
 
 generator = torch.Generator(device=device).manual_seed(seed)
 

examples/minimax_h3_fun/predict_v2v_control_inpaint.py

@@ -1,7 +1,9 @@
 import os
 import sys
 
+import numpy as np
 import torch
+from PIL import Image
 
 current_file_path = os.path.abspath(__file__)
 project_roots = [os.path.dirname(current_file_path), os.path.dirname(os.path.dirname(current_file_path)), os.path.dirname(os.path.dirname(os.path.dirname(current_file_path)))]
@@ -16,10 +18,14 @@ from videox_fun.models import (AutoencoderKLMiniMaxH3,
                                Qwen3VLForConditionalGeneration,
                                Qwen3VLProcessor)
 from videox_fun.pipeline import MiniMaxH3ControlPipeline
-from videox_fun.utils import (MiniMaxH3Scheduler, apply_gpu_memory_mode,
-                              convert_model_weight_to_float8,
-                              get_video_to_video_latent, merge_lora,
-                              save_videos_with_audio_grid, unmerge_lora)
+from videox_fun.utils import (MiniMaxH3Scheduler, register_auto_device_hook,
+                              safe_enable_group_offload)
+from videox_fun.utils.fp8_optimization import (convert_model_weight_to_float8,
+                                               convert_weight_dtype_wrapper)
+from videox_fun.utils.lora_utils import merge_lora, unmerge_lora
+from videox_fun.utils.utils import (get_video_to_video_latent,
+                                    save_videos_with_audio_grid)
+
 
 # GPU memory mode, which can be chosen in [model_full_load, model_full_load_and_qfloat8, model_cpu_offload, model_cpu_offload_and_qfloat8, model_group_offload, sequential_cpu_offload].
 # model_full_load means that the entire model will be moved to the GPU.
@@ -86,7 +92,7 @@ fps                 = 24
 control_context_scale = 1.00
 
 # Use torch.float16 if GPU does not support torch.bfloat16
-# Some graphics cards, such as v100, 2080ti, do not support torch.bfloat16
+# ome graphics cards, such as v100, 2080ti, do not support torch.bfloat16
 weight_dtype        = torch.bfloat16
 # Path of the control (e.g. pose) video; leaving it None zeroes the control channels of the side branch. With
 # inpaint inputs given the mask then guides the run on its own (the layout training reaches when it drops the
@@ -142,7 +148,7 @@ transformer = MiniMaxH3ControlTransformer3DModel.from_pretrained(
 if transformer_path is not None:
     print(f"From checkpoint: {transformer_path}")
     if transformer_path.endswith("safetensors"):
-        from safetensors.torch import load_file
+        from safetensors.torch import load_file, safe_open
         state_dict = load_file(transformer_path)
     else:
         state_dict = torch.load(transformer_path, map_location="cpu")
@@ -163,7 +169,7 @@ vae = AutoencoderKLMiniMaxH3.from_pretrained(
 if vae_path is not None:
     print(f"From checkpoint: {vae_path}")
     if vae_path.endswith("safetensors"):
-        from safetensors.torch import load_file
+        from safetensors.torch import load_file, safe_open
         state_dict = load_file(vae_path)
     else:
         state_dict = torch.load(vae_path, map_location="cpu")
@@ -253,14 +259,21 @@ if compile_dit:
         pipeline.transformer.transformer_blocks[i] = torch.compile(pipeline.transformer.transformer_blocks[i])
     print("Add Compile")
 
-# Quantize (when the mode carries an "_and_<quant>" suffix) and then place the pipeline.
-# The order lives inside the helper: quantization has to happen before the offload hooks
-# are installed.
-# The FP8 conversion above has already run (before the FSDP sharding, on purpose); only the dequant
-# wrapper and the memory placement are left, which is what the preconverted tag installs.
-apply_gpu_memory_mode(pipeline, GPU_memory_mode, device, weight_dtype,
-                      quant_tag="qfloat8_preconverted" if GPU_memory_mode.endswith("_and_qfloat8") else None,
-                      exclude_module_name=[])
+if GPU_memory_mode == "sequential_cpu_offload":
+    pipeline.enable_sequential_cpu_offload(device=device)
+elif GPU_memory_mode == "model_group_offload":
+    register_auto_device_hook(pipeline.transformer)
+    safe_enable_group_offload(pipeline, onload_device=device, offload_device="cpu", offload_type="leaf_level", use_stream=True)
+elif GPU_memory_mode == "model_cpu_offload_and_qfloat8":
+    convert_weight_dtype_wrapper(pipeline.transformer, weight_dtype)
+    pipeline.enable_model_cpu_offload(device=device)
+elif GPU_memory_mode == "model_cpu_offload":
+    pipeline.enable_model_cpu_offload(device=device)
+elif GPU_memory_mode == "model_full_load_and_qfloat8":
+    convert_weight_dtype_wrapper(pipeline.transformer, weight_dtype)
+    pipeline.to(device=device)
+else:
+    pipeline.to(device=device)
 
 generator = torch.Generator(device=device).manual_seed(seed)
 

examples/mova/predict_i2v.py

@@ -17,12 +17,17 @@ from videox_fun.models import (AutoencoderKLMOVAAudio, AutoencoderKLWan,
                                UMT5EncoderModel, WanAudioTransformer3DModel,
                                WanTransformer3DModel)
 from videox_fun.pipeline import MOVAPipeline
-from videox_fun.utils import (FlowDPMSolverMultistepScheduler,
-                              FlowUniPCMultistepScheduler,
-                              apply_gpu_memory_mode, merge_lora,
-                              save_videos_with_audio_grid, unmerge_lora)
+from videox_fun.utils import (register_auto_device_hook,
+                              safe_enable_group_offload)
+from videox_fun.utils.fm_solvers import FlowDPMSolverMultistepScheduler
+from videox_fun.utils.fm_solvers_unipc import FlowUniPCMultistepScheduler
+from videox_fun.utils.fp8_optimization import (convert_model_weight_to_float8,
+                                               convert_weight_dtype_wrapper,
+                                               replace_parameters_by_name)
+from videox_fun.utils.lora_utils import merge_lora, unmerge_lora
+from videox_fun.utils.utils import save_videos_with_audio_grid
 
-# GPU memory mode, which can be chosen in [model_full_load, model_full_load_and_qfloat8, model_cpu_offload, model_cpu_offload_and_qfloat8, model_group_offload, sequential_cpu_offload].
+# GPU memory mode, which can be chosen in [model_full_load, model_full_load_and_qfloat8, model_cpu_offload, model_cpu_offload_and_qfloat8, sequential_cpu_offload].
 # model_full_load means that the entire model will be moved to the GPU.
 # 
 # model_full_load_and_qfloat8 means that the entire model will be moved to the GPU,
@@ -288,11 +293,31 @@ if compile_dit:
             pipeline.transformer_audio.blocks[i] = torch.compile(pipeline.transformer_audio.blocks[i])
         print("Add Compile")
 
-# Quantize (when the mode carries an "_and_<quant>" suffix) and then place the pipeline.
-# The order lives inside the helper: quantization has to happen before the offload hooks
-# are installed, and both transformers of this MoE setup are handled in one call, which
-# is exactly the bookkeeping the old 30-line if/elif chain repeated per script.
-apply_gpu_memory_mode(pipeline, GPU_memory_mode, device, weight_dtype)
+if GPU_memory_mode == "sequential_cpu_offload":
+    replace_parameters_by_name(pipeline.transformer, ["modulation",], device=device)
+    replace_parameters_by_name(pipeline.transformer_2, ["modulation",], device=device)
+    pipeline.transformer.freqs = pipeline.transformer.freqs.to(device=device)
+    pipeline.transformer_2.freqs = pipeline.transformer_2.freqs.to(device=device)
+    pipeline.enable_sequential_cpu_offload(device=device)
+elif GPU_memory_mode == "model_group_offload":
+    register_auto_device_hook(pipeline.transformer)
+    safe_enable_group_offload(pipeline, onload_device=device, offload_device="cpu", offload_type="leaf_level", use_stream=True)
+elif GPU_memory_mode == "model_cpu_offload_and_qfloat8":
+    convert_model_weight_to_float8(pipeline.transformer, exclude_module_name=["modulation",], device=device)
+    convert_model_weight_to_float8(pipeline.transformer_2, exclude_module_name=["modulation",], device=device)
+    convert_weight_dtype_wrapper(pipeline.transformer, weight_dtype)
+    convert_weight_dtype_wrapper(pipeline.transformer_2, weight_dtype)
+    pipeline.enable_model_cpu_offload(device=device)
+elif GPU_memory_mode == "model_cpu_offload":
+    pipeline.enable_model_cpu_offload(device=device)
+elif GPU_memory_mode == "model_full_load_and_qfloat8":
+    convert_model_weight_to_float8(pipeline.transformer, exclude_module_name=["modulation",], device=device)
+    convert_model_weight_to_float8(pipeline.transformer_2, exclude_module_name=["modulation",], device=device)
+    convert_weight_dtype_wrapper(pipeline.transformer, weight_dtype)
+    convert_weight_dtype_wrapper(pipeline.transformer_2, weight_dtype)
+    pipeline.to(device=device)
+else:
+    pipeline.to(device=device)
 
 generator = torch.Generator(device=device).manual_seed(seed)
 

examples/phantom/predict_s2v.py

@@ -18,11 +18,16 @@ from videox_fun.models import (AutoencoderKLWan, AutoTokenizer,
                                WanT5EncoderModel, WanTransformer3DModel)
 from videox_fun.models.cache_utils import get_teacache_coefficients
 from videox_fun.pipeline import WanFunPhantomPipeline
-from videox_fun.utils import (FlowDPMSolverMultistepScheduler,
-                              FlowUniPCMultistepScheduler,
-                              apply_gpu_memory_mode, filter_kwargs,
-                              get_image_latent, merge_lora, save_videos_grid,
-                              unmerge_lora)
+from videox_fun.utils import (register_auto_device_hook,
+                              safe_enable_group_offload)
+from videox_fun.utils.fp8_optimization import (convert_model_weight_to_float8,
+                                               convert_weight_dtype_wrapper,
+                                               replace_parameters_by_name)
+from videox_fun.utils.lora_utils import merge_lora, unmerge_lora
+from videox_fun.utils.utils import (filter_kwargs, get_image_latent,
+                                    save_videos_grid)
+from videox_fun.utils.fm_solvers import FlowDPMSolverMultistepScheduler
+from videox_fun.utils.fm_solvers_unipc import FlowUniPCMultistepScheduler
 
 # GPU memory mode, which can be chosen in [model_full_load, model_full_load_and_qfloat8, model_cpu_offload, model_cpu_offload_and_qfloat8, model_group_offload, sequential_cpu_offload].
 # model_full_load means that the entire model will be moved to the GPU.
@@ -103,7 +108,7 @@ video_length        = 81
 fps                 = 16
 
 # Use torch.float16 if GPU does not support torch.bfloat16
-# Some graphics cards, such as v100, 2080ti, do not support torch.bfloat16
+# ome graphics cards, such as v100, 2080ti, do not support torch.bfloat16
 weight_dtype            = torch.bfloat16
 subject_ref_images      = ["asset/ref_1.png", "asset/ref_2.png"]
 
@@ -135,7 +140,7 @@ transformer = WanTransformer3DModel.from_pretrained(
 if transformer_path is not None:
     print(f"From checkpoint: {transformer_path}")
     if transformer_path.endswith("safetensors"):
-        from safetensors.torch import load_file
+        from safetensors.torch import load_file, safe_open
         state_dict = load_file(transformer_path)
     else:
         state_dict = torch.load(transformer_path, map_location="cpu")
@@ -153,7 +158,7 @@ vae = AutoencoderKLWan.from_pretrained(
 if vae_path is not None:
     print(f"From checkpoint: {vae_path}")
     if vae_path.endswith("safetensors"):
-        from safetensors.torch import load_file
+        from safetensors.torch import load_file, safe_open
         state_dict = load_file(vae_path)
     else:
         state_dict = torch.load(vae_path, map_location="cpu")
@@ -177,7 +182,7 @@ text_encoder = WanT5EncoderModel.from_pretrained(
 text_encoder = text_encoder.eval()
 
 # Get Scheduler
-Chosen_Scheduler = {
+Chosen_Scheduler = scheduler_dict = {
     "Flow": FlowMatchEulerDiscreteScheduler,
     "Flow_Unipc": FlowUniPCMultistepScheduler,
     "Flow_DPM++": FlowDPMSolverMultistepScheduler,
@@ -213,10 +218,25 @@ if compile_dit:
         pipeline.transformer.blocks[i] = torch.compile(pipeline.transformer.blocks[i])
     print("Add Compile")
 
-# Quantize (when the mode carries an "_and_<quant>" suffix) and then place the pipeline.
-# The order lives inside the helper: quantization has to happen before the offload hooks
-# are installed.
-apply_gpu_memory_mode(pipeline, GPU_memory_mode, device, weight_dtype)
+if GPU_memory_mode == "sequential_cpu_offload":
+    replace_parameters_by_name(transformer, ["modulation",], device=device)
+    transformer.freqs = transformer.freqs.to(device=device)
+    pipeline.enable_sequential_cpu_offload(device=device)
+elif GPU_memory_mode == "model_group_offload":
+    register_auto_device_hook(pipeline.transformer)
+    safe_enable_group_offload(pipeline, onload_device=device, offload_device="cpu", offload_type="leaf_level", use_stream=True)
+elif GPU_memory_mode == "model_cpu_offload_and_qfloat8":
+    convert_model_weight_to_float8(transformer, exclude_module_name=["modulation",], device=device)
+    convert_weight_dtype_wrapper(transformer, weight_dtype)
+    pipeline.enable_model_cpu_offload(device=device)
+elif GPU_memory_mode == "model_cpu_offload":
+    pipeline.enable_model_cpu_offload(device=device)
+elif GPU_memory_mode == "model_full_load_and_qfloat8":
+    convert_model_weight_to_float8(transformer, exclude_module_name=["modulation",], device=device)
+    convert_weight_dtype_wrapper(transformer, weight_dtype)
+    pipeline.to(device=device)
+else:
+    pipeline.to(device=device)
 
 coefficients = get_teacache_coefficients(model_name) if enable_teacache else None
 if coefficients is not None:

examples/qwenimage/predict_i2i_layered.py

@@ -17,11 +17,16 @@ from videox_fun.models import (AutoencoderKLQwenImage,
                                QwenImageTransformer2DModel)
 from videox_fun.models.cache_utils import get_teacache_coefficients
 from videox_fun.pipeline import QwenImageLayeredPipeline
-from videox_fun.utils import (FlowDPMSolverMultistepScheduler,
-                              FlowUniPCMultistepScheduler,
-                              apply_gpu_memory_mode, merge_lora, unmerge_lora)
+from videox_fun.utils import (register_auto_device_hook,
+                              safe_enable_group_offload)
+from videox_fun.utils.fm_solvers import FlowDPMSolverMultistepScheduler
+from videox_fun.utils.fm_solvers_unipc import FlowUniPCMultistepScheduler
+from videox_fun.utils.fp8_optimization import (convert_model_weight_to_float8,
+                                               convert_weight_dtype_wrapper)
+from videox_fun.utils.lora_utils import merge_lora, unmerge_lora
+from videox_fun.utils.utils import get_image
 
-# GPU memory mode, which can be chosen in [model_full_load, model_full_load_and_qfloat8, model_cpu_offload, model_cpu_offload_and_qfloat8, model_group_offload, sequential_cpu_offload].
+# GPU memory mode, which can be chosen in [model_full_load, model_full_load_and_qfloat8, model_cpu_offload, model_cpu_offload_and_qfloat8, sequential_cpu_offload].
 # model_full_load means that the entire model will be moved to the GPU.
 # 
 # model_full_load_and_qfloat8 means that the entire model will be moved to the GPU,
@@ -81,7 +86,7 @@ lora_path           = None
 resolution          = 640
 
 # Use torch.float16 if GPU does not support torch.bfloat16
-# Some graphics cards, such as v100, 2080ti, do not support torch.bfloat16
+# ome graphics cards, such as v100, 2080ti, do not support torch.bfloat16
 weight_dtype        = torch.bfloat16
 image               = "asset/yarn-art-pikachu.png"
 
@@ -106,7 +111,7 @@ transformer = QwenImageTransformer2DModel.from_pretrained(
 if transformer_path is not None:
     print(f"From checkpoint: {transformer_path}")
     if transformer_path.endswith("safetensors"):
-        from safetensors.torch import load_file
+        from safetensors.torch import load_file, safe_open
         state_dict = load_file(transformer_path)
     else:
         state_dict = torch.load(transformer_path, map_location="cpu")
@@ -124,7 +129,7 @@ vae = AutoencoderKLQwenImage.from_pretrained(
 if vae_path is not None:
     print(f"From checkpoint: {vae_path}")
     if vae_path.endswith("safetensors"):
-        from safetensors.torch import load_file
+        from safetensors.torch import load_file, safe_open
         state_dict = load_file(vae_path)
     else:
         state_dict = torch.load(vae_path, map_location="cpu")
@@ -148,7 +153,7 @@ processor = Qwen2VLProcessor.from_pretrained(
 )
 
 # Get Scheduler
-Chosen_Scheduler = {
+Chosen_Scheduler = scheduler_dict = {
     "Flow": FlowMatchEulerDiscreteScheduler,
     "Flow_Unipc": FlowUniPCMultistepScheduler,
     "Flow_DPM++": FlowDPMSolverMultistepScheduler,
@@ -175,6 +180,9 @@ if ulysses_degree > 1 or ring_degree > 1:
         pipeline.transformer = shard_fn(pipeline.transformer)
         print("Add FSDP DIT")
     if fsdp_text_encoder:
+        from functools import partial
+
+        from videox_fun.dist import set_multi_gpus_devices, shard_model
         shard_fn = partial(shard_model, device_id=device, param_dtype=weight_dtype, module_to_wrapper=text_encoder.language_model.layers)
         pipeline.text_encoder = shard_fn(pipeline.text_encoder)
         print("Add FSDP TEXT ENCODER")
@@ -184,10 +192,23 @@ if compile_dit:
         pipeline.transformer.transformer_blocks[i] = torch.compile(pipeline.transformer.transformer_blocks[i])
     print("Add Compile")
 
-# Quantize (when the mode carries an "_and_<quant>" suffix) and then place the pipeline.
-# The order lives inside the helper: quantization has to happen before the offload hooks
-# are installed.
-apply_gpu_memory_mode(pipeline, GPU_memory_mode, device, weight_dtype, exclude_module_name=['img_in', 'txt_in', 'timestep'])
+if GPU_memory_mode == "sequential_cpu_offload":
+    pipeline.enable_sequential_cpu_offload(device=device)
+elif GPU_memory_mode == "model_group_offload":
+    register_auto_device_hook(pipeline.transformer)
+    safe_enable_group_offload(pipeline, onload_device=device, offload_device="cpu", offload_type="leaf_level", use_stream=True)
+elif GPU_memory_mode == "model_cpu_offload_and_qfloat8":
+    convert_model_weight_to_float8(transformer, exclude_module_name=["img_in", "txt_in", "timestep"], device=device)
+    convert_weight_dtype_wrapper(transformer, weight_dtype)
+    pipeline.enable_model_cpu_offload(device=device)
+elif GPU_memory_mode == "model_cpu_offload":
+    pipeline.enable_model_cpu_offload(device=device)
+elif GPU_memory_mode == "model_full_load_and_qfloat8":
+    convert_model_weight_to_float8(transformer, exclude_module_name=["img_in", "txt_in", "timestep"], device=device)
+    convert_weight_dtype_wrapper(transformer, weight_dtype)
+    pipeline.to(device=device)
+else:
+    pipeline.to(device=device)
 
 coefficients = get_teacache_coefficients(model_name) if enable_teacache else None
 if coefficients is not None:

examples/qwenimage/predict_t2i.py

@@ -2,7 +2,8 @@ import os
 import sys
 
 import torch
-from diffusers import FlowMatchEulerDiscreteScheduler
+
+from diffusers import (FlowMatchEulerDiscreteScheduler)
 
 current_file_path = os.path.abspath(__file__)
 project_roots = [os.path.dirname(current_file_path), os.path.dirname(os.path.dirname(current_file_path)), os.path.dirname(os.path.dirname(os.path.dirname(current_file_path)))]
@@ -15,11 +16,15 @@ from videox_fun.models import (AutoencoderKLQwenImage,
                                Qwen2Tokenizer, QwenImageTransformer2DModel)
 from videox_fun.models.cache_utils import get_teacache_coefficients
 from videox_fun.pipeline import QwenImagePipeline
-from videox_fun.utils import (FlowDPMSolverMultistepScheduler,
-                              FlowUniPCMultistepScheduler,
-                              apply_gpu_memory_mode, merge_lora, unmerge_lora)
+from videox_fun.utils import (register_auto_device_hook,
+                              safe_enable_group_offload)
+from videox_fun.utils.fm_solvers import FlowDPMSolverMultistepScheduler
+from videox_fun.utils.fm_solvers_unipc import FlowUniPCMultistepScheduler
+from videox_fun.utils.fp8_optimization import (convert_model_weight_to_float8,
+                                               convert_weight_dtype_wrapper)
+from videox_fun.utils.lora_utils import merge_lora, unmerge_lora
 
-# GPU memory mode, which can be chosen in [model_full_load, model_full_load_and_qfloat8, model_cpu_offload, model_cpu_offload_and_qfloat8, model_group_offload, sequential_cpu_offload].
+# GPU memory mode, which can be chosen in [model_full_load, model_full_load_and_qfloat8, model_cpu_offload, model_cpu_offload_and_qfloat8, sequential_cpu_offload].
 # model_full_load means that the entire model will be moved to the GPU.
 # 
 # model_full_load_and_qfloat8 means that the entire model will be moved to the GPU,
@@ -79,7 +84,7 @@ lora_path           = None
 sample_size         = [1344, 768]
 
 # Use torch.float16 if GPU does not support torch.bfloat16
-# Some graphics cards, such as v100, 2080ti, do not support torch.bfloat16
+# ome graphics cards, such as v100, 2080ti, do not support torch.bfloat16
 weight_dtype        = torch.bfloat16
 # Please use as detailed a prompt as possible to describe the object that needs to be generated.
 prompt              = "1girl, black_hair, brown_eyes, earrings, freckles, grey_background, jewelry, lips, long_hair, looking_at_viewer, nose, piercing, realistic, red_lips, solo, upper_body"
@@ -102,7 +107,7 @@ transformer = QwenImageTransformer2DModel.from_pretrained(
 if transformer_path is not None:
     print(f"From checkpoint: {transformer_path}")
     if transformer_path.endswith("safetensors"):
-        from safetensors.torch import load_file
+        from safetensors.torch import load_file, safe_open
         state_dict = load_file(transformer_path)
     else:
         state_dict = torch.load(transformer_path, map_location="cpu")
@@ -120,7 +125,7 @@ vae = AutoencoderKLQwenImage.from_pretrained(
 if vae_path is not None:
     print(f"From checkpoint: {vae_path}")
     if vae_path.endswith("safetensors"):
-        from safetensors.torch import load_file
+        from safetensors.torch import load_file, safe_open
         state_dict = load_file(vae_path)
     else:
         state_dict = torch.load(vae_path, map_location="cpu")
@@ -138,7 +143,7 @@ text_encoder = Qwen2_5_VLForConditionalGeneration.from_pretrained(
 )
 
 # Get Scheduler
-Chosen_Scheduler = {
+Chosen_Scheduler = scheduler_dict = {
     "Flow": FlowMatchEulerDiscreteScheduler,
     "Flow_Unipc": FlowUniPCMultistepScheduler,
     "Flow_DPM++": FlowDPMSolverMultistepScheduler,
@@ -164,6 +169,8 @@ if ulysses_degree > 1 or ring_degree > 1:
         pipeline.transformer = shard_fn(pipeline.transformer)
         print("Add FSDP DIT")
     if fsdp_text_encoder:
+        from functools import partial
+        from videox_fun.dist import set_multi_gpus_devices, shard_model
         shard_fn = partial(shard_model, device_id=device, param_dtype=weight_dtype, module_to_wrapper=text_encoder.language_model.layers)
         pipeline.text_encoder = shard_fn(pipeline.text_encoder)
         print("Add FSDP TEXT ENCODER")
@@ -173,10 +180,23 @@ if compile_dit:
         pipeline.transformer.transformer_blocks[i] = torch.compile(pipeline.transformer.transformer_blocks[i])
     print("Add Compile")
 
-# Quantize (when the mode carries an "_and_<quant>" suffix) and then place the pipeline.
-# The order lives inside the helper: quantization has to happen before the offload hooks
-# are installed.
-apply_gpu_memory_mode(pipeline, GPU_memory_mode, device, weight_dtype, exclude_module_name=['img_in', 'txt_in', 'timestep'])
+if GPU_memory_mode == "sequential_cpu_offload":
+    pipeline.enable_sequential_cpu_offload(device=device)
+elif GPU_memory_mode == "model_group_offload":
+    register_auto_device_hook(pipeline.transformer)
+    safe_enable_group_offload(pipeline, onload_device=device, offload_device="cpu", offload_type="leaf_level", use_stream=True)
+elif GPU_memory_mode == "model_cpu_offload_and_qfloat8":
+    convert_model_weight_to_float8(transformer, exclude_module_name=["img_in", "txt_in", "timestep"], device=device)
+    convert_weight_dtype_wrapper(transformer, weight_dtype)
+    pipeline.enable_model_cpu_offload(device=device)
+elif GPU_memory_mode == "model_cpu_offload":
+    pipeline.enable_model_cpu_offload(device=device)
+elif GPU_memory_mode == "model_full_load_and_qfloat8":
+    convert_model_weight_to_float8(transformer, exclude_module_name=["img_in", "txt_in", "timestep"], device=device)
+    convert_weight_dtype_wrapper(transformer, weight_dtype)
+    pipeline.to(device=device)
+else:
+    pipeline.to(device=device)
 
 coefficients = get_teacache_coefficients(model_name) if enable_teacache else None
 if coefficients is not None:

examples/qwenimage/predict_t2i_edit.py

@@ -3,6 +3,7 @@ import sys
 
 import torch
 from diffusers import FlowMatchEulerDiscreteScheduler
+from PIL import Image
 
 current_file_path = os.path.abspath(__file__)
 project_roots = [os.path.dirname(current_file_path), os.path.dirname(os.path.dirname(current_file_path)), os.path.dirname(os.path.dirname(os.path.dirname(current_file_path)))]
@@ -16,12 +17,16 @@ from videox_fun.models import (AutoencoderKLQwenImage,
                                QwenImageTransformer2DModel)
 from videox_fun.models.cache_utils import get_teacache_coefficients
 from videox_fun.pipeline import QwenImageEditPipeline
-from videox_fun.utils import (FlowDPMSolverMultistepScheduler,
-                              FlowUniPCMultistepScheduler,
-                              apply_gpu_memory_mode, get_image, merge_lora,
-                              unmerge_lora)
+from videox_fun.utils import (register_auto_device_hook,
+                              safe_enable_group_offload)
+from videox_fun.utils.fm_solvers import FlowDPMSolverMultistepScheduler
+from videox_fun.utils.fm_solvers_unipc import FlowUniPCMultistepScheduler
+from videox_fun.utils.fp8_optimization import (convert_model_weight_to_float8,
+                                               convert_weight_dtype_wrapper)
+from videox_fun.utils.lora_utils import merge_lora, unmerge_lora
+from videox_fun.utils.utils import get_image
 
-# GPU memory mode, which can be chosen in [model_full_load, model_full_load_and_qfloat8, model_cpu_offload, model_cpu_offload_and_qfloat8, model_group_offload, sequential_cpu_offload].
+# GPU memory mode, which can be chosen in [model_full_load, model_full_load_and_qfloat8, model_cpu_offload, model_cpu_offload_and_qfloat8, sequential_cpu_offload].
 # model_full_load means that the entire model will be moved to the GPU.
 # 
 # model_full_load_and_qfloat8 means that the entire model will be moved to the GPU,
@@ -81,7 +86,7 @@ lora_path           = None
 sample_size         = [1344, 768]
 
 # Use torch.float16 if GPU does not support torch.bfloat16
-# Some graphics cards, such as v100, 2080ti, do not support torch.bfloat16
+# ome graphics cards, such as v100, 2080ti, do not support torch.bfloat16
 weight_dtype        = torch.bfloat16
 image               = "asset/8.png"
 # Please use as detailed a prompt as possible to describe the object that needs to be generated.
@@ -105,7 +110,7 @@ transformer = QwenImageTransformer2DModel.from_pretrained(
 if transformer_path is not None:
     print(f"From checkpoint: {transformer_path}")
     if transformer_path.endswith("safetensors"):
-        from safetensors.torch import load_file
+        from safetensors.torch import load_file, safe_open
         state_dict = load_file(transformer_path)
     else:
         state_dict = torch.load(transformer_path, map_location="cpu")
@@ -123,7 +128,7 @@ vae = AutoencoderKLQwenImage.from_pretrained(
 if vae_path is not None:
     print(f"From checkpoint: {vae_path}")
     if vae_path.endswith("safetensors"):
-        from safetensors.torch import load_file
+        from safetensors.torch import load_file, safe_open
         state_dict = load_file(vae_path)
     else:
         state_dict = torch.load(vae_path, map_location="cpu")
@@ -147,7 +152,7 @@ processor = Qwen2VLProcessor.from_pretrained(
 )
 
 # Get Scheduler
-Chosen_Scheduler = {
+Chosen_Scheduler = scheduler_dict = {
     "Flow": FlowMatchEulerDiscreteScheduler,
     "Flow_Unipc": FlowUniPCMultistepScheduler,
     "Flow_DPM++": FlowDPMSolverMultistepScheduler,
@@ -174,6 +179,9 @@ if ulysses_degree > 1 or ring_degree > 1:
         pipeline.transformer = shard_fn(pipeline.transformer)
         print("Add FSDP DIT")
     if fsdp_text_encoder:
+        from functools import partial
+
+        from videox_fun.dist import set_multi_gpus_devices, shard_model
         shard_fn = partial(shard_model, device_id=device, param_dtype=weight_dtype, module_to_wrapper=text_encoder.language_model.layers)
         pipeline.text_encoder = shard_fn(pipeline.text_encoder)
         print("Add FSDP TEXT ENCODER")
@@ -183,10 +191,23 @@ if compile_dit:
         pipeline.transformer.transformer_blocks[i] = torch.compile(pipeline.transformer.transformer_blocks[i])
     print("Add Compile")
 
-# Quantize (when the mode carries an "_and_<quant>" suffix) and then place the pipeline.
-# The order lives inside the helper: quantization has to happen before the offload hooks
-# are installed.
-apply_gpu_memory_mode(pipeline, GPU_memory_mode, device, weight_dtype, exclude_module_name=['img_in', 'txt_in', 'timestep'])
+if GPU_memory_mode == "sequential_cpu_offload":
+    pipeline.enable_sequential_cpu_offload(device=device)
+elif GPU_memory_mode == "model_group_offload":
+    register_auto_device_hook(pipeline.transformer)
+    safe_enable_group_offload(pipeline, onload_device=device, offload_device="cpu", offload_type="leaf_level", use_stream=True)
+elif GPU_memory_mode == "model_cpu_offload_and_qfloat8":
+    convert_model_weight_to_float8(transformer, exclude_module_name=["img_in", "txt_in", "timestep"], device=device)
+    convert_weight_dtype_wrapper(transformer, weight_dtype)
+    pipeline.enable_model_cpu_offload(device=device)
+elif GPU_memory_mode == "model_cpu_offload":
+    pipeline.enable_model_cpu_offload(device=device)
+elif GPU_memory_mode == "model_full_load_and_qfloat8":
+    convert_model_weight_to_float8(transformer, exclude_module_name=["img_in", "txt_in", "timestep"], device=device)
+    convert_weight_dtype_wrapper(transformer, weight_dtype)
+    pipeline.to(device=device)
+else:
+    pipeline.to(device=device)
 
 coefficients = get_teacache_coefficients(model_name) if enable_teacache else None
 if coefficients is not None:

examples/qwenimage/predict_t2i_edit_plus.py

@@ -3,6 +3,7 @@ import sys
 
 import torch
 from diffusers import FlowMatchEulerDiscreteScheduler
+from PIL import Image
 
 current_file_path = os.path.abspath(__file__)
 project_roots = [os.path.dirname(current_file_path), os.path.dirname(os.path.dirname(current_file_path)), os.path.dirname(os.path.dirname(os.path.dirname(current_file_path)))]
@@ -16,12 +17,16 @@ from videox_fun.models import (AutoencoderKLQwenImage,
                                QwenImageTransformer2DModel)
 from videox_fun.models.cache_utils import get_teacache_coefficients
 from videox_fun.pipeline import QwenImageEditPlusPipeline
-from videox_fun.utils import (FlowDPMSolverMultistepScheduler,
-                              FlowUniPCMultistepScheduler,
-                              apply_gpu_memory_mode, get_image, merge_lora,
-                              unmerge_lora)
+from videox_fun.utils import (register_auto_device_hook,
+                              safe_enable_group_offload)
+from videox_fun.utils.fm_solvers import FlowDPMSolverMultistepScheduler
+from videox_fun.utils.fm_solvers_unipc import FlowUniPCMultistepScheduler
+from videox_fun.utils.fp8_optimization import (convert_model_weight_to_float8,
+                                               convert_weight_dtype_wrapper)
+from videox_fun.utils.lora_utils import merge_lora, unmerge_lora
+from videox_fun.utils.utils import get_image
 
-# GPU memory mode, which can be chosen in [model_full_load, model_full_load_and_qfloat8, model_cpu_offload, model_cpu_offload_and_qfloat8, model_group_offload, sequential_cpu_offload].
+# GPU memory mode, which can be chosen in [model_full_load, model_full_load_and_qfloat8, model_cpu_offload, model_cpu_offload_and_qfloat8, sequential_cpu_offload].
 # model_full_load means that the entire model will be moved to the GPU.
 # 
 # model_full_load_and_qfloat8 means that the entire model will be moved to the GPU,
@@ -81,7 +86,7 @@ lora_path           = None
 sample_size         = [1344, 768]
 
 # Use torch.float16 if GPU does not support torch.bfloat16
-# Some graphics cards, such as v100, 2080ti, do not support torch.bfloat16
+# ome graphics cards, such as v100, 2080ti, do not support torch.bfloat16
 weight_dtype        = torch.bfloat16
 image               = ["asset/8.png", "asset/ref_1.png"]
 # Please use as detailed a prompt as possible to describe the object that needs to be generated.
@@ -105,7 +110,7 @@ transformer = QwenImageTransformer2DModel.from_pretrained(
 if transformer_path is not None:
     print(f"From checkpoint: {transformer_path}")
     if transformer_path.endswith("safetensors"):
-        from safetensors.torch import load_file
+        from safetensors.torch import load_file, safe_open
         state_dict = load_file(transformer_path)
     else:
         state_dict = torch.load(transformer_path, map_location="cpu")
@@ -123,7 +128,7 @@ vae = AutoencoderKLQwenImage.from_pretrained(
 if vae_path is not None:
     print(f"From checkpoint: {vae_path}")
     if vae_path.endswith("safetensors"):
-        from safetensors.torch import load_file
+        from safetensors.torch import load_file, safe_open
         state_dict = load_file(vae_path)
     else:
         state_dict = torch.load(vae_path, map_location="cpu")
@@ -147,7 +152,7 @@ processor = Qwen2VLProcessor.from_pretrained(
 )
 
 # Get Scheduler
-Chosen_Scheduler = {
+Chosen_Scheduler = scheduler_dict = {
     "Flow": FlowMatchEulerDiscreteScheduler,
     "Flow_Unipc": FlowUniPCMultistepScheduler,
     "Flow_DPM++": FlowDPMSolverMultistepScheduler,
@@ -174,6 +179,9 @@ if ulysses_degree > 1 or ring_degree > 1:
         pipeline.transformer = shard_fn(pipeline.transformer)
         print("Add FSDP DIT")
     if fsdp_text_encoder:
+        from functools import partial
+
+        from videox_fun.dist import set_multi_gpus_devices, shard_model
         shard_fn = partial(shard_model, device_id=device, param_dtype=weight_dtype, module_to_wrapper=text_encoder.language_model.layers)
         pipeline.text_encoder = shard_fn(pipeline.text_encoder)
         print("Add FSDP TEXT ENCODER")
@@ -183,10 +191,23 @@ if compile_dit:
         pipeline.transformer.transformer_blocks[i] = torch.compile(pipeline.transformer.transformer_blocks[i])
     print("Add Compile")
 
-# Quantize (when the mode carries an "_and_<quant>" suffix) and then place the pipeline.
-# The order lives inside the helper: quantization has to happen before the offload hooks
-# are installed.
-apply_gpu_memory_mode(pipeline, GPU_memory_mode, device, weight_dtype, exclude_module_name=['img_in', 'txt_in', 'timestep'])
+if GPU_memory_mode == "sequential_cpu_offload":
+    pipeline.enable_sequential_cpu_offload(device=device)
+elif GPU_memory_mode == "model_group_offload":
+    register_auto_device_hook(pipeline.transformer)
+    safe_enable_group_offload(pipeline, onload_device=device, offload_device="cpu", offload_type="leaf_level", use_stream=True)
+elif GPU_memory_mode == "model_cpu_offload_and_qfloat8":
+    convert_model_weight_to_float8(transformer, exclude_module_name=["img_in", "txt_in", "timestep"], device=device)
+    convert_weight_dtype_wrapper(transformer, weight_dtype)
+    pipeline.enable_model_cpu_offload(device=device)
+elif GPU_memory_mode == "model_cpu_offload":
+    pipeline.enable_model_cpu_offload(device=device)
+elif GPU_memory_mode == "model_full_load_and_qfloat8":
+    convert_model_weight_to_float8(transformer, exclude_module_name=["img_in", "txt_in", "timestep"], device=device)
+    convert_weight_dtype_wrapper(transformer, weight_dtype)
+    pipeline.to(device=device)
+else:
+    pipeline.to(device=device)
 
 coefficients = get_teacache_coefficients(model_name) if enable_teacache else None
 if coefficients is not None:

examples/qwenimage21/predict_t2i.py

@@ -2,6 +2,7 @@ import os
 import sys
 
 import torch
+
 from diffusers import FlowMatchEulerDiscreteScheduler
 
 current_file_path = os.path.abspath(__file__)
@@ -14,7 +15,11 @@ from videox_fun.models import (AutoencoderKLQwenImage21,
                                Qwen3VLForConditionalGeneration,
                                Qwen3VLProcessor, QwenImage21Transformer2DModel)
 from videox_fun.pipeline import QwenImage21Pipeline
-from videox_fun.utils import apply_gpu_memory_mode, merge_lora, unmerge_lora
+from videox_fun.utils import (register_auto_device_hook,
+                              safe_enable_group_offload)
+from videox_fun.utils.fp8_optimization import (convert_model_weight_to_float8,
+                                               convert_weight_dtype_wrapper)
+from videox_fun.utils.lora_utils import merge_lora, unmerge_lora
 
 # GPU memory mode, which can be chosen in [model_full_load, model_full_load_and_qfloat8, model_cpu_offload, model_cpu_offload_and_qfloat8, model_group_offload, sequential_cpu_offload].
 # model_full_load means that the entire model will be moved to the GPU.
@@ -89,7 +94,7 @@ transformer = QwenImage21Transformer2DModel.from_pretrained(
 if transformer_path is not None:
     print(f"From checkpoint: {transformer_path}")
     if transformer_path.endswith("safetensors"):
-        from safetensors.torch import load_file
+        from safetensors.torch import load_file, safe_open
         state_dict = load_file(transformer_path)
     else:
         state_dict = torch.load(transformer_path, map_location="cpu")
@@ -107,7 +112,7 @@ vae = AutoencoderKLQwenImage21.from_pretrained(
 if vae_path is not None:
     print(f"From checkpoint: {vae_path}")
     if vae_path.endswith("safetensors"):
-        from safetensors.torch import load_file
+        from safetensors.torch import load_file, safe_open
         state_dict = load_file(vae_path)
     else:
         state_dict = torch.load(vae_path, map_location="cpu")
@@ -150,6 +155,8 @@ if ulysses_degree > 1 or ring_degree > 1:
         pipeline.transformer = shard_fn(pipeline.transformer)
         print("Add FSDP DIT")
     if fsdp_text_encoder:
+        from functools import partial
+        from videox_fun.dist import set_multi_gpus_devices, shard_model
         shard_fn = partial(shard_model, device_id=device, param_dtype=weight_dtype, module_to_wrapper=text_encoder.model.language_model.layers)
         pipeline.text_encoder = shard_fn(pipeline.text_encoder)
         print("Add FSDP TEXT ENCODER")
@@ -159,10 +166,23 @@ if compile_dit:
         pipeline.transformer.transformer_blocks[i] = torch.compile(pipeline.transformer.transformer_blocks[i])
     print("Add Compile")
 
-# Quantize (when the mode carries an "_and_<quant>" suffix) and then place the pipeline.
-# The order lives inside the helper: quantization has to happen before the offload hooks
-# are installed.
-apply_gpu_memory_mode(pipeline, GPU_memory_mode, device, weight_dtype, exclude_module_name=['img_in', 'txt_in', 'time_text_embed', 'modulation'])
+if GPU_memory_mode == "sequential_cpu_offload":
+    pipeline.enable_sequential_cpu_offload(device=device)
+elif GPU_memory_mode == "model_group_offload":
+    register_auto_device_hook(pipeline.transformer)
+    safe_enable_group_offload(pipeline, onload_device=device, offload_device="cpu", offload_type="leaf_level", use_stream=True)
+elif GPU_memory_mode == "model_cpu_offload_and_qfloat8":
+    convert_model_weight_to_float8(transformer, exclude_module_name=["img_in", "txt_in", "time_text_embed", "modulation"], device=device)
+    convert_weight_dtype_wrapper(transformer, weight_dtype)
+    pipeline.enable_model_cpu_offload(device=device)
+elif GPU_memory_mode == "model_cpu_offload":
+    pipeline.enable_model_cpu_offload(device=device)
+elif GPU_memory_mode == "model_full_load_and_qfloat8":
+    convert_model_weight_to_float8(transformer, exclude_module_name=["img_in", "txt_in", "time_text_embed", "modulation"], device=device)
+    convert_weight_dtype_wrapper(transformer, weight_dtype)
+    pipeline.to(device=device)
+else:
+    pipeline.to(device=device)
 
 for prompt in prompts:
     generator = torch.Generator(device=device).manual_seed(seed)

examples/qwenimage21_fun/predict_i2i_inpaint.py

@@ -2,6 +2,7 @@ import os
 import sys
 
 import torch
+
 from diffusers import FlowMatchEulerDiscreteScheduler
 from omegaconf import OmegaConf
 
@@ -16,8 +17,12 @@ from videox_fun.models import (AutoencoderKLQwenImage21,
                                Qwen3VLProcessor,
                                QwenImage21ControlTransformer2DModel)
 from videox_fun.pipeline import QwenImage21ControlPipeline
-from videox_fun.utils import (apply_gpu_memory_mode, get_image_latent,
-                              merge_lora, unmerge_lora)
+from videox_fun.utils import (register_auto_device_hook,
+                              safe_enable_group_offload)
+from videox_fun.utils.fp8_optimization import (convert_model_weight_to_float8,
+                                               convert_weight_dtype_wrapper)
+from videox_fun.utils.lora_utils import merge_lora, unmerge_lora
+from videox_fun.utils.utils import get_image_latent
 
 # GPU memory mode, which can be chosen in [model_full_load, model_full_load_and_qfloat8, model_cpu_offload, model_cpu_offload_and_qfloat8, model_group_offload, sequential_cpu_offload].
 # model_full_load means that the entire model will be moved to the GPU.
@@ -102,7 +107,7 @@ transformer = QwenImage21ControlTransformer2DModel.from_pretrained(
 if transformer_path is not None:
     print(f"From checkpoint: {transformer_path}")
     if transformer_path.endswith("safetensors"):
-        from safetensors.torch import load_file
+        from safetensors.torch import load_file, safe_open
         state_dict = load_file(transformer_path)
     else:
         state_dict = torch.load(transformer_path, map_location="cpu")
@@ -120,7 +125,7 @@ vae = AutoencoderKLQwenImage21.from_pretrained(
 if vae_path is not None:
     print(f"From checkpoint: {vae_path}")
     if vae_path.endswith("safetensors"):
-        from safetensors.torch import load_file
+        from safetensors.torch import load_file, safe_open
         state_dict = load_file(vae_path)
     else:
         state_dict = torch.load(vae_path, map_location="cpu")
@@ -172,10 +177,23 @@ if compile_dit:
         pipeline.transformer.transformer_blocks[i] = torch.compile(pipeline.transformer.transformer_blocks[i])
     print("Add Compile")
 
-# Quantize (when the mode carries an "_and_<quant>" suffix) and then place the pipeline.
-# The order lives inside the helper: quantization has to happen before the offload hooks
-# are installed.
-apply_gpu_memory_mode(pipeline, GPU_memory_mode, device, weight_dtype, exclude_module_name=['img_in', 'txt_in', 'time_text_embed', 'modulation'])
+if GPU_memory_mode == "sequential_cpu_offload":
+    pipeline.enable_sequential_cpu_offload(device=device)
+elif GPU_memory_mode == "model_group_offload":
+    register_auto_device_hook(pipeline.transformer)
+    safe_enable_group_offload(pipeline, onload_device=device, offload_device="cpu", offload_type="leaf_level", use_stream=True)
+elif GPU_memory_mode == "model_cpu_offload_and_qfloat8":
+    convert_model_weight_to_float8(transformer, exclude_module_name=["img_in", "txt_in", "time_text_embed", "modulation"], device=device)
+    convert_weight_dtype_wrapper(transformer, weight_dtype)
+    pipeline.enable_model_cpu_offload(device=device)
+elif GPU_memory_mode == "model_cpu_offload":
+    pipeline.enable_model_cpu_offload(device=device)
+elif GPU_memory_mode == "model_full_load_and_qfloat8":
+    convert_model_weight_to_float8(transformer, exclude_module_name=["img_in", "txt_in", "time_text_embed", "modulation"], device=device)
+    convert_weight_dtype_wrapper(transformer, weight_dtype)
+    pipeline.to(device=device)
+else:
+    pipeline.to(device=device)
 
 generator = torch.Generator(device=device).manual_seed(seed)
 

examples/qwenimage21_fun/predict_t2i_control.py

@@ -2,6 +2,7 @@ import os
 import sys
 
 import torch
+
 from diffusers import FlowMatchEulerDiscreteScheduler
 from omegaconf import OmegaConf
 
@@ -16,8 +17,12 @@ from videox_fun.models import (AutoencoderKLQwenImage21,
                                Qwen3VLProcessor,
                                QwenImage21ControlTransformer2DModel)
 from videox_fun.pipeline import QwenImage21ControlPipeline
-from videox_fun.utils import (apply_gpu_memory_mode, get_image_latent,
-                              merge_lora, unmerge_lora)
+from videox_fun.utils import (register_auto_device_hook,
+                              safe_enable_group_offload)
+from videox_fun.utils.fp8_optimization import (convert_model_weight_to_float8,
+                                               convert_weight_dtype_wrapper)
+from videox_fun.utils.lora_utils import merge_lora, unmerge_lora
+from videox_fun.utils.utils import get_image_latent
 
 # GPU memory mode, which can be chosen in [model_full_load, model_full_load_and_qfloat8, model_cpu_offload, model_cpu_offload_and_qfloat8, model_group_offload, sequential_cpu_offload].
 # model_full_load means that the entire model will be moved to the GPU.
@@ -99,7 +104,7 @@ transformer = QwenImage21ControlTransformer2DModel.from_pretrained(
 if transformer_path is not None:
     print(f"From checkpoint: {transformer_path}")
     if transformer_path.endswith("safetensors"):
-        from safetensors.torch import load_file
+        from safetensors.torch import load_file, safe_open
         state_dict = load_file(transformer_path)
     else:
         state_dict = torch.load(transformer_path, map_location="cpu")
@@ -117,7 +122,7 @@ vae = AutoencoderKLQwenImage21.from_pretrained(
 if vae_path is not None:
     print(f"From checkpoint: {vae_path}")
     if vae_path.endswith("safetensors"):
-        from safetensors.torch import load_file
+        from safetensors.torch import load_file, safe_open
         state_dict = load_file(vae_path)
     else:
         state_dict = torch.load(vae_path, map_location="cpu")
@@ -169,10 +174,23 @@ if compile_dit:
         pipeline.transformer.transformer_blocks[i] = torch.compile(pipeline.transformer.transformer_blocks[i])
     print("Add Compile")
 
-# Quantize (when the mode carries an "_and_<quant>" suffix) and then place the pipeline.
-# The order lives inside the helper: quantization has to happen before the offload hooks
-# are installed.
-apply_gpu_memory_mode(pipeline, GPU_memory_mode, device, weight_dtype, exclude_module_name=['img_in', 'txt_in', 'time_text_embed', 'modulation'])
+if GPU_memory_mode == "sequential_cpu_offload":
+    pipeline.enable_sequential_cpu_offload(device=device)
+elif GPU_memory_mode == "model_group_offload":
+    register_auto_device_hook(pipeline.transformer)
+    safe_enable_group_offload(pipeline, onload_device=device, offload_device="cpu", offload_type="leaf_level", use_stream=True)
+elif GPU_memory_mode == "model_cpu_offload_and_qfloat8":
+    convert_model_weight_to_float8(transformer, exclude_module_name=["img_in", "txt_in", "time_text_embed", "modulation"], device=device)
+    convert_weight_dtype_wrapper(transformer, weight_dtype)
+    pipeline.enable_model_cpu_offload(device=device)
+elif GPU_memory_mode == "model_cpu_offload":
+    pipeline.enable_model_cpu_offload(device=device)
+elif GPU_memory_mode == "model_full_load_and_qfloat8":
+    convert_model_weight_to_float8(transformer, exclude_module_name=["img_in", "txt_in", "time_text_embed", "modulation"], device=device)
+    convert_weight_dtype_wrapper(transformer, weight_dtype)
+    pipeline.to(device=device)
+else:
+    pipeline.to(device=device)
 
 generator = torch.Generator(device=device).manual_seed(seed)
 

examples/qwenimage_fun/predict_i2i_inpaint.py

@@ -17,12 +17,16 @@ from videox_fun.models import (AutoencoderKLQwenImage,
                                QwenImageControlTransformer2DModel)
 from videox_fun.models.cache_utils import get_teacache_coefficients
 from videox_fun.pipeline import QwenImageControlPipeline
-from videox_fun.utils import (FlowDPMSolverMultistepScheduler,
-                              FlowUniPCMultistepScheduler,
-                              apply_gpu_memory_mode, get_image_latent,
-                              merge_lora, unmerge_lora)
+from videox_fun.utils import (register_auto_device_hook,
+                              safe_enable_group_offload)
+from videox_fun.utils.fm_solvers import FlowDPMSolverMultistepScheduler
+from videox_fun.utils.fm_solvers_unipc import FlowUniPCMultistepScheduler
+from videox_fun.utils.fp8_optimization import (convert_model_weight_to_float8,
+                                               convert_weight_dtype_wrapper)
+from videox_fun.utils.lora_utils import merge_lora, unmerge_lora
+from videox_fun.utils.utils import get_image_latent, save_videos_grid
 
-# GPU memory mode, which can be chosen in [model_full_load, model_full_load_and_qfloat8, model_cpu_offload, model_cpu_offload_and_qfloat8, model_group_offload, sequential_cpu_offload].
+# GPU memory mode, which can be chosen in [model_full_load, model_full_load_and_qfloat8, model_cpu_offload, model_cpu_offload_and_qfloat8, sequential_cpu_offload].
 # model_full_load means that the entire model will be moved to the GPU.
 # 
 # model_full_load_and_qfloat8 means that the entire model will be moved to the GPU,
@@ -84,7 +88,7 @@ lora_path           = None
 sample_size         = [1728, 992]
 
 # Use torch.float16 if GPU does not support torch.bfloat16
-# Some graphics cards, such as v100, 2080ti, do not support torch.bfloat16
+# ome graphics cards, such as v100, 2080ti, do not support torch.bfloat16
 weight_dtype        = torch.bfloat16
 control_image       = "asset/pose.jpg"
 inpaint_image       = "asset/8.png"
@@ -151,7 +155,7 @@ text_encoder = Qwen2_5_VLForConditionalGeneration.from_pretrained(
 )
 
 # Get Scheduler
-Chosen_Scheduler = {
+Chosen_Scheduler = scheduler_dict = {
     "Flow": FlowMatchEulerDiscreteScheduler,
     "Flow_Unipc": FlowUniPCMultistepScheduler,
     "Flow_DPM++": FlowDPMSolverMultistepScheduler,
@@ -177,6 +181,9 @@ if ulysses_degree > 1 or ring_degree > 1:
         pipeline.transformer = shard_fn(pipeline.transformer)
         print("Add FSDP DIT")
     if fsdp_text_encoder:
+        from functools import partial
+
+        from videox_fun.dist import set_multi_gpus_devices, shard_model
         shard_fn = partial(shard_model, device_id=device, param_dtype=weight_dtype, module_to_wrapper=text_encoder.language_model.layers)
         pipeline.text_encoder = shard_fn(pipeline.text_encoder)
         print("Add FSDP TEXT ENCODER")
@@ -186,10 +193,23 @@ if compile_dit:
         pipeline.transformer.transformer_blocks[i] = torch.compile(pipeline.transformer.transformer_blocks[i])
     print("Add Compile")
 
-# Quantize (when the mode carries an "_and_<quant>" suffix) and then place the pipeline.
-# The order lives inside the helper: quantization has to happen before the offload hooks
-# are installed.
-apply_gpu_memory_mode(pipeline, GPU_memory_mode, device, weight_dtype, exclude_module_name=['img_in', 'txt_in', 'timestep'])
+if GPU_memory_mode == "sequential_cpu_offload":
+    pipeline.enable_sequential_cpu_offload(device=device)
+elif GPU_memory_mode == "model_group_offload":
+    register_auto_device_hook(pipeline.transformer)
+    safe_enable_group_offload(pipeline, onload_device=device, offload_device="cpu", offload_type="leaf_level", use_stream=True)
+elif GPU_memory_mode == "model_cpu_offload_and_qfloat8":
+    convert_model_weight_to_float8(transformer, exclude_module_name=["img_in", "txt_in", "timestep"], device=device)
+    convert_weight_dtype_wrapper(transformer, weight_dtype)
+    pipeline.enable_model_cpu_offload(device=device)
+elif GPU_memory_mode == "model_cpu_offload":
+    pipeline.enable_model_cpu_offload(device=device)
+elif GPU_memory_mode == "model_full_load_and_qfloat8":
+    convert_model_weight_to_float8(transformer, exclude_module_name=["img_in", "txt_in", "timestep"], device=device)
+    convert_weight_dtype_wrapper(transformer, weight_dtype)
+    pipeline.to(device=device)
+else:
+    pipeline.to(device=device)
 
 coefficients = get_teacache_coefficients(model_name) if enable_teacache else None
 if coefficients is not None:

examples/qwenimage_fun/predict_t2i_control.py

@@ -17,12 +17,16 @@ from videox_fun.models import (AutoencoderKLQwenImage,
                                QwenImageControlTransformer2DModel)
 from videox_fun.models.cache_utils import get_teacache_coefficients
 from videox_fun.pipeline import QwenImageControlPipeline
-from videox_fun.utils import (FlowDPMSolverMultistepScheduler,
-                              FlowUniPCMultistepScheduler,
-                              apply_gpu_memory_mode, get_image_latent,
-                              merge_lora, unmerge_lora)
+from videox_fun.utils import (register_auto_device_hook,
+                              safe_enable_group_offload)
+from videox_fun.utils.fm_solvers import FlowDPMSolverMultistepScheduler
+from videox_fun.utils.fm_solvers_unipc import FlowUniPCMultistepScheduler
+from videox_fun.utils.fp8_optimization import (convert_model_weight_to_float8,
+                                               convert_weight_dtype_wrapper)
+from videox_fun.utils.lora_utils import merge_lora, unmerge_lora
+from videox_fun.utils.utils import get_image_latent, save_videos_grid
 
-# GPU memory mode, which can be chosen in [model_full_load, model_full_load_and_qfloat8, model_cpu_offload, model_cpu_offload_and_qfloat8, model_group_offload, sequential_cpu_offload].
+# GPU memory mode, which can be chosen in [model_full_load, model_full_load_and_qfloat8, model_cpu_offload, model_cpu_offload_and_qfloat8, sequential_cpu_offload].
 # model_full_load means that the entire model will be moved to the GPU.
 # 
 # model_full_load_and_qfloat8 means that the entire model will be moved to the GPU,
@@ -84,7 +88,7 @@ lora_path           = None
 sample_size         = [1728, 992]
 
 # Use torch.float16 if GPU does not support torch.bfloat16
-# Some graphics cards, such as v100, 2080ti, do not support torch.bfloat16
+# ome graphics cards, such as v100, 2080ti, do not support torch.bfloat16
 weight_dtype        = torch.bfloat16
 control_image       = "asset/pose.jpg"
 inpaint_image       = None
@@ -151,7 +155,7 @@ text_encoder = Qwen2_5_VLForConditionalGeneration.from_pretrained(
 )
 
 # Get Scheduler
-Chosen_Scheduler = {
+Chosen_Scheduler = scheduler_dict = {
     "Flow": FlowMatchEulerDiscreteScheduler,
     "Flow_Unipc": FlowUniPCMultistepScheduler,
     "Flow_DPM++": FlowDPMSolverMultistepScheduler,
@@ -177,6 +181,9 @@ if ulysses_degree > 1 or ring_degree > 1:
         pipeline.transformer = shard_fn(pipeline.transformer)
         print("Add FSDP DIT")
     if fsdp_text_encoder:
+        from functools import partial
+
+        from videox_fun.dist import set_multi_gpus_devices, shard_model
         shard_fn = partial(shard_model, device_id=device, param_dtype=weight_dtype, module_to_wrapper=text_encoder.language_model.layers)
         pipeline.text_encoder = shard_fn(pipeline.text_encoder)
         print("Add FSDP TEXT ENCODER")
@@ -186,10 +193,23 @@ if compile_dit:
         pipeline.transformer.transformer_blocks[i] = torch.compile(pipeline.transformer.transformer_blocks[i])
     print("Add Compile")
 
-# Quantize (when the mode carries an "_and_<quant>" suffix) and then place the pipeline.
-# The order lives inside the helper: quantization has to happen before the offload hooks
-# are installed.
-apply_gpu_memory_mode(pipeline, GPU_memory_mode, device, weight_dtype, exclude_module_name=['img_in', 'txt_in', 'timestep'])
+if GPU_memory_mode == "sequential_cpu_offload":
+    pipeline.enable_sequential_cpu_offload(device=device)
+elif GPU_memory_mode == "model_group_offload":
+    register_auto_device_hook(pipeline.transformer)
+    safe_enable_group_offload(pipeline, onload_device=device, offload_device="cpu", offload_type="leaf_level", use_stream=True)
+elif GPU_memory_mode == "model_cpu_offload_and_qfloat8":
+    convert_model_weight_to_float8(transformer, exclude_module_name=["img_in", "txt_in", "timestep"], device=device)
+    convert_weight_dtype_wrapper(transformer, weight_dtype)
+    pipeline.enable_model_cpu_offload(device=device)
+elif GPU_memory_mode == "model_cpu_offload":
+    pipeline.enable_model_cpu_offload(device=device)
+elif GPU_memory_mode == "model_full_load_and_qfloat8":
+    convert_model_weight_to_float8(transformer, exclude_module_name=["img_in", "txt_in", "timestep"], device=device)
+    convert_weight_dtype_wrapper(transformer, weight_dtype)
+    pipeline.to(device=device)
+else:
+    pipeline.to(device=device)
 
 coefficients = get_teacache_coefficients(model_name) if enable_teacache else None
 if coefficients is not None:

examples/qwenimage_instantx/predict_t2i_control.py

@@ -2,7 +2,9 @@ import os
 import sys
 
 import torch
-from diffusers import FlowMatchEulerDiscreteScheduler
+
+from omegaconf import OmegaConf
+from diffusers import (FlowMatchEulerDiscreteScheduler)
 
 current_file_path = os.path.abspath(__file__)
 project_roots = [os.path.dirname(current_file_path), os.path.dirname(os.path.dirname(current_file_path)), os.path.dirname(os.path.dirname(os.path.dirname(current_file_path)))]
@@ -10,19 +12,23 @@ for project_root in project_roots:
     sys.path.insert(0, project_root) if project_root not in sys.path else None
 
 from videox_fun.dist import set_multi_gpus_devices, shard_model
-from videox_fun.models import (AutoencoderKLQwenImage,
+from videox_fun.models import (AutoencoderKLQwenImage, QwenImageInstantXControlNetModel, 
                                Qwen2_5_VLForConditionalGeneration,
-                               Qwen2Tokenizer,
-                               QwenImageInstantXControlNetModel,
-                               QwenImageTransformer2DModel)
+                               Qwen2Tokenizer, QwenImageTransformer2DModel)
 from videox_fun.models.cache_utils import get_teacache_coefficients
 from videox_fun.pipeline import QwenImageControlNetPipeline
-from videox_fun.utils import (FlowDPMSolverMultistepScheduler,
-                              FlowUniPCMultistepScheduler,
-                              apply_gpu_memory_mode, get_image, merge_lora,
-                              unmerge_lora)
+from videox_fun.utils import (register_auto_device_hook,
+                              safe_enable_group_offload)
+from videox_fun.utils.fm_solvers import FlowDPMSolverMultistepScheduler
+from videox_fun.utils.fm_solvers_unipc import FlowUniPCMultistepScheduler
+from videox_fun.utils.fp8_optimization import (convert_model_weight_to_float8,
+                                               convert_weight_dtype_wrapper)
+from videox_fun.utils.lora_utils import merge_lora, unmerge_lora
+from videox_fun.utils.utils import (filter_kwargs, get_image_to_video_latent, get_image,
+                                    get_video_to_video_latent,
+                                    save_videos_grid)
 
-# GPU memory mode, which can be chosen in [model_full_load, model_full_load_and_qfloat8, model_cpu_offload, model_cpu_offload_and_qfloat8, model_group_offload, sequential_cpu_offload].
+# GPU memory mode, which can be chosen in [model_full_load, model_full_load_and_qfloat8, model_cpu_offload, model_cpu_offload_and_qfloat8, sequential_cpu_offload].
 # model_full_load means that the entire model will be moved to the GPU.
 # 
 # model_full_load_and_qfloat8 means that the entire model will be moved to the GPU,
@@ -85,7 +91,7 @@ lora_path           = None
 sample_size         = [1728, 992]
 
 # Use torch.float16 if GPU does not support torch.bfloat16
-# Some graphics cards, such as v100, 2080ti, do not support torch.bfloat16
+# ome graphics cards, such as v100, 2080ti, do not support torch.bfloat16
 weight_dtype        = torch.bfloat16
 control_image       = "asset/pose.jpg"
 controlnet_conditioning_scale = 0.80
@@ -166,7 +172,7 @@ text_encoder = Qwen2_5_VLForConditionalGeneration.from_pretrained(
 )
 
 # Get Scheduler
-Chosen_Scheduler = {
+Chosen_Scheduler = scheduler_dict = {
     "Flow": FlowMatchEulerDiscreteScheduler,
     "Flow_Unipc": FlowUniPCMultistepScheduler,
     "Flow_DPM++": FlowDPMSolverMultistepScheduler,
@@ -193,6 +199,8 @@ if ulysses_degree > 1 or ring_degree > 1:
         pipeline.transformer = shard_fn(pipeline.transformer)
         print("Add FSDP DIT")
     if fsdp_text_encoder:
+        from functools import partial
+        from videox_fun.dist import set_multi_gpus_devices, shard_model
         shard_fn = partial(shard_model, device_id=device, param_dtype=weight_dtype, module_to_wrapper=text_encoder.language_model.layers)
         pipeline.text_encoder = shard_fn(pipeline.text_encoder)
         print("Add FSDP TEXT ENCODER")
@@ -202,10 +210,23 @@ if compile_dit:
         pipeline.transformer.transformer_blocks[i] = torch.compile(pipeline.transformer.transformer_blocks[i])
     print("Add Compile")
 
-# Quantize (when the mode carries an "_and_<quant>" suffix) and then place the pipeline.
-# The order lives inside the helper: quantization has to happen before the offload hooks
-# are installed.
-apply_gpu_memory_mode(pipeline, GPU_memory_mode, device, weight_dtype, exclude_module_name=['img_in', 'txt_in', 'timestep'])
+if GPU_memory_mode == "sequential_cpu_offload":
+    pipeline.enable_sequential_cpu_offload(device=device)
+elif GPU_memory_mode == "model_group_offload":
+    register_auto_device_hook(pipeline.transformer)
+    safe_enable_group_offload(pipeline, onload_device=device, offload_device="cpu", offload_type="leaf_level", use_stream=True)
+elif GPU_memory_mode == "model_cpu_offload_and_qfloat8":
+    convert_model_weight_to_float8(transformer, exclude_module_name=["img_in", "txt_in", "timestep"], device=device)
+    convert_weight_dtype_wrapper(transformer, weight_dtype)
+    pipeline.enable_model_cpu_offload(device=device)
+elif GPU_memory_mode == "model_cpu_offload":
+    pipeline.enable_model_cpu_offload(device=device)
+elif GPU_memory_mode == "model_full_load_and_qfloat8":
+    convert_model_weight_to_float8(transformer, exclude_module_name=["img_in", "txt_in", "timestep"], device=device)
+    convert_weight_dtype_wrapper(transformer, weight_dtype)
+    pipeline.to(device=device)
+else:
+    pipeline.to(device=device)
 
 coefficients = get_teacache_coefficients(model_name) if enable_teacache else None
 if coefficients is not None:

examples/taomate_h3/predict_t2audio.py

@@ -29,7 +29,8 @@ from videox_fun.pipeline.pipeline_taomate_h3 import (
     TAOMATE_H3_VIDEO_SIGMA_SHIFT, taomate_h3_canonical_continuation_plan,
     taomate_h3_direct_5s_plan, taomate_h3_select_time_shift_sigmas,
     taomate_h3_teacher_geometry)
-from videox_fun.utils import MiniMaxH3Scheduler, apply_gpu_memory_mode
+from videox_fun.utils import (MiniMaxH3Scheduler, register_auto_device_hook,
+                              safe_enable_group_offload)
 
 # GPU memory mode, which can be chosen in [model_full_load, model_cpu_offload, model_group_offload, sequential_cpu_offload].
 # model_full_load means that the entire model will be moved to the GPU.
@@ -168,10 +169,20 @@ if compile_dit:
         pipeline.transformer.transformer_blocks[i] = torch.compile(pipeline.transformer.transformer_blocks[i])
     print("Add Compile")
 
-# Quantize (when the mode carries an "_and_<quant>" suffix) and then place the pipeline.
-# The order lives inside the helper: quantization has to happen before the offload hooks
-# are installed.
-apply_gpu_memory_mode(pipeline, GPU_memory_mode, device, weight_dtype, exclude_module_name=[], strict=True)
+if GPU_memory_mode == "sequential_cpu_offload":
+    pipeline.enable_sequential_cpu_offload(device=device)
+elif GPU_memory_mode == "model_group_offload":
+    register_auto_device_hook(pipeline.transformer)
+    safe_enable_group_offload(pipeline, onload_device=device, offload_device="cpu", offload_type="leaf_level", use_stream=True)
+elif GPU_memory_mode == "model_cpu_offload":
+    pipeline.enable_model_cpu_offload(device=device)
+elif GPU_memory_mode == "model_full_load":
+    pipeline.to(device=device)
+else:
+    raise ValueError(
+        f"`GPU_memory_mode` must be one of ['model_full_load', 'model_cpu_offload', 'model_group_offload', "
+        f"'sequential_cpu_offload'], got {GPU_memory_mode}."
+    )
 
 
 def official_audio_noise(*, video_latent_t, video_latent_h, video_latent_w, audio_latent_t, seed):
@@ -267,6 +278,7 @@ def audio_only_step_timesteps(t_video, t_audio, *, has_reference):
 
     return unique_timesteps, expand
 
+
 _OFFLOAD_MODES = ("model_cpu_offload", "model_group_offload", "sequential_cpu_offload")
 
 
@@ -540,6 +552,7 @@ def write_teacher_artifact(output_dir, *, pipeline, request_records, request_cou
         json.dump(completion, handle, ensure_ascii=False, indent=2)
         handle.write("\n")
 
+
 # One continuous audio timeline 5 seconds at a time: every request after the first denoises the previous
 # request's clean tail as a frozen reference, and the captured 3/6/9 rows are the Base10 teacher artifact.
 audio_rows, request_records = generate_audio_track(

examples/taomate_h3/predict_t2av_streaming.py

@@ -15,9 +15,12 @@ from videox_fun.models import (AutoencoderKLMiniMaxH3,
                                Qwen3VLForConditionalGeneration,
                                Qwen3VLProcessor)
 from videox_fun.pipeline import MiniMaxH3StreamingPipeline
-from videox_fun.utils import (MiniMaxH3Scheduler, apply_gpu_memory_mode,
-                              convert_model_weight_to_float8, merge_lora,
-                              save_videos_with_audio_grid, unmerge_lora)
+from videox_fun.utils import (MiniMaxH3Scheduler, register_auto_device_hook,
+                              safe_enable_group_offload)
+from videox_fun.utils.fp8_optimization import (convert_model_weight_to_float8,
+                                               convert_weight_dtype_wrapper)
+from videox_fun.utils.lora_utils import merge_lora, unmerge_lora
+from videox_fun.utils.utils import save_videos_with_audio_grid
 
 # GPU memory mode, which can be chosen in [model_full_load, model_full_load_and_qfloat8, model_cpu_offload, model_cpu_offload_and_qfloat8, model_group_offload, sequential_cpu_offload].
 # model_full_load means that the entire model will be moved to the GPU.
@@ -130,7 +133,7 @@ if transformer_path is not None:
         )
     else:
         if transformer_path.endswith("safetensors"):
-            from safetensors.torch import load_file
+            from safetensors.torch import load_file, safe_open
             state_dict = load_file(transformer_path)
         else:
             state_dict = torch.load(transformer_path, map_location="cpu")
@@ -156,7 +159,7 @@ vae = AutoencoderKLMiniMaxH3.from_pretrained(
 if vae_path is not None:
     print(f"From checkpoint: {vae_path}")
     if vae_path.endswith("safetensors"):
-        from safetensors.torch import load_file
+        from safetensors.torch import load_file, safe_open
         state_dict = load_file(vae_path)
     else:
         state_dict = torch.load(vae_path, map_location="cpu")
@@ -214,14 +217,21 @@ if compile_dit:
         pipeline.transformer.transformer_blocks[i] = torch.compile(pipeline.transformer.transformer_blocks[i])
     print("Add Compile")
 
-# Quantize (when the mode carries an "_and_<quant>" suffix) and then place the pipeline.
-# The order lives inside the helper: quantization has to happen before the offload hooks
-# are installed.
-# The FP8 conversion above has already run (before the FSDP sharding, on purpose); only the dequant
-# wrapper and the memory placement are left, which is what the preconverted tag installs.
-apply_gpu_memory_mode(pipeline, GPU_memory_mode, device, weight_dtype,
-                      quant_tag="qfloat8_preconverted" if GPU_memory_mode.endswith("_and_qfloat8") else None,
-                      exclude_module_name=[])
+if GPU_memory_mode == "sequential_cpu_offload":
+    pipeline.enable_sequential_cpu_offload(device=device)
+elif GPU_memory_mode == "model_group_offload":
+    register_auto_device_hook(pipeline.transformer)
+    safe_enable_group_offload(pipeline, onload_device=device, offload_device="cpu", offload_type="leaf_level", use_stream=True)
+elif GPU_memory_mode == "model_cpu_offload_and_qfloat8":
+    convert_weight_dtype_wrapper(pipeline.transformer, weight_dtype)
+    pipeline.enable_model_cpu_offload(device=device)
+elif GPU_memory_mode == "model_cpu_offload":
+    pipeline.enable_model_cpu_offload(device=device)
+elif GPU_memory_mode == "model_full_load_and_qfloat8":
+    convert_weight_dtype_wrapper(pipeline.transformer, weight_dtype)
+    pipeline.to(device=device)
+else:
+    pipeline.to(device=device)
 
 # Merge LoRA through the standard entry point: `merge_lora` detects an official TaoMate-H3 adapter
 # directory and converts it to the kohya layout on the fly, and takes a kohya safetensors checkpoint

examples/turbodiffusion/predict_i2v_wan2.2.py

@@ -13,15 +13,19 @@ for project_root in project_roots:
     sys.path.insert(0, project_root) if project_root not in sys.path else None
 
 from videox_fun.dist import set_multi_gpus_devices, shard_model
-from videox_fun.models import (AutoencoderKLWan, AutoencoderKLWan3_8,
-                               AutoTokenizer, TurboWanTransformer3DModel,
-                               WanT5EncoderModel)
+from videox_fun.models import (AutoencoderKLWan, AutoencoderKLWan3_8, AutoTokenizer, CLIPModel,
+                              WanT5EncoderModel, TurboWanTransformer3DModel)
+from videox_fun.models.cache_utils import get_teacache_coefficients
 from videox_fun.pipeline import Wan2_2I2VPipeline
-from videox_fun.utils import (FlowDPMSolverMultistepScheduler,
-                              FlowUniPCMultistepScheduler,
-                              apply_gpu_memory_mode, filter_kwargs,
-                              get_image_to_video_latent, merge_lora,
-                              save_videos_grid, unmerge_lora)
+from videox_fun.utils.fp8_optimization import (convert_model_weight_to_float8, replace_parameters_by_name,
+                                              convert_weight_dtype_wrapper)
+from videox_fun.utils.lora_utils import merge_lora, unmerge_lora
+from videox_fun.utils import (register_auto_device_hook,
+                              safe_enable_group_offload)
+from videox_fun.utils.utils import (filter_kwargs, get_image_to_video_latent,
+                                   save_videos_grid)
+from videox_fun.utils.fm_solvers import FlowDPMSolverMultistepScheduler
+from videox_fun.utils.fm_solvers_unipc import FlowUniPCMultistepScheduler
 
 # GPU memory mode, which can be chosen in [model_full_load, model_full_load_and_qfloat8, model_cpu_offload, model_cpu_offload_and_qfloat8, model_group_offload, sequential_cpu_offload].
 # model_full_load means that the entire model will be moved to the GPU.
@@ -85,7 +89,7 @@ video_length        = 81
 fps                 = 16
 
 # Use torch.float16 if GPU does not support torch.bfloat16
-# Some graphics cards, such as v100, 2080ti, do not support torch.bfloat16
+# ome graphics cards, such as v100, 2080ti, do not support torch.bfloat16
 weight_dtype            = torch.bfloat16
 # If you want to generate from text, please set the validation_image_start = None and validation_image_end = None
 validation_image_start  = "asset/1.png"
@@ -122,7 +126,7 @@ transformer_2 = TurboWanTransformer3DModel.from_pretrained(
 if transformer_path is not None:
     print(f"From checkpoint: {transformer_path}")
     if transformer_path.endswith("safetensors"):
-        from safetensors.torch import load_file
+        from safetensors.torch import load_file, safe_open
         state_dict = load_file(transformer_path)
     else:
         state_dict = torch.load(transformer_path, map_location="cpu")
@@ -140,7 +144,7 @@ if transformer_path is not None:
 if transformer_high_path is not None:
     print(f"From checkpoint: {transformer_high_path}")
     if transformer_high_path.endswith("safetensors"):
-        from safetensors.torch import load_file
+        from safetensors.torch import load_file, safe_open
         state_dict = load_file(transformer_high_path)
     else:
         state_dict = torch.load(transformer_high_path, map_location="cpu")
@@ -168,7 +172,7 @@ vae = Chosen_AutoencoderKL.from_pretrained(
 if vae_path is not None:
     print(f"From checkpoint: {vae_path}")
     if vae_path.endswith("safetensors"):
-        from safetensors.torch import load_file
+        from safetensors.torch import load_file, safe_open
         state_dict = load_file(vae_path)
     else:
         state_dict = torch.load(vae_path, map_location="cpu")
@@ -192,7 +196,7 @@ text_encoder = WanT5EncoderModel.from_pretrained(
 text_encoder = text_encoder.eval()
 
 # Get Scheduler
-Chosen_Scheduler = {
+Chosen_Scheduler = scheduler_dict = {
     "Flow": FlowMatchEulerDiscreteScheduler,
     "Flow_Unipc": FlowUniPCMultistepScheduler,
     "Flow_DPM++": FlowDPMSolverMultistepScheduler,
@@ -233,11 +237,32 @@ if compile_dit:
         pipeline.transformer_2.blocks[i] = torch.compile(pipeline.transformer_2.blocks[i])
     print("Add Compile")
 
-# Quantize (when the mode carries an "_and_<quant>" suffix) and then place the pipeline.
-# The order lives inside the helper: quantization has to happen before the offload hooks
-# are installed, and both transformers of this MoE setup are handled in one call, which
-# is exactly the bookkeeping the old 30-line if/elif chain repeated per script.
-apply_gpu_memory_mode(pipeline, GPU_memory_mode, device, weight_dtype)
+if GPU_memory_mode == "sequential_cpu_offload":
+    replace_parameters_by_name(transformer, ["modulation",], device=device)
+    replace_parameters_by_name(transformer_2, ["modulation",], device=device)
+    transformer.freqs = transformer.freqs.to(device=device)
+    transformer_2.freqs = transformer_2.freqs.to(device=device)
+    pipeline.enable_sequential_cpu_offload(device=device)
+elif GPU_memory_mode == "model_group_offload":
+    register_auto_device_hook(transformer)
+    register_auto_device_hook(transformer_2)
+    safe_enable_group_offload(pipeline, onload_device=device, offload_device="cpu", offload_type="leaf_level", use_stream=True)
+elif GPU_memory_mode == "model_cpu_offload_and_qfloat8":
+    convert_model_weight_to_float8(transformer, exclude_module_name=["modulation",], device=device)
+    convert_model_weight_to_float8(transformer_2, exclude_module_name=["modulation",], device=device)
+    convert_weight_dtype_wrapper(transformer, weight_dtype)
+    convert_weight_dtype_wrapper(transformer_2, weight_dtype)
+    pipeline.enable_model_cpu_offload(device=device)
+elif GPU_memory_mode == "model_cpu_offload":
+    pipeline.enable_model_cpu_offload(device=device)
+elif GPU_memory_mode == "model_full_load_and_qfloat8":
+    convert_model_weight_to_float8(transformer, exclude_module_name=["modulation",], device=device)
+    convert_model_weight_to_float8(transformer_2, exclude_module_name=["modulation",], device=device)
+    convert_weight_dtype_wrapper(transformer, weight_dtype)
+    convert_weight_dtype_wrapper(transformer_2, weight_dtype)
+    pipeline.to(device=device)
+else:
+    pipeline.to(device=device)
 
 generator = torch.Generator(device=device).manual_seed(seed)
 

examples/turbodiffusion/predict_t2v_wan2.1.py

@@ -15,11 +15,18 @@ for project_root in project_roots:
 from videox_fun.dist import set_multi_gpus_devices, shard_model
 from videox_fun.models import (AutoencoderKLWan, AutoTokenizer,
                                TurboWanTransformer3DModel, WanT5EncoderModel)
+from videox_fun.models.cache_utils import get_teacache_coefficients
 from videox_fun.pipeline import WanPipeline
-from videox_fun.utils import (FlowDPMSolverMultistepScheduler,
-                              FlowUniPCMultistepScheduler,
-                              apply_gpu_memory_mode, filter_kwargs, merge_lora,
-                              save_videos_grid, unmerge_lora)
+from videox_fun.utils.fm_solvers import FlowDPMSolverMultistepScheduler
+from videox_fun.utils.fm_solvers_unipc import FlowUniPCMultistepScheduler
+from videox_fun.utils.fp8_optimization import (convert_model_weight_to_float8,
+                                               convert_weight_dtype_wrapper,
+                                               replace_parameters_by_name)
+from videox_fun.utils.lora_utils import merge_lora, unmerge_lora
+from videox_fun.utils import (register_auto_device_hook,
+                              safe_enable_group_offload)
+from videox_fun.utils.utils import (filter_kwargs, get_image_to_video_latent,
+                                    save_videos_grid)
 
 # GPU memory mode, which can be chosen in [model_full_load, model_full_load_and_qfloat8, model_cpu_offload, model_cpu_offload_and_qfloat8, model_group_offload, sequential_cpu_offload].
 # model_full_load means that the entire model will be moved to the GPU.
@@ -81,7 +88,7 @@ video_length        = 81
 fps                 = 16
 
 # Use torch.float16 if GPU does not support torch.bfloat16
-# Some graphics cards, such as v100, 2080ti, do not support torch.bfloat16
+# ome graphics cards, such as v100, 2080ti, do not support torch.bfloat16
 weight_dtype        = torch.bfloat16
 prompt              = "A stylish woman walks down a Tokyo street filled with warm glowing neon and animated city signage. She wears a black leather jacket, a long red dress, and black boots, and carries a black purse. She wears sunglasses and red lipstick. She walks confidently and casually. The street is damp and reflective, creating a mirror effect of the colorful lights. Many pedestrians walk about."
 negative_prompt     = "色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走"
@@ -103,7 +110,7 @@ transformer = TurboWanTransformer3DModel.from_pretrained(
 if transformer_path is not None:
     print(f"From checkpoint: {transformer_path}")
     if transformer_path.endswith("safetensors"):
-        from safetensors.torch import load_file
+        from safetensors.torch import load_file, safe_open
         state_dict = load_file(transformer_path)
     else:
         state_dict = torch.load(transformer_path, map_location="cpu")
@@ -129,7 +136,7 @@ vae = AutoencoderKLWan.from_pretrained(
 if vae_path is not None:
     print(f"From checkpoint: {vae_path}")
     if vae_path.endswith("safetensors"):
-        from safetensors.torch import load_file
+        from safetensors.torch import load_file, safe_open
         state_dict = load_file(vae_path)
     else:
         state_dict = torch.load(vae_path, map_location="cpu")
@@ -152,7 +159,7 @@ text_encoder = WanT5EncoderModel.from_pretrained(
 )
 
 # Get Scheduler
-Chosen_Scheduler = {
+Chosen_Scheduler = scheduler_dict = {
     "Flow": FlowMatchEulerDiscreteScheduler,
     "Flow_Unipc": FlowUniPCMultistepScheduler,
     "Flow_DPM++": FlowDPMSolverMultistepScheduler,
@@ -188,10 +195,25 @@ if compile_dit:
         pipeline.transformer.blocks[i] = torch.compile(pipeline.transformer.blocks[i])
     print("Add Compile")
 
-# Quantize (when the mode carries an "_and_<quant>" suffix) and then place the pipeline.
-# The order lives inside the helper: quantization has to happen before the offload hooks
-# are installed.
-apply_gpu_memory_mode(pipeline, GPU_memory_mode, device, weight_dtype)
+if GPU_memory_mode == "sequential_cpu_offload":
+    replace_parameters_by_name(transformer, ["modulation",], device=device)
+    transformer.freqs = transformer.freqs.to(device=device)
+    pipeline.enable_sequential_cpu_offload(device=device)
+elif GPU_memory_mode == "model_group_offload":
+    register_auto_device_hook(transformer)
+    safe_enable_group_offload(pipeline, onload_device=device, offload_device="cpu", offload_type="leaf_level", use_stream=True)
+elif GPU_memory_mode == "model_cpu_offload_and_qfloat8":
+    convert_model_weight_to_float8(transformer, exclude_module_name=["modulation",], device=device)
+    convert_weight_dtype_wrapper(transformer, weight_dtype)
+    pipeline.enable_model_cpu_offload(device=device)
+elif GPU_memory_mode == "model_cpu_offload":
+    pipeline.enable_model_cpu_offload(device=device)
+elif GPU_memory_mode == "model_full_load_and_qfloat8":
+    convert_model_weight_to_float8(transformer, exclude_module_name=["modulation",], device=device)
+    convert_weight_dtype_wrapper(transformer, weight_dtype)
+    pipeline.to(device=device)
+else:
+    pipeline.to(device=device)
 
 generator = torch.Generator(device=device).manual_seed(seed)
 

examples/wan2.1/app.py

@@ -10,7 +10,7 @@ for project_root in project_roots:
     sys.path.insert(0, project_root) if project_root not in sys.path else None
 
 from videox_fun.api.api import (infer_forward_api,
-                                update_diffusion_transformer_api)
+                               update_diffusion_transformer_api)
 from videox_fun.ui.controller import flow_scheduler_dict
 from videox_fun.ui.wan_ui import ui, ui_client, ui_host
 

examples/wan2.1/launch_api.py

@@ -4,6 +4,7 @@ import sys
 import time
 
 import gradio as gr
+import ray
 import torch
 
 current_file_path = os.path.abspath(__file__)
@@ -12,11 +13,10 @@ for project_root in project_roots:
     sys.path.insert(0, project_root) if project_root not in sys.path else None
 
 from videox_fun.api.api_multi_nodes import (MultiNodesEngine,
-                                            multi_nodes_infer_forward_api)
+                                           multi_nodes_infer_forward_api)
 from videox_fun.ui.controller import flow_scheduler_dict
 from videox_fun.ui.wan_ui import Wan_Controller
 
-
 def main():
     parser = argparse.ArgumentParser(description='xDiT HTTP Service')
     parser.add_argument('--world_size', type=int, default=8, help='Number of parallel workers')

examples/wan2.1/predict_i2v.py

@@ -18,11 +18,16 @@ from videox_fun.models import (AutoencoderKLWan, AutoTokenizer, CLIPModel,
                                WanT5EncoderModel, WanTransformer3DModel)
 from videox_fun.models.cache_utils import get_teacache_coefficients
 from videox_fun.pipeline import WanI2VPipeline
-from videox_fun.utils import (FlowDPMSolverMultistepScheduler,
-                              FlowUniPCMultistepScheduler,
-                              apply_gpu_memory_mode, filter_kwargs,
-                              get_image_to_video_latent, merge_lora,
-                              save_videos_grid, unmerge_lora)
+from videox_fun.utils import (register_auto_device_hook,
+                              safe_enable_group_offload)
+from videox_fun.utils.fm_solvers import FlowDPMSolverMultistepScheduler
+from videox_fun.utils.fm_solvers_unipc import FlowUniPCMultistepScheduler
+from videox_fun.utils.fp8_optimization import (convert_model_weight_to_float8,
+                                               convert_weight_dtype_wrapper,
+                                               replace_parameters_by_name)
+from videox_fun.utils.lora_utils import merge_lora, unmerge_lora
+from videox_fun.utils.utils import (filter_kwargs, get_image_to_video_latent,
+                                    save_videos_grid)
 
 # GPU memory mode, which can be chosen in [model_full_load, model_full_load_and_qfloat8, model_cpu_offload, model_cpu_offload_and_qfloat8, model_group_offload, sequential_cpu_offload].
 # model_full_load means that the entire model will be moved to the GPU.
@@ -103,7 +108,7 @@ video_length        = 81
 fps                 = 16
 
 # Use torch.float16 if GPU does not support torch.bfloat16
-# Some graphics cards, such as v100, 2080ti, do not support torch.bfloat16
+# ome graphics cards, such as v100, 2080ti, do not support torch.bfloat16
 weight_dtype            = torch.bfloat16
 # If you want to generate from text, please set the validation_image_start = None and validation_image_end = None
 validation_image_start  = "asset/1.png"
@@ -130,7 +135,7 @@ transformer = WanTransformer3DModel.from_pretrained(
 if transformer_path is not None:
     print(f"From checkpoint: {transformer_path}")
     if transformer_path.endswith("safetensors"):
-        from safetensors.torch import load_file
+        from safetensors.torch import load_file, safe_open
         state_dict = load_file(transformer_path)
     else:
         state_dict = torch.load(transformer_path, map_location="cpu")
@@ -148,7 +153,7 @@ vae = AutoencoderKLWan.from_pretrained(
 if vae_path is not None:
     print(f"From checkpoint: {vae_path}")
     if vae_path.endswith("safetensors"):
-        from safetensors.torch import load_file
+        from safetensors.torch import load_file, safe_open
         state_dict = load_file(vae_path)
     else:
         state_dict = torch.load(vae_path, map_location="cpu")
@@ -178,7 +183,7 @@ clip_image_encoder = CLIPModel.from_pretrained(
 clip_image_encoder = clip_image_encoder.eval()
 
 # Get Scheduler
-Chosen_Scheduler = {
+Chosen_Scheduler = scheduler_dict = {
     "Flow": FlowMatchEulerDiscreteScheduler,
     "Flow_Unipc": FlowUniPCMultistepScheduler,
     "Flow_DPM++": FlowDPMSolverMultistepScheduler,
@@ -215,10 +220,25 @@ if compile_dit:
         pipeline.transformer.blocks[i] = torch.compile(pipeline.transformer.blocks[i])
     print("Add Compile")
 
-# Quantize (when the mode carries an "_and_<quant>" suffix) and then place the pipeline.
-# The order lives inside the helper: quantization has to happen before the offload hooks
-# are installed.
-apply_gpu_memory_mode(pipeline, GPU_memory_mode, device, weight_dtype)
+if GPU_memory_mode == "sequential_cpu_offload":
+    replace_parameters_by_name(transformer, ["modulation",], device=device)
+    transformer.freqs = transformer.freqs.to(device=device)
+    pipeline.enable_sequential_cpu_offload(device=device)
+elif GPU_memory_mode == "model_group_offload":
+    register_auto_device_hook(pipeline.transformer)
+    safe_enable_group_offload(pipeline, onload_device=device, offload_device="cpu", offload_type="leaf_level", use_stream=True)
+elif GPU_memory_mode == "model_cpu_offload_and_qfloat8":
+    convert_model_weight_to_float8(transformer, exclude_module_name=["modulation",], device=device)
+    convert_weight_dtype_wrapper(transformer, weight_dtype)
+    pipeline.enable_model_cpu_offload(device=device)
+elif GPU_memory_mode == "model_cpu_offload":
+    pipeline.enable_model_cpu_offload(device=device)
+elif GPU_memory_mode == "model_full_load_and_qfloat8":
+    convert_model_weight_to_float8(transformer, exclude_module_name=["modulation",], device=device)
+    convert_weight_dtype_wrapper(transformer, weight_dtype)
+    pipeline.to(device=device)
+else:
+    pipeline.to(device=device)
 
 coefficients = get_teacache_coefficients(model_name) if enable_teacache else None
 if coefficients is not None:

examples/wan2.1/predict_i2v_tae.py

@@ -15,15 +15,20 @@ for project_root in project_roots:
 
 from videox_fun.dist import set_multi_gpus_devices, shard_model
 from videox_fun.models import (AutoencoderKLWan, AutoencoderTinyWan,
-                               AutoTokenizer, CLIPModel, WanT5EncoderModel,
-                               WanTransformer3DModel)
+                               AutoTokenizer, CLIPModel,
+                               WanT5EncoderModel, WanTransformer3DModel)
 from videox_fun.models.cache_utils import get_teacache_coefficients
 from videox_fun.pipeline import WanI2VPipeline
-from videox_fun.utils import (FlowDPMSolverMultistepScheduler,
-                              FlowUniPCMultistepScheduler,
-                              apply_gpu_memory_mode, filter_kwargs,
-                              get_image_to_video_latent, merge_lora,
-                              save_videos_grid, unmerge_lora)
+from videox_fun.utils import (register_auto_device_hook,
+                              safe_enable_group_offload)
+from videox_fun.utils.fm_solvers import FlowDPMSolverMultistepScheduler
+from videox_fun.utils.fm_solvers_unipc import FlowUniPCMultistepScheduler
+from videox_fun.utils.fp8_optimization import (convert_model_weight_to_float8,
+                                               convert_weight_dtype_wrapper,
+                                               replace_parameters_by_name)
+from videox_fun.utils.lora_utils import merge_lora, unmerge_lora
+from videox_fun.utils.utils import (filter_kwargs, get_image_to_video_latent,
+                                    save_videos_grid)
 
 # GPU memory mode, which can be chosen in [model_full_load, model_full_load_and_qfloat8, model_cpu_offload, model_cpu_offload_and_qfloat8, model_group_offload, sequential_cpu_offload].
 # model_full_load means that the entire model will be moved to the GPU.
@@ -114,7 +119,7 @@ video_length        = 81
 fps                 = 16
 
 # Use torch.float16 if GPU does not support torch.bfloat16
-# Some graphics cards, such as v100, 2080ti, do not support torch.bfloat16
+# ome graphics cards, such as v100, 2080ti, do not support torch.bfloat16
 weight_dtype            = torch.bfloat16
 # If you want to generate from text, please set the validation_image_start = None and validation_image_end = None
 validation_image_start  = "asset/1.png"
@@ -147,7 +152,7 @@ transformer = WanTransformer3DModel.from_pretrained(
 if transformer_path is not None:
     print(f"From checkpoint: {transformer_path}")
     if transformer_path.endswith("safetensors"):
-        from safetensors.torch import load_file
+        from safetensors.torch import load_file, safe_open
         state_dict = load_file(transformer_path)
     else:
         state_dict = torch.load(transformer_path, map_location="cpu")
@@ -169,7 +174,7 @@ vae = Chosen_AutoencoderKL.from_pretrained(
 if vae_path is not None:
     print(f"From checkpoint: {vae_path}")
     if vae_path.endswith("safetensors"):
-        from safetensors.torch import load_file
+        from safetensors.torch import load_file, safe_open
         state_dict = load_file(vae_path)
     else:
         state_dict = torch.load(vae_path, map_location="cpu")
@@ -203,7 +208,7 @@ clip_image_encoder = CLIPModel.from_pretrained(
 clip_image_encoder = clip_image_encoder.eval()
 
 # Get Scheduler
-Chosen_Scheduler = {
+Chosen_Scheduler = scheduler_dict = {
     "Flow": FlowMatchEulerDiscreteScheduler,
     "Flow_Unipc": FlowUniPCMultistepScheduler,
     "Flow_DPM++": FlowDPMSolverMultistepScheduler,
@@ -240,10 +245,25 @@ if compile_dit:
         pipeline.transformer.blocks[i] = torch.compile(pipeline.transformer.blocks[i])
     print("Add Compile")
 
-# Quantize (when the mode carries an "_and_<quant>" suffix) and then place the pipeline.
-# The order lives inside the helper: quantization has to happen before the offload hooks
-# are installed.
-apply_gpu_memory_mode(pipeline, GPU_memory_mode, device, weight_dtype)
+if GPU_memory_mode == "sequential_cpu_offload":
+    replace_parameters_by_name(transformer, ["modulation",], device=device)
+    transformer.freqs = transformer.freqs.to(device=device)
+    pipeline.enable_sequential_cpu_offload(device=device)
+elif GPU_memory_mode == "model_group_offload":
+    register_auto_device_hook(pipeline.transformer)
+    safe_enable_group_offload(pipeline, onload_device=device, offload_device="cpu", offload_type="leaf_level", use_stream=True)
+elif GPU_memory_mode == "model_cpu_offload_and_qfloat8":
+    convert_model_weight_to_float8(transformer, exclude_module_name=["modulation",], device=device)
+    convert_weight_dtype_wrapper(transformer, weight_dtype)
+    pipeline.enable_model_cpu_offload(device=device)
+elif GPU_memory_mode == "model_cpu_offload":
+    pipeline.enable_model_cpu_offload(device=device)
+elif GPU_memory_mode == "model_full_load_and_qfloat8":
+    convert_model_weight_to_float8(transformer, exclude_module_name=["modulation",], device=device)
+    convert_weight_dtype_wrapper(transformer, weight_dtype)
+    pipeline.to(device=device)
+else:
+    pipeline.to(device=device)
 
 coefficients = get_teacache_coefficients(model_name) if enable_teacache else None
 if coefficients is not None:

examples/wan2.1/predict_t2v.py

@@ -17,10 +17,16 @@ from videox_fun.models import (AutoencoderKLWan, AutoTokenizer,
                                WanT5EncoderModel, WanTransformer3DModel)
 from videox_fun.models.cache_utils import get_teacache_coefficients
 from videox_fun.pipeline import WanPipeline
-from videox_fun.utils import (FlowDPMSolverMultistepScheduler,
-                              FlowUniPCMultistepScheduler,
-                              apply_gpu_memory_mode, filter_kwargs, merge_lora,
-                              save_videos_grid, unmerge_lora)
+from videox_fun.utils import (register_auto_device_hook,
+                              safe_enable_group_offload)
+from videox_fun.utils.fm_solvers import FlowDPMSolverMultistepScheduler
+from videox_fun.utils.fm_solvers_unipc import FlowUniPCMultistepScheduler
+from videox_fun.utils.fp8_optimization import (convert_model_weight_to_float8,
+                                               convert_weight_dtype_wrapper,
+                                               replace_parameters_by_name)
+from videox_fun.utils.lora_utils import merge_lora, unmerge_lora
+from videox_fun.utils.utils import (filter_kwargs, get_image_to_video_latent,
+                                    save_videos_grid)
 
 # GPU memory mode, which can be chosen in [model_full_load, model_full_load_and_qfloat8, model_cpu_offload, model_cpu_offload_and_qfloat8, model_group_offload, sequential_cpu_offload].
 # model_full_load means that the entire model will be moved to the GPU.
@@ -101,7 +107,7 @@ video_length        = 81
 fps                 = 16
 
 # Use torch.float16 if GPU does not support torch.bfloat16
-# Some graphics cards, such as v100, 2080ti, do not support torch.bfloat16
+# ome graphics cards, such as v100, 2080ti, do not support torch.bfloat16
 weight_dtype        = torch.bfloat16
 prompt              = "一只棕色的狗摇着头，坐在舒适房间里的浅色沙发上。在狗的后面，架子上有一幅镶框的画，周围是粉红色的花朵。房间里柔和温暖的灯光营造出舒适的氛围。"
 negative_prompt     = "色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走"
@@ -124,7 +130,7 @@ transformer = WanTransformer3DModel.from_pretrained(
 if transformer_path is not None:
     print(f"From checkpoint: {transformer_path}")
     if transformer_path.endswith("safetensors"):
-        from safetensors.torch import load_file
+        from safetensors.torch import load_file, safe_open
         state_dict = load_file(transformer_path)
     else:
         state_dict = torch.load(transformer_path, map_location="cpu")
@@ -142,7 +148,7 @@ vae = AutoencoderKLWan.from_pretrained(
 if vae_path is not None:
     print(f"From checkpoint: {vae_path}")
     if vae_path.endswith("safetensors"):
-        from safetensors.torch import load_file
+        from safetensors.torch import load_file, safe_open
         state_dict = load_file(vae_path)
     else:
         state_dict = torch.load(vae_path, map_location="cpu")
@@ -165,7 +171,7 @@ text_encoder = WanT5EncoderModel.from_pretrained(
 )
 
 # Get Scheduler
-Chosen_Scheduler = {
+Chosen_Scheduler = scheduler_dict = {
     "Flow": FlowMatchEulerDiscreteScheduler,
     "Flow_Unipc": FlowUniPCMultistepScheduler,
     "Flow_DPM++": FlowDPMSolverMultistepScheduler,
@@ -201,10 +207,25 @@ if compile_dit:
         pipeline.transformer.blocks[i] = torch.compile(pipeline.transformer.blocks[i])
     print("Add Compile")
 
-# Quantize (when the mode carries an "_and_<quant>" suffix) and then place the pipeline.
-# The order lives inside the helper: quantization has to happen before the offload hooks
-# are installed.
-apply_gpu_memory_mode(pipeline, GPU_memory_mode, device, weight_dtype)
+if GPU_memory_mode == "sequential_cpu_offload":
+    replace_parameters_by_name(transformer, ["modulation",], device=device)
+    transformer.freqs = transformer.freqs.to(device=device)
+    pipeline.enable_sequential_cpu_offload(device=device)
+elif GPU_memory_mode == "model_group_offload":
+    register_auto_device_hook(pipeline.transformer)
+    safe_enable_group_offload(pipeline, onload_device=device, offload_device="cpu", offload_type="leaf_level", use_stream=True)
+elif GPU_memory_mode == "model_cpu_offload_and_qfloat8":
+    convert_model_weight_to_float8(transformer, exclude_module_name=["modulation",], device=device)
+    convert_weight_dtype_wrapper(transformer, weight_dtype)
+    pipeline.enable_model_cpu_offload(device=device)
+elif GPU_memory_mode == "model_cpu_offload":
+    pipeline.enable_model_cpu_offload(device=device)
+elif GPU_memory_mode == "model_full_load_and_qfloat8":
+    convert_model_weight_to_float8(transformer, exclude_module_name=["modulation",], device=device)
+    convert_weight_dtype_wrapper(transformer, weight_dtype)
+    pipeline.to(device=device)
+else:
+    pipeline.to(device=device)
 
 coefficients = get_teacache_coefficients(model_name) if enable_teacache else None
 if coefficients is not None:

examples/wan2.1/predict_t2v_tae.py

@@ -18,10 +18,16 @@ from videox_fun.models import (AutoencoderKLWan, AutoencoderTinyWan,
                                WanTransformer3DModel)
 from videox_fun.models.cache_utils import get_teacache_coefficients
 from videox_fun.pipeline import WanPipeline
-from videox_fun.utils import (FlowDPMSolverMultistepScheduler,
-                              FlowUniPCMultistepScheduler,
-                              apply_gpu_memory_mode, filter_kwargs, merge_lora,
-                              save_videos_grid, unmerge_lora)
+from videox_fun.utils import (register_auto_device_hook,
+                              safe_enable_group_offload)
+from videox_fun.utils.fm_solvers import FlowDPMSolverMultistepScheduler
+from videox_fun.utils.fm_solvers_unipc import FlowUniPCMultistepScheduler
+from videox_fun.utils.fp8_optimization import (convert_model_weight_to_float8,
+                                               convert_weight_dtype_wrapper,
+                                               replace_parameters_by_name)
+from videox_fun.utils.lora_utils import merge_lora, unmerge_lora
+from videox_fun.utils.utils import (filter_kwargs, get_image_to_video_latent,
+                                    save_videos_grid)
 
 # GPU memory mode, which can be chosen in [model_full_load, model_full_load_and_qfloat8, model_cpu_offload, model_cpu_offload_and_qfloat8, model_group_offload, sequential_cpu_offload].
 # model_full_load means that the entire model will be moved to the GPU.
@@ -112,7 +118,7 @@ video_length        = 81
 fps                 = 16
 
 # Use torch.float16 if GPU does not support torch.bfloat16
-# Some graphics cards, such as v100, 2080ti, do not support torch.bfloat16
+# ome graphics cards, such as v100, 2080ti, do not support torch.bfloat16
 weight_dtype        = torch.bfloat16
 prompt              = "一只棕色的狗摇着头，坐在舒适房间里的浅色沙发上。在狗的后面，架子上有一幅镶框的画，周围是粉红色的花朵。房间里柔和温暖的灯光营造出舒适的氛围。"
 negative_prompt     = "色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走"
@@ -141,7 +147,7 @@ transformer = WanTransformer3DModel.from_pretrained(
 if transformer_path is not None:
     print(f"From checkpoint: {transformer_path}")
     if transformer_path.endswith("safetensors"):
-        from safetensors.torch import load_file
+        from safetensors.torch import load_file, safe_open
         state_dict = load_file(transformer_path)
     else:
         state_dict = torch.load(transformer_path, map_location="cpu")
@@ -163,7 +169,7 @@ vae = Chosen_AutoencoderKL.from_pretrained(
 if vae_path is not None:
     print(f"From checkpoint: {vae_path}")
     if vae_path.endswith("safetensors"):
-        from safetensors.torch import load_file
+        from safetensors.torch import load_file, safe_open
         state_dict = load_file(vae_path)
     else:
         state_dict = torch.load(vae_path, map_location="cpu")
@@ -190,7 +196,7 @@ text_encoder = WanT5EncoderModel.from_pretrained(
 )
 
 # Get Scheduler
-Chosen_Scheduler = {
+Chosen_Scheduler = scheduler_dict = {
     "Flow": FlowMatchEulerDiscreteScheduler,
     "Flow_Unipc": FlowUniPCMultistepScheduler,
     "Flow_DPM++": FlowDPMSolverMultistepScheduler,
@@ -226,10 +232,25 @@ if compile_dit:
         pipeline.transformer.blocks[i] = torch.compile(pipeline.transformer.blocks[i])
     print("Add Compile")
 
-# Quantize (when the mode carries an "_and_<quant>" suffix) and then place the pipeline.
-# The order lives inside the helper: quantization has to happen before the offload hooks
-# are installed.
-apply_gpu_memory_mode(pipeline, GPU_memory_mode, device, weight_dtype)
+if GPU_memory_mode == "sequential_cpu_offload":
+    replace_parameters_by_name(transformer, ["modulation",], device=device)
+    transformer.freqs = transformer.freqs.to(device=device)
+    pipeline.enable_sequential_cpu_offload(device=device)
+elif GPU_memory_mode == "model_group_offload":
+    register_auto_device_hook(pipeline.transformer)
+    safe_enable_group_offload(pipeline, onload_device=device, offload_device="cpu", offload_type="leaf_level", use_stream=True)
+elif GPU_memory_mode == "model_cpu_offload_and_qfloat8":
+    convert_model_weight_to_float8(transformer, exclude_module_name=["modulation",], device=device)
+    convert_weight_dtype_wrapper(transformer, weight_dtype)
+    pipeline.enable_model_cpu_offload(device=device)
+elif GPU_memory_mode == "model_cpu_offload":
+    pipeline.enable_model_cpu_offload(device=device)
+elif GPU_memory_mode == "model_full_load_and_qfloat8":
+    convert_model_weight_to_float8(transformer, exclude_module_name=["modulation",], device=device)
+    convert_weight_dtype_wrapper(transformer, weight_dtype)
+    pipeline.to(device=device)
+else:
+    pipeline.to(device=device)
 
 coefficients = get_teacache_coefficients(model_name) if enable_teacache else None
 if coefficients is not None:

examples/wan2.1_causal_forcing/predict_t2v.py

@@ -17,10 +17,16 @@ from videox_fun.models import (AutoencoderKLWan, AutoTokenizer,
                                WanT5EncoderModel,
                                WanTransformer3DModel_SelfForcing)
 from videox_fun.pipeline import WanSelfForcingPipeline
-from videox_fun.utils import (FlowDPMSolverMultistepScheduler,
-                              FlowUniPCMultistepScheduler,
-                              apply_gpu_memory_mode, filter_kwargs, merge_lora,
-                              save_videos_grid, unmerge_lora)
+from videox_fun.utils import (register_auto_device_hook,
+                              safe_enable_group_offload)
+from videox_fun.utils.fm_solvers import FlowDPMSolverMultistepScheduler
+from videox_fun.utils.fm_solvers_unipc import FlowUniPCMultistepScheduler
+from videox_fun.utils.fp8_optimization import (convert_model_weight_to_float8,
+                                               convert_weight_dtype_wrapper,
+                                               replace_parameters_by_name)
+from videox_fun.utils.lora_utils import merge_lora, unmerge_lora
+from videox_fun.utils.utils import (filter_kwargs, get_image_to_video_latent,
+                                    save_videos_grid)
 
 # GPU memory mode, which can be chosen in [model_full_load, model_full_load_and_qfloat8, model_cpu_offload, model_cpu_offload_and_qfloat8, model_group_offload, sequential_cpu_offload].
 # model_full_load means that the entire model will be moved to the GPU.
@@ -160,7 +166,7 @@ if transformer_path is not None:
         print(f"use_ema={use_ema}: resolved {_raw_transformer_path} -> {transformer_path}")
     print(f"From checkpoint: {transformer_path}")
     if transformer_path.endswith("safetensors"):
-        from safetensors.torch import load_file
+        from safetensors.torch import load_file, safe_open
         state_dict = load_file(transformer_path)
     else:
         state_dict = torch.load(transformer_path, map_location="cpu")
@@ -188,7 +194,7 @@ vae = AutoencoderKLWan.from_pretrained(
 if vae_path is not None:
     print(f"From checkpoint: {vae_path}")
     if vae_path.endswith("safetensors"):
-        from safetensors.torch import load_file
+        from safetensors.torch import load_file, safe_open
         state_dict = load_file(vae_path)
     else:
         state_dict = torch.load(vae_path, map_location="cpu")
@@ -211,7 +217,7 @@ text_encoder = WanT5EncoderModel.from_pretrained(
 )
 
 # Get Scheduler
-Chosen_Scheduler = {
+Chosen_Scheduler = scheduler_dict = {
     "Flow": FlowMatchEulerDiscreteScheduler,
     "Flow_Unipc": FlowUniPCMultistepScheduler,
     "Flow_DPM++": FlowDPMSolverMultistepScheduler,
@@ -248,10 +254,25 @@ if compile_dit:
         pipeline.transformer.blocks[i] = torch.compile(pipeline.transformer.blocks[i])
     print("Add Compile")
 
-# Quantize (when the mode carries an "_and_<quant>" suffix) and then place the pipeline.
-# The order lives inside the helper: quantization has to happen before the offload hooks
-# are installed.
-apply_gpu_memory_mode(pipeline, GPU_memory_mode, device, weight_dtype)
+if GPU_memory_mode == "sequential_cpu_offload":
+    replace_parameters_by_name(transformer, ["modulation",], device=device)
+    transformer.freqs = transformer.freqs.to(device=device)
+    pipeline.enable_sequential_cpu_offload(device=device)
+elif GPU_memory_mode == "model_group_offload":
+    register_auto_device_hook(pipeline.transformer)
+    safe_enable_group_offload(pipeline, onload_device=device, offload_device="cpu", offload_type="leaf_level", use_stream=True)
+elif GPU_memory_mode == "model_cpu_offload_and_qfloat8":
+    convert_model_weight_to_float8(transformer, exclude_module_name=["modulation",], device=device)
+    convert_weight_dtype_wrapper(transformer, weight_dtype)
+    pipeline.enable_model_cpu_offload(device=device)
+elif GPU_memory_mode == "model_cpu_offload":
+    pipeline.enable_model_cpu_offload(device=device)
+elif GPU_memory_mode == "model_full_load_and_qfloat8":
+    convert_model_weight_to_float8(transformer, exclude_module_name=["modulation",], device=device)
+    convert_weight_dtype_wrapper(transformer, weight_dtype)
+    pipeline.to(device=device)
+else:
+    pipeline.to(device=device)
 
 generator = torch.Generator(device=device).manual_seed(seed)
 

examples/wan2.1_causal_forcing/predict_t2v_stream.py

@@ -17,11 +17,17 @@ from videox_fun.models import (AutoencoderKLWan, AutoTokenizer,
                                WanT5EncoderModel,
                                WanTransformer3DModel_SelfForcing)
 from videox_fun.pipeline import WanSelfForcingPipeline
-from videox_fun.utils import (FlowDPMSolverMultistepScheduler,
-                              FlowUniPCMultistepScheduler, SegmentVideoSaver,
-                              StreamVideoSaver, apply_gpu_memory_mode,
-                              filter_kwargs, merge_lora, save_videos_grid,
-                              unmerge_lora)
+from videox_fun.utils import (register_auto_device_hook,
+                              safe_enable_group_offload)
+from videox_fun.utils.fm_solvers import FlowDPMSolverMultistepScheduler
+from videox_fun.utils.fm_solvers_unipc import FlowUniPCMultistepScheduler
+from videox_fun.utils.fp8_optimization import (convert_model_weight_to_float8,
+                                               convert_weight_dtype_wrapper,
+                                               replace_parameters_by_name)
+from videox_fun.utils.lora_utils import merge_lora, unmerge_lora
+from videox_fun.utils.utils import (filter_kwargs, get_image_to_video_latent,
+                                    save_videos_grid, StreamVideoSaver,
+                                    SegmentVideoSaver)
 
 # GPU memory mode, which can be chosen in [model_full_load, model_full_load_and_qfloat8, model_cpu_offload, model_cpu_offload_and_qfloat8, model_group_offload, sequential_cpu_offload].
 # model_full_load means that the entire model will be moved to the GPU.
@@ -173,7 +179,7 @@ if transformer_path is not None:
         print(f"use_ema={use_ema}: resolved {_raw_transformer_path} -> {transformer_path}")
     print(f"From checkpoint: {transformer_path}")
     if transformer_path.endswith("safetensors"):
-        from safetensors.torch import load_file
+        from safetensors.torch import load_file, safe_open
         state_dict = load_file(transformer_path)
     else:
         state_dict = torch.load(transformer_path, map_location="cpu")
@@ -201,7 +207,7 @@ vae = AutoencoderKLWan.from_pretrained(
 if vae_path is not None:
     print(f"From checkpoint: {vae_path}")
     if vae_path.endswith("safetensors"):
-        from safetensors.torch import load_file
+        from safetensors.torch import load_file, safe_open
         state_dict = load_file(vae_path)
     else:
         state_dict = torch.load(vae_path, map_location="cpu")
@@ -224,7 +230,7 @@ text_encoder = WanT5EncoderModel.from_pretrained(
 )
 
 # Get Scheduler
-Chosen_Scheduler = {
+Chosen_Scheduler = scheduler_dict = {
     "Flow": FlowMatchEulerDiscreteScheduler,
     "Flow_Unipc": FlowUniPCMultistepScheduler,
     "Flow_DPM++": FlowDPMSolverMultistepScheduler,
@@ -261,10 +267,25 @@ if compile_dit:
         pipeline.transformer.blocks[i] = torch.compile(pipeline.transformer.blocks[i])
     print("Add Compile")
 
-# Quantize (when the mode carries an "_and_<quant>" suffix) and then place the pipeline.
-# The order lives inside the helper: quantization has to happen before the offload hooks
-# are installed.
-apply_gpu_memory_mode(pipeline, GPU_memory_mode, device, weight_dtype)
+if GPU_memory_mode == "sequential_cpu_offload":
+    replace_parameters_by_name(transformer, ["modulation",], device=device)
+    transformer.freqs = transformer.freqs.to(device=device)
+    pipeline.enable_sequential_cpu_offload(device=device)
+elif GPU_memory_mode == "model_group_offload":
+    register_auto_device_hook(pipeline.transformer)
+    safe_enable_group_offload(pipeline, onload_device=device, offload_device="cpu", offload_type="leaf_level", use_stream=True)
+elif GPU_memory_mode == "model_cpu_offload_and_qfloat8":
+    convert_model_weight_to_float8(transformer, exclude_module_name=["modulation",], device=device)
+    convert_weight_dtype_wrapper(transformer, weight_dtype)
+    pipeline.enable_model_cpu_offload(device=device)
+elif GPU_memory_mode == "model_cpu_offload":
+    pipeline.enable_model_cpu_offload(device=device)
+elif GPU_memory_mode == "model_full_load_and_qfloat8":
+    convert_model_weight_to_float8(transformer, exclude_module_name=["modulation",], device=device)
+    convert_weight_dtype_wrapper(transformer, weight_dtype)
+    pipeline.to(device=device)
+else:
+    pipeline.to(device=device)
 
 # Only the main process (rank 0, or single-GPU) writes files to disk.
 if ulysses_degree * ring_degree > 1:

examples/wan2.1_flex_forcing/predict_t2v.py

@@ -18,11 +18,17 @@ from videox_fun.models import (AutoencoderKLWan, AutoTokenizer,
                                WanT5EncoderModel,
                                WanTransformer3DModel_FlexForcing)
 from videox_fun.pipeline import WanFlexForcingPipeline
-from videox_fun.pipeline.pipeline_wan_flex_forcing import PAPER_CHUNK_CONFIGS
-from videox_fun.utils import (FlowDPMSolverMultistepScheduler,
-                              FlowUniPCMultistepScheduler,
-                              apply_gpu_memory_mode, filter_kwargs, merge_lora,
-                              save_videos_grid, unmerge_lora)
+from videox_fun.pipeline.pipeline_wan_flex_forcing import \
+    PAPER_CHUNK_CONFIGS
+from videox_fun.utils import (register_auto_device_hook,
+                              safe_enable_group_offload)
+from videox_fun.utils.fm_solvers import FlowDPMSolverMultistepScheduler
+from videox_fun.utils.fm_solvers_unipc import FlowUniPCMultistepScheduler
+from videox_fun.utils.fp8_optimization import (convert_model_weight_to_float8,
+                                               convert_weight_dtype_wrapper,
+                                               replace_parameters_by_name)
+from videox_fun.utils.lora_utils import merge_lora, unmerge_lora
+from videox_fun.utils.utils import filter_kwargs, save_videos_grid
 
 # GPU memory mode, which can be chosen in [model_full_load, model_full_load_and_qfloat8, model_cpu_offload, model_cpu_offload_and_qfloat8, model_group_offload, sequential_cpu_offload].
 # model_full_load means that the entire model will be moved to the GPU.
@@ -155,7 +161,7 @@ transformer = WanTransformer3DModel_FlexForcing.from_pretrained(
 if transformer_path is not None:
     print(f"From checkpoint: {transformer_path}")
     if transformer_path.endswith("safetensors"):
-        from safetensors.torch import load_file
+        from safetensors.torch import load_file, safe_open
         state_dict = load_file(transformer_path)
     else:
         state_dict = torch.load(transformer_path, map_location="cpu")
@@ -184,7 +190,7 @@ vae = AutoencoderKLWan.from_pretrained(
 if vae_path is not None:
     print(f"From checkpoint: {vae_path}")
     if vae_path.endswith("safetensors"):
-        from safetensors.torch import load_file
+        from safetensors.torch import load_file, safe_open
         state_dict = load_file(vae_path)
     else:
         state_dict = torch.load(vae_path, map_location="cpu")
@@ -207,7 +213,7 @@ text_encoder = WanT5EncoderModel.from_pretrained(
 )
 
 # Get Scheduler
-Chosen_Scheduler = {
+Chosen_Scheduler = scheduler_dict = {
     "Flow": FlowMatchEulerDiscreteScheduler,
     "Flow_Unipc": FlowUniPCMultistepScheduler,
     "Flow_DPM++": FlowDPMSolverMultistepScheduler,
@@ -244,10 +250,25 @@ if compile_dit:
         pipeline.transformer.blocks[i] = torch.compile(pipeline.transformer.blocks[i])
     print("Add Compile")
 
-# Quantize (when the mode carries an "_and_<quant>" suffix) and then place the pipeline.
-# The order lives inside the helper: quantization has to happen before the offload hooks
-# are installed.
-apply_gpu_memory_mode(pipeline, GPU_memory_mode, device, weight_dtype)
+if GPU_memory_mode == "sequential_cpu_offload":
+    replace_parameters_by_name(transformer, ["modulation",], device=device)
+    transformer.freqs = transformer.freqs.to(device=device)
+    pipeline.enable_sequential_cpu_offload(device=device)
+elif GPU_memory_mode == "model_group_offload":
+    register_auto_device_hook(pipeline.transformer)
+    safe_enable_group_offload(pipeline, onload_device=device, offload_device="cpu", offload_type="leaf_level", use_stream=True)
+elif GPU_memory_mode == "model_cpu_offload_and_qfloat8":
+    convert_model_weight_to_float8(transformer, exclude_module_name=["modulation",], device=device)
+    convert_weight_dtype_wrapper(transformer, weight_dtype)
+    pipeline.enable_model_cpu_offload(device=device)
+elif GPU_memory_mode == "model_cpu_offload":
+    pipeline.enable_model_cpu_offload(device=device)
+elif GPU_memory_mode == "model_full_load_and_qfloat8":
+    convert_model_weight_to_float8(transformer, exclude_module_name=["modulation",], device=device)
+    convert_weight_dtype_wrapper(transformer, weight_dtype)
+    pipeline.to(device=device)
+else:
+    pipeline.to(device=device)
 
 print(f"[Flex-Forcing] denoise_mode={denoise_mode}, "
       f"min_num_frame_per_block={min_num_frame_per_block}, "

examples/wan2.1_flex_forcing/predict_t2v_edit.py

@@ -18,11 +18,16 @@ from videox_fun.models import (AutoencoderKLWan, AutoTokenizer,
                                WanT5EncoderModel,
                                WanTransformer3DModel_FlexForcing)
 from videox_fun.pipeline import WanFlexForcingPipeline
-from videox_fun.utils import (FlowDPMSolverMultistepScheduler,
-                              FlowUniPCMultistepScheduler,
-                              apply_gpu_memory_mode, filter_kwargs,
-                              get_video_to_video_latent, merge_lora,
-                              save_videos_grid, unmerge_lora)
+from videox_fun.utils import (register_auto_device_hook,
+                              safe_enable_group_offload)
+from videox_fun.utils.fm_solvers import FlowDPMSolverMultistepScheduler
+from videox_fun.utils.fm_solvers_unipc import FlowUniPCMultistepScheduler
+from videox_fun.utils.fp8_optimization import (convert_model_weight_to_float8,
+                                               convert_weight_dtype_wrapper,
+                                               replace_parameters_by_name)
+from videox_fun.utils.lora_utils import merge_lora, unmerge_lora
+from videox_fun.utils.utils import (filter_kwargs, get_video_to_video_latent,
+                                    save_videos_grid)
 
 # GPU memory mode, which can be chosen in [model_full_load, model_full_load_and_qfloat8, model_cpu_offload, model_cpu_offload_and_qfloat8, model_group_offload, sequential_cpu_offload].
 # model_full_load means that the entire model will be moved to the GPU.
@@ -159,7 +164,7 @@ transformer = WanTransformer3DModel_FlexForcing.from_pretrained(
 if transformer_path is not None:
     print(f"From checkpoint: {transformer_path}")
     if transformer_path.endswith("safetensors"):
-        from safetensors.torch import load_file
+        from safetensors.torch import load_file, safe_open
         state_dict = load_file(transformer_path)
     else:
         state_dict = torch.load(transformer_path, map_location="cpu")
@@ -188,7 +193,7 @@ vae = AutoencoderKLWan.from_pretrained(
 if vae_path is not None:
     print(f"From checkpoint: {vae_path}")
     if vae_path.endswith("safetensors"):
-        from safetensors.torch import load_file
+        from safetensors.torch import load_file, safe_open
         state_dict = load_file(vae_path)
     else:
         state_dict = torch.load(vae_path, map_location="cpu")
@@ -211,7 +216,7 @@ text_encoder = WanT5EncoderModel.from_pretrained(
 )
 
 # Get Scheduler
-Chosen_Scheduler = {
+Chosen_Scheduler = scheduler_dict = {
     "Flow": FlowMatchEulerDiscreteScheduler,
     "Flow_Unipc": FlowUniPCMultistepScheduler,
     "Flow_DPM++": FlowDPMSolverMultistepScheduler,
@@ -248,10 +253,25 @@ if compile_dit:
         pipeline.transformer.blocks[i] = torch.compile(pipeline.transformer.blocks[i])
     print("Add Compile")
 
-# Quantize (when the mode carries an "_and_<quant>" suffix) and then place the pipeline.
-# The order lives inside the helper: quantization has to happen before the offload hooks
-# are installed.
-apply_gpu_memory_mode(pipeline, GPU_memory_mode, device, weight_dtype)
+if GPU_memory_mode == "sequential_cpu_offload":
+    replace_parameters_by_name(transformer, ["modulation",], device=device)
+    transformer.freqs = transformer.freqs.to(device=device)
+    pipeline.enable_sequential_cpu_offload(device=device)
+elif GPU_memory_mode == "model_group_offload":
+    register_auto_device_hook(pipeline.transformer)
+    safe_enable_group_offload(pipeline, onload_device=device, offload_device="cpu", offload_type="leaf_level", use_stream=True)
+elif GPU_memory_mode == "model_cpu_offload_and_qfloat8":
+    convert_model_weight_to_float8(transformer, exclude_module_name=["modulation",], device=device)
+    convert_weight_dtype_wrapper(transformer, weight_dtype)
+    pipeline.enable_model_cpu_offload(device=device)
+elif GPU_memory_mode == "model_cpu_offload":
+    pipeline.enable_model_cpu_offload(device=device)
+elif GPU_memory_mode == "model_full_load_and_qfloat8":
+    convert_model_weight_to_float8(transformer, exclude_module_name=["modulation",], device=device)
+    convert_weight_dtype_wrapper(transformer, weight_dtype)
+    pipeline.to(device=device)
+else:
+    pipeline.to(device=device)
 
 print(f"[Flex-Forcing] local_attn_size={local_attn_size}, sink_size={sink_size}, "
       f"context_noise={context_noise}")

examples/wan2.1_fun/app.py

@@ -10,7 +10,7 @@ for project_root in project_roots:
     sys.path.insert(0, project_root) if project_root not in sys.path else None
 
 from videox_fun.api.api import (infer_forward_api,
-                                update_diffusion_transformer_api)
+                               update_diffusion_transformer_api)
 from videox_fun.ui.controller import flow_scheduler_dict
 from videox_fun.ui.wan_fun_ui import ui, ui_client, ui_host
 

examples/wan2.1_fun/launch_api.py

@@ -4,6 +4,7 @@ import sys
 import time
 
 import gradio as gr
+import ray
 import torch
 
 current_file_path = os.path.abspath(__file__)
@@ -12,11 +13,10 @@ for project_root in project_roots:
     sys.path.insert(0, project_root) if project_root not in sys.path else None
 
 from videox_fun.api.api_multi_nodes import (MultiNodesEngine,
-                                            multi_nodes_infer_forward_api)
+                                           multi_nodes_infer_forward_api)
 from videox_fun.ui.controller import flow_scheduler_dict
 from videox_fun.ui.wan_fun_ui import Wan_Fun_Controller
 
-
 def main():
     parser = argparse.ArgumentParser(description='xDiT HTTP Service')
     parser.add_argument('--world_size', type=int, default=8, help='Number of parallel workers')

examples/wan2.1_fun/predict_i2v.py

@@ -18,11 +18,16 @@ from videox_fun.models import (AutoencoderKLWan, CLIPModel, WanT5EncoderModel,
                                WanTransformer3DModel)
 from videox_fun.models.cache_utils import get_teacache_coefficients
 from videox_fun.pipeline import WanFunInpaintPipeline
-from videox_fun.utils import (FlowDPMSolverMultistepScheduler,
-                              FlowUniPCMultistepScheduler,
-                              apply_gpu_memory_mode, filter_kwargs,
-                              get_image_to_video_latent, merge_lora,
-                              save_videos_grid, unmerge_lora)
+from videox_fun.utils import (register_auto_device_hook,
+                              safe_enable_group_offload)
+from videox_fun.utils.fm_solvers import FlowDPMSolverMultistepScheduler
+from videox_fun.utils.fm_solvers_unipc import FlowUniPCMultistepScheduler
+from videox_fun.utils.fp8_optimization import (convert_model_weight_to_float8,
+                                               convert_weight_dtype_wrapper,
+                                               replace_parameters_by_name)
+from videox_fun.utils.lora_utils import merge_lora, unmerge_lora
+from videox_fun.utils.utils import (filter_kwargs, get_image_to_video_latent,
+                                    save_videos_grid)
 
 # GPU memory mode, which can be chosen in [model_full_load, model_full_load_and_qfloat8, model_cpu_offload, model_cpu_offload_and_qfloat8, model_group_offload, sequential_cpu_offload].
 # model_full_load means that the entire model will be moved to the GPU.
@@ -103,7 +108,7 @@ video_length        = 81
 fps                 = 16
 
 # Use torch.float16 if GPU does not support torch.bfloat16
-# Some graphics cards, such as v100, 2080ti, do not support torch.bfloat16
+# ome graphics cards, such as v100, 2080ti, do not support torch.bfloat16
 weight_dtype            = torch.bfloat16
 # If you want to generate from text, please set the validation_image_start = None and validation_image_end = None
 validation_image_start  = "asset/1.png"
@@ -131,7 +136,7 @@ transformer = WanTransformer3DModel.from_pretrained(
 if transformer_path is not None:
     print(f"From checkpoint: {transformer_path}")
     if transformer_path.endswith("safetensors"):
-        from safetensors.torch import load_file
+        from safetensors.torch import load_file, safe_open
         state_dict = load_file(transformer_path)
     else:
         state_dict = torch.load(transformer_path, map_location="cpu")
@@ -149,7 +154,7 @@ vae = AutoencoderKLWan.from_pretrained(
 if vae_path is not None:
     print(f"From checkpoint: {vae_path}")
     if vae_path.endswith("safetensors"):
-        from safetensors.torch import load_file
+        from safetensors.torch import load_file, safe_open
         state_dict = load_file(vae_path)
     else:
         state_dict = torch.load(vae_path, map_location="cpu")
@@ -179,7 +184,7 @@ clip_image_encoder = CLIPModel.from_pretrained(
 clip_image_encoder = clip_image_encoder.eval()
 
 # Get Scheduler
-Chosen_Scheduler = {
+Chosen_Scheduler = scheduler_dict = {
     "Flow": FlowMatchEulerDiscreteScheduler,
     "Flow_Unipc": FlowUniPCMultistepScheduler,
     "Flow_DPM++": FlowDPMSolverMultistepScheduler,
@@ -216,10 +221,25 @@ if compile_dit:
         pipeline.transformer.blocks[i] = torch.compile(pipeline.transformer.blocks[i])
     print("Add Compile")
 
-# Quantize (when the mode carries an "_and_<quant>" suffix) and then place the pipeline.
-# The order lives inside the helper: quantization has to happen before the offload hooks
-# are installed.
-apply_gpu_memory_mode(pipeline, GPU_memory_mode, device, weight_dtype)
+if GPU_memory_mode == "sequential_cpu_offload":
+    replace_parameters_by_name(transformer, ["modulation",], device=device)
+    transformer.freqs = transformer.freqs.to(device=device)
+    pipeline.enable_sequential_cpu_offload(device=device)
+elif GPU_memory_mode == "model_group_offload":
+    register_auto_device_hook(pipeline.transformer)
+    safe_enable_group_offload(pipeline, onload_device=device, offload_device="cpu", offload_type="leaf_level", use_stream=True)
+elif GPU_memory_mode == "model_cpu_offload_and_qfloat8":
+    convert_model_weight_to_float8(transformer, exclude_module_name=["modulation",], device=device)
+    convert_weight_dtype_wrapper(transformer, weight_dtype)
+    pipeline.enable_model_cpu_offload(device=device)
+elif GPU_memory_mode == "model_cpu_offload":
+    pipeline.enable_model_cpu_offload(device=device)
+elif GPU_memory_mode == "model_full_load_and_qfloat8":
+    convert_model_weight_to_float8(transformer, exclude_module_name=["modulation",], device=device)
+    convert_weight_dtype_wrapper(transformer, weight_dtype)
+    pipeline.to(device=device)
+else:
+    pipeline.to(device=device)
 
 coefficients = get_teacache_coefficients(model_name) if enable_teacache else None
 if coefficients is not None:

examples/wan2.1_fun/predict_t2v.py

@@ -18,11 +18,16 @@ from videox_fun.models import (AutoencoderKLWan, CLIPModel, WanT5EncoderModel,
                                WanTransformer3DModel)
 from videox_fun.models.cache_utils import get_teacache_coefficients
 from videox_fun.pipeline import WanFunInpaintPipeline, WanFunPipeline
-from videox_fun.utils import (FlowDPMSolverMultistepScheduler,
-                              FlowUniPCMultistepScheduler,
-                              apply_gpu_memory_mode, filter_kwargs,
-                              get_image_to_video_latent, merge_lora,
-                              save_videos_grid, unmerge_lora)
+from videox_fun.utils import (register_auto_device_hook,
+                              safe_enable_group_offload)
+from videox_fun.utils.fm_solvers import FlowDPMSolverMultistepScheduler
+from videox_fun.utils.fm_solvers_unipc import FlowUniPCMultistepScheduler
+from videox_fun.utils.fp8_optimization import (convert_model_weight_to_float8,
+                                               convert_weight_dtype_wrapper,
+                                               replace_parameters_by_name)
+from videox_fun.utils.lora_utils import merge_lora, unmerge_lora
+from videox_fun.utils.utils import (filter_kwargs, get_image_to_video_latent,
+                                    save_videos_grid)
 
 # GPU memory mode, which can be chosen in [model_full_load, model_full_load_and_qfloat8, model_cpu_offload, model_cpu_offload_and_qfloat8, model_group_offload, sequential_cpu_offload].
 # model_full_load means that the entire model will be moved to the GPU.
@@ -103,7 +108,7 @@ video_length        = 81
 fps                 = 16
 
 # Use torch.float16 if GPU does not support torch.bfloat16
-# Some graphics cards, such as v100, 2080ti, do not support torch.bfloat16
+# ome graphics cards, such as v100, 2080ti, do not support torch.bfloat16
 weight_dtype        = torch.bfloat16
 prompt              = "一只棕色的狗摇着头，坐在舒适房间里的浅色沙发上。在狗的后面，架子上有一幅镶框的画，周围是粉红色的花朵。房间里柔和温暖的灯光营造出舒适的氛围。"
 negative_prompt     = "色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走"
@@ -126,7 +131,7 @@ transformer = WanTransformer3DModel.from_pretrained(
 if transformer_path is not None:
     print(f"From checkpoint: {transformer_path}")
     if transformer_path.endswith("safetensors"):
-        from safetensors.torch import load_file
+        from safetensors.torch import load_file, safe_open
         state_dict = load_file(transformer_path)
     else:
         state_dict = torch.load(transformer_path, map_location="cpu")
@@ -144,7 +149,7 @@ vae = AutoencoderKLWan.from_pretrained(
 if vae_path is not None:
     print(f"From checkpoint: {vae_path}")
     if vae_path.endswith("safetensors"):
-        from safetensors.torch import load_file
+        from safetensors.torch import load_file, safe_open
         state_dict = load_file(vae_path)
     else:
         state_dict = torch.load(vae_path, map_location="cpu")
@@ -178,7 +183,7 @@ else:
     clip_image_processor = None
 
 # Get Scheduler
-Chosen_Scheduler = {
+Chosen_Scheduler = scheduler_dict = {
     "Flow": FlowMatchEulerDiscreteScheduler,
     "Flow_Unipc": FlowUniPCMultistepScheduler,
     "Flow_DPM++": FlowDPMSolverMultistepScheduler,
@@ -223,10 +228,25 @@ if compile_dit:
         pipeline.transformer.blocks[i] = torch.compile(pipeline.transformer.blocks[i])
     print("Add Compile")
 
-# Quantize (when the mode carries an "_and_<quant>" suffix) and then place the pipeline.
-# The order lives inside the helper: quantization has to happen before the offload hooks
-# are installed.
-apply_gpu_memory_mode(pipeline, GPU_memory_mode, device, weight_dtype)
+if GPU_memory_mode == "sequential_cpu_offload":
+    replace_parameters_by_name(transformer, ["modulation",], device=device)
+    transformer.freqs = transformer.freqs.to(device=device)
+    pipeline.enable_sequential_cpu_offload(device=device)
+elif GPU_memory_mode == "model_group_offload":
+    register_auto_device_hook(pipeline.transformer)
+    safe_enable_group_offload(pipeline, onload_device=device, offload_device="cpu", offload_type="leaf_level", use_stream=True)
+elif GPU_memory_mode == "model_cpu_offload_and_qfloat8":
+    convert_model_weight_to_float8(transformer, exclude_module_name=["modulation",], device=device)
+    convert_weight_dtype_wrapper(transformer, weight_dtype)
+    pipeline.enable_model_cpu_offload(device=device)
+elif GPU_memory_mode == "model_cpu_offload":
+    pipeline.enable_model_cpu_offload(device=device)
+elif GPU_memory_mode == "model_full_load_and_qfloat8":
+    convert_model_weight_to_float8(transformer, exclude_module_name=["modulation",], device=device)
+    convert_weight_dtype_wrapper(transformer, weight_dtype)
+    pipeline.to(device=device)
+else:
+    pipeline.to(device=device)
 
 coefficients = get_teacache_coefficients(model_name) if enable_teacache else None
 if coefficients is not None:

examples/wan2.1_fun/predict_v2v_control.py

@@ -18,12 +18,18 @@ from videox_fun.dist import set_multi_gpus_devices, shard_model
 from videox_fun.models import (AutoencoderKLWan, AutoTokenizer, CLIPModel,
                                WanT5EncoderModel, WanTransformer3DModel)
 from videox_fun.models.cache_utils import get_teacache_coefficients
-from videox_fun.pipeline import WanFunControlPipeline
-from videox_fun.utils import (FlowDPMSolverMultistepScheduler,
-                              FlowUniPCMultistepScheduler,
-                              apply_gpu_memory_mode, filter_kwargs,
-                              get_image_latent, get_video_to_video_latent,
-                              merge_lora, save_videos_grid, unmerge_lora)
+from videox_fun.pipeline import WanFunControlPipeline, WanPipeline
+from videox_fun.utils import (register_auto_device_hook,
+                              safe_enable_group_offload)
+from videox_fun.utils.fm_solvers import FlowDPMSolverMultistepScheduler
+from videox_fun.utils.fm_solvers_unipc import FlowUniPCMultistepScheduler
+from videox_fun.utils.fp8_optimization import (convert_model_weight_to_float8,
+                                               convert_weight_dtype_wrapper,
+                                               replace_parameters_by_name)
+from videox_fun.utils.lora_utils import merge_lora, unmerge_lora
+from videox_fun.utils.utils import (filter_kwargs, get_image_latent,
+                                    get_video_to_video_latent,
+                                    save_videos_grid)
 
 # GPU memory mode, which can be chosen in [model_full_load, model_full_load_and_qfloat8, model_cpu_offload, model_cpu_offload_and_qfloat8, model_group_offload, sequential_cpu_offload].
 # model_full_load means that the entire model will be moved to the GPU.
@@ -104,7 +110,7 @@ video_length        = 49
 fps                 = 16
 
 # Use torch.float16 if GPU does not support torch.bfloat16
-# Some graphics cards, such as v100, 2080ti, do not support torch.bfloat16
+# ome graphics cards, such as v100, 2080ti, do not support torch.bfloat16
 weight_dtype            = torch.bfloat16
 control_video           = "asset/pose.mp4"
 control_camera_txt      = None
@@ -139,7 +145,7 @@ transformer = WanTransformer3DModel.from_pretrained(
 if transformer_path is not None:
     print(f"From checkpoint: {transformer_path}")
     if transformer_path.endswith("safetensors"):
-        from safetensors.torch import load_file
+        from safetensors.torch import load_file, safe_open
         state_dict = load_file(transformer_path)
     else:
         state_dict = torch.load(transformer_path, map_location="cpu")
@@ -157,7 +163,7 @@ vae = AutoencoderKLWan.from_pretrained(
 if vae_path is not None:
     print(f"From checkpoint: {vae_path}")
     if vae_path.endswith("safetensors"):
-        from safetensors.torch import load_file
+        from safetensors.torch import load_file, safe_open
         state_dict = load_file(vae_path)
     else:
         state_dict = torch.load(vae_path, map_location="cpu")
@@ -187,7 +193,7 @@ clip_image_encoder = CLIPModel.from_pretrained(
 clip_image_encoder = clip_image_encoder.eval()
 
 # Get Scheduler
-Chosen_Scheduler = {
+Chosen_Scheduler = scheduler_dict = {
     "Flow": FlowMatchEulerDiscreteScheduler,
     "Flow_Unipc": FlowUniPCMultistepScheduler,
     "Flow_DPM++": FlowDPMSolverMultistepScheduler,
@@ -224,10 +230,25 @@ if compile_dit:
         pipeline.transformer.blocks[i] = torch.compile(pipeline.transformer.blocks[i])
     print("Add Compile")
 
-# Quantize (when the mode carries an "_and_<quant>" suffix) and then place the pipeline.
-# The order lives inside the helper: quantization has to happen before the offload hooks
-# are installed.
-apply_gpu_memory_mode(pipeline, GPU_memory_mode, device, weight_dtype)
+if GPU_memory_mode == "sequential_cpu_offload":
+    replace_parameters_by_name(transformer, ["modulation",], device=device)
+    transformer.freqs = transformer.freqs.to(device=device)
+    pipeline.enable_sequential_cpu_offload(device=device)
+elif GPU_memory_mode == "model_group_offload":
+    register_auto_device_hook(pipeline.transformer)
+    safe_enable_group_offload(pipeline, onload_device=device, offload_device="cpu", offload_type="leaf_level", use_stream=True)
+elif GPU_memory_mode == "model_cpu_offload_and_qfloat8":
+    convert_model_weight_to_float8(transformer, exclude_module_name=["modulation",], device=device)
+    convert_weight_dtype_wrapper(transformer, weight_dtype)
+    pipeline.enable_model_cpu_offload(device=device)
+elif GPU_memory_mode == "model_cpu_offload":
+    pipeline.enable_model_cpu_offload(device=device)
+elif GPU_memory_mode == "model_full_load_and_qfloat8":
+    convert_model_weight_to_float8(transformer, exclude_module_name=["modulation",], device=device)
+    convert_weight_dtype_wrapper(transformer, weight_dtype)
+    pipeline.to(device=device)
+else:
+    pipeline.to(device=device)
 
 coefficients = get_teacache_coefficients(model_name) if enable_teacache else None
 if coefficients is not None:

examples/wan2.1_fun/predict_v2v_control_camera.py

@@ -18,12 +18,19 @@ from videox_fun.dist import set_multi_gpus_devices, shard_model
 from videox_fun.models import (AutoencoderKLWan, AutoTokenizer, CLIPModel,
                                WanT5EncoderModel, WanTransformer3DModel)
 from videox_fun.models.cache_utils import get_teacache_coefficients
-from videox_fun.pipeline import WanFunControlPipeline
-from videox_fun.utils import (FlowDPMSolverMultistepScheduler,
-                              FlowUniPCMultistepScheduler,
-                              apply_gpu_memory_mode, filter_kwargs,
-                              get_image_latent, get_video_to_video_latent,
-                              merge_lora, save_videos_grid, unmerge_lora)
+from videox_fun.pipeline import WanFunControlPipeline, WanPipeline
+from videox_fun.utils import (register_auto_device_hook,
+                              safe_enable_group_offload)
+from videox_fun.utils.fm_solvers import FlowDPMSolverMultistepScheduler
+from videox_fun.utils.fm_solvers_unipc import FlowUniPCMultistepScheduler
+from videox_fun.utils.fp8_optimization import (convert_model_weight_to_float8,
+                                               convert_weight_dtype_wrapper,
+                                               replace_parameters_by_name)
+from videox_fun.utils.lora_utils import merge_lora, unmerge_lora
+from videox_fun.utils.utils import (filter_kwargs, get_image_latent,
+                                    get_image_to_video_latent,
+                                    get_video_to_video_latent,
+                                    save_videos_grid)
 
 # GPU memory mode, which can be chosen in [model_full_load, model_full_load_and_qfloat8, model_cpu_offload, model_cpu_offload_and_qfloat8, model_group_offload, sequential_cpu_offload].
 # model_full_load means that the entire model will be moved to the GPU.
@@ -104,7 +111,7 @@ video_length        = 81
 fps                 = 16
 
 # Use torch.float16 if GPU does not support torch.bfloat16
-# Some graphics cards, such as v100, 2080ti, do not support torch.bfloat16
+# ome graphics cards, such as v100, 2080ti, do not support torch.bfloat16
 weight_dtype            = torch.bfloat16
 control_video           = None
 control_camera_txt      = "asset/Pan_Left.txt"
@@ -139,7 +146,7 @@ transformer = WanTransformer3DModel.from_pretrained(
 if transformer_path is not None:
     print(f"From checkpoint: {transformer_path}")
     if transformer_path.endswith("safetensors"):
-        from safetensors.torch import load_file
+        from safetensors.torch import load_file, safe_open
         state_dict = load_file(transformer_path)
     else:
         state_dict = torch.load(transformer_path, map_location="cpu")
@@ -157,7 +164,7 @@ vae = AutoencoderKLWan.from_pretrained(
 if vae_path is not None:
     print(f"From checkpoint: {vae_path}")
     if vae_path.endswith("safetensors"):
-        from safetensors.torch import load_file
+        from safetensors.torch import load_file, safe_open
         state_dict = load_file(vae_path)
     else:
         state_dict = torch.load(vae_path, map_location="cpu")
@@ -187,7 +194,7 @@ clip_image_encoder = CLIPModel.from_pretrained(
 clip_image_encoder = clip_image_encoder.eval()
 
 # Get Scheduler
-Chosen_Scheduler = {
+Chosen_Scheduler = scheduler_dict = {
     "Flow": FlowMatchEulerDiscreteScheduler,
     "Flow_Unipc": FlowUniPCMultistepScheduler,
     "Flow_DPM++": FlowDPMSolverMultistepScheduler,
@@ -224,10 +231,25 @@ if compile_dit:
         pipeline.transformer.blocks[i] = torch.compile(pipeline.transformer.blocks[i])
     print("Add Compile")
 
-# Quantize (when the mode carries an "_and_<quant>" suffix) and then place the pipeline.
-# The order lives inside the helper: quantization has to happen before the offload hooks
-# are installed.
-apply_gpu_memory_mode(pipeline, GPU_memory_mode, device, weight_dtype)
+if GPU_memory_mode == "sequential_cpu_offload":
+    replace_parameters_by_name(transformer, ["modulation",], device=device)
+    transformer.freqs = transformer.freqs.to(device=device)
+    pipeline.enable_sequential_cpu_offload(device=device)
+elif GPU_memory_mode == "model_group_offload":
+    register_auto_device_hook(pipeline.transformer)
+    safe_enable_group_offload(pipeline, onload_device=device, offload_device="cpu", offload_type="leaf_level", use_stream=True)
+elif GPU_memory_mode == "model_cpu_offload_and_qfloat8":
+    convert_model_weight_to_float8(transformer, exclude_module_name=["modulation",], device=device)
+    convert_weight_dtype_wrapper(transformer, weight_dtype)
+    pipeline.enable_model_cpu_offload(device=device)
+elif GPU_memory_mode == "model_cpu_offload":
+    pipeline.enable_model_cpu_offload(device=device)
+elif GPU_memory_mode == "model_full_load_and_qfloat8":
+    convert_model_weight_to_float8(transformer, exclude_module_name=["modulation",], device=device)
+    convert_weight_dtype_wrapper(transformer, weight_dtype)
+    pipeline.to(device=device)
+else:
+    pipeline.to(device=device)
 
 coefficients = get_teacache_coefficients(model_name) if enable_teacache else None
 if coefficients is not None:

examples/wan2.1_fun/predict_v2v_control_ref.py

@@ -18,12 +18,19 @@ from videox_fun.dist import set_multi_gpus_devices, shard_model
 from videox_fun.models import (AutoencoderKLWan, AutoTokenizer, CLIPModel,
                                WanT5EncoderModel, WanTransformer3DModel)
 from videox_fun.models.cache_utils import get_teacache_coefficients
-from videox_fun.pipeline import WanFunControlPipeline
-from videox_fun.utils import (FlowDPMSolverMultistepScheduler,
-                              FlowUniPCMultistepScheduler,
-                              apply_gpu_memory_mode, filter_kwargs,
-                              get_image_latent, get_video_to_video_latent,
-                              merge_lora, save_videos_grid, unmerge_lora)
+from videox_fun.pipeline import WanFunControlPipeline, WanPipeline
+from videox_fun.utils import (register_auto_device_hook,
+                              safe_enable_group_offload)
+from videox_fun.utils.fm_solvers import FlowDPMSolverMultistepScheduler
+from videox_fun.utils.fm_solvers_unipc import FlowUniPCMultistepScheduler
+from videox_fun.utils.fp8_optimization import (convert_model_weight_to_float8,
+                                               convert_weight_dtype_wrapper,
+                                               replace_parameters_by_name)
+from videox_fun.utils.lora_utils import merge_lora, unmerge_lora
+from videox_fun.utils.utils import (filter_kwargs, get_image_latent,
+                                    get_image_to_video_latent,
+                                    get_video_to_video_latent,
+                                    save_videos_grid)
 
 # GPU memory mode, which can be chosen in [model_full_load, model_full_load_and_qfloat8, model_cpu_offload, model_cpu_offload_and_qfloat8, model_group_offload, sequential_cpu_offload].
 # model_full_load means that the entire model will be moved to the GPU.
@@ -104,7 +111,7 @@ video_length        = 49
 fps                 = 16
 
 # Use torch.float16 if GPU does not support torch.bfloat16
-# Some graphics cards, such as v100, 2080ti, do not support torch.bfloat16
+# ome graphics cards, such as v100, 2080ti, do not support torch.bfloat16
 weight_dtype            = torch.bfloat16
 control_video           = "asset/pose.mp4"
 control_camera_txt      = None
@@ -139,7 +146,7 @@ transformer = WanTransformer3DModel.from_pretrained(
 if transformer_path is not None:
     print(f"From checkpoint: {transformer_path}")
     if transformer_path.endswith("safetensors"):
-        from safetensors.torch import load_file
+        from safetensors.torch import load_file, safe_open
         state_dict = load_file(transformer_path)
     else:
         state_dict = torch.load(transformer_path, map_location="cpu")
@@ -157,7 +164,7 @@ vae = AutoencoderKLWan.from_pretrained(
 if vae_path is not None:
     print(f"From checkpoint: {vae_path}")
     if vae_path.endswith("safetensors"):
-        from safetensors.torch import load_file
+        from safetensors.torch import load_file, safe_open
         state_dict = load_file(vae_path)
     else:
         state_dict = torch.load(vae_path, map_location="cpu")
@@ -187,7 +194,7 @@ clip_image_encoder = CLIPModel.from_pretrained(
 clip_image_encoder = clip_image_encoder.eval()
 
 # Get Scheduler
-Chosen_Scheduler = {
+Chosen_Scheduler = scheduler_dict = {
     "Flow": FlowMatchEulerDiscreteScheduler,
     "Flow_Unipc": FlowUniPCMultistepScheduler,
     "Flow_DPM++": FlowDPMSolverMultistepScheduler,
@@ -224,10 +231,25 @@ if compile_dit:
         pipeline.transformer.blocks[i] = torch.compile(pipeline.transformer.blocks[i])
     print("Add Compile")
 
-# Quantize (when the mode carries an "_and_<quant>" suffix) and then place the pipeline.
-# The order lives inside the helper: quantization has to happen before the offload hooks
-# are installed.
-apply_gpu_memory_mode(pipeline, GPU_memory_mode, device, weight_dtype)
+if GPU_memory_mode == "sequential_cpu_offload":
+    replace_parameters_by_name(transformer, ["modulation",], device=device)
+    transformer.freqs = transformer.freqs.to(device=device)
+    pipeline.enable_sequential_cpu_offload(device=device)
+elif GPU_memory_mode == "model_group_offload":
+    register_auto_device_hook(pipeline.transformer)
+    safe_enable_group_offload(pipeline, onload_device=device, offload_device="cpu", offload_type="leaf_level", use_stream=True)
+elif GPU_memory_mode == "model_cpu_offload_and_qfloat8":
+    convert_model_weight_to_float8(transformer, exclude_module_name=["modulation",], device=device)
+    convert_weight_dtype_wrapper(transformer, weight_dtype)
+    pipeline.enable_model_cpu_offload(device=device)
+elif GPU_memory_mode == "model_cpu_offload":
+    pipeline.enable_model_cpu_offload(device=device)
+elif GPU_memory_mode == "model_full_load_and_qfloat8":
+    convert_model_weight_to_float8(transformer, exclude_module_name=["modulation",], device=device)
+    convert_weight_dtype_wrapper(transformer, weight_dtype)
+    pipeline.to(device=device)
+else:
+    pipeline.to(device=device)
 
 coefficients = get_teacache_coefficients(model_name) if enable_teacache else None
 if coefficients is not None:

examples/wan2.1_self_forcing/predict_t2v.py

@@ -17,10 +17,16 @@ from videox_fun.models import (AutoencoderKLWan, AutoTokenizer,
                                WanT5EncoderModel,
                                WanTransformer3DModel_SelfForcing)
 from videox_fun.pipeline import WanSelfForcingPipeline
-from videox_fun.utils import (FlowDPMSolverMultistepScheduler,
-                              FlowUniPCMultistepScheduler,
-                              apply_gpu_memory_mode, filter_kwargs, merge_lora,
-                              save_videos_grid, unmerge_lora)
+from videox_fun.utils import (register_auto_device_hook,
+                              safe_enable_group_offload)
+from videox_fun.utils.fm_solvers import FlowDPMSolverMultistepScheduler
+from videox_fun.utils.fm_solvers_unipc import FlowUniPCMultistepScheduler
+from videox_fun.utils.fp8_optimization import (convert_model_weight_to_float8,
+                                               convert_weight_dtype_wrapper,
+                                               replace_parameters_by_name)
+from videox_fun.utils.lora_utils import merge_lora, unmerge_lora
+from videox_fun.utils.utils import (filter_kwargs, get_image_to_video_latent,
+                                    save_videos_grid)
 
 # GPU memory mode, which can be chosen in [model_full_load, model_full_load_and_qfloat8, model_cpu_offload, model_cpu_offload_and_qfloat8, model_group_offload, sequential_cpu_offload].
 # model_full_load means that the entire model will be moved to the GPU.
@@ -84,7 +90,7 @@ independent_first_frame = False
 context_noise           = 0.0
 
 # Use torch.float16 if GPU does not support torch.bfloat16
-# Some graphics cards, such as v100, 2080ti, do not support torch.bfloat16
+# ome graphics cards, such as v100, 2080ti, do not support torch.bfloat16
 weight_dtype        = torch.bfloat16
 prompt              = "A stylish woman walks down a Tokyo street filled with warm glowing neon and animated city signage. She wears a black leather jacket, a long red dress, and black boots, and carries a black purse. She wears sunglasses and red lipstick. She walks confidently and casually. The street is damp and reflective, creating a mirror effect of the colorful lights. Many pedestrians walk about."
 negative_prompt     = "色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走"
@@ -111,7 +117,7 @@ transformer = WanTransformer3DModel_SelfForcing.from_pretrained(
 if transformer_path is not None:
     print(f"From checkpoint: {transformer_path}")
     if transformer_path.endswith("safetensors"):
-        from safetensors.torch import load_file
+        from safetensors.torch import load_file, safe_open
         state_dict = load_file(transformer_path)
     else:
         state_dict = torch.load(transformer_path, map_location="cpu")
@@ -134,7 +140,7 @@ vae = AutoencoderKLWan.from_pretrained(
 if vae_path is not None:
     print(f"From checkpoint: {vae_path}")
     if vae_path.endswith("safetensors"):
-        from safetensors.torch import load_file
+        from safetensors.torch import load_file, safe_open
         state_dict = load_file(vae_path)
     else:
         state_dict = torch.load(vae_path, map_location="cpu")
@@ -157,7 +163,7 @@ text_encoder = WanT5EncoderModel.from_pretrained(
 )
 
 # Get Scheduler
-Chosen_Scheduler = {
+Chosen_Scheduler = scheduler_dict = {
     "Flow": FlowMatchEulerDiscreteScheduler,
     "Flow_Unipc": FlowUniPCMultistepScheduler,
     "Flow_DPM++": FlowDPMSolverMultistepScheduler,
@@ -194,10 +200,25 @@ if compile_dit:
         pipeline.transformer.blocks[i] = torch.compile(pipeline.transformer.blocks[i])
     print("Add Compile")
 
-# Quantize (when the mode carries an "_and_<quant>" suffix) and then place the pipeline.
-# The order lives inside the helper: quantization has to happen before the offload hooks
-# are installed.
-apply_gpu_memory_mode(pipeline, GPU_memory_mode, device, weight_dtype)
+if GPU_memory_mode == "sequential_cpu_offload":
+    replace_parameters_by_name(transformer, ["modulation",], device=device)
+    transformer.freqs = transformer.freqs.to(device=device)
+    pipeline.enable_sequential_cpu_offload(device=device)
+elif GPU_memory_mode == "model_group_offload":
+    register_auto_device_hook(pipeline.transformer)
+    safe_enable_group_offload(pipeline, onload_device=device, offload_device="cpu", offload_type="leaf_level", use_stream=True)
+elif GPU_memory_mode == "model_cpu_offload_and_qfloat8":
+    convert_model_weight_to_float8(transformer, exclude_module_name=["modulation",], device=device)
+    convert_weight_dtype_wrapper(transformer, weight_dtype)
+    pipeline.enable_model_cpu_offload(device=device)
+elif GPU_memory_mode == "model_cpu_offload":
+    pipeline.enable_model_cpu_offload(device=device)
+elif GPU_memory_mode == "model_full_load_and_qfloat8":
+    convert_model_weight_to_float8(transformer, exclude_module_name=["modulation",], device=device)
+    convert_weight_dtype_wrapper(transformer, weight_dtype)
+    pipeline.to(device=device)
+else:
+    pipeline.to(device=device)
 
 generator = torch.Generator(device=device).manual_seed(seed)
 

examples/wan2.1_self_forcing/predict_t2v_forcing_kv.py

@@ -18,10 +18,16 @@ from videox_fun.models import (AutoencoderKLWan, AutoTokenizer,
                                WanT5EncoderModel,
                                WanTransformer3DModel_SelfForcing)
 from videox_fun.pipeline import WanSelfForcingPipeline
-from videox_fun.utils import (FlowDPMSolverMultistepScheduler,
-                              FlowUniPCMultistepScheduler,
-                              apply_gpu_memory_mode, filter_kwargs, merge_lora,
-                              save_videos_grid, unmerge_lora)
+from videox_fun.utils import (register_auto_device_hook,
+                              safe_enable_group_offload)
+from videox_fun.utils.fm_solvers import FlowDPMSolverMultistepScheduler
+from videox_fun.utils.fm_solvers_unipc import FlowUniPCMultistepScheduler
+from videox_fun.utils.fp8_optimization import (convert_model_weight_to_float8,
+                                               convert_weight_dtype_wrapper,
+                                               replace_parameters_by_name)
+from videox_fun.utils.lora_utils import merge_lora, unmerge_lora
+from videox_fun.utils.utils import (filter_kwargs, get_image_to_video_latent,
+                                    save_videos_grid)
 
 # GPU memory mode, which can be chosen in [model_full_load, model_full_load_and_qfloat8, model_cpu_offload, model_cpu_offload_and_qfloat8, model_group_offload, sequential_cpu_offload].
 # model_full_load means that the entire model will be moved to the GPU.
@@ -109,7 +115,7 @@ forcing_kv_num_frame_patch         = 6    # token segments per latent frame
 forcing_kv_sim_retention_ratio     = 0.33 # fraction of candidate segments kept
 
 # Use torch.float16 if GPU does not support torch.bfloat16
-# Some graphics cards, such as v100, 2080ti, do not support torch.bfloat16
+# ome graphics cards, such as v100, 2080ti, do not support torch.bfloat16
 weight_dtype        = torch.bfloat16
 prompts = [
     "A stylish woman walks down a Tokyo street filled with warm glowing neon and animated city signage. She wears a black leather jacket, a long red dress, and black boots, and carries a black purse. She wears sunglasses and red lipstick. She walks confidently and casually. The street is damp and reflective, creating a mirror effect of the colorful lights. Many pedestrians walk about.",
@@ -139,7 +145,7 @@ transformer = WanTransformer3DModel_SelfForcing.from_pretrained(
 if transformer_path is not None:
     print(f"From checkpoint: {transformer_path}")
     if transformer_path.endswith("safetensors"):
-        from safetensors.torch import load_file
+        from safetensors.torch import load_file, safe_open
         state_dict = load_file(transformer_path)
     else:
         state_dict = torch.load(transformer_path, map_location="cpu")
@@ -164,7 +170,7 @@ vae = AutoencoderKLWan.from_pretrained(
 if vae_path is not None:
     print(f"From checkpoint: {vae_path}")
     if vae_path.endswith("safetensors"):
-        from safetensors.torch import load_file
+        from safetensors.torch import load_file, safe_open
         state_dict = load_file(vae_path)
     else:
         state_dict = torch.load(vae_path, map_location="cpu")
@@ -187,7 +193,7 @@ text_encoder = WanT5EncoderModel.from_pretrained(
 )
 
 # Get Scheduler
-Chosen_Scheduler = {
+Chosen_Scheduler = scheduler_dict = {
     "Flow": FlowMatchEulerDiscreteScheduler,
     "Flow_Unipc": FlowUniPCMultistepScheduler,
     "Flow_DPM++": FlowDPMSolverMultistepScheduler,
@@ -224,10 +230,25 @@ if compile_dit:
         pipeline.transformer.blocks[i] = torch.compile(pipeline.transformer.blocks[i])
     print("Add Compile")
 
-# Quantize (when the mode carries an "_and_<quant>" suffix) and then place the pipeline.
-# The order lives inside the helper: quantization has to happen before the offload hooks
-# are installed.
-apply_gpu_memory_mode(pipeline, GPU_memory_mode, device, weight_dtype)
+if GPU_memory_mode == "sequential_cpu_offload":
+    replace_parameters_by_name(transformer, ["modulation",], device=device)
+    transformer.freqs = transformer.freqs.to(device=device)
+    pipeline.enable_sequential_cpu_offload(device=device)
+elif GPU_memory_mode == "model_group_offload":
+    register_auto_device_hook(pipeline.transformer)
+    safe_enable_group_offload(pipeline, onload_device=device, offload_device="cpu", offload_type="leaf_level", use_stream=True)
+elif GPU_memory_mode == "model_cpu_offload_and_qfloat8":
+    convert_model_weight_to_float8(transformer, exclude_module_name=["modulation",], device=device)
+    convert_weight_dtype_wrapper(transformer, weight_dtype)
+    pipeline.enable_model_cpu_offload(device=device)
+elif GPU_memory_mode == "model_cpu_offload":
+    pipeline.enable_model_cpu_offload(device=device)
+elif GPU_memory_mode == "model_full_load_and_qfloat8":
+    convert_model_weight_to_float8(transformer, exclude_module_name=["modulation",], device=device)
+    convert_weight_dtype_wrapper(transformer, weight_dtype)
+    pipeline.to(device=device)
+else:
+    pipeline.to(device=device)
 
 if forcing_kv_enable:
     assert forcing_kv_head_profile is not None, \

examples/wan2.1_self_forcing/predict_t2v_stream.py

@@ -17,11 +17,17 @@ from videox_fun.models import (AutoencoderKLWan, AutoTokenizer,
                                WanT5EncoderModel,
                                WanTransformer3DModel_SelfForcing)
 from videox_fun.pipeline import WanSelfForcingPipeline
-from videox_fun.utils import (FlowDPMSolverMultistepScheduler,
-                              FlowUniPCMultistepScheduler, SegmentVideoSaver,
-                              StreamVideoSaver, apply_gpu_memory_mode,
-                              filter_kwargs, merge_lora, save_videos_grid,
-                              unmerge_lora)
+from videox_fun.utils import (register_auto_device_hook,
+                              safe_enable_group_offload)
+from videox_fun.utils.fm_solvers import FlowDPMSolverMultistepScheduler
+from videox_fun.utils.fm_solvers_unipc import FlowUniPCMultistepScheduler
+from videox_fun.utils.fp8_optimization import (convert_model_weight_to_float8,
+                                               convert_weight_dtype_wrapper,
+                                               replace_parameters_by_name)
+from videox_fun.utils.lora_utils import merge_lora, unmerge_lora
+from videox_fun.utils.utils import (filter_kwargs, get_image_to_video_latent,
+                                    save_videos_grid, StreamVideoSaver,
+                                    SegmentVideoSaver)
 
 # GPU memory mode, which can be chosen in [model_full_load, model_full_load_and_qfloat8, model_cpu_offload, model_cpu_offload_and_qfloat8, model_group_offload, sequential_cpu_offload].
 # model_full_load means that the entire model will be moved to the GPU.
@@ -97,7 +103,7 @@ streaming               = True
 save_mode               = "segments"
 
 # Use torch.float16 if GPU does not support torch.bfloat16
-# Some graphics cards, such as v100, 2080ti, do not support torch.bfloat16
+# ome graphics cards, such as v100, 2080ti, do not support torch.bfloat16
 weight_dtype        = torch.bfloat16
 prompt              = "A stylish woman walks down a Tokyo street filled with warm glowing neon and animated city signage. She wears a black leather jacket, a long red dress, and black boots, and carries a black purse. She wears sunglasses and red lipstick. She walks confidently and casually. The street is damp and reflective, creating a mirror effect of the colorful lights. Many pedestrians walk about."
 negative_prompt     = "色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走"
@@ -124,7 +130,7 @@ transformer = WanTransformer3DModel_SelfForcing.from_pretrained(
 if transformer_path is not None:
     print(f"From checkpoint: {transformer_path}")
     if transformer_path.endswith("safetensors"):
-        from safetensors.torch import load_file
+        from safetensors.torch import load_file, safe_open
         state_dict = load_file(transformer_path)
     else:
         state_dict = torch.load(transformer_path, map_location="cpu")
@@ -147,7 +153,7 @@ vae = AutoencoderKLWan.from_pretrained(
 if vae_path is not None:
     print(f"From checkpoint: {vae_path}")
     if vae_path.endswith("safetensors"):
-        from safetensors.torch import load_file
+        from safetensors.torch import load_file, safe_open
         state_dict = load_file(vae_path)
     else:
         state_dict = torch.load(vae_path, map_location="cpu")
@@ -170,7 +176,7 @@ text_encoder = WanT5EncoderModel.from_pretrained(
 )
 
 # Get Scheduler
-Chosen_Scheduler = {
+Chosen_Scheduler = scheduler_dict = {
     "Flow": FlowMatchEulerDiscreteScheduler,
     "Flow_Unipc": FlowUniPCMultistepScheduler,
     "Flow_DPM++": FlowDPMSolverMultistepScheduler,
@@ -207,10 +213,25 @@ if compile_dit:
         pipeline.transformer.blocks[i] = torch.compile(pipeline.transformer.blocks[i])
     print("Add Compile")
 
-# Quantize (when the mode carries an "_and_<quant>" suffix) and then place the pipeline.
-# The order lives inside the helper: quantization has to happen before the offload hooks
-# are installed.
-apply_gpu_memory_mode(pipeline, GPU_memory_mode, device, weight_dtype)
+if GPU_memory_mode == "sequential_cpu_offload":
+    replace_parameters_by_name(transformer, ["modulation",], device=device)
+    transformer.freqs = transformer.freqs.to(device=device)
+    pipeline.enable_sequential_cpu_offload(device=device)
+elif GPU_memory_mode == "model_group_offload":
+    register_auto_device_hook(pipeline.transformer)
+    safe_enable_group_offload(pipeline, onload_device=device, offload_device="cpu", offload_type="leaf_level", use_stream=True)
+elif GPU_memory_mode == "model_cpu_offload_and_qfloat8":
+    convert_model_weight_to_float8(transformer, exclude_module_name=["modulation",], device=device)
+    convert_weight_dtype_wrapper(transformer, weight_dtype)
+    pipeline.enable_model_cpu_offload(device=device)
+elif GPU_memory_mode == "model_cpu_offload":
+    pipeline.enable_model_cpu_offload(device=device)
+elif GPU_memory_mode == "model_full_load_and_qfloat8":
+    convert_model_weight_to_float8(transformer, exclude_module_name=["modulation",], device=device)
+    convert_weight_dtype_wrapper(transformer, weight_dtype)
+    pipeline.to(device=device)
+else:
+    pipeline.to(device=device)
 
 # Only the main process (rank 0, or single-GPU) writes files to disk.
 if ulysses_degree * ring_degree > 1:

examples/wan2.1_self_forcing/profile_forcing_kv_heads.py

@@ -32,7 +32,7 @@ from videox_fun.models import (AutoencoderKLWan, AutoTokenizer,
                                WanT5EncoderModel,
                                WanTransformer3DModel_SelfForcing)
 from videox_fun.pipeline import WanSelfForcingPipeline
-from videox_fun.utils import filter_kwargs
+from videox_fun.utils.utils import filter_kwargs
 
 # Config and model path
 config_path         = "config/wan2.1/wan_civitai.yaml"
@@ -90,7 +90,7 @@ transformer = WanTransformer3DModel_SelfForcing.from_pretrained(
 if transformer_path is not None:
     print(f"From checkpoint: {transformer_path}")
     if transformer_path.endswith("safetensors"):
-        from safetensors.torch import load_file
+        from safetensors.torch import load_file, safe_open
         state_dict = load_file(transformer_path)
     else:
         state_dict = torch.load(transformer_path, map_location="cpu")

examples/wan2.1_vace/predict_i2v.py

@@ -13,17 +13,24 @@ project_roots = [os.path.dirname(current_file_path), os.path.dirname(os.path.dir
 for project_root in project_roots:
     sys.path.insert(0, project_root) if project_root not in sys.path else None
 
+from videox_fun.data import process_pose_file
 from videox_fun.dist import set_multi_gpus_devices, shard_model
 from videox_fun.models import (AutoencoderKLWan, AutoTokenizer,
                                VaceWanTransformer3DModel, WanT5EncoderModel)
 from videox_fun.models.cache_utils import get_teacache_coefficients
-from videox_fun.pipeline import WanVacePipeline
-from videox_fun.utils import (FlowDPMSolverMultistepScheduler,
-                              FlowUniPCMultistepScheduler,
-                              apply_gpu_memory_mode, filter_kwargs,
-                              get_image_latent, get_image_to_video_latent,
-                              get_video_to_video_latent, merge_lora,
-                              save_videos_grid, unmerge_lora)
+from videox_fun.pipeline import WanPipeline, WanVacePipeline
+from videox_fun.utils import (register_auto_device_hook,
+                              safe_enable_group_offload)
+from videox_fun.utils.fm_solvers import FlowDPMSolverMultistepScheduler
+from videox_fun.utils.fm_solvers_unipc import FlowUniPCMultistepScheduler
+from videox_fun.utils.fp8_optimization import (convert_model_weight_to_float8,
+                                               convert_weight_dtype_wrapper,
+                                               replace_parameters_by_name)
+from videox_fun.utils.lora_utils import merge_lora, unmerge_lora
+from videox_fun.utils.utils import (filter_kwargs, get_image_latent,
+                                    get_image_to_video_latent,
+                                    get_video_to_video_latent,
+                                    save_videos_grid)
 
 # GPU memory mode, which can be chosen in [model_full_load, model_full_load_and_qfloat8, model_cpu_offload, model_cpu_offload_and_qfloat8, model_group_offload, sequential_cpu_offload].
 # model_full_load means that the entire model will be moved to the GPU.
@@ -101,7 +108,7 @@ video_length        = 81
 fps                 = 16
 
 # Use torch.float16 if GPU does not support torch.bfloat16
-# Some graphics cards, such as v100, 2080ti, do not support torch.bfloat16
+# ome graphics cards, such as v100, 2080ti, do not support torch.bfloat16
 weight_dtype            = torch.bfloat16
 control_video           = None
 start_image             = "asset/1.png"
@@ -137,7 +144,7 @@ transformer = VaceWanTransformer3DModel.from_pretrained(
 if transformer_path is not None:
     print(f"From checkpoint: {transformer_path}")
     if transformer_path.endswith("safetensors"):
-        from safetensors.torch import load_file
+        from safetensors.torch import load_file, safe_open
         state_dict = load_file(transformer_path)
     else:
         state_dict = torch.load(transformer_path, map_location="cpu")
@@ -155,7 +162,7 @@ vae = AutoencoderKLWan.from_pretrained(
 if vae_path is not None:
     print(f"From checkpoint: {vae_path}")
     if vae_path.endswith("safetensors"):
-        from safetensors.torch import load_file
+        from safetensors.torch import load_file, safe_open
         state_dict = load_file(vae_path)
     else:
         state_dict = torch.load(vae_path, map_location="cpu")
@@ -179,7 +186,7 @@ text_encoder = WanT5EncoderModel.from_pretrained(
 text_encoder = text_encoder.eval()
 
 # Get Scheduler
-Chosen_Scheduler = {
+Chosen_Scheduler = scheduler_dict = {
     "Flow": FlowMatchEulerDiscreteScheduler,
     "Flow_Unipc": FlowUniPCMultistepScheduler,
     "Flow_DPM++": FlowDPMSolverMultistepScheduler,
@@ -215,10 +222,25 @@ if compile_dit:
         pipeline.transformer.blocks[i] = torch.compile(pipeline.transformer.blocks[i])
     print("Add Compile")
 
-# Quantize (when the mode carries an "_and_<quant>" suffix) and then place the pipeline.
-# The order lives inside the helper: quantization has to happen before the offload hooks
-# are installed.
-apply_gpu_memory_mode(pipeline, GPU_memory_mode, device, weight_dtype)
+if GPU_memory_mode == "sequential_cpu_offload":
+    replace_parameters_by_name(transformer, ["modulation",], device=device)
+    transformer.freqs = transformer.freqs.to(device=device)
+    pipeline.enable_sequential_cpu_offload(device=device)
+elif GPU_memory_mode == "model_group_offload":
+    register_auto_device_hook(pipeline.transformer)
+    safe_enable_group_offload(pipeline, onload_device=device, offload_device="cpu", offload_type="leaf_level", use_stream=True)
+elif GPU_memory_mode == "model_cpu_offload_and_qfloat8":
+    convert_model_weight_to_float8(transformer, exclude_module_name=["modulation",], device=device)
+    convert_weight_dtype_wrapper(transformer, weight_dtype)
+    pipeline.enable_model_cpu_offload(device=device)
+elif GPU_memory_mode == "model_cpu_offload":
+    pipeline.enable_model_cpu_offload(device=device)
+elif GPU_memory_mode == "model_full_load_and_qfloat8":
+    convert_model_weight_to_float8(transformer, exclude_module_name=["modulation",], device=device)
+    convert_weight_dtype_wrapper(transformer, weight_dtype)
+    pipeline.to(device=device)
+else:
+    pipeline.to(device=device)
 
 coefficients = get_teacache_coefficients(model_name) if enable_teacache else None
 if coefficients is not None:

examples/wan2.1_vace/predict_s2v.py

@@ -13,17 +13,24 @@ project_roots = [os.path.dirname(current_file_path), os.path.dirname(os.path.dir
 for project_root in project_roots:
     sys.path.insert(0, project_root) if project_root not in sys.path else None
 
+from videox_fun.data import process_pose_file
 from videox_fun.dist import set_multi_gpus_devices, shard_model
 from videox_fun.models import (AutoencoderKLWan, AutoTokenizer,
                                VaceWanTransformer3DModel, WanT5EncoderModel)
 from videox_fun.models.cache_utils import get_teacache_coefficients
-from videox_fun.pipeline import WanVacePipeline
-from videox_fun.utils import (FlowDPMSolverMultistepScheduler,
-                              FlowUniPCMultistepScheduler,
-                              apply_gpu_memory_mode, filter_kwargs,
-                              get_image_latent, get_image_to_video_latent,
-                              get_video_to_video_latent, merge_lora,
-                              save_videos_grid, unmerge_lora)
+from videox_fun.pipeline import WanPipeline, WanVacePipeline
+from videox_fun.utils import (register_auto_device_hook,
+                              safe_enable_group_offload)
+from videox_fun.utils.fm_solvers import FlowDPMSolverMultistepScheduler
+from videox_fun.utils.fm_solvers_unipc import FlowUniPCMultistepScheduler
+from videox_fun.utils.fp8_optimization import (convert_model_weight_to_float8,
+                                               convert_weight_dtype_wrapper,
+                                               replace_parameters_by_name)
+from videox_fun.utils.lora_utils import merge_lora, unmerge_lora
+from videox_fun.utils.utils import (filter_kwargs, get_image_latent,
+                                    get_image_to_video_latent,
+                                    get_video_to_video_latent,
+                                    save_videos_grid)
 
 # GPU memory mode, which can be chosen in [model_full_load, model_full_load_and_qfloat8, model_cpu_offload, model_cpu_offload_and_qfloat8, model_group_offload, sequential_cpu_offload].
 # model_full_load means that the entire model will be moved to the GPU.
@@ -102,7 +109,7 @@ fps                 = 16
 vace_context_scale  = 1.00
 
 # Use torch.float16 if GPU does not support torch.bfloat16
-# Some graphics cards, such as v100, 2080ti, do not support torch.bfloat16
+# ome graphics cards, such as v100, 2080ti, do not support torch.bfloat16
 weight_dtype            = torch.bfloat16
 control_video           = None
 start_image             = None
@@ -137,7 +144,7 @@ transformer = VaceWanTransformer3DModel.from_pretrained(
 if transformer_path is not None:
     print(f"From checkpoint: {transformer_path}")
     if transformer_path.endswith("safetensors"):
-        from safetensors.torch import load_file
+        from safetensors.torch import load_file, safe_open
         state_dict = load_file(transformer_path)
     else:
         state_dict = torch.load(transformer_path, map_location="cpu")
@@ -155,7 +162,7 @@ vae = AutoencoderKLWan.from_pretrained(
 if vae_path is not None:
     print(f"From checkpoint: {vae_path}")
     if vae_path.endswith("safetensors"):
-        from safetensors.torch import load_file
+        from safetensors.torch import load_file, safe_open
         state_dict = load_file(vae_path)
     else:
         state_dict = torch.load(vae_path, map_location="cpu")
@@ -179,7 +186,7 @@ text_encoder = WanT5EncoderModel.from_pretrained(
 text_encoder = text_encoder.eval()
 
 # Get Scheduler
-Chosen_Scheduler = {
+Chosen_Scheduler = scheduler_dict = {
     "Flow": FlowMatchEulerDiscreteScheduler,
     "Flow_Unipc": FlowUniPCMultistepScheduler,
     "Flow_DPM++": FlowDPMSolverMultistepScheduler,
@@ -215,10 +222,25 @@ if compile_dit:
         pipeline.transformer.blocks[i] = torch.compile(pipeline.transformer.blocks[i])
     print("Add Compile")
 
-# Quantize (when the mode carries an "_and_<quant>" suffix) and then place the pipeline.
-# The order lives inside the helper: quantization has to happen before the offload hooks
-# are installed.
-apply_gpu_memory_mode(pipeline, GPU_memory_mode, device, weight_dtype)
+if GPU_memory_mode == "sequential_cpu_offload":
+    replace_parameters_by_name(transformer, ["modulation",], device=device)
+    transformer.freqs = transformer.freqs.to(device=device)
+    pipeline.enable_sequential_cpu_offload(device=device)
+elif GPU_memory_mode == "model_group_offload":
+    register_auto_device_hook(pipeline.transformer)
+    safe_enable_group_offload(pipeline, onload_device=device, offload_device="cpu", offload_type="leaf_level", use_stream=True)
+elif GPU_memory_mode == "model_cpu_offload_and_qfloat8":
+    convert_model_weight_to_float8(transformer, exclude_module_name=["modulation",], device=device)
+    convert_weight_dtype_wrapper(transformer, weight_dtype)
+    pipeline.enable_model_cpu_offload(device=device)
+elif GPU_memory_mode == "model_cpu_offload":
+    pipeline.enable_model_cpu_offload(device=device)
+elif GPU_memory_mode == "model_full_load_and_qfloat8":
+    convert_model_weight_to_float8(transformer, exclude_module_name=["modulation",], device=device)
+    convert_weight_dtype_wrapper(transformer, weight_dtype)
+    pipeline.to(device=device)
+else:
+    pipeline.to(device=device)
 
 coefficients = get_teacache_coefficients(model_name) if enable_teacache else None
 if coefficients is not None:

examples/wan2.1_vace/predict_v2v_control.py

@@ -13,17 +13,24 @@ project_roots = [os.path.dirname(current_file_path), os.path.dirname(os.path.dir
 for project_root in project_roots:
     sys.path.insert(0, project_root) if project_root not in sys.path else None
 
+from videox_fun.data import process_pose_file
 from videox_fun.dist import set_multi_gpus_devices, shard_model
 from videox_fun.models import (AutoencoderKLWan, AutoTokenizer,
                                VaceWanTransformer3DModel, WanT5EncoderModel)
 from videox_fun.models.cache_utils import get_teacache_coefficients
-from videox_fun.pipeline import WanVacePipeline
-from videox_fun.utils import (FlowDPMSolverMultistepScheduler,
-                              FlowUniPCMultistepScheduler,
-                              apply_gpu_memory_mode, filter_kwargs,
-                              get_image_latent, get_image_to_video_latent,
-                              get_video_to_video_latent, merge_lora,
-                              save_videos_grid, unmerge_lora)
+from videox_fun.pipeline import WanPipeline, WanVacePipeline
+from videox_fun.utils import (register_auto_device_hook,
+                              safe_enable_group_offload)
+from videox_fun.utils.fm_solvers import FlowDPMSolverMultistepScheduler
+from videox_fun.utils.fm_solvers_unipc import FlowUniPCMultistepScheduler
+from videox_fun.utils.fp8_optimization import (convert_model_weight_to_float8,
+                                               convert_weight_dtype_wrapper,
+                                               replace_parameters_by_name)
+from videox_fun.utils.lora_utils import merge_lora, unmerge_lora
+from videox_fun.utils.utils import (filter_kwargs, get_image_latent,
+                                    get_image_to_video_latent,
+                                    get_video_to_video_latent,
+                                    save_videos_grid)
 
 # GPU memory mode, which can be chosen in [model_full_load, model_full_load_and_qfloat8, model_cpu_offload, model_cpu_offload_and_qfloat8, model_group_offload, sequential_cpu_offload].
 # model_full_load means that the entire model will be moved to the GPU.
@@ -102,7 +109,7 @@ fps                 = 16
 vace_context_scale  = 1.00
 
 # Use torch.float16 if GPU does not support torch.bfloat16
-# Some graphics cards, such as v100, 2080ti, do not support torch.bfloat16
+# ome graphics cards, such as v100, 2080ti, do not support torch.bfloat16
 weight_dtype            = torch.bfloat16
 control_video           = "asset/pose.mp4"
 start_image             = None
@@ -137,7 +144,7 @@ transformer = VaceWanTransformer3DModel.from_pretrained(
 if transformer_path is not None:
     print(f"From checkpoint: {transformer_path}")
     if transformer_path.endswith("safetensors"):
-        from safetensors.torch import load_file
+        from safetensors.torch import load_file, safe_open
         state_dict = load_file(transformer_path)
     else:
         state_dict = torch.load(transformer_path, map_location="cpu")
@@ -155,7 +162,7 @@ vae = AutoencoderKLWan.from_pretrained(
 if vae_path is not None:
     print(f"From checkpoint: {vae_path}")
     if vae_path.endswith("safetensors"):
-        from safetensors.torch import load_file
+        from safetensors.torch import load_file, safe_open
         state_dict = load_file(vae_path)
     else:
         state_dict = torch.load(vae_path, map_location="cpu")
@@ -179,7 +186,7 @@ text_encoder = WanT5EncoderModel.from_pretrained(
 text_encoder = text_encoder.eval()
 
 # Get Scheduler
-Chosen_Scheduler = {
+Chosen_Scheduler = scheduler_dict = {
     "Flow": FlowMatchEulerDiscreteScheduler,
     "Flow_Unipc": FlowUniPCMultistepScheduler,
     "Flow_DPM++": FlowDPMSolverMultistepScheduler,
@@ -215,10 +222,25 @@ if compile_dit:
         pipeline.transformer.blocks[i] = torch.compile(pipeline.transformer.blocks[i])
     print("Add Compile")
 
-# Quantize (when the mode carries an "_and_<quant>" suffix) and then place the pipeline.
-# The order lives inside the helper: quantization has to happen before the offload hooks
-# are installed.
-apply_gpu_memory_mode(pipeline, GPU_memory_mode, device, weight_dtype)
+if GPU_memory_mode == "sequential_cpu_offload":
+    replace_parameters_by_name(transformer, ["modulation",], device=device)
+    transformer.freqs = transformer.freqs.to(device=device)
+    pipeline.enable_sequential_cpu_offload(device=device)
+elif GPU_memory_mode == "model_group_offload":
+    register_auto_device_hook(pipeline.transformer)
+    safe_enable_group_offload(pipeline, onload_device=device, offload_device="cpu", offload_type="leaf_level", use_stream=True)
+elif GPU_memory_mode == "model_cpu_offload_and_qfloat8":
+    convert_model_weight_to_float8(transformer, exclude_module_name=["modulation",], device=device)
+    convert_weight_dtype_wrapper(transformer, weight_dtype)
+    pipeline.enable_model_cpu_offload(device=device)
+elif GPU_memory_mode == "model_cpu_offload":
+    pipeline.enable_model_cpu_offload(device=device)
+elif GPU_memory_mode == "model_full_load_and_qfloat8":
+    convert_model_weight_to_float8(transformer, exclude_module_name=["modulation",], device=device)
+    convert_weight_dtype_wrapper(transformer, weight_dtype)
+    pipeline.to(device=device)
+else:
+    pipeline.to(device=device)
 
 coefficients = get_teacache_coefficients(model_name) if enable_teacache else None
 if coefficients is not None:

examples/wan2.2/app.py

@@ -10,7 +10,7 @@ for project_root in project_roots:
     sys.path.insert(0, project_root) if project_root not in sys.path else None
 
 from videox_fun.api.api import (infer_forward_api,
-                                update_diffusion_transformer_api)
+                               update_diffusion_transformer_api)
 from videox_fun.ui.controller import flow_scheduler_dict
 from videox_fun.ui.wan2_2_ui import ui, ui_client, ui_host
 

examples/wan2.2/launch_api.py

@@ -4,6 +4,7 @@ import sys
 import time
 
 import gradio as gr
+import ray
 import torch
 
 current_file_path = os.path.abspath(__file__)
@@ -12,11 +13,10 @@ for project_root in project_roots:
     sys.path.insert(0, project_root) if project_root not in sys.path else None
 
 from videox_fun.api.api_multi_nodes import (MultiNodesEngine,
-                                            multi_nodes_infer_forward_api)
+                                           multi_nodes_infer_forward_api)
 from videox_fun.ui.controller import flow_scheduler_dict
 from videox_fun.ui.wan2_2_ui import Wan2_2_Controller
 
-
 def main():
     parser = argparse.ArgumentParser(description='xDiT HTTP Service')
     parser.add_argument('--world_size', type=int, default=8, help='Number of parallel workers')

examples/wan2.2/predict_animate.py

@@ -19,11 +19,18 @@ from videox_fun.models import (AutoencoderKLWan, AutoencoderKLWan3_8,
                                WanT5EncoderModel)
 from videox_fun.models.cache_utils import get_teacache_coefficients
 from videox_fun.pipeline import Wan2_2AnimatePipeline
-from videox_fun.utils import (FlowDPMSolverMultistepScheduler,
-                              FlowUniPCMultistepScheduler,
-                              apply_gpu_memory_mode, filter_kwargs, get_image,
-                              get_video_to_video_latent, merge_lora,
-                              save_videos_grid, unmerge_lora)
+from videox_fun.utils import (register_auto_device_hook,
+                              safe_enable_group_offload)
+from videox_fun.utils.fm_solvers import FlowDPMSolverMultistepScheduler
+from videox_fun.utils.fm_solvers_unipc import FlowUniPCMultistepScheduler
+from videox_fun.utils.fp8_optimization import (convert_model_weight_to_float8,
+                                               convert_weight_dtype_wrapper,
+                                               replace_parameters_by_name)
+from videox_fun.utils.lora_utils import merge_lora, unmerge_lora
+from videox_fun.utils.utils import (filter_kwargs, get_image,
+                                    get_image_to_video_latent,
+                                    get_video_to_video_latent,
+                                    save_videos_grid)
 
 # GPU memory mode, which can be chosen in [model_full_load, model_full_load_and_qfloat8, model_cpu_offload, model_cpu_offload_and_qfloat8, model_group_offload, sequential_cpu_offload].
 # model_full_load means that the entire model will be moved to the GPU.
@@ -112,7 +119,7 @@ segment_frame_length    = 77
 fps                     = 16
 
 # Use torch.float16 if GPU does not support torch.bfloat16
-# Some graphics cards, such as v100, 2080ti, do not support torch.bfloat16
+# ome graphics cards, such as v100, 2080ti, do not support torch.bfloat16
 weight_dtype        = torch.bfloat16
 prompt              = "视频中的人在做动作"
 negative_prompt     = "色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走"
@@ -147,7 +154,7 @@ else:
 if transformer_path is not None:
     print(f"From checkpoint: {transformer_path}")
     if transformer_path.endswith("safetensors"):
-        from safetensors.torch import load_file
+        from safetensors.torch import load_file, safe_open
         state_dict = load_file(transformer_path)
     else:
         state_dict = torch.load(transformer_path, map_location="cpu")
@@ -160,7 +167,7 @@ if transformer_2 is not None:
     if transformer_high_path is not None:
         print(f"From checkpoint: {transformer_high_path}")
         if transformer_high_path.endswith("safetensors"):
-            from safetensors.torch import load_file
+            from safetensors.torch import load_file, safe_open
             state_dict = load_file(transformer_high_path)
         else:
             state_dict = torch.load(transformer_high_path, map_location="cpu")
@@ -182,7 +189,7 @@ vae = Chosen_AutoencoderKL.from_pretrained(
 if vae_path is not None:
     print(f"From checkpoint: {vae_path}")
     if vae_path.endswith("safetensors"):
-        from safetensors.torch import load_file
+        from safetensors.torch import load_file, safe_open
         state_dict = load_file(vae_path)
     else:
         state_dict = torch.load(vae_path, map_location="cpu")
@@ -212,7 +219,7 @@ clip_image_encoder = CLIPModel.from_pretrained(
 clip_image_encoder = clip_image_encoder.eval()
 
 # Get Scheduler
-Chosen_Scheduler = {
+Chosen_Scheduler = scheduler_dict = {
     "Flow": FlowMatchEulerDiscreteScheduler,
     "Flow_Unipc": FlowUniPCMultistepScheduler,
     "Flow_DPM++": FlowDPMSolverMultistepScheduler,
@@ -257,11 +264,36 @@ if compile_dit:
             pipeline.transformer_2.blocks[i] = torch.compile(pipeline.transformer_2.blocks[i])
     print("Add Compile")
 
-# Quantize (when the mode carries an "_and_<quant>" suffix) and then place the pipeline.
-# The order lives inside the helper: quantization has to happen before the offload hooks
-# are installed, and both transformers of this MoE setup are handled in one call, which
-# is exactly the bookkeeping the old 30-line if/elif chain repeated per script.
-apply_gpu_memory_mode(pipeline, GPU_memory_mode, device, weight_dtype)
+if GPU_memory_mode == "sequential_cpu_offload":
+    replace_parameters_by_name(transformer, ["modulation",], device=device)
+    transformer.freqs = transformer.freqs.to(device=device)
+    if transformer_2 is not None:
+        replace_parameters_by_name(transformer_2, ["modulation",], device=device)
+        transformer_2.freqs = transformer_2.freqs.to(device=device)
+    pipeline.enable_sequential_cpu_offload(device=device)
+elif GPU_memory_mode == "model_group_offload":
+    register_auto_device_hook(pipeline.transformer)
+    if transformer_2 is not None:
+        register_auto_device_hook(pipeline.transformer_2)
+    safe_enable_group_offload(pipeline, onload_device=device, offload_device="cpu", offload_type="leaf_level", use_stream=True)
+elif GPU_memory_mode == "model_cpu_offload_and_qfloat8":
+    convert_model_weight_to_float8(transformer, exclude_module_name=["modulation",], device=device)
+    convert_weight_dtype_wrapper(transformer, weight_dtype)
+    if transformer_2 is not None:
+        convert_model_weight_to_float8(transformer_2, exclude_module_name=["modulation",], device=device)
+        convert_weight_dtype_wrapper(transformer_2, weight_dtype)
+    pipeline.enable_model_cpu_offload(device=device)
+elif GPU_memory_mode == "model_cpu_offload":
+    pipeline.enable_model_cpu_offload(device=device)
+elif GPU_memory_mode == "model_full_load_and_qfloat8":
+    convert_model_weight_to_float8(transformer, exclude_module_name=["modulation",], device=device)
+    convert_weight_dtype_wrapper(transformer, weight_dtype)
+    if transformer_2 is not None:
+        convert_model_weight_to_float8(transformer_2, exclude_module_name=["modulation",], device=device)
+        convert_weight_dtype_wrapper(transformer_2, weight_dtype)
+    pipeline.to(device=device)
+else:
+    pipeline.to(device=device)
 
 coefficients = get_teacache_coefficients(model_name) if enable_teacache else None
 if coefficients is not None:

examples/wan2.2/predict_i2v.py

@@ -14,15 +14,20 @@ for project_root in project_roots:
 
 from videox_fun.dist import set_multi_gpus_devices, shard_model
 from videox_fun.models import (AutoencoderKLWan, AutoencoderKLWan3_8,
-                               AutoTokenizer, Wan2_2Transformer3DModel,
-                               WanT5EncoderModel)
+                               AutoTokenizer, CLIPModel,
+                               Wan2_2Transformer3DModel, WanT5EncoderModel)
 from videox_fun.models.cache_utils import get_teacache_coefficients
 from videox_fun.pipeline import Wan2_2I2VPipeline
-from videox_fun.utils import (FlowDPMSolverMultistepScheduler,
-                              FlowUniPCMultistepScheduler,
-                              apply_gpu_memory_mode, filter_kwargs,
-                              get_image_to_video_latent, merge_lora,
-                              save_videos_grid, unmerge_lora)
+from videox_fun.utils import (register_auto_device_hook,
+                              safe_enable_group_offload)
+from videox_fun.utils.fm_solvers import FlowDPMSolverMultistepScheduler
+from videox_fun.utils.fm_solvers_unipc import FlowUniPCMultistepScheduler
+from videox_fun.utils.fp8_optimization import (convert_model_weight_to_float8,
+                                               convert_weight_dtype_wrapper,
+                                               replace_parameters_by_name)
+from videox_fun.utils.lora_utils import merge_lora, unmerge_lora
+from videox_fun.utils.utils import (filter_kwargs, get_image_to_video_latent,
+                                    save_videos_grid)
 
 # GPU memory mode, which can be chosen in [model_full_load, model_full_load_and_qfloat8, model_cpu_offload, model_cpu_offload_and_qfloat8, model_group_offload, sequential_cpu_offload].
 # model_full_load means that the entire model will be moved to the GPU.
@@ -105,7 +110,7 @@ video_length        = 81
 fps                 = 16
 
 # Use torch.float16 if GPU does not support torch.bfloat16
-# Some graphics cards, such as v100, 2080ti, do not support torch.bfloat16
+# ome graphics cards, such as v100, 2080ti, do not support torch.bfloat16
 weight_dtype            = torch.bfloat16
 # If you want to generate from text, please set the validation_image_start = None and validation_image_end = None
 validation_image_start  = "asset/1.png"
@@ -144,7 +149,7 @@ else:
 if transformer_path is not None:
     print(f"From checkpoint: {transformer_path}")
     if transformer_path.endswith("safetensors"):
-        from safetensors.torch import load_file
+        from safetensors.torch import load_file, safe_open
         state_dict = load_file(transformer_path)
     else:
         state_dict = torch.load(transformer_path, map_location="cpu")
@@ -157,7 +162,7 @@ if transformer_2 is not None:
     if transformer_high_path is not None:
         print(f"From checkpoint: {transformer_high_path}")
         if transformer_high_path.endswith("safetensors"):
-            from safetensors.torch import load_file
+            from safetensors.torch import load_file, safe_open
             state_dict = load_file(transformer_high_path)
         else:
             state_dict = torch.load(transformer_high_path, map_location="cpu")
@@ -179,7 +184,7 @@ vae = Chosen_AutoencoderKL.from_pretrained(
 if vae_path is not None:
     print(f"From checkpoint: {vae_path}")
     if vae_path.endswith("safetensors"):
-        from safetensors.torch import load_file
+        from safetensors.torch import load_file, safe_open
         state_dict = load_file(vae_path)
     else:
         state_dict = torch.load(vae_path, map_location="cpu")
@@ -203,7 +208,7 @@ text_encoder = WanT5EncoderModel.from_pretrained(
 text_encoder = text_encoder.eval()
 
 # Get Scheduler
-Chosen_Scheduler = {
+Chosen_Scheduler = scheduler_dict = {
     "Flow": FlowMatchEulerDiscreteScheduler,
     "Flow_Unipc": FlowUniPCMultistepScheduler,
     "Flow_DPM++": FlowDPMSolverMultistepScheduler,
@@ -247,11 +252,36 @@ if compile_dit:
             pipeline.transformer_2.blocks[i] = torch.compile(pipeline.transformer_2.blocks[i])
     print("Add Compile")
 
-# Quantize (when the mode carries an "_and_<quant>" suffix) and then place the pipeline.
-# The order lives inside the helper: quantization has to happen before the offload hooks
-# are installed, and both transformers of this MoE setup are handled in one call, which
-# is exactly the bookkeeping the old 30-line if/elif chain repeated per script.
-apply_gpu_memory_mode(pipeline, GPU_memory_mode, device, weight_dtype)
+if GPU_memory_mode == "sequential_cpu_offload":
+    replace_parameters_by_name(transformer, ["modulation",], device=device)
+    transformer.freqs = transformer.freqs.to(device=device)
+    if transformer_2 is not None:
+        replace_parameters_by_name(transformer_2, ["modulation",], device=device)
+        transformer_2.freqs = transformer_2.freqs.to(device=device)
+    pipeline.enable_sequential_cpu_offload(device=device)
+elif GPU_memory_mode == "model_group_offload":
+    register_auto_device_hook(pipeline.transformer)
+    if transformer_2 is not None:
+        register_auto_device_hook(pipeline.transformer_2)
+    safe_enable_group_offload(pipeline, onload_device=device, offload_device="cpu", offload_type="leaf_level", use_stream=True)
+elif GPU_memory_mode == "model_cpu_offload_and_qfloat8":
+    convert_model_weight_to_float8(transformer, exclude_module_name=["modulation",], device=device)
+    convert_weight_dtype_wrapper(transformer, weight_dtype)
+    if transformer_2 is not None:
+        convert_model_weight_to_float8(transformer_2, exclude_module_name=["modulation",], device=device)
+        convert_weight_dtype_wrapper(transformer_2, weight_dtype)
+    pipeline.enable_model_cpu_offload(device=device)
+elif GPU_memory_mode == "model_cpu_offload":
+    pipeline.enable_model_cpu_offload(device=device)
+elif GPU_memory_mode == "model_full_load_and_qfloat8":
+    convert_model_weight_to_float8(transformer, exclude_module_name=["modulation",], device=device)
+    convert_weight_dtype_wrapper(transformer, weight_dtype)
+    if transformer_2 is not None:
+        convert_model_weight_to_float8(transformer_2, exclude_module_name=["modulation",], device=device)
+        convert_weight_dtype_wrapper(transformer_2, weight_dtype)
+    pipeline.to(device=device)
+else:
+    pipeline.to(device=device)
 
 coefficients = get_teacache_coefficients(model_name) if enable_teacache else None
 if coefficients is not None:

examples/wan2.2/predict_s2v.py

@@ -14,16 +14,23 @@ for project_root in project_roots:
 
 from videox_fun.dist import set_multi_gpus_devices, shard_model
 from videox_fun.models import (AutoencoderKLWan, AutoencoderKLWan3_8,
-                               AutoTokenizer, Wan2_2Transformer3DModel_S2V,
-                               WanAudioEncoder, WanT5EncoderModel)
+                               AutoTokenizer, CLIPModel,
+                               Wan2_2Transformer3DModel_S2V, WanAudioEncoder,
+                               WanT5EncoderModel)
 from videox_fun.models.cache_utils import get_teacache_coefficients
 from videox_fun.pipeline import Wan2_2S2VPipeline
-from videox_fun.utils import (FlowDPMSolverMultistepScheduler,
-                              FlowUniPCMultistepScheduler,
-                              apply_gpu_memory_mode, filter_kwargs,
-                              get_image_latent, get_video_to_video_latent,
-                              merge_lora, merge_video_audio, save_videos_grid,
-                              unmerge_lora)
+from videox_fun.utils import (register_auto_device_hook,
+                              safe_enable_group_offload)
+from videox_fun.utils.fm_solvers import FlowDPMSolverMultistepScheduler
+from videox_fun.utils.fm_solvers_unipc import FlowUniPCMultistepScheduler
+from videox_fun.utils.fp8_optimization import (convert_model_weight_to_float8,
+                                               convert_weight_dtype_wrapper,
+                                               replace_parameters_by_name)
+from videox_fun.utils.lora_utils import merge_lora, unmerge_lora
+from videox_fun.utils.utils import (filter_kwargs, get_image_latent,
+                                    get_image_to_video_latent,
+                                    get_video_to_video_latent,
+                                    merge_video_audio, save_videos_grid)
 
 # GPU memory mode, which can be chosen in [model_full_load, model_full_load_and_qfloat8, model_cpu_offload, model_cpu_offload_and_qfloat8, model_group_offload, sequential_cpu_offload].
 # model_full_load means that the entire model will be moved to the GPU.
@@ -107,7 +114,7 @@ segment_frame_length    = 80
 fps                     = 16
 
 # Use torch.float16 if GPU does not support torch.bfloat16
-# Some graphics cards, such as v100, 2080ti, do not support torch.bfloat16
+# ome graphics cards, such as v100, 2080ti, do not support torch.bfloat16
 weight_dtype            = torch.bfloat16
 # The path of the pose control video
 control_video           = "asset/pose.mp4"  
@@ -152,7 +159,7 @@ else:
 if transformer_path is not None:
     print(f"From checkpoint: {transformer_path}")
     if transformer_path.endswith("safetensors"):
-        from safetensors.torch import load_file
+        from safetensors.torch import load_file, safe_open
         state_dict = load_file(transformer_path)
     else:
         state_dict = torch.load(transformer_path, map_location="cpu")
@@ -165,7 +172,7 @@ if transformer_2 is not None:
     if transformer_high_path is not None:
         print(f"From checkpoint: {transformer_high_path}")
         if transformer_high_path.endswith("safetensors"):
-            from safetensors.torch import load_file
+            from safetensors.torch import load_file, safe_open
             state_dict = load_file(transformer_high_path)
         else:
             state_dict = torch.load(transformer_high_path, map_location="cpu")
@@ -187,7 +194,7 @@ vae = Chosen_AutoencoderKL.from_pretrained(
 if vae_path is not None:
     print(f"From checkpoint: {vae_path}")
     if vae_path.endswith("safetensors"):
-        from safetensors.torch import load_file
+        from safetensors.torch import load_file, safe_open
         state_dict = load_file(vae_path)
     else:
         state_dict = torch.load(vae_path, map_location="cpu")
@@ -216,7 +223,7 @@ audio_encoder = WanAudioEncoder(
 )
 
 # Get Scheduler
-Chosen_Scheduler = {
+Chosen_Scheduler = scheduler_dict = {
     "Flow": FlowMatchEulerDiscreteScheduler,
     "Flow_Unipc": FlowUniPCMultistepScheduler,
     "Flow_DPM++": FlowDPMSolverMultistepScheduler,
@@ -261,11 +268,36 @@ if compile_dit:
             pipeline.transformer_2.blocks[i] = torch.compile(pipeline.transformer_2.blocks[i])
     print("Add Compile")
 
-# Quantize (when the mode carries an "_and_<quant>" suffix) and then place the pipeline.
-# The order lives inside the helper: quantization has to happen before the offload hooks
-# are installed, and both transformers of this MoE setup are handled in one call, which
-# is exactly the bookkeeping the old 30-line if/elif chain repeated per script.
-apply_gpu_memory_mode(pipeline, GPU_memory_mode, device, weight_dtype)
+if GPU_memory_mode == "sequential_cpu_offload":
+    replace_parameters_by_name(transformer, ["modulation",], device=device)
+    transformer.freqs = transformer.freqs.to(device=device)
+    if transformer_2 is not None:
+        replace_parameters_by_name(transformer_2, ["modulation",], device=device)
+        transformer_2.freqs = transformer_2.freqs.to(device=device)
+    pipeline.enable_sequential_cpu_offload(device=device)
+elif GPU_memory_mode == "model_group_offload":
+    register_auto_device_hook(pipeline.transformer)
+    if transformer_2 is not None:
+        register_auto_device_hook(pipeline.transformer_2)
+    safe_enable_group_offload(pipeline, onload_device=device, offload_device="cpu", offload_type="leaf_level", use_stream=True)
+elif GPU_memory_mode == "model_cpu_offload_and_qfloat8":
+    convert_model_weight_to_float8(transformer, exclude_module_name=["modulation",], device=device)
+    convert_weight_dtype_wrapper(transformer, weight_dtype)
+    if transformer_2 is not None:
+        convert_model_weight_to_float8(transformer_2, exclude_module_name=["modulation",], device=device)
+        convert_weight_dtype_wrapper(transformer_2, weight_dtype)
+    pipeline.enable_model_cpu_offload(device=device)
+elif GPU_memory_mode == "model_cpu_offload":
+    pipeline.enable_model_cpu_offload(device=device)
+elif GPU_memory_mode == "model_full_load_and_qfloat8":
+    convert_model_weight_to_float8(transformer, exclude_module_name=["modulation",], device=device)
+    convert_weight_dtype_wrapper(transformer, weight_dtype)
+    if transformer_2 is not None:
+        convert_model_weight_to_float8(transformer_2, exclude_module_name=["modulation",], device=device)
+        convert_weight_dtype_wrapper(transformer_2, weight_dtype)
+    pipeline.to(device=device)
+else:
+    pipeline.to(device=device)
 
 coefficients = get_teacache_coefficients(model_name) if enable_teacache else None
 if coefficients is not None:

examples/wan2.2/predict_t2v.py

@@ -18,10 +18,16 @@ from videox_fun.models import (AutoencoderKLWan, AutoencoderKLWan3_8,
                                WanT5EncoderModel)
 from videox_fun.models.cache_utils import get_teacache_coefficients
 from videox_fun.pipeline import Wan2_2Pipeline
-from videox_fun.utils import (FlowDPMSolverMultistepScheduler,
-                              FlowUniPCMultistepScheduler,
-                              apply_gpu_memory_mode, filter_kwargs, merge_lora,
-                              save_videos_grid, unmerge_lora)
+from videox_fun.utils import (register_auto_device_hook,
+                              safe_enable_group_offload)
+from videox_fun.utils.fm_solvers import FlowDPMSolverMultistepScheduler
+from videox_fun.utils.fm_solvers_unipc import FlowUniPCMultistepScheduler
+from videox_fun.utils.fp8_optimization import (convert_model_weight_to_float8,
+                                               convert_weight_dtype_wrapper,
+                                               replace_parameters_by_name)
+from videox_fun.utils.lora_utils import merge_lora, unmerge_lora
+from videox_fun.utils.utils import (filter_kwargs, get_image_to_video_latent,
+                                    save_videos_grid)
 
 # GPU memory mode, which can be chosen in [model_full_load, model_full_load_and_qfloat8, model_cpu_offload, model_cpu_offload_and_qfloat8, model_group_offload, sequential_cpu_offload].
 # model_full_load means that the entire model will be moved to the GPU.
@@ -104,7 +110,7 @@ video_length        = 81
 fps                 = 16
 
 # Use torch.float16 if GPU does not support torch.bfloat16
-# Some graphics cards, such as v100, 2080ti, do not support torch.bfloat16
+# ome graphics cards, such as v100, 2080ti, do not support torch.bfloat16
 weight_dtype        = torch.bfloat16
 prompt              = "一只棕色的狗摇着头，坐在舒适房间里的浅色沙发上。在狗的后面，架子上有一幅镶框的画，周围是粉红色的花朵。房间里柔和温暖的灯光营造出舒适的氛围。"
 negative_prompt     = "色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走"
@@ -124,20 +130,22 @@ transformer = Wan2_2Transformer3DModel.from_pretrained(
     os.path.join(model_name, config['transformer_additional_kwargs'].get('transformer_low_noise_model_subpath', 'transformer')),
     transformer_additional_kwargs=OmegaConf.to_container(config['transformer_additional_kwargs']),
     low_cpu_mem_usage=True,
-    torch_dtype=weight_dtype)
+    torch_dtype=weight_dtype,
+)
 if config['transformer_additional_kwargs'].get('transformer_combination_type', 'single') == "moe":
     transformer_2 = Wan2_2Transformer3DModel.from_pretrained(
         os.path.join(model_name, config['transformer_additional_kwargs'].get('transformer_high_noise_model_subpath', 'transformer')),
         transformer_additional_kwargs=OmegaConf.to_container(config['transformer_additional_kwargs']),
         low_cpu_mem_usage=True,
-        torch_dtype=weight_dtype)
+        torch_dtype=weight_dtype,
+    )
 else:
     transformer_2 = None
 
 if transformer_path is not None:
     print(f"From checkpoint: {transformer_path}")
     if transformer_path.endswith("safetensors"):
-        from safetensors.torch import load_file
+        from safetensors.torch import load_file, safe_open
         state_dict = load_file(transformer_path)
     else:
         state_dict = torch.load(transformer_path, map_location="cpu")
@@ -150,7 +158,7 @@ if transformer_2 is not None:
     if transformer_high_path is not None:
         print(f"From checkpoint: {transformer_high_path}")
         if transformer_high_path.endswith("safetensors"):
-            from safetensors.torch import load_file
+            from safetensors.torch import load_file, safe_open
             state_dict = load_file(transformer_high_path)
         else:
             state_dict = torch.load(transformer_high_path, map_location="cpu")
@@ -166,12 +174,13 @@ Chosen_AutoencoderKL = {
 }[config['vae_kwargs'].get('vae_type', 'AutoencoderKLWan')]
 vae = Chosen_AutoencoderKL.from_pretrained(
     os.path.join(model_name, config['vae_kwargs'].get('vae_subpath', 'vae')),
-    additional_kwargs=OmegaConf.to_container(config['vae_kwargs'])).to(weight_dtype)
+    additional_kwargs=OmegaConf.to_container(config['vae_kwargs']),
+).to(weight_dtype)
 
 if vae_path is not None:
     print(f"From checkpoint: {vae_path}")
     if vae_path.endswith("safetensors"):
-        from safetensors.torch import load_file
+        from safetensors.torch import load_file, safe_open
         state_dict = load_file(vae_path)
     else:
         state_dict = torch.load(vae_path, map_location="cpu")
@@ -182,17 +191,19 @@ if vae_path is not None:
 
 # Get Tokenizer
 tokenizer = AutoTokenizer.from_pretrained(
-    os.path.join(model_name, config['text_encoder_kwargs'].get('tokenizer_subpath', 'tokenizer')))
+    os.path.join(model_name, config['text_encoder_kwargs'].get('tokenizer_subpath', 'tokenizer')),
+)
 
 # Get Text encoder
 text_encoder = WanT5EncoderModel.from_pretrained(
     os.path.join(model_name, config['text_encoder_kwargs'].get('text_encoder_subpath', 'text_encoder')),
     additional_kwargs=OmegaConf.to_container(config['text_encoder_kwargs']),
     low_cpu_mem_usage=True,
-    torch_dtype=weight_dtype)
+    torch_dtype=weight_dtype,
+)
 
 # Get Scheduler
-Chosen_Scheduler = {
+Chosen_Scheduler = scheduler_dict = {
     "Flow": FlowMatchEulerDiscreteScheduler,
     "Flow_Unipc": FlowUniPCMultistepScheduler,
     "Flow_DPM++": FlowDPMSolverMultistepScheduler,
@@ -210,7 +221,8 @@ pipeline = Wan2_2Pipeline(
     vae=vae,
     tokenizer=tokenizer,
     text_encoder=text_encoder,
-    scheduler=scheduler)
+    scheduler=scheduler,
+)
 if ulysses_degree > 1 or ring_degree > 1:
     from functools import partial
     transformer.enable_multi_gpus_inference()
@@ -235,11 +247,36 @@ if compile_dit:
             pipeline.transformer_2.blocks[i] = torch.compile(pipeline.transformer_2.blocks[i])
     print("Add Compile")
 
-# Quantize (when the mode carries an "_and_<quant>" suffix) and then place the pipeline.
-# The order lives inside the helper: quantization has to happen before the offload hooks
-# are installed, and both transformers of this MoE setup are handled in one call, which
-# is exactly the bookkeeping the old 30-line if/elif chain repeated per script.
-apply_gpu_memory_mode(pipeline, GPU_memory_mode, device, weight_dtype)
+if GPU_memory_mode == "sequential_cpu_offload":
+    replace_parameters_by_name(transformer, ["modulation",], device=device)
+    transformer.freqs = transformer.freqs.to(device=device)
+    if transformer_2 is not None:
+        replace_parameters_by_name(transformer_2, ["modulation",], device=device)
+        transformer_2.freqs = transformer_2.freqs.to(device=device)
+    pipeline.enable_sequential_cpu_offload(device=device)
+elif GPU_memory_mode == "model_group_offload":
+    register_auto_device_hook(pipeline.transformer)
+    if transformer_2 is not None:
+        register_auto_device_hook(pipeline.transformer_2)
+    safe_enable_group_offload(pipeline, onload_device=device, offload_device="cpu", offload_type="leaf_level", use_stream=True)
+elif GPU_memory_mode == "model_cpu_offload_and_qfloat8":
+    convert_model_weight_to_float8(transformer, exclude_module_name=["modulation",], device=device)
+    convert_weight_dtype_wrapper(transformer, weight_dtype)
+    if transformer_2 is not None:
+        convert_model_weight_to_float8(transformer_2, exclude_module_name=["modulation",], device=device)
+        convert_weight_dtype_wrapper(transformer_2, weight_dtype)
+    pipeline.enable_model_cpu_offload(device=device)
+elif GPU_memory_mode == "model_cpu_offload":
+    pipeline.enable_model_cpu_offload(device=device)
+elif GPU_memory_mode == "model_full_load_and_qfloat8":
+    convert_model_weight_to_float8(transformer, exclude_module_name=["modulation",], device=device)
+    convert_weight_dtype_wrapper(transformer, weight_dtype)
+    if transformer_2 is not None:
+        convert_model_weight_to_float8(transformer_2, exclude_module_name=["modulation",], device=device)
+        convert_weight_dtype_wrapper(transformer_2, weight_dtype)
+    pipeline.to(device=device)
+else:
+    pipeline.to(device=device)
 
 coefficients = get_teacache_coefficients(model_name) if enable_teacache else None
 if coefficients is not None:
@@ -282,7 +319,8 @@ with torch.no_grad():
         guidance_scale = guidance_scale,
         num_inference_steps = num_inference_steps,
         boundary = boundary,
-        shift = shift).videos
+        shift = shift,
+    ).videos
 
 if lora_path is not None:
     pipeline = unmerge_lora(pipeline, lora_path, lora_weight, device=device, dtype=weight_dtype)

examples/wan2.2/predict_ti2v.py

@@ -18,11 +18,16 @@ from videox_fun.models import (AutoencoderKLWan, AutoencoderKLWan3_8,
                                WanT5EncoderModel)
 from videox_fun.models.cache_utils import get_teacache_coefficients
 from videox_fun.pipeline import Wan2_2TI2VPipeline
-from videox_fun.utils import (FlowDPMSolverMultistepScheduler,
-                              FlowUniPCMultistepScheduler,
-                              apply_gpu_memory_mode, filter_kwargs,
-                              get_image_to_video_latent, merge_lora,
-                              save_videos_grid, unmerge_lora)
+from videox_fun.utils import (register_auto_device_hook,
+                              safe_enable_group_offload)
+from videox_fun.utils.fm_solvers import FlowDPMSolverMultistepScheduler
+from videox_fun.utils.fm_solvers_unipc import FlowUniPCMultistepScheduler
+from videox_fun.utils.fp8_optimization import (convert_model_weight_to_float8,
+                                               convert_weight_dtype_wrapper,
+                                               replace_parameters_by_name)
+from videox_fun.utils.lora_utils import merge_lora, unmerge_lora
+from videox_fun.utils.utils import (filter_kwargs, get_image_to_video_latent,
+                                    save_videos_grid)
 
 # GPU memory mode, which can be chosen in [model_full_load, model_full_load_and_qfloat8, model_cpu_offload, model_cpu_offload_and_qfloat8, model_group_offload, sequential_cpu_offload].
 # model_full_load means that the entire model will be moved to the GPU.
@@ -107,7 +112,7 @@ video_length        = 121
 fps                 = 24
 
 # Use torch.float16 if GPU does not support torch.bfloat16
-# Some graphics cards, such as v100, 2080ti, do not support torch.bfloat16
+# ome graphics cards, such as v100, 2080ti, do not support torch.bfloat16
 weight_dtype            = torch.bfloat16
 # If you want to generate from text, please set the validation_image_start = None and validation_image_end = None
 validation_image_start  = "asset/1.png"
@@ -146,7 +151,7 @@ else:
 if transformer_path is not None:
     print(f"From checkpoint: {transformer_path}")
     if transformer_path.endswith("safetensors"):
-        from safetensors.torch import load_file
+        from safetensors.torch import load_file, safe_open
         state_dict = load_file(transformer_path)
     else:
         state_dict = torch.load(transformer_path, map_location="cpu")
@@ -159,7 +164,7 @@ if transformer_2 is not None:
     if transformer_high_path is not None:
         print(f"From checkpoint: {transformer_high_path}")
         if transformer_high_path.endswith("safetensors"):
-            from safetensors.torch import load_file
+            from safetensors.torch import load_file, safe_open
             state_dict = load_file(transformer_high_path)
         else:
             state_dict = torch.load(transformer_high_path, map_location="cpu")
@@ -181,7 +186,7 @@ vae = Chosen_AutoencoderKL.from_pretrained(
 if vae_path is not None:
     print(f"From checkpoint: {vae_path}")
     if vae_path.endswith("safetensors"):
-        from safetensors.torch import load_file
+        from safetensors.torch import load_file, safe_open
         state_dict = load_file(vae_path)
     else:
         state_dict = torch.load(vae_path, map_location="cpu")
@@ -204,7 +209,7 @@ text_encoder = WanT5EncoderModel.from_pretrained(
 )
 
 # Get Scheduler
-Chosen_Scheduler = {
+Chosen_Scheduler = scheduler_dict = {
     "Flow": FlowMatchEulerDiscreteScheduler,
     "Flow_Unipc": FlowUniPCMultistepScheduler,
     "Flow_DPM++": FlowDPMSolverMultistepScheduler,
@@ -248,11 +253,36 @@ if compile_dit:
             pipeline.transformer_2.blocks[i] = torch.compile(pipeline.transformer_2.blocks[i])
     print("Add Compile")
 
-# Quantize (when the mode carries an "_and_<quant>" suffix) and then place the pipeline.
-# The order lives inside the helper: quantization has to happen before the offload hooks
-# are installed, and both transformers of this MoE setup are handled in one call, which
-# is exactly the bookkeeping the old 30-line if/elif chain repeated per script.
-apply_gpu_memory_mode(pipeline, GPU_memory_mode, device, weight_dtype)
+if GPU_memory_mode == "sequential_cpu_offload":
+    replace_parameters_by_name(transformer, ["modulation",], device=device)
+    transformer.freqs = transformer.freqs.to(device=device)
+    if transformer_2 is not None:
+        replace_parameters_by_name(transformer_2, ["modulation",], device=device)
+        transformer_2.freqs = transformer_2.freqs.to(device=device)
+    pipeline.enable_sequential_cpu_offload(device=device)
+elif GPU_memory_mode == "model_group_offload":
+    register_auto_device_hook(pipeline.transformer)
+    if transformer_2 is not None:
+        register_auto_device_hook(pipeline.transformer_2)
+    safe_enable_group_offload(pipeline, onload_device=device, offload_device="cpu", offload_type="leaf_level", use_stream=True)
+elif GPU_memory_mode == "model_cpu_offload_and_qfloat8":
+    convert_model_weight_to_float8(transformer, exclude_module_name=["modulation",], device=device)
+    convert_weight_dtype_wrapper(transformer, weight_dtype)
+    if transformer_2 is not None:
+        convert_model_weight_to_float8(transformer_2, exclude_module_name=["modulation",], device=device)
+        convert_weight_dtype_wrapper(transformer_2, weight_dtype)
+    pipeline.enable_model_cpu_offload(device=device)
+elif GPU_memory_mode == "model_cpu_offload":
+    pipeline.enable_model_cpu_offload(device=device)
+elif GPU_memory_mode == "model_full_load_and_qfloat8":
+    convert_model_weight_to_float8(transformer, exclude_module_name=["modulation",], device=device)
+    convert_weight_dtype_wrapper(transformer, weight_dtype)
+    if transformer_2 is not None:
+        convert_model_weight_to_float8(transformer_2, exclude_module_name=["modulation",], device=device)
+        convert_weight_dtype_wrapper(transformer_2, weight_dtype)
+    pipeline.to(device=device)
+else:
+    pipeline.to(device=device)
 
 coefficients = get_teacache_coefficients(model_name) if enable_teacache else None
 if coefficients is not None:

examples/wan2.2/predict_ti2v_tae.py

@@ -13,16 +13,21 @@ for project_root in project_roots:
     sys.path.insert(0, project_root) if project_root not in sys.path else None
 
 from videox_fun.dist import set_multi_gpus_devices, shard_model
-from videox_fun.models import (AutoencoderKLWan, AutoencoderKLWan3_8,
-                               AutoencoderTinyWan, AutoTokenizer,
-                               Wan2_2Transformer3DModel, WanT5EncoderModel)
+from videox_fun.models import (AutoencoderKLWan, AutoencoderKLWan3_8, AutoencoderTinyWan,
+                               AutoTokenizer, Wan2_2Transformer3DModel,
+                               WanT5EncoderModel)
 from videox_fun.models.cache_utils import get_teacache_coefficients
 from videox_fun.pipeline import Wan2_2TI2VPipeline
-from videox_fun.utils import (FlowDPMSolverMultistepScheduler,
-                              FlowUniPCMultistepScheduler,
-                              apply_gpu_memory_mode, filter_kwargs,
-                              get_image_to_video_latent, merge_lora,
-                              save_videos_grid, unmerge_lora)
+from videox_fun.utils import (register_auto_device_hook,
+                              safe_enable_group_offload)
+from videox_fun.utils.fm_solvers import FlowDPMSolverMultistepScheduler
+from videox_fun.utils.fm_solvers_unipc import FlowUniPCMultistepScheduler
+from videox_fun.utils.fp8_optimization import (convert_model_weight_to_float8,
+                                               convert_weight_dtype_wrapper,
+                                               replace_parameters_by_name)
+from videox_fun.utils.lora_utils import merge_lora, unmerge_lora
+from videox_fun.utils.utils import (filter_kwargs, get_image_to_video_latent,
+                                    save_videos_grid)
 
 # GPU memory mode, which can be chosen in [model_full_load, model_full_load_and_qfloat8, model_cpu_offload, model_cpu_offload_and_qfloat8, model_group_offload, sequential_cpu_offload].
 # model_full_load means that the entire model will be moved to the GPU.
@@ -111,7 +116,7 @@ video_length        = 81
 fps                 = 24
 
 # Use torch.float16 if GPU does not support torch.bfloat16
-# Some graphics cards, such as v100, 2080ti, do not support torch.bfloat16
+# ome graphics cards, such as v100, 2080ti, do not support torch.bfloat16
 weight_dtype            = torch.bfloat16
 # If you want to generate from text, please set the validation_image_start = None and validation_image_end = None
 validation_image_start  = "asset/1.png"
@@ -156,7 +161,7 @@ else:
 if transformer_path is not None:
     print(f"From checkpoint: {transformer_path}")
     if transformer_path.endswith("safetensors"):
-        from safetensors.torch import load_file
+        from safetensors.torch import load_file, safe_open
         state_dict = load_file(transformer_path)
     else:
         state_dict = torch.load(transformer_path, map_location="cpu")
@@ -169,7 +174,7 @@ if transformer_2 is not None:
     if transformer_high_path is not None:
         print(f"From checkpoint: {transformer_high_path}")
         if transformer_high_path.endswith("safetensors"):
-            from safetensors.torch import load_file
+            from safetensors.torch import load_file, safe_open
             state_dict = load_file(transformer_high_path)
         else:
             state_dict = torch.load(transformer_high_path, map_location="cpu")
@@ -192,7 +197,7 @@ vae = Chosen_AutoencoderKL.from_pretrained(
 if vae_path is not None:
     print(f"From checkpoint: {vae_path}")
     if vae_path.endswith("safetensors"):
-        from safetensors.torch import load_file
+        from safetensors.torch import load_file, safe_open
         state_dict = load_file(vae_path)
     else:
         state_dict = torch.load(vae_path, map_location="cpu")
@@ -221,7 +226,7 @@ text_encoder = WanT5EncoderModel.from_pretrained(
 )
 
 # Get Scheduler
-Chosen_Scheduler = {
+Chosen_Scheduler = scheduler_dict = {
     "Flow": FlowMatchEulerDiscreteScheduler,
     "Flow_Unipc": FlowUniPCMultistepScheduler,
     "Flow_DPM++": FlowDPMSolverMultistepScheduler,
@@ -265,11 +270,36 @@ if compile_dit:
             pipeline.transformer_2.blocks[i] = torch.compile(pipeline.transformer_2.blocks[i])
     print("Add Compile")
 
-# Quantize (when the mode carries an "_and_<quant>" suffix) and then place the pipeline.
-# The order lives inside the helper: quantization has to happen before the offload hooks
-# are installed, and both transformers of this MoE setup are handled in one call, which
-# is exactly the bookkeeping the old 30-line if/elif chain repeated per script.
-apply_gpu_memory_mode(pipeline, GPU_memory_mode, device, weight_dtype)
+if GPU_memory_mode == "sequential_cpu_offload":
+    replace_parameters_by_name(transformer, ["modulation",], device=device)
+    transformer.freqs = transformer.freqs.to(device=device)
+    if transformer_2 is not None:
+        replace_parameters_by_name(transformer_2, ["modulation",], device=device)
+        transformer_2.freqs = transformer_2.freqs.to(device=device)
+    pipeline.enable_sequential_cpu_offload(device=device)
+elif GPU_memory_mode == "model_group_offload":
+    register_auto_device_hook(pipeline.transformer)
+    if transformer_2 is not None:
+        register_auto_device_hook(pipeline.transformer_2)
+    safe_enable_group_offload(pipeline, onload_device=device, offload_device="cpu", offload_type="leaf_level", use_stream=True)
+elif GPU_memory_mode == "model_cpu_offload_and_qfloat8":
+    convert_model_weight_to_float8(transformer, exclude_module_name=["modulation",], device=device)
+    convert_weight_dtype_wrapper(transformer, weight_dtype)
+    if transformer_2 is not None:
+        convert_model_weight_to_float8(transformer_2, exclude_module_name=["modulation",], device=device)
+        convert_weight_dtype_wrapper(transformer_2, weight_dtype)
+    pipeline.enable_model_cpu_offload(device=device)
+elif GPU_memory_mode == "model_cpu_offload":
+    pipeline.enable_model_cpu_offload(device=device)
+elif GPU_memory_mode == "model_full_load_and_qfloat8":
+    convert_model_weight_to_float8(transformer, exclude_module_name=["modulation",], device=device)
+    convert_weight_dtype_wrapper(transformer, weight_dtype)
+    if transformer_2 is not None:
+        convert_model_weight_to_float8(transformer_2, exclude_module_name=["modulation",], device=device)
+        convert_weight_dtype_wrapper(transformer_2, weight_dtype)
+    pipeline.to(device=device)
+else:
+    pipeline.to(device=device)
 
 coefficients = get_teacache_coefficients(model_name) if enable_teacache else None
 if coefficients is not None:

examples/wan2.2_fun/app.py

@@ -10,7 +10,7 @@ for project_root in project_roots:
     sys.path.insert(0, project_root) if project_root not in sys.path else None
 
 from videox_fun.api.api import (infer_forward_api,
-                                update_diffusion_transformer_api)
+                               update_diffusion_transformer_api)
 from videox_fun.ui.controller import flow_scheduler_dict
 from videox_fun.ui.wan2_2_fun_ui import ui, ui_client, ui_host
 

examples/wan2.2_fun/launch_api.py

@@ -4,6 +4,7 @@ import sys
 import time
 
 import gradio as gr
+import ray
 import torch
 
 current_file_path = os.path.abspath(__file__)
@@ -12,11 +13,10 @@ for project_root in project_roots:
     sys.path.insert(0, project_root) if project_root not in sys.path else None
 
 from videox_fun.api.api_multi_nodes import (MultiNodesEngine,
-                                            multi_nodes_infer_forward_api)
+                                           multi_nodes_infer_forward_api)
 from videox_fun.ui.controller import flow_scheduler_dict
 from videox_fun.ui.wan2_2_fun_ui import Wan2_2_Fun_Controller
 
-
 def main():
     parser = argparse.ArgumentParser(description='xDiT HTTP Service')
     parser.add_argument('--world_size', type=int, default=8, help='Number of parallel workers')

examples/wan2.2_fun/predict_i2v.py

@@ -14,15 +14,20 @@ for project_root in project_roots:
 
 from videox_fun.dist import set_multi_gpus_devices, shard_model
 from videox_fun.models import (AutoencoderKLWan, AutoencoderKLWan3_8,
-                               AutoTokenizer, Wan2_2Transformer3DModel,
-                               WanT5EncoderModel)
+                               AutoTokenizer, CLIPModel,
+                               Wan2_2Transformer3DModel, WanT5EncoderModel)
 from videox_fun.models.cache_utils import get_teacache_coefficients
 from videox_fun.pipeline import Wan2_2FunInpaintPipeline
-from videox_fun.utils import (FlowDPMSolverMultistepScheduler,
-                              FlowUniPCMultistepScheduler,
-                              apply_gpu_memory_mode, filter_kwargs,
-                              get_image_to_video_latent, merge_lora,
-                              save_videos_grid, unmerge_lora)
+from videox_fun.utils import (register_auto_device_hook,
+                              safe_enable_group_offload)
+from videox_fun.utils.fm_solvers import FlowDPMSolverMultistepScheduler
+from videox_fun.utils.fm_solvers_unipc import FlowUniPCMultistepScheduler
+from videox_fun.utils.fp8_optimization import (convert_model_weight_to_float8,
+                                               convert_weight_dtype_wrapper,
+                                               replace_parameters_by_name)
+from videox_fun.utils.lora_utils import merge_lora, unmerge_lora
+from videox_fun.utils.utils import (filter_kwargs, get_image_to_video_latent,
+                                    save_videos_grid)
 
 # GPU memory mode, which can be chosen in [model_full_load, model_full_load_and_qfloat8, model_cpu_offload, model_cpu_offload_and_qfloat8, model_group_offload, sequential_cpu_offload].
 # model_full_load means that the entire model will be moved to the GPU.
@@ -106,7 +111,7 @@ video_length        = 81
 fps                 = 16
 
 # Use torch.float16 if GPU does not support torch.bfloat16
-# Some graphics cards, such as v100, 2080ti, do not support torch.bfloat16
+# ome graphics cards, such as v100, 2080ti, do not support torch.bfloat16
 weight_dtype            = torch.bfloat16
 # If you want to generate from text, please set the validation_image_start = None and validation_image_end = None
 validation_image_start  = "asset/1.png"
@@ -147,7 +152,7 @@ else:
 if transformer_path is not None:
     print(f"From checkpoint: {transformer_path}")
     if transformer_path.endswith("safetensors"):
-        from safetensors.torch import load_file
+        from safetensors.torch import load_file, safe_open
         state_dict = load_file(transformer_path)
     else:
         state_dict = torch.load(transformer_path, map_location="cpu")
@@ -160,7 +165,7 @@ if transformer_2 is not None:
     if transformer_high_path is not None:
         print(f"From checkpoint: {transformer_high_path}")
         if transformer_high_path.endswith("safetensors"):
-            from safetensors.torch import load_file
+            from safetensors.torch import load_file, safe_open
             state_dict = load_file(transformer_high_path)
         else:
             state_dict = torch.load(transformer_high_path, map_location="cpu")
@@ -182,7 +187,7 @@ vae = Chosen_AutoencoderKL.from_pretrained(
 if vae_path is not None:
     print(f"From checkpoint: {vae_path}")
     if vae_path.endswith("safetensors"):
-        from safetensors.torch import load_file
+        from safetensors.torch import load_file, safe_open
         state_dict = load_file(vae_path)
     else:
         state_dict = torch.load(vae_path, map_location="cpu")
@@ -206,7 +211,7 @@ text_encoder = WanT5EncoderModel.from_pretrained(
 text_encoder = text_encoder.eval()
 
 # Get Scheduler
-Chosen_Scheduler = {
+Chosen_Scheduler = scheduler_dict = {
     "Flow": FlowMatchEulerDiscreteScheduler,
     "Flow_Unipc": FlowUniPCMultistepScheduler,
     "Flow_DPM++": FlowDPMSolverMultistepScheduler,
@@ -250,11 +255,36 @@ if compile_dit:
             pipeline.transformer_2.blocks[i] = torch.compile(pipeline.transformer_2.blocks[i])
     print("Add Compile")
 
-# Quantize (when the mode carries an "_and_<quant>" suffix) and then place the pipeline.
-# The order lives inside the helper: quantization has to happen before the offload hooks
-# are installed, and both transformers of this MoE setup are handled in one call, which
-# is exactly the bookkeeping the old 30-line if/elif chain repeated per script.
-apply_gpu_memory_mode(pipeline, GPU_memory_mode, device, weight_dtype)
+if GPU_memory_mode == "sequential_cpu_offload":
+    replace_parameters_by_name(transformer, ["modulation",], device=device)
+    transformer.freqs = transformer.freqs.to(device=device)
+    if transformer_2 is not None:
+        replace_parameters_by_name(transformer_2, ["modulation",], device=device)
+        transformer_2.freqs = transformer_2.freqs.to(device=device)
+    pipeline.enable_sequential_cpu_offload(device=device)
+elif GPU_memory_mode == "model_group_offload":
+    register_auto_device_hook(pipeline.transformer)
+    if transformer_2 is not None:
+        register_auto_device_hook(pipeline.transformer_2)
+    safe_enable_group_offload(pipeline, onload_device=device, offload_device="cpu", offload_type="leaf_level", use_stream=True)
+elif GPU_memory_mode == "model_cpu_offload_and_qfloat8":
+    convert_model_weight_to_float8(transformer, exclude_module_name=["modulation",], device=device)
+    convert_weight_dtype_wrapper(transformer, weight_dtype)
+    if transformer_2 is not None:
+        convert_model_weight_to_float8(transformer_2, exclude_module_name=["modulation",], device=device)
+        convert_weight_dtype_wrapper(transformer_2, weight_dtype)
+    pipeline.enable_model_cpu_offload(device=device)
+elif GPU_memory_mode == "model_cpu_offload":
+    pipeline.enable_model_cpu_offload(device=device)
+elif GPU_memory_mode == "model_full_load_and_qfloat8":
+    convert_model_weight_to_float8(transformer, exclude_module_name=["modulation",], device=device)
+    convert_weight_dtype_wrapper(transformer, weight_dtype)
+    if transformer_2 is not None:
+        convert_model_weight_to_float8(transformer_2, exclude_module_name=["modulation",], device=device)
+        convert_weight_dtype_wrapper(transformer_2, weight_dtype)
+    pipeline.to(device=device)
+else:
+    pipeline.to(device=device)
 
 coefficients = get_teacache_coefficients(model_name) if enable_teacache else None
 if coefficients is not None:

examples/wan2.2_fun/predict_i2v_2.2vae.py

@@ -14,16 +14,21 @@ for project_root in project_roots:
 
 from videox_fun.dist import set_multi_gpus_devices, shard_model
 from videox_fun.models import (AutoencoderKLWan, AutoencoderKLWan3_8,
-                               AutoTokenizer, Wan2_2Transformer3DModel,
+                               AutoTokenizer, CLIPModel,
+                               Wan2_2Transformer3DModel,
                                WanLatentUpsamplerModel, WanT5EncoderModel)
 from videox_fun.models.cache_utils import get_teacache_coefficients
-from videox_fun.pipeline import (Wan2_2FunInpaintPipeline,
-                                 WanLatentUpsamplePipeline)
-from videox_fun.utils import (FlowDPMSolverMultistepScheduler,
-                              FlowUniPCMultistepScheduler,
-                              apply_gpu_memory_mode, filter_kwargs,
-                              get_image_to_video_latent, merge_lora,
-                              save_videos_grid, unmerge_lora)
+from videox_fun.pipeline import Wan2_2FunInpaintPipeline, WanLatentUpsamplePipeline
+from videox_fun.utils import (register_auto_device_hook,
+                              safe_enable_group_offload)
+from videox_fun.utils.fm_solvers import FlowDPMSolverMultistepScheduler
+from videox_fun.utils.fm_solvers_unipc import FlowUniPCMultistepScheduler
+from videox_fun.utils.fp8_optimization import (convert_model_weight_to_float8,
+                                               convert_weight_dtype_wrapper,
+                                               replace_parameters_by_name)
+from videox_fun.utils.lora_utils import merge_lora, unmerge_lora
+from videox_fun.utils.utils import (filter_kwargs, get_image_to_video_latent,
+                                    save_videos_grid)
 
 # GPU memory mode, which can be chosen in [model_full_load, model_full_load_and_qfloat8, model_cpu_offload, model_cpu_offload_and_qfloat8, model_group_offload, sequential_cpu_offload].
 # model_full_load means that the entire model will be moved to the GPU.
@@ -109,7 +114,7 @@ video_length        = 81
 fps                 = 16
 
 # Use torch.float16 if GPU does not support torch.bfloat16
-# Some graphics cards, such as v100, 2080ti, do not support torch.bfloat16
+# ome graphics cards, such as v100, 2080ti, do not support torch.bfloat16
 weight_dtype            = torch.bfloat16
 # If you want to generate from text, please set the validation_image_start = None and validation_image_end = None
 validation_image_start  = "asset/1.png"
@@ -150,7 +155,7 @@ else:
 if transformer_path is not None:
     print(f"From checkpoint: {transformer_path}")
     if transformer_path.endswith("safetensors"):
-        from safetensors.torch import load_file
+        from safetensors.torch import load_file, safe_open
         state_dict = load_file(transformer_path)
     else:
         state_dict = torch.load(transformer_path, map_location="cpu")
@@ -163,7 +168,7 @@ if transformer_2 is not None:
     if transformer_high_path is not None:
         print(f"From checkpoint: {transformer_high_path}")
         if transformer_high_path.endswith("safetensors"):
-            from safetensors.torch import load_file
+            from safetensors.torch import load_file, safe_open
             state_dict = load_file(transformer_high_path)
         else:
             state_dict = torch.load(transformer_high_path, map_location="cpu")
@@ -185,7 +190,7 @@ vae = Chosen_AutoencoderKL.from_pretrained(
 if vae_path is not None:
     print(f"From checkpoint: {vae_path}")
     if vae_path.endswith("safetensors"):
-        from safetensors.torch import load_file
+        from safetensors.torch import load_file, safe_open
         state_dict = load_file(vae_path)
     else:
         state_dict = torch.load(vae_path, map_location="cpu")
@@ -209,7 +214,7 @@ text_encoder = WanT5EncoderModel.from_pretrained(
 text_encoder = text_encoder.eval()
 
 # Get Scheduler
-Chosen_Scheduler = {
+Chosen_Scheduler = scheduler_dict = {
     "Flow": FlowMatchEulerDiscreteScheduler,
     "Flow_Unipc": FlowUniPCMultistepScheduler,
     "Flow_DPM++": FlowDPMSolverMultistepScheduler,
@@ -253,11 +258,36 @@ if compile_dit:
             pipeline.transformer_2.blocks[i] = torch.compile(pipeline.transformer_2.blocks[i])
     print("Add Compile")
 
-# Quantize (when the mode carries an "_and_<quant>" suffix) and then place the pipeline.
-# The order lives inside the helper: quantization has to happen before the offload hooks
-# are installed, and both transformers of this MoE setup are handled in one call, which
-# is exactly the bookkeeping the old 30-line if/elif chain repeated per script.
-apply_gpu_memory_mode(pipeline, GPU_memory_mode, device, weight_dtype)
+if GPU_memory_mode == "sequential_cpu_offload":
+    replace_parameters_by_name(transformer, ["modulation",], device=device)
+    transformer.freqs = transformer.freqs.to(device=device)
+    if transformer_2 is not None:
+        replace_parameters_by_name(transformer_2, ["modulation",], device=device)
+        transformer_2.freqs = transformer_2.freqs.to(device=device)
+    pipeline.enable_sequential_cpu_offload(device=device)
+elif GPU_memory_mode == "model_group_offload":
+    register_auto_device_hook(pipeline.transformer)
+    if transformer_2 is not None:
+        register_auto_device_hook(pipeline.transformer_2)
+    safe_enable_group_offload(pipeline, onload_device=device, offload_device="cpu", offload_type="leaf_level", use_stream=True)
+elif GPU_memory_mode == "model_cpu_offload_and_qfloat8":
+    convert_model_weight_to_float8(transformer, exclude_module_name=["modulation",], device=device)
+    convert_weight_dtype_wrapper(transformer, weight_dtype)
+    if transformer_2 is not None:
+        convert_model_weight_to_float8(transformer_2, exclude_module_name=["modulation",], device=device)
+        convert_weight_dtype_wrapper(transformer_2, weight_dtype)
+    pipeline.enable_model_cpu_offload(device=device)
+elif GPU_memory_mode == "model_cpu_offload":
+    pipeline.enable_model_cpu_offload(device=device)
+elif GPU_memory_mode == "model_full_load_and_qfloat8":
+    convert_model_weight_to_float8(transformer, exclude_module_name=["modulation",], device=device)
+    convert_weight_dtype_wrapper(transformer, weight_dtype)
+    if transformer_2 is not None:
+        convert_model_weight_to_float8(transformer_2, exclude_module_name=["modulation",], device=device)
+        convert_weight_dtype_wrapper(transformer_2, weight_dtype)
+    pipeline.to(device=device)
+else:
+    pipeline.to(device=device)
 
 coefficients = get_teacache_coefficients(model_name) if enable_teacache else None
 if coefficients is not None:

examples/wan2.2_fun/predict_i2v_2.2vae_tae.py

@@ -13,18 +13,22 @@ for project_root in project_roots:
     sys.path.insert(0, project_root) if project_root not in sys.path else None
 
 from videox_fun.dist import set_multi_gpus_devices, shard_model
-from videox_fun.models import (AutoencoderKLWan, AutoencoderKLWan3_8,
-                               AutoencoderTinyWan, AutoTokenizer,
+from videox_fun.models import (AutoencoderKLWan, AutoencoderKLWan3_8, AutoencoderTinyWan,
+                               AutoTokenizer, CLIPModel,
                                Wan2_2Transformer3DModel,
                                WanLatentUpsamplerModel, WanT5EncoderModel)
 from videox_fun.models.cache_utils import get_teacache_coefficients
-from videox_fun.pipeline import (Wan2_2FunInpaintPipeline,
-                                 WanLatentUpsamplePipeline)
-from videox_fun.utils import (FlowDPMSolverMultistepScheduler,
-                              FlowUniPCMultistepScheduler,
-                              apply_gpu_memory_mode, filter_kwargs,
-                              get_image_to_video_latent, merge_lora,
-                              save_videos_grid, unmerge_lora)
+from videox_fun.pipeline import Wan2_2FunInpaintPipeline, WanLatentUpsamplePipeline
+from videox_fun.utils import (register_auto_device_hook,
+                              safe_enable_group_offload)
+from videox_fun.utils.fm_solvers import FlowDPMSolverMultistepScheduler
+from videox_fun.utils.fm_solvers_unipc import FlowUniPCMultistepScheduler
+from videox_fun.utils.fp8_optimization import (convert_model_weight_to_float8,
+                                               convert_weight_dtype_wrapper,
+                                               replace_parameters_by_name)
+from videox_fun.utils.lora_utils import merge_lora, unmerge_lora
+from videox_fun.utils.utils import (filter_kwargs, get_image_to_video_latent,
+                                    save_videos_grid)
 
 # GPU memory mode, which can be chosen in [model_full_load, model_full_load_and_qfloat8, model_cpu_offload, model_cpu_offload_and_qfloat8, model_group_offload, sequential_cpu_offload].
 # model_full_load means that the entire model will be moved to the GPU.
@@ -113,7 +117,7 @@ video_length        = 81
 fps                 = 16
 
 # Use torch.float16 if GPU does not support torch.bfloat16
-# Some graphics cards, such as v100, 2080ti, do not support torch.bfloat16
+# ome graphics cards, such as v100, 2080ti, do not support torch.bfloat16
 weight_dtype            = torch.bfloat16
 # If you want to generate from text, please set the validation_image_start = None and validation_image_end = None
 validation_image_start  = "asset/1.png"
@@ -160,7 +164,7 @@ else:
 if transformer_path is not None:
     print(f"From checkpoint: {transformer_path}")
     if transformer_path.endswith("safetensors"):
-        from safetensors.torch import load_file
+        from safetensors.torch import load_file, safe_open
         state_dict = load_file(transformer_path)
     else:
         state_dict = torch.load(transformer_path, map_location="cpu")
@@ -173,7 +177,7 @@ if transformer_2 is not None:
     if transformer_high_path is not None:
         print(f"From checkpoint: {transformer_high_path}")
         if transformer_high_path.endswith("safetensors"):
-            from safetensors.torch import load_file
+            from safetensors.torch import load_file, safe_open
             state_dict = load_file(transformer_high_path)
         else:
             state_dict = torch.load(transformer_high_path, map_location="cpu")
@@ -196,7 +200,7 @@ vae = Chosen_AutoencoderKL.from_pretrained(
 if vae_path is not None:
     print(f"From checkpoint: {vae_path}")
     if vae_path.endswith("safetensors"):
-        from safetensors.torch import load_file
+        from safetensors.torch import load_file, safe_open
         state_dict = load_file(vae_path)
     else:
         state_dict = torch.load(vae_path, map_location="cpu")
@@ -224,7 +228,7 @@ text_encoder = WanT5EncoderModel.from_pretrained(
 text_encoder = text_encoder.eval()
 
 # Get Scheduler
-Chosen_Scheduler = {
+Chosen_Scheduler = scheduler_dict = {
     "Flow": FlowMatchEulerDiscreteScheduler,
     "Flow_Unipc": FlowUniPCMultistepScheduler,
     "Flow_DPM++": FlowDPMSolverMultistepScheduler,
@@ -268,11 +272,36 @@ if compile_dit:
             pipeline.transformer_2.blocks[i] = torch.compile(pipeline.transformer_2.blocks[i])
     print("Add Compile")
 
-# Quantize (when the mode carries an "_and_<quant>" suffix) and then place the pipeline.
-# The order lives inside the helper: quantization has to happen before the offload hooks
-# are installed, and both transformers of this MoE setup are handled in one call, which
-# is exactly the bookkeeping the old 30-line if/elif chain repeated per script.
-apply_gpu_memory_mode(pipeline, GPU_memory_mode, device, weight_dtype)
+if GPU_memory_mode == "sequential_cpu_offload":
+    replace_parameters_by_name(transformer, ["modulation",], device=device)
+    transformer.freqs = transformer.freqs.to(device=device)
+    if transformer_2 is not None:
+        replace_parameters_by_name(transformer_2, ["modulation",], device=device)
+        transformer_2.freqs = transformer_2.freqs.to(device=device)
+    pipeline.enable_sequential_cpu_offload(device=device)
+elif GPU_memory_mode == "model_group_offload":
+    register_auto_device_hook(pipeline.transformer)
+    if transformer_2 is not None:
+        register_auto_device_hook(pipeline.transformer_2)
+    safe_enable_group_offload(pipeline, onload_device=device, offload_device="cpu", offload_type="leaf_level", use_stream=True)
+elif GPU_memory_mode == "model_cpu_offload_and_qfloat8":
+    convert_model_weight_to_float8(transformer, exclude_module_name=["modulation",], device=device)
+    convert_weight_dtype_wrapper(transformer, weight_dtype)
+    if transformer_2 is not None:
+        convert_model_weight_to_float8(transformer_2, exclude_module_name=["modulation",], device=device)
+        convert_weight_dtype_wrapper(transformer_2, weight_dtype)
+    pipeline.enable_model_cpu_offload(device=device)
+elif GPU_memory_mode == "model_cpu_offload":
+    pipeline.enable_model_cpu_offload(device=device)
+elif GPU_memory_mode == "model_full_load_and_qfloat8":
+    convert_model_weight_to_float8(transformer, exclude_module_name=["modulation",], device=device)
+    convert_weight_dtype_wrapper(transformer, weight_dtype)
+    if transformer_2 is not None:
+        convert_model_weight_to_float8(transformer_2, exclude_module_name=["modulation",], device=device)
+        convert_weight_dtype_wrapper(transformer_2, weight_dtype)
+    pipeline.to(device=device)
+else:
+    pipeline.to(device=device)
 
 coefficients = get_teacache_coefficients(model_name) if enable_teacache else None
 if coefficients is not None:

examples/wan2.2_fun/predict_i2v_5b.py

@@ -14,15 +14,20 @@ for project_root in project_roots:
 
 from videox_fun.dist import set_multi_gpus_devices, shard_model
 from videox_fun.models import (AutoencoderKLWan, AutoencoderKLWan3_8,
-                               AutoTokenizer, Wan2_2Transformer3DModel,
-                               WanT5EncoderModel)
+                               AutoTokenizer, CLIPModel,
+                               Wan2_2Transformer3DModel, WanT5EncoderModel)
 from videox_fun.models.cache_utils import get_teacache_coefficients
 from videox_fun.pipeline import Wan2_2FunInpaintPipeline
-from videox_fun.utils import (FlowDPMSolverMultistepScheduler,
-                              FlowUniPCMultistepScheduler,
-                              apply_gpu_memory_mode, filter_kwargs,
-                              get_image_to_video_latent, merge_lora,
-                              save_videos_grid, unmerge_lora)
+from videox_fun.utils import (register_auto_device_hook,
+                              safe_enable_group_offload)
+from videox_fun.utils.fm_solvers import FlowDPMSolverMultistepScheduler
+from videox_fun.utils.fm_solvers_unipc import FlowUniPCMultistepScheduler
+from videox_fun.utils.fp8_optimization import (convert_model_weight_to_float8,
+                                               convert_weight_dtype_wrapper,
+                                               replace_parameters_by_name)
+from videox_fun.utils.lora_utils import merge_lora, unmerge_lora
+from videox_fun.utils.utils import (filter_kwargs, get_image_to_video_latent,
+                                    save_videos_grid)
 
 # GPU memory mode, which can be chosen in [model_full_load, model_full_load_and_qfloat8, model_cpu_offload, model_cpu_offload_and_qfloat8, model_group_offload, sequential_cpu_offload].
 # model_full_load means that the entire model will be moved to the GPU.
@@ -108,7 +113,7 @@ video_length        = 121
 fps                 = 24
 
 # Use torch.float16 if GPU does not support torch.bfloat16
-# Some graphics cards, such as v100, 2080ti, do not support torch.bfloat16
+# ome graphics cards, such as v100, 2080ti, do not support torch.bfloat16
 weight_dtype            = torch.bfloat16
 # If you want to generate from text, please set the validation_image_start = None and validation_image_end = None
 validation_image_start  = "asset/1.png"
@@ -149,7 +154,7 @@ else:
 if transformer_path is not None:
     print(f"From checkpoint: {transformer_path}")
     if transformer_path.endswith("safetensors"):
-        from safetensors.torch import load_file
+        from safetensors.torch import load_file, safe_open
         state_dict = load_file(transformer_path)
     else:
         state_dict = torch.load(transformer_path, map_location="cpu")
@@ -162,7 +167,7 @@ if transformer_2 is not None:
     if transformer_high_path is not None:
         print(f"From checkpoint: {transformer_high_path}")
         if transformer_high_path.endswith("safetensors"):
-            from safetensors.torch import load_file
+            from safetensors.torch import load_file, safe_open
             state_dict = load_file(transformer_high_path)
         else:
             state_dict = torch.load(transformer_high_path, map_location="cpu")
@@ -184,7 +189,7 @@ vae = Chosen_AutoencoderKL.from_pretrained(
 if vae_path is not None:
     print(f"From checkpoint: {vae_path}")
     if vae_path.endswith("safetensors"):
-        from safetensors.torch import load_file
+        from safetensors.torch import load_file, safe_open
         state_dict = load_file(vae_path)
     else:
         state_dict = torch.load(vae_path, map_location="cpu")
@@ -208,7 +213,7 @@ text_encoder = WanT5EncoderModel.from_pretrained(
 text_encoder = text_encoder.eval()
 
 # Get Scheduler
-Chosen_Scheduler = {
+Chosen_Scheduler = scheduler_dict = {
     "Flow": FlowMatchEulerDiscreteScheduler,
     "Flow_Unipc": FlowUniPCMultistepScheduler,
     "Flow_DPM++": FlowDPMSolverMultistepScheduler,
@@ -252,11 +257,36 @@ if compile_dit:
             pipeline.transformer_2.blocks[i] = torch.compile(pipeline.transformer_2.blocks[i])
     print("Add Compile")
 
-# Quantize (when the mode carries an "_and_<quant>" suffix) and then place the pipeline.
-# The order lives inside the helper: quantization has to happen before the offload hooks
-# are installed, and both transformers of this MoE setup are handled in one call, which
-# is exactly the bookkeeping the old 30-line if/elif chain repeated per script.
-apply_gpu_memory_mode(pipeline, GPU_memory_mode, device, weight_dtype)
+if GPU_memory_mode == "sequential_cpu_offload":
+    replace_parameters_by_name(transformer, ["modulation",], device=device)
+    transformer.freqs = transformer.freqs.to(device=device)
+    if transformer_2 is not None:
+        replace_parameters_by_name(transformer_2, ["modulation",], device=device)
+        transformer_2.freqs = transformer_2.freqs.to(device=device)
+    pipeline.enable_sequential_cpu_offload(device=device)
+elif GPU_memory_mode == "model_group_offload":
+    register_auto_device_hook(pipeline.transformer)
+    if transformer_2 is not None:
+        register_auto_device_hook(pipeline.transformer_2)
+    safe_enable_group_offload(pipeline, onload_device=device, offload_device="cpu", offload_type="leaf_level", use_stream=True)
+elif GPU_memory_mode == "model_cpu_offload_and_qfloat8":
+    convert_model_weight_to_float8(transformer, exclude_module_name=["modulation",], device=device)
+    convert_weight_dtype_wrapper(transformer, weight_dtype)
+    if transformer_2 is not None:
+        convert_model_weight_to_float8(transformer_2, exclude_module_name=["modulation",], device=device)
+        convert_weight_dtype_wrapper(transformer_2, weight_dtype)
+    pipeline.enable_model_cpu_offload(device=device)
+elif GPU_memory_mode == "model_cpu_offload":
+    pipeline.enable_model_cpu_offload(device=device)
+elif GPU_memory_mode == "model_full_load_and_qfloat8":
+    convert_model_weight_to_float8(transformer, exclude_module_name=["modulation",], device=device)
+    convert_weight_dtype_wrapper(transformer, weight_dtype)
+    if transformer_2 is not None:
+        convert_model_weight_to_float8(transformer_2, exclude_module_name=["modulation",], device=device)
+        convert_weight_dtype_wrapper(transformer_2, weight_dtype)
+    pipeline.to(device=device)
+else:
+    pipeline.to(device=device)
 
 coefficients = get_teacache_coefficients(model_name) if enable_teacache else None
 if coefficients is not None:

examples/wan2.2_fun/predict_t2v.py

@@ -14,15 +14,20 @@ for project_root in project_roots:
 
 from videox_fun.dist import set_multi_gpus_devices, shard_model
 from videox_fun.models import (AutoencoderKLWan, AutoencoderKLWan3_8,
-                               AutoTokenizer, Wan2_2Transformer3DModel,
-                               WanT5EncoderModel)
+                               AutoTokenizer, CLIPModel,
+                               Wan2_2Transformer3DModel, WanT5EncoderModel)
 from videox_fun.models.cache_utils import get_teacache_coefficients
 from videox_fun.pipeline import Wan2_2FunInpaintPipeline
-from videox_fun.utils import (FlowDPMSolverMultistepScheduler,
-                              FlowUniPCMultistepScheduler,
-                              apply_gpu_memory_mode, filter_kwargs,
-                              get_image_to_video_latent, merge_lora,
-                              save_videos_grid, unmerge_lora)
+from videox_fun.utils import (register_auto_device_hook,
+                              safe_enable_group_offload)
+from videox_fun.utils.fm_solvers import FlowDPMSolverMultistepScheduler
+from videox_fun.utils.fm_solvers_unipc import FlowUniPCMultistepScheduler
+from videox_fun.utils.fp8_optimization import (convert_model_weight_to_float8,
+                                               convert_weight_dtype_wrapper,
+                                               replace_parameters_by_name)
+from videox_fun.utils.lora_utils import merge_lora, unmerge_lora
+from videox_fun.utils.utils import (filter_kwargs, get_image_to_video_latent,
+                                    save_videos_grid)
 
 # GPU memory mode, which can be chosen in [model_full_load, model_full_load_and_qfloat8, model_cpu_offload, model_cpu_offload_and_qfloat8, model_group_offload, sequential_cpu_offload].
 # model_full_load means that the entire model will be moved to the GPU.
@@ -106,7 +111,7 @@ video_length        = 81
 fps                 = 16
 
 # Use torch.float16 if GPU does not support torch.bfloat16
-# Some graphics cards, such as v100, 2080ti, do not support torch.bfloat16
+# ome graphics cards, such as v100, 2080ti, do not support torch.bfloat16
 weight_dtype            = torch.bfloat16
 # 使用更长的neg prompt如"模糊，突变，变形，失真，画面暗，文本字幕，画面固定，连环画，漫画，线稿，没有主体。"，可以增加稳定性
 # 在neg prompt中添加"安静，固定"等词语可以增加动态性。
@@ -143,7 +148,7 @@ else:
 if transformer_path is not None:
     print(f"From checkpoint: {transformer_path}")
     if transformer_path.endswith("safetensors"):
-        from safetensors.torch import load_file
+        from safetensors.torch import load_file, safe_open
         state_dict = load_file(transformer_path)
     else:
         state_dict = torch.load(transformer_path, map_location="cpu")
@@ -156,7 +161,7 @@ if transformer_2 is not None:
     if transformer_high_path is not None:
         print(f"From checkpoint: {transformer_high_path}")
         if transformer_high_path.endswith("safetensors"):
-            from safetensors.torch import load_file
+            from safetensors.torch import load_file, safe_open
             state_dict = load_file(transformer_high_path)
         else:
             state_dict = torch.load(transformer_high_path, map_location="cpu")
@@ -178,7 +183,7 @@ vae = Chosen_AutoencoderKL.from_pretrained(
 if vae_path is not None:
     print(f"From checkpoint: {vae_path}")
     if vae_path.endswith("safetensors"):
-        from safetensors.torch import load_file
+        from safetensors.torch import load_file, safe_open
         state_dict = load_file(vae_path)
     else:
         state_dict = torch.load(vae_path, map_location="cpu")
@@ -202,7 +207,7 @@ text_encoder = WanT5EncoderModel.from_pretrained(
 text_encoder = text_encoder.eval()
 
 # Get Scheduler
-Chosen_Scheduler = {
+Chosen_Scheduler = scheduler_dict = {
     "Flow": FlowMatchEulerDiscreteScheduler,
     "Flow_Unipc": FlowUniPCMultistepScheduler,
     "Flow_DPM++": FlowDPMSolverMultistepScheduler,
@@ -246,11 +251,36 @@ if compile_dit:
             pipeline.transformer_2.blocks[i] = torch.compile(pipeline.transformer_2.blocks[i])
     print("Add Compile")
 
-# Quantize (when the mode carries an "_and_<quant>" suffix) and then place the pipeline.
-# The order lives inside the helper: quantization has to happen before the offload hooks
-# are installed, and both transformers of this MoE setup are handled in one call, which
-# is exactly the bookkeeping the old 30-line if/elif chain repeated per script.
-apply_gpu_memory_mode(pipeline, GPU_memory_mode, device, weight_dtype)
+if GPU_memory_mode == "sequential_cpu_offload":
+    replace_parameters_by_name(transformer, ["modulation",], device=device)
+    transformer.freqs = transformer.freqs.to(device=device)
+    if transformer_2 is not None:
+        replace_parameters_by_name(transformer_2, ["modulation",], device=device)
+        transformer_2.freqs = transformer_2.freqs.to(device=device)
+    pipeline.enable_sequential_cpu_offload(device=device)
+elif GPU_memory_mode == "model_group_offload":
+    register_auto_device_hook(pipeline.transformer)
+    if transformer_2 is not None:
+        register_auto_device_hook(pipeline.transformer_2)
+    safe_enable_group_offload(pipeline, onload_device=device, offload_device="cpu", offload_type="leaf_level", use_stream=True)
+elif GPU_memory_mode == "model_cpu_offload_and_qfloat8":
+    convert_model_weight_to_float8(transformer, exclude_module_name=["modulation",], device=device)
+    convert_weight_dtype_wrapper(transformer, weight_dtype)
+    if transformer_2 is not None:
+        convert_model_weight_to_float8(transformer_2, exclude_module_name=["modulation",], device=device)
+        convert_weight_dtype_wrapper(transformer_2, weight_dtype)
+    pipeline.enable_model_cpu_offload(device=device)
+elif GPU_memory_mode == "model_cpu_offload":
+    pipeline.enable_model_cpu_offload(device=device)
+elif GPU_memory_mode == "model_full_load_and_qfloat8":
+    convert_model_weight_to_float8(transformer, exclude_module_name=["modulation",], device=device)
+    convert_weight_dtype_wrapper(transformer, weight_dtype)
+    if transformer_2 is not None:
+        convert_model_weight_to_float8(transformer_2, exclude_module_name=["modulation",], device=device)
+        convert_weight_dtype_wrapper(transformer_2, weight_dtype)
+    pipeline.to(device=device)
+else:
+    pipeline.to(device=device)
 
 coefficients = get_teacache_coefficients(model_name) if enable_teacache else None
 if coefficients is not None:

examples/wan2.2_fun/predict_t2v_2.2vae.py

@@ -15,13 +15,20 @@ for project_root in project_roots:
 from videox_fun.dist import set_multi_gpus_devices, shard_model
 from videox_fun.models import (AutoencoderKLWan, AutoencoderKLWan3_8,
                                AutoTokenizer, Wan2_2Transformer3DModel,
-                               WanLatentUpsamplerModel, WanT5EncoderModel)
+                               WanLatentUpsamplerModel,
+                               WanT5EncoderModel)
 from videox_fun.models.cache_utils import get_teacache_coefficients
 from videox_fun.pipeline import Wan2_2Pipeline, WanLatentUpsamplePipeline
-from videox_fun.utils import (FlowDPMSolverMultistepScheduler,
-                              FlowUniPCMultistepScheduler,
-                              apply_gpu_memory_mode, filter_kwargs, merge_lora,
-                              save_videos_grid, unmerge_lora)
+from videox_fun.utils import (register_auto_device_hook,
+                              safe_enable_group_offload)
+from videox_fun.utils.fm_solvers import FlowDPMSolverMultistepScheduler
+from videox_fun.utils.fm_solvers_unipc import FlowUniPCMultistepScheduler
+from videox_fun.utils.fp8_optimization import (convert_model_weight_to_float8,
+                                               convert_weight_dtype_wrapper,
+                                               replace_parameters_by_name)
+from videox_fun.utils.lora_utils import merge_lora, unmerge_lora
+from videox_fun.utils.utils import (filter_kwargs, get_image_to_video_latent,
+                                    save_videos_grid)
 
 # GPU memory mode, which can be chosen in [model_full_load, model_full_load_and_qfloat8, model_cpu_offload, model_cpu_offload_and_qfloat8, model_group_offload, sequential_cpu_offload].
 # model_full_load means that the entire model will be moved to the GPU.
@@ -107,7 +114,7 @@ video_length        = 81
 fps                 = 16
 
 # Use torch.float16 if GPU does not support torch.bfloat16
-# Some graphics cards, such as v100, 2080ti, do not support torch.bfloat16
+# ome graphics cards, such as v100, 2080ti, do not support torch.bfloat16
 weight_dtype            = torch.bfloat16
 # 使用更长的neg prompt如"模糊，突变，变形，失真，画面暗，文本字幕，画面固定，连环画，漫画，线稿，没有主体。"，可以增加稳定性
 # 在neg prompt中添加"安静，固定"等词语可以增加动态性。
@@ -144,7 +151,7 @@ else:
 if transformer_path is not None:
     print(f"From checkpoint: {transformer_path}")
     if transformer_path.endswith("safetensors"):
-        from safetensors.torch import load_file
+        from safetensors.torch import load_file, safe_open
         state_dict = load_file(transformer_path)
     else:
         state_dict = torch.load(transformer_path, map_location="cpu")
@@ -157,7 +164,7 @@ if transformer_2 is not None:
     if transformer_high_path is not None:
         print(f"From checkpoint: {transformer_high_path}")
         if transformer_high_path.endswith("safetensors"):
-            from safetensors.torch import load_file
+            from safetensors.torch import load_file, safe_open
             state_dict = load_file(transformer_high_path)
         else:
             state_dict = torch.load(transformer_high_path, map_location="cpu")
@@ -179,7 +186,7 @@ vae = Chosen_AutoencoderKL.from_pretrained(
 if vae_path is not None:
     print(f"From checkpoint: {vae_path}")
     if vae_path.endswith("safetensors"):
-        from safetensors.torch import load_file
+        from safetensors.torch import load_file, safe_open
         state_dict = load_file(vae_path)
     else:
         state_dict = torch.load(vae_path, map_location="cpu")
@@ -203,7 +210,7 @@ text_encoder = WanT5EncoderModel.from_pretrained(
 text_encoder = text_encoder.eval()
 
 # Get Scheduler
-Chosen_Scheduler = {
+Chosen_Scheduler = scheduler_dict = {
     "Flow": FlowMatchEulerDiscreteScheduler,
     "Flow_Unipc": FlowUniPCMultistepScheduler,
     "Flow_DPM++": FlowDPMSolverMultistepScheduler,
@@ -247,11 +254,36 @@ if compile_dit:
             pipeline.transformer_2.blocks[i] = torch.compile(pipeline.transformer_2.blocks[i])
     print("Add Compile")
 
-# Quantize (when the mode carries an "_and_<quant>" suffix) and then place the pipeline.
-# The order lives inside the helper: quantization has to happen before the offload hooks
-# are installed, and both transformers of this MoE setup are handled in one call, which
-# is exactly the bookkeeping the old 30-line if/elif chain repeated per script.
-apply_gpu_memory_mode(pipeline, GPU_memory_mode, device, weight_dtype)
+if GPU_memory_mode == "sequential_cpu_offload":
+    replace_parameters_by_name(transformer, ["modulation",], device=device)
+    transformer.freqs = transformer.freqs.to(device=device)
+    if transformer_2 is not None:
+        replace_parameters_by_name(transformer_2, ["modulation",], device=device)
+        transformer_2.freqs = transformer_2.freqs.to(device=device)
+    pipeline.enable_sequential_cpu_offload(device=device)
+elif GPU_memory_mode == "model_group_offload":
+    register_auto_device_hook(pipeline.transformer)
+    if transformer_2 is not None:
+        register_auto_device_hook(pipeline.transformer_2)
+    safe_enable_group_offload(pipeline, onload_device=device, offload_device="cpu", offload_type="leaf_level", use_stream=True)
+elif GPU_memory_mode == "model_cpu_offload_and_qfloat8":
+    convert_model_weight_to_float8(transformer, exclude_module_name=["modulation",], device=device)
+    convert_weight_dtype_wrapper(transformer, weight_dtype)
+    if transformer_2 is not None:
+        convert_model_weight_to_float8(transformer_2, exclude_module_name=["modulation",], device=device)
+        convert_weight_dtype_wrapper(transformer_2, weight_dtype)
+    pipeline.enable_model_cpu_offload(device=device)
+elif GPU_memory_mode == "model_cpu_offload":
+    pipeline.enable_model_cpu_offload(device=device)
+elif GPU_memory_mode == "model_full_load_and_qfloat8":
+    convert_model_weight_to_float8(transformer, exclude_module_name=["modulation",], device=device)
+    convert_weight_dtype_wrapper(transformer, weight_dtype)
+    if transformer_2 is not None:
+        convert_model_weight_to_float8(transformer_2, exclude_module_name=["modulation",], device=device)
+        convert_weight_dtype_wrapper(transformer_2, weight_dtype)
+    pipeline.to(device=device)
+else:
+    pipeline.to(device=device)
 
 coefficients = get_teacache_coefficients(model_name) if enable_teacache else None
 if coefficients is not None:

examples/wan2.2_fun/predict_t2v_2.2vae_tae.py

@@ -13,16 +13,22 @@ for project_root in project_roots:
     sys.path.insert(0, project_root) if project_root not in sys.path else None
 
 from videox_fun.dist import set_multi_gpus_devices, shard_model
-from videox_fun.models import (AutoencoderKLWan, AutoencoderKLWan3_8,
-                               AutoencoderTinyWan, AutoTokenizer,
-                               Wan2_2Transformer3DModel,
-                               WanLatentUpsamplerModel, WanT5EncoderModel)
+from videox_fun.models import (AutoencoderKLWan, AutoencoderKLWan3_8, AutoencoderTinyWan,
+                               AutoTokenizer, Wan2_2Transformer3DModel,
+                               WanLatentUpsamplerModel,
+                               WanT5EncoderModel)
 from videox_fun.models.cache_utils import get_teacache_coefficients
 from videox_fun.pipeline import Wan2_2Pipeline, WanLatentUpsamplePipeline
-from videox_fun.utils import (FlowDPMSolverMultistepScheduler,
-                              FlowUniPCMultistepScheduler,
-                              apply_gpu_memory_mode, filter_kwargs, merge_lora,
-                              save_videos_grid, unmerge_lora)
+from videox_fun.utils import (register_auto_device_hook,
+                              safe_enable_group_offload)
+from videox_fun.utils.fm_solvers import FlowDPMSolverMultistepScheduler
+from videox_fun.utils.fm_solvers_unipc import FlowUniPCMultistepScheduler
+from videox_fun.utils.fp8_optimization import (convert_model_weight_to_float8,
+                                               convert_weight_dtype_wrapper,
+                                               replace_parameters_by_name)
+from videox_fun.utils.lora_utils import merge_lora, unmerge_lora
+from videox_fun.utils.utils import (filter_kwargs, get_image_to_video_latent,
+                                    save_videos_grid)
 
 # GPU memory mode, which can be chosen in [model_full_load, model_full_load_and_qfloat8, model_cpu_offload, model_cpu_offload_and_qfloat8, model_group_offload, sequential_cpu_offload].
 # model_full_load means that the entire model will be moved to the GPU.
@@ -112,7 +118,7 @@ video_length        = 81
 fps                 = 16
 
 # Use torch.float16 if GPU does not support torch.bfloat16
-# Some graphics cards, such as v100, 2080ti, do not support torch.bfloat16
+# ome graphics cards, such as v100, 2080ti, do not support torch.bfloat16
 weight_dtype        = torch.bfloat16
 # 使用更长的neg prompt如"模糊，突变，变形，失真，画面暗，文本字幕，画面固定，连环画，漫画，线稿，没有主体。"，可以增加稳定性
 # 在neg prompt中添加"安静，固定"等词语可以增加动态性。
@@ -155,7 +161,7 @@ else:
 if transformer_path is not None:
     print(f"From checkpoint: {transformer_path}")
     if transformer_path.endswith("safetensors"):
-        from safetensors.torch import load_file
+        from safetensors.torch import load_file, safe_open
         state_dict = load_file(transformer_path)
     else:
         state_dict = torch.load(transformer_path, map_location="cpu")
@@ -168,7 +174,7 @@ if transformer_2 is not None:
     if transformer_high_path is not None:
         print(f"From checkpoint: {transformer_high_path}")
         if transformer_high_path.endswith("safetensors"):
-            from safetensors.torch import load_file
+            from safetensors.torch import load_file, safe_open
             state_dict = load_file(transformer_high_path)
         else:
             state_dict = torch.load(transformer_high_path, map_location="cpu")
@@ -191,7 +197,7 @@ vae = Chosen_AutoencoderKL.from_pretrained(
 if vae_path is not None:
     print(f"From checkpoint: {vae_path}")
     if vae_path.endswith("safetensors"):
-        from safetensors.torch import load_file
+        from safetensors.torch import load_file, safe_open
         state_dict = load_file(vae_path)
     else:
         state_dict = torch.load(vae_path, map_location="cpu")
@@ -219,7 +225,7 @@ text_encoder = WanT5EncoderModel.from_pretrained(
 text_encoder = text_encoder.eval()
 
 # Get Scheduler
-Chosen_Scheduler = {
+Chosen_Scheduler = scheduler_dict = {
     "Flow": FlowMatchEulerDiscreteScheduler,
     "Flow_Unipc": FlowUniPCMultistepScheduler,
     "Flow_DPM++": FlowDPMSolverMultistepScheduler,
@@ -263,11 +269,36 @@ if compile_dit:
             pipeline.transformer_2.blocks[i] = torch.compile(pipeline.transformer_2.blocks[i])
     print("Add Compile")
 
-# Quantize (when the mode carries an "_and_<quant>" suffix) and then place the pipeline.
-# The order lives inside the helper: quantization has to happen before the offload hooks
-# are installed, and both transformers of this MoE setup are handled in one call, which
-# is exactly the bookkeeping the old 30-line if/elif chain repeated per script.
-apply_gpu_memory_mode(pipeline, GPU_memory_mode, device, weight_dtype)
+if GPU_memory_mode == "sequential_cpu_offload":
+    replace_parameters_by_name(transformer, ["modulation",], device=device)
+    transformer.freqs = transformer.freqs.to(device=device)
+    if transformer_2 is not None:
+        replace_parameters_by_name(transformer_2, ["modulation",], device=device)
+        transformer_2.freqs = transformer_2.freqs.to(device=device)
+    pipeline.enable_sequential_cpu_offload(device=device)
+elif GPU_memory_mode == "model_group_offload":
+    register_auto_device_hook(pipeline.transformer)
+    if transformer_2 is not None:
+        register_auto_device_hook(pipeline.transformer_2)
+    safe_enable_group_offload(pipeline, onload_device=device, offload_device="cpu", offload_type="leaf_level", use_stream=True)
+elif GPU_memory_mode == "model_cpu_offload_and_qfloat8":
+    convert_model_weight_to_float8(transformer, exclude_module_name=["modulation",], device=device)
+    convert_weight_dtype_wrapper(transformer, weight_dtype)
+    if transformer_2 is not None:
+        convert_model_weight_to_float8(transformer_2, exclude_module_name=["modulation",], device=device)
+        convert_weight_dtype_wrapper(transformer_2, weight_dtype)
+    pipeline.enable_model_cpu_offload(device=device)
+elif GPU_memory_mode == "model_cpu_offload":
+    pipeline.enable_model_cpu_offload(device=device)
+elif GPU_memory_mode == "model_full_load_and_qfloat8":
+    convert_model_weight_to_float8(transformer, exclude_module_name=["modulation",], device=device)
+    convert_weight_dtype_wrapper(transformer, weight_dtype)
+    if transformer_2 is not None:
+        convert_model_weight_to_float8(transformer_2, exclude_module_name=["modulation",], device=device)
+        convert_weight_dtype_wrapper(transformer_2, weight_dtype)
+    pipeline.to(device=device)
+else:
+    pipeline.to(device=device)
 
 coefficients = get_teacache_coefficients(model_name) if enable_teacache else None
 if coefficients is not None:

examples/wan2.2_fun/predict_t2v_5b.py

@@ -14,15 +14,20 @@ for project_root in project_roots:
 
 from videox_fun.dist import set_multi_gpus_devices, shard_model
 from videox_fun.models import (AutoencoderKLWan, AutoencoderKLWan3_8,
-                               AutoTokenizer, Wan2_2Transformer3DModel,
-                               WanT5EncoderModel)
+                               AutoTokenizer, CLIPModel,
+                               Wan2_2Transformer3DModel, WanT5EncoderModel)
 from videox_fun.models.cache_utils import get_teacache_coefficients
 from videox_fun.pipeline import Wan2_2FunInpaintPipeline
-from videox_fun.utils import (FlowDPMSolverMultistepScheduler,
-                              FlowUniPCMultistepScheduler,
-                              apply_gpu_memory_mode, filter_kwargs,
-                              get_image_to_video_latent, merge_lora,
-                              save_videos_grid, unmerge_lora)
+from videox_fun.utils import (register_auto_device_hook,
+                              safe_enable_group_offload)
+from videox_fun.utils.fm_solvers import FlowDPMSolverMultistepScheduler
+from videox_fun.utils.fm_solvers_unipc import FlowUniPCMultistepScheduler
+from videox_fun.utils.fp8_optimization import (convert_model_weight_to_float8,
+                                               convert_weight_dtype_wrapper,
+                                               replace_parameters_by_name)
+from videox_fun.utils.lora_utils import merge_lora, unmerge_lora
+from videox_fun.utils.utils import (filter_kwargs, get_image_to_video_latent,
+                                    save_videos_grid)
 
 # GPU memory mode, which can be chosen in [model_full_load, model_full_load_and_qfloat8, model_cpu_offload, model_cpu_offload_and_qfloat8, model_group_offload, sequential_cpu_offload].
 # model_full_load means that the entire model will be moved to the GPU.
@@ -106,7 +111,7 @@ video_length        = 121
 fps                 = 21
 
 # Use torch.float16 if GPU does not support torch.bfloat16
-# Some graphics cards, such as v100, 2080ti, do not support torch.bfloat16
+# ome graphics cards, such as v100, 2080ti, do not support torch.bfloat16
 weight_dtype            = torch.bfloat16
 # 使用更长的neg prompt如"模糊，突变，变形，失真，画面暗，文本字幕，画面固定，连环画，漫画，线稿，没有主体。"，可以增加稳定性
 # 在neg prompt中添加"安静，固定"等词语可以增加动态性。
@@ -143,7 +148,7 @@ else:
 if transformer_path is not None:
     print(f"From checkpoint: {transformer_path}")
     if transformer_path.endswith("safetensors"):
-        from safetensors.torch import load_file
+        from safetensors.torch import load_file, safe_open
         state_dict = load_file(transformer_path)
     else:
         state_dict = torch.load(transformer_path, map_location="cpu")
@@ -156,7 +161,7 @@ if transformer_2 is not None:
     if transformer_high_path is not None:
         print(f"From checkpoint: {transformer_high_path}")
         if transformer_high_path.endswith("safetensors"):
-            from safetensors.torch import load_file
+            from safetensors.torch import load_file, safe_open
             state_dict = load_file(transformer_high_path)
         else:
             state_dict = torch.load(transformer_high_path, map_location="cpu")
@@ -178,7 +183,7 @@ vae = Chosen_AutoencoderKL.from_pretrained(
 if vae_path is not None:
     print(f"From checkpoint: {vae_path}")
     if vae_path.endswith("safetensors"):
-        from safetensors.torch import load_file
+        from safetensors.torch import load_file, safe_open
         state_dict = load_file(vae_path)
     else:
         state_dict = torch.load(vae_path, map_location="cpu")
@@ -202,7 +207,7 @@ text_encoder = WanT5EncoderModel.from_pretrained(
 text_encoder = text_encoder.eval()
 
 # Get Scheduler
-Chosen_Scheduler = {
+Chosen_Scheduler = scheduler_dict = {
     "Flow": FlowMatchEulerDiscreteScheduler,
     "Flow_Unipc": FlowUniPCMultistepScheduler,
     "Flow_DPM++": FlowDPMSolverMultistepScheduler,
@@ -246,11 +251,36 @@ if compile_dit:
             pipeline.transformer_2.blocks[i] = torch.compile(pipeline.transformer_2.blocks[i])
     print("Add Compile")
 
-# Quantize (when the mode carries an "_and_<quant>" suffix) and then place the pipeline.
-# The order lives inside the helper: quantization has to happen before the offload hooks
-# are installed, and both transformers of this MoE setup are handled in one call, which
-# is exactly the bookkeeping the old 30-line if/elif chain repeated per script.
-apply_gpu_memory_mode(pipeline, GPU_memory_mode, device, weight_dtype)
+if GPU_memory_mode == "sequential_cpu_offload":
+    replace_parameters_by_name(transformer, ["modulation",], device=device)
+    transformer.freqs = transformer.freqs.to(device=device)
+    if transformer_2 is not None:
+        replace_parameters_by_name(transformer_2, ["modulation",], device=device)
+        transformer_2.freqs = transformer_2.freqs.to(device=device)
+    pipeline.enable_sequential_cpu_offload(device=device)
+elif GPU_memory_mode == "model_group_offload":
+    register_auto_device_hook(pipeline.transformer)
+    if transformer_2 is not None:
+        register_auto_device_hook(pipeline.transformer_2)
+    safe_enable_group_offload(pipeline, onload_device=device, offload_device="cpu", offload_type="leaf_level", use_stream=True)
+elif GPU_memory_mode == "model_cpu_offload_and_qfloat8":
+    convert_model_weight_to_float8(transformer, exclude_module_name=["modulation",], device=device)
+    convert_weight_dtype_wrapper(transformer, weight_dtype)
+    if transformer_2 is not None:
+        convert_model_weight_to_float8(transformer_2, exclude_module_name=["modulation",], device=device)
+        convert_weight_dtype_wrapper(transformer_2, weight_dtype)
+    pipeline.enable_model_cpu_offload(device=device)
+elif GPU_memory_mode == "model_cpu_offload":
+    pipeline.enable_model_cpu_offload(device=device)
+elif GPU_memory_mode == "model_full_load_and_qfloat8":
+    convert_model_weight_to_float8(transformer, exclude_module_name=["modulation",], device=device)
+    convert_weight_dtype_wrapper(transformer, weight_dtype)
+    if transformer_2 is not None:
+        convert_model_weight_to_float8(transformer_2, exclude_module_name=["modulation",], device=device)
+        convert_weight_dtype_wrapper(transformer_2, weight_dtype)
+    pipeline.to(device=device)
+else:
+    pipeline.to(device=device)
 
 coefficients = get_teacache_coefficients(model_name) if enable_teacache else None
 if coefficients is not None:

examples/wan2.2_fun/predict_v2v_control.py

@@ -16,16 +16,22 @@ for project_root in project_roots:
 from videox_fun.data import process_pose_file
 from videox_fun.dist import set_multi_gpus_devices, shard_model
 from videox_fun.models import (AutoencoderKLWan, AutoencoderKLWan3_8,
-                               AutoTokenizer, Wan2_2Transformer3DModel,
-                               WanT5EncoderModel)
+                               AutoTokenizer, CLIPModel,
+                               Wan2_2Transformer3DModel, WanT5EncoderModel)
 from videox_fun.models.cache_utils import get_teacache_coefficients
-from videox_fun.pipeline import Wan2_2FunControlPipeline
-from videox_fun.utils import (FlowDPMSolverMultistepScheduler,
-                              FlowUniPCMultistepScheduler,
-                              apply_gpu_memory_mode, filter_kwargs,
-                              get_image_latent, get_image_to_video_latent,
-                              get_video_to_video_latent, merge_lora,
-                              save_videos_grid, unmerge_lora)
+from videox_fun.pipeline import Wan2_2FunControlPipeline, WanPipeline
+from videox_fun.utils import (register_auto_device_hook,
+                              safe_enable_group_offload)
+from videox_fun.utils.fm_solvers import FlowDPMSolverMultistepScheduler
+from videox_fun.utils.fm_solvers_unipc import FlowUniPCMultistepScheduler
+from videox_fun.utils.fp8_optimization import (convert_model_weight_to_float8,
+                                               convert_weight_dtype_wrapper,
+                                               replace_parameters_by_name)
+from videox_fun.utils.lora_utils import merge_lora, unmerge_lora
+from videox_fun.utils.utils import (filter_kwargs, get_image_latent,
+                                    get_image_to_video_latent,
+                                    get_video_to_video_latent,
+                                    save_videos_grid)
 
 # GPU memory mode, which can be chosen in [model_full_load, model_full_load_and_qfloat8, model_cpu_offload, model_cpu_offload_and_qfloat8, model_group_offload, sequential_cpu_offload].
 # model_full_load means that the entire model will be moved to the GPU.
@@ -111,7 +117,7 @@ video_length        = 81
 fps                 = 16
 
 # Use torch.float16 if GPU does not support torch.bfloat16
-# Some graphics cards, such as v100, 2080ti, do not support torch.bfloat16
+# ome graphics cards, such as v100, 2080ti, do not support torch.bfloat16
 weight_dtype            = torch.bfloat16
 control_video           = "asset/pose.mp4"
 control_camera_txt      = None
@@ -159,7 +165,7 @@ else:
 if transformer_path is not None:
     print(f"From checkpoint: {transformer_path}")
     if transformer_path.endswith("safetensors"):
-        from safetensors.torch import load_file
+        from safetensors.torch import load_file, safe_open
         state_dict = load_file(transformer_path)
     else:
         state_dict = torch.load(transformer_path, map_location="cpu")
@@ -172,7 +178,7 @@ if transformer_2 is not None:
     if transformer_high_path is not None:
         print(f"From checkpoint: {transformer_high_path}")
         if transformer_high_path.endswith("safetensors"):
-            from safetensors.torch import load_file
+            from safetensors.torch import load_file, safe_open
             state_dict = load_file(transformer_high_path)
         else:
             state_dict = torch.load(transformer_high_path, map_location="cpu")
@@ -194,7 +200,7 @@ vae = Chosen_AutoencoderKL.from_pretrained(
 if vae_path is not None:
     print(f"From checkpoint: {vae_path}")
     if vae_path.endswith("safetensors"):
-        from safetensors.torch import load_file
+        from safetensors.torch import load_file, safe_open
         state_dict = load_file(vae_path)
     else:
         state_dict = torch.load(vae_path, map_location="cpu")
@@ -218,7 +224,7 @@ text_encoder = WanT5EncoderModel.from_pretrained(
 text_encoder = text_encoder.eval()
 
 # Get Scheduler
-Chosen_Scheduler = {
+Chosen_Scheduler = scheduler_dict = {
     "Flow": FlowMatchEulerDiscreteScheduler,
     "Flow_Unipc": FlowUniPCMultistepScheduler,
     "Flow_DPM++": FlowDPMSolverMultistepScheduler,
@@ -262,11 +268,36 @@ if compile_dit:
             pipeline.transformer_2.blocks[i] = torch.compile(pipeline.transformer_2.blocks[i])
     print("Add Compile")
 
-# Quantize (when the mode carries an "_and_<quant>" suffix) and then place the pipeline.
-# The order lives inside the helper: quantization has to happen before the offload hooks
-# are installed, and both transformers of this MoE setup are handled in one call, which
-# is exactly the bookkeeping the old 30-line if/elif chain repeated per script.
-apply_gpu_memory_mode(pipeline, GPU_memory_mode, device, weight_dtype)
+if GPU_memory_mode == "sequential_cpu_offload":
+    replace_parameters_by_name(transformer, ["modulation",], device=device)
+    transformer.freqs = transformer.freqs.to(device=device)
+    if transformer_2 is not None:
+        replace_parameters_by_name(transformer_2, ["modulation",], device=device)
+        transformer_2.freqs = transformer_2.freqs.to(device=device)
+    pipeline.enable_sequential_cpu_offload(device=device)
+elif GPU_memory_mode == "model_group_offload":
+    register_auto_device_hook(pipeline.transformer)
+    if transformer_2 is not None:
+        register_auto_device_hook(pipeline.transformer_2)
+    safe_enable_group_offload(pipeline, onload_device=device, offload_device="cpu", offload_type="leaf_level", use_stream=True)
+elif GPU_memory_mode == "model_cpu_offload_and_qfloat8":
+    convert_model_weight_to_float8(transformer, exclude_module_name=["modulation",], device=device)
+    convert_weight_dtype_wrapper(transformer, weight_dtype)
+    if transformer_2 is not None:
+        convert_model_weight_to_float8(transformer_2, exclude_module_name=["modulation",], device=device)
+        convert_weight_dtype_wrapper(transformer_2, weight_dtype)
+    pipeline.enable_model_cpu_offload(device=device)
+elif GPU_memory_mode == "model_cpu_offload":
+    pipeline.enable_model_cpu_offload(device=device)
+elif GPU_memory_mode == "model_full_load_and_qfloat8":
+    convert_model_weight_to_float8(transformer, exclude_module_name=["modulation",], device=device)
+    convert_weight_dtype_wrapper(transformer, weight_dtype)
+    if transformer_2 is not None:
+        convert_model_weight_to_float8(transformer_2, exclude_module_name=["modulation",], device=device)
+        convert_weight_dtype_wrapper(transformer_2, weight_dtype)
+    pipeline.to(device=device)
+else:
+    pipeline.to(device=device)
 
 coefficients = get_teacache_coefficients(model_name) if enable_teacache else None
 if coefficients is not None:

examples/wan2.2_fun/predict_v2v_control_5b.py

@@ -16,16 +16,22 @@ for project_root in project_roots:
 from videox_fun.data import process_pose_file
 from videox_fun.dist import set_multi_gpus_devices, shard_model
 from videox_fun.models import (AutoencoderKLWan, AutoencoderKLWan3_8,
-                               AutoTokenizer, Wan2_2Transformer3DModel,
-                               WanT5EncoderModel)
+                               AutoTokenizer, CLIPModel,
+                               Wan2_2Transformer3DModel, WanT5EncoderModel)
 from videox_fun.models.cache_utils import get_teacache_coefficients
-from videox_fun.pipeline import Wan2_2FunControlPipeline
-from videox_fun.utils import (FlowDPMSolverMultistepScheduler,
-                              FlowUniPCMultistepScheduler,
-                              apply_gpu_memory_mode, filter_kwargs,
-                              get_image_latent, get_image_to_video_latent,
-                              get_video_to_video_latent, merge_lora,
-                              save_videos_grid, unmerge_lora)
+from videox_fun.pipeline import Wan2_2FunControlPipeline, WanPipeline
+from videox_fun.utils import (register_auto_device_hook,
+                              safe_enable_group_offload)
+from videox_fun.utils.fm_solvers import FlowDPMSolverMultistepScheduler
+from videox_fun.utils.fm_solvers_unipc import FlowUniPCMultistepScheduler
+from videox_fun.utils.fp8_optimization import (convert_model_weight_to_float8,
+                                               convert_weight_dtype_wrapper,
+                                               replace_parameters_by_name)
+from videox_fun.utils.lora_utils import merge_lora, unmerge_lora
+from videox_fun.utils.utils import (filter_kwargs, get_image_latent,
+                                    get_image_to_video_latent,
+                                    get_video_to_video_latent,
+                                    save_videos_grid)
 
 # GPU memory mode, which can be chosen in [model_full_load, model_full_load_and_qfloat8, model_cpu_offload, model_cpu_offload_and_qfloat8, model_group_offload, sequential_cpu_offload].
 # model_full_load means that the entire model will be moved to the GPU.
@@ -111,7 +117,7 @@ video_length        = 121
 fps                 = 24
 
 # Use torch.float16 if GPU does not support torch.bfloat16
-# Some graphics cards, such as v100, 2080ti, do not support torch.bfloat16
+# ome graphics cards, such as v100, 2080ti, do not support torch.bfloat16
 weight_dtype            = torch.bfloat16
 control_video           = "asset/pose.mp4"
 control_camera_txt      = None
@@ -159,7 +165,7 @@ else:
 if transformer_path is not None:
     print(f"From checkpoint: {transformer_path}")
     if transformer_path.endswith("safetensors"):
-        from safetensors.torch import load_file
+        from safetensors.torch import load_file, safe_open
         state_dict = load_file(transformer_path)
     else:
         state_dict = torch.load(transformer_path, map_location="cpu")
@@ -172,7 +178,7 @@ if transformer_2 is not None:
     if transformer_high_path is not None:
         print(f"From checkpoint: {transformer_high_path}")
         if transformer_high_path.endswith("safetensors"):
-            from safetensors.torch import load_file
+            from safetensors.torch import load_file, safe_open
             state_dict = load_file(transformer_high_path)
         else:
             state_dict = torch.load(transformer_high_path, map_location="cpu")
@@ -194,7 +200,7 @@ vae = Chosen_AutoencoderKL.from_pretrained(
 if vae_path is not None:
     print(f"From checkpoint: {vae_path}")
     if vae_path.endswith("safetensors"):
-        from safetensors.torch import load_file
+        from safetensors.torch import load_file, safe_open
         state_dict = load_file(vae_path)
     else:
         state_dict = torch.load(vae_path, map_location="cpu")
@@ -218,7 +224,7 @@ text_encoder = WanT5EncoderModel.from_pretrained(
 text_encoder = text_encoder.eval()
 
 # Get Scheduler
-Chosen_Scheduler = {
+Chosen_Scheduler = scheduler_dict = {
     "Flow": FlowMatchEulerDiscreteScheduler,
     "Flow_Unipc": FlowUniPCMultistepScheduler,
     "Flow_DPM++": FlowDPMSolverMultistepScheduler,
@@ -262,11 +268,36 @@ if compile_dit:
             pipeline.transformer_2.blocks[i] = torch.compile(pipeline.transformer_2.blocks[i])
     print("Add Compile")
 
-# Quantize (when the mode carries an "_and_<quant>" suffix) and then place the pipeline.
-# The order lives inside the helper: quantization has to happen before the offload hooks
-# are installed, and both transformers of this MoE setup are handled in one call, which
-# is exactly the bookkeeping the old 30-line if/elif chain repeated per script.
-apply_gpu_memory_mode(pipeline, GPU_memory_mode, device, weight_dtype)
+if GPU_memory_mode == "sequential_cpu_offload":
+    replace_parameters_by_name(transformer, ["modulation",], device=device)
+    transformer.freqs = transformer.freqs.to(device=device)
+    if transformer_2 is not None:
+        replace_parameters_by_name(transformer_2, ["modulation",], device=device)
+        transformer_2.freqs = transformer_2.freqs.to(device=device)
+    pipeline.enable_sequential_cpu_offload(device=device)
+elif GPU_memory_mode == "model_group_offload":
+    register_auto_device_hook(pipeline.transformer)
+    if transformer_2 is not None:
+        register_auto_device_hook(pipeline.transformer_2)
+    safe_enable_group_offload(pipeline, onload_device=device, offload_device="cpu", offload_type="leaf_level", use_stream=True)
+elif GPU_memory_mode == "model_cpu_offload_and_qfloat8":
+    convert_model_weight_to_float8(transformer, exclude_module_name=["modulation",], device=device)
+    convert_weight_dtype_wrapper(transformer, weight_dtype)
+    if transformer_2 is not None:
+        convert_model_weight_to_float8(transformer_2, exclude_module_name=["modulation",], device=device)
+        convert_weight_dtype_wrapper(transformer_2, weight_dtype)
+    pipeline.enable_model_cpu_offload(device=device)
+elif GPU_memory_mode == "model_cpu_offload":
+    pipeline.enable_model_cpu_offload(device=device)
+elif GPU_memory_mode == "model_full_load_and_qfloat8":
+    convert_model_weight_to_float8(transformer, exclude_module_name=["modulation",], device=device)
+    convert_weight_dtype_wrapper(transformer, weight_dtype)
+    if transformer_2 is not None:
+        convert_model_weight_to_float8(transformer_2, exclude_module_name=["modulation",], device=device)
+        convert_weight_dtype_wrapper(transformer_2, weight_dtype)
+    pipeline.to(device=device)
+else:
+    pipeline.to(device=device)
 
 coefficients = get_teacache_coefficients(model_name) if enable_teacache else None
 if coefficients is not None:

examples/wan2.2_fun/predict_v2v_control_camera.py

@@ -16,16 +16,22 @@ for project_root in project_roots:
 from videox_fun.data import process_pose_file
 from videox_fun.dist import set_multi_gpus_devices, shard_model
 from videox_fun.models import (AutoencoderKLWan, AutoencoderKLWan3_8,
-                               AutoTokenizer, Wan2_2Transformer3DModel,
-                               WanT5EncoderModel)
+                               AutoTokenizer, CLIPModel,
+                               Wan2_2Transformer3DModel, WanT5EncoderModel)
 from videox_fun.models.cache_utils import get_teacache_coefficients
-from videox_fun.pipeline import Wan2_2FunControlPipeline
-from videox_fun.utils import (FlowDPMSolverMultistepScheduler,
-                              FlowUniPCMultistepScheduler,
-                              apply_gpu_memory_mode, filter_kwargs,
-                              get_image_latent, get_image_to_video_latent,
-                              get_video_to_video_latent, merge_lora,
-                              save_videos_grid, unmerge_lora)
+from videox_fun.pipeline import Wan2_2FunControlPipeline, WanPipeline
+from videox_fun.utils import (register_auto_device_hook,
+                              safe_enable_group_offload)
+from videox_fun.utils.fm_solvers import FlowDPMSolverMultistepScheduler
+from videox_fun.utils.fm_solvers_unipc import FlowUniPCMultistepScheduler
+from videox_fun.utils.fp8_optimization import (convert_model_weight_to_float8,
+                                               convert_weight_dtype_wrapper,
+                                               replace_parameters_by_name)
+from videox_fun.utils.lora_utils import merge_lora, unmerge_lora
+from videox_fun.utils.utils import (filter_kwargs, get_image_latent,
+                                    get_image_to_video_latent,
+                                    get_video_to_video_latent,
+                                    save_videos_grid)
 
 # GPU memory mode, which can be chosen in [model_full_load, model_full_load_and_qfloat8, model_cpu_offload, model_cpu_offload_and_qfloat8, model_group_offload, sequential_cpu_offload].
 # model_full_load means that the entire model will be moved to the GPU.
@@ -111,7 +117,7 @@ video_length        = 81
 fps                 = 16
 
 # Use torch.float16 if GPU does not support torch.bfloat16
-# Some graphics cards, such as v100, 2080ti, do not support torch.bfloat16
+# ome graphics cards, such as v100, 2080ti, do not support torch.bfloat16
 weight_dtype            = torch.bfloat16
 control_video           = None
 control_camera_txt      = "asset/Zoom_In.txt"
@@ -159,7 +165,7 @@ else:
 if transformer_path is not None:
     print(f"From checkpoint: {transformer_path}")
     if transformer_path.endswith("safetensors"):
-        from safetensors.torch import load_file
+        from safetensors.torch import load_file, safe_open
         state_dict = load_file(transformer_path)
     else:
         state_dict = torch.load(transformer_path, map_location="cpu")
@@ -172,7 +178,7 @@ if transformer_2 is not None:
     if transformer_high_path is not None:
         print(f"From checkpoint: {transformer_high_path}")
         if transformer_high_path.endswith("safetensors"):
-            from safetensors.torch import load_file
+            from safetensors.torch import load_file, safe_open
             state_dict = load_file(transformer_high_path)
         else:
             state_dict = torch.load(transformer_high_path, map_location="cpu")
@@ -194,7 +200,7 @@ vae = Chosen_AutoencoderKL.from_pretrained(
 if vae_path is not None:
     print(f"From checkpoint: {vae_path}")
     if vae_path.endswith("safetensors"):
-        from safetensors.torch import load_file
+        from safetensors.torch import load_file, safe_open
         state_dict = load_file(vae_path)
     else:
         state_dict = torch.load(vae_path, map_location="cpu")
@@ -218,7 +224,7 @@ text_encoder = WanT5EncoderModel.from_pretrained(
 text_encoder = text_encoder.eval()
 
 # Get Scheduler
-Chosen_Scheduler = {
+Chosen_Scheduler = scheduler_dict = {
     "Flow": FlowMatchEulerDiscreteScheduler,
     "Flow_Unipc": FlowUniPCMultistepScheduler,
     "Flow_DPM++": FlowDPMSolverMultistepScheduler,
@@ -262,11 +268,36 @@ if compile_dit:
             pipeline.transformer_2.blocks[i] = torch.compile(pipeline.transformer_2.blocks[i])
     print("Add Compile")
 
-# Quantize (when the mode carries an "_and_<quant>" suffix) and then place the pipeline.
-# The order lives inside the helper: quantization has to happen before the offload hooks
-# are installed, and both transformers of this MoE setup are handled in one call, which
-# is exactly the bookkeeping the old 30-line if/elif chain repeated per script.
-apply_gpu_memory_mode(pipeline, GPU_memory_mode, device, weight_dtype)
+if GPU_memory_mode == "sequential_cpu_offload":
+    replace_parameters_by_name(transformer, ["modulation",], device=device)
+    transformer.freqs = transformer.freqs.to(device=device)
+    if transformer_2 is not None:
+        replace_parameters_by_name(transformer_2, ["modulation",], device=device)
+        transformer_2.freqs = transformer_2.freqs.to(device=device)
+    pipeline.enable_sequential_cpu_offload(device=device)
+elif GPU_memory_mode == "model_group_offload":
+    register_auto_device_hook(pipeline.transformer)
+    if transformer_2 is not None:
+        register_auto_device_hook(pipeline.transformer_2)
+    safe_enable_group_offload(pipeline, onload_device=device, offload_device="cpu", offload_type="leaf_level", use_stream=True)
+elif GPU_memory_mode == "model_cpu_offload_and_qfloat8":
+    convert_model_weight_to_float8(transformer, exclude_module_name=["modulation",], device=device)
+    convert_weight_dtype_wrapper(transformer, weight_dtype)
+    if transformer_2 is not None:
+        convert_model_weight_to_float8(transformer_2, exclude_module_name=["modulation",], device=device)
+        convert_weight_dtype_wrapper(transformer_2, weight_dtype)
+    pipeline.enable_model_cpu_offload(device=device)
+elif GPU_memory_mode == "model_cpu_offload":
+    pipeline.enable_model_cpu_offload(device=device)
+elif GPU_memory_mode == "model_full_load_and_qfloat8":
+    convert_model_weight_to_float8(transformer, exclude_module_name=["modulation",], device=device)
+    convert_weight_dtype_wrapper(transformer, weight_dtype)
+    if transformer_2 is not None:
+        convert_model_weight_to_float8(transformer_2, exclude_module_name=["modulation",], device=device)
+        convert_weight_dtype_wrapper(transformer_2, weight_dtype)
+    pipeline.to(device=device)
+else:
+    pipeline.to(device=device)
 
 coefficients = get_teacache_coefficients(model_name) if enable_teacache else None
 if coefficients is not None:

examples/wan2.2_fun/predict_v2v_control_camera_5b.py

@@ -16,16 +16,22 @@ for project_root in project_roots:
 from videox_fun.data import process_pose_file
 from videox_fun.dist import set_multi_gpus_devices, shard_model
 from videox_fun.models import (AutoencoderKLWan, AutoencoderKLWan3_8,
-                               AutoTokenizer, Wan2_2Transformer3DModel,
-                               WanT5EncoderModel)
+                               AutoTokenizer, CLIPModel,
+                               Wan2_2Transformer3DModel, WanT5EncoderModel)
 from videox_fun.models.cache_utils import get_teacache_coefficients
-from videox_fun.pipeline import Wan2_2FunControlPipeline
-from videox_fun.utils import (FlowDPMSolverMultistepScheduler,
-                              FlowUniPCMultistepScheduler,
-                              apply_gpu_memory_mode, filter_kwargs,
-                              get_image_latent, get_image_to_video_latent,
-                              get_video_to_video_latent, merge_lora,
-                              save_videos_grid, unmerge_lora)
+from videox_fun.pipeline import Wan2_2FunControlPipeline, WanPipeline
+from videox_fun.utils import (register_auto_device_hook,
+                              safe_enable_group_offload)
+from videox_fun.utils.fm_solvers import FlowDPMSolverMultistepScheduler
+from videox_fun.utils.fm_solvers_unipc import FlowUniPCMultistepScheduler
+from videox_fun.utils.fp8_optimization import (convert_model_weight_to_float8,
+                                               convert_weight_dtype_wrapper,
+                                               replace_parameters_by_name)
+from videox_fun.utils.lora_utils import merge_lora, unmerge_lora
+from videox_fun.utils.utils import (filter_kwargs, get_image_latent,
+                                    get_image_to_video_latent,
+                                    get_video_to_video_latent,
+                                    save_videos_grid)
 
 # GPU memory mode, which can be chosen in [model_full_load, model_full_load_and_qfloat8, model_cpu_offload, model_cpu_offload_and_qfloat8, model_group_offload, sequential_cpu_offload].
 # model_full_load means that the entire model will be moved to the GPU.
@@ -111,7 +117,7 @@ video_length        = 81
 fps                 = 24
 
 # Use torch.float16 if GPU does not support torch.bfloat16
-# Some graphics cards, such as v100, 2080ti, do not support torch.bfloat16
+# ome graphics cards, such as v100, 2080ti, do not support torch.bfloat16
 weight_dtype            = torch.bfloat16
 control_video           = None
 control_camera_txt      = "asset/Zoom_In.txt"
@@ -159,7 +165,7 @@ else:
 if transformer_path is not None:
     print(f"From checkpoint: {transformer_path}")
     if transformer_path.endswith("safetensors"):
-        from safetensors.torch import load_file
+        from safetensors.torch import load_file, safe_open
         state_dict = load_file(transformer_path)
     else:
         state_dict = torch.load(transformer_path, map_location="cpu")
@@ -172,7 +178,7 @@ if transformer_2 is not None:
     if transformer_high_path is not None:
         print(f"From checkpoint: {transformer_high_path}")
         if transformer_high_path.endswith("safetensors"):
-            from safetensors.torch import load_file
+            from safetensors.torch import load_file, safe_open
             state_dict = load_file(transformer_high_path)
         else:
             state_dict = torch.load(transformer_high_path, map_location="cpu")
@@ -194,7 +200,7 @@ vae = Chosen_AutoencoderKL.from_pretrained(
 if vae_path is not None:
     print(f"From checkpoint: {vae_path}")
     if vae_path.endswith("safetensors"):
-        from safetensors.torch import load_file
+        from safetensors.torch import load_file, safe_open
         state_dict = load_file(vae_path)
     else:
         state_dict = torch.load(vae_path, map_location="cpu")
@@ -218,7 +224,7 @@ text_encoder = WanT5EncoderModel.from_pretrained(
 text_encoder = text_encoder.eval()
 
 # Get Scheduler
-Chosen_Scheduler = {
+Chosen_Scheduler = scheduler_dict = {
     "Flow": FlowMatchEulerDiscreteScheduler,
     "Flow_Unipc": FlowUniPCMultistepScheduler,
     "Flow_DPM++": FlowDPMSolverMultistepScheduler,
@@ -262,11 +268,36 @@ if compile_dit:
             pipeline.transformer_2.blocks[i] = torch.compile(pipeline.transformer_2.blocks[i])
     print("Add Compile")
 
-# Quantize (when the mode carries an "_and_<quant>" suffix) and then place the pipeline.
-# The order lives inside the helper: quantization has to happen before the offload hooks
-# are installed, and both transformers of this MoE setup are handled in one call, which
-# is exactly the bookkeeping the old 30-line if/elif chain repeated per script.
-apply_gpu_memory_mode(pipeline, GPU_memory_mode, device, weight_dtype)
+if GPU_memory_mode == "sequential_cpu_offload":
+    replace_parameters_by_name(transformer, ["modulation",], device=device)
+    transformer.freqs = transformer.freqs.to(device=device)
+    if transformer_2 is not None:
+        replace_parameters_by_name(transformer_2, ["modulation",], device=device)
+        transformer_2.freqs = transformer_2.freqs.to(device=device)
+    pipeline.enable_sequential_cpu_offload(device=device)
+elif GPU_memory_mode == "model_group_offload":
+    register_auto_device_hook(pipeline.transformer)
+    if transformer_2 is not None:
+        register_auto_device_hook(pipeline.transformer_2)
+    safe_enable_group_offload(pipeline, onload_device=device, offload_device="cpu", offload_type="leaf_level", use_stream=True)
+elif GPU_memory_mode == "model_cpu_offload_and_qfloat8":
+    convert_model_weight_to_float8(transformer, exclude_module_name=["modulation",], device=device)
+    convert_weight_dtype_wrapper(transformer, weight_dtype)
+    if transformer_2 is not None:
+        convert_model_weight_to_float8(transformer_2, exclude_module_name=["modulation",], device=device)
+        convert_weight_dtype_wrapper(transformer_2, weight_dtype)
+    pipeline.enable_model_cpu_offload(device=device)
+elif GPU_memory_mode == "model_cpu_offload":
+    pipeline.enable_model_cpu_offload(device=device)
+elif GPU_memory_mode == "model_full_load_and_qfloat8":
+    convert_model_weight_to_float8(transformer, exclude_module_name=["modulation",], device=device)
+    convert_weight_dtype_wrapper(transformer, weight_dtype)
+    if transformer_2 is not None:
+        convert_model_weight_to_float8(transformer_2, exclude_module_name=["modulation",], device=device)
+        convert_weight_dtype_wrapper(transformer_2, weight_dtype)
+    pipeline.to(device=device)
+else:
+    pipeline.to(device=device)
 
 coefficients = get_teacache_coefficients(model_name) if enable_teacache else None
 if coefficients is not None:

examples/wan2.2_fun/predict_v2v_control_ref.py

@@ -16,16 +16,22 @@ for project_root in project_roots:
 from videox_fun.data import process_pose_file
 from videox_fun.dist import set_multi_gpus_devices, shard_model
 from videox_fun.models import (AutoencoderKLWan, AutoencoderKLWan3_8,
-                               AutoTokenizer, Wan2_2Transformer3DModel,
-                               WanT5EncoderModel)
+                               AutoTokenizer, CLIPModel,
+                               Wan2_2Transformer3DModel, WanT5EncoderModel)
 from videox_fun.models.cache_utils import get_teacache_coefficients
-from videox_fun.pipeline import Wan2_2FunControlPipeline
-from videox_fun.utils import (FlowDPMSolverMultistepScheduler,
-                              FlowUniPCMultistepScheduler,
-                              apply_gpu_memory_mode, filter_kwargs,
-                              get_image_latent, get_image_to_video_latent,
-                              get_video_to_video_latent, merge_lora,
-                              save_videos_grid, unmerge_lora)
+from videox_fun.pipeline import Wan2_2FunControlPipeline, WanPipeline
+from videox_fun.utils import (register_auto_device_hook,
+                              safe_enable_group_offload)
+from videox_fun.utils.fm_solvers import FlowDPMSolverMultistepScheduler
+from videox_fun.utils.fm_solvers_unipc import FlowUniPCMultistepScheduler
+from videox_fun.utils.fp8_optimization import (convert_model_weight_to_float8,
+                                               convert_weight_dtype_wrapper,
+                                               replace_parameters_by_name)
+from videox_fun.utils.lora_utils import merge_lora, unmerge_lora
+from videox_fun.utils.utils import (filter_kwargs, get_image_latent,
+                                    get_image_to_video_latent,
+                                    get_video_to_video_latent,
+                                    save_videos_grid)
 
 # GPU memory mode, which can be chosen in [model_full_load, model_full_load_and_qfloat8, model_cpu_offload, model_cpu_offload_and_qfloat8, model_group_offload, sequential_cpu_offload].
 # model_full_load means that the entire model will be moved to the GPU.
@@ -111,7 +117,7 @@ video_length        = 81
 fps                 = 16
 
 # Use torch.float16 if GPU does not support torch.bfloat16
-# Some graphics cards, such as v100, 2080ti, do not support torch.bfloat16
+# ome graphics cards, such as v100, 2080ti, do not support torch.bfloat16
 weight_dtype            = torch.bfloat16
 control_video           = "asset/pose.mp4"
 control_camera_txt      = None
@@ -159,7 +165,7 @@ else:
 if transformer_path is not None:
     print(f"From checkpoint: {transformer_path}")
     if transformer_path.endswith("safetensors"):
-        from safetensors.torch import load_file
+        from safetensors.torch import load_file, safe_open
         state_dict = load_file(transformer_path)
     else:
         state_dict = torch.load(transformer_path, map_location="cpu")
@@ -172,7 +178,7 @@ if transformer_2 is not None:
     if transformer_high_path is not None:
         print(f"From checkpoint: {transformer_high_path}")
         if transformer_high_path.endswith("safetensors"):
-            from safetensors.torch import load_file
+            from safetensors.torch import load_file, safe_open
             state_dict = load_file(transformer_high_path)
         else:
             state_dict = torch.load(transformer_high_path, map_location="cpu")
@@ -194,7 +200,7 @@ vae = Chosen_AutoencoderKL.from_pretrained(
 if vae_path is not None:
     print(f"From checkpoint: {vae_path}")
     if vae_path.endswith("safetensors"):
-        from safetensors.torch import load_file
+        from safetensors.torch import load_file, safe_open
         state_dict = load_file(vae_path)
     else:
         state_dict = torch.load(vae_path, map_location="cpu")
@@ -218,7 +224,7 @@ text_encoder = WanT5EncoderModel.from_pretrained(
 text_encoder = text_encoder.eval()
 
 # Get Scheduler
-Chosen_Scheduler = {
+Chosen_Scheduler = scheduler_dict = {
     "Flow": FlowMatchEulerDiscreteScheduler,
     "Flow_Unipc": FlowUniPCMultistepScheduler,
     "Flow_DPM++": FlowDPMSolverMultistepScheduler,
@@ -262,11 +268,36 @@ if compile_dit:
             pipeline.transformer_2.blocks[i] = torch.compile(pipeline.transformer_2.blocks[i])
     print("Add Compile")
 
-# Quantize (when the mode carries an "_and_<quant>" suffix) and then place the pipeline.
-# The order lives inside the helper: quantization has to happen before the offload hooks
-# are installed, and both transformers of this MoE setup are handled in one call, which
-# is exactly the bookkeeping the old 30-line if/elif chain repeated per script.
-apply_gpu_memory_mode(pipeline, GPU_memory_mode, device, weight_dtype)
+if GPU_memory_mode == "sequential_cpu_offload":
+    replace_parameters_by_name(transformer, ["modulation",], device=device)
+    transformer.freqs = transformer.freqs.to(device=device)
+    if transformer_2 is not None:
+        replace_parameters_by_name(transformer_2, ["modulation",], device=device)
+        transformer_2.freqs = transformer_2.freqs.to(device=device)
+    pipeline.enable_sequential_cpu_offload(device=device)
+elif GPU_memory_mode == "model_group_offload":
+    register_auto_device_hook(pipeline.transformer)
+    if transformer_2 is not None:
+        register_auto_device_hook(pipeline.transformer_2)
+    safe_enable_group_offload(pipeline, onload_device=device, offload_device="cpu", offload_type="leaf_level", use_stream=True)
+elif GPU_memory_mode == "model_cpu_offload_and_qfloat8":
+    convert_model_weight_to_float8(transformer, exclude_module_name=["modulation",], device=device)
+    convert_weight_dtype_wrapper(transformer, weight_dtype)
+    if transformer_2 is not None:
+        convert_model_weight_to_float8(transformer_2, exclude_module_name=["modulation",], device=device)
+        convert_weight_dtype_wrapper(transformer_2, weight_dtype)
+    pipeline.enable_model_cpu_offload(device=device)
+elif GPU_memory_mode == "model_cpu_offload":
+    pipeline.enable_model_cpu_offload(device=device)
+elif GPU_memory_mode == "model_full_load_and_qfloat8":
+    convert_model_weight_to_float8(transformer, exclude_module_name=["modulation",], device=device)
+    convert_weight_dtype_wrapper(transformer, weight_dtype)
+    if transformer_2 is not None:
+        convert_model_weight_to_float8(transformer_2, exclude_module_name=["modulation",], device=device)
+        convert_weight_dtype_wrapper(transformer_2, weight_dtype)
+    pipeline.to(device=device)
+else:
+    pipeline.to(device=device)
 
 coefficients = get_teacache_coefficients(model_name) if enable_teacache else None
 if coefficients is not None: