feat: ✨ 3D Tools

endpoint.py

@@ -12,6 +12,7 @@
     backup_file,
     build_glob_patterns,
     glob_multiple,
+    here,
     import_install,
     input_dir,
     output_dir,
@@ -23,7 +24,93 @@
 endlog = mklog("mtb endpoint")
 
 # - ACTIONS
+import asyncio
+import platform
+import sys
+from pathlib import Path
+
+try:
+    import websockets.server
+except ModuleNotFoundError:
+    endlog.warning(
+        "You do not have websockets installed, the video server won't work"
+    )
+    websockets = False
+
 import_install("requirements")
+import io
+
+import numpy as np
+from PIL import Image
+
+
+def generate_random_frame():
+    # Generate a random image frame
+    width, height = 640, 480
+    image = np.random.randint(0, 256, (height, width, 3), dtype=np.uint8)
+    pil_image = Image.fromarray(image)
+    byte_buffer = io.BytesIO()
+    pil_image.save(byte_buffer, format="JPEG")
+    frame_data = byte_buffer.getvalue()
+    return frame_data
+
+
+class VideoStreamingManager:
+    def __init__(self):
+        self.video_servers = {}
+        self.next_port = (
+            8767  # Start with a default port and increment for each server
+        )
+
+    async def start_video_streaming_server(self, video_id):
+        if video_id not in self.video_servers:
+            # Create and start a new video streaming server for the specified video
+            video_server = await self.create_video_streaming_server(video_id)
+            self.video_servers[video_id] = video_server
+
+            return video_server
+
+    async def video_stream(self, websocket, path):
+        # Implement the logic to continuously capture and send video frames here
+        while True:
+            # frame_data = capture_and_encode_frame()  # Implement this function
+            frame_data = generate_random_frame()
+            await websocket.send(frame_data)
+            await asyncio.sleep(0.033)  # Adjust the frame rate as needed
+
+    async def create_video_streaming_server(self, video_id):
+        # Create and start a new WebSocket server for the specified video
+        port = self.next_port
+        self.next_port += 1  # Increment port number for the next server
+
+        server = await websockets.server.serve(
+            self.video_stream, "localhost", port
+        )
+
+        return server
+
+    async def stop_video_streaming_server(self, video_id):
+        if video_id in self.video_servers:
+            # Terminate and remove the video streaming server for the specified video
+            video_server = self.video_servers[video_id]
+            video_server.close()
+            await video_server.wait_closed()
+            del self.video_servers[video_id]
+
+
+async def start_video_streaming_server():
+    async def video_stream(websocket, path):
+        # Continuously capture and send video frames here
+        while True:
+            frame_data = capture_and_encode_frame()  # Implement this function
+            await websocket.send(frame_data)
+            await asyncio.sleep(0.033)  # Adjust the frame rate as needed
+
+    start_server = websockets.server.serve(
+        video_stream, "localhost", 8766
+    )  # Use a different port (e.g., 8766)
+
+    return await start_server
 
 
 def ACTIONS_installDependency(dependency_names=None):
