Support Torch Compile (#1368)

* support simple compile

* add support for compile

* minor fix

* minor fix

* minor fix
This commit is contained in:
Hong Zhang
2026-03-24 11:19:43 +08:00
committed by GitHub
parent e2a3a987da
commit ae8cb139e8
18 changed files with 59 additions and 2 deletions

View File

@@ -39,6 +39,7 @@ class AnimaImagePipeline(BasePipeline):
AnimaUnit_PromptEmbedder(),
]
self.model_fn = model_fn_anima
self.compilable_models = ["dit"]
@staticmethod

View File

@@ -42,6 +42,7 @@ class Flux2ImagePipeline(BasePipeline):
Flux2Unit_ImageIDs(),
]
self.model_fn = model_fn_flux2
self.compilable_models = ["dit"]
@staticmethod

View File

@@ -103,6 +103,7 @@ class FluxImagePipeline(BasePipeline):
FluxImageUnit_LoRAEncode(),
]
self.model_fn = model_fn_flux_image
self.compilable_models = ["dit"]
self.lora_loader = FluxLoRALoader
def enable_lora_merger(self):

View File

@@ -76,6 +76,7 @@ class LTX2AudioVideoPipeline(BasePipeline):
LTX2AudioVideoUnit_SetScheduleStage2(),
]
self.model_fn = model_fn_ltx2
self.compilable_models = ["dit"]
self.default_negative_prompt = {
"LTX-2": (

View File

@@ -52,6 +52,7 @@ class MovaAudioVideoPipeline(BasePipeline):
MovaAudioVideoUnit_UnifiedSequenceParallel(),
]
self.model_fn = model_fn_mova_audio_video
self.compilable_models = ["video_dit", "video_dit2", "audio_dit"]
def enable_usp(self):
from ..utils.xfuser import get_sequence_parallel_world_size, usp_attn_forward

View File

@@ -56,6 +56,7 @@ class QwenImagePipeline(BasePipeline):
QwenImageUnit_BlockwiseControlNet(),
]
self.model_fn = model_fn_qwen_image
self.compilable_models = ["dit"]
@staticmethod

View File

@@ -83,6 +83,7 @@ class WanVideoPipeline(BasePipeline):
WanVideoPostUnit_S2V(),
]
self.model_fn = model_fn_wan_video
self.compilable_models = ["dit", "dit2"]
def enable_usp(self):

View File

@@ -54,6 +54,7 @@ class ZImagePipeline(BasePipeline):
ZImageUnit_PAIControlNet(),
]
self.model_fn = model_fn_z_image
self.compilable_models = ["dit"]
@staticmethod