sn0w12-ComfyUI-Syntax-Highlighting/__init__.py

import json
import os
from pathlib import Path
import folder_paths
from server import PromptServer
from aiohttp import web
import aiohttp

from .src.config_reader import ConfigReader
from .src.save_preview_image import SavePreviewImage
from .src.util import get_preview_image_path

WEB_DIRECTORY = "./web"
API_PREFIX = "/SyntaxHighlighting"
NODE_PREFIX = "SyntaxHighlighting_"
NODE_CLASS_MAPPINGS = {
    f"{NODE_PREFIX}SavePreviewImage": SavePreviewImage,
}
NODE_DISPLAY_NAME_MAPPINGS = {
    f"{NODE_PREFIX}SavePreviewImage": "Save Preview Image",
}


image_index = {}


def load_image_index():
    images_dir = Path(get_preview_image_path())
    json_path = images_dir.joinpath("images.json")
    if not json_path.exists():
        return {}
    with open(json_path, "r", encoding="utf-8") as f:
        data = json.load(f)
    return {img["filename"]: img["path"] for img in data.get("images", [])}


def reload_image_index():
    global image_index
    image_index.clear()
    image_index.update(load_image_index())


def index_images():
    """
    Create a JSON index of all images in the web/images directory and its subfolders.
    """
    try:
        images_dir = Path(get_preview_image_path())
        if not images_dir.exists():
            print(f"Warning: Directory not found: {images_dir}")
            return {"success": False}

        image_files = []
        for ext in [".png", ".jpg", ".jpeg", ".webp"]:
            image_files.extend(images_dir.rglob(f"*{ext}"))
        print(f"Found {len(image_files)} images in {images_dir}")

        images = []
        for img_path in image_files:
            images.append(
                {
                    "filename": img_path.stem,
                    "path": str(img_path),
                }
            )

        output = {"images": images, "count": len(images)}
        json_path = images_dir.joinpath("images.json")
        with open(json_path, "w", encoding="utf-8") as f:
            json.dump(output, f, indent=2)

        print(f"Created index with {len(images)} images at {json_path}")
        return_data = {"success": True, "count": len(images)}
        reload_image_index()
        return return_data

    except Exception as e:
        print(f"Error creating image index: {str(e)}")
        return {"success": False}


index_images()


@PromptServer.instance.routes.get(f"{API_PREFIX}/loras")
async def get_loras(request):
    loras = folder_paths.get_filename_list("loras")
    return web.json_response(list(map(lambda a: os.path.splitext(a)[0], loras)))


@PromptServer.instance.routes.get(f"{API_PREFIX}/embeddings")
async def get_embeddings(request):
    embeddings = folder_paths.get_filename_list("embeddings")
    return web.json_response(list(map(lambda a: os.path.splitext(a)[0], embeddings)))


@PromptServer.instance.routes.get(f"{API_PREFIX}/enabled")
async def get_enabled(request):
    return web.json_response({"enabled": True})


@PromptServer.instance.routes.get(f"{API_PREFIX}/index")
async def index_images_endpoint(request):
    data = index_images()
    reload_image_index()
    return web.json_response(data)


@PromptServer.instance.routes.get(f"{API_PREFIX}/favorites")
async def favorites_endpoint(request):
    favorites = ConfigReader.get_setting("SyntaxHighlighting.favorites", [])
    return web.json_response({"favorites": favorites})


@PromptServer.instance.routes.get(f"{API_PREFIX}/images_json")
async def serve_images_json(request):
    images_dir = Path(get_preview_image_path())
    json_path = images_dir.joinpath("images.json")
    if json_path.exists():
        return web.FileResponse(path=json_path)
    return web.Response(status=404, text="images.json not found")


@PromptServer.instance.routes.get(f"{API_PREFIX}/images/{{filename}}")
async def serve_image(request):
    filename = request.match_info["filename"]
    if not filename:
        return web.Response(status=400, text="Filename is required")

    img_path = image_index.get(filename)
    if img_path and Path(img_path).exists():
        return web.FileResponse(path=img_path)
    return web.Response(status=404, text="Image not found")


@PromptServer.instance.routes.get(f"{API_PREFIX}/wiki/{{tag}}")
async def fetch_wiki_page(request):
    """
    Proxy endpoint to fetch Danbooru wiki pages.
    This avoids CSP violations by making the request server-side.
    """
    tag = request.match_info["tag"]
    if not tag:
        return web.Response(status=400, text="Tag is required")

    # Clean the tag similar to frontend logic
    tag = tag.strip()

    url = f"https://danbooru.donmai.us/wiki_pages/{tag}"

    try:
        async with aiohttp.ClientSession() as session:
            async with session.get(url, timeout=aiohttp.ClientTimeout(total=10)) as response:
                if response.status != 200:
                    return web.json_response({
                        "success": False,
                        "error": {
                            "status": response.status,
                            "message": f"HTTP error! status: {response.status}",
                            "url": url,
                        }
                    }, status=response.status)

                html_content = await response.text()
                return web.json_response({
                    "success": True,
                    "data": html_content
                })
    except aiohttp.ClientError as e:
        return web.json_response({ "success": False, "error": { "message": str(e), "url": url, }
        }, status=500)
    except Exception as e:
        return web.json_response({
            "success": False,
            "error": {
                "message": f"Unexpected error: {str(e)}",
                "url": url,
            }
        }, status=500)