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

Code before:
            self.encoder = pipe.vae.encoder

        # Because CoreMLTools does not support the torch.randn op, we pass in both
        # the diagonal Noise for the `DiagonalGaussianDistribution` operation and
        # the noise tensor combined with precalculated `sqrtAlphasCumprod` and `sqrtOneMinusAlphasCumprod`
        # for faster computation.
        def forward(self, sample, diagonalNoise, noise, sqrtAlphasCumprod, sqrtOneMinusAlphasCumprod):
            h = self.encoder(sample)
            moments = self.quant_conv(h)
            diagonalNoise = diagonalNoise.to(sample.device)
            posterior = CoreMLDiagonalGaussianDistribution(moments, diagonalNoise)
            posteriorSample = posterior.sample()
            
            # Add the scaling operation and the latent noise for faster computation
            init_latents = 0.18215 * posteriorSample
            result = self.add_noise(init_latents, noise, sqrtAlphasCumprod, sqrtOneMinusAlphasCumprod)
            return result
        
        def add_noise(
            self,
            original_samples: torch.FloatTensor,
            noise: torch.FloatTensor,
            sqrtAlphasCumprod: torch.FloatTensor,
            sqrtOneMinusAlphasCumprod: torch.FloatTensor
        ) -> torch.FloatTensor:
            noise = noise.to(original_samples.device)
            sqrtAlphasCumprod = sqrtAlphasCumprod.to(original_samples.device)
            sqrtOneMinusAlphasCumprod = sqrtOneMinusAlphasCumprod.to(original_samples.device)
            noisy_samples = sqrtAlphasCumprod * original_samples + sqrtOneMinusAlphasCumprod * noise
            return noisy_samples
         

    baseline_encoder = VAEEncoder().eval()

    # No optimization needed for the VAE Encoder as it is a pure ConvNet
    traced_vae_encoder = torch.jit.trace(
        baseline_encoder, (
            sample_vae_encoder_inputs["sample"].to(torch.float32),
            sample_vae_encoder_inputs["diagonalNoise"].to(torch.float32),
            sample_vae_encoder_inputs["noise"].to(torch.float32),
            sqrtAlphasCumprodTorchShape.to(torch.float32),
            sqrtOneMinusAlphasCumprodTorchShape.to(torch.float32)
        ))

    modify_coremltools_torch_frontend_badbmm()
    coreml_vae_encoder, out_path = _convert_to_coreml(
        "vae_encoder", traced_vae_encoder, sample_vae_encoder_inputs,
        ["latent_dist"], args)

    # Set model metadata
    coreml_vae_encoder.author = f"Please refer to the Model Card available at huggingface.co/{args.model_version}"

Review comment (refactor):
Could you simplify `add_noise` by removing all three of the device placement lines? Exported models will be agnostic to these placements and it will help improve readability

Code after:
        # Because CoreMLTools does not support the torch.randn op, we pass in both
        # the diagonal Noise for the `DiagonalGaussianDistribution` operation and
        # the noise tensor combined with precalculated `sqrt_alphas_cumprod` and `sqrt_one_minus_alphas_cumprod`
        # for faster computation.
        def forward(self, sample, diagonal_noise, noise, sqrt_alphas_cumprod, sqrt_one_minus_alphas_cumprod):
            h = self.encoder(sample)
            moments = self.quant_conv(h)
            posterior = CoreMLDiagonalGaussianDistribution(moments, diagonal_noise)
            posteriorSample = posterior.sample()
            
            # Add the scaling operation and the latent noise for faster computation
            init_latents = 0.18215 * posteriorSample
            result = self.add_noise(init_latents, noise, sqrt_alphas_cumprod, sqrt_one_minus_alphas_cumprod)
            return result
        
        def add_noise(
            self,
            original_samples: torch.FloatTensor,
            noise: torch.FloatTensor,
            sqrt_alphas_cumprod: torch.FloatTensor,
            sqrt_one_minus_alphas_cumprod: torch.FloatTensor
        ) -> torch.FloatTensor:
            noisy_samples = sqrt_alphas_cumprod * original_samples + sqrt_one_minus_alphas_cumprod * noise
            return noisy_samples
         

    baseline_encoder = VAEEncoder().eval()

    # No optimization needed for the VAE Encoder as it is a pure ConvNet
    traced_vae_encoder = torch.jit.trace(
        baseline_encoder, (
            sample_vae_encoder_inputs["sample"].to(torch.float32),
            sample_vae_encoder_inputs["diagonal_noise"].to(torch.float32),
            sample_vae_encoder_inputs["noise"].to(torch.float32),
            sqrt_alphas_cumprod_torch_shape.to(torch.float32),
            sqrt_one_minus_alphas_cumprod_torch_shape.to(torch.float32)
        ))

    modify_coremltools_torch_frontend_badbmm()
    coreml_vae_encoder, out_path = _convert_to_coreml(
        "vae_encoder", traced_vae_encoder, sample_vae_encoder_inputs,
        ["latent_dist"], args)

    # Set model metadata
    coreml_vae_encoder.author = f"Please refer to the Model Card available at huggingface.co/{args.model_version}"
    coreml_vae_encoder.license = "OpenRAIL (https://huggingface.co/spaces/CompVis/stable-diffusion-license)"
    coreml_vae_encoder.version = args.model_version
    coreml_vae_encoder.short_description = \
        "Stable Diffusion generates images conditioned on text and/or other images as input through the diffusion process. " \
        "Please refer to https://arxiv.org/abs/2112.10752 for details."