Merge fe89d39c1b into b08debceca

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

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:

folder_paths.py

@@ -473,7 +473,12 @@ def cached_filename_list_(folder_name: str) -> tuple[list[str], dict[str, float]
     for x in out[1]:
         time_modified = out[1][x]
         folder = x
-        if os.path.getmtime(folder) != time_modified:
+        try:
+            if os.path.getmtime(folder) != time_modified:
+                return None
+        except OSError:
+            # A tracked folder was deleted or became inaccessible; treat the
+            # cache as stale so it gets rebuilt instead of raising.
             return None
 
     folders = folder_names_and_paths[folder_name]

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/folder_paths_test/cache_invalidation_test.py

@@ -0,0 +1,46 @@
+import os
+import shutil
+import tempfile
+
+import pytest
+
+import folder_paths
+
+
+@pytest.fixture
+def model_folder():
+    """Register a temporary model category with a tracked subfolder."""
+    folder_name = "cache_invalidation_test_cat"
+    with tempfile.TemporaryDirectory() as base:
+        category = os.path.join(base, "category")
+        subfolder = os.path.join(category, "sub")
+        os.makedirs(subfolder)
+        open(os.path.join(category, "a.safetensors"), "w").close()
+        open(os.path.join(subfolder, "b.safetensors"), "w").close()
+
+        folder_paths.folder_names_and_paths[folder_name] = (
+            [category],
+            {".safetensors"},
+        )
+        try:
+            yield folder_name, category, subfolder
+        finally:
+            folder_paths.folder_names_and_paths.pop(folder_name, None)
+            folder_paths.filename_list_cache.pop(folder_name, None)
+            folder_paths.cache_helper.clear()
+
+
+def test_rebuilds_when_tracked_subfolder_deleted(model_folder):
+    folder_name, _category, subfolder = model_folder
+
+    # Populate the filename cache, which records the mtime of every subfolder.
+    initial = folder_paths.get_filename_list(folder_name)
+    assert sorted(initial) == ["a.safetensors", "sub/b.safetensors"]
+
+    # Remove a tracked subfolder at runtime (e.g. user deletes a model folder).
+    shutil.rmtree(subfolder)
+
+    # The cache must be treated as stale and rebuilt, not crash with
+    # FileNotFoundError when probing the deleted folder's mtime.
+    refreshed = folder_paths.get_filename_list(folder_name)
+    assert refreshed == ["a.safetensors"]