diffsynth-engine 0.6.1.dev38__tar.gz → 0.6.1.dev39__tar.gz

{diffsynth_engine-0.6.1.dev38 → diffsynth_engine-0.6.1.dev39}/PKG-INFO

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

{diffsynth_engine-0.6.1.dev38 → diffsynth_engine-0.6.1.dev39}/diffsynth_engine/__init__.py

@@ -27,6 +27,7 @@ from .pipelines import (
     SDXLImagePipeline,
     FluxImagePipeline,
     WanVideoPipeline,
+    WanDMDPipeline,
     QwenImagePipeline,
     Hunyuan3DShapePipeline,
 )
@@ -81,6 +82,7 @@ __all__ = [
     "FluxIPAdapter",
     "FluxRedux",
     "WanVideoPipeline",
+    "WanDMDPipeline",
     "QwenImagePipeline",
     "Hunyuan3DShapePipeline",
     "FluxInpaintingTool",

{diffsynth_engine-0.6.1.dev38 → diffsynth_engine-0.6.1.dev39}/diffsynth_engine/pipelines/__init__.py

@@ -4,6 +4,7 @@ from .sdxl_image import SDXLImagePipeline
 from .sd_image import SDImagePipeline
 from .wan_video import WanVideoPipeline
 from .wan_s2v import WanSpeech2VideoPipeline
+from .wan_dmd import WanDMDPipeline
 from .qwen_image import QwenImagePipeline
 from .hunyuan3d_shape import Hunyuan3DShapePipeline
 from .z_image import ZImagePipeline
@@ -16,6 +17,7 @@ __all__ = [
     "SDImagePipeline",
     "WanVideoPipeline",
     "WanSpeech2VideoPipeline",
+    "WanDMDPipeline",
     "QwenImagePipeline",
     "Hunyuan3DShapePipeline",
     "ZImagePipeline",

{diffsynth_engine-0.6.1.dev38 → diffsynth_engine-0.6.1.dev39}/diffsynth_engine/pipelines/base.py

@@ -145,7 +145,7 @@ class BasePipeline:
         self.load_loras([(path, scale)], fused, save_original_weight)
 
     def apply_scheduler_config(self, scheduler_config: Dict):
-        pass
+        self.noise_scheduler.update_config(scheduler_config)
 
     def unload_loras(self):
         raise NotImplementedError()

{diffsynth_engine-0.6.1.dev38 → diffsynth_engine-0.6.1.dev39}/diffsynth_engine/pipelines/qwen_image.py

@@ -393,9 +393,6 @@ class QwenImagePipeline(BasePipeline):
         self.dit.unload_loras()
         self.noise_scheduler.restore_config()
 
-    def apply_scheduler_config(self, scheduler_config: Dict):
-        self.noise_scheduler.update_config(scheduler_config)
-
     def prepare_latents(
         self,
         latents: torch.Tensor,

diffsynth_engine-0.6.1.dev39/diffsynth_engine/pipelines/wan_dmd.py

@@ -0,0 +1,111 @@
+import torch
+import torch.distributed as dist
+from typing import Callable, List, Optional
+from tqdm import tqdm
+from PIL import Image
+
+from diffsynth_engine.pipelines.wan_video import WanVideoPipeline
+
+
+class WanDMDPipeline(WanVideoPipeline):
+    def prepare_latents(
+        self,
+        latents,
+        denoising_step_list,
+    ):
+        height, width = latents.shape[-2:]
+        height, width = height * self.upsampling_factor, width * self.upsampling_factor
+        sigmas, timesteps = self.noise_scheduler.schedule(num_inference_steps=1000)
+        sigmas = sigmas[[1000 - t for t in denoising_step_list] + [-1]]
+        timesteps = timesteps[[1000 - t for t in denoising_step_list]]
+        init_latents = latents.clone()
+
+        return init_latents, latents, sigmas, timesteps
+
+    @torch.no_grad()
+    def __call__(
+        self,
+        prompt,
+        input_image: Image.Image | None = None,
+        seed=None,
+        height=480,
+        width=832,
+        num_frames=81,
+        denoising_step_list: List[int] = None,
+        progress_callback: Optional[Callable] = None,  # def progress_callback(current, total, status)
+    ):
+        denoising_step_list = [1000, 750, 500, 250] if denoising_step_list is None else denoising_step_list
+        divisor = 32 if self.vae.z_dim == 48 else 16  # 32 for wan2.2 vae, 16 for wan2.1 vae
+        assert height % divisor == 0 and width % divisor == 0, f"height and width must be divisible by {divisor}"
+        assert (num_frames - 1) % 4 == 0, "num_frames must be 4X+1"
+
+        # Initialize noise
+        if dist.is_initialized() and seed is None:
+            raise ValueError("must provide a seed when parallelism is enabled")
+        noise = self.generate_noise(
+            (
+                1,
+                self.vae.z_dim,
+                (num_frames - 1) // 4 + 1,
+                height // self.upsampling_factor,
+                width // self.upsampling_factor,
+            ),
+            seed=seed,
+            device="cpu",
+            dtype=torch.float32,
+        ).to(self.device)
+        init_latents, latents, sigmas, timesteps = self.prepare_latents(noise, denoising_step_list)
+        mask = torch.ones((1, 1, *latents.shape[2:]), dtype=latents.dtype, device=latents.device)
+
+        # Encode prompts
+        self.load_models_to_device(["text_encoder"])
+        prompt_emb_posi = self.encode_prompt(prompt)
+        prompt_emb_nega = None
+
+        # Encode image
+        image_clip_feature = self.encode_clip_feature(input_image, height, width)
+        image_y = self.encode_vae_feature(input_image, num_frames, height, width)
+        image_latents = self.encode_image_latents(input_image, height, width)
+        if image_latents is not None:
+            latents[:, :, : image_latents.shape[2], :, :] = image_latents
+            init_latents = latents.clone()
+            mask[:, :, : image_latents.shape[2], :, :] = 0
+
+        # Initialize sampler
+        self.sampler.initialize(sigmas=sigmas)
+
+        # Denoise
+        hide_progress = dist.is_initialized() and dist.get_rank() != 0
+        for i, timestep in enumerate(tqdm(timesteps, disable=hide_progress)):
+            if timestep.item() / 1000 >= self.config.boundary:
+                self.load_models_to_device(["dit"])
+                model = self.dit
+            else:
+                self.load_models_to_device(["dit2"])
+                model = self.dit2
+
+            timestep = timestep * mask[:, :, :, ::2, ::2].flatten()  # seq_len
+            timestep = timestep.to(dtype=self.dtype, device=self.device)
+            # Classifier-free guidance
+            noise_pred = self.predict_noise_with_cfg(
+                model=model,
+                latents=latents,
+                timestep=timestep,
+                positive_prompt_emb=prompt_emb_posi,
+                negative_prompt_emb=prompt_emb_nega,
+                image_clip_feature=image_clip_feature,
+                image_y=image_y,
+                cfg_scale=1.0,
+                batch_cfg=self.config.batch_cfg,
+            )
+            # Scheduler
+            latents = self.sampler.step(latents, noise_pred, i)
+            latents = latents * mask + init_latents * (1 - mask)
+            if progress_callback is not None:
+                progress_callback(i + 1, len(timesteps), "DENOISING")
+
+        # Decode
+        self.load_models_to_device(["vae"])
+        frames = self.decode_video(latents, progress_callback=progress_callback)
+        frames = self.vae_output_to_image(frames)
+        return frames

{diffsynth_engine-0.6.1.dev38 → diffsynth_engine-0.6.1.dev39}/diffsynth_engine/pipelines/wan_video.py

@@ -43,6 +43,24 @@ class WanLoRAConverter(LoRAStateDictConverter):
             dit_dict[key] = lora_args
         return {"dit": dit_dict}
 
+    def _from_diffusers(self, state_dict):
+        dit_dict = {}
+        for key, param in state_dict.items():
+            if ".lora_down.weight" not in key:
+                continue
+
+            lora_args = {}
+            lora_args["up"] = state_dict[key.replace(".lora_down.weight", ".lora_up.weight")]
+            lora_args["down"] = param
+            lora_args["rank"] = lora_args["up"].shape[1]
+            if key.replace(".lora_down.weight", ".alpha") in state_dict:
+                lora_args["alpha"] = state_dict[key.replace(".lora_down.weight", ".alpha")]
+            else:
+                lora_args["alpha"] = lora_args["rank"]
+            key = key.replace("diffusion_model.", "").replace(".lora_down.weight", "")
+            dit_dict[key] = lora_args
+        return {"dit": dit_dict}
+
     def _from_civitai(self, state_dict):
         dit_dict = {}
         for key, param in state_dict.items():
@@ -86,6 +104,9 @@ class WanLoRAConverter(LoRAStateDictConverter):
         if "lora_unet_blocks_0_cross_attn_k.lora_down.weight" in state_dict:
             state_dict = self._from_fun(state_dict)
             logger.info("use fun format state dict")
+        elif "diffusion_model.blocks.0.cross_attn.k.lora_down.weight" in state_dict:
+            state_dict = self._from_diffusers(state_dict)
+            logger.info("use diffusers format state dict")
         elif "diffusion_model.blocks.0.cross_attn.k.lora_A.weight" in state_dict:
             state_dict = self._from_civitai(state_dict)
             logger.info("use civitai format state dict")
@@ -480,8 +501,8 @@ class WanVideoPipeline(BasePipeline):
 
         dit_state_dict, dit2_state_dict = None, None
         if isinstance(config.model_path, list):
-            high_noise_model_ckpt = [path for path in config.model_path if "high_noise_model" in path]
-            low_noise_model_ckpt = [path for path in config.model_path if "low_noise_model" in path]
+            high_noise_model_ckpt = [path for path in config.model_path if "high_noise" in path]
+            low_noise_model_ckpt = [path for path in config.model_path if "low_noise" in path]
             if high_noise_model_ckpt and low_noise_model_ckpt:
                 logger.info(f"loading high noise model state dict from {high_noise_model_ckpt} ...")
                 dit_state_dict = cls.load_model_checkpoint(
@@ -681,8 +702,9 @@ class WanVideoPipeline(BasePipeline):
                 config.attn_params = VideoSparseAttentionParams(sparsity=0.9)
 
     def update_weights(self, state_dicts: WanStateDicts) -> None:
-        is_dual_model_state_dict = (isinstance(state_dicts.model, dict) and
-                                     ("high_noise_model" in state_dicts.model or "low_noise_model" in state_dicts.model))
+        is_dual_model_state_dict = isinstance(state_dicts.model, dict) and (
+            "high_noise_model" in state_dicts.model or "low_noise_model" in state_dicts.model
+        )
         is_dual_model_pipeline = self.dit2 is not None
 
         if is_dual_model_state_dict != is_dual_model_pipeline:
@@ -694,15 +716,21 @@ class WanVideoPipeline(BasePipeline):
 
         if is_dual_model_state_dict:
             if "high_noise_model" in state_dicts.model:
-                self.update_component(self.dit, state_dicts.model["high_noise_model"], self.config.device, self.config.model_dtype)
+                self.update_component(
+                    self.dit, state_dicts.model["high_noise_model"], self.config.device, self.config.model_dtype
+                )
             if "low_noise_model" in state_dicts.model:
-                self.update_component(self.dit2, state_dicts.model["low_noise_model"], self.config.device, self.config.model_dtype)
+                self.update_component(
+                    self.dit2, state_dicts.model["low_noise_model"], self.config.device, self.config.model_dtype
+                )
         else:
             self.update_component(self.dit, state_dicts.model, self.config.device, self.config.model_dtype)
 
         self.update_component(self.text_encoder, state_dicts.t5, self.config.device, self.config.t5_dtype)
         self.update_component(self.vae, state_dicts.vae, self.config.device, self.config.vae_dtype)
-        self.update_component(self.image_encoder, state_dicts.image_encoder, self.config.device, self.config.image_encoder_dtype)
+        self.update_component(
+            self.image_encoder, state_dicts.image_encoder, self.config.device, self.config.image_encoder_dtype
+        )
 
     def compile(self):
         self.dit.compile_repeated_blocks()

{diffsynth_engine-0.6.1.dev38 → diffsynth_engine-0.6.1.dev39}/diffsynth_engine.egg-info/PKG-INFO

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

{diffsynth_engine-0.6.1.dev38 → diffsynth_engine-0.6.1.dev39}/diffsynth_engine.egg-info/SOURCES.txt

@@ -174,6 +174,7 @@ diffsynth_engine/pipelines/qwen_image.py
 diffsynth_engine/pipelines/sd_image.py
 diffsynth_engine/pipelines/sdxl_image.py
 diffsynth_engine/pipelines/utils.py
+diffsynth_engine/pipelines/wan_dmd.py
 diffsynth_engine/pipelines/wan_s2v.py
 diffsynth_engine/pipelines/wan_video.py
 diffsynth_engine/pipelines/z_image.py