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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: diffsynth_engine
-Version: 0.4.3.dev6
+Version: 0.4.3.dev8
 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.dev8}/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.dev8}/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.dev8}/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.dev8}/diffsynth_engine/tools/flux_inpainting_tool.py

@@ -1,5 +1,12 @@
-from diffsynth_engine import fetch_model, FluxPipelineConfig, FluxControlNet, ControlNetParams, FluxImagePipeline
-from typing import List, Tuple, Optional, Callable
+from diffsynth_engine import (
+    fetch_model,
+    FluxPipelineConfig,
+    FluxControlNet,
+    ControlNetParams,
+    FluxImagePipeline,
+    FluxStateDicts,
+)
+from typing import List, Tuple, Optional, Callable, Dict
 from PIL import Image
 import torch
 
@@ -7,6 +14,15 @@ import torch
 class FluxInpaintingTool:
     def __init__(
         self,
+        flux_pipe: FluxImagePipeline,
+        controlnet: FluxControlNet
+    ):
+        self.pipe = flux_pipe
+        self.controlnet = controlnet
+
+    @classmethod
+    def from_pretrained(
+        cls,
         flux_model_path: str,
         device: str = "cuda:0",
         dtype: torch.dtype = torch.bfloat16,
@@ -18,14 +34,35 @@ class FluxInpaintingTool:
             device=device,
             offload_mode=offload_mode,
         )
