From 3eccfd82e8fb96293b9ae3cc1a72912aa97a0f7e Mon Sep 17 00:00:00 2001
From: Royce Schultz <royceschultz1@gmail.com>
Date: Thu, 16 May 2024 15:27:48 -0600
Subject: [PATCH] Adds unified notification to send all types of notifications
 in the same node.

---
 __init__.py                   |  3 ++
 nodes/unified_notification.py | 29 +++++++++++++++++++
 web/unified_notification.js   | 36 ++++++++++++++++++++++++
 web/util.js                   | 52 +++++++++++++++++++++++++++++++++++
 4 files changed, 120 insertions(+)
 create mode 100644 nodes/unified_notification.py
 create mode 100644 web/unified_notification.js
 create mode 100644 web/util.js

diff --git a/__init__.py b/__init__.py
index 861603b..4f82269 100644
--- a/__init__.py
+++ b/__init__.py
@@ -1,14 +1,17 @@
 from .nodes.play_sound import PlaySound
 from .nodes.system_notification import SystemNotification
+from .nodes.unified_notification import UnifiedNotification
 
 NODE_CLASS_MAPPINGS = {
     "Notif-PlaySound": PlaySound,
     "Notif-SystemNotification": SystemNotification,
+    "Notif-UnifiedNotification": UnifiedNotification,
 }
 
 NODE_DISPLAY_NAME_MAPPINGS = {
     "Notif-PlaySound": "Play Sound",
     "Notif-SystemNotification": "System Notification",
+    "Notif-UnifiedNotification": "Unified Notification",
 }
 
 WEB_DIRECTORY = "./web"
diff --git a/nodes/unified_notification.py b/nodes/unified_notification.py
new file mode 100644
index 0000000..ef91a6a
--- /dev/null
+++ b/nodes/unified_notification.py
@@ -0,0 +1,29 @@
+from .util import ComfyAnyType
+import json
+
+class UnifiedNotification:
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "any": (ComfyAnyType("*"), {}),
+                "mode": (["always", "on empty queue"], {}),
+                "system_notification": ("BOOLEAN", {"default": True}),
+                "notification_text": ('STRING', {'default': 'Your notification has triggered.'}),
+                "play_sound": ("BOOLEAN", {"default": True}),
+                "volume": ("FLOAT", {"min": 0, "max": 1, "step": 0.1, "default": 0.5}),
+                "file": ("STRING", { "default": "notify.mp3" }),
+            },
+        }
+
+    FUNCTION = "nop"
+    OUTPUT_NODE = True
+    RETURN_TYPES = tuple()
+    CATEGORY = "notifications"
+
+    def nop(self, *args, **kwargs):
+        for k, v in kwargs.items():
+            kwargs[k] = [v] # UI expects an iterable for some reason
+        del kwargs['any'] # Avoid passing large objects
+        kwargs['args'] = args
+        return {"ui": kwargs, "result": (0, )}
diff --git a/web/unified_notification.js b/web/unified_notification.js
new file mode 100644
index 0000000..b43c0d7
--- /dev/null
+++ b/web/unified_notification.js
@@ -0,0 +1,36 @@
+import { app } from '../../../scripts/app.js'
+import {
+    notificationSetup, sendNotification, getNotificationTextFromArguments,
+    playSound,
+    appQueueIsEmpty,
+} from './util.js'
+
+
+app.registerExtension({
+    name: 'Notifications.UnifiedNotification',
+    async beforeRegisterNodeDef(nodeType, nodeData, app) {
+        if (nodeData.name === 'Notif-UnifiedNotification') {
+            const onExecuted = nodeType.prototype.onExecuted
+            nodeType.prototype.onExecuted = async function () {
+                onExecuted?.apply(this, arguments)
+                const args = arguments[0] ?? {}
+                if (args.mode[0] === 'on empty queue') {
+                    if(!await appQueueIsEmpty(app)) return
+                }
+                if (args.system_notification[0]) {
+                    const notif_text = getNotificationTextFromArguments(arguments)
+                    sendNotification('ComfyUI', notif_text)
+                }
+                if (args.play_sound[0]) {
+                    playSound(args.file[0], args.volume[0] ?? 0.5)
+                }
+            }
+
+            const onNodeCreated = nodeType.prototype.onNodeCreated
+            nodeType.prototype.onNodeCreated = function () {
+                onNodeCreated?.apply(this, arguments)
+                notificationSetup()
+            }
+        }
+    },
+})
diff --git a/web/util.js b/web/util.js
new file mode 100644
index 0000000..ff223ae
--- /dev/null
+++ b/web/util.js
@@ -0,0 +1,52 @@
+export const notificationSetup = () => {
+    if (!('Notification' in window)) {
+        console.log('This browser does not support notifications.')
+        alert('This browser does not support notifications.')
+        return
+    }
+    if (Notification.permission === 'denied') {
+        console.log('Notifications are blocked. Please enable them in your browser settings.')
+        alert('Notifications are blocked. Please enable them in your browser settings.')
+        return
+    }
+    if (Notification.permission !== 'granted') {
+        Notification.requestPermission()
+    }
+    return true
+}
+
+export const sendNotification = (title, body) => {
+    if (!notificationSetup()) return
+    const notification = new Notification(title, { body })
+    return notification
+}
+
+export const getNotificationTextFromArguments = (args) => {
+    return args[0]?.notification_text[0] ?? 'Your notification has triggered.'
+}
+
+export const playSound = (file, volume) => {
+    if (!file) {
+        file = 'notify.mp3'
+    }
+    if (!file.startsWith('http')) {
+        if (!file.includes('/')) {
+            file = 'assets/' + file
+        }
+        file = new URL(file, import.meta.url)
+    }
+    const url = new URL(file)
+    const audio = new Audio(url)
+    audio.volume = volume
+    audio.play()
+}
+
+export const appQueueIsEmpty = async (app) => {
+    if (app.ui.lastQueueSize !== 0) {
+        await new Promise((r) => setTimeout(r, 500))
+    }
+    if (app.ui.lastQueueSize !== 0) {
+        return false
+    }
+    return true
+}