From c65b7300203be9360d09d5b00aceb6a3b49f1817 Mon Sep 17 00:00:00 2001
From: Radionic <khpluslt@gmail.com>
Date: Fri, 29 Dec 2023 11:52:48 +0800
Subject: [PATCH] feat: upload avatar avatar_main_output node

---
 blender/avatar_main_output.py | 26 +++++++++++++-------------
 blender/mesh_utils.py         | 18 ++++++++----------
 routes.py                     | 10 ++++++----
 3 files changed, 27 insertions(+), 27 deletions(-)

diff --git a/blender/avatar_main_output.py b/blender/avatar_main_output.py
index e9e3fc7..b62af49 100644
--- a/blender/avatar_main_output.py
+++ b/blender/avatar_main_output.py
@@ -56,7 +56,7 @@ class AvatarMainOutput(blender_node.ObjectOps):
     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='', upload_to_cloud=False):
+    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='', upload_to_cloud=False, SHAPE_FLOW=''):
         if open_in_blender:
             p = blender_path_override if blender_path_override else global_blender_path
             output_file = self.output_dir + '/tmp.blend'
@@ -75,22 +75,22 @@ class AvatarMainOutput(blender_node.ObjectOps):
         global_bpy.set_should_reset_scene(True)
 
         outputs = {
-            "ui": {
-                "gltfFilename": {filepath.replace(f"{self.output_dir}/", "")}, 
-                "files": [{
-                    "filename": filepath.replace(f"{self.output_dir}/", ""),
-                    "content_type": "model/gltf+json",
-                },],
-                "SHAPE_FLOW": {SHAPE_FLOW}, 
-                "auto_save": {'true' if auto_save else 'false'}, 
-            }
+            "gltfFilename": [filepath.replace(f"{self.output_dir}/", "")],
+            "files": [{
+                "filename": filepath.replace(f"{self.output_dir}/", ""),
+                "content_type": "model/gltf+json",
+            },],
+            "SHAPE_FLOW": {SHAPE_FLOW}, 
+            "auto_save": {'true' if auto_save else 'false'}, 
         }
         if upload_to_cloud:
             avatarId = upload_avatar_file(outputs)
             return {
                 "ui": {
-                    **outputs["ui"],
-                    "avatarId": {avatarId},
+                    **outputs,
+                    "avatarId": [avatarId],
                 },
             }
-        return outputs
+        return {
+            "ui": outputs
+        }
diff --git a/blender/mesh_utils.py b/blender/mesh_utils.py
index 6af949e..a5b0e6a 100644
--- a/blender/mesh_utils.py
+++ b/blender/mesh_utils.py
@@ -250,17 +250,15 @@ def export_gltf(output_dir, bpy_objects, filename, model_type, write_mode, metad
 
     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 get_avatar_file(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)
+def upload_avatar_file(output):
+    file = get_avatar_file(output)
     response = requests.get("https://labs.avatech.ai/api/share")
     labData = response.json()
     modelId = labData["modelId"]
diff --git a/routes.py b/routes.py
index 6ce1372..3bcac52 100644
--- a/routes.py
+++ b/routes.py
@@ -281,10 +281,12 @@ async def post_prompt_block(request):
         if history:
             # file = get_avatar_file(history[prompt_id]["outputs"])
             # return web.Response(body=file)
-
-            modelId = upload_avatar_file(history[prompt_id]["outputs"])
-            print("model id", modelId)
-            return web.json_response({"id": modelId}, status=200)
+            outputs = history[prompt_id]["outputs"]
+            for node_id, output in outputs.items():
+                if "gltfFilename" in output:
+                    modelId = upload_avatar_file(output)
+                    print("model id", modelId)
+                    return web.json_response({"id": modelId}, status=200)
         time.sleep(0.5)