93 lines
3.3 KiB
Python
93 lines
3.3 KiB
Python
"""Test-only IMAGE source. Copy into an isolated ComfyUI custom_nodes folder."""
|
|
import os
|
|
import json
|
|
from pathlib import Path
|
|
|
|
from aiohttp import web
|
|
import numpy as np
|
|
from PIL import Image, ImageOps
|
|
import torch
|
|
from server import PromptServer
|
|
import sys
|
|
import folder_paths
|
|
from comfy.cli_args import args, LatentPreviewMethod
|
|
from comfy.utils import ProgressBar
|
|
|
|
|
|
class PreviewTestImage:
|
|
executions = 0
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {"required": {
|
|
"count": ("INT", {"default": 100, "min": 1, "max": 256}),
|
|
"size": ("INT", {"default": 1024, "min": 8, "max": 4096}),
|
|
"fail_at": ("INT", {"default": -1, "min": -1, "max": 255}),
|
|
}}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
FUNCTION = "create"
|
|
CATEGORY = "testing"
|
|
|
|
def create(self, count, size, fail_at):
|
|
type(self).executions += 1
|
|
source = os.environ.get("ENCRYPT_PREVIEW_TEST_IMAGE")
|
|
if source:
|
|
with Image.open(source) as image:
|
|
with ImageOps.fit(image.convert("RGB"), (size, size)) as resized:
|
|
base = np.array(resized, dtype=np.float32) / 255.0
|
|
else:
|
|
y, x = np.mgrid[:size, :size].astype(np.float32) / size
|
|
base = np.stack([x, y, (x + y) / 2], axis=-1)
|
|
images = torch.empty((count, size, size, 3), dtype=torch.float32)
|
|
for index in range(count):
|
|
images[index] = torch.from_numpy(np.roll(base, index % 4, axis=1).copy())
|
|
images[index, 0, 0, 0] = index / 255.0
|
|
if 0 <= fail_at < count:
|
|
images[fail_at, 0, 0, 0] = float("nan")
|
|
if args.preview_method != LatentPreviewMethod.NoPreviews:
|
|
with Image.fromarray(np.clip(base * 255, 0, 255).astype(np.uint8)) as preview:
|
|
ProgressBar(1).update_absolute(1, 1, ("JPEG", preview, 384))
|
|
return (images,)
|
|
|
|
|
|
WEB_DIRECTORY = "./web"
|
|
NODE_CLASS_MAPPINGS = {"PreviewTestImage": PreviewTestImage}
|
|
|
|
|
|
@PromptServer.instance.routes.post("/encrypt-preview-test-result")
|
|
async def save_browser_result(request):
|
|
destination = Path(os.environ["ENCRYPT_PREVIEW_TEST_RESULTS"])
|
|
destination.parent.mkdir(parents=True, exist_ok=True)
|
|
destination.write_text(json.dumps(await request.json(), ensure_ascii=False, indent=2), encoding="utf-8")
|
|
return web.json_response({"saved": True})
|
|
|
|
|
|
def preview_store():
|
|
import nodes
|
|
return sys.modules[nodes.NODE_CLASS_MAPPINGS["EncryptPreviewImage"].__module__].STORE
|
|
|
|
|
|
@PromptServer.instance.routes.get("/encrypt-preview-test/status")
|
|
async def test_status(request):
|
|
store = preview_store()
|
|
return web.json_response({
|
|
"fixture_executions": PreviewTestImage.executions,
|
|
"cache_bytes": store.size,
|
|
"cache_items": store.item_count,
|
|
"preview_method": args.preview_method.value,
|
|
"temp_pngs": len(list(Path(folder_paths.get_temp_directory()).rglob("*.png"))),
|
|
"output_pngs": len(list(Path(folder_paths.get_output_directory()).rglob("*.png"))),
|
|
})
|
|
|
|
|
|
@PromptServer.instance.routes.post("/encrypt-preview-test/expire")
|
|
async def expire_preview(request):
|
|
batch_id = (await request.json())["batch_id"]
|
|
store = preview_store()
|
|
with store.lock:
|
|
batch = store.batches.get(batch_id)
|
|
if batch and not batch.active:
|
|
batch.expires = 0
|
|
store.expire()
|
|
return web.json_response({"expired": True})
|