88 lines
3.0 KiB
Python
88 lines
3.0 KiB
Python
"""
|
||
ZImage Model Sampling 节点
|
||
|
||
用于调整 ZImage/Lumina2 模型的采样参数。
|
||
关键区别：ZImage 使用 multiplier=1.0，而 SD3 使用 multiplier=1000。
|
||
"""
|
||
|
||
import torch
|
||
import logging
|
||
import comfy.model_sampling
|
||
|
||
|
||
class ModelSamplingZImage:
|
||
"""
|
||
ZImage/Lumina2 采样参数调整节点
|
||
|
||
与 ModelSamplingSD3 的关键区别：
|
||
- ZImage 使用 multiplier=1.0（timestep 范围 0-1）
|
||
- SD3 使用 multiplier=1000（timestep 范围 0-1000）
|
||
|
||
shift 参数控制噪声调度的偏移（通过 time_snr_shift 函数）：
|
||
- shift=1.0: 无偏移，线性调度
|
||
- shift>1.0: 向高噪声偏移，早期步骤更激进
|
||
- ZImage 默认 shift=3.0
|
||
"""
|
||
|
||
@classmethod
|
||
def INPUT_TYPES(s):
|
||
return {
|
||
"required": {
|
||
"model": ("MODEL",),
|
||
"shift": ("FLOAT", {
|
||
"default": 3.0,
|
||
"min": 0.0,
|
||
"max": 100.0,
|
||
"step": 0.01,
|
||
"tooltip": "噪声调度偏移量。ZImage 默认 3.0。shift=1.0 为线性，>1.0 向高噪声偏移"
|
||
}),
|
||
"multiplier": ("FLOAT", {
|
||
"default": 1.0,
|
||
"min": 0.001,
|
||
"max": 10000.0,
|
||
"step": 0.001,
|
||
"tooltip": "Timestep 乘数。ZImage/AuraFlow=1.0，SD3/Flux=1000"
|
||
}),
|
||
}
|
||
}
|
||
|
||
RETURN_TYPES = ("MODEL",)
|
||
FUNCTION = "patch"
|
||
CATEGORY = "advanced/model"
|
||
DESCRIPTION = "调整 ZImage/Lumina2 模型的采样参数。可设置 shift 和 multiplier。"
|
||
|
||
def patch(self, model, shift, multiplier):
|
||
m = model.clone()
|
||
|
||
sampling_base = comfy.model_sampling.ModelSamplingDiscreteFlow
|
||
sampling_type = comfy.model_sampling.CONST
|
||
|
||
class ModelSamplingAdvanced(sampling_base, sampling_type):
|
||
pass
|
||
|
||
model_sampling = ModelSamplingAdvanced(model.model.model_config)
|
||
model_sampling.set_parameters(shift=shift, multiplier=multiplier)
|
||
m.add_object_patch("model_sampling", model_sampling)
|
||
|
||
# 调试输出
|
||
logging.info(f"[ModelSamplingZImage] Applied: shift={shift}, multiplier={multiplier}")
|
||
logging.info(f"[ModelSamplingZImage] sigma_min={model_sampling.sigma_min:.6f}, sigma_max={model_sampling.sigma_max:.6f}")
|
||
logging.info(f"[ModelSamplingZImage] sigmas[0:5]={model_sampling.sigmas[:5].tolist()}")
|
||
|
||
# 验证 patch 是否成功
|
||
patched_ms = m.get_model_object("model_sampling")
|
||
logging.info(f"[ModelSamplingZImage] Patch verification: patched shift={patched_ms.shift}, multiplier={patched_ms.multiplier}")
|
||
logging.info(f"[ModelSamplingZImage] Patched sigmas[0:5]={patched_ms.sigmas[:5].tolist()}")
|
||
|
||
return (m,)
|
||
|
||
|
||
# 注册节点
|
||
NODE_CLASS_MAPPINGS = {
|
||
"ModelSamplingZImage": ModelSamplingZImage,
|
||
}
|
||
|
||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||
"ModelSamplingZImage": "Model Sampling ZImage",
|
||
}
|