feat(@captn/react): add comfyui helpers

adds several helpful hooks, that make stable-diffusion more flexible to app developers

packages/react/package.json

@@ -33,6 +33,26 @@
       "require": "./dist/cjs/use-captain-action/index.js",
       "import": "./dist/esm/use-captain-action/index.js"
     },
+    "./use-comfyui": {
+      "types": "./dist/types/use-comfyui/index.d.ts",
+      "require": "./dist/cjs/use-comfyui/index.js",
+      "import": "./dist/esm/use-comfyui/index.js"
+    },
+    "./use-history-state": {
+      "types": "./dist/types/use-history-state/index.d.ts",
+      "require": "./dist/cjs/use-history-state/index.js",
+      "import": "./dist/esm/use-history-state/index.js"
+    },
+    "./use-inventory": {
+      "types": "./dist/types/use-inventory/index.d.ts",
+      "require": "./dist/cjs/use-inventory/index.js",
+      "import": "./dist/esm/use-inventory/index.js"
+    },
+    "./use-seed": {
+      "types": "./dist/types/use-seed/index.d.ts",
+      "require": "./dist/cjs/use-seed/index.js",
+      "import": "./dist/esm/use-seed/index.js"
+    },
     "./use-resettable-state": {
       "types": "./dist/types/use-resettable-state/index.d.ts",
       "require": "./dist/cjs/use-resettable-state/index.js",
@@ -91,7 +111,8 @@
     "dot-prop": "^8.0.2",
     "lodash.isequal": "^4.5.0",
     "use-debounce": "^10.0.0",
-    "uuid": "^9.0.1"
+    "uuid": "^9.0.1",
+    "swr": "^2.2.5"
   },
   "devDependencies": {
     "@types/lodash.isequal": "^4.5.8",

packages/react/src/index.ts

@@ -1,12 +1,15 @@
 export * from "./types";
-export * from "./use-sdk";
 export * from "./use-captain-action";
+export * from "./use-comfyui";
+export * from "./use-history-state";
+export * from "./use-inventory";
 export * from "./use-language";
 export * from "./use-object";
 export * from "./use-required-downloads";
 export * from "./use-resettable-state";
 export * from "./use-save-image";
 export * from "./use-sdk";
+export * from "./use-seed";
 export * from "./use-text-to-image";
 export * from "./use-theme";
 export * from "./use-unload";

packages/react/src/use-comfyui/index.ts

@@ -0,0 +1,131 @@
+import { ComfyUIUpdate } from "@captn/utils/types";
+import { useCallback, useEffect, useState } from "react";
+
+import { StableDiffusion, useInventory } from "../use-inventory";
+import { useSDK } from "../use-sdk";
+import { useUnload } from "../use-unload";
+
+/**
+ * A hook to interact with the ComfyUI application. This hook manages the state of the image generation queue,
+ * interacts with the Stable Diffusion inventory, and handles message communication with the SDK.
+ *
+ * @param {string} appId - The ID of the application using the ComfyUI.
+ * @param {(image: string, meta: { filename: string; subfolder: string }) => void} onGenerated - A callback function that gets called when an image is generated. It receives the generated image's temporary URL and metadata.
+ *
+ * @returns {Object} - An object containing the following properties and functions:
+ *   - `models`: An object containing arrays of different model types (loras, checkpoints, vae, upscalers) from the Stable Diffusion inventory.
+ *   - `generate`: A function to queue an image generation workflow.
+ *   - `isGenerating`: A boolean indicating if an image is currently being generated.
+ *   - `image`: A string URL of the currently generated image.
+ *   - `queueSize`: A number indicating the current size of the image generation queue.
+ *
+ * @example
+ * ```tsx
+ * const { models, generate, isGenerating, image, queueSize } = useComfyUI(appId, (image, meta) => {
+ *   console.log("Generated image:", image, "with metadata:", meta);
+ * });
+ *
+ * // To generate an image:
+ * generate("workflow_string_here");
+ *
+ * // To access models:
+ * console.log(models.loras, models.checkpoints, models.vae, models.upscalers);
+ *
+ * // To check if an image is being generated:
+ * console.log(isGenerating);
+ *
+ * // To access the currently generated image:
+ * console.log(image);
+ *
+ * // To check the size of the generation queue:
+ * console.log(queueSize);
+ * ```
+ */
+export function useComfyUI(
+	appId: string,
+	onGenerated: (_image: string, meta: { filename: string; subfolder: string }) => void
+) {
+	const [queueSize, setQueueSize] = useState(0);
+	const [isGenerating, setIsGenerating] = useState(false);
+	const [image, setImage] = useState<string | null>(null);
+	const { data: inventoryData } = useInventory<StableDiffusion>("stable-diffusion");
+
+	const loras = inventoryData?.loras ?? [];
+	const checkpoints = inventoryData?.checkpoints ?? [];
+	const vae = inventoryData?.vae ?? [];
+	const upscalers = inventoryData?.upscalers ?? [];
+
+	const { send } = useSDK<unknown, object>(appId, {
+		async onMessage(message) {
+			switch (message.action) {
+				case "comfyui:update": {
+					const { data, comfyUITemporaryPath } = message.payload as {
+						data: ComfyUIUpdate;
+						comfyUITemporaryPath: string;
+					};
+					if (data.type === "status") {
+						setQueueSize(data.data.status.exec_info.queue_remaining);
+						setIsGenerating(data.data.status.exec_info.queue_remaining > 0);
+					} else if (data.type === "executed") {
+						const { filename, subfolder } = data.data.output.images[0];
+						const temporaryImage =
+							`${comfyUITemporaryPath}/${subfolder}/${filename}`.replaceAll(
+								/\?+/g,
+								"/"
+							);
+						setImage(temporaryImage);
+						onGenerated(temporaryImage, { filename, subfolder });
+					}
+
+					break;
+				}
+
+				default: {
+					break;
+				}
+			}
+		},
+	});
+
+	const generate = useCallback(
+		(workflow: string) => {
+			setIsGenerating(true);
+
+			send({
+				action: "comfyui:queue",
+				payload: {
+					workflow,
+				},
+			});
+		},
+		[send]
+	);
+
+	useUnload(appId, "comfyui:unregisterClient", {});
+
+	useEffect(() => {
+		send({
+			action: "comfyui:registerClient",
+			payload: {},
+		});
+		return () => {
+			send({
+				action: "comfyui:unregisterClient",
+				payload: {},
+			});
+		};
+	}, [send]);
+
+	return {
+		models: {
+			loras,
+			checkpoints,
+			vae,
+			upscalers,
+		},
+		generate,
+		isGenerating,
+		image,
+		queueSize,
+	};
+}

packages/react/src/use-history-state/index.ts

@@ -0,0 +1,139 @@
+import { isEqual } from "lodash";
+import { useCallback, useRef, useState } from "react";
+
+/**
+ * A custom React hook to manage state with history, allowing undo, redo, and reset operations.
+ *
+ * @template T - The type of the state.
+ *
+ * @param {T} [initialState_] - The initial state value.
+ *
+ * @returns {[T, (newState: T) => void, Object]} - Returns the current state, a function to update the state,
+ * and an object with functions and properties for managing the state history.
+ *
+ * The returned object contains:
+ * - `undo`: A function to revert the state to the previous state.
+ * - `redo`: A function to advance the state to the next state, if available.
+ * - `reset`: A function to reset the state to the initial state.
+ * - `flush`: A function to reset the history, keeping the current state as the only entry.
+ * - `hasPast`: A boolean indicating if there is a previous state in the history.
+ * - `hasFuture`: A boolean indicating if there is a next state in the history.
+ * - `isPresent`: A boolean indicating if the current state is the latest state.
+ * - `isPast`: A boolean indicating if the current state is not the latest state.
+ * - `initialState`: A boolean indicating if the current state is the initial state.
+ * - `historySize`: The total number of states in the history.
+ * - `pointerPosition`: The current position in the history.
+ * - `canReset`: A boolean indicating if the state can be reset to the initial state.
+ * - `modifiedSinceFlush`: A boolean indicating if the state has been modified since the last flush.
+ *
+ * @example
+ * ```js
+ * const [state, setState, history] = useHistoryState({ count: 0 });
+ *
+ * // Update state
+ * setState({ count: state.count + 1 });
+ *
+ * // Undo the last state change
+ * history.undo();
+ *
+ * // Redo the undone state change
+ * history.redo();
+ *
+ * // Reset state to the initial value
+ * history.reset();
+ *
+ * // Flush the history, keeping the current state as the only entry
+ * history.flush();
+ *
+ * // Check if there are past states
+ * console.log(history.hasPast); // true or false
+ *
+ * // Check if there are future states
+ * console.log(history.hasFuture); // true or false
+ * ```
+ */
+export function useHistoryState<T>(initialState_?: T) {
+	const [state, _setState] = useState<T>(initialState_);
+	const [history, setHistory] = useState({
+		states: initialState_ ? [initialState_] : [],
+		pointer: 0,
+		baseState: initialState_,
+	});
+	const lastFlushedPointer = useRef(0);
+
+	const flush = useCallback(() => {
+		// Set the current state as the only entry in the history
+		setHistory({
+			states: state ? [state] : [],
+			pointer: 0,
+			baseState: state,
+		});
+		lastFlushedPointer.current = 0;
+	}, [state]);
+
+	const setState = useCallback(
+		(newState: T) => {
+			if (!isEqual(state, newState)) {
+				// Update the state and add it to the history only if the new state is different
+				_setState(newState);
+				setHistory(previousState => ({
+					...previousState,
+					states: [...previousState.states.slice(0, previousState.pointer + 1), newState],
+					pointer: previousState.pointer + 1,
+				}));
+			}
+		},
+		[state, flush]
+	);
+
+	const undo = useCallback(() => {
+		if (history.pointer > 0) {
+			_setState(history.states[history.pointer - 1]);
+			setHistory({ ...history, pointer: history.pointer - 1 });
+		}
+	}, [history]);
+
+	const redo = useCallback(() => {
+		if (history.pointer !== history.states.length - 1) {
+			_setState(history.states[history.pointer + 1]);
+			setHistory({ ...history, pointer: history.pointer + 1 });
+		}
+	}, [history]);
+
+	const reset = useCallback(() => {
+		if (history.baseState) {
+			_setState(history.baseState);
+			setHistory({ states: [history.baseState], pointer: 0, baseState: history.baseState });
+		}
+	}, [history]);
+
+	const hasPast = history.pointer > 0;
+	const hasFuture = history.pointer < history.states.length - 1;
+	const isPresent = history.pointer === history.states.length - 1;
+	const isPast = history.pointer < history.states.length - 1;
+	const initialState = history.pointer === 0;
+	const historySize = history.states.length;
+	const pointerPosition = history.pointer;
+	const canReset = history.pointer !== 0;
+	const modifiedSinceFlush = history.pointer !== lastFlushedPointer.current;
+
+	return [
+		state,
+		setState,
+		{
+			undo,
+			redo,
+			reset,
+			flush,
+			hasPast,
+			hasFuture,
+			isPresent,
+			isPast,
+			initialState,
+			historySize,
+			pointerPosition,
+			canReset,
+			modifiedSinceFlush,
+		},
+	] as const;
+}