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

invokeai/app/services/invocation_stats/invocation_stats_default.py

@@ -52,8 +52,9 @@ class InvocationStatsService(InvocationStatsServiceBase):
         # Record state before the invocation.
         start_time = time.time()
         start_ram = psutil.Process().memory_info().rss
-        if torch.cuda.is_available():
-            torch.cuda.reset_peak_memory_stats()
+
+        # Remember current VRAM usage
+        vram_in_use = torch.cuda.memory_allocated() if torch.cuda.is_available() else 0.0
 
         assert services.model_manager.load is not None
         services.model_manager.load.ram_cache.stats = self._cache_stats[graph_execution_state_id]
@@ -62,14 +63,16 @@ class InvocationStatsService(InvocationStatsServiceBase):
             # Let the invocation run.
             yield None
         finally:
-            # Record state after the invocation.
+            # Record delta VRAM
+            delta_vram_gb = ((torch.cuda.memory_allocated() - vram_in_use) / GB) if torch.cuda.is_available() else 0.0
+
             node_stats = NodeExecutionStats(
                 invocation_type=invocation.get_type(),
                 start_time=start_time,
                 end_time=time.time(),
                 start_ram_gb=start_ram / GB,
                 end_ram_gb=psutil.Process().memory_info().rss / GB,
-                peak_vram_gb=torch.cuda.max_memory_allocated() / GB if torch.cuda.is_available() else 0.0,
+                delta_vram_gb=delta_vram_gb,
             )
             self._stats[graph_execution_state_id].add_node_execution_stats(node_stats)
 
@@ -81,6 +84,8 @@ class InvocationStatsService(InvocationStatsServiceBase):
         graph_stats_summary = self._get_graph_summary(graph_execution_state_id)
         node_stats_summaries = self._get_node_summaries(graph_execution_state_id)
         model_cache_stats_summary = self._get_model_cache_summary(graph_execution_state_id)
