remove compute_clip_image_embedding overloads in image_prompt.py + improve docstring slightly

Author: Laurent

src/refiners/foundationals/latent_diffusion/image_prompt.py

@@ -1,5 +1,5 @@
 import math
-from typing import TYPE_CHECKING, Any, Generic, TypeVar, overload
+from typing import TYPE_CHECKING, Any, Generic, TypeVar
 
 import torch
 from jaxtyping import Float
@@ -454,37 +454,28 @@ def set_clip_image_embedding(self, image_embedding: Tensor) -> None:
         """
         self.set_context("ip_adapter", {"clip_image_embedding": image_embedding})
 
-    @overload
-    def compute_clip_image_embedding(self, image_prompt: Tensor, weights: list[float] | None = None) -> Tensor: ...
-
-    @overload
-    def compute_clip_image_embedding(self, image_prompt: Image.Image) -> Tensor: ...
-
-    @overload
-    def compute_clip_image_embedding(
-        self, image_prompt: list[Image.Image], weights: list[float] | None = None
-    ) -> Tensor: ...
-
     def compute_clip_image_embedding(
         self,
         image_prompt: Tensor | Image.Image | list[Image.Image],
         weights: list[float] | None = None,
         concat_batches: bool = True,
     ) -> Tensor:
-        """Compute the CLIP image embedding.
+        """Compute CLIP image embedding(s).
 
         Args:
-            image_prompt: The image prompt to use.
-            weights: The scale to use for the image prompt.
-            concat_batches: Whether to concatenate the batches.
+            image_prompt: The image prompt(s) to use.
+            weights: The scale(s) to use for the image prompt(s).
+            concat_batches: Whether to concatenate the image embeddings along the feature dimension.
 
         Returns:
-            The CLIP image embedding.
+            The CLIP image embedding(s).
         """
         if isinstance(image_prompt, Image.Image):
             image_prompt = self.preprocess_image(image_prompt)
         elif isinstance(image_prompt, list):
-            assert all(isinstance(image, Image.Image) for image in image_prompt)
+            assert all(
+                isinstance(image, Image.Image) for image in image_prompt
+            ), "All elements of `image_prompt` must be of type `Image.Image`"
             image_prompt = torch.cat([self.preprocess_image(image) for image in image_prompt])
 
         negative_embedding, conditional_embedding = self._compute_clip_image_embedding(image_prompt)