137 lines
6.9 KiB
Python
Executable File
137 lines
6.9 KiB
Python
Executable File
import importlib.util
|
|
|
|
from diffusers import AutoencoderKL
|
|
from transformers import (AutoProcessor, AutoTokenizer, CLIPImageProcessor,
|
|
CLIPTextModel, CLIPTokenizer,
|
|
CLIPVisionModelWithProjection, LlamaModel,
|
|
LlamaTokenizerFast, LlavaForConditionalGeneration,
|
|
Mistral3ForConditionalGeneration, PixtralProcessor,
|
|
Qwen3Config, Qwen3ForCausalLM, T5EncoderModel,
|
|
T5Tokenizer, T5TokenizerFast, UMT5EncoderModel)
|
|
|
|
try:
|
|
from transformers import (Qwen2_5_VLConfig,
|
|
Qwen2_5_VLForConditionalGeneration,
|
|
Qwen2Tokenizer, Qwen2VLProcessor)
|
|
except:
|
|
Qwen2_5_VLForConditionalGeneration, Qwen2Tokenizer = None, None
|
|
Qwen2VLProcessor, Qwen2_5_VLConfig = None, None
|
|
print("Your transformers version is too old to load Qwen2_5_VLForConditionalGeneration and Qwen2Tokenizer. If you wish to use QwenImage, please upgrade your transformers package to the latest version.")
|
|
|
|
from .cogvideox_transformer3d import CogVideoXTransformer3DModel
|
|
from .cogvideox_vae import AutoencoderKLCogVideoX
|
|
from .fantasytalking_audio_encoder import FantasyTalkingAudioEncoder
|
|
from .fantasytalking_transformer3d import FantasyTalkingTransformer3DModel
|
|
from .flux2_image_processor import Flux2ImageProcessor
|
|
from .flux2_transformer2d import Flux2Transformer2DModel
|
|
from .flux2_transformer2d_control import Flux2ControlTransformer2DModel
|
|
from .flux2_vae import AutoencoderKLFlux2
|
|
from .flux_transformer2d import FluxTransformer2DModel
|
|
from .hunyuanvideo_transformer3d import HunyuanVideoTransformer3DModel
|
|
from .hunyuanvideo_vae import AutoencoderKLHunyuanVideo
|
|
from .longcatvideo_transformer3d import LongCatVideoTransformer3DModel
|
|
from .longcatvideo_vae import AutoencoderKLLongCatVideo
|
|
from .qwenimage_transformer2d import QwenImageTransformer2DModel
|
|
from .qwenimage_transformer2d_control import QwenImageControlTransformer2DModel
|
|
from .qwenimage_transformer2d_instantx import QwenImageInstantXControlNetModel
|
|
from .qwenimage_vae import AutoencoderKLQwenImage
|
|
from .turbowan_transformer3d import TurboWanTransformer3DModel
|
|
from .wan_audio_encoder import WanAudioEncoder
|
|
from .wan_image_encoder import CLIPModel
|
|
from .wan_text_encoder import WanT5EncoderModel
|
|
from .wan_transformer3d import (Wan2_2Transformer3DModel, WanRMSNorm,
|
|
WanSelfAttention, WanTransformer3DModel)
|
|
from .wan_transformer3d_animate import Wan2_2Transformer3DModel_Animate
|
|
from .wan_transformer3d_lingbot import LingBotWanTransformer3DModel
|
|
from .wan_transformer3d_s2v import Wan2_2Transformer3DModel_S2V
|
|
from .wan_transformer3d_vace import VaceWanTransformer3DModel
|
|
from .wan_vae import AutoencoderKLWan, AutoencoderKLWan_
|
|
from .wan_vae3_8 import AutoencoderKLWan2_2_, AutoencoderKLWan3_8
|
|
from .z_image_transformer2d import ZImageTransformer2DModel
|
|
from .z_image_transformer2d_control import ZImageControlTransformer2DModel
|
|
|
|
# The pai_fuser is an internally developed acceleration package, which can be used on PAI.
|
|
if importlib.util.find_spec("paifuser") is not None:
|
|
# --------------------------------------------------------------- #
|
|
# The simple_wrapper is used to solve the problem
|
|
# about conflicts between cython and torch.compile
|
|
# --------------------------------------------------------------- #
|
|
def simple_wrapper(func):
|
|
def inner(*args, **kwargs):
|
|
return func(*args, **kwargs)
|
|
return inner
|
|
|
|
# --------------------------------------------------------------- #
|
|
# VAE Parallel Kernel
|
|
# --------------------------------------------------------------- #
|
|
from ..dist import parallel_magvit_vae
|
|
AutoencoderKLWan_.decode = simple_wrapper(parallel_magvit_vae(0.4, 8)(AutoencoderKLWan_.decode))
|
|
AutoencoderKLWan2_2_.decode = simple_wrapper(parallel_magvit_vae(0.4, 16)(AutoencoderKLWan2_2_.decode))
|
|
|
|
# --------------------------------------------------------------- #
|
|
# Sparse Attention
|
|
# --------------------------------------------------------------- #
|
|
import torch
|
|
from paifuser.ops import wan_sparse_attention_wrapper
|
|
|
|
WanSelfAttention.forward = simple_wrapper(wan_sparse_attention_wrapper()(WanSelfAttention.forward))
|
|
print("Import Sparse Attention")
|
|
|
|
WanTransformer3DModel.forward = simple_wrapper(WanTransformer3DModel.forward)
|
|
|
|
# --------------------------------------------------------------- #
|
|
# CFG Skip Turbo
|
|
# --------------------------------------------------------------- #
|
|
import os
|
|
|
|
if importlib.util.find_spec("paifuser.accelerator") is not None:
|
|
from paifuser.accelerator import (cfg_skip_turbo, disable_cfg_skip,
|
|
enable_cfg_skip, share_cfg_skip)
|
|
else:
|
|
from paifuser import (cfg_skip_turbo, disable_cfg_skip,
|
|
enable_cfg_skip, share_cfg_skip)
|
|
|
|
WanTransformer3DModel.enable_cfg_skip = enable_cfg_skip()(WanTransformer3DModel.enable_cfg_skip)
|
|
WanTransformer3DModel.disable_cfg_skip = disable_cfg_skip()(WanTransformer3DModel.disable_cfg_skip)
|
|
WanTransformer3DModel.share_cfg_skip = share_cfg_skip()(WanTransformer3DModel.share_cfg_skip)
|
|
|
|
QwenImageTransformer2DModel.enable_cfg_skip = enable_cfg_skip()(QwenImageTransformer2DModel.enable_cfg_skip)
|
|
QwenImageTransformer2DModel.disable_cfg_skip = disable_cfg_skip()(QwenImageTransformer2DModel.disable_cfg_skip)
|
|
print("Import CFG Skip Turbo")
|
|
|
|
# --------------------------------------------------------------- #
|
|
# RMS Norm Kernel
|
|
# --------------------------------------------------------------- #
|
|
from paifuser.ops import rms_norm_forward
|
|
WanRMSNorm.forward = rms_norm_forward
|
|
print("Import PAI RMS Fuse")
|
|
|
|
# --------------------------------------------------------------- #
|
|
# Fast Rope Kernel
|
|
# --------------------------------------------------------------- #
|
|
import types
|
|
|
|
import torch
|
|
from paifuser.ops import (ENABLE_KERNEL, fast_rope_apply_qk,
|
|
rope_apply_real_qk)
|
|
|
|
from . import wan_transformer3d
|
|
|
|
def deepcopy_function(f):
|
|
return types.FunctionType(f.__code__, f.__globals__, name=f.__name__, argdefs=f.__defaults__,closure=f.__closure__)
|
|
|
|
local_rope_apply_qk = deepcopy_function(wan_transformer3d.rope_apply_qk)
|
|
|
|
if ENABLE_KERNEL:
|
|
def adaptive_fast_rope_apply_qk(q, k, grid_sizes, freqs):
|
|
if torch.is_grad_enabled():
|
|
return local_rope_apply_qk(q, k, grid_sizes, freqs)
|
|
else:
|
|
return fast_rope_apply_qk(q, k, grid_sizes, freqs)
|
|
else:
|
|
def adaptive_fast_rope_apply_qk(q, k, grid_sizes, freqs):
|
|
return rope_apply_real_qk(q, k, grid_sizes, freqs)
|
|
|
|
wan_transformer3d.rope_apply_qk = adaptive_fast_rope_apply_qk
|
|
rope_apply_qk = adaptive_fast_rope_apply_qk
|
|
print("Import PAI Fast rope") |