Update install.py

pytorch · Jul 13, 2023 · 623ac4b · 623ac4b
1 parent 9efa3ee
commit 623ac4b
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/torchbenchmark/models/stable_diffusion/install.py b/torchbenchmark/models/stable_diffusion/install.py
@@ -6,15 +6,15 @@
 MODEL_NAME = "stabilityai/stable-diffusion-2"
 
 def load_model_checkpoint():
+    from diffusers import StableDiffusionPipeline                   │
     StableDiffusionPipeline.from_pretrained(MODEL_NAME, torch_dtype=torch.float16, safety_checker=None)
 
 def main():
     if not 'HUGGING_FACE_HUB_TOKEN' in os.environ:
         return NotImplementedError("Make sure to set `HUGGINGFACE_HUB_TOKEN` so you can download weights")
     else:
         install_diffusers()
-        from diffusers import StableDiffusionPipeline
         load_model_checkpoint()
 
 if __name__ == "__main__":
-    main()
+    main()