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

@@ -1928,6 +1928,64 @@ class CustomHiresFix(scripts.Script):
                 pass
         self._orig_opt_vae_tiling = None
 
+    def _decode_first_stage_safe(self, samples: torch.Tensor) -> torch.Tensor:
+        """
+        Обёртка над processing.decode_first_stage с защитой от OOM:
+          1) Агрессивная очистка VRAM перед декодом (torch_gc + empty_cache).
+          2) Авто-включение VAE tiling при первом OOM и повтор с теми же данными.
+          3) Второй OOM: повтор с clamp(-3,3) — ВНИМАНИЕ: незначительно меняет
+             результат, используется только как последний аварийный шанс.
+          4) NaN guard: если декод вернул NaN — clamp + повтор.
+        Перевод результата на CPU выполняется в вызывающем коде, не здесь.
+        """
+        # Шаг 1: максимально освободить VRAM перед тяжёлым декодом
+        devices.torch_gc()
+        try:
+            torch.cuda.empty_cache()
+        except Exception:
+            pass
+
+        def _do_decode(s: torch.Tensor) -> torch.Tensor:
+            return processing.decode_first_stage(shared.sd_model, s)
+
+        # Шаг 2: попытка декода; при OOM — включаем VAE tiling и повторяем
+        # с теми же данными (результат не меняется)
+        try:
+            decoded = _do_decode(samples)
+        except torch.cuda.OutOfMemoryError:
+            logger.warning(
+                "OOM during decode_first_stage — enabling VAE tiling and retrying "
+                "(same data, result unchanged). Consider enabling 'VAE tiling (low VRAM)' "
+                "in extension settings to avoid this overhead."
+            )
+            devices.torch_gc()
+            try:
+                torch.cuda.empty_cache()
+            except Exception:
+                pass
+            self._set_vae_tiling(True)
+            try:
+                decoded = _do_decode(samples)
+            except torch.cuda.OutOfMemoryError:
+                # Шаг 3: последний аварийный шанс — clamp(-3,3) + повтор.
+                # ВНИМАНИЕ: clamp незначительно изменяет экстремальные значения латента,
+                # что может чуть повлиять на детали картинки. Применяется только при
+                # втором подряд OOM, когда даже VAE tiling не помог.
+                logger.warning(
+                    "OOM again after VAE tiling — clamping latents to [-3, 3] and retrying. "
+                    "This is a last-resort fallback and may slightly alter the output image."
+                )
+                devices.torch_gc()
+                decoded = _do_decode(torch.clamp(samples, -3, 3))
+
+        # Шаг 4: NaN guard (отдельно от OOM — может возникнуть без нехватки памяти)
+        if torch.isnan(decoded).any().item():
+            logger.warning("NaN detected in decoded sample — clamping latents and retrying.")
+            devices.torch_gc()
+            decoded = _do_decode(torch.clamp(samples, -3, 3))
+
+        return decoded
+
 
     # --- safe upscaler wrapper (PIL fallback) ---
     def _resize_with_upscaler(self, img: Image.Image, w: int, h: int, upscaler_name: str) -> Image.Image:
@@ -2558,11 +2616,7 @@ class CustomHiresFix(scripts.Script):
             ).to(devices.dtype_vae)
     
         devices.torch_gc()
-        decoded_sample = processing.decode_first_stage(shared.sd_model, samples)
-        if torch.isnan(decoded_sample).any().item():
-            devices.torch_gc()
-            samples = torch.clamp(samples, -3, 3)
-            decoded_sample = processing.decode_first_stage(shared.sd_model, samples)
+        decoded_sample = self._decode_first_stage_safe(samples)
     
         decoded_sample = torch.clamp((decoded_sample + 1.0) / 2.0, min=0.0, max=1.0).squeeze()
         x_np = 255.0 * np.moveaxis(decoded_sample.to(torch.float32).cpu().numpy(), 0, 2)
@@ -2790,11 +2844,7 @@ class CustomHiresFix(scripts.Script):
             ).to(devices.dtype_vae)
 
         devices.torch_gc()
-        decoded_sample = processing.decode_first_stage(shared.sd_model, samples)
-        if torch.isnan(decoded_sample).any().item():
-            devices.torch_gc()
-            samples = torch.clamp(samples, -3, 3)
-            decoded_sample = processing.decode_first_stage(shared.sd_model, samples)
+        decoded_sample = self._decode_first_stage_safe(samples)
     
         decoded_sample = torch.clamp((decoded_sample + 1.0) / 2.0, min=0.0, max=1.0).squeeze()
         x_np = 255.0 * np.moveaxis(decoded_sample.to(torch.float32).cpu().numpy(), 0, 2)