benchmark: add revsiion when loading pretrianed models for half preci…

…sion benchmark
LambdaLabsML · Oct 3, 2022 · 2469d61 · 2469d61
1 parent 37b5afd
commit 2469d61
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/scripts/benchmark.py b/scripts/benchmark.py
@@ -26,6 +26,7 @@ def get_inference_pipeline(precision, backend):
     if backend == "pytorch":
         pipe = StableDiffusionPipeline.from_pretrained(
             "CompVis/stable-diffusion-v1-4",
+            revision="main" if precision == "single" else "fp16",
             use_auth_token=os.environ["ACCESS_TOKEN"],
             torch_dtype=torch.float32 if precision == "single" else torch.float16,
         )