Skip to content

Commit

Permalink
Merge pull request #134 from saddam213/StableCascade
Browse files Browse the repository at this point in the history 
StableCascade Pipeline
  • Loading branch information
@saddam213
saddam213 authored Apr 30, 2024
2 parents 9b84ca2 + dfda689 commit c8a58fe
Show file tree
Hide file tree
Showing 25 changed files with 1,392 additions and 267 deletions.
54 changes: 54 additions & 0 deletions OnnxStack.Console/Examples/StableCascadeExample.cs
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1,54 @@
using OnnxStack.Core.Image;
using OnnxStack.StableDiffusion.Config;
using OnnxStack.StableDiffusion.Enums;
using OnnxStack.StableDiffusion.Pipelines;
using SixLabors.ImageSharp;
using System.Diagnostics;

namespace OnnxStack.Console.Runner
{
public sealed class StableCascadeExample : IExampleRunner
{
private readonly string _outputDirectory;
private readonly StableDiffusionConfig _configuration;

public StableCascadeExample(StableDiffusionConfig configuration)
{
_configuration = configuration;
_outputDirectory = Path.Combine(Directory.GetCurrentDirectory(), "Examples", nameof(StableCascadeExample));
Directory.CreateDirectory(_outputDirectory);
}

public int Index => 20;

public string Name => "Stable Cascade Demo";

public string Description => "Creates images from the text prompt provided";

/// <summary>
/// Stable Diffusion Demo
/// </summary>
public async Task RunAsync()
{
// Prompt
var promptOptions = new PromptOptions
{
Prompt = "an image of a cat, donning a spacesuit and helmet",
DiffuserType = DiffuserType.TextToImage,
//InputImage = await OnnxImage.FromFileAsync("Input.png"),
};

// Create Pipeline
var pipeline = StableCascadePipeline.CreatePipeline("D:\\Models\\stable-cascade-onnx", memoryMode: MemoryModeType.Minimum);

// Run pipeline
var result = await pipeline.GenerateImageAsync(promptOptions, progressCallback: OutputHelpers.ProgressCallback);

// Save Image File
await result.SaveAsync(Path.Combine(_outputDirectory, $"output.png"));

// Unload
await pipeline.UnloadAsync();
}
}
}
36 changes: 18 additions & 18 deletions OnnxStack.Console/appsettings.json
View file
Original file line number Diff line number Diff line change
Expand Up @@ -25,29 +25,29 @@
"ExecutionProvider": "DirectML",
"SchedulerOptions": {
"InferenceSteps": 22,
"GuidanceScale": 8
"GuidanceScale": 8
},
"TokenizerConfig": {
"PadTokenId": 49407,
"BlankTokenId": 49407,
"TokenizerLimit": 77,
"TokenizerLength": 768,
"OnnxModelPath": "D:\\Repositories\\stable-diffusion-v1-5\\cliptokenizer.onnx"
"OnnxModelPath": "D:\\Repositories\\stable-diffusion-v1-5-onnx\\cliptokenizer.onnx"
},
"TextEncoderConfig": {
"OnnxModelPath": "D:\\Repositories\\stable-diffusion-v1-5\\text_encoder\\model.onnx"
"OnnxModelPath": "D:\\Repositories\\stable-diffusion-v1-5-onnx\\text_encoder\\model.onnx"
},
"UnetConfig": {
"ModelType": "Base",
"OnnxModelPath": "D:\\Repositories\\stable-diffusion-v1-5\\unet\\model.onnx"
"OnnxModelPath": "D:\\Repositories\\stable-diffusion-v1-5-onnx\\unet\\model.onnx"
},
"VaeDecoderConfig": {
"ScaleFactor": 0.18215,
"OnnxModelPath": "D:\\Repositories\\stable-diffusion-v1-5\\vae_decoder\\model.onnx"
"OnnxModelPath": "D:\\Repositories\\stable-diffusion-v1-5-onnx\\vae_decoder\\model.onnx"
},
"VaeEncoderConfig": {
"ScaleFactor": 0.18215,
"OnnxModelPath": "D:\\Repositories\\stable-diffusion-v1-5\\vae_encoder\\model.onnx"
"OnnxModelPath": "D:\\Repositories\\stable-diffusion-v1-5-onnx\\vae_encoder\\model.onnx"
}
},
{
Expand All @@ -70,22 +70,22 @@
"BlankTokenId": 49407,
"TokenizerLimit": 77,
"TokenizerLength": 768,
"OnnxModelPath": "D:\\Repositories\\stable-diffusion-v1-5\\cliptokenizer.onnx"
"OnnxModelPath": "D:\\Repositories\\LCM_Dreamshaper_v7-onnx\\tokenizer\\model.onnx"
},
"TextEncoderConfig": {
"OnnxModelPath": "D:\\Repositories\\lcm-dreamshaper-v7-f16\\text_encoder\\model.onnx"
"OnnxModelPath": "D:\\Repositories\\LCM_Dreamshaper_v7-onnx\\text_encoder\\model.onnx"
},
"UnetConfig": {
"ModelType": "Base",
"OnnxModelPath": "D:\\Repositories\\lcm-dreamshaper-v7-f16\\unet\\model.onnx"
"OnnxModelPath": "D:\\Repositories\\LCM_Dreamshaper_v7-onnx\\unet\\model.onnx"
},
"VaeDecoderConfig": {
"ScaleFactor": 0.18215,
"OnnxModelPath": "D:\\Repositories\\lcm-dreamshaper-v7-f16\\vae_decoder\\model.onnx"
"OnnxModelPath": "D:\\Repositories\\LCM_Dreamshaper_v7-onnx\\vae_decoder\\model.onnx"
},
"VaeEncoderConfig": {
"ScaleFactor": 0.18215,
"OnnxModelPath": "D:\\Repositories\\lcm-dreamshaper-v7-f16\\vae_encoder\\model.onnx"
"OnnxModelPath": "D:\\Repositories\\LCM_Dreamshaper_v7-onnx\\vae_encoder\\model.onnx"
}
},
{
Expand All @@ -108,32 +108,32 @@
"BlankTokenId": 49407,
"TokenizerLimit": 77,
"TokenizerLength": 768,
"OnnxModelPath": "D:\\Repositories\\stable-diffusion-v1-5\\cliptokenizer.onnx"
"OnnxModelPath": "D:\\Repositories\\stable-diffusion-xl-base-1.0-onnx\\tokenizer\\model.onnx"
},
"Tokenizer2Config": {
"PadTokenId": 1,
"BlankTokenId": 49407,
"TokenizerLimit": 77,
"TokenizerLength": 1280,
"OnnxModelPath": "D:\\Repositories\\stable-diffusion-v1-5\\cliptokenizer.onnx"
"OnnxModelPath": "D:\\Repositories\\stable-diffusion-xl-base-1.0-onnx\\tokenizer_2\\model.onnx"
},
"TextEncoderConfig": {
"OnnxModelPath": "D:\\Repositories\\stable-diffusion-xl-base-1.0-Olive-Onnx\\text_encoder\\model.onnx"
"OnnxModelPath": "D:\\Repositories\\stable-diffusion-xl-base-1.0-onnx\\text_encoder\\model.onnx"
},
"TextEncoder2Config": {
"OnnxModelPath": "D:\\Repositories\\stable-diffusion-xl-base-1.0-Olive-Onnx\\text_encoder_2\\model.onnx"
"OnnxModelPath": "D:\\Repositories\\stable-diffusion-xl-base-1.0-onnx\\text_encoder_2\\model.onnx"
},
"UnetConfig": {
"ModelType": "Base",
"OnnxModelPath": "D:\\Repositories\\stable-diffusion-xl-base-1.0-Olive-Onnx\\unet\\model.onnx"
"OnnxModelPath": "D:\\Repositories\\stable-diffusion-xl-base-1.0-onnx\\unet\\model.onnx"
},
"VaeDecoderConfig": {
"ScaleFactor": 0.13025,
"OnnxModelPath": "D:\\Repositories\\stable-diffusion-xl-base-1.0-Olive-Onnx\\vae_decoder\\model.onnx"
"OnnxModelPath": "D:\\Repositories\\stable-diffusion-xl-base-1.0-onnx\\vae_decoder\\model.onnx"
},
"VaeEncoderConfig": {
"ScaleFactor": 0.13025,
"OnnxModelPath": "D:\\Repositories\\stable-diffusion-xl-base-1.0-Olive-Onnx\\vae_encoder\\model.onnx"
"OnnxModelPath": "D:\\Repositories\\stable-diffusion-xl-base-1.0-onnx\\vae_encoder\\model.onnx"
}
}
]
Expand Down
94 changes: 79 additions & 15 deletions OnnxStack.Converter/stable_cascade/config_text_encoder.json
View file
Original file line number Diff line number Diff line change
Expand Up @@ -5,10 +5,84 @@
"model_path": "stabilityai/stable-cascade",
"model_loader": "text_encoder_load",
"model_script": "models.py",
"io_config": {
"input_names": [ "input_ids" ],
"output_names": [ "last_hidden_state", "pooler_output" ],
"dynamic_axes": { "input_ids": { "0": "batch", "1": "sequence" } }
"io_config": {
"input_names": [ "input_ids", "attention_mask", "output_hidden_states" ],
"output_names": [
"text_embeds",
"last_hidden_state",
"hidden_states.0",
"hidden_states.1",
"hidden_states.2",
"hidden_states.3",
"hidden_states.4",
"hidden_states.5",
"hidden_states.6",
"hidden_states.7",
"hidden_states.8",
"hidden_states.9",
"hidden_states.10",
"hidden_states.11",
"hidden_states.12",
"hidden_states.13",
"hidden_states.14",
"hidden_states.15",
"hidden_states.16",
"hidden_states.17",
"hidden_states.18",
"hidden_states.19",
"hidden_states.20",
"hidden_states.21",
"hidden_states.22",
"hidden_states.23",
"hidden_states.24",
"hidden_states.25",
"hidden_states.26",
"hidden_states.27",
"hidden_states.28",
"hidden_states.29",
"hidden_states.30",
"hidden_states.31",
"hidden_states.32"
],
"dynamic_axes": {
"input_ids": { "0": "batch_size", "1": "sequence_length" },
"attention_mask": { "0": "batch_size", "1": "sequence_length" },
"text_embeds": { "0": "batch_size"},
"last_hidden_state": { "0": "batch_size", "1": "sequence_length" },
"hidden_states.0": { "0": "batch_size", "1": "sequence_length" },
"hidden_states.1": { "0": "batch_size", "1": "sequence_length" },
"hidden_states.2": { "0": "batch_size", "1": "sequence_length" },
"hidden_states.3": { "0": "batch_size", "1": "sequence_length" },
"hidden_states.4": { "0": "batch_size", "1": "sequence_length" },
"hidden_states.5": { "0": "batch_size", "1": "sequence_length" },
"hidden_states.6": { "0": "batch_size", "1": "sequence_length" },
"hidden_states.7": { "0": "batch_size", "1": "sequence_length" },
"hidden_states.8": { "0": "batch_size", "1": "sequence_length" },
"hidden_states.9": { "0": "batch_size", "1": "sequence_length" },
"hidden_states.10": { "0": "batch_size", "1": "sequence_length" },
"hidden_states.11": { "0": "batch_size", "1": "sequence_length" },
"hidden_states.12": { "0": "batch_size", "1": "sequence_length" },
"hidden_states.13": { "0": "batch_size", "1": "sequence_length" },
"hidden_states.14": { "0": "batch_size", "1": "sequence_length" },
"hidden_states.15": { "0": "batch_size", "1": "sequence_length" },
"hidden_states.16": { "0": "batch_size", "1": "sequence_length" },
"hidden_states.17": { "0": "batch_size", "1": "sequence_length" },
"hidden_states.18": { "0": "batch_size", "1": "sequence_length" },
"hidden_states.19": { "0": "batch_size", "1": "sequence_length" },
"hidden_states.20": { "0": "batch_size", "1": "sequence_length" },
"hidden_states.21": { "0": "batch_size", "1": "sequence_length" },
"hidden_states.22": { "0": "batch_size", "1": "sequence_length" },
"hidden_states.23": { "0": "batch_size", "1": "sequence_length" },
"hidden_states.24": { "0": "batch_size", "1": "sequence_length" },
"hidden_states.25": { "0": "batch_size", "1": "sequence_length" },
"hidden_states.26": { "0": "batch_size", "1": "sequence_length" },
"hidden_states.27": { "0": "batch_size", "1": "sequence_length" },
"hidden_states.28": { "0": "batch_size", "1": "sequence_length" },
"hidden_states.29": { "0": "batch_size", "1": "sequence_length" },
"hidden_states.30": { "0": "batch_size", "1": "sequence_length" },
"hidden_states.31": { "0": "batch_size", "1": "sequence_length" },
"hidden_states.32": { "0": "batch_size", "1": "sequence_length" }
}
},
"dummy_inputs_func": "text_encoder_conversion_inputs"
}
Expand Down Expand Up @@ -58,7 +132,7 @@
"opt_level": 0,
"float16": true,
"use_gpu": true,
"keep_io_types": true,
"keep_io_types": false,
"optimization_options": {
"enable_gelu": true,
"enable_layer_norm": true,
Expand All @@ -85,16 +159,6 @@
"GroupNorm": [0, 1, 2]
}
}
},
"optimize_cuda": {
"type": "OrtTransformersOptimization",
"config": {
"model_type": "clip",
"opt_level": 0,
"float16": true,
"use_gpu": true,
"keep_io_types": false
}
}
},
"pass_flows": [
Expand Down
59 changes: 31 additions & 28 deletions OnnxStack.Converter/stable_cascade/convert.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -27,7 +27,7 @@ def optimize(
model_input: str,
model_output: Path,
provider: str,
controlnet: bool
submodel_names: list[str]
):
from google.protobuf import __version__ as protobuf_version

Expand All @@ -51,18 +51,10 @@ def optimize(
# config.unet_sample_size = pipeline.unet.config.sample_size

model_info = {}
submodel_names = [ "tokenizer", "text_encoder", "decoder", "prior", "vqgan" ]
has_safety_checker = getattr(pipeline, "safety_checker", None) is not None

if has_safety_checker:
submodel_names.append("safety_checker")

if controlnet:
submodel_names.append("controlnet")

for submodel_name in submodel_names:
if submodel_name == "tokenizer":
save_onnx_tokenizer_model(script_dir, model_input, model_info)
if submodel_name == "tokenizer" or submodel_name == "tokenizer_2":
save_onnx_tokenizer_model(script_dir, model_input, submodel_name, model_info)
continue

print(f"\nOptimizing {submodel_name}")
Expand All @@ -85,16 +77,18 @@ def save_onnx_Models(model_dir, model_info, model_output, submodel_names):
conversion_dir = model_output / conversion_type
conversion_dir.mkdir(parents=True, exist_ok=True)

only_unet = "prior" in submodel_names and len(submodel_names) <= 2
# Copy the config and other files required by some applications
model_index_path = model_dir / "model_index.json"
if os.path.exists(model_index_path):
shutil.copy(model_index_path, conversion_dir)
if os.path.exists(model_dir / "tokenizer"):
shutil.copytree(model_dir / "tokenizer",
conversion_dir / "tokenizer")
if os.path.exists(model_dir / "scheduler"):
shutil.copytree(model_dir / "scheduler",
conversion_dir / "scheduler")
if only_unet is False:
model_index_path = model_dir / "model_index.json"
if os.path.exists(model_index_path):
shutil.copy(model_index_path, conversion_dir)
if os.path.exists(model_dir / "tokenizer"):
shutil.copytree(model_dir / "tokenizer", conversion_dir / "tokenizer")
if os.path.exists(model_dir / "tokenizer_2"):
shutil.copytree(model_dir / "tokenizer_2", conversion_dir / "tokenizer_2")
if os.path.exists(model_dir / "scheduler"):
shutil.copytree(model_dir / "scheduler", conversion_dir / "scheduler")

# Save models files
for submodel_name in submodel_names:
Expand Down Expand Up @@ -159,10 +153,10 @@ def save_onnx_submodel(script_dir, submodel_name, model_info, provider):
print(f"Unoptimized Model : {model_info[submodel_name]['unoptimized']['path']}")


def save_onnx_tokenizer_model(script_dir, model_dir, model_info, max_length=-1, attention_mask=True, offset_map=False):
def save_onnx_tokenizer_model(script_dir, model_dir, submodel_name, model_info, max_length=-1, attention_mask=True, offset_map=False):
model_dir = Path(model_dir)
vocab_file = model_dir / "tokenizer" / "vocab.json"
merges_file = model_dir / "tokenizer" / "merges.txt"
vocab_file = model_dir / submodel_name / "vocab.json"
merges_file = model_dir / submodel_name / "merges.txt"

input1 = helper.make_tensor_value_info('string_input', onnx_proto.TensorProto.STRING, [None])
output1 = helper.make_tensor_value_info('input_ids', onnx_proto.TensorProto.INT64, ["batch_size", "num_input_ids"])
Expand Down Expand Up @@ -194,11 +188,11 @@ def save_onnx_tokenizer_model(script_dir, model_dir, model_info, max_length=-1,
graph = helper.make_graph(node, 'main_graph', inputs, outputs)
model = make_onnx_model(graph)

output_dir = script_dir / "cache" / "tokenizer"
output_dir = script_dir / "cache" / submodel_name
output_dir.mkdir(parents=True, exist_ok=True)
output_path = output_dir / 'model.onnx'
onnx.save(model, output_path)
model_info["tokenizer"] = {
model_info[submodel_name] = {
"optimized": {"path": output_path},
"unoptimized": {"path": output_path}
}
Expand All @@ -213,9 +207,11 @@ def parse_common_args(raw_args):
parser = argparse.ArgumentParser("Common arguments")
parser.add_argument("--model_input", default="stable-diffusion-v1-5", type=str)
parser.add_argument("--model_output", default=None, type=Path)
parser.add_argument("--controlnet", action="store_true", help="Create ControlNet Unet Model")
parser.add_argument("--image_encoder", action="store_true", help="Create image_encoder Model")
parser.add_argument("--clean", action="store_true", help="Deletes the Olive cache")
parser.add_argument("--tempdir", default=None, type=str, help="Root directory for tempfile directories and files")
parser.add_argument("--only_unet", action="store_true", help="Only convert UNET model")

return parser.parse_known_args(raw_args)


Expand All @@ -241,10 +237,17 @@ def main(raw_args=None):

set_tempdir(common_args.tempdir)

submodel_names = ["tokenizer", "text_encoder", "prior", "decoder", "vqgan"]

if common_args.image_encoder:
submodel_names.append("image_encoder")

if common_args.only_unet:
submodel_names = ["prior", "decoder"]

with warnings.catch_warnings():
warnings.simplefilter("ignore")
optimize(script_dir, common_args.model_input,
model_output, provider, common_args.controlnet)
optimize(script_dir, common_args.model_input, model_output, provider, submodel_names)


if __name__ == "__main__":
Expand Down
Loading

0 comments on commit c8a58fe

Please sign in to comment.