Merge branch 'comfyanonymous:master' into master

comfy/ldm/modules/attention.py

@@ -335,7 +335,7 @@ def attention_xformers(q, k, v, heads, mask=None, attn_precision=None, skip_resh
             disabled_xformers = True
 
     if disabled_xformers:
-        return attention_pytorch(q, k, v, heads, mask)
+        return attention_pytorch(q, k, v, heads, mask, skip_reshape=skip_reshape)
 
     if skip_reshape:
          q, k, v = map(

main.py

@@ -261,6 +261,7 @@ if __name__ == "__main__":
         call_on_start = startup_server
 
     try:
+        loop.run_until_complete(server.setup())
         loop.run_until_complete(run(server, address=args.listen, port=args.port, verbose=not args.dont_print_server, call_on_start=call_on_start))
     except KeyboardInterrupt:
         logging.info("\nStopped server")

model_filemanager/__init__.py

@@ -0,0 +1,2 @@
+# model_manager/__init__.py
+from .download_models import download_model, DownloadModelStatus, DownloadStatusType, create_model_path, check_file_exists, track_download_progress, validate_model_subdirectory, validate_filename

model_filemanager/download_models.py

@@ -0,0 +1,240 @@
+from __future__ import annotations
+import aiohttp
+import os
+import traceback
+import logging
+from folder_paths import models_dir
+import re
+from typing import Callable, Any, Optional, Awaitable, Dict
+from enum import Enum
+import time
+from dataclasses import dataclass
+
+
+class DownloadStatusType(Enum):
+    PENDING = "pending"
+    IN_PROGRESS = "in_progress"
+    COMPLETED = "completed"
+    ERROR = "error"
+
+@dataclass
+class DownloadModelStatus():
+    status: str
+    progress_percentage: float
+    message: str
+    already_existed: bool = False
+
+    def __init__(self, status: DownloadStatusType, progress_percentage: float, message: str, already_existed: bool):
+        self.status = status.value  # Store the string value of the Enum
+        self.progress_percentage = progress_percentage
+        self.message = message
+        self.already_existed = already_existed
+    
+    def to_dict(self) -> Dict[str, Any]:
+        return {
+            "status": self.status,
+            "progress_percentage": self.progress_percentage,
+            "message": self.message,
+            "already_existed": self.already_existed
+        }
+
+async def download_model(model_download_request: Callable[[str], Awaitable[aiohttp.ClientResponse]],
+                         model_name: str,  
+                         model_url: str, 
+                         model_sub_directory: str,
+                         progress_callback: Callable[[str, DownloadModelStatus], Awaitable[Any]],
+                         progress_interval: float = 1.0) -> DownloadModelStatus:
+    """
+    Download a model file from a given URL into the models directory.
+
+    Args:
+        model_download_request (Callable[[str], Awaitable[aiohttp.ClientResponse]]): 
+            A function that makes an HTTP request. This makes it easier to mock in unit tests.
+        model_name (str): 
+            The name of the model file to be downloaded. This will be the filename on disk.
+        model_url (str): 
+            The URL from which to download the model.
+        model_sub_directory (str): 
+            The subdirectory within the main models directory where the model 
+            should be saved (e.g., 'checkpoints', 'loras', etc.).
+        progress_callback (Callable[[str, DownloadModelStatus], Awaitable[Any]]): 
+            An asynchronous function to call with progress updates.
+
+    Returns:
+        DownloadModelStatus: The result of the download operation.
+    """
+    if not validate_model_subdirectory(model_sub_directory):
+        return DownloadModelStatus(
+            DownloadStatusType.ERROR, 
+            0,
+            "Invalid model subdirectory", 
+            False
+        )
+
+    if not validate_filename(model_name):
+        return DownloadModelStatus(
+            DownloadStatusType.ERROR, 
+            0,
+            "Invalid model name", 
+            False
+        )
+
+    file_path, relative_path = create_model_path(model_name, model_sub_directory, models_dir)
+    existing_file = await check_file_exists(file_path, model_name, progress_callback, relative_path)
+    if existing_file:
+        return existing_file
+
+    try:
+        status = DownloadModelStatus(DownloadStatusType.PENDING, 0, f"Starting download of {model_name}", False)
+        await progress_callback(relative_path, status)
+
+        response = await model_download_request(model_url)
+        if response.status != 200:
+            error_message = f"Failed to download {model_name}. Status code: {response.status}"
+            logging.error(error_message)
+            status = DownloadModelStatus(DownloadStatusType.ERROR, 0, error_message, False)
+            await progress_callback(relative_path, status)
+            return DownloadModelStatus(DownloadStatusType.ERROR, 0, error_message, False)
+
+        return await track_download_progress(response, file_path, model_name, progress_callback, relative_path, progress_interval)
+
+    except Exception as e:
+        logging.error(f"Error in downloading model: {e}")
+        return await handle_download_error(e, model_name, progress_callback, relative_path)
+    
+
+def create_model_path(model_name: str, model_directory: str, models_base_dir: str) -> tuple[str, str]:
+    full_model_dir = os.path.join(models_base_dir, model_directory)
+    os.makedirs(full_model_dir, exist_ok=True)
+    file_path = os.path.join(full_model_dir, model_name)
+
+    # Ensure the resulting path is still within the base directory
+    abs_file_path = os.path.abspath(file_path)
+    abs_base_dir = os.path.abspath(str(models_base_dir))
+    if os.path.commonprefix([abs_file_path, abs_base_dir]) != abs_base_dir:
+        raise Exception(f"Invalid model directory: {model_directory}/{model_name}")
+
+
+    relative_path = '/'.join([model_directory, model_name])
+    return file_path, relative_path
+
+async def check_file_exists(file_path: str, 
+                            model_name: str, 
+                            progress_callback: Callable[[str, DownloadModelStatus], Awaitable[Any]], 
+                            relative_path: str) -> Optional[DownloadModelStatus]:
+    if os.path.exists(file_path):
+        status = DownloadModelStatus(DownloadStatusType.COMPLETED, 100, f"{model_name} already exists", True)
+        await progress_callback(relative_path, status)
+        return status
+    return None
+
+
+async def track_download_progress(response: aiohttp.ClientResponse, 
+                                  file_path: str, 
+                                  model_name: str, 
+                                  progress_callback: Callable[[str, DownloadModelStatus], Awaitable[Any]], 
+                                  relative_path: str, 
+                                  interval: float = 1.0) -> DownloadModelStatus:
+    try:
+        total_size = int(response.headers.get('Content-Length', 0))
+        downloaded = 0
+        last_update_time = time.time()
+
+        async def update_progress():
+            nonlocal last_update_time
+            progress = (downloaded / total_size) * 100 if total_size > 0 else 0
+            status = DownloadModelStatus(DownloadStatusType.IN_PROGRESS, progress, f"Downloading {model_name}", False)
+            await progress_callback(relative_path, status)
+            last_update_time = time.time()
+
+        with open(file_path, 'wb') as f:
+            chunk_iterator = response.content.iter_chunked(8192)
+            while True:
+                try:
+                    chunk = await chunk_iterator.__anext__()
+                except StopAsyncIteration:
+                    break
+                f.write(chunk)
+                downloaded += len(chunk)
+                
+                if time.time() - last_update_time >= interval:
+                    await update_progress()
+
+        await update_progress()
+        
+        logging.info(f"Successfully downloaded {model_name}. Total downloaded: {downloaded}")
+        status = DownloadModelStatus(DownloadStatusType.COMPLETED, 100, f"Successfully downloaded {model_name}", False)
+        await progress_callback(relative_path, status)
+
+        return status
+    except Exception as e:
+        logging.error(f"Error in track_download_progress: {e}")
+        logging.error(traceback.format_exc())
+        return await handle_download_error(e, model_name, progress_callback, relative_path)
+
+async def handle_download_error(e: Exception, 
+                                model_name: str, 
+                                progress_callback: Callable[[str, DownloadModelStatus], Any], 
+                                relative_path: str) -> DownloadModelStatus:
+    error_message = f"Error downloading {model_name}: {str(e)}"
+    status = DownloadModelStatus(DownloadStatusType.ERROR, 0, error_message, False)
+    await progress_callback(relative_path, status)
+    return status
+
+def validate_model_subdirectory(model_subdirectory: str) -> bool:
+    """
+    Validate that the model subdirectory is safe to install into. 
+    Must not contain relative paths, nested paths or special characters
+    other than underscores and hyphens.
+
+    Args:
+        model_subdirectory (str): The subdirectory for the specific model type.
+
+    Returns:
+        bool: True if the subdirectory is safe, False otherwise.
+    """
+    if len(model_subdirectory) > 50:
+        return False
+
+    if '..' in model_subdirectory or '/' in model_subdirectory:
+        return False
+
+    if not re.match(r'^[a-zA-Z0-9_-]+$', model_subdirectory):
+        return False
+
+    return True
+
+def validate_filename(filename: str)-> bool:
+    """
+    Validate a filename to ensure it's safe and doesn't contain any path traversal attempts.
+    
+    Args:
+    filename (str): The filename to validate
+
+    Returns:
+    bool: True if the filename is valid, False otherwise
+    """
+    if not filename.lower().endswith(('.sft', '.safetensors')):
+        return False
+
+    # Check if the filename is empty, None, or just whitespace
+    if not filename or not filename.strip():
+        return False
+
+    # Check for any directory traversal attempts or invalid characters
+    if any(char in filename for char in ['..', '/', '\\', '\n', '\r', '\t', '\0']):
+        return False
+
+    # Check if the filename starts with a dot (hidden file)
+    if filename.startswith('.'):
+        return False
+
+    # Use a whitelist of allowed characters
+    if not re.match(r'^[a-zA-Z0-9_\-. ]+$', filename):
+        return False
+
+    # Ensure the filename isn't too long
+    if len(filename) > 255:
+        return False
+
+    return True

nodes.py

@@ -47,11 +47,18 @@ MAX_RESOLUTION=16384
 class CLIPTextEncode:
     @classmethod
     def INPUT_TYPES(s):
-        return {"required": {"text": ("STRING", {"multiline": True, "dynamicPrompts": True}), "clip": ("CLIP", )}}
+        return {
+            "required": {
+                "text": ("STRING", {"multiline": True, "dynamicPrompts": True, "tooltip": "The text to be encoded."}), 
+                "clip": ("CLIP", {"tooltip": "The CLIP model used for encoding the text."})
+            }
+        }
     RETURN_TYPES = ("CONDITIONING",)
+    OUTPUT_TOOLTIPS = ("A conditioning containing the embedded text used to guide the diffusion model.",)
     FUNCTION = "encode"
 
     CATEGORY = "conditioning"
+    DESCRIPTION = "Encodes a text prompt using a CLIP model into an embedding that can be used to guide the diffusion model towards generating specific images."
 
     def encode(self, clip, text):
         tokens = clip.tokenize(text)
@@ -260,11 +267,18 @@ class ConditioningSetTimestepRange:
 class VAEDecode:
     @classmethod
     def INPUT_TYPES(s):
-        return {"required": { "samples": ("LATENT", ), "vae": ("VAE", )}}
+        return {
+            "required": { 
+                "samples": ("LATENT", {"tooltip": "The latent to be decoded."}), 
+                "vae": ("VAE", {"tooltip": "The VAE model used for decoding the latent."})
+            }
+        }
     RETURN_TYPES = ("IMAGE",)
+    OUTPUT_TOOLTIPS = ("The decoded image.",)
     FUNCTION = "decode"
 
     CATEGORY = "latent"
+    DESCRIPTION = "Decodes latent images back into pixel space images."
 
     def decode(self, vae, samples):
         return (vae.decode(samples["samples"]), )
@@ -506,12 +520,19 @@ class CheckpointLoader:
 class CheckpointLoaderSimple:
     @classmethod
     def INPUT_TYPES(s):
-        return {"required": { "ckpt_name": (folder_paths.get_filename_list("checkpoints"), ),
+        return {
-                             }}
+            "required": { 
+                "ckpt_name": (folder_paths.get_filename_list("checkpoints"), {"tooltip": "The name of the checkpoint (model) to load."}),
+            }
+        }
     RETURN_TYPES = ("MODEL", "CLIP", "VAE")
+    OUTPUT_TOOLTIPS = ("The model used for denoising latents.", 
+                       "The CLIP model used for encoding text prompts.", 
+                       "The VAE model used for encoding and decoding images to and from latent space.")
     FUNCTION = "load_checkpoint"
 
     CATEGORY = "loaders"
+    DESCRIPTION = "Loads a diffusion model checkpoint, diffusion models are used to denoise latents."
 
     def load_checkpoint(self, ckpt_name):
         ckpt_path = folder_paths.get_full_path("checkpoints", ckpt_name)
@@ -582,16 +603,22 @@ class LoraLoader:
 
     @classmethod
     def INPUT_TYPES(s):
-        return {"required": { "model": ("MODEL",),
+        return {
-                              "clip": ("CLIP", ),
+            "required": { 
-                              "lora_name": (folder_paths.get_filename_list("loras"), ),
+                "model": ("MODEL", {"tooltip": "The diffusion model the LoRA will be applied to."}),
-                              "strength_model": ("FLOAT", {"default": 1.0, "min": -100.0, "max": 100.0, "step": 0.01}),
+                "clip": ("CLIP", {"tooltip": "The CLIP model the LoRA will be applied to."}),
-                              "strength_clip": ("FLOAT", {"default": 1.0, "min": -100.0, "max": 100.0, "step": 0.01}),
+                "lora_name": (folder_paths.get_filename_list("loras"), {"tooltip": "The name of the LoRA."}),
-                              }}
+                "strength_model": ("FLOAT", {"default": 1.0, "min": -100.0, "max": 100.0, "step": 0.01, "tooltip": "How strongly to modify the diffusion model. This value can be negative."}),
+                "strength_clip": ("FLOAT", {"default": 1.0, "min": -100.0, "max": 100.0, "step": 0.01, "tooltip": "How strongly to modify the CLIP model. This value can be negative."}),
+            }
+        }
+    
     RETURN_TYPES = ("MODEL", "CLIP")
+    OUTPUT_TOOLTIPS = ("The modified diffusion model.", "The modified CLIP model.")
     FUNCTION = "load_lora"
 
     CATEGORY = "loaders"
+    DESCRIPTION = "LoRAs are used to modify diffusion and CLIP models, altering the way in which latents are denoised such as applying styles. Multiple LoRA nodes can be linked together."
 
     def load_lora(self, model, clip, lora_name, strength_model, strength_clip):
         if strength_model == 0 and strength_clip == 0:
@@ -1033,13 +1060,19 @@ class EmptyLatentImage:
 
     @classmethod
     def INPUT_TYPES(s):
-        return {"required": { "width": ("INT", {"default": 512, "min": 16, "max": MAX_RESOLUTION, "step": 8}),
+        return {
-                              "height": ("INT", {"default": 512, "min": 16, "max": MAX_RESOLUTION, "step": 8}),
+            "required": { 
-                              "batch_size": ("INT", {"default": 1, "min": 1, "max": 4096})}}
+                "width": ("INT", {"default": 512, "min": 16, "max": MAX_RESOLUTION, "step": 8, "tooltip": "The width of the latent images in pixels."}),
+                "height": ("INT", {"default": 512, "min": 16, "max": MAX_RESOLUTION, "step": 8, "tooltip": "The height of the latent images in pixels."}),
+                "batch_size": ("INT", {"default": 1, "min": 1, "max": 4096, "tooltip": "The number of latent images in the batch."})
+            }
+        }
     RETURN_TYPES = ("LATENT",)
+    OUTPUT_TOOLTIPS = ("The empty latent image batch.",)
     FUNCTION = "generate"
 
     CATEGORY = "latent"
+    DESCRIPTION = "Create a new batch of empty latent images to be denoised via sampling."
 
     def generate(self, width, height, batch_size=1):
         latent = torch.zeros([batch_size, 4, height // 8, width // 8], device=self.device)
@@ -1359,24 +1392,27 @@ def common_ksampler(model, seed, steps, cfg, sampler_name, scheduler, positive,
 class KSampler:
     @classmethod
     def INPUT_TYPES(s):
-        return {"required":
+        return {
-                    {"model": ("MODEL",),
+            "required": {
-                    "seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}),
+                "model": ("MODEL", {"tooltip": "The model used for denoising the input latent."}),
-                    "steps": ("INT", {"default": 20, "min": 1, "max": 10000}),
+                "seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff, "tooltip": "The random seed used for creating the noise."}),
-                    "cfg": ("FLOAT", {"default": 8.0, "min": 0.0, "max": 100.0, "step":0.1, "round": 0.01}),
+                "steps": ("INT", {"default": 20, "min": 1, "max": 10000, "tooltip": "The number of steps used in the denoising process."}),
-                    "sampler_name": (comfy.samplers.KSampler.SAMPLERS, ),
+                "cfg": ("FLOAT", {"default": 8.0, "min": 0.0, "max": 100.0, "step":0.1, "round": 0.01, "tooltip": "The Classifier-Free Guidance scale balances creativity and adherence to the prompt. Higher values result in images more closely matching the prompt however too high values will negatively impact quality."}),
-                    "scheduler": (comfy.samplers.KSampler.SCHEDULERS, ),
+                "sampler_name": (comfy.samplers.KSampler.SAMPLERS, {"tooltip": "The algorithm used when sampling, this can affect the quality, speed, and style of the generated output."}),
-                    "positive": ("CONDITIONING", ),
+                "scheduler": (comfy.samplers.KSampler.SCHEDULERS, {"tooltip": "The scheduler controls how noise is gradually removed to form the image."}),
-                    "negative": ("CONDITIONING", ),
+                "positive": ("CONDITIONING", {"tooltip": "The conditioning describing the attributes you want to include in the image."}),
-                    "latent_image": ("LATENT", ),
+                "negative": ("CONDITIONING", {"tooltip": "The conditioning describing the attributes you want to exclude from the image."}),
-                    "denoise": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}),
+                "latent_image": ("LATENT", {"tooltip": "The latent image to denoise."}),
-                     }
+                "denoise": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01, "tooltip": "The amount of denoising applied, lower values will maintain the structure of the initial image allowing for image to image sampling."}),
-                }
+            }
+        }
 
     RETURN_TYPES = ("LATENT",)
+    OUTPUT_TOOLTIPS = ("The denoised latent.",)
     FUNCTION = "sample"
 
     CATEGORY = "sampling"
+    DESCRIPTION = "Uses the provided model, positive and negative conditioning to denoise the latent image."
 
     def sample(self, model, seed, steps, cfg, sampler_name, scheduler, positive, negative, latent_image, denoise=1.0):
         return common_ksampler(model, seed, steps, cfg, sampler_name, scheduler, positive, negative, latent_image, denoise=denoise)
@@ -1424,11 +1460,15 @@ class SaveImage:
 
     @classmethod
     def INPUT_TYPES(s):
-        return {"required": 
+        return {
-                    {"images": ("IMAGE", ),
+            "required": {
-                     "filename_prefix": ("STRING", {"default": "ComfyUI"})},
+                "images": ("IMAGE", {"tooltip": "The images to save."}),
-                "hidden": {"prompt": "PROMPT", "extra_pnginfo": "EXTRA_PNGINFO"},
+                "filename_prefix": ("STRING", {"default": "ComfyUI", "tooltip": "The prefix for the file to save. This may include formatting information such as %date:yyyy-MM-dd% or %Empty Latent Image.width% to include values from nodes."})
-                }
+            },
+            "hidden": {
+                "prompt": "PROMPT", "extra_pnginfo": "EXTRA_PNGINFO"
+            },
+        }
 
     RETURN_TYPES = ()
     FUNCTION = "save_images"
@@ -1436,6 +1476,7 @@ class SaveImage:
     OUTPUT_NODE = True
 
     CATEGORY = "image"
+    DESCRIPTION = "Saves the input images to your ComfyUI output directory."
 
     def save_images(self, images, filename_prefix="ComfyUI", prompt=None, extra_pnginfo=None):
         filename_prefix += self.prefix_append

server.py

@@ -12,7 +12,6 @@ import json
 import glob
 import struct
 import ssl
-import hashlib
 from PIL import Image, ImageOps
 from PIL.PngImagePlugin import PngInfo
 from io import BytesIO
@@ -28,7 +27,8 @@ import comfy.model_management
 import node_helpers
 from app.frontend_management import FrontendManager
 from app.user_manager import UserManager
-
+from model_filemanager import download_model, DownloadModelStatus
+from typing import Optional
 
 class BinaryEventTypes:
     PREVIEW_IMAGE = 1
@@ -76,6 +76,7 @@ class PromptServer():
         self.prompt_queue = None
         self.loop = loop
         self.messages = asyncio.Queue()
+        self.client_session:Optional[aiohttp.ClientSession] = None
         self.number = 0
 
         middlewares = [cache_control]
@@ -437,6 +438,9 @@ class PromptServer():
 
             if hasattr(obj_class, 'CATEGORY'):
                 info['category'] = obj_class.CATEGORY
+
+            if hasattr(obj_class, 'OUTPUT_TOOLTIPS'):
+                info['output_tooltips'] = obj_class.OUTPUT_TOOLTIPS
             return info
 
         @routes.get("/object_info")
@@ -560,6 +564,36 @@ class PromptServer():
 
             return web.Response(status=200)
         
+        # Internal route. Should not be depended upon and is subject to change at any time.
+        # TODO(robinhuang): Move to internal route table class once we refactor PromptServer to pass around Websocket.
+        @routes.post("/internal/models/download")
+        async def download_handler(request):
+            async def report_progress(filename: str, status: DownloadModelStatus):
+                await self.send_json("download_progress", status.to_dict())
+
+            data = await request.json()
+            url = data.get('url')
+            model_directory = data.get('model_directory')
+            model_filename = data.get('model_filename')
+            progress_interval = data.get('progress_interval', 1.0) # In seconds, how often to report download progress.
+
+            if not url or not model_directory or not model_filename:
+                return web.json_response({"status": "error", "message": "Missing URL or folder path or filename"}, status=400)
+
+            session = self.client_session
+            if session is None:
+                logging.error("Client session is not initialized")
+                return web.Response(status=500)
+            
+            task = asyncio.create_task(download_model(lambda url: session.get(url), model_filename, url, model_directory, report_progress, progress_interval))
+            await task
+
+            return web.json_response(task.result().to_dict())
+
+    async def setup(self):
+        timeout = aiohttp.ClientTimeout(total=None) # no timeout
+        self.client_session = aiohttp.ClientSession(timeout=timeout)
+
     def add_routes(self):
         self.user_manager.add_routes(self.routes)
 

tests-unit/prompt_server_test/download_models_test.py

@@ -0,0 +1,321 @@
+import pytest
+import aiohttp
+from aiohttp import ClientResponse
+import itertools
+import os 
+from unittest.mock import AsyncMock, patch, MagicMock
+from model_filemanager import download_model, validate_model_subdirectory, track_download_progress, create_model_path, check_file_exists, DownloadStatusType, DownloadModelStatus, validate_filename
+
+class AsyncIteratorMock:
+    """
+    A mock class that simulates an asynchronous iterator.
+    This is used to mimic the behavior of aiohttp's content iterator.
+    """
+    def __init__(self, seq):
+        # Convert the input sequence into an iterator
+        self.iter = iter(seq)
+
+    def __aiter__(self):
+        # This method is called when 'async for' is used
+        return self
+
+    async def __anext__(self):
+        # This method is called for each iteration in an 'async for' loop
+        try:
+            return next(self.iter)
+        except StopIteration:
+            # This is the asynchronous equivalent of StopIteration
+            raise StopAsyncIteration
+
+class ContentMock:
+    """
+    A mock class that simulates the content attribute of an aiohttp ClientResponse.
+    This class provides the iter_chunked method which returns an async iterator of chunks.
+    """
+    def __init__(self, chunks):
+        # Store the chunks that will be returned by the iterator
+        self.chunks = chunks
+
+    def iter_chunked(self, chunk_size):
+        # This method mimics aiohttp's content.iter_chunked()
+        # For simplicity in testing, we ignore chunk_size and just return our predefined chunks
+        return AsyncIteratorMock(self.chunks)
+
+@pytest.mark.asyncio
+async def test_download_model_success():
+    mock_response = AsyncMock(spec=aiohttp.ClientResponse)
+    mock_response.status = 200
+    mock_response.headers = {'Content-Length': '1000'}
+    # Create a mock for content that returns an async iterator directly
+    chunks = [b'a' * 500, b'b' * 300, b'c' * 200]
+    mock_response.content = ContentMock(chunks)
+
+    mock_make_request = AsyncMock(return_value=mock_response)
+    mock_progress_callback = AsyncMock()
+
+    # Mock file operations
+    mock_open = MagicMock()
+    mock_file = MagicMock()
+    mock_open.return_value.__enter__.return_value = mock_file
+    time_values = itertools.count(0, 0.1)
+
+    with patch('model_filemanager.create_model_path', return_value=('models/checkpoints/model.sft', 'checkpoints/model.sft')), \
+         patch('model_filemanager.check_file_exists', return_value=None), \
+         patch('builtins.open', mock_open), \
+         patch('time.time', side_effect=time_values):  # Simulate time passing
+
+        result = await download_model(
+            mock_make_request,
+            'model.sft',
+            'http://example.com/model.sft',
+            'checkpoints',
+            mock_progress_callback
+        )
+
+    # Assert the result
+    assert isinstance(result, DownloadModelStatus)
+    assert result.message == 'Successfully downloaded model.sft'
+    assert result.status == 'completed'
+    assert result.already_existed is False
+
+    # Check progress callback calls
+    assert mock_progress_callback.call_count >= 3  # At least start, one progress update, and completion
+    
+    # Check initial call
+    mock_progress_callback.assert_any_call(
+        'checkpoints/model.sft',
+        DownloadModelStatus(DownloadStatusType.PENDING, 0, "Starting download of model.sft", False)
+    )
+
+    # Check final call
+    mock_progress_callback.assert_any_call(
+        'checkpoints/model.sft',
+        DownloadModelStatus(DownloadStatusType.COMPLETED, 100, "Successfully downloaded model.sft", False)
+    )
+
+    # Verify file writing
+    mock_file.write.assert_any_call(b'a' * 500)
+    mock_file.write.assert_any_call(b'b' * 300)
+    mock_file.write.assert_any_call(b'c' * 200)
+
+    # Verify request was made
+    mock_make_request.assert_called_once_with('http://example.com/model.sft')
+
+@pytest.mark.asyncio
+async def test_download_model_url_request_failure():
+    # Mock dependencies
+    mock_response = AsyncMock(spec=ClientResponse)
+    mock_response.status = 404  # Simulate a "Not Found" error
+    mock_get = AsyncMock(return_value=mock_response)
+    mock_progress_callback = AsyncMock()
+
+    # Mock the create_model_path function
+    with patch('model_filemanager.create_model_path', return_value=('/mock/path/model.safetensors', 'mock/path/model.safetensors')):
+        # Mock the check_file_exists function to return None (file doesn't exist)
+        with patch('model_filemanager.check_file_exists', return_value=None):
+            # Call the function
+            result = await download_model(
+                mock_get,
+                'model.safetensors',
+                'http://example.com/model.safetensors',
+                'mock_directory',
+                mock_progress_callback
+            )
+
+    # Assert the expected behavior
+    assert isinstance(result, DownloadModelStatus)
+    assert result.status == 'error'
+    assert result.message == 'Failed to download model.safetensors. Status code: 404'
+    assert result.already_existed is False
+
+    # Check that progress_callback was called with the correct arguments
+    mock_progress_callback.assert_any_call(
+        'mock_directory/model.safetensors',
+        DownloadModelStatus(
+            status=DownloadStatusType.PENDING,
+            progress_percentage=0,
+            message='Starting download of model.safetensors',
+            already_existed=False
+        )
+    )
+    mock_progress_callback.assert_called_with(
+        'mock_directory/model.safetensors',
+        DownloadModelStatus(
+            status=DownloadStatusType.ERROR,
+            progress_percentage=0,
+            message='Failed to download model.safetensors. Status code: 404',
+            already_existed=False
+        )
+    )
+
+    # Verify that the get method was called with the correct URL
+    mock_get.assert_called_once_with('http://example.com/model.safetensors')
+
+@pytest.mark.asyncio
+async def test_download_model_invalid_model_subdirectory():
+    
+    mock_make_request = AsyncMock()
+    mock_progress_callback = AsyncMock()
+
+    
+    result = await download_model(
+        mock_make_request,
+        'model.sft',
+        'http://example.com/model.sft',
+        '../bad_path',
+        mock_progress_callback
+    )
+
+    # Assert the result
+    assert isinstance(result, DownloadModelStatus)
+    assert result.message == 'Invalid model subdirectory'
+    assert result.status == 'error'
+    assert result.already_existed is False
+
+
+# For create_model_path function
+def test_create_model_path(tmp_path, monkeypatch):
+    mock_models_dir = tmp_path / "models"
+    monkeypatch.setattr('folder_paths.models_dir', str(mock_models_dir))
+    
+    model_name = "test_model.sft"
+    model_directory = "test_dir"
+    
+    file_path, relative_path = create_model_path(model_name, model_directory, mock_models_dir)
+    
+    assert file_path == str(mock_models_dir / model_directory / model_name)
+    assert relative_path == f"{model_directory}/{model_name}"
+    assert os.path.exists(os.path.dirname(file_path))
+
+
+@pytest.mark.asyncio
+async def test_check_file_exists_when_file_exists(tmp_path):
+    file_path = tmp_path / "existing_model.sft"
+    file_path.touch()  # Create an empty file
+    
+    mock_callback = AsyncMock()
+    
+    result = await check_file_exists(str(file_path), "existing_model.sft", mock_callback, "test/existing_model.sft")
+    
+    assert result is not None
+    assert result.status == "completed"
+    assert result.message == "existing_model.sft already exists"
+    assert result.already_existed is True
+    
+    mock_callback.assert_called_once_with(
+        "test/existing_model.sft",
+        DownloadModelStatus(DownloadStatusType.COMPLETED, 100, "existing_model.sft already exists", already_existed=True)
+    )
+
+@pytest.mark.asyncio
+async def test_check_file_exists_when_file_does_not_exist(tmp_path):
+    file_path = tmp_path / "non_existing_model.sft"
+    
+    mock_callback = AsyncMock()
+    
+    result = await check_file_exists(str(file_path), "non_existing_model.sft", mock_callback, "test/non_existing_model.sft")
+    
+    assert result is None
+    mock_callback.assert_not_called()
+
+@pytest.mark.asyncio
+async def test_track_download_progress_no_content_length():
+    mock_response = AsyncMock(spec=aiohttp.ClientResponse)
+    mock_response.headers = {}  # No Content-Length header
+    mock_response.content.iter_chunked.return_value = AsyncIteratorMock([b'a' * 500, b'b' * 500])
+
+    mock_callback = AsyncMock()
+    mock_open = MagicMock(return_value=MagicMock())
+
+    with patch('builtins.open', mock_open):
+        result = await track_download_progress(
+            mock_response, '/mock/path/model.sft', 'model.sft',
+            mock_callback, 'models/model.sft', interval=0.1
+        )
+
+    assert result.status == "completed"
+    # Check that progress was reported even without knowing the total size
+    mock_callback.assert_any_call(
+        'models/model.sft',
+        DownloadModelStatus(DownloadStatusType.IN_PROGRESS, 0, "Downloading model.sft", already_existed=False)
+    )
+
+@pytest.mark.asyncio
+async def test_track_download_progress_interval():
+    mock_response = AsyncMock(spec=aiohttp.ClientResponse)
+    mock_response.headers = {'Content-Length': '1000'}
+    mock_response.content.iter_chunked.return_value = AsyncIteratorMock([b'a' * 100] * 10)
+
+    mock_callback = AsyncMock()
+    mock_open = MagicMock(return_value=MagicMock())
+
+    # Create a mock time function that returns incremental float values
+    mock_time = MagicMock()
+    mock_time.side_effect = [i * 0.5 for i in range(30)]  # This should be enough for 10 chunks
+
+    with patch('builtins.open', mock_open), \
+         patch('time.time', mock_time):
+        await track_download_progress(
+            mock_response, '/mock/path/model.sft', 'model.sft',
+            mock_callback, 'models/model.sft', interval=1.0
+        )
+
+    # Print out the actual call count and the arguments of each call for debugging
+    print(f"mock_callback was called {mock_callback.call_count} times")
+    for i, call in enumerate(mock_callback.call_args_list):
+        args, kwargs = call
+        print(f"Call {i + 1}: {args[1].status}, Progress: {args[1].progress_percentage:.2f}%")
+
+    # Assert that progress was updated at least 3 times (start, at least one interval, and end)
+    assert mock_callback.call_count >= 3, f"Expected at least 3 calls, but got {mock_callback.call_count}"
+
+    # Verify the first and last calls
+    first_call = mock_callback.call_args_list[0]
+    assert first_call[0][1].status == "in_progress"
+    # Allow for some initial progress, but it should be less than 50%
+    assert 0 <= first_call[0][1].progress_percentage < 50, f"First call progress was {first_call[0][1].progress_percentage}%"
+
+    last_call = mock_callback.call_args_list[-1]
+    assert last_call[0][1].status == "completed"
+    assert last_call[0][1].progress_percentage == 100
+
+def test_valid_subdirectory():
+    assert validate_model_subdirectory("valid-model123") is True
+
+def test_subdirectory_too_long():
+    assert validate_model_subdirectory("a" * 51) is False
+
+def test_subdirectory_with_double_dots():
+    assert validate_model_subdirectory("model/../unsafe") is False
+
+def test_subdirectory_with_slash():
+    assert validate_model_subdirectory("model/unsafe") is False
+
+def test_subdirectory_with_special_characters():
+    assert validate_model_subdirectory("model@unsafe") is False
+
+def test_subdirectory_with_underscore_and_dash():
+    assert validate_model_subdirectory("valid_model-name") is True
+
+def test_empty_subdirectory():
+    assert validate_model_subdirectory("") is False
+
+@pytest.mark.parametrize("filename, expected", [
+    ("valid_model.safetensors", True),
+    ("valid_model.sft", True),
+    ("valid model.safetensors", True), # Test with space
+    ("UPPERCASE_MODEL.SAFETENSORS", True),
+    ("model_with.multiple.dots.pt", False),
+    ("", False),  # Empty string
+    ("../../../etc/passwd", False),  # Path traversal attempt
+    ("/etc/passwd", False),  # Absolute path
+    ("\\windows\\system32\\config\\sam", False),  # Windows path
+    (".hidden_file.pt", False),  # Hidden file
+    ("invalid<char>.ckpt", False),  # Invalid character
+    ("invalid?.ckpt", False),  # Another invalid character
+    ("very" * 100 + ".safetensors", False),  # Too long filename
+    ("\nmodel_with_newline.pt", False),  # Newline character
+    ("model_with_emoji😊.pt", False),  # Emoji in filename
+])
+def test_validate_filename(filename, expected):
+    assert validate_filename(filename) == expected

tests-unit/requirements.txt

@@ -1 +1,3 @@
 pytest>=7.8.0
+pytest-aiohttp
+pytest-asyncio

web/extensions/core/tooltips.js

@@ -0,0 +1,122 @@
+import { app } from "../../scripts/app.js";
+import { $el } from "../../scripts/ui.js";
+
+// Adds support for tooltips
+
+function getHoveredWidget() {
+	if (!app) {
+		return;
+	}
+
+	const node = app.canvas.node_over;
+	if (!node.widgets) return;
+
+	const graphPos = app.canvas.graph_mouse;
+
+	const x = graphPos[0] - node.pos[0];
+	const y = graphPos[1] - node.pos[1];
+
+	for (const w of node.widgets) {
+		let widgetWidth, widgetHeight;
+		if (w.computeSize) {
+			const sz = w.computeSize();
+			widgetWidth = sz[0];
+			widgetHeight = sz[1];
+		} else {
+			widgetWidth = w.width || node.size[0];
+			widgetHeight = LiteGraph.NODE_WIDGET_HEIGHT;
+		}
+
+		if (w.last_y !== undefined && x >= 6 && x <= widgetWidth - 12 && y >= w.last_y && y <= w.last_y + widgetHeight) {
+			return w;
+		}
+	}
+}
+
+app.registerExtension({
+	name: "Comfy.Tooltips",
+	setup() {
+		const tooltipEl = $el("div.comfy-graph-tooltip", {
+			parent: document.body,
+		});
+		let idleTimeout;
+
+		const hideTooltip = () => {
+			tooltipEl.style.display = "none";
+		};
+		const showTooltip = (tooltip) => {
+			if (!tooltip) return;
+
+			tooltipEl.textContent = tooltip;
+			tooltipEl.style.display = "block";
+			tooltipEl.style.left = app.canvas.mouse[0] + "px";
+			tooltipEl.style.top = app.canvas.mouse[1] + "px";
+			const rect = tooltipEl.getBoundingClientRect();
+			if (rect.right > window.innerWidth) {
+				tooltipEl.style.left = app.canvas.mouse[0] - rect.width + "px";
+			}
+
+			if (rect.top < 0) {
+				tooltipEl.style.top = app.canvas.mouse[1] + rect.height + "px";
+			}
+		};
+		const getInputTooltip = (nodeData, name) => {
+			const inputDef = nodeData.input?.required?.[name] ?? nodeData.input?.optional?.[name];
+			return inputDef?.[1]?.tooltip;
+		};
+		const onIdle = () => {
+			const { canvas } = app;
+			const node = canvas.node_over;
+			if (!node) return;
+
+			const nodeData = node.constructor.nodeData ?? {};
+
+			if (node.constructor.title_mode !== LiteGraph.NO_TITLE && canvas.graph_mouse[1] < node.pos[1]) {
+				return showTooltip(nodeData.description);
+			}
+
+			if (node.flags?.collapsed) return;
+
+			const inputSlot = canvas.isOverNodeInput(node, canvas.graph_mouse[0], canvas.graph_mouse[1], [0, 0]);
+			if (inputSlot !== -1) {
+				const inputName = node.inputs[inputSlot].name;
+				return showTooltip(getInputTooltip(nodeData, inputName));
+			}
+
+			const outputSlot = canvas.isOverNodeOutput(node, canvas.graph_mouse[0], canvas.graph_mouse[1], [0, 0]);
+			if (outputSlot !== -1) {
+				return showTooltip(nodeData.output_tooltips?.[outputSlot]);
+			}
+
+			const widget = getHoveredWidget();
+			// Dont show for DOM widgets, these use native browser tooltips as we dont get proper mouse events on these
+			if (widget && !widget.element) {
+				return showTooltip(widget.tooltip ?? getInputTooltip(nodeData, widget.name));
+			}
+		};
+
+		const onMouseMove = (e) => {
+			hideTooltip();
+			clearTimeout(idleTimeout);
+			
+			if(e.target.nodeName !== "CANVAS") return
+			idleTimeout = setTimeout(onIdle, 500);
+		};
+
+		app.ui.settings.addSetting({
+			id: "Comfy.EnableTooltips",
+			name: "Enable Tooltips",
+			type: "boolean",
+			defaultValue: true,
+			onChange(value) {
+				if (value) {
+					window.addEventListener("mousemove", onMouseMove);
+					window.addEventListener("click", hideTooltip);
+				} else {
+					window.removeEventListener("mousemove", onMouseMove);
+					window.removeEventListener("click", hideTooltip);
+				}
+			},
+		});
+	},
+});

web/extensions/core/widgetInputs.js

@@ -181,7 +181,7 @@ export function mergeIfValid(output, config2, forceUpdate, recreateWidget, confi
 
 	const isNumber = config1[0] === "INT" || config1[0] === "FLOAT";
 	for (const k of keys.values()) {
-		if (k !== "default" && k !== "forceInput" && k !== "defaultInput" && k !== "control_after_generate" && k !== "multiline") {
+		if (k !== "default" && k !== "forceInput" && k !== "defaultInput" && k !== "control_after_generate" && k !== "multiline" && k !== "tooltip") {
 			let v1 = config1[1][k];
 			let v2 = config2[1]?.[k];
 

web/scripts/app.js

@@ -1713,9 +1713,10 @@ export class ComfyApp {
 				for (const o in nodeData["output"]) {
 					let output = nodeData["output"][o];
 					if(output instanceof Array) output = "COMBO";
+					const outputTooltip = nodeData["output_tooltips"]?.[o];
 					const outputName = nodeData["output_name"][o] || output;
 					const outputShape = nodeData["output_is_list"][o] ? LiteGraph.GRID_SHAPE : LiteGraph.CIRCLE_SHAPE ;
-					this.addOutput(outputName, output, { shape: outputShape });
+					this.addOutput(outputName, output, { shape: outputShape, tooltip: outputTooltip });
 				}
 
 				const s = this.computeSize();

web/scripts/domWidget.js

@@ -223,6 +223,12 @@ LGraphNode.prototype.addDOMWidget = function (name, type, element, options) {
 		document.addEventListener("mousedown", mouseDownHandler);
 	}
 
+	const { nodeData } = this.constructor;
+	const tooltip = (nodeData?.input.required?.[name] ?? nodeData?.input.optional?.[name])?.[1]?.tooltip;
+	if (tooltip && !element.title) {
+		element.title = tooltip;
+	}
+
 	const widget = {
 		type,
 		name,

web/scripts/widgets.js

@@ -75,6 +75,7 @@ export function addValueControlWidgets(node, targetWidget, defaultValue = "rando
 			serialize: false, // Don't include this in prompt.
 		}
 	);
+	valueControl.tooltip = "Allows the linked widget to be changed automatically, for example randomizing the noise seed.";
 	valueControl[IS_CONTROL_WIDGET] = true;
 	updateControlWidgetLabel(valueControl);
 	widgets.push(valueControl);
@@ -95,6 +96,7 @@ export function addValueControlWidgets(node, targetWidget, defaultValue = "rando
 			}
 		);
 		updateControlWidgetLabel(comboFilter);
+		comboFilter.tooltip = "Allows for filtering the list of values when changing the value via the control generate mode. Allows for RegEx matches in the format /abc/ to only filter to values containing 'abc'."
 
 		widgets.push(comboFilter);
 	}

web/style.css

@@ -645,3 +645,20 @@ dialog::backdrop {
 audio.comfy-audio.empty-audio-widget {
 	display: none;
 }
+
+.comfy-graph-tooltip {
+	background: var(--comfy-input-bg);
+	border-radius: 5px;
+	box-shadow: 0 0 5px rgba(0, 0, 0, 0.4);
+	color: var(--input-text);
+	display: none;
+	font-family: sans-serif;
+	left: 0;
+	max-width: 30vw;
+	padding: 4px 8px;
+	position: absolute;
+	top: 0;
+	transform: translate(5px, calc(-100% - 5px));
+	white-space: pre-wrap;
+	z-index: 99999;
+}