ONNX Support

installer/lib/installer.py

@@ -451,7 +451,7 @@ def get_torch_source() -> (Union[str, None], str):
     device = graphical_accelerator()
 
     url = None
-    optional_modules = None
+    optional_modules = "[onnx]"
     if OS == "Linux":
         if device == "rocm":
             url = "https://download.pytorch.org/whl/rocm5.4.2"
@@ -460,7 +460,10 @@ def get_torch_source() -> (Union[str, None], str):
 
     if device == "cuda":
         url = "https://download.pytorch.org/whl/cu117"
-        optional_modules = "[xformers]"
+        optional_modules = "[xformers,onnx-cuda]"
+    if device == "cuda_and_dml":
+        url = "https://download.pytorch.org/whl/cu117"
+        optional_modules = "[xformers,onnx-directml]"
 
     # in all other cases, Torch wheels should be coming from PyPi as of Torch 1.13
 

installer/lib/messages.py

@@ -167,6 +167,10 @@ def graphical_accelerator():
         "an [gold1 b]NVIDIA[/] GPU (using CUDA™)",
         "cuda",
     )
+    nvidia_with_dml = (
+        "an [gold1 b]NVIDIA[/] GPU (using CUDA™, and DirectML™ for ONNX)",
+        "cuda_and_dml",
+    )
     amd = (
         "an [gold1 b]AMD[/] GPU (using ROCm™)",
         "rocm",
@@ -181,7 +185,7 @@ def graphical_accelerator():
     )
 
     if OS == "Windows":
-        options = [nvidia, cpu]
+        options = [nvidia, nvidia_with_dml, cpu]
     if OS == "Linux":
         options = [nvidia, amd, cpu]
     elif OS == "Darwin":

invokeai/app/invocations/compel.py

@@ -1,6 +1,14 @@
 from typing import Literal, Optional, Union, List, Annotated
 from pydantic import BaseModel, Field
 import re
+
+from .baseinvocation import BaseInvocation, BaseInvocationOutput, InvocationContext, InvocationConfig
+from .model import ClipField
+
+from ...backend.util.devices import torch_dtype
+from ...backend.stable_diffusion.diffusion import InvokeAIDiffuserComponent
+from ...backend.model_management import BaseModelType, ModelType, SubModelType, ModelPatcher
+
 import torch
 from compel import Compel, ReturnedEmbeddingsType
 from compel.prompt_parser import Blend, Conjunction, CrossAttentionControlSubstitute, FlattenedPrompt, Fragment

invokeai/app/invocations/latent.py

@@ -24,6 +24,7 @@
 )
 from ...backend.stable_diffusion.diffusion.shared_invokeai_diffusion import PostprocessingSettings
 from ...backend.stable_diffusion.schedulers import SCHEDULER_MAP
+from ...backend.model_management import ModelPatcher
 from ...backend.util.devices import choose_torch_device, torch_dtype, choose_precision
 from ..models.image import ImageCategory, ImageField, ResourceOrigin
 from .baseinvocation import BaseInvocation, BaseInvocationOutput, InvocationConfig, InvocationContext

invokeai/app/invocations/model.py

@@ -53,6 +53,7 @@ class MainModelField(BaseModel):
 
     model_name: str = Field(description="Name of the model")
     base_model: BaseModelType = Field(description="Base model")
+    model_type: ModelType = Field(description="Model Type")
 
 
 class LoRAModelField(BaseModel):