From edd76e341cbd0eefae5c0f70ceaa031fd15a9b74 Mon Sep 17 00:00:00 2001 From: Wouter Verweirder Date: Thu, 18 Sep 2025 16:06:57 +0200 Subject: [PATCH] cleanup on cancel prompt --- requirements.txt | 3 + src/penplotter/nodes.py | 201 +++++++++++++++++++++++++++++++--------- 2 files changed, 162 insertions(+), 42 deletions(-) diff --git a/requirements.txt b/requirements.txt index 477e903..f24b8c5 100644 --- a/requirements.txt +++ b/requirements.txt @@ -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 \ No newline at end of file diff --git a/src/penplotter/nodes.py b/src/penplotter/nodes.py index d2e4025..36b4755 100644 --- a/src/penplotter/nodes.py +++ b/src/penplotter/nodes.py @@ -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 \ No newline at end of file