Update pipeline_stable_diffusion_3_ipa_clip.py

pipeline_stable_diffusion_3_ipa_clip.py

@@ -14,6 +14,7 @@
 
 import inspect
 from typing import Any, Callable, Dict, List, Optional, Union
+from PIL import Image
 
 import torch
 import torch.nn as nn
@@ -922,8 +923,9 @@ class StableDiffusion3Pipeline(DiffusionPipeline, SD3LoraLoaderMixin, FromSingle
 
     @torch.inference_mode()
     def encode_clip_image_emb(self, clip_image, device, dtype):
-
-        # clip
+        if isinstance(clip_image, Image.Image):
+            clip_image = [clip_image]
+            # clip
         clip_image_tensor = self.clip_image_processor(images=clip_image, return_tensors="pt").pixel_values
         clip_image_tensor = clip_image_tensor.to(device, dtype=dtype)
         clip_image_embeds = self.image_encoder(clip_image_tensor, output_hidden_states=True).hidden_states[-2]
@@ -965,7 +967,17 @@ class StableDiffusion3Pipeline(DiffusionPipeline, SD3LoraLoaderMixin, FromSingle
 
         # ipa
         clip_image=None,
+        clip_image_2=None,
+        clip_image_3=None,
+        clip_image_4=None,
+        clip_image_5=None,
+        text_scale=1.0,
         ipadapter_scale=1.0,
+        scale_1=1.0,
+        scale_2=1.0,
+        scale_3=1.0,
+        scale_4=1.0,
+        scale_5=1.0,
     ):
         r"""
         Function invoked when calling the pipeline for generation.
@@ -1126,11 +1138,82 @@ class StableDiffusion3Pipeline(DiffusionPipeline, SD3LoraLoaderMixin, FromSingle
         if self.do_classifier_free_guidance:
             prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds], dim=0)
             pooled_prompt_embeds = torch.cat([negative_pooled_prompt_embeds, pooled_prompt_embeds], dim=0)
-
+            
+        prompt_embeds = prompt_embeds * text_scale
+        
+        image_prompt_embeds_list = []
+        
         # 3. prepare clip emb
-        clip_image = clip_image.resize((max(clip_image.size), max(clip_image.size)))
-        clip_image_embeds = self.encode_clip_image_emb(clip_image, device, dtype)
-
+        if clip_image != None:
+            print('Using primary image.')
+            clip_image = clip_image.resize((max(clip_image.size), max(clip_image.size)))
+            #clip_image_embeds_1 = self.encode_clip_image_emb(clip_image, device, dtype)
+            #with torch.no_grad():
+            clip_image_embeds_1 = self.clip_image_processor(images=clip_image, return_tensors="pt").pixel_values
+            print('clip output shape: ', clip_image_embeds_1.shape)
+            clip_image_embeds_1 = clip_image_embeds_1.to(device, dtype=dtype)
+            clip_image_embeds_1 = self.image_encoder(clip_image_embeds_1, output_hidden_states=True).hidden_states[-2]
+            print('encoder output shape: ', clip_image_embeds_1.shape)
+            clip_image_embeds_1 = clip_image_embeds_1 * scale_1
+            image_prompt_embeds_list.append(clip_image_embeds_1)
+        if clip_image_2 != None:
+            print('Using secondary image.')
+            clip_image_2 = clip_image_2.resize((max(clip_image_2.size), max(clip_image_2.size)))
+            #with torch.no_grad():
+            clip_image_embeds_2 = self.clip_image_processor(images=clip_image_2, return_tensors="pt").pixel_values
+            clip_image_embeds_2 = clip_image_embeds_2.to(device, dtype=dtype)
+            clip_image_embeds_2 = self.image_encoder(clip_image_embeds_2, output_hidden_states=True).hidden_states[-2]
+            clip_image_embeds_2 = clip_image_embeds_2 * scale_2
+            image_prompt_embeds_list.append(clip_image_embeds_2)
+        if clip_image_3 != None:
+            print('Using tertiary image.')
+            clip_image_3 = clip_image_3.resize((max(clip_image_3.size), max(clip_image_3.size)))
+            #with torch.no_grad():
+            clip_image_embeds_3 = self.clip_image_processor(images=clip_image_3, return_tensors="pt").pixel_values
+            clip_image_embeds_3 = clip_image_embeds_3.to(device, dtype=dtype)
+            clip_image_embeds_3 = self.image_encoder(clip_image_embeds_3, output_hidden_states=True).hidden_states[-2]
+            clip_image_embeds_3 = clip_image_embeds_3 * scale_3
+            image_prompt_embeds_list.append(clip_image_embeds_3)
+        if clip_image_4 != None:
+            print('Using quaternary image.')
+            clip_image_4 = clip_image_4.resize((max(clip_image_4.size), max(clip_image_4.size)))
+            #with torch.no_grad():
+            clip_image_embeds_4 = self.clip_image_processor(images=clip_image_4, return_tensors="pt").pixel_values
+            clip_image_embeds_4 = clip_image_embeds_4.to(device, dtype=dtype)
+            clip_image_embeds_4 = self.image_encoder(clip_image_embeds_4, output_hidden_states=True).hidden_states[-2]
+            clip_image_embeds_4 = clip_image_embeds_4 * scale_4
+            image_prompt_embeds_list.append(clip_image_embeds_4)
+        if clip_image_5 != None:
+            print('Using quinary image.')
+            clip_image_5 = clip_image_5.resize((max(clip_image_5.size), max(clip_image_5.size)))
+            #with torch.no_grad():
+            clip_image_embeds_5 = self.clip_image_processor(images=clip_image_5, return_tensors="pt").pixel_values
+            clip_image_embeds_5 = clip_image_embeds_5.to(device, dtype=dtype)
+            clip_image_embeds_5 = self.image_encoder(clip_image_embeds_5, output_hidden_states=True).hidden_states[-2]
+            clip_image_embeds_5 = clip_image_embeds_5 * scale_5
+            image_prompt_embeds_list.append(clip_image_embeds_5)
+
+            # with cat and mean
+        clip_image_embeds_cat_list = torch.cat(image_prompt_embeds_list)
+        clip_image_embeds_cat_list = torch.mean(clip_image_embeds_cat_list,dim=0,keepdim=True)
+        print('catted embeds list: ',clip_image_embeds_cat_list.shape)
+        zeros_tensor = torch.zeros_like(clip_image_embeds_cat_list)
+        clip_image_embeds = torch.cat([zeros_tensor, clip_image_embeds_cat_list], dim=0)
+        print('catted embeds: ',clip_image_embeds.shape)
+        
+        '''
+        clip_image_embeds_cat_list = torch.cat(image_prompt_embeds_list).mean(dim=0)
+        print('catted embeds list with mean: ',clip_image_embeds_cat_list.shape)
+        seq_len, _ = clip_image_embeds_cat_list.shape
+        clip_image_embeds_cat_list_repeat = clip_image_embeds_cat_list.repeat(1, 1, 1)
+        print('catted embeds repeat: ',clip_image_embeds_cat_list_repeat.shape)
+        clip_image_embeds_view = clip_image_embeds_cat_list_repeat.view(1, seq_len, -1)
+        print('catted viewed: ',clip_image_embeds_view.shape)
+        zeros_tensor = torch.zeros_like(clip_image_embeds_view)
+        print('zeros: ',zeros_tensor.shape)
+        clip_image_embeds = torch.cat([zeros_tensor, clip_image_embeds_view], dim=0)
+        print('embeds shape: ', clip_image_embeds.shape)
+        '''
         # 4. Prepare timesteps
         timesteps, num_inference_steps = retrieve_timesteps(self.scheduler, num_inference_steps, device, timesteps)
         num_warmup_steps = max(len(timesteps) - num_inference_steps * self.scheduler.order, 0)
@@ -1223,7 +1306,8 @@ class StableDiffusion3Pipeline(DiffusionPipeline, SD3LoraLoaderMixin, FromSingle
         else:
             latents = (latents / self.vae.config.scaling_factor) + self.vae.config.shift_factor
 
-            image = self.vae.decode(latents, return_dict=False)[0]
+            image = self.vae.decode(latents.to(torch.float32), return_dict=False)[0]
+            
             image = self.image_processor.postprocess(image, output_type=output_type)
 
         # Offload all models