fix: onnx runtime version

# Conflicts:
#	blender/avatar_main_output.py

__init__.py

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

blender/avatar_main_output.py

@@ -79,6 +79,7 @@ class AvatarMainOutput(blender_node.ObjectOps):
             "files": [{
                 "filename": filepath.replace(f"{self.output_dir}/", ""),
                 "content_type": "model/gltf+json",
+                "type": "output"
             },],
             "SHAPE_FLOW": {SHAPE_FLOW}, 
             "auto_save": {'true' if auto_save else 'false'}, 

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,11 +14,12 @@ 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:
-        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).")
+        print("Warning: No contours found. The image may have 0 segment.")
+        black_image = torch.zeros(1, *image.shape)
+        return (black_image, None)
 
     # Get the largest contour
     areas = [cv2.contourArea(contour) for contour in contours]
@@ -38,11 +39,12 @@ 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))
@@ -63,12 +65,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)
 
@@ -80,7 +82,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()
 
@@ -92,7 +94,8 @@ 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:
@@ -109,7 +112,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()
 
@@ -124,28 +127,27 @@ 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:
@@ -160,21 +162,30 @@ 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:
@@ -187,8 +198,8 @@ def open_in_blender(blender_process, blender_path, output_file, camera_location=
         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)
@@ -198,18 +209,20 @@ def open_in_blender(blender_process, blender_path, output_file, camera_location=
 
 # 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)
@@ -230,36 +243,51 @@ 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")
+    response = requests.get("https://labs.avatech.ai/api/share?version=v2")
     labData = response.json()
     modelId = labData["modelId"]
 
@@ -272,15 +300,15 @@ def upload_avatar_file(output):
     requests.put(labData["url"], headers=headers, data=file)
 
     # send notification
-    webhook_url = os.getenv("DISCORD_WEBHOOK_URL")
-    data = {
-        "username": "Avabot",
-        "avatar_url": "https://avatech-avatar-dev1.nyc3.cdn.digitaloceanspaces.com/avatechai.png",
-        "content": "[API Call] New register!",
-    }
-    headers = {
-        "Content-Type": "application/json",
-    }
-    response = requests.post(webhook_url, headers=headers, data=json.dumps(data))
+    # webhook_url = os.getenv("DISCORD_WEBHOOK_URL")
+    # data = {
+    #     "username": "Avabot",
+    #     "avatar_url": "https://avatech-avatar-dev1.nyc3.cdn.digitaloceanspaces.com/avatechai.png",
+    #     "content": "[API Call] New register!",
+    # }
+    # headers = {
+    #     "Content-Type": "application/json",
+    # }
+    # response = requests.post(webhook_url, headers=headers, data=json.dumps(data))
 
     return modelId

blender/ops_create_mesh_layer.py

@@ -24,6 +24,8 @@ 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
 
@@ -41,7 +43,7 @@ class Object_CreateMeshLayer(blender_node.ObjectOps):
         bpy.ops.mesh.select_all(action='SELECT')
         bpy.ops.mesh.edge_face_add()
 
-        bpy.ops.transform.resize(value=(scale_x, scale_y, 1))
+        bpy.ops.transform.resize(value=(float(scale_x), float(scale_y), 1))
 
         bpy.context.object.vertex_groups.new(name=mesh_layer_name)
         bpy.ops.object.vertex_group_assign()

blender/ops_create_mesh_layer_advanced.py

@@ -28,6 +28,8 @@ 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,14 +2,12 @@ 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)
@@ -18,4 +16,3 @@ class Mesh_JoinMesh(blender_node.ObjectOps):
         bpy.ops.object.join()
 
         return (BPY_OBJ,)
-

blender/ops_set_shape_key_value.py

@@ -0,0 +1,26 @@
+import blender_node
+
+
+class Mesh_SetShapeKeyValue(blender_node.ObjectOps):
+
+    CUSTOM_NAME = "Set Shape Key Value"
+
+    EXTRA_INPUT_TYPES = {
+        "shape_key_name": ("STRING", {
+            "multiline": False,
+            "default": "my_shape_key",
+        }),
+        "value": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.01, "display": "number"}),
+    }
+
+    def blender_process(self, bpy, BPY_OBJ, shape_key_name, value):
+        # Check if the object has shape keys
+        if BPY_OBJ.data.shape_keys:
+            # Check if the specified shape key exists
+            if shape_key_name in BPY_OBJ.data.shape_keys.key_blocks:
+                BPY_OBJ.data.shape_keys.key_blocks[shape_key_name].value = float(value)
+            else:
+                print(f"The shape key {shape_key_name} does not exist on the object.")
+        else:
+            print("The object does not have any shape keys.")
+        return (BPY_OBJ,)

blender/render_image.py

