merge: 🔀 pull request #50 from melMass/dev/august-refactor

Author: melMass

__init__.py

@@ -45,19 +45,15 @@ def extract_nodes_from_source(filename):
                 if isinstance(target, ast.Name) and target.id == "__nodes__":
                     value = ast.get_source_segment(source_code, node.value)
                     node_value = ast.parse(value).body[0].value
-                    if isinstance(node_value, ast.List) or isinstance(
-                        node_value, ast.Tuple
-                    ):
-                        for element in node_value.elts:
-                            if isinstance(element, ast.Name):
-                                print(element.id)
-                                nodes.append(element.id)
-
+                    if isinstance(node_value, (ast.List, ast.Tuple)):
+                        nodes.extend(
+                            element.id
+                            for element in node_value.elts
+                            if isinstance(element, ast.Name)
+                        )
                     break
     except SyntaxError:
         log.error("Failed to parse")
-        pass  # File couldn't be parsed
-
     return nodes
 
 
@@ -187,6 +183,14 @@ def load_nodes():
 from .endpoint import endlog
 
 if hasattr(PromptServer, "instance"):
+    restore_deps = ["basicsr"]
+    swap_deps = ["insightface", "onnxruntime"]
+
+    node_dependency_mapping = {
+        "FaceSwap": swap_deps,
+        "LoadFaceSwapModel": swap_deps,
+        "LoadFaceAnalysisModel": restore_deps,
+    }
 
     @PromptServer.instance.routes.get("/mtb/status")
     async def get_full_library(request):
