Run ruff, fix bug in hf downloading code which failed to download parts of a model

invokeai/app/invocations/fields.py

@@ -41,6 +41,7 @@ class UIType(str, Enum, metaclass=MetaEnum):
     # region Model Field Types
     MainModel = "MainModelField"
     FluxMainModel = "FluxMainModelField"
+    SD3MainModel = "SD3MainModelField"
     SDXLMainModel = "SDXLMainModelField"
     SDXLRefinerModel = "SDXLRefinerModelField"
     ONNXModel = "ONNXModelField"
@@ -133,6 +134,7 @@ class FieldDescriptions:
     clip_embed_model = "CLIP Embed loader"
     unet = "UNet (scheduler, LoRAs)"
     transformer = "Transformer"
+    mmditx = "MMDiTX"
     vae = "VAE"
     cond = "Conditioning tensor"
     controlnet_model = "ControlNet model to load"
@@ -140,6 +142,7 @@ class FieldDescriptions:
     lora_model = "LoRA model to load"
     main_model = "Main model (UNet, VAE, CLIP) to load"
     flux_model = "Flux model (Transformer) to load"
+    sd3_model = "SD3 model (MMDiTX) to load"
     sdxl_main_model = "SDXL Main model (UNet, VAE, CLIP1, CLIP2) to load"
     sdxl_refiner_model = "SDXL Refiner Main Modde (UNet, VAE, CLIP2) to load"
     onnx_main_model = "ONNX Main model (UNet, VAE, CLIP) to load"
@@ -246,6 +249,12 @@ class FluxConditioningField(BaseModel):
     conditioning_name: str = Field(description="The name of conditioning tensor")
 
 
+class SD3ConditioningField(BaseModel):
+    """A conditioning tensor primitive value"""
+
+    conditioning_name: str = Field(description="The name of conditioning tensor")
+
+
 class ConditioningField(BaseModel):
     """A conditioning tensor primitive value"""
 

invokeai/app/invocations/flux_model_loader.py

@@ -0,0 +1,89 @@
+from typing import Literal
+
+from invokeai.app.invocations.baseinvocation import (
+    BaseInvocation,
+    BaseInvocationOutput,
+    Classification,
+    invocation,
+    invocation_output,
+)
+from invokeai.app.invocations.fields import FieldDescriptions, Input, InputField, OutputField, UIType
+from invokeai.app.invocations.model import CLIPField, ModelIdentifierField, T5EncoderField, TransformerField, VAEField
+from invokeai.app.services.shared.invocation_context import InvocationContext
+from invokeai.backend.flux.util import max_seq_lengths
+from invokeai.backend.model_manager.config import (
+    CheckpointConfigBase,
+    SubModelType,
+)
+
+
+@invocation_output("flux_model_loader_output")
+class FluxModelLoaderOutput(BaseInvocationOutput):
+    """Flux base model loader output"""
+
+    transformer: TransformerField = OutputField(description=FieldDescriptions.transformer, title="Transformer")
+    clip: CLIPField = OutputField(description=FieldDescriptions.clip, title="CLIP")
+    t5_encoder: T5EncoderField = OutputField(description=FieldDescriptions.t5_encoder, title="T5 Encoder")
+    vae: VAEField = OutputField(description=FieldDescriptions.vae, title="VAE")
+    max_seq_len: Literal[256, 512] = OutputField(
+        description="The max sequence length to used for the T5 encoder. (256 for schnell transformer, 512 for dev transformer)",
+        title="Max Seq Length",
+    )
+
+
+@invocation(
+    "flux_model_loader",
+    title="Flux Main Model",
+    tags=["model", "flux"],
+    category="model",
+    version="1.0.4",
+    classification=Classification.Prototype,
+)
+class FluxModelLoaderInvocation(BaseInvocation):
+    """Loads a flux base model, outputting its submodels."""
+
+    model: ModelIdentifierField = InputField(
+        description=FieldDescriptions.flux_model,
+        ui_type=UIType.FluxMainModel,
+        input=Input.Direct,
+    )
+
+    t5_encoder_model: ModelIdentifierField = InputField(
+        description=FieldDescriptions.t5_encoder, ui_type=UIType.T5EncoderModel, input=Input.Direct, title="T5 Encoder"
+    )
+
+    clip_embed_model: ModelIdentifierField = InputField(
+        description=FieldDescriptions.clip_embed_model,
+        ui_type=UIType.CLIPEmbedModel,
+        input=Input.Direct,
+        title="CLIP Embed",
+    )
+
+    vae_model: ModelIdentifierField = InputField(
+        description=FieldDescriptions.vae_model, ui_type=UIType.FluxVAEModel, title="VAE"
+    )
+
+    def invoke(self, context: InvocationContext) -> FluxModelLoaderOutput:
+        for key in [self.model.key, self.t5_encoder_model.key, self.clip_embed_model.key, self.vae_model.key]:
+            if not context.models.exists(key):
+                raise ValueError(f"Unknown model: {key}")
+
+        transformer = self.model.model_copy(update={"submodel_type": SubModelType.Transformer})
+        vae = self.vae_model.model_copy(update={"submodel_type": SubModelType.VAE})
+
+        tokenizer = self.clip_embed_model.model_copy(update={"submodel_type": SubModelType.Tokenizer})
+        clip_encoder = self.clip_embed_model.model_copy(update={"submodel_type": SubModelType.TextEncoder})
+
+        tokenizer2 = self.t5_encoder_model.model_copy(update={"submodel_type": SubModelType.Tokenizer2})
+        t5_encoder = self.t5_encoder_model.model_copy(update={"submodel_type": SubModelType.TextEncoder2})
+
+        transformer_config = context.models.get_config(transformer)
+        assert isinstance(transformer_config, CheckpointConfigBase)
+
+        return FluxModelLoaderOutput(
+            transformer=TransformerField(transformer=transformer, loras=[]),
+            clip=CLIPField(tokenizer=tokenizer, text_encoder=clip_encoder, loras=[], skipped_layers=0),
+            t5_encoder=T5EncoderField(tokenizer=tokenizer2, text_encoder=t5_encoder),
+            vae=VAEField(vae=vae),
+            max_seq_len=max_seq_lengths[transformer_config.config_path],
+        )

invokeai/app/invocations/model.py

@@ -1,5 +1,5 @@
 import copy
-from typing import List, Literal, Optional
+from typing import List, Optional
 
 from pydantic import BaseModel, Field
 
