65 lines
1.8 KiB
Python
65 lines
1.8 KiB
Python
import folder_paths
|
|
import torch
|
|
import os
|
|
from segment_anything import SamAutomaticMaskGenerator, sam_model_registry
|
|
from einops import rearrange, repeat
|
|
|
|
class SAM_Save_Embedding:
|
|
def __init__(self):
|
|
self.output_dir = folder_paths.get_output_directory()
|
|
self.type = "output"
|
|
self.prefix_append = ""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"embeddings": ("EMBEDDINGS",),
|
|
"filename": ("STRING", {
|
|
"multiline": False,
|
|
"default": "embeddings"
|
|
}),
|
|
"write_mode": (["Overwrite", "Increment"],),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ()
|
|
RETURN_NAMES = ()
|
|
|
|
OUTPUT_NODE = True
|
|
|
|
FUNCTION = "process"
|
|
|
|
CATEGORY = "image"
|
|
|
|
def process(self, embeddings, filename, write_mode):
|
|
import json
|
|
import numpy as np
|
|
|
|
filepath = self.output_dir + "/" + filename + '.json'
|
|
|
|
if write_mode == "Increment":
|
|
count = 0
|
|
# while file exists, increment count
|
|
while os.path.exists(self.output_dir + "/" + filename + '_' + str(count) + '.json'):
|
|
count += 1
|
|
|
|
filepath = self.output_dir + "/" + filename + '_' + str(count) + '.json'
|
|
|
|
# print(embeddings)
|
|
if isinstance(embeddings['image_embedding'], np.ndarray):
|
|
embeddings['image_embedding'] = embeddings['image_embedding'].tolist()
|
|
with open(filepath, 'w') as f:
|
|
json.dump(embeddings, f)
|
|
|
|
return { "ui" : { "file": { filepath.replace(f"{self.output_dir}/", "") } } }
|
|
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"SAM_Save_Embedding": SAM_Save_Embedding
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"SAM_Save_Embedding": "SAM_Save_Embedding "
|
|
}
|