ConsistentID/lib/pipeline_ConsistentID.py

from typing import Any, Callable, Dict, List, Optional, Union, Tuple
import cv2
import PIL
import numpy as np 
from PIL import Image
import torch
from torchvision import transforms
from insightface.app import FaceAnalysis 
### insight-face installation can be found at https://github.com/deepinsight/insightface
from safetensors import safe_open
from huggingface_hub.utils import validate_hf_hub_args
from transformers import CLIPImageProcessor, CLIPVisionModelWithProjection
from diffusers.pipelines.stable_diffusion import StableDiffusionPipelineOutput
from diffusers.pipelines.stable_diffusion import StableDiffusionPipeline
from .functions import insert_markers_to_prompt, masks_for_unique_values, apply_mask_to_raw_image, tokenize_and_mask_noun_phrases_ends, prepare_image_token_idx
from .functions import ProjPlusModel, masks_for_unique_values
from .attention import Consistent_IPAttProcessor, Consistent_AttProcessor, FacialEncoder
from easydict import EasyDict as edict
from huggingface_hub import hf_hub_download
### Model can be imported from https://github.com/zllrunning/face-parsing.PyTorch?tab=readme-ov-file
### We use the ckpt of 79999_iter.pth: https://drive.google.com/open?id=154JgKpzCPW82qINcVieuPH3fZ2e0P812
### Thanks for the open source of face-parsing model.
from .BiSeNet.model import BiSeNet
import os

PipelineImageInput = Union[
    PIL.Image.Image,
    torch.FloatTensor,
    List[PIL.Image.Image],
    List[torch.FloatTensor],
]

