diffsynth-engine 0.4.4.dev5__py3-none-any.whl → 0.5.1.dev1__py3-none-any.whl

diffsynth_engine/utils/fp8_linear.py

@@ -72,39 +72,33 @@ def fp8_inference(enabled=True):
     ) -> torch.Tensor:
         device = input.device
         origin_dtype = input.dtype
-        scale_a = 1.0
+        origin_shape = input.shape
+        input = input.reshape(-1, origin_shape[-1])
+
+        x_max = torch.max(torch.abs(input), dim=-1, keepdim=True).values
+        fp8_max = 448.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
+            fp8_max = fp8_max / 2.0
+        scale_a = torch.clamp(x_max / fp8_max, min=1.0).float().to(device=device)
+        scale_b = torch.ones((weight.shape[0], 1)).float().to(device=device)
+        input = input / scale_a
         input = input.to(DTYPE_FP8)
         weight = weight.to(DTYPE_FP8)
 
-        if len(input.shape) > 2:
-            origin_shape = input.shape
-            input = input.reshape(-1, origin_shape[-1])
-            result = torch._scaled_mm(
-                input,
-                weight.T,
-                scale_a=torch.tensor(scale_a).to(device=device),
-                scale_b=torch.tensor(1.0).to(device=device),
-                bias=bias,
-                out_dtype=origin_dtype,
-            )
-            new_shape = origin_shape[:-1] + result.shape[-1:]
-            result = result.reshape(new_shape)
-        else:
-            result = torch._scaled_mm(
-                input,
-                weight.T,
-                scale_a=torch.tensor(scale_a).to(device=device),
-                scale_b=torch.tensor(1.0).to(device=device),
-                bias=bias,
-                out_dtype=origin_dtype,
-            )
+        result = torch._scaled_mm(
+            input,
+            weight.T,
+            scale_a=scale_a,
+            scale_b=scale_b.T,
+            bias=bias,
+            out_dtype=origin_dtype,
+        )
+        new_shape = origin_shape[:-1] + result.shape[-1:]
+        result = result.reshape(new_shape)
         return result
 
     F.linear = fp8_linear

{diffsynth_engine-0.4.4.dev5.dist-info → diffsynth_engine-0.5.1.dev1.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: diffsynth_engine
-Version: 0.4.4.dev5
+Version: 0.5.1.dev1
 Author: MuseAI x ModelScope
 Classifier: Programming Language :: Python :: 3
 Classifier: Operating System :: OS Independent

{diffsynth_engine-0.4.4.dev5.dist-info → diffsynth_engine-0.5.1.dev1.dist-info}/RECORD

@@ -166,7 +166,7 @@ diffsynth_engine/utils/constants.py,sha256=csKPtXA9YFoRPnhUBmnw3zf1KlTV963a43Nu4
 diffsynth_engine/utils/download.py,sha256=8O56zQr2taY4BnCBEKVk3YGorTqTTCqTblf7tI9otXA,6721
 diffsynth_engine/utils/env.py,sha256=43x-kBjt5zI2cwZ9G4BOeTbedi2k6TuBzHGOBeFbFvU,280
 diffsynth_engine/utils/flag.py,sha256=6zQLnoEaU69pBEyhavCgydQfP0khw5ppCU7sue4yRqg,1370
-diffsynth_engine/utils/fp8_linear.py,sha256=NosnWMoAr_IpFcLn-OYbAx-vXySphjxutDZqmXLNjJI,4064
+diffsynth_engine/utils/fp8_linear.py,sha256=k34YFWo2dc3t8aKjHaCW9CbQMOTqXxaDHk8aw8aKif4,3857
 diffsynth_engine/utils/gguf.py,sha256=ZWvw46V4g4uVyAR_oCq-4K5nPdKVrYk3u47uXMgA9lU,14092
 diffsynth_engine/utils/image.py,sha256=xZ_bEU-DdoSwMPG7jpP1daAro2tsy9ddeXKbXqXaeC0,9335
 diffsynth_engine/utils/loader.py,sha256=Z5v1WNDWFY0OrVubB70j5VU3zeaAfEK_j8c1KrGI4yM,1240
@@ -181,8 +181,8 @@ diffsynth_engine/utils/video.py,sha256=Ne0rd2lb59UT1q5EotpjlY7OT8F9oTCFDyo1ST77u
 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.4.4.dev5.dist-info/licenses/LICENSE,sha256=x7aBqQuVI0IYnftgoTPI_A0I_rjdjPPQkjnU6N2nikM,11346
-diffsynth_engine-0.4.4.dev5.dist-info/METADATA,sha256=9da4omhF3UQxWXQbIf47qR8KqeGL7vhp-CvlEnhYpUA,1117
-diffsynth_engine-0.4.4.dev5.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-diffsynth_engine-0.4.4.dev5.dist-info/top_level.txt,sha256=6zgbiIzEHLbhgDKRyX0uBJOV3F6VnGGBRIQvSiYYn6w,17
-diffsynth_engine-0.4.4.dev5.dist-info/RECORD,,
+diffsynth_engine-0.5.1.dev1.dist-info/licenses/LICENSE,sha256=x7aBqQuVI0IYnftgoTPI_A0I_rjdjPPQkjnU6N2nikM,11346
+diffsynth_engine-0.5.1.dev1.dist-info/METADATA,sha256=myookCu04gVV_XbTrYo9J798-pmfmzpBernyCIl_5NQ,1117
+diffsynth_engine-0.5.1.dev1.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+diffsynth_engine-0.5.1.dev1.dist-info/top_level.txt,sha256=6zgbiIzEHLbhgDKRyX0uBJOV3F6VnGGBRIQvSiYYn6w,17
+diffsynth_engine-0.5.1.dev1.dist-info/RECORD,,