69 lines
2.2 KiB
Python
69 lines
2.2 KiB
Python
from dataclasses import dataclass
|
|
import os
|
|
|
|
from torch import device
|
|
|
|
from .download_utils import load_file_from_url
|
|
from ..model.modules.flow_comp_raft import RAFT_bi
|
|
from ..model.recurrent_flow_completion import RecurrentFlowCompleteNet
|
|
from ..model.propainter import InpaintGenerator
|
|
|
|
|
|
@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 download_model(model_url: str, model_name: str) -> str:
|
|
"""Downloads a model from a URL and returns the local path to the downloaded model."""
|
|
return load_file_from_url(
|
|
url=os.path.join(model_url, model_name),
|
|
model_dir="custom_nodes/ComfyUI_ProPainter_Nodes/weights",
|
|
progress=True,
|
|
file_name=None,
|
|
)
|
|
|
|
|
|
def load_raft_model(device: 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: 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: 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: 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)
|