Update !1-custom-hires-fix-mod-for-automatic1111-2.9.5/scripts/!!custom_hires_fix.py

!1-custom-hires-fix-mod-for-automatic1111-2.9.5/scripts/!!custom_hires_fix.py

@@ -45,10 +45,10 @@ def _safe_import(modname: str, pipname: str | None = None) -> bool:
         return False
 
 _safe_import("omegaconf", "omegaconf")
-_safe_import("kornia", "kornia")
-_safe_import("k_diffusion", "k-diffusion")
-_safe_import("skimage", "scikit-image")
-_safe_import("cv2", "opencv-python")
+_HAS_KORNIA  = _safe_import("kornia",     "kornia")
+_HAS_KDIFF   = _safe_import("k_diffusion", "k-diffusion")
+_HAS_SKIMAGE = _safe_import("skimage",    "scikit-image")
+_HAS_CV2     = _safe_import("cv2",        "opencv-python")
 
 try:
     from omegaconf import OmegaConf, DictConfig  # type: ignore
@@ -62,21 +62,23 @@ except Exception:  # graceful fallback if OmegaConf not available
         @staticmethod
         def create(obj):
             return DictConfig(obj)
-try:
-    import kornia  # type: ignore
-    _HAS_KORNIA = True
-except Exception as e:
-    logger.warning(e)
+if _HAS_KORNIA:
+    try:
+        import kornia  # type: ignore
+    except Exception:
+        kornia = None
+        _HAS_KORNIA = False
+else:
     kornia = None
-    _HAS_KORNIA = False
 
-try:
-    import k_diffusion as K  # type: ignore
-    _HAS_KDIFF = True
-except Exception as e:
-    logger.warning(e)
-    _HAS_KDIFF = False
+if _HAS_KDIFF:
+    try:
+        import k_diffusion as K  # type: ignore
+    except Exception as e:
+        logger.warning(e)
+        _HAS_KDIFF = False
 
+if not _HAS_KDIFF:
     class _KStub:
         class sampling:
             @staticmethod
