Missing model downloader

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.

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

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()