49 lines
1.1 KiB
Python
49 lines
1.1 KiB
Python
"""
|
|
@author: Fuou Marinas
|
|
@title: FM Nodes
|
|
@nickname: FM_nodes
|
|
@description: A collection of nodes. WFEN Face Super Resolution.
|
|
"""
|
|
|
|
from pathlib import Path
|
|
|
|
import torch
|
|
|
|
from .constants import WFEN_MODEL_PATH
|
|
from .utils import ensure_size
|
|
from .wfen_module.wfen_model import WFENModel
|
|
|
|
base_dir = Path(__file__).resolve().parent
|
|
|
|
|
|
class WFEN:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"src_img": ("IMAGE",),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
RETURN_NAMES = ("res_img",)
|
|
FUNCTION = "todo"
|
|
CATEGORY = "FM_nodes"
|
|
|
|
def todo(
|
|
self,
|
|
src_img: torch.Tensor,
|
|
):
|
|
print(f"{src_img.shape=}")
|
|
source_tensor = ensure_size(src_img.permute(0, 3, 1, 2))
|
|
print(f"{source_tensor.shape=}")
|
|
wfening = WFENModel()
|
|
wfening.for_load_pretrain_model(f"{base_dir}/{WFEN_MODEL_PATH}")
|
|
|
|
out_tensor = wfening.netG.forward(source_tensor)
|
|
# out_tensor = ensure_size(out_tensor, size=src_img.shape[1:3])
|
|
print(f"{out_tensor.shape=}")
|
|
out_tensor = out_tensor.permute(0, 2, 3, 1)
|
|
|
|
return (out_tensor,)
|