diffsynth-engine 0.6.1.dev36__py3-none-any.whl → 0.6.1.dev38__py3-none-any.whl

diffsynth_engine/configs/pipeline.py

@@ -301,11 +301,11 @@ class HunyuanPipelineConfig(BaseConfig):
 @dataclass
 class ZImagePipelineConfig(AttentionConfig, OptimizationConfig, ParallelConfig, BaseConfig):
     model_path: str | os.PathLike | List[str | os.PathLike]
-    model_dtype: torch.dtype = torch.float16
+    model_dtype: torch.dtype = torch.bfloat16
     vae_path: Optional[str | os.PathLike | List[str | os.PathLike]] = None
-    vae_dtype: torch.dtype = torch.float16
+    vae_dtype: torch.dtype = torch.bfloat16
     encoder_path: Optional[str | os.PathLike | List[str | os.PathLike]] = None
-    encoder_dtype: torch.dtype = torch.float16
+    encoder_dtype: torch.dtype = torch.bfloat16
 
     @classmethod
     def basic_config(

diffsynth_engine/pipelines/z_image.py

@@ -40,20 +40,54 @@ class ZImageLoRAConverter(LoRAStateDictConverter):
         for key, param in lora_state_dict.items():
             if "lora_A.weight" in key:
                 lora_b_key = key.replace("lora_A.weight", "lora_B.weight")
-                target_key = key.replace(".lora_A.weight", "").replace("transformer.", "")
+                target_key = key.replace(".lora_A.weight", "").replace("diffusion_model.", "")
 
-                if "attn.to_out.0" in target_key:
-                    target_key = target_key.replace("attn.to_out.0", "attn.to_out")
+                if "attention.to_out.0" in target_key:
+                    target_key = target_key.replace("attention.to_out.0", "attention.to_out")
+                if "adaLN_modulation.0" in target_key:
+                    target_key = target_key.replace("adaLN_modulation.0", "adaLN_modulation")
+
+                up = lora_state_dict[lora_b_key]
+                rank = up.shape[1]
 
                 dit_dict[target_key] = {
                     "down": param,
-                    "up": lora_state_dict[lora_b_key],
-                    "alpha": lora_state_dict.get(key.replace("lora_A.weight", "alpha"), None),
+                    "up": up,
+                    "rank": rank,
+                    "alpha": lora_state_dict.get(key.replace("lora_A.weight", "alpha"), rank),
                 }
+
         return {"dit": dit_dict}
 
+    def _from_diffsynth(self, lora_state_dict: Dict[str, torch.Tensor]) -> Dict[str, Dict[str, torch.Tensor]]:
+        dit_dict = {}
+        for key, param in lora_state_dict.items():
+            if "lora_A.default.weight" in key:
+                lora_b_key = key.replace("lora_A.default.weight", "lora_B.default.weight")
+                target_key = key.replace(".lora_A.default.weight", "")
+
+                if "attention.to_out.0" in target_key:
+                    target_key = target_key.replace("attention.to_out.0", "attention.to_out")
+
+                up = lora_state_dict[lora_b_key]
+                rank = up.shape[1]
+
+                dit_dict[target_key] = {
+                    "down": param,
+                    "up": up,
+                    "rank": rank,
+                    "alpha": lora_state_dict.get(key.replace("lora_A.default.weight", "alpha"), rank),
+                }
+
+        return {"dit": dit_dict}
+
+
     def convert(self, lora_state_dict: Dict[str, torch.Tensor]) -> Dict[str, Dict[str, torch.Tensor]]:
-        return self._from_diffusers(lora_state_dict)
+        key = list(lora_state_dict.keys())[0]
+        if key.startswith("diffusion_model."):
+            return self._from_diffusers(lora_state_dict)
+        else:
+            return self._from_diffsynth(lora_state_dict)
 
 
 class ZImagePipeline(BasePipeline):
@@ -180,7 +214,7 @@ class ZImagePipeline(BasePipeline):
     def update_weights(self, state_dicts: ZImageStateDicts) -> None:
         self.update_component(self.dit, state_dicts.model, self.config.device, self.config.model_dtype)
         self.update_component(
-            self.text_encoder, state_dicts.text_encoder, self.config.device, self.config.encoder_dtype
+            self.text_encoder, state_dicts.encoder, self.config.device, self.config.encoder_dtype
         )
         self.update_component(self.vae_decoder, state_dicts.vae, self.config.device, self.config.vae_dtype)
 
@@ -276,8 +310,8 @@ class ZImagePipeline(BasePipeline):
             comb_pred = self.predict_noise(latents, t, prompt_emb)[0]
         else:
             if not batch_cfg:
-                positive_noise_pred = self.predict_noise(latents, t, prompt_emb)
-                negative_noise_pred = self.predict_noise(latents, t, negative_prompt_emb)
+                positive_noise_pred = self.predict_noise(latents, t, prompt_emb)[0]
+                negative_noise_pred = self.predict_noise(latents, t, negative_prompt_emb)[0]
             else:
                 latents_input = torch.cat([latents, latents], dim=0)
                 t = torch.cat([t, t], dim=0)
@@ -360,6 +394,7 @@ class ZImagePipeline(BasePipeline):
                 prompt_emb=prompt_embeds,
                 negative_prompt_emb=negative_prompt_embeds,
                 batch_cfg=self.config.batch_cfg,
+                cfg_scale=cfg_scale,
                 cfg_truncation=cfg_truncation,
                 cfg_normalization=cfg_normalization,
             )

