From 8c4d49e4cd0bb1e1b875fa837e6a043ac9b65fc3 Mon Sep 17 00:00:00 2001
From: smthemex <138738845+smthemex@users.noreply.github.com>
Date: Mon, 7 Oct 2024 08:44:21 +0800
Subject: [PATCH] Update Storydiffusion_node.py

fix bug
---
 Storydiffusion_node.py | 1082 ++++++++++++++++++++++++++++++++++++++--
 1 file changed, 1052 insertions(+), 30 deletions(-)

diff --git a/Storydiffusion_node.py b/Storydiffusion_node.py
index 1725a5c..5bf108d 100644
--- a/Storydiffusion_node.py
+++ b/Storydiffusion_node.py
@@ -1,14 +1,12 @@
 # !/usr/bin/env python
 # -*- coding: UTF-8 -*-
-
+import random
 import gc
 import logging
 import numpy as np
 import torch
 import os
-
-from PIL import ImageFont
-from huggingface_hub import hf_hub_download
+from PIL import ImageFont,Image
 from diffusers import (StableDiffusionXLPipeline, DiffusionPipeline,EulerDiscreteScheduler, UNet2DConditionModel,UniPCMultistepScheduler, AutoencoderKL,)
 from transformers import CLIPVisionModelWithProjection
 from transformers import CLIPImageProcessor
@@ -19,14 +17,20 @@ from comfy.clip_vision import load as clip_load
 from comfy.model_management import total_vram
 
 from .utils.utils import get_comic
-from .utils.gradio_utils import character_to_dict
-from .utils.load_models_utils import load_models, get_instance_path
+from .utils.load_models_utils import load_models
 from .model_loader_utils import  (story_maker_loader,kolor_loader,phi2narry,
                                   extract_content_from_brackets,narry_list,remove_punctuation_from_strings,phi_list,center_crop_s,center_crop,
                                   narry_list_pil,setup_seed,find_directories,
-                                  apply_style,get_scheduler,set_attention_processor,load_character_files_on_running,
-                                  save_results,nomarl_upscale,process_generation,SAMPLER_NAMES,SCHEDULER_NAMES,lora_lightning_list,pre_checkpoint,get_easy_function)
-
+                                  apply_style,get_scheduler,load_character_files_on_running,apply_style_positive,
+                                  save_results,nomarl_upscale,SAMPLER_NAMES,SCHEDULER_NAMES,lora_lightning_list,pre_checkpoint,get_easy_function)
+from .utils.gradio_utils import cal_attn_indice_xl_effcient_memory,is_torch2_available,process_original_prompt,get_ref_character,character_to_dict
+from .ip_adapter.attention_processor import IPAttnProcessor2_0
+if is_torch2_available():
+    from .utils.gradio_utils import AttnProcessor2_0 as AttnProcessor
+else:
+    from .utils.gradio_utils import AttnProcessor
+import torch.nn.functional as F
+import copy
 global total_count, attn_count, cur_step, mask1024, mask4096, attn_procs, unet
 global sa32, sa64
 global write
@@ -45,6 +49,1045 @@ base_pt = os.path.join(photomaker_dir,"pt")
 if not os.path.exists(base_pt):
     os.makedirs(base_pt)
     
