Merge 15d49a61b8 into b08debceca

Signed-off-by: bigcat88 <bigcat88@icloud.com>

.coderabbit.yaml

@@ -4,12 +4,12 @@ early_access: false
 tone_instructions: "Only comment on issues introduced by this PR's changes. Do not flag pre-existing problems in moved, re-indented, or reformatted code."
 
 reviews:
-  profile: "chill"
-  request_changes_workflow: false
+  profile: "assertive"
+  request_changes_workflow: true
   high_level_summary: false
   poem: false
   review_status: false
-  review_details: false
+  review_details: true
   commit_status: true
   collapse_walkthrough: true
   changed_files_summary: false
@@ -39,6 +39,14 @@ reviews:
     - path: "**"
       instructions: |
         IMPORTANT: Only comment on issues directly introduced by this PR's code changes.
+        Treat AGENTS.md as mandatory repository policy, not optional style guidance.
+        Flag PR changes that violate AGENTS.md even when the code is otherwise functional.
+        In particular, enforce architecture boundaries, dtype/device/memory rules,
+        interface contracts, import style, no unnecessary try/except blocks, no inline
+        imports, no outbound internet paths in core ComfyUI, and narrow scoped fixes.
+        Prefer direct findings over suggestions when a rule is violated. Only ignore
+        AGENTS.md when it clearly conflicts with a newer explicit maintainer instruction
+        in the PR.
         Do NOT flag pre-existing issues in code that was merely moved, re-indented,
         de-indented, or reformatted without logic changes. If code appears in the diff
         only due to whitespace or structural reformatting (e.g., removing a `with:` block),
@@ -123,5 +131,10 @@ chat:
 
 knowledge_base:
   opt_out: false
+  code_guidelines:
+    enabled: true
+    filePatterns:
+      - files: "AGENTS.md"
+        applyTo: "**"
   learnings:
     scope: "auto"

api_server/routes/internal/internal_routes.py

@@ -2,7 +2,9 @@ from aiohttp import web
 from typing import Optional
 from folder_paths import folder_names_and_paths, get_directory_by_type
 from api_server.services.terminal_service import TerminalService
+from app.model_download import ModelDownloadError, download_model_to_destination, parse_model_download_request, resolve_model_download_destination
 import app.logger
+import logging
 import os
 
 class InternalRoutes:
@@ -51,6 +53,32 @@ class InternalRoutes:
                 response[key] = folder_names_and_paths[key][0]
             return web.json_response(response)
 
+        @self.routes.post('/models/download')
+        async def download_model(request):
+            try:
+                try:
+                    json_data = await request.json()
+                except Exception as err:
+                    raise ModelDownloadError("Expected a JSON request body.") from err
+
+                download_request = parse_model_download_request(json_data)
+                destination = resolve_model_download_destination(download_request)
+                if self.prompt_server.client_session is None:
+                    raise ModelDownloadError("HTTP client session is not ready.", status=503)
+                result = await download_model_to_destination(
+                    self.prompt_server.client_session,
+                    download_request,
+                    destination,
+                )
+            except ModelDownloadError as err:
+                return web.json_response({"error": str(err)}, status=err.status)
+            except Exception as err:
+                logging.exception("Failed to download model: %s", err)
+                return web.json_response({"error": "Failed to download model."}, status=500)
+
+            response_status = 200 if result["status"] == "already_exists" else 201
+            return web.json_response(result, status=response_status)
+
         @self.routes.get('/files/{directory_type}')
         async def get_files(request: web.Request) -> web.Response:
             directory_type = request.match_info['directory_type']

app/model_download.py

