Dev/flux tool from state dict (#153)

* FluxOutpaintingTool & FluxReplaceByControlTool implement from_state_dict

* flux tool add from_pretrained & from_state_dict

* fix redux model path

Author: qzzz95

diffsynth_engine/tools/flux_inpainting_tool.py

@@ -1,12 +1,28 @@
-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
 
 
 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 @@ def __init__(
             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/tools/flux_outpainting_tool.py

@@ -1,12 +1,28 @@
-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
 
 
 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 @@ def __init__(
             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/tools/flux_reference_tool.py

@@ -5,8 +5,9 @@
     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 @@ def __init__(
             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/tools/flux_replace_tool.py

@@ -5,8 +5,9 @@
     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 @@ def __init__(
             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)