Fix mixed bucket key types during metadata cache refresh

simpletuner/helpers/data_backend/factory.py

@@ -10,7 +10,7 @@
 from math import sqrt
 from pathlib import Path
 from types import SimpleNamespace
-from typing import Any, Dict, Iterable, List, Optional, Tuple, Union
+from typing import Any, Dict, List, Optional, Tuple, Union
 
 from .runtime import BatchFetcher
 
@@ -66,26 +66,6 @@ def _normalise_vae_cache_config(config: Dict[str, Any]) -> Tuple[bool, bool]:
     return vae_cache_disable, vae_cache_ondemand
 
 
-def _coerce_bucket_keys(indices: Dict[Any, Iterable]) -> Dict[Any, list]:
-    """Return a copy of aspect ratio bucket indices with numeric keys coerced to float."""
-    coerced: Dict[Any, list] = {}
-    for key, values in (indices or {}).items():
-        try:
-            coerced_key: Any = float(key)
-        except (TypeError, ValueError):
-            coerced_key = key
-        if isinstance(values, dict):
-            iterable_values = [values]
-        elif isinstance(values, str):
-            iterable_values = [values]
-        elif isinstance(values, Iterable):
-            iterable_values = list(values)
-        else:
-            iterable_values = [values]
-        coerced.setdefault(coerced_key, []).extend(iterable_values)
-    return coerced
-
-
 import numpy as np
 import pandas as pd
 import torch
@@ -3179,8 +3159,6 @@ def _configure_metadata_backend(self, backend: Dict[str, Any], init_backend: Dic
         # Restore the live-authoritative runtime config after metadata cache loading.
         StateTracker.set_data_backend_config(init_backend["id"], init_backend["config"])
         metadata_backend = init_backend["metadata_backend"]
-        if isinstance(getattr(metadata_backend, "aspect_ratio_bucket_indices", None), dict):
-            metadata_backend.aspect_ratio_bucket_indices = _coerce_bucket_keys(metadata_backend.aspect_ratio_bucket_indices)
         if hasattr(metadata_backend, "attach_bucket_report"):
             metadata_backend.attach_bucket_report(init_backend.get("bucket_report"))
         if hasattr(metadata_backend, "_mock_children"):

simpletuner/helpers/metadata/backends/base.py

@@ -235,7 +235,10 @@ def aspect_ratio_bucket_indices(self):
 
     @aspect_ratio_bucket_indices.setter
     def aspect_ratio_bucket_indices(self, value):
-        """Set aspect ratio bucket indices with debug tracking."""
+        """Keep bucket keys identical during discovery, cache loading, and resume."""
+        normalized = {}
+        for key, samples in value.items():
+            normalized.setdefault(str(key), []).extend(samples)
         if hasattr(self, "_aspect_ratio_bucket_indices"):
             old_count = sum(len(v) for v in self._aspect_ratio_bucket_indices.values())
             new_count = sum(len(v) for v in value.values()) if value else 0
@@ -246,7 +249,7 @@ def aspect_ratio_bucket_indices(self, value):
                     f"Old buckets: {list(self._aspect_ratio_bucket_indices.keys())}, "
                     f"New buckets: {list(value.keys()) if value else []}"
                 )
-        self._aspect_ratio_bucket_indices = value
+        self._aspect_ratio_bucket_indices = normalized
 
     def _extract_audio_config(self) -> Dict[str, Any]:
         if self.dataset_config is None:

simpletuner/helpers/metadata/backends/discovery.py

@@ -18,19 +18,6 @@
 from simpletuner.helpers.training.multi_process import should_log
 from simpletuner.helpers.training.state_tracker import StateTracker
 
-
-def _coerce_bucket_keys_to_float(indices: dict) -> dict:
-    """Coerce bucket keys from strings to floats (fixes JSON serialization issue)."""
-    coerced = {}
-    for key, values in (indices or {}).items():
-        try:
-            coerced_key = float(key)
-        except (TypeError, ValueError):
-            coerced_key = key
-        coerced[coerced_key] = list(values) if not isinstance(values, list) else values
-    return coerced
-
-
 logger = logging.getLogger("DiscoveryMetadataBackend")
 if should_log():
     target_level = os.environ.get("SIMPLETUNER_LOG_LEVEL", "INFO")
@@ -340,9 +327,7 @@ def reload_cache(self, set_config: bool = True):
             except Exception as e:
                 logger.warning(f"Error loading aspect bucket cache, creating new one: {e}")
                 cache_data = {}
-            # Coerce bucket keys from strings to floats (JSON serialization converts float keys to strings)
-            loaded_indices = cache_data.get("aspect_ratio_bucket_indices", {})
-            self.aspect_ratio_bucket_indices = _coerce_bucket_keys_to_float(loaded_indices)
+            self.aspect_ratio_bucket_indices = cache_data.get("aspect_ratio_bucket_indices", {})
             if set_config:
                 self.config = cache_data.get("config", {})
                 if self.config != {}:

simpletuner/helpers/metadata/backends/huggingface.py

@@ -21,18 +21,6 @@
 from simpletuner.helpers.training.state_tracker import StateTracker
 
 