@@ -0,0 +1,281 @@
+from __future__ import annotations
+
+import os
+import posixpath
+from dataclasses import dataclass
+from pathlib import PurePosixPath
+from urllib.parse import urljoin, urlparse
+
+from aiohttp import ClientSession, ClientTimeout
+
+import folder_paths
+
+
+ALLOWED_DOWNLOAD_HOSTS = {"huggingface.co", "civitai.com", "civitai.red"}
+ALLOWED_DOWNLOAD_SUFFIXES = (".safetensors", ".sft", ".ckpt", ".pth", ".pt")
+BLOCKED_MODEL_FOLDERS = {"configs", "custom_nodes"}
+CHUNK_SIZE = 1024 * 1024
+
+# Bound the network call so a hung remote eventually surfaces an error
+# instead of blocking the request handler forever. ``sock_read`` is the
+# inter-chunk read timeout, which is the right knob for long downloads:
+# a slow-but-progressing transfer keeps making progress, while a stalled
+# socket fails predictably.
+DOWNLOAD_TIMEOUT = ClientTimeout(total=None, connect=30, sock_connect=30, sock_read=300)
+
+# Maximum number of redirects we follow manually. Hugging Face typically
+# redirects ``/resolve/main/...`` to a single CDN URL, so a small budget
+# is enough while still preventing redirect loops.
+MAX_DOWNLOAD_REDIRECTS = 5
+
+WHITE_LISTED_DOWNLOAD_URLS = {
+    "https://huggingface.co/stabilityai/stable-zero123/resolve/main/stable_zero123.ckpt",
+    "https://huggingface.co/TencentARC/T2I-Adapter/resolve/main/models/t2iadapter_depth_sd14v1.pth?download=true",
+    "https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.0/RealESRGAN_x4plus.pth",
+}
+
+
+@dataclass(frozen=True)
+class ModelDownloadRequest:
+    name: str
+    url: str
+    directory: str
+
+
+@dataclass(frozen=True)
+class ModelDownloadDestination:
+    directory: str
+    relative_path: str
+    full_path: str
+    already_exists: bool = False
+
+
+class ModelDownloadError(Exception):
+    def __init__(self, message: str, status: int = 400):
+        super().__init__(message)
+        self.status = status
+
+
+def parse_model_download_request(data) -> ModelDownloadRequest:
+    if not isinstance(data, dict):
+        raise ModelDownloadError("Expected a JSON object.")
+
+    name = data.get("name")
+    url = data.get("url")
+    directory = data.get("directory")
+    if not isinstance(name, str) or not isinstance(url, str) or not isinstance(directory, str):
+        raise ModelDownloadError("Missing model name, URL, or directory.")
+
+    name = name.strip()
+    url = url.strip()
+    directory = directory.strip()
+    if not name or not url or not directory:
+        raise ModelDownloadError("Model name, URL, and directory are required.")
+
+    if not is_allowed_model_download_url(url):
+        raise ModelDownloadError("Model download URL is not allowed.")
+
+    relative_path = normalize_model_relative_path(name)
+    if not relative_path.lower().endswith(ALLOWED_DOWNLOAD_SUFFIXES):
+        raise ModelDownloadError("Model filename extension is not allowed.")
+
+    return ModelDownloadRequest(name=relative_path, url=url, directory=directory)
+
+
+def is_allowed_model_download_url(url: str) -> bool:
+    """Return True for URLs we are willing to fetch on behalf of the user.
+
+    The same predicate is applied to the user-supplied URL and to every
+    redirect target, so SSRF via redirects on an allowed host is contained
+    to the same allowlist. Subdomains of allowlisted hosts are accepted
+    because Hugging Face and Civitai both serve actual file payloads from
+    CDN subdomains (e.g. ``cdn-lfs.huggingface.co``).
+    """
+    if url in WHITE_LISTED_DOWNLOAD_URLS:
+        return True
+
+    try:
+        parsed = urlparse(url)
+    except ValueError:
+        return False
+
+    if parsed.scheme != "https":
+        return False
+
+    host = (parsed.hostname or "").lower()
+    if not host:
+        return False
+
+    for allowed in ALLOWED_DOWNLOAD_HOSTS:
+        if host == allowed or host.endswith("." + allowed):
+            return True
+    return False
+
+
+def normalize_model_relative_path(name: str) -> str:
+    if "\x00" in name:
+        raise ModelDownloadError("Model filename is invalid.")
+
+    candidate = name.replace("\\", "/")
+    if candidate.startswith("/"):
+        raise ModelDownloadError("Model filename must be relative.")
+
+    parts = PurePosixPath(candidate).parts
+    if not parts or any(part in ("", ".", "..") for part in parts):
+        raise ModelDownloadError("Model filename must stay inside the model folder.")
+
+    normalized = posixpath.normpath(candidate)
+    if normalized in ("", ".") or normalized.startswith("../"):
+        raise ModelDownloadError("Model filename must stay inside the model folder.")
+
+    return normalized
+
+
+def resolve_model_download_destination(request: ModelDownloadRequest) -> ModelDownloadDestination:
+    directory = folder_paths.map_legacy(request.directory)
+    if directory in BLOCKED_MODEL_FOLDERS or directory not in folder_paths.folder_names_and_paths:
+        raise ModelDownloadError("Model directory is not allowed.", status=404)
+
+    existing_path = folder_paths.get_full_path(directory, request.name)
+    if existing_path is not None:
+        return ModelDownloadDestination(
+            directory=directory,
+            relative_path=request.name,
+            full_path=existing_path,
+            already_exists=True,
+        )
+
+    destination_root = find_writable_model_root(directory)
+    full_path = safe_join(destination_root, request.name)
+    return ModelDownloadDestination(
+        directory=directory,
+        relative_path=request.name,
+        full_path=full_path,
+    )
+
+
+def find_writable_model_root(directory: str) -> str:
+    for root in folder_paths.get_folder_paths(directory):
+        try:
+            os.makedirs(root, exist_ok=True)
+        except OSError:
+            continue
+
+        if is_writable_directory(root):
+            return os.path.abspath(root)
+
+    raise ModelDownloadError("No writable model folder is configured.", status=403)
+
+
+def is_writable_directory(path: str) -> bool:
+    probe = os.path.join(path, ".comfy-download-write-test")
+    try:
+        with open(probe, "xb"):
+            pass
+        os.remove(probe)
+        return True
+    except OSError:
+        try:
+            if os.path.exists(probe):
+                os.remove(probe)
+        except OSError:
+            pass
+        return False
+
+
+def safe_join(root: str, relative_path: str) -> str:
+    root = os.path.abspath(root)
+    full_path = os.path.abspath(os.path.join(root, relative_path))
+    if os.path.commonpath((root, full_path)) != root:
+        raise ModelDownloadError("Model filename must stay inside the model folder.")
+    return full_path
+
+
+async def open_model_download_response(session: ClientSession, url: str):
+    """GET ``url`` with explicit timeout and an allowlist-checked redirect chain.
+
+    aiohttp follows redirects by default, which would let an allowed host
+    redirect to an arbitrary internal target (SSRF). We disable automatic
+    following and validate every ``Location`` against the same allowlist
+    used for the initial URL.
+    """
+    current_url = url
+    for _ in range(MAX_DOWNLOAD_REDIRECTS + 1):
+        response = await session.get(
+            current_url,
+            allow_redirects=False,
+            timeout=DOWNLOAD_TIMEOUT,
+        )
+        if response.status not in (301, 302, 303, 307, 308):
+            return response
+
+        location = response.headers.get("Location", "").strip()
+        response.release()
+        if not location:
+            raise ModelDownloadError("Redirect response missing Location header.", status=502)
+        next_url = urljoin(current_url, location)
+        if not is_allowed_model_download_url(next_url):
+            raise ModelDownloadError("Model download redirect target is not allowed.", status=403)
+        current_url = next_url
+
+    raise ModelDownloadError("Too many redirects while downloading model.", status=502)
+
+
+async def download_model_to_destination(
+    session: ClientSession,
+    request: ModelDownloadRequest,
+    destination: ModelDownloadDestination,
+) -> dict:
+    if destination.already_exists:
+        return download_response("already_exists", request, destination)
+
+    os.makedirs(os.path.dirname(destination.full_path), exist_ok=True)
+    partial_path = f"{destination.full_path}.part"
+
+    try:
+        fd = os.open(partial_path, os.O_WRONLY | os.O_CREAT | os.O_EXCL, 0o644)
+    except FileExistsError as err:
+        raise ModelDownloadError("A download for this model is already in progress.", status=409) from err
+
+    bytes_written = 0
+    try:
+        with os.fdopen(fd, "wb") as output:
+            async with await open_model_download_response(session, request.url) as response:
+                if response.status >= 400:
+                    raise ModelDownloadError(f"Model download failed with HTTP {response.status}.", status=502)
+
+                expected_size = response.content_length
+                async for chunk in response.content.iter_chunked(CHUNK_SIZE):
+                    output.write(chunk)
+                    bytes_written += len(chunk)
+
+                if expected_size is not None and bytes_written != expected_size:
+                    raise ModelDownloadError("Model download ended before all bytes were received.", status=502)
+
+        os.replace(partial_path, destination.full_path)
+    except Exception:
+        try:
+            if os.path.exists(partial_path):
+                os.remove(partial_path)
+        except OSError:
+            pass
+        raise
+
+    return download_response("downloaded", request, destination, bytes_written)
+
+
+def download_response(
+    status: str,
+    request: ModelDownloadRequest,
+    destination: ModelDownloadDestination,
+    size: int | None = None,
+) -> dict:
+    response = {
+        "status": status,
+        "name": request.name,
+        "directory": destination.directory,
+        "path": destination.full_path,
+    }
+    if size is not None:
+        response["size"] = size
+    return response

