nodes/Raykosan-ComfyUI_RaykoStudio
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

replacing the code in Rayko_LoRA_Loader.py

Browse Source
This commit is contained in:
Raykosan
2026-09-11 16:48:47 +03:00
parent 510d918958
commit bba80db1de
1 changed files with 102 additions and 53 deletions
Download Patch File Download Diff File Expand all files Collapse all files
Rayko_LoRA_Loader.py
+102 -53
View File
@@ -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))