@@ -0,0 +1,132 @@
+import blender_node
+import math
+import folder_paths
+import torch
+import numpy as np
+import os
+from PIL import Image, ImageOps
+
+def get_incremented_filename(folder_path, base_filename):
+    # Initialize the counter and create the full initial path
+    counter = 0
+    output_path = f"{folder_path}/{base_filename}.png"
+
+    # Check if the file exists and increment the counter until the file does not exist
+    while os.path.exists(output_path):
+        counter += 1
+        output_path = f"{folder_path}/{base_filename}_{counter}.png"
+
+    return output_path
+
+class BlenderRenderImage(blender_node.ObjectOps):
+    def __init__(self):
+        pass
+
+    EXTRA_INPUT_TYPES = {
+    }
+
+    # OUTPUT_NODE = True
+    RETURN_TYPES = ("BPY_OBJ", "IMAGE")
+
+    def add_light(self, bpy):
+        # Check if there is at least one light source in the scene
+        light_exists = any(ob for ob in bpy.data.objects if ob.type == 'LIGHT')
+        if not light_exists:
+            # Create a new Area light datablock for ambient light
+            light_data = bpy.data.lights.new(name='AmbientLight', type='AREA')
+            light_object = bpy.data.objects.new(name='AmbientLight', object_data=light_data)
+            bpy.context.collection.objects.link(light_object)
+            # Position the light in the scene
+            light_object.location = (0, 0, 10)
+            # Set light size for soft shadows and ambient effect
+            light_data.size = 10
+            light_data.energy = 1000
+            print("Added an ambient light source to the scene.")
+    
+    def add_camera(self, bpy):
+        # Check if there is a camera in the scene
+        if bpy.context.scene.camera:
+            return bpy.context.scene.camera 
+
+        # If not, create a new camera
+        cam_data = bpy.data.cameras.new(name='Camera')
+        cam = bpy.data.objects.new(name='Camera', object_data=cam_data)
+        bpy.context.collection.objects.link(cam)
+        # Set the new camera to the active camera
+        bpy.context.scene.camera = cam
+        # Position the camera to a default view
+        cam.location = (0, 0, 10)
+        return cam
+    
+    def get_texture_size(self, obj):
+        # Get the first material slot
+        mat = obj.data.materials[0]
+        
+        # Check if the material has a node tree
+        if mat.node_tree:
+            nodes = mat.node_tree.nodes
+            # Find an image texture node in the node tree
+            for node in nodes:
+                if node.type == 'TEX_IMAGE':
+                    texture = node.image
+                    if texture:
+                        return texture.size
+            print("No image texture node found in the material's node tree.")
+        else:
+            print("Material has no node tree.")
+
+
+    def blender_process(self, bpy, BPY_OBJ=None):
+        cam = self.add_camera(bpy)
+
+        plane = BPY_OBJ
+        if plane:
+            tex_width, tex_height = self.get_texture_size(plane)
+            # Calculate the aspect ratio of the plane
+            aspect_ratio_plane = tex_width / tex_height
+
+            # Set the render resolution to match the plane's aspect ratio
+            # Choose an arbitrary resolution for the longer side of the plane
+            base_resolution = 512
+
+            if aspect_ratio_plane > 1:
+                # Plane is wider than it is tall
+                bpy.context.scene.render.resolution_x = base_resolution
+                bpy.context.scene.render.resolution_y = int(base_resolution / aspect_ratio_plane)
+            else:
+                # Plane is taller than it is wide
+                bpy.context.scene.render.resolution_x = int(base_resolution * aspect_ratio_plane)
+                bpy.context.scene.render.resolution_y = base_resolution
+            bpy.context.scene.render.resolution_percentage = 100
+
+
+        ortho_scale = max(plane.dimensions.x, plane.dimensions.y)
+        cam.data.type = 'ORTHO'
+        cam.data.ortho_scale = ortho_scale
+
+        self.add_light(bpy)
+
+        # Update the scene to reflect changes
+        bpy.context.view_layer.update()
+
+        # Set render engine (e.g., 'BLENDER_EEVEE', 'CYCLES', 'BLENDER_WORKBENCH')
+        bpy.context.scene.render.engine = "BLENDER_EEVEE"
+
+        # Specify the render output path
+        
+        output_path = get_incremented_filename(folder_paths.get_output_directory(), "render")
+        bpy.context.scene.render.filepath = output_path
+
+        # Render the image
+        bpy.ops.render.render(write_still=True)
+
+        # Load the image
+        i = Image.open(output_path)
+        i = ImageOps.exif_transpose(i)
+        image = i.convert("RGB")
+        image = np.array(image).astype(np.float32) / 255.0
+        image = torch.from_numpy(image)[None,]
+
+        # print(image.shape)
+
+        return (BPY_OBJ, image)

js/AvatarPreview.js

@@ -173,10 +173,6 @@ async function prepareImageFromUrlRedirect(stage) {
 
 export function AvatarPreview() {
   console.log("getting workflow json now");
-  loadJSONWorkflow("default").then(() => {
-    console.log("done loading");
-    jsonWorkflowLoading.val = false;
-  });
 
   const loading = van.state(false);
   const shareLoading = van.state("share"); // share, loading, shared
@@ -465,7 +461,7 @@ export function AvatarPreview() {
       name: "avatech-viewer-iframe",
       allow: "cross-origin-isolated",
       class: () =>
-        "w-full h-full min-w-[400px] min-h-[400px] z-[100] pointer-events-auto flex border-none overflow-hidden" +
+        "w-full h-full min-w-[350px] min-h-[350px] z-[100] pointer-events-auto flex border-none overflow-hidden bg-transparent" +
         (showPreview.val ? "" : "hidden"),
       // src: "https://labs.avatech.ai/viewer/default",
       // src: "http://localhost:3000/viewer/default",
@@ -710,51 +706,12 @@ export function AvatarPreview() {
         console.log(showPreview);
 
         return (
-          (showPreview.val ? "" : "hidden ") +
-          "w-full h-full absolute left-0 top-0 z-[99] pointer-events-auto flex border-none bg-white"
+          (showPreview.val && !showEditor.val ? "" : "hidden ") +
+          "absolute w-[360px] h-[360px] rounded-xl overflow-hidden right-0 top-0 z-[99] pointer-events-auto flex border-none bg-transparent"
         );
       },
     },
-    div(
-      { class: "overflow-y-auto overflow-x-hidden w-full h-full" },
-      div(
-        {
-          class: "absolute top-4 left-4 flex flex-row gap-2",
-        },
-        renderCloseButton(),
-        renderRestartButton(),
-        renderChangeWorkflowButton()
-      ),
-      renderTwitter(),
-      () => {
-        if (isMobileDevice()) {
-          return div(
-            {
-              class: () =>
-                "flex flex-col w-full h-fit bg-white justify-center items-center py-16 px-4 gap-2" +
-                (showPreview.val ? "" : "hidden"),
-            },
-            renderIFrame(),
-            renderSteps(),
-            renderShareLink()
-          );
-        } else {
-          return div(
-            {
-              class: () =>
-                "flex w-full h-full bg-white justify-around items-center p-24" +
-                (showPreview.val ? "" : "hidden"),
-            },
-            renderSteps(),
-            div(
-              { class: () => "flex flex-col" },
-              renderIFrame(),
-              renderShareLink()
-            )
-          );
-        }
-      }
-    )
+    renderIFrame(),
   );
 }
 

js/CombinePointsDialog.js

