Changes
diff --git a/diffusion_cli/cli.py b/diffusion_cli/cli.py
index 3f7461c..3c8cab9 100644
--- a/diffusion_cli/cli.py
+++ b/diffusion_cli/cli.py
@@ -23,97 +23,242 @@ from diffusion_cli.sampling import sampleLatents
from diffusion_cli.server import serve, validateServerConfig
from diffusion_cli.generation_service import GenerationService
+MODEL_PROFILE_LIST_VALUE = "__list_model_profiles__"
-def buildParser() -> argparse.ArgumentParser:
+
+def _modelProfileHelp(profile_names: tuple[str, ...]) -> str:
+ """Describe profile selection and the values available in the config."""
+
+ if profile_names:
+ available = ", ".join(profile_names)
+ return (
+ "Select a named local model profile. Available configured "
+ f"values: {available}. Omit NAME to list them and exit."
+ )
+ return (
+ "Select a named local model profile. Values come from "
+ "[model_profiles] in ~/.config/diffusion.toml; legacy is added "
+ "when the legacy [models] table is configured. Omit NAME to list "
+ "the configured profiles and exit."
+ )
+
+
+def buildParser(
+ model_profile_names: tuple[str, ...] | None = None,
+) -> argparse.ArgumentParser:
"""Build the argparse command line parser."""
+ profile_names = tuple(model_profile_names or ())
parser = argparse.ArgumentParser(
prog="diffusion-cli",
- description="Standalone local diffusion model CLI.",
+ description=(
+ "Generate images from local diffusion checkpoints, or start the "
+ "HTTP server."
+ ),
+ epilog=(
+ "Examples:\n"
+ " diffusion-cli --prompt \"a ceramic mug\"\n"
+ " diffusion-cli --model-profile krea2-turbo "
+ "--prompt \"a fox in snow\"\n"
+ " diffusion-cli --inspect-models"
+ ),
+ formatter_class=argparse.RawDescriptionHelpFormatter,
)
parser.add_argument(
"--model-profile",
- help="Named local model profile to use.",
+ metavar="NAME",
+ nargs="?",
+ const=MODEL_PROFILE_LIST_VALUE,
+ help=_modelProfileHelp(profile_names),
+ )
+ parser.add_argument(
+ "--prompt",
+ metavar="TEXT",
+ help="Positive text prompt to generate.",
)
- parser.add_argument("--prompt", help="Positive prompt.")
parser.add_argument(
"--negative-prompt",
- help="Negative prompt.",
- )
- parser.add_argument("--seed", type=int, help="Noise seed.")
- parser.add_argument("--width", type=int)
- parser.add_argument("--height", type=int)
- parser.add_argument("--batch-size", type=int)
- parser.add_argument("--steps", type=int)
- parser.add_argument("--cfg", type=float)
- parser.add_argument("--mu", type=float)
- parser.add_argument("--shift-y1", type=float)
- parser.add_argument("--shift-y2", type=float)
- parser.add_argument("--output", type=Path)
- parser.add_argument("--output-extension")
- parser.add_argument("--output-quality", type=int)
- parser.add_argument("--device")
+ metavar="TEXT",
+ help="Negative text prompt; overrides the configured default.",
+ )
+ parser.add_argument(
+ "--seed",
+ type=int,
+ metavar="INTEGER",
+ help="Noise seed. Omit it to choose a random seed.",
+ )
+ parser.add_argument(
+ "--width",
+ type=int,
+ metavar="PIXELS",
+ help=(
+ "Output width in pixels. Z-Image requires a multiple of 8; "
+ "Krea aligns dimensions to a multiple of 16."
+ ),
+ )
+ parser.add_argument(
+ "--height",
+ type=int,
+ metavar="PIXELS",
+ help=(
+ "Output height in pixels. Z-Image requires a multiple of 8; "
+ "Krea aligns dimensions to a multiple of 16."
+ ),
+ )
+ parser.add_argument(
+ "--batch-size",
+ type=int,
+ metavar="INTEGER",
+ help="Number of images to generate in the batch.",
+ )
+ parser.add_argument(
+ "--steps",
+ type=int,
+ metavar="INTEGER",
+ help=(
+ "Number of diffusion steps. Krea Turbo defaults to 8 and "
+ "Krea Raw defaults to 52."
+ ),
+ )
+ parser.add_argument(
+ "--cfg",
+ type=float,
+ metavar="FLOAT",
+ help=(
+ "Classifier-free guidance scale. Krea Turbo defaults to 0; "
+ "Krea Raw defaults to 3.5."
+ ),
+ )
+ parser.add_argument(
+ "--mu",
+ type=float,
+ metavar="FLOAT",
+ help=(
+ "Krea flow timestep shift. Turbo defaults to 1.15; Raw "
+ "derives it from the image resolution."
+ ),
+ )
+ parser.add_argument(
+ "--shift-y1",
+ type=float,
+ metavar="FLOAT",
+ help="Krea Raw minimum-resolution timestep shift, default 0.5.",
+ )
+ parser.add_argument(
+ "--shift-y2",
+ type=float,
+ metavar="FLOAT",
+ help="Krea Raw maximum-resolution timestep shift, default 1.15.",
+ )
+ parser.add_argument(
+ "--output",
+ type=Path,
+ metavar="PATH",
+ help="Output image path; the configured default is output.png.",
+ )
+ parser.add_argument(
+ "--output-extension",
+ metavar="FORMAT",
+ help="Output format: png, jpg, webp, or avif.",
+ )
+ parser.add_argument(
+ "--output-quality",
+ type=int,
+ metavar="INTEGER",
+ help="Lossy output quality from 1 through 100.",
+ )
+ parser.add_argument(
+ "--device",
+ metavar="CUDA_DEVICE",
+ help="CUDA device, for example cuda or cuda:0.",
+ )
parser.add_argument(
"--checkpoint",
type=Path,
- help="Local all-in-one safetensors checkpoint file.",
+ metavar="PATH",
+ help=(
+ "Legacy local all-in-one safetensors file. Prefer separate "
+ "profile component paths for Krea."
+ ),
)
parser.add_argument(
"--diffusion-model",
type=Path,
- help="Local Z-Image diffusion model safetensors file.",
+ metavar="PATH",
+ help=(
+ "Override the selected profile's local diffusion safetensors "
+ "file."
+ ),
)
parser.add_argument(
"--text-encoder",
type=Path,
- help="Local Qwen text encoder safetensors file.",
+ metavar="PATH",
+ help="Override the selected profile's local text encoder file.",
)
parser.add_argument(
"--vae",
type=Path,
- help="Local VAE safetensors file.",
+ metavar="PATH",
+ help="Override the selected profile's local VAE safetensors file.",
)
parser.add_argument(
"--tokenizer-path",
type=Path,
- help="Local Qwen tokenizer directory.",
+ metavar="PATH",
+ help="Override the selected profile's local tokenizer directory.",
)
parser.add_argument(
"--dtype",
choices=("auto", "bf16", "fp16", "fp32"),
+ help=(
+ "Activation dtype: auto selects a supported default; bf16, "
+ "fp16, and fp32 request a specific dtype."
+ ),
)
parser.add_argument(
"--inspect-models",
action="store_true",
- help="Inspect local safetensors metadata and exit.",
+ help=(
+ "Inspect selected local checkpoints, dtypes, FP8 metadata, and "
+ "profile compatibility, then exit."
+ ),
)
subparsers = parser.add_subparsers(dest="command")
serve_parser = subparsers.add_parser(
"serve",
- help="Start a long-running HTTP API server.",
+ help="Start the local HTTP API server.",
+ description="Start the local HTTP API server for one selected profile.",
)
serve_parser.add_argument(
"--api-profile",
required=True,
choices=tuple(API_PROFILES),
- help="HTTP API profile to expose.",
+ help=(
+ "HTTP compatibility profile to expose. Available values: "
+ + ", ".join(API_PROFILES)
+ + "."
+ ),
)
serve_parser.add_argument(
"--host",
default="127.0.0.1",
- help="Host interface to bind.",
+ help="Host interface to bind; default: 127.0.0.1.",
)
serve_parser.add_argument(
"--port",
default=7860,
type=int,
- help="TCP port to bind.",
+ help="TCP port to bind; default: 7860.",
)
serve_parser.add_argument(
"--model-residency",
default="cpu-cache",
choices=("staged", "cpu-cache"),
- help="How server mode keeps model components resident.",
+ help=(
+ "Model residency policy: staged reloads each component per "
+ "request; cpu-cache keeps components in system memory."
+ ),
)
return parser
@@ -132,11 +277,27 @@ def inspectModels(args, user_config: UserConfig) -> None:
("VAE", inspectModelSource(model_sources.vae)),
]
output = "\n\n".join(
- formatSummary(name, summary, profile=profile) for name, summary in summaries
+ formatSummary(name, summary, profile=profile)
+ for name, summary in summaries
)
print(output)
+def listModelProfiles(user_config: UserConfig) -> None:
+ """Print configured model profiles for the profile-listing command."""
+
+ print("Available model profiles:")
+ if not user_config.model_profiles:
+ print(" (none configured)")
+ return
+ for name in sorted(user_config.model_profiles):
+ profile = user_config.model_profiles[name]
+ default = " (default)" if name == user_config.default_model else ""
+ print(
+ f" {name}: {profile.architecture}/{profile.variant}{default}"
+ )
+
+
def releaseMemory() -> None:
"""Release Python and CUDA caches between large model stages."""
@@ -192,11 +353,35 @@ def generate(args, user_config: UserConfig) -> list[Path]:
def main(argv: list[str] | None = None) -> int:
"""Run the CLI entry point."""
- parser = buildParser()
+ config_error = None
+ try:
+ user_config = loadUserConfig()
+ except DiffusionCliError as exc:
+ user_config = None
+ config_error = exc
+
+ parser = buildParser(
+ tuple(sorted(user_config.model_profiles))
+ if user_config is not None
+ else None
+ )
args = parser.parse_args(argv)
try:
- user_config = loadUserConfig()
+ if config_error is not None:
+ raise config_error
+
+ if args.model_profile == MODEL_PROFILE_LIST_VALUE:
+ listModelProfiles(user_config)
+ return 0
+ if (
+ args.model_profile is not None
+ and args.model_profile not in user_config.model_profiles
+ ):
+ raise DiffusionCliError(
+ f"Unknown model profile: {args.model_profile}"
+ )
+
if args.command == "serve":
server_config = validateServerConfig(
args.api_profile,
diff --git a/diffusion_cli/krea2_text_encoder.py b/diffusion_cli/krea2_text_encoder.py
index 3f38d2c..56247fc 100644
--- a/diffusion_cli/krea2_text_encoder.py
+++ b/diffusion_cli/krea2_text_encoder.py
@@ -27,6 +27,7 @@ KREA2_PROMPT_PREFIX = (
"<|im_end|>\n<|im_start|>user\n"
)
KREA2_PROMPT_SUFFIX = "<|im_end|>\n<|im_start|>assistant\n"
+KREA2_MAX_TEXT_TOKENS = 512
KREA2_PAD_TOKEN_ID = 151643
QWEN3_VL_4B_VOCAB_SIZE = 151936
QWEN3_VL_4B_HIDDEN_SIZE = 2560
@@ -169,20 +170,6 @@ class Krea2TextEncoder:
if not isinstance(prompt, str):
raise DiffusionCliError("Krea prompt must be a string")
- text = KREA2_PROMPT_PREFIX + prompt + KREA2_PROMPT_SUFFIX
- token_batch = self.tokenizer(
- [text],
- return_tensors="pt",
- padding=True,
- truncation=True,
- max_length=512,
- add_special_tokens=False,
- )
- input_ids = _asTensor(token_batch["input_ids"], dtype=torch.long)
- attention_mask = _asTensor(
- token_batch["attention_mask"],
- dtype=torch.bool,
- )
prefix_batch = self.tokenizer(
KREA2_PROMPT_PREFIX,
return_tensors="pt",
@@ -192,6 +179,43 @@ class Krea2TextEncoder:
)
prefix_ids = _asTensor(prefix_batch["input_ids"], dtype=torch.long)
prefix_ids = prefix_ids.reshape(-1).tolist()
+ suffix_batch = self.tokenizer(
+ KREA2_PROMPT_SUFFIX,
+ return_tensors="pt",
+ padding=False,
+ truncation=False,
+ add_special_tokens=False,
+ )
+ suffix_ids = _asTensor(
+ suffix_batch["input_ids"],
+ dtype=torch.long,
+ )
+ suffix_mask = _asTensor(
+ suffix_batch["attention_mask"],
+ dtype=torch.bool,
+ )
+ prompt_batch = self.tokenizer(
+ [KREA2_PROMPT_PREFIX + prompt],
+ return_tensors="pt",
+ padding="max_length",
+ truncation=True,
+ max_length=(
+ KREA2_MAX_TEXT_TOKENS
+ + len(prefix_ids)
+ - suffix_ids.shape[1]
+ ),
+ add_special_tokens=False,
+ )
+ input_ids = _asTensor(
+ prompt_batch["input_ids"],
+ dtype=torch.long,
+ )
+ attention_mask = _asTensor(
+ prompt_batch["attention_mask"],
+ dtype=torch.bool,
+ )
+ input_ids = torch.cat((input_ids, suffix_ids), dim=1)
+ attention_mask = torch.cat((attention_mask, suffix_mask), dim=1)
boundary = _findTokenBoundary(
input_ids[0].tolist(),
prefix_ids,
diff --git a/tests/test_cli.py b/tests/test_cli.py
index a06ab72..fc49fbc 100644
--- a/tests/test_cli.py
+++ b/tests/test_cli.py
@@ -1,10 +1,78 @@
import unittest
from unittest.mock import patch
-from diffusion_cli.cli import buildParser, main
+from diffusion_cli.cli import MODEL_PROFILE_LIST_VALUE, buildParser, main
class CliTest(unittest.TestCase):
+ def testHelpDescribesOptionsAndConfiguredProfiles(self):
+ help_text = buildParser(("krea2-turbo", "z-image-turbo")).format_help()
+
+ self.assertIn(
+ "configured values: krea2-turbo, z-image-turbo.",
+ help_text,
+ )
+ for option in (
+ "--prompt",
+ "--width",
+ "--height",
+ "--batch-size",
+ "--steps",
+ "--cfg",
+ "--mu",
+ "--shift-y1",
+ "--shift-y2",
+ "--output",
+ "--device",
+ "--dtype",
+ "--inspect-models",
+ ):
+ self.assertIn(option, help_text)
+
+ self.assertIn("Krea Turbo defaults to 8", help_text)
+ self.assertIn("Krea aligns dimensions to a multiple of 16", help_text)
+
+ def testParserAcceptsProfileNamesAndListingForm(self):
+ parser = buildParser(("krea2-turbo", "z-image-turbo"))
+
+ args = parser.parse_args([
+ "--model-profile",
+ "krea2-turbo",
+ "--prompt",
+ "test",
+ ])
+ self.assertEqual(args.model_profile, "krea2-turbo")
+
+ args = parser.parse_args(["--model-profile"])
+ self.assertEqual(args.model_profile, MODEL_PROFILE_LIST_VALUE)
+
+ args = parser.parse_args([
+ "--model-profile",
+ "unknown",
+ "--prompt",
+ "test",
+ ])
+ self.assertEqual(args.model_profile, "unknown")
+
+ def testUnknownModelProfileFailsClearly(self):
+ with patch("sys.stderr"):
+ result = main(["--model-profile", "unknown", "--prompt", "test"])
+
+ self.assertEqual(result, 2)
+
+ def testModelProfileWithoutValueListsProfilesAndExits(self):
+ with patch("sys.stdout") as stdout:
+ result = main(["--model-profile"])
+
+ self.assertEqual(result, 0)
+ output = "".join(
+ call.args[0]
+ for call in stdout.write.call_args_list
+ if call.args
+ )
+ self.assertIn("Available model profiles:", output)
+ self.assertIn("legacy: z-image/turbo", output)
+
def testParserLeavesConfigurableDefaultsUnset(self):
args = buildParser().parse_args(["--prompt", "test"])
diff --git a/tests/test_krea2.py b/tests/test_krea2.py
index 9927ef2..f01260f 100644
--- a/tests/test_krea2.py
+++ b/tests/test_krea2.py
@@ -1,6 +1,7 @@
import tempfile
import unittest
from pathlib import Path
+from types import SimpleNamespace
from unittest.mock import patch
import torch
@@ -23,6 +24,12 @@ from diffusion_cli.krea2_model import (
Krea2Config,
Krea2Model,
)
+from diffusion_cli.krea2_text_encoder import (
+ KREA2_MAX_TEXT_TOKENS,
+ KREA2_PROMPT_PREFIX,
+ KREA2_PROMPT_SUFFIX,
+ Krea2TextEncoder,
+)
from diffusion_cli.krea2_sampling import sampleKrea2, timesteps
from diffusion_cli.quantization import (
QuantizedLayerSpec,
@@ -38,6 +45,65 @@ from diffusion_cli.model_inspect import inspectModelSource
class Krea2Test(unittest.TestCase):
+ def testLongPromptsKeepFixedWindowAndAssistantSuffix(self):
+ prefix_ids = list(range(34))
+ suffix_ids = list(range(900, 905))
+
+ class FakeTokenizer:
+ def __call__(
+ self,
+ text,
+ *,
+ return_tensors,
+ padding,
+ truncation,
+ max_length=None,
+ add_special_tokens,
+ ):
+ del return_tensors, truncation, add_special_tokens
+ value = text[0] if isinstance(text, list) else text
+ if value == KREA2_PROMPT_PREFIX:
+ ids = prefix_ids
+ elif value == KREA2_PROMPT_SUFFIX:
+ ids = suffix_ids
+ else:
+ ids = prefix_ids + list(range(100, 800))
+ if max_length is not None:
+ ids = ids[:max_length]
+ mask = [1] * len(ids)
+ if padding == "max_length" and max_length is not None:
+ mask.extend([0] * (max_length - len(ids)))
+ ids.extend([0] * (max_length - len(ids)))
+ return {
+ "input_ids": torch.tensor([ids]),
+ "attention_mask": torch.tensor([mask]),
+ }
+
+ class FakeModel(nn.Module):
+ def forward(self, input_ids, **kwargs):
+ del kwargs
+ hidden = torch.zeros(
+ input_ids.shape[0],
+ input_ids.shape[1],
+ 2560,
+ )
+ return SimpleNamespace(hidden_states=[hidden] * 36)
+
+ encoder = Krea2TextEncoder(
+ None,
+ Path("unused"),
+ "cpu",
+ model=FakeModel(),
+ tokenizer=FakeTokenizer(),
+ )
+ conditioning = encoder.encodePrompt("long prompt")
+
+ self.assertEqual(
+ conditioning.hidden_states.shape,
+ (1, KREA2_MAX_TEXT_TOKENS, 12, 2560),
+ )
+ self.assertTrue(conditioning.attention_mask[0, -5:].all())
+
def testNamedProfilesCoexistWithLegacyConfiguration(self):
with tempfile.TemporaryDirectory() as temp_dir:
root = Path(temp_dir)