diff --git a/README.md b/README.md index cd7cd6a..ed790ee 100644 --- a/README.md +++ b/README.md @@ -254,7 +254,7 @@ mflux-generate \ - **`--lora-scales`** (optional, `[float]`, default: `None`): The scale for each respective [LoRA](#-LoRA) (will default to `1.0` if not specified and only one LoRA weight is loaded.) -- **`--metadata`** (optional): Exports a `.json` file containing the metadata for the image with the same name. (Even without this flag, the image metadata is saved and can be viewed using `exiftool image.png`) +- **`--metadata`** (optional): Exports a `.json` file containing the metadata for the image with the same name. (Even without this flag, the image metadata is saved and can be viewed using `mflux-info image.png` or `exiftool image.png`) - **`--image-path`** (optional, `str`, default: `None`): Local path to the initial image for image-to-image generation. @@ -643,7 +643,7 @@ system_profiler SPHardwareDataType SPDisplaysDataType *Note that these numbers includes starting the application from scratch, which means doing model i/o, setting/quantizing weights etc. -If we assume that the model is already loaded, you can inspect the image metadata using `exiftool image.png` and see the total duration of the denoising loop (excluding text embedding).* +If we assume that the model is already loaded, you can inspect the image metadata using `mflux-info image.png` (or `exiftool image.png`) and see the total duration of the denoising loop (excluding text embedding).* *These benchmarks are not very scientific and is only intended to give ballpark numbers. They were performed during different times with different MFLUX and MLX-versions etc. Additional hardware information such as number of GPU cores, Mac device etc. are not always known.* diff --git a/pyproject.toml b/pyproject.toml index 5d28898..987f132 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -91,6 +91,7 @@ mflux-save-depth = "mflux.save_depth:main" mflux-train = "mflux.train:main" mflux-upscale = "mflux.upscale:main" mflux-lora-library = "mflux.lora_library:main" +mflux-info = "mflux.info:main" mflux-completions = "mflux.ui.cli.completions.install:main" diff --git a/src/mflux/info.py b/src/mflux/info.py new file mode 100644 index 0000000..d88b37b --- /dev/null +++ b/src/mflux/info.py @@ -0,0 +1,124 @@ +"""Display metadata information from MFLUX generated images.""" + +import sys +from datetime import datetime +from pathlib import Path + +from mflux.post_processing.metadata_reader import MetadataReader +from mflux.ui.cli.parsers import CommandLineParser + + +def format_metadata(metadata: dict) -> str: + """Format metadata in a clean, readable format.""" + exif = metadata.get("exif", {}) + if not exif: + return "No metadata found" + + lines = [] + lines.append("=" * 60) + lines.append("MFLUX Image Information") + lines.append("=" * 60) + + # Prompt + if prompt := exif.get("prompt"): + lines.append(f"\nPrompt: {prompt}") + + if negative_prompt := exif.get("negative_prompt"): + lines.append(f"Negative Prompt: {negative_prompt}") + + # Model information + lines.append("") + if model := exif.get("model"): + lines.append(f"Model: {model}") + + # Image dimensions + if width := exif.get("width"): + lines.append(f"Width: {width}") + if height := exif.get("height"): + lines.append(f"Height: {height}") + + # Generation parameters + lines.append("") + if seed := exif.get("seed"): + lines.append(f"Seed: {seed}") + if steps := exif.get("steps"): + lines.append(f"Steps: {steps}") + if guidance := exif.get("guidance"): + lines.append(f"Guidance: {guidance}") + + # Technical settings + if quantize := exif.get("quantize"): + lines.append(f"Quantization: {quantize}-bit") + if precision := exif.get("precision"): + lines.append(f"Precision: {precision}") + + # LoRA information + if lora_paths := exif.get("lora_paths"): + lines.append("") + lines.append(f"LoRAs ({len(lora_paths)}):") + lora_scales = exif.get("lora_scales", []) + for i, lora in enumerate(lora_paths): + scale = lora_scales[i] if i < len(lora_scales) else 1.0 + lora_name = Path(lora).name + lines.append(f" - {lora_name} (scale: {scale})") + + # Image-to-image parameters + if image_path := exif.get("image_path"): + lines.append("") + lines.append(f"Source Image: {Path(image_path).name}") + if image_strength := exif.get("image_strength"): + lines.append(f"Image Strength: {image_strength}") + + # ControlNet parameters + if controlnet_path := exif.get("controlnet_image_path"): + lines.append("") + lines.append(f"ControlNet Image: {Path(controlnet_path).name}") + if controlnet_strength := exif.get("controlnet_strength"): + lines.append(f"ControlNet Strength: {controlnet_strength}") + + # Generation metadata + lines.append("") + if gen_time := exif.get("generation_time_seconds"): + lines.append(f"Generation Time: {gen_time:.2f}s") + + if created_at := exif.get("created_at"): + try: + dt = datetime.fromisoformat(created_at) + lines.append(f"Created: {dt.strftime('%Y-%m-%d %H:%M:%S')}") + except (ValueError, AttributeError): + lines.append(f"Created: {created_at}") + + if version := exif.get("mflux_version"): + lines.append(f"MFLUX Version: {version}") + + lines.append("=" * 60) + return "\n".join(lines) + + +def main(): + # Parse command line arguments + parser = CommandLineParser(description="Display metadata from MFLUX generated images") + parser.add_info_arguments() + args = parser.parse_args() + + # Check if file exists + image_path = Path(args.image_path) + if not image_path.exists(): + print(f"Error: Image file not found: {image_path}") + sys.exit(1) + + # Read metadata + metadata = MetadataReader.read_all_metadata(image_path) + + # Check if metadata was found + if not metadata or (not metadata.get("exif") and not metadata.get("xmp")): + print("No metadata found") + sys.exit(1) + + # Format and display + print(format_metadata(metadata)) + + +if __name__ == "__main__": + main() + diff --git a/src/mflux/post_processing/generated_image.py b/src/mflux/post_processing/generated_image.py index e682009..2861013 100644 --- a/src/mflux/post_processing/generated_image.py +++ b/src/mflux/post_processing/generated_image.py @@ -1,3 +1,4 @@ +from datetime import datetime from pathlib import Path import mlx.core as mx @@ -22,6 +23,8 @@ class GeneratedImage: generation_time: float, lora_paths: list[str], lora_scales: list[float], + height: int | None = None, + width: int | None = None, controlnet_image_path: str | Path | None = None, controlnet_strength: float | None = None, image_path: str | Path | None = None, @@ -44,6 +47,8 @@ class GeneratedImage: self.generation_time = generation_time self.lora_paths = lora_paths self.lora_scales = lora_scales + self.height = height + self.width = width self.controlnet_image_path = controlnet_image_path self.controlnet_strength = controlnet_strength self.image_path = image_path @@ -73,6 +78,8 @@ class GeneratedImage: generation_time=self.generation_time, lora_paths=self.lora_paths, lora_scales=self.lora_scales, + height=self.height, + width=self.width, controlnet_image_path=self.controlnet_image_path, controlnet_strength=self.controlnet_strength, image_path=self.image_path, @@ -139,9 +146,12 @@ class GeneratedImage: "seed": self.seed, "steps": self.steps, "guidance": self.guidance if self.model_config.supports_guidance else None, + "height": self.height, + "width": self.width, "precision": str(self.precision), "quantize": self.quantization, "generation_time_seconds": round(self.generation_time, 2), + "created_at": datetime.now().isoformat(), "lora_paths": [str(p) for p in self.lora_paths] if self.lora_paths else None, "lora_scales": [round(scale, 2) for scale in self.lora_scales] if self.lora_scales else None, "image_path": str(self.image_path) if self.image_path else None, diff --git a/src/mflux/post_processing/image_util.py b/src/mflux/post_processing/image_util.py index ce9f483..6204112 100644 --- a/src/mflux/post_processing/image_util.py +++ b/src/mflux/post_processing/image_util.py @@ -12,6 +12,7 @@ from PIL._typing import StrOrBytesPath from mflux.config.runtime_config import RuntimeConfig from mflux.models.flux.variants.concept_attention.attention_data import ConceptHeatmap from mflux.post_processing.generated_image import GeneratedImage +from mflux.post_processing.metadata_builder import MetadataBuilder from mflux.ui.box_values import AbsoluteBoxValues, BoxValues log = logging.getLogger(__name__) @@ -53,6 +54,8 @@ class ImageUtil: generation_time=generation_time, lora_paths=lora_paths, lora_scales=lora_scales, + height=config.height, + width=config.width, image_path=image_path, image_strength=image_strength, controlnet_image_path=controlnet_image_path, @@ -242,21 +245,24 @@ class ImageUtil: with open(f"{file_path.with_suffix('.json')}", "w") as json_file: json.dump(metadata, json_file, indent=4) - # Embed metadata + # Embed metadata in multiple formats for maximum compatibility if metadata is not None: ImageUtil._embed_metadata(metadata, file_path) + MetadataBuilder.embed_metadata(metadata, file_path) log.info(f"Metadata embedded successfully at: {file_path}") except Exception as e: # noqa: BLE001 log.error(f"Error saving image: {e}") @staticmethod def _embed_metadata(metadata: dict, path: str | Path) -> None: + """Original EXIF metadata embedding - preserved for compatibility""" try: # Convert metadata dictionary to a string metadata_str = json.dumps(metadata) # Convert the string to bytes (using UTF-8 encoding) - user_comment_bytes = metadata_str.encode("utf-8") + # Add the ASCII character code prefix required by EXIF spec + user_comment_bytes = b"ASCII\x00\x00\x00" + metadata_str.encode("utf-8") # Define the UserComment tag ID USER_COMMENT_TAG_ID = 0x9286 @@ -273,7 +279,7 @@ class ImageUtil: image.save(path, exif=exif_bytes) except Exception as e: # noqa: BLE001 - log.error(f"Error embedding metadata: {e}") + log.error(f"Error embedding EXIF metadata: {e}") @staticmethod def preprocess_for_model( diff --git a/src/mflux/post_processing/metadata_builder.py b/src/mflux/post_processing/metadata_builder.py new file mode 100644 index 0000000..a0dffbd --- /dev/null +++ b/src/mflux/post_processing/metadata_builder.py @@ -0,0 +1,231 @@ +"""Metadata builder for XMP and IPTC formats.""" + +import logging +from pathlib import Path + +import PIL.Image + +log = logging.getLogger(__name__) + + +class MetadataBuilder: + """Builds XMP and IPTC metadata packets for image embedding.""" + + @staticmethod + def embed_metadata(metadata: dict, path: str | Path) -> None: + """ + Embed XMP and IPTC metadata into an image file without touching existing EXIF. + Only supports PNG format. + + Args: + metadata: Dictionary containing image generation metadata + path: Path to the image file to embed metadata into + + Raises: + Exception: If there's an error during metadata embedding + """ + # Check if file is PNG format + path_obj = Path(path) if isinstance(path, str) else path + if path_obj.suffix.lower() != ".png": + log.warning(f"XMP/IPTC metadata embedding is only supported for PNG files, skipping: {path_obj}") + return + + try: + from PIL import PngImagePlugin + + # Load the image preserving existing metadata + image = PIL.Image.open(path) + + # Get existing PNG info to preserve it + existing_info = image.info if hasattr(image, "info") else {} + + # Preserve existing EXIF separately (if it exists) + existing_exif = existing_info.get("exif") + + # Create new PngInfo preserving existing data + pnginfo = PngImagePlugin.PngInfo() + + # Copy existing metadata + for key, value in existing_info.items(): + if key not in ["XML:com.adobe.xmp", "IPTC", "exif"]: # Handle these separately + pnginfo.add_text(key, str(value)) + + # Build XMP and IPTC metadata using builder methods + xmp_packet = MetadataBuilder.build_xmp_packet(metadata) + iptc_binary = MetadataBuilder.build_iptc_binary(metadata) + + # Add XMP and IPTC to PNG info + pnginfo.add_text("XML:com.adobe.xmp", xmp_packet) + if iptc_binary: + pnginfo.add_text("IPTC", iptc_binary.hex()) + + # Save preserving ALL existing metadata + adding XMP/IPTC + # Pass exif separately to preserve it correctly + if existing_exif: + image.save(path, pnginfo=pnginfo, exif=existing_exif) + else: + image.save(path, pnginfo=pnginfo) + + except Exception as e: # noqa: BLE001 + log.error(f"Error embedding XMP/IPTC metadata: {e}") + + @staticmethod + def build_xmp_packet(metadata: dict) -> str: + """ + Build an XMP metadata packet from the provided metadata dictionary. + + Args: + metadata: Dictionary containing image generation metadata + + Returns: + XMP packet as a string + """ + # Escape prompt for XML + prompt_escaped = metadata.get("prompt", "").replace("&", "&").replace("<", "<").replace(">", ">") + + # Build LoRA info for XMP + lora_info = MetadataBuilder._build_lora_string(metadata) + + # Get version from metadata + version = metadata.get("mflux_version", "unknown") + + xmp_packet = f""" + + + + {prompt_escaped} + MFLUX + AI Generated Content + MFLUX {version} + ART + Generated by MFLUX""" + + # Add technical parameters to XMP + if "seed" in metadata: + xmp_packet += f"\n {metadata['seed']}" + if "steps" in metadata: + xmp_packet += f"\n {metadata['steps']}" + if "guidance" in metadata: + xmp_packet += f"\n {metadata['guidance']}" + if "model_config" in metadata: + xmp_packet += f"\n {metadata['model_config']}" + if lora_info: + xmp_packet += f"\n {lora_info}" + if "generation_time" in metadata: + xmp_packet += f"\n {metadata['generation_time']}" + + xmp_packet += """ + + + +""" + + return xmp_packet + + @staticmethod + def build_iptc_binary(metadata: dict) -> bytes: + """ + Build IPTC metadata in binary format from the provided metadata dictionary. + + Args: + metadata: Dictionary containing image generation metadata + + Returns: + IPTC binary data + """ + # Build LoRA info for IPTC + lora_info = MetadataBuilder._build_lora_string(metadata) + + iptc_data = {} + + # Add prompt information + if "prompt" in metadata: + prompt = metadata["prompt"] + prompt_encoded = prompt.encode("utf-8") + if len(prompt_encoded) > 2000: + log.warning(f"Prompt is too long ({len(prompt_encoded)} bytes), truncating to 2000 bytes for IPTC") + iptc_data[120] = prompt_encoded[:2000] # Caption/Description + else: + iptc_data[120] = prompt_encoded # Caption/Description + iptc_data[5] = f"AI: {prompt[:50]}...".encode("utf-8") # Object Name/Title + iptc_data[105] = f"AI Generated: {prompt[:80]}...".encode("utf-8") # Headline + + # Add standard fields + iptc_data[80] = b"MFLUX" # By-line (Creator) + iptc_data[85] = b"AI Artist" # By-line Title + iptc_data[15] = b"ART" # Category + iptc_data[110] = b"Generated by MFLUX" # Credit + iptc_data[115] = b"AI Generation" # Source + iptc_data[116] = b"AI Generated Content" # Copyright Notice + iptc_data[118] = b"AI Generated using MFLUX" # Contact + iptc_data[103] = b"AI" # Instructions/Special Instructions + + # Add seed and model info in specific fields + if "seed" in metadata: + iptc_data[122] = f"Seed: {metadata['seed']}".encode("utf-8") # Writer/Editor + + if "model_config" in metadata: + iptc_data[90] = f"Model: {metadata['model_config']}".encode("utf-8") # City + + # Add LoRA info in Province/State field + if lora_info: + iptc_data[95] = f"LoRA: {lora_info}".encode("utf-8") # Province/State + + # Add generation parameters in Country field + if "steps" in metadata and "guidance" in metadata: + iptc_data[101] = f"Steps:{metadata['steps']} CFG:{metadata['guidance']}".encode("utf-8") # Country + + # Build keywords including LoRA info + keywords = ["AI", "Generated", "MFLUX"] + if "seed" in metadata: + keywords.append(f"seed-{metadata['seed']}") + if "steps" in metadata: + keywords.append(f"steps-{metadata['steps']}") + if "guidance" in metadata: + keywords.append(f"guidance-{metadata['guidance']}") + if "model_config" in metadata: + keywords.append(f"model-{metadata['model_config']}") + if lora_info: + keywords.append(f"loras-{lora_info}") + + iptc_data[25] = ";".join(keywords).encode("utf-8") # Keywords + + # Build IPTC binary (IPTC field length limit is 32767 bytes) + iptc_binary = b"" + for tag_id, value in iptc_data.items(): + length = len(value) + if length < 32768: + iptc_binary += bytes([0x1C, 0x02, tag_id]) + length.to_bytes(2, "big") + value + + return iptc_binary + + @staticmethod + def _build_lora_string(metadata: dict) -> str: + """ + Build a LoRA information string from metadata. + + Args: + metadata: Dictionary containing lora_paths and lora_scales + + Returns: + Comma-separated string of LoRA names and scales, or empty string + """ + if "lora_paths" not in metadata or not metadata["lora_paths"]: + return "" + + lora_list = [] + lora_paths = metadata["lora_paths"] + lora_scales = metadata.get("lora_scales", []) + + for i, lora_path in enumerate(lora_paths): + # Use Path for OS-agnostic path handling (works on Windows, Mac, Linux) + lora_name = Path(lora_path).name + scale = lora_scales[i] if i < len(lora_scales) else "1.0" + lora_list.append(f"{lora_name}:{scale}") + + return ", ".join(lora_list) + diff --git a/src/mflux/post_processing/metadata_reader.py b/src/mflux/post_processing/metadata_reader.py new file mode 100644 index 0000000..0fb53d4 --- /dev/null +++ b/src/mflux/post_processing/metadata_reader.py @@ -0,0 +1,126 @@ +"""Metadata reader for extracting and parsing image metadata.""" + +import json +import logging +from pathlib import Path + +import piexif +import PIL.Image + +log = logging.getLogger(__name__) + + +class MetadataReader: + """Reads and parses metadata from MFLUX generated images.""" + + @staticmethod + def read_exif_metadata(image_path: str | Path) -> dict | None: + """ + Extract EXIF metadata from an image. + + Args: + image_path: Path to the image file + + Returns: + Dictionary containing parsed EXIF metadata, or None if not found + """ + try: + img = PIL.Image.open(image_path) + exif_bytes = img.info.get("exif") + + if not exif_bytes: + return None + + exif_dict = piexif.load(exif_bytes) + user_comment = exif_dict["Exif"].get(0x9286, b"") + + if user_comment: + # Try to parse as JSON (strip the ASCII prefix if present) + if user_comment.startswith(b"ASCII\x00\x00\x00"): + metadata_str = user_comment[8:].decode("utf-8") + else: + metadata_str = user_comment.decode("utf-8") + return json.loads(metadata_str) + + return None + + except Exception as e: + log.debug(f"Error reading EXIF metadata: {e}") + return None + + @staticmethod + def read_xmp_metadata(image_path: str | Path) -> dict | None: + """ + Extract XMP metadata from an image. + + Args: + image_path: Path to the image file + + Returns: + Dictionary containing parsed XMP metadata, or None if not found + """ + try: + img = PIL.Image.open(image_path) + xmp_data = img.info.get("XML:com.adobe.xmp") + + if not xmp_data: + return None + + # Parse XMP XML to extract key fields + xmp_dict = {} + + # Simple XML parsing for common fields + fields = { + "description": "", + "creator": "", + "rights": "", + "creator_tool": "", + "category": "", + "credit": "", + "seed": "", + "steps": "", + "guidance": "", + "model": "", + "loras": "", + "generation_time": "", + } + + for key, start_tag in fields.items(): + if start_tag in xmp_data: + start_idx = xmp_data.index(start_tag) + len(start_tag) + # Find the closing tag + if key in ["description", "rights"]: + end_tag = "" + elif key == "creator": + end_tag = "" + else: + # Extract tag name from start_tag + tag_name = start_tag.split(":")[1].rstrip(">") + end_tag = f"" + + end_idx = xmp_data.index(end_tag, start_idx) + value = xmp_data[start_idx:end_idx] + xmp_dict[key] = value + + return xmp_dict if xmp_dict else None + + except Exception as e: + log.debug(f"Error reading XMP metadata: {e}") + return None + + @staticmethod + def read_all_metadata(image_path: str | Path) -> dict: + """ + Read all available metadata from an image. + + Args: + image_path: Path to the image file + + Returns: + Dictionary with 'exif' and 'xmp' keys containing metadata + """ + return { + "exif": MetadataReader.read_exif_metadata(image_path), + "xmp": MetadataReader.read_xmp_metadata(image_path), + } + diff --git a/src/mflux/ui/cli/parsers.py b/src/mflux/ui/cli/parsers.py index e0c9e30..4b39e4b 100644 --- a/src/mflux/ui/cli/parsers.py +++ b/src/mflux/ui/cli/parsers.py @@ -192,6 +192,9 @@ class CommandLineParser(argparse.ArgumentParser): self.add_argument("--train-config", type=str, required=False, help="Local path of the training configuration file") self.add_argument("--train-checkpoint", type=str, required=False, help="Local path of the checkpoint file which specifies how to continue the training process") + def add_info_arguments(self) -> None: + self.add_argument("image_path", type=str, help="Path to the image file to inspect") + def parse_args(self) -> argparse.Namespace: # type: ignore namespace = super().parse_args() diff --git a/tests/metadata/__init__.py b/tests/metadata/__init__.py new file mode 100644 index 0000000..5a367b0 --- /dev/null +++ b/tests/metadata/__init__.py @@ -0,0 +1,2 @@ +# Metadata tests + diff --git a/tests/metadata/test_metadata.py b/tests/metadata/test_metadata.py new file mode 100644 index 0000000..59c2707 --- /dev/null +++ b/tests/metadata/test_metadata.py @@ -0,0 +1,184 @@ +""" +Test metadata embedding and reading functionality. + +This test generates a single small image using the schnell model and verifies +that all metadata fields are correctly embedded and can be read back. +""" + +import json +import os +import tempfile +from datetime import datetime +from pathlib import Path + +from mflux.config.config import Config +from mflux.config.model_config import ModelConfig +from mflux.models.flux.variants.txt2img.flux import Flux1 +from mflux.post_processing.metadata_reader import MetadataReader + + +class TestMetadata: + """Test suite for image metadata functionality.""" + + def test_metadata_complete(self): + """ + Comprehensive test that generates one image using img2img and verifies: + - EXIF metadata (all fields including img2img-specific ones) + - mflux-info command output formatting + - Creation timestamp validity + - Optional fields handling (both None and populated) + - Metadata reader handles missing files gracefully + """ + # Use a temporary file - don't keep it open to avoid conflicts + fd, temp_path = tempfile.mkstemp(suffix=".png") + os.close(fd) # Close the file descriptor immediately + output_path = Path(temp_path) + + try: + # Use img2img to test image_path and image_strength metadata fields + reference_image = Path(__file__).parent.parent / "resources" / "reference_schnell.png" + + # Generate a small, fast image with schnell img2img (256x256, 2 steps, quantized) + flux = Flux1( + model_config=ModelConfig.schnell(), + quantize=8, + ) + + config = Config( + num_inference_steps=2, + height=256, + width=256, + image_path=reference_image, + image_strength=0.3, + ) + + image = flux.generate_image( + seed=42, + prompt="A simple test image", + config=config, + ) + + # Save with metadata (overwrite=True since mkstemp creates an empty file) + image.save(path=output_path, overwrite=True) + + # ================================================================= + # Test 1: Read metadata and verify structure + # ================================================================= + metadata = MetadataReader.read_all_metadata(output_path) + + assert metadata is not None, "Metadata should not be None" + assert "exif" in metadata, "EXIF metadata should be present" + + exif = metadata["exif"] + + # ================================================================= + # Test 2: Core generation parameters + # ================================================================= + assert exif.get("seed") == 42, "Seed should match" + assert exif.get("steps") == 2, "Steps should match" + assert exif.get("prompt") == "A simple test image", "Prompt should match" + assert exif.get("model") == "black-forest-labs/FLUX.1-schnell", "Model should match" + + # ================================================================= + # Test 3: Dimensions (NEW FEATURE) + # ================================================================= + assert exif.get("width") == 256, "Width should be saved" + assert exif.get("height") == 256, "Height should be saved" + + # ================================================================= + # Test 4: Technical parameters + # ================================================================= + assert exif.get("quantize") == 8, "Quantization should match" + assert exif.get("precision") is not None, "Precision should be set" + + # ================================================================= + # Test 5: MFLUX version (not hardcoded) + # ================================================================= + assert exif.get("mflux_version") is not None, "MFLUX version should be present" + assert exif.get("mflux_version") != "unknown", "MFLUX version should not be unknown" + + # ================================================================= + # Test 6: Generation time + # ================================================================= + assert exif.get("generation_time_seconds") is not None, "Generation time should be present" + assert isinstance(exif.get("generation_time_seconds"), (int, float)), "Generation time should be numeric" + assert exif.get("generation_time_seconds") > 0, "Generation time should be positive" + + # ================================================================= + # Test 7: Creation timestamp (NEW FEATURE) + # ================================================================= + assert exif.get("created_at") is not None, "Creation timestamp should be present" + created_at = exif.get("created_at") + + # Verify it's valid ISO format + dt = datetime.fromisoformat(created_at) + now = datetime.now() + time_diff = abs((now - dt).total_seconds()) + assert time_diff < 120, f"Timestamp should be recent (within 120s), but was {time_diff}s ago" + + # ================================================================= + # Test 8: Img2img-specific fields are populated + # ================================================================= + assert exif.get("image_path") is not None, "Image path should be set for img2img" + assert str(reference_image) in exif.get("image_path"), "Image path should contain reference image name" + assert exif.get("image_strength") == 0.3, "Image strength should match config" + + # ================================================================= + # Test 9: Optional fields are None when not used + # ================================================================= + assert exif.get("lora_paths") is None, "LoRA paths should be None when not used" + assert exif.get("lora_scales") is None, "LoRA scales should be None when not used" + assert exif.get("controlnet_image_path") is None, "ControlNet path should be None when not used" + assert exif.get("negative_prompt") is None, "Negative prompt should be None when not used" + assert exif.get("guidance") is None, "Guidance should be None for schnell model" + + # ================================================================= + # Test 10: EXIF JSON is valid and parseable + # ================================================================= + import piexif + from PIL import Image + + img = Image.open(output_path) + exif_bytes = img.info.get("exif") + assert exif_bytes is not None, "EXIF bytes should be present" + + exif_dict = piexif.load(exif_bytes) + user_comment = exif_dict.get("Exif", {}).get(piexif.ExifIFD.UserComment) + assert user_comment is not None, "UserComment should be present" + + # Decode JSON + if user_comment.startswith(b"ASCII\x00\x00\x00"): + json_str = user_comment[8:].decode("utf-8") + else: + json_str = user_comment.decode("utf-8") + + metadata_parsed = json.loads(json_str) + # Just verify it's valid JSON + assert metadata_parsed.get("prompt") == "A simple test image", "EXIF JSON should contain correct prompt" + + # ================================================================= + # Test 11: mflux-info command output + # ================================================================= + from mflux.info import format_metadata + + output = format_metadata(metadata) + assert "A simple test image" in output, "Prompt should be in output" + assert "MFLUX" in output, "MFLUX should be mentioned" + assert "42" in output, "Seed should be in output" + assert "256" in output, "Dimensions should be in output" + assert "Generation Time:" in output, "Generation time should be shown" + assert "Created:" in output, "Creation timestamp should be shown" + assert "Source Image:" in output, "Source image should be shown for img2img" + assert "Image Strength:" in output, "Image strength should be shown for img2img" + + # ================================================================= + # Test 12: Metadata reader handles nonexistent files + # ================================================================= + nonexistent_metadata = MetadataReader.read_all_metadata(Path("/nonexistent/file.png")) + assert nonexistent_metadata.get("exif") is None, "Nonexistent file should return None for EXIF" + assert nonexistent_metadata.get("xmp") is None, "Nonexistent file should return None for XMP" + + finally: + # Cleanup: Always remove the temporary file + if output_path.exists(): + output_path.unlink()