diffsynth-engine 0.6.1.dev2__py3-none-any.whl → 0.6.1.dev3__py3-none-any.whl

diffsynth_engine/pipelines/base.py

@@ -51,6 +51,20 @@ class BasePipeline:
     def from_state_dict(cls, state_dicts: BaseStateDicts, pipeline_config: BaseConfig) -> "BasePipeline":
         raise NotImplementedError()
 
+    def update_weights(self, state_dicts: BaseStateDicts) -> None:
+        raise NotImplementedError()
+    
+    @staticmethod
+    def update_component(
+        component: torch.nn.Module,
+        state_dict: Dict[str, torch.Tensor],
+        device: str,
+        dtype: torch.dtype,
+    ) -> None:
+        if component and state_dict:
+            component.load_state_dict(state_dict, assign=True)
+            component.to(device=device, dtype=dtype, non_blocking=True)
+
     def load_loras(self, lora_list: List[Tuple[str, float]], fused: bool = True, save_original_weight: bool = False):
         for lora_path, lora_scale in lora_list:
             logger.info(f"loading lora from {lora_path} with scale {lora_scale}")

diffsynth_engine/pipelines/flux_image.py

@@ -573,6 +573,13 @@ class FluxImagePipeline(BasePipeline):
             pipe.compile()
         return pipe
 
+    def update_weights(self, state_dicts: FluxStateDicts) -> None:
+        self.update_component(self.dit, state_dicts.model, self.config.device, self.config.model_dtype)
+        self.update_component(self.text_encoder_1, state_dicts.clip, self.config.device, self.config.clip_dtype)
+        self.update_component(self.text_encoder_2, state_dicts.t5, self.config.device, self.config.t5_dtype)
+        self.update_component(self.vae_decoder, state_dicts.vae, self.config.device, self.config.vae_dtype)
+        self.update_component(self.vae_encoder, state_dicts.vae, self.config.device, self.config.vae_dtype)
+
     def compile(self):
         self.dit.compile_repeated_blocks(dynamic=True)
 

diffsynth_engine/pipelines/qwen_image.py

@@ -254,6 +254,11 @@ class QwenImagePipeline(BasePipeline):
             pipe.compile()
         return pipe
 
+    def update_weights(self, state_dicts: QwenImageStateDicts) -> None:
+        self.update_component(self.dit, state_dicts.model, self.config.device, self.config.model_dtype)
+        self.update_component(self.encoder, state_dicts.encoder, self.config.device, self.config.encoder_dtype)
+        self.update_component(self.vae, state_dicts.vae, self.config.device, self.config.vae_dtype)
+
     def compile(self):
         self.dit.compile_repeated_blocks(dynamic=True)
 

{diffsynth_engine-0.6.1.dev2.dist-info → diffsynth_engine-0.6.1.dev3.dist-info}/METADATA

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

{diffsynth_engine-0.6.1.dev2.dist-info → diffsynth_engine-0.6.1.dev3.dist-info}/RECORD

@@ -139,10 +139,10 @@ diffsynth_engine/models/wan/wan_s2v_dit.py,sha256=sOJsSs1snI-ZGPJS8utstmgj0wcYwl
 diffsynth_engine/models/wan/wan_text_encoder.py,sha256=OERlmwOqthAFPNnnT2sXJ4OjyyRmsRLx7VGp1zlBkLU,11021
 diffsynth_engine/models/wan/wan_vae.py,sha256=ogXrVlwmzXR4iLxjSCkBPtYW8KWebnvvd2UtPZeoziY,38853
 diffsynth_engine/pipelines/__init__.py,sha256=jh-4LSJ0vqlXiT8BgFgRIQxuAr2atEPyHrxXWj-Ud1U,604
-diffsynth_engine/pipelines/base.py,sha256=RTkVwWaWXr5ujqn5-UBHvdPddYwr-uvChj9-fmoXrms,13729
-diffsynth_engine/pipelines/flux_image.py,sha256=a-MaHuguV7Z6LJukC_Tvp7d9_2dnrAaJZ4MZH_sKsKo,49116
+diffsynth_engine/pipelines/base.py,sha256=7x7gEdCk_DRnGDMdPGLvNPlk-Yn2p0yQ8pNvr59i-hU,14199
+diffsynth_engine/pipelines/flux_image.py,sha256=LET1gPlkXJN2xE22GRVjUgWJH0ZZBFnAe3bIvJrjb1s,49726
 diffsynth_engine/pipelines/hunyuan3d_shape.py,sha256=fwNKET54KjCiWDpW2S1Fk-p3nfJreZ-RH7p46VLawEQ,7911
-diffsynth_engine/pipelines/qwen_image.py,sha256=EAYoq1QkdOSie_yVZG9enxJJRcncwVFPfDftMo-3zBA,23745
+diffsynth_engine/pipelines/qwen_image.py,sha256=vWe1M3FU-aqckvrVeJlZgMIN55qnRz9hZ2AxGCxro1Y,24134
 diffsynth_engine/pipelines/sd_image.py,sha256=nr-Nhsnomq8CsUqhTM3i2l2zG01YjwXdfRXgr_bC3F0,17891
 diffsynth_engine/pipelines/sdxl_image.py,sha256=FaihRd9Rt_qtqup2xEbHViVIFwFZVyvekYW4lCodNKY,21692
 diffsynth_engine/pipelines/utils.py,sha256=lk7sFGEk-fGjgadLpwwppHKG-yZ0RC-4ZmHW7pRRe8A,473
@@ -185,8 +185,8 @@ diffsynth_engine/utils/video.py,sha256=GoMyc2as4_VqfWX4pjQyAWh9QObsFMov42zADVZNa
 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.dev2.dist-info/licenses/LICENSE,sha256=x7aBqQuVI0IYnftgoTPI_A0I_rjdjPPQkjnU6N2nikM,11346
-diffsynth_engine-0.6.1.dev2.dist-info/METADATA,sha256=6bwBcpLV0q3yBPFYlDecEtAJP1ga-zi7u9tioW41D7k,1163
-diffsynth_engine-0.6.1.dev2.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-diffsynth_engine-0.6.1.dev2.dist-info/top_level.txt,sha256=6zgbiIzEHLbhgDKRyX0uBJOV3F6VnGGBRIQvSiYYn6w,17
-diffsynth_engine-0.6.1.dev2.dist-info/RECORD,,
+diffsynth_engine-0.6.1.dev3.dist-info/licenses/LICENSE,sha256=x7aBqQuVI0IYnftgoTPI_A0I_rjdjPPQkjnU6N2nikM,11346
+diffsynth_engine-0.6.1.dev3.dist-info/METADATA,sha256=CShZsvCWGPzvwFuc_85oCrlb-KBIAEdvV2iA7MJACRQ,1163
+diffsynth_engine-0.6.1.dev3.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+diffsynth_engine-0.6.1.dev3.dist-info/top_level.txt,sha256=6zgbiIzEHLbhgDKRyX0uBJOV3F6VnGGBRIQvSiYYn6w,17
+diffsynth_engine-0.6.1.dev3.dist-info/RECORD,,