import csv
import secrets
import sys
from pathlib import Path
from typing import Any, Literal
from aiohttp import web
from .log import mklog
from .utils import (
SortMode,
backup_file,
build_glob_patterns,
glob_multiple,
here,
import_install,
input_dir,
output_dir,
reqs_map,
run_command,
styles_dir,
)
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):
if dependency_names is None:
# return web.Response(text="No dependency name provided", status=400)
return {"error": "No dependency name provided"}
endlog.debug(f"Received Install Dependency request for {dependency_names}")
# reqs = []
resolved_names = [reqs_map.get(name, name) for name in dependency_names]
try:
run_command(
[Path(sys.executable), "-m", "pip", "install"] + resolved_names
)
return {"success": True}
except Exception as e:
return {"error": f"Failed to install dependencies: {e}"}
# if platform.system() == "Windows":
# reqs = list(requirements.parse((here / "reqs_windows.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
def ACTIONS_getUserImages(
mode: Literal["input", "output"],
count=200,
offset=0,
sort: str | None = None,
include_subfolders: bool = False,
):
# enabled = "MTB_EXPOSE" in os.environ
# if not enabled:
# return {"error": "Session not authorized to getInputs"}
imgs = {}
entry_dir = input_dir if mode == "input" else output_dir
supported = ["png", "jpg", "jpeg", "webp", "gif"]
entries = {}
patterns = build_glob_patterns(supported, recursive=include_subfolders)
entries = glob_multiple(entry_dir, patterns)
sort_mode = SortMode.from_str(sort)
if sort_mode:
sort_key = {
SortMode.MODIFIED: lambda x: x.stat().st_mtime,
SortMode.MODIFIED_REVERSE: lambda x: x.stat().st_mtime,
SortMode.NAME: lambda x: x.name,
SortMode.NAME_REVERSE: lambda x: x.name,
}.get(sort_mode)
if sort_key:
reverse = sort_mode in (SortMode.MODIFIED, SortMode.NAME_REVERSE)
entries = sorted(entries, key=sort_key, reverse=reverse)
imgs = {
img.name: (
f"/mtb/view?filename={img.name}&width=512&type={mode}&subfolder="
f"{img.parent.relative_to(entry_dir) if include_subfolders else ''}"
f"&preview=&rand={secrets.randbelow(424242)}"
)
for i, img in enumerate(entries)
if offset <= i < offset + count
}
return imgs
def ACTIONS_getStyles(style_name=None):
from .nodes.conditions import MTB_StylesLoader
styles = MTB_StylesLoader.options
match_list = ["name"]
if styles:
filtered_styles = {
key: value
for key, value in styles.items()
if not key.startswith("__") and key not in match_list
}
if style_name:
return filtered_styles.get(
style_name, {"error": "Style not found"}
)
return filtered_styles
return {"error": "No styles found"}
def ACTIONS_saveStyle(data):
# endlog.debug(f"Received Save Styles for {data.keys()}")
# endlog.debug(data)
styles = [f.name for f in styles_dir.iterdir() if f.suffix == ".csv"]
target = None
rows = []
for fp, content in data.items():
if fp in styles:
endlog.debug(f"Overwriting {fp}")
target = styles_dir / fp
rows = content
break
if not target:
endlog.warning(
f"Could not determine the target file for {data.keys()}"
)
return {"error": "Could not determine the target file for the style"}
backup_file(target)
with target.open("w", newline="", encoding="utf-8") as file:
csv_writer = csv.writer(file, quoting=csv.QUOTE_ALL)
for row in rows:
csv_writer.writerow(row)
async def do_action(request: web.Request) -> web.Response:
endlog.debug("Init action request")
request_data = await request.json()
name = request_data.get("name")
args = request_data.get("args")
endlog.debug(f"Received action request: {name} {args}")
method_name = f"ACTIONS_{name}"
method = globals().get(method_name)
if callable(method):
result = None
if args:
result = method(*args) if isinstance(args, list) else method(args)
else:
result = method()
endlog.debug(f"Action result: {result}")
return web.json_response({"result": result})
available_methods = [
attr[len("ACTIONS_") :]
for attr in globals()
if attr.startswith("ACTIONS_")
]
return web.json_response(
{
"error": "Invalid method name.",
"available_methods": available_methods,
}
)
# - HTML UTILS
def dependencies_button(name: str, dependencies: list[str]) -> str:
deps = ",".join([f"'{x}'" for x in dependencies])
return f"""
"""
def csv_editor():
inputs = [f for f in styles_dir.iterdir() if f.suffix == ".csv"]
# rows = {f.stem: list(csv.reader(f.read_text("utf8"))) for f in styles}
style_files = {}
for file in inputs:
with open(file, encoding="utf8") as f:
parsed = csv.reader(f)
style_files[file.name] = []
for row in parsed:
endlog.debug(f"Adding style {row[0]}")
style_files[file.name].append((row[0], row[1], row[2]))
html_out = """
Style Editor
"""
for current, styles in style_files.items():
current_out = f"
{current}
"
table_rows = []
for index, style in enumerate(styles):
table_rows += (
(["