@@ -7,12 +7,21 @@ van.derive(() => {
     combinePointsNode.val != undefined &&
     combinePointsNode.val.type === "Combine Points"
   ) {
-    const inputNames =
-      combinePointsNode.val.inputs?.map((x) => x.name).slice(1) || [];
+    const inputNames = combinePointsNode.val.inputs?.map((x) => x.name) || [];
     const record = Object.keys(samPrompts.val);
 
+    const diff = inputNames.filter((x) => !record.includes(x));
     const missingDiff = record.filter((x) => !inputNames.includes(x));
 
+    if (diff.length > 0) {
+      diff.forEach((x) => {
+        combinePointsNode.val.removeInput(
+          combinePointsNode.val.findInputSlot(x)
+        );
+      });
+      combinePointsNode.val.graph.change();
+    }
+
     if (missingDiff.length > 0) {
       missingDiff.forEach((x) => {
         combinePointsNode.val.addInput(x, "POINTS");

js/SideBar.js

@@ -6,6 +6,7 @@ import {
   targetNode,
   showImageEditor,
   allImagePrompts,
+  samPrompts,
 } from "./state.js";
 import { van } from "./van.js";
 const {
@@ -24,6 +25,33 @@ const {
   span,
 } = van.tags;
 
+export const updateOutputs = () => {
+  const outputNames = targetNode.val.outputs.map((x) => x.name).slice(1);
+  const record = Object.keys(imagePromptsMulti.val);
+
+  const diff = outputNames.filter((x) => !record.includes(x));
+  const missingDiff = record.filter((x) => !outputNames.includes(x));
+
+  if (diff.length > 0) {
+    console.log("Cleaning up missing output slots", diff);
+    diff.forEach((x) => {
+      targetNode.val.removeOutput(targetNode.val.findOutputSlot(x));
+    });
+    targetNode.val.graph.change();
+  }
+
+  if (missingDiff.length > 0) {
+    console.log("Adding missing output slots", diff);
+    missingDiff.forEach((x) => {
+      targetNode.val.addOutput(
+        x,
+        targetNode.val.type === "SAM MultiLayer" ? "IMAGE" : "SAM_PROMPT"
+      );
+    });
+    targetNode.val.graph.change();
+  }
+};
+
 van.derive(() => {
   if (
     showImageEditor.val &&
@@ -31,30 +59,7 @@ van.derive(() => {
     targetNode.val.outputs &&
     targetNode.val.type === "SAM MultiLayer"
   ) {
-    const outputNames = targetNode.val.outputs.map((x) => x.name).slice(1);
-    const record = Object.keys(imagePromptsMulti.val);
-
-    const diff = outputNames.filter((x) => !record.includes(x));
-    const missingDiff = record.filter((x) => !outputNames.includes(x));
-
-    if (diff.length > 0) {
-      console.log("Cleaning up missing output slots", diff);
-      diff.forEach((x) => {
-        targetNode.val.removeOutput(targetNode.val.findOutputSlot(x));
-      });
-      targetNode.val.graph.change();
-    }
-
-    if (missingDiff.length > 0) {
-      console.log("Adding missing output slots", diff);
-      missingDiff.forEach((x) => {
-        targetNode.val.addOutput(
-          x,
-          targetNode.val.type === "SAM MultiLayer" ? "IMAGE" : "SAM_PROMPT"
-        );
-      });
-      targetNode.val.graph.change();
-    }
+    updateOutputs();
   }
 });
 

js/index.js

@@ -18,6 +18,7 @@ import {
   previewModelId,
   embeddingID,
   enableAutoSegment,
+  samPrompts,
 } from "./state.js";
 import { van } from "./van.js";
 import { app } from "./app.js";
@@ -34,6 +35,7 @@ import {
 import { infoDialog } from "./dialog.js";
 import { sharedAvatarLink } from "./AvatarPreview.js";
 import { updateImagePrompts } from "./LayerEditor.js";
+import { updateOutputs } from "./SideBar.js";
 
 export const generatedImages = {};
 
@@ -319,7 +321,6 @@ function showMyImageEditor(node) {
           isGeneratedImage,
           embedding_id: id,
           ckpt,
-          // remote: true,
         }),
       })
       .then(() => {
@@ -358,7 +359,6 @@ function showMyImageEditor(node) {
           drawSegment(getClicks());
           updateImagePrompts();
         });
-        targetNode.val = node;
       })
       .catch((err) => {
         console.log(err);
@@ -386,6 +386,15 @@ const ext = {
           showMyImageEditor(node);
           btn.serialize = false;
         });
+
+        targetNode.val = node;
+        node.onConnectInput = (node, slot, targetSlot) => {
+          if (targetSlot.name === "SAM_PROMPTS") {
+            imagePromptsMulti.val = samPrompts.val;
+            updateOutputs();
+          }
+        };
+
         return {
           widget: btn,
         };

js/onnx.js

@@ -1,16 +1,18 @@
-import "https://cdn.jsdelivr.net/npm/onnxruntime-web/dist/ort.min.js";
+import "https://cdn.jsdelivr.net/npm/onnxruntime-web@1.16.3/dist/ort.min.js";
 import npyjs from "https://esm.sh/npyjs";
 import { imageSize } from "./state.js";
 import { modelData, onnxMaskToImage } from "./onnx_helper.js";
 
-ort.env.wasm.wasmPaths = "https://cdn.jsdelivr.net/npm/onnxruntime-web/dist/";
+ort.env.wasm.wasmPaths = "https://cdn.jsdelivr.net/npm/onnxruntime-web@1.16.3/dist/";
 
 export let model = null;
+let modelType = null;
 
 // Initialize the ONNX model
-export const initModel = async (modelType) => {
+export const initModel = async (type) => {
   try {
-    if (!model) {
+    if (!model || modelType !== type) {
+      modelType = type;
       model = await ort.InferenceSession.create(
         `${location.protocol}//${location.host}/sam_model?type=${modelType}`
       );

js/state.js

@@ -19,13 +19,15 @@
  */
 
 import { van } from "./van.js";
-
 export const iframeSrc = van.state("https://editor.avatech.ai?comfyui=true");
 export const showEditor = van.state(false);
 // localStorage.getItem("showPreview") == 'true'
-export const showPreview = van.state(true);
+console.log(localStorage.getItem("showPreview"));
+if (localStorage.getItem("showPreview") == null)
+  localStorage.setItem("showPreview", 'true')
+export const showPreview = van.state(localStorage.getItem("showPreview") == 'true');
 export const previewUrl = van.state(
-  "https://editor.avatech.ai/viewer?avatarId=default&debug=false&width=400&height=400&hideTrigger=true&voiceSelection=true&hideUI=true"
+  "https://editor.avatech.ai/viewer?avatarId=default&debug=true&width=350&height=350&hideTrigger=true&voiceSelection=true&hideUI=true"
 );
 export const previewImg = van.state("");
 export const previewImgLoading = van.state(false);

js/tw-styles.css

@@ -1,7 +1,7 @@
 @import url('https://fonts.googleapis.com/css2?family=Gabarito&display=swap');
 
 /*
-! tailwindcss v3.3.3 | MIT License | https://tailwindcss.com
+! tailwindcss v3.3.5 | MIT License | https://tailwindcss.com
 */
 
 /*
@@ -501,7 +501,7 @@ html{
 :root{
   color-scheme: light;
   --in: 0.7206 0.191 231.6;
-  --su: 0.7441 0.213 164.75;
+  --su: 64.8% 0.150 160;
   --wa: 0.8471 0.199 83.87;
   --er: 0.7176 0.221 22.18;
   --pc: 0.89824 0.06192 275.75;
@@ -742,7 +742,7 @@ html{
   justify-content: center;
   border-radius: var(--rounded-btn, 0.5rem);
   border-color: transparent;
-  border-color: oklch(var(--btn-color, var(--b2)) / var(--tw-border-opacity, 1));
+  border-color: oklch(var(--btn-color, var(--b2)) / var(--tw-border-opacity));
   padding-left: 1rem;
   padding-right: 1rem;
   text-align: center;
@@ -762,7 +762,9 @@ html{
   --tw-shadow-colored: 0 1px 2px 0 var(--tw-shadow-color);
   box-shadow: var(--tw-ring-offset-shadow, 0 0 #0000), var(--tw-ring-shadow, 0 0 #0000), var(--tw-shadow);
   outline-color: var(--fallback-bc,oklch(var(--bc)/1));
-  background-color: oklch(var(--btn-color, var(--b2)) / var(--tw-bg-opacity, 1));
+  background-color: oklch(var(--btn-color, var(--b2)) / var(--tw-bg-opacity));
+  --tw-bg-opacity: 1;
+  --tw-border-opacity: 1;
 }
 
 .btn-disabled,
@@ -870,8 +872,7 @@ html{
   flex-grow: 1;
   --tw-content: '';
   content: var(--tw-content);
-  --tw-bg-opacity: 1;
-  background-color: var(--fallback-b3,oklch(var(--b3)/var(--tw-bg-opacity)));
+  background-color: var(--fallback-bc,oklch(var(--bc)/0.1));
 }
 
 .dropdown{
@@ -959,18 +960,23 @@ html{
   .btn:hover{
     --tw-border-opacity: 1;
     border-color: var(--fallback-b3,oklch(var(--b3)/var(--tw-border-opacity)));
-    border-color: color-mix(
-          in oklab,
-          oklch(var(--btn-color, var(--b2)) / var(--tw-border-opacity, 1)) 90%,
-          black
-        );
     --tw-bg-opacity: 1;
     background-color: var(--fallback-b3,oklch(var(--b3)/var(--tw-bg-opacity)));
-    background-color: color-mix(
-          in oklab,
-          oklch(var(--btn-color, var(--b2)) / var(--tw-bg-opacity, 1)) 90%,
-          black
-        );
+  }
+
+  @supports (color: color-mix(in oklab, black, black)){
+    .btn:hover{
+      background-color: color-mix(
+            in oklab,
+            oklch(var(--btn-color, var(--b2)) / var(--tw-bg-opacity, 1)) 90%,
+            black
+          );
+      border-color: color-mix(
+            in oklab,
+            oklch(var(--btn-color, var(--b2)) / var(--tw-border-opacity, 1)) 90%,
+            black
+          );
+    }
   }
 
   @supports not (color: oklch(0 0 0)){
@@ -1007,50 +1013,85 @@ html{
   .btn-outline.btn-primary:hover{
     --tw-text-opacity: 1;
     color: var(--fallback-pc,oklch(var(--pc)/var(--tw-text-opacity)));
-    background-color: color-mix(in oklab, var(--fallback-p,oklch(var(--p)/1)) 90%, black);
-    border-color: color-mix(in oklab, var(--fallback-p,oklch(var(--p)/1)) 90%, black);
+  }
+
+  @supports (color: color-mix(in oklab, black, black)){
+    .btn-outline.btn-primary:hover{
+      background-color: color-mix(in oklab, var(--fallback-p,oklch(var(--p)/1)) 90%, black);
+      border-color: color-mix(in oklab, var(--fallback-p,oklch(var(--p)/1)) 90%, black);
+    }
   }
 
   .btn-outline.btn-secondary:hover{
     --tw-text-opacity: 1;
     color: var(--fallback-sc,oklch(var(--sc)/var(--tw-text-opacity)));
-    background-color: color-mix(in oklab, var(--fallback-s,oklch(var(--s)/1)) 90%, black);
-    border-color: color-mix(in oklab, var(--fallback-s,oklch(var(--s)/1)) 90%, black);
+  }
+
+  @supports (color: color-mix(in oklab, black, black)){
+    .btn-outline.btn-secondary:hover{
+      background-color: color-mix(in oklab, var(--fallback-s,oklch(var(--s)/1)) 90%, black);
+      border-color: color-mix(in oklab, var(--fallback-s,oklch(var(--s)/1)) 90%, black);
+    }
   }
 
   .btn-outline.btn-accent:hover{
     --tw-text-opacity: 1;
     color: var(--fallback-ac,oklch(var(--ac)/var(--tw-text-opacity)));
-    background-color: color-mix(in oklab, var(--fallback-a,oklch(var(--a)/1)) 90%, black);
-    border-color: color-mix(in oklab, var(--fallback-a,oklch(var(--a)/1)) 90%, black);
+  }
+
+  @supports (color: color-mix(in oklab, black, black)){
+    .btn-outline.btn-accent:hover{
+      background-color: color-mix(in oklab, var(--fallback-a,oklch(var(--a)/1)) 90%, black);
+      border-color: color-mix(in oklab, var(--fallback-a,oklch(var(--a)/1)) 90%, black);
+    }
   }
 
   .btn-outline.btn-success:hover{
     --tw-text-opacity: 1;
     color: var(--fallback-suc,oklch(var(--suc)/var(--tw-text-opacity)));
-    background-color: color-mix(in oklab, var(--fallback-su,oklch(var(--su)/1)) 90%, black);
-    border-color: color-mix(in oklab, var(--fallback-su,oklch(var(--su)/1)) 90%, black);
+  }
+
+  @supports (color: color-mix(in oklab, black, black)){
+    .btn-outline.btn-success:hover{
+      background-color: color-mix(in oklab, var(--fallback-su,oklch(var(--su)/1)) 90%, black);
+      border-color: color-mix(in oklab, var(--fallback-su,oklch(var(--su)/1)) 90%, black);
+    }
   }
 
   .btn-outline.btn-info:hover{
     --tw-text-opacity: 1;
     color: var(--fallback-inc,oklch(var(--inc)/var(--tw-text-opacity)));
-    background-color: color-mix(in oklab, var(--fallback-in,oklch(var(--in)/1)) 90%, black);
-    border-color: color-mix(in oklab, var(--fallback-in,oklch(var(--in)/1)) 90%, black);
+  }
+
+  @supports (color: color-mix(in oklab, black, black)){
+    .btn-outline.btn-info:hover{
+      background-color: color-mix(in oklab, var(--fallback-in,oklch(var(--in)/1)) 90%, black);
+      border-color: color-mix(in oklab, var(--fallback-in,oklch(var(--in)/1)) 90%, black);
+    }
   }
 
   .btn-outline.btn-warning:hover{
     --tw-text-opacity: 1;
     color: var(--fallback-wac,oklch(var(--wac)/var(--tw-text-opacity)));
-    background-color: color-mix(in oklab, var(--fallback-wa,oklch(var(--wa)/1)) 90%, black);
-    border-color: color-mix(in oklab, var(--fallback-wa,oklch(var(--wa)/1)) 90%, black);
+  }
+
+  @supports (color: color-mix(in oklab, black, black)){
+    .btn-outline.btn-warning:hover{
+      background-color: color-mix(in oklab, var(--fallback-wa,oklch(var(--wa)/1)) 90%, black);
+      border-color: color-mix(in oklab, var(--fallback-wa,oklch(var(--wa)/1)) 90%, black);
+    }
   }
 
   .btn-outline.btn-error:hover{
     --tw-text-opacity: 1;
     color: var(--fallback-erc,oklch(var(--erc)/var(--tw-text-opacity)));
-    background-color: color-mix(in oklab, var(--fallback-er,oklch(var(--er)/1)) 90%, black);
-    border-color: color-mix(in oklab, var(--fallback-er,oklch(var(--er)/1)) 90%, black);
+  }
+
+  @supports (color: color-mix(in oklab, black, black)){
+    .btn-outline.btn-error:hover{
+      background-color: color-mix(in oklab, var(--fallback-er,oklch(var(--er)/1)) 90%, black);
+      border-color: color-mix(in oklab, var(--fallback-er,oklch(var(--er)/1)) 90%, black);
+    }
   }
 
   .btn-disabled:hover,
@@ -1063,9 +1104,11 @@ html{
     --tw-text-opacity: 0.2;
   }
 
-  .btn:is(input[type="checkbox"]:checked):hover, .btn:is(input[type="radio"]:checked):hover{
-    background-color: color-mix(in oklab, var(--fallback-p,oklch(var(--p)/1)) 90%, black);
-    border-color: color-mix(in oklab, var(--fallback-p,oklch(var(--p)/1)) 90%, black);
+  @supports (color: color-mix(in oklab, black, black)){
+    .btn:is(input[type="checkbox"]:checked):hover, .btn:is(input[type="radio"]:checked):hover{
+      background-color: color-mix(in oklab, var(--fallback-p,oklch(var(--p)/1)) 90%, black);
+      border-color: color-mix(in oklab, var(--fallback-p,oklch(var(--p)/1)) 90%, black);
+    }
   }
 
   .dropdown.dropdown-hover:hover .dropdown-content{
@@ -1076,8 +1119,6 @@ html{
 
   :where(.menu li:not(.menu-title):not(.disabled) > *:not(ul):not(details):not(.menu-title)):not(.active):hover, :where(.menu li:not(.menu-title):not(.disabled) > details > summary:not(.menu-title)):not(.active):hover{
     cursor: pointer;
-    --tw-text-opacity: 1;
-    color: var(--fallback-bc,oklch(var(--bc)/var(--tw-text-opacity)));
     outline: 2px solid transparent;
     outline-offset: 2px;
   }
@@ -1108,6 +1149,9 @@ html{
 
 .input{
   flex-shrink: 1;
+  -webkit-appearance: none;
+     -moz-appearance: none;
+          appearance: none;
   height: 3rem;
   padding-left: 1rem;
   padding-right: 1rem;
@@ -1210,10 +1254,8 @@ html{
 .menu :where(li ul){
   position: relative;
   white-space: nowrap;
-  -webkit-margin-start: 1rem;
-          margin-inline-start: 1rem;
-  -webkit-padding-start: 0.5rem;
-          padding-inline-start: 0.5rem;
+  margin-inline-start: 1rem;
+  padding-inline-start: 0.5rem;
 }
 
 .menu :where(li:not(.menu-title) > *:not(ul):not(details):not(.menu-title)),
@@ -1394,6 +1436,43 @@ html{
   }
 }
 
+@supports (color: color-mix(in oklab, black, black)){
+  .btn-outline.btn-primary.btn-active{
+    background-color: color-mix(in oklab, var(--fallback-p,oklch(var(--p)/1)) 90%, black);
+    border-color: color-mix(in oklab, var(--fallback-p,oklch(var(--p)/1)) 90%, black);
+  }
+
+  .btn-outline.btn-secondary.btn-active{
+    background-color: color-mix(in oklab, var(--fallback-s,oklch(var(--s)/1)) 90%, black);
+    border-color: color-mix(in oklab, var(--fallback-s,oklch(var(--s)/1)) 90%, black);
+  }
+
+  .btn-outline.btn-accent.btn-active{
+    background-color: color-mix(in oklab, var(--fallback-a,oklch(var(--a)/1)) 90%, black);
+    border-color: color-mix(in oklab, var(--fallback-a,oklch(var(--a)/1)) 90%, black);
+  }
+
+  .btn-outline.btn-success.btn-active{
+    background-color: color-mix(in oklab, var(--fallback-su,oklch(var(--su)/1)) 90%, black);
+    border-color: color-mix(in oklab, var(--fallback-su,oklch(var(--su)/1)) 90%, black);
+  }
+
+  .btn-outline.btn-info.btn-active{
+    background-color: color-mix(in oklab, var(--fallback-in,oklch(var(--in)/1)) 90%, black);
+    border-color: color-mix(in oklab, var(--fallback-in,oklch(var(--in)/1)) 90%, black);
+  }
+
+  .btn-outline.btn-warning.btn-active{
+    background-color: color-mix(in oklab, var(--fallback-wa,oklch(var(--wa)/1)) 90%, black);
+    border-color: color-mix(in oklab, var(--fallback-wa,oklch(var(--wa)/1)) 90%, black);
+  }
+
+  .btn-outline.btn-error.btn-active{
+    background-color: color-mix(in oklab, var(--fallback-er,oklch(var(--er)/1)) 90%, black);
+    border-color: color-mix(in oklab, var(--fallback-er,oklch(var(--er)/1)) 90%, black);
+  }
+}
+
 .btn:focus-visible{
   outline-style: solid;
   outline-width: 2px;
@@ -1467,8 +1546,6 @@ html{
 .btn-outline.btn-primary.btn-active{
   --tw-text-opacity: 1;
   color: var(--fallback-pc,oklch(var(--pc)/var(--tw-text-opacity)));
-  background-color: color-mix(in oklab, var(--fallback-p,oklch(var(--p)/1)) 90%, black);
-  border-color: color-mix(in oklab, var(--fallback-p,oklch(var(--p)/1)) 90%, black);
 }
 
 .btn-outline.btn-secondary{
@@ -1479,8 +1556,6 @@ html{
 .btn-outline.btn-secondary.btn-active{
   --tw-text-opacity: 1;
   color: var(--fallback-sc,oklch(var(--sc)/var(--tw-text-opacity)));
-  background-color: color-mix(in oklab, var(--fallback-s,oklch(var(--s)/1)) 90%, black);
-  border-color: color-mix(in oklab, var(--fallback-s,oklch(var(--s)/1)) 90%, black);
 }
 
 .btn-outline.btn-accent{
@@ -1491,8 +1566,6 @@ html{
 .btn-outline.btn-accent.btn-active{
   --tw-text-opacity: 1;
   color: var(--fallback-ac,oklch(var(--ac)/var(--tw-text-opacity)));
-  background-color: color-mix(in oklab, var(--fallback-a,oklch(var(--a)/1)) 90%, black);
-  border-color: color-mix(in oklab, var(--fallback-a,oklch(var(--a)/1)) 90%, black);
 }
 
 .btn-outline.btn-success{
@@ -1503,8 +1576,6 @@ html{
 .btn-outline.btn-success.btn-active{
   --tw-text-opacity: 1;
   color: var(--fallback-suc,oklch(var(--suc)/var(--tw-text-opacity)));
-  background-color: color-mix(in oklab, var(--fallback-su,oklch(var(--su)/1)) 90%, black);
-  border-color: color-mix(in oklab, var(--fallback-su,oklch(var(--su)/1)) 90%, black);
 }
 
 .btn-outline.btn-info{
@@ -1515,8 +1586,6 @@ html{
 .btn-outline.btn-info.btn-active{
   --tw-text-opacity: 1;
   color: var(--fallback-inc,oklch(var(--inc)/var(--tw-text-opacity)));
-  background-color: color-mix(in oklab, var(--fallback-in,oklch(var(--in)/1)) 90%, black);
-  border-color: color-mix(in oklab, var(--fallback-in,oklch(var(--in)/1)) 90%, black);
 }
 
 .btn-outline.btn-warning{
@@ -1527,8 +1596,6 @@ html{
 .btn-outline.btn-warning.btn-active{
   --tw-text-opacity: 1;
   color: var(--fallback-wac,oklch(var(--wac)/var(--tw-text-opacity)));
-  background-color: color-mix(in oklab, var(--fallback-wa,oklch(var(--wa)/1)) 90%, black);
-  border-color: color-mix(in oklab, var(--fallback-wa,oklch(var(--wa)/1)) 90%, black);
 }
 
 .btn-outline.btn-error{
@@ -1539,8 +1606,6 @@ html{
 .btn-outline.btn-error.btn-active{
   --tw-text-opacity: 1;
   color: var(--fallback-erc,oklch(var(--erc)/var(--tw-text-opacity)));
-  background-color: color-mix(in oklab, var(--fallback-er,oklch(var(--er)/1)) 90%, black);
-  border-color: color-mix(in oklab, var(--fallback-er,oklch(var(--er)/1)) 90%, black);
 }
 
 .btn.btn-disabled,
@@ -1646,8 +1711,7 @@ details.collapse summary::-webkit-details-marker{
 :where(.collapse > input[type="radio"]){
   width: 100%;
   padding: 1rem;
-  -webkit-padding-end: 3rem;
-          padding-inline-end: 3rem;
+  padding-inline-end: 3rem;
   min-height: 3.75rem;
   transition: background-color 0.2s ease-out;
 }
@@ -1750,11 +1814,14 @@ details.collapse summary::-webkit-details-marker{
   --tw-placeholder-opacity: 0.2;
 }
 
+.input::-webkit-date-and-time-value{
+  text-align: inherit;
+}
+
 .join > :where(*:not(:first-child)){
   margin-top: 0px;
   margin-bottom: 0px;
-  -webkit-margin-start: -1px;
-          margin-inline-start: -1px;
+  margin-inline-start: -1px;
 }
 
 .link:focus{
@@ -1797,7 +1864,9 @@ details.collapse summary::-webkit-details-marker{
 }
 
 :where(.menu li:empty){
-  background-color: var(--fallback-bc,oklch(var(--bc)/0.1));
+  --tw-bg-opacity: 1;
+  background-color: var(--fallback-bc,oklch(var(--bc)/var(--tw-bg-opacity)));
+  opacity: 0.1;
   margin: 0.5rem 1rem;
   height: 1px;
 }
@@ -1808,7 +1877,9 @@ details.collapse summary::-webkit-details-marker{
   inset-inline-start: 0px;
   top: 0.75rem;
   width: 1px;
-  background-color: var(--fallback-bc,oklch(var(--bc)/0.1));
+  --tw-bg-opacity: 1;
+  background-color: var(--fallback-bc,oklch(var(--bc)/var(--tw-bg-opacity)));
+  opacity: 0.1;
   content: "";
 }
 
@@ -1933,7 +2004,7 @@ details.collapse summary::-webkit-details-marker{
 
 .modal:not(dialog:not(.modal-open)),
   .modal::backdrop{
-  background-color: rgba(0, 0, 0, 0.3);
+  background-color: #0006;
   animation: modal-pop 0.2s ease-out;
 }
 
@@ -2256,8 +2327,7 @@ details.collapse summary::-webkit-details-marker{
 .join.join-horizontal > :where(*:not(:first-child)){
   margin-top: 0px;
   margin-bottom: 0px;
-  -webkit-margin-start: -1px;
-          margin-inline-start: -1px;
+  margin-inline-start: -1px;
 }
 
 .modal-top :where(.modal-box){
@@ -2399,6 +2469,10 @@ details.collapse summary::-webkit-details-marker{
   z-index: 200;
 }
 
+.z-\[999\]{
+  z-index: 999;
+}
+
 .z-\[99\]{
   z-index: 99;
 }
@@ -2465,21 +2539,20 @@ details.collapse summary::-webkit-details-marker{
   height: 24rem;
 }
 
-.h-\[394px\]{
-  height: 394px;
+.h-\[360px\]{
+  height: 360px;
 }
 
-.h-fit{
-  height: -moz-fit-content;
-  height: fit-content;
+.h-\[394px\]{
+  height: 394px;
 }
 
 .h-full{
   height: 100%;
 }
 
-.min-h-\[400px\]{
-  min-height: 400px;
+.min-h-\[350px\]{
+  min-height: 350px;
 }
 
 .w-20{
@@ -2506,6 +2579,10 @@ details.collapse summary::-webkit-details-marker{
   width: 32rem;
 }
 
+.w-\[360px\]{
+  width: 360px;
+}
+
 .w-fit{
   width: -moz-fit-content;
   width: fit-content;
@@ -2515,8 +2592,8 @@ details.collapse summary::-webkit-details-marker{
   width: 100%;
 }
 
-.min-w-\[400px\]{
-  min-width: 400px;
+.min-w-\[350px\]{
+  min-width: 350px;
 }
 
 .max-w-\[100\%\]{
@@ -2569,10 +2646,6 @@ details.collapse summary::-webkit-details-marker{
   justify-content: space-between;
 }
 
-.justify-around{
-  justify-content: space-around;
-}
-
 .gap-2{
   gap: 0.5rem;
 }
@@ -2589,14 +2662,6 @@ details.collapse summary::-webkit-details-marker{
   overflow: hidden;
 }
 
-.overflow-y-auto{
-  overflow-y: auto;
-}
-
-.overflow-x-hidden{
-  overflow-x: hidden;
-}
-
 .rounded{
   border-radius: 0.25rem;
 }
@@ -2621,6 +2686,10 @@ details.collapse summary::-webkit-details-marker{
   border-radius: 0.125rem;
 }
 
+.rounded-xl{
+  border-radius: 0.75rem;
+}
+
 .rounded-b-md{
   border-bottom-right-radius: 0.375rem;
   border-bottom-left-radius: 0.375rem;
@@ -2834,10 +2903,6 @@ details.collapse summary::-webkit-details-marker{
   padding: 0.5rem;
 }
 
-.p-24{
-  padding: 6rem;
-}
-
 .p-4{
   padding: 1rem;
 }
@@ -2857,11 +2922,6 @@ details.collapse summary::-webkit-details-marker{
   padding-right: 1rem;
 }
 
-.py-16{
-  padding-top: 4rem;
-  padding-bottom: 4rem;
-}
-
 .py-2{
   padding-top: 0.5rem;
   padding-bottom: 0.5rem;

routes.py

@@ -1,18 +1,19 @@
 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
@@ -67,80 +68,19 @@ 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")
-    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):
+    model_type = sam_ckpt_to_type[ckpt]
+    if not check_embedding_exists(emb_id, model_type):
         image = load_image(post.get("image"), is_generated_image)
-        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).to("cuda")
-            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,
-                )
+        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)
     print("Finished embedding")
     return web.json_response({})
 
@@ -222,9 +162,6 @@ def load_workflow(workflow_name):
         return "\n".join(f.readlines())
 
 
-default_workflow = load_workflow("avatar_generation_mask_api_v11(FaceToon)")
-
-
 @server.PromptServer.instance.routes.post("/avatar_generation")
 async def post_prompt_block(request):
     prompt_server = server.PromptServer.instance
@@ -235,8 +172,6 @@ async def post_prompt_block(request):
         workflow = uploaded_workflow
     elif workflow_name is not None:
         workflow = load_workflow(workflow_name)
-    else:
-        workflow = default_workflow
 
     ref_image = post.get("ref_image")
     base_image = post.get("base_image")
@@ -290,6 +225,54 @@ async def post_prompt_block(request):
         time.sleep(0.5)
 
 
+# TODO: refactor the code
+@server.PromptServer.instance.routes.post("/rendering_generation")
+async def post_data_generation(request):
+    prompt_server = server.PromptServer.instance
+    post = await request.json()
+    workflow_name = post.get("workflow_name")
+    workflow = load_workflow(workflow_name)
+    workflow = workflow.replace("SEED", str(randomSeed()))
+
+    inputs = post.get("inputs")
+
+    for key, value in inputs.items():
+        workflow = workflow.replace(f'"{key}"', f'"{str(value)}"')
+
+    res = post_prompt({"prompt": json.loads(workflow)})
+    prompt_id = json.loads(res.text)["prompt_id"]
+    while True:
+        history = prompt_server.prompt_queue.get_history(prompt_id=prompt_id)
+        if history:
+            outputs = history[prompt_id]["outputs"]
+            for node_id, output in outputs.items():
+                if "images" in output:
+                    filename = output["images"][0]["filename"]
+                    if filename.startswith("rendered"):
+                        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
+    workflow = load_workflow("generation")
+    workflow = workflow.replace("SEED", str(randomSeed()))
+
+    res = post_prompt({"prompt": json.loads(workflow)})
+    prompt_id = json.loads(res.text)["prompt_id"]
+    while True:
+        history = prompt_server.prompt_queue.get_history(prompt_id=prompt_id)
+        if history:
+            outputs = history[prompt_id]["outputs"]
+            for node_id, output in outputs.items():
+                if "images" in output:
+                    filename = output["images"][0]["filename"]
+                    if filename.startswith("avatar"):
+                        return web.json_response({"image": filename}, status=200)
+        time.sleep(0.5)
+
+
 # @server.PromptServer.instance.routes.get("/get_default_workflow")
 # async def get_default_workflow(request):
 #     # json_link = "https://cdn.discordapp.com/attachments/1119102674437156984/1172255632586448987/workflow_boy_2_1.json?ex=655fa722&is=654d3222&hm=463fa6a3c6ea60f7471196ff45382c729d3b856e86282f905d37a0398711860e&" # YP workflow

sam/combine_points.py

@@ -1,3 +1,4 @@
+import json
 
 class CombinePoints:
     @classmethod
@@ -7,7 +8,8 @@ class CombinePoints:
             },
         }
 
-    RETURN_TYPES = ("SAM_PROMPTS",)
+    RETURN_NAMES = ("SAM_PROMPTS",)
+    RETURN_TYPES = ("STRING",)
 
     FUNCTION = "run"
 
@@ -16,9 +18,8 @@ class CombinePoints:
     # OUTPUT_NODE = True
 
     def run(self, *args, **kwargs):
-        print("args", args)
-        print("kwargs", kwargs)
-        return ([])
+        sam_prompts = json.dumps(kwargs, default=str)
+        return (sam_prompts,)
 
 
 NODE_CLASS_MAPPINGS = {"Combine Points": CombinePoints}

sam/load_value_from_request.py

@@ -0,0 +1,31 @@
+class LoadValueFromRequest:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "name": (
+                    "STRING",
+                    {"multiline": False, "default": "key_name"},
+                ),
+            },
+            "optional": {
+                "value": ("FLOAT", {"default": 1.0, "min": -10.0, "max": 10.0, "step": 0.01, "display": "number"}),
+            }
+        }
+
+    RETURN_TYPES = ("FLOAT",)
+    RETURN_NAMES = ("value",)
+
+    FUNCTION = "run"
+
+    CATEGORY = "image"
+
+    def run(self, name, value=None):
+        if name:
+            value = name
+        return (value,)
+
+
+NODE_CLASS_MAPPINGS = {"LoadValueFromRequest": LoadValueFromRequest}
+
+NODE_DISPLAY_NAME_MAPPINGS = {"LoadValueFromRequest": "Load Value From Request"}

sam/mediapipe_utils.py

@@ -0,0 +1,254 @@
+# 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

@@ -4,105 +4,20 @@ import numpy as np
 import torch
 import re
 import json
-from segment_anything import sam_model_registry, SamPredictor
+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 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 {
@@ -119,183 +34,31 @@ class SAMMultiLayer:
 
     CATEGORY = "image"
 
-    RETURN_TYPES = ("SAM_PROMPT",)
-    FUNCTION = "load_image"
+    RETURN_TYPES = ["SAM_PROMPT"] # + ["IMAGE"] * 100
+    FUNCTION = "run"
 
-    def load_models(self, ckpt, model_type):
-        global global_predictor, face_landmarker, pose_landmarker
+    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("'", '"'))
 
-        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"
+        order_file = (
+            f"{folder_paths.get_output_directory()}/segments_{embedding_id}/order.json"
         )
-        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
-        face_landmarks = face_landmarks[0] if len(face_landmarks) > 0 else None
-        pose_landmarks = pose_landmarker.detect(mp_image).pose_landmarks
-        pose_landmarks = pose_landmarks[0] if len(pose_landmarks) > 0 else None
-        imagePromptsMulti, boxesMulti = self.auto_segment(
-            np_image, face_landmarks, pose_landmarks
-        )
-
-        return imagePromptsMulti, boxesMulti
-
-    def load_image(self, image, ckpt, embedding_id, image_prompts_json):
-        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 = [image_prompts]
-
+            result = [layer_points]
             for segment in order:
-                image = Image.open(
-                    f"{self.output_dir}/segments_{embedding_id}/{segment}.png"
+                image = load_image(
+                    f"{folder_paths.get_output_directory()}/segments_{embedding_id}/{segment}.png",
+                    comfyui_format=True,
                 )
-                image = np.array(image).astype(np.float32) / 255.0
-                image = torch.from_numpy(image)[None,]
                 result.append(image)
 
             return result
@@ -303,68 +66,60 @@ 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]
 
-            global global_predictor
-            if global_predictor is None:
-                global_predictor, _, _ = self.load_models(ckpt, model_type)
+            # get first image from batch
+            image = image[0]
+            if image.shape[2] == 4:
+                # to RGB
+                image = image[:, :, :3]
 
-            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)
+            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,
+                )
             else:
-                emb = np.load(emb_filename)
+                load_embdding(embedding_id, model_type)
 
-                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"]
+            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]
 
-            imagePromptsMulti, boxesMulti = self.detect_face(image[0].numpy())
+                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
 
-            image_prompts = json.loads(image_prompts_json.replace("'", '"'))
-            result = [image_prompts]
+                # 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
+                )
 
-            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 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 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
+                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
 
 
 NODE_CLASS_MAPPINGS = {"SAM MultiLayer": SAMMultiLayer}

sam/sam_utils.py

@@ -0,0 +1,97 @@
+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