# !/usr/bin/env python
# -*- coding: UTF-8 -*-

import logging
import os
import re
import random
import torch
import gc
from omegaconf import OmegaConf
from PIL import Image
import numpy as np
import cv2
from diffusers import ( DDIMScheduler, 
                       KDPM2AncestralDiscreteScheduler, LMSDiscreteScheduler,
                        DPMSolverMultistepScheduler, DPMSolverSinglestepScheduler,
                       EulerDiscreteScheduler, HeunDiscreteScheduler,
                       KDPM2DiscreteScheduler,
                       EulerAncestralDiscreteScheduler, UniPCMultistepScheduler,
                        DDPMScheduler, LCMScheduler)

from .msdiffusion.models.projection import Resampler
from .msdiffusion.models.modelWrapper import MSAdapter as MSAdapterWarpper
from .utils.style_template import styles
from .utils.load_models_utils import  get_lora_dict
from .PuLID.pulid.utils import resize_numpy_image_long
from transformers import AutoModel, AutoTokenizer
from comfy.utils import common_upscale,ProgressBar
import folder_paths
from comfy.clip_vision import load as clip_load
from .utils.gradio_utils import process_original_text,character_to_dict
import json

cur_path = os.path.dirname(os.path.abspath(__file__))
photomaker_dir=os.path.join(folder_paths.models_dir, "photomaker")
base_pt = os.path.join(photomaker_dir,"pt")
device = "cuda" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu"

lora_get = get_lora_dict()
lora_lightning_list = lora_get["lightning_xl_lora"]


SAMPLER_NAMES = ["euler", "euler_cfg_pp", "euler_ancestral", "euler_ancestral_cfg_pp", "heun", "heunpp2","dpm_2", "dpm_2_ancestral",
                  "lms", "dpm_fast", "dpm_adaptive", "dpmpp_2s_ancestral", "dpmpp_2s_ancestral_cfg_pp", "dpmpp_sde", "dpmpp_sde_gpu",
                  "dpmpp_2m", "dpmpp_2m_sde", "dpmpp_2m_sde_gpu", "dpmpp_3m_sde", "dpmpp_3m_sde_gpu", "ddpm", "lcm",
                  "ipndm", "ipndm_v", "deis","ddim", "uni_pc", "uni_pc_bh2"]

SCHEDULER_NAMES = ["normal", "karras", "exponential", "sgm_uniform", "simple", "ddim_uniform", "beta"]


def gc_cleanup():
    gc.collect()
    torch.cuda.empty_cache()


def get_scheduler(name,scheduler_):
    scheduler = False
    if name == "euler" or name =="euler_cfg_pp":
        scheduler = EulerDiscreteScheduler()
    elif name == "euler_ancestral" or name =="euler_ancestral_cfg_pp":
        scheduler = EulerAncestralDiscreteScheduler()
    elif name == "ddim":
        scheduler = DDIMScheduler()
    elif name == "ddpm":
        scheduler = DDPMScheduler()
    elif name == "dpmpp_2m":
        scheduler = DPMSolverMultistepScheduler()
    elif name == "dpmpp_2m" and scheduler_=="karras":
        scheduler = DPMSolverMultistepScheduler(use_karras_sigmas=True)
    elif name == "dpmpp_2m_sde":
        scheduler = DPMSolverMultistepScheduler(algorithm_type="sde-dpmsolver++")
    elif name == "dpmpp_2m" and scheduler_=="karras":
        scheduler = DPMSolverMultistepScheduler(use_karras_sigmas=True, algorithm_type="sde-dpmsolver++")
    elif name == "dpmpp_sde" or name == "dpmpp_sde_gpu":
        scheduler = DPMSolverSinglestepScheduler()
    elif (name == "dpmpp_sde" or name == "dpmpp_sde_gpu") and scheduler_=="karras":
        scheduler = DPMSolverSinglestepScheduler(use_karras_sigmas=True)
    elif name == "dpm_2":
        scheduler = KDPM2DiscreteScheduler()
    elif name == "dpm_2" and scheduler_=="karras":
        scheduler = KDPM2DiscreteScheduler(use_karras_sigmas=True)
    elif name == "dpm_2_ancestral":
        scheduler = KDPM2AncestralDiscreteScheduler()
    elif name == "dpm_2_ancestral" and scheduler_=="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" and scheduler_=="karras":
        scheduler = LMSDiscreteScheduler(use_karras_sigmas=True)
    elif name == "uni_pc":
        scheduler = UniPCMultistepScheduler()
    else:
        scheduler = EulerDiscreteScheduler()
    return scheduler



def get_extra_function(extra_function,extra_param,photomake_ckpt_path,image,infer_mode):
    auraface=False
    use_photov2=False
    cached=False
    inject=False
    onnx_provider="gpu"
    img2img_mode = True if isinstance(image, torch.Tensor) else False
    trigger_words_dual="best"
    dual_lora_scale=1.0
    dreamo_mode="ip"

    if extra_function:
        extra_function = extra_function.strip().lower()
        if "auraface" in extra_function:
            auraface=True  
       
    
    if extra_param:
        extra_param = extra_param.strip().lower()
        if "cache" in extra_param:
            cached=True
        if "inject" in extra_param:
            inject=True
        if "cpu" in extra_param:
            onnx_provider="cpu"
        if "id" in extra_param:
            dreamo_mode="id"
        elif "style" in extra_param:
            dreamo_mode="style"
        if "[" in extra_param:
            trigger_words_param=extra_param.split("[")[1].split("]")[0]
            trigger_words_dual=trigger_words_param.split(",")[0]
            dual_lora_scale=float(trigger_words_param.split(",")[1])


    if isinstance(photomake_ckpt_path, str) and img2img_mode:
        use_photov2 = True if "v2" in photomake_ckpt_path else False

    return auraface,use_photov2,img2img_mode,cached,inject,onnx_provider,dreamo_mode,trigger_words_dual,dual_lora_scale

def extract_content_from_brackets(text):
    # 正则表达式匹配多对方括号内的内容
    return re.findall(r'\[(.*?)\]', text)

def phi2narry(img):
    img = torch.from_numpy(np.array(img).astype(np.float32) / 255.0).unsqueeze(0)
    return img

def tensor_to_image(tensor):
    image_np = tensor.squeeze().mul(255).clamp(0, 255).byte().numpy()
    image = Image.fromarray(image_np, mode='RGB')
    return image

def tensortopil_list(tensor_in):
    d1, _, _, _ = tensor_in.size()
    if d1 == 1:
        img_list = [tensor_to_image(tensor_in)]
    else:
        tensor_list = torch.chunk(tensor_in, chunks=d1)
        img_list=[tensor_to_image(i) for i in tensor_list]
    return img_list

def tensortopil_list_upscale(tensor_in,width,height):
    d1, _, _, _ = tensor_in.size()
    if d1 == 1:
        img_list = [nomarl_upscale(tensor_in,width,height)]
    else:
        tensor_list = torch.chunk(tensor_in, chunks=d1)
        img_list=[nomarl_upscale(i,width,height) for i in tensor_list]
    return img_list

def tensortolist(tensor_in,width,height):
    d1, _, _, _ = tensor_in.size()
    if d1 == 1:
        tensor_list = [nomarl_tensor_upscale(tensor_in,width,height)]
    else:
        tensor_list_ = torch.chunk(tensor_in, chunks=d1)
        tensor_list=[nomarl_tensor_upscale(i,width,height) for i in tensor_list_]
    return tensor_list



def nomarl_tensor_upscale(tensor, width, height):
    samples = tensor.movedim(-1, 1)
    samples = common_upscale(samples, width, height, "nearest-exact", "center")
    samples = samples.movedim(1, -1)
    return samples
def nomarl_upscale(img, width, height):
    samples = img.movedim(-1, 1)
    img = common_upscale(samples, width, height, "nearest-exact", "center")
    samples = img.movedim(1, -1)
    img = tensor_to_image(samples)
    return img
def nomarl_upscale_tensor(img, width, height):
    samples = img.movedim(-1, 1)
    img = common_upscale(samples, width, height, "nearest-exact", "center")
    samples = img.movedim(1, -1)
    return samples
    
def center_crop(img):
    width, height = img.size
    square = min(width, height)
    left = (width - square) / 2
    top = (height - square) / 2
    right = (width + square) / 2
    bottom = (height + square) / 2
    return img.crop((left, top, right, bottom))

def center_crop_s(img, new_width, new_height):
    width, height = img.size
    left = (width - new_width) / 2
    top = (height - new_height) / 2
    right = (width + new_width) / 2
    bottom = (height + new_height) / 2
    return img.crop((left, top, right, bottom))


def contains_brackets(s):
    return '[' in s or ']' in s

def has_parentheses(s):
    return bool(re.search(r'\(.*?\)', s))
def extract_content_from_brackets_(text):
    # 正则表达式匹配多对圆括号内的内容
    return re.findall(r'\((.*?)\)', text)

def narry_list(list_in):
    for i in range(len(list_in)):
        value = list_in[i]
        modified_value = phi2narry(value)
        list_in[i] = modified_value
    return list_in
def remove_punctuation_from_strings(lst):
    pattern = r"[\W]+$"  # 匹配字符串末尾的所有非单词字符
    return [re.sub(pattern, '', s) for s in lst]

def phi_list(list_in):
    for i in range(len(list_in)):
        value = list_in[i]
        list_in[i] = value
    return list_in

def narry_list_pil(list_in):
    for i in range(len(list_in)):
        value = list_in[i]
        modified_value = tensor_to_image(value)
        list_in[i] = modified_value
    return list_in


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 apply_style_positive(style_name: str, positive: str):
    p, n = styles.get(style_name, styles[style_name])
    #print(p, "test0", n)
    return p.replace("{prompt}", positive),n

def apply_style(style_name: str, positives: list, negative: str = ""):
    p, n = styles.get(style_name, styles[style_name])
    #print(p,"test1",n)
    return [
        p.replace("{prompt}", positive) for positive in positives
    ], n + " " + negative





def face_bbox_to_square(bbox):
    ## l, t, r, b to square l, t, r, b
    l,t,r,b = bbox
    cent_x = (l + r) / 2
    cent_y = (t + b) / 2
    w, h = r - l, b - t
    r = max(w, h) / 2

    l0 = cent_x - r
    r0 = cent_x + r
    t0 = cent_y - r
    b0 = cent_y + r

    return [l0, t0, r0, b0]



def insight_face_loader(infer_mode,use_photov2,auraface,onnx_provider="cpu",mask_repo="briaai/RMBG-1.4"):
    insightface_root_path= folder_paths.base_path
    if infer_mode=="story" and use_photov2:
        from .utils.insightface_package import FaceAnalysis2
        if auraface:
            from huggingface_hub import snapshot_download
            snapshot_download(
                "fal/AuraFace-v1",
                local_dir="models/auraface",
            )
            app_face = FaceAnalysis2(name="auraface",
                                     providers=["CUDAExecutionProvider", "CPUExecutionProvider"], root=insightface_root_path,
                                     allowed_modules=['detection', 'recognition'])
        else:
            app_face = FaceAnalysis2(providers=['CUDAExecutionProvider'],
                                     allowed_modules=['detection', 'recognition'])
        app_face.prepare(ctx_id=0, det_size=(640, 640))
        pipeline_mask = None
        app_face_ = None
    elif infer_mode=="kolor_face" :
        from .kolors.models.sample_ipadapter_faceid_plus import FaceInfoGenerator
        from huggingface_hub import snapshot_download
        snapshot_download(
            'DIAMONIK7777/antelopev2',
            local_dir='models/antelopev2',
        )
        app_face = FaceInfoGenerator(root_dir=insightface_root_path)
        pipeline_mask = None
        app_face_ = None
    elif infer_mode=="story_maker":
        from insightface.app import FaceAnalysis
        from transformers import pipeline
        pipeline_mask = pipeline("image-segmentation", model=mask_repo,
                                 trust_remote_code=True)

        app_face = FaceAnalysis(name='buffalo_l', root=insightface_root_path,
                                providers=['CUDAExecutionProvider', 'CPUExecutionProvider'])
        app_face.prepare(ctx_id=0, det_size=(640, 640))
        app_face_ = None

    elif infer_mode=="story_and_maker":
        from insightface.app import FaceAnalysis
        from transformers import pipeline
        pipeline_mask = pipeline("image-segmentation", model=mask_repo,
                                 trust_remote_code=True)
        if use_photov2:
            from .utils.insightface_package import FaceAnalysis2
            if auraface:
                from huggingface_hub import snapshot_download
                snapshot_download(
                    "fal/AuraFace-v1",
                    local_dir="models/auraface",
                )
                app_face = FaceAnalysis2(name="auraface",
                                            providers=["CUDAExecutionProvider", "CPUExecutionProvider"],
                                            root=insightface_root_path,
                                            allowed_modules=['detection', 'recognition'])
            else:
                app_face = FaceAnalysis2(providers=['CUDAExecutionProvider'],
                                            allowed_modules=['detection', 'recognition'])
            app_face.prepare(ctx_id=0, det_size=(640, 640))
            app_face_ = FaceAnalysis(name='buffalo_l', root=insightface_root_path,
                                        providers=['CUDAExecutionProvider', 'CPUExecutionProvider'])
            app_face_.prepare(ctx_id=0, det_size=(640, 640))
        else:
            app_face = FaceAnalysis(name='buffalo_l', root=insightface_root_path,
                                    providers=['CUDAExecutionProvider', 'CPUExecutionProvider'])
            app_face.prepare(ctx_id=0, det_size=(640, 640))
            app_face_ = None
    elif infer_mode=="flux_pulid":
        import insightface
        from facexlib.parsing import init_parsing_model
        from facexlib.utils.face_restoration_helper import FaceRestoreHelper
        from insightface.app import FaceAnalysis
        from huggingface_hub import snapshot_download
        # antelopev2
        snapshot_download('DIAMONIK7777/antelopev2', local_dir='models/antelopev2')
        
        providers = ['CPUExecutionProvider'] if onnx_provider == 'cpu' else ['CUDAExecutionProvider', 'CPUExecutionProvider']
        app_face = FaceAnalysis(name='antelopev2', root=insightface_root_path, providers=providers)
        app_face.prepare(ctx_id=0, det_size=(640, 640))
        # face_helper
        face_helper = FaceRestoreHelper(
                            upscale_factor=1,
                            face_size=512,
                            crop_ratio=(1, 1),
                            det_model='retinaface_resnet50',
                            save_ext='png',
                            device=device,)
        face_helper.face_parse = None
        face_helper.face_parse = init_parsing_model(model_name='bisenet', device=device)
        app_face_ = face_helper
        handler_ante = insightface.model_zoo.get_model('models/antelopev2/glintr100.onnx',providers=providers)
        handler_ante.prepare(ctx_id=0)
        pipeline_mask = handler_ante
    elif infer_mode=="infiniteyou":   
        from facexlib.recognition import init_recognition_model
        from insightface.app import FaceAnalysis
         # Load face encoder
        
        app_face = FaceAnalysis(name='antelopev2', 
                                root=insightface_root_path, providers=['CUDAExecutionProvider', 'CPUExecutionProvider'])
        app_face.prepare(ctx_id=0, det_size=(640, 640))
        app_face_ = init_recognition_model('arcface', device='cuda')
        pipeline_mask = None
    else:
        app_face = None
        pipeline_mask = None
        app_face_ = None
    return app_face,pipeline_mask,app_face_


def get_float(str_in):
    list_str=str_in.split(",")
    float_box=[float(x) for x in list_str]
    return float_box


def adjust_indices(original_indices, deleted_indices):
    """根据删除的索引列表调整原索引"""
    # 处理删除列表为空的特殊情况
    if not deleted_indices:
        return original_indices.copy()  # 直接返回原索引，无需调整
    
    deleted_sorted = sorted(deleted_indices)
    adjusted = []
    for idx in original_indices:
        # 计算偏移量：有多少个删除索引 < 当前索引
        offset = sum(1 for d in deleted_sorted if d < idx)
        new_idx = idx - offset
        # 如果原索引未被删除，则保留
        if idx not in deleted_sorted:
            adjusted.append(new_idx)
        else:
            adjusted.append(-1)  # 标记为无效
    return adjusted


