feat(backend): add ascii animation for parallax run/parallax join

src/backend/main.py

@@ -13,6 +13,7 @@
 from backend.server.scheduler_manage import SchedulerManage
 from backend.server.server_args import parse_args
 from backend.server.static_config import get_model_list, get_node_join_command
+from parallax_utils.ascii_anime import display_parallax_run
 from parallax_utils.logging_config import get_logger
 
 app = FastAPI()
@@ -115,6 +116,7 @@ async def openai_v1_chat_completions(raw_request: Request):
     args = parse_args()
     logger.info(f"args: {args}")
 
+    display_parallax_run()
     host_maddrs = args.host_maddrs
     dht_port = args.dht_port
     if args.dht_port is not None:

src/parallax/server/executor.py

@@ -56,6 +56,7 @@
     pad_inputs,
     pad_prefix_caches,
 )
+from parallax_utils.ascii_anime import display_parallax_join
 from parallax_utils.logging_config import get_logger
 from parallax_utils.utils import compute_max_batch_size
 
@@ -273,6 +274,7 @@ def __init__(
             self.send_to_ipc_socket = get_zmq_socket(
                 self.zmq_context, zmq.PUSH, executor_output_ipc_addr, bind=False
             )
+        display_parallax_join(model_repo)
 
     @classmethod
     def create_from_args(cls, args: argparse.Namespace):

src/parallax_utils/ascii_anime.py

@@ -0,0 +1,203 @@
+import json
+import math
+import os
+import time
+
+
+class HexColorPrinter:
+    COLOR_MAP = {
+        "#000000": ("\033[30m", (0, 0, 0)),
+        "#800000": ("\033[31m", (128, 0, 0)),
+        "#008000": ("\033[32m", (0, 128, 0)),
+        "#808000": ("\033[33m", (128, 128, 0)),
+        "#000080": ("\033[34m", (0, 0, 128)),
+        "#800080": ("\033[35m", (128, 0, 128)),
+        "#008080": ("\033[36m", (0, 128, 128)),
+        "#c0c0c0": ("\033[37m", (192, 192, 192)),
+        "#808080": ("\033[90m", (128, 128, 128)),
+        "#ff0000": ("\033[91m", (255, 0, 0)),
+        "#00ff00": ("\033[92m", (0, 255, 0)),
+        "#ffff00": ("\033[93m", (255, 255, 0)),
+        "#0000ff": ("\033[94m", (0, 0, 255)),
+        "#ff00ff": ("\033[95m", (255, 0, 255)),
+        "#00ffff": ("\033[96m", (0, 255, 255)),
+        "#ffffff": ("\033[97m", (255, 255, 255)),
+    }
+
+    RESET = "\033[0m"
+    SHOW = "\033[97m"
+    WHITE = "\033[97m"
+
+    @classmethod
+    def hex_to_rgb(cls, hex_color):
+        hex_color = hex_color.lstrip("#")
+        return tuple(int(hex_color[i : i + 2], 16) for i in (0, 2, 4))
+
+    @classmethod
+    def color_distance(cls, rgb1, rgb2):
+        return math.sqrt(sum((c1 - c2) ** 2 for c1, c2 in zip(rgb1, rgb2)))
+
+    @classmethod
+    def find_closest_color(cls, target_hex):
+        target_rgb = cls.hex_to_rgb(target_hex)
+        min_distance = float("inf")
+        closest_color = "\033[97m"
+
+        for _, (ansi_code, rgb) in cls.COLOR_MAP.items():
+            distance = cls.color_distance(target_rgb, rgb)
+            if distance < min_distance:
+                min_distance = distance
+                closest_color = ansi_code
+        if closest_color == "\033[37m":
+            closest_color = "\033[35m"
+        if closest_color == "\033[90m":
+            closest_color = "\033[95m"
+
+        return closest_color
+
+
+def clear_screen():
+    # Clear screen command for different operating systems
+    os.system("cls" if os.name == "nt" else "clear")
+
+
+def handle_colors_data(raw_data):
+    color_dict = {}
+    if raw_data is not None:
+        config = json.loads(raw_data)
+        for key, value in config.items():
+            if isinstance(value, str) and value.startswith("#"):
+                color_dict[key] = value
+    return color_dict
+
+
+def process_context_color_run(content, colors):
+    res = []
+    for row, row_str in enumerate(content):
+        processed_row = ""
+        for column, text in enumerate(row_str):
+            if text in (" ", "#"):
+                processed_row += text
+                continue
+            position_str = str(column) + "," + str(row)
+            if row == 11 and text not in ("▝", "#", " "):
+                color = HexColorPrinter.WHITE
+                processed_row += color
+            else:
+                hex_color = colors.get(position_str, None)
+                if hex_color:
+                    color = HexColorPrinter.find_closest_color(hex_color)
+                    processed_row += color
+            processed_row += text
+        processed_row += HexColorPrinter.RESET
+        res.append(processed_row)
+    return res
+
+
+def process_context_color_join(content, colors, model_name):
+    res = []
+    if len(model_name) > 22:
+        model_name = model_name[:22]
+    name_len = len(model_name)
+    for row, row_str in enumerate(content):
+        processed_row = ""
+        for column, text in enumerate(row_str):
+            if text in (" ", "#"):
+                processed_row += text
+                continue
+            if row == 5 and 10 < column < 32:
+                pos = column - 11
+                if pos < name_len:
+                    text = model_name[pos]
+                    processed_row += HexColorPrinter.RESET
+                else:
+                    text = " "
+                    position_str = str(column) + "," + str(row)
+                    hex_color = colors.get(position_str, None)
+                    if hex_color:
+                        color = HexColorPrinter.find_closest_color(hex_color)
+                        processed_row += color
+            else:
+                position_str = str(column) + "," + str(row)
+                hex_color = colors.get(position_str, None)
+                if hex_color:
+                    color = HexColorPrinter.find_closest_color(hex_color)
+                    processed_row += color
+            processed_row += text
+        processed_row += HexColorPrinter.RESET
+        res.append(processed_row)
+    return res
+
+
+def display_ascii_animation_run(animation_data):
+    frames = animation_data.get("frames", [])
+    # loop = animation_data.get('loop', False)
+
+    if not frames:
+        print("No animation frames found in the JSON data.")
+        return
+
+    for frame_data in frames:
+        content = frame_data.get("content", None)
+        delay = frame_data.get("duration", 30) / 1000.0
+        colors_data = frame_data.get("colors", None)
+        foreground = colors_data.get("foreground", None)
+        colors = handle_colors_data(foreground)
+
+        if content:
+            res = process_context_color_run(content, colors)
+            res = "\n".join(res).replace("#", " ")
+            clear_screen()
+            print(res)
+            time.sleep(delay)
+
+
+def display_ascii_animation_join(animation_data, model_name):
+    frames = animation_data.get("frames", [])
+    # loop = animation_data.get('loop', False)
+
+    if not frames:
+        print("No animation frames found in the JSON data.")
+        return
+
+    for frame_data in frames:
+        content = frame_data.get("content", None)
+        delay = frame_data.get("duration", 30) / 1000.0
+        colors_data = frame_data.get("colors", None)
+        foreground = colors_data.get("foreground", None)
+        colors = handle_colors_data(foreground)
+
+        if content:
+            res = process_context_color_join(content, colors, model_name)
+            res = "\n".join(res).replace("#", " ")
+            clear_screen()
+            print(res)
+            time.sleep(delay)
+
+
+def display_parallax_run():
+    file_path = "./src/parallax_utils/anime/parallax_run.json"
+    try:
+        with open(file_path, "r") as f:
+            animation_data = json.load(f)
+    except FileNotFoundError:
+        print(f"Error: The file '{file_path}' was not found.")
+        return
+    except json.JSONDecodeError:
+        print(f"Error: The file '{file_path}' contains invalid JSON.")
+        return
+    display_ascii_animation_run(animation_data)
+
+
+def display_parallax_join(model_name):
+    file_path = "./src/parallax_utils/anime/parallax_join.json"
+    try:
+        with open(file_path, "r") as f:
+            animation_data = json.load(f)
+    except FileNotFoundError:
+        print(f"Error: The file '{file_path}' was not found.")
+        return
+    except json.JSONDecodeError:
+        print(f"Error: The file '{file_path}' contains invalid JSON.")
+        return
+    display_ascii_animation_join(animation_data, model_name)