Changes
diff --git a/diffusion_cli/checkpoint.py b/diffusion_cli/checkpoint.py
new file mode 100644
index 0000000..c128c67
--- /dev/null
+++ b/diffusion_cli/checkpoint.py
@@ -0,0 +1,61 @@
+"""Selective safetensors loading for component and checkpoint sources."""
+
+from __future__ import annotations
+
+import torch
+from safetensors import safe_open
+from safetensors.torch import load_file
+
+from diffusion_cli.config import ModelSource
+from diffusion_cli.errors import DiffusionCliError
+
+SAFETENSOR_DTYPE_MAP = {
+ "BF16": torch.bfloat16,
+ "F16": torch.float16,
+ "F32": torch.float32,
+}
+
+
+def loadStateDict(source: ModelSource) -> dict[str, torch.Tensor]:
+ """Load tensors for one model component from a safetensors source."""
+
+ if source.checkpoint_prefix is None:
+ return load_file(source.path, device="cpu")
+
+ state_dict: dict[str, torch.Tensor] = {}
+ with safe_open(source.path, framework="pt", device="cpu") as tensors:
+ for key in tensors.keys():
+ if key.startswith(source.checkpoint_prefix):
+ stripped_key = key.removeprefix(source.checkpoint_prefix)
+ state_dict[stripped_key] = tensors.get_tensor(key)
+
+ if not state_dict:
+ raise DiffusionCliError(
+ f"Checkpoint does not contain {source.role}: {source.path}"
+ )
+ return state_dict
+
+
+def inspectSourceTorchDtype(source: ModelSource) -> torch.dtype | None:
+ """Return a single floating torch dtype from source metadata if clear."""
+
+ dtypes: set[str] = set()
+ matched = False
+ with safe_open(source.path, framework="pt", device="cpu") as tensors:
+ for key in tensors.keys():
+ if source.checkpoint_prefix is not None:
+ if not key.startswith(source.checkpoint_prefix):
+ continue
+ matched = True
+ dtypes.add(tensors.get_slice(key).get_dtype())
+
+ if source.checkpoint_prefix is not None and not matched:
+ raise DiffusionCliError(
+ f"Checkpoint does not contain {source.role}: {source.path}"
+ )
+
+ mapped = {SAFETENSOR_DTYPE_MAP[dtype] for dtype in dtypes
+ if dtype in SAFETENSOR_DTYPE_MAP}
+ if len(mapped) == 1 and len(mapped) == len(dtypes):
+ return next(iter(mapped))
+ return None
diff --git a/diffusion_cli/cli.py b/diffusion_cli/cli.py
index 199f180..88c9b15 100644
--- a/diffusion_cli/cli.py
+++ b/diffusion_cli/cli.py
@@ -7,6 +7,7 @@ import gc
from pathlib import Path
import sys
+from diffusion_cli.checkpoint import inspectSourceTorchDtype
from diffusion_cli.config import (
UserConfig,
buildGenerationConfig,
@@ -14,8 +15,8 @@ from diffusion_cli.config import (
)
from diffusion_cli.errors import DiffusionCliError
from diffusion_cli.image_io import saveImages
-from diffusion_cli.model_inspect import formatSummary, inspectSafetensors
-from diffusion_cli.paths import resolveModelFiles
+from diffusion_cli.model_inspect import formatSummary, inspectModelSource
+from diffusion_cli.paths import resolveModelSources
from diffusion_cli.sampling import sampleLatents
from diffusion_cli.text_encoder import ZImageTextEncoder
from diffusion_cli.vae import ZImageVae
@@ -42,6 +43,11 @@ def buildParser() -> argparse.ArgumentParser:
parser.add_argument("--cfg", type=float)
parser.add_argument("--output", type=Path)
parser.add_argument("--device")
+ parser.add_argument(
+ "--checkpoint",
+ type=Path,
+ help="Local all-in-one safetensors checkpoint file.",
+ )
parser.add_argument(
"--diffusion-model",
type=Path,
@@ -64,7 +70,7 @@ def buildParser() -> argparse.ArgumentParser:
)
parser.add_argument(
"--dtype",
- choices=("auto", "bf16", "fp32"),
+ choices=("auto", "bf16", "fp16", "fp32"),
)
parser.add_argument(
"--inspect-models",
@@ -77,11 +83,11 @@ def buildParser() -> argparse.ArgumentParser:
def inspectModels(args, user_config: UserConfig) -> None:
"""Print metadata summaries for the required model files."""
- model_files = resolveModelFiles(args, user_config)
+ model_sources = resolveModelSources(args, user_config)
summaries = [
- ("diffusion model", inspectSafetensors(model_files.diffusion_model)),
- ("text encoder", inspectSafetensors(model_files.text_encoder)),
- ("VAE", inspectSafetensors(model_files.vae)),
+ ("diffusion model", inspectModelSource(model_sources.diffusion_model)),
+ ("text encoder", inspectModelSource(model_sources.text_encoder)),
+ ("VAE", inspectModelSource(model_sources.vae)),
]
output = "\n\n".join(
formatSummary(name, summary) for name, summary in summaries
@@ -102,17 +108,28 @@ def releaseMemory() -> None:
pass
+def componentDtype(config, source):
+ """Return the runtime dtype for one component source."""
+
+ if config.dtype_name != "auto":
+ return config.dtype
+ return inspectSourceTorchDtype(source) or config.dtype
+
+
def generate(args, user_config: UserConfig) -> list[Path]:
"""Validate generation arguments and run the current milestone."""
config = buildGenerationConfig(args, user_config)
- model_files = resolveModelFiles(args, user_config)
+ model_sources = resolveModelSources(args, user_config)
+ text_dtype = componentDtype(config, model_sources.text_encoder)
+ diffusion_dtype = componentDtype(config, model_sources.diffusion_model)
+ vae_dtype = componentDtype(config, model_sources.vae)
text_encoder = ZImageTextEncoder(
- model_files.text_encoder,
+ model_sources.text_encoder,
config.tokenizer_path,
config.device,
- config.dtype,
+ text_dtype,
)
conditioning = text_encoder.encodePrompts(
config.prompt,
@@ -122,9 +139,9 @@ def generate(args, user_config: UserConfig) -> list[Path]:
releaseMemory()
model = ZImageModel(
- model_files.diffusion_model,
+ model_sources.diffusion_model,
config.device,
- config.dtype,
+ diffusion_dtype,
)
latent = sampleLatents(
model,
@@ -136,12 +153,12 @@ def generate(args, user_config: UserConfig) -> list[Path]:
steps=config.steps,
cfg=config.cfg,
device=config.device,
- dtype=config.dtype,
+ dtype=diffusion_dtype,
)
del model
releaseMemory()
- vae = ZImageVae(model_files.vae, config.device, config.dtype)
+ vae = ZImageVae(model_sources.vae, config.device, vae_dtype)
images = vae.decode(latent)
return saveImages(images, config.output)
diff --git a/diffusion_cli/config.py b/diffusion_cli/config.py
index 37cfa30..df657a6 100644
--- a/diffusion_cli/config.py
+++ b/diffusion_cli/config.py
@@ -26,7 +26,13 @@ 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"}
+MODEL_CONFIG_KEYS = {
+ "checkpoint",
+ "diffusion_model",
+ "text_encoder",
+ "vae",
+ "tokenizer",
+}
GENERATION_CONFIG_KEYS = {
"negative_prompt",
"width",
@@ -38,6 +44,12 @@ GENERATION_CONFIG_KEYS = {
"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)
@@ -49,10 +61,29 @@ class ModelFiles:
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
@@ -96,6 +127,7 @@ class GenerationConfig:
cfg: float
device: object
dtype: object
+ dtype_name: str
output: Path
tokenizer_path: Path
@@ -217,6 +249,7 @@ def loadUserConfig(path: Path | None = None) -> UserConfig:
return UserConfig(
models=ModelPathConfig(
+ checkpoint=_optionalPath(models, "models", "checkpoint"),
diffusion_model=_optionalPath(
models,
"models",
@@ -324,13 +357,15 @@ def selectDtype(dtype_name: str, device) -> object:
import torch
- if dtype_name not in {"auto", "bf16", "fp32"}:
+ if dtype_name not in {"auto", "bf16", "fp16", "fp32"}:
raise DiffusionCliError(
- f"dtype must be one of auto, bf16, fp32: got {dtype_name}"
+ 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():
@@ -409,6 +444,7 @@ def buildGenerationConfig(
cfg=cfg,
device=device,
dtype=dtype,
+ dtype_name=dtype_name,
output=output,
tokenizer_path=tokenizer_path,
)
diff --git a/diffusion_cli/model_inspect.py b/diffusion_cli/model_inspect.py
index efac0cf..0cf6832 100644
--- a/diffusion_cli/model_inspect.py
+++ b/diffusion_cli/model_inspect.py
@@ -8,12 +8,16 @@ from pathlib import Path
from safetensors import safe_open
+from diffusion_cli.config import ModelSource
+from diffusion_cli.errors import DiffusionCliError
+
@dataclass(frozen=True)
class TensorSummary:
"""Small metadata summary for a safetensors checkpoint."""
path: Path
+ source_prefix: str | None
tensor_count: int
dtype_counts: dict[str, int]
top_level_counts: dict[str, int]
@@ -33,7 +37,13 @@ def guessArchitecture(keys: list[str]) -> str:
return "vae"
if "model.embed_tokens.weight" in key_set:
return "qwen_text_encoder"
- if "img_in.weight" in key_set or "x_embedder.proj.weight" in key_set:
+ if "embed_tokens.weight" in key_set:
+ return "qwen_text_encoder"
+ if (
+ "img_in.weight" in key_set
+ or "x_embedder.proj.weight" in key_set
+ or "x_embedder.weight" in key_set
+ ):
return "z_image_diffusion"
if "diffusion_model.img_in.weight" in key_set:
return "z_image_diffusion"
@@ -47,19 +57,37 @@ def guessArchitecture(keys: list[str]) -> str:
def inspectSafetensors(path: Path) -> TensorSummary:
"""Read safetensors metadata without loading tensor data into memory."""
+ return inspectModelSource(ModelSource(path=path, role="unknown"))
+
+
+def inspectModelSource(source: ModelSource) -> TensorSummary:
+ """Read metadata for one resolved model source."""
+
dtype_counts: Counter[str] = Counter()
top_counts: Counter[str] = Counter()
keys: list[str] = []
- with safe_open(path, framework="pt", device="cpu") as tensors:
+ with safe_open(source.path, framework="pt", device="cpu") as tensors:
for key in tensors.keys():
- keys.append(key)
+ if source.checkpoint_prefix is not None:
+ if not key.startswith(source.checkpoint_prefix):
+ continue
+ summary_key = key.removeprefix(source.checkpoint_prefix)
+ else:
+ summary_key = key
+ keys.append(summary_key)
metadata = tensors.get_slice(key)
dtype_counts[metadata.get_dtype()] += 1
- top_counts[_topLevel(key)] += 1
+ top_counts[_topLevel(summary_key)] += 1
+
+ if source.checkpoint_prefix is not None and not keys:
+ raise DiffusionCliError(
+ f"Checkpoint does not contain {source.role}: {source.path}"
+ )
return TensorSummary(
- path=path,
+ path=source.path,
+ source_prefix=source.checkpoint_prefix,
tensor_count=len(keys),
dtype_counts=dict(sorted(dtype_counts.items())),
top_level_counts=dict(top_counts.most_common(12)),
@@ -78,11 +106,16 @@ def formatSummary(name: str, summary: TensorSummary) -> str:
for prefix, count in summary.top_level_counts.items()
)
return "\n".join(
- [
+ [line for line in [
f"{name}: {summary.path}",
+ (
+ f" source prefix: {summary.source_prefix}"
+ if summary.source_prefix is not None
+ else None
+ ),
f" tensors: {summary.tensor_count}",
f" dtypes: {dtype_text}",
f" top-level keys: {top_text}",
f" architecture guess: {summary.architecture_guess}",
- ]
+ ] if line is not None]
)
diff --git a/diffusion_cli/paths.py b/diffusion_cli/paths.py
index 3b3b9f3..741a178 100644
--- a/diffusion_cli/paths.py
+++ b/diffusion_cli/paths.py
@@ -4,7 +4,18 @@ from __future__ import annotations
from pathlib import Path
-from diffusion_cli.config import ModelFiles, UserConfig
+from diffusion_cli.config import (
+ CHECKPOINT_DIFFUSION_PREFIX,
+ CHECKPOINT_TEXT_ENCODER_PREFIX,
+ CHECKPOINT_VAE_PREFIX,
+ DIFFUSION_ROLE,
+ TEXT_ENCODER_ROLE,
+ VAE_ROLE,
+ ModelFiles,
+ ModelSource,
+ ModelSources,
+ UserConfig,
+)
from diffusion_cli.errors import DiffusionCliError
@@ -20,21 +31,96 @@ def requireFile(path: Path | None, role: str, config_key: str) -> Path:
return resolved_path
-def resolveModelFiles(args, user_config: UserConfig) -> ModelFiles:
- """Resolve and validate explicit model paths for the active workflow."""
+def _pathFromArgs(args, name: str) -> Path | None:
+ return getattr(args, name, None)
+
+
+def _requireSelectedFile(path: Path, role: str) -> Path:
+ resolved_path = path.expanduser().resolve()
+ if not resolved_path.is_file():
+ raise DiffusionCliError(f"Missing {role}: {resolved_path}")
+ return resolved_path
+
+
+def resolveComponentSource(
+ cli_path: Path | None,
+ config_path: Path | None,
+ checkpoint_path: Path | None,
+ role: str,
+ checkpoint_prefix: str,
+ config_key: str,
+) -> ModelSource:
+ """Resolve one component path using component-before-checkpoint order."""
+
+ path = cli_path or config_path
+ if path is not None:
+ return ModelSource(
+ path=_requireSelectedFile(path, role),
+ role=role,
+ checkpoint_prefix=None,
+ )
+
+ if checkpoint_path is None:
+ raise DiffusionCliError(f"Missing {config_key} or models.checkpoint")
+
+ return ModelSource(
+ path=_requireSelectedFile(checkpoint_path, "checkpoint"),
+ role=role,
+ checkpoint_prefix=checkpoint_prefix,
+ )
+
+
+def resolveModelSources(args, user_config: UserConfig) -> ModelSources:
+ """Resolve and validate model sources for the active workflow."""
models = user_config.models
+ checkpoint = _pathFromArgs(args, "checkpoint") or models.checkpoint
- return ModelFiles(
- diffusion_model=requireFile(
- args.diffusion_model or models.diffusion_model,
- "diffusion model",
+ return ModelSources(
+ diffusion_model=resolveComponentSource(
+ _pathFromArgs(args, "diffusion_model"),
+ models.diffusion_model,
+ checkpoint,
+ DIFFUSION_ROLE,
+ CHECKPOINT_DIFFUSION_PREFIX,
"models.diffusion_model",
),
- text_encoder=requireFile(
- args.text_encoder or models.text_encoder,
- "text encoder",
+ text_encoder=resolveComponentSource(
+ _pathFromArgs(args, "text_encoder"),
+ models.text_encoder,
+ checkpoint,
+ TEXT_ENCODER_ROLE,
+ CHECKPOINT_TEXT_ENCODER_PREFIX,
"models.text_encoder",
),
- vae=requireFile(args.vae or models.vae, "VAE", "models.vae"),
+ vae=resolveComponentSource(
+ _pathFromArgs(args, "vae"),
+ models.vae,
+ checkpoint,
+ VAE_ROLE,
+ CHECKPOINT_VAE_PREFIX,
+ "models.vae",
+ ),
+ )
+
+
+def resolveModelFiles(args, user_config: UserConfig) -> ModelFiles:
+ """Resolve standalone model paths for compatibility with old callers."""
+
+ sources = resolveModelSources(args, user_config)
+ if any(
+ source.checkpoint_prefix is not None
+ for source in (
+ sources.diffusion_model,
+ sources.text_encoder,
+ sources.vae,
+ )
+ ):
+ raise DiffusionCliError(
+ "resolveModelFiles cannot return checkpoint-backed sources"
+ )
+ return ModelFiles(
+ diffusion_model=sources.diffusion_model.path,
+ text_encoder=sources.text_encoder.path,
+ vae=sources.vae.path,
)
diff --git a/diffusion_cli/text_encoder.py b/diffusion_cli/text_encoder.py
index 27f3fa1..06cdd70 100644
--- a/diffusion_cli/text_encoder.py
+++ b/diffusion_cli/text_encoder.py
@@ -5,8 +5,8 @@ from __future__ import annotations
from dataclasses import dataclass
from pathlib import Path
-from safetensors.torch import load_file
-
+from diffusion_cli.checkpoint import loadStateDict
+from diffusion_cli.config import ModelSource, TEXT_ENCODER_ROLE
QWEN3_4B_VOCAB_SIZE = 151936
QWEN3_4B_HIDDEN_SIZE = 2560
@@ -81,12 +81,36 @@ def normalizeQwenStateDict(state_dict: dict[str, object]) -> dict[str, object]:
return normalized
+def initializeQwenRotaryEmbedding(model) -> None:
+ """Initialize non-persistent Qwen RoPE buffers after meta allocation."""
+
+ import torch
+
+ config = model.config
+ head_dim = getattr(config, "head_dim", None)
+ if head_dim is None:
+ head_dim = config.hidden_size // config.num_attention_heads
+ rope_parameters = getattr(config, "rope_parameters", {})
+ rope_theta = rope_parameters.get("rope_theta", QWEN3_4B_ROPE_THETA)
+ device = model.rotary_emb.inv_freq.device
+ inv_freq = 1.0 / (
+ rope_theta
+ ** (
+ torch.arange(0, head_dim, 2, dtype=torch.float32, device=device)
+ / head_dim
+ )
+ )
+ model.rotary_emb.inv_freq = inv_freq
+ if hasattr(model.rotary_emb, "original_inv_freq"):
+ model.rotary_emb.original_inv_freq = inv_freq.clone()
+
+
class ZImageTextEncoder:
"""Load the local Qwen3-4B encoder and produce text conditioning."""
def __init__(
self,
- model_path: Path,
+ model_path: ModelSource | Path | None,
tokenizer_path: Path,
device,
dtype,
@@ -154,15 +178,22 @@ class ZImageTextEncoder:
tokenizer.pad_token_id = QWEN3_4B_PAD_TOKEN_ID
return tokenizer
- def _loadModel(self, model_path: Path):
+ def _loadModel(self, model_path: ModelSource | Path | None):
import torch
from transformers import Qwen3Model
+ if model_path is None:
+ raise ValueError("model_path is required when model is absent")
+ if isinstance(model_path, ModelSource):
+ model_source = model_path
+ else:
+ model_source = ModelSource(model_path, TEXT_ENCODER_ROLE)
+
+ state_dict = normalizeQwenStateDict(loadStateDict(model_source))
with torch.device("meta"):
model = Qwen3Model(buildQwen3_4BConfig())
model.to_empty(device="cpu")
- state_dict = normalizeQwenStateDict(load_file(model_path, device="cpu"))
missing_keys, unexpected_keys = model.load_state_dict(
state_dict,
strict=False,
@@ -181,4 +212,5 @@ class ZImageTextEncoder:
details,
))
+ initializeQwenRotaryEmbedding(model)
return model.to(device=self.device, dtype=self.dtype)
diff --git a/diffusion_cli/vae.py b/diffusion_cli/vae.py
index 1e4e262..5a67d1b 100644
--- a/diffusion_cli/vae.py
+++ b/diffusion_cli/vae.py
@@ -6,11 +6,16 @@ from dataclasses import dataclass
from pathlib import Path
import torch
-from safetensors.torch import load_file
from torch import nn
from torch.nn import functional as F
-from diffusion_cli.config import LATENT_CHANNELS, LATENT_DOWNSCALE
+from diffusion_cli.checkpoint import loadStateDict
+from diffusion_cli.config import (
+ LATENT_CHANNELS,
+ LATENT_DOWNSCALE,
+ VAE_ROLE,
+ ModelSource,
+)
from diffusion_cli.errors import DiffusionCliError
@@ -322,7 +327,7 @@ class ZImageVae:
def __init__(
self,
- model_path: Path | None,
+ model_path: ModelSource | Path | None,
device,
dtype,
*,
@@ -337,7 +342,12 @@ class ZImageVae:
if model_path is None:
raise ValueError("model_path is required when model is absent")
- state_dict = load_file(model_path, device="cpu")
+ if isinstance(model_path, ModelSource):
+ model_source = model_path
+ else:
+ model_source = ModelSource(model_path, VAE_ROLE)
+
+ state_dict = loadStateDict(model_source)
config = decoder_config or detectVaeConfig(state_dict)
decoder = VaeDecoder(config, dtype=dtype, device=device)
decoder_state = {
diff --git a/diffusion_cli/zimage_model.py b/diffusion_cli/zimage_model.py
index 1bbe252..2ff61e6 100644
--- a/diffusion_cli/zimage_model.py
+++ b/diffusion_cli/zimage_model.py
@@ -9,10 +9,11 @@ import re
from typing import Iterable
import torch
-from safetensors.torch import load_file
from torch import nn
from torch.nn import functional as F
+from diffusion_cli.checkpoint import loadStateDict
+from diffusion_cli.config import DIFFUSION_ROLE, ModelSource
from diffusion_cli.errors import DiffusionCliError
@@ -755,7 +756,7 @@ class ZImageModel:
def __init__(
self,
- model_path: Path,
+ model_path: ModelSource | Path | None,
device,
dtype,
*,
@@ -816,9 +817,15 @@ class ZImageModel:
raise DiffusionCliError("Z-Image forward produced NaN or Inf")
return output
- def _loadModel(self, model_path: Path) -> nn.Module:
- state_dict = normalizeZImageStateDict(load_file(model_path,
- device="cpu"))
+ def _loadModel(self, model_path: ModelSource | Path | None) -> nn.Module:
+ if model_path is None:
+ raise ValueError("model_path is required when model is absent")
+ if isinstance(model_path, ModelSource):
+ model_source = model_path
+ else:
+ model_source = ModelSource(model_path, DIFFUSION_ROLE)
+
+ state_dict = normalizeZImageStateDict(loadStateDict(model_source))
config = detectZImageConfig(state_dict)
with torch.device("meta"):
model = NextDiT(config, dtype=self.dtype)
diff --git a/tests/test_checkpoint.py b/tests/test_checkpoint.py
new file mode 100644
index 0000000..ce037e7
--- /dev/null
+++ b/tests/test_checkpoint.py
@@ -0,0 +1,115 @@
+from pathlib import Path
+import tempfile
+import unittest
+
+import torch
+from safetensors.torch import save_file
+
+from diffusion_cli.checkpoint import inspectSourceTorchDtype, loadStateDict
+from diffusion_cli.config import (
+ CHECKPOINT_DIFFUSION_PREFIX,
+ DIFFUSION_ROLE,
+ TEXT_ENCODER_ROLE,
+ ModelSource,
+)
+from diffusion_cli.errors import DiffusionCliError
+
+
+class CheckpointTest(unittest.TestCase):
+ def testLoadsOnlyMatchingCheckpointPrefix(self):
+ with tempfile.TemporaryDirectory() as temp_dir:
+ path = Path(temp_dir) / "aio.safetensors"
+ save_file(
+ {
+ "model.diffusion_model.cap_embedder.1.weight":
+ torch.ones(1, 1),
+ "model.diffusion_model.x_embedder.weight":
+ torch.ones(1, 1) * 2,
+ "text_encoders.qwen3_4b.transformer.model."
+ "embed_tokens.weight": torch.ones(1, 1) * 3,
+ "unrelated.weight": torch.ones(1, 1) * 4,
+ },
+ path,
+ )
+ source = ModelSource(
+ path=path,
+ role=DIFFUSION_ROLE,
+ checkpoint_prefix=CHECKPOINT_DIFFUSION_PREFIX,
+ )
+
+ state_dict = loadStateDict(source)
+
+ self.assertEqual(
+ sorted(state_dict),
+ ["cap_embedder.1.weight", "x_embedder.weight"],
+ )
+ self.assertTrue(
+ torch.equal(
+ state_dict["x_embedder.weight"],
+ torch.ones(1, 1) * 2,
+ )
+ )
+
+ def testMissingCheckpointPrefixFailsClearly(self):
+ with tempfile.TemporaryDirectory() as temp_dir:
+ path = Path(temp_dir) / "aio.safetensors"
+ save_file({"unrelated.weight": torch.ones(1, 1)}, path)
+ source = ModelSource(
+ path=path,
+ role=TEXT_ENCODER_ROLE,
+ checkpoint_prefix="missing.",
+ )
+
+ with self.assertRaises(DiffusionCliError) as context:
+ loadStateDict(source)
+
+ self.assertIn(
+ "Checkpoint does not contain text_encoder",
+ str(context.exception),
+ )
+
+ def testInspectsSingleCheckpointSourceDtype(self):
+ with tempfile.TemporaryDirectory() as temp_dir:
+ path = Path(temp_dir) / "aio.safetensors"
+ save_file(
+ {
+ "model.diffusion_model.x_embedder.weight":
+ torch.ones(1, 1, dtype=torch.float16),
+ "text_encoders.qwen3_4b.transformer.model."
+ "embed_tokens.weight": torch.ones(
+ 1,
+ 1,
+ dtype=torch.bfloat16,
+ ),
+ },
+ path,
+ )
+ source = ModelSource(
+ path=path,
+ role=DIFFUSION_ROLE,
+ checkpoint_prefix=CHECKPOINT_DIFFUSION_PREFIX,
+ )
+
+ dtype = inspectSourceTorchDtype(source)
+
+ self.assertEqual(dtype, torch.float16)
+
+ def testMixedSourceDtypeIsNotInferred(self):
+ with tempfile.TemporaryDirectory() as temp_dir:
+ path = Path(temp_dir) / "mixed.safetensors"
+ save_file(
+ {
+ "a.weight": torch.ones(1, 1, dtype=torch.float16),
+ "b.weight": torch.ones(1, 1, dtype=torch.bfloat16),
+ },
+ path,
+ )
+ source = ModelSource(path=path, role=DIFFUSION_ROLE)
+
+ dtype = inspectSourceTorchDtype(source)
+
+ self.assertIsNone(dtype)
+
+
+if __name__ == "__main__":
+ unittest.main()
diff --git a/tests/test_cli.py b/tests/test_cli.py
index 32e8a22..a97848b 100644
--- a/tests/test_cli.py
+++ b/tests/test_cli.py
@@ -16,6 +16,7 @@ class CliTest(unittest.TestCase):
self.assertIsNone(args.cfg)
self.assertIsNone(args.output)
self.assertIsNone(args.device)
+ self.assertIsNone(args.checkpoint)
self.assertIsNone(args.diffusion_model)
self.assertIsNone(args.text_encoder)
self.assertIsNone(args.vae)
@@ -32,13 +33,16 @@ class CliTest(unittest.TestCase):
"--cfg",
"1.5",
"--dtype",
- "bf16",
+ "fp16",
+ "--checkpoint",
+ "aio.safetensors",
]
)
self.assertEqual(args.steps, 12)
self.assertEqual(args.cfg, 1.5)
- self.assertEqual(args.dtype, "bf16")
+ self.assertEqual(args.dtype, "fp16")
+ self.assertEqual(args.checkpoint.name, "aio.safetensors")
def testModelRootIsNoLongerAccepted(self):
with patch("sys.stderr"):
diff --git a/tests/test_config.py b/tests/test_config.py
index 278bb6b..41474b5 100644
--- a/tests/test_config.py
+++ b/tests/test_config.py
@@ -61,6 +61,7 @@ class ConfigTest(unittest.TestCase):
"\n".join(
(
"[models]",
+ f'checkpoint = "{root / "aio.safetensors"}"',
f'diffusion_model = "{root / "diffusion.safetensors"}"',
f'text_encoder = "{root / "text_encoder.safetensors"}"',
f'vae = "{root / "ae.safetensors"}"',
@@ -83,6 +84,10 @@ class ConfigTest(unittest.TestCase):
config = loadUserConfig(config_path)
+ self.assertEqual(
+ config.models.checkpoint.name,
+ "aio.safetensors",
+ )
self.assertEqual(
config.models.diffusion_model.name,
"diffusion.safetensors",
@@ -165,6 +170,22 @@ class ConfigTest(unittest.TestCase):
str(context.exception),
)
+ def testWrongCheckpointTypeFails(self):
+ with tempfile.TemporaryDirectory() as temp_dir:
+ config_path = Path(temp_dir) / "diffusion.toml"
+ config_path.write_text(
+ "[models]\ncheckpoint = 12\n",
+ encoding="utf-8",
+ )
+
+ with self.assertRaises(DiffusionCliError) as context:
+ loadUserConfig(config_path)
+
+ self.assertIn(
+ "models.checkpoint must be a string path",
+ str(context.exception),
+ )
+
def testGenerationConfigUsesConfigDefaults(self):
with tempfile.TemporaryDirectory() as temp_dir:
root = Path(temp_dir)
diff --git a/tests/test_model_inspect.py b/tests/test_model_inspect.py
index 9d9f5e5..7bbf09c 100644
--- a/tests/test_model_inspect.py
+++ b/tests/test_model_inspect.py
@@ -1,6 +1,16 @@
+from pathlib import Path
+import tempfile
import unittest
-from diffusion_cli.model_inspect import guessArchitecture
+import torch
+from safetensors.torch import save_file
+
+from diffusion_cli.config import (
+ CHECKPOINT_DIFFUSION_PREFIX,
+ DIFFUSION_ROLE,
+ ModelSource,
+)
+from diffusion_cli.model_inspect import guessArchitecture, inspectModelSource
class ModelInspectTest(unittest.TestCase):
@@ -9,6 +19,33 @@ class ModelInspectTest(unittest.TestCase):
self.assertEqual(guessArchitecture(keys), "z_image_diffusion")
+ def testInspectCheckpointSourceUsesStrippedKeys(self):
+ with tempfile.TemporaryDirectory() as temp_dir:
+ path = Path(temp_dir) / "aio.safetensors"
+ save_file(
+ {
+ "model.diffusion_model.x_embedder.weight":
+ torch.ones(1, 1),
+ "model.diffusion_model.layers.0.attention.qkv.weight":
+ torch.ones(1, 1),
+ "unrelated.weight": torch.ones(1, 1),
+ },
+ path,
+ )
+ source = ModelSource(
+ path=path,
+ role=DIFFUSION_ROLE,
+ checkpoint_prefix=CHECKPOINT_DIFFUSION_PREFIX,
+ )
+
+ summary = inspectModelSource(source)
+
+ self.assertEqual(summary.tensor_count, 2)
+ self.assertEqual(summary.source_prefix, CHECKPOINT_DIFFUSION_PREFIX)
+ self.assertEqual(summary.top_level_counts, {"x_embedder": 1,
+ "layers": 1})
+ self.assertEqual(summary.architecture_guess, "z_image_diffusion")
+
if __name__ == "__main__":
unittest.main()
diff --git a/tests/test_paths.py b/tests/test_paths.py
index 5c31090..1db03f6 100644
--- a/tests/test_paths.py
+++ b/tests/test_paths.py
@@ -4,8 +4,15 @@ import unittest
from types import SimpleNamespace
from diffusion_cli.errors import DiffusionCliError
-from diffusion_cli.config import ModelPathConfig, UserConfig, GenerationDefaults
-from diffusion_cli.paths import resolveModelFiles
+from diffusion_cli.config import (
+ CHECKPOINT_DIFFUSION_PREFIX,
+ CHECKPOINT_TEXT_ENCODER_PREFIX,
+ CHECKPOINT_VAE_PREFIX,
+ GenerationDefaults,
+ ModelPathConfig,
+ UserConfig,
+)
+from diffusion_cli.paths import resolveModelFiles, resolveModelSources
class PathsTest(unittest.TestCase):
@@ -19,6 +26,7 @@ class PathsTest(unittest.TestCase):
path.write_bytes(b"placeholder")
args = SimpleNamespace(
+ checkpoint=None,
diffusion_model=None,
text_encoder=None,
vae=None,
@@ -40,8 +48,111 @@ class PathsTest(unittest.TestCase):
self.assertEqual(files.text_encoder.name, "text_encoder.safetensors")
self.assertEqual(files.vae.name, "ae.safetensors")
+ def testConfigCheckpointResolvesAllComponents(self):
+ with tempfile.TemporaryDirectory() as temp_dir:
+ checkpoint = Path(temp_dir) / "aio.safetensors"
+ checkpoint.write_bytes(b"placeholder")
+ args = SimpleNamespace(
+ checkpoint=None,
+ diffusion_model=None,
+ text_encoder=None,
+ vae=None,
+ )
+ user_config = UserConfig(
+ models=ModelPathConfig(checkpoint=checkpoint),
+ generation=GenerationDefaults(),
+ )
+
+ sources = resolveModelSources(args, user_config)
+
+ self.assertEqual(sources.diffusion_model.path, checkpoint.resolve())
+ self.assertEqual(sources.text_encoder.path, checkpoint.resolve())
+ self.assertEqual(sources.vae.path, checkpoint.resolve())
+ self.assertEqual(
+ sources.diffusion_model.checkpoint_prefix,
+ CHECKPOINT_DIFFUSION_PREFIX,
+ )
+ self.assertEqual(
+ sources.text_encoder.checkpoint_prefix,
+ CHECKPOINT_TEXT_ENCODER_PREFIX,
+ )
+ self.assertEqual(
+ sources.vae.checkpoint_prefix,
+ CHECKPOINT_VAE_PREFIX,
+ )
+
+ def testCliCheckpointOverridesConfigCheckpoint(self):
+ with tempfile.TemporaryDirectory() as temp_dir:
+ root = Path(temp_dir)
+ config_checkpoint = root / "config.safetensors"
+ cli_checkpoint = root / "cli.safetensors"
+ config_checkpoint.write_bytes(b"placeholder")
+ cli_checkpoint.write_bytes(b"placeholder")
+ args = SimpleNamespace(
+ checkpoint=cli_checkpoint,
+ diffusion_model=None,
+ text_encoder=None,
+ vae=None,
+ )
+ user_config = UserConfig(
+ models=ModelPathConfig(checkpoint=config_checkpoint),
+ generation=GenerationDefaults(),
+ )
+
+ sources = resolveModelSources(args, user_config)
+
+ self.assertEqual(sources.diffusion_model.path, cli_checkpoint.resolve())
+
+ def testComponentPathOverridesCheckpointForOneRole(self):
+ with tempfile.TemporaryDirectory() as temp_dir:
+ root = Path(temp_dir)
+ checkpoint = root / "aio.safetensors"
+ vae = root / "ae.safetensors"
+ checkpoint.write_bytes(b"placeholder")
+ vae.write_bytes(b"placeholder")
+ args = SimpleNamespace(
+ checkpoint=checkpoint,
+ diffusion_model=None,
+ text_encoder=None,
+ vae=vae,
+ )
+ user_config = UserConfig(
+ models=ModelPathConfig(),
+ generation=GenerationDefaults(),
+ )
+
+ sources = resolveModelSources(args, user_config)
+
+ self.assertEqual(sources.vae.path, vae.resolve())
+ self.assertIsNone(sources.vae.checkpoint_prefix)
+ self.assertEqual(
+ sources.diffusion_model.checkpoint_prefix,
+ CHECKPOINT_DIFFUSION_PREFIX,
+ )
+
def testMissingFileFailsClearly(self):
args = SimpleNamespace(
+ checkpoint=None,
+ diffusion_model=None,
+ text_encoder=None,
+ vae=None,
+ )
+ user_config = UserConfig(
+ models=ModelPathConfig(),
+ generation=GenerationDefaults(),
+ )
+
+ with self.assertRaises(DiffusionCliError) as context:
+ resolveModelSources(args, user_config)
+
+ self.assertIn(
+ "Missing models.diffusion_model or models.checkpoint",
+ str(context.exception),
+ )
+
+ def testMissingCheckpointFileFailsClearly(self):
+ args = SimpleNamespace(
+ checkpoint=Path("/missing/aio.safetensors"),
diffusion_model=None,
text_encoder=None,
vae=None,
@@ -52,9 +163,9 @@ class PathsTest(unittest.TestCase):
)
with self.assertRaises(DiffusionCliError) as context:
- resolveModelFiles(args, user_config)
+ resolveModelSources(args, user_config)
- self.assertIn("Missing models.diffusion_model", str(context.exception))
+ self.assertIn("Missing checkpoint", str(context.exception))
if __name__ == "__main__":
diff --git a/tests/test_text_encoder.py b/tests/test_text_encoder.py
index 2e65bbe..bccc1a4 100644
--- a/tests/test_text_encoder.py
+++ b/tests/test_text_encoder.py
@@ -1,12 +1,15 @@
import unittest
+from types import SimpleNamespace
import torch
from diffusion_cli.text_encoder import (
QWEN3_4B_LAYER_INDEX,
+ QWEN3_4B_ROPE_THETA,
ZImageTextEncoder,
applyZImagePromptTemplate,
buildQwen3_4BConfig,
+ initializeQwenRotaryEmbedding,
normalizeQwenStateDict,
)
@@ -73,6 +76,29 @@ class TextEncoderTest(unittest.TestCase):
self.assertEqual(normalized["embed_tokens.weight"], "a")
self.assertEqual(normalized["layers.0.input_layernorm.weight"], "b")
+ def testInitializeRotaryEmbeddingRebuildsMetaAllocatedBuffer(self):
+ class FakeRotary:
+ pass
+
+ model = SimpleNamespace()
+ model.config = buildQwen3_4BConfig()
+ model.rotary_emb = FakeRotary()
+ model.rotary_emb.inv_freq = torch.empty(64)
+ model.rotary_emb.original_inv_freq = torch.empty(64)
+
+ initializeQwenRotaryEmbedding(model)
+
+ self.assertTrue(torch.all(torch.isfinite(model.rotary_emb.inv_freq)))
+ self.assertEqual(model.rotary_emb.inv_freq[0].item(), 1.0)
+ self.assertAlmostEqual(
+ model.rotary_emb.inv_freq[-1].item(),
+ 1.0 / (QWEN3_4B_ROPE_THETA ** (126 / 128)),
+ )
+ self.assertTrue(torch.equal(
+ model.rotary_emb.inv_freq,
+ model.rotary_emb.original_inv_freq,
+ ))
+
def testQwenConfigMatchesComfyTextEncoderShape(self):
config = buildQwen3_4BConfig()