257 lines
8.2 KiB
Python
257 lines
8.2 KiB
Python
# "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 = "svg_to_png"
|
|
API_URL = "https://augmentstudio.app/api"
|
|
|
|
UPLOAD_DIR = folder_paths.get_input_directory()
|
|
|
|
|
|
class AugmentSVGToPNG:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
files = sorted([
|
|
f for f in os.listdir(UPLOAD_DIR)
|
|
if f.lower().endswith(".svg")
|
|
])
|
|
return {
|
|
"required": {
|
|
"svg_file": (files, {"image_upload": True}),
|
|
"api_key": ("STRING", {"default": ""}),
|
|
},
|
|
"optional": {
|
|
"trigger": ("TRIGGER",),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE", "MASK", "TRIGGER")
|
|
RETURN_NAMES = ("image", "alpha", "trigger")
|
|
FUNCTION = "execute"
|
|
CATEGORY = "Augment/Vector"
|
|
|
|
@classmethod
|
|
def IS_CHANGED(cls, svg_file="", **kwargs):
|
|
path = os.path.join(UPLOAD_DIR, svg_file)
|
|
if os.path.isfile(path):
|
|
return os.path.getmtime(path)
|
|
return float("nan")
|
|
|
|
@classmethod
|
|
def VALIDATE_INPUTS(cls, svg_file="", **kwargs):
|
|
path = os.path.join(UPLOAD_DIR, svg_file)
|
|
if not os.path.isfile(path):
|
|
return f"File not found: {svg_file}"
|
|
return True
|
|
|
|
def execute(self, svg_file, api_key, trigger=None):
|
|
api_url = API_URL.rstrip("/")
|
|
auth = {"Authorization": f"Bearer {api_key}"}
|
|
|
|
path = os.path.join(UPLOAD_DIR, svg_file)
|
|
with open(path, "rb") as f:
|
|
svg_bytes = f.read()
|
|
|
|
print(f"[Augment API] Submitting SVG: {svg_file} ({len(svg_bytes)} bytes)")
|
|
try:
|
|
r = requests.post(
|
|
f"{api_url}/process",
|
|
files={"image": (svg_file, svg_bytes, "image/svg+xml")},
|
|
data={"node_id": NODE_ID},
|
|
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 = 60
|
|
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("RGBA")
|
|
arr = np.array(img).astype(np.float32) / 255.0
|
|
|
|
image_tensor = torch.from_numpy(arr[:, :, :3]).unsqueeze(0)
|
|
mask_tensor = torch.from_numpy(arr[:, :, 3]).unsqueeze(0)
|
|
|
|
print(f"[Augment API] Done! {img.size[0]}x{img.size[1]}")
|
|
return (image_tensor, mask_tensor, "done")
|
|
|
|
class AugmentPNGToSVG:
|
|
@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": ""}),
|
|
},
|
|
"optional": {
|
|
"image": ("IMAGE",),
|
|
"image_file": (files, {"image_upload": True}),
|
|
"trigger": ("TRIGGER",),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("SVG", "AUGMENT_JSON", "TRIGGER")
|
|
RETURN_NAMES = ("svg", "json", "trigger")
|
|
FUNCTION = "execute"
|
|
CATEGORY = "Augment/Vector"
|
|
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, image=None, image_file=None, trigger=None):
|
|
api_url = API_URL.rstrip("/")
|
|
auth = {"Authorization": f"Bearer {api_key}"}
|
|
|
|
if image is not None:
|
|
from PIL import Image as PILImage
|
|
img_np = (image[0].cpu().numpy() * 255).astype(np.uint8)
|
|
pil_img = PILImage.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: {source_name} ({len(img_bytes)} bytes)")
|
|
try:
|
|
r = requests.post(
|
|
f"{api_url}/process",
|
|
files={"image": (source_name, img_bytes, mime)},
|
|
data={"node_id": "png_to_svg"},
|
|
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 = 60
|
|
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")
|
|
|
|
svg_r = requests.get(
|
|
f"{api_url}/job/{request_id}/image",
|
|
headers=auth, timeout=60,
|
|
)
|
|
if svg_r.status_code != 200:
|
|
raise RuntimeError(f"SVG fetch error: {svg_r.status_code}")
|
|
|
|
from comfy_api.latest._util.image_types import SVG
|
|
svg_bytes = io.BytesIO(svg_r.content)
|
|
size = svg_bytes.getbuffer().nbytes
|
|
print(f"[Augment API] Done! Received SVG ({size} bytes)")
|
|
json_result = json.dumps({"node": "AugmentPNGToSVG", "size_bytes": size})
|
|
return {"ui": {"text": [svg_r.text[:500]]}, "result": (SVG([svg_bytes]), json_result, "done")}
|
|
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"AugmentSVGToPNG": AugmentSVGToPNG,
|
|
"AugmentPNGToSVG": AugmentPNGToSVG,
|
|
}
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"AugmentSVGToPNG": "Augment SVG to PNG",
|
|
"AugmentPNGToSVG": "Augment PNG to SVG",
|
|
}
|
|
|