Compare commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
431941fc9d | ||
|
|
f5eab116c0 | ||
|
|
6008d31653 | ||
|
|
434817bc76 | ||
|
|
cdd98b13a3 | ||
|
|
f9683347be | ||
|
|
a2853a0ae1 | ||
|
|
b2eb11996b | ||
|
|
36fcf9db4f | ||
|
|
cb21a14424 | ||
|
|
f9ca96f648 | ||
|
|
2d7dd74fe2 | ||
|
|
8a379940bd | ||
|
|
a87ffb84b7 | ||
|
|
89bda6ba11 | ||
|
|
e191f0fc40 | ||
|
|
090dcdad62 | ||
|
|
47d1c4f739 | ||
|
|
a92017cff4 | ||
|
|
799d731ab5 | ||
|
|
5fc8cffc9f | ||
|
|
4dffaf3991 | ||
|
|
bc474f1532 | ||
|
|
c206d9e9da | ||
|
|
eba880cdee | ||
|
|
b6a67a4c34 | ||
|
|
27525ffdd8 | ||
|
|
eb668c448c | ||
|
|
9b3cd7c769 | ||
|
|
13c303178b | ||
|
|
f945feabb6 | ||
|
|
24ce4ffb06 | ||
|
|
980ed5db8b | ||
|
|
bf59887643 | ||
|
|
60be0ee865 | ||
|
|
146e3cfa8b | ||
|
|
e22ddb90d5 |
@@ -7,6 +7,10 @@ on:
|
||||
description: Full commit SHA from ComfyUI-AusBoss-Model-Downloader main
|
||||
required: true
|
||||
type: string
|
||||
changelog:
|
||||
description: Release note for this version, shown in the Registry's Updates section
|
||||
required: true
|
||||
type: string
|
||||
|
||||
# The existing AusBoss publisher credential stays in this repository. This
|
||||
# manual job publishes only the separately maintained Model Downloader addon.
|
||||
@@ -62,12 +66,16 @@ jobs:
|
||||
id: publish
|
||||
env:
|
||||
REGISTRY_ACCESS_TOKEN: ${{ secrets.REGISTRY_ACCESS_TOKEN }}
|
||||
COMFY_NODE_CHANGELOG: 'Initial release: model downloads from Civitai and Hugging Face, installed-file checks, saved model lists, and AusBoss Workflow Note integration.'
|
||||
COMFY_NODE_CHANGELOG: ${{ inputs.changelog }}
|
||||
run: |
|
||||
if [ -z "$REGISTRY_ACCESS_TOKEN" ]; then
|
||||
echo '::error::REGISTRY_ACCESS_TOKEN is not configured.'
|
||||
exit 1
|
||||
fi
|
||||
if [ -z "${COMFY_NODE_CHANGELOG// /}" ]; then
|
||||
echo '::error::The changelog input is empty; the Registry shows it as the release note.'
|
||||
exit 1
|
||||
fi
|
||||
comfy --skip-prompt --no-enable-telemetry node publish --token "$REGISTRY_ACCESS_TOKEN"
|
||||
- name: Report Registry review status
|
||||
if: always() && steps.publish.outcome != 'skipped'
|
||||
|
||||
@@ -2,6 +2,98 @@
|
||||
|
||||
All notable changes to ComfyUI-AusBoss are documented here.
|
||||
|
||||
## 2.24.0 - 2026-11-12
|
||||
|
||||
- **New node: Edit Lock.** Keeps every part of your picture that an edit
|
||||
did not change exactly where it was, with the colors it had. Edit models
|
||||
redraw the whole picture, so untouched parts often come back a little
|
||||
zoomed, shifted or tinted. Put Edit Lock between the model and the
|
||||
sampler and give it the picture you are editing. It needs no mask. It
|
||||
steps aside when the edit changes the whole picture, and lets rain,
|
||||
snow and haze through while holding the picture in place. Tested on
|
||||
MiniMax-H3-Image.
|
||||
|
||||
- **New node: Edit Lock Restore.** Puts your picture's own pixels back in
|
||||
every part the edit did not change. Add it after the VAE Decode with
|
||||
Edit Lock's `edit_lock` output. It also gives a mask of what the edit
|
||||
changed and a one-line report.
|
||||
|
||||
- **New node: Text Lines.** It runs what comes after once for each line of
|
||||
your text. Type five prompts, one on each line, and one run makes five
|
||||
pictures. It also takes a list a text model wrote: empty lines are
|
||||
skipped, a leading number or bullet comes off, and **Limit** stops a list
|
||||
that ran long.
|
||||
|
||||
- **Save Image shows every picture of a run at once.** When a run saves
|
||||
several pictures, they now show side by side in the node, the way the
|
||||
core Save Image does. The arrows still go to each picture by itself, and
|
||||
round again to all of them. Select Frame, Mask Refine, Mask by Name and
|
||||
LaMa Inpaint show a batch the same way.
|
||||
|
||||
- **Save Image can wait for a click before it saves.** Switch **Auto save**
|
||||
off and a run only shows its pictures on the node. **SAVE ALL** over them
|
||||
keeps every one; step to one picture with the arrows and **SAVE** keeps
|
||||
just that one. Good for making a batch and keeping the best. Auto save is
|
||||
on unless you switch it off, so nothing changes for workflows you have.
|
||||
|
||||
- **New node: Prompt Variations.** A prompt box that picks one choice from
|
||||
each set in curly brackets, such as `a {red|blue|green} {sweater|jacket}`,
|
||||
on every run. The seed decides the picks, so the same seed gives the same
|
||||
prompt again. Queue 50 runs with a random seed and get 50 different
|
||||
prompts, with no wildcard pack.
|
||||
|
||||
- **Mask by Touch: type a name to take it out of the mask.** The name field
|
||||
under the picture now has a sign in front. Click it to switch between **+**
|
||||
(add by name) and **−** (take away by name), then type "face" or "hair" and
|
||||
press Enter. Typing "-face" takes it away whatever the sign says. Before,
|
||||
a typed name was only taken away while Erase was chosen, which nothing on
|
||||
the node said.
|
||||
|
||||
- **Mask by Touch: Ctrl-click only takes away what you point at.** With a
|
||||
sweater picked, a Ctrl-click on the face offered the whole person, and the
|
||||
sweater went with it. Now a spot where nothing is picked says so and
|
||||
changes nothing.
|
||||
|
||||
- **Mask by Touch works on big pictures.** A look at a 3840 pixel picture took
|
||||
twice as long as at 2048, and at 8192 pixels it ran a 32 GB graphics card
|
||||
out of memory. The node now asks SAM 3 about a copy of at most 2048 pixels
|
||||
and brings the answer back to full size. Pictures up to 2048 pixels come
|
||||
out exactly as before. On bigger ones a mask's edge can sit a pixel or two
|
||||
differently.
|
||||
|
||||
- **Mask by Touch no longer fills the queue.** With Image Crop + Rotate + Pad
|
||||
or Load Image + Pad wired in, every change to the loader (each step of its
|
||||
megapixels, for one) queued one more look at a picture that was already
|
||||
replaced. Now only the newest look waits. The button that pauses the
|
||||
looking also says **Pause** and **Look on**.
|
||||
|
||||
- **New node: FlashVSR Upscale 🆎.** Makes a video larger and sharper with
|
||||
the FlashVSR v1.1 model, from 1x to 4x. It needs one model file and no
|
||||
other pack: the help page has the download link. A long clip needs no more
|
||||
graphics memory than a short one, and a result too large for the free
|
||||
memory is made in tiles.
|
||||
The new example `FlashVSR Video Upscale (AusBoss)` shows it between Load
|
||||
Video and Save Video.
|
||||
|
||||
- **Two errors from outside a workflow now say what to do.** When a MiniMax H3
|
||||
run stops with `KeyError: 'latent'` because the Muse MiniMax Director pack is
|
||||
installed, a message names that pack and says to turn it off. When an older
|
||||
ComfyUI cannot read a `w6a8` model file, a message says to update ComfyUI.
|
||||
ComfyUI's own error report stays as it is.
|
||||
|
||||
- **ComfyUI starts when the pack is installed twice.** A Manager install
|
||||
(`custom_nodes/ausboss-nodes`) beside a git clone
|
||||
(`custom_nodes/ComfyUI-AusBoss`) stopped ComfyUI while it started. Now one
|
||||
copy loads and the other steps aside, and the console says which folder to
|
||||
delete. A copy from 2.23.0 or older is the one that loads when it sits
|
||||
beside a newer one, so delete the older folder.
|
||||
|
||||
- **A Workflow Note popup closes when you click outside it.** Opening an
|
||||
info tile (a Tips note, say) showed the full note, and only its Close
|
||||
button or Escape put it away. A click on the dark area around it now
|
||||
closes it too. Selecting text in the note, or a drag that starts on it,
|
||||
keeps it open, and the Edit form still asks before it drops a change.
|
||||
|
||||
## 2.23.0 - 2026-11-09
|
||||
|
||||
- **Mask by Touch keeps your picks when the picture changes.** Your clicks
|
||||
|
||||
@@ -20,6 +20,7 @@ Most of them are for image and video work. ComfyUI's core nodes already cover th
|
||||
- **[Mask by Touch](#mask-by-touch):** click a thing on the picture and get its mask. Hover shows what a click would take, and **Fine edges** follows hair and fur.
|
||||
- **[Mask by Name](#mask-by-name):** type "the red jacket" and get its mask and a cut-out.
|
||||
- **[Image Crop + Rotate + Pad](#image-crop--rotate--pad):** turn, crop and pad a picture by dragging handles on the node, for outpainting.
|
||||
- **[Edit Lock](#edit-lock):** an edit changes what you asked for, and the rest of the picture stays exactly where it was. No mask needed.
|
||||
- **[Stitch Inpaint](#stitch-inpaint):** after an inpaint or outpaint, everything you didn't ask to change is your original pixels again.
|
||||
- **[Video Crop + Rotate + Pad → Clip](#video-crop--rotate--pad--clip):** trim a clip on a timeline, and pad or crop it to a new shape.
|
||||
- **[Tiled Upscale](#tiled-upscale):** upscale a picture in tiles to the megapixels you set.
|
||||
@@ -42,6 +43,8 @@ git clone https://github.com/ausboss/ComfyUI-AusBoss.git
|
||||
|
||||
Restart ComfyUI and search the node library for **AusBoss**.
|
||||
|
||||
Install it one way, not both. Manager puts the pack in `custom_nodes/ausboss-nodes` and a clone in `custom_nodes/ComfyUI-AusBoss`. They are the same pack, so if both folders are there, delete one. The console says which copy it skipped.
|
||||
|
||||
After an update, through Manager or `git pull`, restart ComfyUI and hard-refresh the browser with **Ctrl+Shift+R**. Otherwise the browser keeps the old JavaScript, and the pack warns you about it. Manager can install a version or two behind while the Comfy Registry reviews a new release. To get the newest one, pick it in Manager's version list, or update by hand: delete the pack's folder in `ComfyUI/custom_nodes` and clone the repository as shown above.
|
||||
|
||||
The pack uses the Pillow, NumPy, Torch and PyAV that come with ComfyUI. Each example lists the models it needs on its Workflow Note. LaMa Inpaint also needs [`big-lama.pt`](https://github.com/Sanster/models/releases/download/add_big_lama/big-lama.pt) in `ComfyUI/models/lama/`. The pack never downloads weights by itself.
|
||||
@@ -137,6 +140,18 @@ The same watercolor edit without and with Realign to Source: 61 px off before, 0
|
||||
|
||||
Qwen drew this watercolor 5.6% taller, and Realign to Source lined it back up with the original.
|
||||
|
||||
### Edit Lock
|
||||
|
||||
Keeps every part of your picture that an edit did not change exactly where it was, with the colors it had. An edit model redraws the whole picture, so the parts you did not ask about often come back a little zoomed in, shifted or tinted, and the picture jumps when you flip between before and after. Put Edit Lock between the model and the sampler and give it the picture you are editing. It needs no mask: it sees what the edit changes while the picture is being made, and holds on to everything else.
|
||||
|
||||
When the edit changes the whole picture (a new style, new light), there is nothing to keep, so the node steps aside and the edit runs as usual. A soft effect over the whole picture, such as rain or falling snow, is let through as well: the node then only holds the picture in place. **Lock** (`lock`) set to `off` does nothing at all, `always` locks even a whole-picture edit, and an optional `mask` says where the edit may happen. **Tolerance** (`tolerance`) sets how big a change must be to count as part of the edit.
|
||||
|
||||
On MiniMax-H3-Image, 4 megapixel edits moved the picture by up to 14 px without the lock (measured at 1 megapixel) and by 0.2 px or less with it (0.08 px or less on 35 of the 36 locked runs), on 19 local edits at two seeds each. After five edits in a row on one picture the plain chain had drifted 1.65 px and turned visibly pink; the locked chain sat at 0.04 px with no tint.
|
||||
|
||||
### Edit Lock Restore
|
||||
|
||||
Puts your picture's own pixels back in every part the edit did not change. Edit Lock keeps those parts in place, but they still pass through the VAE, which softens fine detail a little. Add this right after the VAE Decode and wire in Edit Lock's `edit_lock` output. **Feather** (`feather`) sets how softly the original fades into the edit. It also returns `edit_mask`, white where the edit changed the picture, and a one-line `report` of how much was kept.
|
||||
|
||||
### Image Compare A/B
|
||||
|
||||
Compares two pictures with a slider, or flips between them whole. Nothing covers the picture, and its size shows underneath. The panel grows with the node, and images A and B pass through, so you can put it anywhere on an image wire. **Erase** mixes the two by hand: rub parts of one picture away with a brush so the other shows through, then save the mix as a PNG under a name you type, with a layered PSD for Photoshop or GIMP if you want it.
|
||||
@@ -197,6 +212,12 @@ Changes a clip's frame rate by making in-between frames. You set the source and
|
||||
|
||||
It returns the new frames and their fps. It doesn't touch audio, so carry the source audio to Save Video separately. The [help page](js/docs/AUSBOSS_NODES_FrameInterpolate.md) covers the weights and timing.
|
||||
|
||||
### FlashVSR Upscale
|
||||
|
||||
Makes a video larger and sharper with the FlashVSR model. Set **Scale** from 1 to 4 and it returns every frame at the new size, with fine detail drawn back in. It reads the clip a few frames at a time, so a long clip needs no more graphics memory than a short one, and a result too large for the free memory is made in tiles.
|
||||
|
||||
It needs one model file in `ComfyUI/models/diffusion_models` and no other pack. The [help page](js/docs/AUSBOSS_NODES_FlashVSRUpscale.md) has the download link, the memory it needs and the speed. [Try the FlashVSR Video Upscale example →](example_workflows/FlashVSR%20Video%20Upscale%20%28AusBoss%29.json)
|
||||
|
||||
### Save Video
|
||||
|
||||
Saves a frame batch or a core `VIDEO` as MP4 (H.264 or H.265), WebM (VP9 or AV1), ProRes MOV, lossless FFV1 MKV, GIF or WebP. Which encoders you get depends on your FFmpeg and PyAV build, and H.264 and H.265 come in CPU and NVENC versions. The videos carry color tags, and optionally audio and the workflow.
|
||||
@@ -342,6 +363,14 @@ Puts batch `b` after `a`. When the sizes match, nothing is resized. When they di
|
||||
|
||||
A text box that grows with the node and sends its text, unchanged, down a `STRING` wire. Use it for a prompt, caption, path or other text that several nodes share. Spaces and line breaks are kept. Ctrl + Up and Ctrl + Down change the weight of the selected words, as in ComfyUI's own prompt boxes.
|
||||
|
||||
### Text Lines
|
||||
|
||||
Runs what comes after once for each line of your text. Type five prompts, one on each line, and one run makes five pictures. It also takes a list another node wrote, skips empty lines, and can take a leading `1.` or `-` off each line. `number` counts the lines from 1, to give each its own seed.
|
||||
|
||||
### Prompt Variations
|
||||
|
||||
A prompt box that changes on every run. Write choices in curly brackets, split by `|`, such as `a {red|blue|green} {sweater|jacket}`, and the node picks one from each set. Wire in the seed your sampler uses: a new seed gives a new mix, and the same seed gives the same prompt again. Set the seed to randomize and queue 50 runs to get 50 different prompts. Text without brackets comes out exactly as typed, so a JSON prompt is safe.
|
||||
|
||||
### H3 Prompt Writer
|
||||
|
||||
Type what happens, one line for each moment, and leave a blank line to cut to a new place. Connect up to four pictures and give each one a job: a person, a place, an outfit, a thing or a style. Copy the moves from a video, keep a clip going, or add a voice or a song. A timeline on the node shows when each moment plays before you run it. It needs no extra model and you never see H3's prompt format. Connect a small text model and it fills in a one-line idea.
|
||||
@@ -394,7 +423,7 @@ Turns parts of a workflow off and on from one small card. Select the nodes of on
|
||||
|
||||
## Example workflows
|
||||
|
||||
These are the workflows where the nodes do something core ComfyUI doesn't do by itself: outpainting and inpainting that put your original picture back, lining an edit back up with its source, upscaling a picture in tiles, masking a thing by its name or with a click, and removing a watermark from a clip. Open a JSON file from [`example_workflows/`](example_workflows), or find them under this pack in ComfyUI's template browser. Each one has a thumbnail, numbered groups and a Workflow Note with the setup steps and model downloads.
|
||||
These are the workflows where the nodes do something core ComfyUI doesn't do by itself: outpainting and inpainting that put your original picture back, lining an edit back up with its source, upscaling a picture in tiles, masking a thing by its name or with a click, removing a watermark from a clip, and making a clip larger and sharper. Open a JSON file from [`example_workflows/`](example_workflows), or find them under this pack in ComfyUI's template browser. Each one has a thumbnail, numbered groups and a Workflow Note with the setup steps and model downloads.
|
||||
|
||||

|
||||
|
||||
@@ -415,6 +444,7 @@ The workflows open with their image and video loaders empty, so you load your ow
|
||||
| [Mask by Name](example_workflows/Mask%20by%20Name%20%28AusBoss%29.json) | Type what to find in a picture, and get its mask and a cut-out with a see-through background | ComfyUI's own SAM 3.1 (ComfyUI 0.20 or newer) |
|
||||
| [Mask by Touch](example_workflows/Mask%20by%20Touch%20%28AusBoss%29.json) | Click a thing in a picture, or its whole subject, and get its mask and a cut-out with hair-fine edges | ComfyUI's own SAM 3.1 and its background-removal model (ComfyUI 0.21 or newer) |
|
||||
| [Simple Video Watermark Remover](example_workflows/Simple%20Video%20Watermark%20Remover%20%28AusBoss%29.json) | Find an overlay in a clip and remove it, with a one-frame comparison branch | ComfyUI's own SAM 3.1 (ComfyUI 0.20 or newer), plus `big-lama.pt` |
|
||||
| [FlashVSR Video Upscale](example_workflows/FlashVSR%20Video%20Upscale%20%28AusBoss%29.json) | Make a clip larger and sharper with one node | `flashvsr_v1.1_fp16.safetensors` |
|
||||
|
||||
Every example uses only core nodes and this pack. Speed and memory use depend on the models, the sizes, the frame count and whatever else is using the GPU.
|
||||
|
||||
|
||||
+38
-7
@@ -22,6 +22,8 @@ NODE_MODULES = [
|
||||
"node_callout",
|
||||
"node_color_match",
|
||||
"node_compare",
|
||||
"node_edit_lock",
|
||||
"node_flashvsr_upscale",
|
||||
"node_frame_interpolate",
|
||||
"node_free_memory",
|
||||
"node_h3_prompt_writer",
|
||||
@@ -42,6 +44,7 @@ NODE_MODULES = [
|
||||
"node_mask_by_name",
|
||||
"node_mask_by_touch",
|
||||
"node_math",
|
||||
"node_prompt_variations",
|
||||
"node_realign_to_source",
|
||||
"node_refine_mask",
|
||||
"node_resolution",
|
||||
@@ -52,6 +55,7 @@ NODE_MODULES = [
|
||||
"node_select_frame",
|
||||
"node_show_text",
|
||||
"node_step_preview",
|
||||
"node_text_lines",
|
||||
"node_tiled_upscale",
|
||||
"node_video_crop_rotate_pad",
|
||||
"node_video_crop_rotate_pad_clip",
|
||||
@@ -63,7 +67,19 @@ NODE_CLASS_MAPPINGS = {}
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {}
|
||||
_failed_modules = []
|
||||
|
||||
# Two copies of the pack (a Manager install beside a git clone) add every
|
||||
# route twice and ComfyUI stops while starting. Only one copy loads.
|
||||
_other_copy = None
|
||||
try:
|
||||
from .nodes._single_copy import other_copy_that_loads
|
||||
|
||||
_other_copy = other_copy_that_loads(os.path.dirname(__file__))
|
||||
except Exception:
|
||||
traceback.print_exc()
|
||||
|
||||
for _module_name in NODE_MODULES:
|
||||
if _other_copy:
|
||||
break
|
||||
try:
|
||||
_module = importlib.import_module(f".nodes.{_module_name}", __name__)
|
||||
NODE_CLASS_MAPPINGS.update(getattr(_module, "NODE_CLASS_MAPPINGS", {}))
|
||||
@@ -76,7 +92,7 @@ for _module_name in NODE_MODULES:
|
||||
|
||||
# Folder ComfyUI serves to the browser; every .js file in it is loaded as a
|
||||
# frontend extension. Shared modules use .mjs so they are import-only.
|
||||
WEB_DIRECTORY = "./js"
|
||||
WEB_DIRECTORY = None if _other_copy else "./js"
|
||||
|
||||
__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS", "WEB_DIRECTORY"]
|
||||
|
||||
@@ -112,11 +128,26 @@ def _print_banner():
|
||||
print(bar)
|
||||
|
||||
|
||||
try:
|
||||
from .nodes._frontend_cache import register_frontend_cache_guard
|
||||
def _print_second_copy():
|
||||
# ASCII only, as in the banner.
|
||||
red = "\033[0;31m"
|
||||
reset = "\033[0m"
|
||||
own = os.path.basename(os.path.dirname(os.path.realpath(__file__)))
|
||||
other = os.path.basename(_other_copy["folder"])
|
||||
print(f"{red}AusBoss: this pack is installed twice.{reset}")
|
||||
print(f" Using: custom_nodes/{other} (v{_other_copy['version']})")
|
||||
print(f" Skipping: custom_nodes/{own} (v{_read_version()})")
|
||||
print(" Delete one of the two folders, then restart ComfyUI.")
|
||||
|
||||
register_frontend_cache_guard()
|
||||
except Exception:
|
||||
traceback.print_exc()
|
||||
|
||||
_print_banner()
|
||||
if _other_copy:
|
||||
_print_second_copy()
|
||||
else:
|
||||
try:
|
||||
from .nodes._frontend_cache import register_frontend_cache_guard
|
||||
|
||||
register_frontend_cache_guard()
|
||||
except Exception:
|
||||
traceback.print_exc()
|
||||
|
||||
_print_banner()
|
||||
|
||||
@@ -1,9 +1,9 @@
|
||||
<svg xmlns="http://www.w3.org/2000/svg" width="105" height="20" role="img" aria-label="release: v2.23.0">
|
||||
<title>release: v2.23.0</title>
|
||||
<svg xmlns="http://www.w3.org/2000/svg" width="105" height="20" role="img" aria-label="release: v2.24.0">
|
||||
<title>release: v2.24.0</title>
|
||||
<path fill="#555" d="M0 0h48v20H0z"/>
|
||||
<path fill="#00b4aa" d="M48 0h57v20H48z"/>
|
||||
<g fill="#fff" text-anchor="middle" font-family="Verdana,Geneva,DejaVu Sans,sans-serif" font-size="11">
|
||||
<text x="24" y="14">release</text>
|
||||
<text x="76.5" y="14">v2.23.0</text>
|
||||
<text x="76.5" y="14">v2.24.0</text>
|
||||
</g>
|
||||
</svg>
|
||||
|
||||
|
Before Width: | Height: | Size: 425 B After Width: | Height: | Size: 425 B |
Binary file not shown.
|
After Width: | Height: | Size: 127 KiB |
@@ -0,0 +1,526 @@
|
||||
{
|
||||
"id": "0b455c1f-6b0c-4f69-8a7e-86f8de61df5c",
|
||||
"revision": 0,
|
||||
"last_node_id": 5,
|
||||
"last_link_id": 4,
|
||||
"nodes": [
|
||||
{
|
||||
"id": 1,
|
||||
"type": "AUSBOSS_NODES_WorkflowNote",
|
||||
"pos": [
|
||||
0,
|
||||
30
|
||||
],
|
||||
"size": [
|
||||
440,
|
||||
640
|
||||
],
|
||||
"flags": {},
|
||||
"order": 0,
|
||||
"mode": 0,
|
||||
"inputs": [
|
||||
{
|
||||
"name": "note",
|
||||
"type": "STRING",
|
||||
"link": null,
|
||||
"widget": {
|
||||
"name": "note"
|
||||
}
|
||||
}
|
||||
],
|
||||
"outputs": [],
|
||||
"title": "🆎 AusBoss",
|
||||
"properties": {
|
||||
"cnr_id": "ausboss-nodes",
|
||||
"ver": "2.24.0",
|
||||
"Node name for S&R": "AUSBOSS_NODES_WorkflowNote"
|
||||
},
|
||||
"widgets_values": [
|
||||
"{\"v\": 1, \"title\": \"FlashVSR Video Upscale\", \"subtitle\": \"Make a video larger and sharper with one node and one model file.\", \"author\": \"AusBoss\", \"layout\": \"card\", \"accent\": \"\", \"body\": \"[▶ AusBoss on YouTube](https://www.youtube.com/@ausbossdev) · videos on these nodes and workflows\\n\\n---\\n\\n1. Load your video.\\n2. Set **Scale**: how many times larger each side gets. At 1 the video keeps its size and is only sharpened.\\n3. Press **Run**, then play the result in **Save Video**.\\n\\nTo try a short piece first, drag the right end of the bar under the player in Load Video.\\n\\nA long video needs no more graphics memory than a short one. A large result is made in tiles, so it fits the memory you have and takes longer.\\n\\nNeeds AusBoss nodes 2.24.0 or newer. No other pack is needed.\", \"models\": [{\"name\": \"flashvsr_v1.1_fp16.safetensors\", \"dir\": \"diffusion_models\", \"size\": \"3.5 GB\", \"url\": \"https://huggingface.co/ausboss/FlashVSR-v1.1-ComfyUI/resolve/main/flashvsr_v1.1_fp16.safetensors\", \"note\": \"FlashVSR v1.1 by Junhao Zhuang and the OpenImagingLab team, in one file\"}], \"packs\": [{\"name\": \"ComfyUI-AusBoss\", \"url\": \"https://github.com/ausboss/ComfyUI-AusBoss\", \"node\": \"AUSBOSS_NODES_FlashVSRUpscale\"}], \"links\": [{\"label\": \"YouTube · ausbossdev\", \"url\": \"https://www.youtube.com/@ausbossdev\"}, {\"label\": \"GitHub · ausboss\", \"url\": \"https://github.com/ausboss\"}, {\"label\": \"Civitai · AusBoss\", \"url\": \"https://civitai.com/user/AusBoss\"}, {\"label\": \"X · Zanzibased\", \"url\": \"https://x.com/Zanzibased\"}]}",
|
||||
""
|
||||
],
|
||||
"widgets_values_named": {
|
||||
"note": "{\"v\": 1, \"title\": \"FlashVSR Video Upscale\", \"subtitle\": \"Make a video larger and sharper with one node and one model file.\", \"author\": \"AusBoss\", \"layout\": \"card\", \"accent\": \"\", \"body\": \"[▶ AusBoss on YouTube](https://www.youtube.com/@ausbossdev) · videos on these nodes and workflows\\n\\n---\\n\\n1. Load your video.\\n2. Set **Scale**: how many times larger each side gets. At 1 the video keeps its size and is only sharpened.\\n3. Press **Run**, then play the result in **Save Video**.\\n\\nTo try a short piece first, drag the right end of the bar under the player in Load Video.\\n\\nA long video needs no more graphics memory than a short one. A large result is made in tiles, so it fits the memory you have and takes longer.\\n\\nNeeds AusBoss nodes 2.24.0 or newer. No other pack is needed.\", \"models\": [{\"name\": \"flashvsr_v1.1_fp16.safetensors\", \"dir\": \"diffusion_models\", \"size\": \"3.5 GB\", \"url\": \"https://huggingface.co/ausboss/FlashVSR-v1.1-ComfyUI/resolve/main/flashvsr_v1.1_fp16.safetensors\", \"note\": \"FlashVSR v1.1 by Junhao Zhuang and the OpenImagingLab team, in one file\"}], \"packs\": [{\"name\": \"ComfyUI-AusBoss\", \"url\": \"https://github.com/ausboss/ComfyUI-AusBoss\", \"node\": \"AUSBOSS_NODES_FlashVSRUpscale\"}], \"links\": [{\"label\": \"YouTube · ausbossdev\", \"url\": \"https://www.youtube.com/@ausbossdev\"}, {\"label\": \"GitHub · ausboss\", \"url\": \"https://github.com/ausboss\"}, {\"label\": \"Civitai · AusBoss\", \"url\": \"https://civitai.com/user/AusBoss\"}, {\"label\": \"X · Zanzibased\", \"url\": \"https://x.com/Zanzibased\"}]}",
|
||||
"ausboss_workflow_note_panel": ""
|
||||
},
|
||||
"color": "#14424d",
|
||||
"bgcolor": "#161f21"
|
||||
},
|
||||
{
|
||||
"id": 2,
|
||||
"type": "AUSBOSS_NODES_LoadVideo",
|
||||
"pos": [
|
||||
500,
|
||||
80
|
||||
],
|
||||
"size": [
|
||||
420,
|
||||
760
|
||||
],
|
||||
"flags": {},
|
||||
"order": 0,
|
||||
"mode": 0,
|
||||
"inputs": [
|
||||
{
|
||||
"name": "video",
|
||||
"type": "COMBO",
|
||||
"link": null,
|
||||
"widget": {
|
||||
"name": "video"
|
||||
}
|
||||
},
|
||||
{
|
||||
"name": "start_seconds",
|
||||
"type": "FLOAT",
|
||||
"link": null,
|
||||
"widget": {
|
||||
"name": "start_seconds"
|
||||
}
|
||||
},
|
||||
{
|
||||
"name": "end_seconds",
|
||||
"type": "FLOAT",
|
||||
"link": null,
|
||||
"widget": {
|
||||
"name": "end_seconds"
|
||||
}
|
||||
},
|
||||
{
|
||||
"name": "custom_width",
|
||||
"type": "INT",
|
||||
"link": null,
|
||||
"widget": {
|
||||
"name": "custom_width"
|
||||
}
|
||||
},
|
||||
{
|
||||
"name": "custom_height",
|
||||
"type": "INT",
|
||||
"link": null,
|
||||
"widget": {
|
||||
"name": "custom_height"
|
||||
}
|
||||
},
|
||||
{
|
||||
"name": "every_nth",
|
||||
"type": "INT",
|
||||
"link": null,
|
||||
"widget": {
|
||||
"name": "every_nth"
|
||||
}
|
||||
},
|
||||
{
|
||||
"name": "max_frames",
|
||||
"type": "INT",
|
||||
"link": null,
|
||||
"widget": {
|
||||
"name": "max_frames"
|
||||
}
|
||||
},
|
||||
{
|
||||
"name": "single_frame",
|
||||
"type": "BOOLEAN",
|
||||
"link": null,
|
||||
"widget": {
|
||||
"name": "single_frame"
|
||||
}
|
||||
},
|
||||
{
|
||||
"name": "upload",
|
||||
"type": "IMAGEUPLOAD",
|
||||
"link": null,
|
||||
"widget": {
|
||||
"name": "upload"
|
||||
}
|
||||
}
|
||||
],
|
||||
"outputs": [
|
||||
{
|
||||
"name": "frames",
|
||||
"type": "IMAGE",
|
||||
"links": [
|
||||
1
|
||||
]
|
||||
},
|
||||
{
|
||||
"name": "audio",
|
||||
"type": "AUDIO",
|
||||
"links": [
|
||||
3
|
||||
]
|
||||
},
|
||||
{
|
||||
"name": "frame_count",
|
||||
"type": "INT",
|
||||
"links": []
|
||||
},
|
||||
{
|
||||
"name": "fps",
|
||||
"type": "FLOAT",
|
||||
"links": [
|
||||
4
|
||||
]
|
||||
},
|
||||
{
|
||||
"name": "width",
|
||||
"type": "INT",
|
||||
"links": []
|
||||
},
|
||||
{
|
||||
"name": "height",
|
||||
"type": "INT",
|
||||
"links": []
|
||||
},
|
||||
{
|
||||
"name": "duration",
|
||||
"type": "FLOAT",
|
||||
"links": []
|
||||
},
|
||||
{
|
||||
"name": "video",
|
||||
"type": "VIDEO",
|
||||
"links": []
|
||||
}
|
||||
],
|
||||
"title": "Load Video - Your video",
|
||||
"properties": {
|
||||
"cnr_id": "ausboss-nodes",
|
||||
"ver": "2.24.0",
|
||||
"Node name for S&R": "AUSBOSS_NODES_LoadVideo",
|
||||
"ausboss_load_video_loop": true
|
||||
},
|
||||
"widgets_values": {
|
||||
"video": "",
|
||||
"start_seconds": 0,
|
||||
"end_seconds": 0,
|
||||
"custom_width": 0,
|
||||
"custom_height": 0,
|
||||
"every_nth": 1,
|
||||
"max_frames": 0,
|
||||
"single_frame": false
|
||||
},
|
||||
"widgets_values_named": {
|
||||
"video": "",
|
||||
"start_seconds": 0,
|
||||
"end_seconds": 0,
|
||||
"custom_width": 0,
|
||||
"custom_height": 0,
|
||||
"every_nth": 1,
|
||||
"max_frames": 0,
|
||||
"single_frame": false,
|
||||
"upload": "image",
|
||||
"ausboss_widget_card": "",
|
||||
"ausboss_load_video_viewer": ""
|
||||
},
|
||||
"color": "#14424d",
|
||||
"bgcolor": "#161f21"
|
||||
},
|
||||
{
|
||||
"id": 3,
|
||||
"type": "AUSBOSS_NODES_FlashVSRUpscale",
|
||||
"pos": [
|
||||
980,
|
||||
80
|
||||
],
|
||||
"size": [
|
||||
360,
|
||||
190
|
||||
],
|
||||
"flags": {},
|
||||
"order": 0,
|
||||
"mode": 0,
|
||||
"inputs": [
|
||||
{
|
||||
"name": "frames",
|
||||
"type": "IMAGE",
|
||||
"link": 1
|
||||
},
|
||||
{
|
||||
"name": "vae",
|
||||
"type": "VAE",
|
||||
"shape": 7,
|
||||
"link": null
|
||||
},
|
||||
{
|
||||
"name": "model",
|
||||
"type": "COMBO",
|
||||
"link": null,
|
||||
"widget": {
|
||||
"name": "model"
|
||||
}
|
||||
},
|
||||
{
|
||||
"name": "scale",
|
||||
"type": "FLOAT",
|
||||
"link": null,
|
||||
"widget": {
|
||||
"name": "scale"
|
||||
}
|
||||
},
|
||||
{
|
||||
"name": "color_match",
|
||||
"type": "BOOLEAN",
|
||||
"link": null,
|
||||
"widget": {
|
||||
"name": "color_match"
|
||||
}
|
||||
}
|
||||
],
|
||||
"outputs": [
|
||||
{
|
||||
"name": "frames",
|
||||
"type": "IMAGE",
|
||||
"links": [
|
||||
2
|
||||
]
|
||||
}
|
||||
],
|
||||
"title": "FlashVSR Upscale",
|
||||
"properties": {
|
||||
"cnr_id": "ausboss-nodes",
|
||||
"ver": "2.24.0",
|
||||
"Node name for S&R": "AUSBOSS_NODES_FlashVSRUpscale",
|
||||
"models": [
|
||||
{
|
||||
"name": "flashvsr_v1.1_fp16.safetensors",
|
||||
"url": "https://huggingface.co/ausboss/FlashVSR-v1.1-ComfyUI/resolve/main/flashvsr_v1.1_fp16.safetensors",
|
||||
"directory": "diffusion_models"
|
||||
}
|
||||
]
|
||||
},
|
||||
"widgets_values": [
|
||||
"flashvsr_v1.1_fp16.safetensors",
|
||||
2,
|
||||
true,
|
||||
""
|
||||
],
|
||||
"widgets_values_named": {
|
||||
"model": "flashvsr_v1.1_fp16.safetensors",
|
||||
"scale": 2,
|
||||
"color_match": true,
|
||||
"ausboss_widget_card": ""
|
||||
},
|
||||
"color": "#14424d",
|
||||
"bgcolor": "#161f21"
|
||||
},
|
||||
{
|
||||
"id": 4,
|
||||
"type": "AUSBOSS_NODES_SaveVideo",
|
||||
"pos": [
|
||||
1400,
|
||||
80
|
||||
],
|
||||
"size": [
|
||||
440,
|
||||
640
|
||||
],
|
||||
"flags": {},
|
||||
"order": 0,
|
||||
"mode": 0,
|
||||
"inputs": [
|
||||
{
|
||||
"name": "frames",
|
||||
"type": "IMAGE",
|
||||
"link": 2
|
||||
},
|
||||
{
|
||||
"name": "audio",
|
||||
"type": "AUDIO",
|
||||
"link": 3
|
||||
},
|
||||
{
|
||||
"name": "video",
|
||||
"type": "VIDEO",
|
||||
"link": null
|
||||
},
|
||||
{
|
||||
"name": "fps",
|
||||
"type": "FLOAT",
|
||||
"link": 4
|
||||
},
|
||||
{
|
||||
"name": "filename_prefix",
|
||||
"type": "STRING",
|
||||
"link": null,
|
||||
"widget": {
|
||||
"name": "filename_prefix"
|
||||
}
|
||||
},
|
||||
{
|
||||
"name": "crf",
|
||||
"type": "INT",
|
||||
"link": null
|
||||
},
|
||||
{
|
||||
"name": "format",
|
||||
"type": "COMBO",
|
||||
"link": null,
|
||||
"widget": {
|
||||
"name": "format"
|
||||
}
|
||||
},
|
||||
{
|
||||
"name": "pingpong",
|
||||
"type": "BOOLEAN",
|
||||
"link": null,
|
||||
"widget": {
|
||||
"name": "pingpong"
|
||||
}
|
||||
},
|
||||
{
|
||||
"name": "save_metadata",
|
||||
"type": "BOOLEAN",
|
||||
"link": null,
|
||||
"widget": {
|
||||
"name": "save_metadata"
|
||||
}
|
||||
}
|
||||
],
|
||||
"outputs": [
|
||||
{
|
||||
"name": "file_path",
|
||||
"type": "STRING",
|
||||
"links": []
|
||||
}
|
||||
],
|
||||
"title": "Save Video",
|
||||
"properties": {
|
||||
"cnr_id": "ausboss-nodes",
|
||||
"ver": "2.24.0",
|
||||
"Node name for S&R": "AUSBOSS_NODES_SaveVideo",
|
||||
"ausboss_save_video_loop": true
|
||||
},
|
||||
"widgets_values": {
|
||||
"fps": 24,
|
||||
"filename_prefix": "flashvsr_upscaled",
|
||||
"crf": 19,
|
||||
"format": "mp4 h264",
|
||||
"pingpong": false,
|
||||
"save_metadata": true
|
||||
},
|
||||
"widgets_values_named": {
|
||||
"fps": 24,
|
||||
"filename_prefix": "flashvsr_upscaled",
|
||||
"crf": 19,
|
||||
"format": "mp4 h264",
|
||||
"pingpong": false,
|
||||
"save_metadata": true,
|
||||
"ausboss_widget_card": "",
|
||||
"ausboss_save_video_viewer": ""
|
||||
},
|
||||
"color": "#14424d",
|
||||
"bgcolor": "#161f21"
|
||||
},
|
||||
{
|
||||
"id": 5,
|
||||
"type": "AUSBOSS_NODES_RunTimer",
|
||||
"pos": [
|
||||
1515,
|
||||
770
|
||||
],
|
||||
"size": [
|
||||
210,
|
||||
67
|
||||
],
|
||||
"flags": {},
|
||||
"order": 0,
|
||||
"mode": 0,
|
||||
"inputs": [],
|
||||
"outputs": [],
|
||||
"title": "Run Timer",
|
||||
"properties": {
|
||||
"cnr_id": "ausboss-nodes",
|
||||
"ver": "2.24.0",
|
||||
"Node name for S&R": "AUSBOSS_NODES_RunTimer"
|
||||
},
|
||||
"color": "#14424d",
|
||||
"bgcolor": "#161f21"
|
||||
}
|
||||
],
|
||||
"links": [
|
||||
[
|
||||
1,
|
||||
2,
|
||||
0,
|
||||
3,
|
||||
0,
|
||||
"IMAGE"
|
||||
],
|
||||
[
|
||||
2,
|
||||
3,
|
||||
0,
|
||||
4,
|
||||
0,
|
||||
"IMAGE"
|
||||
],
|
||||
[
|
||||
3,
|
||||
2,
|
||||
1,
|
||||
4,
|
||||
1,
|
||||
"AUDIO"
|
||||
],
|
||||
[
|
||||
4,
|
||||
2,
|
||||
3,
|
||||
4,
|
||||
3,
|
||||
"FLOAT"
|
||||
]
|
||||
],
|
||||
"groups": [
|
||||
{
|
||||
"id": 1,
|
||||
"title": "1 · Your video",
|
||||
"bounding": [
|
||||
480,
|
||||
0,
|
||||
460,
|
||||
880
|
||||
],
|
||||
"color": "#1d5560",
|
||||
"flags": {}
|
||||
},
|
||||
{
|
||||
"id": 2,
|
||||
"title": "2 · Make it larger",
|
||||
"bounding": [
|
||||
960,
|
||||
0,
|
||||
400,
|
||||
880
|
||||
],
|
||||
"color": "#1d5560",
|
||||
"flags": {}
|
||||
},
|
||||
{
|
||||
"id": 3,
|
||||
"title": "3 · Result",
|
||||
"bounding": [
|
||||
1380,
|
||||
0,
|
||||
480,
|
||||
880
|
||||
],
|
||||
"color": "#1d5560",
|
||||
"flags": {}
|
||||
}
|
||||
],
|
||||
"definitions": {
|
||||
"subgraphs": []
|
||||
},
|
||||
"config": {},
|
||||
"extra": {
|
||||
"ds": {
|
||||
"scale": 0.73,
|
||||
"offset": [
|
||||
137,
|
||||
20
|
||||
]
|
||||
}
|
||||
},
|
||||
"version": 0.4
|
||||
}
|
||||
@@ -0,0 +1,142 @@
|
||||
# Edit Lock
|
||||
|
||||
Keeps every part of your picture that an edit did not change exactly where
|
||||
it was, with the colors it had.
|
||||
|
||||
An edit model redraws the whole picture, also the parts you did not ask it
|
||||
to change. Those parts often come back a little zoomed in, shifted or
|
||||
tinted. Flip between the original and the edit and the whole picture
|
||||
jumps. Edit Lock stops that: only the edit is new.
|
||||
|
||||
It needs no mask. It sees what the edit changes while the picture is being
|
||||
made, and holds on to everything else.
|
||||
|
||||
## Quick start
|
||||
|
||||
1. Put **Edit Lock 🆎** between the model and the sampler: the model (after
|
||||
any LoRA loader) goes in, and its `model` output goes to the sampler or
|
||||
guider.
|
||||
2. Wire the VAE into `vae` and the picture you are editing into `image`.
|
||||
Use the same picture the edit model gets, at the size of the result.
|
||||
3. After the VAE Decode, add **Edit Lock Restore 🆎**. Wire the decoded
|
||||
picture into `image` and Edit Lock's `edit_lock` output into
|
||||
`edit_lock`. This puts your own pixels back in the untouched parts.
|
||||
4. Wire Restore's `report` into **Show Text 🆎** to read how much was kept.
|
||||
|
||||
Run your edit as usual. Prompt, steps and seed stay as they were.
|
||||
|
||||
## What it does with each kind of edit
|
||||
|
||||
- **A local edit** (add glasses, repaint the car, change the sign): the
|
||||
edit happens, and the rest of the picture stays put.
|
||||
- **An edit of the whole picture** (a new style, new light, fog):
|
||||
there is nothing to keep, so Edit Lock steps aside and the edit runs as
|
||||
if the node were not there. The report says so.
|
||||
- **A soft effect over the whole picture** (rain, falling snow, haze):
|
||||
it touches everything a little, so there is no untouched part to put
|
||||
back. Edit Lock lets the effect through and only holds the picture in
|
||||
place, so it does not come back zoomed or shifted. The report says so.
|
||||
- **A soft effect and a clear change in one edit** (rain and an umbrella):
|
||||
the node may lock on the umbrella and hold the rain back everywhere
|
||||
else. If the effect is missing outside the changed part, set **Lock** to
|
||||
`off`.
|
||||
- **Something in between** (a new hair color on a close-up, a new
|
||||
background behind a person): it may lock or step aside. The report says
|
||||
which. These edits often change the light on the face too. `auto` steps
|
||||
aside for most of them, so the new light comes with the new hair.
|
||||
`always` keeps the old light on everything it keeps, which can look
|
||||
wrong on a face. A mask settles it: paint the hair and the face, and
|
||||
the rest is locked.
|
||||
- **The model ignores your picture and draws a new one:** that reads as a
|
||||
whole-picture edit, so `auto` steps aside. `always` holds the model to
|
||||
your picture.
|
||||
|
||||
## Controls
|
||||
|
||||
- **model**: the edit model. Patch last, after any LoRAs.
|
||||
- **vae**: the same VAE the workflow decodes with.
|
||||
- **image**: the picture you are editing, at the size of the result. If
|
||||
the sizes differ by more than a few percent the lock does nothing, and
|
||||
the report says so.
|
||||
- **Lock** (`lock`): `auto` (the default) locks the untouched parts and
|
||||
steps aside when the whole picture changes. `always` locks even then.
|
||||
`off` does nothing, which makes a quick before and after.
|
||||
- **Tolerance** (`tolerance`): how big a change must be to count as part of
|
||||
the edit. 1 is the default. Raise it if small unwanted changes get
|
||||
through. Lower it if the lock holds back part of your edit.
|
||||
- **mask** (optional): white marks where the edit may happen; everything
|
||||
else is locked. With a mask the node does not look for the change
|
||||
itself. A mask from **Mask by Touch 🆎** or **Mask by Name 🆎** works.
|
||||
|
||||
## Outputs
|
||||
|
||||
- **model**: the model with the lock on.
|
||||
- **edit_lock**: what the lock kept, for Edit Lock Restore 🆎.
|
||||
|
||||
## Good to know
|
||||
|
||||
- The lock belongs to one picture. For an edit with two pictures, give it
|
||||
the one the result should line up with.
|
||||
- It works on single pictures, not video.
|
||||
- The node watches the first steps of a run and locks for the rest, so
|
||||
the model finishes the edit on top of your real picture. Shadows and
|
||||
reflections the edit adds next to the changed part still come through.
|
||||
- If your edit is missing and the report says nearly all of the picture
|
||||
was kept, set **Lock** to `off`.
|
||||
- Tested on MiniMax-H3-Image. The lock itself is not tied to one model, but
|
||||
the line between "local edit" and "whole picture" was measured there.
|
||||
|
||||
## Technical details
|
||||
|
||||
At every step an edit model predicts the finished picture. Edit Lock
|
||||
compares that prediction with your picture in latent space.
|
||||
|
||||
**Watching.** The first prediction of a run is a blur, and rain, snow or
|
||||
the light a new hair color throws on a face only show a few steps in. So
|
||||
the node watches until the noise is down to 0.8 of where it began: step 4
|
||||
of a 6-step MiniMax-H3-Image run. A run that never gets that low decides
|
||||
at its last step. While it watches, the model draws its own picture and
|
||||
only one thing is corrected. The model often draws the whole picture about
|
||||
one percent larger, or a little shifted. The node measures that zoom and
|
||||
shift at every step, on the calmer half of the picture, and moves the
|
||||
prediction back so it sits where your picture does.
|
||||
|
||||
**Deciding.** Change is measured against the quiet part of the picture:
|
||||
the tenth where the prediction differs least. An even tint over the whole
|
||||
picture is taken out first, so a color cast is not read as a change. A
|
||||
spot counts as changed when it differs by 3 times the quiet level, or by 2
|
||||
times while it touches a spot that does. So a faint change beside a clear
|
||||
one belongs to the edit, and a faint change alone is redraw noise.
|
||||
**Tolerance** multiplies both numbers.
|
||||
|
||||
**Locking.** From the decision on, wherever the picture did not change,
|
||||
the original latent is put back before the sampler takes its next step. A
|
||||
part once read as changed stays changed for the rest of the run. The last
|
||||
step ends on the original there, so what was kept is the original itself,
|
||||
not a redrawn copy.
|
||||
|
||||
**Stepping aside.** `auto` does not lock in three cases:
|
||||
|
||||
- At the first step the whole picture is already changing (a new style,
|
||||
new light). The node reads the calmest fifth of a smoothed copy, where a
|
||||
slightly larger redraw nearly vanishes and a new style does not. On
|
||||
MiniMax-H3-Image, local edits at 1, 4 and 16 megapixels read 0.22 or
|
||||
less there and whole-picture edits 0.27 or more; the line sits at 0.243.
|
||||
The model then runs untouched.
|
||||
- At the decision no place has changed clearly: the strongest spot of the
|
||||
picture is under 4 times the quiet level. That is light rain or fog.
|
||||
Local edits read 4.7 or more there, glasses and makeup included.
|
||||
- At the decision the lock would hand a tenth of the picture or more to
|
||||
the edit, yet the strongest hundredth of the picture is under 4 times
|
||||
the quiet level. That is heavier rain or snow, and a lock would let it
|
||||
through in patches. Local edits that wide read 4.4 or more; rain and
|
||||
snow read 3.6 or less.
|
||||
|
||||
In the last two cases the node keeps undoing the zoom and shift to the end
|
||||
of the run, and puts nothing back. An edit close to one of these lines (a
|
||||
hair color change on a close-up, a new background, smoke) can fall on
|
||||
either side. The report says which, and **Lock** overrides it.
|
||||
|
||||
This is the target the TailSteer paper trains a model toward (the source
|
||||
where nothing changed, the model's own edit elsewhere), applied while
|
||||
sampling. It needs no LoRA and no training.
|
||||
@@ -0,0 +1,48 @@
|
||||
# Edit Lock Restore
|
||||
|
||||
Puts your picture's own pixels back in every part the edit did not change.
|
||||
|
||||
**Edit Lock 🆎** keeps the untouched parts in place while the picture is
|
||||
made. They still pass through the VAE, which softens fine detail a little.
|
||||
This node swaps your original pixels in for those parts, so they match
|
||||
your picture pixel for pixel. Only the edit is new.
|
||||
|
||||
## Quick start
|
||||
|
||||
1. Add it right after the VAE Decode, before any upscaler.
|
||||
2. Wire the decoded picture into `image`.
|
||||
3. Wire the `edit_lock` output of Edit Lock 🆎 into `edit_lock`.
|
||||
4. Wire `report` into **Show Text 🆎** to read what was kept.
|
||||
|
||||
## Controls
|
||||
|
||||
- **image**: the edited picture, straight from the VAE Decode.
|
||||
- **edit_lock**: from Edit Lock 🆎. It carries your original picture and
|
||||
the map of what the edit changed, so there is nothing else to wire.
|
||||
- **Feather** (`feather`): how softly your original fades into the edit,
|
||||
in pixels. 16 is the default. The fade lies outside the edit, so the
|
||||
edit itself is never mixed with the old picture.
|
||||
|
||||
## Outputs
|
||||
|
||||
- **image**: the edit, with your original pixels in the untouched parts.
|
||||
- **edit_mask**: white where the edit changed the picture. Use it to see
|
||||
what the lock decided, or for a later step that should only touch the
|
||||
edit.
|
||||
- **report**: one line, such as "Edit Lock kept 82% of your picture as it
|
||||
was and left 18% to the edit".
|
||||
|
||||
## When nothing was locked
|
||||
|
||||
The picture passes through unchanged, the mask is all white, and the
|
||||
report says why: the edit changed the whole picture, the lock was off, or
|
||||
the picture and the result were different sizes.
|
||||
|
||||
## Technical details
|
||||
|
||||
The map of what changed is made while sampling, on the latent grid. Here
|
||||
it is enlarged to the picture's size, widened by the feather and blurred
|
||||
by the same amount, and used to blend the decoded picture with your
|
||||
original. If the decoded picture is not the size of your original (after
|
||||
an upscale), the original is resized to fit first, which costs it some
|
||||
sharpness. That is why this node goes before an upscaler.
|
||||
@@ -0,0 +1,69 @@
|
||||
# FlashVSR Upscale
|
||||
|
||||
Makes a video larger and sharper. Give it the frames of a clip, set how much
|
||||
larger, and it returns the same clip with more detail. Use it as the last
|
||||
step before Save Video, for a generated clip or one you loaded.
|
||||
|
||||
It needs one model file and no other pack.
|
||||
|
||||
## The model file
|
||||
|
||||
Download `flashvsr_v1.1_fp16.safetensors` (3.5 GB) from
|
||||
[ausboss/FlashVSR-v1.1-ComfyUI](https://huggingface.co/ausboss/FlashVSR-v1.1-ComfyUI)
|
||||
and put it in `ComfyUI/models/diffusion_models`. Then refresh the page, or
|
||||
press `R`, and pick it under **Model**.
|
||||
|
||||
The file holds the official FlashVSR v1.1 weights in one piece. Another
|
||||
FlashVSR file will not load here: the node stops and says which part is
|
||||
missing.
|
||||
|
||||
## Controls
|
||||
|
||||
- **frames**: The clip, as a batch of frames. Load Video's frames output
|
||||
goes here.
|
||||
- **Model** (`model`): The model file above. The list shows the files in
|
||||
`diffusion_models` that have FlashVSR in their name.
|
||||
- **Scale** (`scale`, starts at `2`): How many times larger each side gets,
|
||||
from 1 to 4. At `1` the clip keeps its size and is only sharpened.
|
||||
- **Keep colors** (`color_match`, starts on): Brings each frame's colours
|
||||
back to those of your clip. Leave it on unless you want the model's own
|
||||
colours.
|
||||
- **vae** (optional): Connect the Wan 2.1 VAE for a slower, finer finish.
|
||||
With nothing connected, the fast decoder inside the model file makes the
|
||||
picture.
|
||||
|
||||
## Outputs
|
||||
|
||||
- **frames**: The larger frames, as many as went in, in the same order.
|
||||
The frame rate and the sound are not touched, so wire those from Load
|
||||
Video to Save Video as before.
|
||||
|
||||
## What to expect
|
||||
|
||||
FlashVSR redraws fine detail: skin, hair, cloth and edges come back
|
||||
sharper than the clip had them. It keeps the scene, the motion and the
|
||||
timing.
|
||||
|
||||
A long clip needs no more graphics memory than a short one. The node reads
|
||||
the clip a few frames at a time and keeps only the last few in memory.
|
||||
|
||||
## Technical details
|
||||
|
||||
- **Size.** The result is exactly your size times **Scale**, rounded to
|
||||
even numbers. Inside, the model works at a multiple of 128 pixels each
|
||||
way; the node pads the picture for that and cuts the padding off again,
|
||||
so nothing of your picture is cropped.
|
||||
- **Graphics memory.** About 10 GB for a 576 x 1024 result. A result too
|
||||
large for the free memory is made in overlapping tiles and blended, and
|
||||
the tile size is chosen from the memory that is free when the node runs.
|
||||
1152 x 2048 ran in two tiles at about 16 GB. On an 8 GB card (RTX 3070
|
||||
Laptop) the same result ran in smaller tiles at about 6.5 GB, at about 3
|
||||
seconds a frame.
|
||||
- **Speed.** On an RTX 5090: about 7 frames a second for a 576 x 1024
|
||||
result and about 2 frames a second for 1152 x 2048.
|
||||
- **One step.** The model runs once per group of frames. There are no
|
||||
steps, sampler, prompt or seed to set.
|
||||
- **The model.** FlashVSR v1.1 by Junhao Zhuang and the OpenImagingLab
|
||||
team, released under Apache-2.0. It is built on the Wan 2.1 1.3B video
|
||||
model. This node is the pack's own code in plain PyTorch.
|
||||
- The node makes no network requests and writes no files.
|
||||
@@ -19,8 +19,8 @@ What is lit in teal is what the mask will be. When it is right, press Run
|
||||
for the whole workflow.
|
||||
|
||||
The first time the pointer rests on a picture, the node looks at it. A
|
||||
small button under the picture shows how far along it is, **❚❚ 40/104**,
|
||||
filling up as it goes. The first things light up after a few seconds, and
|
||||
button under the picture shows how far along it is, **❚❚ Pause · looking
|
||||
40/104**, filling up as it goes. The first things light up after a few seconds, and
|
||||
it has seen the whole picture after about 20.
|
||||
|
||||
You do not have to wait for that. Whatever is lit under the pointer can be
|
||||
@@ -28,7 +28,7 @@ picked: a click takes it at once, also while the node is still looking.
|
||||
|
||||
Click that button to pause the looking when what you need is already lit:
|
||||
what is lit stays, and you go on picking and painting. It turns into
|
||||
**▶**; click it to look on from where it stopped. A click on a spot where
|
||||
**▶ Look on**; click it to look on from where it stopped. A click on a spot where
|
||||
nothing is lit yet looks on as well.
|
||||
|
||||
A click on a spot where nothing is lit yet is not lost: a dot marks it, the
|
||||
@@ -50,6 +50,9 @@ at that spot.
|
||||
A stroke across two things, like both shoes, takes both.
|
||||
- Did the thing not light up at all? Type its name into the field under the
|
||||
picture, like "screen" or "earring", and press Enter.
|
||||
- Is something in the mask that should not be, like the face or the hair of
|
||||
a person you picked? Click the **+** in front of that field so it shows
|
||||
**−**, type "face" and press Enter. Or hold Ctrl and click the face.
|
||||
- When nothing fits, switch to **Paint** and paint the area yourself.
|
||||
|
||||
## The whole subject, or the background
|
||||
@@ -59,6 +62,10 @@ click: the person, the animal, the product. **Background** picks everything
|
||||
else. They are picks like any other: click more things to add them, or
|
||||
Ctrl-click a thing to take it away again.
|
||||
|
||||
Ctrl-click takes away the thing under the pointer, where it is in the mask.
|
||||
On a spot where nothing is picked it changes nothing, and the picture says
|
||||
so.
|
||||
|
||||
Both need ComfyUI's background-removal model, see "The model for Fine
|
||||
edges" below.
|
||||
|
||||
@@ -115,8 +122,11 @@ or an area you painted.
|
||||
it.
|
||||
- **all 3** on a chip takes every thing of that name instead of the one you
|
||||
clicked: all five birds, both shoes.
|
||||
- The field at the end, **+ type a name**, picks every thing of the name
|
||||
you type. With Erase chosen it takes them away.
|
||||
- The field at the end picks every thing of the name you type. The sign in
|
||||
front of it says what happens to the name: **+** adds it, **−** takes it
|
||||
away. Click the sign to change it. A name typed with a minus in front,
|
||||
like "-face", is taken away whatever the sign shows. With Erase chosen the
|
||||
sign starts on **−**.
|
||||
|
||||
## The tools on the node
|
||||
|
||||
@@ -273,6 +283,15 @@ then.
|
||||
|
||||
### How SAM 3 is asked
|
||||
|
||||
- A picture whose long side is over 2048 pixels is asked about on a copy
|
||||
with its long side at 2048. SAM 3 reads every picture at 1008 pixels
|
||||
anyway, and returns each answer at the size it was given, on the graphics
|
||||
card: at 8192 pixels one word with many matches needed over 27 GB. The
|
||||
answer is brought back to the picture's own size, so the mask, the cut-out
|
||||
and your touches stay in pixels of the picture. Paint is never scaled.
|
||||
Measured on a 3840 pixel picture, five words: the new masks share 95 to
|
||||
99.6 % of their area with the old ones, the edge sits a median 0 to 2
|
||||
pixels differently, and a look takes 31 seconds in place of 53.
|
||||
- The search is ComfyUI's own SAM 3: SAM3 Detect, asked about points, or
|
||||
about a word that CLIP Text Encode has read. Both use ComfyUI's own
|
||||
settings (threshold 0.5, two edge passes). The node needs ComfyUI 0.20 or
|
||||
|
||||
@@ -0,0 +1,54 @@
|
||||
# Prompt Variations
|
||||
|
||||
A prompt box that changes on every run. Write choices in curly brackets and
|
||||
the node picks one from each set:
|
||||
|
||||
```text
|
||||
A woman in a {plain|striped|patterned} {tank top|sweater|button-up shirt}
|
||||
and {jeans|a pleated skirt|cargo shorts}, {standing|sitting on a stool|leaning on the wall}.
|
||||
```
|
||||
|
||||
One run might give "A woman in a striped sweater and jeans, sitting on a
|
||||
stool." The next gives another mix.
|
||||
|
||||
Use it to make many different pictures in one go. Wire in the seed, set the
|
||||
Seed node to **randomize**, and set the number of runs next to the Run
|
||||
button to 50. Each run gets a new seed, so each picture gets new choices.
|
||||
|
||||
## Controls
|
||||
|
||||
- The text box (`text`): your prompt. Put each set of choices in curly
|
||||
brackets and split the choices with `|`. Write as many sets as you like,
|
||||
each with as many choices as you like.
|
||||
- `seed` (a socket): the seed your sampler uses. A new seed picks new
|
||||
choices. The same seed picks the same ones again, so a picture you like
|
||||
can be made again from its seed.
|
||||
|
||||
## Output
|
||||
|
||||
- **prompt**: your text with one choice picked from each set. Wire it to
|
||||
Show Text to see what was picked.
|
||||
|
||||
## Tips
|
||||
|
||||
- An empty choice lets a part come and go: `a {|striped} shirt` gives "a
|
||||
shirt" or "a striped shirt".
|
||||
- Sets can go inside sets: `{a {red|blue} hat|a scarf}`.
|
||||
- Text without curly brackets comes out exactly as typed.
|
||||
- To make one choice come up more often, write it more than once:
|
||||
`{jeans|jeans|skirt}`.
|
||||
|
||||
## Technical details
|
||||
|
||||
Only brackets that hold a `|` are choices. Brackets without one are left as
|
||||
they are, so a JSON prompt such as `{"light": "soft"}` passes through
|
||||
unchanged. A `|` inside JSON brackets still counts as a choice.
|
||||
|
||||
The innermost sets are picked first, from left to right, with Python's
|
||||
`random.Random(seed)`. The same text and seed always give the same prompt,
|
||||
on any machine. Spaces around each choice are trimmed, and an empty pick
|
||||
also drops the one space in front of it.
|
||||
|
||||
The picks happen when the prompt runs, on the server, not in the browser.
|
||||
ComfyUI's own `{a|b}` choices in some text boxes are picked by the browser
|
||||
when you queue, with no seed; this node does not use them.
|
||||
@@ -50,8 +50,18 @@ file's name. A caption on the `caption_text` input writes the paired `.txt`.
|
||||
shorter; on again, the node returns to its height with the last saved
|
||||
image. It is display only: the file is saved either way and the queue
|
||||
still lists it. When a run saves more than one image (a batch, or a list
|
||||
such as a folder of images), arrows beside the switch (‹ 2 / 5 ›) step
|
||||
through them.
|
||||
such as a folder of images), they all show side by side. The arrows
|
||||
beside the switch (‹ all 5 ›) go on to each one by itself (‹ 2 / 5 ›) and
|
||||
back round to all of them. A run of more than 48 shows one at a time.
|
||||
- **Auto save** (`autosave`, on by default): on, every run saves its
|
||||
pictures. Off, a run only shows them here and saves nothing. A **SAVE**
|
||||
button then sits over the picture: **SAVE ALL** keeps every picture of
|
||||
the run, and with the arrows on one picture, **SAVE** keeps that one.
|
||||
Use it to make a batch, look, and keep only the good ones. The pictures
|
||||
are saved under the name and folder set above, in the format picked,
|
||||
with the workflow inside if **Embed workflow** is on. They wait in
|
||||
ComfyUI's temp folder, which is emptied when ComfyUI restarts, so save
|
||||
what you want before that.
|
||||
- **Preview steps** (`step_preview`, off by default): the switch left of
|
||||
**preview** shows the picture while it is being made, one step at a time,
|
||||
where the saved image goes, with a **STOP** button on the bar. Stop a run
|
||||
|
||||
@@ -0,0 +1,60 @@
|
||||
# Text Lines
|
||||
|
||||
Runs what comes after once for each line of your text. Type five prompts,
|
||||
one on each line, press Run once, and you get five pictures.
|
||||
|
||||
Use it for a list of prompts you typed, or for a list another node wrote,
|
||||
such as a text model asked for "five different poses".
|
||||
|
||||
## How to use it
|
||||
|
||||
1. Type your lines in the box, or wire text into **text**.
|
||||
2. Wire **line** to where one prompt would go, such as a text encoder's
|
||||
prompt.
|
||||
3. Run. Every node after it runs once for each line, and a Save Image at the
|
||||
end saves one picture for each.
|
||||
|
||||
Empty lines are skipped. Models and pictures that do not depend on the line
|
||||
are loaded once, not once for each line.
|
||||
|
||||
## Controls
|
||||
|
||||
- The text box (`text`): One line for each run.
|
||||
- **Split at** (`split`): **each line** makes every line its own run.
|
||||
**blank line** keeps lines together and starts the next run at an empty
|
||||
line. Pick it when one prompt is several lines long.
|
||||
- **Numbering** (`remove_numbering`): **remove** takes a number or bullet
|
||||
off the start of each line, such as `1.`, `2)`, `-` or `[prompt 3]`.
|
||||
Lists a text model writes usually have them. **keep** leaves every line
|
||||
as written.
|
||||
- **Limit** (`limit`): The most lines to run. 0 runs them all. Set it when
|
||||
a text model might write more lines than you asked for.
|
||||
|
||||
## Outputs
|
||||
|
||||
- **line**: Each line, one at a time.
|
||||
- **number**: The line's place in the list, starting at 1. Add it to a seed
|
||||
with Math Expression (`a + b`) to give each line its own seed.
|
||||
- **count**: How many lines run.
|
||||
|
||||
## Notes
|
||||
|
||||
- With no lines at all, the run stops at this node and says so.
|
||||
- Every run uses the same seed unless you add **number** to it. Different
|
||||
prompts still give different pictures.
|
||||
- A line such as "2 women on a beach" or "1.5x zoom" keeps its number: only
|
||||
a number followed by a period, bracket or colon counts as numbering.
|
||||
|
||||
## Technical details
|
||||
|
||||
**line** and **number** are ComfyUI lists (`OUTPUT_IS_LIST`), the same way
|
||||
Image Folder hands out pictures. ComfyUI runs each later node once for each
|
||||
item and keeps the results in order. A node that takes a whole list at once,
|
||||
such as core Rebatch Images, gets all of them together.
|
||||
|
||||
Numbering that is removed: a number in square, round or CJK brackets with an
|
||||
optional word before it (`[prompt 1]`, `(2)`), a word and a number followed by
|
||||
a colon, period, bracket or dash (`Shot 3:`), a bare number followed by a
|
||||
period, bracket or colon that has no digit right after it, and the bullets
|
||||
`-`, `*` and `•` followed by a space. Only one marker is removed from each
|
||||
line.
|
||||
@@ -24,6 +24,7 @@ help panel.
|
||||
| Get a mask by typing what to find | Mask by Name (`AUSBOSS_NODES_MaskByName`) |
|
||||
| Get a mask by clicking the thing | Mask by Touch (`AUSBOSS_NODES_MaskByTouch`) |
|
||||
| Grow, shrink, fill or smooth a mask | Mask Refine (`AUSBOSS_NODES_RefineMask`) |
|
||||
| Keep the untouched parts of a picture exactly as they were while an edit model edits it | Edit Lock (`AUSBOSS_NODES_EditLock`) then Edit Lock Restore (`AUSBOSS_NODES_EditLockRestore`) |
|
||||
| Line an edited picture back up with the original | Realign to Source (`AUSBOSS_NODES_RealignToSource`, experimental) |
|
||||
| Match one picture's colors to another | Color Match (`AUSBOSS_NODES_ColorMatch`) |
|
||||
| Upscale big pictures in tiles | Tiled Upscale (`AUSBOSS_NODES_TiledUpscale`) then Tiled Upscale Stitch (`AUSBOSS_NODES_TiledUpscaleStitch`) |
|
||||
@@ -37,8 +38,11 @@ help panel.
|
||||
| Load a sound file and keep only part of it, up to a set length | Load Audio (`AUSBOSS_NODES_LoadAudio`) |
|
||||
| Turn, crop and pad a video: one frame, or the whole clip | Video Crop + Rotate + Pad → Frame (`AUSBOSS_NODES_VideoCropRotatePad`), → Clip (`AUSBOSS_NODES_VideoCropRotatePadClip`) |
|
||||
| Change a clip's frame rate | Frame Interpolate (`AUSBOSS_NODES_FrameInterpolate`) |
|
||||
| Make a clip larger and sharper | FlashVSR Upscale (`AUSBOSS_NODES_FlashVSRUpscale`) |
|
||||
| Work with frame batches | Select Frame, Select Every Nth, Split Batch, Merge Batches (`AUSBOSS_NODES_SelectFrame`, `_SelectEveryNth`, `_SplitBatch`, `_MergeBatches`) |
|
||||
| Run a workflow once for each line of text, such as a list of prompts | Text Lines (`AUSBOSS_NODES_TextLines`) |
|
||||
| One shared seed, number or text | Seed, Integer, Float, Text (`AUSBOSS_NODES_Seed`, `_Integer`, `_Float`, `_Text`) |
|
||||
| A prompt that picks from `{a\|b\|c}` choices on each run, by seed | Prompt Variations (`AUSBOSS_NODES_PromptVariations`) |
|
||||
| Do maths on numbers | Math Expression (`AUSBOSS_NODES_MathExpression`) |
|
||||
| Show text on the node, or the picture while it is made | Show Text (`AUSBOSS_NODES_ShowText`), Step Preview (`AUSBOSS_NODES_StepPreview`) |
|
||||
| Free graphics memory between steps | Free Memory (`AUSBOSS_NODES_FreeMemory`) |
|
||||
|
||||
+199
-10
@@ -26,7 +26,7 @@ import { suppressCoreImagePreview } from "../shared/core_preview.mjs";
|
||||
import { autoMaskValues } from "../shared/mask_auto.mjs";
|
||||
import { hideInputsInDef, hideWidget, setWidgetVisible } from "../shared/widget_visibility.mjs";
|
||||
import { makePictureStepper } from "../shared/picture_stepper.mjs";
|
||||
import { keepIndex, stepIndex } from "../shared/picture_steps.mjs";
|
||||
import { ALL, allLabel, gridShape, heldPlaces, keepIndex, stepIndex, stepWithAll } from "../shared/picture_steps.mjs";
|
||||
import { addStepConsumer, mountStepOverlay, showsStepsOn, stepRun, stopStepRun, watchSteps } from "../shared/step_preview.mjs";
|
||||
import {
|
||||
describeNodePreview,
|
||||
@@ -55,6 +55,11 @@ const BAR_PADDING = 2;
|
||||
const PANEL_HEIGHT = STAGE_HEIGHT + BAR_HEIGHT + 4 + 16;
|
||||
const OFF_HEIGHT = BAR_HEIGHT + 16;
|
||||
const INPUT_SIDE = 1; // LiteGraph.INPUT
|
||||
// A run that left this many pictures or fewer opens with all of them side
|
||||
// by side. More than that and each would be a speck, so the arrows go
|
||||
// through them one at a time, as before.
|
||||
const GRID_MAX = 48;
|
||||
const GRID_GAP = 4;
|
||||
// The node's height from before its preview was switched off, so switching
|
||||
// it back on returns the node to it (saved with the workflow).
|
||||
const HEIGHT_PROPERTY = "ausboss_preview_height";
|
||||
@@ -114,7 +119,17 @@ const NODE_CONFIG = {
|
||||
},
|
||||
AUSBOSS_NODES_SelectFrame: { inputName: "frames", noun: "frames" },
|
||||
// Its result is the file it just saved; the card sits above the panel.
|
||||
AUSBOSS_NODES_SaveImage: { inputName: "images", noun: "an image", switchTitle: "Show the saved image here. The file is saved either way.", steps: true },
|
||||
// SAVE shows only for a run made with Auto save off, which holds its
|
||||
// pictures until this button keeps them.
|
||||
AUSBOSS_NODES_SaveImage: {
|
||||
inputName: "images",
|
||||
noun: "an image",
|
||||
switchTitle: "Show the saved image here. The file is saved either way.",
|
||||
steps: true,
|
||||
tools: [
|
||||
{ label: "SAVE", title: "Save what is showing: every picture, or the one the arrows picked.", action: keepHeld, held: true },
|
||||
],
|
||||
},
|
||||
};
|
||||
|
||||
// The picture sits out of flow, centred in the stage, so it can never set
|
||||
@@ -138,10 +153,16 @@ function ensureCss() {
|
||||
.ausboss-input-preview-stage img,.ausboss-input-preview-stage video{display:none;position:absolute;inset:0;margin:auto;max-width:100%;max-height:100%;object-fit:contain;}
|
||||
.ausboss-input-preview-stage.show-image img{display:block;}
|
||||
.ausboss-input-preview-stage.show-video video{display:block;}
|
||||
.ausboss-input-preview-grid{display:none;position:absolute;inset:0;box-sizing:border-box;gap:${GRID_GAP}px;padding:${GRID_GAP}px;justify-content:center;align-content:center;}
|
||||
.ausboss-input-preview-stage.show-grid .ausboss-input-preview-grid{display:grid;}
|
||||
.ausboss-input-preview-stage .ausboss-input-preview-grid img{display:block;position:static;margin:0;width:100%;height:100%;min-width:0;min-height:0;max-width:none;max-height:none;object-fit:contain;}
|
||||
.ausboss-input-preview-hint{display:none;max-width:86%;color:${T.faint};font:11px/1.35 ${FONT_UI};text-align:center;}
|
||||
.ausboss-input-preview-stage:not(.show-image):not(.show-video) .ausboss-input-preview-hint{display:block;}
|
||||
.ausboss-input-preview-stage:not(.show-image):not(.show-video):not(.show-grid) .ausboss-input-preview-hint{display:block;}
|
||||
.ausboss-input-preview-tools{display:flex;gap:4px;pointer-events:auto;}
|
||||
.ausboss-input-preview-tool{box-sizing:border-box;height:${BAR_HEIGHT}px;min-width:24px;padding:0 6px;border:1px solid transparent;border-radius:${RADIUS_XS}px;background:${T.fill};color:${T.label};font:600 9.5px/18px ${FONT_UI};letter-spacing:.02em;cursor:pointer;}
|
||||
.ausboss-input-preview-tool[hidden]{display:none;}
|
||||
.ausboss-input-preview-tool.keep{color:${T.brandInk};background:${T.brandSoft};}
|
||||
.ausboss-input-preview-stage .ausboss-input-preview-grid img.kept{outline:2px solid ${T.brand};outline-offset:-2px;}
|
||||
.ausboss-input-preview-tool:hover{color:${T.ink};background:${T.fillHover};}
|
||||
.ausboss-input-preview-tool.active{color:${T.brandInk};background:${T.brandSoft};}
|
||||
.ausboss-input-preview-tool:focus-visible{outline:2px solid ${T.brand};outline-offset:1px;}
|
||||
@@ -488,13 +509,14 @@ function buildTools(state, tools, signal) {
|
||||
tool.action(state);
|
||||
}, { signal });
|
||||
state.toolButtons[tool.label] = button;
|
||||
if (tool.held) { button.hidden = true; state.keepButton = button; }
|
||||
bar.append(button);
|
||||
}
|
||||
return bar;
|
||||
}
|
||||
|
||||
function showHint(state, text) {
|
||||
state.stage.classList.remove("show-image", "show-video");
|
||||
state.stage.classList.remove("show-image", "show-video", "show-grid");
|
||||
state.hint.textContent = text;
|
||||
state.placeholder = null;
|
||||
releaseSteps(state);
|
||||
@@ -509,6 +531,7 @@ function showPlaceholder(state, connected) {
|
||||
function showMedia(state, kind) {
|
||||
state.stage.classList.toggle("show-image", kind === "image");
|
||||
state.stage.classList.toggle("show-video", kind === "video");
|
||||
state.stage.classList.toggle("show-grid", kind === "grid");
|
||||
state.placeholder = null;
|
||||
releaseSteps(state);
|
||||
}
|
||||
@@ -517,6 +540,53 @@ function clearMedia(state) {
|
||||
state.img.removeAttribute("src");
|
||||
state.video.pause?.();
|
||||
state.video.removeAttribute("src");
|
||||
clearGrid(state);
|
||||
}
|
||||
|
||||
function clearGrid(state) {
|
||||
if (!state.gridKey) return;
|
||||
state.gridKey = "";
|
||||
state.grid?.replaceChildren();
|
||||
}
|
||||
|
||||
// Every picture of the run at once, each as big as the stage allows.
|
||||
function layoutGrid(state) {
|
||||
const cells = [...(state.grid?.children ?? [])];
|
||||
if (!cells.length) return;
|
||||
const loaded = cells.find((cell) => cell.naturalWidth > 0 && cell.naturalHeight > 0);
|
||||
const aspect = loaded ? loaded.naturalWidth / loaded.naturalHeight : 1;
|
||||
const { columns, rows } = gridShape(cells.length, state.stage.clientWidth, state.stage.clientHeight, aspect);
|
||||
// Tracks the size of one picture, so the pictures sit together in the
|
||||
// middle of the stage with no bands between the rows.
|
||||
const wide = (state.stage.clientWidth - GRID_GAP * (columns + 1)) / columns;
|
||||
const tall = (state.stage.clientHeight - GRID_GAP * (rows + 1)) / rows;
|
||||
if (!(wide > 0) || !(tall > 0)) return;
|
||||
const width = Math.floor(Math.min(wide, tall * aspect));
|
||||
const height = Math.floor(width / aspect);
|
||||
state.grid.style.gridTemplateColumns = `repeat(${columns}, ${width}px)`;
|
||||
state.grid.style.gridTemplateRows = `repeat(${rows}, ${height}px)`;
|
||||
}
|
||||
|
||||
function showGrid(state, count) {
|
||||
const urls = Array.from({ length: count }, (_, index) => storedResultUrl(state.node, index) ?? "");
|
||||
const key = urls.join("\n");
|
||||
if (state.gridKey !== key) {
|
||||
state.gridKey = key;
|
||||
state.grid.replaceChildren(...urls.map((url, index) => {
|
||||
const cell = document.createElement("img");
|
||||
cell.alt = `Picture ${index + 1} of ${count}`;
|
||||
cell.decoding = "async";
|
||||
// The first picture to load tells the grid the shape to lay out for.
|
||||
cell.addEventListener("load", () => layoutGrid(state), { once: true });
|
||||
cell.src = url;
|
||||
return cell;
|
||||
}));
|
||||
}
|
||||
state.img.removeAttribute("src");
|
||||
state.video.pause?.();
|
||||
state.video.removeAttribute("src");
|
||||
layoutGrid(state);
|
||||
showMedia(state, "grid");
|
||||
}
|
||||
|
||||
// Each source has one hook, with subscribers that can leave in any order.
|
||||
@@ -569,8 +639,9 @@ function storedResultUrl(node, index = null) {
|
||||
}
|
||||
|
||||
// Several pictures from one run (a list, or a batch Save Image saved one by
|
||||
// one): the arrows in the bar pick which one shows. One picture keeps the
|
||||
// panel as it was, with no arrows.
|
||||
// one): they all show side by side, and the arrows in the bar go on to each
|
||||
// one by itself and back round to all of them. One picture keeps the panel
|
||||
// as it was, with no arrows.
|
||||
function resultCount(state) {
|
||||
if (state.stale || state.resultGone) return 0;
|
||||
return outputRecordCount(app.nodeOutputs?.[outputLocatorId(state.node)]);
|
||||
@@ -586,10 +657,107 @@ function previewMinWidth(state) {
|
||||
return Math.max(PANEL_MIN_WIDTH, Math.ceil(state.node.size[0] - state.bar.clientWidth + needed));
|
||||
}
|
||||
|
||||
// --- pictures held until Save (Save Image with Auto save off) ---------------
|
||||
|
||||
// The pictures this node is holding, in the order it shows them. Empty when
|
||||
// the last run saved its pictures itself.
|
||||
function heldPictures(state) {
|
||||
if (state.stale || state.resultGone) return [];
|
||||
return heldPlaces(app.nodeOutputs?.[outputLocatorId(state.node)]);
|
||||
}
|
||||
|
||||
// Which of the shown pictures are kept, by their place on the node.
|
||||
function keptOf(state, held) {
|
||||
const key = held.map((picture) => `${picture.token}_${picture.index}`).join(" ");
|
||||
if (state.kept?.key !== key) state.kept = { key, indexes: new Set() };
|
||||
return state.kept.indexes;
|
||||
}
|
||||
|
||||
// A click keeps every picture, or only the one the arrows picked. With the
|
||||
// preview off there is no picture to pick, so it keeps them all.
|
||||
function keepsAll(state, count) {
|
||||
return count < 2 || showsAll(state, count) || !previewEnabled(state);
|
||||
}
|
||||
|
||||
// The button says what a click will do: SAVE ALL over the grid, SAVE over
|
||||
// one picture, SAVED once there is nothing left to keep there.
|
||||
function syncHeld(state, count = resultCount(state)) {
|
||||
const button = state.keepButton;
|
||||
if (!button) return;
|
||||
const held = heldPictures(state);
|
||||
button.hidden = !held.length || count < 1;
|
||||
if (button.hidden) return;
|
||||
const kept = keptOf(state, held);
|
||||
const all = keepsAll(state, count);
|
||||
const done = all ? kept.size >= held.length : kept.has(state.index);
|
||||
button.textContent = state.keeping ? "SAVING" : done ? "SAVED" : all && count > 1 ? "SAVE ALL" : "SAVE";
|
||||
button.classList.toggle("keep", !done);
|
||||
button.disabled = Boolean(state.keeping);
|
||||
[...(state.grid?.children ?? [])].forEach((cell, index) => cell.classList.toggle("kept", kept.has(index)));
|
||||
}
|
||||
|
||||
async function keepHeld(state) {
|
||||
const held = heldPictures(state);
|
||||
if (!held.length || state.keeping) return;
|
||||
const count = resultCount(state);
|
||||
const wanted = keepsAll(state, count) ? held.map((_, place) => place) : [state.index];
|
||||
// One request for each run, in the order the pictures show.
|
||||
const runs = new Map();
|
||||
for (const place of wanted) {
|
||||
const picture = held[place];
|
||||
if (!picture) continue;
|
||||
if (!runs.has(picture.token)) runs.set(picture.token, []);
|
||||
runs.get(picture.token).push(picture.index);
|
||||
}
|
||||
state.keeping = true;
|
||||
syncHeld(state);
|
||||
const kept = keptOf(state, held);
|
||||
const saved = [];
|
||||
let skipped = 0;
|
||||
try {
|
||||
for (const [token, picks] of runs) {
|
||||
const response = await api.fetchApi("/ausboss/save_image/keep", {
|
||||
method: "POST",
|
||||
headers: { "Content-Type": "application/json" },
|
||||
body: JSON.stringify({ token, picks }),
|
||||
});
|
||||
const answer = await response.json().catch(() => ({}));
|
||||
if (!response.ok) throw new Error(answer.error || "Could not save.");
|
||||
if (!state.alive) return;
|
||||
const done = new Set(answer.kept ?? []);
|
||||
held.forEach((picture, place) => {
|
||||
if (picture.token === token && done.has(picture.index)) kept.add(place);
|
||||
});
|
||||
saved.push(...(answer.saved ?? []));
|
||||
skipped += answer.skipped?.length ?? 0;
|
||||
}
|
||||
toast(state, saved.length === 1
|
||||
? `Saved ${saved[0].path}`
|
||||
: saved.length
|
||||
? `Saved ${saved.length} pictures`
|
||||
: skipped ? "Not saved: the file already exists" : "Already saved");
|
||||
} catch (error) {
|
||||
if (state.alive) toast(state, error?.message || "Could not save.");
|
||||
} finally {
|
||||
state.keeping = false;
|
||||
if (state.alive) syncHeld(state);
|
||||
}
|
||||
}
|
||||
|
||||
function showsAll(state, count) {
|
||||
return count > 1 && count <= GRID_MAX && state.all !== false;
|
||||
}
|
||||
|
||||
function stepPicture(state, step) {
|
||||
const count = resultCount(state);
|
||||
if (count < 2) return;
|
||||
state.index = stepIndex(state.index, count, step);
|
||||
if (count > GRID_MAX) {
|
||||
state.index = stepIndex(state.index, count, step);
|
||||
} else {
|
||||
const at = stepWithAll(showsAll(state, count) ? ALL : state.index, count, step);
|
||||
state.all = at === ALL;
|
||||
if (!state.all) state.index = at;
|
||||
}
|
||||
refresh(state);
|
||||
}
|
||||
|
||||
@@ -597,6 +765,7 @@ function refresh(state) {
|
||||
if (!state.alive) return;
|
||||
if (!previewEnabled(state)) {
|
||||
clearMedia(state);
|
||||
syncHeld(state);
|
||||
return;
|
||||
}
|
||||
const source = upstreamNode(state.node, state.inputName);
|
||||
@@ -604,7 +773,9 @@ function refresh(state) {
|
||||
const count = resultCount(state);
|
||||
const hidden = state.stepper.root.hidden;
|
||||
state.index = keepIndex(state.index, count);
|
||||
state.stepper.set(state.index, count);
|
||||
const all = showsAll(state, count);
|
||||
state.stepper.set(state.index, count, all ? allLabel(count) : undefined);
|
||||
syncHeld(state, count);
|
||||
if (state.stepper.root.hidden !== hidden) growToNodeFloor(state.widget);
|
||||
if (state.stale) {
|
||||
// The node ran while its preview was off, so the picture it still has
|
||||
@@ -616,6 +787,13 @@ function refresh(state) {
|
||||
// A node that just stopped the run has no result for this picture: the
|
||||
// one it still holds is from an earlier run (Mask by Name's last find,
|
||||
// tinted), so show what feeds it instead.
|
||||
if (all) {
|
||||
showGrid(state, count);
|
||||
// The grid's cells are new: mark the ones already kept.
|
||||
syncHeld(state, count);
|
||||
return;
|
||||
}
|
||||
clearGrid(state);
|
||||
const described = state.resultGone
|
||||
? describeSourcePreview(source)
|
||||
: count > 1
|
||||
@@ -701,7 +879,9 @@ function buildPanel(node, config) {
|
||||
const toastEl = document.createElement("div");
|
||||
toastEl.className = "ausboss-input-preview-toast";
|
||||
toastEl.style.display = "none";
|
||||
stage.append(img, video, hint, toastEl);
|
||||
const grid = document.createElement("div");
|
||||
grid.className = "ausboss-input-preview-grid";
|
||||
stage.append(img, video, grid, hint, toastEl);
|
||||
const bar = document.createElement("div");
|
||||
bar.className = "ausboss-input-preview-bar";
|
||||
root.append(bar, stage);
|
||||
@@ -735,7 +915,10 @@ function buildPanel(node, config) {
|
||||
noun: config.noun,
|
||||
switchTitle: config.switchTitle,
|
||||
advanced: config.advanced,
|
||||
root, bar, stage, img, video, hint, widget, abort,
|
||||
root, bar, stage, img, video, grid, hint, widget, abort,
|
||||
// Several pictures open side by side; the arrows pick one (false).
|
||||
all: true,
|
||||
gridKey: "",
|
||||
toast: toastEl,
|
||||
toastTimer: 0,
|
||||
toolButtons: null,
|
||||
@@ -769,6 +952,12 @@ function buildPanel(node, config) {
|
||||
installPreviewSwitch(state, abort.signal);
|
||||
if (config.steps) installStepsSwitch(state, abort.signal);
|
||||
|
||||
// The grid is laid out for the stage's shape, so a resized node lays it out again.
|
||||
if (typeof ResizeObserver === "function") {
|
||||
const watcher = new ResizeObserver(() => { if (state.alive && state.gridKey) layoutGrid(state); });
|
||||
watcher.observe(stage);
|
||||
abort.signal.addEventListener("abort", () => watcher.disconnect(), { once: true });
|
||||
}
|
||||
img.addEventListener("load", () => showMedia(state, "image"));
|
||||
img.addEventListener("error", () => {
|
||||
if (img.getAttribute("src")) showPlaceholder(state, true);
|
||||
|
||||
@@ -47,6 +47,24 @@ async function stopLook(node) {
|
||||
return Boolean(running || waiting.length);
|
||||
}
|
||||
|
||||
// Take this node's own runs that still wait out of the queue: the looks and
|
||||
// touches its picture asked for (they carry an `ask`), never a run of the
|
||||
// workflow. Says how many went. The one that is running is left alone.
|
||||
async function dropWaiting(node) {
|
||||
const id = executionIdOf(app.rootGraph ?? app.graph, node);
|
||||
if (!id) return 0;
|
||||
const queue = await (await api.fetchApi("/queue")).json();
|
||||
const ours = (entry) => {
|
||||
const step = entry?.[2]?.[id];
|
||||
return step?.class_type === NODE_CLASS && String(step.inputs?.ask ?? "") !== "";
|
||||
};
|
||||
const waiting = (queue?.queue_pending ?? []).filter(ours).map((entry) => entry[1]);
|
||||
if (waiting.length) {
|
||||
await api.fetchApi("/queue", { method: "POST", headers: { "Content-Type": "application/json" }, body: JSON.stringify({ delete: waiting }) });
|
||||
}
|
||||
return waiting.length;
|
||||
}
|
||||
|
||||
// The picture goes above the settings: the stage widget sits ahead of the
|
||||
// node's card, whichever of the two was built first.
|
||||
function stageFirst(node) {
|
||||
@@ -92,6 +110,7 @@ function build(node) {
|
||||
const face = mountTouchStage(node, {
|
||||
lookOnButton: () => lookOnButton,
|
||||
stopLook: () => stopLook(node).catch(() => false),
|
||||
dropWaiting: () => dropWaiting(node).catch(() => 0),
|
||||
// A workflow is being opened: nothing is run on the person's behalf.
|
||||
loading: () => Boolean(app.configuringGraph),
|
||||
// Bypassed or muted, or inside a box that is: nothing is run for it.
|
||||
|
||||
@@ -26,6 +26,7 @@ import { nodeByExecutionId } from "../shared/graph_ids.mjs";
|
||||
import {
|
||||
collapseRepeatedErrors,
|
||||
emptySourceMessage,
|
||||
explainRunError,
|
||||
findEmptySources,
|
||||
nodeLacksSource,
|
||||
} from "../shared/run_check.mjs";
|
||||
@@ -174,10 +175,20 @@ function collapseAnswers() {
|
||||
api.queuePrompt = collapsing;
|
||||
}
|
||||
|
||||
// A run that stops on an error the pack knows the cause of gets one toast
|
||||
// that says what to do. ComfyUI's own error report is left as it is.
|
||||
function explainKnownErrors() {
|
||||
api.addEventListener("execution_error", ({ detail }) => {
|
||||
const known = explainRunError(detail);
|
||||
if (known) showToast({ severity: "warn", summary: known.summary, detail: known.detail, life: 30000 });
|
||||
});
|
||||
}
|
||||
|
||||
app.registerExtension({
|
||||
name: "ausboss.run_check",
|
||||
setup() {
|
||||
guardRuns();
|
||||
collapseAnswers();
|
||||
explainKnownErrors();
|
||||
},
|
||||
});
|
||||
|
||||
+31
-3
@@ -33,7 +33,7 @@ const CARD_HEIGHT = 288;
|
||||
const WRAPPER_INSET = 16;
|
||||
const HIDDEN = [
|
||||
"filename_prefix", "format", "save_metadata", "exact_name", "on_existing",
|
||||
"output_dir", "caption", "name_counter", "name_date", "name_time", "name_size", "name_batch",
|
||||
"output_dir", "caption", "name_counter", "name_date", "name_time", "name_size", "name_batch", "autosave",
|
||||
];
|
||||
const FONT = `-apple-system,BlinkMacSystemFont,"Segoe UI",sans-serif`;
|
||||
const MONO = `ui-monospace,SFMono-Regular,Menlo,Consolas,monospace`;
|
||||
@@ -60,7 +60,10 @@ function ensureCss() {
|
||||
.ausboss-save-chip:hover{color:#fff;border-color:${BRAND}}
|
||||
.ausboss-save-chip.on{background:${BRAND};border-color:${BRAND};color:#04201d}
|
||||
.ausboss-save-chips.off .ausboss-save-chip{opacity:.4;pointer-events:none}
|
||||
.ausboss-save-embed-label{flex:0 0 30%;color:#8ba3a1;font-size:11px;user-select:none}
|
||||
.ausboss-save-embed-label{flex:none;color:#8ba3a1;font-size:11px;white-space:nowrap;user-select:none}
|
||||
.ausboss-save-half{display:flex;align-items:center;gap:6px;flex:1 1 auto;min-width:0}
|
||||
.ausboss-save-half .ausboss-card-seg{flex:1 1 0;min-width:0}
|
||||
.ausboss-save-switches{gap:12px}
|
||||
.ausboss-save-browse{position:fixed;z-index:10000;width:260px;max-height:320px;overflow:auto;padding:4px;border:1px solid #3a4047;border-radius:7px;background:#1c1f23;box-shadow:0 8px 28px rgba(0,0,0,.5);font:12px/1.3 ${FONT};color:#d8ecea}
|
||||
.ausboss-save-browse-head{padding:5px 8px 6px;border-bottom:1px solid #2c3238;color:#78908e;font:10px/1 ${FONT};letter-spacing:.08em;text-transform:uppercase;white-space:nowrap;overflow:hidden;text-overflow:ellipsis}
|
||||
.ausboss-save-browse button{display:block;width:100%;padding:6px 8px;border:none;border-radius:5px;background:transparent;color:inherit;font:inherit;text-align:left;cursor:pointer}
|
||||
@@ -282,9 +285,29 @@ function buildCard(node) {
|
||||
button.addEventListener("click", () => { if (setValue(node, "save_metadata", on)) notifyAusbossChange(); state.refresh(); });
|
||||
embedSeg.append(button); embedButtons.set(on, button);
|
||||
}
|
||||
embedRow.append(el("span", "ausboss-save-embed-label", PANEL_LABELS.embed), embedSeg);
|
||||
const embedHalf = el("div", "ausboss-save-half");
|
||||
embedHalf.append(el("span", "ausboss-save-embed-label", PANEL_LABELS.embed), embedSeg);
|
||||
embedRow.classList.add("ausboss-save-switches");
|
||||
embedRow.append(embedHalf);
|
||||
root.append(embedRow);
|
||||
|
||||
// Auto save: off shows the pictures without saving them; the Save button
|
||||
// over the picture (js/input_preview) then keeps the ones that are wanted.
|
||||
// It shares the Embed workflow row, so the card is no taller than before.
|
||||
const autoHalf = el("div", "ausboss-save-half");
|
||||
const autoSeg = el("div", "ausboss-card-seg ausboss-card-bool");
|
||||
autoSeg.setAttribute("role", "radiogroup");
|
||||
const autoButtons = new Map();
|
||||
const autoTitle = "On: every run saves its pictures. Off: the pictures only show here, and the Save button over them keeps the ones you want.";
|
||||
for (const [on, text] of [[false, PANEL_LABELS.off], [true, PANEL_LABELS.on]]) {
|
||||
const button = el("button", "", text);
|
||||
button.type = "button"; button.title = autoTitle;
|
||||
button.addEventListener("click", () => { if (setValue(node, "autosave", on)) { state.savedPath = null; notifyAusbossChange(); } state.refresh(); });
|
||||
autoSeg.append(button); autoButtons.set(on, button);
|
||||
}
|
||||
autoHalf.append(el("span", "ausboss-save-embed-label", "Auto save"), autoSeg);
|
||||
embedRow.append(autoHalf);
|
||||
|
||||
root.addEventListener("pointerdown", (event) => {
|
||||
if (event.target.closest("button,input,select")) event.stopPropagation();
|
||||
});
|
||||
@@ -330,6 +353,11 @@ function buildCard(node) {
|
||||
button.classList.toggle("on", on === embed);
|
||||
button.setAttribute("aria-checked", String(on === embed));
|
||||
}
|
||||
const auto = values.autosave !== false;
|
||||
for (const [on, button] of autoButtons) {
|
||||
button.classList.toggle("on", on === auto);
|
||||
button.setAttribute("aria-checked", String(on === auto));
|
||||
}
|
||||
};
|
||||
|
||||
const domWidget = node.addDOMWidget(WIDGET_NAME, WIDGET_NAME, root, {
|
||||
|
||||
@@ -0,0 +1,24 @@
|
||||
// A click on the dark area around a modal <dialog> closes it, the way people
|
||||
// expect any popup to work. The browser reports a click on a dialog's
|
||||
// backdrop as a click on the dialog element itself, so a click counts as
|
||||
// outside when it lands on the dialog element beyond its own box. Both the
|
||||
// press and the release must be outside: a text selection or a drag that
|
||||
// starts on the card and ends on the backdrop keeps the popup open.
|
||||
|
||||
export function pointOutside(rect, x, y) {
|
||||
return x < rect.left || x > rect.right || y < rect.top || y > rect.bottom;
|
||||
}
|
||||
|
||||
export function closeOnBackdropClick(dialog, close = () => dialog.close()) {
|
||||
const outside = (event) => event.target === dialog
|
||||
&& pointOutside(dialog.getBoundingClientRect(), event.clientX, event.clientY);
|
||||
let pressedOutside = false;
|
||||
dialog.addEventListener("pointerdown", (event) => {
|
||||
pressedOutside = outside(event);
|
||||
});
|
||||
dialog.addEventListener("click", (event) => {
|
||||
const started = pressedOutside;
|
||||
pressedOutside = false;
|
||||
if (started && outside(event)) close();
|
||||
});
|
||||
}
|
||||
+1
-1
@@ -9,7 +9,7 @@ import { app } from "../../../scripts/app.js";
|
||||
// Must match `version` in pyproject.toml — scripts/release_preflight.py
|
||||
// enforces it. The backend serves its copy at /ausboss/pack_version; a
|
||||
// mismatch means the browser cached JavaScript from an older install.
|
||||
export const AUSBOSS_JS_VERSION = "2.23.0";
|
||||
export const AUSBOSS_JS_VERSION = "2.24.0";
|
||||
|
||||
export const BRAND = "#00b4aa"; // AusBoss teal — keep node accents consistent.
|
||||
export const BRAND_DARK = "#007f78";
|
||||
|
||||
@@ -52,9 +52,10 @@ export function makePictureStepper({ noun = "picture", height = 24, onStep, sign
|
||||
root.append(arrow("‹", -1, `Previous ${noun}`), label, arrow("›", 1, `Next ${noun}`));
|
||||
return {
|
||||
root,
|
||||
set(index, count) {
|
||||
// `text` replaces the "2 / 5" for a stop that is not one picture.
|
||||
set(index, count, text) {
|
||||
root.hidden = !(count > 1);
|
||||
label.textContent = stepperLabel(index, count);
|
||||
label.textContent = text ?? stepperLabel(index, count);
|
||||
},
|
||||
};
|
||||
}
|
||||
|
||||
@@ -19,3 +19,53 @@ export function stepIndex(index, count, step) {
|
||||
export function stepperLabel(index, count) {
|
||||
return count > 1 ? `${keepIndex(index, count) + 1} / ${count}` : "";
|
||||
}
|
||||
|
||||
// The preview panel has one more stop, before picture 1: every picture at
|
||||
// once. ALL stands for it where an index would be.
|
||||
export const ALL = -1;
|
||||
|
||||
// One step back or on through "all, 1, 2 ... count", round at both ends.
|
||||
export function stepWithAll(index, count, step) {
|
||||
if (!(count > 0)) return ALL;
|
||||
const stops = count + 1;
|
||||
const at = index === ALL ? 0 : keepIndex(index, count) + 1;
|
||||
const next = (((at + Math.sign(Number(step) || 0)) % stops) + stops) % stops;
|
||||
return next === 0 ? ALL : next - 1;
|
||||
}
|
||||
|
||||
export function allLabel(count) {
|
||||
return count > 1 ? `all ${count}` : "";
|
||||
}
|
||||
|
||||
// How to lay `count` pictures out in a box so each one draws as big as it
|
||||
// can: the columns and rows that leave a picture of this shape (`aspect`,
|
||||
// width over height) the most room.
|
||||
export function gridShape(count, width, height, aspect = 1) {
|
||||
const total = Math.max(1, Math.trunc(Number(count) || 0));
|
||||
const shape = Number.isFinite(aspect) && aspect > 0 ? aspect : 1;
|
||||
if (!(width > 0) || !(height > 0)) {
|
||||
const columns = Math.ceil(Math.sqrt(total));
|
||||
return { columns, rows: Math.ceil(total / columns) };
|
||||
}
|
||||
let best = { columns: 1, rows: total, size: 0 };
|
||||
for (let columns = 1; columns <= total; columns += 1) {
|
||||
const rows = Math.ceil(total / columns);
|
||||
const size = Math.min(width / columns, (height / rows) * shape);
|
||||
if (size > best.size) best = { columns, rows, size };
|
||||
}
|
||||
return { columns: best.columns, rows: best.rows };
|
||||
}
|
||||
|
||||
// The held pictures in a Save Image result, in the order they show: each
|
||||
// one's run token and its place in that run. A list run saves once for each
|
||||
// item, so one result can hold several tokens. Empty when the result holds
|
||||
// none, or when a picture in it is not one of the held ones.
|
||||
export function heldPlaces(record) {
|
||||
const tokens = new Set(Array.isArray(record?.ausboss_held) ? record.ausboss_held : []);
|
||||
if (!tokens.size) return [];
|
||||
const held = (Array.isArray(record.images) ? record.images : []).map((image) => {
|
||||
const match = /^([0-9a-f]{16})_(\d{3})\./.exec(String(image?.filename ?? ""));
|
||||
return match && tokens.has(match[1]) ? { token: match[1], index: Number(match[2]) } : null;
|
||||
});
|
||||
return held.length && held.every(Boolean) ? held : [];
|
||||
}
|
||||
|
||||
@@ -160,3 +160,51 @@ export function collapseRepeatedErrors(nodeErrors) {
|
||||
}
|
||||
return nodeErrors;
|
||||
}
|
||||
|
||||
// ------------------------------------------------------------ known run errors
|
||||
//
|
||||
// Some errors that stop a run come from outside the workflow and say nothing
|
||||
// a person can act on ("KeyError: 'latent'"). Each entry here names one such
|
||||
// error by what is in ComfyUI's own report, and says in plain words what to
|
||||
// do. An entry is added only after the cause was found for real.
|
||||
const KNOWN_RUN_ERRORS = Object.freeze([
|
||||
{
|
||||
// Muse MiniMax Director replaces part of H3 for every workflow, and that
|
||||
// part reads a picture from every guide. ComfyUI's own H3 nodes also make
|
||||
// guides that hold sound alone, which is how a clip keeps its sound.
|
||||
matches: ({ type, message, traceback }) =>
|
||||
type === "KeyError" && message.includes("'latent'") && traceback.includes("muse_minimax_hybrid_patch"),
|
||||
summary: "Another node pack stopped this run",
|
||||
detail:
|
||||
"This error comes from the Muse MiniMax Director pack, not from this workflow. Turn that pack off in "
|
||||
+ "Manager, restart ComfyUI and run again.",
|
||||
},
|
||||
{
|
||||
// The w6a8 H3 files load from ComfyUI 0.38.0 on; older versions stop in
|
||||
// the model loader on the name of the format.
|
||||
matches: ({ message }) => message.includes("w6a8_int8"),
|
||||
summary: "This model needs a newer ComfyUI",
|
||||
detail:
|
||||
"Your ComfyUI is too old to read this model file. Update ComfyUI to 0.38 or newer and run again.",
|
||||
},
|
||||
]);
|
||||
|
||||
/**
|
||||
* Plain advice for a run error the pack knows the cause of, from the
|
||||
* `execution_error` event's detail. `null` for every other error.
|
||||
*/
|
||||
export function explainRunError(detail) {
|
||||
const report = {
|
||||
type: String(detail?.exception_type ?? "").split(".").pop(),
|
||||
message: String(detail?.exception_message ?? ""),
|
||||
traceback: [].concat(detail?.traceback ?? []).join("\n"),
|
||||
};
|
||||
for (const known of KNOWN_RUN_ERRORS) {
|
||||
try {
|
||||
if (known.matches(report)) return { summary: known.summary, detail: known.detail };
|
||||
} catch (_error) {
|
||||
// A report in a shape this entry did not expect is simply not a match.
|
||||
}
|
||||
}
|
||||
return null;
|
||||
}
|
||||
|
||||
@@ -74,6 +74,25 @@ export function loaderSignature(widgets) {
|
||||
return JSON.stringify(settings);
|
||||
}
|
||||
|
||||
// What the picture says when a run stopped at the node. ComfyUI's own text
|
||||
// for a full graphics card is ten lines of numbers; here it is one plain
|
||||
// line. Anything else is said as it came.
|
||||
export const OUT_OF_MEMORY = "The graphics card ran out of memory, so the node could not look. Free some memory on it, then click again.";
|
||||
export function stopMessage(message) {
|
||||
const text = String(message ?? "");
|
||||
return /out of memory|OutOfMemory/i.test(text) ? OUT_OF_MEMORY : text;
|
||||
}
|
||||
|
||||
// A name as it is typed under the picture: "-face" (or the minus sign a
|
||||
// phone types) takes the name away, "+face" adds it; without a sign it does
|
||||
// what `takes` says. Returns { word, add }, the word cleaned.
|
||||
export function typedName(text, takes = false) {
|
||||
const typed = String(text ?? "").trim();
|
||||
const sign = /^[-\u2212\u2013+]/.exec(typed)?.[0] ?? "";
|
||||
const word = cleanWord(sign ? typed.slice(1) : typed);
|
||||
return { word, add: sign ? sign === "+" : !takes };
|
||||
}
|
||||
|
||||
// A word as the node sends it to SAM 3: commas, colons and brackets are the
|
||||
// text reader's own syntax and go; spaces are collapsed.
|
||||
export function cleanWord(word) {
|
||||
@@ -1044,19 +1063,46 @@ export function strokePieces(planes, list, sizes, pixels, stroke, ranks, wholes
|
||||
// picked area, and as untouched while almost none of it is.
|
||||
const NEARLY = 0.03;
|
||||
|
||||
// Which pieces hold a whole pick: nearly all of one picked area (`areas`,
|
||||
// on/off masks, one per chip) lies inside the piece. A person holds the
|
||||
// sweater picked on them; hair that a picked jacket only laps over does not.
|
||||
export const HOLDS = 0.9;
|
||||
export function piecesHolding(planes, count, areas) {
|
||||
const holds = new Uint8Array(count);
|
||||
for (const area of areas ?? []) {
|
||||
if (!area) continue;
|
||||
let size = 0;
|
||||
for (let index = 0; index < area.length; index += 1) if (area[index]) size += 1;
|
||||
if (!size) continue;
|
||||
const inside = piecesInside(planes, count, area);
|
||||
for (let piece = 0; piece < count; piece += 1) if (inside[piece] >= HOLDS * size) holds[piece] = 1;
|
||||
}
|
||||
return holds;
|
||||
}
|
||||
|
||||
// Which piece a pointer at pixel `index` lights up: the first piece there,
|
||||
// in the order they are offered (`ranks`; without it, smallest first), that
|
||||
// would still change something. Adding, that is a piece not yet (nearly)
|
||||
// all picked, so clicking the same spot again goes on from the glasses to
|
||||
// the face to the person. Taking away, it is a piece that has some of the
|
||||
// picked area in it. -1 when there is none.
|
||||
export function pieceAt(planes, index, { sizes, inside, adding = true, ranks = null } = {}) {
|
||||
//
|
||||
// Taking away from a spot that is not picked itself (`here` false) never
|
||||
// means a thing that holds a whole pick somewhere else (`holds`, from
|
||||
// piecesHolding): with a sweater picked, pointing at the face does not
|
||||
// offer the person, which would take the sweater along.
|
||||
export function pieceAt(planes, index, { sizes, inside, adding = true, ranks = null, here = true, holds = null } = {}) {
|
||||
for (const piece of candidatesAt(planes, index, ranks)) {
|
||||
const size = sizes?.[piece] ?? 0;
|
||||
if (!size) continue;
|
||||
const picked = inside?.[piece] ?? 0;
|
||||
const slack = Math.max(4, size * NEARLY);
|
||||
if (adding ? size - picked > slack : picked > slack) return piece;
|
||||
if (adding) {
|
||||
if (size - picked > slack) return piece;
|
||||
continue;
|
||||
}
|
||||
if (picked <= slack || (!here && holds?.[piece])) continue;
|
||||
return piece;
|
||||
}
|
||||
return -1;
|
||||
}
|
||||
|
||||
+106
-19
@@ -77,6 +77,7 @@ import {
|
||||
pieceRanks,
|
||||
pieceWholes,
|
||||
pieceSizes,
|
||||
piecesHolding,
|
||||
piecesInside,
|
||||
prefixStamps,
|
||||
previewPixels,
|
||||
@@ -85,9 +86,11 @@ import {
|
||||
sharedPixels,
|
||||
shrinkMask,
|
||||
stampOf,
|
||||
stopMessage,
|
||||
strokePieces,
|
||||
thinStroke,
|
||||
touchFits,
|
||||
typedName,
|
||||
} from "./touch_math.mjs";
|
||||
|
||||
const CSS_ID = "ausboss-touch-css";
|
||||
@@ -190,9 +193,17 @@ function ensureCss() {
|
||||
.ausboss-touch-chip>button:hover{opacity:1;background:rgba(255,255,255,.14);}
|
||||
.ausboss-touch-chip>button.more{padding:0 6px;background:rgba(0,0,0,.35);opacity:.95;}
|
||||
.ausboss-touch-chip.open>button.more{background:rgba(0,0,0,.2);}
|
||||
.ausboss-touch-name{box-sizing:border-box;flex:none;width:104px;height:20px;padding:0 9px;border:1px dashed rgba(0,180,170,.5);border-radius:10px;background:transparent;color:#dffcf9;font:600 10px/18px -apple-system,BlinkMacSystemFont,"Segoe UI",sans-serif;outline:none;}
|
||||
.ausboss-touch-namebox{box-sizing:border-box;flex:none;display:inline-flex;align-items:stretch;height:20px;border:1px dashed rgba(0,180,170,.5);border-radius:10px;overflow:hidden;}
|
||||
.ausboss-touch-namebox:focus-within{border-style:solid;border-color:${BRAND};background:rgba(0,0,0,.4);}
|
||||
.ausboss-touch-namebox.away{border-color:rgba(${LIT_AWAY.join(",")},.6);}
|
||||
.ausboss-touch-namebox.away:focus-within{border-color:rgb(${LIT_AWAY.join(",")});}
|
||||
.ausboss-touch-sign{box-sizing:border-box;flex:none;width:22px;padding:0;border:0;border-right:1px solid rgba(0,180,170,.35);background:rgba(0,180,170,.16);color:#dffcf9;font:700 13px/18px -apple-system,BlinkMacSystemFont,"Segoe UI",sans-serif;cursor:pointer;}
|
||||
.ausboss-touch-sign:hover{background:rgba(0,180,170,.34);color:#fff;}
|
||||
.ausboss-touch-sign:disabled{opacity:.4;cursor:default;}
|
||||
.ausboss-touch-namebox.away .ausboss-touch-sign{border-right-color:rgba(${LIT_AWAY.join(",")},.4);background:rgba(${LIT_AWAY.join(",")},.24);color:#ffe9e2;}
|
||||
.ausboss-touch-namebox.away .ausboss-touch-sign:hover{background:rgba(${LIT_AWAY.join(",")},.42);}
|
||||
.ausboss-touch-name{box-sizing:border-box;flex:none;width:118px;height:18px;padding:0 8px 0 6px;border:0;background:transparent;color:#dffcf9;font:600 10px/18px -apple-system,BlinkMacSystemFont,"Segoe UI",sans-serif;outline:none;}
|
||||
.ausboss-touch-name::placeholder{color:#78908e;font-weight:500;}
|
||||
.ausboss-touch-name:focus{border-style:solid;border-color:${BRAND};background:rgba(0,0,0,.4);}
|
||||
.ausboss-touch-gap{flex:1 1 auto;min-width:0;}
|
||||
.ausboss-touch-stage{position:relative;flex:1 1 auto;min-height:0;width:100%;overflow:hidden;border:1px solid rgba(0,180,170,.27);border-radius:6px;background:rgba(0,0,0,.28);}
|
||||
[data-node-id] .ausboss-touch-stage{box-sizing:border-box;min-height:${STAGE_MIN}px;}
|
||||
@@ -224,9 +235,12 @@ function ensureCss() {
|
||||
.ausboss-touch.fining .ausboss-touch-fining{display:flex;}
|
||||
.ausboss-touch-fining::before{content:"";flex:none;width:7px;height:7px;border-radius:50%;background:${BRAND};animation:ausboss-touch-pulse 1s ease-in-out infinite;}
|
||||
@keyframes ausboss-touch-pulse{0%,100%{opacity:.25;}50%{opacity:1;}}
|
||||
.ausboss-touch-btn.lookpill{display:inline-flex;align-items:center;gap:5px;height:${BAR_HEIGHT}px;padding:0 9px 0 8px;border:1px solid rgba(0,180,170,.6);border-radius:10px;letter-spacing:0;text-transform:none;font:600 10px/18px -apple-system,BlinkMacSystemFont,"Segoe UI",sans-serif;color:#dffcf9;background:linear-gradient(90deg,rgba(0,180,170,.38) var(--done,0%),rgba(0,180,170,.08) var(--done,0%));}
|
||||
.ausboss-touch-btn.lookpill{display:inline-flex;align-items:center;gap:5px;height:${BAR_HEIGHT}px;padding:0 9px 0 0;overflow:hidden;border:1px solid rgba(0,180,170,.8);border-radius:10px;letter-spacing:0;text-transform:none;font:600 10px/18px -apple-system,BlinkMacSystemFont,"Segoe UI",sans-serif;color:#dffcf9;background:linear-gradient(90deg,rgba(0,180,170,.38) var(--done,0%),rgba(0,180,170,.08) var(--done,0%));}
|
||||
.ausboss-touch-btn.lookpill:hover{border-color:${BRAND};color:#fff;background:linear-gradient(90deg,rgba(0,180,170,.5) var(--done,0%),rgba(0,180,170,.16) var(--done,0%));}
|
||||
.ausboss-touch-btn.lookpill .icon{flex:none;font-size:8px;letter-spacing:-1px;opacity:.9;}
|
||||
.ausboss-touch-btn.lookpill .key{flex:none;display:inline-flex;align-items:center;gap:4px;align-self:stretch;padding:0 7px 0 8px;background:${BRAND};color:#04201e;font-weight:700;}
|
||||
.ausboss-touch-btn.lookpill:hover .key{background:#35d6cb;}
|
||||
.ausboss-touch-btn.lookpill.paused .key{background:#9fb4b2;}
|
||||
.ausboss-touch-btn.lookpill .icon{flex:none;font-size:8px;letter-spacing:-1px;}
|
||||
.ausboss-touch-btn.lookpill.paused{border-color:rgba(120,144,142,.6);color:#c8dddd;background:linear-gradient(90deg,rgba(120,144,142,.3) var(--done,0%),rgba(120,144,142,.06) var(--done,0%));}
|
||||
.ausboss-touch-btn.lookpill.paused .icon{letter-spacing:0;}
|
||||
.ausboss-touch-btn.lookpill .count{opacity:.8;}
|
||||
@@ -527,6 +541,7 @@ function resolve(face) {
|
||||
face.raw = raw;
|
||||
face.exact = exact;
|
||||
face.chipAreas = null;
|
||||
face.holds = null;
|
||||
face.inside = pic.pieces ? piecesInside(pic.pieces.planes, pic.pieces.list.length, raw) : null;
|
||||
face.hoverMask = null;
|
||||
}
|
||||
@@ -697,6 +712,11 @@ function requestRun(face, ask) {
|
||||
return false;
|
||||
}
|
||||
const run = { ask, stamp: face.stamp, epoch: face.sourceEpoch, queued: false };
|
||||
// Runs of this node that are at a picture it no longer has. The one that
|
||||
// is running ends by itself once the next job waits; the ones that still
|
||||
// wait in the queue would each run for nothing, one after the other, so
|
||||
// they are taken out before this one is sent.
|
||||
const stale = face.flying.filter((other) => other.queued && other.epoch !== run.epoch);
|
||||
face.flying.push(run);
|
||||
face.inFlight = true;
|
||||
face.askedAt = performance.now();
|
||||
@@ -707,7 +727,15 @@ function requestRun(face, ask) {
|
||||
face.flying = face.flying.filter((other) => other !== run);
|
||||
if (!face.flying.length) settle(face);
|
||||
};
|
||||
Promise.resolve(face.host.run(ask === "scope" ? scopeAsk(face) : ask)).then((queued) => {
|
||||
const cleared = stale.length && face.host.dropWaiting
|
||||
? Promise.resolve(face.host.dropWaiting()).then((gone) => {
|
||||
// What was taken out never answers: the newest of the stale ones,
|
||||
// because a run of this node goes to the front of the queue.
|
||||
const out = new Set(stale.slice(Math.max(0, stale.length - (Number(gone) || 0))));
|
||||
if (out.size) face.flying = face.flying.filter((other) => !out.has(other));
|
||||
}, () => {})
|
||||
: Promise.resolve();
|
||||
cleared.then(() => face.host.run(ask === "scope" ? scopeAsk(face) : ask)).then((queued) => {
|
||||
if (queued === false) {
|
||||
failed();
|
||||
return;
|
||||
@@ -1258,11 +1286,23 @@ function takeAll(face, index) {
|
||||
refreshHover(face);
|
||||
}
|
||||
|
||||
// Whether a typed name is taken away: what the sign in front of the field
|
||||
// was set to, or else what the tool does right now (Erase, or Ctrl held).
|
||||
function nameTakes(face) {
|
||||
return face.nameAway ?? face.away;
|
||||
}
|
||||
|
||||
// A typed name, as every thing of that name. "-face" takes it away and
|
||||
// "+face" adds it, whatever the sign says.
|
||||
function addName(face, text) {
|
||||
const word = cleanWord(text);
|
||||
if (!word || !face.pic || face.held) return false;
|
||||
const typed = typedName(text, nameTakes(face));
|
||||
if (!typed.word || !face.pic || face.held) return false;
|
||||
flushPending(face);
|
||||
addItem(face, { t: "all", op: face.away ? 0 : 1, w: word });
|
||||
if (!typed.add && !face.items.some((item) => item.op === 1)) {
|
||||
setNote(face, "Nothing is picked yet. Pick something first, then take a name away from it.", { life: 3200 });
|
||||
return false;
|
||||
}
|
||||
addItem(face, { t: "all", op: typed.add ? 1 : 0, w: typed.word });
|
||||
return true;
|
||||
}
|
||||
|
||||
@@ -1413,9 +1453,19 @@ function syncTools(face) {
|
||||
face.lookButton.style.display = offer ? "" : "none";
|
||||
}
|
||||
if (face.nameField) {
|
||||
face.nameField.disabled = !live || face.held > 0;
|
||||
// A typed name is added, or taken away while the tool takes away.
|
||||
face.nameField.placeholder = face.away ? "\u2212 type a name" : "+ type a name";
|
||||
const off = !live || face.held > 0;
|
||||
const takes = nameTakes(face);
|
||||
face.nameField.disabled = off;
|
||||
face.nameSign.disabled = off;
|
||||
face.nameBox.classList.toggle("away", takes);
|
||||
face.nameSign.textContent = takes ? "\u2212" : "+";
|
||||
face.nameSign.title = takes
|
||||
? "A name you type is taken away from the mask. Click to add by name instead."
|
||||
: "A name you type is added to the mask. Click to take away by name instead.";
|
||||
face.nameField.placeholder = takes ? "take away by name" : "add by name";
|
||||
face.nameField.title = takes
|
||||
? "Type what to take out of the mask, like \"face\" or \"hair\", and press Enter. Every thing of that name is taken away."
|
||||
: "Type what to pick, like \"screen\" or \"earring\", and press Enter. Every thing of that name is taken. Start with a minus, like \"-face\", to take it away instead.";
|
||||
}
|
||||
if (face.newButton) face.newButton.disabled = !live || face.held > 0 || Boolean(face.fresh);
|
||||
// The switch shows what a press does right now: with Ctrl held, Paint
|
||||
@@ -1949,12 +1999,31 @@ function pointOf(face, event) {
|
||||
return [clamp(x, 0, face.pic.width - 1), clamp(y, 0, face.pic.height - 1)];
|
||||
}
|
||||
|
||||
// The pieces that hold a whole pick, worked out once per change to the
|
||||
// picks and only when something is being taken away.
|
||||
function holdsOf(face) {
|
||||
const pic = face.pic;
|
||||
if (!pic?.pieces || !face.raw) return null;
|
||||
if (!face.holds || face.holds.length !== pic.pieces.list.length) {
|
||||
const areas = chipGroups(face.items).filter((chip) => !chip.away).map((chip) => chipArea(face, chip));
|
||||
face.holds = piecesHolding(pic.pieces.planes, pic.pieces.list.length, areas);
|
||||
}
|
||||
return face.holds;
|
||||
}
|
||||
|
||||
function pieceUnder(face, point, adding) {
|
||||
const pic = face.pic;
|
||||
if (!pic?.pieces || !point) return -1;
|
||||
const x = clamp(Math.floor(point[0] * pic.scale), 0, pic.pw - 1);
|
||||
const y = clamp(Math.floor(point[1] * pic.scale), 0, pic.ph - 1);
|
||||
return pieceAt(pic.pieces.planes, y * pic.pw + x, { sizes: pic.pieces.sizes, inside: face.inside, adding, ranks: pic.pieces.ranks });
|
||||
const at = y * pic.pw + x;
|
||||
if (adding) return pieceAt(pic.pieces.planes, at, { sizes: pic.pieces.sizes, inside: face.inside, ranks: pic.pieces.ranks });
|
||||
// Taking away from a spot that is not picked: never a thing that holds a
|
||||
// whole pick somewhere else (the person, for the sweater picked on them).
|
||||
return pieceAt(pic.pieces.planes, at, {
|
||||
sizes: pic.pieces.sizes, inside: face.inside, adding, ranks: pic.pieces.ranks,
|
||||
here: face.raw?.[at] >= 0.5, holds: holdsOf(face),
|
||||
});
|
||||
}
|
||||
|
||||
// The pixel of the small copy under a point of the picture.
|
||||
@@ -2324,7 +2393,8 @@ function syncLooking(face) {
|
||||
pill.classList.toggle("paused", paused);
|
||||
pill.style.setProperty("--done", `${total ? Math.round((done / total) * 100) : 0}%`);
|
||||
face.lookIcon.textContent = paused ? "\u25b6" : "\u275a\u275a";
|
||||
face.lookWord.textContent = paused ? "Paused" : "Looking";
|
||||
face.lookVerb.textContent = paused ? "Look on" : "Pause";
|
||||
face.lookWord.textContent = paused ? "paused" : "looking";
|
||||
face.lookCount.textContent = total ? `${done}/${total}` : "\u2026";
|
||||
const far = total ? ` ${done} of ${total}` : "";
|
||||
pill.title = paused
|
||||
@@ -2717,6 +2787,8 @@ function setSoft(face, soft) {
|
||||
|
||||
// Pick takes whole things; Paint and Erase are the brush.
|
||||
function setTool(face, tool) {
|
||||
// The sign in front of the name field follows the tool again.
|
||||
face.nameAway = null;
|
||||
if (!TOOLS.includes(tool)) return;
|
||||
face.tool = tool;
|
||||
face.snap = tool === "pick";
|
||||
@@ -2773,16 +2845,21 @@ export function mountTouchStage(node, host) {
|
||||
const bar = el("div", "ausboss-touch-bar");
|
||||
const chips = el("div", "ausboss-touch-chips");
|
||||
const chipList = el("span", "ausboss-touch-chiplist");
|
||||
// A name typed here picks every thing of that name, or takes it away:
|
||||
// the sign in front says which, and a press on it changes it.
|
||||
const nameBox = el("span", "ausboss-touch-namebox");
|
||||
const nameSign = el("button", "ausboss-touch-sign", "+");
|
||||
nameSign.type = "button";
|
||||
const nameField = el("input", "ausboss-touch-name");
|
||||
nameField.type = "text";
|
||||
nameField.placeholder = "+ type a name";
|
||||
nameField.placeholder = "add by name";
|
||||
nameField.maxLength = 60;
|
||||
nameField.spellcheck = false;
|
||||
nameField.title = "Type what to pick, like \"screen\" or \"earring\", and press Enter. Every thing of that name is taken.";
|
||||
nameBox.append(nameSign, nameField);
|
||||
const newButton = el("button", "ausboss-touch-new", "+ new");
|
||||
newButton.type = "button";
|
||||
newButton.title = "Start a new area: what you paint next gets a chip of its own, also where it touches another area.";
|
||||
chips.append(chipList, newButton, nameField);
|
||||
chips.append(chipList, newButton, nameBox);
|
||||
root.append(tools, stage, bar, chips);
|
||||
|
||||
const frame = el("div", "ausboss-touch-frame");
|
||||
@@ -2827,7 +2904,7 @@ export function mountTouchStage(node, host) {
|
||||
|
||||
const face = node.__ausbossTouch = {
|
||||
node, host, root, stage, frame, img, tint, lit, ink, path, dots, ringBack, ring, coreRing, dab, brushLabel, tag, hint, noteEl, abort,
|
||||
chips, chipList, newButton, nameField, chipsShown: null,
|
||||
chips, chipList, newButton, nameField, nameBox, nameSign, nameAway: null, chipsShown: null,
|
||||
flying: [],
|
||||
pendingTimer: 0,
|
||||
alive: true,
|
||||
@@ -3031,10 +3108,15 @@ export function mountTouchStage(node, host) {
|
||||
if (face.pic?.paused && !face.inFlight) resumeLook(face);
|
||||
else pauseLook(face);
|
||||
}, signal);
|
||||
// The press it takes comes first, as a button of its own look: "Pause"
|
||||
// while the node looks, "Look on" once it is paused.
|
||||
face.lookIcon = el("span", "icon");
|
||||
face.lookVerb = el("span", "verb");
|
||||
face.lookWord = el("span", "word");
|
||||
face.lookCount = el("span", "count");
|
||||
face.lookPill.append(face.lookIcon, face.lookWord, face.lookCount);
|
||||
const lookKey = el("span", "key");
|
||||
lookKey.append(face.lookIcon, face.lookVerb);
|
||||
face.lookPill.append(lookKey, face.lookWord, face.lookCount);
|
||||
face.lookPill.style.display = "none";
|
||||
// First in the line under the picture: what the node is doing with the
|
||||
// picture, then what is picked; the buttons that pick come after.
|
||||
@@ -3044,6 +3126,11 @@ export function mountTouchStage(node, host) {
|
||||
for (const type of ["keydown", "keyup", "keypress", "pointerdown", "mousedown", "dblclick"]) {
|
||||
chips.addEventListener(type, (event) => event.stopPropagation(), { signal });
|
||||
}
|
||||
nameSign.addEventListener("click", () => {
|
||||
face.nameAway = !nameTakes(face);
|
||||
syncTools(face);
|
||||
nameField.focus();
|
||||
}, { signal });
|
||||
nameField.addEventListener("keydown", (event) => {
|
||||
if (event.key === "Enter") {
|
||||
event.preventDefault();
|
||||
@@ -3096,7 +3183,7 @@ export function mountTouchStage(node, host) {
|
||||
// The run stopped at this node with a message for the person.
|
||||
face.stopped = (message) => {
|
||||
settle(face);
|
||||
setNote(face, String(message ?? ""));
|
||||
setNote(face, stopMessage(message));
|
||||
// It stays up until something else is said or the next result comes in.
|
||||
face.stopNote = true;
|
||||
};
|
||||
|
||||
@@ -121,6 +121,15 @@ const CARDS = {
|
||||
{ widget: "interpolation", label: "Filter", kind: "select" },
|
||||
],
|
||||
},
|
||||
AUSBOSS_NODES_EditLock: {
|
||||
minWidth: 300,
|
||||
rows: [
|
||||
{ widget: "lock", label: "Lock", kind: "segment",
|
||||
titles: { auto: "Lock the untouched parts; step aside when the whole picture changes", always: "Lock even when the whole picture changes", off: "Do nothing" } },
|
||||
{ widget: "tolerance", label: "Tolerance", decimals: 2, suffix: "×" },
|
||||
],
|
||||
},
|
||||
AUSBOSS_NODES_EditLockRestore: { minWidth: 300, rows: [{ widget: "feather", label: "Feather", suffix: "px" }] },
|
||||
AUSBOSS_NODES_ColorMatch: {
|
||||
minWidth: 300,
|
||||
rows: [
|
||||
@@ -281,6 +290,14 @@ const CARDS = {
|
||||
{ widget: "batch_size", label: "Batch" },
|
||||
],
|
||||
},
|
||||
AUSBOSS_NODES_FlashVSRUpscale: {
|
||||
minWidth: 340,
|
||||
rows: [
|
||||
{ widget: "model", label: "Model", kind: "select" },
|
||||
{ widget: "scale", label: "Scale", suffix: "×", decimals: 2 },
|
||||
{ widget: "color_match", label: "Keep colors" },
|
||||
],
|
||||
},
|
||||
AUSBOSS_NODES_AlignImage: {
|
||||
minWidth: 300,
|
||||
rows: [
|
||||
@@ -331,6 +348,23 @@ const CARDS = {
|
||||
rows: [{ widget: "text", kind: "textarea", placeholder: "Prompt, caption, or shared text", height: 100, grow: true,
|
||||
looks: { big: { height: 62, placeholder: "type here" } } }],
|
||||
},
|
||||
// One line for each run. The three settings under the box say where a run
|
||||
// ends, whether a leading "1." comes off, and how many lines run at most.
|
||||
AUSBOSS_NODES_TextLines: {
|
||||
minWidth: 320,
|
||||
rows: [
|
||||
{ widget: "text", kind: "textarea", placeholder: "One line for each run", height: 100, grow: true },
|
||||
{ widget: "split", label: "Split at", kind: "segment" },
|
||||
{ widget: "remove_numbering", label: "Numbering", onText: "remove", offText: "keep" },
|
||||
{ widget: "limit", label: "Limit" },
|
||||
],
|
||||
},
|
||||
// The seed arrives on a socket (forceInput), so the card is the prompt
|
||||
// alone, with its {a|b} choice sets.
|
||||
AUSBOSS_NODES_PromptVariations: {
|
||||
minWidth: 300,
|
||||
rows: [{ widget: "text", kind: "textarea", placeholder: "A woman in a {red|blue|green} {sweater|jacket}. One choice is picked from each set.", height: 100, grow: true }],
|
||||
},
|
||||
// All the words go in the box, and the strip under it shows when each
|
||||
// moment plays. The pictures, the video, the voice and the song arrive on
|
||||
// sockets, and each picture's job sits under the words.
|
||||
|
||||
@@ -19,6 +19,7 @@ import { BRAND, chainCallback, keepDomWidgetWidthAuto, notifyAusbossChange } fro
|
||||
import { collectGraphNodes } from "../shared/appearance.mjs";
|
||||
import { copyToClipboard } from "../shared/clipboard.mjs";
|
||||
import { confirmDiscard } from "../shared/discard_prompt.mjs";
|
||||
import { closeOnBackdropClick } from "../shared/dialog_backdrop.mjs";
|
||||
import { WIDGET_FRAME, fillNodeHeight, holdNodeMinHeight } from "../shared/panel_layout.mjs";
|
||||
import { hideInputsInDef, hideWidget } from "../shared/widget_visibility.mjs";
|
||||
import { suppressCoreImagePreview } from "../shared/core_preview.mjs";
|
||||
@@ -551,6 +552,7 @@ function render(state) {
|
||||
if (note.author) header.append(el("p", "", `by ${note.author}`));
|
||||
dialog.append(header);
|
||||
renderSections({ note, rows: [] }, dialog);
|
||||
closeOnBackdropClick(dialog);
|
||||
document.body.append(dialog);
|
||||
state.reader = dialog;
|
||||
dialog.addEventListener("close", () => {
|
||||
|
||||
@@ -172,6 +172,64 @@
|
||||
}
|
||||
}
|
||||
},
|
||||
"AUSBOSS_NODES_EditLock": {
|
||||
"display_name": "编辑锁定 🆎",
|
||||
"description": "让图片里编辑没有改动的部分保持在原位,颜色也不变。编辑模型会把整张图重画一遍,没改的部分常常略微放大、偏移或偏色。把这个节点放在模型和采样器之间,并接上你要编辑的图片。不需要遮罩:它在出图过程中看出编辑改了哪里。如果编辑改变的是整张图(新风格、新光线、新天气),它会自动让开。解码后再接上编辑锁定还原 🆎,未改动的部分就能换回你原图的像素。",
|
||||
"inputs": {
|
||||
"model": {
|
||||
"name": "模型"
|
||||
},
|
||||
"vae": {
|
||||
"name": "vae"
|
||||
},
|
||||
"image": {
|
||||
"name": "图像"
|
||||
},
|
||||
"lock": {
|
||||
"name": "锁定"
|
||||
},
|
||||
"tolerance": {
|
||||
"name": "容差"
|
||||
},
|
||||
"mask": {
|
||||
"name": "遮罩"
|
||||
}
|
||||
},
|
||||
"outputs": {
|
||||
"0": {
|
||||
"name": "模型"
|
||||
},
|
||||
"1": {
|
||||
"name": "编辑锁"
|
||||
}
|
||||
}
|
||||
},
|
||||
"AUSBOSS_NODES_EditLockRestore": {
|
||||
"display_name": "编辑锁定还原 🆎",
|
||||
"description": "把编辑没有改动的部分换回你原图自己的像素。放在 VAE 解码之后,并接上编辑锁定 🆎 的 edit_lock 输出。未改动的部分与原图逐像素一致,只有编辑的部分是新的。如果锁定让开了,图片原样通过。",
|
||||
"inputs": {
|
||||
"image": {
|
||||
"name": "图像"
|
||||
},
|
||||
"edit_lock": {
|
||||
"name": "编辑锁"
|
||||
},
|
||||
"feather": {
|
||||
"name": "羽化"
|
||||
}
|
||||
},
|
||||
"outputs": {
|
||||
"0": {
|
||||
"name": "图像"
|
||||
},
|
||||
"1": {
|
||||
"name": "编辑遮罩"
|
||||
},
|
||||
"2": {
|
||||
"name": "报告"
|
||||
}
|
||||
}
|
||||
},
|
||||
"AUSBOSS_NODES_Float": {
|
||||
"display_name": "浮点数 🆎",
|
||||
"description": "一个独立的 FLOAT 值。在多个节点之间共用同一个强度、缩放或 CFG,只需在一处修改。",
|
||||
@@ -1052,6 +1110,9 @@
|
||||
},
|
||||
"step_preview": {
|
||||
"name": "步骤预览"
|
||||
},
|
||||
"autosave": {
|
||||
"name": "自动保存"
|
||||
}
|
||||
},
|
||||
"outputs": {
|
||||
@@ -1240,6 +1301,35 @@
|
||||
}
|
||||
}
|
||||
},
|
||||
"AUSBOSS_NODES_TextLines": {
|
||||
"display_name": "逐行文本 🆎",
|
||||
"description": "文本的每一行都会让后面的节点运行一次。输入五条提示词,每行一条,运行一次就得到五张图。空行会跳过,行首的序号或项目符号可以去掉。",
|
||||
"inputs": {
|
||||
"text": {
|
||||
"name": "文本"
|
||||
},
|
||||
"split": {
|
||||
"name": "分隔方式"
|
||||
},
|
||||
"remove_numbering": {
|
||||
"name": "去掉序号"
|
||||
},
|
||||
"limit": {
|
||||
"name": "上限"
|
||||
}
|
||||
},
|
||||
"outputs": {
|
||||
"0": {
|
||||
"name": "行"
|
||||
},
|
||||
"1": {
|
||||
"name": "序号"
|
||||
},
|
||||
"2": {
|
||||
"name": "数量"
|
||||
}
|
||||
}
|
||||
},
|
||||
"AUSBOSS_NODES_TiledUpscale": {
|
||||
"display_name": "分块放大 🆎",
|
||||
"description": "以百万像素设置成品尺寸,并把图片切成模型可以逐块重画的图块。一块就能装下的图片保持完整。把图块送进采样器,再和 stitcher 一起送进分块放大拼合 🆎。小图不会被重画到远超其细节的程度:最多重画到原百万像素的 max_growth 倍,再用普通方式放大。report 会说明对你的图片做了什么。",
|
||||
@@ -1591,5 +1681,48 @@
|
||||
"AUSBOSS_NODES_WorkflowNoteTile": {
|
||||
"display_name": "工作流说明卡片 🆎",
|
||||
"description": "点击打开完整说明的小按钮。点击铅笔编辑,拖动角落调整大小。旧版节点包会将它显示为缺失节点,请更新 ComfyUI-AusBoss。"
|
||||
},
|
||||
"AUSBOSS_NODES_FlashVSRUpscale": {
|
||||
"display_name": "FlashVSR 视频放大 🆎",
|
||||
"description": "用 FlashVSR 模型把视频放大并变清晰。选好模型文件,设定放大倍数,然后运行。节点每次只读取几帧,所以长视频和短视频占用的显存一样多,每一帧都会返回。接上 Wan 2.1 VAE 可以得到更慢但更细致的结果。",
|
||||
"inputs": {
|
||||
"frames": {
|
||||
"name": "帧"
|
||||
},
|
||||
"model": {
|
||||
"name": "模型"
|
||||
},
|
||||
"scale": {
|
||||
"name": "放大倍数"
|
||||
},
|
||||
"color_match": {
|
||||
"name": "保留颜色"
|
||||
},
|
||||
"vae": {
|
||||
"name": "vae"
|
||||
}
|
||||
},
|
||||
"outputs": {
|
||||
"0": {
|
||||
"name": "帧"
|
||||
}
|
||||
}
|
||||
},
|
||||
"AUSBOSS_NODES_PromptVariations": {
|
||||
"display_name": "提示词变体 🆎",
|
||||
"description": "每次运行都会变化的提示词框。把选项写在花括号里,例如 {红色|蓝色|绿色} 的 {毛衣|夹克},节点会从每组里选一个。种子决定选哪个:同一个种子得到同一个提示词,每次换新种子就得到新的组合。没有花括号的文字原样输出。",
|
||||
"inputs": {
|
||||
"text": {
|
||||
"name": "文本"
|
||||
},
|
||||
"seed": {
|
||||
"name": "种子"
|
||||
}
|
||||
},
|
||||
"outputs": {
|
||||
"0": {
|
||||
"name": "提示词"
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -0,0 +1,455 @@
|
||||
"""Edit Lock: keep what an edit did not change. Pure torch, no ComfyUI imports.
|
||||
|
||||
At every sampling step an edit model predicts the finished picture. Where
|
||||
that prediction sits close to the original, the edit is not happening
|
||||
there, so the original is put back and the model carries on from the real
|
||||
picture. This is the target TailSteer trains a model toward
|
||||
(original where nothing changed, the model's edit elsewhere), applied while
|
||||
sampling instead of through training.
|
||||
|
||||
What counts as changed is measured against the quiet part of the picture:
|
||||
the tenth of it where the prediction differs least, which a local edit
|
||||
leaves near the level of the model's own redraw noise. The node watches the
|
||||
first steps, where it only undoes the model's zoom and shift, and decides
|
||||
once the picture has formed. An edit of the whole picture (a style, the
|
||||
light) has no quiet part: the first look checks for that and the lock steps
|
||||
aside. A soft change over the whole picture (rain, snow, haze) has no clear
|
||||
peak: the decision checks for that, and the picture is only held in place.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
|
||||
LOCK_MODES = ("auto", "always", "off")
|
||||
|
||||
# A spot counts as changed when it differs from the original by STRONG times
|
||||
# the quiet level, or by WEAK times while it touches a spot that does. So a
|
||||
# faint change next to a clear one belongs to the edit, and a faint change
|
||||
# alone is redraw noise. Set against what 14 plain MiniMax-H3-Image edits
|
||||
# changed in the end: this frees about 80% of it at the first look and 3%
|
||||
# of what stayed.
|
||||
WEAK = 2.0
|
||||
STRONG = 3.0
|
||||
# The model often draws the whole picture a little larger or shifted. That
|
||||
# is measured and discounted before anything is compared, up to these
|
||||
# limits; a bigger move is treated as part of the edit.
|
||||
# The node watches the start of a run and decides once the noise is down to
|
||||
# this share of where it began. The first prediction is a blur: rain, snow,
|
||||
# smoke and the light a new hair color throws on a face only show a few
|
||||
# steps in. While it watches it only undoes the zoom and shift.
|
||||
DECIDE_AT = 0.8
|
||||
# A local edit has a peak: somewhere the picture clearly changes. Rain, snow
|
||||
# and haze change the whole picture a little and no part clearly. A lock
|
||||
# would hold them back, or let them through in patches where it hands a wide
|
||||
# part of the picture to the edit. So two things are read at the decision,
|
||||
# in units of the quiet level: the strongest spot, and the strongest
|
||||
# hundredth when the lock would free SOFT_WIDE of the picture or more. On
|
||||
# recorded MiniMax-H3-Image runs the strongest spot read 4.7 or more for
|
||||
# every local edit, glasses and makeup included, and 3.7 or less for light
|
||||
# rain, fog and snow; with a wide part freed, the strongest hundredth read
|
||||
# 4.4 or more for local edits and 3.6 or less for rain and snow.
|
||||
SOFT_WIDE = 0.10
|
||||
SOFT_PEAK = 4.0
|
||||
# Kept this much or more: the report adds what to do if the edit is missing.
|
||||
# A soft effect over the whole picture (rain, snow, smoke) looks like a small
|
||||
# edit at the first look, and under the lock the model then drops it.
|
||||
NEARLY_ALL = 0.95
|
||||
MAX_ZOOM = 0.05
|
||||
MAX_SHIFT = 4.0
|
||||
# The whole picture is changing when even its calmest fifth sits this far
|
||||
# from the original at the first look, in units of the original's own
|
||||
# spread. Measured on MiniMax-H3-Image: 28 local edits at 1, 4 and 16 MP came
|
||||
# in at 0.22 or less (zoomed ones included), 12 whole-picture edits (styles,
|
||||
# light, weather) at 0.27 or more.
|
||||
WHOLE_PICTURE = 0.243
|
||||
# Blur and growth of the change map, as a share of the short side.
|
||||
REACH = 1.0 / 54.0
|
||||
# A result this much larger or smaller than the picture is another picture.
|
||||
SIZE_SLACK = 0.04
|
||||
|
||||
|
||||
def box_blur(x: torch.Tensor, radius: int) -> torch.Tensor:
|
||||
if radius <= 0:
|
||||
return x
|
||||
size = 2 * radius + 1
|
||||
return F.avg_pool2d(F.pad(x, (radius,) * 4, mode="replicate"), size, stride=1)
|
||||
|
||||
|
||||
def reach(height: int, width: int) -> int:
|
||||
return max(1, round(min(height, width) * REACH))
|
||||
|
||||
|
||||
def _moved(source: torch.Tensor, move: torch.Tensor) -> torch.Tensor:
|
||||
"""The original sampled at (zoom_x * x + shift_x, zoom_y * y + shift_y), x and y in cells from the centre.
|
||||
|
||||
move = [zoom_x - 1, zoom_y - 1, shift_x, shift_y].
|
||||
"""
|
||||
height, width = source.shape[-2:]
|
||||
theta = torch.zeros(source.shape[0], 2, 3, device=source.device, dtype=source.dtype)
|
||||
theta[:, 0, 0] = 1.0 + move[0]
|
||||
theta[:, 1, 1] = 1.0 + move[1]
|
||||
theta[:, 0, 2] = 2.0 * move[2] / width
|
||||
theta[:, 1, 2] = 2.0 * move[3] / height
|
||||
grid = F.affine_grid(theta, list(source.shape), align_corners=False)
|
||||
return F.grid_sample(source, grid, mode="bilinear", padding_mode="border", align_corners=False)
|
||||
|
||||
|
||||
def measure_redraw(pred: torch.Tensor, source: torch.Tensor) -> torch.Tensor:
|
||||
"""How much larger and how far shifted the prediction is drawn, as the move that lays the original over it.
|
||||
|
||||
pred and source are [1, C, H, W]. Fitted on the calmer half of the
|
||||
picture, so the edit itself does not pull the answer, from a coarse grid
|
||||
down to the full one. Returns [zoom_x - 1, zoom_y - 1, shift_x, shift_y];
|
||||
zeros when nothing can be measured.
|
||||
"""
|
||||
pred, source = pred.float(), source.float()
|
||||
move = torch.zeros(4, device=pred.device, dtype=torch.float32)
|
||||
limit = torch.tensor([MAX_ZOOM, MAX_ZOOM, MAX_SHIFT, MAX_SHIFT], device=pred.device)
|
||||
for level in (4, 2, 1):
|
||||
if min(pred.shape[-2:]) // level < 16:
|
||||
continue
|
||||
a = box_blur(F.avg_pool2d(pred, level) if level > 1 else pred, 1)
|
||||
b = box_blur(F.avg_pool2d(source, level) if level > 1 else source, 1)
|
||||
height, width = a.shape[-2:]
|
||||
ys = (torch.arange(height, device=a.device, dtype=torch.float32) + 0.5 - height / 2)[:, None].expand(height, width)
|
||||
xs = (torch.arange(width, device=a.device, dtype=torch.float32) + 0.5 - width / 2)[None, :].expand(height, width)
|
||||
here = move.clone()
|
||||
here[2:] = here[2:] / level
|
||||
for _ in range(6):
|
||||
laid = _moved(b, here)
|
||||
rest = a - laid
|
||||
rest = rest - rest.flatten(2).median(dim=2).values[:, :, None, None]
|
||||
wide = F.pad(laid, (1, 1, 1, 1), mode="replicate")
|
||||
gx = (wide[..., 1:-1, 2:] - wide[..., 1:-1, :-2]) / 2.0
|
||||
gy = (wide[..., 2:, 1:-1] - wide[..., :-2, 1:-1]) / 2.0
|
||||
size = rest.abs().mean(dim=1, keepdim=True)
|
||||
calm = (size <= size.flatten().median()).float()
|
||||
calm[..., :2, :] = 0.0
|
||||
calm[..., -2:, :] = 0.0
|
||||
calm[..., :, :2] = 0.0
|
||||
calm[..., :, -2:] = 0.0
|
||||
jac = (torch.stack([gx * xs, gy * ys, gx, gy], dim=-1) * calm[..., None]).reshape(-1, 4)
|
||||
normal = jac.T @ jac + 1e-6 * torch.eye(4, device=a.device)
|
||||
try:
|
||||
change = torch.linalg.solve(normal, jac.T @ (rest * calm).reshape(-1))
|
||||
except RuntimeError:
|
||||
return torch.zeros(4, device=pred.device, dtype=torch.float32)
|
||||
if not bool(torch.isfinite(change).all()):
|
||||
return torch.zeros(4, device=pred.device, dtype=torch.float32)
|
||||
here = here + change
|
||||
if float(change.abs().max()) < 1e-4:
|
||||
break
|
||||
move = here.clone()
|
||||
move[2:] = move[2:] * level
|
||||
move = torch.maximum(torch.minimum(move, limit), -limit)
|
||||
return move
|
||||
|
||||
|
||||
def redraws(pred: torch.Tensor, source: torch.Tensor) -> list[torch.Tensor]:
|
||||
"""measure_redraw for each picture of a batch."""
|
||||
return [measure_redraw(pred[index:index + 1], source[:1]) for index in range(pred.shape[0])]
|
||||
|
||||
|
||||
def difference(pred: torch.Tensor, source: torch.Tensor, moves: list[torch.Tensor] | None = None) -> torch.Tensor:
|
||||
"""How far each spot of the prediction is from the original, [B, 1, H, W].
|
||||
|
||||
Two things are not a change. A redraw of the whole picture slightly
|
||||
larger or shifted: the original is laid over the prediction first. And
|
||||
an even tint: each channel's middle difference is taken out. The result
|
||||
is in units of the original's own spread, which keeps the numbers alike
|
||||
between models.
|
||||
"""
|
||||
pred = pred.float()
|
||||
source = source.float()
|
||||
moves = moves if moves is not None else redraws(pred, source)
|
||||
out = []
|
||||
for index in range(pred.shape[0]):
|
||||
delta = pred[index:index + 1] - _moved(source[:1], moves[index])
|
||||
delta = delta - delta.flatten(2).median(dim=2).values[:, :, None, None]
|
||||
out.append(delta.abs().mean(dim=1, keepdim=True))
|
||||
magnitude = torch.cat(out) / source.std().clamp_min(1e-6)
|
||||
return box_blur(magnitude, reach(*magnitude.shape[-2:]))
|
||||
|
||||
|
||||
def moved_back(pred: torch.Tensor, moves: list[torch.Tensor]) -> torch.Tensor:
|
||||
"""The prediction with its zoom and shift undone, so it sits where the original does.
|
||||
|
||||
Resampling softens a latent, so a move under a tenth of a cell at the
|
||||
corner is left alone, and the rest is resampled bicubic: on recorded
|
||||
runs that kept 85% of the fine detail where bilinear kept 65%.
|
||||
"""
|
||||
out = []
|
||||
for index in range(pred.shape[0]):
|
||||
one = pred[index:index + 1]
|
||||
move = moves[index]
|
||||
height, width = one.shape[-2:]
|
||||
corner = (move[0].abs() * width / 2 + move[2].abs()).maximum(move[1].abs() * height / 2 + move[3].abs())
|
||||
if float(corner) < 0.1:
|
||||
out.append(one)
|
||||
continue
|
||||
theta = torch.zeros(1, 2, 3, device=one.device, dtype=torch.float32)
|
||||
theta[:, 0, 0] = 1.0 / (1.0 + move[0])
|
||||
theta[:, 1, 1] = 1.0 / (1.0 + move[1])
|
||||
theta[:, 0, 2] = -2.0 * move[2] / (1.0 + move[0]) / width
|
||||
theta[:, 1, 2] = -2.0 * move[3] / (1.0 + move[1]) / height
|
||||
grid = F.affine_grid(theta, list(one.shape), align_corners=False)
|
||||
back = F.grid_sample(one.float(), grid, mode="bicubic", padding_mode="border", align_corners=False)
|
||||
out.append(back.to(one.dtype))
|
||||
return torch.cat(out)
|
||||
|
||||
|
||||
def broad_change(pred: torch.Tensor, source: torch.Tensor) -> torch.Tensor:
|
||||
"""How far the calmest fifth of each picture sits from the original, [B].
|
||||
|
||||
Read on a smoothed copy, where a redraw one percent larger nearly
|
||||
vanishes and a new style, new light or new weather does not. An even
|
||||
tint is taken out first: the model adds one by itself, a stronger one
|
||||
on large pictures, and that is no edit.
|
||||
"""
|
||||
delta = (pred - source).float()
|
||||
delta = delta - delta.flatten(2).median(dim=2).values[:, :, None, None]
|
||||
delta = box_blur(delta, max(1, round(min(delta.shape[-2:]) / 30.0)))
|
||||
level = (delta * delta).mean(dim=1).sqrt().flatten(1)
|
||||
stride = max(1, level.shape[1] // 200000)
|
||||
spread = source.float().std().clamp_min(1e-6)
|
||||
return torch.quantile(level[:, ::stride], 0.2, dim=1) / spread
|
||||
|
||||
|
||||
def quiet_level(magnitude: torch.Tensor) -> torch.Tensor:
|
||||
"""The tenth of each picture that differs least, [B]."""
|
||||
flat = magnitude.flatten(1)
|
||||
stride = max(1, flat.shape[1] // 200000)
|
||||
return torch.quantile(flat[:, ::stride], 0.1, dim=1)
|
||||
|
||||
|
||||
def peaks(magnitude: torch.Tensor, quiet: torch.Tensor) -> tuple[float, float]:
|
||||
"""How clear the clearest change is, in units of the quiet level: (strongest spot, strongest hundredth)."""
|
||||
ratio = (magnitude / quiet.clamp_min(1e-6)[:, None, None, None]).flatten(1)
|
||||
stride = max(1, ratio.shape[1] // 200000)
|
||||
return float(ratio.max(dim=1).values.min()), float(torch.quantile(ratio[:, ::stride], 0.99, dim=1).min())
|
||||
|
||||
|
||||
def change_core(magnitude: torch.Tensor, quiet: torch.Tensor, tolerance: float = 1.0,
|
||||
held: torch.Tensor | None = None) -> torch.Tensor:
|
||||
"""1 where the edit changed the picture, 0 where it did not, hard-edged, [B, 1, H, W].
|
||||
|
||||
`held` is what earlier steps called changed: it stays changed, and
|
||||
faint change that touches it joins it.
|
||||
"""
|
||||
ratio = magnitude / quiet.clamp_min(1e-6)[:, None, None, None]
|
||||
weak = (ratio > WEAK * tolerance).float()
|
||||
core = (ratio > STRONG * tolerance).float()
|
||||
if held is not None:
|
||||
weak = torch.maximum(weak, held)
|
||||
core = torch.maximum(core, held)
|
||||
for _ in range(sum(core.shape[-2:])):
|
||||
grown = F.max_pool2d(core, 3, stride=1, padding=1) * weak
|
||||
if torch.equal(grown, core):
|
||||
break
|
||||
core = grown
|
||||
return core
|
||||
|
||||
|
||||
def soften(core: torch.Tensor) -> torch.Tensor:
|
||||
"""The hard change map with a margin and a soft edge, so the edit has room to blend in."""
|
||||
grow = reach(*core.shape[-2:])
|
||||
return box_blur(F.max_pool2d(core, 2 * grow + 1, stride=1, padding=grow), grow)
|
||||
|
||||
|
||||
def fit_latent(latent: torch.Tensor, height: int, width: int) -> torch.Tensor | None:
|
||||
"""The original's latent on the result's grid, or None when they are different pictures."""
|
||||
h, w = latent.shape[-2:]
|
||||
if (h, w) == (height, width):
|
||||
return latent
|
||||
if abs(h / height - 1.0) > SIZE_SLACK or abs(w / width - 1.0) > SIZE_SLACK:
|
||||
return None
|
||||
return F.interpolate(latent, size=(height, width), mode="bilinear", align_corners=False)
|
||||
|
||||
|
||||
class EditLock:
|
||||
"""One lock: the original, the settings, and what the last run kept."""
|
||||
|
||||
def __init__(self, mode: str = "auto", tolerance: float = 1.0):
|
||||
self.mode = mode
|
||||
self.tolerance = tolerance
|
||||
self.source: torch.Tensor | None = None # IMAGE, [1, H, W, 3]
|
||||
self.latent: torch.Tensor | None = None # [1, C, H, W] or [1, C, 1, H, W], model space
|
||||
self.user_mask: torch.Tensor | None = None # MASK, white = the edit may happen
|
||||
self.status = "waiting"
|
||||
self.mask: torch.Tensor | None = None # [B, 1, h, w], 1 = changed, soft-edged
|
||||
self.core: torch.Tensor | None = None # the same before its margin, hard-edged
|
||||
self.quiet: torch.Tensor | None = None # [B], the quiet level at the decision
|
||||
self.broad: float | None = None # the first look's whole-picture reading
|
||||
self.wide: float | None = None # share of the picture the lock would free, at the decision
|
||||
self.top: float | None = None # the strongest spot then, in units of the quiet level
|
||||
self.peak: float | None = None # and the strongest hundredth
|
||||
self._sigma: float | None = None
|
||||
self._first = True
|
||||
self._first_sigma = 1.0
|
||||
self._decided = False
|
||||
|
||||
def _begin(self) -> None:
|
||||
self.status = "waiting"
|
||||
self.mask = None
|
||||
self.core = None
|
||||
self.quiet = None
|
||||
self.broad = None
|
||||
self.wide = None
|
||||
self.top = None
|
||||
self.peak = None
|
||||
self._first = True
|
||||
self._decided = False
|
||||
|
||||
def _decides_now(self, args: dict, sigma: float) -> bool:
|
||||
"""True once the noise is down far enough, or this is the run's last step."""
|
||||
steps = (args.get("model_options") or {}).get("transformer_options", {}).get("sample_sigmas")
|
||||
if steps is None or len(steps) < 2:
|
||||
return True # the sampler does not say where the run ends: decide at once
|
||||
if sigma <= float(steps[-2]) * 1.0001:
|
||||
return True
|
||||
return sigma <= DECIDE_AT * self._first_sigma
|
||||
|
||||
def step(self, args: dict) -> torch.Tensor:
|
||||
"""ComfyUI's post-CFG hook: the prediction comes in, the locked prediction goes out."""
|
||||
pred = args["denoised"]
|
||||
sigma = float(args["sigma"].flatten()[0])
|
||||
if self._sigma is None or sigma >= self._sigma:
|
||||
self._begin() # the noise level went up: a new run
|
||||
self._sigma = sigma
|
||||
if self.status in ("whole picture", "size", "video", "vae") or self.latent is None:
|
||||
return pred
|
||||
video = pred.ndim == 5
|
||||
if video and pred.shape[2] != 1:
|
||||
self.status = "video"
|
||||
return pred
|
||||
flat = pred[:, :, 0] if video else pred
|
||||
latent = self.latent[:, :, 0] if self.latent.ndim == 5 else self.latent
|
||||
if latent.shape[1] != flat.shape[1]:
|
||||
self.status = "vae"
|
||||
return pred
|
||||
source = fit_latent(latent.to(device=flat.device, dtype=torch.float32), *flat.shape[-2:])
|
||||
if source is None:
|
||||
self.status = "size"
|
||||
return pred
|
||||
if self.status == "soft":
|
||||
# nothing is put back, the picture is only held where the original sits
|
||||
back = moved_back(flat, redraws(flat, source))
|
||||
return back[:, :, None] if video else back
|
||||
if self.user_mask is not None:
|
||||
mask = self._given_mask(flat)
|
||||
else:
|
||||
if self._first:
|
||||
self._first_sigma = sigma
|
||||
self.broad = float(broad_change(flat, source).max())
|
||||
if self.mode == "auto" and self.broad > WHOLE_PICTURE:
|
||||
self.status = "whole picture"
|
||||
return pred
|
||||
moves = redraws(flat, source)
|
||||
magnitude = difference(flat, source, moves)
|
||||
if not self._decided and not self._decides_now(args, sigma):
|
||||
# still watching: the model draws its own picture, only its zoom and shift are undone
|
||||
self._first = False
|
||||
self.status = "watching"
|
||||
if self.broad > WHOLE_PICTURE:
|
||||
return pred # too much is changing to trust the measured move
|
||||
back = moved_back(flat, moves)
|
||||
return back[:, :, None] if video else back
|
||||
if not self._decided:
|
||||
self._decided = True
|
||||
# the yardstick is set once, here: later steps are read against it
|
||||
self.quiet = quiet_level(magnitude)
|
||||
self.top, self.peak = peaks(magnitude, self.quiet)
|
||||
self.wide = float(change_core(magnitude, self.quiet, self.tolerance).flatten(1).mean(dim=1).max())
|
||||
thin = self.top < SOFT_PEAK or (self.wide >= SOFT_WIDE and self.peak < SOFT_PEAK)
|
||||
if self.mode == "auto" and not self._first and thin:
|
||||
# a soft change over the whole picture: a lock would hold it back or let it through in patches
|
||||
self.status = "soft"
|
||||
back = moved_back(flat, moves)
|
||||
return back[:, :, None] if video else back
|
||||
held = self.core if self.core is not None and self.core.shape[0] == flat.shape[0] else None
|
||||
self.core = change_core(magnitude, self.quiet.to(magnitude.device), self.tolerance,
|
||||
held.to(magnitude.device) if held is not None else None).detach()
|
||||
mask = soften(self.core)
|
||||
self._first = False
|
||||
self.mask = mask.detach()
|
||||
self.status = "locked"
|
||||
keep = (1.0 - mask).to(flat.dtype)
|
||||
locked = flat * (1.0 - keep) + source.to(flat.dtype) * keep
|
||||
return locked[:, :, None] if video else locked
|
||||
|
||||
def _given_mask(self, flat: torch.Tensor) -> torch.Tensor:
|
||||
mask = self.user_mask
|
||||
if mask.ndim == 2:
|
||||
mask = mask[None]
|
||||
mask = mask[:1, None].to(device=flat.device, dtype=torch.float32)
|
||||
mask = F.interpolate(mask, size=flat.shape[-2:], mode="bilinear", align_corners=False).clamp(0.0, 1.0)
|
||||
return soften(mask).expand(flat.shape[0], -1, -1, -1)
|
||||
|
||||
def kept_share(self) -> float | None:
|
||||
if self.status != "locked" or self.mask is None:
|
||||
return None
|
||||
return float((self.mask < 0.5).float().mean())
|
||||
|
||||
def pixel_mask(self, height: int, width: int) -> torch.Tensor | None:
|
||||
"""The change mask at the picture's size, [B, H, W], or None when nothing was locked."""
|
||||
if self.status != "locked" or self.mask is None:
|
||||
return None
|
||||
mask = F.interpolate(self.mask.float().cpu(), size=(height, width), mode="bilinear", align_corners=False)
|
||||
return mask[:, 0].clamp(0.0, 1.0)
|
||||
|
||||
|
||||
def restore_pixels(result: torch.Tensor, source: torch.Tensor, mask: torch.Tensor,
|
||||
feather: int) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
"""The original's own pixels where the mask is black. result [B, H, W, C], source [1, h, w, 3], mask [b, H, W].
|
||||
|
||||
The fade lies outside the edit: the mask is widened by the feather first.
|
||||
"""
|
||||
frames, height, width = int(result.shape[0]), int(result.shape[1]), int(result.shape[2])
|
||||
edit = mask[:, None].float()
|
||||
if edit.shape[0] not in (1, frames):
|
||||
edit = edit[:1]
|
||||
if feather > 0:
|
||||
size = 2 * feather + 1
|
||||
edit = F.max_pool2d(edit, size, stride=1, padding=feather)
|
||||
edit = box_blur(edit, feather)
|
||||
original = source[..., :3].float()
|
||||
if tuple(original.shape[1:3]) != (height, width):
|
||||
original = F.interpolate(original.movedim(-1, 1), size=(height, width), mode="bicubic",
|
||||
align_corners=False, antialias=True).movedim(1, -1).clamp(0.0, 1.0)
|
||||
weight = edit.movedim(1, -1).to(result.device, result.dtype)
|
||||
out = result.clone()
|
||||
out[..., :3] = result[..., :3] * weight + original.to(result.device, result.dtype) * (1.0 - weight)
|
||||
edit_mask = edit[:, 0].expand(frames, -1, -1).clone()
|
||||
return out, edit_mask
|
||||
|
||||
|
||||
def describe(lock: EditLock) -> str:
|
||||
if lock.status == "locked":
|
||||
kept = lock.kept_share() or 0.0
|
||||
how = "your mask" if lock.user_mask is not None else "what the edit changed"
|
||||
text = (f"Edit Lock kept {kept * 100:.0f}% of your picture as it was and left "
|
||||
f"{(1.0 - kept) * 100:.0f}% to the edit ({how}).")
|
||||
if kept >= NEARLY_ALL and lock.user_mask is None:
|
||||
# rain, snow or smoke over the whole picture reads as a small edit and is held back
|
||||
text += (" If your edit is missing, it is one that spreads thinly over the whole picture "
|
||||
"(rain, snow, smoke): set lock to off.")
|
||||
return text
|
||||
if lock.status == "whole picture":
|
||||
return ("Edit Lock stepped aside: the edit changes the whole picture, so there was nothing to keep. "
|
||||
"Set lock to always to lock anyway.")
|
||||
if lock.status == "soft":
|
||||
return ("Edit Lock stepped aside: the edit spreads thinly over the whole picture (rain, snow, haze), "
|
||||
"so no part of it stayed as it was. It only held the picture in place. "
|
||||
"Set lock to always to lock anyway.")
|
||||
if lock.status == "off":
|
||||
return "Edit Lock is off."
|
||||
if lock.status == "size":
|
||||
return ("Edit Lock did nothing: the picture and the result are different sizes. "
|
||||
"Give it the picture at the size of the result.")
|
||||
if lock.status == "video":
|
||||
return "Edit Lock did nothing: it works on single pictures, not video."
|
||||
if lock.status == "vae":
|
||||
return "Edit Lock did nothing: this VAE does not belong to this model."
|
||||
return "Edit Lock has not run yet: plug its model output into the sampler."
|
||||
@@ -0,0 +1,946 @@
|
||||
"""FlashVSR video upscaling: the model's parts and the loop that runs a clip.
|
||||
|
||||
FlashVSR (OpenImagingLab, Apache-2.0) sharpens a clip in one model step.
|
||||
It is a Wan 2.1 1.3B transformer that is also shown the low-quality frames,
|
||||
a small projector that turns those frames into tokens, and a small decoder
|
||||
that reads the frames again while it draws the result.
|
||||
|
||||
Everything here is plain PyTorch. The model looks at a clip a few frames at
|
||||
a time and keeps a short memory of the frames before, so a long clip costs
|
||||
no more graphics memory than a short one. Module and parameter names follow
|
||||
the names in the published weights, which is what lets them load.
|
||||
|
||||
Shapes used below: a token is 16x16 output pixels of one latent frame, a
|
||||
latent frame is 4 video frames, and an attention block is 2 latent frames
|
||||
by 8x8 tokens (128 tokens, 128x128 pixels over 8 frames).
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import math
|
||||
import os
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
PIXEL_STEP = 128 # the working size is a multiple of this in both directions
|
||||
BLOCK_TOKENS = 128
|
||||
KEPT_GROUPS = 3 # earlier 2-frame groups each layer remembers
|
||||
LOCAL_REACH = 11 # a block looks at blocks inside an 11x11 neighbourhood
|
||||
SPARSE_RATIO = 2.0
|
||||
SPARSE_REFERENCE_AREA = 768 * 1280
|
||||
FIRST_PASS_LATENTS = 6
|
||||
PASS_LATENTS = 2
|
||||
MIN_GROUPS = 3 # the first pass alone needs 25 frames
|
||||
TIMESTEP = 1000.0
|
||||
DEFAULT_SLOT_BUDGET = 12000
|
||||
|
||||
|
||||
def _cast(parameter: torch.Tensor, like: torch.Tensor) -> torch.Tensor:
|
||||
return parameter.to(device=like.device, dtype=like.dtype)
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------
|
||||
# Sizes and frame counts
|
||||
# --------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _even(value: float) -> int:
|
||||
return max(2, int(round(value / 2.0)) * 2)
|
||||
|
||||
|
||||
def plan_sizes(height: int, width: int, scale: float) -> tuple[int, int, int, int]:
|
||||
"""Output size (even numbers) and the padded size the model works at."""
|
||||
out_h, out_w = _even(height * scale), _even(width * scale)
|
||||
work_h = math.ceil(out_h / PIXEL_STEP) * PIXEL_STEP
|
||||
work_w = math.ceil(out_w / PIXEL_STEP) * PIXEL_STEP
|
||||
return out_h, out_w, work_h, work_w
|
||||
|
||||
|
||||
def plan_groups(frame_count: int) -> int:
|
||||
"""How many 8-frame groups the padded clip has.
|
||||
|
||||
The model reads 8n+1 frames and returns 8n-3, so the clip is padded by
|
||||
repeating its last frame until 8n-3 covers every frame it has.
|
||||
"""
|
||||
return max(MIN_GROUPS, math.ceil((int(frame_count) + 3) / 8))
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------
|
||||
# Position encoding and attention blocks
|
||||
# --------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _axis_angles(dim: int, positions: torch.Tensor) -> torch.Tensor:
|
||||
steps = torch.arange(0, dim, 2, dtype=torch.float64)[: dim // 2] / dim
|
||||
return torch.outer(positions.double(), 1.0 / (10000.0**steps))
|
||||
|
||||
|
||||
def rope_table(head_dim: int, first_frame: int, frames: int, rows: int, cols: int, device):
|
||||
"""Rotation angles for every token, as (cos, sin) of shape [tokens, 1, head_dim / 2]."""
|
||||
side = head_dim // 3
|
||||
t = _axis_angles(head_dim - 2 * side, torch.arange(first_frame, first_frame + frames))
|
||||
r = _axis_angles(side, torch.arange(rows))
|
||||
c = _axis_angles(side, torch.arange(cols))
|
||||
angles = torch.cat(
|
||||
[
|
||||
t[:, None, None, :].expand(frames, rows, cols, -1),
|
||||
r[None, :, None, :].expand(frames, rows, cols, -1),
|
||||
c[None, None, :, :].expand(frames, rows, cols, -1),
|
||||
],
|
||||
dim=-1,
|
||||
).reshape(frames * rows * cols, 1, head_dim // 2)
|
||||
return angles.cos().float().to(device), angles.sin().float().to(device)
|
||||
|
||||
|
||||
def apply_rope(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor, heads: int) -> torch.Tensor:
|
||||
"""Rotate each pair of channels of x [tokens, dim] by its token's angle."""
|
||||
tokens, dim = x.shape
|
||||
a, b = x.float().view(tokens, heads, dim // heads // 2, 2).unbind(-1)
|
||||
return torch.stack((a * cos - b * sin, a * sin + b * cos), dim=-1).view(tokens, dim).to(x.dtype)
|
||||
|
||||
|
||||
def to_blocks(x: torch.Tensor, frames: int, rows: int, cols: int) -> torch.Tensor:
|
||||
"""[tokens, dim] in (frame, row, column) order -> [blocks, 128, dim]."""
|
||||
dim = x.shape[-1]
|
||||
x = x.view(frames // 2, 2, rows // 8, 8, cols // 8, 8, dim).permute(0, 2, 4, 1, 3, 5, 6)
|
||||
return x.reshape(-1, BLOCK_TOKENS, dim)
|
||||
|
||||
|
||||
def from_blocks(blocks: torch.Tensor, frames: int, rows: int, cols: int) -> torch.Tensor:
|
||||
dim = blocks.shape[-1]
|
||||
x = blocks.view(frames // 2, rows // 8, cols // 8, 2, 8, 8, dim).permute(0, 3, 1, 4, 2, 5, 6)
|
||||
return x.reshape(frames * rows * cols, dim)
|
||||
|
||||
|
||||
def local_block_mask(block_rows: int, block_cols: int, reach: int, device) -> torch.Tensor:
|
||||
"""[blocks, blocks] True where a key block is inside a query block's neighbourhood."""
|
||||
r = torch.arange(block_rows, device=device).repeat_interleave(block_cols)
|
||||
c = torch.arange(block_cols, device=device).repeat(block_rows)
|
||||
dr, dc = r[None, :] - r[:, None], c[None, :] - c[:, None]
|
||||
low, high = -(reach // 2), reach - 1 - reach // 2
|
||||
return (dr >= low) & (dr <= high) & (dc >= low) & (dc <= high)
|
||||
|
||||
|
||||
def choose_key_blocks(
|
||||
q_blocks: torch.Tensor,
|
||||
k_blocks: torch.Tensor,
|
||||
heads: int,
|
||||
local: torch.Tensor,
|
||||
query_groups: int,
|
||||
keep: int,
|
||||
) -> torch.Tensor:
|
||||
"""Pick, for every head, which key blocks each query block reads.
|
||||
|
||||
A block's mean query is scored against every block's mean key. Blocks
|
||||
outside the neighbourhood are ruled out, the scores of a query block are
|
||||
turned into shares that sum to one, and for each head and each 2-frame
|
||||
group of queries the `keep` largest shares win. Returns two
|
||||
[heads, query blocks, key blocks] tensors: True where a pair is kept,
|
||||
and the shares themselves (0 outside the neighbourhood).
|
||||
"""
|
||||
nq, _, dim = q_blocks.shape
|
||||
nk = k_blocks.shape[0]
|
||||
head_dim = dim // heads
|
||||
q_mean = q_blocks.float().mean(1).view(nq, heads, head_dim).transpose(0, 1)
|
||||
k_mean = k_blocks.float().mean(1).view(nk, heads, head_dim).transpose(0, 1)
|
||||
scores = q_mean @ k_mean.transpose(1, 2) / math.sqrt(head_dim)
|
||||
per_group = local.shape[0]
|
||||
allowed = local.repeat(nq // per_group, nk // per_group)
|
||||
shares = scores.masked_fill(~allowed, float("-inf")).softmax(dim=-1)
|
||||
flat = shares.view(heads, query_groups, per_group * nk)
|
||||
keep = min(flat.shape[-1] - 1, int(keep))
|
||||
threshold = flat.topk(keep + 1, dim=-1).values[..., -1:]
|
||||
return (flat > threshold).view(heads, nq, nk), shares
|
||||
|
||||
|
||||
def sparse_block_attention(
|
||||
q: torch.Tensor,
|
||||
k: torch.Tensor,
|
||||
v: torch.Tensor,
|
||||
chosen: torch.Tensor,
|
||||
shares: torch.Tensor | None = None,
|
||||
slot_budget: int = DEFAULT_SLOT_BUDGET,
|
||||
) -> torch.Tensor:
|
||||
"""Attention where each query block only reads its chosen key blocks.
|
||||
|
||||
q is [heads, query blocks, 128, head_dim], k and v are
|
||||
[heads, key blocks, 128, head_dim] and chosen is
|
||||
[heads, query blocks, key blocks]. Each (head, query block) is a row
|
||||
that gathers the keys it chose. Rows with a similar number of chosen
|
||||
blocks run together, all at the width of the row that chose most. A
|
||||
row that chose nothing returns zeros. `slot_budget` caps
|
||||
(rows x chosen blocks) per call, which caps memory.
|
||||
|
||||
With `shares` (from choose_key_blocks) a narrower row is filled up
|
||||
with its next-best blocks from its neighbourhood, so a call needs no
|
||||
padding mask, which is what makes it fast: a row then reads a few
|
||||
blocks more than it chose, never fewer. Without `shares` the padding
|
||||
is masked and every row reads exactly what it chose.
|
||||
"""
|
||||
heads, nq, tokens, head_dim = q.shape
|
||||
nk = k.shape[1]
|
||||
rows = chosen.reshape(heads * nq, nk)
|
||||
counts = rows.sum(dim=1)
|
||||
if shares is None:
|
||||
rank, usable = rows.float(), counts
|
||||
else:
|
||||
flat_shares = shares.reshape(heads * nq, nk).float()
|
||||
rank, usable = rows.float() * 2 + flat_shares, (flat_shares > 0).sum(dim=1)
|
||||
order = torch.argsort(counts)
|
||||
picks = torch.argsort(rank, dim=1, descending=True, stable=True)
|
||||
picks = picks + (torch.arange(heads, device=q.device).repeat_interleave(nq) * nk)[:, None]
|
||||
q_rows = q.reshape(heads * nq, tokens, head_dim)
|
||||
k_rows = k.reshape(heads * nk, tokens, head_dim)
|
||||
v_rows = v.reshape(heads * nk, tokens, head_dim)
|
||||
out = torch.zeros_like(q_rows)
|
||||
sorted_counts = counts[order].tolist()
|
||||
total = len(sorted_counts)
|
||||
start = 0
|
||||
while start < total and sorted_counts[start] == 0:
|
||||
start += 1
|
||||
while start < total:
|
||||
end = start + 1
|
||||
while end < total and (end - start + 1) * sorted_counts[end] <= slot_budget:
|
||||
end += 1
|
||||
rows_here = order[start:end]
|
||||
widest = sorted_counts[end - 1]
|
||||
index = picks[rows_here, :widest]
|
||||
keys = k_rows[index].view(end - start, widest * tokens, head_dim)
|
||||
values = v_rows[index].view(end - start, widest * tokens, head_dim)
|
||||
mask = None
|
||||
real = torch.arange(widest, device=q.device)[None, :] < usable[rows_here][:, None]
|
||||
if not bool(real.all()):
|
||||
mask = real.repeat_interleave(tokens, dim=1)[:, None, :]
|
||||
out[rows_here] = F.scaled_dot_product_attention(q_rows[rows_here], keys, values, attn_mask=mask)
|
||||
start = end
|
||||
return out.view(heads, nq, tokens, head_dim)
|
||||
|
||||
|
||||
class AttentionPlan:
|
||||
"""What stays the same for every layer and pass at one working size."""
|
||||
|
||||
def __init__(self, rows: int, cols: int, device, slot_budget: int = DEFAULT_SLOT_BUDGET):
|
||||
self.rows, self.cols = rows, cols
|
||||
self.per_group = (rows // 8) * (cols // 8)
|
||||
self.local = local_block_mask(rows // 8, cols // 8, LOCAL_REACH, device)
|
||||
ratio = SPARSE_RATIO * SPARSE_REFERENCE_AREA / (rows * 16 * cols * 16)
|
||||
self.keep = int(self.per_group * self.per_group * ratio) - 1
|
||||
self.slot_budget = slot_budget
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------
|
||||
# The transformer
|
||||
# --------------------------------------------------------------------------
|
||||
|
||||
|
||||
class _RMSNorm(nn.Module):
|
||||
def __init__(self, dim: int, eps: float, device=None, dtype=None):
|
||||
super().__init__()
|
||||
self.eps = eps
|
||||
self.weight = nn.Parameter(torch.ones(dim, device=device, dtype=dtype))
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
y = x.float()
|
||||
y = y * torch.rsqrt(y.pow(2).mean(dim=-1, keepdim=True) + self.eps)
|
||||
return y.to(x.dtype) * _cast(self.weight, x)
|
||||
|
||||
|
||||
def _heads_first(blocks: torch.Tensor, heads: int) -> torch.Tensor:
|
||||
count, tokens, dim = blocks.shape
|
||||
return blocks.view(count, tokens, heads, dim // heads).permute(2, 0, 1, 3).contiguous()
|
||||
|
||||
|
||||
class _SelfAttention(nn.Module):
|
||||
def __init__(self, dim, heads, eps, operations, device=None, dtype=None):
|
||||
super().__init__()
|
||||
self.heads = heads
|
||||
kw = {"device": device, "dtype": dtype}
|
||||
self.q = operations.Linear(dim, dim, **kw)
|
||||
self.k = operations.Linear(dim, dim, **kw)
|
||||
self.v = operations.Linear(dim, dim, **kw)
|
||||
self.o = operations.Linear(dim, dim, **kw)
|
||||
self.norm_q = _RMSNorm(dim, eps, **kw)
|
||||
self.norm_k = _RMSNorm(dim, eps, **kw)
|
||||
|
||||
def forward(self, x, rope, grid, memory: dict, plan: AttentionPlan):
|
||||
frames, rows, cols = grid
|
||||
cos, sin = rope
|
||||
q = to_blocks(apply_rope(self.norm_q(self.q(x))[0], cos, sin, self.heads), *grid)
|
||||
k = to_blocks(apply_rope(self.norm_k(self.k(x))[0], cos, sin, self.heads), *grid)
|
||||
v = to_blocks(self.v(x)[0], *grid)
|
||||
if memory.get("k") is not None:
|
||||
k = torch.cat([memory["k"], k], dim=0)
|
||||
v = torch.cat([memory["v"], v], dim=0)
|
||||
chosen, shares = choose_key_blocks(q, k, self.heads, plan.local, frames // 2, plan.keep)
|
||||
out = sparse_block_attention(
|
||||
_heads_first(q, self.heads),
|
||||
_heads_first(k, self.heads),
|
||||
_heads_first(v, self.heads),
|
||||
chosen,
|
||||
shares,
|
||||
plan.slot_budget,
|
||||
)
|
||||
if k.shape[0] // plan.per_group > KEPT_GROUPS:
|
||||
k, v = k[plan.per_group :], v[plan.per_group :]
|
||||
memory["k"], memory["v"] = k, v
|
||||
out = out.permute(1, 2, 0, 3).reshape(-1, BLOCK_TOKENS, x.shape[-1])
|
||||
return self.o(from_blocks(out, *grid)[None])
|
||||
|
||||
|
||||
class _CrossAttention(nn.Module):
|
||||
def __init__(self, dim, heads, eps, operations, device=None, dtype=None):
|
||||
super().__init__()
|
||||
self.heads = heads
|
||||
kw = {"device": device, "dtype": dtype}
|
||||
self.q = operations.Linear(dim, dim, **kw)
|
||||
self.k = operations.Linear(dim, dim, **kw)
|
||||
self.v = operations.Linear(dim, dim, **kw)
|
||||
self.o = operations.Linear(dim, dim, **kw)
|
||||
self.norm_q = _RMSNorm(dim, eps, **kw)
|
||||
self.norm_k = _RMSNorm(dim, eps, **kw)
|
||||
|
||||
def _split(self, x: torch.Tensor) -> torch.Tensor:
|
||||
batch, tokens, dim = x.shape
|
||||
return x.view(batch, tokens, self.heads, dim // self.heads).transpose(1, 2)
|
||||
|
||||
def forward(self, x, context, memory: dict):
|
||||
if memory.get("k") is None: # the prompt never changes during a run
|
||||
memory["k"] = self._split(self.norm_k(self.k(context)))
|
||||
memory["v"] = self._split(self.v(context))
|
||||
out = F.scaled_dot_product_attention(self._split(self.norm_q(self.q(x))), memory["k"], memory["v"])
|
||||
return self.o(out.transpose(1, 2).reshape(x.shape))
|
||||
|
||||
|
||||
class _Block(nn.Module):
|
||||
def __init__(self, dim, heads, ffn_dim, eps, operations, device=None, dtype=None):
|
||||
super().__init__()
|
||||
self.dim, self.eps = dim, eps
|
||||
kw = {"device": device, "dtype": dtype}
|
||||
self.self_attn = _SelfAttention(dim, heads, eps, operations, **kw)
|
||||
self.cross_attn = _CrossAttention(dim, heads, eps, operations, **kw)
|
||||
self.norm3 = operations.LayerNorm(dim, eps=eps, elementwise_affine=True, **kw)
|
||||
self.ffn = nn.Sequential(
|
||||
operations.Linear(dim, ffn_dim, **kw), nn.GELU(approximate="tanh"), operations.Linear(ffn_dim, dim, **kw)
|
||||
)
|
||||
self.modulation = nn.Parameter(torch.zeros(1, 6, dim, **kw))
|
||||
|
||||
def _plain_norm(self, x: torch.Tensor) -> torch.Tensor:
|
||||
return F.layer_norm(x, (self.dim,), eps=self.eps)
|
||||
|
||||
def forward(self, x, context, time_mod, rope, grid, memory: dict, plan: AttentionPlan):
|
||||
shift_a, scale_a, gate_a, shift_m, scale_m, gate_m = (_cast(self.modulation, x) + time_mod).chunk(6, dim=1)
|
||||
x = x + gate_a * self.self_attn(self._plain_norm(x) * (1 + scale_a) + shift_a, rope, grid, memory["self"], plan)
|
||||
x = x + self.cross_attn(self.norm3(x), context, memory["cross"])
|
||||
return x + gate_m * self.ffn(self._plain_norm(x) * (1 + scale_m) + shift_m)
|
||||
|
||||
|
||||
class _Head(nn.Module):
|
||||
def __init__(self, dim, out_dim, eps, operations, device=None, dtype=None):
|
||||
super().__init__()
|
||||
self.dim, self.eps = dim, eps
|
||||
self.head = operations.Linear(dim, out_dim * 4, device=device, dtype=dtype)
|
||||
self.modulation = nn.Parameter(torch.zeros(1, 2, dim, device=device, dtype=dtype))
|
||||
|
||||
def forward(self, x, time_embed):
|
||||
shift, scale = (_cast(self.modulation, x) + time_embed).chunk(2, dim=1)
|
||||
return self.head(F.layer_norm(x, (self.dim,), eps=self.eps) * (1 + scale) + shift)
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------
|
||||
# The low-quality projector
|
||||
# --------------------------------------------------------------------------
|
||||
|
||||
|
||||
class _ChannelNorm(nn.Module):
|
||||
def __init__(self, dim: int, device=None, dtype=None):
|
||||
super().__init__()
|
||||
self.scale = dim**0.5
|
||||
self.gamma = nn.Parameter(torch.ones(dim, 1, 1, 1, device=device, dtype=dtype))
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
return F.normalize(x.float(), dim=1).to(x.dtype) * self.scale * _cast(self.gamma, x)
|
||||
|
||||
|
||||
class LowQualityProjector(nn.Module):
|
||||
"""Turns enlarged low-quality frames into one set of tokens per 4 frames.
|
||||
|
||||
Two convolutions that only look backwards in time, each halving the
|
||||
frame rate, then one linear layer. `push` is fed the first frame alone,
|
||||
then 4 frames at a time, and remembers what the next call needs.
|
||||
"""
|
||||
|
||||
def __init__(self, out_dim=1536, hidden=(2048, 3072), patch=16, operations=nn, device=None, dtype=None):
|
||||
super().__init__()
|
||||
self.patch = patch
|
||||
kw = {"device": device, "dtype": dtype}
|
||||
self.conv1 = operations.Conv3d(3 * patch * patch, hidden[0], (4, 3, 3), stride=(2, 1, 1), **kw)
|
||||
self.norm1 = _ChannelNorm(hidden[0], **kw)
|
||||
self.conv2 = operations.Conv3d(hidden[0], hidden[1], (4, 3, 3), stride=(2, 1, 1), **kw)
|
||||
self.norm2 = _ChannelNorm(hidden[1], **kw)
|
||||
self.linear_layers = nn.ModuleList([operations.Linear(hidden[1], out_dim, **kw)])
|
||||
|
||||
def _fold(self, clip: torch.Tensor) -> torch.Tensor:
|
||||
batch, channels, frames, height, width = clip.shape
|
||||
p = self.patch
|
||||
x = clip.view(batch, channels, frames, height // p, p, width // p, p).permute(0, 1, 4, 6, 2, 3, 5)
|
||||
return x.reshape(batch, channels * p * p, frames, height // p, width // p)
|
||||
|
||||
@staticmethod
|
||||
def _causal(conv, x: torch.Tensor, earlier) -> torch.Tensor:
|
||||
front = 2
|
||||
if earlier is not None:
|
||||
x = torch.cat([earlier, x], dim=2)
|
||||
front -= earlier.shape[2]
|
||||
return conv(F.pad(x, (1, 1, 1, 1, front, 0), mode="replicate"))
|
||||
|
||||
def push(self, memory: dict, clip: torch.Tensor):
|
||||
"""clip is [1, 3, frames, height, width] in -1..1. Returns tokens, or None for the first frame."""
|
||||
first = not memory
|
||||
if first:
|
||||
clip = torch.cat([clip[:, :, :1].repeat(1, 1, 3, 1, 1), clip], dim=2)
|
||||
x = self._fold(clip)
|
||||
tail = x[:, :, -2:].clone()
|
||||
x = F.silu(self.norm1(self._causal(self.conv1, x, memory.get("conv1"))))
|
||||
memory["conv1"] = tail
|
||||
tail = x[:, :, -2:].clone()
|
||||
if first:
|
||||
memory["conv2"] = tail
|
||||
return None
|
||||
x = F.silu(self.norm2(self._causal(self.conv2, x, memory["conv2"])))
|
||||
memory["conv2"] = tail
|
||||
return self.linear_layers[0](x.flatten(2).transpose(1, 2))
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------
|
||||
# The fast decoder
|
||||
# --------------------------------------------------------------------------
|
||||
|
||||
|
||||
class _SoftClamp(nn.Module):
|
||||
def forward(self, x):
|
||||
return torch.tanh(x / 3) * 3
|
||||
|
||||
|
||||
class _MemoryBlock(nn.Module):
|
||||
"""A residual block that also sees the frame before."""
|
||||
|
||||
def __init__(self, width: int, operations, device=None, dtype=None):
|
||||
super().__init__()
|
||||
kw = {"device": device, "dtype": dtype}
|
||||
self.conv = nn.Sequential(
|
||||
operations.Conv2d(width * 2, width, 3, padding=1, **kw),
|
||||
nn.ReLU(),
|
||||
operations.Conv2d(width, width, 3, padding=1, **kw),
|
||||
nn.ReLU(),
|
||||
operations.Conv2d(width, width, 3, padding=1, **kw),
|
||||
)
|
||||
|
||||
def forward(self, x, past):
|
||||
return F.relu(self.conv(torch.cat([x, past], dim=1)) + x)
|
||||
|
||||
|
||||
class _TimeGrow(nn.Module):
|
||||
"""Makes `stride` frames out of each frame."""
|
||||
|
||||
def __init__(self, width: int, stride: int, operations, device=None, dtype=None):
|
||||
super().__init__()
|
||||
self.conv = operations.Conv2d(width, width * stride, 1, bias=False, device=device, dtype=dtype)
|
||||
|
||||
def forward(self, x):
|
||||
_, channels, height, width = x.shape
|
||||
return self.conv(x).reshape(-1, channels, height, width)
|
||||
|
||||
|
||||
class _Upsample(nn.Module):
|
||||
def forward(self, x):
|
||||
return F.interpolate(x, scale_factor=2, mode="nearest")
|
||||
|
||||
|
||||
class ConditionedDecoder(nn.Module):
|
||||
"""Latent frames plus the low-quality frames in, 4 video frames per latent frame out."""
|
||||
|
||||
def __init__(self, latent_channels=16, widths=(512, 256, 128, 128), operations=nn, device=None, dtype=None):
|
||||
super().__init__()
|
||||
kw = {"device": device, "dtype": dtype}
|
||||
a, b, c, d = widths
|
||||
|
||||
def conv(n_in, n_out, bias=True):
|
||||
return operations.Conv2d(n_in, n_out, 3, padding=1, bias=bias, **kw)
|
||||
|
||||
def memory_blocks(width):
|
||||
return [_MemoryBlock(width, operations, **kw) for _ in range(3)]
|
||||
|
||||
self.latent_channels = latent_channels
|
||||
self.decoder = nn.Sequential(
|
||||
_SoftClamp(),
|
||||
conv(latent_channels + 3 * 4 * 8 * 8, a),
|
||||
nn.ReLU(),
|
||||
conv(a, a, bias=False),
|
||||
nn.ReLU(),
|
||||
*memory_blocks(a),
|
||||
_Upsample(),
|
||||
_TimeGrow(a, 1, operations, **kw),
|
||||
conv(a, b, bias=False),
|
||||
*memory_blocks(b),
|
||||
_Upsample(),
|
||||
_TimeGrow(b, 2, operations, **kw),
|
||||
conv(b, c, bias=False),
|
||||
*memory_blocks(c),
|
||||
_Upsample(),
|
||||
_TimeGrow(c, 2, operations, **kw),
|
||||
conv(c, d, bias=False),
|
||||
nn.ReLU(),
|
||||
conv(d, d, bias=False),
|
||||
nn.ReLU(),
|
||||
conv(d, 3),
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def fold_frames(clip: torch.Tensor) -> torch.Tensor:
|
||||
"""[3, 4, height, width] -> [768, height / 8, width / 8]."""
|
||||
channels, frames, height, width = clip.shape
|
||||
x = clip.view(channels, frames, height // 8, 8, width // 8, 8).permute(0, 1, 3, 5, 2, 4)
|
||||
return x.reshape(channels * frames * 64, height // 8, width // 8)
|
||||
|
||||
def decode(self, memory: dict, x: torch.Tensor) -> torch.Tensor:
|
||||
"""x is [latent frames, 768 + latent channels, h, w]. Returns [4 x latent frames, 3, 8h, 8w] in 0..1.
|
||||
|
||||
The first 3 frames of the first call are lead-in and are dropped.
|
||||
"""
|
||||
first = not memory
|
||||
memory["started"] = True
|
||||
for index, layer in enumerate(self.decoder):
|
||||
if isinstance(layer, _MemoryBlock):
|
||||
before = memory.get(index)
|
||||
if before is None:
|
||||
before = torch.zeros_like(x[:1])
|
||||
memory[index] = x[-1:].clone()
|
||||
x = layer(x, torch.cat([before, x[:-1]], dim=0))
|
||||
else:
|
||||
x = layer(x)
|
||||
return x[3:] if first else x
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------
|
||||
# The whole model
|
||||
# --------------------------------------------------------------------------
|
||||
|
||||
|
||||
class FlashVSRModel(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
dim=1536,
|
||||
heads=12,
|
||||
ffn_dim=8960,
|
||||
layers=30,
|
||||
latent_channels=16,
|
||||
text_dim=4096,
|
||||
text_len=512,
|
||||
freq_dim=256,
|
||||
eps=1e-6,
|
||||
projector_hidden=(2048, 3072),
|
||||
decoder_widths=(512, 256, 128, 128),
|
||||
operations=nn,
|
||||
device=None,
|
||||
dtype=None,
|
||||
):
|
||||
super().__init__()
|
||||
self.dim, self.heads, self.freq_dim, self.latent_channels = dim, heads, freq_dim, latent_channels
|
||||
kw = {"device": device, "dtype": dtype}
|
||||
self.patch_embedding = operations.Conv3d(latent_channels, dim, (1, 2, 2), stride=(1, 2, 2), **kw)
|
||||
self.text_embedding = nn.Sequential(
|
||||
operations.Linear(text_dim, dim, **kw), nn.GELU(approximate="tanh"), operations.Linear(dim, dim, **kw)
|
||||
)
|
||||
self.time_embedding = nn.Sequential(
|
||||
operations.Linear(freq_dim, dim, **kw), nn.SiLU(), operations.Linear(dim, dim, **kw)
|
||||
)
|
||||
self.time_projection = nn.Sequential(nn.SiLU(), operations.Linear(dim, dim * 6, **kw))
|
||||
self.blocks = nn.ModuleList([_Block(dim, heads, ffn_dim, eps, operations, **kw) for _ in range(layers)])
|
||||
self.head = _Head(dim, latent_channels, eps, operations, **kw)
|
||||
self.lq_proj_in = LowQualityProjector(dim, projector_hidden, operations=operations, **kw)
|
||||
self.tc_decoder = ConditionedDecoder(latent_channels, decoder_widths, operations=operations, **kw)
|
||||
self.register_buffer("positive_conditioning", torch.zeros(text_len, text_dim, **kw))
|
||||
|
||||
def start(self, like: torch.Tensor) -> dict:
|
||||
"""The values that never change during one clip."""
|
||||
half = self.freq_dim // 2
|
||||
steps = torch.pow(10000.0, -torch.arange(half, dtype=torch.float64) / half) * TIMESTEP
|
||||
wave = torch.cat([torch.cos(steps), torch.sin(steps)])[None].to(device=like.device, dtype=like.dtype)
|
||||
time_embed = self.time_embedding(wave)
|
||||
return {
|
||||
"time_embed": time_embed[:, None],
|
||||
"time_mod": self.time_projection(time_embed).unflatten(1, (6, self.dim)),
|
||||
"context": self.text_embedding(_cast(self.positive_conditioning, like)[None]),
|
||||
"layers": [{"self": {}, "cross": {}} for _ in self.blocks],
|
||||
"projector": {},
|
||||
"decoder": {},
|
||||
}
|
||||
|
||||
def denoise(self, state: dict, noise: torch.Tensor, lq_tokens: torch.Tensor, first_frame: int, plan: AttentionPlan):
|
||||
"""One pass over a few latent frames. Returns the clean latent frames."""
|
||||
x = self.patch_embedding(noise)
|
||||
frames, rows, cols = x.shape[2:]
|
||||
grid = (frames, rows, cols)
|
||||
x = x.flatten(2).transpose(1, 2) + lq_tokens
|
||||
rope = rope_table(self.dim // self.heads, first_frame, frames, rows, cols, x.device)
|
||||
for block, memory in zip(self.blocks, state["layers"]):
|
||||
x = block(x, state["context"], state["time_mod"], rope, grid, memory, plan)
|
||||
x = self.head(x, state["time_embed"])
|
||||
x = x.view(1, frames, rows, cols, 2, 2, self.latent_channels).permute(0, 6, 1, 2, 4, 3, 5)
|
||||
return noise - x.reshape(1, self.latent_channels, frames, rows * 2, cols * 2)
|
||||
|
||||
|
||||
def pass_layout(groups: int):
|
||||
"""(first latent frame, latent frames) of each pass over a clip of `groups` 8-frame groups."""
|
||||
yield 0, FIRST_PASS_LATENTS
|
||||
for index in range(1, groups - 2):
|
||||
yield 4 + 2 * index, PASS_LATENTS
|
||||
|
||||
|
||||
def match_colors(frames: torch.Tensor, source: torch.Tensor) -> torch.Tensor:
|
||||
"""Give each frame its source frame's mean and spread, per colour channel. Both [n, 3, h, w]."""
|
||||
f, s = frames.float(), source.float()
|
||||
f_mean, s_mean = f.mean(dim=(2, 3), keepdim=True), s.mean(dim=(2, 3), keepdim=True)
|
||||
f_std, s_std = f.std(dim=(2, 3), keepdim=True) + 1e-5, s.std(dim=(2, 3), keepdim=True) + 1e-5
|
||||
return (f - f_mean) / f_std * s_std + s_mean
|
||||
|
||||
|
||||
TILE_OVERLAP = 128
|
||||
|
||||
|
||||
def plan_tiles(work_h: int, work_w: int, max_pixels: int | None) -> list[tuple[int, int, int, int]]:
|
||||
"""Cut the working size into as few overlapping tiles as fit `max_pixels` each.
|
||||
|
||||
Returns (top, left, height, width) boxes. Tile sides are multiples of
|
||||
128 and corners sit on the 16-pixel token grid. One box, the whole
|
||||
picture, when it fits or when no limit is given.
|
||||
"""
|
||||
if not max_pixels or work_h * work_w <= max_pixels:
|
||||
return [(0, 0, work_h, work_w)]
|
||||
|
||||
def side(total: int, parts: int) -> int:
|
||||
if parts == 1:
|
||||
return total
|
||||
return min(total, math.ceil((total + (parts - 1) * TILE_OVERLAP) / parts / PIXEL_STEP) * PIXEL_STEP)
|
||||
|
||||
best = None
|
||||
for rows in range(1, work_h // PIXEL_STEP + 1):
|
||||
for cols in range(1, work_w // PIXEL_STEP + 1):
|
||||
tile_h, tile_w = side(work_h, rows), side(work_w, cols)
|
||||
if tile_h * tile_w <= max_pixels and (best is None or rows * cols < best[0] * best[1]):
|
||||
best = (rows, cols, tile_h, tile_w)
|
||||
if best is None: # even one 128x128 tile is over the limit: use the smallest tiles there are
|
||||
best = (work_h // PIXEL_STEP, work_w // PIXEL_STEP, PIXEL_STEP, PIXEL_STEP)
|
||||
rows, cols, tile_h, tile_w = best
|
||||
|
||||
def starts(total: int, tile: int, parts: int) -> list[int]:
|
||||
if parts == 1:
|
||||
return [0]
|
||||
return [round(i * (total - tile) / (parts - 1) / 16) * 16 for i in range(parts)]
|
||||
|
||||
return [(top, left, tile_h, tile_w) for top in starts(work_h, tile_h, rows) for left in starts(work_w, tile_w, cols)]
|
||||
|
||||
|
||||
def tile_weight(box: tuple[int, int, int, int], work_h: int, work_w: int) -> torch.Tensor:
|
||||
"""[height, width] blend weights for a tile: 1 inside, fading over the overlap on edges that meet another tile."""
|
||||
top, left, height, width = box
|
||||
|
||||
def ramp(length: int, fade_start: bool, fade_end: bool) -> torch.Tensor:
|
||||
w = torch.ones(length)
|
||||
fade = (torch.arange(TILE_OVERLAP) + 0.5) / TILE_OVERLAP
|
||||
if fade_start:
|
||||
w[:TILE_OVERLAP] = fade
|
||||
if fade_end:
|
||||
w[-TILE_OVERLAP:] = torch.minimum(w[-TILE_OVERLAP:], fade.flip(0))
|
||||
return w
|
||||
|
||||
rows = ramp(height, top > 0, top + height < work_h)
|
||||
cols = ramp(width, left > 0, left + width < work_w)
|
||||
return rows[:, None] * cols[None, :]
|
||||
|
||||
|
||||
@torch.no_grad()
|
||||
def upscale_frames(
|
||||
model: FlashVSRModel,
|
||||
frames: torch.Tensor,
|
||||
scale: float,
|
||||
*,
|
||||
device,
|
||||
dtype,
|
||||
color_match: bool = True,
|
||||
seed: int = 0,
|
||||
slot_budget: int = DEFAULT_SLOT_BUDGET,
|
||||
max_tile_pixels: int | None = None,
|
||||
decode_latents=None,
|
||||
on_pass=None,
|
||||
) -> torch.Tensor:
|
||||
"""Upscale a clip. `frames` is [count, height, width, 3] in 0..1; so is the result, at the new size.
|
||||
|
||||
A picture larger than `max_tile_pixels` is done in overlapping tiles,
|
||||
one after the other, and blended; that is what bounds graphics memory.
|
||||
`decode_latents`, when given, takes the clean latent frames of a tile
|
||||
([1, 16, frames, h, w]) and returns [frames, height, width, 3]; it is how
|
||||
another decoder (the Wan 2.1 VAE) is used in place of the fast one.
|
||||
`on_pass(done, total)` is called after each pass.
|
||||
"""
|
||||
count, height, width, _ = frames.shape
|
||||
out_h, out_w, work_h, work_w = plan_sizes(height, width, scale)
|
||||
groups = plan_groups(count)
|
||||
layout = list(pass_layout(groups))
|
||||
tiles = plan_tiles(work_h, work_w, max_tile_pixels)
|
||||
|
||||
def enlarged(start: int, stop: int) -> torch.Tensor:
|
||||
"""Frames start..stop-1 of the padded clip at the working size: [n, 3, work_h, work_w] in 0..1."""
|
||||
index = torch.arange(start, stop).clamp(max=count - 1)
|
||||
x = frames[index][..., :3].to(device).movedim(-1, 1).float()
|
||||
if (height, width) != (out_h, out_w):
|
||||
x = F.interpolate(x, size=(out_h, out_w), mode="bicubic", align_corners=False).clamp(0, 1)
|
||||
return F.pad(x, (0, work_w - out_w, 0, work_h - out_h), mode="replicate")
|
||||
|
||||
noise = torch.randn(
|
||||
(1, model.latent_channels, 2 * groups, work_h // 8, work_w // 8), generator=torch.Generator().manual_seed(seed)
|
||||
)
|
||||
result = torch.zeros((count, out_h, out_w, 3), dtype=torch.float32)
|
||||
total_weight = torch.zeros((work_h, work_w), dtype=torch.float32)
|
||||
like = torch.zeros(1, device=device, dtype=dtype)
|
||||
done = 0
|
||||
for box in tiles:
|
||||
top, left, tile_h, tile_w = box
|
||||
weight = tile_weight(box, work_h, work_w)
|
||||
total_weight[top : top + tile_h, left : left + tile_w] += weight
|
||||
weight = weight.to(device)
|
||||
|
||||
def crop(start: int, stop: int) -> torch.Tensor:
|
||||
return enlarged(start, stop)[..., top : top + tile_h, left : left + tile_w]
|
||||
|
||||
def signed_clip(start: int, stop: int) -> torch.Tensor:
|
||||
return (crop(start, stop) * 2 - 1).movedim(0, 1)[None].to(dtype)
|
||||
|
||||
def store(decoded: torch.Tensor, written: int) -> int:
|
||||
"""Add decoded frames [n, 3, tile_h, tile_w] to the result; returns how many frames the tile has given."""
|
||||
decoded = decoded[: count - written].float()
|
||||
if decoded.shape[0] == 0:
|
||||
return written
|
||||
stop = written + decoded.shape[0]
|
||||
if color_match:
|
||||
decoded = match_colors(decoded, crop(written, stop))
|
||||
bottom, right = min(top + tile_h, out_h), min(left + tile_w, out_w)
|
||||
if bottom > top and right > left:
|
||||
piece = (decoded.clamp(0, 1) * weight)[..., : bottom - top, : right - left]
|
||||
result[written:stop, top:bottom, left:right] += piece.movedim(1, -1).cpu()
|
||||
return stop
|
||||
|
||||
state = model.start(like)
|
||||
plan = AttentionPlan(tile_h // 16, tile_w // 16, device, slot_budget)
|
||||
tile_noise = noise[..., top // 8 : (top + tile_h) // 8, left // 8 : (left + tile_w) // 8]
|
||||
kept_latents = []
|
||||
written = 0
|
||||
model.lq_proj_in.push(state["projector"], signed_clip(0, 1))
|
||||
for first, latent_count in layout:
|
||||
tokens = torch.cat(
|
||||
[
|
||||
model.lq_proj_in.push(state["projector"], signed_clip(4 * i + 1, 4 * i + 5))
|
||||
for i in range(first, first + latent_count)
|
||||
],
|
||||
dim=1,
|
||||
)
|
||||
chunk = tile_noise[:, :, first : first + latent_count].to(device=device, dtype=dtype)
|
||||
latents = model.denoise(state, chunk, tokens, first, plan)
|
||||
if decode_latents is not None:
|
||||
kept_latents.append(latents.float().cpu())
|
||||
else:
|
||||
for i in range(first, first + latent_count):
|
||||
if written >= count:
|
||||
break
|
||||
source = signed_clip(0, 1)[0].repeat(1, 4, 1, 1) if i == 0 else signed_clip(4 * i - 3, 4 * i + 1)[0]
|
||||
both = torch.cat([ConditionedDecoder.fold_frames(source), latents[0, :, i - first]], dim=0)
|
||||
written = store(model.tc_decoder.decode(state["decoder"], both[None]), written)
|
||||
done += 1
|
||||
if on_pass is not None:
|
||||
on_pass(done, len(layout) * len(tiles))
|
||||
state.clear() # the layers' memory is no longer needed; the next tile or another decoder wants the room
|
||||
if decode_latents is not None:
|
||||
decoded = decode_latents(torch.cat(kept_latents, dim=2)).movedim(-1, 1)
|
||||
for start in range(0, min(count, decoded.shape[0]), 8):
|
||||
written = store(decoded[start : start + 8].to(device), written)
|
||||
if len(tiles) > 1:
|
||||
result /= total_weight[:out_h, :out_w, None]
|
||||
return result
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------
|
||||
# ComfyUI side: loading the file and running a clip. ComfyUI is imported
|
||||
# inside the functions so the code above runs without it.
|
||||
# --------------------------------------------------------------------------
|
||||
|
||||
MODEL_PARTS = {
|
||||
"blocks.0.self_attn.q.weight": "the transformer",
|
||||
"lq_proj_in.conv1.weight": "the low-quality projector",
|
||||
"tc_decoder.decoder.1.weight": "the fast decoder",
|
||||
"positive_conditioning": "the stored prompt",
|
||||
}
|
||||
NOT_A_FLASHVSR_FILE = (
|
||||
"This is not a FlashVSR model file for this node: it has no {missing}. "
|
||||
"Pick the FlashVSR file named in the node's help page."
|
||||
)
|
||||
|
||||
_loaded: dict = {}
|
||||
|
||||
|
||||
# Measured on one clip (640x1024, bf16): about 9.6 KB of graphics memory per working pixel on top of
|
||||
# the weights. Most of it is the keys and values every layer keeps of the last few frames.
|
||||
WORKING_BYTES_PER_PIXEL = 12_000
|
||||
MEMORY_TO_ASK_FOR = 16 * 1024**3
|
||||
MEMORY_HEADROOM = 1024**3
|
||||
|
||||
|
||||
def working_memory(pixels: int, bytes_per_value: int = 2) -> int:
|
||||
"""A rough count of the graphics memory a picture of `pixels` needs besides the weights, in bytes."""
|
||||
return int(pixels * WORKING_BYTES_PER_PIXEL * bytes_per_value / 2)
|
||||
|
||||
|
||||
def tile_pixels_that_fit(free_bytes: int, bytes_per_value: int = 2) -> int:
|
||||
"""The largest tile, in pixels, that `free_bytes` of graphics memory can hold."""
|
||||
return max(PIXEL_STEP * PIXEL_STEP, int((free_bytes - MEMORY_HEADROOM) / working_memory(1_000_000, bytes_per_value) * 1_000_000))
|
||||
|
||||
|
||||
def load_flashvsr(path: str, dtype: torch.dtype | None = None):
|
||||
"""The model behind ComfyUI's memory manager, and the number type it runs in. Kept between runs."""
|
||||
import comfy.model_management as mm
|
||||
import comfy.model_patcher
|
||||
import comfy.ops
|
||||
import comfy.utils
|
||||
|
||||
stamp = (path, os.path.getmtime(path), dtype)
|
||||
if _loaded.get("stamp") == stamp:
|
||||
return _loaded["patcher"], _loaded["dtype"]
|
||||
_loaded.clear()
|
||||
weights = comfy.utils.load_torch_file(path, safe_load=True)
|
||||
missing = [what for key, what in MODEL_PARTS.items() if key not in weights]
|
||||
if missing:
|
||||
raise ValueError(NOT_A_FLASHVSR_FILE.format(missing=" or ".join(missing)))
|
||||
device, offload = mm.get_torch_device(), mm.unet_offload_device()
|
||||
if dtype is None:
|
||||
# The file is 16-bit, so that is the first choice: no weight has to be converted.
|
||||
if mm.should_use_fp16(device):
|
||||
dtype = torch.float16
|
||||
elif mm.should_use_bf16(device):
|
||||
dtype = torch.bfloat16
|
||||
else:
|
||||
dtype = torch.float32
|
||||
# manual_cast: every layer brings its weights to the device and number type of what it is given,
|
||||
# so a model ComfyUI keeps partly, or wholly, outside graphics memory still runs.
|
||||
model = FlashVSRModel(operations=comfy.ops.manual_cast, device=offload, dtype=dtype)
|
||||
model.eval()
|
||||
# ComfyUI's own wrapper for a model: the one that loads weights as they are needed, where it has one.
|
||||
wrapper = getattr(comfy.model_patcher, "CoreModelPatcher", comfy.model_patcher.ModelPatcher)
|
||||
patcher = wrapper(model, load_device=device, offload_device=offload)
|
||||
assign = bool(getattr(patcher, "is_dynamic", lambda: False)())
|
||||
model.load_state_dict(weights, strict=True, assign=assign)
|
||||
model.requires_grad_(False)
|
||||
_loaded.update(stamp=stamp, patcher=patcher, dtype=dtype)
|
||||
return patcher, dtype
|
||||
|
||||
|
||||
def _memory_for_work(patcher, device) -> int:
|
||||
"""Graphics memory left for the picture once all of the model's weights are in it."""
|
||||
import comfy.model_management as mm
|
||||
|
||||
free = mm.get_free_memory(device)
|
||||
try:
|
||||
still_to_load = max(0, int(patcher.model_size()) - int(patcher.loaded_size()))
|
||||
except Exception:
|
||||
still_to_load = 0
|
||||
return free - still_to_load
|
||||
|
||||
|
||||
def run_flashvsr(
|
||||
path: str,
|
||||
frames: torch.Tensor,
|
||||
scale: float,
|
||||
color_match: bool = True,
|
||||
vae=None,
|
||||
seed: int = 0,
|
||||
max_tile_pixels: int | None = None,
|
||||
):
|
||||
"""Upscale `frames` with the model file at `path`. `vae`, when given, decodes in place of the fast decoder."""
|
||||
import comfy.latent_formats
|
||||
import comfy.model_management as mm
|
||||
|
||||
from ._execution_helpers import progress_bar, raise_if_interrupted
|
||||
|
||||
patcher, dtype = load_flashvsr(path)
|
||||
device = patcher.load_device
|
||||
_, _, work_h, work_w = plan_sizes(frames.shape[1], frames.shape[2], scale)
|
||||
wanted = working_memory(work_h * work_w, dtype.itemsize)
|
||||
mm.load_models_gpu([patcher], memory_required=min(wanted, MEMORY_TO_ASK_FOR))
|
||||
if max_tile_pixels is None and device.type != "cpu":
|
||||
max_tile_pixels = tile_pixels_that_fit(_memory_for_work(patcher, device), dtype.itemsize)
|
||||
passes = len(list(pass_layout(plan_groups(frames.shape[0]))))
|
||||
|
||||
decode = None
|
||||
if vae is not None:
|
||||
latent_format = comfy.latent_formats.Wan21()
|
||||
|
||||
def decode(latents: torch.Tensor) -> torch.Tensor:
|
||||
images = vae.decode(latent_format.process_out(latents))
|
||||
return images.reshape(-1, *images.shape[-3:])
|
||||
|
||||
while True:
|
||||
progress = progress_bar(passes * len(plan_tiles(work_h, work_w, max_tile_pixels)))
|
||||
|
||||
def on_pass(done: int, total: int) -> None:
|
||||
raise_if_interrupted()
|
||||
if progress is not None:
|
||||
progress.update_absolute(done, total)
|
||||
|
||||
try:
|
||||
return upscale_frames(
|
||||
patcher.model,
|
||||
frames,
|
||||
scale,
|
||||
device=device,
|
||||
dtype=dtype,
|
||||
color_match=color_match,
|
||||
seed=seed,
|
||||
max_tile_pixels=max_tile_pixels,
|
||||
decode_latents=decode,
|
||||
on_pass=on_pass,
|
||||
)
|
||||
except torch.cuda.OutOfMemoryError:
|
||||
# Out of graphics memory: start again in tiles half the size, down to the smallest tile.
|
||||
tried = max_tile_pixels if max_tile_pixels is not None else work_h * work_w
|
||||
if tried <= PIXEL_STEP * PIXEL_STEP:
|
||||
raise
|
||||
max_tile_pixels = max(PIXEL_STEP * PIXEL_STEP, min(tried, work_h * work_w) // 2)
|
||||
mm.soft_empty_cache()
|
||||
|
||||
|
||||
MODEL_FOLDER = "diffusion_models"
|
||||
EXPECTED_MODEL_NAME = "flashvsr_v1.1_fp16.safetensors"
|
||||
MODEL_NOT_FOUND = (
|
||||
"FlashVSR Upscale could not find {name} in ComfyUI/models/diffusion_models. "
|
||||
"Download the FlashVSR model file named in the node's help page and put it there."
|
||||
)
|
||||
|
||||
|
||||
def list_flashvsr_models() -> list[str]:
|
||||
"""FlashVSR files in the diffusion models folder; every file there when none is named FlashVSR."""
|
||||
try:
|
||||
import folder_paths
|
||||
except ImportError: # Offline tests import this module without ComfyUI.
|
||||
return [EXPECTED_MODEL_NAME]
|
||||
names = list(folder_paths.get_filename_list(MODEL_FOLDER))
|
||||
ours = [name for name in names if "flashvsr" in name.lower().replace("_", "").replace("-", "")]
|
||||
return ours or names or [EXPECTED_MODEL_NAME]
|
||||
|
||||
|
||||
def resolve_flashvsr_model(name: str) -> str:
|
||||
import folder_paths
|
||||
|
||||
path = folder_paths.get_full_path(MODEL_FOLDER, name)
|
||||
if not path:
|
||||
raise FileNotFoundError(MODEL_NOT_FOUND.format(name=name))
|
||||
return path
|
||||
@@ -0,0 +1,170 @@
|
||||
"""Pictures Save Image holds until they are kept.
|
||||
|
||||
With Auto save off, a run writes its pictures to ComfyUI's temp folder and
|
||||
shows them on the node. A click on the node's Save button keeps the ones
|
||||
that are wanted: they are copied into the output folder under the names the
|
||||
node would have given them. Everything here works from names this module
|
||||
made itself; the only thing a caller sends is the run's token and which
|
||||
pictures to keep.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from datetime import datetime
|
||||
import json
|
||||
from pathlib import Path
|
||||
import re
|
||||
import secrets
|
||||
import shutil
|
||||
|
||||
from ._image_save_helpers import (
|
||||
IMAGE_SAVE_LOCK,
|
||||
existing_action,
|
||||
name_stem,
|
||||
next_free_counter,
|
||||
plan_exact_names,
|
||||
plan_local_names,
|
||||
require_output_path,
|
||||
resolve_output_root,
|
||||
sidecar_path,
|
||||
split_prefix,
|
||||
strip_image_extension,
|
||||
)
|
||||
|
||||
HELD_FOLDER = "ausboss_held"
|
||||
_TOKEN = re.compile(r"^[0-9a-f]{16}$")
|
||||
GONE = "These pictures are no longer held. Run the workflow again."
|
||||
|
||||
|
||||
def plan_destinations(settings: dict, count: int, now: datetime, root: Path, output_root: Path, picks=None):
|
||||
"""Where each picture goes, as (path, what to do if it exists).
|
||||
|
||||
Call it with IMAGE_SAVE_LOCK held: the counter it reads stays free only
|
||||
until the lock is let go. `picks` keeps some of a run's pictures; an
|
||||
exact name then keeps the number it would have had in the whole run.
|
||||
"""
|
||||
wanted = list(range(count)) if picks is None else list(picks)
|
||||
extension = settings["extension"]
|
||||
exact = settings.get("exact") or ""
|
||||
if exact:
|
||||
names = plan_exact_names(strip_image_extension(exact), extension, count)
|
||||
return [(root / names[index], settings["on_existing"]) for index in wanted]
|
||||
subfolder, base = split_prefix(settings["filename_prefix"])
|
||||
folder = root / subfolder if subfolder else root
|
||||
require_output_path(folder, output_root)
|
||||
stem = name_stem(
|
||||
base,
|
||||
{"date": bool(settings["name_date"]), "time": bool(settings["name_time"]), "size": bool(settings["name_size"])},
|
||||
now,
|
||||
int(settings["width"]),
|
||||
int(settings["height"]),
|
||||
)
|
||||
counter = bool(settings["name_counter"])
|
||||
existing = [entry.name for entry in folder.iterdir()] if folder.is_dir() else []
|
||||
names = plan_local_names(
|
||||
stem, extension, len(wanted), counter=counter, batch=bool(settings["name_batch"]),
|
||||
next_counter=next_free_counter(existing, stem),
|
||||
)
|
||||
# With a counter the name is free by construction. Without one, the
|
||||
# legacy policy decides whether to overwrite.
|
||||
policy = "overwrite" if counter else settings["on_existing"]
|
||||
return [(folder / name, policy) for name in names]
|
||||
|
||||
|
||||
def settle_destinations(planned, has_caption: bool, output_root: Path) -> list[str]:
|
||||
"""Check every destination, then say 'write' or 'skip' for each.
|
||||
|
||||
The whole batch is settled before any write: an "error" on a later
|
||||
picture stops the save before the first one lands.
|
||||
"""
|
||||
for path, _policy in planned:
|
||||
require_output_path(path, output_root)
|
||||
if has_caption:
|
||||
require_output_path(sidecar_path(path), output_root)
|
||||
return [
|
||||
existing_action(path.exists() or (has_caption and sidecar_path(path).exists()), policy)
|
||||
for path, policy in planned
|
||||
]
|
||||
|
||||
|
||||
def _held_name(token: str, index: int, extension: str) -> str:
|
||||
return f"{token}_{index:03d}.{extension}"
|
||||
|
||||
|
||||
def hold_pictures(temp_root, frames, settings: dict, encode) -> tuple[str, list[dict]]:
|
||||
"""Write a run's pictures to the temp folder and remember how to name them.
|
||||
|
||||
`encode(path, frame)` writes one picture in the run's format. Returns
|
||||
the run's token and the entries the node shows them from.
|
||||
"""
|
||||
folder = Path(temp_root) / HELD_FOLDER
|
||||
folder.mkdir(parents=True, exist_ok=True)
|
||||
token = secrets.token_hex(8)
|
||||
previews = []
|
||||
for index, frame in enumerate(frames):
|
||||
name = _held_name(token, index, settings["extension"])
|
||||
encode(folder / name, frame)
|
||||
previews.append({"filename": name, "subfolder": HELD_FOLDER, "type": "temp"})
|
||||
manifest = dict(settings, count=len(previews), kept={})
|
||||
(folder / f"{token}.json").write_text(json.dumps(manifest), encoding="utf-8")
|
||||
return token, previews
|
||||
|
||||
|
||||
def keep_pictures(temp_root, output_root, token, picks=None, now: datetime | None = None) -> dict:
|
||||
"""Copy held pictures into the output folder. Returns what was saved.
|
||||
|
||||
`picks` lists which pictures of the run to keep, counted from 0; None
|
||||
keeps them all. A picture kept before is not saved a second time.
|
||||
"""
|
||||
if not isinstance(token, str) or not _TOKEN.match(token):
|
||||
raise ValueError("Save Image: that is not a run it is holding.")
|
||||
folder = Path(temp_root) / HELD_FOLDER
|
||||
manifest_path = folder / f"{token}.json"
|
||||
if not manifest_path.is_file():
|
||||
raise LookupError(GONE)
|
||||
manifest = json.loads(manifest_path.read_text(encoding="utf-8"))
|
||||
count = int(manifest["count"])
|
||||
kept = {int(index): str(path) for index, path in dict(manifest.get("kept") or {}).items()}
|
||||
|
||||
if picks is None:
|
||||
wanted = list(range(count))
|
||||
else:
|
||||
wanted = []
|
||||
for pick in picks:
|
||||
if isinstance(pick, bool) or not isinstance(pick, int) or not 0 <= pick < count:
|
||||
raise ValueError("Save Image: that picture is not in this run.")
|
||||
if pick not in wanted:
|
||||
wanted.append(pick)
|
||||
todo = [index for index in wanted if index not in kept]
|
||||
|
||||
output_root = Path(output_root)
|
||||
caption = str(manifest.get("caption") or "")
|
||||
has_caption = bool(caption.strip())
|
||||
skipped: list[int] = []
|
||||
saved: list[int] = []
|
||||
with IMAGE_SAVE_LOCK:
|
||||
root = resolve_output_root(manifest.get("output_dir", ""), output_root)
|
||||
planned = plan_destinations(manifest, count, now or datetime.now(), root, output_root, picks=todo)
|
||||
actions = settle_destinations(planned, has_caption, output_root)
|
||||
resolved_root = output_root.resolve()
|
||||
for (path, _policy), action, index in zip(planned, actions, todo):
|
||||
if action == "skip":
|
||||
skipped.append(index)
|
||||
continue
|
||||
source = folder / _held_name(token, index, manifest["extension"])
|
||||
if not source.is_file():
|
||||
raise LookupError(GONE)
|
||||
path.parent.mkdir(parents=True, exist_ok=True)
|
||||
shutil.copyfile(source, path)
|
||||
if has_caption:
|
||||
sidecar_path(path).write_text(caption, encoding="utf-8")
|
||||
kept[index] = path.resolve().relative_to(resolved_root).as_posix()
|
||||
saved.append(index)
|
||||
manifest["kept"] = {str(index): path for index, path in kept.items()}
|
||||
manifest_path.write_text(json.dumps(manifest), encoding="utf-8")
|
||||
return {
|
||||
"count": count,
|
||||
"saved": [{"index": index, "path": kept[index]} for index in saved],
|
||||
"kept": sorted(kept),
|
||||
"skipped": skipped,
|
||||
}
|
||||
@@ -41,6 +41,16 @@ RADIUS_MAX = 4096.0
|
||||
# The face's own picture: the long side is at most this.
|
||||
SMALL_SIDE = 1024
|
||||
|
||||
# SAM 3 is asked about a copy of a big picture: its long side is at most
|
||||
# this. ComfyUI's SAM 3 reads every picture at 1008 pixels and returns each
|
||||
# answer at the size of the picture it was given, on the graphics card. At
|
||||
# 8192 pixels one word with sixteen matches needed over 27 GB there, and a
|
||||
# look at 3840 took twice as long as at 2048, all for outlines the model
|
||||
# never saw that finely. Twice what the model reads keeps its second pass
|
||||
# (a close-up of each thing it found) as sharp as before on pictures up to
|
||||
# this size, which come out exactly as they did.
|
||||
WORK_SIDE = 2048
|
||||
|
||||
# How a find stroke is asked (find_points, choose_found). Measured on 27
|
||||
# strokes over three pictures, each drawn 13 slightly different ways: four
|
||||
# points alone returned the stroked thing in 87 % of 351 tries, two points
|
||||
@@ -386,6 +396,73 @@ def read_ask(ask, height: int, width: int) -> tuple:
|
||||
return "look", None
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# The copy of a big picture that SAM 3 is asked about
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class WorkView:
|
||||
"""One picture as SAM 3 is asked about it.
|
||||
|
||||
A picture whose long side is over WORK_SIDE is asked about on a smaller
|
||||
copy (`frame`, averaged down, so nothing shimmers), and what SAM 3
|
||||
answers there is brought back to the picture's own size with `full`.
|
||||
A smaller picture is its own copy: `same` is true and nothing changes.
|
||||
|
||||
Touches and pieces keep their places in pixels of the picture. `point`
|
||||
gives a place on the copy, `full_point` a place on the picture that
|
||||
`point` leads back to the same pixel of the copy.
|
||||
"""
|
||||
|
||||
def __init__(self, frame: torch.Tensor):
|
||||
self.height, self.width = int(frame.shape[1]), int(frame.shape[2])
|
||||
long_side = max(self.height, self.width)
|
||||
self.same = long_side <= WORK_SIDE
|
||||
if self.same:
|
||||
self.frame, self.work_height, self.work_width = frame, self.height, self.width
|
||||
return
|
||||
scale = WORK_SIDE / long_side
|
||||
self.work_height = max(1, int(math.floor(self.height * scale + 0.5)))
|
||||
self.work_width = max(1, int(math.floor(self.width * scale + 0.5)))
|
||||
small = F.interpolate(
|
||||
frame.detach().float().cpu().movedim(-1, 1), size=(self.work_height, self.work_width), mode="area",
|
||||
)
|
||||
self.frame = small.movedim(1, -1).clamp(0.0, 1.0).contiguous()
|
||||
|
||||
def point(self, point) -> tuple:
|
||||
x = min(self.work_width - 1, max(0, int(point[0]) * self.work_width // self.width))
|
||||
y = min(self.work_height - 1, max(0, int(point[1]) * self.work_height // self.height))
|
||||
return (x, y)
|
||||
|
||||
def full_point(self, point) -> tuple:
|
||||
x = min(self.width - 1, -(-int(point[0]) * self.width // self.work_width))
|
||||
y = min(self.height - 1, -(-int(point[1]) * self.height // self.work_height))
|
||||
return (x, y)
|
||||
|
||||
def points(self, points) -> tuple:
|
||||
return tuple(self.point(point) for point in points)
|
||||
|
||||
def length(self, value: float) -> float:
|
||||
"""A length of the picture (a brush's radius) on the copy."""
|
||||
return float(value) * self.work_width / self.width
|
||||
|
||||
def full_box(self, box) -> list:
|
||||
if self.same:
|
||||
return list(box)
|
||||
x0, y0, x1, y1 = (int(value) for value in box)
|
||||
return [
|
||||
x0 * self.width // self.work_width, y0 * self.height // self.work_height,
|
||||
min(self.width, -(-x1 * self.width // self.work_width)), min(self.height, -(-y1 * self.height // self.work_height)),
|
||||
]
|
||||
|
||||
def full(self, mask: torch.Tensor) -> torch.Tensor:
|
||||
"""An answer from the copy at the picture's own size (bool)."""
|
||||
if self.same or tuple(mask.shape) == (self.height, self.width):
|
||||
return mask
|
||||
soft = F.interpolate(mask[None, None].float(), size=(self.height, self.width), mode="bilinear", align_corners=False)
|
||||
return soft[0, 0] >= 0.5
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Strokes
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
@@ -0,0 +1,35 @@
|
||||
"""Pick one choice from each {a|b|c} set in a prompt, decided by a seed.
|
||||
|
||||
Only a set that holds a ``|`` is a choice. Braces without one are left as
|
||||
they are, so a JSON prompt passes through unchanged. Sets can sit inside
|
||||
other sets; the innermost ones are picked first.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import random
|
||||
import re
|
||||
|
||||
# An innermost {...} that holds at least one "|", with the one space in
|
||||
# front of it, so an empty pick does not leave two spaces behind.
|
||||
_CHOICE = re.compile(r"( ?)\{([^{}]*\|[^{}]*)\}")
|
||||
|
||||
|
||||
def pick_variations(text: str, seed: int) -> tuple[str, list[str]]:
|
||||
"""Return the prompt with every choice set replaced, and the picks in order."""
|
||||
rng = random.Random(int(seed))
|
||||
picks: list[str] = []
|
||||
|
||||
def choose(match: re.Match) -> str:
|
||||
options = [option.strip() for option in match.group(2).split("|")]
|
||||
pick = rng.choice(options)
|
||||
picks.append(pick)
|
||||
return match.group(1) + pick if pick else ""
|
||||
|
||||
while True:
|
||||
text, count = _CHOICE.subn(choose, text)
|
||||
if not count:
|
||||
return text, picks
|
||||
|
||||
|
||||
__all__ = ["pick_variations"]
|
||||
@@ -0,0 +1,97 @@
|
||||
"""Only one copy of the pack loads.
|
||||
|
||||
A Manager install lands in ``custom_nodes/ausboss-nodes`` and a git clone in
|
||||
``custom_nodes/ComfyUI-AusBoss``. With both in place every route is added
|
||||
twice and ComfyUI stops while starting. So each copy looks for the others
|
||||
and all of them agree on the one that loads; the rest step aside and say
|
||||
which folder to delete.
|
||||
|
||||
A copy from before this file existed cannot step aside, so it always loads
|
||||
and a newer copy yields to it. Among copies that can, the newest version
|
||||
loads.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
import tomllib
|
||||
|
||||
PACK_NAME = "ausboss-nodes"
|
||||
# A copy that has this file knows how to step aside.
|
||||
_GUARD_FILE = os.path.join("nodes", "_single_copy.py")
|
||||
|
||||
|
||||
def _read_pack(folder: str) -> dict | None:
|
||||
"""Name the pack copy in ``folder``, or None when it is another pack."""
|
||||
try:
|
||||
with open(os.path.join(folder, "pyproject.toml"), "rb") as f:
|
||||
project = tomllib.load(f).get("project", {})
|
||||
except Exception:
|
||||
return None
|
||||
if project.get("name") != PACK_NAME:
|
||||
return None
|
||||
if not os.path.isfile(os.path.join(folder, "__init__.py")):
|
||||
return None
|
||||
version = project.get("version")
|
||||
return {
|
||||
"folder": folder,
|
||||
"version": version if isinstance(version, str) else "unknown",
|
||||
"guarded": os.path.isfile(os.path.join(folder, _GUARD_FILE)),
|
||||
}
|
||||
|
||||
|
||||
def _version_key(version: str) -> tuple:
|
||||
parts = []
|
||||
for part in version.split("."):
|
||||
digits = "".join(ch for ch in part if ch.isdigit())
|
||||
parts.append(int(digits) if digits else -1)
|
||||
return tuple(parts)
|
||||
|
||||
|
||||
def find_copies(own_folder: str, roots: list[str]) -> list[dict]:
|
||||
"""Every copy of the pack below ``roots``, this one included."""
|
||||
own = os.path.realpath(own_folder)
|
||||
seen = {own}
|
||||
copies = [_read_pack(own) or {"folder": own, "version": "unknown", "guarded": True}]
|
||||
for root in roots:
|
||||
try:
|
||||
names = sorted(os.listdir(root))
|
||||
except OSError:
|
||||
continue
|
||||
for name in names:
|
||||
# ComfyUI skips these too.
|
||||
if name.endswith(".disabled") or name.startswith("."):
|
||||
continue
|
||||
folder = os.path.realpath(os.path.join(root, name))
|
||||
if folder in seen or not os.path.isdir(folder):
|
||||
continue
|
||||
seen.add(folder)
|
||||
pack = _read_pack(folder)
|
||||
if pack:
|
||||
copies.append(pack)
|
||||
return copies
|
||||
|
||||
|
||||
def pick_copy(copies: list[dict]) -> dict:
|
||||
"""The one copy that loads. Every copy reaches the same answer."""
|
||||
unguarded = [c for c in copies if not c["guarded"]]
|
||||
pool = unguarded or copies
|
||||
return max(pool, key=lambda c: (_version_key(c["version"]), c["folder"]))
|
||||
|
||||
|
||||
def other_copy_that_loads(own_folder: str, roots: list[str] | None = None) -> dict | None:
|
||||
"""The copy this one steps aside for, or None when this one loads."""
|
||||
if roots is None:
|
||||
# Outside ComfyUI nothing is served, so there is nothing to clash.
|
||||
try:
|
||||
import folder_paths
|
||||
|
||||
roots = list(folder_paths.get_folder_paths("custom_nodes"))
|
||||
except Exception:
|
||||
return None
|
||||
roots = list(dict.fromkeys(os.path.realpath(r) for r in roots))
|
||||
copies = find_copies(own_folder, roots)
|
||||
if len(copies) < 2:
|
||||
return None
|
||||
chosen = pick_copy(copies)
|
||||
return None if chosen["folder"] == os.path.realpath(own_folder) else chosen
|
||||
@@ -0,0 +1,67 @@
|
||||
"""Cut a block of text into the items Text Lines runs one at a time.
|
||||
|
||||
Pure text handling with no ComfyUI imports, so it runs under the
|
||||
standard-library tests.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import re
|
||||
|
||||
SPLIT_EACH_LINE = "each line"
|
||||
SPLIT_BLANK_LINE = "blank line"
|
||||
SPLIT_MODES = (SPLIT_EACH_LINE, SPLIT_BLANK_LINE)
|
||||
|
||||
# The words a list writer puts before an item's number: "Prompt 3:", "Shot 2 -".
|
||||
# Not "image" or "picture": an edit prompt may open with "Image 1: ...".
|
||||
_LABEL = r"(?:prompt|shot|scene|line|step|item)"
|
||||
|
||||
# A marker at the start of an item that only counts or bullets it. A bare
|
||||
# number needs punctuation after it, and that punctuation must not be followed
|
||||
# by a digit, so "2 women on a beach", "1.5x zoom" and "10:30 at night" stay
|
||||
# whole. A bullet needs a space after it, so "*bold* text" stays whole.
|
||||
_MARKER = re.compile(
|
||||
r"""^\s*(?:
|
||||
[\[【(]\s*(?:%(label)s[\s_#-]*)?\d+\s*[\]】)]\s*[:.\-–—]? # [prompt 1] (2)
|
||||
| %(label)s[\s_#-]*\d+\s*[:.)\-–—] # Prompt 1: Shot 2 -
|
||||
| \#?\d+\s*[.):](?!\d) # 1. 2) #3:
|
||||
| [-*•·–—](?=\s|$) # - bullet
|
||||
)\s*"""
|
||||
% {"label": _LABEL},
|
||||
re.IGNORECASE | re.VERBOSE,
|
||||
)
|
||||
|
||||
_BLANK_LINE = re.compile(r"\n[ \t]*(?:\n[ \t]*)+")
|
||||
|
||||
|
||||
def strip_marker(item: str) -> str:
|
||||
"""Take one leading number or bullet off an item."""
|
||||
return _MARKER.sub("", item, count=1).strip()
|
||||
|
||||
|
||||
def split_items(text, split=SPLIT_EACH_LINE, remove_numbering=True, limit=0) -> list[str]:
|
||||
"""The items of `text`, in order, with empty ones left out.
|
||||
|
||||
`split` picks where one item ends: at every line break, or only at an
|
||||
empty line (so an item can hold several lines). `remove_numbering` takes
|
||||
a leading "1." or "-" off each item. `limit` keeps the first so many;
|
||||
0 keeps them all.
|
||||
"""
|
||||
text = str(text or "").replace("\r\n", "\n").replace("\r", "\n")
|
||||
if split == SPLIT_BLANK_LINE:
|
||||
parts = _BLANK_LINE.split(text)
|
||||
else:
|
||||
parts = text.split("\n")
|
||||
|
||||
items = []
|
||||
for part in parts:
|
||||
item = part.strip()
|
||||
if remove_numbering:
|
||||
item = strip_marker(item)
|
||||
if item:
|
||||
items.append(item)
|
||||
|
||||
limit = int(limit or 0)
|
||||
if limit > 0:
|
||||
items = items[:limit]
|
||||
return items
|
||||
@@ -0,0 +1,213 @@
|
||||
"""Edit Lock 🆎 and Edit Lock Restore 🆎: an edit leaves the rest of the picture alone.
|
||||
|
||||
An edit model redraws the whole picture, so the parts nobody asked to change
|
||||
come back a little moved and a little tinted. Edit Lock watches the model's
|
||||
own prediction at every step, reads where it differs from the original, and
|
||||
puts the original back everywhere else, so the model finishes the edit on
|
||||
top of the real picture. Edit Lock Restore then swaps the original's own
|
||||
pixels in for those parts after the decode. The deciding and blending live
|
||||
in _edit_lock_helpers.py.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import torch
|
||||
|
||||
from ._edit_lock_helpers import (
|
||||
LOCK_MODES,
|
||||
EditLock,
|
||||
describe,
|
||||
restore_pixels,
|
||||
)
|
||||
|
||||
|
||||
class AusBossEditLock:
|
||||
CATEGORY = "🆎 AusBoss/Image"
|
||||
DESCRIPTION = (
|
||||
"Keeps every part of your picture that an edit did not change "
|
||||
"exactly where it was, with the colors it had. Edit models redraw "
|
||||
"the whole picture, so untouched parts often come back slightly "
|
||||
"zoomed, shifted or tinted. Put this between the model and the "
|
||||
"sampler and give it the picture you are editing. It needs no "
|
||||
"mask: it sees what the edit changes while the picture is made. "
|
||||
"When the edit changes the whole picture (a new style, new light, "
|
||||
"new weather) it steps aside. Add Edit Lock Restore 🆎 after the "
|
||||
"decode to get your own pixels back in the untouched parts."
|
||||
)
|
||||
SEARCH_ALIASES = [
|
||||
"edit lock",
|
||||
"pixel drift",
|
||||
"pixel shift",
|
||||
"color shift",
|
||||
"consistency",
|
||||
"keep unchanged",
|
||||
"preserve original",
|
||||
"tail steer",
|
||||
"no drift edit",
|
||||
"ausboss",
|
||||
]
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"model": (
|
||||
"MODEL",
|
||||
{"tooltip": "The edit model. Patch last, after any LoRAs."},
|
||||
),
|
||||
"vae": (
|
||||
"VAE",
|
||||
{"tooltip": "The same VAE the workflow decodes with."},
|
||||
),
|
||||
"image": (
|
||||
"IMAGE",
|
||||
{
|
||||
"tooltip": (
|
||||
"The picture you are editing, at the size of the "
|
||||
"result. Use the same picture the edit model gets."
|
||||
),
|
||||
},
|
||||
),
|
||||
"lock": (
|
||||
list(LOCK_MODES),
|
||||
{
|
||||
"default": LOCK_MODES[0],
|
||||
"tooltip": (
|
||||
"auto locks the untouched parts and steps aside "
|
||||
"when the whole picture changes. always locks "
|
||||
"even then. off does nothing."
|
||||
),
|
||||
},
|
||||
),
|
||||
"tolerance": (
|
||||
"FLOAT",
|
||||
{
|
||||
"default": 1.0,
|
||||
"min": 0.25,
|
||||
"max": 4.0,
|
||||
"step": 0.05,
|
||||
"tooltip": (
|
||||
"How big a change must be to count as part of the "
|
||||
"edit. Raise it if the lock lets small unwanted "
|
||||
"changes through; lower it if it holds back part "
|
||||
"of your edit."
|
||||
),
|
||||
},
|
||||
),
|
||||
},
|
||||
"optional": {
|
||||
"mask": (
|
||||
"MASK",
|
||||
{
|
||||
"tooltip": (
|
||||
"Optional. White marks where the edit may happen; "
|
||||
"everything else is locked. With a mask the node "
|
||||
"does not look for the change itself."
|
||||
),
|
||||
},
|
||||
),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("MODEL", "AUSBOSS_EDIT_LOCK")
|
||||
RETURN_NAMES = ("model", "edit_lock")
|
||||
OUTPUT_TOOLTIPS = (
|
||||
"The model with the lock on. Plug it into the sampler.",
|
||||
"What the lock kept. Plug it into Edit Lock Restore 🆎 after the decode.",
|
||||
)
|
||||
FUNCTION = "patch"
|
||||
|
||||
def patch(self, model, vae, image, lock, tolerance, mask=None):
|
||||
if image.ndim != 4:
|
||||
raise ValueError("Edit Lock: image must be an IMAGE batch.")
|
||||
state = EditLock(mode=str(lock), tolerance=float(tolerance))
|
||||
source = image[:1, ..., :3]
|
||||
state.source = source.detach().float().cpu()
|
||||
if mask is not None:
|
||||
state.user_mask = mask.detach().float().cpu()
|
||||
if state.mode == "off":
|
||||
state.status = "off"
|
||||
return (model, state)
|
||||
latent = vae.encode(source)
|
||||
patched = model.clone()
|
||||
state.latent = patched.model.process_latent_in(latent).detach().float().cpu()
|
||||
patched.set_model_sampler_post_cfg_function(state.step)
|
||||
return (patched, state)
|
||||
|
||||
|
||||
class AusBossEditLockRestore:
|
||||
CATEGORY = "🆎 AusBoss/Image"
|
||||
DESCRIPTION = (
|
||||
"Puts your picture's own pixels back in every part the edit did not "
|
||||
"change. Use it right after the VAE decode, with the edit_lock "
|
||||
"output of Edit Lock 🆎. The untouched parts then match your "
|
||||
"original pixel for pixel; only the edit is new. When the lock "
|
||||
"stepped aside the picture passes through unchanged."
|
||||
)
|
||||
SEARCH_ALIASES = [
|
||||
"edit lock restore",
|
||||
"restore original pixels",
|
||||
"paste back",
|
||||
"keep unchanged",
|
||||
"ausboss",
|
||||
]
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"image": (
|
||||
"IMAGE",
|
||||
{"tooltip": "The edited picture, straight from the VAE decode."},
|
||||
),
|
||||
"edit_lock": (
|
||||
"AUSBOSS_EDIT_LOCK",
|
||||
{"tooltip": "From Edit Lock 🆎."},
|
||||
),
|
||||
"feather": (
|
||||
"INT",
|
||||
{
|
||||
"default": 16,
|
||||
"min": 0,
|
||||
"max": 256,
|
||||
"step": 1,
|
||||
"tooltip": (
|
||||
"How softly your original fades into the edit, in "
|
||||
"pixels. The fade stays outside the edit."
|
||||
),
|
||||
},
|
||||
),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("IMAGE", "MASK", "STRING")
|
||||
RETURN_NAMES = ("image", "edit_mask", "report")
|
||||
OUTPUT_TOOLTIPS = (
|
||||
"The edit with your original pixels in the untouched parts.",
|
||||
"White where the edit changed the picture.",
|
||||
"How much of the picture was kept. Plug it into Show Text 🆎 to read it.",
|
||||
)
|
||||
FUNCTION = "restore"
|
||||
|
||||
def restore(self, image, edit_lock, feather):
|
||||
if image.ndim != 4:
|
||||
raise ValueError("Edit Lock Restore: image must be an IMAGE batch.")
|
||||
if not isinstance(edit_lock, EditLock):
|
||||
raise ValueError("Edit Lock Restore: edit_lock needs the output of Edit Lock.")
|
||||
height, width = int(image.shape[1]), int(image.shape[2])
|
||||
mask = edit_lock.pixel_mask(height, width)
|
||||
if mask is None or edit_lock.source is None:
|
||||
return (image, torch.ones((int(image.shape[0]), height, width)), describe(edit_lock))
|
||||
result, edit_mask = restore_pixels(image, edit_lock.source, mask, int(feather))
|
||||
return (result, edit_mask, describe(edit_lock))
|
||||
|
||||
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"AUSBOSS_NODES_EditLock": AusBossEditLock,
|
||||
"AUSBOSS_NODES_EditLockRestore": AusBossEditLockRestore,
|
||||
}
|
||||
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"AUSBOSS_NODES_EditLock": "Edit Lock 🆎",
|
||||
"AUSBOSS_NODES_EditLockRestore": "Edit Lock Restore 🆎",
|
||||
}
|
||||
@@ -0,0 +1,87 @@
|
||||
"""FlashVSR Upscale 🆎."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from ._flashvsr_helpers import EXPECTED_MODEL_NAME, list_flashvsr_models, resolve_flashvsr_model, run_flashvsr
|
||||
|
||||
|
||||
class AusBossFlashVSRUpscale:
|
||||
CATEGORY = "🆎 AusBoss/Video"
|
||||
DESCRIPTION = (
|
||||
"Makes a video larger and sharper with the FlashVSR model. Pick the "
|
||||
"model file, set how much larger, and run. It reads the clip a few "
|
||||
"frames at a time, so a long clip needs no more graphics memory than "
|
||||
"a short one, and every frame comes back. Connect the Wan 2.1 VAE "
|
||||
"for a slower, finer finish."
|
||||
)
|
||||
SEARCH_ALIASES = [
|
||||
"flashvsr",
|
||||
"video upscale",
|
||||
"upscale video",
|
||||
"video super resolution",
|
||||
"vsr",
|
||||
"sharpen video",
|
||||
"ausboss",
|
||||
]
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
models = list_flashvsr_models()
|
||||
return {
|
||||
"required": {
|
||||
"frames": (
|
||||
"IMAGE",
|
||||
{"tooltip": "The video frames to make larger."},
|
||||
),
|
||||
"model": (
|
||||
models,
|
||||
{
|
||||
"default": EXPECTED_MODEL_NAME if EXPECTED_MODEL_NAME in models else models[0],
|
||||
"tooltip": "The FlashVSR model file, from ComfyUI/models/diffusion_models.",
|
||||
},
|
||||
),
|
||||
"scale": (
|
||||
"FLOAT",
|
||||
{
|
||||
"default": 2.0,
|
||||
"min": 1.0,
|
||||
"max": 4.0,
|
||||
"step": 0.25,
|
||||
"tooltip": "How many times larger each side gets. 1 keeps the size and only sharpens.",
|
||||
},
|
||||
),
|
||||
"color_match": (
|
||||
"BOOLEAN",
|
||||
{
|
||||
"default": True,
|
||||
"tooltip": "Keeps each frame's colours the same as your video. Leave it on.",
|
||||
},
|
||||
),
|
||||
},
|
||||
"optional": {
|
||||
"vae": (
|
||||
"VAE",
|
||||
{
|
||||
"tooltip": (
|
||||
"Optional. Connect the Wan 2.1 VAE for a slower, finer finish. "
|
||||
"Left empty, the fast decoder inside the model file is used."
|
||||
)
|
||||
},
|
||||
),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
RETURN_NAMES = ("frames",)
|
||||
OUTPUT_TOOLTIPS = ("The larger, sharper frames: as many as went in.",)
|
||||
FUNCTION = "upscale"
|
||||
|
||||
def upscale(self, frames, model, scale, color_match, vae=None):
|
||||
output = run_flashvsr(resolve_flashvsr_model(model), frames, float(scale), bool(color_match), vae)
|
||||
return (output,)
|
||||
|
||||
|
||||
NODE_CLASS_MAPPINGS = {"AUSBOSS_NODES_FlashVSRUpscale": AusBossFlashVSRUpscale}
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {"AUSBOSS_NODES_FlashVSRUpscale": "FlashVSR Upscale 🆎"}
|
||||
|
||||
__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"]
|
||||
+65
-30
@@ -42,6 +42,7 @@ from ._mask_by_touch_helpers import (
|
||||
WORDS_MORE,
|
||||
Memory,
|
||||
PieceSieve,
|
||||
WorkView,
|
||||
apply_item,
|
||||
box_of,
|
||||
choose_found,
|
||||
@@ -126,6 +127,8 @@ QUICK = True
|
||||
|
||||
# Shared by every Mask by Touch node: what SAM 3 answered, per picture.
|
||||
_MEMORY = Memory()
|
||||
# How many of those smaller copies a node keeps: this picture's and the last one's.
|
||||
VIEWS_KEPT = 2
|
||||
# What the text reader made of each word, per text reader. A word reads the
|
||||
# same for every picture, so it is read once.
|
||||
_WORDS = OrderedDict()
|
||||
@@ -836,25 +839,44 @@ class AusBossMaskByTouch:
|
||||
# What Fine edges is set to in the run at hand: the subject is asked
|
||||
# of that model.
|
||||
self._edge_choice = FINE_OFF
|
||||
# The smaller copies of big pictures SAM 3 is asked about, newest last.
|
||||
self._views = OrderedDict()
|
||||
|
||||
# -- what SAM 3 finds ---------------------------------------------------
|
||||
|
||||
def _tap(self, session: _Session, frame, frame_key: str, point):
|
||||
"""The thing at one point, without its stray specks. None for nothing.
|
||||
def _view(self, frame, frame_key: str) -> WorkView:
|
||||
"""The picture as SAM 3 is asked about it: itself, or a smaller copy
|
||||
of a big one, kept for the next run on the same picture."""
|
||||
held = self._views.get(frame_key)
|
||||
if held is not None and (held.height, held.width) == (int(frame.shape[1]), int(frame.shape[2])):
|
||||
self._views.move_to_end(frame_key)
|
||||
return held
|
||||
view = WorkView(frame)
|
||||
if not view.same:
|
||||
self._views[frame_key] = view
|
||||
while len(self._views) > VIEWS_KEPT:
|
||||
self._views.popitem(last=False)
|
||||
return view
|
||||
|
||||
def _tap(self, session: _Session, view: WorkView, frame_key: str, point):
|
||||
"""The thing at one point, without its stray specks, at the size SAM 3
|
||||
was asked at (`view`). None for nothing. `point` is a pixel of the
|
||||
picture.
|
||||
|
||||
Hovering's pieces and taps share this, and what it returned is
|
||||
remembered, so a tap on a piece that was shown gives that piece."""
|
||||
key = (session.token, frame_key, ("tap", int(point[0]), int(point[1])))
|
||||
found, answer = _MEMORY.get(key)
|
||||
if not found:
|
||||
answer = drop_specks(session.ask(frame, [point]), point)
|
||||
spot = view.point(point)
|
||||
answer = drop_specks(session.ask(view.frame, [spot]), spot)
|
||||
answer = answer if bool(answer.any()) else None
|
||||
_MEMORY.put(key, answer)
|
||||
return answer
|
||||
|
||||
def _named(self, session: _Session, frame, frame_key: str, word: str) -> list:
|
||||
def _named(self, session: _Session, view: WorkView, frame_key: str, word: str) -> list:
|
||||
"""Everything in one picture that SAM 3 calls by a word, tidied:
|
||||
[(H x W bool, how sure), ...].
|
||||
[(bool at the size SAM 3 was asked at, how sure), ...].
|
||||
|
||||
Hovering's named pieces and `name` touches share this, and what it
|
||||
returned is remembered, so a touch on a named piece that was shown
|
||||
@@ -862,7 +884,7 @@ class AusBossMaskByTouch:
|
||||
key = (session.token, session.reader, frame_key, ("word", word))
|
||||
found, matches = _MEMORY.get_matches(key)
|
||||
if not found:
|
||||
matches = tidy_matches(*session.ask_word(frame, word))
|
||||
matches = tidy_matches(*session.ask_word(view.frame, word))
|
||||
_MEMORY.put_matches(key, matches)
|
||||
return matches
|
||||
|
||||
@@ -873,6 +895,9 @@ class AusBossMaskByTouch:
|
||||
The area of a `paint` touch made with a soft brush is float, 0..1.
|
||||
"""
|
||||
height, width = int(frame.shape[1]), int(frame.shape[2])
|
||||
# What SAM 3 is asked comes back at the size it was asked at, and is
|
||||
# brought to the picture's own size here.
|
||||
view = self._view(frame, frame_key)
|
||||
if item.kind in ("subject", "background"):
|
||||
# The picture's main subject, as ComfyUI's background-removal
|
||||
# model sees it, or everything but it.
|
||||
@@ -881,25 +906,25 @@ class AusBossMaskByTouch:
|
||||
return (area, "thing", ()) if bool(area.any()) else (None, "none", ())
|
||||
if item.kind == "all":
|
||||
# Every thing called by that word, joined.
|
||||
matches = self._named(session, frame, frame_key, item.word) if item.word else []
|
||||
matches = self._named(session, view, frame_key, item.word) if item.word else []
|
||||
if not matches:
|
||||
return None, "none", ()
|
||||
area = torch.zeros(height, width, dtype=torch.bool)
|
||||
area = torch.zeros(view.work_height, view.work_width, dtype=torch.bool)
|
||||
for mask, _ in matches:
|
||||
area |= mask
|
||||
return area, "thing", tuple(range(len(matches)))
|
||||
return view.full(area), "thing", tuple(range(len(matches)))
|
||||
if item.kind in ("tap", "name"):
|
||||
if not (0 <= item.x < width and 0 <= item.y < height):
|
||||
return None, "none", ()
|
||||
if item.kind == "tap":
|
||||
answer = self._tap(session, frame, frame_key, (item.x, item.y))
|
||||
return (answer, "thing", ()) if answer is not None else (None, "none", ())
|
||||
answer = self._tap(session, view, frame_key, (item.x, item.y))
|
||||
return (view.full(answer), "thing", ()) if answer is not None else (None, "none", ())
|
||||
if not item.word:
|
||||
return None, "none", ()
|
||||
# Of the things called by that word, the smallest at the point.
|
||||
matches = self._named(session, frame, frame_key, item.word)
|
||||
index = match_under(matches, (item.x, item.y))
|
||||
return (None, "none", ()) if index is None else (matches[index][0], "thing", (index,))
|
||||
matches = self._named(session, view, frame_key, item.word)
|
||||
index = match_under(matches, view.point((item.x, item.y)))
|
||||
return (None, "none", ()) if index is None else (view.full(matches[index][0]), "thing", (index,))
|
||||
if item.kind not in ("find", "paint"):
|
||||
return None, "none", ()
|
||||
if item.kind == "paint" and item.soft > 0:
|
||||
@@ -918,14 +943,18 @@ class AusBossMaskByTouch:
|
||||
found, answer = _MEMORY.get(key)
|
||||
if not found:
|
||||
many, pair = find_points(item.pts, height, width)
|
||||
# The stroke as it lies on the copy SAM 3 is asked about.
|
||||
under = stroke if view.same else stroke_area(
|
||||
view.points(item.pts), max(1.0, view.length(item.r)), view.work_height, view.work_width,
|
||||
)
|
||||
answer = choose_found(
|
||||
session.ask(frame, many) if many else None,
|
||||
session.ask(frame, pair) if pair else None,
|
||||
stroke,
|
||||
session.ask(view.frame, view.points(many)) if many else None,
|
||||
session.ask(view.frame, view.points(pair)) if pair else None,
|
||||
under,
|
||||
)
|
||||
_MEMORY.put(key, answer)
|
||||
# Over empty space the paint is kept as it is.
|
||||
return (answer, "thing", ()) if answer is not None else (stroke, "paint", ())
|
||||
return (view.full(answer), "thing", ()) if answer is not None else (stroke, "paint", ())
|
||||
|
||||
def _look(self, session: _Session, frame, key: str, pieces: dict, pointer=None, tell=None) -> dict:
|
||||
"""Find a picture's pieces for hovering.
|
||||
@@ -943,6 +972,7 @@ class AusBossMaskByTouch:
|
||||
piece.
|
||||
"""
|
||||
height, width = int(frame.shape[1]), int(frame.shape[2])
|
||||
view = self._view(frame, key)
|
||||
queue = _scope_queue(width, height)
|
||||
sieve = pieces["sieve"]
|
||||
fresh = {} # (name or None, point) -> (area, is it tidy and remembered)
|
||||
@@ -967,7 +997,7 @@ class AusBossMaskByTouch:
|
||||
if pointer is not None:
|
||||
since = time.perf_counter()
|
||||
known = _MEMORY.has((session.token, key, ("tap", pointer[0], pointer[1])))
|
||||
answer = self._tap(session, frame, key, pointer)
|
||||
answer = self._tap(session, view, key, pointer)
|
||||
answered(known, since)
|
||||
if answer is not None and sieve.offer(pointer, answer, pointed=True):
|
||||
fresh[(None, pointer)] = (answer, True)
|
||||
@@ -993,8 +1023,10 @@ class AusBossMaskByTouch:
|
||||
break
|
||||
ahead.append(later)
|
||||
session.read_ahead(ahead, stop=_jobs_waiting)
|
||||
matches = self._named(session, frame, key, what)
|
||||
for (area, sure), point in zip(matches, match_points(matches)):
|
||||
matches = self._named(session, view, key, what)
|
||||
for (area, sure), spot in zip(matches, match_points(matches)):
|
||||
# The match's own point, as a pixel of the picture.
|
||||
point = None if spot is None else view.full_point(spot)
|
||||
if point is not None and sieve.offer(point, area, name=what, sure=sure):
|
||||
fresh[(what, point)] = (area, True)
|
||||
pieces["words"] += 1
|
||||
@@ -1002,7 +1034,7 @@ class AusBossMaskByTouch:
|
||||
memory_key = (session.token, key, ("tap", what[0], what[1]))
|
||||
known, answer = _MEMORY.get(memory_key)
|
||||
if not known:
|
||||
answer = session.ask(frame, [what])
|
||||
answer = session.ask(view.frame, [view.point(what)])
|
||||
if not bool(answer.any()):
|
||||
_MEMORY.put(memory_key, None)
|
||||
answer = None
|
||||
@@ -1025,6 +1057,7 @@ class AusBossMaskByTouch:
|
||||
first, and the small copies packed into planes. `fresh` holds the
|
||||
areas found since the last time, and is emptied."""
|
||||
height, width = int(frame.shape[1]), int(frame.shape[2])
|
||||
view = self._view(frame, key)
|
||||
small_width, small_height, _ = small_size(width, height)
|
||||
sieve = pieces["sieve"]
|
||||
before = {ident: slot for slot, ident in enumerate(pieces["ids"])}
|
||||
@@ -1036,18 +1069,19 @@ class AusBossMaskByTouch:
|
||||
else:
|
||||
# A piece kept by a look that was stopped half way is asked
|
||||
# for again: its answer was not held on to.
|
||||
answer, tidy = fresh.get(ident) or (self._piece(session, frame, key, name, point), True)
|
||||
answer, tidy = fresh.get(ident) or (self._piece(session, view, key, name, point), True)
|
||||
if answer is None:
|
||||
continue
|
||||
if not tidy:
|
||||
# What a tap at this point gives: the same answer, tidied
|
||||
# the same way, remembered so the tap asks nothing.
|
||||
answer = drop_specks(answer, point)
|
||||
answer = drop_specks(answer, view.point(point))
|
||||
_MEMORY.put((session.token, key, ("tap", point[0], point[1])), answer)
|
||||
sure = sieve.sure_of(name, point)
|
||||
entry = {
|
||||
"x": point[0], "y": point[1],
|
||||
"share": round(float(answer.sum()) / (height * width), 5), "box": box_of(answer),
|
||||
"share": round(float(answer.sum()) / (view.work_height * view.work_width), 5),
|
||||
"box": view.full_box(box_of(answer)),
|
||||
"kind": "point" if name is None else "name", "name": name,
|
||||
"sure": None if sure is None else round(sure, 3),
|
||||
}
|
||||
@@ -1064,13 +1098,13 @@ class AusBossMaskByTouch:
|
||||
pieces["list"] = [row[4] for row in rows]
|
||||
pieces["planes"] = pack_planes(masks)
|
||||
|
||||
def _piece(self, session: _Session, frame, key: str, name, point):
|
||||
def _piece(self, session: _Session, view: WorkView, key: str, name, point):
|
||||
"""The area of a piece from what is remembered: a point's answer, or
|
||||
the match of a word that the piece's own point stands for."""
|
||||
if name is None:
|
||||
return self._tap(session, frame, key, point)
|
||||
matches = self._named(session, frame, key, name)
|
||||
index = match_under(matches, point)
|
||||
return self._tap(session, view, key, point)
|
||||
matches = self._named(session, view, key, name)
|
||||
index = match_under(matches, view.point(point))
|
||||
return None if index is None else matches[index][0]
|
||||
|
||||
# -- files for the node's face -------------------------------------------
|
||||
@@ -1385,8 +1419,9 @@ class AusBossMaskByTouch:
|
||||
if asked == "scope" and (pieces is None or pointer is not None or pieces["next"] < pieces["total"]):
|
||||
if pieces is None:
|
||||
words = len(WORDS_FIRST) + len(WORDS_MORE)
|
||||
view = self._view(image[:1], key)
|
||||
pieces = {
|
||||
"sieve": PieceSieve(height, width), "next": 0, "total": len(_scope_queue(width, height)),
|
||||
"sieve": PieceSieve(view.work_height, view.work_width), "next": 0, "total": len(_scope_queue(width, height)),
|
||||
"words": 0, "words_total": words, "number": 0, "ids": [], "list": [], "planes": [],
|
||||
}
|
||||
# Remembered from the start: a look that is stopped half way
|
||||
|
||||
@@ -0,0 +1,73 @@
|
||||
"""Prompt Variations 🆎 — one prompt, a different mix of choices on every run."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from ._prompt_variations_helpers import pick_variations
|
||||
|
||||
|
||||
class AusBossPromptVariations:
|
||||
CATEGORY = "🆎 AusBoss/Utility"
|
||||
DESCRIPTION = (
|
||||
"A prompt box that changes on every run. Write choices in curly "
|
||||
"brackets, such as a {red|blue|green} {sweater|jacket}, and the node "
|
||||
"picks one from each set. The seed decides the picks: the same seed "
|
||||
"gives the same prompt again, and a new seed on each run gives a new "
|
||||
"mix. Text without brackets comes out exactly as typed."
|
||||
)
|
||||
SEARCH_ALIASES = [
|
||||
"random prompt", "wildcard", "dynamic prompt", "variations",
|
||||
"random choice", "prompt mixer", "ausboss",
|
||||
]
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"text": (
|
||||
"STRING",
|
||||
{
|
||||
"default": "",
|
||||
"multiline": True,
|
||||
"dynamicPrompts": False,
|
||||
"tooltip": (
|
||||
"Your prompt. Put choices in curly brackets, split "
|
||||
"by |, such as {tank top|sweater|shirt}. One is "
|
||||
"picked from each set."
|
||||
),
|
||||
},
|
||||
),
|
||||
"seed": (
|
||||
"INT",
|
||||
{
|
||||
"default": 0,
|
||||
"min": 0,
|
||||
"max": 0xFFFFFFFFFFFFFFFF,
|
||||
"forceInput": True,
|
||||
"tooltip": (
|
||||
"Wire in the seed your sampler uses. A new seed "
|
||||
"picks new choices; the same seed picks the same "
|
||||
"ones again."
|
||||
),
|
||||
},
|
||||
),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("STRING",)
|
||||
RETURN_NAMES = ("prompt",)
|
||||
OUTPUT_TOOLTIPS = ("The prompt with one choice picked from each set.",)
|
||||
FUNCTION = "pick"
|
||||
|
||||
def pick(self, text, seed):
|
||||
prompt, _picks = pick_variations(text, seed)
|
||||
return (prompt,)
|
||||
|
||||
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"AUSBOSS_NODES_PromptVariations": AusBossPromptVariations,
|
||||
}
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"AUSBOSS_NODES_PromptVariations": "Prompt Variations 🆎",
|
||||
}
|
||||
|
||||
__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"]
|
||||
+78
-44
@@ -6,26 +6,20 @@ from datetime import datetime
|
||||
from pathlib import Path
|
||||
|
||||
from ._execution_helpers import ignore_extra_inputs
|
||||
from ._image_hold_helpers import hold_pictures, keep_pictures, plan_destinations, settle_destinations
|
||||
from ._image_save_helpers import (
|
||||
EXISTING_POLICIES,
|
||||
FORMAT_EXTENSIONS,
|
||||
IMAGE_FORMATS,
|
||||
IMAGE_SAVE_LOCK,
|
||||
encode_image,
|
||||
existing_action,
|
||||
fill_name_tokens,
|
||||
metadata_disabled,
|
||||
name_stem,
|
||||
next_free_counter,
|
||||
plan_exact_names,
|
||||
plan_local_names,
|
||||
resolve_output_root,
|
||||
require_output_path,
|
||||
sanitize_exact_name,
|
||||
sidecar_path,
|
||||
split_prefix,
|
||||
strip_image_extension,
|
||||
)
|
||||
from ._route_errors import safe_route_error
|
||||
from ._step_preview_helpers import STEP_PREVIEW_DEFAULT, register_step_preview_handler
|
||||
from ._video_save_helpers import workflow_metadata
|
||||
|
||||
@@ -68,6 +62,33 @@ def _register_folder_route() -> None:
|
||||
shown = "" if target == root else target.relative_to(root).as_posix()
|
||||
return web.json_response({"path": shown, "folders": folders})
|
||||
|
||||
@PromptServer.instance.routes.post("/ausboss/save_image/keep")
|
||||
async def keep_held(request):
|
||||
"""Save pictures a run is holding (Auto save off). The body names the
|
||||
run by the token the node was given and, optionally, which of its
|
||||
pictures to keep. Names and folders come from what the run stored,
|
||||
never from the request."""
|
||||
try:
|
||||
body = await request.json()
|
||||
except Exception:
|
||||
return web.json_response({"error": "bad request"}, status=400)
|
||||
if not isinstance(body, dict):
|
||||
return web.json_response({"error": "bad request"}, status=400)
|
||||
picks = body.get("picks")
|
||||
if picks is not None and not isinstance(picks, list):
|
||||
return web.json_response({"error": "bad request"}, status=400)
|
||||
try:
|
||||
answer = keep_pictures(
|
||||
folder_paths.get_temp_directory(), folder_paths.get_output_directory(), body.get("token"), picks,
|
||||
)
|
||||
except LookupError as exc:
|
||||
return web.json_response({"error": safe_route_error(exc, "Those pictures are gone.")}, status=404)
|
||||
except ValueError as exc:
|
||||
return web.json_response({"error": safe_route_error(exc, "Could not save.")}, status=400)
|
||||
except Exception as exc:
|
||||
return web.json_response({"error": safe_route_error(exc, "Could not save.")}, status=500)
|
||||
return web.json_response(answer)
|
||||
|
||||
|
||||
class AusBossSaveImage:
|
||||
CATEGORY = "🆎 AusBoss/Image"
|
||||
@@ -277,6 +298,17 @@ class AusBossSaveImage:
|
||||
),
|
||||
},
|
||||
),
|
||||
"autosave": (
|
||||
"BOOLEAN",
|
||||
{
|
||||
"default": True,
|
||||
"tooltip": (
|
||||
"On: every run saves its pictures. Off: the "
|
||||
"pictures only show on the node, and its Save "
|
||||
"button keeps the ones you want."
|
||||
),
|
||||
},
|
||||
),
|
||||
},
|
||||
"hidden": {"prompt": "PROMPT", "extra_pnginfo": "EXTRA_PNGINFO"},
|
||||
}
|
||||
@@ -310,6 +342,7 @@ class AusBossSaveImage:
|
||||
name_batch=False,
|
||||
preview=True,
|
||||
step_preview=STEP_PREVIEW_DEFAULT,
|
||||
autosave=True,
|
||||
prompt=None,
|
||||
extra_pnginfo=None,
|
||||
):
|
||||
@@ -343,45 +376,46 @@ class AusBossSaveImage:
|
||||
else:
|
||||
exact = sanitize_exact_name(exact_name)
|
||||
|
||||
with IMAGE_SAVE_LOCK:
|
||||
planned: list[tuple[Path, str]] = [] # (path, collision policy)
|
||||
if exact:
|
||||
base = strip_image_extension(exact)
|
||||
for name in plan_exact_names(base, extension, count):
|
||||
planned.append((root / name, on_existing))
|
||||
else:
|
||||
subfolder, base = split_prefix(filename_prefix)
|
||||
folder = root / subfolder if subfolder else root
|
||||
require_output_path(folder, output_root)
|
||||
stem = name_stem(
|
||||
base,
|
||||
{"date": bool(name_date), "time": bool(name_time), "size": bool(name_size)},
|
||||
now,
|
||||
width,
|
||||
height,
|
||||
)
|
||||
counter = bool(name_counter)
|
||||
existing = [entry.name for entry in folder.iterdir()] if folder.is_dir() else []
|
||||
names = plan_local_names(
|
||||
stem, extension, count, counter=counter, batch=bool(name_batch),
|
||||
next_counter=next_free_counter(existing, stem),
|
||||
)
|
||||
# The shared lock keeps this counter free until writing finishes.
|
||||
# Without it, the legacy policy decides whether to overwrite.
|
||||
policy = "overwrite" if counter else on_existing
|
||||
planned.extend((folder / name, policy) for name in names)
|
||||
settings = {
|
||||
"format": format,
|
||||
"extension": extension,
|
||||
"width": width,
|
||||
"height": height,
|
||||
"exact": exact,
|
||||
"filename_prefix": filename_prefix,
|
||||
"on_existing": on_existing,
|
||||
"output_dir": str(output_dir or ""),
|
||||
"name_counter": bool(name_counter),
|
||||
"name_date": bool(name_date),
|
||||
"name_time": bool(name_time),
|
||||
"name_size": bool(name_size),
|
||||
"name_batch": bool(name_batch),
|
||||
"caption": caption_value,
|
||||
}
|
||||
has_caption = bool(caption_value.strip())
|
||||
|
||||
if not autosave:
|
||||
# Nothing is saved yet. The pictures wait in the temp folder, in
|
||||
# their final format, and the node's Save button keeps the ones
|
||||
# that are wanted (the keep route above). A bad name or folder
|
||||
# still stops the run here, not at the click.
|
||||
with IMAGE_SAVE_LOCK:
|
||||
settle_destinations(plan_destinations(settings, count, now, root, output_root), has_caption, output_root)
|
||||
token, previews = hold_pictures(
|
||||
folder_paths.get_temp_directory(), images, settings,
|
||||
lambda path, frame: encode_image(path, frame, format, metadata),
|
||||
)
|
||||
return {
|
||||
"ui": {"images": previews, "ausboss_saved_path": [], "ausboss_held": [token]},
|
||||
"result": ("", images),
|
||||
}
|
||||
|
||||
with IMAGE_SAVE_LOCK:
|
||||
# Validate the whole batch before any write, including caption targets,
|
||||
# then settle every collision: an "error" on a later frame must stop
|
||||
# the run before the first frame lands, not after.
|
||||
for path, _policy in planned:
|
||||
require_output_path(path, output_root)
|
||||
if caption_value.strip():
|
||||
require_output_path(sidecar_path(path), output_root)
|
||||
actions = [
|
||||
existing_action(path.exists() or (bool(caption_value.strip()) and sidecar_path(path).exists()), policy)
|
||||
for path, policy in planned
|
||||
]
|
||||
planned = plan_destinations(settings, count, now, root, output_root)
|
||||
actions = settle_destinations(planned, has_caption, output_root)
|
||||
|
||||
saved: list[Path] = []
|
||||
for (path, _policy), action, frame in zip(planned, actions, images):
|
||||
@@ -389,7 +423,7 @@ class AusBossSaveImage:
|
||||
continue
|
||||
path.parent.mkdir(parents=True, exist_ok=True)
|
||||
encode_image(path, frame, format, metadata)
|
||||
if caption_value.strip():
|
||||
if has_caption:
|
||||
sidecar_path(path).write_text(caption_value, encoding="utf-8")
|
||||
saved.append(path)
|
||||
|
||||
|
||||
@@ -0,0 +1,100 @@
|
||||
"""Text Lines 🆎 — run what comes after once for each line of text."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from ._text_lines_helpers import SPLIT_EACH_LINE, SPLIT_MODES, split_items
|
||||
|
||||
MAX_LIMIT = 1000
|
||||
|
||||
|
||||
class AusBossTextLines:
|
||||
CATEGORY = "🆎 AusBoss/Utility"
|
||||
DESCRIPTION = (
|
||||
"Runs what comes after once for each line of your text. Type five "
|
||||
"prompts, one on each line, and one run makes five pictures. Empty "
|
||||
"lines are skipped, and a number or bullet at the start of a line "
|
||||
"can be taken off."
|
||||
)
|
||||
SEARCH_ALIASES = [
|
||||
"prompt list",
|
||||
"one prompt per line",
|
||||
"each line",
|
||||
"split text",
|
||||
"text to list",
|
||||
"batch prompts",
|
||||
"queue prompts",
|
||||
"ausboss",
|
||||
]
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"text": (
|
||||
"STRING",
|
||||
{
|
||||
"default": "",
|
||||
"multiline": True,
|
||||
"tooltip": "One line for each run. Type it here, or wire in text another node wrote.",
|
||||
},
|
||||
),
|
||||
"split": (
|
||||
list(SPLIT_MODES),
|
||||
{
|
||||
"default": SPLIT_EACH_LINE,
|
||||
"tooltip": (
|
||||
"each line: every line is its own run. blank line: lines stay "
|
||||
"together, and an empty line starts the next run."
|
||||
),
|
||||
},
|
||||
),
|
||||
"remove_numbering": (
|
||||
"BOOLEAN",
|
||||
{
|
||||
"default": True,
|
||||
"tooltip": (
|
||||
'On: a number or bullet at the start is taken off, like "1.", "-" '
|
||||
'or "[prompt 2]". Off: every line stays as written.'
|
||||
),
|
||||
},
|
||||
),
|
||||
"limit": (
|
||||
"INT",
|
||||
{
|
||||
"default": 0,
|
||||
"min": 0,
|
||||
"max": MAX_LIMIT,
|
||||
"step": 1,
|
||||
"tooltip": "The most lines to run. 0 runs them all.",
|
||||
},
|
||||
),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("STRING", "INT", "INT")
|
||||
RETURN_NAMES = ("line", "number", "count")
|
||||
OUTPUT_IS_LIST = (True, True, False)
|
||||
OUTPUT_TOOLTIPS = (
|
||||
"Each line, one at a time. What comes after runs once for each.",
|
||||
"The line's place in the list, starting at 1. Add it to a seed to give each run its own.",
|
||||
"How many lines run.",
|
||||
)
|
||||
FUNCTION = "split"
|
||||
|
||||
def split(self, text, split=SPLIT_EACH_LINE, remove_numbering=True, limit=0):
|
||||
items = split_items(text, split, remove_numbering, limit)
|
||||
if not items:
|
||||
raise ValueError(
|
||||
"Text Lines has no lines to run. Type at least one line, or check the node that writes its text."
|
||||
)
|
||||
return (items, list(range(1, len(items) + 1)), len(items))
|
||||
|
||||
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"AUSBOSS_NODES_TextLines": AusBossTextLines,
|
||||
}
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"AUSBOSS_NODES_TextLines": "Text Lines 🆎",
|
||||
}
|
||||
|
||||
__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"]
|
||||
+3
-2
@@ -2,8 +2,8 @@
|
||||
name = "ausboss-nodes"
|
||||
# Manager/registry search matches against this text, so it names the actual
|
||||
# nodes instead of describing the collection abstractly.
|
||||
description = "ComfyUI-AusBoss (AusBoss nodes): image and video nodes with compact controls: Image and Video Crop + Rotate + Pad (frame or clip), Load Image + Pad, Image Folder, Image Resize, Align Image, Image Size, Color Match, Image Compare A/B, Realign to Source (experimental), Save Image (PNG, WebP/JXL, captions), Load Video, Load Audio, Save Video, Select Frame, Frame Interpolate, Mask Refine, Mask by Name and Mask by Touch (type a name or click a thing, get its mask), LaMa Inpaint, Crop For Inpaint, Stitch Inpaint, Tiled Upscale and Tiled Upscale Stitch, LoRA Loader, Krea 2 Encode and Outpaint Model Patch, Latent Size, Layout Boxes, Seed, Select Every Nth, Split Batch, Merge Batches, Text, H3 Prompt Writer, Integer, Float, Math Expression, Show Text, Free Memory, Workflow Note, Workflow Note Tile, Callout, Callout Tile, Workflow Switches, Run Timer, and Step Preview (stop bad runs). Example workflows for Qwen Image 2.1, Krea 2, Klein 9B and LTX 2.3."
|
||||
version = "2.23.0"
|
||||
description = "ComfyUI-AusBoss (AusBoss nodes): image and video nodes with compact controls: Image and Video Crop + Rotate + Pad (frame or clip), Load Image + Pad, Image Folder, Image Resize, Align Image, Image Size, Color Match, Image Compare A/B, Realign to Source (experimental), Save Image (PNG, WebP/JXL, captions), Load Video, Load Audio, Save Video, Select Frame, Frame Interpolate, Mask Refine, Mask by Name and Mask by Touch (by name or click), LaMa Inpaint, Crop For Inpaint, Stitch Inpaint, Tiled Upscale and Tiled Upscale Stitch, FlashVSR Upscale, Edit Lock and Edit Lock Restore, LoRA Loader, Krea 2 Encode and Outpaint Model Patch, Latent Size, Layout Boxes, Seed, Select Every Nth, Split Batch, Merge Batches, Text, Text Lines, Prompt Variations, H3 Prompt Writer, Integer, Float, Math Expression, Show Text, Free Memory, Workflow Note, Workflow Note Tile, Callout, Callout Tile, Workflow Switches, Run Timer, and Step Preview. Example workflows for Qwen Image 2.1, Krea 2, Klein 9B and LTX 2.3."
|
||||
version = "2.24.0"
|
||||
# comfy-cli serializes this to JSON verbatim, so `file` publishes the literal
|
||||
# string {"file": "LICENSE"} to the registry listing. `text` puts a readable
|
||||
# name there instead. Both forms are supported; the LICENSE file is unchanged.
|
||||
@@ -21,6 +21,7 @@ keywords = [
|
||||
"image folder", "batch folder", "callout", "sticky note", "arrows",
|
||||
"tiled upscale", "upscale", "tiles", "restore", "layout boxes", "bounding boxes", "moss image", "replace text", "rotate box",
|
||||
"step preview", "live preview", "stop run",
|
||||
"text lines", "prompt list", "prompt variations", "wildcards", "flashvsr", "video upscale", "edit lock",
|
||||
]
|
||||
|
||||
# Read by comfy-cli into the registry's supported_os field. Nothing in the pack
|
||||
|
||||
@@ -56,6 +56,9 @@ PUBLIC_NODE_IDS = {
|
||||
"AUSBOSS_NODES_ColorMatch",
|
||||
"AUSBOSS_NODES_Compare",
|
||||
"AUSBOSS_NODES_CropForInpaint",
|
||||
"AUSBOSS_NODES_EditLock",
|
||||
"AUSBOSS_NODES_EditLockRestore",
|
||||
"AUSBOSS_NODES_FlashVSRUpscale",
|
||||
"AUSBOSS_NODES_Float",
|
||||
"AUSBOSS_NODES_FrameInterpolate",
|
||||
"AUSBOSS_NODES_FreeMemory",
|
||||
@@ -77,6 +80,7 @@ PUBLIC_NODE_IDS = {
|
||||
"AUSBOSS_NODES_MaskByTouch",
|
||||
"AUSBOSS_NODES_MathExpression",
|
||||
"AUSBOSS_NODES_MergeBatches",
|
||||
"AUSBOSS_NODES_PromptVariations",
|
||||
"AUSBOSS_NODES_RealignToSource",
|
||||
"AUSBOSS_NODES_RefineMask",
|
||||
"AUSBOSS_NODES_Resolution",
|
||||
@@ -92,6 +96,7 @@ PUBLIC_NODE_IDS = {
|
||||
"AUSBOSS_NODES_StitchInpaint",
|
||||
"AUSBOSS_NODES_TiledUpscaleStitch",
|
||||
"AUSBOSS_NODES_Text",
|
||||
"AUSBOSS_NODES_TextLines",
|
||||
"AUSBOSS_NODES_TiledUpscale",
|
||||
"AUSBOSS_NODES_VideoCropRotatePad",
|
||||
"AUSBOSS_NODES_VideoCropRotatePadClip",
|
||||
|
||||
@@ -0,0 +1,73 @@
|
||||
import assert from "node:assert/strict";
|
||||
import test from "node:test";
|
||||
import { readFileSync } from "node:fs";
|
||||
import { closeOnBackdropClick, pointOutside } from "../js/shared/dialog_backdrop.mjs";
|
||||
|
||||
// A stand-in for a modal <dialog>: a box from (100, 100) to (500, 400) with
|
||||
// one child, the card's content.
|
||||
function fakeDialog() {
|
||||
const dialog = new EventTarget();
|
||||
dialog.closed = 0;
|
||||
dialog.close = () => { dialog.closed += 1; };
|
||||
dialog.getBoundingClientRect = () => ({ left: 100, top: 100, right: 500, bottom: 400 });
|
||||
const child = new EventTarget();
|
||||
closeOnBackdropClick(dialog);
|
||||
return { dialog, child };
|
||||
}
|
||||
|
||||
// Events fired at the dialog itself, with the target the browser would give.
|
||||
function fire(dialog, type, target, x, y) {
|
||||
const event = new Event(type);
|
||||
Object.defineProperties(event, {
|
||||
target: { value: target },
|
||||
clientX: { value: x },
|
||||
clientY: { value: y },
|
||||
});
|
||||
dialog.dispatchEvent(event);
|
||||
}
|
||||
|
||||
function click(dialog, down, up) {
|
||||
fire(dialog, "pointerdown", down.target, down.x, down.y);
|
||||
fire(dialog, "click", up.target, up.x, up.y);
|
||||
}
|
||||
|
||||
test("pointOutside is true only beyond the box", () => {
|
||||
const rect = { left: 10, top: 10, right: 20, bottom: 20 };
|
||||
assert.equal(pointOutside(rect, 15, 15), false);
|
||||
assert.equal(pointOutside(rect, 10, 20), false);
|
||||
assert.equal(pointOutside(rect, 9, 15), true);
|
||||
assert.equal(pointOutside(rect, 15, 21), true);
|
||||
});
|
||||
|
||||
test("a click on the backdrop closes the popup", () => {
|
||||
const { dialog } = fakeDialog();
|
||||
click(dialog, { target: dialog, x: 40, y: 40 }, { target: dialog, x: 40, y: 40 });
|
||||
assert.equal(dialog.closed, 1);
|
||||
});
|
||||
|
||||
test("a click inside the card keeps it open", () => {
|
||||
const { dialog, child } = fakeDialog();
|
||||
click(dialog, { target: child, x: 200, y: 200 }, { target: child, x: 200, y: 200 });
|
||||
// The card's own padding and scroll bar are the dialog element itself.
|
||||
click(dialog, { target: dialog, x: 495, y: 200 }, { target: dialog, x: 495, y: 200 });
|
||||
assert.equal(dialog.closed, 0);
|
||||
});
|
||||
|
||||
test("selecting text or dragging out of the card keeps it open", () => {
|
||||
const { dialog, child } = fakeDialog();
|
||||
// The browser sends the click to the dialog, the common parent.
|
||||
click(dialog, { target: child, x: 200, y: 200 }, { target: dialog, x: 40, y: 40 });
|
||||
assert.equal(dialog.closed, 0);
|
||||
// A press on the backdrop that ends on the card does not close it either.
|
||||
click(dialog, { target: dialog, x: 40, y: 40 }, { target: dialog, x: 200, y: 200 });
|
||||
assert.equal(dialog.closed, 0);
|
||||
// A click with no press before it (a key on a focused button) is not outside.
|
||||
fire(dialog, "click", dialog, 40, 40);
|
||||
assert.equal(dialog.closed, 0);
|
||||
});
|
||||
|
||||
test("the Workflow Note popup closes on a backdrop click", () => {
|
||||
const source = readFileSync(new URL("../js/workflow_note/index.js", import.meta.url), "utf8");
|
||||
const reader = source.slice(source.indexOf('el("dialog", "ausboss-note-reader")'), source.indexOf("dialog.showModal()"));
|
||||
assert.match(reader, /closeOnBackdropClick\(dialog\)/);
|
||||
});
|
||||
Vendored
+99
-1
@@ -116,6 +116,67 @@
|
||||
]
|
||||
]
|
||||
},
|
||||
"AUSBOSS_NODES_EditLock": {
|
||||
"required": [
|
||||
"model",
|
||||
"vae",
|
||||
"image",
|
||||
"lock",
|
||||
"tolerance"
|
||||
],
|
||||
"optional": [
|
||||
"mask"
|
||||
],
|
||||
"outputs": [
|
||||
[
|
||||
"MODEL",
|
||||
"model"
|
||||
],
|
||||
[
|
||||
"AUSBOSS_EDIT_LOCK",
|
||||
"edit_lock"
|
||||
]
|
||||
]
|
||||
},
|
||||
"AUSBOSS_NODES_EditLockRestore": {
|
||||
"required": [
|
||||
"image",
|
||||
"edit_lock",
|
||||
"feather"
|
||||
],
|
||||
"optional": [],
|
||||
"outputs": [
|
||||
[
|
||||
"IMAGE",
|
||||
"image"
|
||||
],
|
||||
[
|
||||
"MASK",
|
||||
"edit_mask"
|
||||
],
|
||||
[
|
||||
"STRING",
|
||||
"report"
|
||||
]
|
||||
]
|
||||
},
|
||||
"AUSBOSS_NODES_FlashVSRUpscale": {
|
||||
"required": [
|
||||
"frames",
|
||||
"model",
|
||||
"scale",
|
||||
"color_match"
|
||||
],
|
||||
"optional": [
|
||||
"vae"
|
||||
],
|
||||
"outputs": [
|
||||
[
|
||||
"IMAGE",
|
||||
"frames"
|
||||
]
|
||||
]
|
||||
},
|
||||
"AUSBOSS_NODES_Float": {
|
||||
"required": [
|
||||
"value"
|
||||
@@ -739,6 +800,19 @@
|
||||
]
|
||||
]
|
||||
},
|
||||
"AUSBOSS_NODES_PromptVariations": {
|
||||
"required": [
|
||||
"text",
|
||||
"seed"
|
||||
],
|
||||
"optional": [],
|
||||
"outputs": [
|
||||
[
|
||||
"STRING",
|
||||
"prompt"
|
||||
]
|
||||
]
|
||||
},
|
||||
"AUSBOSS_NODES_RealignToSource": {
|
||||
"required": [
|
||||
"edited",
|
||||
@@ -842,7 +916,8 @@
|
||||
"name_size",
|
||||
"name_batch",
|
||||
"preview",
|
||||
"step_preview"
|
||||
"step_preview",
|
||||
"autosave"
|
||||
],
|
||||
"outputs": [
|
||||
[
|
||||
@@ -992,6 +1067,29 @@
|
||||
]
|
||||
]
|
||||
},
|
||||
"AUSBOSS_NODES_TextLines": {
|
||||
"required": [
|
||||
"text",
|
||||
"split",
|
||||
"remove_numbering",
|
||||
"limit"
|
||||
],
|
||||
"optional": [],
|
||||
"outputs": [
|
||||
[
|
||||
"STRING",
|
||||
"line"
|
||||
],
|
||||
[
|
||||
"INT",
|
||||
"number"
|
||||
],
|
||||
[
|
||||
"INT",
|
||||
"count"
|
||||
]
|
||||
]
|
||||
},
|
||||
"AUSBOSS_NODES_TiledUpscale": {
|
||||
"required": [
|
||||
"image",
|
||||
|
||||
@@ -741,3 +741,11 @@ test("every DOM widget stays visible when the graph is zoomed out", () => {
|
||||
}
|
||||
assert.ok(calls >= 14, `expected the pack's DOM widgets, found ${calls}`);
|
||||
});
|
||||
|
||||
// Save Image's settings card keeps the height released workflows were laid
|
||||
// out around. A new row made it 32 px taller, and a node saved at its
|
||||
// smallest size grew over the node under it. A new control shares a row.
|
||||
test("Save Image's settings card stays the height saved workflows expect", () => {
|
||||
const source = readFileSync(join(JS_ROOT, "save_image", "index.js"), "utf-8");
|
||||
assert.match(source, /const CARD_HEIGHT = 288;/, "js/save_image/index.js: the card's height changed; saved nodes at their smallest size would grow");
|
||||
});
|
||||
|
||||
@@ -26,3 +26,51 @@ test("the label counts from one and is empty for a single picture", () => {
|
||||
assert.equal(stepperLabel(0, 1), "");
|
||||
assert.equal(stepperLabel(0, 0), "");
|
||||
});
|
||||
|
||||
test("all the pictures are one more stop before the first", async () => {
|
||||
const { ALL, allLabel, stepWithAll } = await import("../js/shared/picture_steps.mjs");
|
||||
assert.equal(stepWithAll(ALL, 3, 1), 0);
|
||||
assert.equal(stepWithAll(0, 3, 1), 1);
|
||||
assert.equal(stepWithAll(2, 3, 1), ALL);
|
||||
assert.equal(stepWithAll(ALL, 3, -1), 2);
|
||||
assert.equal(stepWithAll(0, 3, -1), ALL);
|
||||
assert.equal(stepWithAll(1, 3, 0), 1);
|
||||
assert.equal(stepWithAll(0, 0, 1), ALL);
|
||||
assert.equal(allLabel(5), "all 5");
|
||||
assert.equal(allLabel(1), "");
|
||||
});
|
||||
|
||||
test("the grid gives each picture the most room its shape allows", async () => {
|
||||
const { gridShape } = await import("../js/shared/picture_steps.mjs");
|
||||
// Five portraits in a tall stage: two across beats one row or one column.
|
||||
assert.deepEqual(gridShape(5, 440, 1100, 3 / 4), { columns: 2, rows: 3 });
|
||||
// The same five in a wide stage go in one row.
|
||||
assert.deepEqual(gridShape(5, 1600, 440, 3 / 4), { columns: 5, rows: 1 });
|
||||
// Wide pictures in a tall stage stack.
|
||||
assert.deepEqual(gridShape(3, 400, 900, 16 / 9), { columns: 1, rows: 3 });
|
||||
assert.deepEqual(gridShape(1, 400, 400, 1), { columns: 1, rows: 1 });
|
||||
// A stage that has no size yet still gets a whole grid.
|
||||
assert.deepEqual(gridShape(5, 0, 0, 1), { columns: 3, rows: 2 });
|
||||
assert.deepEqual(gridShape(4, 400, 400, Number.NaN), { columns: 2, rows: 2 });
|
||||
});
|
||||
|
||||
test("held pictures map to their run and their place in it", async () => {
|
||||
const { heldPlaces } = await import("../js/shared/picture_steps.mjs");
|
||||
const a = "0123456789abcdef";
|
||||
const b = "fedcba9876543210";
|
||||
const name = (token, index) => ({ filename: `${token}_00${index}.png`, subfolder: "ausboss_held", type: "temp" });
|
||||
// One run of three pictures.
|
||||
assert.deepEqual(
|
||||
heldPlaces({ images: [name(a, 0), name(a, 1), name(a, 2)], ausboss_held: [a] }),
|
||||
[{ token: a, index: 0 }, { token: a, index: 1 }, { token: a, index: 2 }],
|
||||
);
|
||||
// A list run: one picture from each of two runs.
|
||||
assert.deepEqual(
|
||||
heldPlaces({ images: [name(a, 0), name(b, 0)], ausboss_held: [a, b] }),
|
||||
[{ token: a, index: 0 }, { token: b, index: 0 }],
|
||||
);
|
||||
// A saved run holds nothing, and a picture from another run is not trusted.
|
||||
assert.deepEqual(heldPlaces({ images: [{ filename: "shot_00001.png" }] }), []);
|
||||
assert.deepEqual(heldPlaces({ images: [name(a, 0), name(b, 0)], ausboss_held: [a] }), []);
|
||||
assert.deepEqual(heldPlaces(undefined), []);
|
||||
});
|
||||
|
||||
@@ -4,6 +4,7 @@ import test from "node:test";
|
||||
import {
|
||||
collapseRepeatedErrors,
|
||||
emptySourceMessage,
|
||||
explainRunError,
|
||||
findEmptySources,
|
||||
nodeLacksSource,
|
||||
nodesThatRun,
|
||||
@@ -199,3 +200,40 @@ test("an empty Load Audio asks for a sound file", () => {
|
||||
prompt[1].inputs.audio = "song.mp3";
|
||||
assert.deepEqual(findEmptySources(prompt, outputs), []);
|
||||
});
|
||||
|
||||
// ------------------------------------------------------------ known run errors
|
||||
|
||||
// The report a user posted: ComfyUI 0.44.0, H3 keeping the clip's sound, with
|
||||
// Muse MiniMax Director installed.
|
||||
const MUSE_REPORT = {
|
||||
node_type: "SamplerCustomAdvanced",
|
||||
exception_type: "KeyError",
|
||||
exception_message: "'latent'",
|
||||
traceback: [
|
||||
' File "G:\\ComfyUI\\ComfyUI\\comfy\\samplers.py", line 953, in encode_model_conds\n out = model_function(**params)\n',
|
||||
' File "G:\\ComfyUI\\ComfyUI\\custom_nodes\\muse-minimax-director-v1-2\\muse_minimax_hybrid_patch.py", line 30, in extra_conds_with_hybrid\n visual_latents.extend(item["latent"] for item in keyframes)\n',
|
||||
],
|
||||
};
|
||||
|
||||
test("the Muse MiniMax Director error is explained and names the pack", () => {
|
||||
const known = explainRunError(MUSE_REPORT);
|
||||
assert.match(known.detail, /Muse MiniMax Director/);
|
||||
assert.match(known.detail, /Turn that pack off/);
|
||||
});
|
||||
|
||||
test("the same KeyError from anywhere else is left alone", () => {
|
||||
assert.equal(explainRunError({ ...MUSE_REPORT, traceback: [' File "comfy/samplers.py", line 953\n'] }), null);
|
||||
assert.equal(explainRunError({ ...MUSE_REPORT, exception_message: "'samples'" }), null);
|
||||
});
|
||||
|
||||
test("an old ComfyUI loading a w6a8 model is told to update", () => {
|
||||
const known = explainRunError({ exception_type: "KeyError", exception_message: "'w6a8_int8'", traceback: [] });
|
||||
assert.match(known.detail, /0\.38/);
|
||||
});
|
||||
|
||||
test("an unknown error, or no report at all, gives nothing", () => {
|
||||
assert.equal(explainRunError({ exception_type: "RuntimeError", exception_message: "CUDA out of memory", traceback: [] }), null);
|
||||
assert.equal(explainRunError(null), null);
|
||||
assert.equal(explainRunError({}), null);
|
||||
assert.equal(explainRunError({ exception_type: "KeyError", exception_message: "'latent'", traceback: "muse_minimax_hybrid_patch" }) !== null, true);
|
||||
});
|
||||
|
||||
@@ -0,0 +1,373 @@
|
||||
"""Edit Lock: a local edit keeps the rest of the picture, a whole-picture edit is left alone."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import sys
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
|
||||
import torch
|
||||
|
||||
ROOT = Path(__file__).resolve().parent.parent
|
||||
sys.path.insert(0, str(ROOT))
|
||||
if "nodes" in sys.modules and not hasattr(sys.modules["nodes"], "__path__"):
|
||||
del sys.modules["nodes"]
|
||||
|
||||
from nodes._edit_lock_helpers import ( # noqa: E402
|
||||
WHOLE_PICTURE,
|
||||
EditLock,
|
||||
box_blur,
|
||||
broad_change,
|
||||
change_core,
|
||||
describe,
|
||||
difference,
|
||||
measure_redraw,
|
||||
quiet_level,
|
||||
restore_pixels,
|
||||
soften,
|
||||
)
|
||||
from nodes.node_edit_lock import ( # noqa: E402
|
||||
NODE_CLASS_MAPPINGS,
|
||||
NODE_DISPLAY_NAME_MAPPINGS,
|
||||
AusBossEditLock,
|
||||
AusBossEditLockRestore,
|
||||
)
|
||||
|
||||
H, W = 96, 80
|
||||
BOX = (slice(30, 60), slice(20, 50)) # where the planted edit sits
|
||||
|
||||
|
||||
def latents(seed: int = 0, video: bool = True):
|
||||
gen = torch.Generator().manual_seed(seed)
|
||||
# a real latent is smooth from cell to cell, so these are too
|
||||
source, other = (box_blur(torch.randn(1, 24, H, W, generator=gen), 4) for _ in range(2))
|
||||
source, other = source / source.std(), other / other.std()
|
||||
redraw = 0.15 * torch.randn(1, 24, H, W, generator=gen) # the model's own redraw noise
|
||||
if video:
|
||||
return source[:, :, None], other[:, :, None], redraw[:, :, None]
|
||||
return source, other, redraw
|
||||
|
||||
|
||||
def local_prediction(source, other, redraw, tint: float = 0.0):
|
||||
pred = source + redraw + tint
|
||||
pred[..., BOX[0], BOX[1]] = other[..., BOX[0], BOX[1]]
|
||||
return pred
|
||||
|
||||
|
||||
def lock_for(source, mode="auto", tolerance=1.0):
|
||||
lock = EditLock(mode=mode, tolerance=tolerance)
|
||||
lock.latent = source.clone()
|
||||
lock.source = torch.zeros(1, H * 16, W * 16, 3)
|
||||
return lock
|
||||
|
||||
|
||||
def run(lock, pred, sigma=1.0, steps=None):
|
||||
args = {"denoised": pred, "sigma": torch.tensor([sigma])}
|
||||
if steps is not None:
|
||||
# the sampler's own list of noise levels, the way ComfyUI hands it to the hook
|
||||
args["model_options"] = {"transformer_options": {"sample_sigmas": torch.tensor(steps)}}
|
||||
return lock.step(args)
|
||||
|
||||
|
||||
SIX = [1.0, 0.97, 0.89, 0.75, 0.53, 0.23, 0.0]
|
||||
|
||||
|
||||
class EditLockTests(unittest.TestCase):
|
||||
def test_local_edit_keeps_the_rest(self):
|
||||
for video in (True, False):
|
||||
source, other, redraw = latents(video=video)
|
||||
pred = local_prediction(source, other, redraw, tint=0.2)
|
||||
lock = lock_for(source)
|
||||
out = run(lock, pred)
|
||||
self.assertEqual(lock.status, "locked")
|
||||
self.assertEqual(out.shape, pred.shape)
|
||||
# far from the edit the original is back exactly, tint and redraw noise gone
|
||||
self.assertTrue(torch.equal(out[..., :12, :], source[..., :12, :].to(out.dtype)))
|
||||
self.assertTrue(torch.equal(out[..., 80:, :], source[..., 80:, :].to(out.dtype)))
|
||||
# inside the edit the model's prediction is untouched
|
||||
inner = (slice(36, 54), slice(26, 44))
|
||||
self.assertTrue(torch.equal(out[..., inner[0], inner[1]], pred[..., inner[0], inner[1]]))
|
||||
kept = lock.kept_share()
|
||||
self.assertGreater(kept, 0.7)
|
||||
self.assertLess(kept, 0.92)
|
||||
self.assertIn("kept", describe(lock))
|
||||
|
||||
def test_report_says_what_to_do_when_nearly_all_was_kept(self):
|
||||
source, other, redraw = latents(video=False)
|
||||
lock = lock_for(source)
|
||||
run(lock, local_prediction(source, other, redraw))
|
||||
self.assertNotIn("set lock to off", describe(lock))
|
||||
# a faint change everywhere (rain) and one small clear spot: nearly all of it is kept
|
||||
pred = source + redraw
|
||||
pred[..., 40:46, 30:36] = other[..., 40:46, 30:36]
|
||||
lock = lock_for(source)
|
||||
run(lock, pred)
|
||||
self.assertGreater(lock.kept_share(), 0.95)
|
||||
self.assertIn("set lock to off", describe(lock))
|
||||
|
||||
def test_it_watches_first_and_locks_when_the_noise_is_down(self):
|
||||
source, other, redraw = latents(video=False)
|
||||
lock = lock_for(source)
|
||||
# the model's own redraw noise fades from step to step, as it does in a real run
|
||||
for sigma, fade in ((1.0, 1.0), (0.97, 0.8), (0.89, 0.7)):
|
||||
pred = local_prediction(source, other, redraw * fade)
|
||||
out = run(lock, pred, sigma, SIX)
|
||||
self.assertEqual(lock.status, "watching")
|
||||
self.assertIsNone(lock.pixel_mask(H * 16, W * 16))
|
||||
# nothing is put back yet, and with no zoom to undo the prediction passes as it is
|
||||
self.assertTrue(torch.equal(out, pred))
|
||||
pred = local_prediction(source, other, redraw * 0.65)
|
||||
out = run(lock, pred, 0.75, SIX)
|
||||
self.assertEqual(lock.status, "locked")
|
||||
self.assertGreater(lock.top, 4.0)
|
||||
self.assertTrue(torch.equal(out[..., :12, :], source[..., :12, :]))
|
||||
inner = (slice(36, 54), slice(26, 44))
|
||||
self.assertTrue(torch.equal(out[..., inner[0], inner[1]], pred[..., inner[0], inner[1]]))
|
||||
|
||||
def test_a_faint_change_everywhere_is_let_through(self):
|
||||
source, _, redraw = latents(video=False)
|
||||
rain = 0.25 * box_blur(torch.randn(1, 24, H, W, generator=torch.Generator().manual_seed(5)), 1)
|
||||
for mode, status in (("auto", "soft"), ("always", "locked")):
|
||||
lock = lock_for(source, mode=mode)
|
||||
# nothing shows at the first look; then the rain comes in everywhere, and no spot changes clearly
|
||||
for sigma, fade, wet in ((1.0, 1.0, 0.0), (0.97, 0.9, 0.5), (0.89, 0.8, 1.0)):
|
||||
run(lock, source + redraw * fade + rain * wet, sigma, SIX)
|
||||
self.assertEqual(lock.status, "watching")
|
||||
pred = source + redraw * 0.7 + rain * 1.5
|
||||
out = run(lock, pred, 0.75, SIX)
|
||||
self.assertEqual(lock.status, status)
|
||||
self.assertLess(lock.top, 4.0)
|
||||
if mode == "auto":
|
||||
self.assertLess(float((out - pred).abs().max()), 1e-4)
|
||||
self.assertIn("stepped aside", describe(lock))
|
||||
|
||||
def test_a_wide_faint_change_with_no_peak_is_not_locked_in_patches(self):
|
||||
source, _, redraw = latents(video=False)
|
||||
gen = torch.Generator().manual_seed(9)
|
||||
snow = torch.zeros(1, 24, H, W)
|
||||
snow[..., :, :34] = 0.25 * torch.randn(1, 24, H, 34, generator=gen) # faint, over two fifths of the picture
|
||||
theta = torch.tensor([[[1 / 1.015, 0.0, 0.0], [0.0, 1 / 1.015, 0.0]]])
|
||||
grid = torch.nn.functional.affine_grid(theta, list(source.shape), align_corners=False)
|
||||
zoomed = torch.nn.functional.grid_sample(source, grid, mode="bilinear", padding_mode="border", align_corners=False)
|
||||
for mode, status in (("auto", "soft"), ("always", "locked")):
|
||||
lock = lock_for(source, mode=mode)
|
||||
for sigma, fade in ((1.0, 1.0), (0.97, 0.8), (0.89, 0.7)):
|
||||
run(lock, source + redraw * fade, sigma, SIX)
|
||||
pred = source + redraw * 0.65 + snow
|
||||
out = run(lock, pred, 0.75, SIX)
|
||||
self.assertEqual(lock.status, status)
|
||||
self.assertGreater(lock.wide, 0.10)
|
||||
self.assertLess(lock.peak, 4.0)
|
||||
if mode == "always":
|
||||
self.assertLess(lock.kept_share(), 0.8)
|
||||
continue
|
||||
self.assertIsNone(lock.pixel_mask(H * 16, W * 16))
|
||||
self.assertLess(float((out - pred).abs().max()), 1e-4)
|
||||
self.assertIn("held the picture in place", describe(lock))
|
||||
# for the rest of the run only the zoom is undone
|
||||
out = run(lock, zoomed, 0.53, SIX)
|
||||
self.assertEqual(lock.status, "soft")
|
||||
smooth = lambda x: box_blur(x, 2)[..., 6:-6, 6:-6]
|
||||
self.assertLess(float((smooth(out) - smooth(source)).abs().mean()),
|
||||
float((smooth(zoomed) - smooth(source)).abs().mean()) * 0.35)
|
||||
|
||||
def test_the_last_step_decides_when_the_noise_never_gets_that_low(self):
|
||||
source, other, redraw = latents(video=False)
|
||||
lock = lock_for(source)
|
||||
two = [1.0, 0.92, 0.0]
|
||||
run(lock, local_prediction(source, other, redraw), 1.0, two)
|
||||
self.assertEqual(lock.status, "watching")
|
||||
run(lock, local_prediction(source, other, redraw * 0.8), 0.92, two)
|
||||
self.assertEqual(lock.status, "locked")
|
||||
# a one-step run decides at once
|
||||
lock = lock_for(source)
|
||||
run(lock, local_prediction(source, other, redraw), 1.0, [1.0, 0.0])
|
||||
self.assertEqual(lock.status, "locked")
|
||||
|
||||
def test_the_zoom_is_undone_while_it_watches(self):
|
||||
source, other, redraw = latents(video=False)
|
||||
theta = torch.tensor([[[1 / 1.015, 0.0, 0.0], [0.0, 1 / 1.015, 0.0]]])
|
||||
grid = torch.nn.functional.affine_grid(theta, list(source.shape), align_corners=False)
|
||||
zoomed = torch.nn.functional.grid_sample(source, grid, mode="bilinear", padding_mode="border", align_corners=False)
|
||||
lock = lock_for(source)
|
||||
out = run(lock, zoomed, 1.0, SIX)
|
||||
self.assertEqual(lock.status, "watching")
|
||||
# read on smoothed copies: this is about where things sit, and resampling costs fine detail either way
|
||||
smooth = lambda x: box_blur(x, 2)[..., 6:-6, 6:-6]
|
||||
off_before = float((smooth(zoomed) - smooth(source)).abs().mean())
|
||||
off_after = float((smooth(out) - smooth(source)).abs().mean())
|
||||
self.assertLess(off_after, off_before * 0.35)
|
||||
|
||||
def test_whole_picture_edit_is_left_alone(self):
|
||||
source, other, _ = latents()
|
||||
pred = 0.3 * source + other # a new picture over the old one
|
||||
lock = lock_for(source)
|
||||
out = run(lock, pred)
|
||||
self.assertEqual(lock.status, "whole picture")
|
||||
self.assertTrue(torch.equal(out, pred))
|
||||
self.assertIsNone(lock.pixel_mask(H * 16, W * 16))
|
||||
# and it stays aside for the rest of the run
|
||||
self.assertTrue(torch.equal(run(lock, pred, 0.8), pred))
|
||||
self.assertIn("stepped aside", describe(lock))
|
||||
|
||||
def test_a_slightly_larger_redraw_is_measured_and_is_not_an_edit(self):
|
||||
source, other, redraw = latents(video=False)
|
||||
# drawn 1.5% wider, 1% taller and half a cell to the right
|
||||
theta = torch.tensor([[[1 / 1.015, 0.0, -1.0 / W], [0.0, 1 / 1.01, 0.0]]])
|
||||
grid = torch.nn.functional.affine_grid(theta, list(source.shape), align_corners=False)
|
||||
zoomed = torch.nn.functional.grid_sample(source, grid, mode="bilinear", padding_mode="border", align_corners=False)
|
||||
move = measure_redraw(local_prediction(zoomed, other, redraw), source)
|
||||
self.assertAlmostEqual(float(move[0]), 1 / 1.015 - 1, delta=0.002)
|
||||
self.assertAlmostEqual(float(move[1]), 1 / 1.01 - 1, delta=0.002)
|
||||
self.assertAlmostEqual(float(move[2]), -0.5, delta=0.15)
|
||||
self.assertAlmostEqual(float(move[3]), 0.0, delta=0.15)
|
||||
self.assertLess(float(measure_redraw(source + redraw, source).abs().max()), 0.02)
|
||||
self.assertLess(float(broad_change(zoomed + redraw, source)), WHOLE_PICTURE / 2)
|
||||
lock = lock_for(source)
|
||||
out = run(lock, local_prediction(zoomed, other, redraw))
|
||||
self.assertEqual(lock.status, "locked")
|
||||
# the zoom did not read as an edit: only the planted box is free, and the rest is the original, unzoomed
|
||||
self.assertGreater(lock.kept_share(), 0.7)
|
||||
self.assertTrue(torch.equal(out[..., :12, :], source[..., :12, :]))
|
||||
# an even tint is not one either; new light across the picture is
|
||||
self.assertLess(float(broad_change(source + redraw + 0.6, source)), WHOLE_PICTURE / 2)
|
||||
self.assertGreater(float(broad_change(source + redraw + 0.6 * other, source)), WHOLE_PICTURE)
|
||||
|
||||
def test_always_locks_a_whole_picture_edit(self):
|
||||
source, other, _ = latents()
|
||||
lock = lock_for(source, mode="always")
|
||||
run(lock, 0.3 * source + other)
|
||||
self.assertEqual(lock.status, "locked")
|
||||
|
||||
def test_changed_stays_changed_and_a_new_run_starts_clean(self):
|
||||
source, other, redraw = latents()
|
||||
lock = lock_for(source)
|
||||
run(lock, local_prediction(source, other, redraw), 1.0)
|
||||
first = lock.mask.clone()
|
||||
run(lock, source + redraw, 0.8) # a later step that agrees with the original everywhere
|
||||
self.assertTrue(torch.all(lock.mask >= first))
|
||||
run(lock, source + redraw, 1.0) # the noise level went up: a new run
|
||||
self.assertLess(float(lock.mask.max()), 0.5)
|
||||
|
||||
def test_given_mask_is_used_as_is(self):
|
||||
source, other, redraw = latents()
|
||||
lock = lock_for(source)
|
||||
mask = torch.zeros(1, H * 16, W * 16)
|
||||
mask[:, : H * 8] = 1.0 # the top half may change
|
||||
lock.user_mask = mask
|
||||
pred = 0.3 * source + other # would count as a whole-picture edit
|
||||
out = run(lock, pred)
|
||||
self.assertEqual(lock.status, "locked")
|
||||
self.assertTrue(torch.equal(out[..., :40, :], pred[..., :40, :]))
|
||||
self.assertTrue(torch.equal(out[..., 56:, :], source[..., 56:, :]))
|
||||
self.assertIn("your mask", describe(lock))
|
||||
|
||||
def test_other_sizes_video_and_wrong_vae_do_nothing(self):
|
||||
source, other, redraw = latents()
|
||||
pred = local_prediction(source, other, redraw)
|
||||
lock = lock_for(source[..., :48, :40].clone())
|
||||
self.assertTrue(torch.equal(run(lock, pred), pred))
|
||||
self.assertEqual(lock.status, "size")
|
||||
lock = lock_for(source)
|
||||
clip = pred.expand(-1, -1, 5, -1, -1)
|
||||
self.assertTrue(torch.equal(run(lock, clip), clip))
|
||||
self.assertEqual(lock.status, "video")
|
||||
lock = lock_for(source[:, :16].clone())
|
||||
self.assertTrue(torch.equal(run(lock, pred), pred))
|
||||
self.assertEqual(lock.status, "vae")
|
||||
for status in ("size", "video", "vae"):
|
||||
lock.status = status
|
||||
self.assertIn("did nothing", describe(lock))
|
||||
|
||||
def test_a_grid_one_cell_off_still_locks(self):
|
||||
source, other, redraw = latents(video=False)
|
||||
lock = lock_for(torch.nn.functional.interpolate(source, size=(H + 2, W + 1), mode="bilinear"))
|
||||
run(lock, local_prediction(source, other, redraw))
|
||||
self.assertEqual(lock.status, "locked")
|
||||
self.assertGreater(lock.kept_share(), 0.7)
|
||||
|
||||
def test_tolerance_moves_the_line(self):
|
||||
source, other, redraw = latents(video=False)
|
||||
pred = source + redraw
|
||||
pred[..., BOX[0], BOX[1]] += 0.45 * other[..., BOX[0], BOX[1]] # a faint change
|
||||
magnitude = difference(pred, source)
|
||||
quiet = quiet_level(magnitude)
|
||||
loose = change_core(magnitude, quiet, 0.5)
|
||||
strict = change_core(magnitude, quiet, 2.0)
|
||||
self.assertGreater(float(loose[..., 40:50, 30:40].mean()), 0.9)
|
||||
self.assertEqual(float(strict.max()), 0.0)
|
||||
|
||||
def test_faint_change_counts_only_beside_a_clear_one(self):
|
||||
quiet = torch.tensor([1.0])
|
||||
magnitude = torch.ones(1, 1, 40, 40)
|
||||
magnitude[..., 5:15, 5:15] = 2.2 # faint, alone: redraw noise
|
||||
magnitude[..., 22:34, 22:34] = 2.2 # faint ...
|
||||
magnitude[..., 26:30, 26:30] = 4.0 # ... around a clear change
|
||||
core = change_core(magnitude, quiet)
|
||||
self.assertEqual(float(core[..., 5:15, 5:15].max()), 0.0)
|
||||
self.assertEqual(float(core[..., 22:34, 22:34].min()), 1.0)
|
||||
self.assertEqual(float(core.sum()), 144.0)
|
||||
# what an earlier step held stays, and faint change that touches it joins
|
||||
held = torch.zeros_like(core)
|
||||
held[..., 8:10, 8:10] = 1.0
|
||||
again = change_core(torch.ones(1, 1, 40, 40), quiet, held=held)
|
||||
self.assertEqual(float(again.sum()), 4.0)
|
||||
joined = change_core(magnitude, quiet, held=held)
|
||||
self.assertEqual(float(joined[..., 5:15, 5:15].min()), 1.0)
|
||||
soft = soften(core)
|
||||
self.assertTrue(torch.all(soft[..., 22:34, 22:34] == 1.0))
|
||||
self.assertGreater(float(soft[0, 0, 21, 28]), 0.5)
|
||||
self.assertEqual(float(soft[0, 0, 2, 2]), 0.0)
|
||||
|
||||
def test_restore_gives_the_original_pixels_back(self):
|
||||
gen = torch.Generator().manual_seed(1)
|
||||
original = torch.rand(1, 128, 160, 3, generator=gen)
|
||||
result = torch.rand(2, 128, 160, 3, generator=gen)
|
||||
mask = torch.zeros(1, 128, 160)
|
||||
mask[:, 40:80, 60:100] = 1.0
|
||||
out, edit = restore_pixels(result, original, mask, 8)
|
||||
self.assertEqual(out.shape, result.shape)
|
||||
self.assertEqual(edit.shape, (2, 128, 160))
|
||||
self.assertTrue(torch.equal(out[:, :16], original[:, :16].expand(2, -1, -1, -1)))
|
||||
self.assertTrue(torch.equal(out[:, 40:80, 60:100], result[:, 40:80, 60:100]))
|
||||
# the fade lies outside the edit
|
||||
self.assertTrue(torch.all(edit[:, 40:80, 60:100] == 1.0))
|
||||
self.assertGreater(float(edit[0, 36, 80]), 0.0)
|
||||
# another size: the original is resized to the result
|
||||
small = torch.rand(1, 64, 80, 3, generator=gen)
|
||||
out, _ = restore_pixels(result, small, mask, 0)
|
||||
self.assertEqual(out.shape, result.shape)
|
||||
|
||||
def test_restore_node_passes_through_when_nothing_was_locked(self):
|
||||
lock = EditLock()
|
||||
lock.status = "whole picture"
|
||||
image = torch.rand(1, 32, 48, 3)
|
||||
out, mask, report = AusBossEditLockRestore().restore(image, lock, 16)
|
||||
self.assertTrue(torch.equal(out, image))
|
||||
self.assertEqual(mask.shape, (1, 32, 48))
|
||||
self.assertTrue(torch.all(mask == 1.0))
|
||||
self.assertIn("stepped aside", report)
|
||||
|
||||
def test_restore_node_after_a_locked_run(self):
|
||||
source, other, redraw = latents()
|
||||
lock = lock_for(source)
|
||||
lock.source = torch.rand(1, H * 4, W * 4, 3)
|
||||
run(lock, local_prediction(source, other, redraw))
|
||||
image = torch.rand(1, H * 4, W * 4, 3)
|
||||
out, mask, report = AusBossEditLockRestore().restore(image, lock, 4)
|
||||
self.assertTrue(torch.equal(out[:, :24], lock.source[:, :24]))
|
||||
self.assertTrue(torch.equal(out[:, 150:210, 110:170], image[:, 150:210, 110:170]))
|
||||
self.assertGreater(float(mask[0, 180, 140]), 0.99)
|
||||
self.assertIn("kept", report)
|
||||
|
||||
def test_mappings(self):
|
||||
keys = {"AUSBOSS_NODES_EditLock", "AUSBOSS_NODES_EditLockRestore"}
|
||||
self.assertEqual(set(NODE_CLASS_MAPPINGS), keys)
|
||||
self.assertEqual(set(NODE_DISPLAY_NAME_MAPPINGS), keys)
|
||||
self.assertIs(NODE_CLASS_MAPPINGS["AUSBOSS_NODES_EditLock"], AusBossEditLock)
|
||||
for name in NODE_DISPLAY_NAME_MAPPINGS.values():
|
||||
self.assertTrue(name.endswith(" 🆎"))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,225 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from pathlib import Path
|
||||
import importlib.util
|
||||
import math
|
||||
import unittest
|
||||
|
||||
import torch
|
||||
|
||||
ROOT = Path(__file__).resolve().parent.parent
|
||||
# Loaded by path: the model code is plain PyTorch and needs nothing the pack imports.
|
||||
_spec = importlib.util.spec_from_file_location("ausboss_flashvsr_helpers", ROOT / "nodes" / "_flashvsr_helpers.py")
|
||||
vsr = importlib.util.module_from_spec(_spec)
|
||||
_spec.loader.exec_module(vsr)
|
||||
|
||||
|
||||
def tiny_model(seed: int = 0):
|
||||
torch.manual_seed(seed)
|
||||
model = vsr.FlashVSRModel(
|
||||
dim=24,
|
||||
heads=2,
|
||||
ffn_dim=48,
|
||||
layers=2,
|
||||
text_dim=8,
|
||||
text_len=4,
|
||||
freq_dim=16,
|
||||
projector_hidden=(8, 12),
|
||||
decoder_widths=(8, 8, 4, 4),
|
||||
)
|
||||
for parameter in model.parameters():
|
||||
torch.nn.init.normal_(parameter, std=0.2)
|
||||
model.positive_conditioning.normal_()
|
||||
return model.eval()
|
||||
|
||||
|
||||
class SizeTests(unittest.TestCase):
|
||||
def test_output_is_even_and_the_working_size_is_a_multiple_of_128(self):
|
||||
self.assertEqual(vsr.plan_sizes(360, 640, 2.0), (720, 1280, 768, 1280))
|
||||
self.assertEqual(vsr.plan_sizes(271, 479, 1.5), (406, 718, 512, 768))
|
||||
self.assertEqual(vsr.plan_sizes(64, 64, 2.0), (128, 128, 128, 128))
|
||||
|
||||
def test_the_padded_clip_always_covers_every_frame(self):
|
||||
for count in (1, 5, 21, 22, 29, 30, 81, 200):
|
||||
groups = vsr.plan_groups(count)
|
||||
self.assertGreaterEqual(groups, 3)
|
||||
self.assertGreaterEqual(8 * groups - 3, count)
|
||||
latents = sum(n for _, n in vsr.pass_layout(groups))
|
||||
self.assertEqual(latents, 2 * groups)
|
||||
|
||||
def test_passes_follow_one_another_without_a_gap(self):
|
||||
expected = 0
|
||||
for first, count in vsr.pass_layout(6):
|
||||
self.assertEqual(first, expected)
|
||||
expected += count
|
||||
|
||||
|
||||
class BlockTests(unittest.TestCase):
|
||||
def test_blocks_round_trip(self):
|
||||
x = torch.randn(4 * 16 * 24, 6)
|
||||
back = vsr.from_blocks(vsr.to_blocks(x, 4, 16, 24), 4, 16, 24)
|
||||
self.assertTrue(torch.equal(x, back))
|
||||
|
||||
def test_a_block_holds_two_frames_of_an_8_by_8_patch(self):
|
||||
frames, rows, cols = 2, 16, 16
|
||||
t, r, c = torch.meshgrid(torch.arange(frames), torch.arange(rows), torch.arange(cols), indexing="ij")
|
||||
ids = torch.stack([t, r // 8, c // 8], dim=-1).reshape(-1, 3).float()
|
||||
blocks = vsr.to_blocks(ids, frames, rows, cols)
|
||||
self.assertEqual(blocks.shape, (4, 128, 3))
|
||||
for block in blocks:
|
||||
self.assertEqual(block[:, 1:].unique(dim=0).shape[0], 1)
|
||||
self.assertEqual(sorted(block[:, 0].unique().tolist()), [0.0, 1.0])
|
||||
|
||||
def test_the_neighbourhood_is_centred_and_cut_at_the_edges(self):
|
||||
mask = vsr.local_block_mask(3, 20, 11, "cpu")
|
||||
row = mask[20 + 10].view(3, 20) # the block at row 1, column 10
|
||||
self.assertTrue(row[:, 5:16].all())
|
||||
self.assertFalse(row[:, :5].any() or row[:, 16:].any())
|
||||
self.assertEqual(int(mask[0].sum()), 3 * 6) # the corner block: all 3 rows, columns 0-5
|
||||
|
||||
def test_selection_keeps_the_asked_number_inside_the_neighbourhood(self):
|
||||
torch.manual_seed(1)
|
||||
heads, per_group = 2, 6
|
||||
local = vsr.local_block_mask(2, 3, 3, "cpu")
|
||||
q = torch.randn(2 * per_group, 128, 8)
|
||||
k = torch.randn(3 * per_group, 128, 8)
|
||||
chosen, shares = vsr.choose_key_blocks(q, k, heads, local, 2, keep=20)
|
||||
self.assertTrue((shares[chosen] > 0).all())
|
||||
self.assertTrue((shares[:, ~local.repeat(2, 3)] == 0).all())
|
||||
self.assertEqual(chosen.shape, (heads, 12, 18))
|
||||
self.assertTrue((chosen.view(heads, 2, -1).sum(-1) == 20).all())
|
||||
self.assertFalse((chosen & ~local.repeat(2, 3)).any())
|
||||
|
||||
|
||||
class SparseAttentionTests(unittest.TestCase):
|
||||
def dense(self, q, k, v, chosen):
|
||||
heads, nq, tokens, dim = q.shape
|
||||
nk = k.shape[1]
|
||||
scores = torch.einsum("hqtd,hksd->hqtks", q, k) / math.sqrt(dim)
|
||||
scores = scores.masked_fill(~chosen[:, :, None, :, None], float("-inf"))
|
||||
weights = scores.reshape(heads, nq, tokens, nk * tokens).softmax(-1).nan_to_num(0.0)
|
||||
return weights @ v.reshape(heads, 1, nk * tokens, dim)
|
||||
|
||||
def test_it_matches_plain_masked_attention(self):
|
||||
torch.manual_seed(2)
|
||||
q, k, v = torch.randn(3, 7, 128, 4), torch.randn(3, 11, 128, 4), torch.randn(3, 11, 128, 4)
|
||||
chosen = torch.rand(3, 7, 11) < 0.4
|
||||
chosen[0, 0] = False # a row that chose nothing
|
||||
chosen[1, 2] = True # a row that chose everything
|
||||
want = self.dense(q, k, v, chosen)
|
||||
for budget in (1, 8, 40, 100000):
|
||||
got = vsr.sparse_block_attention(q, k, v, chosen, slot_budget=budget)
|
||||
self.assertTrue(torch.allclose(got, want, atol=1e-5), budget)
|
||||
self.assertTrue((got[0, 0] == 0).all())
|
||||
|
||||
def test_with_shares_a_row_is_filled_up_from_its_neighbourhood_and_never_reads_less(self):
|
||||
torch.manual_seed(4)
|
||||
q, k, v = torch.randn(2, 6, 128, 4), torch.randn(2, 9, 128, 4), torch.randn(2, 9, 128, 4)
|
||||
allowed = torch.rand(2, 6, 9) < 0.8
|
||||
allowed[..., 0] = True
|
||||
shares = torch.rand(2, 6, 9) * allowed
|
||||
chosen = (shares > 0.5) & allowed
|
||||
chosen[..., 0] = True # every row chose something
|
||||
chosen[1, 3] = allowed[1, 3] # the widest row: everything it may read
|
||||
widest = int(chosen.sum(-1).max())
|
||||
# one call at the widest row's width: each row reads its chosen blocks, then its best other ones
|
||||
rank = chosen.float() * 2 + shares
|
||||
top = torch.zeros_like(chosen).scatter(-1, rank.argsort(-1, descending=True)[..., :widest], True)
|
||||
want = self.dense(q, k, v, top & allowed)
|
||||
got = vsr.sparse_block_attention(q, k, v, chosen, shares, slot_budget=100000)
|
||||
self.assertTrue(torch.allclose(got, want, atol=1e-5))
|
||||
self.assertTrue(((top & allowed) | ~chosen).all())
|
||||
|
||||
|
||||
class DecoderTests(unittest.TestCase):
|
||||
def test_frames_one_at_a_time_match_all_at_once(self):
|
||||
decoder = tiny_model().tc_decoder.double() # double precision: random weights make large numbers
|
||||
x = torch.randn(5, 16 + 768, 4, 6, dtype=torch.float64)
|
||||
with torch.no_grad():
|
||||
whole = decoder.decode({}, x)
|
||||
memory = {}
|
||||
parts = torch.cat([decoder.decode(memory, x[i : i + 1]) for i in range(5)])
|
||||
self.assertEqual(whole.shape, (5 * 4 - 3, 3, 32, 48))
|
||||
self.assertTrue(torch.allclose(whole, parts, rtol=1e-9, atol=1e-9))
|
||||
|
||||
|
||||
class ProjectorTests(unittest.TestCase):
|
||||
def test_one_set_of_tokens_per_four_frames(self):
|
||||
model = tiny_model()
|
||||
memory = {}
|
||||
clip = torch.rand(1, 3, 9, 32, 48) * 2 - 1
|
||||
with torch.no_grad():
|
||||
self.assertIsNone(model.lq_proj_in.push(memory, clip[:, :, :1]))
|
||||
first = model.lq_proj_in.push(memory, clip[:, :, 1:5])
|
||||
second = model.lq_proj_in.push(memory, clip[:, :, 5:9])
|
||||
self.assertEqual(first.shape, (1, 2 * 3, 24))
|
||||
self.assertEqual(second.shape, first.shape)
|
||||
self.assertFalse(torch.allclose(first, second))
|
||||
|
||||
|
||||
class WholeClipTests(unittest.TestCase):
|
||||
def run_clip(self, count, height, width, scale, **kw):
|
||||
torch.manual_seed(3)
|
||||
frames = torch.rand(count, height, width, 3)
|
||||
return vsr.upscale_frames(tiny_model(), frames, scale, device="cpu", dtype=torch.float32, **kw), frames
|
||||
|
||||
def test_every_frame_comes_back_at_the_new_size(self):
|
||||
for count, height, width, scale in ((1, 64, 64, 2.0), (23, 50, 90, 2.0), (30, 128, 128, 1.0)):
|
||||
out, _ = self.run_clip(count, height, width, scale)
|
||||
self.assertEqual(out.shape, (count, vsr._even(height * scale), vsr._even(width * scale), 3))
|
||||
self.assertTrue(torch.isfinite(out).all())
|
||||
self.assertTrue(0.0 <= float(out.min()) and float(out.max()) <= 1.0)
|
||||
|
||||
def test_the_same_seed_gives_the_same_clip(self):
|
||||
a, _ = self.run_clip(9, 64, 64, 2.0, seed=5)
|
||||
b, _ = self.run_clip(9, 64, 64, 2.0, seed=5)
|
||||
c, _ = self.run_clip(9, 64, 64, 2.0, seed=6)
|
||||
self.assertTrue(torch.equal(a, b))
|
||||
self.assertFalse(torch.equal(a, c))
|
||||
|
||||
def test_colour_match_brings_the_means_back_to_the_source(self):
|
||||
out, frames = self.run_clip(5, 64, 64, 2.0, color_match=True)
|
||||
self.assertLess(float((out.mean(dim=(1, 2)) - frames.mean(dim=(1, 2))).abs().max()), 0.05)
|
||||
|
||||
def test_another_decoder_gets_every_latent_frame(self):
|
||||
seen = {}
|
||||
|
||||
def decode(latents):
|
||||
seen["shape"] = tuple(latents.shape)
|
||||
frames = 4 * (latents.shape[2] - 1) + 1
|
||||
return torch.full((frames, latents.shape[3] * 8, latents.shape[4] * 8, 3), 0.5)
|
||||
|
||||
out, _ = self.run_clip(23, 64, 64, 2.0, decode_latents=decode, color_match=False)
|
||||
self.assertEqual(seen["shape"], (1, 16, 8, 16, 16))
|
||||
self.assertEqual(out.shape, (23, 128, 128, 3))
|
||||
self.assertTrue(torch.allclose(out, torch.full_like(out, 0.5)))
|
||||
|
||||
def test_a_large_picture_is_done_in_tiles_that_cover_it(self):
|
||||
self.assertEqual(vsr.plan_tiles(1152, 2048, None), [(0, 0, 1152, 2048)])
|
||||
self.assertEqual(vsr.plan_tiles(640, 1024, 640 * 1024), [(0, 0, 640, 1024)])
|
||||
for work_h, work_w, limit in ((1152, 2048, 1_300_000), (2176, 3840, 2_000_000), (768, 768, 100_000), (128, 1280, 50_000)):
|
||||
tiles = vsr.plan_tiles(work_h, work_w, limit)
|
||||
self.assertGreater(len(tiles), 1)
|
||||
cover = torch.zeros(work_h, work_w)
|
||||
for top, left, height, width in tiles:
|
||||
self.assertEqual((height % 128, width % 128, top % 16, left % 16), (0, 0, 0, 0))
|
||||
self.assertTrue(top + height <= work_h and left + width <= work_w)
|
||||
self.assertLessEqual(height * width, max(limit, 128 * 128))
|
||||
cover[top : top + height, left : left + width] += vsr.tile_weight((top, left, height, width), work_h, work_w)
|
||||
self.assertGreater(float(cover.min()), 0.5)
|
||||
|
||||
def test_tiles_give_every_frame_at_the_same_size(self):
|
||||
whole, _ = self.run_clip(9, 128, 192, 2.0)
|
||||
tiled, _ = self.run_clip(9, 128, 192, 2.0, max_tile_pixels=256 * 256)
|
||||
self.assertEqual(tiled.shape, whole.shape)
|
||||
self.assertTrue(torch.isfinite(tiled).all())
|
||||
self.assertTrue(0.0 <= float(tiled.min()) and float(tiled.max()) <= 1.0)
|
||||
|
||||
def test_progress_counts_the_passes(self):
|
||||
calls = []
|
||||
self.run_clip(40, 64, 64, 2.0, on_pass=lambda done, total: calls.append((done, total)))
|
||||
self.assertEqual(calls, [(1, 4), (2, 4), (3, 4), (4, 4)])
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,110 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from datetime import datetime
|
||||
from pathlib import Path
|
||||
import sys
|
||||
import tempfile
|
||||
import unittest
|
||||
|
||||
ROOT = Path(__file__).resolve().parent.parent
|
||||
sys.path.insert(0, str(ROOT))
|
||||
if "nodes" in sys.modules and not hasattr(sys.modules["nodes"], "__path__"):
|
||||
del sys.modules["nodes"]
|
||||
|
||||
from nodes._image_hold_helpers import GONE, HELD_FOLDER, hold_pictures, keep_pictures
|
||||
|
||||
NOW = datetime(2026, 11, 10, 9, 30, 0)
|
||||
|
||||
|
||||
def settings(**changes):
|
||||
base = {
|
||||
"format": "png", "extension": "png", "width": 64, "height": 96, "exact": "",
|
||||
"filename_prefix": "shoot", "on_existing": "overwrite", "output_dir": "",
|
||||
"name_counter": True, "name_date": False, "name_time": False, "name_size": False,
|
||||
"name_batch": False, "caption": "",
|
||||
}
|
||||
base.update(changes)
|
||||
return base
|
||||
|
||||
|
||||
class HoldTests(unittest.TestCase):
|
||||
def setUp(self):
|
||||
self.box = tempfile.TemporaryDirectory()
|
||||
self.temp = Path(self.box.name) / "temp"
|
||||
self.output = Path(self.box.name) / "output"
|
||||
self.temp.mkdir()
|
||||
self.output.mkdir()
|
||||
|
||||
def tearDown(self):
|
||||
self.box.cleanup()
|
||||
|
||||
def hold(self, count=3, **changes):
|
||||
frames = [f"picture {index}".encode() for index in range(count)]
|
||||
return hold_pictures(self.temp, frames, settings(**changes), lambda path, frame: path.write_bytes(frame))
|
||||
|
||||
def keep(self, token, picks=None):
|
||||
return keep_pictures(self.temp, self.output, token, picks, now=NOW)
|
||||
|
||||
def test_a_held_run_saves_nothing_until_it_is_kept(self):
|
||||
token, previews = self.hold()
|
||||
self.assertEqual(len(previews), 3)
|
||||
self.assertTrue(all(entry["type"] == "temp" and entry["subfolder"] == HELD_FOLDER for entry in previews))
|
||||
self.assertEqual(list(self.output.iterdir()), [])
|
||||
answer = self.keep(token)
|
||||
self.assertEqual([item["path"] for item in answer["saved"]], ["shoot_00001.png", "shoot_00002.png", "shoot_00003.png"])
|
||||
self.assertEqual((self.output / "shoot_00002.png").read_bytes(), b"picture 1")
|
||||
|
||||
def test_one_picture_can_be_kept_and_is_not_saved_twice(self):
|
||||
token, _ = self.hold()
|
||||
first = self.keep(token, [2])
|
||||
self.assertEqual(first["saved"], [{"index": 2, "path": "shoot_00001.png"}])
|
||||
self.assertEqual((self.output / "shoot_00001.png").read_bytes(), b"picture 2")
|
||||
again = self.keep(token, [2])
|
||||
self.assertEqual(again["saved"], [])
|
||||
rest = self.keep(token)
|
||||
self.assertEqual([item["index"] for item in rest["saved"]], [0, 1])
|
||||
self.assertEqual(rest["kept"], [0, 1, 2])
|
||||
self.assertEqual(sorted(entry.name for entry in self.output.iterdir()), ["shoot_00001.png", "shoot_00002.png", "shoot_00003.png"])
|
||||
|
||||
def test_the_folder_and_caption_of_the_run_are_used(self):
|
||||
token, _ = self.hold(count=1, output_dir="sets/portraits", caption="a portrait")
|
||||
answer = self.keep(token)
|
||||
self.assertEqual(answer["saved"][0]["path"], "sets/portraits/shoot_00001.png")
|
||||
self.assertEqual((self.output / "sets/portraits/shoot_00001.txt").read_text(), "a portrait")
|
||||
|
||||
def test_an_exact_name_keeps_its_place_in_the_run(self):
|
||||
token, _ = self.hold(count=3, exact="portrait")
|
||||
whole = [item["path"] for item in self.keep(self.hold(count=3, exact="other")[0])["saved"]]
|
||||
picked = self.keep(token, [1])["saved"][0]["path"]
|
||||
self.assertEqual(picked, whole[1].replace("other", "portrait"))
|
||||
|
||||
def test_a_request_cannot_name_a_place(self):
|
||||
token, _ = self.hold()
|
||||
for bad in ["../" + token, token + "/..", "", None, 5, token.upper(), "g" * 16]:
|
||||
with self.assertRaises(ValueError):
|
||||
self.keep(bad)
|
||||
for picks in [[3], [-1], ["0"], [True], [0.0]]:
|
||||
with self.assertRaises(ValueError):
|
||||
self.keep(token, picks)
|
||||
self.assertEqual(list(self.output.iterdir()), [])
|
||||
|
||||
def test_pictures_that_are_gone_say_so(self):
|
||||
with self.assertRaises(LookupError) as caught:
|
||||
self.keep("0" * 16)
|
||||
self.assertEqual(str(caught.exception), GONE)
|
||||
token, previews = self.hold(count=1)
|
||||
(self.temp / HELD_FOLDER / previews[0]["filename"]).unlink()
|
||||
with self.assertRaises(LookupError):
|
||||
self.keep(token)
|
||||
|
||||
def test_a_folder_outside_the_output_folder_is_refused_at_the_click_too(self):
|
||||
token, _ = self.hold(count=1)
|
||||
manifest = self.temp / HELD_FOLDER / f"{token}.json"
|
||||
manifest.write_text(manifest.read_text().replace('"output_dir": ""', '"output_dir": "../elsewhere"'))
|
||||
with self.assertRaises(ValueError):
|
||||
self.keep(token)
|
||||
self.assertFalse((Path(self.box.name) / "elsewhere").exists())
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,235 @@
|
||||
"""Mask by Touch on a big picture: SAM 3 is asked about a smaller copy, and
|
||||
what it answers comes back at the picture's own size. ComfyUI's detector and
|
||||
text reader are stood in for by fakes, so this runs without ComfyUI and
|
||||
without a model."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from pathlib import Path
|
||||
import shutil
|
||||
import sys
|
||||
import tempfile
|
||||
import types
|
||||
import unittest
|
||||
|
||||
import torch
|
||||
|
||||
ROOT = Path(__file__).resolve().parent.parent
|
||||
sys.path.insert(0, str(ROOT))
|
||||
if "nodes" in sys.modules and not hasattr(sys.modules["nodes"], "__path__"):
|
||||
del sys.modules["nodes"]
|
||||
|
||||
from nodes import _mask_by_touch_helpers as helpers # noqa: E402
|
||||
from nodes import node_mask_by_touch as node_module # noqa: E402
|
||||
from nodes._mask_by_name_helpers import IF_NOTHING # noqa: E402
|
||||
from nodes._mask_by_touch_helpers import WORK_SIDE, WorkView # noqa: E402
|
||||
|
||||
UI = "ausboss_mask_by_touch"
|
||||
# The big picture, and the long side of the copy SAM 3 is asked about here.
|
||||
H, W = 300, 240
|
||||
SIDE = 60
|
||||
# Things as shares of the picture: (left, top, right, bottom).
|
||||
PERSON = (0.25, 0.1, 0.75, 0.95)
|
||||
FACE = (0.4, 0.15, 0.6, 0.3)
|
||||
THINGS = {"person": [(PERSON, 0.98)], "face": [(FACE, 0.9)]}
|
||||
|
||||
|
||||
def box(thing, height, width):
|
||||
left, top, right, bottom = thing
|
||||
mask = torch.zeros(height, width)
|
||||
mask[round(top * height):round(bottom * height), round(left * width):round(right * width)] = 1.0
|
||||
return mask
|
||||
|
||||
|
||||
def same_area(one, other):
|
||||
one, other = one.bool(), other.bool()
|
||||
return float((one & other).sum()) / max(1.0, float((one | other).sum()))
|
||||
|
||||
|
||||
class SizedSam:
|
||||
"""Stands in for ComfyUI's SAM3_Detect. It answers at the size of the
|
||||
picture it is given, as the real one does."""
|
||||
|
||||
calls: list = []
|
||||
|
||||
@classmethod
|
||||
def execute(cls, model, image, conditioning=None, positive_coords=None, negative_coords=None, threshold=0.5, refine_iterations=2, individual_masks=False):
|
||||
height, width = int(image.shape[1]), int(image.shape[2])
|
||||
call = {"size": (height, width)}
|
||||
cls.calls.append(call)
|
||||
if conditioning is not None:
|
||||
call["word"] = conditioning[0][2].rpartition(":")[0]
|
||||
found = THINGS.get(call["word"], [])
|
||||
masks = torch.stack([box(thing, height, width) for thing, _ in found]) if found else torch.zeros(0, height, width)
|
||||
return types.SimpleNamespace(result=(masks, [[{"score": sure} for _, sure in found]]))
|
||||
call["points"] = points = [(point["x"], point["y"]) for point in json.loads(positive_coords)]
|
||||
answer = torch.zeros(height, width)
|
||||
for thing in (FACE, PERSON): # the smallest thing that holds every point
|
||||
area = box(thing, height, width)
|
||||
if all(0 <= y < height and 0 <= x < width and float(area[y, x]) > 0 for x, y in points):
|
||||
answer = area
|
||||
break
|
||||
return types.SimpleNamespace(result=(answer[None], [[]]))
|
||||
|
||||
|
||||
class Encode:
|
||||
def encode(self, clip, text):
|
||||
return ([["read by", clip, text]],)
|
||||
|
||||
|
||||
class Folders:
|
||||
def __init__(self, folder):
|
||||
self.folder = folder
|
||||
|
||||
def get_temp_directory(self):
|
||||
return self.folder
|
||||
|
||||
def get_save_image_path(self, prefix, folder, _width, _height):
|
||||
counter = 1 + sum(1 for name in Path(folder).iterdir() if name.name.startswith(prefix + "_"))
|
||||
return folder, prefix, counter, "", prefix
|
||||
|
||||
|
||||
class WorkViewTests(unittest.TestCase):
|
||||
def test_a_picture_up_to_the_work_side_is_asked_about_as_it_is(self):
|
||||
frame = torch.rand(1, 40, WORK_SIDE, 3)
|
||||
view = WorkView(frame)
|
||||
self.assertTrue(view.same)
|
||||
self.assertIs(view.frame, frame)
|
||||
self.assertEqual((view.work_height, view.work_width), (40, WORK_SIDE))
|
||||
self.assertEqual(view.point((17, 5)), (17, 5))
|
||||
self.assertEqual(view.full_point((17, 5)), (17, 5))
|
||||
mask = torch.zeros(40, WORK_SIDE, dtype=torch.bool)
|
||||
self.assertIs(view.full(mask), mask)
|
||||
self.assertEqual(view.full_box([1, 2, 3, 4]), [1, 2, 3, 4])
|
||||
|
||||
def test_a_bigger_picture_gets_a_copy_with_the_long_side_at_the_work_side(self):
|
||||
frame = torch.rand(1, WORK_SIDE * 2, WORK_SIDE + 500, 3)
|
||||
view = WorkView(frame)
|
||||
self.assertFalse(view.same)
|
||||
self.assertEqual(view.work_height, WORK_SIDE)
|
||||
self.assertEqual(tuple(view.frame.shape), (1, view.work_height, view.work_width, 3))
|
||||
self.assertEqual(view.work_width, round((WORK_SIDE + 500) / 2))
|
||||
# Averaged down: the copy's mean is the picture's.
|
||||
self.assertAlmostEqual(float(view.frame.mean()), float(frame.mean()), places=3)
|
||||
|
||||
def test_a_place_on_the_copy_leads_back_to_itself(self):
|
||||
for height, width in ((4096, 4096), (3840, 2160), (2049, 8192), (5000, 3333)):
|
||||
view = WorkView(torch.zeros(1, height, width, 1))
|
||||
for x in (0, 1, 7, view.work_width // 3, view.work_width - 2, view.work_width - 1):
|
||||
for y in (0, 1, 5, view.work_height // 2, view.work_height - 1):
|
||||
full = view.full_point((x, y))
|
||||
self.assertTrue(0 <= full[0] < width and 0 <= full[1] < height)
|
||||
self.assertEqual(view.point(full), (x, y), f"{width}x{height} at {x},{y}")
|
||||
self.assertEqual(view.point((width - 1, height - 1)), (view.work_width - 1, view.work_height - 1))
|
||||
|
||||
def test_an_answer_comes_back_at_the_pictures_own_size(self):
|
||||
view = WorkView(torch.zeros(1, 4096, 3000, 3))
|
||||
small = box(FACE, view.work_height, view.work_width) > 0.5
|
||||
full = view.full(small)
|
||||
self.assertEqual(tuple(full.shape), (4096, 3000))
|
||||
self.assertEqual(full.dtype, torch.bool)
|
||||
self.assertGreater(same_area(full, box(FACE, 4096, 3000)), 0.98)
|
||||
x0, y0, x1, y1 = view.full_box(helpers.box_of(small))
|
||||
wanted = helpers.box_of(box(FACE, 4096, 3000) > 0.5)
|
||||
for got, want in zip((x0, y0, x1, y1), wanted):
|
||||
self.assertLessEqual(abs(got - want), 3)
|
||||
|
||||
|
||||
class BigPictureTests(unittest.TestCase):
|
||||
def setUp(self):
|
||||
SizedSam.calls = []
|
||||
node_module._MEMORY.clear()
|
||||
node_module._WORDS.clear()
|
||||
node_module._TRUST.clear()
|
||||
node_module._UNSTEADY.clear()
|
||||
node_module._NOTED.clear()
|
||||
self.swap(node_module, "_core_nodes", lambda: {"SAM3_Detect": SizedSam, "CLIPTextEncode": Encode})
|
||||
self.swap(node_module, "SCOPE_SPACING", 12)
|
||||
self.swap(node_module, "SCOPE_SECONDS_BLIND", 1e9)
|
||||
self.swap(helpers, "WORK_SIDE", SIDE)
|
||||
folder = tempfile.mkdtemp(prefix="ausboss_touch_big_")
|
||||
self.addCleanup(shutil.rmtree, folder, True)
|
||||
self.swap(node_module, "folder_paths", Folders(folder))
|
||||
self.picture = torch.rand(1, H, W, 3) * 0.9
|
||||
self.node = node_module.AusBossMaskByTouch()
|
||||
|
||||
def swap(self, module, attribute, value):
|
||||
self.addCleanup(setattr, module, attribute, getattr(module, attribute))
|
||||
setattr(module, attribute, value)
|
||||
|
||||
def run_node(self, text="", image=None, **overrides):
|
||||
values = {"grow": 0, "soften": 0.0, "smooth": 0, "fill_holes": False, "if_nothing": IF_NOTHING[1]}
|
||||
values.update(overrides)
|
||||
return self.node.pick(self.picture if image is None else image, "sam3 model", "sam3 clip", text, **values)
|
||||
|
||||
def touches(self, *items):
|
||||
return json.dumps({"v": 1, "items": list(items)})
|
||||
|
||||
def test_sam3_is_never_given_the_big_picture(self):
|
||||
self.run_node(self.touches({"t": "tap", "op": 1, "x": 120, "y": 66}, {"t": "all", "op": 1, "w": "person"}))
|
||||
self.assertTrue(SizedSam.calls)
|
||||
for call in SizedSam.calls:
|
||||
self.assertEqual(call["size"], (SIDE, 48), "the copy: 300 x 240 with its long side at 60")
|
||||
|
||||
def test_a_tap_takes_the_thing_at_its_place_and_the_mask_is_full_size(self):
|
||||
out = self.run_node(self.touches({"t": "tap", "op": 1, "x": 120, "y": 66})) # on the face
|
||||
mask = out["result"][0]
|
||||
self.assertEqual(tuple(mask.shape), (1, H, W))
|
||||
self.assertGreater(same_area(mask[0] > 0.5, box(FACE, H, W)), 0.9)
|
||||
# The point was asked where it lies on the copy.
|
||||
self.assertEqual(SizedSam.calls[0]["points"], [(120 * 48 // W, 66 * SIDE // H)])
|
||||
|
||||
def test_a_word_takes_every_thing_of_that_name_at_full_size(self):
|
||||
out = self.run_node(self.touches({"t": "all", "op": 1, "w": "person"}, {"t": "all", "op": 0, "w": "face"}))
|
||||
mask = out["result"][0][0] > 0.5
|
||||
wanted = (box(PERSON, H, W) > 0.5) & ~(box(FACE, H, W) > 0.5)
|
||||
self.assertGreater(same_area(mask, wanted), 0.93)
|
||||
self.assertEqual(out["result"][2], "person")
|
||||
|
||||
def test_a_stroke_is_asked_about_on_the_copy(self):
|
||||
stroke = {"t": "find", "op": 1, "r": 8, "pts": [[110, 60], [130, 70]]} # across the face
|
||||
out = self.run_node(self.touches(stroke))
|
||||
self.assertGreater(same_area(out["result"][0][0] > 0.5, box(FACE, H, W)), 0.9)
|
||||
for call in SizedSam.calls:
|
||||
for x, y in call.get("points", []):
|
||||
self.assertTrue(0 <= x < 48 and 0 <= y < SIDE)
|
||||
|
||||
def test_what_a_look_lights_up_is_what_a_click_on_it_takes(self):
|
||||
note = self.run_node(ask="scope")["ui"][UI][0]
|
||||
self.assertEqual((note["width"], note["height"]), (W, H))
|
||||
pieces = {piece["name"]: piece for piece in note["pieces"]["list"] if piece["name"]}
|
||||
self.assertEqual(set(pieces), {"person", "face"})
|
||||
face = pieces["face"]
|
||||
# Its own point is a pixel of the picture, inside the thing, and its box is the picture's.
|
||||
self.assertGreater(float(box(FACE, H, W)[face["y"], face["x"]]), 0)
|
||||
for got, want in zip(face["box"], helpers.box_of(box(FACE, H, W) > 0.5)):
|
||||
self.assertLessEqual(abs(got - want), 6)
|
||||
self.assertAlmostEqual(face["share"], float(box(FACE, H, W).mean()), delta=0.01)
|
||||
# The face picks it as a named touch at that point: nothing is asked again.
|
||||
before = len(SizedSam.calls)
|
||||
out = self.run_node(self.touches({"t": "name", "op": 1, "w": "face", "x": face["x"], "y": face["y"]}), ask="look")
|
||||
self.assertEqual(len(SizedSam.calls), before)
|
||||
self.assertGreater(same_area(out["result"][0][0] > 0.5, box(FACE, H, W)), 0.9)
|
||||
# A point piece is remembered under the picture's own pixel as well.
|
||||
spot = next((piece for piece in note["pieces"]["list"] if not piece["name"]), None)
|
||||
if spot is not None:
|
||||
self.run_node(self.touches({"t": "tap", "op": 1, "x": spot["x"], "y": spot["y"]}), ask="look")
|
||||
self.assertEqual(len(SizedSam.calls), before)
|
||||
|
||||
def test_a_small_picture_is_asked_about_as_it_is(self):
|
||||
small = torch.rand(1, 50, 40, 3) * 0.9
|
||||
out = self.run_node(self.touches({"t": "tap", "op": 1, "x": 20, "y": 11}), image=small)
|
||||
self.assertEqual([call["size"] for call in SizedSam.calls], [(50, 40)])
|
||||
self.assertTrue(torch.equal(out["result"][0][0] > 0.5, box(FACE, 50, 40) > 0.5), "exactly what SAM 3 answered")
|
||||
|
||||
def test_paint_stays_exactly_as_painted_on_the_big_picture(self):
|
||||
stroke = [[20, 20], [200, 20]]
|
||||
out = self.run_node(self.touches({"t": "paint", "op": 1, "r": 5, "pts": stroke}))
|
||||
wanted = helpers.stroke_area(tuple(tuple(point) for point in stroke), 5.0, H, W)
|
||||
self.assertTrue(torch.equal(out["result"][0][0] > 0.5, wanted))
|
||||
self.assertEqual(SizedSam.calls, [], "paint asks SAM 3 nothing")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,71 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from pathlib import Path
|
||||
import sys
|
||||
import unittest
|
||||
|
||||
ROOT = Path(__file__).resolve().parent.parent
|
||||
sys.path.insert(0, str(ROOT))
|
||||
if "nodes" in sys.modules and not hasattr(sys.modules["nodes"], "__path__"):
|
||||
del sys.modules["nodes"]
|
||||
|
||||
from nodes._prompt_variations_helpers import pick_variations
|
||||
|
||||
|
||||
class PickVariationsTests(unittest.TestCase):
|
||||
def test_one_pick_from_each_set(self):
|
||||
text, picks = pick_variations("a {red|blue} {sweater|jacket|coat}", 7)
|
||||
self.assertEqual(len(picks), 2)
|
||||
self.assertIn(picks[0], ("red", "blue"))
|
||||
self.assertIn(picks[1], ("sweater", "jacket", "coat"))
|
||||
self.assertEqual(text, f"a {picks[0]} {picks[1]}")
|
||||
|
||||
def test_same_seed_same_prompt(self):
|
||||
prompt = "{a|b|c|d|e|f} {g|h|i|j|k|l} {m|n|o|p|q|r}"
|
||||
self.assertEqual(pick_variations(prompt, 123), pick_variations(prompt, 123))
|
||||
|
||||
def test_seeds_give_different_mixes(self):
|
||||
prompt = "{a|b|c|d|e|f} {g|h|i|j|k|l} {m|n|o|p|q|r}"
|
||||
mixes = {pick_variations(prompt, seed)[0] for seed in range(50)}
|
||||
self.assertGreater(len(mixes), 30)
|
||||
|
||||
def test_every_option_can_come_up(self):
|
||||
seen = {pick_variations("{one|two|three}", seed)[0] for seed in range(200)}
|
||||
self.assertEqual(seen, {"one", "two", "three"})
|
||||
|
||||
def test_text_without_sets_is_unchanged(self):
|
||||
for text in ("", "plain words, two spaces\nand a line", '{"subject": "a cat", "light": {"key": "soft"}}'):
|
||||
self.assertEqual(pick_variations(text, 5), (text, []))
|
||||
|
||||
def test_spaces_around_options_are_trimmed(self):
|
||||
text, _ = pick_variations("a { red | red } hat", 1)
|
||||
self.assertEqual(text, "a red hat")
|
||||
|
||||
def test_empty_choice_leaves_no_double_space(self):
|
||||
texts = {pick_variations("a {|striped} shirt", seed)[0] for seed in range(40)}
|
||||
self.assertEqual(texts, {"a shirt", "a striped shirt"})
|
||||
|
||||
def test_sets_inside_sets(self):
|
||||
texts = {pick_variations("{a {red|blue} hat|a scarf}", seed)[0] for seed in range(80)}
|
||||
self.assertEqual(texts, {"a red hat", "a blue hat", "a scarf"})
|
||||
|
||||
def test_choice_inside_json_braces(self):
|
||||
text, picks = pick_variations('{"top": "{tank top|tank top}"}', 3)
|
||||
self.assertEqual(text, '{"top": "tank top"}')
|
||||
self.assertEqual(picks, ["tank top"])
|
||||
|
||||
def test_large_seed(self):
|
||||
text, _ = pick_variations("{x|y}", 0xFFFFFFFFFFFFFFFF)
|
||||
self.assertIn(text, ("x", "y"))
|
||||
|
||||
|
||||
class NodeTests(unittest.TestCase):
|
||||
def test_node_returns_a_tuple(self):
|
||||
from nodes.node_prompt_variations import AusBossPromptVariations
|
||||
|
||||
(prompt,) = AusBossPromptVariations().pick("no sets here", 0)
|
||||
self.assertEqual(prompt, "no sets here")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,79 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from pathlib import Path
|
||||
import sys
|
||||
import tempfile
|
||||
import unittest
|
||||
|
||||
ROOT = Path(__file__).resolve().parent.parent
|
||||
sys.path.insert(0, str(ROOT))
|
||||
if "nodes" in sys.modules and not hasattr(sys.modules["nodes"], "__path__"):
|
||||
del sys.modules["nodes"]
|
||||
|
||||
from nodes._single_copy import other_copy_that_loads
|
||||
|
||||
|
||||
def make_copy(root: Path, name: str, version: str, guarded: bool = True, pack: str = "ausboss-nodes") -> str:
|
||||
folder = root / name
|
||||
(folder / "nodes").mkdir(parents=True)
|
||||
(folder / "__init__.py").write_text("")
|
||||
(folder / "pyproject.toml").write_text(f'[project]\nname = "{pack}"\nversion = "{version}"\n')
|
||||
if guarded:
|
||||
(folder / "nodes" / "_single_copy.py").write_text("")
|
||||
return str(folder)
|
||||
|
||||
|
||||
def loads(folder: str, root: Path) -> bool:
|
||||
return other_copy_that_loads(folder, [str(root)]) is None
|
||||
|
||||
|
||||
class SingleCopyTests(unittest.TestCase):
|
||||
def setUp(self):
|
||||
self._tmp = tempfile.TemporaryDirectory()
|
||||
self.root = Path(self._tmp.name).resolve()
|
||||
|
||||
def tearDown(self):
|
||||
self._tmp.cleanup()
|
||||
|
||||
def test_a_lone_copy_loads(self):
|
||||
a = make_copy(self.root, "ComfyUI-AusBoss", "2.24.0")
|
||||
make_copy(self.root, "some-other-pack", "9.0.0", pack="other")
|
||||
self.assertTrue(loads(a, self.root))
|
||||
|
||||
def test_exactly_one_of_two_copies_loads_and_it_is_the_newer(self):
|
||||
a = make_copy(self.root, "ComfyUI-AusBoss", "2.24.0")
|
||||
b = make_copy(self.root, "ausboss-nodes", "2.25.0")
|
||||
self.assertFalse(loads(a, self.root))
|
||||
self.assertTrue(loads(b, self.root))
|
||||
self.assertEqual(other_copy_that_loads(a, [str(self.root)])["version"], "2.25.0")
|
||||
|
||||
def test_same_version_still_picks_one(self):
|
||||
a = make_copy(self.root, "ComfyUI-AusBoss", "2.24.0")
|
||||
b = make_copy(self.root, "ausboss-nodes", "2.24.0")
|
||||
self.assertEqual([loads(a, self.root), loads(b, self.root)].count(True), 1)
|
||||
|
||||
def test_yields_to_an_older_copy_that_cannot_step_aside(self):
|
||||
new = make_copy(self.root, "ComfyUI-AusBoss", "2.24.0")
|
||||
make_copy(self.root, "ausboss-nodes", "2.22.0", guarded=False)
|
||||
self.assertFalse(loads(new, self.root))
|
||||
|
||||
def test_a_disabled_copy_does_not_count(self):
|
||||
a = make_copy(self.root, "ComfyUI-AusBoss", "2.24.0")
|
||||
make_copy(self.root, "ausboss-nodes.disabled", "2.25.0")
|
||||
self.assertTrue(loads(a, self.root))
|
||||
|
||||
def test_a_link_to_the_same_folder_is_one_copy(self):
|
||||
a = make_copy(self.root, "ComfyUI-AusBoss", "2.24.0")
|
||||
(self.root / "link").symlink_to(a)
|
||||
self.assertTrue(loads(a, self.root))
|
||||
self.assertTrue(loads(str(self.root / "link"), self.root))
|
||||
|
||||
def test_ten_is_newer_than_nine(self):
|
||||
a = make_copy(self.root, "ComfyUI-AusBoss", "2.9.0")
|
||||
b = make_copy(self.root, "ausboss-nodes", "2.10.0")
|
||||
self.assertFalse(loads(a, self.root))
|
||||
self.assertTrue(loads(b, self.root))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,100 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from pathlib import Path
|
||||
import sys
|
||||
import unittest
|
||||
|
||||
ROOT = Path(__file__).resolve().parent.parent
|
||||
sys.path.insert(0, str(ROOT))
|
||||
if "nodes" in sys.modules and not hasattr(sys.modules["nodes"], "__path__"):
|
||||
del sys.modules["nodes"]
|
||||
|
||||
from nodes._text_lines_helpers import split_items, strip_marker
|
||||
from nodes.node_text_lines import AusBossTextLines
|
||||
|
||||
|
||||
class SplitItemsTests(unittest.TestCase):
|
||||
def test_each_line_is_an_item_and_empty_lines_are_skipped(self):
|
||||
text = "a red car\n\n \na blue boat \r\n\tgreen hills\n"
|
||||
self.assertEqual(split_items(text), ["a red car", "a blue boat", "green hills"])
|
||||
|
||||
def test_blank_line_keeps_lines_together(self):
|
||||
text = "first line\nsecond line\n\n\nnext item\n \nlast"
|
||||
self.assertEqual(
|
||||
split_items(text, "blank line"),
|
||||
["first line\nsecond line", "next item", "last"],
|
||||
)
|
||||
|
||||
def test_limit_keeps_the_first_items(self):
|
||||
self.assertEqual(split_items("a\nb\nc\nd", limit=2), ["a", "b"])
|
||||
self.assertEqual(split_items("a\nb", limit=0), ["a", "b"])
|
||||
self.assertEqual(split_items("a\nb", limit=9), ["a", "b"])
|
||||
|
||||
def test_numbers_and_bullets_come_off(self):
|
||||
cases = {
|
||||
"1. close-up, soft smile": "close-up, soft smile",
|
||||
"2) medium shot": "medium shot",
|
||||
"#3: from above": "from above",
|
||||
"- a bullet": "a bullet",
|
||||
"* a star": "a star",
|
||||
"• a dot": "a dot",
|
||||
"[prompt 4] from the side": "from the side",
|
||||
"【prompt_5】 looking back": "looking back",
|
||||
"(6) arms crossed": "arms crossed",
|
||||
"Prompt 7: laughing": "laughing",
|
||||
"Shot 8 - wide": "wide",
|
||||
"12. twelve": "twelve",
|
||||
}
|
||||
for line, want in cases.items():
|
||||
self.assertEqual(strip_marker(line), want, line)
|
||||
|
||||
def test_text_that_only_starts_like_a_marker_stays(self):
|
||||
for line in [
|
||||
"2 women on a beach",
|
||||
"1.5x zoom on the face",
|
||||
"10:30 at night, a diner",
|
||||
"3D render of a robot",
|
||||
"*bold* light from the left",
|
||||
"-90 degrees, seen from below",
|
||||
"1920s street, a man in a hat",
|
||||
"shot from below",
|
||||
"Image 1: the woman, in the coat from image 2",
|
||||
]:
|
||||
self.assertEqual(strip_marker(line), line, line)
|
||||
|
||||
def test_numbering_stays_when_asked(self):
|
||||
self.assertEqual(split_items("1. a\n- b", remove_numbering=False), ["1. a", "- b"])
|
||||
|
||||
def test_a_line_that_is_only_a_marker_is_skipped(self):
|
||||
self.assertEqual(split_items("1.\n2. real\n-"), ["real"])
|
||||
|
||||
def test_nothing_gives_nothing(self):
|
||||
self.assertEqual(split_items(""), [])
|
||||
self.assertEqual(split_items(None), [])
|
||||
self.assertEqual(split_items("\n \n"), [])
|
||||
|
||||
|
||||
class TextLinesNodeTests(unittest.TestCase):
|
||||
def test_outputs_lines_numbers_and_count(self):
|
||||
lines, numbers, count = AusBossTextLines().split("1. one\n2. two\n\n3. three")
|
||||
self.assertEqual(lines, ["one", "two", "three"])
|
||||
self.assertEqual(numbers, [1, 2, 3])
|
||||
self.assertEqual(count, 3)
|
||||
|
||||
def test_the_line_and_number_outputs_are_lists(self):
|
||||
self.assertEqual(AusBossTextLines.OUTPUT_IS_LIST, (True, True, False))
|
||||
self.assertEqual(AusBossTextLines.RETURN_NAMES, ("line", "number", "count"))
|
||||
|
||||
def test_no_lines_stops_the_run_with_plain_words(self):
|
||||
with self.assertRaises(ValueError) as caught:
|
||||
AusBossTextLines().split(" \n")
|
||||
self.assertIn("no lines to run", str(caught.exception))
|
||||
|
||||
def test_inputs_and_their_order(self):
|
||||
inputs = AusBossTextLines.INPUT_TYPES()
|
||||
self.assertEqual(list(inputs), ["required"])
|
||||
self.assertEqual(list(inputs["required"]), ["text", "split", "remove_numbering", "limit"])
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -60,6 +60,7 @@ import {
|
||||
thinStroke,
|
||||
touchFits,
|
||||
} from "../js/shared/touch_math.mjs";
|
||||
import { OUT_OF_MEMORY, goesBehind as goesBehindAgain, piecesHolding, stopMessage, typedName } from "../js/shared/touch_math.mjs";
|
||||
import {
|
||||
GROUP_MAX,
|
||||
areaKey,
|
||||
@@ -1062,3 +1063,68 @@ test("picks the node moved onto a new picture are taken only while nothing was p
|
||||
assert.equal(movedTouches({ moved: { from: 7, touches: after } }, before), null, "a note that is not this shape moves nothing");
|
||||
assert.equal(movedTouches(null, before), null);
|
||||
});
|
||||
|
||||
test("taking away from a spot that is not picked never offers a thing that holds a whole pick", () => {
|
||||
const { width, masks, planes, sizes, ranks } = scene();
|
||||
// The jacket is picked. The pointer is on the person's legs, which are not.
|
||||
const inside = piecesInside(planes, 5, masks.jacket);
|
||||
const holds = piecesHolding(planes, 5, [masks.jacket]);
|
||||
assert.deepEqual([...holds], [0, 1, 1, 1, 1], "the jacket, the dress over it, the person and the wall hold all of it");
|
||||
const legs = 32 * width + 14;
|
||||
assert.equal(pieceAt(planes, legs, { sizes, inside, ranks, adding: false }), 3, "as before, when nothing is said about the spot");
|
||||
assert.equal(pieceAt(planes, legs, { sizes, inside, ranks, adding: false, here: false, holds }), -1, "not the person: the jacket would go with it");
|
||||
// On the jacket itself the spot is picked: the jacket is offered.
|
||||
const onJacket = 14 * width + 14;
|
||||
assert.equal(pieceAt(planes, onJacket, { sizes, inside, ranks, adding: false, here: true, holds }), 1);
|
||||
// The person is picked, and the pointer is on the jacket: the jacket is a part of the pick, and is offered.
|
||||
const person = piecesInside(planes, 5, masks.person);
|
||||
const personHolds = piecesHolding(planes, 5, [masks.person]);
|
||||
assert.equal(pieceAt(planes, onJacket, { sizes, inside: person, ranks, adding: false, here: true, holds: personHolds }), 1);
|
||||
// A thing a pick only laps over is still offered from a spot that is not picked.
|
||||
const lap = Uint8Array.from(masks.person, (value, index) => (value && Math.floor(index / width) >= 30 ? 1 : 0));
|
||||
const lapInside = piecesInside(planes, 5, lap);
|
||||
const lapHolds = piecesHolding(planes, 5, [masks.wall]);
|
||||
assert.equal(lapHolds[3], 0, "the person does not hold a pick of the wall");
|
||||
assert.equal(pieceAt(planes, onJacket + 0, { sizes, inside: lapInside, ranks, adding: false, here: false, holds: lapHolds }), 3, "the person has picked area in it and holds no whole pick");
|
||||
// Adding is not changed by any of it.
|
||||
assert.equal(pieceAt(planes, legs, { sizes, inside, ranks, here: false, holds }), 3);
|
||||
assert.deepEqual([...piecesHolding(planes, 5, [null, new Uint8Array(masks.jacket.length)])], [0, 0, 0, 0, 0], "no area, nothing held");
|
||||
});
|
||||
|
||||
test("a typed name is added or taken away by the sign, and a sign in the text wins", () => {
|
||||
assert.deepEqual(typedName("face"), { word: "face", add: true });
|
||||
assert.deepEqual(typedName("face", true), { word: "face", add: false });
|
||||
assert.deepEqual(typedName("-face"), { word: "face", add: false });
|
||||
assert.deepEqual(typedName(" - fur collar "), { word: "fur collar", add: false });
|
||||
assert.deepEqual(typedName("\u2212hair"), { word: "hair", add: false }, "the minus a phone types");
|
||||
assert.deepEqual(typedName("+hat", true), { word: "hat", add: true });
|
||||
assert.deepEqual(typedName("t-shirt", false), { word: "t-shirt", add: true }, "a dash inside a word is the word's");
|
||||
assert.deepEqual(typedName("-"), { word: "", add: false });
|
||||
assert.deepEqual(typedName(null), { word: "", add: true });
|
||||
});
|
||||
|
||||
test("looks at a picture the node no longer has are taken out of the queue before the next one is sent", () => {
|
||||
const here = dirname(fileURLToPath(import.meta.url));
|
||||
const stage = readFileSync(join(here, "..", "js", "shared", "touch_stage.mjs"), "utf8");
|
||||
const entry = readFileSync(join(here, "..", "js", "mask_by_touch", "index.js"), "utf8");
|
||||
const run = stage.slice(stage.indexOf("function requestRun"), stage.indexOf("function settle(face)"));
|
||||
// the stale ones are counted before this run joins them, and dropped before it is sent
|
||||
assert.ok(run.indexOf("const stale = face.flying.filter(") < run.indexOf("face.flying.push(run)"));
|
||||
assert.match(run, /other\.queued && other\.epoch !== run\.epoch/);
|
||||
assert.ok(run.indexOf("face.host.dropWaiting()") < run.indexOf("face.host.run("));
|
||||
// only this node's own runs (they carry an `ask`) that still wait: never the one that runs, never a workflow run
|
||||
const drop = entry.slice(entry.indexOf("async function dropWaiting"), entry.indexOf("// The picture goes above the settings"));
|
||||
assert.match(drop, /queue_pending/);
|
||||
assert.doesNotMatch(drop, /queue_running|interrupt/);
|
||||
assert.match(drop, /class_type === NODE_CLASS && String\(step\.inputs\?\.ask \?\? ""\) !== ""/);
|
||||
// a look still goes behind looks at an older picture, as before
|
||||
assert.equal(goesBehindAgain([{ ask: "scope", queued: true, epoch: 1 }], "scope", 2), true);
|
||||
});
|
||||
|
||||
test("a full graphics card is said in one plain line, anything else as it came", () => {
|
||||
const comfy = "Allocation on device 0 would exceed allowed memory. (out of memory)\nCurrently allocated : 2.07 GiB\nRequested : 40.50 MiB";
|
||||
assert.equal(stopMessage(comfy), OUT_OF_MEMORY);
|
||||
assert.equal(stopMessage("torch.OutOfMemoryError: CUDA"), OUT_OF_MEMORY);
|
||||
assert.equal(stopMessage("Nothing is picked on this picture."), "Nothing is picked on this picture.");
|
||||
assert.equal(stopMessage(null), "");
|
||||
});
|
||||
|
||||
@@ -39,7 +39,7 @@ test("a stopped run's message stays up until the next result, then goes", () =>
|
||||
const { face, setNote, shown } = noteFace();
|
||||
const stopped = SOURCE.match(/face\.stopped = \(message\) => \{([\s\S]*?)\n {2}\};/)?.[1];
|
||||
assert.ok(stopped, "touch_stage.mjs has no face.stopped");
|
||||
assert.match(stopped, /setNote\(face, String\(message \?\? ""\)\);\s*(?:\/\/[^\n]*\n\s*)*face\.stopNote = true;/,
|
||||
assert.match(stopped, /setNote\(face, stopMessage\(message\)\);\s*(?:\/\/[^\n]*\n\s*)*face\.stopNote = true;/,
|
||||
"a stop marks its message as a stop message, after it is set");
|
||||
// What face.stopped does to the note.
|
||||
setNote(face, "Mask by Touch: nothing is picked.");
|
||||
|
||||
Reference in New Issue
Block a user