mirror of
https://git.datalinker.icu/comfyanonymous/ComfyUI
synced 2026-11-14 01:27:00 +08:00
When generating workflow_api.json in API mode, add logic to get model hash values based on the property status of the custom node (SaveAllModelHashesNode).
This commit is contained in:
parent
b33cd61070
commit
890034e999
77
model_mng/model_hash.py
Normal file
77
model_mng/model_hash.py
Normal file
@ -0,0 +1,77 @@
|
||||
import hashlib
|
||||
import logging
|
||||
import os
|
||||
import json
|
||||
from aiohttp import web
|
||||
from pathlib import Path
|
||||
|
||||
import os
|
||||
import time
|
||||
import logging
|
||||
import folder_paths
|
||||
|
||||
class ModelHash():
|
||||
def __init__(self):
|
||||
self.model_prefix_names = {"ckpt_name": "checkpoints",
|
||||
"control_net_name": "controlnet",
|
||||
"lora_name": "loras",
|
||||
"unet_name":"unet",
|
||||
"clip_name": "clip",
|
||||
"style_model_name": "style_models",
|
||||
"gligen_name": "gligen",
|
||||
"vae_name": "vae" }
|
||||
pass
|
||||
|
||||
|
||||
def get_MD5_hash(self, path):
|
||||
path = Path(path)
|
||||
path.parent.mkdir(parents=True, exist_ok=True)
|
||||
if not path.exists():
|
||||
return None
|
||||
|
||||
file_digest = hashlib.sha256(open(path, 'rb').read()).hexdigest()
|
||||
return file_digest
|
||||
|
||||
async def get_hash(self, request):
|
||||
|
||||
json_data = await request.json()
|
||||
|
||||
bPromtFlag = False
|
||||
prompt = {}
|
||||
resp = True
|
||||
|
||||
try:
|
||||
if "prompt" in json_data:
|
||||
prompt = json_data["prompt"]
|
||||
|
||||
hash_node = {}
|
||||
for (node_key, node_values) in prompt.items():
|
||||
if "class_type" in node_values and node_values["class_type"] =="SaveAllModelHashesNode" :
|
||||
hash_node = node_values
|
||||
if "digest" not in node_values:
|
||||
hash_node['digest'] = {}
|
||||
bPromtFlag = True
|
||||
break
|
||||
|
||||
if bPromtFlag:
|
||||
for (node_key, node_values) in prompt.items():
|
||||
for (model_key, model_values) in self.model_prefix_names.items():
|
||||
digest = None
|
||||
if 'inputs' in node_values and model_key in node_values['inputs']:
|
||||
ckpt_path = folder_paths.get_full_path(model_values, node_values['inputs'][model_key])
|
||||
digest = self.get_MD5_hash(ckpt_path)
|
||||
|
||||
hash_node['digest'][node_values['inputs'][model_key]] = digest
|
||||
|
||||
except Exception as e:
|
||||
print(f"Failed to parse the json data: {json_data} \n {e}")
|
||||
resp = False
|
||||
|
||||
res = {"output": prompt, "res_status": resp}
|
||||
return res
|
||||
|
||||
def add_routes(self, routes):
|
||||
@routes.post("/api/get_hash")
|
||||
async def post_gethash(request):
|
||||
res = await self.get_hash(request)
|
||||
return web.json_response(res)
|
||||
@ -12,6 +12,7 @@ import json
|
||||
import glob
|
||||
import struct
|
||||
import ssl
|
||||
import hashlib
|
||||
from PIL import Image, ImageOps
|
||||
from PIL.PngImagePlugin import PngInfo
|
||||
from io import BytesIO
|
||||
@ -31,6 +32,7 @@ from model_filemanager import download_model, DownloadModelStatus
|
||||
from typing import Optional
|
||||
from api_server.routes.internal.internal_routes import InternalRoutes
|
||||
|
||||
from model_mng.model_hash import ModelHash
|
||||
|
||||
class BinaryEventTypes:
|
||||
PREVIEW_IMAGE = 1
|
||||
@ -81,6 +83,7 @@ class PromptServer():
|
||||
self.messages = asyncio.Queue()
|
||||
self.client_session:Optional[aiohttp.ClientSession] = None
|
||||
self.number = 0
|
||||
self.model_hash = ModelHash()
|
||||
|
||||
middlewares = [cache_control]
|
||||
if args.enable_cors_header:
|
||||
@ -629,6 +632,8 @@ class PromptServer():
|
||||
if isinstance(route, web.RouteDef):
|
||||
api_routes.route(route.method, "/api" + route.path)(route.handler, **route.kwargs)
|
||||
self.app.add_routes(api_routes)
|
||||
self.model_hash.add_routes(self.routes)
|
||||
|
||||
self.app.add_routes(self.routes)
|
||||
|
||||
for name, dir in nodes.EXTENSION_WEB_DIRS.items():
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user