Merge 3c3d7feefe into b08debceca

Co-authored-by: coderabbitai[bot] <136622811+coderabbitai[bot]@users.noreply.github.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)

folder_paths.py

@@ -4,7 +4,7 @@ import mimetypes
 import logging
 from typing import Literal, List
 from collections.abc import Collection
-
+import json
 from comfy.cli_args import args
 
 supported_pt_extensions: set[str] = {'.ckpt', '.pt', '.pt2', '.bin', '.pth', '.safetensors', '.pkl', '.sft'}
@@ -13,11 +13,38 @@ folder_names_and_paths: dict[str, tuple[list[str], set[str]]] = {}
 
 # --base-directory - Resets all default paths configured in folder_paths with a new base path
 if args.base_directory:
-    base_path = os.path.abspath(args.base_directory)
+    raw_path = os.path.abspath(args.base_directory)
 else:
-    base_path = os.path.dirname(os.path.realpath(__file__))
+    raw_path = os.path.dirname(os.path.realpath(__file__))
 
-models_dir = os.path.join(base_path, "models")
+
+def valid_path(saved_path,description = "Instance Folder"):
+    if os.path.exists(saved_path):
+        return saved_path
+
+    redirect_path = os.path.join(os.path.dirname(__file__),"path.redirects.json")
+    if os.path.exists(redirect_path):
+        with open(redirect_path,"r") as f:
+            redirects = json.load(f)
+            if description in redirects and os.path.exists(redirects[description]):
+                return redirects[description]
+    print(f"\n[!] ALERT: {description} not found at: {saved_path}")
+    new_path= input(f"Please paste the updated absolute path for '{description}': ").strip()        
+    redirects = {}
+    if os.path.exists(redirect_path):
+        with open(redirect_path, "r") as f:
+            redirects = json.load(f)
+    redirects[description] = new_path
+    with open(redirect_path, "w") as f:
+        json.dump(redirects, f, indent=4)
+        
+    return new_path
+
+
+
+base_path = valid_path(raw_path,description = "ComfyUI Core Directory")
+
+models_dir =valid_path( os.path.join(base_path,"models"),description = "Model Directory")
 folder_names_and_paths["checkpoints"] = ([os.path.join(models_dir, "checkpoints")], supported_pt_extensions)
 folder_names_and_paths["configs"] = ([os.path.join(models_dir, "configs")], [".yaml"])
 
@@ -60,7 +87,7 @@ folder_names_and_paths["optical_flow"] = ([os.path.join(models_dir, "optical_flo
 
 folder_names_and_paths["detection"] = ([os.path.join(models_dir, "detection")], supported_pt_extensions)
 
-output_directory = os.path.join(base_path, "output")
+output_directory =valid_path( os.path.join(base_path, "output"),"Output Directory")
 temp_directory = os.path.join(base_path, "temp")
 input_directory = os.path.join(base_path, "input")
 user_directory = os.path.join(base_path, "user")

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