InvokeAI 6.10.0rc2__py3-none-any.whl → 6.11.0rc1__py3-none-any.whl

invokeai/backend/model_manager/configs/vae.py

@@ -33,6 +33,25 @@ REGEX_TO_BASE: dict[str, BaseModelType] = {
 }
 
 
+def _is_flux2_vae(state_dict: dict[str | int, Any]) -> bool:
+    """Check if state dict is a FLUX.2 VAE (AutoencoderKLFlux2).
+
+    FLUX.2 VAE can be identified by:
+    1. Batch Normalization layers (bn.running_mean, bn.running_var) - unique to FLUX.2
+    2. 32-dimensional latent space (decoder.conv_in has 32 input channels)
+
+    FLUX.1 VAE has 16-dimensional latent space and no BatchNorm layers.
+    """
+    # Check for BN layer which is unique to FLUX.2 VAE
+    has_bn = "bn.running_mean" in state_dict or "bn.running_var" in state_dict
+
+    # Check for 32-channel latent space (FLUX.2 has 32, FLUX.1 has 16)
+    decoder_conv_in_key = "decoder.conv_in.weight"
+    has_32_latent_channels = decoder_conv_in_key in state_dict and state_dict[decoder_conv_in_key].shape[1] == 32
+
+    return has_bn or has_32_latent_channels
+
+
 class VAE_Checkpoint_Config_Base(Checkpoint_Config_Base):
     """Model config for standalone VAE models."""
 
@@ -61,8 +80,9 @@ class VAE_Checkpoint_Config_Base(Checkpoint_Config_Base):
 
     @classmethod
     def _validate_looks_like_vae(cls, mod: ModelOnDisk) -> None:
