diffsynth-engine 0.3.6.dev4__tar.gz → 0.3.6.dev5__tar.gz

{diffsynth_engine-0.3.6.dev4 → diffsynth_engine-0.3.6.dev5}/PKG-INFO

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

{diffsynth_engine-0.3.6.dev4 → diffsynth_engine-0.3.6.dev5}/diffsynth_engine/models/wan/wan_dit.py

@@ -334,9 +334,10 @@ class WanDiT(PreTrainedModel):
         clip_feature: Optional[torch.Tensor] = None,  # clip_vision_encoder(img)
         y: Optional[torch.Tensor] = None,  # vae_encoder(img)
     ):
+        use_cfg = x.shape[0] > 1
         with (
             gguf_inference(),
-            cfg_parallel((x, context, timestep, clip_feature, y)),
+            cfg_parallel((x, context, timestep, clip_feature, y), use_cfg=use_cfg),
         ):
             t = self.time_embedding(sinusoidal_embedding_1d(self.freq_dim, timestep))
             t_mod = self.time_projection(t).unflatten(1, (6, self.dim))
@@ -365,7 +366,7 @@ class WanDiT(PreTrainedModel):
                 x = self.head(x, t)
                 (x,) = sequence_parallel_unshard((x,), seq_dims=(1,), seq_lens=(f * h * w,))
             x = self.unpatchify(x, (f, h, w))
-            (x,) = cfg_parallel_unshard((x,))
+            (x,) = cfg_parallel_unshard((x,), use_cfg=use_cfg)
             return x
 
     @classmethod

{diffsynth_engine-0.3.6.dev4 → diffsynth_engine-0.3.6.dev5}/diffsynth_engine/models/wan/wan_vae.py

@@ -515,7 +515,7 @@ class WanVideoVAEStateDictConverter(StateDictConverter):
 class WanVideoVAE(PreTrainedModel):
     converter = WanVideoVAEStateDictConverter()
 
-    def __init__(self, z_dim=16, parallelism: int = 1, device: str = "cuda:0", dtype: torch.dtype = torch.float32):
+    def __init__(self, z_dim=16, device: str = "cuda:0", dtype: torch.dtype = torch.float32):
         super().__init__()
 
         mean = [
@@ -561,12 +561,11 @@ class WanVideoVAE(PreTrainedModel):
         # init model
         self.model = VideoVAE(z_dim=z_dim).eval().requires_grad_(False)
         self.upsampling_factor = 8
-        self.parallelism = parallelism
 
     @classmethod
-    def from_state_dict(cls, state_dict, parallelism=1, device="cuda:0", dtype=torch.float32) -> "WanVideoVAE":
+    def from_state_dict(cls, state_dict, device="cuda:0", dtype=torch.float32) -> "WanVideoVAE":
         with no_init_weights():
-            model = torch.nn.utils.skip_init(cls, parallelism=parallelism, device=device, dtype=dtype)
+            model = torch.nn.utils.skip_init(cls, device=device, dtype=dtype)
         model.load_state_dict(state_dict, assign=True)
         model.to(device=device, dtype=dtype, non_blocking=True)
         return model
@@ -607,7 +606,7 @@ class WanVideoVAE(PreTrainedModel):
                 h_, w_ = h + size_h, w + size_w
                 tasks.append((h, h_, w, w_))
 
-        data_device = device if self.parallelism > 1 else "cpu"
+        data_device = device if dist.is_initialized() else "cpu"
         computation_device = device
 
         out_T = T * 4 - 3
@@ -622,9 +621,9 @@ class WanVideoVAE(PreTrainedModel):
             device=data_device,
         )
 
-        hide_progress_bar = self.parallelism > 1 and dist.get_rank() != 0
-        for i, (h, h_, w, w_) in enumerate(tqdm(tasks, desc="VAE DECODING", disable=hide_progress_bar)):
-            if self.parallelism > 1 and (i % dist.get_world_size() != dist.get_rank()):
+        hide_progress = dist.is_initialized() and dist.get_rank() != 0
+        for i, (h, h_, w, w_) in enumerate(tqdm(tasks, desc="VAE DECODING", disable=hide_progress)):
+            if dist.is_initialized() and (i % dist.get_world_size() != dist.get_rank()):
                 continue
             hidden_states_batch = hidden_states[:, :, :, h:h_, w:w_].to(computation_device)
             hidden_states_batch = self.model.decode(hidden_states_batch, self.scale).to(data_device)
@@ -654,11 +653,11 @@ class WanVideoVAE(PreTrainedModel):
                 target_h : target_h + hidden_states_batch.shape[3],
                 target_w : target_w + hidden_states_batch.shape[4],
             ] += mask
-            if progress_callback is not None and not hide_progress_bar:
+            if progress_callback is not None and not hide_progress:
                 progress_callback(i + 1, len(tasks), "VAE DECODING")
-        if progress_callback is not None and not hide_progress_bar:
+        if progress_callback is not None and not hide_progress:
             progress_callback(len(tasks), len(tasks), "VAE DECODING")
-        if self.parallelism > 1:
+        if dist.is_initialized():
             dist.all_reduce(values)
             dist.all_reduce(weight)
         values = values / weight
@@ -681,7 +680,7 @@ class WanVideoVAE(PreTrainedModel):
                 h_, w_ = h + size_h, w + size_w
                 tasks.append((h, h_, w, w_))
 
-        data_device = device if self.parallelism > 1 else "cpu"
+        data_device = device if dist.is_initialized() else "cpu"
         computation_device = device
 
         out_T = (T + 3) // 4
@@ -696,9 +695,9 @@ class WanVideoVAE(PreTrainedModel):
             device=data_device,
         )
 
-        hide_progress_bar = self.parallelism > 1 and dist.get_rank() != 0
+        hide_progress_bar = dist.is_initialized() and dist.get_rank() != 0
         for i, (h, h_, w, w_) in enumerate(tqdm(tasks, desc="VAE ENCODING", disable=hide_progress_bar)):
-            if self.parallelism > 1 and (i % dist.get_world_size() != dist.get_rank()):
+            if dist.is_initialized() and (i % dist.get_world_size() != dist.get_rank()):
                 continue
             hidden_states_batch = video[:, :, :, h:h_, w:w_].to(computation_device)
             hidden_states_batch = self.model.encode(hidden_states_batch, self.scale).to(data_device)
@@ -732,7 +731,7 @@ class WanVideoVAE(PreTrainedModel):
                 progress_callback(i + 1, len(tasks), "VAE ENCODING")
         if progress_callback is not None and not hide_progress_bar:
             progress_callback(len(tasks), len(tasks), "VAE ENCODING")
-        if self.parallelism > 1:
+        if dist.is_initialized():
             dist.all_reduce(values)
             dist.all_reduce(weight)
         values = values / weight

{diffsynth_engine-0.3.6.dev4 → diffsynth_engine-0.3.6.dev5}/diffsynth_engine.egg-info/PKG-INFO

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