Fix validation artifact asset formats

Author: bghira

simpletuner/helpers/publishing/huggingface.py

@@ -8,6 +8,7 @@
 
 from simpletuner.helpers.publishing.metadata import save_model_card, save_training_config
 from simpletuner.helpers.training.state_tracker import StateTracker
+from simpletuner.helpers.training.validation_images import save_image_with_validation_format
 
 logger = logging.getLogger(__name__)
 from simpletuner.helpers.training.multi_process import should_log
@@ -21,6 +22,10 @@
 LORA_SAFETENSORS_FILENAME = "pytorch_lora_weights.safetensors"
 EMA_SAFETENSORS_FILENAME = "ema_model.safetensors"
 SLA_ATTENTION_FILENAME = "sla_attention.pt"
+VALIDATION_IMAGE_SUFFIXES = {".png", ".jpg", ".jpeg", ".webp", ".gif"}
+VALIDATION_VIDEO_SUFFIXES = {".mp4", ".avi", ".mov", ".webm"}
+VALIDATION_AUDIO_SUFFIXES = {".wav", ".flac", ".mp3", ".ogg", ".m4a"}
+VALIDATION_MEDIA_SUFFIXES = VALIDATION_IMAGE_SUFFIXES | VALIDATION_VIDEO_SUFFIXES | VALIDATION_AUDIO_SUFFIXES
 
 
 class HubManager:
@@ -409,6 +414,59 @@ def find_latest_checkpoint(self):
 
         return highest_checkpoint
 
