Merge pull request #287 from TobyRoseman/minor-fix

Fix pipeline.py for Python 3.8

Author: TobyRoseman

python_coreml_stable_diffusion/pipeline.py

@@ -5,6 +5,7 @@
 
 import argparse
 
+from diffusers import StableDiffusionPipeline, StableDiffusionXLPipeline
 from diffusers.pipelines.pipeline_utils import DiffusionPipeline
 from diffusers.pipelines.stable_diffusion import StableDiffusionPipelineOutput
 from diffusers.schedulers import (
@@ -517,7 +518,7 @@ def __call__(
                 control_net_additional_residuals = {}
 
             # predict the noise residual
-            unet_additional_kwargs = unet_additional_kwargs | control_net_additional_residuals
+            unet_additional_kwargs.update(control_net_additional_residuals)
 
             noise_pred = self.unet(
                 sample=latent_model_input.astype(np.float16),
@@ -697,7 +698,6 @@ def main(args):
 
     logger.info("Initializing PyTorch pipe for reference configuration")
 
-    from diffusers import StableDiffusionPipeline, StableDiffusionXLPipeline
     SDP = StableDiffusionXLPipeline if 'xl' in args.model_version else StableDiffusionPipeline
 
     pytorch_pipe = SDP.from_pretrained(

python_coreml_stable_diffusion/torch2coreml.py

@@ -480,7 +480,7 @@ def convert_vae_decoder(pipe, args):
         1,  # B
         pipe.vae.config.latent_channels,  # C
         args.latent_h or pipe.unet.config.sample_size,  # H
-        args.latent_w or pipe.unet.config.sample_size,  # w
+        args.latent_w or pipe.unet.config.sample_size,  # W
     )
 
     if args.custom_vae_version is None and args.xl_version: