feat: add memory mode

README.md

@@ -125,8 +125,10 @@ Each top-level `.txt` file appears in the `system_prompt` dropdown. Choose
 | `top_k` | Top-K sampling cutoff. |
 | `repeat_penalty` | Penalty for repeated tokens. |
 | `ctx_size` | llama.cpp context size. |
-| `n_gpu_layers` | Layers offloaded to GPU. Use `0` for CPU. |
-| `seed` | Random seed. |
+| `memory_mode` | Advanced memory placement mode. `auto` passes no layer flags. |
+| `n_gpu_layers` | Used in GPU layer modes. Passes `--gpu-layers` / `-ngl`. |
+| `n_cpu_moe_layers` | Used in CPU MoE modes. Passes `--n-cpu-moe`. |
+| `seed` | Random seed. Use `-1` for a random seed. |
 | `enable_thinking` | Uses `/think` when enabled and `/no_think` when disabled. |
 | `timeout_seconds` | Maximum runtime before the subprocess is stopped. |
 | `image` | Optional ComfyUI image input. Uses the first image in a batch. |

llama_cli.py

@@ -9,6 +9,13 @@ from pathlib import Path
 
 
 THINK_BLOCK_RE = re.compile(r"<think[^>]*>.*?</think>", flags=re.IGNORECASE | re.DOTALL)
+MAX_LLAMA_SEED = 2**32 - 1
+MEMORY_MODES = (
+    "auto",
+    "gpu_layers",
+    "cpu_moe_layers",
+    "gpu_and_cpu_moe_layers",
+)
 
 
 def tensor_to_temp_png(image) -> Path:
@@ -51,6 +58,17 @@ def _split_extra_args(extra_args: str) -> list[str]:
     return [part.strip("\"'") for part in parts]
 
 
+def normalize_seed(seed: int) -> int:
+    # llama.cpp uses -1 as "random seed"; otherwise Windows builds parse the
+    # value as an unsigned 32-bit integer.
+    seed = int(seed)
+    if seed == -1:
+        return -1
+    if not 0 <= seed <= MAX_LLAMA_SEED:
+        raise ValueError(f"seed must be -1 or between 0 and {MAX_LLAMA_SEED}")
+    return seed
+
+
 def build_command(
     cli_path: Path,
     model_path: Path,
@@ -63,10 +81,17 @@ def build_command(
     top_k: int,
     repeat_penalty: float,
     ctx_size: int,
+    memory_mode: str,
     n_gpu_layers: int,
+    n_cpu_moe_layers: int,
     seed: int,
     extra_args: str = "",
 ) -> list[str]:
+    if memory_mode not in MEMORY_MODES:
+        raise ValueError(f"Unsupported memory_mode: {memory_mode}")
+    if memory_mode in {"cpu_moe_layers", "gpu_and_cpu_moe_layers"} and n_cpu_moe_layers < 1:
+        raise ValueError("n_cpu_moe_layers must be at least 1 when CPU MoE layers are enabled.")
+
     command = [
         str(cli_path),
         "-m", str(model_path),
@@ -77,10 +102,16 @@ def build_command(
         "--top-k", str(top_k),
         "--repeat-penalty", str(repeat_penalty),
         "-c", str(ctx_size),
-        "-ngl", str(n_gpu_layers),
-        "--seed", str(seed),
+        "--seed", str(normalize_seed(seed)),
     ]
 
+    # In auto mode llama.cpp receives neither flag and uses its own placement
+    # defaults. Other modes pass only the user-selected memory controls.
+    if memory_mode in {"gpu_layers", "gpu_and_cpu_moe_layers"}:
+        command.extend(["-ngl", str(n_gpu_layers)])
+    if memory_mode in {"cpu_moe_layers", "gpu_and_cpu_moe_layers"}:
+        command.extend(["--n-cpu-moe", str(n_cpu_moe_layers)])
+
     if mmproj_path is not None:
         command.extend(["--mmproj", str(mmproj_path)])
     if image_path is not None:
@@ -103,7 +134,9 @@ def run_llama_cli(
     top_k: int,
     repeat_penalty: float,
     ctx_size: int,
+    memory_mode: str,
     n_gpu_layers: int,
+    n_cpu_moe_layers: int,
     seed: int,
     enable_thinking: bool,
     timeout_seconds: int,
@@ -130,7 +163,9 @@ def run_llama_cli(
             top_k=top_k,
             repeat_penalty=repeat_penalty,
             ctx_size=ctx_size,
+            memory_mode=memory_mode,
             n_gpu_layers=n_gpu_layers,
+            n_cpu_moe_layers=n_cpu_moe_layers,
             seed=seed,
             extra_args=extra_args,
         )

nodes.py

@@ -10,7 +10,7 @@ from .folder_registry import (
     system_prompt_options,
 )
 from .llama_binary import ensure_llama_cli
-from .llama_cli import extract_thinking, run_llama_cli
+from .llama_cli import MAX_LLAMA_SEED, MEMORY_MODES, extract_thinking, run_llama_cli
 
 
 class QwenGGUF:
@@ -74,17 +74,30 @@ class QwenGGUF:
                     "step": 512,
                     "tooltip": "llama.cpp context window size in tokens.",
                 }),
+                "memory_mode": (list(MEMORY_MODES), {
+                    "default": "auto",
+                    "tooltip": "auto passes no layer flags; other modes pass GPU layers, CPU MoE layers, or both.",
+                    "advanced": True,
+                }),
                 "n_gpu_layers": ("INT", {
                     "default": 99,
                     "min": -1,
                     "max": 999,
-                    "tooltip": "Number of layers to offload to GPU. 0 runs on CPU.",
+                    "tooltip": "Used only in gpu_layers and gpu_and_cpu_moe_layers modes. Number of layers to offload to GPU.",
+                    "advanced": True,
+                }),
+                "n_cpu_moe_layers": ("INT", {
+                    "default": 1,
+                    "min": 1,
+                    "max": 999,
+                    "tooltip": "Used only in cpu_moe_layers and gpu_and_cpu_moe_layers modes. Keeps MoE weights of the first N layers on CPU.",
+                    "advanced": True,
                 }),
                 "seed": ("INT", {
                     "default": 1,
-                    "min": 0,
-                    "max": 0xffffffffffffffff,
-                    "tooltip": "Random seed used by llama.cpp.",
+                    "min": -1,
+                    "max": MAX_LLAMA_SEED,
+                    "tooltip": "Random seed used by llama.cpp. Use -1 for a random seed.",
                 }),
                 "enable_thinking": ("BOOLEAN", {
                     "default": False,
@@ -131,7 +144,9 @@ class QwenGGUF:
         top_k: int,
         repeat_penalty: float,
         ctx_size: int,
+        memory_mode: str,
         n_gpu_layers: int,
+        n_cpu_moe_layers: int,
         seed: int,
         enable_thinking: bool,
         timeout_seconds: int,
@@ -159,7 +174,9 @@ class QwenGGUF:
             top_k=top_k,
             repeat_penalty=repeat_penalty,
             ctx_size=ctx_size,
+            memory_mode=memory_mode,
             n_gpu_layers=n_gpu_layers,
+            n_cpu_moe_layers=n_cpu_moe_layers,
             seed=seed,
             enable_thinking=enable_thinking,
             timeout_seconds=timeout_seconds,