Skip to content
This repository was archived by the owner on Nov 27, 2024. It is now read-only.

Commit c8a58fe

Browse files
saddam213saddam213
authored
Merge pull request #134 from saddam213/StableCascade
StableCascade Pipeline
2 parents 9b84ca2 + dfda689 commit c8a58fe

25 files changed

+1392
-267
lines changed

‎OnnxStack.Console/Examples/StableCascadeExample.cs

+54
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,54 @@
1+
using OnnxStack.Core.Image;
2+
using OnnxStack.StableDiffusion.Config;
3+
using OnnxStack.StableDiffusion.Enums;
4+
using OnnxStack.StableDiffusion.Pipelines;
5+
using SixLabors.ImageSharp;
6+
using System.Diagnostics;
7+
8+
namespace OnnxStack.Console.Runner
9+
{
10+
public sealed class StableCascadeExample : IExampleRunner
11+
{
12+
private readonly string _outputDirectory;
13+
private readonly StableDiffusionConfig _configuration;
14+
15+
public StableCascadeExample(StableDiffusionConfig configuration)
16+
{
17+
_configuration = configuration;
18+
_outputDirectory = Path.Combine(Directory.GetCurrentDirectory(), "Examples", nameof(StableCascadeExample));
19+
Directory.CreateDirectory(_outputDirectory);
20+
}
21+
22+
public int Index => 20;
23+
24+
public string Name => "Stable Cascade Demo";
25+
26+
public string Description => "Creates images from the text prompt provided";
27+
28+
/// <summary>
29+
/// Stable Diffusion Demo
30+
/// </summary>
31+
public async Task RunAsync()
32+
{
33+
// Prompt
34+
var promptOptions = new PromptOptions
35+
{
36+
Prompt = "an image of a cat, donning a spacesuit and helmet",
37+
DiffuserType = DiffuserType.TextToImage,
38+
//InputImage = await OnnxImage.FromFileAsync("Input.png"),
39+
};
40+
41+
// Create Pipeline
42+
var pipeline = StableCascadePipeline.CreatePipeline("D:\\Models\\stable-cascade-onnx", memoryMode: MemoryModeType.Minimum);
43+
44+
// Run pipeline
45+
var result = await pipeline.GenerateImageAsync(promptOptions, progressCallback: OutputHelpers.ProgressCallback);
46+
47+
// Save Image File
48+
await result.SaveAsync(Path.Combine(_outputDirectory, $"output.png"));
49+
50+
// Unload
51+
await pipeline.UnloadAsync();
52+
}
53+
}
54+
}

‎OnnxStack.Console/appsettings.json

