More configuration + image gen script using cached model

Author: jbschlosser

gen_image.py

@@ -0,0 +1,29 @@
+import random
+import time
+import torch
+from torch.profiler import profile, record_function, ProfilerActivity
+from utils.benchmark_utils import annotate, create_parser
+from utils.pipeline_utils import load_pipeline  # noqa: E402
+
+
+def set_rand_seeds(seed):
+    random.seed(seed)
+    torch.manual_seed(seed)
+
+
+def main(args):
+    pipeline = load_pipeline(args)
+    set_rand_seeds(args.seed)
+
+    image = pipeline(
+        args.prompt, num_inference_steps=args.num_inference_steps, guidance_scale=0.0
+    ).images[0]
+    image.save(args.output_file)
+
+
+if __name__ == "__main__":
+    parser = create_parser()
+    args = parser.parse_args()
+    # use the cached model to minimize latency
+    args.use_cached_model = True
+    main(args)

run_benchmark.py

@@ -6,8 +6,15 @@
 from utils.pipeline_utils import load_pipeline  # noqa: E402
 
 
+def set_rand_seeds(seed):
+    random.seed(seed)
+    torch.manual_seed(seed)
+
+
 def main(args):
+    set_rand_seeds(args.seed)
     pipeline = load_pipeline(args)
+    set_rand_seeds(args.seed)
 
     # warmup
     for _ in range(3):
@@ -52,10 +59,6 @@ def main(args):
 
 
 if __name__ == "__main__":
-    seed = 42
-    random.seed(seed)
-    torch.manual_seed(seed)
-
     parser = create_parser()
     args = parser.parse_args()
     main(args)

utils/benchmark_utils.py

@@ -15,12 +15,16 @@ def create_parser():
                         help="Text prompt")
     parser.add_argument("--cache-dir", type=str, default=os.path.expandvars("$HOME/.cache/flux-fast"),
                         help="Cache directory for storing exported models")
+    parser.add_argument("--use-cached-model", action="store_true",
+                        help="Attempt to use cached model only (don't re-export)")
     parser.add_argument("--device", type=str, choices=["cuda", "cpu"], default="cuda",
                         help="Device to use")
     parser.add_argument("--num_inference_steps", type=int, default=4,
                         help="Number of denoising steps")
     parser.add_argument("--output-file", type=str, default="output.png",
                         help="Output image file path")
+    parser.add_argument("--seed", type=int, default=42,
+                        help="Random seed to use")
     # file path for optional output PyTorch Profiler trace
     parser.add_argument("--trace-file", type=str, default=None,
                         help="Output PyTorch Profiler trace file path")

utils/pipeline_utils.py

@@ -375,8 +375,11 @@ def optimize(pipeline, args):
     if args.compile_export_mode == "compile":
         pipeline = use_compile(pipeline)
     elif args.compile_export_mode == "export_aoti":
-        # NB: Using a cached export + AOTI model is not supported yet
-        pipeline = use_export_aoti(pipeline, cache_dir=args.cache_dir, serialize=True)
+        pipeline = use_export_aoti(
+            pipeline,
+            cache_dir=args.cache_dir,
+            serialize=(not args.use_cached_model),
+        )
     elif args.compile_export_mode == "disabled":
         pass
     else: