Fix diffusion (#267)

mila-iqia · Sep 5, 2024 · 2893f94 · 2893f94
1 parent 6cb2c92
commit 2893f94
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 2 deletions.
diff --git a/benchmarks/diffusion/benchfile.py b/benchmarks/diffusion/benchfile.py
@@ -1,3 +1,4 @@
+import os
 from milabench.pack import Package
 
 from milabench.commands import AccelerateAllNodes
@@ -26,9 +27,12 @@ async def prepare(self):
     def build_run_plan(self):
         from milabench.commands import PackCommand
 
+        if "HF_TOKEN" in os.environ or "MILABENCH_HF_TOKEN" in os.environ:
+            os.environ["HF_TOKEN"] = os.environ.get("HF_TOKEN", os.environ["MILABENCH_HF_TOKEN"])
+
         main = self.dirs.code / self.main_script
         plan = PackCommand(self, *self.argv, lazy=True)
-        
+
         if False:
             plan = VoirCommand(plan, cwd=main.parent)
 

diff --git a/benchmarks/diffusion/main.py b/benchmarks/diffusion/main.py
@@ -21,7 +21,7 @@
 
 @dataclass
 class Arguments:
-    model: str = "runwayml/stable-diffusion-v1-5"
+    model: str = "benjamin-paine/stable-diffusion-v1-5"
     dataset: str = "lambdalabs/naruto-blip-captions"
     batch_size: int = 16
     num_workers: int = 8