diffsynth-engine 0.4.3.dev6__tar.gz → 0.4.3.dev7__tar.gz

{diffsynth_engine-0.4.3.dev6 → diffsynth_engine-0.4.3.dev7}/PKG-INFO

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

{diffsynth_engine-0.4.3.dev6 → diffsynth_engine-0.4.3.dev7}/diffsynth_engine/__init__.py

@@ -4,6 +4,10 @@ from .configs import (
     FluxPipelineConfig,
     WanPipelineConfig,
     QwenImagePipelineConfig,
+    SDStateDicts,
+    SDXLStateDicts,
+    FluxStateDicts,
+    QwenImageStateDicts,
     ControlNetParams,
     ControlType,
 )
@@ -38,6 +42,10 @@ __all__ = [
     "SDXLPipelineConfig",
     "FluxPipelineConfig",
     "WanPipelineConfig",
+    "SDStateDicts",
+    "SDXLStateDicts",
+    "FluxStateDicts",
+    "QwenImageStateDicts",
     "FluxImagePipeline",
     "QwenImagePipelineConfig",
     "FluxControlNet",

{diffsynth_engine-0.4.3.dev6 → diffsynth_engine-0.4.3.dev7}/diffsynth_engine/pipelines/sd_image.py

@@ -172,31 +172,40 @@ class SDImagePipeline(BasePipeline):
         else:
             config = model_path_or_config
 
-        logger.info(f"loading state dict from {config.model_path} ...")
-        unet_state_dict = cls.load_model_checkpoint(config.model_path, device="cpu", dtype=config.model_dtype)
-
-        if config.vae_path is not None:
-            logger.info(f"loading state dict from {config.vae_path} ...")
-            vae_state_dict = cls.load_model_checkpoint(config.vae_path, device="cpu", dtype=config.vae_dtype)
-        else:
-            vae_state_dict = unet_state_dict
-
-        if config.clip_path is not None:
-            logger.info(f"loading state dict from {config.clip_path} ...")
-            clip_state_dict = cls.load_model_checkpoint(config.clip_path, device="cpu", dtype=config.clip_dtype)
-        else:
-            clip_state_dict = unet_state_dict
+        return cls.from_state_dict(SDStateDicts(), config)
 
+    @classmethod
+    def from_state_dict(cls, state_dicts: SDStateDicts, config: SDPipelineConfig) -> "SDImagePipeline":
+        if state_dicts.model is None:
+            if config.model_path is None:
+                raise ValueError("`model_path` cannot be empty")
+            logger.info(f"loading state dict from {config.model_path} ...")
+            state_dicts.model = cls.load_model_checkpoint(config.model_path, device="cpu", dtype=config.model_dtype)
+        
+        if state_dicts.vae is None:
+            if config.vae_path is None:
+                state_dicts.vae = state_dicts.model
+            else:
+                logger.info(f"loading state dict from {config.vae_path} ...")
+                state_dicts.vae = cls.load_model_checkpoint(config.vae_path, device="cpu", dtype=config.vae_dtype)
+        
+        if state_dicts.clip is None:
+            if config.clip_path is None:
+                state_dicts.clip = state_dicts.model
+            else:
+                logger.info(f"loading state dict from {config.clip_path} ...")
+                state_dicts.clip = cls.load_model_checkpoint(config.clip_path, device="cpu", dtype=config.clip_dtype)
+        
         init_device = "cpu" if config.offload_mode is not None else config.device
         tokenizer = CLIPTokenizer.from_pretrained(SDXL_TOKENIZER_CONF_PATH)
         with LoRAContext():
-            text_encoder = SDTextEncoder.from_state_dict(clip_state_dict, device=init_device, dtype=config.clip_dtype)
-            unet = SDUNet.from_state_dict(unet_state_dict, device=init_device, dtype=config.model_dtype)
+            text_encoder = SDTextEncoder.from_state_dict(state_dicts.clip, device=init_device, dtype=config.clip_dtype)
+            unet = SDUNet.from_state_dict(state_dicts.model, device=init_device, dtype=config.model_dtype)
         vae_decoder = SDVAEDecoder.from_state_dict(
-            vae_state_dict, device=init_device, dtype=config.vae_dtype, attn_impl="sdpa"
+            state_dicts.vae, device=init_device, dtype=config.vae_dtype, attn_impl="sdpa"
         )
         vae_encoder = SDVAEEncoder.from_state_dict(
-            vae_state_dict, device=init_device, dtype=config.vae_dtype, attn_impl="sdpa"
+            state_dicts.vae, device=init_device, dtype=config.vae_dtype, attn_impl="sdpa"
         )
 
         pipe = cls(
@@ -213,10 +222,6 @@ class SDImagePipeline(BasePipeline):
             pipe.enable_cpu_offload(config.offload_mode)
         return pipe
 
-    @classmethod
-    def from_state_dict(cls, state_dicts: SDStateDicts, pipeline_config: SDPipelineConfig) -> "SDImagePipeline":
-        raise NotImplementedError()
-
     def denoising_model(self):
         return self.unet
 

{diffsynth_engine-0.4.3.dev6 → diffsynth_engine-0.4.3.dev7}/diffsynth_engine/pipelines/sdxl_image.py

@@ -150,43 +150,53 @@ class SDXLImagePipeline(BasePipeline):
         else:
             config = model_path_or_config
 
-        logger.info(f"loading state dict from {config.model_path} ...")
-        unet_state_dict = cls.load_model_checkpoint(config.model_path, device="cpu", dtype=config.model_dtype)
-
-        if config.vae_path is not None:
-            logger.info(f"loading state dict from {config.vae_path} ...")
-            vae_state_dict = cls.load_model_checkpoint(config.vae_path, device="cpu", dtype=config.vae_dtype)
-        else:
-            vae_state_dict = unet_state_dict
-
-        if config.clip_l_path is not None:
-            logger.info(f"loading state dict from {config.clip_l_path} ...")
-            clip_l_state_dict = cls.load_model_checkpoint(config.clip_l_path, device="cpu", dtype=config.clip_l_dtype)
-        else:
-            clip_l_state_dict = unet_state_dict
-
-        if config.clip_g_path is not None:
-            logger.info(f"loading state dict from {config.clip_g_path} ...")
-            clip_g_state_dict = cls.load_model_checkpoint(config.clip_g_path, device="cpu", dtype=config.clip_g_dtype)
-        else:
-            clip_g_state_dict = unet_state_dict
+        return cls.from_state_dict(SDXLStateDicts(), config)
 
+    @classmethod
+    def from_state_dict(cls, state_dicts: SDXLStateDicts, config: SDXLPipelineConfig) -> "SDXLImagePipeline":
+        if state_dicts.model is None:
+            if config.model_path is None:
+                raise ValueError("`model_path` cannot be empty")
+            logger.info(f"loading state dict from {config.model_path} ...")
+            state_dicts.model = cls.load_model_checkpoint(config.model_path, device="cpu", dtype=config.model_dtype)
+            
+        if state_dicts.vae is None:
+            if config.vae_path is None:
+                state_dicts.vae = state_dicts.model
+            else:
+                logger.info(f"loading state dict from {config.vae_path} ...")
+                state_dicts.vae = cls.load_model_checkpoint(config.vae_path, device="cpu", dtype=config.vae_dtype)
+        
+        if state_dicts.clip_l is None:
+            if config.clip_l_path is None:
+                state_dicts.clip_l = state_dicts.model
+            else:
+                logger.info(f"loading state dict from {config.clip_l_path} ...")
+                state_dicts.clip_l = cls.load_model_checkpoint(config.clip_l_path, device="cpu", dtype=config.clip_l_dtype)
+                
+        if state_dicts.clip_g is None:
+            if config.clip_g_path is None:
+                state_dicts.clip_g = state_dicts.model
+            else:
+                logger.info(f"loading state dict from {config.clip_g_path} ...")
+                state_dicts.clip_g = cls.load_model_checkpoint(config.clip_g_path, device="cpu", dtype=config.clip_g_dtype)
+                
         init_device = "cpu" if config.offload_mode else config.device
         tokenizer = CLIPTokenizer.from_pretrained(SDXL_TOKENIZER_CONF_PATH)
         tokenizer_2 = CLIPTokenizer.from_pretrained(SDXL_TOKENIZER_2_CONF_PATH)
         with LoRAContext():
             text_encoder = SDXLTextEncoder.from_state_dict(
-                clip_l_state_dict, device=init_device, dtype=config.clip_l_dtype
+                state_dicts.clip_l, device=init_device, dtype=config.clip_l_dtype
             )
             text_encoder_2 = SDXLTextEncoder2.from_state_dict(
-                clip_g_state_dict, device=init_device, dtype=config.clip_g_dtype
+                state_dicts.clip_g, device=init_device, dtype=config.clip_g_dtype
             )
-            unet = SDXLUNet.from_state_dict(unet_state_dict, device=init_device, dtype=config.model_dtype)
+            unet = SDXLUNet.from_state_dict(state_dicts.model, device=init_device, dtype=config.model_dtype)
         vae_decoder = SDXLVAEDecoder.from_state_dict(
-            vae_state_dict, device=init_device, dtype=config.vae_dtype, attn_impl="sdpa"
+            state_dicts.vae, device=init_device, dtype=config.vae_dtype, attn_impl="sdpa"
         )
         vae_encoder = SDXLVAEEncoder.from_state_dict(
-            vae_state_dict, device=init_device, dtype=config.vae_dtype, attn_impl="sdpa"
+            state_dicts.vae, device=init_device, dtype=config.vae_dtype, attn_impl="sdpa"
         )
 
         pipe = cls(
@@ -205,10 +215,6 @@ class SDXLImagePipeline(BasePipeline):
             pipe.enable_cpu_offload(config.offload_mode)
         return pipe
 
-    @classmethod
-    def from_state_dict(cls, state_dicts: SDXLStateDicts, pipeline_config: SDXLPipelineConfig) -> "SDXLImagePipeline":
-        raise NotImplementedError()
-
     def denoising_model(self):
         return self.unet
 

{diffsynth_engine-0.4.3.dev6 → diffsynth_engine-0.4.3.dev7}/diffsynth_engine.egg-info/PKG-INFO

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