Repository: apple/ml-stable-diffusion
File: python_coreml_stable_diffusion/torch2coreml.py (Python)

Code before:
        report_correctness(baseline_out, coreml_out,
                           "vae_decoder baseline PyTorch to baseline CoreML")

    del traced_vae_decoder, pipe.vae.decoder, coreml_vae_decoder
    gc.collect()


def convert_vae_encoder(pipe, args):
    """ Converts the VAE Encoder component of Stable Diffusion
    """
    out_path = _get_out_path(args, "vae_encoder")
    if os.path.exists(out_path):
        logger.info(
            f"`vae_encoder` already exists at {out_path}, skipping conversion."
        )
        return

    if not hasattr(pipe, "unet"):
        raise RuntimeError(
            "convert_unet() deletes pipe.unet to save RAM. "
            "Please use convert_vae_encoder() before convert_unet()")

    sample_shape = (
        1,  # B
        3,  # C (RGB range from -1 to 1)
        args.latent_h or pipe.unet.config.sample_size * 8,  # H
        args.latent_w or pipe.unet.config.sample_size * 8,  # w
    )
    
    noise_shape = (
        1,  # B
        4,  # C
        pipe.unet.config.sample_size,  # H
        pipe.unet.config.sample_size,  # w
    )

    float_value_shape = (
        1,
        1,
    )

    sqrtAlphasCumprodTorchShape = torch.tensor([[0.2,]])
    sqrtOneMinusAlphasCumprodTorchShape = torch.tensor([[0.8,]])

    sample_vae_encoder_inputs = {
        "sample": torch.rand(*sample_shape, dtype=torch.float16),
        "diagonalNoise": torch.rand(*noise_shape, dtype=torch.float16),
        "noise": torch.rand(*noise_shape, dtype=torch.float16),
        "sqrtAlphasCumprod": torch.rand(*float_value_shape, dtype=torch.float16),
        "sqrtOneMinusAlphasCumprod": torch.rand(*float_value_shape, dtype=torch.float16),
    }

Review comment (bug):
This is the shape logic for the `vae_decoder`'s input which would match `vae_encoder` output. If `--latent-h` or `--latent-w` were to be specified, the model will have incorrect input shape for `sample`. The multiplier 8 needs to apply to the output of the `or` operation.

Code after:

    del traced_vae_decoder, pipe.vae.decoder, coreml_vae_decoder
    gc.collect()


def convert_vae_encoder(pipe, args):
    """ Converts the VAE Encoder component of Stable Diffusion
    """
    out_path = _get_out_path(args, "vae_encoder")
    if os.path.exists(out_path):
        logger.info(
            f"`vae_encoder` already exists at {out_path}, skipping conversion."
        )
        return

    if not hasattr(pipe, "unet"):
        raise RuntimeError(
            "convert_unet() deletes pipe.unet to save RAM. "
            "Please use convert_vae_encoder() before convert_unet()")

    sample_shape = (
        1,  # B
        3,  # C (RGB range from -1 to 1)
        (args.latent_h or pipe.unet.config.sample_size) * 8,  # H
        (args.latent_w or pipe.unet.config.sample_size) * 8,  # w
    )
    
    noise_shape = (
        1,  # B
        4,  # C
        pipe.unet.config.sample_size,  # H
        pipe.unet.config.sample_size,  # w
    )

    float_value_shape = (
        1,
        1,
    )

    sqrt_alphas_cumprod_torch_shape = torch.tensor([[0.2,]])
    sqrt_one_minus_alphas_cumprod_torch_shape = torch.tensor([[0.8,]])

    sample_vae_encoder_inputs = {
        "sample": torch.rand(*sample_shape, dtype=torch.float16),
        "diagonal_noise": torch.rand(*noise_shape, dtype=torch.float16),
        "noise": torch.rand(*noise_shape, dtype=torch.float16),
        "sqrt_alphas_cumprod": torch.rand(*float_value_shape, dtype=torch.float16),
        "sqrt_one_minus_alphas_cumprod": torch.rand(*float_value_shape, dtype=torch.float16),
    }

    class VAEEncoder(nn.Module):