fix: 🐛 prepend MTB_ to all classes

to avoid any future clash.
This commit is contained in:
Mel Massadian
2024-04-27 01:13:32 +02:00
parent 652ac3f3b9
commit b1d74adb15
16 changed files with 192 additions and 152 deletions
+29 -4
View File
@@ -70,7 +70,7 @@ def extract_nodes_from_source(filename: Path):
def load_nodes():
errors = []
errors: list[str] = []
nodes = []
nodes_failed = []
@@ -87,9 +87,11 @@ def load_nodes():
log.debug(f"Imported {module_name} nodes")
except AttributeError:
log.debug(f"Skipping wip module {module_name}")
pass # wip nodes
except Exception:
error_message = traceback.format_exc().splitlines()[-1]
errors.append(
f"Failed to import module {module_name} because {error_message}"
)
@@ -136,6 +138,18 @@ def wiki_to_classname(s: str):
)
def classname_to_wiki(s: str):
classname = s.replace("MTB_", "")
parts = []
start = 0
for i in range(1, len(classname)):
if classname[i].isupper():
parts.append(classname[start:i].lower())
start = i
parts.append(classname[start:].lower())
return "nodes-" + "-".join(parts)
wiki = here / "wiki"
node_docs = {}
if wiki.exists() and wiki.is_dir():
@@ -146,20 +160,31 @@ if wiki.exists() and wiki.is_dir():
# - REGISTER NODES
MTB_EXPORT = os.environ.get("MTB_EXPORT")
nodes, failed = load_nodes()
for node_class in nodes:
class_name = node_class.__name__
class_name: str = node_class.__name__
linked_doc = node_docs.get(class_name)
if not hasattr(node_class, "DESCRIPTION"):
if linked_doc:
log.debug(f"Found linked doc for {class_name}")
log.debug(f"Found linked doc for {class_name}, using it")
node_class.DESCRIPTION = linked_doc
elif node_class.__doc__:
log.debug(f"Using __doc__ as description for {class_name}")
node_class.DESCRIPTION = node_class.__doc__
if MTB_EXPORT:
wiki_name = classname_to_wiki(class_name)
(wiki / "nodes" / wiki_name + ".md").write_text(
node_class.__doc__, encoding="utf-8"
)
else:
log.debug(
f"None of the methods could retrive documentation for {class_name}"
f"None of the methods could retrieve documentation for {class_name}"
)
node_label = f"{get_label(class_name)} (mtb)"
+35 -13
View File
@@ -3,12 +3,17 @@ import csv
from aiohttp import web
from .log import mklog
from .utils import backup_file, here, import_install, reqs_map, run_command, styles_dir
from .utils import (
backup_file,
import_install,
reqs_map,
run_command,
styles_dir,
)
endlog = mklog("mtb endpoint")
# - ACTIONS
import platform
import sys
from pathlib import Path
@@ -22,7 +27,9 @@ def ACTIONS_installDependency(dependency_names=None):
# reqs = []
resolved_names = [reqs_map.get(name, name) for name in dependency_names]
try:
run_command([Path(sys.executable), "-m", "pip", "install"] + resolved_names)
run_command(
[Path(sys.executable), "-m", "pip", "install"] + resolved_names
)
return {"success": True}
except Exception as e:
@@ -44,9 +51,9 @@ def ACTIONS_installDependency(dependency_names=None):
def ACTIONS_getStyles(style_name=None):
from .nodes.conditions import StylesLoader
from .nodes.conditions import MTB_StylesLoader
styles = StylesLoader.options
styles = MTB_StylesLoader.options
match_list = ["name"]
if styles:
filtered_styles = {
@@ -55,7 +62,9 @@ def ACTIONS_getStyles(style_name=None):
if not key.startswith("__") and key not in match_list
}
if style_name:
return filtered_styles.get(style_name, {"error": "Style not found"})
return filtered_styles.get(
style_name, {"error": "Style not found"}
)
return filtered_styles
return {"error": "No styles found"}
@@ -75,7 +84,9 @@ def ACTIONS_saveStyle(data):
break
if not target:
endlog.warning(f"Could not determine the target file for {data.keys()}")
endlog.warning(
f"Could not determine the target file for {data.keys()}"
)
return {"error": "Could not determine the target file for the style"}
backup_file(target)
@@ -103,11 +114,16 @@ async def do_action(request) -> web.Response:
return web.json_response({"result": result})
available_methods = [
attr[len("ACTIONS_") :] for attr in globals() if attr.startswith("ACTIONS_")
attr[len("ACTIONS_") :]
for attr in globals()
if attr.startswith("ACTIONS_")
]
return web.json_response(
{"error": "Invalid method name.", "available_methods": available_methods}
{
"error": "Invalid method name.",
"available_methods": available_methods,
}
)
@@ -127,7 +143,7 @@ def csv_editor():
style_files = {}
for file in inputs:
with open(file, "r", encoding="utf8") as f:
with open(file, encoding="utf8") as f:
parsed = csv.reader(f)
style_files[file.name] = []
for row in parsed:
@@ -235,7 +251,9 @@ def add_split_pane(left_content, right_content, vertical=True):
def add_dropdown(title, options):
option_str = "\n".join([f"<option value='{opt}'>{opt}</option>" for opt in options])
option_str = "\n".join(
[f"<option value='{opt}'>{opt}</option>" for opt in options]
)
return f"""
<select>
<option disabled selected>{title}</option>
@@ -254,11 +272,15 @@ def render_table(table_dict, sort=True, title=None):
if isinstance(item, dict):
if "dependencies" in item:
table_rows += f"<tr><td>{name}</td><td>"
table_rows += f"{dependencies_button(name,item['dependencies'])}"
table_rows += (
f"{dependencies_button(name,item['dependencies'])}"
)
table_rows += "</td></tr>"
else:
table_rows += f"<tr><td>{name}</td><td>{render_table(item)}</td></tr>"
table_rows += (
f"<tr><td>{name}</td><td>{render_table(item)}</td></tr>"
)
# elif isinstance(item, str):
# table_rows += f"<tr><td>{name}</td><td>{item}</td></tr>"
else:
+2 -2
View File
@@ -7,7 +7,7 @@ from PIL import Image
from ..log import log
from ..utils import apply_easing, pil2tensor
from .transform import TransformImage
from .transform import MTB_TransformImage
def hex_to_rgb(hex_color, bgr=False):
@@ -415,7 +415,7 @@ class MTB_Batch2dTransform:
if count == 0:
keyframes[name] = [default_vals[name]] * image.shape[0]
transformer = TransformImage()
transformer = MTB_TransformImage()
res = [
transformer.transform(
image[i].unsqueeze(0),
+34 -14
View File
@@ -1,4 +1,5 @@
import csv, shutil
import csv
import shutil
from pathlib import Path
import folder_paths
@@ -7,7 +8,7 @@ from ..log import log
from ..utils import here
class InterpolateClipSequential:
class MTB_InterpolateClipSequential:
@classmethod
def INPUT_TYPES(cls):
return {
@@ -28,7 +29,12 @@ class InterpolateClipSequential:
CATEGORY = "mtb/conditioning"
def interpolate_encodings_sequential(
self, base_text, text_to_replace, clip, interpolation_strength, **replacements
self,
base_text,
text_to_replace,
clip,
interpolation_strength,
**replacements,
):
log.debug(f"Received interpolation_strength: {interpolation_strength}")
@@ -63,20 +69,30 @@ class InterpolateClipSequential:
log.debug("Using the base text a the base blend")
# - Start with the base_text condition
tokens = clip.tokenize(base_text)
cond_from, pooled_from = clip.encode_from_tokens(tokens, return_pooled=True)
cond_from, pooled_from = clip.encode_from_tokens(
tokens, return_pooled=True
)
else:
base_replace = list(replacements.values())[segment_index - 1]
log.debug(f"Using {base_replace} a the base blend")
# - Start with the base_text condition replaced by the closest replacement
tokens = clip.tokenize(base_text.replace(text_to_replace, base_replace))
cond_from, pooled_from = clip.encode_from_tokens(tokens, return_pooled=True)
tokens = clip.tokenize(
base_text.replace(text_to_replace, base_replace)
)
cond_from, pooled_from = clip.encode_from_tokens(
tokens, return_pooled=True
)
replacement_text = list(replacements.values())[segment_index]
interpolated_text = base_text.replace(text_to_replace, replacement_text)
interpolated_text = base_text.replace(
text_to_replace, replacement_text
)
tokens = clip.tokenize(interpolated_text)
cond_to, pooled_to = clip.encode_from_tokens(tokens, return_pooled=True)
cond_to, pooled_to = clip.encode_from_tokens(
tokens, return_pooled=True
)
# - Linearly interpolate between the two conditions
interpolated_condition = (
@@ -86,10 +102,12 @@ class InterpolateClipSequential:
1.0 - local_strength
) * pooled_from + local_strength * pooled_to
return ([[interpolated_condition, {"pooled_output": interpolated_pooled}]],)
return (
[[interpolated_condition, {"pooled_output": interpolated_pooled}]],
)
class SmartStep:
class MTB_SmartStep:
"""Utils to control the steps start/stop of the KAdvancedSampler in percentage"""
@classmethod
@@ -136,7 +154,7 @@ def install_default_styles(force=False):
return dest_style
class StylesLoader:
class MTB_StylesLoader:
"""Load csv files and populate a dropdown from the rows (à la A111)"""
options = {}
@@ -148,13 +166,15 @@ class StylesLoader:
if not input_dir.exists():
install_default_styles()
if not (files := [f for f in input_dir.iterdir() if f.suffix == ".csv"]):
if not (
files := [f for f in input_dir.iterdir() if f.suffix == ".csv"]
):
log.warn(
"No styles found in the styles folder, place at least one csv file in the styles folder at the root of ComfyUI (for instance ComfyUI/styles/mystyle.csv)"
)
for file in files:
with open(file, "r", encoding="utf8") as f:
with open(file, encoding="utf8") as f:
parsed = csv.reader(f)
for i, row in enumerate(parsed):
log.debug(f"Adding style {row[0]}")
@@ -193,4 +213,4 @@ class StylesLoader:
return (self.options[style_name][0], self.options[style_name][1])
__nodes__ = [SmartStep, StylesLoader, InterpolateClipSequential]
__nodes__ = [MTB_SmartStep, MTB_StylesLoader, MTB_InterpolateClipSequential]
+5 -5
View File
@@ -6,7 +6,7 @@ from ..log import log
from ..utils import np2tensor, pil2tensor, tensor2np, tensor2pil
class Bbox:
class MTB_Bbox:
"""The bounding box (BBOX) custom type used by other nodes"""
@classmethod
@@ -41,7 +41,7 @@ class Bbox:
return ((x, y, width, height),)
class BboxFromMask:
class MTB_BboxFromMask:
"""From a mask extract the bounding box"""
@classmethod
@@ -110,7 +110,7 @@ class BboxFromMask:
)
class Crop:
class MTB_Crop:
"""Crops an image and an optional mask to a given bounding box
The bounding box can be given as a tuple of (x, y, width, height) or as a BBOX type
@@ -218,7 +218,7 @@ def bbox_to_region(bbox, target_size=None):
return (bbox[0], bbox[1], bbox[0] + bbox[2], bbox[1] + bbox[3])
class Uncrop:
class MTB_Uncrop:
"""Uncrops an image to a given bounding box
The bounding box can be given as a tuple of (x, y, width, height) or as a BBOX type
@@ -324,4 +324,4 @@ class Uncrop:
return (pil2tensor(out_images),)
__nodes__ = [BboxFromMask, Bbox, Crop, Uncrop]
__nodes__ = [MTB_BboxFromMask, MTB_Bbox, MTB_Crop, MTB_Uncrop]
+2 -2
View File
@@ -303,7 +303,7 @@ def normals_to_height(normals_img, seamless, progress_callback):
# - ADDON
class DeepBump:
class MTB_DeepBump:
"""Normal & height maps generation from single pictures"""
@classmethod
@@ -386,4 +386,4 @@ class DeepBump:
return (torch.cat(out_images, dim=0),)
__nodes__ = [DeepBump]
__nodes__ = [MTB_DeepBump]
+8 -9
View File
@@ -1,6 +1,4 @@
import os
from pathlib import Path
from typing import Tuple
import comfy
import comfy.utils
@@ -9,14 +7,13 @@ import folder_paths
import numpy as np
import torch
from comfy import model_management
from gfpgan import GFPGANer
from PIL import Image
from ..log import NullWriter, log
from ..utils import get_model_path, np2tensor, pil2tensor, tensor2np
class LoadFaceEnhanceModel:
class MTB_LoadFaceEnhanceModel:
"""Loads a GFPGan or RestoreFormer model for face enhancement."""
def __init__(self) -> None:
@@ -81,6 +78,8 @@ class LoadFaceEnhanceModel:
CATEGORY = "mtb/facetools"
def load_model(self, model_name, upscale=2, bg_upsampler=None):
from gfpgan import GFPGANer
basic = "RestoreFormer" not in model_name
fr_root, um_root = self.get_models_root()
@@ -153,7 +152,7 @@ class BGUpscaleWrapper:
import sys
class RestoreFace:
class MTB_RestoreFace:
"""Uses GFPGan to restore faces"""
def __init__(self) -> None:
@@ -182,7 +181,7 @@ class RestoreFace:
def do_restore(
self,
image: torch.Tensor,
model: GFPGANer,
model,
aligned,
only_center_face,
weight,
@@ -220,12 +219,12 @@ class RestoreFace:
def restore(
self,
image: torch.Tensor,
model: GFPGANer,
model,
aligned=False,
only_center_face=False,
weight=0.5,
save_tmp_steps=True,
) -> Tuple[torch.Tensor]:
) -> tuple[torch.Tensor]:
out = [
self.do_restore(
image[i],
@@ -275,4 +274,4 @@ class RestoreFace:
cv2.imwrite(file, cmp_img)
__nodes__ = [RestoreFace, LoadFaceEnhanceModel]
__nodes__ = [MTB_RestoreFace, MTB_LoadFaceEnhanceModel]
+4 -4
View File
@@ -22,7 +22,7 @@ from ..utils import download_antelopev2, get_model_path, pil2tensor, tensor2pil
log = mklog(__name__)
class LoadFaceAnalysisModel:
class MTB_LoadFaceAnalysisModel:
"""Loads a face analysis model"""
models = []
@@ -53,7 +53,7 @@ class LoadFaceAnalysisModel:
return (face_analyser,)
class LoadFaceSwapModel:
class MTB_LoadFaceSwapModel:
"""Loads a faceswap model"""
@staticmethod
@@ -97,7 +97,7 @@ class LoadFaceSwapModel:
# region roop node
class FaceSwap:
class MTB_FaceSwap:
"""Face swap using deepinsight/insightface models"""
model = None
@@ -239,4 +239,4 @@ def swap_face(
# endregion face swap utils
__nodes__ = [FaceSwap, LoadFaceSwapModel, LoadFaceAnalysisModel]
__nodes__ = [MTB_FaceSwap, MTB_LoadFaceSwapModel, MTB_LoadFaceAnalysisModel]
+4 -4
View File
@@ -52,7 +52,7 @@ from ..utils import comfy_dir, font_path, pil2tensor
# return m.digest().hex()
class UnsplashImage:
class MTB_UnsplashImage:
"""Unsplash Image given a keyword and a size"""
@classmethod
@@ -113,7 +113,7 @@ class UnsplashImage:
return (None,)
class QrCode:
class MTB_QrCode:
"""Basic QR Code generator"""
@classmethod
@@ -364,8 +364,8 @@ by default it fallsback to a default font.
__nodes__ = [
QrCode,
UnsplashImage,
MTB_QrCode,
MTB_UnsplashImage,
MTB_TextToImage,
# MtbExamples,
]
+9 -8
View File
@@ -1,12 +1,9 @@
import glob
import os
from pathlib import Path
from typing import List
import comfy
import comfy.model_management as model_management
import comfy.utils
import folder_paths
import numpy as np
import tensorflow as tf
import torch
@@ -17,7 +14,7 @@ from ..log import log
from ..utils import get_model_path
class LoadFilmModel:
class MTB_LoadFilmModel:
"""Loads a FILM model"""
@staticmethod
@@ -58,7 +55,7 @@ class LoadFilmModel:
return (interpolator.Interpolator(model_path.as_posix(), None),)
class FilmInterpolation:
class MTB_FilmInterpolation:
"""Google Research FILM frame interpolation for large motion"""
@classmethod
@@ -107,12 +104,16 @@ class FilmInterpolation:
in_frames, interpolate, film_model
):
out_tensors.append(
torch.from_numpy(frame) if isinstance(frame, np.ndarray) else frame
torch.from_numpy(frame)
if isinstance(frame, np.ndarray)
else frame
)
model_management.throw_exception_if_processing_interrupted()
pbar.update(1)
out_tensors = torch.cat([tens.unsqueeze(0) for tens in out_tensors], dim=0)
out_tensors = torch.cat(
[tens.unsqueeze(0) for tens in out_tensors], dim=0
)
log.debug(f"Returning {len(out_tensors)} tensors")
log.debug(f"Output shape {out_tensors.shape}")
@@ -120,4 +121,4 @@ class FilmInterpolation:
return (out_tensors,)
__nodes__ = [LoadFilmModel, FilmInterpolation]
__nodes__ = [MTB_LoadFilmModel, MTB_FilmInterpolation]
+2 -2
View File
@@ -21,7 +21,7 @@ def get_playlist_path(playlist_name: str, persistant_playlist=False):
return output_dir / "playlists" / session_id / f"{playlist_name}.json"
class ReadPlaylist:
class MTB_ReadPlaylist:
"""Read a playlist"""
@classmethod
@@ -399,5 +399,5 @@ __nodes__ = [
MTB_SaveGif,
MTB_ExportWithFfmpeg,
MTB_AddToPlaylist,
ReadPlaylist,
MTB_ReadPlaylist,
]
+6 -3
View File
@@ -1,7 +1,7 @@
import torch
class LatentLerp:
class MTB_LatentLerp:
"""Linear interpolation (blend) between two latent vectors"""
@classmethod
@@ -10,7 +10,10 @@ class LatentLerp:
"required": {
"A": ("LATENT",),
"B": ("LATENT",),
"t": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.01}),
"t": (
"FLOAT",
{"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.01},
),
}
}
@@ -29,5 +32,5 @@ class LatentLerp:
__nodes__ = [
LatentLerp,
MTB_LatentLerp,
]
+2 -2
View File
@@ -80,7 +80,7 @@ def conv_forward(lyr, tensor, weight, bias):
)
class ModelPatchSeamless:
class MTB_ModelPatchSeamless:
"""Uses the stable diffusion 'hack' to infer seamless images by setting the model layers padding mode to circular (experimental)"""
@classmethod
@@ -152,4 +152,4 @@ class ModelPatchSeamless:
return (model, hacked_model)
__nodes__ = [ModelPatchSeamless, MTB_VaeDecode]
__nodes__ = [MTB_ModelPatchSeamless, MTB_VaeDecode]
-60
View File
@@ -22,66 +22,6 @@ class MTB_IntToBool:
return (bool(int),)
class MTB_IntToNumber:
"""Node addon for the WAS Suite. Converts a "comfy" INT to a NUMBER."""
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"int": (
"INT",
{
"default": 0,
"min": -1e9,
"max": 1e9,
"step": 1,
"forceInput": True,
},
),
}
}
RETURN_TYPES = ("NUMBER",)
FUNCTION = "int_to_number"
CATEGORY = "mtb/number"
def int_to_number(self, int):
return (int,)
class MTB_FloatToNumber:
"""Node addon for the WAS Suite. Converts a "comfy" FLOAT to a NUMBER."""
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"float": (
"FLOAT",
{
"default": 0,
"min": -1e9,
"max": 1e9,
"step": 1,
"forceInput": True,
},
),
}
}
RETURN_TYPES = ("NUMBER",)
FUNCTION = "float_to_number"
CATEGORY = "mtb/number"
def float_to_number(self, float):
return (float,)
return (int,)
__nodes__ = [
MTB_FloatToNumber,
MTB_IntToBool,
MTB_IntToNumber,
]
+42 -15
View File
@@ -1,15 +1,16 @@
from math import ceil, sqrt
from typing import cast
import torch
import torchvision.transforms.functional as TF
from ..utils import log, hex_to_rgb, tensor2pil, pil2tensor
from math import sqrt, ceil
from typing import cast
from PIL import Image
from ..utils import hex_to_rgb, log, pil2tensor, tensor2pil
class TransformImage:
class MTB_TransformImage:
"""Save torch tensors (image, mask or latent) to disk, useful to debug things outside comfy
it return a tensor representing the transformed images with the same shape as the input tensor
"""
@@ -18,10 +19,22 @@ class TransformImage:
return {
"required": {
"image": ("IMAGE",),
"x": ("FLOAT", {"default": 0, "step": 1, "min": -4096, "max": 4096}),
"y": ("FLOAT", {"default": 0, "step": 1, "min": -4096, "max": 4096}),
"zoom": ("FLOAT", {"default": 1.0, "min": 0.001, "step": 0.01}),
"angle": ("FLOAT", {"default": 0, "step": 1, "min": -360, "max": 360}),
"x": (
"FLOAT",
{"default": 0, "step": 1, "min": -4096, "max": 4096},
),
"y": (
"FLOAT",
{"default": 0, "step": 1, "min": -4096, "max": 4096},
),
"zoom": (
"FLOAT",
{"default": 1.0, "min": 0.001, "step": 0.01},
),
"angle": (
"FLOAT",
{"default": 0, "step": 1, "min": -360, "max": 360},
),
"shear": (
"FLOAT",
{"default": 0, "step": 1, "min": -4096, "max": 4096},
@@ -53,14 +66,21 @@ class TransformImage:
y = int(y)
angle = int(angle)
log.debug(f"Zoom: {zoom} | x: {x}, y: {y}, angle: {angle}, shear: {shear}")
log.debug(
f"Zoom: {zoom} | x: {x}, y: {y}, angle: {angle}, shear: {shear}"
)
if image.size(0) == 0:
return (torch.zeros(0),)
transformed_images = []
frames_count, frame_height, frame_width, frame_channel_count = image.size()
frames_count, frame_height, frame_width, frame_channel_count = (
image.size()
)
new_height, new_width = int(frame_height * zoom), int(frame_width * zoom)
new_height, new_width = (
int(frame_height * zoom),
int(frame_width * zoom),
)
log.debug(f"New height: {new_height}, New width: {new_width}")
@@ -74,7 +94,12 @@ class TransformImage:
pw += abs(max_padding)
ph += abs(max_padding)
padding = [max(0, pw + x), max(0, ph + y), max(0, pw - x), max(0, ph - y)]
padding = [
max(0, pw + x),
max(0, ph + y),
max(0, pw - x),
max(0, ph - y),
]
constant_color = hex_to_rgb(constant_color)
log.debug(f"Fill Tuple: {constant_color}")
@@ -89,7 +114,9 @@ class TransformImage:
img = cast(
Image.Image,
TF.affine(img, angle=angle, scale=zoom, translate=[x, y], shear=shear),
TF.affine(
img, angle=angle, scale=zoom, translate=[x, y], shear=shear
),
)
left = abs(padding[0])
@@ -107,4 +134,4 @@ class TransformImage:
return (pil2tensor(transformed_images),)
__nodes__ = [TransformImage]
__nodes__ = [MTB_TransformImage]
+8 -5
View File
@@ -1,4 +1,7 @@
import hashlib, json, os, re
import hashlib
import json
import os
import re
from pathlib import Path
import folder_paths
@@ -10,7 +13,7 @@ from PIL.PngImagePlugin import PngInfo
from ..log import log
class LoadImageSequence:
class MTB_LoadImageSequence:
"""Load an image sequence from a folder. The current frame is used to determine which image to load.
Usually used in conjunction with the `Primitive` node set to increment to load a sequence of images from a folder.
@@ -155,7 +158,7 @@ def resolve_path(path, frame):
return re.sub("#+", padded_number, path)
class SaveImageSequence:
class MTB_SaveImageSequence:
"""Save an image sequence to a folder. The current frame is used to determine which image to save.
This is merely a wrapper around the `save_images` function with formatting for the output folder and filename.
@@ -251,6 +254,6 @@ class SaveImageSequence:
__nodes__ = [
LoadImageSequence,
SaveImageSequence,
MTB_LoadImageSequence,
MTB_SaveImageSequence,
]