import torch
import gradio as gr
from modules import scripts, script_callbacks, sd_samplers_cfg_denoiser, shared
# ==============================================================================
# ЧАСТЬ 1: Математическое ядро NRS
# ==============================================================================
def calc_nrs(x_orig, cond, uncond, sigma, skew, stretch, squash):
# 1. Проверка режима V-prediction
is_v_pred = False
if hasattr(shared.sd_model, 'parameterization'):
is_v_pred = shared.sd_model.parameterization == "v"
# 2. Подготовка тензора Sigma
if isinstance(sigma, torch.Tensor):
sig_tens = sigma[0]
else:
sig_tens = torch.tensor(sigma, device=cond.device, dtype=cond.dtype)
if sig_tens.dtype != cond.dtype:
sig_tens = sig_tens.to(dtype=cond.dtype)
sig_tens = sig_tens.view(1, 1, 1, 1)
sig_root = (sig_tens ** 2 + 1).sqrt()
# 3. Конвертация EPS -> V
if is_v_pred:
nrs_cond, nrs_uncond = cond, uncond
x_div = None
else:
x_div = x_orig / (sig_tens ** 2 + 1)
factor = sig_tens / sig_root
nrs_cond = x_orig - (x_div - cond * factor)
nrs_uncond = x_orig - (x_div - uncond * factor)
# 4. Векторные операции (Core NRS Math)
def _dot(a, b): return (a * b).sum(dim=1, keepdim=True)
def _nrm2(v): return _dot(v, v)
eps_safe = 1e-6
# Проекция Uncond на Cond
c_dot_c = _nrm2(nrs_cond) + eps_safe
u_dot_c = _dot(nrs_uncond, nrs_cond)
u_on_c = (u_dot_c / c_dot_c) * nrs_cond
# Stretch: Усиление
proj_diff = nrs_cond - u_on_c
stretched = nrs_cond + (stretch * proj_diff)
# Skew: Отклонение
u_rej_c = nrs_uncond - u_on_c
skewed = stretched - (skew * u_rej_c)
# Squash: Нормализация
cond_len = nrs_cond.norm(dim=1, keepdim=True)
nrs_len = skewed.norm(dim=1, keepdim=True) + eps_safe
squash_scale = (1 - squash) + (squash * (cond_len / nrs_len))
x_final = skewed * squash_scale
# 5. Возврат в исходное пространство
if is_v_pred:
return x_final
else:
return (x_div - (x_orig - x_final)) * (sig_root / sig_tens)
# ==============================================================================
# ЧАСТЬ 1.5: Утилиты для управления шагами
# ==============================================================================
def should_apply_at_step(current_step, total_steps, start_step, end_step, start_frac, end_frac, step_mode):
"""
Определяет, нужно ли применять эффект на текущем шаге
Args:
current_step: Текущий шаг сэмплера
total_steps: Общее количество шагов
start_step: Начальный абсолютный шаг
end_step: Конечный абсолютный шаг
start_frac: Начальная доля (0.0-1.0)
end_frac: Конечная доля (0.0-1.0)
step_mode: Режим ("Absolute Steps" или "Fraction of Steps")
Returns:
bool: True если эффект должен применяться
"""
if step_mode == "Absolute Steps":
# Абсолютные шаги
effective_start = max(0, start_step)
effective_end = min(total_steps, end_step) if end_step > 0 else total_steps
return effective_start <= current_step < effective_end
else:
# Доли от общего количества шагов
effective_start = int(total_steps * max(0.0, min(1.0, start_frac)))
effective_end = int(total_steps * max(0.0, min(1.0, end_frac)))
if effective_end == 0:
effective_end = total_steps
return effective_start <= current_step < effective_end
def get_param_value_at_step(base_value, current_step, total_steps, start_step, end_step,
start_frac, end_frac, step_mode, enabled):
"""
Возвращает значение параметра для текущего шага
Если шаг вне диапазона - возвращает 0.0
"""
if not enabled:
return base_value
if should_apply_at_step(current_step, total_steps, start_step, end_step,
start_frac, end_frac, step_mode):
return base_value
else:
return 0.0
# ==============================================================================
# ЧАСТЬ 2: Перехват управления (Hooking)
# ==============================================================================
# Хук 1: Сохраняем Sigma, Input Latent и текущий шаг
def hook_cfg_denoiser_params(params):
if hasattr(params.denoiser, 'p') and getattr(params.denoiser.p, '_nrs_enabled', False):
params.denoiser.p._nrs_current_sigma = params.sigma
params.denoiser.p._nrs_current_x_in = params.x
# Определяем текущий шаг
if hasattr(params, 'sampling_step'):
params.denoiser.p._nrs_current_step = params.sampling_step
elif hasattr(params.denoiser, 'step'):
params.denoiser.p._nrs_current_step = params.denoiser.step
else:
# Fallback: пытаемся определить по sigma
params.denoiser.p._nrs_current_step = getattr(params.denoiser.p, '_nrs_current_step', 0)
script_callbacks.on_cfg_denoiser(hook_cfg_denoiser_params)
# Бэкап оригинала
if not hasattr(sd_samplers_cfg_denoiser.CFGDenoiser, 'original_combine_denoised_nrs_backup'):
sd_samplers_cfg_denoiser.CFGDenoiser.original_combine_denoised_nrs_backup = sd_samplers_cfg_denoiser.CFGDenoiser.combine_denoised
# Хук 2: Подменная функция с поддержкой step control
def hijacked_combine_denoised(self, x_out, conds_list, uncond, cond_scale):
# Проверка активности NRS
if not getattr(self, 'p', None) or not getattr(self.p, '_nrs_enabled', False):
return sd_samplers_cfg_denoiser.CFGDenoiser.original_combine_denoised_nrs_backup(self, x_out, conds_list, uncond, cond_scale)
if not hasattr(self.p, '_nrs_current_sigma') or not hasattr(self.p, '_nrs_current_x_in'):
return sd_samplers_cfg_denoiser.CFGDenoiser.original_combine_denoised_nrs_backup(self, x_out, conds_list, uncond, cond_scale)
try:
# Получаем базовые параметры
base_skew, base_stretch, base_squash = self.p._nrs_params
# Получаем настройки step control
step_control_enabled = getattr(self.p, '_nrs_step_control_enabled', False)
step_control_mode = getattr(self.p, '_nrs_step_control_mode', 'Global')
# Определяем текущий и общий шаги
current_step = getattr(self.p, '_nrs_current_step', 0)
total_steps = getattr(self.p, 'steps', 20)
# Если step control включен, вычисляем эффективные значения параметров
if step_control_enabled:
if step_control_mode == 'Global':
# Глобальный режим - одни настройки для всех
global_settings = getattr(self.p, '_nrs_global_step_settings', {})
start_step = global_settings.get('start_step', 0)
end_step = global_settings.get('end_step', total_steps)
start_frac = global_settings.get('start_frac', 0.0)
end_frac = global_settings.get('end_frac', 1.0)
step_mode = global_settings.get('step_mode', 'Absolute Steps')
# Проверяем, нужно ли применять эффекты
if not should_apply_at_step(current_step, total_steps, start_step, end_step,
start_frac, end_frac, step_mode):
# Вне диапазона - используем fallback
return sd_samplers_cfg_denoiser.CFGDenoiser.original_combine_denoised_nrs_backup(
self, x_out, conds_list, uncond, cond_scale)
skew, stretch, squash = base_skew, base_stretch, base_squash
else:
# Индивидуальный режим - отдельные настройки для каждого параметра
individual_settings = getattr(self.p, '_nrs_individual_step_settings', {})
skew_settings = individual_settings.get('skew', {})
skew = get_param_value_at_step(
base_skew, current_step, total_steps,
skew_settings.get('start_step', 0),
skew_settings.get('end_step', total_steps),
skew_settings.get('start_frac', 0.0),
skew_settings.get('end_frac', 1.0),
skew_settings.get('step_mode', 'Absolute Steps'),
skew_settings.get('enabled', True)
)
stretch_settings = individual_settings.get('stretch', {})
stretch = get_param_value_at_step(
base_stretch, current_step, total_steps,
stretch_settings.get('start_step', 0),
stretch_settings.get('end_step', total_steps),
stretch_settings.get('start_frac', 0.0),
stretch_settings.get('end_frac', 1.0),
stretch_settings.get('step_mode', 'Absolute Steps'),
stretch_settings.get('enabled', True)
)
squash_settings = individual_settings.get('squash', {})
squash = get_param_value_at_step(
base_squash, current_step, total_steps,
squash_settings.get('start_step', 0),
squash_settings.get('end_step', total_steps),
squash_settings.get('start_frac', 0.0),
squash_settings.get('end_frac', 1.0),
squash_settings.get('step_mode', 'Absolute Steps'),
squash_settings.get('enabled', True)
)
else:
# Step control выключен - используем базовые значения
skew, stretch, squash = base_skew, base_stretch, base_squash
# Основная логика NRS
denoised_uncond = x_out[-uncond.shape[0]:]
denoised = torch.clone(denoised_uncond)
x_orig_all = self.p._nrs_current_x_in
x_orig_uncond = x_orig_all[-uncond.shape[0]:]
for i, conds in enumerate(conds_list):
for idx, (cond_index, weight) in enumerate(conds):
current_cond = x_out[cond_index]
# NRS только для основного промпта
if idx == 0:
current_x_orig = x_orig_uncond[i].unsqueeze(0)
c_in = current_cond.unsqueeze(0)
u_in = denoised_uncond[i].unsqueeze(0)
nrs_result = calc_nrs(
current_x_orig, c_in, u_in,
self.p._nrs_current_sigma,
skew, stretch, squash
)
if len(conds) == 1:
denoised[i] = nrs_result.squeeze(0)
else:
delta = nrs_result.squeeze(0) - denoised_uncond[i]
denoised[i] += delta * weight
else:
denoised[i] += (current_cond - denoised_uncond[i]) * (weight * cond_scale)
return denoised
except Exception as e:
print(f"!!! NRS Error (Fallback): {e}")
return sd_samplers_cfg_denoiser.CFGDenoiser.original_combine_denoised_nrs_backup(self, x_out, conds_list, uncond, cond_scale)
# ==============================================================================
# ЧАСТЬ 3: Интерфейс (UI) с расширенными настройками
# ==============================================================================
class NRSScript(scripts.Script):
def title(self):
return "Negative Rejection Steering (Enhanced)"
def show(self, is_img2img):
return scripts.AlwaysVisible
def ui(self, is_img2img):
with gr.Accordion("Negative Rejection Steering", open=False):
with gr.Row():
enabled = gr.Checkbox(label="Включить NRS (Enable)", value=False)
# --- БЛОК СПРАВКИ ---
with gr.Accordion("❓ Как этим пользоваться (Инструкция)", open=False):
gr.Markdown("""
### Что это такое?
**NRS** — это замена стандартному CFG Scale. Вместо простого усиления промпта, он дает вам 3 рычага управления для более точного контроля.
### Параметры:
* **Skew (Сдвиг/Руление):** *Влияет на:* **Композицию и Геометрию**.
Отклоняет генерацию *прочь* от того, что указано в Negative prompt. Если картинка "ломается" или лезут лишние объекты — крутите это.
* **Stretch (Растяжение):** *Влияет на:* **Цвета и Текстуры**.
Усиливает элементы, которые совпадают с вашим Positive prompt. Работает как "газ" для ваших идей.
* **Squash (Сплющивание):** *Влияет на:* **Детализацию и Контраст**.
Ограничитель скорости. При `0.0` эффект максимальный (может "пережарить" картинку). Увеличение добавляет микро-деталей и смягчает цвета.
### 🚀 С чего начать (Советы от автора):
1. **Skew** ≈ Половина вашего обычного CFG (например, **3.0** - **4.0**).
2. **Stretch** ≈ Ваш обычный CFG (например, **6.0** - **7.0**).
3. **Squash** ≈ Оставьте **0.0** для начала.
### ⏱️ Step Control (Управление шагами):
Позволяет применять NRS только на определенных шагах генерации.
- **Global Mode**: Одинаковые настройки для всех трех параметров
- **Individual Mode**: Отдельные настройки для каждого параметра (Skew, Stretch, Squash)
""")
# --- ОСНОВНЫЕ ПАРАМЕТРЫ ---
gr.HTML("Основные настройки
")
with gr.Row():
skew = gr.Slider(
label="Skew (Композиция)",
minimum=-30.0, maximum=30.0, step=0.05, value=4.0,
info="Сила отклонения от Negative prompt. Аналог 'силы' CFG для структуры."
)
stretch = gr.Slider(
label="Stretch (Цвета/Текстура)",
minimum=-30.0, maximum=30.0, step=0.05, value=2.0,
info="Сила притяжения к Positive prompt. Усиливает цвета и стиль."
)
squash = gr.Slider(
label="Squash (Защита от пережарки)",
minimum=0.0, maximum=1.0, step=0.01, value=0.0,
info="0.0 = Макс. эффект. 1.0 = Ослабление (больше деталей, меньше контраста)."
)
# --- STEP CONTROL ---
with gr.Accordion("⏱️ Step Control (Управление шагами)", open=False):
with gr.Row():
step_control_enabled = gr.Checkbox(
label="Включить Step Control",
value=False,
info="Применять NRS только на определенных шагах"
)
step_control_mode = gr.Radio(
label="Режим управления",
choices=["Global", "Individual"],
value="Global",
info="Global: общие настройки | Individual: настройки для каждого параметра"
)
# ГЛОБАЛЬНЫЕ НАСТРОЙКИ
with gr.Group(visible=True) as global_group:
gr.HTML("Глобальные настройки (для всех параметров)
")
global_step_mode = gr.Radio(
label="Режим шагов",
choices=["Absolute Steps", "Fraction of Steps"],
value="Absolute Steps",
info="Absolute: номера шагов | Fraction: доли от общего числа"
)
with gr.Row():
global_start_step = gr.Slider(
label="Start Step (абсолютный)",
minimum=0, maximum=150, step=1, value=0,
visible=True
)
global_end_step = gr.Slider(
label="End Step (абсолютный, 0=конец)",
minimum=0, maximum=150, step=1, value=0,
visible=True
)
with gr.Row():
global_start_frac = gr.Slider(
label="Start (доля от общего числа)",
minimum=0.0, maximum=1.0, step=0.01, value=0.0,
visible=False
)
global_end_frac = gr.Slider(
label="End (доля от общего числа)",
minimum=0.0, maximum=1.0, step=0.01, value=1.0,
visible=False
)
# ИНДИВИДУАЛЬНЫЕ НАСТРОЙКИ
with gr.Group(visible=False) as individual_group:
gr.HTML("Индивидуальные настройки
")
# SKEW
with gr.Accordion("Skew (Композиция) - Step Settings", open=False):
skew_step_enabled = gr.Checkbox(label="Включить управление шагами для Skew", value=True)
skew_step_mode = gr.Radio(
label="Режим",
choices=["Absolute Steps", "Fraction of Steps"],
value="Absolute Steps"
)
with gr.Row():
skew_start_step = gr.Slider(label="Start Step", minimum=0, maximum=150, step=1, value=0, visible=True)
skew_end_step = gr.Slider(label="End Step (0=конец)", minimum=0, maximum=150, step=1, value=0, visible=True)
with gr.Row():
skew_start_frac = gr.Slider(label="Start (fraction)", minimum=0.0, maximum=1.0, step=0.01, value=0.0, visible=False)
skew_end_frac = gr.Slider(label="End (fraction)", minimum=0.0, maximum=1.0, step=0.01, value=1.0, visible=False)
# STRETCH
with gr.Accordion("Stretch (Цвета/Текстура) - Step Settings", open=False):
stretch_step_enabled = gr.Checkbox(label="Включить управление шагами для Stretch", value=True)
stretch_step_mode = gr.Radio(
label="Режим",
choices=["Absolute Steps", "Fraction of Steps"],
value="Absolute Steps"
)
with gr.Row():
stretch_start_step = gr.Slider(label="Start Step", minimum=0, maximum=150, step=1, value=0, visible=True)
stretch_end_step = gr.Slider(label="End Step (0=конец)", minimum=0, maximum=150, step=1, value=0, visible=True)
with gr.Row():
stretch_start_frac = gr.Slider(label="Start (fraction)", minimum=0.0, maximum=1.0, step=0.01, value=0.0, visible=False)
stretch_end_frac = gr.Slider(label="End (fraction)", minimum=0.0, maximum=1.0, step=0.01, value=1.0, visible=False)
# SQUASH
with gr.Accordion("Squash (Защита от пережарки) - Step Settings", open=False):
squash_step_enabled = gr.Checkbox(label="Включить управление шагами для Squash", value=True)
squash_step_mode = gr.Radio(
label="Режим",
choices=["Absolute Steps", "Fraction of Steps"],
value="Absolute Steps"
)
with gr.Row():
squash_start_step = gr.Slider(label="Start Step", minimum=0, maximum=150, step=1, value=0, visible=True)
squash_end_step = gr.Slider(label="End Step (0=конец)", minimum=0, maximum=150, step=1, value=0, visible=True)
with gr.Row():
squash_start_frac = gr.Slider(label="Start (fraction)", minimum=0.0, maximum=1.0, step=0.01, value=0.0, visible=False)
squash_end_frac = gr.Slider(label="End (fraction)", minimum=0.0, maximum=1.0, step=0.01, value=1.0, visible=False)
# Переключение видимости групп при изменении режима
def update_mode_visibility(mode):
return {
global_group: gr.update(visible=(mode == "Global")),
individual_group: gr.update(visible=(mode == "Individual"))
}
step_control_mode.change(
fn=update_mode_visibility,
inputs=[step_control_mode],
outputs=[global_group, individual_group]
)
# Переключение между абсолютными и дробными шагами (GLOBAL)
def update_global_step_inputs(mode):
is_absolute = (mode == "Absolute Steps")
return {
global_start_step: gr.update(visible=is_absolute),
global_end_step: gr.update(visible=is_absolute),
global_start_frac: gr.update(visible=not is_absolute),
global_end_frac: gr.update(visible=not is_absolute)
}
global_step_mode.change(
fn=update_global_step_inputs,
inputs=[global_step_mode],
outputs=[global_start_step, global_end_step, global_start_frac, global_end_frac]
)
# Переключение для SKEW
def update_skew_step_inputs(mode):
is_absolute = (mode == "Absolute Steps")
return {
skew_start_step: gr.update(visible=is_absolute),
skew_end_step: gr.update(visible=is_absolute),
skew_start_frac: gr.update(visible=not is_absolute),
skew_end_frac: gr.update(visible=not is_absolute)
}
skew_step_mode.change(
fn=update_skew_step_inputs,
inputs=[skew_step_mode],
outputs=[skew_start_step, skew_end_step, skew_start_frac, skew_end_frac]
)
# Переключение для STRETCH
def update_stretch_step_inputs(mode):
is_absolute = (mode == "Absolute Steps")
return {
stretch_start_step: gr.update(visible=is_absolute),
stretch_end_step: gr.update(visible=is_absolute),
stretch_start_frac: gr.update(visible=not is_absolute),
stretch_end_frac: gr.update(visible=not is_absolute)
}
stretch_step_mode.change(
fn=update_stretch_step_inputs,
inputs=[stretch_step_mode],
outputs=[stretch_start_step, stretch_end_step, stretch_start_frac, stretch_end_frac]
)
# Переключение для SQUASH
def update_squash_step_inputs(mode):
is_absolute = (mode == "Absolute Steps")
return {
squash_start_step: gr.update(visible=is_absolute),
squash_end_step: gr.update(visible=is_absolute),
squash_start_frac: gr.update(visible=not is_absolute),
squash_end_frac: gr.update(visible=not is_absolute)
}
squash_step_mode.change(
fn=update_squash_step_inputs,
inputs=[squash_step_mode],
outputs=[squash_start_step, squash_end_step, squash_start_frac, squash_end_frac]
)
return [
enabled, skew, stretch, squash,
step_control_enabled, step_control_mode,
# Global
global_step_mode, global_start_step, global_end_step, global_start_frac, global_end_frac,
# Skew
skew_step_enabled, skew_step_mode, skew_start_step, skew_end_step, skew_start_frac, skew_end_frac,
# Stretch
stretch_step_enabled, stretch_step_mode, stretch_start_step, stretch_end_step, stretch_start_frac, stretch_end_frac,
# Squash
squash_step_enabled, squash_step_mode, squash_start_step, squash_end_step, squash_start_frac, squash_end_frac
]
def process(self, p, enabled, skew, stretch, squash,
step_control_enabled, step_control_mode,
global_step_mode, global_start_step, global_end_step, global_start_frac, global_end_frac,
skew_step_enabled, skew_step_mode, skew_start_step, skew_end_step, skew_start_frac, skew_end_frac,
stretch_step_enabled, stretch_step_mode, stretch_start_step, stretch_end_step, stretch_start_frac, stretch_end_frac,
squash_step_enabled, squash_step_mode, squash_start_step, squash_end_step, squash_start_frac, squash_end_frac):
p._nrs_enabled = enabled
p._nrs_params = (skew, stretch, squash)
p._nrs_step_control_enabled = step_control_enabled
p._nrs_step_control_mode = step_control_mode
# Сохраняем глобальные настройки
p._nrs_global_step_settings = {
'step_mode': global_step_mode,
'start_step': global_start_step,
'end_step': global_end_step,
'start_frac': global_start_frac,
'end_frac': global_end_frac
}
# Сохраняем индивидуальные настройки
p._nrs_individual_step_settings = {
'skew': {
'enabled': skew_step_enabled,
'step_mode': skew_step_mode,
'start_step': skew_start_step,
'end_step': skew_end_step,
'start_frac': skew_start_frac,
'end_frac': skew_end_frac
},
'stretch': {
'enabled': stretch_step_enabled,
'step_mode': stretch_step_mode,
'start_step': stretch_start_step,
'end_step': stretch_end_step,
'start_frac': stretch_start_frac,
'end_frac': stretch_end_frac
},
'squash': {
'enabled': squash_step_enabled,
'step_mode': squash_step_mode,
'start_step': squash_start_step,
'end_step': squash_end_step,
'start_frac': squash_start_frac,
'end_frac': squash_end_frac
}
}
if enabled:
sd_samplers_cfg_denoiser.CFGDenoiser.combine_denoised = hijacked_combine_denoised
# Инициализируем счетчик шагов
p._nrs_current_step = 0
def postprocess(self, p, processed, *args):
pass