inference with regular ModelPatcher works correctly

torchvulkan_loader.py

@@ -1,4 +1,20 @@
 from .common import *
+import comfy
+import comfy.model_patcher
+import comfy.sd
+
+from .vulkan_monkeypatch import *
+
+def load_ModelPatcher(device, mp):
+	assert isinstance(mp.model, torch.nn.Module)
+	
+	# move the model to the device
+	mp.model.to(device)
+	
+	# overwrite load device
+	mp.load_device = device
+	
+	# this will at least function, I hope
 
 class TorchVulkanLoader:
 	def __init__(self):
@@ -9,7 +25,7 @@ class TorchVulkanLoader:
 		devices = [i for i in range(torchvulkan.device_count())]
 		return {
 			"required": {
-				"device": (devices, {"default": "default ( local-task:// )"})
+				"device": (devices, {"default": 0})
 			},
 			"optional": {
 				"model": ("MODEL",),
@@ -25,10 +41,26 @@ class TorchVulkanLoader:
 	CATEGORY = "Backends"
 	
 	def load_to_torchvulkan(self, device, model = None, clip = None, vae = None):
-		raise NotImplementedError
+		# only do monkeypatching if node is invoked or whatever
+		comfy.model_management.get_torch_device = get_torch_device_patched
+		comfy.model_management.get_free_memory = get_free_memory_patched
+		
 		device = torch.device(f"vulkan:{device}")
-		if model: model.to(device)
-		if clip: clip.to(device)
-		if vae: vae.to(device)
+		if not model is None:
+			if isinstance(model, comfy.model_patcher.ModelPatcher):
+				load_ModelPatcher(device, model)
+			else:
+				raise NotImplementedError
+		
+		if not clip is None:
+			if isinstance(clip, comfy.sd.CLIP):
+				load_ModelPatcher(device, clip.patcher)
+			else:
+				raise NotImplementedError
+		if not vae is None:
+			if isinstance(vae, comfy.sd.VAE):
+				load_ModelPatcher(device, vae.patcher)
+			else:
+				raise NotImplementedError
 		return model, clip, vae
 		

vulkan_monkeypatch.py

@@ -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