323 lines
13 KiB
Python
323 lines
13 KiB
Python
# !/usr/bin/env python
|
|
# -*- coding: UTF-8 -*-
|
|
|
|
import numpy as np
|
|
import torch
|
|
import os
|
|
from .model_loader_utils import tensor2image,phi2narry,nomarl_upscale
|
|
from .OmniPart.scripts.inference_omnipart import (
|
|
inference_omnipart,prepare_inputs,apply_base_model,save_state_data,load_state_data,
|
|
process_image,apply_merge)
|
|
import folder_paths
|
|
import comfy
|
|
import node_helpers
|
|
from typing_extensions import override
|
|
from comfy_api.latest import ComfyExtension, io
|
|
from pathlib import PureWindowsPath
|
|
from datetime import datetime
|
|
import nodes
|
|
os.environ["OPENCV_IO_ENABLE_OPENEXR"] = "1"
|
|
|
|
MAX_SEED = np.iinfo(np.int32).max
|
|
node_cr_path = os.path.dirname(os.path.abspath(__file__))
|
|
|
|
device = torch.device(
|
|
"cuda:0") if torch.cuda.is_available() else torch.device(
|
|
"mps") if torch.backends.mps.is_available() else torch.device(
|
|
"cpu")
|
|
|
|
weigths_OmniPart_current_path = os.path.join(folder_paths.models_dir, "OmniPart")
|
|
if not os.path.exists(weigths_OmniPart_current_path):
|
|
os.makedirs(weigths_OmniPart_current_path)
|
|
|
|
weigths_dinov2_current_path = os.path.join(folder_paths.models_dir, "dinov2")
|
|
if not os.path.exists(weigths_dinov2_current_path):
|
|
os.makedirs(weigths_dinov2_current_path)
|
|
|
|
folder_paths.add_model_folder_path("OmniPart", weigths_OmniPart_current_path) # OmniPart dir
|
|
folder_paths.add_model_folder_path("dinov2", weigths_dinov2_current_path)
|
|
|
|
class OmniPart_SM_Model(io.ComfyNode):
|
|
@classmethod
|
|
def define_schema(cls):
|
|
|
|
return io.Schema(
|
|
node_id="OmniPart_SM_Model",
|
|
display_name="OmniPart_SM_Model",
|
|
category="OmniPart",
|
|
inputs=[
|
|
io.String.Input("custom_path", multiline=False,default="omnipart/OmniPart",),
|
|
io.Combo.Input("box_ckpt",options= ["none"] + [i for i in folder_paths.get_filename_list("OmniPart") if "box" in i.lower()]),
|
|
io.Combo.Input("partfield_encoder",options= ["none"] + [i for i in folder_paths.get_filename_list("OmniPart") if "encoder" in i.lower()]),
|
|
io.Combo.Input("dino",options= ["none"] + folder_paths.get_filename_list("dinov2") ),
|
|
io.Combo.Input("dino_l",options= ["none"] + folder_paths.get_filename_list("dinov2") ),
|
|
],
|
|
outputs=[
|
|
io.Custom("OmniPart_SM_Model").Output(),
|
|
],
|
|
)
|
|
@classmethod
|
|
def execute(cls, custom_path,box_ckpt,partfield_encoder,dino,dino_l) -> io.NodeOutput:
|
|
part_synthesis_ckpt=PureWindowsPath(custom_path).as_posix() if custom_path else None
|
|
bbox_gen_ckpt=folder_paths.get_full_path("OmniPart", box_ckpt) if box_ckpt != "none" else None
|
|
partfield_encoder_path=folder_paths.get_full_path("OmniPart", partfield_encoder) if box_ckpt != "none" else None
|
|
dino_path=folder_paths.get_full_path("dinov2", dino) if dino != "none" else None
|
|
dino_l_path=folder_paths.get_full_path("dinov2", dino_l) if dino_l != "none" else None
|
|
assert bbox_gen_ckpt is not None and partfield_encoder is not None and dino_path is not None and dino_l_path is not None, "Please select the BboxGen and partfield_encoder and dino and dino l checkpoint file"
|
|
config_path=os.path.join(node_cr_path,"OmniPart/configs/bbox_gen.yaml")
|
|
if part_synthesis_ckpt is None:
|
|
part_synthesis_ckpt="omnipart/OmniPart"
|
|
part_synthesis_pipeline,bbox_gen_model=apply_base_model(part_synthesis_ckpt,partfield_encoder_path,bbox_gen_ckpt,dino_path,config_path,node_cr_path,dino_l_path)
|
|
model={"part_synthesis_pipeline":part_synthesis_pipeline,"bbox_gen_model":bbox_gen_model}
|
|
return io.NodeOutput(model)
|
|
|
|
|
|
class OmniPart_SM_PreModel(io.ComfyNode):
|
|
@classmethod
|
|
def define_schema(cls):
|
|
return io.Schema(
|
|
node_id="OmniPart_SM_PreModel",
|
|
display_name="OmniPart_SM_PreModel",
|
|
category="OmniPart",
|
|
inputs=[
|
|
io.Combo.Input("rmgb",options= ["none"] + folder_paths.get_filename_list("OmniPart")),
|
|
io.Combo.Input("sam",options= ["none"] + [i for i in folder_paths.get_filename_list("OmniPart") if "sam" in i.lower()]),
|
|
],
|
|
outputs=[
|
|
io.Custom("OmniPart_SM_SamModel").Output(display_name="sam_model"),
|
|
io.Custom("OmniPart_SM_SegModel").Output(display_name="rmbg_model"),
|
|
],
|
|
)
|
|
@classmethod
|
|
def execute(cls, rmgb,sam,) -> io.NodeOutput:
|
|
sam_ckpt_path=folder_paths.get_full_path("OmniPart", sam) if sam != "none" else None
|
|
rmgb_ckpt_path=folder_paths.get_full_path("OmniPart", rmgb) if rmgb != "none" else None
|
|
assert sam_ckpt_path is not None and rmgb_ckpt_path is not None, "Please select the SAM and RMBG 2.0 checkpoint file"
|
|
|
|
print("Loading SAM model...")
|
|
from segment_anything import SamAutomaticMaskGenerator, build_sam #seg
|
|
sam_model = build_sam(checkpoint=sam_ckpt_path)
|
|
|
|
print("Loading BriaRMBG 2.0 model...")
|
|
from transformers import AutoConfig
|
|
from .OmniPart.configs.RMBG.birefnet import BiRefNet
|
|
rmgb_repo=os.path.join(node_cr_path,"OmniPart/configs/RMBG")
|
|
config = AutoConfig.from_pretrained(rmgb_repo, trust_remote_code=True)
|
|
rmbg_model = BiRefNet(config=config)
|
|
sd=comfy.utils.load_torch_file(rmgb_ckpt_path)
|
|
rmbg_model.load_state_dict(sd,strict=False)
|
|
del sd
|
|
rmbg_model.eval()
|
|
# rmbg_model = AutoModelForImageSegmentation.from_pretrained(rmgb_repo, trust_remote_code=True)
|
|
|
|
return io.NodeOutput (sam_model,rmbg_model)
|
|
|
|
class OmniPart_SM_PreImg(io.ComfyNode):
|
|
@classmethod
|
|
def define_schema(cls):
|
|
return io.Schema(
|
|
node_id="OmniPart_SM_PreImg",
|
|
display_name="OmniPart_SM_PreImg",
|
|
category="OmniPart",
|
|
inputs=[
|
|
io.Custom("OmniPart_SM_SamModel").Input("sam_model"),
|
|
io.Custom("OmniPart_SM_SegModel").Input("rmbg_model"),
|
|
io.Image.Input("image"),
|
|
io.Int.Input("threshold", default=1600, min=600, max=4000,step=200),
|
|
],
|
|
outputs=[
|
|
io.Image.Output(display_name="image"),
|
|
io.String.Output( display_name="dir_prefix"),
|
|
],
|
|
)
|
|
@classmethod
|
|
def execute(cls, sam_model,rmbg_model,image,threshold) -> io.NodeOutput:
|
|
phi_in=tensor2image(image)
|
|
|
|
dir_prefix = datetime.now().strftime("%y%m%d%H%M%S")[-6:]
|
|
user_dir=os.path.join(folder_paths.get_output_directory(), f"OmniPart/temp_{dir_prefix}")
|
|
if not os.path.exists(user_dir):
|
|
os.makedirs(user_dir)
|
|
|
|
image_path=os.path.join(user_dir,f"img_{dir_prefix}.png")
|
|
phi_in.save(image_path)
|
|
|
|
sam_model.to(device)
|
|
from segment_anything import SamAutomaticMaskGenerator
|
|
sam_mask_generator = SamAutomaticMaskGenerator(sam_model)
|
|
rmbg_model.to(device)
|
|
|
|
init_seg_path, pre_merge_image, state=process_image(image_path, threshold, sam_mask_generator,rmbg_model,user_dir,device)
|
|
save_state_data(state, user_dir,threshold)
|
|
|
|
sam_model.to("cpu")
|
|
del sam_mask_generator
|
|
rmbg_model.to("cpu")
|
|
|
|
torch.cuda.empty_cache()
|
|
image=phi2narry(pre_merge_image)
|
|
return io.NodeOutput(image, dir_prefix)
|
|
|
|
|
|
class OmniPart_SM_MergeImg(io.ComfyNode):
|
|
@classmethod
|
|
def define_schema(cls):
|
|
return io.Schema(
|
|
node_id="OmniPart_SM_MergeImg",
|
|
display_name="OmniPart_SM_MergeImg",
|
|
category="OmniPart",
|
|
inputs=[
|
|
io.Custom("OmniPart_SM_SamModel").Input("sam_model"),
|
|
io.String.Input( "dir_prefix", multiline=False,default="",),
|
|
io.String.Input( "merge_input", multiline=False,default="1,3;2,4",),
|
|
io.Image.Input("image",optional=True),
|
|
],
|
|
outputs=[
|
|
io.Conditioning.Output(display_name="state"),
|
|
],
|
|
)
|
|
@classmethod
|
|
def execute(cls, sam_model,dir_prefix,merge_input,image=None) -> io.NodeOutput:
|
|
if image is not None:
|
|
image=tensor2image(image)
|
|
if not dir_prefix:
|
|
dir_prefix_list=extract_loadimage_filenames(recent_workflow_data_sm)
|
|
if not dir_prefix_list:
|
|
raise ValueError("dir_prefix is required if no active workflow found input a valid dir_prefix")
|
|
else:
|
|
print(dir_prefix_list)
|
|
for prefix in dir_prefix_list:
|
|
if isinstance(prefix, str) and "_" in prefix:
|
|
dir_prefix_ = prefix.split("_",1)[0]
|
|
print(dir_prefix_)
|
|
if len(dir_prefix_)==6:
|
|
dir_prefix = dir_prefix_
|
|
break
|
|
else:
|
|
continue
|
|
|
|
else:
|
|
if len(dir_prefix)!=6:
|
|
raise ValueError("dir_prefix is required a number of 6 digits ")
|
|
user_dir=os.path.join(folder_paths.get_output_directory(), f"OmniPart/temp_{dir_prefix}")
|
|
img_name=f"img_{dir_prefix}"
|
|
state={
|
|
"processed_image": os.path.join(user_dir, f"{img_name}_processed.png"),
|
|
"img_name": img_name,
|
|
"pre_split_path": os.path.join(user_dir, f"{img_name}_pre_split.png"),
|
|
"user_dir": user_dir,
|
|
}
|
|
state=load_state_data(state, user_dir)
|
|
|
|
sam_model.to(device)
|
|
from segment_anything import SamAutomaticMaskGenerator
|
|
sam_mask_generator = SamAutomaticMaskGenerator(sam_model)
|
|
merged_seg_path, cond=apply_merge(merge_input, sam_mask_generator,state, user_dir,int(state.get("threshold",2000)))
|
|
cond.update(state)
|
|
sam_model.to("cpu")
|
|
del sam_mask_generator
|
|
torch.cuda.empty_cache()
|
|
return io.NodeOutput(state)
|
|
|
|
|
|
class OmniPart_SM_Prebox(io.ComfyNode):
|
|
@classmethod
|
|
def define_schema(cls):
|
|
return io.Schema(
|
|
node_id="OmniPart_SM_Prebox",
|
|
display_name="OmniPart_SM_Prebox",
|
|
category="OmniPart",
|
|
inputs=[
|
|
io.Custom("OmniPart_SM_Model").Input("model"),
|
|
io.Conditioning.Input("state"),
|
|
io.Int.Input("seed", default=0, min=0, max=MAX_SEED),
|
|
],
|
|
outputs=[
|
|
io.Conditioning.Output(display_name="cond"),
|
|
],
|
|
)
|
|
@classmethod
|
|
def execute(cls, model, state,seed) -> io.NodeOutput:
|
|
cond=prepare_inputs(state,model["part_synthesis_pipeline"],model["bbox_gen_model"],user_dir=state.get("user_dir"),seed=seed,output_dir=folder_paths.get_output_directory(),device=device)
|
|
cond.update({"seed":seed})
|
|
return io.NodeOutput(cond,)
|
|
|
|
|
|
class OmniPart_SM_KSampler(io.ComfyNode):
|
|
@classmethod
|
|
def define_schema(cls):
|
|
return io.Schema(
|
|
node_id="OmniPart_SM_KSampler",
|
|
display_name="OmniPart_SM_KSampler",
|
|
category="OmniPart",
|
|
inputs=[
|
|
io.Custom("OmniPart_SM_Model").Input("model"),
|
|
io.Int.Input("steps", default=20, min=1, max=10000),
|
|
io.Float.Input("cfg", default=8.0, min=0.0, max=100.0, step=0.1, round=0.01,),
|
|
io.Conditioning.Input("cond"),
|
|
io.Float.Input("simplify_ratio", default=0.3, min=0.0, max=1.0, step=0.1,display_mode=io.NumberDisplay.slider),
|
|
|
|
],
|
|
outputs=[
|
|
io.String.Output(display_name="scene_texture"),
|
|
io.String.Output(display_name="mesh_segment"),
|
|
],
|
|
)
|
|
|
|
@classmethod
|
|
def execute(cls, model, steps, cfg, cond,simplify_ratio ) -> io.NodeOutput:
|
|
scene_texture,mesh_segment=inference_omnipart(model["part_synthesis_pipeline"],cond,steps,cfg,simplify_ratio,device)
|
|
return io.NodeOutput(scene_texture,mesh_segment)
|
|
|
|
from aiohttp import web
|
|
from server import PromptServer
|
|
@PromptServer.instance.routes.get("/OmniPart_SM_Extension")
|
|
async def get_hello(request):
|
|
return web.json_response("OmniPart_SM_Extension")
|
|
|
|
class OmniPart_SM_Extension(ComfyExtension):
|
|
@override
|
|
async def get_node_list(self) -> list[type[io.ComfyNode]]:
|
|
return [
|
|
OmniPart_SM_Model,
|
|
OmniPart_SM_PreModel,
|
|
OmniPart_SM_PreImg,
|
|
OmniPart_SM_MergeImg,
|
|
OmniPart_SM_Prebox,
|
|
OmniPart_SM_KSampler,
|
|
]
|
|
async def comfy_entrypoint() -> OmniPart_SM_Extension: # ComfyUI calls this to load your extension and its nodes.
|
|
return OmniPart_SM_Extension()
|
|
|
|
|
|
recent_workflow_data_sm = None
|
|
|
|
def on_prompt_handler(json_data):
|
|
global recent_workflow_data_sm
|
|
recent_workflow_data_sm = json_data
|
|
return json_data
|
|
|
|
PromptServer.instance.add_on_prompt_handler(on_prompt_handler)
|
|
|
|
def extract_loadimage_filenames(json_data):
|
|
filenames = []
|
|
if not json_data or not isinstance(json_data, dict):
|
|
return filenames
|
|
|
|
nodes_data = json_data.get("prompt", {}).get("nodes", [])
|
|
if not nodes_data and "nodes" in json_data:
|
|
nodes_data = json_data.get("nodes", [])
|
|
|
|
|
|
for node in nodes_data:
|
|
if node.get("type") == "LoadImage":
|
|
widgets_values = node.get("widgets_values", [])
|
|
if widgets_values and len(widgets_values) > 0:
|
|
filename = widgets_values[0]
|
|
if isinstance(filename, str):
|
|
filenames.append(filename)
|
|
|
|
return filenames
|
|
|
|
|