-        self.pipe = FluxImagePipeline.from_pretrained(config)
-        self.controlnet = FluxControlNet.from_pretrained(
+        flux_pipe = FluxImagePipeline.from_pretrained(config)
+        controlnet = FluxControlNet.from_pretrained(
             fetch_model(
                 "alimama-creative/FLUX.1-dev-Controlnet-Inpainting-Beta", path="diffusion_pytorch_model.safetensors"
             ),
             device=device,
             dtype=torch.bfloat16,
         )
+        return cls(flux_pipe, controlnet)
+
+    @classmethod
+    def from_state_dict(
+        cls,
+        flux_state_dicts: FluxStateDicts,
+        controlnet_state_dict: Dict[str, torch.Tensor],
+        device: str = "cuda:0",
+        dtype: torch.dtype = torch.bfloat16,
+        offload_mode: Optional[str] = None,
+    ):
+        config = FluxPipelineConfig(
+            model_path="",
+            model_dtype=dtype,
+            device=device,
+            offload_mode=offload_mode,
+        )
+        flux_pipe = FluxImagePipeline.from_state_dict(flux_state_dicts, config)
+        controlnet = FluxControlNet.from_state_dict(controlnet_state_dict, device, dtype)
+        return cls(flux_pipe, controlnet)
+
 
     def load_loras(self, lora_list: List[Tuple[str, float]], fused: bool = True, save_original_weight: bool = False):
         self.pipe.load_loras(lora_list, fused, save_original_weight)

{diffsynth_engine-0.4.3.dev6 → diffsynth_engine-0.4.3.dev8}/diffsynth_engine/tools/flux_outpainting_tool.py

@@ -1,5 +1,12 @@
-from diffsynth_engine import fetch_model, FluxPipelineConfig, FluxControlNet, ControlNetParams, FluxImagePipeline
-from typing import List, Tuple, Optional, Callable
+from diffsynth_engine import (
+    fetch_model,
+    FluxPipelineConfig,
+    FluxControlNet,
+    ControlNetParams,
+    FluxImagePipeline,
+    FluxStateDicts
+)
+from typing import List, Tuple, Optional, Callable, Dict
 from PIL import Image
 import torch
 
@@ -7,6 +14,15 @@ import torch
 class FluxOutpaintingTool:
     def __init__(
         self,
+        flux_pipe: FluxImagePipeline,
+        controlnet: FluxControlNet,
+    ):
+        self.pipe = flux_pipe
+        self.controlnet = controlnet
+
+    @classmethod
+    def from_pretrained(
+        cls,
         flux_model_path: str,
         device: str = "cuda:0",
         dtype: torch.dtype = torch.bfloat16,
@@ -18,14 +34,35 @@ class FluxOutpaintingTool:
             device=device,
             offload_mode=offload_mode,
         )
-        self.pipe = FluxImagePipeline.from_pretrained(config)
-        self.controlnet = FluxControlNet.from_pretrained(
+        flux_pipe = FluxImagePipeline.from_pretrained(config)
+        controlnet = FluxControlNet.from_pretrained(
             fetch_model(
-                "alimama-creative/FLUX.1-dev-Controlnet-Inpainting-Beta", path="diffusion_pytorch_model.safetensors"
+                "alimama-creative/FLUX.1-dev-Controlnet-Inpainting-Beta",
+                path="diffusion_pytorch_model.safetensors"
             ),
             device=device,
-            dtype=torch.bfloat16,
+            dtype=torch.bfloat16
+        )
+        return cls(flux_pipe, controlnet)
+
+    @classmethod
+    def from_state_dict(
+        cls,
+        flux_state_dicts: FluxStateDicts,
+        controlnet_state_dict: Dict[str, torch.Tensor],
+        device: str = "cuda:0",
+        dtype: torch.dtype = torch.bfloat16,
+        offload_mode: Optional[str] = None,
+    ):
+        config = FluxPipelineConfig(
+            model_path="",
+            model_dtype=dtype,
+            device=device,
+            offload_mode=offload_mode,
         )
+        flux_pipe = FluxImagePipeline.from_state_dict(flux_state_dicts, config)
+        controlnet = FluxControlNet.from_state_dict(controlnet_state_dict, device, dtype)
+        return cls(flux_pipe, controlnet)
 
     def load_loras(self, lora_list: List[Tuple[str, float]], fused: bool = True, save_original_weight: bool = False):
         self.pipe.load_loras(lora_list, fused, save_original_weight)

{diffsynth_engine-0.4.3.dev6 → diffsynth_engine-0.4.3.dev8}/diffsynth_engine/tools/flux_reference_tool.py

@@ -5,8 +5,9 @@ from diffsynth_engine import (
     FluxIPAdapter,
     FluxRedux,
     fetch_model,
+    FluxStateDicts
 )
-from typing import List, Tuple, Optional
+from typing import List, Tuple, Optional, Dict
 from PIL import Image
 import torch
 
@@ -18,8 +19,17 @@ class FluxReduxRefTool:
 
     def __init__(
         self,
+        flux_pipe: FluxImagePipeline,
+        redux: FluxRedux,
+    ):
+        self.pipe = flux_pipe
+        self.pipe.load_redux(redux)
+
+    @classmethod
+    def from_pretrained(
+        cls,
         flux_model_path: str,
-        load_text_encoder=True,
+        load_text_encoder: bool = True,
         device: str = "cuda:0",
         dtype: torch.dtype = torch.bfloat16,
         offload_mode: Optional[str] = None,
@@ -31,10 +41,31 @@ class FluxReduxRefTool:
             device=device,
             offload_mode=offload_mode,
         )
-        self.pipe: FluxImagePipeline = FluxImagePipeline.from_pretrained(config)
+        flux_pipe = FluxImagePipeline.from_pretrained(config)
         redux_model_path = fetch_model("muse/flux1-redux-dev", path="flux1-redux-dev.safetensors", revision="v1")
-        flux_redux = FluxRedux.from_pretrained(redux_model_path, device=device)
-        self.pipe.load_redux(flux_redux)
+        redux = FluxRedux.from_pretrained(redux_model_path, device=device)
+        return cls(flux_pipe, redux)
+
+    @classmethod
+    def from_state_dict(
+        cls,
+        flux_state_dicts: FluxStateDicts,
+        redux_state_dict: Dict[str, torch.Tensor],
+        load_text_encoder: bool = True,
+        device: str = "cuda:0",
+        dtype: torch.dtype = torch.bfloat16,
+        offload_mode: Optional[str] = None,
+    ):
+        config = FluxPipelineConfig(
+            model_path="",
+            model_dtype=dtype,
+            load_text_encoder=load_text_encoder,
+            device=device,
+            offload_mode=offload_mode,
+        )
+        flux_pipe = FluxImagePipeline.from_state_dict(flux_state_dicts, config)
+        redux = FluxRedux.from_state_dict(redux_state_dict, device=device, dtype=dtype)
+        return cls(flux_pipe, redux)
 
     def load_loras(self, lora_list: List[Tuple[str, float]], fused: bool = True, save_original_weight: bool = False):
         self.pipe.load_loras(lora_list, fused, save_original_weight)

{diffsynth_engine-0.4.3.dev6 → diffsynth_engine-0.4.3.dev8}/diffsynth_engine/tools/flux_replace_tool.py

@@ -5,8 +5,9 @@ from diffsynth_engine import (
     FluxImagePipeline,
     FluxRedux,
     fetch_model,
+    FluxStateDicts
 )
-from typing import List, Tuple, Optional, Callable
+from typing import List, Tuple, Optional, Callable, Dict
 from PIL import Image
 import torch
 
@@ -19,8 +20,19 @@ class FluxReplaceByControlTool:
 
     def __init__(
         self,
+        flux_pipe: FluxImagePipeline,
+        redux: FluxRedux,
+        controlnet: FluxControlNet,
+    ):
+        self.pipe = flux_pipe
+        self.pipe.load_redux(redux)
+        self.controlnet = controlnet
+
+    @classmethod
+    def from_pretrained(
+        cls,
         flux_model_path: str,
-        load_text_encoder=True,
+        load_text_encoder: bool = True,
         device: str = "cuda:0",
         dtype: torch.dtype = torch.bfloat16,
         offload_mode: Optional[str] = None,
@@ -32,17 +44,42 @@ class FluxReplaceByControlTool:
             device=device,
             offload_mode=offload_mode,
         )
-        self.pipe: FluxImagePipeline = FluxImagePipeline.from_pretrained(config)
+        flux_pipe = FluxImagePipeline.from_pretrained(config)
         redux_model_path = fetch_model("muse/flux1-redux-dev", path="flux1-redux-dev.safetensors", revision="v1")
-        flux_redux = FluxRedux.from_pretrained(redux_model_path, device=device)
-        self.pipe.load_redux(flux_redux)
-        self.controlnet = FluxControlNet.from_pretrained(
+        redux = FluxRedux.from_pretrained(redux_model_path, device=device, dtype=dtype)
+        controlnet = FluxControlNet.from_pretrained(
             fetch_model(
-                "alimama-creative/FLUX.1-dev-Controlnet-Inpainting-Beta", path="diffusion_pytorch_model.safetensors"
+                "alimama-creative/FLUX.1-dev-Controlnet-Inpainting-Beta",
+                path="diffusion_pytorch_model.safetensors"
             ),
             device=device,
             dtype=torch.bfloat16,
         )
+        return cls(flux_pipe, redux, controlnet)
+
+    @classmethod
+    def from_state_dict(
+        cls,
+        flux_state_dicts: FluxStateDicts,
+        redux_state_dict: Dict[str, torch.Tensor],
+        controlnet_state_dict: Dict[str, torch.Tensor],
+        load_text_encoder: bool = True,
+        device: str = "cuda:0",
+        dtype: torch.dtype = torch.bfloat16,
+        offload_mode: Optional[str] = None,
+    ):
+        config = FluxPipelineConfig(
+            model_path="",
+            model_dtype=dtype,
+            load_text_encoder=load_text_encoder,
+            device=device,
+            offload_mode=offload_mode,
+        )
+        flux_pipe = FluxImagePipeline.from_state_dict(flux_state_dicts, config)
+        redux = FluxRedux.from_state_dict(redux_state_dict, device=device, dtype=dtype)
+        controlnet = FluxControlNet.from_state_dict(controlnet_state_dict, device=device, dtype=dtype)
+        return cls(flux_pipe, redux, controlnet)
+
 
     def load_loras(self, lora_list: List[Tuple[str, float]], fused: bool = True, save_original_weight: bool = False):
         self.pipe.load_loras(lora_list, fused, save_original_weight)

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