-def _coerce_bucket_keys_to_float(indices: dict) -> dict:
-    """Coerce bucket keys from strings to floats (fixes JSON serialization issue)."""
-    coerced = {}
-    for key, values in (indices or {}).items():
-        try:
-            coerced_key = float(key)
-        except (TypeError, ValueError):
-            coerced_key = key
-        coerced[coerced_key] = list(values) if not isinstance(values, list) else values
-    return coerced
-
-
 def _dataset_type_value(dataset_type: Any) -> str:
     return str(getattr(dataset_type, "value", dataset_type)).lower()
 
@@ -385,9 +373,7 @@ def reload_cache(self, set_config: bool = True):
             except Exception as e:
                 logger.warning(f"Error loading aspect ratio bucket cache, creating new one: {e}")
                 cache_data = {}
-            # Coerce bucket keys from strings to floats (JSON serialization converts float keys to strings)
-            loaded_indices = cache_data.get("aspect_ratio_bucket_indices", {})
-            self.aspect_ratio_bucket_indices = _coerce_bucket_keys_to_float(loaded_indices)
+            self.aspect_ratio_bucket_indices = cache_data.get("aspect_ratio_bucket_indices", {})
             if set_config:
                 self.config = cache_data.get("config", {})
                 if self.config != {}:

simpletuner/helpers/metadata/backends/parquet.py

@@ -18,19 +18,6 @@
 from simpletuner.helpers.training import audio_file_extensions, image_file_extensions, video_file_extensions
 from simpletuner.helpers.training.state_tracker import StateTracker
 
-
-def _coerce_bucket_keys_to_float(indices: dict) -> dict:
-    """Coerce bucket keys from strings to floats (fixes JSON serialization issue)."""
-    coerced = {}
-    for key, values in (indices or {}).items():
-        try:
-            coerced_key = float(key)
-        except (TypeError, ValueError):
-            coerced_key = key
-        coerced[coerced_key] = list(values) if not isinstance(values, list) else values
-    return coerced
-
-
 logger = logging.getLogger("ParquetMetadataBackend")
 from simpletuner.helpers.training.multi_process import should_log
 
@@ -278,9 +265,7 @@ def reload_cache(self, set_config: bool = True):
             except Exception as e:
                 logger.warning(f"Error loading aspect ratio bucket cache, creating new one: {e}")
                 cache_data = {}
-            # Coerce bucket keys from strings to floats (JSON serialization converts float keys to strings)
-            loaded_indices = cache_data.get("aspect_ratio_bucket_indices", {})
-            self.aspect_ratio_bucket_indices = _coerce_bucket_keys_to_float(loaded_indices)
+            self.aspect_ratio_bucket_indices = cache_data.get("aspect_ratio_bucket_indices", {})
             if set_config:
                 self.config = cache_data.get("config", {})
                 if self.config != {}:

simpletuner/helpers/metadata/backends/webshart.py

@@ -27,17 +27,6 @@
     logger.setLevel("ERROR")
 
 
-def _coerce_bucket_keys_to_float(indices: dict) -> dict:
-    coerced = {}
-    for key, values in (indices or {}).items():
-        try:
-            coerced_key = float(key)
-        except (TypeError, ValueError):
-            coerced_key = key
-        coerced[coerced_key] = list(values) if not isinstance(values, list) else values
-    return coerced
-
-
 class WebshartMetadataBackend(MetadataBackend):
     def __init__(
         self,
@@ -162,9 +151,7 @@ def reload_cache(self, set_config: bool = True):
             except Exception as exc:
                 logger.warning("Error loading webshart aspect bucket cache, creating new one: %s", exc)
                 cache_data = {}
-            self.aspect_ratio_bucket_indices = _coerce_bucket_keys_to_float(
-                cache_data.get("aspect_ratio_bucket_indices", {})
-            )
+            self.aspect_ratio_bucket_indices = cache_data.get("aspect_ratio_bucket_indices", {})
             self._sync_image_files_with_buckets()
             if set_config:
                 self.config = cache_data.get("config", {})
@@ -354,15 +341,15 @@ def _prepare_bucket_entry(
         shard_metadata: dict,
         entry: dict,
         sample_path: str,
-    ) -> tuple[dict, Optional[tuple[float, dict]], Optional[Exception]]:
+    ) -> tuple[dict, Optional[tuple[str, dict]], Optional[Exception]]:
         try:
             filename = str(entry["filename"])
             sample_metadata = self._metadata_for_entry(shard_metadata, filename, entry, sample_path)
             return sample_metadata, self._prepare_metadata(sample_path, sample_metadata), None
         except Exception as exc:
             return {}, None, exc
 
-    def _prepare_metadata(self, sample_path: str, sample_metadata: dict) -> Optional[tuple[float, dict]]:
+    def _prepare_metadata(self, sample_path: str, sample_metadata: dict) -> Optional[tuple[str, dict]]:
         if not sample_metadata or "original_size" not in sample_metadata:
             return None
         if not self.meets_resolution_requirements(image_metadata=sample_metadata):