{diffsynth_engine-0.6.1.dev36.dist-info → diffsynth_engine-0.6.1.dev38.dist-info}/METADATA

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

{diffsynth_engine-0.6.1.dev36.dist-info → diffsynth_engine-0.6.1.dev38.dist-info}/RECORD

@@ -86,7 +86,7 @@ diffsynth_engine/conf/tokenizers/z_image/tokenizer/tokenizer_config.json,sha256=
 diffsynth_engine/conf/tokenizers/z_image/tokenizer/vocab.json,sha256=yhDX6fs-0YV13R4neiV5wW0QjjLydDloSvoOELFECRA,2776833
 diffsynth_engine/configs/__init__.py,sha256=biluGSEw78PPwO7XFlms16iuWXDiM0Eg_qsOMMTY0NQ,1409
 diffsynth_engine/configs/controlnet.py,sha256=f3vclyP3lcAjxDGD9C1vevhqqQ7W2LL_c6Wye0uxk3Q,1180
-diffsynth_engine/configs/pipeline.py,sha256=0WmKz_mykmJkRCGwv9DjuN8s27LppkD_Ier4VtovZSg,15307
+diffsynth_engine/configs/pipeline.py,sha256=RqhPAZOCpIMkFk-OsfiNYlqpqM-7B52ny0Zcr9Ix7wY,15310
 diffsynth_engine/kernels/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 diffsynth_engine/models/__init__.py,sha256=8Ze7cSE8InetgXWTNb0neVA2Q44K7WlE-h7O-02m2sY,119
 diffsynth_engine/models/base.py,sha256=svao__9WH8VNcyXz5o5dzywYXDcGV0YV9IfkLzDKews,2558
@@ -160,7 +160,7 @@ diffsynth_engine/pipelines/sdxl_image.py,sha256=v7ZACGPb6EcBunL6e5E9jynSQjE7GQx8
 diffsynth_engine/pipelines/utils.py,sha256=HZbJHErNJS1DhlwJKvZ9dY7Kh8Zdlsw3zE2e88TYGRY,2277
 diffsynth_engine/pipelines/wan_s2v.py,sha256=QHlCLMqlmnp55iYm2mzg4qCq4jceRAP3Zt5Mubz3mAM,29384
 diffsynth_engine/pipelines/wan_video.py,sha256=9xjSvQ4mlVEDdaL6QuUURj4iyxhJ2xABBphQjkfzK8s,31323
-diffsynth_engine/pipelines/z_image.py,sha256=gSBhKV7TBL9xvCUrABdZA0kNqQzPuawmEv8OcI6KTcs,14756
+diffsynth_engine/pipelines/z_image.py,sha256=VvqjxsKRsmP2tfWg9nDlcQu5oEzIRFa2wtuArzjQAlk,16151
 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
@@ -199,8 +199,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.dev36.dist-info/licenses/LICENSE,sha256=x7aBqQuVI0IYnftgoTPI_A0I_rjdjPPQkjnU6N2nikM,11346
-diffsynth_engine-0.6.1.dev36.dist-info/METADATA,sha256=ShMi7F-NEzgV7Tse-BFIzIa5N7t8FLVpAHA1FkDWp14,1164
-diffsynth_engine-0.6.1.dev36.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-diffsynth_engine-0.6.1.dev36.dist-info/top_level.txt,sha256=6zgbiIzEHLbhgDKRyX0uBJOV3F6VnGGBRIQvSiYYn6w,17
-diffsynth_engine-0.6.1.dev36.dist-info/RECORD,,
+diffsynth_engine-0.6.1.dev38.dist-info/licenses/LICENSE,sha256=x7aBqQuVI0IYnftgoTPI_A0I_rjdjPPQkjnU6N2nikM,11346
+diffsynth_engine-0.6.1.dev38.dist-info/METADATA,sha256=0fI0prUJox3z_sDzvhl-wh6wlCCYCA7N-naxpobysL0,1164
+diffsynth_engine-0.6.1.dev38.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+diffsynth_engine-0.6.1.dev38.dist-info/top_level.txt,sha256=6zgbiIzEHLbhgDKRyX0uBJOV3F6VnGGBRIQvSiYYn6w,17
+diffsynth_engine-0.6.1.dev38.dist-info/RECORD,,