fix: onnx runtime version

__init__.py

@@ -9,6 +9,7 @@ import sys
 
 sys.path.append(os.path.join(os.path.dirname(__file__)))
 
+import routes
 import inspect
 import sys
 import importlib
@@ -114,7 +115,6 @@ for path in paths:
 NODE_CLASS_MAPPINGS = {}
 NODE_DISPLAY_NAME_MAPPINGS = {}
 
-import routes
 import blender_node
 
 base_class = blender_node.ObjectOps

blender/mesh_utils.py

@@ -3,7 +3,7 @@ import subprocess
 import os
 import folder_paths
 import requests
-
+import json
 
 def genreate_mesh_from_texture(bpy, image):
     import torch
@@ -14,12 +14,11 @@ def genreate_mesh_from_texture(bpy, image):
     gray = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY)
     gray = (gray * 255).astype(np.uint8)
     # Find contours
-    contours, _ = cv2.findContours(gray, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
+    contours, _ = cv2.findContours(
+        gray, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
 
     if len(contours) == 0:
-        print("Warning: No contours found. The image may have 0 segment.")
-        black_image = torch.zeros(1, *image.shape)
-        return (black_image, None)
+        raise Exception("No contours found. Please ensure that the image has the correct segments (e.g. when you click on the mouth, it should display a proper blue area over the mouth region).")
 
     # Get the largest contour
     areas = [cv2.contourArea(contour) for contour in contours]
@@ -39,12 +38,11 @@ def genreate_mesh_from_texture(bpy, image):
     for contour in contours:
         normalized_contour = []
         for vertex in contour:
-            normalized_vertex = [
-                normalize_vertices(vertex[0][0], width),
-                normalize_vertices(vertex[0][1], height) * -1,
-            ]
+            normalized_vertex = [normalize_vertices(
+                vertex[0][0], width), normalize_vertices(vertex[0][1], height) * -1]
             normalized_contour.append(normalized_vertex)
-        normalized_contours.append(np.array(normalized_contour, dtype=np.float32))
+        normalized_contours.append(
+            np.array(normalized_contour, dtype=np.float32))
 
     meshes = []
     # print(len(normalized_contours))
@@ -65,12 +63,12 @@ def genreate_mesh_from_texture(bpy, image):
         mesh.from_pydata(ordered_vertices, [], [face])
 
         # Create a default shape key for the mesh
-        sk_basis = obj.shape_key_add(name="Basis")
+        sk_basis = obj.shape_key_add(name='Basis')
 
         meshes.append(obj)  # Add the object to the list of meshes
 
     # Draw contours on the original image
-    if not image.flags["C_CONTIGUOUS"]:
+    if not image.flags['C_CONTIGUOUS']:
         image = np.ascontiguousarray(image)
     cv2.drawContours(image, contours, -1, (0, 255, 0), 3)
 
@@ -82,7 +80,7 @@ def genreate_mesh_from_texture(bpy, image):
 def assign_texture(bpy, BPY_OBJ, texture, texture_name):
     import numpy as np
     import time
-
+    
     # Start the timer
     start_time = time.time()
 
@@ -94,8 +92,7 @@ def assign_texture(bpy, BPY_OBJ, texture, texture_name):
 
     # Create an image with the required dimensions
     img = bpy.data.images.new(
-        texture_name, width=texture.shape[1], height=texture.shape[0], alpha=True
-    )
+        texture_name, width=texture.shape[1], height=texture.shape[0], alpha = True)
 
     # If there is no alpha channel, append one full of 1's
     if texture.shape[2] == 3:
@@ -112,7 +109,7 @@ def assign_texture(bpy, BPY_OBJ, texture, texture_name):
     print(f"Time taken (texture.ravel) : {end_time - start_time} seconds")
 
     # End the timer and print the time taken
-
+  
     # Pack image to store it within .blend file
     img.pack()
 
@@ -127,27 +124,28 @@ def assign_texture(bpy, BPY_OBJ, texture, texture_name):
     # Create a material
     mat = bpy.data.materials.new("MaterialName")
     mat.use_nodes = True
-    mat.blend_method = "BLEND"
+    mat.blend_method = 'BLEND'
     nodes = mat.node_tree.nodes
     for node in nodes:
         nodes.remove(node)
 
     # Add a new texture node
-    texture_node = nodes.new(type="ShaderNodeTexImage")
+    texture_node = nodes.new(type='ShaderNodeTexImage')
     texture_node.image = img
 
     # Add a new BSDF node
-    bsdf_node = nodes.new(type="ShaderNodeBsdfPrincipled")
+    bsdf_node = nodes.new(type='ShaderNodeBsdfPrincipled')
 
     # Add a new output node
-    output_node = nodes.new(type="ShaderNodeOutputMaterial")
+    output_node = nodes.new(type='ShaderNodeOutputMaterial')
 
     # Link nodes together
     links = mat.node_tree.links
-    links.new(bsdf_node.inputs["Base Color"], texture_node.outputs["Color"])
-    links.new(output_node.inputs["Surface"], bsdf_node.outputs["BSDF"])
+    links.new(bsdf_node.inputs['Base Color'],
+              texture_node.outputs['Color'])
+    links.new(output_node.inputs['Surface'], bsdf_node.outputs['BSDF'])
 
-    links.new(bsdf_node.inputs["Alpha"], texture_node.outputs["Alpha"])
+    links.new(bsdf_node.inputs['Alpha'], texture_node.outputs['Alpha'])
 
     # Assign the material to the active object
     if obj.data.materials:
@@ -162,30 +160,21 @@ def assign_texture(bpy, BPY_OBJ, texture, texture_name):
 blender_process_global = []
 
 
-def open_in_blender(
-    blender_process,
-    blender_path,
-    output_file,
-    camera_location=(0, 0, 0),
-    camera_rotation=(0, 0, 0),
-    shading="Material",
-):
+def open_in_blender(blender_process, blender_path, output_file, camera_location=(0, 0, 0), camera_rotation=(0, 0, 0), shading="Material"):
     import global_bpy
     import mathutils
-
     bpy = global_bpy.get_bpy()
 
     # Change shading mode and viewport
     for area in bpy.context.screen.areas:
-        if area.type == "VIEW_3D":
+        if area.type == 'VIEW_3D':
             for space in area.spaces:
-                if space.type == "VIEW_3D":
+                if space.type == 'VIEW_3D':
                     space.shading.type = shading.upper()
                     rv3d = space.region_3d
                     rv3d.view_location = camera_location
                     rv3d.view_rotation = mathutils.Euler(
-                        camera_rotation
-                    ).to_quaternion()
+                        camera_rotation).to_quaternion()
 
     # Open blender
     if blender_process != None:
@@ -198,8 +187,8 @@ def open_in_blender(
         os.remove(output_file)
     bpy.ops.wm.save_as_mainfile(filepath=output_file)
 
-    print("blender_path", blender_path)
-    print("output_file", output_file)
+    print('blender_path', blender_path)
+    print('output_file', output_file)
     blender_process = subprocess.Popen([blender_path, output_file])
     # append to global list so it doesn't get garbage collected
     blender_process_global.append(blender_process)
@@ -209,20 +198,18 @@ def open_in_blender(
 
 # detects when the python process is killed, and kills the blender process
 
-
 @atexit.register
 def kill_blender_process():
-    print("blender_process_global", blender_process_global)
+    print('blender_process_global', blender_process_global)
     for process in blender_process_global:
         process.kill()
 
 
 def export_gltf(output_dir, bpy_objects, filename, model_type, write_mode, metadata):
     import global_bpy
-
     bpy = global_bpy.get_bpy()
     # print(bpy, bpy_objects)
-
+    
     # deselect all objects
     override = bpy.context.copy()
     override["selected_objects"] = list(bpy_objects)
@@ -243,48 +230,33 @@ def export_gltf(output_dir, bpy_objects, filename, model_type, write_mode, metad
             return ".ava"
 
     ext = get_file_extension(model_type)
-    filepath = (
-        output_dir + "/" + filename + ext + (".glb" if model_type == "AVA" else "")
-    )
+    filepath = output_dir + "/" + filename + ext + (".glb" if model_type == "AVA" else "")
 
     if write_mode == "Increment":
         count = 0
         # while file exists, increment count
-        while os.path.exists(output_dir + "/" + filename + "_" + str(count) + ext):
+        while os.path.exists(output_dir + "/" + filename + '_' + str(count) + ext):
             count += 1
 
-        filepath = (
-            output_dir
-            + "/"
-            + filename
-            + "_"
-            + str(count)
-            + ext
-            + (".glb" if model_type == "AVA" else "")
-        )
-
+        filepath = output_dir + "/" + filename + '_' + str(count) + ext + (".glb" if model_type == "AVA" else "")
+    
     with bpy.context.temp_override(**override):
-        bpy.ops.export_scene.gltf(
-            filepath=filepath,
-            export_format="GLB" if model_type == "AVA" else model_type,
-            use_selection=True,
-            export_extras=True,
-        )
+        bpy.ops.export_scene.gltf(filepath=filepath, export_format="GLB" if model_type == "AVA" else model_type, use_selection=True, export_extras=True)
         # print(filepath)
-        if filepath.endswith(".ava.glb"):
-            new_filepath = filepath.replace(".ava.glb", ".ava")
+        if filepath.endswith('.ava.glb'):
+            new_filepath = filepath.replace('.ava.glb', '.ava')
             os.replace(filepath, new_filepath)
             filepath = new_filepath
 
     return filepath
 
-
 def get_avatar_file(output):
     avatar_filename = output["gltfFilename"][0]
-    with open(f"{folder_paths.get_output_directory()}/{avatar_filename}", "rb") as f:
+    with open(
+        f"{folder_paths.get_output_directory()}/{avatar_filename}", "rb"
+    ) as f:
         return f.read()
 
-
 def upload_avatar_file(output):
     file = get_avatar_file(output)
     response = requests.get("https://labs.avatech.ai/api/share?version=v2")

blender/ops_create_mesh_layer.py

@@ -24,8 +24,6 @@ class Object_CreateMeshLayer(blender_node.ObjectOps):
 
     def blender_process(self, bpy, image, convex_hull, shape_threshold, mesh_layer_name, scale_x,scale_y , extrude_x, extrude_y, seed):
         image, BPY_OBJ = genreate_mesh_from_texture(bpy, image)
-        if BPY_OBJ is None:
-            return (None, image)
 
         bpy.context.view_layer.objects.active = BPY_OBJ
 

blender/ops_create_mesh_layer_advanced.py

@@ -28,8 +28,6 @@ class Object_CreateMeshLayer_Advanced(blender_node.ObjectOps):
 
     def blender_process(self, bpy, image, convex_hull, shape_threshold, mesh_layer_name, scale_x,scale_y , extrude_x, extrude_y, inner_translate_x, inner_translate_y, outer_translate_x, outer_translate_y, seed):
         image, BPY_OBJ = genreate_mesh_from_texture(bpy, image)
-        if BPY_OBJ is None:
-            return (None, image)
 
         bpy.context.view_layer.objects.active = BPY_OBJ
 

blender/ops_join.py

@@ -2,12 +2,14 @@ import blender_node
 
 
 class Mesh_JoinMesh(blender_node.ObjectOps):
-    EXTRA_INPUT_TYPES = {"BPY_OBJ2": (blender_node.BPY_OBJ,)}
-
+    EXTRA_INPUT_TYPES = {
+        "BPY_OBJ2": (blender_node.BPY_OBJ,)
+    }
+    
     CUSTOM_NAME = "Join Meshes"
 
     def blender_process(self, bpy, BPY_OBJ, **props):
-        prop_values = props.values()
+        prop_values = props.values()        
         for obj in list(prop_values) + [BPY_OBJ]:
             if obj is not None:
                 obj.select_set(True)
@@ -16,3 +18,4 @@ class Mesh_JoinMesh(blender_node.ObjectOps):
         bpy.ops.object.join()
 
         return (BPY_OBJ,)
+

js/index.js

@@ -321,6 +321,7 @@ function showMyImageEditor(node) {
           isGeneratedImage,
           embedding_id: id,
           ckpt,
+          // remote: true,
         }),
       })
       .then(() => {

routes.py

@@ -1,19 +1,18 @@
 from aiohttp import web
+from segment_anything import sam_model_registry, SamPredictor
+from PIL import Image, ImageOps
 from dotenv import load_dotenv
 from blender.mesh_utils import upload_avatar_file
-from sam_utils import (
-    sam_ckpt_to_type,
-    compute_image_embedding,
-    check_embedding_exists,
-    save_embedding,
-    load_image,
-)
 import os
 import requests
 import folder_paths
 import json
+import numpy as np
 import server
+import re
 import base64
+from PIL import Image
+import io
 import time
 import execution
 import random
@@ -68,19 +67,80 @@ async def get_sam_model(request):
     return web.FileResponse(filename)
 
 
+def load_image(image, is_generated_image):
+    if is_generated_image:
+        image_path = f"{folder_paths.get_output_directory()}/{image}"
+    else:
+        image_path = folder_paths.get_annotated_filepath(image)
+    i = Image.open(image_path)
+    i = ImageOps.exif_transpose(i)
+    image = i.convert("RGB")
+    image = np.array(image).astype(np.float32) / 255.0
+    return image
+
+
 @server.PromptServer.instance.routes.post("/sam_model")
 async def post_sam_model(request):
     post = await request.json()
     is_generated_image = post.get("isGeneratedImage")
     emb_id = post.get("embedding_id")
     ckpt = post.get("ckpt")
-    model_type = sam_ckpt_to_type[ckpt]
-    if not check_embedding_exists(emb_id, model_type):
+    ckpt = folder_paths.get_full_path("sams", ckpt)
+    remote = post.get("remote")
+    model_type = re.findall(r"vit_[lbh]", ckpt)[0]
+    emb_filename = f"{folder_paths.get_output_directory()}/{emb_id}_{model_type}.npy"
+    output_json_filename = (
+        f"{folder_paths.get_output_directory()}/{emb_id}_{model_type}.json"
+    )
+    if not os.path.exists(emb_filename):
         image = load_image(post.get("image"), is_generated_image)
-        emb, img_model_input_size, img_original_size = compute_image_embedding(
-            image, model_type
-        )
-        save_embedding(emb_id, model_type, emb, img_model_input_size, img_original_size)
+        if remote:
+            # Run embed in remote server
+            image = Image.fromarray((image * 255).astype(np.uint8))
+            buffered = io.BytesIO()
+            image.save(buffered, format="PNG")
+            image = base64.b64encode(buffered.getvalue()).decode()
+            res = requests.post(
+                "https://avatechgg--sam-embed.modal.run",
+                headers={
+                    "Content-type": "application/json",
+                    "Accept": "application/json",
+                },
+                data=json.dumps(
+                    {
+                        "image": image,
+                    }
+                ),
+            ).json()
+            emb, input_size, original_size = (
+                res["emb"],
+                res["input_size"],
+                res["original_size"],
+            )
+            emb = np.array(emb).astype(np.float32)
+            np.save(emb_filename, emb)
+            with open(output_json_filename, "w") as f:
+                data = {
+                    "input_size": input_size,
+                    "original_size": original_size,
+                }
+                json.dump(data, f)
+        else:
+            sam = sam_model_registry[model_type](checkpoint=ckpt)
+            predictor = SamPredictor(sam)
+
+            image_np = (image * 255).astype(np.uint8)
+            predictor.set_image(image_np)
+            emb = predictor.get_image_embedding().cpu().numpy()
+            np.save(emb_filename, emb)
+            with open(output_json_filename, "w") as f:
+                json.dump(
+                    {
+                        "input_size": predictor.input_size,
+                        "original_size": predictor.original_size,
+                    },
+                    f,
+                )
     print("Finished embedding")
     return web.json_response({})
 
@@ -162,6 +222,7 @@ def load_workflow(workflow_name):
         return "\n".join(f.readlines())
 
 
+
 @server.PromptServer.instance.routes.post("/avatar_generation")
 async def post_prompt_block(request):
     prompt_server = server.PromptServer.instance
@@ -224,7 +285,6 @@ async def post_prompt_block(request):
                     return web.json_response({"id": modelId}, status=200)
         time.sleep(0.5)
 
-
 # TODO: refactor the code
 @server.PromptServer.instance.routes.post("/rendering_generation")
 async def post_data_generation(request):
@@ -252,7 +312,6 @@ async def post_data_generation(request):
                         return web.json_response({"image": filename}, status=200)
         time.sleep(0.5)
 
-
 @server.PromptServer.instance.routes.post("/image_generation")
 async def post_image_generation(request):
     prompt_server = server.PromptServer.instance

sam/mediapipe_utils.py

@@ -1,254 +0,0 @@
-# For auto-segmentation
-import mediapipe as mp
-import numpy as np
-import os
-from math import sqrt
-
-face_landmarker = None
-pose_landmarker = None
-
-BaseOptions = mp.tasks.BaseOptions
-FaceLandmarker = mp.tasks.vision.FaceLandmarker
-FaceLandmarkerOptions = mp.tasks.vision.FaceLandmarkerOptions
-PoseLandmarker = mp.tasks.vision.PoseLandmarker
-PoseLandmarkerOptions = mp.tasks.vision.PoseLandmarkerOptions
-VisionRunningMode = mp.tasks.vision.RunningMode
-
-layerMapping = {
-    "L_eye": {
-        "useMiddle": False,
-        "positiveOffsetX": 0,
-        "positiveOffsetY": 0,
-        "negativeOffsetX": 0,
-        "negativeOffsetY": 0,
-        "positiveScale": 0,
-        "negativeScale": 0.5,
-        "indices": mp.solutions.face_mesh.FACEMESH_LEFT_EYE,
-    },
-    "R_eye": {
-        "useMiddle": False,
-        "positiveOffsetX": 0,
-        "positiveOffsetY": 0,
-        "negativeOffsetX": 0,
-        "negativeOffsetY": 0,
-        "positiveScale": 0,
-        "negativeScale": 0.5,
-        "indices": mp.solutions.face_mesh.FACEMESH_RIGHT_EYE,
-    },
-    "L_iris": {
-        "useMiddle": False,
-        "positiveOffsetX": 0,
-        "positiveOffsetY": 0,
-        "negativeOffsetX": 0,
-        "negativeOffsetY": 0,
-        "positiveScale": -0.2,
-        "negativeScale": 0.5,
-        "indices": mp.solutions.face_mesh.FACEMESH_LEFT_IRIS,
-    },
-    "R_iris": {
-        "useMiddle": False,
-        "positiveOffsetX": 0,
-        "positiveOffsetY": 0,
-        "negativeOffsetX": 0,
-        "negativeOffsetY": 0,
-        "positiveScale": -0.2,
-        "negativeScale": 0.5,
-        "indices": mp.solutions.face_mesh.FACEMESH_RIGHT_IRIS,
-    },
-    "face": {
-        "useMiddle": False,
-        "positiveOffsetX": 0,
-        "positiveOffsetY": 40,
-        "negativeOffsetX": 0,
-        "negativeOffsetY": 60,
-        "positiveScale": 0.2,
-        "negativeScale": 0.6,
-        "indices": mp.solutions.face_mesh.FACEMESH_FACE_OVAL,
-    },
-    "mouth": {
-        "useMiddle": False,
-        "positiveOffsetX": 0,
-        "positiveOffsetY": 0,
-        "negativeOffsetX": 0,
-        "negativeOffsetY": 0,
-        "positiveScale": -0.3,
-        "negativeScale": 0.3,
-        # https://stackoverflow.com/questions/66649492/how-to-get-specific-landmark-of-face-like-lips-or-eyes-using-tensorflow-js-face
-        "indices": [[x, x] for x in [61, 37, 270, 91, 314]],
-    },
-    "mouth_in": {
-        "useMiddle": False,
-        "positiveOffsetX": 0,
-        "positiveOffsetY": 0,
-        "negativeOffsetX": 0,
-        "negativeOffsetY": 0,
-        "positiveScale": -0.5,
-        "negativeScale": 0.5,
-        # https://stackoverflow.com/questions/66649492/how-to-get-specific-landmark-of-face-like-lips-or-eyes-using-tensorflow-js-face
-        "indices": [[x, x] for x in [310, 88]],
-    },
-}
-
-
-def load_mediapipe_models():
-    global face_landmarker, pose_landmarker
-    if face_landmarker is None and pose_landmarker is None:
-        face_landmarker_model_path = os.path.join(
-            os.path.dirname(__file__), "../mediapipe_models/face_landmarker.task"
-        )
-        face_landmarker_options = FaceLandmarkerOptions(
-            base_options=BaseOptions(model_asset_path=face_landmarker_model_path),
-            running_mode=VisionRunningMode.IMAGE,
-        )
-        face_landmarker = FaceLandmarker.create_from_options(face_landmarker_options)
-
-        pose_landmarker_model_path = os.path.join(
-            os.path.dirname(__file__), "../mediapipe_models/pose_landmarker_full.task"
-        )
-        pose_landmarker_options = PoseLandmarkerOptions(
-            base_options=BaseOptions(model_asset_path=pose_landmarker_model_path),
-            running_mode=VisionRunningMode.IMAGE,
-        )
-        pose_landmarker = PoseLandmarker.create_from_options(pose_landmarker_options)
-    return face_landmarker, pose_landmarker
-
-
-def auto_segment_face(image, face_landmarks):
-    H, W, C = image.shape
-    layer_points = {}
-    layer_bboxes = {}
-
-    for key, value in layerMapping.items():
-        positivePoints = []
-        middlePoints = []
-        negativePoints = []
-
-        for index in value["indices"]:
-            start, end = index
-            startPoint = face_landmarks[start]
-
-            startX = startPoint.x * W
-            startY = startPoint.y * H
-
-            if len(middlePoints) == 0:
-                middlePoints.append({"x": startX, "y": startY, "label": 1})
-            else:
-                middlePoints[0]["x"] += startX
-                middlePoints[0]["y"] += startY
-
-            positivePoints.append({"x": startX, "y": startY, "label": 1})
-
-        len_indices = len(value["indices"])
-        middlePoints[0]["x"] /= len_indices
-        middlePoints[0]["y"] /= len_indices
-
-        if value["useMiddle"]:
-            layer_points[key] = middlePoints
-        else:
-            for i, index in enumerate(value["indices"]):
-                start, end = index
-                startPoint = face_landmarks[start]
-
-                startX = startPoint.x * W
-                startY = startPoint.y * H
-
-                middlePoint = middlePoints[0]
-                directionVector = {
-                    "x": middlePoint["x"] - startX,
-                    "y": middlePoint["y"] - startY,
-                }
-                directionVectorLength = sqrt(
-                    directionVector["x"] * directionVector["x"]
-                    + directionVector["y"] * directionVector["y"]
-                )
-
-                if value["negativeScale"] != 0:
-                    negativePointDistance = (
-                        value["negativeScale"] * directionVectorLength
-                    )
-                    negativePoint = {
-                        "x": startX
-                        - (negativePointDistance * directionVector["x"])
-                        / directionVectorLength
-                        - value["negativeOffsetX"],
-                        "y": startY
-                        - (negativePointDistance * directionVector["y"])
-                        / directionVectorLength
-                        - value["negativeOffsetY"],
-                        "label": 0,
-                    }
-                    negativePoints.append(negativePoint)
-
-                positivePointDistance = value["positiveScale"] * directionVectorLength
-                positivePoints[i] = {
-                    "x": positivePoints[i]["x"]
-                    - (positivePointDistance * directionVector["x"])
-                    / directionVectorLength
-                    - value["positiveOffsetX"],
-                    "y": positivePoints[i]["y"]
-                    - (positivePointDistance * directionVector["y"])
-                    / directionVectorLength
-                    - value["positiveOffsetY"],
-                    "label": 1,
-                }
-
-            layer_points[key] = positivePoints + negativePoints
-
-        points = negativePoints if len(negativePoints) > 0 else positivePoints
-        box = np.array(
-            [
-                min(x["x"] for x in points),
-                min(x["y"] for x in points),
-                max(x["x"] for x in points),
-                max(x["y"] for x in points),
-            ]
-        )
-        layer_bboxes[key] = box
-
-    return layer_points, layer_bboxes
-
-
-def auto_segment_pose(image, pose_landmarks):
-    H, W, C = image.shape
-    layer_points = {}
-    if pose_landmarks is not None:
-        positiveBreathX = ((pose_landmarks[11].x + pose_landmarks[12].x) / 2) * W
-        positiveBreathY = ((pose_landmarks[11].y + pose_landmarks[12].y) / 2) * H
-        negativeBreathX1 = pose_landmarks[0].x * W
-        negativeBreathY1 = pose_landmarks[0].y * H
-        negativeBreathX2 = pose_landmarks[9].x * W
-        negativeBreathY2 = pose_landmarks[9].y * H
-        negativeBreathX3 = pose_landmarks[10].x * W
-        negativeBreathY3 = pose_landmarks[10].y * H
-        layer_points["breath"] = [
-            {"x": positiveBreathX, "y": positiveBreathY, "label": 1},
-            {"x": negativeBreathX1, "y": negativeBreathY1, "label": 0},
-            {"x": negativeBreathX2, "y": negativeBreathY2, "label": 0},
-            {"x": negativeBreathX3, "y": negativeBreathY3, "label": 0},
-        ]
-    return layer_points
-
-
-def detect_face(np_image):
-    face_landmarker, pose_landmarker = load_mediapipe_models()
-    mp_image = mp.Image(
-        image_format=mp.ImageFormat.SRGB, data=(np_image * 255).astype(np.uint8)
-    )
-
-    face_landmarks = face_landmarker.detect(mp_image).face_landmarks
-    if len(face_landmarks) > 0:
-        face_points, face_bboxes = auto_segment_face(np_image, face_landmarks[0])
-    else:
-        face_points, face_bboxes = {}, {}
-        print("Warning: no face detected")
-
-    pose_landmarks = pose_landmarker.detect(mp_image).pose_landmarks
-    if len(pose_landmarks) > 0:
-        pose_points = auto_segment_pose(np_image, pose_landmarks[0])
-    else:
-        pose_points = {}
-        print("Warning: no pose detected")
-
-    layer_points = {**face_points, **pose_points}
-    layer_bboxes = {**face_bboxes}
-    return layer_points, layer_bboxes

sam/sam_multilayer.py

@@ -5,19 +5,105 @@ import torch
 import re
 import json
 import uuid
-from sam_utils import (
-    load_model,
-    check_embedding_exists,
-    compute_image_embedding,
-    save_embedding,
-    load_embdding,
-    load_image,
-)
-from mediapipe_utils import detect_face
+from segment_anything import sam_model_registry, SamPredictor
 from einops import rearrange, repeat
+from PIL import Image
+import mediapipe as mp
+from math import sqrt
+
+BaseOptions = mp.tasks.BaseOptions
+FaceLandmarker = mp.tasks.vision.FaceLandmarker
+FaceLandmarkerOptions = mp.tasks.vision.FaceLandmarkerOptions
+PoseLandmarker = mp.tasks.vision.PoseLandmarker
+PoseLandmarkerOptions = mp.tasks.vision.PoseLandmarkerOptions
+VisionRunningMode = mp.tasks.vision.RunningMode
+
+global_predictor = None
+face_landmarker = None
+pose_landmarker = None
+
+# For auto-segmentation
+layerMapping = {
+    "L_eye": {
+        "useMiddle": False,
+        "positiveOffsetX": 0,
+        "positiveOffsetY": 0,
+        "negativeOffsetX": 0,
+        "negativeOffsetY": 0,
+        "positiveScale": 0,
+        "negativeScale": 0.5,
+        "indices": mp.solutions.face_mesh.FACEMESH_LEFT_EYE,
+    },
+    "R_eye": {
+        "useMiddle": False,
+        "positiveOffsetX": 0,
+        "positiveOffsetY": 0,
+        "negativeOffsetX": 0,
+        "negativeOffsetY": 0,
+        "positiveScale": 0,
+        "negativeScale": 0.5,
+        "indices": mp.solutions.face_mesh.FACEMESH_RIGHT_EYE,
+    },
+    "L_iris": {
+        "useMiddle": False,
+        "positiveOffsetX": 0,
+        "positiveOffsetY": 0,
+        "negativeOffsetX": 0,
+        "negativeOffsetY": 0,
+        "positiveScale": -0.2,
+        "negativeScale": 0.5,
+        "indices": mp.solutions.face_mesh.FACEMESH_LEFT_IRIS,
+    },
+    "R_iris": {
+        "useMiddle": False,
+        "positiveOffsetX": 0,
+        "positiveOffsetY": 0,
+        "negativeOffsetX": 0,
+        "negativeOffsetY": 0,
+        "positiveScale": -0.2,
+        "negativeScale": 0.5,
+        "indices": mp.solutions.face_mesh.FACEMESH_RIGHT_IRIS,
+    },
+    "face": {
+        "useMiddle": False,
+        "positiveOffsetX": 0,
+        "positiveOffsetY": 40,
+        "negativeOffsetX": 0,
+        "negativeOffsetY": 60,
+        "positiveScale": 0.2,
+        "negativeScale": 0.6,
+        "indices": mp.solutions.face_mesh.FACEMESH_FACE_OVAL,
+    },
+    "mouth": {
+        "useMiddle": False,
+        "positiveOffsetX": 0,
+        "positiveOffsetY": 0,
+        "negativeOffsetX": 0,
+        "negativeOffsetY": 0,
+        "positiveScale": -0.3,
+        "negativeScale": 0.3,
+        # https://stackoverflow.com/questions/66649492/how-to-get-specific-landmark-of-face-like-lips-or-eyes-using-tensorflow-js-face
+        "indices": [[x, x] for x in [61, 37, 270, 91, 314]],
+    },
+    "mouth_in": {
+        "useMiddle": False,
+        "positiveOffsetX": 0,
+        "positiveOffsetY": 0,
+        "negativeOffsetX": 0,
+        "negativeOffsetY": 0,
+        "positiveScale": -0.5,
+        "negativeScale": 0.5,
+        # https://stackoverflow.com/questions/66649492/how-to-get-specific-landmark-of-face-like-lips-or-eyes-using-tensorflow-js-face
+        "indices": [[x, x] for x in [310, 88]],
+    },
+}
 
 
 class SAMMultiLayer:
+    def __init__(self):
+        self.predictor = None
+        self.output_dir = folder_paths.get_output_directory()
+
     @classmethod
     def INPUT_TYPES(s):
         return {
@@ -34,31 +120,191 @@ class SAMMultiLayer:
 
     CATEGORY = "image"
 
-    RETURN_TYPES = ["SAM_PROMPT"] # + ["IMAGE"] * 100
-    FUNCTION = "run"
+    RETURN_TYPES = ("SAM_PROMPT",)
+    FUNCTION = "load_image"
 
-    def run(self, image, ckpt, embedding_id, image_prompts_json):
-        if (
-            "COMFY_DEPLOY" in os.environ
-            and os.getenv("COMFY_DEPLOY", "FALSE") == "TRUE"
-        ):
-            embedding_id = str(uuid.uuid4())
-        layer_points = json.loads(image_prompts_json.replace("'", '"'))
+    def load_models(self, ckpt, model_type):
+        global global_predictor, face_landmarker, pose_landmarker
 
-        order_file = (
-            f"{folder_paths.get_output_directory()}/segments_{embedding_id}/order.json"
+        ckpt = folder_paths.get_full_path("sams", ckpt)
+        sam = sam_model_registry[model_type](checkpoint=ckpt) # .to("cuda")
+        global_predictor = SamPredictor(sam)
+
+        face_landmarker_model_path = os.path.join(
+            os.path.dirname(__file__), "../mediapipe_models/face_landmarker.task"
         )
+        face_landmarker_options = FaceLandmarkerOptions(
+            base_options=BaseOptions(model_asset_path=face_landmarker_model_path),
+            running_mode=VisionRunningMode.IMAGE,
+        )
+        face_landmarker = FaceLandmarker.create_from_options(face_landmarker_options)
+
+        pose_landmarker_model_path = os.path.join(
+            os.path.dirname(__file__), "../mediapipe_models/pose_landmarker_full.task"
+        )
+        pose_landmarker_options = PoseLandmarkerOptions(
+            base_options=BaseOptions(model_asset_path=pose_landmarker_model_path),
+            running_mode=VisionRunningMode.IMAGE,
+        )
+        pose_landmarker = PoseLandmarker.create_from_options(pose_landmarker_options)
+        return global_predictor, face_landmarker, pose_landmarker
+
+    def auto_segment(self, image, face_landmarks, pose_landmarks):
+        H, W, C = image.shape
+        imagePromptsMulti = {}
+        boxesMulti = {}
+
+        for key, value in layerMapping.items():
+            positivePoints = []
+            middlePoints = []
+            negativePoints = []
+
+            for index in value["indices"]:
+                start, end = index
+                startPoint = face_landmarks[start]
+
+                startX = startPoint.x * W
+                startY = startPoint.y * H
+
+                if len(middlePoints) == 0:
+                    middlePoints.append({"x": startX, "y": startY, "label": 1})
+                else:
+                    middlePoints[0]["x"] += startX
+                    middlePoints[0]["y"] += startY
+
+                positivePoints.append({"x": startX, "y": startY, "label": 1})
+
+            len_indices = len(value["indices"])
+            middlePoints[0]["x"] /= len_indices
+            middlePoints[0]["y"] /= len_indices
+
+            if value["useMiddle"]:
+                imagePromptsMulti[key] = middlePoints
+            else:
+                for i, index in enumerate(value["indices"]):
+                    start, end = index
+                    startPoint = face_landmarks[start]
+
+                    startX = startPoint.x * W
+                    startY = startPoint.y * H
+
+                    middlePoint = middlePoints[0]
+                    directionVector = {
+                        "x": middlePoint["x"] - startX,
+                        "y": middlePoint["y"] - startY,
+                    }
+                    directionVectorLength = sqrt(
+                        directionVector["x"] * directionVector["x"]
+                        + directionVector["y"] * directionVector["y"]
+                    )
+
+                    if value["negativeScale"] != 0:
+                        negativePointDistance = (
+                            value["negativeScale"] * directionVectorLength
+                        )
+                        negativePoint = {
+                            "x": startX
+                            - (negativePointDistance * directionVector["x"])
+                            / directionVectorLength
+                            - value["negativeOffsetX"],
+                            "y": startY
+                            - (negativePointDistance * directionVector["y"])
+                            / directionVectorLength
+                            - value["negativeOffsetY"],
+                            "label": 0,
+                        }
+                        negativePoints.append(negativePoint)
+
+                    positivePointDistance = (
+                        value["positiveScale"] * directionVectorLength
+                    )
+                    positivePoints[i] = {
+                        "x": positivePoints[i]["x"]
+                        - (positivePointDistance * directionVector["x"])
+                        / directionVectorLength
+                        - value["positiveOffsetX"],
+                        "y": positivePoints[i]["y"]
+                        - (positivePointDistance * directionVector["y"])
+                        / directionVectorLength
+                        - value["positiveOffsetY"],
+                        "label": 1,
+                    }
+
+                imagePromptsMulti[key] = positivePoints + negativePoints
+
+            points = negativePoints if len(negativePoints) > 0 else positivePoints
+            box = np.array(
+                [
+                    min(x["x"] for x in points),
+                    min(x["y"] for x in points),
+                    max(x["x"] for x in points),
+                    max(x["y"] for x in points),
+                ]
+            )
+            boxesMulti[key] = box
+
+            if pose_landmarks is not None:
+                positiveBreathX = (
+                    (pose_landmarks[11].x + pose_landmarks[12].x) / 2
+                ) * W
+                positiveBreathY = (
+                    (pose_landmarks[11].y + pose_landmarks[12].y) / 2
+                ) * H
+                negativeBreathX1 = pose_landmarks[0].x * W
+                negativeBreathY1 = pose_landmarks[0].y * H
+                negativeBreathX2 = pose_landmarks[9].x * W
+                negativeBreathY2 = pose_landmarks[9].y * H
+                negativeBreathX3 = pose_landmarks[10].x * W
+                negativeBreathY3 = pose_landmarks[10].y * H
+                imagePromptsMulti["breath"] = [
+                    {"x": positiveBreathX, "y": positiveBreathY, "label": 1},
+                    {"x": negativeBreathX1, "y": negativeBreathY1, "label": 0},
+                    {"x": negativeBreathX2, "y": negativeBreathY2, "label": 0},
+                    {"x": negativeBreathX3, "y": negativeBreathY3, "label": 0},
+                ]
+
+        return imagePromptsMulti, boxesMulti
+
+    def detect_face(self, np_image):
+        global face_landmarker, pose_landmarker
+        mp_image = mp.Image(
+            image_format=mp.ImageFormat.SRGB, data=(np_image * 255).astype(np.uint8)
+        )
+        face_landmarks = face_landmarker.detect(mp_image).face_landmarks
+        if len(face_landmarks) == 0:
+            print("Warning: no face detected")
+            return None, None
+
+        pose_landmarks = pose_landmarker.detect(mp_image).pose_landmarks
+        if len(pose_landmarks) == 0:
+            print("Warning: no pose detected")
+            return None, None
+
+        imagePromptsMulti, boxesMulti = self.auto_segment(
+            np_image, face_landmarks[0], pose_landmarks[0]
+        )
+
+        return imagePromptsMulti, boxesMulti
+
+    def load_image(self, image, ckpt, embedding_id, image_prompts_json):
+        if 'COMFY_DEPLOY' in os.environ and os.getenv('COMFY_DEPLOY', 'FALSE') == "TRUE":
+            embedding_id = str(uuid.uuid4())
+        image_prompts = json.loads(image_prompts_json.replace("'", '"'))
+
+        order_file = f"{self.output_dir}/segments_{embedding_id}/order.json"
         if os.path.exists(order_file):
             # Frontend uploads segments images to backend => backend reads all segments images and passes them to next nodes
             with open(order_file) as f:
                 order = json.load(f)
 
-            result = [layer_points]
+            result = [image_prompts]
+
             for segment in order:
-                image = load_image(
-                    f"{folder_paths.get_output_directory()}/segments_{embedding_id}/{segment}.png",
-                    comfyui_format=True,
+                image = Image.open(
+                    f"{self.output_dir}/segments_{embedding_id}/{segment}.png"
                 )
+                image = np.array(image).astype(np.float32) / 255.0
+                image = torch.from_numpy(image)[None,]
                 result.append(image)
 
             return result
@@ -66,60 +312,72 @@ class SAMMultiLayer:
             # Frontend uploads clicks coordinates to backend => backend runs SAM and passes the segments to next nodes
             model_type = re.findall(r"vit_[lbh]", ckpt)[0]
 
-            # get first image from batch
-            image = image[0]
-            if image.shape[2] == 4:
-                # to RGB
-                image = image[:, :, :3]
+            global global_predictor
+            if global_predictor is None:
+                global_predictor, _, _ = self.load_models(ckpt, model_type)
 
-            if not check_embedding_exists(embedding_id, model_type):
-                emb, img_model_input_size, img_original_size = compute_image_embedding(
-                    image, model_type
-                )
-                save_embedding(
-                    embedding_id,
-                    model_type,
-                    emb,
-                    img_model_input_size,
-                    img_original_size,
-                )
+            if image.shape[3] == 4:
+                image = image[:, :, :, :3]
+
+            emb_filename = f"{self.output_dir}/{embedding_id}_{model_type}.npy"
+            if not os.path.exists(emb_filename):
+                image_np = (image[0].numpy() * 255).astype(np.uint8)
+                global_predictor.set_image(image_np)
+                emb = global_predictor.get_image_embedding().cpu().numpy()
+                np.save(emb_filename, emb)
+
+                with open(
+                    f"{self.output_dir}/{embedding_id}_{model_type}.json", "w"
+                ) as f:
+                    data = {
+                        "input_size": global_predictor.input_size,
+                        "original_size": global_predictor.original_size,
+                    }
+                    json.dump(data, f)
             else:
-                load_embdding(embedding_id, model_type)
+                emb = np.load(emb_filename)
 
-            detected_points, detected_bboxes = detect_face(image.numpy())
-            result = [layer_points]
-            for layer, points in layer_points.items():
-                if detected_points is not None and layer in detected_points:
-                    # use detected points by mediapipe
-                    points = detected_points[layer]
+                with open(f"{self.output_dir}/{embedding_id}_{model_type}.json") as f:
+                    data = json.load(f)
+                    global_predictor.input_size = data["input_size"]
+                    global_predictor.features = torch.from_numpy(emb)
+                    global_predictor.is_image_set = True
+                    global_predictor.original_size = data["original_size"]
 
-                if len(points) == 0:
-                    # no points, append a black image
-                    h, w, c = image.shape
-                    result.append(torch.zeros(1, h, w, c))
-                    print("No points for layer", layer)
-                    continue
+            imagePromptsMulti, boxesMulti = self.detect_face(image[0].numpy())
 
-                # prepare for SAM inferencing
-                point_coords = np.array([[p["x"], p["y"]] for p in points])
-                point_labels = np.array([p["label"] for p in points])
-                bbox = (
-                    detected_bboxes[layer]
-                    if detected_bboxes is not None and layer in detected_bboxes
-                    else None
-                )
+            image_prompts = json.loads(image_prompts_json.replace("'", '"'))
+            result = [image_prompts]  # use imagePromptsMulti
 
-                sam_predictor = load_model(model_type)["predictor"]
-                masks, _, _ = sam_predictor.predict(
-                    point_coords=point_coords,
-                    point_labels=point_labels,
-                    box=bbox,
-                )
-                masks = torch.from_numpy(masks)
-                masks = rearrange(masks[0], "h w -> 1 h w")
-                out_image = repeat(masks, "1 h w -> 1 h w c", c=3) * image.unsqueeze(0)
-                result.append(out_image)
-        return result
+            if isinstance(image_prompts, list):
+                pass
+            elif all(isinstance(item, list) for item in image_prompts.values()):
+                for key, item in image_prompts.items():
+                    if len(item) == 0:
+                        h, w, c = image[0].shape
+                        result.append(torch.zeros(1, h, w, c))
+                        continue
+
+                    points = (
+                        imagePromptsMulti[key]
+                        if imagePromptsMulti is not None and key in imagePromptsMulti
+                        else item
+                    )
+                    point_coords = np.array([[p["x"], p["y"]] for p in points])
+                    point_labels = np.array([p["label"] for p in points])
+
+                    masks, _, _ = global_predictor.predict(
+                        point_coords=point_coords,
+                        point_labels=point_labels,
+                        box=boxesMulti[key]
+                        if boxesMulti is not None and key in boxesMulti
+                        else None,
+                    )
+                    masks = torch.from_numpy(masks)
+                    masks = rearrange(masks[0], "h w -> 1 h w")
+                    out_image = repeat(masks, "1 h w -> 1 h w c", c=3) * image
+                    result.append(out_image)
+            return result
 
 
 NODE_CLASS_MAPPINGS = {"SAM MultiLayer": SAMMultiLayer}

sam/sam_utils.py

@@ -1,97 +0,0 @@
-from segment_anything import sam_model_registry, SamPredictor
-from PIL import Image, ImageOps
-import folder_paths
-import numpy as np
-import torch
-import os
-import json
-
-sam_type_to_ckpt = {
-    "vit_h": "sam_vit_h_4b8939.pth",
-    "vit_l": "sam_vit_l_0b3195.pth",
-    "vit_b": "sam_vit_b_01ec64.pth",
-}
-sam_ckpt_to_type = {v: k for k, v in sam_type_to_ckpt.items()}
-
-sam_instance = {"model_type": None, "model": None, "predictor": None}
-
-
-def load_model(model_type):
-    global sam_instance
-    if sam_instance["model"] is None or sam_instance["model_type"] != model_type:
-        ckpt = sam_type_to_ckpt[model_type]
-        ckpt = folder_paths.get_full_path("sams", ckpt)
-        sam_instance["model_type"] = model_type
-        sam_instance["model"] = sam_model_registry[model_type](checkpoint=ckpt)
-        if torch.cuda.is_available():
-            sam_instance["model"].cuda()
-        sam_instance["predictor"] = SamPredictor(sam_instance["model"])
-    return sam_instance
-
-
-def check_embedding_exists(emb_id, model_type):
-    emb_filename = f"{folder_paths.get_output_directory()}/{emb_id}_{model_type}.npy"
-    return os.path.exists(emb_filename)
-
-
-def save_embedding(emb_id, model_type, emb, img_input_size, img_original_size):
-    emb_filename = f"{folder_paths.get_output_directory()}/{emb_id}_{model_type}.npy"
-    np.save(emb_filename, emb)
-
-    json_filename = f"{folder_paths.get_output_directory()}/{emb_id}_{model_type}.json"
-    with open(json_filename, "w") as f:
-        json.dump(
-            {
-                "input_size": img_input_size,
-                "original_size": img_original_size,
-            },
-            f,
-        )
-
-
-def load_embdding(emb_id, model_type):
-    emb_filename = f"{folder_paths.get_output_directory()}/{emb_id}_{model_type}.npy"
-    emb = np.load(emb_filename)
-
-    json_filename = f"{folder_paths.get_output_directory()}/{emb_id}_{model_type}.json"
-    with open(json_filename, "r") as f:
-        sizes = json.load(f)
-
-    predictor = load_model(model_type)["predictor"]
-    predictor.input_size = sizes["input_size"]
-    predictor.features = torch.from_numpy(emb)
-    predictor.is_image_set = True
-    predictor.original_size = sizes["original_size"]
-
-
-@torch.no_grad()
-def compute_image_embedding(image, model_type="vit_h"):
-    sam = load_model(model_type)
-    predictor = sam["predictor"]
-
-    # if image.shape[3] == 4:
-    #     image = image[:, :, :, :3]
-
-    if torch.is_tensor(image):
-        image = image.numpy()
-
-    image_np = (image * 255).astype(np.uint8)
-    predictor.set_image(image_np)
-    emb = predictor.get_image_embedding().cpu().numpy()
-
-    return emb, predictor.input_size, predictor.original_size
-
-
-def load_image(image, is_generated_image=False, comfyui_format=False):
-    if is_generated_image:
-        image_path = f"{folder_paths.get_output_directory()}/{image}"
-    else:
-        image_path = folder_paths.get_annotated_filepath(image)
-    i = Image.open(image_path)
-    i = ImageOps.exif_transpose(i)
-    image = i.convert("RGB")
-    image = np.array(image).astype(np.float32) / 255.0
-    if comfyui_format:
-        # to torch and create batch dimension
-        image = torch.from_numpy(image)[None,]
-    return image