Changes
diff --git a/README.md b/README.md
index 847a02c..dd8a8d3 100644
--- a/README.md
+++ b/README.md
@@ -11,13 +11,43 @@ uv run diffusion-cli --help
uv run diffusion-cli --inspect-models
```
-The default model root is:
+Model paths and stable generation defaults can be configured in:
```text
-/home/mw/documents/ai/diffusion/models
+~/.config/diffusion.toml
```
-Override it with `--model-root` or `DIFFUSION_CLI_MODEL_ROOT`.
+Example:
+
+```toml
+[models]
+diffusion_model = "/models/z-image/diffusion.safetensors"
+text_encoder = "/models/z-image/qwen_3_4b.safetensors"
+vae = "/models/z-image/ae.safetensors"
+tokenizer = "/home/mw/programs/ComfyUI/comfy/text_encoders/qwen25_tokenizer"
+
+[generation]
+negative_prompt = "text, watermark, full-body"
+width = 832
+height = 1248
+batch_size = 1
+steps = 10
+cfg = 1.0
+output = "output.png"
+device = "cuda"
+dtype = "auto"
+```
+
+Command line options override config values for the current run. Model
+paths can also be passed with `--diffusion-model`, `--text-encoder`,
+`--vae`, and `--tokenizer-path`.
+
+```bash
+uv run diffusion-cli --prompt "a ceramic mug" --steps 12
+```
+
+In this example, `--steps 12` overrides any `generation.steps` value in
+the config file for that run only.
Generation is intentionally guarded until the local Z-Image/Lumina2
model port is implemented.
diff --git a/diffusion_cli/cli.py b/diffusion_cli/cli.py
index 4709c55..199f180 100644
--- a/diffusion_cli/cli.py
+++ b/diffusion_cli/cli.py
@@ -8,16 +8,9 @@ from pathlib import Path
import sys
from diffusion_cli.config import (
- DEFAULT_BATCH_SIZE,
- DEFAULT_CFG,
- DEFAULT_DEVICE,
- DEFAULT_DTYPE,
- DEFAULT_HEIGHT,
- DEFAULT_NEGATIVE_PROMPT,
- DEFAULT_OUTPUT,
- DEFAULT_STEPS,
- DEFAULT_WIDTH,
+ UserConfig,
buildGenerationConfig,
+ loadUserConfig,
)
from diffusion_cli.errors import DiffusionCliError
from diffusion_cli.image_io import saveImages
@@ -39,28 +32,38 @@ def buildParser() -> argparse.ArgumentParser:
parser.add_argument("--prompt", help="Positive prompt.")
parser.add_argument(
"--negative-prompt",
- default=DEFAULT_NEGATIVE_PROMPT,
help="Negative prompt.",
)
parser.add_argument("--seed", type=int, help="Noise seed.")
- parser.add_argument("--width", type=int, default=DEFAULT_WIDTH)
- parser.add_argument("--height", type=int, default=DEFAULT_HEIGHT)
- parser.add_argument("--batch-size", type=int, default=DEFAULT_BATCH_SIZE)
- parser.add_argument("--steps", type=int, default=DEFAULT_STEPS)
- parser.add_argument("--cfg", type=float, default=DEFAULT_CFG)
- parser.add_argument("--output", type=Path, default=DEFAULT_OUTPUT)
- parser.add_argument("--device", default=DEFAULT_DEVICE)
- parser.add_argument("--model-root", type=Path)
+ 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("--output", type=Path)
+ parser.add_argument("--device")
+ parser.add_argument(
+ "--diffusion-model",
+ type=Path,
+ help="Local Z-Image diffusion model safetensors file.",
+ )
+ parser.add_argument(
+ "--text-encoder",
+ type=Path,
+ help="Local Qwen text encoder safetensors file.",
+ )
+ parser.add_argument(
+ "--vae",
+ type=Path,
+ help="Local VAE safetensors file.",
+ )
parser.add_argument(
"--tokenizer-path",
type=Path,
- default=Path("/home/mw/programs/ComfyUI/comfy/text_encoders")
- / "qwen25_tokenizer",
help="Local Qwen tokenizer directory.",
)
parser.add_argument(
"--dtype",
- default=DEFAULT_DTYPE,
choices=("auto", "bf16", "fp32"),
)
parser.add_argument(
@@ -71,10 +74,10 @@ def buildParser() -> argparse.ArgumentParser:
return parser
-def inspectModels(model_root: Path | None) -> None:
+def inspectModels(args, user_config: UserConfig) -> None:
"""Print metadata summaries for the required model files."""
- model_files = resolveModelFiles(model_root)
+ model_files = resolveModelFiles(args, user_config)
summaries = [
("diffusion model", inspectSafetensors(model_files.diffusion_model)),
("text encoder", inspectSafetensors(model_files.text_encoder)),
@@ -99,11 +102,11 @@ def releaseMemory() -> None:
pass
-def generate(args) -> list[Path]:
+def generate(args, user_config: UserConfig) -> list[Path]:
"""Validate generation arguments and run the current milestone."""
- config = buildGenerationConfig(args)
- model_files = resolveModelFiles(args.model_root)
+ config = buildGenerationConfig(args, user_config)
+ model_files = resolveModelFiles(args, user_config)
text_encoder = ZImageTextEncoder(
model_files.text_encoder,
@@ -150,11 +153,12 @@ def main(argv: list[str] | None = None) -> int:
args = parser.parse_args(argv)
try:
+ user_config = loadUserConfig()
if args.inspect_models:
- inspectModels(args.model_root)
+ inspectModels(args, user_config)
return 0
- paths = generate(args)
+ paths = generate(args, user_config)
except DiffusionCliError as exc:
print(f"error: {exc}", file=sys.stderr)
return 2
diff --git a/diffusion_cli/config.py b/diffusion_cli/config.py
index 0effebe..37cfa30 100644
--- a/diffusion_cli/config.py
+++ b/diffusion_cli/config.py
@@ -5,9 +5,12 @@ from __future__ import annotations
from dataclasses import dataclass
from pathlib import Path
from secrets import randbits
+import tomllib
+from typing import Any
from diffusion_cli.errors import DiffusionCliError
+DEFAULT_CONFIG_PATH = Path("~/.config/diffusion.toml")
DEFAULT_NEGATIVE_PROMPT = "text, watermark, full-body"
DEFAULT_WIDTH = 832
DEFAULT_HEIGHT = 1248
@@ -22,6 +25,19 @@ DEFAULT_MULTIPLIER = 1.0
LATENT_CHANNELS = 16
LATENT_DOWNSCALE = 8
TOKENIZER_FILES = ("vocab.json", "merges.txt", "tokenizer_config.json")
+TOP_LEVEL_CONFIG_KEYS = {"models", "generation"}
+MODEL_CONFIG_KEYS = {"diffusion_model", "text_encoder", "vae", "tokenizer"}
+GENERATION_CONFIG_KEYS = {
+ "negative_prompt",
+ "width",
+ "height",
+ "batch_size",
+ "steps",
+ "cfg",
+ "output",
+ "device",
+ "dtype",
+}
@dataclass(frozen=True)
@@ -33,6 +49,39 @@ class ModelFiles:
vae: Path
+@dataclass(frozen=True)
+class ModelPathConfig:
+ """Optional model path defaults loaded from user configuration."""
+
+ diffusion_model: Path | None = None
+ text_encoder: Path | None = None
+ vae: Path | None = None
+ tokenizer: Path | None = None
+
+
+@dataclass(frozen=True)
+class GenerationDefaults:
+ """Optional generation defaults loaded from user configuration."""
+
+ negative_prompt: str | None = None
+ width: int | None = None
+ height: int | None = None
+ batch_size: int | None = None
+ steps: int | None = None
+ cfg: float | None = None
+ output: Path | None = None
+ device: str | None = None
+ dtype: str | None = None
+
+
+@dataclass(frozen=True)
+class UserConfig:
+ """User-provided model paths and generation defaults."""
+
+ models: ModelPathConfig
+ generation: GenerationDefaults
+
+
@dataclass(frozen=True)
class GenerationConfig:
"""Validated user intent for one text-to-image generation request."""
@@ -57,6 +106,153 @@ def randomSeed() -> int:
return randbits(64)
+def configPath() -> Path:
+ """Return the fixed user config path."""
+
+ return DEFAULT_CONFIG_PATH.expanduser()
+
+
+def _optionalTable(data: dict[str, Any], name: str) -> dict[str, Any]:
+ value = data.get(name, {})
+ if not isinstance(value, dict):
+ raise DiffusionCliError(f"Config table must be a table: [{name}]")
+ return value
+
+
+def _rejectUnknownKeys(
+ data: dict[str, Any],
+ allowed_keys: set[str],
+ label: str,
+) -> None:
+ unknown_keys = sorted(set(data) - allowed_keys)
+ if unknown_keys:
+ unknown = unknown_keys[0]
+ raise DiffusionCliError(f"Unknown config key {label}.{unknown}")
+
+
+def _optionalPath(
+ table: dict[str, Any],
+ table_name: str,
+ key: str,
+) -> Path | None:
+ value = table.get(key)
+ if value is None:
+ return None
+ if not isinstance(value, str):
+ raise DiffusionCliError(
+ f"Config value {table_name}.{key} must be a string path"
+ )
+ return Path(value).expanduser().resolve()
+
+
+def _optionalString(
+ table: dict[str, Any],
+ table_name: str,
+ key: str,
+) -> str | None:
+ value = table.get(key)
+ if value is None:
+ return None
+ if not isinstance(value, str):
+ raise DiffusionCliError(
+ f"Config value {table_name}.{key} must be a string"
+ )
+ return value
+
+
+def _optionalInt(
+ table: dict[str, Any],
+ table_name: str,
+ key: str,
+) -> int | None:
+ value = table.get(key)
+ if value is None:
+ return None
+ if not isinstance(value, int) or isinstance(value, bool):
+ raise DiffusionCliError(
+ f"Config value {table_name}.{key} must be an integer"
+ )
+ return value
+
+
+def _optionalFloat(
+ table: dict[str, Any],
+ table_name: str,
+ key: str,
+) -> float | None:
+ value = table.get(key)
+ if value is None:
+ return None
+ if not isinstance(value, int | float) or isinstance(value, bool):
+ raise DiffusionCliError(
+ f"Config value {table_name}.{key} must be a number"
+ )
+ return float(value)
+
+
+def loadUserConfig(path: Path | None = None) -> UserConfig:
+ """Load optional defaults from the fixed TOML config file."""
+
+ config_file = configPath() if path is None else path.expanduser()
+ if not config_file.exists():
+ return UserConfig(ModelPathConfig(), GenerationDefaults())
+
+ try:
+ with config_file.open("rb") as file:
+ data = tomllib.load(file)
+ except tomllib.TOMLDecodeError as exc:
+ raise DiffusionCliError(
+ f"Invalid TOML config {config_file}: {exc}"
+ ) from exc
+ except OSError as exc:
+ raise DiffusionCliError(
+ f"Could not read config: {config_file}"
+ ) from exc
+
+ _rejectUnknownKeys(data, TOP_LEVEL_CONFIG_KEYS, "top-level")
+ models = _optionalTable(data, "models")
+ generation = _optionalTable(data, "generation")
+ _rejectUnknownKeys(models, MODEL_CONFIG_KEYS, "models")
+ _rejectUnknownKeys(generation, GENERATION_CONFIG_KEYS, "generation")
+
+ return UserConfig(
+ models=ModelPathConfig(
+ diffusion_model=_optionalPath(
+ models,
+ "models",
+ "diffusion_model",
+ ),
+ text_encoder=_optionalPath(models, "models", "text_encoder"),
+ vae=_optionalPath(models, "models", "vae"),
+ tokenizer=_optionalPath(models, "models", "tokenizer"),
+ ),
+ generation=GenerationDefaults(
+ negative_prompt=_optionalString(
+ generation,
+ "generation",
+ "negative_prompt",
+ ),
+ width=_optionalInt(generation, "generation", "width"),
+ height=_optionalInt(generation, "generation", "height"),
+ batch_size=_optionalInt(generation, "generation", "batch_size"),
+ steps=_optionalInt(generation, "generation", "steps"),
+ cfg=_optionalFloat(generation, "generation", "cfg"),
+ output=_optionalPath(generation, "generation", "output"),
+ device=_optionalString(generation, "generation", "device"),
+ dtype=_optionalString(generation, "generation", "dtype"),
+ ),
+ )
+
+
+def coalesce(*values):
+ """Return the first value that is not None."""
+
+ for value in values:
+ if value is not None:
+ return value
+ raise AssertionError("coalesce requires at least one non-None value")
+
+
def validateDimensions(width: int, height: int) -> None:
"""Validate that image dimensions are positive latent multiples."""
@@ -148,36 +344,69 @@ def selectDtype(dtype_name: str, device) -> object:
raise AssertionError("unreachable dtype branch")
-def buildGenerationConfig(args) -> GenerationConfig:
+def buildGenerationConfig(
+ args,
+ user_config: UserConfig | None = None,
+) -> GenerationConfig:
"""Validate parsed CLI arguments and build a generation config."""
+ if user_config is None:
+ user_config = loadUserConfig()
+ generation = user_config.generation
+ models = user_config.models
+
if not args.prompt:
raise DiffusionCliError("--prompt is required for generation")
- validateDimensions(args.width, args.height)
- if args.batch_size < 1:
+
+ negative_prompt = coalesce(
+ args.negative_prompt,
+ generation.negative_prompt,
+ DEFAULT_NEGATIVE_PROMPT,
+ )
+ width = coalesce(args.width, generation.width, DEFAULT_WIDTH)
+ height = coalesce(args.height, generation.height, DEFAULT_HEIGHT)
+ batch_size = coalesce(
+ args.batch_size,
+ generation.batch_size,
+ DEFAULT_BATCH_SIZE,
+ )
+ steps = coalesce(args.steps, generation.steps, DEFAULT_STEPS)
+ cfg = coalesce(args.cfg, generation.cfg, DEFAULT_CFG)
+ device_name = coalesce(args.device, generation.device, DEFAULT_DEVICE)
+ dtype_name = coalesce(args.dtype, generation.dtype, DEFAULT_DTYPE)
+ output_path = coalesce(args.output, generation.output, DEFAULT_OUTPUT)
+ tokenizer_path = args.tokenizer_path
+ if tokenizer_path is None:
+ tokenizer_path = models.tokenizer
+
+ if tokenizer_path is None:
+ raise DiffusionCliError("Missing models.tokenizer")
+
+ validateDimensions(width, height)
+ if batch_size < 1:
raise DiffusionCliError(
- f"Batch size must be at least 1: got {args.batch_size}"
+ f"Batch size must be at least 1: got {batch_size}"
)
- if args.steps < 1:
- raise DiffusionCliError(f"Steps must be at least 1: got {args.steps}")
- if args.cfg < 0:
- raise DiffusionCliError(f"CFG must be non-negative: got {args.cfg}")
-
- device = selectDevice(args.device)
- dtype = selectDtype(args.dtype, device)
- output = validateOutputPath(args.output)
- tokenizer_path = validateTokenizerPath(args.tokenizer_path)
+ if steps < 1:
+ raise DiffusionCliError(f"Steps must be at least 1: got {steps}")
+ if cfg < 0:
+ raise DiffusionCliError(f"CFG must be non-negative: got {cfg}")
+
+ device = selectDevice(device_name)
+ dtype = selectDtype(dtype_name, device)
+ output = validateOutputPath(output_path)
+ tokenizer_path = validateTokenizerPath(tokenizer_path)
seed = args.seed if args.seed is not None else randomSeed()
return GenerationConfig(
prompt=args.prompt,
- negative_prompt=args.negative_prompt,
+ negative_prompt=negative_prompt,
seed=seed,
- width=args.width,
- height=args.height,
- batch_size=args.batch_size,
- steps=args.steps,
- cfg=args.cfg,
+ width=width,
+ height=height,
+ batch_size=batch_size,
+ steps=steps,
+ cfg=cfg,
device=device,
dtype=dtype,
output=output,
diff --git a/diffusion_cli/paths.py b/diffusion_cli/paths.py
index e27b87a..3b3b9f3 100644
--- a/diffusion_cli/paths.py
+++ b/diffusion_cli/paths.py
@@ -2,56 +2,39 @@
from __future__ import annotations
-import os
from pathlib import Path
-from diffusion_cli.config import ModelFiles
+from diffusion_cli.config import ModelFiles, UserConfig
from diffusion_cli.errors import DiffusionCliError
-MODEL_ROOT_ENV = "DIFFUSION_CLI_MODEL_ROOT"
-DEFAULT_MODEL_ROOT = Path("/home/mw/documents/ai/diffusion/models")
-DIFFUSION_MODEL_RELATIVE = Path(
- "diffusion_models/z-image_turbo/moodyPornMix_zitV3.safetensors"
-)
-TEXT_ENCODER_RELATIVE = Path(
- "text_encoders/z-image_turbo/qwen_3_4b.safetensors"
-)
-VAE_RELATIVE = Path("vae/z-image_turbo/ae.safetensors")
+def requireFile(path: Path | None, role: str, config_key: str) -> Path:
+ """Return an existing regular file or raise a direct CLI error."""
-def resolveModelRoot(model_root: Path | None) -> Path:
- """Resolve the model root from CLI, environment, or default."""
-
- if model_root is not None:
- return model_root.expanduser().resolve()
-
- env_root = os.environ.get(MODEL_ROOT_ENV)
- if env_root:
- return Path(env_root).expanduser().resolve()
-
- return DEFAULT_MODEL_ROOT
-
+ if path is None:
+ raise DiffusionCliError(f"Missing {config_key}")
-def requireFile(path: Path, role: str) -> Path:
- """Return an existing regular file or raise a direct CLI error."""
+ resolved_path = path.expanduser().resolve()
+ if not resolved_path.is_file():
+ raise DiffusionCliError(f"Missing {role}: {resolved_path}")
+ return resolved_path
- if not path.is_file():
- raise DiffusionCliError(f"Missing {role}: {path}")
- return path
+def resolveModelFiles(args, user_config: UserConfig) -> ModelFiles:
+ """Resolve and validate explicit model paths for the active workflow."""
-def resolveModelFiles(model_root: Path | None = None) -> ModelFiles:
- """Resolve and validate the three model files for the active workflow."""
+ models = user_config.models
- root = resolveModelRoot(model_root)
return ModelFiles(
diffusion_model=requireFile(
- root / DIFFUSION_MODEL_RELATIVE,
+ args.diffusion_model or models.diffusion_model,
"diffusion model",
+ "models.diffusion_model",
),
text_encoder=requireFile(
- root / TEXT_ENCODER_RELATIVE,
+ args.text_encoder or models.text_encoder,
"text encoder",
+ "models.text_encoder",
),
- vae=requireFile(root / VAE_RELATIVE, "VAE"),
+ vae=requireFile(args.vae or models.vae, "VAE", "models.vae"),
)
diff --git a/tests/test_cli.py b/tests/test_cli.py
index 1806c87..32e8a22 100644
--- a/tests/test_cli.py
+++ b/tests/test_cli.py
@@ -1,22 +1,53 @@
import unittest
+from unittest.mock import patch
from diffusion_cli.cli import buildParser, main
-from diffusion_cli.config import DEFAULT_HEIGHT, DEFAULT_WIDTH
class CliTest(unittest.TestCase):
- def testParserDefaults(self):
+ def testParserLeavesConfigurableDefaultsUnset(self):
args = buildParser().parse_args(["--prompt", "test"])
self.assertEqual(args.prompt, "test")
- self.assertEqual(args.width, DEFAULT_WIDTH)
- self.assertEqual(args.height, DEFAULT_HEIGHT)
- self.assertEqual(args.batch_size, 1)
- self.assertEqual(args.steps, 10)
- self.assertEqual(args.cfg, 1.0)
+ self.assertIsNone(args.width)
+ self.assertIsNone(args.height)
+ self.assertIsNone(args.batch_size)
+ self.assertIsNone(args.steps)
+ self.assertIsNone(args.cfg)
+ self.assertIsNone(args.output)
+ self.assertIsNone(args.device)
+ self.assertIsNone(args.diffusion_model)
+ self.assertIsNone(args.text_encoder)
+ self.assertIsNone(args.vae)
+ self.assertIsNone(args.tokenizer_path)
+ self.assertIsNone(args.dtype)
+
+ def testParserAcceptsExplicitConfigurableValues(self):
+ args = buildParser().parse_args(
+ [
+ "--prompt",
+ "test",
+ "--steps",
+ "12",
+ "--cfg",
+ "1.5",
+ "--dtype",
+ "bf16",
+ ]
+ )
+
+ self.assertEqual(args.steps, 12)
+ self.assertEqual(args.cfg, 1.5)
+ self.assertEqual(args.dtype, "bf16")
+
+ def testModelRootIsNoLongerAccepted(self):
+ with patch("sys.stderr"):
+ with self.assertRaises(SystemExit):
+ buildParser().parse_args(["--model-root", "/tmp/models"])
def testMissingPromptFailsForGeneration(self):
- self.assertEqual(main([]), 2)
+ with patch("sys.stderr"):
+ self.assertEqual(main([]), 2)
if __name__ == "__main__":
diff --git a/tests/test_config.py b/tests/test_config.py
index 4893393..278bb6b 100644
--- a/tests/test_config.py
+++ b/tests/test_config.py
@@ -1,11 +1,18 @@
from pathlib import Path
import tempfile
import unittest
+from types import SimpleNamespace
+from unittest.mock import patch
from diffusion_cli.config import (
DEFAULT_HEIGHT,
DEFAULT_WIDTH,
+ GenerationDefaults,
+ ModelPathConfig,
+ UserConfig,
TOKENIZER_FILES,
+ buildGenerationConfig,
+ loadUserConfig,
validateDimensions,
validateOutputPath,
validateTokenizerPath,
@@ -39,6 +46,213 @@ class ConfigTest(unittest.TestCase):
self.assertEqual(result, output)
+ def testMissingUserConfigReturnsEmptyConfig(self):
+ with tempfile.TemporaryDirectory() as temp_dir:
+ config = loadUserConfig(Path(temp_dir) / "missing.toml")
+
+ self.assertIsNone(config.models.diffusion_model)
+ self.assertIsNone(config.generation.steps)
+
+ def testLoadUserConfig(self):
+ with tempfile.TemporaryDirectory() as temp_dir:
+ root = Path(temp_dir)
+ config_path = root / "diffusion.toml"
+ config_path.write_text(
+ "\n".join(
+ (
+ "[models]",
+ f'diffusion_model = "{root / "diffusion.safetensors"}"',
+ f'text_encoder = "{root / "text_encoder.safetensors"}"',
+ f'vae = "{root / "ae.safetensors"}"',
+ f'tokenizer = "{root / "tokenizer"}"',
+ "",
+ "[generation]",
+ 'negative_prompt = "low quality"',
+ "width = 512",
+ "height = 768",
+ "batch_size = 2",
+ "steps = 12",
+ "cfg = 1.5",
+ f'output = "{root / "output.png"}"',
+ 'device = "cuda:0"',
+ 'dtype = "bf16"',
+ )
+ ),
+ encoding="utf-8",
+ )
+
+ config = loadUserConfig(config_path)
+
+ self.assertEqual(
+ config.models.diffusion_model.name,
+ "diffusion.safetensors",
+ )
+ self.assertEqual(config.models.tokenizer.name, "tokenizer")
+ self.assertEqual(config.generation.width, 512)
+ self.assertEqual(config.generation.height, 768)
+ self.assertEqual(config.generation.cfg, 1.5)
+ self.assertEqual(config.generation.dtype, "bf16")
+
+ def testMalformedUserConfigFailsClearly(self):
+ with tempfile.TemporaryDirectory() as temp_dir:
+ config_path = Path(temp_dir) / "diffusion.toml"
+ config_path.write_text("[generation\n", encoding="utf-8")
+
+ with self.assertRaises(DiffusionCliError) as context:
+ loadUserConfig(config_path)
+
+ self.assertIn("Invalid TOML config", str(context.exception))
+ self.assertIn("diffusion.toml", str(context.exception))
+
+ def testUnknownTopLevelConfigTableFails(self):
+ with tempfile.TemporaryDirectory() as temp_dir:
+ config_path = Path(temp_dir) / "diffusion.toml"
+ config_path.write_text("[model]\n", encoding="utf-8")
+
+ with self.assertRaises(DiffusionCliError) as context:
+ loadUserConfig(config_path)
+
+ self.assertIn(
+ "Unknown config key top-level.model",
+ str(context.exception),
+ )
+
+ def testUnknownModelConfigKeyFails(self):
+ with tempfile.TemporaryDirectory() as temp_dir:
+ config_path = Path(temp_dir) / "diffusion.toml"
+ config_path.write_text(
+ "[models]\ndiffusion = '/tmp/model.safetensors'\n",
+ encoding="utf-8",
+ )
+
+ with self.assertRaises(DiffusionCliError) as context:
+ loadUserConfig(config_path)
+
+ self.assertIn(
+ "Unknown config key models.diffusion",
+ str(context.exception),
+ )
+
+ def testUnknownGenerationConfigKeyFails(self):
+ with tempfile.TemporaryDirectory() as temp_dir:
+ config_path = Path(temp_dir) / "diffusion.toml"
+ config_path.write_text(
+ "[generation]\nstep = 10\n",
+ encoding="utf-8",
+ )
+
+ with self.assertRaises(DiffusionCliError) as context:
+ loadUserConfig(config_path)
+
+ self.assertIn(
+ "Unknown config key generation.step",
+ str(context.exception),
+ )
+
+ def testWrongConfigScalarTypeFails(self):
+ with tempfile.TemporaryDirectory() as temp_dir:
+ config_path = Path(temp_dir) / "diffusion.toml"
+ config_path.write_text(
+ '[generation]\nsteps = "10"\n',
+ encoding="utf-8",
+ )
+
+ with self.assertRaises(DiffusionCliError) as context:
+ loadUserConfig(config_path)
+
+ self.assertIn(
+ "generation.steps must be an integer",
+ str(context.exception),
+ )
+
+ def testGenerationConfigUsesConfigDefaults(self):
+ with tempfile.TemporaryDirectory() as temp_dir:
+ root = Path(temp_dir)
+ tokenizer = root / "tokenizer"
+ tokenizer.mkdir()
+ for file_name in TOKENIZER_FILES:
+ (tokenizer / file_name).write_text("{}", encoding="utf-8")
+
+ user_config = UserConfig(
+ models=ModelPathConfig(tokenizer=tokenizer),
+ generation=GenerationDefaults(
+ negative_prompt="low quality",
+ width=512,
+ height=768,
+ batch_size=2,
+ steps=12,
+ cfg=1.5,
+ output=root / "image.png",
+ device="cuda:0",
+ dtype="bf16",
+ ),
+ )
+ args = SimpleNamespace(
+ prompt="a mug",
+ negative_prompt=None,
+ seed=123,
+ width=None,
+ height=None,
+ batch_size=None,
+ steps=None,
+ cfg=None,
+ output=None,
+ device=None,
+ dtype=None,
+ tokenizer_path=None,
+ )
+
+ with patch("diffusion_cli.config.selectDevice") as select_device:
+ with patch("diffusion_cli.config.selectDtype") as select_dtype:
+ select_device.return_value = "cuda:0"
+ select_dtype.return_value = "bf16"
+ config = buildGenerationConfig(args, user_config)
+
+ self.assertEqual(config.negative_prompt, "low quality")
+ self.assertEqual(config.width, 512)
+ self.assertEqual(config.height, 768)
+ self.assertEqual(config.batch_size, 2)
+ self.assertEqual(config.steps, 12)
+ self.assertEqual(config.cfg, 1.5)
+ self.assertEqual(config.seed, 123)
+
+ def testGenerationConfigCliOverridesConfig(self):
+ with tempfile.TemporaryDirectory() as temp_dir:
+ root = Path(temp_dir)
+ tokenizer = root / "tokenizer"
+ tokenizer.mkdir()
+ for file_name in TOKENIZER_FILES:
+ (tokenizer / file_name).write_text("{}", encoding="utf-8")
+
+ user_config = UserConfig(
+ models=ModelPathConfig(tokenizer=tokenizer),
+ generation=GenerationDefaults(width=512, steps=12),
+ )
+ args = SimpleNamespace(
+ prompt="a mug",
+ negative_prompt=None,
+ seed=123,
+ width=640,
+ height=None,
+ batch_size=None,
+ steps=4,
+ cfg=None,
+ output=root / "image.png",
+ device="cuda",
+ dtype="fp32",
+ tokenizer_path=None,
+ )
+
+ with patch("diffusion_cli.config.selectDevice") as select_device:
+ with patch("diffusion_cli.config.selectDtype") as select_dtype:
+ select_device.return_value = "cuda"
+ select_dtype.return_value = "fp32"
+ config = buildGenerationConfig(args, user_config)
+
+ self.assertEqual(config.width, 640)
+ self.assertEqual(config.height, DEFAULT_HEIGHT)
+ self.assertEqual(config.steps, 4)
+
if __name__ == "__main__":
unittest.main()
diff --git a/tests/test_paths.py b/tests/test_paths.py
index 65c44f0..5c31090 100644
--- a/tests/test_paths.py
+++ b/tests/test_paths.py
@@ -1,44 +1,60 @@
from pathlib import Path
import tempfile
import unittest
+from types import SimpleNamespace
from diffusion_cli.errors import DiffusionCliError
-from diffusion_cli.paths import (
- DIFFUSION_MODEL_RELATIVE,
- TEXT_ENCODER_RELATIVE,
- VAE_RELATIVE,
- resolveModelFiles,
-)
+from diffusion_cli.config import ModelPathConfig, UserConfig, GenerationDefaults
+from diffusion_cli.paths import resolveModelFiles
class PathsTest(unittest.TestCase):
def testResolveModelFiles(self):
with tempfile.TemporaryDirectory() as temp_dir:
root = Path(temp_dir)
- for relative in (
- DIFFUSION_MODEL_RELATIVE,
- TEXT_ENCODER_RELATIVE,
- VAE_RELATIVE,
- ):
- path = root / relative
- path.parent.mkdir(parents=True, exist_ok=True)
+ diffusion_model = root / "diffusion.safetensors"
+ text_encoder = root / "text_encoder.safetensors"
+ vae = root / "ae.safetensors"
+ for path in (diffusion_model, text_encoder, vae):
path.write_bytes(b"placeholder")
- files = resolveModelFiles(root)
+ args = SimpleNamespace(
+ diffusion_model=None,
+ text_encoder=None,
+ vae=None,
+ )
+ user_config = UserConfig(
+ models=ModelPathConfig(
+ diffusion_model=diffusion_model,
+ text_encoder=text_encoder,
+ vae=vae,
+ ),
+ generation=GenerationDefaults(),
+ )
+ files = resolveModelFiles(args, user_config)
self.assertEqual(
files.diffusion_model.name,
- "moodyPornMix_zitV3.safetensors",
+ "diffusion.safetensors",
)
- self.assertEqual(files.text_encoder.name, "qwen_3_4b.safetensors")
+ self.assertEqual(files.text_encoder.name, "text_encoder.safetensors")
self.assertEqual(files.vae.name, "ae.safetensors")
def testMissingFileFailsClearly(self):
- with tempfile.TemporaryDirectory() as temp_dir:
- with self.assertRaises(DiffusionCliError) as context:
- resolveModelFiles(Path(temp_dir))
+ args = SimpleNamespace(
+ diffusion_model=None,
+ text_encoder=None,
+ vae=None,
+ )
+ user_config = UserConfig(
+ models=ModelPathConfig(),
+ generation=GenerationDefaults(),
+ )
+
+ with self.assertRaises(DiffusionCliError) as context:
+ resolveModelFiles(args, user_config)
- self.assertIn("Missing diffusion model", str(context.exception))
+ self.assertIn("Missing models.diffusion_model", str(context.exception))
if __name__ == "__main__":