120 lines
4.5 KiB
Python
120 lines
4.5 KiB
Python
import os
|
|
from aiohttp import web
|
|
from server import PromptServer
|
|
import folder_paths
|
|
|
|
WEB_DIRECTORY = "./js"
|
|
|
|
# Required for ComfyUI to recognize this as a valid extension
|
|
NODE_CLASS_MAPPINGS = {}
|
|
NODE_DISPLAY_NAME_MAPPINGS = {}
|
|
|
|
class NotesManager:
|
|
def __init__(self):
|
|
self.notes_dir = os.path.join(os.path.dirname(__file__), "notes")
|
|
os.makedirs(self.notes_dir, exist_ok=True)
|
|
|
|
def get_notes_list(self) -> list[str]:
|
|
notes : list[str] = []
|
|
for filename in os.listdir(self.notes_dir):
|
|
if filename.endswith('.md'):
|
|
notes.append(filename[:-3]) # Remove .md extension
|
|
return sorted(notes)
|
|
|
|
def get_note_content(self, note_name: str):
|
|
file_path = os.path.join(self.notes_dir, f"{note_name}.md")
|
|
if os.path.exists(file_path):
|
|
with open(file_path, 'r', encoding='utf-8') as f:
|
|
return f.read()
|
|
return ""
|
|
|
|
def save_note(self, note_name: str, content: str):
|
|
file_path = os.path.join(self.notes_dir, f"{note_name}.md")
|
|
with open(file_path, 'w', encoding='utf-8') as f:
|
|
f.write(content)
|
|
|
|
def delete_note(self, note_name: str):
|
|
file_path = os.path.join(self.notes_dir, f"{note_name}.md")
|
|
if os.path.exists(file_path):
|
|
os.remove(file_path)
|
|
|
|
notes_manager = NotesManager()
|
|
|
|
@PromptServer.instance.routes.get("/notes_sidebar/list")
|
|
async def get_notes_list(request):
|
|
try:
|
|
notes = notes_manager.get_notes_list()
|
|
return web.json_response({"notes": notes})
|
|
except Exception as e:
|
|
return web.json_response({"error": str(e)}, status=500)
|
|
|
|
@PromptServer.instance.routes.get("/notes_sidebar/note/{note_name}")
|
|
async def get_note(request):
|
|
try:
|
|
note_name = request.match_info['note_name']
|
|
content = notes_manager.get_note_content(note_name)
|
|
return web.json_response({"content": content})
|
|
except Exception as e:
|
|
return web.json_response({"error": str(e)}, status=500)
|
|
|
|
@PromptServer.instance.routes.post("/notes_sidebar/note/{note_name}")
|
|
async def save_note(request):
|
|
try:
|
|
note_name = request.match_info['note_name']
|
|
data = await request.json()
|
|
content = data.get('content', '')
|
|
notes_manager.save_note(note_name, content)
|
|
return web.json_response({"success": True})
|
|
except Exception as e:
|
|
return web.json_response({"error": str(e)}, status=500)
|
|
|
|
@PromptServer.instance.routes.delete("/notes_sidebar/note/{note_name}")
|
|
async def delete_note(request):
|
|
try:
|
|
note_name = request.match_info['note_name']
|
|
notes_manager.delete_note(note_name)
|
|
return web.json_response({"success": True})
|
|
except Exception as e:
|
|
return web.json_response({"error": str(e)}, status=500)
|
|
|
|
@PromptServer.instance.routes.get("/notes_sidebar/loras")
|
|
async def get_loras_list(request):
|
|
try:
|
|
loras = folder_paths.get_filename_list("loras")
|
|
return web.json_response({"loras": loras})
|
|
except Exception as e:
|
|
return web.json_response({"error": str(e)}, status=500)
|
|
|
|
@PromptServer.instance.routes.get("/notes_sidebar/loras_with_metadata")
|
|
async def get_loras_with_metadata(request):
|
|
try:
|
|
loras = folder_paths.get_filename_list("loras")
|
|
|
|
loras_with_metadata = []
|
|
for lora in loras:
|
|
# Get the full path to the lora file
|
|
lora_full_path = folder_paths.get_full_path("loras", lora)
|
|
|
|
metadata = ""
|
|
if lora_full_path:
|
|
# Get the directory and base name
|
|
lora_dir = os.path.dirname(lora_full_path)
|
|
base_name = os.path.splitext(os.path.basename(lora_full_path))[0]
|
|
txt_file_path = os.path.join(lora_dir, f"{base_name}.txt")
|
|
|
|
if os.path.exists(txt_file_path):
|
|
try:
|
|
with open(txt_file_path, 'r', encoding='utf-8') as f:
|
|
lines = f.read().strip().split('\n')
|
|
# Join multiple lines with " | "
|
|
metadata = " | ".join(line.strip() for line in lines if line.strip())
|
|
except Exception as e:
|
|
print(f"Error reading metadata for {lora}: {e}")
|
|
|
|
loras_with_metadata.append({ "name": lora, "metadata": metadata })
|
|
|
|
return web.json_response({"loras": loras_with_metadata})
|
|
except Exception as e:
|
|
return web.json_response({"error": str(e)}, status=500)
|
|
|
|
__all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS', 'WEB_DIRECTORY'] |