"""Immutable configuration objects and validation helpers."""
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
DEFAULT_BATCH_SIZE = 1
DEFAULT_STEPS = 10
DEFAULT_CFG = 1.0
DEFAULT_OUTPUT = Path("output.png")
DEFAULT_OUTPUT_EXTENSION = "png"
DEFAULT_OUTPUT_QUALITY = 95
DEFAULT_DEVICE = "cuda"
DEFAULT_DTYPE = "auto"
DEFAULT_SHIFT = 3.0
DEFAULT_MULTIPLIER = 1.0
OUTPUT_EXTENSION_ALIASES = {
"png": "png",
"jpg": "jpg",
"jpeg": "jpg",
"webp": "webp",
"avif": "avif",
}
OUTPUT_MIME_TYPES = {
"png": "image/png",
"jpg": "image/jpeg",
"webp": "image/webp",
"avif": "image/avif",
}
OUTPUT_FORMAT_OPTIONS = (
{"value": "png", "label": "PNG"},
{"value": "jpg", "label": "JPEG"},
{"value": "webp", "label": "WebP"},
{"value": "avif", "label": "AVIF"},
)
LATENT_CHANNELS = 16
LATENT_DOWNSCALE = 8
TOKENIZER_FILES = ("vocab.json", "merges.txt", "tokenizer_config.json")
# UI-visible Z Image Turbo resolution presets, grouped by base resolution.
UI_RESOLUTION_PRESETS = (
(1024, 1024, "1024", "1:1"),
(1152, 896, "1024", "9:7"),
(896, 1152, "1024", "7:9"),
(1152, 864, "1024", "4:3"),
(864, 1152, "1024", "3:4"),
(1248, 832, "1024", "3:2"),
(832, 1248, "1024", "2:3"),
(1280, 720, "1024", "16:9"),
(720, 1280, "1024", "9:16"),
(1344, 576, "1024", "21:9"),
(576, 1344, "1024", "9:21"),
(1280, 1280, "1280", "1:1"),
(1440, 1120, "1280", "9:7"),
(1120, 1440, "1280", "7:9"),
(1472, 1104, "1280", "4:3"),
(1104, 1472, "1280", "3:4"),
(1536, 1024, "1280", "3:2"),
(1024, 1536, "1280", "2:3"),
(1536, 864, "1280", "16:9"),
(864, 1536, "1280", "9:16"),
(1680, 720, "1280", "21:9"),
(720, 1680, "1280", "9:21"),
(1536, 1536, "1536", "1:1"),
(1728, 1344, "1536", "9:7"),
(1344, 1728, "1536", "7:9"),
(1728, 1296, "1536", "4:3"),
(1296, 1728, "1536", "3:4"),
(1872, 1248, "1536", "3:2"),
(1248, 1872, "1536", "2:3"),
(2048, 1152, "1536", "16:9"),
(1152, 2048, "1536", "9:16"),
(2016, 864, "1536", "21:9"),
(864, 2016, "1536", "9:21"),
)
TOP_LEVEL_CONFIG_KEYS = {"models", "generation"}
MODEL_CONFIG_KEYS = {
"checkpoint",
"diffusion_model",
"text_encoder",
"vae",
"tokenizer",
}
GENERATION_CONFIG_KEYS = {
"negative_prompt",
"width",
"height",
"batch_size",
"steps",
"cfg",
"output",
"output_extension",
"output_quality",
"device",
"dtype",
}
DIFFUSION_ROLE = "diffusion_model"
TEXT_ENCODER_ROLE = "text_encoder"
VAE_ROLE = "vae"
CHECKPOINT_DIFFUSION_PREFIX = "model.diffusion_model."
CHECKPOINT_TEXT_ENCODER_PREFIX = "text_encoders.qwen3_4b.transformer."
CHECKPOINT_VAE_PREFIX = "vae."
@dataclass(frozen=True)
class ModelFiles:
"""Absolute paths to the model files needed for generation."""
diffusion_model: Path
text_encoder: Path
vae: Path
@dataclass(frozen=True)
class ModelSource:
"""A local safetensors source for one model component."""
path: Path
role: str
checkpoint_prefix: str | None = None
@dataclass(frozen=True)
class ModelSources:
"""Resolved local sources for all model components."""
diffusion_model: ModelSource
text_encoder: ModelSource
vae: ModelSource
@dataclass(frozen=True)
class ModelPathConfig:
"""Optional model path defaults loaded from user configuration."""
checkpoint: Path | None = None
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
output_extension: str | None = None
output_quality: int | 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."""
prompt: str
negative_prompt: str
seed: int
width: int
height: int
batch_size: int
steps: int
cfg: float
device: object
dtype: object
dtype_name: str
output: Path
output_extension: str
output_quality: int
output_mime_type: str
tokenizer_path: Path
@dataclass(frozen=True)
class ImageGenerationRequest:
"""Validated user intent for one text-to-image request."""
prompt: str
negative_prompt: str | None = None
seed: int | 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
output_extension: str | None = None
output_quality: int | None = None
device: str | None = None
dtype: str | None = None
tokenizer_path: Path | None = None
def randomSeed() -> int:
"""Return a random 64-bit seed suitable for torch generators."""
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(
checkpoint=_optionalPath(models, "models", "checkpoint"),
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"),
output_extension=_optionalString(
generation,
"generation",
"output_extension",
),
output_quality=_optionalInt(
generation,
"generation",
"output_quality",
),
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."""
if width <= 0 or width % LATENT_DOWNSCALE != 0:
raise DiffusionCliError(
f"Width must be a positive multiple of 8: got {width}"
)
if height <= 0 or height % LATENT_DOWNSCALE != 0:
raise DiffusionCliError(
f"Height must be a positive multiple of 8: got {height}"
)
def validateTokenizerPath(tokenizer_path: Path) -> Path:
"""Validate the local tokenizer directory expected by transformers."""
path = tokenizer_path.expanduser().resolve()
if not path.is_dir():
raise DiffusionCliError(f"Tokenizer path is not a directory: {path}")
for file_name in TOKENIZER_FILES:
token_file = path / file_name
if not token_file.is_file():
raise DiffusionCliError(f"Missing tokenizer file: {token_file}")
return path
def validateOutputPath(output: Path) -> Path:
"""Create the output parent directory when possible."""
path = output.expanduser()
parent = path.parent if path.parent != Path("") else Path(".")
try:
parent.mkdir(parents=True, exist_ok=True)
except OSError as exc:
raise DiffusionCliError(
f"Could not create output directory: {parent}"
) from exc
if not parent.is_dir():
raise DiffusionCliError(f"Output parent is not a directory: {parent}")
return path
def normalizeOutputExtension(value: str) -> str:
"""Return a supported canonical output extension."""
if not isinstance(value, str):
raise DiffusionCliError("Output extension must be a string")
normalized = value.strip()
if normalized.startswith("."):
normalized = normalized[1:]
normalized = normalized.lower()
extension = OUTPUT_EXTENSION_ALIASES.get(normalized)
if extension is None:
accepted = ", ".join(OUTPUT_MIME_TYPES)
raise DiffusionCliError(
f"Output extension must be one of {accepted}: got {normalized}"
)
return extension
def validateOutputQuality(value: int) -> int:
"""Return a valid ImageMagick output quality value."""
if not isinstance(value, int) or isinstance(value, bool):
raise DiffusionCliError("Output quality must be an integer")
if value < 1 or value > 100:
raise DiffusionCliError(
f"Output quality must be between 1 and 100: got {value}"
)
return value
def outputMimeType(extension: str) -> str:
"""Return the MIME type for a normalized output extension."""
return OUTPUT_MIME_TYPES[normalizeOutputExtension(extension)]
def selectDevice(device_name: str):
"""Validate and return the requested CUDA device."""
import torch
if not device_name.startswith("cuda"):
raise DiffusionCliError(
"Only CUDA devices are supported for this milestone"
)
if not torch.cuda.is_available():
raise DiffusionCliError(
"CUDA requested but torch.cuda.is_available() is false"
)
device = torch.device(device_name)
try:
torch.cuda.get_device_properties(device)
except Exception as exc:
raise DiffusionCliError(f"Invalid CUDA device: {device_name}") from exc
return device
def selectDtype(dtype_name: str, device) -> object:
"""Validate and return the requested inference dtype."""
import torch
if dtype_name not in {"auto", "bf16", "fp16", "fp32"}:
raise DiffusionCliError(
f"dtype must be one of auto, bf16, fp16, fp32: got {dtype_name}"
)
if dtype_name == "fp32":
return torch.float32
if dtype_name == "fp16":
return torch.float16
if dtype_name in {"auto", "bf16"}:
if torch.cuda.is_bf16_supported():
return torch.bfloat16
if dtype_name == "bf16":
raise DiffusionCliError(
"bf16 requested but selected device does not support bf16"
)
return torch.float32
raise AssertionError("unreachable dtype branch")
def buildGenerationConfigFromRequest(
request: ImageGenerationRequest,
user_config: UserConfig | None = None,
) -> GenerationConfig:
"""Validate an internal request and build a generation config."""
if user_config is None:
user_config = loadUserConfig()
generation = user_config.generation
models = user_config.models
if not request.prompt:
raise DiffusionCliError("--prompt is required for generation")
negative_prompt = coalesce(
request.negative_prompt,
generation.negative_prompt,
DEFAULT_NEGATIVE_PROMPT,
)
width = coalesce(request.width, generation.width, DEFAULT_WIDTH)
height = coalesce(request.height, generation.height, DEFAULT_HEIGHT)
batch_size = coalesce(
request.batch_size,
generation.batch_size,
DEFAULT_BATCH_SIZE,
)
steps = coalesce(request.steps, generation.steps, DEFAULT_STEPS)
cfg = coalesce(request.cfg, generation.cfg, DEFAULT_CFG)
device_name = coalesce(request.device, generation.device, DEFAULT_DEVICE)
dtype_name = coalesce(request.dtype, generation.dtype, DEFAULT_DTYPE)
output_path = coalesce(request.output, generation.output, DEFAULT_OUTPUT)
raw_extension = coalesce(
request.output_extension,
generation.output_extension,
DEFAULT_OUTPUT_EXTENSION,
)
raw_quality = coalesce(
request.output_quality,
generation.output_quality,
DEFAULT_OUTPUT_QUALITY,
)
tokenizer_path = request.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 {batch_size}"
)
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)
output_extension = normalizeOutputExtension(raw_extension)
output_quality = validateOutputQuality(raw_quality)
output_mime_type = outputMimeType(output_extension)
tokenizer_path = validateTokenizerPath(tokenizer_path)
seed = request.seed if request.seed is not None else randomSeed()
return GenerationConfig(
prompt=request.prompt,
negative_prompt=negative_prompt,
seed=seed,
width=width,
height=height,
batch_size=batch_size,
steps=steps,
cfg=cfg,
device=device,
dtype=dtype,
dtype_name=dtype_name,
output=output,
output_extension=output_extension,
output_quality=output_quality,
output_mime_type=output_mime_type,
tokenizer_path=tokenizer_path,
)
def buildDefaultGenerationRequest(
user_config: UserConfig,
) -> ImageGenerationRequest:
"""Build UI-visible generation defaults from config and fallbacks."""
generation = user_config.generation
output_extension = normalizeOutputExtension(
coalesce(generation.output_extension, DEFAULT_OUTPUT_EXTENSION)
)
output_quality = validateOutputQuality(
coalesce(generation.output_quality, DEFAULT_OUTPUT_QUALITY)
)
return ImageGenerationRequest(
prompt="",
negative_prompt=generation.negative_prompt or "",
seed=-1,
width=coalesce(generation.width, DEFAULT_WIDTH),
height=coalesce(generation.height, DEFAULT_HEIGHT),
batch_size=coalesce(generation.batch_size, DEFAULT_BATCH_SIZE),
steps=coalesce(generation.steps, DEFAULT_STEPS),
cfg=coalesce(generation.cfg, DEFAULT_CFG),
output_extension=output_extension,
output_quality=output_quality,
)
def buildGenerationConfig(
args,
user_config: UserConfig | None = None,
) -> GenerationConfig:
"""Validate parsed CLI arguments and build a generation config."""
return buildGenerationConfigFromRequest(
ImageGenerationRequest(
prompt=args.prompt,
negative_prompt=args.negative_prompt,
seed=args.seed,
width=args.width,
height=args.height,
batch_size=args.batch_size,
steps=args.steps,
cfg=args.cfg,
output=args.output,
output_extension=getattr(args, "output_extension", None),
output_quality=getattr(args, "output_quality", None),
device=args.device,
dtype=args.dtype,
tokenizer_path=args.tokenizer_path,
),
user_config,
)