def get_insight_dict(app_face,app_face_,pipeline_mask,infer_mode,use_photov2,image_list,
                     character_list_,condition_image,width, height,model=None,image_proj_model=None):
    input_id_emb_s_dict = {}
    input_id_img_s_dict = {}
    input_id_emb_un_dict = {}

    for ind, img in enumerate(image_list): # 最大只有2个ID
        if infer_mode == "story" and use_photov2:
            from .utils.insightface_package import analyze_faces
            img_ = np.array(img)
            img_ = cv2.cvtColor(img_, cv2.COLOR_RGB2BGR)
            faces = analyze_faces(app_face, img_, )
            id_embed_list = torch.from_numpy((faces[0]['embedding']))
            crop_image = img_
            uncond_id_embeddings = None
        elif infer_mode == "kolor_face":
            device = (
                "cuda"
                if torch.cuda.is_available()
                else "mps" if torch.backends.mps.is_available() else "cpu"
            )
            face_info = app_face.get_faceinfo_one_img(img)
            face_bbox_square = face_bbox_to_square(face_info["bbox"])
            crop_image = img.crop(face_bbox_square)
            crop_image = crop_image.resize((336, 336))
            face_embeds = torch.from_numpy(np.array([face_info["embedding"]]))
            id_embed_list = face_embeds.to(device, dtype=torch.float16)
            uncond_id_embeddings = None
        elif infer_mode=="story_maker":
            crop_image = pipeline_mask(img, return_mask=True).convert('RGB')  # outputs a pillow mask
            # timestamp = datetime.datetime.now().strftime("%Y%m%d-%H%M%S")
            # crop_image.copy().save(os.path.join(folder_paths.get_output_directory(),f"{timestamp}_mask.png"))
            face_info = app_face.get(cv2.cvtColor(np.array(img), cv2.COLOR_RGB2BGR))
            id_embed_list = sorted(face_info,key=lambda x: (x['bbox'][2] - x['bbox'][0]) * (x['bbox'][3] - x['bbox'][1]))[-1]  # only use the maximum face
            
            uncond_id_embeddings = img
        elif infer_mode=="story_and_maker":
            if use_photov2:
                from .utils.insightface_package import analyze_faces
                img = np.array(img)
                img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
                faces = analyze_faces(app_face, img, )
                id_embed_list = torch.from_numpy((faces[0]['embedding']))
                crop_image = pipeline_mask(img, return_mask=True).convert(
                    'RGB')  # outputs a pillow mask
                face_info = app_face_.get(cv2.cvtColor(np.array(img), cv2.COLOR_RGB2BGR))
                uncond_id_embeddings = sorted(face_info,key=lambda x: (x['bbox'][2] - x['bbox'][0]) * (x['bbox'][3] - x['bbox'][1]))[-1]  # only use the maximum face
                #photomake_mode = "v2"
                # make+v2模式下，emb存v2的向量，corp 和 unemb 存make的向量
            else:  # V1不需要调用emb
                crop_image = pipeline_mask(img, return_mask=True).convert(
                    'RGB')  # outputs a pillow mask
                face_info = app_face.get(cv2.cvtColor(np.array(img), cv2.COLOR_RGB2BGR))
                id_embed_list = sorted(face_info,key=lambda x: (x['bbox'][2] - x['bbox'][0]) * (x['bbox'][3] - x['bbox'][1]))[-1]  # only use the maximum face
                uncond_id_embeddings = img

        elif infer_mode=="flux_pulid":
            id_image = resize_numpy_image_long(img, 1024)
            use_true_cfg = abs(1.0 - 1.0) > 1e-2
           
            id_embed_list, uncond_id_embeddings=model.pulid_model.get_id_embedding_( id_image,app_face_,app_face,pipeline_mask,cal_uncond=use_true_cfg)


            #id_embed_list, uncond_id_embeddings = pipe.pulid_model.get_id_embedding(id_image,cal_uncond=use_true_cfg)
            crop_image = img

        elif infer_mode=="infiniteyou":
    
            def _detect_face(app_face, id_image_cv2):
                face_info = app_face.get(id_image_cv2)
                if face_info:
                    return face_info
                else:
                    print("No face detected in the input ID image")
                    return []
            from .pipelines.pipeline_infu_flux import extract_arcface_bgr_embedding,resize_and_pad_image,draw_kps
             # Extract ID embeddings
            print('Preparing ID embeddings')
            id_image_cv2 = cv2.cvtColor(np.array(img), cv2.COLOR_RGB2BGR)
            face_info = _detect_face(app_face,id_image_cv2)
            if len(face_info) == 0:
                raise ValueError('No face detected in the input ID image')
            
            face_info = sorted(face_info, key=lambda x:(x['bbox'][2]-x['bbox'][0])*(x['bbox'][3]-x['bbox'][1]))[-1] # only use the maximum face
            landmark = face_info['kps']
            id_embed = extract_arcface_bgr_embedding(id_image_cv2, landmark, app_face_)
            id_embed = id_embed.clone().unsqueeze(0).float().cuda()
            id_embed = id_embed.reshape([1, -1, 512])
            id_embed = id_embed.to(device='cuda', dtype=torch.bfloat16)
            with torch.no_grad():
                id_embed = image_proj_model(id_embed)
                bs_embed, seq_len, _ = id_embed.shape
                id_embed = id_embed.repeat(1, 1, 1)
                id_embed = id_embed.view(bs_embed * 1, seq_len, -1)
                id_embed = id_embed.to(device='cuda', dtype=torch.bfloat16)
            
            # Load control image
            print('Preparing the control image')
            if isinstance(condition_image, torch.Tensor):
                e1, _, _, _ = condition_image.size()
                if e1 == 1:
                    cn_image_load = [nomarl_upscale(condition_image, width, height)]
                else:
                    img_list = list(torch.chunk(condition_image, chunks=e1))
                    cn_image_load = [nomarl_upscale(img, width, height) for img in img_list]
                # control_image = control_image.convert("RGB")
                # control_image = resize_and_pad_image(control_image, (width, height))
                control_image=[] #如果是单人多张控制图，可能导致失序，所以统一改成列表，并在采样的时候依次调用
                for img in cn_image_load:
                    face_info = _detect_face(app_face,cv2.cvtColor(np.array(img), cv2.COLOR_RGB2BGR)) 
                    if len(face_info) == 0:
                        print('No face detected in the control image,use empty image,无法识别到面部,加载全黑图片替代.') #卡通人物很难识别，所以避免反复加载， 直接用黑图
                        face_cn=Image.fromarray(np.zeros([height, width, 3]).astype(np.uint8))
                    else:
                        face_info = sorted(face_info, key=lambda x:(x['bbox'][2]-x['bbox'][0])*(x['bbox'][3]-x['bbox'][1]))[-1] # only use the maximum face
                        face_cn = draw_kps(img, face_info['kps'])
                    control_image.append(face_cn)
            else:
                out_img = np.zeros([height, width, 3])
                control_image = Image.fromarray(out_img.astype(np.uint8))
            id_embed_list=id_embed
            crop_image = control_image  # inf use crop to control img
            uncond_id_embeddings = None 
        else:
            id_embed_list = None
            uncond_id_embeddings = None
            crop_image = None
        input_id_img_s_dict[character_list_[ind]] = [crop_image]
        input_id_emb_s_dict[character_list_[ind]] = [id_embed_list]
        input_id_emb_un_dict[character_list_[ind]] = [uncond_id_embeddings]
    
    
    app_face = None
    app_face_= None
    pipeline_mask =None
    gc_cleanup()
    
    if isinstance(condition_image, torch.Tensor) and infer_mode=="story_maker":
        e1, _, _, _ = condition_image.size()
        if e1 == 1:
            cn_image_load = [nomarl_upscale(condition_image, width, height)]
        else:
            img_list = list(torch.chunk(condition_image, chunks=e1))
            cn_image_load = [nomarl_upscale(img, width, height) for img in img_list]
        input_id_cloth_dict = {}
        if len(cn_image_load)>2:
            cn_image_load_role=cn_image_load[0:2]
        else:
            cn_image_load_role=cn_image_load
        for ind, img in enumerate(cn_image_load_role):
            input_id_cloth_dict[character_list_[ind]] = [img]
        if len(cn_image_load)>2: #处理多张control img
            my_list=cn_image_load[2:]
            for ind,img in enumerate(my_list):
                input_id_cloth_dict[f"dual{ind}"] = [img]
    else:
        input_id_cloth_dict = {}
    return input_id_emb_s_dict,input_id_img_s_dict,input_id_emb_un_dict,input_id_cloth_dict


def load_model_tag(repo,device,select_method):
    if "flor" in select_method.lower():#"thwri/CogFlorence-2-Large-Freeze"
        #pip install flash_attn
        from transformers import AutoModelForCausalLM, AutoProcessor, AutoConfig
        model = AutoModelForCausalLM.from_pretrained(repo, trust_remote_code=True).to(
            device)
        processor = AutoProcessor.from_pretrained(repo, trust_remote_code=True)
    else:
        model = AutoModel.from_pretrained(repo, trust_remote_code=True)
        processor = AutoTokenizer.from_pretrained(repo, trust_remote_code=True)#tokenizer
    model.eval()
    return model,processor

class StoryLiteTag:
    def __init__(self, device,temperature,select_method,repo="pzc163/MiniCPMv2_6-prompt-generator",):
        self.device = device
        self.repo = repo
        self.select_method=select_method
        self.model, self.processor=load_model_tag(self.repo, self.device,self.select_method)
        self.temperature=temperature
    def run_tag(self,image):
        if "flor" in self.select_method.lower():
            inputs = self.processor(text="<MORE_DETAILED_CAPTION>" , images=image, return_tensors="pt").to(device)
            generated_ids = self.model.generate(
                input_ids=inputs["input_ids"],
                pixel_values=inputs["pixel_values"],
                max_new_tokens=1024,
                num_beams=3,
                do_sample=True
            )
            generated_text = self.processor.batch_decode(generated_ids, skip_special_tokens=False)[0]
            parsed_answer = self.processor.post_process_generation(generated_text, task="<MORE_DETAILED_CAPTION>" ,
                                                              image_size=(image.width, image.height))
            res=parsed_answer["<MORE_DETAILED_CAPTION>"]
        else:
            question = 'Provide a detailed description of the details and content contained in the image, and generate a short prompt that can be used for image generation tasks in Stable Diffusion,remind you only need respons prompt itself and no other information.'
            msgs = [{'role': 'user', 'content': [image, question]}]
            res = self.model.chat(
                image=None,
                msgs=msgs,
                tokenizer=self.processor,# tokenizer
                temperature=self.temperature
            )
            res=res.split(":",1)[1].strip('"')
        s=res.strip()
        res=re.sub(r'^\n+|\n+$', '', s)
        res.strip("'")
        logging.info(f"{res}")
        return res


def images_generator(img_list: list, ):
    # get img size
    sizes = {}
    for image_ in img_list:
        if isinstance(image_, Image.Image):
            count = sizes.get(image_.size, 0)
            sizes[image_.size] = count + 1
        elif isinstance(image_, np.ndarray):
            count = sizes.get(image_.shape[:2][::-1], 0)
            sizes[image_.shape[:2][::-1]] = count + 1
        else:
            raise "unsupport image list,must be pil or cv2!!!"
    size = max(sizes.items(), key=lambda x: x[1])[0]
    yield size[0], size[1]
    
    # any to tensor
    def load_image(img_in):
        if isinstance(img_in, Image.Image):
            img_in = img_in.convert("RGB")
            i = np.array(img_in, dtype=np.float32)
            i = torch.from_numpy(i).div_(255)
            if i.shape[0] != size[1] or i.shape[1] != size[0]:
                i = torch.from_numpy(i).movedim(-1, 0).unsqueeze(0)
                i = common_upscale(i, size[0], size[1], "lanczos", "center")
                i = i.squeeze(0).movedim(0, -1).numpy()
            return i
        elif isinstance(img_in, np.ndarray):
            i = cv2.cvtColor(img_in, cv2.COLOR_BGR2RGB).astype(np.float32)
            i = torch.from_numpy(i).div_(255)
            print(i.shape)
            return i
        else:
            raise "unsupport image list,must be pil,cv2 or tensor!!!"
    
    total_images = len(img_list)
    processed_images = 0
    pbar = ProgressBar(total_images)
    images = map(load_image, img_list)
    try:
        prev_image = next(images)
        while True:
            next_image = next(images)
            yield prev_image
            processed_images += 1
            pbar.update_absolute(processed_images, total_images)
            prev_image = next_image
    except StopIteration:
        pass
    if prev_image is not None:
        yield prev_image


def load_images_list(img_list: list, ):
    gen = images_generator(img_list)
    (width, height) = next(gen)
    images = torch.from_numpy(np.fromiter(gen, np.dtype((np.float32, (height, width, 3)))))
    if len(images) == 0:
        raise FileNotFoundError(f"No images could be loaded .")
    return images
    

def fitter_cf_model_type(model):
    VALUE_=model.model.model_type.value
    if VALUE_ == 8:
        cf_model_type = "FLUX"
    elif VALUE_ == 4:
        cf_model_type = "CASCADE"
    elif VALUE_ == 6:
        cf_model_type = "FLOW"
    elif VALUE_ == 5:
        cf_model_type = "EDM"
    elif VALUE_ == 1:
        cf_model_type = "EPS"
    elif VALUE_ == 2:
        cf_model_type = "V_PREDICTION"
    elif VALUE_ == 3:
        cf_model_type = "V_PREDICTION_EDM"
    elif VALUE_ == 7:
        cf_model_type = "PREDICTION_CONTINUOUS"
    else:
         raise "unsupport checkpoints"
    return cf_model_type

def pre_text2infer(role_text,scene_text,lora_trigger_words,use_lora,tag_list):
    '''
    
    Args:
        role_index_dict: {'[Taylor]': [0, 1], '[Lecun]': [2, 3]} 
        invert_role_index_dict{0: ['[Taylor]'], 1: ['[Taylor]'], 2: ['[Lecun]'], 3: ['[Lecun]']}
        ref_role_index_dict {'[Taylor]': [0, 1], '[Lecun]': [2, 3]} 
        ref_role_totals[0, 1, 2, 3]
        role_list ['[Taylor]', '[Lecun]']
        role_dict {'[Taylor]': ' a woman img, wearing a white T-shirt, blue loose hair.', '[Lecun]': ' a man img,wearing a suit,black hair.'} 
        nc_txt_list [' a panda'] 
        nc_indexs[4]
        positions_index_dual[]
        positions_index_char_1 0 2 
        positions_index_char_2 [] 
        prompts_dual[] #['[A] play whith [B] in the garden'] 
        index_char_1_list [0, 1] 
        index_char_2_list[2, 3]

    Returns:
      character_index_dict:{'[Taylor]': [0, 3], '[sam]': [1, 2]},if 1 role {'[Taylor]': [0, 1, 2]}
      
    '''
    add_trigger_words = " " + lora_trigger_words + " style "
    # pre role_text
    role_dict, role_list = character_to_dict(role_text, use_lora, add_trigger_words)
    
    id_len=len(role_dict)
    #pre scene_text
    scene_text_origin = [i.strip() for i in scene_text.splitlines()]
    scene_text_origin = [i for i in scene_text_origin if '[' in i]  # 删除空行
    prompt_sence = [i for i in scene_text_origin if not len(extract_content_from_brackets(i)) >= 2]  # 剔除双角色的场景词
    
    positions_index_dual = [index for index, prompt in enumerate(scene_text_origin) if
                      len(extract_content_from_brackets(prompt)) >= 2]   #获取双角色出现的位置
    prompts_dual = [prompt for prompt in scene_text_origin if len(extract_content_from_brackets(prompt)) >= 2] # 改成单句中双方括号方法，利于MS组句，[A]... [B]...[C]
    
    #print(id_len,role_list)
    if id_len == 2:
        index_char_1_list=[index for index, prompt in enumerate(scene_text_origin) if role_list[0] in prompt]
        positions_index_char_1 = index_char_1_list[0]  # 获取角色出现的索引列表，并获取首次出现的位置
        index_char_2_list=[index for index, prompt in enumerate(scene_text_origin) if role_list[1] in prompt]
        positions_index_char_2 = index_char_2_list[0]  # 获取角色出现的索引列表，并获取首次出现的位置
    
       # print(positions_index_char_1, positions_index_char_2) #0,2
    else:
        index_char_1_list=[index for index, prompt in enumerate(scene_text_origin) if role_list[0] in prompt]
        index_char_2_list=[]
        positions_index_char_1=[]
        positions_index_char_2=[]

    if index_char_1_list and index_char_2_list: #dual 情况下，该列表有误，须排除重复元素
        common_ = set(index_char_1_list) & set(index_char_2_list)
        if common_:
            index_char_1_list=[x for x in index_char_1_list if x not in common_]
            index_char_2_list = [x for x in index_char_2_list if x not in common_]


    clipped_prompts = prompt_sence[:]  # copy
   
    nc_indexs = []
    for ind, prompt in enumerate(clipped_prompts): #获取NC的index
        if "[NC]" in prompt:
            nc_indexs.append(ind)
            if ind < id_len:
                raise f"The first [role] row need be a id prompts, cannot use [NC]!"
    prompts = [
        i if "[NC]" not in i else i.replace("[NC]", "") for i in clipped_prompts]
    #去除#
    prompts = [prompt.rpartition("#")[0] if "#" in prompt else prompt for prompt in prompts]
    
    # character_dict:{'[Taylor]': ' a woman img, wearing a white T-shirt, blue loose hair.'},character_list:['[Taylor]'] 1role
    #character_dict:{'[Taylor]': ' a woman img, wearing a white T-shirt, blue loose hair.', '[Lecun]': ' a man img,wearing a suit,black hair.'},character_list:['[Taylor]', '[Lecun]'] 2 role
    #获取字典，实际用词等
    role_index_dict, invert_role_index_dict, replace_prompts, ref_role_index_dict, ref_role_totals= process_original_text(role_dict, prompts)
    
    if tag_list: #tag方法
        if len(tag_list) < len(replace_prompts):
            raise "The number of input condition images is less than the number of scene prompts!"
        replace_prompts = [prompt + " " + tag_list[i] for i, prompt in enumerate(replace_prompts)]
    
    if nc_indexs:
        for x in nc_indexs:  # 获取NC列表
            nc_txt_list = [item for i, item in enumerate(replace_prompts) if i == x]
        
        for x in nc_indexs:  # 去除NC列表
            replace_prompts = [item for i, item in enumerate(replace_prompts) if i != x]
    else:
        nc_txt_list = []
        
   
    replace_prompts=[i for i in replace_prompts]
    
    return 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


def convert_cf2diffuser(model):
    from diffusers.pipelines.stable_diffusion.convert_from_ckpt import convert_ldm_unet_checkpoint
    from diffusers import UNet2DConditionModel
    config_file = os.path.join(cur_path,"local_repo/unet/config.json")
    cf_state_dict = model.diffusion_model.state_dict()
    unet_state_dict = model.model_config.process_unet_state_dict_for_saving(cf_state_dict)
    unet_config = UNet2DConditionModel.load_config(config_file)
    Unet = UNet2DConditionModel.from_config(unet_config).to(device, torch.float16)
    cf_state_dict = convert_ldm_unet_checkpoint(unet_state_dict, Unet.config)
    Unet.load_state_dict(cf_state_dict, strict=False)
    del cf_state_dict,unet_state_dict
    gc.collect()
    torch.cuda.empty_cache()
    return Unet