+def set_attention_processor(unet, id_length, is_ipadapter=False):
+    global attn_procs
+    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:
+            if name.startswith("up_blocks"):
+                attn_procs[name] = SpatialAttnProcessor2_0(id_length=id_length)
+            else:
+                attn_procs[name] = AttnProcessor()
+        else:
+            if is_ipadapter:
+                attn_procs[name] = IPAttnProcessor2_0(
+                    hidden_size=hidden_size,
+                    cross_attention_dim=cross_attention_dim,
+                    scale=1,
+                    num_tokens=4,
+                ).to(unet.device, dtype=torch.float16)
+            else:
+                attn_procs[name] = AttnProcessor()
+
+    unet.set_attn_processor(copy.deepcopy(attn_procs))
+
+class SpatialAttnProcessor2_0(torch.nn.Module):
+    r"""
+    Attention processor for IP-Adapater for PyTorch 2.0.
+    Args:
+        hidden_size (`int`):
+            The hidden size of the attention layer.
+        cross_attention_dim (`int`):
+            The number of channels in the `encoder_hidden_states`.
+        text_context_len (`int`, defaults to 77):
+            The context length of the text features.
+        scale (`float`, defaults to 1.0):
+            the weight scale of image prompt.
+    """
+
+    def __init__(
+            self,
+            hidden_size=None,
+            cross_attention_dim=None,
+            id_length=4,
+            device=device,
+            dtype=torch.float16,
+    ):
+        super().__init__()
+        if not hasattr(F, "scaled_dot_product_attention"):
+            raise ImportError(
+                "AttnProcessor2_0 requires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0."
+            )
+        self.device = device
+        self.dtype = dtype
+        self.hidden_size = hidden_size
+        self.cross_attention_dim = cross_attention_dim
+        self.total_length = id_length + 1
+        self.id_length = id_length
+        self.id_bank = {}
+
+    def __call__(
+            self,
+            attn,
+            hidden_states,
+            encoder_hidden_states=None,
+            attention_mask=None,
+            temb=None,
+    ):
+        # un_cond_hidden_states, cond_hidden_states = hidden_states.chunk(2)
+        # un_cond_hidden_states = self.__call2__(attn, un_cond_hidden_states,encoder_hidden_states,attention_mask,temb)
+        # 生成一个0到1之间的随机数
+        global total_count, attn_count, cur_step, indices1024, indices4096
+        global sa32, sa64
+        global write
+        global height_s, width_s
+        global character_dict
+        global  character_index_dict, invert_character_index_dict, cur_character, ref_indexs_dict, ref_totals, cur_character
+        if attn_count == 0 and cur_step == 0:
+            indices1024, indices4096 = cal_attn_indice_xl_effcient_memory(
+                self.total_length,
+                self.id_length,
+                sa32,
+                sa64,
+                height_s,
+                width_s,
+                device=self.device,
+                dtype=self.dtype,
+            )
+        if write:
+            assert len(cur_character) == 1
+            if hidden_states.shape[1] == (height_s // 32) * (width_s // 32):
+                indices = indices1024
+            else:
+                indices = indices4096
+            # print(f"white:{cur_step}")
+            total_batch_size, nums_token, channel = hidden_states.shape
+            img_nums = total_batch_size // 2
+            hidden_states = hidden_states.reshape(-1, img_nums, nums_token, channel)
+            # print(img_nums,len(indices),hidden_states.shape,self.total_length)
+            if cur_character[0] not in self.id_bank:
+                self.id_bank[cur_character[0]] = {}
+            self.id_bank[cur_character[0]][cur_step] = [
+                hidden_states[:, img_ind, indices[img_ind], :]
+                .reshape(2, -1, channel)
+                .clone()
+                for img_ind in range(img_nums)
+            ]
+            hidden_states = hidden_states.reshape(-1, nums_token, channel)
+            # self.id_bank[cur_step] = [hidden_states[:self.id_length].clone(), hidden_states[self.id_length:].clone()]
+        else:
+            # encoder_hidden_states = torch.cat((self.id_bank[cur_step][0].to(self.device),self.id_bank[cur_step][1].to(self.device)))
+            # TODO: ADD Multipersion Control
+            encoder_arr = []
+            for character in cur_character:
+                encoder_arr = encoder_arr + [
+                    tensor.to(self.device)
+                    for tensor in self.id_bank[character][cur_step]
+                ]
+        # 判断随机数是否大于0.5
+        if cur_step < 1:
+            hidden_states = self.__call2__(
+                attn, hidden_states, None, attention_mask, temb
+            )
+        else:  # 256 1024 4096
+            random_number = random.random()
+            if cur_step < 20:
+                rand_num = 0.3
+            else:
+                rand_num = 0.1
+            # print(f"hidden state shape {hidden_states.shape[1]}")
+            if random_number > rand_num:
+                if hidden_states.shape[1] == (height_s // 32) * (width_s // 32):
+                    indices = indices1024
+                else:
+                    indices = indices4096
+                # print("before attention",hidden_states.shape,attention_mask.shape,encoder_hidden_states.shape if encoder_hidden_states is not None else "None")
+                if write:
+                    total_batch_size, nums_token, channel = hidden_states.shape
+                    img_nums = total_batch_size // 2
+                    hidden_states = hidden_states.reshape(
+                        -1, img_nums, nums_token, channel
+                    )
+                    encoder_arr = [
+                        hidden_states[:, img_ind, indices[img_ind], :].reshape(
+                            2, -1, channel
+                        )
+                        for img_ind in range(img_nums)
+                    ]
+                    for img_ind in range(img_nums):
+                        # print(img_nums)
+                        # assert img_nums != 1
+                        img_ind_list = [i for i in range(img_nums)]
+                        # print(img_ind_list,img_ind)
+                        img_ind_list.remove(img_ind)
+                        # print(img_ind,invert_character_index_dict[img_ind])
+                        # print(character_index_dict[invert_character_index_dict[img_ind]])
+                        # print(img_ind_list)
+                        # print(img_ind,img_ind_list)
+                        encoder_hidden_states_tmp = torch.cat(
+                            [encoder_arr[img_ind] for img_ind in img_ind_list]
+                            + [hidden_states[:, img_ind, :, :]],
+                            dim=1,
+                        )
+
+                        hidden_states[:, img_ind, :, :] = self.__call2__(
+                            attn,
+                            hidden_states[:, img_ind, :, :],
+                            encoder_hidden_states_tmp,
+                            None,
+                            temb,
+                        )
+                else:
+                    _, nums_token, channel = hidden_states.shape
+                    # img_nums = total_batch_size // 2
+                    # encoder_hidden_states = encoder_hidden_states.reshape(-1,img_nums,nums_token,channel)
+                    hidden_states = hidden_states.reshape(2, -1, nums_token, channel)
+
+                    # encoder_arr = [encoder_hidden_states[:,img_ind,indices[img_ind],:].reshape(2,-1,channel) for img_ind in range(img_nums)]
+                    encoder_hidden_states_tmp = torch.cat(
+                        encoder_arr + [hidden_states[:, 0, :, :]], dim=1
+                    )
+
+                    hidden_states[:, 0, :, :] = self.__call2__(
+                        attn,
+                        hidden_states[:, 0, :, :],
+                        encoder_hidden_states_tmp,
+                        None,
+                        temb,
+                    )
+                hidden_states = hidden_states.reshape(-1, nums_token, channel)
+            else:
+                hidden_states = self.__call2__(
+                    attn, hidden_states, None, attention_mask, temb
+                )
+        attn_count += 1
+        if attn_count == total_count:
+            attn_count = 0
+            cur_step += 1
+            indices1024, indices4096 = cal_attn_indice_xl_effcient_memory(
+                self.total_length,
+                self.id_length,
+                sa32,
+                sa64,
+                height_s,
+                width_s,
+                device=self.device,
+                dtype=self.dtype,
+            )
+
+        return hidden_states
+
+    def __call2__(
+            self,
+            attn,
+            hidden_states,
+            encoder_hidden_states=None,
+            attention_mask=None,
+            temb=None,
+    ):
+        residual = hidden_states
+
+        if attn.spatial_norm is not None:
+            hidden_states = attn.spatial_norm(hidden_states, temb)
+
+        input_ndim = hidden_states.ndim
+
+        if input_ndim == 4:
+            batch_size, channel, height_s, width_s = hidden_states.shape
+            hidden_states = hidden_states.view(
+                batch_size, channel, height_s * width_s
+            ).transpose(1, 2)
+
+        batch_size, sequence_length, channel = hidden_states.shape
+
+        if attention_mask is not None:
+            attention_mask = attn.prepare_attention_mask(
+                attention_mask, sequence_length, batch_size
+            )
+            # scaled_dot_product_attention expects attention_mask shape to be
+            # (batch, heads, source_length, target_length)
+            attention_mask = attention_mask.view(
+                batch_size, attn.heads, -1, attention_mask.shape[-1]
+            )
+
+        if attn.group_norm is not None:
+            hidden_states = attn.group_norm(hidden_states.transpose(1, 2)).transpose(
+                1, 2
+            )
+
+        query = attn.to_q(hidden_states)
+
+        if encoder_hidden_states is None:
+            encoder_hidden_states = hidden_states  # B, N, C
+        # else:
+        #     encoder_hidden_states = encoder_hidden_states.view(-1,self.id_length+1,sequence_length,channel).reshape(-1,(self.id_length+1) * sequence_length,channel)
+
+        key = attn.to_k(encoder_hidden_states)
+        value = attn.to_v(encoder_hidden_states)
+
+        inner_dim = key.shape[-1]
+        head_dim = inner_dim // attn.heads
+
+        query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+
+        key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+        value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+
+        # the output of sdp = (batch, num_heads, seq_len, head_dim)
+        # TODO: add support for attn.scale when we move to Torch 2.1
+        hidden_states = F.scaled_dot_product_attention(
+            query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False
+        )
+
+        hidden_states = hidden_states.transpose(1, 2).reshape(
+            batch_size, -1, attn.heads * head_dim
+        )
+        hidden_states = hidden_states.to(query.dtype)
+
+        # linear proj
+        hidden_states = attn.to_out[0](hidden_states)
+        # dropout
+        hidden_states = attn.to_out[1](hidden_states)
+
+        if input_ndim == 4:
+            hidden_states = hidden_states.transpose(-1, -2).reshape(
+                batch_size, channel, height_s, width_s
+            )
+
+        if attn.residual_connection:
+            hidden_states = hidden_states + residual
+
+        hidden_states = hidden_states / attn.rescale_output_factor
+
+        return hidden_states
+
+
+def process_generation(
+        pipe,
+        upload_images,
+        model_type,
+        _num_steps,
+        style_name,
+        denoise_or_ip_sacle,
+        _style_strength_ratio,
+        cfg,
+        seed_,
+        id_length,
+        general_prompt,
+        negative_prompt,
+        prompt_array,
+        width,
+        height,
+        load_chars,
+        lora,
+        trigger_words, photomake_mode, use_kolor, use_flux, make_dual_only, kolor_face, pulid, story_maker,
+        input_id_emb_s_dict, input_id_img_s_dict, input_id_emb_un_dict, input_id_cloth_dict, guidance, condition_image,
+        empty_emb_zero, use_cf, cf_scheduler, controlnet_path, controlnet_scale, cn_dict
+):  # Corrected font_choice usage
+    
+    if len(general_prompt.splitlines()) >= 3:
+        raise "Support for more than three characters is temporarily unavailable due to VRAM limitations, but this issue will be resolved soon."
+    # _model_type = "Photomaker" if _model_type == "Using Ref Images" else "original"
+    
+    if not use_kolor and not use_flux and not story_maker:
+        if model_type == "img2img" and "img" not in general_prompt:
+            raise 'if using normal SDXL img2img ,need add the triger word " img "  behind the class word you want to customize, such as: man img or woman img'
+    
+    global total_length, attn_procs, cur_model_type
+    global write
+    global cur_step, attn_count
+    
+    # load_chars = load_character_files_on_running(unet, character_files=char_files)
+    
+    prompts_origin = prompt_array.splitlines()
+    prompts_origin = [i.strip() for i in prompts_origin]
+    prompts_origin = [i for i in prompts_origin if '[' in i]  # 删除空行
+    # print(prompts_origin)
+    prompts = [prompt for prompt in prompts_origin if not len(extract_content_from_brackets(prompt)) >= 2]  # 剔除双角色
+    
+    add_trigger_words = " " + trigger_words + " style "
+    if lora:
+        if lora in lora_lightning_list:
+            prompts = remove_punctuation_from_strings(prompts)
+        else:
+            prompts = remove_punctuation_from_strings(prompts)
+            prompts = [item + add_trigger_words for item in prompts]
+    
+    global character_index_dict, invert_character_index_dict, ref_indexs_dict, ref_totals
+    global character_dict
+    
+    character_dict, character_list = character_to_dict(general_prompt, lora, add_trigger_words)
+    # print(character_dict)
+    start_merge_step = int(float(_style_strength_ratio) / 100 * _num_steps)
+    if start_merge_step > 30:
+        start_merge_step = 30
+    print(f"start_merge_step:{start_merge_step}")
+    # generator = torch.Generator(device=device).manual_seed(seed_)
+    clipped_prompts = prompts[:]
+    nc_indexs = []
+    for ind, prompt in enumerate(clipped_prompts):
+        if "[NC]" in prompt:
+            nc_indexs.append(ind)
+            if ind < id_length:
+                raise f"The first {id_length} row is id prompts, cannot use [NC]!"
+    
+    prompts = [
+        prompt if "[NC]" not in prompt else prompt.replace("[NC]", "")
+        for prompt in clipped_prompts
+    ]
+    
+    if lora:
+        if lora in lora_lightning_list:
+            prompts = [
+                prompt.rpartition("#")[0] if "#" in prompt else prompt for prompt in prompts
+            ]
+        else:
+            prompts = [
+                prompt.rpartition("#")[0] + add_trigger_words if "#" in prompt else prompt for prompt in prompts
+            ]
+    else:
+        prompts = [
+            prompt.rpartition("#")[0] if "#" in prompt else prompt for prompt in prompts
+        ]
+    img_mode = False
+    if kolor_face or pulid or (story_maker and not make_dual_only) or model_type == "img2img":
+        img_mode = True
+    # id_prompts = prompts[:id_length]
+    (
+        character_index_dict,
+        invert_character_index_dict,
+        replace_prompts,
+        ref_indexs_dict,
+        ref_totals,
+    ) = process_original_prompt(character_dict, prompts, id_length, img_mode)
+    # print(character_index_dict,invert_character_index_dict,replace_prompts,ref_indexs_dict,ref_totals)
+    # character_index_dict：{'[Taylor]': [0, 3], '[sam]': [1, 2]},if 1 role {'[Taylor]': [0, 1, 2]}
+    # invert_character_index_dict:{0: ['[Taylor]'], 1: ['[sam]'], 2: ['[sam]'], 3: ['[Taylor]']},if 1 role  {0: ['[Taylor]'], 1: ['[Taylor]'], 2: ['[Taylor]']}
+    # ref_indexs_dict:{'[Taylor]': [0, 3], '[sam]': [1, 2]},if 1 role {'[Taylor]': [0]}
+    # ref_totals: [0, 3, 1, 2]  if 1 role [0]
+    if model_type == "img2img":
+        # _upload_images = [_upload_images]
+        input_id_images_dict = {}
+        if len(upload_images) != len(character_dict.keys()):
+            raise f"You upload images({len(upload_images)}) is not equal to the number of characters({len(character_dict.keys())})!"
+        for ind, img in enumerate(upload_images):
+            input_id_images_dict[character_list[ind]] = [img]  # 已经pil转化了 不用load {a:[img],b:[img]}
+            # input_id_images_dict[character_list[ind]] = [load_image(img)]
+    # real_prompts = prompts[id_length:]
+    # if device == "cuda":
+    #     torch.cuda.empty_cache()
+    write = True
+    cur_step = 0
+    attn_count = 0
+    total_results = []
+    id_images = []
+    results_dict = {}
+    p_num = 0
+    
+    global cur_character
+    if not load_chars:
+        for character_key in character_dict.keys():  # 先生成角色对应第一句场景提示词的图片,图生图是批次生成
+            character_key_str = character_key
+            cur_character = [character_key]
+            ref_indexs = ref_indexs_dict[character_key]
+            current_prompts = [replace_prompts[ref_ind] for ref_ind in ref_indexs]
+            if model_type == "txt2img":
+                setup_seed(seed_)
+            generator = torch.Generator(device=device).manual_seed(seed_)
+            cur_step = 0
+            cur_positive_prompts, cur_negative_prompt = apply_style(
+                style_name, current_prompts, negative_prompt
+            )
+            print(f"Sampler  {character_key_str} 's cur_positive_prompts :{cur_positive_prompts}")
+            if model_type == "txt2img":
+                if use_flux:
+                    if not use_cf:
+                        id_images = pipe(
+                            prompt=cur_positive_prompts,
+                            num_inference_steps=_num_steps,
+                            guidance_scale=guidance,
+                            output_type="pil",
+                            max_sequence_length=256,
+                            height=height,
+                            width=width,
+                            generator=generator
+                        ).images
+                    else:
+                        cur_negative_prompt = [cur_negative_prompt]
+                        cur_negative_prompt = cur_negative_prompt * len(cur_positive_prompts) if len(
+                            cur_negative_prompt) != len(cur_positive_prompts) else cur_negative_prompt
+                        id_images = []
+                        cfg = 1.0
+                        for index, text in enumerate(cur_positive_prompts):
+                            id_image = pipe.generate_image(
+                                width=width,
+                                height=height,
+                                num_steps=_num_steps,
+                                cfg=cfg,
+                                guidance=guidance,
+                                seed=seed_,
+                                prompt=text,
+                                negative_prompt=cur_negative_prompt[index],
+                                cf_scheduler=cf_scheduler,
+                                denoise=denoise_or_ip_sacle,
+                                image=None
+                            )
+                            id_images.append(id_image)
+                else:
+                    if use_cf:
+                        cur_negative_prompt = [cur_negative_prompt]
+                        cur_negative_prompt = cur_negative_prompt * len(cur_positive_prompts) if len(
+                            cur_negative_prompt) != len(cur_positive_prompts) else cur_negative_prompt
+                        id_images = []
+                        for index, text in enumerate(cur_positive_prompts):
+                            id_image = pipe.generate_image(
+                                width=width,
+                                height=height,
+                                num_steps=_num_steps,
+                                cfg=cfg,
+                                guidance=guidance,
+                                seed=seed_,
+                                prompt=text,
+                                negative_prompt=cur_negative_prompt[index],
+                                cf_scheduler=cf_scheduler,
+                                denoise=denoise_or_ip_sacle,
+                                image=None
+                            )
+                            id_images.append(id_image)
+                    else:
+                        if use_kolor:
+                            cur_negative_prompt = [cur_negative_prompt]
+                            cur_negative_prompt = cur_negative_prompt * len(cur_positive_prompts) if len(
+                                cur_negative_prompt) != len(cur_positive_prompts) else cur_negative_prompt
+                        id_images = pipe(
+                            cur_positive_prompts,
+                            num_inference_steps=_num_steps,
+                            guidance_scale=cfg,
+                            height=height,
+                            width=width,
+                            negative_prompt=cur_negative_prompt,
+                            generator=generator
+                        ).images
+            
+            
+            elif model_type == "img2img":
+                if use_kolor:
+                    cur_negative_prompt = [cur_negative_prompt]
+                    cur_negative_prompt = cur_negative_prompt * len(cur_positive_prompts) if len(
+                        cur_negative_prompt) != len(cur_positive_prompts) else cur_negative_prompt
+                    if kolor_face:
+                        crop_image = input_id_img_s_dict[character_key_str][0]
+                        face_embeds = input_id_emb_s_dict[character_key_str][0]
+                        face_embeds = face_embeds.to(device, dtype=torch.float16)
+                        if id_length > 1:
+                            id_images = []
+                            for index, i in enumerate(cur_positive_prompts):
+                                id_image = pipe(
+                                    prompt=i,
+                                    negative_prompt=cur_negative_prompt[index],
+                                    height=height,
+                                    width=width,
+                                    num_inference_steps=_num_steps,
+                                    guidance_scale=cfg,
+                                    num_images_per_prompt=1,
+                                    generator=generator,
+                                    face_crop_image=crop_image,
+                                    face_insightface_embeds=face_embeds,
+                                ).images
+                                id_images.append(id_image)
+                        else:
+                            id_images = pipe(
+                                prompt=cur_positive_prompts,
+                                negative_prompt=cur_negative_prompt,
+                                height=height,
+                                width=width,
+                                num_inference_steps=_num_steps,
+                                guidance_scale=cfg,
+                                num_images_per_prompt=1,
+                                generator=generator,
+                                face_crop_image=crop_image,
+                                face_insightface_embeds=face_embeds,
+                            ).images
+                    else:
+                        pipe.set_ip_adapter_scale([denoise_or_ip_sacle])
+                        id_images = pipe(
+                            prompt=cur_positive_prompts,
+                            ip_adapter_image=input_id_images_dict[character_key],
+                            negative_prompt=cur_negative_prompt,
+                            num_inference_steps=_num_steps,
+                            height=height,
+                            width=width,
+                            guidance_scale=cfg,
+                            num_images_per_prompt=1,
+                            generator=generator,
+                        ).images
+                elif use_flux:
+                    if pulid:
+                        id_embeddings = input_id_emb_s_dict[character_key_str][0]
+                        uncond_id_embeddings = input_id_emb_un_dict[character_key_str][0]
+                        if id_length > 1:
+                            id_images = []
+                            for index, i in enumerate(cur_positive_prompts):
+                                id_image = pipe.generate_image(
+                                    prompt=i,
+                                    seed=seed_,
+                                    start_step=2,
+                                    num_steps=_num_steps,
+                                    height=height,
+                                    width=width,
+                                    id_embeddings=id_embeddings,
+                                    uncond_id_embeddings=uncond_id_embeddings,
+                                    id_weight=1,
+                                    guidance=guidance,
+                                    true_cfg=1.0,
+                                    max_sequence_length=128,
+                                )
+                                id_images.append(id_image)
+                        
+                        else:
+                            id_images = pipe.generate_image(
+                                prompt=cur_positive_prompts,
+                                seed=seed_,
+                                start_step=2,
+                                num_steps=_num_steps,
+                                height=height,
+                                width=width,
+                                id_embeddings=id_embeddings,
+                                uncond_id_embeddings=uncond_id_embeddings,
+                                id_weight=1,
+                                guidance=guidance,
+                                true_cfg=1.0,
+                                max_sequence_length=128,
+                            )
+                            id_images = [id_images]
+                    elif use_cf:
+                        cfg = 1.0
+                        cur_negative_prompt = [cur_negative_prompt]
+                        cur_negative_prompt = cur_negative_prompt * len(cur_positive_prompts) if len(
+                            cur_negative_prompt) != len(cur_positive_prompts) else cur_negative_prompt
+                        id_images = []
+                        for index, text in enumerate(cur_positive_prompts):
+                            id_image = pipe.generate_image(
+                                width=width,
+                                height=height,
+                                num_steps=_num_steps,
+                                cfg=cfg,
+                                guidance=guidance,
+                                seed=seed_,
+                                prompt=text,
+                                negative_prompt=cur_negative_prompt[index],
+                                cf_scheduler=cf_scheduler,
+                                denoise=denoise_or_ip_sacle,
+                                image=input_id_images_dict[character_key][0]
+                            )
+                            id_images.append(id_image)
+                    else:
+                        cfg = cfg if cfg <= 1 else cfg / 10 if 1 < cfg <= 10 else cfg / 100
+                        id_images = pipe(
+                            prompt=cur_positive_prompts,
+                            image=input_id_images_dict[character_key],
+                            strength=cfg,
+                            latents=None,
+                            num_inference_steps=_num_steps,
+                            height=height,
+                            width=width,
+                            output_type="pil",
+                            max_sequence_length=256,
+                            guidance_scale=guidance,
+                            generator=generator,
+                        ).images
+                
+                elif story_maker and not make_dual_only:
+                    img = input_id_images_dict[character_key][0]
+                    # print(character_key_str,input_id_images_dict)
+                    mask_image = input_id_img_s_dict[character_key_str][0]
+                    face_info = input_id_emb_s_dict[character_key_str][0]
+                    cloth_info = None
+                    if isinstance(condition_image, torch.Tensor):
+                        cloth_info = input_id_cloth_dict[character_key_str][0]
+                    cur_negative_prompt = [cur_negative_prompt]
+                    cur_negative_prompt = cur_negative_prompt * len(cur_positive_prompts) if len(
+                        cur_negative_prompt) != len(cur_positive_prompts) else cur_negative_prompt
+                    if id_length > 1:
+                        id_images = []
+                        for index, i in enumerate(cur_positive_prompts):
+                            id_image = pipe(
+                                image=img if not controlnet_path else [img, cn_dict[ref_indexs[index]][
+                                    0]] if cn_dict else img,
+                                mask_image=mask_image,
+                                face_info=face_info,
+                                prompt=i,
+                                negative_prompt=cur_negative_prompt[index],
+                                ip_adapter_scale=denoise_or_ip_sacle, lora_scale=0.8,
+                                controlnet_conditioning_scale=controlnet_scale,
+                                num_inference_steps=_num_steps,
+                                guidance_scale=cfg,
+                                height=height, width=width,
+                                generator=generator,
+                                cloth=cloth_info,
+                            ).images
+                            id_images.append(id_image)
+                    else:
+                        id_images = pipe(
+                            image=img if not controlnet_path else [img, cn_dict[ref_indexs[0]][0]] if cn_dict else img,
+                            mask_image=mask_image,
+                            face_info=face_info,
+                            prompt=cur_positive_prompts,
+                            negative_prompt=cur_negative_prompt,
+                            ip_adapter_scale=denoise_or_ip_sacle, lora_scale=0.8,
+                            num_inference_steps=_num_steps,
+                            guidance_scale=cfg,
+                            height=height, width=width,
+                            generator=generator,
+                            cloth=cloth_info,
+                        ).images
+                
+                else:
+                    if use_cf:
+                        cur_negative_prompt = [cur_negative_prompt]
+                        cur_negative_prompt = cur_negative_prompt * len(cur_positive_prompts) if len(
+                            cur_negative_prompt) != len(cur_positive_prompts) else cur_negative_prompt
+                        id_images = []
+                        for index, text in enumerate(cur_positive_prompts):
+                            id_image = pipe.generate_image(
+                                width=width,
+                                height=height,
+                                num_steps=_num_steps,
+                                cfg=cfg,
+                                guidance=guidance,
+                                seed=seed_,
+                                prompt=text,
+                                negative_prompt=cur_negative_prompt[index],
+                                cf_scheduler=cf_scheduler,
+                                denoise=denoise_or_ip_sacle,
+                                image=input_id_images_dict[character_key][0]
+                            )
+                            id_images.append(id_image)
+                    else:
+                        if photomake_mode == "v2":
+                            id_embeds = input_id_emb_s_dict[character_key_str][0]
+                            id_images = pipe(
+                                cur_positive_prompts,
+                                input_id_images=input_id_images_dict[character_key],
+                                num_inference_steps=_num_steps,
+                                guidance_scale=cfg,
+                                start_merge_step=start_merge_step,
+                                height=height,
+                                width=width,
+                                negative_prompt=cur_negative_prompt,
+                                id_embeds=id_embeds,
+                                generator=generator
+                            ).images
+                        else:
+                            # print("v1 mode,load_chars", cur_positive_prompts, negative_prompt,character_key )
+                            id_images = pipe(
+                                cur_positive_prompts,
+                                input_id_images=input_id_images_dict[character_key],
+                                num_inference_steps=_num_steps,
+                                guidance_scale=cfg,
+                                start_merge_step=start_merge_step,
+                                height=height,
+                                width=width,
+                                negative_prompt=cur_negative_prompt,
+                                generator=generator
+                            ).images
+            
+            else:
+                raise NotImplementedError(
+                    "You should choice between original and Photomaker!",
+                    f"But you choice {model_type}",
+                )
+            p_num += 1
+            # total_results = id_images + total_results
+            # yield total_results
+            if story_maker and not make_dual_only and id_length > 1 and model_type == "img2img":
+                for index, ind in enumerate(character_index_dict[character_key]):
+                    results_dict[ref_totals[ind]] = id_images[index]
+            elif pulid and id_length > 1 and model_type == "img2img":
+                for index, ind in enumerate(character_index_dict[character_key]):
+                    results_dict[ref_totals[ind]] = id_images[index]
+            elif kolor_face and id_length > 1 and model_type == "img2img":
+                for index, ind in enumerate(character_index_dict[character_key]):
+                    results_dict[ref_totals[ind]] = id_images[index]
+            elif use_flux and use_cf and id_length > 1 and model_type == "img2img":
+                for index, ind in enumerate(character_index_dict[character_key]):
+                    results_dict[ref_totals[ind]] = id_images[index]
+            elif use_cf and not use_flux and id_length > 1 and model_type == "img2img":
+                for index, ind in enumerate(character_index_dict[character_key]):
+                    results_dict[ref_totals[ind]] = id_images[index]
+            else:
+                for ind, img in enumerate(id_images):
+                    results_dict[ref_indexs[ind]] = img
+            # real_images = []
+            # print(results_dict)
+            yield [results_dict[ind] for ind in results_dict.keys()]
+    
+    write = False
+    if not load_chars:
+        real_prompts_inds = [
+            ind for ind in range(len(prompts)) if ind not in ref_totals
+        ]
+    else:
+        real_prompts_inds = [ind for ind in range(len(prompts))]
+    
+    real_prompt_no, negative_prompt_style = apply_style_positive(style_name, "real_prompt")
+    negative_prompt = str(negative_prompt) + str(negative_prompt_style)
+    # print(f"real_prompts_inds is {real_prompts_inds}")
+    for real_prompts_ind in real_prompts_inds:  #
+        real_prompt = replace_prompts[real_prompts_ind]
+        cur_character = get_ref_character(prompts[real_prompts_ind], character_dict)
+        
+        if model_type == "txt2img":
+            setup_seed(seed_)
+        generator = torch.Generator(device=device).manual_seed(seed_)
+        
+        if len(cur_character) > 1 and model_type == "img2img":
+            raise "Temporarily Not Support Multiple character in Ref Image Mode!"
+        cur_step = 0
+        real_prompt, negative_prompt_style_no = apply_style_positive(style_name, real_prompt)
+        print(f"Sample real_prompt : {real_prompt}")
+        if model_type == "txt2img":
+            # print(results_dict,real_prompts_ind)
+            if use_flux:
+                if not use_cf:
+                    results_dict[real_prompts_ind] = pipe(
+                        prompt=real_prompt,
+                        num_inference_steps=_num_steps,
+                        guidance_scale=guidance,
+                        output_type="pil",
+                        max_sequence_length=256,
+                        height=height,
+                        width=width,
+                        generator=torch.Generator("cpu").manual_seed(seed_)
+                    ).images[0]
+                else:
+                    cfg = 1.0
+                    results_dict[real_prompts_ind] = pipe.generate_image(
+                        width=width,
+                        height=height,
+                        num_steps=_num_steps,
+                        cfg=cfg,
+                        guidance=guidance,
+                        seed=seed_,
+                        prompt=real_prompt,
+                        negative_prompt=negative_prompt,
+                        cf_scheduler=cf_scheduler,
+                        denoise=denoise_or_ip_sacle,
+                        image=None
+                    )
+            else:
+                if use_cf:
+                    results_dict[real_prompts_ind] = pipe.generate_image(
+                        width=width,
+                        height=height,
+                        num_steps=_num_steps,
+                        cfg=cfg,
+                        guidance=guidance,
+                        seed=seed_,
+                        prompt=real_prompt,
+                        negative_prompt=negative_prompt,
+                        cf_scheduler=cf_scheduler,
+                        denoise=denoise_or_ip_sacle,
+                        image=None
+                    )
+                else:
+                    results_dict[real_prompts_ind] = pipe(
+                        real_prompt,
+                        num_inference_steps=_num_steps,
+                        guidance_scale=cfg,
+                        height=height,
+                        width=width,
+                        negative_prompt=negative_prompt,
+                        generator=generator,
+                    ).images[0]
+        
+        elif model_type == "img2img":
+            empty_img = Image.new('RGB', (height, width), (255, 255, 255))
+            if use_kolor:
+                if kolor_face:
+                    empty_image = Image.new('RGB', (336, 336), (255, 255, 255))
+                    crop_image = input_id_img_s_dict[
+                        cur_character[0]] if real_prompts_ind not in nc_indexs else empty_image
+                    face_embeds = input_id_emb_s_dict[cur_character[0]][
+                        0] if real_prompts_ind not in nc_indexs else empty_emb_zero
+                    face_embeds = face_embeds.to(device, dtype=torch.float16)
+                    results_dict[real_prompts_ind] = pipe(
+                        prompt=real_prompt,
+                        negative_prompt=negative_prompt,
+                        height=height,
+                        width=width,
+                        num_inference_steps=_num_steps,
+                        guidance_scale=cfg,
+                        num_images_per_prompt=1,
+                        generator=generator,
+                        face_crop_image=crop_image,
+                        face_insightface_embeds=face_embeds,
+                    ).images[0]
+                else:
+                    results_dict[real_prompts_ind] = pipe(
+                        prompt=real_prompt,
+                        ip_adapter_image=(
+                            input_id_images_dict[cur_character[0]]
+                            if real_prompts_ind not in nc_indexs
+                            else empty_img
+                        ),
+                        negative_prompt=negative_prompt,
+                        height=height,
+                        width=width,
+                        num_inference_steps=_num_steps,
+                        guidance_scale=cfg,
+                        num_images_per_prompt=1,
+                        generator=generator,
+                        nc_flag=True if real_prompts_ind in nc_indexs else False,  # nc_flag，用索引标记，主要控制非角色人物的生成，默认false
+                    ).images[0]
+            elif use_flux:
+                if pulid:
+                    id_embeddings = input_id_emb_s_dict[cur_character[0]][
+                        0] if real_prompts_ind not in nc_indexs else empty_emb_zero
+                    uncond_id_embeddings = input_id_emb_un_dict[cur_character[0]][
+                        0] if real_prompts_ind not in nc_indexs else empty_emb_zero
+                    results_dict[real_prompts_ind] = pipe.generate_image(
+                        prompt=real_prompt,
+                        seed=seed_,
+                        start_step=2,
+                        num_steps=_num_steps,
+                        height=height,
+                        width=width,
+                        id_embeddings=id_embeddings,
+                        uncond_id_embeddings=uncond_id_embeddings,
+                        id_weight=1,
+                        guidance=guidance,
+                        true_cfg=1.0,
+                        max_sequence_length=128,
+                    )
+                if use_cf:
+                    cfg = 1.0
+                    results_dict[real_prompts_ind] = pipe.generate_image(
+                        width=width,
+                        height=height,
+                        num_steps=_num_steps,
+                        cfg=cfg,
+                        guidance=guidance,
+                        seed=seed_,
+                        prompt=real_prompt,
+                        negative_prompt=negative_prompt,
+                        cf_scheduler=cf_scheduler,
+                        denoise=denoise_or_ip_sacle,
+                        image=(input_id_images_dict[cur_character[0]]
+                               if real_prompts_ind not in nc_indexs
+                               else empty_img
+                               )
+                    )
+                else:
+                    cfg = cfg if cfg <= 1 else cfg / 10 if 1 < cfg <= 10 else cfg / 100
+                    results_dict[real_prompts_ind] = pipe(
+                        prompt=real_prompt,
+                        image=(
+                            input_id_images_dict[cur_character[0]]
+                            if real_prompts_ind not in nc_indexs
+                            else empty_img
+                        ),
+                        latents=None,
+                        strength=cfg,
+                        num_inference_steps=_num_steps,
+                        height=height,
+                        width=width,
+                        output_type="pil",
+                        max_sequence_length=256,
+                        guidance_scale=guidance,
+                        generator=generator,
+                    ).images[0]
+            elif story_maker and not make_dual_only:
+                mask_image = input_id_img_s_dict[cur_character[0]][0]
+                img_2 = input_id_images_dict[cur_character[0]][0] if real_prompts_ind not in nc_indexs else empty_img
+                cloth_info = None
+                if isinstance(condition_image, torch.Tensor):
+                    if controlnet_path:
+                        cn_img = input_id_cloth_dict[cur_character[0]][0]
+                        img_2 = [img_2, cn_img]
+                    else:
+                        cloth_info = input_id_cloth_dict[cur_character[0]][0]
+                face_info = input_id_emb_s_dict[cur_character[0]][
+                    0] if real_prompts_ind not in nc_indexs else empty_emb_zero
+                
+                results_dict[real_prompts_ind] = pipe(
+                    image=img_2 if not controlnet_path else [img_2, cn_dict[real_prompts_ind][0]] if cn_dict else img_2,
+                    mask_image=mask_image,
+                    face_info=face_info,
+                    prompt=real_prompt,
+                    negative_prompt=negative_prompt,
+                    ip_adapter_scale=denoise_or_ip_sacle, lora_scale=0.8,
+                    num_inference_steps=_num_steps,
+                    guidance_scale=cfg,
+                    controlnet_conditioning_scale=controlnet_scale,
+                    height=height, width=width,
+                    generator=generator,
+                    cloth=cloth_info,
+                ).images[0]
+            else:
+                if use_cf:
+                    results_dict[real_prompts_ind] = pipe.generate_image(
+                        width=width,
+                        height=height,
+                        num_steps=_num_steps,
+                        cfg=cfg,
+                        guidance=guidance,
+                        seed=seed_,
+                        prompt=real_prompt,
+                        negative_prompt=negative_prompt,
+                        cf_scheduler=cf_scheduler,
+                        denoise=denoise_or_ip_sacle,
+                        image=(input_id_images_dict[cur_character[0]]
+                               if real_prompts_ind not in nc_indexs
+                               else empty_img
+                               )
+                    )
+                else:
+                    if photomake_mode == "v2":
+                        # V2版本必须要有id_embeds，只能用input_id_images作为风格参考
+                        print(cur_character)
+                        id_embeds = input_id_emb_s_dict[cur_character[0]][
+                            0] if real_prompts_ind not in nc_indexs else empty_emb_zero
+                        results_dict[real_prompts_ind] = pipe(
+                            real_prompt,
+                            input_id_images=(
+                                input_id_images_dict[cur_character[0]]
+                                if real_prompts_ind not in nc_indexs
+                                else input_id_images_dict[character_list[0]]
+                            ),
+                            num_inference_steps=_num_steps,
+                            guidance_scale=cfg,
+                            start_merge_step=start_merge_step,
+                            height=height,
+                            width=width,
+                            negative_prompt=negative_prompt,
+                            generator=generator,
+                            id_embeds=id_embeds,
+                            nc_flag=True if real_prompts_ind in nc_indexs else False,
+                        ).images[0]
+                    else:
+                        # print(real_prompts_ind, real_prompt, "v1 mode", )
+                        results_dict[real_prompts_ind] = pipe(
+                            real_prompt,
+                            input_id_images=(
+                                input_id_images_dict[cur_character[0]]
+                                if real_prompts_ind not in nc_indexs
+                                else input_id_images_dict[character_list[0]]
+                            ),
+                            num_inference_steps=_num_steps,
+                            guidance_scale=cfg,
+                            start_merge_step=start_merge_step,
+                            height=height,
+                            width=width,
+                            negative_prompt=negative_prompt,
+                            generator=generator,
+                            nc_flag=True if real_prompts_ind in nc_indexs else False,
+                            # nc_flag，用索引标记，主要控制非角色人物的生成，默认false
+                        ).images[0]
+        
+        else:
+            raise NotImplementedError(
+                "You should choice between original and Photomaker!",
+                f"But you choice {model_type}",
+            )
+        yield [results_dict[ind] for ind in results_dict.keys()]
+    total_results = [results_dict[ind] for ind in range(len(prompts))]
+    torch.cuda.empty_cache()
+    yield total_results
+
 
 class Storydiffusion_Model_Loader:
     def __init__(self):
@@ -747,32 +1790,12 @@ class Pre_Translate_prompt:
         scene_prompts = ''.join(captions)
         return (scene_prompts,)
 
-class Story_Easy_Function:
-    @classmethod
-    def INPUT_TYPES(s):
-        return {"required": {"function_repo": ("STRING", {"default": ""}),
-                             "function_ckpt": (["none"] + folder_paths.get_filename_list("checkpoints"),),
-                             "function_cli": ("STRING", {"default": ""}),
-                             }}
-
-    RETURN_TYPES = ("CLIP",)
-    ETURN_NAMES = ("clip",)
-    FUNCTION = "funcion_main"
-    CATEGORY = "Storydiffusion"
-
-    def funcion_main(self, function_repo, function_ckpt,function_cli):
-        if "glm" in function_cli:
-           pass
-        #print("test")
-        return (function_repo,)
-
 
 NODE_CLASS_MAPPINGS = {
     "Storydiffusion_Model_Loader": Storydiffusion_Model_Loader,
     "Storydiffusion_Sampler": Storydiffusion_Sampler,
     "Pre_Translate_prompt": Pre_Translate_prompt,
     "Comic_Type": Comic_Type,
-    "Story_Easy_Function":Story_Easy_Function
 }
 
 NODE_DISPLAY_NAME_MAPPINGS = {
@@ -780,5 +1803,4 @@ NODE_DISPLAY_NAME_MAPPINGS = {
     "Storydiffusion_Sampler": "Storydiffusion_Sampler",
     "Pre_Translate_prompt": "Pre_Translate_prompt",
     "Comic_Type": "Comic_Type",
-    "Story_Easy_Function":"Story_Easy_Function"
 }