diff --git a/custom_routes.py b/custom_routes.py index 5ef78bc..7296884 100755 --- a/custom_routes.py +++ b/custom_routes.py @@ -122,6 +122,10 @@ async def shellagent_get_file(request): async def shellagent_export(request): data = await request.json() prompt = data["prompt"] + custom_dependencies = data.get("custom_dependencies", { + "models": {}, + "custom_nodes": {} + }) # extra_data = data["extra_data"] workflow_id = str(uuid.uuid4()) @@ -137,7 +141,7 @@ async def shellagent_export(request): try: schemas = schema_validator(prompt) # custom_node.json - dependency_results = resolve_dependencies(prompt) + dependency_results = resolve_dependencies(prompt, custom_dependencies) # save_root = os.path.join(WORKFLOW_ROOT, workflow_id) # os.makedirs(save_root, exist_ok=True) @@ -162,6 +166,6 @@ async def shellagent_export(request): status = 400 return_dict = { "success": False, - "message": str(traceback.print_exc()) + "message": str(traceback.format_exc()), } return web.json_response(return_dict, status=status) \ No newline at end of file diff --git a/dependency_checker.py b/dependency_checker.py index 092ddc3..4ea5669 100644 --- a/dependency_checker.py +++ b/dependency_checker.py @@ -26,7 +26,9 @@ ComfyUIModelLoaders = { 'ImageOnlyCheckpointLoader': (["ckpt_name"], "checkpoints"), "UpscaleModelLoader": (["model_name"], "upscale_models"), "TripleCLIPLoader": (["clip_name1", "clip_name2", "clip_name3"], "clip"), - "HypernetworkLoader": (["hypernetwork_name"], "hypernetworks") + "HypernetworkLoader": (["hypernetwork_name"], "hypernetworks"), + "SUPIR_model_loader_v2": (["supir_model"], "checkpoints"), + "SUPIR_model_loader_v2_clip": (["supir_model"], "checkpoints"), } @@ -114,8 +116,21 @@ def inspect_repo_version(module_path): } return result +def fetch_model_searcher_results(model_ids): + import requests + url = "https://shellagent.myshell.ai/models_searcher/search_urls" + headers = { + "Content-Type": "application/json" + } + data = { + "sha256": model_ids + } -def resolve_dependencies(prompt): # resolve custom nodes and models at the same time + response = requests.post(url, headers=headers, json=data) + results = [item[:10] for item in response.json()] + return results + +def resolve_dependencies(prompt, custom_dependencies): # resolve custom nodes and models at the same time from nodes import NODE_CLASS_MAPPINGS custom_nodes = [] ckpt_paths = [] @@ -144,22 +159,37 @@ def resolve_dependencies(prompt): # resolve custom nodes and models at the same try: repo_info = inspect_repo_version(custom_node.replace(".", "/")) custom_nodes_list.append(repo_info) + if repo_info["repo"] == "": + repo_info["require_recheck"] = True + if repo_info["name"] in custom_dependencies["custom_nodes"]: + repo_info["repo"] = custom_dependencies["custom_nodes"][repo_info["name"]].get("repo", "") + repo_info["commit"] = custom_dependencies["custom_nodes"][repo_info["name"]].get("commit", "") except: print(f"failed to resolve repo info of {custom_node}") # step 2: models models_dict = {} + missing_model_ids = [] for ckpt_path in ckpt_paths: model_id, item = handle_model_info(ckpt_path) models_dict[model_id] = item + if len(item["urls"]) == 0: + item["require_recheck"] = True + if model_id in custom_dependencies["models"]: + item["urls"] = custom_dependencies["models"][model_id].get("urls", []) + missing_model_ids.append(model_id) + + # try to fetch from myshell model searcher + missing_model_results_myshell = fetch_model_searcher_results(missing_model_ids) + for missing_model_id, missing_model_urls in zip(missing_model_ids, missing_model_results_myshell): + if len(missing_model_urls) > 0: + models_dict[missing_model_id]["require_recheck"] = False + models_dict[missing_model_id]["urls"] = missing_model_urls + print("successfully fetch results from myshell", models_dict[missing_model_id]) # step 3: handle local files process_local_file_path_async(file_mapping_dict, max_workers=20) files_dict = {v[0]: {"filename": v[2], "urls": [v[1]]} for v in file_mapping_dict.values()} - dependencies = { - "models": models_dict, - "files": files_dict - } results = { "comfyui_version": comfyui_version,