Merge 9789de07b8 into b08debceca

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

app/database/db.py

@@ -57,19 +57,62 @@ def get_alembic_config():
 
     config = Config(config_path)
     config.set_main_option("script_location", scripts_path)
-    config.set_main_option("sqlalchemy.url", args.database_url)
+    config.set_main_option("sqlalchemy.url", get_database_url())
 
     return config
 
 
-def get_db_path():
+def get_database_url():
+    if getattr(args, "database_url_explicit", False):
+        return args.database_url
+
+    import folder_paths
+
+    db_path = os.path.join(folder_paths.get_user_directory(), "comfyui.db")
+    return f"sqlite:///{db_path}"
+
+
+def get_legacy_default_db_path():
     url = args.database_url
     if url.startswith("sqlite:///"):
-        return url.split("///")[1]
+        return url.split("///", 1)[1]
+    return None
+
+
+def get_db_path():
+    url = get_database_url()
+    if url.startswith("sqlite:///"):
+        return url.split("///", 1)[1]
     else:
         raise ValueError(f"Unsupported database URL '{url}'.")
 
 
+def copy_legacy_default_db(db_path):
+    if getattr(args, "database_url_explicit", False):
+        return
+
+    legacy_db_path = get_legacy_default_db_path()
+    if legacy_db_path is None:
+        return
+
+    if os.path.abspath(legacy_db_path) == os.path.abspath(db_path):
+        return
+
+    if os.path.exists(db_path) or not os.path.exists(legacy_db_path):
+        return
+
+    shutil.copy(legacy_db_path, db_path)
+    logging.info(f"Copied legacy database from '{legacy_db_path}' to '{db_path}'")
+
+
+def prepare_file_db_path(db_path):
+    db_dir = os.path.dirname(db_path)
+    if db_dir:
+        os.makedirs(db_dir, exist_ok=True)
+
+    copy_legacy_default_db(db_path)
+
+
 _db_lock = None
 
 def _acquire_file_lock(db_path):
@@ -97,7 +140,7 @@ def _is_memory_db(db_url):
 
 
 def init_db():
-    db_url = args.database_url
+    db_url = get_database_url()
     logging.debug(f"Database URL: {db_url}")
 
     if _is_memory_db(db_url):
@@ -134,6 +177,7 @@ def _init_memory_db(db_url):
 def _init_file_db(db_url):
     """Initialize a file-backed SQLite database using Alembic migrations."""
     db_path = get_db_path()
+    prepare_file_db_path(db_path)
     db_exists = os.path.exists(db_path)
 
     config = get_alembic_config()

comfy/cli_args.py

@@ -1,6 +1,7 @@
 import argparse
 import enum
 import os
+import sys
 import comfy.options
 
 
@@ -246,8 +247,13 @@ parser.add_argument("--list-feature-flags", action="store_true", help="Print the
 
 if comfy.options.args_parsing:
     args = parser.parse_args()
+    args.database_url_explicit = any(
+        arg == "--database-url" or arg.startswith("--database-url=")
+        for arg in sys.argv[1:]
+    )
 else:
     args = parser.parse_args([])
+    args.database_url_explicit = False
 
 if args.cache_ram is not None and len(args.cache_ram) > 2:
     parser.error("--cache-ram accepts at most two values: active GB and inactive GB")

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:

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/database_path_test.py

@@ -0,0 +1,78 @@
+import os
+
+from app.database import db
+
+
+def test_default_database_url_uses_effective_user_directory(monkeypatch, tmp_path):
+    user_dir = tmp_path / "custom_user"
+    user_dir.mkdir()
+
+    monkeypatch.setattr(db.args, "database_url_explicit", False, raising=False)
+    monkeypatch.setattr("folder_paths.get_user_directory", lambda: str(user_dir))
+
+    assert db.get_database_url() == f"sqlite:///{user_dir / 'comfyui.db'}"
+
+
+def test_explicit_database_url_is_preserved(monkeypatch):
+    database_url = "sqlite:///:memory:"
+
+    monkeypatch.setattr(db.args, "database_url", database_url)
+    monkeypatch.setattr(db.args, "database_url_explicit", True, raising=False)
+
+    assert db.get_database_url() == database_url
+
+
+def test_legacy_default_database_is_copied_to_effective_user_directory(monkeypatch, tmp_path):
+    legacy_db = tmp_path / "install" / "user" / "comfyui.db"
+    user_dir = tmp_path / "custom_user"
+    legacy_db.parent.mkdir(parents=True)
+    user_dir.mkdir()
+    legacy_db.write_bytes(b"legacy db")
+
+    monkeypatch.setattr(db.args, "database_url_explicit", False, raising=False)
+    monkeypatch.setattr("folder_paths.get_user_directory", lambda: str(user_dir))
+    monkeypatch.setattr(db, "get_legacy_default_db_path", lambda: str(legacy_db))
+
+    db.copy_legacy_default_db(str(user_dir / "comfyui.db"))
+
+    assert (user_dir / "comfyui.db").read_bytes() == b"legacy db"
+    assert legacy_db.read_bytes() == b"legacy db"
+
+
+def test_legacy_default_database_does_not_overwrite_existing_effective_db(monkeypatch, tmp_path):
+    legacy_db = tmp_path / "install" / "user" / "comfyui.db"
+    user_db = tmp_path / "custom_user" / "comfyui.db"
+    legacy_db.parent.mkdir(parents=True)
+    user_db.parent.mkdir(parents=True)
+    legacy_db.write_bytes(b"legacy db")
+    user_db.write_bytes(b"user db")
+
+    monkeypatch.setattr(db.args, "database_url_explicit", False, raising=False)
+    monkeypatch.setattr(db, "get_legacy_default_db_path", lambda: str(legacy_db))
+
+    db.copy_legacy_default_db(str(user_db))
+
+    assert user_db.read_bytes() == b"user db"
+    assert legacy_db.read_bytes() == b"legacy db"
+
+
+def test_prepare_file_database_creates_parent_directory(monkeypatch, tmp_path):
+    db_path = tmp_path / "nested" / "comfyui.db"
+
+    monkeypatch.setattr(db.args, "database_url_explicit", False, raising=False)
+    monkeypatch.setattr(db, "copy_legacy_default_db", lambda path: None)
+
+    db.prepare_file_db_path(str(db_path))
+
+    assert db_path.parent.is_dir()
+
+
+def test_prepare_file_database_accepts_relative_database_path(monkeypatch, tmp_path):
+    monkeypatch.chdir(tmp_path)
+    monkeypatch.setattr(db.args, "database_url_explicit", True, raising=False)
+    monkeypatch.setattr(db, "copy_legacy_default_db", lambda path: None)
+
+    db.prepare_file_db_path("relative.db")
+
+    assert os.getcwd() == str(tmp_path)
+    assert not list(tmp_path.iterdir())