diff --git a/README.md b/README.md index 61877aa..237c840 100644 --- a/README.md +++ b/README.md @@ -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 diff --git a/example_outputs/example_unist_video.png b/example_outputs/example_unist_video.png new file mode 100644 index 0000000..bd1b14b Binary files /dev/null and b/example_outputs/example_unist_video.png differ diff --git a/example_outputs/unist_image.png b/example_outputs/unist_image.png new file mode 100644 index 0000000..a3baf85 Binary files /dev/null and b/example_outputs/unist_image.png differ diff --git a/module_unist/models.py b/module_unist/models.py index 38fe54c..957c521 100644 --- a/module_unist/models.py +++ b/module_unist/models.py @@ -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() diff --git a/module_unist/unist_model.py b/module_unist/unist_model.py index f1c8e7a..ae5f1e7 100644 --- a/module_unist/unist_model.py +++ b/module_unist/unist_model.py @@ -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] diff --git a/run.py b/run.py index 67b44be..94c482f 100644 --- a/run.py +++ b/run.py @@ -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]): diff --git a/workflows/wf_unist_image.png b/workflows/wf_unist_image.png new file mode 100644 index 0000000..c6efa5e Binary files /dev/null and b/workflows/wf_unist_image.png differ diff --git a/workflows/workflow_unist_image.json b/workflows/workflow_unist_image.json new file mode 100644 index 0000000..df93e7b --- /dev/null +++ b/workflows/workflow_unist_image.json @@ -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 +} \ No newline at end of file diff --git a/workflows/workflow_unist_video.json b/workflows/workflow_unist_video.json new file mode 100644 index 0000000..008f483 --- /dev/null +++ b/workflows/workflow_unist_video.json @@ -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 +} \ No newline at end of file