From 0d60b7ef9a542595b5ae3a7ec32cc6f34f7bbb1d Mon Sep 17 00:00:00 2001 From: Fuou Marinas <57849916+FuouM@users.noreply.github.com> Date: Tue, 13 Aug 2024 13:18:05 +0700 Subject: [PATCH] Performance fix (#2) * move model load to outside * fix oom by changing to inference mode --- README.md | 24 +++++++++++++++----- module_cast/cast_model.py | 20 +++------------- module_efdm/efdm_model.py | 23 ++++--------------- run.py | 48 ++++++++++++++++++++++++++++++++------- 4 files changed, 65 insertions(+), 50 deletions(-) diff --git a/README.md b/README.md index f7cbfc3..e46c1dc 100644 --- a/README.md +++ b/README.md @@ -6,7 +6,7 @@ Click name to jump to workflow 1. [**Neural Neighbor**](#neural-neighbor). Paper: [Neural Neighbor Style Transfer](https://github.com/nkolkin13/NeuralNeighborStyleTransfer) 2. [**CAST**](#cast). Paper: [Domain Enhanced Arbitrary Image Style Transfer via Contrastive Learning](https://github.com/zyxElsa/CAST_pytorch) 3. [**EFDM**](#efdm). Paper: [Exact Feature Distribution Matching for Arbitrary Style Transfer and Domain Generalization](https://github.com/YBZh/EFDM) - +4. [**MicroAST**](#microast). Paper: [Towards Super-Fast Ultra-Resolution Arbitrary Style Transfer](https://github.com/EndyWon/MicroAST) (Under construction) More will come soon @@ -17,7 +17,7 @@ All nodes support batched input (i.e video) but is generally not recommended. On | Neural Neighbor (2020-2022) - Slowest | CAST (2022-2023) - Fast | |:-:|:-:| |![neural_neighbor](example_outputs/neural_neighbor.png) | ![cast](example_outputs/cast.png) | -| **EFDM (2022) - Fast but OOM on larger resolutions** | X | +| **EFDM (2022) - Fast but OOM on larger resolutions** | **MicroAST (2023)** | | ![efdm](example_outputs/efdm.png) | X | ### Neural Neighbor @@ -45,7 +45,7 @@ Download the CAST and UCAST models in the [Test](https://github.com/zyxElsa/CAST Download the `vgg_normalized` model in the [Train](https://github.com/zyxElsa/CAST_pytorch?tab=readme-ov-file#train) section ``` -models\ +models/ │ .gitkeep │ vgg_normalised.pth ─CAST_model @@ -68,7 +68,7 @@ models\ Download the `hm_decoder_iter_160000.pth` in https://github.com/YBZh/EFDM/tree/main/ArbitraryStyleTransfer/models and put it in -`models\hm_decoder_iter_160000.pth` +`models/hm_decoder_iter_160000.pth` Download `vgg_normalized.pth` as instructed in [**CAST**](#cast) @@ -81,12 +81,15 @@ Arguments: - `size`: Size of the height of the images after resizing. The aspect ratio of content is kept, while the styles will be resized to match content's height and width. The higher the `size`, the better the result, but also the more VRAM it will take. You may try `use_cpu` to circumvent OOM. - `use_cpu`: Whether to use CPU or GPU for this. CPU takes very long! (default workflow = ~21 seconds). Defaults to False. -Important: You should not use batched content image input (i.e video) for this node. It is very easy to get OOM and is not optmized for videos. - [workflow_efdm.json](workflows/workflow_efdm.json) ![wf_efdm](workflows/wf_efdm.png) + +### MicroAST + +Download `vgg_normalized.pth` as instructed in [**CAST**](#cast) + ## Credits ``` @@ -113,3 +116,12 @@ year = {2022}} year={2022} } ``` + +``` +@inproceedings{wang2023microast, + title={MicroAST: Towards Super-Fast Ultra-Resolution Arbitrary Style Transfer}, + author={Wang, Zhizhong and Zhao, Lei and Zuo, Zhiwen and Li, Ailin and Chen, Haibo and Xing, Wei and Lu, Dongming}, + booktitle={Proceedings of the AAAI Conference on Artificial Intelligence}, + year={2023} +} +``` diff --git a/module_cast/cast_model.py b/module_cast/cast_model.py index 989334d..3de6cc1 100644 --- a/module_cast/cast_model.py +++ b/module_cast/cast_model.py @@ -13,31 +13,17 @@ def load_a_ckpt(ckpt_path: str): def inference_ucast( src_img: torch.Tensor, style_img: torch.Tensor, - ckpt_ae: str, - ckpt_decB: str, - vgg_path, + netAE, + netDec_B, device, ): - vgg = net.vgg - vgg.load_state_dict(torch.load(vgg_path)) - vgg = torch.nn.Sequential(*list(vgg.children())[:31]) - - netAE = net.ADAIN_Encoder(vgg) - netDec_B = net.Decoder() - - netAE.load_state_dict(load_a_ckpt(ckpt_ae)) - netDec_B.load_state_dict(load_a_ckpt(ckpt_decB)) - - netAE = netAE.to(device).eval() - netDec_B = netDec_B.to(device).eval() - real_A = src_img.permute(0, 3, 1, 2).to(device) real_B = style_img.permute(0, 3, 1, 2).to(device) with torch.no_grad(): real_A_feat = netAE.forward(real_A, real_B) fake_B = netDec_B.forward(real_A_feat) - + if device.type != "cpu": torch.cuda.empty_cache() diff --git a/module_efdm/efdm_model.py b/module_efdm/efdm_model.py index b6e7c25..a8c6925 100644 --- a/module_efdm/efdm_model.py +++ b/module_efdm/efdm_model.py @@ -2,7 +2,6 @@ import torch import torch.nn.functional as F from torchvision import transforms -from . import net from .function import ( adaptive_instance_normalization, adaptive_mean_normalization, @@ -17,8 +16,8 @@ def inference_efdm( src_img: torch.Tensor, style_img: torch.Tensor, device, - decoder_path: str, - vgg_path: str, + decoder, + vgg, alpha: float, size: int, style_type: str, @@ -39,20 +38,6 @@ def inference_efdm( else: style_interpolation_weights = None - decoder = net.decoder - vgg = net.vgg - - decoder.eval() - vgg.eval() - - decoder.load_state_dict(torch.load(decoder_path)) - vgg.load_state_dict(torch.load(vgg_path)) - - vgg = torch.nn.Sequential(*list(vgg.children())[:31]) - - vgg.to(device) - decoder.to(device) - content_tf = test_transform(size, do_crop) content = content_tf(src_img).to(device) @@ -73,7 +58,7 @@ def inference_efdm( tmp_styles = [coral(stl, tmp_content).unsqueeze(0) for stl in style] style = torch.cat(tmp_styles, dim=0) style = style.to(device) - + if device.type != "cpu": torch.cuda.empty_cache() @@ -87,7 +72,7 @@ def inference_efdm( style_interpolation_weights, style_type=style_type, ) - + if device.type != "cpu": torch.cuda.empty_cache() diff --git a/run.py b/run.py index a5b168d..3068d45 100644 --- a/run.py +++ b/run.py @@ -22,7 +22,9 @@ from .constants import ( EFDM_PATH, EFDM_STYLE_TYPES, ) -from .module_cast.cast_model import inference_ucast +from .module_cast import net as net_cast +from .module_cast.cast_model import inference_ucast, load_a_ckpt +from .module_efdm import net as net_efdm from .module_efdm.efdm_model import inference_efdm from .module_neural_neighbor.neural_neighbor_model import inference_neural_neighbor @@ -134,19 +136,35 @@ class CAST: print(f"{style_img.shape=}") device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + vgg = net_cast.vgg + vgg.load_state_dict(torch.load(f"{base_dir}/{CAST_VGG_PATH}")) + vgg = torch.nn.Sequential(*list(vgg.children())[:31]) + + netAE = net_cast.ADAIN_Encoder(vgg) + netDec_B = net_cast.Decoder() + + netAE.load_state_dict( + load_a_ckpt(f"{base_dir}/models/{model_arch}_model/{CAST_NET_AE_PATH}") + ) + netDec_B.load_state_dict( + load_a_ckpt(f"{base_dir}/models/{model_arch}_model/{CAST_NET_DEC_B_PATH}") + ) + + netAE = netAE.to(device).eval() + netDec_B = netDec_B.to(device).eval() + params = { "style_img": style_img, "device": device, - "ckpt_ae": f"{base_dir}/models/{model_arch}_model/{CAST_NET_AE_PATH}", - "ckpt_decB": f"{base_dir}/models/{model_arch}_model/{CAST_NET_DEC_B_PATH}", - "vgg_path": f"{base_dir}/{CAST_VGG_PATH}", + "netAE": netAE, + "netDec_B": netDec_B, } num_frames = src_img.size(0) pbar = ProgressBar(num_frames) result: list[torch.Tensor] = [] - with torch.inference_mode(mode=False): + with torch.no_grad(): for i in range(num_frames): params["src_img"] = src_img[i].unsqueeze(0) res_tensor = inference_ucast(**params) @@ -211,11 +229,25 @@ class EFDM: "cuda" if torch.cuda.is_available() and not use_cpu else "cpu" ) + decoder = net_efdm.decoder + vgg = net_efdm.vgg + + decoder.eval() + vgg.eval() + + decoder.load_state_dict(torch.load(f"{base_dir}/{EFDM_PATH}")) + vgg.load_state_dict(torch.load(f"{base_dir}/{CAST_VGG_PATH}")) + + vgg = torch.nn.Sequential(*list(vgg.children())[:31]) + + vgg.to(device) + decoder.to(device) + params = { "style_img": style_img.permute(0, 3, 1, 2), "device": device, - "decoder_path": f"{base_dir}/{EFDM_PATH}", - "vgg_path": f"{base_dir}/{CAST_VGG_PATH}", + "decoder": decoder, + "vgg": vgg, "alpha": style_strength, "size": size, "style_type": model_arch, @@ -231,7 +263,7 @@ class EFDM: pbar = ProgressBar(num_frames) result: list[torch.Tensor] = [] - with torch.inference_mode(mode=False): + with torch.no_grad(): for i in range(num_frames): params["src_img"] = src_img[i].unsqueeze(0).permute(0, 3, 1, 2) res_tensor = inference_efdm(**params)