From da20a11927b115fb4e77de09a67f0f7c65a94528 Mon Sep 17 00:00:00 2001 From: smthemex <138738845+smthemex@users.noreply.github.com> Date: Mon, 1 Jul 2024 20:37:56 +0800 Subject: [PATCH] Update Storydiffusion_node.py add --- Storydiffusion_node.py | 74 +++++++++++++++++------------------------- 1 file changed, 30 insertions(+), 44 deletions(-) diff --git a/Storydiffusion_node.py b/Storydiffusion_node.py index 37b4a2a..4c75a41 100644 --- a/Storydiffusion_node.py +++ b/Storydiffusion_node.py @@ -99,7 +99,11 @@ for search_path in folder_paths.get_folder_paths("diffusers"): + [p for p in paths_a if "MistoLine" in p] + [o for o in paths_a if "lcm-sdxl" in o] + [Q for Q in paths_a if "controlnet-openpose-sdxl-1.0" in Q] - + [Z for Z in paths_a if "controlnet-scribble-sdxl-1.0" in Z]) + + [Z for Z in paths_a if "controlnet-scribble-sdxl-1.0" in Z] + + [a for a in paths_a if "controlnet-depth-sdxl-1.0" in a] + +[b for b in paths_a if "controlnet-tile-sdxl-1.0" in b] + +[c for c in paths_a if "controlnet-zoe-depth-sdxl-1.0" in c] + +[d for d in paths_a if "sdxl-controlnet-seg " in d]) if control_paths != [] or paths_a != []: control_paths = ["none"] + [x for x in control_paths if x] + [y for y in paths_a if y] @@ -112,7 +116,6 @@ scheduler_list = [ "DPM2 Karras", "DPM2 a", "DPM2 a Karras", "Heun", "LCM", "LMS", "LMS Karras", "UniPC" ] - def phi2narry(img): img = torch.from_numpy(np.array(img).astype(np.float32) / 255.0).unsqueeze(0) return img @@ -458,7 +461,8 @@ class SpatialAttnProcessor2_0(torch.nn.Module): 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 + 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, @@ -721,15 +725,16 @@ def process_generation( prompts = remove_punctuation_from_strings(prompts) prompts = [item + add_trigger_words for item in prompts] - global character_dict, character_index_dict, invert_character_index_dict, ref_indexs_dict, ref_totals + 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_) - # sa32, sa64 = sa32_, sa64_ clipped_prompts = prompts[:] nc_indexs = [] for ind, prompt in enumerate(clipped_prompts): @@ -762,7 +767,8 @@ def process_generation( replace_prompts, ref_indexs_dict, ref_totals, - ) = process_original_prompt(character_dict, prompts.copy(), id_length) + ) = process_original_prompt(character_dict, prompts, id_length) + if model_type == "img2img": # _upload_images = [_upload_images] input_id_images_dict = {} @@ -795,8 +801,8 @@ def process_generation( cur_positive_prompts, negative_prompt = apply_style( style_name, current_prompts, negative_prompt ) - #print("test2",cur_positive_prompts,"test2",negative_prompt,"test2") - #print("21",cur_positive_prompts[p_num],"21") + # print("test2",cur_positive_prompts,"test2",negative_prompt,"test2") + # print("21",cur_positive_prompts[p_num],"21") if model_type == "txt2img": id_images = pipe( cur_positive_prompts, @@ -894,23 +900,11 @@ def nomarl_upscale(img, width, height): samples = img.movedim(1, -1) img = tensor_to_image(samples) return img -def input_size_adaptation_output(img_tensor,base_in, width, height): - #basein=512 - if width == height: - img_pil = nomarl_upscale(img_tensor, base_in, base_in) # 2pil - else: - if min(1,width/ height)==1: #高 - r=height/base_in - img_pil = nomarl_upscale(img_tensor, round(width/r), base_in) # 2pil - else: #宽 - r=width/base_in - img_pil = nomarl_upscale(img_tensor, base_in, round(height/r)) # 2pil - img_pil.resize((width,height)) - return img_pil + def msdiffusion_main(pipe, image_1, image_2, prompts_dual, width, height, steps, seed, style_name,char_describe,char_origin,negative_prompt, encoder_repo, _model_type, _sd_type, lora, lora_path, lora_scale, trigger_words, ckpt_path, - original_config_file, role_scale, mask_threshold, start_step,controlnet_model_path,control_image,controlnet_scale,): + original_config_file, role_scale, mask_threshold, start_step,controlnet_model_path,control_image,controlnet_scale,layout_guidance): def get_phrases_idx(tokenizer, phrases, prompt): res = [] phrase_cnt = {} @@ -1015,7 +1009,7 @@ def msdiffusion_main(pipe, image_1, image_2, prompts_dual, width, height, steps, ms_model.to(device, dtype=torch.float16) input_images = [image_1, image_2] - #input_images = [x.convert("RGB").resize((width, height)) for x in input_images] + input_images = [x.resize((width, height)) for x in input_images] # generation configs num_samples = 1 image_ouput = [] @@ -1039,8 +1033,15 @@ def msdiffusion_main(pipe, image_1, image_2, prompts_dual, width, height, steps, item_x+=x prompts_dual = [item_x+item for item in prompts_dual ] prompts_dual = [item.replace("("," ",1).replace(")"," ",1) for item in prompts_dual] - + print(prompts_dual) torch.cuda.empty_cache() + if layout_guidance: + boxes = [[[0., 0.25, 0.4, 0.75], [0.6, 0.25, 1., 0.75]]] # man+women + else: + boxes = [[[0., 0., 0., 0.], [0., 0., 0., 0.]]] # used if you want no layout guidance + phrases = [[role_a, role_b]] + drop_grounding_tokens = [0] # set to 1 if you want to drop the grounding tokens + if controlnet_model_path!="none": d1, _, _, _ = control_image.size() if d1==1: @@ -1048,18 +1049,10 @@ def msdiffusion_main(pipe, image_1, image_2, prompts_dual, width, height, steps, else: control_img_list = torch.chunk(control_image, chunks=d1) j=0 - for i, prompt in enumerate(prompts_dual): control_image=control_img_list[j] - control_image=input_size_adaptation_output(control_image, 512, width, height) + control_image=nomarl_upscale(control_image, width, height) j+=1 - # print(i, prompt) - # prompt = prompt.replace("(", "").replace(")", "") - # boxes = [[[0., 0.25, 0.4, 0.75], [0.6, 0.25, 1., 0.75]]] # man+women - boxes = [[[0., 0., 0., 0.], [0., 0., 0., 0.]]] # used if you want no layout guidance - phrases = [[role_a, role_b]] - drop_grounding_tokens = [0] # set to 1 if you want to drop the grounding tokens - # used to get the attention map, return zero if the phrase is not in the prompt phrase_idxes = [get_phrases_idx(pipe.tokenizer, phrases[0], prompt)] eot_idxes = [[get_eot_idx(pipe.tokenizer, prompt)] * len(phrases[0])] @@ -1079,13 +1072,6 @@ def msdiffusion_main(pipe, image_1, image_2, prompts_dual, width, height, steps, image_ouput += images else: for i, prompt in enumerate(prompts_dual): - # print(i, prompt) - # prompt = prompt.replace("(", "").replace(")", "") - # boxes = [[[0., 0.25, 0.4, 0.75], [0.6, 0.25, 1., 0.75]]] # man+women - boxes = [[[0., 0., 0., 0.], [0., 0., 0., 0.]]] # used if you want no layout guidance - phrases = [[role_a, role_b]] - drop_grounding_tokens = [0] # set to 1 if you want to drop the grounding tokens - # used to get the attention map, return zero if the phrase is not in the prompt phrase_idxes = [get_phrases_idx(pipe.tokenizer, phrases[0], prompt)] eot_idxes = [[get_eot_idx(pipe.tokenizer, prompt)] * len(phrases[0])] @@ -1200,7 +1186,6 @@ class Storydiffusion_Model_Loader: sd_model_path = ckpt_path else: sd_model_path = models_dict[sd_type]["path"] # diffuser models - single_files = models_dict[sd_type]["single_files"] if sd_type != "Playground_v2p5" or sd_type != "Playground_v2p5": # load stable model model_info = models_dict[sd_type] model_info["model_type"] = model_type @@ -1300,6 +1285,7 @@ class Storydiffusion_Sampler: "controlnet_model_path": (control_paths,), "controlnet_scale": ( "FLOAT", {"default": 0.8, "min": 0.0, "max": 1.0, "step": 0.1, "round": 0.01}), + "layout_guidance": ("BOOLEAN", {"default": True},), }, "optional": {"image": ("IMAGE",), "control_image": ("IMAGE",),} @@ -1312,7 +1298,7 @@ class Storydiffusion_Sampler: def story_sampler(self, info,pipe, character_prompt,scene_prompts,split_prompt, negative_prompt, img_style, seed, steps, cfg, ip_adapter_strength, style_strength_ratio, encoder_repo, - role_scale, mask_threshold, start_step,save_character,controlnet_model_path,controlnet_scale,**kwargs): + role_scale, mask_threshold, start_step,save_character,controlnet_model_path,controlnet_scale,layout_guidance,**kwargs): model_type, ckpt_path, lora_path, original_config_file, lora, trigger_words, sd_type, lora_scale,char_files = info.split( ";") @@ -1408,9 +1394,9 @@ class Storydiffusion_Sampler: image_dual = msdiffusion_main(pipe, image_a, image_b, prompts_dual, width, height, steps, seed, img_style, char_describe,char_origin,negative_prompt, encoder_repo, model_type, sd_type, lora, lora_path,lora_scale, trigger_words, ckpt_path, original_config_file, role_scale, - mask_threshold, start_step,controlnet_model_path,control_image,controlnet_scale) + mask_threshold, start_step,controlnet_model_path,control_image,controlnet_scale,layout_guidance) j = 0 - for i in positions_dual: + for i in positions_dual: #重新将双人场景插入原序列 img = image_dual[j] image_pil_list.insert(int(i), img) j += 1