comfy/sd1_clip.py

@@ -543,18 +543,24 @@ class SDTokenizer:
     def _try_get_embedding(self, embedding_name:str):
         '''
         Takes a potential embedding name and tries to retrieve it.
-        Returns a Tuple consisting of the embedding and any leftover string, embedding can be None.
+        Returns a Tuple consisting of the embedding, the cleaned embedding name, and any leftover string, embedding can be None.
         '''
         split_embed = embedding_name.split()
         embedding_name = split_embed[0]
         leftover = ' '.join(split_embed[1:])
+
+        match = re.search(r'[<\[]', embedding_name)
+        if match is not None:
+            leftover = embedding_name[match.start():] + (" " + leftover if leftover else "")
+            embedding_name = embedding_name[:match.start()]
+
         embed = load_embed(embedding_name, self.embedding_directory, self.embedding_size, self.embedding_key)
         if embed is None:
             stripped = embedding_name.strip(',')
             if len(stripped) < len(embedding_name):
                 embed = load_embed(stripped, self.embedding_directory, self.embedding_size, self.embedding_key)
-                return (embed, "{} {}".format(embedding_name[len(stripped):], leftover))
-        return (embed, leftover)
+                return (embed, embedding_name, "{} {}".format(embedding_name[len(stripped):], leftover))
+        return (embed, embedding_name, leftover)
 
     def pad_tokens(self, tokens, amount):
         if self.pad_left:
