383 lines
14 KiB
Python
383 lines
14 KiB
Python
from __future__ import annotations
|
|
from aiohttp import web
|
|
from typing import NamedTuple
|
|
from pathlib import Path
|
|
import json
|
|
import traceback
|
|
import re
|
|
import logging
|
|
import itertools
|
|
|
|
from comfy import model_detection
|
|
import comfy.utils
|
|
import folder_paths
|
|
import server
|
|
|
|
from .translation import available_languages, translate
|
|
from .krita import WorkflowExchange
|
|
from .nodes import image_cache
|
|
|
|
input_block_name = "model.diffusion_model.input_blocks.0.0.weight"
|
|
|
|
model_names = {
|
|
"SD15": "sd15",
|
|
"SD20": "sd20",
|
|
"SD21UnclipL": "sd21",
|
|
"SD21UnclipH": "sd21",
|
|
"SDXLRefiner": "sdxl-refiner",
|
|
"SDXL": "sdxl",
|
|
"SSD1B": "ssd1b",
|
|
"SVD_img2vid": "svd",
|
|
"Stable_Cascade_B": "cascade-b",
|
|
"Stable_Cascade_C": "cascade-c",
|
|
"SD3": "sd3",
|
|
"AuraFlow": "aura-flow",
|
|
"HunyuanDiT": "hunyuan-dit",
|
|
"HunyuanDiT1": "hunyuan-dit",
|
|
"Flux": "flux",
|
|
"FluxInpaint": "flux",
|
|
"FluxSchnell": "flux-schnell",
|
|
"GenmoMochi": "mochi",
|
|
"LTXV": "ltxv",
|
|
"HunyuanVideo": "hunyuan-video",
|
|
"CosmosT2V": "cosmos",
|
|
"CosmosI2V": "cosmos",
|
|
"CosmosT2IPredict2": "cosmos-predict2",
|
|
"CosmosI2VPredict2": "cosmos-predict2",
|
|
"ZImage": "z-image",
|
|
"Lumina2": "lumina2",
|
|
"WAN21_T2V": "wan21",
|
|
"WAN21_I2V": "wan21",
|
|
"WAN21_FunControl2V": "wan21-fun",
|
|
"WAN21_Vace": "wan21-vace",
|
|
"WAN21_Camera": "wan21-camera",
|
|
"HiDream": "hi-dream",
|
|
"Chroma": "chroma",
|
|
"ACEStep": "ace-step",
|
|
"Omnigen2": "omnigen2",
|
|
"QwenImage": "qwen-image",
|
|
"Flux2": "flux2",
|
|
}
|
|
|
|
gguf_architectures = {
|
|
"sd1": "sd15",
|
|
"qwen_image": "qwen-image",
|
|
}
|
|
|
|
|
|
class FakeTensor(NamedTuple):
|
|
shape: tuple
|
|
|
|
@staticmethod
|
|
def from_dict(d):
|
|
try:
|
|
return FakeTensor(tuple(d["shape"]))
|
|
except KeyError:
|
|
return d
|
|
|
|
|
|
def inspect_safetensors(filename: str, model_type: str, is_checkpoint: bool):
|
|
try:
|
|
# Read header of safetensors file
|
|
path = folder_paths.get_full_path(model_type, filename)
|
|
header = comfy.utils.safetensors_header(path)
|
|
if header:
|
|
cfg = json.loads(header.decode("utf-8"))
|
|
|
|
# Build a fake "state_dict" from the header info to avoid reading the full weights
|
|
for key in cfg:
|
|
if not key == "__metadata__":
|
|
cfg[key] = FakeTensor.from_dict(cfg[key])
|
|
|
|
# Reuse Comfy's model detection
|
|
prefix = model_detection.unet_prefix_from_state_dict(cfg)
|
|
if not is_checkpoint:
|
|
cfg = comfy.utils.state_dict_prefix_replace(cfg, {prefix: ""}, filter_keys=False)
|
|
prefix = ""
|
|
try: # latest ComfyUI takes 2 args
|
|
unet_config = model_detection.detect_unet_config(cfg, prefix)
|
|
except TypeError as e: # older ComfyUI versions take 3 args
|
|
raise TypeError(f"{e} when calling detect_unet_config - old version of ComfyUI?")
|
|
|
|
# Get input count to detect inpaint models
|
|
if input_block := cfg.get(input_block_name, None):
|
|
input_count = input_block.shape[1]
|
|
else:
|
|
input_count = 4
|
|
|
|
# Find a matching base model depending on unet config
|
|
base_model = None
|
|
model_type = None
|
|
model_quant = None
|
|
|
|
# Check if it's a Nunchaku SVDQ model by inspecting metadata
|
|
raw_name = detect_svdq(cfg)
|
|
if raw_name:
|
|
model_quant = "svdq"
|
|
# Otherwise try ComfyUI's model detection
|
|
elif unet_config is not None:
|
|
base_model = model_detection.model_config_from_unet_config(unet_config)
|
|
if base_model:
|
|
raw_name = base_model.__class__.__name__
|
|
if raw_name == "SDXL":
|
|
model_type = base_model.model_type(cfg).name.lower().replace("_", "-")
|
|
|
|
if not raw_name:
|
|
return {"base_model": "unknown"}
|
|
|
|
base_model_name = model_names.get(raw_name, "unknown")
|
|
result = {"base_model": base_model_name}
|
|
result["is_inpaint"] = (
|
|
base_model_name in ["sd15", "sdxl"] and input_count > 4
|
|
) or raw_name == "FluxInpaint"
|
|
if model_quant:
|
|
result["quant"] = model_quant
|
|
if model_type:
|
|
result["type"] = model_type
|
|
elif "T2I" in raw_name:
|
|
result["type"] = "t2i"
|
|
elif "I2V" in raw_name:
|
|
result["type"] = "i2v"
|
|
elif "T2V" in raw_name:
|
|
result["type"] = "t2v"
|
|
elif "Control2V" in raw_name:
|
|
result["type"] = "control2v"
|
|
return result
|
|
return {"base_model": "unknown"}
|
|
except Exception as e:
|
|
traceback.print_exc()
|
|
return {"base_model": "unknown", "error": f"Failed to detect base model: {e}"}
|
|
|
|
|
|
def detect_svdq(cfg: dict) -> str | None:
|
|
if md := cfg.get("__metadata__"):
|
|
if comfy_config := md.get("comfy_config"):
|
|
if isinstance(comfy_config, str):
|
|
comfy_config = json.loads(comfy_config)
|
|
return comfy_config.get("model_class")
|
|
model_class = md.get("model_class")
|
|
if model_class == "NunchakuFluxTransformer2dModel":
|
|
return "Flux"
|
|
if model_class == "NunchakuQwenImageTransformer2DModel":
|
|
return "QwenImage"
|
|
return None
|
|
|
|
|
|
def inspect_gguf(filename: str, model_type: str):
|
|
try:
|
|
import gguf
|
|
except ImportError:
|
|
return {"base_model": "unknown", "error": "GGUF module not found"}
|
|
|
|
try:
|
|
path = folder_paths.get_full_path(model_type, filename)
|
|
reader = gguf.GGUFReader(path)
|
|
arch_field = reader.get_field("general.architecture")
|
|
if arch_field is not None:
|
|
if len(arch_field.types) != 1 or arch_field.types[0] != gguf.GGUFValueType.STRING:
|
|
raise TypeError(
|
|
f"Bad type for GGUF general.architecture key: expected string, got {arch_field.types!r}"
|
|
)
|
|
arch_str = str(arch_field.parts[arch_field.data[-1]], encoding="utf-8")
|
|
else: # stable-diffusion.cpp, requires conversion. not handled for now
|
|
return {"base_model": "flux", "is_inpaint": False}
|
|
if arch_str == "flux" and any(
|
|
t.name.startswith("distilled_guidance_layer")
|
|
for t in itertools.islice(reader.tensors, 5)
|
|
):
|
|
arch_str = "chroma"
|
|
|
|
result = {
|
|
"base_model": gguf_architectures.get(arch_str, arch_str),
|
|
"is_inpaint": False,
|
|
}
|
|
try:
|
|
result["quant"] = reader.get_field("general.file_type").lower()
|
|
except Exception as e:
|
|
result["quant"] = "gguf"
|
|
return result
|
|
|
|
except Exception as e:
|
|
# traceback.print_exc()
|
|
return {"base_model": "unknown", "error": f"Failed to detect base model: {e}"}
|
|
|
|
|
|
def inspect_diffusion_model(filename: str, model_type: str, is_checkpoint: bool):
|
|
if filename.endswith(".gguf"):
|
|
return inspect_gguf(filename, model_type)
|
|
return inspect_safetensors(filename, model_type, is_checkpoint)
|
|
|
|
|
|
def inspect_models(model_type: str, params: dict[str, str]):
|
|
try:
|
|
try:
|
|
files = folder_paths.get_filename_list(model_type)
|
|
except KeyError:
|
|
return web.json_response({"error": f"Model folder not found: {model_type}"})
|
|
limit = int(params.get("limit", "1000"))
|
|
offset = int(params.get("offset", "0"))
|
|
files_range = files[offset : offset + limit]
|
|
is_checkpoint = model_type == "checkpoints"
|
|
info = {
|
|
filename: inspect_diffusion_model(filename, model_type, is_checkpoint)
|
|
for filename in files_range
|
|
}
|
|
if "limit" in params:
|
|
info["_meta"] = dict(offset=offset, count=len(files_range), total=len(files))
|
|
return web.json_response(info)
|
|
except Exception as e:
|
|
traceback.print_exc()
|
|
return web.json_response(dict(error=str(e)), status=500)
|
|
|
|
|
|
def has_invalid_folder_name(folder_name: str):
|
|
valid_names = list(folder_paths.folder_names_and_paths.keys())
|
|
if folder_name not in valid_names:
|
|
return web.json_response(
|
|
dict(error=f"Invalid folder path, must be one of {', '.join(valid_names)}"),
|
|
status=400,
|
|
)
|
|
return None
|
|
|
|
|
|
def has_invalid_filename(filename: str):
|
|
if not filename.lower().endswith((".sft", ".safetensors")):
|
|
return web.json_response(dict(error="File extension must be .safetensors"), status=400)
|
|
if not filename or not filename.strip() or len(filename) > 255:
|
|
return web.json_response(dict(error="Invalid filename"), status=400)
|
|
if any(char in filename for char in ["..", "/", "\\", "\n", "\r", "\t", "\0"]):
|
|
return web.json_response(dict(error="Invalid filename"), status=400)
|
|
if filename.startswith(".") or not re.match(r"^[a-zA-Z0-9_\-. ]+$", filename):
|
|
return web.json_response(dict(error="Invalid filename"), status=400)
|
|
return None
|
|
|
|
|
|
async def image_sender(data: bytes):
|
|
mem = memoryview(data)
|
|
csize = 2**14
|
|
for i in range(0, len(mem), csize):
|
|
yield mem[i : i + csize]
|
|
|
|
|
|
_server: server.PromptServer | None = getattr(server.PromptServer, "instance", None)
|
|
if _server is not None:
|
|
_workflow_exchange = WorkflowExchange(_server)
|
|
|
|
@_server.routes.get("/api/etn/model_info/{folder_name}")
|
|
async def model_info(request: web.Request):
|
|
folder_name = request.match_info.get("folder_name", "checkpoints")
|
|
error = has_invalid_folder_name(folder_name)
|
|
if error is not None:
|
|
return error
|
|
return inspect_models(folder_name, request.rel_url.query)
|
|
|
|
@_server.routes.get("/api/etn/model_info")
|
|
async def api_model_info(request):
|
|
return inspect_models("checkpoints")
|
|
|
|
@_server.routes.get("/api/etn/languages")
|
|
async def languages(request):
|
|
try:
|
|
result = [dict(name=name, code=code) for code, name in available_languages()]
|
|
return web.json_response(result)
|
|
except Exception as e:
|
|
return web.json_response(dict(error=str(e)), status=500)
|
|
|
|
@_server.routes.get("/api/etn/translate/{lang}/{text}")
|
|
async def translate_text(request):
|
|
try:
|
|
language = request.match_info.get("lang", "en")
|
|
text = request.match_info.get("text", "")
|
|
result = translate(f"lang:{language} {text}")
|
|
return web.json_response(result)
|
|
except Exception as e:
|
|
return web.json_response(dict(error=str(e)), status=500)
|
|
|
|
@_server.routes.get("/api/etn/image/{id}")
|
|
async def get_image(request: web.Request):
|
|
try:
|
|
id = request.match_info.get("id", "")
|
|
data, content_type = image_cache.get(id)
|
|
if data is None or content_type is None:
|
|
return web.json_response(dict(error="Image not found"), status=404)
|
|
response = web.Response(
|
|
body=image_sender(data),
|
|
content_type=content_type,
|
|
headers={"Content-Length": str(len(data))},
|
|
)
|
|
return response
|
|
except Exception as e:
|
|
return web.json_response(dict(error=str(e)), status=500)
|
|
|
|
@_server.routes.put("/api/etn/image/{id}")
|
|
async def put_image(request: web.Request):
|
|
try:
|
|
id = request.match_info.get("id", "")
|
|
if id in image_cache:
|
|
return web.json_response(dict(status="cached"), status=200)
|
|
|
|
content_type = request.headers.get("Content-Type", "application/octet-stream")
|
|
data = bytearray()
|
|
async for chunk, _ in request.content.iter_chunks():
|
|
data.extend(chunk)
|
|
|
|
image_cache.insert(id, bytes(data), content_type)
|
|
return web.json_response(dict(status="success"), status=201)
|
|
except Exception as e:
|
|
return web.json_response(dict(error=str(e)), status=500)
|
|
|
|
@_server.routes.put("/api/etn/upload/{folder_name}/{filename}")
|
|
async def upload(request: web.Request):
|
|
folder_name = request.match_info.get("folder_name", "")
|
|
error = has_invalid_folder_name(folder_name)
|
|
if error is not None:
|
|
return error
|
|
|
|
filename = request.match_info.get("filename", "")
|
|
error = has_invalid_filename(filename)
|
|
if error is not None:
|
|
return error
|
|
|
|
try:
|
|
if folder_paths.get_full_path(folder_name, filename) is not None:
|
|
return web.json_response(dict(status="cached"), status=200)
|
|
|
|
folder = Path(folder_paths.folder_names_and_paths[folder_name][0][0])
|
|
total_size = int(request.headers.get("Content-Length", "0"))
|
|
logging.info(
|
|
f"Uploading {filename} ({total_size / (1024**2):.1f} MB) to {folder} folder"
|
|
)
|
|
|
|
with open(folder / filename, "wb") as f:
|
|
async for chunk, _ in request.content.iter_chunks():
|
|
f.write(chunk)
|
|
|
|
return web.json_response(dict(status="success"), status=201)
|
|
except Exception as e:
|
|
return web.json_response(dict(error=str(e)), status=500)
|
|
|
|
async def _handle_workflow_request(request: web.Request, handler, *arg_keys):
|
|
try:
|
|
data = await request.json()
|
|
args = [data[key] for key in arg_keys]
|
|
await handler(*args)
|
|
return web.json_response(dict(status="success"), status=200)
|
|
except KeyError as e:
|
|
return web.json_response(dict(error=str(e)), status=400)
|
|
except Exception as e:
|
|
return web.json_response(dict(error=str(e)), status=500)
|
|
|
|
@_server.routes.post("/api/etn/workflow/publish")
|
|
async def publish_workflow(request: web.Request):
|
|
return await _handle_workflow_request(
|
|
request, _workflow_exchange.publish, "name", "client_id", "workflow"
|
|
)
|
|
|
|
@_server.routes.post("/api/etn/workflow/subscribe")
|
|
async def subscribe_workflow(request: web.Request):
|
|
return await _handle_workflow_request(request, _workflow_exchange.subscribe, "client_id")
|
|
|
|
@_server.routes.post("/api/etn/workflow/unsubscribe")
|
|
async def unsubscribe_workflow(request: web.Request):
|
|
return await _handle_workflow_request(request, _workflow_exchange.unsubscribe, "client_id")
|