diffsynth-engine 0.6.1.dev19__py3-none-any.whl → 0.6.1.dev21__py3-none-any.whl

diffsynth_engine/pipelines/base.py

@@ -2,7 +2,7 @@ import os
 import torch
 import numpy as np
 from einops import rearrange
-from typing import Dict, List, Tuple, Union
+from typing import Dict, List, Tuple, Union, Optional
 from PIL import Image
 
 from diffsynth_engine.configs import BaseConfig, BaseStateDicts, LoraConfig
@@ -70,7 +70,11 @@ class BasePipeline:
         lora_list: List[Tuple[str, Union[float, LoraConfig]]],
         fused: bool = True,
         save_original_weight: bool = False,
+        lora_converter: Optional[LoRAStateDictConverter] = None,
     ):
+        if not lora_converter:
+            lora_converter = self.lora_converter
+
         for lora_path, lora_item in lora_list:
             if isinstance(lora_item, float):
                 lora_scale = lora_item
@@ -86,7 +90,7 @@ class BasePipeline:
                 self.apply_scheduler_config(scheduler_config)
                 logger.info(f"Applied scheduler args from LoraConfig: {scheduler_config}")
 
-            lora_state_dict = self.lora_converter.convert(state_dict)
+            lora_state_dict = lora_converter.convert(state_dict)
             for model_name, state_dict in lora_state_dict.items():
                 model = getattr(self, model_name)
                 lora_args = []

diffsynth_engine/pipelines/flux_image.py

@@ -830,7 +830,7 @@ class FluxImagePipeline(BasePipeline):
                 masked_image = image.clone()
                 masked_image[(mask > 0.5).repeat(1, 3, 1, 1)] = -1
                 latent = self.encode_image(masked_image)
-                mask = torch.nn.functional.interpolate(mask, size=(latent.shape[2], latent.shape[3]))
+                mask = torch.nn.functional.interpolate(mask, size=(latent.shape[2], latent.shape[3])).to(latent.dtype)
                 mask = 1 - mask
                 latent = torch.cat([latent, mask], dim=1)
             elif self.config.control_type == ControlType.bfl_fill:

diffsynth_engine/pipelines/wan_video.py

@@ -95,8 +95,14 @@ class WanLoRAConverter(LoRAStateDictConverter):
         return state_dict
 
 
+class WanLowNoiseLoRAConverter(WanLoRAConverter):
+    def convert(self, state_dict):
+        return {"dit2": super().convert(state_dict)["dit"]}
+
+
 class WanVideoPipeline(BasePipeline):
     lora_converter = WanLoRAConverter()
+    low_noise_lora_converter = WanLowNoiseLoRAConverter()
 
     def __init__(
         self,
@@ -133,7 +139,13 @@ class WanVideoPipeline(BasePipeline):
         self.image_encoder = image_encoder
         self.model_names = ["text_encoder", "dit", "dit2", "vae", "image_encoder"]
 
-    def load_loras(self, lora_list: List[Tuple[str, float]], fused: bool = True, save_original_weight: bool = False):
+    def load_loras(
+        self,
+        lora_list: List[Tuple[str, float]],
+        fused: bool = True,
+        save_original_weight: bool = False,
+        lora_converter: Optional[WanLoRAConverter] = None
+    ):
         assert self.config.tp_degree is None or self.config.tp_degree == 1, (
             "load LoRA is not allowed when tensor parallel is enabled; "
             "set tp_degree=None or tp_degree=1 during pipeline initialization"
@@ -142,10 +154,20 @@ class WanVideoPipeline(BasePipeline):
             "load fused LoRA is not allowed when fully sharded data parallel is enabled; "
             "either load LoRA with fused=False or set use_fsdp=False during pipeline initialization"
         )
-        super().load_loras(lora_list, fused, save_original_weight)
+        super().load_loras(lora_list, fused, save_original_weight, lora_converter)
+
+    def load_loras_low_noise(self, lora_list: List[Tuple[str, float]], fused: bool = True, save_original_weight: bool = False):
+        assert self.dit2 is not None, "low noise LoRA can only be applied to Wan2.2"
+        self.load_loras(lora_list, fused, save_original_weight, self.low_noise_lora_converter)
+
+    def load_loras_high_noise(self, lora_list: List[Tuple[str, float]], fused: bool = True, save_original_weight: bool = False):
+        assert self.dit2 is not None, "high noise LoRA can only be applied to Wan2.2"
+        self.load_loras(lora_list, fused, save_original_weight)
 
     def unload_loras(self):
         self.dit.unload_loras()
+        if self.dit2 is not None:
+            self.dit2.unload_loras()
         self.text_encoder.unload_loras()
 
     def get_default_fps(self) -> int:

{diffsynth_engine-0.6.1.dev19.dist-info → diffsynth_engine-0.6.1.dev21.dist-info}/METADATA

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

{diffsynth_engine-0.6.1.dev19.dist-info → diffsynth_engine-0.6.1.dev21.dist-info}/RECORD

@@ -139,15 +139,15 @@ diffsynth_engine/models/wan/wan_s2v_dit.py,sha256=j63ulcWLY4XGITOKUMGX292LtSEtP-
 diffsynth_engine/models/wan/wan_text_encoder.py,sha256=OERlmwOqthAFPNnnT2sXJ4OjyyRmsRLx7VGp1zlBkLU,11021
 diffsynth_engine/models/wan/wan_vae.py,sha256=dC7MoUFeXRL7SIY0LG1OOUiZW-pp9IbXCghutMxpXr4,38889
 diffsynth_engine/pipelines/__init__.py,sha256=jh-4LSJ0vqlXiT8BgFgRIQxuAr2atEPyHrxXWj-Ud1U,604
-diffsynth_engine/pipelines/base.py,sha256=B6Md10eeAK4itILjx3biRCFwYk2usgSv7v2V9vd4fjA,14842
-diffsynth_engine/pipelines/flux_image.py,sha256=Dpy8AkwywuLAhvJ6cjg5TgzhSUgFQtv6p2JTTkzUHbo,50919
+diffsynth_engine/pipelines/base.py,sha256=BWW7LW0E2qwu8G-6bP3nmeO7VCQxC8srOo8tE4aKA4o,14993
+diffsynth_engine/pipelines/flux_image.py,sha256=vJKvnYmeeQVX2O1Zjtm4NLrltBp66VSZ-KjAUqJ8zJ8,50936
 diffsynth_engine/pipelines/hunyuan3d_shape.py,sha256=TNV0Wr09Dj2bzzlpua9WioCClOj3YiLfE6utI9aWL8A,8164
 diffsynth_engine/pipelines/qwen_image.py,sha256=jt4rg-U5qWsFD0kUeDwKzgIiTAC80Cj8aq1YQOR1_-k,33052
 diffsynth_engine/pipelines/sd_image.py,sha256=nr-Nhsnomq8CsUqhTM3i2l2zG01YjwXdfRXgr_bC3F0,17891
 diffsynth_engine/pipelines/sdxl_image.py,sha256=v7ZACGPb6EcBunL6e5E9jynSQjE7GQx8etEV-ZLP91g,21704
 diffsynth_engine/pipelines/utils.py,sha256=lk7sFGEk-fGjgadLpwwppHKG-yZ0RC-4ZmHW7pRRe8A,473
 diffsynth_engine/pipelines/wan_s2v.py,sha256=3Lkdwf5CYH2fyiD2XeZIqHUfjThsNKV9F_tQXQ-7uoU,29559
-diffsynth_engine/pipelines/wan_video.py,sha256=x4xnP_4VAwGW04Ja78eecfLqyzMnqdgO1J9cK-DZpv4,28173
+diffsynth_engine/pipelines/wan_video.py,sha256=CF8098TIvhYTrrdfuFR7K4GpgFUezONROFJG2LL7wQk,29151
 diffsynth_engine/processor/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 diffsynth_engine/processor/canny_processor.py,sha256=hV30NlblTkEFUAmF_O-LJrNlGVM2SFrqq6okfF8VpOo,602
 diffsynth_engine/processor/depth_processor.py,sha256=dQvs3JsnyMbz4dyI9QoR8oO-mMFBFAgNvgqeCoaU5jk,1532
@@ -185,8 +185,8 @@ diffsynth_engine/utils/video.py,sha256=8FCaeqIdUsWMgWI_6SO9SPynsToGcLCQAVYFTc4CD
 diffsynth_engine/utils/memory/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 diffsynth_engine/utils/memory/linear_regression.py,sha256=oW_EQEw13oPoyUrxiL8A7Ksa5AuJ2ynI2qhCbfAuZbg,3930
 diffsynth_engine/utils/memory/memory_predcit_model.py,sha256=EXprSl_zlVjgfMWNXP-iw83Ot3hyMcgYaRPv-dvyL84,3943
-diffsynth_engine-0.6.1.dev19.dist-info/licenses/LICENSE,sha256=x7aBqQuVI0IYnftgoTPI_A0I_rjdjPPQkjnU6N2nikM,11346
-diffsynth_engine-0.6.1.dev19.dist-info/METADATA,sha256=KQ9a1ITP4r5RnWNUKEGJnnt5dduwknR3rCU2K5ETBC4,1164
-diffsynth_engine-0.6.1.dev19.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-diffsynth_engine-0.6.1.dev19.dist-info/top_level.txt,sha256=6zgbiIzEHLbhgDKRyX0uBJOV3F6VnGGBRIQvSiYYn6w,17
-diffsynth_engine-0.6.1.dev19.dist-info/RECORD,,
+diffsynth_engine-0.6.1.dev21.dist-info/licenses/LICENSE,sha256=x7aBqQuVI0IYnftgoTPI_A0I_rjdjPPQkjnU6N2nikM,11346
+diffsynth_engine-0.6.1.dev21.dist-info/METADATA,sha256=tdKUjrwahEQ72SA-YSPu8LsaswLKJuDrjEZI_6nYySM,1164
+diffsynth_engine-0.6.1.dev21.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+diffsynth_engine-0.6.1.dev21.dist-info/top_level.txt,sha256=6zgbiIzEHLbhgDKRyX0uBJOV3F6VnGGBRIQvSiYYn6w,17
+diffsynth_engine-0.6.1.dev21.dist-info/RECORD,,