@@ -85,8 +87,6 @@ except Exception as e:
                 Простейшая замена: логарифмическое пространство от sigma_max до sigma_min
                 + финальный нулевой шаг, как ожидает WebUI.
                 """
-                import math
-                import torch
                 n = int(n)
                 sigma_min = float(sigma_min)
                 sigma_max = float(sigma_max)
@@ -108,7 +108,6 @@ except Exception as e:
                 """
                 Экспоненциальный (геометрический) ряд между sigma_max и sigma_min + 0.0.
                 """
-                import torch
                 n = int(n)
                 sigma_min = float(sigma_min)
                 sigma_max = float(sigma_max)
@@ -158,27 +157,49 @@ def _interpolate_latent(tensor: torch.Tensor, size_hw: tuple[int, int], mode_nam
         # Старые версии — без параметра antialias
         return F.interpolate(tensor, size=size_hw, **kwargs)
 
-try:
-    from skimage.exposure import match_histograms, equalize_adapthist  # type: ignore
-    from skimage import color as skcolor  # type: ignore
-    _SKIMAGE_OK = True
-except Exception as e:
-    logger.warning(e)
-    _SKIMAGE_OK = False
+_SKIMAGE_OK = False
+if _HAS_SKIMAGE:
+    try:
+        from skimage.exposure import match_histograms, equalize_adapthist  # type: ignore
+        from skimage import color as skcolor  # type: ignore
+        _SKIMAGE_OK = True
+    except Exception as e:
+        logger.warning(f"skimage import failed: {e}")
 
 # OpenCV (optional)
-try:
-    import cv2  # type: ignore
-    _CV2_OK = True
-except Exception as e:
-    logger.warning(e)
-    _CV2_OK = False
+_CV2_OK = False
+if _HAS_CV2:
+    try:
+        import cv2  # type: ignore
+        _CV2_OK = True
+    except Exception as e:
+        logger.warning(f"cv2 import failed: {e}")
 
 QUOTE_SINGLE_TO_DOUBLE = {ord("'"): ord('"')}
 config_path = (Path(__file__).parent.resolve() / "../config.yaml").resolve()
 
 
 class CustomHiresFix(scripts.Script):
+    """Two-stage img2img upscaling with optional latent mixing and prompt overrides.
+
+    Features:
+    - Ratio/width/height or Megapixels target (+ quick MP buttons)
+    - Compact preset panel (global presets)
+    - Separate steps for 1st/2nd pass
+    - Per-pass sampler + scheduler
+    - CFG base + optional delta on 2nd pass
+    - Reuse seed/noise on 2nd pass
+    - Conditioning cache (LRU) with capacity
+    - Second-pass prompt (append/replace)
+    - Per-pass LoRA weight scaling
+    - Seamless tiling (+ overlap)
+    - VAE tiling toggle (low VRAM)
+    - Color match to original (strength) with presets
+    - Post-FX presets: CLAHE (local contrast), Unsharp Mask
+    - PNG-info serialization + paste support
+    - Final ×4 upscale (optional), with its own upscaler and tiling overlap for seam-safe stitching.
+    """
+
     def _get_free_vram_bytes(self):
         """Return free VRAM bytes if available; else None."""
         try:
@@ -213,11 +234,11 @@ class CustomHiresFix(scripts.Script):
         except Exception as e:
             # 1) PyTorch<2.0: попробуем без antialias
             if "antialias" in kwargs:
-                aa = kwargs.pop("antialias")
+                kwargs.pop("antialias")
                 try:
                     return F.interpolate(*args, **kwargs)
                 except Exception:
-                    kwargs["antialias"] = aa  # восстановим для следующих веток
+                    pass
 
             # 2) Старые torch: 'nearest-exact' не поддерживается
             if kwargs.get("mode") == "nearest-exact":
@@ -237,38 +258,22 @@ class CustomHiresFix(scripts.Script):
             except Exception:
                 pass
 
-            logger.warning(f"_interp fallback: {e}")
+            # Все фолбэки исчерпаны — возвращаем исходный тензор без ресайза
+            logger.warning(f"_interp: all fallbacks failed ({e}), returning tensor unchanged (shape may mismatch)")
             try:
-                return args[0]  # мягкий фолбэк без ресайза
+                return args[0]
             except Exception:
                 return None
 
 
-    """Two-stage img2img upscaling with optional latent mixing and prompt overrides.
-
-    Features:
-    - Ratio/width/height or Megapixels target (+ quick MP buttons)
-    - Compact preset panel (global presets)
-    - Separate steps for 1st/2nd pass
-    - Per-pass sampler + scheduler
-    - CFG base + optional delta on 2nd pass
-    - Reuse seed/noise on 2nd pass
-    - Conditioning cache (LRU) with capacity
-    - Second-pass prompt (append/replace)
-    - Per-pass LoRA weight scaling
-    - Seamless tiling (+ overlap)
-    - VAE tiling toggle (low VRAM)
-    - Color match to original (strength) with presets
-    - Post-FX presets: CLAHE (local contrast), Unsharp Mask
-    - PNG-info serialization + paste support
-    - Final ×4 upscale (optional), with its own upscaler and tiling overlap for seam-safe stitching.
-    """
     def __init__(self):
         super().__init__()
         # Load or init config
+        _config_loaded_from_disk = False
         if config_path.exists():
             try:
                 self.config: DictConfig = OmegaConf.load(str(config_path)) or OmegaConf.create({})  # type: ignore
+                _config_loaded_from_disk = True
             except Exception as e:
                 logger.warning(e)
                 self.config = OmegaConf.create({})  # type: ignore
@@ -326,6 +331,13 @@ class CustomHiresFix(scripts.Script):
         # Cache fingerprint
         self._last_cache_fp = None
 
+        # Dirty flag: True means config needs to be written to disk.
+        # False when config was just loaded from disk (already in sync).
+        # True on fresh install (no config file yet) so defaults get persisted on first run.
+        self._config_dirty = not _config_loaded_from_disk
+        # JSON snapshot of config taken at start of postprocess_image for dirty detection.
+        self._config_snapshot = None
+
     def _get_model_fingerprint(self) -> str | None:
         """Build a fingerprint of current model/checkpoint + clip_skip to detect changes."""
         parts = []
@@ -1156,7 +1168,8 @@ class CustomHiresFix(scripts.Script):
         """
         Преобразует подпись из UI в аргументы torch.nn.functional.interpolate:
         (mode, antialias). Поддерживает: "nearest", "nearest-exact", "bilinear",