@@ -585,7 +591,7 @@ class SDTokenizer:
         tokens = []
         for weighted_segment, weight in parsed_weights:
             to_tokenize = unescape_important(weighted_segment)
-            split = re.split(' {0}|\n{0}'.format(self.embedding_identifier), to_tokenize)
+            split = re.split(r'(?<=\s){}'.format(re.escape(self.embedding_identifier)), to_tokenize)
             to_tokenize = [split[0]]
             for i in range(1, len(split)):
                 to_tokenize.append("{}{}".format(self.embedding_identifier, split[i]))
@@ -595,7 +601,7 @@ class SDTokenizer:
                 # if we find an embedding, deal with the embedding
                 if word.startswith(self.embedding_identifier) and self.embedding_directory is not None:
                     embedding_name = word[len(self.embedding_identifier):].strip('\n')
-                    embed, leftover = self._try_get_embedding(embedding_name)
+                    embed, embedding_name, leftover = self._try_get_embedding(embedding_name)
                     if embed is None:
                         logging.warning(f"warning, embedding:{embedding_name} does not exist, ignoring")
                     else:

comfy/text_encoders/llama.py

@@ -937,22 +937,41 @@ class BaseGenerate:
             return torch.argmax(logits, dim=-1, keepdim=True)
 
         # Sampling mode
