copy from diffusers

latent_consistency_controlnet.py

@@ -25,7 +25,6 @@ from transformers import CLIPImageProcessor, CLIPTextModel, CLIPTokenizer
 
 from diffusers import (
     AutoencoderKL,
-    AutoencoderTiny,
     ConfigMixin,
     DiffusionPipeline,
     SchedulerMixin,
@@ -50,6 +49,17 @@ import PIL.Image
 
 logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
 
+
+# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents
+def retrieve_latents(encoder_output, generator):
+    if hasattr(encoder_output, "latent_dist"):
+        return encoder_output.latent_dist.sample(generator)
+    elif hasattr(encoder_output, "latents"):
+        return encoder_output.latents
+    else:
+        raise AttributeError("Could not access latents of provided encoder_output")
+
+
 class LatentConsistencyModelPipeline_controlnet(DiffusionPipeline):
     _optional_components = ["scheduler"]
 
@@ -276,22 +286,17 @@ class LatentConsistencyModelPipeline_controlnet(DiffusionPipeline):
                 )
 
             elif isinstance(generator, list):
-                if isinstance(self.vae, AutoencoderTiny):
-                    init_latents = [
-                        self.vae.encode(image[i : i + 1]).latents
-                        for i in range(batch_size)
-                    ]
-                else:
-                    init_latents = [
-                        self.vae.encode(image[i : i + 1]).latent_dist.sample(generator[i])
-                        for i in range(batch_size)
-                    ]
+                init_latents = [
+                    retrieve_latents(
+                        self.vae.encode(image[i : i + 1]), generator=generator[i]
+                    )
+                    for i in range(batch_size)
+                ]
                 init_latents = torch.cat(init_latents, dim=0)
             else:
-                if isinstance(self.vae, AutoencoderTiny):
-                    init_latents = self.vae.encode(image).latents
-                else:
-                    init_latents = self.vae.encode(image).latent_dist.sample(generator)
+                init_latents = retrieve_latents(
+                    self.vae.encode(image), generator=generator
+                )
 
             init_latents = self.vae.config.scaling_factor * init_latents