-        "bicubic", "lanczos", "area" и варианты с суффиксом "-antialiased".
+        "bicubic", "area" и варианты с суффиксом "-antialiased".
+        Примечание: "lanczos" не поддерживается F.interpolate — маппится на bicubic+antialias.
         Для старых torch, где "nearest-exact" недоступен, _interp уже содержит фолбэк.
         """
         try:
@@ -1164,7 +1177,11 @@ class CustomHiresFix(scripts.Script):
             aa = lab.endswith("-antialiased")
             if aa:
                 lab = lab.replace("-antialiased", "")
-            if lab in ("nearest", "nearest-exact", "bilinear", "bicubic", "lanczos", "area"):
+            # lanczos не поддерживается torch.nn.functional.interpolate →
+            # маппим на bicubic с antialias для схожего качества
+            if lab == "lanczos":
+                return "bicubic", True
+            if lab in ("nearest", "nearest-exact", "bilinear", "bicubic", "area"):
                 return lab, aa
             return "bicubic", aa
         except Exception:
@@ -1387,10 +1404,23 @@ class CustomHiresFix(scripts.Script):
                           # NEW
                           noise_cache_max):
         if not enable:
+            # Persist the disabled state so next startup respects it
+            self.config["enable"] = False
+            self._config_dirty = True
+            self._save_config()
             return
 
         # Save config chosen in UI
         self.pp = pp
+        # Snapshot current config BEFORE writing UI values — used below to detect real changes.
+        try:
+            import json as _j
+            self._config_snapshot = _j.dumps(
+                {k: (v if not hasattr(v, '__iter__') or isinstance(v, str)
+                     else list(v)) for k, v in self.config.items()},
+                sort_keys=True, default=str)
+        except Exception:
+            self._config_snapshot = None
         self.config["enable"] = bool(enable)
         self.config["ratio"] = float(ratio)
         self.config["width"] = int(width)
@@ -1513,6 +1543,19 @@ class CustomHiresFix(scripts.Script):
             logger.warning(_e)
         # Обновить PNG-info уже с актуальным self.config
         p.extra_generation_params["Custom Hires Fix"] = self.create_infotext(p)
+        # Mark config dirty only if something actually changed vs the snapshot taken
+        # at the start of this function (set earlier via _config_snapshot).
+        try:
+            import json as _j
+            new_snap = _j.dumps(
+                {k: (v if not hasattr(v, '__iter__') or isinstance(v, str)
+                     else list(v)) for k, v in self.config.items()},
+                sort_keys=True, default=str)
+            if new_snap != getattr(self, "_config_snapshot", None):
+                self._config_dirty = True
+                self._config_snapshot = new_snap
+        except Exception:
+            self._config_dirty = True  # safe fallback
 
         # Короткий сводный лог запуска — удобно для багрепортов
         try:
@@ -1678,7 +1721,11 @@ class CustomHiresFix(scripts.Script):
 
     # ---- Helpers ----
     def _save_config(self):
-        """Сохранение конфигурации: YAML (если доступно) или JSON с явными предупреждениями."""
+        """Сохранение конфигурации: YAML (если доступно) или JSON с явными предупреждениями.
+        Запись на диск происходит только если конфиг был изменён из UI (_config_dirty=True).
+        """
+        if not getattr(self, "_config_dirty", True):
+            return  # nothing changed since last save
         try:
             config_path.parent.mkdir(parents=True, exist_ok=True)
         except Exception as e:
@@ -1708,8 +1755,11 @@ class CustomHiresFix(scripts.Script):
                 with open(json_path, "w", encoding="utf-8") as f:
                     _json.dump(data, f, ensure_ascii=False, indent=2)
                 logger.info(f"Config saved to JSON: {json_path}")
+                saved = True
             except Exception as e:
                 logger.warning(f"Config save failed: {e}")
+        if saved:
+            self._config_dirty = False
 
 
     def _vae_down_factor(self) -> int:
@@ -1948,15 +1998,10 @@ class CustomHiresFix(scripts.Script):
         scaled_prompt = self._scale_lora_in_prompt(base_prompt, self._current_lora_factor)
     
         clip_skip = int(self.config.get("clip_skip", 0))
-    
-        # Cache lookup
-        cache_key = self._cond_key(width, height, steps_for_cond, scaled_prompt, negative_base, clip_skip)
-        cached = self._cond_cache_get(cache_key)
-        if cached is not None:
-            self.cond, self.uncond = cached
-            return
-    
-        # Parse extra networks and build cond
+
+        # Parse and activate extra networks (LoRA etc.) BEFORE cache lookup.
+        # This must happen unconditionally — even on cache hit the model weights
+        # need to be patched, because extras are deactivated at the end of each run.
         prompt_text = scaled_prompt
         if not getattr(self.p, "disable_extra_networks", False):
             try:
@@ -1967,10 +2012,15 @@ class CustomHiresFix(scripts.Script):
                         self._activated_extras.append(extra)
                     except Exception as e:
                         logger.warning(e)
-                        pass
             except Exception as e:
                 logger.warning(e)
-                pass
+
+        # Cache lookup — cond/uncond tensors can be reused, but extras must be active (done above)
+        cache_key = self._cond_key(width, height, steps_for_cond, scaled_prompt, negative_base, clip_skip)
+        cached = self._cond_cache_get(cache_key)
+        if cached is not None:
+            self.cond, self.uncond = cached
+            return
     
         if width and height and hasattr(prompt_parser, "SdConditioning"):
             c = prompt_parser.SdConditioning([prompt_text], False, width, height)
@@ -2273,6 +2323,13 @@ class CustomHiresFix(scripts.Script):
         """
         cb_ref = None
         limit_steps = int(max(1, limit_steps))
