Update diffusion model checkpoint to be milabench checkpoint (#379)

* Update diffusion model checkpoint to be milabench checkpoint

* Update main.py

* Update prepare.py

Author: Delaunay

benchmarks/diffusion/main.py

@@ -21,7 +21,7 @@
 
 @dataclass
 class Arguments:
-    model: str = "stabilityai/stable-diffusion-2"
+    model: str = "Milabench/stable-diffusion-2"
     dataset: str = "lambdalabs/naruto-blip-captions"
     batch_size: int = 16
     num_workers: int = 8
@@ -252,4 +252,4 @@ def main():
 
 
 if __name__ == "__main__":
-    main()
+    main()

benchmarks/diffusion/prepare.py

@@ -6,7 +6,7 @@
 
 @dataclass
 class TrainingConfig:
-    model: str = "stabilityai/stable-diffusion-2"
+    model: str = "Milabench/stable-diffusion-2"
     dataset: str = "lambdalabs/naruto-blip-captions"
     revision: str = None
     variant: str = None

config/base.yaml

@@ -420,6 +420,7 @@ _diffusion:
     method: per_gpu
 
   argv:
+    --model: "Milabench/stable-diffusion-2"
     --num_epochs: 5
     --batch_size: "auto_batch(32)"
     --num_workers: "auto({n_worker}, 8)"