def cf_clip(txt_list, clip, infer_mode,role_list,input_split=True):
   
    if infer_mode == "classic":
        input_split= False
    if len(role_list)==1:
        input_split= False
    if input_split:
        role_emb_dict={}
        for role,role_text_list in zip(role_list,txt_list):
            pos_cond_list=[]
            for j in role_text_list:
                tokens_p = clip.tokenize(j)
                output_p = clip.encode_from_tokens(tokens_p, return_dict=True)  # {"pooled_output":tensor}
                cond_p = output_p.pop("cond")
                if cond_p.shape[1] / 77 > 1 and infer_mode != "classic":
                # logging.warning("prompt'tokens length is abvoe 77,split it")
                    cond_p = torch.chunk(cond_p, cond_p.shape[1] // 77, dim=1)[0]
                positive = [cond_p, output_p]
                pos_cond_list.append(positive)
            role_emb_dict[role]=pos_cond_list
        return role_emb_dict
    
    pos_cond_list = []
    for i in txt_list[0]:
        tokens_p = clip.tokenize(i)
        output_p = clip.encode_from_tokens(tokens_p, return_dict=True)  # {"pooled_output":tensor}
        cond_p = output_p.pop("cond")
        if cond_p.shape[1] / 77 > 1 and infer_mode != "classic":
            # logging.warning("prompt'tokens length is abvoe 77,split it")
            cond_p = torch.chunk(cond_p, cond_p.shape[1] // 77, dim=1)[0]
        if infer_mode == "classic":
            positive = [[cond_p, output_p]]
        else:
            positive = [cond_p, output_p]
        # logging.info(f"sampler text is {i}")
        pos_cond_list.append(positive)
    return pos_cond_list

def get_eot_idx_cf(tokenizer, prompt):
    words = prompt.split()
    start = 1
    for w in words:
        start += len(tokenizer.encode(w)) - 2
    return start

def get_phrase_idx_cf(tokenizer, phrase, prompt, get_last_word=False, num=0):
    def is_equal_words(pr_words, ph_words):
        if len(pr_words) != len(ph_words):
            return False
        for pr_word, ph_word in zip(pr_words, ph_words):
            if "-"+ph_word not in pr_word and ph_word != re.sub(r'[.!?,:]$', '', pr_word):
                return False
        return True

    phrase_words = phrase.split()
    if len(phrase_words) == 0:
        return [0, 0], None
    if get_last_word:
        phrase_words = phrase_words[-1:]
    # prompt_words = re.findall(r'\b[\w\'-]+\b', prompt)
    prompt_words = prompt.split()
    start = 1
    end = 0
    res_words = phrase_words
    for i in range(len(prompt_words)):
        if is_equal_words(prompt_words[i:i+len(phrase_words)], phrase_words):
            if num != 0:
                # skip this one
                num -= 1
                continue
            end = start
            res_words = prompt_words[i:i+len(phrase_words)]
            res_words = [re.sub(r'[.!?,:]$', '', w) for w in res_words]
            prompt_words[i+len(phrase_words)-1] = res_words[-1]  # remove the last punctuation
            for j in range(i, i+len(phrase_words)):
                end += len(tokenizer.encode(prompt_words[j])) - 2
            break
        else:
            start += len(tokenizer.encode(prompt_words[i])) - 2

    if end == 0:
        return [0, 0], None

    return [start, end], res_words

def get_phrases_idx_cf(tokenizer, phrases, prompt):
    res = []
    phrase_cnt = {}
    for phrase in phrases:
        if phrase in phrase_cnt:
            cur_cnt = phrase_cnt[phrase]
            phrase_cnt[phrase] += 1
        else:
            cur_cnt = 0
            phrase_cnt[phrase] = 1
        res.append(get_phrase_idx_cf(tokenizer, phrase, prompt, num=cur_cnt)[0])
    return res

def replicate_data_by_indices(data_list, index_list1, index_list2):
    # 确定新列表的长度
    max_index = max(max(index_list1), max(index_list2)) if (index_list1 and index_list2) else 0
    new_list = [None] * (max_index + 1)
    
    # 根据索引填充数据
    for idx in index_list1:
        new_list[idx] = data_list[0]
    for idx in index_list2:
        new_list[idx] = data_list[1]
    
    return new_list



def get_ms_phrase_emb(boxes, device, weight_dtype, drop_grounding_tokens, bsz, phrase_idxes,
                      num_samples, eot_idxes,phrases,clip,tokenizer):
    cross_attention_kwargs = None
    grounding_kwargs = None
    if boxes is not None:
        boxes = torch.tensor(boxes).to(device, weight_dtype) #torch.Size([1, 2, 4])
        if phrases is not None:
            drop_grounding_tokens = drop_grounding_tokens if drop_grounding_tokens is not None else [0] * bsz
            batch_boxes = boxes.view(bsz * boxes.shape[1], -1).to(device)
            phrase_input_ids=[]
            for phrase in phrases:
                phrase_input_id = tokenizer(phrase, max_length=tokenizer.model_max_length,
                                                     padding="max_length", truncation=True,
                                                     return_tensors="pt").input_ids
                int_list = phrase_input_id.tolist()[0]
                clean_input_ids_1=[(i,1.0) for i in int_list]
                clean_input_ids_2=[(i,1.0) for i in int_list if i==int_list[0] or i==49407 ]
                clean_input_ids_2=clean_input_ids_2[:77] if len(clean_input_ids_2) >=77 else clean_input_ids_2 + [(0,1.0)]*(77 - len(clean_input_ids_2))
                phrase_input_ids.append([clean_input_ids_1,clean_input_ids_2])

            phrase_embeds_list = []
            for i in phrase_input_ids:
                _, pooled_prompt_embed_= clip.cond_stage_model.clip_l.encode_token_weights(i)
                pooled_prompt_embed_= pooled_prompt_embed_.to(device,dtype=weight_dtype)
                phrase_embeds_list.append(pooled_prompt_embed_)
                #phrase_input_id = clip.tokenize(phrase, return_word_ids=False)["l"]
                #output_=clip.cond_stage_model.clip_l.encode_token_weights(phrase_input_id)[1]
               
            phrase_embeds=torch.cat(phrase_embeds_list,dim=0)
            #print(phrase_embeds.shape,phrase_embeds.is_cuda,batch_boxes.shape,batch_boxes.is_cuda)#torch.Size([2, 768]) torch.Size([2 4])
            grounding_kwargs = {"boxes": batch_boxes, "phrase_embeds": phrase_embeds,"drop_grounding_tokens": drop_grounding_tokens}
        else:
            grounding_kwargs = None
        boxes = torch.repeat_interleave(boxes, repeats=num_samples, dim=0)
        uncond_boxes = torch.zeros_like(boxes)
        boxes = torch.cat([uncond_boxes, boxes], dim=0)
        cross_attention_kwargs = {"boxes": boxes}
    
    if phrase_idxes is not None:
        # phrase_idxes = torch.tensor(phrase_idxes).to(device, torch.int)
        # eot_idxes = torch.tensor(eot_idxes).to(device, torch.int)
        phrase_idxes = torch.tensor(phrase_idxes).to(device, weight_dtype)
        eot_idxes = torch.tensor(eot_idxes).to(device, weight_dtype)
        
        phrase_idxes = torch.repeat_interleave(phrase_idxes, repeats=num_samples, dim=0)
        eot_idxes = torch.repeat_interleave(eot_idxes, repeats=num_samples, dim=0)
        uncond_phrase_idxes = torch.zeros_like(phrase_idxes)
        uncond_eot_idxes = torch.zeros_like(eot_idxes)
        phrase_idxes = torch.cat([uncond_phrase_idxes, phrase_idxes], dim=0)
        eot_idxes = torch.cat([uncond_eot_idxes, eot_idxes], dim=0)
        if cross_attention_kwargs is None:
            cross_attention_kwargs = {"phrase_idxes": phrase_idxes, "eot_idxes": eot_idxes}
        else:
            cross_attention_kwargs["phrase_idxes"] = phrase_idxes
            cross_attention_kwargs["eot_idxes"] = eot_idxes
    
    return cross_attention_kwargs,grounding_kwargs
def Infer_MSdiffusion(model, photomake_or_ipadapter_path,image_embeds,  prompt_embeds_, negative_prompt_embeds_,grounding_kwargs,cross_attention_kwargs, bsz,mask_threshold,height,width,steps,seed,cfg,pooled_prompt_embeds,negative_pooled_prompt_embeds,):
    image_proj_model = Resampler(
        dim=1280,
        depth=4,
        dim_head=64,
        heads=20,
        num_queries=16,
        embedding_dim=1664,
        output_dim=model.unet.config.cross_attention_dim,
        ff_mult=4,
        latent_init_mode="grounding",
        phrase_embeddings_dim=768,
    ).to(device, dtype=torch.float16)
    PIPE = MSAdapterWarpper(model, image_proj_model, ckpt_path=photomake_or_ipadapter_path, device=device, num_tokens=16)
    samples=PIPE.generate(model, image_embeds,  prompt_embeds_, negative_prompt_embeds_,grounding_kwargs,cross_attention_kwargs, bsz,height,width,steps,seed,cfg,pooled_prompt_embeds,negative_pooled_prompt_embeds,scale=1.0,
                 num_samples=bsz,  weight_dtype=torch.float16, subject_scales=None, mask_threshold=mask_threshold, start_step=5)
    return samples

def convert_clip2diffuser(clip):
   
    from transformers import (
            CLIPTextModel,CLIPTextModelWithProjection,
            CLIPTextConfig,
                                            )
    from contextlib import nullcontext

    try:
        from accelerate import init_empty_weights
        from accelerate.utils import set_module_tensor_to_device
        is_accelerate_available = True
    except:
        pass
    
    import comfy.model_management as mm
   
    offload_device=mm.unet_offload_device()
   
    text_encoder_config=CLIPTextConfig.from_pretrained(os.path.join(cur_path, 'local_repo/text_encoder'), local_files_only=True)
    ctx = init_empty_weights if is_accelerate_available else nullcontext
    with ctx():
        text_encoder = CLIPTextModel(text_encoder_config)
    text_encoder_sd = clip.get_sd()
    text_encoder_2_sd = clip.get_sd()

    from diffusers.pipelines.stable_diffusion.convert_from_ckpt import convert_open_clip_checkpoint,convert_ldm_clip_checkpoint
    text_encoder = convert_ldm_clip_checkpoint(text_encoder_sd, text_encoder=text_encoder)
    text_encoder=None
    text_encoder_2 = convert_open_clip_checkpoint(text_encoder_2_sd,config_name=os.path.join(cur_path, 'local_repo/text_encoder_2'))    
   
    return text_encoder,text_encoder_2
    


def convert_cfvae2diffuser(VAE,use_flux=False):
    from diffusers import AutoencoderKL
    from diffusers.pipelines.stable_diffusion.convert_from_ckpt import convert_ldm_vae_checkpoint
    if use_flux:
        vae_config = os.path.join(cur_path, "config/FLUX.1-dev/vae/config.json")
    else:
        vae_config = os.path.join(cur_path, "local_repo/vae/config.json")
    vae_state_dict=VAE.get_sd()
    ae_config = AutoencoderKL.load_config(vae_config)
    
    if not use_flux:
        AE = AutoencoderKL.from_config(ae_config).to(device, torch.float16)
        cf_state_dict = convert_ldm_vae_checkpoint(vae_state_dict, AE.config)
        AE.load_state_dict(cf_state_dict, strict=False)
        del cf_state_dict,vae_state_dict
    else:
        AE = AutoencoderKL.from_config(ae_config).to(device, torch.bfloat16)
        AE.load_state_dict(vae_state_dict, strict=False)
        del vae_state_dict
    torch.cuda.empty_cache()
    return AE


def Loader_storydiffusion(cf_model,photomake_ckpt_path,VAE,UNET=None):
    sdxl_repo = os.path.join(cur_path, "local_repo")
    vae_config=OmegaConf.load(os.path.join(cur_path,"local_repo/vae/config.json"))
    
    from .utils.pipeline_wrapper import StableDiffusionXLPipelineWrapper
    from .utils.pipeline_img import PhotoMakerStableDiffusionXLPipeline as PhotoMakerStableDiffusionXLPipelineWrapper
    from .utils.pipeline_v2_warrper import  PhotoMakerStableDiffusionXLPipeline as PhotoMakerStableDiffusionXLPipelineV2
    if cf_model is not None:
        Unet=convert_cf2diffuser(cf_model.model)
        if photomake_ckpt_path is not None: # guess img2img
            if VAE is None:
                raise "need link a SDXL vae when use img2img!"
            vae = convert_cfvae2diffuser(VAE)
            if "v1" in photomake_ckpt_path:
                model = PhotoMakerStableDiffusionXLPipelineWrapper.from_pretrained(sdxl_repo, unet=Unet, vae=vae,text_encoder=None,text_encoder_2=None,
                                                                     torch_dtype=torch.float16, use_safetensors=True)
                model.load_photomaker_adapter(
                photomake_ckpt_path,
                subfolder="",
                weight_name="photomaker-v1.bin",
                trigger_word="img"  # define the trigger word
            )
            
            else:
                model = PhotoMakerStableDiffusionXLPipelineV2.from_pretrained(sdxl_repo, unet=Unet, vae=vae,text_encoder=None,text_encoder_2=None,
                                                                     torch_dtype=torch.float16, use_safetensors=True)
                model.load_photomaker_adapter(
                photomake_ckpt_path,
                subfolder="",
                weight_name="photomaker-v2.bin",
                trigger_word="img",
                pm_version='v2',
            )
        else:
            model = StableDiffusionXLPipelineWrapper.from_pretrained(sdxl_repo, unet=Unet, vae_config=vae_config,
                                                                     torch_dtype=torch.float16, use_safetensors=True)
    elif UNET is not None:  
         model = StableDiffusionXLPipelineWrapper.from_pretrained(sdxl_repo, unet=UNET.unet, vae_config=vae_config,
                                                                     torch_dtype=torch.float16, use_safetensors=True)     
    else:
        raise "need  link a comfyUI checkpoint model!"

    model.to(device)
    return model




def Loader_story_maker(cf_model,ipadapter_ckpt_path,VAE,low_vram,lora_scale,controlnet=None,UNET=None):
    sdxl_repo = os.path.join(cur_path, "local_repo")
    #vae_config = OmegaConf.load(vae_config)
    if cf_model is not None:
        Unet = convert_cf2diffuser(cf_model.model)
        AE=convert_cfvae2diffuser(VAE)
        from .StoryMaker.pipeline_sdxl_storymaker_wrapper import StableDiffusionXLStoryMakerPipeline as StableDiffusionXLStoryMakerPipeline_wapper
        model = StableDiffusionXLStoryMakerPipeline_wapper.from_pretrained(sdxl_repo, vae=AE,unet=Unet,text_encoder=None,text_encoder_2=None, 
                                                                 torch_dtype=torch.float16, use_safetensors=True)
    elif UNET is not None:
        AE=convert_cfvae2diffuser(VAE)
        from .StoryMaker.pipeline_sdxl_storymaker_wrapper import StableDiffusionXLStoryMakerPipeline as StableDiffusionXLStoryMakerPipeline_wapper
        model = StableDiffusionXLStoryMakerPipeline_wapper.from_pretrained(sdxl_repo, vae=AE,unet=UNET.unet,text_encoder=None,text_encoder_2=None,
                                                                               torch_dtype=torch.float16, use_safetensors=True)
    else:
        raise "need  link a comfyUI checkpoint model!"
       
    
    if device != "mps":
        if not low_vram:
            model.cuda()
    
    model.load_storymaker_adapter(ipadapter_ckpt_path, scale=0.8, lora_scale=lora_scale, controlnet=controlnet)
    #model.scheduler = UniPCMultistepScheduler.from_config(model.scheduler.config)
    return model


def Loader_Flux_Pulid(model,cf_model, ipadapter_ckpt_path,quantized_mode,aggressive_offload,offload,if_repo,clip_vision_path):
    
    logging.info("start flux-pulid processing...")
    from .PuLID.app_flux import FluxGenerator
    
    pipe = FluxGenerator(model, "cuda:0", offload,aggressive_offload, pretrained_model=ipadapter_ckpt_path,
                            quantized_mode=quantized_mode, if_repo=if_repo,clip_vision_path=clip_vision_path)

    return pipe

def Loader_InfiniteYou(extra_info,VAE,quantize_mode):
    logging.info("start InfiniteYou mode processing...")    
    from .pipelines.pipeline_infu_flux import InfUFluxPipeline
    from diffusers import BitsAndBytesConfig as DiffusersBitsAndBytesConfig,FluxTransformer2DModel
    from transformers import BitsAndBytesConfig as TransformersBitsAndBytesConfig,T5EncoderModel
    from .pipelines.resampler import Resampler
    
    
    # quantize T5 
    # quant_config = TransformersBitsAndBytesConfig(
    #         load_in_4bit=True,
    #         bnb_4bit_quant_type="nf4",
    #     )

    # text_encoder_2_4bit = T5EncoderModel.from_pretrained(
    #         repo_id,
    #         subfolder="text_encoder_2",
    #         quantization_config=quant_config,
    #         torch_dtype=torch.bfloat16,
    #     ) # init nf4 
    repo_list=[i for i  in [extra_info.get("repo1_path") , extra_info.get("repo2_path")] if i is not None]
    assert repo_list is not [] ,"repo_list is None"
    find_infusenet=[i for i in repo_list if "sim_stage1" in i or "aes_stage2" in i]

    assert len(find_infusenet)>=1 ,"no infusenet"

    repo_id=os.path.join(cur_path, "config/FLUX.1-dev")
    use_svdq=False
    if extra_info.get("svdq_repo") is not None:
        print("use svdq quantization")   
        from nunchaku import NunchakuFluxTransformer2dModel
        use_svdq=True
        transformer = NunchakuFluxTransformer2dModel.from_pretrained(extra_info.get("svdq_repo"),offload=True)
        try:
            transformer.set_attention_impl("nunchaku-fp16")
        except:
            pass
        vae=convert_cfvae2diffuser(VAE,use_flux=True)
    elif extra_info.get("gguf_path") is not None:
        print("use gguf quantization")   
        from diffusers import  GGUFQuantizationConfig
        transformer = FluxTransformer2DModel.from_single_file(
            extra_info.get("gguf_path"),
            config=os.path.join(repo_id, "transformer"),
            quantization_config=GGUFQuantizationConfig(compute_dtype=torch.bfloat16),
            torch_dtype=torch.bfloat16,
        )
        vae=convert_cfvae2diffuser(VAE,use_flux=True)
    elif extra_info.get("unet_path") is not None:

        if quantize_mode=="fp8":
            transformer = FluxTransformer2DModel.from_single_file(extra_info.get("unet_path"), config=os.path.join(repo_id,"transformer/config.json"),
                                                                            torch_dtype=torch.bfloat16)
        else:
            from safetensors.torch import load_file
            t_state_dict=load_file(extra_info.get("unet_path"))
            unet_config = FluxTransformer2DModel.load_config(os.path.join(repo_id,"transformer/config.json"))
            transformer = FluxTransformer2DModel.from_config(unet_config).to(torch.bfloat16)
            transformer.load_state_dict(t_state_dict, strict=False)
            del t_state_dict
            gc_cleanup()
    else:
        vae = None
        find_flux=[i for i in repo_list if "flux" in i ]
        if not find_flux:
            raise "need fill flux repo  in EasyFunction_Lite repo1 or repo2!"
        print("use nf4 quantization")   
        if quantize_mode=="fp8":
            transformer = FluxTransformer2DModel.from_pretrained(
                find_flux[0],
                subfolder="transformer",
                quantization_config=DiffusersBitsAndBytesConfig(load_in_8bit=True,),
                torch_dtype=torch.bfloat16,
            )
                   
        elif quantize_mode=="nf4": #nf4
            transformer = FluxTransformer2DModel.from_pretrained(
                find_flux[0],
                subfolder="transformer",
                quantization_config=DiffusersBitsAndBytesConfig(
                    load_in_4bit=True, bnb_4bit_quant_type="nf4", bnb_4bit_compute_dtype=torch.bfloat16
                ),
                torch_dtype=torch.bfloat16,)
        else:
            transformer = FluxTransformer2DModel.from_pretrained(
                find_flux[0],
                subfolder="transformer",
                 torch_dtype=torch.bfloat16,
                )

        
    infusenet_path = os.path.join(find_infusenet[0], 'InfuseNetModel')

     
    pipe = InfUFluxPipeline(
            repo_id,infusenet_path,
            text_encoder_2=None,
            text_encoder=None,
            transformer=transformer,
            vae=vae,
            use_svdq=use_svdq,
        )
    # Load image proj model
    num_tokens = 8 # image_proj_num_tokens
    image_emb_dim = 512
    image_proj_model = Resampler(
        dim=1280,
        depth=4,
        dim_head=64,
        heads=20,
        num_queries=num_tokens,
        embedding_dim=image_emb_dim,
        output_dim=4096,
        ff_mult=4,
    )
    image_proj_model_path = os.path.join(find_infusenet[0], 'image_proj_model.bin')
    ipm_state_dict = torch.load(image_proj_model_path, map_location="cpu",weights_only=False)
    image_proj_model.load_state_dict(ipm_state_dict['image_proj'])
    del ipm_state_dict
    torch.cuda.empty_cache()
    image_proj_model.to('cuda', torch.bfloat16)
    image_proj_model.eval()

    pipe.pipe.enable_model_cpu_offload()
    return pipe,image_proj_model

def load_pipeline_consistory(cf_model,VAE):
    from .consistory.consistory_unet_sdxl import ConsistorySDXLUNet2DConditionModel
    from .consistory.consistory_pipeline_wapper import ConsistoryExtendAttnSDXLPipeline as ConsistoryExtendAttnSDXLPipeline_wapper
    float_type = torch.float16
    
    device = torch.device(f'cuda:0') if torch.cuda.is_available() else torch.device('cpu')

    if cf_model is not None:
        config_file = os.path.join(cur_path,"local_repo/unet/config.json")
        sdxl_repo = os.path.join(cur_path, "local_repo")

        from diffusers.pipelines.stable_diffusion.convert_from_ckpt import convert_ldm_unet_checkpoint
        cf_state_dict = cf_model.model.diffusion_model.state_dict()
        unet_state_dict = cf_model.model.model_config.process_unet_state_dict_for_saving(cf_state_dict)

        unet_config = ConsistorySDXLUNet2DConditionModel.load_config(config_file)
        Unet = ConsistorySDXLUNet2DConditionModel.from_config(unet_config).to(torch.float16)
        cf_state_dict = convert_ldm_unet_checkpoint(unet_state_dict, Unet.config)
        Unet.load_state_dict(cf_state_dict, strict=False)
        del cf_state_dict,unet_state_dict
        gc_cleanup()
        
        AE=convert_cfvae2diffuser(VAE)
       
        story_pipeline=ConsistoryExtendAttnSDXLPipeline_wapper.from_pretrained(sdxl_repo,unet=Unet,vae=AE,text_encoder=None,text_encoder_2=None, torch_dtype=float_type, )
        story_pipeline.to(device)
    else:
        raise "need link a sdxl checkpoints"
    
    story_pipeline.enable_freeu(s1=0.6, s2=0.4, b1=1.1, b2=1.2)
    return story_pipeline

def Loader_KOLOR(repo_id,clip_vision_path,face_ckpt,):

    from .kolors.models.modeling_chatglm import ChatGLMModel
    from .kolors.models.tokenization_chatglm import ChatGLMTokenizer
    from transformers import CLIPVisionModelWithProjection,CLIPImageProcessor
    from diffusers import  EulerDiscreteScheduler, UNet2DConditionModel, AutoencoderKL

    logging.info("loader kolor processing...")

    # text_encoder = ChatGLMModel.from_pretrained(
    #     f'{repo_id}/text_encoder', torch_dtype=torch.float16).half()
    vae = AutoencoderKL.from_pretrained(f"{repo_id}/vae", revision=None).half()
    #tokenizer = ChatGLMTokenizer.from_pretrained(f'{repo_id}/text_encoder')
    scheduler = EulerDiscreteScheduler.from_pretrained(f"{repo_id}/scheduler")

   
    from .kolors.pipelines.pipeline_stable_diffusion_xl_chatglm_256_ipadapter_FaceID import \
        StableDiffusionXLPipeline as StableDiffusionXLPipelineFaceID
    unet = UNet2DConditionModel.from_pretrained(f'{repo_id}/unet', revision=None).half()
    
    if clip_vision_path:

        clip_image_encoder = clip_vision_path
        clip_image_processor = CLIPImageProcessor(size=224, crop_size=224)
        use_singel_clip = True
    else:
        clip_image_encoder = CLIPVisionModelWithProjection.from_pretrained(
            f'{repo_id}/clip-vit-large-patch14-336', ignore_mismatched_sizes=True)
        clip_image_encoder.to("cuda")
        clip_image_processor = CLIPImageProcessor(size=336, crop_size=336)
        use_singel_clip = False
    
    pipe = StableDiffusionXLPipelineFaceID(
        vae=vae,
        text_encoder=None,
        tokenizer=None,
        unet=unet,
        scheduler=scheduler,
        face_clip_encoder=clip_image_encoder,
        face_clip_processor=clip_image_processor,
        force_zeros_for_empty_prompt=False,
        use_single_clip=use_singel_clip,
    )
    pipe = pipe.to("cuda")
    pipe.load_ip_adapter_faceid_plus(face_ckpt, device="cuda")
    pipe.set_face_fidelity_scale(0.8)
    return pipe

def GLM_clip(cur_path,chatglm3_path): #https://github.com/kijai/ComfyUI-KwaiKolorsWrapper/blob/main/nodes.py @kijai
    from .kolors.models.modeling_chatglm import ChatGLMModel, ChatGLMConfig
    from contextlib import nullcontext
    from comfy.utils import ProgressBar, load_torch_file
    try:
        from accelerate import init_empty_weights
        from accelerate.utils import set_module_tensor_to_device
        is_accelerate_available = True
    except:
        pass
    
    import comfy.model_management as mm
    import json
    offload_device=mm.unet_offload_device()
    text_encoder_config = os.path.join(cur_path, 'config/glm/text_encoder_config.json')
    with open(text_encoder_config, 'r') as file:
        config = json.load(file)

    text_encoder_config = ChatGLMConfig(**config)
    with (init_empty_weights() if is_accelerate_available else nullcontext()):
        text_encoder = ChatGLMModel(text_encoder_config)
        if '4bit' in chatglm3_path:
            text_encoder.quantize(4)
        elif '8bit' in chatglm3_path:
            text_encoder.quantize(8)

    text_encoder_sd = load_torch_file(chatglm3_path)
    if is_accelerate_available:
        for key in text_encoder_sd:
            set_module_tensor_to_device(text_encoder, key, device=offload_device, value=text_encoder_sd[key])
    else:
        text_encoder.load_state_dict()
    

    return text_encoder

def glm_single_encode(chatglm3_model, prompt_list,role_list, negative_prompt_, num_images_per_prompt,nc_mode=False): #https://github.com/kijai/ComfyUI-KwaiKolorsWrapper/blob/main/nodes.py @kijai
    import comfy.model_management as mm
    device = mm.get_torch_device()
    offload_device = mm.unet_offload_device()
    mm.unload_all_models()
    mm.soft_empty_cache()
        # Function to randomly select an option from the brackets
    # def choose_random_option(match):
    #     options = match.group(1).split('|')
    #     return random.choice(options)

    # # Randomly choose between options in brackets for prompt and negative_prompt
    # prompt = re.sub(r'\{([^{}]*)\}', choose_random_option, prompt)
    # negative_prompt = re.sub(r'\{([^{}]*)\}', choose_random_option, negative_prompt)

    # if "|" in prompt:
    #     prompt = prompt.split("|")
    #     negative_prompt = [negative_prompt] * len(prompt)  # Replicate negative_prompt to match length of prompt list

    do_classifier_free_guidance = True
    #print(prompt)
    if nc_mode:
        only_role_emb,only_role_emb_ne = [],[]
        for prompt,negative_prompt in zip(prompt_list,[negative_prompt_]*len(prompt_list)):
                
            if prompt is not None and isinstance(prompt, str):
                batch_size = 1
            elif prompt is not None and isinstance(prompt, list):
                batch_size = len(prompt)

            # Define tokenizers and text encoders
            tokenizer = chatglm3_model['tokenizer']
            text_encoder = chatglm3_model['text_encoder']

            text_encoder.to(device)

            text_inputs = tokenizer(
                prompt,
                padding="max_length",
                max_length=256,
                truncation=True,
                return_tensors="pt",
            ).to(device)

            output = text_encoder(
                    input_ids=text_inputs['input_ids'] ,
                    attention_mask=text_inputs['attention_mask'],
                    position_ids=text_inputs['position_ids'],
                    output_hidden_states=True)
            
            prompt_embeds = output.hidden_states[-2].permute(1, 0, 2).clone() # [batch_size, 77, 4096]
            text_proj = output.hidden_states[-1][-1, :, :].clone() # [batch_size, 4096]
            bs_embed, seq_len, _ = prompt_embeds.shape
            prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1)
            prompt_embeds = prompt_embeds.view(bs_embed * num_images_per_prompt, seq_len, -1)


            if do_classifier_free_guidance:
                uncond_tokens = []
                if negative_prompt is None:
                    uncond_tokens = [""] * batch_size
                elif prompt is not None and type(prompt) is not type(negative_prompt):
                    raise TypeError(
                        f"`negative_prompt` should be the same type to `prompt`, but got {type(negative_prompt)} !="
                        f" {type(prompt)}."
                    )
                elif isinstance(negative_prompt, str):
                    uncond_tokens = [negative_prompt]
                elif batch_size != len(negative_prompt):
                    raise ValueError(
                        f"`negative_prompt`: {negative_prompt} has batch size {len(negative_prompt)}, but `prompt`:"
                        f" {prompt} has batch size {batch_size}. Please make sure that passed `negative_prompt` matches"
                        " the batch size of `prompt`."
                    )
                else:
                    uncond_tokens = negative_prompt
            

                max_length = prompt_embeds.shape[1]
                uncond_input = tokenizer(
                    uncond_tokens,
                    padding="max_length",
                    max_length=max_length,
                    truncation=True,
                    return_tensors="pt",
                ).to(device)
                output = text_encoder(
                        input_ids=uncond_input['input_ids'] ,
                        attention_mask=uncond_input['attention_mask'],
                        position_ids=uncond_input['position_ids'],
                        output_hidden_states=True)
                negative_prompt_embeds = output.hidden_states[-2].permute(1, 0, 2).clone() # [batch_size, 77, 4096]
                negative_text_proj = output.hidden_states[-1][-1, :, :].clone() # [batch_size, 4096]

                if do_classifier_free_guidance:
                    # duplicate unconditional embeddings for each generation per prompt, using mps friendly method
                    seq_len = negative_prompt_embeds.shape[1]

                    negative_prompt_embeds = negative_prompt_embeds.to(dtype=text_encoder.dtype, device=device)

                    negative_prompt_embeds = negative_prompt_embeds.repeat(1, num_images_per_prompt, 1)
                    negative_prompt_embeds = negative_prompt_embeds.view(
                        batch_size * num_images_per_prompt, seq_len, -1
                    )

            bs_embed = text_proj.shape[0]
            text_proj = text_proj.repeat(1, num_images_per_prompt).view(
                bs_embed * num_images_per_prompt, -1
            )
            negative_text_proj = negative_text_proj.repeat(1, num_images_per_prompt).view(
                bs_embed * num_images_per_prompt, -1
            )
            math_emb=[prompt_embeds,text_proj]
            only_role_emb.append(math_emb)
            math_emb_n=[negative_prompt_embeds,negative_text_proj]
            only_role_emb_ne.append(math_emb_n)
        
        text_encoder.to(offload_device)
        mm.soft_empty_cache()
        gc.collect()
        return only_role_emb,only_role_emb_ne
    else:
        role_emb_dict={}
        for key,prompts in zip(role_list,prompt_list):
            only_role_emb,only_role_emb_ne = [],[]
            for prompt,negative_prompt in zip(prompts,[negative_prompt_]*len(prompts)):
                
                if prompt is not None and isinstance(prompt, str):
                    batch_size = 1
                elif prompt is not None and isinstance(prompt, list):
                    batch_size = len(prompt)

                # Define tokenizers and text encoders
                tokenizer = chatglm3_model['tokenizer']
                text_encoder = chatglm3_model['text_encoder']

                text_encoder.to(device)

                text_inputs = tokenizer(
                    prompt,
                    padding="max_length",
                    max_length=256,
                    truncation=True,
                    return_tensors="pt",
                ).to(device)

                output = text_encoder(
                        input_ids=text_inputs['input_ids'] ,
                        attention_mask=text_inputs['attention_mask'],
                        position_ids=text_inputs['position_ids'],
                        output_hidden_states=True)
                
                prompt_embeds = output.hidden_states[-2].permute(1, 0, 2).clone() # [batch_size, 77, 4096]
                text_proj = output.hidden_states[-1][-1, :, :].clone() # [batch_size, 4096]
                bs_embed, seq_len, _ = prompt_embeds.shape
                prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1)
                prompt_embeds = prompt_embeds.view(bs_embed * num_images_per_prompt, seq_len, -1)


                if do_classifier_free_guidance:
                    uncond_tokens = []
                    if negative_prompt is None:
                        uncond_tokens = [""] * batch_size
                    elif prompt is not None and type(prompt) is not type(negative_prompt):
                        raise TypeError(
                            f"`negative_prompt` should be the same type to `prompt`, but got {type(negative_prompt)} !="
                            f" {type(prompt)}."
                        )
                    elif isinstance(negative_prompt, str):
                        uncond_tokens = [negative_prompt]
                    elif batch_size != len(negative_prompt):
                        raise ValueError(
                            f"`negative_prompt`: {negative_prompt} has batch size {len(negative_prompt)}, but `prompt`:"
                            f" {prompt} has batch size {batch_size}. Please make sure that passed `negative_prompt` matches"
                            " the batch size of `prompt`."
                        )
                    else:
                        uncond_tokens = negative_prompt
                

                    max_length = prompt_embeds.shape[1]
                    uncond_input = tokenizer(
                        uncond_tokens,
                        padding="max_length",
                        max_length=max_length,
                        truncation=True,
                        return_tensors="pt",
                    ).to(device)
                    output = text_encoder(
                            input_ids=uncond_input['input_ids'] ,
                            attention_mask=uncond_input['attention_mask'],
                            position_ids=uncond_input['position_ids'],
                            output_hidden_states=True)
                    negative_prompt_embeds = output.hidden_states[-2].permute(1, 0, 2).clone() # [batch_size, 77, 4096]
                    negative_text_proj = output.hidden_states[-1][-1, :, :].clone() # [batch_size, 4096]

                    if do_classifier_free_guidance:
                        # duplicate unconditional embeddings for each generation per prompt, using mps friendly method
                        seq_len = negative_prompt_embeds.shape[1]

                        negative_prompt_embeds = negative_prompt_embeds.to(dtype=text_encoder.dtype, device=device)

                        negative_prompt_embeds = negative_prompt_embeds.repeat(1, num_images_per_prompt, 1)
                        negative_prompt_embeds = negative_prompt_embeds.view(
                            batch_size * num_images_per_prompt, seq_len, -1
                        )

                bs_embed = text_proj.shape[0]
                text_proj = text_proj.repeat(1, num_images_per_prompt).view(
                    bs_embed * num_images_per_prompt, -1
                )
                negative_text_proj = negative_text_proj.repeat(1, num_images_per_prompt).view(
                    bs_embed * num_images_per_prompt, -1
                )
                math_emb=[prompt_embeds,text_proj]
                only_role_emb.append(math_emb)
                math_emb_n=[negative_prompt_embeds,negative_text_proj]
                only_role_emb_ne.append(math_emb_n)
            role_emb_dict[key]=only_role_emb
    
        text_encoder.to(offload_device)
        mm.soft_empty_cache()
        gc.collect()

        return role_emb_dict,only_role_emb_ne


