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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: diffsynth_engine
-Version: 0.6.1.dev2
+Version: 0.6.1.dev4
 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.dev4}/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.dev4}/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.dev4}/diffsynth_engine/pipelines/hunyuan3d_shape.py

@@ -1,4 +1,5 @@
 import torch
+from typing import Optional, Callable
 from tqdm import tqdm
 from PIL import Image
 from diffsynth_engine.algorithm.noise_scheduler.flow_match.recifited_flow import RecifitedFlowScheduler
@@ -179,6 +180,7 @@ class Hunyuan3DShapePipeline(BasePipeline):
         num_inference_steps: int = 50,
         guidance_scale: float = 7.5,
         seed: int = 42,
+        progress_callback: Optional[Callable] = None,  # def progress_callback(current, total, status)
     ):
         image_emb = self.encode_image(image)
 
@@ -197,4 +199,6 @@ class Hunyuan3DShapePipeline(BasePipeline):
             noise_pred, noise_pred_uncond = model_outputs.chunk(2)
             model_outputs = noise_pred_uncond + guidance_scale * (noise_pred - noise_pred_uncond)
             latents = self.sampler.step(latents, model_outputs, i)
+            if progress_callback is not None:
+                progress_callback(i, len(timesteps), "DENOISING")            
         return self.decode_latents(latents)

{diffsynth_engine-0.6.1.dev2 → diffsynth_engine-0.6.1.dev4}/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.dev4}/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.dev4
 Author: MuseAI x ModelScope
 Classifier: Programming Language :: Python :: 3
 Classifier: Operating System :: OS Independent