diff --git a/Rayko_LoRA_Loader.py b/Rayko_LoRA_Loader.py index a577c48..d919277 100644 --- a/Rayko_LoRA_Loader.py +++ b/Rayko_LoRA_Loader.py @@ -32,6 +32,7 @@ os.makedirs(PRESETS_DIR, exist_ok=True) os.makedirs(LORA_PRESETS_DIR, exist_ok=True) os.makedirs(RAYKO_LORA_DATA_DIR, exist_ok=True) + def clean_html(text): if not text: return "" @@ -39,6 +40,38 @@ def clean_html(text): clean = re.sub(r'\s+', ' ', clean).strip() return clean + +def sanitize_preset_name(name): + """Sanitize a preset name to a safe filename component. + + Strips all characters except alphanumerics, spaces, underscores and dashes. + Returns an empty string if the result is empty. + """ + if not isinstance(name, str): + return "" + return "".join(c for c in name.strip() if c.isalnum() or c in " _-").strip() + + +def _safe_preset_path(name): + """Return an absolute path inside LORA_PRESETS_DIR for the given preset name. + + Returns None if the name is invalid or if the resulting path escapes + LORA_PRESETS_DIR (defense in depth against path traversal). + """ + safe_name = sanitize_preset_name(name) + if not safe_name: + return None + base = os.path.realpath(LORA_PRESETS_DIR) + candidate = os.path.realpath(os.path.join(base, f"{safe_name}.json")) + try: + if os.path.commonpath([base, candidate]) != base: + return None + except ValueError: + # Different drives on Windows / incompatible paths + return None + return candidate + + def extract_metadata_from_safetensors(lora_full_path): try: with open(lora_full_path, "rb") as f: @@ -56,6 +89,7 @@ def extract_metadata_from_safetensors(lora_full_path): pass return None + def compute_sha256(file_path): sha256_hash = hashlib.sha256() with open(file_path, "rb") as f: @@ -63,6 +97,7 @@ def compute_sha256(file_path): sha256_hash.update(chunk) return sha256_hash.hexdigest() + def save_to_rayko_db(file_hash, data, source, overwrite=False): db_path = os.path.join(RAYKO_LORA_DATA_DIR, f"{file_hash}.json") if not overwrite and os.path.exists(db_path): @@ -79,6 +114,7 @@ def save_to_rayko_db(file_hash, data, source, overwrite=False): json.dump(db_data, f, indent=2, ensure_ascii=False) return True + class RaykoLoRALoader: @classmethod def INPUT_TYPES(cls): @@ -141,27 +177,30 @@ class RaykoLoRALoader: else: return (model, None) + NODE_CLASS_MAPPINGS = {"RaykoLoRALoader": RaykoLoRALoader} NODE_DISPLAY_NAME_MAPPINGS = {"RaykoLoRALoader": "🦊 RS LoRA Loader"} + @PromptServer.instance.routes.get("/rayko_lora_loader/get_loras") async def get_loras(request): return aiohttp.web.json_response(sorted(folder_paths.get_filename_list("loras"), key=lambda x: x.lower())) + @PromptServer.instance.routes.post("/rayko_lora_loader/save_preset") async def rayko_lora_loader_save_preset(request): try: data = await request.json() - name = "".join(c for c in data.get("name", "").strip() if c.isalnum() or c in " _-").strip() - if not name: + filepath = _safe_preset_path(data.get("name", "")) + if not filepath: return aiohttp.web.Response(status=400, text="Invalid name") - filepath = os.path.join(LORA_PRESETS_DIR, f"{name}.json") with open(filepath, 'w', encoding='utf-8') as f: json.dump({"lora_rows": data.get("lora_rows", [])}, f, indent=2) return aiohttp.web.Response(status=200, text="OK") except Exception as e: return aiohttp.web.Response(status=500, text=str(e)) + @PromptServer.instance.routes.post("/rayko_lora_loader/list_presets") async def rayko_lora_loader_list_presets(request): try: @@ -170,11 +209,14 @@ async def rayko_lora_loader_list_presets(request): except Exception as e: return aiohttp.web.Response(status=500, text=str(e)) + @PromptServer.instance.routes.post("/rayko_lora_loader/load_preset") async def rayko_lora_loader_load_preset(request): try: - name = (await request.json()).get("name") - filepath = os.path.join(LORA_PRESETS_DIR, f"{name}.json") + data = await request.json() + filepath = _safe_preset_path(data.get("name", "")) + if not filepath: + return aiohttp.web.Response(status=400, text="Invalid name") if os.path.exists(filepath): with open(filepath, 'r', encoding='utf-8') as f: return aiohttp.web.json_response(json.load(f)) @@ -182,11 +224,14 @@ async def rayko_lora_loader_load_preset(request): except Exception as e: return aiohttp.web.Response(status=500, text=str(e)) + @PromptServer.instance.routes.post("/rayko_lora_loader/delete_preset") async def rayko_lora_loader_delete_preset(request): try: - name = (await request.json()).get("name") - filepath = os.path.join(LORA_PRESETS_DIR, f"{name}.json") + data = await request.json() + filepath = _safe_preset_path(data.get("name", "")) + if not filepath: + return aiohttp.web.Response(status=400, text="Invalid name") if os.path.exists(filepath): os.remove(filepath) return aiohttp.web.Response(status=200, text="OK") @@ -194,6 +239,7 @@ async def rayko_lora_loader_delete_preset(request): except Exception as e: return aiohttp.web.Response(status=500, text=str(e)) + @PromptServer.instance.routes.post("/rayko_lora_loader/get_lora_info") async def rayko_lora_loader_get_lora_info(request): try: @@ -201,15 +247,15 @@ async def rayko_lora_loader_get_lora_info(request): lora_name = data.get("name", "") if not lora_name: return aiohttp.web.Response(status=400, text="No name provided") - + lora_relative_path = lora_name.replace("\\", "/") lora_full_path = folder_paths.get_full_path("loras", lora_relative_path) - + if not lora_full_path or not os.path.exists(lora_full_path): return aiohttp.web.json_response({"error": "File not found"}) - + file_hash = compute_sha256(lora_full_path) - + db_path = os.path.join(RAYKO_LORA_DATA_DIR, f"{file_hash}.json") if os.path.exists(db_path): try: @@ -223,7 +269,7 @@ async def rayko_lora_loader_get_lora_info(request): }) except Exception: pass - + safetensors_meta = extract_metadata_from_safetensors(lora_full_path) if safetensors_meta and (safetensors_meta["trained_words"] or safetensors_meta["description"]): save_to_rayko_db(file_hash, safetensors_meta, "safetensors") @@ -233,17 +279,18 @@ async def rayko_lora_loader_get_lora_info(request): "description": safetensors_meta["description"], "source": "safetensors" }) - + return aiohttp.web.json_response({ "full_name": lora_name, "trained_words": [], "description": "", "source": "none" }) - + except Exception as e: return aiohttp.web.Response(status=500, text=str(e)) + @PromptServer.instance.routes.post("/rayko_lora_loader/fetch_civitai_info") async def rayko_lora_loader_fetch_civitai_info(request): try: @@ -251,16 +298,16 @@ async def rayko_lora_loader_fetch_civitai_info(request): lora_name = data.get("name", "") if not lora_name: return aiohttp.web.Response(status=400, text="No name provided") - + lora_relative_path = lora_name.replace("\\", "/") lora_full_path = folder_paths.get_full_path("loras", lora_relative_path) - + if not lora_full_path or not os.path.exists(lora_full_path): return aiohttp.web.json_response({"error": "File not found"}) - + file_hash = compute_sha256(lora_full_path) civitai_url = f"https://civitai.com/api/v1/model-versions/by-hash/{file_hash}" - + async with aiohttp.ClientSession() as session: async with session.get(civitai_url, headers={"User-Agent": "ComfyUI-RaykoStudio/1.0"}) as response: if response.status == 404: @@ -269,12 +316,12 @@ async def rayko_lora_loader_fetch_civitai_info(request): return aiohttp.web.json_response({"error": "rate_limit", "message": "Rate limit reached."}) elif response.status != 200: return aiohttp.web.json_response({"error": "api_error", "message": f"API status {response.status}"}) - + civitai_data = await response.json() - + model_id = civitai_data.get("modelId") model_tags = [] - + if model_id: try: model_url = f"https://civitai.com/api/v1/models/{model_id}" @@ -284,23 +331,23 @@ async def rayko_lora_loader_fetch_civitai_info(request): model_tags = model_data.get("tags", []) except Exception as e: print(f"[Rayko] Error fetching full model details: {e}") - + raw_trained = civitai_data.get("trainedWords", []) - + trained_words = [] if isinstance(raw_trained, list): trained_words.extend([str(t).strip() for t in raw_trained if str(t).strip()]) if isinstance(model_tags, list): trained_words.extend([str(t).strip() for t in model_tags if str(t).strip()]) - + trained_words = list(dict.fromkeys(trained_words)) - + raw_description = civitai_data.get("description", "") clean_description = clean_html(raw_description) - + model_name = civitai_data.get("model", {}).get("name", "") version_name = civitai_data.get("name", "") - + if model_name: if version_name and version_name.lower() not in ["v1", "v1.0", "default", "latest", ""]: full_name = f"{model_name} ({version_name})" @@ -308,48 +355,49 @@ async def rayko_lora_loader_fetch_civitai_info(request): full_name = model_name else: full_name = lora_name - + cache_data = { "full_name": full_name, "trained_words": trained_words, "description": clean_description } - + save_to_rayko_db(file_hash, cache_data, "civitai", overwrite=True) - + return aiohttp.web.json_response({ "full_name": full_name, "trained_words": trained_words, "description": clean_description, "source": "civitai" }) - + except aiohttp.ClientError as e: return aiohttp.web.json_response({"error": "network", "message": f"Network error: {str(e)}"}) except Exception as e: return aiohttp.web.Response(status=500, text=str(e)) + @PromptServer.instance.routes.post("/rayko_lora_loader/add_tags") async def rayko_add_tags(request): try: data = await request.json() lora_name = data.get("name", "") new_tags = data.get("tags", []) - + if not lora_name: return aiohttp.web.Response(status=400, text="No name provided") if not isinstance(new_tags, list) or not new_tags: return aiohttp.web.json_response({"error": "No tags provided"}) - + lora_relative_path = lora_name.replace("\\", "/") lora_full_path = folder_paths.get_full_path("loras", lora_relative_path) - + if not lora_full_path or not os.path.exists(lora_full_path): return aiohttp.web.json_response({"error": "File not found"}) - + file_hash = compute_sha256(lora_full_path) db_path = os.path.join(RAYKO_LORA_DATA_DIR, f"{file_hash}.json") - + existing_data = {} if os.path.exists(db_path): try: @@ -357,74 +405,75 @@ async def rayko_add_tags(request): existing_data = json.load(f) except Exception: pass - + existing_tags = existing_data.get("trained_words", []) if not isinstance(existing_tags, list): existing_tags = [] - + combined_tags = existing_tags.copy() for tag in new_tags: tag = str(tag).strip() if tag and tag not in combined_tags: combined_tags.append(tag) - + save_data = { "full_name": existing_data.get("name", ""), "trained_words": combined_tags, "description": existing_data.get("description", "") } - + save_to_rayko_db(file_hash, save_data, "manual", overwrite=True) - + return aiohttp.web.json_response({ "trained_words": combined_tags, "added": len(combined_tags) - len(existing_tags) }) - + except Exception as e: return aiohttp.web.Response(status=500, text=str(e)) + @PromptServer.instance.routes.post("/rayko_lora_loader/remove_tag") async def rayko_remove_tag(request): try: data = await request.json() lora_name = data.get("name", "") tag_to_remove = data.get("tag", "") - + if not lora_name or not tag_to_remove: return aiohttp.web.Response(status=400, text="Missing parameters") - + lora_relative_path = lora_name.replace("\\", "/") lora_full_path = folder_paths.get_full_path("loras", lora_relative_path) - + if not lora_full_path or not os.path.exists(lora_full_path): return aiohttp.web.json_response({"error": "File not found"}) - + file_hash = compute_sha256(lora_full_path) db_path = os.path.join(RAYKO_LORA_DATA_DIR, f"{file_hash}.json") - + if not os.path.exists(db_path): return aiohttp.web.json_response({"error": "No data found"}) - + with open(db_path, 'r', encoding='utf-8') as f: existing_data = json.load(f) - + existing_tags = existing_data.get("trained_words", []) if tag_to_remove in existing_tags: existing_tags.remove(tag_to_remove) - + save_data = { "full_name": existing_data.get("name", ""), "trained_words": existing_tags, "description": existing_data.get("description", "") } - + save_to_rayko_db(file_hash, save_data, existing_data.get("source", "manual"), overwrite=True) - + return aiohttp.web.json_response({ "trained_words": existing_tags, "removed": tag_to_remove }) - + except Exception as e: return aiohttp.web.Response(status=500, text=str(e)) \ No newline at end of file