322 lines
12 KiB
Python
322 lines
12 KiB
Python
import json
|
|
import math
|
|
from pathlib import Path
|
|
|
|
import numpy as np
|
|
import torch
|
|
|
|
try:
|
|
from aiohttp import web
|
|
from server import PromptServer
|
|
except Exception: # pragma: no cover
|
|
web = None
|
|
PromptServer = None
|
|
|
|
|
|
ROOT_DIR = Path(__file__).resolve().parent
|
|
PRESETS_DIR = ROOT_DIR / "presets_gradients"
|
|
PRESETS_DIR.mkdir(parents=True, exist_ok=True)
|
|
|
|
DEFAULT_SETTINGS = {
|
|
"width": 1024,
|
|
"height": 1024,
|
|
"fill_mode": "solid",
|
|
"color_count": "1",
|
|
"color_1": "#FFFFFF",
|
|
"color_2": "#FFFFFF",
|
|
"color_3": "#FFFFFF",
|
|
"color_4": "#FFFFFF",
|
|
"invert": False,
|
|
"angle": 0.0,
|
|
"center_x": 0.5,
|
|
"center_y": 0.5,
|
|
"inner_border_px": 0,
|
|
"border_color": "#FFFFFF",
|
|
"border_opacity": 1.0,
|
|
}
|
|
|
|
|
|
def _safe_preset_name(value: str) -> str:
|
|
value = str(value or "").strip()
|
|
if not value:
|
|
return ""
|
|
forbidden = '<>:"/\\|?*'
|
|
for ch in forbidden:
|
|
value = value.replace(ch, "_")
|
|
value = value.replace("..", "_")
|
|
return value[:120].strip()
|
|
|
|
|
|
def _preset_path(name: str) -> Path:
|
|
safe = _safe_preset_name(name)
|
|
if not safe:
|
|
raise ValueError("Preset name is empty")
|
|
return PRESETS_DIR / f"{safe}.json"
|
|
|
|
|
|
def _normalize_hex(value, fallback="#FFFFFF"):
|
|
value = str(value or "").strip()
|
|
if not value.startswith("#"):
|
|
value = "#" + value
|
|
raw = value[1:]
|
|
try:
|
|
if len(raw) == 3:
|
|
raw = "".join(ch * 2 for ch in raw)
|
|
if len(raw) != 6:
|
|
raise ValueError("invalid hex length")
|
|
int(raw, 16)
|
|
return f"#{raw.upper()}"
|
|
except Exception:
|
|
return fallback.upper()
|
|
|
|
|
|
def _hex_to_rgb01(value, fallback="#FFFFFF"):
|
|
value = _normalize_hex(value, fallback)
|
|
raw = value[1:]
|
|
return np.array([
|
|
int(raw[0:2], 16),
|
|
int(raw[2:4], 16),
|
|
int(raw[4:6], 16),
|
|
], dtype=np.float32) / 255.0
|
|
|
|
|
|
def _active_colors(color_count, color_1, color_2, color_3, color_4):
|
|
try:
|
|
count = int(str(color_count).strip())
|
|
except Exception:
|
|
count = 2
|
|
count = max(1, min(count, 4))
|
|
|
|
colors_hex = [
|
|
_normalize_hex(color_1, "#FFFFFF"),
|
|
_normalize_hex(color_2, "#FF7A18"),
|
|
_normalize_hex(color_3, "#FFFFFF"),
|
|
_normalize_hex(color_4, "#111827"),
|
|
][:count]
|
|
|
|
colors_rgb = np.stack([_hex_to_rgb01(c) for c in colors_hex], axis=0)
|
|
return colors_hex, colors_rgb
|
|
|
|
|
|
def _interpolate_stops(t, colors_rgb):
|
|
t = np.clip(t.astype(np.float32), 0.0, 1.0)
|
|
if len(colors_rgb) == 1:
|
|
return np.broadcast_to(colors_rgb[0].reshape(1, 1, 3), (*t.shape, 3)).copy().astype(np.float32)
|
|
|
|
scaled = t * (len(colors_rgb) - 1)
|
|
idx = np.floor(scaled).astype(np.int32)
|
|
idx = np.clip(idx, 0, len(colors_rgb) - 2)
|
|
frac = (scaled - idx).astype(np.float32)
|
|
|
|
c0 = colors_rgb[idx]
|
|
c1 = colors_rgb[idx + 1]
|
|
return (c0 * (1.0 - frac[..., None]) + c1 * frac[..., None]).astype(np.float32)
|
|
|
|
|
|
class Orion4D_GradientFillV2:
|
|
MODES = ["solid", "linear", "radial", "angular", "reflected", "diamond"]
|
|
COLOR_COUNTS = ["1", "2", "3", "4"]
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"width": ("INT", {"default": DEFAULT_SETTINGS["width"], "min": 1, "max": 8192, "step": 1}),
|
|
"height": ("INT", {"default": DEFAULT_SETTINGS["height"], "min": 1, "max": 8192, "step": 1}),
|
|
"fill_mode": (cls.MODES, {"default": DEFAULT_SETTINGS["fill_mode"]}),
|
|
"color_count": (cls.COLOR_COUNTS, {"default": DEFAULT_SETTINGS["color_count"]}),
|
|
"color_1": ("STRING", {"default": DEFAULT_SETTINGS["color_1"]}),
|
|
"color_2": ("STRING", {"default": DEFAULT_SETTINGS["color_2"]}),
|
|
"color_3": ("STRING", {"default": DEFAULT_SETTINGS["color_3"]}),
|
|
"color_4": ("STRING", {"default": DEFAULT_SETTINGS["color_4"]}),
|
|
"invert": ("BOOLEAN", {"default": DEFAULT_SETTINGS["invert"]}),
|
|
"angle": ("FLOAT", {"default": DEFAULT_SETTINGS["angle"], "min": 0.0, "max": 360.0, "step": 1.0}),
|
|
"center_x": ("FLOAT", {"default": DEFAULT_SETTINGS["center_x"], "min": 0.0, "max": 1.0, "step": 0.01}),
|
|
"center_y": ("FLOAT", {"default": DEFAULT_SETTINGS["center_y"], "min": 0.0, "max": 1.0, "step": 0.01}),
|
|
"inner_border_px": ("INT", {"default": DEFAULT_SETTINGS["inner_border_px"], "min": 0, "max": 200, "step": 1}),
|
|
"border_color": ("STRING", {"default": DEFAULT_SETTINGS["border_color"]}),
|
|
"border_opacity": ("FLOAT", {"default": DEFAULT_SETTINGS["border_opacity"], "min": 0.0, "max": 1.0, "step": 0.01}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE", "MASK", "MASK", "STRING")
|
|
RETURN_NAMES = ("image", "full_mask", "border_mask", "settings_json")
|
|
FUNCTION = "generate"
|
|
CATEGORY = "Orion4D_Layer"
|
|
|
|
def generate(
|
|
self,
|
|
width,
|
|
height,
|
|
fill_mode,
|
|
color_count,
|
|
color_1,
|
|
color_2,
|
|
color_3,
|
|
color_4,
|
|
invert,
|
|
angle,
|
|
center_x,
|
|
center_y,
|
|
inner_border_px,
|
|
border_color,
|
|
border_opacity,
|
|
):
|
|
w = max(1, int(width))
|
|
h = max(1, int(height))
|
|
mode = fill_mode if fill_mode in self.MODES else "linear"
|
|
|
|
colors_hex, colors_rgb = _active_colors(color_count, color_1, color_2, color_3, color_4)
|
|
border_color_hex = _normalize_hex(border_color, "#FFFFFF")
|
|
border_rgb = _hex_to_rgb01(border_color_hex)
|
|
|
|
cx = float(max(0.0, min(float(center_x), 1.0)))
|
|
cy = float(max(0.0, min(float(center_y), 1.0)))
|
|
theta = math.radians(float(angle) % 360.0)
|
|
border = int(max(0, min(int(inner_border_px), 200)))
|
|
opacity = float(max(0.0, min(float(border_opacity), 1.0)))
|
|
inv = bool(invert)
|
|
|
|
yy, xx = np.mgrid[0:h, 0:w].astype(np.float32)
|
|
nx = xx / (w - 1) if w > 1 else np.zeros((h, w), dtype=np.float32)
|
|
ny = yy / (h - 1) if h > 1 else np.zeros((h, w), dtype=np.float32)
|
|
dx = nx - cx
|
|
dy = ny - cy
|
|
|
|
if mode == "solid" or len(colors_rgb) == 1:
|
|
img = np.broadcast_to(colors_rgb[0].reshape(1, 1, 3), (h, w, 3)).copy().astype(np.float32)
|
|
else:
|
|
vx = math.cos(theta)
|
|
vy = math.sin(theta)
|
|
|
|
if mode == "linear":
|
|
proj = nx * vx + ny * vy
|
|
mn = float(proj.min())
|
|
mx = float(proj.max())
|
|
t = (proj - mn) / max(mx - mn, 1e-6)
|
|
elif mode == "reflected":
|
|
proj = dx * vx + dy * vy
|
|
t = np.abs(proj) / max(float(np.max(np.abs(proj))), 1e-6)
|
|
elif mode == "radial":
|
|
dist = np.sqrt(dx * dx + dy * dy)
|
|
corners = np.array([[0.0 - cx, 0.0 - cy], [1.0 - cx, 0.0 - cy], [0.0 - cx, 1.0 - cy], [1.0 - cx, 1.0 - cy]], dtype=np.float32)
|
|
t = dist / max(float(np.sqrt((corners * corners).sum(axis=1)).max()), 1e-6)
|
|
elif mode == "diamond":
|
|
dist = np.abs(dx) + np.abs(dy)
|
|
corners = np.array([[0.0 - cx, 0.0 - cy], [1.0 - cx, 0.0 - cy], [0.0 - cx, 1.0 - cy], [1.0 - cx, 1.0 - cy]], dtype=np.float32)
|
|
t = dist / max(float((np.abs(corners[:, 0]) + np.abs(corners[:, 1])).max()), 1e-6)
|
|
elif mode == "angular":
|
|
ang = np.arctan2(dy, dx) - theta
|
|
t = np.mod(ang, 2.0 * math.pi) / (2.0 * math.pi)
|
|
else:
|
|
t = np.zeros((h, w), dtype=np.float32)
|
|
|
|
t = np.clip(t, 0.0, 1.0).astype(np.float32)
|
|
if inv:
|
|
t = 1.0 - t
|
|
img = _interpolate_stops(t, colors_rgb)
|
|
|
|
full_mask = np.ones((h, w), dtype=np.float32)
|
|
border_mask = np.zeros((h, w), dtype=np.float32)
|
|
|
|
if border > 0:
|
|
bx = min(border, w)
|
|
by = min(border, h)
|
|
border_mask[:by, :] = 1.0
|
|
border_mask[h - by:, :] = 1.0
|
|
border_mask[:, :bx] = 1.0
|
|
border_mask[:, w - bx:] = 1.0
|
|
if opacity > 0.0:
|
|
alpha = border_mask[..., None] * opacity
|
|
img = img * (1.0 - alpha) + border_rgb.reshape(1, 1, 3) * alpha
|
|
|
|
image_tensor = torch.from_numpy(np.ascontiguousarray(img.astype(np.float32))).unsqueeze(0)
|
|
full_mask_tensor = torch.from_numpy(full_mask).unsqueeze(0)
|
|
border_mask_tensor = torch.from_numpy(border_mask).unsqueeze(0)
|
|
|
|
settings = {
|
|
"width": w,
|
|
"height": h,
|
|
"fill_mode": mode,
|
|
"color_count": str(len(colors_hex)),
|
|
"color_1": colors_hex[0] if len(colors_hex) > 0 else DEFAULT_SETTINGS["color_1"],
|
|
"color_2": colors_hex[1] if len(colors_hex) > 1 else DEFAULT_SETTINGS["color_2"],
|
|
"color_3": colors_hex[2] if len(colors_hex) > 2 else DEFAULT_SETTINGS["color_3"],
|
|
"color_4": colors_hex[3] if len(colors_hex) > 3 else DEFAULT_SETTINGS["color_4"],
|
|
"invert": inv,
|
|
"angle": float(angle),
|
|
"center_x": cx,
|
|
"center_y": cy,
|
|
"inner_border_px": border,
|
|
"border_color": border_color_hex,
|
|
"border_opacity": opacity,
|
|
}
|
|
|
|
return (
|
|
image_tensor,
|
|
full_mask_tensor,
|
|
border_mask_tensor,
|
|
json.dumps(settings, ensure_ascii=False),
|
|
)
|
|
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"Orion4D_GradientFillV2": Orion4D_GradientFillV2,
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"Orion4D_GradientFillV2": "🎨 Fill / Gradient Generator v2",
|
|
}
|
|
|
|
|
|
if PromptServer is not None and web is not None:
|
|
@PromptServer.instance.routes.get("/orion4d/gradient_presets")
|
|
async def orion4d_gradient_presets_list(request):
|
|
presets = sorted(p.stem for p in PRESETS_DIR.glob("*.json"))
|
|
return web.json_response({"presets": presets, "directory": str(PRESETS_DIR)})
|
|
|
|
|
|
@PromptServer.instance.routes.get("/orion4d/gradient_presets/{name}")
|
|
async def orion4d_gradient_presets_get(request):
|
|
name = request.match_info.get("name", "")
|
|
try:
|
|
path = _preset_path(name)
|
|
except Exception as e:
|
|
return web.json_response({"error": str(e)}, status=400)
|
|
if not path.exists():
|
|
return web.json_response({"error": "Preset not found"}, status=404)
|
|
try:
|
|
data = json.loads(path.read_text(encoding="utf-8"))
|
|
except Exception as e:
|
|
return web.json_response({"error": f"Invalid preset file: {e}"}, status=500)
|
|
return web.json_response({"name": path.stem, "data": data})
|
|
|
|
|
|
@PromptServer.instance.routes.post("/orion4d/gradient_presets/save")
|
|
async def orion4d_gradient_presets_save(request):
|
|
payload = await request.json()
|
|
name = payload.get("name", "")
|
|
data = payload.get("data", {})
|
|
try:
|
|
path = _preset_path(name)
|
|
except Exception as e:
|
|
return web.json_response({"error": str(e)}, status=400)
|
|
merged = dict(DEFAULT_SETTINGS)
|
|
if isinstance(data, dict):
|
|
merged.update(data)
|
|
path.write_text(json.dumps(merged, ensure_ascii=False, indent=2), encoding="utf-8")
|
|
return web.json_response({"ok": True, "name": path.stem})
|
|
|
|
|
|
@PromptServer.instance.routes.post("/orion4d/gradient_presets/delete")
|
|
async def orion4d_gradient_presets_delete(request):
|
|
payload = await request.json()
|
|
name = payload.get("name", "")
|
|
try:
|
|
path = _preset_path(name)
|
|
except Exception as e:
|
|
return web.json_response({"error": str(e)}, status=400)
|
|
if path.exists():
|
|
path.unlink()
|
|
return web.json_response({"ok": True, "name": path.stem})
|