diffsynth-engine 0.4.2.dev4__py3-none-any.whl → 0.4.2.dev6__py3-none-any.whl

diffsynth_engine/models/qwen_image/qwen_image_dit.py

@@ -315,6 +315,7 @@ class QwenImageTransformerBlock(nn.Module):
 
 class QwenImageDiT(PreTrainedModel):
     converter = QwenImageDiTStateDictConverter()
+    _supports_parallelization = True
 
     def __init__(
         self,
@@ -423,3 +424,6 @@ class QwenImageDiT(PreTrainedModel):
         model.load_state_dict(state_dict, assign=True)
         model.to(device=device, dtype=dtype, non_blocking=True)
         return model
+
+    def get_fsdp_modules(self):
+        return ["transformer_blocks"]

diffsynth_engine/pipelines/base.py

@@ -6,7 +6,7 @@ from typing import Dict, List, Tuple
 from PIL import Image
 
 from diffsynth_engine.configs import BaseConfig, BaseStateDicts
-from diffsynth_engine.utils.offload import enable_sequential_cpu_offload
+from diffsynth_engine.utils.offload import enable_sequential_cpu_offload, offload_model_to_dict, restore_model_from_dict
 from diffsynth_engine.utils.fp8_linear import enable_fp8_autocast
 from diffsynth_engine.utils.gguf import load_gguf_checkpoint
 from diffsynth_engine.utils import logging
@@ -40,6 +40,7 @@ class BasePipeline:
         self.dtype = dtype
         self.offload_mode = None
         self.model_names = []
+        self._offload_param_dict = {}
 
     @classmethod
     def from_pretrained(cls, model_path_or_config: str | BaseConfig) -> "BasePipeline":
@@ -243,14 +244,13 @@ class BasePipeline:
         for model_name in self.model_names:
             model = getattr(self, model_name)
             if model is not None:
-                model.to("cpu")
+                self._offload_param_dict[model_name] = offload_model_to_dict(model)
         self.offload_mode = "cpu_offload"
 
     def _enable_sequential_cpu_offload(self):
         for model_name in self.model_names:
             model = getattr(self, model_name)
             if model is not None:
-                model.to("cpu")
                 enable_sequential_cpu_offload(model, self.device)
         self.offload_mode = "sequential_cpu_offload"
 
@@ -277,20 +277,12 @@ class BasePipeline:
         for model_name in self.model_names:
             if model_name not in load_model_names:
                 model = getattr(self, model_name)
-                if (
-                    model is not None
-                    and (p := next(model.parameters(), None)) is not None
-                    and p.device != torch.device("cpu")
-                ):
-                    model.to("cpu")
+                if model is not None and (p := next(model.parameters(), None)) is not None and p.device.type != "cpu":
+                    restore_model_from_dict(model, self._offload_param_dict[model_name])
         # load the needed models to device
         for model_name in load_model_names:
             model = getattr(self, model_name)
-            if (
-                model is not None
-                and (p := next(model.parameters(), None)) is not None
-                and p.device != torch.device(self.device)
-            ):
+            if model is not None and (p := next(model.parameters(), None)) is not None and p.device.type != self.device:
                 model.to(self.device)
         # fresh the cuda cache
         empty_cache()

diffsynth_engine/pipelines/wan_video.py

@@ -584,4 +584,11 @@ class WanVideoPipeline(BasePipeline):
                 use_fsdp=config.use_fsdp,
                 device="cuda",
             )
+        if config.use_torch_compile:
+            pipe.compile()
         return pipe
+
+    def compile(self):
+        self.dit.compile()
+        if self.dit2 is not None:
+            self.dit2.compile()

diffsynth_engine/utils/offload.py

@@ -1,8 +1,10 @@
 import torch
 import torch.nn as nn
+from typing import Dict
 
 
 def enable_sequential_cpu_offload(module: nn.Module, device: str = "cuda"):
+    module = module.to("cpu")
     if len(list(module.children())) == 0:
         if len(list(module.parameters())) > 0 or len(list(module.buffers())) > 0:
             # leaf module with parameters or buffers