-        if repetition_penalty != 1.0:
-            for i in range(logits.shape[0]):
-                for token_id in set(token_history):
-                    logits[i, token_id] *= repetition_penalty if logits[i, token_id] < 0 else 1/repetition_penalty
-
-        if presence_penalty is not None and presence_penalty != 0.0:
-            for i in range(logits.shape[0]):
-                for token_id in set(token_history):
-                    logits[i, token_id] -= presence_penalty
+        if len(token_history) > 0 and (repetition_penalty != 1.0 or (presence_penalty is not None and presence_penalty != 0.0)):
+            token_ids = torch.tensor(list(set(token_history)), device=logits.device)
+            token_logits = logits[:, token_ids]
+            if repetition_penalty != 1.0:
+                token_logits = torch.where(token_logits < 0, token_logits * repetition_penalty, token_logits / repetition_penalty)
+            if presence_penalty is not None and presence_penalty != 0.0:
+                token_logits = token_logits - presence_penalty
+            logits[:, token_ids] = token_logits
 
         if temperature != 1.0:
             logits = logits / temperature
 
         if top_k > 0:
-            indices_to_remove = logits < torch.topk(logits, top_k)[0][..., -1, None]
-            logits[indices_to_remove] = torch.finfo(logits.dtype).min
+            top_k = min(top_k, logits.shape[-1])
+            logits, top_indices = torch.topk(logits, top_k)
+
+            if min_p > 0.0:
+                probs_before_filter = torch.nn.functional.softmax(logits, dim=-1)
+                top_probs, _ = probs_before_filter.max(dim=-1, keepdim=True)
+                min_threshold = min_p * top_probs
+                indices_to_remove = probs_before_filter < min_threshold
+                logits[indices_to_remove] = torch.finfo(logits.dtype).min
+
+            if top_p < 1.0:
+                sorted_logits, sorted_indices = torch.sort(logits, descending=True)
+                cumulative_probs = torch.cumsum(torch.nn.functional.softmax(sorted_logits, dim=-1), dim=-1)
+                sorted_indices_to_remove = cumulative_probs > top_p
+                sorted_indices_to_remove[..., 0] = False
+                indices_to_remove = torch.zeros_like(logits, dtype=torch.bool)
+                indices_to_remove.scatter_(1, sorted_indices, sorted_indices_to_remove)
+                logits[indices_to_remove] = torch.finfo(logits.dtype).min
+
+            probs = torch.nn.functional.softmax(logits, dim=-1)
+            next_token = torch.multinomial(probs, num_samples=1, generator=generator)
+            return top_indices.gather(1, next_token)
 
         if min_p > 0.0:
             probs_before_filter = torch.nn.functional.softmax(logits, dim=-1)

comfy_api_nodes/util/request_logger.py

@@ -9,6 +9,7 @@ from typing import Any
 import folder_paths
 
 logger = logging.getLogger(__name__)
+_SENSITIVE_HEADERS = {"authorization", "x-api-key"}
 
 
 def get_log_directory():
@@ -73,6 +74,10 @@ def _format_data_for_logging(data: Any) -> str:
     return str(data)
 
 
