Merge pull request microsoft#393 from keshavv27/rel/v0.4.2

[StableDiffusion Sample][NvTensorRTRTX EP] Add free dimension overrides for SD1.4 components

Author: nmetulev

AIDevGallery/Samples/SharedCode/StableDiffusionCode/SafetyChecker.cs

@@ -58,6 +58,11 @@ private Task<InferenceSession> GetInferenceSession(string modelPath, ExecutionPr
             SessionOptions sessionOptions = new();
             sessionOptions.RegisterOrtExtensions();
 
+            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 +93,7 @@ public bool IsNotSafe(Tensor<float> resultImage, StableDiffusionConfig config)
         var inputTensor = ClipImageFeatureExtractor(resultImage, config);
 
         // images input
-        var inputImagesTensor = ReorderTensor(inputTensor);
+        var inputImagesTensor = inputTensor;
 
         var input = new List<NamedOnnxValue>
         {

AIDevGallery/Samples/SharedCode/StableDiffusionCode/StableDiffusion.cs

@@ -83,6 +83,12 @@ private Task<InferenceSession> GetInferenceSession(string modelPath, ExecutionPr
             SessionOptions sessionOptions = new();
             sessionOptions.RegisterOrtExtensions();
 
+            sessionOptions.AddFreeDimensionOverrideByName("batch", 2);
+            sessionOptions.AddFreeDimensionOverrideByName("channels", 4);
+            sessionOptions.AddFreeDimensionOverrideByName("height", config.Height / 8);
+            sessionOptions.AddFreeDimensionOverrideByName("width", config.Width / 8);
+            sessionOptions.AddFreeDimensionOverrideByName("sequence", 77);
+
             if (policy != null)
             {
                 sessionOptions.SetEpSelectionPolicy(policy.Value);
@@ -108,7 +114,7 @@ public static List<NamedOnnxValue> CreateUnetModelInput(Tensor<float> encoderHid
         {
             NamedOnnxValue.CreateFromTensor("encoder_hidden_states", encoderHiddenStates),
             NamedOnnxValue.CreateFromTensor("sample", sample),
-            NamedOnnxValue.CreateFromTensor("timestep", new DenseTensor<long>(new long[] { timeStep }, [ 1 ]))
+            NamedOnnxValue.CreateFromTensor("timestep", new DenseTensor<long>(new long[] { timeStep, timeStep }, [ 2 ]))
         };
 
         return input;

AIDevGallery/Samples/SharedCode/StableDiffusionCode/TextProcessing.cs

@@ -14,6 +14,15 @@ namespace AIDevGallery.Samples.SharedCode.StableDiffusionCode;
 
 internal class TextProcessing : IDisposable
 {
+    private readonly StableDiffusionConfig config = new()
+    {
+        // Number of denoising steps
+        NumInferenceSteps = 15,
+
+        // Scale for classifier-free guidance
+        GuidanceScale = 7.5
+    };
+
     private InferenceSession? tokenizerInferenceSession;
     private InferenceSession? encoderInferenceSession;
     private bool disposedValue;
@@ -60,6 +69,11 @@ private Task<InferenceSession> GetInferenceSession(string modelPath, ExecutionPr
             SessionOptions sessionOptions = new();
             sessionOptions.RegisterOrtExtensions();
 
+            sessionOptions.AddFreeDimensionOverrideByName("batch", 1);
+            sessionOptions.AddFreeDimensionOverrideByName("channels", 3);
+            sessionOptions.AddFreeDimensionOverrideByName("height", config.Height);
+            sessionOptions.AddFreeDimensionOverrideByName("width", config.Width);
+
             if (policy != null)
             {
                 sessionOptions.SetEpSelectionPolicy(policy.Value);

AIDevGallery/Samples/SharedCode/StableDiffusionCode/VaeDecoder.cs

@@ -14,6 +14,15 @@ namespace AIDevGallery.Samples.SharedCode.StableDiffusionCode;
 
 internal class VaeDecoder : IDisposable
 {
+    private readonly StableDiffusionConfig config = new()
+    {
+        // Number of denoising steps
+        NumInferenceSteps = 15,
+
+        // Scale for classifier-free guidance
+        GuidanceScale = 7.5
+    };
+
     private InferenceSession? vaeDecoderInferenceSession;
     private bool disposedValue;
 
@@ -57,6 +66,11 @@ private Task<InferenceSession> GetInferenceSession(string modelPath, ExecutionPr
             SessionOptions sessionOptions = new();
             sessionOptions.RegisterOrtExtensions();
 
+            sessionOptions.AddFreeDimensionOverrideByName("batch", 1);
+            sessionOptions.AddFreeDimensionOverrideByName("channels", 4);
+            sessionOptions.AddFreeDimensionOverrideByName("height", config.Height / 8);
+            sessionOptions.AddFreeDimensionOverrideByName("width", config.Width / 8);
+
             if (policy != null)
             {
                 sessionOptions.SetEpSelectionPolicy(policy.Value);