Changed to default to browser storage and add server-storage arg

comfyanonymous · Jan 6, 2024 · f698982 · f698982
1 parent fa3127c
commit f698982
Show file tree

Hide file tree

Showing 9 changed files with 159 additions and 56 deletions.
diff --git a/app/user_manager.py b/app/user_manager.py
@@ -91,10 +91,13 @@ def add_routes(self, routes):
         @routes.get("/users")
         async def get_users(request):
             if args.multi_user:
-                return web.json_response(self.users)
+                return web.json_response({"storage": "server", "users": self.users})
             else:
                 user_dir = self.get_request_user_filepath(request, None, create_dir=False)
-                return web.json_response(os.path.exists(user_dir))
+                return web.json_response({
+                    "storage": "server" if args.server_storage else "browser", 
+                    "migrated": os.path.exists(user_dir)
+                })
 
         @routes.post("/users")
         async def post_users(request):

diff --git a/comfy/cli_args.py b/comfy/cli_args.py
@@ -112,7 +112,8 @@ class LatentPreviewMethod(enum.Enum):
 
 parser.add_argument("--disable-metadata", action="store_true", help="Disable saving prompt metadata in files.")
 
-parser.add_argument("--multi-user", action="store_true", help="Enables per-user settings.")
+parser.add_argument("--server-storage", action="store_true", help="Saves settings and other user configuration on the server instead of in browser storage.")
+parser.add_argument("--multi-user", action="store_true", help="Enables per-user storage. If enabled, server-storage will be unconditionally enabled.")
 
 if comfy.options.args_parsing:
     args = parser.parse_args()
@@ -124,3 +125,6 @@ class LatentPreviewMethod(enum.Enum):
 
 if args.disable_auto_launch:
     args.auto_launch = False
