diffsynth-engine 0.2.8__tar.gz → 0.2.9__tar.gz

{diffsynth_engine-0.2.8 → diffsynth_engine-0.2.9}/PKG-INFO

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

{diffsynth_engine-0.2.8 → diffsynth_engine-0.2.9}/diffsynth_engine/models/wan/wan_vae.py

@@ -1,6 +1,7 @@
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
+import torch.distributed as dist
 from einops import rearrange, repeat
 from tqdm import tqdm
 
@@ -514,7 +515,7 @@ class WanVideoVAEStateDictConverter(StateDictConverter):
 class WanVideoVAE(PreTrainedModel):
     converter = WanVideoVAEStateDictConverter()
 
-    def __init__(self, z_dim=16, device: str = "cuda:0", dtype: torch.dtype = torch.float32):
+    def __init__(self, z_dim=16, parallelism: int = 1, device: str = "cuda:0", dtype: torch.dtype = torch.float32):
         super().__init__()
 
         mean = [
@@ -560,11 +561,12 @@ 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, device="cuda:0", dtype=torch.float32) -> "WanVideoVAE":
+    def from_state_dict(cls, state_dict, parallelism=1, device="cuda:0", dtype=torch.float32) -> "WanVideoVAE":
         with no_init_weights():
-            model = torch.nn.utils.skip_init(cls, device=device, dtype=dtype)
+            model = torch.nn.utils.skip_init(cls, parallelism=parallelism, device=device, dtype=dtype)
         model.load_state_dict(state_dict, assign=True)
         model.to(device=device, dtype=dtype, non_blocking=True)
         return model
@@ -605,7 +607,7 @@ class WanVideoVAE(PreTrainedModel):
                 h_, w_ = h + size_h, w + size_w
                 tasks.append((h, h_, w, w_))
 
-        data_device = "cpu"
+        data_device = device if self.parallelism > 1 else "cpu"
         computation_device = device
 
         out_T = T * 4 - 3
@@ -620,7 +622,10 @@ class WanVideoVAE(PreTrainedModel):
             device=data_device,
         )
 
-        for i, (h, h_, w, w_) in enumerate(tqdm(tasks, desc="VAE DECODING")):
+        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()):
+                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)
 
@@ -649,8 +654,13 @@ 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:
+            if progress_callback is not None and not hide_progress_bar:
                 progress_callback(i + 1, len(tasks), "VAE DECODING")
+        if progress_callback is not None and not hide_progress_bar:
+            progress_callback(len(tasks), len(tasks), "VAE DECODING")
+        if self.parallelism > 1:
+            dist.all_reduce(values)
+            dist.all_reduce(weight)
         values = values / weight
         values = values.float().clamp_(-1, 1)
         return values
@@ -671,7 +681,7 @@ class WanVideoVAE(PreTrainedModel):
                 h_, w_ = h + size_h, w + size_w
                 tasks.append((h, h_, w, w_))
 
-        data_device = "cpu"
+        data_device = device if self.parallelism > 1 else "cpu"
         computation_device = device
 
         out_T = (T + 3) // 4
@@ -686,7 +696,10 @@ class WanVideoVAE(PreTrainedModel):
             device=data_device,
         )
 
-        for i, (h, h_, w, w_) in enumerate(tqdm(tasks, desc="VAE ENCODING")):
+        hide_progress_bar = self.parallelism > 1 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()):
+                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)
 
@@ -715,8 +728,13 @@ 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:
+            if progress_callback is not None and not hide_progress_bar:
                 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:
+            dist.all_reduce(values)
+            dist.all_reduce(weight)
         values = values / weight
         values = values.float()
         return values

{diffsynth_engine-0.2.8 → diffsynth_engine-0.2.9}/diffsynth_engine/pipelines/wan_video.py

@@ -442,8 +442,6 @@ class WanVideoPipeline(BasePipeline):
         tokenizer = WanT5Tokenizer(WAN_TOKENIZER_CONF_PATH, seq_len=512, clean="whitespace")
         text_encoder = WanTextEncoder.from_state_dict(t5_state_dict, device=init_device, dtype=model_config.t5_dtype)
 
