Snapshot: full project state
This commit is contained in:
153
cover_gen.py
Normal file
153
cover_gen.py
Normal file
@@ -0,0 +1,153 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Generate a cover image for a THREATMARKET intel pack via ComfyUI SDXL.
|
||||
|
||||
Zero-failure design: tries ComfyUI (juggernautXL, then dreamshaperXL, then
|
||||
sdxl-base), then falls back to a locally-rendered SVG->PNG-style placeholder
|
||||
(a real, deterministic dark radar-brand cover - not fake data, just brand art).
|
||||
Always writes /opt/threatmarket/static/covers/<pack_id>.png (or .svg) and
|
||||
returns the web path, or None only if the filesystem itself is unwritable.
|
||||
|
||||
Usage: from cover_gen import generate_cover
|
||||
path = generate_cover("chewy.com", "2026-09-25")
|
||||
"""
|
||||
import base64
|
||||
import json
|
||||
import os
|
||||
import urllib.request
|
||||
|
||||
COMFY = os.environ.get("COMFY_URL", "http://10.30.20.29:8188")
|
||||
COVER_DIR = os.environ.get("COVER_DIR", "/opt/threatmarket/static/covers")
|
||||
WEB_PREFIX = "/covers"
|
||||
# Order = preference. Lightning models = few steps = fast + reliable.
|
||||
CKPTS = ["flux1-schnell-fp8.safetensors",
|
||||
"juggernautXL_ragnarok.safetensors",
|
||||
"dreamshaperXL_lightningDPMSDE.safetensors"]
|
||||
TIMEOUT_SUBMIT = 20
|
||||
TIMEOUT_POLL = 150
|
||||
TIMEOUT_FETCH = 30
|
||||
|
||||
|
||||
def _prompt_text(domain, date):
|
||||
return ('dark hacker operations room at night, glowing green matrix rain code '
|
||||
'walls, multiple monitors showing world-map network dashboards and '
|
||||
'breach data, empty desk chair, dramatic cyan and green lighting, '
|
||||
'cinematic depth of field, gritty cyberpunk atmosphere, wide dark '
|
||||
'area at top for title text, professional threat-intelligence dossier '
|
||||
'cover art, no text, no letters, no words, 4k')
|
||||
|
||||
|
||||
def _workflow(ckpt, seed):
|
||||
return {
|
||||
"3": {"class_type": "KSampler", "inputs": {
|
||||
"seed": seed, "steps": 20, "cfg": 1.8, "sampler_name": "euler",
|
||||
"scheduler": "sgm_uniform", "denoise": 1.0,
|
||||
"model": ["4", 0], "positive": ["6", 0], "negative": ["7", 0],
|
||||
"latent_image": ["5", 0]}},
|
||||
"4": {"class_type": "CheckpointLoaderSimple", "inputs": {"ckpt_name": ckpt}},
|
||||
"5": {"class_type": "EmptyLatentImage", "inputs": {"width": 1024, "height": 640, "batch_size": 1}},
|
||||
"6": {"class_type": "CLIPTextEncode", "inputs": {"text": "placeholder", "clip": ["4", 1]}},
|
||||
"7": {"class_type": "CLIPTextEncode", "inputs": {"text": "text, watermark, blurry, low quality, oversaturated", "clip": ["4", 1]}},
|
||||
"8": {"class_type": "VAEDecode", "inputs": {"samples": ["3", 0], "vae": ["4", 2]}},
|
||||
"9": {"class_type": "SaveImage", "inputs": {"filename_prefix": "tmcover", "images": ["8", 0]}},
|
||||
}
|
||||
|
||||
|
||||
def _comfy_generate(domain, date, seed):
|
||||
"""Try each checkpoint; return PNG bytes or None."""
|
||||
for ckpt in CKPTS:
|
||||
try:
|
||||
wf = _workflow(ckpt, seed)
|
||||
# fix node 6 text properly (dict-literal scoping hack above)
|
||||
wf["6"]["inputs"]["text"] = _prompt_text(domain, date)
|
||||
req = urllib.request.Request(
|
||||
COMFY + "/prompt",
|
||||
data=json.dumps({"prompt": wf}).encode(),
|
||||
headers={"Content-Type": "application/json"})
|
||||
with urllib.request.urlopen(req, timeout=TIMEOUT_SUBMIT) as r:
|
||||
resp = json.loads(r.read())
|
||||
pid = resp.get("prompt_id")
|
||||
if not pid:
|
||||
continue
|
||||
import time
|
||||
deadline = time.time() + TIMEOUT_POLL
|
||||
while time.time() < deadline:
|
||||
with urllib.request.urlopen(COMFY + "/history/" + pid, timeout=15) as r:
|
||||
hist = json.loads(r.read())
|
||||
if pid in hist and hist[pid].get("outputs"):
|
||||
outs = hist[pid]["outputs"]
|
||||
for node in outs.values():
|
||||
for img in node.get("images", []):
|
||||
fn = img["filename"]
|
||||
sub = img.get("subfolder", "")
|
||||
imgtype = img.get("type", "output")
|
||||
url = f"{COMFY}/view?filename={fn}&subfolder={sub}&type={imgtype}"
|
||||
with urllib.request.urlopen(url, timeout=TIMEOUT_FETCH) as fr:
|
||||
return fr.read()
|
||||
time.sleep(3)
|
||||
except Exception:
|
||||
continue
|
||||
return None
|
||||
|
||||
|
||||
def _fallback_svg(domain, date):
|
||||
"""Deterministic branded SVG cover - always succeeds if fs writable."""
|
||||
seed = sum(ord(c) for c in domain) % 360
|
||||
svg = f"""<svg xmlns='http://www.w3.org/2000/svg' width='1024' height='640'>
|
||||
<defs><radialGradient id='g' cx='75%' cy='20%'><stop offset='0%' stop-color='hsl({seed},80%,22%)'/>
|
||||
<stop offset='100%' stop-color='#0a0e1a'/></radialGradient></defs>
|
||||
<rect width='1024' height='640' fill='url(#g)'/>
|
||||
<g stroke='hsl({seed},90%,60%)' stroke-opacity='0.25' fill='none'>"""
|
||||
for i in range(6):
|
||||
svg += f"<circle cx='780' cy='200' r='{40+i*45}' stroke-width='1'/>"
|
||||
svg += f"</g><g fill='hsl({seed},90%,65%)'>"
|
||||
import math
|
||||
for i in range(14):
|
||||
x = 120 + (i * 61) % 700
|
||||
y = 420 + (i * 37) % 160
|
||||
r = 2 + (i % 3)
|
||||
svg += f"<circle cx='{x}' cy='{y}' r='{r}' opacity='{0.35 + (i % 5) * 0.12}'/>"
|
||||
svg += f"""</g>
|
||||
<text x='48' y='90' font-family='monospace' font-size='30' fill='#e8f4ff' font-weight='bold'>INTEL PACK</text>
|
||||
<text x='48' y='130' font-family='monospace' font-size='22' fill='hsl({seed},90%,70%)'>{domain}</text>
|
||||
<text x='48' y='162' font-family='monospace' font-size='16' fill='#8fa3bf'>{date} — THREATMARKET</text>
|
||||
</svg>"""
|
||||
return svg.encode()
|
||||
|
||||
|
||||
def generate_cover(domain, date):
|
||||
"""Returns web path of cover image, or None. Never raises."""
|
||||
os.makedirs(COVER_DIR, exist_ok=True)
|
||||
pid = (domain + "-" + date.replace("-", "")).replace("/", "_")
|
||||
png_path = os.path.join(COVER_DIR, pid + ".png")
|
||||
svg_path = os.path.join(COVER_DIR, pid + ".svg")
|
||||
if os.path.exists(png_path):
|
||||
return f"{WEB_PREFIX}/{pid}.png"
|
||||
if os.path.exists(svg_path):
|
||||
return f"{WEB_PREFIX}/{pid}.svg"
|
||||
seed = sum(ord(c) for c in pid) % (2 ** 31)
|
||||
png = _comfy_generate(domain, date, seed)
|
||||
if png:
|
||||
try:
|
||||
with open(png_path, "wb") as f:
|
||||
f.write(png)
|
||||
try:
|
||||
from text_stamp import stamp
|
||||
stamp(png_path, domain)
|
||||
except Exception:
|
||||
pass # art still fine even if stamp fails
|
||||
return f"{WEB_PREFIX}/{pid}.png"
|
||||
except OSError:
|
||||
pass
|
||||
try:
|
||||
with open(svg_path, "wb") as f:
|
||||
f.write(_fallback_svg(domain, date))
|
||||
return f"{WEB_PREFIX}/{pid}.svg"
|
||||
except OSError:
|
||||
return None
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import sys
|
||||
d = sys.argv[1] if len(sys.argv) > 1 else "mozilla.org"
|
||||
dt = sys.argv[2] if len(sys.argv) > 2 else "2026-09-24"
|
||||
print(generate_cover(d, dt) or "FAILED")
|
||||
Reference in New Issue
Block a user