feat: upload avatar in output node

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

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
+    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

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