Add sparse attention methods to H3 Sparse AAttention

Author: WASasquatch

NODES.md

@@ -1962,7 +1962,7 @@ Pick a frame cache written by Video Cache or H3 Decode Video and use it as a cli
 | Name | Type | Required | Default | Choices | What it does |
 |---|---|---|---|---|---|
 | `cache` | `COMBO` | Yes |  |  | Which frame cache to read, as `frame_cache/scene_00001/cache.wasframes [output]`. A cache in temp is gone after a restart. |
-| `delete_after_save` | `BOOLEAN` | Yes | False |  | `true` = the cache is deleted once a save has written all of it; `false` = it is kept for the next save. |
+| `delete_after_save` | `BOOLEAN` | Yes | False |  | `true` = the cache's own files are deleted once a save has written all of it, for a cache in output, temp or a `paths.allow_write` folder; `false` = it is kept for the next save. |
 
 **Outputs**
 
@@ -2088,7 +2088,7 @@ Keep frames on disk as a clip instead of holding them in memory, for a long rend
 | `name` | `STRING` | Yes | frame_cache/clip |  | The cache's name below root, numbered on each run, as `frame_cache/scene` for `frame_cache/scene_00001`. |
 | `delete_after_save` | `BOOLEAN` | Yes | False |  | `true` = the cached frames are deleted once a save has written all of them, and the node runs again on every queue; `false` = they are kept. |
 | `audio` | `AUDIO` | No |  |  | Sound for the frames, added after any already in the cache. |
-| `append_to` | `VIDEO` | No |  |  | A clip from Video Cache, H3 Decode Video or Load Video Cache to add these frames to, as a loop's carried value. Left empty, a new cache is made. |
+| `append_to` | `VIDEO` | No |  |  | A clip from Video Cache, H3 Decode Video or Load Video Cache to add these frames to, as a loop's carried value, for a cache in output, temp or a `paths.allow_write` folder. Left empty, a new cache is made. |
 | `depth` | `COMBO` | No | auto | `auto`, `16 bit` | How frames are kept: 'auto' = 8-bit codes, half floats for a batch outside 0 to 1; '16 bit' = half floats always, for a 10-bit save. |
 
 **Outputs**

SECURITY.md

