648 lines
24 KiB
Python
648 lines
24 KiB
Python
import importlib
|
|
import logging
|
|
import re
|
|
import sys
|
|
import traceback
|
|
from dataclasses import dataclass, field
|
|
from typing import Any, Optional
|
|
|
|
import torch
|
|
|
|
import comfy.ops
|
|
import comfy.sd
|
|
from comfy.comfy_types.node_typing import IO
|
|
import folder_paths
|
|
|
|
|
|
BACKENDS = [
|
|
"auto_torch_npu_with_fallback",
|
|
"torch_npu_strict",
|
|
"fallback_dequant_only",
|
|
]
|
|
|
|
SCALE_MODES = [
|
|
"per_channel",
|
|
"per_tensor",
|
|
]
|
|
|
|
_RUNTIME_STATS = {}
|
|
|
|
|
|
def _reset_runtime_stats():
|
|
_RUNTIME_STATS.clear()
|
|
_RUNTIME_STATS.update(
|
|
{
|
|
"linear_forward_calls": 0,
|
|
"torch_npu_int8_calls": 0,
|
|
"auto_fallback_dequant_calls": 0,
|
|
"forced_dequant_calls": 0,
|
|
"weight_patch_dequant_calls": 0,
|
|
"nan_input_calls": 0,
|
|
"inf_input_calls": 0,
|
|
"nan_output_calls": 0,
|
|
"inf_output_calls": 0,
|
|
"last_error": "",
|
|
"layers": {},
|
|
}
|
|
)
|
|
|
|
|
|
def _bump_runtime(key: str, layer_name: str = ""):
|
|
_RUNTIME_STATS[key] = _RUNTIME_STATS.get(key, 0) + 1
|
|
if layer_name:
|
|
layer_stats = _RUNTIME_STATS.setdefault("layers", {}).setdefault(
|
|
layer_name,
|
|
{
|
|
"calls": 0,
|
|
"torch_npu_int8": 0,
|
|
"auto_fallback_dequant": 0,
|
|
"forced_dequant": 0,
|
|
"weight_patch_dequant": 0,
|
|
},
|
|
)
|
|
layer_stats["calls"] += 1
|
|
if key == "torch_npu_int8_calls":
|
|
layer_stats["torch_npu_int8"] += 1
|
|
elif key == "auto_fallback_dequant_calls":
|
|
layer_stats["auto_fallback_dequant"] += 1
|
|
elif key == "forced_dequant_calls":
|
|
layer_stats["forced_dequant"] += 1
|
|
elif key == "weight_patch_dequant_calls":
|
|
layer_stats["weight_patch_dequant"] += 1
|
|
|
|
|
|
def _record_tensor_health(prefix: str, tensor: torch.Tensor):
|
|
if tensor is None or not tensor.dtype.is_floating_point:
|
|
return
|
|
try:
|
|
if torch.isnan(tensor).any().item():
|
|
_RUNTIME_STATS[f"nan_{prefix}_calls"] = _RUNTIME_STATS.get(f"nan_{prefix}_calls", 0) + 1
|
|
if torch.isinf(tensor).any().item():
|
|
_RUNTIME_STATS[f"inf_{prefix}_calls"] = _RUNTIME_STATS.get(f"inf_{prefix}_calls", 0) + 1
|
|
except Exception as e:
|
|
_RUNTIME_STATS["last_error"] = f"tensor health check failed: {e.__class__.__name__}: {e}"
|
|
|
|
|
|
def _format_runtime_stats() -> str:
|
|
layer_lines = []
|
|
layers = _RUNTIME_STATS.get("layers", {})
|
|
for name, stats in sorted(layers.items(), key=lambda item: item[1]["calls"], reverse=True)[:40]:
|
|
layer_lines.append(
|
|
f" {name}: calls={stats['calls']}, "
|
|
f"torch_npu_int8={stats['torch_npu_int8']}, "
|
|
f"auto_fallback_dequant={stats['auto_fallback_dequant']}, "
|
|
f"forced_dequant={stats['forced_dequant']}, "
|
|
f"weight_patch_dequant={stats['weight_patch_dequant']}"
|
|
)
|
|
if len(layers) > 40:
|
|
layer_lines.append(f" ... {len(layers) - 40} more")
|
|
if not layer_lines:
|
|
layer_lines.append(" -")
|
|
|
|
return "\n".join(
|
|
[
|
|
"Ascend INT8 runtime report",
|
|
f"linear_forward_calls: {_RUNTIME_STATS.get('linear_forward_calls', 0)}",
|
|
f"torch_npu_int8_calls: {_RUNTIME_STATS.get('torch_npu_int8_calls', 0)}",
|
|
f"auto_fallback_dequant_calls: {_RUNTIME_STATS.get('auto_fallback_dequant_calls', 0)}",
|
|
f"forced_dequant_calls: {_RUNTIME_STATS.get('forced_dequant_calls', 0)}",
|
|
f"weight_patch_dequant_calls: {_RUNTIME_STATS.get('weight_patch_dequant_calls', 0)}",
|
|
f"nan_input_calls: {_RUNTIME_STATS.get('nan_input_calls', 0)}",
|
|
f"inf_input_calls: {_RUNTIME_STATS.get('inf_input_calls', 0)}",
|
|
f"nan_output_calls: {_RUNTIME_STATS.get('nan_output_calls', 0)}",
|
|
f"inf_output_calls: {_RUNTIME_STATS.get('inf_output_calls', 0)}",
|
|
f"last_error: {_RUNTIME_STATS.get('last_error', '') or '-'}",
|
|
"top_layers:",
|
|
*layer_lines,
|
|
]
|
|
)
|
|
|
|
|
|
_reset_runtime_stats()
|
|
|
|
|
|
@dataclass
|
|
class Int8LoadStats:
|
|
seen_linear: int = 0
|
|
quantized_linear: int = 0
|
|
skipped_linear: int = 0
|
|
original_weight_bytes: int = 0
|
|
quantized_weight_bytes: int = 0
|
|
quantized_layers: list[str] = field(default_factory=list)
|
|
skipped_layers: list[str] = field(default_factory=list)
|
|
|
|
|
|
@dataclass
|
|
class Int8Config:
|
|
backend: str
|
|
scale_mode: str
|
|
include_regex: str
|
|
exclude_regex: str
|
|
min_in_features: int
|
|
min_out_features: int
|
|
stats: Int8LoadStats = field(default_factory=Int8LoadStats)
|
|
|
|
def __post_init__(self):
|
|
self.include_re = re.compile(self.include_regex or ".*")
|
|
self.exclude_re = re.compile(self.exclude_regex) if self.exclude_regex else None
|
|
|
|
def should_quantize(self, layer_name: str, in_features: int, out_features: int) -> tuple[bool, str]:
|
|
self.stats.seen_linear += 1
|
|
if in_features < self.min_in_features:
|
|
return False, f"in_features<{self.min_in_features}"
|
|
if out_features < self.min_out_features:
|
|
return False, f"out_features<{self.min_out_features}"
|
|
if not self.include_re.search(layer_name):
|
|
return False, "include_regex"
|
|
if self.exclude_re is not None and self.exclude_re.search(layer_name):
|
|
return False, "exclude_regex"
|
|
return True, "quantized"
|
|
|
|
|
|
def _dtype_nbytes(dtype: torch.dtype) -> int:
|
|
try:
|
|
return torch.tensor([], dtype=dtype).element_size()
|
|
except Exception:
|
|
return 4
|
|
|
|
|
|
def _tensor_nbytes(tensor: Optional[torch.Tensor]) -> int:
|
|
if tensor is None:
|
|
return 0
|
|
return tensor.numel() * _dtype_nbytes(tensor.dtype)
|
|
|
|
|
|
def _short_list(items: list[str], limit: int = 40) -> str:
|
|
if not items:
|
|
return "-"
|
|
shown = items[:limit]
|
|
suffix = "" if len(items) <= limit else f"\n ... {len(items) - limit} more"
|
|
return "\n " + "\n ".join(shown) + suffix
|
|
|
|
|
|
def _safe_version(module_name: str) -> str:
|
|
try:
|
|
module = importlib.import_module(module_name)
|
|
return str(getattr(module, "__version__", "unknown"))
|
|
except Exception as e:
|
|
return f"not importable: {e.__class__.__name__}: {e}"
|
|
|
|
|
|
def _torch_npu_module():
|
|
try:
|
|
return importlib.import_module("torch_npu")
|
|
except Exception:
|
|
return None
|
|
|
|
|
|
def _linear_weight_quant_op():
|
|
torch_npu = _torch_npu_module()
|
|
if torch_npu is None:
|
|
return None
|
|
return getattr(torch_npu, "npu_weight_quant_batchmatmul", None)
|
|
|
|
|
|
def _linear_weight_quant_module_available() -> bool:
|
|
try:
|
|
module = importlib.import_module("torch_npu.contrib.module")
|
|
return hasattr(module, "LinearWeightQuant")
|
|
except Exception:
|
|
return False
|
|
|
|
|
|
def _quantize_weight_to_int8(weight: torch.Tensor, scale_mode: str) -> tuple[torch.Tensor, torch.Tensor]:
|
|
# torch_npu.npu_weight_quant_batchmatmul computes x @ antiquant(weight).
|
|
# F.linear computes x @ original_weight.T, so store original_weight.T as K,N.
|
|
weight_t = torch.nan_to_num(weight.detach().to(torch.float32).t().contiguous())
|
|
eps = torch.tensor(1.0e-8, dtype=torch.float32, device=weight_t.device)
|
|
|
|
if scale_mode == "per_tensor":
|
|
scale = torch.clamp(weight_t.abs().amax() / 127.0, min=eps)
|
|
qweight = torch.clamp(torch.round(weight_t / scale), -127, 127).to(torch.int8)
|
|
return qweight.contiguous(), scale.reshape(1).contiguous()
|
|
|
|
if scale_mode != "per_channel":
|
|
raise ValueError(f"Unsupported scale_mode: {scale_mode}")
|
|
|
|
scale = torch.clamp(weight_t.abs().amax(dim=0) / 127.0, min=eps)
|
|
qweight = torch.clamp(torch.round(weight_t / scale), -127, 127).to(torch.int8)
|
|
return qweight.contiguous(), scale.contiguous()
|
|
|
|
|
|
def _dequantized_linear(
|
|
x_2d: torch.Tensor,
|
|
qweight: torch.Tensor,
|
|
antiquant_scale: torch.Tensor,
|
|
bias: Optional[torch.Tensor],
|
|
) -> torch.Tensor:
|
|
scale = antiquant_scale.to(device=x_2d.device, dtype=torch.float32)
|
|
q = qweight.to(device=x_2d.device)
|
|
weight_t = q.to(torch.float32) * scale
|
|
weight = weight_t.t().to(dtype=x_2d.dtype)
|
|
b = bias.to(device=x_2d.device, dtype=x_2d.dtype) if bias is not None else None
|
|
return torch.nn.functional.linear(x_2d, weight, b)
|
|
|
|
|
|
def _torch_npu_weight_quant_linear(
|
|
x_2d: torch.Tensor,
|
|
qweight: torch.Tensor,
|
|
antiquant_scale: torch.Tensor,
|
|
bias: Optional[torch.Tensor],
|
|
) -> torch.Tensor:
|
|
op = _linear_weight_quant_op()
|
|
if op is None:
|
|
raise RuntimeError("torch_npu.npu_weight_quant_batchmatmul is not available")
|
|
if x_2d.device.type != "npu":
|
|
raise RuntimeError(f"input device is {x_2d.device.type}, expected npu")
|
|
|
|
qw = qweight.to(device=x_2d.device)
|
|
scale = antiquant_scale.to(device=x_2d.device, dtype=torch.float32)
|
|
b = bias.to(device=x_2d.device, dtype=x_2d.dtype) if bias is not None else None
|
|
|
|
try:
|
|
return op(x_2d, qw, scale, None, None, None, b, 0, 0)
|
|
except TypeError:
|
|
return op(
|
|
x_2d,
|
|
qw,
|
|
scale,
|
|
antiquant_offset=None,
|
|
quant_scale=None,
|
|
quant_offset=None,
|
|
bias=b,
|
|
antiquant_group_size=0,
|
|
inner_precise=0,
|
|
)
|
|
|
|
|
|
def _format_stats(config: Int8Config, title: str) -> str:
|
|
stats = config.stats
|
|
saved = stats.original_weight_bytes - stats.quantized_weight_bytes
|
|
ratio = 0.0
|
|
if stats.original_weight_bytes:
|
|
ratio = stats.quantized_weight_bytes / stats.original_weight_bytes
|
|
|
|
return "\n".join(
|
|
[
|
|
title,
|
|
f"backend: {config.backend}",
|
|
f"scale_mode: {config.scale_mode}",
|
|
f"include_regex: {config.include_regex or '.*'}",
|
|
f"exclude_regex: {config.exclude_regex or '-'}",
|
|
f"min_in_features: {config.min_in_features}",
|
|
f"min_out_features: {config.min_out_features}",
|
|
f"linear_seen: {stats.seen_linear}",
|
|
f"linear_quantized: {stats.quantized_linear}",
|
|
f"linear_skipped: {stats.skipped_linear}",
|
|
f"original_linear_weight_bytes: {stats.original_weight_bytes}",
|
|
f"int8_weight_plus_scale_bytes: {stats.quantized_weight_bytes}",
|
|
f"rough_weight_ratio: {ratio:.4f}",
|
|
f"rough_bytes_saved: {saved}",
|
|
"quantized_layers:" + _short_list(stats.quantized_layers),
|
|
"skipped_layers:" + _short_list(stats.skipped_layers),
|
|
]
|
|
)
|
|
|
|
|
|
def make_ascend_int8_ops(config: Int8Config):
|
|
class AscendInt8Ops(comfy.ops.manual_cast):
|
|
class Linear(comfy.ops.manual_cast.Linear):
|
|
comfy_cast_weights = True
|
|
|
|
def __init__(self, in_features, out_features, bias=True, device=None, dtype=None):
|
|
super().__init__(in_features, out_features, bias, device=device, dtype=dtype)
|
|
self._ascend_int8_enabled = False
|
|
self._ascend_int8_config = config
|
|
self._ascend_int8_layer_name = ""
|
|
|
|
def _load_from_state_dict(
|
|
self,
|
|
state_dict,
|
|
prefix,
|
|
local_metadata,
|
|
strict,
|
|
missing_keys,
|
|
unexpected_keys,
|
|
error_msgs,
|
|
):
|
|
layer_name = prefix.rstrip(".")
|
|
should_quantize, reason = config.should_quantize(layer_name, self.in_features, self.out_features)
|
|
if not should_quantize:
|
|
config.stats.skipped_linear += 1
|
|
config.stats.skipped_layers.append(f"{layer_name}: {reason}")
|
|
return super()._load_from_state_dict(
|
|
state_dict,
|
|
prefix,
|
|
local_metadata,
|
|
strict,
|
|
missing_keys,
|
|
unexpected_keys,
|
|
error_msgs,
|
|
)
|
|
|
|
weight_key = f"{prefix}weight"
|
|
bias_key = f"{prefix}bias"
|
|
weight = state_dict.pop(weight_key, None)
|
|
bias = state_dict.pop(bias_key, None)
|
|
|
|
if weight is None:
|
|
config.stats.skipped_linear += 1
|
|
config.stats.skipped_layers.append(f"{layer_name}: missing weight")
|
|
missing_keys.append(weight_key)
|
|
return
|
|
|
|
try:
|
|
qweight, antiquant_scale = _quantize_weight_to_int8(weight, config.scale_mode)
|
|
except Exception as e:
|
|
config.stats.skipped_linear += 1
|
|
config.stats.skipped_layers.append(f"{layer_name}: quantize failed: {e}")
|
|
error_msgs.append(f"{layer_name}: quantize failed: {traceback.format_exc()}")
|
|
return
|
|
|
|
config.stats.quantized_linear += 1
|
|
config.stats.original_weight_bytes += _tensor_nbytes(weight)
|
|
config.stats.quantized_weight_bytes += _tensor_nbytes(qweight) + _tensor_nbytes(antiquant_scale)
|
|
config.stats.quantized_layers.append(layer_name)
|
|
|
|
self.weight = None
|
|
self.register_buffer("qweight", qweight, persistent=True)
|
|
self.register_buffer("antiquant_scale", antiquant_scale, persistent=True)
|
|
|
|
if bias is not None:
|
|
self.bias = torch.nn.Parameter(bias.detach().clone(), requires_grad=False)
|
|
else:
|
|
self.bias = None
|
|
|
|
self._ascend_int8_enabled = True
|
|
self._ascend_int8_layer_name = layer_name
|
|
|
|
def _bias_for_input(self, input: torch.Tensor) -> Optional[torch.Tensor]:
|
|
bias = self.bias
|
|
if bias is not None:
|
|
for f in self.bias_function:
|
|
bias = f(bias)
|
|
bias = bias.to(device=input.device, dtype=input.dtype)
|
|
return bias
|
|
|
|
def _dequant_weight_for_input(self, input: torch.Tensor) -> torch.Tensor:
|
|
scale = self.antiquant_scale.to(device=input.device, dtype=torch.float32)
|
|
q = self.qweight.to(device=input.device)
|
|
weight_t = q.to(torch.float32) * scale
|
|
return weight_t.t().to(dtype=input.dtype)
|
|
|
|
def forward(self, input, *args, **kwargs):
|
|
if not self._ascend_int8_enabled:
|
|
return super().forward(input, *args, **kwargs)
|
|
|
|
comfy.ops.run_every_op()
|
|
_RUNTIME_STATS["linear_forward_calls"] = _RUNTIME_STATS.get("linear_forward_calls", 0) + 1
|
|
_record_tensor_health("input", input)
|
|
|
|
input_shape = tuple(input.shape)
|
|
if input.ndim < 2:
|
|
raise RuntimeError(
|
|
f"Ascend INT8 Linear expects input rank >= 2, got {input.ndim} "
|
|
f"at {self._ascend_int8_layer_name}"
|
|
)
|
|
x_2d = input.reshape(-1, input_shape[-1])
|
|
bias = self._bias_for_input(input)
|
|
|
|
# Weight patch functions, such as LoRA, need a real dequantized weight.
|
|
if len(self.weight_function) > 0 or len(self.bias_function) > 0:
|
|
_bump_runtime("weight_patch_dequant_calls", self._ascend_int8_layer_name)
|
|
weight = self._dequant_weight_for_input(input)
|
|
for f in self.weight_function:
|
|
weight = f(weight)
|
|
out = torch.nn.functional.linear(x_2d, weight, bias)
|
|
elif config.backend == "fallback_dequant_only":
|
|
_bump_runtime("forced_dequant_calls", self._ascend_int8_layer_name)
|
|
out = _dequantized_linear(x_2d, self.qweight, self.antiquant_scale, bias)
|
|
elif config.backend == "torch_npu_strict":
|
|
out = _torch_npu_weight_quant_linear(x_2d, self.qweight, self.antiquant_scale, bias)
|
|
_bump_runtime("torch_npu_int8_calls", self._ascend_int8_layer_name)
|
|
else:
|
|
try:
|
|
out = _torch_npu_weight_quant_linear(x_2d, self.qweight, self.antiquant_scale, bias)
|
|
_bump_runtime("torch_npu_int8_calls", self._ascend_int8_layer_name)
|
|
except Exception as e:
|
|
_RUNTIME_STATS["last_error"] = f"{self._ascend_int8_layer_name}: {e.__class__.__name__}: {e}"
|
|
_bump_runtime("auto_fallback_dequant_calls", self._ascend_int8_layer_name)
|
|
logging.debug(
|
|
"Ascend INT8 fallback for %s: %s",
|
|
self._ascend_int8_layer_name,
|
|
e,
|
|
)
|
|
out = _dequantized_linear(x_2d, self.qweight, self.antiquant_scale, bias)
|
|
|
|
_record_tensor_health("output", out)
|
|
return out.reshape(input_shape[:-1] + (self.out_features,))
|
|
|
|
def extra_repr(self):
|
|
base = f"in_features={self.in_features}, out_features={self.out_features}, bias={self.bias is not None}"
|
|
if self._ascend_int8_enabled:
|
|
return base + f", ascend_int8=True, backend={config.backend}, scale_mode={config.scale_mode}"
|
|
return base
|
|
|
|
return AscendInt8Ops
|
|
|
|
|
|
def _make_config(backend, scale_mode, include_regex, exclude_regex, min_in_features, min_out_features) -> Int8Config:
|
|
if backend not in BACKENDS:
|
|
raise ValueError(f"Unsupported backend: {backend}")
|
|
if scale_mode not in SCALE_MODES:
|
|
raise ValueError(f"Unsupported scale_mode: {scale_mode}")
|
|
return Int8Config(
|
|
backend=backend,
|
|
scale_mode=scale_mode,
|
|
include_regex=include_regex,
|
|
exclude_regex=exclude_regex,
|
|
min_in_features=int(min_in_features),
|
|
min_out_features=int(min_out_features),
|
|
)
|
|
|
|
|
|
def _load_model_options(config: Int8Config) -> dict[str, Any]:
|
|
return {"custom_operations": make_ascend_int8_ops(config)}
|
|
|
|
|
|
class AscendInt8EnvironmentReport:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"probe_torch_npu": ("BOOLEAN", {"default": True}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("STRING",)
|
|
RETURN_NAMES = ("report",)
|
|
FUNCTION = "run"
|
|
OUTPUT_NODE = True
|
|
CATEGORY = "ascend/int8"
|
|
|
|
@classmethod
|
|
def IS_CHANGED(cls, **kwargs):
|
|
return float("nan")
|
|
|
|
def run(self, probe_torch_npu=True):
|
|
lines = [
|
|
"Ascend INT8 environment report",
|
|
f"python: {sys.executable}",
|
|
f"torch: {torch.__version__}",
|
|
f"torch_cuda: {torch.version.cuda}",
|
|
f"torch_npu: {_safe_version('torch_npu') if probe_torch_npu else 'not probed'}",
|
|
f"has torch.npu: {hasattr(torch, 'npu')}",
|
|
]
|
|
|
|
if hasattr(torch, "npu"):
|
|
try:
|
|
lines.append(f"torch.npu.is_available: {torch.npu.is_available()}")
|
|
lines.append(f"torch.npu.device_count: {torch.npu.device_count()}")
|
|
if torch.npu.is_available():
|
|
lines.append(f"current_npu: {torch.npu.current_device()}")
|
|
lines.append(f"npu_name: {torch.npu.get_device_name(torch.npu.current_device())}")
|
|
except Exception as e:
|
|
lines.append(f"torch.npu probe_error: {e.__class__.__name__}: {e}")
|
|
|
|
if probe_torch_npu:
|
|
lines.extend(
|
|
[
|
|
f"npu_weight_quant_batchmatmul: {_linear_weight_quant_op() is not None}",
|
|
f"npu_anti_quant: {hasattr(_torch_npu_module(), 'npu_anti_quant') if _torch_npu_module() else False}",
|
|
f"LinearWeightQuant: {_linear_weight_quant_module_available()}",
|
|
]
|
|
)
|
|
|
|
report = "\n".join(lines)
|
|
return {"ui": {"text": (report,)}, "result": (report,)}
|
|
|
|
|
|
class AscendInt8RuntimeReport:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"source": (IO.ANY, {}),
|
|
"reset_after_report": ("BOOLEAN", {"default": False}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("STRING",)
|
|
RETURN_NAMES = ("report",)
|
|
FUNCTION = "run"
|
|
OUTPUT_NODE = True
|
|
CATEGORY = "ascend/int8"
|
|
|
|
@classmethod
|
|
def IS_CHANGED(cls, **kwargs):
|
|
return float("nan")
|
|
|
|
def run(self, source=None, reset_after_report=False):
|
|
report = _format_runtime_stats()
|
|
if reset_after_report:
|
|
_reset_runtime_stats()
|
|
return {"ui": {"text": (report,)}, "result": (report,)}
|
|
|
|
|
|
class AscendInt8UNETLoader:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"unet_name": (folder_paths.get_filename_list("diffusion_models"),),
|
|
"backend": (BACKENDS, {"default": "auto_torch_npu_with_fallback"}),
|
|
"scale_mode": (SCALE_MODES, {"default": "per_channel"}),
|
|
"include_regex": ("STRING", {"default": ".*", "multiline": False}),
|
|
"exclude_regex": ("STRING", {"default": "", "multiline": False}),
|
|
"min_in_features": ("INT", {"default": 16, "min": 1, "max": 65536}),
|
|
"min_out_features": ("INT", {"default": 16, "min": 1, "max": 65536}),
|
|
"reset_runtime_stats": ("BOOLEAN", {"default": True}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("MODEL", "STRING")
|
|
RETURN_NAMES = ("model", "report")
|
|
FUNCTION = "load_unet"
|
|
CATEGORY = "ascend/int8"
|
|
|
|
def load_unet(
|
|
self,
|
|
unet_name,
|
|
backend,
|
|
scale_mode,
|
|
include_regex,
|
|
exclude_regex,
|
|
min_in_features,
|
|
min_out_features,
|
|
reset_runtime_stats=True,
|
|
):
|
|
if reset_runtime_stats:
|
|
_reset_runtime_stats()
|
|
config = _make_config(backend, scale_mode, include_regex, exclude_regex, min_in_features, min_out_features)
|
|
unet_path = folder_paths.get_full_path_or_raise("diffusion_models", unet_name)
|
|
model = comfy.sd.load_diffusion_model(unet_path, model_options=_load_model_options(config))
|
|
report = _format_stats(config, f"Ascend INT8 UNET load report: {unet_name}")
|
|
return (model, report)
|
|
|
|
|
|
class AscendInt8CheckpointLoader:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"ckpt_name": (folder_paths.get_filename_list("checkpoints"),),
|
|
"backend": (BACKENDS, {"default": "auto_torch_npu_with_fallback"}),
|
|
"scale_mode": (SCALE_MODES, {"default": "per_channel"}),
|
|
"include_regex": ("STRING", {"default": ".*", "multiline": False}),
|
|
"exclude_regex": ("STRING", {"default": "", "multiline": False}),
|
|
"min_in_features": ("INT", {"default": 16, "min": 1, "max": 65536}),
|
|
"min_out_features": ("INT", {"default": 16, "min": 1, "max": 65536}),
|
|
"reset_runtime_stats": ("BOOLEAN", {"default": True}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("MODEL", "CLIP", "VAE", "STRING")
|
|
RETURN_NAMES = ("model", "clip", "vae", "report")
|
|
FUNCTION = "load_checkpoint"
|
|
CATEGORY = "ascend/int8"
|
|
|
|
def load_checkpoint(
|
|
self,
|
|
ckpt_name,
|
|
backend,
|
|
scale_mode,
|
|
include_regex,
|
|
exclude_regex,
|
|
min_in_features,
|
|
min_out_features,
|
|
reset_runtime_stats=True,
|
|
):
|
|
if reset_runtime_stats:
|
|
_reset_runtime_stats()
|
|
config = _make_config(backend, scale_mode, include_regex, exclude_regex, min_in_features, min_out_features)
|
|
ckpt_path = folder_paths.get_full_path_or_raise("checkpoints", ckpt_name)
|
|
model, clip, vae = comfy.sd.load_checkpoint_guess_config(
|
|
ckpt_path,
|
|
output_vae=True,
|
|
output_clip=True,
|
|
embedding_directory=folder_paths.get_folder_paths("embeddings"),
|
|
model_options=_load_model_options(config),
|
|
)
|
|
report = _format_stats(config, f"Ascend INT8 checkpoint load report: {ckpt_name}")
|
|
return (model, clip, vae, report)
|
|
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"AscendInt8EnvironmentReport": AscendInt8EnvironmentReport,
|
|
"AscendInt8RuntimeReport": AscendInt8RuntimeReport,
|
|
"AscendInt8UNETLoader": AscendInt8UNETLoader,
|
|
"AscendInt8CheckpointLoader": AscendInt8CheckpointLoader,
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"AscendInt8EnvironmentReport": "Ascend INT8 Environment Report",
|
|
"AscendInt8RuntimeReport": "Ascend INT8 Runtime Report",
|
|
"AscendInt8UNETLoader": "Load Diffusion Model (Ascend INT8)",
|
|
"AscendInt8CheckpointLoader": "Load Checkpoint (Ascend INT8)",
|
|
}
|