Update pipeline.py

pipeline.py

@@ -18,12 +18,7 @@ import einops
 import PIL.Image
 import numpy as np
 import torch as th
-import torch.nn as nn
-from torchvision import transforms
 
-from diffusers import ModelMixin
-from transformers import AutoModel, AutoConfig, SiglipVisionConfig, Dinov2Config, Dinov2Model
-from transformers import SiglipVisionModel
 from diffusers import DiffusionPipeline
 from diffusers.image_processor import VaeImageProcessor
 from diffusers.models import AutoencoderKL, UNet2DConditionModel
@@ -31,8 +26,6 @@ from diffusers.schedulers import KarrasDiffusionSchedulers
 from diffusers.utils.torch_utils import randn_tensor
 from diffusers.pipelines.pipeline_utils import DiffusionPipeline, ImagePipelineOutput
 
-from diffusers.configuration_utils import ConfigMixin, register_to_config
-# REf: https://github.com/tatp22/multidim-positional-encoding/tree/master
 from analogy_encoder import AnalogyEncoder
 from analogy_projector import AnalogyProjector
 from analogy_input_processor import AnalogyInputProcessor
@@ -259,8 +252,8 @@ class PatternAnalogyTrifuser(DiffusionPipeline):
         The call function to the pipeline for generation.
 
         Args:
-            image (`PIL.Image.Image`, `List[PIL.Image.Image]` or `torch.Tensor`):
-                The image prompt or prompts to guide the image generation.
+            analogy_prompt (`List[Tuple[PIL.Image.Image]]'):
+                The analogy sequence A, A*, B which is our model's prompt for generating B* the analogical pattern satisfying A:A*::B:B*.
             height (`int`, *optional*, defaults to `self.image_unet.config.sample_size * self.vae_scale_factor`):
                 The height in pixels of the generated image.
             width (`int`, *optional*, defaults to `self.image_unet.config.sample_size * self.vae_scale_factor`):
@@ -301,32 +294,77 @@ class PatternAnalogyTrifuser(DiffusionPipeline):
         Examples:
 
         ```py
-        >>> from diffusers import VersatileDiffusionImageVariationPipeline
-        >>> import torch
-        >>> import requests
-        >>> from io import BytesIO
-        >>> from PIL import Image
-
-        >>> # let's download an initial image
-        >>> url = "https://huggingface.co/datasets/diffusers/images/resolve/main/benz.jpg"
-
-        >>> response = requests.get(url)
-        >>> image = Image.open(BytesIO(response.content)).convert("RGB")
-
-        >>> pipe = VersatileDiffusionImageVariationPipeline.from_pretrained(
-        ...     "shi-labs/versatile-diffusion", torch_dtype=torch.float16
-        ... )
-        >>> pipe = pipe.to("cuda")
-
-        >>> generator = torch.Generator(device="cuda").manual_seed(0)
-        >>> image = pipe(image, generator=generator).images[0]
-        >>> image.save("./car_variation.png")
+        import requests
+        import torch as th
+        from PIL import Image
+        from io import BytesIO
+        import matplotlib.pyplot as plt
+        from PIL import Image, ImageOps
+        from diffusers import DiffusionPipeline
+
+        SEED = 1729
+        DEVICE = th.device("cuda")
+        DTYPE = th.float16
+        FIG_K = 3
+        EXAMPLE_ID = 0
+
+        # Now we need to do the trick
+        pretrained_path = "bardofcodes/pattern_analogies"
+        new_pipe = DiffusionPipeline.from_pretrained(
+            pretrained_path, 
+            custom_pipeline=pretrained_path, 
+            trust_remote_code=True
+        )
+
+        img_urls = [
+            f"https://huggingface.co/bardofcodes/pattern_analogies/resolve/main/examples/{EXAMPLE_ID}_a.png",
+            f"https://huggingface.co/bardofcodes/pattern_analogies/resolve/main/examples/{EXAMPLE_ID}_a_star.png",
+            f"https://huggingface.co/bardofcodes/pattern_analogies/resolve/main/examples/{EXAMPLE_ID}_b.png",
+        ]
+        images = []
+        for url in img_urls:
+            response = requests.get(url)
+            image = Image.open(BytesIO(response.content)).convert("RGB")
+            images.append(image)
+
+        pipe_input = [tuple(images)]
+
+        pipe = new_pipe.to(DEVICE, DTYPE)
+        var_images = pipe(pipe_input, num_inference_steps=50, num_images_per_prompt=3,).images
+
+        plt.figure(figsize=(3*FIG_K, 2*FIG_K))
+        plt.axis('off')
+        plt.legend(framealpha=1)
+        plt.rcParams['legend.fontsize'] = 'large'
+        for i in range(6):
+            if i == 0:
+                plt.subplot(2, 3, i+1)
+                val_image = img1
+                label_str = "A"
+            elif i == 1:
+                plt.subplot(2, 3, i+1)
+                val_image = alt_img1
+                label_str = "A*"
+            elif i == 2:
+                plt.subplot(2, 3, i+1)
+                val_image = img2
+                label_str = "Target"
+            else:
+                plt.subplot(2, 3,i + 1)
+                val_image = var_images[i-3]
+                label_str = f"Variation {i-2}"
+            
+            val_image = ImageOps.expand(val_image,border=2,fill='black')
+            plt.imshow(val_image)
+            plt.scatter([], [], c="r", label=label_str)
+            plt.legend(loc="lower right")
+            plt.axis('off')
+        plt.subplots_adjust(wspace=0.01, hspace=0.01)
         ```
 
         Returns:
-            [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] or `tuple`:
-                If `return_dict` is `True`, [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] is returned,
-                otherwise a `tuple` is returned where the first element is a list with the generated images.
+            [`~ImagePipelineOutput`] or `tuple`
+            The generated image(s) as a [`~ImagePipelineOutput`] or a tuple of images.
         """
 
         # 1. Check inputs. Raise error if not correct