diffsynth-engine 0.3.6.dev7__tar.gz → 0.3.6.dev9__tar.gz

{diffsynth_engine-0.3.6.dev7 → diffsynth_engine-0.3.6.dev9}/PKG-INFO

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

{diffsynth_engine-0.3.6.dev7 → diffsynth_engine-0.3.6.dev9}/diffsynth_engine/models/flux/flux_vae.py

@@ -53,7 +53,8 @@ class FluxVAEEncoder(VAEEncoder):
     def from_state_dict(cls, state_dict: Dict[str, torch.Tensor], device: str, dtype: torch.dtype):
         with no_init_weights():
             model = torch.nn.utils.skip_init(cls, device=device, dtype=dtype)
-        model.load_state_dict(state_dict)
+        model.load_state_dict(state_dict, assign=True)
+        model.to(device=device, dtype=dtype, non_blocking=True)
         return model
 
 
@@ -74,5 +75,6 @@ class FluxVAEDecoder(VAEDecoder):
     def from_state_dict(cls, state_dict: Dict[str, torch.Tensor], device: str, dtype: torch.dtype):
         with no_init_weights():
             model = torch.nn.utils.skip_init(cls, device=device, dtype=dtype)
-        model.load_state_dict(state_dict)
+        model.load_state_dict(state_dict, assign=True)
+        model.to(device=device, dtype=dtype, non_blocking=True)
         return model

{diffsynth_engine-0.3.6.dev7 → diffsynth_engine-0.3.6.dev9}/diffsynth_engine/pipelines/base.py

@@ -43,6 +43,7 @@ class BasePipeline:
         self.dtype = dtype
         self.offload_mode = None
         self.model_names = []
+        self._models_offload_params = {}
 
     @classmethod
     def from_pretrained(
@@ -288,6 +289,10 @@ class BasePipeline:
             model = getattr(self, model_name)
             if model is not None:
                 model.to("cpu")
+                self._models_offload_params[model_name] = {}
+                for name, param in model.named_parameters(recurse=True):
+                    param.data = param.data.pin_memory()
+                    self._models_offload_params[model_name][name] = param.data
         self.offload_mode = "cpu_offload"
 
     def _enable_sequential_cpu_offload(self):
@@ -321,12 +326,14 @@ class BasePipeline:
         for model_name in self.model_names:
             if model_name not in load_model_names:
                 model = getattr(self, model_name)
-                if model is not None and (p := next(model.parameters(), None)) is not None and p.device != "cpu":
-                    model.to("cpu")
+                if model is not None and (p := next(model.parameters(), None)) is not None and p.device != torch.device("cpu"):
+                    param_cache = self._models_offload_params[model_name]
+                    for name, param in model.named_parameters(recurse=True):
+                        param.data = param_cache[name]
         # load the needed models to device
         for model_name in load_model_names:
             model = getattr(self, model_name)
-            if model is not None and (p := next(model.parameters(), None)) is not None and p.device != self.device:
+            if model is not None and (p := next(model.parameters(), None)) is not None and p.device != torch.device(self.device):
                 model.to(self.device)
         # fresh the cuda cache
         empty_cache()

{diffsynth_engine-0.3.6.dev7 → diffsynth_engine-0.3.6.dev9}/diffsynth_engine/pipelines/flux_image.py

@@ -604,7 +604,8 @@ class FluxImagePipeline(BasePipeline):
             device=device,
             dtype=model_config.dit_dtype,
         )
-        pipe.enable_cpu_offload(offload_mode)
+        if offload_mode is not None:
+            pipe.enable_cpu_offload(offload_mode)
         if model_config.dit_dtype == torch.float8_e4m3fn:
             pipe.dtype = torch.bfloat16  # running dtype
             pipe.enable_fp8_autocast(

{diffsynth_engine-0.3.6.dev7 → diffsynth_engine-0.3.6.dev9}/diffsynth_engine/utils/fp8_linear.py

@@ -2,6 +2,7 @@ import torch
 import torch.nn as nn
 import torch.nn.functional as F
 from contextlib import contextmanager
+from diffsynth_engine.utils.platform import DTYPE_FP8
 
 
 def enable_fp8_autocast(module: nn.Module, compute_dtype: torch.dtype = torch.bfloat16, use_fp8_linear: bool = False):
@@ -51,7 +52,7 @@ def enable_fp8_linear(module: nn.Module):
 def _enable_fp8_linear(module: nn.Module):
     if isinstance(module, nn.Linear) and torch.is_floating_point(module.weight.data):
         # avoid conversion for int weights like GGUF
-        module.weight.data = module.weight.data.to(torch.float8_e4m3fn)
+        module.weight.data = module.weight.data.to(DTYPE_FP8)
     for submodule in module.children():
         _enable_fp8_linear(submodule)
 
@@ -71,8 +72,16 @@ def fp8_inference(enabled=True):
     ) -> torch.Tensor:
         device = input.device
         origin_dtype = input.dtype
