diff --git a/CLAUDE.md b/CLAUDE.md new file mode 100644 index 000000000000..47da538e4bfc --- /dev/null +++ b/CLAUDE.md @@ -0,0 +1,19 @@ +# CLAUDE.md + +This file provides guidance to Claude Code (claude.ai/code) when working with code in this repository. + +## Important Instructions + +**ALWAYS commit changes to git after completing any task.** After making any code changes, modifications, or additions, run: +```bash +# If not in a git repository, initialize it first: +git init + +# Then commit changes: +git add . +git commit -m "Description of changes" +``` + +## Development Workflow + +Repository-specific build, test, and lint commands will be added here as needed. \ No newline at end of file diff --git a/PR_INSTRUCTIONS.md b/PR_INSTRUCTIONS.md new file mode 100644 index 000000000000..f2446f3f34bc --- /dev/null +++ b/PR_INSTRUCTIONS.md @@ -0,0 +1,125 @@ +# Pull Request Instructions for Missing Models Downloader + +## Frontend Changes + +The frontend changes have been prepared in `/tmp/ComfyUI_frontend` on the branch `add-missing-models-downloader`. + +### To submit the PR to ComfyUI_frontend: + +1. **Fork the ComfyUI_frontend repository** on GitHub: + - Go to https://github.com/Comfy-Org/ComfyUI_frontend + - Click "Fork" in the top right + +2. **Push the changes to your fork**: + ```bash + cd /tmp/ComfyUI_frontend + git remote add fork https://github.com/YOUR_USERNAME/ComfyUI_frontend.git + git push fork add-missing-models-downloader + ``` + +3. **Create the Pull Request**: + - Go to your fork on GitHub + - Click "Pull requests" → "New pull request" + - Set base repository: `Comfy-Org/ComfyUI_frontend` base: `main` + - Set head repository: `YOUR_USERNAME/ComfyUI_frontend` compare: `add-missing-models-downloader` + - Click "Create pull request" + +4. **PR Title**: "Add Missing Models Downloader extension" + +5. **PR Description**: + ```markdown + ## Summary + + This PR adds automatic download functionality to the "Missing Models" dialog, allowing users to download missing models directly from the interface without manually searching and moving files. + + ## Features + + - ✅ Automatically adds "Download" buttons to each missing model in the dialog + - ✅ Pre-configured URLs for popular models: + - SDXL & SD 1.5 checkpoints + - VAE models (sdxl_vae, vae-ft-mse) + - LoRA models (LCM LoRAs) + - ControlNet models (canny, openpose, depth) + - Upscale models (ESRGAN, RealESRGAN) + - CLIP encoders + - Flux models + - ✅ Real-time download progress shown as percentage in button + - ✅ Custom URL prompt for unknown models + - ✅ "Download All" button for bulk downloads + - ✅ TypeScript implementation with proper typing + + ## How it works + + 1. The extension monitors for the "Missing Models" dialog + 2. When detected, it adds a download button next to each missing model + 3. Known models download immediately from pre-configured sources + 4. Unknown models prompt for a custom URL + 5. Progress is shown in real-time (0% → 100%) + 6. Models are automatically placed in the correct folders + + ## Backend Requirements + + This extension requires the ComfyUI backend to have the model download API endpoints: + - `POST /api/models/download` - Start download + - `GET /api/models/download/{task_id}` - Check status + - `POST /api/models/download/{task_id}/cancel` - Cancel download + + These endpoints are available in ComfyUI with the model_downloader module. + + ## Testing + + 1. Load a workflow with missing models + 2. The "Missing Models" dialog should appear with download buttons + 3. Click a button to download the model + 4. Progress should update in real-time + 5. Once complete, the model is ready to use + + ## Screenshots + + [Would add screenshots here of the dialog with download buttons] + ``` + +## Backend Changes + +The backend changes are already committed to your ComfyUI repository: + +### Files Added/Modified: +- `app/model_downloader.py` - Core download functionality +- `comfy_config/download_config.py` - Configuration system +- `server.py` - API endpoints +- `nodes.py` - ModelDownloader node + +### To use the complete system: + +1. **Backend (ComfyUI)**: Your changes are ready in the current branch +2. **Frontend (ComfyUI_frontend)**: Submit the PR as described above + +## Alternative: Direct Installation + +If you want to use this immediately without waiting for PR approval: + +### Frontend: +```bash +# Clone ComfyUI_frontend +git clone https://github.com/Comfy-Org/ComfyUI_frontend.git +cd ComfyUI_frontend + +# Apply the changes +git remote add temp /tmp/ComfyUI_frontend +git fetch temp +git cherry-pick 8c2f919ba128 + +# Build and use +npm install +npm run build +``` + +### Backend: +Your ComfyUI already has all necessary backend components installed and ready. + +## Notes + +- The extension is fully TypeScript compliant +- It integrates seamlessly with the existing Missing Models dialog +- No modifications to existing code, only additions +- Backwards compatible - if backend endpoints don't exist, buttons simply won't work \ No newline at end of file diff --git a/app/simple_downloader.py b/app/simple_downloader.py new file mode 100644 index 000000000000..ab972103b258 --- /dev/null +++ b/app/simple_downloader.py @@ -0,0 +1,222 @@ +"""Simple model downloader for ComfyUI.""" + +import os +import json +import uuid +import threading +import time +import folder_paths +from typing import Dict, Any, Optional +import urllib.request +import urllib.error + + +class SimpleDownloader: + """Simple downloader for ComfyUI models.""" + + def __init__(self): + self.downloads = {} + self.lock = threading.Lock() + + def create_download(self, url: str, model_type: str, filename: str) -> str: + """Create a new download task.""" + task_id = str(uuid.uuid4()) + + # SECURITY: Validate and sanitize inputs to prevent path traversal + # Sanitize model_type - remove dangerous characters but keep underscores + import re + model_type = re.sub(r'[./\\]', '', model_type) + model_type = model_type.replace('..', '') + + # Sanitize filename - use os.path.basename and remove traversal attempts + filename = os.path.basename(filename) + if '..' in filename or filename.startswith('.'): + raise ValueError("Invalid filename - no hidden files or path traversal") + + # Validate filename has allowed extension + allowed_extensions = ['.safetensors', '.ckpt', '.pt', '.pth', '.bin', '.sft'] + if not any(filename.lower().endswith(ext) for ext in allowed_extensions): + raise ValueError(f"Invalid file extension. Allowed extensions: {allowed_extensions}") + + # Whitelist of allowed model types + allowed_types = ['checkpoints', 'vae', 'loras', 'controlnet', 'clip', 'unet', + 'upscale_models', 'text_encoders', 'diffusion_models', 'embeddings'] + + # Map alternative names + type_mapping = { + 'text_encoders': 'clip', + 'diffusion_models': 'unet' + } + model_type = type_mapping.get(model_type, model_type) + + if model_type not in allowed_types: + raise ValueError(f"Invalid model type. Allowed types: {allowed_types}") + + # Determine destination folder + folder_map = { + 'checkpoints': folder_paths.get_folder_paths('checkpoints')[0], + 'vae': folder_paths.get_folder_paths('vae')[0], + 'loras': folder_paths.get_folder_paths('loras')[0], + 'controlnet': folder_paths.get_folder_paths('controlnet')[0], + 'clip': folder_paths.get_folder_paths('clip')[0], + 'unet': folder_paths.get_folder_paths('diffusion_models')[0], + 'upscale_models': folder_paths.get_folder_paths('upscale_models')[0], + 'embeddings': folder_paths.get_folder_paths('embeddings')[0] if folder_paths.get_folder_paths('embeddings') else os.path.join(folder_paths.models_dir, 'embeddings') + } + + dest_folder = folder_map.get(model_type) + if not dest_folder: + # Only allow creating folders for whitelisted types + if model_type in allowed_types: + dest_folder = os.path.join(folder_paths.models_dir, model_type) + os.makedirs(dest_folder, exist_ok=True) + else: + raise ValueError(f"Cannot find or create folder for model type: {model_type}") + + # Use safe path joining and verify result + dest_path = os.path.abspath(os.path.join(dest_folder, filename)) + + # SECURITY: Ensure destination path is within the models directory + models_base = os.path.abspath(folder_paths.models_dir) + if not dest_path.startswith(models_base): + raise ValueError("Invalid destination path - outside models directory") + + with self.lock: + self.downloads[task_id] = { + 'task_id': task_id, + 'url': url, + 'dest_path': dest_path, + 'filename': filename, + 'model_type': model_type, + 'status': 'pending', + 'progress': 0, + 'total_size': 0, + 'downloaded_size': 0, + 'error': None, + 'thread': None + } + + # Start download in background + thread = threading.Thread(target=self._download_file, args=(task_id,)) + thread.daemon = True + thread.start() + + with self.lock: + self.downloads[task_id]['thread'] = thread + self.downloads[task_id]['status'] = 'downloading' + + return task_id + + def _download_file(self, task_id: str): + """Download file in background.""" + with self.lock: + task = self.downloads.get(task_id) + if not task: + return + url = task['url'] + dest_path = task['dest_path'] + + try: + # SECURITY: Basic URL validation + from urllib.parse import urlparse + parsed = urlparse(url) + + # Allow both HTTP and HTTPS (many model URLs use HTTP redirects) + if parsed.scheme not in ['http', 'https']: + raise ValueError("Only HTTP/HTTPS URLs are allowed") + + # Basic hostname check - only block obvious local addresses + if parsed.hostname: + hostname_lower = parsed.hostname.lower() + # Block obvious local hostnames + if hostname_lower in ['localhost', '127.0.0.1', '0.0.0.0', '::1']: + raise ValueError("Downloads from localhost are not allowed") + # Block local IP ranges by pattern (not DNS lookup which can fail) + if hostname_lower.startswith(('127.', '10.', '192.168.', '172.16.', '172.17.', + '172.18.', '172.19.', '172.20.', '172.21.', + '172.22.', '172.23.', '172.24.', '172.25.', + '172.26.', '172.27.', '172.28.', '172.29.', + '172.30.', '172.31.')): + raise ValueError("Downloads from private IP ranges are not allowed") + + # Create request with headers + req = urllib.request.Request(url) + req.add_header('User-Agent', 'ComfyUI/1.0') + + # Open URL + response = urllib.request.urlopen(req, timeout=30) + + # Get total size + total_size = int(response.headers.get('Content-Length', 0)) + + with self.lock: + self.downloads[task_id]['total_size'] = total_size + + # Download in chunks + chunk_size = 8192 + downloaded = 0 + + os.makedirs(os.path.dirname(dest_path), exist_ok=True) + + with open(dest_path, 'wb') as f: + while True: + with self.lock: + if self.downloads[task_id]['status'] == 'cancelled': + break + + chunk = response.read(chunk_size) + if not chunk: + break + + f.write(chunk) + downloaded += len(chunk) + + # Update progress + with self.lock: + self.downloads[task_id]['downloaded_size'] = downloaded + if total_size > 0: + self.downloads[task_id]['progress'] = (downloaded / total_size) * 100 + + # Mark as completed + with self.lock: + if self.downloads[task_id]['status'] != 'cancelled': + self.downloads[task_id]['status'] = 'completed' + self.downloads[task_id]['progress'] = 100 + + except Exception as e: + with self.lock: + self.downloads[task_id]['status'] = 'failed' + self.downloads[task_id]['error'] = str(e) + + def get_status(self, task_id: str) -> Optional[Dict[str, Any]]: + """Get download status.""" + with self.lock: + task = self.downloads.get(task_id) + if task: + return { + 'task_id': task['task_id'], + 'status': task['status'], + 'progress': task['progress'], + 'total_size': task['total_size'], + 'downloaded_size': task['downloaded_size'], + 'error': task['error'], + 'filename': task['filename'] + } + return None + + def cancel_download(self, task_id: str) -> bool: + """Cancel a download.""" + with self.lock: + if task_id in self.downloads: + self.downloads[task_id]['status'] = 'cancelled' + return True + return False + + def get_all_downloads(self) -> list: + """Get all download statuses.""" + with self.lock: + return [self.get_status(task_id) for task_id in self.downloads.keys()] + + +# Global instance +simple_downloader = SimpleDownloader() \ No newline at end of file diff --git a/frontend_extensions/missingModelsDownloader/README.md b/frontend_extensions/missingModelsDownloader/README.md new file mode 100644 index 000000000000..c41746b0bc55 --- /dev/null +++ b/frontend_extensions/missingModelsDownloader/README.md @@ -0,0 +1,99 @@ +# Missing Models Downloader for ComfyUI + +This extension adds automatic download functionality to ComfyUI's "Missing Models" dialog, allowing users to download missing models directly from the interface with a single click. + +## Features + +- **Automatic Download Buttons**: Adds a "Download" button next to each missing model in the dialog +- **Known Model Repository**: Pre-configured URLs for popular models (SDXL, SD1.5, VAEs, LoRAs, ControlNet, etc.) +- **Custom URL Support**: Prompts for URL if model is not in the repository +- **Real-time Progress**: Shows download percentage directly in the button +- **Bulk Downloads**: "Download All" button for multiple missing models +- **Smart Detection**: Automatically detects model type and places files in correct folders + +## Installation + +### For ComfyUI Frontend Repository + +If you're using the separate ComfyUI frontend repository: + +1. Clone the frontend repository: +```bash +git clone https://github.com/Comfy-Org/ComfyUI_frontend.git +cd ComfyUI_frontend +``` + +2. Copy this extension to the extensions folder: +```bash +cp -r path/to/frontend_extensions/missingModelsDownloader web/extensions/ +``` + +3. Build and run the frontend as usual + +### For ComfyUI with Built-in Frontend + +If your ComfyUI still has the built-in frontend: + +1. Copy the extension files to ComfyUI's web extensions: +```bash +cp -r frontend_extensions/missingModelsDownloader ComfyUI/web/extensions/core/ +``` + +2. Restart ComfyUI + +## Backend Requirements + +The backend (ComfyUI server) must have the model downloader API endpoints installed. These are included in the `easy-download` branch or can be added manually: + +1. Ensure `app/model_downloader.py` exists +2. Ensure `comfy_config/download_config.py` exists +3. Ensure `server.py` includes the download API endpoints + +## How It Works + +1. When ComfyUI shows the "Missing Models" dialog, the extension automatically detects it +2. Each missing model gets a "Download" button +3. For known models, clicking downloads immediately from the pre-configured source +4. For unknown models, you'll be prompted to enter the download URL +5. Download progress is shown as a percentage in the button +6. Once complete, the model is ready to use (refresh the node to see it) + +## Supported Model Sources + +The extension includes pre-configured URLs for models from: + +- **HuggingFace**: Stable Diffusion models, VAEs, LoRAs +- **GitHub**: Upscale models (ESRGAN, RealESRGAN) +- **ComfyAnonymous**: Flux text encoders + +## Configuration + +Edit `missingModelsDownloader.js` to add more models to the repository: + +```javascript +this.modelRepositories = { + "checkpoints": { + "your_model.safetensors": "https://url/to/model" + } +} +``` + +## API Endpoints + +The extension uses these backend API endpoints: + +- `POST /api/models/download` - Start a download +- `GET /api/models/download/{task_id}` - Check download status +- `POST /api/models/download/{task_id}/cancel` - Cancel a download + +## Troubleshooting + +If the download buttons don't appear: + +1. Check the browser console for errors +2. Ensure the backend API endpoints are working: `curl http://localhost:8188/api/models/downloads` +3. Verify the extension is loaded (should see `[MissingModelsDownloader]` in console) + +## License + +MIT \ No newline at end of file diff --git a/frontend_extensions/missingModelsDownloader/missingModelsDownloader.js b/frontend_extensions/missingModelsDownloader/missingModelsDownloader.js new file mode 100644 index 000000000000..f37812afcf43 --- /dev/null +++ b/frontend_extensions/missingModelsDownloader/missingModelsDownloader.js @@ -0,0 +1,435 @@ +import { app } from "/scripts/app.js"; +import { api } from "/scripts/api.js"; + +// Missing Models Dialog Enhancer - Adds download buttons to the missing models dialog +app.registerExtension({ + name: "Comfy.MissingModelsDownloader", + + async setup() { + console.log("[MissingModelsDownloader] Extension loading..."); + + // Model repository with known URLs + this.modelRepositories = { + "checkpoints": { + "sd_xl_base_1.0.safetensors": "https://huggingface.co/stabilityai/stable-diffusion-xl-base-1.0/resolve/main/sd_xl_base_1.0.safetensors", + "sd_xl_refiner_1.0.safetensors": "https://huggingface.co/stabilityai/stable-diffusion-xl-refiner-1.0/resolve/main/sd_xl_refiner_1.0.safetensors", + "v1-5-pruned-emaonly.safetensors": "https://huggingface.co/runwayml/stable-diffusion-v1-5/resolve/main/v1-5-pruned-emaonly.safetensors", + "v1-5-pruned.safetensors": "https://huggingface.co/runwayml/stable-diffusion-v1-5/resolve/main/v1-5-pruned.safetensors", + "v2-1_768-ema-pruned.safetensors": "https://huggingface.co/stabilityai/stable-diffusion-2-1/resolve/main/v2-1_768-ema-pruned.safetensors" + }, + "vae": { + "vae-ft-mse-840000-ema-pruned.safetensors": "https://huggingface.co/stabilityai/sd-vae-ft-mse-original/resolve/main/vae-ft-mse-840000-ema-pruned.safetensors", + "sdxl_vae.safetensors": "https://huggingface.co/stabilityai/sdxl-vae/resolve/main/sdxl_vae.safetensors", + "sdxl.vae.safetensors": "https://huggingface.co/stabilityai/sdxl-vae/resolve/main/sdxl_vae.safetensors" + }, + "clip": { + "clip_l.safetensors": "https://huggingface.co/comfyanonymous/flux_text_encoders/resolve/main/clip_l.safetensors", + "t5xxl_fp16.safetensors": "https://huggingface.co/comfyanonymous/flux_text_encoders/resolve/main/t5xxl_fp16.safetensors", + "t5xxl_fp8_e4m3fn.safetensors": "https://huggingface.co/comfyanonymous/flux_text_encoders/resolve/main/t5xxl_fp8_e4m3fn.safetensors" + }, + "loras": { + "lcm-lora-sdv1-5.safetensors": "https://huggingface.co/latent-consistency/lcm-lora-sdv1-5/resolve/main/pytorch_lora_weights.safetensors", + "lcm-lora-sdxl.safetensors": "https://huggingface.co/latent-consistency/lcm-lora-sdxl/resolve/main/pytorch_lora_weights.safetensors" + }, + "controlnet": { + "control_sd15_canny.pth": "https://huggingface.co/lllyasviel/ControlNet/resolve/main/models/control_sd15_canny.pth", + "control_sd15_openpose.pth": "https://huggingface.co/lllyasviel/ControlNet/resolve/main/models/control_sd15_openpose.pth", + "control_sd15_depth.pth": "https://huggingface.co/lllyasviel/ControlNet/resolve/main/models/control_sd15_depth.pth", + "control_v11p_sd15_canny.pth": "https://huggingface.co/lllyasviel/ControlNet-v1-1/resolve/main/control_v11p_sd15_canny.pth", + "control_v11f1p_sd15_depth.pth": "https://huggingface.co/lllyasviel/ControlNet-v1-1/resolve/main/control_v11f1p_sd15_depth.pth" + }, + "upscale_models": { + "RealESRGAN_x4plus.pth": "https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.0/RealESRGAN_x4plus.pth", + "RealESRGAN_x4plus_anime_6B.pth": "https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.2.4/RealESRGAN_x4plus_anime_6B.pth", + "realesr-general-x4v3.pth": "https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.5.0/realesr-general-x4v3.pth", + "4x-UltraSharp.pth": "https://huggingface.co/lokCX/4x-Ultrasharp/resolve/main/4x-UltraSharp.pth" + }, + "unet": { + "flux1-dev.safetensors": "https://huggingface.co/black-forest-labs/FLUX.1-dev/resolve/main/flux1-dev.safetensors", + "flux1-schnell.safetensors": "https://huggingface.co/black-forest-labs/FLUX.1-schnell/resolve/main/flux1-schnell.safetensors" + } + }; + + // Active downloads tracking + this.activeDownloads = new Map(); + + // Hook into the app to monitor for missing models dialog + this.setupDialogMonitoring(); + }, + + setupDialogMonitoring() { + const self = this; + + // Monitor DOM mutations for dialog creation + const observer = new MutationObserver((mutations) => { + mutations.forEach((mutation) => { + mutation.addedNodes.forEach((node) => { + if (node.nodeType === 1) { // Element node + // Check for dialog containers + self.checkForMissingModelsDialog(node); + } + }); + }); + }); + + // Start observing + observer.observe(document.body, { + childList: true, + subtree: true + }); + + console.log("[MissingModelsDownloader] Dialog monitoring active"); + }, + + checkForMissingModelsDialog(element) { + // Look for the missing models dialog by its content + const isDialog = element.classList && ( + element.classList.contains('p-dialog') || + element.classList.contains('comfy-modal') || + element.tagName === 'DIALOG' + ); + + if (!isDialog && element.querySelector) { + const dialogs = element.querySelectorAll('dialog, .p-dialog, .comfy-modal'); + dialogs.forEach(dialog => this.checkForMissingModelsDialog(dialog)); + return; + } + + const textContent = element.textContent || ""; + + // Check for missing models dialog indicators + if (textContent.includes("Missing Models") || + textContent.includes("When loading the graph") || + textContent.includes("models were not found")) { + + console.log("[MissingModelsDownloader] Found missing models dialog"); + + // Add a small delay to ensure dialog is fully rendered + setTimeout(() => { + this.enhanceMissingModelsDialog(element); + }, 100); + } + }, + + enhanceMissingModelsDialog(dialogElement) { + // Don't enhance twice + if (dialogElement.dataset.enhancedWithDownloads) { + return; + } + dialogElement.dataset.enhancedWithDownloads = "true"; + + // Find model entries in the dialog + const modelEntries = this.findModelEntries(dialogElement); + + if (modelEntries.length === 0) { + console.log("[MissingModelsDownloader] No model entries found in dialog"); + return; + } + + console.log(`[MissingModelsDownloader] Found ${modelEntries.length} missing models`); + + // Add download button to each model + modelEntries.forEach(entry => { + this.addDownloadButton(entry); + }); + + // Add "Download All" button if multiple models + if (modelEntries.length > 1) { + this.addDownloadAllButton(dialogElement, modelEntries); + } + }, + + findModelEntries(dialogElement) { + const entries = []; + + // Look for list items containing model paths + const listItems = dialogElement.querySelectorAll('li, .model-item, [class*="missing"]'); + + listItems.forEach(item => { + const text = item.textContent || ""; + // Pattern: folder.filename or folder/filename + if (text.match(/\w+[\.\/]\w+/)) { + entries.push({ + element: item, + text: text.trim() + }); + } + }); + + // Also check for any divs or spans that might contain model names + if (entries.length === 0) { + const textElements = dialogElement.querySelectorAll('div, span, p'); + textElements.forEach(elem => { + const text = elem.textContent || ""; + if (text.match(/\w+\.\w+/) && !elem.querySelector('button')) { + // Check if this looks like a model filename + const parts = text.split(/[\.\/]/); + if (parts.length >= 2 && this.looksLikeModelName(parts[parts.length - 1])) { + entries.push({ + element: elem, + text: text.trim() + }); + } + } + }); + } + + return entries; + }, + + looksLikeModelName(filename) { + const modelExtensions = ['safetensors', 'ckpt', 'pt', 'pth', 'bin']; + const lower = filename.toLowerCase(); + return modelExtensions.some(ext => lower.includes(ext)); + }, + + addDownloadButton(entry) { + const { element, text } = entry; + + // Parse model info from text + const modelInfo = this.parseModelInfo(text); + if (!modelInfo) return; + + // Create download button + const btn = document.createElement('button'); + btn.textContent = 'Download'; + btn.style.cssText = ` + margin-left: 10px; + padding: 4px 12px; + background: #4CAF50; + color: white; + border: none; + border-radius: 4px; + cursor: pointer; + font-size: 12px; + font-weight: 500; + `; + + // Check if we have a known URL + const knownUrl = this.getKnownUrl(modelInfo.folder, modelInfo.filename); + if (knownUrl) { + btn.style.background = '#2196F3'; + btn.title = 'Download from known source'; + } + + btn.onclick = () => this.startDownload(modelInfo, btn); + + element.appendChild(btn); + entry.button = btn; + }, + + parseModelInfo(text) { + // Try different patterns + const patterns = [ + /(\w+)\.(\w+(?:\.\w+)*)/, // folder.filename + /(\w+)\/(\w+(?:\.\w+)*)/, // folder/filename + /^(\w+(?:\.\w+)*)$/ // just filename + ]; + + for (const pattern of patterns) { + const match = text.match(pattern); + if (match) { + if (match.length === 2) { + // Just filename, try to guess folder + return { + folder: this.guessFolder(match[1]), + filename: match[1] + }; + } else { + return { + folder: match[1], + filename: match[2] + }; + } + } + } + + return null; + }, + + guessFolder(filename) { + const lower = filename.toLowerCase(); + if (lower.includes('vae')) return 'vae'; + if (lower.includes('lora')) return 'loras'; + if (lower.includes('control')) return 'controlnet'; + if (lower.includes('upscale') || lower.includes('esrgan')) return 'upscale_models'; + if (lower.includes('clip')) return 'clip'; + if (lower.includes('unet') || lower.includes('flux')) return 'unet'; + return 'checkpoints'; + }, + + getKnownUrl(folder, filename) { + const repo = this.modelRepositories[folder]; + if (repo && repo[filename]) { + return repo[filename]; + } + + // Try alternate folders + const alternateFolders = { + 'text_encoders': 'clip', + 'diffusion_models': 'unet' + }; + + const altFolder = alternateFolders[folder]; + if (altFolder) { + const altRepo = this.modelRepositories[altFolder]; + if (altRepo && altRepo[filename]) { + return altRepo[filename]; + } + } + + return null; + }, + + async startDownload(modelInfo, button) { + const knownUrl = this.getKnownUrl(modelInfo.folder, modelInfo.filename); + + let url = knownUrl; + if (!url) { + // Prompt for custom URL + url = prompt( + `Enter download URL for:\n${modelInfo.filename}\n\n` + + `Model type: ${modelInfo.folder}\n\n` + + `You can find models at:\n` + + `• HuggingFace: https://huggingface.co/models\n` + + `• CivitAI: https://civitai.com/models` + ); + + if (!url || !url.trim()) return; + url = url.trim(); + } + + // Update button state + button.textContent = 'Starting...'; + button.disabled = true; + button.style.background = '#FF9800'; + + try { + const response = await api.fetchApi("/api/models/download", { + method: "POST", + headers: { "Content-Type": "application/json" }, + body: JSON.stringify({ + url: url, + model_type: modelInfo.folder, + filename: modelInfo.filename + }) + }); + + const data = await response.json(); + + if (response.ok) { + console.log(`[MissingModelsDownloader] Started download: ${data.task_id}`); + this.activeDownloads.set(data.task_id, { button, modelInfo }); + this.monitorDownload(data.task_id, button); + } else { + button.textContent = 'Failed'; + button.style.background = '#F44336'; + button.disabled = false; + alert(`Download failed: ${data.error || 'Unknown error'}`); + } + } catch (error) { + console.error('[MissingModelsDownloader] Download error:', error); + button.textContent = 'Error'; + button.style.background = '#F44336'; + button.disabled = false; + } + }, + + async monitorDownload(taskId, button) { + const checkStatus = async () => { + try { + const response = await api.fetchApi(`/api/models/download/${taskId}`); + const status = await response.json(); + + if (!response.ok) { + button.textContent = 'Failed'; + button.style.background = '#F44336'; + button.disabled = false; + this.activeDownloads.delete(taskId); + return; + } + + switch (status.status) { + case 'completed': + button.textContent = '✓ Downloaded'; + button.style.background = '#4CAF50'; + button.disabled = true; + this.activeDownloads.delete(taskId); + + // Refresh model lists + if (app.refreshComboInNodes) { + app.refreshComboInNodes(); + } + break; + + case 'downloading': + const progress = Math.round(status.progress || 0); + button.textContent = `${progress}%`; + button.style.background = '#2196F3'; + setTimeout(checkStatus, 2000); + break; + + case 'failed': + button.textContent = 'Failed'; + button.style.background = '#F44336'; + button.disabled = false; + this.activeDownloads.delete(taskId); + break; + + default: + button.textContent = status.status; + setTimeout(checkStatus, 2000); + } + } catch (error) { + console.error('[MissingModelsDownloader] Status check error:', error); + button.textContent = 'Error'; + button.style.background = '#F44336'; + button.disabled = false; + this.activeDownloads.delete(taskId); + } + }; + + checkStatus(); + }, + + addDownloadAllButton(dialogElement, modelEntries) { + // Find dialog footer or create button container + let buttonContainer = dialogElement.querySelector('.p-dialog-footer, .dialog-footer'); + + if (!buttonContainer) { + buttonContainer = document.createElement('div'); + buttonContainer.style.cssText = ` + padding: 15px; + border-top: 1px solid #444; + text-align: center; + margin-top: 15px; + `; + dialogElement.appendChild(buttonContainer); + } + + const downloadAllBtn = document.createElement('button'); + downloadAllBtn.textContent = `Download All (${modelEntries.length} models)`; + downloadAllBtn.style.cssText = ` + padding: 8px 16px; + background: #FF9800; + color: white; + border: none; + border-radius: 4px; + cursor: pointer; + font-weight: bold; + margin: 0 5px; + `; + + downloadAllBtn.onclick = () => { + modelEntries.forEach(entry => { + if (entry.button && !entry.button.disabled) { + entry.button.click(); + } + }); + downloadAllBtn.disabled = true; + downloadAllBtn.textContent = 'Downloads Started'; + }; + + buttonContainer.appendChild(downloadAllBtn); + } +}); \ No newline at end of file diff --git a/frontend_extensions/missingModelsDownloader/package.json b/frontend_extensions/missingModelsDownloader/package.json new file mode 100644 index 000000000000..1f764964dd72 --- /dev/null +++ b/frontend_extensions/missingModelsDownloader/package.json @@ -0,0 +1,11 @@ +{ + "name": "comfyui-missing-models-downloader", + "version": "1.0.0", + "description": "Adds download buttons to the Missing Models dialog in ComfyUI", + "main": "missingModelsDownloader.js", + "author": "ComfyUI", + "license": "MIT", + "comfyui": { + "extension_type": "web" + } +} \ No newline at end of file diff --git a/server.py b/server.py index 80e9d3fa78a0..9cf5dc1ca915 100644 --- a/server.py +++ b/server.py @@ -35,6 +35,7 @@ from app.user_manager import UserManager from app.model_manager import ModelFileManager from app.custom_node_manager import CustomNodeManager +from app.simple_downloader import simple_downloader from typing import Optional, Union from api_server.routes.internal.internal_routes import InternalRoutes from protocol import BinaryEventTypes @@ -789,6 +790,102 @@ async def post_history(request): return web.Response(status=200) + @routes.post("/models/download") + async def start_model_download(request): + """Start a new model download.""" + try: + json_data = await request.json() + url = json_data.get("url") + model_type = json_data.get("model_type", "checkpoints") + filename = json_data.get("filename") + + if not url: + return web.json_response({"error": "URL is required"}, status=400) + + # SECURITY: Validate URL format + from urllib.parse import urlparse + try: + parsed_url = urlparse(url) + if parsed_url.scheme not in ['http', 'https']: + return web.json_response({"error": "Only HTTP/HTTPS URLs are allowed"}, status=400) + except Exception: + return web.json_response({"error": "Invalid URL format"}, status=400) + + # SECURITY: Sanitize model_type + import re + if not re.match(r'^[a-zA-Z0-9_-]+$', model_type): + return web.json_response({"error": "Invalid model type format"}, status=400) + + if not filename: + # Extract filename from URL + filename = url.split('/')[-1].split('?')[0] + if not filename: + filename = "model.safetensors" + + # SECURITY: Sanitize filename - allow more characters but still safe + import os + filename = os.path.basename(filename) + # Block path traversal attempts but allow spaces, parens, brackets, etc + if '..' in filename or '/' in filename or '\\' in filename or filename.startswith('.'): + return web.json_response({"error": "Invalid filename format"}, status=400) + + # Create download task (simple_downloader now has additional validation) + task_id = simple_downloader.create_download(url, model_type, filename) + + # Return task ID and initial status + status = simple_downloader.get_status(task_id) + return web.json_response(status) + + except ValueError as e: + # Return validation errors from simple_downloader + return web.json_response({"error": str(e)}, status=400) + except Exception as e: + logging.error(f"Error starting download: {e}") + return web.json_response({"error": "Internal server error"}, status=500) + + @routes.get("/models/download/{task_id}") + async def get_download_status(request): + """Get status of a specific download.""" + task_id = request.match_info.get("task_id") + status = simple_downloader.get_status(task_id) + + if status is None: + return web.json_response({"error": "Download task not found"}, status=404) + + return web.json_response(status) + + @routes.get("/models/downloads") + async def get_all_downloads(request): + """Get status of all downloads.""" + downloads = simple_downloader.get_all_downloads() + return web.json_response(downloads) + + @routes.post("/models/download/{task_id}/pause") + async def pause_download(request): + """Pause a download - placeholder.""" + return web.json_response({"error": "Model download functionality not available"}, status=501) + + @routes.post("/models/download/{task_id}/resume") + async def resume_download(request): + """Resume a paused download - placeholder.""" + return web.json_response({"error": "Model download functionality not available"}, status=501) + + @routes.post("/models/download/{task_id}/cancel") + async def cancel_download(request): + """Cancel a download.""" + task_id = request.match_info.get("task_id") + success = simple_downloader.cancel_download(task_id) + + if not success: + return web.json_response({"error": "Failed to cancel download"}, status=400) + + return web.json_response({"success": True}) + + @routes.get("/models/download/history") + async def get_download_history(request): + """Get download history - placeholder.""" + return web.json_response([]) + async def setup(self): timeout = aiohttp.ClientTimeout(total=None) # no timeout self.client_session = aiohttp.ClientSession(timeout=timeout) diff --git a/start_with_custom_frontend.sh b/start_with_custom_frontend.sh new file mode 100755 index 000000000000..ccdb4d516465 --- /dev/null +++ b/start_with_custom_frontend.sh @@ -0,0 +1,4 @@ +#!/bin/bash + +# Start ComfyUI with custom frontend +python main.py --front-end-root /Users/fragmede/projects/llm/ComfyUI_frontend/dist "$@" \ No newline at end of file diff --git a/startup_config.json b/startup_config.json new file mode 100644 index 000000000000..e64d85aaf02d --- /dev/null +++ b/startup_config.json @@ -0,0 +1,3 @@ +{ + "force_mps": true +} \ No newline at end of file