diff --git a/blender/avatar_main_output.py b/blender/avatar_main_output.py index 3ceab47..e9e3fc7 100644 --- a/blender/avatar_main_output.py +++ b/blender/avatar_main_output.py @@ -1,6 +1,6 @@ import platform import blender_node -from mesh_utils import assign_texture, open_in_blender as open_blender, export_gltf +from mesh_utils import upload_avatar_file, open_in_blender as open_blender, export_gltf import folder_paths global_blender_path = '' @@ -47,14 +47,16 @@ class AvatarMainOutput(blender_node.ObjectOps): }), "model_type": (["AVA","GLB", "GLTF_EMBEDDED"],), "write_mode": (["Overwrite", "Increment"],), - + "upload_to_cloud": ("BOOLEAN", { + "default": False + }), "SHAPE_FLOW": ("SHAPE_FLOW",), } OUTPUT_NODE = True RETURN_TYPES = () - def blender_process(self, bpy, BPY_OBJ=None, BPY_OBJS=None, open_in_blender=False, auto_save=False, blender_path_override='', filename='', model_type='', write_mode='', SHAPE_FLOW=''): + def blender_process(self, bpy, BPY_OBJ=None, BPY_OBJS=None, open_in_blender=False, auto_save=False, blender_path_override='', filename='', model_type='', write_mode='', SHAPE_FLOW='', upload_to_cloud=False): if open_in_blender: p = blender_path_override if blender_path_override else global_blender_path output_file = self.output_dir + '/tmp.blend' @@ -72,7 +74,7 @@ class AvatarMainOutput(blender_node.ObjectOps): import global_bpy global_bpy.set_should_reset_scene(True) - return { + outputs = { "ui": { "gltfFilename": {filepath.replace(f"{self.output_dir}/", "")}, "files": [{ @@ -83,3 +85,12 @@ class AvatarMainOutput(blender_node.ObjectOps): "auto_save": {'true' if auto_save else 'false'}, } } + if upload_to_cloud: + avatarId = upload_avatar_file(outputs) + return { + "ui": { + **outputs["ui"], + "avatarId": {avatarId}, + }, + } + return outputs diff --git a/blender/mesh_utils.py b/blender/mesh_utils.py index 3b755e3..6af949e 100644 --- a/blender/mesh_utils.py +++ b/blender/mesh_utils.py @@ -1,6 +1,9 @@ import atexit import subprocess import os +import folder_paths +import requests +import json def genreate_mesh_from_texture(bpy, image): import torch @@ -245,4 +248,41 @@ def export_gltf(output_dir, bpy_objects, filename, model_type, write_mode, metad os.replace(filepath, new_filepath) filepath = new_filepath - return filepath \ No newline at end of file + return filepath + +def get_avatar_file(outputs): + for node_id, output in outputs.items(): + if "gltfFilename" in output: + avatar_filename = output["gltfFilename"][0] + with open( + f"{folder_paths.get_output_directory()}/{avatar_filename}", "rb" + ) as f: + return f.read() + +def upload_avatar_file(outputs): + file = get_avatar_file(outputs) + response = requests.get("https://labs.avatech.ai/api/share") + labData = response.json() + modelId = labData["modelId"] + + # upload model + headers = { + "x-amz-acl": "public-read", + "Content-Type": "model/gltf-binary", + "Content-Length": str(len(file)), + } + 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)) + + return modelId diff --git a/routes.py b/routes.py index 7abf5be..6ce1372 100644 --- a/routes.py +++ b/routes.py @@ -2,6 +2,7 @@ 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 import os import requests import folder_paths @@ -205,45 +206,6 @@ def post_prompt(json_data): return web.json_response({"error": "no prompt", "node_errors": []}, status=400) -def get_avatar_file(outputs): - for node_id, output in outputs.items(): - if "gltfFilename" in output: - avatar_filename = output["gltfFilename"][0] - with open( - f"{folder_paths.get_output_directory()}/{avatar_filename}", "rb" - ) as f: - return f.read() - - -def upload_avatar_file(outputs): - file = get_avatar_file(outputs) - response = requests.get("https://labs.avatech.ai/api/share") - labData = response.json() - modelId = labData["modelId"] - - # upload model - headers = { - "x-amz-acl": "public-read", - "Content-Type": "model/gltf-binary", - "Content-Length": str(len(file)), - } - 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)) - - return modelId - - def randomSeed(num_digits=15): range_start = 10 ** (num_digits - 1) range_end = (10**num_digits) - 1