Refactor/parallel (#157)

* compile repeated blocks

* pin memory

* refactor parallelism & model init

* fix

* fix cycle import

* fix qwen image

* fix examples

* check wan2.2 image shape

* fix lora

Author: akaitsuki-ii

diffsynth_engine/__init__.py

@@ -4,24 +4,33 @@
     FluxPipelineConfig,
     WanPipelineConfig,
     QwenImagePipelineConfig,
+    HunyuanPipelineConfig,
     SDStateDicts,
     SDXLStateDicts,
     FluxStateDicts,
+    WanStateDicts,
     QwenImageStateDicts,
     ControlNetParams,
     ControlType,
 )
 from .pipelines import (
-    FluxImagePipeline,
-    SDXLImagePipeline,
     SDImagePipeline,
+    SDXLImagePipeline,
+    FluxImagePipeline,
     WanVideoPipeline,
     QwenImagePipeline,
     Hunyuan3DShapePipeline,
 )
 from .models.flux import FluxControlNet, FluxIPAdapter, FluxRedux
 from .models.sd import SDControlNet
 from .models.sdxl import SDXLControlNetUnion
+from .tools import (
+    FluxInpaintingTool,
+    FluxOutpaintingTool,
+    FluxIPAdapterRefTool,
+    FluxReduxRefTool,
+    FluxReplaceByControlTool,
+)
 from .utils.download import (
     fetch_model,
     fetch_modelscope_model,
@@ -30,32 +39,29 @@
     reset_fetch_modelscope_model,
 )
 from .utils.video import load_video, save_video
-from .tools import (
-    FluxInpaintingTool,
-    FluxOutpaintingTool,
-    FluxIPAdapterRefTool,
-    FluxReduxRefTool,
-    FluxReplaceByControlTool,
-)
 
 __all__ = [
     "SDPipelineConfig",
     "SDXLPipelineConfig",
     "FluxPipelineConfig",
     "WanPipelineConfig",
+    "QwenImagePipelineConfig",
+    "HunyuanPipelineConfig",
     "SDStateDicts",
     "SDXLStateDicts",
     "FluxStateDicts",
+    "WanStateDicts",
     "QwenImageStateDicts",
+    "ControlNetParams",
+    "ControlType",
+    "SDImagePipeline",
+    "SDControlNet",
+    "SDXLImagePipeline",
+    "SDXLControlNetUnion",
     "FluxImagePipeline",
-    "QwenImagePipelineConfig",
     "FluxControlNet",
     "FluxIPAdapter",
     "FluxRedux",
-    "SDControlNet",
-    "SDXLControlNetUnion",
-    "SDXLImagePipeline",
-    "SDImagePipeline",
     "WanVideoPipeline",
     "QwenImagePipeline",
     "Hunyuan3DShapePipeline",
@@ -64,8 +70,6 @@
     "FluxIPAdapterRefTool",
     "FluxReplaceByControlTool",
     "FluxReduxRefTool",
-    "ControlNetParams",
-    "ControlType",
     "fetch_model",
     "fetch_modelscope_model",
     "register_fetch_modelscope_model",

diffsynth_engine/configs/__init__.py

@@ -8,10 +8,12 @@
     FluxPipelineConfig,
     WanPipelineConfig,
     QwenImagePipelineConfig,
+    HunyuanPipelineConfig,
     BaseStateDicts,
     SDStateDicts,
     SDXLStateDicts,
     FluxStateDicts,
+    WanStateDicts,
     QwenImageStateDicts,
 )
 from .controlnet import ControlType, ControlNetParams
@@ -26,11 +28,13 @@
     "FluxPipelineConfig",
     "WanPipelineConfig",
     "QwenImagePipelineConfig",
-    "ControlType",
-    "ControlNetParams",
+    "HunyuanPipelineConfig",
     "BaseStateDicts",
     "SDStateDicts",
     "SDXLStateDicts",
     "FluxStateDicts",
+    "WanStateDicts",
     "QwenImageStateDicts",
+    "ControlType",
+    "ControlNetParams",
 ]

diffsynth_engine/configs/pipeline.py

@@ -1,7 +1,7 @@
 import os
 import torch
 from dataclasses import dataclass, field
-from typing import List, Tuple, Optional, Dict
+from typing import List, Dict, Tuple, Optional
 
 from diffsynth_engine.configs.controlnet import ControlType
 
@@ -127,7 +127,7 @@ def basic_config(
             model_path=model_path,
             device=device,
             parallelism=parallelism,
-            use_fsdp=True,
+            use_fsdp=True if parallelism > 1 else False,
             offload_mode=offload_mode,
             offload_to_disk=offload_to_disk,
         )
@@ -174,8 +174,8 @@ def basic_config(
             image_encoder_path=image_encoder_path,
             device=device,
             parallelism=parallelism,
-            use_cfg_parallel=True,
-            use_fsdp=True,
+            use_cfg_parallel=True if parallelism > 1 else False,
+            use_fsdp=True if parallelism > 1 else False,
             offload_mode=offload_mode,
             offload_to_disk=offload_to_disk,
         )