@@ -50,3 +52,24 @@ def add_cpu_offload_hook(module: nn.Module, device: str = "cuda", recurse: bool
     module.register_forward_pre_hook(_forward_pre_hook)
     module.register_forward_hook(_forward_hook)
     setattr(module, "_cpu_offload_enabled", True)
+
+
+def offload_model_to_dict(module: nn.Module) -> Dict[str, torch.Tensor]:
+    module = module.to("cpu")
+    offload_param_dict = {}
+    for name, param in module.named_parameters(recurse=True):
+        param.data = param.data.pin_memory()
+        offload_param_dict[name] = param.data
+    for name, buffer in module.named_buffers(recurse=True):
+        buffer.data = buffer.data.pin_memory()
+        offload_param_dict[name] = buffer.data
+    return offload_param_dict
+
+
+def restore_model_from_dict(module: nn.Module, offload_param_dict: Dict[str, torch.Tensor]):
+    for name, param in module.named_parameters(recurse=True):
+        if name in offload_param_dict:
+            param.data = offload_param_dict[name]
+    for name, buffer in module.named_buffers(recurse=True):
+        if name in offload_param_dict:
+            buffer.data = offload_param_dict[name]

{diffsynth_engine-0.4.2.dev4.dist-info → diffsynth_engine-0.4.2.dev6.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: diffsynth_engine
-Version: 0.4.2.dev4
+Version: 0.4.2.dev6
 Author: MuseAI x ModelScope
 Classifier: Programming Language :: Python :: 3
 Classifier: Operating System :: OS Independent

{diffsynth_engine-0.4.2.dev4.dist-info → diffsynth_engine-0.4.2.dev6.dist-info}/RECORD

@@ -100,7 +100,7 @@ diffsynth_engine/models/flux/flux_text_encoder.py,sha256=Qcs277RIPP-O5AkcAb5Fb0j
 diffsynth_engine/models/flux/flux_vae.py,sha256=YoeZZSTbXo5rpYQ5-XY3A-Bd1swb8H1yOmtVjm0g5ZI,2994
 diffsynth_engine/models/qwen_image/__init__.py,sha256=X5pig621WEsDZ6L7HVkmYspV53-GDfs_la1ncaq_NFw,417
 diffsynth_engine/models/qwen_image/qwen2_5_vl.py,sha256=vpBo6eo_96iLky9YV5MX0nbmOleY2EX97TrJoRBNnw4,56511
-diffsynth_engine/models/qwen_image/qwen_image_dit.py,sha256=CzYtYrqBzaHcn6OE4Z4yYqYdAoGRbfQtZ5KQVrN03to,16839
+diffsynth_engine/models/qwen_image/qwen_image_dit.py,sha256=bmsVuwMU26i34lCDDpKsxY1OVRwtWawn44r_E_Pd1Ws,16947
 diffsynth_engine/models/qwen_image/qwen_image_dit_fbcache.py,sha256=tgCmD4MFNgd3HmLyoYnt8HZCzPBTgQ4zCQjV5qZSW_I,4870
 diffsynth_engine/models/qwen_image/qwen_image_vae.py,sha256=m455iJfJx8KVEsrkinjClokkEqd0RshSZoKZ_QAdRyk,38509
 diffsynth_engine/models/sd/__init__.py,sha256=hjoKRnwoXOLD0wude-w7I6wK5ak7ACMbnbkPuBB2oU0,380
@@ -129,13 +129,13 @@ diffsynth_engine/models/wan/wan_image_encoder.py,sha256=LYwcfCcQmXf9FP08DGaU2bfa
 diffsynth_engine/models/wan/wan_text_encoder.py,sha256=bkphxtqNNwXcEA_OaUrwV9CvICV-s16awu5Z9gjjzsM,10912
 diffsynth_engine/models/wan/wan_vae.py,sha256=AmBuqyPwZCFY0e8lUThlJoNHmpmTm2_dE1XYzXBCaAI,38937
 diffsynth_engine/pipelines/__init__.py,sha256=9QZVhZeRm_5m7yxie08yBtgM26NB4mfVBOjj1Prlv-k,447
-diffsynth_engine/pipelines/base.py,sha256=AUoB07rCvrfVyied-qQBIsexyf5eE8hqyXr-2IIWYwU,12339
+diffsynth_engine/pipelines/base.py,sha256=730FzAn1cxZ2SM2qr9SXk7qNLpAau9PKEqwfMSW_8Aw,12336
 diffsynth_engine/pipelines/flux_image.py,sha256=gWuZaMeupB_Wz3AY97eE1eEVSAmAm14aXIxkAqNXY7E,49224
 diffsynth_engine/pipelines/qwen_image.py,sha256=S15FAxmiWcF3qItBigXR9jVJa4h9Qqs4_8cT3KXi8Ec,17506
 diffsynth_engine/pipelines/sd_image.py,sha256=5cIIknh2M-fOqj7urKi9nZ40yc1LnvepbH_Af7SF4UA,17789
 diffsynth_engine/pipelines/sdxl_image.py,sha256=otv1T_0fhX3UcIoKbKCqb47Yge6xg0fPM0ry-uPEanI,21548
 diffsynth_engine/pipelines/utils.py,sha256=lk7sFGEk-fGjgadLpwwppHKG-yZ0RC-4ZmHW7pRRe8A,473
-diffsynth_engine/pipelines/wan_video.py,sha256=wbCHPDgs4BmyX1DsawaXqxeCoVAcISNcXNnFr2qcTx0,25424
+diffsynth_engine/pipelines/wan_video.py,sha256=stoYKm0wHf_pxZ_WHRTGHTR61KVG_U21yBUaUrDjSqw,25605
 diffsynth_engine/processor/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 diffsynth_engine/processor/canny_processor.py,sha256=hV30NlblTkEFUAmF_O-LJrNlGVM2SFrqq6okfF8VpOo,602
 diffsynth_engine/processor/depth_processor.py,sha256=dQvs3JsnyMbz4dyI9QoR8oO-mMFBFAgNvgqeCoaU5jk,1532
@@ -162,14 +162,14 @@ diffsynth_engine/utils/image.py,sha256=_46CVs1Qe7GdZNulWWJISnR_Y6FotC2tZGLKtr04g
 diffsynth_engine/utils/loader.py,sha256=Z5v1WNDWFY0OrVubB70j5VU3zeaAfEK_j8c1KrGI4yM,1240
 diffsynth_engine/utils/lock.py,sha256=1Ipgst9eEFfFdViAvD5bxdB6HnHHBcqWYOb__fGaPUI,1601
 diffsynth_engine/utils/logging.py,sha256=XB0xTT8PBN6btkOjFtOvjlrOCRVgDGT8PFAp1vmse28,467
-diffsynth_engine/utils/offload.py,sha256=CYIIDr9CDGE3YN1kCOUnd1BBGvYxjfBCR3BMzqCx4RQ,2580
+diffsynth_engine/utils/offload.py,sha256=qEiqbeMQqeV1DtHF-6OuUO8Akdr1enfpPYwF0OBpx98,3500
 diffsynth_engine/utils/onnx.py,sha256=jeWUudJHnESjuiEAHyUZYUZz7dCj34O9aGjHCe8yjWo,1149
 diffsynth_engine/utils/parallel.py,sha256=TtB6FzP2qo4VQqSenVnV-ZaKmp9xHaIWJ8D1ZjHtukE,17064
 diffsynth_engine/utils/platform.py,sha256=2lXdw6YkqcRONCeT98n4cyg1Ii8Ybbyj2Ns72Se9tlk,496
 diffsynth_engine/utils/prompt.py,sha256=YItMchoVzsG6y-LB4vzzDUWrkhKRVlt1HfVhxZjSxMQ,280
 diffsynth_engine/utils/video.py,sha256=Ne0rd2lb59UT1q5EotpjlY7OT8F9oTCFDyo1ST77uoQ,1004
-diffsynth_engine-0.4.2.dev4.dist-info/licenses/LICENSE,sha256=x7aBqQuVI0IYnftgoTPI_A0I_rjdjPPQkjnU6N2nikM,11346
-diffsynth_engine-0.4.2.dev4.dist-info/METADATA,sha256=-RN0vnmcItmjNVo9q-OV8oMAyfLk8MMvy03DUdZ516Q,1110
-diffsynth_engine-0.4.2.dev4.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-diffsynth_engine-0.4.2.dev4.dist-info/top_level.txt,sha256=6zgbiIzEHLbhgDKRyX0uBJOV3F6VnGGBRIQvSiYYn6w,17
-diffsynth_engine-0.4.2.dev4.dist-info/RECORD,,
+diffsynth_engine-0.4.2.dev6.dist-info/licenses/LICENSE,sha256=x7aBqQuVI0IYnftgoTPI_A0I_rjdjPPQkjnU6N2nikM,11346
+diffsynth_engine-0.4.2.dev6.dist-info/METADATA,sha256=xy1w-I4tAN0UahyFnr4ZP9EoMzm8C1uFGSrC4Y1P0m4,1110
+diffsynth_engine-0.4.2.dev6.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+diffsynth_engine-0.4.2.dev6.dist-info/top_level.txt,sha256=6zgbiIzEHLbhgDKRyX0uBJOV3F6VnGGBRIQvSiYYn6w,17
+diffsynth_engine-0.4.2.dev6.dist-info/RECORD,,