+
+        # Save interrupt state so we can reliably restore it even if we set it ourselves.
+        try:
+            prev_interrupt = bool(getattr(shared.state, "interrupt", False))
+        except Exception:
+            prev_interrupt = False
+
         try:
             def _cb(params):
                 try:
@@ -2304,12 +2361,26 @@ class CustomHiresFix(scripts.Script):
         finally:
             if cb_ref is not None:
                 try:
-                    # В новых WebUI есть явная отписка; если нет — игнор
+                    # В новых WebUI есть явная отписка; если нет — пробуем ручное удаление
                     remove = getattr(script_callbacks, "remove_callbacks_for_function", None)
                     if callable(remove):
                         remove(cb_ref)
+                    else:
+                        # Ручной fallback: ищем колбэк в карте и удаляем
+                        try:
+                            cbs = getattr(script_callbacks, "callback_map", {}).get("callbacks_cfg_denoiser", [])
+                            for cb in list(cbs):
+                                if getattr(cb, "callback", None) is cb_ref:
+                                    cbs.remove(cb)
+                        except Exception:
+                            pass
                 except Exception as e:
                     logger.warning(f"stop-hook remove warning: {e}")
+            # Restore interrupt flag to avoid leaking into subsequent generations
+            try:
+                shared.state.interrupt = prev_interrupt
+            except Exception:
+                pass
 
     def _first_pass(self, x: Image.Image) -> Image.Image:
         # Determine target size
@@ -2416,15 +2487,10 @@ class CustomHiresFix(scripts.Script):
             steps = min(steps, sfirst)
         # NEW: LRU-кэш шума для первой стадии
         try:
-             try:
-                 seeds = getattr(self.p, "seeds", None)
-                 key_seed = seeds[0] if (seeds and len(seeds) > 0) else getattr(self.p, "seed", None)
-             except (IndexError, TypeError, AttributeError) as e:
-                 logger.warning(f"Seed extraction warning: {e}")
-                 key_seed = None
-
-        except Exception as e:
-            logger.warning(e)
+            seeds = getattr(self.p, "seeds", None)
+            key_seed = seeds[0] if (seeds and len(seeds) > 0) else getattr(self.p, "seed", None)
+        except (IndexError, TypeError, AttributeError) as e:
+            logger.warning(f"Seed extraction warning: {e}")
             key_seed = None
         use_cache = bool(self.config.get("reuse_noise_cache", True))
         sam_name = str(self.config.get("sampler_first", self.config.get("sampler", "")))
@@ -2636,14 +2702,10 @@ class CustomHiresFix(scripts.Script):
             else:
                 # NEW: если размеры не совпали — попробуем кэш
                 try:
-                     try:
-                         seeds = getattr(self.p, "seeds", None)
-                         key_seed = seeds[0] if (seeds and len(seeds) > 0) else getattr(self.p, "seed", None)
-                     except (IndexError, TypeError, AttributeError) as e:
-                         logger.warning(f"Seed extraction warning: {e}")
-                         key_seed = None
-                except Exception as e:
-                    logger.warning(e)
+                    seeds = getattr(self.p, "seeds", None)
+                    key_seed = seeds[0] if (seeds and len(seeds) > 0) else getattr(self.p, "seed", None)
+                except (IndexError, TypeError, AttributeError) as e:
+                    logger.warning(f"Seed extraction warning: {e}")
                     key_seed = None
                 use_cache = bool(self.config.get("reuse_noise_cache", True))
                 sam_name = str(self.config.get("sampler_second", self.config.get("sampler", "")))
@@ -2665,14 +2727,10 @@ class CustomHiresFix(scripts.Script):
         else:
             # NEW: обычный путь — используем LRU-кэш
             try:
