# !/usr/bin/env python # -*- coding: UTF-8 -*- import random import logging import numpy as np import torch import os from PIL import ImageFont,Image import torch.nn.functional as F import copy from pathlib import PureWindowsPath from tqdm import tqdm from .utils.utils import get_comic from .model_loader_utils import (phi2narry,replicate_data_by_indices,get_float,gc_cleanup,tensor_to_image,photomaker_clip,tensortopil_list_upscale,tensortopil_list,extract_content_from_brackets_, narry_list_pil,pre_text2infer,cf_clip,get_phrases_idx_cf,get_eot_idx_cf,get_ms_phrase_emb,get_extra_function,photomaker_clip_v2,adjust_indices,load_clip_clipvsion, get_scheduler,apply_style_positive,load_lora_for_unet_only,tensortolist, nomarl_upscale,SAMPLER_NAMES,SCHEDULER_NAMES,lora_lightning_list) from .utils.gradio_utils import cal_attn_indice_xl_effcient_memory,is_torch2_available 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 folder_paths from comfy.model_management import total_vram import comfy import latent_preview device = "cuda" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu" MAX_SEED = np.iinfo(np.int32).max dir_path = os.path.dirname(os.path.abspath(__file__)) weigths_gguf_current_path = os.path.join(folder_paths.models_dir, "gguf") if not os.path.exists(weigths_gguf_current_path): os.makedirs(weigths_gguf_current_path) folder_paths.add_model_folder_path("gguf", weigths_gguf_current_path) # use gguf dir global total_count, attn_count_, cur_step, mask1024, mask4096, attn_procs_, unet_,sa32, sa64,write,height_s, width_s infer_type_g=torch.float16 if device=="cuda" else torch.float32 #TODO class EasyFunction_Lite: @classmethod def INPUT_TYPES(s): return {"required": { "repo1": ("STRING", { "default": ""}), "repo2": ("STRING", { "default": ""}), "unet":(["none"] +folder_paths.get_filename_list("diffusion_models"),), "gguf": (["none"] + folder_paths.get_filename_list("gguf"),), "clip1": (["none"] + folder_paths.get_filename_list("clip"),), "clip2": (["none"] + folder_paths.get_filename_list("clip"),), "clip_vision1": (["none"] + folder_paths.get_filename_list("clip_vision"),), "clip_vision2": (["none"] + folder_paths.get_filename_list("clip_vision"),), "lora1": (["none"] +folder_paths.get_filename_list("loras"),), "lora2": (["none"] +folder_paths.get_filename_list("loras"),), "controlnet":(["none"] +folder_paths.get_filename_list("controlnet"),), "special_mode": (["none","tag", "glm"],), "tag_temperature": ( "FLOAT", {"default": 0.7, "min": 0.1, "max": 1.0, "step": 0.1, "round": 0.01}), }} RETURN_TYPES = ("MODEL","CLIP","STORY_CONDITIONING_1") RETURN_NAMES = ("model","clip","info") FUNCTION = "easy_function_main" CATEGORY = "Storydiffusion" def easy_function_main(self, repo1,repo2,unet,gguf,clip1,clip2,clip_vision1,clip_vision2,lora1,lora2,controlnet,special_mode,tag_temperature): repo1_path=PureWindowsPath(repo1).as_posix() if repo1 else None repo2_path=PureWindowsPath(repo2).as_posix() if repo2 else None gguf_path=folder_paths.get_full_path("gguf", gguf) if gguf != "none" else None unet_path=folder_paths.get_full_path("unet", unet) if unet != "none" else None clip1_path=folder_paths.get_full_path("clip", clip1) if clip1 != "none" else None clip2_path=folder_paths.get_full_path("clip", clip2) if clip2 != "none" else None lora1_path =folder_paths.get_full_path("loras", lora1) if lora1 != "none" else None lora2_path =folder_paths.get_full_path("loras", lora2) if lora2 != "none" else None clip_vision1_path=folder_paths.get_full_path("clip_vision", clip_vision1) if clip_vision1 != "none" else None clip_vision2_path=folder_paths.get_full_path("clip_vision", clip_vision2) if clip_vision2 != "none" else None controlnet_path=folder_paths.get_full_path("controlnet", controlnet) if controlnet != "none" else None clip_glm,tag_model,pipe,svdq_repo=None,None,None,None repo_list=[i for i in [repo1_path,repo2_path] if i is not None] if repo1_path is not None or repo2_path is not None: find_svdq=[i for i in repo_list if "svdq" in i ] svdq_repo=find_svdq[0] if find_svdq else None if special_mode=="tag": from .model_loader_utils import StoryLiteTag tag_model = StoryLiteTag(device, tag_temperature,repo2_path, repo1_path) #No repo will load default elif special_mode=="glm": #kolor only from .model_loader_utils import GLM_clip if clip1_path is not None: clip_glm=GLM_clip(dir_path,clip1_path) elif clip2_path is not None and clip1_path is None: clip_glm=GLM_clip(dir_path,clip2_path) else: clip_glm=None else: pass info={"gguf_path":gguf_path,"unet_path":unet_path,"tag_model":tag_model,"clip_vision1_path":clip_vision1_path,"controlnet_path":controlnet_path, "clip1_path":clip1_path,"clip2_path":clip2_path,"repo1_path":repo1_path,"repo2_path":repo2_path,"svdq_repo":svdq_repo, "lora1_path":lora1_path,"lora2_path":lora2_path,"clip_vision2_path":clip_vision2_path,} return (pipe,clip_glm,info) class StoryDiffusion_Apply: @classmethod def INPUT_TYPES(s): return {"required": { "model": ("MODEL",), "vae": ("VAE",), "infer_mode": (["story", "classic","flux_pulid","infiniteyou","uno","realcustom","instant_character","dreamo","qwen_image","flux_omi","bagel_edit","story_maker","story_and_maker","consistory","kolor_face","msdiffusion" ],), "photomake_ckpt": (["none"] + [i for i in folder_paths.get_filename_list("photomaker") if "v1" in i or "v2" in i],), "ipadapter_ckpt": (["none"] + folder_paths.get_filename_list("photomaker"),), "quantize_mode": ([ "fp8", "nf4","fp16", ],), "lora_scale": ("FLOAT", {"default": 0.8, "min": 0.1, "max": 1.0, "step": 0.1}), "extra_function":("STRING", {"default": ""}), }, "optional":{ "info": ("STORY_CONDITIONING_1",), "CLIP_VISION": ("CLIP_VISION",), } } RETURN_TYPES = ("MODEL","DIFFCONDI",) RETURN_NAMES = ("model","switch",) FUNCTION = "main_apply" CATEGORY = "Storydiffusion" def main_apply(self,model,vae,infer_mode,photomake_ckpt,ipadapter_ckpt,quantize_mode,lora_scale,extra_function, **kwargs): print(f"infer model is {infer_mode}") extra_info=kwargs.get("info",{}) # pre data photomake_ckpt_path = None if photomake_ckpt == "none" else folder_paths.get_full_path("photomaker", photomake_ckpt) ipadapter_ckpt_path = None if ipadapter_ckpt == "none" else folder_paths.get_full_path("photomaker", ipadapter_ckpt) if extra_function: extra_function=PureWindowsPath(extra_function).as_posix() save_quantezed=True if "save" in extra_function and quantize_mode=="fp8" else False clip_vision1_path=extra_info.get("clip_vision1_path") if extra_info else None clip_vision2_path=extra_info.get("clip_vision2_path") if extra_info else None repo1_path=extra_info.get("repo1_path",None) repo2_path=extra_info.get("repo2_path",None) lora1_path=extra_info.get("lora1_path") if extra_info else None lora2_path=extra_info.get("lora2_path") if extra_info else None unet_path=extra_info.get("unet_path") if extra_info else None gguf_path=extra_info.get("gguf_path") if extra_info else None repo_list=[i for i in [repo1_path,repo2_path] if i is not None] lora_list=[i for i in [lora1_path,lora2_path] if i is not None] dreamo_version="v1.0" if "v1.0" in extra_function else "v1.1" vae_encoder,vae_downsample_factor,vae_config,vision_model_config_ar,image_proj_model,no_dif_quantization,find_Kolors=None,None,None,None,None,False,None # per clip vision CLIP_VISION=kwargs.get("CLIP_VISION") unet_type=torch.float16 #use for sdxl if infer_mode=="flux_pulid" or infer_mode=="kolor_face":# 2种加载clip vision的方式 from comfy.clip_vision import load as clip_load if CLIP_VISION is not None: clip_vision_path=CLIP_VISION elif clip_vision1_path is not None: clip_vision_path=clip_load(clip_vision1_path).model elif clip_vision2_path is not None: clip_vision_path=clip_load(clip_vision2_path).model else: if infer_mode=="kolor_face": pass else: raise ValueError("Please specify one of CLIP_VISION or clip_vision1_path or clip_vision2_path") if infer_mode=="msdiffusion" or infer_mode in ["story_maker" ,"story_and_maker"]: if CLIP_VISION is not None: pass else: from comfy.clip_vision import load as clip_load if clip_vision1_path is not None: CLIP_VISION=clip_load(clip_vision1_path).model elif clip_vision2_path is not None: CLIP_VISION=clip_load(clip_vision2_path).model else: raise ValueError("Please provide a CLIP_VISION or CLIP_VISION1 or CLIP_VISION2,Msdiffusion need a clipvison g model,story_maker need a clipvison H model") # pre dreamo lora if infer_mode =="dreamo" and lora1_path is not None and lora2_path is not None: if "distill" in lora1_path.lower(): cfg_distill_path=lora1_path dreamo_lora_path=lora2_path else: cfg_distill_path=lora2_path dreamo_lora_path=lora1_path else: cfg_distill_path=None dreamo_lora_path=None if infer_mode in ["story_maker" ,"story_and_maker"] and ipadapter_ckpt_path is None: raise "story_maker need a mask.bin" # check vram only using in flux pulid or UNO if total_vram > 45000.0: aggressive_offload = False offload = False elif 17000.0 < total_vram < 45000.0: aggressive_offload = False offload = True else: aggressive_offload = True offload = True edit_mode=None logging.info(f"total_vram is {total_vram},aggressive_offload is {aggressive_offload},offload is {offload}") if infer_mode in["story", "story_and_maker","msdiffusion"]:# mix mode,use maker or ms to make 2 roles in on image from .model_loader_utils import Loader_storydiffusion model = Loader_storydiffusion(model,photomake_ckpt_path,vae) elif infer_mode =="story_maker": from .model_loader_utils import Loader_story_maker model = Loader_story_maker(model,ipadapter_ckpt_path,vae,False,lora_scale) elif infer_mode == "flux_pulid": from .PuLID.app_flux import get_models from .model_loader_utils import Loader_Flux_Pulid if unet_path is None: raise "PuLID can't link a normal comfyui model,you need load a flux unet model " model_=get_models("flux-dev",unet_path,False,aggressive_offload,device=device,offload=offload,quantized_mode=quantize_mode,) model = Loader_Flux_Pulid(model_,model,ipadapter_ckpt_path,quantize_mode,aggressive_offload,offload,False,clip_vision_path) elif infer_mode == "infiniteyou": from .model_loader_utils import Loader_InfiniteYou assert extra_info ,"you need to provide extra_info" model,image_proj_model = Loader_InfiniteYou(extra_info,vae,quantize_mode) elif infer_mode == "consistory": from .model_loader_utils import load_pipeline_consistory model = load_pipeline_consistory(model,vae) elif infer_mode == "instant_character": from .model_loader_utils import load_pipeline_instant_character assert extra_info ,"you need to provide extra_info" model = load_pipeline_instant_character(extra_info,ipadapter_ckpt_path,vae,quantize_mode) elif infer_mode == "realcustom": from .model_loader_utils import load_pipeline_realcustom,load_realcustom_vae if ipadapter_ckpt_path is None: raise "realcustom need a realcustom model which in photomaker folder, and chocie it in ipadapter_ckpt_path" model,vision_model_config_ar,_ = load_pipeline_realcustom(model,ipadapter_ckpt_path) vae_encoder,vae_downsample_factor,vae_config=load_realcustom_vae(vae,device) elif infer_mode == "kolor_face": from .model_loader_utils import Loader_KOLOR find_Kolors =[i for i in repo_list if "kolor" in i.lower()] if not find_Kolors: raise ValueError("No Kolor model found in the repo") model = Loader_KOLOR(find_Kolors[0],clip_vision_path,ipadapter_ckpt_path) elif infer_mode == "uno": from .model_loader_utils import Loader_UNO model = Loader_UNO(extra_info,offload,quantize_mode,save_quantezed,lora_rank=512) elif infer_mode == "dreamo": from.model_loader_utils import Loader_Dreamo if dreamo_lora_path is None or cfg_distill_path is None or ipadapter_ckpt_path is None: raise "dreamo need a dreamo lora and cfg distill and turbo lora in ipadapter menu" model = Loader_Dreamo(extra_info,vae,quantize_mode,dreamo_lora_path,cfg_distill_path,ipadapter_ckpt_path,device,dreamo_version) elif infer_mode == "bagel_edit": from .Bagel.app import load_bagel_model if not repo_list : raise "EasyFunction_Lite node repo1 or repo2 must fill bagel repo" max_mem_per_gpu=str(int(total_vram/1000))+"GIB" model = load_bagel_model(repo_list[0],quantize_mode,max_mem_per_gpu) elif infer_mode == "flux_omi": from .model_loader_utils import Loader_Flux_Diffuser no_dif_quantization=True if extra_info.get("unet_path") or extra_info.get("svdq_repo") or extra_info.get("gguf_path") else False model = Loader_Flux_Diffuser(extra_info,ipadapter_ckpt_path,vae,quantize_mode) elif infer_mode == "qwen_image": from.qwen_image.inferencer import load_quwen_image df_repo=repo_list[0] if repo_list else None model,edit_mode = load_quwen_image(cpu_offload=True,cpu_offload_blocks=16,no_pin_memory=True, dir_path =dir_path, repo=df_repo,unet_path=unet_path,gguf_path=gguf_path,lora_path=lora_list[0] if lora_list else None) else: # can not choice a mode print("infer use comfyui classic mode") story_img=True if photomake_ckpt_path and infer_mode in["story","story_maker","story_and_maker","msdiffusion"] else False model_=model if infer_mode=="flux_pulid" or story_img else None switch={"infer_mode":infer_mode,"ipadapter_ckpt_path":ipadapter_ckpt_path,"photomake_ckpt_path":photomake_ckpt_path, "vision_model_config_ar":vision_model_config_ar,"no_dif_quantization":no_dif_quantization,"edit_mode":edit_mode, "lora_scale":lora_scale,"image_proj_model":image_proj_model, "vae_encoder":vae_encoder,"vae_downsample_factor":vae_downsample_factor,"vae_config":vae_config,"dreamo_version":dreamo_version, "CLIP_VISION":CLIP_VISION,"VAE":vae,"find_Kolors":find_Kolors,"model_":model_,"unet_type":unet_type,"extra_function":extra_function,} switch.update(extra_info) return (model,switch,) class StoryDiffusion_CLIPTextEncode: @classmethod def INPUT_TYPES(s): return { "required": { "clip": ("CLIP",), "switch": ("DIFFCONDI", ), "width": ("INT", {"default": 768, "min": 256, "max": 2048, "step": 16, "display": "number"}), "height": ("INT", {"default": 768, "min": 256, "max": 2048, "step": 16, "display": "number"}), "role_text": ("STRING", {"multiline": True,"default": "[Taylor] a woman img, wearing a white T-shirt, blue loose hair.\n""[Lecun] a man img,wearing a suit,black hair."}), "scene_text":("STRING", {"multiline": True, "default": "[Taylor] wake up in the bed ;\n[Taylor] have breakfast by the window;\n[Lecun] driving a car;\n[Lecun] is working."}), "pos_text": ("STRING", {"multiline": True,"default": ",best"}), "neg_text": ("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,amputation, disconnected limbs"}), "lora_trigger_words": ("STRING", {"default": "best quality"}), "add_style": (["No_style", "Realistic", "Japanese_Anime", "Digital_Oil_Painting", "Pixar_Disney_Character","Photographic", "Comic_book","Line_art", "Black_and_White_Film_Noir", "Isometric_Rooms"],), "mask_threshold": ( "FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.1, "round": 0.01}), "extra_param":("STRING", {"default": ""}), "guidance_list": ("STRING", {"multiline": True, "default": "0., 0.25, 0.4, 0.75;0.6, 0.25, 1., 0.75"}), }, "optional": { "image":("IMAGE",), "control_image":("IMAGE",), } } RETURN_TYPES = ("CONDITIONING","CONDITIONING","DIFFINFO","INT","INT",) RETURN_NAMES = ("positive", "negative","condition","width","height",) FUNCTION = "encode" CATEGORY = "Storydiffusion" def encode(self, clip,switch,width,height, role_text,scene_text,pos_text,neg_text,lora_trigger_words,add_style,mask_threshold,extra_param,guidance_list,**kwargs): infer_mode=switch.get("infer_mode") CLIP_VISION = switch.get("CLIP_VISION") extra_function=switch.get("extra_function") photomake_ckpt_path=switch.get("photomake_ckpt_path") unet_type=switch.get("unet_type") model_=switch.get("model_") vae=switch.get("VAE") image=kwargs.get("image",None) control_image=kwargs.get("control_image",None) lora_list=[i for i in [switch.get("lora1_path"),switch.get("lora2_path")] if i is not None] repo_list=[i for i in [switch.get("repo1_path"),switch.get("repo2_path")] if i is not None] use_lora =True if lora_list else False if extra_function: extra_function=PureWindowsPath(extra_function).as_posix() tag_list,text_model,vision_model,siglip_path,dino_path=None,None,None,None,None, # 反推功能 if switch.get("tag_model") is not None and isinstance(image,torch.Tensor): tag_img_list=tensortopil_list(image) tag_list=[] for i in tag_img_list: tag_text=switch.get("tag_model").run_tag(i) tag_list.append(tag_text) siglip_path_=switch.get("clip_vision1_path") dino_path_=switch.get("clip_vision2_path") clip1_path=switch.get("clip1_path") clip2_path=switch.get("clip2_path") siglip_path, dino_path = siglip_path_, dino_path_ if (siglip_path_ is not None and dino_path_ is not None and "sig" in dino_path_ and "dino" in siglip_path_): siglip_path, dino_path = dino_path_, siglip_path_i if infer_mode == "realcustom" and (siglip_path is None or dino_path is None): print ("if use realcustom mode, u must linke add_function to your node,if not ,will auto load clip_vision_path and clip_path") siglip_path="" if siglip_path is None else siglip_path dino_path="" if dino_path is None else dino_path auraface,use_photov2,img2img_mode,cached,inject,onnx_provider,dreamo_mode,trigger_words_dual,dual_lora_scale=get_extra_function(extra_function,extra_param,photomake_ckpt_path,image,infer_mode) (replace_prompts,role_index_dict,invert_role_index_dict,ref_role_index_dict,ref_role_totals,role_list,role_dict, nc_txt_list,nc_indexs,positions_index_char_1,positions_index_char_2,positions_index_dual,prompts_dual,index_char_1_list,index_char_2_list)=pre_text2infer(role_text,scene_text,lora_trigger_words,use_lora,tag_list) global character_index_dict, invert_character_index_dict, cur_character, ref_indexs_dict, ref_totals, character_dict character_index_dict=role_index_dict invert_character_index_dict=invert_role_index_dict ref_indexs_dict=ref_role_index_dict ref_totals=ref_role_totals character_dict=role_dict _, style_neg = apply_style_positive(add_style, " ") #get n neg_text = neg_text + style_neg replace_prompts=[i+pos_text for i in replace_prompts] only_role_list=[apply_style_positive(add_style,i)[0] for i in replace_prompts] if len(role_list)==1: role_key_list=role_list*len(only_role_list) else: role_key_list=replicate_data_by_indices(role_list, index_char_1_list, index_char_2_list) role_key_list=[i for i in role_key_list if i is not None] if len(role_list)>1: #重新整理prompt排序，便于后续emb和ID的对应 nc_dual_list=[] for i in nc_indexs: nc_dual_list.append(i) for i in positions_index_dual: nc_dual_list.append(i) adjusted_a = adjust_indices(index_char_1_list, nc_dual_list) adjusted_b = adjust_indices(index_char_2_list, nc_dual_list) a_list = [only_role_list[i] for i in adjusted_a if 0 <= i < len(only_role_list)] b_list = [only_role_list[i] for i in adjusted_b if 0 <= i < len(only_role_list)] inf_list_split=[a_list,b_list] else: inf_list_split=[only_role_list] if not img2img_mode and infer_mode=="flux_pulid": raise "flux_pulid mode only support image2image" if infer_mode=="msdiffusion" and not prompts_dual: raise "use msdiffusion mode need have [role1] and [role2] in sence txt." if img2img_mode and photomake_ckpt_path is None and infer_mode=="story": raise "need chocie photomake v1 or v2 when use img2img mode" if infer_mode=="story_maker" and image is None: raise "story_maker only support img2img mode,you can use story_and_maker or link a iamge" if infer_mode=="msdiffusion" and photomake_ckpt_path and not img2img_mode: raise "if use msdiffusion txt2img mode,can not use photomake_ckpt_path" uno_pe="d" #Literal['d', 'h', 'w', 'o'] only_role_emb_ne,x_1_refs_dict,image_emb,x_1_refs_dual=None,None,None,None input_id_images_dict={} image_list=[] if img2img_mode and image is not None : if len(role_list)==1: image_pil=tensortopil_list_upscale(image, width, height) input_id_images_dict[role_list[0]]=image_pil[0] image_list=[image_pil[0]] else: f1, _, _, _ = image.size() img_list = list(torch.chunk(image, chunks=f1)) image_list = [nomarl_upscale(img, width, height) for img in img_list] for index, key in enumerate(role_list): input_id_images_dict[key]=image_list[index] if img2img_mode: if infer_mode == "story_and_maker" or infer_mode == "story_maker": k1, _, _, _ = image.size() if k1 == 1: image_emb = [CLIP_VISION.encode_image(image)["penultimate_hidden_states"]] else: img_list = list(torch.chunk(image, chunks=k1)) image_emb=[] for i in img_list: image_emb.append(CLIP_VISION.encode_image(i)["penultimate_hidden_states"].to(device, dtype=unet_type)) elif infer_mode=="story": image_emb = None # story的图片模式的emb要从ip的模型里单独拿， elif infer_mode=="msdiffusion": image_emb = CLIP_VISION.encode_image(image)["penultimate_hidden_states"] #MS分情况，图生图直接在前面拿，文生图在在sample拿 if not image_emb.is_cuda:#确保emb在cuda,以及type的正确 image_emb = image_emb.to(device,dtype=unet_type) elif infer_mode == "instant_character": from .model_loader_utils import load_dual_clip,instant_character_id_clip if extra_function is not None and extra_param: if "sig" in extra_function and "dino" in extra_param: #TODO siglip_path_i=extra_function dino_path_i=extra_param else: siglip_path_i=extra_param dino_path_i=extra_function else: dino_path_i="facebook/dinov2-giant" siglip_path_i="google/siglip-so400m-patch14-384" siglip_image_encoder,siglip_image_processor,dino_image_encoder_2,dino_image_processor_2=load_dual_clip(siglip_path_i,dino_path_i,device,torch.bfloat16) image_emb=[] for img in tensortopil_list(image): id_emb=instant_character_id_clip(img,siglip_image_encoder,siglip_image_processor,dino_image_encoder_2,dino_image_processor_2,device,torch.bfloat16) image_emb.append(id_emb) siglip_image_encoder.to("cpu") dino_image_encoder_2.to("cpu") gc_cleanup() elif infer_mode == "flux_pulid": # get emb use insightface pass elif infer_mode == "bagel_edit": image_emb=input_id_images_dict elif infer_mode == "flux_omi": image_emb=input_id_images_dict elif infer_mode == "dreamo": from .model_loader_utils import Dreamo_image_encoder from huggingface_hub import hf_hub_download BEN2_path_list=[i for i in repo_list if 'BEN2_Base.pth' in i] if not BEN2_path_list: BEN2_path= hf_hub_download(repo_id='PramaLLC/BEN2', filename='BEN2_Base.pth', local_dir='ComfyUI/models') else: BEN2_path=BEN2_path_list[0] ref_list=tensortopil_list_upscale(image,width,height) if control_image is not None and dreamo_mode=="id":#如果是id加ip模式，角色可以多次换装，但是要考虑衣服的数量 control_pil_list=tensortopil_list_upscale(control_image,width,height) if len(control_pil_list) != len(only_role_list): raise "when use dreamo id + ip,control image must same size as prompts,dreamo的id模式,多少句角色prompt,就要有多少件衣服在control节点输入" else: if len(role_list)>1: nc_dual_list=[] for i in nc_indexs: nc_dual_list.append(i) for i in positions_index_dual: nc_dual_list.append(i) adjusted_a = adjust_indices(index_char_1_list, nc_dual_list) adjusted_b = adjust_indices(index_char_2_list, nc_dual_list) # print("adjusted_a",adjusted_a) # print("adjusted_b",adjusted_b) a_list = [control_pil_list[i] for i in adjusted_a if 0 <= i < len(control_pil_list)] b_list = [control_pil_list[i] for i in adjusted_b if 0 <= i < len(control_pil_list)] control_pil_list=[a_list,b_list] else: control_pil_list=[control_pil_list] else: control_pil_list=None #task_list=["ip", "id", "style"] #TODO, image_emb={} if control_pil_list is not None and dreamo_mode=="id": for key,role_img ,control_list in zip(role_list,ref_list,control_pil_list): id_role_emb=[] for i in control_list: role_emb=Dreamo_image_encoder(BEN2_path,role_img,i,dreamo_mode,"ip",ref_res=512) #TODO id_role_emb.append(role_emb) image_emb[key]=id_role_emb else: for key,role_img in zip(role_list,ref_list): role_emb=Dreamo_image_encoder(BEN2_path,role_img,None,dreamo_mode,"ip",ref_res=512) #TODO image_emb[key]=[role_emb]*len(only_role_list) #改成列表方便协同id模式 elif infer_mode == "realcustom": if "g" in os.path.basename(clip1_path).lower(): clip1_path, clip2_path = clip2_path, clip1_path text_model,vision_model=load_clip_clipvsion([clip1_path,clip2_path], [os.path.join(dir_path, "config/clip_1"),os.path.join(dir_path, "config/clip_2")], dino_path,siglip_path,switch.get("vision_model_config_ar")) elif infer_mode == "uno": from .UNO.uno.flux.pipeline import preprocess_ref ref_pil_list=tensortopil_list(image) #原方法需要图片缩放，先转回pil if prompts_dual: x_1_refs_dual_=[phi2narry(preprocess_ref(i, 320)) for i in ref_pil_list ] x_1_refs_dual=[[vae.encode(i[:,:,:,:3]).to(torch.bfloat16) for i in x_1_refs_dual_]*len(prompts_dual)]#[[1,2],[1,2]] 不考虑cn if control_image is not None: control_pil_list=tensortopil_list(control_image) else: control_pil_list=None if control_pil_list is None:# 无控制图时,默认时单图模式,x_1_refs用key取值 x_1_refs_dict={} for key ,role_pil,prompts in zip(role_list,ref_pil_list,inf_list_split): role_tensor=phi2narry(preprocess_ref(role_pil, 512))#单图默认512 x_1_refs_dict[key]=[[vae.encode(role_tensor[:,:,:,:3]).to(torch.bfloat16)]*len(prompts)] # {key:[[1],[2],key2:[[3],[4]]} else: if len(control_pil_list) != len(only_role_list): raise "control image must same size as prompts,多少句话就多少张图" else: if len(role_list)>1: nc_dual_list=[] for i in nc_indexs: nc_dual_list.append(i) for i in positions_index_dual: nc_dual_list.append(i) adjusted_a = adjust_indices(index_char_1_list, nc_dual_list) adjusted_b = adjust_indices(index_char_2_list, nc_dual_list) # print("adjusted_a",adjusted_a) # print("adjusted_b",adjusted_b) a_list = [control_pil_list[i] for i in adjusted_a if 0 <= i < len(control_pil_list)] b_list = [control_pil_list[i] for i in adjusted_b if 0 <= i < len(control_pil_list)] control_pil_list=[a_list,b_list] else: control_pil_list=[control_pil_list] x_1_refs_dict={} for key ,role_pil,control_pil in zip(role_list,ref_pil_list,control_pil_list): mix_list=[] role_tensor=phi2narry(preprocess_ref(role_pil, 320) )#多图默认320 for c in (control_pil): c_tensor=phi2narry(preprocess_ref(c, 320)) mix_list.append([vae.encode(role_tensor[:,:,:,:3]).to(torch.bfloat16),vae.encode(c_tensor[:,:,:,:3]).to(torch.bfloat16)]) x_1_refs_dict[key]=mix_list #{key:[[1,2],[3,4]]} # pre insight face model and emb if img2img_mode: if infer_mode in ["story_and_maker","story_maker","flux_pulid","kolor_face","infiniteyou"] or (use_photov2 and infer_mode=="story"): from .model_loader_utils import insight_face_loader,get_insight_dict find_mask=[i for i in repo_list if "rmgb" in i.lower()] if repo_list else None if find_mask is not None: mask_repo=find_mask[0] else: mask_repo="briaai/RMBG-1.4" app_face,pipeline_mask,app_face_=insight_face_loader(infer_mode,use_photov2, auraface,onnx_provider,mask_repo) image_list=tensortopil_list_upscale(image, 640, 640) input_id_emb_s_dict,input_id_img_s_dict,input_id_emb_un_dict,input_id_cloth_dict=get_insight_dict(app_face,app_face_,pipeline_mask,infer_mode,use_photov2,image_list, role_list,control_image,width, height,model_,switch.get("image_proj_model")) # CHECK role_list else: input_id_emb_s_dict,input_id_img_s_dict,input_id_emb_un_dict,input_id_cloth_dict={}, {}, {}, {} else: input_id_emb_s_dict,input_id_img_s_dict,input_id_emb_un_dict,input_id_cloth_dict={}, {}, {}, {} # pre clip txt emb noise_x,inp_neg_list,letent_real=[],[],{} if infer_mode=="consistory": only_role_emb=None elif infer_mode=="instant_character": from .model_loader_utils import cf_flux_prompt_clip only_role_emb={} for key ,prompts in zip(role_list,inf_list_split): emb_list_=[] for prompt in prompts: p_,pool_,ind_=cf_flux_prompt_clip(clip,prompt) emb_list_.append([p_,pool_,ind_]) only_role_emb[key]=emb_list_ elif infer_mode=="realcustom": from .model_loader_utils import realcustom_clip_emb samples_per_prompt=1 guidance_weight=3.5 role_text_list = role_text.splitlines() roel_text_c=''.join(role_text_list) if '(' in roel_text_c and ')' in roel_text_c: object_prompt = extract_content_from_brackets_(roel_text_c) # 提取prompt的object list #print(f"object_prompt:{object_prompt}") object_prompt=[i.strip() for i in object_prompt] for i in object_prompt: if " " in i: raise "when using [object],object must be a word,any blank in it will cause error." object_prompt=[i for i in object_prompt ] target_phrases = sorted(list(set(object_prompt)),key=lambda x: list(object_prompt).index(x)) # 清除同名物体,保持原有顺序 #print(f"object_prompt:{phrases}",len(phrases)) assert len(target_phrases)>=2,"when using msdiffusion ,object must be more than 2." if len(target_phrases)>2: target_phrases=target_phrases[:2] #只取前两个物体 else: raise "when using realcustom ,(objectA) and (objectA) must be in the role prompt." print(f"object_prompt:{target_phrases}") image_list=tensortopil_list_upscale(image, width, height) only_role_emb,letent_real={},{} for key ,prompts,role_image,target_phrase in zip(role_list,inf_list_split,image_list,target_phrases): emb_dict_real_list,latent_dict_list=[],[] for p,n in zip(prompts,[neg_text]*len(prompts)): emb_dict_real,latent_dict=realcustom_clip_emb(text_model,vision_model,switch.get("vae_config"),switch.get("vae_downsample_factor"),p,n,role_image,target_phrase, width,height,device,samples_per_prompt,guidance_weight) emb_dict_real_list.append(emb_dict_real) latent_dict_list.append(latent_dict) only_role_emb[key]=emb_dict_real_list letent_real[key]=latent_dict_list vision_model.to("cpu") gc_cleanup() elif infer_mode=="flux_pulid": from .PuLID.flux.util import load_clip, load_t5 from .PuLID.app_flux import get_emb_flux_pulid #repo_in="flux-dev" if not repo else repo if_repo =False t5_ = load_t5("flux-dev",clip,if_repo,device, max_length=128) clip_ = load_clip("flux-dev",clip,if_repo,device) only_role_emb,noise_x,inp_neg_list={},{},{} for key ,prompts in zip(role_list,inf_list_split): ip_emb,inp_n=[],[] for p,n in zip(prompts,[neg_text]*len(prompts)): inp,inp_neg=get_emb_flux_pulid(t5_,clip_,if_repo,p,n,width,height,num_steps=20,guidance=3.5,device=device) ip_emb.append(inp) inp_n.append(inp_neg) only_role_emb[key]=ip_emb inp_neg_list[key]=inp_n elif infer_mode == "uno": only_role_emb={} from .UNO.uno.flux.sampling import prepare_multi_ip_wrapper for key ,prompts in zip(role_list,inf_list_split): ip_emb=[] for p,x_1 in zip(prompts,x_1_refs_dict[key]): inp = prepare_multi_ip_wrapper(clip,prompt=p, ref_imgs=x_1, pe=uno_pe,device=device,h=height,w=width) ip_emb.append(inp) only_role_emb[key]=ip_emb elif infer_mode=="kolor_face": from .model_loader_utils import glm_single_encode from .kolors.models.tokenization_chatglm import ChatGLMTokenizer tokenizer = ChatGLMTokenizer.from_pretrained(os.path.join(switch.get("find_Kolors"),'text_encoder')) assert clip is not None, "clip is None,check your clip path" chatglm3_model = { 'text_encoder': clip, 'tokenizer': tokenizer } only_role_emb,only_role_emb_ne=glm_single_encode(chatglm3_model, inf_list_split,role_list, neg_text, 1) elif infer_mode=="flux_omi" and switch.get("no_dif_quantization"): #need comfyclip from .model_loader_utils import cf_flux_prompt_clip only_role_emb={} for key ,prompts in zip(role_list,inf_list_split): emb_list_=[] for prompt in prompts: p_,pool_,ind_=cf_flux_prompt_clip(clip,prompt) emb_list_.append([p_,pool_,ind_]) only_role_emb[key]=emb_list_ else: if photomake_ckpt_path is not None and img2img_mode and infer_mode in["story","story_and_maker","msdiffusion"]: #img2img模式下SDXL的story的clip要特殊处理，有2个imgencoder进程，所以分离出来 TODO if use_photov2: if len(role_list)==1: emb_dict=photomaker_clip_v2(clip,model_,only_role_list,neg_text,image_list,input_id_emb_s_dict[role_key_list[0]][0]) only_role_emb=[emb_dict] else: only_role_emb=[] for role_list_s,role_list_id,key in zip(inf_list_split,image_list,role_key_list): #print(input_id_emb_s_dict[key][0].shape) #torch.Size([512]) emb_dict=photomaker_clip_v2(clip,model_,role_list_s,neg_text,[role_list_id],input_id_emb_s_dict[key][0]) only_role_emb.append(emb_dict) else: if len(role_list)==1: emb_dict=photomaker_clip(clip,model_,only_role_list,neg_text,image_list) only_role_emb=[emb_dict] else: only_role_emb=[] for role_list_s,role_list_id in zip(inf_list_split,image_list): emb_dict=photomaker_clip(clip,model_,role_list_s,neg_text,[role_list_id]) only_role_emb.append(emb_dict) else: if infer_mode=="classic": #TODO 逆序的角色会出现iD不匹配，受影响的有story文生图 only_role_emb= cf_clip([only_role_list], clip, infer_mode,role_list) #story模式需要拆分prompt，所以这里需要传入role_list elif infer_mode=="dreamo" or infer_mode=="bagel_edit" or (infer_mode=="flux_omi" and not switch.get("no_dif_quantization")): pass # TODO 暂时不支持dreamo elif infer_mode=="qwen_image": from .qwen_image.inferencer import get_emb_data image_list_=tensortolist(image,width,height) only_role_emb= get_emb_data(clip,vae,inf_list_split,image_list_,role_list) else: only_role_emb= cf_clip(inf_list_split, clip, infer_mode,role_list) #story,story_maker,story_and_maker,msdiffusion,infinite if len (role_list)==1 : only_role_emb_dict={} only_role_emb_dict[role_list[0]]=only_role_emb # [ [cond_p, output_p],...] only_role_emb=only_role_emb_dict # pre nc txt emb if nc_txt_list and not infer_mode=="consistory": nc_txt_list=[i+pos_text for i in nc_txt_list] if photomake_ckpt_path is not None and img2img_mode and infer_mode in["story","story_maker","story_and_maker","msdiffusion"]: #img2img模式下SDXL的story的clip要特殊处理，有2个imgencoder进程，所以分离出来 TODO nc_emb=[] for i in nc_txt_list: if use_photov2: empty_emb_zero = torch.zeros_like(input_id_emb_s_dict[role_list[0]][0]).to(device,dtype=torch.float16) emb_dict_=photomaker_clip_v2(clip,model_,[i],neg_text,image_list,empty_emb_zero,nc_flag=True) else: emb_dict_=photomaker_clip(clip,model_,[i],neg_text,image_list,nc_flag=True) nc_emb.append(emb_dict_) else: if infer_mode!="kolor_face": nc_emb=cf_clip(nc_txt_list, clip, infer_mode,role_list,input_split=False) elif infer_mode=="dreamo" or infer_mode=="bagel_edit" or infer_mode=="flux_omi": pass # TODO 暂时不支持dreamo else: nc_emb,_= glm_single_encode(chatglm3_model, nc_txt_list,role_list, neg_text, 1,nc_mode=True) else: nc_emb=None # pre dual role txt emb grounding_kwargs=None cross_attention_kwargs=None if prompts_dual and infer_mode in["story_maker","story_and_maker","msdiffusion"] : #忽略不支持的模式 if infer_mode=="msdiffusion": #[A] a (pig) play whith [B] a (doll) in the garden prompts_dual=[i.replace(role_list[0] ,role_dict[role_list[0]]) for i in prompts_dual if role_list[0] in i ] prompts_dual = [i.replace(role_list[1], role_dict[role_list[1]]) for i in prompts_dual if role_list[1] in i] if '(' in prompts_dual[0] and ')' in prompts_dual[0]: object_prompt = extract_content_from_brackets_(prompts_dual[0]) # 提取prompt的object list #print(f"object_prompt:{object_prompt}") object_prompt=[i.strip() for i in object_prompt] for i in object_prompt: if " " in i: raise "when using [object],object must be a word,any blank in it will cause error." object_prompt=[i for i in object_prompt ] phrases = sorted(list(set(object_prompt)),key=lambda x: list(object_prompt).index(x)) # 清除同名物体,保持原有顺序 #print(f"object_prompt:{phrases}",len(phrases)) assert len(phrases)>=2,"when using msdiffusion ,object must be more than 2." if len(phrases)>2: phrases=phrases[:2] #只取前两个物体 else: raise "when using msdiffusion ,(objectA) and (objectA) must be in the prompt." if use_lora: prompts_dual=[i+lora_trigger_words for i in prompts_dual] prompts_dual=[apply_style_positive(add_style,i+pos_text)[0] for i in prompts_dual] #[' T a (pig) play whith a (doll) in the garden,best 8k,RAW'] prompts_dual=[i.replace("("," ").replace(")"," ") for i in prompts_dual] #clear the bracket box_add = [] # 获取预设box guidance_list = guidance_list.strip().split(";") for i in range(len(guidance_list)): box_add.append(get_float(guidance_list[i])) if mask_threshold == 0: mask_threshold = None if mask_threshold: boxes = [box_add[:2]] # 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 print(f"Roles position on {boxes}") from transformers import CLIPTokenizer tokenizer_=CLIPTokenizer.from_pretrained(os.path.join(dir_path, "local_repo/tokenizer")) for i in prompts_dual: phrase_idxes = [get_phrases_idx_cf(tokenizer_, phrases[0], i)] eot_idxes = [[get_eot_idx_cf(tokenizer_, i)] * len(phrases[0])] cross_attention_kwargs, grounding_kwargs = get_ms_phrase_emb(boxes, device, infer_type_g, [0], 1, phrase_idxes, 1, eot_idxes, phrases, clip,tokenizer_) daul_emb = cf_clip(prompts_dual, clip, infer_mode,role_list,input_split=False) else: prompts_dual=[i.replace(role_list[0] ,role_dict[role_list[0]]) for i in prompts_dual if role_list[0] in i ] prompts_dual = [i.replace(role_list[1], role_dict[role_list[1]]) for i in prompts_dual if role_list[1] in i] if use_lora: prompts_dual=[i+lora_trigger_words for i in prompts_dual] prompts_dual=[apply_style_positive(add_style,i+pos_text)[0] for i in prompts_dual] #[' The figurine play whith The pig in the garden,best 8k,RAW'] daul_emb=cf_clip(prompts_dual, clip, infer_mode,role_list,input_split=False) # maker elif prompts_dual and infer_mode == "uno": #UNO双角色图片要单独处理 prompts_dual=[i.replace(role_list[0] ,role_dict[role_list[0]]) for i in prompts_dual if role_list[0] in i ] prompts_dual = [i.replace(role_list[1], role_dict[role_list[1]]) for i in prompts_dual if role_list[1] in i] prompts_dual=[apply_style_positive(add_style,i+pos_text)[0] for i in prompts_dual] #[' The figurine play whith The pig in the garden,best 8k,RAW'] daul_emb=[] for dual_t,x_1 in zip(prompts_dual,x_1_refs_dual): # dual_t:The figurine play whith The pig in the garden best 8k,RAW inp = prepare_multi_ip_wrapper(clip,prompt=dual_t, ref_imgs=x_1, pe=uno_pe,device=device,h=height,w=width) daul_emb.append(inp) elif prompts_dual and infer_mode == "dreamo": from .model_loader_utils import Dreamo_image_encoder from huggingface_hub import hf_hub_download BEN2_path_list=[i for i in repo_list if "BEN2_Base.pth" in i] if not BEN2_path_list: BEN2_path= hf_hub_download(repo_id='PramaLLC/BEN2', filename='BEN2_Base.pth', local_dir='ComfyUI/models') else: BEN2_path=BEN2_path_list[0] ref_list=tensortopil_list_upscale(image,width,height) #task_list=["ip", "id", "style"] #TODO images_emb=Dreamo_image_encoder(BEN2_path,ref_list[0],ref_list[1],"ip","ip",ref_res=512) #TODO prompts_dual=[i.replace(role_list[0] ,role_dict[role_list[0]]) for i in prompts_dual if role_list[0] in i ] prompts_dual = [i.replace(role_list[1], role_dict[role_list[1]]) for i in prompts_dual if role_list[1] in i] prompts_dual=[apply_style_positive(add_style,i+pos_text)[0] for i in prompts_dual] #[' The figurine play whith The pig in the garden,best 8k,RAW'] daul_emb=[images_emb,prompts_dual] else: daul_emb=None # neg if infer_mode=="consistory": negative = None postive_dict={} elif infer_mode=="instant_character": postive_dict= {"role": only_role_emb, "nc": None, "daul": daul_emb} negative = None elif infer_mode=="dreamo": only_role_emb={} for key ,prompts in zip(role_list,inf_list_split): only_role_emb[key]=prompts postive_dict= {"role": only_role_emb, "nc": None, "daul": daul_emb} negative = neg_text # TODO elif infer_mode=="bagel_edit": only_role_emb={} for key ,prompts in zip(role_list,inf_list_split): only_role_emb[key]=prompts postive_dict= {"role": only_role_emb, "nc": None, "daul": None} negative = neg_text # TODO elif infer_mode=="flux_omi": if switch.get("no_dif_quantization"): postive_dict= {"role": only_role_emb, "nc": None, "daul": None} #only_role_emb:p_,pool_,ind_ negative = neg_text # TODO else: only_role_emb={} for key,prompts in zip(role_list,inf_list_split): only_role_emb[key]=prompts postive_dict= {"role": only_role_emb, "nc": None, "daul": None} negative = neg_text # TODO elif infer_mode=="realcustom": postive_dict= {"role": only_role_emb, "nc": None, "daul": daul_emb} negative=[letent_real] elif infer_mode=="flux_pulid": postive_dict= {"role": only_role_emb, "nc": None, "daul": daul_emb} #不支持NC negative = [inp_neg_list,noise_x] elif infer_mode=="uno": postive_dict= {"role": only_role_emb, "nc": None, "daul": daul_emb} #TODO negative=None elif infer_mode=="kolor_face": postive_dict = {"role": only_role_emb, "nc": nc_emb, "daul": daul_emb} negative = only_role_emb_ne[0] elif infer_mode=="qwen_image": postive_dict = {"role": only_role_emb, "nc": nc_emb, "daul": daul_emb} from .qwen_image.inferencer import get_emb_data image_list_=tensortolist(image, width, height) if not switch.get("edit_mode") else None neg_list=[[neg_text]] if len(role_list)==1 else [[neg_text],[neg_text]] negative= get_emb_data(clip,vae,neg_list,image_list_,role_list) else: tokens_n = clip.tokenize(neg_text) output_n = clip.encode_from_tokens(tokens_n, return_pooled=True, return_dict=True) #{"pooled_output":tensor} cond_n = output_n.pop("cond") if cond_n.shape[1] /77>1 and infer_mode != "classic": logging.warning("nagetive prompt'tokens length is abvoe 77,will split it.") cond_n=torch.chunk(cond_n,cond_n.shape[1] //77,dim=1)[0] if infer_mode == "classic": if nc_emb is not None: for index,i in zip(nc_indexs,nc_emb): only_role_emb.insert(index,i) if daul_emb is not None: for index,i in zip(positions_index_dual,daul_emb): only_role_emb.insert(index,i) postive_dict = only_role_emb negative = [[cond_n, output_n]] else: postive_dict = {"role": only_role_emb, "nc": nc_emb, "daul": daul_emb} negative = [cond_n, output_n] # Pre emb for maker if img2img_mode and infer_mode in ["story_and_maker","story_maker"]: from .StoryMaker.pipeline_sdxl_storymaker_wrapper import encode_prompt_image_emb_ num_images_per_prompt=1 make_img,make_mask_img,make_face_info,make_cloth_info=[],[],[],[] for key in role_list: img_ = input_id_emb_un_dict[key][0] # print(character_key_str,input_id_images_dict) mask_image_ = input_id_img_s_dict[key][0] #mask_image face_info_ = input_id_emb_s_dict[key][0] cloth_info_ = None if isinstance(control_image, torch.Tensor): cloth_info_ = input_id_cloth_dict[key][0] make_img.append(img_) make_mask_img.append(mask_image_) make_face_info.append(face_info_) make_cloth_info.append(cloth_info_) mask_image_2=None face_info_2=None cloth_2=None image_2=None prompt_image_emb_dual=None if len(role_list)>1: if isinstance(control_image, torch.Tensor): cn_image_list=tensortopil_list_upscale(control_image,width,height) if len(cn_image_list) 30: start_merge_step = 30 if infer_mode =="classic": samples_list=[] for i in positive: seed_random=random.randint(0, seed) samples=self.common_ksampler(model, seed_random, steps, cfg, sampler_name, scheduler, i, negative, latent_image, denoise=denoise) #torch.Size([1, 4, 64, 64]) samples_list.append(samples) out = latent_image.copy() out["samples"] = torch.cat(samples_list,dim=0) return (out,) else: lora_list=condition.get("lora_list",[]) trigger_words=condition.get("trigger_words","best") trigger_words_dual=condition.get("trigger_words_dual","best") lora_scale=condition.get("lora_scale",1.0) dual_lora_scale=condition.get("dual_lora_scale",1.0) if infer_mode in ["story_maker" ,"msdiffusion" ,"story_and_maker" ,"story",]: if lora_list: if len(lora_list)==2: lora_list_l = [ i for i in lora_list if os.path.basename(i) in lora_lightning_list] if lora_list_l: if 1==len(lora_list_l): model=load_lora_for_unet_only(model,lora_list_l[0],trigger_words) else: for i,j,k in zip(lora_list_l,[trigger_words,trigger_words_dual],[lora_scale,dual_lora_scale]): model=load_lora_for_unet_only(model,i,j,k) else: for i,j ,k in zip(lora_list,[trigger_words,trigger_words_dual],[lora_scale,dual_lora_scale]): model=load_lora_for_unet_only(model,i,j,k) else: model=load_lora_for_unet_only(model,lora_list[0],trigger_words,lora_scale) if condition.get("controlnet_path") is not None and infer_mode!="story": controlnet_path =condition.get("controlnet") from diffusers import ControlNetModel from safetensors.torch import load_file controlnet = ControlNetModel.from_unet(model.unet) cn_state_dict = load_file(controlnet_path, device="cpu") controlnet.load_state_dict(cn_state_dict, strict=False) del cn_state_dict gc_cleanup() controlnet.to(torch.float16) if infer_mode=="story_maker" : model.controlnet = controlnet #get emb only_role_emb=positive.get("role") nc_emb=positive.get("nc") daul_emb=positive.get("daul") daul_emb_ms=daul_emb if infer_mode in["story" ,"msdiffusion","story_and_maker"]: #三者都调用story的unet方法，只是双角色引入ms或者maker if ipadapter_ckpt_path is None and infer_mode=="msdiffusion": raise "msdiffusion need a ms_adapter.bin file at ipadapter_ckpt menu." global attn_procs_,sa32, sa64, write, height_s, width_s,attn_count_, total_count, id_length, total_length, cur_step,cur_character sa32 = sa32_degree sa64 = sa64_degree attn_count_ = 0 total_count = 0 cur_step = 0 id_length = id_len total_length = 5 attn_procs_ = {} write = False height_s = height width_s = width set_attention_processor(model.unet, id_len, is_ipadapter=False) model.scheduler = scheduler_choice.from_config(model.scheduler.config) model.enable_freeu(s1=0.6, s2=0.4, b1=1.1, b2=1.2) samples_list = [] if img2img_mode and photomake_ckpt_path: generator=torch.Generator(device=device).manual_seed(seed) #print(invert_role_index_dict) for emb_dict,r_key,key in zip(only_role_emb,role_list,invert_role_index_dict.keys()): write = True cur_character = invert_role_index_dict[key] #全局角色名称，需要 #print(input_id_images_dict) #{'[Taylor]': , '[Lecun]': } #print(r_key)# '[Taylor]' # print(emb_dict) samples = model( input_id_images=input_id_images_dict[r_key], num_inference_steps=steps, guidance_scale=cfg, start_merge_step=start_merge_step, height=height, width=width, negative_prompt=neg_text, generator=generator, emb_dict=emb_dict, )[0] # torch.Size([1, 4, 64, 64]) #print(samples.shape) samples_list.append(samples) if nc_emb: #no role 无角色的emb nc_list=[] for i in (nc_emb): write = False samples = model(input_id_images=[input_id_images_dict[role_list[0]]], num_inference_steps=steps, guidance_scale=cfg, start_merge_step=start_merge_step, height=height, width=width, negative_prompt=neg_text, generator=generator, emb_dict=i, nc_flag=True)[0] # torch.Size([1, 4, 64, 64]) nc_list.append(samples) for index_,sample_nc in zip(nc_index,nc_list): samples_list.insert(int(index_), sample_nc) else: for key in role_list: # i:{a:[[tensor,tensor],[tensor,tensor]],b:[[tensor,tensor],[tensor,tensor]]} 只跑第一张图生成ID seed_random = random.randint(0, seed) write = True cur_character = [key] samples = model(height=height, width=width, num_inference_steps=steps, guidance_scale=cfg, generator=torch.Generator(device=device).manual_seed(seed_random), prompt_embeds=only_role_emb[key][0][0], negative_prompt_embeds=negative[0], pooled_prompt_embeds=only_role_emb[key][0][1].get("pooled_output"), negative_pooled_prompt_embeds=negative[1].get("pooled_output"))[ 0] # torch.Size([1, 4, 64, 64]) samples_list.append(samples) write = False #关闭角色生成 for index,emb in enumerate(only_role_emb[key]): if index==0: continue samples_ = model(height=height, width=width, num_inference_steps=steps, guidance_scale=cfg, generator=torch.Generator(device=device).manual_seed(seed_random), prompt_embeds=emb[0], negative_prompt_embeds=negative[0], pooled_prompt_embeds=emb[1].get("pooled_output"), negative_pooled_prompt_embeds=negative[1].get("pooled_output"))[ 0] # torch.Size([1, 4, 64, 64]) samples_list.append(samples_) if nc_emb: #no role 无角色的emb for index, i in zip(nc_index, nc_emb): seed_random = random.randint(0, seed) write = False samples = model(height=height, width=width, num_inference_steps=steps, guidance_scale=cfg, generator=torch.Generator(device=device).manual_seed(seed_random), prompt_embeds=i[0], negative_prompt_embeds=negative[0], pooled_prompt_embeds=i[1].get("pooled_output"), negative_pooled_prompt_embeds=negative[1].get("pooled_output"))[ 0] # torch.Size([1, 4, 64, 64]) samples_list.insert(index, samples) if infer_mode=="msdiffusion" or infer_mode =="story_and_maker": daul_emb = None #确保pass后段dual的执行 if daul_emb_ms: # del model # gc.collect() # torch.cuda.empty_cache() VAE=condition.get("VAE") for j ,(index, i) in enumerate(zip(dual_index, daul_emb_ms)): seed_random = random.randint(0, seed) write = False if not img2img_mode: #文生图模式，以文生图第一张为ID参考拿emb CLIP_VISION = condition.get("CLIP_VISION") out_1, out_2 = {}, {} out_1["samples"]=samples_list[int(condition.get("positions_index_char_1"))] out_2["samples"] =samples_list[int(condition.get("positions_index_char_2"))] role_1=VAE.decode(out_1["samples"]) role_2 = VAE.decode(out_2["samples"]) if infer_mode=="msdiffusion": from .model_loader_utils import Infer_MSdiffusion role_tensor=torch.cat((role_1,role_2),dim=0) image_embeds=CLIP_VISION.encode_image(role_tensor)["penultimate_hidden_states"] image_embeds = image_embeds.to(device, dtype=model.unet.dtype) if controlnet: model=model #TO DO #write = False samples = Infer_MSdiffusion(model,ipadapter_ckpt_path,image_embeds, i[0],negative[0],grounding_kwargs,cross_attention_kwargs, 1,mask_threshold,height,width,steps,seed_random,cfg,i[1].get("pooled_output"),negative[1].get("pooled_output") ) else: #maker的emb要单独处理，用列表形式传入 make from .StoryMaker.pipeline_sdxl_storymaker_wrapper import encode_prompt_image_emb_ # image_embeds=[] # for X in [role_1,role_2]: # image_embeds.append(clip_vision.encode_image(X)["penultimate_hidden_states"].to(device, dtype=model.unet.dtype)) make_img,make_mask_img,make_face_info,make_cloth_info=[],[],[],[] from .model_loader_utils import insight_face_loader,get_insight_dict app_face,pipeline_mask,app_face_=insight_face_loader(infer_mode,False, False) image_list = [tensor_to_image(role_1),tensor_to_image(role_2)] input_id_emb_s_dict,input_id_img_s_dict,input_id_emb_un_dict,input_id_cloth_dict=get_insight_dict(app_face,app_face_,pipeline_mask,infer_mode,False,image_list,role_list,condition.get("control_image"),width, height) for key in role_list: img_ = input_id_emb_un_dict[key][0] # print(character_key_str,input_id_images_dict) mask_image_ = input_id_img_s_dict[key][0] #mask_image face_info_ = input_id_emb_s_dict[key][0] cloth_info_ = None if isinstance(condition.get("control_image"), torch.Tensor): cloth_info_ = input_id_cloth_dict[key][0] make_img.append(img_) make_mask_img.append(mask_image_) make_face_info.append(face_info_) make_cloth_info.append(cloth_info_) prompt_image_emb_dual,maker_control_image_dual=encode_prompt_image_emb_(make_img[0], make_img[1], make_mask_img[0], make_mask_img[1], make_face_info[0], make_face_info[1], make_cloth_info[0], make_cloth_info[1],device, num_images_per_prompt, condition.get("unet_type"), CLIP_VISION,VAE,do_classifier_free_guidance=True) prompt_image_emb_dual=len(prompts_dual)*[prompt_image_emb_dual] maker_control_image_dual=len(prompts_dual)*[maker_control_image_dual] samples = model(height=height, width=width, num_inference_steps=steps, guidance_scale=cfg, generator=torch.Generator(device=device).manual_seed(seed_random), prompt_embeds=i[0], negative_prompt_embeds=negative[0], pooled_prompt_embeds=i[1].get("pooled_output"), negative_pooled_prompt_embeds=negative[1].get("pooled_output"), prompt_image_emb=prompt_image_emb_dual[j], control_image=maker_control_image_dual[j], )[0] # torch.Size([1, 4, 64, 64]) else: if infer_mode=="msdiffusion": #del model from .model_loader_utils import Loader_storydiffusion,Infer_MSdiffusion model=Loader_storydiffusion(None,None,None,model) gc_cleanup() samples = Infer_MSdiffusion(model,ipadapter_ckpt_path,image_embeds, i[0],negative[0],grounding_kwargs,cross_attention_kwargs, 1,mask_threshold,height,width,steps,seed_random,cfg,i[1].get("pooled_output"),negative[1].get("pooled_output") ) #del model else: from .model_loader_utils import Loader_story_maker print("reload maker") model=Loader_story_maker(None,ipadapter_ckpt_path,VAE,False,condition.get("lora_scale"),UNET=model) gc_cleanup() if controlnet: model.controlnet = controlnet seed_random = random.randint(0, seed) #print(i) samples = model(height=height, width=width, num_inference_steps=steps, guidance_scale=cfg, generator=torch.Generator(device=device).manual_seed(seed_random), prompt_embeds=i[0], negative_prompt_embeds=negative[0], pooled_prompt_embeds=i[1].get("pooled_output"), negative_pooled_prompt_embeds=negative[1].get("pooled_output"), prompt_image_emb=prompt_image_emb_dual[j], control_image=maker_control_image_dual[j], )[0] # torch.Size([1, 4, 64, 64]) samples_list.insert(index,samples) if daul_emb: #双角色的emb，即便不是msdiffusion，其他方法也能用，只是ID不一致而已 for index, i in zip(dual_index, daul_emb): seed_random = random.randint(0, seed) write = False samples = model(height=height, width=width, num_inference_steps=steps, guidance_scale=cfg, generator=torch.Generator(device=device).manual_seed(seed_random), prompt_embeds=i[0], negative_prompt_embeds=negative[0], pooled_prompt_embeds=i[1].get("pooled_output"), negative_pooled_prompt_embeds=negative[1].get("pooled_output"), )[0] # torch.Size([1, 4, 64, 64]) samples_list.insert(index, samples) out = {} out["samples"] = torch.cat(samples_list, dim=0) return (out,) elif infer_mode == "story_maker" : #单纯使用maker，兼容单体双人及双人同框，目前需要修改源码，将cn的图片与衣服的emb拿出来提前处理，再传入pipe model.scheduler = scheduler_choice.from_config(model.scheduler.config) seed_random = random.randint(0, seed) samples_list = [] for key in role_list : for i,emb in enumerate(only_role_emb[key]): samples = model(height=height, width=width, num_inference_steps=steps, guidance_scale=cfg, generator=torch.Generator(device=device).manual_seed(seed_random), prompt_embeds=emb[0], negative_prompt_embeds=negative[0], pooled_prompt_embeds=emb[1].get("pooled_output"), negative_pooled_prompt_embeds=negative[1].get("pooled_output"), prompt_image_emb=prompt_image_emb[key][i], control_image=maker_control_image[key][i], )[0] # torch.Size([1, 4, 64, 64]) samples_list.append(samples) if nc_emb: #no role 无角色的emb id_embeds_z, clip_image_embeds_z, clip_face_embeds_z = torch.zeros_like(prompt_image_emb[0]), torch.zeros_like(prompt_image_emb[1]), torch.zeros_like(prompt_image_emb[2]) for index, i in zip(nc_index, nc_emb): seed_random = random.randint(0, seed) write = False samples = model(height=height, width=width, num_inference_steps=steps, guidance_scale=cfg, generator=torch.Generator(device=device).manual_seed(seed_random), prompt_embeds=i[0], negative_prompt_embeds=negative[0], pooled_prompt_embeds=i[1].get("pooled_output"), negative_pooled_prompt_embeds=negative[1].get("pooled_output"), prompt_image_emb=(id_embeds_z, clip_image_embeds_z, clip_face_embeds_z), control_image=empty_img_init, )[0] # torch.Size([1, 4, 64, 64]) samples_list.insert(index, samples) if daul_emb: for j ,(index, i) in enumerate(zip(dual_index, daul_emb)): seed_random = random.randint(0, seed) write = False samples = model(height=height, width=width, num_inference_steps=steps, guidance_scale=cfg, generator=torch.Generator(device=device).manual_seed(seed_random), prompt_embeds=i[0], negative_prompt_embeds=negative[0], pooled_prompt_embeds=i[1].get("pooled_output"), negative_pooled_prompt_embeds=negative[1].get("pooled_output"), prompt_image_emb=prompt_image_emb_dual[j], control_image=maker_control_image_dual[j], )[0] # torch.Size([1, 4, 64, 64]) samples_list.insert(index, samples) out = {} out["samples"] = torch.cat(samples_list, dim=0) return (out,) elif infer_mode == "consistory": from .consistory.consistory_run import run_batch_generation,run_anchor_generation, run_extra_generation mask_dropout = 0.5 same_latent = False n_achors = 2 role_input=condition.get("role_text").splitlines()[0] main_role = role_input.replace("]", "").replace("[", "") concept_token = [main_role] if ")" in role_input: object_role = role_input.split(")")[0].split("(")[-1] concept_token=[main_role,object_role] style = "A photo of " subject=f"a {main_role} " replace_prompts= [f'{style}{subject} {i}' for i in only_role_list] gpu = 0 torch.cuda.reset_max_memory_allocated(gpu) model.scheduler = scheduler_choice.from_config(model.scheduler.config) seed_random = random.randint(0, seed) if not cached: if not inject: model.enable_vae_slicing() model.enable_model_cpu_offload() else: model.to(torch.float16) samples_list = run_batch_generation(model, replace_prompts, concept_token,neg_text, seed,n_steps=steps,mask_dropout=mask_dropout, same_latent=same_latent, perform_injection=inject,n_achors=n_achors,cf_clip=condition.get("cf_clip")) out = {} out["samples"] = samples_list return (out,zero_tensor) else: samples_list = [] if len(replace_prompts)>2: spilit_prompt=replace_prompts[:2] else: spilit_prompt=replace_prompts anchor_out_images, anchor_cache_first_stage, anchor_cache_second_stage = run_anchor_generation( model, spilit_prompt, concept_token,neg_text, seed=seed, n_steps=steps, mask_dropout=mask_dropout, same_latent=same_latent,perform_injection=inject, cache_cpu_offloading=True,cf_clip=condition.get("cf_clip")) samples_list.append(anchor_out_images) if len(replace_prompts) > 2: left_prompt=replace_prompts[2:] else: left_prompt=replace_prompts[:1] # use default for extra_prompt in left_prompt: extra_image = run_extra_generation(model, [extra_prompt], concept_token,neg_text, anchor_cache_first_stage, anchor_cache_second_stage, seed=seed, n_steps=steps, mask_dropout=mask_dropout, same_latent=same_latent, perform_injection=inject, cache_cpu_offloading=True,cf_clip=condition.get("cf_clip")) samples_list.append(extra_image) out = {} out["samples"] = torch.cat(samples_list, dim=0) return (out,) elif infer_mode == "kolor_face": samples_list = [] for key in role_list: for emb in only_role_emb[key]: samples = model( prompt=None, prompt_embeds=emb[0], negative_prompt_embeds=negative[0], pooled_prompt_embeds=emb[1], negative_pooled_prompt_embeds=negative[1], height=height, width=width, num_inference_steps=steps, guidance_scale=cfg, num_images_per_prompt=1, generator=torch.Generator(device=device).manual_seed(seed), face_crop_image=input_id_img_s_dict[key][0], face_insightface_embeds=input_id_emb_s_dict[key][0].to(device, dtype=torch.float16), ).images #print(samples.shape) samples_list.append(samples) if nc_emb: for emb in nc_emb: samples = model( prompt=None, prompt_embeds=emb[0], negative_prompt_embeds=negative[0], pooled_prompt_embeds=emb[1], negative_pooled_prompt_embeds=negative[1], height=height, width=width, num_inference_steps=steps, guidance_scale=cfg, num_images_per_prompt=1, generator=torch.Generator(device=device).manual_seed(seed), face_crop_image=Image.new('RGB', (336, 336), (255, 255, 255)), face_insightface_embeds=torch.zeros_like(input_id_emb_s_dict[role_list[0]][0]).to(device, dtype=torch.float16), ).images #print(samples.shape) samples_list.append(samples) out = {} out["samples"] = torch.cat(samples_list, dim=0) return (out,) elif infer_mode == "infiniteyou": samples_list = [] for key in role_list: if isinstance(input_id_img_s_dict[key][0], list) and (len(input_id_img_s_dict[key][0]) < len(only_role_emb[key])):#输入的图片数量小于emb数量 cn_img=input_id_img_s_dict[key][0]+(len(only_role_emb[key])-len(input_id_img_s_dict[key][0]))*[control_image_zero] else: cn_img=input_id_img_s_dict[key][0] for index,emb in enumerate(only_role_emb[key]): samples = model (id_embed=input_id_emb_s_dict[key][0], prompt_embeds=emb[0], pooled_prompt_embeds=emb[1].get("pooled_output"), control_image=cn_img[index] if isinstance(cn_img , list) else input_id_img_s_dict[key][0], guidance_scale=cfg, num_steps=steps, seed=random.randint(0, seed), infusenet_conditioning_scale=1.0, infusenet_guidance_start=0, infusenet_guidance_end=1.0, height=height, width=width, ) samples_list.append(samples) out = {} out["samples"] = torch.cat(samples_list, dim=0) return (out,) elif infer_mode == "flux_pulid": samples_list = [] for key in role_list: for index,emb in enumerate(only_role_emb[key]): samples = model.generate_image( width=width, height=height, num_steps=steps, start_step=2, guidance=cfg, seed=random.randint(0, seed), inp=emb, inp_neg=negative[0][key][index], id_embeddings=input_id_emb_s_dict[key][0], uncond_id_embeddings=input_id_emb_un_dict[key][0], ) # torch.Size([1, 4, 64, 64]) samples_list.append(samples) out = {} out["samples"] = torch.cat(samples_list, dim=0) return (out,) elif infer_mode=="uno": samples_list = [] for key in role_list: for index,emb in enumerate(only_role_emb[key]): #print(emb) samples = model( width=width, height=height, guidance=cfg, num_steps=steps, inp_cond=emb, seed=random.randint(0, seed), ) # torch.Size([1, 4, 64, 64]) samples_list.append(samples) if daul_emb: for index, emb in zip(dual_index, daul_emb): samples = model( width=width, height=height, guidance=cfg, num_steps=steps, inp_cond=emb, seed=random.randint(0, seed) ) # torch.Size([1, 4, 64, 64]) samples_list.insert(index, samples) out = {} out["samples"] = torch.cat(samples_list, dim=0) return (out,) elif infer_mode=="realcustom": from .model_loader_utils import realcustom_infer write = False samples_list = [] for key in role_list: for emb,latent_ in zip(only_role_emb[key],negative[0][key]): samples = realcustom_infer( model, sample_steps=steps, mask_reused_step=12, emb_dict=emb, latent_dict=latent_, mask_scope=0.20, mask_strategy="max_norm", #"min_max_per_channel", "max_norm" guidance_weight=cfg, height=height, width=width, seed=random.randint(0, seed), device=device, ) samples_list.append(samples) out = {} out["samples"] = torch.cat(samples_list, dim=0) return (out,) elif infer_mode=="instant_character": samples_list = [] #model.to(device) for key,id_emb in zip(role_list,image_embeds): for emb_list in only_role_emb[key]: samples = model( prompt_embeds=emb_list[0], pooled_prompt_embeds=emb_list[1], num_inference_steps=steps, guidance_scale=cfg, subject_image=True, subject_scale=0.9, generator=torch.manual_seed(random.randint(0, seed)), text_ids=emb_list[2], subject_image_embeds_dict=id_emb, )[0] # torch.Size([1, 4, 64, 64]) #print(samples.shape) samples_list.append(samples) out = {} out["samples"] = torch.cat(samples_list, dim=0) return (out,) elif infer_mode=="dreamo": samples_list = [] cfg=4.5 if condition.get("dreamo_version")=="v1.1" else 3.5 first_step_guidance=0 for key in role_list: for index,prompt in tqdm(enumerate(only_role_emb[key]),desc=f"Processing {key}"): samples = model(prompt=prompt, width=width, height=height, num_inference_steps=steps, guidance_scale=cfg, ref_conds=image_embeds[key][index], generator=torch.Generator(device="cpu").manual_seed(seed), true_cfg_scale=1, true_cfg_start_step=0, true_cfg_end_step=0, negative_prompt=neg_text, neg_guidance_scale=cfg, first_step_guidance_scale=first_step_guidance if first_step_guidance > 0 else cfg, ).images #print(samples.shape) samples_list.append(samples) if daul_emb: for index, prompt in tqdm(zip(dual_index, daul_emb[1]),desc="Processing dual prompts"): #daul_emb [emb,prompts] samples = model(prompt=prompt, width=width, height=height, num_inference_steps=steps, guidance_scale=cfg, ref_conds=daul_emb[0], generator=torch.Generator(device="cpu").manual_seed(seed), true_cfg_scale=1, true_cfg_start_step=0, true_cfg_end_step=0, negative_prompt=neg_text, neg_guidance_scale=cfg, first_step_guidance_scale=first_step_guidance if first_step_guidance > 0 else cfg, ).images samples_list.insert(index, samples) out = {} out["samples"] = torch.cat(samples_list, dim=0) return (out,) elif infer_mode =="bagel_edit": VAE=condition.get("VAE") samples_list = [] if img2img_mode: from .Bagel.app import edit_image for key in role_list: for prompt in only_role_emb[key]: samples,text_=edit_image( model, image=image_embeds[key], prompt=prompt, show_thinking=False, cfg_text_scale=cfg, cfg_img_scale=2.0, cfg_interval=0.0, timestep_shift=3.0, num_timesteps=steps, cfg_renorm_min=0.0, cfg_renorm_type="text_channel", max_think_token_n=1024, do_sample=False, text_temperature=0.3, seed=seed) # tuple (image,str) print(f'thinking text is {text_}') samples_list.append(samples) else: from .Bagel.app import text_to_image if width==height: #todo image_ratio="1:1" elif width/height==0.75: image_ratio="4:3" elif width/height==3/4: image_ratio="3:4" elif width/height==9/16: image_ratio="9:16" elif width/height==16/9: image_ratio="16:9" else: image_ratio="1:1" for key in role_list: for prompt in only_role_emb[key]: samples,text_=text_to_image( model, prompt=prompt, show_thinking=False, cfg_text_scale=cfg, cfg_interval=0.4, timestep_shift=3.0, num_timesteps=steps, cfg_renorm_min=0.0, cfg_renorm_type="global", max_think_token_n=1024, do_sample=False, text_temperature=0.3, seed=seed, image_ratio=image_ratio) # tuple (image,str) print(f'thinking text is {text_}') samples_list.append(samples) out = {} samples_list=[VAE.encode(phi2narry(pixels)[:,:,:,:3]) for pixels in samples_list] # TODO pil to tensor out["samples"] = torch.cat(samples_list, dim=0) return (out,) elif infer_mode =="flux_omi": samples_list = [] if condition.get("no_dif_quantization"): for key in role_list: for ebm_list in only_role_emb[key]: samples=model( prompt=None, height=height, width=width, guidance_scale=cfg if cfg==3.5 else 3.5, num_inference_steps=steps, max_sequence_length=512, generator=torch.Generator("cpu").manual_seed(seed), prompt_embeds=ebm_list[0], pooled_prompt_embeds=ebm_list[1], spatial_images=[image_embeds[key]], subject_images=[], #TODO TRY ON cond_size=512, ).images #print(samples.shape) samples_list.append(samples) else: for key in role_list: for prompt in only_role_emb[key]: samples=model( prompt=prompt, height=height, width=width, guidance_scale=cfg if cfg==3.5 else 3.5, num_inference_steps=steps, max_sequence_length=512, generator=torch.Generator("cpu").manual_seed(seed), spatial_images=[image_embeds[key]], subject_images=[], #TODO TRY ON cond_size=512, ).images #print(samples.shape) samples_list.append(samples) out = {} out["samples"] = torch.cat(samples_list, dim=0) # Clear cache after generation from .OmniConsistency.infer import clear_cache clear_cache(model.transformer) return (out,) elif infer_mode =="qwen_image": from .qwen_image.inferencer import infer_qwen_image_edit samples_list = [] # import comfy.model_management as model_management # model_management.unload_all_models() image_list=condition.get("image_list") if condition.get("edit_mode") else [None,None] for key,image in zip(role_list,image_list): negative_prompt_embeds,_=negative[key][0] for prompt in only_role_emb[key]: samples=infer_qwen_image_edit( model, image=image, prompt_embeds=prompt[0], prompt_embeds_mask=torch.zeros_like(prompt[0]), negative_prompt_embeds=negative_prompt_embeds, negative_prompt_embeds_mask=torch.zeros_like(negative_prompt_embeds), seed=seed, num_inference_steps=steps, true_cfg_scale=cfg, image_latents=prompt[1], edit_mode=condition.get("edit_mode"), latents=latent_init, ) #print(samples.shape) samples_list.append(samples) out = {} out["samples"] = torch.cat(samples_list, dim=0) return (out,) else: #none: return class Comic_Type: @classmethod def INPUT_TYPES(s): return {"required": {"image": ("IMAGE",), "scene_prompts": ("STRING", {"multiline": True, "forceInput": True, "default": ""}), "fonts_list": (os.listdir(os.path.join(dir_path, "fonts")),), "text_size": ("INT", {"default": 40, "min": 1, "max": 100}), "comic_type": (["Four_Pannel", "Classic_Comic_Style"],), "split_lines": ("STRING", {"default": "；"}), }} RETURN_TYPES = ("IMAGE",) RETURN_NAMES = ("image",) FUNCTION = "comic_gen" CATEGORY = "Storydiffusion" def comic_gen(self, image, scene_prompts, fonts_list, text_size, comic_type, split_lines): result = [item for index, item in enumerate(image)] total_results = narry_list_pil(result) font_choice = os.path.join(dir_path, "fonts", fonts_list) captions = scene_prompts.splitlines() if len(captions) > 1: captions = [caption.replace("(", "").replace(")", "") if "(" or ")" in caption else caption for caption in captions] # del () captions = [caption.replace("[NC]", "") for caption in captions] captions = [caption.replace("]", "").replace("[", "") for caption in captions] captions = [ caption.split("#")[-1] if "#" in caption else caption for caption in captions ] else: prompt_array = scene_prompts.replace(split_lines, "\n") captions = prompt_array.splitlines() font = ImageFont.truetype(font_choice, text_size) images = ( get_comic(total_results, comic_type, captions=captions, font=font) + total_results ) images = phi2narry(images[0]) return (images,) class Pre_Translate_prompt: @classmethod def INPUT_TYPES(s): return {"required": {"scene_prompts": ("STRING", {"forceInput": True, "default": ""}), "keep_character_name": ("BOOLEAN", {"default": False},) }} RETURN_TYPES = ("STRING",) ETURN_NAMES = ("prompt_array",) FUNCTION = "translate_prompt" CATEGORY = "Storydiffusion" def translate_prompt(self, scene_prompts, keep_character_name): captions = scene_prompts.splitlines() if not keep_character_name: captions = [caption.split(")", 1)[-1] if ")" in caption else caption for caption in captions] # del character else: captions = [caption.replace("(", "").replace(")", "") if "(" or ")" in caption else caption for caption in captions] # del () captions = [caption.replace("[NC]", "") for caption in captions] if not keep_character_name: captions = [caption.split("]", 1)[-1] for caption in captions] else: captions = [caption.replace("]", "").replace("[", "") for caption in captions] captions = [ caption.split("#")[-1] if "#" in caption else caption for caption in captions ] scene_prompts = ''.join(captions) return (scene_prompts,) 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 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 NODE_CLASS_MAPPINGS = { "Pre_Translate_prompt": Pre_Translate_prompt, "Comic_Type": Comic_Type, "EasyFunction_Lite":EasyFunction_Lite, "StoryDiffusion_Apply":StoryDiffusion_Apply, "StoryDiffusion_CLIPTextEncode":StoryDiffusion_CLIPTextEncode, "StoryDiffusion_KSampler":StoryDiffusion_KSampler, } NODE_DISPLAY_NAME_MAPPINGS = { "Pre_Translate_prompt": "Pre_Translate_prompt", "Comic_Type": "Comic_Type", "EasyFunction_Lite":"EasyFunction_Lite", "StoryDiffusion_Apply":"StoryDiffusion_Apply", "StoryDiffusion_CLIPTextEncode":"StoryDiffusion_CLIPTextEncode", "StoryDiffusion_KSampler":"StoryDiffusion_KSampler", }