inference with regular ModelPatcher works correctly

2026-11-03 09:31:33 -08:00
parent 01e3eb6e36
commit 3a6f92f612
2 changed files with 92 additions and 26 deletions
@@ -1,28 +1,62 @@
+from .common import *
 import comfy
 import comfy.model_management

-# Patched version of comfy.model_management.get_torch_device that allows vulkan devices to be treated as default
 def get_torch_device_patched():
+	# just do 0 for now. need proof of concept
+	return torch.device(f"vulkan:{0}")
+
+
+
+def get_free_memory_patched(dev=None, torch_free_too=False):
 	global directml_enabled
-	global cpu_state
-	if directml_enabled:
-		global directml_device
-		return directml_device
-	if cpu_state == CPUState.MPS:
-		return torch.device("mps")
-	if cpu_state == CPUState.CPU:
-		return torch.device("cpu")
+	if dev is None:
+		dev = get_torch_device()
+		
+	# this is necessary to find the vulkan device at all
+	if "vulkan" in str(dev):
+		# no idea if torchvulkan exposes memory budget.
+		# will have to investigate later
+		mem_free_total = 1024 * 1024 * 1024
+		mem_free_torch = mem_free_total
+		directml_enabled = False
+	elif hasattr(dev, 'type') and (dev.type == 'cpu' or dev.type == 'mps'):
+		mem_free_total = comfy.system_memory.virtual_memory_available()
+		mem_free_torch = mem_free_total
 	else:
-		if True:
-			# for now, just see if this works c:
-			return torch.device(f"vulkan:0")
+		if directml_enabled:
+			mem_free_total = 1024 * 1024 * 1024 #TODO
+			mem_free_torch = mem_free_total
+		elif is_intel_xpu():
+			stats = torch.xpu.memory_stats(dev)
+			mem_active = stats['active_bytes.all.current']
+			mem_reserved = stats['reserved_bytes.all.current']
+			mem_free_xpu = torch.xpu.get_device_properties(dev).total_memory - mem_reserved
+			mem_free_torch = mem_reserved - mem_active
+			mem_free_total = mem_free_xpu + mem_free_torch
+		elif is_ascend_npu():
+			stats = torch.npu.memory_stats(dev)
+			mem_active = stats['active_bytes.all.current']
+			mem_reserved = stats['reserved_bytes.all.current']
+			mem_free_npu, _ = torch.npu.mem_get_info(dev)
+			mem_free_torch = mem_reserved - mem_active
+			mem_free_total = mem_free_npu + mem_free_torch
+		elif is_mlu():
+			stats = torch.mlu.memory_stats(dev)
+			mem_active = stats['active_bytes.all.current']
+			mem_reserved = stats['reserved_bytes.all.current']
+			mem_free_mlu, _ = torch.mlu.mem_get_info(dev)
+			mem_free_torch = mem_reserved - mem_active
+			mem_free_total = mem_free_mlu + mem_free_torch
 		else:
-			if is_intel_xpu():
-				return torch.device("xpu", torch.xpu.current_device())
-			elif is_ascend_npu():
-				return torch.device("npu", torch.npu.current_device())
-			elif is_mlu():
-				return torch.device("mlu", torch.mlu.current_device())
-			else:
-				return torch.device(torch.cuda.current_device())
-comfy.model_management.get_torch_device = get_torch_device_patched
+			stats = torch.cuda.memory_stats(dev)
+			mem_active = stats['active_bytes.all.current']
+			mem_reserved = stats['reserved_bytes.all.current']
+			mem_free_cuda, _ = torch.cuda.mem_get_info(dev)
+			mem_free_torch = mem_reserved - mem_active
+			mem_free_total = mem_free_cuda + mem_free_torch
+
+	if torch_free_too:
+		return (mem_free_total, mem_free_torch)
+	else:
+		return mem_free_total