def encode_prompt_with_trigger_word(clip,model_,
    prompt: str,
    prompt_2 = None,
    num_id_images= 1,
    device = None,
    prompt_embeds = None,
    pooled_prompt_embeds= None,
    class_tokens_mask = None,
    nc_flag: bool = False,
    trigger_word="img",
    unet_dtype=torch.float16
):
    device = device 
    tokenizer_2=model_.tokenizer_2
    tokenizer=model_.tokenizer
    #text_encoders=clip.encode_from_tokens

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

    # Find the token id of the trigger word
    image_token_id = tokenizer_2.convert_tokens_to_ids(trigger_word)

    # Define tokenizers and text encoders
    tokenizers = [tokenizer, tokenizer_2] if tokenizer is not None else [tokenizer_2]
    # text_encoders = (
    #     [self.text_encoder, self.text_encoder_2] if self.text_encoder is not None else [self.text_encoder_2]
    # )

    if prompt_embeds is None:
        prompt_2 = prompt_2 or prompt
        prompt_embeds_list = []
        prompts = [prompt, prompt_2]
       
        for prompt, tokenizer,text_en in zip(prompts, tokenizers,[clip.cond_stage_model.clip_l.encode_token_weights,clip.cond_stage_model.clip_g.encode_token_weights]):
            input_ids = tokenizer.encode(prompt) # TODO: batch encode
            clean_index = 0
            clean_input_ids = []
            class_token_index = []
            # Find out the corresponding class word token based on the newly added trigger word token
            for i, token_id in enumerate(input_ids):
                if token_id == image_token_id:
                    class_token_index.append(clean_index - 1)
                else:
                    clean_input_ids.append(token_id)
                    clean_index += 1
            if nc_flag:
                return None, None, None
            if len(class_token_index) > 1:
                raise ValueError(
                    f"PhotoMaker currently does not support multiple trigger words in a single prompt.\
                        Trigger word: {trigger_word}, Prompt: {prompt}."
                )
            elif len(class_token_index) == 0 and not nc_flag:
                raise ValueError(
                    f"PhotoMaker currently does not support multiple trigger words in a single prompt.\
                        Trigger word: {trigger_word}, Prompt: {prompt}."
                )
            class_token_index = class_token_index[0]

            # Expand the class word token and corresponding mask
            class_token = clean_input_ids[class_token_index]
            clean_input_ids = clean_input_ids[:class_token_index] + [class_token] * num_id_images + \
                clean_input_ids[class_token_index+1:]

            # Truncation or padding
            max_len = tokenizer.model_max_length
            if len(clean_input_ids) > max_len:
                clean_input_ids = clean_input_ids[:max_len]
            else:
                clean_input_ids = clean_input_ids + [tokenizer.pad_token_id] * (
                    max_len - len(clean_input_ids)
                )

            class_tokens_mask = [True if class_token_index <= i < class_token_index+num_id_images else False \
                    for i in range(len(clean_input_ids))]
            #print(clean_input_ids)
            #print(max_len)
            
            clean_input_ids=clean_input_ids[:77] if len(clean_input_ids) > 77 else clean_input_ids + [0]*(77 - len(clean_input_ids))
                
            clean_input_ids_1=[(i,1.0) for i in clean_input_ids]
            clean_input_ids_2=[(i,1.0) for i in clean_input_ids if i==clean_input_ids[0] or i==49407 ]
            clean_input_ids_2=clean_input_ids_2[:77] if len(clean_input_ids_2) > 77 else clean_input_ids_2 + [(0,1.0)]*(77 - len(clean_input_ids_2))
            clean_input_ids_list=[clean_input_ids_1,clean_input_ids_2]
            #print(clean_input_ids)
            #clean_input_ids_ = torch.tensor(clean_input_ids, dtype=torch.long).unsqueeze(0)
            class_tokens_mask = torch.tensor(class_tokens_mask, dtype=torch.bool).unsqueeze(0)

           
            prompt_embeds, pooled_prompt_embeds = text_en(clean_input_ids_list) # print(prompt_embeds.shape,pooled_prompt_embeds.shape)torch.Size([1, 154, 768]) torch.Size([1, 768]) torch.Size([1, 154, 1280]) torch.Size([1, 1280])


            prompt_embeds_list.append((prompt_embeds,pooled_prompt_embeds))

        #prompt_embeds = torch.concat(prompt_embeds_list, dim=-1)
        #prompt_embeds=prompt_embeds_list[1]
    pooled_prompt_embeds=prompt_embeds_list[1][1]
    # cut_to = min(prompt_embeds_list[0][0].shape[1], prompt_embeds_list[1][0].shape[1])
    # print(cut_to)
    cut_to=77 
    prompt_embeds= torch.cat([prompt_embeds_list[0][0][:,:cut_to], prompt_embeds_list[1][0][:,:cut_to]], dim=-1)
    prompt_embeds = prompt_embeds.to(device=device,dtype=unet_dtype, )
    class_tokens_mask = class_tokens_mask.to(device=device) # TODO: ignoring two-prompt case

    return prompt_embeds, pooled_prompt_embeds, class_tokens_mask

