Update sd-webui-chunk-weights/scripts/chunk_weighting.py

sd-webui-chunk-weights/scripts/chunk_weighting.py

@@ -11,7 +11,7 @@ from modules.script_callbacks import on_app_started, on_script_unloaded
 from modules.ui_components import InputAccordion
 
 # ==============================================================================
-# ЧАСТЬ 1: ЛОГГЕР (Встроено, чтобы избежать ошибок импорта)
+# ЧАСТЬ 1: ЛОГГЕР
 # ==============================================================================
 class ColorCode:
     RESET = "\033[0m"
@@ -47,15 +47,12 @@ if not logger.handlers:
 # ЧАСТЬ 2: СКРИПТ (Логика весов)
 # ==============================================================================
 
-# Поиск целевого класса для патчинга.
-# Мы ищем базовый класс, от которого наследуются и SD1.5, и SDXL эмбеддеры.
 target_classes = []
 
 try:
     from modules.sd_hijack_clip import FrozenCLIPEmbedderWithCustomWordsBase
     target_classes.append(FrozenCLIPEmbedderWithCustomWordsBase)
 except ImportError:
-    # Фолбэк для очень старых версий или если архитектура изменится
     logger.warning("Base class FrozenCLIPEmbedderWithCustomWordsBase not found. Trying individual classes.")
     try:
         from modules.sd_hijack_clip import FrozenCLIPEmbedderWithCustomWords
@@ -102,7 +99,6 @@ class ChunkWeight(scripts.Script):
         if not enable:
             return
 
-        # Парсинг весов
         for v in weights.split(","):
             v = v.strip()
             if not v: continue
@@ -114,18 +110,14 @@ class ChunkWeight(scripts.Script):
 
         p.extra_generation_params["Chunk Weights"] = ", ".join(str(v) for v in WEIGHTS)
 
-        # Сброс кэшей A1111. Это критически важно, иначе изменение весов
-        # не применится к повторным генерациям с тем же промптом.
         p.cached_c = [None, None]
         p.cached_uc = [None, None]
         
-        # Поддержка Hires Fix (Highres. fix кэширует свои кондишены отдельно)
         if hasattr(p, 'cached_hr_c'):
             p.cached_hr_c = [None, None]
             p.cached_hr_uc = [None, None]
 
     def postprocess(self, *args):
-        # Очистка глобальных кэшей класса после генерации для надежности
         StableDiffusionProcessing.cached_c = [None, None]
         StableDiffusionProcessing.cached_uc = [None, None]
         StableDiffusionProcessingTxt2Img.cached_hr_c = [None, None]
@@ -133,12 +125,10 @@ class ChunkWeight(scripts.Script):
 
 
 # ==============================================================================
-# ЧАСТЬ 3: ПАТЧИНГ (Внедрение в ядро)
+# ЧАСТЬ 3: ПАТЧИНГ
 # ==============================================================================
 
 def patch_embedder(cls):
-    """Патчит класс Embedder, добавляя умножение весов."""
-    
     if hasattr(cls, '_chunk_weight_patched'):
         return
 
@@ -149,16 +139,11 @@ def patch_embedder(cls):
     def patched_process_texts(self, texts: List[str]):
         global IS_NEGATIVE_PROMPT, INDEX
 
-        # A1111 передает флаг is_negative_prompt как атрибут списка texts.
-        # Это работает благодаря классу SdConditioning в modules/prompt_parser.py
         if hasattr(texts, "is_negative_prompt"):
             IS_NEGATIVE_PROMPT = texts.is_negative_prompt
         else:
-            # Если атрибут не найден (редкий случай), считаем это позитивным промптом
-            # чтобы веса применились.
             IS_NEGATIVE_PROMPT = False
 
-        # Сбрасываем индекс чанка перед обработкой нового текста
         INDEX = 0
         return original_process_texts(self, texts)
 
@@ -166,11 +151,9 @@ def patch_embedder(cls):
     def patched_process_tokens(self, remade_batch_tokens: list, batch_multipliers: list):
         global INDEX, WEIGHTS, IS_NEGATIVE_PROMPT
 
-        # Если веса не заданы пользователем, работаем как оригинальный метод
         if not WEIGHTS:
             return original_process_tokens(self, remade_batch_tokens, batch_multipliers)
 
-        # Применяем веса только к позитивным промптам
         if INDEX >= 0 and not IS_NEGATIVE_PROMPT:
             batches: int = len(batch_multipliers)
             
@@ -179,31 +162,30 @@ def patch_embedder(cls):
                 logger.debug(f"Applying weight {current_weight}x to Chunk {INDEX}")
                 
                 for b in range(batches):
-                    # batch_multipliers[b] - это список множителей (по умолчанию 1.0) для каждого токена
                     for i in range(len(batch_multipliers[b])):
                          batch_multipliers[b][i] *= current_weight
             else:
-                 # Если весов меньше, чем чанков, остальные получают 1.0 (стандарт)
                  if not ChunkWeight._error_logged:
                     logger.warning(f"Not enough weights provided! Chunk {INDEX} uses default weight 1.0.")
                     ChunkWeight._error_logged = True
             
-            # Увеличиваем индекс, так как process_tokens вызывается для каждого чанка (BREAK) по очереди
             INDEX += 1
 
         return original_process_tokens(self, remade_batch_tokens, batch_multipliers)
 
-    # Применяем подмену методов
     cls.process_texts = patched_process_texts
     cls.process_tokens = patched_process_tokens
     cls._chunk_weight_patched = True
     
-    # Сохраняем оригиналы для возможности отключения
     cls._original_process_texts = original_process_texts
     cls._original_process_tokens = original_process_tokens
 
 
-def apply_patches():
+def apply_patches(*args, **kwargs):
+    """
+    Применяет патчи.
+    Принимает *args, **kwargs, так как A1111 передает (demo, app) в этот коллбек.
+    """
     global PATCHED
     if PATCHED: return
     
@@ -237,6 +219,5 @@ def remove_patches():
     PATCHED = False
     logger.info("Chunk Weight patches removed.")
 
-# Регистрируем коллбэки для автозапуска
 on_app_started(apply_patches)
 on_script_unloaded(remove_patches)