diff --git a/CUSTOM_ROUTES.md b/CUSTOM_ROUTES.md deleted file mode 100644 index be1b50e..0000000 --- a/CUSTOM_ROUTES.md +++ /dev/null @@ -1,250 +0,0 @@ -# SageUtils Custom API Routes - -SageUtils adds custom HTTP routes to ComfyUI that expose the SageCache data. These routes allow you to access model metadata, file hashes, and cache statistics programmatically. - -## Available Endpoints - -All endpoints are available under the `/sage_cache/` prefix and can also be accessed with the `/api/sage_cache/` prefix. - -### 1. Get Cache Info -**GET** `/sage_cache/info` - -Returns the complete contents of SageCache.info as JSON. This contains model metadata, Civitai information, and cache details for all models. - -**Example Response:** -```json -{ - "abc123hash": { - "hash": "abc123hash", - "civitai": "true", - "modelId": "12345", - "name": "Model Name", - "baseModel": "SDXL", - "lastUsed": "2025-06-22T10:30:00" - } -} -``` - -### 2. Get Cache Hash Mapping -**GET** `/sage_cache/hash` - -Returns the complete contents of SageCache.hash as JSON. This contains the mapping from file paths to their SHA256 hashes. - -**Example Response:** -```json -{ - "/path/to/model1.safetensors": "abc123hash", - "/path/to/model2.ckpt": "def456hash" -} -``` - -### 3. Get Cache Statistics -**GET** `/sage_cache/stats` - -Returns statistics about the SageCache including file counts, Civitai status, and cache file locations. - -**Example Response:** -```json -{ - "total_files": 25, - "total_info_entries": 23, - "cache_files": { - "hash_path": "/path/to/sage_cache_hash.json", - "info_path": "/path/to/sage_cache_info.json", - "main_path": "/path/to/sage_cache.json", - "ollama_models_path": "/path/to/sage_cache_ollama.json" - }, - "civitai_stats": { - "found_on_civitai": 18, - "not_found_on_civitai": 5 - } -} -``` - -### 4. Get File Information by Hash -**GET** `/sage_cache/file/{file_hash}` - -Returns detailed information for a specific file hash, including all file paths that use this hash. - -**Example Response:** -```json -{ - "hash": "abc123hash", - "info": { - "hash": "abc123hash", - "civitai": "true", - "modelId": "12345", - "name": "Model Name", - "baseModel": "SDXL", - "lastUsed": "2025-06-22T10:30:00" - }, - "file_paths": [ - "/path/to/model1.safetensors" - ] -} -``` - -### 5. Get File Information by Path -**GET** `/sage_cache/path?file_path={file_path}` - -Returns detailed information for a file by its path. Looks up the hash from the file path in cache.hash, then retrieves the info from cache.info. - -**Parameters:** -- `file_path` (query parameter, required): The full file path to look up - -**Example Response:** -```json -{ - "file_path": "/path/to/model1.safetensors", - "hash": "abc123hash", - "info": { - "hash": "abc123hash", - "civitai": "true", - "modelId": "12345", - "name": "Model Name", - "baseModel": "SDXL", - "lastUsed": "2025-06-22T10:30:00" - }, - "all_paths_with_same_hash": [ - "/path/to/model1.safetensors" - ] -} -``` - -## Usage Examples - -### Using cURL - -```bash -# Get all cache info -curl http://localhost:8188/sage_cache/info - -# Get hash mappings -curl http://localhost:8188/sage_cache/hash - -# Get cache statistics -curl http://localhost:8188/sage_cache/stats - -# Get specific file info by hash -curl http://localhost:8188/sage_cache/file/abc123hash456def - -# Get specific file info by path -curl "http://localhost:8188/sage_cache/path?file_path=/path/to/model1.safetensors" -``` - -### Using JavaScript/Fetch - -```javascript -// Get cache statistics -async function getCacheStats() { - const response = await fetch('/api/sage_cache/stats'); - const stats = await response.json(); - console.log('Cache stats:', stats); -} - -// Get info for a specific hash -async function getFileInfo(hash) { - const response = await fetch(`/api/sage_cache/file/${hash}`); - const info = await response.json(); - console.log('File info:', info); -} - -// Get info for a specific file path -async function getFileInfoByPath(filePath) { - const response = await fetch(`/api/sage_cache/path?file_path=${encodeURIComponent(filePath)}`); - const info = await response.json(); - console.log('File info by path:', info); -} -``` - -### Using Python - -```python -import requests - -# Get all cache info -response = requests.get('http://localhost:8188/sage_cache/info') -cache_info = response.json() - -# Get cache statistics -response = requests.get('http://localhost:8188/sage_cache/stats') -stats = response.json() -print(f"Total files: {stats['total_files']}") -print(f"Found on Civitai: {stats['civitai_stats']['found_on_civitai']}") - -# Get file info by hash -file_hash = 'abc123hash' -response = requests.get(f'http://localhost:8188/sage_cache/file/{file_hash}') -file_info = response.json() -print(f"File info for hash {file_hash}: {file_info}") - -# Get file info by path -file_path = '/path/to/model1.safetensors' -response = requests.get('http://localhost:8188/sage_cache/path', params={'file_path': file_path}) -file_info = response.json() -print(f"File info for path {file_path}: {file_info}") -``` - -## Error Handling - -All endpoints return appropriate HTTP status codes: -- **200**: Success -- **400**: Bad Request (e.g., missing file hash) -- **404**: Not Found (e.g., hash doesn't exist) -- **500**: Internal Server Error - -Error responses include an `error` field with a descriptive message: - -```json -{ - "error": "No information found for hash: invalid_hash" -} -``` - -The implementation includes graceful fallback if PromptServer is not available during development or ComfyUI startup issues. -``` - -## Implementation Details - -### Route Registration - -The custom routes are registered using `PromptServer.instance.routes` to add new endpoints to ComfyUI's existing web server. Routes are automatically registered when the custom node imports during ComfyUI startup. - -### Cache Access - -Each request calls `cache.load()` to ensure fresh data is returned, respecting any changes made to the cache files while ComfyUI is running. - -### File Structure - -```text -comfyui_sageutils/ -├── server_routes.py # Custom route implementations -├── CUSTOM_ROUTES.md # This documentation -├── test_routes.py # Testing script -├── demo_routes.py # Demonstration script -├── __init__.py # Modified to import routes -└── utils/ - └── model_cache.py # Contains SageCache class -``` - -### Technical Notes - -- The implementation is read-only (no routes modify cache data) -- Routes respect ComfyUI's security model -- Compatible with existing ComfyUI routing infrastructure -- No additional dependencies required beyond what SageUtils already uses -- Routes only register if PromptServer is available (fails silently otherwise) -- All routes are also available with the `/api/` prefix for consistency with ComfyUI standards - -### Testing and Validation - -When ComfyUI is running properly, you should see the message "SageUtils custom routes loaded successfully!" in the console, indicating the routes are available. - -Use the provided test script to validate functionality: - -```bash -cd /path/to/comfyui_sageutils -python test_routes.py -``` - -The cache is automatically loaded when endpoints are accessed, and all endpoints are read-only and do not modify the cache. The endpoints respect the same security model as other ComfyUI routes. diff --git a/FILE_MANAGEMENT.md b/FILE_MANAGEMENT.md deleted file mode 100644 index 0e39684..0000000 --- a/FILE_MANAGEMENT.md +++ /dev/null @@ -1,118 +0,0 @@ -# SageUtils File and Path Management - -This document explains the reorganized file and path management system in SageUtils. - -## Overview - -The file and path management has been centralized into three main modules: - -1. **`path_manager.py`** - Centralized path management -2. **`config_manager.py`** - Configuration file management (refactored) -3. **`model_cache.py`** - Model cache management (refactored) -4. **`sage_utils.py`** - Convenience module with commonly used functions - -## Key Improvements - -### 1. Centralized Path Management -- All paths are now managed in `SagePathManager` -- Consistent directory creation and path resolution -- Single source of truth for all file locations - -### 2. Unified File Operations -- `SageFileManager` provides atomic JSON operations -- Consistent error handling and backup strategies -- Reusable file I/O methods - -### 3. Consistent Initialization -- Config files are automatically copied from assets to user directory if missing -- Cache files are created in user directory as needed -- All necessary directories are created automatically - -## File Locations - -### Default Asset Files (read-only templates) -- `/assets/config.json` - Default settings -- `/assets/sage_styles.json` - Default styles -- `/assets/llm_prompts.json` - Default LLM prompts - -### User Files (modifiable copies) -- `/user/default/SageUtils/config.json` - User settings -- `/user/default/SageUtils/sage_styles.json` - User styles -- `/user/default/SageUtils/llm_prompts.json` - User LLM prompts - -### User Override Files (optional) -- `/user/default/SageUtils/config_user.json` - User setting overrides -- `/user/default/SageUtils/sage_styles_user.json` - User style overrides -- `/user/default/SageUtils/llm_prompts_user.json` - User prompt overrides - -### Cache Files -- `/user/default/SageUtils/sage_cache.json` - Legacy cache (converted automatically) -- `/user/default/SageUtils/sage_cache_hash.json` - File path to hash mapping -- `/user/default/SageUtils/sage_cache_info.json` - Hash to metadata mapping -- `/user/default/SageUtils/sage_cache_ollama.json` - Ollama model cache - -### Backup Files -- `/user/default/SageUtils/backup/` - Automatic backups of all files - -### Wildcard Files -- `/user/default/SageUtils/wildcards/` - User wildcard text files - -### Notes Files -- `/user/default/SageUtils/notes/` - User notes and documentation files - -## Usage Examples - -### Using the convenience module: -```python -from .utils.sage_utils import get_user_path, load_json, save_json - -# Get paths -user_config = get_user_path("my_config.json") - -# Load/save JSON -data = load_json(user_config, "my config") -save_json(user_config, {"key": "value"}, "my config") -``` - -### Using path manager directly: -```python -from .utils.path_manager import path_manager, file_manager - -# Get specific paths -backup_path = path_manager.get_backup_file_path("backup-2024.json") -asset_path = path_manager.get_asset_file_path("defaults.json") - -# File operations -file_manager.ensure_user_config_file("new_config") -config = file_manager.load_config_with_overrides("config") -``` - -### Using config manager: -```python -from .utils.config_manager import ConfigManager - -# Create and use a config manager -my_config = ConfigManager("my_config") -data = my_config.load() # Loads with user overrides -my_config.save(data, user_override=True) # Save to override file -``` - -## Migration Notes - -### For Existing Code -- Replace direct file path references with path_manager calls -- Use file_manager for JSON operations instead of manual file handling -- Config files will be automatically migrated on first load - -### Breaking Changes -- Old global path variables (`sage_users_path`, etc.) are now in `sage_utils.py` -- Manual directory creation is no longer needed -- File operations should use the centralized managers - -## Benefits - -1. **Consistency** - All files are handled the same way -2. **Reliability** - Atomic operations and error handling -3. **Maintainability** - Single place to modify path/file logic -4. **Backup Safety** - Automatic backups with deduplication -5. **User Experience** - Automatic setup and migration diff --git a/demo_routes.py b/demo_routes.py deleted file mode 100644 index 93b3ac4..0000000 --- a/demo_routes.py +++ /dev/null @@ -1,88 +0,0 @@ -#!/usr/bin/env python3 -""" -Simple demonstration of the SageUtils custom routes implementation. -This shows the structure and functionality without requiring a full ComfyUI environment. -""" - -def demonstrate_routes(): - """Demonstrate the custom routes implementation for SageUtils""" - - print("=== SageUtils Custom Routes Demo ===\n") - - print("✓ Successfully implemented custom routes for SageUtils in ComfyUI!") - print("✓ Routes are registered via server_routes.py module") - print("✓ Routes are imported in __init__.py during custom node loading") - print("✓ Implementation follows ComfyUI documentation patterns\n") - - print("📁 Files created:") - print(" - server_routes.py: Custom route implementations") - print(" - CUSTOM_ROUTES.md: Documentation for the new routes") - print(" - test_routes.py: Test script for validation") - print(" - Modified __init__.py: Added route import\n") - - print("🌐 Available Routes (when ComfyUI is running):") - routes = [ - ("GET", "/sage_cache/info", "Returns complete SageCache.info as JSON"), - ("GET", "/sage_cache/hash", "Returns complete SageCache.hash mapping as JSON"), - ("GET", "/sage_cache/stats", "Returns cache statistics and counts"), - ("GET", "/sage_cache/file/{hash}", "Returns info for specific file hash"), - ("GET", "/sage_cache/path?file_path={path}", "Returns info for file by path") - ] - - for method, path, description in routes: - print(f" {method:4} {path:25} - {description}") - - print(f"\n🔗 All routes also available with /api prefix") - print(f" Example: http://localhost:8188/api/sage_cache/stats") - - print("\n📊 Sample Response Structure:") - - # Sample stats response - sample_stats = { - "total_files": 25, - "total_info_entries": 23, - "cache_files": { - "hash_path": "/user/default/SageUtils/sage_cache_hash.json", - "info_path": "/user/default/SageUtils/sage_cache_info.json" - }, - "civitai_stats": { - "found_on_civitai": 18, - "not_found_on_civitai": 5 - } - } - - import json - print(" /sage_cache/stats response:") - print(" " + json.dumps(sample_stats, indent=4).replace('\n', '\n ')) - - print("\n🛡️ Error Handling:") - print(" - Graceful import handling (warns if PromptServer unavailable)") - print(" - Try/catch blocks around cache operations") - print(" - Proper HTTP status codes (400, 404, 500)") - print(" - Descriptive error messages in JSON responses") - - print("\n🔧 Implementation Details:") - print(" - Uses PromptServer.instance.routes for registration") - print(" - Imports happen at custom node load time") - print(" - Cache is loaded on each request to ensure freshness") - print(" - Compatible with ComfyUI's existing route structure") - - print("\n📖 Usage Examples:") - print(" # Get cache statistics") - print(" curl http://localhost:8188/sage_cache/stats") - print("") - print(" # Get all model hashes") - print(" curl http://localhost:8188/sage_cache/hash") - print("") - print(" # Get info for specific model") - print(" curl http://localhost:8188/sage_cache/file/abc123def456") - print("") - print(" # Get info for model by file path") - print(' curl "http://localhost:8188/sage_cache/path?file_path=/path/to/model.safetensors"') - - print("\n✅ Implementation Complete!") - print(" The custom routes are ready to use once ComfyUI starts successfully.") - print(" They will appear alongside ComfyUI's built-in routes like /prompt, /queue, etc.") - -if __name__ == "__main__": - demonstrate_routes() diff --git a/js/sage.js b/js/sage.js index 13a8177..8a567ae 100644 --- a/js/sage.js +++ b/js/sage.js @@ -14,6 +14,9 @@ import { setupViewTextOrAnythingNode } from "./nodes/viewAnything.js"; import { setupViewNotesNode } from "./nodes/viewNotes.js"; import { setupModelInfoDisplayNode } from "./nodes/modelInfoDisplay.js"; +// Import settings UI +import "./settingsUI.js"; + app.registerExtension({ name: "arcum42.sage.utils", async setup() { diff --git a/js/settingsUI.js b/js/settingsUI.js new file mode 100644 index 0000000..f025c3e --- /dev/null +++ b/js/settingsUI.js @@ -0,0 +1,181 @@ +// SageUtils Settings Integration +// Integrates SageUtils settings with ComfyUI's native settings system + +import { app } from "../../../../scripts/app.js"; +import { api } from "../../../../scripts/api.js"; + +// All of the setting types are based on PrimeVue components. +// https://primevue.org/ + +// Props described in the PrimeVue documentation can be defined +// for ComfyUI settings by adding them in an attrs field. + +// Helper functions for API communication +async function loadSageSettings() { + try { + const response = await api.fetchApi('/sage_utils/settings'); + if (response.ok) { + const data = await response.json(); + if (data.success) { + return data.settings; + } else { + console.error('Server returned error:', data.error); + if (data.details) { + console.error('Error details:', data.details); + } + } + } else { + console.error('HTTP error:', response.status, response.statusText); + } + throw new Error('Failed to load settings from server'); + } catch (error) { + console.error('Error loading SageUtils settings:', error); + // Return fallback settings structure + return { + enable_ollama: { current_value: true }, + enable_lmstudio: { current_value: true }, + ollama_use_custom_url: { current_value: false }, + ollama_custom_url: { current_value: "" }, + lmstudio_use_custom_url: { current_value: false }, + lmstudio_custom_url: { current_value: "" } + }; + } +} + +async function saveSageSetting(key, value) { + try { + const response = await api.fetchApi('/sage_utils/settings', { + method: 'POST', + headers: { + 'Content-Type': 'application/json', + }, + body: JSON.stringify({ [key]: value }) + }); + + const data = await response.json(); + if (data.success) { + console.log(`Successfully saved setting ${key}: ${value}`); + if (data.errors && data.errors.length > 0) { + console.warn('Warnings during setting save:', data.errors); + } + } else { + console.error(`Failed to save setting ${key}:`, data.error); + } + return data.success; + } catch (error) { + console.error(`Error saving SageUtils setting ${key}:`, error); + return false; + } +} + +// Register SageUtils settings with ComfyUI's native settings system +app.registerExtension({ + name: "arcum42.sage.utils.settings", + + async setup() { + console.log("Setting up SageUtils settings integration..."); + + // Load current settings from server + const serverSettings = await loadSageSettings(); + if (!serverSettings) { + console.warn('Could not load SageUtils settings from server - settings will use defaults'); + return; + } + + console.log('Loaded SageUtils settings from server:', serverSettings); + + // Set current values for all settings from server + // Map backend keys to the new frontend setting IDs + const keyToIdMap = { + 'enable_lmstudio': 'SageUtils.LLM Providers.enable_lmstudio', + 'enable_ollama': 'SageUtils.LLM Providers.enable_ollama', + 'ollama_custom_url': 'SageUtils.Local Custom Ollama URL.ollama_custom_url', + 'ollama_use_custom_url': 'SageUtils.Local Custom Ollama URL.ollama_use_custom_url', + 'lmstudio_custom_url': 'SageUtils.Local Custom LM Studio URL.lmstudio_custom_url', + 'lmstudio_use_custom_url': 'SageUtils.Local Custom LM Studio URL.lmstudio_use_custom_url' + }; + + for (const [key, settingInfo] of Object.entries(serverSettings)) { + const settingId = keyToIdMap[key]; + if (settingId && settingInfo.current_value !== undefined) { + try { + await app.extensionManager.setting.set(settingId, settingInfo.current_value); + console.log(`Set initial value for ${settingId}:`, settingInfo.current_value); + } catch (error) { + console.warn(`Could not set initial value for ${settingId}:`, error); + } + } + } + + console.log(`SageUtils settings integration completed`); + }, + + settings: [ + { + id: "SageUtils.LLM Providers.enable_lmstudio", + name: "Enable LM Studio Integration", + type: "boolean", + defaultValue: true, + tooltip: "Enable LM Studio LLM integration", + onChange: async (newVal, oldVal) => { + console.log(`LM Studio integration changed from ${oldVal} to ${newVal}`); + await saveSageSetting('enable_lmstudio', newVal); + } + }, + { + id: "SageUtils.LLM Providers.enable_ollama", + name: "Enable Ollama Integration", + type: "boolean", + defaultValue: true, + tooltip: "Enable Ollama LLM integration", + onChange: async (newVal, oldVal) => { + console.log(`Ollama integration changed from ${oldVal} to ${newVal}`); + await saveSageSetting('enable_ollama', newVal); + } + }, + { + id: "SageUtils.Local Custom Ollama URL.ollama_custom_url", + name: "Address", + type: "text", + defaultValue: "", + tooltip: "Custom URL for Ollama service (e.g., 'http://localhost:11434')", + onChange: async (newVal, oldVal) => { + console.log(`Ollama custom URL changed from '${oldVal}' to '${newVal}'`); + await saveSageSetting('ollama_custom_url', newVal); + } + }, + { + id: "SageUtils.Local Custom Ollama URL.ollama_use_custom_url", + name: "Enable Custom URL", + type: "boolean", + defaultValue: false, + tooltip: "Use a custom URL for Ollama instead of the default", + onChange: async (newVal, oldVal) => { + console.log(`Ollama custom URL setting changed from ${oldVal} to ${newVal}`); + await saveSageSetting('ollama_use_custom_url', newVal); + } + }, + { + id: "SageUtils.Local Custom LM Studio URL.lmstudio_custom_url", + name: "Address", + type: "text", + defaultValue: "", + tooltip: "Custom URL for LM Studio service (e.g., 'http://localhost:1234')", + onChange: async (newVal, oldVal) => { + console.log(`LM Studio custom URL changed from '${oldVal}' to '${newVal}'`); + await saveSageSetting('lmstudio_custom_url', newVal); + } + }, + { + id: "SageUtils.Local Custom LM Studio URL.lmstudio_use_custom_url", + name: "Enable Custom URL", + type: "boolean", + defaultValue: false, + tooltip: "Use a custom URL for LM Studio instead of the default", + onChange: async (newVal, oldVal) => { + console.log(`LM Studio custom URL setting changed from ${oldVal} to ${newVal}`); + await saveSageSetting('lmstudio_use_custom_url', newVal); + } + } + ] +}); diff --git a/server_routes.py b/server_routes.py index b0a26c1..bb5d505 100644 --- a/server_routes.py +++ b/server_routes.py @@ -1,17 +1,128 @@ """ -Custom routes for SageUtils to expose cache data via HTTP endpoints. +Custom routes for SageUtils to expose cache data and settings via HTTP endpoints. """ try: from server import PromptServer from aiohttp import web from .utils.model_cache import cache + from .utils.settings import get_settings, SETTINGS_SCHEMA # Check if PromptServer instance is available if hasattr(PromptServer, 'instance') and PromptServer.instance is not None: # Get the PromptServer instance routes = PromptServer.instance.routes + # Settings management routes + @routes.get('/sage_utils/settings') + async def get_sage_settings(request): + """ + Returns all SageUtils settings with their current values and schema information. + """ + try: + import json + settings = get_settings() + settings_info = settings.list_all_settings() + + # Double-check that the result is JSON serializable + json.dumps(settings_info) # This will raise an exception if not serializable + + return web.json_response({ + "success": True, + "settings": settings_info + }) + except (TypeError, ValueError) as e: + return web.json_response( + {"success": False, "error": f"JSON serialization error: {str(e)}"}, + status=500 + ) + except Exception as e: + import traceback + error_details = traceback.format_exc() + print(f"SageUtils settings error: {error_details}") + return web.json_response( + {"success": False, "error": f"Failed to retrieve settings: {str(e)}", "details": error_details}, + status=500 + ) + + @routes.post('/sage_utils/settings') + async def update_sage_settings(request): + """ + Updates SageUtils settings. Expects JSON body with setting key-value pairs. + """ + try: + data = await request.json() + settings = get_settings() + + updated_settings = [] + errors = [] + + for key, value in data.items(): + if key in SETTINGS_SCHEMA: + try: + if settings.set(key, value): + updated_settings.append(key) + except Exception as e: + errors.append(f"Failed to set '{key}': {str(e)}") + else: + errors.append(f"Unknown setting: '{key}'") + + # Save if any settings were updated + if updated_settings: + if settings.save(): + # Check if LLM-related settings were updated and trigger lazy initialization + llm_settings = {'enable_ollama', 'enable_lmstudio', 'custom_ollama_url', 'custom_lmstudio_url'} + if any(setting in llm_settings for setting in updated_settings): + try: + from .utils.llm_wrapper import ensure_llm_initialized + ensure_llm_initialized() + except Exception as llm_e: + errors.append(f"Warning: Failed to initialize LLM services: {str(llm_e)}") + + return web.json_response({ + "success": True, + "updated": updated_settings, + "errors": errors, + "message": f"Updated {len(updated_settings)} setting(s)" + }) + else: + return web.json_response({ + "success": False, + "error": "Failed to save settings", + "updated": updated_settings, + "errors": errors + }, status=500) + else: + return web.json_response({ + "success": False, + "error": "No valid settings to update", + "errors": errors + }, status=400) + + except Exception as e: + return web.json_response( + {"success": False, "error": f"Failed to update settings: {str(e)}"}, + status=500 + ) + + @routes.post('/sage_utils/settings/reset') + async def reset_sage_settings(request): + """ + Resets all SageUtils settings to their default values. + """ + try: + settings = get_settings() + settings.reset_to_defaults() + return web.json_response({ + "success": True, + "message": "All settings reset to defaults" + }) + except Exception as e: + return web.json_response( + {"success": False, "error": f"Failed to reset settings: {str(e)}"}, + status=500 + ) + @routes.get('/sage_cache/info') async def get_sage_cache_info(request): """ diff --git a/test_routes.py b/test_routes.py deleted file mode 100644 index 1b5b41e..0000000 --- a/test_routes.py +++ /dev/null @@ -1,351 +0,0 @@ -#!/usr/bin/env python3 -""" -Test script for SageUtils custom routes. -This demonstrates how the custom routes would work when ComfyUI is running. -""" - -import sys -import os -import json -from typing import Dict, Any -from unittest.mock import Mock, MagicMock - -# Add ComfyUI path -sys.path.insert(0, '/home/ai/programs/comfyui') - -class MockRequest: - """Mock aiohttp request object""" - def __init__(self, match_info=None): - self.match_info = match_info or {} - -class MockWeb: - """Mock aiohttp web module""" - @staticmethod - def json_response(data, status=200): - return { - "status": status, - "content_type": "application/json", - "data": data - } - -class MockRoutes: - """Mock routes object that stores registered routes""" - def __init__(self): - self.routes = {} - - def get(self, path): - def decorator(func): - self.routes[f"GET {path}"] = func - print(f"✓ Registered route: GET {path}") - return func - return decorator - - async def call_route(self, method, path, request=None): - """Call a registered route for testing""" - route_key = f"{method} {path}" - if route_key in self.routes: - return await self.routes[route_key](request or MockRequest()) - else: - return MockWeb.json_response({"error": "Route not found"}, 404) - -class MockPromptServer: - """Mock PromptServer class""" - def __init__(self): - self.routes = MockRoutes() - - @classmethod - def create_instance(cls): - cls.instance = cls() - return cls.instance - -def test_sage_routes(): - """Test the SageUtils custom routes""" - print("=== Testing SageUtils Custom Routes ===\n") - - # Create mock server instance - prompt_server = MockPromptServer.create_instance() - - # Mock the modules that would be imported - sys.modules['server'] = Mock() - sys.modules['server'].PromptServer = MockPromptServer - sys.modules['aiohttp'] = Mock() - sys.modules['aiohttp'].web = MockWeb - - # Import and setup our cache - print("1. Loading SageUtils cache...") - try: - # Suppress hash warnings - import warnings - warnings.filterwarnings("ignore") - - from custom_nodes.comfyui_sageutils.utils.model_cache import cache - cache.load() - print(f" ✓ Cache loaded: {len(cache.hash)} files, {len(cache.info)} info entries") - - # Add some dummy data for testing if cache is empty - if len(cache.hash) == 0: - print(" Adding dummy test data...") - test_hash = "abc123def456789" - test_path = "/models/test_model.safetensors" - cache.hash[test_path] = test_hash - cache.info[test_hash] = { - "hash": test_hash, - "civitai": "true", - "modelId": "12345", - "name": "Test Model", - "baseModel": "SDXL", - "lastUsed": "2025-06-22T10:30:00" - } - print(f" ✓ Added test data: 1 file, 1 info entry") - - except Exception as e: - print(f" ✗ Error loading cache: {e}") - import traceback - traceback.print_exc() - return False - - # Mock the imports for our routes module - original_modules = {} - mock_modules = { - 'server': Mock(), - 'aiohttp': Mock(), - '.utils.model_cache': Mock() - } - - # Setup mocks - mock_modules['server'].PromptServer = MockPromptServer - mock_modules['aiohttp'].web = MockWeb - mock_modules['.utils.model_cache'].cache = cache - - print("\n2. Testing route registration...") - - # Manually register routes like our module would - routes = prompt_server.routes - - @routes.get('/sage_cache/info') - async def get_sage_cache_info(request): - try: - cache.load() - return MockWeb.json_response(cache.info) - except Exception as e: - return MockWeb.json_response( - {"error": f"Failed to retrieve cache info: {str(e)}"}, - status=500 - ) - - @routes.get('/sage_cache/hash') - async def get_sage_cache_hash(request): - try: - cache.load() - return MockWeb.json_response(cache.hash) - except Exception as e: - return MockWeb.json_response( - {"error": f"Failed to retrieve cache hash: {str(e)}"}, - status=500 - ) - - @routes.get('/sage_cache/stats') - async def get_sage_cache_stats(request): - try: - cache.load() - - stats = { - "total_files": len(cache.hash), - "total_info_entries": len(cache.info), - "cache_files": { - "hash_path": str(cache.hash_path), - "info_path": str(cache.info_path), - "main_path": str(cache.main_path), - "ollama_models_path": str(cache.ollama_models_path) - } - } - - # Count files by civitai status - civitai_found = 0 - civitai_not_found = 0 - - for info_data in cache.info.values(): - civitai_status = info_data.get("civitai", "False") - if isinstance(civitai_status, str): - if civitai_status.lower() in ["true", "1"]: - civitai_found += 1 - else: - civitai_not_found += 1 - elif civitai_status is True: - civitai_found += 1 - else: - civitai_not_found += 1 - - stats["civitai_stats"] = { - "found_on_civitai": civitai_found, - "not_found_on_civitai": civitai_not_found - } - - return MockWeb.json_response(stats) - except Exception as e: - return MockWeb.json_response( - {"error": f"Failed to retrieve cache stats: {str(e)}"}, - status=500 - ) - - @routes.get('/sage_cache/file/{file_hash}') - async def get_sage_cache_file_info(request): - try: - file_hash = request.match_info.get('file_hash', '') - if not file_hash: - return MockWeb.json_response( - {"error": "No file hash provided"}, - status=400 - ) - - cache.load() - - file_info = cache.info.get(file_hash) - if file_info is None: - return MockWeb.json_response( - {"error": f"No information found for hash: {file_hash}"}, - status=404 - ) - - # Also include which file paths use this hash - file_paths = [path for path, hash_val in cache.hash.items() if hash_val == file_hash] - - result = { - "hash": file_hash, - "info": file_info, - "file_paths": file_paths - } - - return MockWeb.json_response(result) - except Exception as e: - return MockWeb.json_response( - {"error": f"Failed to retrieve file info: {str(e)}"}, - status=500 - ) - - @routes.get('/sage_cache/path') - async def get_sage_cache_path_info(request): - try: - # Get file path from query parameter (mocked) - file_path = getattr(request, 'file_path', '/models/test_model.safetensors') - if not file_path: - return MockWeb.json_response( - {"error": "No file_path parameter provided"}, - status=400 - ) - - cache.load() - - # Look up hash for this file path - file_hash = cache.hash.get(file_path) - if file_hash is None: - return MockWeb.json_response( - {"error": f"No hash found for file path: {file_path}"}, - status=404 - ) - - # Get the info for this hash - file_info = cache.info.get(file_hash) - if file_info is None: - return MockWeb.json_response( - {"error": f"No information found for hash {file_hash} (path: {file_path})"}, - status=404 - ) - - # Also include all file paths that use this same hash - all_file_paths = [path for path, hash_val in cache.hash.items() if hash_val == file_hash] - - result = { - "file_path": file_path, - "hash": file_hash, - "info": file_info, - "all_paths_with_same_hash": all_file_paths - } - - return MockWeb.json_response(result) - except Exception as e: - return MockWeb.json_response( - {"error": f"Failed to retrieve path info: {str(e)}"}, - status=500 - ) - - print("\n3. Testing routes...") - - # Test each route - async def test_routes(): - # Test /sage_cache/stats - print("\n Testing GET /sage_cache/stats") - response = await routes.call_route("GET", "/sage_cache/stats") - print(f" Status: {response['status']}") - if response['status'] == 200: - data = response['data'] - print(f" Total files: {data['total_files']}") - print(f" Total info entries: {data['total_info_entries']}") - print(f" Civitai found: {data['civitai_stats']['found_on_civitai']}") - print(f" Civitai not found: {data['civitai_stats']['not_found_on_civitai']}") - - # Test /sage_cache/hash - print("\n Testing GET /sage_cache/hash") - response = await routes.call_route("GET", "/sage_cache/hash") - print(f" Status: {response['status']}") - if response['status'] == 200: - data = response['data'] - print(f" Hash entries: {len(data)}") - if data: - first_path = list(data.keys())[0] - print(f" Sample: {first_path[:50]}... -> {data[first_path][:16]}...") - - # Test /sage_cache/info - print("\n Testing GET /sage_cache/info") - response = await routes.call_route("GET", "/sage_cache/info") - print(f" Status: {response['status']}") - if response['status'] == 200: - data = response['data'] - print(f" Info entries: {len(data)}") - if data: - first_hash = list(data.keys())[0] - print(f" Sample hash: {first_hash}") - info = data[first_hash] - print(f" Sample info: civitai={info.get('civitai')}, name={info.get('name', 'N/A')}") - - # Test /sage_cache/file/{hash} - if cache.info: - test_hash = list(cache.info.keys())[0] - print(f"\n Testing GET /sage_cache/file/{test_hash}") - request = MockRequest({"file_hash": test_hash}) - response = await routes.call_route("GET", "/sage_cache/file/{file_hash}", request) - print(f" Status: {response['status']}") - if response['status'] == 200: - data = response['data'] - print(f" Hash: {data['hash']}") - print(f" File paths: {len(data['file_paths'])}") - print(f" Info keys: {list(data['info'].keys())}") - - # Test /sage_cache/path - print("\n Testing GET /sage_cache/path") - response = await routes.call_route("GET", "/sage_cache/path") - print(f" Status: {response['status']}") - if response['status'] == 200: - data = response['data'] - print(f" File path: {data['file_path']}") - print(f" Hash: {data['hash']}") - print(f" Info keys: {list(data['info'].keys())}") - print(f" All paths with same hash: {len(data['all_paths_with_same_hash'])}") - - # Run the async tests - import asyncio - asyncio.run(test_routes()) - - print("\n=== Test Complete ===") - print("✓ All routes registered and tested successfully!") - print("\nWhen ComfyUI is running, these routes will be available at:") - print(" - http://localhost:8188/sage_cache/info") - print(" - http://localhost:8188/sage_cache/hash") - print(" - http://localhost:8188/sage_cache/stats") - print(" - http://localhost:8188/sage_cache/file/{hash}") - print(" - Or with /api prefix for each route") - - return True - -if __name__ == "__main__": - test_sage_routes() diff --git a/utils/llm_wrapper.py b/utils/llm_wrapper.py index 6d65ba7..0016c1c 100644 --- a/utils/llm_wrapper.py +++ b/utils/llm_wrapper.py @@ -2,6 +2,10 @@ import logging from .helpers_image import tensor_to_base64, tensor_to_temp_image from .llm_cache import get_llm_cache +# Initialization flags to track if services have been initialized +_ollama_initialized = False +_lmstudio_initialized = False + # Attempt to import ollama, if available. Set a flag if it is not available. try: import ollama @@ -153,6 +157,9 @@ def get_ollama_models() -> list[str]: def ollama_generate_vision(model: str, prompt: str, keep_alive: float = 0.0, images=None, options=None) -> str: """Generate a response from an Ollama vision model.""" + # Ensure Ollama is initialized before use + ensure_ollama_initialized() + if not OLLAMA_AVAILABLE or ollama_client is None: raise ImportError("Ollama is not available. Please install it to use this function.") vision_models = get_ollama_vision_models() @@ -191,6 +198,9 @@ def ollama_generate_vision(model: str, prompt: str, keep_alive: float = 0.0, ima def ollama_generate(model: str, prompt: str, keep_alive: float = 0.0, options=None) -> str: """Generate a response from an Ollama model.""" + # Ensure Ollama is initialized before use + ensure_ollama_initialized() + if not OLLAMA_AVAILABLE or ollama_client is None: raise ImportError("Ollama is not available. Please install it to use this function.") models = get_ollama_models() @@ -357,6 +367,9 @@ def get_lmstudio_vision_models() -> list[str]: def lmstudio_generate_vision(model: str, prompt: str, keep_alive: int = 0, images=None, options=None) -> str: """Generate a response from an LM Studio vision model.""" + # Ensure LM Studio is initialized before use + ensure_lmstudio_initialized() + if not LMSTUDIO_AVAILABLE or lms is None: raise ImportError("LM Studio is not available. Please install it to use this function.") model_list = get_lmstudio_vision_models() @@ -391,6 +404,9 @@ def lmstudio_generate_vision(model: str, prompt: str, keep_alive: int = 0, image def lmstudio_generate(model: str, prompt: str, keep_alive: int = 0, options=None) -> str: """Generate a response from an LM Studio model.""" + # Ensure LM Studio is initialized before use + ensure_lmstudio_initialized() + if not LMSTUDIO_AVAILABLE or lms is None: raise ImportError("LM Studio is not available. Please install it to use this function.") model_list = get_lmstudio_models() @@ -472,73 +488,71 @@ def lmstudio_generate_vision_refine(model: str, prompt: str, images=None, option return ("", "") def init_ollama(): - """Initialize Ollama if available. Print config values for Ollama.""" - global ollama_client + """Initialize Ollama client""" + global ollama_client, _ollama_initialized + from .settings import get_setting - if not OLLAMA_AVAILABLE or ollama is None: - logging.info("Ollama is not available; skipping Ollama initialization.") - return + # Check if Ollama is available and enabled + if not OLLAMA_AVAILABLE: + logging.warning("Ollama library is not available.") + return False + + if not get_setting("enable_ollama", False): + logging.info("Ollama is disabled in settings.") + _ollama_initialized = False + ollama_client = None + return False try: - # Try to use new settings system - try: - from .settings import get_setting, is_feature_enabled - - if not is_feature_enabled('enable_ollama'): - logging.info("Ollama is disabled in settings; skipping initialization.") - return - - use_custom_url = get_setting('ollama_use_custom_url', False) - custom_url = get_setting('ollama_custom_url', '') - except ImportError: - # Fallback to old config system - from . import config_manager - config = config_manager.settings_manager.data or {} - use_custom_url = config.get('ollama_use_custom_url', False) - custom_url = config.get('ollama_custom_url', '') - - if use_custom_url and custom_url: + # Get custom URL or use default + custom_url = get_setting("custom_ollama_url", "http://localhost:11434") + if custom_url and custom_url.strip(): ollama_client = ollama.Client(host=custom_url) - logging.info(f"Ollama client initialized with custom host: {custom_url}") + logging.info(f"Ollama client initialized with custom URL: {custom_url}") else: ollama_client = ollama.Client() - logging.info("Ollama client initialized with default settings.") + logging.info("Ollama client initialized with default URL") + + _ollama_initialized = True + return True except Exception as e: - ollama_client = None logging.error(f"Failed to initialize Ollama client: {e}") + _ollama_initialized = False + ollama_client = None + return False def init_lmstudio(): """Initialize LM Studio if available. Print config values for LM Studio.""" + global _lmstudio_initialized + from .settings import get_setting + if not LMSTUDIO_AVAILABLE or lms is None: logging.info("LM Studio is not available; skipping LM Studio initialization.") - return + _lmstudio_initialized = False + return False + + # Check if LM Studio is enabled + if not get_setting("enable_lmstudio", False): + logging.info("LM Studio is disabled in settings; skipping initialization.") + _lmstudio_initialized = False + return False try: - # Try to use new settings system - try: - from .settings import get_setting, is_feature_enabled - - if not is_feature_enabled('enable_lmstudio'): - logging.info("LM Studio is disabled in settings; skipping initialization.") - return - - use_custom_url = get_setting('lmstudio_use_custom_url', False) - custom_url = get_setting('lmstudio_custom_url', '') - except ImportError: - # Fallback to old config system - from . import config_manager - config = config_manager.settings_manager.data or {} - use_custom_url = config.get('lmstudio_use_custom_url', False) - custom_url = config.get('lmstudio_custom_url', '') + custom_url = get_setting('custom_lmstudio_url', '') - if use_custom_url and custom_url: + if custom_url and custom_url.strip(): lm_client = lms.get_default_client(custom_url) logging.info(f"LM Studio client configured with custom URL: {custom_url}") else: logging.info("LM Studio using default configuration.") + + _lmstudio_initialized = True + return True except Exception as e: logging.error(f"Failed to configure LM Studio: {e}") + _lmstudio_initialized = False + return False def init_llm(): @@ -546,3 +560,32 @@ def init_llm(): init_ollama() init_lmstudio() logging.info("LLM clients initialized.") + + +def ensure_ollama_initialized(): + """Ensure Ollama is initialized if it's enabled in settings and not already initialized.""" + global _ollama_initialized + from .settings import get_setting + + if get_setting("enable_ollama", False) and not _ollama_initialized: + logging.info("Ollama is enabled but not initialized, initializing now...") + return init_ollama() + return _ollama_initialized + + +def ensure_lmstudio_initialized(): + """Ensure LM Studio is initialized if it's enabled in settings and not already initialized.""" + global _lmstudio_initialized + from .settings import get_setting + + if get_setting("enable_lmstudio", False) and not _lmstudio_initialized: + logging.info("LM Studio is enabled but not initialized, initializing now...") + return init_lmstudio() + return _lmstudio_initialized + + +def ensure_llm_initialized(): + """Ensure all enabled LLM services are initialized.""" + ollama_ok = ensure_ollama_initialized() + lmstudio_ok = ensure_lmstudio_initialized() + return ollama_ok or lmstudio_ok diff --git a/utils/settings.py b/utils/settings.py index b613fd2..658d4b3 100644 --- a/utils/settings.py +++ b/utils/settings.py @@ -178,6 +178,9 @@ class SageSettings: schema_entry = SETTINGS_SCHEMA[key].copy() schema_entry["current_value"] = self._settings.get(key) + # Convert type to string representation for JSON serialization + if "type" in schema_entry: + schema_entry["type"] = schema_entry["type"].__name__ return schema_entry def list_all_settings(self) -> Dict[str, Dict[str, Any]]: @@ -186,6 +189,7 @@ class SageSettings: for key in SETTINGS_SCHEMA.keys(): setting_info = self.get_setting_info(key) if setting_info is not None: + # The setting_info is already JSON-safe from get_setting_info() result[key] = setting_info return result