-        vae = WanVideoVAE.from_state_dict(vae_state_dict, device=init_device, dtype=model_config.vae_dtype)
-
         image_encoder = None
         if model_config.image_encoder_path is not None:
             logger.info(f"loading state dict from {model_config.image_encoder_path} ...")
@@ -480,6 +478,20 @@ class WanVideoPipeline(BasePipeline):
             tp_degree = parallel_config["tp_degree"]
             use_fsdp = parallel_config["use_fsdp"]
             batch_cfg = True if use_cfg_parallel else batch_cfg
+
+            vae = WanVideoVAE.from_state_dict(
+                vae_state_dict, parallelism=parallelism, device="cpu", dtype=model_config.vae_dtype
+            )
+            vae = ParallelModel(
+                vae,
+                cfg_degree=1,
+                sp_ulysses_degree=parallelism,  # not real sequence parallel
+                sp_ring_degree=1,
+                tp_degree=1,
+                master_port=29501,
+                device="cuda",
+            )
+
             with LoRAContext():
                 dit = WanDiT.from_state_dict(
                     dit_state_dict,
@@ -499,6 +511,8 @@ class WanVideoPipeline(BasePipeline):
                     device="cuda",
                 )
         else:
+            vae = WanVideoVAE.from_state_dict(vae_state_dict, device=init_device, dtype=model_config.vae_dtype)
+
             with LoRAContext():
                 dit = WanDiT.from_state_dict(
                     dit_state_dict,
@@ -526,4 +540,5 @@ class WanVideoPipeline(BasePipeline):
         return pipe
 
     def __del__(self):
+        del self.vae
         del self.dit

{diffsynth_engine-0.2.8 → diffsynth_engine-0.2.9}/diffsynth_engine/utils/parallel.py

@@ -228,7 +228,6 @@ def parallelize_module(
 
 NCCL_TIMEOUT_SEC = int(os.environ.get("NCCL_TIMEOUT_SEC", 600))
 PARALLEL_FWD_TIMEOUT_SEC = int(os.environ.get("PARALLEL_FWD_TIMEOUT_SEC", 300))
-PARALLEL_LORA_TIMEOUT_SEC = int(os.environ.get("PARALLEL_LORA_TIMEOUT_SEC ", 60))
 
 
 def _worker_loop(
@@ -285,28 +284,26 @@ def _worker_loop(
 
         while True:
             if rank == 0:
-                kwargs = queue_in.get()
-                data = [kwargs]
+                data = queue_in.get()
             else:
-                data = [None]
+                data = [None, None, None]
             dist.broadcast_object_list(data, src=0)
-            kwargs = clone(data[0])
+            name, args, kwargs = clone(data)
             del data
 
-            y = None
-            if kwargs.get("method", None) == "load_loras":
-                module.load_loras(lora_args=kwargs["lora_args"], fused=kwargs["fused"])
-            elif kwargs.get("method", None) == "unload_loras":
-                module.unload_loras()
-            else:
-                kwargs = to_device(kwargs, device)
-                kwargs = split_and_get(kwargs, get_cfg_world_size(), 0, get_cfg_rank())
-                with torch.no_grad():
-                    y = module(**kwargs)
-                if get_sp_rank() == 0 and get_tp_rank() == 0:
-                    gathered = torch.zeros((get_cfg_world_size(), *y.shape[1:]), dtype=y.dtype, device=y.device)
-                    dist.all_gather_into_tensor(gathered, y, group=get_cfg_group())
-                    y = gathered
+            if not hasattr(module, name):
+                raise AttributeError(f'{module.__class__.__name__} has no attribute "{name}"')
+            if name != "load_loras":
+                args = split_and_get(to_device(args, device), get_cfg_world_size(), 0, get_cfg_rank())
+                kwargs = split_and_get(to_device(kwargs, device), get_cfg_world_size(), 0, get_cfg_rank())
+            with torch.no_grad():
+                y = getattr(module, name)(*args, **kwargs)
+
+            is_rank_zero_in_cfg_group = get_cfg_world_size() > 1 and get_sp_rank() == 0 and get_tp_rank() == 0
+            if y is not None and is_rank_zero_in_cfg_group:
+                gathered = torch.zeros((get_cfg_world_size(), *y.shape[1:]), dtype=y.dtype, device=y.device)
+                dist.all_gather_into_tensor(gathered, y, group=get_cfg_group())
+                y = gathered
 
             if rank == 0:
                 queue_out.put(y)
@@ -359,42 +356,9 @@ class ParallelModel(nn.Module):
             join=False,
         )
 
-    def load_loras(self, lora_args: List[Dict[str, any]], fused: bool = True):
-        self.queue_in.put(
-            {
-                "method": "load_loras",
-                "lora_args": lora_args,
-                "fused": fused,
-            }
-        )
-        try:
-            res = self.queue_out.get(timeout=PARALLEL_LORA_TIMEOUT_SEC)
-            if isinstance(res, Exception):
-                raise res
-        except Empty:
-            logger.error("ParallelModel load LoRA timeout")
-            raise RuntimeError("ParallelModel load LoRA timeout")
-        except Exception as e:
-            logger.error(f"ParallelModel load LoRA error: {e}")
-            raise RuntimeError(f"ParallelModel load LoRA error: {e}")
-        logger.info("ParallelModel load LoRA done")
-
-    def unload_loras(self):
-        self.queue_in.put({"method": "unload_loras"})
-        try:
-            res = self.queue_out.get(timeout=PARALLEL_LORA_TIMEOUT_SEC)
-            if isinstance(res, Exception):
-                raise res
-        except Empty:
-            logger.error("ParallelModel unload LoRA timeout")
-            raise RuntimeError("ParallelModel unload LoRA timeout")
-        except Exception as e:
-            logger.error(f"ParallelModel unload LoRA error: {e}")
-            raise RuntimeError(f"ParallelModel unload LoRA error: {e}")
-        logger.info("ParallelModel unload LoRA done")
-
-    def forward(self, **kwargs):
-        self.queue_in.put(kwargs)
+    def forward(self, *args, **kwargs):
+        data = ["forward", args, kwargs]
+        self.queue_in.put(data)
         try:
             res = self.queue_out.get(timeout=PARALLEL_FWD_TIMEOUT_SEC)
             if isinstance(res, Exception):
@@ -407,6 +371,25 @@ class ParallelModel(nn.Module):
             raise RuntimeError(f"ParallelModel forward error: {e}")
         return res
 
+    def __getattr__(self, name):
+        def wrapped_func(*args, **kwargs):
+            data = [name, args, kwargs]
+            self.queue_in.put(data)
+            try:
+                res = self.queue_out.get(timeout=PARALLEL_FWD_TIMEOUT_SEC)
+                if isinstance(res, Exception):
+                    raise res
+            except Empty:
+                logger.error(f"ParallelModel {name} timeout")
+                raise RuntimeError(f"ParallelModel {name} timeout")
+            except Exception as e:
+                logger.error(f"ParallelModel {name} error: {e}")
+                raise RuntimeError(f"ParallelModel {name} error: {e}")
+            logger.info(f"ParallelModel {name} done")
+            return res
+
+        return wrapped_func
+
     def __del__(self):
         # Send terminate signal to all workers
         for p in self.ctx.processes:

{diffsynth_engine-0.2.8 → diffsynth_engine-0.2.9}/diffsynth_engine.egg-info/PKG-INFO

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

{diffsynth_engine-0.2.8 → diffsynth_engine-0.2.9}/diffsynth_engine/pipelines/flux_image.py

@@ -801,13 +801,13 @@ class FluxImagePipeline(BasePipeline):
         controlnet_params: List[ControlNetParams] | ControlNetParams = [],
         progress_callback: Optional[Callable] = None,  # def progress_callback(current, total, status)
     ):
+        if not isinstance(controlnet_params, list):
+            controlnet_params = [controlnet_params]
         if self.control_type != ControlType.normal:
             assert controlnet_params and len(controlnet_params) == 1, "bfl_controlnet must have one controlnet"
 
         if input_image is not None:
             width, height = input_image.size
-        if not isinstance(controlnet_params, list):
-            controlnet_params = [controlnet_params]
         self.validate_image_size(height, width, minimum=64, multiple_of=16)
 
         noise = self.generate_noise((1, 16, height // 8, width // 8), seed=seed, device="cpu", dtype=self.dtype).to(