def cf_clip_single(clip,prompt,
            prompt_2,
            device,
            num_images_per_prompt,
            do_classifier_free_guidance,
            negative_prompt,
            negative_prompt_2,
            prompt_embeds,
            negative_prompt_embeds,
            pooled_prompt_embeds,
            negative_pooled_prompt_embeds,):
    
    tokens_p = clip.tokenize(prompt)
    output_p = clip.encode_from_tokens(tokens_p, return_dict=True)  # {"pooled_output":tensor}
    cond_p = output_p.pop("cond")
    if cond_p.shape[1] / 77 > 1 :
        logging.warning("prompt'tokens length is abvoe 77,split it")
        cond_p = torch.chunk(cond_p, cond_p.shape[1] // 77, dim=1)[0]
    pool_p=output_p.pop("pooled_output")

    tokens_n = clip.tokenize(negative_prompt)
    output_n = clip.encode_from_tokens(tokens_n, return_dict=True)  # {"pooled_output":tensor}
    cond_n = output_n.pop("cond")
    if cond_n.shape[1] / 77 > 1 :
        logging.warning("prompt'tokens length is abvoe 77,split it")
        cond_n = torch.chunk(cond_n, cond_n.shape[1] // 77, dim=1)[0]
    pool_n = output_n.pop("pooled_output")

    return cond_p,cond_n,pool_p,pool_n # TODO: replace the pooled_prompt_embeds with text only prompt

    

def photomaker_clip(clip,model_,prompt_list,negative_prompt,input_id_images,trigger_word="img",num_images_per_prompt=1,nc_flag=False,prompt_2=None,prompt_embeds=None,pooled_prompt_embeds=None,class_tokens_mask=None,prompt_embeds_text_only=None,negative_pooled_prompt_embeds=None):
    id_encoder=model_.id_encoder
    id_image_processor=model_.id_image_processor

    tokenizer=model_.tokenizer_2
    batch_size=len(prompt_list)

    num_id_images = len(input_id_images) #双角色时，须确保输入的图片是列表，且跟prompt对应上
    prompt_arr = prompt_list
    negative_prompt_embeds_arr = []
    prompt_embeds_text_only_arr = []
    prompt_embeds_arr = []
    latents_arr = []
    add_time_ids_arr = []
    negative_pooled_prompt_embeds_arr = []
    pooled_prompt_embeds_text_only_arr = []
    pooled_prompt_embeds_arr = []
    for prompt in prompt_arr:
        (
            prompt_embeds,
            pooled_prompt_embeds,
            class_tokens_mask,
        ) = encode_prompt_with_trigger_word(clip,model_,
            prompt=prompt,
            prompt_2=prompt_2,
            device=device,
            num_id_images=num_id_images,
            prompt_embeds=prompt_embeds,
            pooled_prompt_embeds=pooled_prompt_embeds,
            class_tokens_mask=class_tokens_mask,
            nc_flag = nc_flag,
        )

        # 4. Encode input prompt without the trigger word for delayed conditioning
        # encode, remove trigger word token, then decode
        tokens_text_only = tokenizer.encode(prompt, add_special_tokens=False)
        trigger_word_token = tokenizer.convert_tokens_to_ids(trigger_word)
        if not nc_flag:
            tokens_text_only.remove(trigger_word_token)
        prompt_text_only = tokenizer.decode(tokens_text_only, add_special_tokens=False)
        #print(prompt_text_only)
        (
            prompt_embeds_text_only,
            negative_prompt_embeds,
            pooled_prompt_embeds_text_only, # TODO: replace the pooled_prompt_embeds with text only prompt
            negative_pooled_prompt_embeds,
        ) = cf_clip_single(clip,
            prompt=prompt_text_only,
            prompt_2=prompt_2,
            device=device,
            num_images_per_prompt=num_images_per_prompt,
            do_classifier_free_guidance=True,
            negative_prompt=negative_prompt,
            negative_prompt_2=None,
            prompt_embeds=prompt_embeds_text_only,
            negative_prompt_embeds=None,
            pooled_prompt_embeds=None,
            negative_pooled_prompt_embeds=negative_pooled_prompt_embeds,
        )

        # 5. Prepare the input ID images
        dtype = next(id_encoder.parameters()).dtype
        if not isinstance(input_id_images[0], torch.Tensor):
            id_pixel_values = id_image_processor(input_id_images, return_tensors="pt").pixel_values

        id_pixel_values = id_pixel_values.unsqueeze(0).to(device=device, dtype=dtype) # TODO: multiple prompts

        # duplicate text embeddings for each generation per prompt, using mps friendly method
        if not nc_flag:
            # 6. Get the update text embedding with the stacked ID embedding
            prompt_embeds = id_encoder(id_pixel_values, prompt_embeds, class_tokens_mask)

            bs_embed, seq_len, _ = prompt_embeds.shape
            prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1)
            prompt_embeds = prompt_embeds.view(bs_embed * num_images_per_prompt, seq_len, -1)
            pooled_prompt_embeds = pooled_prompt_embeds.repeat(1, num_images_per_prompt).view(
                bs_embed * num_images_per_prompt, -1
            )
            pooled_prompt_embeds_arr.append(pooled_prompt_embeds)
            pooled_prompt_embeds = None

        negative_prompt_embeds_arr.append(negative_prompt_embeds)
        negative_prompt_embeds = None
        negative_pooled_prompt_embeds_arr.append(negative_pooled_prompt_embeds)
        negative_pooled_prompt_embeds = None
        prompt_embeds_text_only_arr.append(prompt_embeds_text_only)
        prompt_embeds_text_only = None
        prompt_embeds_arr.append(prompt_embeds)
        prompt_embeds = None
        pooled_prompt_embeds_text_only_arr.append(pooled_prompt_embeds_text_only)
        pooled_prompt_embeds_text_only = None

    emb_dict={"negative_prompt_embeds_arr":negative_prompt_embeds_arr,
              "negative_prompt_embeds":negative_prompt_embeds,
              "negative_pooled_prompt_embeds_arr":negative_pooled_prompt_embeds_arr,
              "negative_pooled_prompt_embeds":negative_pooled_prompt_embeds,
              "prompt_embeds_text_only_arr":prompt_embeds_text_only_arr,
              "prompt_embeds_text_only":prompt_embeds_text_only,
              "prompt_embeds_arr":prompt_embeds_arr,
              "prompt_embeds":prompt_embeds,
              "pooled_prompt_embeds":pooled_prompt_embeds,
              "pooled_prompt_embeds_text_only_arr":pooled_prompt_embeds_text_only_arr,
              "pooled_prompt_embeds_text_only":pooled_prompt_embeds_text_only,
              "pooled_prompt_embeds_arr":pooled_prompt_embeds_arr,
                "batch_size":batch_size,
              }    
    return emb_dict


from typing import Any, Callable, Dict, List, Optional, Tuple, Union
def encode_prompt_with_trigger_word_v2(clip,model_,
    prompt: str,
    prompt_2 = None,
    num_id_images= 1,
    device = None,
    prompt_embeds = None,
    negative_prompt = None,
    negative_prompt_2 = None,
    negative_prompt_embeds= None,
    pooled_prompt_embeds= None,
    negative_pooled_prompt_embeds= None,
    num_images_per_prompt=1,
    class_tokens_mask = None,
    nc_flag: bool = False,
    trigger_word="img",
    unet_dtype=torch.float16,
    clip_skip=None,
    lora_scale=None,
    do_classifier_free_guidance=True,
    force_zeros_for_empty_prompts=False,
    num_tokens=2,

):
    
    device = device 
    tokenizer_2=model_.tokenizer_2
    tokenizer=model_.tokenizer
    #text_encoders=clip.encode_from_tokens

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

    # Find the token id of the trigger word
    image_token_id = tokenizer_2.convert_tokens_to_ids(trigger_word)

    # Define tokenizers and text encoders
    tokenizers = [tokenizer, tokenizer_2] if tokenizer is not None else [tokenizer_2]
    # text_encoders = (
    #     [self.text_encoder, self.text_encoder_2] if self.text_encoder is not None else [self.text_encoder_2]
    # )

    if prompt_embeds is None:
        prompt_2 = prompt_2 or prompt
        prompt_embeds_list = []
        prompts = [prompt, prompt_2]
       
        for prompt, tokenizer,text_en in zip(prompts, tokenizers,[clip.cond_stage_model.clip_l.encode_token_weights,clip.cond_stage_model.clip_g.encode_token_weights]):
            # if isinstance(self, TextualInversionLoaderMixin):
            #     prompt = self.maybe_convert_prompt(prompt, tokenizer)
            
            input_ids = tokenizer.encode(prompt) # TODO: batch encode
            clean_index = 0
            clean_input_ids = []
            class_token_index = []
            # Find out the corresponding class word token based on the newly added trigger word token
            for i, token_id in enumerate(input_ids):
                if token_id == image_token_id:
                    class_token_index.append(clean_index - 1)
                else:
                    clean_input_ids.append(token_id)
                    clean_index += 1
            if nc_flag:
                return None, None, None,None, None
            if len(class_token_index) > 1:
                raise ValueError(
                    f"PhotoMaker currently does not support multiple trigger words in a single prompt.\
                        Trigger word: {trigger_word}, Prompt: {prompt}."
                )
            elif len(class_token_index) == 0 and not nc_flag:
                raise ValueError(
                    f"PhotoMaker currently does not support multiple trigger words in a single prompt.\
                        Trigger word: {trigger_word}, Prompt: {prompt}."
                )
            class_token_index = class_token_index[0]

            # Expand the class word token and corresponding mask
            class_token = clean_input_ids[class_token_index]
            clean_input_ids = clean_input_ids[:class_token_index] + [class_token] * num_id_images * num_tokens + clean_input_ids[class_token_index + 1:]
            # Truncation or padding
            max_len = tokenizer.model_max_length
            if len(clean_input_ids) > max_len:
                clean_input_ids = clean_input_ids[:max_len]
            else:
                clean_input_ids = clean_input_ids + [tokenizer.pad_token_id] * (
                    max_len - len(clean_input_ids)
                )

            #class_tokens_mask = [True if class_token_index <= i < class_token_index+num_id_images else False for i in range(len(clean_input_ids))]
            class_tokens_mask = [True if class_token_index <= i < class_token_index + (num_id_images * num_tokens) else False for i in range(len(clean_input_ids))]
            #print(clean_input_ids)
            #print(max_len)
            
            clean_input_ids=clean_input_ids[:77] if len(clean_input_ids) > 77 else clean_input_ids + [0]*(77 - len(clean_input_ids))
                
            clean_input_ids_1=[(i,1.0) for i in clean_input_ids]
            clean_input_ids_2=[(i,1.0) for i in clean_input_ids if i==clean_input_ids[0] or i==49407 ]
            clean_input_ids_2=clean_input_ids_2[:77] if len(clean_input_ids_2) > 77 else clean_input_ids_2 + [(0,1.0)]*(77 - len(clean_input_ids_2))
            clean_input_ids_list=[clean_input_ids_1,clean_input_ids_2]
            #print(clean_input_ids)
            #clean_input_ids_ = torch.tensor(clean_input_ids, dtype=torch.long).unsqueeze(0)
            class_tokens_mask = torch.tensor(class_tokens_mask, dtype=torch.bool).unsqueeze(0)

           
            prompt_embeds, pooled_prompt_embeds = text_en(clean_input_ids_list) # print(prompt_embeds.shape,pooled_prompt_embeds.shape)torch.Size([1, 154, 768]) torch.Size([1, 768]) torch.Size([1, 154, 1280]) torch.Size([1, 1280])


            prompt_embeds_list.append((prompt_embeds,pooled_prompt_embeds))

        #prompt_embeds = torch.concat(prompt_embeds_list, dim=-1)
        #prompt_embeds=prompt_embeds_list[1]
    pooled_prompt_embeds=prompt_embeds_list[1][1]
    # cut_to = min(prompt_embeds_list[0][0].shape[1], prompt_embeds_list[1][0].shape[1])
    # print(cut_to)
    cut_to=77 
    prompt_embeds= torch.cat([prompt_embeds_list[0][0][:,:cut_to], prompt_embeds_list[1][0][:,:cut_to]], dim=-1)
    prompt_embeds = prompt_embeds.to(device=device,dtype=unet_dtype, )
    class_tokens_mask = class_tokens_mask.to(device=device) # TODO: ignoring two-prompt case
    # get unconditional embeddings for classifier free guidance
    zero_out_negative_prompt = negative_prompt is None and force_zeros_for_empty_prompts
    if do_classifier_free_guidance and negative_prompt_embeds is None and zero_out_negative_prompt:
        negative_prompt_embeds = torch.zeros_like(prompt_embeds)
        negative_pooled_prompt_embeds = torch.zeros_like(pooled_prompt_embeds)
    elif do_classifier_free_guidance and negative_prompt_embeds is None:
        negative_prompt = negative_prompt or ""
        negative_prompt_2 = negative_prompt_2 or negative_prompt
        
        # normalize str to list
        #negative_prompt = batch_size * [negative_prompt] if isinstance(negative_prompt, str) else negative_prompt
        negative_prompt = str(negative_prompt) if isinstance(negative_prompt, list) else negative_prompt
        negative_prompt_2 = (
            batch_size * [negative_prompt_2] if isinstance(negative_prompt_2, str) else negative_prompt_2
        )
        
        uncond_tokens: List[str]
        if prompt is not None and type(prompt) is not type(negative_prompt):
            raise TypeError(
                f"`negative_prompt` should be the same type to `prompt`, but got {type(negative_prompt)} !="
                f" {type(prompt)}."
            )
        # elif batch_size != len(negative_prompt):
        #     raise ValueError(
        #         f"`negative_prompt`: {negative_prompt} has batch size {len(negative_prompt)}, but `prompt`:"
        #         f" {prompt} has batch size {batch_size}. Please make sure that passed `negative_prompt` matches"
        #         " the batch size of `prompt`."
        #     )
        else:
            uncond_tokens = [negative_prompt, negative_prompt_2]
        # ng 没用上
        negative_prompt_embeds_list = []
        for negative_prompt, tokenizer, text_encoder in zip(uncond_tokens, tokenizers, [clip.cond_stage_model.clip_l.encode_token_weights,clip.cond_stage_model.clip_g.encode_token_weights]):
            # if isinstance(self, TextualInversionLoaderMixin):
            #     negative_prompt = self.maybe_convert_prompt(negative_prompt, tokenizer)
            
            max_length = prompt_embeds.shape[1]
            uncond_input=tokenizer.encode(negative_prompt)
            # uncond_input = tokenizer(
            #     negative_prompt,
            #     padding="max_length",
            #     max_length=max_length,
            #     truncation=True,
            #     return_tensors="pt",
            # )
            # negative_prompt_embeds = text_encoder(
            #     uncond_input.input_ids.to(device),
            #     output_hidden_states=True,
            # )
            clean_input_ids_n=uncond_input[:77] if len(uncond_input) > 77 else uncond_input + [0]*(77 - len(uncond_input))
                
            clean_input_ids_1_=[(i,1.0) for i in clean_input_ids_n]
            clean_input_ids_2_=[(i,1.0) for i in clean_input_ids_n if i==clean_input_ids_n[0] or i==49407 ]
            clean_input_ids_2_=clean_input_ids_2_[:77] if len(clean_input_ids_2_) > 77 else clean_input_ids_2_ + [(0,1.0)]*(77 - len(clean_input_ids_2_))
            clean_input_ids_list_n=[clean_input_ids_1_,clean_input_ids_2_]
            prompt_embeds_n, pooled_prompt_embeds_n = text_en(clean_input_ids_list_n)
            negative_prompt_embeds_list.append((prompt_embeds_n,pooled_prompt_embeds_n))


          
            # We are only ALWAYS interested in the pooled output of the final text encoder
            # negative_pooled_prompt_embeds = negative_prompt_embeds[0]
            # negative_prompt_embeds = negative_prompt_embeds.hidden_states[-2]
            
            # negative_prompt_embeds_list.append(negative_prompt_embeds)

        negative_pooled_prompt_embeds=negative_prompt_embeds_list[1][1]
        #negative_prompt_embeds = torch.concat(negative_prompt_embeds_list, dim=-1)
        cut_to=77 
        negative_prompt_embeds= torch.cat([negative_prompt_embeds_list[0][0][:,:cut_to], negative_prompt_embeds_list[1][0][:,:cut_to]], dim=-1)
        negative_prompt_embeds = negative_prompt_embeds.to(device=device,dtype=unet_dtype, )
    
    
    bs_embed, seq_len, _ = prompt_embeds.shape
    
    if do_classifier_free_guidance:
        # duplicate unconditional embeddings for each generation per prompt, using mps friendly method
        seq_len = negative_prompt_embeds.shape[1]
        
        
        negative_prompt_embeds = negative_prompt_embeds.to(dtype=unet_dtype, device=device)
        
        negative_prompt_embeds = negative_prompt_embeds.repeat(1, num_images_per_prompt, 1)
        negative_prompt_embeds = negative_prompt_embeds.view(batch_size * num_images_per_prompt, seq_len, -1)
    
    pooled_prompt_embeds = pooled_prompt_embeds.repeat(1, num_images_per_prompt).view(
        bs_embed * num_images_per_prompt, -1
    )
    if do_classifier_free_guidance:
        negative_pooled_prompt_embeds = negative_pooled_prompt_embeds.repeat(1, num_images_per_prompt).view(
            bs_embed * num_images_per_prompt, -1
        )
    

    return prompt_embeds, negative_prompt_embeds,pooled_prompt_embeds, negative_pooled_prompt_embeds, class_tokens_mask


def photomaker_clip_v2(clip,model_,prompt_list,negative_prompt,input_id_images,id_embeds,trigger_word="img",num_images_per_prompt=1,nc_flag=False,prompt_2=None,prompt_embeds=None,pooled_prompt_embeds=None,class_tokens_mask=None,prompt_embeds_text_only=None,negative_pooled_prompt_embeds=None):
    id_encoder=model_.id_encoder
    id_image_processor=model_.id_image_processor

    tokenizer=model_.tokenizer_2
    batch_size=len(prompt_list)

    num_id_images = len(input_id_images) #双角色时，须确保输入的图片是列表，且跟prompt对应上
    prompt_arr = prompt_list
    negative_prompt_embeds_arr = []
    prompt_embeds_text_only_arr = []
    prompt_embeds_arr = []
    latents_arr = []
    add_time_ids_arr = []
    negative_pooled_prompt_embeds_arr = []
    pooled_prompt_embeds_text_only_arr = []
    pooled_prompt_embeds_arr = []
    for prompt in prompt_arr:
        (
            prompt_embeds,
            _,
            pooled_prompt_embeds,
            _,
            class_tokens_mask
        ) = encode_prompt_with_trigger_word_v2(clip,model_,
            prompt=prompt,
            prompt_2=prompt_2,
            device=device,
            num_id_images=num_id_images,
            class_tokens_mask=class_tokens_mask,
            num_images_per_prompt=num_images_per_prompt,
            do_classifier_free_guidance=True,
            negative_prompt=negative_prompt,
            negative_prompt_2=None,
            prompt_embeds=None,
            negative_prompt_embeds=None,
            pooled_prompt_embeds=None,
            negative_pooled_prompt_embeds=None,
            nc_flag=nc_flag,
        )

        # 4. Encode input prompt without the trigger word for delayed conditioning
        # encode, remove trigger word token, then decode
        tokens_text_only = tokenizer.encode(prompt, add_special_tokens=False)
        trigger_word_token = tokenizer.convert_tokens_to_ids(trigger_word)
        if not nc_flag:
            tokens_text_only.remove(trigger_word_token)
        prompt_text_only = tokenizer.decode(tokens_text_only, add_special_tokens=False)
        #print(prompt_text_only)
        (
            prompt_embeds_text_only,
            negative_prompt_embeds,
            pooled_prompt_embeds_text_only, # TODO: replace the pooled_prompt_embeds with text only prompt
            negative_pooled_prompt_embeds,
        ) = cf_clip_single(clip,
            prompt=prompt_text_only,
            prompt_2=prompt_2,
            device=device,
            num_images_per_prompt=num_images_per_prompt,
            do_classifier_free_guidance=True,
            negative_prompt=negative_prompt,
            negative_prompt_2=None,
            prompt_embeds=prompt_embeds_text_only,
            negative_prompt_embeds=None,
            pooled_prompt_embeds=None,
            negative_pooled_prompt_embeds=negative_pooled_prompt_embeds,
        )

        # 5. Prepare the input ID images
        dtype = next(id_encoder.parameters()).dtype
        if not isinstance(input_id_images[0], torch.Tensor):
            id_pixel_values = id_image_processor(input_id_images, return_tensors="pt").pixel_values

        id_pixel_values = id_pixel_values.unsqueeze(0).to(device=device, dtype=dtype) # TODO: multiple prompts

        # duplicate text embeddings for each generation per prompt, using mps friendly method
        if not nc_flag:
            if id_embeds is not None:
                id_embeds = id_embeds.unsqueeze(0).to(device=device, dtype=dtype)
               
                prompt_embeds = id_encoder(id_pixel_values, prompt_embeds, class_tokens_mask, id_embeds) #torch.Size([1, 1, 3, 224, 224]) torch.Size([1, 77, 2048]) torch.Size([1, 77]) torch.Size([1, 512])
            else:
                prompt_embeds = id_encoder(id_pixel_values, prompt_embeds, class_tokens_mask)
            # 6. Get the update text embedding with the stacked ID embedding
           
            bs_embed, seq_len, _ = prompt_embeds.shape
            prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1)
            prompt_embeds = prompt_embeds.view(bs_embed * num_images_per_prompt, seq_len, -1)
            pooled_prompt_embeds = pooled_prompt_embeds.repeat(1, num_images_per_prompt).view(
                bs_embed * num_images_per_prompt, -1
            )
            pooled_prompt_embeds_arr.append(pooled_prompt_embeds)
            pooled_prompt_embeds = None

        negative_prompt_embeds_arr.append(negative_prompt_embeds)
        negative_prompt_embeds = None
        negative_pooled_prompt_embeds_arr.append(negative_pooled_prompt_embeds)
        negative_pooled_prompt_embeds = None
        prompt_embeds_text_only_arr.append(prompt_embeds_text_only)
        prompt_embeds_text_only = None
        prompt_embeds_arr.append(prompt_embeds)
        prompt_embeds = None
        pooled_prompt_embeds_text_only_arr.append(pooled_prompt_embeds_text_only)
        pooled_prompt_embeds_text_only = None

    emb_dict={"negative_prompt_embeds_arr":negative_prompt_embeds_arr,
              "negative_prompt_embeds":negative_prompt_embeds,
              "negative_pooled_prompt_embeds_arr":negative_pooled_prompt_embeds_arr,
              "negative_pooled_prompt_embeds":negative_pooled_prompt_embeds,
              "prompt_embeds_text_only_arr":prompt_embeds_text_only_arr,
              "prompt_embeds_text_only":prompt_embeds_text_only,
              "prompt_embeds_arr":prompt_embeds_arr,
              "prompt_embeds":prompt_embeds,
              "pooled_prompt_embeds":pooled_prompt_embeds,
              "pooled_prompt_embeds_text_only_arr":pooled_prompt_embeds_text_only_arr,
              "pooled_prompt_embeds_text_only":pooled_prompt_embeds_text_only,
              "pooled_prompt_embeds_arr":pooled_prompt_embeds_arr,
              "batch_size":batch_size,
              }    
    return emb_dict