+    @staticmethod
+    def _validation_media_shortname(filename: str, checkpoint_step: int) -> str | None:
+        stem = Path(filename).stem
+        prefix = f"step_{checkpoint_step}_"
+        if not stem.startswith(prefix):
+            return None
+
+        remainder = stem[len(prefix) :]
+        image_or_video_parts = remainder.rsplit("_", 2)
+        if len(image_or_video_parts) == 3 and image_or_video_parts[1].isdigit():
+            return image_or_video_parts[0]
+
+        audio_parts = remainder.rsplit("_", 1)
+        if len(audio_parts) == 2 and audio_parts[1].isdigit():
+            return audio_parts[0]
+
+        return None
+
+    def _filter_checkpoint_validation_media(
+        self,
+        checkpoint_step: int,
+        validation_images: dict | None,
+        validation_audios: dict | None,
+    ) -> tuple[dict, dict]:
+        filtered_images = {}
+        filtered_audios = {}
+        validation_dir = Path(self.config.output_dir) / "validation_images"
+        if not validation_dir.exists():
+            return filtered_images, filtered_audios
+
+        image_shortnames = set(validation_images.keys()) if validation_images else set()
+        audio_shortnames = set(validation_audios.keys()) if validation_audios else set()
+        shortnames = image_shortnames | audio_shortnames
+        if not shortnames:
+            return filtered_images, filtered_audios
+
+        for media_path in sorted(validation_dir.iterdir(), key=lambda path: path.name):
+            if not media_path.is_file() or media_path.suffix.lower() not in VALIDATION_MEDIA_SUFFIXES:
+                continue
+            shortname = self._validation_media_shortname(media_path.name, checkpoint_step)
+            if shortname not in shortnames:
+                continue
+
+            suffix = media_path.suffix.lower()
+            if suffix in VALIDATION_AUDIO_SUFFIXES:
+                if shortname in audio_shortnames:
+                    filtered_audios.setdefault(shortname, []).append(str(media_path))
+                continue
+            if shortname in image_shortnames:
+                filtered_images.setdefault(shortname, []).append(str(media_path))
+
+        return filtered_images, filtered_audios
+
     def upload_latest_checkpoint(
         self,
         validation_images: dict,
@@ -435,30 +493,11 @@ def upload_latest_checkpoint(
                 filtered_images = {}
                 filtered_audios = {}
                 if (validation_images or validation_audios) and checkpoint_step is not None:
-                    validation_dir = os.path.join(self.config.output_dir, "validation_images")
-                    if os.path.exists(validation_dir):
-                        shortnames = set()
-                        if validation_images:
-                            shortnames.update(validation_images.keys())
-                        if validation_audios:
-                            shortnames.update(validation_audios.keys())
-                        for shortname in shortnames:
-                            step_pattern = f"step_{checkpoint_step}_"
-                            for img_file in os.listdir(validation_dir):
-                                if step_pattern in img_file and shortname in img_file:
-                                    img_path = os.path.join(validation_dir, img_file)
-                                    try:
-                                        if img_path.endswith((".mp4", ".avi", ".mov", ".webm")):
-                                            filtered_images.setdefault(shortname, []).append(img_path)
-                                        elif img_path.endswith((".wav", ".flac", ".mp3", ".ogg", ".m4a")):
-                                            filtered_audios.setdefault(shortname, []).append(img_path)
-                                        else:
-                                            from PIL import Image
-
-                                            img = Image.open(img_path)
-                                            filtered_images.setdefault(shortname, []).append(img)
-                                    except Exception as e:
-                                        logger.warning(f"Could not load validation asset {img_path}: {e}")
+                    filtered_images, filtered_audios = self._filter_checkpoint_validation_media(
+                        checkpoint_step,
+                        validation_images,
+                        validation_audios,
+                    )
 
                 # Only use media generated at this checkpoint step. Previous-step benchmark samples
                 # should not be associated with the checkpoint model card.
@@ -503,12 +542,13 @@ def upload_validation_images(self, validation_images, webhook_handler=None, over
                     images = [images]
                 sub_idx = 0
                 for image in images:
-                    image_path = os.path.join(
-                        override_path or self.config.output_dir,
-                        "assets",
-                        f"image_{idx}_{sub_idx}.png",
+                    assets_dir = os.path.join(override_path or self.config.output_dir, "assets")
+                    os.makedirs(assets_dir, exist_ok=True)
+                    image_path, image_extension = save_image_with_validation_format(
+                        image,
+                        os.path.join(assets_dir, f"image_{idx}_{sub_idx}"),
+                        self.config,
                     )
-                    image.save(image_path, format="PNG")
                     if not self.config.push_to_hub:
                         continue
                     attempt = 0
@@ -517,7 +557,7 @@ def upload_validation_images(self, validation_images, webhook_handler=None, over
                         try:
                             self._hub_api.upload_file(
                                 repo_id=self._repo_id,
-                                path_in_repo=f"/assets/image_{idx}_{sub_idx}.png",
+                                path_in_repo=f"/assets/image_{idx}_{sub_idx}.{image_extension}",
                                 path_or_fileobj=image_path,
                                 commit_message="Validation image auto-generated by SimpleTuner",
                                 token=self.hub_token,

simpletuner/helpers/publishing/metadata.py

@@ -13,6 +13,7 @@
 from simpletuner.helpers.data_backend.dataset_types import DatasetType, ensure_dataset_type
 from simpletuner.helpers.models.common import AudioModelFoundation, ModelFoundation
 from simpletuner.helpers.training.state_tracker import StateTracker
+from simpletuner.helpers.training.validation_images import save_image_with_validation_format
 
 logger = logging.getLogger(__name__)
 from simpletuner.helpers.training.multi_process import should_log
@@ -488,6 +489,7 @@ def save_metadata_sample(
     image_path: str,
     image: Union[Image.Image, np.ndarray, list, str, torch.Tensor],
     sample_rate: int = 44100,
+    config: Any = None,
 ):
     if isinstance(image, str):
         import shutil
@@ -519,9 +521,7 @@ def save_metadata_sample(
             fps=StateTracker.get_args().framerate,
         )
     elif isinstance(image, Image.Image):
-        file_extension = "png"
-        output_path = f"{image_path}.{file_extension}"
-        image.save(output_path, format="PNG")
+        output_path, file_extension = save_image_with_validation_format(image, image_path, config)
     else:
         raise ValueError(f"Cannot export sample type {type(image)} yet.")
 
@@ -772,6 +772,7 @@ def _add_widget_entries(media, asset_prefix: str):
                     image_path=os.path.join(assets_folder, f"{asset_prefix}_{idx}_{sub_idx}"),
                     image=media_sample,
                     sample_rate=audio_sample_rate,
+                    config=args,
                 )
                 asset_filename = os.path.basename(output_path)
                 if media_extension in {"mp4", "avi", "mov", "webm"}:

simpletuner/helpers/training/validation_images.py

@@ -2,40 +2,54 @@
 import os
 
 import numpy as np
-import wandb
 from PIL import Image
 
+import wandb
 from simpletuner.helpers.training.local_metrics import record_validation_media
 from simpletuner.helpers.training.state_tracker import StateTracker
 
 logger = logging.getLogger(__name__)
 
+VALIDATION_IMAGE_FORMATS = {"png", "webp", "jpeg"}
 
-def save_validation_image(
-    image,
-    save_dir,
-    filename_stem,
-    config,
-    *,
-    label,
-    index,
-    resolution,
-):
+
+def validation_image_file_extension(config) -> str:
     image_format = str(getattr(config, "validation_image_format", "png") or "png").lower()
-    if image_format not in {"png", "webp", "jpeg"}:
+    if image_format not in VALIDATION_IMAGE_FORMATS:
+        raise ValueError("validation_image_format must be one of: png, webp, jpeg.")
+    return "jpg" if image_format == "jpeg" else image_format
+
+
+def save_image_with_validation_format(image, image_path_stem, config):
+    image_format = str(getattr(config, "validation_image_format", "png") or "png").lower()
+    if image_format not in VALIDATION_IMAGE_FORMATS:
         raise ValueError("validation_image_format must be one of: png, webp, jpeg.")
     configured_quality = getattr(config, "validation_image_quality", 90)
     quality = 90 if configured_quality is None else int(configured_quality)
     if not 1 <= quality <= 100:
         raise ValueError("validation_image_quality must be within [1, 100].")
 
-    extension = "jpg" if image_format == "jpeg" else image_format
-    save_path = os.path.join(save_dir, f"{filename_stem}.{extension}")
+    extension = validation_image_file_extension(config)
+    save_path = f"{image_path_stem}.{extension}"
     save_image = image.convert("RGB") if image_format == "jpeg" and getattr(image, "mode", None) != "RGB" else image
     save_kwargs = {"format": image_format.upper()}
     if image_format in {"webp", "jpeg"}:
         save_kwargs["quality"] = quality
     save_image.save(save_path, **save_kwargs)
+    return save_path, extension
+
+
+def save_validation_image(
+    image,
+    save_dir,
+    filename_stem,
+    config,
+    *,
+    label,
+    index,
+    resolution,
+):
+    save_path, _extension = save_image_with_validation_format(image, os.path.join(save_dir, filename_stem), config)
     record_validation_media(
         config,
         save_path,

tests/test_model_card.py

@@ -7,6 +7,7 @@
 from unittest.mock import MagicMock, patch
 
 import torch
+from PIL import Image
 
 from simpletuner.helpers.publishing.huggingface import HubManager
 from simpletuner.helpers.publishing.metadata import *
@@ -72,6 +73,8 @@ def setUp(self):
         self.args.validation_guidance_skip_layers = None
         self.args.validation_seed = 1234
         self.args.validation_noise_scheduler = "ddim"
+        self.args.validation_image_format = "png"
+        self.args.validation_image_quality = 90
         self.args.model_card_safe_for_work = True
         self.args.learning_rate = 1e-4
         self.args.max_grad_norm = 1.0
@@ -83,6 +86,7 @@ def setUp(self):
         self.args.base_model_precision = "no_change"
         self.args.flux_guidance_mode = "constant"
         self.args.flux_guidance_value = 1.0
+        self.args.peft_lora_mode = "standard"
         self.args.t5_padding = "unmodified"
         self.args.enable_xformers_memory_efficient_attention = False
         self.args.attention_mechanism = "diffusers"
@@ -561,6 +565,85 @@ def test_upload_full_model_uses_empty_repo_path_for_top_level_uploads(self):
         hub_manager._hub_api.upload_folder.assert_called_once()
         self.assertEqual(hub_manager._hub_api.upload_folder.call_args.kwargs["path_in_repo"], "")
 
+    def test_save_model_card_honors_validation_image_format_for_assets(self):
+        self.args.model_family = "sdxl"
+        self.args.model_type = "lora"
+        self.args.lora_type = "standard"
+        self.args.validation_image_format = "jpeg"
+        self.args.validation_image_quality = 81
+        self.args.controlnet = False
+        self.args.control = False
+        model = MagicMock(
+            MODEL_LICENSE="other",
+            PREDICTION_TYPE=SimpleNamespace(value="epsilon"),
+            MODEL_TYPE=SimpleNamespace(value="unet"),
+            gligen=False,
+        )
+        model.validation_audio_sample_rate.return_value = 44100
+        model.custom_model_card_schedule_info.return_value = ""
+
+        with tempfile.TemporaryDirectory() as tmpdir:
+            with (
+                patch("simpletuner.helpers.publishing.metadata.StateTracker.get_model_family", return_value="sdxl"),
+                patch("simpletuner.helpers.publishing.metadata.StateTracker.get_data_backends", return_value={}),
+                patch("simpletuner.helpers.publishing.metadata.StateTracker.get_weight_dtype", return_value=torch.bfloat16),
+                patch(
+                    "simpletuner.helpers.publishing.metadata.StateTracker.get_accelerator",
+                    return_value=MagicMock(num_processes=1),
+                ),
+                patch("simpletuner.helpers.publishing.metadata.StateTracker.get_args", return_value=self.args),
+                patch("simpletuner.helpers.publishing.metadata.StateTracker.get_model", return_value=model),
+            ):
+                save_model_card(
+                    repo_id="test-repo",
+                    images={"prompt": [Image.new("RGBA", (8, 8), color=(255, 0, 0, 128))]},
+                    base_model="test-base-model",
+                    train_text_encoder=False,
+                    prompt="Test prompt",
+                    validation_prompts=["Test prompt"],
+                    validation_shortnames=["prompt"],
+                    repo_folder=tmpdir,
+                    model=model,
+                    global_step=1000,
+                    epoch=1,
+                )
+
+            readme = Path(tmpdir, "README.md").read_text(encoding="utf-8")
+            self.assertIn("url: ./assets/image_0_0.jpg", readme)
+            self.assertTrue(Path(tmpdir, "assets", "image_0_0.jpg").exists())
+            self.assertFalse(Path(tmpdir, "assets", "image_0_0.png").exists())
+
+    def test_checkpoint_validation_media_filter_matches_shortnames_exactly(self):
+        with tempfile.TemporaryDirectory() as tmpdir:
+            output_dir = Path(tmpdir)
+            validation_dir = output_dir / "validation_images"
+            validation_dir.mkdir()
+            (validation_dir / "step_50_prompt_0_64x64.jpg").write_bytes(b"jpg")
+            (validation_dir / "step_50_prompt_adapter_0_64x64.jpg").write_bytes(b"adapter")
+            (validation_dir / "step_49_prompt_0_64x64.jpg").write_bytes(b"old")
+            (validation_dir / "step_50_prompt_0.wav").write_bytes(b"wav")
+
+            hub_manager = object.__new__(HubManager)
+            hub_manager.config = SimpleNamespace(output_dir=str(output_dir))
+
+            images, audios = hub_manager._filter_checkpoint_validation_media(
+                50,
+                {"prompt": [], "prompt_adapter": []},
+                {"prompt": []},
+            )
+
+            self.assertEqual(
+                {key: [Path(path).name for path in paths] for key, paths in images.items()},
+                {
+                    "prompt": ["step_50_prompt_0_64x64.jpg"],
+                    "prompt_adapter": ["step_50_prompt_adapter_0_64x64.jpg"],
+                },
+            )
+            self.assertEqual(
+                {key: [Path(path).name for path in paths] for key, paths in audios.items()},
+                {"prompt": ["step_50_prompt_0.wav"]},
+            )
+
     def test_save_training_config_sanitizes_public_export(self):
         config = SimpleNamespace(
             output_dir="output/test",