Changes
diff --git a/MANIFEST.in b/MANIFEST.in
new file mode 100644
index 0000000..bd7bbf8
--- /dev/null
+++ b/MANIFEST.in
@@ -0,0 +1,3 @@
+include THIRD_PARTY_NOTICES.md
+include scripts/krea2_integration.py
+recursive-include licenses *
diff --git a/README.md b/README.md
index e227b56..a32fd14 100644
--- a/README.md
+++ b/README.md
@@ -1,10 +1,9 @@
# diffusion-cli
-Standalone Python CLI scaffolding for local Z-Image Turbo inference.
+Standalone Python CLI for local diffusion model inference.
-The first implemented milestone validates CLI inputs, resolves the local
-model files, and inspects safetensors metadata without importing ComfyUI
-or downloading files.
+The CLI supports the existing Z-Image backend and local Krea 2 Turbo
+profiles. It never imports ComfyUI or downloads model files.
```bash
uv run diffusion-cli --help
@@ -71,6 +70,38 @@ Non-PNG output uses the `magick` executable from ImageMagick after image
generation completes. AVIF support depends on the local ImageMagick
build and installed delegates.
+## Named model profiles
+
+The legacy `[models]` table remains supported and is treated as an implicit
+`legacy` Z-Image profile. Named profiles can coexist:
+
+```toml
+default_model = "krea2-turbo"
+
+[model_profiles.krea2-turbo]
+architecture = "krea2"
+variant = "turbo"
+diffusion_model = "/models/krea2/krea2_turbo_fp8_scaled.safetensors"
+text_encoder = "/models/krea2/qwen3vl_4b_fp8_scaled.safetensors"
+vae = "/models/krea2/qwen_image_vae.safetensors"
+tokenizer = "/home/mw/programs/ComfyUI/comfy/text_encoders/qwen25_tokenizer"
+
+[model_profiles.z-image-turbo]
+architecture = "z-image"
+variant = "turbo"
+diffusion_model = "/models/z-image/diffusion.safetensors"
+text_encoder = "/models/z-image/qwen_3_4b.safetensors"
+vae = "/models/z-image/ae.safetensors"
+tokenizer = "/models/qwen25_tokenizer"
+```
+
+Select a profile for one-shot generation with `--model-profile krea2-turbo`.
+Krea 2 Turbo uses eight steps, CFG zero, and timestep shift `mu = 1.15` by
+default. Raw profiles use 52 steps and CFG 3.5. All model and tokenizer paths
+must be local regular files/directories; missing paths fail without checking a
+Hub cache. `HF_HUB_OFFLINE`, `TRANSFORMERS_OFFLINE`, and `DIFFUSERS_OFFLINE`
+are also enabled at process startup.
+
## HTTP API server
Start the SillyTavern compatibility API with an explicit profile:
diff --git a/THIRD_PARTY_NOTICES.md b/THIRD_PARTY_NOTICES.md
new file mode 100644
index 0000000..f8ed217
--- /dev/null
+++ b/THIRD_PARTY_NOTICES.md
@@ -0,0 +1,19 @@
+# Third-party notices
+
+## Krea 2 reference implementation
+
+`diffusion_cli/krea2_model.py` is an independent adaptation of the Krea 2
+single-stream MMDiT implementation in the Krea 2 repository:
+
+<https://github.com/krea-ai/krea-2>
+
+The reference source is licensed under Apache License 2.0. This file and the
+module docstring identify the source and state that the implementation was
+modified for local safetensors loading, the `InferenceBackend` boundary, and
+scaled-FP8 layers. The applicable license text is shipped at
+`licenses/KREA2-APACHE-2.0.txt`.
+
+Krea 2 model weights are governed by the separate Krea 2 Community License
+Agreement. That model license is not replaced or expanded by this source
+code notice. Deployments using Krea 2 weights remain responsible for the
+agreement's content-filtering and attribution requirements.
diff --git a/diffusion_cli/__init__.py b/diffusion_cli/__init__.py
index 1b0cca1..81898c7 100644
--- a/diffusion_cli/__init__.py
+++ b/diffusion_cli/__init__.py
@@ -1,5 +1,13 @@
"""Standalone diffusion CLI package."""
+import os
+
+# Local checkpoint loading is an explicit product guarantee. Set these before
+# any optional Transformers or Diffusers import as defense in depth.
+os.environ.setdefault("HF_HUB_OFFLINE", "1")
+os.environ.setdefault("TRANSFORMERS_OFFLINE", "1")
+os.environ.setdefault("DIFFUSERS_OFFLINE", "1")
+
__all__ = ["__version__"]
__version__ = "0.1.0"
diff --git a/diffusion_cli/api_profiles.py b/diffusion_cli/api_profiles.py
index 14a155b..ce61f40 100644
--- a/diffusion_cli/api_profiles.py
+++ b/diffusion_cli/api_profiles.py
@@ -56,7 +56,7 @@ def _optionalFloat(data: dict, key: str) -> float | None:
return float(value)
-def _txt2imgRequest(data: dict) -> ImageGenerationRequest:
+def _txt2imgRequest(data: dict, model_profile=None) -> ImageGenerationRequest:
for field in UNSUPPORTED_IMAGE_FIELDS:
if field in data:
raise DiffusionCliError(f"{field} is not supported")
@@ -70,8 +70,17 @@ def _txt2imgRequest(data: dict) -> ImageGenerationRequest:
seed = None
width = _optionalInt(data, "width")
height = _optionalInt(data, "height")
- if width is not None and height is not None:
+ is_krea = (
+ model_profile is not None
+ and model_profile.architecture == "krea2"
+ )
+ if width is not None and height is not None and not is_krea:
validateDimensions(width, height)
+ elif is_krea:
+ if width is not None and width <= 0:
+ raise DiffusionCliError(f"Width must be positive: got {width}")
+ if height is not None and height <= 0:
+ raise DiffusionCliError(f"Height must be positive: got {height}")
elif width is not None and (width <= 0 or width % 8 != 0):
raise DiffusionCliError(
f"Width must be a positive multiple of 8: got {width}"
@@ -125,9 +134,17 @@ def registerSillyTavernSdcppRoutes(app, context) -> None:
@app.route("/health", methods=["GET"])
def health():
+ profile = getattr(context.generation_service, "modelProfile", None)
return jsonify({
"status": "ok",
"api_profile": context.server_config.api_profile,
+ "model_profile": (
+ profile.name if profile is not None else "z-image-local"
+ ),
+ "architecture": (
+ profile.architecture if profile is not None else "z-image"
+ ),
+ "variant": profile.variant if profile is not None else "turbo",
})
@app.route("/v1/images/generations", methods=["OPTIONS"])
@@ -138,10 +155,15 @@ def registerSillyTavernSdcppRoutes(app, context) -> None:
@app.route("/v1/models", methods=["GET"])
def models():
+ model_id = getattr(
+ context.generation_service,
+ "modelId",
+ "z-image-local",
+ )
return jsonify({
"data": [
{
- "id": "z-image-local",
+ "id": model_id,
"object": "model",
"owned_by": "local",
}
@@ -155,7 +177,10 @@ def registerSillyTavernSdcppRoutes(app, context) -> None:
return _badRequest("JSON object body required")
try:
- generation_request = _txt2imgRequest(data)
+ generation_request = _txt2imgRequest(
+ data,
+ getattr(context.generation_service, "modelProfile", None),
+ )
except DiffusionCliError as exc:
return _badRequest(str(exc))
diff --git a/diffusion_cli/backends.py b/diffusion_cli/backends.py
new file mode 100644
index 0000000..c5efa22
--- /dev/null
+++ b/diffusion_cli/backends.py
@@ -0,0 +1,388 @@
+"""Inference backend boundary shared by one-shot and server generation."""
+
+from __future__ import annotations
+
+from pathlib import Path
+from typing import Protocol
+
+import torch
+
+from diffusion_cli.config import GenerationConfig, ModelProfile, ModelSources
+from diffusion_cli.errors import DiffusionCliError
+
+
+class InferenceBackend(Protocol):
+ """Generate NCHW RGB image tensors for one local model profile."""
+
+ @property
+ def modelId(self) -> str:
+ """Return the stable identifier exposed to clients."""
+
+ def generate(self, config: GenerationConfig) -> torch.Tensor:
+ """Generate RGB image tensors in the range zero through one."""
+
+
+def _stageError(profile: ModelProfile | None, stage: str, error: Exception):
+ """Turn a CUDA OOM into an actionable stage-specific CLI error."""
+
+ if isinstance(error, DiffusionCliError):
+ raise error
+ if isinstance(error, RuntimeError) and "out of memory" in str(error).lower():
+ name = profile.name if profile is not None else "local"
+ raise DiffusionCliError(
+ f"CUDA out of memory in {stage} for model profile {name}"
+ ) from error
+ raise error
+
+
+class ZImageBackend:
+ """Inference backend retaining the existing Z-Image generation path."""
+
+ def __init__(
+ self,
+ profile: ModelProfile | None,
+ sources: ModelSources,
+ residency: str,
+ *,
+ text_encoder_factory=None,
+ model_factory=None,
+ vae_factory=None,
+ sampler=None,
+ component_dtype=None,
+ ) -> None:
+ from diffusion_cli.text_encoder import ZImageTextEncoder
+ from diffusion_cli.vae import ZImageVae
+ from diffusion_cli.zimage_model import ZImageModel
+ from diffusion_cli.sampling import sampleLatents
+
+ self.profile = profile
+ self.sources = sources
+ self.residency = residency
+ self.text_encoder_factory = text_encoder_factory or ZImageTextEncoder
+ self.model_factory = model_factory or ZImageModel
+ self.vae_factory = vae_factory or ZImageVae
+ self.sampler = sampler or sampleLatents
+ self.component_dtype = component_dtype
+ self._text_encoder = None
+ self._model = None
+ self._vae = None
+
+ @property
+ def modelId(self) -> str:
+ """Return the configured profile name or legacy identity."""
+
+ return self.profile.name if self.profile is not None else "z-image-local"
+
+ def generate(self, config: GenerationConfig):
+ """Generate RGB tensors using staged or CPU-cached Z-Image residency."""
+
+ if self.residency == "cpu-cache":
+ return self._generateCpuCache(config)
+ return self._generateStaged(config)
+
+ def _dtype(self, config, source):
+ if self.component_dtype is not None:
+ return self.component_dtype(config, source)
+ if config.dtype_name != "auto":
+ return config.dtype
+ from diffusion_cli.checkpoint import inspectSourceTorchDtype
+
+ return inspectSourceTorchDtype(source) or config.dtype
+
+ def _generateStaged(self, config):
+ text_dtype = self._dtype(config, self.sources.text_encoder)
+ diffusion_dtype = self._dtype(config, self.sources.diffusion_model)
+ vae_dtype = self._dtype(config, self.sources.vae)
+ try:
+ text_encoder = self.text_encoder_factory(
+ self.sources.text_encoder,
+ config.tokenizer_path,
+ config.device,
+ text_dtype,
+ )
+ conditioning = text_encoder.encodePrompts(
+ config.prompt,
+ config.negative_prompt,
+ )
+ del text_encoder
+ _releaseMemory()
+ except Exception as exc:
+ _stageError(self.profile, "text encoder", exc)
+
+ try:
+ model = self.model_factory(
+ self.sources.diffusion_model,
+ config.device,
+ diffusion_dtype,
+ )
+ latent = self.sampler(
+ model,
+ conditioning,
+ batch_size=config.batch_size,
+ height=config.height,
+ width=config.width,
+ seed=config.seed,
+ steps=config.steps,
+ cfg=config.cfg,
+ device=config.device,
+ dtype=diffusion_dtype,
+ )
+ del model
+ _releaseMemory()
+ except Exception as exc:
+ _stageError(self.profile, "diffusion model", exc)
+
+ try:
+ vae = self.vae_factory(
+ self.sources.vae,
+ config.device,
+ vae_dtype,
+ )
+ images = vae.decode(latent)
+ del vae
+ _releaseMemory()
+ return images
+ except Exception as exc:
+ _stageError(self.profile, "VAE", exc)
+
+ def _generateCpuCache(self, config):
+ cpu = torch.device("cpu")
+ text_dtype = self._dtype(config, self.sources.text_encoder)
+ diffusion_dtype = self._dtype(config, self.sources.diffusion_model)
+ vae_dtype = self._dtype(config, self.sources.vae)
+ try:
+ if self._text_encoder is None:
+ self._text_encoder = self.text_encoder_factory(
+ self.sources.text_encoder,
+ config.tokenizer_path,
+ cpu,
+ text_dtype,
+ )
+ self._text_encoder.toDevice(config.device, text_dtype)
+ try:
+ conditioning = self._text_encoder.encodePrompts(
+ config.prompt,
+ config.negative_prompt,
+ )
+ finally:
+ self._text_encoder.toCpu()
+ _releaseMemory()
+ except Exception as exc:
+ _stageError(self.profile, "text encoder", exc)
+
+ try:
+ if self._model is None:
+ self._model = self.model_factory(
+ self.sources.diffusion_model,
+ cpu,
+ diffusion_dtype,
+ )
+ self._model.toDevice(config.device, diffusion_dtype)
+ try:
+ latent = self.sampler(
+ self._model,
+ conditioning,
+ batch_size=config.batch_size,
+ height=config.height,
+ width=config.width,
+ seed=config.seed,
+ steps=config.steps,
+ cfg=config.cfg,
+ device=config.device,
+ dtype=diffusion_dtype,
+ )
+ finally:
+ self._model.toCpu()
+ _releaseMemory()
+ except Exception as exc:
+ _stageError(self.profile, "diffusion model", exc)
+
+ try:
+ if self._vae is None:
+ self._vae = self.vae_factory(
+ self.sources.vae,
+ cpu,
+ vae_dtype,
+ )
+ self._vae.toDevice(config.device, vae_dtype)
+ try:
+ images = self._vae.decode(latent)
+ finally:
+ self._vae.toCpu()
+ _releaseMemory()
+ return images
+ except Exception as exc:
+ _stageError(self.profile, "VAE", exc)
+
+
+class Krea2Backend:
+ """Local Krea 2 Turbo/Raw backend with model-specific sampling."""
+
+ def __init__(self, profile: ModelProfile, sources: ModelSources, residency: str):
+ self.profile = profile
+ self.sources = sources
+ self.residency = residency
+ self._text_encoder = None
+ self._model = None
+ self._vae = None
+
+ @property
+ def modelId(self) -> str:
+ """Return the selected Krea profile name."""
+
+ return self.profile.name
+
+ def _constructText(self, source, tokenizer, device, dtype):
+ from diffusion_cli.krea2_text_encoder import Krea2TextEncoder
+
+ return Krea2TextEncoder(source, tokenizer, device, dtype)
+
+ def _constructModel(self, source, device, dtype):
+ from diffusion_cli.krea2_model import buildKrea2Model
+
+ return buildKrea2Model(source, device=device, dtype=dtype)
+
+ def _constructVae(self, source, device, dtype):
+ from diffusion_cli.qwen_image_vae import QwenImageVae
+
+ return QwenImageVae(source, device, dtype)
+
+ def _sample(self, model, positive, negative, config, dtype):
+ from diffusion_cli.krea2_sampling import sampleKrea2
+
+ return sampleKrea2(
+ model,
+ positive,
+ negative_conditioning=negative,
+ batch_size=config.batch_size,
+ height=config.height,
+ width=config.width,
+ seed=config.seed,
+ steps=config.steps,
+ cfg=config.cfg,
+ device=config.device,
+ dtype=dtype,
+ mu=config.mu,
+ shift_y1=config.shift_y1,
+ shift_y2=config.shift_y2,
+ )
+
+ def _generateStages(self, config, *, cache: bool):
+ import gc
+
+ cpu = torch.device("cpu")
+ device = config.device
+ dtype = config.dtype
+ try:
+ if cache and self._text_encoder is not None:
+ text_encoder = self._text_encoder
+ else:
+ text_encoder = self._constructText(
+ self.sources.text_encoder,
+ config.tokenizer_path,
+ cpu if cache else device,
+ dtype,
+ )
+ if cache:
+ self._text_encoder = text_encoder
+ if cache:
+ text_encoder.toDevice(device, dtype)
+ positive = text_encoder.encodePrompt(config.prompt)
+ negative = (
+ text_encoder.encodePrompt(config.negative_prompt)
+ if config.cfg > 0
+ else None
+ )
+ if cache:
+ text_encoder.toCpu()
+ else:
+ del text_encoder
+ gc.collect()
+ except Exception as exc:
+ _stageError(self.profile, "text encoder", exc)
+
+ try:
+ if cache and self._model is not None:
+ model = self._model
+ else:
+ model = self._constructModel(
+ self.sources.diffusion_model,
+ cpu if cache else device,
+ dtype,
+ )
+ if cache:
+ self._model = model
+ if cache:
+ model.toDevice(device, dtype)
+ latent = self._sample(model, positive, negative, config, dtype)
+ if cache:
+ model.toCpu()
+ else:
+ del model
+ gc.collect()
+ except Exception as exc:
+ _stageError(self.profile, "diffusion model", exc)
+
+ try:
+ if cache and self._vae is not None:
+ vae = self._vae
+ else:
+ vae = self._constructVae(
+ self.sources.vae,
+ cpu if cache else device,
+ dtype,
+ )
+ if cache:
+ self._vae = vae
+ if cache:
+ vae.toDevice(device, dtype)
+ images = vae.decode(latent)
+ if cache:
+ vae.toCpu()
+ else:
+ del vae
+ gc.collect()
+ return images
+ except Exception as exc:
+ _stageError(self.profile, "VAE", exc)
+
+ def generate(self, config: GenerationConfig):
+ """Generate RGB images with Krea's flow-matching sampler."""
+
+ return self._generateStages(config, cache=self.residency == "cpu-cache")
+
+
+def _releaseMemory() -> None:
+ """Release Python and CUDA caches between staged components."""
+
+ import gc
+
+ gc.collect()
+ if torch.cuda.is_available():
+ torch.cuda.empty_cache()
+
+
+BACKEND_FACTORIES = {
+ "z-image": ZImageBackend,
+ "krea2": Krea2Backend,
+}
+
+
+def createBackend(
+ profile: ModelProfile,
+ sources: ModelSources,
+ residency: str,
+ **kwargs,
+) -> InferenceBackend:
+ """Construct the registered backend for one configured profile."""
+
+ factory = BACKEND_FACTORIES.get(profile.architecture)
+ if factory is None:
+ raise DiffusionCliError(
+ f"Unknown model architecture: {profile.architecture}"
+ )
+ if factory is ZImageBackend:
+ return factory(profile, sources, residency, **kwargs)
+ if kwargs:
+ raise TypeError("Krea2Backend does not accept factory overrides")
+ return factory(profile, sources, residency)
diff --git a/diffusion_cli/checkpoint.py b/diffusion_cli/checkpoint.py
index c128c67..d2ea607 100644
--- a/diffusion_cli/checkpoint.py
+++ b/diffusion_cli/checkpoint.py
@@ -13,6 +13,8 @@ SAFETENSOR_DTYPE_MAP = {
"BF16": torch.bfloat16,
"F16": torch.float16,
"F32": torch.float32,
+ "F8_E4M3": torch.float8_e4m3fn,
+ "F8_E4M3FN": torch.float8_e4m3fn,
}
diff --git a/diffusion_cli/cli.py b/diffusion_cli/cli.py
index a30f1eb..3f7461c 100644
--- a/diffusion_cli/cli.py
+++ b/diffusion_cli/cli.py
@@ -10,9 +10,10 @@ import sys
from diffusion_cli.checkpoint import inspectSourceTorchDtype
from diffusion_cli.api_profiles import API_PROFILES
from diffusion_cli.config import (
+ ImageGenerationRequest,
UserConfig,
- buildGenerationConfig,
loadUserConfig,
+ selectModelProfile,
)
from diffusion_cli.errors import DiffusionCliError
from diffusion_cli.image_io import saveImages
@@ -20,9 +21,7 @@ from diffusion_cli.model_inspect import formatSummary, inspectModelSource
from diffusion_cli.paths import resolveModelSources
from diffusion_cli.sampling import sampleLatents
from diffusion_cli.server import serve, validateServerConfig
-from diffusion_cli.text_encoder import ZImageTextEncoder
-from diffusion_cli.vae import ZImageVae
-from diffusion_cli.zimage_model import ZImageModel
+from diffusion_cli.generation_service import GenerationService
def buildParser() -> argparse.ArgumentParser:
@@ -30,7 +29,11 @@ def buildParser() -> argparse.ArgumentParser:
parser = argparse.ArgumentParser(
prog="diffusion-cli",
- description="Standalone local Z-Image Turbo CLI.",
+ description="Standalone local diffusion model CLI.",
+ )
+ parser.add_argument(
+ "--model-profile",
+ help="Named local model profile to use.",
)
parser.add_argument("--prompt", help="Positive prompt.")
parser.add_argument(
@@ -43,6 +46,9 @@ def buildParser() -> argparse.ArgumentParser:
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)
@@ -115,6 +121,10 @@ def buildParser() -> argparse.ArgumentParser:
def inspectModels(args, user_config: UserConfig) -> None:
"""Print metadata summaries for the required model files."""
+ profile = selectModelProfile(
+ user_config,
+ getattr(args, "model_profile", None),
+ )
model_sources = resolveModelSources(args, user_config)
summaries = [
("diffusion model", inspectModelSource(model_sources.diffusion_model)),
@@ -122,7 +132,7 @@ def inspectModels(args, user_config: UserConfig) -> None:
("VAE", inspectModelSource(model_sources.vae)),
]
output = "\n\n".join(
- formatSummary(name, summary) for name, summary in summaries
+ formatSummary(name, summary, profile=profile) for name, summary in summaries
)
print(output)
@@ -149,55 +159,34 @@ def componentDtype(config, source):
def generate(args, user_config: UserConfig) -> list[Path]:
- """Validate generation arguments and run the current milestone."""
-
- config = buildGenerationConfig(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_sources.text_encoder,
- config.tokenizer_path,
- config.device,
- text_dtype,
- )
- conditioning = text_encoder.encodePrompts(
- config.prompt,
- config.negative_prompt,
- )
- del text_encoder
- releaseMemory()
-
- model = ZImageModel(
- model_sources.diffusion_model,
- config.device,
- diffusion_dtype,
- )
- latent = sampleLatents(
- model,
- conditioning,
- batch_size=config.batch_size,
- height=config.height,
- width=config.width,
- seed=config.seed,
- steps=config.steps,
- cfg=config.cfg,
- device=config.device,
- dtype=diffusion_dtype,
- )
- del model
- releaseMemory()
-
- vae = ZImageVae(model_sources.vae, config.device, vae_dtype)
- images = vae.decode(latent)
- return saveImages(
- images,
- config.output,
- config.output_extension,
- config.output_quality,
- )
+ """Run one-shot generation through the shared generation service."""
+
+ request = 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=args.output_extension,
+ output_quality=args.output_quality,
+ device=args.device,
+ dtype=args.dtype,
+ tokenizer_path=args.tokenizer_path,
+ mu=args.mu,
+ shift_y1=args.shift_y1,
+ shift_y2=args.shift_y2,
+ )
+ service = GenerationService(
+ user_config,
+ model_residency="staged",
+ model_profile_name=args.model_profile,
+ path_overrides=args,
+ )
+ return service.generateToFiles(request)
def main(argv: list[str] | None = None) -> int:
@@ -215,7 +204,11 @@ def main(argv: list[str] | None = None) -> int:
args.port,
args.model_residency,
)
- serve(server_config, user_config)
+ serve(
+ server_config,
+ user_config,
+ model_profile_name=args.model_profile,
+ )
return 0
if args.inspect_models:
diff --git a/diffusion_cli/config.py b/diffusion_cli/config.py
index f02b7d7..e01293e 100644
--- a/diffusion_cli/config.py
+++ b/diffusion_cli/config.py
@@ -2,7 +2,7 @@
from __future__ import annotations
-from dataclasses import dataclass
+from dataclasses import dataclass, field
from pathlib import Path
from secrets import randbits
import tomllib
@@ -17,6 +17,9 @@ DEFAULT_HEIGHT = 1248
DEFAULT_BATCH_SIZE = 1
DEFAULT_STEPS = 10
DEFAULT_CFG = 1.0
+DEFAULT_MU = None
+DEFAULT_SHIFT_Y1 = 0.5
+DEFAULT_SHIFT_Y2 = 1.15
DEFAULT_OUTPUT = Path("output.png")
DEFAULT_OUTPUT_EXTENSION = "png"
DEFAULT_OUTPUT_QUALITY = 95
@@ -82,7 +85,12 @@ UI_RESOLUTION_PRESETS = (
(2016, 864, "1536", "21:9"),
(864, 2016, "1536", "9:21"),
)
-TOP_LEVEL_CONFIG_KEYS = {"models", "generation"}
+TOP_LEVEL_CONFIG_KEYS = {
+ "default_model",
+ "models",
+ "model_profiles",
+ "generation",
+}
MODEL_CONFIG_KEYS = {
"checkpoint",
"diffusion_model",
@@ -90,6 +98,15 @@ MODEL_CONFIG_KEYS = {
"vae",
"tokenizer",
}
+MODEL_PROFILE_CONFIG_KEYS = {
+ "architecture",
+ "variant",
+ "checkpoint",
+ "diffusion_model",
+ "text_encoder",
+ "vae",
+ "tokenizer",
+}
GENERATION_CONFIG_KEYS = {
"negative_prompt",
"width",
@@ -102,6 +119,14 @@ GENERATION_CONFIG_KEYS = {
"output_quality",
"device",
"dtype",
+ "mu",
+ "shift_y1",
+ "shift_y2",
+}
+SUPPORTED_ARCHITECTURES = {"z-image", "krea2"}
+SUPPORTED_VARIANTS = {
+ "z-image": {"turbo"},
+ "krea2": {"raw", "turbo"},
}
DIFFUSION_ROLE = "diffusion_model"
TEXT_ENCODER_ROLE = "text_encoder"
@@ -136,6 +161,7 @@ class ModelSources:
diffusion_model: ModelSource
text_encoder: ModelSource
vae: ModelSource
+ tokenizer: Path | None = None
@dataclass(frozen=True)
@@ -149,6 +175,34 @@ class ModelPathConfig:
tokenizer: Path | None = None
+@dataclass(frozen=True)
+class ModelProfile:
+ """Local components and behavior for one named backend."""
+
+ name: str
+ architecture: str
+ variant: str
+ checkpoint: Path | None = None
+ diffusion_model: Path | None = None
+ text_encoder: Path | None = None
+ vae: Path | None = None
+ tokenizer: Path | None = None
+
+ def hasComponentPath(self) -> bool:
+ """Return whether the profile names at least one model component."""
+
+ return any(
+ path is not None
+ for path in (
+ self.checkpoint,
+ self.diffusion_model,
+ self.text_encoder,
+ self.vae,
+ self.tokenizer,
+ )
+ )
+
+
@dataclass(frozen=True)
class GenerationDefaults:
"""Optional generation defaults loaded from user configuration."""
@@ -164,14 +218,24 @@ class GenerationDefaults:
output_quality: int | None = None
device: str | None = None
dtype: str | None = None
+ mu: float | None = None
+ shift_y1: float | None = None
+ shift_y2: float | None = None
@dataclass(frozen=True)
class UserConfig:
- """User-provided model paths and generation defaults."""
+ """Validated profiles and generation defaults.
- models: ModelPathConfig
- generation: GenerationDefaults
+ ``models`` remains as a compatibility view for configuration files using
+ the original ``[models]`` table. New callers should use ``model_profiles``
+ and ``default_model``.
+ """
+
+ models: ModelPathConfig | None = None
+ generation: GenerationDefaults = field(default_factory=GenerationDefaults)
+ model_profiles: dict[str, ModelProfile] = field(default_factory=dict)
+ default_model: str | None = None
@dataclass(frozen=True)
@@ -194,6 +258,10 @@ class GenerationConfig:
output_quality: int
output_mime_type: str
tokenizer_path: Path
+ model_profile: ModelProfile | None = None
+ mu: float | None = None
+ shift_y1: float = DEFAULT_SHIFT_Y1
+ shift_y2: float = DEFAULT_SHIFT_Y2
@dataclass(frozen=True)
@@ -214,6 +282,9 @@ class ImageGenerationRequest:
device: str | None = None
dtype: str | None = None
tokenizer_path: Path | None = None
+ mu: float | None = None
+ shift_y1: float | None = None
+ shift_y2: float | None = None
def randomSeed() -> int:
@@ -306,6 +377,102 @@ def _optionalFloat(
return float(value)
+def _validateProfileIdentity(
+ name: str,
+ architecture: str,
+ variant: str,
+) -> None:
+ """Validate the backend identity declared by one profile."""
+
+ if not name.strip():
+ raise DiffusionCliError("Model profile name must not be empty")
+ if architecture not in SUPPORTED_ARCHITECTURES:
+ raise DiffusionCliError(
+ f"Unknown model architecture in {name}: {architecture}"
+ )
+ if variant not in SUPPORTED_VARIANTS[architecture]:
+ if architecture == "krea2":
+ raise DiffusionCliError(
+ f"Unsupported Krea 2 variant: {variant}"
+ )
+ raise DiffusionCliError(
+ f"Unsupported {architecture} variant in {name}: {variant}"
+ )
+
+
+def _profileFromTable(name: str, table: dict[str, Any]) -> ModelProfile:
+ """Parse and validate one named model profile."""
+
+ _rejectUnknownKeys(table, MODEL_PROFILE_CONFIG_KEYS, f"model_profiles.{name}")
+ architecture = _optionalString(
+ table, f"model_profiles.{name}", "architecture"
+ )
+ variant = _optionalString(table, f"model_profiles.{name}", "variant")
+ if architecture is None:
+ raise DiffusionCliError(
+ f"Missing model_profiles.{name}.architecture"
+ )
+ if variant is None:
+ raise DiffusionCliError(f"Missing model_profiles.{name}.variant")
+ _validateProfileIdentity(name, architecture, variant)
+
+ profile = ModelProfile(
+ name=name,
+ architecture=architecture,
+ variant=variant,
+ checkpoint=_optionalPath(
+ table, f"model_profiles.{name}", "checkpoint"
+ ),
+ diffusion_model=_optionalPath(
+ table, f"model_profiles.{name}", "diffusion_model"
+ ),
+ text_encoder=_optionalPath(
+ table, f"model_profiles.{name}", "text_encoder"
+ ),
+ vae=_optionalPath(table, f"model_profiles.{name}", "vae"),
+ tokenizer=_optionalPath(
+ table, f"model_profiles.{name}", "tokenizer"
+ ),
+ )
+ if not profile.hasComponentPath():
+ raise DiffusionCliError(
+ f"Model profile {name} has no component paths or checkpoint"
+ )
+ if architecture == "krea2" and profile.checkpoint is not None:
+ raise DiffusionCliError(
+ f"Krea profile {name} requires separate diffusion_model, "
+ "text_encoder, vae, and tokenizer paths"
+ )
+ return profile
+
+
+def _legacyProfile(models: ModelPathConfig) -> ModelProfile | None:
+ """Convert legacy model paths into an in-memory Z-Image profile."""
+
+ if models is None:
+ return None
+ if not models.checkpoint and not any(
+ path is not None
+ for path in (
+ models.diffusion_model,
+ models.text_encoder,
+ models.vae,
+ models.tokenizer,
+ )
+ ):
+ return None
+ return ModelProfile(
+ name="legacy",
+ architecture="z-image",
+ variant="turbo",
+ checkpoint=models.checkpoint,
+ diffusion_model=models.diffusion_model,
+ text_encoder=models.text_encoder,
+ vae=models.vae,
+ tokenizer=models.tokenizer,
+ )
+
+
def loadUserConfig(path: Path | None = None) -> UserConfig:
"""Load optional defaults from the fixed TOML config file."""
@@ -327,22 +494,45 @@ def loadUserConfig(path: Path | None = None) -> UserConfig:
_rejectUnknownKeys(data, TOP_LEVEL_CONFIG_KEYS, "top-level")
models = _optionalTable(data, "models")
+ model_profiles_table = _optionalTable(data, "model_profiles")
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"),
+ parsed_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"),
+ )
+ named_profiles: dict[str, ModelProfile] = {}
+ for name, table in model_profiles_table.items():
+ if not isinstance(name, str) or not name.strip():
+ raise DiffusionCliError("Model profile name must not be empty")
+ if not isinstance(table, dict):
+ raise DiffusionCliError(
+ f"Config table must be a table: [model_profiles.{name}]"
+ )
+ if name in named_profiles:
+ raise DiffusionCliError(f"Duplicate model profile: {name}")
+ named_profiles[name] = _profileFromTable(name, table)
+ legacy = _legacyProfile(parsed_models)
+ if legacy is not None:
+ named_profiles.setdefault(legacy.name, legacy)
+
+ default_model = data.get("default_model")
+ if default_model is not None and not isinstance(default_model, str):
+ raise DiffusionCliError("Config value default_model must be a string")
+ if default_model is not None and default_model not in named_profiles:
+ raise DiffusionCliError(f"Unknown default_model: {default_model}")
+
+ return UserConfig(
+ models=parsed_models,
generation=GenerationDefaults(
negative_prompt=_optionalString(
generation,
@@ -367,7 +557,12 @@ def loadUserConfig(path: Path | None = None) -> UserConfig:
),
device=_optionalString(generation, "generation", "device"),
dtype=_optionalString(generation, "generation", "dtype"),
+ mu=_optionalFloat(generation, "generation", "mu"),
+ shift_y1=_optionalFloat(generation, "generation", "shift_y1"),
+ shift_y2=_optionalFloat(generation, "generation", "shift_y2"),
),
+ model_profiles=named_profiles,
+ default_model=default_model,
)
@@ -380,6 +575,106 @@ def coalesce(*values):
raise AssertionError("coalesce requires at least one non-None value")
+def legacyUserConfig(user_config: UserConfig) -> UserConfig:
+ """Return a config with the legacy profile materialized when needed."""
+
+ if "legacy" in user_config.model_profiles:
+ return user_config
+ profile = _legacyProfile(user_config.models)
+ if profile is None:
+ return user_config
+ profiles = dict(user_config.model_profiles)
+ profiles[profile.name] = profile
+ return UserConfig(
+ models=user_config.models,
+ generation=user_config.generation,
+ model_profiles=profiles,
+ default_model=user_config.default_model,
+ )
+
+
+def selectModelProfile(
+ user_config: UserConfig,
+ name: str | None = None,
+) -> ModelProfile:
+ """Select one configured model profile using the documented precedence."""
+
+ user_config = legacyUserConfig(user_config)
+ selected_name = name or user_config.default_model
+ if selected_name is None and "legacy" in user_config.model_profiles:
+ selected_name = "legacy"
+ if selected_name is None:
+ raise DiffusionCliError(
+ "No model is configured; set default_model or define a model "
+ "profile"
+ )
+ profile = user_config.model_profiles.get(selected_name)
+ if profile is None:
+ raise DiffusionCliError(f"Unknown model profile: {selected_name}")
+ return profile
+
+
+def profileDefaults(profile: ModelProfile) -> dict[str, object]:
+ """Return architecture and variant generation fallbacks."""
+
+ if profile.architecture == "krea2":
+ if profile.variant == "turbo":
+ return {
+ "width": 1024,
+ "height": 1024,
+ "steps": 8,
+ "cfg": 0.0,
+ "mu": 1.15,
+ }
+ return {
+ "width": 1024,
+ "height": 1024,
+ "steps": 52,
+ "cfg": 3.5,
+ "mu": None,
+ }
+ return {
+ "width": DEFAULT_WIDTH,
+ "height": DEFAULT_HEIGHT,
+ "steps": DEFAULT_STEPS,
+ "cfg": DEFAULT_CFG,
+ "mu": None,
+ }
+
+
+def _profileGenerationValue(
+ request_value,
+ configured_value,
+ profile: ModelProfile,
+ key: str,
+):
+ """Resolve one generation value including profile-specific defaults."""
+
+ if request_value is not None:
+ return request_value
+ if configured_value is not None:
+ return configured_value
+ return profileDefaults(profile)[key]
+
+
+def validateProfileOverrides(
+ profile: ModelProfile,
+ request: ImageGenerationRequest,
+) -> None:
+ """Reject options that a selected backend cannot interpret."""
+
+ if profile.architecture == "z-image":
+ for name, value in (
+ ("mu", request.mu),
+ ("shift-y1", request.shift_y1),
+ ("shift-y2", request.shift_y2),
+ ):
+ if value is not None:
+ raise DiffusionCliError(
+ f"{name} is not supported by the z-image backend"
+ )
+
+
def validateDimensions(width: int, height: int) -> None:
"""Validate that image dimensions are positive latent multiples."""
@@ -513,13 +808,28 @@ def selectDtype(dtype_name: str, device) -> object:
def buildGenerationConfigFromRequest(
request: ImageGenerationRequest,
user_config: UserConfig | None = None,
+ profile: ModelProfile | 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 profile is None:
+ profile = selectModelProfile(user_config)
+ validateProfileOverrides(profile, request)
+ if profile.architecture == "z-image" and any(
+ value is not None
+ for value in (
+ generation.mu,
+ generation.shift_y1,
+ generation.shift_y2,
+ )
+ ):
+ raise DiffusionCliError(
+ "Krea timestep shift options are not supported by the z-image "
+ "backend"
+ )
if not request.prompt:
raise DiffusionCliError("--prompt is required for generation")
@@ -529,15 +839,35 @@ def buildGenerationConfigFromRequest(
generation.negative_prompt,
DEFAULT_NEGATIVE_PROMPT,
)
- width = coalesce(request.width, generation.width, DEFAULT_WIDTH)
- height = coalesce(request.height, generation.height, DEFAULT_HEIGHT)
+ width = _profileGenerationValue(
+ request.width,
+ generation.width,
+ profile,
+ "width",
+ )
+ height = _profileGenerationValue(
+ request.height,
+ generation.height,
+ profile,
+ "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)
+ steps = _profileGenerationValue(
+ request.steps,
+ generation.steps,
+ profile,
+ "steps",
+ )
+ cfg = _profileGenerationValue(
+ request.cfg,
+ generation.cfg,
+ profile,
+ "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)
@@ -551,14 +881,17 @@ def buildGenerationConfigFromRequest(
generation.output_quality,
DEFAULT_OUTPUT_QUALITY,
)
- tokenizer_path = request.tokenizer_path
- if tokenizer_path is None:
- tokenizer_path = models.tokenizer
+ tokenizer_path = request.tokenizer_path or profile.tokenizer
if tokenizer_path is None:
raise DiffusionCliError("Missing models.tokenizer")
- validateDimensions(width, height)
+ if profile.architecture == "z-image":
+ validateDimensions(width, height)
+ elif width <= 0 or height <= 0:
+ raise DiffusionCliError(
+ f"Width and height must be positive: got {width}x{height}"
+ )
if batch_size < 1:
raise DiffusionCliError(
f"Batch size must be at least 1: got {batch_size}"
@@ -576,6 +909,22 @@ def buildGenerationConfigFromRequest(
output_mime_type = outputMimeType(output_extension)
tokenizer_path = validateTokenizerPath(tokenizer_path)
seed = request.seed if request.seed is not None else randomSeed()
+ raw_mu = _profileGenerationValue(
+ request.mu,
+ generation.mu,
+ profile,
+ "mu",
+ )
+ shift_y1 = coalesce(
+ request.shift_y1,
+ generation.shift_y1,
+ DEFAULT_SHIFT_Y1,
+ )
+ shift_y2 = coalesce(
+ request.shift_y2,
+ generation.shift_y2,
+ DEFAULT_SHIFT_Y2,
+ )
return GenerationConfig(
prompt=request.prompt,
@@ -594,15 +943,26 @@ def buildGenerationConfigFromRequest(
output_quality=output_quality,
output_mime_type=output_mime_type,
tokenizer_path=tokenizer_path,
+ model_profile=profile,
+ mu=raw_mu,
+ shift_y1=shift_y1,
+ shift_y2=shift_y2,
)
def buildDefaultGenerationRequest(
user_config: UserConfig,
+ profile: ModelProfile | None = None,
) -> ImageGenerationRequest:
"""Build UI-visible generation defaults from config and fallbacks."""
generation = user_config.generation
+ if profile is None:
+ try:
+ profile = selectModelProfile(user_config)
+ except DiffusionCliError:
+ profile = ModelProfile("legacy", "z-image", "turbo")
+ defaults = profileDefaults(profile)
output_extension = normalizeOutputExtension(
coalesce(generation.output_extension, DEFAULT_OUTPUT_EXTENSION)
)
@@ -614,13 +974,20 @@ def buildDefaultGenerationRequest(
prompt="",
negative_prompt=generation.negative_prompt or "",
seed=-1,
- width=coalesce(generation.width, DEFAULT_WIDTH),
- height=coalesce(generation.height, DEFAULT_HEIGHT),
+ width=coalesce(generation.width, defaults["width"]),
+ height=coalesce(generation.height, defaults["height"]),
batch_size=coalesce(generation.batch_size, DEFAULT_BATCH_SIZE),
- steps=coalesce(generation.steps, DEFAULT_STEPS),
- cfg=coalesce(generation.cfg, DEFAULT_CFG),
+ steps=coalesce(generation.steps, defaults["steps"]),
+ cfg=coalesce(generation.cfg, defaults["cfg"]),
output_extension=output_extension,
output_quality=output_quality,
+ mu=(
+ generation.mu
+ if generation.mu is not None
+ else defaults["mu"]
+ ),
+ shift_y1=coalesce(generation.shift_y1, DEFAULT_SHIFT_Y1),
+ shift_y2=coalesce(generation.shift_y2, DEFAULT_SHIFT_Y2),
)
@@ -646,6 +1013,13 @@ def buildGenerationConfig(
device=args.device,
dtype=args.dtype,
tokenizer_path=args.tokenizer_path,
+ mu=getattr(args, "mu", None),
+ shift_y1=getattr(args, "shift_y1", None),
+ shift_y2=getattr(args, "shift_y2", None),
),
user_config,
+ selectModelProfile(
+ user_config,
+ getattr(args, "model_profile", None),
+ ),
)
diff --git a/diffusion_cli/generation_service.py b/diffusion_cli/generation_service.py
index 01d1b88..b32a2e6 100644
--- a/diffusion_cli/generation_service.py
+++ b/diffusion_cli/generation_service.py
@@ -8,14 +8,20 @@ from pathlib import Path
from threading import Lock
from diffusion_cli.checkpoint import inspectSourceTorchDtype
+from diffusion_cli.backends import ZImageBackend, createBackend
from diffusion_cli.config import (
GenerationConfig,
ImageGenerationRequest,
UserConfig,
buildGenerationConfigFromRequest,
+ selectModelProfile,
)
from diffusion_cli.image_io import encodeImages, saveImages
-from diffusion_cli.paths import resolveModelSourcesFromConfig
+from diffusion_cli.paths import (
+ resolveModelSourcesFromConfig,
+ resolveModelSourcesFromProfile,
+)
+from diffusion_cli.errors import DiffusionCliError
from diffusion_cli.sampling import sampleLatents
from diffusion_cli.text_encoder import ZImageTextEncoder
from diffusion_cli.vae import ZImageVae
@@ -62,15 +68,43 @@ class GenerationService:
user_config: UserConfig,
*,
model_residency: str = "staged",
+ model_profile_name: str | None = None,
+ path_overrides=None,
) -> None:
if model_residency not in MODEL_RESIDENCY_VALUES:
raise ValueError(f"Unknown model residency: {model_residency}")
self.user_config = user_config
self.model_residency = model_residency
+ self.model_profile_name = model_profile_name
+ self.path_overrides = path_overrides
self._lock = Lock()
- self._text_encoder: ZImageTextEncoder | None = None
- self._model: ZImageModel | None = None
- self._vae: ZImageVae | None = None
+ self._backend = None
+ if model_profile_name is not None:
+ selectModelProfile(user_config, model_profile_name)
+
+ @property
+ def modelProfile(self):
+ """Return the selected profile when one is configured."""
+
+ if self.model_profile_name is not None:
+ return selectModelProfile(
+ self.user_config,
+ self.model_profile_name,
+ )
+ try:
+ return selectModelProfile(
+ self.user_config,
+ self.model_profile_name,
+ )
+ except DiffusionCliError:
+ return None
+
+ @property
+ def modelId(self) -> str:
+ """Return the selected backend's client-visible model identifier."""
+
+ profile = self.modelProfile
+ return profile.name if profile is not None else "z-image-local"
def generateImages(
self,
@@ -79,9 +113,11 @@ class GenerationService:
"""Generate final-format image bytes for one request."""
with self._lock:
+ profile = self.modelProfile
config = buildGenerationConfigFromRequest(
request,
self.user_config,
+ profile,
)
images = self._generateTensor(config)
encoded_images = encodeImages(
@@ -102,9 +138,11 @@ class GenerationService:
"""Generate images and write them to the configured output path."""
with self._lock:
+ profile = self.modelProfile
config = buildGenerationConfigFromRequest(
request,
self.user_config,
+ profile,
)
images = self._generateTensor(config)
return saveImages(
@@ -115,111 +153,43 @@ class GenerationService:
)
def _generateTensor(self, config: GenerationConfig):
- if self.model_residency == "cpu-cache":
- return self._generateCpuCache(config)
- return self._generateStaged(config)
-
- def _generateStaged(self, config: GenerationConfig):
- model_sources = resolveModelSourcesFromConfig(self.user_config.models)
- 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_sources.text_encoder,
- config.tokenizer_path,
- config.device,
- text_dtype,
- )
- conditioning = text_encoder.encodePrompts(
- config.prompt,
- config.negative_prompt,
- )
- del text_encoder
- releaseMemory()
-
- model = ZImageModel(
- model_sources.diffusion_model,
- config.device,
- diffusion_dtype,
- )
- latent = sampleLatents(
- model,
- conditioning,
- batch_size=config.batch_size,
- height=config.height,
- width=config.width,
- seed=config.seed,
- steps=config.steps,
- cfg=config.cfg,
- device=config.device,
- dtype=diffusion_dtype,
- )
- del model
- releaseMemory()
-
- vae = ZImageVae(model_sources.vae, config.device, vae_dtype)
- images = vae.decode(latent)
- del vae
- releaseMemory()
- return images
-
- def _generateCpuCache(self, config: GenerationConfig):
- import torch
-
- model_sources = resolveModelSourcesFromConfig(self.user_config.models)
- text_dtype = componentDtype(config, model_sources.text_encoder)
- diffusion_dtype = componentDtype(config, model_sources.diffusion_model)
- vae_dtype = componentDtype(config, model_sources.vae)
- cpu = torch.device("cpu")
-
- if self._text_encoder is None:
- self._text_encoder = ZImageTextEncoder(
- model_sources.text_encoder,
- config.tokenizer_path,
- cpu,
- text_dtype,
- )
- self._text_encoder.toDevice(config.device, text_dtype)
- try:
- conditioning = self._text_encoder.encodePrompts(
- config.prompt,
- config.negative_prompt,
- )
- finally:
- self._text_encoder.toCpu()
- releaseMemory()
-
- if self._model is None:
- self._model = ZImageModel(
- model_sources.diffusion_model,
- cpu,
- diffusion_dtype,
+ profile = getattr(config, "model_profile", None) or self.modelProfile
+ if self._backend is None:
+ if profile is None:
+ # This compatibility path is used by older callers and tests
+ # that inject the source resolver directly.
+ sources = resolveModelSourcesFromConfig(self.user_config.models)
+ else:
+ if self.path_overrides is not None:
+ from diffusion_cli.paths import resolveModelSourcesForProfile
+
+ sources = resolveModelSourcesForProfile(
+ profile,
+ args=self.path_overrides,
+ )
+ else:
+ sources = resolveModelSourcesFromProfile(
+ self.user_config,
+ profile.name,
+ )
+ kwargs = {}
+ if profile is None or profile.architecture == "z-image":
+ kwargs = {
+ "text_encoder_factory": ZImageTextEncoder,
+ "model_factory": ZImageModel,
+ "vae_factory": ZImageVae,
+ "sampler": sampleLatents,
+ "component_dtype": componentDtype,
+ }
+ self._backend = createBackend(
+ profile
+ or type("LegacyProfile", (), {
+ "name": "legacy",
+ "architecture": "z-image",
+ "variant": "turbo",
+ })(),
+ sources,
+ self.model_residency,
+ **kwargs,
)
- self._model.toDevice(config.device, diffusion_dtype)
- try:
- latent = sampleLatents(
- self._model,
- conditioning,
- batch_size=config.batch_size,
- height=config.height,
- width=config.width,
- seed=config.seed,
- steps=config.steps,
- cfg=config.cfg,
- device=config.device,
- dtype=diffusion_dtype,
- )
- finally:
- self._model.toCpu()
- releaseMemory()
-
- if self._vae is None:
- self._vae = ZImageVae(model_sources.vae, cpu, vae_dtype)
- self._vae.toDevice(config.device, vae_dtype)
- try:
- images = self._vae.decode(latent)
- finally:
- self._vae.toCpu()
- releaseMemory()
- return images
+ return self._backend.generate(config)
diff --git a/diffusion_cli/krea2_backend.py b/diffusion_cli/krea2_backend.py
new file mode 100644
index 0000000..b4a9060
--- /dev/null
+++ b/diffusion_cli/krea2_backend.py
@@ -0,0 +1,7 @@
+"""Public Krea 2 backend module."""
+
+# The implementation lives in ``backends`` so the registry and interface
+# remain colocated. Keep this module as the stable family-specific import.
+from diffusion_cli.backends import Krea2Backend
+
+__all__ = ["Krea2Backend"]
diff --git a/diffusion_cli/krea2_model.py b/diffusion_cli/krea2_model.py
new file mode 100644
index 0000000..1cee831
--- /dev/null
+++ b/diffusion_cli/krea2_model.py
@@ -0,0 +1,840 @@
+"""Krea 2 single-stream diffusion model.
+
+The architecture is adapted from Krea's Apache-2.0 reference implementation
+in ``~/programs/krea-2/mmdit.py``. It is modified to use local checkpoint
+loading, strict state-key accounting, and the local scaled-FP8 layer.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+import math
+import re
+from typing import Iterable
+
+import torch
+from torch import nn
+from torch.nn import functional as F
+
+from diffusion_cli.checkpoint import loadStateDict
+from diffusion_cli.config import ModelSource
+from diffusion_cli.errors import DiffusionCliError
+from diffusion_cli.quantization import (
+ QuantizationManifest,
+ ScaledFp8Linear,
+ manifestFromStateDict,
+ loadQuantizationManifest,
+ moveModulePreservingQuantization,
+ validateQuantizedStateDict,
+)
+
+KREA2_FEATURES = 6144
+KREA2_TIMESTEP_WIDTH = 256
+KREA2_TEXT_WIDTH = 2560
+KREA2_HEADS = 48
+KREA2_KV_HEADS = 12
+KREA2_BLOCKS = 28
+KREA2_MLP_MULTIPLIER = 4
+KREA2_PATCH_SIZE = 2
+KREA2_LATENT_CHANNELS = 16
+KREA2_TEXT_LAYERS = 12
+KREA2_TEXT_HEADS = 20
+KREA2_TEXT_KV_HEADS = 20
+KREA2_ROPE_THETA = 1000.0
+
+
+@dataclass(frozen=True)
+class Krea2Conditioning:
+ """Twelve Qwen hidden-state taps and their token mask."""
+
+ hidden_states: torch.Tensor
+ attention_mask: torch.Tensor
+
+
+@dataclass(frozen=True)
+class Krea2Config:
+ """Architecture settings for one Krea 2 diffusion model."""
+
+ features: int = KREA2_FEATURES
+ timestep_width: int = KREA2_TIMESTEP_WIDTH
+ text_width: int = KREA2_TEXT_WIDTH
+ heads: int = KREA2_HEADS
+ kv_heads: int = KREA2_KV_HEADS
+ blocks: int = KREA2_BLOCKS
+ mlp_multiplier: int = KREA2_MLP_MULTIPLIER
+ patch_size: int = KREA2_PATCH_SIZE
+ latent_channels: int = KREA2_LATENT_CHANNELS
+ text_layers: int = KREA2_TEXT_LAYERS
+ text_heads: int = KREA2_TEXT_HEADS
+ text_kv_heads: int = KREA2_TEXT_KV_HEADS
+ rope_theta: float = KREA2_ROPE_THETA
+ bias: bool = False
+
+
+def _finite(value: torch.Tensor, label: str) -> None:
+ """Reject non-finite tensors at a model boundary."""
+
+ if not torch.isfinite(value).all():
+ raise DiffusionCliError(f"Krea diffusion {label} contains NaN or Inf")
+
+
+class KreaRmsNorm(nn.Module):
+ """RMS normalization with the reference implementation's additive scale."""
+
+ def __init__(self, features: int, *, device=None, dtype=torch.float32):
+ super().__init__()
+ self.features = features
+ self.eps = 1e-5
+ self.scale = nn.Parameter(
+ torch.zeros(features, device=device, dtype=dtype),
+ requires_grad=False,
+ )
+
+ def forward(self, value: torch.Tensor) -> torch.Tensor:
+ """Normalize the last feature dimension."""
+
+ original_dtype = value.dtype
+ normalized = F.rms_norm(
+ value.float(),
+ (self.features,),
+ eps=self.eps,
+ weight=self.scale.float() + 1.0,
+ )
+ return normalized.to(original_dtype)
+
+
+class KreaAttention(nn.Module):
+ """Grouped-query attention used by text fusion and image blocks."""
+
+ def __init__(
+ self,
+ features: int,
+ heads: int,
+ kv_heads: int,
+ *,
+ bias: bool,
+ device=None,
+ dtype=None,
+ ) -> None:
+ super().__init__()
+ if features % heads != 0 or heads % kv_heads != 0:
+ raise ValueError("Krea attention heads must divide evenly")
+ self.heads = heads
+ self.kv_heads = kv_heads
+ self.head_dim = features // heads
+ self.wq = nn.Linear(
+ features,
+ features,
+ bias=bias,
+ device=device,
+ dtype=dtype,
+ )
+ self.wk = nn.Linear(
+ features,
+ self.head_dim * kv_heads,
+ bias=bias,
+ device=device,
+ dtype=dtype,
+ )
+ self.wv = nn.Linear(
+ features,
+ self.head_dim * kv_heads,
+ bias=bias,
+ device=device,
+ dtype=dtype,
+ )
+ self.gate = nn.Linear(
+ features,
+ features,
+ bias=bias,
+ device=device,
+ dtype=dtype,
+ )
+ self.qknorm = KreaQkNorm(
+ self.head_dim,
+ device=device,
+ )
+ self.wo = nn.Linear(
+ features,
+ features,
+ bias=bias,
+ device=device,
+ dtype=dtype,
+ )
+
+ def forward(
+ self,
+ value: torch.Tensor,
+ frequencies: tuple[torch.Tensor, torch.Tensor] | None = None,
+ mask: torch.Tensor | None = None,
+ ) -> torch.Tensor:
+ """Apply normalized, optionally rotary, grouped-query attention."""
+
+ batch, sequence, _ = value.shape
+ query = self.wq(value).reshape(
+ batch, sequence, self.heads, self.head_dim
+ ).transpose(1, 2)
+ key = self.wk(value).reshape(
+ batch, sequence, self.kv_heads, self.head_dim
+ ).transpose(1, 2)
+ val = self.wv(value).reshape(
+ batch, sequence, self.kv_heads, self.head_dim
+ ).transpose(1, 2)
+ query, key = self.qknorm(query, key)
+ if frequencies is not None:
+ query = _applyRotary(query, frequencies)
+ key = _applyRotary(key, frequencies)
+ if self.heads != self.kv_heads:
+ repeat = self.heads // self.kv_heads
+ key = key.repeat_interleave(repeat, dim=1)
+ val = val.repeat_interleave(repeat, dim=1)
+ attention_mask = _attentionMask(mask)
+ attended = F.scaled_dot_product_attention(
+ query,
+ key,
+ val,
+ attn_mask=attention_mask,
+ )
+ attended = attended.transpose(1, 2).reshape(batch, sequence, -1)
+ return self.wo(attended * torch.sigmoid(self.gate(value)))
+
+
+class KreaQkNorm(nn.Module):
+ """Apply independent RMS normalization to query and key heads."""
+
+ def __init__(self, head_dim: int, *, device=None):
+ super().__init__()
+ self.qnorm = KreaRmsNorm(head_dim, device=device)
+ self.knorm = KreaRmsNorm(head_dim, device=device)
+
+ def forward(
+ self,
+ query: torch.Tensor,
+ key: torch.Tensor,
+ ) -> tuple[torch.Tensor, torch.Tensor]:
+ """Normalize query and key while preserving their shapes."""
+
+ return self.qnorm(query), self.knorm(key)
+
+
+class KreaSwiGlu(nn.Module):
+ """SwiGLU MLP with the Krea checkpoint's dimensions."""
+
+ def __init__(
+ self,
+ features: int,
+ multiplier: int,
+ *,
+ bias: bool,
+ device=None,
+ dtype=None,
+ ) -> None:
+ super().__init__()
+ hidden = int(2 * features / 3) * multiplier
+ hidden = 128 * ((hidden + 127) // 128)
+ self.gate = nn.Linear(features, hidden, bias=bias, device=device, dtype=dtype)
+ self.up = nn.Linear(features, hidden, bias=bias, device=device, dtype=dtype)
+ self.down = nn.Linear(hidden, features, bias=bias, device=device, dtype=dtype)
+
+ def forward(self, value: torch.Tensor) -> torch.Tensor:
+ """Apply gated SiLU projection."""
+
+ return self.down(F.silu(self.gate(value)) * self.up(value))
+
+
+class KreaSimpleModulation(nn.Module):
+ """Final two-way AdaLN modulation."""
+
+ def __init__(self, features: int, *, device=None):
+ super().__init__()
+ self.lin = nn.Parameter(
+ torch.zeros(2, features, device=device, dtype=torch.float32),
+ requires_grad=False,
+ )
+
+ def forward(self, value: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+ """Return final scale and shift."""
+
+ value = value[:, None, :] + self.lin.to(value.dtype)[None, :, :]
+ scale, shift = value.unbind(dim=1)
+ return scale, shift
+
+
+class KreaDoubleModulation(nn.Module):
+ """Six-way modulation shared by one single-stream block."""
+
+ def __init__(self, features: int, *, device=None):
+ super().__init__()
+ self.lin = nn.Parameter(
+ torch.zeros(6 * features, device=device, dtype=torch.float32),
+ requires_grad=False,
+ )
+
+ def forward(self, value: torch.Tensor) -> tuple[torch.Tensor, ...]:
+ """Return pre/post scale, shift, and residual gates."""
+
+ return tuple((value + self.lin.to(value.dtype)).chunk(6, dim=-1))
+
+
+class KreaTextFusionBlock(nn.Module):
+ """One attention and MLP block in the text fusion adapter."""
+
+ def __init__(self, config: Krea2Config, *, device=None, dtype=None):
+ super().__init__()
+ self.prenorm = KreaRmsNorm(config.text_width, device=device)
+ self.postnorm = KreaRmsNorm(config.text_width, device=device)
+ self.attn = KreaAttention(
+ config.text_width,
+ config.text_heads,
+ config.text_kv_heads,
+ bias=config.bias,
+ device=device,
+ dtype=dtype,
+ )
+ self.mlp = KreaSwiGlu(
+ config.text_width,
+ config.mlp_multiplier,
+ bias=config.bias,
+ device=device,
+ dtype=dtype,
+ )
+
+ def forward(self, value: torch.Tensor, mask=None) -> torch.Tensor:
+ """Apply a residual text-fusion block."""
+
+ value = value + self.attn(self.prenorm(value), mask=mask)
+ return value + self.mlp(self.postnorm(value))
+
+
+class KreaTextFusion(nn.Module):
+ """Fuse twelve Qwen layer taps into one text sequence."""
+
+ def __init__(self, config: Krea2Config, *, device=None, dtype=None):
+ super().__init__()
+ self.layerwise_blocks = nn.ModuleList(
+ [
+ KreaTextFusionBlock(config, device=device, dtype=dtype)
+ for _ in range(2)
+ ]
+ )
+ self.projector = nn.Linear(
+ config.text_layers,
+ 1,
+ bias=False,
+ device=device,
+ dtype=dtype,
+ )
+ self.refiner_blocks = nn.ModuleList(
+ [
+ KreaTextFusionBlock(config, device=device, dtype=dtype)
+ for _ in range(2)
+ ]
+ )
+
+ def forward(self, value: torch.Tensor, mask=None) -> torch.Tensor:
+ """Project the layer axis and refine the resulting text sequence."""
+
+ batch, sequence, layers, width = value.shape
+ value = value.reshape(batch * sequence, layers, width)
+ for block in self.layerwise_blocks:
+ value = block(value)
+ value = value.reshape(batch, sequence, layers, width).permute(0, 1, 3, 2)
+ value = self.projector(value).squeeze(-1)
+ for block in self.refiner_blocks:
+ value = block(value, mask=mask)
+ return value
+
+
+class KreaSingleBlock(nn.Module):
+ """One single-stream MMDiT block."""
+
+ def __init__(self, config: Krea2Config, *, device=None, dtype=None):
+ super().__init__()
+ self.mod = KreaDoubleModulation(config.features, device=device)
+ self.prenorm = KreaRmsNorm(config.features, device=device)
+ self.postnorm = KreaRmsNorm(config.features, device=device)
+ self.attn = KreaAttention(
+ config.features,
+ config.heads,
+ config.kv_heads,
+ bias=config.bias,
+ device=device,
+ dtype=dtype,
+ )
+ self.mlp = KreaSwiGlu(
+ config.features,
+ config.mlp_multiplier,
+ bias=config.bias,
+ device=device,
+ dtype=dtype,
+ )
+
+ def forward(self, value, modulation, frequencies, mask=None):
+ """Apply modulated attention and MLP residuals."""
+
+ pre_scale, pre_shift, pre_gate, post_scale, post_shift, post_gate = (
+ self.mod(modulation)
+ )
+ attention_input = (1 + pre_scale) * self.prenorm(value) + pre_shift
+ value = value + pre_gate * self.attn(
+ attention_input,
+ frequencies,
+ mask,
+ )
+ mlp_input = (1 + post_scale) * self.postnorm(value) + post_shift
+ return value + post_gate * self.mlp(mlp_input)
+
+
+class KreaLastLayer(nn.Module):
+ """Final normalized patch projection with checkpoint-compatible names."""
+
+ def __init__(self, config: Krea2Config, *, device=None, dtype=None):
+ super().__init__()
+ self.norm = KreaRmsNorm(config.features, device=device)
+ self.linear = nn.Linear(
+ config.features,
+ config.patch_size ** 2 * config.latent_channels,
+ bias=True,
+ device=device,
+ dtype=dtype,
+ )
+ self.modulation = KreaSimpleModulation(
+ config.features,
+ device=device,
+ )
+
+ def forward(self, value: torch.Tensor, timestep_value: torch.Tensor):
+ """Apply final AdaLN modulation and patch projection."""
+
+ scale, shift = self.modulation(timestep_value)
+ return self.linear((1 + scale[:, None, :]) * self.norm(value) + shift[:, None, :])
+
+
+def _attentionMask(mask: torch.Tensor | None) -> torch.Tensor | None:
+ """Expand a token mask into the boolean mask expected by SDPA."""
+
+ if mask is None:
+ return None
+ return mask[:, None, :, None] & mask[:, None, None, :]
+
+
+def _rotaryFrequencies(
+ positions: torch.Tensor,
+ axis_dimensions: Iterable[int],
+ theta: float,
+) -> tuple[torch.Tensor, torch.Tensor]:
+ """Build three-axis rotary frequencies for image and text positions."""
+
+ cos_parts = []
+ sin_parts = []
+ for axis, dimension in enumerate(axis_dimensions):
+ frequencies = torch.arange(
+ 0,
+ dimension,
+ 2,
+ device=positions.device,
+ dtype=torch.float32,
+ )
+ frequencies = 1.0 / theta ** (frequencies / dimension)
+ angles = positions[..., axis, None].float() * frequencies
+ cos_parts.append(torch.cos(angles))
+ sin_parts.append(torch.sin(angles))
+ return torch.cat(cos_parts, dim=-1), torch.cat(sin_parts, dim=-1)
+
+
+def _applyRotary(
+ value: torch.Tensor,
+ frequencies: tuple[torch.Tensor, torch.Tensor],
+) -> torch.Tensor:
+ """Apply pairwise rotary rotation to a head-major tensor."""
+
+ cos, sin = frequencies
+ original_dtype = value.dtype
+ value = value.float().reshape(*value.shape[:-1], -1, 2)
+ cos = cos[:, None, :, :].unsqueeze(-1)
+ sin = sin[:, None, :, :].unsqueeze(-1)
+ first, second = value.unbind(dim=-1)
+ rotated = torch.stack((first * cos[..., 0] - second * sin[..., 0],
+ first * sin[..., 0] + second * cos[..., 0]), dim=-1)
+ return rotated.reshape(*rotated.shape[:-2], -1).to(original_dtype)
+
+
+def _patchify(value: torch.Tensor, patch: int) -> torch.Tensor:
+ """Convert an NCHW latent into row-major patch tokens."""
+
+ batch, channels, height, width = value.shape
+ return value.reshape(
+ batch,
+ channels,
+ height // patch,
+ patch,
+ width // patch,
+ patch,
+ ).permute(0, 2, 4, 1, 3, 5).reshape(
+ batch,
+ (height // patch) * (width // patch),
+ channels * patch * patch,
+ )
+
+
+def _unpatchify(value: torch.Tensor, height: int, width: int, patch: int, channels: int):
+ """Convert row-major patch tokens back to an NCHW latent."""
+
+ batch = value.shape[0]
+ return value.reshape(
+ batch,
+ height // patch,
+ width // patch,
+ channels,
+ patch,
+ patch,
+ ).permute(0, 3, 1, 4, 2, 5).reshape(batch, channels, height, width)
+
+
+def patchify(value: torch.Tensor, patch: int = KREA2_PATCH_SIZE) -> torch.Tensor:
+ """Convert an NCHW latent into row-major Krea patch tokens."""
+
+ return _patchify(value, patch)
+
+
+def unpatchify(
+ value: torch.Tensor,
+ height: int,
+ width: int,
+ channels: int = KREA2_LATENT_CHANNELS,
+ patch: int = KREA2_PATCH_SIZE,
+) -> torch.Tensor:
+ """Convert row-major Krea patch tokens back to an NCHW latent."""
+
+ return _unpatchify(value, height, width, patch, channels)
+
+
+class Krea2Model(nn.Module):
+ """Krea 2 MMDiT wrapper accepting NCHW latents and conditioning."""
+
+ def __init__(
+ self,
+ config: Krea2Config = Krea2Config(),
+ *,
+ device=None,
+ dtype=torch.bfloat16,
+ quantization_manifest: QuantizationManifest | None = None,
+ ) -> None:
+ super().__init__()
+ self.config = config
+ self._manifest = quantization_manifest or QuantizationManifest({})
+ self._device = device
+ self.first = nn.Linear(
+ config.latent_channels * config.patch_size ** 2,
+ config.features,
+ bias=True,
+ device=device,
+ dtype=dtype,
+ )
+ self.blocks = nn.ModuleList(
+ [
+ KreaSingleBlock(config, device=device, dtype=dtype)
+ for _ in range(config.blocks)
+ ]
+ )
+ self.tmlp = nn.Sequential(
+ nn.Linear(config.timestep_width, config.features, device=device, dtype=dtype),
+ nn.GELU(approximate="tanh"),
+ nn.Linear(config.features, config.features, device=device, dtype=dtype),
+ )
+ self.txtfusion = KreaTextFusion(config, device=device, dtype=dtype)
+ self.txtmlp = nn.Sequential(
+ KreaRmsNorm(config.text_width, device=device),
+ nn.Linear(config.text_width, config.features, device=device, dtype=dtype),
+ nn.GELU(approximate="tanh"),
+ nn.Linear(config.features, config.features, device=device, dtype=dtype),
+ )
+ self.tproj = nn.Sequential(
+ nn.GELU(approximate="tanh"),
+ nn.Linear(config.features, config.features * 6, device=device, dtype=dtype),
+ )
+ self.last = KreaLastLayer(config, device=device, dtype=dtype)
+
+ def _replaceQuantizedLinear(self) -> None:
+ """Replace manifest-selected ordinary linears with scaled FP8 linears."""
+
+ for name in self._manifest.layers:
+ module = self.get_submodule(name)
+ if not isinstance(module, nn.Linear):
+ raise DiffusionCliError(
+ f"Quantized Krea layer is not a linear module: {name}"
+ )
+ replacement = ScaledFp8Linear(
+ module.in_features,
+ module.out_features,
+ bias=module.bias is not None,
+ device=module.weight.device,
+ dtype=module.bias.dtype if module.bias is not None else torch.bfloat16,
+ )
+ parent_name, _, child_name = name.rpartition(".")
+ parent = self.get_submodule(parent_name) if parent_name else self
+ setattr(parent, child_name, replacement)
+
+ def forward(
+ self,
+ latent: torch.Tensor,
+ timestep: torch.Tensor,
+ conditioning: Krea2Conditioning,
+ ) -> torch.Tensor:
+ """Return flow velocity with the same NCHW shape as ``latent``."""
+
+ config = self.config
+ if latent.ndim != 4 or latent.shape[1] != config.latent_channels:
+ raise DiffusionCliError(
+ f"Krea latent must be NCHW with {config.latent_channels} channels"
+ )
+ if latent.shape[2] % config.patch_size or latent.shape[3] % config.patch_size:
+ raise DiffusionCliError("Krea latent dimensions must be patch aligned")
+ hidden = conditioning.hidden_states
+ mask = conditioning.attention_mask
+ if hidden.ndim != 4:
+ raise DiffusionCliError("Krea conditioning must have four dimensions")
+ if hidden.shape[2] != config.text_layers or hidden.shape[3] != config.text_width:
+ raise DiffusionCliError(
+ "Krea conditioning must contain 12 layers of 2560 features"
+ )
+ if mask.ndim != 2 or mask.shape[:2] != hidden.shape[:2]:
+ raise DiffusionCliError("Krea conditioning mask shape does not match")
+ if latent.shape[0] != hidden.shape[0]:
+ raise DiffusionCliError("Krea latent and conditioning batch mismatch")
+ timestep = timestep.reshape(-1)
+ if timestep.numel() == 1:
+ timestep = timestep.expand(latent.shape[0])
+ if timestep.numel() != latent.shape[0]:
+ raise DiffusionCliError("Krea timestep and latent batch mismatch")
+ hidden = hidden.to(device=latent.device)
+ mask = mask.to(device=latent.device, dtype=torch.bool)
+ _finite(latent, "input")
+ _finite(hidden, "conditioning")
+
+ image = self.first(_patchify(latent, config.patch_size))
+ time_embedding = _timestepEmbedding(
+ timestep,
+ config.timestep_width,
+ image.dtype,
+ )
+ modulation = self.tproj(self.tmlp(time_embedding))
+ text = self.txtfusion(hidden.to(image.dtype), mask=mask)
+ text = self.txtmlp(text)
+
+ batch, _, height, width = latent.shape
+ image_height = height // config.patch_size
+ image_width = width // config.patch_size
+ text_length = text.shape[1]
+ text_positions = torch.zeros(
+ batch,
+ text_length,
+ 3,
+ device=latent.device,
+ dtype=torch.float32,
+ )
+ image_positions = torch.zeros(
+ batch,
+ image_height,
+ image_width,
+ 3,
+ device=latent.device,
+ dtype=torch.float32,
+ )
+ image_positions[..., 1] = torch.arange(
+ image_height,
+ device=latent.device,
+ )[None, :, None]
+ image_positions[..., 2] = torch.arange(
+ image_width,
+ device=latent.device,
+ )[None, None, :]
+ image_positions = image_positions.reshape(batch, -1, 3)
+ value = torch.cat((text, image), dim=1)
+ positions = torch.cat((text_positions, image_positions), dim=1)
+ mask = torch.cat(
+ (
+ mask,
+ torch.ones(batch, image.shape[1], device=latent.device, dtype=torch.bool),
+ ),
+ dim=1,
+ )
+ pad = (-value.shape[1]) % 256
+ if pad:
+ value = F.pad(value, (0, 0, 0, pad))
+ positions = F.pad(positions, (0, 0, 0, pad))
+ mask = F.pad(mask, (0, pad), value=False)
+ axes = _axisDimensions(config)
+ frequencies = _rotaryFrequencies(positions, axes, config.rope_theta)
+ for block in self.blocks:
+ value = block(value, modulation, frequencies, mask)
+ value = value[:, text_length : text_length + image.shape[1]]
+ output = self.last(value, self.tmlp(time_embedding))
+ output = _unpatchify(
+ output,
+ height,
+ width,
+ config.patch_size,
+ config.latent_channels,
+ )
+ _finite(output, "output")
+ return output
+
+ def toDevice(self, device, dtype=None) -> None:
+ """Move model components while preserving quantized storage dtypes."""
+
+ moveModulePreservingQuantization(
+ self,
+ device,
+ dtype or torch.bfloat16,
+ )
+ self._device = device
+
+ def toCpu(self) -> None:
+ """Move the model back to CPU without expanding FP8 weights."""
+
+ self.toDevice(torch.device("cpu"))
+
+
+def _axisDimensions(config: Krea2Config) -> tuple[int, int, int]:
+ """Return the three RoPE axis widths from the official layout."""
+
+ head_dim = config.features // config.heads
+ unit = head_dim // 16
+ return head_dim - 12 * unit, 6 * unit, 6 * unit
+
+
+def _timestepEmbedding(
+ timestep: torch.Tensor,
+ width: int,
+ dtype: torch.dtype,
+) -> torch.Tensor:
+ """Build the official sinusoidal timestep embedding."""
+
+ timestep = timestep.reshape(-1)
+ half = width // 2
+ frequencies = torch.exp(
+ -math.log(10000.0)
+ * torch.arange(half, device=timestep.device, dtype=torch.float32)
+ / half
+ )
+ angles = timestep.float()[:, None] * 1000.0 * frequencies[None]
+ return torch.cat((torch.cos(angles), torch.sin(angles)), dim=-1).to(dtype)
+
+
+def normalizeKreaStateDict(
+ state_dict: dict[str, torch.Tensor],
+) -> dict[str, torch.Tensor]:
+ """Strip known ComfyUI diffusion prefixes and metadata tensors."""
+
+ normalized = {}
+ for key, value in state_dict.items():
+ if key.endswith(".comfy_quant"):
+ continue
+ for prefix in ("model.diffusion_model.", "diffusion_model.", "model."):
+ if key.startswith(prefix):
+ key = key.removeprefix(prefix)
+ break
+ normalized[key] = value
+ return normalized
+
+
+def detectKrea2Config(state_dict: dict[str, torch.Tensor]) -> Krea2Config:
+ """Validate supported Krea checkpoint shapes and infer block count."""
+
+ state_dict = normalizeKreaStateDict(state_dict)
+ try:
+ first = state_dict["first.weight"]
+ text_projector = state_dict["txtfusion.projector.weight"]
+ except KeyError as exc:
+ raise DiffusionCliError(
+ "Checkpoint is not a Krea 2 diffusion model"
+ ) from exc
+ blocks = sorted(
+ {
+ int(match.group(1))
+ for key in state_dict
+ if (match := re.match(r"blocks\.(\d+)\.", key))
+ }
+ )
+ if (
+ text_projector.ndim != 2
+ or tuple(text_projector.shape) != (1, KREA2_TEXT_LAYERS)
+ ):
+ raise DiffusionCliError("Unsupported Krea text-fusion projector shape")
+ config = Krea2Config(
+ features=first.shape[0],
+ latent_channels=KREA2_LATENT_CHANNELS,
+ patch_size=KREA2_PATCH_SIZE,
+ text_layers=text_projector.shape[1],
+ blocks=len(blocks),
+ )
+ if first.shape[1] != config.latent_channels * config.patch_size ** 2:
+ raise DiffusionCliError("Unsupported Krea latent patch shape")
+ if blocks != list(range(KREA2_BLOCKS)):
+ raise DiffusionCliError(
+ "Unsupported Krea block numbering: expected blocks 0 through 27"
+ )
+ if config.features != KREA2_FEATURES or config.blocks != KREA2_BLOCKS:
+ raise DiffusionCliError(
+ "Unsupported Krea 2 architecture constants: expected 6144 features "
+ "and 28 blocks"
+ )
+ return config
+
+
+def loadKrea2StateDict(
+ model: Krea2Model,
+ state_dict: dict[str, torch.Tensor],
+ manifest: QuantizationManifest | None = None,
+) -> Krea2Model:
+ """Strictly load a Krea state dict into an existing, including reduced, model."""
+
+ normalized = normalizeKreaStateDict(state_dict)
+ if manifest is None:
+ manifest = manifestFromStateDict(normalized)
+ manifest = validateQuantizedStateDict(normalized, manifest)
+ model._manifest = manifest
+ model._replaceQuantizedLinear()
+ missing, unexpected = model.load_state_dict(
+ normalized,
+ strict=False,
+ assign=True,
+ )
+ if missing or unexpected:
+ raise DiffusionCliError(
+ "Krea diffusion checkpoint key mismatch: "
+ f"missing={list(missing)[:3]}, unexpected={list(unexpected)[:3]}"
+ )
+ return model
+
+
+def buildKrea2Model(
+ source: ModelSource,
+ *,
+ device=None,
+ dtype=torch.bfloat16,
+) -> Krea2Model:
+ """Load a local Krea diffusion checkpoint with strict key accounting."""
+
+ state_dict = loadStateDict(source)
+ manifest = validateQuantizedStateDict(
+ normalizeKreaStateDict(state_dict),
+ loadQuantizationManifest(source),
+ )
+ config = detectKrea2Config(state_dict)
+ with torch.device("meta"):
+ model = Krea2Model(
+ config,
+ device="meta",
+ dtype=dtype,
+ quantization_manifest=manifest,
+ )
+ loadKrea2StateDict(model, state_dict, manifest)
+ model.toDevice(device or torch.device("cpu"), dtype)
+ model.eval()
+ return model
+
+
+# Descriptive aliases used by integration callers.
+loadKrea2Model = buildKrea2Model
diff --git a/diffusion_cli/krea2_sampling.py b/diffusion_cli/krea2_sampling.py
new file mode 100644
index 0000000..b2a6c6b
--- /dev/null
+++ b/diffusion_cli/krea2_sampling.py
@@ -0,0 +1,281 @@
+"""Flow-matching sampling for Krea 2 Turbo and Raw checkpoints."""
+
+from __future__ import annotations
+
+import math
+import warnings
+
+import torch
+
+from diffusion_cli.errors import DiffusionCliError
+from diffusion_cli.krea2_model import Krea2Conditioning
+
+KREA2_LATENT_DOWNSCALE = 8
+KREA2_PATCH_SIZE = 2
+
+
+def roundUp(value: int, multiple: int, name: str) -> int:
+ """Round a dimension up and issue the required visible warning."""
+
+ if value <= 0:
+ raise DiffusionCliError(f"{name} must be positive: got {value}")
+ aligned = ((value + multiple - 1) // multiple) * multiple
+ if aligned != value:
+ warnings.warn(
+ f"Krea {name}={value} is not a multiple of {multiple}; "
+ f"rounding up to {aligned}",
+ UserWarning,
+ stacklevel=2,
+ )
+ return aligned
+
+
+def prepareConditioning(
+ latent: torch.Tensor,
+ conditioning: Krea2Conditioning,
+ patch_size: int = KREA2_PATCH_SIZE,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+ """Patchify a latent and build image/text positions and masks."""
+
+ if latent.ndim != 4:
+ raise DiffusionCliError("Krea sampler expects an NCHW latent")
+ batch, channels, height, width = latent.shape
+ if height % patch_size or width % patch_size:
+ raise DiffusionCliError("Krea latent is not patch aligned")
+ image_height = height // patch_size
+ image_width = width // patch_size
+ image = latent.reshape(
+ batch,
+ channels,
+ image_height,
+ patch_size,
+ image_width,
+ patch_size,
+ ).permute(0, 2, 4, 1, 3, 5).reshape(
+ batch,
+ image_height * image_width,
+ channels * patch_size * patch_size,
+ )
+ image_positions = torch.zeros(
+ batch,
+ image_height,
+ image_width,
+ 3,
+ device=latent.device,
+ dtype=torch.float32,
+ )
+ image_positions[..., 1] = torch.arange(
+ image_height,
+ device=latent.device,
+ )[None, :, None]
+ image_positions[..., 2] = torch.arange(
+ image_width,
+ device=latent.device,
+ )[None, None, :]
+ text_positions = torch.zeros(
+ batch,
+ conditioning.hidden_states.shape[1],
+ 3,
+ device=latent.device,
+ dtype=torch.float32,
+ )
+ positions = torch.cat(
+ (text_positions, image_positions.reshape(batch, -1, 3)),
+ dim=1,
+ )
+ mask = torch.cat(
+ (
+ conditioning.attention_mask.to(device=latent.device, dtype=torch.bool),
+ torch.ones(
+ batch,
+ image.shape[1],
+ device=latent.device,
+ dtype=torch.bool,
+ ),
+ ),
+ dim=1,
+ )
+ return image, positions, mask
+
+
+def timestepMu(
+ image_token_count: int,
+ *,
+ min_resolution: int = 256,
+ max_resolution: int = 1280,
+ y1: float = 0.5,
+ y2: float = 1.15,
+) -> float:
+ """Interpolate the Raw timestep shift from image-token resolution."""
+
+ minimum = (min_resolution // KREA2_LATENT_DOWNSCALE) ** 2
+ maximum = (max_resolution // KREA2_LATENT_DOWNSCALE) ** 2
+ if maximum == minimum:
+ return y1
+ slope = (y2 - y1) / (maximum - minimum)
+ return slope * image_token_count + (y1 - slope * minimum)
+
+
+def timesteps(
+ image_token_count: int,
+ steps: int,
+ *,
+ y1: float = 0.5,
+ y2: float = 1.15,
+ sigma: float = 1.0,
+ mu: float | None = None,
+) -> list[float]:
+ """Return a resolution-shifted descending flow-matching grid."""
+
+ if steps < 1:
+ raise DiffusionCliError(f"Steps must be at least 1: got {steps}")
+ if mu is None:
+ mu = timestepMu(image_token_count, y1=y1, y2=y2)
+ values = torch.linspace(1.0, 0.0, steps + 1, dtype=torch.float64)
+ exp_mu = math.exp(mu)
+ shifted = exp_mu / (
+ exp_mu + (1.0 / values - 1.0).pow(sigma)
+ )
+ shifted[-1] = 0.0
+ return shifted.tolist()
+
+
+def buildInitialNoise(
+ batch_size: int,
+ channels: int,
+ height: int,
+ width: int,
+ seed: int,
+ *,
+ device,
+ dtype,
+) -> torch.Tensor:
+ """Create independent per-image Gaussian noise using ``seed + index``."""
+
+ if batch_size < 1:
+ raise DiffusionCliError("Batch size must be at least 1")
+ values = []
+ for index in range(batch_size):
+ generator = torch.Generator(device=device).manual_seed(seed + index)
+ values.append(
+ torch.randn(
+ 1,
+ channels,
+ height // KREA2_LATENT_DOWNSCALE,
+ width // KREA2_LATENT_DOWNSCALE,
+ device=device,
+ dtype=dtype,
+ generator=generator,
+ )
+ )
+ return torch.cat(values, dim=0)
+
+
+@torch.no_grad()
+def sampleKrea2(
+ model,
+ conditioning: Krea2Conditioning,
+ *,
+ negative_conditioning: Krea2Conditioning | None = None,
+ batch_size: int,
+ height: int,
+ width: int,
+ seed: int,
+ steps: int,
+ cfg: float,
+ device,
+ dtype,
+ latent_channels: int = 16,
+ patch_size: int = 2,
+ mu: float | None = None,
+ shift_y1: float = 0.5,
+ shift_y2: float = 1.15,
+ progress_callback=None,
+) -> torch.Tensor:
+ """Run Krea Euler flow integration and return an NCHW latent."""
+
+ if cfg < 0:
+ raise DiffusionCliError(f"CFG must be non-negative: got {cfg}")
+ height = roundUp(height, KREA2_LATENT_DOWNSCALE * patch_size, "height")
+ width = roundUp(width, KREA2_LATENT_DOWNSCALE * patch_size, "width")
+ if conditioning.hidden_states.shape[0] != batch_size:
+ if conditioning.hidden_states.shape[0] == 1:
+ conditioning = Krea2Conditioning(
+ conditioning.hidden_states.expand(
+ batch_size,
+ -1,
+ -1,
+ -1,
+ ),
+ conditioning.attention_mask.expand(batch_size, -1),
+ )
+ else:
+ raise DiffusionCliError("Krea conditioning batch mismatch")
+ if cfg > 0 and negative_conditioning is None:
+ raise DiffusionCliError(
+ "Krea CFG requires negative conditioning when cfg is nonzero"
+ )
+ if cfg > 0 and negative_conditioning.hidden_states.shape[0] == 1:
+ negative_conditioning = Krea2Conditioning(
+ negative_conditioning.hidden_states.expand(
+ batch_size,
+ -1,
+ -1,
+ -1,
+ ),
+ negative_conditioning.attention_mask.expand(batch_size, -1),
+ )
+ elif cfg > 0 and negative_conditioning.hidden_states.shape[0] != batch_size:
+ raise DiffusionCliError("Krea negative conditioning batch mismatch")
+ latent = buildInitialNoise(
+ batch_size,
+ latent_channels,
+ height,
+ width,
+ seed,
+ device=device,
+ dtype=dtype,
+ )
+ image_tokens = (height // KREA2_LATENT_DOWNSCALE // patch_size) * (
+ width // KREA2_LATENT_DOWNSCALE // patch_size
+ )
+ schedule = timesteps(
+ image_tokens,
+ steps,
+ y1=shift_y1,
+ y2=shift_y2,
+ mu=mu,
+ )
+ for index, (current, following) in enumerate(
+ zip(schedule[:-1], schedule[1:])
+ ):
+ time = torch.full(
+ (batch_size,),
+ current,
+ device=device,
+ dtype=dtype,
+ )
+ conditional = model(latent, time, conditioning)
+ if cfg > 0:
+ unconditional = model(latent, time, negative_conditioning)
+ velocity = conditional + cfg * (conditional - unconditional)
+ else:
+ velocity = conditional
+ if not torch.isfinite(velocity).all():
+ raise DiffusionCliError(
+ f"Krea diffusion output contains NaN or Inf at step {index}"
+ )
+ latent = latent + (following - current) * velocity
+ if not torch.isfinite(latent).all():
+ raise DiffusionCliError(
+ f"Krea latent contains NaN or Inf at step {index}"
+ )
+ if progress_callback is not None:
+ progress_callback(index + 1, steps)
+ return latent
+
+
+# Compatibility aliases make the sampling boundary easy to discover.
+sample = sampleKrea2
+sampleLatents = sampleKrea2
+prepare = prepareConditioning
diff --git a/diffusion_cli/krea2_text_encoder.py b/diffusion_cli/krea2_text_encoder.py
new file mode 100644
index 0000000..3f38d2c
--- /dev/null
+++ b/diffusion_cli/krea2_text_encoder.py
@@ -0,0 +1,327 @@
+"""Local Qwen3-VL conditioning for Krea 2."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from pathlib import Path
+
+import torch
+
+from diffusion_cli.checkpoint import loadStateDict
+from diffusion_cli.config import ModelSource
+from diffusion_cli.errors import DiffusionCliError
+from diffusion_cli.krea2_model import Krea2Conditioning
+from diffusion_cli.quantization import (
+ QuantizationManifest,
+ ScaledFp8Linear,
+ loadQuantizationManifest,
+ moveModulePreservingQuantization,
+ validateQuantizedStateDict,
+)
+
+KREA2_TAP_LAYERS = (2, 5, 8, 11, 14, 17, 20, 23, 26, 29, 32, 35)
+KREA2_PROMPT_PREFIX = (
+ "<|im_start|>system\n"
+ "Describe the image by detailing the color, shape, size, texture, "
+ "quantity, text, spatial relationships of the objects and background:"
+ "<|im_end|>\n<|im_start|>user\n"
+)
+KREA2_PROMPT_SUFFIX = "<|im_end|>\n<|im_start|>assistant\n"
+KREA2_PAD_TOKEN_ID = 151643
+QWEN3_VL_4B_VOCAB_SIZE = 151936
+QWEN3_VL_4B_HIDDEN_SIZE = 2560
+QWEN3_VL_4B_INTERMEDIATE_SIZE = 9728
+QWEN3_VL_4B_LAYERS = 36
+QWEN3_VL_4B_HEADS = 32
+QWEN3_VL_4B_KV_HEADS = 8
+QWEN3_VL_4B_HEAD_DIM = 128
+QWEN3_VL_4B_MAX_POSITION = 262144
+QWEN3_VL_4B_ROPE_THETA = 5000000.0
+
+
+@dataclass(frozen=True)
+class Krea2PromptConditioning:
+ """Positive and optional negative Krea conditioning."""
+
+ positive: Krea2Conditioning
+ negative: Krea2Conditioning | None = None
+
+
+def buildQwen3Vl4BConfig():
+ """Build the text-only Qwen3-VL-4B configuration in code."""
+
+ from transformers import Qwen3VLTextConfig
+
+ return Qwen3VLTextConfig(
+ vocab_size=QWEN3_VL_4B_VOCAB_SIZE,
+ hidden_size=QWEN3_VL_4B_HIDDEN_SIZE,
+ intermediate_size=QWEN3_VL_4B_INTERMEDIATE_SIZE,
+ num_hidden_layers=QWEN3_VL_4B_LAYERS,
+ num_attention_heads=QWEN3_VL_4B_HEADS,
+ num_key_value_heads=QWEN3_VL_4B_KV_HEADS,
+ head_dim=QWEN3_VL_4B_HEAD_DIM,
+ max_position_embeddings=QWEN3_VL_4B_MAX_POSITION,
+ rope_parameters={
+ "rope_theta": QWEN3_VL_4B_ROPE_THETA,
+ "rope_type": "default",
+ },
+ rms_norm_eps=1e-6,
+ attention_bias=False,
+ tie_word_embeddings=True,
+ pad_token_id=KREA2_PAD_TOKEN_ID,
+ use_cache=False,
+ )
+
+
+def normalizeKreaTextStateDict(
+ state_dict: dict[str, torch.Tensor],
+) -> dict[str, torch.Tensor]:
+ """Normalize ComfyUI Qwen keys and consume ``comfy_quant`` metadata."""
+
+ normalized = {}
+ for key, value in state_dict.items():
+ if key.endswith(".comfy_quant"):
+ continue
+ if key.startswith("model."):
+ key = key.removeprefix("model.")
+ if key.startswith("visual."):
+ # Qwen3-VL files may carry a vision tower; Krea text conditioning
+ # deliberately documents it as an ignored component.
+ continue
+ normalized[key] = value
+ return normalized
+
+
+def _findTokenBoundary(
+ token_ids: list[int],
+ prefix_ids: list[int],
+) -> int:
+ """Find the prefix boundary using token IDs rather than character offsets."""
+
+ for start in range(len(token_ids) - len(prefix_ids) + 1):
+ if token_ids[start : start + len(prefix_ids)] == prefix_ids:
+ return start + len(prefix_ids)
+ raise DiffusionCliError(
+ "Krea prompt template prefix is missing its known token boundary"
+ )
+
+
+def _asTensor(value, *, dtype=None) -> torch.Tensor:
+ """Convert tokenizer output from lists or tensors to a tensor."""
+
+ if isinstance(value, torch.Tensor):
+ return value if dtype is None else value.to(dtype=dtype)
+ return torch.tensor(value, dtype=dtype)
+
+
+def initializeKreaRotaryEmbedding(model) -> None:
+ """Initialize the non-persistent Qwen RoPE buffer after meta loading."""
+
+ 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_VL_4B_ROPE_THETA,
+ )
+ parameter = next(model.parameters())
+ inv_freq = 1.0 / (
+ rope_theta
+ ** (
+ torch.arange(
+ 0,
+ head_dim,
+ 2,
+ dtype=torch.float32,
+ device=parameter.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 Krea2TextEncoder:
+ """Load a local Qwen3-VL text backbone and emit twelve hidden-state taps."""
+
+ def __init__(
+ self,
+ model_path: ModelSource | Path,
+ tokenizer_path: Path,
+ device,
+ dtype=torch.bfloat16,
+ *,
+ model=None,
+ tokenizer=None,
+ ) -> None:
+ self.device = device
+ self.dtype = dtype
+ self.tokenizer = tokenizer or self._loadTokenizer(tokenizer_path)
+ self.model = model or self._loadModel(model_path)
+ self.model.eval()
+
+ def encodePrompt(self, prompt: str) -> Krea2Conditioning:
+ """Encode one prompt with the exact Krea chat template."""
+
+ 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",
+ padding=False,
+ truncation=False,
+ add_special_tokens=False,
+ )
+ prefix_ids = _asTensor(prefix_batch["input_ids"], dtype=torch.long)
+ prefix_ids = prefix_ids.reshape(-1).tolist()
+ boundary = _findTokenBoundary(
+ input_ids[0].tolist(),
+ prefix_ids,
+ )
+ input_ids = input_ids.to(self.device)
+ attention_mask = attention_mask.to(self.device)
+ with torch.inference_mode():
+ output = self.model(
+ input_ids=input_ids,
+ attention_mask=attention_mask,
+ output_hidden_states=True,
+ use_cache=False,
+ )
+ try:
+ hidden = torch.stack(
+ [output.hidden_states[index] for index in KREA2_TAP_LAYERS],
+ dim=2,
+ )
+ except (AttributeError, IndexError) as exc:
+ raise DiffusionCliError(
+ "Text encoder checkpoint does not expose the 36 Qwen layers "
+ "required by Krea 2"
+ ) from exc
+ hidden = hidden[:, boundary:]
+ attention_mask = attention_mask[:, boundary:]
+ if hidden.shape[2] != len(KREA2_TAP_LAYERS):
+ raise DiffusionCliError("Krea text encoder returned invalid tap layout")
+ if hidden.shape[-1] != QWEN3_VL_4B_HIDDEN_SIZE:
+ raise DiffusionCliError(
+ "Text encoder checkpoint is not Qwen3-VL-4B"
+ )
+ return Krea2Conditioning(hidden, attention_mask)
+
+ def encode(self, prompt: str) -> Krea2Conditioning:
+ """Encode one positive prompt."""
+
+ return self.encodePrompt(prompt)
+
+ def encodePrompts(
+ self,
+ prompt: str,
+ negative_prompt: str = "",
+ *,
+ cfg: float = 0.0,
+ ) -> Krea2PromptConditioning:
+ """Encode positive and, only when requested, negative prompts."""
+
+ positive = self.encodePrompt(prompt)
+ negative = self.encodePrompt(negative_prompt) if cfg > 0 else None
+ return Krea2PromptConditioning(positive, negative)
+
+ def toDevice(self, device, dtype=None) -> None:
+ """Move the Qwen backbone without expanding its FP8 weights."""
+
+ self.device = device
+ self.dtype = dtype or self.dtype
+ moveModulePreservingQuantization(
+ self.model,
+ device,
+ self.dtype,
+ )
+ self.model.eval()
+
+ def toCpu(self) -> None:
+ """Move the Qwen backbone back to CPU."""
+
+ self.toDevice(torch.device("cpu"))
+
+ def _loadTokenizer(self, tokenizer_path: Path):
+ """Load only tokenizer assets from an explicit local directory."""
+
+ from transformers import Qwen2Tokenizer
+
+ tokenizer = Qwen2Tokenizer.from_pretrained(
+ tokenizer_path,
+ local_files_only=True,
+ padding_side="right",
+ )
+ tokenizer.pad_token_id = KREA2_PAD_TOKEN_ID
+ tokenizer.padding_side = "right"
+ return tokenizer
+
+ def _loadModel(self, model_path: ModelSource | Path):
+ """Construct the text-only model and strictly load local tensors."""
+
+ from transformers import Qwen3VLTextModel
+
+ source = model_path if isinstance(model_path, ModelSource) else ModelSource(
+ model_path,
+ "text_encoder",
+ )
+ raw_state = loadStateDict(source)
+ manifest = validateQuantizedStateDict(
+ normalizeKreaTextStateDict(raw_state),
+ loadQuantizationManifest(source),
+ )
+ normalized = normalizeKreaTextStateDict(raw_state)
+ with torch.device("meta"):
+ model = Qwen3VLTextModel(buildQwen3Vl4BConfig())
+ for name in manifest.layers:
+ module = model.get_submodule(name)
+ if not isinstance(module, torch.nn.Linear):
+ raise DiffusionCliError(
+ f"Quantized Qwen layer is not linear: {name}"
+ )
+ replacement = ScaledFp8Linear(
+ module.in_features,
+ module.out_features,
+ bias=module.bias is not None,
+ device="meta",
+ dtype=self.dtype,
+ )
+ parent_name, _, child_name = name.rpartition(".")
+ parent = model.get_submodule(parent_name) if parent_name else model
+ setattr(parent, child_name, replacement)
+ missing, unexpected = model.load_state_dict(
+ normalized,
+ strict=False,
+ assign=True,
+ )
+ allowed_missing = {"rotary_emb.inv_freq"}
+ missing = [key for key in missing if key not in allowed_missing]
+ if missing or unexpected:
+ raise DiffusionCliError(
+ "Krea text encoder checkpoint key mismatch: "
+ f"missing={missing[:3]}, unexpected={unexpected[:3]}"
+ )
+ initializeKreaRotaryEmbedding(model)
+ moveModulePreservingQuantization(model, self.device, self.dtype)
+ return model
+
+
+loadKrea2TextEncoder = Krea2TextEncoder
diff --git a/diffusion_cli/model_inspect.py b/diffusion_cli/model_inspect.py
index 0cf6832..dcf864e 100644
--- a/diffusion_cli/model_inspect.py
+++ b/diffusion_cli/model_inspect.py
@@ -10,6 +10,7 @@ from safetensors import safe_open
from diffusion_cli.config import ModelSource
from diffusion_cli.errors import DiffusionCliError
+from diffusion_cli.quantization import loadQuantizationManifest
@dataclass(frozen=True)
@@ -22,6 +23,13 @@ class TensorSummary:
dtype_counts: dict[str, int]
top_level_counts: dict[str, int]
architecture_guess: str
+ quantized_layer_count: int = 0
+ missing_scale_count: int = 0
+ unsupported_formats: tuple[str, ...] = ()
+ profile: str | None = None
+ architecture: str | None = None
+ variant: str | None = None
+ compatibility: str | None = None
def _topLevel(key: str) -> str:
@@ -45,6 +53,10 @@ def guessArchitecture(keys: list[str]) -> str:
or "x_embedder.weight" in key_set
):
return "z_image_diffusion"
+ if any(key.startswith("blocks.") for key in key_set) and any(
+ key.startswith("txtfusion.") for key in key_set
+ ):
+ return "krea2_diffusion"
if "diffusion_model.img_in.weight" in key_set:
return "z_image_diffusion"
if any(key.startswith("model.diffusion_model.") for key in key_set):
@@ -60,7 +72,7 @@ def inspectSafetensors(path: Path) -> TensorSummary:
return inspectModelSource(ModelSource(path=path, role="unknown"))
-def inspectModelSource(source: ModelSource) -> TensorSummary:
+def inspectModelSource(source: ModelSource, profile=None) -> TensorSummary:
"""Read metadata for one resolved model source."""
dtype_counts: Counter[str] = Counter()
@@ -85,17 +97,62 @@ def inspectModelSource(source: ModelSource) -> TensorSummary:
f"Checkpoint does not contain {source.role}: {source.path}"
)
+ architecture_guess = guessArchitecture(keys)
+ quantized_layer_count = 0
+ missing_scale_count = 0
+ unsupported_formats: set[str] = set()
+ try:
+ manifest = loadQuantizationManifest(source)
+ quantized_layer_count = len(manifest.layers)
+ for name, spec in manifest.layers.items():
+ if spec.format != "float8_e4m3fn":
+ unsupported_formats.add(spec.format)
+ if not any(
+ key in keys
+ for key in (f"{name}.weight_scale", f"model.{name}.weight_scale")
+ ):
+ missing_scale_count += 1
+ except DiffusionCliError:
+ raise
+
+ compatibility = None
+ if profile is not None:
+ expected = {
+ "z-image": {
+ "diffusion_model": "z_image_diffusion",
+ "text_encoder": "qwen_text_encoder",
+ "vae": "vae",
+ },
+ "krea2": {
+ "diffusion_model": "krea2_diffusion",
+ "text_encoder": "qwen_text_encoder",
+ "vae": "vae",
+ },
+ }[profile.architecture].get(source.role)
+ compatibility = (
+ "compatible"
+ if expected is None or architecture_guess == expected
+ else f"incompatible (expected {expected})"
+ )
+
return TensorSummary(
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)),
- architecture_guess=guessArchitecture(keys),
+ architecture_guess=architecture_guess,
+ quantized_layer_count=quantized_layer_count,
+ missing_scale_count=missing_scale_count,
+ unsupported_formats=tuple(sorted(unsupported_formats)),
+ profile=profile.name if profile is not None else None,
+ architecture=profile.architecture if profile is not None else None,
+ variant=profile.variant if profile is not None else None,
+ compatibility=compatibility,
)
-def formatSummary(name: str, summary: TensorSummary) -> str:
+def formatSummary(name: str, summary: TensorSummary, profile=None) -> str:
"""Format a checkpoint summary for terminal output."""
dtype_text = ", ".join(
@@ -117,5 +174,33 @@ def formatSummary(name: str, summary: TensorSummary) -> str:
f" dtypes: {dtype_text}",
f" top-level keys: {top_text}",
f" architecture guess: {summary.architecture_guess}",
+ f" quantized layers: {summary.quantized_layer_count}",
+ f" missing scales: {summary.missing_scale_count}",
+ (
+ " unsupported quantization formats: "
+ + ", ".join(summary.unsupported_formats)
+ if summary.unsupported_formats
+ else None
+ ),
+ (
+ f" profile: {profile.name}"
+ if profile is not None
+ else (f" profile: {summary.profile}" if summary.profile else None)
+ ),
+ (
+ f" configured architecture: {profile.architecture}"
+ if profile is not None
+ else None
+ ),
+ (
+ f" configured variant: {profile.variant}"
+ if profile is not None
+ else None
+ ),
+ (
+ f" compatibility: {summary.compatibility}"
+ if summary.compatibility is not None
+ else None
+ ),
] if line is not None]
)
diff --git a/diffusion_cli/paths.py b/diffusion_cli/paths.py
index 271da8a..291f727 100644
--- a/diffusion_cli/paths.py
+++ b/diffusion_cli/paths.py
@@ -13,9 +13,12 @@ from diffusion_cli.config import (
VAE_ROLE,
ModelPathConfig,
ModelFiles,
+ ModelProfile,
ModelSource,
ModelSources,
UserConfig,
+ validateTokenizerPath,
+ selectModelProfile,
)
from diffusion_cli.errors import DiffusionCliError
@@ -73,35 +76,102 @@ def resolveComponentSource(
def resolveModelSources(args, user_config: UserConfig) -> ModelSources:
"""Resolve and validate model sources for the active workflow."""
+ if not user_config.model_profiles and user_config.default_model is None:
+ models = user_config.models
+ if models is None:
+ raise DiffusionCliError(
+ "No model is configured; set default_model or define a model "
+ "profile"
+ )
+ checkpoint = _pathFromArgs(args, "checkpoint") or models.checkpoint
+ return ModelSources(
+ diffusion_model=resolveComponentSource(
+ _pathFromArgs(args, "diffusion_model"),
+ models.diffusion_model,
+ checkpoint,
+ DIFFUSION_ROLE,
+ CHECKPOINT_DIFFUSION_PREFIX,
+ "models.diffusion_model",
+ ),
+ text_encoder=resolveComponentSource(
+ _pathFromArgs(args, "text_encoder"),
+ models.text_encoder,
+ checkpoint,
+ TEXT_ENCODER_ROLE,
+ CHECKPOINT_TEXT_ENCODER_PREFIX,
+ "models.text_encoder",
+ ),
+ vae=resolveComponentSource(
+ _pathFromArgs(args, "vae"),
+ models.vae,
+ checkpoint,
+ VAE_ROLE,
+ CHECKPOINT_VAE_PREFIX,
+ "models.vae",
+ ),
+ tokenizer=models.tokenizer,
+ )
+ profile = selectModelProfile(
+ user_config,
+ getattr(args, "model_profile", None),
+ )
+ return resolveModelSourcesForProfile(
+ profile,
+ args=args,
+ )
- models = user_config.models
- checkpoint = _pathFromArgs(args, "checkpoint") or models.checkpoint
+
+def resolveModelSourcesForProfile(
+ profile: ModelProfile,
+ *,
+ args=None,
+) -> ModelSources:
+ """Resolve component paths and CLI overrides for one profile."""
+
+ args = args or object()
+ profile_checkpoint = _pathFromArgs(args, "checkpoint") or profile.checkpoint
+ diffusion_path = (
+ _pathFromArgs(args, "diffusion_model") or profile.diffusion_model
+ )
+ text_path = _pathFromArgs(args, "text_encoder") or profile.text_encoder
+ vae_path = _pathFromArgs(args, "vae") or profile.vae
+ tokenizer = (
+ _pathFromArgs(args, "tokenizer_path") or profile.tokenizer
+ )
+ if tokenizer is not None:
+ tokenizer = validateTokenizerPath(tokenizer)
+
+ if profile.architecture == "krea2" and profile_checkpoint is not None:
+ raise DiffusionCliError(
+ f"Krea profile {profile.name} requires separate model files"
+ )
return ModelSources(
diffusion_model=resolveComponentSource(
_pathFromArgs(args, "diffusion_model"),
- models.diffusion_model,
- checkpoint,
+ diffusion_path,
+ profile_checkpoint,
DIFFUSION_ROLE,
CHECKPOINT_DIFFUSION_PREFIX,
- "models.diffusion_model",
+ f"model_profiles.{profile.name}.diffusion_model",
),
text_encoder=resolveComponentSource(
_pathFromArgs(args, "text_encoder"),
- models.text_encoder,
- checkpoint,
+ text_path,
+ profile_checkpoint,
TEXT_ENCODER_ROLE,
CHECKPOINT_TEXT_ENCODER_PREFIX,
- "models.text_encoder",
+ f"model_profiles.{profile.name}.text_encoder",
),
vae=resolveComponentSource(
_pathFromArgs(args, "vae"),
- models.vae,
- checkpoint,
+ vae_path,
+ profile_checkpoint,
VAE_ROLE,
CHECKPOINT_VAE_PREFIX,
- "models.vae",
+ f"model_profiles.{profile.name}.vae",
),
+ tokenizer=tokenizer,
)
@@ -134,9 +204,20 @@ def resolveModelSourcesFromConfig(models: ModelPathConfig) -> ModelSources:
CHECKPOINT_VAE_PREFIX,
"models.vae",
),
+ tokenizer=models.tokenizer,
)
+def resolveModelSourcesFromProfile(
+ user_config: UserConfig,
+ profile_name: str | None = None,
+) -> ModelSources:
+ """Resolve model files for a selected configured profile."""
+
+ profile = selectModelProfile(user_config, profile_name)
+ return resolveModelSourcesForProfile(profile)
+
+
def resolveModelFiles(args, user_config: UserConfig) -> ModelFiles:
"""Resolve standalone model paths for compatibility with old callers."""
diff --git a/diffusion_cli/quantization.py b/diffusion_cli/quantization.py
new file mode 100644
index 0000000..50911fe
--- /dev/null
+++ b/diffusion_cli/quantization.py
@@ -0,0 +1,355 @@
+"""Local scaled-FP8 checkpoint support for the Krea backends."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+import json
+from pathlib import Path
+from typing import Any
+
+import torch
+from torch import nn
+from torch.nn import functional as F
+from safetensors import safe_open
+
+from diffusion_cli.config import ModelSource
+from diffusion_cli.errors import DiffusionCliError
+
+SUPPORTED_FP8_FORMAT = "float8_e4m3fn"
+FP8_DTYPES = {torch.float8_e4m3fn}
+
+
+@dataclass(frozen=True)
+class QuantizedLayerSpec:
+ """Quantization settings for one parameterized layer."""
+
+ format: str
+ full_precision_matrix_mult: bool = False
+
+
+@dataclass(frozen=True)
+class QuantizationManifest:
+ """Quantization settings indexed by module path."""
+
+ layers: dict[str, QuantizedLayerSpec]
+
+
+def _decodeJson(value: Any, label: str) -> Any:
+ """Decode JSON metadata stored as a string or uint8 tensor."""
+
+ if isinstance(value, torch.Tensor):
+ value = bytes(value.detach().cpu().tolist())
+ if isinstance(value, bytes):
+ try:
+ value = value.decode("utf-8")
+ except UnicodeDecodeError as exc:
+ raise DiffusionCliError(
+ f"Invalid quantization metadata in {label}"
+ ) from exc
+ if isinstance(value, str):
+ try:
+ return json.loads(value)
+ except json.JSONDecodeError as exc:
+ raise DiffusionCliError(
+ f"Invalid quantization metadata in {label}: {exc}"
+ ) from exc
+ if isinstance(value, dict):
+ return value
+ raise DiffusionCliError(f"Invalid quantization metadata in {label}")
+
+
+def _normalizeLayerName(name: str) -> str:
+ """Normalize ComfyUI and Transformers module prefixes."""
+
+ name = name.removesuffix(".weight")
+ name = name.removesuffix(".comfy_quant")
+ changed = True
+ while changed:
+ changed = False
+ for prefix in (
+ "model.diffusion_model.",
+ "model.",
+ "diffusion_model.",
+ ):
+ if name.startswith(prefix):
+ name = name.removeprefix(prefix)
+ changed = True
+ break
+ return name
+
+
+def _manifestEntries(value: Any) -> dict[str, Any]:
+ """Find the layer mapping in the observed ComfyUI metadata shapes."""
+
+ if not isinstance(value, dict):
+ raise DiffusionCliError("Quantization metadata must be a JSON object")
+ for key in ("layers", "quantization", "modules"):
+ nested = value.get(key)
+ if isinstance(nested, dict):
+ return nested
+ return value
+
+
+def _specFromValue(value: Any, label: str) -> QuantizedLayerSpec:
+ """Parse one quantized-layer metadata object."""
+
+ if isinstance(value, str):
+ format_name = value
+ full_precision = False
+ elif isinstance(value, dict):
+ format_name = value.get("format") or value.get("quant_format")
+ full_precision = bool(value.get("full_precision_matrix_mult", False))
+ else:
+ raise DiffusionCliError(f"Invalid quantization metadata for {label}")
+ if not isinstance(format_name, str) or not format_name:
+ raise DiffusionCliError(
+ f"Missing quantization format for {label}"
+ )
+ return QuantizedLayerSpec(format_name, full_precision)
+
+
+def manifestFromMetadata(value: Any, label: str = "checkpoint") -> QuantizationManifest:
+ """Normalize a JSON quantization metadata object into a manifest."""
+
+ entries = _manifestEntries(_decodeJson(value, label))
+ layers: dict[str, QuantizedLayerSpec] = {}
+ for raw_name, raw_spec in entries.items():
+ if not isinstance(raw_name, str):
+ raise DiffusionCliError(f"Invalid quantization layer in {label}")
+ name = _normalizeLayerName(raw_name)
+ if name.startswith("_"):
+ continue
+ layers[name] = _specFromValue(raw_spec, name)
+ return QuantizationManifest(layers)
+
+
+def mergeManifests(*manifests: QuantizationManifest) -> QuantizationManifest:
+ """Merge manifests while rejecting inconsistent layer declarations."""
+
+ layers: dict[str, QuantizedLayerSpec] = {}
+ for manifest in manifests:
+ for name, spec in manifest.layers.items():
+ previous = layers.get(name)
+ if previous is not None and previous != spec:
+ raise DiffusionCliError(
+ f"Inconsistent quantization metadata for {name}"
+ )
+ layers[name] = spec
+ return QuantizationManifest(layers)
+
+
+def manifestFromStateDict(
+ state_dict: dict[str, torch.Tensor],
+) -> QuantizationManifest:
+ """Read per-layer ``comfy_quant`` tensors from a state dict."""
+
+ manifests = []
+ for key, value in state_dict.items():
+ if key.endswith(".comfy_quant"):
+ manifests.append(manifestFromMetadata(value, key))
+ return mergeManifests(*manifests)
+
+
+def loadQuantizationManifest(source: ModelSource) -> QuantizationManifest:
+ """Read file-level and per-layer Krea quantization metadata locally."""
+
+ manifests: list[QuantizationManifest] = []
+ with safe_open(source.path, framework="pt", device="cpu") as tensors:
+ metadata = tensors.metadata() or {}
+ raw_metadata = metadata.get("_quantization_metadata")
+ if raw_metadata is not None:
+ manifests.append(manifestFromMetadata(raw_metadata, str(source.path)))
+ for key in tensors.keys():
+ if source.checkpoint_prefix is not None:
+ if not key.startswith(source.checkpoint_prefix):
+ continue
+ if key.endswith(".comfy_quant"):
+ decoded = _decodeJson(tensors.get_tensor(key), key)
+ if isinstance(decoded, dict) and "format" in decoded:
+ manifests.append(
+ QuantizationManifest(
+ {
+ _normalizeLayerName(key): _specFromValue(
+ decoded,
+ _normalizeLayerName(key),
+ )
+ }
+ )
+ )
+ else:
+ manifests.append(manifestFromMetadata(decoded, key))
+ return mergeManifests(*manifests)
+
+
+def validateQuantizedStateDict(
+ state_dict: dict[str, torch.Tensor],
+ manifest: QuantizationManifest | None = None,
+) -> QuantizationManifest:
+ """Validate scaled-FP8 weights, scalar scales, and supported metadata."""
+
+ if manifest is None:
+ manifest = manifestFromStateDict(state_dict)
+ for name, spec in manifest.layers.items():
+ if spec.format != SUPPORTED_FP8_FORMAT:
+ raise DiffusionCliError(
+ f"Unsupported quantization format in {name}: {spec.format}"
+ )
+ weight_key = f"{name}.weight"
+ scale_key = f"{name}.weight_scale"
+ weight = next(
+ (
+ state_dict.get(candidate)
+ for candidate in (
+ weight_key,
+ f"model.{weight_key}",
+ f"diffusion_model.{weight_key}",
+ )
+ if state_dict.get(candidate) is not None
+ ),
+ None,
+ )
+ if weight is None:
+ raise DiffusionCliError(f"Quantized layer {name} is missing weight")
+ if weight.dtype not in FP8_DTYPES:
+ raise DiffusionCliError(
+ f"Unsupported quantized weight dtype in {name}: {weight.dtype}"
+ )
+ scale = next(
+ (
+ state_dict.get(candidate)
+ for candidate in (
+ scale_key,
+ f"model.{scale_key}",
+ f"diffusion_model.{scale_key}",
+ )
+ if state_dict.get(candidate) is not None
+ ),
+ None,
+ )
+ if scale is None:
+ raise DiffusionCliError(
+ f"Krea layer {name} is missing weight_scale"
+ )
+ if scale.numel() != 1:
+ raise DiffusionCliError(
+ "Expected scalar weight scale, got shape "
+ f"{tuple(scale.shape)}"
+ )
+ if scale.dtype != torch.float32:
+ raise DiffusionCliError(
+ f"Expected F32 weight scale in {name}, got {scale.dtype}"
+ )
+ bias = next(
+ (
+ state_dict.get(candidate)
+ for candidate in (f"{name}.bias", f"model.{name}.bias")
+ if state_dict.get(candidate) is not None
+ ),
+ None,
+ )
+ if bias is not None and (
+ bias.dtype in FP8_DTYPES or not bias.dtype.is_floating_point
+ ):
+ raise DiffusionCliError(f"Quantized bias is unsupported in {name}")
+ for key, value in state_dict.items():
+ if key.endswith(".weight") and (
+ value.dtype in {torch.int8, torch.uint8, torch.int32, torch.int64}
+ ):
+ raise DiffusionCliError(
+ f"Unsupported integer quantized weight in {_normalizeLayerName(key)}"
+ )
+ if key.endswith(".weight") and value.dtype in FP8_DTYPES:
+ name = _normalizeLayerName(key)
+ if name not in manifest.layers:
+ raise DiffusionCliError(
+ f"FP8 layer {name} has no quantization metadata"
+ )
+ return manifest
+
+
+class ScaledFp8Linear(nn.Module):
+ """A correctness-first linear layer for scalar-scaled FP8 weights."""
+
+ def __init__(
+ self,
+ in_features: int,
+ out_features: int,
+ *,
+ bias: bool = True,
+ device=None,
+ dtype=torch.bfloat16,
+ ) -> None:
+ super().__init__()
+ self.in_features = in_features
+ self.out_features = out_features
+ self.weight = nn.Parameter(
+ torch.empty(
+ out_features,
+ in_features,
+ device=device,
+ dtype=torch.float8_e4m3fn,
+ ),
+ requires_grad=False,
+ )
+ self.register_buffer(
+ "weight_scale",
+ torch.ones((), device=device, dtype=torch.float32),
+ )
+ if bias:
+ self.bias = nn.Parameter(
+ torch.zeros(out_features, device=device, dtype=dtype),
+ requires_grad=False,
+ )
+ else:
+ self.register_parameter("bias", None)
+
+ def forward(self, input: torch.Tensor) -> torch.Tensor:
+ """Dequantize this layer's weight for one matrix operation."""
+
+ compute_dtype = input.dtype
+ weight = self.weight.to(device=input.device, dtype=compute_dtype)
+ scale = self.weight_scale.to(device=input.device, dtype=compute_dtype)
+ weight = weight * scale.reshape(1)
+ bias = None if self.bias is None else self.bias.to(input.dtype)
+ return F.linear(input, weight, bias)
+
+ def toDevice(self, device, dtype=None) -> None:
+ """Move the layer without expanding its FP8 storage weight."""
+
+ self.weight.data = self.weight.data.to(device=device)
+ self.weight_scale.data = self.weight_scale.data.to(device=device)
+ if self.bias is not None:
+ self.bias.data = self.bias.data.to(device=device)
+ if dtype is not None and self.bias.dtype.is_floating_point:
+ self.bias.data = self.bias.data.to(dtype=dtype)
+
+
+def moveModulePreservingQuantization(
+ module: nn.Module,
+ device,
+ runtime_dtype: torch.dtype,
+) -> nn.Module:
+ """Move a module while retaining FP8 and F32 storage dtypes."""
+
+ for child in module.modules():
+ if isinstance(child, ScaledFp8Linear):
+ child.toDevice(device, runtime_dtype)
+ else:
+ for parameter in child.parameters(recurse=False):
+ if parameter.dtype in FP8_DTYPES:
+ parameter.data = parameter.data.to(device=device)
+ else:
+ parameter.data = parameter.data.to(
+ device=device,
+ dtype=(
+ runtime_dtype
+ if parameter.dtype.is_floating_point
+ and parameter.dtype != torch.float32
+ else parameter.dtype
+ ),
+ )
+ for name, buffer in child.named_buffers(recurse=False):
+ if buffer is not None:
+ buffer = buffer.to(device=device)
+ child._buffers[name] = buffer
+ return module
diff --git a/diffusion_cli/qwen_image_vae.py b/diffusion_cli/qwen_image_vae.py
new file mode 100644
index 0000000..acb7693
--- /dev/null
+++ b/diffusion_cli/qwen_image_vae.py
@@ -0,0 +1,308 @@
+"""Explicit local construction and decoding for the Qwen Image VAE."""
+
+from __future__ import annotations
+
+from pathlib import Path
+import re
+
+import torch
+
+from diffusion_cli.checkpoint import loadStateDict
+from diffusion_cli.config import ModelSource
+from diffusion_cli.errors import DiffusionCliError
+
+QWEN_IMAGE_LATENT_MEAN = (
+ -0.7571,
+ -0.7089,
+ -0.9113,
+ 0.1075,
+ -0.1745,
+ 0.9653,
+ -0.1517,
+ 1.5508,
+ 0.4134,
+ -0.0715,
+ 0.5517,
+ -0.3632,
+ -0.1922,
+ -0.9497,
+ 0.2503,
+ -0.2921,
+)
+QWEN_IMAGE_LATENT_STD = (
+ 2.8184,
+ 1.4541,
+ 2.3275,
+ 2.6558,
+ 1.2196,
+ 1.7708,
+ 2.6052,
+ 2.0743,
+ 3.2687,
+ 2.1526,
+ 2.8652,
+ 1.5579,
+ 1.6382,
+ 1.1253,
+ 2.8251,
+ 1.9160,
+)
+
+
+def _residualKey(key: str) -> str:
+ """Convert an original residual block's sequential module names."""
+
+ return (
+ key.replace(".residual.0.gamma", ".norm1.gamma")
+ .replace(".residual.2.bias", ".conv1.bias")
+ .replace(".residual.2.weight", ".conv1.weight")
+ .replace(".residual.3.gamma", ".norm2.gamma")
+ .replace(".residual.6.bias", ".conv2.bias")
+ .replace(".residual.6.weight", ".conv2.weight")
+ .replace(".shortcut.bias", ".conv_shortcut.bias")
+ .replace(".shortcut.weight", ".conv_shortcut.weight")
+ )
+
+
+def convertWanVaeToDiffusers(
+ state_dict: dict[str, torch.Tensor],
+) -> dict[str, torch.Tensor]:
+ """Convert original Wan/Qwen VAE keys to Diffusers Qwen Image keys.
+
+ This mapping is attributed to Diffusers' ``convert_wan_vae_to_diffusers``
+ behavior and is maintained here so runtime loading never imports that
+ private helper.
+ """
+
+ converted: dict[str, torch.Tensor] = {}
+ middle_mapping = {
+ "encoder.middle.0.residual.0.gamma": "encoder.mid_block.resnets.0.norm1.gamma",
+ "encoder.middle.0.residual.2.bias": "encoder.mid_block.resnets.0.conv1.bias",
+ "encoder.middle.0.residual.2.weight": "encoder.mid_block.resnets.0.conv1.weight",
+ "encoder.middle.0.residual.3.gamma": "encoder.mid_block.resnets.0.norm2.gamma",
+ "encoder.middle.0.residual.6.bias": "encoder.mid_block.resnets.0.conv2.bias",
+ "encoder.middle.0.residual.6.weight": "encoder.mid_block.resnets.0.conv2.weight",
+ "encoder.middle.2.residual.0.gamma": "encoder.mid_block.resnets.1.norm1.gamma",
+ "encoder.middle.2.residual.2.bias": "encoder.mid_block.resnets.1.conv1.bias",
+ "encoder.middle.2.residual.2.weight": "encoder.mid_block.resnets.1.conv1.weight",
+ "encoder.middle.2.residual.3.gamma": "encoder.mid_block.resnets.1.norm2.gamma",
+ "encoder.middle.2.residual.6.bias": "encoder.mid_block.resnets.1.conv2.bias",
+ "encoder.middle.2.residual.6.weight": "encoder.mid_block.resnets.1.conv2.weight",
+ "decoder.middle.0.residual.0.gamma": "decoder.mid_block.resnets.0.norm1.gamma",
+ "decoder.middle.0.residual.2.bias": "decoder.mid_block.resnets.0.conv1.bias",
+ "decoder.middle.0.residual.2.weight": "decoder.mid_block.resnets.0.conv1.weight",
+ "decoder.middle.0.residual.3.gamma": "decoder.mid_block.resnets.0.norm2.gamma",
+ "decoder.middle.0.residual.6.bias": "decoder.mid_block.resnets.0.conv2.bias",
+ "decoder.middle.0.residual.6.weight": "decoder.mid_block.resnets.0.conv2.weight",
+ "decoder.middle.2.residual.0.gamma": "decoder.mid_block.resnets.1.norm1.gamma",
+ "decoder.middle.2.residual.2.bias": "decoder.mid_block.resnets.1.conv1.bias",
+ "decoder.middle.2.residual.2.weight": "decoder.mid_block.resnets.1.conv1.weight",
+ "decoder.middle.2.residual.3.gamma": "decoder.mid_block.resnets.1.norm2.gamma",
+ "decoder.middle.2.residual.6.bias": "decoder.mid_block.resnets.1.conv2.bias",
+ "decoder.middle.2.residual.6.weight": "decoder.mid_block.resnets.1.conv2.weight",
+ }
+ attention_mapping = {
+ "encoder.middle.1.norm.gamma": "encoder.mid_block.attentions.0.norm.gamma",
+ "encoder.middle.1.to_qkv.weight": "encoder.mid_block.attentions.0.to_qkv.weight",
+ "encoder.middle.1.to_qkv.bias": "encoder.mid_block.attentions.0.to_qkv.bias",
+ "encoder.middle.1.proj.weight": "encoder.mid_block.attentions.0.proj.weight",
+ "encoder.middle.1.proj.bias": "encoder.mid_block.attentions.0.proj.bias",
+ "decoder.middle.1.norm.gamma": "decoder.mid_block.attentions.0.norm.gamma",
+ "decoder.middle.1.to_qkv.weight": "decoder.mid_block.attentions.0.to_qkv.weight",
+ "decoder.middle.1.to_qkv.bias": "decoder.mid_block.attentions.0.to_qkv.bias",
+ "decoder.middle.1.proj.weight": "decoder.mid_block.attentions.0.proj.weight",
+ "decoder.middle.1.proj.bias": "decoder.mid_block.attentions.0.proj.bias",
+ }
+ head_mapping = {
+ "encoder.head.0.gamma": "encoder.norm_out.gamma",
+ "encoder.head.2.bias": "encoder.conv_out.bias",
+ "encoder.head.2.weight": "encoder.conv_out.weight",
+ "decoder.head.0.gamma": "decoder.norm_out.gamma",
+ "decoder.head.2.bias": "decoder.conv_out.bias",
+ "decoder.head.2.weight": "decoder.conv_out.weight",
+ "conv1.weight": "quant_conv.weight",
+ "conv1.bias": "quant_conv.bias",
+ "conv2.weight": "post_quant_conv.weight",
+ "conv2.bias": "post_quant_conv.bias",
+ }
+ for key, value in state_dict.items():
+ if key in middle_mapping:
+ converted[middle_mapping[key]] = value
+ elif key in attention_mapping:
+ converted[attention_mapping[key]] = value
+ elif key in head_mapping:
+ converted[head_mapping[key]] = value
+ elif key == "encoder.conv1.weight":
+ converted["encoder.conv_in.weight"] = value
+ elif key == "encoder.conv1.bias":
+ converted["encoder.conv_in.bias"] = value
+ elif key == "decoder.conv1.weight":
+ converted["decoder.conv_in.weight"] = value
+ elif key == "decoder.conv1.bias":
+ converted["decoder.conv_in.bias"] = value
+ elif key.startswith("encoder.downsamples."):
+ new_key = key.replace("encoder.downsamples.", "encoder.down_blocks.")
+ converted[_residualKey(new_key)] = value
+ elif key.startswith("decoder.upsamples."):
+ new_key = _convertDecoderUpsampleKey(key)
+ converted[new_key] = value
+ else:
+ converted[key] = value
+ return converted
+
+
+def _convertDecoderUpsampleKey(key: str) -> str:
+ """Convert one decoder upsample family key."""
+
+ match = re.match(r"decoder\.upsamples\.(\d+)\.(.*)", key)
+ if match is None:
+ return key
+ block_index = int(match.group(1))
+ suffix = match.group(2)
+ if suffix.startswith("residual.") or ".residual." in suffix:
+ groups = {0: (0, 0), 1: (0, 1), 2: (0, 2), 4: (1, 0), 5: (1, 1),
+ 6: (1, 2), 8: (2, 0), 9: (2, 1), 10: (2, 2),
+ 12: (3, 0), 13: (3, 1), 14: (3, 2)}
+ if block_index in groups:
+ up, residual = groups[block_index]
+ residual_suffix = suffix.removeprefix("residual.")
+ converted_suffix = _residualKey(
+ f"block.residual.{residual_suffix}"
+ ).removeprefix("block.")
+ return f"decoder.up_blocks.{up}.resnets.{residual}.{converted_suffix}"
+ if suffix.startswith("shortcut.") or ".shortcut." in suffix:
+ if block_index == 4:
+ suffix = suffix.replace("shortcut.", "resnets.0.conv_shortcut.")
+ return f"decoder.up_blocks.1.{suffix}"
+ suffix = suffix.replace("shortcut.", "conv_shortcut.")
+ return f"decoder.up_blocks.{block_index}.{suffix}"
+ if suffix.startswith("resample.") or suffix.startswith("time_conv.") \
+ or ".resample." in suffix or ".time_conv." in suffix:
+ block = {3: 0, 7: 1, 11: 2}.get(block_index, block_index)
+ return f"decoder.up_blocks.{block}.upsamplers.0.{suffix}"
+ return f"decoder.up_blocks.{block_index}.{suffix}"
+
+
+def normalizeQwenVaeStateDict(
+ state_dict: dict[str, torch.Tensor],
+) -> dict[str, torch.Tensor]:
+ """Remove common component prefixes before conversion."""
+
+ normalized = {}
+ for key, value in state_dict.items():
+ if key.startswith("vae."):
+ key = key.removeprefix("vae.")
+ normalized[key] = value
+ return normalized
+
+
+class QwenImageVae:
+ """Decoder-only wrapper around an explicitly constructed Qwen Image VAE."""
+
+ compression = 8
+ channels = 16
+
+ def __init__(
+ self,
+ model_path: ModelSource | Path,
+ device,
+ dtype=torch.bfloat16,
+ *,
+ model=None,
+ ) -> None:
+ self.device = device
+ self.dtype = dtype
+ if model is not None:
+ self.model = model
+ else:
+ self.model = self._loadModel(model_path)
+ self.model.eval()
+
+ def decode(self, latent: torch.Tensor) -> torch.Tensor:
+ """Decode an NCHW normalized latent into an RGB [0, 1] tensor."""
+
+ if latent.ndim != 4 or latent.shape[1] != self.channels:
+ raise DiffusionCliError("Qwen Image VAE expects NCHW 16-channel latent")
+ if not torch.isfinite(latent).all():
+ raise DiffusionCliError("Qwen Image VAE input contains NaN or Inf")
+ mean = torch.tensor(
+ QWEN_IMAGE_LATENT_MEAN,
+ device=latent.device,
+ dtype=latent.dtype,
+ ).view(1, self.channels, 1, 1, 1)
+ std = torch.tensor(
+ QWEN_IMAGE_LATENT_STD,
+ device=latent.device,
+ dtype=latent.dtype,
+ ).view(1, self.channels, 1, 1, 1)
+ value = (latent.unsqueeze(2) * std) + mean
+ with torch.inference_mode():
+ result = self.model.decode(value)
+ result = result.sample if hasattr(result, "sample") else result
+ if result.ndim != 5 or result.shape[1] != 3 or result.shape[2] != 1:
+ raise DiffusionCliError("Qwen Image VAE returned an invalid RGB shape")
+ result = result[:, :, 0]
+ result = result.add(1.0).div(2.0).clamp(0.0, 1.0)
+ if not torch.isfinite(result).all():
+ raise DiffusionCliError("Qwen Image VAE output contains NaN or Inf")
+ return result
+
+ def toDevice(self, device, dtype=None) -> None:
+ """Move the VAE to the active device."""
+
+ self.device = device
+ if dtype is not None:
+ self.dtype = dtype
+ self.model.to(device=device, dtype=self.dtype)
+ self.model.eval()
+
+ def toCpu(self) -> None:
+ """Move the VAE back to CPU."""
+
+ self.toDevice(torch.device("cpu"))
+
+ def _loadModel(self, model_path: ModelSource | Path):
+ """Construct and strictly load the local VAE without Hub loaders."""
+
+ try:
+ from diffusers import AutoencoderKLQwenImage
+ except ImportError as exc:
+ raise DiffusionCliError(
+ "diffusers is required for the Krea Qwen Image VAE"
+ ) from exc
+ source = model_path if isinstance(model_path, ModelSource) else ModelSource(
+ model_path,
+ "vae",
+ )
+ raw_state = normalizeQwenVaeStateDict(loadStateDict(source))
+ if any(
+ value.dtype in {torch.float8_e4m3fn, torch.float8_e5m2}
+ for value in raw_state.values()
+ ):
+ raise DiffusionCliError(
+ "Unsupported FP8 convolution in Qwen Image VAE"
+ )
+ converted = convertWanVaeToDiffusers(raw_state)
+ model = AutoencoderKLQwenImage(
+ base_dim=96,
+ z_dim=16,
+ dim_mult=[1, 2, 4, 4],
+ num_res_blocks=2,
+ attn_scales=[],
+ temperal_downsample=[False, True, True],
+ dropout=0.0,
+ input_channels=3,
+ latents_mean=list(QWEN_IMAGE_LATENT_MEAN),
+ latents_std=list(QWEN_IMAGE_LATENT_STD),
+ )
+ try:
+ model.load_state_dict(converted, strict=True)
+ except RuntimeError as exc:
+ raise DiffusionCliError(
+ f"Qwen Image VAE conversion left unexpected keys: {exc}"
+ ) from exc
+ model.to(device=self.device, dtype=self.dtype)
+ return model
+
+
+loadQwenImageVae = QwenImageVae
diff --git a/diffusion_cli/server.py b/diffusion_cli/server.py
index 8791bf3..d514f5c 100644
--- a/diffusion_cli/server.py
+++ b/diffusion_cli/server.py
@@ -9,6 +9,8 @@ from flask import Flask, request
from diffusion_cli.api_profiles import API_PROFILES
from diffusion_cli.errors import DiffusionCliError
+from diffusion_cli.config import selectModelProfile
+from diffusion_cli.paths import resolveModelSourcesFromProfile
from diffusion_cli.generation_service import (
MODEL_RESIDENCY_VALUES,
GenerationService,
@@ -99,12 +101,20 @@ def createApp(server_config: ServerConfig, generation_service) -> Flask:
return app
-def serve(server_config: ServerConfig, user_config) -> None:
+def serve(
+ server_config: ServerConfig,
+ user_config,
+ *,
+ model_profile_name: str | None = None,
+) -> None:
"""Run the configured Flask development server."""
+ profile = selectModelProfile(user_config, model_profile_name)
+ resolveModelSourcesFromProfile(user_config, profile.name)
generation_service = GenerationService(
user_config,
model_residency=server_config.model_residency,
+ model_profile_name=model_profile_name,
)
app = createApp(server_config, generation_service)
app.run(host=server_config.host, port=server_config.port)
diff --git a/diffusion_cli/static/ui/app.js b/diffusion_cli/static/ui/app.js
index 7517162..99b9aaa 100644
--- a/diffusion_cli/static/ui/app.js
+++ b/diffusion_cli/static/ui/app.js
@@ -464,7 +464,9 @@ class App extends Component {
? "Generating"
: state.is_loading_defaults ? "Loading" : "Ready";
const health_text = state.health
- ? state.health.api_profile + " / " + state.health.model_residency
+ ? (state.health.model_profile || "z-image-local")
+ + " / " + state.health.api_profile
+ + " / " + state.health.model_residency
: "Server";
return h(
diff --git a/diffusion_cli/ui.py b/diffusion_cli/ui.py
index faff2ca..14aaf6b 100644
--- a/diffusion_cli/ui.py
+++ b/diffusion_cli/ui.py
@@ -57,15 +57,26 @@ def _optionalFloat(data: dict, key: str) -> float | None:
return float(value)
-def _validateOptionalDimensions(width: int | None, height: int | None) -> None:
+def _validateOptionalDimensions(
+ width: int | None,
+ height: int | None,
+ model_profile=None,
+) -> None:
if width is None and height is None:
return
effective_width = 8 if width is None else width
effective_height = 8 if height is None else height
+ if (
+ model_profile is not None
+ and model_profile.architecture == "krea2"
+ ):
+ if effective_width <= 0 or effective_height <= 0:
+ raise DiffusionCliError("Krea width and height must be positive")
+ return
validateDimensions(effective_width, effective_height)
-def _txt2imgRequest(data: dict) -> ImageGenerationRequest:
+def _txt2imgRequest(data: dict, model_profile=None) -> ImageGenerationRequest:
prompt = data.get("prompt")
if not isinstance(prompt, str) or not prompt.strip():
raise DiffusionCliError("prompt required")
@@ -80,7 +91,7 @@ def _txt2imgRequest(data: dict) -> ImageGenerationRequest:
width = _optionalInt(data, "width")
height = _optionalInt(data, "height")
- _validateOptionalDimensions(width, height)
+ _validateOptionalDimensions(width, height, model_profile)
batch_size = _optionalInt(data, "batch_size")
steps = _optionalInt(data, "steps")
@@ -115,8 +126,10 @@ def _txt2imgRequest(data: dict) -> ImageGenerationRequest:
def _defaultResponse(context) -> dict:
+ profile = getattr(context.generation_service, "modelProfile", None)
default_request = buildDefaultGenerationRequest(
context.generation_service.user_config,
+ profile,
)
return {
"negative_prompt": default_request.negative_prompt,
@@ -125,8 +138,14 @@ def _defaultResponse(context) -> dict:
"batch_size": default_request.batch_size,
"steps": default_request.steps,
"cfg": default_request.cfg,
+ "mu": default_request.mu,
+ "shift_y1": default_request.shift_y1,
+ "shift_y2": default_request.shift_y2,
"output_extension": default_request.output_extension,
"output_quality": default_request.output_quality,
+ "model_profile": profile.name if profile is not None else "z-image-local",
+ "architecture": profile.architecture if profile is not None else "z-image",
+ "variant": profile.variant if profile is not None else "turbo",
"output_formats": list(OUTPUT_FORMAT_OPTIONS),
"seed": default_request.seed,
"resolution_presets": [
@@ -174,10 +193,18 @@ def registerUiRoutes(app, context) -> None:
@app.route("/ui/api/health", methods=["GET"])
def uiHealth():
+ profile = getattr(context.generation_service, "modelProfile", None)
return jsonify({
"status": "ok",
"api_profile": context.server_config.api_profile,
"model_residency": context.server_config.model_residency,
+ "model_profile": (
+ profile.name if profile is not None else "z-image-local"
+ ),
+ "architecture": (
+ profile.architecture if profile is not None else "z-image"
+ ),
+ "variant": profile.variant if profile is not None else "turbo",
})
@app.route("/ui/api/defaults", methods=["GET"])
@@ -191,7 +218,10 @@ def registerUiRoutes(app, context) -> None:
return _badRequest("JSON object body required")
try:
- generation_request = _txt2imgRequest(data)
+ generation_request = _txt2imgRequest(
+ data,
+ getattr(context.generation_service, "modelProfile", None),
+ )
except DiffusionCliError as exc:
return _badRequest(str(exc))
diff --git a/licenses/KREA2-APACHE-2.0.txt b/licenses/KREA2-APACHE-2.0.txt
new file mode 100644
index 0000000..a463d4e
--- /dev/null
+++ b/licenses/KREA2-APACHE-2.0.txt
@@ -0,0 +1,200 @@
+Apache License
+ Version 2.0, January 2004
+ http://www.apache.org/licenses/
+
+ TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+ 1. Definitions.
+
+ "License" shall mean the terms and conditions for use, reproduction,
+ and distribution as defined by Sections 1 through 9 of this document.
+
+ "Licensor" shall mean the copyright owner or entity authorized by
+ the copyright owner that is granting the License.
+
+ "Legal Entity" shall mean the union of the acting entity and all
+ other entities that control, are controlled by, or are under common
+ control with that entity. For the purposes of this definition,
+ "control" means (i) the power, direct or indirect, to cause the
+ direction or management of such entity, whether by contract or
+ otherwise, or (ii) ownership of fifty percent (50%) or more of the
+ outstanding shares, or (iii) beneficial ownership of such entity.
+
+ "You" (or "Your") shall mean an individual or Legal Entity
+ exercising permissions granted by this License.
+
+ "Source" form shall mean the preferred form for making modifications,
+ including but not limited to software source code, documentation
+ source, and configuration files.
+
+ "Object" form shall mean any form resulting from mechanical
+ transformation or translation of a Source form, including but not
+ limited to compiled object code, generated documentation,
+ and conversions to other media types.
+
+ "Work" shall mean the work of authorship, whether in Source or
+ Object form, made available under the License, as indicated by a
+ copyright notice that is included in or attached to the work
+ (an example is provided in the Appendix below).
+
+ "Derivative Works" shall mean any work, whether in Source or
+ Object form, that is based on (or derived from) the Work and for which
+ the editorial revisions, annotations, elaborations, or other
+ modifications represent, as a whole, an original work of authorship.
+ For the purposes of this License, Derivative Works shall not include
+ works that remain separable from, or merely link (or bind by name) to
+ the interfaces of, the Work and Derivative Works thereof.
+
+ "Contribution" shall mean any work of authorship, including
+ the original version of the Work and any modifications or additions
+ to that Work or Derivative Works thereof, that is intentionally
+ submitted to Licensor for inclusion in the Work by the copyright owner
+ or by an individual or Legal Entity authorized to submit on behalf of
+ the copyright owner. For the purposes of this definition,
+ "submitted" means any form of electronic, verbal, or written
+ communication sent to the Licensor or its representatives, including
+ but not limited to communication on electronic mailing lists, source
+ code control systems, and issue tracking systems that are managed by,
+ or on behalf of, the Licensor for the purpose of discussing and
+ improving the Work, but excluding communication that is conspicuously
+ marked or otherwise designated in writing by the copyright owner as
+ "Not a Contribution."
+
+ "Contributor" shall mean Licensor and any individual or Legal Entity
+ on behalf of whom a Contribution has been received by Licensor and
+ subsequently incorporated within the Work.
+
+ 2. Grant of Copyright License. Subject to the terms and conditions of
+ this License, each Contributor hereby grants to You a perpetual,
+ worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+ copyright license to reproduce, prepare Derivative Works of,
+ publicly display, publicly perform, sublicense, and distribute the
+ Work and such Derivative Works in Source or Object form.
+
+ 3. Grant of Patent License. Subject to the terms and conditions of
+ this License, each Contributor hereby grants to You a perpetual,
+ worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+ (except as stated in this section) patent license to make, have made,
+ use, offer to sell, sell, import, and otherwise transfer the Work,
+ where such license applies only to those patent claims licensable
+ by such Contributor that are necessarily infringed by their
+ Contribution(s) alone or by combination of their Contribution(s)
+ with the Work to which such Contribution(s) was submitted. If You
+ institute patent litigation against any entity (including a
+ cross-claim or counterclaim in a lawsuit) alleging that the Work
+ or a Contribution incorporated within the Work constitutes direct
+ or contributory patent infringement, then any patent licenses
+ granted to You under this License for that Work shall terminate
+ as of the date such litigation is filed.
+
+ 4. Redistribution. You may reproduce and distribute copies of the
+ Work or Derivative Works thereof in any medium, with or without
+ modifications, and in Source or Object form, provided that You meet
+ the following conditions:
+
+ (a) You must give any other recipients of the Work or Derivative
+ Works a copy of this License; and
+
+ (b) You must cause any modified files to carry prominent notices
+ stating that You changed the files; and
+
+ (c) You must retain, in the Source form of any Derivative Works that
+ You distribute, all copyright, patent, trademark, and attribution
+ notices from the Source form of the Work, excluding those notices
+ that do not pertain to any part of the Derivative Works; and
+
+ (d) If the Work includes a "NOTICE" text file as part of its
+ distribution, then any Derivative Works that You distribute must
+ include a readable copy of the attribution notices contained
+ within such NOTICE file, excluding those notices that do not
+ pertain to any part of the Derivative Works, in at least one of
+ the following places: within a NOTICE text file distributed as
+ part of the Derivative Works; within the Source form or
+ documentation, if provided along with the Derivative Works; or,
+ within a display generated by the Derivative Works, if and
+ wherever such third-party notices normally appear. The contents
+ of the NOTICE file are for informational purposes only and do not
+ modify the License. You may add Your own attribution notices
+ within Derivative Works alongside or as an addendum to the NOTICE
+ text from the Work, provided that such additional attribution
+ notices cannot be construed as modifying the License.
+
+ You may add Your own copyright statement to Your modifications and
+ may provide additional or different license terms and conditions for
+ use, reproduction, and distribution of Your modifications, or for any
+ such Derivative Works as a whole, provided Your use, reproduction, and
+ distribution of the Work otherwise complies with the conditions stated
+ in this License.
+
+ 5. Submission of Contributions. Unless You explicitly state otherwise,
+ any Contribution intentionally submitted for inclusion in the Work
+ by You to the Licensor shall be under the terms and conditions of
+ this License, without any additional terms or conditions.
+ Notwithstanding the above, nothing herein shall supersede or modify
+ the terms of any separate license agreement you may have executed
+ with Licensor regarding such Contributions.
+
+ 6. Trademarks. This License does not grant permission to use the trade
+ names, trademarks, service marks, or product names of the Licensor,
+ except as required for reasonable and customary use in describing the
+ origin of the Work and reproducing the content of the NOTICE file.
+
+ 7. Disclaimer of Warranty. Unless required by applicable law or agreed
+ to in writing, Licensor provides the Work (and each Contributor
+ provides its Contributions) on an "AS IS" BASIS, WITHOUT WARRANTIES
+ OR CONDITIONS OF ANY KIND, either express or implied, including,
+ without limitation, any warranties or conditions of TITLE,
+ NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A PARTICULAR
+ PURPOSE. You are solely responsible for determining the
+ appropriateness of using or redistributing the Work and assume any
+ risks associated with Your exercise of permissions under this License.
+
+ 8. Limitation of Liability. In no event and under no legal theory,
+ whether in tort (including negligence), contract, or otherwise,
+ unless required by applicable law (such as deliberate and grossly
+ negligent acts) or agreed to in writing, shall any Contributor be
+ liable to You for damages, including any direct, indirect, special,
+ incidental, or consequential damages of any character arising as a
+ result of this License or out of the use or inability to use the Work
+ (including but not limited to damages for loss of goodwill, work
+ stoppage, computer failure or malfunction, or any and all other
+ commercial damages or losses), even if such Contributor has been
+ advised of the possibility of such damages.
+
+ 9. Accepting Warranty or Additional Liability. While redistributing the
+ Work or Derivative Works thereof, You may choose to offer, and charge
+ a fee for, acceptance of support, warranty, indemnity, or other
+ liability obligations and/or rights consistent with this License.
+ However, in accepting such obligations, You may act only on Your own
+ behalf and on Your sole responsibility, not on behalf of any other
+ Contributor, and only if You agree to indemnify, defend, and hold each
+ Contributor harmless for any liability incurred by, or claims asserted
+ against, such Contributor by reason of your accepting any such warranty
+ or additional liability.
+
+ END OF TERMS AND CONDITIONS
+
+ APPENDIX: How to apply the Apache License to your work.
+
+ To apply the Apache License to your work, attach the following
+ boilerplate notice, with the fields enclosed by brackets "[]"
+ replaced with your own identifying information. (Don't include the
+ brackets!) The text should be enclosed in the appropriate comment
+ syntax for the file format. We also recommend that a file or class
+ name and description of purpose be included on the same "printed page"
+ as the copyright notice for easier identification within third-party
+ archives.
+
+ Copyright [yyyy] [name of copyright owner]
+
+ Licensed under the Apache License, Version 2.0 (the "License");
+ you may not use this file except in compliance with the License.
+ You may obtain a copy of the License at
+
+ http://www.apache.org/licenses/LICENSE-2.0
+
+ Unless required by applicable law or agreed to in writing, software
+ distributed under the License is distributed on an "AS IS" BASIS,
+ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ See the License for the specific language governing permissions and
+ limitations under the License.
diff --git a/pyproject.toml b/pyproject.toml
index 8b35129..af5c1e2 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -3,6 +3,7 @@ name = "diffusion-cli"
version = "0.1.0"
description = "Standalone CLI experiments for local diffusion model inference."
readme = "README.md"
+license-files = ["licenses/*"]
requires-python = ">=3.11"
dependencies = [
"flask",
@@ -13,6 +14,7 @@ dependencies = [
"torch",
"tqdm",
"transformers",
+ "diffusers",
]
[project.scripts]
diff --git a/scripts/krea2_integration.py b/scripts/krea2_integration.py
new file mode 100644
index 0000000..e295cd6
--- /dev/null
+++ b/scripts/krea2_integration.py
@@ -0,0 +1,90 @@
+#!/usr/bin/env python3
+"""Opt-in local Krea 2 Turbo integration check.
+
+This script only accepts explicit local paths. It never resolves a Hub model
+identifier and is intentionally outside ordinary unit-test discovery.
+"""
+
+from __future__ import annotations
+
+import argparse
+from pathlib import Path
+import time
+
+from PIL import Image
+import torch
+
+from diffusion_cli.config import (
+ GenerationDefaults,
+ ImageGenerationRequest,
+ ModelProfile,
+ UserConfig,
+)
+from diffusion_cli.generation_service import GenerationService
+
+
+def buildParser() -> argparse.ArgumentParser:
+ """Build arguments for the explicit local-checkpoint integration run."""
+
+ parser = argparse.ArgumentParser()
+ parser.add_argument("--diffusion-model", type=Path, required=True)
+ parser.add_argument("--text-encoder", type=Path, required=True)
+ parser.add_argument("--vae", type=Path, required=True)
+ parser.add_argument("--tokenizer", type=Path, required=True)
+ parser.add_argument("--output", type=Path, default=Path("krea2-test.png"))
+ parser.add_argument("--seed", type=int, default=0)
+ return parser
+
+
+def main() -> int:
+ """Generate one fixed-default Turbo image and report basic statistics."""
+
+ args = buildParser().parse_args()
+ profile = ModelProfile(
+ name="krea2-turbo",
+ architecture="krea2",
+ variant="turbo",
+ diffusion_model=args.diffusion_model,
+ text_encoder=args.text_encoder,
+ vae=args.vae,
+ tokenizer=args.tokenizer,
+ )
+ user_config = UserConfig(
+ models=None,
+ generation=GenerationDefaults(
+ output=args.output,
+ device="cuda",
+ dtype="auto",
+ ),
+ model_profiles={profile.name: profile},
+ default_model=profile.name,
+ )
+ request = ImageGenerationRequest(
+ prompt="a fox walking through fresh snow",
+ seed=args.seed,
+ )
+ service = GenerationService(user_config, model_residency="staged")
+ if torch.cuda.is_available():
+ torch.cuda.reset_peak_memory_stats()
+ started = time.perf_counter()
+ paths = service.generateToFiles(request)
+ elapsed = time.perf_counter() - started
+ peak_memory = (
+ torch.cuda.max_memory_allocated() / (1024 ** 3)
+ if torch.cuda.is_available()
+ else 0.0
+ )
+ for path in paths:
+ with Image.open(path) as image:
+ extrema = image.convert("RGB").getextrema()
+ print(
+ f"generated: {path} size={image.size} "
+ f"extrema={extrema}"
+ )
+ print(f"elapsed_seconds: {elapsed:.2f}")
+ print(f"peak_cuda_gib: {peak_memory:.2f}")
+ return 0
+
+
+if __name__ == "__main__":
+ raise SystemExit(main())
diff --git a/tests/test_krea2.py b/tests/test_krea2.py
new file mode 100644
index 0000000..9927ef2
--- /dev/null
+++ b/tests/test_krea2.py
@@ -0,0 +1,238 @@
+import tempfile
+import unittest
+from pathlib import Path
+from unittest.mock import patch
+
+import torch
+from torch import nn
+from safetensors.torch import save_file
+
+from diffusion_cli.config import (
+ GenerationDefaults,
+ ImageGenerationRequest,
+ ModelProfile,
+ ModelSource,
+ UserConfig,
+ TOKENIZER_FILES,
+ buildGenerationConfigFromRequest,
+ loadUserConfig,
+ selectModelProfile,
+)
+from diffusion_cli.krea2_model import (
+ Krea2Conditioning,
+ Krea2Config,
+ Krea2Model,
+)
+from diffusion_cli.krea2_sampling import sampleKrea2, timesteps
+from diffusion_cli.quantization import (
+ QuantizedLayerSpec,
+ QuantizationManifest,
+ ScaledFp8Linear,
+)
+from diffusion_cli.qwen_image_vae import (
+ QWEN_IMAGE_LATENT_MEAN,
+ QwenImageVae,
+ convertWanVaeToDiffusers,
+)
+from diffusion_cli.model_inspect import inspectModelSource
+
+
+class Krea2Test(unittest.TestCase):
+ def testNamedProfilesCoexistWithLegacyConfiguration(self):
+ with tempfile.TemporaryDirectory() as temp_dir:
+ root = Path(temp_dir)
+ config_path = root / "diffusion.toml"
+ config_path.write_text(
+ "\n".join(
+ (
+ 'default_model = "krea2-turbo"',
+ "",
+ "[models]",
+ f'tokenizer = "{root / "tokenizer"}"',
+ "",
+ "[model_profiles.krea2-turbo]",
+ 'architecture = "krea2"',
+ 'variant = "turbo"',
+ f'diffusion_model = "{root / "diffusion.safetensors"}"',
+ f'text_encoder = "{root / "text.safetensors"}"',
+ f'vae = "{root / "vae.safetensors"}"',
+ f'tokenizer = "{root / "tokenizer"}"',
+ )
+ ),
+ encoding="utf-8",
+ )
+ config = loadUserConfig(config_path)
+
+ self.assertIn("legacy", config.model_profiles)
+ self.assertEqual(selectModelProfile(config).name, "krea2-turbo")
+ self.assertEqual(config.model_profiles["legacy"].architecture, "z-image")
+
+ def testProfileInspectionReportsCompatibility(self):
+ with tempfile.TemporaryDirectory() as temp_dir:
+ path = Path(temp_dir) / "diffusion.safetensors"
+ save_file(
+ {
+ "blocks.0.weight": torch.zeros(1),
+ "txtfusion.projector.weight": torch.zeros(1),
+ },
+ path,
+ )
+ summary = inspectModelSource(
+ ModelSource(path, "diffusion_model"),
+ ModelProfile("krea2-turbo", "krea2", "turbo"),
+ )
+
+ self.assertEqual(summary.architecture_guess, "krea2_diffusion")
+ self.assertEqual(summary.compatibility, "compatible")
+
+ def testTurboProfileDefaults(self):
+ with tempfile.TemporaryDirectory() as temp_dir:
+ root = Path(temp_dir)
+ tokenizer = root / "tokenizer"
+ tokenizer.mkdir()
+ for name in TOKENIZER_FILES:
+ (tokenizer / name).write_text("{}", encoding="utf-8")
+ profile = ModelProfile(
+ "krea2-turbo",
+ "krea2",
+ "turbo",
+ diffusion_model=root / "diffusion.safetensors",
+ text_encoder=root / "text.safetensors",
+ vae=root / "vae.safetensors",
+ tokenizer=tokenizer,
+ )
+ config = UserConfig(
+ generation=GenerationDefaults(),
+ model_profiles={profile.name: profile},
+ default_model=profile.name,
+ )
+ with patch(
+ "diffusion_cli.config.selectDevice",
+ return_value=torch.device("cpu"),
+ ), patch(
+ "diffusion_cli.config.selectDtype",
+ return_value=torch.float32,
+ ):
+ generation = buildGenerationConfigFromRequest(
+ ImageGenerationRequest("a mug"),
+ config,
+ )
+ self.assertEqual(generation.steps, 8)
+ self.assertEqual(generation.cfg, 0.0)
+ self.assertEqual(generation.mu, 1.15)
+
+ def testScaledFp8LinearMatchesExplicitDequantization(self):
+ layer = ScaledFp8Linear(2, 1, bias=False, dtype=torch.float32)
+ stored = torch.tensor([[1.0, 2.0]], dtype=torch.float8_e4m3fn)
+ layer.weight.data.copy_(stored)
+ layer.weight_scale.data.fill_(2.0)
+ value = torch.tensor([[3.0, 4.0]])
+ expected = value @ (stored.float() * 2.0).transpose(0, 1)
+ self.assertTrue(torch.allclose(layer(value), expected))
+ self.assertEqual(layer.weight.dtype, torch.float8_e4m3fn)
+
+ def testReducedModelReturnsFiniteVelocity(self):
+ config = Krea2Config(
+ features=32,
+ timestep_width=8,
+ text_width=16,
+ heads=4,
+ kv_heads=2,
+ blocks=1,
+ mlp_multiplier=1,
+ patch_size=2,
+ latent_channels=4,
+ text_layers=3,
+ text_heads=2,
+ text_kv_heads=2,
+ )
+ model = Krea2Model(config, dtype=torch.float32)
+ conditioning = Krea2Conditioning(
+ torch.randn(1, 3, 3, 16),
+ torch.ones(1, 3, dtype=torch.bool),
+ )
+ output = model(
+ torch.randn(1, 4, 4, 4),
+ torch.tensor([0.5]),
+ conditioning,
+ )
+ self.assertEqual(output.shape, (1, 4, 4, 4))
+ self.assertTrue(torch.isfinite(output).all())
+
+ def testSamplerSkipsUnconditionalAtZeroCfg(self):
+ calls = []
+
+ class FakeModel:
+ def __call__(self, latent, timestep, conditioning):
+ calls.append(conditioning)
+ return torch.ones_like(latent)
+
+ conditioning = Krea2Conditioning(
+ torch.zeros(1, 1, 12, 2560),
+ torch.ones(1, 1, dtype=torch.bool),
+ )
+ output = sampleKrea2(
+ FakeModel(),
+ conditioning,
+ batch_size=1,
+ height=16,
+ width=16,
+ seed=7,
+ steps=2,
+ cfg=0.0,
+ device="cpu",
+ dtype=torch.float32,
+ )
+ self.assertEqual(len(calls), 2)
+ self.assertEqual(output.shape, (1, 16, 2, 2))
+
+ def testTimestepsDescendAndTurboMuIsStable(self):
+ values = timesteps(16, 8, mu=1.15)
+ self.assertEqual(len(values), 9)
+ self.assertEqual(values[0], 1.0)
+ self.assertEqual(values[-1], 0.0)
+ self.assertTrue(all(a >= b for a, b in zip(values, values[1:])))
+
+ def testVaeDecodeNormalizesTimeAxisAndRange(self):
+ class Result:
+ def __init__(self, sample):
+ self.sample = sample
+
+ class FakeVae(nn.Module):
+ def decode(self, value):
+ return Result(
+ torch.zeros(
+ value.shape[0],
+ 3,
+ 1,
+ value.shape[3] * 8,
+ value.shape[4] * 8,
+ )
+ )
+
+ vae = QwenImageVae(None, "cpu", torch.float32, model=FakeVae())
+ output = vae.decode(torch.zeros(1, 16, 2, 2))
+ self.assertEqual(output.shape, (1, 3, 16, 16))
+ self.assertEqual(float(output.mean()), 0.5)
+
+ def testVaeConversionCoversRepresentativeFamilies(self):
+ state = {
+ "encoder.conv1.weight": torch.zeros(1),
+ "decoder.middle.1.proj.weight": torch.zeros(1),
+ "encoder.downsamples.0.residual.2.weight": torch.zeros(1),
+ "decoder.upsamples.4.shortcut.weight": torch.zeros(1),
+ "conv1.weight": torch.zeros(1),
+ }
+ converted = convertWanVaeToDiffusers(state)
+ self.assertIn("encoder.conv_in.weight", converted)
+ self.assertIn("decoder.mid_block.attentions.0.proj.weight", converted)
+ self.assertIn("encoder.down_blocks.0.conv1.weight", converted)
+ self.assertIn(
+ "decoder.up_blocks.1.resnets.0.conv_shortcut.weight",
+ converted,
+ )
+ self.assertIn("quant_conv.weight", converted)
+
+
+if __name__ == "__main__":
+ unittest.main()
diff --git a/uv.lock b/uv.lock
index 9db8b4e..4b23623 100644
--- a/uv.lock
+++ b/uv.lock
@@ -46,6 +46,80 @@ wheels = [
{ url = "https://files.pythonhosted.org/packages/ef/2f/c5464532e965badff2f4c4c1a3a83f5697f0d7c407ed0cda44aaa99bb451/certifi-2026.6.17-py3-none-any.whl", hash = "sha256:2227dcbaafe0d2f59279d1762ddddc37783ed4354594f194ffc31d20f41fc3db", size = 133289, upload-time = "2026-06-17T10:31:06.348Z" },
]
+[[package]]
+name = "charset-normalizer"
+version = "3.4.9"
+source = { registry = "https://pypi.org/simple" }
+sdist = { url = "https://files.pythonhosted.org/packages/bd/2a/23f34ec9d04624958e137efdc394888716353190e75f25dd22c7a2c7a8aa/charset_normalizer-3.4.9.tar.gz", hash = "sha256:673611bbd43f0810bec0b0f028ddeaaa501190339cac411f347ac76917c3ae7b", size = 152439, upload-time = "2026-07-07T14:34:58.454Z" }
+wheels = [
+ { url = "https://files.pythonhosted.org/packages/0b/e3/85ec501f206fb049259288c1f3506e53876937fb00edb47009348e66756b/charset_normalizer-3.4.9-cp311-cp311-macosx_10_9_universal2.whl", hash = "sha256:0e94703ec9684807f20cfb5eed95c70f67f2a8f21ad620146d7b5a13677b93e5", size = 317075, upload-time = "2026-07-07T14:32:56.021Z" },
+ { url = "https://files.pythonhosted.org/packages/c3/69/2a5385192e67175f7d8bd5ce4f57c24bc956439adeae5c13a99aa28a53d1/charset_normalizer-3.4.9-cp311-cp311-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:2a441ea71902098ffe78c5abe6c494f44160b4af614ed16c3d9a3b1d17fd8ee2", size = 213837, upload-time = "2026-07-07T14:32:57.78Z" },
+ { url = "https://files.pythonhosted.org/packages/b3/46/03ddc7da576d814fe0a36dd1f0fd3258e95404b4b2e3c026b7923d7e133f/charset_normalizer-3.4.9-cp311-cp311-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:304b13570067b2547562e308af560b3963857b1fa90bd6afd978130130fe2d6a", size = 235503, upload-time = "2026-07-07T14:32:59.205Z" },
+ { url = "https://files.pythonhosted.org/packages/4e/6e/de0229a7ef40f6f9d28a837eebf4ec47bdca5dab4e900c84f22919af636a/charset_normalizer-3.4.9-cp311-cp311-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:4773092f8019072343a7447203308b176e10199920eb02d6195e81bbb3274c29", size = 229944, upload-time = "2026-07-07T14:33:00.803Z" },
+ { url = "https://files.pythonhosted.org/packages/a5/34/49b9060e8418b14fb5cba9cf6bfb383111e2538a03a1fb18e66a95aeb3d5/charset_normalizer-3.4.9-cp311-cp311-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:04ce310cb89c15df659582aee80a0603788732a5e017d5bd5c81158106ce249c", size = 221276, upload-time = "2026-07-07T14:33:02.199Z" },
+ { url = "https://files.pythonhosted.org/packages/44/95/80282cce0fae9c3061203d723ee87da996aed79679e65d8935050ee7ca1f/charset_normalizer-3.4.9-cp311-cp311-manylinux_2_31_armv7l.whl", hash = "sha256:c0323c9daef75ef2e5083624b4585018a0c9d5e3b40f607eed81a311270b934b", size = 205260, upload-time = "2026-07-07T14:33:03.698Z" },
+ { url = "https://files.pythonhosted.org/packages/0c/74/2f62c8821b969ea3bd67cc2e6976834f48ca5d12664d2559ebcd9bcfbed7/charset_normalizer-3.4.9-cp311-cp311-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:871ff67ea1aad4dfd91736464934d56b32dac49f9fbe16cddba36198a7b3a0db", size = 217786, upload-time = "2026-07-07T14:33:05.12Z" },
+ { url = "https://files.pythonhosted.org/packages/d9/8d/feabb82cb49fcad14515b1d7d1ca4787b0da7fc723a212bf89bc9e0fac52/charset_normalizer-3.4.9-cp311-cp311-musllinux_1_2_aarch64.whl", hash = "sha256:67830fc78e67501f47bb950471b2dcb9b35b140084429318e862895a8e89c993", size = 216798, upload-time = "2026-07-07T14:33:06.629Z" },
+ { url = "https://files.pythonhosted.org/packages/a5/ff/c946d63bc3786d5b84d960b0f7ab7e25b828486a946b5aa997625bcaf6a6/charset_normalizer-3.4.9-cp311-cp311-musllinux_1_2_armv7l.whl", hash = "sha256:3d92613ec25e43b05f042302531ec0f00b8445190e43325880cbd6ab7c2581da", size = 206429, upload-time = "2026-07-07T14:33:08.006Z" },
+ { url = "https://files.pythonhosted.org/packages/af/ba/5e5007c370702f85d2ef75791fac7943ed41e080364a673b20142e430e3e/charset_normalizer-3.4.9-cp311-cp311-musllinux_1_2_x86_64.whl", hash = "sha256:280081916dc341820640489a66e4696049401ef1cf6dd672f672e70ad915aca3", size = 223066, upload-time = "2026-07-07T14:33:09.783Z" },
+ { url = "https://files.pythonhosted.org/packages/83/d5/9096aa3cf532dfad237861544eb47a0f20d5adbf1039760fed8eaae935d9/charset_normalizer-3.4.9-cp311-cp311-win32.whl", hash = "sha256:ac351b3b8014eead140e77e9717e2992c6bbe30b63bc3422422eb84865412e3d", size = 150456, upload-time = "2026-07-07T14:33:11.217Z" },
+ { url = "https://files.pythonhosted.org/packages/ed/a1/e29995109e455dc8eff8d0fac6ae509be39561318a7cfeac5d33ad029213/charset_normalizer-3.4.9-cp311-cp311-win_amd64.whl", hash = "sha256:6366a16e1a25018694d6a5d784d09b046edc9eac40ea2b54065c3052672516a1", size = 161410, upload-time = "2026-07-07T14:33:12.743Z" },
+ { url = "https://files.pythonhosted.org/packages/4f/8d/1569f4d0032d6ba2a4fe4591c35bf87868c600c41a71eb5c2e1ffa8464c2/charset_normalizer-3.4.9-cp311-cp311-win_arm64.whl", hash = "sha256:1d22856ffbe153a602df38e4a5464f0b748a54002e0d69ac6d2ad0a197cc99ec", size = 152649, upload-time = "2026-07-07T14:33:14.173Z" },
+ { url = "https://files.pythonhosted.org/packages/70/4a/ecbd131485c07fcdfad54e28946d513e3da22ef3b4bd854dcafae54ec739/charset_normalizer-3.4.9-cp312-cp312-macosx_10_13_universal2.whl", hash = "sha256:45b0cc4e3556cd875e09102988d1ab8356c998b596c9fced84547c8138b487a0", size = 319300, upload-time = "2026-07-07T14:33:15.666Z" },
+ { url = "https://files.pythonhosted.org/packages/ec/96/5d9364e3342d69f3a045e1777bc47c85c383e6e9466d561b33fdb419d1f9/charset_normalizer-3.4.9-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:9b2aff1c7b3884512b9512c3eaadd9bab39fb45042ffaaa1dd08ff2b9f8109d9", size = 215802, upload-time = "2026-07-07T14:33:17.031Z" },
+ { url = "https://files.pythonhosted.org/packages/4b/4c/5361f9aa7f2cb58d94f2ab831b3d493f69efb1d239654b4744e3c09527cb/charset_normalizer-3.4.9-cp312-cp312-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:9104ed0bd76a429d46f9ec0dbc9b08ad1d2dcdf2b00a5a0daa1c145329b35b44", size = 237171, upload-time = "2026-07-07T14:33:18.576Z" },
+ { url = "https://files.pythonhosted.org/packages/50/78/ce342ca4ff30b2eb49fe6d9578df85974f90c67d294113e94efdd9664cbd/charset_normalizer-3.4.9-cp312-cp312-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:7b86a2b16095d250c6f58b3d9b2eee6f4147754344f3dab0922f7c9bf7d226c9", size = 233075, upload-time = "2026-07-07T14:33:20.084Z" },
+ { url = "https://files.pythonhosted.org/packages/01/c4/4fa4c8b3097a11f3c5f09a35b72ed6855fb1d332469504962ab7bafcc702/charset_normalizer-3.4.9-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:5e226f6218febc71f6c1fc2fafb91c226f75bdc1d8fb12d66823716e891608fd", size = 224256, upload-time = "2026-07-07T14:33:21.747Z" },
+ { url = "https://files.pythonhosted.org/packages/87/3a/ad914516df7e358a81aae018caa5e0470ba827fa6d763b1d2e87d920a5f6/charset_normalizer-3.4.9-cp312-cp312-manylinux_2_31_armv7l.whl", hash = "sha256:90c44bc373b7687f6948b693cceaea1348ae0975d7474746559494468e3c1d84", size = 208784, upload-time = "2026-07-07T14:33:23.313Z" },
+ { url = "https://files.pythonhosted.org/packages/d7/74/3c12f9755717dfe5c5c87da63f35d765fa0c00382ec26bf23f7fae34f2ba/charset_normalizer-3.4.9-cp312-cp312-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:9cdef90ae47919cae358d8ab15797a800ed41da7aba5d72419fb510729e2ed4b", size = 219928, upload-time = "2026-07-07T14:33:24.814Z" },
+ { url = "https://files.pythonhosted.org/packages/33/9a/895095b83e7907abd6d3d99aad3a38ad0d9686cc186cb0c94c24320fe63e/charset_normalizer-3.4.9-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:60f44ade2cf573dad7a277e6f8ca9a51a21dda572b13bd7d8539bb3cd5dbedde", size = 218489, upload-time = "2026-07-07T14:33:26.42Z" },
+ { url = "https://files.pythonhosted.org/packages/a1/34/ef5c05f412f42520d7709b7d3784d19640839eb7366ded1755511585429f/charset_normalizer-3.4.9-cp312-cp312-musllinux_1_2_armv7l.whl", hash = "sha256:a1786910334ed46ab1dd73222f2cd1e05c2c3bb39f6dddb4f8b36fc382058a39", size = 210267, upload-time = "2026-07-07T14:33:27.952Z" },
+ { url = "https://files.pythonhosted.org/packages/83/dc/9b29fa4412b318bf3bfea985c35d67eb55e04b59a7c3f2237168b0e0be6f/charset_normalizer-3.4.9-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:03d07803992c6c7bbc976327f34b18b6160327fc81cb82c9d504720ac0be3b62", size = 226030, upload-time = "2026-07-07T14:33:29.397Z" },
+ { url = "https://files.pythonhosted.org/packages/0e/42/6dbc00b8cd16011691203e33570fa42ed5746599a2e878112d16eab403a3/charset_normalizer-3.4.9-cp312-cp312-win32.whl", hash = "sha256:78841cccf1af7b40f6f716338d50c0902dbe88d9f800b3c973b7a9a0a693a642", size = 151185, upload-time = "2026-07-07T14:33:30.781Z" },
+ { url = "https://files.pythonhosted.org/packages/80/cc/f920afd1a23c58ccd53c1d36085a71893a4737ff5e66e0371efab6809850/charset_normalizer-3.4.9-cp312-cp312-win_amd64.whl", hash = "sha256:4b3dac63058cc36820b0dd072f89898604e2d39686fe05321729d00d8ac185a0", size = 162557, upload-time = "2026-07-07T14:33:32.176Z" },
+ { url = "https://files.pythonhosted.org/packages/f0/e6/0386d43a261ff4e4b30c5857af7df877254b46bec7b9d1b74b6bf969a90b/charset_normalizer-3.4.9-cp312-cp312-win_arm64.whl", hash = "sha256:78fa18e436a1a0e58dbd7e02fc4473f3f32cceb12df9dfca542d075961c307d2", size = 152665, upload-time = "2026-07-07T14:33:33.711Z" },
+ { url = "https://files.pythonhosted.org/packages/b2/06/97ec2aeae780b31d742b6352218b43841a6871e2564578ca522dce4a45c3/charset_normalizer-3.4.9-cp313-cp313-macosx_10_13_universal2.whl", hash = "sha256:440eede837960000d74978f0eba527be106b5b9aee0daf779d395276ed0b0614", size = 317688, upload-time = "2026-07-07T14:33:35.408Z" },
+ { url = "https://files.pythonhosted.org/packages/d0/39/8ff066c672434225f8d25f8b739f992af250944392173dcc88362681c9bf/charset_normalizer-3.4.9-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:21e764fd1e70b6a3e205a0e46f3051701f98a8cb3fad66eeb80e48bb502f8698", size = 214982, upload-time = "2026-07-07T14:33:36.996Z" },
+ { url = "https://files.pythonhosted.org/packages/92/8f/3a47a3667c83c2df9483d91644c6c107de3bf8874aa1793da9d3012eb986/charset_normalizer-3.4.9-cp313-cp313-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:e4fd89cc178bced6ad29cb3e6dd4aa63fa5017c3524dbd0b25998fb64a87cc8b", size = 236460, upload-time = "2026-07-07T14:33:38.536Z" },
+ { url = "https://files.pythonhosted.org/packages/f1/60/b22cdbee7e4013dab8b0d7647fc6181120fbbbc8f7025c226d15bd5a47fc/charset_normalizer-3.4.9-cp313-cp313-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:bd47ba7fc3ca94896759ea0109775132d3e7ab921fbf54038e1bab2e46c313c9", size = 232003, upload-time = "2026-07-07T14:33:40.059Z" },
+ { url = "https://files.pythonhosted.org/packages/ea/f8/72eb13dcabe7257035cea8aefd922caad2f110d252bf9f67c4c2ca763aee/charset_normalizer-3.4.9-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:84fd18bcc17526fc2b3c1af7d2b9217d32c9c04448c16ec693b9b4f1985c3d33", size = 223149, upload-time = "2026-07-07T14:33:41.631Z" },
+ { url = "https://files.pythonhosted.org/packages/b0/3e/faee8f9de92b14ee1198e9163252bb15efee7301b31256a3b6d9ebfdd0dd/charset_normalizer-3.4.9-cp313-cp313-manylinux_2_31_armv7l.whl", hash = "sha256:5b10cd92fc5c498b35a8635df6d5a100207f88b63a4dc1de7ef9a548e1e2cd63", size = 207901, upload-time = "2026-07-07T14:33:43.209Z" },
+ { url = "https://files.pythonhosted.org/packages/3a/25/45f30093ae27dd7b92a793b61882a38685f993700113ca36e0c9c14965e1/charset_normalizer-3.4.9-cp313-cp313-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:a4fbdde9dd4a9ce5fd52c2b3a347bb50cc89483ef783f1cb00d408c13f7a96c0", size = 219176, upload-time = "2026-07-07T14:33:44.725Z" },
+ { url = "https://files.pythonhosted.org/packages/48/18/c8f397329c35e32f6a837e488986f4ae03bd2abebc453b48714991630c2f/charset_normalizer-3.4.9-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:416c229f77e5ea25b3dfd4b582f8d73d7e43c22320302b9ab128a2d3a0b38efe", size = 217356, upload-time = "2026-07-07T14:33:46.192Z" },
+ { url = "https://files.pythonhosted.org/packages/86/7e/5ce0bba863470fd1902d5e5843968951bddf38abe4742fc97116ef4598b3/charset_normalizer-3.4.9-cp313-cp313-musllinux_1_2_armv7l.whl", hash = "sha256:75286256590a6320cf106a0d28970d3560aad9ee09aa7b34fb40524792436d35", size = 209614, upload-time = "2026-07-07T14:33:47.705Z" },
+ { url = "https://files.pythonhosted.org/packages/6c/ef/2473d3c4d869155be4af1191111d59c4d5c4e0173026f7e85b176e23bf65/charset_normalizer-3.4.9-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:69b157c5d3292bcd443faca052f3096f637f1e074b98212a933c074ae23dc3b8", size = 224991, upload-time = "2026-07-07T14:33:49.238Z" },
+ { url = "https://files.pythonhosted.org/packages/d0/a3/53ddae3db108a088156aa8ddfafd411ebbc1340f48c5573f697b27f69a39/charset_normalizer-3.4.9-cp313-cp313-win32.whl", hash = "sha256:51307f5c71007673a2bf8232ad973483d281e74cb99c8c5a990af1eefa6277d9", size = 150622, upload-time = "2026-07-07T14:33:50.711Z" },
+ { url = "https://files.pythonhosted.org/packages/e8/ef/6953a77c7cf2c2ff9998e6f575ab3e380119f100223381565a4f94c1f836/charset_normalizer-3.4.9-cp313-cp313-win_amd64.whl", hash = "sha256:fe2c7201c642b7c308f1675355ad7ff7b66acfe3541625efe5a3ad38f29d6115", size = 161947, upload-time = "2026-07-07T14:33:52.197Z" },
+ { url = "https://files.pythonhosted.org/packages/6e/fb/d560d1d1555debbfe7849d9cac6145c1b537709d79576bf22557ed803b82/charset_normalizer-3.4.9-cp313-cp313-win_arm64.whl", hash = "sha256:611057cc5d5c0afc743ba8be6bd828c17e0aaa8643f9d0a9b9bb7dea80eb8012", size = 152594, upload-time = "2026-07-07T14:33:53.486Z" },
+ { url = "https://files.pythonhosted.org/packages/7e/8d/496817fa0944239ecae662dd57ea765cfeaec6a735f9f025d4b7b72e7143/charset_normalizer-3.4.9-cp314-cp314-macosx_10_15_universal2.whl", hash = "sha256:0327fcd59a935777d83410750c50600ee9571af2846f71ce40f25b13da1ef380", size = 317253, upload-time = "2026-07-07T14:33:54.994Z" },
+ { url = "https://files.pythonhosted.org/packages/2b/f9/ef4a69ea338ad3c0deceea0f5f7d2380ae8b52132b06d652cb0d2cd86706/charset_normalizer-3.4.9-cp314-cp314-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:8a79d9f4d8001473a30c163556b3c3bfebec837495a412dde78b51672f6134f9", size = 215898, upload-time = "2026-07-07T14:33:56.334Z" },
+ { url = "https://files.pythonhosted.org/packages/8c/e7/5ddfd76fc061eb52de219658a4aa431cbacadf0a0219c8854f00da50d289/charset_normalizer-3.4.9-cp314-cp314-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:33bdcc2a32c0a0e861f60841a512c8acc658c87c2ac59d89e3a46dacf7d866e4", size = 236718, upload-time = "2026-07-07T14:33:57.9Z" },
+ { url = "https://files.pythonhosted.org/packages/49/ba/768fa3f36048d81c477a0ce61f813bc1454d80917ccfe550abd9f44f5e24/charset_normalizer-3.4.9-cp314-cp314-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:f840ed6d8ecba8255df8c42b87fadeda98ddfc6eeec05e2dc66e26d46dd6f58a", size = 232519, upload-time = "2026-07-07T14:33:59.811Z" },
+ { url = "https://files.pythonhosted.org/packages/f4/c4/b3e049d2aa3766180c78507110543d9d50894cc97f57de543f1be521dcdc/charset_normalizer-3.4.9-cp314-cp314-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:c25fe15c70c59eb7c5ce8c06a1f3fa1da0ecc5ea1e7a5922c40fd2fa9b0d5046", size = 223143, upload-time = "2026-07-07T14:34:01.517Z" },
+ { url = "https://files.pythonhosted.org/packages/19/79/55c32d06d76ae4feafe053f061f3e3ab70bcf19f4007797ce8c3efda7830/charset_normalizer-3.4.9-cp314-cp314-manylinux_2_31_armv7l.whl", hash = "sha256:f7fb7d750cfa0a070d2c24e831fd3481019a60dd317ea2b39acbcebc08b6ed81", size = 206742, upload-time = "2026-07-07T14:34:03.04Z" },
+ { url = "https://files.pythonhosted.org/packages/10/e0/47c079dd82d217c807479cd59ffd30af56307ea31c108b75758970459ad3/charset_normalizer-3.4.9-cp314-cp314-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:4d1c96a7a18b9690a4d46df09e3e3382406ae3213727cd1019ebade1c4a81917", size = 219191, upload-time = "2026-07-07T14:34:04.657Z" },
+ { url = "https://files.pythonhosted.org/packages/42/ab/b9bc2e77d6b44a7e46ef62ec5cac1c9a6ba7b9135a5d560f002696ec9995/charset_normalizer-3.4.9-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:a4cfde78a9f2880208d16a93b795726a3017d5977e08d1e162a7a31322479c41", size = 218328, upload-time = "2026-07-07T14:34:06.115Z" },
+ { url = "https://files.pythonhosted.org/packages/f1/78/c9c71d599f5aa2d42bcdd35cbbd46d7f535351a57e40ff7d8e5a7e219401/charset_normalizer-3.4.9-cp314-cp314-musllinux_1_2_armv7l.whl", hash = "sha256:d4d6fcde76f94f5cb9e43e9e9a61f16dacefd228cbbf6f1a09bd9b219a92f1a1", size = 207406, upload-time = "2026-07-07T14:34:07.554Z" },
+ { url = "https://files.pythonhosted.org/packages/f6/39/c914445c321a845097ce4f6ac7de9a18228a77b766272125a1ce00d851eb/charset_normalizer-3.4.9-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:898f0e9068ca27d37f8e83a5b962821df851532e6c4a7d615c1c033f9da6eedf", size = 225157, upload-time = "2026-07-07T14:34:09.061Z" },
+ { url = "https://files.pythonhosted.org/packages/9b/f2/c0d4b8508565a36bc5c624e88ed297f5b0b1095011034d7f5b83a69908b5/charset_normalizer-3.4.9-cp314-cp314-win32.whl", hash = "sha256:c1c948747b03be832dceed96ca815cef7360de9aa19d37c730f8e3f6101aca48", size = 151095, upload-time = "2026-07-07T14:34:10.901Z" },
+ { url = "https://files.pythonhosted.org/packages/49/fd/a1d26144398c67486422a72bf5812cda22cb4ccfcd95a290fb41ceb4b8e2/charset_normalizer-3.4.9-cp314-cp314-win_amd64.whl", hash = "sha256:16b65ea0f2465b6fb52aa22de5eca612aa964ddfec00a912e26f4656cbef890b", size = 162796, upload-time = "2026-07-07T14:34:12.47Z" },
+ { url = "https://files.pythonhosted.org/packages/20/95/d75e82f8ce9fd323ebf059c16c9aadefb22a1ecde13b7840b35835e4886c/charset_normalizer-3.4.9-cp314-cp314-win_arm64.whl", hash = "sha256:40a126142a56b2dfc0aacbad1de8310cbf60da7656db0e6b16eebd48e3e93519", size = 153334, upload-time = "2026-07-07T14:34:14.044Z" },
+ { url = "https://files.pythonhosted.org/packages/00/5e/17398df3a139985ba9d11ed072531986f408c8fca952835ef1ab1820c02b/charset_normalizer-3.4.9-cp314-cp314t-macosx_10_15_universal2.whl", hash = "sha256:609b3ba8fcc0fb5ab7af00719d0fb6ad0cb518e48e7712d12fd68f1327951198", size = 338848, upload-time = "2026-07-07T14:34:15.688Z" },
+ { url = "https://files.pythonhosted.org/packages/cd/91/7253a32e86b7e1d1239b1b36ba6dd0f021a21107ab33054b53119cc083b9/charset_normalizer-3.4.9-cp314-cp314t-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:51447e9aa2684679af07ca5021c3db526e0284347ebf4ffcec1154c3350cfe32", size = 223022, upload-time = "2026-07-07T14:34:17.248Z" },
+ { url = "https://files.pythonhosted.org/packages/cb/32/2e64bd2be10e89c61e57ebe6a93fd98ae88eb7ebe414b5121f22c96c69eb/charset_normalizer-3.4.9-cp314-cp314t-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:cc1b0fff8ead343dae06305f954eb8468ba0ec1a97881f42489d198e4ce3c632", size = 241590, upload-time = "2026-07-07T14:34:18.813Z" },
+ { url = "https://files.pythonhosted.org/packages/3d/ef/d96ec496cfea0c21db43b0ad03891308b02388d054cc902cf0e5a1ad6a88/charset_normalizer-3.4.9-cp314-cp314t-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:fa36ec09ef71d158186bc79e359ff5fdd6e7996fe8ab638f00d6b93139ba4fcf", size = 239584, upload-time = "2026-07-07T14:34:20.52Z" },
+ { url = "https://files.pythonhosted.org/packages/d4/ce/9af95f7876194bd7a14e3dfe4a4de2e0bff02666a3910d72beafd06cc297/charset_normalizer-3.4.9-cp314-cp314t-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:df115d4d83168fdf2cae48ef1ff6d1cb4c466364e30861b37121de0f3bf1b990", size = 230224, upload-time = "2026-07-07T14:34:22.189Z" },
+ { url = "https://files.pythonhosted.org/packages/52/94/af74dde74a3996bd959c350709bfe50e297823d70a8c1cbd54b838880863/charset_normalizer-3.4.9-cp314-cp314t-manylinux_2_31_armv7l.whl", hash = "sha256:f86c6358749bd4fda175388691e3ba8c46e24c5347d0afd20f9b7edfc9faf07d", size = 212667, upload-time = "2026-07-07T14:34:23.857Z" },
+ { url = "https://files.pythonhosted.org/packages/ee/f0/f1c4fe746c395922961b5916ed1d7d6e7d4c84851d19ed43cc89980ec953/charset_normalizer-3.4.9-cp314-cp314t-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:32286a2c8d167e897177b673176c1e3e00d4057caf5d2b64eef9a3666b03018e", size = 227179, upload-time = "2026-07-07T14:34:25.586Z" },
+ { url = "https://files.pythonhosted.org/packages/e4/56/6c745619ac397e8871e2bcd3cea1eec86b877488f33888b3aef5c3ed506e/charset_normalizer-3.4.9-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:83aed2c10721ddd90f68140685391b50811a880af20654c59af6b6c66c40513c", size = 225372, upload-time = "2026-07-07T14:34:27.212Z" },
+ { url = "https://files.pythonhosted.org/packages/78/ad/98aae8630ac71f16711968e38a5acfecce41b778bf2f0312851020f565a8/charset_normalizer-3.4.9-cp314-cp314t-musllinux_1_2_armv7l.whl", hash = "sha256:cd6c3d4b783c556fa00bf540854e42f135e2f256abd29669fcd0da0f2dec79c2", size = 215222, upload-time = "2026-07-07T14:34:28.774Z" },
+ { url = "https://files.pythonhosted.org/packages/f7/40/9593d54209765207a7f11073c06494c1721e4ca4a0a426c597679bf7f91e/charset_normalizer-3.4.9-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:ee2f2a527e3c1a6e6411eb4209642e138b544a2d72fe5d0d76daf77b24063534", size = 231958, upload-time = "2026-07-07T14:34:30.345Z" },
+ { url = "https://files.pythonhosted.org/packages/b1/27/693ee5e8a18191eb38647360c51cd505013e2bd3b366aa43fd5344c21e3c/charset_normalizer-3.4.9-cp314-cp314t-win32.whl", hash = "sha256:0d861473f743244d349b50f850d10eb87aeb22bbdcc8e64f79273c94af5a8226", size = 155580, upload-time = "2026-07-07T14:34:31.884Z" },
+ { url = "https://files.pythonhosted.org/packages/80/3f/bd97d3d9c613013d07cb7733d299385b41df37f0471310f5a73dc359f0b8/charset_normalizer-3.4.9-cp314-cp314t-win_amd64.whl", hash = "sha256:9b8e0f3107e2200b76f6054de99016eac3ee6762713587b36baaa7e4bd2ae177", size = 167620, upload-time = "2026-07-07T14:34:33.438Z" },
+ { url = "https://files.pythonhosted.org/packages/3d/c6/eee9dca4439b1061f76373f06ea855678cc4a64c1c3c90b50e479edbb8eb/charset_normalizer-3.4.9-cp314-cp314t-win_arm64.whl", hash = "sha256:19ac87f93086ce37b86e098888555c4b4bc48102279bae3350098c0ed664b501", size = 158037, upload-time = "2026-07-07T14:34:35.018Z" },
+ { url = "https://files.pythonhosted.org/packages/98/2b/f97f1c193fb855c345d678f5077d6926034db0722df74c8f057020e05a25/charset_normalizer-3.4.9-py3-none-any.whl", hash = "sha256:68e5f26a1ad57ded6d1cfb85331d1c1a195314756471d97758c48498bb4dcdf5", size = 64538, upload-time = "2026-07-07T14:34:56.993Z" },
+]
+
[[package]]
name = "click"
version = "8.4.2"
@@ -135,11 +209,33 @@ nvtx = [
{ name = "nvidia-nvtx", marker = "sys_platform == 'linux' or sys_platform == 'win32'" },
]
+[[package]]
+name = "diffusers"
+version = "0.39.0"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+ { name = "filelock" },
+ { name = "httpx" },
+ { name = "huggingface-hub" },
+ { name = "importlib-metadata" },
+ { name = "numpy", version = "2.4.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.12'" },
+ { name = "numpy", version = "2.5.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.12'" },
+ { name = "pillow" },
+ { name = "regex" },
+ { name = "requests" },
+ { name = "safetensors" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/1a/81/6095237b86a3116c4789f28c4435d5296c00c0fc74ffde99008fd6b3a36c/diffusers-0.39.0.tar.gz", hash = "sha256:14bb1d98c85a0e463d734c99aaa73b480a7bc9bad22af30fbf730ef8f09c1d67", size = 4651240, upload-time = "2026-07-03T08:48:47.904Z" }
+wheels = [
+ { url = "https://files.pythonhosted.org/packages/3f/3f/7469c46e9d22307ea686bab687d70e6bf328722952f9d10339f5e913e608/diffusers-0.39.0-py3-none-any.whl", hash = "sha256:912aca51b5787365110806e984d5555735bf8a461073bb8459029d0bca7870ef", size = 5631176, upload-time = "2026-07-03T08:48:45.337Z" },
+]
+
[[package]]
name = "diffusion-cli"
version = "0.1.0"
source = { editable = "." }
dependencies = [
+ { name = "diffusers" },
{ name = "flask" },
{ name = "numpy", version = "2.4.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.12'" },
{ name = "numpy", version = "2.5.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.12'" },
@@ -154,6 +250,7 @@ dependencies = [
[package.metadata]
requires-dist = [
+ { name = "diffusers" },
{ name = "flask" },
{ name = "numpy" },
{ name = "pillow" },
@@ -297,6 +394,18 @@ wheels = [
{ url = "https://files.pythonhosted.org/packages/1e/5e/d4e9f1a599fb8e573b7b87160658329fbf28d19eac2718f51fc3def3aa5a/idna-3.18-py3-none-any.whl", hash = "sha256:7f952cbe720b688055e3f87de14f5c3e5fdaa8bc3928985c4077ca689de849a2", size = 65455, upload-time = "2026-06-02T14:34:06.319Z" },
]
+[[package]]
+name = "importlib-metadata"
+version = "9.0.0"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+ { name = "zipp" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/a9/01/15bb152d77b21318514a96f43af312635eb2500c96b55398d020c93d86ea/importlib_metadata-9.0.0.tar.gz", hash = "sha256:a4f57ab599e6a2e3016d7595cfd72eb4661a5106e787a95bcc90c7105b831efc", size = 56405, upload-time = "2026-03-20T06:42:56.999Z" }
+wheels = [
+ { url = "https://files.pythonhosted.org/packages/38/3d/2d244233ac4f76e38533cfcb2991c9eb4c7bf688ae0a036d30725b8faafe/importlib_metadata-9.0.0-py3-none-any.whl", hash = "sha256:2d21d1cc5a017bd0559e36150c21c830ab1dc304dedd1b7ea85d20f45ef3edd7", size = 27789, upload-time = "2026-03-20T06:42:55.665Z" },
+]
+
[[package]]
name = "itsdangerous"
version = "2.2.0"
@@ -981,6 +1090,21 @@ wheels = [
{ url = "https://files.pythonhosted.org/packages/4d/f3/f5ec86839bbabe33b6dee649b62ff9a445d43de6b0ad780cf6b83c56f61e/regex-2026.6.28-cp314-cp314t-win_arm64.whl", hash = "sha256:4da6f6a72f8700b97a1a765e837fb7d5750bfd9f13acea7bae498f573e3a70a8", size = 283338, upload-time = "2026-06-28T19:56:52.879Z" },
]
+[[package]]
+name = "requests"
+version = "2.34.2"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+ { name = "certifi" },
+ { name = "charset-normalizer" },
+ { name = "idna" },
+ { name = "urllib3" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/ac/c3/e2a2b89f2d3e2179abd6d00ebd70bff6273f37fb3e0cc209f48b39d00cbf/requests-2.34.2.tar.gz", hash = "sha256:f288924cae4e29463698d6d60bc6a4da69c89185ad1e0bcc4104f584e960b9ed", size = 142856, upload-time = "2026-05-14T19:25:27.735Z" }
+wheels = [
+ { url = "https://files.pythonhosted.org/packages/a0/f4/c67b0b3f1b9245e8d266f0f112c500d50e5b4e83cb6f3b71b6528104182a/requests-2.34.2-py3-none-any.whl", hash = "sha256:2a0d60c172f83ac6ab31e4554906c0f3b3588d37b5cb939b1c061f4907e278e0", size = 73075, upload-time = "2026-05-14T19:25:26.443Z" },
+]
+
[[package]]
name = "rich"
version = "15.0.0"
@@ -1320,6 +1444,15 @@ wheels = [
{ url = "https://files.pythonhosted.org/packages/49/d3/b8441a820a491ddfc024b0b0cf0393375b75ea13866d9c66727e54c2fc80/typing_extensions-4.16.0-py3-none-any.whl", hash = "sha256:481caa481374e813c1b176ada14e97f1f67a4539ce9cfeb3f350d78d6370c2e8", size = 45571, upload-time = "2026-07-02T08:40:04.659Z" },
]
+[[package]]
+name = "urllib3"
+version = "2.7.0"
+source = { registry = "https://pypi.org/simple" }
+sdist = { url = "https://files.pythonhosted.org/packages/53/0c/06f8b233b8fd13b9e5ee11424ef85419ba0d8ba0b3138bf360be2ff56953/urllib3-2.7.0.tar.gz", hash = "sha256:231e0ec3b63ceb14667c67be60f2f2c40a518cb38b03af60abc813da26505f4c", size = 433602, upload-time = "2026-05-07T16:13:18.596Z" }
+wheels = [
+ { url = "https://files.pythonhosted.org/packages/7f/3e/5db95bcf282c52709639744ca2a8b149baccf648e39c8cc87553df9eae0c/urllib3-2.7.0-py3-none-any.whl", hash = "sha256:9fb4c81ebbb1ce9531cce37674bbc6f1360472bc18ca9a553ede278ef7276897", size = 131087, upload-time = "2026-05-07T16:13:17.151Z" },
+]
+
[[package]]
name = "werkzeug"
version = "3.1.8"
@@ -1331,3 +1464,12 @@ sdist = { url = "https://files.pythonhosted.org/packages/dd/b2/381be8cfdee792dd1
wheels = [
{ url = "https://files.pythonhosted.org/packages/93/8c/2e650f2afeb7ee576912636c23ddb621c91ac6a98e66dc8d29c3c69446e1/werkzeug-3.1.8-py3-none-any.whl", hash = "sha256:63a77fb8892bf28ebc3178683445222aa500e48ebad5ec77b0ad80f8726b1f50", size = 226459, upload-time = "2026-04-02T18:49:12.72Z" },
]
+
+[[package]]
+name = "zipp"
+version = "4.1.0"
+source = { registry = "https://pypi.org/simple" }
+sdist = { url = "https://files.pythonhosted.org/packages/b9/d8/eab98a517c14134c0b2eb4e2387bc5f457334293ec5d2dd3857ec2966802/zipp-4.1.0.tar.gz", hash = "sha256:4cb57381f544315db7688e976e922a2b18cdb513d21cc194eb42232ba2a3e602", size = 26214, upload-time = "2026-05-18T20:08:57.967Z" }
+wheels = [
+ { url = "https://files.pythonhosted.org/packages/3a/13/547360d81e6d88d58492968ffda9f9542854f11310ee556fef14260cc886/zipp-4.1.0-py3-none-any.whl", hash = "sha256:25ad4e16390cd314347dd8f1de67a2ac538ae658ed4ab9db16029c07c188e97f", size = 10238, upload-time = "2026-05-18T20:08:57.045Z" },
+]