-                try:
-                    seeds = getattr(self.p, "seeds", None)
-                    key_seed = seeds[0] if (seeds and len(seeds) > 0) else getattr(self.p, "seed", None)
-                except (IndexError, TypeError, AttributeError) as e:
-                    logger.warning(f"Seed extraction warning: {e}")
-                    key_seed = None
-            except Exception as e:
-                logger.warning(e)
+                seeds = getattr(self.p, "seeds", None)
+                key_seed = seeds[0] if (seeds and len(seeds) > 0) else getattr(self.p, "seed", None)
+            except (IndexError, TypeError, AttributeError) as e:
+                logger.warning(f"Seed extraction warning: {e}")
                 key_seed = None
             use_cache = bool(self.config.get("reuse_noise_cache", True))
             sam_name = str(self.config.get("sampler_second", self.config.get("sampler", "")))
@@ -2905,6 +2963,7 @@ class CustomHiresFix(scripts.Script):
                 crop = img.crop((x0, y0, x1, y1))
     
                 # Reflect-padding по краям кадра, если запрошено
+                pad_left = pad_top = pad_right = pad_bottom = 0
                 if use_reflect and (x0 == 0 or y0 == 0 or x1 == W or y1 == H):
                     arr = np.array(crop)
                     # считаем недостающий «выход» за границы слева/сверху
@@ -2918,6 +2977,19 @@ class CustomHiresFix(scripts.Script):
                 # Апскейлим тайл
                 up = _upscale(crop, final_scale)
                 up_np = np.array(up).astype(np.float32)
+
+                # Обрезаем reflect-padding после апскейла, чтобы размер и позиция вставки
+                # соответствовали реальным координатам холста (без отражённых краёв).
+                if pad_left or pad_top or pad_right or pad_bottom:
+                    pt = int(round(pad_top    * final_scale))
+                    pb = int(round(pad_bottom * final_scale))
+                    pl = int(round(pad_left   * final_scale))
+                    pr = int(round(pad_right  * final_scale))
+                    h_up, w_up = up_np.shape[:2]
+                    # Оставляем только центральную (реальную) часть
+                    up_np = up_np[pt : h_up - pb if pb else h_up,
+                                  pl : w_up - pr if pr else w_up]
+
                 uh, uw = up_np.shape[0], up_np.shape[1]
     
                 # Координаты вставки на апскейленном полотне
@@ -2980,8 +3052,9 @@ class CustomHiresFix(scripts.Script):
                         cap = int(self.config.get("cn_proc_res_cap", 1024))
                         cap = max(256, min(4096, cap))
                         unit.processor_res = max(256, min(cap, min_side))
-                        if getattr(unit, "image", None) is None:
-                            unit.image = image_np
+                        # Always update unit.image so the 2nd pass uses the current reference,
+                        # not the stale one set during the 1st pass.
+                        unit.image = image_np
                 self.p.width = image_np.shape[1]
                 self.p.height = image_np.shape[0]
             except Exception as e:
@@ -3013,8 +3086,17 @@ def parse_infotext(infotext, params):
                 try:
                     data = json.loads(block)
                 except json.JSONDecodeError:
-                    # Fallback для старых строк с одинарными кавычками
-                    data = json.loads(block.translate(QUOTE_SINGLE_TO_DOUBLE))
+                    # Fallback для старых строк с одинарными кавычками.
+                    # Используем ast.literal_eval вместо translate(), чтобы не ломать
+                    # промпты с апострофами (напр. "it's a dog" → невалидный JSON после translate).
+                    try:
+                        import ast
+                        data = ast.literal_eval(block)
+                        if not isinstance(data, dict):
+                            raise ValueError("not a dict")
+                    except Exception:
+                        # Последний резерв — старый translate-хак, только если ast тоже упал
+                        data = json.loads(block.translate(QUOTE_SINGLE_TO_DOUBLE))
         else:
             logger.warning(f"Unexpected infotext type: {type(block)}")
             return
@@ -3146,10 +3228,6 @@ def parse_infotext(infotext, params):
         logger.warning(e)
         return
 # Register paste-params hook (guard against double registration)
-try:
-    _INFOTEXT_HOOK_REGISTERED
-except NameError:
-    _INFOTEXT_HOOK_REGISTERED = False
 if not _INFOTEXT_HOOK_REGISTERED:
     script_callbacks.on_infotext_pasted(parse_infotext)
-    _INFOTEXT_HOOK_REGISTERED = True
+    _INFOTEXT_HOOK_REGISTERED = True