@@ -63,10 +63,11 @@ prefix `\??\`, checked after `~` is expanded. A link along a path is read withou
 and one pointing at another machine is refused the same way. ComfyUI's own path helpers resolve
 what they are given, so the pack calls them only through wrappers that refuse first.
 
-A glob pattern, a file name prefix, and a name or folder typed below a folder the node has
-already chosen are refused before anything is listed or resolved when they name another
-machine, carry a drive, start at a root, or hold a segment of dots and spaces alone such as
-`..`. Whatever they match stays inside that folder.
+A glob pattern, a file name prefix, a frame cache name, and a name or folder typed below a
+folder the node has already chosen are refused before anything is listed or resolved when they
+name another machine, carry a drive, start at a root, or hold a segment of dots and spaces alone
+such as `..`. A frame cache name is refused as well when it holds a folder separator. Whatever
+they match stays inside that folder.
 
 A node input naming a file or directory is resolved and checked against an allowlist before it
 is opened, listed or written, whatever value arrives: a menu offers labels, and a linked input
@@ -83,6 +84,25 @@ No node writes the pack's settings. `config.yaml`, `config.json`, the config fil
 names, and the `viewer-extensions/` folder are refused as write targets, even inside the state
 directory. [docs/CONFIG.md](docs/CONFIG.md#containment) has the table.
 
+A frame cache picked from a folder no node writes to, such as `input/` or a `paths.allow_read`
+folder, is read and never added to or deleted.
+
+### Deleting files
+
+Files are deleted, moved or replaced in these modules and no others. Each removes files it
+wrote, by name. `shutil.rmtree` is called only on a checkpoint's download folder, which the
+fetch made inside the model folder.
+
+| Module | What it removes |
+|---|---|
+| `modules/media/frame_cache.py` | A frame cache's own `cache.wasframes`, `frames.bin` and `audio.npy` once a save has written it, then its numbered folder if that is left empty. Only inside `output/`, `temp/` or a `paths.allow_write` folder, and never a permitted root itself. The manifest is replaced through a `.part` file |
+| `modules/media/video.py` | The video Write to Video adds a frame to, replaced by the copy holding that frame |
+| `modules/media/temporal_upscale.py` | The files in `temp/` a failed Video Temporal Upscale run was writing |
+| `modules/archive/save.py` | The archive Zip Save was writing, when the write fails |
+| `modules/latent/h3_clip_file.py` | The `.part` file H3 Save Clip writes, renamed over the clip once complete |
+| `modules/interface/preview.py` | Mask editor uploads, as the route table below states |
+| `modules/model/h3_latent_upscaler.py`, `modules/model/flow_networks.py` | The download folder a checkpoint lands in, once the checkpoint is moved out of it |
+
 ## HTTP routes
 
 Every route the pack registers takes a key. None takes a path.

modules/media/frame_cache.py

@@ -12,25 +12,28 @@
 import json
 import os
 import re
-import shutil
 from fractions import Fraction
+from pathlib import Path
 
 import numpy as np
 import torch
 from comfy_api.latest import Input, InputImpl, Types
 
 from .. import log
 from ..image import scratch
+from ..util import sandbox
 
 __all__ = [
     "DATA",
     "DEPTHS",
     "EXTENSION",
     "MANIFEST",
+    "OWN_FILES",
     "SOUND",
     "CachedVideo",
     "FrameCache",
     "cached",
+    "writable",
 ]
 
 logger = log.get_logger("media.frame_cache")
@@ -60,16 +63,36 @@
 #: Folder names a cache takes, the name then a five digit number.
 FOLDER = re.compile(r"^(?P<name>.+)_(?P<number>\d{5})$")
 
+#: The files a cache writes in its folder, the only ones deleting it removes.
+OWN_FILES = (MANIFEST, MANIFEST + ".part", DATA, SOUND)
+
+
+def writable(folder: str) -> bool:
+    """Whether a cache folder lies inside a folder nodes may write to, so it may be added to or deleted."""
+    resolved = Path(os.path.abspath(folder)).resolve(strict=False)
+    return any(sandbox.contains(root, resolved) for root in sandbox.write_roots())
+
 
 def _next_folder(parent: str, name: str) -> str:
-    """The next unused ``name_NNNNN`` folder below ``parent``."""
+    """The next unused ``name_NNNNN`` folder below ``parent``, inside a permitted write folder.
+
+    Raises:
+        PathNotAllowed: ``name`` is not one folder name, or ``parent`` is outside every write root.
+    """
+    reason = sandbox.leaves_folder(name)
+    if reason is None and ("/" in name or "\\" in name):
+        reason = "holds a folder separator"
+    if reason is not None:
+        raise sandbox.PathNotAllowed(
+            f"refusing the frame cache name `{name}`: it is the name of one folder inside "
+            f"{parent}, and this one {reason}. Use a plain name such as `clip`.")
     highest = 0
     with contextlib.suppress(OSError):
         for entry in os.listdir(parent):
             found = FOLDER.match(entry)
             if found and found.group("name") == name:
                 highest = max(highest, int(found.group("number")))
-    return os.path.join(parent, f"{name}_{highest + 1:05d}")
+    return str(sandbox.resolve_write_file(parent, f"{name}_{highest + 1:05d}"))
 
 
 class FrameCache:
@@ -98,7 +121,12 @@ def create(cls, parent: str, name: str, rate, depth: str = "auto",
 
         Returns:
             The :class:`FrameCache`.
+
+        Raises:
+            PathNotAllowed: ``name`` is not one folder name, or the folder would land outside
+                every permitted write folder.
         """
+        parent = str(sandbox.resolve_write(parent))
         os.makedirs(parent, exist_ok=True)
         folder = _next_folder(parent, name)
         os.makedirs(folder)
@@ -196,6 +224,19 @@ def _frame_bytes(self, dtype: str) -> int:
         width, height = self.size
         return width * height * self.channels * (1 if dtype == "uint8" else 2)
 
+    def _require_writable(self) -> None:
+        """Refuse to add to a cache outside every folder nodes may write to.
+
+        Raises:
+            PathNotAllowed: The cache's folder is outside every permitted write folder.
+        """
+        if not writable(self.folder):
+            raise sandbox.PathNotAllowed(
+                f"the frame cache at {self.folder} is outside every folder nodes write to "
+                f"(output, temp and paths.allow_write), so nothing is added to it. Add to a "
+                f"cache in output or temp, or start a new cache."
+            )
+
     def append(self, images, scene: bool = False, device=None) -> int:
         """Write a batch of frames after the last.
 
@@ -209,11 +250,13 @@ def append(self, images, scene: bool = False, device=None) -> int:
 
         Raises:
             ValueError: The batch's size or channel count differs from the cache's.
+            PathNotAllowed: The cache is outside every permitted write folder.
         """
         count, height, width, channels = (int(side) for side in images.shape)
         first = self.frames
         if count == 0:
             return first
+        self._require_writable()
         if first == 0:
             self.manifest.update(width=width, height=height, channels=channels)
         elif (width, height, channels) != (*self.size, self.channels):
@@ -263,9 +306,11 @@ def add_audio(self, audio) -> None:
 
         Raises:
             ValueError: The sample rate or channel count differs from the sound kept.
+            PathNotAllowed: The cache is outside every permitted write folder.
         """
         if not isinstance(audio, dict) or audio.get("waveform") is None:
             return
+        self._require_writable()
         wave = audio["waveform"][0].detach().float().cpu().numpy()
         rate = int(audio.get("sample_rate", 44100))
         held = self.manifest.get("audio")
@@ -349,9 +394,27 @@ def audio(self, start: int = 0, stop: int | None = None) -> dict | None:
         return {"waveform": piece.unsqueeze(0), "sample_rate": int(held["sample_rate"])}
 
     def delete(self) -> None:
-        """Remove the cache's folder and everything in it."""
-        shutil.rmtree(self.folder, ignore_errors=True)
-        logger.info("deleted the frame cache at %s", self.folder)
+        """Remove the cache's own files and its numbered folder once empty; a cache outside every write folder stays."""
+        folder = Path(os.path.abspath(self.folder)).resolve(strict=False)
+        permitted = sandbox.write_roots()
+        if not any(sandbox.contains(root, folder) for root in permitted):
+            logger.warning(
+                "the frame cache at %s was not deleted: it is outside every folder a node "
+                "writes to (%s). Add its folder to paths.allow_write in config.yaml to let "
+                "the cache be deleted after saving.", folder, ", ".join(map(str, permitted)))
+            return
+        for name in OWN_FILES:
+            try:
+                (folder / name).unlink()
+            except FileNotFoundError:
+                pass
+            except OSError as error:
+                logger.warning("could not delete %s: %s", folder / name, error)
+        roots = {os.path.normcase(str(root)) for root in (*permitted, *sandbox.read_roots())}
+        if FOLDER.match(folder.name) and os.path.normcase(str(folder)) not in roots:
+            with contextlib.suppress(OSError):
+                folder.rmdir()
+        logger.info("deleted the frame cache at %s", folder)
 
 
 class _Frames:

nodes/io/load_video_cache.py

@@ -70,8 +70,9 @@ def define_schema(cls) -> io.Schema:
                     "delete_after_save",
                     default=False,
                     tooltip=(
-                        "`true` = the cache is deleted once a save has written all of it; "
-                        "`false` = it is kept for the next save."
+                        "`true` = the cache's own files are deleted once a save has written "
+                        "all of it, for a cache in output, temp or a `paths.allow_write` "
+                        "folder; `false` = it is kept for the next save."
                     ),
                 ),
             ],
@@ -98,7 +99,8 @@ def execute(cls, cache="", delete_after_save=False) -> io.NodeOutput:
 
         Raises:
             ValueError: No cache is chosen, or the chosen one is no longer there.
-            PathNotAllowed: The cache sits outside every folder this pack may read.
+            PathNotAllowed: The cache sits outside every folder this pack may read, or
+                ``delete_after_save`` is on for a cache outside every folder it may write.
         """
         path = _named(cache)
         if not path:
@@ -112,6 +114,12 @@ def execute(cls, cache="", delete_after_save=False) -> io.NodeOutput:
                 f"{NODE_NAME} cannot find {cache}; it was deleted or moved. Pick another cache."
             )
         folder = os.path.dirname(str(manifest))
+        if delete_after_save and not frame_cache.writable(folder):
+            raise sandbox.PathNotAllowed(
+                f"{NODE_NAME} deletes a cache after saving only from output, temp or a "
+                f"paths.allow_write folder, and {folder} is in none of them. Turn "
+                f"delete_after_save off to read it, or move the cache into output."
+            )
         video = frame_cache.CachedVideo(folder, delete_after_save=bool(delete_after_save))
         rate = video.get_frame_rate()
         logger.info("opened %d cached frame(s) at %s", video.get_frame_count(), folder)

nodes/io/video_cache.py

@@ -80,8 +80,8 @@ def define_schema(cls) -> io.Schema:
                     optional=True,
                     tooltip=(
                         "A clip from Video Cache, H3 Decode Video or Load Video Cache to add "
-                        "these frames to, as a loop's carried value. Left empty, a new cache "
-                        "is made."
+                        "these frames to, as a loop's carried value, for a cache in output, "
+                        "temp or a `paths.allow_write` folder. Left empty, a new cache is made."
                     ),
                 ),
                 io.Combo.Input(
@@ -121,7 +121,8 @@ def execute(
 
         Raises:
             ValueError: append_to is not a cached clip, or the frames differ from it in size.
-            PathNotAllowed: root and name settle outside every permitted write folder.
+            PathNotAllowed: root and name settle outside every permitted write folder, or
+                append_to's cache is outside every one.
         """
         if append_to is not None:
             if not frame_cache.cached(append_to):

pyproject.toml

@@ -1,7 +1,7 @@
 [project]
 name = "was-node-suite-comfyui"
 description = "WAS-NS reborn; Tools for image processing, filters, masking, text, logic, numbers, latents, files, 3D scenes, and animation."
-version = "3.5.2"
+version = "3.5.3"
 license = { file = "LICENSE" }
 requires-python = ">=3.10"