diff --git a/server.py b/server.py index 82a035a85..abeda9612 100644 --- a/server.py +++ b/server.py @@ -148,6 +148,11 @@ class PromptServer(): return web.Response(status=404) images = [f for f in os.listdir(output_dir) if f.endswith('.png')] + # Check the length of the images list + if len(images) > 50: + # If there are more than 50 images, raise an exception + return web.json_response({"error": "Too many images, load from output disabled."}, status=400) + return web.json_response({"images": images}) @routes.get("/view") diff --git a/web/extensions/core/imageFeed.js b/web/extensions/core/imageFeed.js index b528229c8..7c5147dff 100644 --- a/web/extensions/core/imageFeed.js +++ b/web/extensions/core/imageFeed.js @@ -153,17 +153,6 @@ app.registerExtension({ } } - api.getOutput().then(data => { - try { - var images = data.filenames[0].map((filename) => { - return { filename: filename, type: 'output', subfolder: '' }; - }); - var output = {images: images} - var detail = {output: output} - loadImages(detail); - } catch(err){} - }); - // stop resize function function stopResize() { document.removeEventListener("mousemove", resize); @@ -225,7 +214,24 @@ app.registerExtension({ allImages = [] imageList.replaceChildren(menu, resizeHandle); }; - + api.getOutput().then(data => { + try { + if (data.message == "Success"){ + var images = data.filenames[0].map((filename) => { + return { filename: filename, type: 'output', subfolder: '' }; + }); + var output = {images: images} + var detail = {output: output} + loadImages(detail); + } + else { + deleteAllButton.setAttribute("disabled", true); + } + } catch(err){ + deleteAllButton.setAttribute("disabled", true); + console.error(err); + } + }); showButton.classList.add("comfy-settings-btn"); showButton.style.right = "16px"; showButton.style.cursor = "pointer"; diff --git a/web/scripts/api.js b/web/scripts/api.js index 5c128122b..6af88ced0 100644 --- a/web/scripts/api.js +++ b/web/scripts/api.js @@ -230,15 +230,21 @@ class ComfyApi extends EventTarget { * Gets the prompt execution history * @returns Prompt history including node outputs */ - async getOutput() { - try { - const res = await fetch("/output/images"); - return { filenames: Object.values(await res.json()) }; - } catch (error) { - console.error(error); - return { output: [] }; - } - } + async getOutput() { + try { + const res = await fetch("/output/images"); + if (res.status == 200){ + var values = await res.json(); + return { filenames: Object.values(values), message: "Success" }; + } + else { + return { filenames: [], message: "Failed to fetch output images" }; + } + } catch (error) { + console.error(error); + return { filenames: [], message: "Failed to fetch output images" }; + } + } /** * Sends a POST request to the API