Stable Diffusion XL

torchbenchmark/canary_models/stable_diffusion_xl/__init__.py

@@ -0,0 +1,50 @@
+"""
+Stable Diffusion XL model
+It requires users to specify "HUGGINGFACE_AUTH_TOKEN" in environment variable
+to authorize login and agree HuggingFace terms and conditions.
+"""
+from torchbenchmark.tasks import COMPUTER_VISION
+from torchbenchmark.util.model import BenchmarkModel
+from torchbenchmark.util.framework.huggingface.model_factory import HuggingFaceAuthMixin
+
+import torch
+from diffusers import DiffusionPipeline
+
+
+class Model(BenchmarkModel, HuggingFaceAuthMixin):
+    task = COMPUTER_VISION.GENERATION
+
+    DEFAULT_TRAIN_BSIZE = 1
+    DEFAULT_EVAL_BSIZE = 1
+    ALLOW_CUSTOMIZE_BSIZE = False
+    # Default eval precision on CUDA device is fp16
+    DEFAULT_EVAL_CUDA_PRECISION = "fp16"
+
+    def __init__(self, test, device, batch_size=None, extra_args=[]):
+        HuggingFaceAuthMixin.__init__(self)
+        super().__init__(test=test, device=device,
+                         batch_size=batch_size, extra_args=extra_args)
+        model_id = "stabilityai/stable-diffusion-xl-base-1.0"
+        self.pipe = DiffusionPipeline.from_pretrained(model_id).to(device)
+        random_input = torch.randn(1, 4, 256, 256).to(device)
+        timestep = torch.tensor([1.0]).to(device)
+        encoder_hidden_states = torch.randn(1, 1, 2048).to(device)
+        added_cond_kwargs = {
+            "text_embeds": torch.randn(1, 2560).to(device),  # Example tensor, adjust shape as needed
+            "time_ids": torch.tensor([1]).to(device)  # Replace 'some_value' with the appropriate value or tensor shape for time_ids
+        }
+
+        self.args_tuple = (random_input, timestep, encoder_hidden_states, added_cond_kwargs)
+
+    def enable_fp16_half(self):
+        pass
+
+    def get_module(self):
+        self.pipe.unet, self.args_tuple
+
+    def train(self):
+        raise NotImplementedError("Train is not implemented for the stable diffusion XL model.")
+
+    def eval(self):
+        image = self.pipe.unet(*self.args_tuple)
+        return (image, )

torchbenchmark/canary_models/stable_diffusion_xl/install.py

@@ -0,0 +1,17 @@
+from torchbenchmark.util.framework.diffusers import install_diffusers
+from torchbenchmark.util.framework.huggingface.model_factory import HuggingFaceAuthMixin
+import torch
+import os
+import warnings
+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)
+
+if __name__ == "__main__":
+    if not 'HUGGING_FACE_HUB_TOKEN' in os.environ:
+        warnings.warn("Make sure to set `HUGGINGFACE_HUB_TOKEN` so you can download weights")
+    else:
+        install_diffusers()
+        load_model_checkpoint()

torchbenchmark/canary_models/stable_diffusion_xl/metadata.yaml

@@ -0,0 +1,10 @@
+devices:
+  NVIDIA A100-SXM4-40GB:
+    eval_batch_size: 32
+eval_benchmark: false
+eval_deterministic: false
+eval_nograd: true
+train_benchmark: false
+train_deterministic: false
+not_implemented:
+- device: cpu

torchbenchmark/util/framework/diffusers/requirements.txt

@@ -1 +1 @@
-git+https://github.com/huggingface/diffusers.git@56958e1
+diffusers==0.20.2