Class StableDiffusion

examples/stable_diffusion.py:162–236

                                "num_groups":16, "st_norm_eps":1e-6}

class StableDiffusion:
  def __init__(self, version:str|None=None, pretrained:str|None=None):
    self.alphas_cumprod = get_alphas_cumprod()
    if version != "v2-mlperf-train":
      self.first_stage_model = AutoencoderKL() # only needed for decoding generated latents to images; not needed in mlperf training from preprocessed moments

    if not version:
      self.cond_stage_model = namedtuple("CondStageModel", ["transformer"])(transformer = namedtuple("Transformer", ["text_model"])(text_model = Closed.ClipTextTransformer()))
      unet_init_params = unet_params
    elif version in {"v2-mlperf-train", "v2-mlperf-eval"}:
      unet_init_params = mlperf_params
      clip.gelu = gelu_erf
      self.cond_stage_model = FrozenOpenClipEmbedder(**{"dims": 1024, "n_heads": 16, "layers": 24, "return_pooled": False, "ln_penultimate": True,
                                                        "clip_tokenizer_version": "sd_mlperf_v5_0"})
      unet.Linear, unet.Conv2d, unet.GroupNorm, unet.LayerNorm = AutocastLinear, AutocastConv2d, AutocastGroupNorm, AutocastLayerNorm
      unet.attention, unet.gelu, unet.mixed_precision_dtype = attn_f32_softmax, gelu_erf, dtypes.bfloat16
      if pretrained:
        print("loading text encoder")
        weights: dict[str,Tensor] = {k.replace("cond_stage_model.", "", 1):v for k,v in torch_load(pretrained)["state_dict"].items() if k.startswith("cond_stage_model.")}
        weights["model.attn_mask"] = Tensor.full((77, 77), fill_value=float("-inf")).triu(1)
        load_state_dict(self.cond_stage_model, weights)
        # only the eval model needs the decoder
        if version == "v2-mlperf-eval":
          print("loading image latent encoder")
          weights = {k.replace("first_stage_model.", "", 1):v for k,v in torch_load(pretrained)["state_dict"].items() if k.startswith("first_stage_model.")}
          load_state_dict(self.first_stage_model, weights)

    self.model = namedtuple("DiffusionModel", ["diffusion_model"])(diffusion_model = UNetModel(**unet_init_params))
    if version == "v2-mlperf-train":
      # the mlperf reference inits certain weights as zeroes
      for bb in flatten(self.model.diffusion_model.input_blocks) + self.model.diffusion_model.middle_block + flatten(self.model.diffusion_model.output_blocks):
        if isinstance(bb, unet.ResBlock):
          zero_module(bb.out_layers[3])
        elif isinstance(bb, unet.SpatialTransformer):
          zero_module(bb.proj_out)
      zero_module(self.model.diffusion_model.out[2])

  def get_x_prev_and_pred_x0(self, x, e_t, a_t, a_prev):
    temperature = 1
    sigma_t = 0
    sqrt_one_minus_at = (1-a_t).sqrt()
    #print(a_t, a_prev, sigma_t, sqrt_one_minus_at)

    pred_x0 = (x - sqrt_one_minus_at * e_t) / a_t.sqrt()

    # direction pointing to x_t
    dir_xt = (1. - a_prev - sigma_t**2).sqrt() * e_t

    x_prev = a_prev.sqrt() * pred_x0 + dir_xt
    return x_prev, pred_x0

  def get_model_output(self, unconditional_context, context, latent, timestep, unconditional_guidance_scale):
    # put into diffuser
    latents = self.model.diffusion_model(latent.expand(2, *latent.shape[1:]), timestep, unconditional_context.cat(context, dim=0))
    unconditional_latent, latent = latents[0:1], latents[1:2]

    e_t = unconditional_latent + unconditional_guidance_scale * (latent - unconditional_latent)
    return e_t