Function eval_stable_diffusion

examples/mlperf/model_eval.py:274–497

()

  return x

def eval_stable_diffusion():
  import csv, PIL, sys
  from tqdm import tqdm
  from examples.mlperf.initializers import init_stable_diffusion, gelu_erf
  from examples.stable_diffusion import AutoencoderKL
  from extra.models.unet import UNetModel
  from tinygrad.nn.state import load_state_dict, torch_load
  from tinygrad.helpers import BEAM
  from extra.models import clip
  from extra.models.clip import FrozenOpenClipEmbedder
  from extra.models.clip import OpenClipEncoder
  from extra.models.inception import FidInceptionV3

  config = {}
  GPUS               = config["GPUS"]                   = [f"{Device.DEFAULT}:{i}" for i in range(getenv("GPUS", 1))]
  for x in GPUS: Device[x]
  print(f"running eval on {GPUS}")
  seed               = config["seed"]                   = getenv("SEED", 12345)
  CKPTDIR            = config["CKPTDIR"]                = Path(getenv("CKPTDIR", "./checkpoints"))
  DATADIR            = config["DATADIR"]                = Path(getenv("DATADIR", "./datasets"))
  CONTEXT_BS         = config["CONTEXT_BS"]             = getenv("CONTEXT_BS", 1 * len(GPUS))
  DENOISE_BS         = config["DENOISE_BS"]             = getenv("DENOISE_BS", 1 * len(GPUS))
  DECODE_BS          = config["DECODE_BS"]              = getenv("DECODE_BS", 1 * len(GPUS))
  INCEPTION_BS       = config["INCEPTION_BS"]           = getenv("INCEPTION_BS", 1 * len(GPUS))
  CLIP_BS            = config["CLIP_BS"]                = getenv("CLIP_BS", 1 * len(GPUS))
  EVAL_CKPT_DIR      = config["EVAL_CKPT_DIR"]          = getenv("EVAL_CKPT_DIR", "")
  STOP_IF_CONVERGED  = config["STOP_IF_CONVERGED"]      = getenv("STOP_IF_CONVERGED", 0)

  if (WANDB := getenv("WANDB", "")):
    import wandb
    wandb.init(config=config, project="MLPerf-Stable-Diffusion")

  assert EVAL_CKPT_DIR != "", "provide a directory with checkpoints to be evaluated"
  print(f"running eval on checkpoints in {EVAL_CKPT_DIR}\nSEED={seed}")
  eval_queue:list[tuple[int, Path]] = []
  for p in Path(EVAL_CKPT_DIR).iterdir():
    if p.name.endswith(".safetensors"):
      ckpt_iteration = p.name.split(".safetensors")[0]
      assert ckpt_iteration.isdigit(), f"invalid checkpoint name: {p.name}, expected <digits>.safetensors"
      eval_queue.append((int(ckpt_iteration), p))
  assert len(eval_queue), f'no files ending with ".safetensors" were found in {EVAL_CKPT_DIR}'
  print(sorted(eval_queue, reverse=True))

  Tensor.manual_seed(seed)  # seed for weight initialization
  model, unet, sqrt_alphas_cumprod, sqrt_one_minus_alphas_cumprod = init_stable_diffusion("v2-mlperf-eval", CKPTDIR / "sd" / "512-base-ema.ckpt", GPUS)

  # load prompts for generating images for validation; 2 MB of data total
  with open(DATADIR / "coco2014" / "val2014_30k.tsv") as f:
    reader = csv.DictReader(f, delimiter="\t")
    eval_inputs:list[dict] = [{"image_id": int(row["image_id"]), "id": int(row["id"]), "caption": row["caption"]} for row in reader]
  assert len(eval_inputs) == 30_000
  # NOTE: the clip weights are the same between model.cond_stage_model and clip_encoder
  eval_timesteps = list(reversed(range(1, 1000, 20)))

  with Context(DEV="CPU"):
    # The choice of alphas_prev[0] = alphas_cumprod[0] seems arbitrary, but it's how the mlperf ref does it:
    #   alphas_prev = np.asarray([alphacums[0]] + alphacums[ddim_timesteps[:-1]].tolist())
    eval_alphas_prev = model.alphas_cumprod[0:1].cat(model.alphas_cumprod[list(range(1, 1000, 20))[:-1]]).to(GPUS).realize()