+        # Note: We use memory_allocated() here (not memory_reserved()) because we want to show
+        # the current actively-used VRAM, not the total reserved memory including PyTorch's cache.
         vram_usage_gb = torch.cuda.memory_allocated() / GB if torch.cuda.is_available() else None
 
         return InvocationStatsSummary(

invokeai/app/services/model_manager/model_manager_default.py

@@ -60,6 +60,10 @@ class ModelManagerService(ModelManagerServiceBase):
                 service.start(invoker)
 
     def stop(self, invoker: Invoker) -> None:
+        # Shutdown the model cache to cancel any pending timers
+        if hasattr(self._load, "ram_cache"):
+            self._load.ram_cache.shutdown()
+
         for service in [self._store, self._install, self._load]:
             if hasattr(service, "stop"):
                 service.stop(invoker)
@@ -88,7 +92,10 @@ class ModelManagerService(ModelManagerServiceBase):
             max_ram_cache_size_gb=app_config.max_cache_ram_gb,
             max_vram_cache_size_gb=app_config.max_cache_vram_gb,
             execution_device=execution_device or TorchDevice.choose_torch_device(),
+            storage_device="cpu",
+            log_memory_usage=app_config.log_memory_usage,
             logger=logger,
+            keep_alive_minutes=app_config.model_cache_keep_alive_min,
         )
         loader = ModelLoadService(
             app_config=app_config,

invokeai/app/services/model_records/model_records_base.py

@@ -19,11 +19,13 @@ from invokeai.backend.model_manager.configs.main import MainModelDefaultSettings
 from invokeai.backend.model_manager.taxonomy import (
     BaseModelType,
     ClipVariantType,
+    Flux2VariantType,
     FluxVariantType,
     ModelFormat,
     ModelSourceType,
     ModelType,
     ModelVariantType,
+    Qwen3VariantType,
     SchedulerPredictionType,
 )
 
@@ -89,8 +91,8 @@ class ModelRecordChanges(BaseModelExcludeNull):
 
     # Checkpoint-specific changes
     # TODO(MM2): Should we expose these? Feels footgun-y...
-    variant: Optional[ModelVariantType | ClipVariantType | FluxVariantType] = Field(
-        description="The variant of the model.", default=None
+    variant: Optional[ModelVariantType | ClipVariantType | FluxVariantType | Flux2VariantType | Qwen3VariantType] = (
+        Field(description="The variant of the model.", default=None)
     )
     prediction_type: Optional[SchedulerPredictionType] = Field(
         description="The prediction type of the model.", default=None

invokeai/app/services/shared/invocation_context.py

@@ -630,6 +630,21 @@ class UtilInterface(InvocationContextInterface):
             is_canceled=self.is_canceled,
         )
 
+    def flux2_step_callback(self, intermediate_state: PipelineIntermediateState) -> None:
+        """
+        The step callback for FLUX.2 Klein models (32-channel VAE).
+
+        Args:
+            intermediate_state: The intermediate state of the diffusion pipeline.
+        """
+
+        diffusion_step_callback(
+            signal_progress=self.signal_progress,
+            intermediate_state=intermediate_state,
+            base_model=BaseModelType.Flux2,
+            is_canceled=self.is_canceled,
+        )
+
     def signal_progress(
         self,
         message: str,

invokeai/app/services/shared/sqlite/sqlite_util.py

@@ -27,6 +27,7 @@ from invokeai.app.services.shared.sqlite_migrator.migrations.migration_21 import
 from invokeai.app.services.shared.sqlite_migrator.migrations.migration_22 import build_migration_22
 from invokeai.app.services.shared.sqlite_migrator.migrations.migration_23 import build_migration_23
 from invokeai.app.services.shared.sqlite_migrator.migrations.migration_24 import build_migration_24
+from invokeai.app.services.shared.sqlite_migrator.migrations.migration_25 import build_migration_25
 from invokeai.app.services.shared.sqlite_migrator.sqlite_migrator_impl import SqliteMigrator
 
 
@@ -71,6 +72,7 @@ def init_db(config: InvokeAIAppConfig, logger: Logger, image_files: ImageFileSto
     migrator.register_migration(build_migration_22(app_config=config, logger=logger))
     migrator.register_migration(build_migration_23(app_config=config, logger=logger))
     migrator.register_migration(build_migration_24(app_config=config, logger=logger))
+    migrator.register_migration(build_migration_25(app_config=config, logger=logger))
     migrator.run_migrations()
 
     return db

invokeai/app/services/shared/sqlite_migrator/migrations/migration_25.py

@@ -0,0 +1,61 @@
+import json
+import sqlite3
+from logging import Logger
+from typing import Any
+
+from invokeai.app.services.config import InvokeAIAppConfig
+from invokeai.app.services.shared.sqlite_migrator.sqlite_migrator_common import Migration
+from invokeai.backend.model_manager.taxonomy import ModelType, Qwen3VariantType
+
+
+class Migration25Callback:
+    def __init__(self, app_config: InvokeAIAppConfig, logger: Logger) -> None:
+        self._app_config = app_config
+        self._logger = logger
+
+    def __call__(self, cursor: sqlite3.Cursor) -> None:
+        cursor.execute("SELECT id, config FROM models;")
+        rows = cursor.fetchall()
+
+        migrated_count = 0
+
+        for model_id, config_json in rows:
+            try:
+                config_dict: dict[str, Any] = json.loads(config_json)
+
+                if config_dict.get("type") != ModelType.Qwen3Encoder.value:
+                    continue
+
+                if "variant" in config_dict:
+                    continue
+
+                config_dict["variant"] = Qwen3VariantType.Qwen3_4B.value
+
+                cursor.execute(
+                    "UPDATE models SET config = ? WHERE id = ?;",
+                    (json.dumps(config_dict), model_id),
+                )
+                migrated_count += 1
+
+            except json.JSONDecodeError as e:
+                self._logger.error("Invalid config JSON for model %s: %s", model_id, e)
+                raise
+
+        if migrated_count > 0:
+            self._logger.info(f"Migration complete: {migrated_count} Qwen3 encoder configs updated with variant field")
+        else:
+            self._logger.info("Migration complete: no Qwen3 encoder configs needed migration")
+
+
+def build_migration_25(app_config: InvokeAIAppConfig, logger: Logger) -> Migration:
+    """Builds the migration object for migrating from version 24 to version 25.
+
+    This migration adds the variant field to existing Qwen3 encoder models.
+    Models installed before the variant field was added will default to Qwen3_4B (for Z-Image compatibility).
+    """
+
+    return Migration(
+        from_version=24,
+        to_version=25,
+        callback=Migration25Callback(app_config=app_config, logger=logger),
+    )

invokeai/app/util/step_callback.py

@@ -93,14 +93,60 @@ COGVIEW4_LATENT_RGB_FACTORS = [
     [-0.00955853, -0.00980067, -0.00977842],
 ]
 
+# FLUX.2 uses 32 latent channels.
+# Factors from ComfyUI: https://github.com/Comfy-Org/ComfyUI/blob/main/comfy/latent_formats.py
+FLUX2_LATENT_RGB_FACTORS = [
+    #   R        G        B
+    [0.0058, 0.0113, 0.0073],
+    [0.0495, 0.0443, 0.0836],
+    [-0.0099, 0.0096, 0.0644],
+    [0.2144, 0.3009, 0.3652],
+    [0.0166, -0.0039, -0.0054],
+    [0.0157, 0.0103, -0.0160],
+    [-0.0398, 0.0902, -0.0235],
+    [-0.0052, 0.0095, 0.0109],
+    [-0.3527, -0.2712, -0.1666],
+    [-0.0301, -0.0356, -0.0180],
+    [-0.0107, 0.0078, 0.0013],
+    [0.0746, 0.0090, -0.0941],
+    [0.0156, 0.0169, 0.0070],
+    [-0.0034, -0.0040, -0.0114],
+    [0.0032, 0.0181, 0.0080],
+    [-0.0939, -0.0008, 0.0186],
+    [0.0018, 0.0043, 0.0104],
+    [0.0284, 0.0056, -0.0127],
+    [-0.0024, -0.0022, -0.0030],
+    [0.1207, -0.0026, 0.0065],
+    [0.0128, 0.0101, 0.0142],
+    [0.0137, -0.0072, -0.0007],
+    [0.0095, 0.0092, -0.0059],
+    [0.0000, -0.0077, -0.0049],
+    [-0.0465, -0.0204, -0.0312],
+    [0.0095, 0.0012, -0.0066],
+    [0.0290, -0.0034, 0.0025],
+    [0.0220, 0.0169, -0.0048],
+    [-0.0332, -0.0457, -0.0468],
+    [-0.0085, 0.0389, 0.0609],
+    [-0.0076, 0.0003, -0.0043],
+    [-0.0111, -0.0460, -0.0614],
+]
+
+FLUX2_LATENT_RGB_BIAS = [-0.0329, -0.0718, -0.0851]
+
 
 def sample_to_lowres_estimated_image(
-    samples: torch.Tensor, latent_rgb_factors: torch.Tensor, smooth_matrix: Optional[torch.Tensor] = None
+    samples: torch.Tensor,
+    latent_rgb_factors: torch.Tensor,
+    smooth_matrix: Optional[torch.Tensor] = None,
+    latent_rgb_bias: Optional[torch.Tensor] = None,
 ):
     if samples.dim() == 4:
         samples = samples[0]
     latent_image = samples.permute(1, 2, 0) @ latent_rgb_factors
 
+    if latent_rgb_bias is not None:
+        latent_image = latent_image + latent_rgb_bias
+
     if smooth_matrix is not None:
         latent_image = latent_image.unsqueeze(0).permute(3, 0, 1, 2)
         latent_image = torch.nn.functional.conv2d(latent_image, smooth_matrix.reshape((1, 1, 3, 3)), padding=1)
@@ -153,6 +199,7 @@ def diffusion_step_callback(
         sample = intermediate_state.latents
 
     smooth_matrix: list[list[float]] | None = None
+    latent_rgb_bias: list[float] | None = None
     if base_model in [BaseModelType.StableDiffusion1, BaseModelType.StableDiffusion2]:
         latent_rgb_factors = SD1_5_LATENT_RGB_FACTORS
     elif base_model in [BaseModelType.StableDiffusionXL, BaseModelType.StableDiffusionXLRefiner]:
@@ -164,6 +211,9 @@ def diffusion_step_callback(
         latent_rgb_factors = COGVIEW4_LATENT_RGB_FACTORS
     elif base_model == BaseModelType.Flux:
         latent_rgb_factors = FLUX_LATENT_RGB_FACTORS
+    elif base_model == BaseModelType.Flux2:
+        latent_rgb_factors = FLUX2_LATENT_RGB_FACTORS
+        latent_rgb_bias = FLUX2_LATENT_RGB_BIAS
     elif base_model == BaseModelType.ZImage:
         # Z-Image uses FLUX-compatible VAE with 16 latent channels
         latent_rgb_factors = FLUX_LATENT_RGB_FACTORS
@@ -174,8 +224,14 @@ def diffusion_step_callback(
     smooth_matrix_torch = (
         torch.tensor(smooth_matrix, dtype=sample.dtype, device=sample.device) if smooth_matrix else None
     )
+    latent_rgb_bias_torch = (
+        torch.tensor(latent_rgb_bias, dtype=sample.dtype, device=sample.device) if latent_rgb_bias else None
+    )
     image = sample_to_lowres_estimated_image(
-        samples=sample, latent_rgb_factors=latent_rgb_factors_torch, smooth_matrix=smooth_matrix_torch
+        samples=sample,
+        latent_rgb_factors=latent_rgb_factors_torch,
+        smooth_matrix=smooth_matrix_torch,
+        latent_rgb_bias=latent_rgb_bias_torch,
     )
 
     width = image.width * 8