@@ -184,16 +184,6 @@ def __post_init__(self):
         init_parallel_config(self)
 
 
-@dataclass
-class HunyuanPipelineConfig(BaseConfig):
-    model_path: str | os.PathLike | List[str | os.PathLike]
-    model_dtype: torch.dtype = torch.float16
-    vae_path: Optional[str | os.PathLike | List[str | os.PathLike]] = None
-    vae_dtype: torch.dtype = torch.float16
-    image_encoder_path: Optional[str | os.PathLike | List[str | os.PathLike]] = None
-    image_encoder_dtype: torch.dtype = torch.float16
-
-
 @dataclass
 class QwenImagePipelineConfig(AttentionConfig, OptimizationConfig, ParallelConfig, BaseConfig):
     model_path: str | os.PathLike | List[str | os.PathLike]
@@ -228,8 +218,8 @@ def basic_config(
             encoder_path=encoder_path,
             vae_path=vae_path,
             parallelism=parallelism,
-            use_cfg_parallel=True,
-            use_fsdp=True,
+            use_cfg_parallel=True if parallelism > 1 else False,
+            use_fsdp=True if parallelism > 1 else False,
             offload_mode=offload_mode,
             offload_to_disk=offload_to_disk,
         )
@@ -238,32 +228,57 @@ def __post_init__(self):
         init_parallel_config(self)
 
 
+@dataclass
+class HunyuanPipelineConfig(BaseConfig):
+    model_path: str | os.PathLike | List[str | os.PathLike]
+    model_dtype: torch.dtype = torch.float16
+    vae_path: Optional[str | os.PathLike | List[str | os.PathLike]] = None
+    vae_dtype: torch.dtype = torch.float16
+    image_encoder_path: Optional[str | os.PathLike | List[str | os.PathLike]] = None
+    image_encoder_dtype: torch.dtype = torch.float16
+
+
 @dataclass
 class BaseStateDicts:
-    model: Optional[Dict[str, torch.Tensor]] = None
-    vae: Optional[Dict[str, torch.Tensor]] = None
+    pass
+
+
+@dataclass
+class SDStateDicts:
+    model: Dict[str, torch.Tensor]
+    clip: Dict[str, torch.Tensor]
+    vae: Dict[str, torch.Tensor]
 
 
 @dataclass
-class SDStateDicts(BaseStateDicts):
-    clip: Optional[Dict[str, torch.Tensor]] = None
+class SDXLStateDicts:
+    model: Dict[str, torch.Tensor]
+    clip_l: Dict[str, torch.Tensor]
+    clip_g: Dict[str, torch.Tensor]
+    vae: Dict[str, torch.Tensor]
 
 
 @dataclass
-class SDXLStateDicts(BaseStateDicts):
-    clip_l: Optional[Dict[str, torch.Tensor]] = None
-    clip_g: Optional[Dict[str, torch.Tensor]] = None
+class FluxStateDicts:
+    model: Dict[str, torch.Tensor]
+    t5: Dict[str, torch.Tensor]
+    clip: Dict[str, torch.Tensor]
+    vae: Dict[str, torch.Tensor]
 
 
 @dataclass
-class FluxStateDicts(BaseStateDicts):
-    t5: Optional[Dict[str, torch.Tensor]] = None
-    clip: Optional[Dict[str, torch.Tensor]] = None
+class WanStateDicts:
+    model: Dict[str, torch.Tensor] | Dict[str, Dict[str, torch.Tensor]]
+    t5: Dict[str, torch.Tensor]
+    vae: Dict[str, torch.Tensor]
+    image_encoder: Optional[Dict[str, torch.Tensor]] = None
 
 
 @dataclass
-class QwenImageStateDicts(BaseStateDicts):
-    encoder: Optional[Dict[str, torch.Tensor]] = None
+class QwenImageStateDicts:
+    model: Dict[str, torch.Tensor]
+    encoder: Dict[str, torch.Tensor]
+    vae: Dict[str, torch.Tensor]
 
 
 def init_parallel_config(config: FluxPipelineConfig | QwenImagePipelineConfig | WanPipelineConfig):

diffsynth_engine/models/base.py

@@ -4,7 +4,6 @@
 from typing import Dict, Union, List, Any
 from diffsynth_engine.utils.loader import load_file
 from diffsynth_engine.models.basic.lora import LoRALinear, LoRAConv2d
-from diffsynth_engine.models.utils import no_init_weights
 
 
 class StateDictConverter:
@@ -33,10 +32,9 @@ def from_pretrained(
 
     @classmethod
     def from_state_dict(cls, state_dict: Dict[str, torch.Tensor], device: str, dtype: torch.dtype, **kwargs):
-        with no_init_weights():
-            model = torch.nn.utils.skip_init(cls, device=device, dtype=dtype, **kwargs)
-        model.to_empty(device=device)
-        model.load_state_dict(state_dict)
+        model = cls(device="meta", dtype=dtype, **kwargs)
+        model.requires_grad_(False)
+        model.load_state_dict(state_dict, assign=True)
         model.to(device=device, dtype=dtype, non_blocking=True)
         return model
 

diffsynth_engine/models/basic/lora.py

@@ -74,14 +74,13 @@ def __init__(
 
     @staticmethod
     def from_linear(linear: nn.Linear):
-        lora_linear = torch.nn.utils.skip_init(
-            LoRALinear,
+        lora_linear = LoRALinear(
             linear.in_features,
             linear.out_features,
             linear.bias is not None,
-            device=linear.weight.device,
+            device="meta",
             dtype=linear.weight.dtype,
-        )
+        ).to_empty(device=linear.weight.device)
         lora_linear.weight = linear.weight
         lora_linear.bias = linear.bias
         return lora_linear
@@ -98,12 +97,20 @@ def add_lora(
         dtype: torch.dtype,
         **kwargs,
     ):
-        up_linear = torch.nn.utils.skip_init(
-            nn.Linear, up.shape[1], up.shape[0], bias=False, device=device, dtype=dtype
-        )
-        down_linear = torch.nn.utils.skip_init(
-            nn.Linear, down.shape[0], down.shape[1], bias=False, device=device, dtype=dtype
-        )
+        up_linear = nn.Linear(
+            up.shape[1],
+            up.shape[0],
+            bias=False,
+            device="meta",
+            dtype=dtype,
+        ).to_empty(device=device)
+        down_linear = nn.Linear(
+            down.shape[0],
+            down.shape[1],
+            bias=False,
+            device="meta",
+            dtype=dtype,
+        ).to_empty(device=device)
         up_linear.weight.data = up
         down_linear.weight.data = down
         lora = LoRA(scale, rank, alpha, up_linear, down_linear, device, dtype)
@@ -182,8 +189,7 @@ def __init__(
 
     @staticmethod
     def from_conv2d(conv2d: nn.Conv2d):
-        lora_conv2d = torch.nn.utils.skip_init(
-            LoRAConv2d,
+        lora_conv2d = LoRAConv2d(
             conv2d.in_channels,
             conv2d.out_channels,
             conv2d.kernel_size,
@@ -193,9 +199,9 @@ def from_conv2d(conv2d: nn.Conv2d):
             conv2d.groups,
             conv2d.bias is not None,
             conv2d.padding_mode,
-            device=conv2d.weight.device,
+            device="meta",
             dtype=conv2d.weight.dtype,
-        )
+        ).to_empty(device=conv2d.weight.device)
         lora_conv2d.weight = conv2d.weight
         lora_conv2d.bias = conv2d.bias
         return lora_conv2d
@@ -211,31 +217,29 @@ def _construct_lora(
         device: str,
         dtype: torch.dtype,
     ):
-        down_conv = torch.nn.utils.skip_init(
-            nn.Conv2d,
+        down_conv = nn.Conv2d(
             self.in_channels,
             rank,
             kernel_size=self.kernel_size,
             stride=self.stride,
             padding=self.padding,
             bias=False,
-            device=device,
+            device="meta",
             dtype=dtype,
-        )
+        ).to_empty(device=device)
         down_conv.weight.data = down
         # according to the official kohya_ss trainer kernel_size are always fixed for the up layer
         # see: https://github.com/bmaltais/kohya_ss/blob/2accb1305979ba62f5077a23aabac23b4c37e935/networks/lora_diffusers.py#L129
         # refer from diffusers
-        up_conv = torch.nn.utils.skip_init(
-            nn.Conv2d,
+        up_conv = nn.Conv2d(
             rank,
             self.out_channels,
             kernel_size=(1, 1),
             stride=(1, 1),
             bias=False,
-            device=device,
+            device="meta",
             dtype=dtype,
-        )
+        ).to_empty(device=device)
         up_conv.weight.data = up
 
         lora = LoRA(scale, rank, alpha, up_conv, down_conv, device, dtype)

diffsynth_engine/models/flux/flux_controlnet.py

@@ -8,7 +8,6 @@
     RoPEEmbedding,
     TimestepEmbeddings,
 )
-from diffsynth_engine.models.utils import no_init_weights
 
 
 class FluxControlNetStateDictConverter(StateDictConverter):
@@ -164,10 +163,13 @@ def from_state_dict(
         else:
             condition_channels = 64
 
-        with no_init_weights():
-            model = torch.nn.utils.skip_init(
-                cls, condition_channels=condition_channels, attn_kwargs=attn_kwargs, device=device, dtype=dtype
-            )
-        model.load_state_dict(state_dict)
+        model = cls(
+            condition_channels=condition_channels,
+            attn_kwargs=attn_kwargs,
+            device="meta",
+            dtype=dtype,
+        )
+        model.requires_grad_(False)
+        model.load_state_dict(state_dict, assign=True)
         model.to(device=device, dtype=dtype, non_blocking=True)
         return model