Changes
diff --git a/README.md b/README.md
index fd4e076..e227b56 100644
--- a/README.md
+++ b/README.md
@@ -34,6 +34,8 @@ batch_size = 1
steps = 10
cfg = 1.0
output = "output.png"
+output_extension = "png"
+output_quality = 95
device = "cuda"
dtype = "auto"
```
@@ -49,6 +51,26 @@ uv run diffusion-cli --prompt "a ceramic mug" --steps 12
In this example, `--steps 12` overrides any `generation.steps` value in
the config file for that run only.
+PNG is the default output format and does not require ImageMagick. To
+write a smaller lossy or compressed format, request a final extension
+and encoder quality:
+
+```bash
+uv run diffusion-cli \
+ --prompt "a ceramic mug" \
+ --output output.png \
+ --output-extension jpg \
+ --output-quality 85
+```
+
+This writes `output.jpg`; the configured output path still controls the
+directory and base name. Supported extensions are `png`, `jpg`, `jpeg`,
+`webp`, and `avif`. The `jpeg` alias is normalized to `jpg`.
+
+Non-PNG output uses the `magick` executable from ImageMagick after image
+generation completes. AVIF support depends on the local ImageMagick
+build and installed delegates.
+
## HTTP API server
Start the SillyTavern compatibility API with an explicit profile:
diff --git a/designs/design-5-format.md b/designs/design-5-format.md
new file mode 100644
index 0000000..7643650
--- /dev/null
+++ b/designs/design-5-format.md
@@ -0,0 +1,1187 @@
+# Output Format Design
+
+## Purpose
+
+This document describes how `diffusion-cli` should let users choose the
+final image file format and compression quality.
+
+The current generation path produces RGB image tensors and writes or
+returns PNG images. PNG is a good default because it is lossless and is
+widely supported, but it can be inconveniently large. Users should be
+able to request smaller formats such as JPEG, WebP, or AVIF without
+changing the model inference path.
+
+The design keeps inference and compression separate:
+
+- Model inference still produces image tensors.
+- The internal no-loss handoff format remains PNG when an external
+ compressor is needed.
+- A post-processing step invokes ImageMagick only when the requested
+ final format requires it.
+- CLI and web UI users receive the final requested format.
+
+This separation matters because image generation and image compression
+have different failure modes. A model failure means inference did not
+complete. A compression failure means inference completed, but the final
+encoding could not be produced. The user should see a clear error for
+the exact failing stage.
+
+## Goals
+
+The feature must:
+
+- Add configurable output format and quality values.
+- Support config file defaults.
+- Support CLI overrides for one-shot generation.
+- Expose format and quality controls in the web UI.
+- Keep PNG as the default output format.
+- Avoid requiring ImageMagick for default PNG output.
+- Use ImageMagick for non-PNG formats.
+- Surface ImageMagick failures clearly to CLI and web UI users.
+- Return correct MIME types from first-party UI generation responses.
+- Preserve existing base64 response shapes while encoding the final
+ compressed image bytes.
+- Add focused tests for validation, command construction, errors, and
+ UI request handling.
+
+The feature should not:
+
+- Change the diffusion sampling algorithm.
+- Change VAE decode behavior.
+- Add image-to-image support.
+- Add upload support.
+- Add an async conversion job queue.
+- Add a frontend build system.
+- Replace Pillow for the existing PNG path.
+- Let arbitrary user strings become output command arguments without
+ validation.
+
+## External References
+
+The implementation should be checked against these references:
+
+- ImageMagick command line usage:
+ [ImageMagick command-line processing](https://imagemagick.org/script/command-line-processing.php).
+- ImageMagick quality option:
+ [ImageMagick command-line options: quality](https://imagemagick.org/script/command-line-options.php#quality).
+- Python subprocess safety and error handling:
+ [subprocess documentation](https://docs.python.org/3/library/subprocess.html).
+- Python temporary file handling:
+ [tempfile documentation](https://docs.python.org/3/library/tempfile.html).
+- MIME type registration background:
+ [IANA media types](https://www.iana.org/assignments/media-types/media-types.xhtml).
+- HTTP response semantics:
+ [RFC 9110](https://www.rfc-editor.org/rfc/rfc9110.html).
+
+## Terminology
+
+### Output Format
+
+The output format is the final image encoding given to the user.
+
+Examples:
+
+- `png`
+- `jpg`
+- `webp`
+- `avif`
+
+The output format is not the same thing as the model output. The model
+does not directly create JPEG or AVIF files. The model creates numeric
+image data, and the image I/O layer encodes that data.
+
+### Extension Name
+
+The extension name is the suffix used for output file names.
+
+For example, if the output path is:
+
+```text
+image.png
+```
+
+and the requested extension is:
+
+```text
+jpg
+```
+
+then the final output path should be:
+
+```text
+image.jpg
+```
+
+The extension name must be normalized before use. This avoids treating
+`JPG`, `.jpg`, and `jpeg` as different internal formats.
+
+### Quality
+
+Quality is an integer passed to ImageMagick's `-quality` option.
+
+For lossy formats, higher values usually mean larger files and fewer
+visible compression artifacts. For some formats and ImageMagick builds,
+the exact interpretation can vary. The application should therefore
+describe this setting as an encoder quality value rather than promising
+a specific file size.
+
+### Final Image
+
+The final image is the image after optional compression.
+
+For PNG output, the final image is produced directly by the existing
+Pillow path. For non-PNG output, the final image is produced by
+ImageMagick from a temporary PNG input.
+
+## User-Facing Behavior
+
+### Default Behavior
+
+Existing users should not need ImageMagick unless they opt into a
+non-PNG format.
+
+With no new config and no new CLI options:
+
+```bash
+uv run diffusion-cli --prompt "a ceramic mug"
+```
+
+the command should continue to write:
+
+```text
+output.png
+```
+
+This behavior avoids turning ImageMagick into a hard runtime dependency
+for the default case.
+
+### CLI Behavior
+
+Add two generation options:
+
+```text
+--output-extension EXTENSION
+--output-quality QUALITY
+```
+
+`--output-extension`
+
+: Selects the final file extension and image format. The value should
+ accept either a bare extension or a dotted extension. For example,
+ `jpg` and `.jpg` should both be valid.
+
+`--output-quality`
+
+: Selects the ImageMagick quality value for non-PNG formats. The value
+ must be an integer between 1 and 100.
+
+Example:
+
+```bash
+uv run diffusion-cli \
+ --prompt "a ceramic mug" \
+ --output output.png \
+ --output-extension jpg \
+ --output-quality 85
+```
+
+The command should write:
+
+```text
+output.jpg
+```
+
+The user-provided output path still controls the directory and base
+name. The selected extension controls the suffix. This means users can
+keep existing config such as `output = "output.png"` and only change
+the format setting.
+
+### Config File Behavior
+
+Add two keys under `[generation]`:
+
+```toml
+[generation]
+output_extension = "jpg"
+output_quality = 85
+```
+
+These values should follow the same precedence model as other
+generation defaults:
+
+1. Command line option.
+2. Config file value.
+3. Built-in default.
+
+The built-in defaults should be:
+
+```text
+output_extension = "png"
+output_quality = 95
+```
+
+The exact default quality does not matter for PNG when PNG bypasses
+ImageMagick. It does matter as the default for users who select a
+non-PNG format without setting quality. A value of `95` preserves high
+visual quality while still allowing formats such as JPEG, WebP, and
+AVIF to compress better than PNG in many cases.
+
+### Web UI Behavior
+
+The web UI should expose format and quality next to the other generation
+settings.
+
+The controls should be:
+
+- A `select` for output format.
+- A numeric input or slider for quality.
+
+The format options should be visible labels backed by normalized values:
+
+```text
+PNG -> png
+JPEG -> jpg
+WebP -> webp
+AVIF -> avif
+```
+
+The UI should send these fields to `/ui/api/txt2img`:
+
+```json
+{
+ "prompt": "a ceramic mug",
+ "output_extension": "jpg",
+ "output_quality": 85
+}
+```
+
+The UI defaults endpoint should include:
+
+```json
+{
+ "output_extension": "png",
+ "output_quality": 95,
+ "output_formats": [
+ {"value": "png", "label": "PNG"},
+ {"value": "jpg", "label": "JPEG"},
+ {"value": "webp", "label": "WebP"},
+ {"value": "avif", "label": "AVIF"}
+ ]
+}
+```
+
+The generated image response should use the final MIME type:
+
+```json
+{
+ "images": [
+ {
+ "mime_type": "image/jpeg",
+ "data": "...base64...",
+ "seed": 123
+ }
+ ]
+}
+```
+
+The browser already builds data URLs from `mime_type` and `data`, so the
+gallery can keep using the same display flow.
+
+## Supported Formats
+
+The first implementation should use an explicit allowlist:
+
+```text
+png
+jpg
+jpeg
+webp
+avif
+```
+
+Internally, normalize this list to canonical extensions:
+
+```text
+png -> png
+jpg -> jpg
+jpeg -> jpg
+webp -> webp
+avif -> avif
+```
+
+The application should reject unsupported extensions before running
+ImageMagick. This gives the user a predictable validation error and
+prevents arbitrary strings from becoming output file suffixes.
+
+The initial MIME type mapping should be:
+
+```text
+png -> image/png
+jpg -> image/jpeg
+webp -> image/webp
+avif -> image/avif
+```
+
+The mapping should live in Python code rather than relying only on
+`mimetypes.guess_type()`. A small explicit mapping is easier to test and
+keeps behavior stable across platforms.
+
+## Data Model Changes
+
+### GenerationDefaults
+
+Add optional config defaults:
+
+```python
+output_extension: str | None = None
+output_quality: int | None = None
+```
+
+These are optional because a config file may omit either key.
+
+### ImageGenerationRequest
+
+Add optional per-request values:
+
+```python
+output_extension: str | None = None
+output_quality: int | None = None
+```
+
+These values are request-level because the web UI and CLI can override
+config defaults for a single generation.
+
+### GenerationConfig
+
+Add validated final values:
+
+```python
+output_extension: str
+output_quality: int
+output_mime_type: str
+```
+
+The extension should already be normalized in `GenerationConfig`. Code
+outside the config layer should not need to understand aliases such as
+`jpeg`.
+
+The MIME type can either be stored directly or computed by a helper.
+Storing it on the config is acceptable because it is derived from a
+validated extension and is useful at response boundaries.
+
+### GeneratedImage
+
+`GeneratedImage` already has:
+
+```python
+data: bytes
+format: str
+seed: int
+```
+
+The `format` field should contain the normalized final extension. The UI
+route can derive MIME type from this value, or the service can include
+MIME type if the response model is expanded later.
+
+## Validation
+
+### Extension Validation
+
+Validation should:
+
+1. Require a string.
+2. Strip surrounding whitespace.
+3. Remove one leading dot if present.
+4. Convert to lowercase.
+5. Map aliases to canonical values.
+6. Reject unsupported values.
+
+Examples:
+
+```text
+"jpg" -> "jpg"
+".jpg" -> "jpg"
+"JPEG" -> "jpg"
+" webp" -> "webp"
+"gif" -> error
+"" -> error
+```
+
+The error should mention the accepted values:
+
+```text
+Output extension must be one of png, jpg, webp, avif: got gif
+```
+
+### Quality Validation
+
+Validation should:
+
+1. Require an integer.
+2. Reject booleans.
+3. Require the range `1..100`.
+
+Examples:
+
+```text
+85 -> valid
+1 -> valid
+100 -> valid
+0 -> error
+101 -> error
+true -> error
+"85" -> error in TOML and JSON parsing layers
+```
+
+The error should be:
+
+```text
+Output quality must be between 1 and 100: got 0
+```
+
+## Image I/O Design
+
+### Current Boundary
+
+The current `image_io.py` module already owns:
+
+- Batch output path calculation.
+- Tensor-to-uint8 conversion.
+- File saving.
+- Byte encoding for HTTP responses.
+
+This is the right module for final image format handling because it is
+already the boundary between tensors and encoded images.
+
+### Output Path Calculation
+
+Batch naming should keep the existing behavior of inserting the batch
+index before the suffix:
+
+```text
+output.jpg
+output_0000.jpg
+output_0001.jpg
+```
+
+The output suffix should be determined by the normalized output
+extension, not by the suffix on the configured output path.
+
+Recommended helper:
+
+```python
+def outputPathWithExtension(output: Path, extension: str) -> Path:
+ """Return output with the normalized image extension."""
+```
+
+For `output = Path("result.png")` and `extension = "webp"`, the helper
+returns:
+
+```text
+result.webp
+```
+
+For `output = Path("result")` and `extension = "jpg"`, the helper
+returns:
+
+```text
+result.jpg
+```
+
+### PNG Path
+
+For `png`, keep using Pillow directly:
+
+```python
+Image.fromarray(..., mode="RGB").save(path)
+```
+
+This preserves current behavior and keeps the default path independent
+from ImageMagick.
+
+### Non-PNG File Path
+
+For non-PNG file output:
+
+1. Convert the tensor image to an RGB uint8 array.
+2. Save that array to a temporary PNG file.
+3. Run ImageMagick:
+
+ ```text
+ magick TEMP.png -quality QUALITY OUTPUT.EXT
+ ```
+
+4. Return the final output path.
+5. Clean up the temporary PNG.
+
+The temporary PNG should be created in a temporary directory. This
+avoids leaving intermediate files next to user output and avoids
+collisions between concurrent conversions.
+
+Use `subprocess.run()` with a list of arguments:
+
+```python
+subprocess.run(
+ [
+ "magick",
+ str(temp_png),
+ "-quality",
+ str(quality),
+ str(output_path),
+ ],
+ check=True,
+ capture_output=True,
+ text=True,
+)
+```
+
+Do not use `shell=True`. A list of arguments avoids shell parsing and
+prevents spaces or special characters in paths from changing command
+meaning.
+
+### Non-PNG Byte Path
+
+For web UI responses:
+
+1. Convert the tensor image to an RGB uint8 array.
+2. Encode the array to temporary PNG bytes or a temporary PNG file.
+3. Run ImageMagick to write the requested final format to a temporary
+ output file.
+4. Read the final output bytes.
+5. Return those bytes in `GeneratedImage`.
+
+Using temporary files for both input and output is simpler than piping
+binary data into and out of ImageMagick. It also matches the CLI command
+shape and keeps error handling consistent.
+
+## ImageMagick Error Handling
+
+### Missing Executable
+
+If `magick` is missing, `subprocess.run()` raises `FileNotFoundError`.
+
+The application should catch it and raise `DiffusionCliError`:
+
+```text
+ImageMagick executable not found: magick
+```
+
+This tells the user the dependency is missing.
+
+### Failed Conversion
+
+If ImageMagick exits non-zero, `subprocess.run(..., check=True)` raises
+`subprocess.CalledProcessError`.
+
+The application should include the most useful available diagnostic:
+
+1. `stderr`, if present.
+2. `stdout`, if present.
+3. A generic message with the exit code.
+
+Example final error:
+
+```text
+ImageMagick failed to write avif output: no encode delegate for this
+image format
+```
+
+This is important for AVIF because ImageMagick support depends on the
+local build and installed delegates.
+
+### Empty Output
+
+After ImageMagick succeeds, the code should verify that the final output
+file exists and is a regular file.
+
+If the file is missing, raise:
+
+```text
+ImageMagick did not create output file: /path/to/output.avif
+```
+
+This check catches unusual cases where a command exits successfully but
+does not produce the expected file.
+
+## Service Design
+
+### GenerationService.generateImages
+
+`generateImages()` should return final-format bytes.
+
+Current behavior:
+
+```text
+tensor -> PNG bytes -> GeneratedImage(format="png")
+```
+
+New behavior:
+
+```text
+tensor -> configured output bytes -> GeneratedImage(format=extension)
+```
+
+The service already builds a full `GenerationConfig` before calling the
+model stack. That config should include `output_extension` and
+`output_quality`, so the service can pass validated values to
+`encodeImages()`.
+
+### GenerationService.generateToFiles
+
+`generateToFiles()` should write final-format files.
+
+Current behavior:
+
+```text
+tensor -> save PNG files -> list[Path]
+```
+
+New behavior:
+
+```text
+tensor -> save configured format files -> list[Path]
+```
+
+The returned paths must be the final paths. If the user requested
+`jpg`, the printed CLI path should end in `.jpg`.
+
+## API Behavior
+
+### First-Party UI API
+
+`POST /ui/api/txt2img` should accept:
+
+```json
+{
+ "output_extension": "jpg",
+ "output_quality": 85
+}
+```
+
+The route should validate JSON scalar types before constructing the
+internal request:
+
+- `output_extension` must be a string when present.
+- `output_quality` must be an integer when present.
+
+The config layer should perform final normalization and range checking.
+The UI route may also reject obviously invalid scalar types early so the
+user gets a `400 bad_request` instead of a later generation error.
+
+On ImageMagick failure, the response should be:
+
+```json
+{
+ "error": "generation_failed",
+ "message": "ImageMagick failed to write avif output: ..."
+}
+```
+
+Use status code `500` because inference may have succeeded, but the
+server failed to complete the requested operation.
+
+### Base64 API Responses
+
+All API responses that return base64 image data should base64 encode the
+final compressed image bytes.
+
+For example, if the selected output extension is `jpg`, the generation
+flow should produce JPEG bytes first, and then encode those JPEG bytes
+as base64. The base64 string should not imply PNG. It is only a text
+transport encoding for whatever image bytes were produced.
+
+The first-party UI response already includes `mime_type`, so it can tell
+the browser exactly what the base64 data contains:
+
+```json
+{
+ "images": [
+ {
+ "mime_type": "image/jpeg",
+ "data": "...base64...",
+ "seed": 123
+ }
+ ]
+}
+```
+
+Compatibility responses such as `/sdapi/v1/txt2img` may not have a MIME
+type field in their existing response shape. For now, they should still
+base64 encode the compressed bytes and assume the consumer can handle
+the image format. This keeps one consistent output-format path across
+CLI, UI, and compatibility APIs.
+
+### SillyTavern Compatibility API
+
+The compatibility API should keep returning base64 strings in the same
+shape it returns today.
+
+The first implementation may add parsing for output format fields only
+if a compatibility client sends them, but it does not need to invent a
+new compatibility-specific response shape. If the service returns JPEG,
+WebP, or AVIF bytes, the route should base64 encode those bytes directly
+and return them in the existing `images` array.
+
+This assumes the consumer can handle the compressed image format. If a
+specific client later proves it only accepts PNG, that client-specific
+behavior should be handled in that API profile deliberately.
+
+## Frontend Design
+
+### Form State
+
+Add fields to `INITIAL_FORM`:
+
+```javascript
+output_extension: "png",
+output_quality: 95,
+```
+
+When defaults load, merge server-provided values the same way as other
+generation defaults.
+
+### Controls
+
+Add a format select:
+
+```text
+Format: [PNG | JPEG | WebP | AVIF]
+```
+
+Add a quality input:
+
+```text
+Quality: [95]
+```
+
+The quality input should use:
+
+```text
+min = 1
+max = 100
+step = 1
+```
+
+The quality control can remain enabled for PNG for implementation
+simplicity, but the UI may disable it when PNG is selected because PNG
+does not use ImageMagick in the first implementation. If it is disabled
+for PNG, the request should still include the current value or omit it;
+the backend default will handle either case.
+
+### Gallery
+
+The current gallery uses:
+
+```javascript
+"data:" + image.mime_type + ";base64," + image.data
+```
+
+That behavior should remain. The only required change is that
+`mime_type` must come from the backend's final format.
+
+## Security and Robustness
+
+### Avoid Shell Invocation
+
+The ImageMagick command must be run with argument lists and
+`shell=False`, which is the default for `subprocess.run()`.
+
+This matters because output paths and config values are user-controlled.
+Even though the extension is validated, paths may contain spaces or
+characters that have meaning in a shell. Passing an argument list avoids
+shell interpretation.
+
+### Validate Before Command Construction
+
+Extension and quality should be validated before any temporary file or
+subprocess work starts.
+
+This keeps bad input cheap to reject and makes error messages
+deterministic.
+
+### Temporary Files
+
+Temporary files should live in a temporary directory created by
+`tempfile.TemporaryDirectory()`.
+
+The temporary directory should be scoped to one conversion or one batch
+save operation. When the context manager exits, intermediates are
+removed automatically.
+
+### Concurrency
+
+Server generation is already serialized by `GenerationService`'s lock.
+Temporary directories still matter because future code may relax the
+lock or add more server concurrency. Unique temporary directories avoid
+file name collisions.
+
+## Implementation Plan
+
+### Step 1: Add Config Constants and Validation
+
+Add constants:
+
+```python
+DEFAULT_OUTPUT_EXTENSION = "png"
+DEFAULT_OUTPUT_QUALITY = 95
+OUTPUT_EXTENSION_ALIASES = {
+ "png": "png",
+ "jpg": "jpg",
+ "jpeg": "jpg",
+ "webp": "webp",
+ "avif": "avif",
+}
+OUTPUT_MIME_TYPES = {
+ "png": "image/png",
+ "jpg": "image/jpeg",
+ "webp": "image/webp",
+ "avif": "image/avif",
+}
+```
+
+Add helpers:
+
+```python
+def normalizeOutputExtension(value: str) -> str:
+ """Return a supported canonical output extension."""
+
+def validateOutputQuality(value: int) -> int:
+ """Return a valid output quality value."""
+
+def outputMimeType(extension: str) -> str:
+ """Return the MIME type for a normalized output extension."""
+```
+
+These helpers should live in `config.py` unless the project later grows
+a dedicated output-format module.
+
+### Step 2: Extend TOML Loading
+
+Add `output_extension` and `output_quality` to
+`GENERATION_CONFIG_KEYS`.
+
+Parse them with existing scalar helpers:
+
+- `_optionalString()` for `output_extension`
+- `_optionalInt()` for `output_quality`
+
+Do not normalize TOML values during raw loading. Normalization belongs
+in the config build step where CLI, UI, and config values are merged.
+
+### Step 3: Extend Request and Config Dataclasses
+
+Add the fields described in the data model section.
+
+Update `buildGenerationConfigFromRequest()` to coalesce:
+
+```python
+raw_extension = coalesce(
+ request.output_extension,
+ generation.output_extension,
+ DEFAULT_OUTPUT_EXTENSION,
+)
+raw_quality = coalesce(
+ request.output_quality,
+ generation.output_quality,
+ DEFAULT_OUTPUT_QUALITY,
+)
+```
+
+Then validate:
+
+```python
+output_extension = normalizeOutputExtension(raw_extension)
+output_quality = validateOutputQuality(raw_quality)
+```
+
+### Step 4: Extend CLI Parser
+
+Add parser options with no argparse defaults:
+
+```python
+parser.add_argument("--output-extension")
+parser.add_argument("--output-quality", type=int)
+```
+
+Update `buildGenerationConfig()` to pass the parsed values into
+`ImageGenerationRequest`.
+
+### Step 5: Add Image I/O Conversion
+
+Update `saveImages()` and `encodeImages()` to accept:
+
+```python
+output_extension: str = "png"
+output_quality: int = 95
+```
+
+For PNG, keep existing behavior.
+
+For non-PNG, call a helper that writes a temporary PNG and runs
+ImageMagick.
+
+Recommended helper names:
+
+```python
+def convertPngFile(
+ png_path: Path,
+ output_path: Path,
+ output_extension: str,
+ output_quality: int,
+) -> None:
+ """Convert one PNG file to the requested output format."""
+
+def encodeArray(
+ image_array: np.ndarray,
+ output_extension: str,
+ output_quality: int,
+) -> bytes:
+ """Encode one RGB image array to final image bytes."""
+```
+
+### Step 6: Update Generation Service
+
+Pass config output fields to image I/O:
+
+```python
+encoded_images = encodeImages(
+ images,
+ config.output_extension,
+ config.output_quality,
+)
+```
+
+Return:
+
+```python
+GeneratedImage(
+ data=data,
+ format=config.output_extension,
+ seed=config.seed,
+)
+```
+
+Update file output similarly:
+
+```python
+saveImages(
+ images,
+ config.output,
+ config.output_extension,
+ config.output_quality,
+)
+```
+
+### Step 7: Update UI Routes
+
+Add JSON parsing for:
+
+```text
+output_extension
+output_quality
+```
+
+Include defaults and format options in `/ui/api/defaults`.
+
+Use the image format to set MIME type in `/ui/api/txt2img`:
+
+```python
+"mime_type": outputMimeType(image.format)
+```
+
+### Step 8: Update Frontend
+
+Add form state, controls, and request fields.
+
+The UI should not hard-code MIME behavior. It should continue to trust
+the backend response's `mime_type`.
+
+### Step 9: Update Documentation
+
+Update `README.md` with:
+
+- Config keys.
+- CLI examples.
+- Note that ImageMagick is required for non-PNG formats.
+- Note that AVIF support depends on the local ImageMagick build.
+
+## Testing Plan
+
+### Config Tests
+
+Add tests that verify:
+
+- Missing config uses `png` and `95`.
+- TOML `output_extension = "jpeg"` normalizes to `jpg`.
+- CLI `--output-extension webp` overrides config `jpg`.
+- Invalid extension fails clearly.
+- Invalid quality fails clearly.
+- Boolean quality is rejected.
+- UI defaults include output format and quality.
+
+### CLI Tests
+
+Add tests that verify:
+
+- Parser leaves new configurable defaults as `None`.
+- Parser accepts explicit `--output-extension`.
+- Parser accepts explicit `--output-quality`.
+
+### Image I/O Tests
+
+Add tests that verify:
+
+- `saveImages()` still writes PNG by default.
+- Non-PNG output changes the file suffix.
+- ImageMagick command uses argument lists in the expected order.
+- Missing `magick` raises `DiffusionCliError`.
+- Non-zero ImageMagick exit raises `DiffusionCliError` with stderr.
+- `encodeImages()` returns final-format bytes for non-PNG by reading
+ the converted output file.
+
+The ImageMagick tests should mock `subprocess.run()` rather than require
+ImageMagick in the test environment.
+
+### Generation Service Tests
+
+Add tests that verify:
+
+- `generateImages()` passes output extension and quality to
+ `encodeImages()`.
+- Returned `GeneratedImage.format` matches the configured extension.
+- `generateToFiles()` passes output extension and quality to
+ `saveImages()`.
+
+### UI Route Tests
+
+Add tests that verify:
+
+- Defaults endpoint includes format options.
+- UI request includes output format values in the internal request.
+- Invalid `output_extension` returns `400`.
+- Invalid `output_quality` returns `400`.
+- Response MIME type changes when the generated image format is `jpg`.
+- ImageMagick-related `DiffusionCliError` returns `500` with the error
+ message.
+
+### Frontend Tests
+
+The project currently has lightweight static frontend coverage through
+served asset tests. The first implementation can keep this level and
+verify that `app.js` contains the new request fields through existing
+route tests.
+
+If browser-level tests are added later, cover:
+
+- Changing the format select changes request JSON.
+- Changing quality changes request JSON.
+- A JPEG response displays in the gallery through a data URL.
+
+## Edge Cases
+
+### Output Path Has No Suffix
+
+If the output path is:
+
+```text
+result
+```
+
+and the extension is `jpg`, write:
+
+```text
+result.jpg
+```
+
+### Output Path Has Different Suffix
+
+If the output path is:
+
+```text
+result.png
+```
+
+and the extension is `avif`, write:
+
+```text
+result.avif
+```
+
+The configured path's suffix should not override the explicit format
+setting.
+
+### Batch Output
+
+If batch size is 2, output path is `result.png`, and extension is `webp`,
+write:
+
+```text
+result_0000.webp
+result_0001.webp
+```
+
+### AVIF Unsupported Locally
+
+If the local ImageMagick build lacks AVIF support, the user should see
+the stderr message from ImageMagick. The application should not attempt
+to detect delegate support ahead of time in the first implementation.
+Running the conversion is simpler and produces the most accurate local
+diagnostic.
+
+### PNG Quality
+
+The first implementation should ignore `output_quality` for PNG because
+PNG remains on the existing Pillow path.
+
+If PNG compression tuning becomes important later, add a separate
+setting with PNG-specific semantics. Do not overload JPEG-style quality
+without documenting the difference.
+
+## Open Questions
+
+### Should Compatibility APIs Force PNG?
+
+No for the first implementation.
+
+The compatibility endpoint currently returns base64 image strings
+without MIME metadata. The first implementation should base64 encode the
+compressed image bytes and assume the consumer can handle the actual
+format. This keeps behavior simple and avoids special-casing output
+format by route.
+
+If a specific consumer later proves it assumes PNG, the relevant API
+profile can force PNG or add a profile-specific compatibility option.
+
+### Should Quality Default Be 90 or 95?
+
+This design recommends `95` because it favors visual quality for image
+generation output. Users who want smaller files can lower it in config
+or the UI.
+
+`90` would produce smaller files by default for lossy formats, but it
+would also make the first opt-in conversion more visibly lossy. Since
+the feature exposes quality directly, the higher default is the more
+conservative choice.
+
+### Should ImageMagick Be Checked At Server Startup?
+
+No for the first implementation.
+
+The server can run normally for PNG output without ImageMagick. Checking
+at startup would create a warning or failure for users who never request
+non-PNG output. The conversion path should report the missing executable
+only when it is needed.
+
+## Acceptance Criteria
+
+The feature is complete when:
+
+- Existing PNG generation still works without ImageMagick installed.
+- CLI users can request `jpg`, `webp`, and `avif` output.
+- Config users can set default output format and quality.
+- Web UI users can select output format and quality.
+- UI image responses include correct MIME types.
+- ImageMagick missing or conversion failures are surfaced as clear
+ user-facing errors.
+- Tests cover validation, command construction, route parsing, and
+ service integration.
diff --git a/diffusion_cli/api_profiles.py b/diffusion_cli/api_profiles.py
index f25c3c8..14a155b 100644
--- a/diffusion_cli/api_profiles.py
+++ b/diffusion_cli/api_profiles.py
@@ -9,7 +9,12 @@ from typing import Callable
from flask import Response, jsonify, request
-from diffusion_cli.config import ImageGenerationRequest, validateDimensions
+from diffusion_cli.config import (
+ ImageGenerationRequest,
+ normalizeOutputExtension,
+ validateDimensions,
+ validateOutputQuality,
+)
from diffusion_cli.errors import DiffusionCliError
UNSUPPORTED_IMAGE_FIELDS = ("init_images", "mask", "extra_images")
@@ -78,6 +83,12 @@ def _txt2imgRequest(data: dict) -> ImageGenerationRequest:
steps = _optionalInt(data, "steps")
batch_size = _optionalInt(data, "batch_size")
cfg = _optionalFloat(data, "cfg_scale")
+ output_extension = _optionalString(data, "output_extension")
+ output_quality = _optionalInt(data, "output_quality")
+ if output_extension is not None:
+ output_extension = normalizeOutputExtension(output_extension)
+ if output_quality is not None:
+ output_quality = validateOutputQuality(output_quality)
if steps is not None and steps < 1:
raise DiffusionCliError(f"Steps must be at least 1: got {steps}")
if batch_size is not None and batch_size < 1:
@@ -96,6 +107,8 @@ def _txt2imgRequest(data: dict) -> ImageGenerationRequest:
batch_size=batch_size,
steps=steps,
cfg=cfg,
+ output_extension=output_extension,
+ output_quality=output_quality,
)
diff --git a/diffusion_cli/cli.py b/diffusion_cli/cli.py
index 69e11b3..a30f1eb 100644
--- a/diffusion_cli/cli.py
+++ b/diffusion_cli/cli.py
@@ -44,6 +44,8 @@ def buildParser() -> argparse.ArgumentParser:
parser.add_argument("--steps", type=int)
parser.add_argument("--cfg", type=float)
parser.add_argument("--output", type=Path)
+ parser.add_argument("--output-extension")
+ parser.add_argument("--output-quality", type=int)
parser.add_argument("--device")
parser.add_argument(
"--checkpoint",
@@ -190,7 +192,12 @@ def generate(args, user_config: UserConfig) -> list[Path]:
vae = ZImageVae(model_sources.vae, config.device, vae_dtype)
images = vae.decode(latent)
- return saveImages(images, config.output)
+ return saveImages(
+ images,
+ config.output,
+ config.output_extension,
+ config.output_quality,
+ )
def main(argv: list[str] | None = None) -> int:
diff --git a/diffusion_cli/config.py b/diffusion_cli/config.py
index c7a6f0f..f02b7d7 100644
--- a/diffusion_cli/config.py
+++ b/diffusion_cli/config.py
@@ -18,10 +18,31 @@ DEFAULT_BATCH_SIZE = 1
DEFAULT_STEPS = 10
DEFAULT_CFG = 1.0
DEFAULT_OUTPUT = Path("output.png")
+DEFAULT_OUTPUT_EXTENSION = "png"
+DEFAULT_OUTPUT_QUALITY = 95
DEFAULT_DEVICE = "cuda"
DEFAULT_DTYPE = "auto"
DEFAULT_SHIFT = 3.0
DEFAULT_MULTIPLIER = 1.0
+OUTPUT_EXTENSION_ALIASES = {
+ "png": "png",
+ "jpg": "jpg",
+ "jpeg": "jpg",
+ "webp": "webp",
+ "avif": "avif",
+}
+OUTPUT_MIME_TYPES = {
+ "png": "image/png",
+ "jpg": "image/jpeg",
+ "webp": "image/webp",
+ "avif": "image/avif",
+}
+OUTPUT_FORMAT_OPTIONS = (
+ {"value": "png", "label": "PNG"},
+ {"value": "jpg", "label": "JPEG"},
+ {"value": "webp", "label": "WebP"},
+ {"value": "avif", "label": "AVIF"},
+)
LATENT_CHANNELS = 16
LATENT_DOWNSCALE = 8
TOKENIZER_FILES = ("vocab.json", "merges.txt", "tokenizer_config.json")
@@ -77,6 +98,8 @@ GENERATION_CONFIG_KEYS = {
"steps",
"cfg",
"output",
+ "output_extension",
+ "output_quality",
"device",
"dtype",
}
@@ -137,6 +160,8 @@ class GenerationDefaults:
steps: int | None = None
cfg: float | None = None
output: Path | None = None
+ output_extension: str | None = None
+ output_quality: int | None = None
device: str | None = None
dtype: str | None = None
@@ -165,6 +190,9 @@ class GenerationConfig:
dtype: object
dtype_name: str
output: Path
+ output_extension: str
+ output_quality: int
+ output_mime_type: str
tokenizer_path: Path
@@ -181,6 +209,8 @@ class ImageGenerationRequest:
steps: int | None = None
cfg: float | None = None
output: Path | None = None
+ output_extension: str | None = None
+ output_quality: int | None = None
device: str | None = None
dtype: str | None = None
tokenizer_path: Path | None = None
@@ -325,6 +355,16 @@ def loadUserConfig(path: Path | None = None) -> UserConfig:
steps=_optionalInt(generation, "generation", "steps"),
cfg=_optionalFloat(generation, "generation", "cfg"),
output=_optionalPath(generation, "generation", "output"),
+ output_extension=_optionalString(
+ generation,
+ "generation",
+ "output_extension",
+ ),
+ output_quality=_optionalInt(
+ generation,
+ "generation",
+ "output_quality",
+ ),
device=_optionalString(generation, "generation", "device"),
dtype=_optionalString(generation, "generation", "dtype"),
),
@@ -384,6 +424,43 @@ def validateOutputPath(output: Path) -> Path:
return path
+def normalizeOutputExtension(value: str) -> str:
+ """Return a supported canonical output extension."""
+
+ if not isinstance(value, str):
+ raise DiffusionCliError("Output extension must be a string")
+
+ normalized = value.strip()
+ if normalized.startswith("."):
+ normalized = normalized[1:]
+ normalized = normalized.lower()
+ extension = OUTPUT_EXTENSION_ALIASES.get(normalized)
+ if extension is None:
+ accepted = ", ".join(OUTPUT_MIME_TYPES)
+ raise DiffusionCliError(
+ f"Output extension must be one of {accepted}: got {normalized}"
+ )
+ return extension
+
+
+def validateOutputQuality(value: int) -> int:
+ """Return a valid ImageMagick output quality value."""
+
+ if not isinstance(value, int) or isinstance(value, bool):
+ raise DiffusionCliError("Output quality must be an integer")
+ if value < 1 or value > 100:
+ raise DiffusionCliError(
+ f"Output quality must be between 1 and 100: got {value}"
+ )
+ return value
+
+
+def outputMimeType(extension: str) -> str:
+ """Return the MIME type for a normalized output extension."""
+
+ return OUTPUT_MIME_TYPES[normalizeOutputExtension(extension)]
+
+
def selectDevice(device_name: str):
"""Validate and return the requested CUDA device."""
@@ -464,6 +541,16 @@ def buildGenerationConfigFromRequest(
device_name = coalesce(request.device, generation.device, DEFAULT_DEVICE)
dtype_name = coalesce(request.dtype, generation.dtype, DEFAULT_DTYPE)
output_path = coalesce(request.output, generation.output, DEFAULT_OUTPUT)
+ raw_extension = coalesce(
+ request.output_extension,
+ generation.output_extension,
+ DEFAULT_OUTPUT_EXTENSION,
+ )
+ raw_quality = coalesce(
+ request.output_quality,
+ generation.output_quality,
+ DEFAULT_OUTPUT_QUALITY,
+ )
tokenizer_path = request.tokenizer_path
if tokenizer_path is None:
tokenizer_path = models.tokenizer
@@ -484,6 +571,9 @@ def buildGenerationConfigFromRequest(
device = selectDevice(device_name)
dtype = selectDtype(dtype_name, device)
output = validateOutputPath(output_path)
+ output_extension = normalizeOutputExtension(raw_extension)
+ output_quality = validateOutputQuality(raw_quality)
+ output_mime_type = outputMimeType(output_extension)
tokenizer_path = validateTokenizerPath(tokenizer_path)
seed = request.seed if request.seed is not None else randomSeed()
@@ -500,6 +590,9 @@ def buildGenerationConfigFromRequest(
dtype=dtype,
dtype_name=dtype_name,
output=output,
+ output_extension=output_extension,
+ output_quality=output_quality,
+ output_mime_type=output_mime_type,
tokenizer_path=tokenizer_path,
)
@@ -510,6 +603,13 @@ def buildDefaultGenerationRequest(
"""Build UI-visible generation defaults from config and fallbacks."""
generation = user_config.generation
+ output_extension = normalizeOutputExtension(
+ coalesce(generation.output_extension, DEFAULT_OUTPUT_EXTENSION)
+ )
+ output_quality = validateOutputQuality(
+ coalesce(generation.output_quality, DEFAULT_OUTPUT_QUALITY)
+ )
+
return ImageGenerationRequest(
prompt="",
negative_prompt=generation.negative_prompt or "",
@@ -519,6 +619,8 @@ def buildDefaultGenerationRequest(
batch_size=coalesce(generation.batch_size, DEFAULT_BATCH_SIZE),
steps=coalesce(generation.steps, DEFAULT_STEPS),
cfg=coalesce(generation.cfg, DEFAULT_CFG),
+ output_extension=output_extension,
+ output_quality=output_quality,
)
@@ -539,6 +641,8 @@ def buildGenerationConfig(
steps=args.steps,
cfg=args.cfg,
output=args.output,
+ output_extension=getattr(args, "output_extension", None),
+ output_quality=getattr(args, "output_quality", None),
device=args.device,
dtype=args.dtype,
tokenizer_path=args.tokenizer_path,
diff --git a/diffusion_cli/generation_service.py b/diffusion_cli/generation_service.py
index 8f0fefd..01d1b88 100644
--- a/diffusion_cli/generation_service.py
+++ b/diffusion_cli/generation_service.py
@@ -76,7 +76,7 @@ class GenerationService:
self,
request: ImageGenerationRequest,
) -> list[GeneratedImage]:
- """Generate PNG image bytes for one request."""
+ """Generate final-format image bytes for one request."""
with self._lock:
config = buildGenerationConfigFromRequest(
@@ -84,9 +84,17 @@ class GenerationService:
self.user_config,
)
images = self._generateTensor(config)
- encoded_images = encodeImages(images, "png")
+ encoded_images = encodeImages(
+ images,
+ config.output_extension,
+ config.output_quality,
+ )
return [
- GeneratedImage(data=data, format="png", seed=config.seed)
+ GeneratedImage(
+ data=data,
+ format=config.output_extension,
+ seed=config.seed,
+ )
for data in encoded_images
]
@@ -99,7 +107,12 @@ class GenerationService:
self.user_config,
)
images = self._generateTensor(config)
- return saveImages(images, config.output)
+ return saveImages(
+ images,
+ config.output,
+ config.output_extension,
+ config.output_quality,
+ )
def _generateTensor(self, config: GenerationConfig):
if self.model_residency == "cpu-cache":
diff --git a/diffusion_cli/image_io.py b/diffusion_cli/image_io.py
index 4b24d7f..23afb8e 100644
--- a/diffusion_cli/image_io.py
+++ b/diffusion_cli/image_io.py
@@ -4,15 +4,34 @@ from __future__ import annotations
from io import BytesIO
from pathlib import Path
+import subprocess
+import tempfile
import numpy as np
from PIL import Image
import torch
+from diffusion_cli.config import (
+ DEFAULT_OUTPUT_EXTENSION,
+ DEFAULT_OUTPUT_QUALITY,
+)
+from diffusion_cli.errors import DiffusionCliError
-def outputPaths(output: Path, batch_size: int) -> list[Path]:
+
+def outputPathWithExtension(output: Path, extension: str) -> Path:
+ """Return output with the normalized image extension."""
+
+ return output.with_suffix(f".{extension}")
+
+
+def outputPaths(
+ output: Path,
+ batch_size: int,
+ output_extension: str = DEFAULT_OUTPUT_EXTENSION,
+) -> list[Path]:
"""Return output file names for a batch."""
+ output = outputPathWithExtension(output, output_extension)
if batch_size == 1:
return [output]
return [
@@ -21,14 +40,32 @@ def outputPaths(output: Path, batch_size: int) -> list[Path]:
]
-def saveImages(images: torch.Tensor, output: Path) -> list[Path]:
- """Save an NCHW image tensor in [0, 1] as PNG files."""
+def saveImages(
+ images: torch.Tensor,
+ output: Path,
+ output_extension: str = DEFAULT_OUTPUT_EXTENSION,
+ output_quality: int = DEFAULT_OUTPUT_QUALITY,
+) -> list[Path]:
+ """Save an NCHW image tensor in [0, 1] as final image files."""
image_arrays = imageArrays(images)
- paths = outputPaths(output, len(image_arrays))
+ paths = outputPaths(output, len(image_arrays), output_extension)
for image_array, path in zip(image_arrays, paths, strict=True):
- Image.fromarray(np.asarray(image_array), mode="RGB").save(path)
+ if output_extension == "png":
+ Image.fromarray(np.asarray(image_array), mode="RGB").save(path)
+ else:
+ with tempfile.TemporaryDirectory() as temp_dir:
+ temp_png = Path(temp_dir) / "input.png"
+ Image.fromarray(np.asarray(image_array), mode="RGB").save(
+ temp_png,
+ )
+ convertPngFile(
+ temp_png,
+ path,
+ output_extension,
+ output_quality,
+ )
return paths
@@ -56,18 +93,82 @@ def imageArrays(images: torch.Tensor) -> np.ndarray:
def encodeImages(
images: torch.Tensor,
- image_format: str = "png",
+ output_extension: str = DEFAULT_OUTPUT_EXTENSION,
+ output_quality: int = DEFAULT_OUTPUT_QUALITY,
) -> list[bytes]:
- """Encode an NCHW image tensor in [0, 1] to image bytes."""
+ """Encode an NCHW image tensor in [0, 1] to final image bytes."""
image_arrays = imageArrays(images)
- encoded = []
- pil_format = image_format.upper()
- for image_array in image_arrays:
+ return [
+ encodeArray(image_array, output_extension, output_quality)
+ for image_array in image_arrays
+ ]
+
+
+def convertPngFile(
+ png_path: Path,
+ output_path: Path,
+ output_extension: str,
+ output_quality: int,
+) -> None:
+ """Convert one PNG file to the requested output format."""
+
+ try:
+ subprocess.run(
+ [
+ "magick",
+ str(png_path),
+ "-quality",
+ str(output_quality),
+ str(output_path),
+ ],
+ check=True,
+ capture_output=True,
+ text=True,
+ )
+ except FileNotFoundError as exc:
+ raise DiffusionCliError(
+ "ImageMagick executable not found: magick"
+ ) from exc
+ except subprocess.CalledProcessError as exc:
+ message = (exc.stderr or "").strip() or (exc.stdout or "").strip()
+ if not message:
+ message = f"exit code {exc.returncode}"
+ raise DiffusionCliError(
+ f"ImageMagick failed to write {output_extension} output: "
+ f"{message}"
+ ) from exc
+
+ if not output_path.is_file():
+ raise DiffusionCliError(
+ f"ImageMagick did not create output file: {output_path}"
+ )
+
+
+def encodeArray(
+ image_array: np.ndarray,
+ output_extension: str,
+ output_quality: int,
+) -> bytes:
+ """Encode one RGB image array to final image bytes."""
+
+ if output_extension == "png":
output = BytesIO()
Image.fromarray(np.asarray(image_array), mode="RGB").save(
output,
- format=pil_format,
+ format="PNG",
+ )
+ return output.getvalue()
+
+ with tempfile.TemporaryDirectory() as temp_dir:
+ temp_path = Path(temp_dir)
+ input_png = temp_path / "input.png"
+ output_path = temp_path / f"output.{output_extension}"
+ Image.fromarray(np.asarray(image_array), mode="RGB").save(input_png)
+ convertPngFile(
+ input_png,
+ output_path,
+ output_extension,
+ output_quality,
)
- encoded.append(output.getvalue())
- return encoded
+ return output_path.read_bytes()
diff --git a/diffusion_cli/static/ui/app.js b/diffusion_cli/static/ui/app.js
index 9430dc8..7517162 100644
--- a/diffusion_cli/static/ui/app.js
+++ b/diffusion_cli/static/ui/app.js
@@ -9,6 +9,8 @@ const INITIAL_FORM = {
steps: 10,
cfg: 1.0,
seed: -1,
+ output_extension: "png",
+ output_quality: 95,
};
function apiMessage(error, fallback) {
@@ -52,6 +54,7 @@ function NumberField(props) {
type: "number",
name: props.name,
min: props.min,
+ max: props.max,
step: props.step,
value: props.value,
disabled: props.disabled,
@@ -60,6 +63,28 @@ function NumberField(props) {
);
}
+function SelectField(props) {
+ return h(
+ "label",
+ { class: "field" },
+ h("span", { class: "field-label" }, props.label),
+ h(
+ "select",
+ {
+ name: props.name,
+ value: props.value,
+ disabled: props.disabled,
+ onInput: props.onInput,
+ },
+ props.options.map((option) => h(
+ "option",
+ { key: option.value, value: option.value },
+ option.label,
+ )),
+ ),
+ );
+}
+
function TextAreaField(props) {
return h(
"label",
@@ -205,6 +230,24 @@ function GenerationForm(props) {
disabled,
onInput: props.onInput,
})),
+ h(SelectField, {
+ name: "output_extension",
+ label: "Format",
+ value: form.output_extension,
+ options: props.output_formats,
+ disabled,
+ onInput: props.onInput,
+ }),
+ h(NumberField, {
+ name: "output_quality",
+ label: "Quality",
+ min: 1,
+ max: 100,
+ step: 1,
+ value: form.output_quality,
+ disabled,
+ onInput: props.onInput,
+ }),
h(
"div",
{ class: "seed-row" },
@@ -249,6 +292,12 @@ class App extends Component {
health: null,
form: { ...INITIAL_FORM },
resolution_presets: [],
+ output_formats: [
+ { value: "png", label: "PNG" },
+ { value: "jpg", label: "JPEG" },
+ { value: "webp", label: "WebP" },
+ { value: "avif", label: "AVIF" },
+ ],
is_loading_defaults: true,
is_generating: false,
error_message: "",
@@ -278,8 +327,11 @@ class App extends Component {
throw new Error("Could not load defaults");
}
const resolution_presets = defaults.resolution_presets || [];
+ const output_formats = defaults.output_formats
+ || this.state.output_formats;
const form_defaults = { ...defaults };
delete form_defaults.resolution_presets;
+ delete form_defaults.output_formats;
const form = { ...INITIAL_FORM, ...form_defaults, prompt: "" };
const is_valid_resolution = resolution_presets.some(
(preset) => (
@@ -297,6 +349,7 @@ class App extends Component {
health: health_response.ok ? health : null,
form,
resolution_presets,
+ output_formats,
is_loading_defaults: false,
});
}
@@ -361,6 +414,8 @@ class App extends Component {
steps: fieldValue(form.steps, null),
cfg: fieldValue(form.cfg, null),
seed: fieldValue(form.seed, -1),
+ output_extension: form.output_extension,
+ output_quality: fieldValue(form.output_quality, null),
};
}
@@ -439,6 +494,7 @@ class App extends Component {
h(GenerationForm, {
form: state.form,
resolution_presets: state.resolution_presets,
+ output_formats: state.output_formats,
disabled,
is_generating: state.is_generating,
onInput: this.handleInput,
diff --git a/diffusion_cli/ui.py b/diffusion_cli/ui.py
index a32a706..faff2ca 100644
--- a/diffusion_cli/ui.py
+++ b/diffusion_cli/ui.py
@@ -9,9 +9,13 @@ from flask import jsonify, redirect, request, send_from_directory
from diffusion_cli.config import (
ImageGenerationRequest,
+ OUTPUT_FORMAT_OPTIONS,
UI_RESOLUTION_PRESETS,
buildDefaultGenerationRequest,
+ normalizeOutputExtension,
+ outputMimeType,
validateDimensions,
+ validateOutputQuality,
)
from diffusion_cli.errors import DiffusionCliError
@@ -81,6 +85,12 @@ def _txt2imgRequest(data: dict) -> ImageGenerationRequest:
batch_size = _optionalInt(data, "batch_size")
steps = _optionalInt(data, "steps")
cfg = _optionalFloat(data, "cfg")
+ output_extension = _optionalString(data, "output_extension")
+ output_quality = _optionalInt(data, "output_quality")
+ if output_extension is not None:
+ output_extension = normalizeOutputExtension(output_extension)
+ if output_quality is not None:
+ output_quality = validateOutputQuality(output_quality)
if batch_size is not None and batch_size < 1:
raise DiffusionCliError(
f"Batch size must be at least 1: got {batch_size}"
@@ -99,6 +109,8 @@ def _txt2imgRequest(data: dict) -> ImageGenerationRequest:
batch_size=batch_size,
steps=steps,
cfg=cfg,
+ output_extension=output_extension,
+ output_quality=output_quality,
)
@@ -113,6 +125,9 @@ def _defaultResponse(context) -> dict:
"batch_size": default_request.batch_size,
"steps": default_request.steps,
"cfg": default_request.cfg,
+ "output_extension": default_request.output_extension,
+ "output_quality": default_request.output_quality,
+ "output_formats": list(OUTPUT_FORMAT_OPTIONS),
"seed": default_request.seed,
"resolution_presets": [
{
@@ -194,7 +209,7 @@ def registerUiRoutes(app, context) -> None:
return jsonify({
"images": [
{
- "mime_type": "image/png",
+ "mime_type": outputMimeType(image.format),
"data": b64encode(image.data).decode("ascii"),
"seed": image.seed,
}
diff --git a/tests/test_api_profiles.py b/tests/test_api_profiles.py
index cceb14f..44ceeaf 100644
--- a/tests/test_api_profiles.py
+++ b/tests/test_api_profiles.py
@@ -76,6 +76,8 @@ class ApiProfilesTest(unittest.TestCase):
"cfg_scale": 1.5,
"seed": 42,
"batch_size": 2,
+ "output_extension": "jpeg",
+ "output_quality": 80,
"sampler_name": "euler",
"scheduler": "normal",
"clip_skip": 1,
@@ -92,6 +94,8 @@ class ApiProfilesTest(unittest.TestCase):
self.assertEqual(request.cfg, 1.5)
self.assertEqual(request.seed, 42)
self.assertEqual(request.batch_size, 2)
+ self.assertEqual(request.output_extension, "jpg")
+ self.assertEqual(request.output_quality, 80)
def testTxt2ImgSeedMinusOneRequestsRandomSeed(self):
client, service = self.makeClient()
diff --git a/tests/test_cli.py b/tests/test_cli.py
index 7c2a22d..a06ab72 100644
--- a/tests/test_cli.py
+++ b/tests/test_cli.py
@@ -15,6 +15,8 @@ class CliTest(unittest.TestCase):
self.assertIsNone(args.steps)
self.assertIsNone(args.cfg)
self.assertIsNone(args.output)
+ self.assertIsNone(args.output_extension)
+ self.assertIsNone(args.output_quality)
self.assertIsNone(args.device)
self.assertIsNone(args.checkpoint)
self.assertIsNone(args.diffusion_model)
@@ -36,6 +38,10 @@ class CliTest(unittest.TestCase):
"fp16",
"--checkpoint",
"aio.safetensors",
+ "--output-extension",
+ ".jpeg",
+ "--output-quality",
+ "85",
]
)
@@ -43,6 +49,8 @@ class CliTest(unittest.TestCase):
self.assertEqual(args.cfg, 1.5)
self.assertEqual(args.dtype, "fp16")
self.assertEqual(args.checkpoint.name, "aio.safetensors")
+ self.assertEqual(args.output_extension, ".jpeg")
+ self.assertEqual(args.output_quality, 85)
def testModelRootIsNoLongerAccepted(self):
with patch("sys.stderr"):
diff --git a/tests/test_config.py b/tests/test_config.py
index 626819c..9ac5c77 100644
--- a/tests/test_config.py
+++ b/tests/test_config.py
@@ -6,6 +6,8 @@ from unittest.mock import patch
from diffusion_cli.config import (
DEFAULT_HEIGHT,
+ DEFAULT_OUTPUT_EXTENSION,
+ DEFAULT_OUTPUT_QUALITY,
DEFAULT_WIDTH,
GenerationDefaults,
ModelPathConfig,
@@ -14,7 +16,10 @@ from diffusion_cli.config import (
buildDefaultGenerationRequest,
buildGenerationConfig,
loadUserConfig,
+ normalizeOutputExtension,
+ outputMimeType,
validateDimensions,
+ validateOutputQuality,
validateOutputPath,
validateTokenizerPath,
)
@@ -47,6 +52,39 @@ class ConfigTest(unittest.TestCase):
self.assertEqual(result, output)
+ def testOutputExtensionValidationNormalizesAliases(self):
+ self.assertEqual(normalizeOutputExtension("jpg"), "jpg")
+ self.assertEqual(normalizeOutputExtension(".JPEG"), "jpg")
+ self.assertEqual(normalizeOutputExtension(" webp"), "webp")
+ self.assertEqual(outputMimeType("jpg"), "image/jpeg")
+
+ def testOutputExtensionValidationRejectsUnsupportedValues(self):
+ with self.assertRaises(DiffusionCliError) as context:
+ normalizeOutputExtension("gif")
+
+ self.assertIn(
+ "Output extension must be one of png, jpg, webp, avif: got gif",
+ str(context.exception),
+ )
+
+ def testOutputQualityValidationRejectsOutOfRangeAndBool(self):
+ self.assertEqual(validateOutputQuality(1), 1)
+ self.assertEqual(validateOutputQuality(100), 100)
+
+ with self.assertRaises(DiffusionCliError) as context:
+ validateOutputQuality(0)
+ self.assertIn(
+ "Output quality must be between 1 and 100",
+ str(context.exception),
+ )
+
+ with self.assertRaises(DiffusionCliError) as context:
+ validateOutputQuality(True)
+ self.assertIn(
+ "Output quality must be an integer",
+ str(context.exception),
+ )
+
def testMissingUserConfigReturnsEmptyConfig(self):
with tempfile.TemporaryDirectory() as temp_dir:
config = loadUserConfig(Path(temp_dir) / "missing.toml")
@@ -76,6 +114,8 @@ class ConfigTest(unittest.TestCase):
"steps = 12",
"cfg = 1.5",
f'output = "{root / "output.png"}"',
+ 'output_extension = "jpeg"',
+ "output_quality = 85",
'device = "cuda:0"',
'dtype = "bf16"',
)
@@ -97,6 +137,8 @@ class ConfigTest(unittest.TestCase):
self.assertEqual(config.generation.width, 512)
self.assertEqual(config.generation.height, 768)
self.assertEqual(config.generation.cfg, 1.5)
+ self.assertEqual(config.generation.output_extension, "jpeg")
+ self.assertEqual(config.generation.output_quality, 85)
self.assertEqual(config.generation.dtype, "bf16")
def testMalformedUserConfigFailsClearly(self):
@@ -205,6 +247,8 @@ class ConfigTest(unittest.TestCase):
steps=12,
cfg=1.5,
output=root / "image.png",
+ output_extension="jpeg",
+ output_quality=85,
device="cuda:0",
dtype="bf16",
),
@@ -237,6 +281,9 @@ class ConfigTest(unittest.TestCase):
self.assertEqual(config.steps, 12)
self.assertEqual(config.cfg, 1.5)
self.assertEqual(config.seed, 123)
+ self.assertEqual(config.output_extension, "jpg")
+ self.assertEqual(config.output_quality, 85)
+ self.assertEqual(config.output_mime_type, "image/jpeg")
def testGenerationConfigCliOverridesConfig(self):
with tempfile.TemporaryDirectory() as temp_dir:
@@ -248,7 +295,12 @@ class ConfigTest(unittest.TestCase):
user_config = UserConfig(
models=ModelPathConfig(tokenizer=tokenizer),
- generation=GenerationDefaults(width=512, steps=12),
+ generation=GenerationDefaults(
+ width=512,
+ steps=12,
+ output_extension="jpg",
+ output_quality=75,
+ ),
)
args = SimpleNamespace(
prompt="a mug",
@@ -260,6 +312,8 @@ class ConfigTest(unittest.TestCase):
steps=4,
cfg=None,
output=root / "image.png",
+ output_extension="webp",
+ output_quality=80,
device="cuda",
dtype="fp32",
tokenizer_path=None,
@@ -274,6 +328,44 @@ class ConfigTest(unittest.TestCase):
self.assertEqual(config.width, 640)
self.assertEqual(config.height, DEFAULT_HEIGHT)
self.assertEqual(config.steps, 4)
+ self.assertEqual(config.output_extension, "webp")
+ self.assertEqual(config.output_quality, 80)
+
+ def testGenerationConfigUsesBuiltInOutputDefaults(self):
+ with tempfile.TemporaryDirectory() as temp_dir:
+ root = Path(temp_dir)
+ tokenizer = root / "tokenizer"
+ tokenizer.mkdir()
+ for file_name in TOKENIZER_FILES:
+ (tokenizer / file_name).write_text("{}", encoding="utf-8")
+
+ user_config = UserConfig(
+ models=ModelPathConfig(tokenizer=tokenizer),
+ generation=GenerationDefaults(),
+ )
+ args = SimpleNamespace(
+ prompt="a mug",
+ negative_prompt=None,
+ seed=123,
+ width=None,
+ height=None,
+ batch_size=None,
+ steps=None,
+ cfg=None,
+ output=root / "image.png",
+ device="cuda",
+ dtype="fp32",
+ tokenizer_path=None,
+ )
+
+ with patch("diffusion_cli.config.selectDevice") as select_device:
+ with patch("diffusion_cli.config.selectDtype") as select_dtype:
+ select_device.return_value = "cuda"
+ select_dtype.return_value = "fp32"
+ config = buildGenerationConfig(args, user_config)
+
+ self.assertEqual(config.output_extension, DEFAULT_OUTPUT_EXTENSION)
+ self.assertEqual(config.output_quality, DEFAULT_OUTPUT_QUALITY)
def testDefaultGenerationRequestUsesUiVisibleFallbacks(self):
user_config = UserConfig(
@@ -289,6 +381,8 @@ class ConfigTest(unittest.TestCase):
self.assertEqual(request.height, DEFAULT_HEIGHT)
self.assertEqual(request.steps, 12)
self.assertEqual(request.seed, -1)
+ self.assertEqual(request.output_extension, "png")
+ self.assertEqual(request.output_quality, 95)
if __name__ == "__main__":
diff --git a/tests/test_generation_service.py b/tests/test_generation_service.py
index 9d3a414..24fdaff 100644
--- a/tests/test_generation_service.py
+++ b/tests/test_generation_service.py
@@ -55,6 +55,9 @@ class GenerationServiceTest(unittest.TestCase):
device="cuda",
dtype="dtype",
dtype_name="fp32",
+ output="output.png",
+ output_extension="webp",
+ output_quality=80,
tokenizer_path="tokenizer",
)
sources = SimpleNamespace(
@@ -95,12 +98,15 @@ class GenerationServiceTest(unittest.TestCase):
),
patch(
"diffusion_cli.generation_service.encodeImages",
- return_value=[b"png"],
- ),
+ return_value=[b"webp"],
+ ) as encode_images,
):
- service.generateImages(request)
+ first = service.generateImages(request)
service.generateImages(request)
+ self.assertEqual(first[0].format, "webp")
+ encode_images.assert_called_with("images", "webp", 80)
+
def testCpuCacheLoadsComponentsOnceAcrossRequests(self):
self._runTwoRequests("cpu-cache")
@@ -125,6 +131,31 @@ class GenerationServiceTest(unittest.TestCase):
self.assertEqual(FakeModel.load_count, 2)
self.assertEqual(FakeVae.load_count, 2)
+ def testGenerateToFilesPassesOutputFormatToImageIo(self):
+ config = SimpleNamespace(
+ output="output.png",
+ output_extension="jpg",
+ output_quality=85,
+ )
+ service = GenerationService(UserConfig(models=None, generation=None))
+
+ with (
+ patch(
+ "diffusion_cli.generation_service."
+ "buildGenerationConfigFromRequest",
+ return_value=config,
+ ),
+ patch.object(service, "_generateTensor", return_value="images"),
+ patch(
+ "diffusion_cli.generation_service.saveImages",
+ return_value=["output.jpg"],
+ ) as save_images,
+ ):
+ paths = service.generateToFiles(ImageGenerationRequest("a mug"))
+
+ self.assertEqual(paths, ["output.jpg"])
+ save_images.assert_called_with("images", "output.png", "jpg", 85)
+
if __name__ == "__main__":
unittest.main()
diff --git a/tests/test_image_io.py b/tests/test_image_io.py
index d5a3bab..1540bc6 100644
--- a/tests/test_image_io.py
+++ b/tests/test_image_io.py
@@ -1,19 +1,38 @@
import tempfile
import unittest
+from subprocess import CalledProcessError
from pathlib import Path
+from unittest.mock import patch
from PIL import Image
import torch
-from diffusion_cli.image_io import encodeImages, outputPaths, saveImages
+from diffusion_cli.errors import DiffusionCliError
+from diffusion_cli.image_io import (
+ convertPngFile,
+ encodeImages,
+ outputPathWithExtension,
+ outputPaths,
+ saveImages,
+)
class ImageIoTest(unittest.TestCase):
def testOutputPathsAppendBatchIndexBeforeSuffix(self):
- paths = outputPaths(Path("output.png"), 2)
+ paths = outputPaths(Path("output.png"), 2, "jpg")
- self.assertEqual(paths, [Path("output_0000.png"),
- Path("output_0001.png")])
+ self.assertEqual(paths, [Path("output_0000.jpg"),
+ Path("output_0001.jpg")])
+
+ def testOutputPathWithExtensionReplacesOrAddsSuffix(self):
+ self.assertEqual(
+ outputPathWithExtension(Path("result.png"), "webp"),
+ Path("result.webp"),
+ )
+ self.assertEqual(
+ outputPathWithExtension(Path("result"), "jpg"),
+ Path("result.jpg"),
+ )
def testSaveImagesWritesPngWithExpectedSize(self):
images = torch.zeros(1, 3, 4, 5)
@@ -29,6 +48,87 @@ class ImageIoTest(unittest.TestCase):
self.assertEqual(size, (5, 4))
self.assertEqual(mode, "RGB")
+ def testSaveImagesUsesFinalSuffixForNonPngOutput(self):
+ images = torch.zeros(1, 3, 4, 5)
+
+ with tempfile.TemporaryDirectory() as temp_dir:
+ output = Path(temp_dir) / "image.png"
+
+ def fake_run(args, **_kwargs):
+ Path(args[-1]).write_bytes(b"jpg")
+
+ with patch("diffusion_cli.image_io.subprocess.run", fake_run):
+ paths = saveImages(images, output, "jpg", 85)
+
+ self.assertEqual(paths, [Path(temp_dir) / "image.jpg"])
+
+ def testConvertPngFileUsesImageMagickArgumentList(self):
+ with tempfile.TemporaryDirectory() as temp_dir:
+ root = Path(temp_dir)
+ input_path = root / "input.png"
+ output_path = root / "output.webp"
+ input_path.write_bytes(b"png")
+
+ def fake_run(args, **kwargs):
+ self.assertEqual(
+ args,
+ [
+ "magick",
+ str(input_path),
+ "-quality",
+ "80",
+ str(output_path),
+ ],
+ )
+ self.assertTrue(kwargs["check"])
+ self.assertTrue(kwargs["capture_output"])
+ self.assertTrue(kwargs["text"])
+ output_path.write_bytes(b"webp")
+
+ with patch("diffusion_cli.image_io.subprocess.run", fake_run):
+ convertPngFile(input_path, output_path, "webp", 80)
+
+ def testConvertPngFileReportsMissingImageMagick(self):
+ with tempfile.TemporaryDirectory() as temp_dir:
+ input_path = Path(temp_dir) / "input.png"
+ output_path = Path(temp_dir) / "output.avif"
+ input_path.write_bytes(b"png")
+
+ with patch(
+ "diffusion_cli.image_io.subprocess.run",
+ side_effect=FileNotFoundError(),
+ ):
+ with self.assertRaises(DiffusionCliError) as context:
+ convertPngFile(input_path, output_path, "avif", 95)
+
+ self.assertIn(
+ "ImageMagick executable not found",
+ str(context.exception),
+ )
+
+ def testConvertPngFileReportsImageMagickFailureStderr(self):
+ with tempfile.TemporaryDirectory() as temp_dir:
+ input_path = Path(temp_dir) / "input.png"
+ output_path = Path(temp_dir) / "output.avif"
+ input_path.write_bytes(b"png")
+ error = CalledProcessError(
+ 1,
+ ["magick"],
+ stderr="no encode delegate\n",
+ )
+
+ with patch(
+ "diffusion_cli.image_io.subprocess.run",
+ side_effect=error,
+ ):
+ with self.assertRaises(DiffusionCliError) as context:
+ convertPngFile(input_path, output_path, "avif", 95)
+
+ self.assertIn(
+ "ImageMagick failed to write avif output: no encode delegate",
+ str(context.exception),
+ )
+
def testEncodeImagesReturnsPngBytes(self):
images = torch.zeros(1, 3, 4, 5)
@@ -37,6 +137,17 @@ class ImageIoTest(unittest.TestCase):
self.assertEqual(len(encoded), 1)
self.assertTrue(encoded[0].startswith(b"\x89PNG\r\n\x1a\n"))
+ def testEncodeImagesReturnsConvertedBytesForNonPng(self):
+ images = torch.zeros(1, 3, 4, 5)
+
+ def fake_run(args, **_kwargs):
+ Path(args[-1]).write_bytes(b"converted")
+
+ with patch("diffusion_cli.image_io.subprocess.run", fake_run):
+ encoded = encodeImages(images, "webp", 90)
+
+ self.assertEqual(encoded, [b"converted"])
+
if __name__ == "__main__":
unittest.main()
diff --git a/tests/test_ui.py b/tests/test_ui.py
index 09bdf6d..0e482ff 100644
--- a/tests/test_ui.py
+++ b/tests/test_ui.py
@@ -15,16 +15,22 @@ class FakeGenerationService:
def __init__(self):
self.user_config = UserConfig(
models=ModelPathConfig(),
- generation=GenerationDefaults(width=512, steps=12),
+ generation=GenerationDefaults(
+ width=512,
+ steps=12,
+ output_extension="jpeg",
+ output_quality=85,
+ ),
)
self.requests = []
self.error = None
+ self.image = GeneratedImage(b"\x89PNG\r\n\x1a\nimage", "png", 321)
def generateImages(self, request):
if self.error:
raise self.error
self.requests.append(request)
- return [GeneratedImage(b"\x89PNG\r\n\x1a\nimage", "png", 321)]
+ return [self.image]
class UiTest(unittest.TestCase):
@@ -58,6 +64,17 @@ class UiTest(unittest.TestCase):
self.assertEqual(response.json["batch_size"], 1)
self.assertEqual(response.json["steps"], 12)
self.assertEqual(response.json["cfg"], 1.0)
+ self.assertEqual(response.json["output_extension"], "jpg")
+ self.assertEqual(response.json["output_quality"], 85)
+ self.assertEqual(
+ response.json["output_formats"],
+ [
+ {"value": "png", "label": "PNG"},
+ {"value": "jpg", "label": "JPEG"},
+ {"value": "webp", "label": "WebP"},
+ {"value": "avif", "label": "AVIF"},
+ ],
+ )
self.assertEqual(response.json["seed"], -1)
presets = response.json["resolution_presets"]
self.assertEqual(len(presets), 33)
@@ -115,6 +132,8 @@ class UiTest(unittest.TestCase):
"cfg": 1.5,
"seed": -1,
"batch_size": 2,
+ "output_extension": "jpg",
+ "output_quality": 85,
},
)
@@ -128,6 +147,8 @@ class UiTest(unittest.TestCase):
self.assertEqual(request.cfg, 1.5)
self.assertIsNone(request.seed)
self.assertEqual(request.batch_size, 2)
+ self.assertEqual(request.output_extension, "jpg")
+ self.assertEqual(request.output_quality, 85)
image = response.json["images"][0]
self.assertEqual(image["mime_type"], "image/png")
self.assertEqual(
@@ -136,6 +157,17 @@ class UiTest(unittest.TestCase):
)
self.assertEqual(image["seed"], 321)
+ def testUiTxt2ImgUsesGeneratedImageMimeType(self):
+ client, service = self.makeClient()
+ service.image = GeneratedImage(b"jpeg", "jpg", 321)
+
+ response = client.post("/ui/api/txt2img", json={"prompt": "a mug"})
+
+ self.assertEqual(response.status_code, 200)
+ image = response.json["images"][0]
+ self.assertEqual(image["mime_type"], "image/jpeg")
+ self.assertEqual(base64.b64decode(image["data"]), b"jpeg")
+
def testUiTxt2ImgRejectsEmptyPrompt(self):
client, _service = self.makeClient()
@@ -156,6 +188,50 @@ class UiTest(unittest.TestCase):
self.assertEqual(response.status_code, 400)
self.assertIn("CFG must be non-negative", response.json["message"])
+ def testUiTxt2ImgRejectsInvalidOutputFieldTypes(self):
+ client, _service = self.makeClient()
+
+ response = client.post(
+ "/ui/api/txt2img",
+ json={"prompt": "a mug", "output_extension": 12},
+ )
+
+ self.assertEqual(response.status_code, 400)
+ self.assertEqual(
+ response.json["message"],
+ "output_extension must be a string",
+ )
+
+ response = client.post(
+ "/ui/api/txt2img",
+ json={"prompt": "a mug", "output_quality": True},
+ )
+
+ self.assertEqual(response.status_code, 400)
+ self.assertEqual(
+ response.json["message"],
+ "output_quality must be an integer",
+ )
+
+ def testUiTxt2ImgRejectsInvalidOutputValues(self):
+ client, _service = self.makeClient()
+
+ response = client.post(
+ "/ui/api/txt2img",
+ json={"prompt": "a mug", "output_extension": "gif"},
+ )
+
+ self.assertEqual(response.status_code, 400)
+ self.assertIn("Output extension must be", response.json["message"])
+
+ response = client.post(
+ "/ui/api/txt2img",
+ json={"prompt": "a mug", "output_quality": 0},
+ )
+
+ self.assertEqual(response.status_code, 400)
+ self.assertIn("Output quality must be", response.json["message"])
+
def testUiTxt2ImgRejectsZeroWidth(self):
client, _service = self.makeClient()