+def _redact_headers(headers: dict) -> dict:
+    return {k: ("***" if k.lower() in _SENSITIVE_HEADERS else v) for k, v in headers.items()}
+
+
 def log_request_response(
     operation_id: str,
     request_method: str,
@@ -101,7 +106,7 @@ def log_request_response(
         log_content.append(f"Method: {request_method}")
         log_content.append(f"URL: {request_url}")
         if request_headers:
-            log_content.append(f"Headers:\n{_format_data_for_logging(request_headers)}")
+            log_content.append(f"Headers:\n{_format_data_for_logging(_redact_headers(request_headers))}")
         if request_params:
             log_content.append(f"Params:\n{_format_data_for_logging(request_params)}")
         if request_data is not None:

comfy_extras/nodes_color.py

@@ -16,23 +16,30 @@ class ColorToRGBInt(io.ComfyNode):
             ],
             outputs=[
                 io.Int.Output(display_name="rgb_int"),
-                io.Color.Output(display_name="hex")
+                io.Color.Output(display_name="hex"),
+                io.Float.Output(display_name="alpha"),
             ],
         )
 
     @classmethod
     def execute(cls, color: str) -> io.NodeOutput:
-        # expect format #RRGGBB
-        if len(color) != 7 or color[0] != "#":
-            raise ValueError("Color must be in format #RRGGBB")
+        # expect format #RRGGBB or #RRGGBBAA
+        if len(color) not in (7, 9) or color[0] != "#":
+            raise ValueError("Color must be in format #RRGGBB or #RRGGBBAA")
         try:
             int(color[1:], 16)
         except ValueError:
-            raise ValueError("Color must be in format #RRGGBB") from None
+            raise ValueError("Color must be in format #RRGGBB or #RRGGBBAA") from None
+
+        alpha = 1.0
+        if len(color) == 9:
+            alpha = int(color[7:9], 16) / 255.0
+            color = color[:7]
+
         r, g, b = hex_to_rgb(color)
 
         rgb_int = r * 256 * 256 + g * 256 + b
-        return io.NodeOutput(rgb_int, color)
+        return io.NodeOutput(rgb_int, color, alpha)
 
 
 class ColorExtension(ComfyExtension):

requirements.txt

@@ -1,6 +1,6 @@
 comfyui-frontend-package==1.45.20
 comfyui-workflow-templates==0.11.2
-comfyui-embedded-docs==0.5.6
+comfyui-embedded-docs==0.5.7
 torch
 torchsde
 torchvision

tests-unit/app_test/model_download_test.py

