From d2488d73f63f8c1e4a5df0f6c2ffdd7e2eb4b22b Mon Sep 17 00:00:00 2001
From: City <125218114+city96@users.noreply.github.com>
Date: Sun, 28 Jan 2024 17:47:02 +0100
Subject: [PATCH] Quick FP8 fix.

---
 DiT/loader.py    | 14 +++++++++++---
 PixArt/loader.py | 12 ++++++++++--
 2 files changed, 21 insertions(+), 5 deletions(-)

diff --git a/DiT/loader.py b/DiT/loader.py
index daad506..83db9b7 100644
--- a/DiT/loader.py
+++ b/DiT/loader.py
@@ -26,6 +26,14 @@ def load_dit(model_path, model_conf):
 	state_dict = state_dict.get("model", state_dict)
 	parameters = comfy.utils.calculate_parameters(state_dict)
 	unet_dtype = model_management.unet_dtype(model_params=parameters)
+	load_device = comfy.model_management.get_torch_device()
+	offload_device = comfy.model_management.unet_offload_device()
+
+	# ignore fp8/etc and use directly for now
+	manual_cast_dtype = model_management.unet_manual_cast(unet_dtype, load_device)
+	if manual_cast_dtype:
+		print(f"DiT: falling back to {manual_cast_dtype}")
+		unet_dtype = manual_cast_dtype
 
 	model_conf["unet_config"]["num_classes"] = state_dict["y_embedder.embedding_table.weight"].shape[0] - 1 # adj. for empty
 
@@ -46,8 +54,8 @@ def load_dit(model_path, model_conf):
 
 	model_patcher = comfy.model_patcher.ModelPatcher(
 		model,
-		load_device=comfy.model_management.get_torch_device(),
-		offload_device=comfy.model_management.unet_offload_device(),
-		current_device="cpu"
+		load_device = load_device,
+		offload_device = offload_device,
+		current_device = "cpu",
 	)
 	return model_patcher
diff --git a/PixArt/loader.py b/PixArt/loader.py
index e84fe14..1d66fa3 100644
--- a/PixArt/loader.py
+++ b/PixArt/loader.py
@@ -38,6 +38,14 @@ def load_pixart(model_path, model_conf):
 
 	parameters = comfy.utils.calculate_parameters(state_dict)
 	unet_dtype = model_management.unet_dtype(model_params=parameters)
+	load_device = comfy.model_management.get_torch_device()
+	offload_device = comfy.model_management.unet_offload_device()
+
+	# ignore fp8/etc and use directly for now
+	manual_cast_dtype = model_management.unet_manual_cast(unet_dtype, load_device)
+	if manual_cast_dtype:
+		print(f"PixArt: falling back to {manual_cast_dtype}")
+		unet_dtype = manual_cast_dtype
 
 	model_conf = EXM_PixArt(model_conf) # convert to object
 	model = comfy.model_base.BaseModel(
@@ -64,8 +72,8 @@ def load_pixart(model_path, model_conf):
 
 	model_patcher = comfy.model_patcher.ModelPatcher(
 		model,
-		load_device    = comfy.model_management.get_torch_device(),
-		offload_device = comfy.model_management.unet_offload_device(),
+		load_device = load_device,
+		offload_device = offload_device,
 		current_device = "cpu",
 	)
 	return model_patcher