Add free dimension overrides for SD 1.4 components

Author: keshavv27

AIDevGallery/Samples/SharedCode/StableDiffusionCode/SafetyChecker.cs

@@ -58,6 +58,14 @@ private Task<InferenceSession> GetInferenceSession(string modelPath, ExecutionPr
             SessionOptions sessionOptions = new();
             sessionOptions.RegisterOrtExtensions();
 
+            if(device == "NvTensorRTRTXExecutionProvider")
+            {
+                sessionOptions.AddFreeDimensionOverrideByName("batch", 1);
+                sessionOptions.AddFreeDimensionOverrideByName("channels", 3);
+                sessionOptions.AddFreeDimensionOverrideByName("height", 224);
+                sessionOptions.AddFreeDimensionOverrideByName("width", 224);
+            }
+
             if (policy != null)
             {
                 sessionOptions.SetEpSelectionPolicy(policy.Value);
@@ -88,7 +96,14 @@ public bool IsNotSafe(Tensor<float> resultImage, StableDiffusionConfig config)
         var inputTensor = ClipImageFeatureExtractor(resultImage, config);
 
         // images input
-        var inputImagesTensor = ReorderTensor(inputTensor);
+        if(device == "NvTensorRTRTXExecutionProvider")
+        {
+            var inputImagesTensor = inputTensor;
+        }
+        else
+        {
+            var inputImagesTensor = ReorderTensor(inputTensor);
+        }
 
         var input = new List<NamedOnnxValue>
         {

AIDevGallery/Samples/SharedCode/StableDiffusionCode/StableDiffusion.cs

@@ -83,6 +83,17 @@ private Task<InferenceSession> GetInferenceSession(string modelPath, ExecutionPr
             SessionOptions sessionOptions = new();
             sessionOptions.RegisterOrtExtensions();
 
+            if(device == "NvTensorRTRTXExecutionProvider")
+            {
+                sessionOptions.AddFreeDimensionOverrideByName("batch", 2);
+                sessionOptions.AddFreeDimensionOverrideByName("time_batch", 1);
+                sessionOptions.AddFreeDimensionOverrideByName("channels", 4);
+                sessionOptions.AddFreeDimensionOverrideByName("height", 64);
+                sessionOptions.AddFreeDimensionOverrideByName("width", 64);
+                sessionOptions.AddFreeDimensionOverrideByName("sequence", 77);
+            }
+            
+
             if (policy != null)
             {
                 sessionOptions.SetEpSelectionPolicy(policy.Value);

AIDevGallery/Samples/SharedCode/StableDiffusionCode/TextProcessing.cs

@@ -60,6 +60,14 @@ private Task<InferenceSession> GetInferenceSession(string modelPath, ExecutionPr
             SessionOptions sessionOptions = new();
             sessionOptions.RegisterOrtExtensions();
 
+            if(device == "NvTensorRTRTXExecutionProvider")
+            {
+                sessionOptions.AddFreeDimensionOverrideByName("batch", 1);
+                sessionOptions.AddFreeDimensionOverrideByName("channels", 4);
+                sessionOptions.AddFreeDimensionOverrideByName("height", 512);
+                sessionOptions.AddFreeDimensionOverrideByName("width", 512);
+            }
+
             if (policy != null)
             {
                 sessionOptions.SetEpSelectionPolicy(policy.Value);

AIDevGallery/Samples/SharedCode/StableDiffusionCode/VaeDecoder.cs

@@ -57,6 +57,14 @@ private Task<InferenceSession> GetInferenceSession(string modelPath, ExecutionPr
             SessionOptions sessionOptions = new();
             sessionOptions.RegisterOrtExtensions();
 
+            if(device == "NvTensorRTRTXExecutionProvider")
+            {
+                sessionOptions.AddFreeDimensionOverrideByName("batch", 1);
+                sessionOptions.AddFreeDimensionOverrideByName("channels", 4);
+                sessionOptions.AddFreeDimensionOverrideByName("height", 64);
+                sessionOptions.AddFreeDimensionOverrideByName("width", 64);
+            }            
+
             if (policy != null)
             {
                 sessionOptions.SetEpSelectionPolicy(policy.Value);