def Loader_UNO(extra_info,offload,quantize_mode,save_quantezed,lora_rank):
    from .UNO.uno.flux.pipeline import UNOPipeline
    from accelerate import Accelerator
    accelerator = Accelerator()
    device= accelerator.device
    if extra_info is not None:
        model_path = extra_info["unet_path"]
        lora_list=[i for i in [extra_info.get("lora1_path"),extra_info.get("lora2_path")] if i is not None]
        if lora_list:
            lora_ckpt_path=lora_list[0]
            if "dit_lora" not in lora_ckpt_path.lower():
                raise ValueError("lora_ckpt_path must be a dit_lora checkpoint")
            only_lora=True
        else:
            only_lora=False
    
        
        use_fp8=True if quantize_mode=="fp8"  else False

        if  "schnell" in model_path:
            model_type="flux-schnell"
        else:
            model_type="flux-dev"
     
        
        if "schnell" in model_path:
            model_type="flux-schnell"
        else:
            model_type="flux-dev"
        
        if only_lora :
            from .UNO.uno.flux.util import load_flow_model_only_lora
            model_ = load_flow_model_only_lora(
                model_type,model_path,lora_ckpt_path,use_fp8, device="cpu" if offload else device,save_quantezed=save_quantezed, lora_rank=lora_rank)
        else:
            from .UNO.uno.flux.util import load_flow_model
            model_ = load_flow_model(model_type,model_path,use_fp8, device="cpu" if offload else device,save_quantezed=save_quantezed)
        model =  UNOPipeline(
            model_,
            device,
            offload,
            use_fp8,
        )
    else:
        raise ValueError("must use EasyFunction_Lite node")
    return model

def load_pipeline_realcustom(cf_model,realcustom_checkpoint):
    from .RealCustom.models.unet_2d_condition_custom import UNet2DConditionModel as UNet2DConditionModelDiffusers
    if "shallow" in realcustom_checkpoint.lower():
        unet_config_path=os.path.join(cur_path, "RealCustom/configs/realcustom_sigdino_highres_shallow.json")
    else:
        unet_config_path=os.path.join(cur_path, "RealCustom/configs/realcustom_sigdino_highres.json")
    with open(unet_config_path, 'r') as f:
        unet_config = json.load(f)
    
    # Settings for image encoder
    vision_model_config = unet_config.pop("vision_model_config", None)
    vision_model_config_ar = vision_model_config.pop("vision_model_config", None)



    unet_type = unet_config.pop("type", None)
    unet_model = UNet2DConditionModelDiffusers(**unet_config)
    unet_model.eval()

    #cf_state_dict=torch.load(unet_checkpoint,weights_only=False, map_location="cpu")
    cf_state_dict = cf_model.model.diffusion_model.state_dict()
    unet_state_dict = cf_model.model.model_config.process_unet_state_dict_for_saving(cf_state_dict)

    from diffusers.pipelines.stable_diffusion.convert_from_ckpt import convert_ldm_unet_checkpoint
    cf_state_dict = convert_ldm_unet_checkpoint(unet_state_dict, unet_model.config)
    unet_model.load_state_dict(cf_state_dict, strict=False)

    realcustom_DICT=torch.load(realcustom_checkpoint,weights_only=False, map_location="cpu")
    unet_model.load_state_dict(realcustom_DICT, strict=False)
    unet_model = torch.compile(unet_model, disable=True)# CHECK THIS

    del cf_state_dict,realcustom_DICT
    gc.collect()
    torch.cuda.empty_cache()
    print("loading unet base model finished.")

    return unet_model,vision_model_config_ar,unet_type




def load_clip_clipvsion(clip_paths,token_paths,dino_path,siglip_path,vision_model_config):
    from .RealCustom.models.text import TextModel
    from transformers import CLIPTextModel,CLIPTextConfig
    from .RealCustom.utils import instantiate_from_config
    
    # text_encoder=CLIPTextModel.from_pretrained(clip_paths[0], local_files_only=True)
    # text_encoder_2=CLIPTextModel.from_pretrained(clip_paths[1], local_files_only=True)

    from safetensors.torch import load_file
    text_encoder_config=CLIPTextConfig.from_pretrained(token_paths[0], local_files_only=True)
    text_encoder = CLIPTextModel(text_encoder_config)
    text_encoder_sd_1 = load_file(clip_paths[0])
    text_encoder.load_state_dict(text_encoder_sd_1, strict=False)
    
    
    text_encoder_config_2=CLIPTextConfig.from_pretrained(token_paths[1], local_files_only=True)
    text_encoder_2 = CLIPTextModel(text_encoder_config_2)
    text_encoder_sd_2 = load_file(clip_paths[1])
    text_encoder_2.load_state_dict(text_encoder_sd_2, strict=False)
    del text_encoder_sd_2,text_encoder_sd_1
    gc.collect()


    text_model = TextModel([text_encoder,text_encoder_2],token_paths,["penultimate_nonorm"])
    text_model.eval()

    siglip_path = os.path.normpath(siglip_path).replace("\\", "/")
    dino_path = os.path.normpath(dino_path).replace("\\", "/")

    vision_model_config["params"]["siglip_path"] = siglip_path
    vision_model_config["params"]["dino_path"] = dino_path

    vision_model = instantiate_from_config(vision_model_config)
    vision_model = vision_model.eval().to(device)

    print("loading image model and text_model finished.")
    return text_model,vision_model




def realcustom_clip_emb(text_model,vision_model,vae_config,vae_downsample_factor,positive_prompt,negative_prompt,positive_image,target_phrase,width,height,device,samples_per_prompt,guidance_weight=7.5):
    from  .RealCustom.inference.inference_utils import find_phrase_positions_in_text
    import torchvision
    with torch.no_grad():
        image_metadata_validate = torch.tensor(
            data=[
                    width,     # original_height
                    height,    # original_width
                    0,              # coordinate top
                    0,              # coordinate left
                    width,     # target_height
                    height,    # target_width
                ],
                device=device,
                dtype=torch.float32
            ).view(1, -1).repeat(samples_per_prompt, 1)
       
        if guidance_weight != 1:
            text_negative_output = text_model(negative_prompt)
        # Compute target phrases
        target_token = torch.zeros(1, 77).to(device)
        
        positions = find_phrase_positions_in_text(positive_prompt, target_phrase)
        for position in positions:
            prompt_before = positive_prompt[:position] # NOTE We do not need -1 here because the SDXL text encoder does not encode the trailing space.
            prompt_include = positive_prompt[:position+len(target_phrase)]
            #print("prompt before: ", prompt_before, ", prompt_include: ", prompt_include)
            prompt_before_length = text_model.get_vaild_token_length(prompt_before) + 1
            prompt_include_length = text_model.get_vaild_token_length(prompt_include) + 1
           # print("prompt_before_length: ", prompt_before_length, ", prompt_include_length: ", prompt_include_length)
            target_token[:, prompt_before_length:prompt_include_length] = 1

        # Text used for progress bar
        pbar_text = positive_prompt[:40]

        # Compute text embeddings
        text_positive_output = text_model(positive_prompt)
        text_positive_embeddings = text_positive_output.embeddings.repeat_interleave(samples_per_prompt, dim=0)
        text_positive_pooled = text_positive_output.pooled[-1].repeat_interleave(samples_per_prompt, dim=0)
        if guidance_weight != 1:
            text_negative_embeddings = text_negative_output.embeddings.repeat_interleave(samples_per_prompt, dim=0)
            text_negative_pooled = text_negative_output.pooled[-1].repeat_interleave(samples_per_prompt, dim=0)
        
        # Compute image embeddings
    # positive_image = Image.open(positive_promt_image_path).convert("RGB")
        positive_image = torchvision.transforms.ToTensor()(positive_image)

        positive_image = positive_image.unsqueeze(0).repeat_interleave(samples_per_prompt, dim=0)
        positive_image = torch.nn.functional.interpolate(
            positive_image, 
            size=(768, 768), 
            mode="bilinear", 
            align_corners=False
        )
        negative_image = torch.zeros_like(positive_image)
        #print("positive_image:",positive_image.size(), negative_image.size()) #torch.Size([1, 3, 768, 768]) torch.Size([1, 3, 768, 768])
        positive_image = positive_image.to(device)
        negative_image = negative_image.to(device)

        positive_image_dict = {"image_ref": positive_image}
        positive_image_output = vision_model(positive_image_dict, device=device)
        #positive_image_output = daul_encoder(vision_model,positive_image_dict)

        negative_image_dict = {"image_ref": negative_image}
        negative_image_output = vision_model(negative_image_dict, device=device)
        #negative_image_output=daul_encoder(vision_model,negative_image_dict)
        # Initialize latent with input latent + noise (i2i) / pure noise (t2i)
        # latent = torch.randn(
        #     size=[
        #         samples_per_prompt,
        #         vae_config["latent_channels"],
        #         height // vae_downsample_factor,
        #         width // vae_downsample_factor
        #     ],
        #     device=device,
        #     generator=torch.Generator(device).manual_seed(seed))
        target_h = (height // vae_downsample_factor) // 2
        target_w = (width // vae_downsample_factor) // 2

        emb_dict={
            "text_positive_embeddings":text_positive_embeddings.to(device),
            "text_positive_pooled":text_positive_pooled.to(device),
            "image_metadata_validate":image_metadata_validate.to(device),
            "positive_image_output":positive_image_output,
            "text_negative_embeddings":text_negative_embeddings.to(device),
            "text_negative_pooled":text_negative_pooled.to(device),
            "negative_image_output":negative_image_output,
        }

        latent_dict={
            "vae_config":vae_config,
            "vae_downsample_factor":vae_downsample_factor,
            "target_h":target_h,
            "target_w":target_w,
            "pbar_text":pbar_text,
            "guidance_weight":guidance_weight,
            "target_token":target_token,
            
        }

        return emb_dict,latent_dict


def realcustom_infer(unet_model,sample_steps,mask_reused_step,emb_dict,latent_dict,mask_scope,mask_strategy,guidance_weight,height,width,seed,
                     device,schedule_type="squared_linear",schedule_shift_snr=1.0,scheduler_type="ddim",unet_prediction="epsilon"):
    from .RealCustom.schedulers.ddim import DDIMScheduler
    from .RealCustom.schedulers.dpm_s import DPMSolverSingleStepScheduler
    from .RealCustom.schedulers.utils import get_betas
    from .RealCustom.inference.mask_generation import mask_generation
    from  .RealCustom.inference.inference_utils import classifier_free_guidance_image_prompt_cascade
    from torchvision.transforms.functional import to_pil_image
    from tqdm import tqdm 
    # Initialize ddim scheduler
    ddim_train_steps = 1000
    ddim_betas = get_betas(name=schedule_type, num_steps=ddim_train_steps, shift_snr=schedule_shift_snr, terminal_pure_noise=False)
    scheduler_class = DPMSolverSingleStepScheduler if scheduler_type == 'dpm' else DDIMScheduler
    scheduler = scheduler_class(betas=ddim_betas, num_train_timesteps=ddim_train_steps, num_inference_timesteps=sample_steps, device=device)
    infer_timesteps = scheduler.timesteps
    #print(infer_timesteps)
    # Real Reverse diffusion process.
    text2image_crossmap_2d_all_timesteps_list = []
    current_step = 0
    pbar_text = latent_dict.get("pbar_text")
    #guidance_weight = latent_dict.get("guidance_weight")
    unet_model.to(device)
    latent=latent_dict.get("latent")
    latent = torch.randn(
            size=[
                1,
                latent_dict.get("vae_config")["latent_channels"],
                height // latent_dict.get("vae_downsample_factor"),
                width // latent_dict.get("vae_downsample_factor")
            ],
            device=device,
            generator=torch.Generator(device).manual_seed(seed))
    

    with torch.no_grad():
        for timestep in tqdm(iterable=infer_timesteps, desc=f"[{pbar_text}]", dynamic_ncols=True):
            
            if current_step < mask_reused_step:
                pred_cond, pred_cond_dict = unet_model(
                    sample=latent,
                    timestep=timestep,
                    encoder_hidden_states=emb_dict.get("text_positive_embeddings"),
                    encoder_attention_mask=None,
                    added_cond_kwargs=dict(
                        text_embeds=emb_dict.get("text_positive_pooled"),
                        time_ids=emb_dict.get("image_metadata_validate")
                    ),
                    vision_input_dict=None,
                    vision_guided_mask=None,
                    return_as_origin=False,
                    return_text2image_mask=True,
                )
                
                crossmap_2d_avg = mask_generation(
                    crossmap_2d_list=pred_cond_dict.get("text2image_crossmap_2d",[]), selfmap_2d_list=pred_cond_dict.get("self_attention_map", []), 
                    target_token=latent_dict.get("target_token"), mask_scope=mask_scope,
                    mask_target_h=latent_dict.get("target_h"), mask_target_w=latent_dict.get("target_w"), mask_mode=mask_strategy,
                )
            else:
                # using previous step's mask
                crossmap_2d_avg = text2image_crossmap_2d_all_timesteps_list[-1].squeeze(1)
            if crossmap_2d_avg.dim() == 5: # Means that each layer uses a separate mask weight.
                text2image_crossmap_2d_all_timesteps_list.append(crossmap_2d_avg.mean(dim=2).unsqueeze(1))
            else:
                text2image_crossmap_2d_all_timesteps_list.append(crossmap_2d_avg.unsqueeze(1))

            pred_cond, pred_cond_dict = unet_model(
                sample=latent,
                timestep=timestep,
                encoder_hidden_states=emb_dict.get("text_positive_embeddings"),
                encoder_attention_mask=None,
                added_cond_kwargs=dict(
                    text_embeds=emb_dict.get("text_positive_pooled"),
                    time_ids=emb_dict.get("image_metadata_validate")
                ),
                vision_input_dict=emb_dict.get("positive_image_output"),
                vision_guided_mask=crossmap_2d_avg,
                return_as_origin=False,
                return_text2image_mask=True,
                multiple_reference_image=False
            )

            crossmap_2d_avg_neg = crossmap_2d_avg.mean(dim=1, keepdim=True)
            pred_negative, pred_negative_dict = unet_model(
                sample=latent,
                timestep=timestep,
                encoder_hidden_states=emb_dict.get("text_negative_embeddings"),
                encoder_attention_mask=None,
                added_cond_kwargs=dict(
                    text_embeds=emb_dict.get("text_negative_pooled"),
                    time_ids=emb_dict.get("image_metadata_validate")
                ),
                vision_input_dict=emb_dict.get("negative_image_output"),
                vision_guided_mask=crossmap_2d_avg,
                return_as_origin=False,
                return_text2image_mask=True,
                multiple_reference_image=False
            )

            pred = classifier_free_guidance_image_prompt_cascade(
                pred_t_cond=None, pred_ti_cond=pred_cond, pred_uncond=pred_negative, 
                guidance_weight_t=guidance_weight, guidance_weight_i=guidance_weight, 
                guidance_stdev_rescale_factor=0, cfg_rescale_mode="naive_global_direct"
            )
            step = scheduler.step(
                model_output=pred,
                model_output_type=unet_prediction,
                timestep=timestep,
                sample=latent,
                )

            latent = step.prev_sample

            current_step += 1

        sample=step.pred_original_sample/0.13025
       
        return sample
    

def load_realcustom_vae(vae,device):
    from .RealCustom.models.vae import AutoencoderKL
    from diffusers.pipelines.stable_diffusion.convert_from_ckpt import convert_ldm_vae_checkpoint

    vae_config_path = os.path.join(cur_path, "local_repo/vae/config.json")
    # Initialize vae model
    with open(vae_config_path, 'r') as vae_config_file:
        vae_config = json.load(vae_config_file)
    vae_downsample_factor = 2 ** (len(vae_config["block_out_channels"]) - 1) # 2 ** 3 = 8
    vae_model = AutoencoderKL(**vae_config)
    vae_model.eval().to(device)

    vae_state_dict=vae.get_sd()
    #AE = AutoencoderKL.from_config(ae_config).to(device, torch.float16)
    cf_state_dict = convert_ldm_vae_checkpoint(vae_state_dict, vae_model.config)
    vae_model.load_state_dict(cf_state_dict, strict=False)
    del cf_state_dict,vae_state_dict

   
    #vae_decoder = torch.compile(lambda x: vae_model.decode(x / vae_model.scaling_factor).sample.clip(-1, 1), disable=True)
    vae_encoder = torch.compile(lambda x: vae_model.encode(x).latent_dist.mode().mul_(vae_model.scaling_factor), disable=True)
    print("loading vae finished.")
    return vae_encoder,vae_downsample_factor,vae_config



def load_pipeline_instant_character(extra_info,ip_adapter_path,VAE,quantize_mode):
    from .InstantCharacter.pipeline_wrapper import InstantCharacterFluxPipeline
    from diffusers import BitsAndBytesConfig as DiffusersBitsAndBytesConfig,FluxTransformer2DModel
    if extra_info.get("svdq_repo") is not None:
        print("use svdq")
        from nunchaku import NunchakuFluxTransformer2dModel    
        transformer = NunchakuFluxTransformer2dModel.from_pretrained(extra_info.get("svdq_repo"),offload=True)
        try:
            transformer.set_attention_impl("nunchaku-fp16")
        except:
            pass
        vae=convert_cfvae2diffuser(VAE,use_flux=True)
        pipe = InstantCharacterFluxPipeline.from_pretrained(os.path.join(cur_path,"config/FLUX.1-dev"),vae=vae,transformer=transformer,text_encoder=None,text_encoder_2=None, torch_dtype=torch.bfloat16)
    elif extra_info.get("gguf_path") is not None:
        print("use gguf quantization")   
        from diffusers import  GGUFQuantizationConfig
        transformer = FluxTransformer2DModel.from_single_file(
            extra_info.get("gguf_path") ,
            config=os.path.join(cur_path, "config/FLUX.1-dev/transformer"),
            quantization_config=GGUFQuantizationConfig(compute_dtype=torch.bfloat16),
            torch_dtype=torch.bfloat16,
        )
        vae=convert_cfvae2diffuser(VAE,use_flux=True)
        pipe = InstantCharacterFluxPipeline.from_pretrained(os.path.join(cur_path,"config/FLUX.1-dev"),vae=vae,transformer=transformer,text_encoder=None,text_encoder_2=None, torch_dtype=torch.bfloat16)
    elif extra_info.get("unet_path"):

        if quantize_mode=="fp8":
            transformer = FluxTransformer2DModel.from_single_file(extra_info.get("unet_path"), config=os.path.join(cur_path,"config/FLUX.1-dev/transformer/config.json"),
                                                                            torch_dtype=torch.bfloat16)
        else:
            from safetensors.torch import load_file
            t_state_dict=load_file(extra_info.get("unet_path"))
            unet_config = FluxTransformer2DModel.load_config(os.path.join(cur_path,"config/FLUX.1-dev/transformer/config.json"))
            transformer = FluxTransformer2DModel.from_config(unet_config).to(torch.bfloat16)
            transformer.load_state_dict(t_state_dict, strict=False)
            del t_state_dict
            gc_cleanup()
        vae=convert_cfvae2diffuser(VAE,use_flux=True)
        pipe = InstantCharacterFluxPipeline.from_pretrained(os.path.join(cur_path,"config/FLUX.1-dev"),transformer=transformer,vae=vae,text_encoder=None,text_encoder_2=None, torch_dtype=torch.bfloat16)
    else:
        print(f"use {quantize_mode} quantization") 
        find_flux=[i for i in [extra_info.get("repo1_path") , extra_info.get("repo2_path")] if "flux" in i ]
        if  not find_flux:
            raise ValueError("can not find flux repo")
        if quantize_mode=="fp8":
            transformer = FluxTransformer2DModel.from_pretrained(
                find_flux[0],
                subfolder="transformer",
                quantization_config=DiffusersBitsAndBytesConfig(load_in_8bit=True,),
                torch_dtype=torch.bfloat16,
            )
                   
        elif quantize_mode=="nf4": #nf4
            transformer = FluxTransformer2DModel.from_pretrained(
                find_flux[0],
                subfolder="transformer",
                quantization_config=DiffusersBitsAndBytesConfig(
                    load_in_4bit=True, bnb_4bit_quant_type="nf4", bnb_4bit_compute_dtype=torch.bfloat16
                ),
                torch_dtype=torch.bfloat16,)
        else:
            transformer = FluxTransformer2DModel.from_pretrained(
                find_flux[0],
                subfolder="transformer",
                 torch_dtype=torch.bfloat16,
                )

        pipe = InstantCharacterFluxPipeline.from_pretrained(find_flux[0],transformer=transformer,text_encoder=None,text_encoder_2=None, torch_dtype=torch.bfloat16)
    # if not cf_model.get("use_svdq"):
    #     pipe.enable_model_cpu_offload()
    pipe.to(device)
    pipe.init_adapter(
        subject_ipadapter_cfg=dict(subject_ip_adapter_path=ip_adapter_path, nb_token=1024), 
    )
    
    return pipe

def cf_flux_prompt_clip(cf_clip,prompt):
    tokens = cf_clip.tokenize(prompt)
    tokens["t5xxl"] = cf_clip.tokenize(prompt)["t5xxl"]
    prompt_embeds = cf_clip.encode_from_tokens(tokens, return_pooled=True, return_dict=True).pop("cond")
    tokens["l"] = cf_clip.tokenize(prompt)["l"]
    pooled_prompt_embeds = cf_clip.encode_from_tokens(tokens, return_dict=True).pop("pooled_output")
    prompt_embeds=prompt_embeds.to(device,torch.bfloat16)
    pooled_prompt_embeds=pooled_prompt_embeds.to(device,torch.bfloat16)
    text_ids = torch.zeros(prompt_embeds.shape[1], 3).to(device=device, dtype=torch.bfloat16)
    return prompt_embeds,pooled_prompt_embeds,text_ids


def load_dual_clip(image_encoder_path,image_encoder_2_path,device,dtype):
    from transformers import SiglipVisionModel, SiglipImageProcessor, AutoModel, AutoImageProcessor
    # image encoder
    print(f"=> loading image_encoder_1: {image_encoder_path}")
    image_encoder = SiglipVisionModel.from_pretrained(image_encoder_path)
    image_processor = SiglipImageProcessor.from_pretrained(image_encoder_path)
    image_encoder.eval()
    image_encoder.to(device, dtype=dtype)


    # image encoder 2
    print(f"=> loading image_encoder_2: {image_encoder_2_path}")
    image_encoder_2 = AutoModel.from_pretrained(image_encoder_2_path)
    image_processor_2 = AutoImageProcessor.from_pretrained(image_encoder_2_path)
    image_encoder_2.eval()
    image_encoder_2.to(device, dtype=dtype)
    image_processor_2.crop_size = dict(height=384, width=384)
    image_processor_2.size = dict(shortest_edge=384)


    return image_encoder,image_processor,image_encoder_2,image_processor_2



@torch.inference_mode()
def instant_character_id_clip(subject_image,siglip_image_encoder,siglip_image_processor,dino_image_encoder_2,dino_image_processor_2,device,dtype):

   
    from einops import rearrange
    
    def encode_siglip_image_emb(siglip_image, device, dtype):
        siglip_image = siglip_image.to(device, dtype=dtype)
        res = siglip_image_encoder(siglip_image, output_hidden_states=True)

        siglip_image_embeds = res.last_hidden_state

        siglip_image_shallow_embeds = torch.cat([res.hidden_states[i] for i in [7, 13, 26]], dim=1)
        
        return siglip_image_embeds, siglip_image_shallow_embeds


    def encode_dinov2_image_emb(dinov2_image, device, dtype):
        dinov2_image = dinov2_image.to(device, dtype=dtype)
        res = dino_image_encoder_2(dinov2_image, output_hidden_states=True)

        dinov2_image_embeds = res.last_hidden_state[:, 1:]

        dinov2_image_shallow_embeds = torch.cat([res.hidden_states[i][:, 1:] for i in [9, 19, 29]], dim=1)

        return dinov2_image_embeds, dinov2_image_shallow_embeds
    


    def encode_image_emb(siglip_image, device, dtype):
        object_image_pil = siglip_image
        object_image_pil_low_res = [object_image_pil.resize((384, 384))]
        object_image_pil_high_res = object_image_pil.resize((768, 768))
        object_image_pil_high_res = [
            object_image_pil_high_res.crop((0, 0, 384, 384)),
            object_image_pil_high_res.crop((384, 0, 768, 384)),
            object_image_pil_high_res.crop((0, 384, 384, 768)),
            object_image_pil_high_res.crop((384, 384, 768, 768)),
        ]
        nb_split_image = len(object_image_pil_high_res)

        siglip_image_embeds = encode_siglip_image_emb(
            siglip_image_processor(images=object_image_pil_low_res, return_tensors="pt").pixel_values, 
            device, 
            dtype
        )
        dinov2_image_embeds = encode_dinov2_image_emb(
            dino_image_processor_2(images=object_image_pil_low_res, return_tensors="pt").pixel_values, 
            device, 
            dtype
        )

        image_embeds_low_res_deep = torch.cat([siglip_image_embeds[0], dinov2_image_embeds[0]], dim=2)
        image_embeds_low_res_shallow = torch.cat([siglip_image_embeds[1], dinov2_image_embeds[1]], dim=2)

        siglip_image_high_res = siglip_image_processor(images=object_image_pil_high_res, return_tensors="pt").pixel_values
        siglip_image_high_res = siglip_image_high_res[None]
        siglip_image_high_res = rearrange(siglip_image_high_res, 'b n c h w -> (b n) c h w')
        siglip_image_high_res_embeds = encode_siglip_image_emb(siglip_image_high_res, device, dtype)
        siglip_image_high_res_deep = rearrange(siglip_image_high_res_embeds[0], '(b n) l c -> b (n l) c', n=nb_split_image)
        dinov2_image_high_res = dino_image_processor_2(images=object_image_pil_high_res, return_tensors="pt").pixel_values
        dinov2_image_high_res = dinov2_image_high_res[None]
        dinov2_image_high_res = rearrange(dinov2_image_high_res, 'b n c h w -> (b n) c h w')
        dinov2_image_high_res_embeds = encode_dinov2_image_emb(dinov2_image_high_res, device, dtype)
        dinov2_image_high_res_deep = rearrange(dinov2_image_high_res_embeds[0], '(b n) l c -> b (n l) c', n=nb_split_image)
        image_embeds_high_res_deep = torch.cat([siglip_image_high_res_deep, dinov2_image_high_res_deep], dim=2)

        image_embeds_dict = dict(
            image_embeds_low_res_shallow=image_embeds_low_res_shallow.to(device),
            image_embeds_low_res_deep=image_embeds_low_res_deep.to(device),
            image_embeds_high_res_deep=image_embeds_high_res_deep.to(device),
        )

        return image_embeds_dict

    subject_image = subject_image.resize((max(subject_image.size), max(subject_image.size)))
    subject_image_embeds_dict = encode_image_emb(subject_image, device, dtype)

    return subject_image_embeds_dict

def Loader_Dreamo(extra_info,VAE,quantize_mode,dreamo_lora_path,cfg_distill_path,Turbo_path,device,dreamo_version):
    from .DreamO.dreamo.dreamo_pipeline import DreamOPipeline
    from diffusers import BitsAndBytesConfig as DiffusersBitsAndBytesConfig,FluxTransformer2DModel
    #vae=convert_cfvae2diffuser(VAE,use_flux=True)
    repo_list=[i for i  in [extra_info.get("repo1_path"),extra_info.get("repo2_path")] if i is not None]
    if repo_list:
        flux_repo=[i for i in repo_list if "flux" in i.lower()]
    else:
        raise "you must fill a flux repo"
    if extra_info.get("gguf_path") is not None: #get error
        print("use gguf quantization")
        from diffusers import  GGUFQuantizationConfig
        transformer = FluxTransformer2DModel.from_single_file(
            extra_info.get("gguf_path"),
            config=os.path.join(cur_path, "config/FLUX.1-dev/transformer"),
            quantization_config=GGUFQuantizationConfig(compute_dtype=torch.bfloat16),
            torch_dtype=torch.bfloat16,
        )
        dreamo_pipeline = DreamOPipeline.from_pretrained(flux_repo, transformer=transformer,torch_dtype=torch.bfloat16)
        dreamo_pipeline.load_dreamo_model(device,dreamo_lora_path,cfg_distill_path,Turbo_path, use_turbo=True,dreamo_version=dreamo_version)
    elif extra_info.get("svdq_repo") is not None:#get error
        print("use svdq")
        from nunchaku import NunchakuFluxTransformer2dModel    
        transformer = NunchakuFluxTransformer2dModel.from_pretrained(extra_info.get("svdq_repo"),offload=True)
        try:
            transformer.set_attention_impl("nunchaku-fp16")
        except:
            pass
        dreamo_pipeline = DreamOPipeline.from_pretrained(flux_repo, transformer=transformer,torch_dtype=torch.bfloat16)
        dreamo_pipeline.load_dreamo_model(device,dreamo_lora_path,cfg_distill_path,Turbo_path, use_turbo=True,use_svdq=True,dreamo_version=dreamo_version)
    elif extra_info.get("unet_path") is not None:
        print("use single unet")
        if quantize_mode=="fp8":
            transformer = FluxTransformer2DModel.from_single_file(extra_info.get("unet_path"), config=os.path.join(cur_path,"config/FLUX.1-dev/transformer/config.json"),
                                                                            torch_dtype=torch.bfloat16)
        else:
            from safetensors.torch import load_file
            t_state_dict=load_file(extra_info.get("unet_path"))
            unet_config = FluxTransformer2DModel.load_config(os.path.join(cur_path,"config/FLUX.1-dev/transformer/config.json"))
            transformer = FluxTransformer2DModel.from_config(unet_config).to(torch.bfloat16)
            transformer.load_state_dict(t_state_dict, strict=False)
            del t_state_dict
            gc_cleanup()
        dreamo_pipeline = DreamOPipeline.from_pretrained(flux_repo, transformer=transformer,torch_dtype=torch.bfloat16)
        dreamo_pipeline.load_dreamo_model(device,dreamo_lora_path,cfg_distill_path,Turbo_path, use_turbo=True,dreamo_version=dreamo_version)
    else:
        print(f"use {quantize_mode} quantization") 
        if quantize_mode=="fp8":
            # transformer = FluxTransformer2DModel.from_pretrained(
            #     cf_model.get("extra_repo"),
            #     subfolder="transformer",
            #     quantization_config=DiffusersBitsAndBytesConfig(load_in_8bit=True,),
            #     torch_dtype=torch.bfloat16,
            # )
            dreamo_pipeline = DreamOPipeline.from_pretrained(flux_repo, torch_dtype=torch.bfloat16)
            dreamo_pipeline.load_dreamo_model(device,dreamo_lora_path,cfg_distill_path,Turbo_path, use_turbo=True,dreamo_version=dreamo_version)
            from optimum.quanto import freeze, qint8, quantize
            quantize(dreamo_pipeline.transformer, qint8)
            freeze(dreamo_pipeline.transformer)
            quantize(dreamo_pipeline.text_encoder_2, qint8)
            freeze(dreamo_pipeline.text_encoder_2)  
        
        else:    
            if quantize_mode=="nf4": #nf4
                transformer = FluxTransformer2DModel.from_pretrained(
                    flux_repo,
                    subfolder="transformer",
                    quantization_config=DiffusersBitsAndBytesConfig(
                        load_in_4bit=True, bnb_4bit_quant_type="nf4", bnb_4bit_compute_dtype=torch.bfloat16
                    ),
                    torch_dtype=torch.bfloat16,)
            else:
                transformer = FluxTransformer2DModel.from_pretrained(
                   flux_repo,
                    subfolder="transformer",
                        torch_dtype=torch.bfloat16,
                    )
            dreamo_pipeline = DreamOPipeline.from_pretrained(flux_repo, transformer=transformer,torch_dtype=torch.bfloat16)
            dreamo_pipeline.load_dreamo_model(device,dreamo_lora_path,cfg_distill_path,Turbo_path, use_turbo=True,dreamo_version=dreamo_version)
    dreamo_pipeline.enable_model_cpu_offload()
    return dreamo_pipeline

@torch.no_grad()
def get_align_face(img,face_helper):
    from .DreamO.dreamo.utils import img2tensor,tensor2img
    from torchvision.transforms.functional import normalize
    # the face preprocessing code is same as PuLID
    face_helper.clean_all()
    image_bgr = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
    face_helper.read_image(image_bgr)
    face_helper.get_face_landmarks_5(only_center_face=True)
    face_helper.align_warp_face()
    if len(face_helper.cropped_faces) == 0:
        return None
    align_face = face_helper.cropped_faces[0]

    input = img2tensor(align_face, bgr2rgb=True).unsqueeze(0) / 255.0
    input = input.to(torch.device("cuda"))
    parsing_out = face_helper.face_parse(normalize(input, [0.485, 0.456, 0.406], [0.229, 0.224, 0.225]))[0]
    parsing_out = parsing_out.argmax(dim=1, keepdim=True)
    bg_label = [0, 16, 18, 7, 8, 9, 14, 15]
    bg = sum(parsing_out == i for i in bg_label).bool()
    white_image = torch.ones_like(input)
    # only keep the face features
    face_features_image = torch.where(bg, white_image, input)
    face_features_image = tensor2img(face_features_image, rgb2bgr=False)

    return face_features_image

def Dreamo_image_encoder(BEN2_path,ref_image1,ref_image2,ref_task1,ref_task2,ref_res):
    from .DreamO.dreamo.utils import img2tensor, resize_numpy_image_area
    ref_conds = []
    #debug_images = []

    ref_images = [ref_image1, ref_image2]
    ref_tasks = [ref_task1, ref_task2]

    for idx, (ref_image, ref_task) in enumerate(zip(ref_images, ref_tasks)):
        if ref_image is not None:
            if ref_task == "id":
                from facexlib.utils.face_restoration_helper import FaceRestoreHelper
                face_helper = FaceRestoreHelper(
                    upscale_factor=1,
                    face_size=512,
                    crop_ratio=(1, 1),
                    det_model='retinaface_resnet50',
                    save_ext='png',
                    device=device,
                )
                ref_image = resize_numpy_image_long(ref_image, 1024)
                ref_image = get_align_face(ref_image,face_helper)
            elif ref_task != "style": # ip
                from .DreamO.tools import BEN2
                bg_rm_model = BEN2.BEN_Base().to(device).eval() #TODO 反复加载 需要修复
                #hf_hub_download(repo_id='PramaLLC/BEN2', filename='BEN2_Base.pth', local_dir='models')
                bg_rm_model.loadcheckpoints(BEN2_path)
                ref_image = bg_rm_model.inference(ref_image) #NEED TO CHECK
                bg_rm_model.to(torch.device('cpu'))
            if ref_task != "id":
                ref_image = resize_numpy_image_area(np.array(ref_image), ref_res * ref_res)
                
            #debug_images.append(ref_image)
            ref_image = img2tensor(ref_image, bgr2rgb=False).unsqueeze(0) / 255.0
            ref_image = 2 * ref_image - 1.0
            ref_conds.append(
                {
                    'img': ref_image,
                    'task': ref_task,
                    'idx': idx + 1,
                }
            )
    return ref_conds

def Loader_Flux_Diffuser(extra_info,omi_lora_path,VAE,quantize_mode):
    from .OmniConsistency.src_inference.pipeline import FluxPipeline as FluxPipeline_dif
    from .OmniConsistency.src_inference.pipeline_ import FluxPipeline as FluxPipeline_dif_original
    from diffusers import BitsAndBytesConfig as DiffusersBitsAndBytesConfig,FluxTransformer2DModel
    from .OmniConsistency.src_inference.lora_helper import set_single_lora
    flux_dif_repo=os.path.join(cur_path,"config/FLUX.1-dev")
    if extra_info.get("gguf_path") is not None: 
        vae=convert_cfvae2diffuser(VAE,use_flux=True)
        print("use gguf quantization")
     
        from diffusers import  GGUFQuantizationConfig
        transformer = FluxTransformer2DModel.from_single_file(
            extra_info.get("gguf_path"),
            config=os.path.join(cur_path, "config/FLUX.1-dev/transformer"),
            quantization_config=GGUFQuantizationConfig(compute_dtype=torch.bfloat16),
            torch_dtype=torch.bfloat16,
        )
        pipeline = FluxPipeline_dif.from_pretrained(flux_dif_repo, vae=vae,transformer=transformer,text_encoder=None,text_encoder_2=None,torch_dtype=torch.bfloat16)
       
    elif extra_info.get("svdq_repo") is not None:
        print("use svdq")
        vae=convert_cfvae2diffuser(VAE,use_flux=True)
        from nunchaku import NunchakuFluxTransformer2dModel    
        transformer = NunchakuFluxTransformer2dModel.from_pretrained(extra_info.get("svdq_repo"),offload=True)
        try:
            transformer.set_attention_impl("nunchaku-fp16")
        except:
            pass
        pipeline = FluxPipeline_dif.from_pretrained(flux_dif_repo, vae=vae,transformer=transformer,text_encoder=None,text_encoder_2=None,torch_dtype=torch.bfloat16)
       
    elif extra_info.get("unet_path") is not None:
        print("use single unet")
        vae=convert_cfvae2diffuser(VAE,use_flux=True)
        if quantize_mode=="fp8":
            transformer = FluxTransformer2DModel.from_single_file(extra_info.get("unet_path"), config=os.path.join(cur_path,"config/FLUX.1-dev/transformer/config.json"),
                                                                            torch_dtype=torch.bfloat16)
        else:
            from safetensors.torch import load_file
            t_state_dict=load_file(extra_info.get("unet_path"))
            unet_config = FluxTransformer2DModel.load_config(os.path.join(cur_path,"config/FLUX.1-dev/transformer/config.json"))
            transformer = FluxTransformer2DModel.from_config(unet_config).to(torch.bfloat16)
            transformer.load_state_dict(t_state_dict, strict=False)
            del t_state_dict
            gc_cleanup()
        pipeline = FluxPipeline_dif.from_pretrained(flux_dif_repo, vae=vae,transformer=transformer,text_encoder=None,text_encoder_2=None,torch_dtype=torch.bfloat16)
        
    else:
        print(f"use {quantize_mode} quantization") 
        repo_list=[i for i  in [extra_info.get("repo1_path"),extra_info.get("repo2_path")] if i is not None]
        if repo_list:
            flux_repo=[i for i in repo_list if "flux" in i.lower()]
        else:
            raise "you must fill a flux repo"
        if quantize_mode=="nf4": #nf4
            transformer = FluxTransformer2DModel.from_pretrained(
                flux_repo[0],
                subfolder="transformer",
                quantization_config=DiffusersBitsAndBytesConfig(
                    load_in_4bit=True, bnb_4bit_quant_type="nf4", bnb_4bit_compute_dtype=torch.bfloat16
                ),
                torch_dtype=torch.bfloat16,)
        elif quantize_mode=="fp8":
            transformer = FluxTransformer2DModel.from_pretrained(
                flux_repo[0],
                subfolder="transformer",
                quantization_config=DiffusersBitsAndBytesConfig(load_in_8bit=True,),
                torch_dtype=torch.bfloat16,
            )
        else:
            transformer = FluxTransformer2DModel.from_pretrained(
                flux_repo[0],
                subfolder="transformer",
                    torch_dtype=torch.bfloat16,
                )
        pipeline = FluxPipeline_dif_original.from_pretrained(flux_repo[0],transformer=transformer,torch_dtype=torch.bfloat16)
    
    multi_lora=[]
    if extra_info.get("lora1_path") is not None:
        multi_lora.append(extra_info.get("lora1_path"))
    if extra_info.get("lora2_path") is not None:
        multi_lora.append(extra_info.get("lora2_path"))
    if multi_lora:
        if len(multi_lora)==1:
            set_single_lora(pipeline.transformer, omi_lora_path, lora_weights=[1], cond_size=512)
            pipeline.unload_lora_weights() 
            pipeline.load_lora_weights(os.path.dirname(multi_lora[0]),weight_name=os.path.basename(multi_lora[0]))      
            
        else:
            from .OmniConsistency.src_inference.lora_helper  import set_multi_lora
            lora_weights = [1]*len(multi_lora)
            set_multi_lora(pipeline.transformer, multi_lora, lora_weights=[[lora_weights]], cond_size=512)
            pipeline.unload_lora_weights() 
    else:
        raise "need chocie a  lora model in EasyFunction_Lite!"
    pipeline.enable_model_cpu_offload()
    return pipeline



def load_lora_for_unet_only(pipeline, lora_path, adapter_name="default", lora_scale=1.0):

    try:
      
        pipeline.load_lora_weights(lora_path, adapter_name=adapter_name)
        pipeline.set_adapters(adapter_name, adapter_weights=lora_scale)
        print(f"成功加载LoRA权重: {adapter_name} (scale: {lora_scale})")
    except Exception as e:
        print(f"加载LoRA权重失败: {e}")
    
    return pipeline