+        state_dict = mod.load_state_dict()
         if not state_dict_has_any_keys_starting_with(
-            mod.load_state_dict(),
+            state_dict,
             {
                 "encoder.conv_in",
                 "decoder.conv_in",
@@ -70,9 +90,30 @@ class VAE_Checkpoint_Config_Base(Checkpoint_Config_Base):
         ):
             raise NotAMatchError("model does not match Checkpoint VAE heuristics")
 
+        # Exclude FLUX.2 VAEs - they have their own config class
+        if _is_flux2_vae(state_dict):
+            raise NotAMatchError("model is a FLUX.2 VAE, not a standard VAE")
+
     @classmethod
     def _get_base_or_raise(cls, mod: ModelOnDisk) -> BaseModelType:
-        # Heuristic: VAEs of all architectures have a similar structure; the best we can do is guess based on name
+        # First, try to identify by latent space dimensions (most reliable)
+        state_dict = mod.load_state_dict()
+        decoder_conv_in_key = "decoder.conv_in.weight"
+        if decoder_conv_in_key in state_dict:
+            latent_channels = state_dict[decoder_conv_in_key].shape[1]
+            if latent_channels == 16:
+                # Flux1 VAE has 16-dimensional latent space
+                return BaseModelType.Flux
+            elif latent_channels == 4:
+                # SD/SDXL VAE has 4-dimensional latent space
+                # Try to distinguish SD1/SD2/SDXL by name, fallback to SD1
+                for regexp, base in REGEX_TO_BASE.items():
+                    if re.search(regexp, mod.path.name, re.IGNORECASE):
+                        return base
+                # Default to SD1 if we can't determine from name
+                return BaseModelType.StableDiffusion1
+
+        # Fallback: guess based on name
         for regexp, base in REGEX_TO_BASE.items():
             if re.search(regexp, mod.path.name, re.IGNORECASE):
                 return base
@@ -96,6 +137,44 @@ class VAE_Checkpoint_FLUX_Config(VAE_Checkpoint_Config_Base, Config_Base):
     base: Literal[BaseModelType.Flux] = Field(default=BaseModelType.Flux)
 
 
+class VAE_Checkpoint_Flux2_Config(Checkpoint_Config_Base, Config_Base):
+    """Model config for FLUX.2 VAE checkpoint models (AutoencoderKLFlux2)."""
+
+    type: Literal[ModelType.VAE] = Field(default=ModelType.VAE)
+    format: Literal[ModelFormat.Checkpoint] = Field(default=ModelFormat.Checkpoint)
+    base: Literal[BaseModelType.Flux2] = Field(default=BaseModelType.Flux2)
+
+    @classmethod
+    def from_model_on_disk(cls, mod: ModelOnDisk, override_fields: dict[str, Any]) -> Self:
+        raise_if_not_file(mod)
+
+        raise_for_override_fields(cls, override_fields)
+
+        cls._validate_looks_like_vae(mod)
+
+        cls._validate_is_flux2_vae(mod)
+
+        return cls(**override_fields)
+
+    @classmethod
+    def _validate_looks_like_vae(cls, mod: ModelOnDisk) -> None:
+        if not state_dict_has_any_keys_starting_with(
+            mod.load_state_dict(),
+            {
+                "encoder.conv_in",
+                "decoder.conv_in",
+            },
+        ):
+            raise NotAMatchError("model does not match Checkpoint VAE heuristics")
+
+    @classmethod
+    def _validate_is_flux2_vae(cls, mod: ModelOnDisk) -> None:
+        """Validate that this is a FLUX.2 VAE, not FLUX.1."""
+        state_dict = mod.load_state_dict()
+        if not _is_flux2_vae(state_dict):
+            raise NotAMatchError("state dict does not look like a FLUX.2 VAE")
+
+
 class VAE_Diffusers_Config_Base(Diffusers_Config_Base):
     """Model config for standalone VAE models (diffusers version)."""
 
@@ -161,3 +240,26 @@ class VAE_Diffusers_SD1_Config(VAE_Diffusers_Config_Base, Config_Base):
 
 class VAE_Diffusers_SDXL_Config(VAE_Diffusers_Config_Base, Config_Base):
     base: Literal[BaseModelType.StableDiffusionXL] = Field(default=BaseModelType.StableDiffusionXL)
+
+
+class VAE_Diffusers_Flux2_Config(Diffusers_Config_Base, Config_Base):
+    """Model config for FLUX.2 VAE models in diffusers format (AutoencoderKLFlux2)."""
+
+    type: Literal[ModelType.VAE] = Field(default=ModelType.VAE)
+    format: Literal[ModelFormat.Diffusers] = Field(default=ModelFormat.Diffusers)
+    base: Literal[BaseModelType.Flux2] = Field(default=BaseModelType.Flux2)
+
+    @classmethod
+    def from_model_on_disk(cls, mod: ModelOnDisk, override_fields: dict[str, Any]) -> Self:
+        raise_if_not_dir(mod)
+
+        raise_for_override_fields(cls, override_fields)
+
+        raise_for_class_name(
+            common_config_paths(mod.path),
+            {
+                "AutoencoderKLFlux2",
+            },
+        )
+
+        return cls(**override_fields)

invokeai/backend/model_manager/load/load_default.py

@@ -75,7 +75,6 @@ class ModelLoader(ModelLoaderBase):
 
         config.path = str(self._get_model_path(config))
         self._ram_cache.make_room(self.get_size_fs(config, Path(config.path), submodel_type))
-        self._logger.info(f"Loading model '{stats_name}' into RAM cache..., config={config}")
         loaded_model = self._load_model(config, submodel_type)
 
         self._ram_cache.put(

invokeai/backend/model_manager/load/model_cache/model_cache.py

@@ -55,6 +55,21 @@ def synchronized(method: Callable[..., Any]) -> Callable[..., Any]:
     return wrapper
 
 
+def record_activity(method: Callable[..., Any]) -> Callable[..., Any]:
+    """A decorator that records activity after a method completes successfully.
+
+    Note: This decorator should be applied to methods that already hold self._lock.
+    """
+
+    @wraps(method)
+    def wrapper(self, *args, **kwargs):
+        result = method(self, *args, **kwargs)
+        self._record_activity()
+        return result
+
+    return wrapper
+
+
 @dataclass
 class CacheEntrySnapshot:
     cache_key: str
@@ -132,6 +147,7 @@ class ModelCache:
         storage_device: torch.device | str = "cpu",
         log_memory_usage: bool = False,
         logger: Optional[Logger] = None,
+        keep_alive_minutes: float = 0,
     ):
         """Initialize the model RAM cache.
 
@@ -151,6 +167,7 @@ class ModelCache:
             snapshots, so it is recommended to disable this feature unless you are actively inspecting the model cache's
             behaviour.
         :param logger: InvokeAILogger to use (otherwise creates one)
+        :param keep_alive_minutes: How long to keep models in cache after last use (in minutes). 0 means keep indefinitely.
         """
         self._enable_partial_loading = enable_partial_loading
         self._keep_ram_copy_of_weights = keep_ram_copy_of_weights
@@ -182,6 +199,12 @@ class ModelCache:
         self._on_cache_miss_callbacks: set[CacheMissCallback] = set()
         self._on_cache_models_cleared_callbacks: set[CacheModelsClearedCallback] = set()
 
+        # Keep-alive timeout support
+        self._keep_alive_minutes = keep_alive_minutes
+        self._last_activity_time: Optional[float] = None
+        self._timeout_timer: Optional[threading.Timer] = None
+        self._shutdown_event = threading.Event()
+
     def on_cache_hit(self, cb: CacheHitCallback) -> Callable[[], None]:
         self._on_cache_hit_callbacks.add(cb)
 
@@ -190,7 +213,7 @@ class ModelCache:
 
         return unsubscribe
 
-    def on_cache_miss(self, cb: CacheHitCallback) -> Callable[[], None]:
+    def on_cache_miss(self, cb: CacheMissCallback) -> Callable[[], None]:
         self._on_cache_miss_callbacks.add(cb)
 
         def unsubscribe() -> None:
@@ -217,8 +240,82 @@ class ModelCache:
     def stats(self, stats: CacheStats) -> None:
         """Set the CacheStats object for collecting cache statistics."""
         self._stats = stats
+        # Populate the cache size in the stats object when it's set
+        if self._stats is not None:
+            self._stats.cache_size = self._ram_cache_size_bytes
+
+    def _record_activity(self) -> None:
+        """Record model activity and reset the timeout timer if configured.
+
+        Note: This method should only be called when self._lock is already held.
+        """
+        if self._keep_alive_minutes <= 0:
+            return
+
+        self._last_activity_time = time.time()
+
+        # Cancel any existing timer
+        if self._timeout_timer is not None:
+            self._timeout_timer.cancel()
+
+        # Start a new timer
+        timeout_seconds = self._keep_alive_minutes * 60
+        self._timeout_timer = threading.Timer(timeout_seconds, self._on_timeout)
+        # Set as daemon so it doesn't prevent application shutdown
+        self._timeout_timer.daemon = True
+        self._timeout_timer.start()
+        self._logger.debug(f"Model cache activity recorded. Timeout set to {self._keep_alive_minutes} minutes.")
 
     @synchronized
+    @record_activity
+    def _on_timeout(self) -> None:
+        """Called when the keep-alive timeout expires. Clears the model cache."""
+        if self._shutdown_event.is_set():
+            return
+
+        # Double-check if there has been activity since the timer was set
+        # This handles the race condition where activity occurred just before the timer fired
+        if self._last_activity_time is not None and self._keep_alive_minutes > 0:
+            elapsed_minutes = (time.time() - self._last_activity_time) / 60
+            if elapsed_minutes < self._keep_alive_minutes:
+                # Activity occurred, don't clear cache
+                self._logger.debug(
+                    f"Model cache timeout fired but activity detected {elapsed_minutes:.2f} minutes ago. "
+                    f"Skipping cache clear."
+                )
+                return
+
+        # Check if there are any unlocked models that can be cleared
+        unlocked_models = [key for key, entry in self._cached_models.items() if not entry.is_locked]
+
+        if len(unlocked_models) > 0:
+            self._logger.info(
+                f"Model cache keep-alive timeout of {self._keep_alive_minutes} minutes expired. "
+                f"Clearing {len(unlocked_models)} unlocked model(s) from cache."
+            )
+            # Clear the cache by requesting a very large amount of space.
+            # This is the same logic used by the "Clear Model Cache" button.
+            # Using 1000 GB ensures all unlocked models are removed.
+            self._make_room_internal(1000 * GB)
+        elif len(self._cached_models) > 0:
+            # All models are locked, don't log at info level
+            self._logger.debug(
+                f"Model cache timeout fired but all {len(self._cached_models)} model(s) are locked. "
+                f"Skipping cache clear."
+            )
+        else:
+            self._logger.debug("Model cache timeout fired but cache is already empty.")
+
+    @synchronized
+    def shutdown(self) -> None:
+        """Shutdown the model cache, cancelling any pending timers."""
+        self._shutdown_event.set()
+        if self._timeout_timer is not None:
+            self._timeout_timer.cancel()
+            self._timeout_timer = None
+
+    @synchronized
+    @record_activity
     def put(self, key: str, model: AnyModel) -> None:
         """Add a model to the cache."""
         if key in self._cached_models:
@@ -228,7 +325,7 @@ class ModelCache:
             return
 
         size = calc_model_size_by_data(self._logger, model)
-        self.make_room(size)
+        self._make_room_internal(size)
 
         # Inject custom modules into the model.
         if isinstance(model, torch.nn.Module):
@@ -272,6 +369,7 @@ class ModelCache:
         return overview
 
     @synchronized
+    @record_activity
     def get(self, key: str, stats_name: Optional[str] = None) -> CacheRecord:
         """Retrieve a model from the cache.
 
@@ -309,9 +407,11 @@ class ModelCache:
         self._logger.debug(f"Cache hit: {key} (Type: {cache_entry.cached_model.model.__class__.__name__})")
         for cb in self._on_cache_hit_callbacks:
             cb(model_key=key, cache_snapshot=self._get_cache_snapshot())
+
         return cache_entry
 
     @synchronized
+    @record_activity
     def lock(self, cache_entry: CacheRecord, working_mem_bytes: Optional[int]) -> None:
         """Lock a model for use and move it into VRAM."""
         if cache_entry.key not in self._cached_models:
@@ -348,6 +448,7 @@ class ModelCache:
         self._log_cache_state()
 
     @synchronized
+    @record_activity
     def unlock(self, cache_entry: CacheRecord) -> None:
         """Unlock a model."""
         if cache_entry.key not in self._cached_models:
@@ -691,6 +792,10 @@ class ModelCache:
         external references to the model, there's nothing that the cache can do about it, and those models will not be
         garbage-collected.
         """
+        self._make_room_internal(bytes_needed)
+
+    def _make_room_internal(self, bytes_needed: int) -> None:
+        """Internal implementation of make_room(). Assumes the lock is already held."""
         self._logger.debug(f"Making room for {bytes_needed / MB:.2f}MB of RAM.")
         self._log_cache_state(title="Before dropping models:")