Merge main: FlashVSR Upscale and Prompt Variations share the zh locale file
This commit is contained in:
ausboss
@@ -67,6 +67,14 @@ All notable changes to ComfyUI-AusBoss are documented here.
|
||||
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
|
||||
|
||||
@@ -212,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.
|
||||
@@ -417,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.
|
||||
|
||||
|
||||
|
||||
@@ -438,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.
|
||||
|
||||
|
||||
@@ -23,6 +23,7 @@ NODE_MODULES = [
|
||||
"node_color_match",
|
||||
"node_compare",
|
||||
"node_edit_lock",
|
||||
"node_flashvsr_upscale",
|
||||
"node_frame_interpolate",
|
||||
"node_free_memory",
|
||||
"node_h3_prompt_writer",
|
||||
|
||||
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,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.
|
||||
@@ -38,6 +38,7 @@ 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`) |
|
||||
|
||||
@@ -290,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: [
|
||||
|
||||
@@ -1682,6 +1682,32 @@
|
||||
"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": "每次运行都会变化的提示词框。把选项写在花括号里，例如 {红色|蓝色|绿色} 的 {毛衣|夹克}，节点会从每组里选一个。种子决定选哪个：同一个种子得到同一个提示词，每次换新种子就得到新的组合。没有花括号的文字原样输出。",
|
||||
|
||||
@@ -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,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"]
|
||||
@@ -58,6 +58,7 @@ PUBLIC_NODE_IDS = {
|
||||
"AUSBOSS_NODES_CropForInpaint",
|
||||
"AUSBOSS_NODES_EditLock",
|
||||
"AUSBOSS_NODES_EditLockRestore",
|
||||
"AUSBOSS_NODES_FlashVSRUpscale",
|
||||
"AUSBOSS_NODES_Float",
|
||||
"AUSBOSS_NODES_FrameInterpolate",
|
||||
"AUSBOSS_NODES_FreeMemory",
|
||||
|
||||
Vendored
+17
@@ -160,6 +160,23 @@
|
||||
]
|
||||
]
|
||||
},
|
||||
"AUSBOSS_NODES_FlashVSRUpscale": {
|
||||
"required": [
|
||||
"frames",
|
||||
"model",
|
||||
"scale",
|
||||
"color_match"
|
||||
],
|
||||
"optional": [
|
||||
"vae"
|
||||
],
|
||||
"outputs": [
|
||||
[
|
||||
"IMAGE",
|
||||
"frames"
|
||||
]
|
||||
]
|
||||
},
|
||||
"AUSBOSS_NODES_Float": {
|
||||
"required": [
|
||||
"value"
|
||||
|
||||
@@ -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()
|
||||
Reference in New Issue
Block a user