+
+if args.multi_user:
+    args.server_storage = True
diff --git a/tests-ui/tests/users.test.js b/tests-ui/tests/users.test.js
@@ -21,20 +21,48 @@ describe("users", () => {
 	}
 
 	describe("multi-user", () => {
+		function mockAddStylesheet() {
+			const utils = require("../../web/scripts/utils");
+			utils.addStylesheet = jest.fn().mockReturnValue(Promise.resolve());
+		}
+
+		async function waitForUserScreenShow() {
+			mockAddStylesheet();
+
+			// Wait for "show" to be called
+			const { UserSelectionScreen } = require("../../web/scripts/ui/userSelection");
+			let resolve, reject;
+			const fn = UserSelectionScreen.prototype.show;
+			const p = new Promise((res, rej) => {
+				resolve = res;
+				reject = rej;
+			});
+			jest.spyOn(UserSelectionScreen.prototype, "show").mockImplementation(async (...args) => {
+				const res = fn(...args);
+				await new Promise(process.nextTick); // wait for promises to resolve
+				resolve();
+				return res;
+			});
+			// @ts-ignore
+			setTimeout(() => reject("timeout waiting for UserSelectionScreen to be shown."), 500);
+			await p;
+			await new Promise(process.nextTick); // wait for promises to resolve
+		}
+
 		async function testUserScreen(onShown, users) {
 			if (!users) {
 				users = {};
 			}
 			const starting = start({
 				resetEnv: true,
-				users,
+				userConfig: { storage: "server", users },
 			});
 
 			// Ensure no current user
 			expect(localStorage["Comfy.userId"]).toBeFalsy();
 			expect(localStorage["Comfy.userName"]).toBeFalsy();
 
-			await new Promise(process.nextTick); // wait for promises to resolve
+			await waitForUserScreenShow();
 
 			const selection = document.querySelectorAll("#comfy-user-selection")?.[0];
 			expect(selection).toBeTruthy();
@@ -153,8 +181,11 @@ describe("users", () => {
 		it("doesnt show user screen if current user", async () => {
 			const starting = start({
 				resetEnv: true,
-				users: {
-					"User!": "User",
+				userConfig: {
+					storage: "server",
+					users: {
+						"User!": "User",
+					},
 				},
 				localStorage: {
 					"Comfy.userId": "User!",
@@ -170,8 +201,11 @@ describe("users", () => {
 		it("allows user switching", async () => {
 			const { app } = await start({
 				resetEnv: true,
-				users: {
-					"User!": "User",
+				userConfig: {
+					storage: "server",
+					users: {
+						"User!": "User",
+					},
 				},
 				localStorage: {
 					"Comfy.userId": "User!",
@@ -187,9 +221,9 @@ describe("users", () => {
 		it("doesnt show user creation if no default user", async () => {
 			const { app } = await start({
 				resetEnv: true,
-				users: false,
+				userConfig: { migrated: false, storage: "server" },
 			});
-            expectNoUserScreen();
+			expectNoUserScreen();
 
 			// It should store the settings
 			const { api } = require("../../web/scripts/api");
@@ -201,9 +235,9 @@ describe("users", () => {
 		it("doesnt show user creation if default user", async () => {
 			const { app } = await start({
 				resetEnv: true,
-				users: true,
+				userConfig: { migrated: true, storage: "server" },
 			});
-            expectNoUserScreen();
+			expectNoUserScreen();
 
 			// It should store the settings
 			const { api } = require("../../web/scripts/api");
@@ -214,8 +248,46 @@ describe("users", () => {
 		it("doesnt allow user switching", async () => {
 			const { app } = await start({
 				resetEnv: true,
+				userConfig: { migrated: true, storage: "server" },
 			});
-            expectNoUserScreen();
+			expectNoUserScreen();
+
+			expect(app.ui.settings.settingsLookup["Comfy.SwitchUser"]).toBeFalsy();
+		});
+	});
+	describe("browser-user", () => {
+		it("doesnt show user creation if no default user", async () => {
+			const { app } = await start({
+				resetEnv: true,
+				userConfig: { migrated: false, storage: "browser" },
+			});
+			expectNoUserScreen();
+
+			// It should store the settings
+			const { api } = require("../../web/scripts/api");
+			expect(api.storeSettings).toHaveBeenCalledTimes(0);
+			expect(api.storeUserData).toHaveBeenCalledTimes(0);
+			expect(app.isNewUserSession).toBeFalsy();
+		});
+		it("doesnt show user creation if default user", async () => {
+			const { app } = await start({
+				resetEnv: true,
+				userConfig: { migrated: true, storage: "server" },
+			});
+			expectNoUserScreen();
+
+			// It should store the settings
+			const { api } = require("../../web/scripts/api");
+			expect(api.storeSettings).toHaveBeenCalledTimes(0);
+			expect(api.storeUserData).toHaveBeenCalledTimes(0);
+			expect(app.isNewUserSession).toBeFalsy();
+		});
+		it("doesnt allow user switching", async () => {
+			const { app } = await start({
+				resetEnv: true,
+				userConfig: { migrated: true, storage: "browser" },
+			});
+			expectNoUserScreen();
 
 			expect(app.ui.settings.settingsLookup["Comfy.SwitchUser"]).toBeFalsy();
 		});

diff --git a/tests-ui/utils/index.js b/tests-ui/utils/index.js
@@ -27,18 +27,12 @@ export async function start(config = {}) {
 	Object.assign(localStorage, config.localStorage ?? {});
 	document.body.innerHTML = html;
 
-	mockUtils();
 	mockApi(config);
 	const { app } = require("../../web/scripts/app");
 	config.preSetup?.(app);
 	await app.setup();
-	return { ...Ez.graph(app, global["LiteGraph"], global["LGraphCanvas"]), app };
-}
 
-function mockUtils() {
-	jest.mock("../../web/scripts/utils", () => ({
-		addStylesheet: () => Promise.resolve()
-	}));
+	return { ...Ez.graph(app, global["LiteGraph"], global["LGraphCanvas"]), app };
 }
 
 /**

diff --git a/tests-ui/utils/setup.js b/tests-ui/utils/setup.js
@@ -21,14 +21,14 @@ function* walkSync(dir) {
  * @param {{ 
  * 	mockExtensions?: string[], 
  * 	mockNodeDefs?: Record<string, ComfyObjectInfo>,
- * 	users?: boolean | Record<string, string>
 * 	settings?: Record<string, string>
+* 	userConfig?: {storage: "server" | "browser", users?: Record<string, any>, migrated?: boolean },
 * 	userData?: Record<string, any>
  * }} config
  */
 export function mockApi(config = {}) {
-	let { mockExtensions, mockNodeDefs, users, settings, userData } = {
-		users: true,
+	let { mockExtensions, mockNodeDefs, userConfig, settings, userData } = {
+		userConfig,
 		settings: {},
 		userData: {},
 		...config,
@@ -53,13 +53,13 @@ export function mockApi(config = {}) {
 		init: jest.fn(),
 		apiURL: jest.fn((x) => "../../web/" + x),
 		createUser: jest.fn((username) => {
-			if(username in users) {
+			if(username in userConfig.users) {
 				return { status: 400, json: () => "Duplicate" }
 			}
-			users[username + "!"] = username;
+			userConfig.users[username + "!"] = username;
 			return { status: 200, json: () => username + "!" }
 		}),
-		getUsers: jest.fn(() => users),
+		getUserConfig: jest.fn(() => userConfig ?? { storage: "browser", migrated: false }),
 		getSettings: jest.fn(() => settings),
 		storeSettings: jest.fn((v) => Object.assign(settings, v)),
 		getUserData: jest.fn((f) => {

diff --git a/web/extensions/core/nodeTemplates.js b/web/extensions/core/nodeTemplates.js
@@ -74,33 +74,44 @@ class ManageTemplates extends ComfyDialog {
 
 	async load() {
 		let templates;
-		if (app.isNewUserSession) {
-			// New user so migrate existing templates
+		if (app.storageLocation === "server") {
+			if (app.isNewUserSession) {
+				// New user so migrate existing templates
+				const json = localStorage.getItem(id);
+				if (json) {
+					templates = JSON.parse(json);
+				}
+				await api.storeUserData(file, json, { stringify: false });
+			} else {
+				const res = await api.getUserData(file);
+				if (res.status === 200) {
+					templates = await res.json();
+				} else if (res.status !== 404) {
+					console.error(res.status + " " + res.statusText);
+				}
+			}
+		} else {
 			const json = localStorage.getItem(id);
 			if (json) {
 				templates = JSON.parse(json);
 			}
-			await api.storeUserData(file, json, { stringify: false });
-		} else {
-			const res = await api.getUserData(file);
-			if (res.status === 200) {
-				templates = await res.json();
-			} else if (res.status !== 404) {
-				console.error(res.status + " " + res.statusText);
-			}
 		}
 
 		return templates ?? [];
 	}
 
 	async store() {
-		const templates = JSON.stringify(this.templates, undefined, 4);
-		localStorage.setItem(id, templates); // Backwards compatibility
-		try {
-			await api.storeUserData(file, templates, { stringify: false });
-		} catch (error) {
-			console.error(error);
-			alert(error.message);
+		if(app.storageLocation === "server") {
+			const templates = JSON.stringify(this.templates, undefined, 4);
+			localStorage.setItem(id, templates); // Backwards compatibility
+			try {
+				await api.storeUserData(file, templates, { stringify: false });
+			} catch (error) {
+				console.error(error);
+				alert(error.message);
+			}
+		} else {
+			localStorage.setItem(id, JSON.stringify(this.templates));
 		}
 	}
 

diff --git a/web/scripts/api.js b/web/scripts/api.js
@@ -324,10 +324,10 @@ class ComfyApi extends EventTarget {
 	}
 
 	/**
-	 * Gets a list of users or true if single user mode and default user is created, or false if single user mode and default user is not created.
-	 * @returns { Promise<string, unknown> | boolean } If multi-user, a dictionary of id -> value, else whether the default user is created
+	 * Gets user configuration data and where data should be stored
+	 * @returns { Promise<{ storage: "server" | "browser", users?: Promise<string, unknown>, migrated?: boolean }> } 
 	 */
-	async getUsers() {
+	async getUserConfig() {
 		return (await this.fetchApi("/users")).json();
 	}
 

diff --git a/web/scripts/app.js b/web/scripts/app.js
@@ -1308,10 +1308,11 @@ export class ComfyApp {
 	}
 
 	async #setUser() {
-		const users = await api.getUsers();
-		if(typeof users === "boolean") {
-			// Single user mode returns true/false for if the default user is created
-			if(!users) {
+		const userConfig = await api.getUserConfig();
+		this.storageLocation = userConfig.storage;
+		if (typeof userConfig.migrated == "boolean") {
+			// Single user mode migrated true/false for if the default user is created
+			if (!userConfig.migrated && this.storageLocation === "server") {
 				// Default user not created yet
 				await this.#migrateSettings();
 			}
@@ -1320,6 +1321,7 @@ export class ComfyApp {
 
 		this.multiUserServer = true;
 		let user = localStorage["Comfy.userId"];
+		const users = userConfig.users ?? {};
 		if (!user || !users[user]) {
 			// This will rarely be hit so move the loading to on demand
 			const { UserSelectionScreen } = await import("./ui/userSelection.js");

diff --git a/web/scripts/ui/settings.js b/web/scripts/ui/settings.js
@@ -38,19 +38,36 @@ export class ComfySettingsDialog extends ComfyDialog {
 	}
 
 	async load() {
-		this.settingsValues = await api.getSettings();
+		if (this.app.storageLocation === "browser") {
+			this.settingsValues = localStorage;
+		} else {
+			this.settingsValues = await api.getSettings();
+		}
+	}
+
+	getId(id) {
+		if (this.app.storageLocation === "browser") {
+			id = "Comfy.Settings." + id;
+		}
+		return id;
 	}
 
 	getSettingValue(id, defaultValue) {
-		return this.settingsValues[id] ?? defaultValue;
+		let value = this.settingsValues[this.getId(id)];
+		if(value != null) {
+			if(this.app.storageLocation === "browser") {
+				value = JSON.parse(value);
+			}
+		}
+		return value ?? defaultValue;
 	}
 
 	async setSettingValueAsync(id, value) {
 		const json = JSON.stringify(value);
 		localStorage["Comfy.Settings." + id] = json; // backwards compatibility for extensions keep setting in storage
 
-		let oldValue = this.settingsValues[id];
-		this.settingsValues[id] = value;
+		let oldValue = this.getSettingValue(id, undefined);
+		this.settingsValues[this.getId(id)] = value;
 
 		if (id in this.settingsLookup) {
 			this.settingsLookup[id].onChange?.(value, oldValue);
@@ -76,7 +93,7 @@ export class ComfySettingsDialog extends ComfyDialog {
 		}
 
 		let skipOnChange = false;
-		let value = this.settingsValues[id];
+		let value = this.getSettingValue(id);
 		if (value == null) {
 			if (this.app.isNewUserSession) {
 				// Check if we have a localStorage value but not a setting value and we are a new user