nodes/smthemex-ComfyUI_StoryDiffusion
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

init

Browse Source
This commit is contained in:
smthem
2025-04-05 08:22:39 +08:00
parent 5c3f1945b0
commit 6762cde3eb
66 changed files with 249306 additions and 3468 deletions
Download Patch File Download Diff File Expand all files Collapse all files
PuLID/app_flux.py
+132 -67
View File
@@ -8,50 +8,101 @@ from .flux.util import load_ae, load_clip, load_flow_model, load_t5, load_flow_m
from .pulid.pipeline_flux import PuLIDPipeline
def get_models(name: str,ckpt_path,vae_cf, clip_cf,if_repo,use_quantize,device: torch.device, offload: bool,quantized_mode):
def get_models_(name: str,ckpt_path,vae_cf, clip_cf,if_repo,aggressive_offload,device: torch.device, offload: bool,quantized_mode):
t5 = load_t5(name,clip_cf,if_repo,device, max_length=128)
clip = load_clip(name,clip_cf,if_repo,device)
if if_repo:
name=name.rsplit("/")[-1].lower()
name=name.lower()
if "dev" in name:
name="flux-dev"
else:
name = "flux-schnell"
if quantized_mode=="fp8" or quantized_mode=="nf4":
model = load_flow_model_quintized(name, ckpt_path,quantized_mode,use_quantize,device="cpu" if offload else device)
model = load_flow_model_quintized(name, ckpt_path,quantized_mode,aggressive_offload,device="cpu" if offload else device)
else:
model = load_flow_model(name, ckpt_path,device="cpu" if offload else device)
model.eval()
ae = load_ae(name,vae_cf, device="cpu" if offload else device)
if isinstance(vae_cf,str):
ae = load_ae(name,vae_cf, device="cpu" if offload else device)
else:
ae = None
return model, ae, t5, clip
def get_models(name: str,ckpt_path,if_repo,aggressive_offload,device: torch.device, offload: bool,quantized_mode):
if if_repo:
name=name.lower()
if "dev" in name:
name="flux-dev"
else:
name = "flux-schnell"
if quantized_mode=="fp8" or quantized_mode=="nf4":
model = load_flow_model_quintized(name, ckpt_path,quantized_mode,aggressive_offload,device="cpu" if offload else device)
else:
model = load_flow_model(name, ckpt_path,device="cpu" if offload else device)
model.eval()
return model
def get_emb_flux_pulid(t5,clip,if_repo,seed,prompt,neg_prompt,width,height,num_steps,guidance,device,true_cfg=1.0):
seed = int(seed)
if seed == -1:
seed = None
if isinstance(prompt,str):
text=prompt
elif isinstance(prompt,list):
text=prompt[0]
else:
raise "prompt must be list or str"
opts = SamplingOptions(
text=text,
width=width,
height=height,
num_steps=num_steps,
guidance=guidance,
seed=seed,
)
if opts.seed is None:
opts.seed = torch.Generator(device="cpu").seed()
opts.text = opts.text if isinstance(opts.text,str) else str(opts.text)
print(f"Generating '{opts.text}' with seed {opts.seed}")
use_true_cfg = abs(true_cfg - 1.0) > 1e-2
# prepare input
x = get_noise(
1,
opts.height,
opts.width,
device=device,
dtype=torch.bfloat16,
seed=opts.seed,
)
inp = prepare(t5=t5, clip=clip, img=x, prompt=opts.text,if_repo=if_repo)
inp_neg = prepare(t5=t5, clip=clip, img=x, prompt=neg_prompt,if_repo=if_repo) if use_true_cfg else None
return inp,inp_neg,x
class FluxGenerator:
def __init__(self, model_name: str, ckpt_path,device: str, offload: bool,aggressive_offload: bool, pretrained_model,quantized_mode,clip_vision_path,clip_cf,vae_cf,if_repo,onnx_provider,use_quantize):
def __init__(self, model,device: str, offload: bool,aggressive_offload: bool, pretrained_model,quantized_mode,if_repo,clip_vision_path,ae=None,t5=None,clip=None):
self.device = torch.device(device)
self.offload = offload
self.model_name = model_name
self.aggressive_offload= aggressive_offload
self.ckpt_path=ckpt_path
self.clip_vision_path=clip_vision_path
self.clip_cf=clip_cf
self.use_quantize=use_quantize
self.vae_cf=vae_cf
self.if_repo=if_repo
self.onnx_provider=onnx_provider
self.quantized_mode=quantized_mode
self.model, self.ae, self.t5, self.clip = get_models(
model_name,self.ckpt_path,self.vae_cf,self.clip_cf,self.if_repo,self.use_quantize,
device=self.device,
offload=self.offload,quantized_mode=self.quantized_mode,
)
self.pulid_model = PuLIDPipeline(self.model, device, self.clip_vision_path,weight_dtype=torch.bfloat16,onnx_provider=self.onnx_provider)
self.ae = ae
self.t5=t5
self.clip=clip
self.model = model
self.pulid_model = PuLIDPipeline(self.model, device,clip_vision_path,weight_dtype=torch.bfloat16,)
if self.offload :
if self.onnx_provider=="gpu":
self.pulid_model.face_helper.face_det.mean_tensor = self.pulid_model.face_helper.face_det.mean_tensor.to(
torch.device("cuda"))
self.pulid_model.face_helper.face_det.device = torch.device("cuda")
self.pulid_model.face_helper.device = torch.device("cuda")
if not self.aggressive_offload:
self.pulid_model.device = torch.device("cuda")
else:
@@ -67,7 +118,10 @@ class FluxGenerator:
start_step,
guidance,
seed,
prompt,
inp,
inp_neg,
x=None,
prompt="none",
id_embeddings = None,
uncond_id_embeddings=None,
id_weight=1.0,
@@ -76,62 +130,63 @@ class FluxGenerator:
timestep_to_start_cfg=1,
max_sequence_length=128,
):
if self.if_repo:
self.t5.max_length = max_sequence_length
# if self.if_repo:
# self.t5.max_length = max_sequence_length
seed = int(seed)
if seed == -1:
seed = None
if isinstance(prompt,str):
text=prompt
elif isinstance(prompt,list):
text=prompt[0]
else:
raise "prompt must be list or str"
# if isinstance(prompt,str):
# text=prompt
# elif isinstance(prompt,list):
# text=prompt[0]
# else:
# raise "prompt must be list or str"
opts = SamplingOptions(
text=text,
width=width,
height=height,
num_steps=num_steps,
guidance=guidance,
seed=seed,
)
# opts = SamplingOptions(
# text=text,
# width=width,
# height=height,
# num_steps=num_steps,
# guidance=guidance,
# seed=seed,
# )
if opts.seed is None:
opts.seed = torch.Generator(device="cpu").seed()
opts.text = opts.text if isinstance(opts.text,str) else str(opts.text)
print(f"Generating '{opts.text}' with seed {opts.seed}")
# if opts.seed is None:
# opts.seed = torch.Generator(device="cpu").seed()
# opts.text = opts.text if isinstance(opts.text,str) else str(opts.text)
# print(f"Generating '{opts.text}' with seed {opts.seed}")
t0 = time.perf_counter()
use_true_cfg = abs(true_cfg - 1.0) > 1e-2
# prepare input
x = get_noise(
1,
opts.height,
opts.width,
device=self.device,
dtype=torch.bfloat16,
seed=opts.seed,
)
# x = get_noise(
# 1,
# height,
# width,
# device=self.device,
# dtype=torch.bfloat16,
# seed=seed,
# )
timesteps = get_schedule(
opts.num_steps,
num_steps,
x.shape[-1] * x.shape[-2] // 4,
shift=True,
)
if self.offload and self.if_repo:
self.t5, self.clip = self.t5.to(self.device), self.clip.to(self.device)
inp = prepare(t5=self.t5, clip=self.clip, img=x, prompt=opts.text,if_repo=self.if_repo)
inp_neg = prepare(t5=self.t5, clip=self.clip, img=x, prompt=neg_prompt,if_repo=self.if_repo) if use_true_cfg else None
# if self.offload and self.if_repo:
# self.t5, self.clip = self.t5.to(self.device), self.clip.to(self.device)
# inp = prepare(t5=self.t5, clip=self.clip, img=x, prompt=opts.text,if_repo=self.if_repo)
# inp_neg = prepare(t5=self.t5, clip=self.clip, img=x, prompt=neg_prompt,if_repo=self.if_repo) if use_true_cfg else None
# offload TEs to CPU, offload pulid_model to cpu, load model to gpu
if self.offload and self.if_repo:
self.t5, self.clip = self.t5.cpu(), self.clip.cpu()
torch.cuda.empty_cache()
# if self.offload and self.if_repo:
# self.t5, self.clip = self.t5.cpu(), self.clip.cpu()
# torch.cuda.empty_cache()
if self.offload:
if self.aggressive_offload:
self.pulid_model.components_to_device(torch.device("cpu"))
@@ -145,7 +200,7 @@ class FluxGenerator:
# denoise initial noise
print("start denoise...")
x = denoise(
self.model, **inp, timesteps=timesteps, guidance=opts.guidance, id=id_embeddings, id_weight=id_weight,
self.model, **inp, timesteps=timesteps, guidance=guidance, id=id_embeddings, id_weight=id_weight,
start_step=start_step, uncond_id=uncond_id_embeddings, true_cfg=true_cfg,
timestep_to_start_cfg=timestep_to_start_cfg,
neg_txt=inp_neg["txt"] if use_true_cfg else None,
@@ -154,16 +209,26 @@ class FluxGenerator:
aggressive_offload=self.aggressive_offload,
)
# offload model, load autoencoder to gpu
print("start decoder...")
if self.offload :
if self.aggressive_offload:
self.model.cpu()
torch.cuda.empty_cache()
self.ae.decoder.to(x.device)
if self.ae is not None:
if self.offload :
if self.aggressive_offload:
self.model.cpu()
torch.cuda.empty_cache()
self.ae.decoder.to(x.device)
# decode latents to pixel space
x = unpack(x.float(), opts.height, opts.width)
x = unpack(x.float(), height, width)
if self.ae is None:
t1 = time.perf_counter()
print(f"Done in {t1 - t0:.1f}s.")
if self.offload :
if self.aggressive_offload:
self.model.cpu()
torch.cuda.empty_cache()
return (x/0.3611)+0.1159
with torch.autocast(device_type=self.device.type, dtype=torch.bfloat16):
x = self.ae.decode(x)
PuLID/flux/sampling.py
+2
View File
@@ -56,6 +56,8 @@ def prepare(t5: HFEmbedder, clip: HFEmbedder, img: Tensor, prompt: str,if_repo:b
if vec.shape[0] == 1 and bs > 1:
vec = repeat(vec, "1 ... -> bs ...", bs=bs)
#print(txt.shape,vec.shape) #torch.Size([1, 256, 4096]) torch.Size([1, 768])
if if_repo:
return {
"img": img,
PuLID/flux/util.py
+57 -20
View File
@@ -9,7 +9,7 @@ from .model import Flux, FluxParams
from .modules.autoencoder import AutoEncoder, AutoEncoderParams
from .modules.conditioner import HFEmbedder
import folder_paths
import logging
@dataclass
class SamplingOptions:
@@ -143,44 +143,81 @@ def print_load_warning(missing: list[str], unexpected: list[str]) -> None:
print(f"Got {len(unexpected)} unexpected keys:\n\t" + "\n\t".join(unexpected))
# from XLabs-AI https://github.com/XLabs-AI/x-flux/blob/1f8ef54972105ad9062be69fe6b7f841bce02a08/src/flux/util.py#L330
def load_flow_model_quintized(name: str, ckpt_path,quantized_mode,use_quantize=True,device: str = "cuda"):
def load_flow_model_quintized(name: str, ckpt_path,quantized_mode,aggressive_offload,device: str = "cuda"):
from optimum.quanto import requantize,qint4,quantize,freeze
# Loading Flux
print(f"Init model in {quantized_mode}")
device=torch.device('cpu')
model = Flux(configs[name].params).to(torch.bfloat16)
sd = load_sft(ckpt_path, device="cpu")
if quantized_mode == "nf4": #nf4 is now useful.
if use_quantize:
from optimum.quanto import requantize
json_path = os.path.join(folder_paths.base_path, "custom_nodes", "ComfyUI_StoryDiffusion", "utils","config.json")
json_path = os.path.join(folder_paths.base_path, "custom_nodes", "ComfyUI_StoryDiffusion", "utils","config.json")
test=False
if test:
from accelerate.utils import set_module_tensor_to_device
from accelerate import init_empty_weights
from .convert_nf4_flux import _replace_with_bnb_linear, create_quantized_param, \
check_quantized_param
import safetensors.torch
import gc
from diffusers import FluxTransformer2DModel
dtype = torch.bfloat16
is_torch_e4m3fn_available = hasattr(torch, "float8_e4m3fn")
original_state_dict = safetensors.torch.load_file(ckpt_path)
print("Start a quantization process...")
with init_empty_weights():
config = FluxTransformer2DModel.load_config(json_path)
model = FluxTransformer2DModel.from_config(config).to(dtype)
expected_state_dict_keys = list(model.state_dict().keys())
_replace_with_bnb_linear(model, "nf4")
for param_name, param in original_state_dict.items():
if param_name not in expected_state_dict_keys:
continue
is_param_float8_e4m3fn = is_torch_e4m3fn_available and param.dtype == torch.float8_e4m3fn
if torch.is_floating_point(param) and not is_param_float8_e4m3fn:
param = param.to(dtype)
if not check_quantized_param(model, param_name):
set_module_tensor_to_device(model, param_name, device=0, value=param)
else:
create_quantized_param(
model, param, param_name, target_device=0, state_dict=original_state_dict,
pre_quantized=True
)
del original_state_dict
gc.collect()
else:
sd = load_sft(ckpt_path, device='cpu')
with open(json_path) as f:
quantization_map = json.load(f)
print(f"Start a {quantized_mode} requantization process...")
print("Start a requantization process...")
requantize(model, sd, quantization_map,device=device)
print("Model is requantized!")
else:
print("normal loading ...")
missing, unexpected = model.load_state_dict(sd, strict=False)
print_load_warning(missing, unexpected)
print("Model is requantized!")
else: #fp8
if use_quantize:
from optimum.quanto import requantize
sd = load_sft(ckpt_path, device="cpu")
if "xlab" in ckpt_path.lower():#"XLabs-AI/flux-dev-fp8" need high env torch 24.1
from comfy import model_detection
from comfy.utils import weight_dtype,load_torch_file
import logging
sd= load_torch_file(ckpt_path)
diffusion_model_prefix = model_detection.unet_prefix_from_state_dict(sd)
weight_dtype = weight_dtype( sd, diffusion_model_prefix)
logging.info("loading weight_dtype is {}".format(weight_dtype))
if weight_dtype is None:
if str(torch.__version__).split(".")[-1]=="4.1" and "xlab" in ckpt_path.lower():#"XLabs-AI/flux-dev-fp8" need high env torch 24.1
json_path = os.path.join(folder_paths.base_path, "custom_nodes", "ComfyUI_StoryDiffusion", "utils","flux_dev_quantization_map.json")
else: # Kijai/flux-fp8,Shakker-Labs/AWPortrait-FL
json_path = os.path.join(folder_paths.base_path, "custom_nodes", "ComfyUI_StoryDiffusion", "utils","config.json") #config is for pass block
with open(json_path,'r') as f:
quantization_map = json.load(f)
print(f"Start a {quantized_mode} requantization process...")
print("Start a requantization process...")
requantize(model, sd, quantization_map, device=device)
print("Model is requantized!")
else: #not useful
torch.cuda.empty_cache()
print("normal loading ...")
sd = load_sft(ckpt_path, device="cpu")
missing, unexpected = model.load_state_dict(sd, strict=False)
print_load_warning(missing, unexpected)
del sd
torch.cuda.empty_cache()
return model
def load_flow_model(name: str,ckpt_path, device: str = "cuda"):
PuLID/pulid/pipeline_flux.py
+104 -22
View File
@@ -112,7 +112,7 @@ def img2tensor(imgs, bgr2rgb=True, float32=True):
class PuLIDPipeline(nn.Module):
def __init__(self, dit, device,clip_vision_path, weight_dtype=torch.bfloat16, onnx_provider='gpu',*args, **kwargs):
def __init__(self, dit, device,clip_vision_path, weight_dtype=torch.bfloat16,*args, **kwargs):
super().__init__()
self.device = device
self.weight_dtype = weight_dtype
@@ -120,7 +120,7 @@ class PuLIDPipeline(nn.Module):
double_interval = 2
single_interval = 4
# init encoder
#init encoder
self.pulid_encoder = IDFormer().to(self.device, self.weight_dtype)
num_ca = 19 // double_interval + 38 // single_interval
@@ -138,18 +138,18 @@ class PuLIDPipeline(nn.Module):
# preprocessors
# face align and parsing
self.face_helper = FaceRestoreHelper(
upscale_factor=1,
face_size=512,
crop_ratio=(1, 1),
det_model='retinaface_resnet50',
save_ext='png',
device=self.device,
)
self.face_helper.face_parse = None
self.face_helper.face_parse = init_parsing_model(model_name='bisenet', device=self.device)
# self.face_helper = FaceRestoreHelper(
# upscale_factor=1,
# face_size=512,
# crop_ratio=(1, 1),
# det_model='retinaface_resnet50',
# save_ext='png',
# device=self.device,
# )
# self.face_helper.face_parse = None
# self.face_helper.face_parse = init_parsing_model(model_name='bisenet', device=self.device)
# clip-vit backbone
model, _, _ = create_model_and_transforms('EVA02-CLIP-L-14-336', clip_vision_path, force_custom_clip=True)
model, _, _ = create_model_and_transforms('EVA02-CLIP-L-14-336', self.clip_vision_path, force_custom_clip=True)
model = model.visual
self.clip_vision_model = model.to(self.device, dtype=self.weight_dtype)
eva_transform_mean = getattr(self.clip_vision_model, 'image_mean', OPENAI_DATASET_MEAN)
@@ -162,14 +162,14 @@ class PuLIDPipeline(nn.Module):
self.eva_transform_std = eva_transform_std
# antelopev2
snapshot_download('DIAMONIK7777/antelopev2', local_dir='models/antelopev2')
# snapshot_download('DIAMONIK7777/antelopev2', local_dir='models/antelopev2')
providers = ['CPUExecutionProvider'] if onnx_provider == 'cpu' else ['CUDAExecutionProvider', 'CPUExecutionProvider']
self.app = FaceAnalysis(name='antelopev2', root='.', providers=providers)
self.app.prepare(ctx_id=0, det_size=(640, 640))
# providers = ['CPUExecutionProvider'] if onnx_provider == 'cpu' else ['CUDAExecutionProvider', 'CPUExecutionProvider']
# self.app = FaceAnalysis(name='antelopev2', root='.', providers=providers)
# self.app.prepare(ctx_id=0, det_size=(640, 640))
self.handler_ante = insightface.model_zoo.get_model('models/antelopev2/glintr100.onnx',providers=providers)
self.handler_ante.prepare(ctx_id=0)
# self.handler_ante = insightface.model_zoo.get_model('models/antelopev2/glintr100.onnx',providers=providers)
# self.handler_ante.prepare(ctx_id=0)
gc.collect()
torch.cuda.empty_cache()
@@ -181,10 +181,10 @@ class PuLIDPipeline(nn.Module):
def components_to_device(self, device):
# everything but pulid_ca
self.face_helper.face_det = self.face_helper.face_det.to(device)
self.face_helper.face_parse = self.face_helper.face_parse.to(device)
# self.face_helper.face_det = self.face_helper.face_det.to(device)
# self.face_helper.face_parse = self.face_helper.face_parse.to(device)
self.clip_vision_model = self.clip_vision_model.to(device)
self.pulid_encoder = self.pulid_encoder.to(device)
# self.pulid_encoder = self.pulid_encoder.to(device)
def load_pretrain(self, pretrain_path=None):
if pretrain_path is not None:
@@ -288,3 +288,85 @@ class PuLIDPipeline(nn.Module):
uncond_id_embedding = self.pulid_encoder(id_uncond, id_vit_hidden_uncond)
return id_embedding, uncond_id_embedding
@torch.no_grad()
def get_id_embedding_(self, image,face_helper,app,handler_ante, cal_uncond=False):
"""
Args:
image: numpy rgb image, range [0, 255]
"""
self.face_helper=face_helper
self.app=app
self.handler_ante=handler_ante
self.face_helper.clean_all()
self.debug_img_list = []
image_bgr = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
# get antelopev2 embedding
face_info = self.app.get(image_bgr)
if len(face_info) > 0:
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
id_ante_embedding = face_info['embedding']
self.debug_img_list.append(
image[
int(face_info['bbox'][1]) : int(face_info['bbox'][3]),
int(face_info['bbox'][0]) : int(face_info['bbox'][2]),
]
)
else:
id_ante_embedding = None
# using facexlib to detect and align face
self.face_helper.read_image(image_bgr)
self.face_helper.get_face_landmarks_5(only_center_face=True)
self.face_helper.align_warp_face()
if len(self.face_helper.cropped_faces) == 0:
raise RuntimeError('facexlib align face fail')
align_face = self.face_helper.cropped_faces[0]
# incase insightface didn't detect face
if id_ante_embedding is None:
print('fail to detect face using insightface, extract embedding on align face')
id_ante_embedding = self.handler_ante.get_feat(align_face)
id_ante_embedding = torch.from_numpy(id_ante_embedding).to(self.device, self.weight_dtype)
if id_ante_embedding.ndim == 1:
id_ante_embedding = id_ante_embedding.unsqueeze(0)
# parsing
input = img2tensor(align_face, bgr2rgb=True).unsqueeze(0) / 255.0
input = input.to(self.device)
parsing_out = self.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, self.to_gray(input))
self.debug_img_list.append(tensor2img(face_features_image, rgb2bgr=False))
# transform img before sending to eva-clip-vit
face_features_image = resize(face_features_image, self.clip_vision_model.image_size, InterpolationMode.BICUBIC)
face_features_image = normalize(face_features_image, self.eva_transform_mean, self.eva_transform_std)
id_cond_vit, id_vit_hidden = self.clip_vision_model(
face_features_image.to(self.weight_dtype), return_all_features=False, return_hidden=True, shuffle=False
)
id_cond_vit_norm = torch.norm(id_cond_vit, 2, 1, True)
id_cond_vit = torch.div(id_cond_vit, id_cond_vit_norm)
id_cond = torch.cat([id_ante_embedding, id_cond_vit], dim=-1)
id_embedding = self.pulid_encoder(id_cond, id_vit_hidden)
if not cal_uncond:
return id_embedding, None
id_uncond = torch.zeros_like(id_cond)
id_vit_hidden_uncond = []
for layer_idx in range(0, len(id_vit_hidden)):
id_vit_hidden_uncond.append(torch.zeros_like(id_vit_hidden[layer_idx]))
uncond_id_embedding = self.pulid_encoder(id_uncond, id_vit_hidden_uncond)
return id_embedding, uncond_id_embedding
README-CN.md
+12 -58
View File
@@ -1,11 +1,12 @@
# ComfyUI_StoryDiffusion
本节点主要方法来源于[StoryDiffusion](https://github.com/HVision-NKU/StoryDiffusion) ，部分内容也来源于[MS-Diffusion](https://github.com/MS-Diffusion/MS-Diffusion),[StoryMaker](https://github.com/RedAIGC/StoryMaker)，[consistory](https://github.com/NVlabs/consistory),[kolor](https://github.com/Kwai-Kolors/Kolors),[pulid](https://github.com/ToTheBeginning/PuLID),[flux](https://github.com/black-forest-labs/flux),[photomaker](https://github.com/TencentARC/PhotoMaker),[IP-Adapter](https://github.com/tencent-ailab/IP-Adapter) [InfiniteYou](https://github.com/bytedance/InfiniteYou) 感谢他们的开源！
本节点主要方法来源于[StoryDiffusion](https://github.com/HVision-NKU/StoryDiffusion) ，部分内容也来源于[MS-Diffusion](https://github.com/MS-Diffusion/MS-Diffusion),[StoryMaker](https://github.com/RedAIGC/StoryMaker)，[consistory](https://github.com/NVlabs/consistory),[kolor](https://github.com/Kwai-Kolors/Kolors),[pulid](https://github.com/ToTheBeginning/PuLID),[flux](https://github.com/black-forest-labs/flux),[photomaker](https://github.com/TencentARC/PhotoMaker),[IP-Adapter](https://github.com/tencent-ailab/IP-Adapter) 感谢他们的开源！
## 更新:
**2025/03/26**
* 加入了gguf和svdquant对InfiniteYou的量化支持，svdquant需要torch>=2.6，且python版本有要求（windows），测试svdquant的量化似乎会导致InfiniteYou失效
* 加入了InfiniteYou的支持，类IP方法人脸特征迁移，用于flux模型，可用lora，所有模型地址在[此](https://huggingface.co/ByteDance/InfiniteYou)，有两种风格的cn模型
**2024/11/15**
* 加入了consistory的支持，你可以在easyfunction输入consi开启此功能(cache,inject2个附属功能，显存大的可以试试)
* consistory模式只支持单主体，当然你也可以使用(cat)或者（boy）或者（hat）,来创造2个主体，比如在角色栏输入 a cure [girl],wearing a (hat).可查看示例图片
* lora功能已加入
## 特色功能
@@ -35,10 +36,7 @@
**Flux and PULID-FLUX**
* flux支持图生图和文生图，量化fp8和nf4运行，推荐用最快的nf4，但是这两模式都需要本地flux diffuse模型并在ckpt name选模型，在填写 'repo_id'格式"X:/xxx/xxx/black-forest-labs/FLUX.1-dev";开启nf4，需要在easy-function输入nf4，如果你想保存自己的量化fp8模型，可以只用repo跑，然后在easy-function输入save，会保存一个fp8量化模型到你的checkpoints目录；
* PULID-FLUX 改成散装式的，clip接comfyUI标准的双clip（1个l 1个T5），ckpt_name选类kj大佬的flux Unet模型（模型名字里必须有flux），clip-vision选pulid的'EVA02_CLIP_L_336_psz14_s6B.pt“模型，vae选择flux标准的 'ae.safetensors'模型，然后在easy-function输入 'pulid, fp8, cpu' 就可以跑，当然你是4090可以试试去掉cpu；
**InfiniteYou**
* 模型加载节点填写本地或者在线flux dev 的repo地址，easyfunction填写infinite,注意不能有大写，前置model连节点easyfunction_lite（插件自带），easyfunction_lite的repo填写InfiniteYou的cn模型地址（见示例及说明）。
* 因为模型显存要求较大，默认对flux模型使用nf4的量化，新增gguf和svd量化加速，12G目前可用，8G未测试
1.安装
-----
@@ -267,43 +265,6 @@ RMBG-1.4 下载至 [link](https://huggingface.co/briaai/RMBG-1.4/tree/main)#自
| ├── w600k_r50.onnx
```
**3.4 使用字节的InfiniteYou**
* 3.4.1 flux-dev 的diffusers 模型地址 [here](https://huggingface.co/black-forest-labs/FLUX.1-dev)
* 3.4.2 infinite 的cn和IP模型地址 from [here ](https://huggingface.co/ByteDance/InfiniteYou) ,sim_stage1 ， aes_stage2两个可选，
```
├── 任意地址/sim_stage1/
| ├── image_proj_model.bin
| ├── InfuseNetModel/
| ├── diffusion_pytorch_model-00001-of-00002.safetensors
| ├── diffusion_pytorch_model-00002-of-00002.safetensors
| ├── diffusion_pytorch_model.safetensors.index.json
| ├── config.json
```
或者
```
├── 任意地址/aes_stage2/
| ├── image_proj_model.bin
| ├── InfuseNetModel/
| ├── diffusion_pytorch_model-00001-of-00002.safetensors
| ├── diffusion_pytorch_model-00002-of-00002.safetensors
| ├── diffusion_pytorch_model.safetensors.index.json
| ├── config.json
```
* 3.4.3 lora 可选 from [here](https://huggingface.co/ByteDance/InfiniteYou)
* 3.4.4 insightface
```
├── ComfyUI/models/antelopev2/
| ├──1k3d68.onnx
| ├──2d106det.onnx
| ├──genderage.onnx
| ├──glintr100.onnx
| ├──scrfd_10g_bnkps.onnx
```
* 3.4.5 recognition_arcface_ir_se50.pth from [here](https://github.com/xinntao/facexlib/releases/download/v0.1.0/recognition_arcface_ir_se50.pth) 自带下载，如果网络不通，手动下载放在控制台提示的地址，一般是C盘XXX，便携包的地址在便携的python目录下的"Lib\site-packages\facexlib\weights"，秋叶类似
* 3.4.6 如果使用gguf量化
下载 gguf 模型 from [here](https://huggingface.co/city96/FLUX.1-dev-gguf/tree/main),在'easyfunction_lite' 节点的'select_method'栏目里填写你下载的模型的绝对地址，反斜杠
* 3.4.7 if use svdquant(optional)
下载 svdquant 的repo文件夹模型 from [here](https://huggingface.co/mit-han-lab/svdq-fp4-flux.1-dev/tree/main) 在'easyfunction_lite' 节点的'select_method'栏目里填写你下载的模型的绝对地址，反斜杠
4 Example
----
@@ -440,6 +401,10 @@ sd1.5
图生图 纯storymaker生成，非最新示例 (outdated version examples)
![](https://github.com/smthemex/ComfyUI_StoryDiffusion/blob/main/examples/maker2role.png)
**flux-pulid**
图生图 flux pulid 12G Vram,cpu Flux使用PULID功能,不需要diffuser模型，非最新示例(outdated version example)
![](https://github.com/smthemex/ComfyUI_StoryDiffusion/blob/main/examples/flux.png)
**kolor-face**
图生图 kolor face，参数输入没变化，非最新示例 (outdated version examples)
![](https://github.com/smthemex/ComfyUI_StoryDiffusion/blob/main/examples/kolor.png)
@@ -449,8 +414,8 @@ sd1.5
![](https://github.com/smthemex/ComfyUI_StoryDiffusion/blob/main/examples/nf4.png)
**ms-diffusion**
* img2img2role in 1 image，双角色同图，下图1为最新示例
![](https://github.com/smthemex/ComfyUI_StoryDiffusion/blob/main/examples/msdiffuion.png)
* img2img2role in 1 image，双角色同图，非最新示例 (Outdated version examples)
![](https://github.com/smthemex/ComfyUI_StoryDiffusion/blob/main/examples/2rolein1img.png)
* ControlNet added dual role co frame (Role 1 and Role 2) (Outdated version examples)
![](https://github.com/smthemex/ComfyUI_StoryDiffusion/blob/main/examples/controlnet.png)
@@ -536,16 +501,5 @@ Consistory
publisher={ACM New York, NY, USA}
}
```
```
@article{jiang2025infiniteyou,
title={{InfiniteYou}: Flexible Photo Recrafting While Preserving Your Identity},
author={Jiang, Liming and Yan, Qing and Jia, Yumin and Liu, Zichuan and Kang, Hao and Lu, Xin},
journal={arXiv preprint},
volume={arXiv:2503.16418},
year={2025}
}
````
FLUX
![LICENSE](https://huggingface.co/black-forest-labs/FLUX.1-dev/blob/main/LICENSE.md)
README.md
+13 -84
View File
@@ -2,13 +2,13 @@
* [中文说明](https://github.com/smthemex/ComfyUI_StoryDiffusion/blob/main/README-CN.md)
* StoryDiffusion origin From: [link](https://github.com/HVision-NKU/StoryDiffusion)
* The project also uses the following open-source projects:[MS-Diffusion](https://github.com/MS-Diffusion/MS-Diffusion),[StoryMaker](https://github.com/RedAIGC/StoryMaker)，[consistory](https://github.com/NVlabs/consistory),[kolor](https://github.com/Kwai-Kolors/Kolors),[pulid](https://github.com/ToTheBeginning/PuLID),[flux](https://github.com/black-forest-labs/flux),[photomaker](https://github.com/TencentARC/PhotoMaker),[IP-Adapter](https://github.com/tencent-ailab/IP-Adapter) [InfiniteYou](https://github.com/bytedance/InfiniteYou)
* The project also uses the following open-source projects:[MS-Diffusion](https://github.com/MS-Diffusion/MS-Diffusion),[StoryMaker](https://github.com/RedAIGC/StoryMaker)，[consistory](https://github.com/NVlabs/consistory),[kolor](https://github.com/Kwai-Kolors/Kolors),[pulid](https://github.com/ToTheBeginning/PuLID),[flux](https://github.com/black-forest-labs/flux),[photomaker](https://github.com/TencentARC/PhotoMaker),[IP-Adapter](https://github.com/tencent-ailab/IP-Adapter)
## Updates:
**2025/03/27**
* fix some bug ,you can link condition_image for control infinite now (4 prompt use 4 condition is better)
* and [gguf](https://github.com/city96/ComfyUI-GGUF) and [SVDQuant](https://github.com/mit-han-lab/nunchaku?tab=readme-ov-file#Installation) quantization model support for 'InfiniteYou',svdquant need torch>=2.6,use wheel to install it.
* add InfiniteYou function from [here](https://github.com/bytedance/InfiniteYou)
**2024/11/15**
* Added support for 'consistory', you can enable this feature by typing 'consi' in easy_function ('cache' and 'inject' are two additional features, you can try with larger VRAM);
* The 'consistory' mode only supports single subjects, but you can also use (cat), (boy), or (hat) to create two subjects, such as entering:' a curve [girl] and wearing a (hat) 'in the character bar，Example images can be viewed；
* NOW,you can use lora when using consistory mode;
## Function introduction
**story-diffusion**
@@ -38,11 +38,7 @@
**Flux and PULID-FLUX**
* Flux supports img2img and txt2img, and supports FP8 and NF4 (recommended) quantization models；To enable it, enter the local path of flux diffuser in 'repo_id' and select the corresponding model in 'ckpt-name';example fill "X:/xxx/xxx/black-forest-labs/FLUX.1-dev";
* PULID-FLUX needs to connect to the dual clip nodes of comfy in clip, and select 'EVA02_CLIP_L_336_psz14_s6B.pt ' in clip-vision, Select a Flux FP8('Kijai/flux-fp8' and "Shakker-Labs/AWPortrait-FL") model (with flux in the name),'ae.safetensors' in vae menu, fill in 'pulid, fp8' in 'easy-function' ;The accompanying 'insightface' model will be automatically downloaded;
**InfiniteYou**
* InfiniteYou only support img2img,download checkpoints from [here](https://huggingface.co/ByteDance/InfiniteYou),need link Easyfunction_lite,and fill in repo
* only support repo mode now,default quant type is nf4 (if use fp16,need VRAM 40G)
* PULID-FLUX needs to connect to the dual clip nodes of comfy in clip, and select 'EVA02_CLIP_L_336_psz14_s6B.pt ' in clip-vision, Select a Flux FP8('Kijai/flux-fp8' and "Shakker-Labs/AWPortrait-FL") model (with flux in the name),'ae.safetensors' in vae menu, fill in 'pulid, fp8' in 'easy-function' ;The accompanying 'insightface' model will be automatically downloaded;
1.Installation
-----
@@ -265,58 +261,8 @@ RMBG-1.4 from [link](https://huggingface.co/briaai/RMBG-1.4/tree/main)#自动
| ├── genderage.onnx
| ├── w600k_r50.onnx
```
**3.6 if using InfiniteYou**
* 3.6.1 flux-dev repo frome [here](https://huggingface.co/black-forest-labs/FLUX.1-dev)
* 3.6.2 infinite controlnet from [here ](https://huggingface.co/ByteDance/InfiniteYou) ,you can use sim_stage1 or aes_stage2
```
├── any_path/sim_stage1/
| ├── image_proj_model.bin
| ├── InfuseNetModel/
| ├── diffusion_pytorch_model-00001-of-00002.safetensors
| ├── diffusion_pytorch_model-00002-of-00002.safetensors
| ├── diffusion_pytorch_model.safetensors.index.json
| ├── config.json
```
or
```
├── any_path/aes_stage2/
| ├── image_proj_model.bin
| ├── InfuseNetModel/
| ├── diffusion_pytorch_model-00001-of-00002.safetensors
| ├── diffusion_pytorch_model-00002-of-00002.safetensors
| ├── diffusion_pytorch_model.safetensors.index.json
| ├── config.json
```
* 3.6.3 lora optional from [here](https://huggingface.co/ByteDance/InfiniteYou)
* 3.6.4 insightface
```
├── ComfyUI/models/antelopev2/
| ├──1k3d68.onnx
| ├──2d106det.onnx
| ├──genderage.onnx
| ├──glintr100.onnx
| ├──scrfd_10g_bnkps.onnx
```
* 3.6.5 recognition_arcface_ir_se50.pth from [here](https://github.com/xinntao/facexlib/releases/download/v0.1.0/recognition_arcface_ir_se50.pth) auto download,which embeded comfyui in "Lib\site-packages\facexlib\weights" dir
* 3.6.6 if use gguf quatization (optional)
download gguf from [here](https://huggingface.co/city96/FLUX.1-dev-gguf/tree/main),and fill local path in 'easyfunction_lite' node's 'select_method'
* 3.6.7 if use svdquant(optional)
download svdquant repo from [here](https://huggingface.co/mit-han-lab/svdq-fp4-flux.1-dev/tree/main) and fill local path in 'easyfunction_lite' node's 'select_method'
4 Example
----
**InfiniteYou**
* 2 role 4 prompt 2 control_image(need fit prompt)
![](https://github.com/smthemex/ComfyUI_StoryDiffusion/blob/main/examples/infinite_control.png)
* normal use nf4 default,fast,quality best
![](https://github.com/smthemex/ComfyUI_StoryDiffusion/blob/main/examples/infinite.png)
* use gguf quant,slowly(Q6,VRAM12G) quality best
![](https://github.com/smthemex/ComfyUI_StoryDiffusion/blob/main/examples/gguf.png)
* use svd quant,fast,quality normal(look like not useful)
![](https://github.com/smthemex/ComfyUI_StoryDiffusion/blob/main/examples/svd.png)
**consistory**
* whtn fill 'consi' in easyfunction is enable.. (Latest version)
![](https://github.com/smthemex/ComfyUI_StoryDiffusion/blob/main/examples/consitstory.png)
@@ -352,6 +298,10 @@ img2img using controlnet and 2roles in 1 img 纯storymaker生成，最新示例
img2img 纯storymaker生成，非最新示例 (outdated version examples)
![](https://github.com/smthemex/ComfyUI_StoryDiffusion/blob/main/examples/maker2role.png)
**flux-pulid**
img2img mode use flux pulid 12G Vram,cpu Flux使用PULID功能,非最新示例(outdated version examples)
![](https://github.com/smthemex/ComfyUI_StoryDiffusion/blob/main/examples/flux.png)
**kolor-face**
img2img kolor face，参数输入没变化，非最新示例 (outdated version examples)
![](https://github.com/smthemex/ComfyUI_StoryDiffusion/blob/main/examples/kolor.png)
@@ -361,8 +311,8 @@ img2img kolor face，参数输入没变化，非最新示例 (outdated version
![](https://github.com/smthemex/ComfyUI_StoryDiffusion/blob/main/examples/nf4.png)
**ms-diffusion**
* img2img2role in 1 image，双角色同图，最新示例 (new version examples)
![](https://github.com/smthemex/ComfyUI_StoryDiffusion/blob/main/examples/msdiffuion.png)
* img2img2role in 1 image，双角色同图，非最新示例 (Outdated version examples)
![](https://github.com/smthemex/ComfyUI_StoryDiffusion/blob/main/examples/2rolein1img.png)
* ControlNet added dual role co frame (Role 1 and Role 2) (Outdated version examples)
![](https://github.com/smthemex/ComfyUI_StoryDiffusion/blob/main/examples/controlnet.png)
@@ -433,9 +383,6 @@ Function Description of Nodes
Previous updates
----
* change some ms-diffusion'codes,if you use "[A] a (man) img " in role prompts,will get better of MS mode.
* Added support for 'consistory', you can enable this feature by typing 'consi' in easy_function ('cache' and 'inject' are two additional features, you can try with larger VRAM);
* The 'consistory' mode only supports single subjects, but you can also use (cat), (boy), or (hat) to create two subjects, such as entering:' a curve [girl] and wearing a (hat) 'in the character bar，Example images can be viewed；you can use lora when using consistory mode;
* if use comfyUI sd3.5 clip and sd 3.5vae( from sd3.5 repo),can load single checkpoint(fp16,nf4 ) which can infer in nf4 mode.（need newest diffusers）
* add sd3.5 large support,can infer in normal or nf4 mode,nf4 mode has two chocie: fill in all local sd3.5 repo(need pip install -U diffusers) or fill local sd3.5 repo and chocie [nf4](https://huggingface.co/sayakpaul/sd35-large-nf4/tree/main/transformer) checkpoint([example](https://github.com/smthemex/ComfyUI_StoryDiffusion/blob/main/examples/sd35nf4singlefile.png)). if use nf4 need fill in 'nf4' in easyfunction.
* add easy_function_lite node,you can use img2tag instead of scene prompts. The current models used are [pzc163/MiniCPMv2_6-prompt-generator](https://huggingface.co/pzc163/MiniCPMv2_6-prompt-generator) and [thwri/CogFlorence-2-Large-Freeze](https://huggingface.co/thwri/CogFlorence-2-Large-Freeze) . Using "thwri/CogFlorence-2-Large-Freeze" requires inputting "flor" in the lite node's easy_function,Temporarily run CUDA during the testing phase.
@@ -549,24 +496,6 @@ Consistory
publisher={ACM New York, NY, USA}
}
```
```
@article{jiang2025infiniteyou,
title={{InfiniteYou}: Flexible Photo Recrafting While Preserving Your Identity},
author={Jiang, Liming and Yan, Qing and Jia, Yumin and Liu, Zichuan and Kang, Hao and Lu, Xin},
journal={arXiv preprint},
volume={arXiv:2503.16418},
year={2025}
}
```
```
@inproceedings{
li2024svdquant,
title={SVDQuant: Absorbing Outliers by Low-Rank Components for 4-Bit Diffusion Models},
author={Li*, Muyang and Lin*, Yujun and Zhang*, Zhekai and Cai, Tianle and Li, Xiuyu and Guo, Junxian and Xie, Enze and Meng, Chenlin and Zhu, Jun-Yan and Han, Song},
booktitle={The Thirteenth International Conference on Learning Representations},
year={2025}
}
```
[GGUF](https://github.com/city96/ComfyUI-GGUF)
FLUX
![LICENSE](https://huggingface.co/black-forest-labs/FLUX.1-dev/blob/main/LICENSE.md)
SD35_Large_IPAdapter/attention_processor.py
+141
View File
@@ -0,0 +1,141 @@
import torch
import torch.nn as nn
import torch.nn.functional as F
from diffusers.models.normalization import FP32LayerNorm, RMSNorm
from typing import Callable, List, Optional, Tuple, Union
import math
import numpy as np
from PIL import Image
class IPAFluxAttnProcessor2_0(nn.Module):
"""Attention processor used typically in processing the SD3-like self-attention projections."""
def __init__(self, hidden_size, cross_attention_dim=None, scale=1.0, num_tokens=4):
super().__init__()
self.hidden_size = hidden_size # 3072
self.cross_attention_dim = cross_attention_dim # 4096
self.scale = scale
self.num_tokens = num_tokens
self.to_k_ip = nn.Linear(cross_attention_dim or hidden_size, hidden_size, bias=False)
self.to_v_ip = nn.Linear(cross_attention_dim or hidden_size, hidden_size, bias=False)
self.norm_added_k = RMSNorm(128, eps=1e-5, elementwise_affine=False)
#self.norm_added_v = RMSNorm(128, eps=1e-5, elementwise_affine=False)
def __call__(
self,
attn,
hidden_states: torch.FloatTensor,
image_emb: torch.FloatTensor,
encoder_hidden_states: torch.FloatTensor = None,
attention_mask: Optional[torch.FloatTensor] = None,
image_rotary_emb: Optional[torch.Tensor] = None,
mask: Optional[torch.Tensor] = None,
) -> torch.FloatTensor:
batch_size, _, _ = hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape
# `sample` projections.
query = attn.to_q(hidden_states)
key = attn.to_k(hidden_states)
value = attn.to_v(hidden_states)
inner_dim = key.shape[-1]
head_dim = inner_dim // attn.heads
query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) # torch.Size([1, 24, 4800, 128])
key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
if attn.norm_q is not None:
query = attn.norm_q(query)
if attn.norm_k is not None:
key = attn.norm_k(key)
if image_emb is not None:
# `ip-adapter` projections
ip_hidden_states = image_emb
ip_hidden_states_key_proj = self.to_k_ip(ip_hidden_states)
ip_hidden_states_value_proj = self.to_v_ip(ip_hidden_states)
ip_hidden_states_key_proj = ip_hidden_states_key_proj.view(
batch_size, -1, attn.heads, head_dim
).transpose(1, 2)
ip_hidden_states_value_proj = ip_hidden_states_value_proj.view(
batch_size, -1, attn.heads, head_dim
).transpose(1, 2)
ip_hidden_states_key_proj = self.norm_added_k(ip_hidden_states_key_proj)
#ip_hidden_states_valye_proj = self.norm_added_v(ip_hidden_states_value_proj)
ip_hidden_states = F.scaled_dot_product_attention(query,
ip_hidden_states_key_proj,
ip_hidden_states_value_proj,
dropout_p=0.0, is_causal=False)
ip_hidden_states = ip_hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)
ip_hidden_states = ip_hidden_states.to(query.dtype)
# the attention in FluxSingleTransformerBlock does not use `encoder_hidden_states`
if encoder_hidden_states is not None:
# `context` projections.
encoder_hidden_states_query_proj = attn.add_q_proj(encoder_hidden_states)
encoder_hidden_states_key_proj = attn.add_k_proj(encoder_hidden_states)
encoder_hidden_states_value_proj = attn.add_v_proj(encoder_hidden_states)
encoder_hidden_states_query_proj = encoder_hidden_states_query_proj.view(
batch_size, -1, attn.heads, head_dim
).transpose(1, 2)
encoder_hidden_states_key_proj = encoder_hidden_states_key_proj.view(
batch_size, -1, attn.heads, head_dim
).transpose(1, 2)
encoder_hidden_states_value_proj = encoder_hidden_states_value_proj.view(
batch_size, -1, attn.heads, head_dim
).transpose(1, 2)
if attn.norm_added_q is not None:
encoder_hidden_states_query_proj = attn.norm_added_q(encoder_hidden_states_query_proj)
if attn.norm_added_k is not None:
encoder_hidden_states_key_proj = attn.norm_added_k(encoder_hidden_states_key_proj)
# attention
query = torch.cat([encoder_hidden_states_query_proj, query], dim=2)
key = torch.cat([encoder_hidden_states_key_proj, key], dim=2)
value = torch.cat([encoder_hidden_states_value_proj, value], dim=2) # (512+3840,128)
if image_rotary_emb is not None:
from diffusers.models.embeddings import apply_rotary_emb
query = apply_rotary_emb(query, image_rotary_emb)
key = apply_rotary_emb(key, image_rotary_emb)
hidden_states = F.scaled_dot_product_attention(query, key, value, dropout_p=0.0, is_causal=False)
hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)
hidden_states = hidden_states.to(query.dtype)
if encoder_hidden_states is not None:
encoder_hidden_states, hidden_states = (
hidden_states[:, : encoder_hidden_states.shape[1]],
hidden_states[:, encoder_hidden_states.shape[1] :],
)
if image_emb is not None:
hidden_states = hidden_states + self.scale * ip_hidden_states
# linear proj
hidden_states = attn.to_out[0](hidden_states)
# dropout
hidden_states = attn.to_out[1](hidden_states)
encoder_hidden_states = attn.to_add_out(encoder_hidden_states)
return hidden_states, encoder_hidden_states
else:
if image_emb is not None:
hidden_states = hidden_states + self.scale * ip_hidden_states
return hidden_states
SD35_Large_IPAdapter/infer_flux_ipa_siglip.py
+190
View File
@@ -0,0 +1,190 @@
import os
import glob
import numpy as np
from PIL import Image
import torch
import torch.nn as nn
from pipeline_flux_ipa import FluxPipeline
from transformer_flux import FluxTransformer2DModel
from attention_processor import IPAFluxAttnProcessor2_0
from transformers import AutoProcessor, SiglipVisionModel
def resize_img(input_image, max_side=1280, min_side=1024, size=None,
pad_to_max_side=False, mode=Image.BILINEAR, base_pixel_number=64):
w, h = input_image.size
if size is not None:
w_resize_new, h_resize_new = size
else:
ratio = min_side / min(h, w)
w, h = round(ratio*w), round(ratio*h)
ratio = max_side / max(h, w)
input_image = input_image.resize([round(ratio*w), round(ratio*h)], mode)
w_resize_new = (round(ratio * w) // base_pixel_number) * base_pixel_number
h_resize_new = (round(ratio * h) // base_pixel_number) * base_pixel_number
input_image = input_image.resize([w_resize_new, h_resize_new], mode)
if pad_to_max_side:
res = np.ones([max_side, max_side, 3], dtype=np.uint8) * 255
offset_x = (max_side - w_resize_new) // 2
offset_y = (max_side - h_resize_new) // 2
res[offset_y:offset_y+h_resize_new, offset_x:offset_x+w_resize_new] = np.array(input_image)
input_image = Image.fromarray(res)
return input_image
class MLPProjModel(torch.nn.Module):
def __init__(self, cross_attention_dim=768, id_embeddings_dim=512, num_tokens=4):
super().__init__()
self.cross_attention_dim = cross_attention_dim
self.num_tokens = num_tokens
self.proj = torch.nn.Sequential(
torch.nn.Linear(id_embeddings_dim, id_embeddings_dim*2),
torch.nn.GELU(),
torch.nn.Linear(id_embeddings_dim*2, cross_attention_dim*num_tokens),
)
self.norm = torch.nn.LayerNorm(cross_attention_dim)
def forward(self, id_embeds):
x = self.proj(id_embeds)
x = x.reshape(-1, self.num_tokens, self.cross_attention_dim)
x = self.norm(x)
return x
class IPAdapter:
def __init__(self, sd_pipe, image_encoder_path, ip_ckpt, device, num_tokens=4):
self.device = device
self.image_encoder_path = image_encoder_path
self.ip_ckpt = ip_ckpt
self.num_tokens = num_tokens
self.pipe = sd_pipe.to(self.device)
self.set_ip_adapter()
# load image encoder
self.image_encoder = SiglipVisionModel.from_pretrained(image_encoder_path).to(self.device, dtype=torch.bfloat16)
self.clip_image_processor = AutoProcessor.from_pretrained(self.image_encoder_path)
# image proj model
self.image_proj_model = self.init_proj()
self.load_ip_adapter()
def init_proj(self):
image_proj_model = MLPProjModel(
cross_attention_dim=self.pipe.transformer.config.joint_attention_dim, # 4096
id_embeddings_dim=1152,
num_tokens=self.num_tokens,
).to(self.device, dtype=torch.bfloat16)
return image_proj_model
def set_ip_adapter(self):
transformer = self.pipe.transformer
ip_attn_procs = {} # 19+38=57
for name in transformer.attn_processors.keys():
if name.startswith("transformer_blocks.") or name.startswith("single_transformer_blocks"):
ip_attn_procs[name] = IPAFluxAttnProcessor2_0(
hidden_size=transformer.config.num_attention_heads * transformer.config.attention_head_dim,
cross_attention_dim=transformer.config.joint_attention_dim,
num_tokens=self.num_tokens,
).to(self.device, dtype=torch.bfloat16)
else:
ip_attn_procs[name] = transformer.attn_processors[name]
transformer.set_attn_processor(ip_attn_procs)
def load_ip_adapter(self):
state_dict = torch.load(self.ip_ckpt, map_location="cpu")
self.image_proj_model.load_state_dict(state_dict["image_proj"], strict=True)
ip_layers = torch.nn.ModuleList(self.pipe.transformer.attn_processors.values())
ip_layers.load_state_dict(state_dict["ip_adapter"], strict=False)
@torch.inference_mode()
def get_image_embeds(self, pil_image=None, clip_image_embeds=None):
if pil_image is not None:
if isinstance(pil_image, Image.Image):
pil_image = [pil_image]
clip_image = self.clip_image_processor(images=pil_image, return_tensors="pt").pixel_values
clip_image_embeds = self.image_encoder(clip_image.to(self.device, dtype=self.image_encoder.dtype)).pooler_output
clip_image_embeds = clip_image_embeds.to(dtype=torch.bfloat16)
else:
clip_image_embeds = clip_image_embeds.to(self.device, dtype=torch.bfloat16)
image_prompt_embeds = self.image_proj_model(clip_image_embeds)
return image_prompt_embeds
def set_scale(self, scale):
for attn_processor in self.pipe.transformer.attn_processors.values():
if isinstance(attn_processor, IPAFluxAttnProcessor2_0):
attn_processor.scale = scale
def generate(
self,
pil_image=None,
clip_image_embeds=None,
prompt=None,
scale=1.0,
num_samples=1,
seed=None,
guidance_scale=3.5,
num_inference_steps=24,
**kwargs,
):
self.set_scale(scale)
image_prompt_embeds = self.get_image_embeds(
pil_image=pil_image, clip_image_embeds=clip_image_embeds
)
if seed is None:
generator = None
else:
generator = torch.Generator(self.device).manual_seed(seed)
images = self.pipe(
prompt=prompt,
image_emb=image_prompt_embeds,
guidance_scale=guidance_scale,
num_inference_steps=num_inference_steps,
generator=generator,
**kwargs,
).images
return images
if __name__ == '__main__':
model_path = "black-forest-labs/FLUX.1-dev"
image_encoder_path = "google/siglip-so400m-patch14-384"
ipadapter_path = "./ip-adapter.bin"
transformer = FluxTransformer2DModel.from_pretrained(
model_path, subfolder="transformer", torch_dtype=torch.bfloat16
)
pipe = FluxPipeline.from_pretrained(
model_path, transformer=transformer, torch_dtype=torch.bfloat16
)
ip_model = IPAdapter(pipe, image_encoder_path, ipadapter_path, device="cuda", num_tokens=128)
image_dir = "./assets/images/2.jpg"
image_name = image_dir.split("/")[-1]
image = Image.open(image_dir).convert("RGB")
image = resize_img(image)
prompt = "a young girl"
images = ip_model.generate(
pil_image=image,
prompt=prompt,
scale=0.7,
width=960, height=1280,
seed=42
)
images[0].save(f"results/{image_name}")
SD35_Large_IPAdapter/models/__init__.py
View File
SD35_Large_IPAdapter/models/attention.py
+1245
View File
File diff suppressed because it is too large Load Diff
SD35_Large_IPAdapter/models/resampler.py
+304
View File
@@ -0,0 +1,304 @@
# modified from https://github.com/mlfoundations/open_flamingo/blob/main/open_flamingo/src/helpers.py
import math
import torch
import torch.nn as nn
from diffusers.models.embeddings import Timesteps, TimestepEmbedding
def get_timestep_embedding(
timesteps: torch.Tensor,
embedding_dim: int,
flip_sin_to_cos: bool = False,
downscale_freq_shift: float = 1,
scale: float = 1,
max_period: int = 10000,
):
"""
This matches the implementation in Denoising Diffusion Probabilistic Models: Create sinusoidal timestep embeddings.
:param timesteps: a 1-D Tensor of N indices, one per batch element.
These may be fractional.
:param embedding_dim: the dimension of the output. :param max_period: controls the minimum frequency of the
embeddings. :return: an [N x dim] Tensor of positional embeddings.
"""
assert len(timesteps.shape) == 1, "Timesteps should be a 1d-array"
half_dim = embedding_dim // 2
exponent = -math.log(max_period) * torch.arange(
start=0, end=half_dim, dtype=torch.float32, device=timesteps.device
)
exponent = exponent / (half_dim - downscale_freq_shift)
emb = torch.exp(exponent)
emb = timesteps[:, None].float() * emb[None, :]
# scale embeddings
emb = scale * emb
# concat sine and cosine embeddings
emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=-1)
# flip sine and cosine embeddings
if flip_sin_to_cos:
emb = torch.cat([emb[:, half_dim:], emb[:, :half_dim]], dim=-1)
# zero pad
if embedding_dim % 2 == 1:
emb = torch.nn.functional.pad(emb, (0, 1, 0, 0))
return emb
# FFN
def FeedForward(dim, mult=4):
inner_dim = int(dim * mult)
return nn.Sequential(
nn.LayerNorm(dim),
nn.Linear(dim, inner_dim, bias=False),
nn.GELU(),
nn.Linear(inner_dim, dim, bias=False),
)
def reshape_tensor(x, heads):
bs, length, width = x.shape
#(bs, length, width) --> (bs, length, n_heads, dim_per_head)
x = x.view(bs, length, heads, -1)
# (bs, length, n_heads, dim_per_head) --> (bs, n_heads, length, dim_per_head)
x = x.transpose(1, 2)
# (bs, n_heads, length, dim_per_head) --> (bs*n_heads, length, dim_per_head)
x = x.reshape(bs, heads, length, -1)
return x
class PerceiverAttention(nn.Module):
def __init__(self, *, dim, dim_head=64, heads=8):
super().__init__()
self.scale = dim_head**-0.5
self.dim_head = dim_head
self.heads = heads
inner_dim = dim_head * heads
self.norm1 = nn.LayerNorm(dim)
self.norm2 = nn.LayerNorm(dim)
self.to_q = nn.Linear(dim, inner_dim, bias=False)
self.to_kv = nn.Linear(dim, inner_dim * 2, bias=False)
self.to_out = nn.Linear(inner_dim, dim, bias=False)
def forward(self, x, latents, shift=None, scale=None):
"""
Args:
x (torch.Tensor): image features
shape (b, n1, D)
latent (torch.Tensor): latent features
shape (b, n2, D)
"""
x = self.norm1(x)
latents = self.norm2(latents)
if shift is not None and scale is not None:
latents = latents * (1 + scale.unsqueeze(1)) + shift.unsqueeze(1)
b, l, _ = latents.shape
q = self.to_q(latents)
kv_input = torch.cat((x, latents), dim=-2)
k, v = self.to_kv(kv_input).chunk(2, dim=-1)
q = reshape_tensor(q, self.heads)
k = reshape_tensor(k, self.heads)
v = reshape_tensor(v, self.heads)
# attention
scale = 1 / math.sqrt(math.sqrt(self.dim_head))
weight = (q * scale) @ (k * scale).transpose(-2, -1) # More stable with f16 than dividing afterwards
weight = torch.softmax(weight.float(), dim=-1).type(weight.dtype)
out = weight @ v
out = out.permute(0, 2, 1, 3).reshape(b, l, -1)
return self.to_out(out)
class Resampler(nn.Module):
def __init__(
self,
dim=1024,
depth=8,
dim_head=64,
heads=16,
num_queries=8,
embedding_dim=768,
output_dim=1024,
ff_mult=4,
*args,
**kwargs,
):
super().__init__()
self.latents = nn.Parameter(torch.randn(1, num_queries, dim) / dim**0.5)
self.proj_in = nn.Linear(embedding_dim, dim)
self.proj_out = nn.Linear(dim, output_dim)
self.norm_out = nn.LayerNorm(output_dim)
self.layers = nn.ModuleList([])
for _ in range(depth):
self.layers.append(
nn.ModuleList(
[
PerceiverAttention(dim=dim, dim_head=dim_head, heads=heads),
FeedForward(dim=dim, mult=ff_mult),
]
)
)
def forward(self, x):
latents = self.latents.repeat(x.size(0), 1, 1)
x = self.proj_in(x)
for attn, ff in self.layers:
latents = attn(x, latents) + latents
latents = ff(latents) + latents
latents = self.proj_out(latents)
return self.norm_out(latents)
class TimeResampler(nn.Module):
def __init__(
self,
dim=1024,
depth=8,
dim_head=64,
heads=16,
num_queries=8,
embedding_dim=768,
output_dim=1024,
ff_mult=4,
timestep_in_dim=320,
timestep_flip_sin_to_cos=True,
timestep_freq_shift=0,
):
super().__init__()
self.latents = nn.Parameter(torch.randn(1, num_queries, dim) / dim**0.5)
self.proj_in = nn.Linear(embedding_dim, dim)
self.proj_out = nn.Linear(dim, output_dim)
self.norm_out = nn.LayerNorm(output_dim)
self.layers = nn.ModuleList([])
for _ in range(depth):
self.layers.append(
nn.ModuleList(
[
# msa
PerceiverAttention(dim=dim, dim_head=dim_head, heads=heads),
# ff
FeedForward(dim=dim, mult=ff_mult),
# adaLN
nn.Sequential(nn.SiLU(), nn.Linear(dim, 4 * dim, bias=True))
]
)
)
# time
self.time_proj = Timesteps(timestep_in_dim, timestep_flip_sin_to_cos, timestep_freq_shift)
self.time_embedding = TimestepEmbedding(timestep_in_dim, dim, act_fn="silu")
# adaLN
# self.adaLN_modulation = nn.Sequential(
# nn.SiLU(),
# nn.Linear(timestep_out_dim, 6 * timestep_out_dim, bias=True)
# )
def forward(self, x, timestep, need_temb=False):
timestep_emb = self.embedding_time(x, timestep) # bs, dim
latents = self.latents.repeat(x.size(0), 1, 1)
x = self.proj_in(x)
x = x + timestep_emb[:, None]
for attn, ff, adaLN_modulation in self.layers:
shift_msa, scale_msa, shift_mlp, scale_mlp = adaLN_modulation(timestep_emb).chunk(4, dim=1)
latents = attn(x, latents, shift_msa, scale_msa) + latents
res = latents
for idx_ff in range(len(ff)):
layer_ff = ff[idx_ff]
latents = layer_ff(latents)
if idx_ff == 0 and isinstance(layer_ff, nn.LayerNorm): # adaLN
latents = latents * (1 + scale_mlp.unsqueeze(1)) + shift_mlp.unsqueeze(1)
latents = latents + res
# latents = ff(latents) + latents
latents = self.proj_out(latents)
latents = self.norm_out(latents)
if need_temb:
return latents, timestep_emb
else:
return latents
def embedding_time(self, sample, timestep):
# 1. time
timesteps = timestep
if not torch.is_tensor(timesteps):
# TODO: this requires sync between CPU and GPU. So try to pass timesteps as tensors if you can
# This would be a good case for the `match` statement (Python 3.10+)
is_mps = sample.device.type == "mps"
if isinstance(timestep, float):
dtype = torch.float32 if is_mps else torch.float64
else:
dtype = torch.int32 if is_mps else torch.int64
timesteps = torch.tensor([timesteps], dtype=dtype, device=sample.device)
elif len(timesteps.shape) == 0:
timesteps = timesteps[None].to(sample.device)
# broadcast to batch dimension in a way that's compatible with ONNX/Core ML
timesteps = timesteps.expand(sample.shape[0])
t_emb = self.time_proj(timesteps)
# timesteps does not contain any weights and will always return f32 tensors
# but time_embedding might actually be running in fp16. so we need to cast here.
# there might be better ways to encapsulate this.
t_emb = t_emb.to(dtype=sample.dtype)
emb = self.time_embedding(t_emb, None)
return emb
if __name__ == '__main__':
model = TimeResampler(
dim=1280,
depth=4,
dim_head=64,
heads=20,
num_queries=16,
embedding_dim=512,
output_dim=2048,
ff_mult=4,
timestep_in_dim=320,
timestep_flip_sin_to_cos=True,
timestep_freq_shift=0,
in_channel_extra_emb=2048,
)
SD35_Large_IPAdapter/models/transformer_sd3.py
+375
View File
@@ -0,0 +1,375 @@
# Copyright 2024 Stability AI, The HuggingFace Team and The InstantX Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from typing import Any, Dict, List, Optional, Tuple, Union
import torch
import torch.nn as nn
from diffusers.configuration_utils import ConfigMixin, register_to_config
from diffusers.loaders import FromOriginalModelMixin, PeftAdapterMixin
from .attention import JointTransformerBlock
from diffusers.models.attention_processor import Attention, AttentionProcessor, FusedJointAttnProcessor2_0
from diffusers.models.modeling_utils import ModelMixin
from diffusers.models.normalization import AdaLayerNormContinuous
from diffusers.utils import USE_PEFT_BACKEND, is_torch_version, logging, scale_lora_layers, unscale_lora_layers
from diffusers.models.embeddings import CombinedTimestepTextProjEmbeddings, PatchEmbed
from diffusers.models.modeling_outputs import Transformer2DModelOutput
logger = logging.get_logger(__name__) # pylint: disable=invalid-name
class SD3Transformer2DModel(ModelMixin, ConfigMixin, PeftAdapterMixin, FromOriginalModelMixin):
"""
The Transformer model introduced in Stable Diffusion 3.
Reference: https://arxiv.org/abs/2403.03206
Parameters:
sample_size (`int`): The width of the latent images. This is fixed during training since
it is used to learn a number of position embeddings.
patch_size (`int`): Patch size to turn the input data into small patches.
in_channels (`int`, *optional*, defaults to 16): The number of channels in the input.
num_layers (`int`, *optional*, defaults to 18): The number of layers of Transformer blocks to use.
attention_head_dim (`int`, *optional*, defaults to 64): The number of channels in each head.
num_attention_heads (`int`, *optional*, defaults to 18): The number of heads to use for multi-head attention.
cross_attention_dim (`int`, *optional*): The number of `encoder_hidden_states` dimensions to use.
caption_projection_dim (`int`): Number of dimensions to use when projecting the `encoder_hidden_states`.
pooled_projection_dim (`int`): Number of dimensions to use when projecting the `pooled_projections`.
out_channels (`int`, defaults to 16): Number of output channels.
"""
_supports_gradient_checkpointing = True
@register_to_config
def __init__(
self,
sample_size: int = 128,
patch_size: int = 2,
in_channels: int = 16,
num_layers: int = 18,
attention_head_dim: int = 64,
num_attention_heads: int = 18,
joint_attention_dim: int = 4096,
caption_projection_dim: int = 1152,
pooled_projection_dim: int = 2048,
out_channels: int = 16,
pos_embed_max_size: int = 96,
dual_attention_layers: Tuple[
int, ...
] = (), # () for sd3.0; (0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12) for sd3.5
qk_norm: Optional[str] = None,
):
super().__init__()
default_out_channels = in_channels
self.out_channels = out_channels if out_channels is not None else default_out_channels
self.inner_dim = self.config.num_attention_heads * self.config.attention_head_dim
self.pos_embed = PatchEmbed(
height=self.config.sample_size,
width=self.config.sample_size,
patch_size=self.config.patch_size,
in_channels=self.config.in_channels,
embed_dim=self.inner_dim,
pos_embed_max_size=pos_embed_max_size, # hard-code for now.
)
self.time_text_embed = CombinedTimestepTextProjEmbeddings(
embedding_dim=self.inner_dim, pooled_projection_dim=self.config.pooled_projection_dim
)
self.context_embedder = nn.Linear(self.config.joint_attention_dim, self.config.caption_projection_dim)
# `attention_head_dim` is doubled to account for the mixing.
# It needs to crafted when we get the actual checkpoints.
self.transformer_blocks = nn.ModuleList(
[
JointTransformerBlock(
dim=self.inner_dim,
num_attention_heads=self.config.num_attention_heads,
attention_head_dim=self.config.attention_head_dim,
context_pre_only=i == num_layers - 1,
qk_norm=qk_norm,
use_dual_attention=True if i in dual_attention_layers else False,
)
for i in range(self.config.num_layers)
]
)
self.norm_out = AdaLayerNormContinuous(self.inner_dim, self.inner_dim, elementwise_affine=False, eps=1e-6)
self.proj_out = nn.Linear(self.inner_dim, patch_size * patch_size * self.out_channels, bias=True)
self.gradient_checkpointing = False
# Copied from diffusers.models.unets.unet_3d_condition.UNet3DConditionModel.enable_forward_chunking
def enable_forward_chunking(self, chunk_size: Optional[int] = None, dim: int = 0) -> None:
"""
Sets the attention processor to use [feed forward
chunking](https://huggingface.co/blog/reformer#2-chunked-feed-forward-layers).
Parameters:
chunk_size (`int`, *optional*):
The chunk size of the feed-forward layers. If not specified, will run feed-forward layer individually
over each tensor of dim=`dim`.
dim (`int`, *optional*, defaults to `0`):
The dimension over which the feed-forward computation should be chunked. Choose between dim=0 (batch)
or dim=1 (sequence length).
"""
if dim not in [0, 1]:
raise ValueError(f"Make sure to set `dim` to either 0 or 1, not {dim}")
# By default chunk size is 1
chunk_size = chunk_size or 1
def fn_recursive_feed_forward(module: torch.nn.Module, chunk_size: int, dim: int):
if hasattr(module, "set_chunk_feed_forward"):
module.set_chunk_feed_forward(chunk_size=chunk_size, dim=dim)
for child in module.children():
fn_recursive_feed_forward(child, chunk_size, dim)
for module in self.children():
fn_recursive_feed_forward(module, chunk_size, dim)
# Copied from diffusers.models.unets.unet_3d_condition.UNet3DConditionModel.disable_forward_chunking
def disable_forward_chunking(self):
def fn_recursive_feed_forward(module: torch.nn.Module, chunk_size: int, dim: int):
if hasattr(module, "set_chunk_feed_forward"):
module.set_chunk_feed_forward(chunk_size=chunk_size, dim=dim)
for child in module.children():
fn_recursive_feed_forward(child, chunk_size, dim)
for module in self.children():
fn_recursive_feed_forward(module, None, 0)
@property
# Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.attn_processors
def attn_processors(self) -> Dict[str, AttentionProcessor]:
r"""
Returns:
`dict` of attention processors: A dictionary containing all attention processors used in the model with
indexed by its weight name.
"""
# set recursively
processors = {}
def fn_recursive_add_processors(name: str, module: torch.nn.Module, processors: Dict[str, AttentionProcessor]):
if hasattr(module, "get_processor"):
processors[f"{name}.processor"] = module.get_processor()
for sub_name, child in module.named_children():
fn_recursive_add_processors(f"{name}.{sub_name}", child, processors)
return processors
for name, module in self.named_children():
fn_recursive_add_processors(name, module, processors)
return processors
# Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.set_attn_processor
def set_attn_processor(self, processor: Union[AttentionProcessor, Dict[str, AttentionProcessor]]):
r"""
Sets the attention processor to use to compute attention.
Parameters:
processor (`dict` of `AttentionProcessor` or only `AttentionProcessor`):
The instantiated processor class or a dictionary of processor classes that will be set as the processor
for **all** `Attention` layers.
If `processor` is a dict, the key needs to define the path to the corresponding cross attention
processor. This is strongly recommended when setting trainable attention processors.
"""
count = len(self.attn_processors.keys())
if isinstance(processor, dict) and len(processor) != count:
raise ValueError(
f"A dict of processors was passed, but the number of processors {len(processor)} does not match the"
f" number of attention layers: {count}. Please make sure to pass {count} processor classes."
)
def fn_recursive_attn_processor(name: str, module: torch.nn.Module, processor):
if hasattr(module, "set_processor"):
if not isinstance(processor, dict):
module.set_processor(processor)
else:
module.set_processor(processor.pop(f"{name}.processor"))
for sub_name, child in module.named_children():
fn_recursive_attn_processor(f"{name}.{sub_name}", child, processor)
for name, module in self.named_children():
fn_recursive_attn_processor(name, module, processor)
# Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.fuse_qkv_projections with FusedAttnProcessor2_0->FusedJointAttnProcessor2_0
def fuse_qkv_projections(self):
"""
Enables fused QKV projections. For self-attention modules, all projection matrices (i.e., query, key, value)
are fused. For cross-attention modules, key and value projection matrices are fused.
<Tip warning={true}>
This API is 🧪 experimental.
</Tip>
"""
self.original_attn_processors = None
for _, attn_processor in self.attn_processors.items():
if "Added" in str(attn_processor.__class__.__name__):
raise ValueError("`fuse_qkv_projections()` is not supported for models having added KV projections.")
self.original_attn_processors = self.attn_processors
for module in self.modules():
if isinstance(module, Attention):
module.fuse_projections(fuse=True)
self.set_attn_processor(FusedJointAttnProcessor2_0())
# Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.unfuse_qkv_projections
def unfuse_qkv_projections(self):
"""Disables the fused QKV projection if enabled.
<Tip warning={true}>
This API is 🧪 experimental.
</Tip>
"""
if self.original_attn_processors is not None:
self.set_attn_processor(self.original_attn_processors)
def _set_gradient_checkpointing(self, module, value=False):
if hasattr(module, "gradient_checkpointing"):
module.gradient_checkpointing = value
def forward(
self,
hidden_states: torch.FloatTensor,
encoder_hidden_states: torch.FloatTensor = None,
pooled_projections: torch.FloatTensor = None,
timestep: torch.LongTensor = None,
block_controlnet_hidden_states: List = None,
joint_attention_kwargs: Optional[Dict[str, Any]] = None,
return_dict: bool = True,
) -> Union[torch.FloatTensor, Transformer2DModelOutput]:
"""
The [`SD3Transformer2DModel`] forward method.
Args:
hidden_states (`torch.FloatTensor` of shape `(batch size, channel, height, width)`):
Input `hidden_states`.
encoder_hidden_states (`torch.FloatTensor` of shape `(batch size, sequence_len, embed_dims)`):
Conditional embeddings (embeddings computed from the input conditions such as prompts) to use.
pooled_projections (`torch.FloatTensor` of shape `(batch_size, projection_dim)`): Embeddings projected
from the embeddings of input conditions.
timestep ( `torch.LongTensor`):
Used to indicate denoising step.
block_controlnet_hidden_states: (`list` of `torch.Tensor`):
A list of tensors that if specified are added to the residuals of transformer blocks.
joint_attention_kwargs (`dict`, *optional*):
A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under
`self.processor` in
[diffusers.models.attention_processor](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py).
return_dict (`bool`, *optional*, defaults to `True`):
Whether or not to return a [`~models.transformer_2d.Transformer2DModelOutput`] instead of a plain
tuple.
Returns:
If `return_dict` is True, an [`~models.transformer_2d.Transformer2DModelOutput`] is returned, otherwise a
`tuple` where the first element is the sample tensor.
"""
if joint_attention_kwargs is not None:
joint_attention_kwargs = joint_attention_kwargs.copy()
lora_scale = joint_attention_kwargs.pop("scale", 1.0)
else:
lora_scale = 1.0
if USE_PEFT_BACKEND:
# weight the lora layers by setting `lora_scale` for each PEFT layer
scale_lora_layers(self, lora_scale)
else:
if joint_attention_kwargs is not None and joint_attention_kwargs.get("scale", None) is not None:
logger.warning(
"Passing `scale` via `joint_attention_kwargs` when not using the PEFT backend is ineffective."
)
height, width = hidden_states.shape[-2:]
hidden_states = self.pos_embed(hidden_states) # takes care of adding positional embeddings too.
temb = self.time_text_embed(timestep, pooled_projections)
encoder_hidden_states = self.context_embedder(encoder_hidden_states)
for index_block, block in enumerate(self.transformer_blocks):
if self.training and self.gradient_checkpointing:
def create_custom_forward(module, return_dict=None):
def custom_forward(*inputs):
if return_dict is not None:
return module(*inputs, return_dict=return_dict)
else:
return module(*inputs)
return custom_forward
ckpt_kwargs: Dict[str, Any] = {"use_reentrant": False} if is_torch_version(">=", "1.11.0") else {}
encoder_hidden_states, hidden_states = torch.utils.checkpoint.checkpoint(
create_custom_forward(block),
hidden_states,
encoder_hidden_states,
temb,
joint_attention_kwargs,
**ckpt_kwargs,
)
else:
encoder_hidden_states, hidden_states = block(
hidden_states=hidden_states, encoder_hidden_states=encoder_hidden_states, temb=temb,
joint_attention_kwargs=joint_attention_kwargs,
)
# controlnet residual
if block_controlnet_hidden_states is not None and block.context_pre_only is False:
interval_control = len(self.transformer_blocks) // len(block_controlnet_hidden_states)
hidden_states = hidden_states + block_controlnet_hidden_states[index_block // interval_control]
hidden_states = self.norm_out(hidden_states, temb)
hidden_states = self.proj_out(hidden_states)
# unpatchify
patch_size = self.config.patch_size
height = height // patch_size
width = width // patch_size
hidden_states = hidden_states.reshape(
shape=(hidden_states.shape[0], height, width, patch_size, patch_size, self.out_channels)
)
hidden_states = torch.einsum("nhwpqc->nchpwq", hidden_states)
output = hidden_states.reshape(
shape=(hidden_states.shape[0], self.out_channels, height * patch_size, width * patch_size)
)
if USE_PEFT_BACKEND:
# remove `lora_scale` from each PEFT layer
unscale_lora_layers(self, lora_scale)
if not return_dict:
return (output,)
return Transformer2DModelOutput(sample=output)
SD35_Large_IPAdapter/pipeline_stable_diffusion_3_ipa.py
+1238
View File
File diff suppressed because it is too large Load Diff
SD35_Large_IPAdapter/transformer_flux.py
+567
View File
@@ -0,0 +1,567 @@
# Copyright 2024 Black Forest Labs, The HuggingFace Team and The InstantX Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from typing import Any, Dict, Optional, Tuple, Union
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from diffusers.configuration_utils import ConfigMixin, register_to_config
from diffusers.loaders import FromOriginalModelMixin, PeftAdapterMixin
from diffusers.models.attention import FeedForward
from diffusers.models.attention_processor import (
Attention,
AttentionProcessor,
FluxAttnProcessor2_0,
FusedFluxAttnProcessor2_0,
)
from diffusers.models.modeling_utils import ModelMixin
from diffusers.models.normalization import AdaLayerNormContinuous, AdaLayerNormZero, AdaLayerNormZeroSingle
from diffusers.utils import USE_PEFT_BACKEND, is_torch_version, logging, scale_lora_layers, unscale_lora_layers
from diffusers.utils.torch_utils import maybe_allow_in_graph
from diffusers.models.embeddings import CombinedTimestepGuidanceTextProjEmbeddings, CombinedTimestepTextProjEmbeddings, FluxPosEmbed
from diffusers.models.modeling_outputs import Transformer2DModelOutput
logger = logging.get_logger(__name__) # pylint: disable=invalid-name
@maybe_allow_in_graph
class FluxSingleTransformerBlock(nn.Module):
r"""
A Transformer block following the MMDiT architecture, introduced in Stable Diffusion 3.
Reference: https://arxiv.org/abs/2403.03206
Parameters:
dim (`int`): The number of channels in the input and output.
num_attention_heads (`int`): The number of heads to use for multi-head attention.
attention_head_dim (`int`): The number of channels in each head.
context_pre_only (`bool`): Boolean to determine if we should add some blocks associated with the
processing of `context` conditions.
"""
def __init__(self, dim, num_attention_heads, attention_head_dim, mlp_ratio=4.0):
super().__init__()
self.mlp_hidden_dim = int(dim * mlp_ratio)
self.norm = AdaLayerNormZeroSingle(dim)
self.proj_mlp = nn.Linear(dim, self.mlp_hidden_dim)
self.act_mlp = nn.GELU(approximate="tanh")
self.proj_out = nn.Linear(dim + self.mlp_hidden_dim, dim)
processor = FluxAttnProcessor2_0()
self.attn = Attention(
query_dim=dim,
cross_attention_dim=None,
dim_head=attention_head_dim,
heads=num_attention_heads,
out_dim=dim,
bias=True,
processor=processor,
qk_norm="rms_norm",
eps=1e-6,
pre_only=True,
)
def forward(
self,
hidden_states: torch.FloatTensor,
temb: torch.FloatTensor,
image_emb=None,
image_rotary_emb=None,
):
residual = hidden_states
norm_hidden_states, gate = self.norm(hidden_states, emb=temb)
mlp_hidden_states = self.act_mlp(self.proj_mlp(norm_hidden_states))
attn_output = self.attn(
hidden_states=norm_hidden_states,
image_rotary_emb=image_rotary_emb,
image_emb=image_emb,
)
hidden_states = torch.cat([attn_output, mlp_hidden_states], dim=2)
gate = gate.unsqueeze(1) # torch.Size([1, 1, 3072])
hidden_states = gate * self.proj_out(hidden_states) # torch.Size([1, 4352, 3072])
hidden_states = residual + hidden_states
if hidden_states.dtype == torch.float16:
hidden_states = hidden_states.clip(-65504, 65504)
return hidden_states
@maybe_allow_in_graph
class FluxTransformerBlock(nn.Module):
r"""
A Transformer block following the MMDiT architecture, introduced in Stable Diffusion 3.
Reference: https://arxiv.org/abs/2403.03206
Parameters:
dim (`int`): The number of channels in the input and output.
num_attention_heads (`int`): The number of heads to use for multi-head attention.
attention_head_dim (`int`): The number of channels in each head.
context_pre_only (`bool`): Boolean to determine if we should add some blocks associated with the
processing of `context` conditions.
"""
def __init__(self, dim, num_attention_heads, attention_head_dim, qk_norm="rms_norm", eps=1e-6):
super().__init__()
self.norm1 = AdaLayerNormZero(dim)
self.norm1_context = AdaLayerNormZero(dim)
if hasattr(F, "scaled_dot_product_attention"):
processor = FluxAttnProcessor2_0()
else:
raise ValueError(
"The current PyTorch version does not support the `scaled_dot_product_attention` function."
)
self.attn = Attention(
query_dim=dim,
cross_attention_dim=None,
added_kv_proj_dim=dim,
dim_head=attention_head_dim,
heads=num_attention_heads,
out_dim=dim,
context_pre_only=False,
bias=True,
processor=processor,
qk_norm=qk_norm,
eps=eps,
)
self.norm2 = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6)
self.ff = FeedForward(dim=dim, dim_out=dim, activation_fn="gelu-approximate")
self.norm2_context = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6)
self.ff_context = FeedForward(dim=dim, dim_out=dim, activation_fn="gelu-approximate")
# let chunk size default to None
self._chunk_size = None
self._chunk_dim = 0
def forward(
self,
hidden_states: torch.FloatTensor,
encoder_hidden_states: torch.FloatTensor,
temb: torch.FloatTensor,
image_emb=None,
image_rotary_emb=None,
):
norm_hidden_states, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.norm1(hidden_states, emb=temb)
norm_encoder_hidden_states, c_gate_msa, c_shift_mlp, c_scale_mlp, c_gate_mlp = self.norm1_context(
encoder_hidden_states, emb=temb
)
# Attention.
attn_output, context_attn_output = self.attn(
hidden_states=norm_hidden_states,
encoder_hidden_states=norm_encoder_hidden_states,
image_rotary_emb=image_rotary_emb,
image_emb=image_emb,
)
# Process attention outputs for the `hidden_states`.
attn_output = gate_msa.unsqueeze(1) * attn_output
hidden_states = hidden_states + attn_output
norm_hidden_states = self.norm2(hidden_states)
norm_hidden_states = norm_hidden_states * (1 + scale_mlp[:, None]) + shift_mlp[:, None]
ff_output = self.ff(norm_hidden_states)
ff_output = gate_mlp.unsqueeze(1) * ff_output
hidden_states = hidden_states + ff_output
# Process attention outputs for the `encoder_hidden_states`.
context_attn_output = c_gate_msa.unsqueeze(1) * context_attn_output
encoder_hidden_states = encoder_hidden_states + context_attn_output
norm_encoder_hidden_states = self.norm2_context(encoder_hidden_states)
norm_encoder_hidden_states = norm_encoder_hidden_states * (1 + c_scale_mlp[:, None]) + c_shift_mlp[:, None]
context_ff_output = self.ff_context(norm_encoder_hidden_states)
encoder_hidden_states = encoder_hidden_states + c_gate_mlp.unsqueeze(1) * context_ff_output
if encoder_hidden_states.dtype == torch.float16:
encoder_hidden_states = encoder_hidden_states.clip(-65504, 65504)
return encoder_hidden_states, hidden_states
class FluxTransformer2DModel(ModelMixin, ConfigMixin, PeftAdapterMixin, FromOriginalModelMixin):
"""
The Transformer model introduced in Flux.
Reference: https://blackforestlabs.ai/announcing-black-forest-labs/
Parameters:
patch_size (`int`): Patch size to turn the input data into small patches.
in_channels (`int`, *optional*, defaults to 16): The number of channels in the input.
num_layers (`int`, *optional*, defaults to 18): The number of layers of MMDiT blocks to use.
num_single_layers (`int`, *optional*, defaults to 18): The number of layers of single DiT blocks to use.
attention_head_dim (`int`, *optional*, defaults to 64): The number of channels in each head.
num_attention_heads (`int`, *optional*, defaults to 18): The number of heads to use for multi-head attention.
joint_attention_dim (`int`, *optional*): The number of `encoder_hidden_states` dimensions to use.
pooled_projection_dim (`int`): Number of dimensions to use when projecting the `pooled_projections`.
guidance_embeds (`bool`, defaults to False): Whether to use guidance embeddings.
"""
_supports_gradient_checkpointing = True
_no_split_modules = ["FluxTransformerBlock", "FluxSingleTransformerBlock"]
@register_to_config
def __init__(
self,
patch_size: int = 1,
in_channels: int = 64,
num_layers: int = 19,
num_single_layers: int = 38,
attention_head_dim: int = 128,
num_attention_heads: int = 24,
joint_attention_dim: int = 4096,
pooled_projection_dim: int = 768,
guidance_embeds: bool = False,
axes_dims_rope: Tuple[int] = (16, 56, 56),
):
super().__init__()
self.out_channels = in_channels
self.inner_dim = self.config.num_attention_heads * self.config.attention_head_dim
self.pos_embed = FluxPosEmbed(theta=10000, axes_dim=axes_dims_rope)
text_time_guidance_cls = (
CombinedTimestepGuidanceTextProjEmbeddings if guidance_embeds else CombinedTimestepTextProjEmbeddings
)
self.time_text_embed = text_time_guidance_cls(
embedding_dim=self.inner_dim, pooled_projection_dim=self.config.pooled_projection_dim
)
self.context_embedder = nn.Linear(self.config.joint_attention_dim, self.inner_dim)
self.x_embedder = torch.nn.Linear(self.config.in_channels, self.inner_dim)
self.transformer_blocks = nn.ModuleList(
[
FluxTransformerBlock(
dim=self.inner_dim,
num_attention_heads=self.config.num_attention_heads,
attention_head_dim=self.config.attention_head_dim,
)
for i in range(self.config.num_layers)
]
)
self.single_transformer_blocks = nn.ModuleList(
[
FluxSingleTransformerBlock(
dim=self.inner_dim,
num_attention_heads=self.config.num_attention_heads,
attention_head_dim=self.config.attention_head_dim,
)
for i in range(self.config.num_single_layers)
]
)
self.norm_out = AdaLayerNormContinuous(self.inner_dim, self.inner_dim, elementwise_affine=False, eps=1e-6)
self.proj_out = nn.Linear(self.inner_dim, patch_size * patch_size * self.out_channels, bias=True)
self.gradient_checkpointing = False
@property
# Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.attn_processors
def attn_processors(self) -> Dict[str, AttentionProcessor]:
r"""
Returns:
`dict` of attention processors: A dictionary containing all attention processors used in the model with
indexed by its weight name.
"""
# set recursively
processors = {}
def fn_recursive_add_processors(name: str, module: torch.nn.Module, processors: Dict[str, AttentionProcessor]):
if hasattr(module, "get_processor"):
processors[f"{name}.processor"] = module.get_processor()
for sub_name, child in module.named_children():
fn_recursive_add_processors(f"{name}.{sub_name}", child, processors)
return processors
for name, module in self.named_children():
fn_recursive_add_processors(name, module, processors)
return processors
# Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.set_attn_processor
def set_attn_processor(self, processor: Union[AttentionProcessor, Dict[str, AttentionProcessor]]):
r"""
Sets the attention processor to use to compute attention.
Parameters:
processor (`dict` of `AttentionProcessor` or only `AttentionProcessor`):
The instantiated processor class or a dictionary of processor classes that will be set as the processor
for **all** `Attention` layers.
If `processor` is a dict, the key needs to define the path to the corresponding cross attention
processor. This is strongly recommended when setting trainable attention processors.
"""
count = len(self.attn_processors.keys())
if isinstance(processor, dict) and len(processor) != count:
raise ValueError(
f"A dict of processors was passed, but the number of processors {len(processor)} does not match the"
f" number of attention layers: {count}. Please make sure to pass {count} processor classes."
)
def fn_recursive_attn_processor(name: str, module: torch.nn.Module, processor):
if hasattr(module, "set_processor"):
if not isinstance(processor, dict):
module.set_processor(processor)
else:
module.set_processor(processor.pop(f"{name}.processor"))
for sub_name, child in module.named_children():
fn_recursive_attn_processor(f"{name}.{sub_name}", child, processor)
for name, module in self.named_children():
fn_recursive_attn_processor(name, module, processor)
# Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.fuse_qkv_projections with FusedAttnProcessor2_0->FusedFluxAttnProcessor2_0
def fuse_qkv_projections(self):
"""
Enables fused QKV projections. For self-attention modules, all projection matrices (i.e., query, key, value)
are fused. For cross-attention modules, key and value projection matrices are fused.
<Tip warning={true}>
This API is 🧪 experimental.
</Tip>
"""
self.original_attn_processors = None
for _, attn_processor in self.attn_processors.items():
if "Added" in str(attn_processor.__class__.__name__):
raise ValueError("`fuse_qkv_projections()` is not supported for models having added KV projections.")
self.original_attn_processors = self.attn_processors
for module in self.modules():
if isinstance(module, Attention):
module.fuse_projections(fuse=True)
self.set_attn_processor(FusedFluxAttnProcessor2_0())
# Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.unfuse_qkv_projections
def unfuse_qkv_projections(self):
"""Disables the fused QKV projection if enabled.
<Tip warning={true}>
This API is 🧪 experimental.
</Tip>
"""
if self.original_attn_processors is not None:
self.set_attn_processor(self.original_attn_processors)
def _set_gradient_checkpointing(self, module, value=False):
if hasattr(module, "gradient_checkpointing"):
module.gradient_checkpointing = value
def forward(
self,
hidden_states: torch.Tensor,
encoder_hidden_states: torch.Tensor = None,
image_emb: torch.FloatTensor = None,
pooled_projections: torch.Tensor = None,
timestep: torch.LongTensor = None,
img_ids: torch.Tensor = None,
txt_ids: torch.Tensor = None,
guidance: torch.Tensor = None,
joint_attention_kwargs: Optional[Dict[str, Any]] = None,
controlnet_block_samples=None,
controlnet_single_block_samples=None,
return_dict: bool = True,
) -> Union[torch.FloatTensor, Transformer2DModelOutput]:
"""
The [`FluxTransformer2DModel`] forward method.
Args:
hidden_states (`torch.FloatTensor` of shape `(batch size, channel, height, width)`):
Input `hidden_states`.
encoder_hidden_states (`torch.FloatTensor` of shape `(batch size, sequence_len, embed_dims)`):
Conditional embeddings (embeddings computed from the input conditions such as prompts) to use.
pooled_projections (`torch.FloatTensor` of shape `(batch_size, projection_dim)`): Embeddings projected
from the embeddings of input conditions.
timestep ( `torch.LongTensor`):
Used to indicate denoising step.
block_controlnet_hidden_states: (`list` of `torch.Tensor`):
A list of tensors that if specified are added to the residuals of transformer blocks.
joint_attention_kwargs (`dict`, *optional*):
A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under
`self.processor` in
[diffusers.models.attention_processor](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py).
return_dict (`bool`, *optional*, defaults to `True`):
Whether or not to return a [`~models.transformer_2d.Transformer2DModelOutput`] instead of a plain
tuple.
Returns:
If `return_dict` is True, an [`~models.transformer_2d.Transformer2DModelOutput`] is returned, otherwise a
`tuple` where the first element is the sample tensor.
"""
if joint_attention_kwargs is not None:
joint_attention_kwargs = joint_attention_kwargs.copy()
lora_scale = joint_attention_kwargs.pop("scale", 1.0)
else:
lora_scale = 1.0
if USE_PEFT_BACKEND:
# weight the lora layers by setting `lora_scale` for each PEFT layer
scale_lora_layers(self, lora_scale)
else:
if joint_attention_kwargs is not None and joint_attention_kwargs.get("scale", None) is not None:
logger.warning(
"Passing `scale` via `joint_attention_kwargs` when not using the PEFT backend is ineffective."
)
hidden_states = self.x_embedder(hidden_states)
timestep = timestep.to(hidden_states.dtype) * 1000
if guidance is not None:
guidance = guidance.to(hidden_states.dtype) * 1000
else:
guidance = None
temb = (
self.time_text_embed(timestep, pooled_projections)
if guidance is None
else self.time_text_embed(timestep, guidance, pooled_projections)
)
# torch.Size([1, 512*num_prompt, 4096]) -> torch.Size([1, 512*num_prompt, 3072])
encoder_hidden_states = self.context_embedder(encoder_hidden_states)
if txt_ids.ndim == 3:
logger.warning(
"Passing `txt_ids` 3d torch.Tensor is deprecated."
"Please remove the batch dimension and pass it as a 2d torch Tensor"
)
txt_ids = txt_ids[0]
if img_ids.ndim == 3:
logger.warning(
"Passing `img_ids` 3d torch.Tensor is deprecated."
"Please remove the batch dimension and pass it as a 2d torch Tensor"
)
img_ids = img_ids[0]
ids = torch.cat((txt_ids, img_ids), dim=0)
image_rotary_emb = self.pos_embed(ids)
for index_block, block in enumerate(self.transformer_blocks):
if self.training and self.gradient_checkpointing:
def create_custom_forward(module, return_dict=None):
def custom_forward(*inputs):
if return_dict is not None:
return module(*inputs, return_dict=return_dict)
else:
return module(*inputs)
return custom_forward
ckpt_kwargs: Dict[str, Any] = {"use_reentrant": False} if is_torch_version(">=", "1.11.0") else {}
encoder_hidden_states, hidden_states = torch.utils.checkpoint.checkpoint(
create_custom_forward(block),
hidden_states,
encoder_hidden_states,
temb,
image_emb,
image_rotary_emb,
**ckpt_kwargs,
)
else:
encoder_hidden_states, hidden_states = block(
hidden_states=hidden_states,
encoder_hidden_states=encoder_hidden_states,
temb=temb,
image_emb=image_emb,
image_rotary_emb=image_rotary_emb,
)
# controlnet residual
if controlnet_block_samples is not None:
interval_control = len(self.transformer_blocks) / len(controlnet_block_samples)
interval_control = int(np.ceil(interval_control))
hidden_states = hidden_states + controlnet_block_samples[index_block // interval_control]
hidden_states = torch.cat([encoder_hidden_states, hidden_states], dim=1)
for index_block, block in enumerate(self.single_transformer_blocks):
if self.training and self.gradient_checkpointing:
def create_custom_forward(module, return_dict=None):
def custom_forward(*inputs):
if return_dict is not None:
return module(*inputs, return_dict=return_dict)
else:
return module(*inputs)
return custom_forward
ckpt_kwargs: Dict[str, Any] = {"use_reentrant": False} if is_torch_version(">=", "1.11.0") else {}
hidden_states = torch.utils.checkpoint.checkpoint(
create_custom_forward(block),
hidden_states,
temb,
image_emb,
image_rotary_emb,
**ckpt_kwargs,
)
else:
hidden_states = block(
hidden_states=hidden_states,
temb=temb,
image_emb=image_emb,
image_rotary_emb=image_rotary_emb,
)
# controlnet residual
if controlnet_single_block_samples is not None:
interval_control = len(self.single_transformer_blocks) / len(controlnet_single_block_samples)
interval_control = int(np.ceil(interval_control))
hidden_states[:, encoder_hidden_states.shape[1] :, ...] = (
hidden_states[:, encoder_hidden_states.shape[1] :, ...]
+ controlnet_single_block_samples[index_block // interval_control]
)
hidden_states = hidden_states[:, encoder_hidden_states.shape[1] :, ...]
hidden_states = self.norm_out(hidden_states, temb)
output = self.proj_out(hidden_states)
if USE_PEFT_BACKEND:
# remove `lora_scale` from each PEFT layer
unscale_lora_layers(self, lora_scale)
if not return_dict:
return (output,)
return Transformer2DModelOutput(sample=output)
StoryMaker/pipeline_sdxl_storymaker_wrapper.py
+1039
View File
File diff suppressed because it is too large Load Diff
Storydiffusion_node.py
+1260 -1981
View File
File diff suppressed because it is too large Load Diff
config/FLUX.1-dev/LICENSE.md
+82
View File
@@ -0,0 +1,82 @@
**FLUX.1 \[dev\] Non-Commercial License**
Black Forest Labs, Inc. (“**we**” or “**our**” or “**Company**”) is pleased to make available the weights, parameters and inference code for the FLUX.1 \[dev\] Model (as defined below) freely available for your non-commercial and non-production use as set forth in this FLUX.1 \[dev\] Non-Commercial License (“**License**”). The “**FLUX.1 \[dev\] Model**” means the FLUX.1 \[dev\] text-to-image AI model and its elements which includes algorithms, software, checkpoints, parameters, source code (inference code, evaluation code, and if applicable, fine-tuning code) and any other materials associated with the FLUX.1 \[dev\] AI model made available by Company under this License, including if any, the technical documentation, manuals and instructions for the use and operation thereof (collectively, “**FLUX.1 \[dev\] Model**”).
By downloading, accessing, use, Distributing (as defined below), or creating a Derivative (as defined below) of the FLUX.1 \[dev\] Model, you agree to the terms of this License. If you do not agree to this License, then you do not have any rights to access, use, Distribute or create a Derivative of the FLUX.1 \[dev\] Model and you must immediately cease using the FLUX.1 \[dev\] Model. If you are agreeing to be bound by the terms of this License on behalf of your employer or other entity, you represent and warrant to us that you have full legal authority to bind your employer or such entity to this License. If you do not have the requisite authority, you may not accept the License or access the FLUX.1 \[dev\] Model on behalf of your employer or other entity.
1. **Definitions**. Capitalized terms used in this License but not defined herein have the following meanings:
1. “**Derivative**” means any (i) modified version of the FLUX.1 \[dev\] Model (including but not limited to any customized or fine-tuned version thereof), (ii) work based on the FLUX.1 \[dev\] Model, or (iii) any other derivative work thereof. For the avoidance of doubt, Outputs are not considered Derivatives under this License.
1. “**Distribution**” or “**Distribute**” or “**Distributing**” means providing or making available, by any means, a copy of the FLUX.1 \[dev\] Models and/or the Derivatives as the case may be.
1. “**Non-Commercial Purpose**” means any of the following uses, but only so far as you do not receive any direct or indirect payment arising from the use of the model or its output: (i) personal use for research, experiment, and testing for the benefit of public knowledge, personal study, private entertainment, hobby projects, or otherwise not directly or indirectly connected to any commercial activities, business operations, or employment responsibilities; (ii) use by commercial or for-profit entities for testing, evaluation, or non-commercial research and development in a non-production environment, (iii) use by any charitable organization for charitable purposes, or for testing or evaluation. For clarity, use for revenue-generating activity or direct interactions with or impacts on end users, or use to train, fine tune or distill other models for commercial use is not a Non-Commercial purpose.
1. “**Outputs**” means any content generated by the operation of the FLUX.1 \[dev\] Models or the Derivatives from a prompt (i.e., text instructions) provided by users. For the avoidance of doubt, Outputs do not include any components of a FLUX.1 \[dev\] Models, such as any fine-tuned versions of the FLUX.1 \[dev\] Models, the weights, or parameters.
1. “**you**” or “**your**” means the individual or entity entering into this License with Company.
1. **License Grant**.
1. License. Subject to your compliance with this License, Company grants you a non-exclusive, worldwide, non-transferable, non-sublicensable, revocable, royalty free and limited license to access, use, create Derivatives of, and Distribute the FLUX.1 \[dev\] Models solely for your Non-Commercial Purposes. The foregoing license is personal to you, and you may not assign or sublicense this License or any other rights or obligations under this License without Company’s prior written consent; any such assignment or sublicense will be void and will automatically and immediately terminate this License. Any restrictions set forth herein in regarding the FLUX.1 \[dev\] Model also applies to any Derivative you create or that are created on your behalf.
1. Non-Commercial Use Only. You may only access, use, Distribute, or creative Derivatives of or the FLUX.1 \[dev\] Model or Derivatives for Non-Commercial Purposes. If You want to use a FLUX.1 \[dev\] Model a Derivative for any purpose that is not expressly authorized under this License, such as for a commercial activity, you must request a license from Company, which Company may grant to you in Company’s sole discretion and which additional use may be subject to a fee, royalty or other revenue share. Please contact Company at the following e-mail address if you want to discuss such a license: info@blackforestlabs.ai.
1. Reserved Rights. The grant of rights expressly set forth in this License are the complete grant of rights to you in the FLUX.1 \[dev\] Model, and no other licenses are granted, whether by waiver, estoppel, implication, equity or otherwise. Company and its licensors reserve all rights not expressly granted by this License.
1. Outputs. We claim no ownership rights in and to the Outputs. You are solely responsible for the Outputs you generate and their subsequent uses in accordance with this License. You may use Output for any purpose (including for commercial purposes), except as expressly prohibited herein. You may not use the Output to train, fine-tune or distill a model that is competitive with the FLUX.1 \[dev\] Model.
1. **Distribution**. Subject to this License, you may Distribute copies of the FLUX.1 \[dev\] Model and/or Derivatives made by you, under the following conditions:
1. you must make available a copy of this License to third-party recipients of the FLUX.1 \[dev\] Models and/or Derivatives you Distribute, and specify that any rights to use the FLUX.1 \[dev\] Models and/or Derivatives shall be directly granted by Company to said third-party recipients pursuant to this License;
1. you must make prominently display the following notice alongside the Distribution of the FLUX.1 \[dev\] Model or Derivative (such as via a “Notice” text file distributed as part of such FLUX.1 \[dev\] Model or Derivative) (the “**Attribution Notice**”):
“The FLUX.1 \[dev\] Model is licensed by Black Forest Labs. Inc. under the FLUX.1 \[dev\] Non-Commercial License. Copyright Black Forest Labs. Inc.
IN NO EVENT SHALL BLACK FOREST LABS, INC. BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH USE OF THIS MODEL.”
1. in the case of Distribution of Derivatives made by you, you must also include in the Attribution Notice a statement that you have modified the applicable FLUX.1 \[dev\] Model; and
1. in the case of Distribution of Derivatives made by you, any terms and conditions you impose on any third-party recipients relating to Derivatives made by or for you shall neither limit such third-party recipients’ use of the FLUX.1 \[dev\] Model or any Derivatives made by or for Company in accordance with this License nor conflict with any of its terms and conditions.
1. In the case of Distribution of Derivatives made by you, you must not misrepresent or imply, through any means, that the Derivatives made by or for you and/or any modified version of the FLUX.1 \[dev\] Model you Distribute under your name and responsibility is an official product of the Company or has been endorsed, approved or validated by the Company, unless you are authorized by Company to do so in writing.
1. **Restrictions**. You will not, and will not permit, assist or cause any third party to
1. use, modify, copy, reproduce, create Derivatives of, or Distribute the FLUX.1 \[dev\] Model (or any Derivative thereof, or any data produced by the FLUX.1 \[dev\] Model), in whole or in part, for (i) any commercial or production purposes, (ii) military purposes, (iii) purposes of surveillance, including any research or development relating to surveillance, (iv) biometric processing, (v) in any manner that infringes, misappropriates, or otherwise violates any third-party rights, or (vi) in any manner that violates any applicable law and violating any privacy or security laws, rules, regulations, directives, or governmental requirements (including the General Data Privacy Regulation (Regulation (EU) 2016/679), the California Consumer Privacy Act, and any and all laws governing the processing of biometric information), as well as all amendments and successor laws to any of the foregoing;
1. alter or remove copyright and other proprietary notices which appear on or in any portion of the FLUX.1 \[dev\] Model;
1. utilize any equipment, device, software, or other means to circumvent or remove any security or protection used by Company in connection with the FLUX.1 \[dev\] Model, or to circumvent or remove any usage restrictions, or to enable functionality disabled by FLUX.1 \[dev\] Model; or
1. offer or impose any terms on the FLUX.1 \[dev\] Model that alter, restrict, or are inconsistent with the terms of this License.
1. violate any applicable U.S. and non-U.S. export control and trade sanctions laws (“**Export Laws**”) in connection with your use or Distribution of any FLUX.1 \[dev\] Model;
1. directly or indirectly Distribute, export, or otherwise transfer FLUX.1 \[dev\] Model (a) to any individual, entity, or country prohibited by Export Laws; (b) to anyone on U.S. or non-U.S. government restricted parties lists; or (c) for any purpose prohibited by Export Laws, including nuclear, chemical or biological weapons, or missile technology applications; 3\) use or download FLUX.1 \[dev\] Model if you or they are (a) located in a comprehensively sanctioned jurisdiction, (b) currently listed on any U.S. or non-U.S. restricted parties list, or (c) for any purpose prohibited by Export Laws; and (4) will not disguise your location through IP proxying or other methods.
1. **DISCLAIMERS**. THE FLUX.1 \[dev\] MODEL IS PROVIDED “AS IS” AND “WITH ALL FAULTS” WITH NO WARRANTY OF ANY KIND, EXPRESS OR IMPLIED. COMPANY EXPRESSLY DISCLAIMS ALL REPRESENTATIONS AND WARRANTIES, EXPRESS OR IMPLIED, WHETHER BY STATUTE, CUSTOM, USAGE OR OTHERWISE AS TO ANY MATTERS RELATED TO THE FLUX.1 \[dev\] MODEL, INCLUDING BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE, TITLE, SATISFACTORY QUALITY, OR NON-INFRINGEMENT. COMPANY MAKES NO WARRANTIES OR REPRESENTATIONS THAT THE FLUX.1 \[dev\] MODEL WILL BE ERROR FREE OR FREE OF VIRUSES OR OTHER HARMFUL COMPONENTS, OR PRODUCE ANY PARTICULAR RESULTS.
1. **LIMITATION OF LIABILITY**. TO THE FULLEST EXTENT PERMITTED BY LAW, IN NO EVENT WILL COMPANY BE LIABLE TO YOU OR YOUR EMPLOYEES, AFFILIATES, USERS, OFFICERS OR DIRECTORS (A) UNDER ANY THEORY OF LIABILITY, WHETHER BASED IN CONTRACT, TORT, NEGLIGENCE, STRICT LIABILITY, WARRANTY, OR OTHERWISE UNDER THIS LICENSE, OR (B) FOR ANY INDIRECT, CONSEQUENTIAL, EXEMPLARY, INCIDENTAL, PUNITIVE OR SPECIAL DAMAGES OR LOST PROFITS, EVEN IF COMPANY HAS BEEN ADVISED OF THE POSSIBILITY OF SUCH DAMAGES. THE FLUX.1 \[dev\] MODEL, ITS CONSTITUENT COMPONENTS, AND ANY OUTPUT (COLLECTIVELY, “MODEL MATERIALS”) ARE NOT DESIGNED OR INTENDED FOR USE IN ANY APPLICATION OR SITUATION WHERE FAILURE OR FAULT OF THE MODEL MATERIALS COULD REASONABLY BE ANTICIPATED TO LEAD TO SERIOUS INJURY OF ANY PERSON, INCLUDING POTENTIAL DISCRIMINATION OR VIOLATION OF AN INDIVIDUAL’S PRIVACY RIGHTS, OR TO SEVERE PHYSICAL, PROPERTY, OR ENVIRONMENTAL DAMAGE (EACH, A “**HIGH-RISK USE**”). IF YOU ELECT TO USE ANY OF THE MODEL MATERIALS FOR A HIGH-RISK USE, YOU DO SO AT YOUR OWN RISK. YOU AGREE TO DESIGN AND IMPLEMENT APPROPRIATE DECISION-MAKING AND RISK-MITIGATION PROCEDURES AND POLICIES IN CONNECTION WITH A HIGH-RISK USE SUCH THAT EVEN IF THERE IS A FAILURE OR FAULT IN ANY OF THE MODEL MATERIALS, THE SAFETY OF PERSONS OR PROPERTY AFFECTED BY THE ACTIVITY STAYS AT A LEVEL THAT IS REASONABLE, APPROPRIATE, AND LAWFUL FOR THE FIELD OF THE HIGH-RISK USE.
1. **INDEMNIFICATION**
You will indemnify, defend and hold harmless Company and our subsidiaries and affiliates, and each of our respective shareholders, directors, officers, employees, agents, successors, and assigns (collectively, the “**Company Parties**”) from and against any losses, liabilities, damages, fines, penalties, and expenses (including reasonable attorneys’ fees) incurred by any Company Party in connection with any claim, demand, allegation, lawsuit, proceeding, or investigation (collectively, “**Claims**”) arising out of or related to (a) your access to or use of the FLUX.1 \[dev\] Model (as well as any Output, results or data generated from such access or use), including any High-Risk Use (defined below); (b) your violation of this License; or (c) your violation, misappropriation or infringement of any rights of another (including intellectual property or other proprietary rights and privacy rights). You will promptly notify the Company Parties of any such Claims, and cooperate with Company Parties in defending such Claims. You will also grant the Company Parties sole control of the defense or settlement, at Company’s sole option, of any Claims. This indemnity is in addition to, and not in lieu of, any other indemnities or remedies set forth in a written agreement between you and Company or the other Company Parties.
1. **Termination; Survival**.
1. This License will automatically terminate upon any breach by you of the terms of this License.
1. We may terminate this License, in whole or in part, at any time upon notice (including electronic) to you.
1. If You initiate any legal action or proceedings against Company or any other entity (including a cross-claim or counterclaim in a lawsuit), alleging that the FLUX.1 \[dev\] Model or any Derivative, or any part thereof, infringe upon intellectual property or other rights owned or licensable by you, then any licenses granted to you under this License will immediately terminate as of the date such legal action or claim is filed or initiated.
1. Upon termination of this License, you must cease all use, access or Distribution of the FLUX.1 \[dev\] Model and any Derivatives. The following sections survive termination of this License 2(c), 2(d), 4-11.
1. **Third Party Materials**. The FLUX.1 \[dev\] Model may contain third-party software or other components (including free and open source software) (all of the foregoing, “**Third Party Materials**”), which are subject to the license terms of the respective third-party licensors. Your dealings or correspondence with third parties and your use of or interaction with any Third Party Materials are solely between you and the third party. Company does not control or endorse, and makes no representations or warranties regarding, any Third Party Materials, and your access to and use of such Third Party Materials are at your own risk.
1. **Trademarks**. You have not been granted any trademark license as part of this License and may not use any name or mark associated with Company without the prior written permission of Company, except to the extent necessary to make the reference required in the Attribution Notice as specified above or as is reasonably necessary in describing the FLUX.1 \[dev\] Model and its creators.
1. **General**. This License will be governed and construed under the laws of the State of Delaware without regard to conflicts of law provisions. If any provision or part of a provision of this License is unlawful, void or unenforceable, that provision or part of the provision is deemed severed from this License, and will not affect the validity and enforceability of any remaining provisions. The failure of Company to exercise or enforce any right or provision of this License will not operate as a waiver of such right or provision. This License does not confer any third-party beneficiary rights upon any other person or entity. This License, together with the Documentation, contains the entire understanding between you and Company regarding the subject matter of this License, and supersedes all other written or oral agreements and understandings between you and Company regarding such subject matter. No change or addition to any provision of this License will be binding unless it is in writing and signed by an authorized representative of both you and Company.
config/FLUX.1-dev/README.md
+93
View File
@@ -0,0 +1,93 @@
---
language:
- en
license: other
license_name: flux-1-dev-non-commercial-license
license_link: LICENSE.md
extra_gated_prompt: By clicking "Agree", you agree to the [FluxDev Non-Commercial License Agreement](https://huggingface.co/black-forest-labs/FLUX.1-dev/blob/main/LICENSE.md)
and acknowledge the [Acceptable Use Policy](https://huggingface.co/black-forest-labs/FLUX.1-dev/blob/main/POLICY.md).
tags:
- text-to-image
- image-generation
- flux
---
![FLUX.1 [dev] Grid](./dev_grid.jpg)
`FLUX.1 [dev]` is a 12 billion parameter rectified flow transformer capable of generating images from text descriptions.
For more information, please read our [blog post](https://blackforestlabs.ai/announcing-black-forest-labs/).
# Key Features
1. Cutting-edge output quality, second only to our state-of-the-art model `FLUX.1 [pro]`.
2. Competitive prompt following, matching the performance of closed source alternatives .
3. Trained using guidance distillation, making `FLUX.1 [dev]` more efficient.
4. Open weights to drive new scientific research, and empower artists to develop innovative workflows.
5. Generated outputs can be used for personal, scientific, and commercial purposes as described in the [flux-1-dev-non-commercial-license](./licence.md).
# Usage
We provide a reference implementation of `FLUX.1 [dev]`, as well as sampling code, in a dedicated [github repository](https://github.com/black-forest-labs/flux).
Developers and creatives looking to build on top of `FLUX.1 [dev]` are encouraged to use this as a starting point.
## API Endpoints
The FLUX.1 models are also available via API from the following sources
1. [bfl.ml](https://docs.bfl.ml/) (currently `FLUX.1 [pro]`)
2. [replicate.com](https://replicate.com/collections/flux)
3. [fal.ai](https://fal.ai/models/fal-ai/flux/dev)
## ComfyUI
`FLUX.1 [dev]` is also available in [Comfy UI](https://github.com/comfyanonymous/ComfyUI) for local inference with a node-based workflow.
## Diffusers
To use `FLUX.1 [dev]` with the 🧨 diffusers python library, first install or upgrade diffusers
```shell
pip install git+https://github.com/huggingface/diffusers.git
```
Then you can use `FluxPipeline` to run the model
```python
import torch
from diffusers import FluxPipeline
pipe = FluxPipeline.from_pretrained("black-forest-labs/FLUX.1-dev", torch_dtype=torch.bfloat16)
pipe.enable_model_cpu_offload() #save some VRAM by offloading the model to CPU. Remove this if you have enough GPU power
prompt = "A cat holding a sign that says hello world"
image = pipe(
prompt,
height=1024,
width=1024,
guidance_scale=3.5,
output_type="pil",
num_inference_steps=50,
max_sequence_length=512,
generator=torch.Generator("cpu").manual_seed(0)
).images[0]
image.save("flux-dev.png")
```
To learn more check out the [diffusers](https://huggingface.co/docs/diffusers/main/en/api/pipelines/flux) documentation
---
# Limitations
- This model is not intended or able to provide factual information.
- As a statistical model this checkpoint might amplify existing societal biases.
- The model may fail to generate output that matches the prompts.
- Prompt following is heavily influenced by the prompting-style.
# Out-of-Scope Use
The model and its derivatives may not be used
- In any way that violates any applicable national, federal, state, local or international law or regulation.
- For the purpose of exploiting, harming or attempting to exploit or harm minors in any way; including but not limited to the solicitation, creation, acquisition, or dissemination of child exploitative content.
- To generate or disseminate verifiably false information and/or content with the purpose of harming others.
- To generate or disseminate personal identifiable information that can be used to harm an individual.
- To harass, abuse, threaten, stalk, or bully individuals or groups of individuals.
- To create non-consensual nudity or illegal pornographic content.
- For fully automated decision making that adversely impacts an individual's legal rights or otherwise creates or modifies a binding, enforceable obligation.
- Generating or facilitating large-scale disinformation campaigns.
# License
This model falls under the [`FLUX.1 [dev]` Non-Commercial License](https://huggingface.co/black-forest-labs/FLUX.1-dev/blob/main/LICENSE.md).
config/FLUX.1-dev/dev_grid.jpg
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 1.2 MiB

config/FLUX.1-dev/model_index.json
+32
View File
@@ -0,0 +1,32 @@
{
"_class_name": "FluxPipeline",
"_diffusers_version": "0.30.0.dev0",
"scheduler": [
"diffusers",
"FlowMatchEulerDiscreteScheduler"
],
"text_encoder": [
"transformers",
"CLIPTextModel"
],
"text_encoder_2": [
"transformers",
"T5EncoderModel"
],
"tokenizer": [
"transformers",
"CLIPTokenizer"
],
"tokenizer_2": [
"transformers",
"T5TokenizerFast"
],
"transformer": [
"diffusers",
"FluxTransformer2DModel"
],
"vae": [
"diffusers",
"AutoencoderKL"
]
}
config/FLUX.1-dev/scheduler/scheduler_config.json
+11
View File
@@ -0,0 +1,11 @@
{
"_class_name": "FlowMatchEulerDiscreteScheduler",
"_diffusers_version": "0.30.0.dev0",
"base_image_seq_len": 256,
"base_shift": 0.5,
"max_image_seq_len": 4096,
"max_shift": 1.15,
"num_train_timesteps": 1000,
"shift": 3.0,
"use_dynamic_shifting": true
}
config/FLUX.1-dev/text_encoder/config.json
+25
View File
@@ -0,0 +1,25 @@
{
"_name_or_path": "openai/clip-vit-large-patch14",
"architectures": [
"CLIPTextModel"
],
"attention_dropout": 0.0,
"bos_token_id": 0,
"dropout": 0.0,
"eos_token_id": 2,
"hidden_act": "quick_gelu",
"hidden_size": 768,
"initializer_factor": 1.0,
"initializer_range": 0.02,
"intermediate_size": 3072,
"layer_norm_eps": 1e-05,
"max_position_embeddings": 77,
"model_type": "clip_text_model",
"num_attention_heads": 12,
"num_hidden_layers": 12,
"pad_token_id": 1,
"projection_dim": 768,
"torch_dtype": "bfloat16",
"transformers_version": "4.43.3",
"vocab_size": 49408
}
config/FLUX.1-dev/text_encoder_2/config.json
+32
View File
@@ -0,0 +1,32 @@
{
"_name_or_path": "google/t5-v1_1-xxl",
"architectures": [
"T5EncoderModel"
],
"classifier_dropout": 0.0,
"d_ff": 10240,
"d_kv": 64,
"d_model": 4096,
"decoder_start_token_id": 0,
"dense_act_fn": "gelu_new",
"dropout_rate": 0.1,
"eos_token_id": 1,
"feed_forward_proj": "gated-gelu",
"initializer_factor": 1.0,
"is_encoder_decoder": true,
"is_gated_act": true,
"layer_norm_epsilon": 1e-06,
"model_type": "t5",
"num_decoder_layers": 24,
"num_heads": 64,
"num_layers": 24,
"output_past": true,
"pad_token_id": 0,
"relative_attention_max_distance": 128,
"relative_attention_num_buckets": 32,
"tie_word_embeddings": false,
"torch_dtype": "bfloat16",
"transformers_version": "4.43.3",
"use_cache": true,
"vocab_size": 32128
}
config/FLUX.1-dev/text_encoder_2/model.safetensors.index.json
+226
View File
@@ -0,0 +1,226 @@
{
"metadata": {
"total_size": 9524621312
},
"weight_map": {
"encoder.block.0.layer.0.SelfAttention.k.weight": "model-00001-of-00002.safetensors",
"encoder.block.0.layer.0.SelfAttention.o.weight": "model-00001-of-00002.safetensors",
"encoder.block.0.layer.0.SelfAttention.q.weight": "model-00001-of-00002.safetensors",
"encoder.block.0.layer.0.SelfAttention.relative_attention_bias.weight": "model-00001-of-00002.safetensors",
"encoder.block.0.layer.0.SelfAttention.v.weight": "model-00001-of-00002.safetensors",
"encoder.block.0.layer.0.layer_norm.weight": "model-00001-of-00002.safetensors",
"encoder.block.0.layer.1.DenseReluDense.wi_0.weight": "model-00001-of-00002.safetensors",
"encoder.block.0.layer.1.DenseReluDense.wi_1.weight": "model-00001-of-00002.safetensors",
"encoder.block.0.layer.1.DenseReluDense.wo.weight": "model-00001-of-00002.safetensors",
"encoder.block.0.layer.1.layer_norm.weight": "model-00001-of-00002.safetensors",
"encoder.block.1.layer.0.SelfAttention.k.weight": "model-00001-of-00002.safetensors",
"encoder.block.1.layer.0.SelfAttention.o.weight": "model-00001-of-00002.safetensors",
"encoder.block.1.layer.0.SelfAttention.q.weight": "model-00001-of-00002.safetensors",
"encoder.block.1.layer.0.SelfAttention.v.weight": "model-00001-of-00002.safetensors",
"encoder.block.1.layer.0.layer_norm.weight": "model-00001-of-00002.safetensors",
"encoder.block.1.layer.1.DenseReluDense.wi_0.weight": "model-00001-of-00002.safetensors",
"encoder.block.1.layer.1.DenseReluDense.wi_1.weight": "model-00001-of-00002.safetensors",
"encoder.block.1.layer.1.DenseReluDense.wo.weight": "model-00001-of-00002.safetensors",
"encoder.block.1.layer.1.layer_norm.weight": "model-00001-of-00002.safetensors",
"encoder.block.10.layer.0.SelfAttention.k.weight": "model-00001-of-00002.safetensors",
"encoder.block.10.layer.0.SelfAttention.o.weight": "model-00001-of-00002.safetensors",
"encoder.block.10.layer.0.SelfAttention.q.weight": "model-00001-of-00002.safetensors",
"encoder.block.10.layer.0.SelfAttention.v.weight": "model-00001-of-00002.safetensors",
"encoder.block.10.layer.0.layer_norm.weight": "model-00001-of-00002.safetensors",
"encoder.block.10.layer.1.DenseReluDense.wi_0.weight": "model-00001-of-00002.safetensors",
"encoder.block.10.layer.1.DenseReluDense.wi_1.weight": "model-00001-of-00002.safetensors",
"encoder.block.10.layer.1.DenseReluDense.wo.weight": "model-00001-of-00002.safetensors",
"encoder.block.10.layer.1.layer_norm.weight": "model-00001-of-00002.safetensors",
"encoder.block.11.layer.0.SelfAttention.k.weight": "model-00001-of-00002.safetensors",
"encoder.block.11.layer.0.SelfAttention.o.weight": "model-00001-of-00002.safetensors",
"encoder.block.11.layer.0.SelfAttention.q.weight": "model-00001-of-00002.safetensors",
"encoder.block.11.layer.0.SelfAttention.v.weight": "model-00001-of-00002.safetensors",
"encoder.block.11.layer.0.layer_norm.weight": "model-00001-of-00002.safetensors",
"encoder.block.11.layer.1.DenseReluDense.wi_0.weight": "model-00001-of-00002.safetensors",
"encoder.block.11.layer.1.DenseReluDense.wi_1.weight": "model-00001-of-00002.safetensors",
"encoder.block.11.layer.1.DenseReluDense.wo.weight": "model-00001-of-00002.safetensors",
"encoder.block.11.layer.1.layer_norm.weight": "model-00001-of-00002.safetensors",
"encoder.block.12.layer.0.SelfAttention.k.weight": "model-00001-of-00002.safetensors",
"encoder.block.12.layer.0.SelfAttention.o.weight": "model-00002-of-00002.safetensors",
"encoder.block.12.layer.0.SelfAttention.q.weight": "model-00001-of-00002.safetensors",
"encoder.block.12.layer.0.SelfAttention.v.weight": "model-00001-of-00002.safetensors",
"encoder.block.12.layer.0.layer_norm.weight": "model-00002-of-00002.safetensors",
"encoder.block.12.layer.1.DenseReluDense.wi_0.weight": "model-00002-of-00002.safetensors",
"encoder.block.12.layer.1.DenseReluDense.wi_1.weight": "model-00002-of-00002.safetensors",
"encoder.block.12.layer.1.DenseReluDense.wo.weight": "model-00002-of-00002.safetensors",
"encoder.block.12.layer.1.layer_norm.weight": "model-00002-of-00002.safetensors",
"encoder.block.13.layer.0.SelfAttention.k.weight": "model-00002-of-00002.safetensors",
"encoder.block.13.layer.0.SelfAttention.o.weight": "model-00002-of-00002.safetensors",
"encoder.block.13.layer.0.SelfAttention.q.weight": "model-00002-of-00002.safetensors",
"encoder.block.13.layer.0.SelfAttention.v.weight": "model-00002-of-00002.safetensors",
"encoder.block.13.layer.0.layer_norm.weight": "model-00002-of-00002.safetensors",
"encoder.block.13.layer.1.DenseReluDense.wi_0.weight": "model-00002-of-00002.safetensors",
"encoder.block.13.layer.1.DenseReluDense.wi_1.weight": "model-00002-of-00002.safetensors",
"encoder.block.13.layer.1.DenseReluDense.wo.weight": "model-00002-of-00002.safetensors",
"encoder.block.13.layer.1.layer_norm.weight": "model-00002-of-00002.safetensors",
"encoder.block.14.layer.0.SelfAttention.k.weight": "model-00002-of-00002.safetensors",
"encoder.block.14.layer.0.SelfAttention.o.weight": "model-00002-of-00002.safetensors",
"encoder.block.14.layer.0.SelfAttention.q.weight": "model-00002-of-00002.safetensors",
"encoder.block.14.layer.0.SelfAttention.v.weight": "model-00002-of-00002.safetensors",
"encoder.block.14.layer.0.layer_norm.weight": "model-00002-of-00002.safetensors",
"encoder.block.14.layer.1.DenseReluDense.wi_0.weight": "model-00002-of-00002.safetensors",
"encoder.block.14.layer.1.DenseReluDense.wi_1.weight": "model-00002-of-00002.safetensors",
"encoder.block.14.layer.1.DenseReluDense.wo.weight": "model-00002-of-00002.safetensors",
"encoder.block.14.layer.1.layer_norm.weight": "model-00002-of-00002.safetensors",
"encoder.block.15.layer.0.SelfAttention.k.weight": "model-00002-of-00002.safetensors",
"encoder.block.15.layer.0.SelfAttention.o.weight": "model-00002-of-00002.safetensors",
"encoder.block.15.layer.0.SelfAttention.q.weight": "model-00002-of-00002.safetensors",
"encoder.block.15.layer.0.SelfAttention.v.weight": "model-00002-of-00002.safetensors",
"encoder.block.15.layer.0.layer_norm.weight": "model-00002-of-00002.safetensors",
"encoder.block.15.layer.1.DenseReluDense.wi_0.weight": "model-00002-of-00002.safetensors",
"encoder.block.15.layer.1.DenseReluDense.wi_1.weight": "model-00002-of-00002.safetensors",
"encoder.block.15.layer.1.DenseReluDense.wo.weight": "model-00002-of-00002.safetensors",
"encoder.block.15.layer.1.layer_norm.weight": "model-00002-of-00002.safetensors",
"encoder.block.16.layer.0.SelfAttention.k.weight": "model-00002-of-00002.safetensors",
"encoder.block.16.layer.0.SelfAttention.o.weight": "model-00002-of-00002.safetensors",
"encoder.block.16.layer.0.SelfAttention.q.weight": "model-00002-of-00002.safetensors",
"encoder.block.16.layer.0.SelfAttention.v.weight": "model-00002-of-00002.safetensors",
"encoder.block.16.layer.0.layer_norm.weight": "model-00002-of-00002.safetensors",
"encoder.block.16.layer.1.DenseReluDense.wi_0.weight": "model-00002-of-00002.safetensors",
"encoder.block.16.layer.1.DenseReluDense.wi_1.weight": "model-00002-of-00002.safetensors",
"encoder.block.16.layer.1.DenseReluDense.wo.weight": "model-00002-of-00002.safetensors",
"encoder.block.16.layer.1.layer_norm.weight": "model-00002-of-00002.safetensors",
"encoder.block.17.layer.0.SelfAttention.k.weight": "model-00002-of-00002.safetensors",
"encoder.block.17.layer.0.SelfAttention.o.weight": "model-00002-of-00002.safetensors",
"encoder.block.17.layer.0.SelfAttention.q.weight": "model-00002-of-00002.safetensors",
"encoder.block.17.layer.0.SelfAttention.v.weight": "model-00002-of-00002.safetensors",
"encoder.block.17.layer.0.layer_norm.weight": "model-00002-of-00002.safetensors",
"encoder.block.17.layer.1.DenseReluDense.wi_0.weight": "model-00002-of-00002.safetensors",
"encoder.block.17.layer.1.DenseReluDense.wi_1.weight": "model-00002-of-00002.safetensors",
"encoder.block.17.layer.1.DenseReluDense.wo.weight": "model-00002-of-00002.safetensors",
"encoder.block.17.layer.1.layer_norm.weight": "model-00002-of-00002.safetensors",
"encoder.block.18.layer.0.SelfAttention.k.weight": "model-00002-of-00002.safetensors",
"encoder.block.18.layer.0.SelfAttention.o.weight": "model-00002-of-00002.safetensors",
"encoder.block.18.layer.0.SelfAttention.q.weight": "model-00002-of-00002.safetensors",
"encoder.block.18.layer.0.SelfAttention.v.weight": "model-00002-of-00002.safetensors",
"encoder.block.18.layer.0.layer_norm.weight": "model-00002-of-00002.safetensors",
"encoder.block.18.layer.1.DenseReluDense.wi_0.weight": "model-00002-of-00002.safetensors",
"encoder.block.18.layer.1.DenseReluDense.wi_1.weight": "model-00002-of-00002.safetensors",
"encoder.block.18.layer.1.DenseReluDense.wo.weight": "model-00002-of-00002.safetensors",
"encoder.block.18.layer.1.layer_norm.weight": "model-00002-of-00002.safetensors",
"encoder.block.19.layer.0.SelfAttention.k.weight": "model-00002-of-00002.safetensors",
"encoder.block.19.layer.0.SelfAttention.o.weight": "model-00002-of-00002.safetensors",
"encoder.block.19.layer.0.SelfAttention.q.weight": "model-00002-of-00002.safetensors",
"encoder.block.19.layer.0.SelfAttention.v.weight": "model-00002-of-00002.safetensors",
"encoder.block.19.layer.0.layer_norm.weight": "model-00002-of-00002.safetensors",
"encoder.block.19.layer.1.DenseReluDense.wi_0.weight": "model-00002-of-00002.safetensors",
"encoder.block.19.layer.1.DenseReluDense.wi_1.weight": "model-00002-of-00002.safetensors",
"encoder.block.19.layer.1.DenseReluDense.wo.weight": "model-00002-of-00002.safetensors",
"encoder.block.19.layer.1.layer_norm.weight": "model-00002-of-00002.safetensors",
"encoder.block.2.layer.0.SelfAttention.k.weight": "model-00001-of-00002.safetensors",
"encoder.block.2.layer.0.SelfAttention.o.weight": "model-00001-of-00002.safetensors",
"encoder.block.2.layer.0.SelfAttention.q.weight": "model-00001-of-00002.safetensors",
"encoder.block.2.layer.0.SelfAttention.v.weight": "model-00001-of-00002.safetensors",
"encoder.block.2.layer.0.layer_norm.weight": "model-00001-of-00002.safetensors",
"encoder.block.2.layer.1.DenseReluDense.wi_0.weight": "model-00001-of-00002.safetensors",
"encoder.block.2.layer.1.DenseReluDense.wi_1.weight": "model-00001-of-00002.safetensors",
"encoder.block.2.layer.1.DenseReluDense.wo.weight": "model-00001-of-00002.safetensors",
"encoder.block.2.layer.1.layer_norm.weight": "model-00001-of-00002.safetensors",
"encoder.block.20.layer.0.SelfAttention.k.weight": "model-00002-of-00002.safetensors",
"encoder.block.20.layer.0.SelfAttention.o.weight": "model-00002-of-00002.safetensors",
"encoder.block.20.layer.0.SelfAttention.q.weight": "model-00002-of-00002.safetensors",
"encoder.block.20.layer.0.SelfAttention.v.weight": "model-00002-of-00002.safetensors",
"encoder.block.20.layer.0.layer_norm.weight": "model-00002-of-00002.safetensors",
"encoder.block.20.layer.1.DenseReluDense.wi_0.weight": "model-00002-of-00002.safetensors",
"encoder.block.20.layer.1.DenseReluDense.wi_1.weight": "model-00002-of-00002.safetensors",
"encoder.block.20.layer.1.DenseReluDense.wo.weight": "model-00002-of-00002.safetensors",
"encoder.block.20.layer.1.layer_norm.weight": "model-00002-of-00002.safetensors",
"encoder.block.21.layer.0.SelfAttention.k.weight": "model-00002-of-00002.safetensors",
"encoder.block.21.layer.0.SelfAttention.o.weight": "model-00002-of-00002.safetensors",
"encoder.block.21.layer.0.SelfAttention.q.weight": "model-00002-of-00002.safetensors",
"encoder.block.21.layer.0.SelfAttention.v.weight": "model-00002-of-00002.safetensors",
"encoder.block.21.layer.0.layer_norm.weight": "model-00002-of-00002.safetensors",
"encoder.block.21.layer.1.DenseReluDense.wi_0.weight": "model-00002-of-00002.safetensors",
"encoder.block.21.layer.1.DenseReluDense.wi_1.weight": "model-00002-of-00002.safetensors",
"encoder.block.21.layer.1.DenseReluDense.wo.weight": "model-00002-of-00002.safetensors",
"encoder.block.21.layer.1.layer_norm.weight": "model-00002-of-00002.safetensors",
"encoder.block.22.layer.0.SelfAttention.k.weight": "model-00002-of-00002.safetensors",
"encoder.block.22.layer.0.SelfAttention.o.weight": "model-00002-of-00002.safetensors",
"encoder.block.22.layer.0.SelfAttention.q.weight": "model-00002-of-00002.safetensors",
"encoder.block.22.layer.0.SelfAttention.v.weight": "model-00002-of-00002.safetensors",
"encoder.block.22.layer.0.layer_norm.weight": "model-00002-of-00002.safetensors",
"encoder.block.22.layer.1.DenseReluDense.wi_0.weight": "model-00002-of-00002.safetensors",
"encoder.block.22.layer.1.DenseReluDense.wi_1.weight": "model-00002-of-00002.safetensors",
"encoder.block.22.layer.1.DenseReluDense.wo.weight": "model-00002-of-00002.safetensors",
"encoder.block.22.layer.1.layer_norm.weight": "model-00002-of-00002.safetensors",
"encoder.block.23.layer.0.SelfAttention.k.weight": "model-00002-of-00002.safetensors",
"encoder.block.23.layer.0.SelfAttention.o.weight": "model-00002-of-00002.safetensors",
"encoder.block.23.layer.0.SelfAttention.q.weight": "model-00002-of-00002.safetensors",
"encoder.block.23.layer.0.SelfAttention.v.weight": "model-00002-of-00002.safetensors",
"encoder.block.23.layer.0.layer_norm.weight": "model-00002-of-00002.safetensors",
"encoder.block.23.layer.1.DenseReluDense.wi_0.weight": "model-00002-of-00002.safetensors",
"encoder.block.23.layer.1.DenseReluDense.wi_1.weight": "model-00002-of-00002.safetensors",
"encoder.block.23.layer.1.DenseReluDense.wo.weight": "model-00002-of-00002.safetensors",
"encoder.block.23.layer.1.layer_norm.weight": "model-00002-of-00002.safetensors",
"encoder.block.3.layer.0.SelfAttention.k.weight": "model-00001-of-00002.safetensors",
"encoder.block.3.layer.0.SelfAttention.o.weight": "model-00001-of-00002.safetensors",
"encoder.block.3.layer.0.SelfAttention.q.weight": "model-00001-of-00002.safetensors",
"encoder.block.3.layer.0.SelfAttention.v.weight": "model-00001-of-00002.safetensors",
"encoder.block.3.layer.0.layer_norm.weight": "model-00001-of-00002.safetensors",
"encoder.block.3.layer.1.DenseReluDense.wi_0.weight": "model-00001-of-00002.safetensors",
"encoder.block.3.layer.1.DenseReluDense.wi_1.weight": "model-00001-of-00002.safetensors",
"encoder.block.3.layer.1.DenseReluDense.wo.weight": "model-00001-of-00002.safetensors",
"encoder.block.3.layer.1.layer_norm.weight": "model-00001-of-00002.safetensors",
"encoder.block.4.layer.0.SelfAttention.k.weight": "model-00001-of-00002.safetensors",
"encoder.block.4.layer.0.SelfAttention.o.weight": "model-00001-of-00002.safetensors",
"encoder.block.4.layer.0.SelfAttention.q.weight": "model-00001-of-00002.safetensors",
"encoder.block.4.layer.0.SelfAttention.v.weight": "model-00001-of-00002.safetensors",
"encoder.block.4.layer.0.layer_norm.weight": "model-00001-of-00002.safetensors",
"encoder.block.4.layer.1.DenseReluDense.wi_0.weight": "model-00001-of-00002.safetensors",
"encoder.block.4.layer.1.DenseReluDense.wi_1.weight": "model-00001-of-00002.safetensors",
"encoder.block.4.layer.1.DenseReluDense.wo.weight": "model-00001-of-00002.safetensors",
"encoder.block.4.layer.1.layer_norm.weight": "model-00001-of-00002.safetensors",
"encoder.block.5.layer.0.SelfAttention.k.weight": "model-00001-of-00002.safetensors",
"encoder.block.5.layer.0.SelfAttention.o.weight": "model-00001-of-00002.safetensors",
"encoder.block.5.layer.0.SelfAttention.q.weight": "model-00001-of-00002.safetensors",
"encoder.block.5.layer.0.SelfAttention.v.weight": "model-00001-of-00002.safetensors",
"encoder.block.5.layer.0.layer_norm.weight": "model-00001-of-00002.safetensors",
"encoder.block.5.layer.1.DenseReluDense.wi_0.weight": "model-00001-of-00002.safetensors",
"encoder.block.5.layer.1.DenseReluDense.wi_1.weight": "model-00001-of-00002.safetensors",
"encoder.block.5.layer.1.DenseReluDense.wo.weight": "model-00001-of-00002.safetensors",
"encoder.block.5.layer.1.layer_norm.weight": "model-00001-of-00002.safetensors",
"encoder.block.6.layer.0.SelfAttention.k.weight": "model-00001-of-00002.safetensors",
"encoder.block.6.layer.0.SelfAttention.o.weight": "model-00001-of-00002.safetensors",
"encoder.block.6.layer.0.SelfAttention.q.weight": "model-00001-of-00002.safetensors",
"encoder.block.6.layer.0.SelfAttention.v.weight": "model-00001-of-00002.safetensors",
"encoder.block.6.layer.0.layer_norm.weight": "model-00001-of-00002.safetensors",
"encoder.block.6.layer.1.DenseReluDense.wi_0.weight": "model-00001-of-00002.safetensors",
"encoder.block.6.layer.1.DenseReluDense.wi_1.weight": "model-00001-of-00002.safetensors",
"encoder.block.6.layer.1.DenseReluDense.wo.weight": "model-00001-of-00002.safetensors",
"encoder.block.6.layer.1.layer_norm.weight": "model-00001-of-00002.safetensors",
"encoder.block.7.layer.0.SelfAttention.k.weight": "model-00001-of-00002.safetensors",
"encoder.block.7.layer.0.SelfAttention.o.weight": "model-00001-of-00002.safetensors",
"encoder.block.7.layer.0.SelfAttention.q.weight": "model-00001-of-00002.safetensors",
"encoder.block.7.layer.0.SelfAttention.v.weight": "model-00001-of-00002.safetensors",
"encoder.block.7.layer.0.layer_norm.weight": "model-00001-of-00002.safetensors",
"encoder.block.7.layer.1.DenseReluDense.wi_0.weight": "model-00001-of-00002.safetensors",
"encoder.block.7.layer.1.DenseReluDense.wi_1.weight": "model-00001-of-00002.safetensors",
"encoder.block.7.layer.1.DenseReluDense.wo.weight": "model-00001-of-00002.safetensors",
"encoder.block.7.layer.1.layer_norm.weight": "model-00001-of-00002.safetensors",
"encoder.block.8.layer.0.SelfAttention.k.weight": "model-00001-of-00002.safetensors",
"encoder.block.8.layer.0.SelfAttention.o.weight": "model-00001-of-00002.safetensors",
"encoder.block.8.layer.0.SelfAttention.q.weight": "model-00001-of-00002.safetensors",
"encoder.block.8.layer.0.SelfAttention.v.weight": "model-00001-of-00002.safetensors",
"encoder.block.8.layer.0.layer_norm.weight": "model-00001-of-00002.safetensors",
"encoder.block.8.layer.1.DenseReluDense.wi_0.weight": "model-00001-of-00002.safetensors",
"encoder.block.8.layer.1.DenseReluDense.wi_1.weight": "model-00001-of-00002.safetensors",
"encoder.block.8.layer.1.DenseReluDense.wo.weight": "model-00001-of-00002.safetensors",
"encoder.block.8.layer.1.layer_norm.weight": "model-00001-of-00002.safetensors",
"encoder.block.9.layer.0.SelfAttention.k.weight": "model-00001-of-00002.safetensors",
"encoder.block.9.layer.0.SelfAttention.o.weight": "model-00001-of-00002.safetensors",
"encoder.block.9.layer.0.SelfAttention.q.weight": "model-00001-of-00002.safetensors",
"encoder.block.9.layer.0.SelfAttention.v.weight": "model-00001-of-00002.safetensors",
"encoder.block.9.layer.0.layer_norm.weight": "model-00001-of-00002.safetensors",
"encoder.block.9.layer.1.DenseReluDense.wi_0.weight": "model-00001-of-00002.safetensors",
"encoder.block.9.layer.1.DenseReluDense.wi_1.weight": "model-00001-of-00002.safetensors",
"encoder.block.9.layer.1.DenseReluDense.wo.weight": "model-00001-of-00002.safetensors",
"encoder.block.9.layer.1.layer_norm.weight": "model-00001-of-00002.safetensors",
"encoder.final_layer_norm.weight": "model-00002-of-00002.safetensors",
"shared.weight": "model-00001-of-00002.safetensors"
}
}
config/FLUX.1-dev/tokenizer/merges.txt
+48895
View File
File diff suppressed because it is too large Load Diff
config/FLUX.1-dev/tokenizer/special_tokens_map.json
+30
View File
@@ -0,0 +1,30 @@
{
"bos_token": {
"content": "<|startoftext|>",
"lstrip": false,
"normalized": true,
"rstrip": false,
"single_word": false
},
"eos_token": {
"content": "<|endoftext|>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false
},
"pad_token": {
"content": "<|endoftext|>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false
},
"unk_token": {
"content": "<|endoftext|>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false
}
}
config/FLUX.1-dev/tokenizer/tokenizer_config.json
+30
View File
@@ -0,0 +1,30 @@
{
"add_prefix_space": false,
"added_tokens_decoder": {
"49406": {
"content": "<|startoftext|>",
"lstrip": false,
"normalized": true,
"rstrip": false,
"single_word": false,
"special": true
},
"49407": {
"content": "<|endoftext|>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": true
}
},
"bos_token": "<|startoftext|>",
"clean_up_tokenization_spaces": true,
"do_lower_case": true,
"eos_token": "<|endoftext|>",
"errors": "replace",
"model_max_length": 77,
"pad_token": "<|endoftext|>",
"tokenizer_class": "CLIPTokenizer",
"unk_token": "<|endoftext|>"
}
config/FLUX.1-dev/tokenizer/vocab.json
+49410
View File
File diff suppressed because it is too large Load Diff
config/FLUX.1-dev/tokenizer_2/special_tokens_map.json
+125
View File
@@ -0,0 +1,125 @@
{
"additional_special_tokens": [
"<extra_id_0>",
"<extra_id_1>",
"<extra_id_2>",
"<extra_id_3>",
"<extra_id_4>",
"<extra_id_5>",
"<extra_id_6>",
"<extra_id_7>",
"<extra_id_8>",
"<extra_id_9>",
"<extra_id_10>",
"<extra_id_11>",
"<extra_id_12>",
"<extra_id_13>",
"<extra_id_14>",
"<extra_id_15>",
"<extra_id_16>",
"<extra_id_17>",
"<extra_id_18>",
"<extra_id_19>",
"<extra_id_20>",
"<extra_id_21>",
"<extra_id_22>",
"<extra_id_23>",
"<extra_id_24>",
"<extra_id_25>",
"<extra_id_26>",
"<extra_id_27>",
"<extra_id_28>",
"<extra_id_29>",
"<extra_id_30>",
"<extra_id_31>",
"<extra_id_32>",
"<extra_id_33>",
"<extra_id_34>",
"<extra_id_35>",
"<extra_id_36>",
"<extra_id_37>",
"<extra_id_38>",
"<extra_id_39>",
"<extra_id_40>",
"<extra_id_41>",
"<extra_id_42>",
"<extra_id_43>",
"<extra_id_44>",
"<extra_id_45>",
"<extra_id_46>",
"<extra_id_47>",
"<extra_id_48>",
"<extra_id_49>",
"<extra_id_50>",
"<extra_id_51>",
"<extra_id_52>",
"<extra_id_53>",
"<extra_id_54>",
"<extra_id_55>",
"<extra_id_56>",
"<extra_id_57>",
"<extra_id_58>",
"<extra_id_59>",
"<extra_id_60>",
"<extra_id_61>",
"<extra_id_62>",
"<extra_id_63>",
"<extra_id_64>",
"<extra_id_65>",
"<extra_id_66>",
"<extra_id_67>",
"<extra_id_68>",
"<extra_id_69>",
"<extra_id_70>",
"<extra_id_71>",
"<extra_id_72>",
"<extra_id_73>",
"<extra_id_74>",
"<extra_id_75>",
"<extra_id_76>",
"<extra_id_77>",
"<extra_id_78>",
"<extra_id_79>",
"<extra_id_80>",
"<extra_id_81>",
"<extra_id_82>",
"<extra_id_83>",
"<extra_id_84>",
"<extra_id_85>",
"<extra_id_86>",
"<extra_id_87>",
"<extra_id_88>",
"<extra_id_89>",
"<extra_id_90>",
"<extra_id_91>",
"<extra_id_92>",
"<extra_id_93>",
"<extra_id_94>",
"<extra_id_95>",
"<extra_id_96>",
"<extra_id_97>",
"<extra_id_98>",
"<extra_id_99>"
],
"eos_token": {
"content": "</s>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false
},
"pad_token": {
"content": "<pad>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false
},
"unk_token": {
"content": "<unk>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false
}
}
config/FLUX.1-dev/tokenizer_2/spiece.model
BIN
View File
Binary file not shown.
config/FLUX.1-dev/tokenizer_2/tokenizer.json
+129428
View File
File diff suppressed because one or more lines are too long
config/FLUX.1-dev/tokenizer_2/tokenizer_config.json
+940
View File
@@ -0,0 +1,940 @@
{
"add_prefix_space": true,
"added_tokens_decoder": {
"0": {
"content": "<pad>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": true
},
"1": {
"content": "</s>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": true
},
"2": {
"content": "<unk>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": true
},
"32000": {
"content": "<extra_id_99>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": true
},
"32001": {
"content": "<extra_id_98>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": true
},
"32002": {
"content": "<extra_id_97>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": true
},
"32003": {
"content": "<extra_id_96>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": true
},
"32004": {
"content": "<extra_id_95>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": true
},
"32005": {
"content": "<extra_id_94>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": true
},
"32006": {
"content": "<extra_id_93>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": true
},
"32007": {
"content": "<extra_id_92>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": true
},
"32008": {
"content": "<extra_id_91>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": true
},
"32009": {
"content": "<extra_id_90>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": true
},
"32010": {
"content": "<extra_id_89>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": true
},
"32011": {
"content": "<extra_id_88>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": true
},
"32012": {
"content": "<extra_id_87>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": true
},
"32013": {
"content": "<extra_id_86>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": true
},
"32014": {
"content": "<extra_id_85>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": true
},
"32015": {
"content": "<extra_id_84>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": true
},
"32016": {
"content": "<extra_id_83>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": true
},
"32017": {
"content": "<extra_id_82>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": true
},
"32018": {
"content": "<extra_id_81>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": true
},
"32019": {
"content": "<extra_id_80>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": true
},
"32020": {
"content": "<extra_id_79>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": true
},
"32021": {
"content": "<extra_id_78>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": true
},
"32022": {
"content": "<extra_id_77>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": true
},
"32023": {
"content": "<extra_id_76>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": true
},
"32024": {
"content": "<extra_id_75>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": true
},
"32025": {
"content": "<extra_id_74>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": true
},
"32026": {
"content": "<extra_id_73>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": true
},
"32027": {
"content": "<extra_id_72>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": true
},
"32028": {
"content": "<extra_id_71>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": true
},
"32029": {
"content": "<extra_id_70>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": true
},
"32030": {
"content": "<extra_id_69>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": true
},
"32031": {
"content": "<extra_id_68>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": true
},
"32032": {
"content": "<extra_id_67>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": true
},
"32033": {
"content": "<extra_id_66>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": true
},
"32034": {
"content": "<extra_id_65>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": true
},
"32035": {
"content": "<extra_id_64>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": true
},
"32036": {
"content": "<extra_id_63>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": true
},
"32037": {
"content": "<extra_id_62>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": true
},
"32038": {
"content": "<extra_id_61>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": true
},
"32039": {
"content": "<extra_id_60>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": true
},
"32040": {
"content": "<extra_id_59>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": true
},
"32041": {
"content": "<extra_id_58>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": true
},
"32042": {
"content": "<extra_id_57>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": true
},
"32043": {
"content": "<extra_id_56>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": true
},
"32044": {
"content": "<extra_id_55>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": true
},
"32045": {
"content": "<extra_id_54>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": true
},
"32046": {
"content": "<extra_id_53>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": true
},
"32047": {
"content": "<extra_id_52>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": true
},
"32048": {
"content": "<extra_id_51>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": true
},
"32049": {
"content": "<extra_id_50>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": true
},
"32050": {
"content": "<extra_id_49>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": true
},
"32051": {
"content": "<extra_id_48>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": true
},
"32052": {
"content": "<extra_id_47>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": true
},
"32053": {
"content": "<extra_id_46>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": true
},
"32054": {
"content": "<extra_id_45>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": true
},
"32055": {
"content": "<extra_id_44>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": true
},
"32056": {
"content": "<extra_id_43>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": true
},
"32057": {
"content": "<extra_id_42>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": true
},
"32058": {
"content": "<extra_id_41>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": true
},
"32059": {
"content": "<extra_id_40>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": true
},
"32060": {
"content": "<extra_id_39>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": true
},
"32061": {
"content": "<extra_id_38>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": true
},
"32062": {
"content": "<extra_id_37>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": true
},
"32063": {
"content": "<extra_id_36>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": true
},
"32064": {
"content": "<extra_id_35>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": true
},
"32065": {
"content": "<extra_id_34>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": true
},
"32066": {
"content": "<extra_id_33>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": true
},
"32067": {
"content": "<extra_id_32>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": true
},
"32068": {
"content": "<extra_id_31>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": true
},
"32069": {
"content": "<extra_id_30>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": true
},
"32070": {
"content": "<extra_id_29>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": true
},
"32071": {
"content": "<extra_id_28>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": true
},
"32072": {
"content": "<extra_id_27>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": true
},
"32073": {
"content": "<extra_id_26>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": true
},
"32074": {
"content": "<extra_id_25>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": true
},
"32075": {
"content": "<extra_id_24>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": true
},
"32076": {
"content": "<extra_id_23>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": true
},
"32077": {
"content": "<extra_id_22>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": true
},
"32078": {
"content": "<extra_id_21>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": true
},
"32079": {
"content": "<extra_id_20>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": true
},
"32080": {
"content": "<extra_id_19>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": true
},
"32081": {
"content": "<extra_id_18>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": true
},
"32082": {
"content": "<extra_id_17>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": true
},
"32083": {
"content": "<extra_id_16>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": true
},
"32084": {
"content": "<extra_id_15>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": true
},
"32085": {
"content": "<extra_id_14>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": true
},
"32086": {
"content": "<extra_id_13>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": true
},
"32087": {
"content": "<extra_id_12>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": true
},
"32088": {
"content": "<extra_id_11>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": true
},
"32089": {
"content": "<extra_id_10>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": true
},
"32090": {
"content": "<extra_id_9>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": true
},
"32091": {
"content": "<extra_id_8>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": true
},
"32092": {
"content": "<extra_id_7>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": true
},
"32093": {
"content": "<extra_id_6>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": true
},
"32094": {
"content": "<extra_id_5>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": true
},
"32095": {
"content": "<extra_id_4>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": true
},
"32096": {
"content": "<extra_id_3>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": true
},
"32097": {
"content": "<extra_id_2>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": true
},
"32098": {
"content": "<extra_id_1>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": true
},
"32099": {
"content": "<extra_id_0>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": true
}
},
"additional_special_tokens": [
"<extra_id_0>",
"<extra_id_1>",
"<extra_id_2>",
"<extra_id_3>",
"<extra_id_4>",
"<extra_id_5>",
"<extra_id_6>",
"<extra_id_7>",
"<extra_id_8>",
"<extra_id_9>",
"<extra_id_10>",
"<extra_id_11>",
"<extra_id_12>",
"<extra_id_13>",
"<extra_id_14>",
"<extra_id_15>",
"<extra_id_16>",
"<extra_id_17>",
"<extra_id_18>",
"<extra_id_19>",
"<extra_id_20>",
"<extra_id_21>",
"<extra_id_22>",
"<extra_id_23>",
"<extra_id_24>",
"<extra_id_25>",
"<extra_id_26>",
"<extra_id_27>",
"<extra_id_28>",
"<extra_id_29>",
"<extra_id_30>",
"<extra_id_31>",
"<extra_id_32>",
"<extra_id_33>",
"<extra_id_34>",
"<extra_id_35>",
"<extra_id_36>",
"<extra_id_37>",
"<extra_id_38>",
"<extra_id_39>",
"<extra_id_40>",
"<extra_id_41>",
"<extra_id_42>",
"<extra_id_43>",
"<extra_id_44>",
"<extra_id_45>",
"<extra_id_46>",
"<extra_id_47>",
"<extra_id_48>",
"<extra_id_49>",
"<extra_id_50>",
"<extra_id_51>",
"<extra_id_52>",
"<extra_id_53>",
"<extra_id_54>",
"<extra_id_55>",
"<extra_id_56>",
"<extra_id_57>",
"<extra_id_58>",
"<extra_id_59>",
"<extra_id_60>",
"<extra_id_61>",
"<extra_id_62>",
"<extra_id_63>",
"<extra_id_64>",
"<extra_id_65>",
"<extra_id_66>",
"<extra_id_67>",
"<extra_id_68>",
"<extra_id_69>",
"<extra_id_70>",
"<extra_id_71>",
"<extra_id_72>",
"<extra_id_73>",
"<extra_id_74>",
"<extra_id_75>",
"<extra_id_76>",
"<extra_id_77>",
"<extra_id_78>",
"<extra_id_79>",
"<extra_id_80>",
"<extra_id_81>",
"<extra_id_82>",
"<extra_id_83>",
"<extra_id_84>",
"<extra_id_85>",
"<extra_id_86>",
"<extra_id_87>",
"<extra_id_88>",
"<extra_id_89>",
"<extra_id_90>",
"<extra_id_91>",
"<extra_id_92>",
"<extra_id_93>",
"<extra_id_94>",
"<extra_id_95>",
"<extra_id_96>",
"<extra_id_97>",
"<extra_id_98>",
"<extra_id_99>"
],
"clean_up_tokenization_spaces": true,
"eos_token": "</s>",
"extra_ids": 100,
"legacy": true,
"model_max_length": 512,
"pad_token": "<pad>",
"sp_model_kwargs": {},
"tokenizer_class": "T5Tokenizer",
"unk_token": "<unk>"
}
config/FLUX.1-dev/transformer/config.json
+14
View File
@@ -0,0 +1,14 @@
{
"_class_name": "FluxTransformer2DModel",
"_diffusers_version": "0.30.0.dev0",
"_name_or_path": "../checkpoints/flux-dev/transformer",
"attention_head_dim": 128,
"guidance_embeds": true,
"in_channels": 64,
"joint_attention_dim": 4096,
"num_attention_heads": 24,
"num_layers": 19,
"num_single_layers": 38,
"patch_size": 1,
"pooled_projection_dim": 768
}
config/FLUX.1-dev/transformer/diffusion_pytorch_model.safetensors.index.json
+1167
View File
File diff suppressed because it is too large Load Diff
config/FLUX.1-dev/vae/config.json
+38
View File
@@ -0,0 +1,38 @@
{
"_class_name": "AutoencoderKL",
"_diffusers_version": "0.30.0.dev0",
"_name_or_path": "../checkpoints/flux-dev",
"act_fn": "silu",
"block_out_channels": [
128,
256,
512,
512
],
"down_block_types": [
"DownEncoderBlock2D",
"DownEncoderBlock2D",
"DownEncoderBlock2D",
"DownEncoderBlock2D"
],
"force_upcast": true,
"in_channels": 3,
"latent_channels": 16,
"latents_mean": null,
"latents_std": null,
"layers_per_block": 2,
"mid_block_add_attention": true,
"norm_num_groups": 32,
"out_channels": 3,
"sample_size": 1024,
"scaling_factor": 0.3611,
"shift_factor": 0.1159,
"up_block_types": [
"UpDecoderBlock2D",
"UpDecoderBlock2D",
"UpDecoderBlock2D",
"UpDecoderBlock2D"
],
"use_post_quant_conv": false,
"use_quant_conv": false
}
config/glm/text_encoder_config.json
+42
View File
@@ -0,0 +1,42 @@
{
"_name_or_path": "THUDM/chatglm3-6b-base",
"model_type": "chatglm",
"architectures": [
"ChatGLMModel"
],
"auto_map": {
"AutoConfig": "configuration_chatglm.ChatGLMConfig",
"AutoModel": "modeling_chatglm.ChatGLMForConditionalGeneration",
"AutoModelForCausalLM": "modeling_chatglm.ChatGLMForConditionalGeneration",
"AutoModelForSeq2SeqLM": "modeling_chatglm.ChatGLMForConditionalGeneration",
"AutoModelForSequenceClassification": "modeling_chatglm.ChatGLMForSequenceClassification"
},
"add_bias_linear": false,
"add_qkv_bias": true,
"apply_query_key_layer_scaling": true,
"apply_residual_connection_post_layernorm": false,
"attention_dropout": 0.0,
"attention_softmax_in_fp32": true,
"bias_dropout_fusion": true,
"ffn_hidden_size": 13696,
"fp32_residual_connection": false,
"hidden_dropout": 0.0,
"hidden_size": 4096,
"kv_channels": 128,
"layernorm_epsilon": 1e-05,
"multi_query_attention": true,
"multi_query_group_num": 2,
"num_attention_heads": 32,
"num_layers": 28,
"original_rope": true,
"padded_vocab_size": 65024,
"post_layer_norm": true,
"rmsnorm": true,
"seq_length": 32768,
"use_cache": true,
"torch_dtype": "float16",
"transformers_version": "4.30.2",
"tie_word_embeddings": false,
"eos_token_id": 2,
"pad_token_id": 0
}
consistory/consistory_pipeline_wapper.py
+619
View File
@@ -0,0 +1,619 @@
# Copyright 2023 The HuggingFace Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# Not a contribution
# Changes made by NVIDIA CORPORATION & AFFILIATES enabling ConsiStory or otherwise documented as NVIDIA-proprietary
# are not a contribution and subject to the license under the LICENSE file located at the root directory.
import torch
from diffusers.pipelines.stable_diffusion_xl.pipeline_output import StableDiffusionXLPipelineOutput
from diffusers.pipelines.stable_diffusion_xl.pipeline_stable_diffusion_xl import StableDiffusionXLPipeline, \
rescale_noise_cfg, EXAMPLE_DOC_STRING
from diffusers.models.attention_processor import (
AttnProcessor2_0,
FusedAttnProcessor2_0,
XFormersAttnProcessor,)
from diffusers.loaders import FromSingleFileMixin
from diffusers.utils import (
deprecate,
is_torch_xla_available,
logging,
replace_example_docstring,
)
from typing import Any, Callable, Dict, List, Optional, Tuple, Union
from .attention_processor import register_extended_self_attn
from .consistory_utils import FeatureInjector, AnchorCache, QueryStore
from .utils.ptp_utils import AttentionStore
if is_torch_xla_available():
# import torch_xla.core.xla_model as xm
XLA_AVAILABLE = True
else:
XLA_AVAILABLE = False
logger = logging.get_logger(__name__) # pylint: disable=invalid-name
T = torch.Tensor
class ConsistoryExtendAttnSDXLPipeline(
StableDiffusionXLPipeline
):
model_cpu_offload_seq = "text_encoder->text_encoder_2->image_encoder->unet->vae"
_optional_components = [
"tokenizer",
"tokenizer_2",
"text_encoder",
"text_encoder_2",
"image_encoder",
"feature_extractor",
]
_callback_tensor_inputs = [
"latents",
"prompt_embeds",
"negative_prompt_embeds",
"add_text_embeds",
"add_time_ids",
"negative_pooled_prompt_embeds",
"negative_add_time_ids",
]
def upcast_vae(self):
dtype = self.vae.dtype
self.vae.to(dtype=torch.float32)
use_torch_2_0_or_xformers = isinstance(
self.vae.decoder.mid_block.attentions[0].processor,
(
AttnProcessor2_0,
XFormersAttnProcessor,
FusedAttnProcessor2_0,
),
)
# if xformers or torch_2_0 is used attention block does not need
# to be in float32 which can save lots of memory
if use_torch_2_0_or_xformers:
self.vae.post_quant_conv.to(dtype)
self.vae.decoder.conv_in.to(dtype)
self.vae.decoder.mid_block.to(dtype)
@torch.no_grad()
@replace_example_docstring(EXAMPLE_DOC_STRING)
def __call__(
self,
prompt: Union[str, List[str]] = None,
prompt_2: Optional[Union[str, List[str]]] = None,
height: Optional[int] = None,
width: Optional[int] = None,
num_inference_steps: int = 50,
denoising_end: Optional[float] = None,
guidance_scale: float = 5.0,
negative_prompt: Optional[Union[str, List[str]]] = None,
negative_prompt_2: Optional[Union[str, List[str]]] = None,
num_images_per_prompt: Optional[int] = 1,
eta: float = 0.0,
generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
latents: Optional[torch.FloatTensor] = None,
prompt_embeds: Optional[torch.FloatTensor] = None,
negative_prompt_embeds: Optional[torch.FloatTensor] = None,
pooled_prompt_embeds: Optional[torch.FloatTensor] = None,
negative_pooled_prompt_embeds: Optional[torch.FloatTensor] = None,
output_type: Optional[str] = "latent",
return_dict: bool = True,
cross_attention_kwargs: Optional[Dict[str, Any]] = None,
guidance_rescale: float = 0.0,
original_size: Optional[Tuple[int, int]] = None,
crops_coords_top_left: Tuple[int, int] = (0, 0),
target_size: Optional[Tuple[int, int]] = None,
negative_original_size: Optional[Tuple[int, int]] = None,
negative_crops_coords_top_left: Tuple[int, int] = (0, 0),
negative_target_size: Optional[Tuple[int, int]] = None,
clip_skip: Optional[int] = None,
callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None,
callback_on_step_end_tensor_inputs: List[str] = ["latents"],
attention_store_kwargs: Optional[Dict] = None,
extended_attn_kwargs: Optional[Dict] = None,
share_queries: bool = False,
query_store_kwargs: Optional[Dict] = {},
feature_injector: Optional[FeatureInjector] = None,
anchors_cache: Optional[AnchorCache] = None,
instance_latents: Optional[torch.FloatTensor] = None,
cf_clip=None,
**kwargs,
):
r"""
Function invoked when calling the pipeline for generation.
Args:
prompt (`str` or `List[str]`, *optional*):
The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.
instead.
prompt_2 (`str` or `List[str]`, *optional*):
The prompt or prompts to be sent to the `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is
used in both text-encoders
height (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor):
The height in pixels of the generated image. This is set to 1024 by default for the best results.
Anything below 512 pixels won't work well for
[stabilityai/stable-diffusion-xl-base-1.0](https://huggingface.co/stabilityai/stable-diffusion-xl-base-1.0)
and checkpoints that are not specifically fine-tuned on low resolutions.
width (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor):
The width in pixels of the generated image. This is set to 1024 by default for the best results.
Anything below 512 pixels won't work well for
[stabilityai/stable-diffusion-xl-base-1.0](https://huggingface.co/stabilityai/stable-diffusion-xl-base-1.0)
and checkpoints that are not specifically fine-tuned on low resolutions.
num_inference_steps (`int`, *optional*, defaults to 50):
The number of denoising steps. More denoising steps usually lead to a higher quality image at the
expense of slower inference.
denoising_end (`float`, *optional*):
When specified, determines the fraction (between 0.0 and 1.0) of the total denoising process to be
completed before it is intentionally prematurely terminated. As a result, the returned sample will
still retain a substantial amount of noise as determined by the discrete timesteps selected by the
scheduler. The denoising_end parameter should ideally be utilized when this pipeline forms a part of a
"Mixture of Denoisers" multi-pipeline setup, as elaborated in [**Refining the Image
Output**](https://huggingface.co/docs/diffusers/api/pipelines/stable_diffusion/stable_diffusion_xl#refining-the-image-output)
guidance_scale (`float`, *optional*, defaults to 5.0):
Guidance scale as defined in [Classifier-Free Diffusion Guidance](https://arxiv.org/abs/2207.12598).
`guidance_scale` is defined as `w` of equation 2. of [Imagen
Paper](https://arxiv.org/pdf/2205.11487.pdf). Guidance scale is enabled by setting `guidance_scale >
1`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`,
usually at the expense of lower image quality.
negative_prompt (`str` or `List[str]`, *optional*):
The prompt or prompts not to guide the image generation. If not defined, one has to pass
`negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
less than `1`).
negative_prompt_2 (`str` or `List[str]`, *optional*):
The prompt or prompts not to guide the image generation to be sent to `tokenizer_2` and
`text_encoder_2`. If not defined, `negative_prompt` is used in both text-encoders
num_images_per_prompt (`int`, *optional*, defaults to 1):
The number of images to generate per prompt.
eta (`float`, *optional*, defaults to 0.0):
Corresponds to parameter eta (η) in the DDIM paper: https://arxiv.org/abs/2010.02502. Only applies to
[`schedulers.DDIMScheduler`], will be ignored for others.
generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)
to make generation deterministic.
latents (`torch.FloatTensor`, *optional*):
Pre-generated noisy latents, sampled from a Gaussian distribution, to be used as inputs for image
generation. Can be used to tweak the same generation with different prompts. If not provided, a latents
tensor will ge generated by sampling using the supplied random `generator`.
prompt_embeds (`torch.FloatTensor`, *optional*):
Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not
provided, text embeddings will be generated from `prompt` input argument.
negative_prompt_embeds (`torch.FloatTensor`, *optional*):
Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt
weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input
argument.
pooled_prompt_embeds (`torch.FloatTensor`, *optional*):
Pre-generated pooled text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting.
If not provided, pooled text embeddings will be generated from `prompt` input argument.
negative_pooled_prompt_embeds (`torch.FloatTensor`, *optional*):
Pre-generated negative pooled text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt
weighting. If not provided, pooled negative_prompt_embeds will be generated from `negative_prompt`
input argument.
output_type (`str`, *optional*, defaults to `"pil"`):
The output format of the generate image. Choose between
[PIL](https://pillow.readthedocs.io/en/stable/): `PIL.Image.Image` or `np.array`.
return_dict (`bool`, *optional*, defaults to `True`):
Whether or not to return a [`~pipelines.stable_diffusion_xl.StableDiffusionXLPipelineOutput`] instead
of a plain tuple.
cross_attention_kwargs (`dict`, *optional*):
A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under
`self.processor` in
[diffusers.models.attention_processor](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py).
guidance_rescale (`float`, *optional*, defaults to 0.0):
Guidance rescale factor proposed by [Common Diffusion Noise Schedules and Sample Steps are
Flawed](https://arxiv.org/pdf/2305.08891.pdf) `guidance_scale` is defined as `φ` in equation 16. of
[Common Diffusion Noise Schedules and Sample Steps are Flawed](https://arxiv.org/pdf/2305.08891.pdf).
Guidance rescale factor should fix overexposure when using zero terminal SNR.
original_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)):
If `original_size` is not the same as `target_size` the image will appear to be down- or upsampled.
`original_size` defaults to `(height, width)` if not specified. Part of SDXL's micro-conditioning as
explained in section 2.2 of
[https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952).
crops_coords_top_left (`Tuple[int]`, *optional*, defaults to (0, 0)):
`crops_coords_top_left` can be used to generate an image that appears to be "cropped" from the position
`crops_coords_top_left` downwards. Favorable, well-centered images are usually achieved by setting
`crops_coords_top_left` to (0, 0). Part of SDXL's micro-conditioning as explained in section 2.2 of
[https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952).
target_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)):
For most cases, `target_size` should be set to the desired height and width of the generated image. If
not specified it will default to `(height, width)`. Part of SDXL's micro-conditioning as explained in
section 2.2 of [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952).
negative_original_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)):
To negatively condition the generation process based on a specific image resolution. Part of SDXL's
micro-conditioning as explained in section 2.2 of
[https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). For more
information, refer to this issue thread: https://github.com/huggingface/diffusers/issues/4208.
negative_crops_coords_top_left (`Tuple[int]`, *optional*, defaults to (0, 0)):
To negatively condition the generation process based on a specific crop coordinates. Part of SDXL's
micro-conditioning as explained in section 2.2 of
[https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). For more
information, refer to this issue thread: https://github.com/huggingface/diffusers/issues/4208.
negative_target_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)):
To negatively condition the generation process based on a target image resolution. It should be as same
as the `target_size` for most cases. Part of SDXL's micro-conditioning as explained in section 2.2 of
[https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). For more
information, refer to this issue thread: https://github.com/huggingface/diffusers/issues/4208.
callback_on_step_end (`Callable`, *optional*):
A function that calls at the end of each denoising steps during the inference. The function is called
with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int,
callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by
`callback_on_step_end_tensor_inputs`.
callback_on_step_end_tensor_inputs (`List`, *optional*):
The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list
will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the
`._callback_tensor_inputs` attribute of your pipeine class.
Examples:
Returns:
[`~pipelines.stable_diffusion_xl.StableDiffusionXLPipelineOutput`] or `tuple`:
[`~pipelines.stable_diffusion_xl.StableDiffusionXLPipelineOutput`] if `return_dict` is True, otherwise a
`tuple`. When returning a tuple, the first element is a list with the generated images.
"""
callback = kwargs.pop("callback", None)
callback_steps = kwargs.pop("callback_steps", None)
if callback is not None:
deprecate(
"callback",
"1.0.0",
"Passing `callback` as an input argument to `__call__` is deprecated, consider use `callback_on_step_end`",
)
if callback_steps is not None:
deprecate(
"callback_steps",
"1.0.0",
"Passing `callback_steps` as an input argument to `__call__` is deprecated, consider use `callback_on_step_end`",
)
# 0. Default height and width to unet
height = height or self.default_sample_size * self.vae_scale_factor
width = width or self.default_sample_size * self.vae_scale_factor
original_size = original_size or (height, width)
target_size = target_size or (height, width)
# 1. Check inputs. Raise error if not correct
# self.check_inputs(
# prompt,
# prompt_2,
# height,
# width,
# callback_steps,
# negative_prompt,
# negative_prompt_2,
# prompt_embeds,
# negative_prompt_embeds,
# pooled_prompt_embeds,
# negative_pooled_prompt_embeds,
# callback_on_step_end_tensor_inputs,
# )
self._guidance_scale = guidance_scale
self._guidance_rescale = guidance_rescale
self._clip_skip = clip_skip
self._cross_attention_kwargs = cross_attention_kwargs
self._denoising_end = denoising_end
# 2. Define call parameters
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]
device = self._execution_device
# 3. Encode input prompt
lora_scale = (
self.cross_attention_kwargs.get("scale", None) if self.cross_attention_kwargs is not None else None
)
if isinstance(prompt, list):
p_emb,po_emb=[],[]
for i in prompt:
tokens_p = cf_clip.tokenize(i)
output_p = cf_clip.encode_from_tokens(tokens_p, return_dict=True) # {"pooled_output":tensor}
prompt_embeds_,pooled_prompt_embeds_ = output_p.pop("cond").to(device,self.unet.dtype),output_p.pop("pooled_output").to(device,self.unet.dtype)
p_emb.append(prompt_embeds_),
po_emb.append(pooled_prompt_embeds_)
prompt_embeds = torch.cat(p_emb,dim=0)
pooled_prompt_embeds = torch.cat(po_emb,dim=0)
else:
tokens_p = cf_clip.tokenize(prompt)
output_p = cf_clip.encode_from_tokens(tokens_p, return_dict=True) # {"pooled_output":tensor}
prompt_embeds,pooled_prompt_embeds = output_p.pop("cond").to(device,self.unet.dtype),output_p.pop("pooled_output").to(device,self.unet.dtype)
if negative_prompt is not None:
tokens_n = cf_clip.tokenize(negative_prompt)
output_n = cf_clip.encode_from_tokens(tokens_n, return_dict=True) # {"pooled_output":tensor}
negative_prompt_embeds,negative_pooled_prompt_embeds = output_n.pop("cond").to(device,self.unet.dtype),output_n.pop("pooled_output").to(device,self.unet.dtype)
if isinstance(prompt, list):
negative_prompt_embeds = torch.cat(batch_size*[negative_prompt_embeds] ,dim=0)
negative_pooled_prompt_embeds = torch.cat(batch_size*[negative_pooled_prompt_embeds] ,dim=0)
else:
negative_prompt_embeds = torch.zeros_like(prompt_embeds)
negative_pooled_prompt_embeds = torch.zeros_like(pooled_prompt_embeds)
# print("prompt_embeds",prompt_embeds.shape) #prompt_embeds torch.Size([2, 77, 2048])
# print("prompt_embeds",prompt_embeds.shape) #prompt_embeds torch.Size([2, 77, 2048])
# print("negative_prompt_embeds",negative_prompt_embeds.shape)#negative_prompt_embeds torch.Size([1, 77, 2048])
# print("pooled_prompt_embeds",pooled_prompt_embeds.shape)#pooled_prompt_embeds torch.Size([2, 1280])
# print("negative_pooled_prompt_embeds",negative_pooled_prompt_embeds.shape)#negative_pooled_prompt_embeds torch.Size([1, 1280])
# if prompt is not None: #PASS it
# (
# prompt_embeds,
# negative_prompt_embeds,
# pooled_prompt_embeds,
# negative_pooled_prompt_embeds,
# ) = self.encode_prompt(
# prompt=prompt,
# prompt_2=prompt_2,
# device=device,
# num_images_per_prompt=num_images_per_prompt,
# do_classifier_free_guidance=self.do_classifier_free_guidance,
# negative_prompt=negative_prompt,
# negative_prompt_2=negative_prompt_2,
# prompt_embeds=prompt_embeds,
# negative_prompt_embeds=negative_prompt_embeds,
# pooled_prompt_embeds=pooled_prompt_embeds,
# negative_pooled_prompt_embeds=negative_pooled_prompt_embeds,
# lora_scale=lora_scale,
# clip_skip=self.clip_skip,
# )
# 4. Prepare timesteps
self.scheduler.set_timesteps(num_inference_steps, device=device)
timesteps = self.scheduler.timesteps
# 5. Prepare latent variables
num_channels_latents = self.unet.config.in_channels
latents = self.prepare_latents(
batch_size * num_images_per_prompt,
num_channels_latents,
height,
width,
prompt_embeds.dtype,
device,
generator,
latents,
)
# 6. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline
extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta)
if share_queries:
query_store = QueryStore(**query_store_kwargs)
else:
query_store = None
self.attention_store = AttentionStore(attention_store_kwargs)
register_extended_self_attn(self.unet, self.attention_store, extended_attn_kwargs)
# 7. Prepare added time ids & embeddings
add_text_embeds = pooled_prompt_embeds
if self.text_encoder_2 is None:
text_encoder_projection_dim = int(pooled_prompt_embeds.shape[-1])
else:
text_encoder_projection_dim = self.text_encoder_2.config.projection_dim
add_time_ids = self._get_add_time_ids(
original_size,
crops_coords_top_left,
target_size,
dtype=prompt_embeds.dtype,
text_encoder_projection_dim=text_encoder_projection_dim,
)
if negative_original_size is not None and negative_target_size is not None:
negative_add_time_ids = self._get_add_time_ids(
negative_original_size,
negative_crops_coords_top_left,
negative_target_size,
dtype=prompt_embeds.dtype,
text_encoder_projection_dim=text_encoder_projection_dim,
)
else:
negative_add_time_ids = add_time_ids
if self.do_classifier_free_guidance:
prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds], dim=0)
add_text_embeds = torch.cat([negative_pooled_prompt_embeds, add_text_embeds], dim=0)
add_time_ids = torch.cat([negative_add_time_ids, add_time_ids], dim=0)
prompt_embeds = prompt_embeds.to(device)
add_text_embeds = add_text_embeds.to(device)
add_time_ids = add_time_ids.to(device).repeat(batch_size * num_images_per_prompt, 1)
# 8. Denoising loop
num_warmup_steps = max(len(timesteps) - num_inference_steps * self.scheduler.order, 0)
# 8.1 Apply denoising_end
if (
self.denoising_end is not None
and isinstance(self.denoising_end, float)
and self.denoising_end > 0
and self.denoising_end < 1
):
discrete_timestep_cutoff = int(
round(
self.scheduler.config.num_train_timesteps
- (self.denoising_end * self.scheduler.config.num_train_timesteps)
)
)
num_inference_steps = len(list(filter(lambda ts: ts >= discrete_timestep_cutoff, timesteps)))
timesteps = timesteps[:num_inference_steps]
# 9. Optionally get Guidance Scale Embedding
timestep_cond = None
if self.unet.config.time_cond_proj_dim is not None:
guidance_scale_tensor = torch.tensor(self.guidance_scale - 1).repeat(batch_size * num_images_per_prompt)
timestep_cond = self.get_guidance_scale_embedding(
guidance_scale_tensor, embedding_dim=self.unet.config.time_cond_proj_dim
).to(device=device, dtype=latents.dtype)
self._num_timesteps = len(timesteps)
if instance_latents is not None:
n_instances = instance_latents.shape[0]
instance_noise = latents[:n_instances].clone()
with self.progress_bar(total=num_inference_steps) as progress_bar:
for i, t in enumerate(timesteps):
self.attention_store.curr_iter = i
if instance_latents is not None:
noised_instances = self.scheduler.add_noise(instance_latents, instance_noise, t.repeat(n_instances).long())
latents[:n_instances] = noised_instances
# expand the latents if we are doing classifier free guidance
latent_model_input = torch.cat([latents] * 2) if self.do_classifier_free_guidance else latents
latent_model_input = self.scheduler.scale_model_input(latent_model_input, t)
# predict the noise residual
added_cond_kwargs = {"text_embeds": add_text_embeds, "time_ids": add_time_ids}
if share_queries and (i >= query_store.t_range[0] and i <= query_store.t_range[1]):
query_store.set_mode('cache')
# print(latent_model_input.shape) #[4, 4, 128, 128]
# print(prompt_embeds.shape)#([4, 77, 2048])
noise_pred_vanilla = self.unet(
latent_model_input,
t,
encoder_hidden_states=prompt_embeds,
timestep_cond=timestep_cond,
cross_attention_kwargs={'query_store': query_store,
'perform_extend_attn': False,
'record_attention': False},
added_cond_kwargs=added_cond_kwargs,
return_dict=False,
)[0]
query_store.set_mode('inject')
noise_pred = self.unet(
latent_model_input,
t,
encoder_hidden_states=prompt_embeds,
timestep_cond=timestep_cond,
cross_attention_kwargs={'query_store': query_store,
'perform_extend_attn': True,
'record_attention': True,
'feature_injector': feature_injector,
'anchors_cache': anchors_cache},
added_cond_kwargs=added_cond_kwargs,
return_dict=False,
)[0]
# perform guidance
if self.do_classifier_free_guidance:
noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
noise_pred = noise_pred_uncond + self.guidance_scale * (noise_pred_text - noise_pred_uncond)
if self.do_classifier_free_guidance and self.guidance_rescale > 0.0:
# Based on 3.4. in https://arxiv.org/pdf/2305.08891.pdf
noise_pred = rescale_noise_cfg(noise_pred, noise_pred_text, guidance_rescale=self.guidance_rescale)
# compute the previous noisy sample x_t -> x_t-1
latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs, return_dict=False)[0]
if callback_on_step_end is not None:
callback_kwargs = {}
for k in callback_on_step_end_tensor_inputs:
callback_kwargs[k] = locals()[k]
callback_outputs = callback_on_step_end(self, i, t, callback_kwargs)
latents = callback_outputs.pop("latents", latents)
prompt_embeds = callback_outputs.pop("prompt_embeds", prompt_embeds)
negative_prompt_embeds = callback_outputs.pop("negative_prompt_embeds", negative_prompt_embeds)
add_text_embeds = callback_outputs.pop("add_text_embeds", add_text_embeds)
negative_pooled_prompt_embeds = callback_outputs.pop(
"negative_pooled_prompt_embeds", negative_pooled_prompt_embeds
)
add_time_ids = callback_outputs.pop("add_time_ids", add_time_ids)
negative_add_time_ids = callback_outputs.pop("negative_add_time_ids", negative_add_time_ids)
# call the callback, if provided
if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0):
progress_bar.update()
if callback is not None and i % callback_steps == 0:
step_idx = i // getattr(self.scheduler, "order", 1)
callback(step_idx, t, latents)
if XLA_AVAILABLE:
# xm.mark_step()
pass
# Update attention store mask
self.attention_store.aggregate_last_steps_attention()
if not output_type == "latent":
# make sure the VAE is in float32 mode, as it overflows in float16
#print(self.vae.config.force_upcast)
needs_upcasting = self.vae.dtype == torch.float16 and self.vae.config.force_upcast
#print(needs_upcasting,112233)
if needs_upcasting:
self.upcast_vae()
latents = latents.to(next(iter(self.vae.post_quant_conv.parameters())).dtype)
elif latents.dtype != self.vae.dtype:
if torch.backends.mps.is_available():
# some platforms (eg. apple mps) misbehave due to a pytorch bug: https://github.com/pytorch/pytorch/pull/99272
self.vae = self.vae.to(latents.dtype)
# unscale/denormalize the latents
# denormalize with the mean and std if available and not None
has_latents_mean = hasattr(self.vae.config, "latents_mean") and self.vae.config.latents_mean is not None
has_latents_std = hasattr(self.vae.config, "latents_std") and self.vae.config.latents_std is not None
if has_latents_mean and has_latents_std:
latents_mean = (
torch.tensor(self.vae.config.latents_mean).view(1, 4, 1, 1).to(latents.device, latents.dtype)
)
latents_std = (
torch.tensor(self.vae.config.latents_std).view(1, 4, 1, 1).to(latents.device, latents.dtype)
)
latents = latents * latents_std / self.vae.config.scaling_factor + latents_mean
else:
latents = latents / self.vae.config.scaling_factor
image = self.vae.decode(latents, return_dict=False)[0]
#image = self.vae.decode(latents / self.vae.config.scaling_factor, return_dict=False)[0]
# cast back to fp16 if needed
if needs_upcasting:
self.vae.to(dtype=torch.float16)
else:
image = latents/ 0.13025
if not output_type == "latent":
# apply watermark if available
# if self.watermark is not None:
# image = self.watermark.apply_watermark(image)
image = self.image_processor.postprocess(image, output_type=output_type)
# Offload all models
self.maybe_free_model_hooks()
if not return_dict:
return (image,)
return StableDiffusionXLPipelineOutput(images=image)
consistory/consistory_run.py
+11 -9
View File
@@ -20,6 +20,8 @@ def clear_memory():
gc.collect()
def load_pipeline(repo_id,unet_path,gpu_id=0):
float_type = torch.float16
#sd_id = "stabilityai/stable-diffusion-xl-base-1.0"
@@ -108,7 +110,7 @@ def run_batch_generation(story_pipeline, prompts, concept_token,negative_prompt,
seed=40, n_steps=50, mask_dropout=0.5,
same_latent=False, share_queries=True,
perform_sdsa=True, perform_injection=True,
downscale_rate=4, n_achors=2):
downscale_rate=4, n_achors=2,cf_clip=None):
device = story_pipeline.device
tokenizer = story_pipeline.tokenizer
float_type = story_pipeline.dtype
@@ -145,7 +147,7 @@ def run_batch_generation(story_pipeline, prompts, concept_token,negative_prompt,
extended_attn_kwargs=extended_attn_kwargs,
share_queries=share_queries,
query_store_kwargs=query_store_kwargs,
num_inference_steps=n_steps)
num_inference_steps=n_steps,cf_clip=cf_clip)
# Extended attention with nn_map #
@@ -171,7 +173,7 @@ def run_batch_generation(story_pipeline, prompts, concept_token,negative_prompt,
share_queries=share_queries,
query_store_kwargs=query_store_kwargs,
feature_injector=feature_injector,
num_inference_steps=n_steps)
num_inference_steps=n_steps,cf_clip=cf_clip)
#img_all = view_images([np.array(x) for x in out.images], display_image=False, downscale_rate=downscale_rate)
# display_attn_maps(story_pipeline.attention_store.last_mask, out.images)
@@ -184,7 +186,7 @@ def run_anchor_generation(story_pipeline, prompts, concept_token,negative_prompt
seed=40, n_steps=50, mask_dropout=0.5,
same_latent=False, share_queries=True,
perform_sdsa=True, perform_injection=True,
downscale_rate=4, cache_cpu_offloading=False):
downscale_rate=4, cache_cpu_offloading=False,cf_clip=None):
device = story_pipeline.device
tokenizer = story_pipeline.tokenizer
float_type = story_pipeline.dtype
@@ -224,7 +226,7 @@ def run_anchor_generation(story_pipeline, prompts, concept_token,negative_prompt
share_queries=share_queries,
query_store_kwargs=query_store_kwargs,
anchors_cache=anchor_cache_first_stage,
num_inference_steps=n_steps)
num_inference_steps=n_steps,cf_clip=cf_clip)
last_masks = story_pipeline.attention_store.last_mask
@@ -256,7 +258,7 @@ def run_anchor_generation(story_pipeline, prompts, concept_token,negative_prompt
query_store_kwargs=query_store_kwargs,
feature_injector=feature_injector,
anchors_cache=anchor_cache_second_stage,
num_inference_steps=n_steps)
num_inference_steps=n_steps,cf_clip=cf_clip)
#img_all = view_images([np.array(x) for x in out.images], display_image=False, downscale_rate=downscale_rate)
# display_attn_maps(story_pipeline.attention_store.last_mask, out.images)
@@ -277,7 +279,7 @@ def run_extra_generation(story_pipeline, prompts, concept_token, negative_prompt
seed=40, n_steps=50, mask_dropout=0.5,
same_latent=False, share_queries=True,
perform_sdsa=True, perform_injection=True,
downscale_rate=4, cache_cpu_offloading=False):
downscale_rate=4, cache_cpu_offloading=False,cf_clip=None):
device = story_pipeline.device
tokenizer = story_pipeline.tokenizer
float_type = story_pipeline.dtype
@@ -323,7 +325,7 @@ def run_extra_generation(story_pipeline, prompts, concept_token, negative_prompt
share_queries=share_queries,
query_store_kwargs=query_store_kwargs,
anchors_cache=anchor_cache_first_stage,
num_inference_steps=n_steps)
num_inference_steps=n_steps,cf_clip=cf_clip)
# ------------------ #
@@ -360,7 +362,7 @@ def run_extra_generation(story_pipeline, prompts, concept_token, negative_prompt
query_store_kwargs=query_store_kwargs,
feature_injector=feature_injector,
anchors_cache=anchor_cache_second_stage,
num_inference_steps=n_steps)
num_inference_steps=n_steps,cf_clip=cf_clip)
#img_all = view_images([np.array(x) for x in out.images], display_image=False, downscale_rate=downscale_rate)
examples/2rolein1img.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 484 KiB

examples/flux.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 312 KiB

examples/new_0927.json
+2321
View File
File diff suppressed because it is too large Load Diff
examples/outdated.json
+663
View File
@@ -0,0 +1,663 @@
{
"last_node_id": 25,
"last_link_id": 43,
"nodes": [
{
"id": 23,
"type": "SaveImage",
"pos": {
"0": 1213,
"1": -578
},
"size": [
723.8321337686805,
498.9358674218496
],
"flags": {},
"order": 10,
"mode": 0,
"inputs": [
{
"name": "images",
"type": "IMAGE",
"link": 31,
"label": "images"
}
],
"outputs": [],
"properties": {},
"widgets_values": [
"ComfyUI"
]
},
{
"id": 21,
"type": "ImageBatch",
"pos": {
"0": 62,
"1": -166
},
"size": {
"0": 210,
"1": 46
},
"flags": {},
"order": 4,
"mode": 0,
"inputs": [
{
"name": "image1",
"type": "IMAGE",
"link": 39,
"label": "image1"
},
{
"name": "image2",
"type": "IMAGE",
"link": 40,
"label": "image2"
}
],
"outputs": [
{
"name": "IMAGE",
"type": "IMAGE",
"links": [
29
],
"shape": 3,
"label": "IMAGE",
"slot_index": 0
}
],
"properties": {
"Node name for S&R": "ImageBatch"
}
},
{
"id": 20,
"type": "LoadImage",
"pos": {
"0": -428,
"1": 5
},
"size": {
"0": 315,
"1": 314
},
"flags": {},
"order": 0,
"mode": 0,
"inputs": [],
"outputs": [
{
"name": "IMAGE",
"type": "IMAGE",
"links": [
36,
39
],
"shape": 3,
"label": "IMAGE",
"slot_index": 0
},
{
"name": "MASK",
"type": "MASK",
"links": null,
"shape": 3,
"label": "MASK"
}
],
"properties": {
"Node name for S&R": "LoadImage"
},
"widgets_values": [
"4.jpg",
"image"
]
},
{
"id": 19,
"type": "LoadImage",
"pos": {
"0": -429,
"1": -432
},
"size": [
315,
314
],
"flags": {},
"order": 1,
"mode": 0,
"inputs": [],
"outputs": [
{
"name": "IMAGE",
"type": "IMAGE",
"links": [
38,
40
],
"shape": 3,
"label": "IMAGE",
"slot_index": 0
},
{
"name": "MASK",
"type": "MASK",
"links": null,
"shape": 3,
"label": "MASK"
}
],
"properties": {
"Node name for S&R": "LoadImage"
},
"widgets_values": [
"1 (1).jpg",
"image"
]
},
{
"id": 9,
"type": "Storydiffusion_Sampler",
"pos": {
"0": 776,
"1": -671
},
"size": [
402.78569644628124,
620.1786676859502
],
"flags": {},
"order": 8,
"mode": 0,
"inputs": [
{
"name": "model",
"type": "STORY_DICT",
"link": 30,
"label": "model"
}
],
"outputs": [
{
"name": "image",
"type": "IMAGE",
"links": [
31
],
"shape": 3,
"label": "image",
"slot_index": 0
},
{
"name": "prompt_array",
"type": "STRING",
"links": null,
"shape": 3,
"label": "prompt_array"
}
],
"properties": {
"Node name for S&R": "Storydiffusion_Sampler"
},
"widgets_values": [
"[Taylor] wake up in the bed ;\n[Lecun] is working.",
"bad anatomy, bad hands, missing fingers, extra fingers, three hands, three legs, bad arms, missing legs, missing arms, poorly drawn face, bad face, fused face, cloned face, three crus, fused feet, fused thigh, extra crus, ugly fingers, horn,amputation, disconnected limbs",
"No_style",
1753524959,
"randomize",
20,
7,
1,
20,
3.5,
0.5,
5,
false,
0.8,
"0., 0.25, 0.4, 0.75;0.6, 0.25, 1., 0.75"
]
},
{
"id": 13,
"type": "DualCLIPLoader",
"pos": {
"0": -8,
"1": 291
},
"size": {
"0": 315,
"1": 106
},
"flags": {},
"order": 2,
"mode": 0,
"inputs": [],
"outputs": [
{
"name": "CLIP",
"type": "CLIP",
"links": [
22
],
"shape": 3,
"label": "CLIP",
"slot_index": 0
}
],
"properties": {
"Node name for S&R": "DualCLIPLoader"
},
"widgets_values": [
"t5xxl_fp8_e4m3fn.safetensors",
"clip_l.safetensors",
"flux"
]
},
{
"id": 24,
"type": "ImageBatch",
"pos": {
"0": 61,
"1": 33
},
"size": {
"0": 210,
"1": 46
},
"flags": {},
"order": 3,
"mode": 0,
"inputs": [
{
"name": "image1",
"type": "IMAGE",
"link": 38,
"label": "image1"
},
{
"name": "image2",
"type": "IMAGE",
"link": 36,
"label": "image2"
}
],
"outputs": [
{
"name": "IMAGE",
"type": "IMAGE",
"links": [
43
],
"shape": 3,
"label": "IMAGE",
"slot_index": 0
}
],
"properties": {
"Node name for S&R": "ImageBatch"
}
},
{
"id": 25,
"type": "Storydiffusion_Sampler",
"pos": {
"0": 781,
"1": -12
},
"size": {
"0": 402.78570556640625,
"1": 620.1786499023438
},
"flags": {},
"order": 7,
"mode": 0,
"inputs": [
{
"name": "model",
"type": "STORY_DICT",
"link": 41,
"label": "model"
}
],
"outputs": [
{
"name": "image",
"type": "IMAGE",
"links": [
42
],
"shape": 3,
"label": "image",
"slot_index": 0
},
{
"name": "prompt_array",
"type": "STRING",
"links": null,
"shape": 3,
"label": "prompt_array",
"slot_index": 1
}
],
"properties": {
"Node name for S&R": "Storydiffusion_Sampler"
},
"widgets_values": [
"[Taylor] wake up in the bed ;\n[Lecun] is working.",
"bad anatomy, bad hands, missing fingers, extra fingers, three hands, three legs, bad arms, missing legs, missing arms, poorly drawn face, bad face, fused face, cloned face, three crus, fused feet, fused thigh, extra crus, ugly fingers, horn,amputation, disconnected limbs",
"No_style",
1753524959,
"randomize",
20,
7,
1,
20,
3.5,
0.5,
5,
false,
0.8,
"0., 0.25, 0.4, 0.75;0.6, 0.25, 1., 0.75"
]
},
{
"id": 8,
"type": "Storydiffusion_Model_Loader",
"pos": {
"0": 337,
"1": -15
},
"size": [
415.4856964462813,
579.8786676859503
],
"flags": {},
"order": 5,
"mode": 0,
"inputs": [
{
"name": "image",
"type": "IMAGE",
"link": 43,
"label": "image"
},
{
"name": "control_image",
"type": "IMAGE",
"link": null,
"label": "control_image"
},
{
"name": "model",
"type": "MODEL",
"link": null,
"label": "model"
},
{
"name": "clip",
"type": "CLIP",
"link": 22,
"label": "clip"
},
{
"name": "vae",
"type": "VAE",
"link": null,
"label": "vae"
}
],
"outputs": [
{
"name": "model",
"type": "STORY_DICT",
"links": [
41
],
"shape": 3,
"label": "model",
"slot_index": 0
}
],
"properties": {
"Node name for S&R": "Storydiffusion_Model_Loader"
},
"widgets_values": [
"[Taylor] a woman img, wearing a white T-shirt, blue loose hair.\n[Lecun] a man img,wearing a suit,black hair.",
"",
"flux1-kj-dev-fp8.safetensors",
"ae.safetensors",
"none",
"none",
0.8,
"none",
"EVA02_CLIP_L_336_psz14_s6B.pt",
"best quality",
"euler",
"normal",
0.5,
0.5,
768,
768,
"v1",
"pulid,fp8,cpu"
]
},
{
"id": 10,
"type": "SaveImage",
"pos": {
"0": 1224,
"1": -10
},
"size": [
713.6932337686808,
472.37407242184963
],
"flags": {},
"order": 9,
"mode": 0,
"inputs": [
{
"name": "images",
"type": "IMAGE",
"link": 42,
"label": "images"
}
],
"outputs": [],
"properties": {},
"widgets_values": [
"ComfyUI"
]
},
{
"id": 22,
"type": "Storydiffusion_Model_Loader",
"pos": {
"0": 330,
"1": -678
},
"size": {
"0": 415.4856872558594,
"1": 579.878662109375
},
"flags": {},
"order": 6,
"mode": 0,
"inputs": [
{
"name": "image",
"type": "IMAGE",
"link": 29,
"label": "image"
},
{
"name": "control_image",
"type": "IMAGE",
"link": null,
"label": "control_image"
},
{
"name": "model",
"type": "MODEL",
"link": null,
"label": "model"
},
{
"name": "clip",
"type": "CLIP",
"link": null,
"label": "clip"
},
{
"name": "vae",
"type": "VAE",
"link": null,
"label": "vae"
}
],
"outputs": [
{
"name": "model",
"type": "STORY_DICT",
"links": [
30
],
"shape": 3,
"label": "model",
"slot_index": 0
}
],
"properties": {
"Node name for S&R": "Storydiffusion_Model_Loader"
},
"widgets_values": [
"[Taylor] a woman img, wearing a white T-shirt, blue loose hair.\n[Lecun] a man img,wearing a suit,black hair.",
"F:/test/ComfyUI/models/diffusers/black-forest-labs/FLUX.1-dev",
"flux1-kj-dev-fp8.safetensors",
"ae.safetensors",
"none",
"none",
0.8,
"none",
"EVA02_CLIP_L_336_psz14_s6B.pt",
"best quality",
"euler",
"normal",
0.5,
0.5,
768,
768,
"v1",
"pulid,fp8,cpu"
]
}
],
"links": [
[
22,
13,
0,
8,
3,
"CLIP"
],
[
29,
21,
0,
22,
0,
"IMAGE"
],
[
30,
22,
0,
9,
0,
"STORY_DICT"
],
[
31,
9,
0,
23,
0,
"IMAGE"
],
[
36,
20,
0,
24,
1,
"IMAGE"
],
[
38,
19,
0,
24,
0,
"IMAGE"
],
[
39,
20,
0,
21,
0,
"IMAGE"
],
[
40,
19,
0,
21,
1,
"IMAGE"
],
[
41,
8,
0,
25,
0,
"STORY_DICT"
],
[
42,
25,
0,
10,
0,
"IMAGE"
],
[
43,
24,
0,
8,
0,
"IMAGE"
]
],
"groups": [],
"config": {},
"extra": {
"ds": {
"scale": 0.620921323059155,
"offset": [
722.8692862313196,
817.9813775781506
]
}
},
"version": 0.4
}
kolors/models/modeling_chatglm.py
+35 -19
View File
@@ -10,11 +10,11 @@ import torch
import torch.utils.checkpoint
import torch.nn.functional as F
from torch import nn
from torch.nn import CrossEntropyLoss, LayerNorm
from torch.nn import CrossEntropyLoss, LayerNorm, MSELoss, BCEWithLogitsLoss
from torch.nn.utils import skip_init
from typing import Optional, Tuple, Union, List, Callable, Dict, Any
from copy import deepcopy
import transformers
from transformers.modeling_outputs import (
BaseModelOutputWithPast,
@@ -26,11 +26,7 @@ from transformers.utils import logging
from transformers.generation.logits_process import LogitsProcessor
from transformers.generation.utils import LogitsProcessorList, StoppingCriteriaList, GenerationConfig, ModelOutput
try:
from .configuration_chatglm import ChatGLMConfig
except:
from configuration_chatglm import ChatGLMConfig
from .configuration_chatglm import ChatGLMConfig
# flags required to enable jit fusion kernels
@@ -43,13 +39,16 @@ if sys.platform != 'darwin':
logger = logging.get_logger(__name__)
_CHECKPOINT_FOR_DOC = "THUDM/ChatGLM"
_CONFIG_FOR_DOC = "ChatGLM6BConfig"
_CONFIG_FOR_DOC = "ChatGLMConfig"
CHATGLM_6B_PRETRAINED_MODEL_ARCHIVE_LIST = [
"THUDM/chatglm3-6b-base",
"THUDM/chatglm3-6b",
# See all ChatGLM models at https://huggingface.co/models?filter=chatglm
]
is_transformers_4_42_or_higher = int(transformers.__version__.split(".")[1]) >= 42
is_transformers_4_44_or_higher = int(transformers.__version__.split(".")[1]) >= 44
def default_init(cls, *args, **kwargs):
return cls(*args, **kwargs)
@@ -186,7 +185,7 @@ def apply_rotary_pos_emb(x: torch.Tensor, rope_cache: torch.Tensor) -> torch.Ten
class RMSNorm(torch.nn.Module):
def __init__(self, normalized_shape, eps=1e-5, device=None, dtype=None, **kwargs):
super().__init__()
self.weight = torch.nn.Parameter(torch.empty(normalized_shape, device=device, dtype=dtype))
self.weight = torch.nn.Parameter(torch.ones(normalized_shape, device=device, dtype=dtype))
self.eps = eps
def forward(self, hidden_states: torch.Tensor):
@@ -639,7 +638,8 @@ class GLMTransformer(torch.nn.Module):
attention_mask,
rotary_pos_emb,
kv_caches[index],
use_cache
use_cache,
use_reentrant=False
)
else:
layer_ret = layer(
@@ -702,9 +702,9 @@ class ChatGLMPreTrainedModel(PreTrainedModel):
position_ids = torch.arange(seq_length, dtype=torch.long, device=device).unsqueeze(0).repeat(batch_size, 1)
return position_ids
def _set_gradient_checkpointing(self, module, value=False):
if isinstance(module, GLMTransformer):
module.gradient_checkpointing = value
def gradient_checkpointing_enable(self, gradient_checkpointing_kwargs=None):
if not self.supports_gradient_checkpointing:
raise ValueError(f"{self.__class__.__name__} does not support gradient checkpointing.")
class Embedding(torch.nn.Module):
@@ -773,6 +773,9 @@ class ChatGLMModel(ChatGLMPreTrainedModel):
def get_input_embeddings(self):
return self.embedding.word_embeddings
def set_input_embeddings(self, value):
self.embedding.word_embeddings = value
def get_prompt(self, batch_size, device, dtype=torch.half):
prefix_tokens = self.prefix_tokens.unsqueeze(0).expand(batch_size, -1).to(device)
past_key_values = self.prefix_encoder(prefix_tokens).type(dtype)
@@ -873,9 +876,19 @@ class ChatGLMForConditionalGeneration(ChatGLMPreTrainedModel):
standardize_cache_format: bool = False,
) -> Dict[str, Any]:
# update past_key_values
model_kwargs["past_key_values"] = self._extract_past_from_model_output(
outputs, standardize_cache_format=standardize_cache_format
)
if is_transformers_4_44_or_higher:
model_kwargs["past_key_values"] = self._extract_past_from_model_output(
outputs
)[1]
elif is_transformers_4_42_or_higher:
# update past_key_values
model_kwargs["past_key_values"] = self._extract_past_from_model_output(
outputs, standardize_cache_format=standardize_cache_format
)[1]
else:
model_kwargs["past_key_values"] = self._extract_past_from_model_output(
outputs, standardize_cache_format=standardize_cache_format
)
# update attention mask
if "attention_mask" in model_kwargs:
@@ -1005,7 +1018,10 @@ class ChatGLMForConditionalGeneration(ChatGLMPreTrainedModel):
content = ""
history = deepcopy(history)
for response in output.split("<|assistant|>"):
metadata, content = response.split("\n", maxsplit=1)
if "\n" in response:
metadata, content = response.split("\n", maxsplit=1)
else:
metadata, content = "", response
if not metadata.strip():
content = content.strip()
history.append({"role": "assistant", "metadata": metadata, "content": content})
@@ -1023,7 +1039,7 @@ class ChatGLMForConditionalGeneration(ChatGLMPreTrainedModel):
return content, history
@torch.inference_mode()
def chat(self, tokenizer, query: str, history: List[Tuple[str, str]] = None, role: str = "user",
def chat(self, tokenizer, query: str, history: List[Dict] = None, role: str = "user",
max_length: int = 8192, num_beams=1, do_sample=True, top_p=0.8, temperature=0.8, logits_processor=None,
**kwargs):
if history is None:
@@ -1045,7 +1061,7 @@ class ChatGLMForConditionalGeneration(ChatGLMPreTrainedModel):
return response, history
@torch.inference_mode()
def stream_chat(self, tokenizer, query: str, history: List[Tuple[str, str]] = None, role: str = "user",
def stream_chat(self, tokenizer, query: str, history: List[Dict] = None, role: str = "user",
past_key_values=None,max_length: int = 8192, do_sample=True, top_p=0.8, temperature=0.8,
logits_processor=None, return_past_key_values=False, **kwargs):
if history is None:
kolors/models/quantization.py
+188
View File
File diff suppressed because one or more lines are too long
kolors/models/tokenization_chatglm.py
+53 -23
View File
@@ -8,6 +8,9 @@ from transformers.utils import logging, PaddingStrategy
from transformers.tokenization_utils_base import EncodedInput, BatchEncoding
logger = logging.get_logger(__name__)
class SPTokenizer:
def __init__(self, model_path: str):
# reload tokenizer
@@ -29,7 +32,7 @@ class SPTokenizer:
self.special_tokens[token] = self.n_words
self.index_special_tokens[self.n_words] = token
self.n_words += 1
self.role_special_token_expression = "|".join([re.escape(token) for token in role_special_tokens])
self.role_special_token_expression = "|".join([re.escape(token) for token in special_tokens]) # for apply_chat_template
def tokenize(self, s: str, encode_special_tokens=False):
if encode_special_tokens:
@@ -83,31 +86,40 @@ class SPTokenizer:
"""Converts an index (integer) in a token (str) using the vocab."""
if index in self.index_special_tokens:
return self.index_special_tokens[index]
if index in [self.eos_id, self.bos_id, self.pad_id] or index < 0:
if index in [self.eos_id, self.bos_id, self.pad_id] or index < 0 or index > self.sp_model.vocab_size():
return ""
return self.sp_model.IdToPiece(index)
class ChatGLMTokenizer(PreTrainedTokenizer):
vocab_files_names = {"vocab_file": "tokenizer.model"}
vocab_files_names = {"vocab_file": "tokenizer.model"}
model_input_names = ["input_ids", "attention_mask", "position_ids"]
def __init__(self, vocab_file, padding_side="left", clean_up_tokenization_spaces=False, encode_special_tokens=False,
**kwargs):
def __init__(
self,
vocab_file,
padding_side="left",
clean_up_tokenization_spaces=False,
encode_special_tokens=False,
**kwargs
):
self.name = "GLMTokenizer"
self.vocab_file = vocab_file
self.tokenizer = SPTokenizer(vocab_file)
self.special_tokens = {
"<bos>": self.tokenizer.bos_id,
"<eos>": self.tokenizer.eos_id,
"<unk>": self.tokenizer.pad_id,
"<pad>": self.tokenizer.pad_id
}
self.encode_special_tokens = encode_special_tokens
super().__init__(padding_side=padding_side, clean_up_tokenization_spaces=clean_up_tokenization_spaces,
encode_special_tokens=encode_special_tokens,
**kwargs)
super().__init__(
padding_side=padding_side,
clean_up_tokenization_spaces=clean_up_tokenization_spaces,
**kwargs
)
def get_command(self, token):
if token in self.special_tokens:
@@ -117,24 +129,40 @@ class ChatGLMTokenizer(PreTrainedTokenizer):
@property
def unk_token(self) -> str:
return "<unk>"
return self.tokenizer.sp_model.IdToPiece(self.get_command("<unk>"))
@property
def pad_token(self) -> str:
return "<unk>"
@property
def pad_token_id(self):
return self.get_command("<pad>")
return self.tokenizer.sp_model.IdToPiece(self.get_command("<pad>"))
@property
def eos_token(self) -> str:
return "</s>"
return self.tokenizer.sp_model.IdToPiece(self.get_command("<eos>"))
@property
def unk_token_id(self) -> int:
return self.get_command("<unk>")
@property
def pad_token_id(self) -> int:
return self.get_command("<pad>")
@property
def eos_token_id(self):
return self.get_command("<eos>")
@unk_token.setter
def unk_token(self, value):
logger.warning("Setting unk_token is not supported, use the default one.")
@pad_token.setter
def pad_token(self, value):
logger.warning("Setting pad_token is not supported, use the default one.")
@eos_token.setter
def eos_token(self, value):
logger.warning("Setting eos_token is not supported, use the default one.")
@property
def vocab_size(self):
return self.tokenizer.n_words
@@ -212,7 +240,7 @@ class ChatGLMTokenizer(PreTrainedTokenizer):
return self.batch_encode_plus([input_ids], return_tensors="pt", is_split_into_words=True)
def build_inputs_with_special_tokens(
self, token_ids_0: List[int], token_ids_1: Optional[List[int]] = None
self, token_ids_0: List[int], token_ids_1: Optional[List[int]] = None
) -> List[int]:
"""
Build model inputs from a sequence or a pair of sequence for sequence classification tasks by concatenating and
@@ -237,12 +265,13 @@ class ChatGLMTokenizer(PreTrainedTokenizer):
return token_ids_0
def _pad(
self,
encoded_inputs: Union[Dict[str, EncodedInput], BatchEncoding],
max_length: Optional[int] = None,
padding_strategy: PaddingStrategy = PaddingStrategy.DO_NOT_PAD,
pad_to_multiple_of: Optional[int] = None,
return_attention_mask: Optional[bool] = None,
self,
encoded_inputs: Union[Dict[str, EncodedInput], BatchEncoding],
max_length: Optional[int] = None,
padding_side: str = "left",
padding_strategy: PaddingStrategy = PaddingStrategy.DO_NOT_PAD,
pad_to_multiple_of: Optional[int] = None,
return_attention_mask: Optional[bool] = None,
) -> dict:
"""
Pad encoded inputs (on left/right and up to predefined length or max length in the batch)
@@ -269,6 +298,7 @@ class ChatGLMTokenizer(PreTrainedTokenizer):
"""
# Load from model defaults
assert self.padding_side == "left"
assert padding_side is None or padding_side == "left"
required_input = encoded_inputs[self.model_input_names[0]]
seq_length = len(required_input)
kolors/pipelines/__pycache__/__init__.cpython-311.pyc
BIN
View File
Binary file not shown.
kolors/pipelines/__pycache__/pipeline_stable_diffusion_xl_chatglm_256_ipadapter_FaceID.cpython-311.pyc
BIN
View File
Binary file not shown.
kolors/pipelines/pipeline_stable_diffusion_xl_chatglm_256_ipadapter_FaceID.py
+21 -20
View File
@@ -680,7 +680,7 @@ class StableDiffusionXLPipeline(DiffusionPipeline, FromSingleFileMixin, LoraLoad
negative_prompt_embeds: Optional[torch.FloatTensor] = None,
pooled_prompt_embeds: Optional[torch.FloatTensor] = None,
negative_pooled_prompt_embeds: Optional[torch.FloatTensor] = None,
output_type: Optional[str] = "pil",
output_type: Optional[str] = "latent",
return_dict: bool = True,
callback: Optional[Callable[[int, int, torch.FloatTensor], None]] = None,
callback_steps: int = 1,
@@ -818,24 +818,24 @@ class StableDiffusionXLPipeline(DiffusionPipeline, FromSingleFileMixin, LoraLoad
text_encoder_lora_scale = (
cross_attention_kwargs.get("scale", None) if cross_attention_kwargs is not None else None
)
(
prompt_embeds,
negative_prompt_embeds,
pooled_prompt_embeds,
negative_pooled_prompt_embeds,
) = self.encode_prompt(
prompt,
device,
num_images_per_prompt,
do_classifier_free_guidance,
negative_prompt,
prompt_embeds=prompt_embeds,
negative_prompt_embeds=negative_prompt_embeds,
pooled_prompt_embeds=pooled_prompt_embeds,
negative_pooled_prompt_embeds=negative_pooled_prompt_embeds,
lora_scale=text_encoder_lora_scale,
)
if prompt is not None:
(
prompt_embeds,
negative_prompt_embeds,
pooled_prompt_embeds,
negative_pooled_prompt_embeds,
) = self.encode_prompt(
prompt,
device,
num_images_per_prompt,
do_classifier_free_guidance,
negative_prompt,
prompt_embeds=prompt_embeds,
negative_prompt_embeds=negative_prompt_embeds,
pooled_prompt_embeds=pooled_prompt_embeds,
negative_pooled_prompt_embeds=negative_pooled_prompt_embeds,
lora_scale=text_encoder_lora_scale,
)
##### prepare fused face embeds
if face_crop_image is not None and face_insightface_embeds is not None:
@@ -945,7 +945,8 @@ class StableDiffusionXLPipeline(DiffusionPipeline, FromSingleFileMixin, LoraLoad
latents = latents.to(next(iter(self.vae.post_quant_conv.parameters())).dtype)
image = self.vae.decode(latents / self.vae.config.scaling_factor, return_dict=False)[0]
else:
image = latents
image = latents/ self.vae.config.scaling_factor
#print("image:",image.shape)
return StableDiffusionXLPipelineOutput(images=image)
# image = self.watermark.apply_watermark(image)
model_loader_utils.py
+1687 -1100
View File
File diff suppressed because it is too large Load Diff
msdiffusion/models/attention_processor.py
+3
View File
@@ -273,10 +273,13 @@ class MaskedIPAttnProcessor2_0(nn.Module):
# (batch, heads, source_length, target_length)
attention_mask = attention_mask.view(batch_size, attn.heads, -1, attention_mask.shape[-1])
rf_attention_mask = None
#print(boxes, phrase_idxes, hidden_states.shape[1],self.text_tokens, batch_size, attn.heads,
#hidden_states.dtype, hidden_states.device) #one None 2304 77 2 10 torch.float16 cuda:0
custom_attention_masks = self.prepare_attention_mask_qk(boxes, phrase_idxes, hidden_states.shape[1],
self.text_tokens, batch_size, attn.heads,
hidden_states.dtype, hidden_states.device,
use_masked_text_attention=False)
#print(custom_attention_masks)
attention_mask_qk_image, attention_mask_qk_text, dummy_attention_mask = custom_attention_masks
if attention_mask_qk_image is not None:
attention_mask_qk_image = attention_mask_qk_image.view(batch_size, attn.heads, -1, attention_mask_qk_image.shape[-1])
msdiffusion/models/model.py
+12 -8
View File
@@ -62,7 +62,7 @@ class MSAdapter(torch.nn.Module):
def set_ms_adapter(self, weight_dtype=torch.float16, cache_attention_maps=True):
# set attention processor
attn_procs_ = {}
attn_procs = {}
for name in self.unet.attn_processors.keys():
cross_attention_dim = None if name.endswith("attn1.processor") else self.unet.config.cross_attention_dim
if name.startswith("mid_block"):
@@ -74,15 +74,15 @@ class MSAdapter(torch.nn.Module):
block_id = int(name[len("down_blocks.")])
hidden_size = self.unet.config.block_out_channels[block_id]
if cross_attention_dim is None:
attn_procs_[name] = AttnProcessor()
attn_procs[name] = AttnProcessor()
else:
attn_procs_[name] = IPAttnProcessor(
attn_procs[name] = IPAttnProcessor(
hidden_size=hidden_size,
cross_attention_dim=cross_attention_dim,
num_tokens=self.num_tokens,
text_tokens=self.text_tokens,
).to(self.device, dtype=weight_dtype)
self.unet.set_attn_processor(attn_procs_)
self.unet.set_attn_processor(attn_procs)
self.adapter_modules = torch.nn.ModuleList(self.unet.attn_processors.values())
if self.controlnet is not None:
if isinstance(self.controlnet, MultiControlNetModel):
@@ -109,11 +109,11 @@ class MSAdapter(torch.nn.Module):
output_hidden_states=True).hidden_states[
-2] # (bsz*rn, num_tokens, embedding_dim)
else:
print(processed_images.shape,1233333)
image_encoder = image_encoder.encode_image
processed_images.to("cpu")
image_embeds = image_encoder(processed_images)["penultimate_hidden_states"]
# print(image_embeds.shape)
print(image_embeds.shape,1234567)
else:
image_embeds = image_encoder(
@@ -174,6 +174,7 @@ class MSAdapter(torch.nn.Module):
grounding_kwargs = None
if boxes is not None:
boxes = torch.tensor(boxes).to(self.device, weight_dtype)
print(boxes.shape,123)
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(self.device)
@@ -187,7 +188,7 @@ class MSAdapter(torch.nn.Module):
phrase_input_ids = torch.stack(phrase_input_ids)
phrase_input_ids = phrase_input_ids.view(-1, phrase_input_ids.shape[-1])
phrase_embeds = pipe.text_encoder(phrase_input_ids.to(self.device)).pooler_output
# print(phrase_embeds.shape,batch_boxes.shape)#torch.Size([1, 768]) torch.Size([1, 4])
print(phrase_embeds.shape,batch_boxes.shape) # torch.Size([2, 768]) torch.Size([2, 4])
grounding_kwargs = {"boxes": batch_boxes, "phrase_embeds": phrase_embeds,
"drop_grounding_tokens": drop_grounding_tokens}
else:
@@ -219,12 +220,13 @@ class MSAdapter(torch.nn.Module):
use_repo=use_repo)
del image_encoder
torch.cuda.empty_cache()
# print(image_embeds.shape) #torch.Size([1, 257, 1664])
# print(image_embeds.shape) #torch.Size([2, 257, 1664])
if not use_repo:
# print(image_embeds.device,self.device,)
image_embeds=image_embeds.clone().detach().to(self.device, weight_dtype)
#print(image_embeds.device)
image_prompt_embeds = self.image_proj_model(image_embeds, grounding_kwargs=grounding_kwargs)
del image_embeds
torch.cuda.empty_cache()
image_prompt_embeds = image_prompt_embeds.view(bsz, -1, image_prompt_embeds.shape[-2],
@@ -236,6 +238,7 @@ class MSAdapter(torch.nn.Module):
image_prompt_embeds.shape[
-1]) # (bsz, total_num_tokens*rn, cross_attention_dim)
image_prompt_embeds = torch.cat([self.dummy_image_tokens, image_prompt_embeds], dim=1)
uncond_image_prompt_embeds = torch.zeros_like(image_prompt_embeds)
bs_embed, seq_len, _ = image_prompt_embeds.shape
image_prompt_embeds = image_prompt_embeds.repeat(1, num_samples, 1)
@@ -250,6 +253,7 @@ class MSAdapter(torch.nn.Module):
do_classifier_free_guidance=True,
negative_prompt=negative_prompt,
)
prompt_embeds = torch.cat([prompt_embeds_, image_prompt_embeds], dim=1)
negative_prompt_embeds = torch.cat([negative_prompt_embeds_, uncond_image_prompt_embeds], dim=1)
msdiffusion/models/modelWrapper.py
+186
View File
@@ -0,0 +1,186 @@
import os
import math
from typing import List
import torch
import torch.nn as nn
from diffusers.pipelines.controlnet import MultiControlNetModel
from .attention_processor import MaskedIPAttnProcessor2_0 as IPAttnProcessor, AttnProcessor2_0 as AttnProcessor, \
CNAttnProcessor2_0 as CNAttnProcessor
class MSAdapter(torch.nn.Module):
def __init__(self, pipe, image_proj_model, adapter_modules=None, ckpt_path=None,
num_tokens=4, text_tokens=77, max_rn=4, num_dummy_tokens=4,
device="cuda", controlnet=None):
super().__init__()
self.unet = pipe.unet
self.pipe = pipe
self.image_proj_model = image_proj_model
self.adapter_modules = adapter_modules
self.num_tokens = num_tokens
self.num_dummy_tokens = num_dummy_tokens
self.text_tokens = text_tokens
self.max_rn = max_rn
self.device = device
self.controlnet = controlnet
self.cross_attention_dim = self.unet.config.cross_attention_dim
# set attention processor when inference
if self.adapter_modules is None:
self.set_ms_adapter()
# dummy image tokens
self.dummy_image_tokens = nn.Parameter(torch.randn(1, self.num_dummy_tokens, self.cross_attention_dim))
if ckpt_path is not None:
self.load_from_checkpoint(ckpt_path)
def load_from_checkpoint(self, ckpt_path: str):
# Calculate original checksums
orig_ip_proj_sum = torch.sum(torch.stack([torch.sum(p) for p in self.image_proj_model.parameters()]))
orig_adapter_sum = torch.sum(torch.stack([torch.sum(p) for p in self.adapter_modules.parameters()]))
state_dict = torch.load(ckpt_path, map_location="cpu")
# Load state dict for image_proj_model and adapter_modules when using resampler
self.image_proj_model.load_state_dict(state_dict["image_proj"], strict=False)
self.adapter_modules.load_state_dict(state_dict["ms_adapter"], strict=False)
self.load_state_dict({"dummy_image_tokens": state_dict["dummy_image_tokens"]}, strict=False)
# Calculate new checksums
new_ip_proj_sum = torch.sum(torch.stack([torch.sum(p) for p in self.image_proj_model.parameters()]))
new_adapter_sum = torch.sum(torch.stack([torch.sum(p) for p in self.adapter_modules.parameters()]))
# Verify if the weights have changed
assert orig_ip_proj_sum != new_ip_proj_sum, "Weights of image_proj_model did not change!"
assert orig_adapter_sum != new_adapter_sum, "Weights of adapter_modules did not change!"
print(f"Successfully loaded weights from checkpoint {ckpt_path}")
def set_ms_adapter(self, weight_dtype=torch.float16, cache_attention_maps=True):
# set attention processor
attn_procs_ = {}
for name in self.unet.attn_processors.keys():
cross_attention_dim = None if name.endswith("attn1.processor") else self.unet.config.cross_attention_dim
if name.startswith("mid_block"):
hidden_size = self.unet.config.block_out_channels[-1]
elif name.startswith("up_blocks"):
block_id = int(name[len("up_blocks.")])
hidden_size = list(reversed(self.unet.config.block_out_channels))[block_id]
elif name.startswith("down_blocks"):
block_id = int(name[len("down_blocks.")])
hidden_size = self.unet.config.block_out_channels[block_id]
if cross_attention_dim is None:
attn_procs_[name] = AttnProcessor()
else:
attn_procs_[name] = IPAttnProcessor(
hidden_size=hidden_size,
cross_attention_dim=cross_attention_dim,
num_tokens=self.num_tokens,
text_tokens=self.text_tokens,
).to(self.device, dtype=weight_dtype)
self.unet.set_attn_processor(attn_procs_)
self.adapter_modules = torch.nn.ModuleList(self.unet.attn_processors.values())
if self.controlnet is not None:
if isinstance(self.controlnet, MultiControlNetModel):
for controlnet in self.controlnet.nets:
controlnet.set_attn_processor(
CNAttnProcessor(text_tokens=self.text_tokens, num_tokens=self.num_tokens))
else:
self.controlnet.set_attn_processor(
CNAttnProcessor(text_tokens=self.text_tokens, num_tokens=self.num_tokens))
@torch.inference_mode()
def get_image_embeds(self, processed_images, image_encoder=None, image_proj_type="linear",
image_encoder_type="clip", weight_dtype=torch.float16, use_repo=None):
# get image embeds
# processed_images: [bsz, rn, ...]
if use_repo:
processed_images = processed_images.view(-1, processed_images.shape[-3], processed_images.shape[-2],
processed_images.shape[-1]) # (bsz*rn, ...)
# print(processed_images.shape) # torch.Size([1, 3, 224, 224])
if image_proj_type == "resampler":
if use_repo:
image_embeds = image_encoder(processed_images.to(self.device, dtype=weight_dtype),
output_hidden_states=True).hidden_states[
-2] # (bsz*rn, num_tokens, embedding_dim)
else:
image_encoder = image_encoder.encode_image
processed_images.to("cpu")
image_embeds = image_encoder(processed_images)["penultimate_hidden_states"]
# print(image_embeds.shape)
else:
image_embeds = image_encoder(
processed_images.to(self.device, dtype=weight_dtype)).image_embeds # (bsz*rn, embedding_dim)
return image_embeds # [bsz*rn, ...]
def set_scale(self, scale, subject_scales):
for attn_processor in self.pipe.unet.attn_processors.values():
if isinstance(attn_processor, IPAttnProcessor):
attn_processor.scale = scale
attn_processor.subject_scales = subject_scales
def enable_psuedo_attention_mask(self, mask_threshold=0.5, start_step=5):
for attn_processor in self.pipe.unet.attn_processors.values():
if isinstance(attn_processor, IPAttnProcessor):
attn_processor.mask_threshold = mask_threshold
attn_processor.start_step = start_step
attn_processor.use_psuedo_attention_mask = True
attn_processor.need_text_attention_map = True
attn_processor.attention_maps = [] # clear attention maps
def generate(self,pipe,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=4, weight_dtype=torch.float16, subject_scales=None, mask_threshold=None, start_step=5,**kwargs):
self.set_scale(scale, subject_scales)
if mask_threshold is not None:
self.enable_psuedo_attention_mask(mask_threshold, start_step)
with (torch.inference_mode()):
#print(image_embeds.is_cuda,grounding_kwargs),
image_prompt_embeds = self.image_proj_model(image_embeds, grounding_kwargs=grounding_kwargs) #([2, 16, 2048])
image_prompt_embeds = image_prompt_embeds.view(bsz, -1, image_prompt_embeds.shape[-2],
image_prompt_embeds.shape[
-1]) # (bsz, rn, num_tokens, cross_attention_dim)
image_prompt_embeds = image_prompt_embeds.view(bsz,
image_prompt_embeds.shape[-3] * image_prompt_embeds.shape[
-2],
image_prompt_embeds.shape[
-1]) # (bsz, total_num_tokens*rn, cross_attention_dim)
dummy_image_tokens=self.dummy_image_tokens.to(self.device, weight_dtype)
image_prompt_embeds = torch.cat([dummy_image_tokens, image_prompt_embeds], dim=1)
uncond_image_prompt_embeds = torch.zeros_like(image_prompt_embeds)
bs_embed, seq_len, _ = image_prompt_embeds.shape
image_prompt_embeds = image_prompt_embeds.repeat(1, num_samples, 1)
image_prompt_embeds = image_prompt_embeds.view(bs_embed * num_samples, seq_len, -1) #([1, 36, 2048])
uncond_image_prompt_embeds = uncond_image_prompt_embeds.repeat(1, num_samples, 1)
uncond_image_prompt_embeds = uncond_image_prompt_embeds.view(bs_embed * num_samples, seq_len, -1)
prompt_embeds_=prompt_embeds_.to(self.device, weight_dtype)#([1, 77, 2048])
negative_prompt_embeds_=negative_prompt_embeds_.to(self.device, weight_dtype)
generator=torch.Generator(device=self.device).manual_seed(seed)
prompt_embeds = torch.cat([prompt_embeds_, image_prompt_embeds], dim=1)
negative_prompt_embeds = torch.cat([negative_prompt_embeds_, uncond_image_prompt_embeds], dim=1)
samples = pipe(height=height, width=width, num_inference_steps=steps, guidance_scale=cfg,
generator=generator,
prompt_embeds=prompt_embeds,
negative_prompt_embeds=negative_prompt_embeds,
pooled_prompt_embeds=pooled_prompt_embeds,
negative_pooled_prompt_embeds=negative_pooled_prompt_embeds,
cross_attention_kwargs=cross_attention_kwargs,
**kwargs
)[0] # torch.Size([1, 4, 64, 64])
return samples
msdiffusion/models/projection.py
+2 -1
View File
@@ -200,8 +200,9 @@ class Resampler(nn.Module):
boxes = grounding_kwargs["boxes"]
phrase_embeds = grounding_kwargs["phrase_embeds"]
fourier_embeds = self.fourier_embedder(boxes)
#print(phrase_embeds.shape,fourier_embeds.shape)
grounding_embeds = torch.cat((phrase_embeds, fourier_embeds), dim=-1)
drop_grounding_tokens = grounding_kwargs["drop_grounding_tokens"]
num_ref = x.shape[0] // len(drop_grounding_tokens)
drop_grounding_tokens = [item for item in drop_grounding_tokens for _ in range(num_ref)]
pipelines/pipeline_flux_infusenet.py
+26 -16
View File
@@ -131,7 +131,7 @@ class FluxInfuseNetPipeline(FluxControlNetPipeline):
latents: Optional[torch.FloatTensor] = None,
prompt_embeds: Optional[torch.FloatTensor] = None,
pooled_prompt_embeds: Optional[torch.FloatTensor] = None,
output_type: Optional[str] = "pil",
output_type: Optional[str] = "latent",
return_dict: bool = True,
joint_attention_kwargs: Optional[Dict[str, Any]] = None,
callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None,
@@ -304,23 +304,27 @@ class FluxInfuseNetPipeline(FluxControlNetPipeline):
device = self._execution_device
dtype = self.transformer.dtype
prompt_embeds = prompt_embeds.to(device=device, dtype=dtype)
pooled_prompt_embeds = pooled_prompt_embeds.to(device=device, dtype=dtype)
lora_scale = (
self.joint_attention_kwargs.get("scale", None) if self.joint_attention_kwargs is not None else None
)
(
prompt_embeds,
pooled_prompt_embeds,
text_ids,
) = self.encode_prompt(
prompt=prompt,
prompt_2=prompt_2,
prompt_embeds=prompt_embeds,
pooled_prompt_embeds=pooled_prompt_embeds,
device=device,
num_images_per_prompt=num_images_per_prompt,
max_sequence_length=max_sequence_length,
lora_scale=lora_scale,
)
text_ids = torch.zeros(prompt_embeds.shape[1], 3).to(device=device, dtype=dtype)
if prompt is not None:
(
prompt_embeds,
pooled_prompt_embeds,
text_ids,
) = self.encode_prompt(
prompt=prompt,
prompt_2=prompt_2,
prompt_embeds=prompt_embeds,
pooled_prompt_embeds=pooled_prompt_embeds,
device=device,
num_images_per_prompt=num_images_per_prompt,
max_sequence_length=max_sequence_length,
lora_scale=lora_scale,
)
if negative_prompt is not None or (negative_prompt_embeds is not None and negative_pooled_prompt_embeds is not None):
(
negative_prompt_embeds,
@@ -514,6 +518,11 @@ class FluxInfuseNetPipeline(FluxControlNetPipeline):
cond_scale = controlnet_cond_scale * controlnet_keep[i]
# controlnet
# print(pooled_prompt_embeds.dtype)
# print(controlnet_prompt_embeds.dtype)
# print(controlnet_text_ids.dtype)
# print(latent_image_ids.dtype)
# print(latents.dtype)
controlnet_block_samples, controlnet_single_block_samples = self.controlnet(
hidden_states=latents,
controlnet_cond=control_image,
@@ -594,7 +603,8 @@ class FluxInfuseNetPipeline(FluxControlNetPipeline):
xm.mark_step()
if output_type == "latent":
image = latents
latents = self._unpack_latents(latents, height, width, 8)
image = (latents/0.3611)+0.1159
else:
latents = self._unpack_latents(latents, height, width, self.vae_scale_factor)
pipelines/pipeline_infu_flux.py
+21 -8
View File
@@ -131,7 +131,9 @@ class InfUFluxPipeline:
infu_flux_version='v1.0',
model_version='aes_stage2',
text_encoder_2=None,
text_encoder=None,
transformer=None,
vae=None,
):
@@ -140,11 +142,18 @@ class InfUFluxPipeline:
# Load pipeline
self.infusenet = FluxControlNetModel.from_pretrained(infu_model_path, torch_dtype=torch.bfloat16)
pipe = FluxInfuseNetPipeline.from_pretrained(
base_model_path,transformer=transformer,text_encoder_2=text_encoder_2,
controlnet=self.infusenet,
torch_dtype=torch.bfloat16,
)
if vae is not None:
pipe = FluxInfuseNetPipeline.from_pretrained(
base_model_path,vae=vae,transformer=transformer,text_encoder_2=text_encoder_2,text_encoder=text_encoder,
controlnet=self.infusenet,
torch_dtype=torch.bfloat16,
)
else:
pipe = FluxInfuseNetPipeline.from_pretrained(
base_model_path,transformer=transformer,text_encoder_2=text_encoder_2,text_encoder=text_encoder,
controlnet=self.infusenet,
torch_dtype=torch.bfloat16,
)
# try:
# infusenet_path = os.path.join(infu_model_path, 'InfuseNetModel')
# self.infusenet = FluxControlNetModel.from_pretrained(infusenet_path, torch_dtype=torch.bfloat16)
@@ -243,7 +252,8 @@ class InfUFluxPipeline:
def __call__(
self,
id_embed: Image.Image, # PIL.Image.Image (RGB)
prompt: str,
prompt_embeds,
pooled_prompt_embeds,
control_image: Optional[Image.Image] = None, # PIL.Image.Image (RGB) or None
width = 864,
height = 1152,
@@ -292,7 +302,7 @@ class InfUFluxPipeline:
print('Generating image')
seed_everything(seed)
image = self.pipe(
prompt=prompt,
prompt=None,
controlnet_prompt_embeds=id_embed,
control_image=control_image,
guidance_scale=guidance_scale,
@@ -303,6 +313,9 @@ class InfUFluxPipeline:
control_guidance_end=infusenet_guidance_end,
height=height,
width=width,
).images[0]
prompt_embeds = prompt_embeds,
pooled_prompt_embeds = pooled_prompt_embeds,
).images
return image
pyproject.toml
+3 -3
View File
@@ -1,9 +1,9 @@
[project]
name = "comfyui_storydiffusion"
description = "you can using sotry-diffusion in comfyui"
version = "1.3.0"
description = "you can using sotry-diffusion and lot's of ID reservation methods in comfyui"
version = "2.0.0"
license = { file = "LICENSE" }
dependencies = ["diffusers", "spaces", "peft"]
dependencies = ["diffusers", "peft"]
[project.urls]
Repository = "https://github.com/smthemex/ComfyUI_StoryDiffusion"
requirements.txt
-1
View File
@@ -1,4 +1,3 @@
spaces
peft
diffusers
omegaconf
utils/gradio_utils.py
+63 -3
View File
@@ -511,9 +511,9 @@ def process_original_prompt(character_dict,prompts,id_length,img_mode):
raise f"{character_key} not have prompt description, please remove it"
index_list = character_index_dict[character_key]
index_list = [index for index in index_list if len(invert_character_index_dict[index]) == 1]
if not img_mode:
if len(index_list) < id_length:
raise f"{character_key} not have enough prompt description, need no less than {id_length}, but you give {len(index_list)}"
# if not img_mode:
# if len(index_list) < id_length:
# raise f"{character_key} not have enough prompt description, need no less than {id_length}, but you give {len(index_list)}"
ref_index_dict[character_key] = index_list[:id_length]
ref_totals = ref_totals + index_list[:id_length]
return character_index_dict,invert_character_index_dict,replace_prompts,ref_index_dict,ref_totals
@@ -523,6 +523,66 @@ def process_original_prompt(character_dict,prompts,id_length,img_mode):
#ref_indexs_dict:{'[Taylor]': [0, 3], '[sam]': [1, 2]},if 1 role {'[Taylor]': [0]}
#ref_totals: [0, 3, 1, 2] if 1 role [0]
def process_original_text(role_dict,prompts):
'''
Args:
role_dict: # 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.'} 2role
prompts: #[ list of scense txt]
Returns:
1 role
role_index_dict:{'[Taylor]': [0, 1]},
invert_role_index_dict:{0: ['[Taylor]'], 1: ['[Taylor]']},
replace_prompts:[' a woman img, wearing a white T-shirt, blue loose hair. wake up in the bed ;', ' a woman img, wearing a white T-shirt, blue loose hair. have breakfast by the window;'],
ref_role_index_dict:{'[Taylor]': [0]},
ref_role_totals:[0]
2 role
role_index_dict:{'[Taylor]': [0, 1], '[Lecun]': [2, 3]},
invert_role_index_dict:{0: ['[Taylor]'], 1: ['[Taylor]'], 2: ['[Lecun]'], 3: ['[Lecun]']},
replace_prompts:[' a woman img, wearing a white T-shirt, blue loose hair. wake up in the bed ;', ' a woman img, wearing a white T-shirt, blue loose hair. have breakfast by the window;', ' a man img,wearing a suit,black hair. driving a car;', ' a man img,wearing a suit,black hair. is working.'],
ref_role_index_dict:{'[Taylor]': [0, 1], '[Lecun]': [2, 3]},
ref_role_totals:[0, 1, 2, 3]
'''
replace_prompts = []
role_index_dict = {}
invert_role_index_dict = {}
for ind,prompt in enumerate(prompts):
for key in role_dict.keys(): # [Taylor],[Lecun]
if key in prompt: # [Taylor] a woman img
if key not in role_index_dict:
role_index_dict[key] = [] #character_index_dict:{'[Taylor]': [], '[Lecun]': []},
role_index_dict[key].append(ind) #character_index_dict:{'[Taylor]': [0, 1], '[Lecun]': [2, 3]},
if ind not in invert_role_index_dict:
invert_role_index_dict[ind] = [] #invert_character_index_dict:{0: []},
invert_role_index_dict[ind].append(key) # invert_character_index_dict:{0: ['[Taylor]'], 1: ['[Taylor]'], 2: ['[Lecun]'], 3: ['[Lecun]']},
cur_prompt = prompt
if ind in invert_role_index_dict:
for key in invert_role_index_dict[ind]:
cur_prompt = cur_prompt.replace(key,role_dict[key] + " ") #去掉[Taylor]
replace_prompts.append(cur_prompt)
ref_role_index_dict = {}
ref_role_totals = []
id_len=len(role_dict)
for character_key in role_index_dict.keys():
if character_key not in role_index_dict:
raise f"{character_key} not have prompt description, please remove it"
index_list = role_index_dict[character_key] #[0, 1],[3, 4] 如果2是NC
index_list = [index for index in index_list if len(invert_role_index_dict[index]) == 1]
# if not img_mode:
# if len(index_list) < id_length:
# raise f"{character_key} not have enough prompt description, need no less than {id_length}, but you give {len(index_list)}"
ref_role_index_dict[character_key] = index_list[:id_len]
ref_role_totals = ref_role_totals + index_list[:id_len]
return role_index_dict,invert_role_index_dict,replace_prompts,ref_role_index_dict,ref_role_totals
#character_index_dict：{'[Taylor]': [0, 3], '[sam]': [1, 2]},if 1 role {'[Taylor]': [0, 1, 2]}
#invert_character_index_dict:{0: ['[Taylor]'], 1: ['[sam]'], 2: ['[sam]'], 3: ['[Taylor]']},if 1 role {0: ['[Taylor]'], 1: ['[Taylor]'], 2: ['[Taylor]']}
#ref_indexs_dict:{'[Taylor]': [0, 3], '[sam]': [1, 2]},if 1 role {'[Taylor]': [0]}
#ref_totals: [0, 3, 1, 2] if 1 role [0]
def get_ref_character(real_prompt,character_dict):
list_arr = []
for keys in character_dict.keys():
utils/load_models_utils.py
+8 -18
View File
@@ -38,23 +38,18 @@ def get_lora_dict():
datas = get_lora_dict()
lora_lightning_list = datas["lightning_xl_lora"]
def load_models(path,model_type,single_files,use_safetensors,photomake_mode,photomaker_path,lora,lora_path,trigger_words,lora_scale):
def load_models(path,model_type,single_files,use_safetensors,photo_vesion,photomaker_path,lora,lora_path,trigger_words,lora_scale):
path=get_instance_path(path)
if model_type == "txt2img":
if single_files:
try:
pipe = StableDiffusionXLPipeline.from_single_file(
path,config=add_config, original_config=original_config_file, torch_dtype=torch.float16)
except:
try:
pipe = StableDiffusionXLPipeline.from_single_file(
pipe = StableDiffusionXLPipeline.from_single_file(
path,config=add_config, original_config_file=original_config_file,
torch_dtype=torch.float16
)
except:
raise "load pipe error"
else:
pipe = StableDiffusionXLPipeline.from_pretrained(
path, torch_dtype=torch.float16,use_safetensors=use_safetensors
@@ -69,7 +64,7 @@ def load_models(path,model_type,single_files,use_safetensors,photomake_mode,phot
pipe.fuse_lora(adapter_names=[trigger_words, ], lora_scale=lora_scale)
elif model_type == "img2img":
if photomake_mode=="v1":
if photo_vesion== "v1" :
if single_files:
# print("loading from a single_files")
try:
@@ -77,13 +72,10 @@ def load_models(path,model_type,single_files,use_safetensors,photomake_mode,phot
path,config=add_config, original_config=original_config_file,
torch_dtype=torch.float16, use_safetensors=use_safetensors)
except:
try:
pipe = PhotoMakerStableDiffusionXLPipeline.from_single_file(
pipe = PhotoMakerStableDiffusionXLPipeline.from_single_file(
path, config=add_config,original_config_file=original_config_file,
torch_dtype=torch.float16, use_safetensors=use_safetensors
)
except:
raise "load pipe error"
else:
pipe = PhotoMakerStableDiffusionXLPipeline.from_pretrained(
@@ -104,20 +96,17 @@ def load_models(path,model_type,single_files,use_safetensors,photomake_mode,phot
path,config=add_config, original_config=original_config_file,
torch_dtype=torch.float16, use_safetensors=use_safetensors)
except:
try:
pipe = PhotoMakerStableDiffusionXLPipelineV2.from_single_file(
pipe = PhotoMakerStableDiffusionXLPipelineV2.from_single_file(
path, config=add_config,original_config_file=original_config_file,
torch_dtype=torch.float16, use_safetensors=use_safetensors
)
except:
raise "load pipe error"
else:
pipe = PhotoMakerStableDiffusionXLPipelineV2.from_pretrained(
path, torch_dtype=torch.float16, use_safetensors=use_safetensors
)
# define the trigger word
pipe.load_photomaker_adapter(
photomaker_path,
"F:/test/ComfyUI/models/photomaker/photomaker-v2.bin",
subfolder="",
weight_name="photomaker-v2.bin",
trigger_word="img",
@@ -133,6 +122,7 @@ def load_models(path,model_type,single_files,use_safetensors,photomake_mode,phot
pipe.fuse_lora(adapter_names=[trigger_words, ], lora_scale=lora_scale)
else:
raise f"using{model_type}node,must choice{model_type}type in model_loader node"
print(f"using{model_type}node,must choice{model_type}type in model_loader node")
raise "load error"
return pipe
utils/model_v2.py
+7 -1
View File
@@ -47,7 +47,7 @@ class QFormerPerceiver(nn.Module):
self.num_tokens = num_tokens
self.cross_attention_dim = cross_attention_dim
self.use_residual = use_residual
print(cross_attention_dim*num_tokens)
#print(cross_attention_dim*num_tokens)
self.token_proj = nn.Sequential(
nn.Linear(id_embeddings_dim, id_embeddings_dim*ratio),
nn.GELU(),
@@ -82,6 +82,7 @@ class FuseModule(nn.Module):
self.layer_norm = nn.LayerNorm(embed_dim)
def fuse_fn(self, prompt_embeds, id_embeds):
stacked_id_embeds = torch.cat([prompt_embeds, id_embeds], dim=-1)
stacked_id_embeds = self.mlp1(stacked_id_embeds) + prompt_embeds
stacked_id_embeds = self.mlp2(stacked_id_embeds)
@@ -95,6 +96,7 @@ class FuseModule(nn.Module):
class_tokens_mask,
) -> torch.Tensor:
# id_embeds shape: [b, max_num_inputs, 1, 2048]
id_embeds = id_embeds.to(prompt_embeds.dtype)
num_inputs = class_tokens_mask.sum().unsqueeze(0) # TODO: check for training case
batch_size, max_num_inputs = id_embeds.shape[:2]
@@ -112,10 +114,13 @@ class FuseModule(nn.Module):
valid_id_embeds = flat_id_embeds[valid_id_mask.flatten()]
prompt_embeds = prompt_embeds.view(-1, prompt_embeds.shape[-1])
class_tokens_mask = class_tokens_mask.view(-1)
valid_id_embeds = valid_id_embeds.view(-1, valid_id_embeds.shape[-1])
# slice out the image token embeddings
image_token_embeds = prompt_embeds[class_tokens_mask]
stacked_id_embeds = self.fuse_fn(image_token_embeds, valid_id_embeds)
assert class_tokens_mask.sum() == stacked_id_embeds.shape[0], f"{class_tokens_mask.sum()} != {stacked_id_embeds.shape[0]}"
prompt_embeds.masked_scatter_(class_tokens_mask[:, None], stacked_id_embeds.to(prompt_embeds.dtype))
@@ -148,6 +153,7 @@ class PhotoMakerIDEncoder_CLIPInsightfaceExtendtoken(CLIPVisionModelWithProjecti
id_embeds = self.qformer_perceiver(id_embeds, last_hidden_state)
id_embeds = id_embeds.view(b, num_inputs, self.num_tokens, -1)
updated_prompt_embeds = self.fuse_module(prompt_embeds, id_embeds, class_tokens_mask)
return updated_prompt_embeds
utils/pipeline_controlnet_sdxl_wrapper.py
+1611
View File
File diff suppressed because it is too large Load Diff
utils/pipeline_img.py
+658
View File
@@ -0,0 +1,658 @@
from typing import Any, Callable, Dict, List, Optional, Union, Tuple
from collections import OrderedDict
import os
import PIL
import numpy as np
import torch
from torchvision import transforms as T
from safetensors import safe_open
from huggingface_hub.utils import validate_hf_hub_args
from transformers import CLIPImageProcessor, CLIPTokenizer
from diffusers import StableDiffusionXLPipeline
try:
from diffusers.pipelines.stable_diffusion_xl.pipeline_stable_diffusion_xl import StableDiffusionXLPipelineOutput,rescale_noise_cfg
except:
try:
from diffusers.pipelines.stable_diffusion_xl import StableDiffusionXLPipelineOutput,rescale_noise_cfg
except:
try:
from diffusers.pipelines.stable_diffusion_xl.pipeline_output import StableDiffusionXLPipelineOutput,rescale_noise_cfg
except:
raise "import StableDiffusionXLPipelineOutput,rescale_noise_cfg error"
from diffusers.utils import (
_get_model_file,
is_transformers_available,
logging,
)
from . import PhotoMakerIDEncoder
PipelineImageInput = Union[
PIL.Image.Image,
torch.FloatTensor,
List[PIL.Image.Image],
List[torch.FloatTensor],
]
class PhotoMakerStableDiffusionXLPipeline(StableDiffusionXLPipeline):
@validate_hf_hub_args
def load_photomaker_adapter(
self,
pretrained_model_name_or_path_or_dict: Union[str, Dict[str, torch.Tensor]],
weight_name: str,
subfolder: str = '',
trigger_word: str = 'img',
**kwargs,
):
"""
Parameters:
pretrained_model_name_or_path_or_dict (`str` or `os.PathLike` or `dict`):
Can be either:
- A string, the *model id* (for example `google/ddpm-celebahq-256`) of a pretrained model hosted on
the Hub.
- A path to a *directory* (for example `./my_model_directory`) containing the model weights saved
with [`ModelMixin.save_pretrained`].
- A [torch state
dict](https://pytorch.org/tutorials/beginner/saving_loading_models.html#what-is-a-state-dict).
weight_name (`str`):
The weight name NOT the path to the weight.
subfolder (`str`, defaults to `""`):
The subfolder location of a model file within a larger model repository on the Hub or locally.
trigger_word (`str`, *optional*, defaults to `"img"`):
The trigger word is used to identify the position of class word in the text prompt,
and it is recommended not to set it as a common word.
This trigger word must be placed after the class word when used, otherwise, it will affect the performance of the personalized generation.
"""
# Load the main state dict first.
cache_dir = kwargs.pop("cache_dir", None)
force_download = kwargs.pop("force_download", False)
resume_download = kwargs.pop("resume_download", False)
proxies = kwargs.pop("proxies", None)
local_files_only = kwargs.pop("local_files_only", None)
token = kwargs.pop("token", None)
revision = kwargs.pop("revision", None)
user_agent = {
"file_type": "attn_procs_weights",
"framework": "pytorch",
}
if not isinstance(pretrained_model_name_or_path_or_dict, dict):
try:
model_file = _get_model_file(
pretrained_model_name_or_path_or_dict,
weights_name=weight_name,
cache_dir=cache_dir,
force_download=force_download,
proxies=proxies,
local_files_only=local_files_only,
token=token,
revision=revision,
subfolder=subfolder,
user_agent=user_agent,
)
except:
try:
model_file = _get_model_file(
pretrained_model_name_or_path_or_dict,
weights_name=weight_name,
cache_dir=cache_dir,
force_download=force_download,
resume_download=resume_download,
proxies=proxies,
local_files_only=local_files_only,
token=token,
revision=revision,
subfolder=subfolder,
user_agent=user_agent,)
except:
pass
if weight_name.endswith(".safetensors"):
state_dict = {"id_encoder": {}, "lora_weights": {}}
with safe_open(model_file, framework="pt", device="cpu") as f:
for key in f.keys():
if key.startswith("id_encoder."):
state_dict["id_encoder"][key.replace("id_encoder.", "")] = f.get_tensor(key)
elif key.startswith("lora_weights."):
state_dict["lora_weights"][key.replace("lora_weights.", "")] = f.get_tensor(key)
else:
state_dict = torch.load(model_file, map_location="cpu")
else:
state_dict = pretrained_model_name_or_path_or_dict
keys = list(state_dict.keys())
if keys != ["id_encoder", "lora_weights"]:
raise ValueError("Required keys are (`id_encoder` and `lora_weights`) missing from the state dict.")
self.trigger_word = trigger_word
# load finetuned CLIP image encoder and fuse module here if it has not been registered to the pipeline yet
print(f"Loading PhotoMaker components [1] id_encoder from [{pretrained_model_name_or_path_or_dict}]...")
id_encoder = PhotoMakerIDEncoder()
id_encoder.load_state_dict(state_dict["id_encoder"], strict=False)
device = (
"cuda"
if torch.cuda.is_available()
else "mps" if torch.backends.mps.is_available() else "cpu"
)
if device!="mps":
id_encoder = id_encoder.to( "cuda", dtype=self.unet.dtype)
else:
id_encoder = id_encoder.to(device)
self.id_encoder = id_encoder
self.id_image_processor = CLIPImageProcessor()
# load lora into models
print(f"Loading PhotoMaker components [2] lora_weights from [{pretrained_model_name_or_path_or_dict}]")
self.load_lora_weights(state_dict["lora_weights"], adapter_name="photomaker")
self.tokenizer_2=self.tokenizer_2
self.tokenizer =self.tokenizer
# Add trigger word token
if self.tokenizer is not None:
self.tokenizer.add_tokens([self.trigger_word], special_tokens=True)
self.tokenizer_2.add_tokens([self.trigger_word], special_tokens=True)
def encode_prompt_with_trigger_word(
self,
prompt: str,
prompt_2: Optional[str] = None,
num_id_images: int = 1,
device: Optional[torch.device] = None,
prompt_embeds: Optional[torch.FloatTensor] = None,
pooled_prompt_embeds: Optional[torch.FloatTensor] = None,
class_tokens_mask: Optional[torch.LongTensor] = None,
nc_flag: bool = False,
):
device = device or self._execution_device
print(f"text_encoder_2 device: {self.text_encoder_2.device}") # 检查设备是否为 meta # cpu
print(f"text_encoder_2 param is meta: {next(self.text_encoder_2.parameters()).is_meta}") #False
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 = self.tokenizer_2.convert_tokens_to_ids(self.trigger_word)
# Define tokenizers and text encoders
tokenizers = [self.tokenizer, self.tokenizer_2] if self.tokenizer is not None else [self.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_encoder in zip(prompts, tokenizers, text_encoders):
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: {self.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: {self.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))]
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 = text_encoder(
clean_input_ids.to(device),
output_hidden_states=True,
)
# We are only ALWAYS interested in the pooled output of the final text encoder
pooled_prompt_embeds = prompt_embeds[0]
prompt_embeds = prompt_embeds.hidden_states[-2]
prompt_embeds_list.append(prompt_embeds)
prompt_embeds = torch.concat(prompt_embeds_list, dim=-1)
print(prompt_embeds.shape,1234) #torch.Size([1, 77, 2048])
print(pooled_prompt_embeds.shape,pooled_prompt_embeds)#([1, 1280])
prompt_embeds = prompt_embeds.to(device=device,dtype=self.text_encoder_2.dtype, )
class_tokens_mask = class_tokens_mask.to(device=device) # TODO: ignoring two-prompt case
return prompt_embeds, pooled_prompt_embeds, class_tokens_mask
@property
def interrupt(self):
return self._interrupt
@torch.no_grad()
def __call__(
self,
prompt: Union[str, List[str]] = None,
prompt_2: Optional[Union[str, List[str]]] = None,
height: Optional[int] = None,
width: Optional[int] = None,
num_inference_steps: int = 50,
denoising_end: Optional[float] = None,
guidance_scale: float = 5.0,
negative_prompt: Optional[Union[str, List[str]]] = None,
negative_prompt_2: Optional[Union[str, List[str]]] = None,
num_images_per_prompt: Optional[int] = 1,
eta: float = 0.0,
generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
latents: Optional[torch.FloatTensor] = None,
prompt_embeds: Optional[torch.FloatTensor] = None,
negative_prompt_embeds: Optional[torch.FloatTensor] = None,
pooled_prompt_embeds: Optional[torch.FloatTensor] = None,
negative_pooled_prompt_embeds: Optional[torch.FloatTensor] = None,
output_type: Optional[str] = "latent",
return_dict: bool = True,
cross_attention_kwargs: Optional[Dict[str, Any]] = None,
guidance_rescale: float = 0.0,
original_size: Optional[Tuple[int, int]] = None,
crops_coords_top_left: Tuple[int, int] = (0, 0),
target_size: Optional[Tuple[int, int]] = None,
callback: Optional[Callable[[int, int, torch.FloatTensor], None]] = None,
callback_steps: int = 1,
callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None,
callback_on_step_end_tensor_inputs: List[str] = ["latents"],
# Added parameters (for PhotoMaker)
input_id_images: PipelineImageInput = None,
start_merge_step: int = 0, # TODO: change to `style_strength_ratio` in the future
class_tokens_mask: Optional[torch.LongTensor] = None,
prompt_embeds_text_only: Optional[torch.FloatTensor] = None,
pooled_prompt_embeds_text_only: Optional[torch.FloatTensor] = None,
nc_flag = False,
emb_dict={},
):
r"""
Function invoked when calling the pipeline for generation.
Only the parameters introduced by PhotoMaker are discussed here.
For explanations of the previous parameters in StableDiffusionXLPipeline, please refer to https://github.com/huggingface/diffusers/blob/v0.25.0/src/diffusers/pipelines/stable_diffusion_xl/pipeline_stable_diffusion_xl.py
Args:
input_id_images (`PipelineImageInput`, *optional*):
Input ID Image to work with PhotoMaker.
class_tokens_mask (`torch.LongTensor`, *optional*):
Pre-generated class token. When the `prompt_embeds` parameter is provided in advance, it is necessary to prepare the `class_tokens_mask` beforehand for marking out the position of class word.
prompt_embeds_text_only (`torch.FloatTensor`, *optional*):
Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not
provided, text embeddings will be generated from `prompt` input argument.
pooled_prompt_embeds_text_only (`torch.FloatTensor`, *optional*):
Pre-generated pooled text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting.
If not provided, pooled text embeddings will be generated from `prompt` input argument.
Returns:
[`~pipelines.stable_diffusion_xl.StableDiffusionXLPipelineOutput`] or `tuple`:
[`~pipelines.stable_diffusion_xl.StableDiffusionXLPipelineOutput`] if `return_dict` is True, otherwise a
`tuple`. When returning a tuple, the first element is a list with the generated images.
"""
# 0. Default height and width to unet
height = height or self.unet.config.sample_size * self.vae_scale_factor
width = width or self.unet.config.sample_size * self.vae_scale_factor
original_size = original_size or (height, width)
target_size = target_size or (height, width)
# 1. Check inputs. Raise error if not correct
# self.check_inputs(
# prompt,
# prompt_2,
# height,
# width,
# callback_steps,
# negative_prompt,
# negative_prompt_2,
# prompt_embeds,
# negative_prompt_embeds,
# pooled_prompt_embeds,
# negative_pooled_prompt_embeds,
# callback_on_step_end_tensor_inputs,
# )
self._interrupt = False
#
if prompt_embeds is not None and class_tokens_mask is None:
raise ValueError(
"If `prompt_embeds` are provided, `class_tokens_mask` also have to be passed. Make sure to generate `class_tokens_mask` from the same tokenizer that was used to generate `prompt_embeds`."
)
# check the input id images
if input_id_images is None:
raise ValueError(
"Provide `input_id_images`. Cannot leave `input_id_images` undefined for PhotoMaker pipeline."
)
if not isinstance(input_id_images, list):
input_id_images = [input_id_images]
device = self._execution_device
# here `guidance_scale` is defined analog to the guidance weight `w` of equation (2)
# of the Imagen paper: https://arxiv.org/pdf/2205.11487.pdf . `guidance_scale = 1`
# corresponds to doing no classifier free guidance.
do_classifier_free_guidance = guidance_scale >= 1.0
assert do_classifier_free_guidance
# 3. Encode input prompt
# num_id_images = len(input_id_images)
# if isinstance(prompt, list):
# prompt_arr = prompt
# 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,
# ) = self.encode_prompt_with_trigger_word(
# 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 = self.tokenizer.encode(prompt, add_special_tokens=False)
# trigger_word_token = self.tokenizer.convert_tokens_to_ids(self.trigger_word)
# if not nc_flag:
# tokens_text_only.remove(trigger_word_token)
# prompt_text_only = self.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,
# ) = self.encode_prompt(
# 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=negative_prompt_2,
# prompt_embeds=prompt_embeds_text_only,
# negative_prompt_embeds=negative_prompt_embeds,
# pooled_prompt_embeds=pooled_prompt_embeds_text_only,
# negative_pooled_prompt_embeds=negative_pooled_prompt_embeds,
# )
# # 5. Prepare the input ID images
# dtype = next(self.id_encoder.parameters()).dtype
# if not isinstance(input_id_images[0], torch.Tensor):
# id_pixel_values = self.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 = self.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
negative_prompt_embeds_arr=emb_dict['negative_prompt_embeds_arr']
prompt_embeds_arr=emb_dict['prompt_embeds_arr']
pooled_prompt_embeds_arr=emb_dict['pooled_prompt_embeds_arr']
prompt_embeds_text_only_arr=emb_dict['prompt_embeds_text_only_arr']
pooled_prompt_embeds_text_only_arr=emb_dict['pooled_prompt_embeds_text_only_arr']
negative_pooled_prompt_embeds_arr=emb_dict['negative_pooled_prompt_embeds_arr']
negative_prompt_embeds=emb_dict['negative_pooled_prompt_embeds']
pooled_prompt_embeds=emb_dict['pooled_prompt_embeds']
negative_pooled_prompt_embeds=emb_dict['negative_pooled_prompt_embeds']
prompt_embeds=emb_dict['prompt_embeds']
# 2. Define call parameters
if prompt is not None and isinstance(prompt, str):
batch_size = 1
prompt = [prompt]
elif prompt is not None and isinstance(prompt, list):
batch_size = len(prompt)
else:
batch_size = emb_dict["batch_size"]
# 7. Prepare timesteps
self.scheduler.set_timesteps(num_inference_steps, device=device)
timesteps = self.scheduler.timesteps
negative_prompt_embeds = torch.cat(negative_prompt_embeds_arr ,dim =0)
#print(negative_prompt_embeds.shape)
if not nc_flag:
prompt_embeds = torch.cat(prompt_embeds_arr ,dim = 0)
# print(prompt_embeds.shape)
pooled_prompt_embeds = torch.cat(pooled_prompt_embeds_arr,dim = 0)
#print(pooled_prompt_embeds.shape)
prompt_embeds_text_only = torch.cat(prompt_embeds_text_only_arr ,dim = 0)
#print(prompt_embeds_text_only.shape)
pooled_prompt_embeds_text_only = torch.cat(pooled_prompt_embeds_text_only_arr ,dim = 0)
#print(pooled_prompt_embeds_text_only.shape)
negative_pooled_prompt_embeds = torch.cat(negative_pooled_prompt_embeds_arr ,dim = 0)
#print(negative_pooled_prompt_embeds.shape)
# 8. Prepare latent variables
num_channels_latents = self.unet.config.in_channels
latents = self.prepare_latents(
batch_size * num_images_per_prompt,
num_channels_latents,
height,
width,
prompt_embeds.dtype if not nc_flag else prompt_embeds_text_only.dtype,
device,
generator,
latents,
)
# 9. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline
extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta)
# 10. Prepare added time ids & embeddings
if self.text_encoder_2 is None:
text_encoder_projection_dim = int(pooled_prompt_embeds_text_only_arr[0].shape[-1])
else:
text_encoder_projection_dim = self.text_encoder_2.config.projection_dim
add_time_ids = self._get_add_time_ids(
original_size,
crops_coords_top_left,
target_size,
dtype=prompt_embeds.dtype if not nc_flag else prompt_embeds_text_only.dtype,
text_encoder_projection_dim=text_encoder_projection_dim,
)
add_time_ids = torch.cat([add_time_ids, add_time_ids], dim=0)
add_time_ids = add_time_ids.to(device).repeat(batch_size * num_images_per_prompt, 1)
#print(latents.shape)
#print(add_time_ids.shape)
# 11. Denoising loop
num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order
with self.progress_bar(total=num_inference_steps) as progress_bar:
for i, t in enumerate(timesteps):
if self.interrupt:
continue
latent_model_input = (
torch.cat([latents] * 2) if do_classifier_free_guidance else latents
)
latent_model_input = self.scheduler.scale_model_input(latent_model_input, t)
if nc_flag:
latent_model_input=latent_model_input.to(self.unet.dtype)
if i <= start_merge_step or nc_flag:
negative_prompt_embeds=negative_prompt_embeds.to(latent_model_input.device,dtype=self.unet.dtype)
prompt_embeds_text_only=prompt_embeds_text_only.to(latent_model_input.device,dtype=self.unet.dtype)
negative_pooled_prompt_embeds=negative_pooled_prompt_embeds.to(latent_model_input.device,dtype=self.unet.dtype)
pooled_prompt_embeds_text_only=pooled_prompt_embeds_text_only.to(latent_model_input.device,dtype=self.unet.dtype)
current_prompt_embeds = torch.cat(
[negative_prompt_embeds, prompt_embeds_text_only], dim=0
)
add_text_embeds = torch.cat([negative_pooled_prompt_embeds, pooled_prompt_embeds_text_only], dim=0)
else:
negative_prompt_embeds=negative_prompt_embeds.to(latent_model_input.device,dtype=self.unet.dtype)
prompt_embeds=prompt_embeds.to(latent_model_input.device,dtype=self.unet.dtype)
negative_pooled_prompt_embeds=negative_pooled_prompt_embeds.to(latent_model_input.device,dtype=self.unet.dtype)
pooled_prompt_embeds=pooled_prompt_embeds.to(latent_model_input.device,dtype=self.unet.dtype)
current_prompt_embeds = torch.cat(
[negative_prompt_embeds, prompt_embeds], dim=0
)
add_text_embeds = torch.cat([negative_pooled_prompt_embeds, pooled_prompt_embeds], dim=0)
# predict the noise residual
add_text_embeds = add_text_embeds.to(latent_model_input.device,dtype=self.unet.dtype)
add_time_ids= add_time_ids.to(latent_model_input.device,dtype=self.unet.dtype)
added_cond_kwargs = {"text_embeds": add_text_embeds, "time_ids": add_time_ids}
# print(latent_model_input.shape)
# print(t)
# print(current_prompt_embeds.shape)
# print(add_text_embeds.shape)
# print(add_time_ids.shape)
#zeros_matrix =
#global_mask1024 = torch.cat([torch.randn(1, 1024, 1, 1, device=device) for random_number])
#global_mask4096 =
current_prompt_embeds=current_prompt_embeds.to(latent_model_input.device,dtype=self.unet.dtype)
noise_pred = self.unet(
latent_model_input,
t,
encoder_hidden_states=current_prompt_embeds,
cross_attention_kwargs=cross_attention_kwargs,
added_cond_kwargs=added_cond_kwargs,
return_dict=False,
)[0]
# print(noise_pred.shape)
# perform guidance
if do_classifier_free_guidance:
noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)
if do_classifier_free_guidance and guidance_rescale > 0.0:
# Based on 3.4. in https://arxiv.org/pdf/2305.08891.pdf
noise_pred = rescale_noise_cfg(noise_pred, noise_pred_text, guidance_rescale=guidance_rescale)
# compute the previous noisy sample x_t -> x_t-1
latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs, return_dict=False)[0]
if callback_on_step_end is not None:
callback_kwargs = {}
for k in callback_on_step_end_tensor_inputs:
callback_kwargs[k] = locals()[k]
ck_outputs = callback_on_step_end(self, i, t, callback_kwargs)
latents = callback_outputs.pop("latents", latents)
prompt_embeds = callback_outputs.pop("prompt_embeds", prompt_embeds)
negative_prompt_embeds = callback_outputs.pop("negative_prompt_embeds", negative_prompt_embeds)
add_text_embeds = callback_outputs.pop("add_text_embeds", add_text_embeds)
# negative_pooled_prompt_embeds = callback_outputs.pop(
# "negative_pooled_prompt_embeds", negative_pooled_prompt_embeds
# )
# add_time_ids = callback_outputs.pop("add_time_ids", add_time_ids)
# negative_add_time_ids = callback_outputs.pop("negative_add_time_ids", negative_add_time_ids)
# call the callback, if provided
if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0):
progress_bar.update()
if callback is not None and i % callback_steps == 0:
step_idx = i // getattr(self.scheduler, "order", 1)
callback(step_idx, t, latents)
# make sure the VAE is in float32 mode, as it overflows in float16
if self.vae.dtype == torch.float16 and self.vae.config.force_upcast:
self.upcast_vae()
latents = latents.to(next(iter(self.vae.post_quant_conv.parameters())).dtype)
if not output_type == "latent":
image = self.vae.decode(latents / self.vae.config.scaling_factor, return_dict=False)[0]
else:
image = latents/self.vae.config.scaling_factor
return StableDiffusionXLPipelineOutput(images=image)
# apply watermark if available
# if self.watermark is not None:
# image = self.watermark.apply_watermark(image)
image = self.image_processor.postprocess(image, output_type=output_type)
# Offload all models
self.maybe_free_model_hooks()
if not return_dict:
return (image,)
return StableDiffusionXLPipelineOutput(images=image)
utils/pipeline_photho_wrapper.py
+1340
View File
File diff suppressed because it is too large Load Diff
utils/pipeline_v2.py
+4 -6
View File
@@ -406,9 +406,7 @@ class PhotoMakerStableDiffusionXLPipeline(StableDiffusionXLPipeline):
# 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 * self.num_tokens + \
clean_input_ids[class_token_index + 1:]
clean_input_ids = clean_input_ids[:class_token_index] + [class_token] * num_id_images * self.num_tokens + clean_input_ids[class_token_index + 1:]
# Truncation or padding
max_len = tokenizer.model_max_length
@@ -419,9 +417,7 @@ class PhotoMakerStableDiffusionXLPipeline(StableDiffusionXLPipeline):
max_len - len(clean_input_ids)
)
class_tokens_mask = [
True if class_token_index <= i < class_token_index + (num_id_images * self.num_tokens) 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 * self.num_tokens) else False for i in range(len(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)
@@ -800,6 +796,8 @@ class PhotoMakerStableDiffusionXLPipeline(StableDiffusionXLPipeline):
if not nc_flag: #v2必须要有id_embeds
if id_embeds is not None:
id_embeds = id_embeds.unsqueeze(0).to(device=device, dtype=dtype)
#print(id_pixel_values.shape, id_embeds.shape, prompt_embeds.shape, class_tokens_mask.shape) #torch.Size([1, 1, 3, 224, 224]) torch.Size([1, 512]) torch.Size([1, 77, 2048]) torch.Size([1, 77])
prompt_embeds = self.id_encoder(id_pixel_values, prompt_embeds, class_tokens_mask, id_embeds)
else:
prompt_embeds = self.id_encoder(id_pixel_values, prompt_embeds, class_tokens_mask)
utils/pipeline_v2_warrper.py
+1125
View File
File diff suppressed because it is too large Load Diff
utils/pipeline_wrapper.py
+1340
View File
File diff suppressed because it is too large Load Diff