diff --git a/comfy/ldm/modules/attention.py b/comfy/ldm/modules/attention.py index 4750a3a9a..853e4ec32 100644 --- a/comfy/ldm/modules/attention.py +++ b/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( diff --git a/main.py b/main.py index 479643b84..b878b3e01 100644 --- a/main.py +++ b/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") diff --git a/model_filemanager/__init__.py b/model_filemanager/__init__.py new file mode 100644 index 000000000..e318351c0 --- /dev/null +++ b/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 diff --git a/model_filemanager/download_models.py b/model_filemanager/download_models.py new file mode 100644 index 000000000..712d59328 --- /dev/null +++ b/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 diff --git a/nodes.py b/nodes.py index 525b28d85..16f5c9b00 100644 --- a/nodes.py +++ b/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",), - "clip": ("CLIP", ), - "lora_name": (folder_paths.get_filename_list("loras"), ), - "strength_model": ("FLOAT", {"default": 1.0, "min": -100.0, "max": 100.0, "step": 0.01}), - "strength_clip": ("FLOAT", {"default": 1.0, "min": -100.0, "max": 100.0, "step": 0.01}), - }} + return { + "required": { + "model": ("MODEL", {"tooltip": "The diffusion model the LoRA will be applied to."}), + "clip": ("CLIP", {"tooltip": "The CLIP model the LoRA will be applied to."}), + "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}), - "height": ("INT", {"default": 512, "min": 16, "max": MAX_RESOLUTION, "step": 8}), - "batch_size": ("INT", {"default": 1, "min": 1, "max": 4096})}} + return { + "required": { + "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": - {"model": ("MODEL",), - "seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}), - "steps": ("INT", {"default": 20, "min": 1, "max": 10000}), - "cfg": ("FLOAT", {"default": 8.0, "min": 0.0, "max": 100.0, "step":0.1, "round": 0.01}), - "sampler_name": (comfy.samplers.KSampler.SAMPLERS, ), - "scheduler": (comfy.samplers.KSampler.SCHEDULERS, ), - "positive": ("CONDITIONING", ), - "negative": ("CONDITIONING", ), - "latent_image": ("LATENT", ), - "denoise": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}), - } - } + return { + "required": { + "model": ("MODEL", {"tooltip": "The model used for denoising the input latent."}), + "seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff, "tooltip": "The random seed used for creating the noise."}), + "steps": ("INT", {"default": 20, "min": 1, "max": 10000, "tooltip": "The number of steps used in the denoising process."}), + "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."}), + "sampler_name": (comfy.samplers.KSampler.SAMPLERS, {"tooltip": "The algorithm used when sampling, this can affect the quality, speed, and style of the generated output."}), + "scheduler": (comfy.samplers.KSampler.SCHEDULERS, {"tooltip": "The scheduler controls how noise is gradually removed to form the image."}), + "positive": ("CONDITIONING", {"tooltip": "The conditioning describing the attributes you want to include in the image."}), + "negative": ("CONDITIONING", {"tooltip": "The conditioning describing the attributes you want to exclude from the image."}), + "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": - {"images": ("IMAGE", ), - "filename_prefix": ("STRING", {"default": "ComfyUI"})}, - "hidden": {"prompt": "PROMPT", "extra_pnginfo": "EXTRA_PNGINFO"}, - } + return { + "required": { + "images": ("IMAGE", {"tooltip": "The images to save."}), + "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 diff --git a/server.py b/server.py index 1c07f9781..9d8269f22 100644 --- a/server.py +++ b/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") @@ -559,6 +563,36 @@ class PromptServer(): self.prompt_queue.delete_history_item(id_to_delete) 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) diff --git a/tests-unit/prompt_server_test/__init__.py b/tests-unit/prompt_server_test/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests-unit/prompt_server_test/download_models_test.py b/tests-unit/prompt_server_test/download_models_test.py new file mode 100644 index 000000000..66150a468 --- /dev/null +++ b/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.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 diff --git a/tests-unit/requirements.txt b/tests-unit/requirements.txt index 0587502f8..d70d00f4b 100644 --- a/tests-unit/requirements.txt +++ b/tests-unit/requirements.txt @@ -1 +1,3 @@ pytest>=7.8.0 +pytest-aiohttp +pytest-asyncio diff --git a/web/extensions/core/tooltips.js b/web/extensions/core/tooltips.js new file mode 100644 index 000000000..08ab05efa --- /dev/null +++ b/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); + } + }, + }); + }, +}); diff --git a/web/extensions/core/widgetInputs.js b/web/extensions/core/widgetInputs.js index f1a1d22cd..081554968 100644 --- a/web/extensions/core/widgetInputs.js +++ b/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]; diff --git a/web/scripts/app.js b/web/scripts/app.js index 8b4478a32..df9245097 100644 --- a/web/scripts/app.js +++ b/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(); diff --git a/web/scripts/domWidget.js b/web/scripts/domWidget.js index d97122f9d..03ed2ff10 100644 --- a/web/scripts/domWidget.js +++ b/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, diff --git a/web/scripts/widgets.js b/web/scripts/widgets.js index 6a6899705..59b17fd13 100644 --- a/web/scripts/widgets.js +++ b/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); } diff --git a/web/style.css b/web/style.css index 8ef1d0dd1..7774bfede 100644 --- a/web/style.css +++ b/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; +}