diff --git a/__init__.py b/__init__.py index 40d3fe2..0ca6551 100644 --- a/__init__.py +++ b/__init__.py @@ -19,12 +19,64 @@ NODE_DISPLAY_NAME_MAPPINGS.update(IMAGE_NODE_DISPLAY_NAME_MAPPINGS) WEB_DIRECTORY = "./js" try: + import asyncio import os + import shutil import subprocess import sys + import threading from aiohttp import web from server import PromptServer + _DUMAS_DIALOG_LOCK = threading.Lock() + + def _dumas_dialog_available(): + return sys.platform == "win32" and shutil.which("powershell") is not None + + def _dumas_dialog_windows(start_path): + ps = ( + "Add-Type -AssemblyName System.Windows.Forms;" + "$r='';" + "$o=New-Object System.Windows.Forms.Form;" + "$o.TopMost=$true;$o.ShowInTaskbar=$false;$o.FormBorderStyle='None';" + "$o.Width=1;$o.Height=1;$o.Opacity=0;$o.StartPosition='CenterScreen';" + "$o.Add_Shown({" + "$o.Activate();" + "$d=New-Object System.Windows.Forms.FolderBrowserDialog;" + '$d.Description="Select a folder";$d.ShowNewFolderButton=$true;' + "if($env:DUMAS_START){try{$d.SelectedPath=$env:DUMAS_START}catch{}};" + "if($d.ShowDialog($o) -eq [System.Windows.Forms.DialogResult]::OK){$script:r=$d.SelectedPath};" + "$o.Close()" + "});" + "[void]$o.ShowDialog();" + "[Console]::Out.Write($r)" + ) + env = dict(os.environ) + env["DUMAS_START"] = start_path or "" + result = subprocess.run( + ["powershell", "-NoProfile", "-STA", "-Command", ps], + capture_output=True, + text=True, + timeout=300, + env=env, + creationflags=0x08000000, + check=False, + ) + return (result.stdout or "").strip() + + def _dumas_native_folder_dialog(start_path=""): + if not _DUMAS_DIALOG_LOCK.acquire(blocking=False): + return None + try: + if sys.platform == "win32": + return _dumas_dialog_windows(start_path) + return "" + finally: + try: + _DUMAS_DIALOG_LOCK.release() + except Exception: + pass + @PromptServer.instance.routes.get("/dumas/api/save_image/file") async def dumas_save_image_file(request): path = resolve_serve_token(request.query.get("t", "")) @@ -35,48 +87,17 @@ try: @PromptServer.instance.routes.get("/dumas/api/pick_directory") async def dumas_pick_directory(request): try: - if sys.platform != "win32": - return web.json_response({"ok": False, "message": "Native folder picker is only available on Windows right now."}) + if not _dumas_dialog_available(): + return web.json_response({"ok": False, "message": "Native folder picker is unavailable here.", "unavailable": True}) start_path = str(request.query.get("path", "") or "").strip() if start_path and not os.path.isdir(start_path): start_path = "" - ps_lines = [ - "Add-Type -AssemblyName System.Windows.Forms", - "$dialog = New-Object System.Windows.Forms.FolderBrowserDialog", - '$dialog.Description = "Select a folder"', - "$dialog.ShowNewFolderButton = $true", - ] - if start_path: - escaped = start_path.replace("'", "''") - ps_lines.append(f"$dialog.SelectedPath = '{escaped}'") - ps_lines.extend( - [ - "if ($dialog.ShowDialog() -eq [System.Windows.Forms.DialogResult]::OK) {", - " [Console]::OutputEncoding = [System.Text.Encoding]::UTF8", - " Write-Output $dialog.SelectedPath", - "}", - ] - ) - command = "; ".join(ps_lines) - result = subprocess.run( - [ - "powershell", - "-NoProfile", - "-STA", - "-ExecutionPolicy", - "Bypass", - "-Command", - command, - ], - capture_output=True, - text=True, - encoding="utf-8", - errors="replace", - check=False, - ) - selected = (result.stdout or "").strip() + loop = asyncio.get_running_loop() + selected = await loop.run_in_executor(None, _dumas_native_folder_dialog, start_path) + if selected is None: + return web.json_response({"ok": False, "busy": True}) if not selected: return web.json_response({"ok": False, "cancelled": True}) return web.json_response({"ok": True, "path": selected})