Support models with Float timestep tensor input

saddam213 · Nov 9, 2023 · 38f60b5 · 38f60b5
1 parent e56a7b6
commit 38f60b5
Show file tree

Hide file tree

Showing 2 changed files with 20 additions and 2 deletions.
diff --git a/OnnxStack.StableDiffusion/Diffusers/LatentConsistency/LatentConsistencyDiffuser.cs b/OnnxStack.StableDiffusion/Diffusers/LatentConsistency/LatentConsistencyDiffuser.cs
@@ -193,9 +193,18 @@ protected virtual async Task<DenseTensor<float>> DecodeLatents(IModelOptions mod
         protected virtual IReadOnlyList<NamedOnnxValue> CreateUnetInputParams(IModelOptions model, DenseTensor<float> inputTensor, DenseTensor<float> promptEmbeddings, DenseTensor<float> guidanceEmbeddings, int timestep)
         {
             var inputNames = _onnxModelService.GetInputNames(model, OnnxModelType.Unet);
+            var inputMetaData = _onnxModelService.GetInputMetadata(model, OnnxModelType.Unet);
+
+            // Some models support Long or Float, could be more but fornow just support these 2
+            var timesepMetaKey = inputNames[1];
+            var timestepMetaData = inputMetaData[timesepMetaKey];
+            var timestepNamedOnnxValue = timestepMetaData.ElementDataType == TensorElementType.Int64
+                ? NamedOnnxValue.CreateFromTensor(timesepMetaKey, new DenseTensor<long>(new long[] { timestep }, new int[] { 1 }))
+                : NamedOnnxValue.CreateFromTensor(timesepMetaKey, new DenseTensor<float>(new float[] { timestep }, new int[] { 1 }));
+
             return CreateInputParameters(
                  NamedOnnxValue.CreateFromTensor(inputNames[0], inputTensor),
-                 NamedOnnxValue.CreateFromTensor(inputNames[1], new DenseTensor<long>(new long[] { timestep }, new int[] { 1 })),
+                 timestepNamedOnnxValue,
                  NamedOnnxValue.CreateFromTensor(inputNames[2], promptEmbeddings),
                  NamedOnnxValue.CreateFromTensor(inputNames[3], guidanceEmbeddings));
         }

diff --git a/OnnxStack.StableDiffusion/Diffusers/StableDiffusion/StableDiffusionDiffuser.cs b/OnnxStack.StableDiffusion/Diffusers/StableDiffusion/StableDiffusionDiffuser.cs
@@ -213,9 +213,18 @@ protected virtual DenseTensor<float> PerformGuidance(DenseTensor<float> noisePre
         protected virtual IReadOnlyList<NamedOnnxValue> CreateUnetInputParams(IModelOptions model, DenseTensor<float> inputTensor, DenseTensor<float> promptEmbeddings, int timestep)
         {
             var inputNames = _onnxModelService.GetInputNames(model, OnnxModelType.Unet);
+            var inputMetaData = _onnxModelService.GetInputMetadata(model, OnnxModelType.Unet);
+
+            // Some models support Long or Float, could be more but fornow just support these 2
+            var timesepMetaKey = inputNames[1];
+            var timestepMetaData = inputMetaData[timesepMetaKey];
+            var timestepNamedOnnxValue = timestepMetaData.ElementDataType == TensorElementType.Int64
+                ? NamedOnnxValue.CreateFromTensor(timesepMetaKey, new DenseTensor<long>(new long[] { timestep }, new int[] { 1 }))
+                : NamedOnnxValue.CreateFromTensor(timesepMetaKey, new DenseTensor<float>(new float[] { timestep }, new int[] { 1 }));
+
             return CreateInputParameters(
                  NamedOnnxValue.CreateFromTensor(inputNames[0], inputTensor),
-                 NamedOnnxValue.CreateFromTensor(inputNames[1], new DenseTensor<long>(new long[] { timestep }, new int[] { 1 })),
+                 timestepNamedOnnxValue,
                  NamedOnnxValue.CreateFromTensor(inputNames[2], promptEmbeddings));
         }