# !/usr/bin/env python # -*- coding: UTF-8 -*- import gradio as gr import numpy as np import torch import gc import copy import os import random import datetime from PIL import ImageFont from ip_adapter.attention_processor import IPAttnProcessor2_0 import sys from .utils.gradio_utils import ( character_to_dict, process_original_prompt, get_ref_character, cal_attn_mask_xl, cal_attn_indice_xl_effcient_memory, is_torch2_available, ) from PIL import Image if is_torch2_available(): from .utils.gradio_utils import AttnProcessor2_0 as AttnProcessor else: from .utils.gradio_utils import AttnProcessor from huggingface_hub import hf_hub_download from diffusers.pipelines.stable_diffusion_xl.pipeline_stable_diffusion_xl import ( StableDiffusionXLPipeline, ) from diffusers import (StableDiffusionXLPipeline, DiffusionPipeline, DDIMScheduler, ControlNetModel, KDPM2AncestralDiscreteScheduler, LMSDiscreteScheduler, AutoPipelineForInpainting, DPMSolverMultistepScheduler, DPMSolverSinglestepScheduler, EulerDiscreteScheduler, HeunDiscreteScheduler, UNet2DConditionModel, AutoPipelineForText2Image, StableDiffusionXLControlNetImg2ImgPipeline, KDPM2DiscreteScheduler, EulerAncestralDiscreteScheduler, UniPCMultistepScheduler, AutoencoderKL, StableDiffusionXLControlNetPipeline, DDPMScheduler, TCDScheduler, LCMScheduler) from diffusers.schedulers.scheduling_ddim import DDIMScheduler import torch.nn.functional as F from diffusers.utils.loading_utils import load_image from .utils.utils import get_comic from .utils.style_template import styles from .utils.load_models_utils import get_models_dict, load_models, get_instance_path import folder_paths global total_count, attn_count, cur_step, mask1024, mask4096, attn_procs, unet global sa32, sa64 global write global height, width STYLE_NAMES = list(styles.keys()) DEFAULT_STYLE_NAME = "Japanese Anime" global models_dict models_dict = get_models_dict() # 获取模型信息 device = ( "cuda" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu" ) # print(f"@@device:{device}") 自动选择GPU或CPU MAX_SEED = np.iinfo(np.int32).max dir_path = os.path.dirname(os.path.abspath(__file__)) path_dir = os.path.dirname(dir_path) file_path = os.path.dirname(path_dir) scheduler_list = [ "Euler", "Euler a", "DDIM", "DDPM", "DPM++ 2M", "DPM++ 2M Karras", "DPM++ 2M SDE", "DPM++ 2M SDE Karras", "DPM++ SDE", "DPM++ SDE Karras", "DPM2", "DPM2 Karras", "DPM2 a", "DPM2 a Karras", "Heun", "LCM", "LMS", "LMS Karras", "UniPC", "TCD" ] def phi2narry(img): img = torch.from_numpy(np.array(img).astype(np.float32) / 255.0).unsqueeze(0) return img def tensor_to_image(tensor): tensor = tensor.cpu() image_np = tensor.squeeze().mul(255).clamp(0, 255).byte().numpy() image = Image.fromarray(image_np, mode='RGB') return image def get_local_path(file_path, model_path): path = os.path.join(file_path, "models", "diffusers", model_path) model_path = os.path.normpath(path) if sys.platform.startswith('win32'): model_path = model_path.replace('\\', "/") return model_path def setup_seed(seed): torch.manual_seed(seed) if device == "cuda": torch.cuda.manual_seed_all(seed) np.random.seed(seed) random.seed(seed) torch.backends.cudnn.deterministic = True def get_image_path_list(folder_name): image_basename_list = os.listdir(folder_name) image_path_list = sorted( [os.path.join(folder_name, basename) for basename in image_basename_list] ) return image_path_list def swap_to_gallery(images): return ( gr.update(value=images, visible=True), gr.update(visible=True), gr.update(visible=False), ) def upload_example_to_gallery(images, prompt, style, negative_prompt): return ( gr.update(value=images, visible=True), gr.update(visible=True), gr.update(visible=False), ) def remove_back_to_files(): return gr.update(visible=False), gr.update(visible=False), gr.update(visible=True) def remove_tips(): return gr.update(visible=False) def apply_style_positive(style_name: str, positive: str): p, n = styles.get(style_name, styles[DEFAULT_STYLE_NAME]) return p.replace("{prompt}", positive) def apply_style(style_name: str, positives: list, negative: str = ""): p, n = styles.get(style_name, styles[DEFAULT_STYLE_NAME]) return [ p.replace("{prompt}", positive) for positive in positives ], n + " " + negative def change_visiale_by_model_type(_model_type): if _model_type == "Only Using Textual Description": return ( gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), ) elif _model_type == "Using Ref Images": return ( gr.update(visible=True), gr.update(visible=True), gr.update(visible=False), ) else: raise ValueError("Invalid model type", _model_type) def load_character_files(character_files: str): if character_files == "": raise gr.Error("Please set a character file!") character_files_arr = character_files.splitlines() primarytext = [] for character_file_name in character_files_arr: character_file = torch.load( character_file_name, map_location=torch.device("cpu") ) primarytext.append(character_file["character"] + character_file["description"]) return array2string(primarytext) def array2string(arr): stringtmp = "" for i, part in enumerate(arr): if i != len(arr) - 1: stringtmp += part + "\n" else: stringtmp += part return stringtmp def load_character_files_on_running(unet, character_files: str): if character_files == "": return False character_files_arr = character_files.splitlines() for character_file in character_files_arr: load_single_character_weights(unet, character_file) return True def save_single_character_weights(unet, character, description, filepath): """ 保存 attention_processor 类中的 id_bank GPU Tensor 列表到指定文件中。 参数: - model: 包含 attention_processor 类实例的模型。 - filepath: 权重要保存到的文件路径。 """ weights_to_save = {} weights_to_save["description"] = description weights_to_save["character"] = character for attn_name, attn_processor in unet.attn_processors.items(): if isinstance(attn_processor, SpatialAttnProcessor2_0): # 将每个 Tensor 转到 CPU 并转为列表，以确保它可以被序列化 weights_to_save[attn_name] = {} for step_key in attn_processor.id_bank[character].keys(): weights_to_save[attn_name][step_key] = [ tensor.cpu() for tensor in attn_processor.id_bank[character][step_key] ] # 使用torch.save保存权重 torch.save(weights_to_save, filepath) def load_single_character_weights(unet, filepath): """ 从指定文件中加载权重到 attention_processor 类的 id_bank 中。 参数: - model: 包含 attention_processor 类实例的模型。 - filepath: 权重文件的路径。 """ # 使用torch.load来读取权重 weights_to_load = torch.load(filepath, map_location=torch.device("cpu")) character = weights_to_load["character"] description = weights_to_load["description"] for attn_name, attn_processor in unet.attn_processors.items(): if isinstance(attn_processor, SpatialAttnProcessor2_0): # 转移权重到GPU（如果GPU可用的话）并赋值给id_bank attn_processor.id_bank[character] = {} for step_key in weights_to_load[attn_name].keys(): attn_processor.id_bank[character][step_key] = [ tensor.to(unet.device) for tensor in weights_to_load[attn_name][step_key] ] def get_scheduler(name): scheduler = False if name == "Euler": scheduler = EulerDiscreteScheduler() elif name == "Euler a": scheduler = EulerAncestralDiscreteScheduler() elif name == "DDIM": scheduler = DDIMScheduler() elif name == "DDPM": scheduler = DDPMScheduler() elif name == "DPM++ 2M": scheduler = DPMSolverMultistepScheduler() elif name == "DPM++ 2M Karras": scheduler = DPMSolverMultistepScheduler(use_karras_sigmas=True) elif name == "DPM++ 2M SDE": scheduler = DPMSolverMultistepScheduler(algorithm_type="sde-dpmsolver++") elif name == "DPM++ 2M SDE Karras": scheduler = DPMSolverMultistepScheduler(use_karras_sigmas=True, algorithm_type="sde-dpmsolver++") elif name == "DPM++ SDE": scheduler = DPMSolverSinglestepScheduler() elif name == "DPM++ SDE Karras": scheduler = DPMSolverSinglestepScheduler(use_karras_sigmas=True) elif name == "DPM2": scheduler = KDPM2DiscreteScheduler() elif name == "DPM2 Karras": scheduler = KDPM2DiscreteScheduler(use_karras_sigmas=True) elif name == "DPM2 a": scheduler = KDPM2AncestralDiscreteScheduler() elif name == "DPM2 a Karras": scheduler = KDPM2AncestralDiscreteScheduler(use_karras_sigmas=True) elif name == "Heun": scheduler = HeunDiscreteScheduler() elif name == "LCM": scheduler = LCMScheduler() elif name == "LMS": scheduler = LMSDiscreteScheduler() elif name == "LMS Karras": scheduler = LMSDiscreteScheduler(use_karras_sigmas=True) elif name == "UniPC": scheduler = UniPCMultistepScheduler() elif name == "TCD": scheduler = TCDScheduler() return scheduler 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)) def save_results(unet, img_list): # timestamp = datetime.datetime.now().strftime("%Y%m%d-%H%M%S") # folder_name = f"{timestamp}" weight_folder_name = os.path.join(dir_path, "weights") # weight_folder_name = f"{folder_name}/weights" # os.makedirs(weight_folder_name) # 创建文件夹 # if not os.path.exists(folder_name): # os.makedirs(folder_name) # os.makedirs(weight_folder_name) file_prefix = ''.join(random.choice("0123456789") for _ in range(5)) for idx, img in enumerate(img_list): file_path_res = os.path.join(file_path, "output", f"image_{file_prefix}_{idx}.png") # 图片文件名 img.save(file_path_res) file_path_res = img global character_dict # return img # for char in character_dict: # description = character_dict[char] # save_single_character_weights(unet,char,description,os.path.join(weight_folder_name, f'{char}.pt')) 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, width global character_dict, 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, width, device=self.device, dtype=self.dtype, ) if write: assert len(cur_character) == 1 if hidden_states.shape[1] == (height // 32) * (width // 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 // 32) * (width // 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) # print(len(indices)) # 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 ) # print(len(encoder_arr),encoder_hidden_states_tmp.shape) 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, width, 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, width = hidden_states.shape hidden_states = hidden_states.view( batch_size, channel, height * width ).transpose(1, 2) batch_size, sequence_length, channel = hidden_states.shape # print(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, width ) if attn.residual_connection: hidden_states = hidden_states + residual hidden_states = hidden_states / attn.rescale_output_factor return hidden_states def process_generation( _sd_type, ckpt_path, _model_type, _upload_images, _num_steps, style_name, _Ip_Adapter_Strength, _style_strength_ratio, guidance_scale, seed_, sa32_, sa64_, id_length_, general_prompt, negative_prompt, prompt_array, G_height, G_width, _comic_type, font_choice, _char_files, photomaker_path, scheduler_choice, ): # Corrected font_choice usage if len(general_prompt.splitlines()) >= 3: raise gr.Error( "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 _model_type == "Photomaker" and "img" not in general_prompt: raise gr.Error( 'Please add the triger word " img " behind the class word you want to customize, such as: man img or woman img' ) if _upload_images is None and _model_type != "original": raise gr.Error(f"Cannot find any input face image!") global sa32, sa64, id_length, total_length, attn_procs, unet, cur_model_type global write global cur_step, attn_count global height, width height = G_height width = G_width global pipe global sd_model_path, models_dict sd_model_path = models_dict[_sd_type] use_safe_tensor = True for attn_processor in pipe.unet.attn_processors.values(): if isinstance(attn_processor, SpatialAttnProcessor2_0): for values in attn_processor.id_bank.values(): del values attn_processor.id_bank = {} attn_processor.id_length = id_length attn_processor.total_length = id_length + 1 gc.collect() torch.cuda.empty_cache() if cur_model_type != _sd_type + "-" + _model_type: # load stable model # apply the style template ##### load pipe del pipe gc.collect() if device == "cuda": torch.cuda.empty_cache() model_info = models_dict[_sd_type] model_info["model_type"] = _model_type if _sd_type == "Use_Single_XL_Model": model_info["path"] = ckpt_path pipe = load_models(model_info, device=device, photomaker_path=photomaker_path) set_attention_processor(pipe.unet, id_length_, is_ipadapter=False) ##### ######################## pipe.scheduler = scheduler_choice.from_config(pipe.scheduler.config) # pipe.scheduler = DDIMScheduler.from_config(pipe.scheduler.config) pipe.enable_freeu(s1=0.6, s2=0.4, b1=1.1, b2=1.2) cur_model_type = _sd_type + "-" + _model_type pipe.enable_vae_slicing() if device != "mps": pipe.enable_model_cpu_offload() else: unet = pipe.unet # unet.set_attn_processor(copy.deepcopy(attn_procs)) load_chars = load_character_files_on_running(unet, character_files=_char_files) prompts = prompt_array.splitlines() global character_dict, character_index_dict, invert_character_index_dict, ref_indexs_dict, ref_totals character_dict, character_list = character_to_dict(general_prompt) # print(character_dict,character_list,general_prompt) 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_) sa32, sa64 = sa32_, sa64_ id_length = id_length_ 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 gr.Error( 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 ] prompts = [ prompt.rpartition("#")[0] if "#" in prompt else prompt for prompt in prompts ] # print(prompts) # 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.copy(), id_length) if _model_type != "original": input_id_images_dict = {} # if len(_upload_images) != len(character_dict.keys()): # raise gr.Error( # f"You upload images({len(_upload_images)}) is not equal to the number of characters({len(character_dict.keys())})!" # ) # _upload_images= list(_upload_images) # print(character_list) _upload_images = [_upload_images] # print(_upload_images,type(_upload_images)) for ind, img in enumerate(_upload_images): input_id_images_dict[character_list[ind]] = [img] # 已经pil转化了 不用load # input_id_images_dict[character_list[ind]] = [load_image(img)] # print(character_dict) # print(character_index_dict) # print(invert_character_index_dict) # real_prompts = prompts[id_length:] if device == "cuda": torch.cuda.empty_cache() write = True cur_step = 0 attn_count = 0 # id_prompts, negative_prompt = apply_style(style_name, id_prompts, negative_prompt) # print(id_prompts) setup_seed(seed_) total_results = [] id_images = [] results_dict = {} global cur_character if not load_chars: for character_key in character_dict.keys(): cur_character = [character_key] ref_indexs = ref_indexs_dict[character_key] #print(character_key, ref_indexs) current_prompts = [replace_prompts[ref_ind] for ref_ind in ref_indexs] #print(current_prompts) setup_seed(seed_) generator = torch.Generator(device=device).manual_seed(seed_) cur_step = 0 cur_positive_prompts, negative_prompt = apply_style( style_name, current_prompts, negative_prompt ) if _model_type == "original": id_images = pipe( cur_positive_prompts, num_inference_steps=_num_steps, guidance_scale=guidance_scale, height=height, width=width, negative_prompt=negative_prompt, generator=generator, ).images elif _model_type == "Photomaker": id_images = pipe( cur_positive_prompts, input_id_images=input_id_images_dict[character_key], num_inference_steps=_num_steps, guidance_scale=guidance_scale, start_merge_step=start_merge_step, height=height, width=width, negative_prompt=negative_prompt, generator=generator, ).images else: raise NotImplementedError( "You should choice between original and Photomaker!", f"But you choice {_model_type}", ) # total_results = id_images + total_results # yield total_results # print(id_images) for ind, img in enumerate(id_images): results_dict[ref_indexs[ind]] = img # real_images = [] 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))] #print(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) #print(cur_character, real_prompt) setup_seed(seed_) if len(cur_character) > 1 and _model_type == "Photomaker": raise gr.Error( "Temporarily Not Support Multiple character in Ref Image Mode!" ) generator = torch.Generator(device=device).manual_seed(seed_) cur_step = 0 real_prompt = apply_style_positive(style_name, real_prompt) if _model_type == "original": results_dict[real_prompts_ind] = pipe( real_prompt, num_inference_steps=_num_steps, guidance_scale=guidance_scale, height=height, width=width, negative_prompt=negative_prompt, generator=generator, ).images[0] elif _model_type == "Photomaker": 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=guidance_scale, 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, ).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))] if _comic_type != "No typesetting (default)": captions = prompt_array.splitlines() captions = [caption.replace("[NC]", "") for caption in captions] captions = [ caption.split("#")[-1] if "#" in caption else caption for caption in captions ] # font_path = os.path.join("fonts", font_choice) font = ImageFont.truetype(font_choice, int(45)) total_results = ( get_comic(total_results, _comic_type, captions=captions, font=font) + total_results ) save_results(pipe.unet, total_results) yield total_results class Storydiffusion_Text2Img: def __init__(self): pass @classmethod def INPUT_TYPES(cls): return { "required": { "sd_type": (["SDXL", "Use_Single_XL_Model","SDXL_Flash","RealVision", "Unstable", ],), "ckpt_name": (folder_paths.get_filename_list("checkpoints"),), "photomake_model": (["none"] + folder_paths.get_filename_list("photomaker"),), "character_prompt": ("STRING", {"multiline": True, "default": "[Bob] A man, wearing a black sui,\n" "[Alice]a woman, wearing a white shirt."}), "scene_prompt": ("STRING", {"multiline": True, "default": "[Bob] at home, read new paper #at home, The newspaper says " "there is a treasure house in the forest.\n[Bob] on the road, " "near the forest\n[Alice] is make a call at home \n[NC]A tiger " "appeared in the forest, at night \n[NC] The car on the road, " "near the forest #They drives to the forest in search of " "treasure.at night #He is overjoyed inside the house.\n[Alice] very frightened, open mouth, in the forest, at night."}), "character_id_number": ("INT", {"default": 2, "min": 1, "max": 4, "step": 1, "display": "number"}), "negative_prompt": ("STRING", {"multiline": True, "default": "bad anatomy, bad hands, missing fingers, extra fingers, " "three hands, three legs, bad arms, missing legs, " "missing arms, poorly drawn face, bad face, fused face, " "cloned face, three crus, fused feet, fused thigh, " "extra crus, ugly fingers, horn, cartoon, cg, 3d, unreal, " "animate, amputation, disconnected limbs"}), "scheduler": (scheduler_list,), "font": (["Inkfree", "ariali"],), "img_style": ( ["No style", "Japanese Anime", "Cinematic", "Disney Charactor", "Photographic", "Comic book", "Line art"],), "typesetting_style": (["No typesetting (default)", "Four_Pannel", "Classic_Comic_Style"],), "sa32_degree": ( "FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.1, "round": 0.01}), "sa64_degree": ( "FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.1, "round": 0.01}), "seed": ("INT", {"default": 0, "min": 0, "max": MAX_SEED}), "steps": ("INT", {"default": 20, "min": 1, "max": 100}), "cfg": ("FLOAT", {"default": 7, "min": 0.1, "max": 10.0, "step": 0.1, "round": 0.01}), "img_height": ("INT", {"default": 768, "min": 256, "max": 2048, "step": 32, "display": "number"}), "img_width": ("INT", {"default": 768, "min": 256, "max": 2048, "step": 32, "display": "number"}) } } RETURN_TYPES = ("IMAGE",) RETURN_NAMES = ("image",) FUNCTION = "text2image" CATEGORY = "Storydiffusion" def text2image(self, sd_type, ckpt_name, steps, img_style, photomake_model, scheduler, cfg, seed, sa32_degree, sa64_degree, character_id_number, character_prompt, negative_prompt, scene_prompt, img_height, img_width, typesetting_style, font): # Automatically select the device scheduler_choice = get_scheduler(scheduler) local_dir = os.path.join(file_path, "models", "photomaker") photomaker_local_path = os.path.join(local_dir, photomake_model) photomaker_local_path = get_instance_path(photomaker_local_path) style_name = img_style font_path = os.path.join(dir_path, "fonts", f"{font}.ttf") ckpt_path = folder_paths.get_full_path("checkpoints", ckpt_name) ckpt_path = get_instance_path(ckpt_path) # print(ckpt_path) if photomake_model == "none": photomaker_path = hf_hub_download( repo_id="TencentARC/PhotoMaker", filename="photomaker-v1.bin", repo_type="model", local_dir=local_dir, ) else: photomaker_path = photomaker_local_path global total_count, attn_count, cur_step, mask1024, mask4096, attn_procs, unet global sa32, sa64 global write global height, width global attn_count, total_count, id_length, total_length, cur_step, cur_model_type # # assert ( # len(character_prompt.strip()) > 0 # ), "Please provide the description of the character." height = img_height width = img_width id_length = character_id_number id_length_ = character_id_number sa32 = sa32_degree sa64 = sa64_degree seed_ = seed _model_type = "not Using Ref Images" _sd_type = sd_type _upload_images = None _num_steps = steps guidance_scale = cfg general_prompt = character_prompt prompt_array = scene_prompt _comic_type = typesetting_style G_height = height G_width = width font_choice = font_path _char_files = "" img_output_pass = None _Ip_Adapter_Strength = 0.5 _style_strength_ratio = 20 sa32_ = sa32 sa64_ = sa64 attn_count = 0 total_count = 0 cur_step = 0 #id_length = 4 total_length = 5 cur_model_type = "" global attn_procs, unet attn_procs = {} ### write = False ### # sa32 = 0.5 # sa64 = 0.5 # height = 768 # width = 768 ### global pipe global sd_model_path pipe = None original_config_file = os.path.join(dir_path, 'config', 'sd_xl_base.yaml') original_config_file = get_instance_path(original_config_file) #print(original_config_file) if sd_type == "Use_Single_XL_Model": sd_model_path = ckpt_path else: sd_model_path = models_dict[sd_type]["path"] # diffuser models single_files = models_dict[sd_type]["single_files"] ### LOAD Stable Diffusion Pipeline if single_files: pipe = StableDiffusionXLPipeline.from_single_file( sd_model_path, original_config_file=original_config_file,torch_dtype=torch.float16, ) else: pipe = StableDiffusionXLPipeline.from_pretrained( sd_model_path, torch_dtype=torch.float16, use_safetensors=False ) pipe = pipe.to(device) pipe.enable_freeu(s1=0.6, s2=0.4, b1=1.1, b2=1.2) # pipe.scheduler = DDIMScheduler.from_config(pipe.scheduler.config) pipe.scheduler.set_timesteps(50) pipe.enable_vae_slicing() if device != "mps": pipe.enable_model_cpu_offload() unet = pipe.unet cur_model_type = "Unstable" + "-" + "original" ### Insert PairedAttention 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 and (name.startswith("up_blocks")): attn_procs[name] = SpatialAttnProcessor2_0(id_length=id_length) total_count += 1 else: attn_procs[name] = AttnProcessor() print("successsfully load paired self-attention") print(f"number of the processor : {total_count}") unet.set_attn_processor(copy.deepcopy(attn_procs)) global mask1024, mask4096 mask1024, mask4096 = cal_attn_mask_xl( total_length, id_length, sa32, sa64, height, width, device=device, dtype=torch.float16, ) generator_with_yield = process_generation(_sd_type, ckpt_path, _model_type, _upload_images, _num_steps, style_name, _Ip_Adapter_Strength, _style_strength_ratio, guidance_scale, seed_, sa32_, sa64_, id_length_, general_prompt, negative_prompt, prompt_array, G_height, G_width, _comic_type, font_choice, _char_files, photomaker_path, scheduler_choice) for value in generator_with_yield: pass_value = value image = phi2narry(value[0]) return (image,) class Storydiffusion_Img2Img: def __init__(self): pass @classmethod def INPUT_TYPES(cls): return { "required": { "image": ("IMAGE",), "sd_type": (["SDXL", "Use_Single_XL_Model","SDXL_Flash","RealVision", "Unstable", ],), "ckpt_name": (folder_paths.get_filename_list("checkpoints"),), "photomake_model": (["none"] + folder_paths.get_filename_list("photomaker"),), "character_prompt": ("STRING", {"multiline": True, "default": "[Taylor]a woman img, wearing a white T-shirt, blue loose hair,"}), "scene_prompt": ("STRING", {"multiline": True, "default": "[Taylor]wake up in the bed,\n[Taylor]have breakfast by the window,\n[Taylor]is walking on the road, go to company,\n[Taylor]work in the company."}), "character_id_number": ("INT", {"default": 1, "min": 1, "max": 4, "step": 1, "display": "number"}), "negative_prompt": ("STRING", {"multiline": True, "default": "bad anatomy, bad hands, missing fingers, extra fingers, " "three hands, three legs, bad arms, missing legs, " "missing arms, poorly drawn face, bad face, fused face, " "cloned face, three crus, fused feet, fused thigh, " "extra crus, ugly fingers, horn, cartoon, cg, 3d, unreal, " "animate, amputation, disconnected limbs"}), "scheduler": (scheduler_list,), "font": (["Inkfree", "ariali"],), "img_style": ( ["No style", "Japanese Anime", "Cinematic", "Disney Charactor", "Photographic", "Comic book", "Line art"],), "typesetting_style": (["No typesetting (default)", "Four_Pannel", "Classic_Comic_Style"],), "sa32_degree": ( "FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.1, "round": 0.01}), "sa64_degree": ( "FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.1, "round": 0.01}), "seed": ("INT", {"default": 0, "min": 0, "max": MAX_SEED}), "steps": ("INT", {"default": 20, "min": 1, "max": 100}), "cfg": ("FLOAT", {"default": 7, "min": 0.1, "max": 10.0, "step": 0.1, "round": 0.01}), "ip_adapter_strength": ("FLOAT", {"default": 0.5, "min": 0.1, "max": 1.0, "step": 0.1, "round": 0.01}), "style_strength_ratio": ("INT", {"default": 20, "min": 10, "max": 50, "step": 1, "display": "number"}), "img_height": ("INT", {"default": 768, "min": 256, "max": 2048, "step": 32, "display": "number"}), "img_width": ("INT", {"default": 768, "min": 256, "max": 2048, "step": 32, "display": "number"}) } } RETURN_TYPES = ("IMAGE",) RETURN_NAMES = ("image",) FUNCTION = "img2image" CATEGORY = "Storydiffusion" def img2image(self, image, sd_type,ckpt_name, photomake_model, character_prompt, scene_prompt, character_id_number, negative_prompt, scheduler, font, img_style, typesetting_style, sa32_degree, sa64_degree, seed, steps, cfg, ip_adapter_strength, style_strength_ratio, img_height, img_width, ): image_load = tensor_to_image(image) # Automatically select the device scheduler_choice = get_scheduler(scheduler) local_dir = os.path.join(file_path, "models", "photomaker") photomaker_local_path = os.path.join(local_dir, photomake_model) photomaker_local_path = get_instance_path(photomaker_local_path) ckpt_path = folder_paths.get_full_path("checkpoints", ckpt_name) ckpt_path = get_instance_path(ckpt_path) style_name = img_style font_path = os.path.join(dir_path, "fonts", f"{font}.ttf") if photomake_model == "none": photomaker_path = hf_hub_download( repo_id="TencentARC/PhotoMaker", filename="photomaker-v1.bin", repo_type="model", local_dir=local_dir, ) else: photomaker_path = photomaker_local_path global total_count, attn_count, cur_step, mask1024, mask4096, attn_procs, unet global sa32, sa64 global write global height, width global attn_count, total_count, id_length, total_length, cur_step, cur_model_type # # assert ( # len(character_prompt.strip()) > 0 # ), "Please provide the description of the character." height = img_height width = img_width id_length = character_id_number id_length_ = character_id_number sa32 = sa32_degree sa64 = sa64_degree seed_ = seed _model_type = "Using Ref Images" _sd_type = sd_type _upload_images = image_load _num_steps = steps guidance_scale = cfg general_prompt = character_prompt prompt_array = scene_prompt _comic_type = typesetting_style G_height = height G_width = width font_choice = font_path _char_files = "" img_output_pass = None _Ip_Adapter_Strength = ip_adapter_strength _style_strength_ratio = style_strength_ratio sa32_ = sa32 sa64_ = sa64 attn_count = 0 total_count = 0 cur_step = 0 # id_length = 4 total_length = 5 cur_model_type = "" global attn_procs, unet attn_procs = {} ### write = False ### # sa32 = 0.5 # sa64 = 0.5 # height = 768 # width = 768 ### global pipe global sd_model_path pipe = None if sd_type == "Use_Single_XL_Model": sd_model_path = ckpt_path else: sd_model_path = models_dict[sd_type]["path"] # diffuser models single_files = models_dict[sd_type]["single_files"] original_config_file = os.path.join(dir_path, 'config', 'sd_xl_base.yaml') original_config_file = get_instance_path(original_config_file) ### LOAD Stable Diffusion Pipeline if single_files: pipe = StableDiffusionXLPipeline.from_single_file( sd_model_path,original_config_file=original_config_file, torch_dtype=torch.float16 ) else: pipe = StableDiffusionXLPipeline.from_pretrained( sd_model_path, torch_dtype=torch.float16, use_safetensors=False ) pipe = pipe.to(device) pipe.enable_freeu(s1=0.6, s2=0.4, b1=1.1, b2=1.2) # pipe.scheduler = DDIMScheduler.from_config(pipe.scheduler.config) pipe.scheduler.set_timesteps(50) pipe.enable_vae_slicing() if device != "mps": pipe.enable_model_cpu_offload() unet = pipe.unet cur_model_type = sd_type + "-" + "original" ### Insert PairedAttention 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 and (name.startswith("up_blocks")): attn_procs[name] = SpatialAttnProcessor2_0(id_length=id_length) total_count += 1 else: attn_procs[name] = AttnProcessor() print("successsfully load paired self-attention") print(f"number of the processor : {total_count}") unet.set_attn_processor(copy.deepcopy(attn_procs)) global mask1024, mask4096 mask1024, mask4096 = cal_attn_mask_xl( total_length, id_length, sa32, sa64, height, width, device=device, dtype=torch.float16, ) generator_with_yield = process_generation(_sd_type, ckpt_path, _model_type, _upload_images, _num_steps, style_name, _Ip_Adapter_Strength, _style_strength_ratio, guidance_scale, seed_, sa32_, sa64_, id_length_, general_prompt, negative_prompt, prompt_array, G_height, G_width, _comic_type, font_choice, _char_files, photomaker_path, scheduler_choice) for value in generator_with_yield: pass_value=value image = phi2narry(value[0]) return (image,) NODE_CLASS_MAPPINGS = { "Storydiffusion_Text2Img": Storydiffusion_Text2Img, "Storydiffusion_Img2Img": Storydiffusion_Img2Img } NODE_DISPLAY_NAME_MAPPINGS = { "Storydiffusion_Text2Img": "Storydiffusion_Text2Img", "Storydiffusion_Img2Img": "Storydiffusion_Img2Img" }