cleanup on cancel prompt

This commit is contained in:
Wouter Verweirder
2025-09-18 16:06:57 +02:00
parent c5243865d6
commit edd76e341c
2 changed files with 162 additions and 42 deletions
+3
View File
@@ -5,5 +5,8 @@ Pillow>=9.0.0
# Vector drawing and plotting library
vpype>=1.0.0
# WebSocket client for ComfyUI integration
websockets>=11.0.0
# AxiDraw API for pen plotter control
https://cdn.evilmadscientist.com/dl/ad/public/AxiDraw_API.zip
+159 -42
View File
@@ -3,12 +3,12 @@ import numpy as np
import tempfile
import subprocess
from pathlib import Path
from server import PromptServer
from PIL import Image
import comfy.model_management
try:
from PIL import Image
except ImportError:
print("Warning: PIL (Pillow) not found. ImageToCenterline node will not work without it.")
Image = None
# Global variable to track the currently running subprocess for interruption
_current_subprocess = None
class ImageToCenterline:
CATEGORY = "Pen Plotter"
@@ -70,6 +70,7 @@ class ImageToCenterline:
def run_raster_retrace(self, ppm_path, svg_path, optimize_exhaustive=True):
"""Run the raster-retrace command on a PPM file to generate SVG."""
global _current_subprocess
try:
# Build command with proper flags
cmd = [
@@ -82,18 +83,23 @@ class ImageToCenterline:
if optimize_exhaustive:
cmd.append('--optimize-exhaustive')
result = subprocess.run(cmd, capture_output=True, text=True)
_current_subprocess = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
stdout, stderr = _current_subprocess.communicate()
if result.returncode == 0:
if _current_subprocess.returncode == 0:
_current_subprocess = None
return True
else:
print(f"Error running raster-retrace: {result.stderr}")
print(f"Error running raster-retrace: {stderr}")
_current_subprocess = None
return False
except FileNotFoundError:
print("Error: raster-retrace command not found. Please ensure it's installed and in your PATH.")
_current_subprocess = None
return False
except Exception as e:
print(f"Error running raster-retrace: {e}")
_current_subprocess = None
return False
def convert_to_centerline(self, image, optimize_exhaustive=True):
@@ -160,6 +166,7 @@ class OptimizeSVG:
def run_vpype_command(self, input_path, output_path, fit_to_margins, landscape, linemerge,
tolerance, linesort, reloop, linesimplify, page_width, page_height):
"""Run vpype command with specified parameters."""
global _current_subprocess
try:
# Build the vpype command
cmd = ['vpype', 'read', str(input_path)]
@@ -189,22 +196,30 @@ class OptimizeSVG:
cmd.append('linesimplify')
# Add write command with page size
cmd.extend(['write', '--page-size', f'{page_width}x{page_height}mm', str(output_path)])
if landscape:
cmd.extend(['write', '--page-size', f'{page_height}x{page_width}mm', str(output_path)])
else:
cmd.extend(['write', '--page-size', f'{page_width}x{page_height}mm', str(output_path)])
# Execute the command
result = subprocess.run(cmd, capture_output=True, text=True)
# Execute the command and store subprocess reference
_current_subprocess = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
stdout, stderr = _current_subprocess.communicate()
if result.returncode == 0:
if _current_subprocess.returncode == 0:
_current_subprocess = None
return True
else:
print(f"Error running vpype: {result.stderr}")
print(f"Error running vpype: {stderr}")
_current_subprocess = None
return False
except FileNotFoundError:
print("Error: vpype command not found. Please ensure vpype is installed and in your PATH.")
_current_subprocess = None
return False
except Exception as e:
print(f"Error running vpype: {e}")
_current_subprocess = None
return False
def optimize_svg(self, svg_string, fit_to_margins=0.0, landscape=True, linemerge=True,
@@ -253,22 +268,22 @@ class PlotSVG:
},
"optional": {
"layer": ("INT", {"default": 2, "min": 1, "max": 8, "step": 1}),
"preview_only": ("BOOLEAN", {"default": False}),
"pen_up_speed": ("INT", {"default": 75, "min": 1, "max": 100, "step": 1}),
"pen_down_speed": ("INT", {"default": 25, "min": 1, "max": 100, "step": 1}),
"pen_up_delay": ("INT", {"default": 0, "min": 0, "max": 5000, "step": 10}),
"pen_down_delay": ("INT", {"default": 0, "min": 0, "max": 5000, "step": 10}),
"copies": ("INT", {"default": 1, "min": 1, "max": 100, "step": 1}),
}
}
RETURN_TYPES = ("STRING",)
RETURN_NAMES = ("status",)
FUNCTION = "plot_svg"
OUTPUT_NODE = True
def run_axicli_command(self, svg_path, layer, preview_only, pen_up_speed, pen_down_speed,
pen_up_delay, pen_down_delay, copies):
def run_axicli_plot(self, svg_path, layer, pen_up_speed, pen_down_speed,
pen_up_delay, pen_down_delay):
"""Run axicli command with specified parameters."""
global _current_subprocess
try:
# Build the axicli command
cmd = ['axicli', str(svg_path)]
@@ -276,10 +291,6 @@ class PlotSVG:
# Add layer parameter
cmd.extend(['-L', str(layer)])
# Add preview mode if enabled
if preview_only:
cmd.append('-m') # Preview mode
# Add speed settings
cmd.extend(['-S', str(pen_up_speed)]) # Pen-up speed
cmd.extend(['-s', str(pen_down_speed)]) # Pen-down speed
@@ -290,30 +301,31 @@ class PlotSVG:
if pen_down_delay > 0:
cmd.extend(['-D', str(pen_down_delay)]) # Pen-down delay
# Add copies
if copies > 1:
cmd.extend(['-C', str(copies)])
# Execute the command and store subprocess reference
_current_subprocess = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
stdout, stderr = _current_subprocess.communicate()
# Execute the command
result = subprocess.run(cmd, capture_output=True, text=True)
if result.returncode == 0:
return True, result.stdout
if _current_subprocess.returncode == 0:
_current_subprocess = None
return True, stdout
else:
print(f"Error running axicli: {result.stderr}")
return False, result.stderr
print(f"Error running axicli: {stderr}")
_current_subprocess = None
return False, stderr
except FileNotFoundError:
error_msg = "Error: axicli command not found. Please ensure AxiDraw software is installed and axicli is in your PATH."
print(error_msg)
_current_subprocess = None
return False, error_msg
except Exception as e:
error_msg = f"Error running axicli: {e}"
print(error_msg)
_current_subprocess = None
return False, error_msg
def plot_svg(self, svg_string, layer=2, preview_only=False, pen_up_speed=75, pen_down_speed=25,
pen_up_delay=0, pen_down_delay=0, copies=1):
def plot_svg(self, svg_string, layer=2, pen_up_speed=75, pen_down_speed=25,
pen_up_delay=0, pen_down_delay=0):
"""Send SVG to plotter using axicli."""
try:
# Create temporary file for the SVG
@@ -326,34 +338,139 @@ class PlotSVG:
f.write(svg_string)
# Run axicli command
success, output = self.run_axicli_command(
svg_path, layer, preview_only, pen_up_speed, pen_down_speed,
pen_up_delay, pen_down_delay, copies
success, output = self.run_axicli_plot(
svg_path, layer, pen_up_speed, pen_down_speed,
pen_up_delay, pen_down_delay
)
if success:
if preview_only:
return ("Preview completed successfully. Check AxiDraw software for preview.",)
else:
return (f"Plot completed successfully. Copies: {copies}",)
run_axicli_disengage() # Ensure plotter is disengaged after plotting
return (f"Plot completed successfully.",)
else:
return (f"Error: {output}",)
except Exception as e:
return (f"Error: {str(e)}",)
class DisengagePlotter:
CATEGORY = "Pen Plotter"
@classmethod
def INPUT_TYPES(s):
return {
"required": {
}
}
RETURN_TYPES = ("STRING",)
RETURN_NAMES = ("status",)
FUNCTION = "disengage_plotter"
OUTPUT_NODE = True
def run_axicli_disengage(self):
global _current_subprocess
try:
# Build the axicli command
cmd = ['axicli', '--mode', 'align']
# Execute the command and store subprocess reference
_current_subprocess = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
stdout, stderr = _current_subprocess.communicate()
if _current_subprocess.returncode == 0:
_current_subprocess = None
return True, stdout
else:
print(f"Error running axicli: {stderr}")
_current_subprocess = None
return False, stderr
except FileNotFoundError:
error_msg = "Error: axicli command not found. Please ensure AxiDraw software is installed and axicli is in your PATH."
print(error_msg)
_current_subprocess = None
return False, error_msg
except Exception as e:
error_msg = f"Error running axicli: {e}"
print(error_msg)
_current_subprocess = None
return False, error_msg
def disengage_plotter(self):
"""Disengage the plotter using axicli."""
try:
success, output = self.run_axicli_disengage()
if success:
return ("Plotter disengaged successfully.",)
else:
return (f"Error: {output}",)
except Exception as e:
return (f"Error: {str(e)}",)
def run_axicli_disengage():
"""Function to disengage the plotter using axicli."""
try:
cmd = ['axicli', '--mode', 'align']
result = subprocess.run(cmd, capture_output=True, text=True)
if result.returncode == 0:
print("Plotter disengaged successfully.")
return True
else:
print(f"Error disengaging plotter: {result.stderr}")
return False
except FileNotFoundError:
print("Error: axicli command not found. Please ensure AxiDraw software is installed and axicli is in your PATH.")
return False
except Exception as e:
print(f"Error disengaging plotter: {e}")
return False
# A dictionary that contains all nodes you want to export with their names
# NOTE: names should be globally unique
NODE_CLASS_MAPPINGS = {
"ImageToCenterline": ImageToCenterline,
"OptimizeSVG": OptimizeSVG,
"PlotSVG": PlotSVG
"PlotSVG": PlotSVG,
"DisengagePlotter": DisengagePlotter
}
# A dictionary that contains the friendly/humanly readable titles for the nodes
NODE_DISPLAY_NAME_MAPPINGS = {
"ImageToCenterline": "Image to Centerline SVG",
"OptimizeSVG": "Optimize SVG with vpype",
"PlotSVG": "Plot SVG with AxiDraw"
"PlotSVG": "Plot SVG with AxiDraw",
"DisengagePlotter": "Disengage Plotter"
}
# Overwrite or wrap comfy.model_management.interrupt_current_processing
_original_interrupt = comfy.model_management.interrupt_current_processing
def wrapped_interrupt_current_processing(*args, **kwargs):
global _current_subprocess
result = _original_interrupt(*args, **kwargs)
if args and args[0] is True:
print("[PenPlotter] Processing interrupted.")
# Interrupt the running subprocess if needed
if _current_subprocess is not None:
try:
print("[PenPlotter] Terminating running subprocess...")
_current_subprocess.terminate()
# Give the process a moment to terminate gracefully
try:
_current_subprocess.wait(timeout=2)
except subprocess.TimeoutExpired:
# If it doesn't terminate gracefully, force kill it
print("[PenPlotter] Force killing subprocess...")
_current_subprocess.kill()
_current_subprocess.wait()
print("[PenPlotter] Subprocess terminated.")
_current_subprocess = None
except Exception as e:
print(f"[PenPlotter] Error terminating subprocess: {e}")
_current_subprocess = None
# run the disengage command to put the plotter in a safe state
run_axicli_disengage()
return result
comfy.model_management.interrupt_current_processing = wrapped_interrupt_current_processing