@@ -0,0 +1,244 @@
+import os
+
+import pytest
+
+import folder_paths
+from app.model_download import (
+    ModelDownloadError,
+    ModelDownloadRequest,
+    is_allowed_model_download_url,
+    normalize_model_relative_path,
+    open_model_download_response,
+    parse_model_download_request,
+    resolve_model_download_destination,
+)
+
+
+class _FakeResponse:
+    """Minimal stand-in for ``aiohttp.ClientResponse`` for the redirect tests."""
+
+    def __init__(self, status, headers=None):
+        self.status = status
+        self.headers = headers or {}
+        self.released = False
+
+    def release(self):
+        self.released = True
+
+    async def __aenter__(self):
+        return self
+
+    async def __aexit__(self, exc_type, exc, tb):
+        self.released = True
+
+
+class _FakeSession:
+    """Hands out queued ``_FakeResponse`` objects in order."""
+
+    def __init__(self, responses):
+        self._responses = list(responses)
+        self.calls = []
+
+    async def get(self, url, allow_redirects, timeout):
+        self.calls.append((url, allow_redirects))
+        if not self._responses:
+            raise AssertionError("Unexpected extra session.get call")
+        return self._responses.pop(0)
+
+
+def test_parse_model_download_request_allows_huggingface_model_url():
+    request = parse_model_download_request({
+        "name": "nested/model.safetensors",
+        "url": "https://huggingface.co/org/repo/resolve/main/model.safetensors",
+        "directory": "checkpoints",
+    })
+
+    assert request == ModelDownloadRequest(
+        name="nested/model.safetensors",
+        url="https://huggingface.co/org/repo/resolve/main/model.safetensors",
+        directory="checkpoints",
+    )
+
+
+@pytest.mark.parametrize(
+    "url",
+    [
+        "http://localhost:8000/model.safetensors",
+        "http://huggingface.co/org/repo/resolve/main/model.safetensors",
+        "https://example.com/model.safetensors",
+        "https://huggingface.co.evil.com/model.safetensors",
+    ],
+)
+def test_download_url_allowlist_rejects_untrusted_or_plain_http_urls(url):
+    assert is_allowed_model_download_url(url) is False
+
+
+@pytest.mark.parametrize(
+    "url",
+    [
+        # Direct HF model URLs.
+        "https://huggingface.co/org/repo/resolve/main/model.safetensors",
+        # HF LFS CDN subdomains: this is where `/resolve/main/...` redirects
+        # land, so the allowlist must accept them or downloads break.
+        "https://cdn-lfs.huggingface.co/repos/abc/def/model.safetensors",
+        "https://cdn-lfs-us-1.huggingface.co/repos/abc/def/model.safetensors",
+        # Civitai download endpoints (PR objective: support Civitai too).
+        "https://civitai.com/api/download/models/12345",
+        "https://civitai.red/api/download/models/12345",
+    ],
+)
+def test_download_url_allowlist_accepts_huggingface_and_civitai_urls(url):
+    assert is_allowed_model_download_url(url) is True
+
+
+@pytest.mark.parametrize(
+    "name, expected",
+    [
+        ("model.safetensors", "model.safetensors"),
+        ("sub/model.safetensors", "sub/model.safetensors"),
+        ("nested/dir/model.safetensors", "nested/dir/model.safetensors"),
+        # Backslashes are normalized to forward slashes so Windows-style
+        # paths land in the same place as the POSIX equivalents.
+        ("nested\\dir\\model.safetensors", "nested/dir/model.safetensors"),
+    ],
+)
+def test_normalize_model_relative_path_accepts_safe_paths(name, expected):
+    assert normalize_model_relative_path(name) == expected
+
+
+@pytest.mark.parametrize(
+    "name",
+    [
+        "../model.safetensors",
+        "nested/../../model.safetensors",
+        "/absolute/model.safetensors",
+        "model.safetensors\x00",
+    ],
+)
+def test_normalize_model_relative_path_rejects_unsafe_paths(name):
+    with pytest.raises(ModelDownloadError):
+        normalize_model_relative_path(name)
+
+
+def test_parse_model_download_request_rejects_unsupported_extensions():
+    with pytest.raises(ModelDownloadError):
+        parse_model_download_request({
+            "name": "model.gguf",
+            "url": "https://huggingface.co/org/repo/resolve/main/model.gguf",
+            "directory": "checkpoints",
+        })
+
+
+def test_resolve_model_download_destination_uses_configured_model_folder(tmp_path, monkeypatch):
+    model_root = tmp_path / "models" / "checkpoints"
+    monkeypatch.setattr(folder_paths, "folder_names_and_paths", {
+        "checkpoints": ([str(model_root)], {".safetensors"}),
+    })
+
+    destination = resolve_model_download_destination(ModelDownloadRequest(
+        name="sub/model.safetensors",
+        url="https://huggingface.co/org/repo/resolve/main/model.safetensors",
+        directory="checkpoints",
+    ))
+
+    assert destination.directory == "checkpoints"
+    assert destination.relative_path == "sub/model.safetensors"
+    assert destination.full_path == os.path.join(str(model_root), "sub", "model.safetensors")
+    assert destination.already_exists is False
+
+
+def test_resolve_model_download_destination_reuses_existing_model(tmp_path, monkeypatch):
+    model_root = tmp_path / "models" / "checkpoints"
+    model_root.mkdir(parents=True)
+    existing = model_root / "model.safetensors"
+    existing.write_bytes(b"model")
+    monkeypatch.setattr(folder_paths, "folder_names_and_paths", {
+        "checkpoints": ([str(model_root)], {".safetensors"}),
+    })
+
+    destination = resolve_model_download_destination(ModelDownloadRequest(
+        name="model.safetensors",
+        url="https://huggingface.co/org/repo/resolve/main/model.safetensors",
+        directory="checkpoints",
+    ))
+
+    assert destination.full_path == str(existing)
+    assert destination.already_exists is True
+
+
+@pytest.mark.parametrize("directory", ["configs", "custom_nodes", "unknown"])
+def test_resolve_model_download_destination_rejects_blocked_or_unknown_directories(tmp_path, monkeypatch, directory):
+    monkeypatch.setattr(folder_paths, "folder_names_and_paths", {
+        "configs": ([str(tmp_path / "configs")], {".yaml"}),
+        "custom_nodes": ([str(tmp_path / "custom_nodes")], set()),
+    })
+
+    with pytest.raises(ModelDownloadError):
+        resolve_model_download_destination(ModelDownloadRequest(
+            name="model.safetensors",
+            url="https://huggingface.co/org/repo/resolve/main/model.safetensors",
+            directory=directory,
+        ))
+
+
+@pytest.mark.asyncio
+async def test_open_model_download_response_follows_allowed_subdomain_redirect():
+    """HF redirects /resolve/main/... to cdn-lfs.huggingface.co; that must work."""
+    session = _FakeSession([
+        _FakeResponse(302, {"Location": "https://cdn-lfs.huggingface.co/repos/abc/model.safetensors"}),
+        _FakeResponse(200),
+    ])
+
+    response = await open_model_download_response(
+        session, "https://huggingface.co/org/repo/resolve/main/model.safetensors"
+    )
+
+    assert response.status == 200
+    assert session.calls == [
+        ("https://huggingface.co/org/repo/resolve/main/model.safetensors", False),
+        ("https://cdn-lfs.huggingface.co/repos/abc/model.safetensors", False),
+    ]
+
+
+@pytest.mark.asyncio
+async def test_open_model_download_response_rejects_offsite_redirect():
+    """A redirect leaving the allowlist must surface as a 403 instead of being followed."""
+    session = _FakeSession([
+        _FakeResponse(302, {"Location": "https://attacker.example.com/payload"}),
+    ])
+
+    with pytest.raises(ModelDownloadError) as exc_info:
+        await open_model_download_response(
+            session, "https://huggingface.co/org/repo/resolve/main/model.safetensors"
+        )
+
+    assert exc_info.value.status == 403
+    # The initial request was issued with redirects disabled, otherwise
+    # the validation above would be a no-op.
+    assert session.calls[0][1] is False
+
+
+@pytest.mark.asyncio
+async def test_open_model_download_response_rejects_redirect_without_location():
+    session = _FakeSession([_FakeResponse(302)])
+
+    with pytest.raises(ModelDownloadError) as exc_info:
+        await open_model_download_response(
+            session, "https://huggingface.co/org/repo/resolve/main/model.safetensors"
+        )
+
+    assert exc_info.value.status == 502
+
+
+@pytest.mark.asyncio
+async def test_open_model_download_response_stops_after_too_many_redirects():
+    session = _FakeSession(
+        [_FakeResponse(302, {"Location": "https://cdn-lfs.huggingface.co/loop"})] * 10
+    )
+
+    with pytest.raises(ModelDownloadError) as exc_info:
+        await open_model_download_response(
+            session, "https://huggingface.co/org/repo/resolve/main/model.safetensors"
+        )
+
+    assert exc_info.value.status == 502