@@ -315,7 +402,7 @@ def add_split_pane(
     """
 
 
-def add_dropdown(title: str, options: list[str]):
+def add_dropdown(title, options):
     option_str = "\n".join(
         [f"<option value='{opt}'>{opt}</option>" for opt in options]
     )

env.nu

@@ -1,5 +1,85 @@
 # NOTE: This file is only use for development you can ignore it
 
+
+# NOTE: for CI it's easier to extract parts of my cli for now
+
+const THREE_VERSION = "0.171.0"
+# Update the external web extensions
+export def "comfy mtb update-web" [] {
+
+  let async_dir = $"($env.COMFY_MTB)/web_async"
+  let three_base = $"https://cdn.jsdelivr.net/npm/three@($THREE_VERSION)"
+  let three = {
+    "." : [
+      "build/three.module.js",
+      "build/three.core.js",
+    ],
+    three_addons/capabilities: [
+      "examples/jsm/capabilities/WebGPU.js",
+      "examples/jsm/controls/ArcballControls.js",
+      "examples/jsm/controls/DragControls.js",
+      "examples/jsm/controls/FirstPersonControls.js",
+      "examples/jsm/controls/FlyControls.js",
+      "examples/jsm/controls/MapControls.js",
+      "examples/jsm/controls/OrbitControls.js",
+      "examples/jsm/controls/PointerLockControls.js",
+      "examples/jsm/controls/TrackballControls.js",
+      "examples/jsm/controls/TransformControls.js",
+    ],
+    three_addons/offscreen: [
+      "jank.js",
+      "offscreen.js",
+      "scene.js",
+    ],
+    thee_addons/exporters : [
+      "examples/jsm/exporters/DRACOExporter.js",
+      "examples/jsm/exporters/EXRExporter.js",
+      "examples/jsm/exporters/GLTFExporter.js",
+      "examples/jsm/exporters/KTX2Exporter.js",
+      "examples/jsm/exporters/MMDExporter.js",
+      "examples/jsm/exporters/OBJExporter.js",
+      "examples/jsm/exporters/PLYExporter.js",
+      "examples/jsm/exporters/STLExporter.js",
+      "examples/jsm/exporters/USDZExporter.js"
+    ],
+
+    three_addons/loaders : [
+      "examples/jsm/loaders/3DMLoader.js",
+      "examples/jsm/loaders/BVHLoader.js",
+      "examples/jsm/loaders/ColladaLoader.js",
+      "examples/jsm/loaders/DRACOLoader.js",
+      "examples/jsm/loaders/EXRLoader.js",
+      "examples/jsm/loaders/FBXLoader.js",
+      "examples/jsm/loaders/FontLoader.js",
+      "examples/jsm/loaders/GLTFLoader.js",
+      "examples/jsm/loaders/HDRCubeTextureLoader.js",
+      "examples/jsm/loaders/MaterialXLoader.js",
+      "examples/jsm/loaders/MTLLoader.js",
+      "examples/jsm/loaders/OBJLoader.js",
+      "examples/jsm/loaders/PCDLoader.js",
+      "examples/jsm/loaders/PDBLoader.js",
+      "examples/jsm/loaders/PLYLoader.js",
+      "examples/jsm/loaders/STLLoader.js",
+      "examples/jsm/loaders/UltraHDRLoader.js",
+      "examples/jsm/loaders/USDZLoader.js",
+      "examples/jsm/loaders/VOXLoader.js"
+    ]
+  }
+  $three | items {|root,urls|
+    let dest =  $async_dir | path join $root
+    mkdir $dest
+
+    $urls | par-each {|url|
+      let url = $"($three_base)/($url)"
+      let local = ($dest | path join ($url | path basename))
+      wget -c $url -O ($local)
+    }
+  }
+
+  # $three
+}
+
+
 def get_root [--clean] {
     if $clean {
         $env.COMFY_CLEAN_ROOT
@@ -151,7 +231,7 @@ def --env path-add [pth] {
 
 
 export-env {
-  $env.COMFY_MTB = ("." | path expand)
+  $env.COMFY_MTB = ("." | path expand | str replace -a '\' '/')
   # $env.CUDA_ROOT =  'C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v12.1\'
 
   $env.CUDA_HOME = $env.CUDA_ROOT

log.py

@@ -7,36 +7,55 @@
 
 # Custom object that discards the output
 class NullWriter:
+    """Custom object that discards the output."""
+
     def write(self, text):
         pass
 
 
-class Formatter(logging.Formatter):
+class ConsoleFormatter(logging.Formatter):
+    """Formatter for console based log, using base ansi colors."""
+
     grey = "\x1b[38;20m"
     cyan = "\x1b[36;20m"
     purple = "\x1b[35;20m"
     yellow = "\x1b[33;20m"
     red = "\x1b[31;20m"
     bold_red = "\x1b[31;1m"
     reset = "\x1b[0m"
-    # format = "%(asctime)s - [%(name)s] - %(levelname)s - %(message)s (%(filename)s:%(lineno)d)"
-    format = "[%(name)s] | %(levelname)s -> %(message)s"
+    # format = ("%(asctime)s - [%(name)s] - %(levelname)s "
+    # "- %(message)s (%(filename)s:%(lineno)d)")
+    fmt = "[%(name)s] | %(levelname)s -> %(message)s"
 
     FORMATS = {
-        logging.DEBUG: purple + format + reset,
-        logging.INFO: cyan + format + reset,
-        logging.WARNING: yellow + format + reset,
-        logging.ERROR: red + format + reset,
-        logging.CRITICAL: bold_red + format + reset,
+        logging.DEBUG: f"{purple}{fmt}{reset}",
+        logging.INFO: f"{cyan}{fmt}{reset}",
+        logging.WARNING: f"{yellow}{fmt}{reset}",
+        logging.ERROR: f"{red}{fmt}{reset}",
+        logging.CRITICAL: f"{bold_red}{fmt}{reset}",
     }
 
     def format(self, record):
         log_fmt = self.FORMATS.get(record.levelno)
+
         formatter = logging.Formatter(log_fmt)
         return formatter.format(record)
 
 
-def mklog(name: str, level: int = base_log_level):
+class FileFormatter(logging.Formatter):
+    """Formatter for file base logs."""
+
+    # File specific formatting
+    fmt = (
+        "%(asctime)s - [%(name)s] - "
+        "%(levelname)s - %(message)s (%(filename)s:%(lineno)d)"
+    )
+
+    def __init__(self):
+        super().__init__(self.fmt, "%Y-%m-%d %H:%M:%S")
+
+
+def mklog(name: str, level: int = base_log_level, log_file: str | None = None):
     logger = logging.getLogger(name)
     logger.setLevel(level)
 
@@ -45,9 +64,16 @@ def mklog(name: str, level: int = base_log_level):
 
     ch = logging.StreamHandler()
     ch.setLevel(level)
-    ch.setFormatter(Formatter())
+    ch.setFormatter(ConsoleFormatter())
     logger.addHandler(ch)
 
+    if log_file:
+        # file handler
+        fh = logging.FileHandler(log_file)
+        fh.setLevel(level)
+        fh.setFormatter(FileFormatter())
+        logger.addHandler(fh)
+
     # Disable log propagation
     logger.propagate = False
 

node_list.json

@@ -23,12 +23,20 @@
     "Film Interpolation (mtb)": "Google Research FILM frame interpolation for large motion",
     "Fit Number (mtb)": "Fit the input float using a source and target range",
     "Float To Number (mtb)": "Node addon for the WAS Suite. Converts a \"comfy\" FLOAT to a NUMBER.",
+    "Geometry Box (mtb)": "Makes a Box 3D geometry",
+    "Geometry Decimater (mtb)": "Optimized the geometry to match the target number of triangles",
+    "Geometry Info (mtb)": "Retrieve information about a 3D geometry",
+    "Geometry Sphere (mtb)": "Makes a Sphere 3D geometry",
+    "Geometry Test (mtb)": "Fetches an Open3D data geometry",
     "Get Batch From History (mtb)": "Very experimental node to load images from the history of the server.\n\n    Queue items without output are ignored in the count.",
     "Image Compare (mtb)": "Compare two images and return a difference image",
+    "Image Distort With Uv (mtb)": "Distorts an image based on a UV map.",
     "Image Premultiply (mtb)": "Premultiply image with mask",
     "Image Remove Background Rembg (mtb)": "Removes the background from the input using Rembg.",
     "Image Resize Factor (mtb)": "Extracted mostly from WAS Node Suite, with a few edits (most notably multiple image support) and less features.",
     "Image Tile Offset (mtb)": "Mimics an old photoshop technique to check for seamless textures",
+    "Image To Uv (mtb)": "Turn an image back into a UV map. (Shallow converter)",
+    "Image Tile Offset (mtb)": "Mimics an old photoshop technique to check for seamless textures",
     "Int To Bool (mtb)": "Basic int to bool conversion",
     "Int To Number (mtb)": "Node addon for the WAS Suite. Converts a \"comfy\" INT to a NUMBER.",
     "Interpolate Clip Sequential (mtb)": null,
@@ -37,12 +45,14 @@
     "Load Face Enhance Model (mtb)": "Loads a GFPGan or RestoreFormer model for face enhancement.",
     "Load Face Swap Model (mtb)": "Loads a faceswap model",
     "Load Film Model (mtb)": "Loads a FILM model",
+    "Load Geometry (mtb)": "Load a 3D geometry",
     "Load Image From Url (mtb)": "Load an image from the given URL",
     "Load Image Sequence (mtb)": "Load an image sequence from a folder. The current frame is used to determine which image to load.\n\n    Usually used in conjunction with the `Primitive` node set to increment to load a sequence of images from a folder.\n    Use -1 to load all matching frames as a batch.\n    ",
     "Mask To Image (mtb)": "Converts a mask (alpha) to an RGB image with a color and background",
+    "Model Patch Seamless (mtb)": "Experimental patcher to enable the circular padding mode of the sd model layers, requires a custom VAE",
     "Math Expression (mtb)": "Node to evaluate a simple math expression string",
     "Model Patch Seamless (mtb)": "Uses the stable diffusion 'hack' to infer seamless images by setting the model layers padding mode to circular (experimental)",
-    "Pick From Batch (mtb)": "Pick a specific number of images from a batch, either from the start or end.",
+    "Pick From Batch (mtb)": "Pick a specific number of images from a batch, either from the start or end.",
     "Qr Code (mtb)": "Basic QR Code generator",
     "Restore Face (mtb)": "Uses GFPGan to restore faces",
     "Save Gif (mtb)": "Save the images from the batch as a GIF",
@@ -55,8 +65,15 @@
     "String Replace (mtb)": "Basic string replacement",
     "Styles Loader (mtb)": "Load csv files and populate a dropdown from the rows (\u00e0 la A111)",
     "Text To Image (mtb)": "Utils to convert text to image using a font\n\n\n    The tool looks for any .ttf file in the Comfy folder hierarchy.\n    ",
+    "Transform Geometry (mtb)": "Transforms the input geometry",
     "Transform Image (mtb)": "Save torch tensors (image, mask or latent) to disk, useful to debug things outside comfy\n\n\n    it return a tensor representing the transformed images with the same shape as the input tensor\n    ",
     "Uncrop (mtb)": "Uncrops an image to a given bounding box\n\n    The bounding box can be given as a tuple of (x, y, width, height) or as a BBOX type\n    The BBOX input takes precedence over the tuple input",
     "Unsplash Image (mtb)": "Unsplash Image given a keyword and a size",
-    "Vae Decode (mtb)": "Wrapper for the 2 core decoders but also adding the sd seamless hack, taken from: FlyingFireCo/tiled_ksampler"
-}
+    "Uv Distort (mtb)": "Applies a polar coordinates or wave distortion to the UV map",
+    "Uv Map (mtb)": "Generates a UV Map tensor given a widht and height",
+    "Uv Remove Seams (mtb)": "Blends values near the UV borders to mitigate visible seams.",
+    "Uv Tile (mtb)": "Tiles the UV map based on the specified number of tiles.",
+    "Uv To Image (mtb)": "Converts the UV map to an image. (Shallow converter)",
+    "Vae Decode (mtb)": "Wrapper for the 2 core decoders (nomarl and tiled) but also adding the sd seamless hack, taken from: FlyingFireCo/tiled_ksampler"
+}
+