+18-18
Original file line numberDiff line numberDiff line change
@@ -25,29 +25,29 @@
2525
"ExecutionProvider": "DirectML",
2626
"SchedulerOptions": {
2727
"InferenceSteps": 22,
28-
"GuidanceScale": 8
28+
"GuidanceScale": 8
2929
},
3030
"TokenizerConfig": {
3131
"PadTokenId": 49407,
3232
"BlankTokenId": 49407,
3333
"TokenizerLimit": 77,
3434
"TokenizerLength": 768,
35-
"OnnxModelPath": "D:\\Repositories\\stable-diffusion-v1-5\\cliptokenizer.onnx"
35+
"OnnxModelPath": "D:\\Repositories\\stable-diffusion-v1-5-onnx\\cliptokenizer.onnx"
3636
},
3737
"TextEncoderConfig": {
38-
"OnnxModelPath": "D:\\Repositories\\stable-diffusion-v1-5\\text_encoder\\model.onnx"
38+
"OnnxModelPath": "D:\\Repositories\\stable-diffusion-v1-5-onnx\\text_encoder\\model.onnx"
3939
},
4040
"UnetConfig": {
4141
"ModelType": "Base",
42-
"OnnxModelPath": "D:\\Repositories\\stable-diffusion-v1-5\\unet\\model.onnx"
42+
"OnnxModelPath": "D:\\Repositories\\stable-diffusion-v1-5-onnx\\unet\\model.onnx"
4343
},
4444
"VaeDecoderConfig": {
4545
"ScaleFactor": 0.18215,
46-
"OnnxModelPath": "D:\\Repositories\\stable-diffusion-v1-5\\vae_decoder\\model.onnx"
46+
"OnnxModelPath": "D:\\Repositories\\stable-diffusion-v1-5-onnx\\vae_decoder\\model.onnx"
4747
},
4848
"VaeEncoderConfig": {
4949
"ScaleFactor": 0.18215,
50-
"OnnxModelPath": "D:\\Repositories\\stable-diffusion-v1-5\\vae_encoder\\model.onnx"
50+
"OnnxModelPath": "D:\\Repositories\\stable-diffusion-v1-5-onnx\\vae_encoder\\model.onnx"
5151
}
5252
},
5353
{
@@ -70,22 +70,22 @@
7070
"BlankTokenId": 49407,
7171
"TokenizerLimit": 77,
7272
"TokenizerLength": 768,
73-
"OnnxModelPath": "D:\\Repositories\\stable-diffusion-v1-5\\cliptokenizer.onnx"
73+
"OnnxModelPath": "D:\\Repositories\\LCM_Dreamshaper_v7-onnx\\tokenizer\\model.onnx"
7474
},
7575
"TextEncoderConfig": {
76-
"OnnxModelPath": "D:\\Repositories\\lcm-dreamshaper-v7-f16\\text_encoder\\model.onnx"
76+
"OnnxModelPath": "D:\\Repositories\\LCM_Dreamshaper_v7-onnx\\text_encoder\\model.onnx"
7777
},
7878
"UnetConfig": {
7979
"ModelType": "Base",
80-
"OnnxModelPath": "D:\\Repositories\\lcm-dreamshaper-v7-f16\\unet\\model.onnx"
80+
"OnnxModelPath": "D:\\Repositories\\LCM_Dreamshaper_v7-onnx\\unet\\model.onnx"
8181
},
8282
"VaeDecoderConfig": {
8383
"ScaleFactor": 0.18215,
84-
"OnnxModelPath": "D:\\Repositories\\lcm-dreamshaper-v7-f16\\vae_decoder\\model.onnx"
84+
"OnnxModelPath": "D:\\Repositories\\LCM_Dreamshaper_v7-onnx\\vae_decoder\\model.onnx"
8585
},
8686
"VaeEncoderConfig": {
8787
"ScaleFactor": 0.18215,
88-
"OnnxModelPath": "D:\\Repositories\\lcm-dreamshaper-v7-f16\\vae_encoder\\model.onnx"
88+
"OnnxModelPath": "D:\\Repositories\\LCM_Dreamshaper_v7-onnx\\vae_encoder\\model.onnx"
8989
}
9090
},
9191
{
@@ -108,32 +108,32 @@
108108
"BlankTokenId": 49407,
109109
"TokenizerLimit": 77,
110110
"TokenizerLength": 768,
111-
"OnnxModelPath": "D:\\Repositories\\stable-diffusion-v1-5\\cliptokenizer.onnx"
111+
"OnnxModelPath": "D:\\Repositories\\stable-diffusion-xl-base-1.0-onnx\\tokenizer\\model.onnx"
112112
},
113113
"Tokenizer2Config": {
114114
"PadTokenId": 1,
115115
"BlankTokenId": 49407,
116116
"TokenizerLimit": 77,
117117
"TokenizerLength": 1280,
118-
"OnnxModelPath": "D:\\Repositories\\stable-diffusion-v1-5\\cliptokenizer.onnx"
118+
"OnnxModelPath": "D:\\Repositories\\stable-diffusion-xl-base-1.0-onnx\\tokenizer_2\\model.onnx"
119119
},
120120
"TextEncoderConfig": {
121-
"OnnxModelPath": "D:\\Repositories\\stable-diffusion-xl-base-1.0-Olive-Onnx\\text_encoder\\model.onnx"
121+
"OnnxModelPath": "D:\\Repositories\\stable-diffusion-xl-base-1.0-onnx\\text_encoder\\model.onnx"
122122
},
123123
"TextEncoder2Config": {
124-
"OnnxModelPath": "D:\\Repositories\\stable-diffusion-xl-base-1.0-Olive-Onnx\\text_encoder_2\\model.onnx"
124+
"OnnxModelPath": "D:\\Repositories\\stable-diffusion-xl-base-1.0-onnx\\text_encoder_2\\model.onnx"
125125
},
126126
"UnetConfig": {
127127
"ModelType": "Base",
128-
"OnnxModelPath": "D:\\Repositories\\stable-diffusion-xl-base-1.0-Olive-Onnx\\unet\\model.onnx"
128+
"OnnxModelPath": "D:\\Repositories\\stable-diffusion-xl-base-1.0-onnx\\unet\\model.onnx"
129129
},
130130
"VaeDecoderConfig": {
131131
"ScaleFactor": 0.13025,
132-
"OnnxModelPath": "D:\\Repositories\\stable-diffusion-xl-base-1.0-Olive-Onnx\\vae_decoder\\model.onnx"
132+
"OnnxModelPath": "D:\\Repositories\\stable-diffusion-xl-base-1.0-onnx\\vae_decoder\\model.onnx"
133133
},
134134
"VaeEncoderConfig": {
135135
"ScaleFactor": 0.13025,
136-
"OnnxModelPath": "D:\\Repositories\\stable-diffusion-xl-base-1.0-Olive-Onnx\\vae_encoder\\model.onnx"
136+
"OnnxModelPath": "D:\\Repositories\\stable-diffusion-xl-base-1.0-onnx\\vae_encoder\\model.onnx"
137137
}
138138
}
139139
]

