Update ip_adapter/ip_adapter.py

ip_adapter/ip_adapter.py

@@ -185,18 +185,25 @@ class IPAdapterXL(IPAdapter):
     
     def generate(
         self,
-        pil_image,
+        pil_image_1,
         pil_image_2=None,
+        pil_image_3=None,
+        pil_image_4=None,
+        pil_image_5=None,
         prompt=None,
         negative_prompt=None,
-        scale=1.0,
+        scale_1=1.0,
+        scale_2=1.0,
+        scale_3=1.0,
+        scale_4=1.0,
+        scale_5=1.0,
         num_samples=4,
         seed=-1,
         num_inference_steps=30,
         guidance_scale=7.5,
         **kwargs,
     ):
-        self.set_scale(scale)
+        self.set_scale(scale_1)
         
         if isinstance(pil_image, Image.Image):
             num_prompts = 1
@@ -212,24 +219,48 @@ class IPAdapterXL(IPAdapter):
             prompt = [prompt] * num_prompts
         if not isinstance(negative_prompt, List):
             negative_prompt = [negative_prompt] * num_prompts
+            
+        image_prompt_embeds_list = []
+        uncond_image_prompt_embeds_list = []
+        
+        image_prompt_embeds_1, uncond_image_prompt_embeds_1 = self.get_image_embeds(pil_image)
+        image_prompt_embeds_list.append(image_prompt_embeds_1)
+        uncond_image_prompt_embeds_list.append(uncond_image_prompt_embeds_1)
+        
         if pil_image_2 != None:
             print('Using secondary image.')
-            image_prompt_embeds_1, uncond_image_prompt_embeds = self.get_image_embeds(pil_image)
+            self.set_scale(scale_2)
             image_prompt_embeds_2, uncond_image_prompt_embeds_2 = self.get_image_embeds(pil_image_2)
-            image_prompt_embeds = torch.cat((image_prompt_embeds_1,image_prompt_embeds_2)).mean(dim=0).unsqueeze(0)  
-            bs_embed, seq_len, _ = image_prompt_embeds.shape
-            image_prompt_embeds = image_prompt_embeds.repeat(1, num_samples, 1)
-            image_prompt_embeds = image_prompt_embeds.view(bs_embed * num_samples, seq_len, -1)
-            uncond_image_prompt_embeds = uncond_image_prompt_embeds.repeat(1, num_samples, 1)
-            uncond_image_prompt_embeds = uncond_image_prompt_embeds.view(bs_embed * num_samples, seq_len, -1)
-        else:
-            image_prompt_embeds, uncond_image_prompt_embeds = self.get_image_embeds(pil_image)
-            bs_embed, seq_len, _ = image_prompt_embeds.shape
-            image_prompt_embeds = image_prompt_embeds.repeat(1, num_samples, 1)
-            image_prompt_embeds = image_prompt_embeds.view(bs_embed * num_samples, seq_len, -1)
-            uncond_image_prompt_embeds = uncond_image_prompt_embeds.repeat(1, num_samples, 1)
-            uncond_image_prompt_embeds = uncond_image_prompt_embeds.view(bs_embed * num_samples, seq_len, -1)
-        
+            image_prompt_embeds_list.append(image_prompt_embeds_2)
+            uncond_image_prompt_embeds_list.append(uncond_image_prompt_embeds_2)
+        if pil_image_3 != None:
+            print('Using secondary image.')
+            self.set_scale(scale_3)
+            image_prompt_embeds_3, uncond_image_prompt_embeds_3 = self.get_image_embeds(pil_image_3)
+            image_prompt_embeds_list.append(image_prompt_embeds_3)
+            uncond_image_prompt_embeds_list.append(uncond_image_prompt_embeds_3)
+        if pil_image_4 != None:
+            print('Using secondary image.')
+            self.set_scale(scale_4)
+            image_prompt_embeds_4, uncond_image_prompt_embeds_4 = self.get_image_embeds(pil_image_4)
+            image_prompt_embeds_list.append(image_prompt_embeds_4)
+            uncond_image_prompt_embeds_list.append(uncond_image_prompt_embeds_4)
+        if pil_image_5 != None:
+            print('Using secondary image.')
+            self.set_scale(scale_5)
+            image_prompt_embeds_5, uncond_image_prompt_embeds_5 = self.get_image_embeds(pil_image_5)
+            image_prompt_embeds_list.append(image_prompt_embeds_5)
+            uncond_image_prompt_embeds_list.append(uncond_image_prompt_embeds_5)
+
+        image_prompt_embeds = torch.cat(image_prompt_embeds_list).mean(dim=0).unsqueeze(0)
+        bs_embed, seq_len, _ = image_prompt_embeds.shape
+        image_prompt_embeds = image_prompt_embeds.repeat(1, num_samples, 1)
+        image_prompt_embeds = image_prompt_embeds.view(bs_embed * num_samples, seq_len, -1)
+
+        uncond_image_prompt_embeds = torch.cat(uncond_image_prompt_embeds_list).mean(dim=0).unsqueeze(0)  
+        uncond_image_prompt_embeds = uncond_image_prompt_embeds.repeat(1, num_samples, 1)
+        uncond_image_prompt_embeds = uncond_image_prompt_embeds.view(bs_embed * num_samples, seq_len, -1)
+
         with torch.inference_mode():
             prompt_embeds, negative_prompt_embeds, pooled_prompt_embeds, negative_pooled_prompt_embeds = self.pipe.encode_prompt(
                 prompt, num_images_per_prompt=num_samples, do_classifier_free_guidance=True, negative_prompt=negative_prompt)