Merge 3b87c2c397 into b08debceca

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_extras/nodes_ideogram4.py

@@ -1,6 +1,7 @@
 """Ideogram 4 sampling helper
 """
 
+import enum
 import math
 
 import torch
@@ -10,6 +11,45 @@ from comfy_api.latest import ComfyExtension, io
 _LOGSNR_MIN = -15.0
 _LOGSNR_MAX = 18.0
 
+class Ideogram4Enum(enum.Enum):
+    QUALITY = "Quality"
+    HIGH = "High"
+    DEFAULT = "Default"
+    FAST = "Fast"
+    TURBO = "Turbo"
+
+IDEOGRAM4_PRESET_CONFIGS = {
+  Ideogram4Enum.QUALITY.value: {
+    "num_steps": 48,
+    "mu": 0.0,
+    "std": 1.5,
+    "preset_id": "V4_QUALITY_48"
+  },
+  Ideogram4Enum.HIGH.value: {
+    "num_steps": 34,
+    "mu": 0.0,
+    "std": 1.6875,
+    "preset_id": "V4_HIGH_34"
+  },
+  Ideogram4Enum.DEFAULT.value: {
+    "num_steps": 20,
+    "mu": 0.0,
+    "std": 1.75,
+    "preset_id": "V4_DEFAULT_20"
+  },
+  Ideogram4Enum.FAST.value: {
+    "num_steps": 16,
+    "mu": 0.25,
+    "std": 1.8375,
+    "preset_id": "V4_FAST_16"
+  },
+  Ideogram4Enum.TURBO.value: {
+    "num_steps": 12,
+    "mu": 0.5,
+    "std": 1.75,
+    "preset_id": "V4_TURBO_12"
+  }
+}
 
 def _logit_normal_schedule(u, mean, std):
     # Reference time (0=noise..1=clean) via the probit/ndtri quantile.
@@ -54,10 +94,41 @@ class Ideogram4Scheduler(io.ComfyNode):
         return io.NodeOutput(ideogram4_sigmas(steps, width, height, mu, std))
 
 
+class Ideogram4SchedulerPreset(Ideogram4Scheduler):
+    @classmethod
+    def define_schema(cls) -> io.Schema:
+        return io.Schema(
+            node_id="Ideogram4SchedulerPreset",
+            display_name="Ideogram 4 Scheduler (Presets)",
+            category="sampling/custom_sampling/schedulers",
+            description="Schedule Presets for Ideogram 4. They are as follows: Quality=48, High=34, Default=20, Fast=16, Turbo=12",
+            inputs=[
+                io.Combo.Input("preset", options=[e.value for e in Ideogram4Enum], default=Ideogram4Enum.DEFAULT.value),
+                io.Int.Input("width", default=1024, min=256, max=8192, step=16),
+                io.Int.Input("height", default=1024, min=256, max=8192, step=16),
+            ],
+            outputs=[io.Sigmas.Output()],
+        )
+
+    @classmethod
+    def execute(cls, preset, width, height) -> io.NodeOutput:
+        config = IDEOGRAM4_PRESET_CONFIGS.get(preset)
+        if not config:
+            raise ValueError(f"Invalid preset: {preset}")
+
+        return super().execute(
+            steps=config["num_steps"],
+            width=width,
+            height=height,
+            mu=config["mu"],
+            std=config["std"]
+        )
+
+
 class Ideogram4Extension(ComfyExtension):
     @override
     async def get_node_list(self) -> list[type[io.ComfyNode]]:
-        return [Ideogram4Scheduler]
+        return [Ideogram4Scheduler, Ideogram4SchedulerPreset]
 
 
 async def comfy_entrypoint() -> Ideogram4Extension:

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