Update negative_rejection_steering/scripts/negative_rejection_steering_script.py

negative_rejection_steering/scripts/negative_rejection_steering_script.py

@@ -66,15 +66,73 @@ def calc_nrs(x_orig, cond, uncond, sigma, skew, stretch, squash):
         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
+# Хук 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)
 
@@ -82,7 +140,7 @@ 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: Подменная функция
+# Хук 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):
@@ -92,8 +150,77 @@ def hijacked_combine_denoised(self, x_out, conds_list, uncond, cond_scale):
         return sd_samplers_cfg_denoiser.CFGDenoiser.original_combine_denoised_nrs_backup(self, x_out, conds_list, uncond, cond_scale)
 
     try:
-        skew, stretch, squash = self.p._nrs_params
+        # Получаем базовые параметры
+        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)
         
@@ -132,11 +259,11 @@ def hijacked_combine_denoised(self, x_out, conds_list, uncond, cond_scale):
 
 
 # ==============================================================================
-# ЧАСТЬ 3: Интерфейс (UI) с пояснениями
+# ЧАСТЬ 3: Интерфейс (UI) с расширенными настройками
 # ==============================================================================
 class NRSScript(scripts.Script):
     def title(self):
-        return "Negative Rejection Steering"
+        return "Negative Rejection Steering (Enhanced)"
 
     def show(self, is_img2img):
         return scripts.AlwaysVisible
@@ -164,10 +291,15 @@ class NRSScript(scripts.Script):
                 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("<div style='margin-bottom: 0.5em; opacity: 0.8; font-size: 0.9em; border-bottom: 1px solid #444;'>Настройки</div>")
+            # --- ОСНОВНЫЕ ПАРАМЕТРЫ ---
+            gr.HTML("<div style='margin-bottom: 0.5em; opacity: 0.8; font-size: 0.9em; border-bottom: 1px solid #444;'>Основные настройки</div>")
             
             with gr.Row():
                 skew = gr.Slider(
@@ -187,14 +319,249 @@ class NRSScript(scripts.Script):
                 info="0.0 = Макс. эффект. 1.0 = Ослабление (больше деталей, меньше контраста)."
             )
             
-        return [enabled, skew, stretch, squash]
+            # --- 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("<div style='margin: 0.5em 0; font-weight: bold;'>Глобальные настройки (для всех параметров)</div>")
+                    
+                    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("<div style='margin: 0.5em 0; font-weight: bold;'>Индивидуальные настройки</div>")
+                    
+                    # 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):
+    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