From 6c6f219d43347cba1893dcb76631f442183690e4 Mon Sep 17 00:00:00 2001 From: Aaron Burton Date: Wed, 18 Mar 2026 16:38:55 -0700 Subject: [PATCH] initial commit --- LICENSE.MD | 9 +++ README.md | 26 +++++++ __init__.py | 44 +++++++++++ logo_upscale.py | 164 ++++++++++++++++++++++++++++++++++++++++ pyproject.toml | 21 +++++ trigger.py | 99 ++++++++++++++++++++++++ web/js/augment_style.js | 65 ++++++++++++++++ 7 files changed, 428 insertions(+) create mode 100644 LICENSE.MD create mode 100644 README.md create mode 100644 __init__.py create mode 100644 logo_upscale.py create mode 100644 pyproject.toml create mode 100644 trigger.py create mode 100644 web/js/augment_style.js diff --git a/LICENSE.MD b/LICENSE.MD new file mode 100644 index 0000000..7286d3e --- /dev/null +++ b/LICENSE.MD @@ -0,0 +1,9 @@ + + +Copyright 2026 Augment Studio - Upscale Pro + +Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the “Software”), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED “AS IS”, WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. diff --git a/README.md b/README.md new file mode 100644 index 0000000..542958a --- /dev/null +++ b/README.md @@ -0,0 +1,26 @@ +# Augment Upscale Pro + +Make your logos and brand assets print-ready. Augment Upscale Pro is a ComfyUI node that takes rasterized images and upscales them 2X while keeping lines crisp and colors clean. No more blurry enlargements. A perfect first step to cleaning up assets before vectorizing them. + +Drop in a logo, icon, or any flat artwork and get back a sharp, high-resolution version you can actually use. + +## Getting started + +1. Grab an API key at [augmentstudio.app/pricing](https://augmentstudio.app/pricing) +2. Install the node from ComfyUI Manager (search "augment Upscale Pro") +3. Add the **Upscale Pro** node to your workflow, paste in your key, and connect an image + +## Also included + +**Trigger** flow nodes for wiring up multi-step workflows. Use them to chain the upscaler with other nodes so everything runs in the right order. + +## Manual install + +``` +cd ComfyUI/custom_nodes +git clone https://github.com/augment-lib/augment-upscale-pro.git augment-upscale-pro +``` + +## License + +See [LICENSE.md](LICENSE.MD) for details. diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..4f89669 --- /dev/null +++ b/__init__.py @@ -0,0 +1,44 @@ +""" +Augment - Upscale Pro +""" + +NODE_CLASS_MAPPINGS = {} +NODE_DISPLAY_NAME_MAPPINGS = {} +_failed = [] + +# ── Flow Control ── +try: + from .trigger import GateNode, make_wait_node, WAIT_TYPES + NODE_CLASS_MAPPINGS["GateNode"] = GateNode + NODE_DISPLAY_NAME_MAPPINGS["GateNode"] = "Trigger" + for _name, _type_str in WAIT_TYPES.items(): + _cls_name = f"AugmentWait{_name}" + try: + NODE_CLASS_MAPPINGS[_cls_name] = make_wait_node(_name, _type_str) + NODE_DISPLAY_NAME_MAPPINGS[_cls_name] = f"{_name} (Improved)" + except Exception as e: + _failed.append((_cls_name, e)) + print(f"[augment-upscale] ⚠ {_cls_name} unavailable: {e}") +except Exception as e: + _failed.append(("Wait nodes", e)) + print(f"[augment-upscale] ⚠ Flow control nodes unavailable: {e}") + +# ── Logo Upscale (Paid) ── +try: + from .logo_upscale import AugmentLogoUpscale + NODE_CLASS_MAPPINGS["AugmentLogoUpscale"] = AugmentLogoUpscale + NODE_DISPLAY_NAME_MAPPINGS["AugmentLogoUpscale"] = "Upscale Pro" +except Exception as e: + _failed.append(("AugmentLogoUpscale", e)) + print(f"[augment-upscale] ⚠ AugmentLogoUpscale unavailable: {e}") + +# ── Summary ── +WEB_DIRECTORY = "./web/js" +__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS", "WEB_DIRECTORY"] + +if _failed: + print(f"[augment-upscale] ✓ Registered {len(NODE_CLASS_MAPPINGS)} nodes ({len(_failed)} failed)") + for node_id, err in _failed: + print(f"[augment-upscale] ✗ {node_id}: {err}") +else: + print(f"[augment-upscale] ✓ Registered {len(NODE_CLASS_MAPPINGS)} nodes") diff --git a/logo_upscale.py b/logo_upscale.py new file mode 100644 index 0000000..8ffdee2 --- /dev/null +++ b/logo_upscale.py @@ -0,0 +1,164 @@ +# "Design-focused node suite for ComfyUI." +# Copyright 2026 Augment Studio +# Paid Node, credit values may change. +# Augmentstudio.app + +import io +import json +import os +import requests +import time +import folder_paths +import numpy as np +import torch +from PIL import Image + + +NODE_ID = "logo_upscale" +API_URL = "https://augmentstudio.app/api" + +UPLOAD_DIR = folder_paths.get_input_directory() + + +class AugmentLogoUpscale: + @classmethod + def INPUT_TYPES(cls): + files = sorted([ + f for f in os.listdir(UPLOAD_DIR) + if f.lower().endswith((".png", ".jpg", ".jpeg", ".webp")) + ]) + return { + "required": { + "api_key": ("STRING", {"default": ""}), + "scale": (["2x"], {"default": "2x"}), + }, + "optional": { + "image": ("IMAGE",), + "image_file": (files, {"image_upload": True}), + "trigger": ("TRIGGER",), + }, + } + + RETURN_TYPES = ("IMAGE", "AUGMENT_JSON", "TRIGGER") + RETURN_NAMES = ("image", "json_result", "trigger") + FUNCTION = "execute" + CATEGORY = "Augment/Enchance" + OUTPUT_NODE = True + + @classmethod + def IS_CHANGED(cls, **kwargs): + return time.time() + + @classmethod + def VALIDATE_INPUTS(cls, **kwargs): + return True + + def execute(self, api_key, scale="2x", image=None, image_file=None, trigger=None): + api_url = API_URL.rstrip("/") + auth = {"Authorization": f"Bearer {api_key}"} + scale_factor = int(scale.replace("x", "")) + + if image is not None: + img_np = (image[0].cpu().numpy() * 255).astype(np.uint8) + pil_img = Image.fromarray(img_np, "RGB") + buf = io.BytesIO() + pil_img.save(buf, format="PNG") + img_bytes = buf.getvalue() + source_name = "image_input.png" + mime = "image/png" + elif image_file: + path = os.path.join(UPLOAD_DIR, image_file) + with open(path, "rb") as f: + img_bytes = f.read() + source_name = image_file + ext = os.path.splitext(image_file)[1].lower() + mime_types = {".png": "image/png", ".jpg": "image/jpeg", ".jpeg": "image/jpeg", ".webp": "image/webp"} + mime = mime_types.get(ext, "image/png") + else: + raise RuntimeError("No image provided — connect an image input or select a file") + + print(f"[Augment API] Submitting image for {scale} upscale: {source_name} ({len(img_bytes)} bytes)") + try: + r = requests.post( + f"{api_url}/process", + files={"image": (source_name, img_bytes, mime)}, + data={"node_id": NODE_ID, "scale": scale_factor}, + headers=auth, + timeout=30, + ) + except requests.exceptions.RequestException as e: + print(f"[Augment API] Submit failed: {e}") + raise + + if r.status_code != 200: + raise RuntimeError(f"Submit error {r.status_code}: {r.text[:500]}") + + request_id = r.json().get("request_id") + if not request_id: + raise RuntimeError(f"No request_id: {r.text[:500]}") + + print(f"[Augment API] Job submitted: {request_id}") + + max_wait = 120 + elapsed = 0 + job_status = "unknown" + + while elapsed < max_wait: + time.sleep(1) + elapsed += 1 + try: + status_r = requests.get( + f"{api_url}/job/{request_id}/status", + headers=auth, timeout=10, + ) + job_status = status_r.json().get("status", "unknown") + except Exception as e: + print(f"[Augment API] Poll error: {e}") + continue + + if job_status == "done": + break + elif job_status == "error": + raise RuntimeError(f"Job failed: {status_r.json().get('error')}") + + if job_status != "done": + raise RuntimeError(f"Timed out after {max_wait}s") + + img_r = requests.get( + f"{api_url}/job/{request_id}/image", + headers=auth, timeout=60, + ) + if img_r.status_code != 200: + raise RuntimeError(f"Image fetch error: {img_r.status_code}") + + img = Image.open(io.BytesIO(img_r.content)).convert("RGB") + arr = np.array(img).astype(np.float32) / 255.0 + + image_tensor = torch.from_numpy(arr).unsqueeze(0) + + json_result = json.dumps({ + "node": "LogoUpscale", + "scale": scale, + "width": img.size[0], + "height": img.size[1], + "source": source_name, + }, indent=2) + + temp_dir = folder_paths.get_temp_directory() + os.makedirs(temp_dir, exist_ok=True) + preview_file = f"augment_upscale_{time.time_ns()}.png" + img.save(os.path.join(temp_dir, preview_file)) + + print(f"[Augment API] Done! {img.size[0]}x{img.size[1]} ({scale} upscale)") + return { + "ui": {"images": [{"filename": preview_file, "subfolder": "", "type": "temp"}]}, + "result": (image_tensor, json_result, "done"), + } + + +NODE_CLASS_MAPPINGS = { + "AugmentLogoUpscale": AugmentLogoUpscale, +} +NODE_DISPLAY_NAME_MAPPINGS = { + "AugmentLogoUpscale": "Logo Upscale", +} diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 0000000..48f3769 --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,21 @@ +[project] +name = "augment-UpscalePro" +description = "✨ Augment Upscale — AI-powered logo and image upscaler for ComfyUI. Upscale logos, icons, and brand assets to high resolution with sharp, clean results. A professional alternative to Magnific AI built directly into your ComfyUI workflow. Supports 2x and 4x upscaling via the Augment Studio API with automatic job polling and tensor output. Ideal for logo design, brand asset pipelines, print-ready artwork, and any workflow that needs crisp detail-preserving enlargement. Includes built-in flow control trigger and wait nodes for chaining multi-step pipelines. Get your API key at [a/augmentstudio.app/pricing](https://augmentstudio.app/pricing)" +version = "1.0.0" +licence = { file = "LICENSE.md" } + +dependencies = [ + "numpy>=1.19.0", + "torch>=1.9.0", + "Pillow>=8.0.0", + "requests>=2.25.0" +] + +[project.urls] +Repository = "https://github.com/augment-lib/augment-upscale-pro" +Website = "https://augmentstudio.app/" + +[tool.comfy] +PublisherId = "augmentstudio" +DisplayName = "✨ augment Upscale Pro - AI Image & Logo Upscaler, Enhancer and Enlarger for ComfyUI" +Icon = "https://avatars.githubusercontent.com/u/243209371?s=400&u=db6e0d310dcf8863dd5c735d3ebbdaa5c317530d&v=4" diff --git a/trigger.py b/trigger.py new file mode 100644 index 0000000..fc6fb60 --- /dev/null +++ b/trigger.py @@ -0,0 +1,99 @@ +# "Design-focused node suite for ComfyUI." +# Copyright 2026 Augment Studio +# Augmentstudio.app + +import json + + +class GateNode: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "input": ("*",), + }, + } + + RETURN_TYPES = ("TRIGGER", "*", "STRING") + RETURN_NAMES = ("trigger", "value", "json") + OUTPUT_NODE = True + FUNCTION = "execute" + CATEGORY = "Augment/Utils" + + @classmethod + def VALIDATE_INPUTS(cls, **kwargs): + return True + + def execute(self, input): + json_result = json.dumps({"node": "Convert Any", "type": type(input).__name__, "value": str(input)}) + return {"ui": {"text": [str(input)]}, "result": ("done", input, json_result)} + + +WAIT_DEFAULTS = { + "INT": {"default": 0}, + "FLOAT": {"default": 0.0, "step": 0.01}, + "STRING": {"default": ""}, + "BOOLEAN": {"default": False}, +} + + +def make_wait_node(type_name, type_str): + class WaitNode: + @classmethod + def INPUT_TYPES(cls): + has_widget = type_str in WAIT_DEFAULTS + inputs = { + "required": {}, + "optional": { + "trigger": ("TRIGGER",), + }, + } + if has_widget: + inputs["required"]["value"] = (type_str, WAIT_DEFAULTS[type_str]) + inputs["optional"]["input"] = (type_str, {"forceInput": True}) + else: + inputs["required"]["value"] = (type_str,) + return inputs + + RETURN_TYPES = (type_str, "AUGMENT_JSON", "TRIGGER") + RETURN_NAMES = ("value", "json", "trigger") + OUTPUT_NODE = True + FUNCTION = "execute" + CATEGORY = "Augment/Utils" + + @classmethod + def VALIDATE_INPUTS(cls, **kwargs): + return True + + def execute(self, value=None, trigger=None, input=None): + v = input if input is not None else value + json_result = json.dumps({"node": type_str, "type": type_str.lower(), "value": str(v)}) + return {"ui": {"text": [str(v)]}, "result": (v, json_result, "done")} + + WaitNode.__name__ = f"Wait{type_name}Node" + WaitNode.__qualname__ = f"Wait{type_name}Node" + return WaitNode + + +WAIT_TYPES = { + "Int": "INT", + "Float": "FLOAT", + "String": "STRING", + "Bool": "BOOLEAN", + "Image": "IMAGE", + "Latent": "LATENT", + "Condition": "CONDITIONING", + "VAE": "VAE", + "Model": "MODEL", + "Clip": "CLIP", + "Mask": "MASK", +} + +NODE_CLASS_MAPPINGS = {"GateNode": GateNode} +NODE_DISPLAY_NAME_MAPPINGS = {"GateNode": "Trigger"} + +for name, type_str in WAIT_TYPES.items(): + cls_name = f"Wait{name}Node" + node_cls = make_wait_node(name, type_str) + NODE_CLASS_MAPPINGS[cls_name] = node_cls + NODE_DISPLAY_NAME_MAPPINGS[cls_name] = f"{name} (Improved)" diff --git a/web/js/augment_style.js b/web/js/augment_style.js new file mode 100644 index 0000000..ba3918c --- /dev/null +++ b/web/js/augment_style.js @@ -0,0 +1,65 @@ +/*"Augment Upscale — AI Image & Logo Upscaler for ComfyUI." +# Copyright 2026 Augment Studio +Augmentstudio.app */ + +import { app } from "../../../scripts/app.js"; + +const TITLE_BG = "#443322"; +const NODE_BG = "#665533"; + +const PAID_NODES = [ + "AugmentLogoUpscale", +]; + +const BADGE_FONT = "bold 9px sans-serif"; +const BADGE_RADIUS = 4; +const BADGE_PAD_X = 6; + +app.registerExtension({ + name: "augment-upscale.NodeStyle", + + async beforeRegisterNodeDef(nodeType, nodeData, app) { + if (!PAID_NODES.includes(nodeData.name)) return; + + const origCreated = nodeType.prototype.onNodeCreated; + nodeType.prototype.onNodeCreated = function () { + origCreated?.apply(this, arguments); + this.addWidget("button", "Get Tokens", null, () => { + window.open("https://augmentstudio.app/pricing", "_blank"); + }); + + const jsonIdx = this.outputs?.findIndex(o => o.name === "json_result"); + if (jsonIdx >= 0) { + this.outputs.splice(jsonIdx, 1); + } + }; + + const origDraw = nodeType.prototype.onDrawForeground; + nodeType.prototype.onDrawForeground = function (ctx) { + this.color = TITLE_BG; + this.bgcolor = NODE_BG; + + origDraw?.apply(this, arguments); + + ctx.save(); + ctx.font = BADGE_FONT; + + const creditsText = "🔸 1 export / run"; + const creditsW = ctx.measureText(creditsText).width + BADGE_PAD_X * 2; + const creditsX = this.size[0] - creditsW - 6; + const y = -22; + + ctx.beginPath(); + ctx.roundRect(creditsX, y, creditsW, 14, BADGE_RADIUS); + ctx.fillStyle = "#8d6932"; + ctx.fill(); + + ctx.fillStyle = "#ffffff"; + ctx.textAlign = "left"; + ctx.textBaseline = "middle"; + ctx.fillText(creditsText, creditsX + BADGE_PAD_X, y + 7); + + ctx.restore(); + }; + }, +});