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

Performance fix (#2)

Browse Source 
* move model load to outside

* fix oom by changing to inference mode
This commit is contained in:
Fuou Marinas
2024-08-13 13:18:05 +07:00
committed by GitHub
parent f264247cc3
commit 0d60b7ef9a
4 changed files with 65 additions and 50 deletions
Download Patch File Download Diff File Expand all files Collapse all files
README.md
+18 -6
View File
@@ -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}
}
```
module_cast/cast_model.py
+3 -17
View File
@@ -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()
module_efdm/efdm_model.py
+4 -19
View File
@@ -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()
run.py
+40 -8
View File
@@ -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)