@@ -13,11 +13,9 @@ from invokeai.app.invocations.baseinvocation import (
 from invokeai.app.invocations.fields import FieldDescriptions, Input, InputField, OutputField, UIType
 from invokeai.app.services.shared.invocation_context import InvocationContext
 from invokeai.app.shared.models import FreeUConfig
-from invokeai.backend.flux.util import max_seq_lengths
 from invokeai.backend.model_manager.config import (
     AnyModelConfig,
     BaseModelType,
-    CheckpointConfigBase,
     ModelType,
     SubModelType,
 )
@@ -139,78 +137,6 @@ class ModelIdentifierInvocation(BaseInvocation):
         return ModelIdentifierOutput(model=self.model)
 
 
-@invocation_output("flux_model_loader_output")
-class FluxModelLoaderOutput(BaseInvocationOutput):
-    """Flux base model loader output"""
-
-    transformer: TransformerField = OutputField(description=FieldDescriptions.transformer, title="Transformer")
-    clip: CLIPField = OutputField(description=FieldDescriptions.clip, title="CLIP")
-    t5_encoder: T5EncoderField = OutputField(description=FieldDescriptions.t5_encoder, title="T5 Encoder")
-    vae: VAEField = OutputField(description=FieldDescriptions.vae, title="VAE")
-    max_seq_len: Literal[256, 512] = OutputField(
-        description="The max sequence length to used for the T5 encoder. (256 for schnell transformer, 512 for dev transformer)",
-        title="Max Seq Length",
-    )
-
-
-@invocation(
-    "flux_model_loader",
-    title="Flux Main Model",
-    tags=["model", "flux"],
-    category="model",
-    version="1.0.4",
-    classification=Classification.Prototype,
-)
-class FluxModelLoaderInvocation(BaseInvocation):
-    """Loads a flux base model, outputting its submodels."""
-
-    model: ModelIdentifierField = InputField(
-        description=FieldDescriptions.flux_model,
-        ui_type=UIType.FluxMainModel,
-        input=Input.Direct,
-    )
-
-    t5_encoder_model: ModelIdentifierField = InputField(
-        description=FieldDescriptions.t5_encoder, ui_type=UIType.T5EncoderModel, input=Input.Direct, title="T5 Encoder"
-    )
-
-    clip_embed_model: ModelIdentifierField = InputField(
-        description=FieldDescriptions.clip_embed_model,
-        ui_type=UIType.CLIPEmbedModel,
-        input=Input.Direct,
-        title="CLIP Embed",
-    )
-
-    vae_model: ModelIdentifierField = InputField(
-        description=FieldDescriptions.vae_model, ui_type=UIType.FluxVAEModel, title="VAE"
-    )
-
-    def invoke(self, context: InvocationContext) -> FluxModelLoaderOutput:
-        for key in [self.model.key, self.t5_encoder_model.key, self.clip_embed_model.key, self.vae_model.key]:
-            if not context.models.exists(key):
-                raise ValueError(f"Unknown model: {key}")
-
-        transformer = self.model.model_copy(update={"submodel_type": SubModelType.Transformer})
-        vae = self.vae_model.model_copy(update={"submodel_type": SubModelType.VAE})
-
-        tokenizer = self.clip_embed_model.model_copy(update={"submodel_type": SubModelType.Tokenizer})
-        clip_encoder = self.clip_embed_model.model_copy(update={"submodel_type": SubModelType.TextEncoder})
-
-        tokenizer2 = self.t5_encoder_model.model_copy(update={"submodel_type": SubModelType.Tokenizer2})
-        t5_encoder = self.t5_encoder_model.model_copy(update={"submodel_type": SubModelType.TextEncoder2})
-
-        transformer_config = context.models.get_config(transformer)
-        assert isinstance(transformer_config, CheckpointConfigBase)
-
-        return FluxModelLoaderOutput(
-            transformer=TransformerField(transformer=transformer, loras=[]),
-            clip=CLIPField(tokenizer=tokenizer, text_encoder=clip_encoder, loras=[], skipped_layers=0),
-            t5_encoder=T5EncoderField(tokenizer=tokenizer2, text_encoder=t5_encoder),
-            vae=VAEField(vae=vae),
-            max_seq_len=max_seq_lengths[transformer_config.config_path],
-        )
-
-
 @invocation(
     "main_model_loader",
     title="Main Model",

invokeai/app/invocations/primitives.py

@@ -18,6 +18,7 @@ from invokeai.app.invocations.fields import (
     InputField,
     LatentsField,
     OutputField,
+    SD3ConditioningField,
     TensorField,
     UIComponent,
 )
@@ -426,6 +427,17 @@ class FluxConditioningOutput(BaseInvocationOutput):
         return cls(conditioning=FluxConditioningField(conditioning_name=conditioning_name))
 
 
+@invocation_output("sd3_conditioning_output")
+class SD3ConditioningOutput(BaseInvocationOutput):
+    """Base class for nodes that output a single SD3 conditioning tensor"""
+
+    conditioning: SD3ConditioningField = OutputField(description=FieldDescriptions.cond)
+
+    @classmethod
+    def build(cls, conditioning_name: str) -> "SD3ConditioningOutput":
+        return cls(conditioning=SD3ConditioningField(conditioning_name=conditioning_name))
+
+
 @invocation_output("conditioning_output")
 class ConditioningOutput(BaseInvocationOutput):
     """Base class for nodes that output a single conditioning tensor"""

invokeai/app/invocations/sd3_denoise.py

@@ -0,0 +1,260 @@
+from typing import Callable, Tuple
+
+import torch
+from diffusers.models.transformers.transformer_sd3 import SD3Transformer2DModel
+from diffusers.schedulers.scheduling_flow_match_euler_discrete import FlowMatchEulerDiscreteScheduler
+from tqdm import tqdm
+
+from invokeai.app.invocations.baseinvocation import BaseInvocation, Classification, invocation
+from invokeai.app.invocations.constants import LATENT_SCALE_FACTOR
+from invokeai.app.invocations.fields import (
+    FieldDescriptions,
+    Input,
+    InputField,
+    SD3ConditioningField,
+    WithBoard,
+    WithMetadata,
+)
+from invokeai.app.invocations.model import TransformerField
+from invokeai.app.invocations.primitives import LatentsOutput
+from invokeai.app.invocations.sd3_text_encoder import SD3_T5_MAX_SEQ_LEN
+from invokeai.app.services.shared.invocation_context import InvocationContext
+from invokeai.backend.model_manager.config import BaseModelType
+from invokeai.backend.stable_diffusion.diffusers_pipeline import PipelineIntermediateState
+from invokeai.backend.stable_diffusion.diffusion.conditioning_data import SD3ConditioningInfo
+from invokeai.backend.util.devices import TorchDevice
+
+
+@invocation(
+    "sd3_denoise",
+    title="SD3 Denoise",
+    tags=["image", "sd3"],
+    category="image",
+    version="1.0.0",
+    classification=Classification.Prototype,
+)
+class SD3DenoiseInvocation(BaseInvocation, WithMetadata, WithBoard):
+    """Run denoising process with a SD3 model."""
+
+    transformer: TransformerField = InputField(
+        description=FieldDescriptions.sd3_model,
+        input=Input.Connection,
+        title="Transformer",
+    )
+    positive_text_conditioning: SD3ConditioningField = InputField(
+        description=FieldDescriptions.positive_cond, input=Input.Connection
+    )
+    negative_text_conditioning: SD3ConditioningField = InputField(
+        description=FieldDescriptions.negative_cond, input=Input.Connection
+    )
+    cfg_scale: float | list[float] = InputField(default=3.5, description=FieldDescriptions.cfg_scale, title="CFG Scale")
+    width: int = InputField(default=1024, multiple_of=16, description="Width of the generated image.")
+    height: int = InputField(default=1024, multiple_of=16, description="Height of the generated image.")
+    num_steps: int = InputField(default=10, gt=0, description=FieldDescriptions.steps)
+    seed: int = InputField(default=0, description="Randomness seed for reproducibility.")
+
+    @torch.no_grad()
+    def invoke(self, context: InvocationContext) -> LatentsOutput:
+        latents = self._run_diffusion(context)
+        latents = latents.detach().to("cpu")
+
+        name = context.tensors.save(tensor=latents)
+        return LatentsOutput.build(latents_name=name, latents=latents, seed=None)
+
+    def _load_text_conditioning(
+        self,
+        context: InvocationContext,
+        conditioning_name: str,
+        joint_attention_dim: int,
+        dtype: torch.dtype,
+        device: torch.device,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        # Load the conditioning data.
+        cond_data = context.conditioning.load(conditioning_name)
+        assert len(cond_data.conditionings) == 1
+        sd3_conditioning = cond_data.conditionings[0]
+        assert isinstance(sd3_conditioning, SD3ConditioningInfo)
+        sd3_conditioning = sd3_conditioning.to(dtype=dtype, device=device)
+
+        t5_embeds = sd3_conditioning.t5_embeds
+        if t5_embeds is None:
+            t5_embeds = torch.zeros(
+                (1, SD3_T5_MAX_SEQ_LEN, joint_attention_dim),
+                device=device,
+                dtype=dtype,
+            )
+
+        clip_prompt_embeds = torch.cat([sd3_conditioning.clip_l_embeds, sd3_conditioning.clip_g_embeds], dim=-1)
+        clip_prompt_embeds = torch.nn.functional.pad(
+            clip_prompt_embeds, (0, t5_embeds.shape[-1] - clip_prompt_embeds.shape[-1])
+        )
+
+        prompt_embeds = torch.cat([clip_prompt_embeds, t5_embeds], dim=-2)
+        pooled_prompt_embeds = torch.cat(
+            [sd3_conditioning.clip_l_pooled_embeds, sd3_conditioning.clip_g_pooled_embeds], dim=-1
+        )
+
+        return prompt_embeds, pooled_prompt_embeds
+
+    def _get_noise(
+        self,
+        num_samples: int,
+        num_channels_latents: int,
+        height: int,
+        width: int,
+        dtype: torch.dtype,
+        device: torch.device,
+        seed: int,
+    ) -> torch.Tensor:
+        # We always generate noise on the same device and dtype then cast to ensure consistency across devices/dtypes.
+        rand_device = "cpu"
+        rand_dtype = torch.float16
+
+        return torch.randn(
+            num_samples,
+            num_channels_latents,
+            int(height) // LATENT_SCALE_FACTOR,
+            int(width) // LATENT_SCALE_FACTOR,
+            device=rand_device,
+            dtype=rand_dtype,
+            generator=torch.Generator(device=rand_device).manual_seed(seed),
+        ).to(device=device, dtype=dtype)
+
+    def _prepare_cfg_scale(self, num_timesteps: int) -> list[float]:
+        """Prepare the CFG scale list.
+
+        Args:
+            num_timesteps (int): The number of timesteps in the scheduler. Could be different from num_steps depending
+            on the scheduler used (e.g. higher order schedulers).
+
+        Returns:
+            list[float]: _description_
+        """
+        if isinstance(self.cfg_scale, float):
+            cfg_scale = [self.cfg_scale] * num_timesteps
+        elif isinstance(self.cfg_scale, list):
+            assert len(self.cfg_scale) == num_timesteps
+            cfg_scale = self.cfg_scale
+        else:
+            raise ValueError(f"Invalid CFG scale type: {type(self.cfg_scale)}")
+
+        return cfg_scale
+
+    def _run_diffusion(
+        self,
+        context: InvocationContext,
+    ):
+        inference_dtype = TorchDevice.choose_torch_dtype()
+        device = TorchDevice.choose_torch_device()
+
+        transformer_info = context.models.load(self.transformer.transformer)
+
+        # Load/process the conditioning data.
+        # TODO(ryand): Make CFG optional.
+        do_classifier_free_guidance = True
+        pos_prompt_embeds, pos_pooled_prompt_embeds = self._load_text_conditioning(
+            context=context,
+            conditioning_name=self.positive_text_conditioning.conditioning_name,
+            joint_attention_dim=transformer_info.model.config.joint_attention_dim,
+            dtype=inference_dtype,
+            device=device,
+        )
+        neg_prompt_embeds, neg_pooled_prompt_embeds = self._load_text_conditioning(
+            context=context,
+            conditioning_name=self.negative_text_conditioning.conditioning_name,
+            joint_attention_dim=transformer_info.model.config.joint_attention_dim,
+            dtype=inference_dtype,
+            device=device,
+        )
+        # TODO(ryand): Support both sequential and batched CFG inference.
+        prompt_embeds = torch.cat([neg_prompt_embeds, pos_prompt_embeds], dim=0)
+        pooled_prompt_embeds = torch.cat([neg_pooled_prompt_embeds, pos_pooled_prompt_embeds], dim=0)
+
+        # Prepare the scheduler.
+        scheduler = FlowMatchEulerDiscreteScheduler()
+        scheduler.set_timesteps(num_inference_steps=self.num_steps, device=device)
+        timesteps = scheduler.timesteps
+        assert isinstance(timesteps, torch.Tensor)
+
+        # Prepare the CFG scale list.
+        cfg_scale = self._prepare_cfg_scale(len(timesteps))
+
+        # Generate initial latent noise.
+        num_channels_latents = transformer_info.model.config.in_channels
+        assert isinstance(num_channels_latents, int)
+        noise = self._get_noise(
+            num_samples=1,
+            num_channels_latents=num_channels_latents,
+            height=self.height,
+            width=self.width,
+            dtype=inference_dtype,
+            device=device,
+            seed=self.seed,
+        )
+        latents: torch.Tensor = noise
+
+        total_steps = len(timesteps)
+        step_callback = self._build_step_callback(context)
+
+        step_callback(
+            PipelineIntermediateState(
+                step=0,
+                order=1,
+                total_steps=total_steps,
+                timestep=int(timesteps[0]),
+                latents=latents,
+            ),
+        )
+
+        with transformer_info.model_on_device() as (cached_weights, transformer):
+            assert isinstance(transformer, SD3Transformer2DModel)
+
+            # 6. Denoising loop
+            for step_idx, t in tqdm(list(enumerate(timesteps))):
+                # Expand the latents if we are doing CFG.
+                latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents
+                # Expand the timestep to match the latent model input.
+                timestep = t.expand(latent_model_input.shape[0])
+
+                noise_pred = transformer(
+                    hidden_states=latent_model_input,
+                    timestep=timestep,
+                    encoder_hidden_states=prompt_embeds,
+                    pooled_projections=pooled_prompt_embeds,
+                    joint_attention_kwargs=None,
+                    return_dict=False,
+                )[0]
+
+                # Apply CFG.
+                if do_classifier_free_guidance:
+                    noise_pred_uncond, noise_pred_cond = noise_pred.chunk(2)
+                    noise_pred = noise_pred_uncond + cfg_scale[step_idx] * (noise_pred_cond - noise_pred_uncond)
+
+                # Compute the previous noisy sample x_t -> x_t-1.
+                latents_dtype = latents.dtype
+                latents = scheduler.step(model_output=noise_pred, timestep=t, sample=latents, return_dict=False)[0]
+
+                # TODO(ryand): This MPS dtype handling was copied from diffusers, I haven't tested to see if it's
+                # needed.
+                if latents.dtype != latents_dtype:
+                    if torch.backends.mps.is_available():
+                        # some platforms (eg. apple mps) misbehave due to a pytorch bug: https://github.com/pytorch/pytorch/pull/99272
+                        latents = latents.to(latents_dtype)
+
+                step_callback(
+                    PipelineIntermediateState(
+                        step=step_idx + 1,
+                        order=1,
+                        total_steps=total_steps,
+                        timestep=int(t),
+                        latents=latents,
+                    ),
+                )
+
+        return latents
+
+    def _build_step_callback(self, context: InvocationContext) -> Callable[[PipelineIntermediateState], None]:
+        def step_callback(state: PipelineIntermediateState) -> None:
+            context.util.sd_step_callback(state, BaseModelType.StableDiffusion3)
+
+        return step_callback

invokeai/app/invocations/sd3_model_loader.py

@@ -0,0 +1,97 @@
+from typing import Optional
+
+from invokeai.app.invocations.baseinvocation import (
+    BaseInvocation,
+    BaseInvocationOutput,
+    Classification,
+    invocation,
+    invocation_output,
+)
+from invokeai.app.invocations.fields import FieldDescriptions, Input, InputField, OutputField, UIType
+from invokeai.app.invocations.model import CLIPField, ModelIdentifierField, T5EncoderField, TransformerField, VAEField
+from invokeai.app.services.shared.invocation_context import InvocationContext
+from invokeai.backend.model_manager.config import SubModelType
+
+
+@invocation_output("sd3_model_loader_output")
+class Sd3ModelLoaderOutput(BaseInvocationOutput):
+    """SD3 base model loader output."""
+
+    mmditx: TransformerField = OutputField(description=FieldDescriptions.mmditx, title="MMDiTX")
+    clip_l: CLIPField = OutputField(description=FieldDescriptions.clip, title="CLIP L")
+    clip_g: CLIPField = OutputField(description=FieldDescriptions.clip, title="CLIP G")
+    t5_encoder: T5EncoderField = OutputField(description=FieldDescriptions.t5_encoder, title="T5 Encoder")
+    vae: VAEField = OutputField(description=FieldDescriptions.vae, title="VAE")
+
+
+@invocation(
+    "sd3_model_loader",
+    title="SD3 Main Model",
+    tags=["model", "sd3"],
+    category="model",
+    version="1.0.0",
+    classification=Classification.Prototype,
+)
+class Sd3ModelLoaderInvocation(BaseInvocation):
+    """Loads a SD3 base model, outputting its submodels."""
+
+    model: ModelIdentifierField = InputField(
+        description=FieldDescriptions.sd3_model,
+        ui_type=UIType.SD3MainModel,
+        input=Input.Direct,
+    )
+
+    t5_encoder_model: Optional[ModelIdentifierField] = InputField(
+        description=FieldDescriptions.t5_encoder,
+        ui_type=UIType.T5EncoderModel,
+        input=Input.Direct,
+        title="T5 Encoder",
+        default=None,
+    )
+
+    # TODO(brandon): Setup UI updates to support selecting a clip l model.
+    # clip_l_model: ModelIdentifierField = InputField(
+    #     description=FieldDescriptions.clip_l_model,
+    #     ui_type=UIType.CLIPEmbedModel,
+    #     input=Input.Direct,
+    #     title="CLIP L Encoder",
+    # )
+
+    # TODO(brandon): Setup UI updates to support selecting a clip g model.
+    # clip_g_model: ModelIdentifierField = InputField(
+    #     description=FieldDescriptions.clip_g_model,
+    #     ui_type=UIType.CLIPGModel,
+    #     input=Input.Direct,
+    #     title="CLIP G Encoder",
+    # )
+
+    # TODO(brandon): Setup UI updates to support selecting an SD3 vae model.
+    # vae_model: ModelIdentifierField = InputField(
+    #     description=FieldDescriptions.vae_model, ui_type=UIType.FluxVAEModel, title="VAE", default=None
+    # )
+
+    def invoke(self, context: InvocationContext) -> Sd3ModelLoaderOutput:
+        mmditx = self.model.model_copy(update={"submodel_type": SubModelType.Transformer})
+        vae = self.model.model_copy(update={"submodel_type": SubModelType.VAE})
+        tokenizer_l = self.model.model_copy(update={"submodel_type": SubModelType.Tokenizer})
+        clip_encoder_l = self.model.model_copy(update={"submodel_type": SubModelType.TextEncoder})
+        tokenizer_g = self.model.model_copy(update={"submodel_type": SubModelType.Tokenizer2})
+        clip_encoder_g = self.model.model_copy(update={"submodel_type": SubModelType.TextEncoder2})
+        tokenizer_t5 = (
+            self.t5_encoder_model.model_copy(update={"submodel_type": SubModelType.Tokenizer3})
+            if self.t5_encoder_model
+            else self.model.model_copy(update={"submodel_type": SubModelType.Tokenizer3})
+        )
+        t5_encoder = (
+            self.t5_encoder_model.model_copy(update={"submodel_type": SubModelType.TextEncoder3})
+            if self.t5_encoder_model
+            else self.model.model_copy(update={"submodel_type": SubModelType.TextEncoder3})
+        )
+
+        return Sd3ModelLoaderOutput(
+            mmditx=TransformerField(transformer=mmditx, loras=[]),
+            clip_l=CLIPField(tokenizer=tokenizer_l, text_encoder=clip_encoder_l, loras=[], skipped_layers=0),
+            clip_g=CLIPField(tokenizer=tokenizer_g, text_encoder=clip_encoder_g, loras=[], skipped_layers=0),
+            t5_encoder=T5EncoderField(tokenizer=tokenizer_t5, text_encoder=t5_encoder),
+            vae=VAEField(vae=vae),
+        )

invokeai/app/invocations/sd3_text_encoder.py

@@ -0,0 +1,199 @@
+from contextlib import ExitStack
+from typing import Iterator, Tuple
+
+import torch
+from transformers import (
+    CLIPTextModel,
+    CLIPTextModelWithProjection,
+    CLIPTokenizer,
+    T5EncoderModel,
+    T5Tokenizer,
+    T5TokenizerFast,
+)
+
+from invokeai.app.invocations.baseinvocation import BaseInvocation, Classification, invocation
+from invokeai.app.invocations.fields import FieldDescriptions, Input, InputField
+from invokeai.app.invocations.model import CLIPField, T5EncoderField
+from invokeai.app.invocations.primitives import SD3ConditioningOutput
+from invokeai.app.services.shared.invocation_context import InvocationContext
+from invokeai.backend.lora.conversions.flux_lora_constants import FLUX_LORA_CLIP_PREFIX
+from invokeai.backend.lora.lora_model_raw import LoRAModelRaw
+from invokeai.backend.lora.lora_patcher import LoRAPatcher
+from invokeai.backend.model_manager.config import ModelFormat
+from invokeai.backend.stable_diffusion.diffusion.conditioning_data import ConditioningFieldData, SD3ConditioningInfo
+
+# The SD3 T5 Max Sequence Length set based on the default in diffusers.
+SD3_T5_MAX_SEQ_LEN = 256
+
+
+@invocation(
+    "sd3_text_encoder",
+    title="SD3 Text Encoding",
+    tags=["prompt", "conditioning", "sd3"],
+    category="conditioning",
+    version="1.0.0",
+    classification=Classification.Prototype,
+)
+class Sd3TextEncoderInvocation(BaseInvocation):
+    """Encodes and preps a prompt for a SD3 image."""
+
+    clip_l: CLIPField = InputField(
+        title="CLIP L",
+        description=FieldDescriptions.clip,
+        input=Input.Connection,
+    )
+    clip_g: CLIPField = InputField(
+        title="CLIP G",
+        description=FieldDescriptions.clip,
+        input=Input.Connection,
+    )
+
+    # The SD3 models were trained with text encoder dropout, so the T5 encoder can be omitted to save time/memory.
+    t5_encoder: T5EncoderField | None = InputField(
+        title="T5Encoder",
+        default=None,
+        description=FieldDescriptions.t5_encoder,
+        input=Input.Connection,
+    )
+    prompt: str = InputField(description="Text prompt to encode.")
+
+    @torch.no_grad()
+    def invoke(self, context: InvocationContext) -> SD3ConditioningOutput:
+        # Note: The text encoding model are run in separate functions to ensure that all model references are locally
+        # scoped. This ensures that earlier models can be freed and gc'd before loading later models (if necessary).
+
+        clip_l_embeddings, clip_l_pooled_embeddings = self._clip_encode(context, self.clip_l)
+        clip_g_embeddings, clip_g_pooled_embeddings = self._clip_encode(context, self.clip_g)
+
+        t5_embeddings: torch.Tensor | None = None
+        if self.t5_encoder is not None:
+            t5_embeddings = self._t5_encode(context, SD3_T5_MAX_SEQ_LEN)
+
+        conditioning_data = ConditioningFieldData(
+            conditionings=[
+                SD3ConditioningInfo(
+                    clip_l_embeds=clip_l_embeddings,
+                    clip_l_pooled_embeds=clip_l_pooled_embeddings,
+                    clip_g_embeds=clip_g_embeddings,
+                    clip_g_pooled_embeds=clip_g_pooled_embeddings,
+                    t5_embeds=t5_embeddings,
+                )
+            ]
+        )
+
+        conditioning_name = context.conditioning.save(conditioning_data)
+        return SD3ConditioningOutput.build(conditioning_name)
+
+    def _t5_encode(self, context: InvocationContext, max_seq_len: int) -> torch.Tensor:
+        assert self.t5_encoder is not None
+        t5_tokenizer_info = context.models.load(self.t5_encoder.tokenizer)
+        t5_text_encoder_info = context.models.load(self.t5_encoder.text_encoder)
+
+        prompt = [self.prompt]
+
+        with (
+            t5_text_encoder_info as t5_text_encoder,
+            t5_tokenizer_info as t5_tokenizer,
+        ):
+            assert isinstance(t5_text_encoder, T5EncoderModel)
+            assert isinstance(t5_tokenizer, (T5Tokenizer, T5TokenizerFast))
+
+            text_inputs = t5_tokenizer(
+                prompt,
+                padding="max_length",
+                max_length=max_seq_len,
+                truncation=True,
+                add_special_tokens=True,
+                return_tensors="pt",
+            )
+            text_input_ids = text_inputs.input_ids
+            untruncated_ids = t5_tokenizer(prompt, padding="longest", return_tensors="pt").input_ids
+            assert isinstance(text_input_ids, torch.Tensor)
+            assert isinstance(untruncated_ids, torch.Tensor)
+            if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not torch.equal(
+                text_input_ids, untruncated_ids
+            ):
+                removed_text = t5_tokenizer.batch_decode(untruncated_ids[:, max_seq_len - 1 : -1])
+                context.logger.warning(
+                    "The following part of your input was truncated because `max_sequence_length` is set to "
+                    f" {max_seq_len} tokens: {removed_text}"
+                )
+
+            prompt_embeds = t5_text_encoder(text_input_ids.to(t5_text_encoder.device))[0]
+
+        assert isinstance(prompt_embeds, torch.Tensor)
+        return prompt_embeds
+
+    def _clip_encode(
+        self, context: InvocationContext, clip_model: CLIPField, tokenizer_max_length: int = 77
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        clip_tokenizer_info = context.models.load(clip_model.tokenizer)
+        clip_text_encoder_info = context.models.load(clip_model.text_encoder)
+
+        prompt = [self.prompt]
+
+        with (
+            clip_text_encoder_info.model_on_device() as (cached_weights, clip_text_encoder),
+            clip_tokenizer_info as clip_tokenizer,
+            ExitStack() as exit_stack,
+        ):
+            assert isinstance(clip_text_encoder, (CLIPTextModel, CLIPTextModelWithProjection))
+            assert isinstance(clip_tokenizer, CLIPTokenizer)
+
+            clip_text_encoder_config = clip_text_encoder_info.config
+            assert clip_text_encoder_config is not None
+
+            # Apply LoRA models to the CLIP encoder.
+            # Note: We apply the LoRA after the transformer has been moved to its target device for faster patching.
+            if clip_text_encoder_config.format in [ModelFormat.Diffusers]:
+                # The model is non-quantized, so we can apply the LoRA weights directly into the model.
+                exit_stack.enter_context(
+                    LoRAPatcher.apply_lora_patches(
+                        model=clip_text_encoder,
+                        patches=self._clip_lora_iterator(context, clip_model),
+                        prefix=FLUX_LORA_CLIP_PREFIX,
+                        cached_weights=cached_weights,
+                    )
+                )
+            else:
+                # There are currently no supported CLIP quantized models. Add support here if needed.
+                raise ValueError(f"Unsupported model format: {clip_text_encoder_config.format}")
+
+            clip_text_encoder = clip_text_encoder.eval().requires_grad_(False)
+
+            text_inputs = clip_tokenizer(
+                prompt,
+                padding="max_length",
+                max_length=tokenizer_max_length,
+                truncation=True,
+                return_tensors="pt",
+            )
+
+            text_input_ids = text_inputs.input_ids
+            untruncated_ids = clip_tokenizer(prompt, padding="longest", return_tensors="pt").input_ids
+            assert isinstance(text_input_ids, torch.Tensor)
+            assert isinstance(untruncated_ids, torch.Tensor)
+            if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not torch.equal(
+                text_input_ids, untruncated_ids
+            ):
+                removed_text = clip_tokenizer.batch_decode(untruncated_ids[:, tokenizer_max_length - 1 : -1])
+                context.logger.warning(
+                    "The following part of your input was truncated because CLIP can only handle sequences up to"
+                    f" {tokenizer_max_length} tokens: {removed_text}"
+                )
+            prompt_embeds = clip_text_encoder(
+                input_ids=text_input_ids.to(clip_text_encoder.device), output_hidden_states=True
+            )
+            pooled_prompt_embeds = prompt_embeds[0]
+            prompt_embeds = prompt_embeds.hidden_states[-2]
+
+            return prompt_embeds, pooled_prompt_embeds
+
+    def _clip_lora_iterator(
+        self, context: InvocationContext, clip_model: CLIPField
+    ) -> Iterator[Tuple[LoRAModelRaw, float]]:
+        for lora in clip_model.loras:
+            lora_info = context.models.load(lora.lora)
+            assert isinstance(lora_info.model, LoRAModelRaw)
+            yield (lora_info.model, lora.weight)
+            del lora_info

invokeai/app/util/step_callback.py

@@ -34,6 +34,25 @@ SD1_5_LATENT_RGB_FACTORS = [
     [-0.1307, -0.1874, -0.7445],  # L4
 ]
 
+SD3_5_LATENT_RGB_FACTORS = [
+    [-0.05240681, 0.03251581, 0.0749016],
+    [-0.0580572, 0.00759826, 0.05729818],
+    [0.16144888, 0.01270368, -0.03768577],
+    [0.14418615, 0.08460266, 0.15941818],
+    [0.04894035, 0.0056485, -0.06686988],
+    [0.05187166, 0.19222395, 0.06261094],
+    [0.1539433, 0.04818359, 0.07103094],
+    [-0.08601796, 0.09013458, 0.10893912],
+    [-0.12398469, -0.06766567, 0.0033688],
+    [-0.0439737, 0.07825329, 0.02258823],
+    [0.03101129, 0.06382551, 0.07753657],
+    [-0.01315361, 0.08554491, -0.08772475],
+    [0.06464487, 0.05914605, 0.13262741],
+    [-0.07863674, -0.02261737, -0.12761454],
+    [-0.09923835, -0.08010759, -0.06264447],
+    [-0.03392309, -0.0804029, -0.06078822],
+]
+
 FLUX_LATENT_RGB_FACTORS = [
     [-0.0412, 0.0149, 0.0521],
     [0.0056, 0.0291, 0.0768],
@@ -110,6 +129,9 @@ def stable_diffusion_step_callback(
         sdxl_latent_rgb_factors = torch.tensor(SDXL_LATENT_RGB_FACTORS, dtype=sample.dtype, device=sample.device)
         sdxl_smooth_matrix = torch.tensor(SDXL_SMOOTH_MATRIX, dtype=sample.dtype, device=sample.device)
         image = sample_to_lowres_estimated_image(sample, sdxl_latent_rgb_factors, sdxl_smooth_matrix)
+    elif base_model == BaseModelType.StableDiffusion3:
+        sd3_latent_rgb_factors = torch.tensor(SD3_5_LATENT_RGB_FACTORS, dtype=sample.dtype, device=sample.device)
+        image = sample_to_lowres_estimated_image(sample, sd3_latent_rgb_factors)
     else:
         v1_5_latent_rgb_factors = torch.tensor(SD1_5_LATENT_RGB_FACTORS, dtype=sample.dtype, device=sample.device)
         image = sample_to_lowres_estimated_image(sample, v1_5_latent_rgb_factors)

invokeai/backend/model_manager/config.py

@@ -53,6 +53,7 @@ class BaseModelType(str, Enum):
     Any = "any"
     StableDiffusion1 = "sd-1"
     StableDiffusion2 = "sd-2"
+    StableDiffusion3 = "sd-3"
     StableDiffusionXL = "sdxl"
     StableDiffusionXLRefiner = "sdxl-refiner"
     Flux = "flux"
@@ -83,8 +84,10 @@ class SubModelType(str, Enum):
     Transformer = "transformer"
     TextEncoder = "text_encoder"
     TextEncoder2 = "text_encoder_2"
+    TextEncoder3 = "text_encoder_3"
     Tokenizer = "tokenizer"
     Tokenizer2 = "tokenizer_2"
+    Tokenizer3 = "tokenizer_3"
     VAE = "vae"
     VAEDecoder = "vae_decoder"
     VAEEncoder = "vae_encoder"
@@ -147,6 +150,11 @@ class ModelSourceType(str, Enum):
 DEFAULTS_PRECISION = Literal["fp16", "fp32"]
 
 
+class SubmodelDefinition(BaseModel):
+    path_or_prefix: str
+    model_type: ModelType
+
+
 class MainModelDefaultSettings(BaseModel):
     vae: str | None = Field(default=None, description="Default VAE for this model (model key)")
     vae_precision: DEFAULTS_PRECISION | None = Field(default=None, description="Default VAE precision for this model")
@@ -193,6 +201,9 @@ class ModelConfigBase(BaseModel):
         schema["required"].extend(["key", "type", "format"])
 
     model_config = ConfigDict(validate_assignment=True, json_schema_extra=json_schema_extra)
+    submodels: Optional[Dict[SubModelType, SubmodelDefinition]] = Field(
+        description="Loadable submodels in this model", default=None
+    )
 
 
 class CheckpointConfigBase(ModelConfigBase):

invokeai/backend/model_manager/load/model_loaders/flux.py

@@ -128,9 +128,9 @@ class BnbQuantizedLlmInt8bCheckpointModel(ModelLoader):
                 "The bnb modules are not available. Please install bitsandbytes if available on your platform."
             )
         match submodel_type:
-            case SubModelType.Tokenizer2:
+            case SubModelType.Tokenizer2 | SubModelType.Tokenizer3:
                 return T5Tokenizer.from_pretrained(Path(config.path) / "tokenizer_2", max_length=512)
-            case SubModelType.TextEncoder2:
+            case SubModelType.TextEncoder2 | SubModelType.TextEncoder3:
                 te2_model_path = Path(config.path) / "text_encoder_2"
                 model_config = AutoConfig.from_pretrained(te2_model_path)
                 with accelerate.init_empty_weights():
@@ -172,9 +172,9 @@ class T5EncoderCheckpointModel(ModelLoader):
             raise ValueError("Only T5EncoderConfig models are currently supported here.")
 
         match submodel_type:
-            case SubModelType.Tokenizer2:
+            case SubModelType.Tokenizer2 | SubModelType.Tokenizer3:
                 return T5Tokenizer.from_pretrained(Path(config.path) / "tokenizer_2", max_length=512)
-            case SubModelType.TextEncoder2:
+            case SubModelType.TextEncoder2 | SubModelType.TextEncoder3:
                 return T5EncoderModel.from_pretrained(Path(config.path) / "text_encoder_2", torch_dtype="auto")
 
         raise ValueError(

invokeai/backend/model_manager/load/model_loaders/stable_diffusion.py

@@ -42,6 +42,7 @@ VARIANT_TO_IN_CHANNEL_MAP = {
 @ModelLoaderRegistry.register(
     base=BaseModelType.StableDiffusionXLRefiner, type=ModelType.Main, format=ModelFormat.Diffusers
 )
+@ModelLoaderRegistry.register(base=BaseModelType.StableDiffusion3, type=ModelType.Main, format=ModelFormat.Diffusers)
 @ModelLoaderRegistry.register(base=BaseModelType.StableDiffusion1, type=ModelType.Main, format=ModelFormat.Checkpoint)
 @ModelLoaderRegistry.register(base=BaseModelType.StableDiffusion2, type=ModelType.Main, format=ModelFormat.Checkpoint)
 @ModelLoaderRegistry.register(base=BaseModelType.StableDiffusionXL, type=ModelType.Main, format=ModelFormat.Checkpoint)
@@ -51,13 +52,6 @@ VARIANT_TO_IN_CHANNEL_MAP = {
 class StableDiffusionDiffusersModel(GenericDiffusersLoader):
     """Class to load main models."""
 
-    model_base_to_model_type = {
-        BaseModelType.StableDiffusion1: "FrozenCLIPEmbedder",
-        BaseModelType.StableDiffusion2: "FrozenOpenCLIPEmbedder",
-        BaseModelType.StableDiffusionXL: "SDXL",
-        BaseModelType.StableDiffusionXLRefiner: "SDXL-Refiner",
-    }
-
     def _load_model(
         self,
         config: AnyModelConfig,

invokeai/backend/model_manager/metadata/fetch/huggingface.py

@@ -20,7 +20,7 @@ from typing import Optional
 
 import requests
 from huggingface_hub import HfApi, configure_http_backend, hf_hub_url
-from huggingface_hub.utils._errors import RepositoryNotFoundError, RevisionNotFoundError
+from huggingface_hub.errors import RepositoryNotFoundError, RevisionNotFoundError
 from pydantic.networks import AnyHttpUrl
 from requests.sessions import Session
 

invokeai/backend/model_manager/probe.py

@@ -19,7 +19,7 @@ from invokeai.backend.lora.conversions.flux_diffusers_lora_conversion_utils impo
     is_state_dict_likely_in_flux_diffusers_format,
 )
 from invokeai.backend.lora.conversions.flux_kohya_lora_conversion_utils import is_state_dict_likely_in_flux_kohya_format
-from invokeai.backend.model_hash.model_hash import HASHING_ALGORITHMS, ModelHash
+from invokeai.backend.model_hash.model_hash import HASHING_ALGORITHMS
 from invokeai.backend.model_manager.config import (
     AnyModelConfig,
     BaseModelType,
@@ -33,7 +33,10 @@ from invokeai.backend.model_manager.config import (
     ModelType,
     ModelVariantType,
     SchedulerPredictionType,
+    SubmodelDefinition,
+    SubModelType,
 )
+from invokeai.backend.model_manager.load.model_loaders.generic_diffusers import ConfigLoader
 from invokeai.backend.model_manager.util.model_util import lora_token_vector_length, read_checkpoint_meta
 from invokeai.backend.quantization.gguf.ggml_tensor import GGMLTensor
 from invokeai.backend.quantization.gguf.loaders import gguf_sd_loader
@@ -112,6 +115,7 @@ class ModelProbe(object):
         "StableDiffusionXLPipeline": ModelType.Main,
         "StableDiffusionXLImg2ImgPipeline": ModelType.Main,
         "StableDiffusionXLInpaintPipeline": ModelType.Main,
+        "StableDiffusion3Pipeline": ModelType.Main,
         "LatentConsistencyModelPipeline": ModelType.Main,
         "AutoencoderKL": ModelType.VAE,
         "AutoencoderTiny": ModelType.VAE,
@@ -122,6 +126,8 @@ class ModelProbe(object):
         "CLIPTextModel": ModelType.CLIPEmbed,
         "T5EncoderModel": ModelType.T5Encoder,
         "FluxControlNetModel": ModelType.ControlNet,
+        "SD3Transformer2DModel": ModelType.Main,
+        "CLIPTextModelWithProjection": ModelType.CLIPEmbed,
     }
 
     @classmethod
@@ -178,7 +184,7 @@ class ModelProbe(object):
             fields.get("description") or f"{fields['base'].value} {model_type.value} model {fields['name']}"
         )
         fields["format"] = ModelFormat(fields.get("format")) if "format" in fields else probe.get_format()
-        fields["hash"] = fields.get("hash") or ModelHash(algorithm=hash_algo).hash(model_path)
+        fields["hash"] = "placeholder"  # fields.get("hash") or ModelHash(algorithm=hash_algo).hash(model_path)
 
         fields["default_settings"] = fields.get("default_settings")
 
@@ -217,6 +223,10 @@ class ModelProbe(object):
                 and fields["prediction_type"] == SchedulerPredictionType.VPrediction
             )
 
+        get_submodels = getattr(probe, "get_submodels", None)
+        if fields["base"] == BaseModelType.StableDiffusion3 and callable(get_submodels):
+            fields["submodels"] = get_submodels()
+
         model_info = ModelConfigFactory.make_config(fields)  # , key=fields.get("key", None))
         return model_info
 
@@ -746,18 +756,33 @@ class FolderProbeBase(ProbeBase):
 
 class PipelineFolderProbe(FolderProbeBase):
     def get_base_type(self) -> BaseModelType:
-        with open(self.model_path / "unet" / "config.json", "r") as file:
-            unet_conf = json.load(file)
-        if unet_conf["cross_attention_dim"] == 768:
-            return BaseModelType.StableDiffusion1
-        elif unet_conf["cross_attention_dim"] == 1024:
-            return BaseModelType.StableDiffusion2
-        elif unet_conf["cross_attention_dim"] == 1280:
-            return BaseModelType.StableDiffusionXLRefiner
-        elif unet_conf["cross_attention_dim"] == 2048:
-            return BaseModelType.StableDiffusionXL
-        else:
-            raise InvalidModelConfigException(f"Unknown base model for {self.model_path}")
+        # Handle pipelines with a UNet (i.e SD 1.x, SD2, SDXL).
+        config_path = self.model_path / "unet" / "config.json"
+        if config_path.exists():
+            with open(config_path) as file:
+                unet_conf = json.load(file)
+            if unet_conf["cross_attention_dim"] == 768:
+                return BaseModelType.StableDiffusion1
+            elif unet_conf["cross_attention_dim"] == 1024:
+                return BaseModelType.StableDiffusion2
+            elif unet_conf["cross_attention_dim"] == 1280:
+                return BaseModelType.StableDiffusionXLRefiner
+            elif unet_conf["cross_attention_dim"] == 2048:
+                return BaseModelType.StableDiffusionXL
+            else:
+                raise InvalidModelConfigException(f"Unknown base model for {self.model_path}")
+
+        # Handle pipelines with a transformer (i.e. SD3).
+        config_path = self.model_path / "transformer" / "config.json"
+        if config_path.exists():
+            with open(config_path) as file:
+                transformer_conf = json.load(file)
+            if transformer_conf["_class_name"] == "SD3Transformer2DModel":
+                return BaseModelType.StableDiffusion3
+            else:
+                raise InvalidModelConfigException(f"Unknown base model for {self.model_path}")
+
+        raise InvalidModelConfigException(f"Unknown base model for {self.model_path}")
 
     def get_scheduler_prediction_type(self) -> SchedulerPredictionType:
         with open(self.model_path / "scheduler" / "scheduler_config.json", "r") as file:
@@ -769,6 +794,21 @@ class PipelineFolderProbe(FolderProbeBase):
         else:
             raise InvalidModelConfigException("Unknown scheduler prediction type: {scheduler_conf['prediction_type']}")
 
+    def get_submodels(self) -> Dict[SubModelType, SubmodelDefinition]:
+        config = ConfigLoader.load_config(self.model_path, config_name="model_index.json")
+        submodels: Dict[SubModelType, SubmodelDefinition] = {}
+        for key, value in config.items():
+            if key.startswith("_") or not (isinstance(value, list) and len(value) == 2):
+                continue
+            model_loader = str(value[1])
+            if model_type := ModelProbe.CLASS2TYPE.get(model_loader):
+                submodels[SubModelType(key)] = SubmodelDefinition(
+                    path_or_prefix=(self.model_path / key).resolve().as_posix(),
+                    model_type=model_type,
+                )
+
+        return submodels
+
     def get_variant_type(self) -> ModelVariantType:
         # This only works for pipelines! Any kind of
         # exception results in our returning the

invokeai/backend/model_manager/util/select_hf_files.py

@@ -129,8 +129,10 @@ def _filter_by_variant(files: List[Path], variant: ModelRepoVariant) -> Set[Path
 
             # Some special handling is needed here if there is not an exact match and if we cannot infer the variant
             # from the file name. In this case, we only give this file a point if the requested variant is FP32 or DEFAULT.
-            if candidate_variant_label == f".{variant}" or (
-                not candidate_variant_label and variant in [ModelRepoVariant.FP32, ModelRepoVariant.Default]
+            if (
+                candidate_variant_label
+                and candidate_variant_label.startswith(f".{variant.value}")
+                or (not candidate_variant_label and variant in [ModelRepoVariant.FP32, ModelRepoVariant.Default])
             ):
                 score += 1
 
@@ -146,7 +148,7 @@ def _filter_by_variant(files: List[Path], variant: ModelRepoVariant) -> Set[Path
         # Check if at least one of the files has the explicit fp16 variant.
         at_least_one_fp16 = False
         for candidate in candidate_list:
-            if len(candidate.path.suffixes) == 2 and candidate.path.suffixes[0] == ".fp16":
+            if len(candidate.path.suffixes) == 2 and candidate.path.suffixes[0].startswith(".fp16"):
                 at_least_one_fp16 = True
                 break
 
@@ -162,7 +164,16 @@ def _filter_by_variant(files: List[Path], variant: ModelRepoVariant) -> Set[Path
         # candidate.
         highest_score_candidate = max(candidate_list, key=lambda candidate: candidate.score)
         if highest_score_candidate:
-            result.add(highest_score_candidate.path)
+            pattern = r"^(.*?)-\d+-of-\d+(\.\w+)$"
+            match = re.match(pattern, highest_score_candidate.path.as_posix())
+            if match:
+                for candidate in candidate_list:
+                    if candidate.path.as_posix().startswith(match.group(1)) and candidate.path.as_posix().endswith(
+                        match.group(2)
+                    ):
+                        result.add(candidate.path)
+            else:
+                result.add(highest_score_candidate.path)
 
     # If one of the architecture-related variants was specified and no files matched other than
     # config and text files then we return an empty list

invokeai/backend/stable_diffusion/diffusion/conditioning_data.py

@@ -49,9 +49,32 @@ class FLUXConditioningInfo:
         return self
 
 
+@dataclass
+class SD3ConditioningInfo:
+    clip_l_pooled_embeds: torch.Tensor
+    clip_l_embeds: torch.Tensor
+    clip_g_pooled_embeds: torch.Tensor
+    clip_g_embeds: torch.Tensor
+    t5_embeds: torch.Tensor | None
+
+    def to(self, device: torch.device | None = None, dtype: torch.dtype | None = None):
+        self.clip_l_pooled_embeds = self.clip_l_pooled_embeds.to(device=device, dtype=dtype)
+        self.clip_l_embeds = self.clip_l_embeds.to(device=device, dtype=dtype)
+        self.clip_g_pooled_embeds = self.clip_g_pooled_embeds.to(device=device, dtype=dtype)
+        self.clip_g_embeds = self.clip_g_embeds.to(device=device, dtype=dtype)
+        if self.t5_embeds is not None:
+            self.t5_embeds = self.t5_embeds.to(device=device, dtype=dtype)
+        return self
+
+
 @dataclass
 class ConditioningFieldData:
-    conditionings: List[BasicConditioningInfo] | List[SDXLConditioningInfo] | List[FLUXConditioningInfo]
+    conditionings: (
+        List[BasicConditioningInfo]
+        | List[SDXLConditioningInfo]
+        | List[FLUXConditioningInfo]
+        | List[SD3ConditioningInfo]
+    )
 
 
 @dataclass

invokeai/backend/util/hotfixes.py

@@ -3,7 +3,7 @@ from typing import Any, Dict, List, Optional, Tuple, Union
 import diffusers
 import torch
 from diffusers.configuration_utils import ConfigMixin, register_to_config
-from diffusers.loaders import FromOriginalControlNetMixin
+from diffusers.loaders.single_file_model import FromOriginalModelMixin
 from diffusers.models.attention_processor import AttentionProcessor, AttnProcessor
 from diffusers.models.controlnet import ControlNetConditioningEmbedding, ControlNetOutput, zero_module
 from diffusers.models.embeddings import (
@@ -32,7 +32,9 @@ from invokeai.backend.util.logging import InvokeAILogger
 logger = InvokeAILogger.get_logger(__name__)
 
 
-class ControlNetModel(ModelMixin, ConfigMixin, FromOriginalControlNetMixin):
+# NOTE(ryand): I'm not the origina author of this code, but for future reference, it appears that this class was copied
+# from diffusers in order to add support for the encoder_attention_mask argument.
+class ControlNetModel(ModelMixin, ConfigMixin, FromOriginalModelMixin):
     """
     A ControlNet model.
 

invokeai/frontend/web/src/features/modelManagerV2/subpanels/ModelManagerPanel/ModelBaseBadge.tsx

@@ -11,6 +11,7 @@ const BASE_COLOR_MAP: Record<BaseModelType, string> = {
   any: 'base',
   'sd-1': 'green',
   'sd-2': 'teal',
+  'sd-3': 'purple',
   sdxl: 'invokeBlue',
   'sdxl-refiner': 'invokeBlue',
   flux: 'gold',

invokeai/frontend/web/src/features/nodes/components/flow/nodes/Invocation/fields/InputFieldRenderer.tsx

@@ -34,6 +34,8 @@ import {
   isModelIdentifierFieldInputTemplate,
   isSchedulerFieldInputInstance,
   isSchedulerFieldInputTemplate,
+  isSD3MainModelFieldInputInstance,
+  isSD3MainModelFieldInputTemplate,
   isSDXLMainModelFieldInputInstance,
   isSDXLMainModelFieldInputTemplate,
   isSDXLRefinerModelFieldInputInstance,
@@ -66,6 +68,7 @@ import MainModelFieldInputComponent from './inputs/MainModelFieldInputComponent'
 import NumberFieldInputComponent from './inputs/NumberFieldInputComponent';
 import RefinerModelFieldInputComponent from './inputs/RefinerModelFieldInputComponent';
 import SchedulerFieldInputComponent from './inputs/SchedulerFieldInputComponent';
+import SD3MainModelFieldInputComponent from './inputs/SD3MainModelFieldInputComponent';
 import SDXLMainModelFieldInputComponent from './inputs/SDXLMainModelFieldInputComponent';
 import SpandrelImageToImageModelFieldInputComponent from './inputs/SpandrelImageToImageModelFieldInputComponent';
 import StringFieldInputComponent from './inputs/StringFieldInputComponent';
@@ -168,10 +171,15 @@ const InputFieldRenderer = ({ nodeId, fieldName }: InputFieldProps) => {
   if (isColorFieldInputInstance(fieldInstance) && isColorFieldInputTemplate(fieldTemplate)) {
     return <ColorFieldInputComponent nodeId={nodeId} field={fieldInstance} fieldTemplate={fieldTemplate} />;
   }
+
   if (isFluxMainModelFieldInputInstance(fieldInstance) && isFluxMainModelFieldInputTemplate(fieldTemplate)) {
     return <FluxMainModelFieldInputComponent nodeId={nodeId} field={fieldInstance} fieldTemplate={fieldTemplate} />;
   }
 
+  if (isSD3MainModelFieldInputInstance(fieldInstance) && isSD3MainModelFieldInputTemplate(fieldTemplate)) {
+    return <SD3MainModelFieldInputComponent nodeId={nodeId} field={fieldInstance} fieldTemplate={fieldTemplate} />;
+  }
+
   if (isSDXLMainModelFieldInputInstance(fieldInstance) && isSDXLMainModelFieldInputTemplate(fieldTemplate)) {
     return <SDXLMainModelFieldInputComponent nodeId={nodeId} field={fieldInstance} fieldTemplate={fieldTemplate} />;
   }

invokeai/frontend/web/src/features/nodes/components/flow/nodes/Invocation/fields/inputs/CLIPEmbedModelFieldInputComponent.tsx

@@ -6,7 +6,7 @@ import type { CLIPEmbedModelFieldInputInstance, CLIPEmbedModelFieldInputTemplate
 import { memo, useCallback } from 'react';
 import { useTranslation } from 'react-i18next';
 import { useCLIPEmbedModels } from 'services/api/hooks/modelsByType';
-import type { CLIPEmbedModelConfig } from 'services/api/types';
+import type { CLIPEmbedModelConfig, MainModelConfig } from 'services/api/types';
 
 import type { FieldComponentProps } from './types';
 
@@ -19,7 +19,7 @@ const CLIPEmbedModelFieldInputComponent = (props: Props) => {
   const dispatch = useAppDispatch();
   const [modelConfigs, { isLoading }] = useCLIPEmbedModels();
   const _onChange = useCallback(
-    (value: CLIPEmbedModelConfig | null) => {
+    (value: CLIPEmbedModelConfig | MainModelConfig | null) => {
       if (!value) {
         return;
       }

invokeai/frontend/web/src/features/nodes/components/flow/nodes/Invocation/fields/inputs/FluxVAEModelFieldInputComponent.tsx

@@ -6,7 +6,7 @@ import type { FluxVAEModelFieldInputInstance, FluxVAEModelFieldInputTemplate } f
 import { memo, useCallback } from 'react';
 import { useTranslation } from 'react-i18next';
 import { useFluxVAEModels } from 'services/api/hooks/modelsByType';
-import type { VAEModelConfig } from 'services/api/types';
+import type { MainModelConfig, VAEModelConfig } from 'services/api/types';
 
 import type { FieldComponentProps } from './types';
 
@@ -19,7 +19,7 @@ const FluxVAEModelFieldInputComponent = (props: Props) => {
   const dispatch = useAppDispatch();
   const [modelConfigs, { isLoading }] = useFluxVAEModels();
   const _onChange = useCallback(
-    (value: VAEModelConfig | null) => {
+    (value: VAEModelConfig | MainModelConfig | null) => {
       if (!value) {
         return;
       }

invokeai/frontend/web/src/features/nodes/components/flow/nodes/Invocation/fields/inputs/SD3MainModelFieldInputComponent.tsx

@@ -0,0 +1,59 @@
+import { Combobox, Flex, FormControl } from '@invoke-ai/ui-library';
+import { useAppDispatch } from 'app/store/storeHooks';
+import { useGroupedModelCombobox } from 'common/hooks/useGroupedModelCombobox';
+import { fieldMainModelValueChanged } from 'features/nodes/store/nodesSlice';
+import type { SD3MainModelFieldInputInstance, SD3MainModelFieldInputTemplate } from 'features/nodes/types/field';
+import { memo, useCallback } from 'react';
+import { useSD3Models } from 'services/api/hooks/modelsByType';
+import type { MainModelConfig } from 'services/api/types';
+
+import type { FieldComponentProps } from './types';
+
+type Props = FieldComponentProps<SD3MainModelFieldInputInstance, SD3MainModelFieldInputTemplate>;
+
+const SD3MainModelFieldInputComponent = (props: Props) => {
+  const { nodeId, field } = props;
+  const dispatch = useAppDispatch();
+  const [modelConfigs, { isLoading }] = useSD3Models();
+  const _onChange = useCallback(
+    (value: MainModelConfig | null) => {
+      if (!value) {
+        return;
+      }
+      dispatch(
+        fieldMainModelValueChanged({
+          nodeId,
+          fieldName: field.name,
+          value,
+        })
+      );
+    },
+    [dispatch, field.name, nodeId]
+  );
+  const { options, value, onChange, placeholder, noOptionsMessage } = useGroupedModelCombobox({
+    modelConfigs,
+    onChange: _onChange,
+    isLoading,
+    selectedModel: field.value,
+  });
+
+  return (
+    <Flex w="full" alignItems="center" gap={2}>
+      <FormControl
+        className="nowheel nodrag"
+        isDisabled={!options.length}
+        isInvalid={!value && props.fieldTemplate.required}
+      >
+        <Combobox
+          value={value}
+          placeholder={placeholder}
+          options={options}
+          onChange={onChange}
+          noOptionsMessage={noOptionsMessage}
+        />
+      </FormControl>
+    </Flex>
+  );
+};
+
+export default memo(SD3MainModelFieldInputComponent);

invokeai/frontend/web/src/features/nodes/components/flow/nodes/Invocation/fields/inputs/T5EncoderModelFieldInputComponent.tsx

@@ -7,7 +7,11 @@ import { selectIsModelsTabDisabled } from 'features/system/store/configSlice';
 import { memo, useCallback } from 'react';
 import { useTranslation } from 'react-i18next';
 import { useT5EncoderModels } from 'services/api/hooks/modelsByType';
-import type { T5EncoderBnbQuantizedLlmInt8bModelConfig, T5EncoderModelConfig } from 'services/api/types';
+import type {
+  MainModelConfig,
+  T5EncoderBnbQuantizedLlmInt8bModelConfig,
+  T5EncoderModelConfig,
+} from 'services/api/types';
 
 import type { FieldComponentProps } from './types';
 
@@ -20,7 +24,7 @@ const T5EncoderModelFieldInputComponent = (props: Props) => {
   const dispatch = useAppDispatch();
   const [modelConfigs, { isLoading }] = useT5EncoderModels();
   const _onChange = useCallback(
-    (value: T5EncoderBnbQuantizedLlmInt8bModelConfig | T5EncoderModelConfig | null) => {
+    (value: T5EncoderBnbQuantizedLlmInt8bModelConfig | T5EncoderModelConfig | MainModelConfig | null) => {
       if (!value) {
         return;
       }
@@ -40,14 +44,14 @@ const T5EncoderModelFieldInputComponent = (props: Props) => {
     isLoading,
     selectedModel: field.value,
   });
-
+  const required = props.fieldTemplate.required;
   return (
     <Flex w="full" alignItems="center" gap={2}>
       <Tooltip label={!isModelsTabDisabled && t('modelManager.starterModelsInModelManager')}>
-        <FormControl className="nowheel nodrag" isDisabled={!options.length} isInvalid={!value}>
+        <FormControl className="nowheel nodrag" isDisabled={!options.length} isInvalid={!value && required}>
           <Combobox
             value={value}
-            placeholder={placeholder}
+            placeholder={required ? placeholder : `(Optional) ${placeholder}`}
             options={options}
             onChange={onChange}
             noOptionsMessage={noOptionsMessage}

invokeai/frontend/web/src/features/nodes/components/flow/nodes/Invocation/fields/inputs/VAEModelFieldInputComponent.tsx

@@ -5,7 +5,7 @@ import { fieldVaeModelValueChanged } from 'features/nodes/store/nodesSlice';
 import type { VAEModelFieldInputInstance, VAEModelFieldInputTemplate } from 'features/nodes/types/field';
 import { memo, useCallback } from 'react';
 import { useVAEModels } from 'services/api/hooks/modelsByType';
-import type { VAEModelConfig } from 'services/api/types';
+import type { MainModelConfig, VAEModelConfig } from 'services/api/types';
 
 import type { FieldComponentProps } from './types';
 
@@ -16,7 +16,7 @@ const VAEModelFieldInputComponent = (props: Props) => {
   const dispatch = useAppDispatch();
   const [modelConfigs, { isLoading }] = useVAEModels();
   const _onChange = useCallback(
-    (value: VAEModelConfig | null) => {
+    (value: VAEModelConfig | MainModelConfig | null) => {
       if (!value) {
         return;
       }

invokeai/frontend/web/src/features/nodes/types/common.ts

@@ -61,8 +61,8 @@ export type SchedulerField = z.infer<typeof zSchedulerField>;
 // #endregion
 
 // #region Model-related schemas
-const zBaseModel = z.enum(['any', 'sd-1', 'sd-2', 'sdxl', 'sdxl-refiner', 'flux']);
-export const zMainModelBase = z.enum(['sd-1', 'sd-2', 'sdxl', 'flux']);
+const zBaseModel = z.enum(['any', 'sd-1', 'sd-2', 'sd-3', 'sdxl', 'sdxl-refiner', 'flux']);
+export const zMainModelBase = z.enum(['sd-1', 'sd-2', 'sd-3', 'sdxl', 'flux']);
 export type MainModelBase = z.infer<typeof zMainModelBase>;
 export const isMainModelBase = (base: unknown): base is MainModelBase => zMainModelBase.safeParse(base).success;
 const zModelType = z.enum([
@@ -84,8 +84,10 @@ const zSubModelType = z.enum([
   'transformer',
   'text_encoder',
   'text_encoder_2',
+  'text_encoder_3',
   'tokenizer',
   'tokenizer_2',
+  'tokenizer_3',
   'vae',
   'vae_decoder',
   'vae_encoder',

invokeai/frontend/web/src/features/nodes/types/constants.ts

@@ -32,6 +32,7 @@ export const MODEL_TYPES = [
   'LoRAModelField',
   'MainModelField',
   'FluxMainModelField',
+  'SD3MainModelField',
   'SDXLMainModelField',
   'SDXLRefinerModelField',
   'VaeModelField',
@@ -65,6 +66,7 @@ export const FIELD_COLORS: { [key: string]: string } = {
   LoRAModelField: 'teal.500',
   MainModelField: 'teal.500',
   FluxMainModelField: 'teal.500',
+  SD3MainModelField: 'teal.500',
   SDXLMainModelField: 'teal.500',
   SDXLRefinerModelField: 'teal.500',
   SpandrelImageToImageModelField: 'teal.500',

invokeai/frontend/web/src/features/nodes/types/field.ts

@@ -115,6 +115,10 @@ const zSDXLMainModelFieldType = zFieldTypeBase.extend({
   name: z.literal('SDXLMainModelField'),
   originalType: zStatelessFieldType.optional(),
 });
+const zSD3MainModelFieldType = zFieldTypeBase.extend({
+  name: z.literal('SD3MainModelField'),
+  originalType: zStatelessFieldType.optional(),
+});
 const zFluxMainModelFieldType = zFieldTypeBase.extend({
   name: z.literal('FluxMainModelField'),
   originalType: zStatelessFieldType.optional(),
@@ -174,6 +178,7 @@ const zStatefulFieldType = z.union([
   zModelIdentifierFieldType,
   zMainModelFieldType,
   zSDXLMainModelFieldType,
+  zSD3MainModelFieldType,
   zFluxMainModelFieldType,
   zSDXLRefinerModelFieldType,
   zVAEModelFieldType,
@@ -467,6 +472,29 @@ export const isSDXLMainModelFieldInputTemplate = (val: unknown): val is SDXLMain
   zSDXLMainModelFieldInputTemplate.safeParse(val).success;
 // #endregion
 
+// #region SD3MainModelField
+
+const zSD3MainModelFieldValue = zMainModelFieldValue; // TODO: Narrow to SDXL models only.
+const zSD3MainModelFieldInputInstance = zFieldInputInstanceBase.extend({
+  value: zSD3MainModelFieldValue,
+});
+const zSD3MainModelFieldInputTemplate = zFieldInputTemplateBase.extend({
+  type: zSD3MainModelFieldType,
+  originalType: zFieldType.optional(),
+  default: zSD3MainModelFieldValue,
+});
+const zSD3MainModelFieldOutputTemplate = zFieldOutputTemplateBase.extend({
+  type: zSD3MainModelFieldType,
+});
+export type SD3MainModelFieldInputInstance = z.infer<typeof zSD3MainModelFieldInputInstance>;
+export type SD3MainModelFieldInputTemplate = z.infer<typeof zSD3MainModelFieldInputTemplate>;
+export const isSD3MainModelFieldInputInstance = (val: unknown): val is SD3MainModelFieldInputInstance =>
+  zSD3MainModelFieldInputInstance.safeParse(val).success;
+export const isSD3MainModelFieldInputTemplate = (val: unknown): val is SD3MainModelFieldInputTemplate =>
+  zSD3MainModelFieldInputTemplate.safeParse(val).success;
+
+// #endregion
+
 // #region FluxMainModelField
 
 const zFluxMainModelFieldValue = zMainModelFieldValue; // TODO: Narrow to SDXL models only.
@@ -806,6 +834,7 @@ export const zStatefulFieldValue = z.union([
   zMainModelFieldValue,
   zSDXLMainModelFieldValue,
   zFluxMainModelFieldValue,
+  zSD3MainModelFieldValue,
   zSDXLRefinerModelFieldValue,
   zVAEModelFieldValue,
   zLoRAModelFieldValue,
@@ -837,6 +866,7 @@ const zStatefulFieldInputInstance = z.union([
   zModelIdentifierFieldInputInstance,
   zMainModelFieldInputInstance,
   zFluxMainModelFieldInputInstance,
+  zSD3MainModelFieldInputInstance,
   zSDXLMainModelFieldInputInstance,
   zSDXLRefinerModelFieldInputInstance,
   zVAEModelFieldInputInstance,
@@ -870,6 +900,7 @@ const zStatefulFieldInputTemplate = z.union([
   zModelIdentifierFieldInputTemplate,
   zMainModelFieldInputTemplate,
   zFluxMainModelFieldInputTemplate,
+  zSD3MainModelFieldInputTemplate,
   zSDXLMainModelFieldInputTemplate,
   zSDXLRefinerModelFieldInputTemplate,
   zVAEModelFieldInputTemplate,
@@ -904,6 +935,7 @@ const zStatefulFieldOutputTemplate = z.union([
   zModelIdentifierFieldOutputTemplate,
   zMainModelFieldOutputTemplate,
   zFluxMainModelFieldOutputTemplate,
+  zSD3MainModelFieldOutputTemplate,
   zSDXLMainModelFieldOutputTemplate,
   zSDXLRefinerModelFieldOutputTemplate,
   zVAEModelFieldOutputTemplate,

invokeai/frontend/web/src/features/nodes/util/schema/buildFieldInputInstance.ts

@@ -16,6 +16,7 @@ const FIELD_VALUE_FALLBACK_MAP: Record<StatefulFieldType['name'], FieldValue> =
   SchedulerField: 'dpmpp_3m_k',
   SDXLMainModelField: undefined,
   FluxMainModelField: undefined,
+  SD3MainModelField: undefined,
   SDXLRefinerModelField: undefined,
   StringField: '',
   T2IAdapterModelField: undefined,

invokeai/frontend/web/src/features/nodes/util/schema/buildFieldInputTemplate.ts

@@ -18,6 +18,7 @@ import type {
   MainModelFieldInputTemplate,
   ModelIdentifierFieldInputTemplate,
   SchedulerFieldInputTemplate,
+  SD3MainModelFieldInputTemplate,
   SDXLMainModelFieldInputTemplate,
   SDXLRefinerModelFieldInputTemplate,
   SpandrelImageToImageModelFieldInputTemplate,
@@ -198,6 +199,20 @@ const buildFluxMainModelFieldInputTemplate: FieldInputTemplateBuilder<FluxMainMo
   return template;
 };
 
+const buildSD3MainModelFieldInputTemplate: FieldInputTemplateBuilder<SD3MainModelFieldInputTemplate> = ({
+  schemaObject,
+  baseField,
+  fieldType,
+}) => {
+  const template: SD3MainModelFieldInputTemplate = {
+    ...baseField,
+    type: fieldType,
+    default: schemaObject.default ?? undefined,
+  };
+
+  return template;
+};
+
 const buildRefinerModelFieldInputTemplate: FieldInputTemplateBuilder<SDXLRefinerModelFieldInputTemplate> = ({
   schemaObject,
   baseField,
@@ -446,6 +461,7 @@ export const TEMPLATE_BUILDER_MAP: Record<StatefulFieldType['name'], FieldInputT
   MainModelField: buildMainModelFieldInputTemplate,
   SchedulerField: buildSchedulerFieldInputTemplate,
   SDXLMainModelField: buildSDXLMainModelFieldInputTemplate,
+  SD3MainModelField: buildSD3MainModelFieldInputTemplate,
   FluxMainModelField: buildFluxMainModelFieldInputTemplate,
   SDXLRefinerModelField: buildRefinerModelFieldInputTemplate,
   StringField: buildStringFieldInputTemplate,

invokeai/frontend/web/src/features/nodes/util/workflow/validateWorkflow.ts

@@ -30,6 +30,7 @@ const MODEL_FIELD_TYPES = [
   'MainModelField',
   'SDXLMainModelField',
   'FluxMainModelField',
+  'SD3MainModelField',
   'SDXLRefinerModelField',
   'VAEModelField',
   'LoRAModelField',

invokeai/frontend/web/src/features/parameters/components/Advanced/ParamCLIPEmbedModelSelect.tsx

@@ -6,7 +6,7 @@ import { zModelIdentifierField } from 'features/nodes/types/common';
 import { memo, useCallback } from 'react';
 import { useTranslation } from 'react-i18next';
 import { useCLIPEmbedModels } from 'services/api/hooks/modelsByType';
-import type { CLIPEmbedModelConfig } from 'services/api/types';
+import type { CLIPEmbedModelConfig, MainModelConfig } from 'services/api/types';
 
 const ParamCLIPEmbedModelSelect = () => {
   const dispatch = useAppDispatch();
@@ -15,7 +15,7 @@ const ParamCLIPEmbedModelSelect = () => {
   const [modelConfigs, { isLoading }] = useCLIPEmbedModels();
 
   const _onChange = useCallback(
-    (clipEmbedModel: CLIPEmbedModelConfig | null) => {
+    (clipEmbedModel: CLIPEmbedModelConfig | MainModelConfig | null) => {
       if (clipEmbedModel) {
         dispatch(clipEmbedModelSelected(zModelIdentifierField.parse(clipEmbedModel)));
       }

invokeai/frontend/web/src/features/parameters/components/Advanced/ParamT5EncoderModelSelect.tsx

@@ -6,7 +6,11 @@ import { zModelIdentifierField } from 'features/nodes/types/common';
 import { memo, useCallback } from 'react';
 import { useTranslation } from 'react-i18next';
 import { useT5EncoderModels } from 'services/api/hooks/modelsByType';
-import type { T5EncoderBnbQuantizedLlmInt8bModelConfig, T5EncoderModelConfig } from 'services/api/types';
+import type {
+  MainModelConfig,
+  T5EncoderBnbQuantizedLlmInt8bModelConfig,
+  T5EncoderModelConfig,
+} from 'services/api/types';
 
 const ParamT5EncoderModelSelect = () => {
   const dispatch = useAppDispatch();
@@ -15,7 +19,7 @@ const ParamT5EncoderModelSelect = () => {
   const [modelConfigs, { isLoading }] = useT5EncoderModels();
 
   const _onChange = useCallback(
-    (t5EncoderModel: T5EncoderBnbQuantizedLlmInt8bModelConfig | T5EncoderModelConfig | null) => {
+    (t5EncoderModel: T5EncoderBnbQuantizedLlmInt8bModelConfig | T5EncoderModelConfig | MainModelConfig | null) => {
       if (t5EncoderModel) {
         dispatch(t5EncoderModelSelected(zModelIdentifierField.parse(t5EncoderModel)));
       }

invokeai/frontend/web/src/features/parameters/components/VAEModel/ParamFLUXVAEModelSelect.tsx

@@ -7,7 +7,7 @@ import { zModelIdentifierField } from 'features/nodes/types/common';
 import { memo, useCallback } from 'react';
 import { useTranslation } from 'react-i18next';
 import { useFluxVAEModels } from 'services/api/hooks/modelsByType';
-import type { VAEModelConfig } from 'services/api/types';
+import type { MainModelConfig, VAEModelConfig } from 'services/api/types';
 
 const ParamFLUXVAEModelSelect = () => {
   const dispatch = useAppDispatch();
@@ -16,7 +16,7 @@ const ParamFLUXVAEModelSelect = () => {
   const [modelConfigs, { isLoading }] = useFluxVAEModels();
 
   const _onChange = useCallback(
-    (vae: VAEModelConfig | null) => {
+    (vae: VAEModelConfig | MainModelConfig | null) => {
       if (vae) {
         dispatch(fluxVAESelected(zModelIdentifierField.parse(vae)));
       }

invokeai/frontend/web/src/features/parameters/components/VAEModel/ParamVAEModelSelect.tsx

@@ -7,7 +7,7 @@ import { zModelIdentifierField } from 'features/nodes/types/common';
 import { memo, useCallback } from 'react';
 import { useTranslation } from 'react-i18next';
 import { useVAEModels } from 'services/api/hooks/modelsByType';
-import type { VAEModelConfig } from 'services/api/types';
+import type { MainModelConfig, VAEModelConfig } from 'services/api/types';
 
 const ParamVAEModelSelect = () => {
   const dispatch = useAppDispatch();
@@ -16,7 +16,7 @@ const ParamVAEModelSelect = () => {
   const vae = useAppSelector(selectVAE);
   const [modelConfigs, { isLoading }] = useVAEModels();
   const getIsDisabled = useCallback(
-    (vae: VAEModelConfig): boolean => {
+    (vae: VAEModelConfig | MainModelConfig): boolean => {
       const isCompatible = base === vae.base;
       const hasMainModel = Boolean(base);
       return !hasMainModel || !isCompatible;
@@ -24,7 +24,7 @@ const ParamVAEModelSelect = () => {
     [base]
   );
   const _onChange = useCallback(
-    (vae: VAEModelConfig | null) => {
+    (vae: VAEModelConfig | MainModelConfig | null) => {
       dispatch(vaeSelected(vae ? zModelIdentifierField.parse(vae) : null));
     },
     [dispatch]

invokeai/frontend/web/src/features/parameters/types/constants.ts

@@ -19,6 +19,7 @@ export const MODEL_TYPE_SHORT_MAP = {
   any: 'Any',
   'sd-1': 'SD1.X',
   'sd-2': 'SD2.X',
+  'sd-3': 'SD3.X',
   sdxl: 'SDXL',
   'sdxl-refiner': 'SDXLR',
   flux: 'FLUX',
@@ -40,6 +41,10 @@ export const CLIP_SKIP_MAP = {
     maxClip: 24,
     markers: [0, 1, 2, 3, 5, 10, 15, 20, 24],
   },
+  'sd-3': {
+    maxClip: 0,
+    markers: [],
+  },
   sdxl: {
     maxClip: 24,
     markers: [0, 1, 2, 3, 5, 10, 15, 20, 24],

invokeai/frontend/web/src/services/api/hooks/modelsByType.ts

@@ -20,6 +20,7 @@ import {
   isNonRefinerMainModelConfig,
   isNonSDXLMainModelConfig,
   isRefinerMainModelModelConfig,
+  isSD3MainModelModelConfig,
   isSDXLMainModelModelConfig,
   isSpandrelImageToImageModelConfig,
   isT2IAdapterModelConfig,
@@ -47,6 +48,7 @@ export const useMainModels = buildModelsHook(isNonRefinerMainModelConfig);
 export const useNonSDXLMainModels = buildModelsHook(isNonSDXLMainModelConfig);
 export const useRefinerModels = buildModelsHook(isRefinerMainModelModelConfig);
 export const useFluxModels = buildModelsHook(isFluxMainModelModelConfig);
+export const useSD3Models = buildModelsHook(isSD3MainModelModelConfig);
 export const useSDXLModels = buildModelsHook(isSDXLMainModelModelConfig);
 export const useLoRAModels = buildModelsHook(isLoRAModelConfig);
 export const useControlNetAndT2IAdapterModels = buildModelsHook(isControlNetOrT2IAdapterModelConfig);

invokeai/frontend/web/src/services/api/types.ts

@@ -75,20 +75,38 @@ export type AnyModelConfig =
   | MainModelConfig
   | CLIPVisionDiffusersConfig;
 
+const check_submodel_model_type = (submodels: AnyModelConfig['submodels'], model_type: string): boolean => {
+  for (const submodel in submodels) {
+    if (submodel && submodels[submodel] && submodels[submodel].model_type === model_type) {
+      return true;
+    }
+  }
+  return false;
+};
+
+const check_submodels = (indentifier: string, config: AnyModelConfig): boolean => {
+  return (
+    (config.type === 'main' &&
+      config.submodels &&
+      (indentifier in config.submodels || check_submodel_model_type(config.submodels, indentifier))) ||
+    false
+  );
+};
+
 export const isLoRAModelConfig = (config: AnyModelConfig): config is LoRAModelConfig => {
   return config.type === 'lora';
 };
 
-export const isVAEModelConfig = (config: AnyModelConfig): config is VAEModelConfig => {
-  return config.type === 'vae';
+export const isVAEModelConfig = (config: AnyModelConfig): config is VAEModelConfig | MainModelConfig => {
+  return config.type === 'vae' || check_submodels('vae', config);
 };
 
-export const isNonFluxVAEModelConfig = (config: AnyModelConfig): config is VAEModelConfig => {
-  return config.type === 'vae' && config.base !== 'flux';
+export const isNonFluxVAEModelConfig = (config: AnyModelConfig): config is VAEModelConfig | MainModelConfig => {
+  return (config.type === 'vae' || check_submodels('vae', config)) && config.base !== 'flux';
 };
 
-export const isFluxVAEModelConfig = (config: AnyModelConfig): config is VAEModelConfig => {
-  return config.type === 'vae' && config.base === 'flux';
+export const isFluxVAEModelConfig = (config: AnyModelConfig): config is VAEModelConfig | MainModelConfig => {
+  return (config.type === 'vae' || check_submodels('vae', config)) && config.base === 'flux';
 };
 
 export const isControlNetModelConfig = (config: AnyModelConfig): config is ControlNetModelConfig => {
@@ -109,12 +127,12 @@ export const isT2IAdapterModelConfig = (config: AnyModelConfig): config is T2IAd
 
 export const isT5EncoderModelConfig = (
   config: AnyModelConfig
-): config is T5EncoderModelConfig | T5EncoderBnbQuantizedLlmInt8bModelConfig => {
-  return config.type === 't5_encoder';
+): config is T5EncoderModelConfig | T5EncoderBnbQuantizedLlmInt8bModelConfig | MainModelConfig => {
+  return config.type === 't5_encoder' || check_submodels('t5_encoder', config);
 };
 
-export const isCLIPEmbedModelConfig = (config: AnyModelConfig): config is CLIPEmbedModelConfig => {
-  return config.type === 'clip_embed';
+export const isCLIPEmbedModelConfig = (config: AnyModelConfig): config is CLIPEmbedModelConfig | MainModelConfig => {
+  return config.type === 'clip_embed' || check_submodels('clip_embed', config);
 };
 
 export const isSpandrelImageToImageModelConfig = (
@@ -145,6 +163,10 @@ export const isSDXLMainModelModelConfig = (config: AnyModelConfig): config is Ma
   return config.type === 'main' && config.base === 'sdxl';
 };
 
+export const isSD3MainModelModelConfig = (config: AnyModelConfig): config is MainModelConfig => {
+  return config.type === 'main' && config.base === 'sd-3';
+};
+
 export const isFluxMainModelModelConfig = (config: AnyModelConfig): config is MainModelConfig => {
   return config.type === 'main' && config.base === 'flux';
 };

pyproject.toml

@@ -33,12 +33,12 @@ classifiers = [
 ]
 dependencies = [
   # Core generation dependencies, pinned for reproducible builds.
-  "accelerate==0.30.1",
+  "accelerate==1.0.1",
   "bitsandbytes==0.43.3; sys_platform!='darwin'",
   "clip_anytorch==2.6.0",       # replacing "clip @ https://github.com/openai/CLIP/archive/eaa22acb90a5876642d0507623e859909230a52d.zip",
   "compel==2.0.2",
   "controlnet-aux==0.0.7",
-  "diffusers[torch]==0.27.2",
+  "diffusers[torch]==0.31.0",
   "gguf==0.10.0",
   "invisible-watermark==0.2.0", # needed to install SDXL base and refiner using their repo_ids
   "mediapipe>=0.10.7",          # needed for "mediapipeface" controlnet model
@@ -61,7 +61,7 @@ dependencies = [
   # Core application dependencies, pinned for reproducible builds.
   "fastapi-events==0.11.1",
   "fastapi==0.111.0",
-  "huggingface-hub==0.23.1",
+  "huggingface-hub==0.26.1",
   "pydantic-settings==2.2.1",
   "pydantic==2.7.2",
   "python-socketio==5.11.1",