Add HF Auth mixin to Stable Diffusion

torchbenchmark/canary_models/stable_diffusion/__init__.py → torchbenchmark/models/stable_diffusion/__init__.py

@@ -10,7 +10,7 @@
 from diffusers import StableDiffusionPipeline, EulerDiscreteScheduler
 
 
-class Model(BenchmarkModel):
+class Model(HuggingFaceAuthMixin, BenchmarkModel):
     task = COMPUTER_VISION.GENERATION
 
     DEFAULT_TRAIN_BSIZE = 1

torchbenchmark/models/stable_diffusion/install.py

@@ -0,0 +1,20 @@
+from torchbenchmark.util.framework.diffusers import install_diffusers
+from torchbenchmark.util.framework.huggingface.model_factory import HuggingFaceAuthMixin
+import torch
+import os
+
+MODEL_NAME = "stabilityai/stable-diffusion-2"
+
+def load_model_checkpoint():
+    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()

torchbenchmark/util/framework/huggingface/model_factory.py

@@ -159,6 +159,11 @@ def eval(self) -> Tuple[torch.Tensor]:
         else:
             return (out["logits"], )
 
+class HuggingFaceAuthMixin:
+    def __init__(self):
+        if not 'HUGGING_FACE_HUB_TOKEN' in os.environ:
+            raise NotImplementedError("Make sure to set `HUGGING_FACE_HUB_TOKEN` so you can download weights")
+
 
 class HuggingFaceGenerationModel(HuggingFaceModel):
     task = NLP.GENERATION

torchbenchmark/util/metadata_utils.py

@@ -11,6 +11,8 @@ def match_item(item_name: str, item_val: str, skip_item: Dict[str, Any]) -> bool
 
 def skip_by_metadata(test: str, device:str, jit: bool, extra_args: List[str], metadata: Dict[str, Any]) -> bool:
     "Check if the test should be skipped based on model metadata."
+    if not metadata:
+        return False
     if not "not_implemented" in metadata:
         return False
     for skip_item in metadata["not_implemented"]:
@@ -20,4 +22,4 @@ def skip_by_metadata(test: str, device:str, jit: bool, extra_args: List[str], me
                 match_item("extra_args", extra_args, skip_item)
         if match:
             return True
-    return False
+    return False