Add custom nodes, Civitai loras (LFS), and vast.ai setup script
Some checks failed
Python Linting / Run Ruff (push) Has been cancelled
Python Linting / Run Pylint (push) Has been cancelled
Full Comfy CI Workflow Runs / test-stable (12.1, , linux, 3.10, [self-hosted Linux], stable) (push) Has been cancelled
Full Comfy CI Workflow Runs / test-stable (12.1, , linux, 3.11, [self-hosted Linux], stable) (push) Has been cancelled
Full Comfy CI Workflow Runs / test-stable (12.1, , linux, 3.12, [self-hosted Linux], stable) (push) Has been cancelled
Full Comfy CI Workflow Runs / test-unix-nightly (12.1, , linux, 3.11, [self-hosted Linux], nightly) (push) Has been cancelled
Execution Tests / test (macos-latest) (push) Has been cancelled
Execution Tests / test (ubuntu-latest) (push) Has been cancelled
Execution Tests / test (windows-latest) (push) Has been cancelled
Test server launches without errors / test (push) Has been cancelled
Unit Tests / test (macos-latest) (push) Has been cancelled
Unit Tests / test (ubuntu-latest) (push) Has been cancelled
Unit Tests / test (windows-2022) (push) Has been cancelled
Some checks failed
Python Linting / Run Ruff (push) Has been cancelled
Python Linting / Run Pylint (push) Has been cancelled
Full Comfy CI Workflow Runs / test-stable (12.1, , linux, 3.10, [self-hosted Linux], stable) (push) Has been cancelled
Full Comfy CI Workflow Runs / test-stable (12.1, , linux, 3.11, [self-hosted Linux], stable) (push) Has been cancelled
Full Comfy CI Workflow Runs / test-stable (12.1, , linux, 3.12, [self-hosted Linux], stable) (push) Has been cancelled
Full Comfy CI Workflow Runs / test-unix-nightly (12.1, , linux, 3.11, [self-hosted Linux], nightly) (push) Has been cancelled
Execution Tests / test (macos-latest) (push) Has been cancelled
Execution Tests / test (ubuntu-latest) (push) Has been cancelled
Execution Tests / test (windows-latest) (push) Has been cancelled
Test server launches without errors / test (push) Has been cancelled
Unit Tests / test (macos-latest) (push) Has been cancelled
Unit Tests / test (ubuntu-latest) (push) Has been cancelled
Unit Tests / test (windows-2022) (push) Has been cancelled
Includes 30 custom nodes committed directly, 7 Civitai-exclusive loras stored via Git LFS, and a setup script that installs all dependencies and downloads HuggingFace-hosted models on vast.ai. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
@@ -0,0 +1,167 @@
|
||||
#credit to comfyanonymous for this module
|
||||
#from https://github.com/comfyanonymous/ComfyUI_bitsandbytes_NF4
|
||||
import comfy.ops
|
||||
import torch
|
||||
import folder_paths
|
||||
from ...libs.utils import install_package
|
||||
|
||||
try:
|
||||
from bitsandbytes.nn.modules import Params4bit, QuantState
|
||||
except ImportError:
|
||||
Params4bit = torch.nn.Parameter
|
||||
raise ImportError("Please install bitsandbytes>=0.43.3")
|
||||
|
||||
def functional_linear_4bits(x, weight, bias):
|
||||
try:
|
||||
install_package("bitsandbytes", "0.43.3", True, "0.43.3")
|
||||
import bitsandbytes as bnb
|
||||
except ImportError:
|
||||
raise ImportError("Please install bitsandbytes>=0.43.3")
|
||||
|
||||
out = bnb.matmul_4bit(x, weight.t(), bias=bias, quant_state=weight.quant_state)
|
||||
out = out.to(x)
|
||||
return out
|
||||
|
||||
|
||||
def copy_quant_state(state, device: torch.device = None):
|
||||
if state is None:
|
||||
return None
|
||||
|
||||
device = device or state.absmax.device
|
||||
|
||||
state2 = (
|
||||
QuantState(
|
||||
absmax=state.state2.absmax.to(device),
|
||||
shape=state.state2.shape,
|
||||
code=state.state2.code.to(device),
|
||||
blocksize=state.state2.blocksize,
|
||||
quant_type=state.state2.quant_type,
|
||||
dtype=state.state2.dtype,
|
||||
)
|
||||
if state.nested
|
||||
else None
|
||||
)
|
||||
|
||||
return QuantState(
|
||||
absmax=state.absmax.to(device),
|
||||
shape=state.shape,
|
||||
code=state.code.to(device),
|
||||
blocksize=state.blocksize,
|
||||
quant_type=state.quant_type,
|
||||
dtype=state.dtype,
|
||||
offset=state.offset.to(device) if state.nested else None,
|
||||
state2=state2,
|
||||
)
|
||||
|
||||
|
||||
class ForgeParams4bit(Params4bit):
|
||||
|
||||
def to(self, *args, **kwargs):
|
||||
device, dtype, non_blocking, convert_to_format = torch._C._nn._parse_to(*args, **kwargs)
|
||||
if device is not None and device.type == "cuda" and not self.bnb_quantized:
|
||||
return self._quantize(device)
|
||||
else:
|
||||
n = ForgeParams4bit(
|
||||
torch.nn.Parameter.to(self, device=device, dtype=dtype, non_blocking=non_blocking),
|
||||
requires_grad=self.requires_grad,
|
||||
quant_state=copy_quant_state(self.quant_state, device),
|
||||
blocksize=self.blocksize,
|
||||
compress_statistics=self.compress_statistics,
|
||||
quant_type=self.quant_type,
|
||||
quant_storage=self.quant_storage,
|
||||
bnb_quantized=self.bnb_quantized,
|
||||
module=self.module
|
||||
)
|
||||
self.module.quant_state = n.quant_state
|
||||
self.data = n.data
|
||||
self.quant_state = n.quant_state
|
||||
return n
|
||||
|
||||
class ForgeLoader4Bit(torch.nn.Module):
|
||||
def __init__(self, *, device, dtype, quant_type, **kwargs):
|
||||
super().__init__()
|
||||
self.dummy = torch.nn.Parameter(torch.empty(1, device=device, dtype=dtype))
|
||||
self.weight = None
|
||||
self.quant_state = None
|
||||
self.bias = None
|
||||
self.quant_type = quant_type
|
||||
|
||||
def _save_to_state_dict(self, destination, prefix, keep_vars):
|
||||
super()._save_to_state_dict(destination, prefix, keep_vars)
|
||||
quant_state = getattr(self.weight, "quant_state", None)
|
||||
if quant_state is not None:
|
||||
for k, v in quant_state.as_dict(packed=True).items():
|
||||
destination[prefix + "weight." + k] = v if keep_vars else v.detach()
|
||||
return
|
||||
|
||||
def _load_from_state_dict(self, state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs):
|
||||
quant_state_keys = {k[len(prefix + "weight."):] for k in state_dict.keys() if k.startswith(prefix + "weight.")}
|
||||
|
||||
if any('bitsandbytes' in k for k in quant_state_keys):
|
||||
quant_state_dict = {k: state_dict[prefix + "weight." + k] for k in quant_state_keys}
|
||||
|
||||
self.weight = ForgeParams4bit().from_prequantized(
|
||||
data=state_dict[prefix + 'weight'],
|
||||
quantized_stats=quant_state_dict,
|
||||
requires_grad=False,
|
||||
device=self.dummy.device,
|
||||
module=self
|
||||
)
|
||||
self.quant_state = self.weight.quant_state
|
||||
|
||||
if prefix + 'bias' in state_dict:
|
||||
self.bias = torch.nn.Parameter(state_dict[prefix + 'bias'].to(self.dummy))
|
||||
|
||||
del self.dummy
|
||||
elif hasattr(self, 'dummy'):
|
||||
if prefix + 'weight' in state_dict:
|
||||
self.weight = ForgeParams4bit(
|
||||
state_dict[prefix + 'weight'].to(self.dummy),
|
||||
requires_grad=False,
|
||||
compress_statistics=True,
|
||||
quant_type=self.quant_type,
|
||||
quant_storage=torch.uint8,
|
||||
module=self,
|
||||
)
|
||||
self.quant_state = self.weight.quant_state
|
||||
|
||||
if prefix + 'bias' in state_dict:
|
||||
self.bias = torch.nn.Parameter(state_dict[prefix + 'bias'].to(self.dummy))
|
||||
|
||||
del self.dummy
|
||||
else:
|
||||
super()._load_from_state_dict(state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs)
|
||||
|
||||
current_device = None
|
||||
current_dtype = None
|
||||
current_manual_cast_enabled = False
|
||||
current_bnb_dtype = None
|
||||
|
||||
class OPS(comfy.ops.manual_cast):
|
||||
class Linear(ForgeLoader4Bit):
|
||||
def __init__(self, *args, device=None, dtype=None, **kwargs):
|
||||
super().__init__(device=device, dtype=dtype, quant_type=current_bnb_dtype)
|
||||
self.parameters_manual_cast = current_manual_cast_enabled
|
||||
|
||||
def forward(self, x):
|
||||
self.weight.quant_state = self.quant_state
|
||||
|
||||
if self.bias is not None and self.bias.dtype != x.dtype:
|
||||
# Maybe this can also be set to all non-bnb ops since the cost is very low.
|
||||
# And it only invokes one time, and most linear does not have bias
|
||||
self.bias.data = self.bias.data.to(x.dtype)
|
||||
|
||||
if not self.parameters_manual_cast:
|
||||
return functional_linear_4bits(x, self.weight, self.bias)
|
||||
elif not self.weight.bnb_quantized:
|
||||
assert x.device.type == 'cuda', 'BNB Must Use CUDA as Computation Device!'
|
||||
layer_original_device = self.weight.device
|
||||
self.weight = self.weight._quantize(x.device)
|
||||
bias = self.bias.to(x.device) if self.bias is not None else None
|
||||
out = functional_linear_4bits(x, self.weight, bias)
|
||||
self.weight = self.weight.to(layer_original_device)
|
||||
return out
|
||||
else:
|
||||
weight, bias, signal = weights_manual_cast(self, x, skip_weight_dtype=True, skip_bias_dtype=True)
|
||||
with main_stream_worker(weight, bias, signal):
|
||||
return functional_linear_4bits(x, weight, bias)
|
||||
Reference in New Issue
Block a user