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:
YingWang 2024-08-29 18:55:44 +08:00
parent b33cd61070
commit 890034e999
2 changed files with 82 additions and 0 deletions

77
model_mng/model_hash.py Normal file
View 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)

View File

@ -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():