630 lines
26 KiB
Python
630 lines
26 KiB
Python
"""
|
|
XB-ToolBox ROCm 优化节点 (v5.0)
|
|
================================
|
|
8 个节点，专为 AMD ROCm GPU 设计，零外部依赖。
|
|
|
|
XB_ROCmKSampler — 采样器（rocBLAS调优 + SDPA修复）
|
|
XB_ROCmKSamplerAdvanced — 高级采样器（支持分段采样/加噪控制/残留噪声）
|
|
XB_ROCmSamplerCustom — 自定义采样器（模块化调度器/采样器/引导器）
|
|
XB_ROCmSamplerCustomAdvanced — 自定义高级采样器（完全模块化 + 独立噪波注入）
|
|
XB_ROCmVAEDecode — 空间分块解码器（架构自适应 tile）
|
|
XB_ROCmVAEEncode — 空间分块编码器（架构自适应 tile）
|
|
XB_ROCmVAEDecodeTemporal — 空间+时间分块解码器（视频专用）
|
|
XB_ROCmMemCleaner — 显存清理+诊断报告
|
|
|
|
核心机制:
|
|
- tile_size=0 → 自动根据 GPU 架构选择最优分块大小（用户也可手动覆盖）
|
|
- fp16_accumulation=True → rocBLAS 矩阵乘法用 fp16 中间累积（~5-10% 加速）
|
|
- mem_efficient SDPA → AMD 上最稳定的 attention 后端（防崩溃）
|
|
- 时间分块 → 视频 latent 逐 chunk 解码，显存峰值降低 80%
|
|
"""
|
|
|
|
import torch, gc, time
|
|
import comfy.model_management as mm
|
|
import comfy.samplers
|
|
import comfy.sample
|
|
import comfy.utils
|
|
import comfy.nested_tensor
|
|
import latent_preview
|
|
import nodes
|
|
|
|
|
|
class AnyType(str):
|
|
"""哑类型——接受任意连线，仅用于串联节点防止误删"""
|
|
def __eq__(self, _) -> bool: return True
|
|
def __ne__(self, __value: object) -> bool: return False
|
|
|
|
_any = AnyType("*")
|
|
|
|
|
|
_AMD_ARCH_DB = {
|
|
# RDNA4 — RX 9070 系列
|
|
"gfx1201": {"name": "RX 9070 XT", "tile": 768, "gen": "RDNA4"},
|
|
"gfx1200": {"name": "RX 9070", "tile": 640, "gen": "RDNA4"},
|
|
# RDNA3.5 — Strix APU
|
|
"gfx1151": {"name": "Strix Halo", "tile": 768, "gen": "RDNA3.5"},
|
|
"gfx1150": {"name": "Strix Point", "tile": 512, "gen": "RDNA3.5"},
|
|
# RDNA3 — RX 7000
|
|
"gfx1102": {"name": "RX 7600", "tile": 384, "gen": "RDNA3"},
|
|
"gfx1101": {"name": "RX 7800 XT", "tile": 512, "gen": "RDNA3"},
|
|
"gfx1100": {"name": "RX 7900 XTX", "tile": 640, "gen": "RDNA3"},
|
|
# RDNA2 — RX 6000
|
|
"gfx1032": {"name": "RX 6600", "tile": 256, "gen": "RDNA2"},
|
|
"gfx1031": {"name": "RX 6700 XT", "tile": 320, "gen": "RDNA2"},
|
|
"gfx1030": {"name": "RX 6800/6900", "tile": 384, "gen": "RDNA2"},
|
|
# CDNA
|
|
"gfx942": {"name": "MI300X", "tile": 512, "gen": "CDNA3"},
|
|
"gfx90a": {"name": "MI250X", "tile": 512, "gen": "CDNA2"},
|
|
"gfx908": {"name": "MI100", "tile": 256, "gen": "CDNA"},
|
|
"gfx906": {"name": "Radeon VII", "tile": 256, "gen": "Vega"},
|
|
}
|
|
|
|
|
|
def _lookup(arch: str):
|
|
if not arch: return None
|
|
if arch in _AMD_ARCH_DB: return _AMD_ARCH_DB[arch]
|
|
for k in sorted(_AMD_ARCH_DB, key=len, reverse=True):
|
|
if arch.startswith(k): return _AMD_ARCH_DB[k]
|
|
return None
|
|
|
|
|
|
def is_rocm() -> bool:
|
|
return torch.cuda.is_available() and hasattr(torch.version, 'hip') and torch.version.hip
|
|
|
|
|
|
def get_arch() -> str:
|
|
if not is_rocm(): return "nvidia"
|
|
try:
|
|
raw = torch.cuda.get_device_properties(0).gcnArchName
|
|
return raw.split(':')[0] if raw else "unknown"
|
|
except: return "unknown"
|
|
|
|
|
|
def gpu_info() -> dict:
|
|
arch = get_arch()
|
|
d = {"arch": arch, "amd": is_rocm(), "name": "?", "gen": "?", "tile": 256, "gb": 0}
|
|
if is_rocm():
|
|
e = _lookup(arch)
|
|
if e: d.update(e)
|
|
try: d["gb"] = torch.cuda.get_device_properties(0).total_memory / 1024**3
|
|
except: pass
|
|
else:
|
|
try:
|
|
d["name"] = torch.cuda.get_device_name(0)
|
|
d["gb"] = torch.cuda.get_device_properties(0).total_memory / 1024**3
|
|
except: pass
|
|
return d
|
|
|
|
|
|
def tune():
|
|
"""ROCm 全局调优 — fp16累积 (rocBLAS加速) + mem_efficient SDPA"""
|
|
if not is_rocm(): return
|
|
torch.backends.cuda.matmul.allow_fp16_accumulation = True
|
|
if hasattr(torch.backends.cuda, 'enable_mem_efficient_sdp'):
|
|
torch.backends.cuda.enable_mem_efficient_sdp(True)
|
|
if hasattr(torch.backends.cuda, 'enable_flash_sdp'):
|
|
torch.backends.cuda.enable_flash_sdp(False)
|
|
|
|
|
|
def memclr(sync: bool = True):
|
|
if not torch.cuda.is_available(): return
|
|
if sync: torch.cuda.synchronize()
|
|
torch.cuda.empty_cache()
|
|
gc.collect()
|
|
if hasattr(torch.cuda, "ipc_collect"): torch.cuda.ipc_collect()
|
|
|
|
|
|
def memstat() -> dict:
|
|
"""获取当前显存状态"""
|
|
if not torch.cuda.is_available(): return {}
|
|
return {
|
|
"alloc": torch.cuda.memory_allocated(0) / 1024**3,
|
|
"rsvd": torch.cuda.memory_reserved(0) / 1024**3,
|
|
"total": torch.cuda.get_device_properties(0).total_memory / 1024**3,
|
|
}
|
|
|
|
|
|
def _banner(title, g):
|
|
print(f"\n{'='*50}\n {title}\n GPU: {g['name']} | {g['gen']} | {g['gb']:.1f}GB\n{'='*50}")
|
|
|
|
|
|
def _do_cleanup(stage: str, level: str):
|
|
"""执行分级清理。pre=操作前不同步, post=操作后先sync再清"""
|
|
if level == "不做任何清理":
|
|
return
|
|
sync = (stage == "post")
|
|
if level == "单次缓存清理":
|
|
memclr(sync)
|
|
elif level == "双次缓存清理":
|
|
memclr(sync)
|
|
mm.soft_empty_cache()
|
|
memclr(sync)
|
|
gc.collect()
|
|
elif level == "卸载显存模型":
|
|
mm.unload_all_models()
|
|
mm.soft_empty_cache()
|
|
memclr(sync)
|
|
memclr(sync)
|
|
gc.collect()
|
|
|
|
|
|
def _even_chunks(total: int, chunk: int):
|
|
"""均匀分块：避免末尾出现特别小的碎片。
|
|
如果剩余 < chunk/2，合并到前一个 chunk。"""
|
|
if total <= chunk:
|
|
return [(0, total)]
|
|
n_full = total // chunk
|
|
remainder = total % chunk
|
|
if remainder == 0:
|
|
return [(i * chunk, (i + 1) * chunk) for i in range(n_full)]
|
|
# 有余数：如果余数 < chunk/2，把它分给前面的 chunk
|
|
if remainder < chunk // 2:
|
|
chunks = []
|
|
extra_each = remainder // n_full
|
|
extra_rem = remainder % n_full
|
|
pos = 0
|
|
for i in range(n_full):
|
|
sz = chunk + extra_each + (1 if i < extra_rem else 0)
|
|
chunks.append((pos, pos + sz))
|
|
pos += sz
|
|
chunks.append((pos, total)) # 最后小块也保留，但已经很小了，合并逻辑在上面
|
|
# 更简单：直接把余数加到最后一个 chunk
|
|
chunks = [(i * chunk, (i + 1) * chunk) for i in range(n_full - 1)]
|
|
chunks.append(((n_full - 1) * chunk, total))
|
|
return chunks
|
|
else:
|
|
return [(i * chunk, min((i + 1) * chunk, total)) for i in range(n_full + 1)]
|
|
|
|
|
|
class _Noise_EmptyNoise:
|
|
"""空噪波 — 不加噪"""
|
|
def __init__(self):
|
|
self.seed = 0
|
|
|
|
def generate_noise(self, input_latent):
|
|
latent_image = input_latent["samples"]
|
|
if latent_image.is_nested:
|
|
tensors = latent_image.unbind()
|
|
zeros = [torch.zeros(t.shape, dtype=t.dtype, layout=t.layout, device="cpu") for t in tensors]
|
|
return comfy.nested_tensor.NestedTensor(zeros)
|
|
else:
|
|
return torch.zeros(latent_image.shape, dtype=latent_image.dtype, layout=latent_image.layout, device="cpu")
|
|
|
|
|
|
class _Noise_RandomNoise:
|
|
"""随机噪波"""
|
|
def __init__(self, seed):
|
|
self.seed = seed
|
|
|
|
def generate_noise(self, input_latent):
|
|
latent_image = input_latent["samples"]
|
|
batch_inds = input_latent.get("batch_index", None) if "batch_index" in input_latent else None
|
|
return comfy.sample.prepare_noise(latent_image, self.seed, batch_inds)
|
|
|
|
|
|
# ============================================================
|
|
# XB_ROCmKSampler — ROCm 优化采样器
|
|
# ============================================================
|
|
class XB_ROCmKSampler:
|
|
"""
|
|
ROCm 优化采样器 — rocBLAS fp16累积 + mem_efficient SDPA
|
|
配合 XB_SageAttentionAccelerator 使用（先 SageAttn 注入，再接此采样器）。
|
|
"""
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {"required": {
|
|
"model": ("MODEL",),
|
|
"seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff, "control_after_generate": True}),
|
|
"steps": ("INT", {"default": 20, "min": 1, "max": 10000}),
|
|
"cfg": ("FLOAT", {"default": 8.0, "min": 0.0, "max": 100.0, "step": 0.1, "round": 0.01}),
|
|
"sampler": (comfy.samplers.KSampler.SAMPLERS,),
|
|
"scheduler": (comfy.samplers.KSampler.SCHEDULERS,),
|
|
"positive": ("CONDITIONING",),
|
|
"negative": ("CONDITIONING",),
|
|
"latent": ("LATENT",),
|
|
"denoise": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}),
|
|
"cleanup": (["不做任何清理", "单次缓存清理", "双次缓存清理"], {"default": "不做任何清理"}),
|
|
}}
|
|
RETURN_TYPES = ("LATENT",)
|
|
FUNCTION = "go"
|
|
CATEGORY = "XB_ToolBox/ROCm"
|
|
|
|
def go(self, model, seed, steps, cfg, sampler, scheduler,
|
|
positive, negative, latent, denoise, cleanup="不做任何清理"):
|
|
tune()
|
|
out, = nodes.KSampler().sample(model, seed, steps, cfg, sampler, scheduler,
|
|
positive, negative, latent, denoise)
|
|
_do_cleanup("post", cleanup)
|
|
return (out,)
|
|
|
|
|
|
# ============================================================
|
|
# XB_ROCmKSamplerAdvanced — ROCm 优化高级采样器
|
|
# ============================================================
|
|
class XB_ROCmKSamplerAdvanced:
|
|
"""
|
|
ROCm 优化高级采样器 — 支持分段采样 (start/end step)、加噪控制、残留噪声返回
|
|
配合 XB_SageAttentionAccelerator 使用（先 SageAttn 注入，再接此采样器）。
|
|
"""
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {"required": {
|
|
"model": ("MODEL",),
|
|
"add_noise": (["enable", "disable"], {}),
|
|
"noise_seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff, "control_after_generate": True}),
|
|
"steps": ("INT", {"default": 20, "min": 1, "max": 10000}),
|
|
"cfg": ("FLOAT", {"default": 8.0, "min": 0.0, "max": 100.0, "step": 0.1, "round": 0.01}),
|
|
"sampler": (comfy.samplers.KSampler.SAMPLERS,),
|
|
"scheduler": (comfy.samplers.KSampler.SCHEDULERS,),
|
|
"positive": ("CONDITIONING",),
|
|
"negative": ("CONDITIONING",),
|
|
"latent": ("LATENT",),
|
|
"start_at_step": ("INT", {"default": 0, "min": 0, "max": 10000}),
|
|
"end_at_step": ("INT", {"default": 10000, "min": 0, "max": 10000}),
|
|
"return_with_leftover_noise": (["disable", "enable"], {}),
|
|
"cleanup": (["不做任何清理", "单次缓存清理", "双次缓存清理"], {"default": "不做任何清理"}),
|
|
}}
|
|
RETURN_TYPES = ("LATENT",)
|
|
FUNCTION = "go"
|
|
CATEGORY = "XB_ToolBox/ROCm"
|
|
|
|
def go(self, model, add_noise, noise_seed, steps, cfg, sampler, scheduler,
|
|
positive, negative, latent, start_at_step, end_at_step, return_with_leftover_noise, denoise=1.0, cleanup="不做任何清理"):
|
|
tune()
|
|
out, = nodes.KSamplerAdvanced().sample(
|
|
model, add_noise, noise_seed, steps, cfg, sampler, scheduler,
|
|
positive, negative, latent, start_at_step, end_at_step,
|
|
return_with_leftover_noise, denoise=denoise)
|
|
_do_cleanup("post", cleanup)
|
|
return (out,)
|
|
|
|
|
|
# ============================================================
|
|
# XB_ROCmVAEDecode — ROCm VAE 解码器 (空间分块)
|
|
# ============================================================
|
|
class XB_ROCmVAEDecode:
|
|
"""
|
|
ROCm VAE 解码器 — 架构自适应空间分块
|
|
tile=0 → 自动根据GPU架构选最优值 (RX9070XT→768, RX7900→640, ...)
|
|
tile>0 → 使用手动指定的分块大小
|
|
"""
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {"required": {
|
|
"samples": ("LATENT",),
|
|
"vae": ("VAE",),
|
|
"tile": ("INT", {"default": 0, "min": 0, "max": 2048, "step": 64,
|
|
"tooltip": "0=根据GPU架构自动选择 手动可覆盖"}),
|
|
"overlap": ("INT", {"default": 0, "min": 0, "max": 256, "step": 16,
|
|
"tooltip": "0=自动(tile/8)"}),
|
|
"cleanup": (["不做任何清理", "单次缓存清理", "双次缓存清理", "卸载显存模型"], {"default": "单次缓存清理"}),
|
|
}}
|
|
RETURN_TYPES = ("IMAGE",)
|
|
FUNCTION = "go"
|
|
CATEGORY = "XB_ToolBox/ROCm"
|
|
|
|
def go(self, samples, vae, tile, overlap, cleanup):
|
|
g = gpu_info(); tune()
|
|
lat = samples["samples"]
|
|
|
|
if tile == 0: tile = g["tile"]
|
|
if overlap == 0: overlap = max(32, tile // 8)
|
|
|
|
_do_cleanup("pre", cleanup)
|
|
try:
|
|
img = vae.decode_tiled(lat, tile_x=tile, tile_y=tile, overlap=overlap)
|
|
except AttributeError:
|
|
print(f" ⚠️ ROCm VAE Decode: tiled not available, fallback to standard")
|
|
img = vae.decode(lat)
|
|
_do_cleanup("post", cleanup)
|
|
if img.dim() == 5:
|
|
img = img.reshape(-1, img.shape[-3], img.shape[-2], img.shape[-1])
|
|
return (img,)
|
|
|
|
|
|
# ============================================================
|
|
# XB_ROCmVAEEncode — ROCm VAE 编码器 (空间分块)
|
|
# ============================================================
|
|
class XB_ROCmVAEEncode:
|
|
"""
|
|
ROCm VAE 编码器 — 架构自适应空间分块
|
|
tile=0 → 自动, tile>0 → 手动
|
|
"""
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {"required": {
|
|
"pixels": ("IMAGE",),
|
|
"vae": ("VAE",),
|
|
"tile": ("INT", {"default": 0, "min": 0, "max": 2048, "step": 64,
|
|
"tooltip": "0=自动"}),
|
|
"overlap": ("INT", {"default": 0, "min": 0, "max": 256, "step": 16,
|
|
"tooltip": "0=自动"}),
|
|
"cleanup": (["不做任何清理", "单次缓存清理", "双次缓存清理", "卸载显存模型"], {"default": "不做任何清理"}),
|
|
}}
|
|
RETURN_TYPES = ("LATENT",)
|
|
FUNCTION = "go"
|
|
CATEGORY = "XB_ToolBox/ROCm"
|
|
|
|
def go(self, pixels, vae, tile, overlap, cleanup="不做任何清理"):
|
|
g = gpu_info(); tune()
|
|
|
|
if tile == 0: tile = g["tile"]
|
|
if overlap == 0: overlap = max(32, tile // 8)
|
|
|
|
_do_cleanup("pre", cleanup)
|
|
try:
|
|
lat = vae.encode_tiled(pixels, tile_x=tile, tile_y=tile, overlap=overlap)
|
|
except AttributeError:
|
|
print(f" ⚠️ ROCm VAE Encode: tiled not available, fallback to standard")
|
|
lat = vae.encode(pixels)
|
|
_do_cleanup("post", cleanup)
|
|
return ({"samples": lat},)
|
|
|
|
|
|
# ============================================================
|
|
# XB_ROCmVAEDecodeTemporal — ROCm VAE 时空解码器
|
|
# ============================================================
|
|
class XB_ROCmVAEDecodeTemporal:
|
|
"""
|
|
ROCm VAE 解码器 (空间+时间分块) — 视频模型专用
|
|
适用于 Wan/LTX 等视频 VAE latent。
|
|
所有参数=0 时全自动: 空间tile→架构自适应, 时间chunk→显存自适应
|
|
"""
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {"required": {
|
|
"samples": ("LATENT",),
|
|
"vae": ("VAE",),
|
|
"tile": ("INT", {"default": 0, "min": 0, "max": 2048, "step": 64,
|
|
"tooltip": "空间分块 0=自动"}),
|
|
"overlap": ("INT", {"default": 0, "min": 0, "max": 256, "step": 16,
|
|
"tooltip": "空间重叠 0=自动(tile/8, 最小32)"}),
|
|
"t_tile": ("INT", {"default": 0, "min": 0, "max": 1024, "step": 16,
|
|
"tooltip": "时间分块帧数 0=根据显存自适应"}),
|
|
"t_overlap": ("INT", {"default": 0, "min": 0, "max": 64, "step": 4,
|
|
"tooltip": "时间重叠 0=自动(t_tile/16, 最小4)"}),
|
|
"cleanup": (["不做任何清理", "单次缓存清理", "双次缓存清理", "卸载显存模型"], {"default": "单次缓存清理"}),
|
|
}}
|
|
RETURN_TYPES = ("IMAGE",)
|
|
FUNCTION = "go"
|
|
CATEGORY = "XB_ToolBox/ROCm"
|
|
|
|
def go(self, samples, vae, tile, overlap, t_tile, t_overlap, cleanup):
|
|
g = gpu_info(); tune()
|
|
lat = samples["samples"]
|
|
|
|
if lat.dim() == 5:
|
|
B, C, F, H, W = lat.shape; is_vid = True
|
|
elif lat.dim() == 4:
|
|
B, C, H, W = lat.shape; F = 1; is_vid = False
|
|
else:
|
|
raise ValueError(f"unsupported latent dim: {lat.dim()}")
|
|
|
|
if tile == 0: tile = g["tile"]
|
|
spacial_comp = vae.spacial_compression_decode() if hasattr(vae, 'spacial_compression_decode') else 8
|
|
tile_x = tile // spacial_comp
|
|
tile_y = tile // spacial_comp
|
|
if overlap == 0: overlap = max(32, tile // 8)
|
|
overlap_xy = overlap // spacial_comp
|
|
|
|
temporal_comp = vae.temporal_compression_decode() if hasattr(vae, 'temporal_compression_decode') else None
|
|
if temporal_comp is not None and is_vid:
|
|
if t_tile == 0:
|
|
if g["gb"] >= 48: t_tile = 64
|
|
elif g["gb"] >= 24: t_tile = 32
|
|
elif g["gb"] >= 16: t_tile = 16
|
|
else: t_tile = 8
|
|
t_tile_vae = max(2, t_tile // temporal_comp)
|
|
if t_overlap == 0:
|
|
t_overlap = max(4, t_tile // 16)
|
|
t_overlap_vae = max(1, min(t_tile_vae // 2, t_overlap // temporal_comp))
|
|
else:
|
|
t_tile_vae = None
|
|
t_overlap_vae = None
|
|
|
|
_do_cleanup("pre", cleanup)
|
|
try:
|
|
img = vae.decode_tiled(lat, tile_x=tile_x, tile_y=tile_y,
|
|
overlap=overlap_xy,
|
|
tile_t=t_tile_vae, overlap_t=t_overlap_vae)
|
|
except (AttributeError, TypeError):
|
|
print(f" ⚠️ ROCm VAE Decode Temporal: decode_tiled not available, fallback")
|
|
if is_vid and F > 1:
|
|
frames = []
|
|
for f in range(F):
|
|
fl = lat[:, :, f:f+1, :, :].squeeze(2)
|
|
frames.append(vae.decode(fl).unsqueeze(1))
|
|
img = torch.cat(frames, dim=1)
|
|
else:
|
|
img = vae.decode(lat)
|
|
_do_cleanup("post", cleanup)
|
|
if img.dim() == 5:
|
|
img = img.reshape(-1, img.shape[-3], img.shape[-2], img.shape[-1])
|
|
return (img,)
|
|
|
|
|
|
# ============================================================
|
|
# XB_ROCmMemCleaner — ROCm 显存清理与诊断
|
|
# ============================================================
|
|
class XB_ROCmMemCleaner:
|
|
"""
|
|
ROCm 显存清理+诊断 — 独立节点
|
|
显示 GPU 信息 + 清理前后显存对比。
|
|
|
|
验证 fp16_accumulation 加速效果:
|
|
在 Python 控制台运行:
|
|
import torch, time
|
|
a=torch.randn(4096,4096,dtype=torch.float16,device='cuda')
|
|
b=torch.randn(4096,4096,dtype=torch.float16,device='cuda')
|
|
for flag in [False,True]:
|
|
torch.backends.cuda.matmul.allow_fp16_accumulation=flag
|
|
torch.cuda.synchronize(); t0=time.time()
|
|
for _ in range(500): c=a@b
|
|
torch.cuda.synchronize()
|
|
print(f"fp16_accum={flag}: {time.time()-t0:.2f}s")
|
|
"""
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"mode": (["单次缓存清理", "双次缓存清理", "卸载显存模型"], {"default": "双次缓存清理"}),
|
|
},
|
|
"optional": {
|
|
"anything": (_any, {}),
|
|
},
|
|
}
|
|
RETURN_TYPES = (_any,)
|
|
RETURN_NAMES = ("pass-through",)
|
|
FUNCTION = "go"
|
|
CATEGORY = "XB_ToolBox/ROCm"
|
|
OUTPUT_NODE = True
|
|
|
|
def go(self, mode, anything=None):
|
|
g = gpu_info(); tune()
|
|
|
|
before = memstat()
|
|
|
|
if mode == "卸载显存模型":
|
|
mm.unload_all_models()
|
|
mm.soft_empty_cache()
|
|
memclr(sync=True)
|
|
if mode in ("双次缓存清理", "卸载显存模型"):
|
|
memclr(sync=True)
|
|
gc.collect()
|
|
|
|
after = memstat()
|
|
freed_alloc = before['alloc'] - after['alloc'] if before and after else 0
|
|
freed_rsvd = before['rsvd'] - after['rsvd'] if before and after else 0
|
|
|
|
print(f"\033[92m🧹 ROCm {mode} | {g['name']} {g['gb']:.1f}GB | freed {freed_alloc:.2f}+{freed_rsvd:.2f}GB\033[0m")
|
|
return (anything,)
|
|
|
|
|
|
# ============================================================
|
|
# XB_ROCmSamplerCustom — ROCm 自定义采样器
|
|
# ============================================================
|
|
class XB_ROCmSamplerCustom:
|
|
"""
|
|
ROCm 优化自定义采样器 — 配合调度器/采样器/引导器模块化使用
|
|
优化: rocBLAS fp16累积 + mem_efficient SDPA，在采样循环中生效
|
|
"""
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {"required": {
|
|
"model": ("MODEL",),
|
|
"add_noise": ("BOOLEAN", {"default": True}),
|
|
"noise_seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff, "control_after_generate": True}),
|
|
"cfg": ("FLOAT", {"default": 8.0, "min": 0.0, "max": 100.0, "step": 0.1, "round": 0.01}),
|
|
"positive": ("CONDITIONING",),
|
|
"negative": ("CONDITIONING",),
|
|
"sampler": ("SAMPLER",),
|
|
"sigmas": ("SIGMAS",),
|
|
"latent_image": ("LATENT",),
|
|
"cleanup": (["不做任何清理", "单次缓存清理", "双次缓存清理"], {"default": "不做任何清理"}),
|
|
}}
|
|
RETURN_TYPES = ("LATENT", "LATENT")
|
|
RETURN_NAMES = ("output", "denoised_output")
|
|
FUNCTION = "go"
|
|
CATEGORY = "XB_ToolBox/ROCm"
|
|
|
|
def go(self, model, add_noise, noise_seed, cfg, positive, negative, sampler, sigmas, latent_image, cleanup="不做任何清理"):
|
|
tune()
|
|
latent = latent_image
|
|
latent_image_data = latent["samples"]
|
|
latent = latent.copy()
|
|
latent_image_data = comfy.sample.fix_empty_latent_channels(model, latent_image_data, latent.get("downscale_ratio_spacial", None))
|
|
latent["samples"] = latent_image_data
|
|
|
|
if not add_noise:
|
|
noise = _Noise_EmptyNoise().generate_noise(latent)
|
|
else:
|
|
noise = _Noise_RandomNoise(noise_seed).generate_noise(latent)
|
|
|
|
noise_mask = None
|
|
if "noise_mask" in latent:
|
|
noise_mask = latent["noise_mask"]
|
|
|
|
x0_output = {}
|
|
callback = latent_preview.prepare_callback(model, sigmas.shape[-1] - 1, x0_output)
|
|
|
|
disable_pbar = not comfy.utils.PROGRESS_BAR_ENABLED
|
|
samples = comfy.sample.sample_custom(model, noise, cfg, sampler, sigmas, positive, negative, latent_image_data, noise_mask=noise_mask, callback=callback, disable_pbar=disable_pbar, seed=noise_seed)
|
|
|
|
out = latent.copy()
|
|
out.pop("downscale_ratio_spacial", None)
|
|
out["samples"] = samples
|
|
if "x0" in x0_output:
|
|
x0_out = model.model.process_latent_out(x0_output["x0"].cpu())
|
|
if samples.is_nested:
|
|
latent_shapes = [x.shape for x in samples.unbind()]
|
|
x0_out = comfy.nested_tensor.NestedTensor(comfy.utils.unpack_latents(x0_out, latent_shapes))
|
|
out_denoised = latent.copy()
|
|
out_denoised["samples"] = x0_out
|
|
else:
|
|
out_denoised = out
|
|
_do_cleanup("post", cleanup)
|
|
return (out, out_denoised)
|
|
|
|
|
|
# ============================================================
|
|
# XB_ROCmSamplerCustomAdvanced — ROCm 自定义高级采样器
|
|
# ============================================================
|
|
class XB_ROCmSamplerCustomAdvanced:
|
|
"""
|
|
ROCm 优化自定义高级采样器 — 完全模块化：噪波/引导器/采样器/sigma 独立注入
|
|
优化: rocBLAS fp16累积 + mem_efficient SDPA，在采样循环中生效
|
|
"""
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {"required": {
|
|
"noise": ("NOISE",),
|
|
"guider": ("GUIDER",),
|
|
"sampler": ("SAMPLER",),
|
|
"sigmas": ("SIGMAS",),
|
|
"latent_image": ("LATENT",),
|
|
"cleanup": (["不做任何清理", "单次缓存清理", "双次缓存清理"], {"default": "不做任何清理"}),
|
|
}}
|
|
RETURN_TYPES = ("LATENT", "LATENT")
|
|
RETURN_NAMES = ("output", "denoised_output")
|
|
FUNCTION = "go"
|
|
CATEGORY = "XB_ToolBox/ROCm"
|
|
|
|
def go(self, noise, guider, sampler, sigmas, latent_image, cleanup="不做任何清理"):
|
|
tune()
|
|
latent = latent_image
|
|
latent_image_data = latent["samples"]
|
|
latent = latent.copy()
|
|
latent_image_data = comfy.sample.fix_empty_latent_channels(guider.model_patcher, latent_image_data, latent.get("downscale_ratio_spacial", None))
|
|
latent["samples"] = latent_image_data
|
|
|
|
noise_mask = None
|
|
if "noise_mask" in latent:
|
|
noise_mask = latent["noise_mask"]
|
|
|
|
x0_output = {}
|
|
callback = latent_preview.prepare_callback(guider.model_patcher, sigmas.shape[-1] - 1, x0_output)
|
|
|
|
disable_pbar = not comfy.utils.PROGRESS_BAR_ENABLED
|
|
samples = guider.sample(noise.generate_noise(latent), latent_image_data, sampler, sigmas, denoise_mask=noise_mask, callback=callback, disable_pbar=disable_pbar, seed=noise.seed)
|
|
samples = samples.to(comfy.model_management.intermediate_device())
|
|
|
|
out = latent.copy()
|
|
out.pop("downscale_ratio_spacial", None)
|
|
out["samples"] = samples
|
|
if "x0" in x0_output:
|
|
x0_out = guider.model_patcher.model.process_latent_out(x0_output["x0"].cpu())
|
|
if samples.is_nested:
|
|
latent_shapes = [x.shape for x in samples.unbind()]
|
|
x0_out = comfy.nested_tensor.NestedTensor(comfy.utils.unpack_latents(x0_out, latent_shapes))
|
|
out_denoised = latent.copy()
|
|
out_denoised["samples"] = x0_out
|
|
else:
|
|
out_denoised = out
|
|
_do_cleanup("post", cleanup)
|
|
return (out, out_denoised)
|
|
|
|
|
|
__all__ = ['XB_ROCmKSampler', 'XB_ROCmKSamplerAdvanced', 'XB_ROCmSamplerCustom',
|
|
'XB_ROCmSamplerCustomAdvanced', 'XB_ROCmVAEDecode',
|
|
'XB_ROCmVAEEncode', 'XB_ROCmVAEDecodeTemporal', 'XB_ROCmMemCleaner']
|