-        input = input.to(torch.float8_e4m3fn)
-        weight = weight.to(torch.float8_e4m3fn)
+        scale_a = 1.0
+        # For float8_e4m3fnuz, the maximum representable value is half of that of e4m3fn.
+        # To avoid overflow and ensure numerical compatibility during FP8 computation,
+        # we scale down the input by 2.0 in advance.
+        # This scaling will be compensated later during the final result scaling.
+        if DTYPE_FP8 == torch.float8_e4m3fnuz:
+            scale_a = 2.0
+            input = input / scale_a
+        input = input.to(DTYPE_FP8)
+        weight = weight.to(DTYPE_FP8)
 
         if len(input.shape) > 2:
             origin_shape = input.shape
@@ -80,7 +89,7 @@ def fp8_inference(enabled=True):
             result = torch._scaled_mm(
                 input,
                 weight.T,
-                scale_a=torch.tensor(1.0).to(device=device),
+                scale_a=torch.tensor(scale_a).to(device=device),
                 scale_b=torch.tensor(1.0).to(device=device),
                 bias=bias,
                 out_dtype=origin_dtype,
@@ -91,7 +100,7 @@ def fp8_inference(enabled=True):
             result = torch._scaled_mm(
                 input,
                 weight.T,
-                scale_a=torch.tensor(1.0).to(device=device),
+                scale_a=torch.tensor(scale_a).to(device=device),
                 scale_b=torch.tensor(1.0).to(device=device),
                 bias=bias,
                 out_dtype=origin_dtype,

{diffsynth_engine-0.3.6.dev7 → diffsynth_engine-0.3.6.dev9}/diffsynth_engine/utils/offload.py

@@ -18,7 +18,7 @@ def add_cpu_offload_hook(module: nn.Module, device: str = "cuda", recurse: bool
     def _forward_pre_hook(module: nn.Module, input):
         offload_params = {}
         for name, param in module.named_parameters(recurse=recurse):
-            offload_params[name] = param.data
+            offload_params[name] = param.data.pin_memory()
             param.data = param.data.to(device=device)
         setattr(module, "_offload_params", offload_params)
         return tuple(x.to(device=device) if isinstance(x, torch.Tensor) else x for x in input)

diffsynth_engine-0.3.6.dev9/diffsynth_engine/utils/platform.py

@@ -0,0 +1,20 @@
+# cross-platform definitions and utilities
+import torch
+import gc
+
+
+# data type
+# AMD only supports float8_e4m3fnuz
+# https://onnx.ai/onnx/technical/float8.html
+if torch.version.hip and "gfx94" in torch.cuda.get_device_properties(0).gcnArchName:
+    DTYPE_FP8 = torch.float8_e4m3fnuz
+else:
+    DTYPE_FP8 = torch.float8_e4m3fn
+
+
+def empty_cache():
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
+    if torch.mps.is_available():
+        torch.mps.empty_cache()
+    gc.collect()

{diffsynth_engine-0.3.6.dev7 → diffsynth_engine-0.3.6.dev9}/diffsynth_engine.egg-info/PKG-INFO

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

diffsynth_engine-0.3.6.dev7/diffsynth_engine/utils/platform.py

@@ -1,12 +0,0 @@
-import torch
-import gc
-
-# 存放跨平台的工具类
-
-
-def empty_cache():
-    if torch.cuda.is_available():
-        torch.cuda.empty_cache()
-    if torch.mps.is_available():
-        torch.mps.empty_cache()
-    gc.collect()