feat: Refactor DisTorch SafeTensor wrapper and expand coverage
This commit is contained in:
John Pollock
+164
-6
@@ -26,7 +26,7 @@ from .nodes import (
|
||||
WanVideoModelLoader, WanVideoModelLoader_2, WanVideoVAELoader, LoadWanVideoT5TextEncoder, LoadWanVideoClipTextEncoder,
|
||||
WanVideoTextEncode, WanVideoBlockSwap, WanVideoSampler
|
||||
)
|
||||
from .core.blockswap import DisTorch
|
||||
# DisTorch import removed - all implementations now in __init__.py
|
||||
|
||||
current_device = mm.get_torch_device()
|
||||
current_text_encoder_device = mm.text_encoder_device()
|
||||
@@ -460,7 +460,148 @@ def override_class_clip(cls):
|
||||
|
||||
return NodeOverride
|
||||
|
||||
def override_class_with_distorch(cls):
|
||||
def override_class_with_distorch_safetensor(cls):
|
||||
"""DisTorch wrapper for SafeTensor models, providing block-swap memory optimization."""
|
||||
|
||||
class NodeOverrideDisTorchSafeTensor(cls):
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
inputs = copy.deepcopy(cls.INPUT_TYPES())
|
||||
devices = get_device_list()
|
||||
compute_device = devices[1] if len(devices) > 1 else devices[0]
|
||||
|
||||
inputs["optional"] = inputs.get("optional", {})
|
||||
|
||||
# Reordered and renamed parameters
|
||||
inputs["optional"]["compute_device"] = (devices, {
|
||||
"default": compute_device,
|
||||
"tooltip": "Primary device for computation."
|
||||
})
|
||||
inputs["optional"]["compute_reserved_swap_gb"] = ("FLOAT", {
|
||||
"default": 1.0,
|
||||
"min": 0.1,
|
||||
"max": 16.0,
|
||||
"step": 0.1,
|
||||
"tooltip": "GB of VRAM to keep reserved on the compute device."
|
||||
})
|
||||
inputs["optional"]["virtualram_swap_device"] = (devices, {
|
||||
"default": "cpu",
|
||||
"tooltip": "Device to offload inactive model blocks to."
|
||||
})
|
||||
inputs["optional"]["virtualram_gb"] = ("FLOAT", {
|
||||
"default": 4.0,
|
||||
"min": 0.1,
|
||||
"max": 64.0,
|
||||
"step": 0.1,
|
||||
"tooltip": "Amount of VRAM (in GB) to offload to the swap device."
|
||||
})
|
||||
return inputs
|
||||
|
||||
CATEGORY = "multigpu"
|
||||
FUNCTION = "override"
|
||||
|
||||
def override(self, *args, compute_device=None, compute_reserved_swap_gb=1.0,
|
||||
virtualram_swap_device="cpu", virtualram_gb=4.0, **kwargs):
|
||||
global current_device
|
||||
|
||||
logging.info(f"[DisTorch SafeTensor] Override called with: compute_device={compute_device}, swap_device={virtualram_swap_device}, virtualram_gb={virtualram_gb}, reserved_gb={compute_reserved_swap_gb}")
|
||||
|
||||
if compute_device is not None:
|
||||
current_device = compute_device
|
||||
|
||||
# Call original loader function
|
||||
fn = getattr(super(), cls.FUNCTION)
|
||||
out = fn(*args, **kwargs)
|
||||
|
||||
# Apply block swap logic to the loaded model
|
||||
model = out[0]
|
||||
if hasattr(model, 'model'):
|
||||
logging.info("[DisTorch SafeTensor] Model has 'model' attribute, applying block swap.")
|
||||
apply_block_swap(
|
||||
model,
|
||||
compute_device=compute_device,
|
||||
swap_device=virtualram_swap_device,
|
||||
virtual_vram_gb=virtualram_gb,
|
||||
reserved_swap_gb=compute_reserved_swap_gb
|
||||
)
|
||||
else:
|
||||
logging.warning("[DisTorch SafeTensor] Loaded object does not have a 'model' attribute, skipping block swap.")
|
||||
|
||||
return out
|
||||
|
||||
return NodeOverrideDisTorchSafeTensor
|
||||
|
||||
|
||||
def apply_block_swap(model_patcher, compute_device="cuda:0", swap_device="cpu",
|
||||
virtual_vram_gb=4.0, reserved_swap_gb=1.0):
|
||||
"""
|
||||
Applies WanVideo-style block swapping by patching the forward method of individual model blocks.
|
||||
This allows for offloading parts of the model to a swap device to conserve VRAM.
|
||||
"""
|
||||
logging.info(f"[DisTorch SafeTensor] Initializing block swap: compute_device={compute_device}, swap_device={swap_device}")
|
||||
|
||||
model_to_patch = None
|
||||
if hasattr(model_patcher, 'model') and hasattr(model_patcher.model, 'diffusion_model'):
|
||||
model_to_patch = model_patcher.model.diffusion_model
|
||||
elif hasattr(model_patcher, 'model'):
|
||||
model_to_patch = model_patcher.model
|
||||
else:
|
||||
logging.error("[DisTorch SafeTensor] Could not find a valid model to patch for block swapping.")
|
||||
return
|
||||
|
||||
all_blocks = []
|
||||
if hasattr(model_to_patch, 'input_blocks'):
|
||||
all_blocks.extend(model_to_patch.input_blocks)
|
||||
if hasattr(model_to_patch, 'middle_block'):
|
||||
if isinstance(model_to_patch.middle_block, torch.nn.Module):
|
||||
all_blocks.append(model_to_patch.middle_block)
|
||||
if hasattr(model_to_patch, 'output_blocks'):
|
||||
all_blocks.extend(model_to_patch.output_blocks)
|
||||
|
||||
if not all_blocks:
|
||||
logging.warning("[DisTorch SafeTensor] No swappable blocks were found in the model.")
|
||||
return
|
||||
|
||||
logging.info(f"[DisTorch SafeTensor] Found {len(all_blocks)} swappable blocks in the model.")
|
||||
|
||||
model_size_gb = sum(p.numel() * p.element_size() for p in model_to_patch.parameters()) / (1024**3)
|
||||
if len(all_blocks) > 0:
|
||||
block_size_gb = model_size_gb / len(all_blocks)
|
||||
blocks_to_offload = int(virtual_vram_gb / block_size_gb) if block_size_gb > 0 else 0
|
||||
blocks_on_compute = len(all_blocks) - blocks_to_offload
|
||||
else:
|
||||
blocks_to_offload = 0
|
||||
blocks_on_compute = 0
|
||||
|
||||
logging.info(f"[DisTorch SafeTensor] Model size: {model_size_gb:.2f} GB, Avg block size: {block_size_gb * 1024:.2f} MB")
|
||||
logging.info(f"[DisTorch SafeTensor] Offloading {blocks_to_offload} blocks to {swap_device}. Keeping {blocks_on_compute} blocks on {compute_device}.")
|
||||
|
||||
for i, block in enumerate(all_blocks):
|
||||
if i < blocks_on_compute:
|
||||
block.to(compute_device)
|
||||
else:
|
||||
block.to(swap_device)
|
||||
|
||||
original_forward = block.forward
|
||||
|
||||
def create_patched_forward(original_f, b, cd, sd):
|
||||
def patched_forward(*args, **kwargs):
|
||||
b.to(cd, non_blocking=True)
|
||||
result = original_f(*args, **kwargs)
|
||||
b.to(sd, non_blocking=True)
|
||||
return result
|
||||
return patched_forward
|
||||
|
||||
block.forward = create_patched_forward(original_forward, block, torch.device(compute_device), torch.device(swap_device))
|
||||
logging.info(f"[DisTorch SafeTensor] Patched forward method for block {i}.")
|
||||
|
||||
logging.info("[DisTorch SafeTensor] Block swap setup complete.")
|
||||
# For backwards compatibility, keep the old name pointing to the new safetensor wrapper
|
||||
override_class_with_distorch_bs = override_class_with_distorch_safetensor
|
||||
|
||||
# EXISTING DisTorch wrapper for GGUF models - DO NOT MODIFY
|
||||
def override_class_with_distorch_gguf(cls):
|
||||
"""DisTorch wrapper for GGUF models - DO NOT MODIFY"""
|
||||
class NodeOverrideDisTorch(cls):
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
@@ -503,7 +644,7 @@ def override_class_with_distorch(cls):
|
||||
|
||||
full_allocation = f"{expert_mode_allocations}#{vram_string}" if expert_mode_allocations or vram_string else ""
|
||||
|
||||
logging.info(f"[DisTorch] Full allocation string: {full_allocation}")
|
||||
logging.info(f"[DisTorch GGUF] Full allocation string: {full_allocation}")
|
||||
|
||||
if hasattr(out[0], 'model'):
|
||||
model_hash = create_model_hash(out[0], "override")
|
||||
@@ -516,6 +657,9 @@ def override_class_with_distorch(cls):
|
||||
|
||||
return NodeOverrideDisTorch
|
||||
|
||||
# Keep old name for compatibility but point to GGUF version
|
||||
override_class_with_distorch = override_class_with_distorch_gguf
|
||||
|
||||
def override_class_with_distorch_clip(cls):
|
||||
class NodeOverrideDisTorch(cls):
|
||||
@classmethod
|
||||
@@ -584,7 +728,6 @@ def check_module_exists(module_path):
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"DeviceSelectorMultiGPU": DeviceSelectorMultiGPU,
|
||||
"HunyuanVideoEmbeddingsAdapter": HunyuanVideoEmbeddingsAdapter,
|
||||
"DisTorch": DisTorch,
|
||||
}
|
||||
|
||||
|
||||
@@ -598,6 +741,21 @@ if "QuadrupleCLIPLoader" in GLOBAL_NODE_CLASS_MAPPINGS:
|
||||
NODE_CLASS_MAPPINGS["QuadrupleCLIPLoaderMultiGPU"] = override_class_clip(GLOBAL_NODE_CLASS_MAPPINGS["QuadrupleCLIPLoader"])
|
||||
NODE_CLASS_MAPPINGS["CLIPVisionLoaderMultiGPU"] = override_class_clip(GLOBAL_NODE_CLASS_MAPPINGS["CLIPVisionLoader"])
|
||||
NODE_CLASS_MAPPINGS["CheckpointLoaderSimpleMultiGPU"] = override_class(GLOBAL_NODE_CLASS_MAPPINGS["CheckpointLoaderSimple"])
|
||||
NODE_CLASS_MAPPINGS["UNETLoaderDisTorchMultiGPU"] = override_class_with_distorch_safetensor(GLOBAL_NODE_CLASS_MAPPINGS["UNETLoader"])
|
||||
NODE_CLASS_MAPPINGS["VAELoaderDisTorchMultiGPU"] = override_class_with_distorch_safetensor(GLOBAL_NODE_CLASS_MAPPINGS["VAELoader"])
|
||||
NODE_CLASS_MAPPINGS["CLIPLoaderDisTorchMultiGPU"] = override_class_with_distorch_safetensor(GLOBAL_NODE_CLASS_MAPPINGS["CLIPLoader"])
|
||||
NODE_CLASS_MAPPINGS["DualCLIPLoaderDisTorchMultiGPU"] = override_class_with_distorch_safetensor(GLOBAL_NODE_CLASS_MAPPINGS["DualCLIPLoader"])
|
||||
if "TripleCLIPLoader" in GLOBAL_NODE_CLASS_MAPPINGS:
|
||||
NODE_CLASS_MAPPINGS["TripleCLIPLoaderDisTorchMultiGPU"] = override_class_with_distorch_safetensor(GLOBAL_NODE_CLASS_MAPPINGS["TripleCLIPLoader"])
|
||||
if "QuadrupleCLIPLoader" in GLOBAL_NODE_CLASS_MAPPINGS:
|
||||
NODE_CLASS_MAPPINGS["QuadrupleCLIPLoaderDisTorchMultiGPU"] = override_class_with_distorch_safetensor(GLOBAL_NODE_CLASS_MAPPINGS["QuadrupleCLIPLoader"])
|
||||
NODE_CLASS_MAPPINGS["CLIPVisionLoaderDisTorchMultiGPU"] = override_class_with_distorch_safetensor(GLOBAL_NODE_CLASS_MAPPINGS["CLIPVisionLoader"])
|
||||
NODE_CLASS_MAPPINGS["CheckpointLoaderSimpleDisTorchMultiGPU"] = override_class_with_distorch_safetensor(GLOBAL_NODE_CLASS_MAPPINGS["CheckpointLoaderSimple"])
|
||||
NODE_CLASS_MAPPINGS["ControlNetLoaderDisTorchMultiGPU"] = override_class_with_distorch_safetensor(GLOBAL_NODE_CLASS_MAPPINGS["ControlNetLoader"])
|
||||
if "DiffusersLoader" in GLOBAL_NODE_CLASS_MAPPINGS:
|
||||
NODE_CLASS_MAPPINGS["DiffusersLoaderDisTorchMultiGPU"] = override_class_with_distorch_safetensor(GLOBAL_NODE_CLASS_MAPPINGS["DiffusersLoader"])
|
||||
if "DiffControlNetLoader" in GLOBAL_NODE_CLASS_MAPPINGS:
|
||||
NODE_CLASS_MAPPINGS["DiffControlNetLoaderDisTorchMultiGPU"] = override_class_with_distorch_safetensor(GLOBAL_NODE_CLASS_MAPPINGS["DiffControlNetLoader"])
|
||||
NODE_CLASS_MAPPINGS["ControlNetLoaderMultiGPU"] = override_class(GLOBAL_NODE_CLASS_MAPPINGS["ControlNetLoader"])
|
||||
if "DiffusersLoader" in GLOBAL_NODE_CLASS_MAPPINGS:
|
||||
NODE_CLASS_MAPPINGS["DiffusersLoaderMultiGPU"] = override_class(GLOBAL_NODE_CLASS_MAPPINGS["DiffusersLoader"])
|
||||
@@ -624,9 +782,9 @@ if check_module_exists("ComfyUI-MMAudio") or check_module_exists("comfyui-mmaudi
|
||||
|
||||
if check_module_exists("ComfyUI-GGUF") or check_module_exists("comfyui-gguf"):
|
||||
NODE_CLASS_MAPPINGS["UnetLoaderGGUFMultiGPU"] = override_class(UnetLoaderGGUF)
|
||||
NODE_CLASS_MAPPINGS["UnetLoaderGGUFDisTorchMultiGPU"] = override_class_with_distorch(UnetLoaderGGUF)
|
||||
NODE_CLASS_MAPPINGS["UnetLoaderGGUFDisTorchMultiGPU"] = override_class_with_distorch_gguf(UnetLoaderGGUF)
|
||||
NODE_CLASS_MAPPINGS["UnetLoaderGGUFAdvancedMultiGPU"] = override_class(UnetLoaderGGUFAdvanced)
|
||||
NODE_CLASS_MAPPINGS["UnetLoaderGGUFAdvancedDisTorchMultiGPU"] = override_class_with_distorch(UnetLoaderGGUFAdvanced)
|
||||
NODE_CLASS_MAPPINGS["UnetLoaderGGUFAdvancedDisTorchMultiGPU"] = override_class_with_distorch_gguf(UnetLoaderGGUFAdvanced)
|
||||
NODE_CLASS_MAPPINGS["CLIPLoaderGGUFMultiGPU"] = override_class_clip(CLIPLoaderGGUF)
|
||||
NODE_CLASS_MAPPINGS["CLIPLoaderGGUFDisTorchMultiGPU"] = override_class_with_distorch_clip(CLIPLoaderGGUF)
|
||||
NODE_CLASS_MAPPINGS["DualCLIPLoaderGGUFMultiGPU"] = override_class_clip(DualCLIPLoaderGGUF)
|
||||
|
||||
Reference in New Issue
Block a user