diff --git a/README.md b/README.md index a278f7a..5920182 100644 --- a/README.md +++ b/README.md @@ -17,16 +17,20 @@ def prompt_worker(q, server): # This is a workaround since using blender as a python module, the bpy module has to be imported after in the custom thread, otherwise it will cause a segfault if imported in the custom nodes. - # Add this line here - import sys - sys.path.append(os.path.join(os.path.dirname(__file__), "custom_nodes", "avatar-graph-comfyui")) + # Add next line import global_bpy while True: item, item_id = q.get() - # Add this line here + # Add next line global_bpy.reset_bpy() +if __name__ == "__main__": + # Add next two lines + import sys + sys.path.append(os.path.join(os.path.dirname(__file__), "custom_nodes", "avatar-graph-comfyui")) + + ``` 5. Restart comfyui diff --git a/__init__.py b/__init__.py index 51b16ca..e0dd9ae 100644 --- a/__init__.py +++ b/__init__.py @@ -4,6 +4,7 @@ @nickname: Avatar Graph @description: Include nodes for sam + bpy operation, that allows workflow creations for generative 2d character rig. """ +import routes import inspect import os diff --git a/routes.py b/routes.py new file mode 100644 index 0000000..fa85b75 --- /dev/null +++ b/routes.py @@ -0,0 +1,85 @@ +from aiohttp import web +from segment_anything import sam_model_registry, SamPredictor +from PIL import Image, ImageOps +import os +import requests +import folder_paths +import json +import numpy as np +import server + +# For speeding up ONNX model, see https://github.com/facebookresearch/segment-anything/tree/main/demo#onnx-multithreading-with-sharedarraybuffer +def inject_headers(original_handler): + async def _handler(request): + res = await original_handler(request) + res.headers["Cross-Origin-Opener-Policy"] = "same-origin" + res.headers["Cross-Origin-Embedder-Policy"] = "credentialless" + return res + + return _handler + + +routes = [] +for item in server.PromptServer.instance.routes._items: + if item.path == "/": + item = web.RouteDef( + method=item.method, + path=item.path, + handler=inject_headers(item.handler), + kwargs=item.kwargs, + ) + routes.append(item) +server.PromptServer.instance.routes._items = routes + + +@server.PromptServer.instance.routes.get("/sam_model") +async def get_sam_model(request): + filename = os.path.join(folder_paths.base_path, "web/models/sam.onnx") + # print(filename) + if not os.path.isfile(filename): + os.makedirs(os.path.dirname(filename), exist_ok=True) + response = requests.get( + "https://avatech-avatar-dev1.nyc3.cdn.digitaloceanspaces.com/models/sam.onnx" + ) + response.raise_for_status() + with open(filename, "wb") as f: + f.write(response.content) + print(f"ONNX model downloaded: {filename}") + return web.FileResponse(filename) + + +def load_image(image): + 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() + emb_id = post.get("embedding_id") + emb_filename = f"{folder_paths.get_output_directory()}/{emb_id}.npy" + if not os.path.exists(emb_filename): + image = load_image(post.get("image")) + ckpt = post.get("ckpt") + model_type = post.get("model_type") + ckpt = folder_paths.get_full_path("sam", ckpt) + 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(f"{folder_paths.get_output_directory()}/{emb_id}.json", "w") as f: + json.dump( + { + "input_size": predictor.input_size, + "original_size": predictor.original_size, + }, + f, + ) + return web.json_response({})