60 lines
2.0 KiB
Python
60 lines
2.0 KiB
Python
from dataclasses import dataclass
|
|
|
|
import torch
|
|
|
|
from ComfyUI_ProPainter_Nodes.model.modules.flow_comp_raft import RAFT_bi
|
|
from ComfyUI_ProPainter_Nodes.model.propainter import InpaintGenerator
|
|
from ComfyUI_ProPainter_Nodes.model.recurrent_flow_completion import (
|
|
RecurrentFlowCompleteNet,
|
|
)
|
|
from ComfyUI_ProPainter_Nodes.utils.download_utils import download_model
|
|
|
|
|
|
@dataclass
|
|
class Models:
|
|
raft_model: RAFT_bi
|
|
flow_model: RecurrentFlowCompleteNet
|
|
inpaint_model: InpaintGenerator
|
|
|
|
|
|
PRETRAIN_MODEL_URL = "https://github.com/sczhou/ProPainter/releases/download/v0.1.0/"
|
|
|
|
|
|
def load_raft_model(device: torch.device) -> RAFT_bi:
|
|
"""Loads the RAFT bi-directional model."""
|
|
model_path = download_model(PRETRAIN_MODEL_URL, "raft-things.pth")
|
|
raft_model = RAFT_bi(model_path, device)
|
|
return raft_model
|
|
|
|
|
|
def load_recurrent_flow_model(device: torch.device) -> RecurrentFlowCompleteNet:
|
|
"""Loads the Recurrent Flow Completion Network model."""
|
|
model_path = download_model(PRETRAIN_MODEL_URL, "recurrent_flow_completion.pth")
|
|
flow_model = RecurrentFlowCompleteNet(model_path)
|
|
for p in flow_model.parameters():
|
|
p.requires_grad = False
|
|
flow_model.to(device)
|
|
flow_model.eval()
|
|
return flow_model
|
|
|
|
|
|
def load_inpaint_model(device: torch.device) -> InpaintGenerator:
|
|
"""Loads the Inpaint Generator model."""
|
|
model_path = download_model(PRETRAIN_MODEL_URL, "ProPainter.pth")
|
|
inpaint_model = InpaintGenerator(model_path=model_path).to(device)
|
|
inpaint_model.eval()
|
|
return inpaint_model
|
|
|
|
|
|
def initialize_models(device: torch.device, use_half: str) -> Models:
|
|
"""Return initialized inference models."""
|
|
raft_model = load_raft_model(device)
|
|
flow_model = load_recurrent_flow_model(device)
|
|
inpaint_model = load_inpaint_model(device)
|
|
|
|
if use_half == "enable":
|
|
# raft_model = raft_model.half()
|
|
flow_model = flow_model.half()
|
|
inpaint_model = inpaint_model.half()
|
|
return Models(raft_model, flow_model, inpaint_model)
|