Skip to content

Commit 16ee249

Browse files
committed
Point to main instead of unify-sd branch
1 parent e328124 commit 16ee249

File tree

2 files changed

+927
-292
lines changed

2 files changed

+927
-292
lines changed
Lines changed: 289 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,289 @@
1+
import argparse
2+
import os
3+
from pathlib import Path
4+
5+
6+
def path_expand(s):
7+
return Path(s).expanduser().resolve()
8+
9+
10+
def is_valid_file(arg):
11+
if not os.path.exists(arg):
12+
return None
13+
else:
14+
return arg
15+
16+
17+
# Note: this is where command-line options for the scripts in this directory
18+
# are defined along with their defaults. Thus, they should not be referenced
19+
# within modelling or inference code, only at the entry point to the script.
20+
21+
# We should consider separating out the options that are "model configs" from
22+
# the options that control the compiler, runtime, and script behavior,
23+
# when applicable, as the former would best be kept in a separate
24+
# config or imported from huggingface.
25+
26+
p = argparse.ArgumentParser(
27+
description=__doc__, formatter_class=argparse.ArgumentDefaultsHelpFormatter
28+
)
29+
30+
##############################################################################
31+
# SDXL Huggingface Options
32+
##############################################################################
33+
34+
p.add_argument(
35+
"--hf_auth_token",
36+
type=str,
37+
help="The Hugging Face auth token, if required",
38+
default=None,
39+
)
40+
p.add_argument(
41+
"--hf_model_name",
42+
type=str,
43+
help="HF model name",
44+
default="Trelis/Llama-2-7b-chat-hf-function-calling-v2",
45+
)
46+
p.add_argument(
47+
"--scheduler_id",
48+
type=str,
49+
help="Scheduler ID",
50+
default="Euler",
51+
)
52+
53+
##############################################################################
54+
# SDXL Inference Options
55+
# These options are used to control runtime parameters for SDXL inference.
56+
##############################################################################
57+
58+
p.add_argument(
59+
"--prompt",
60+
type=str,
61+
default=" a cat under the snow with blue eyes, covered by snow, cinematic style, medium shot, professional photo, animal",
62+
help="Prompt input to stable diffusion.",
63+
)
64+
65+
p.add_argument(
66+
"--negative_prompt",
67+
type=str,
68+
default="Watermark, blurry, oversaturated, low resolution, pollution",
69+
help="Negative prompt input to stable diffusion.",
70+
)
71+
72+
p.add_argument(
73+
"--num_inference_steps", type=int, default=30, help="Number of UNet inference steps"
74+
)
75+
76+
p.add_argument(
77+
"--batch_count",
78+
type=int,
79+
default=1,
80+
help="Number of batches to run for a single prompt",
81+
)
82+
83+
p.add_argument(
84+
"--guidance_scale",
85+
type=float,
86+
default=7.5,
87+
help="Scale by which to adjust prompt guidance to the unconditional noise prediction output of UNet after each iteration.",
88+
)
89+
90+
p.add_argument(
91+
"--seed", type=float, default=0, help="Seed for random number/latents generation."
92+
)
93+
94+
p.add_argument(
95+
"--external_weight_path",
96+
type=str,
97+
default="",
98+
help="Path to external weights file, for jobs with one weights filepath. When importing, this is used to specify where to save the model weights, and at runtime, this is used to specify where to load the model weights from.",
99+
)
100+
101+
p.add_argument(
102+
"--external_weights_dir",
103+
type=str,
104+
default="",
105+
help="Directory containing external weights for a job that requires more than one weights file. When importing, this is used to specify where to save the model weights, and at runtime, this is used to specify where to load the model weights from. Files will then be saved according to the parameters that make them unique, i.e. <hf_model_name>_<precision>_<submodel>_<submodel-specific>.<external_weights>",
106+
)
107+
108+
p.add_argument(
109+
"--vmfb_path", type=str, default="", help="path to vmfb containing compiled module"
110+
)
111+
112+
p.add_argument(
113+
"--pipeline_vmfb_path",
114+
type=str,
115+
default="",
116+
help="path to vmfb containing compiled meta-module",
117+
)
118+
119+
p.add_argument(
120+
"--external_weight_file",
121+
type=str,
122+
default=None,
123+
help="Path to external weights, used in benchmark scripts.",
124+
)
125+
126+
p.add_argument(
127+
"--pipeline_dir",
128+
type=str,
129+
default=None,
130+
help="Directory to save pipeline artifacts",
131+
)
132+
133+
p.add_argument(
134+
"--compiled_pipeline",
135+
default=False,
136+
action="store_true",
137+
help="Do one-shot inference from tokens to image in a shrink-wrapped pipeline binary.",
138+
)
139+
140+
##############################################################################
141+
# SDXL Modelling Options
142+
# These options are used to control model defining parameters for SDXL.
143+
# These are MLIR - changing variables! If you change them, you will need
144+
# to import/download and recompile the model.
145+
##############################################################################
146+
147+
p.add_argument("--batch_size", type=int, default=1, help="Batch size for inference")
148+
p.add_argument(
149+
"--height", type=int, default=1024, help="Height of Stable Diffusion output image."
150+
)
151+
p.add_argument(
152+
"--width", type=int, default=1024, help="Width of Stable Diffusion output image"
153+
)
154+
p.add_argument(
155+
"--precision",
156+
type=str,
157+
default="fp16",
158+
help="Precision of Stable Diffusion weights and graph.",
159+
)
160+
p.add_argument(
161+
"--max_length", type=int, default=64, help="Sequence Length of Stable Diffusion"
162+
)
163+
p.add_argument("--vae_variant", type=str, default="decode", help="encode, decode")
164+
p.add_argument(
165+
"--return_index",
166+
action="store_true",
167+
help="Make scheduled unet compiled module return the step index.",
168+
)
169+
170+
p.add_argument(
171+
"--vae_decomp_attn",
172+
type=bool,
173+
default=False,
174+
help="Decompose attention for VAE decode only at fx graph level",
175+
)
176+
177+
##############################################################################
178+
# SDXL script general options.
179+
##############################################################################
180+
181+
p.add_argument("--compile_to", type=str, default="mlir", help="torch, linalg, vmfb")
182+
183+
p.add_argument(
184+
"--external_weights",
185+
type=str,
186+
default=None,
187+
choices=["safetensors", "irpa", "gguf", None],
188+
help="Externalizes model weights from the torch dialect IR and its successors",
189+
)
190+
191+
# See --external_weight_path and --external_weight_dir to specify where to save the model weights.
192+
193+
p.add_argument(
194+
"--compare_vs_torch",
195+
action="store_true",
196+
help="Runs both turbine vmfb and a torch model to compare results",
197+
)
198+
p.add_argument(
199+
"--decomp_attn",
200+
default=False,
201+
action="store_true",
202+
help="Decompose attention at fx graph level",
203+
)
204+
p.add_argument(
205+
"--exit_on_vmfb",
206+
default=True,
207+
action="store_false",
208+
help="Exit program on vmfb compilation completion. Most scripts will also save .mlir if this is disabled.",
209+
)
210+
p.add_argument(
211+
"--input_mlir",
212+
type=str,
213+
default=None,
214+
help="Path to input mlir file to compile. Comma-separate paths to provide more than one input to pipelines.",
215+
)
216+
p.add_argument(
217+
"--download_mlir",
218+
default=False,
219+
action="store_true",
220+
help="Download missing mlir files from Azure storage.",
221+
)
222+
p.add_argument(
223+
"--container_name",
224+
type=str,
225+
default=None,
226+
help="Azure storage container name to download mlir files from.",
227+
)
228+
229+
230+
##############################################################################
231+
# IREE Compiler Options
232+
##############################################################################
233+
234+
p.add_argument("--device", type=str, default="cpu", help="cpu, cuda, vulkan, rocm")
235+
236+
p.add_argument(
237+
"--rt_device",
238+
type=str,
239+
default="local-task",
240+
help="local-task, local-sync, vulkan://0, rocm://0, cuda://0, etc.",
241+
)
242+
243+
# TODO: Bring in detection for target triple
244+
p.add_argument(
245+
"--iree_target_triple",
246+
type=str,
247+
default="",
248+
help="Specify vulkan target triple or rocm/cuda target device.",
249+
)
250+
251+
p.add_argument("--ireec_flags", type=str, default="", help="extra iree-compile options")
252+
253+
p.add_argument(
254+
"--attn_flags",
255+
type=str,
256+
default="",
257+
help="extra iree-compile options for models with iree_linalg_ext.attention ops.",
258+
)
259+
260+
p.add_argument(
261+
"--attn_spec",
262+
type=str,
263+
default=None,
264+
help="extra iree-compile options for models with iree_linalg_ext.attention ops. Set this to 'default' if you are using mfma-capable hardware with ROCM.",
265+
)
266+
267+
p.add_argument(
268+
"--clip_flags",
269+
type=str,
270+
default="",
271+
help="extra iree-compile options to send for compiling CLIP/prompt_encoder. Only use this for testing bleeding edge flags! Any default options should be added to sd_inference/utils.py",
272+
)
273+
274+
p.add_argument(
275+
"--vae_flags",
276+
type=str,
277+
default="",
278+
help="extra iree-compile options to send for compiling VAE. Only use this for testing bleeding edge flags! Any default options should be added to sd_inference/utils.py",
279+
)
280+
281+
p.add_argument(
282+
"--unet_flags",
283+
type=str,
284+
default="",
285+
help="extra iree-compile options to send for compiling unet. Only use this for testing bleeding edge flags! Any default options should be added to sd_inference/utils.py",
286+
)
287+
288+
289+
args, unknown = p.parse_known_args()

0 commit comments

Comments
 (0)