diff --git a/invokeai/app/invocations/anima_latents_to_image.py b/invokeai/app/invocations/anima_latents_to_image.py index 26e3998ae71..04f966c8d46 100644 --- a/invokeai/app/invocations/anima_latents_to_image.py +++ b/invokeai/app/invocations/anima_latents_to_image.py @@ -28,6 +28,7 @@ from invokeai.app.invocations.primitives import ImageOutput from invokeai.app.services.shared.invocation_context import InvocationContext from invokeai.backend.flux.modules.autoencoder import AutoEncoder as FluxAutoEncoder +from invokeai.backend.model_manager.load.model_cache.utils import get_effective_device from invokeai.backend.util.devices import TorchDevice from invokeai.backend.util.vae_working_memory import ( estimate_vae_working_memory_anima, @@ -126,7 +127,8 @@ def invoke(self, context: InvocationContext) -> ImageOutput: raise TypeError(f"Expected AutoencoderKLWan or FluxAutoEncoder, got {type(vae).__name__}.") vae_dtype = next(iter(vae.parameters())).dtype - latents = latents.to(device=TorchDevice.choose_torch_device(), dtype=vae_dtype) + # Use the VAE's actual device (may be CPU if the model is configured cpu_only). + latents = latents.to(device=get_effective_device(vae), dtype=vae_dtype) TorchDevice.empty_cache() diff --git a/invokeai/app/invocations/cogview4_latents_to_image.py b/invokeai/app/invocations/cogview4_latents_to_image.py index 1b77ed8a1f8..bc9d208b669 100644 --- a/invokeai/app/invocations/cogview4_latents_to_image.py +++ b/invokeai/app/invocations/cogview4_latents_to_image.py @@ -17,6 +17,7 @@ from invokeai.app.invocations.model import VAEField from invokeai.app.invocations.primitives import ImageOutput from invokeai.app.services.shared.invocation_context import InvocationContext +from invokeai.backend.model_manager.load.model_cache.utils import get_effective_device from invokeai.backend.stable_diffusion.extensions.seamless import SeamlessExt from invokeai.backend.util.devices import TorchDevice from invokeai.backend.util.vae_working_memory import estimate_vae_working_memory_cogview4 @@ -54,7 +55,8 @@ def invoke(self, context: InvocationContext) -> ImageOutput: ): context.util.signal_progress("Running VAE") assert isinstance(vae, (AutoencoderKL)) - latents = latents.to(TorchDevice.choose_torch_device()) + # Use the VAE's actual device (may be CPU if the model is configured cpu_only). + latents = latents.to(get_effective_device(vae)) vae.disable_tiling() diff --git a/invokeai/app/invocations/flux2_vae_decode.py b/invokeai/app/invocations/flux2_vae_decode.py index ecbc7d9cb83..25ada406873 100644 --- a/invokeai/app/invocations/flux2_vae_decode.py +++ b/invokeai/app/invocations/flux2_vae_decode.py @@ -20,6 +20,7 @@ from invokeai.app.invocations.primitives import ImageOutput from invokeai.app.services.shared.invocation_context import InvocationContext from invokeai.backend.model_manager.load.load_base import LoadedModel +from invokeai.backend.model_manager.load.model_cache.utils import get_effective_device from invokeai.backend.util.devices import TorchDevice @@ -51,7 +52,8 @@ def _vae_decode(self, vae_info: LoadedModel, latents: torch.Tensor) -> Image.Ima """ with vae_info.model_on_device() as (_, vae): vae_dtype = next(iter(vae.parameters())).dtype - device = TorchDevice.choose_torch_device() + # Use the VAE's actual device (may be CPU if the model is configured cpu_only). + device = get_effective_device(vae) latents = latents.to(device=device, dtype=vae_dtype) # Decode using diffusers API diff --git a/invokeai/app/invocations/flux_vae_decode.py b/invokeai/app/invocations/flux_vae_decode.py index c55dfb539ac..400e36bff45 100644 --- a/invokeai/app/invocations/flux_vae_decode.py +++ b/invokeai/app/invocations/flux_vae_decode.py @@ -16,6 +16,7 @@ from invokeai.app.services.shared.invocation_context import InvocationContext from invokeai.backend.flux.modules.autoencoder import AutoEncoder from invokeai.backend.model_manager.load.load_base import LoadedModel +from invokeai.backend.model_manager.load.model_cache.utils import get_effective_device from invokeai.backend.util.devices import TorchDevice from invokeai.backend.util.vae_working_memory import estimate_vae_working_memory_flux @@ -47,7 +48,8 @@ def _vae_decode(self, vae_info: LoadedModel, latents: torch.Tensor) -> Image.Ima with vae_info.model_on_device(working_mem_bytes=estimated_working_memory) as (_, vae): assert isinstance(vae, AutoEncoder) vae_dtype = next(iter(vae.parameters())).dtype - latents = latents.to(device=TorchDevice.choose_torch_device(), dtype=vae_dtype) + # Use the VAE's actual device (may be CPU if the model is configured cpu_only). + latents = latents.to(device=get_effective_device(vae), dtype=vae_dtype) img = vae.decode(latents) img = img.clamp(-1, 1) diff --git a/invokeai/app/invocations/latents_to_image.py b/invokeai/app/invocations/latents_to_image.py index 608485a078b..5bd094c186d 100644 --- a/invokeai/app/invocations/latents_to_image.py +++ b/invokeai/app/invocations/latents_to_image.py @@ -18,6 +18,7 @@ from invokeai.app.invocations.model import VAEField from invokeai.app.invocations.primitives import ImageOutput from invokeai.app.services.shared.invocation_context import InvocationContext +from invokeai.backend.model_manager.load.model_cache.utils import get_effective_device from invokeai.backend.stable_diffusion.extensions.seamless import SeamlessExt from invokeai.backend.stable_diffusion.vae_tiling import patch_vae_tiling_params from invokeai.backend.util.devices import TorchDevice @@ -69,8 +70,13 @@ def invoke(self, context: InvocationContext) -> ImageOutput: ): context.util.signal_progress("Running VAE decoder") assert isinstance(vae, (AutoencoderKL, AutoencoderTiny)) - latents = latents.to(TorchDevice.choose_torch_device()) - if self.fp32: + # Use the VAE's actual device (may be CPU if the model is configured cpu_only). + device = get_effective_device(vae) + latents = latents.to(device) + # Force fp32 when running on CPU (e.g. when the VAE is configured cpu_only). fp16 conv does run on CPU with + # the pinned torch, but it's much slower than fp32 there, and SD/SDXL VAE has known fp16 overflow issues + # that produce black images — fp32 avoids both. + if self.fp32 or device.type == "cpu": # FP32 mode: convert everything to float32 for maximum precision vae.to(dtype=torch.float32) latents = latents.float() diff --git a/invokeai/app/invocations/qwen_image_latents_to_image.py b/invokeai/app/invocations/qwen_image_latents_to_image.py index 072185f147b..68c92be3310 100644 --- a/invokeai/app/invocations/qwen_image_latents_to_image.py +++ b/invokeai/app/invocations/qwen_image_latents_to_image.py @@ -17,6 +17,7 @@ from invokeai.app.invocations.model import VAEField from invokeai.app.invocations.primitives import ImageOutput from invokeai.app.services.shared.invocation_context import InvocationContext +from invokeai.backend.model_manager.load.model_cache.utils import get_effective_device from invokeai.backend.stable_diffusion.extensions.seamless import SeamlessExt from invokeai.backend.util.devices import TorchDevice from invokeai.backend.util.vae_working_memory import estimate_vae_working_memory_qwen_image @@ -53,7 +54,8 @@ def invoke(self, context: InvocationContext) -> ImageOutput: ): context.util.signal_progress("Running VAE") assert isinstance(vae, AutoencoderKLQwenImage) - latents = latents.to(device=TorchDevice.choose_torch_device(), dtype=vae.dtype) + # Use the VAE's actual device (may be CPU if the model is configured cpu_only). + latents = latents.to(device=get_effective_device(vae), dtype=vae.dtype) vae.disable_tiling() diff --git a/invokeai/app/invocations/sd3_latents_to_image.py b/invokeai/app/invocations/sd3_latents_to_image.py index e6a20d38a9c..38c93305df7 100644 --- a/invokeai/app/invocations/sd3_latents_to_image.py +++ b/invokeai/app/invocations/sd3_latents_to_image.py @@ -17,6 +17,7 @@ from invokeai.app.invocations.model import VAEField from invokeai.app.invocations.primitives import ImageOutput from invokeai.app.services.shared.invocation_context import InvocationContext +from invokeai.backend.model_manager.load.model_cache.utils import get_effective_device from invokeai.backend.stable_diffusion.extensions.seamless import SeamlessExt from invokeai.backend.util.devices import TorchDevice from invokeai.backend.util.vae_working_memory import estimate_vae_working_memory_sd3 @@ -56,7 +57,8 @@ def invoke(self, context: InvocationContext) -> ImageOutput: ): context.util.signal_progress("Running VAE") assert isinstance(vae, (AutoencoderKL)) - latents = latents.to(TorchDevice.choose_torch_device()) + # Use the VAE's actual device (may be CPU if the model is configured cpu_only). + latents = latents.to(get_effective_device(vae)) vae.disable_tiling() diff --git a/invokeai/app/invocations/z_image_latents_to_image.py b/invokeai/app/invocations/z_image_latents_to_image.py index a2e6fdcc077..6ba34632d44 100644 --- a/invokeai/app/invocations/z_image_latents_to_image.py +++ b/invokeai/app/invocations/z_image_latents_to_image.py @@ -19,6 +19,7 @@ from invokeai.app.invocations.primitives import ImageOutput from invokeai.app.services.shared.invocation_context import InvocationContext from invokeai.backend.flux.modules.autoencoder import AutoEncoder as FluxAutoEncoder +from invokeai.backend.model_manager.load.model_cache.utils import get_effective_device from invokeai.backend.stable_diffusion.extensions.seamless import SeamlessExt from invokeai.backend.util.devices import TorchDevice from invokeai.backend.util.vae_working_memory import estimate_vae_working_memory_flux @@ -75,7 +76,8 @@ def invoke(self, context: InvocationContext) -> ImageOutput: ) vae_dtype = next(iter(vae.parameters())).dtype - latents = latents.to(device=TorchDevice.choose_torch_device(), dtype=vae_dtype) + # Use the VAE's actual device (may be CPU if the model is configured cpu_only). + latents = latents.to(device=get_effective_device(vae), dtype=vae_dtype) # Disable tiling for AutoencoderKL if isinstance(vae, AutoencoderKL): diff --git a/invokeai/backend/model_manager/configs/vae.py b/invokeai/backend/model_manager/configs/vae.py index 5a88cf12781..30735b443ed 100644 --- a/invokeai/backend/model_manager/configs/vae.py +++ b/invokeai/backend/model_manager/configs/vae.py @@ -76,6 +76,7 @@ class VAE_Checkpoint_Config_Base(Checkpoint_Config_Base): type: Literal[ModelType.VAE] = Field(default=ModelType.VAE) format: Literal[ModelFormat.Checkpoint] = Field(default=ModelFormat.Checkpoint) + cpu_only: bool | None = Field(default=None, description="Whether this model should run on CPU only") @classmethod def from_model_on_disk(cls, mod: ModelOnDisk, override_fields: dict[str, Any]) -> Self: @@ -166,6 +167,7 @@ class VAE_Checkpoint_Flux2_Config(Checkpoint_Config_Base, Config_Base): type: Literal[ModelType.VAE] = Field(default=ModelType.VAE) format: Literal[ModelFormat.Checkpoint] = Field(default=ModelFormat.Checkpoint) base: Literal[BaseModelType.Flux2] = Field(default=BaseModelType.Flux2) + cpu_only: bool | None = Field(default=None, description="Whether this model should run on CPU only") @classmethod def from_model_on_disk(cls, mod: ModelOnDisk, override_fields: dict[str, Any]) -> Self: @@ -204,6 +206,7 @@ class VAE_Checkpoint_QwenImage_Config(Checkpoint_Config_Base, Config_Base): type: Literal[ModelType.VAE] = Field(default=ModelType.VAE) format: Literal[ModelFormat.Checkpoint] = Field(default=ModelFormat.Checkpoint) base: Literal[BaseModelType.QwenImage] = Field(default=BaseModelType.QwenImage) + cpu_only: bool | None = Field(default=None, description="Whether this model should run on CPU only") @classmethod def from_model_on_disk(cls, mod: ModelOnDisk, override_fields: dict[str, Any]) -> Self: @@ -241,6 +244,7 @@ class VAE_Checkpoint_Anima_Config(Checkpoint_Config_Base, Config_Base): type: Literal[ModelType.VAE] = Field(default=ModelType.VAE) format: Literal[ModelFormat.Checkpoint] = Field(default=ModelFormat.Checkpoint) base: Literal[BaseModelType.Anima] = Field(default=BaseModelType.Anima) + cpu_only: bool | None = Field(default=None, description="Whether this model should run on CPU only") @classmethod def from_model_on_disk(cls, mod: ModelOnDisk, override_fields: dict[str, Any]) -> Self: @@ -260,6 +264,7 @@ class VAE_Diffusers_Config_Base(Diffusers_Config_Base): type: Literal[ModelType.VAE] = Field(default=ModelType.VAE) format: Literal[ModelFormat.Diffusers] = Field(default=ModelFormat.Diffusers) + cpu_only: bool | None = Field(default=None, description="Whether this model should run on CPU only") @classmethod def from_model_on_disk(cls, mod: ModelOnDisk, override_fields: dict[str, Any]) -> Self: @@ -328,6 +333,7 @@ class VAE_Diffusers_Flux2_Config(Diffusers_Config_Base, Config_Base): type: Literal[ModelType.VAE] = Field(default=ModelType.VAE) format: Literal[ModelFormat.Diffusers] = Field(default=ModelFormat.Diffusers) base: Literal[BaseModelType.Flux2] = Field(default=BaseModelType.Flux2) + cpu_only: bool | None = Field(default=None, description="Whether this model should run on CPU only") @classmethod def from_model_on_disk(cls, mod: ModelOnDisk, override_fields: dict[str, Any]) -> Self: diff --git a/invokeai/backend/model_manager/load/model_cache/model_cache.py b/invokeai/backend/model_manager/load/model_cache/model_cache.py index e3a0928e52b..0061d9f162a 100644 --- a/invokeai/backend/model_manager/load/model_cache/model_cache.py +++ b/invokeai/backend/model_manager/load/model_cache/model_cache.py @@ -438,11 +438,24 @@ def lock(self, cache_entry: CacheRecord, working_mem_bytes: Optional[int]) -> No f"Locking model {cache_entry.key} (Type: {cache_entry.cached_model.model.__class__.__name__})" ) - # Check if the model's specific compute_device is CPU, not just the cache's default execution_device + # A CPU compute_device means there's no VRAM load to do. This happens in two distinct situations: + # 1. The model is explicitly configured cpu_only, but the cache's default execution device is a GPU. + # 2. The whole install is CPU-only (no GPU), so every model's compute_device is CPU by default. + # Only case 1 is noteworthy — surface it at INFO with the "(cpu_only)" cause so it mirrors the + # "Loaded model ... onto device" line emitted for GPU loads below. Case 2 would fire for every + # lock of every model and says nothing about a per-model choice, so keep it at DEBUG and drop the wording. model_compute_device = cache_entry.cached_model.compute_device if model_compute_device.type == "cpu": - # Models configured for CPU execution don't need to be loaded into VRAM - self._logger.debug(f"Model {cache_entry.key} is configured for CPU execution, skipping VRAM load") + if self._execution_device.type != "cpu": + self._logger.info( + f"Loaded model '{cache_entry.key}' ({cache_entry.cached_model.model.__class__.__name__}) onto " + f"cpu device (cpu_only); skipping VRAM load" + ) + else: + self._logger.debug( + f"Loaded model '{cache_entry.key}' ({cache_entry.cached_model.model.__class__.__name__}) onto " + f"cpu device; skipping VRAM load" + ) return try: diff --git a/invokeai/frontend/web/openapi.json b/invokeai/frontend/web/openapi.json index 3357f948fcd..c877494d23a 100644 --- a/invokeai/frontend/web/openapi.json +++ b/invokeai/frontend/web/openapi.json @@ -71868,6 +71868,18 @@ "const": "anima", "title": "Base", "default": "anima" + }, + "cpu_only": { + "anyOf": [ + { + "type": "boolean" + }, + { + "type": "null" + } + ], + "title": "Cpu Only", + "description": "Whether this model should run on CPU only" } }, "type": "object", @@ -71886,7 +71898,8 @@ "config_path", "type", "format", - "base" + "base", + "cpu_only" ], "title": "VAE_Checkpoint_Anima_Config", "description": "Model config for Anima QwenImage VAE checkpoint models (AutoencoderKLQwenImage)." @@ -71999,6 +72012,18 @@ "title": "Format", "default": "checkpoint" }, + "cpu_only": { + "anyOf": [ + { + "type": "boolean" + }, + { + "type": "null" + } + ], + "title": "Cpu Only", + "description": "Whether this model should run on CPU only" + }, "base": { "type": "string", "const": "flux", @@ -72022,6 +72047,7 @@ "config_path", "type", "format", + "cpu_only", "base" ], "title": "VAE_Checkpoint_FLUX_Config" @@ -72139,6 +72165,18 @@ "const": "flux2", "title": "Base", "default": "flux2" + }, + "cpu_only": { + "anyOf": [ + { + "type": "boolean" + }, + { + "type": "null" + } + ], + "title": "Cpu Only", + "description": "Whether this model should run on CPU only" } }, "type": "object", @@ -72157,7 +72195,8 @@ "config_path", "type", "format", - "base" + "base", + "cpu_only" ], "title": "VAE_Checkpoint_Flux2_Config", "description": "Model config for FLUX.2 VAE checkpoint models (AutoencoderKLFlux2)." @@ -72275,6 +72314,18 @@ "const": "qwen-image", "title": "Base", "default": "qwen-image" + }, + "cpu_only": { + "anyOf": [ + { + "type": "boolean" + }, + { + "type": "null" + } + ], + "title": "Cpu Only", + "description": "Whether this model should run on CPU only" } }, "type": "object", @@ -72293,7 +72344,8 @@ "config_path", "type", "format", - "base" + "base", + "cpu_only" ], "title": "VAE_Checkpoint_QwenImage_Config", "description": "Model config for Qwen Image VAE checkpoint models (AutoencoderKLQwenImage)." @@ -72406,6 +72458,18 @@ "title": "Format", "default": "checkpoint" }, + "cpu_only": { + "anyOf": [ + { + "type": "boolean" + }, + { + "type": "null" + } + ], + "title": "Cpu Only", + "description": "Whether this model should run on CPU only" + }, "base": { "type": "string", "const": "sd-1", @@ -72429,6 +72493,7 @@ "config_path", "type", "format", + "cpu_only", "base" ], "title": "VAE_Checkpoint_SD1_Config" @@ -72541,6 +72606,18 @@ "title": "Format", "default": "checkpoint" }, + "cpu_only": { + "anyOf": [ + { + "type": "boolean" + }, + { + "type": "null" + } + ], + "title": "Cpu Only", + "description": "Whether this model should run on CPU only" + }, "base": { "type": "string", "const": "sd-2", @@ -72564,6 +72641,7 @@ "config_path", "type", "format", + "cpu_only", "base" ], "title": "VAE_Checkpoint_SD2_Config" @@ -72676,6 +72754,18 @@ "title": "Format", "default": "checkpoint" }, + "cpu_only": { + "anyOf": [ + { + "type": "boolean" + }, + { + "type": "null" + } + ], + "title": "Cpu Only", + "description": "Whether this model should run on CPU only" + }, "base": { "type": "string", "const": "sdxl", @@ -72699,6 +72789,7 @@ "config_path", "type", "format", + "cpu_only", "base" ], "title": "VAE_Checkpoint_SDXL_Config" @@ -72808,6 +72899,18 @@ "const": "flux2", "title": "Base", "default": "flux2" + }, + "cpu_only": { + "anyOf": [ + { + "type": "boolean" + }, + { + "type": "null" + } + ], + "title": "Cpu Only", + "description": "Whether this model should run on CPU only" } }, "type": "object", @@ -72826,7 +72929,8 @@ "format", "repo_variant", "type", - "base" + "base", + "cpu_only" ], "title": "VAE_Diffusers_Flux2_Config", "description": "Model config for FLUX.2 VAE models in diffusers format (AutoencoderKLFlux2)." @@ -72931,6 +73035,18 @@ "title": "Type", "default": "vae" }, + "cpu_only": { + "anyOf": [ + { + "type": "boolean" + }, + { + "type": "null" + } + ], + "title": "Cpu Only", + "description": "Whether this model should run on CPU only" + }, "base": { "type": "string", "const": "sd-1", @@ -72954,6 +73070,7 @@ "format", "repo_variant", "type", + "cpu_only", "base" ], "title": "VAE_Diffusers_SD1_Config" @@ -73058,6 +73175,18 @@ "title": "Type", "default": "vae" }, + "cpu_only": { + "anyOf": [ + { + "type": "boolean" + }, + { + "type": "null" + } + ], + "title": "Cpu Only", + "description": "Whether this model should run on CPU only" + }, "base": { "type": "string", "const": "sdxl", @@ -73081,6 +73210,7 @@ "format", "repo_variant", "type", + "cpu_only", "base" ], "title": "VAE_Diffusers_SDXL_Config" diff --git a/invokeai/frontend/web/public/locales/en.json b/invokeai/frontend/web/public/locales/en.json index 050af54fcf1..363bf0da7e0 100644 --- a/invokeai/frontend/web/public/locales/en.json +++ b/invokeai/frontend/web/public/locales/en.json @@ -1142,6 +1142,7 @@ "cpuOnly": "CPU Only", "fp8Storage": "FP8 Storage (Save VRAM)", "runOnCpu": "Run text encoder model on CPU only", + "runVaeOnCpu": "Run VAE model on CPU only", "noDefaultSettings": "No default settings configured for this model. Visit the Model Manager to add default settings.", "defaultSettings": "Default Settings", "defaultSettingsSaved": "Default Settings Saved", @@ -2464,6 +2465,13 @@ "This saves VRAM for the denoiser while only slightly impacting performance. The conditioning outputs are automatically moved to GPU for the denoiser." ] }, + "cpuOnlyVae": { + "heading": "CPU Only", + "paragraphs": [ + "When enabled, the VAE will run on CPU instead of GPU.", + "This saves VRAM during decoding, but VAE decoding will be slower. The decoded image is automatically moved back to the correct device." + ] + }, "fp8Storage": { "heading": "FP8 Storage", "paragraphs": [ diff --git a/invokeai/frontend/web/src/common/components/InformationalPopover/constants.ts b/invokeai/frontend/web/src/common/components/InformationalPopover/constants.ts index e9d855648ad..d858d630c5d 100644 --- a/invokeai/frontend/web/src/common/components/InformationalPopover/constants.ts +++ b/invokeai/frontend/web/src/common/components/InformationalPopover/constants.ts @@ -78,6 +78,7 @@ export type Feature = | 'optimizedDenoising' | 'fluxDevLicense' | 'cpuOnly' + | 'cpuOnlyVae' | 'fp8Storage'; export type PopoverData = PopoverProps & { diff --git a/invokeai/frontend/web/src/features/modelManagerV2/hooks/useCpuOnlyModelSettings.ts b/invokeai/frontend/web/src/features/modelManagerV2/hooks/useCpuOnlyModelSettings.ts new file mode 100644 index 00000000000..961dbb73974 --- /dev/null +++ b/invokeai/frontend/web/src/features/modelManagerV2/hooks/useCpuOnlyModelSettings.ts @@ -0,0 +1,23 @@ +import type { FormField } from 'features/modelManagerV2/subpanels/ModelPanel/MainModelDefaultSettings/MainModelDefaultSettings'; +import { useMemo } from 'react'; + +export type CpuOnlyModelSettingsFormData = { + cpuOnly: FormField; +}; + +/** + * Computes the default form state for the "run on CPU only" toggle. Shared by every model type that + * exposes a standalone `cpu_only` config field (text encoders and VAEs). + */ +export const useCpuOnlyModelSettings = (modelConfig: { cpu_only?: boolean | null }) => { + return useMemo(() => { + const cpuOnly = modelConfig.cpu_only ?? false; + + return { + cpuOnly: { + value: cpuOnly, + isEnabled: cpuOnly, + }, + }; + }, [modelConfig]); +}; diff --git a/invokeai/frontend/web/src/features/modelManagerV2/hooks/useEncoderModelSettings.ts b/invokeai/frontend/web/src/features/modelManagerV2/hooks/useEncoderModelSettings.ts deleted file mode 100644 index 6b3e9d71010..00000000000 --- a/invokeai/frontend/web/src/features/modelManagerV2/hooks/useEncoderModelSettings.ts +++ /dev/null @@ -1,35 +0,0 @@ -import type { EncoderModelSettingsFormData } from 'features/modelManagerV2/subpanels/ModelPanel/EncoderModelSettings/EncoderModelSettings'; -import { useMemo } from 'react'; -import type { - CLIPEmbedModelConfig, - CLIPVisionModelConfig, - LlavaOnevisionModelConfig, - Qwen3EncoderModelConfig, - SigLIPModelConfig, - T5EncoderModelConfig, - TextLLMModelConfig, -} from 'services/api/types'; - -type EncoderModelConfig = - | CLIPEmbedModelConfig - | T5EncoderModelConfig - | Qwen3EncoderModelConfig - | CLIPVisionModelConfig - | SigLIPModelConfig - | LlavaOnevisionModelConfig - | TextLLMModelConfig; - -export const useEncoderModelSettings = (modelConfig: EncoderModelConfig) => { - const encoderModelSettingsDefaults = useMemo(() => { - const cpuOnly = modelConfig.cpu_only ?? false; - - return { - cpuOnly: { - value: cpuOnly, - isEnabled: cpuOnly, - }, - }; - }, [modelConfig]); - - return encoderModelSettingsDefaults; -}; diff --git a/invokeai/frontend/web/src/features/modelManagerV2/subpanels/ModelPanel/CpuOnlyModelSettings/CpuOnlyModelSettings.tsx b/invokeai/frontend/web/src/features/modelManagerV2/subpanels/ModelPanel/CpuOnlyModelSettings/CpuOnlyModelSettings.tsx new file mode 100644 index 00000000000..7331d3e4622 --- /dev/null +++ b/invokeai/frontend/web/src/features/modelManagerV2/subpanels/ModelPanel/CpuOnlyModelSettings/CpuOnlyModelSettings.tsx @@ -0,0 +1,139 @@ +import { Button, Flex, FormControl, FormLabel, Heading, Switch } from '@invoke-ai/ui-library'; +import { useAppSelector } from 'app/store/storeHooks'; +import type { Feature } from 'common/components/InformationalPopover/constants'; +import { InformationalPopover } from 'common/components/InformationalPopover/InformationalPopover'; +import type { CpuOnlyModelSettingsFormData } from 'features/modelManagerV2/hooks/useCpuOnlyModelSettings'; +import { useCpuOnlyModelSettings } from 'features/modelManagerV2/hooks/useCpuOnlyModelSettings'; +import { selectSelectedModelKey } from 'features/modelManagerV2/store/modelManagerV2Slice'; +import type { FormField } from 'features/modelManagerV2/subpanels/ModelPanel/MainModelDefaultSettings/MainModelDefaultSettings'; +import { toast } from 'features/toast/toast'; +import type { ChangeEvent } from 'react'; +import { memo, useCallback, useEffect, useMemo } from 'react'; +import type { Control, SubmitHandler } from 'react-hook-form'; +import { useController, useForm } from 'react-hook-form'; +import { useTranslation } from 'react-i18next'; +import { PiCheckBold } from 'react-icons/pi'; +import { useUpdateModelMutation } from 'services/api/endpoints/models'; + +const CpuOnlyToggle = memo( + (props: { name: 'cpuOnly'; control: Control; feature: Feature; label: string }) => { + const { field } = useController({ name: props.name, control: props.control }); + const { t } = useTranslation(); + + const onChange = useCallback( + (e: ChangeEvent) => { + const updatedValue = { + ...(field.value as FormField), + value: e.target.checked, + isEnabled: e.target.checked, + }; + field.onChange(updatedValue); + }, + [field] + ); + + const value = useMemo(() => { + return (field.value as FormField).value; + }, [field.value]); + + return ( + + + {t(props.label)} + + + + ); + } +); + +CpuOnlyToggle.displayName = 'CpuOnlyToggle'; + +type Props = { + modelConfig: { cpu_only?: boolean | null }; + /** Popover feature key describing the CPU-only behavior for this model type. */ + feature: Feature; + /** i18n key for the toggle label. */ + label: string; + /** Prefix for the save success/failure toast ids, e.g. `ENCODER_SETTINGS` or `VAE_SETTINGS`. */ + toastIdBase: string; +}; + +/** + * Settings panel exposing the standalone `cpu_only` toggle. Shared by every model type that carries + * a `cpu_only` config field (text encoders and VAEs) — the only per-type differences are the label, + * popover copy, and toast ids, which are passed in as props. + */ +export const CpuOnlyModelSettings = memo(({ modelConfig, feature, label, toastIdBase }: Props) => { + const selectedModelKey = useAppSelector(selectSelectedModelKey); + const { t } = useTranslation(); + + const settingsDefaults = useCpuOnlyModelSettings(modelConfig); + const [updateModel, { isLoading: isLoadingUpdateModel }] = useUpdateModelMutation(); + + const { handleSubmit, control, formState, reset } = useForm({ + defaultValues: settingsDefaults, + }); + + useEffect(() => { + reset(settingsDefaults); + }, [settingsDefaults, reset]); + + const onSubmit = useCallback>( + (data) => { + if (!selectedModelKey) { + return; + } + + const body = { + cpu_only: data.cpuOnly.isEnabled ? data.cpuOnly.value : null, + }; + + updateModel({ + key: selectedModelKey, + body, + }) + .unwrap() + .then((_) => { + toast({ + id: `${toastIdBase}_SAVED`, + title: t('modelManager.settingsSaved'), + status: 'success', + }); + reset(data); + }) + .catch((error) => { + if (error) { + toast({ + id: `${toastIdBase}_SAVE_FAILED`, + title: `${error.data.detail} `, + status: 'error', + }); + } + }); + }, + [selectedModelKey, updateModel, toastIdBase, t, reset] + ); + + return ( + <> + + {t('modelManager.settings')} + + + + + + ); +}); + +CpuOnlyModelSettings.displayName = 'CpuOnlyModelSettings'; diff --git a/invokeai/frontend/web/src/features/modelManagerV2/subpanels/ModelPanel/EncoderModelSettings/EncoderModelSettings.tsx b/invokeai/frontend/web/src/features/modelManagerV2/subpanels/ModelPanel/EncoderModelSettings/EncoderModelSettings.tsx index 9bfe3974ddf..455df7673a9 100644 --- a/invokeai/frontend/web/src/features/modelManagerV2/subpanels/ModelPanel/EncoderModelSettings/EncoderModelSettings.tsx +++ b/invokeai/frontend/web/src/features/modelManagerV2/subpanels/ModelPanel/EncoderModelSettings/EncoderModelSettings.tsx @@ -1,17 +1,5 @@ -import { Button, Flex, FormControl, FormLabel, Heading, Switch } from '@invoke-ai/ui-library'; -import { useAppSelector } from 'app/store/storeHooks'; -import { InformationalPopover } from 'common/components/InformationalPopover/InformationalPopover'; -import { useEncoderModelSettings } from 'features/modelManagerV2/hooks/useEncoderModelSettings'; -import { selectSelectedModelKey } from 'features/modelManagerV2/store/modelManagerV2Slice'; -import type { FormField } from 'features/modelManagerV2/subpanels/ModelPanel/MainModelDefaultSettings/MainModelDefaultSettings'; -import { toast } from 'features/toast/toast'; -import type { ChangeEvent } from 'react'; -import { memo, useCallback, useEffect, useMemo } from 'react'; -import type { Control, SubmitHandler } from 'react-hook-form'; -import { useController, useForm } from 'react-hook-form'; -import { useTranslation } from 'react-i18next'; -import { PiCheckBold } from 'react-icons/pi'; -import { useUpdateModelMutation } from 'services/api/endpoints/models'; +import { CpuOnlyModelSettings } from 'features/modelManagerV2/subpanels/ModelPanel/CpuOnlyModelSettings/CpuOnlyModelSettings'; +import { memo } from 'react'; import type { CLIPEmbedModelConfig, CLIPVisionModelConfig, @@ -22,10 +10,6 @@ import type { TextLLMModelConfig, } from 'services/api/types'; -export type EncoderModelSettingsFormData = { - cpuOnly: FormField; -}; - type EncoderModelConfig = | CLIPEmbedModelConfig | T5EncoderModelConfig @@ -39,107 +23,14 @@ type Props = { modelConfig: EncoderModelConfig; }; -const DefaultCpuOnly = memo((props: { name: 'cpuOnly'; control: Control }) => { - const { field } = useController(props); - const { t } = useTranslation(); - - const onChange = useCallback( - (e: ChangeEvent) => { - const updatedValue = { - ...(field.value as FormField), - value: e.target.checked, - isEnabled: e.target.checked, - }; - field.onChange(updatedValue); - }, - [field] - ); - - const value = useMemo(() => { - return (field.value as FormField).value; - }, [field.value]); - - return ( - - - {t('modelManager.runOnCpu')} - - - - ); -}); - -DefaultCpuOnly.displayName = 'DefaultCpuOnly'; - export const EncoderModelSettings = memo(({ modelConfig }: Props) => { - const selectedModelKey = useAppSelector(selectSelectedModelKey); - const { t } = useTranslation(); - - const settingsDefaults = useEncoderModelSettings(modelConfig); - const [updateModel, { isLoading: isLoadingUpdateModel }] = useUpdateModelMutation(); - - const { handleSubmit, control, formState, reset } = useForm({ - defaultValues: settingsDefaults, - }); - - useEffect(() => { - reset(settingsDefaults); - }, [settingsDefaults, reset]); - - const onSubmit = useCallback>( - (data) => { - if (!selectedModelKey) { - return; - } - - const body = { - cpu_only: data.cpuOnly.isEnabled ? data.cpuOnly.value : null, - }; - - updateModel({ - key: selectedModelKey, - body, - }) - .unwrap() - .then((_) => { - toast({ - id: 'ENCODER_SETTINGS_SAVED', - title: t('modelManager.settingsSaved'), - status: 'success', - }); - reset(data); - }) - .catch((error) => { - if (error) { - toast({ - id: 'ENCODER_SETTINGS_SAVE_FAILED', - title: `${error.data.detail} `, - status: 'error', - }); - } - }); - }, - [selectedModelKey, reset, updateModel, t] - ); - return ( - <> - - {t('modelManager.settings')} - - - - - + ); }); diff --git a/invokeai/frontend/web/src/features/modelManagerV2/subpanels/ModelPanel/ModelView.tsx b/invokeai/frontend/web/src/features/modelManagerV2/subpanels/ModelPanel/ModelView.tsx index 365f7cff4b8..e666ebfd1b4 100644 --- a/invokeai/frontend/web/src/features/modelManagerV2/subpanels/ModelPanel/ModelView.tsx +++ b/invokeai/frontend/web/src/features/modelManagerV2/subpanels/ModelPanel/ModelView.tsx @@ -9,6 +9,7 @@ import { ModelHeader } from 'features/modelManagerV2/subpanels/ModelPanel/ModelH import { ModelSettingsExportButton } from 'features/modelManagerV2/subpanels/ModelPanel/ModelSettingsExportButton'; import { ModelSettingsImportButton } from 'features/modelManagerV2/subpanels/ModelPanel/ModelSettingsImportButton'; import { TriggerPhrases } from 'features/modelManagerV2/subpanels/ModelPanel/TriggerPhrases'; +import { VAEModelSettings } from 'features/modelManagerV2/subpanels/ModelPanel/VAEModelSettings/VAEModelSettings'; import { filesize } from 'filesize'; import { memo, useMemo } from 'react'; import { useTranslation } from 'react-i18next'; @@ -82,6 +83,10 @@ export const ModelView = memo(({ modelConfig }: Props) => { if (isEncoderModel(modelConfig)) { return true; } + // VAE models (cpu_only toggle) + if (modelConfig.type === 'vae') { + return true; + } return false; }, [modelConfig]); @@ -151,6 +156,7 @@ export const ModelView = memo(({ modelConfig }: Props) => { )} {modelConfig.type === 'main' && } {isEncoderModel(modelConfig) && } + {modelConfig.type === 'vae' && } )} diff --git a/invokeai/frontend/web/src/features/modelManagerV2/subpanels/ModelPanel/VAEModelSettings/VAEModelSettings.tsx b/invokeai/frontend/web/src/features/modelManagerV2/subpanels/ModelPanel/VAEModelSettings/VAEModelSettings.tsx new file mode 100644 index 00000000000..ad6d08b6040 --- /dev/null +++ b/invokeai/frontend/web/src/features/modelManagerV2/subpanels/ModelPanel/VAEModelSettings/VAEModelSettings.tsx @@ -0,0 +1,20 @@ +import { CpuOnlyModelSettings } from 'features/modelManagerV2/subpanels/ModelPanel/CpuOnlyModelSettings/CpuOnlyModelSettings'; +import { memo } from 'react'; +import type { VAEModelConfig } from 'services/api/types'; + +type Props = { + modelConfig: VAEModelConfig; +}; + +export const VAEModelSettings = memo(({ modelConfig }: Props) => { + return ( + + ); +}); + +VAEModelSettings.displayName = 'VAEModelSettings'; diff --git a/invokeai/frontend/web/src/services/api/schema.ts b/invokeai/frontend/web/src/services/api/schema.ts index 567cd08f9e8..29d9e6f2d76 100644 --- a/invokeai/frontend/web/src/services/api/schema.ts +++ b/invokeai/frontend/web/src/services/api/schema.ts @@ -31746,6 +31746,11 @@ export type components = { * @constant */ base: "anima"; + /** + * Cpu Only + * @description Whether this model should run on CPU only + */ + cpu_only: boolean | null; }; /** VAE_Checkpoint_FLUX_Config */ VAE_Checkpoint_FLUX_Config: { @@ -31818,6 +31823,11 @@ export type components = { * @constant */ format: "checkpoint"; + /** + * Cpu Only + * @description Whether this model should run on CPU only + */ + cpu_only: boolean | null; /** * Base * @default flux @@ -31905,6 +31915,11 @@ export type components = { * @constant */ base: "flux2"; + /** + * Cpu Only + * @description Whether this model should run on CPU only + */ + cpu_only: boolean | null; }; /** * VAE_Checkpoint_QwenImage_Config @@ -31986,6 +32001,11 @@ export type components = { * @constant */ base: "qwen-image"; + /** + * Cpu Only + * @description Whether this model should run on CPU only + */ + cpu_only: boolean | null; }; /** VAE_Checkpoint_SD1_Config */ VAE_Checkpoint_SD1_Config: { @@ -32058,6 +32078,11 @@ export type components = { * @constant */ format: "checkpoint"; + /** + * Cpu Only + * @description Whether this model should run on CPU only + */ + cpu_only: boolean | null; /** * Base * @default sd-1 @@ -32136,6 +32161,11 @@ export type components = { * @constant */ format: "checkpoint"; + /** + * Cpu Only + * @description Whether this model should run on CPU only + */ + cpu_only: boolean | null; /** * Base * @default sd-2 @@ -32214,6 +32244,11 @@ export type components = { * @constant */ format: "checkpoint"; + /** + * Cpu Only + * @description Whether this model should run on CPU only + */ + cpu_only: boolean | null; /** * Base * @default sdxl @@ -32298,6 +32333,11 @@ export type components = { * @constant */ base: "flux2"; + /** + * Cpu Only + * @description Whether this model should run on CPU only + */ + cpu_only: boolean | null; }; /** VAE_Diffusers_SD1_Config */ VAE_Diffusers_SD1_Config: { @@ -32367,6 +32407,11 @@ export type components = { * @constant */ type: "vae"; + /** + * Cpu Only + * @description Whether this model should run on CPU only + */ + cpu_only: boolean | null; /** * Base * @default sd-1 @@ -32442,6 +32487,11 @@ export type components = { * @constant */ type: "vae"; + /** + * Cpu Only + * @description Whether this model should run on CPU only + */ + cpu_only: boolean | null; /** * Base * @default sdxl diff --git a/tests/backend/model_manager/load/test_load_default_cpu_only.py b/tests/backend/model_manager/load/test_load_default_cpu_only.py new file mode 100644 index 00000000000..d99eb5466d0 --- /dev/null +++ b/tests/backend/model_manager/load/test_load_default_cpu_only.py @@ -0,0 +1,49 @@ +"""Tests for `ModelLoader._get_execution_device` — the helper that forces a model onto the CPU +when its config requests `cpu_only`. + +A VAE (or text encoder) configured with `cpu_only=True` must load onto the CPU so its weights +never occupy VRAM. The loader signals this by returning `torch.device("cpu")` from +`_get_execution_device`, which is then passed to `ModelCache.put(..., execution_device=...)`. +""" + +from types import SimpleNamespace +from typing import Optional + +import torch + +from invokeai.backend.model_manager.load.load_default import ModelLoader +from invokeai.backend.model_manager.taxonomy import SubModelType + + +def _loader() -> ModelLoader: + # `_get_execution_device` only reads the config, so an uninitialized loader is sufficient. + return ModelLoader.__new__(ModelLoader) + + +def _vae_config(cpu_only: Optional[bool]) -> SimpleNamespace: + # Mirrors the relevant surface of a standalone VAE config: a `cpu_only` field and no + # `default_settings` (VAE configs do not carry default settings). + return SimpleNamespace(cpu_only=cpu_only, default_settings=None) + + +def test_vae_cpu_only_true_returns_cpu(): + assert _loader()._get_execution_device(_vae_config(cpu_only=True), None) == torch.device("cpu") + + +def test_vae_cpu_only_false_or_unset_returns_none(): + # Falsy values must not force CPU execution — the cache falls back to its default device. + assert _loader()._get_execution_device(_vae_config(cpu_only=False), None) is None + assert _loader()._get_execution_device(_vae_config(cpu_only=None), None) is None + + +def test_vae_cpu_only_applies_regardless_of_submodel_type(): + # The VAE is loaded as a standalone model (submodel_type=None), but the standalone branch + # must not depend on the submodel type either way. + loader = _loader() + assert loader._get_execution_device(_vae_config(cpu_only=True), SubModelType.VAE) == torch.device("cpu") + + +def test_config_without_cpu_only_attr_returns_none(): + # A config type that has neither `cpu_only` nor `default_settings` must be left on the + # cache default (return None), not crash. + assert _loader()._get_execution_device(SimpleNamespace(), None) is None