diffsynth-engine 0.6.1.dev2__tar.gz → 0.6.1.dev3__tar.gz

{diffsynth_engine-0.6.1.dev2 → diffsynth_engine-0.6.1.dev3}/PKG-INFO

@@ -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 → diffsynth_engine-0.6.1.dev3}/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-0.6.1.dev2 → diffsynth_engine-0.6.1.dev3}/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-0.6.1.dev2 → diffsynth_engine-0.6.1.dev3}/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 → diffsynth_engine-0.6.1.dev3}/diffsynth_engine.egg-info/PKG-INFO

@@ -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