sayakpaul

"""
Examples:
    (1) python benchmark_distilled_sd.py --pipeline_id CompVis/stable-diffusion-v1-4
    (2) python benchmark_distilled_sd.py --pipeline_id CompVis/stable-diffusion-v1-4 --vae_path sayakpaul/taesd-diffusers
    (3) python benchmark_distilled_sd.py --pipeline_id nota-ai/bk-sdm-small
    (4) python benchmark_distilled_sd.py --pipeline_id nota-ai/bk-sdm-small --vae_path sayakpaul/taesd-diffusers 
"""

import argparse
import time

sayakpaul

"""
Examples:
    (1) python benchmark_controlnet_sdxl.py --controlnet_id diffusers/controlnet-depth-sdxl-1.0
    (2) python benchmark_controlnet_sdxl.py --controlnet_id diffusers/controlnet-depth-sdxl-1.0-small
    (3) python benchmark_controlnet_sdxl.py --controlnet_id diffusers/controlnet-depth-sdxl-1.0-mid
"""

import argparse
import time

sayakpaul

sayakpaul

sayakpaul

sayakpaul

# To run the script, first install `diffusers` like so:
# pip install git+https://github.com/huggingface/diffusers@temp/lora-debug
# Then grab the LoRA file: 
# wget https://huggingface.co/datasets/sayakpaul/sample-datasets/resolve/main/vividWatercolors_10.safetensors


import torch

from diffusers import StableDiffusionPipeline, DPMSolverMultistepScheduler

sayakpaul

from imgaug import augmenters as iaa
import imgaug as ia
ia.seed(4)

import tensorflow as tf
tf.random.set_seed(666)

aug = iaa.RandAugment(n=2, m=9)
BATCH_SIZE = 224

sayakpaul

sayakpaul

sayakpaul

# Adapted from https://github.com/borisdayma/dalle-mini/blob/main/tools/train/train.py

import random
from dataclasses import dataclass, field
from functools import partial

import jax
import jax.numpy as jnp
import numpy as np
from datasets import Dataset, load_dataset