### Download the pretrained model from huggingface and put it locally, then place the model in a local directory and specify the directory location.
class ConsistentIDPipeline(StableDiffusionPipeline):
    # to() should be only called after all modules are loaded.
    def to(
        self,
        torch_device: Optional[Union[str, torch.device]] = None,
        dtype: Optional[torch.dtype] = None,
    ):
        super().to(torch_device,                dtype=dtype)
        self.bise_net.to(torch_device,          dtype=dtype)
        self.clip_encoder.to(torch_device,      dtype=dtype)
        self.image_proj_model.to(torch_device,  dtype=dtype)
        self.FacialEncoder.to(torch_device,     dtype=dtype)
        # If the unet is not released, the ip_layers should be moved to the specified device and dtype.
        if not isinstance(self.unet, edict):
            self.ip_layers.to(torch_device,     dtype=dtype)
        return self

    @validate_hf_hub_args
    def load_ConsistentID_model(
        self,
        consistentID_weight_path:   str,
        bise_net_weight_path:       str,
        trigger_word_facial:        str = '<|facial|>',
        # A CLIP ViT-H/14 model trained with the LAION-2B English subset of LAION-5B using OpenCLIP.
        # output dim: 1280.
        image_encoder_path:         str = 'laion/CLIP-ViT-H-14-laion2B-s32B-b79K',  
        torch_dtype = torch.float16,
        num_tokens = 4,
        lora_rank= 128,
        **kwargs,
    ):
        self.lora_rank = lora_rank 
        self.torch_dtype = torch_dtype
        self.num_tokens = num_tokens
        self.set_ip_adapter()
        self.image_encoder_path = image_encoder_path
        self.clip_encoder = CLIPVisionModelWithProjection.from_pretrained(self.image_encoder_path)
        self.clip_preprocessor  = CLIPImageProcessor()
        self.id_image_processor = CLIPImageProcessor()
        self.crop_size = 512

        # face_app: FaceAnalysis object
        self.face_app = FaceAnalysis(name="buffalo_l", root='models/insightface', 
                                     providers=['CPUExecutionProvider'])
        # The original det_size=(640, 640) is too large and face_app often fails to detect faces.
        self.face_app.prepare(ctx_id=0, det_size=(512, 512))

        if not os.path.exists(consistentID_weight_path):
            ### Download pretrained models
            hf_hub_download(repo_id="JackAILab/ConsistentID", repo_type="model",
                            filename=os.path.basename(consistentID_weight_path), 
                            local_dir=os.path.dirname(consistentID_weight_path))
        if not os.path.exists(bise_net_weight_path):
            hf_hub_download(repo_id="JackAILab/ConsistentID", 
                            filename=os.path.basename(bise_net_weight_path), 
                            local_dir=os.path.dirname(bise_net_weight_path))

        bise_net = BiSeNet(n_classes = 19)
        bise_net.load_state_dict(torch.load(bise_net_weight_path, map_location="cpu"))
        bise_net.eval()
        self.bise_net = bise_net

        # Colors for all 20 parts
        self.part_colors = [[255, 0, 0], [255, 85, 0], [255, 170, 0],
                            [255, 0, 85], [255, 0, 170],
                            [0, 255, 0], [85, 255, 0], [170, 255, 0],
                            [0, 255, 85], [0, 255, 170],
                            [0, 0, 255], [85, 0, 255], [170, 0, 255],
                            [0, 85, 255], [0, 170, 255],
                            [255, 255, 0], [255, 255, 85], [255, 255, 170],
                            [255, 0, 255], [255, 85, 255], [255, 170, 255],
                            [0, 255, 255], [85, 255, 255], [170, 255, 255]]

        # image_proj_model maps 1280-dim OpenCLIP embeddings to 768-dim face prompt embeddings.
        self.image_proj_model = ProjPlusModel(
            cross_attention_dim=self.unet.config.cross_attention_dim, 
            id_embeddings_dim=512,
            clip_embeddings_dim=self.clip_encoder.config.hidden_size, 
            num_tokens=self.num_tokens,  # 4 - inspirsed by IPAdapter and Midjourney
        )
        self.FacialEncoder = FacialEncoder()

        if consistentID_weight_path.endswith(".safetensors"):
            state_dict = {"id_encoder": {}, "lora_weights": {}}
            with safe_open(consistentID_weight_path, framework="pt", device="cpu") as f:
                ### TODO safetensors add
                for key in f.keys():
                    if key.startswith("FacialEncoder."):
                        state_dict["FacialEncoder"][key.replace("FacialEncoder.", "")] = f.get_tensor(key)
                    elif key.startswith("image_proj."):
                        state_dict["image_proj"][key.replace("image_proj.", "")] = f.get_tensor(key)
        else:
            state_dict = torch.load(consistentID_weight_path, map_location="cpu")
            
        self.trigger_word_facial = trigger_word_facial

        self.FacialEncoder.load_state_dict(state_dict["FacialEncoder"], strict=True)
        self.image_proj_model.load_state_dict(state_dict["image_proj"], strict=True)
        self.ip_layers = torch.nn.ModuleList(self.unet.attn_processors.values())
        self.ip_layers.load_state_dict(state_dict["adapter_modules"], strict=True)
        print(f"Successfully loaded weights from checkpoint")

        # Add trigger word token
        if self.tokenizer is not None: 
            self.tokenizer.add_tokens([self.trigger_word_facial], special_tokens=True)

    def set_ip_adapter(self):
        unet = self.unet
        attn_procs = {}
        for name in unet.attn_processors.keys():
            cross_attention_dim = None if name.endswith("attn1.processor") else unet.config.cross_attention_dim
            if name.startswith("mid_block"):
                hidden_size = unet.config.block_out_channels[-1]
            elif name.startswith("up_blocks"):
                block_id = int(name[len("up_blocks.")])
                hidden_size = list(reversed(unet.config.block_out_channels))[block_id]
            elif name.startswith("down_blocks"):
                block_id = int(name[len("down_blocks.")])
                hidden_size = unet.config.block_out_channels[block_id]
            if cross_attention_dim is None:
                attn_procs[name] = Consistent_AttProcessor(
                    hidden_size=hidden_size, cross_attention_dim=cross_attention_dim, rank=self.lora_rank,
                )
            else:
                attn_procs[name] = Consistent_IPAttProcessor(
                    hidden_size=hidden_size, cross_attention_dim=cross_attention_dim, scale=1.0, rank=self.lora_rank, num_tokens=self.num_tokens,
                )
        
        unet.set_attn_processor(attn_procs)

    @torch.inference_mode()
    # parsed_image_parts2 is a batched tensor of parsed_image_parts with bs=1. It only contains the facial areas of one input image.
    # clip_encoder maps image parts to image-space diffusion prompts.
    # Then the facial class token embeddings are replaced with the fused (multi_facial_embeds, prompt_embeds[class_tokens_mask]).
    def extract_local_facial_embeds(self, prompt_embeds, uncond_prompt_embeds, parsed_image_parts2, 
                                    facial_token_masks, valid_facial_token_idx_mask, calc_uncond=True):
        
        hidden_states = []
        uncond_hidden_states = []
        for parsed_image_parts in parsed_image_parts2:
            hidden_state = self.clip_encoder(parsed_image_parts.to(self.device, dtype=self.torch_dtype), output_hidden_states=True).hidden_states[-2]
            uncond_hidden_state = self.clip_encoder(torch.zeros_like(parsed_image_parts, dtype=self.torch_dtype).to(self.device), output_hidden_states=True).hidden_states[-2]
            hidden_states.append(hidden_state)
            uncond_hidden_states.append(uncond_hidden_state)
        multi_facial_embeds = torch.stack(hidden_states)       
        uncond_multi_facial_embeds = torch.stack(uncond_hidden_states)   

        # conditional prompt.
        # FacialEncoder maps multi_facial_embeds to facial ID embeddings, and replaces the class tokens in prompt_embeds 
        # with the fused (facial ID embeddings, prompt_embeds[class_tokens_mask]).
        # multi_facial_embeds: [1, 5, 257, 1280].
        facial_prompt_embeds = self.FacialEncoder(prompt_embeds, multi_facial_embeds, facial_token_masks, valid_facial_token_idx_mask)  

        if not calc_uncond:
            return facial_prompt_embeds, None
        # unconditional prompt.
        uncond_facial_prompt_embeds = self.FacialEncoder(uncond_prompt_embeds, uncond_multi_facial_embeds, facial_token_masks, valid_facial_token_idx_mask)  

        return facial_prompt_embeds, uncond_facial_prompt_embeds        

    @torch.inference_mode()
    # Extrat OpenCLIP embeddings from the input image and map them to face prompt embeddings.
    def extract_global_id_embeds(self, face_image_obj, s_scale=1.0, shortcut=False):
        clip_image_ts = self.clip_preprocessor(images=face_image_obj, return_tensors="pt").pixel_values
        clip_image_ts = clip_image_ts.to(self.device, dtype=self.torch_dtype)
        clip_image_embeds = self.clip_encoder(clip_image_ts, output_hidden_states=True).hidden_states[-2]
        uncond_clip_image_embeds = self.clip_encoder(torch.zeros_like(clip_image_ts), output_hidden_states=True).hidden_states[-2]

        faceid_embeds = self.extract_faceid(face_image_obj)
        faceid_embeds = faceid_embeds.to(self.device, dtype=self.torch_dtype)
        # image_proj_model maps 1280-dim OpenCLIP embeddings to 768-dim face prompt embeddings.
        # clip_image_embeds are used as queries to transform faceid_embeds.
        # faceid_embeds -> kv, clip_image_embeds -> q
        global_id_embeds        = self.image_proj_model(faceid_embeds, clip_image_embeds, shortcut=shortcut, scale=s_scale)
        uncond_global_id_embeds = self.image_proj_model(torch.zeros_like(faceid_embeds), uncond_clip_image_embeds, shortcut=shortcut, scale=s_scale)
        
        return global_id_embeds, uncond_global_id_embeds

    def set_scale(self, scale):
        for attn_processor in self.pipe.unet.attn_processors.values():
            if isinstance(attn_processor, Consistent_IPAttProcessor):
                attn_processor.scale = scale

    @torch.inference_mode()
    def extract_faceid(self, face_image_obj):
        faceid_image = np.array(face_image_obj)
        faces = self.face_app.get(faceid_image)
        if faces==[]:
            faceid_embeds = torch.zeros_like(torch.empty((1, 512)))
        else:
            faceid_embeds = torch.from_numpy(faces[0].normed_embedding).unsqueeze(0)

        return faceid_embeds

    @torch.inference_mode()
    def parse_face_mask(self, raw_image_refer):

        to_tensor = transforms.Compose([
            transforms.ToTensor(),
            transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)),
        ])
        to_pil = transforms.ToPILImage()

        with torch.no_grad():
            image = raw_image_refer.resize((512, 512), Image.BILINEAR)
            image_resize_PIL = image
            img = to_tensor(image)
            img = torch.unsqueeze(img, 0)
            img = img.to(self.device, dtype=self.torch_dtype)
            out = self.bise_net(img)[0]
            parsing_anno = out.squeeze(0).cpu().numpy().argmax(0)
        
        im = np.array(image_resize_PIL)
        vis_im = im.copy().astype(np.uint8)
        stride=1
        vis_parsing_anno = parsing_anno.copy().astype(np.uint8)
        vis_parsing_anno = cv2.resize(vis_parsing_anno, None, fx=stride, fy=stride, interpolation=cv2.INTER_NEAREST)
        vis_parsing_anno_color = np.zeros((vis_parsing_anno.shape[0], vis_parsing_anno.shape[1], 3)) + 255

        num_of_class = np.max(vis_parsing_anno)

        for pi in range(1, num_of_class + 1): # num_of_class=17 pi=1~16
            index = np.where(vis_parsing_anno == pi) 
            vis_parsing_anno_color[index[0], index[1], :] = self.part_colors[pi] 

        vis_parsing_anno_color = vis_parsing_anno_color.astype(np.uint8)
        vis_parsing_anno_color = cv2.addWeighted(cv2.cvtColor(vis_im, cv2.COLOR_RGB2BGR), 0.4, vis_parsing_anno_color, 0.6, 0)

        return vis_parsing_anno_color, vis_parsing_anno

    @torch.inference_mode()
    def extract_facemask(self, input_image_obj):
        vis_parsing_anno_color, vis_parsing_anno = self.parse_face_mask(input_image_obj)
        parsing_mask_list = masks_for_unique_values(vis_parsing_anno) 

        key_parsing_mask_dict = {}
        key_list = ["Face", "Left_Ear", "Right_Ear", "Left_Eye", "Right_Eye", "Nose", "Upper_Lip", "Lower_Lip"]
        processed_keys = set()
        for key, mask_image in parsing_mask_list.items():
            if key in key_list:
                if "_" in key:
                    prefix = key.split("_")[1]
                    if prefix in processed_keys:                   
                        continue
                    else:            
                        key_parsing_mask_dict[key] = mask_image 
                        processed_keys.add(prefix)  
            
                key_parsing_mask_dict[key] = mask_image            

        return key_parsing_mask_dict, vis_parsing_anno_color

    def augment_prompt_with_trigger_word(
        self,
        prompt: str,
        face_caption: str,
        key_parsing_mask_dict = None,
        facial_token = "<|facial|>",
        max_num_facials = 5,
        num_id_images: int = 1,
        device: Optional[torch.device] = None,
    ):
        device = device or self._execution_device

        # face_caption_align: 'The person has one nose <|facial|>, two ears <|facial|>, two eyes <|facial|>, and a mouth <|facial|>, '
        face_caption_align, key_parsing_mask_dict_align = insert_markers_to_prompt(face_caption, key_parsing_mask_dict) 
        
        prompt_face = prompt + " Detail: " + face_caption_align

        max_text_length=330      
        if len(self.tokenizer(prompt_face, max_length=self.tokenizer.model_max_length, 
                              padding="max_length", truncation=False, return_tensors="pt").input_ids[0]) != 77:
            # Put face_caption_align at the beginning of the prompt, so that the original prompt is truncated,
            # but the face_caption_align is well kept.
            prompt_face = "Detail: " + face_caption_align + " Caption:" + prompt

        # Remove "<|facial|>" from prompt_face.
        # augmented_prompt: 'A person, police officer, half body shot Detail: 
        # The person has one nose , two ears , two eyes , and a mouth , '
        augmented_prompt = prompt_face.replace("<|facial|>", "")
        tokenizer = self.tokenizer
        facial_token_id = tokenizer.convert_tokens_to_ids(facial_token)
        image_token_id = None

        # image_token_id: the token id of "<|image|>". Disabled, as it's set to None.
        # facial_token_id: the token id of "<|facial|>".
        clean_input_id, image_token_mask, facial_token_mask = \
            tokenize_and_mask_noun_phrases_ends(prompt_face, image_token_id, facial_token_id, tokenizer) 

        image_token_idx, image_token_idx_mask, facial_token_idx, facial_token_idx_mask = \
            prepare_image_token_idx(image_token_mask, facial_token_mask, num_id_images, max_num_facials)

        return augmented_prompt, clean_input_id, key_parsing_mask_dict_align, facial_token_mask, facial_token_idx, facial_token_idx_mask

    @torch.inference_mode()
    def extract_parsed_image_parts(self, input_image_obj, key_parsing_mask_dict, image_size=512, max_num_facials=5):
        facial_masks = []
        parsed_image_parts = []
        key_masked_raw_images_dict = {}
        transform_mask = transforms.Compose([transforms.CenterCrop(size=image_size), transforms.ToTensor(),])
        clip_preprocessor = CLIPImageProcessor()

        num_facial_part = len(key_parsing_mask_dict)

        for key in key_parsing_mask_dict:
            key_mask=key_parsing_mask_dict[key]
            facial_masks.append(transform_mask(key_mask))
            key_masked_raw_image = apply_mask_to_raw_image(input_image_obj, key_mask)
            key_masked_raw_images_dict[key] = key_masked_raw_image
            # clip_preprocessor normalizes key_masked_raw_image, so that (masked) zero pixels become non-zero.
            # It also resizes the image to 224x224.
            parsed_image_part = clip_preprocessor(images=key_masked_raw_image, return_tensors="pt").pixel_values
            parsed_image_parts.append(parsed_image_part)
            
        padding_ficial_clip_image = torch.zeros_like(torch.zeros([1, 3, 224, 224]))
        padding_ficial_mask = torch.zeros_like(torch.zeros([1, image_size, image_size]))

        if num_facial_part < max_num_facials:
            parsed_image_parts  += [ torch.zeros_like(padding_ficial_clip_image) for _ in range(max_num_facials - num_facial_part) ]
            facial_masks        += [ torch.zeros_like(padding_ficial_mask)       for _ in range(max_num_facials - num_facial_part) ]

        parsed_image_parts  = torch.stack(parsed_image_parts, dim=1).squeeze(0)
        facial_masks        = torch.stack(facial_masks,       dim=0).squeeze(dim=1)

        return parsed_image_parts, facial_masks, key_masked_raw_images_dict

    # Release the unet/vae/text_encoder to save memory.
    def release_components(self, released_components=["unet", "vae", "text_encoder"]):
        if "unet" in released_components:
            unet = edict()
            # Only keep the config and in_channels attributes that are used in the pipeline.
            unet.config = self.unet.config
            self.unet = unet

        if "vae" in released_components:
            self.vae = None
        if "text_encoder" in released_components:
            self.text_encoder = None

    # input_subj_image_obj: an Image object.
    def extract_double_id_prompt_embeds(self, prompt, negative_prompt, input_subj_image_obj, device, calc_uncond=True):
        face_caption = "The person has one nose, two eyes, two ears, and a mouth."
        key_parsing_mask_dict, vis_parsing_anno_color = self.extract_facemask(input_subj_image_obj)

        augmented_prompt, clean_input_id, key_parsing_mask_dict_align, \
          facial_token_mask, facial_token_idx, facial_token_idx_mask \
            = self.augment_prompt_with_trigger_word(
                prompt = prompt,
                face_caption = face_caption,
                key_parsing_mask_dict=key_parsing_mask_dict,
                device=device,
                max_num_facials = 5,
                num_id_images = 1
                )

        text_embeds, uncond_text_embeds = self.encode_prompt(
            augmented_prompt,
            device=device,
            num_images_per_prompt=1,
            do_classifier_free_guidance=calc_uncond,
            negative_prompt=negative_prompt,
        )

        # 5. Prepare the input ID images
        # global_id_embeds: [1, 4, 768]
        # extract_global_id_embeds() extrats OpenCLIP embeddings from the input image and map them to global face prompt embeddings.
        global_id_embeds, uncond_global_id_embeds = \
            self.extract_global_id_embeds(face_image_obj=input_subj_image_obj, s_scale=1.0, shortcut=False)

        # parsed_image_parts: [5, 3, 224, 224]. 5 parts, each part is a 3-channel 224x224 image (resized by CLIP Preprocessor).
        parsed_image_parts, facial_masks, key_masked_raw_images_dict = \
            self.extract_parsed_image_parts(input_subj_image_obj, key_parsing_mask_dict_align, image_size=512, max_num_facials=5)
        parsed_image_parts2 = parsed_image_parts.unsqueeze(0).to(device, dtype=self.torch_dtype)
        facial_token_mask = facial_token_mask.to(device)
        facial_token_idx_mask = facial_token_idx_mask.to(device)

        # key_masked_raw_images_dict: ['Right_Eye', 'Right_Ear', 'Nose', 'Upper_Lip']
        # for key in key_masked_raw_images_dict:
        #     key_masked_raw_images_dict[key].save(f"{key}.png")

        # 6. Get the update text embedding
        # parsed_image_parts2: the facial areas of the input image
        # extract_local_facial_embeds() maps parsed_image_parts2 to multi_facial_embeds, and then replaces the class tokens in prompt_embeds 
        # with the fused (id_embeds, prompt_embeds[class_tokens_mask]) whose indices are specified by class_tokens_mask.        
        # parsed_image_parts2: [1, 5, 3, 224, 224]
        text_local_id_embeds, uncond_text_local_id_embeds = \
            self.extract_local_facial_embeds(text_embeds, uncond_text_embeds, \
                                             parsed_image_parts2, facial_token_mask, facial_token_idx_mask,
                                             calc_uncond=calc_uncond)

        # text_global_id_embeds, text_local_global_id_embeds: [1, 81, 768]
        # text_local_id_embeds: [1, 77, 768], only differs with text_embeds on 4 ID embeddings, and is identical
        # to text_embeds on the rest 73 tokens.
        text_global_id_embeds         = torch.cat([text_embeds,          global_id_embeds], dim=1)
        text_local_global_id_embeds   = torch.cat([text_local_id_embeds, global_id_embeds], dim=1)

        if calc_uncond:
            uncond_text_global_id_embeds  = torch.cat([uncond_text_local_id_embeds, uncond_global_id_embeds], dim=1)
            coarse_prompt_embeds = torch.cat([uncond_text_global_id_embeds, text_global_id_embeds], dim=0)
            fine_prompt_embeds   = torch.cat([uncond_text_global_id_embeds, text_local_global_id_embeds], dim=0)
        else:
            coarse_prompt_embeds = text_global_id_embeds
            fine_prompt_embeds   = text_local_global_id_embeds

        # fine_prompt_embeds: the conditional part is 
        # (text_global_id_embeds + text_local_global_id_embeds) / 2.
        fine_prompt_embeds   = (coarse_prompt_embeds + fine_prompt_embeds) / 2

        return coarse_prompt_embeds, fine_prompt_embeds
    
    @torch.no_grad()
    def __call__(
        self,
        prompt: Union[str, List[str]] = None,
        height: Optional[int] = None,
        width: Optional[int] = None,
        num_inference_steps: int = 50,
        guidance_scale: float = 5.0,
        negative_prompt: Optional[Union[str, List[str]]] = None,
        num_images_per_prompt: Optional[int] = 1,
        eta: float = 0.0,
        generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
        latents: Optional[torch.FloatTensor] = None,
        prompt_embeds: Optional[torch.FloatTensor] = None,
        negative_prompt_embeds: Optional[torch.FloatTensor] = None,
        output_type: Optional[str] = "pil",
        return_dict: bool = True,
        cross_attention_kwargs: Optional[Dict[str, Any]] = None,
        original_size: Optional[Tuple[int, int]] = None,
        target_size: Optional[Tuple[int, int]] = None,
        callback: Optional[Callable[[int, int, torch.FloatTensor], None]] = None,
        callback_steps: int = 1,
        input_subj_image_objs: PipelineImageInput = None,
        start_merge_step: int = 0,
    ):
        # 0. Default height and width to unet
        height = height or self.unet.config.sample_size * self.vae_scale_factor
        width  = width  or self.unet.config.sample_size * self.vae_scale_factor

        original_size = original_size or (height, width)
        target_size = target_size or (height, width)

        # 1. Check inputs. Raise error if not correct
        self.check_inputs(
            prompt,
            height,
            width,
            callback_steps,
            negative_prompt,
            prompt_embeds,
            negative_prompt_embeds,
        )

        # 2. Define call parameters
        if prompt is not None and isinstance(prompt, str):
            batch_size = 1
        elif prompt is not None and isinstance(prompt, list):
            batch_size = len(prompt)
        else:
            batch_size = prompt_embeds.shape[0]

        device = self._execution_device
        do_classifier_free_guidance = guidance_scale >= 1.0
        assert do_classifier_free_guidance

        if input_subj_image_objs is not None:
            if not isinstance(input_subj_image_objs, list):
                input_subj_image_objs = [input_subj_image_objs]

            # 3. Encode input prompt
            coarse_prompt_embeds, fine_prompt_embeds = \
                self.extract_double_id_prompt_embeds(prompt, negative_prompt, input_subj_image_objs[0], device)
        else:
            # Replace the coarse_prompt_embeds and fine_prompt_embeds with the input prompt_embeds.
            # This is used when prompt_embeds are computed in advance.
            cfg_prompt_embeds    = torch.cat([negative_prompt_embeds, prompt_embeds], dim=0)
            coarse_prompt_embeds = cfg_prompt_embeds
            fine_prompt_embeds   = cfg_prompt_embeds

        # 7. Prepare timesteps
        self.scheduler.set_timesteps(num_inference_steps, device=device)
        timesteps = self.scheduler.timesteps

        # 8. Prepare latent variables
        num_channels_latents = self.unet.config.in_channels
        latents = self.prepare_latents(
            batch_size * num_images_per_prompt,
            num_channels_latents,
            height,
            width,
            self.dtype,
            device,
            generator,
            latents,
        )

        # {'eta': 0.0, 'generator': None}. eta is 0 for DDIM.
        extra_step_kwargs    = self.prepare_extra_step_kwargs(generator, eta)
        cross_attention_kwargs = {}

        # 9. Denoising loop
        num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order
        with self.progress_bar(total=num_inference_steps) as progress_bar:
            for i, t in enumerate(timesteps):
                latent_model_input = (
                    torch.cat([latents] * 2) if do_classifier_free_guidance else latents
                )
                # DDIM doesn't scale latent_model_input.                
                latent_model_input = self.scheduler.scale_model_input(latent_model_input, t)
                
                if i <= start_merge_step:
                    current_prompt_embeds = coarse_prompt_embeds
                else:
                    current_prompt_embeds = fine_prompt_embeds

                # predict the noise residual
                noise_pred = self.unet(
                    latent_model_input,
                    t,
                    encoder_hidden_states=current_prompt_embeds,
                    cross_attention_kwargs=cross_attention_kwargs,
                ).sample

                # perform guidance
                if do_classifier_free_guidance:
                    noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
                    noise_pred = noise_pred_uncond + guidance_scale * (
                        noise_pred_text - noise_pred_uncond
                    )
                else:
                    assert 0, "Not Implemented"

                # compute the previous noisy sample x_t -> x_t-1
                latents = self.scheduler.step(
                    noise_pred, t, latents, **extra_step_kwargs
                ).prev_sample

                # call the callback, if provided
                if i == len(timesteps) - 1 or \
                  ( (i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0 ):
                    progress_bar.update()
                    if callback is not None and i % callback_steps == 0:
                        callback(i, t, latents)

        if output_type == "latent":
            image = latents
        elif output_type == "pil":
            # 9.1 Post-processing
            image = self.decode_latents(latents)
            # 9.3 Convert to PIL
            image = self.numpy_to_pil(image)
        else:
            # 9.1 Post-processing
            image = self.decode_latents(latents)

        # Offload last model to CPU
        if hasattr(self, "final_offload_hook") and self.final_offload_hook is not None:
            self.final_offload_hook.offload()

        if not return_dict:
            return (image, None)

        return StableDiffusionPipelineOutput(
            images=image, nsfw_content_detected=None
        )