diff --git a/comfy/ldm/models/diffusion/ddim.py b/comfy/ldm/models/diffusion/ddim.py index fe39c76c..5e2d7364 100644 --- a/comfy/ldm/models/diffusion/ddim.py +++ b/comfy/ldm/models/diffusion/ddim.py @@ -18,7 +18,7 @@ class DDIMSampler(object): def register_buffer(self, name, attr): if type(attr) == torch.Tensor: if attr.device != self.device: - attr = attr.to(self.device) + attr = attr.float().to(self.device) setattr(self, name, attr) def make_schedule(self, ddim_num_steps, ddim_discretize="uniform", ddim_eta=0., verbose=True): diff --git a/comfy/model_management.py b/comfy/model_management.py index 809b19ea..0d5702b9 100644 --- a/comfy/model_management.py +++ b/comfy/model_management.py @@ -4,6 +4,7 @@ NO_VRAM = 1 LOW_VRAM = 2 NORMAL_VRAM = 3 HIGH_VRAM = 4 +MPS = 5 accelerate_enabled = False vram_state = NORMAL_VRAM @@ -76,10 +77,16 @@ if set_vram_to == LOW_VRAM or set_vram_to == NO_VRAM: total_vram_available_mb = (total_vram - 1024) // 2 total_vram_available_mb = int(max(256, total_vram_available_mb)) +try: + if torch.backends.mps.is_available(): + vram_state = MPS +except: + pass + if "--cpu" in sys.argv: vram_state = CPU -print("Set vram state to:", ["CPU", "NO VRAM", "LOW VRAM", "NORMAL VRAM", "HIGH VRAM"][vram_state]) +print("Set vram state to:", ["CPU", "NO VRAM", "LOW VRAM", "NORMAL VRAM", "HIGH VRAM", "MPS"][vram_state]) current_loaded_model = None @@ -128,6 +135,10 @@ def load_model_gpu(model): current_loaded_model = model if vram_state == CPU: pass + elif vram_state == MPS: + mps_device = torch.device("mps") + real_model.to(mps_device) + pass elif vram_state == NORMAL_VRAM or vram_state == HIGH_VRAM: model_accelerated = False real_model.cuda() @@ -155,9 +166,10 @@ def load_controlnet_gpu(models): if m not in models: m.cpu() + device = get_torch_device() current_gpu_controlnets = [] for m in models: - current_gpu_controlnets.append(m.cuda()) + current_gpu_controlnets.append(m.to(device)) def load_if_low_vram(model): @@ -173,6 +185,8 @@ def unload_if_low_vram(model): return model def get_torch_device(): + if vram_state == MPS: + return torch.device("mps") if vram_state == CPU: return torch.device("cpu") else: @@ -195,7 +209,7 @@ def get_free_memory(dev=None, torch_free_too=False): if dev is None: dev = get_torch_device() - if hasattr(dev, 'type') and dev.type == 'cpu': + if hasattr(dev, 'type') and (dev.type == 'cpu' or dev.type == 'mps'): mem_free_total = psutil.virtual_memory().available mem_free_torch = mem_free_total else: @@ -224,8 +238,12 @@ def cpu_mode(): global vram_state return vram_state == CPU +def mps_mode(): + global vram_state + return vram_state == MPS + def should_use_fp16(): - if cpu_mode(): + if cpu_mode() or mps_mode(): return False #TODO ? if torch.cuda.is_bf16_supported(): diff --git a/comfy/samplers.py b/comfy/samplers.py index bf4f1796..17201d9d 100644 --- a/comfy/samplers.py +++ b/comfy/samplers.py @@ -450,7 +450,7 @@ class KSampler: noise_mask = None if denoise_mask is not None: noise_mask = 1.0 - denoise_mask - sampler = DDIMSampler(self.model) + sampler = DDIMSampler(self.model, device=self.device) sampler.make_schedule_timesteps(ddim_timesteps=timesteps, verbose=False) z_enc = sampler.stochastic_encode(latent_image, torch.tensor([len(timesteps) - 1] * noise.shape[0]).to(self.device), noise=noise, max_denoise=max_denoise) samples, _ = sampler.sample_custom(ddim_timesteps=timesteps, diff --git a/comfyui_screenshot.png b/comfyui_screenshot.png index 72642438..c357e243 100644 Binary files a/comfyui_screenshot.png and b/comfyui_screenshot.png differ diff --git a/nodes.py b/nodes.py index 0b8be765..a981abb8 100644 --- a/nodes.py +++ b/nodes.py @@ -241,8 +241,8 @@ class LoraLoader: return {"required": { "model": ("MODEL",), "clip": ("CLIP", ), "lora_name": (folder_paths.get_filename_list("loras"), ), - "strength_model": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 10.0, "step": 0.01}), - "strength_clip": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 10.0, "step": 0.01}), + "strength_model": ("FLOAT", {"default": 1.0, "min": -10.0, "max": 10.0, "step": 0.01}), + "strength_clip": ("FLOAT", {"default": 1.0, "min": -10.0, "max": 10.0, "step": 0.01}), }} RETURN_TYPES = ("MODEL", "CLIP") FUNCTION = "load_lora" @@ -752,7 +752,7 @@ class SaveImage: full_output_folder = os.path.join(self.output_dir, subfolder) - if os.path.commonpath((self.output_dir, os.path.realpath(full_output_folder))) != self.output_dir: + if os.path.commonpath((self.output_dir, os.path.abspath(full_output_folder))) != self.output_dir: print("Saving image outside the output folder is not allowed.") return {} @@ -908,6 +908,69 @@ class ImageInvert: return (s,) +class ImagePadForOutpaint: + + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "image": ("IMAGE",), + "left": ("INT", {"default": 0, "min": 0, "max": MAX_RESOLUTION, "step": 64}), + "top": ("INT", {"default": 0, "min": 0, "max": MAX_RESOLUTION, "step": 64}), + "right": ("INT", {"default": 0, "min": 0, "max": MAX_RESOLUTION, "step": 64}), + "bottom": ("INT", {"default": 0, "min": 0, "max": MAX_RESOLUTION, "step": 64}), + "feathering": ("INT", {"default": 40, "min": 0, "max": MAX_RESOLUTION, "step": 1}), + } + } + + RETURN_TYPES = ("IMAGE", "MASK") + FUNCTION = "expand_image" + + CATEGORY = "image" + + def expand_image(self, image, left, top, right, bottom, feathering): + d1, d2, d3, d4 = image.size() + + new_image = torch.zeros( + (d1, d2 + top + bottom, d3 + left + right, d4), + dtype=torch.float32, + ) + new_image[:, top:top + d2, left:left + d3, :] = image + + mask = torch.ones( + (d2 + top + bottom, d3 + left + right), + dtype=torch.float32, + ) + + t = torch.zeros( + (d2, d3), + dtype=torch.float32 + ) + + if feathering > 0 and feathering * 2 < d2 and feathering * 2 < d3: + + for i in range(d2): + for j in range(d3): + dt = i if top != 0 else d2 + db = d2 - i if bottom != 0 else d2 + + dl = j if left != 0 else d3 + dr = d3 - j if right != 0 else d3 + + d = min(dt, db, dl, dr) + + if d >= feathering: + continue + + v = (feathering - d) / feathering + + t[i, j] = v * v + + mask[top:top + d2, left:left + d3] = t + + return (new_image, mask) + + NODE_CLASS_MAPPINGS = { "KSampler": KSampler, "CheckpointLoader": CheckpointLoader, @@ -926,6 +989,7 @@ NODE_CLASS_MAPPINGS = { "LoadImageMask": LoadImageMask, "ImageScale": ImageScale, "ImageInvert": ImageInvert, + "ImagePadForOutpaint": ImagePadForOutpaint, "ConditioningCombine": ConditioningCombine, "ConditioningSetArea": ConditioningSetArea, "KSamplerAdvanced": KSamplerAdvanced, diff --git a/notebooks/comfyui_colab.ipynb b/notebooks/comfyui_colab.ipynb index d9726947..5108ec83 100644 --- a/notebooks/comfyui_colab.ipynb +++ b/notebooks/comfyui_colab.ipynb @@ -1,29 +1,13 @@ { - "nbformat": 4, - "nbformat_minor": 0, - "metadata": { - "colab": { - "provenance": [] - }, - "kernelspec": { - "name": "python3", - "display_name": "Python 3" - }, - "language_info": { - "name": "python" - }, - "accelerator": "GPU", - "gpuClass": "standard" - }, "cells": [ { "cell_type": "markdown", - "source": [ - "Git clone the repo and install the requirements. (ignore the pip errors about protobuf)" - ], "metadata": { "id": "aaaaaaaaaa" - } + }, + "source": [ + "Git clone the repo and install the requirements. (ignore the pip errors about protobuf)" + ] }, { "cell_type": "code", @@ -33,22 +17,55 @@ }, "outputs": [], "source": [ - "!git clone https://github.com/comfyanonymous/ComfyUI\n", - "%cd ComfyUI\n", - "!pip install xformers -r requirements.txt" + "#@title Environment Setup\n", + "\n", + "from pathlib import Path\n", + "\n", + "OPTIONS = {}\n", + "\n", + "USE_GOOGLE_DRIVE = False #@param {type:\"boolean\"}\n", + "UPDATE_COMFY_UI = True #@param {type:\"boolean\"}\n", + "WORKSPACE = 'ComfyUI'\n", + "OPTIONS['USE_GOOGLE_DRIVE'] = USE_GOOGLE_DRIVE\n", + "OPTIONS['UPDATE_COMFY_UI'] = UPDATE_COMFY_UI\n", + "\n", + "if OPTIONS['USE_GOOGLE_DRIVE']:\n", + " !echo \"Mounting Google Drive...\"\n", + " %cd /\n", + " \n", + " from google.colab import drive\n", + " drive.mount('/content/drive')\n", + "\n", + " WORKSPACE = \"/content/drive/MyDrive/ComfyUI\"\n", + " %cd /content/drive/MyDrive\n", + "\n", + "![ ! -d $WORKSPACE ] && echo -= Initial setup ComfyUI =- && git clone https://github.com/comfyanonymous/ComfyUI\n", + "%cd $WORKSPACE\n", + "\n", + "if OPTIONS['UPDATE_COMFY_UI']:\n", + " !echo -= Updating ComfyUI =-\n", + " !git pull\n", + "\n", + "!echo -= Install dependencies =-\n", + "!pip -q install xformers -r requirements.txt" ] }, { "cell_type": "markdown", - "source": [ - "Download some models/checkpoints/vae or custom comfyui nodes (uncomment the commands for the ones you want)" - ], "metadata": { "id": "cccccccccc" - } + }, + "source": [ + "Download some models/checkpoints/vae or custom comfyui nodes (uncomment the commands for the ones you want)" + ] }, { "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "dddddddddd" + }, + "outputs": [], "source": [ "# Checkpoints\n", "\n", @@ -110,26 +127,26 @@ "#!wget -c https://huggingface.co/sberbank-ai/Real-ESRGAN/resolve/main/RealESRGAN_x4.pth -P ./models/upscale_models/\n", "\n", "\n" - ], - "metadata": { - "id": "dddddddddd" - }, - "execution_count": null, - "outputs": [] + ] }, { "cell_type": "markdown", + "metadata": { + "id": "kkkkkkkkkkkkkk" + }, "source": [ "### Run ComfyUI with localtunnel (Recommended Way)\n", "\n", "\n" - ], - "metadata": { - "id": "kkkkkkkkkkkkkk" - } + ] }, { "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "jjjjjjjjjjjjj" + }, + "outputs": [], "source": [ "!npm install -g localtunnel\n", "\n", @@ -154,15 +171,13 @@ "threading.Thread(target=iframe_thread, daemon=True, args=(8188,)).start()\n", "\n", "!python main.py --dont-print-server" - ], - "metadata": { - "id": "jjjjjjjjjjjjj" - }, - "execution_count": null, - "outputs": [] + ] }, { "cell_type": "markdown", + "metadata": { + "id": "gggggggggg" + }, "source": [ "### Run ComfyUI with colab iframe (use only in case the previous way with localtunnel doesn't work)\n", "\n", @@ -171,13 +186,15 @@ "If you want to open it in another window use the link.\n", "\n", "Note that some UI features like live image previews won't work because the colab iframe blocks websockets." - ], - "metadata": { - "id": "gggggggggg" - } + ] }, { "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "hhhhhhhhhh" + }, + "outputs": [], "source": [ "import threading\n", "import time\n", @@ -198,12 +215,23 @@ "threading.Thread(target=iframe_thread, daemon=True, args=(8188,)).start()\n", "\n", "!python main.py --dont-print-server" - ], - "metadata": { - "id": "hhhhhhhhhh" - }, - "execution_count": null, - "outputs": [] + ] } - ] -} \ No newline at end of file + ], + "metadata": { + "accelerator": "GPU", + "colab": { + "provenance": [] + }, + "gpuClass": "standard", + "kernelspec": { + "display_name": "Python 3", + "name": "python3" + }, + "language_info": { + "name": "python" + } + }, + "nbformat": 4, + "nbformat_minor": 0 +} diff --git a/server.py b/server.py index 2593e808..1a370317 100644 --- a/server.py +++ b/server.py @@ -127,7 +127,7 @@ class PromptServer(): output_dir = os.path.join(os.path.dirname(os.path.realpath(__file__)), type) if "subfolder" in request.rel_url.query: full_output_dir = os.path.join(output_dir, request.rel_url.query["subfolder"]) - if os.path.commonpath((os.path.realpath(full_output_dir), output_dir)) != output_dir: + if os.path.commonpath((os.path.abspath(full_output_dir), output_dir)) != output_dir: return web.Response(status=403) output_dir = full_output_dir diff --git a/web/extensions/core/widgetInputs.js b/web/extensions/core/widgetInputs.js new file mode 100644 index 00000000..ff9227d2 --- /dev/null +++ b/web/extensions/core/widgetInputs.js @@ -0,0 +1,362 @@ +import { ComfyWidgets, addRandomizeWidget } from "/scripts/widgets.js"; +import { app } from "/scripts/app.js"; + +const CONVERTED_TYPE = "converted-widget"; +const VALID_TYPES = ["STRING", "combo", "number"]; + +function isConvertableWidget(widget, config) { + return VALID_TYPES.includes(widget.type) || VALID_TYPES.includes(config[0]); +} + +function hideWidget(node, widget, suffix = "") { + widget.origType = widget.type; + widget.origComputeSize = widget.computeSize; + widget.origSerializeValue = widget.serializeValue; + widget.computeSize = () => [0, -4]; // -4 is due to the gap litegraph adds between widgets automatically + widget.type = CONVERTED_TYPE + suffix; + widget.serializeValue = () => { + // Prevent serializing the widget if we have no input linked + const { link } = node.inputs.find((i) => i.widget?.name === widget.name); + if (link == null) { + return undefined; + } + return widget.value; + }; + + // Hide any linked widgets, e.g. seed+randomize + if (widget.linkedWidgets) { + for (const w of widget.linkedWidgets) { + hideWidget(node, w, ":" + widget.name); + } + } +} + +function showWidget(widget) { + widget.type = widget.origType; + widget.computeSize = widget.origComputeSize; + widget.serializeValue = widget.origSerializeValue; + + delete widget.origType; + delete widget.origComputeSize; + delete widget.origSerializeValue; + + // Hide any linked widgets, e.g. seed+randomize + if (widget.linkedWidgets) { + for (const w of widget.linkedWidgets) { + showWidget(w); + } + } +} + +function convertToInput(node, widget, config) { + hideWidget(node, widget); + + const { linkType } = getWidgetType(config); + + // Add input and store widget config for creating on primitive node + const sz = node.size; + node.addInput(widget.name, linkType, { + widget: { name: widget.name, config }, + }); + + // Restore original size but grow if needed + node.setSize([Math.max(sz[0], node.size[0]), Math.max(sz[1], node.size[1])]); +} + +function convertToWidget(node, widget) { + showWidget(widget); + const sz = node.size; + node.removeInput(node.inputs.findIndex((i) => i.widget?.name === widget.name)); + + // Restore original size but grow if needed + node.setSize([Math.max(sz[0], node.size[0]), Math.max(sz[1], node.size[1])]); +} + +function getWidgetType(config) { + // Special handling for COMBO so we restrict links based on the entries + let type = config[0]; + let linkType = type; + if (type instanceof Array) { + type = "COMBO"; + linkType = linkType.join(","); + } + return { type, linkType }; +} + +app.registerExtension({ + name: "Comfy.WidgetInputs", + async beforeRegisterNodeDef(nodeType, nodeData, app) { + // Add menu options to conver to/from widgets + const origGetExtraMenuOptions = nodeType.prototype.getExtraMenuOptions; + nodeType.prototype.getExtraMenuOptions = function (_, options) { + const r = origGetExtraMenuOptions ? origGetExtraMenuOptions.apply(this, arguments) : undefined; + + if (this.widgets) { + let toInput = []; + let toWidget = []; + for (const w of this.widgets) { + if (w.type === CONVERTED_TYPE) { + toWidget.push({ + content: `Convert ${w.name} to widget`, + callback: () => convertToWidget(this, w), + }); + } else { + const config = nodeData?.input?.required[w.name] || [w.type, w.options || {}]; + if (isConvertableWidget(w, config)) { + toInput.push({ + content: `Convert ${w.name} to input`, + callback: () => convertToInput(this, w, config), + }); + } + } + } + if (toInput.length) { + options.push(...toInput, null); + } + + if (toWidget.length) { + options.push(...toWidget, null); + } + } + + return r; + }; + + // On initial configure of nodes hide all converted widgets + const origOnConfigure = nodeType.prototype.onConfigure; + nodeType.prototype.onConfigure = function () { + const r = origOnConfigure ? origOnConfigure.apply(this, arguments) : undefined; + + if (this.inputs) { + for (const input of this.inputs) { + if (input.widget) { + const w = this.widgets.find((w) => w.name === input.widget.name); + if (w) { + hideWidget(this, w); + } else { + convertToWidget(this, input) + } + } + } + } + + return r; + }; + + function isNodeAtPos(pos) { + for (const n of app.graph._nodes) { + if (n.pos[0] === pos[0] && n.pos[1] === pos[1]) { + return true; + } + } + return false; + } + + // Double click a widget input to automatically attach a primitive + const origOnInputDblClick = nodeType.prototype.onInputDblClick; + const ignoreDblClick = Symbol(); + nodeType.prototype.onInputDblClick = function (slot) { + const r = origOnInputDblClick ? origOnInputDblClick.apply(this, arguments) : undefined; + + const input = this.inputs[slot]; + if (input.widget && !input[ignoreDblClick]) { + const node = LiteGraph.createNode("PrimitiveNode"); + app.graph.add(node); + + // Calculate a position that wont directly overlap another node + const pos = [this.pos[0] - node.size[0] - 30, this.pos[1]]; + while (isNodeAtPos(pos)) { + pos[1] += LiteGraph.NODE_TITLE_HEIGHT; + } + + node.pos = pos; + node.connect(0, this, slot); + node.title = input.name; + + // Prevent adding duplicates due to triple clicking + input[ignoreDblClick] = true; + setTimeout(() => { + delete input[ignoreDblClick]; + }, 300); + } + + return r; + }; + }, + registerCustomNodes() { + class PrimitiveNode { + constructor() { + this.addOutput("connect to widget input", "*"); + this.serialize_widgets = true; + this.isVirtualNode = true; + } + + applyToGraph() { + if (!this.outputs[0].links?.length) return; + + // For each output link copy our value over the original widget value + for (const l of this.outputs[0].links) { + const linkInfo = app.graph.links[l]; + const node = this.graph.getNodeById(linkInfo.target_id); + const input = node.inputs[linkInfo.target_slot]; + const widgetName = input.widget.name; + if (widgetName) { + const widget = node.widgets.find((w) => w.name === widgetName); + if (widget) { + widget.value = this.widgets[0].value; + if (widget.callback) { + widget.callback(widget.value, app.canvas, node, app.canvas.graph_mouse, {}); + } + } + } + } + } + + onConnectionsChange(_, index, connected) { + if (connected) { + if (this.outputs[0].links?.length) { + if (!this.widgets?.length) { + this.#onFirstConnection(); + } + if (!this.widgets?.length && this.outputs[0].widget) { + // On first load it often cant recreate the widget as the other node doesnt exist yet + // Manually recreate it from the output info + this.#createWidget(this.outputs[0].widget.config); + } + } + } else if (!this.outputs[0].links?.length) { + this.#onLastDisconnect(); + } + } + + onConnectOutput(slot, type, input, target_node, target_slot) { + // Fires before the link is made allowing us to reject it if it isn't valid + + // No widget, we cant connect + if (!input.widget) return false; + + if (this.outputs[slot].links?.length) { + return this.#isValidConnection(input); + } + } + + #onFirstConnection() { + // First connection can fire before the graph is ready on initial load so random things can be missing + const linkId = this.outputs[0].links[0]; + const link = this.graph.links[linkId]; + if (!link) return; + + const theirNode = this.graph.getNodeById(link.target_id); + if (!theirNode || !theirNode.inputs) return; + + const input = theirNode.inputs[link.target_slot]; + if (!input) return; + + const widget = input.widget; + const { type, linkType } = getWidgetType(widget.config); + + // Update our output to restrict to the widget type + this.outputs[0].type = linkType; + this.outputs[0].name = type; + this.outputs[0].widget = widget; + + this.#createWidget(widget.config, theirNode, widget.name); + } + + #createWidget(inputData, node, widgetName) { + let type = inputData[0]; + + if (type instanceof Array) { + type = "COMBO"; + } + + let widget; + if (type in ComfyWidgets) { + widget = (ComfyWidgets[type](this, "value", inputData, app) || {}).widget; + } else { + widget = this.addWidget(type, "value", null, () => {}, {}); + } + + if (node?.widgets && widget) { + const theirWidget = node.widgets.find((w) => w.name === widgetName); + if (theirWidget) { + widget.value = theirWidget.value; + } + } + + if (widget.type === "number") { + addRandomizeWidget(this, widget, "Random after every gen"); + } + + // When our value changes, update other widgets to reflect our changes + // e.g. so LoadImage shows correct image + const callback = widget.callback; + const self = this; + widget.callback = function () { + const r = callback ? callback.apply(this, arguments) : undefined; + self.applyToGraph(); + return r; + }; + + // Grow our node if required + const sz = this.computeSize(); + if (this.size[0] < sz[0]) { + this.size[0] = sz[0]; + } + if (this.size[1] < sz[1]) { + this.size[1] = sz[1]; + } + + requestAnimationFrame(() => { + if (this.onResize) { + this.onResize(this.size); + } + }); + } + + #isValidConnection(input) { + // Only allow connections where the configs match + const config1 = this.outputs[0].widget.config; + const config2 = input.widget.config; + + if (config1[0] !== config2[0]) return false; + + for (const k in config1[1]) { + if (k !== "default") { + if (config1[1][k] !== config2[1][k]) { + return false; + } + } + } + + return true; + } + + #onLastDisconnect() { + // We cant remove + re-add the output here as if you drag a link over the same link + // it removes, then re-adds, causing it to break + this.outputs[0].type = "*"; + this.outputs[0].name = "connect to widget input"; + delete this.outputs[0].widget; + + if (this.widgets) { + // Allow widgets to cleanup + for (const w of this.widgets) { + if (w.onRemove) { + w.onRemove(); + } + } + this.widgets.length = 0; + } + } + } + + LiteGraph.registerNodeType( + "PrimitiveNode", + Object.assign(PrimitiveNode, { + title: "Primitive", + }) + ); + PrimitiveNode.category = "utils"; + }, +}); diff --git a/web/index.html b/web/index.html index cd95594f..86156a7f 100644 --- a/web/index.html +++ b/web/index.html @@ -1,10 +1,11 @@ + + -