Skip to content

Commit

Permalink
add stable diffusion from huggingface (#254)
Browse files Browse the repository at this point in the history
  • Loading branch information
MarcelWilnicki authored Aug 29, 2024
1 parent 4401f4e commit de4c6e5
Show file tree
Hide file tree
Showing 3 changed files with 58 additions and 2 deletions.
2 changes: 1 addition & 1 deletion benchmark.py
Original file line number Diff line number Diff line change
Expand Up @@ -9,7 +9,7 @@
import urllib.request
from pathlib import Path

LATEST_VERSION = "2.2.0a0+git1bef725"
LATEST_VERSION = "2.2.0a0+git6032a25"
SYSTEMS = {
"Altra": {
"ResNet-50 v1.5": "https://ampereaimodelzoo.s3.eu-central-1.amazonaws.com/lookups_aml/q80_30%40ampere_pytorch_1.10.0%40resnet_50_v1.5.json", # noqa
Expand Down
56 changes: 56 additions & 0 deletions text_to_image/stable_diffusion/run_hf.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,56 @@
# SPDX-License-Identifier: Apache-2.0
# Copyright (c) 2024, Ampere Computing LLC
try:
from utils import misc # noqa
except ModuleNotFoundError:
import os
import sys
filename = "set_env_variables.sh"
directory = os.path.realpath(__file__).split("/")[:-1]
for idx in range(1, len(directory) - 1):
subdir = "/".join(directory[:-idx])
if filename in os.listdir(subdir):
print(f"\nPlease run \033[91m'source {os.path.join(subdir, filename)}'\033[0m first.")
break
else:
print(f"\n\033[91mFAIL: Couldn't find {filename}, are you running this script as part of Ampere Model Library?"
f"\033[0m")
sys.exit(1)


def run_pytorch_bf16(model_name, steps, batch_size, num_runs, timeout, **kwargs):
import torch._dynamo
from diffusers import DiffusionPipeline
torch._dynamo.config.suppress_errors = True

from utils.benchmark import run_model
from utils.pytorch import apply_compile
from utils.pytorch import PyTorchRunnerV2
from utils.text_to_image.stable_diffusion import StableDiffusion

model = DiffusionPipeline.from_pretrained(model_name,
use_safetensors=True,
torch_dtype=torch.bfloat16).to("cpu")

model.unet = apply_compile(model.unet)

def single_pass_pytorch(_runner, _stablediffusion):
prompts = [_stablediffusion.get_input() for _ in range(batch_size)]
x_samples = _runner.run(batch_size * steps, prompt=prompts, num_inference_steps=steps)
_stablediffusion.submit_count(batch_size, x_samples)

runner = PyTorchRunnerV2(model)
stablediffusion = StableDiffusion()
return run_model(single_pass_pytorch, runner, stablediffusion, batch_size, num_runs, timeout)


if __name__ == "__main__":
from utils.helpers import DefaultArgParser

stablediffusion_variants = ["stabilityai/stable-diffusion-xl-base-1.0"]
parser = DefaultArgParser(["pytorch"])
parser.require_model_name(stablediffusion_variants)
parser.ask_for_batch_size()
parser.add_argument("--steps", type=int, default=25, help="steps through which the model processes the input")

run_pytorch_bf16(**vars(parser.parse()))
2 changes: 1 addition & 1 deletion utils/text_to_image/stable_diffusion.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@
# Copyright (c) 2024, Ampere Computing LLC
from random import randint, seed
from utils.helpers import Dataset
seed(42)


class StableDiffusion(Dataset):
Expand All @@ -14,7 +15,6 @@ def get_input(self):
adjectives = ["big", "small", "thin", "wide", "blonde", "pale"]
nouns = ["dog", "cat", "horse", "astronaut", "human", "robot"]
actions = ["sings", "rides a triceratop", "rides a horse", "eats a burger", "washes clothes", "looks at hands"]
seed(42)

a = adjectives[randint(0, len(adjectives) - 1)] + " "
b = nouns[randint(0, len(nouns) - 1)] + " " + actions[randint(0, len(actions) - 1)]
Expand Down

0 comments on commit de4c6e5

Please sign in to comment.