@@ -392,7 +379,7 @@ def _prepare_metadata(self, sample_path: str, sample_metadata: dict) -> Optional
             )
             sample_metadata["bucket_frames"] = rounded_frames
         else:
-            bucket_key = round(aspect_ratio, 2)
+            bucket_key = str(round(aspect_ratio, 2))
         return bucket_key, sample_metadata
 
     def _entries_for_shard(self, shard_idx: int) -> list[dict]:

simpletuner/helpers/multiaspect/sampler.py

@@ -203,14 +203,14 @@ def load_states(self, state_path: str):
         if isinstance(saved_schedule, dict) and self._saved_schedule_is_restorable(previous_state, state_path):
             self.metadata_backend.aspect_ratio_bucket_indices = saved_schedule
             self._val_master_list = sorted(sum(saved_schedule.values(), []))
-        self.buckets = previous_state.get("buckets", self.load_buckets())
+        self.buckets = [str(bucket) for bucket in previous_state.get("buckets", self.load_buckets())]
         if "current_bucket" in previous_state:
             self.current_bucket = previous_state["current_bucket"]
 
         self.exhausted_buckets = []
         if "exhausted_buckets" in previous_state:
             self.logger.info(f"Previous checkpoint had {len(previous_state['exhausted_buckets'])} exhausted buckets.")
-            self.exhausted_buckets = previous_state["exhausted_buckets"]
+            self.exhausted_buckets = [str(bucket) for bucket in previous_state["exhausted_buckets"]]
         self.current_epoch = 1
         if "current_epoch" in previous_state:
             self.logger.info(f"Previous checkpoint was on epoch {previous_state['current_epoch']}.")
@@ -230,7 +230,7 @@ def load_states(self, state_path: str):
             self.metadata_backend.seen_images.update(normalized_seen)
 
     def load_buckets(self):
-        return list(self.metadata_backend.aspect_ratio_bucket_indices.keys())  # These keys are a float value, eg. 1.78.
+        return list(self.metadata_backend.aspect_ratio_bucket_indices.keys())
 
     def retrieve_validation_set(self, batch_size: int):
         """
@@ -436,36 +436,7 @@ def _reset_buckets(self, raise_exhaustion_signal: bool = True):
             raise MultiDatasetExhausted()
 
     def _get_bucket_images(self, bucket):
-        """
-        Safely retrieve bucket images, trying both original type and type conversion.
-
-        Args:
-            bucket: The bucket key (could be float or str)
-
-        Returns:
-            list: List of images in the bucket, or empty list if bucket not found
-        """
-        # Try the original bucket key first
-        if bucket in self.metadata_backend.aspect_ratio_bucket_indices:
-            return self.metadata_backend.aspect_ratio_bucket_indices[bucket]
-
-        # Try converting between str and float
-        try:
-            if isinstance(bucket, str):
-                # Try converting str to float
-                bucket_as_float = float(bucket)
-                if bucket_as_float in self.metadata_backend.aspect_ratio_bucket_indices:
-                    return self.metadata_backend.aspect_ratio_bucket_indices[bucket_as_float]
-            elif isinstance(bucket, (float, int)):
-                # Try converting float/int to str
-                bucket_as_str = str(bucket)
-                if bucket_as_str in self.metadata_backend.aspect_ratio_bucket_indices:
-                    return self.metadata_backend.aspect_ratio_bucket_indices[bucket_as_str]
-        except (ValueError, TypeError):
-            pass
-
-        # Bucket not found with either type
-        return []
+        return self.metadata_backend.aspect_ratio_bucket_indices.get(str(bucket), [])
 
     def _filter_unseen_occurrences(self, images):
         """Filter consumed positions without collapsing duplicate filepaths."""

tests/test_dataset.py

@@ -12,7 +12,7 @@
 from PIL import Image
 
 from simpletuner.helpers.data_backend.base import BaseDataBackend
-from simpletuner.helpers.data_backend.factory import _coerce_bucket_keys, check_column_values
+from simpletuner.helpers.data_backend.factory import check_column_values
 from simpletuner.helpers.metadata.backends.discovery import DiscoveryMetadataBackend
 from simpletuner.helpers.multiaspect.dataset import MultiAspectDataset
 
@@ -169,18 +169,6 @@ def test_invalid_data_type(self):
             check_column_values(column_data, "test_column", "test_file.parquet")
         self.assertIn("Unsupported data type in column", str(context.exception))
 
-    def test_coerce_bucket_keys(self):
-        indices = {"1.0": ["foo"], 1.5: ["bar"], "invalid": ["baz"], "single": "path"}
-        coerced = _coerce_bucket_keys(indices)
-        self.assertIn(1.0, coerced)
-        self.assertEqual(coerced[1.0], ["foo"])
-        self.assertIn(1.5, coerced)
-        self.assertEqual(coerced[1.5], ["bar"])
-        self.assertIn("invalid", coerced)
-        self.assertEqual(coerced["invalid"], ["baz"])
-        self.assertIn("single", coerced)
-        self.assertEqual(coerced["single"], ["path"])
-
 
 if __name__ == "__main__":
     unittest.main()