‎OnnxStack.Converter/stable_cascade/config_text_encoder.json

+79-15
Original file line numberDiff line numberDiff line change
@@ -5,10 +5,84 @@
55
"model_path": "stabilityai/stable-cascade",
66
"model_loader": "text_encoder_load",
77
"model_script": "models.py",
8-
"io_config": {
9-
"input_names": [ "input_ids" ],
10-
"output_names": [ "last_hidden_state", "pooler_output" ],
11-
"dynamic_axes": { "input_ids": { "0": "batch", "1": "sequence" } }
8+
"io_config": {
9+
"input_names": [ "input_ids", "attention_mask", "output_hidden_states" ],
10+
"output_names": [
11+
"text_embeds",
12+
"last_hidden_state",
13+
"hidden_states.0",
14+
"hidden_states.1",
15+
"hidden_states.2",
16+
"hidden_states.3",
17+
"hidden_states.4",
18+
"hidden_states.5",
19+
"hidden_states.6",
20+
"hidden_states.7",
21+
"hidden_states.8",
22+
"hidden_states.9",
23+
"hidden_states.10",
24+
"hidden_states.11",
25+
"hidden_states.12",
26+
"hidden_states.13",
27+
"hidden_states.14",
28+
"hidden_states.15",
29+
"hidden_states.16",
30+
"hidden_states.17",
31+
"hidden_states.18",
32+
"hidden_states.19",
33+
"hidden_states.20",
34+
"hidden_states.21",
35+
"hidden_states.22",
36+
"hidden_states.23",
37+
"hidden_states.24",
38+
"hidden_states.25",
39+
"hidden_states.26",
40+
"hidden_states.27",
41+
"hidden_states.28",
42+
"hidden_states.29",
43+
"hidden_states.30",
44+
"hidden_states.31",
45+
"hidden_states.32"
46+
],
47+
"dynamic_axes": {
48+
"input_ids": { "0": "batch_size", "1": "sequence_length" },
49+
"attention_mask": { "0": "batch_size", "1": "sequence_length" },
50+
"text_embeds": { "0": "batch_size"},
51+
"last_hidden_state": { "0": "batch_size", "1": "sequence_length" },
52+
"hidden_states.0": { "0": "batch_size", "1": "sequence_length" },
53+
"hidden_states.1": { "0": "batch_size", "1": "sequence_length" },
54+
"hidden_states.2": { "0": "batch_size", "1": "sequence_length" },
55+
"hidden_states.3": { "0": "batch_size", "1": "sequence_length" },
56+
"hidden_states.4": { "0": "batch_size", "1": "sequence_length" },
57+
"hidden_states.5": { "0": "batch_size", "1": "sequence_length" },
58+
"hidden_states.6": { "0": "batch_size", "1": "sequence_length" },
59+
"hidden_states.7": { "0": "batch_size", "1": "sequence_length" },
60+
"hidden_states.8": { "0": "batch_size", "1": "sequence_length" },
61+
"hidden_states.9": { "0": "batch_size", "1": "sequence_length" },
62+
"hidden_states.10": { "0": "batch_size", "1": "sequence_length" },
63+
"hidden_states.11": { "0": "batch_size", "1": "sequence_length" },
64+
"hidden_states.12": { "0": "batch_size", "1": "sequence_length" },
65+
"hidden_states.13": { "0": "batch_size", "1": "sequence_length" },
66+
"hidden_states.14": { "0": "batch_size", "1": "sequence_length" },
67+
"hidden_states.15": { "0": "batch_size", "1": "sequence_length" },
68+
"hidden_states.16": { "0": "batch_size", "1": "sequence_length" },
69+
"hidden_states.17": { "0": "batch_size", "1": "sequence_length" },
70+
"hidden_states.18": { "0": "batch_size", "1": "sequence_length" },
71+
"hidden_states.19": { "0": "batch_size", "1": "sequence_length" },
72+
"hidden_states.20": { "0": "batch_size", "1": "sequence_length" },
73+
"hidden_states.21": { "0": "batch_size", "1": "sequence_length" },
74+
"hidden_states.22": { "0": "batch_size", "1": "sequence_length" },
75+
"hidden_states.23": { "0": "batch_size", "1": "sequence_length" },
76+
"hidden_states.24": { "0": "batch_size", "1": "sequence_length" },
77+
"hidden_states.25": { "0": "batch_size", "1": "sequence_length" },
78+
"hidden_states.26": { "0": "batch_size", "1": "sequence_length" },
79+
"hidden_states.27": { "0": "batch_size", "1": "sequence_length" },
80+
"hidden_states.28": { "0": "batch_size", "1": "sequence_length" },
81+
"hidden_states.29": { "0": "batch_size", "1": "sequence_length" },
82+
"hidden_states.30": { "0": "batch_size", "1": "sequence_length" },
83+
"hidden_states.31": { "0": "batch_size", "1": "sequence_length" },
84+
"hidden_states.32": { "0": "batch_size", "1": "sequence_length" }
85+
}
1286
},
1387
"dummy_inputs_func": "text_encoder_conversion_inputs"
1488
}
@@ -58,7 +132,7 @@
58132
"opt_level": 0,
59133
"float16": true,
60134
"use_gpu": true,
61-
"keep_io_types": true,
135+
"keep_io_types": false,
62136
"optimization_options": {
63137
"enable_gelu": true,
64138
"enable_layer_norm": true,
@@ -85,16 +159,6 @@
85159
"GroupNorm": [0, 1, 2]
86160
}
87161
}
88-
},
89-
"optimize_cuda": {
90-
"type": "OrtTransformersOptimization",
91-
"config": {
92-
"model_type": "clip",
93-
"opt_level": 0,
94-
"float16": true,
95-
"use_gpu": true,
96-
"keep_io_types": false
97-
}
98162
}
99163
},
100164
"pass_flows": [

‎OnnxStack.Converter/stable_cascade/convert.py

+31-28
Original file line numberDiff line numberDiff line change
@@ -27,7 +27,7 @@ def optimize(
2727
model_input: str,
2828
model_output: Path,
2929
provider: str,
30-
controlnet: bool
30+
submodel_names: list[str]
3131
):
3232
from google.protobuf import __version__ as protobuf_version
3333

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

5353
model_info = {}
54-
submodel_names = [ "tokenizer", "text_encoder", "decoder", "prior", "vqgan" ]
55-
has_safety_checker = getattr(pipeline, "safety_checker", None) is not None
56-
57-
if has_safety_checker:
58-
submodel_names.append("safety_checker")
59-
60-
if controlnet:
61-
submodel_names.append("controlnet")
6254

6355
for submodel_name in submodel_names:
64-
if submodel_name == "tokenizer":
65-
save_onnx_tokenizer_model(script_dir, model_input, model_info)
56+
if submodel_name == "tokenizer" or submodel_name == "tokenizer_2":
57+
save_onnx_tokenizer_model(script_dir, model_input, submodel_name, model_info)
6658
continue
6759

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

80+
only_unet = "prior" in submodel_names and len(submodel_names) <= 2
8881
# Copy the config and other files required by some applications
89-
model_index_path = model_dir / "model_index.json"
90-
if os.path.exists(model_index_path):
91-
shutil.copy(model_index_path, conversion_dir)
92-
if os.path.exists(model_dir / "tokenizer"):
93-
shutil.copytree(model_dir / "tokenizer",
94-
conversion_dir / "tokenizer")
95-
if os.path.exists(model_dir / "scheduler"):
96-
shutil.copytree(model_dir / "scheduler",
97-
conversion_dir / "scheduler")
82+
if only_unet is False:
83+
model_index_path = model_dir / "model_index.json"
84+
if os.path.exists(model_index_path):
85+
shutil.copy(model_index_path, conversion_dir)
86+
if os.path.exists(model_dir / "tokenizer"):
87+
shutil.copytree(model_dir / "tokenizer", conversion_dir / "tokenizer")
88+
if os.path.exists(model_dir / "tokenizer_2"):
89+
shutil.copytree(model_dir / "tokenizer_2", conversion_dir / "tokenizer_2")
90+
if os.path.exists(model_dir / "scheduler"):
91+
shutil.copytree(model_dir / "scheduler", conversion_dir / "scheduler")
9892

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

161155

162-
def save_onnx_tokenizer_model(script_dir, model_dir, model_info, max_length=-1, attention_mask=True, offset_map=False):
156+
def save_onnx_tokenizer_model(script_dir, model_dir, submodel_name, model_info, max_length=-1, attention_mask=True, offset_map=False):
163157
model_dir = Path(model_dir)
164-
vocab_file = model_dir / "tokenizer" / "vocab.json"
165-
merges_file = model_dir / "tokenizer" / "merges.txt"
158+
vocab_file = model_dir / submodel_name / "vocab.json"
159+
merges_file = model_dir / submodel_name / "merges.txt"
166160

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

197-
output_dir = script_dir / "cache" / "tokenizer"
191+
output_dir = script_dir / "cache" / submodel_name
198192
output_dir.mkdir(parents=True, exist_ok=True)
199193
output_path = output_dir / 'model.onnx'
200194
onnx.save(model, output_path)
201-
model_info["tokenizer"] = {
195+
model_info[submodel_name] = {
202196
"optimized": {"path": output_path},
203197
"unoptimized": {"path": output_path}
204198
}
@@ -213,9 +207,11 @@ def parse_common_args(raw_args):
213207
parser = argparse.ArgumentParser("Common arguments")
214208
parser.add_argument("--model_input", default="stable-diffusion-v1-5", type=str)
215209
parser.add_argument("--model_output", default=None, type=Path)
216-
parser.add_argument("--controlnet", action="store_true", help="Create ControlNet Unet Model")
210+
parser.add_argument("--image_encoder", action="store_true", help="Create image_encoder Model")
217211
parser.add_argument("--clean", action="store_true", help="Deletes the Olive cache")
218212
parser.add_argument("--tempdir", default=None, type=str, help="Root directory for tempfile directories and files")
213+
parser.add_argument("--only_unet", action="store_true", help="Only convert UNET model")
214+
219215
return parser.parse_known_args(raw_args)
220216

221217

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

242238
set_tempdir(common_args.tempdir)
243239

240+
submodel_names = ["tokenizer", "text_encoder", "prior", "decoder", "vqgan"]
241+
242+
if common_args.image_encoder:
243+
submodel_names.append("image_encoder")
244+
245+
if common_args.only_unet:
246+
submodel_names = ["prior", "decoder"]
247+
244248
with warnings.catch_warnings():
245249
warnings.simplefilter("ignore")
246-
optimize(script_dir, common_args.model_input,
247-
model_output, provider, common_args.controlnet)
250+
optimize(script_dir, common_args.model_input, model_output, provider, submodel_names)
248251

249252

250253
if __name__ == "__main__":

0 commit comments

Comments
 (0)