@@ -202,7 +206,13 @@ async def get_full_library(request):
                 NODE_CLASS_MAPPINGS_DEBUG, title="Registered"
             )
             html_response += endpoint.render_table(
-                {k: "-" for k in failed}, title="Failed to load"
+                {
+                    k: {"dependencies": node_dependency_mapping.get(k)}
+                    if node_dependency_mapping.get(k)
+                    else "-"
+                    for k in failed
+                },
+                title="Failed to load",
             )
 
             return web.Response(
@@ -226,11 +236,10 @@ async def set_debug(request):
             log.setLevel(logging.DEBUG)
             log.debug("Debug mode set from API (/mtb/debug POST route)")
 
-        else:
-            if "MTB_DEBUG" in os.environ:
-                # del os.environ["MTB_DEBUG"]
-                os.environ.pop("MTB_DEBUG")
-                log.setLevel(logging.INFO)
+        elif "MTB_DEBUG" in os.environ:
+            # del os.environ["MTB_DEBUG"]
+            os.environ.pop("MTB_DEBUG")
+            log.setLevel(logging.INFO)
 
         return web.json_response(
             {"message": f"Debug mode {'set' if enabled else 'unset'}"}
@@ -244,7 +253,7 @@ async def get_home(request):
         # Check if the request prefers HTML content
         if "text/html" in request.headers.get("Accept", ""):
             # # Return an HTML page
-            html_response = f"""
+            html_response = """
             <div class="flex-container menu">
                 <a href="/mtb/debug">debug</a>
                 <a href="/mtb/status">status</a>
@@ -263,9 +272,7 @@ async def get_debug(request):
         from . import endpoint
 
         reload(endpoint)
-        enabled = False
-        if "MTB_DEBUG" in os.environ:
-            enabled = True
+        enabled = "MTB_DEBUG" in os.environ
         # Check if the request prefers HTML content
         if "text/html" in request.headers.get("Accept", ""):
             # # Return an HTML page
@@ -285,7 +292,7 @@ async def no_route(request):
         from . import endpoint
 
         if "text/html" in request.headers.get("Accept", ""):
-            html_response = f"""
+            html_response = """
             <h1>Actions has no get for now...</h1>
             """
             return web.Response(

endpoint.py

@@ -1,11 +1,35 @@
-from .utils import here
+from .utils import here, run_command, comfy_mode
 from aiohttp import web
 from .log import mklog
-import os
+import sys
 
 endlog = mklog("mtb endpoint")
 
-#- ACTIONS
+# - ACTIONS
+import requirements
+
+
+
+def ACTIONS_installDependency(dependency_names=None):
+    if dependency_names is None:
+        return {"error": "No dependency name provided"}
+    endlog.debug(f"Received Install Dependency request for {dependency_names}")
+    reqs = []
+    if comfy_mode == "embeded":
+        reqs = list(requirements.parse((here / "reqs_portable.txt").read_text()))
+    else:
+        reqs = list(requirements.parse((here / "reqs.txt").read_text()))
+    print([x.specs for x in reqs])
+    print(
+        "\n".join([f"{x.line} {''.join(x.specs[0] if x.specs else '')}" for x in reqs])
+    )
+    for dependency_name in dependency_names:
+        for req in reqs:
+            if req.name == dependency_name:
+                endlog.debug(f"Dependency {dependency_name} installed")
+                break
+    return {"success": True}
+
 
 def ACTIONS_getStyles(style_name=None):
     from .nodes.conditions import StylesLoader
@@ -19,10 +43,7 @@ def ACTIONS_getStyles(style_name=None):
             if not key.startswith("__") and key not in match_list
         }
         if style_name:
-            if style_name in filtered_styles:
-                return filtered_styles[style_name]
-            else:
-                return {"error": "Style not found"}
+            return filtered_styles.get(style_name, {"error": "Style not found"})
         return filtered_styles
     return {"error": "No styles found"}
 
@@ -35,7 +56,7 @@ async def do_action(request) -> web.Response:
 
     endlog.debug(f"Received action request: {name} {args}")
 
-    method_name = "ACTIONS_" + name
+    method_name = f"ACTIONS_{name}"
     method = globals().get(method_name)
 
     if callable(method):
@@ -53,16 +74,36 @@ async def do_action(request) -> web.Response:
 
 
 # - HTML UTILS
+
+
+def dependencies_button(name, dependencies):
+    deps = ",".join([f"'{x}'" for x in dependencies])
+    return f"""
+        <button class="dependency-button" onclick="window.mtb_action('installDependency',[{deps}])">Install {name} deps</button>
+        """
+
+
 def render_table(table_dict, sort=True, title=None):
-    table_rows = ""
     table_dict = sorted(
         table_dict.items(), key=lambda item: item[0]
     )  # Sort the dictionary by keys
 
-    for name, description in table_dict:
-        table_rows += f"<tr><td>{name}</td><td>{description}</td></tr>"
+    table_rows = ""
+    for name, item in table_dict:
+        if isinstance(item, dict):
+            if "dependencies" in item:
+                table_rows += f"<tr><td>{name}</td><td>"
+                table_rows += f"{dependencies_button(name,item['dependencies'])}"
+
+                table_rows += "</td></tr>"
+            else:
+                table_rows += f"<tr><td>{name}</td><td>{render_table(item)}</td></tr>"
+        # elif isinstance(item, str):
+        #     table_rows += f"<tr><td>{name}</td><td>{item}</td></tr>"
+        else:
+            table_rows += f"<tr><td>{name}</td><td>{item}</td></tr>"
 
-    html_response = f"""
+    return f"""
         <div class="table-container">
         {"" if title is None else f"<h1>{title}</h1>"}
         <table>
@@ -78,7 +119,6 @@ def render_table(table_dict, sort=True, title=None):
         </table>      
         </div>
         """
-    return html_response
 
 
 def render_base_template(title, content):
@@ -98,6 +138,29 @@ def render_base_template(title, content):
             {css_content}
         </style>
     </head>
+    <script type="module">
+        import {{ api }} from '/scripts/api.js'
+        const mtb_action = async (action, args) =>{{
+            console.log(`Sending ${{action}} with args: ${{args}}`)
+            }}
+        window.mtb_action = async (action, args) =>{{
+            console.log(`Sending ${{action}} with args: ${{args}} to the API`)
+            const res = await api.fetchApi('/actions', {{
+                method: 'POST',
+                body: JSON.stringify({{
+                  name: action,
+                  args,
+                }}),
+            }})
+
+              const output = await res.json()
+              console.debug(`Received ${{action}} response:`, output)
+              if (output?.result?.error){{
+                  alert(`An error occured: {{output?.result?.error}}`)
+              }}
+              return output?.result
+        }}
+    </script>
     <body>
         <header>
         <a href="/">Back to Comfy</a>
@@ -117,5 +180,6 @@ def render_base_template(title, content):
             <!-- Shared footer content here -->
         </footer>
     </body>
+    
     </html>
     """