nodes/FuouM-ComfyUI-StyleTransferPlus
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

finalize unist

Browse Source
This commit is contained in:
FuouM
2024-08-14 01:10:53 +07:00
parent 40c8f35dc8
commit aac3599bcd
9 changed files with 522 additions and 42 deletions
Download Patch File Download Diff File Expand all files Collapse all files
README.md
+12 -2
View File
@@ -22,8 +22,8 @@ All nodes support batched input (i.e video) but is generally not recommended. On
| ![neural_neighbor](example_outputs/neural_neighbor.png) | ![cast](example_outputs/cast.png) |
| **EFDM (2022) - Fast** | **MicroAST (2023) - Fast** |
| ![efdm](example_outputs/efdm.png) | ![microast](example_outputs/microast.png) |
| **UniST (2023)** | **X** |
| X | X |
| **UniST (2023) - Medium - Only squares** | **X** |
| ![efdm](example_outputs/unist_image.png) | X |
### Neural Neighbor
@@ -119,6 +119,16 @@ models/unist/dec_r41.pth
models/unist/vgg_r41.pth
```
The model only works with square images, so inputs are resized to `size x size`. `do_crop` will resize height, then crop to square.
The Video node is "more native" than the Image node for batched images (video) inputs. The model works with `batch=3` (3 consecutive frames), so we split the video into such.
| UniST | UniST Video |
| :-----------------------------------------------------------: | :--------------------------------------------------------------: |
| [workflow_unist_image.json](workflows/workflow_microast.json) | [workflow_unist_video.json](workflows/workflow_unist_video.json) |
| ![wf_unist_image](workflows/wf_unist_image.png) | ![wf_microast](example_outputs/example_unist_video.png) |
||(Not a workflow-embedded image)|
## Extra nodes
example_outputs/example_unist_video.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 614 KiB

example_outputs/unist_image.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 769 KiB

module_unist/models.py
+4 -5
View File
@@ -240,8 +240,8 @@ class Encoder_cross(nn.Module):
)
def forward(self, dec_output, enc_output):
print(f"{dec_output.shape=}")
print(f"{enc_output.shape=}")
# print(f"{dec_output.shape=}")
# print(f"{enc_output.shape=}")
for enc_layer in self.layer_stack:
enc_output = enc_layer(dec_output, enc_output)
return enc_output
@@ -259,10 +259,9 @@ class Decoder_style_tranfer(nn.Module):
)
def forward(self, dec_output, enc_output):
print(f"{dec_output.shape=}")
print(f"{enc_output.shape=}")
# print(f"{dec_output.shape=}")
# print(f"{enc_output.shape=}")
for dec_layer in self.layer_stack:
# print(dec_layer)
dec_output = dec_layer(dec_output, enc_output)
_, l, _ = dec_output.size()
module_unist/unist_model.py
+12 -28
View File
@@ -1,4 +1,3 @@
import numpy as np
import torch
import torchvision.transforms as transforms
@@ -22,13 +21,7 @@ def inference_unist(
if content_type == "video":
style = match_shape(style, content[0])
print(f"{content.shape=}")
print(f"{style.shape=}")
y_hat = network.forward(
content, style, content_type, tab="inference"
) # [B, C, H, W]
print(f"{y_hat.shape=}")
y_hat = network.forward(content, style, content_type, tab="inference")
output = un_normalize_batch(y_hat) # [B, C, H, W]
return output
@@ -58,24 +51,15 @@ def un_normalize_batch(tensor, mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]):
def match_shape(a: torch.Tensor, b: torch.Tensor):
B_a, C, H, W = a.shape
B_b, _, _, _ = b.shape
B_b = b.shape[0]
if B_a == 1 and B_b > 1:
# Expand a along the batch dimension to match the shape of b
expanded_a = a.expand(B_b, -1, -1, -1)
elif B_a > 1 and B_a < B_b:
# Repeat a along the batch dimension to match the shape of b
repeated_a = a.repeat(B_b // B_a, 1, 1, 1)
if B_b % B_a != 0:
remaining = B_b % B_a
repeated_a = torch.cat([repeated_a, a[-remaining:]], dim=0)
expanded_a = repeated_a
elif B_a > B_b:
# Trim a along the batch dimension to match the shape of b
trimmed_a = a[:B_b]
expanded_a = trimmed_a
else:
# If B_a == B_b, no change is needed
expanded_a = a
return expanded_a
if B_a == B_b:
return a
elif B_a == 1:
return a.expand(B_b, -1, -1, -1)
elif B_a < B_b:
repeat_factor = B_b // B_a
remainder = B_b % B_a
return torch.cat([a.repeat(repeat_factor, 1, 1, 1), a[:remainder]], dim=0)
else: # B_a > B_b
return a[:B_b]
run.py
+20 -7
View File
@@ -394,7 +394,6 @@ class UniST:
style_img: torch.Tensor,
do_crop: bool,
size: int,
src_images: torch.Tensor | None = None,
):
print(f"{src_img.shape=}")
print(f"{style_img.shape=}")
@@ -461,7 +460,6 @@ class UniST_Video:
style_video: torch.Tensor,
do_crop: bool,
size: int,
src_images: torch.Tensor | None = None,
):
print(f"{src_video.shape=}")
print(f"{style_video.shape=}")
@@ -476,8 +474,11 @@ class UniST_Video:
network.eval()
network.to(device)
num_frames = src_video.size(0)
batch_size = 3
result = []
params = {
"src": src_video.permute(0, 3, 1, 2),
"src_style": style_video.permute(0, 3, 1, 2),
"device": device,
"size": size,
@@ -485,11 +486,23 @@ class UniST_Video:
"network": network,
"content_type": "video",
}
with torch.no_grad():
res_tensor = inference_unist(**params).permute(0, 2, 3, 1)
return (res_tensor, )
pbar = ProgressBar(num_frames)
with torch.no_grad():
for i in tqdm.tqdm(range(0, num_frames, batch_size), "Generating"):
batch_end = min(i + batch_size, num_frames)
src_batch = src_video[i:batch_end].permute(0, 3, 1, 2)
params["src"] = src_batch
res_tensor = inference_unist(**params).permute(0, 2, 3, 1)
result.append(res_tensor)
pbar.update_absolute(batch_end, num_frames)
final_result = torch.cat(result, dim=0)
return (final_result,)
def serialize_floats(lst: list[float]):
workflows/wf_unist_image.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 545 KiB

workflows/workflow_unist_image.json
+193
View File
@@ -0,0 +1,193 @@
{
"last_node_id": 13,
"last_link_id": 21,
"nodes": [
{
"id": 12,
"type": "LoadImage",
"pos": [
192,
179
],
"size": {
"0": 315,
"1": 314
},
"flags": {},
"order": 0,
"mode": 0,
"outputs": [
{
"name": "IMAGE",
"type": "IMAGE",
"links": [
20
],
"shape": 3,
"slot_index": 0
},
{
"name": "MASK",
"type": "MASK",
"links": null,
"shape": 3
}
],
"properties": {
"Node name for S&R": "LoadImage"
},
"widgets_values": [
"006.jpg",
"image"
]
},
{
"id": 3,
"type": "LoadImage",
"pos": [
195,
548
],
"size": {
"0": 315,
"1": 314
},
"flags": {},
"order": 1,
"mode": 0,
"outputs": [
{
"name": "IMAGE",
"type": "IMAGE",
"links": [
19
],
"shape": 3,
"slot_index": 0
},
{
"name": "MASK",
"type": "MASK",
"links": null,
"shape": 3
}
],
"properties": {
"Node name for S&R": "LoadImage"
},
"widgets_values": [
"S4.jpg",
"image"
]
},
{
"id": 13,
"type": "SaveImage",
"pos": [
572,
415
],
"size": {
"0": 315,
"1": 270
},
"flags": {},
"order": 3,
"mode": 0,
"inputs": [
{
"name": "images",
"type": "IMAGE",
"link": 21
}
],
"properties": {},
"widgets_values": [
"StyleTransfer"
]
},
{
"id": 11,
"type": "UniST",
"pos": [
568,
246
],
"size": {
"0": 315,
"1": 122
},
"flags": {},
"order": 2,
"mode": 0,
"inputs": [
{
"name": "src_img",
"type": "IMAGE",
"link": 20
},
{
"name": "style_img",
"type": "IMAGE",
"link": 19
}
],
"outputs": [
{
"name": "out_img",
"type": "IMAGE",
"links": [
21
],
"shape": 3,
"slot_index": 0
}
],
"properties": {
"Node name for S&R": "UniST"
},
"widgets_values": [
true,
1024
]
}
],
"links": [
[
19,
3,
0,
11,
1,
"IMAGE"
],
[
20,
12,
0,
11,
0,
"IMAGE"
],
[
21,
11,
0,
13,
0,
"IMAGE"
]
],
"groups": [],
"config": {},
"extra": {
"ds": {
"scale": 1.331000000000001,
"offset": [
-38.8281700653258,
-204.83558799978883
]
}
},
"version": 0.4
}
workflows/workflow_unist_video.json
+281
View File
@@ -0,0 +1,281 @@
{
"last_node_id": 17,
"last_link_id": 33,
"nodes": [
{
"id": 16,
"type": "VHS_VideoCombine",
"pos": [
544,
260
],
"size": [
315,
619
],
"flags": {},
"order": 3,
"mode": 0,
"inputs": [
{
"name": "images",
"type": "IMAGE",
"link": 25
},
{
"name": "audio",
"type": "AUDIO",
"link": null
},
{
"name": "meta_batch",
"type": "VHS_BatchManager",
"link": null
},
{
"name": "vae",
"type": "VAE",
"link": null
}
],
"outputs": [
{
"name": "Filenames",
"type": "VHS_FILENAMES",
"links": null,
"shape": 3
}
],
"properties": {
"Node name for S&R": "VHS_VideoCombine"
},
"widgets_values": {
"frame_rate": 20,
"loop_count": 0,
"filename_prefix": "StyleTransfer",
"format": "video/h264-mp4",
"pix_fmt": "yuv420p",
"crf": 2,
"save_metadata": true,
"pingpong": false,
"save_output": true,
"videopreview": {
"hidden": false,
"paused": false,
"params": {
"filename": "StyleTransfer_00110.mp4",
"subfolder": "",
"type": "output",
"format": "video/h264-mp4",
"frame_rate": 20
}
}
}
},
{
"id": 3,
"type": "LoadImage",
"pos": [
152,
564
],
"size": {
"0": 315,
"1": 314
},
"flags": {},
"order": 0,
"mode": 0,
"outputs": [
{
"name": "IMAGE",
"type": "IMAGE",
"links": [
32
],
"shape": 3,
"slot_index": 0
},
{
"name": "MASK",
"type": "MASK",
"links": null,
"shape": 3
}
],
"properties": {
"Node name for S&R": "LoadImage"
},
"widgets_values": [
"S4.jpg",
"image"
]
},
{
"id": 14,
"type": "VHS_LoadVideo",
"pos": [
238,
114
],
"size": [
235.1999969482422,
398.1134866403979
],
"flags": {},
"order": 1,
"mode": 0,
"inputs": [
{
"name": "meta_batch",
"type": "VHS_BatchManager",
"link": null
},
{
"name": "vae",
"type": "VAE",
"link": null
}
],
"outputs": [
{
"name": "IMAGE",
"type": "IMAGE",
"links": [
33
],
"shape": 3,
"slot_index": 0
},
{
"name": "frame_count",
"type": "INT",
"links": null,
"shape": 3
},
{
"name": "audio",
"type": "AUDIO",
"links": null,
"shape": 3
},
{
"name": "video_info",
"type": "VHS_VIDEOINFO",
"links": null,
"shape": 3
}
],
"properties": {
"Node name for S&R": "VHS_LoadVideo"
},
"widgets_values": {
"video": "cat_full.mp4",
"force_rate": 0,
"force_size": "Disabled",
"custom_width": 512,
"custom_height": 512,
"frame_load_cap": 20,
"skip_first_frames": 0,
"select_every_nth": 1,
"choose video to upload": "image",
"videopreview": {
"hidden": false,
"paused": false,
"params": {
"frame_load_cap": 20,
"skip_first_frames": 0,
"force_rate": 0,
"filename": "cat_full.mp4",
"type": "input",
"format": "video/mp4",
"select_every_nth": 1
}
}
}
},
{
"id": 15,
"type": "UniST_Video",
"pos": [
542,
111
],
"size": {
"0": 315,
"1": 102
},
"flags": {},
"order": 2,
"mode": 0,
"inputs": [
{
"name": "src_video",
"type": "IMAGE",
"link": 33
},
{
"name": "style_video",
"type": "IMAGE",
"link": 32
}
],
"outputs": [
{
"name": "out_img",
"type": "IMAGE",
"links": [
25
],
"shape": 3,
"slot_index": 0
}
],
"properties": {
"Node name for S&R": "UniST_Video"
},
"widgets_values": [
true,
768
]
}
],
"links": [
[
25,
15,
0,
16,
0,
"IMAGE"
],
[
32,
3,
0,
15,
1,
"IMAGE"
],
[
33,
14,
0,
15,
0,
"IMAGE"
]
],
"groups": [],
"config": {},
"extra": {
"ds": {
"scale": 1,
"offset": [
258.79825358315156,
-25.984416674369626
]
}
},
"version": 0.4
}