1102 lines
34 KiB
Python
1102 lines
34 KiB
Python
# ruff: noqa: S102 - Use of `exec` detected
|
|
# ruff: noqa: S307 - Use of possibly insecure function
|
|
|
|
|
|
# region imports
|
|
import ast
|
|
import base64
|
|
import code
|
|
import io
|
|
import re
|
|
import sys
|
|
import textwrap
|
|
import time
|
|
import urllib.parse
|
|
import uuid
|
|
from contextlib import redirect_stderr, redirect_stdout
|
|
from pathlib import Path
|
|
from typing import TYPE_CHECKING, Any
|
|
|
|
import folder_paths
|
|
import numpy as np
|
|
import rich
|
|
import torch
|
|
from aiohttp import web
|
|
from PIL import Image
|
|
from rich.console import Console
|
|
from rich.traceback import Traceback
|
|
|
|
if TYPE_CHECKING:
|
|
from .nodes.audio import AudioTensor
|
|
|
|
from .log import log
|
|
from .utils import singleton
|
|
|
|
# region safe imports
|
|
try:
|
|
import pyflakes.api
|
|
import pyflakes.reporter
|
|
|
|
_HAS_LINT = True
|
|
except ImportError:
|
|
log.error(
|
|
textwrap.dedent("""
|
|
ComfyREPL: pyflakes not found.
|
|
Linting will be disabled.
|
|
Install with 'pip install pyflakes'.
|
|
""")
|
|
)
|
|
_HAS_LINT = False
|
|
|
|
# --- Linting Library ---
|
|
# try:
|
|
# import ruff
|
|
# import ruff.lint
|
|
# import ruff.lint.linter
|
|
# import ruff.settings
|
|
#
|
|
# _HAS_LINT = True
|
|
# except ImportError:
|
|
# print(
|
|
# """ComfyREPL: ruff not found.
|
|
# Linting will be disabled.
|
|
# Install with 'pip install ruff'.
|
|
# """
|
|
# )
|
|
# _HAS_LINT = False
|
|
|
|
# --- Audio/Video Libraries ---
|
|
try:
|
|
import scipy.io.wavfile
|
|
|
|
_HAS_SCIPY = True
|
|
|
|
|
|
except ImportError:
|
|
log.warning(
|
|
textwrap.dedent("""
|
|
ComfyREPL: SciPy not found.
|
|
Audio display will be disabled. Install with 'pip install scipy'.
|
|
""")
|
|
)
|
|
_HAS_SCIPY = False
|
|
|
|
try:
|
|
import imageio
|
|
# import imageio.plugins.ffmpeg
|
|
|
|
_HAS_IMAGEIO = True
|
|
except ImportError:
|
|
log.warning(
|
|
textwrap.dedent(
|
|
"""ComfyREPL: Imageio or imageio-ffmpeg not found.
|
|
Video display will be disabled.
|
|
Install with 'pip install imageio imageio-ffmpeg'
|
|
or even better use uv.
|
|
"""
|
|
)
|
|
)
|
|
_HAS_IMAGEIO = False
|
|
|
|
# endregion
|
|
# endregion
|
|
|
|
|
|
# region constants
|
|
# NOTE: The best way I can think of to make outputs dynamic
|
|
# is to register a large number from the python side and
|
|
# control their display dynamically from js (based on the dynamic inputs)
|
|
# but this is hacky and I would rather wait for V3 to start doing that.
|
|
SOCKET_COUNT = 5
|
|
# endregion
|
|
|
|
|
|
# region media_handlers
|
|
class _AudioDisplay:
|
|
def __init__(self, samples, sample_rate):
|
|
if not _HAS_SCIPY:
|
|
raise ImportError("Audio display requires scipy and numpy.")
|
|
if not isinstance(samples, np.ndarray | torch.Tensor):
|
|
raise TypeError(
|
|
"Audio samples must be a numpy array or torch tensor."
|
|
)
|
|
if isinstance(samples, torch.Tensor):
|
|
samples = samples.detach().cpu().numpy()
|
|
|
|
# Ensure samples are in a format scipy.io.wavfile
|
|
# can handle (e.g., int16, float32)
|
|
if samples.dtype == np.float64:
|
|
samples = samples.astype(np.float32)
|
|
elif samples.dtype == np.int64:
|
|
# Or scale to int32 if range requires
|
|
samples = samples.astype(np.int16)
|
|
|
|
self.samples = samples
|
|
self.sample_rate = sample_rate
|
|
|
|
def _to_wav_base64(self):
|
|
buffer = io.BytesIO()
|
|
try:
|
|
scipy.io.wavfile.write(buffer, self.sample_rate, self.samples)
|
|
audio_base64 = base64.b64encode(buffer.getvalue()).decode("utf-8")
|
|
return audio_base64
|
|
except Exception as e:
|
|
return f"<div style='color: red;'>Error encoding audio: {e}</div>"
|
|
|
|
def _repr_html_(self):
|
|
base64_data = self._to_wav_base64()
|
|
if base64_data.startswith("<div"):
|
|
return base64_data
|
|
return f"""
|
|
<audio
|
|
controls
|
|
src="data:audio/wav;base64,{base64_data}"
|
|
style="margin: 5px 0;"
|
|
/>"""
|
|
|
|
|
|
def render_audio(samples: np.ndarray | torch.Tensor, sample_rate: int):
|
|
"""
|
|
Render audio samples as an HTML audio player.
|
|
|
|
Args:
|
|
samples: Audio samples.
|
|
sample_rate: Sample rate in Hz.
|
|
|
|
Returns
|
|
-------
|
|
AudioDisplay: An object that will render as an HTML audio player.
|
|
"""
|
|
return _AudioDisplay(samples, sample_rate)
|
|
|
|
|
|
class _VideoDisplay:
|
|
def __init__(
|
|
self, frames, fps=24, options=None, audio: "AudioTensor | None" = None
|
|
):
|
|
if not _HAS_IMAGEIO:
|
|
raise ImportError(
|
|
textwrap.dedent(
|
|
"""Video display requires
|
|
imageio, imageio-ffmpeg, and image libraries
|
|
(numpy, Pillow, torch).
|
|
"""
|
|
)
|
|
)
|
|
|
|
self.frames = []
|
|
for frame in frames:
|
|
if isinstance(frame, Image.Image):
|
|
self.frames.append(np.array(frame))
|
|
elif isinstance(frame, np.ndarray):
|
|
# ensure HWC and uint8
|
|
# CHW
|
|
if frame.ndim == 3 and frame.shape[0] in [1, 3, 4]:
|
|
frame = np.transpose(frame, (1, 2, 0))
|
|
if frame.dtype != np.uint8:
|
|
frame = (
|
|
(frame * 255).astype(np.uint8)
|
|
if frame.max() <= 1.0
|
|
else frame.astype(np.uint8)
|
|
)
|
|
self.frames.append(frame)
|
|
elif isinstance(frame, torch.Tensor):
|
|
np_frame = frame.detach().cpu().numpy()
|
|
# CHW
|
|
if np_frame.ndim == 3 and np_frame.shape[0] in [
|
|
1,
|
|
3,
|
|
4,
|
|
]:
|
|
np_frame = np.transpose(np_frame, (1, 2, 0))
|
|
if np_frame.dtype != np.uint8:
|
|
np_frame = (
|
|
(np_frame * 255).astype(np.uint8)
|
|
if np_frame.max() <= 1.0
|
|
else np_frame.astype(np.uint8)
|
|
)
|
|
self.frames.append(np_frame)
|
|
else:
|
|
raise TypeError(
|
|
textwrap.dedent(f"""
|
|
Unsupported frame type: {type(frame)}.
|
|
Must be PIL.Image, numpy.ndarray, or torch.Tensor.
|
|
""")
|
|
)
|
|
|
|
self.fps = fps
|
|
self.options = options if options is not None else {}
|
|
|
|
self.audio_array = None
|
|
self.audio_samplerate = None
|
|
if audio:
|
|
if (
|
|
isinstance(audio, dict)
|
|
and "waveform" in audio
|
|
and "sample_rate" in audio
|
|
):
|
|
waveform_tensor = audio["waveform"]
|
|
self.audio_samplerate = audio["sample_rate"]
|
|
audio_np = waveform_tensor.detach().cpu().numpy()
|
|
if audio_np.ndim == 1:
|
|
audio_np = audio_np[:, np.newaxis]
|
|
elif (
|
|
audio_np.ndim > 1 and audio_np.shape[0] < audio_np.shape[1]
|
|
):
|
|
audio_np = np.transpose(audio_np)
|
|
self.audio_array = audio_np
|
|
else:
|
|
log.warning("Audio provided in an incorrect format, ignoring.")
|
|
|
|
def _save_to_disk_and_get_url(self):
|
|
"""Saves the video to a temp file and returns its URL."""
|
|
output_dir = Path(folder_paths.get_temp_directory(), "mtb_repl_videos")
|
|
output_dir.mkdir(exist_ok=True)
|
|
|
|
filename = f"vid_{time.strftime('%Y%m%d-%H%M%S')}_{uuid.uuid4()}.mp4"
|
|
filepath = output_dir / filename
|
|
|
|
try:
|
|
imageio.mimwrite(
|
|
uri=filepath,
|
|
ims=self.frames,
|
|
format="mp4",
|
|
fps=self.fps,
|
|
codec="libx264",
|
|
quality=8,
|
|
audio_array=self.audio_array,
|
|
audio_samplerate=self.audio_samplerate,
|
|
audio_codec="aac",
|
|
)
|
|
|
|
params = urllib.parse.urlencode(
|
|
{
|
|
"filename": filename,
|
|
"subfolder": "mtb_repl_videos",
|
|
"type": "temp",
|
|
}
|
|
)
|
|
|
|
return f"/view?{params}"
|
|
|
|
except Exception as e:
|
|
return f"<div style='color: red;'>Error saving video to disk: {e}</div>"
|
|
|
|
def _to_mp4_base64(self):
|
|
"""[DEPRECATED] this bloats the client"""
|
|
buffer = io.BytesIO()
|
|
try:
|
|
imageio.mimwrite(
|
|
uri=buffer,
|
|
ims=self.frames,
|
|
format="mp4",
|
|
fps=self.fps,
|
|
codec="libx264",
|
|
quality=8,
|
|
)
|
|
video_base64 = base64.b64encode(buffer.getvalue()).decode("utf-8")
|
|
return video_base64
|
|
except Exception as e:
|
|
return f"<div style='color: red;'>Error encoding video: {e}</div>"
|
|
|
|
def _repr_html_(self, *, old=False):
|
|
if old:
|
|
base64_data = self._to_mp4_base64()
|
|
if base64_data.startswith("<div"):
|
|
return base64_data
|
|
|
|
option_str = ""
|
|
for key, value in self.options.items():
|
|
if isinstance(value, bool) and value:
|
|
option_str += f" {key}"
|
|
elif isinstance(value, str):
|
|
option_str += f' {key}="{value}"'
|
|
else:
|
|
option_str += f' {key}="{value}"'
|
|
|
|
return f"""
|
|
<video
|
|
controls
|
|
src="data:video/mp4;base64,{base64_data}"
|
|
style="max-width: 100%; height: auto; border: 1px solid #555; margin: 5px 0;"{option_str}
|
|
/>"""
|
|
|
|
video_url = self._save_to_disk_and_get_url()
|
|
|
|
if video_url.startswith("<div"):
|
|
return video_url
|
|
|
|
option_str = ""
|
|
for key, value in self.options.items():
|
|
if isinstance(value, bool) and value:
|
|
option_str += f" {key}"
|
|
else:
|
|
option_str += f' {key}="{str(value)}"'
|
|
|
|
return f"""
|
|
<video
|
|
controls
|
|
src="{video_url}"
|
|
style="max-width: 100%; height: auto; border: 1px solid #555; margin: 5px 0;"{option_str}>
|
|
</video>
|
|
"""
|
|
|
|
|
|
def render_video(
|
|
frames: torch.Tensor | list[np.ndarray] | list[Image.Image] | Any,
|
|
fps=24,
|
|
options=None,
|
|
*,
|
|
audio: "AudioTensor | None" = None,
|
|
):
|
|
"""
|
|
Render video frames as an HTML video player.
|
|
|
|
Args:
|
|
frames: A list of frames, or a single batch tensor/array
|
|
(B, H, W, C) or (B, C, H, W).
|
|
fps (int): Frames per second.
|
|
options (dict): Dictionary of HTML <video> tag attributes
|
|
(e.g., {"loop": True, "autoplay": True}).
|
|
|
|
Returns
|
|
-------
|
|
VideoDisplay: An object that will render as an HTML video player.
|
|
"""
|
|
frames_list = []
|
|
if isinstance(frames, np.ndarray | torch.Tensor):
|
|
for i in range(frames.shape[0]):
|
|
frames_list.append(frames[i])
|
|
elif isinstance(frames, list):
|
|
frames_list = frames
|
|
else:
|
|
raise TypeError(
|
|
textwrap.dedent("""
|
|
Wrong input passed to render_video.
|
|
must be a list of frames or a batch tensor/array.
|
|
""")
|
|
)
|
|
|
|
return _VideoDisplay(frames_list, fps, options, audio=audio)
|
|
|
|
|
|
class _ImageDisplay:
|
|
"""Convert and display image-like objects."""
|
|
|
|
def __init__(self, img_data: Any):
|
|
self.pil_img = None
|
|
self.error = None
|
|
|
|
if isinstance(img_data, Image.Image):
|
|
self.pil_img = img_data
|
|
elif isinstance(img_data, np.ndarray):
|
|
try:
|
|
if img_data.ndim == 3 and img_data.shape[0] in [
|
|
1,
|
|
3,
|
|
4,
|
|
]: # CHW
|
|
img_data = np.transpose(img_data, (1, 2, 0))
|
|
if img_data.dtype != np.uint8:
|
|
img_data = (
|
|
(img_data * 255).astype(np.uint8)
|
|
if img_data.max() <= 1.0
|
|
else img_data.astype(np.uint8)
|
|
)
|
|
self.pil_img = Image.fromarray(img_data)
|
|
except Exception as e:
|
|
self.error = f"Error converting NumPy array to image: {e}"
|
|
elif isinstance(img_data, torch.Tensor):
|
|
try:
|
|
np_img = img_data.detach().cpu().numpy()
|
|
if np_img.ndim == 3 and np_img.shape[0] in [1, 3, 4]: # CHW
|
|
np_img = np.transpose(np_img, (1, 2, 0))
|
|
if np_img.dtype != np.uint8:
|
|
np_img = (
|
|
(np_img * 255).astype(np.uint8)
|
|
if np_img.max() <= 1.0
|
|
else np_img.astype(np.uint8)
|
|
)
|
|
self.pil_img = Image.fromarray(np_img)
|
|
except Exception as e:
|
|
self.error = f"Error converting torch.Tensor to image: {e}"
|
|
else:
|
|
self.error = (
|
|
f"Unsupported image type for display: {type(img_data)}"
|
|
)
|
|
|
|
def _repr_html_(self) -> str:
|
|
"""Convert the internal PIL image to a base64 HTML string."""
|
|
if self.error:
|
|
return f"<div style='color: red;'>{self.error}</div>"
|
|
if not self.pil_img:
|
|
return (
|
|
"<div style='color: red;'>Could not process image data.</div>"
|
|
)
|
|
|
|
buffer = io.BytesIO()
|
|
try:
|
|
self.pil_img.save(buffer, format="PNG")
|
|
img_base64 = base64.b64encode(buffer.getvalue()).decode("utf-8")
|
|
return f'<img src="data:image/png;base64,{img_base64}" style="max-width: 100%; height: auto; border: 1px solid #555; margin: 5px 0;"/>'
|
|
except Exception as e:
|
|
return f"<div style='color: red;'>Error saving image: {e}</div>"
|
|
|
|
|
|
def render_image(img_data: Any):
|
|
"""
|
|
Render image-like data (PIL, NumPy, Torch) as an HTML image.
|
|
|
|
This is a user-facing function for explicit display calls.
|
|
"""
|
|
return _ImageDisplay(img_data)
|
|
|
|
|
|
# endregion
|
|
|
|
|
|
class SafeInputList(list):
|
|
"""Wrapper of list for custom error message."""
|
|
|
|
def __getitem__(self, key):
|
|
"""Get item with error rewired."""
|
|
try:
|
|
return super().__getitem__(key)
|
|
except IndexError as e:
|
|
raise IndexError(
|
|
textwrap.dedent("""
|
|
Input index out of range.
|
|
In live mode, the 'inputs' list is empty,
|
|
You can use the guard 'if IS_LIVE:'
|
|
to fill them only for the live mode""")
|
|
) from e
|
|
|
|
|
|
class ComfyReplAPI:
|
|
"""'API' available from the REPL instances."""
|
|
|
|
def __init__(self):
|
|
self.__color = ""
|
|
pass
|
|
|
|
def set_color(self, color: str):
|
|
self.__color = color
|
|
|
|
# def get_input(self, at: int, mock: Any) -> Any:
|
|
# pass
|
|
|
|
def get_gradient_2d(self, start, stop, width, height, is_horizontal):
|
|
if is_horizontal:
|
|
return np.tile(np.linspace(start, stop, width), (height, 1))
|
|
else:
|
|
return np.tile(np.linspace(start, stop, height), (width, 1)).T
|
|
|
|
def get_gradient_3d(
|
|
self,
|
|
width,
|
|
height,
|
|
start_list,
|
|
stop_list,
|
|
is_horizontal_list,
|
|
*,
|
|
as_image=False,
|
|
):
|
|
result = np.zeros((height, width, len(start_list)), dtype=np.float32)
|
|
|
|
for i, (start, stop, is_horizontal) in enumerate(
|
|
zip(start_list, stop_list, is_horizontal_list, strict=False)
|
|
):
|
|
result[:, :, i] = self.get_gradient_2d(
|
|
start, stop, width, height, is_horizontal
|
|
)
|
|
if as_image:
|
|
return Image.fromarray(np.uint8(result))
|
|
else:
|
|
return result
|
|
|
|
|
|
@singleton
|
|
class ComfyREPLBackend:
|
|
"""Singleton "backend" for the REPL editor."""
|
|
|
|
def __init__(self):
|
|
self._repl_consoles: dict[str, code.InteractiveConsole] = {}
|
|
|
|
self.media_outputs = []
|
|
|
|
# self._original_displayhook = sys.displayhook
|
|
self.handlers = []
|
|
self.register_handlers()
|
|
|
|
def register_handlers(self):
|
|
self.handlers.extend(
|
|
[
|
|
self._handle_video,
|
|
self._handle_audio,
|
|
self._handle_image,
|
|
]
|
|
)
|
|
|
|
def _handle_image(self, value):
|
|
"""Handle PIL Images, and Numpy/Torch tensors that represent images."""
|
|
if isinstance(value, _ImageDisplay):
|
|
return value._repr_html_()
|
|
|
|
if isinstance(value, Image.Image | np.ndarray | torch.Tensor):
|
|
return _ImageDisplay(value)._repr_html_()
|
|
|
|
return None
|
|
|
|
def _handle_audio(self, value):
|
|
"""Handle AudioDisplay objects."""
|
|
if isinstance(value, _AudioDisplay):
|
|
return value._repr_html_()
|
|
return None
|
|
|
|
def _handle_video(self, value):
|
|
"""Handle VideoDisplay objects."""
|
|
if isinstance(value, _VideoDisplay):
|
|
return value._repr_html_()
|
|
|
|
if isinstance(value, torch.Tensor) and len(value.shape) == 4:
|
|
return render_video(value)._repr_html_()
|
|
|
|
if (
|
|
isinstance(value, list)
|
|
and len(value) > 0
|
|
and isinstance(value, Image.Image | np.ndarray | torch.Tensor)
|
|
):
|
|
return render_video(value)._repr_html_()
|
|
return None
|
|
|
|
def _init_repl_console(self, *inputs):
|
|
"""Define the globals that will be available in the REPL session."""
|
|
repl_locals = {"__builtins__": __builtins__}
|
|
# repl_globals["plt"] = plt
|
|
repl_locals["np"] = np
|
|
repl_locals["Image"] = Image
|
|
repl_locals["torch"] = torch
|
|
repl_api = ComfyReplAPI()
|
|
repl_locals["repl"] = repl_api
|
|
|
|
# "globals"
|
|
repl_locals["IS_LIVE"] = True
|
|
|
|
filename = Path(folder_paths.get_input_directory()) / "example.png"
|
|
log.debug(
|
|
f"Looking for example in {filename.as_posix()}, {filename.exists()}"
|
|
)
|
|
|
|
if filename.exists():
|
|
img = Image.open(filename)
|
|
repl_locals["EXAMPLE"] = torch.from_numpy(
|
|
np.array(img).astype(np.float32) / 255.0
|
|
).unsqueeze(0)
|
|
else:
|
|
repl_locals["EXAMPLE"] = torch.from_numpy(
|
|
repl_api.get_gradient_3d(
|
|
1024,
|
|
1024,
|
|
(0, 0, 192),
|
|
(255, 255, 64),
|
|
(True, False, False),
|
|
)
|
|
).unsqueeze(0)
|
|
|
|
repl_locals["inputs"] = [None] * SOCKET_COUNT
|
|
repl_locals["outputs"] = [None] * SOCKET_COUNT
|
|
|
|
if len(inputs):
|
|
repl_locals["inputs"] = SafeInputList(inputs)
|
|
|
|
# internals but exposed for debug
|
|
if _HAS_SCIPY:
|
|
repl_locals["render_audio"] = render_audio
|
|
if _HAS_IMAGEIO:
|
|
repl_locals["render_video"] = render_video
|
|
|
|
repl_locals["render_image"] = render_image
|
|
|
|
return code.InteractiveConsole(locals=repl_locals)
|
|
|
|
def _custom_displayhook(self, value):
|
|
"""
|
|
Displayhook that capture and process image, audio, video objects.
|
|
|
|
It can also look for renderable items inside lists, tuples, and dicts.
|
|
For other objects, it fallsback to the original displayhook.
|
|
"""
|
|
log.debug("Custom displayhook called")
|
|
if value is None:
|
|
return
|
|
|
|
# handle with handler
|
|
if self._handle_item(value):
|
|
return
|
|
|
|
# containers
|
|
items_to_process = []
|
|
if isinstance(value, list | tuple):
|
|
items_to_process.extend(value)
|
|
elif isinstance(value, dict):
|
|
items_to_process.extend(value.values())
|
|
|
|
if items_to_process:
|
|
was_anything_rendered = False
|
|
for item in items_to_process:
|
|
if self._handle_item(item):
|
|
was_anything_rendered = True
|
|
# Item handled, move to the next item in the container
|
|
break
|
|
|
|
# If we found any media inside the container, we consider the
|
|
# displayhook's job done. We don't also print the container itself.
|
|
if was_anything_rendered:
|
|
return
|
|
|
|
log.debug("Falling back to original display hook")
|
|
# self._original_displayhook(value)
|
|
sys.displayhook(value)
|
|
|
|
def _handle_item(self, item):
|
|
# handle with handler
|
|
for handler in self.handlers:
|
|
html_output = handler(item)
|
|
if html_output:
|
|
log.debug(f"Object handled by {handler.__name__}")
|
|
self.media_outputs.append(html_output)
|
|
return True
|
|
return False
|
|
|
|
def _console_to_html(
|
|
self, stream: io.StringIO | Traceback, width: int = 120
|
|
) -> str:
|
|
if isinstance(stream, io.StringIO):
|
|
captured_text_output = stream.getvalue()
|
|
else:
|
|
log.debug(
|
|
f"Steam is not an io.StringIO but a {type(stream).__name__}"
|
|
)
|
|
captured_text_output = stream # Traceback
|
|
|
|
html_console = Console(
|
|
file=io.StringIO(), record=True, force_terminal=True, width=width
|
|
)
|
|
html_console.print(captured_text_output)
|
|
|
|
return html_console.export_html(inline_styles=True)
|
|
|
|
def get_console(self, name: str | None = None, *, reset=False):
|
|
"""Get or create a new named console."""
|
|
if name is None:
|
|
log.error("Cannot get console without providing a name (uuid)")
|
|
return
|
|
|
|
console = self._repl_consoles.get(name)
|
|
|
|
if console and not reset:
|
|
return console
|
|
|
|
if reset:
|
|
if console is not None:
|
|
del self._repl_consoles[name]
|
|
console = None
|
|
|
|
console = self._init_repl_console()
|
|
self._repl_consoles[name] = console
|
|
return self._repl_consoles[name]
|
|
|
|
log.debug(
|
|
textwrap.dedent(f"""
|
|
Creating new console for node {name}
|
|
(active: {len(self._repl_consoles.keys())})
|
|
""")
|
|
)
|
|
console = self._init_repl_console()
|
|
self._repl_consoles[name] = console
|
|
return self._repl_consoles[name]
|
|
|
|
def get_outputs(
|
|
self,
|
|
*,
|
|
name: str | None = None,
|
|
console: code.InteractiveConsole | None = None,
|
|
):
|
|
if name is None and console is None:
|
|
raise ValueError("Either id or console must be provided")
|
|
|
|
if name is not None:
|
|
console = self.get_console(name)
|
|
|
|
if console is None:
|
|
raise RuntimeError("No console found")
|
|
|
|
if "outputs" in console.locals:
|
|
outputs = console.locals["outputs"]
|
|
# guards
|
|
if outputs is None:
|
|
outputs = [None] * SOCKET_COUNT
|
|
elif isinstance(outputs, list) and len(outputs) > SOCKET_COUNT:
|
|
raise RuntimeError(
|
|
textwrap.dedent(f"""
|
|
Too many outputs: {len(outputs)} vs {SOCKET_COUNT}
|
|
Outputs can be smaller then {SOCKET_COUNT} but not higher.
|
|
""")
|
|
)
|
|
|
|
# enlist if needed
|
|
if not isinstance(outputs, list):
|
|
outputs = [outputs]
|
|
|
|
# pad the outputs
|
|
if len(outputs) < SOCKET_COUNT:
|
|
outputs = outputs + [None] * (SOCKET_COUNT - len(outputs))
|
|
else:
|
|
outputs = [None] * SOCKET_COUNT
|
|
|
|
return outputs
|
|
|
|
def execute_code(
|
|
self,
|
|
code: str,
|
|
*,
|
|
name: str | None = None,
|
|
console=None,
|
|
reset=False,
|
|
inputs: list[Any] | None = None,
|
|
):
|
|
log.debug(f"Executing code for node {name}")
|
|
|
|
inputs = inputs or []
|
|
|
|
if not console and not name:
|
|
raise ValueError("Either name or console must be provided")
|
|
if console is not None and reset:
|
|
raise ValueError(
|
|
textwrap.dedent("""
|
|
Cannot specify both console and reset,
|
|
use name and reset instead
|
|
""")
|
|
)
|
|
|
|
self.media_outputs = []
|
|
output_html = ""
|
|
error_message = None
|
|
|
|
repl_console = console or self.get_console(name, reset=reset)
|
|
if not repl_console:
|
|
log.error(f"Failed to get console named {name}")
|
|
return
|
|
|
|
string_io = io.StringIO()
|
|
repl_console.locals["inputs"] = SafeInputList(inputs)
|
|
|
|
# Temporarily patch sys.displayhook
|
|
# sys.displayhook = self._custom_displayhook
|
|
|
|
try:
|
|
log.debug(f"CODE SENT:\n\n{code}")
|
|
with redirect_stdout(string_io), redirect_stderr(string_io):
|
|
try:
|
|
tree = ast.parse(code)
|
|
except SyntaxError as e:
|
|
raise e
|
|
|
|
if tree.body:
|
|
# split if last statement is an expr
|
|
if isinstance(tree.body[-1], ast.Expr):
|
|
setup_module = ast.Module(
|
|
body=tree.body[:-1], type_ignores=[]
|
|
)
|
|
setup_code = compile(
|
|
setup_module, "<mtb-repl-setup>", "exec"
|
|
)
|
|
|
|
# the last expr
|
|
last_expr = ast.Expression(body=tree.body[-1].value)
|
|
last_code = compile(
|
|
last_expr, "<mtb-repl-eval>", "eval"
|
|
)
|
|
|
|
exec(setup_code, repl_console.locals)
|
|
result = eval(last_code, repl_console.locals)
|
|
|
|
# required as runsource doesn't trigger displayhook
|
|
self._custom_displayhook(result)
|
|
|
|
if "outputs" in repl_console.locals and isinstance(
|
|
repl_console.locals["outputs"], list | tuple
|
|
):
|
|
log.warning(
|
|
textwrap.dedent("""
|
|
Outputs are defined
|
|
but the display hook will overwrite them""")
|
|
)
|
|
repl_console.locals["outputs"] = result
|
|
|
|
# no expression as last statement
|
|
else:
|
|
exec(code, repl_console.locals)
|
|
|
|
# NOTE: old
|
|
# repl_console.runsource(code, "mtb-repl", "exec")
|
|
|
|
|
|
full_rich_html = self._console_to_html(string_io)
|
|
|
|
match = re.search(
|
|
r"<body.*?>(.*?)</body>", full_rich_html, re.DOTALL
|
|
)
|
|
output_html = match.group(1) if match else full_rich_html
|
|
|
|
output_html += "".join(self.media_outputs)
|
|
|
|
except Exception as e:
|
|
exc_type, exc_value, exc_traceback = sys.exc_info()
|
|
rich_traceback = Traceback.from_exception(
|
|
exc_type,
|
|
exc_value,
|
|
exc_traceback,
|
|
show_locals=True,
|
|
suppress=[__file__],
|
|
)
|
|
# error_console = Console(
|
|
# file=io.StringIO(), record=True, force_terminal=True, width=120
|
|
# )
|
|
# error_console.print(rich_traceback)
|
|
# full_error_html = error_console.export_html(inline_styles=True)
|
|
|
|
full_error_html = self._console_to_html(rich_traceback)
|
|
|
|
match = re.search(
|
|
r"<body.*?>(.*?)</body>", full_error_html, re.DOTALL
|
|
)
|
|
output_html = match.group(1) if match else full_error_html
|
|
|
|
error_message = str(e)
|
|
# finally:
|
|
# sys.displayhook = (
|
|
# self._original_displayhook
|
|
# ) # Always restore original displayhook
|
|
|
|
return {"output_html": output_html, "error": error_message}
|
|
|
|
@property
|
|
def active_consoles(self): # noqa: N805
|
|
return [str(k) for k in self._repl_consoles]
|
|
|
|
def lint_code(self, *, name: str = "", code: str):
|
|
log.debug(f"Linting code for node {name}")
|
|
diagnostics = []
|
|
|
|
if not _HAS_LINT:
|
|
diagnostics.append(
|
|
{
|
|
"row": 0,
|
|
"column": 0,
|
|
"text": textwrap.dedent("""
|
|
Pyflakes not installed.
|
|
Linting disabled.
|
|
Install with 'uv add pyflakes'.
|
|
"""),
|
|
"type": "warning",
|
|
}
|
|
)
|
|
return web.json_response({"diagnostics": diagnostics})
|
|
|
|
custom_globals = {"render_video", "render_audio", "repl_display"}
|
|
|
|
# Use a custom reporter to capture messages
|
|
class PyflakesReporter(pyflakes.reporter.Reporter):
|
|
def __init__(self):
|
|
self.messages = []
|
|
# Suppress stdout/stderr from pyflakes itself
|
|
self._stdout = io.StringIO()
|
|
self._stderr = io.StringIO()
|
|
super().__init__(self._stdout, self._stderr)
|
|
|
|
def flake(self, message):
|
|
# Ace editor expects 0-indexed row
|
|
# pyflakes gives 1-indexed lineno
|
|
|
|
import pyflakes.messages
|
|
|
|
kind = "warning"
|
|
|
|
if isinstance(message, pyflakes.messages.UndefinedName):
|
|
kind = "error"
|
|
|
|
if (
|
|
message.message_args
|
|
and message.message_args[0] in custom_globals
|
|
):
|
|
return
|
|
|
|
log.info("result from flake")
|
|
console = rich.console.Console(stderr=True)
|
|
rich.inspect(message, console=console)
|
|
|
|
self.messages.append(
|
|
{
|
|
"row": message.lineno - 1,
|
|
"column": message.col,
|
|
"text": str(message) + "prout",
|
|
"type": kind,
|
|
}
|
|
)
|
|
|
|
def unexpectedError(self, filename, msg):
|
|
self.messages.append(
|
|
{
|
|
"row": 0,
|
|
"column": 0,
|
|
"text": f"Pyflakes internal error: {msg}",
|
|
"type": "error",
|
|
}
|
|
)
|
|
|
|
def syntaxError(self, filename, msg, lineno, offset, text):
|
|
log.info(f"Received {text} to syntax error")
|
|
self.messages.append(
|
|
{
|
|
"row": lineno - 1, # Ace is 0-indexed
|
|
"column": offset,
|
|
"text": f"Syntax Error: {msg}",
|
|
"type": "error",
|
|
}
|
|
)
|
|
|
|
reporter = PyflakesReporter()
|
|
pyflakes.api.check(code, name, reporter)
|
|
|
|
return {"diagnostics": reporter.messages}
|
|
|
|
def lint_code_ruff(self, code: str):
|
|
diagnostics = []
|
|
|
|
if not _HAS_LINT:
|
|
diagnostics.append(
|
|
{
|
|
"row": 0,
|
|
"column": 0,
|
|
"text": textwrap.dedent("""
|
|
Ruff not installed.
|
|
Linting disabled.
|
|
Install with 'pip install ruff'.
|
|
"""),
|
|
"type": "warning",
|
|
}
|
|
)
|
|
return {"diagnostics": diagnostics}
|
|
|
|
# Define the builtins/globals that Ruff should recognize
|
|
# These are the names we inject into the REPL's scope
|
|
repl_builtins = [
|
|
"repl_display",
|
|
"render_audio",
|
|
"render_video",
|
|
"plt",
|
|
"np",
|
|
"Image",
|
|
"torch",
|
|
]
|
|
|
|
try:
|
|
# Lint the code using Ruff's programmatic API
|
|
result = ruff.lint.linter.lint_stdin(
|
|
code.encode("utf-8"),
|
|
path="<stdin>",
|
|
builtins=repl_builtins,
|
|
)
|
|
|
|
for diagnostic in result.diagnostics:
|
|
diag_type = "warning" # Default
|
|
# Ruff's error codes:
|
|
# F (Pyflakes)
|
|
# E (Pycodestyle)
|
|
# W (Pycodestyle warning)
|
|
# I (isort)
|
|
# N (naming)
|
|
# ...
|
|
# F821: Undefined name
|
|
if (
|
|
diagnostic.kind.code.startswith("E")
|
|
or diagnostic.kind.code == "F821"
|
|
):
|
|
diag_type = "error"
|
|
elif diagnostic.kind.code.startswith("W"):
|
|
diag_type = "warning"
|
|
|
|
diagnostics.append(
|
|
{
|
|
"row": diagnostic.location.row - 1, # Ace is 0-indexed
|
|
"column": diagnostic.location.column
|
|
- 1, # Ace is 0-indexed
|
|
"text": diagnostic.message,
|
|
"type": diag_type,
|
|
}
|
|
)
|
|
|
|
except Exception as e:
|
|
diagnostics.append(
|
|
{
|
|
"row": 0,
|
|
"column": 0,
|
|
"text": f"Ruff internal error: {e}",
|
|
"type": "error",
|
|
}
|
|
)
|
|
|
|
return {"diagnostics": diagnostics}
|
|
|
|
|
|
_comfy_repl_backend = ComfyREPLBackend()
|
|
|
|
|
|
# region endpoint handlers
|
|
async def repl_execute_code_handler(request):
|
|
data = await request.json()
|
|
|
|
name = data.get("name")
|
|
reset = data.get("reset")
|
|
|
|
if reset is None:
|
|
reset = False
|
|
|
|
if name is None:
|
|
return web.Response(
|
|
status=417, reason="Expectation Failed, missing name key"
|
|
)
|
|
|
|
code = data.get("code", "")
|
|
result = _comfy_repl_backend.execute_code(code, name=name, reset=reset)
|
|
|
|
return web.json_response(result)
|
|
|
|
|
|
async def repl_lint_code_handler(request):
|
|
data = await request.json()
|
|
name = data.get("name")
|
|
|
|
if name is None: # we send an error
|
|
return web.Response(
|
|
status=417, reason="Expectation Failed", text="Missing name key"
|
|
)
|
|
# raise web.HTTPExpectationFailed(
|
|
# reason="Missing name key (reason)", text="Missing name key (text)"
|
|
# )
|
|
|
|
code = data.get("code", "")
|
|
result = _comfy_repl_backend.lint_code(name=name, code=code)
|
|
return web.json_response(result)
|
|
|
|
|
|
def setup_custom_web_routes(app: web.Application):
|
|
"""Register REPL routes."""
|
|
log.info("ComfyREPL: Registering repl routes...")
|
|
app.router.add_post("/mtb/execute", repl_execute_code_handler)
|
|
app.router.add_post("/mtb/lint", repl_lint_code_handler)
|
|
|
|
|
|
# endregion
|