Offload patched weight backups from VRAM

torch_compile_lora_safe.py

@@ -23,6 +23,85 @@ except Exception:
     _skip_torch_compile_dict = None
 
 
+_CPU_BACKUP_PATCHER_CLASS_CACHE = {}
+
+
+class _CpuPatchedWeightBackupMixin:
+    """Route rollback copies for patched weights to CPU without changing model residency."""
+
+    def patch_weight_to_device(
+        self,
+        key,
+        device_to=None,
+        inplace_update=False,
+        return_weight=False,
+        force_cast=False,
+    ):
+        patches = getattr(self, "patches", None)
+        backup = getattr(self, "backup", None)
+        offload_device = getattr(self, "offload_device", None)
+        should_route_backup = (
+            isinstance(patches, dict)
+            and key in patches
+            and isinstance(backup, dict)
+            and key not in backup
+            and not return_weight
+            and getattr(offload_device, "type", None) == "cuda"
+        )
+        if not should_route_backup:
+            return super().patch_weight_to_device(
+                key,
+                device_to=device_to,
+                inplace_update=inplace_update,
+                return_weight=return_weight,
+                force_cast=force_cast,
+            )
+
+        # ComfyUI normally stores rollback weights on self.offload_device.
+        # HIGH_VRAM makes that the same CUDA device as the live model, retaining
+        # both the original and patched destination matrices. Override only for
+        # the duration of backup creation, then restore normal residency policy.
+        self.offload_device = torch.device("cpu")
+        try:
+            return super().patch_weight_to_device(
+                key,
+                device_to=device_to,
+                inplace_update=inplace_update,
+                return_weight=return_weight,
+                force_cast=force_cast,
+            )
+        finally:
+            self.offload_device = offload_device
+
+
+def _enable_cpu_patched_weight_backups(model) -> bool:
+    """Install the backup-device shim on one ModelPatcher clone."""
+    model_class = model.__class__
+    if issubclass(model_class, _CpuPatchedWeightBackupMixin):
+        return True
+
+    wrapped_class = _CPU_BACKUP_PATCHER_CLASS_CACHE.get(model_class)
+    if wrapped_class is None:
+        wrapped_class = type(
+            f"LoRASafeCpuBackup_{model_class.__name__}",
+            (_CpuPatchedWeightBackupMixin, model_class),
+            {"__module__": __name__},
+        )
+        _CPU_BACKUP_PATCHER_CLASS_CACHE[model_class] = wrapped_class
+
+    try:
+        model.__class__ = wrapped_class
+    except TypeError as exc:
+        _LOG.warning(
+            "TorchCompileModel_LoRASafe: could not install CPU patched-weight "
+            "backup routing on %s: %s",
+            model_class.__name__,
+            exc,
+        )
+        return False
+    return True
+
+
 class _LazyCompiled(nn.Module):
     """
     Wraps any nn.Module (e.g. a UNet or a single transformer block)
@@ -177,6 +256,16 @@ class TorchCompileModel_LoRASafe:
                         "resolution changes or model reloads.",
                     },
                 ),
+                "cpu_offload_patched_weight_backups": (
+                    "BOOLEAN",
+                    {
+                        "default": True,
+                        "tooltip": "Store original weights retained for LoRA/DoRA rollback "
+                        "on CPU when HIGH_VRAM would otherwise keep both the original and "
+                        "patched destination matrices on CUDA. Trades equivalent system "
+                        "RAM for substantially lower steady-state VRAM.",
+                    },
+                ),
             }
         }
 
@@ -460,7 +549,9 @@ class TorchCompileModel_LoRASafe:
         allow_flux_inferred_targets=False,
         fallback_to_full_model_if_no_targets=False,
         allow_patched_model_cudagraphs=False,
+        cpu_offload_patched_weight_backups=True,
     ):
+        has_weight_patches = self._has_weight_patches(model)
         backend, disable_cudagraphs = self._resolve_patched_model_cudagraph_policy(
             model=model,
             backend=backend,
@@ -495,6 +586,15 @@ class TorchCompileModel_LoRASafe:
             m = model.clone(disable_dynamic=True)
         except TypeError:
             m = model.clone()
+
+        if has_weight_patches and cpu_offload_patched_weight_backups:
+            if _enable_cpu_patched_weight_backups(m):
+                _LOG.info(
+                    "TorchCompileModel_LoRASafe: patched model detected; original "
+                    "destination weights retained for rollback will be stored on CPU. "
+                    "Live patched weights remain on CUDA."
+                )
+
         dm = m.get_model_object("diffusion_model")
 
         compile_kw = self._build_compile_kwargs(