fix cfg parallel (#104)

Author: akaitsuki-ii

diffsynth_engine/models/flux/flux_dit.py

@@ -411,6 +411,7 @@ def forward(
         )
 
         fp8_linear_enabled = getattr(self, "fp8_linear_enabled", False)
+        use_cfg = hidden_states.shape[0] > 1
         with (
             fp8_inference(fp8_linear_enabled),
             gguf_inference(),
@@ -426,7 +427,8 @@ def forward(
                     image_ids,
                     *controlnet_double_block_output,
                     *controlnet_single_block_output,
-                )
+                ),
+                use_cfg=use_cfg,
             ),
         ):
             # warning: keep the order of time_embedding + guidance_embedding + pooled_text_embedding
@@ -483,7 +485,7 @@ def forward(
                 (hidden_states,) = sequence_parallel_unshard((hidden_states,), seq_dims=(1,), seq_lens=(h * w // 4,))
 
             hidden_states = self.unpatchify(hidden_states, h, w)
-            (hidden_states,) = cfg_parallel_unshard((hidden_states,))
+            (hidden_states,) = cfg_parallel_unshard((hidden_states,), use_cfg=use_cfg)
             return hidden_states
 
     @classmethod

diffsynth_engine/utils/parallel.py

@@ -330,9 +330,9 @@ def __init__(
         device: str = "cuda",
     ):
         current_method = mp.get_start_method(allow_none=True)
-        if current_method is None or current_method != 'spawn':
+        if current_method is None or current_method != "spawn":
             try:
-                mp.set_start_method('spawn')
+                mp.set_start_method("spawn")
             except RuntimeError as e:
                 raise RuntimeError("Failed to set start method to spawn:", e)
         super().__init__()
@@ -404,8 +404,8 @@ def __del__(self):
 
 
 @contextmanager
-def cfg_parallel(tensors: List[torch.Tensor]):
-    if get_cfg_world_size() == 1:
+def cfg_parallel(tensors: List[torch.Tensor], use_cfg=True):
+    if get_cfg_world_size() == 1 or not use_cfg:
         yield
         return
 
@@ -426,8 +426,8 @@ def cfg_parallel(tensors: List[torch.Tensor]):
         tensor.copy_(original_tensor)
 
 
-def cfg_parallel_unshard(tensors: List[torch.Tensor]):
-    if get_cfg_world_size() == 1:
+def cfg_parallel_unshard(tensors: List[torch.Tensor], use_cfg=True):
+    if get_cfg_world_size() == 1 or not use_cfg:
         return tensors
 
     unshard_tensors = []