smthemex
+3
-3
@@ -201,9 +201,9 @@ def load_flow_model_quintized(name: str, ckpt_path,quantized_mode,aggressive_off
|
||||
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:
|
||||
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
|
||||
|
||||
@@ -3,14 +3,18 @@
|
||||
Using StoryDiffusion and other methods to make storys in ComfyUI
|
||||
----
|
||||
|
||||
* Origin from [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)
|
||||
* Origin from [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), [UNO](https://github.com/bytedance/UNO)
|
||||
|
||||
|
||||
## Updates:
|
||||
* 2025/04/06
|
||||
* 2025/04/11
|
||||
* Add UNO support，Only the single FLUX model (27G) and UNO's Lora are needed. Please enable FP8 quantization and use storydiffusionw_flowjson workflow testing ，fix a bug，
|
||||
* 新增UNO支持，只需要单体FLUX模型(27G)和UNO的lora，请开启fp8量化和使用storydiffusion_workflow.json工作流测试,修复tokens过长的bug;
|
||||
|
||||
* previous
|
||||
* Add infinite svdq v0.2 support,it'work well when your svdq update v0.2，[download wheel](https://huggingface.co/mit-han-lab/nunchaku/tree/main) 更新 svdq v0.2的支持，infinite工作正常，[轮子](https://huggingface.co/mit-han-lab/nunchaku/tree/main)下载地址。
|
||||
* 2025/04/05
|
||||
* 1.修改了模型加载的流程，更新到V2版本，如果你喜欢旧的，可以下载V1.0版本的,2.请使用example.json，它集成了主要的工作流;3.剔除掉一些过时的功能;
|
||||
* 1.Modified the model loading process.Update to V2 version, If you like the old one, you can download version 1.0，2.Please use 'example.json', which integrates the main workflow，3.Remove some outdated features;
|
||||
* 1.修改了模型加载的流程，更新到V2版本，如果你喜欢旧的，可以下载V1.0版本的,2.请使用storydiffusion_workflow.json，它集成了主要的工作流;3.剔除掉一些过时的功能;
|
||||
* 1.Modified the model loading process.Update to V2 version, If you like the old one, you can download version 1.0，2.Please use 'storydiffusion_workflow.json', which integrates the main workflow，3.Remove some outdated features;
|
||||
|
||||
|
||||
1.Installation
|
||||
@@ -165,6 +169,13 @@ or
|
||||
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'
|
||||
|
||||
|
||||
**3.7 UNO mode**
|
||||
download lora [dit_lora.safetensor](https://huggingface.co/bytedance-research/UNO/tree/main),use fp8,if Vram <24.
|
||||
```
|
||||
├── ComfyUI/models/
|
||||
| ├── diffusion_models/flux1-dev.safetensors #
|
||||
| ├── loras/dit_lora.safetensors #
|
||||
```
|
||||
|
||||
4 Example
|
||||
----
|
||||
@@ -206,7 +217,10 @@ or
|
||||
* svdq，升级到v.2工作正常
|
||||
<img src="https://github.com/smthemex/ComfyUI_StoryDiffusion/blob/main/images/infinite_svdqv2.png" width="50%">
|
||||
|
||||
**4.8 comfyUI classic（comfyUI经典模式，可以接任意适配CF的流程，主要是方便使用多角色的clip）**
|
||||
**4.8 UNO**
|
||||
<img src="https://github.com/smthemex/ComfyUI_StoryDiffusion/blob/main/images/UNO.jpg" width="50%">
|
||||
|
||||
**4.9 comfyUI classic（comfyUI经典模式，可以接任意适配CF的流程，主要是方便使用多角色的clip）**
|
||||
* any mode SD1.5 SDXL SD3.5 FLUX...
|
||||
<img src="https://github.com/smthemex/ComfyUI_StoryDiffusion/blob/main/images/comfyui_classic.png" width="50%">
|
||||
|
||||
@@ -306,3 +320,12 @@ svdquant
|
||||
```
|
||||
[GGUF](https://github.com/city96/ComfyUI-GGUF)
|
||||
[FLUX LICENSE](https://huggingface.co/black-forest-labs/FLUX.1-dev/blob/main/LICENSE.md)
|
||||
|
||||
```
|
||||
@article{wu2025less,
|
||||
title={Less-to-More Generalization: Unlocking More Controllability by In-Context Generation},
|
||||
author={Wu, Shaojin and Huang, Mengqi and Wu, Wenxu and Cheng, Yufeng and Ding, Fei and He, Qian},
|
||||
journal={arXiv preprint arXiv:2504.02160},
|
||||
year={2025}
|
||||
}
|
||||
```
|
||||
+141
-56
@@ -11,9 +11,10 @@ import folder_paths
|
||||
from comfy.model_management import total_vram
|
||||
import comfy
|
||||
import latent_preview
|
||||
import torchvision.transforms.functional as TVF
|
||||
from .utils.utils import get_comic
|
||||
from .model_loader_utils import (phi2narry,Loader_storydiffusion,Loader_Flux_Pulid,load_pipeline_consistory,Loader_KOLOR,replicate_data_by_indices,glm_single_encode,
|
||||
get_float,gc_cleanup,tensor_to_image,photomaker_clip,tensortopil_list,tensortopil_list_upscale,
|
||||
get_float,gc_cleanup,tensor_to_image,photomaker_clip,Loader_UNO,tensortopil_list_upscale,tensortopil_list,
|
||||
narry_list_pil,pre_text2infer,cf_clip,get_phrases_idx_cf,get_eot_idx_cf,get_ms_phrase_emb,get_extra_function,photomaker_clip_v2,adjust_indices,
|
||||
get_scheduler,apply_style_positive,fitter_cf_model_type,Infer_MSdiffusion,Loader_story_maker,Loader_InfiniteYou,
|
||||
nomarl_upscale,SAMPLER_NAMES,SCHEDULER_NAMES,lora_lightning_list)
|
||||
@@ -412,8 +413,7 @@ class Pre_Translate_prompt:
|
||||
def translate_prompt(self, scene_prompts, keep_character_name):
|
||||
captions = scene_prompts.splitlines()
|
||||
if not keep_character_name:
|
||||
captions = [caption.split(")", 1)[-1] if ")" in caption else caption
|
||||
for caption in captions] # del character
|
||||
captions = [caption.split(")", 1)[-1] if ")" in caption else caption for caption in captions] # del character
|
||||
else:
|
||||
captions = [caption.replace("(", "").replace(")", "") if "(" or ")" in caption else caption
|
||||
for caption in captions] # del ()
|
||||
@@ -434,7 +434,7 @@ class EasyFunction_Lite:
|
||||
def INPUT_TYPES(s):
|
||||
return {"required": {"extra_repo": ("STRING", { "default": ""}),
|
||||
"checkpoints": (["none"] + folder_paths.get_filename_list("diffusion_models")+folder_paths.get_filename_list("gguf")+folder_paths.get_filename_list("clip"),),
|
||||
"clip_vision": (["none"] + folder_paths.get_filename_list("clip_vision"),),
|
||||
"clip_vision": (["none"] + folder_paths.get_filename_list("clip_vision")+folder_paths.get_filename_list("loras"),),
|
||||
"function_mode": (["none","tag", "clip","mask","infinite",],),
|
||||
"select_method":("STRING",{ "default":""}),
|
||||
"temperature": (
|
||||
@@ -448,7 +448,7 @@ class EasyFunction_Lite:
|
||||
|
||||
def easy_function_main(self, extra_repo,checkpoints,clip_vision,function_mode,select_method,temperature):
|
||||
use_gguf,use_unet=False,False
|
||||
model_path,clip=None,None
|
||||
model_path,clip,lora_path=None,None,None
|
||||
|
||||
if checkpoints != "none":
|
||||
if checkpoints.endswith(".gguf"):
|
||||
@@ -461,7 +461,12 @@ class EasyFunction_Lite:
|
||||
use_unet=True
|
||||
|
||||
if clip_vision != "none":
|
||||
clip_vision_path = folder_paths.get_full_path("clip_vision", clip_vision)
|
||||
if "clip_vision" in clip_vision:
|
||||
clip_vision_path = folder_paths.get_full_path("clip_vision", clip_vision)
|
||||
else:
|
||||
lora_path = folder_paths.get_full_path("loras", clip_vision)
|
||||
clip_vision_path = None
|
||||
|
||||
else:
|
||||
clip_vision_path=None
|
||||
if function_mode=="tag":
|
||||
@@ -478,24 +483,27 @@ class EasyFunction_Lite:
|
||||
|
||||
use_svdq=True if "svdq" in select_method else False
|
||||
|
||||
pipe={"model":model,"extra_repo":extra_repo,"model_path":model_path,"use_svdq":use_svdq,"use_gguf":use_gguf,"use_unet":use_unet,"extra_easy":function_mode,"select_method":select_method,"clip_vision_path":clip_vision_path}
|
||||
pipe={"model":model,"extra_repo":extra_repo,"model_path":model_path,"use_svdq":use_svdq,"use_gguf":use_gguf,"lora_ckpt_path":lora_path,
|
||||
"use_unet":use_unet,"extra_easy":function_mode,"select_method":select_method,"clip_vision_path":clip_vision_path}
|
||||
return (pipe,clip)
|
||||
|
||||
|
||||
class StoryDiffusion_Apply:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {"required": {
|
||||
"model": ("MODEL",),
|
||||
"vae": ("VAE",),
|
||||
"infer_mode": (["story", "classic","flux_pulid","infiniteyou","story_maker","story_and_maker","consistory","kolor_face","msdiffusion" ],),
|
||||
"photomake_ckpt": (["none"] + [i for i in folder_paths.get_filename_list("photomaker") if "v1" in i or "v2" in i],),
|
||||
"ipadapter_ckpt": (["none"] + folder_paths.get_filename_list("photomaker"),),
|
||||
"quantize_mode": ([ "fp8", "nf4","fp16", ],),
|
||||
"lora_scale": ("FLOAT", {"default": 0.8, "min": 0.1, "max": 1.0, "step": 0.1}),
|
||||
"extra_funtion":("STRING", {"default": ""}),
|
||||
},
|
||||
"optional":{"CLIP_VISION": ("CLIP_VISION",),
|
||||
return {"required":
|
||||
{
|
||||
"model": ("MODEL",),
|
||||
"vae": ("VAE",),
|
||||
"infer_mode": (["story", "classic","flux_pulid","infiniteyou","uno","story_maker","story_and_maker","consistory","kolor_face","msdiffusion" ],),
|
||||
"photomake_ckpt": (["none"] + [i for i in folder_paths.get_filename_list("photomaker") if "v1" in i or "v2" in i],),
|
||||
"ipadapter_ckpt": (["none"] + folder_paths.get_filename_list("photomaker"),),
|
||||
"quantize_mode": ([ "fp8", "nf4","fp16", ],),
|
||||
"lora_scale": ("FLOAT", {"default": 0.8, "min": 0.1, "max": 1.0, "step": 0.1}),
|
||||
"extra_funtion":("STRING", {"default": ""}),
|
||||
},
|
||||
"optional":{
|
||||
"CLIP_VISION": ("CLIP_VISION",),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -512,7 +520,7 @@ class StoryDiffusion_Apply:
|
||||
|
||||
CLIP_VISION=kwargs.get("CLIP_VISION")
|
||||
|
||||
|
||||
|
||||
unet_type=torch.float16 #use for sdxl
|
||||
image_proj_model=None
|
||||
|
||||
@@ -526,6 +534,8 @@ class StoryDiffusion_Apply:
|
||||
clip_vision_path=None
|
||||
repo=None
|
||||
|
||||
|
||||
save_quantezed=True if "save" in extra_funtion and quantize_mode=="fp8" else False
|
||||
|
||||
print(f"infer model type is {cf_model_type}")
|
||||
|
||||
@@ -543,27 +553,24 @@ class StoryDiffusion_Apply:
|
||||
if infer_mode in ["story_maker" ,"story_and_maker"] and not CLIP_VISION and ipadapter_ckpt_path is None:
|
||||
raise "story_maker need a clipvison H model,mask.bin"
|
||||
|
||||
|
||||
# check vram only using in flux pulid or UNO
|
||||
if total_vram > 45000.0:
|
||||
aggressive_offload = False
|
||||
offload = False
|
||||
elif 17000.0 < total_vram < 45000.0:
|
||||
aggressive_offload = False
|
||||
offload = True
|
||||
else:
|
||||
aggressive_offload = True
|
||||
offload = True
|
||||
|
||||
logging.info(f"total_vram is {total_vram},aggressive_offload is {aggressive_offload},offload is {offload}")
|
||||
|
||||
if infer_mode in["story", "story_and_maker","msdiffusion"]:# mix mode,use maker or ms to make 2 roles in on image
|
||||
model = Loader_storydiffusion(model,photomake_ckpt_path,vae)
|
||||
elif infer_mode =="story_maker":
|
||||
model = Loader_story_maker(model,ipadapter_ckpt_path,vae,False,lora_scale)
|
||||
elif infer_mode == "flux_pulid":
|
||||
|
||||
# check vram only using in flux pulid
|
||||
if total_vram > 45000.0:
|
||||
aggressive_offload = False
|
||||
offload = False
|
||||
elif 17000.0 < total_vram < 45000.0:
|
||||
aggressive_offload = False
|
||||
offload = True
|
||||
else:
|
||||
aggressive_offload = True
|
||||
offload = True
|
||||
|
||||
logging.info(f"total_vram is {total_vram},aggressive_offload is {aggressive_offload},offload is {offload}")
|
||||
|
||||
from .PuLID.app_flux import get_models
|
||||
if isinstance(model,dict):
|
||||
ckpt_path=model.get("model_path")
|
||||
@@ -585,9 +592,10 @@ class StoryDiffusion_Apply:
|
||||
if repo is None:
|
||||
raise "EasyFunction_Lite node extra_repo must fill kolor repo"
|
||||
model = Loader_KOLOR(repo,clip_vision_path,ipadapter_ckpt_path)
|
||||
elif infer_mode == "uno":
|
||||
model = Loader_UNO(model,offload,quantize_mode,save_quantezed,lora_rank=512)
|
||||
else: # can not choice a mode
|
||||
print("infer use comfyui classic mode")
|
||||
|
||||
|
||||
story_img=True if photomake_ckpt_path and infer_mode in["story","story_maker","story_and_maker","msdiffusion"] else False
|
||||
model_=model if infer_mode=="flux_pulid" or story_img else None
|
||||
@@ -692,6 +700,22 @@ class StoryDiffusion_CLIPTextEncode:
|
||||
role_key_list=replicate_data_by_indices(role_list, index_char_1_list, index_char_2_list)
|
||||
role_key_list=[i for i in role_key_list if i is not None]
|
||||
|
||||
if len(role_list)>1: #重新整理prompt排序，便于后续emb和ID的对应
|
||||
nc_dual_list=[]
|
||||
for i in nc_indexs:
|
||||
nc_dual_list.append(i)
|
||||
for i in positions_index_dual:
|
||||
nc_dual_list.append(i)
|
||||
adjusted_a = adjust_indices(index_char_1_list, nc_dual_list)
|
||||
adjusted_b = adjust_indices(index_char_2_list, nc_dual_list)
|
||||
a_list = [only_role_list[i] for i in adjusted_a if 0 <= i < len(only_role_list)]
|
||||
b_list = [only_role_list[i] for i in adjusted_b if 0 <= i < len(only_role_list)]
|
||||
|
||||
inf_list_split=[a_list,b_list]
|
||||
else:
|
||||
inf_list_split=only_role_list
|
||||
|
||||
|
||||
# if prompts_dual and infer_mode not in ["story_and_maker","story_maker","msdiffusion"]:
|
||||
# raise "only support prompts_dual role in story maker and msdiffusion"
|
||||
|
||||
@@ -710,10 +734,10 @@ class StoryDiffusion_CLIPTextEncode:
|
||||
if infer_mode=="msdiffusion" and photomake_ckpt_path and not img2img_mode:
|
||||
raise "if use msdiffusion txt2img mode,can not use photomake_ckpt_path"
|
||||
|
||||
uno_pe="d" #Literal['d', 'h', 'w', 'o']
|
||||
|
||||
only_role_emb_ne,x_1_refs_dict,image_emb=None,None,None
|
||||
|
||||
|
||||
only_role_emb_ne=None
|
||||
image_emb = None
|
||||
if img2img_mode:
|
||||
if infer_mode == "story_and_maker" or infer_mode == "story_maker":
|
||||
k1, _, _, _ = image.size()
|
||||
@@ -737,8 +761,53 @@ class StoryDiffusion_CLIPTextEncode:
|
||||
pass
|
||||
elif infer_mode == "infiniteyou":
|
||||
pass
|
||||
elif infer_mode == "uno":
|
||||
from .UNO.uno.flux.sampling import get_noise
|
||||
from .UNO.uno.flux.pipeline import preprocess_ref
|
||||
|
||||
ref_pil_list=tensortopil_list(image) #原方法需要图片缩放，先转回pil
|
||||
|
||||
if control_image is not None:
|
||||
control_pil_list=tensortopil_list(control_image)
|
||||
else:
|
||||
control_pil_list=None
|
||||
|
||||
if control_pil_list is None:# 无控制图时,默认时单图模式,x_1_refs用key取值
|
||||
x_1_refs_dict={}
|
||||
for key ,role_pil,prompts in zip(role_list,ref_pil_list,inf_list_split):
|
||||
role_tensor=phi2narry(preprocess_ref(role_pil, 512))#单图默认512
|
||||
x_1_refs_dict[key]=[[vae.encode(role_tensor[:,:,:,:3]).to(torch.bfloat16)]*len(prompts)] # {key:[[1],[2],key2:[[3],[4]]}
|
||||
else:
|
||||
if len(control_pil_list) != len(only_role_list):
|
||||
raise "control image must same size as prompts,多少句话就多少张图"
|
||||
else:
|
||||
if len(role_list)>1:
|
||||
nc_dual_list=[]
|
||||
for i in nc_indexs:
|
||||
nc_dual_list.append(i)
|
||||
for i in positions_index_dual:
|
||||
nc_dual_list.append(i)
|
||||
adjusted_a = adjust_indices(index_char_1_list, nc_dual_list)
|
||||
adjusted_b = adjust_indices(index_char_2_list, nc_dual_list)
|
||||
a_list = [control_pil_list[i] for i in adjusted_a if 0 <= i < len(control_pil_list)]
|
||||
b_list = [control_pil_list[i] for i in adjusted_b if 0 <= i < len(control_pil_list)]
|
||||
|
||||
control_pil_list=[a_list,b_list]
|
||||
else:
|
||||
control_pil_list=[control_pil_list]
|
||||
x_1_refs_dict={}
|
||||
for key ,role_pil,control_pil in zip(role_list,ref_pil_list,control_pil_list):
|
||||
mix_list=[]
|
||||
role_tensor=phi2narry(preprocess_ref(role_pil, 320) )#多图默认320
|
||||
for c in (control_pil):
|
||||
|
||||
c_tensor=phi2narry(preprocess_ref(c, 320))#多图默认320
|
||||
mix_list.append([vae.encode(role_tensor[:,:,:,:3]).to(torch.bfloat16),vae.encode(c_tensor[:,:,:,:3]).to(torch.bfloat16)])
|
||||
x_1_refs_dict[key]=mix_list #{key:[[1,2],[3,4]]}
|
||||
else:
|
||||
pass
|
||||
else:
|
||||
pass
|
||||
|
||||
# pre insight face model and emb
|
||||
if img2img_mode:
|
||||
@@ -775,23 +844,6 @@ class StoryDiffusion_CLIPTextEncode:
|
||||
image_list = [nomarl_upscale(img, width, height) for img in img_list]
|
||||
for index, key in enumerate(role_list):
|
||||
input_id_images_dict[key]=image_list[index]
|
||||
|
||||
|
||||
|
||||
if len(role_list)>1: #重新整理prompt排序，便于后续emb和ID的对应
|
||||
nc_dual_list=[]
|
||||
for i in nc_indexs:
|
||||
nc_dual_list.append(i)
|
||||
for i in positions_index_dual:
|
||||
nc_dual_list.append(i)
|
||||
adjusted_a = adjust_indices(index_char_1_list, nc_dual_list)
|
||||
adjusted_b = adjust_indices(index_char_2_list, nc_dual_list)
|
||||
a_list = [only_role_list[i] for i in adjusted_a if 0 <= i < len(only_role_list)]
|
||||
b_list = [only_role_list[i] for i in adjusted_b if 0 <= i < len(only_role_list)]
|
||||
|
||||
inf_list_split=[a_list,b_list]
|
||||
else:
|
||||
inf_list_split=only_role_list
|
||||
|
||||
#print("inf_list_split",inf_list_split)
|
||||
# pre clip txt emb
|
||||
@@ -817,8 +869,9 @@ class StoryDiffusion_CLIPTextEncode:
|
||||
|
||||
for key ,prompts in zip(role_list,inf_list_split):
|
||||
ip_emb,noise_,inp_n=[],[],[]
|
||||
seed_random = random.randint(0, MAX_SEED) #pulid的emb需要随机数
|
||||
|
||||
for p,n in zip(prompts,[neg_text]*len(prompts)):
|
||||
seed_random = random.randint(0, MAX_SEED) #pulid 和uno的emb需要随机数
|
||||
inp,inp_neg,x=get_emb_flux_pulid(t5_,clip_,if_repo,seed_random,p,n,width,height,num_steps=20,guidance=3.5,device=device)
|
||||
ip_emb.append(inp)
|
||||
inp_n.append(inp_neg)
|
||||
@@ -826,6 +879,20 @@ class StoryDiffusion_CLIPTextEncode:
|
||||
only_role_emb[key]=ip_emb
|
||||
noise_x[key]=noise_
|
||||
inp_neg_list[key]=inp_n
|
||||
elif infer_mode == "uno":
|
||||
only_role_emb={}
|
||||
from .UNO.uno.flux.sampling import prepare_multi_ip_wrapper
|
||||
for key ,prompts in zip(role_list,inf_list_split):
|
||||
ip_emb=[]
|
||||
|
||||
for p,x_1 in zip(prompts,x_1_refs_dict[key]):
|
||||
seed_random = random.randint(0, MAX_SEED) #pulid 和uno的emb需要随机数
|
||||
uno_x = get_noise(1, height, width, device=device,dtype=torch.bfloat16, seed=seed_random)
|
||||
inp = prepare_multi_ip_wrapper(clip,img=uno_x,prompt=p, ref_imgs=x_1, pe=uno_pe)
|
||||
ip_emb.append(inp)
|
||||
only_role_emb[key]=ip_emb
|
||||
|
||||
|
||||
elif infer_mode=="kolor_face":
|
||||
from .kolors.models.tokenization_chatglm import ChatGLMTokenizer
|
||||
tokenizer = ChatGLMTokenizer.from_pretrained(os.path.join(switch.get("repo"),'text_encoder'))
|
||||
@@ -927,6 +994,9 @@ class StoryDiffusion_CLIPTextEncode:
|
||||
elif infer_mode=="flux_pulid":
|
||||
postive_dict= {"role": only_role_emb, "nc": None, "daul": daul_emb} #不支持NC
|
||||
negative = [inp_neg_list,noise_x]
|
||||
elif infer_mode=="uno":
|
||||
postive_dict= {"role": only_role_emb, "nc": None, "daul": daul_emb} #TODO
|
||||
negative=None
|
||||
elif infer_mode=="kolor_face":
|
||||
postive_dict = {"role": only_role_emb, "nc": nc_emb, "daul": daul_emb}
|
||||
negative = only_role_emb_ne[0]
|
||||
@@ -1615,6 +1685,21 @@ class StoryDiffusion_KSampler:
|
||||
out = {}
|
||||
out["samples"] = torch.cat(samples_list, dim=0)
|
||||
return (out,)
|
||||
elif infer_mode=="uno":
|
||||
samples_list = []
|
||||
for key in role_list:
|
||||
for index,emb in enumerate(only_role_emb[key]):
|
||||
#print(emb)
|
||||
samples = model(width=width,
|
||||
height=height,guidance=cfg,
|
||||
num_steps=steps,
|
||||
inp_cond=emb,
|
||||
) # torch.Size([1, 4, 64, 64])
|
||||
|
||||
samples_list.append(samples)
|
||||
out = {}
|
||||
out["samples"] = torch.cat(samples_list, dim=0)
|
||||
return (out,)
|
||||
else: #none:
|
||||
return
|
||||
|
||||
|
||||
+201
@@ -0,0 +1,201 @@
|
||||
Apache License
|
||||
Version 2.0, January 2004
|
||||
http://www.apache.org/licenses/
|
||||
|
||||
TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
|
||||
|
||||
1. Definitions.
|
||||
|
||||
"License" shall mean the terms and conditions for use, reproduction,
|
||||
and distribution as defined by Sections 1 through 9 of this document.
|
||||
|
||||
"Licensor" shall mean the copyright owner or entity authorized by
|
||||
the copyright owner that is granting the License.
|
||||
|
||||
"Legal Entity" shall mean the union of the acting entity and all
|
||||
other entities that control, are controlled by, or are under common
|
||||
control with that entity. For the purposes of this definition,
|
||||
"control" means (i) the power, direct or indirect, to cause the
|
||||
direction or management of such entity, whether by contract or
|
||||
otherwise, or (ii) ownership of fifty percent (50%) or more of the
|
||||
outstanding shares, or (iii) beneficial ownership of such entity.
|
||||
|
||||
"You" (or "Your") shall mean an individual or Legal Entity
|
||||
exercising permissions granted by this License.
|
||||
|
||||
"Source" form shall mean the preferred form for making modifications,
|
||||
including but not limited to software source code, documentation
|
||||
source, and configuration files.
|
||||
|
||||
"Object" form shall mean any form resulting from mechanical
|
||||
transformation or translation of a Source form, including but
|
||||
not limited to compiled object code, generated documentation,
|
||||
and conversions to other media types.
|
||||
|
||||
"Work" shall mean the work of authorship, whether in Source or
|
||||
Object form, made available under the License, as indicated by a
|
||||
copyright notice that is included in or attached to the work
|
||||
(an example is provided in the Appendix below).
|
||||
|
||||
"Derivative Works" shall mean any work, whether in Source or Object
|
||||
form, that is based on (or derived from) the Work and for which the
|
||||
editorial revisions, annotations, elaborations, or other modifications
|
||||
represent, as a whole, an original work of authorship. For the purposes
|
||||
of this License, Derivative Works shall not include works that remain
|
||||
separable from, or merely link (or bind by name) to the interfaces of,
|
||||
the Work and Derivative Works thereof.
|
||||
|
||||
"Contribution" shall mean any work of authorship, including
|
||||
the original version of the Work and any modifications or additions
|
||||
to that Work or Derivative Works thereof, that is intentionally
|
||||
submitted to Licensor for inclusion in the Work by the copyright owner
|
||||
or by an individual or Legal Entity authorized to submit on behalf of
|
||||
the copyright owner. For the purposes of this definition, "submitted"
|
||||
means any form of electronic, verbal, or written communication sent
|
||||
to the Licensor or its representatives, including but not limited to
|
||||
communication on electronic mailing lists, source code control systems,
|
||||
and issue tracking systems that are managed by, or on behalf of, the
|
||||
Licensor for the purpose of discussing and improving the Work, but
|
||||
excluding communication that is conspicuously marked or otherwise
|
||||
designated in writing by the copyright owner as "Not a Contribution."
|
||||
|
||||
"Contributor" shall mean Licensor and any individual or Legal Entity
|
||||
on behalf of whom a Contribution has been received by Licensor and
|
||||
subsequently incorporated within the Work.
|
||||
|
||||
2. Grant of Copyright License. Subject to the terms and conditions of
|
||||
this License, each Contributor hereby grants to You a perpetual,
|
||||
worldwide, non-exclusive, no-charge, royalty-free, irrevocable
|
||||
copyright license to reproduce, prepare Derivative Works of,
|
||||
publicly display, publicly perform, sublicense, and distribute the
|
||||
Work and such Derivative Works in Source or Object form.
|
||||
|
||||
3. Grant of Patent License. Subject to the terms and conditions of
|
||||
this License, each Contributor hereby grants to You a perpetual,
|
||||
worldwide, non-exclusive, no-charge, royalty-free, irrevocable
|
||||
(except as stated in this section) patent license to make, have made,
|
||||
use, offer to sell, sell, import, and otherwise transfer the Work,
|
||||
where such license applies only to those patent claims licensable
|
||||
by such Contributor that are necessarily infringed by their
|
||||
Contribution(s) alone or by combination of their Contribution(s)
|
||||
with the Work to which such Contribution(s) was submitted. If You
|
||||
institute patent litigation against any entity (including a
|
||||
cross-claim or counterclaim in a lawsuit) alleging that the Work
|
||||
or a Contribution incorporated within the Work constitutes direct
|
||||
or contributory patent infringement, then any patent licenses
|
||||
granted to You under this License for that Work shall terminate
|
||||
as of the date such litigation is filed.
|
||||
|
||||
4. Redistribution. You may reproduce and distribute copies of the
|
||||
Work or Derivative Works thereof in any medium, with or without
|
||||
modifications, and in Source or Object form, provided that You
|
||||
meet the following conditions:
|
||||
|
||||
(a) You must give any other recipients of the Work or
|
||||
Derivative Works a copy of this License; and
|
||||
|
||||
(b) You must cause any modified files to carry prominent notices
|
||||
stating that You changed the files; and
|
||||
|
||||
(c) You must retain, in the Source form of any Derivative Works
|
||||
that You distribute, all copyright, patent, trademark, and
|
||||
attribution notices from the Source form of the Work,
|
||||
excluding those notices that do not pertain to any part of
|
||||
the Derivative Works; and
|
||||
|
||||
(d) If the Work includes a "NOTICE" text file as part of its
|
||||
distribution, then any Derivative Works that You distribute must
|
||||
include a readable copy of the attribution notices contained
|
||||
within such NOTICE file, excluding those notices that do not
|
||||
pertain to any part of the Derivative Works, in at least one
|
||||
of the following places: within a NOTICE text file distributed
|
||||
as part of the Derivative Works; within the Source form or
|
||||
documentation, if provided along with the Derivative Works; or,
|
||||
within a display generated by the Derivative Works, if and
|
||||
wherever such third-party notices normally appear. The contents
|
||||
of the NOTICE file are for informational purposes only and
|
||||
do not modify the License. You may add Your own attribution
|
||||
notices within Derivative Works that You distribute, alongside
|
||||
or as an addendum to the NOTICE text from the Work, provided
|
||||
that such additional attribution notices cannot be construed
|
||||
as modifying the License.
|
||||
|
||||
You may add Your own copyright statement to Your modifications and
|
||||
may provide additional or different license terms and conditions
|
||||
for use, reproduction, or distribution of Your modifications, or
|
||||
for any such Derivative Works as a whole, provided Your use,
|
||||
reproduction, and distribution of the Work otherwise complies with
|
||||
the conditions stated in this License.
|
||||
|
||||
5. Submission of Contributions. Unless You explicitly state otherwise,
|
||||
any Contribution intentionally submitted for inclusion in the Work
|
||||
by You to the Licensor shall be under the terms and conditions of
|
||||
this License, without any additional terms or conditions.
|
||||
Notwithstanding the above, nothing herein shall supersede or modify
|
||||
the terms of any separate license agreement you may have executed
|
||||
with Licensor regarding such Contributions.
|
||||
|
||||
6. Trademarks. This License does not grant permission to use the trade
|
||||
names, trademarks, service marks, or product names of the Licensor,
|
||||
except as required for reasonable and customary use in describing the
|
||||
origin of the Work and reproducing the content of the NOTICE file.
|
||||
|
||||
7. Disclaimer of Warranty. Unless required by applicable law or
|
||||
agreed to in writing, Licensor provides the Work (and each
|
||||
Contributor provides its Contributions) on an "AS IS" BASIS,
|
||||
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
|
||||
implied, including, without limitation, any warranties or conditions
|
||||
of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
|
||||
PARTICULAR PURPOSE. You are solely responsible for determining the
|
||||
appropriateness of using or redistributing the Work and assume any
|
||||
risks associated with Your exercise of permissions under this License.
|
||||
|
||||
8. Limitation of Liability. In no event and under no legal theory,
|
||||
whether in tort (including negligence), contract, or otherwise,
|
||||
unless required by applicable law (such as deliberate and grossly
|
||||
negligent acts) or agreed to in writing, shall any Contributor be
|
||||
liable to You for damages, including any direct, indirect, special,
|
||||
incidental, or consequential damages of any character arising as a
|
||||
result of this License or out of the use or inability to use the
|
||||
Work (including but not limited to damages for loss of goodwill,
|
||||
work stoppage, computer failure or malfunction, or any and all
|
||||
other commercial damages or losses), even if such Contributor
|
||||
has been advised of the possibility of such damages.
|
||||
|
||||
9. Accepting Warranty or Additional Liability. While redistributing
|
||||
the Work or Derivative Works thereof, You may choose to offer,
|
||||
and charge a fee for, acceptance of support, warranty, indemnity,
|
||||
or other liability obligations and/or rights consistent with this
|
||||
License. However, in accepting such obligations, You may act only
|
||||
on Your own behalf and on Your sole responsibility, not on behalf
|
||||
of any other Contributor, and only if You agree to indemnify,
|
||||
defend, and hold each Contributor harmless for any liability
|
||||
incurred by, or claims asserted against, such Contributor by reason
|
||||
of your accepting any such warranty or additional liability.
|
||||
|
||||
END OF TERMS AND CONDITIONS
|
||||
|
||||
APPENDIX: How to apply the Apache License to your work.
|
||||
|
||||
To apply the Apache License to your work, attach the following
|
||||
boilerplate notice, with the fields enclosed by brackets "[]"
|
||||
replaced with your own identifying information. (Don't include
|
||||
the brackets!) The text should be enclosed in the appropriate
|
||||
comment syntax for the file format. We also recommend that a
|
||||
file or class name and description of purpose be included on the
|
||||
same "printed page" as the copyright notice for easier
|
||||
identification within third-party archives.
|
||||
|
||||
Copyright [yyyy] [name of copyright owner]
|
||||
|
||||
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.
|
||||
+131
@@ -0,0 +1,131 @@
|
||||
<h3 align="center">
|
||||
<img src="assets/logo.png" alt="Logo" style="vertical-align: middle; width: 40px; height: 40px;">
|
||||
Less-to-More Generalization: Unlocking More Controllability by In-Context Generation
|
||||
</h3>
|
||||
|
||||
<p align="center">
|
||||
<a href="https://bytedance.github.io/UNO/"><img alt="Build" src="https://img.shields.io/badge/Project%20Page-UNO-yellow"></a>
|
||||
<a href="https://arxiv.org/abs/2504.02160"><img alt="Build" src="https://img.shields.io/badge/arXiv%20paper-UNO-b31b1b.svg"></a>
|
||||
<a href="https://huggingface.co/bytedance-research/UNO"><img src="https://img.shields.io/static/v1?label=%F0%9F%A4%97%20Hugging%20Face&message=Model&color=orange"></a>
|
||||
<a href="https://huggingface.co/spaces/bytedance-research/UNO-FLUX"><img src="https://img.shields.io/static/v1?label=%F0%9F%A4%97%20Hugging%20Face&message=demo&color=orange"></a>
|
||||
</p>
|
||||
|
||||
><p align="center"> <span style="color:#137cf3; font-family: Gill Sans">Shaojin Wu,</span><sup></sup></a> <span style="color:#137cf3; font-family: Gill Sans">Mengqi Huang</span><sup>*</sup>,</a> <span style="color:#137cf3; font-family: Gill Sans">Wenxu Wu,</span><sup></sup></a> <span style="color:#137cf3; font-family: Gill Sans">Yufeng Cheng,</span><sup></sup> </a> <span style="color:#137cf3; font-family: Gill Sans">Fei Ding</span><sup>+</sup>,</a> <span style="color:#137cf3; font-family: Gill Sans">Qian He</span></a> <br>
|
||||
><span style="font-size: 16px">Intelligent Creation Team, ByteDance</span></p>
|
||||
|
||||
<p align="center">
|
||||
<img src="./assets/teaser.jpg" width=95% height=95%
|
||||
class="center">
|
||||
</p>
|
||||
|
||||
## 🔥 News
|
||||
- [04/2025] 🔥 The [demo](https://huggingface.co/spaces/bytedance-research/UNO-FLUX) of UNO is released.
|
||||
- [04/2025] 🔥 The [training code](https://github.com/bytedance/UNO), [inference code](https://github.com/bytedance/UNO), and [model](https://huggingface.co/bytedance-research/UNO) of UNO are released.
|
||||
- [04/2025] 🔥 The [project page](https://bytedance.github.io/UNO) of UNO is created.
|
||||
- [04/2025] 🔥 The arXiv [paper](https://arxiv.org/abs/2504.02160) of UNO is released.
|
||||
|
||||
## 📖 Introduction
|
||||
In this study, we propose a highly-consistent data synthesis pipeline to tackle this challenge. This pipeline harnesses the intrinsic in-context generation capabilities of diffusion transformers and generates high-consistency multi-subject paired data. Additionally, we introduce UNO, which consists of progressive cross-modal alignment and universal rotary position embedding. It is a multi-image conditioned subject-to-image model iteratively trained from a text-to-image model. Extensive experiments show that our method can achieve high consistency while ensuring controllability in both single-subject and multi-subject driven generation.
|
||||
|
||||
|
||||
## ⚡️ Quick Start
|
||||
|
||||
### 🔧 Requirements and Installation
|
||||
|
||||
Install the requirements
|
||||
```bash
|
||||
## create a virtual environment with python >= 3.10 <= 3.12, like
|
||||
# python -m venv uno_env
|
||||
# source uno_env/bin/activate
|
||||
# then install
|
||||
pip install -r requirements.txt
|
||||
```
|
||||
|
||||
then download checkpoints in one of the three ways:
|
||||
1. Directly run the inference scripts, the checkpoints will be downloaded automatically by the `hf_hub_download` function in the code to your `$HF_HOME`(the default value is `~/.cache/huggingface`).
|
||||
2. use `huggingface-cli download <repo name>` to download `black-forest-labs/FLUX.1-dev`, `xlabs-ai/xflux_text_encoders`, `openai/clip-vit-large-patch14`, `bytedance-research/UNO`, then run the inference scripts.
|
||||
3. use `huggingface-cli download <repo name> --local-dir <LOCAL_DIR>` to download all the checkpoints menthioned in 2. to the directories your want. Then set the environment variable `AE`, `FLUX`, `T5`, `CLIP`, `LORA` to the corresponding paths. Finally, run the inference scripts.
|
||||
|
||||
### 🌟 Gradio Demo
|
||||
|
||||
```bash
|
||||
python app.py
|
||||
```
|
||||
|
||||
|
||||
### ✍️ Inference
|
||||
Start from the examples below to explore and spark your creativity. ✨
|
||||
```bash
|
||||
python inference.py --prompt "A clock on the beach is under a red sun umbrella" --image_paths "assets/clock.png" --width 704 --height 704
|
||||
python inference.py --prompt "The figurine is in the crystal ball" --image_paths "assets/figurine.png" "assets/crystal_ball.png" --width 704 --height 704
|
||||
python inference.py --prompt "The logo is printed on the cup" --image_paths "assets/cat_cafe.png" "assets/cup.png" --width 704 --height 704
|
||||
```
|
||||
|
||||
Optional prepreration: If you want to test the inference on dreambench at the first time, you should clone the submodule `dreambench` to download the dataset.
|
||||
|
||||
```bash
|
||||
git submodule update --init
|
||||
```
|
||||
Then running the following scripts:
|
||||
```bash
|
||||
# evaluated on dreambench
|
||||
## for single-subject
|
||||
python inference.py --eval_json_path ./datasets/dreambench_singleip.json
|
||||
## for multi-subject
|
||||
python inference.py --eval_json_path ./datasets/dreambench_multiip.json
|
||||
```
|
||||
|
||||
|
||||
|
||||
### 🚄 Training
|
||||
|
||||
```bash
|
||||
accelerate launch train.py
|
||||
```
|
||||
|
||||
|
||||
### 📌 Tips and Notes
|
||||
We integrate single-subject and multi-subject generation within a unified model. For single-subject scenarios, the longest side of the reference image is set to 512 by default, while for multi-subject scenarios, it is set to 320. UNO demonstrates remarkable flexibility across various aspect ratios, thanks to its training on a multi-scale dataset. Despite being trained within 512 buckets, it can handle higher resolutions, including 512, 568, and 704, among others.
|
||||
|
||||
UNO excels in subject-driven generation but has room for improvement in generalization due to dataset constraints. We are actively developing an enhanced model—stay tuned for updates. Your feedback is valuable, so please feel free to share any suggestions.
|
||||
|
||||
## 🎨 Application Scenarios
|
||||
<p align="center">
|
||||
<img src="./assets/simplecase.jpg" width=95% height=95%
|
||||
class="center">
|
||||
</p>
|
||||
|
||||
## 📄 Disclaimer
|
||||
<p>
|
||||
We open-source this project for academic research. The vast majority of images
|
||||
used in this project are either generated or licensed. If you have any concerns,
|
||||
please contact us, and we will promptly remove any inappropriate content.
|
||||
Our code is released under the Apache 2.0 License,, while our models are under
|
||||
the CC BY-NC 4.0 License. Any models related to <a href="https://huggingface.co/black-forest-labs/FLUX.1-dev" target="_blank">FLUX.1-dev</a>
|
||||
base model must adhere to the original licensing terms.
|
||||
<br><br>This research aims to advance the field of generative AI. Users are free to
|
||||
create images using this tool, provided they comply with local laws and exercise
|
||||
responsible usage. The developers are not liable for any misuse of the tool by users.</p>
|
||||
|
||||
## 🚀 Updates
|
||||
For the purpose of fostering research and the open-source community, we plan to open-source the entire project, encompassing training, inference, weights, etc. Thank you for your patience and support! 🌟
|
||||
- [x] Release github repo.
|
||||
- [x] Release inference code.
|
||||
- [x] Release training code.
|
||||
- [x] Release model checkpoints.
|
||||
- [x] Release arXiv paper.
|
||||
- [x] Release huggingface space demo.
|
||||
- [ ] Release in-context data generation pipelines.
|
||||
|
||||
## Citation
|
||||
If UNO is helpful, please help to ⭐ the repo.
|
||||
|
||||
If you find this project useful for your research, please consider citing our paper:
|
||||
```bibtex
|
||||
@article{wu2025less,
|
||||
title={Less-to-More Generalization: Unlocking More Controllability by In-Context Generation},
|
||||
author={Wu, Shaojin and Huang, Mengqi and Wu, Wenxu and Cheng, Yufeng and Ding, Fei and He, Qian},
|
||||
journal={arXiv preprint arXiv:2504.02160},
|
||||
year={2025}
|
||||
}
|
||||
```
|
||||
+148
@@ -0,0 +1,148 @@
|
||||
# Copyright (c) 2025 Bytedance Ltd. and/or its affiliates. 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.
|
||||
|
||||
import dataclasses
|
||||
import json
|
||||
from pathlib import Path
|
||||
|
||||
import gradio as gr
|
||||
import torch
|
||||
|
||||
from .uno.flux.pipeline import UNOPipeline
|
||||
|
||||
|
||||
def get_examples(examples_dir: str = "assets/examples") -> list:
|
||||
examples = Path(examples_dir)
|
||||
ans = []
|
||||
for example in examples.iterdir():
|
||||
if not example.is_dir():
|
||||
continue
|
||||
with open(example / "config.json") as f:
|
||||
example_dict = json.load(f)
|
||||
|
||||
|
||||
example_list = []
|
||||
|
||||
example_list.append(example_dict["useage"]) # case for
|
||||
example_list.append(example_dict["prompt"]) # prompt
|
||||
|
||||
for key in ["image_ref1", "image_ref2", "image_ref3", "image_ref4"]:
|
||||
if key in example_dict:
|
||||
example_list.append(str(example / example_dict[key]))
|
||||
else:
|
||||
example_list.append(None)
|
||||
|
||||
example_list.append(example_dict["seed"])
|
||||
|
||||
ans.append(example_list)
|
||||
return ans
|
||||
|
||||
|
||||
def create_demo(
|
||||
model_type: str,
|
||||
device: str = "cuda" if torch.cuda.is_available() else "cpu",
|
||||
offload: bool = False,
|
||||
):
|
||||
pipeline = UNOPipeline(model_type, device, offload, only_lora=True, lora_rank=512)
|
||||
|
||||
badges_text = r"""
|
||||
<div style="text-align: center; display: flex; justify-content: left; gap: 5px;">
|
||||
<a href="https://bytedance.github.io/UNO/"><img alt="Build" src="https://img.shields.io/badge/Project%20Page-UNO-yellow"></a>
|
||||
<a href="https://arxiv.org/abs/2504.02160"><img alt="Build" src="https://img.shields.io/badge/arXiv%20paper-UNO-b31b1b.svg"></a>
|
||||
<a href="https://huggingface.co/bytedance-research/UNO"><img src="https://img.shields.io/static/v1?label=%F0%9F%A4%97%20Hugging%20Face&message=Model&color=orange"></a>
|
||||
<a href="https://huggingface.co/spaces/bytedance-research/UNO-FLUX"><img src="https://img.shields.io/static/v1?label=%F0%9F%A4%97%20Hugging%20Face&message=demo&color=orange"></a>
|
||||
</div>
|
||||
""".strip()
|
||||
|
||||
with gr.Blocks() as demo:
|
||||
gr.Markdown(f"# UNO by UNO team")
|
||||
gr.Markdown(badges_text)
|
||||
with gr.Row():
|
||||
with gr.Column():
|
||||
prompt = gr.Textbox(label="Prompt", value="handsome woman in the city")
|
||||
with gr.Row():
|
||||
image_prompt1 = gr.Image(label="Ref Img1", visible=True, interactive=True, type="pil")
|
||||
image_prompt2 = gr.Image(label="Ref Img2", visible=True, interactive=True, type="pil")
|
||||
image_prompt3 = gr.Image(label="Ref Img3", visible=True, interactive=True, type="pil")
|
||||
image_prompt4 = gr.Image(label="Ref img4", visible=True, interactive=True, type="pil")
|
||||
|
||||
with gr.Row():
|
||||
with gr.Column():
|
||||
width = gr.Slider(512, 2048, 512, step=16, label="Gneration Width")
|
||||
height = gr.Slider(512, 2048, 512, step=16, label="Gneration Height")
|
||||
with gr.Column():
|
||||
gr.Markdown("📌 The model trained on 512x512 resolution.\n")
|
||||
gr.Markdown(
|
||||
"The size closer to 512 is more stable,"
|
||||
" and the higher size gives a better visual effect but is less stable"
|
||||
)
|
||||
|
||||
with gr.Accordion("Advanced Options", open=False):
|
||||
with gr.Row():
|
||||
num_steps = gr.Slider(1, 50, 25, step=1, label="Number of steps")
|
||||
guidance = gr.Slider(1.0, 5.0, 4.0, step=0.1, label="Guidance", interactive=True)
|
||||
seed = gr.Number(-1, label="Seed (-1 for random)")
|
||||
|
||||
generate_btn = gr.Button("Generate")
|
||||
|
||||
with gr.Column():
|
||||
output_image = gr.Image(label="Generated Image")
|
||||
download_btn = gr.File(label="Download full-resolution", type="filepath", interactive=False)
|
||||
|
||||
|
||||
inputs = [
|
||||
prompt, width, height, guidance, num_steps,
|
||||
seed, image_prompt1, image_prompt2, image_prompt3, image_prompt4
|
||||
]
|
||||
generate_btn.click(
|
||||
fn=pipeline.gradio_generate,
|
||||
inputs=inputs,
|
||||
outputs=[output_image, download_btn],
|
||||
)
|
||||
|
||||
example_text = gr.Text("", visible=False, label="Case For:")
|
||||
examples = get_examples("./assets/examples")
|
||||
|
||||
gr.Examples(
|
||||
examples=examples,
|
||||
inputs=[
|
||||
example_text, prompt,
|
||||
image_prompt1, image_prompt2, image_prompt3, image_prompt4,
|
||||
seed, output_image
|
||||
],
|
||||
)
|
||||
|
||||
return demo
|
||||
|
||||
if __name__ == "__main__":
|
||||
from typing import Literal
|
||||
|
||||
from transformers import HfArgumentParser
|
||||
|
||||
@dataclasses.dataclass
|
||||
class AppArgs:
|
||||
name: Literal["flux-dev", "flux-dev-fp8", "flux-schnell"] = "flux-dev"
|
||||
device: Literal["cuda", "cpu"] = "cuda" if torch.cuda.is_available() else "cpu"
|
||||
offload: bool = dataclasses.field(
|
||||
default=False,
|
||||
metadata={"help": "If True, sequantial offload the models(ae, dit, text encoder) to CPU if not used."}
|
||||
)
|
||||
port: int = 7860
|
||||
|
||||
parser = HfArgumentParser([AppArgs])
|
||||
args_tuple = parser.parse_args_into_dataclasses() # type: tuple[AppArgs]
|
||||
args = args_tuple[0]
|
||||
|
||||
demo = create_demo(args.name, args.device, args.offload)
|
||||
demo.launch(server_port=args.port)
|
||||
Binary file not shown.
|
After Width: | Height: | Size: 1006 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 1.8 MiB |
Binary file not shown.
|
After Width: | Height: | Size: 1011 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 369 KiB |
@@ -0,0 +1,22 @@
|
||||
{
|
||||
"prompt": "A clock on the beach is under a red sun umbrella",
|
||||
"image_paths": [
|
||||
"assets/clock.png"
|
||||
],
|
||||
"eval_json_path": null,
|
||||
"offload": false,
|
||||
"num_images_per_prompt": 1,
|
||||
"model_type": "flux-dev",
|
||||
"width": 704,
|
||||
"height": 704,
|
||||
"ref_size": 512,
|
||||
"num_steps": 25,
|
||||
"guidance": 4,
|
||||
"seed": 3407,
|
||||
"save_path": "output/inference",
|
||||
"only_lora": true,
|
||||
"concat_refs": false,
|
||||
"lora_rank": 512,
|
||||
"data_resolution": 512,
|
||||
"pe": "d"
|
||||
}
|
||||
Binary file not shown.
|
After Width: | Height: | Size: 490 KiB |
@@ -0,0 +1,23 @@
|
||||
{
|
||||
"prompt": "The figurine is in the crystal ball",
|
||||
"image_paths": [
|
||||
"assets/figurine.png",
|
||||
"assets/crystal_ball.png"
|
||||
],
|
||||
"eval_json_path": null,
|
||||
"offload": false,
|
||||
"num_images_per_prompt": 1,
|
||||
"model_type": "flux-dev",
|
||||
"width": 704,
|
||||
"height": 704,
|
||||
"ref_size": 320,
|
||||
"num_steps": 25,
|
||||
"guidance": 4,
|
||||
"seed": 3407,
|
||||
"save_path": "output/inference",
|
||||
"only_lora": true,
|
||||
"concat_refs": false,
|
||||
"lora_rank": 512,
|
||||
"data_resolution": 512,
|
||||
"pe": "d"
|
||||
}
|
||||
Binary file not shown.
|
After Width: | Height: | Size: 570 KiB |
@@ -0,0 +1,8 @@
|
||||
{
|
||||
"prompt": "A clock on the beach is under a red sun umbrella",
|
||||
"seed": 0,
|
||||
"ref_long_side": 512,
|
||||
"useage": "one2one",
|
||||
"image_ref1": "./ref1.jpg",
|
||||
"image_result": "./result.png"
|
||||
}
|
||||
Binary file not shown.
|
After Width: | Height: | Size: 136 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 281 KiB |
@@ -0,0 +1,23 @@
|
||||
{
|
||||
"prompt": "The logo is printed on the cup",
|
||||
"image_paths": [
|
||||
"assets/cat_cafe.png",
|
||||
"assets/cup.png"
|
||||
],
|
||||
"eval_json_path": null,
|
||||
"offload": false,
|
||||
"num_images_per_prompt": 1,
|
||||
"model_type": "flux-dev",
|
||||
"width": 704,
|
||||
"height": 704,
|
||||
"ref_size": 320,
|
||||
"num_steps": 25,
|
||||
"guidance": 4,
|
||||
"seed": 3407,
|
||||
"save_path": "output/inference",
|
||||
"only_lora": true,
|
||||
"concat_refs": false,
|
||||
"lora_rank": 512,
|
||||
"data_resolution": 512,
|
||||
"pe": "d"
|
||||
}
|
||||
Binary file not shown.
|
After Width: | Height: | Size: 297 KiB |
@@ -0,0 +1,8 @@
|
||||
{
|
||||
"prompt": "A pretty woman wears a flower petal dress, in the flower",
|
||||
"seed": 1,
|
||||
"ref_long_side": 512,
|
||||
"useage": "one2one",
|
||||
"image_ref1": "./ref1.png",
|
||||
"image_result": "./result.png"
|
||||
}
|
||||
Binary file not shown.
|
After Width: | Height: | Size: 983 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 338 KiB |
@@ -0,0 +1,9 @@
|
||||
{
|
||||
"prompt": "The figurine is in the crystal ball",
|
||||
"seed": 0,
|
||||
"ref_long_side": 320,
|
||||
"useage": "two2one",
|
||||
"image_ref1": "./ref1.png",
|
||||
"image_ref2": "./ref2.png",
|
||||
"image_result": "./result.png"
|
||||
}
|
||||
Binary file not shown.
|
After Width: | Height: | Size: 743 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 1011 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 285 KiB |
@@ -0,0 +1,9 @@
|
||||
{
|
||||
"prompt": "The logo is printed on the cup",
|
||||
"seed": 61733557,
|
||||
"ref_long_side": 320,
|
||||
"useage": "two2one",
|
||||
"image_ref1": "./ref1.png",
|
||||
"image_ref2": "./ref2.png",
|
||||
"image_result": "./result.png"
|
||||
}
|
||||
Binary file not shown.
|
After Width: | Height: | Size: 1006 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 369 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 162 KiB |
@@ -0,0 +1,10 @@
|
||||
{
|
||||
"prompt": "A woman wears the dress and holds a bag, in the flowers.",
|
||||
"seed": 37635012,
|
||||
"ref_long_side": 320,
|
||||
"useage": "many2one",
|
||||
"image_ref1": "./ref1.png",
|
||||
"image_ref2": "./ref2.png",
|
||||
"image_ref3": "./ref3.png",
|
||||
"image_result": "./result.png"
|
||||
}
|
||||
Binary file not shown.
|
After Width: | Height: | Size: 1.1 MiB |
Binary file not shown.
|
After Width: | Height: | Size: 709 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 698 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 350 KiB |
@@ -0,0 +1,7 @@
|
||||
{
|
||||
"prompt": "A woman wears the dress and holds a bag, in the flowers.",
|
||||
"seed": 37635012,
|
||||
"ref_long_side": 512,
|
||||
"useage": "t2i",
|
||||
"image_result": "./result.png"
|
||||
}
|
||||
Binary file not shown.
|
After Width: | Height: | Size: 319 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 743 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 56 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 631 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 631 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 691 KiB |
@@ -0,0 +1,15 @@
|
||||
{
|
||||
"bf16": {
|
||||
"enabled": "auto"
|
||||
},
|
||||
"zero_optimization": {
|
||||
"stage": 2,
|
||||
"offload_optimizer": {
|
||||
"device": "none"
|
||||
},
|
||||
"contiguous_gradients": true,
|
||||
"overlap_comm": true
|
||||
},
|
||||
"train_micro_batch_size_per_gpu": 1,
|
||||
"gradient_accumulation_steps": "auto"
|
||||
}
|
||||
@@ -0,0 +1,33 @@
|
||||
{
|
||||
"bf16": {
|
||||
"enabled": "auto",
|
||||
"loss_scale": 0,
|
||||
"loss_scale_window": 1000,
|
||||
"initial_scale_power": 16,
|
||||
"hysteresis": 2,
|
||||
"min_loss_scale": 1
|
||||
},
|
||||
|
||||
"zero_optimization": {
|
||||
"stage": 3,
|
||||
"offload_optimizer": {
|
||||
"device": "cpu",
|
||||
"pin_memory": true
|
||||
},
|
||||
"offload_param": {
|
||||
"device": "cpu",
|
||||
"pin_memory": true
|
||||
},
|
||||
"overlap_comm": true,
|
||||
"contiguous_gradients": true,
|
||||
"reduce_bucket_size": 16777216,
|
||||
"stage3_prefetch_bucket_size": 15099494,
|
||||
"stage3_param_persistence_threshold": 40960,
|
||||
"sub_group_size": 1e9,
|
||||
"stage3_max_live_parameters": 1e9,
|
||||
"stage3_max_reuse_distance": 1e9,
|
||||
"stage3_gather_16bit_weights_on_model_save": true
|
||||
},
|
||||
"gradient_accumulation_steps": "auto",
|
||||
"train_micro_batch_size_per_gpu": 1
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,125 @@
|
||||
# Copyright (c) 2025 Bytedance Ltd. and/or its affiliates. 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.
|
||||
|
||||
import os
|
||||
import dataclasses
|
||||
from typing import Literal
|
||||
|
||||
from accelerate import Accelerator
|
||||
from transformers import HfArgumentParser
|
||||
from PIL import Image
|
||||
import json
|
||||
import itertools
|
||||
|
||||
from .uno.flux.pipeline import UNOPipeline, preprocess_ref
|
||||
|
||||
|
||||
def horizontal_concat(images):
|
||||
widths, heights = zip(*(img.size for img in images))
|
||||
|
||||
total_width = sum(widths)
|
||||
max_height = max(heights)
|
||||
|
||||
new_im = Image.new('RGB', (total_width, max_height))
|
||||
|
||||
x_offset = 0
|
||||
for img in images:
|
||||
new_im.paste(img, (x_offset, 0))
|
||||
x_offset += img.size[0]
|
||||
|
||||
return new_im
|
||||
|
||||
@dataclasses.dataclass
|
||||
class InferenceArgs:
|
||||
prompt: str | None = None
|
||||
image_paths: list[str] | None = None
|
||||
eval_json_path: str | None = None
|
||||
offload: bool = False
|
||||
num_images_per_prompt: int = 1
|
||||
model_type: Literal["flux-dev", "flux-dev-fp8", "flux-schnell"] = "flux-dev"
|
||||
width: int = 512
|
||||
height: int = 512
|
||||
ref_size: int = -1
|
||||
num_steps: int = 25
|
||||
guidance: float = 4
|
||||
seed: int = 3407
|
||||
save_path: str = "output/inference"
|
||||
only_lora: bool = True
|
||||
concat_refs: bool = False
|
||||
lora_rank: int = 512
|
||||
data_resolution: int = 512
|
||||
pe: Literal['d', 'h', 'w', 'o'] = 'd'
|
||||
|
||||
def main(args: InferenceArgs):
|
||||
accelerator = Accelerator()
|
||||
|
||||
pipeline = UNOPipeline(
|
||||
args.model_type,
|
||||
accelerator.device,
|
||||
args.offload,
|
||||
only_lora=args.only_lora,
|
||||
lora_rank=args.lora_rank
|
||||
)
|
||||
|
||||
assert args.prompt is not None or args.eval_json_path is not None, \
|
||||
"Please provide either prompt or eval_json_path"
|
||||
|
||||
if args.eval_json_path is not None:
|
||||
with open(args.eval_json_path, "rt") as f:
|
||||
data_dicts = json.load(f)
|
||||
data_root = os.path.dirname(args.eval_json_path)
|
||||
else:
|
||||
data_root = "./"
|
||||
data_dicts = [{"prompt": args.prompt, "image_paths": args.image_paths}]
|
||||
|
||||
for (i, data_dict), j in itertools.product(enumerate(data_dicts), range(args.num_images_per_prompt)):
|
||||
if (i * args.num_images_per_prompt + j) % accelerator.num_processes != accelerator.process_index:
|
||||
continue
|
||||
|
||||
ref_imgs = [
|
||||
Image.open(os.path.join(data_root, img_path))
|
||||
for img_path in data_dict["image_paths"]
|
||||
]
|
||||
if args.ref_size==-1:
|
||||
args.ref_size = 512 if len(ref_imgs)==1 else 320
|
||||
|
||||
ref_imgs = [preprocess_ref(img, args.ref_size) for img in ref_imgs]
|
||||
|
||||
image_gen = pipeline(
|
||||
prompt=data_dict["prompt"],
|
||||
width=args.width,
|
||||
height=args.height,
|
||||
guidance=args.guidance,
|
||||
num_steps=args.num_steps,
|
||||
seed=args.seed + j,
|
||||
ref_imgs=ref_imgs,
|
||||
pe=args.pe,
|
||||
)
|
||||
if args.concat_refs:
|
||||
image_gen = horizontal_concat([image_gen, *ref_imgs])
|
||||
|
||||
os.makedirs(args.save_path, exist_ok=True)
|
||||
image_gen.save(os.path.join(args.save_path, f"{i}_{j}.png"))
|
||||
|
||||
# save config and image
|
||||
args_dict = vars(args)
|
||||
args_dict['prompt'] = data_dict["prompt"]
|
||||
args_dict['image_paths'] = data_dict["image_paths"]
|
||||
with open(os.path.join(args.save_path, f"{i}_{j}.json"), 'w') as f:
|
||||
json.dump(args_dict, f, indent=4)
|
||||
|
||||
if __name__ == "__main__":
|
||||
parser = HfArgumentParser([InferenceArgs])
|
||||
args = parser.parse_args_into_dataclasses()[0]
|
||||
main(args)
|
||||
@@ -0,0 +1,12 @@
|
||||
accelerate==1.1.1
|
||||
deepspeed==0.14.4
|
||||
einops==0.8.0
|
||||
transformers==4.43.3
|
||||
huggingface-hub
|
||||
diffusers==0.30.1
|
||||
sentencepiece==0.2.0
|
||||
gradio==5.22.0
|
||||
|
||||
--extra-index-url https://download.pytorch.org/whl/cu124
|
||||
torch==2.4.0
|
||||
torchvision==0.19.0
|
||||
+482
@@ -0,0 +1,482 @@
|
||||
# Copyright (c) 2025 Bytedance Ltd. and/or its affiliates. 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.
|
||||
|
||||
import dataclasses
|
||||
import gc
|
||||
import itertools
|
||||
import logging
|
||||
import os
|
||||
import random
|
||||
from copy import deepcopy
|
||||
from typing import TYPE_CHECKING, Literal
|
||||
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
import transformers
|
||||
from accelerate import Accelerator, DeepSpeedPlugin
|
||||
from accelerate.logging import get_logger
|
||||
from accelerate.utils import set_seed
|
||||
from diffusers.optimization import get_scheduler
|
||||
from einops import rearrange
|
||||
from PIL import Image
|
||||
from safetensors.torch import load_file
|
||||
from torch.utils.data import DataLoader
|
||||
from tqdm import tqdm
|
||||
|
||||
from uno.dataset.uno import FluxPairedDatasetV2
|
||||
from uno.flux.sampling import denoise, get_noise, get_schedule, prepare_multi_ip, unpack
|
||||
from uno.flux.util import load_ae, load_clip, load_flow_model, load_t5, set_lora
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from uno.flux.model import Flux
|
||||
from uno.flux.modules.autoencoder import AutoEncoder
|
||||
from uno.flux.modules.conditioner import HFEmbedder
|
||||
|
||||
logger = get_logger(__name__)
|
||||
|
||||
def get_models(name: str, device, offload: bool=False):
|
||||
t5 = load_t5(device, max_length=512)
|
||||
clip = load_clip(device)
|
||||
model = load_flow_model(name, device="cpu")
|
||||
vae = load_ae(name, device="cpu" if offload else device)
|
||||
return model, vae, t5, clip
|
||||
|
||||
def inference(
|
||||
batch: dict,
|
||||
model: "Flux", t5: "HFEmbedder", clip: "HFEmbedder", ae: "AutoEncoder",
|
||||
accelerator: Accelerator,
|
||||
seed: int = 0,
|
||||
pe: Literal["d", "h", "w", "o"] = "d"
|
||||
) -> Image.Image:
|
||||
ref_imgs = batch["ref_imgs"]
|
||||
prompt = batch["txt"]
|
||||
neg_prompt = ''
|
||||
width, height = 512, 512
|
||||
num_steps = 25
|
||||
x = get_noise(
|
||||
1, height, width,
|
||||
device=accelerator.device,
|
||||
dtype=torch.bfloat16,
|
||||
seed=seed + accelerator.process_index
|
||||
)
|
||||
timesteps = get_schedule(
|
||||
num_steps,
|
||||
(width // 8) * (height // 8) // (16 * 16),
|
||||
shift=True,
|
||||
)
|
||||
with torch.no_grad():
|
||||
ref_imgs = [
|
||||
ae.encode(ref_img_.to(accelerator.device, torch.float32)).to(torch.bfloat16)
|
||||
for ref_img_ in ref_imgs
|
||||
]
|
||||
inp_cond = prepare_multi_ip(
|
||||
t5=t5, clip=clip, img=x, prompt=prompt,
|
||||
ref_imgs=ref_imgs,
|
||||
pe=pe
|
||||
)
|
||||
neg_inp_cond = prepare_multi_ip(
|
||||
t5=t5, clip=clip, img=x, prompt=neg_prompt,
|
||||
ref_imgs=ref_imgs,
|
||||
pe=pe
|
||||
)
|
||||
|
||||
x = denoise(
|
||||
model,
|
||||
**inp_cond,
|
||||
timesteps=timesteps,
|
||||
guidance=4,
|
||||
timestep_to_start_cfg=30,
|
||||
neg_txt=neg_inp_cond['txt'],
|
||||
neg_txt_ids=neg_inp_cond['txt_ids'],
|
||||
neg_vec=neg_inp_cond['vec'],
|
||||
true_gs=3.5,
|
||||
image_proj=None,
|
||||
neg_image_proj=None,
|
||||
ip_scale=1,
|
||||
neg_ip_scale=1
|
||||
)
|
||||
|
||||
x = unpack(x.float(), height, width)
|
||||
x = ae.decode(x)
|
||||
|
||||
x1 = x.clamp(-1, 1)
|
||||
x1 = rearrange(x1[-1], "c h w -> h w c")
|
||||
output_img = Image.fromarray((127.5 * (x1 + 1.0)).cpu().byte().numpy())
|
||||
|
||||
return output_img
|
||||
|
||||
|
||||
def resume_from_checkpoint(
|
||||
resume_from_checkpoint: str | None | Literal["latest"],
|
||||
project_dir: str,
|
||||
accelerator: Accelerator,
|
||||
dit: "Flux",
|
||||
optimizer: torch.optim.Optimizer,
|
||||
lr_scheduler: torch.optim.lr_scheduler.LRScheduler,
|
||||
dit_ema_dict: dict | None = None,
|
||||
) -> tuple["Flux", torch.optim.Optimizer, torch.optim.lr_scheduler.LRScheduler, dict | None, int]:
|
||||
# Potentially load in the weights and states from a previous save
|
||||
if resume_from_checkpoint is None:
|
||||
return dit, optimizer, lr_scheduler, dit_ema_dict, 0
|
||||
|
||||
if resume_from_checkpoint == "latest":
|
||||
# Get the most recent checkpoint
|
||||
dirs = os.listdir(project_dir)
|
||||
dirs = [d for d in dirs if d.startswith("checkpoint")]
|
||||
dirs = sorted(dirs, key=lambda x: int(x.split("-")[1]))
|
||||
if len(dirs) == 0:
|
||||
accelerator.print(
|
||||
f"Checkpoint '{resume_from_checkpoint}' does not exist. Starting a new training run."
|
||||
)
|
||||
return dit, optimizer, lr_scheduler, dit_ema_dict, 0
|
||||
path = dirs[-1]
|
||||
else:
|
||||
path = os.path.basename(resume_from_checkpoint)
|
||||
|
||||
|
||||
accelerator.print(f"Resuming from checkpoint {path}")
|
||||
lora_state = load_file(os.path.join(project_dir, path, 'dit_lora.safetensors'), device=accelerator.device)
|
||||
unwarp_dit = accelerator.unwrap_model(dit)
|
||||
unwarp_dit.load_state_dict(lora_state, strict=False)
|
||||
if dit_ema_dict is not None:
|
||||
dit_ema_dict = load_file(
|
||||
os.path.join(project_dir, path, 'dit_lora_ema.safetensors'),
|
||||
device=accelerator.device
|
||||
)
|
||||
if dit is not unwarp_dit:
|
||||
dit_ema_dict = {f"module.{k}": v for k, v in dit_ema_dict.items() if k in unwarp_dit.state_dict()}
|
||||
|
||||
global_step = int(path.split("-")[1])
|
||||
|
||||
return dit, optimizer, lr_scheduler, dit_ema_dict, global_step
|
||||
|
||||
@dataclasses.dataclass
|
||||
class TrainArgs:
|
||||
## accelerator
|
||||
project_dir: str | None = None
|
||||
mixed_precision: Literal["no", "fp16", "bf16"] = "bf16"
|
||||
gradient_accumulation_steps: int = 1,
|
||||
seed: int = 42
|
||||
wandb_project_name: str | None = None
|
||||
wandb_run_name: str | None = None
|
||||
|
||||
## model
|
||||
model_name: Literal["flux", "flux-schnell"] = "flux"
|
||||
lora_rank: int = 512
|
||||
double_blocks_indices: list[int] | None = dataclasses.field(
|
||||
default=None,
|
||||
metadata={"help": "Indices of double blocks to apply LoRA. None means all double blocks."}
|
||||
)
|
||||
single_blocks_indices: list[int] | None = dataclasses.field(
|
||||
default=None,
|
||||
metadata={"help": "Indices of double blocks to apply LoRA. None means all single blocks."}
|
||||
)
|
||||
pe: Literal["d", "h", "w", "o"] = "d",
|
||||
gradient_checkpoint: bool = False
|
||||
|
||||
## ema
|
||||
ema: bool = False
|
||||
ema_interval: int = 1
|
||||
ema_decay: float = 0.99
|
||||
|
||||
|
||||
## optimizer
|
||||
learning_rate: float = 1e-2
|
||||
adam_betas: list[float] = dataclasses.field(default_factory=lambda: [0.9, 0.999])
|
||||
adam_eps: float = 1e-8
|
||||
adam_weight_decay: float = 0.01
|
||||
|
||||
## lr_scheduler
|
||||
lr_scheduler: str = "constant"
|
||||
lr_warmup_steps: int = 100
|
||||
max_train_steps: int = 100000
|
||||
|
||||
## dataloader
|
||||
train_data_json: str = "datasets/dreambench_singleip.json" # TODO: change to your own dataset, or use one data syenthsize pipeline comming in the future. stay tuned
|
||||
batch_size: int = 1
|
||||
text_dropout: float = 0.1
|
||||
resolution: int = 512
|
||||
resolution_ref: int | None = None
|
||||
|
||||
eval_data_json: str = "datasets/dreambench_singleip.json"
|
||||
eval_batch_size: int = 1
|
||||
|
||||
## misc
|
||||
resume_from_checkpoint: str | None | Literal["latest"] = None
|
||||
checkpointing_steps: int = 1000
|
||||
|
||||
def main(
|
||||
args: TrainArgs,
|
||||
):
|
||||
## accelerator
|
||||
deepspeed_plugins = {
|
||||
"dit": DeepSpeedPlugin(hf_ds_config='config/deepspeed/zero2_config.json'),
|
||||
"t5": DeepSpeedPlugin(hf_ds_config='config/deepspeed/zero3_config.json'),
|
||||
"clip": DeepSpeedPlugin(hf_ds_config='config/deepspeed/zero3_config.json')
|
||||
}
|
||||
accelerator = Accelerator(
|
||||
project_dir=args.project_dir,
|
||||
gradient_accumulation_steps=args.gradient_accumulation_steps,
|
||||
mixed_precision=args.mixed_precision,
|
||||
deepspeed_plugins=deepspeed_plugins,
|
||||
log_with="wandb",
|
||||
)
|
||||
set_seed(args.seed, device_specific=True)
|
||||
accelerator.init_trackers(
|
||||
project_name=args.wandb_project_name,
|
||||
config=args.__dict__,
|
||||
init_kwargs={
|
||||
"wandb": {
|
||||
"name": args.wandb_run_name,
|
||||
"dir": accelerator.project_dir,
|
||||
},
|
||||
},
|
||||
)
|
||||
weight_dtype = {
|
||||
"fp16": torch.float16,
|
||||
"bf16": torch.bfloat16,
|
||||
"no": torch.float32,
|
||||
}.get(accelerator.mixed_precision, torch.float32)
|
||||
|
||||
## logger
|
||||
logging.basicConfig(
|
||||
format=f"[RANK {accelerator.process_index}] " + "%(asctime)s - %(levelname)s - %(name)s - %(message)s",
|
||||
datefmt="%m/%d/%Y %H:%M:%S",
|
||||
level=logging.INFO,
|
||||
force=True
|
||||
)
|
||||
logger.info(accelerator.state)
|
||||
logger.info("Training script launched", main_process_only=False)
|
||||
|
||||
## model
|
||||
dit, vae, t5, clip = get_models(
|
||||
name=args.model_name,
|
||||
device=accelerator.device,
|
||||
)
|
||||
|
||||
vae.requires_grad_(False)
|
||||
t5.requires_grad_(False)
|
||||
clip.requires_grad_(False)
|
||||
|
||||
dit.requires_grad_(False)
|
||||
dit = set_lora(dit, args.lora_rank, args.double_blocks_indices, args.single_blocks_indices, accelerator.device)
|
||||
dit.train()
|
||||
dit.gradient_checkpointing = args.gradient_checkpoint
|
||||
|
||||
## optimizer and lr scheduler
|
||||
optimizer = torch.optim.AdamW(
|
||||
[p for p in dit.parameters() if p.requires_grad],
|
||||
lr=args.learning_rate,
|
||||
betas=args.adam_betas,
|
||||
weight_decay=args.adam_weight_decay,
|
||||
eps=args.adam_eps,
|
||||
)
|
||||
|
||||
lr_scheduler = get_scheduler(
|
||||
args.lr_scheduler,
|
||||
optimizer=optimizer,
|
||||
num_warmup_steps=args.lr_warmup_steps * accelerator.num_processes,
|
||||
num_training_steps=args.max_train_steps * accelerator.num_processes,
|
||||
)
|
||||
|
||||
# dataloader
|
||||
dataset = FluxPairedDatasetV2(
|
||||
data_json=args.train_data_json,
|
||||
resolution=args.resolution, resolution_ref=args.resolution_ref
|
||||
)
|
||||
dataloader = DataLoader(dataset, batch_size=args.batch_size, shuffle=True)
|
||||
eval_dataset = FluxPairedDatasetV2(
|
||||
data_json=args.eval_data_json,
|
||||
resolution=args.resolution, resolution_ref=args.resolution_ref
|
||||
)
|
||||
eval_dataloader = DataLoader(eval_dataset, batch_size=args.eval_batch_size, shuffle=False)
|
||||
|
||||
dataloader = accelerator.prepare_data_loader(dataloader)
|
||||
eval_dataloader = accelerator.prepare_data_loader(eval_dataloader)
|
||||
dataloader = itertools.cycle(dataloader) # as infinite fetch data loader
|
||||
|
||||
## parallel
|
||||
dit = accelerator.prepare_model(dit)
|
||||
optimizer = accelerator.prepare_optimizer(optimizer)
|
||||
lr_scheduler = accelerator.prepare_scheduler(lr_scheduler)
|
||||
accelerator.state.select_deepspeed_plugin("t5")
|
||||
t5 = accelerator.prepare_model(t5)
|
||||
accelerator.state.select_deepspeed_plugin("clip")
|
||||
clip = accelerator.prepare_model(clip)
|
||||
|
||||
## ema
|
||||
dit_ema_dict = {
|
||||
k: deepcopy(v).requires_grad_(False) for k, v in dit.named_parameters() if v.requires_grad
|
||||
} if args.ema else None
|
||||
|
||||
## resume
|
||||
(
|
||||
dit,
|
||||
optimizer,
|
||||
lr_scheduler,
|
||||
dit_ema_dict,
|
||||
global_step
|
||||
) = resume_from_checkpoint(
|
||||
args.resume_from_checkpoint,
|
||||
project_dir=args.project_dir,
|
||||
accelerator=accelerator,
|
||||
dit=dit,
|
||||
optimizer=optimizer,
|
||||
lr_scheduler=lr_scheduler,
|
||||
dit_ema_dict=dit_ema_dict
|
||||
)
|
||||
|
||||
## noise scheduler
|
||||
timesteps = get_schedule(
|
||||
999,
|
||||
(args.resolution // 8) * (args.resolution // 8) // 4,
|
||||
shift=True,
|
||||
)
|
||||
timesteps = torch.tensor(timesteps, device=accelerator.device)
|
||||
total_batch_size = args.batch_size * accelerator.num_processes * args.gradient_accumulation_steps
|
||||
|
||||
logger.info("***** Running training *****")
|
||||
logger.info(f" Instantaneous batch size per device = {args.batch_size}")
|
||||
logger.info(f" Total train batch size (w. parallel, distributed & accumulation) = {total_batch_size}")
|
||||
logger.info(f" Gradient Accumulation steps = {args.gradient_accumulation_steps}")
|
||||
logger.info(f" Total optimization steps = {args.max_train_steps}")
|
||||
logger.info(f" Total validation prompts = {len(eval_dataloader)}")
|
||||
|
||||
progress_bar = tqdm(
|
||||
range(0, args.max_train_steps),
|
||||
initial=global_step,
|
||||
desc="Steps",
|
||||
total=args.max_train_steps,
|
||||
disable=not accelerator.is_local_main_process,
|
||||
)
|
||||
|
||||
train_loss = 0.0
|
||||
while global_step < (args.max_train_steps):
|
||||
batch = next(dataloader)
|
||||
prompts = [txt_ if random.random() > args.text_dropout else "" for txt_ in batch["txt"]]
|
||||
img = batch["img"]
|
||||
ref_imgs = batch["ref_imgs"]
|
||||
|
||||
with torch.no_grad():
|
||||
x_1 = vae.encode(img.to(accelerator.device).to(torch.float32))
|
||||
x_ref = [vae.encode(ref_img.to(accelerator.device).to(torch.float32)) for ref_img in ref_imgs]
|
||||
inp = prepare_multi_ip(t5=t5, clip=clip, img=x_1, prompt=prompts, ref_imgs=tuple(x_ref), pe=args.pe)
|
||||
x_1 = rearrange(x_1, "b c (h ph) (w pw) -> b (h w) (c ph pw)", ph=2, pw=2)
|
||||
x_ref = [rearrange(x, "b c (h ph) (w pw) -> b (h w) (c ph pw)", ph=2, pw=2) for x in x_ref]
|
||||
|
||||
bs = img.shape[0]
|
||||
t = torch.randint(0, 1000, (bs,), device=accelerator.device)
|
||||
t = timesteps[t]
|
||||
x_0 = torch.randn_like(x_1, device=accelerator.device)
|
||||
x_t = (1 - t[:, None, None]) * x_1 + t[:, None, None] * x_0
|
||||
guidance_vec = torch.full((x_t.shape[0],), 1, device=x_t.device, dtype=x_t.dtype)
|
||||
|
||||
with accelerator.accumulate(dit):
|
||||
# Predict the noise residual and compute loss
|
||||
model_pred = dit(
|
||||
img=x_t.to(weight_dtype),
|
||||
img_ids=inp['img_ids'].to(weight_dtype),
|
||||
ref_img=[x.to(weight_dtype) for x in x_ref],
|
||||
ref_img_ids=[ref_img_id.to(weight_dtype) for ref_img_id in inp['ref_img_ids']],
|
||||
txt=inp['txt'].to(weight_dtype),
|
||||
txt_ids=inp['txt_ids'].to(weight_dtype),
|
||||
y=inp['vec'].to(weight_dtype),
|
||||
timesteps=t.to(weight_dtype),
|
||||
guidance=guidance_vec.to(weight_dtype)
|
||||
)
|
||||
|
||||
loss = F.mse_loss(model_pred.float(), (x_0 - x_1).float(), reduction="mean")
|
||||
|
||||
# Gather the losses across all processes for logging (if we use distributed training).
|
||||
avg_loss = accelerator.gather(loss.repeat(args.train_batch_size)).mean()
|
||||
train_loss += avg_loss.item() / args.gradient_accumulation_steps
|
||||
|
||||
# Backpropagate
|
||||
accelerator.backward(loss)
|
||||
if accelerator.sync_gradients:
|
||||
accelerator.clip_grad_norm_(dit.parameters(), args.max_grad_norm)
|
||||
optimizer.step()
|
||||
lr_scheduler.step()
|
||||
optimizer.zero_grad()
|
||||
|
||||
# Checks if the accelerator has performed an optimization step behind the scenes
|
||||
if accelerator.sync_gradients:
|
||||
progress_bar.update(1)
|
||||
global_step += 1
|
||||
accelerator.log({"train_loss": train_loss}, step=global_step)
|
||||
train_loss = 0.0
|
||||
|
||||
if accelerator.sync_gradients and dit_ema_dict is not None and global_step % args.ema_interval == 0:
|
||||
src_dict = dit.state_dict()
|
||||
for tgt_name in dit_ema_dict:
|
||||
dit_ema_dict[tgt_name].data.lerp_(src_dict[tgt_name].to(dit_ema_dict[tgt_name]), 1 - args.ema_decay)
|
||||
|
||||
if accelerator.sync_gradients and accelerator.is_main_process and global_step % args.checkpointing_steps == 0:
|
||||
logger.info(f"saving checkpoint in {global_step=}")
|
||||
save_path = os.path.join(args.project_dir, f"checkpoint-{global_step}")
|
||||
os.makedirs(save_path, exist_ok=True)
|
||||
|
||||
# save
|
||||
accelerator.wait_for_everyone()
|
||||
unwrapped_model = accelerator.unwrap_model(dit)
|
||||
unwrapped_model_state = unwrapped_model.state_dict()
|
||||
unwrapped_model_state = {k: v for k, v in unwrapped_model_state.items() if v.requires_grad}
|
||||
|
||||
accelerator.save(
|
||||
unwrapped_model_state,
|
||||
os.path.join(save_path, 'dit_lora.safetensors'),
|
||||
safe_serialization=True
|
||||
)
|
||||
unwrapped_opt = accelerator.unwrap_model(optimizer)
|
||||
accelerator.save(unwrapped_opt.state_dict(), os.path.join(save_path, 'optimizer.bin'))
|
||||
logger.info(f"Saved state to {save_path}")
|
||||
|
||||
if args.ema:
|
||||
accelerator.save(
|
||||
{k.split("module.")[-1]: v for k, v in dit_ema_dict.items()},
|
||||
os.path.join(save_path, 'dit_lora_ema.safetensors')
|
||||
)
|
||||
|
||||
# validate
|
||||
dit.eval()
|
||||
torch.set_grad_enabled(False)
|
||||
for i, batch in enumerate(eval_dataloader):
|
||||
result = inference(batch, dit, t5, clip, vae, accelerator, seed=0)
|
||||
accelerator.log({f"eval_gen_{i}": result}, step=global_step)
|
||||
|
||||
|
||||
if args.ema:
|
||||
original_state_dict = dit.state_dict()
|
||||
dit.load_state_dict(dit_ema_dict, strict=False)
|
||||
for batch in eval_dataloader:
|
||||
result = inference(batch, dit, t5, clip, vae, accelerator, seed=0)
|
||||
accelerator.log({f"eval_ema_gen_{i}": result}, step=global_step)
|
||||
dit.load_state_dict(original_state_dict, strict=False)
|
||||
|
||||
torch.cuda.empty_cache()
|
||||
gc.collect()
|
||||
torch.set_grad_enabled(True)
|
||||
dit.train()
|
||||
accelerator.wait_for_everyone()
|
||||
|
||||
logs = {"step_loss": loss.detach().item(), "lr": lr_scheduler.get_last_lr()[0]}
|
||||
progress_bar.set_postfix(**logs)
|
||||
|
||||
accelerator.wait_for_everyone()
|
||||
accelerator.end_training()
|
||||
|
||||
if __name__ == "__main__":
|
||||
parser = transformers.HfArgumentParser([TrainArgs])
|
||||
args_tuple = parser.parse_args_into_dataclasses(args_file_flag="--config")
|
||||
main(*args_tuple)
|
||||
@@ -0,0 +1,132 @@
|
||||
# Copyright (c) 2025 Bytedance Ltd. and/or its affiliates. All rights reserved.
|
||||
# Copyright (c) 2024 Black Forest Labs and The XLabs-AI 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.
|
||||
|
||||
import json
|
||||
import os
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
import torchvision.transforms.functional as TVF
|
||||
from torch.utils.data import DataLoader, Dataset
|
||||
from torchvision.transforms import Compose, Normalize, ToTensor
|
||||
|
||||
def bucket_images(images: list[torch.Tensor], resolution: int = 512):
|
||||
bucket_override=[
|
||||
# h w
|
||||
(256, 768),
|
||||
(320, 768),
|
||||
(320, 704),
|
||||
(384, 640),
|
||||
(448, 576),
|
||||
(512, 512),
|
||||
(576, 448),
|
||||
(640, 384),
|
||||
(704, 320),
|
||||
(768, 320),
|
||||
(768, 256)
|
||||
]
|
||||
bucket_override = [(int(h / 512 * resolution), int(w / 512 * resolution)) for h, w in bucket_override]
|
||||
bucket_override = [(h // 16 * 16, w // 16 * 16) for h, w in bucket_override]
|
||||
|
||||
aspect_ratios = [image.shape[-2] / image.shape[-1] for image in images]
|
||||
mean_aspect_ratio = np.mean(aspect_ratios)
|
||||
|
||||
new_h, new_w = bucket_override[0]
|
||||
min_aspect_diff = np.abs(new_h / new_w - mean_aspect_ratio)
|
||||
for h, w in bucket_override:
|
||||
aspect_diff = np.abs(h / w - mean_aspect_ratio)
|
||||
if aspect_diff < min_aspect_diff:
|
||||
min_aspect_diff = aspect_diff
|
||||
new_h, new_w = h, w
|
||||
|
||||
images = [TVF.resize(image, (new_h, new_w)) for image in images]
|
||||
images = torch.stack(images, dim=0)
|
||||
return images
|
||||
|
||||
class FluxPairedDatasetV2(Dataset):
|
||||
def __init__(self, json_file: str, resolution: int, resolution_ref: int | None = None):
|
||||
super().__init__()
|
||||
self.json_file = json_file
|
||||
self.resolution = resolution
|
||||
self.resolution_ref = resolution_ref if resolution_ref is not None else resolution
|
||||
self.image_root = os.path.dirname(json_file)
|
||||
|
||||
with open(self.json_file, "rt") as f:
|
||||
self.data_dicts = json.load(f)
|
||||
|
||||
self.transform = Compose([
|
||||
ToTensor(),
|
||||
Normalize([0.5], [0.5]),
|
||||
])
|
||||
|
||||
def __getitem__(self, idx):
|
||||
data_dict = self.data_dicts[idx]
|
||||
image_paths = [data_dict["image_path"]] if "image_path" in data_dict else data_dict["image_paths"]
|
||||
txt = data_dict["prompt"]
|
||||
image_tgt_path = data_dict.get("image_tgt_path", None)
|
||||
ref_imgs = [
|
||||
Image.open(os.path.join(self.image_root, path)).convert("RGB")
|
||||
for path in image_paths
|
||||
]
|
||||
ref_imgs = [self.transform(img) for img in ref_imgs]
|
||||
img = None
|
||||
if image_tgt_path is not None:
|
||||
img = Image.open(os.path.join(self.image_root, image_tgt_path)).convert("RGB")
|
||||
img = self.transform(img)
|
||||
|
||||
return {
|
||||
"img": img,
|
||||
"txt": txt,
|
||||
"ref_imgs": ref_imgs,
|
||||
}
|
||||
|
||||
def __len__(self):
|
||||
return len(self.data_dicts)
|
||||
|
||||
def collate_fn(self, batch):
|
||||
img = [data["img"] for data in batch]
|
||||
txt = [data["txt"] for data in batch]
|
||||
ref_imgs = [data["ref_imgs"] for data in batch]
|
||||
assert all([len(ref_imgs[0]) == len(ref_imgs[i]) for i in range(len(ref_imgs))])
|
||||
|
||||
n_ref = len(ref_imgs[0])
|
||||
|
||||
img = bucket_images(img, self.resolution)
|
||||
ref_imgs_new = []
|
||||
for i in range(n_ref):
|
||||
ref_imgs_i = [refs[i] for refs in ref_imgs]
|
||||
ref_imgs_i = bucket_images(ref_imgs_i, self.resolution_ref)
|
||||
ref_imgs_new.append(ref_imgs_i)
|
||||
|
||||
return {
|
||||
"txt": txt,
|
||||
"img": img,
|
||||
"ref_imgs": ref_imgs_new,
|
||||
}
|
||||
|
||||
if __name__ == '__main__':
|
||||
import argparse
|
||||
from pprint import pprint
|
||||
parser = argparse.ArgumentParser()
|
||||
# parser.add_argument("--json_file", type=str, required=True)
|
||||
parser.add_argument("--json_file", type=str, default="datasets/fake_train_data.json")
|
||||
args = parser.parse_args()
|
||||
dataset = FluxPairedDatasetV2(args.json_file, 512)
|
||||
dataloder = DataLoader(dataset, batch_size=4, collate_fn=dataset.collate_fn)
|
||||
|
||||
for i, data_dict in enumerate(dataloder):
|
||||
pprint(i)
|
||||
pprint(data_dict)
|
||||
breakpoint()
|
||||
@@ -0,0 +1,45 @@
|
||||
# Copyright (c) 2025 Bytedance Ltd. and/or its affiliates. All rights reserved.
|
||||
# Copyright (c) 2024 Black Forest Labs and The XLabs-AI 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.
|
||||
|
||||
import torch
|
||||
from einops import rearrange
|
||||
from torch import Tensor
|
||||
|
||||
|
||||
def attention(q: Tensor, k: Tensor, v: Tensor, pe: Tensor) -> Tensor:
|
||||
q, k = apply_rope(q, k, pe)
|
||||
|
||||
x = torch.nn.functional.scaled_dot_product_attention(q, k, v)
|
||||
x = rearrange(x, "B H L D -> B L (H D)")
|
||||
|
||||
return x
|
||||
|
||||
|
||||
def rope(pos: Tensor, dim: int, theta: int) -> Tensor:
|
||||
assert dim % 2 == 0
|
||||
scale = torch.arange(0, dim, 2, dtype=torch.float64, device=pos.device) / dim
|
||||
omega = 1.0 / (theta**scale)
|
||||
out = torch.einsum("...n,d->...nd", pos, omega)
|
||||
out = torch.stack([torch.cos(out), -torch.sin(out), torch.sin(out), torch.cos(out)], dim=-1)
|
||||
out = rearrange(out, "b n d (i j) -> b n d i j", i=2, j=2)
|
||||
return out.float()
|
||||
|
||||
|
||||
def apply_rope(xq: Tensor, xk: Tensor, freqs_cis: Tensor) -> tuple[Tensor, Tensor]:
|
||||
xq_ = xq.float().reshape(*xq.shape[:-1], -1, 1, 2)
|
||||
xk_ = xk.float().reshape(*xk.shape[:-1], -1, 1, 2)
|
||||
xq_out = freqs_cis[..., 0] * xq_[..., 0] + freqs_cis[..., 1] * xq_[..., 1]
|
||||
xk_out = freqs_cis[..., 0] * xk_[..., 0] + freqs_cis[..., 1] * xk_[..., 1]
|
||||
return xq_out.reshape(*xq.shape).type_as(xq), xk_out.reshape(*xk.shape).type_as(xk)
|
||||
@@ -0,0 +1,223 @@
|
||||
# Copyright (c) 2025 Bytedance Ltd. and/or its affiliates. All rights reserved.
|
||||
# Copyright (c) 2024 Black Forest Labs and The XLabs-AI 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 dataclasses import dataclass
|
||||
|
||||
import torch
|
||||
from torch import Tensor, nn
|
||||
|
||||
from .modules.layers import DoubleStreamBlock, EmbedND, LastLayer, MLPEmbedder, SingleStreamBlock, timestep_embedding
|
||||
|
||||
|
||||
@dataclass
|
||||
class FluxParams:
|
||||
in_channels: int
|
||||
vec_in_dim: int
|
||||
context_in_dim: int
|
||||
hidden_size: int
|
||||
mlp_ratio: float
|
||||
num_heads: int
|
||||
depth: int
|
||||
depth_single_blocks: int
|
||||
axes_dim: list[int]
|
||||
theta: int
|
||||
qkv_bias: bool
|
||||
guidance_embed: bool
|
||||
|
||||
|
||||
class Flux(nn.Module):
|
||||
"""
|
||||
Transformer model for flow matching on sequences.
|
||||
"""
|
||||
_supports_gradient_checkpointing = True
|
||||
|
||||
def __init__(self, params: FluxParams):
|
||||
super().__init__()
|
||||
|
||||
self.params = params
|
||||
self.in_channels = params.in_channels
|
||||
self.out_channels = self.in_channels
|
||||
if params.hidden_size % params.num_heads != 0:
|
||||
raise ValueError(
|
||||
f"Hidden size {params.hidden_size} must be divisible by num_heads {params.num_heads}"
|
||||
)
|
||||
pe_dim = params.hidden_size // params.num_heads
|
||||
if sum(params.axes_dim) != pe_dim:
|
||||
raise ValueError(f"Got {params.axes_dim} but expected positional dim {pe_dim}")
|
||||
self.hidden_size = params.hidden_size
|
||||
self.num_heads = params.num_heads
|
||||
self.pe_embedder = EmbedND(dim=pe_dim, theta=params.theta, axes_dim=params.axes_dim)
|
||||
self.img_in = nn.Linear(self.in_channels, self.hidden_size, bias=True)
|
||||
self.time_in = MLPEmbedder(in_dim=256, hidden_dim=self.hidden_size)
|
||||
self.vector_in = MLPEmbedder(params.vec_in_dim, self.hidden_size)
|
||||
self.guidance_in = (
|
||||
MLPEmbedder(in_dim=256, hidden_dim=self.hidden_size) if params.guidance_embed else nn.Identity()
|
||||
)
|
||||
self.txt_in = nn.Linear(params.context_in_dim, self.hidden_size)
|
||||
|
||||
self.double_blocks = nn.ModuleList(
|
||||
[
|
||||
DoubleStreamBlock(
|
||||
self.hidden_size,
|
||||
self.num_heads,
|
||||
mlp_ratio=params.mlp_ratio,
|
||||
qkv_bias=params.qkv_bias,
|
||||
)
|
||||
for _ in range(params.depth)
|
||||
]
|
||||
)
|
||||
|
||||
self.single_blocks = nn.ModuleList(
|
||||
[
|
||||
SingleStreamBlock(self.hidden_size, self.num_heads, mlp_ratio=params.mlp_ratio)
|
||||
for _ in range(params.depth_single_blocks)
|
||||
]
|
||||
)
|
||||
|
||||
self.final_layer = LastLayer(self.hidden_size, 1, self.out_channels)
|
||||
self.gradient_checkpointing = False
|
||||
|
||||
def _set_gradient_checkpointing(self, module, value=False):
|
||||
if hasattr(module, "gradient_checkpointing"):
|
||||
module.gradient_checkpointing = value
|
||||
|
||||
@property
|
||||
def attn_processors(self):
|
||||
# set recursively
|
||||
processors = {} # type: dict[str, nn.Module]
|
||||
|
||||
def fn_recursive_add_processors(name: str, module: torch.nn.Module, processors):
|
||||
if hasattr(module, "set_processor"):
|
||||
processors[f"{name}.processor"] = module.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
|
||||
|
||||
def set_attn_processor(self, processor):
|
||||
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)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
img: Tensor,
|
||||
img_ids: Tensor,
|
||||
txt: Tensor,
|
||||
txt_ids: Tensor,
|
||||
timesteps: Tensor,
|
||||
y: Tensor,
|
||||
guidance: Tensor | None = None,
|
||||
ref_img: Tensor | None = None,
|
||||
ref_img_ids: Tensor | None = None,
|
||||
) -> Tensor:
|
||||
if img.ndim != 3 or txt.ndim != 3:
|
||||
raise ValueError("Input img and txt tensors must have 3 dimensions.")
|
||||
|
||||
# running on sequences img
|
||||
#print(img)
|
||||
img = self.img_in(img)
|
||||
vec = self.time_in(timestep_embedding(timesteps, 256))
|
||||
if self.params.guidance_embed:
|
||||
if guidance is None:
|
||||
raise ValueError("Didn't get guidance strength for guidance distilled model.")
|
||||
vec = vec + self.guidance_in(timestep_embedding(guidance, 256))
|
||||
vec = vec + self.vector_in(y)
|
||||
txt = self.txt_in(txt)
|
||||
|
||||
ids = torch.cat((txt_ids, img_ids), dim=1)
|
||||
|
||||
# concat ref_img/img
|
||||
img_end = img.shape[1]
|
||||
if ref_img is not None:
|
||||
if isinstance(ref_img, tuple) or isinstance(ref_img, list):
|
||||
img_in = [img] + [self.img_in(ref.to(img.device)) for ref in ref_img] #确保在同一个设备
|
||||
img_ids = [ids] + [ref_ids for ref_ids in ref_img_ids]
|
||||
img = torch.cat(img_in, dim=1)
|
||||
ids = torch.cat(img_ids, dim=1)
|
||||
else:
|
||||
img = torch.cat((img, self.img_in(ref_img)), dim=1)
|
||||
ids = torch.cat((ids, ref_img_ids), dim=1)
|
||||
pe = self.pe_embedder(ids)
|
||||
|
||||
for index_block, block in enumerate(self.double_blocks):
|
||||
if self.training and self.gradient_checkpointing:
|
||||
img, txt = torch.utils.checkpoint.checkpoint(
|
||||
block,
|
||||
img=img,
|
||||
txt=txt,
|
||||
vec=vec,
|
||||
pe=pe,
|
||||
use_reentrant=False,
|
||||
)
|
||||
else:
|
||||
img, txt = block(
|
||||
img=img,
|
||||
txt=txt,
|
||||
vec=vec,
|
||||
pe=pe
|
||||
)
|
||||
|
||||
img = torch.cat((txt, img), 1)
|
||||
for block in self.single_blocks:
|
||||
if self.training and self.gradient_checkpointing:
|
||||
img = torch.utils.checkpoint.checkpoint(
|
||||
block,
|
||||
img, vec=vec, pe=pe,
|
||||
use_reentrant=False
|
||||
)
|
||||
else:
|
||||
img = block(img, vec=vec, pe=pe)
|
||||
img = img[:, txt.shape[1] :, ...]
|
||||
# index img
|
||||
img = img[:, :img_end, ...]
|
||||
|
||||
img = self.final_layer(img, vec) # (N, T, patch_size ** 2 * out_channels)
|
||||
return img
|
||||
@@ -0,0 +1,327 @@
|
||||
# Copyright (c) 2025 Bytedance Ltd. and/or its affiliates. All rights reserved.
|
||||
# Copyright (c) 2024 Black Forest Labs and The XLabs-AI 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 dataclasses import dataclass
|
||||
|
||||
import torch
|
||||
from einops import rearrange
|
||||
from torch import Tensor, nn
|
||||
|
||||
|
||||
@dataclass
|
||||
class AutoEncoderParams:
|
||||
resolution: int
|
||||
in_channels: int
|
||||
ch: int
|
||||
out_ch: int
|
||||
ch_mult: list[int]
|
||||
num_res_blocks: int
|
||||
z_channels: int
|
||||
scale_factor: float
|
||||
shift_factor: float
|
||||
|
||||
|
||||
def swish(x: Tensor) -> Tensor:
|
||||
return x * torch.sigmoid(x)
|
||||
|
||||
|
||||
class AttnBlock(nn.Module):
|
||||
def __init__(self, in_channels: int):
|
||||
super().__init__()
|
||||
self.in_channels = in_channels
|
||||
|
||||
self.norm = nn.GroupNorm(num_groups=32, num_channels=in_channels, eps=1e-6, affine=True)
|
||||
|
||||
self.q = nn.Conv2d(in_channels, in_channels, kernel_size=1)
|
||||
self.k = nn.Conv2d(in_channels, in_channels, kernel_size=1)
|
||||
self.v = nn.Conv2d(in_channels, in_channels, kernel_size=1)
|
||||
self.proj_out = nn.Conv2d(in_channels, in_channels, kernel_size=1)
|
||||
|
||||
def attention(self, h_: Tensor) -> Tensor:
|
||||
h_ = self.norm(h_)
|
||||
q = self.q(h_)
|
||||
k = self.k(h_)
|
||||
v = self.v(h_)
|
||||
|
||||
b, c, h, w = q.shape
|
||||
q = rearrange(q, "b c h w -> b 1 (h w) c").contiguous()
|
||||
k = rearrange(k, "b c h w -> b 1 (h w) c").contiguous()
|
||||
v = rearrange(v, "b c h w -> b 1 (h w) c").contiguous()
|
||||
h_ = nn.functional.scaled_dot_product_attention(q, k, v)
|
||||
|
||||
return rearrange(h_, "b 1 (h w) c -> b c h w", h=h, w=w, c=c, b=b)
|
||||
|
||||
def forward(self, x: Tensor) -> Tensor:
|
||||
return x + self.proj_out(self.attention(x))
|
||||
|
||||
|
||||
class ResnetBlock(nn.Module):
|
||||
def __init__(self, in_channels: int, out_channels: int):
|
||||
super().__init__()
|
||||
self.in_channels = in_channels
|
||||
out_channels = in_channels if out_channels is None else out_channels
|
||||
self.out_channels = out_channels
|
||||
|
||||
self.norm1 = nn.GroupNorm(num_groups=32, num_channels=in_channels, eps=1e-6, affine=True)
|
||||
self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=1, padding=1)
|
||||
self.norm2 = nn.GroupNorm(num_groups=32, num_channels=out_channels, eps=1e-6, affine=True)
|
||||
self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size=3, stride=1, padding=1)
|
||||
if self.in_channels != self.out_channels:
|
||||
self.nin_shortcut = nn.Conv2d(in_channels, out_channels, kernel_size=1, stride=1, padding=0)
|
||||
|
||||
def forward(self, x):
|
||||
h = x
|
||||
h = self.norm1(h)
|
||||
h = swish(h)
|
||||
h = self.conv1(h)
|
||||
|
||||
h = self.norm2(h)
|
||||
h = swish(h)
|
||||
h = self.conv2(h)
|
||||
|
||||
if self.in_channels != self.out_channels:
|
||||
x = self.nin_shortcut(x)
|
||||
|
||||
return x + h
|
||||
|
||||
|
||||
class Downsample(nn.Module):
|
||||
def __init__(self, in_channels: int):
|
||||
super().__init__()
|
||||
# no asymmetric padding in torch conv, must do it ourselves
|
||||
self.conv = nn.Conv2d(in_channels, in_channels, kernel_size=3, stride=2, padding=0)
|
||||
|
||||
def forward(self, x: Tensor):
|
||||
pad = (0, 1, 0, 1)
|
||||
x = nn.functional.pad(x, pad, mode="constant", value=0)
|
||||
x = self.conv(x)
|
||||
return x
|
||||
|
||||
|
||||
class Upsample(nn.Module):
|
||||
def __init__(self, in_channels: int):
|
||||
super().__init__()
|
||||
self.conv = nn.Conv2d(in_channels, in_channels, kernel_size=3, stride=1, padding=1)
|
||||
|
||||
def forward(self, x: Tensor):
|
||||
x = nn.functional.interpolate(x, scale_factor=2.0, mode="nearest")
|
||||
x = self.conv(x)
|
||||
return x
|
||||
|
||||
|
||||
class Encoder(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
resolution: int,
|
||||
in_channels: int,
|
||||
ch: int,
|
||||
ch_mult: list[int],
|
||||
num_res_blocks: int,
|
||||
z_channels: int,
|
||||
):
|
||||
super().__init__()
|
||||
self.ch = ch
|
||||
self.num_resolutions = len(ch_mult)
|
||||
self.num_res_blocks = num_res_blocks
|
||||
self.resolution = resolution
|
||||
self.in_channels = in_channels
|
||||
# downsampling
|
||||
self.conv_in = nn.Conv2d(in_channels, self.ch, kernel_size=3, stride=1, padding=1)
|
||||
|
||||
curr_res = resolution
|
||||
in_ch_mult = (1,) + tuple(ch_mult)
|
||||
self.in_ch_mult = in_ch_mult
|
||||
self.down = nn.ModuleList()
|
||||
block_in = self.ch
|
||||
for i_level in range(self.num_resolutions):
|
||||
block = nn.ModuleList()
|
||||
attn = nn.ModuleList()
|
||||
block_in = ch * in_ch_mult[i_level]
|
||||
block_out = ch * ch_mult[i_level]
|
||||
for _ in range(self.num_res_blocks):
|
||||
block.append(ResnetBlock(in_channels=block_in, out_channels=block_out))
|
||||
block_in = block_out
|
||||
down = nn.Module()
|
||||
down.block = block
|
||||
down.attn = attn
|
||||
if i_level != self.num_resolutions - 1:
|
||||
down.downsample = Downsample(block_in)
|
||||
curr_res = curr_res // 2
|
||||
self.down.append(down)
|
||||
|
||||
# middle
|
||||
self.mid = nn.Module()
|
||||
self.mid.block_1 = ResnetBlock(in_channels=block_in, out_channels=block_in)
|
||||
self.mid.attn_1 = AttnBlock(block_in)
|
||||
self.mid.block_2 = ResnetBlock(in_channels=block_in, out_channels=block_in)
|
||||
|
||||
# end
|
||||
self.norm_out = nn.GroupNorm(num_groups=32, num_channels=block_in, eps=1e-6, affine=True)
|
||||
self.conv_out = nn.Conv2d(block_in, 2 * z_channels, kernel_size=3, stride=1, padding=1)
|
||||
|
||||
def forward(self, x: Tensor) -> Tensor:
|
||||
# downsampling
|
||||
hs = [self.conv_in(x)]
|
||||
for i_level in range(self.num_resolutions):
|
||||
for i_block in range(self.num_res_blocks):
|
||||
h = self.down[i_level].block[i_block](hs[-1])
|
||||
if len(self.down[i_level].attn) > 0:
|
||||
h = self.down[i_level].attn[i_block](h)
|
||||
hs.append(h)
|
||||
if i_level != self.num_resolutions - 1:
|
||||
hs.append(self.down[i_level].downsample(hs[-1]))
|
||||
|
||||
# middle
|
||||
h = hs[-1]
|
||||
h = self.mid.block_1(h)
|
||||
h = self.mid.attn_1(h)
|
||||
h = self.mid.block_2(h)
|
||||
# end
|
||||
h = self.norm_out(h)
|
||||
h = swish(h)
|
||||
h = self.conv_out(h)
|
||||
return h
|
||||
|
||||
|
||||
class Decoder(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
ch: int,
|
||||
out_ch: int,
|
||||
ch_mult: list[int],
|
||||
num_res_blocks: int,
|
||||
in_channels: int,
|
||||
resolution: int,
|
||||
z_channels: int,
|
||||
):
|
||||
super().__init__()
|
||||
self.ch = ch
|
||||
self.num_resolutions = len(ch_mult)
|
||||
self.num_res_blocks = num_res_blocks
|
||||
self.resolution = resolution
|
||||
self.in_channels = in_channels
|
||||
self.ffactor = 2 ** (self.num_resolutions - 1)
|
||||
|
||||
# compute in_ch_mult, block_in and curr_res at lowest res
|
||||
block_in = ch * ch_mult[self.num_resolutions - 1]
|
||||
curr_res = resolution // 2 ** (self.num_resolutions - 1)
|
||||
self.z_shape = (1, z_channels, curr_res, curr_res)
|
||||
|
||||
# z to block_in
|
||||
self.conv_in = nn.Conv2d(z_channels, block_in, kernel_size=3, stride=1, padding=1)
|
||||
|
||||
# middle
|
||||
self.mid = nn.Module()
|
||||
self.mid.block_1 = ResnetBlock(in_channels=block_in, out_channels=block_in)
|
||||
self.mid.attn_1 = AttnBlock(block_in)
|
||||
self.mid.block_2 = ResnetBlock(in_channels=block_in, out_channels=block_in)
|
||||
|
||||
# upsampling
|
||||
self.up = nn.ModuleList()
|
||||
for i_level in reversed(range(self.num_resolutions)):
|
||||
block = nn.ModuleList()
|
||||
attn = nn.ModuleList()
|
||||
block_out = ch * ch_mult[i_level]
|
||||
for _ in range(self.num_res_blocks + 1):
|
||||
block.append(ResnetBlock(in_channels=block_in, out_channels=block_out))
|
||||
block_in = block_out
|
||||
up = nn.Module()
|
||||
up.block = block
|
||||
up.attn = attn
|
||||
if i_level != 0:
|
||||
up.upsample = Upsample(block_in)
|
||||
curr_res = curr_res * 2
|
||||
self.up.insert(0, up) # prepend to get consistent order
|
||||
|
||||
# end
|
||||
self.norm_out = nn.GroupNorm(num_groups=32, num_channels=block_in, eps=1e-6, affine=True)
|
||||
self.conv_out = nn.Conv2d(block_in, out_ch, kernel_size=3, stride=1, padding=1)
|
||||
|
||||
def forward(self, z: Tensor) -> Tensor:
|
||||
# z to block_in
|
||||
h = self.conv_in(z)
|
||||
|
||||
# middle
|
||||
h = self.mid.block_1(h)
|
||||
h = self.mid.attn_1(h)
|
||||
h = self.mid.block_2(h)
|
||||
|
||||
# upsampling
|
||||
for i_level in reversed(range(self.num_resolutions)):
|
||||
for i_block in range(self.num_res_blocks + 1):
|
||||
h = self.up[i_level].block[i_block](h)
|
||||
if len(self.up[i_level].attn) > 0:
|
||||
h = self.up[i_level].attn[i_block](h)
|
||||
if i_level != 0:
|
||||
h = self.up[i_level].upsample(h)
|
||||
|
||||
# end
|
||||
h = self.norm_out(h)
|
||||
h = swish(h)
|
||||
h = self.conv_out(h)
|
||||
return h
|
||||
|
||||
|
||||
class DiagonalGaussian(nn.Module):
|
||||
def __init__(self, sample: bool = True, chunk_dim: int = 1):
|
||||
super().__init__()
|
||||
self.sample = sample
|
||||
self.chunk_dim = chunk_dim
|
||||
|
||||
def forward(self, z: Tensor) -> Tensor:
|
||||
mean, logvar = torch.chunk(z, 2, dim=self.chunk_dim)
|
||||
if self.sample:
|
||||
std = torch.exp(0.5 * logvar)
|
||||
return mean + std * torch.randn_like(mean)
|
||||
else:
|
||||
return mean
|
||||
|
||||
|
||||
class AutoEncoder(nn.Module):
|
||||
def __init__(self, params: AutoEncoderParams):
|
||||
super().__init__()
|
||||
self.encoder = Encoder(
|
||||
resolution=params.resolution,
|
||||
in_channels=params.in_channels,
|
||||
ch=params.ch,
|
||||
ch_mult=params.ch_mult,
|
||||
num_res_blocks=params.num_res_blocks,
|
||||
z_channels=params.z_channels,
|
||||
)
|
||||
self.decoder = Decoder(
|
||||
resolution=params.resolution,
|
||||
in_channels=params.in_channels,
|
||||
ch=params.ch,
|
||||
out_ch=params.out_ch,
|
||||
ch_mult=params.ch_mult,
|
||||
num_res_blocks=params.num_res_blocks,
|
||||
z_channels=params.z_channels,
|
||||
)
|
||||
self.reg = DiagonalGaussian()
|
||||
|
||||
self.scale_factor = params.scale_factor
|
||||
self.shift_factor = params.shift_factor
|
||||
|
||||
def encode(self, x: Tensor) -> Tensor:
|
||||
z = self.reg(self.encoder(x))
|
||||
z = self.scale_factor * (z - self.shift_factor)
|
||||
return z
|
||||
|
||||
def decode(self, z: Tensor) -> Tensor:
|
||||
z = z / self.scale_factor + self.shift_factor
|
||||
return self.decoder(z)
|
||||
|
||||
def forward(self, x: Tensor) -> Tensor:
|
||||
return self.decode(self.encode(x))
|
||||
@@ -0,0 +1,53 @@
|
||||
# Copyright (c) 2025 Bytedance Ltd. and/or its affiliates. All rights reserved.
|
||||
# Copyright (c) 2024 Black Forest Labs and The XLabs-AI 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 torch import Tensor, nn
|
||||
from transformers import (CLIPTextModel, CLIPTokenizer, T5EncoderModel,
|
||||
T5Tokenizer)
|
||||
|
||||
|
||||
class HFEmbedder(nn.Module):
|
||||
def __init__(self, version: str, max_length: int, **hf_kwargs):
|
||||
super().__init__()
|
||||
self.is_clip = version.startswith("openai")
|
||||
self.max_length = max_length
|
||||
self.output_key = "pooler_output" if self.is_clip else "last_hidden_state"
|
||||
|
||||
if self.is_clip:
|
||||
self.tokenizer: CLIPTokenizer = CLIPTokenizer.from_pretrained(version, max_length=max_length)
|
||||
self.hf_module: CLIPTextModel = CLIPTextModel.from_pretrained(version, **hf_kwargs)
|
||||
else:
|
||||
self.tokenizer: T5Tokenizer = T5Tokenizer.from_pretrained(version, max_length=max_length)
|
||||
self.hf_module: T5EncoderModel = T5EncoderModel.from_pretrained(version, **hf_kwargs)
|
||||
|
||||
self.hf_module = self.hf_module.eval().requires_grad_(False)
|
||||
|
||||
def forward(self, text: list[str]) -> Tensor:
|
||||
batch_encoding = self.tokenizer(
|
||||
text,
|
||||
truncation=True,
|
||||
max_length=self.max_length,
|
||||
return_length=False,
|
||||
return_overflowing_tokens=False,
|
||||
padding="max_length",
|
||||
return_tensors="pt",
|
||||
)
|
||||
|
||||
outputs = self.hf_module(
|
||||
input_ids=batch_encoding["input_ids"].to(self.hf_module.device),
|
||||
attention_mask=None,
|
||||
output_hidden_states=False,
|
||||
)
|
||||
return outputs[self.output_key]
|
||||
@@ -0,0 +1,436 @@
|
||||
# Copyright (c) 2025 Bytedance Ltd. and/or its affiliates. All rights reserved.
|
||||
# Copyright (c) 2024 Black Forest Labs and The XLabs-AI 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.
|
||||
|
||||
import math
|
||||
from dataclasses import dataclass
|
||||
|
||||
import torch
|
||||
from einops import rearrange
|
||||
from torch import Tensor, nn
|
||||
|
||||
from ..math import attention, rope
|
||||
import torch.nn.functional as F
|
||||
|
||||
class EmbedND(nn.Module):
|
||||
def __init__(self, dim: int, theta: int, axes_dim: list[int]):
|
||||
super().__init__()
|
||||
self.dim = dim
|
||||
self.theta = theta
|
||||
self.axes_dim = axes_dim
|
||||
|
||||
def forward(self, ids: Tensor) -> Tensor:
|
||||
n_axes = ids.shape[-1]
|
||||
emb = torch.cat(
|
||||
[rope(ids[..., i], self.axes_dim[i], self.theta) for i in range(n_axes)],
|
||||
dim=-3,
|
||||
)
|
||||
|
||||
return emb.unsqueeze(1)
|
||||
|
||||
|
||||
def timestep_embedding(t: Tensor, dim, max_period=10000, time_factor: float = 1000.0):
|
||||
"""
|
||||
Create sinusoidal timestep embeddings.
|
||||
:param t: a 1-D Tensor of N indices, one per batch element.
|
||||
These may be fractional.
|
||||
:param dim: the dimension of the output.
|
||||
:param max_period: controls the minimum frequency of the embeddings.
|
||||
:return: an (N, D) Tensor of positional embeddings.
|
||||
"""
|
||||
t = time_factor * t
|
||||
half = dim // 2
|
||||
freqs = torch.exp(-math.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32) / half).to(
|
||||
t.device
|
||||
)
|
||||
|
||||
args = t[:, None].float() * freqs[None]
|
||||
embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
|
||||
if dim % 2:
|
||||
embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1)
|
||||
if torch.is_floating_point(t):
|
||||
embedding = embedding.to(t)
|
||||
return embedding
|
||||
|
||||
|
||||
class MLPEmbedder(nn.Module):
|
||||
def __init__(self, in_dim: int, hidden_dim: int):
|
||||
super().__init__()
|
||||
self.in_layer = nn.Linear(in_dim, hidden_dim, bias=True)
|
||||
self.silu = nn.SiLU()
|
||||
self.out_layer = nn.Linear(hidden_dim, hidden_dim, bias=True)
|
||||
|
||||
def forward(self, x: Tensor) -> Tensor:
|
||||
return self.out_layer(self.silu(self.in_layer(x)))
|
||||
|
||||
|
||||
class RMSNorm(torch.nn.Module):
|
||||
def __init__(self, dim: int):
|
||||
super().__init__()
|
||||
self.scale = nn.Parameter(torch.ones(dim))
|
||||
|
||||
def forward(self, x: Tensor):
|
||||
x_dtype = x.dtype
|
||||
x = x.float()
|
||||
rrms = torch.rsqrt(torch.mean(x**2, dim=-1, keepdim=True) + 1e-6)
|
||||
#return (x * rrms).to(dtype=x_dtype) * self.scale
|
||||
return ((x * rrms) * self.scale.float()).to(dtype=x_dtype)
|
||||
|
||||
|
||||
class QKNorm(torch.nn.Module):
|
||||
def __init__(self, dim: int):
|
||||
super().__init__()
|
||||
self.query_norm = RMSNorm(dim)
|
||||
self.key_norm = RMSNorm(dim)
|
||||
|
||||
def forward(self, q: Tensor, k: Tensor, v: Tensor) -> tuple[Tensor, Tensor]:
|
||||
q = self.query_norm(q)
|
||||
k = self.key_norm(k)
|
||||
return q.to(v), k.to(v)
|
||||
|
||||
class LoRALinearLayer(nn.Module):
|
||||
def __init__(self, in_features, out_features, rank=4, network_alpha=None, device=None, dtype=None):
|
||||
super().__init__()
|
||||
|
||||
self.down = nn.Linear(in_features, rank, bias=False, device=device, dtype=dtype)
|
||||
self.up = nn.Linear(rank, out_features, bias=False, device=device, dtype=dtype)
|
||||
# This value has the same meaning as the `--network_alpha` option in the kohya-ss trainer script.
|
||||
# See https://github.com/darkstorm2150/sd-scripts/blob/main/docs/train_network_README-en.md#execute-learning
|
||||
self.network_alpha = network_alpha
|
||||
self.rank = rank
|
||||
|
||||
nn.init.normal_(self.down.weight, std=1 / rank)
|
||||
nn.init.zeros_(self.up.weight)
|
||||
|
||||
def forward(self, hidden_states):
|
||||
orig_dtype = hidden_states.dtype
|
||||
dtype = self.down.weight.dtype
|
||||
|
||||
down_hidden_states = self.down(hidden_states.to(dtype))
|
||||
up_hidden_states = self.up(down_hidden_states)
|
||||
|
||||
if self.network_alpha is not None:
|
||||
up_hidden_states *= self.network_alpha / self.rank
|
||||
|
||||
return up_hidden_states.to(orig_dtype)
|
||||
|
||||
class FLuxSelfAttnProcessor:
|
||||
def __call__(self, attn, x, pe, **attention_kwargs):
|
||||
qkv = attn.qkv(x)
|
||||
q, k, v = rearrange(qkv, "B L (K H D) -> K B H L D", K=3, H=self.num_heads)
|
||||
q, k = attn.norm(q, k, v)
|
||||
x = attention(q, k, v, pe=pe)
|
||||
x = attn.proj(x)
|
||||
return x
|
||||
|
||||
class LoraFluxAttnProcessor(nn.Module):
|
||||
|
||||
def __init__(self, dim: int, rank=4, network_alpha=None, lora_weight=1):
|
||||
super().__init__()
|
||||
self.qkv_lora = LoRALinearLayer(dim, dim * 3, rank, network_alpha)
|
||||
self.proj_lora = LoRALinearLayer(dim, dim, rank, network_alpha)
|
||||
self.lora_weight = lora_weight
|
||||
|
||||
|
||||
def __call__(self, attn, x, pe, **attention_kwargs):
|
||||
qkv = attn.qkv(x) + self.qkv_lora(x) * self.lora_weight
|
||||
q, k, v = rearrange(qkv, "B L (K H D) -> K B H L D", K=3, H=self.num_heads)
|
||||
q, k = attn.norm(q, k, v)
|
||||
x = attention(q, k, v, pe=pe)
|
||||
x = attn.proj(x) + self.proj_lora(x) * self.lora_weight
|
||||
return x
|
||||
|
||||
class SelfAttention(nn.Module):
|
||||
def __init__(self, dim: int, num_heads: int = 8, qkv_bias: bool = False):
|
||||
super().__init__()
|
||||
self.num_heads = num_heads
|
||||
head_dim = dim // num_heads
|
||||
|
||||
self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
|
||||
self.norm = QKNorm(head_dim)
|
||||
self.proj = nn.Linear(dim, dim)
|
||||
def forward():
|
||||
pass
|
||||
|
||||
|
||||
@dataclass
|
||||
class ModulationOut:
|
||||
shift: Tensor
|
||||
scale: Tensor
|
||||
gate: Tensor
|
||||
|
||||
|
||||
class Modulation(nn.Module):
|
||||
def __init__(self, dim: int, double: bool):
|
||||
super().__init__()
|
||||
self.is_double = double
|
||||
self.multiplier = 6 if double else 3
|
||||
self.lin = nn.Linear(dim, self.multiplier * dim, bias=True)
|
||||
|
||||
def forward(self, vec: Tensor) -> tuple[ModulationOut, ModulationOut | None]:
|
||||
out = self.lin(nn.functional.silu(vec))[:, None, :].chunk(self.multiplier, dim=-1)
|
||||
|
||||
return (
|
||||
ModulationOut(*out[:3]),
|
||||
ModulationOut(*out[3:]) if self.is_double else None,
|
||||
)
|
||||
|
||||
class DoubleStreamBlockLoraProcessor(nn.Module):
|
||||
def __init__(self, dim: int, rank=4, network_alpha=None, lora_weight=1):
|
||||
super().__init__()
|
||||
self.qkv_lora1 = LoRALinearLayer(dim, dim * 3, rank, network_alpha)
|
||||
self.proj_lora1 = LoRALinearLayer(dim, dim, rank, network_alpha)
|
||||
self.qkv_lora2 = LoRALinearLayer(dim, dim * 3, rank, network_alpha)
|
||||
self.proj_lora2 = LoRALinearLayer(dim, dim, rank, network_alpha)
|
||||
self.lora_weight = lora_weight
|
||||
|
||||
def forward(self, attn, img, txt, vec, pe, **attention_kwargs):
|
||||
img_mod1, img_mod2 = attn.img_mod(vec)
|
||||
txt_mod1, txt_mod2 = attn.txt_mod(vec)
|
||||
|
||||
# prepare image for attention
|
||||
img_modulated = attn.img_norm1(img)
|
||||
img_modulated = (1 + img_mod1.scale) * img_modulated + img_mod1.shift
|
||||
img_qkv = attn.img_attn.qkv(img_modulated) + self.qkv_lora1(img_modulated) * self.lora_weight
|
||||
img_q, img_k, img_v = rearrange(img_qkv, "B L (K H D) -> K B H L D", K=3, H=attn.num_heads)
|
||||
img_q, img_k = attn.img_attn.norm(img_q, img_k, img_v)
|
||||
|
||||
# prepare txt for attention
|
||||
txt_modulated = attn.txt_norm1(txt)
|
||||
txt_modulated = (1 + txt_mod1.scale) * txt_modulated + txt_mod1.shift
|
||||
txt_qkv = attn.txt_attn.qkv(txt_modulated) + self.qkv_lora2(txt_modulated) * self.lora_weight
|
||||
txt_q, txt_k, txt_v = rearrange(txt_qkv, "B L (K H D) -> K B H L D", K=3, H=attn.num_heads)
|
||||
txt_q, txt_k = attn.txt_attn.norm(txt_q, txt_k, txt_v)
|
||||
|
||||
# run actual attention
|
||||
q = torch.cat((txt_q, img_q), dim=2)
|
||||
k = torch.cat((txt_k, img_k), dim=2)
|
||||
v = torch.cat((txt_v, img_v), dim=2)
|
||||
|
||||
attn1 = attention(q, k, v, pe=pe)
|
||||
txt_attn, img_attn = attn1[:, : txt.shape[1]], attn1[:, txt.shape[1] :]
|
||||
|
||||
# calculate the img bloks
|
||||
img = img + img_mod1.gate * (attn.img_attn.proj(img_attn) + self.proj_lora1(img_attn) * self.lora_weight)
|
||||
img = img + img_mod2.gate * attn.img_mlp((1 + img_mod2.scale) * attn.img_norm2(img) + img_mod2.shift)
|
||||
|
||||
# calculate the txt bloks
|
||||
txt = txt + txt_mod1.gate * (attn.txt_attn.proj(txt_attn) + self.proj_lora2(txt_attn) * self.lora_weight)
|
||||
txt = txt + txt_mod2.gate * attn.txt_mlp((1 + txt_mod2.scale) * attn.txt_norm2(txt) + txt_mod2.shift)
|
||||
return img, txt
|
||||
|
||||
class DoubleStreamBlockProcessor:
|
||||
def __call__(self, attn, img, txt, vec, pe, **attention_kwargs):
|
||||
img_mod1, img_mod2 = attn.img_mod(vec)
|
||||
txt_mod1, txt_mod2 = attn.txt_mod(vec)
|
||||
|
||||
# prepare image for attention
|
||||
img_modulated = attn.img_norm1(img)
|
||||
img_modulated = (1 + img_mod1.scale) * img_modulated + img_mod1.shift
|
||||
img_qkv = attn.img_attn.qkv(img_modulated)
|
||||
img_q, img_k, img_v = rearrange(img_qkv, "B L (K H D) -> K B H L D", K=3, H=attn.num_heads, D=attn.head_dim)
|
||||
img_q, img_k = attn.img_attn.norm(img_q, img_k, img_v)
|
||||
|
||||
# prepare txt for attention
|
||||
txt_modulated = attn.txt_norm1(txt)
|
||||
txt_modulated = (1 + txt_mod1.scale) * txt_modulated + txt_mod1.shift
|
||||
txt_qkv = attn.txt_attn.qkv(txt_modulated)
|
||||
txt_q, txt_k, txt_v = rearrange(txt_qkv, "B L (K H D) -> K B H L D", K=3, H=attn.num_heads, D=attn.head_dim)
|
||||
txt_q, txt_k = attn.txt_attn.norm(txt_q, txt_k, txt_v)
|
||||
|
||||
# run actual attention
|
||||
q = torch.cat((txt_q, img_q), dim=2)
|
||||
k = torch.cat((txt_k, img_k), dim=2)
|
||||
v = torch.cat((txt_v, img_v), dim=2)
|
||||
|
||||
attn1 = attention(q, k, v, pe=pe)
|
||||
txt_attn, img_attn = attn1[:, : txt.shape[1]], attn1[:, txt.shape[1] :]
|
||||
|
||||
# calculate the img bloks
|
||||
img = img + img_mod1.gate * attn.img_attn.proj(img_attn)
|
||||
img = img + img_mod2.gate * attn.img_mlp((1 + img_mod2.scale) * attn.img_norm2(img) + img_mod2.shift)
|
||||
|
||||
# calculate the txt bloks
|
||||
txt = txt + txt_mod1.gate * attn.txt_attn.proj(txt_attn)
|
||||
txt = txt + txt_mod2.gate * attn.txt_mlp((1 + txt_mod2.scale) * attn.txt_norm2(txt) + txt_mod2.shift)
|
||||
return img, txt
|
||||
|
||||
class DoubleStreamBlock(nn.Module):
|
||||
def __init__(self, hidden_size: int, num_heads: int, mlp_ratio: float, qkv_bias: bool = False):
|
||||
super().__init__()
|
||||
mlp_hidden_dim = int(hidden_size * mlp_ratio)
|
||||
self.num_heads = num_heads
|
||||
self.hidden_size = hidden_size
|
||||
self.head_dim = hidden_size // num_heads
|
||||
|
||||
self.img_mod = Modulation(hidden_size, double=True)
|
||||
self.img_norm1 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
|
||||
self.img_attn = SelfAttention(dim=hidden_size, num_heads=num_heads, qkv_bias=qkv_bias)
|
||||
|
||||
self.img_norm2 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
|
||||
self.img_mlp = nn.Sequential(
|
||||
nn.Linear(hidden_size, mlp_hidden_dim, bias=True),
|
||||
nn.GELU(approximate="tanh"),
|
||||
nn.Linear(mlp_hidden_dim, hidden_size, bias=True),
|
||||
)
|
||||
|
||||
self.txt_mod = Modulation(hidden_size, double=True)
|
||||
self.txt_norm1 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
|
||||
self.txt_attn = SelfAttention(dim=hidden_size, num_heads=num_heads, qkv_bias=qkv_bias)
|
||||
|
||||
self.txt_norm2 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
|
||||
self.txt_mlp = nn.Sequential(
|
||||
nn.Linear(hidden_size, mlp_hidden_dim, bias=True),
|
||||
nn.GELU(approximate="tanh"),
|
||||
nn.Linear(mlp_hidden_dim, hidden_size, bias=True),
|
||||
)
|
||||
processor = DoubleStreamBlockProcessor()
|
||||
self.set_processor(processor)
|
||||
|
||||
def set_processor(self, processor) -> None:
|
||||
self.processor = processor
|
||||
|
||||
def get_processor(self):
|
||||
return self.processor
|
||||
|
||||
def forward(
|
||||
self,
|
||||
img: Tensor,
|
||||
txt: Tensor,
|
||||
vec: Tensor,
|
||||
pe: Tensor,
|
||||
image_proj: Tensor = None,
|
||||
ip_scale: float =1.0,
|
||||
) -> tuple[Tensor, Tensor]:
|
||||
if image_proj is None:
|
||||
return self.processor(self, img, txt, vec, pe)
|
||||
else:
|
||||
return self.processor(self, img, txt, vec, pe, image_proj, ip_scale)
|
||||
|
||||
|
||||
class SingleStreamBlockLoraProcessor(nn.Module):
|
||||
def __init__(self, dim: int, rank: int = 4, network_alpha = None, lora_weight: float = 1):
|
||||
super().__init__()
|
||||
self.qkv_lora = LoRALinearLayer(dim, dim * 3, rank, network_alpha)
|
||||
self.proj_lora = LoRALinearLayer(15360, dim, rank, network_alpha)
|
||||
self.lora_weight = lora_weight
|
||||
|
||||
def forward(self, attn: nn.Module, x: Tensor, vec: Tensor, pe: Tensor) -> Tensor:
|
||||
|
||||
mod, _ = attn.modulation(vec)
|
||||
x_mod = (1 + mod.scale) * attn.pre_norm(x) + mod.shift
|
||||
qkv, mlp = torch.split(attn.linear1(x_mod), [3 * attn.hidden_size, attn.mlp_hidden_dim], dim=-1)
|
||||
qkv = qkv + self.qkv_lora(x_mod) * self.lora_weight
|
||||
|
||||
q, k, v = rearrange(qkv, "B L (K H D) -> K B H L D", K=3, H=attn.num_heads)
|
||||
q, k = attn.norm(q, k, v)
|
||||
|
||||
# compute attention
|
||||
attn_1 = attention(q, k, v, pe=pe)
|
||||
|
||||
# compute activation in mlp stream, cat again and run second linear layer
|
||||
output = attn.linear2(torch.cat((attn_1, attn.mlp_act(mlp)), 2))
|
||||
output = output + self.proj_lora(torch.cat((attn_1, attn.mlp_act(mlp)), 2)) * self.lora_weight
|
||||
output = x + mod.gate * output
|
||||
return output
|
||||
|
||||
|
||||
class SingleStreamBlockProcessor:
|
||||
def __call__(self, attn: nn.Module, x: Tensor, vec: Tensor, pe: Tensor, **attention_kwargs) -> Tensor:
|
||||
|
||||
mod, _ = attn.modulation(vec)
|
||||
x_mod = (1 + mod.scale) * attn.pre_norm(x) + mod.shift
|
||||
qkv, mlp = torch.split(attn.linear1(x_mod), [3 * attn.hidden_size, attn.mlp_hidden_dim], dim=-1)
|
||||
|
||||
q, k, v = rearrange(qkv, "B L (K H D) -> K B H L D", K=3, H=attn.num_heads)
|
||||
q, k = attn.norm(q, k, v)
|
||||
|
||||
# compute attention
|
||||
attn_1 = attention(q, k, v, pe=pe)
|
||||
|
||||
# compute activation in mlp stream, cat again and run second linear layer
|
||||
output = attn.linear2(torch.cat((attn_1, attn.mlp_act(mlp)), 2))
|
||||
output = x + mod.gate * output
|
||||
return output
|
||||
|
||||
class SingleStreamBlock(nn.Module):
|
||||
"""
|
||||
A DiT block with parallel linear layers as described in
|
||||
https://arxiv.org/abs/2302.05442 and adapted modulation interface.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
hidden_size: int,
|
||||
num_heads: int,
|
||||
mlp_ratio: float = 4.0,
|
||||
qk_scale: float | None = None,
|
||||
):
|
||||
super().__init__()
|
||||
self.hidden_dim = hidden_size
|
||||
self.num_heads = num_heads
|
||||
self.head_dim = hidden_size // num_heads
|
||||
self.scale = qk_scale or self.head_dim**-0.5
|
||||
|
||||
self.mlp_hidden_dim = int(hidden_size * mlp_ratio)
|
||||
# qkv and mlp_in
|
||||
self.linear1 = nn.Linear(hidden_size, hidden_size * 3 + self.mlp_hidden_dim)
|
||||
# proj and mlp_out
|
||||
self.linear2 = nn.Linear(hidden_size + self.mlp_hidden_dim, hidden_size)
|
||||
|
||||
self.norm = QKNorm(self.head_dim)
|
||||
|
||||
self.hidden_size = hidden_size
|
||||
self.pre_norm = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
|
||||
|
||||
self.mlp_act = nn.GELU(approximate="tanh")
|
||||
self.modulation = Modulation(hidden_size, double=False)
|
||||
|
||||
processor = SingleStreamBlockProcessor()
|
||||
self.set_processor(processor)
|
||||
|
||||
|
||||
def set_processor(self, processor) -> None:
|
||||
self.processor = processor
|
||||
|
||||
def get_processor(self):
|
||||
return self.processor
|
||||
|
||||
def forward(
|
||||
self,
|
||||
x: Tensor,
|
||||
vec: Tensor,
|
||||
pe: Tensor,
|
||||
image_proj: Tensor | None = None,
|
||||
ip_scale: float = 1.0,
|
||||
) -> Tensor:
|
||||
if image_proj is None:
|
||||
return self.processor(self, x, vec, pe)
|
||||
else:
|
||||
return self.processor(self, x, vec, pe, image_proj, ip_scale)
|
||||
|
||||
|
||||
|
||||
class LastLayer(nn.Module):
|
||||
def __init__(self, hidden_size: int, patch_size: int, out_channels: int):
|
||||
super().__init__()
|
||||
self.norm_final = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
|
||||
self.linear = nn.Linear(hidden_size, patch_size * patch_size * out_channels, bias=True)
|
||||
self.adaLN_modulation = nn.Sequential(nn.SiLU(), nn.Linear(hidden_size, 2 * hidden_size, bias=True))
|
||||
|
||||
def forward(self, x: Tensor, vec: Tensor) -> Tensor:
|
||||
shift, scale = self.adaLN_modulation(vec).chunk(2, dim=1)
|
||||
x = (1 + scale[:, None, :]) * self.norm_final(x) + shift[:, None, :]
|
||||
x = self.linear(x)
|
||||
return x
|
||||
@@ -0,0 +1,293 @@
|
||||
# Copyright (c) 2025 Bytedance Ltd. and/or its affiliates. All rights reserved.
|
||||
# Copyright (c) 2024 Black Forest Labs and The XLabs-AI 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.
|
||||
|
||||
import os
|
||||
from typing import Literal
|
||||
|
||||
import torch
|
||||
from einops import rearrange
|
||||
from PIL import ExifTags, Image
|
||||
import torchvision.transforms.functional as TVF
|
||||
|
||||
from .modules.layers import (
|
||||
DoubleStreamBlockLoraProcessor,
|
||||
DoubleStreamBlockProcessor,
|
||||
SingleStreamBlockLoraProcessor,
|
||||
SingleStreamBlockProcessor,
|
||||
)
|
||||
from .sampling import denoise, get_noise, get_schedule, prepare_multi_ip, unpack
|
||||
from .util import (
|
||||
get_lora_rank,
|
||||
load_ae,
|
||||
load_checkpoint,
|
||||
load_clip,
|
||||
load_flow_model,
|
||||
load_flow_model_only_lora,
|
||||
load_flow_model_quintized,
|
||||
load_t5,
|
||||
)
|
||||
|
||||
|
||||
def find_nearest_scale(image_h, image_w, predefined_scales):
|
||||
"""
|
||||
根据图片的高度和宽度，找到最近的预定义尺度。
|
||||
|
||||
:param image_h: 图片的高度
|
||||
:param image_w: 图片的宽度
|
||||
:param predefined_scales: 预定义尺度列表 [(h1, w1), (h2, w2), ...]
|
||||
:return: 最近的预定义尺度 (h, w)
|
||||
"""
|
||||
# 计算输入图片的长宽比
|
||||
image_ratio = image_h / image_w
|
||||
|
||||
# 初始化变量以存储最小差异和最近的尺度
|
||||
min_diff = float('inf')
|
||||
nearest_scale = None
|
||||
|
||||
# 遍历所有预定义尺度，找到与输入图片长宽比最接近的尺度
|
||||
for scale_h, scale_w in predefined_scales:
|
||||
predefined_ratio = scale_h / scale_w
|
||||
diff = abs(predefined_ratio - image_ratio)
|
||||
|
||||
if diff < min_diff:
|
||||
min_diff = diff
|
||||
nearest_scale = (scale_h, scale_w)
|
||||
|
||||
return nearest_scale
|
||||
|
||||
def preprocess_ref(raw_image: Image.Image, long_size: int = 512):
|
||||
# 获取原始图像的宽度和高度
|
||||
image_w, image_h = raw_image.size
|
||||
|
||||
# 计算长边和短边
|
||||
if image_w >= image_h:
|
||||
new_w = long_size
|
||||
new_h = int((long_size / image_w) * image_h)
|
||||
else:
|
||||
new_h = long_size
|
||||
new_w = int((long_size / image_h) * image_w)
|
||||
|
||||
# 按新的宽高进行等比例缩放
|
||||
raw_image = raw_image.resize((new_w, new_h), resample=Image.LANCZOS)
|
||||
target_w = new_w // 16 * 16
|
||||
target_h = new_h // 16 * 16
|
||||
|
||||
# 计算裁剪的起始坐标以实现中心裁剪
|
||||
left = (new_w - target_w) // 2
|
||||
top = (new_h - target_h) // 2
|
||||
right = left + target_w
|
||||
bottom = top + target_h
|
||||
|
||||
# 进行中心裁剪
|
||||
raw_image = raw_image.crop((left, top, right, bottom))
|
||||
|
||||
# 转换为 RGB 模式
|
||||
raw_image = raw_image.convert("RGB")
|
||||
return raw_image
|
||||
|
||||
class UNOPipeline:
|
||||
def __init__(
|
||||
self,
|
||||
model_,
|
||||
device: torch.device,
|
||||
offload: bool = False,
|
||||
use_fp8: bool = False,
|
||||
):
|
||||
self.device = device
|
||||
self.offload = offload
|
||||
self.model= model_
|
||||
self.use_fp8= use_fp8
|
||||
#self.clip = load_clip(self.device)
|
||||
# self.t5 = load_t5(self.device, max_length=512)
|
||||
# self.ae = load_ae(model_type, device="cpu" if offload else self.device)
|
||||
|
||||
|
||||
def load_ckpt(self, ckpt_path):
|
||||
if ckpt_path is not None:
|
||||
from safetensors.torch import load_file as load_sft
|
||||
print("Loading checkpoint to replace old keys")
|
||||
# load_sft doesn't support torch.device
|
||||
if ckpt_path.endswith('safetensors'):
|
||||
sd = load_sft(ckpt_path, device='cpu')
|
||||
missing, unexpected = self.model.load_state_dict(sd, strict=False, assign=True)
|
||||
else:
|
||||
dit_state = torch.load(ckpt_path, map_location='cpu')
|
||||
sd = {}
|
||||
for k in dit_state.keys():
|
||||
sd[k.replace('module.','')] = dit_state[k]
|
||||
missing, unexpected = self.model.load_state_dict(sd, strict=False, assign=True)
|
||||
self.model.to(str(self.device))
|
||||
print(f"missing keys: {missing}\n\n\n\n\nunexpected keys: {unexpected}")
|
||||
|
||||
def set_lora(self, local_path: str = None, repo_id: str = None,
|
||||
name: str = None, lora_weight: int = 0.7):
|
||||
checkpoint = load_checkpoint(local_path, repo_id, name)
|
||||
self.update_model_with_lora(checkpoint, lora_weight)
|
||||
|
||||
def set_lora_from_collection(self, lora_type: str = "realism", lora_weight: int = 0.7):
|
||||
checkpoint = load_checkpoint(
|
||||
None, self.hf_lora_collection, self.lora_types_to_names[lora_type]
|
||||
)
|
||||
self.update_model_with_lora(checkpoint, lora_weight)
|
||||
|
||||
def update_model_with_lora(self, checkpoint, lora_weight):
|
||||
rank = get_lora_rank(checkpoint)
|
||||
lora_attn_procs = {}
|
||||
|
||||
for name, _ in self.model.attn_processors.items():
|
||||
lora_state_dict = {}
|
||||
for k in checkpoint.keys():
|
||||
if name in k:
|
||||
lora_state_dict[k[len(name) + 1:]] = checkpoint[k] * lora_weight
|
||||
|
||||
if len(lora_state_dict):
|
||||
if name.startswith("single_blocks"):
|
||||
lora_attn_procs[name] = SingleStreamBlockLoraProcessor(dim=3072, rank=rank)
|
||||
else:
|
||||
lora_attn_procs[name] = DoubleStreamBlockLoraProcessor(dim=3072, rank=rank)
|
||||
lora_attn_procs[name].load_state_dict(lora_state_dict)
|
||||
lora_attn_procs[name].to(self.device)
|
||||
else:
|
||||
if name.startswith("single_blocks"):
|
||||
lora_attn_procs[name] = SingleStreamBlockProcessor()
|
||||
else:
|
||||
lora_attn_procs[name] = DoubleStreamBlockProcessor()
|
||||
|
||||
self.model.set_attn_processor(lora_attn_procs)
|
||||
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
width: int = 512,
|
||||
height: int = 512,
|
||||
guidance: float = 4,
|
||||
num_steps: int = 50,
|
||||
inp_cond={},
|
||||
**kwargs
|
||||
):
|
||||
width = 16 * (width // 16)
|
||||
height = 16 * (height // 16)
|
||||
device_type = self.device if isinstance(self.device, str) else self.device.type
|
||||
with torch.autocast(enabled=self.use_fp8, device_type=device_type, dtype=torch.bfloat16):
|
||||
return self.forward(
|
||||
width,
|
||||
height,
|
||||
guidance,
|
||||
num_steps,
|
||||
inp_cond=inp_cond,
|
||||
**kwargs
|
||||
)
|
||||
|
||||
@torch.inference_mode()
|
||||
def gradio_generate(
|
||||
self,
|
||||
prompt: str,
|
||||
width: int,
|
||||
height: int,
|
||||
guidance: float,
|
||||
num_steps: int,
|
||||
seed: int,
|
||||
image_prompt1: Image.Image,
|
||||
image_prompt2: Image.Image,
|
||||
image_prompt3: Image.Image,
|
||||
image_prompt4: Image.Image,
|
||||
):
|
||||
ref_imgs = [image_prompt1, image_prompt2, image_prompt3, image_prompt4]
|
||||
ref_imgs = [img for img in ref_imgs if isinstance(img, Image.Image)]
|
||||
ref_long_side = 512 if len(ref_imgs) <= 1 else 320
|
||||
ref_imgs = [preprocess_ref(img, ref_long_side) for img in ref_imgs]
|
||||
|
||||
seed = seed if seed != -1 else torch.randint(0, 10 ** 8, (1,)).item()
|
||||
|
||||
img = self(prompt=prompt, width=width, height=height, guidance=guidance,
|
||||
num_steps=num_steps, seed=seed, ref_imgs=ref_imgs)
|
||||
|
||||
filename = f"output/gradio/{seed}_{prompt[:20]}.png"
|
||||
os.makedirs(os.path.dirname(filename), exist_ok=True)
|
||||
exif_data = Image.Exif()
|
||||
exif_data[ExifTags.Base.Make] = "UNO"
|
||||
exif_data[ExifTags.Base.Model] = self.model_type
|
||||
info = f"{prompt=}, {seed=}, {width=}, {height=}, {guidance=}, {num_steps=}"
|
||||
exif_data[ExifTags.Base.ImageDescription] = info
|
||||
img.save(filename, format="png", exif=exif_data)
|
||||
return img, filename
|
||||
|
||||
@torch.inference_mode
|
||||
def forward(
|
||||
self,
|
||||
width: int,
|
||||
height: int,
|
||||
guidance: float,
|
||||
num_steps: int,
|
||||
inp_cond={},
|
||||
):
|
||||
# x = get_noise(
|
||||
# 1, height, width, device=self.device,
|
||||
# dtype=torch.bfloat16, seed=seed
|
||||
# )
|
||||
timesteps = get_schedule(
|
||||
num_steps,
|
||||
(width // 8) * (height // 8) // (16 * 16),
|
||||
shift=True,
|
||||
)
|
||||
# if self.offload:
|
||||
# self.ae.encoder = self.ae.encoder.to(self.device)
|
||||
# x_1_refs = [
|
||||
# self.ae.encode(
|
||||
# (TVF.to_tensor(ref_img) * 2.0 - 1.0)
|
||||
# .unsqueeze(0).to(self.device, torch.float32)
|
||||
# ).to(torch.bfloat16)
|
||||
# for ref_img in ref_imgs
|
||||
# ]
|
||||
|
||||
# if self.offload:
|
||||
# self.ae.encoder = self.offload_model_to_cpu(self.ae.encoder)
|
||||
# self.t5, self.clip = self.t5.to(self.device), self.clip.to(self.device)
|
||||
# inp_cond = prepare_multi_ip(
|
||||
# t5=self.t5, clip=self.clip,
|
||||
# img=x,
|
||||
# prompt=prompt, ref_imgs=x_1_refs, pe=pe
|
||||
# )
|
||||
|
||||
# if self.offload:
|
||||
# self.offload_model_to_cpu(self.t5, self.clip)
|
||||
#print(self.device)
|
||||
self.model.to(self.device)
|
||||
|
||||
x = denoise(
|
||||
self.model,
|
||||
**inp_cond,
|
||||
timesteps=timesteps,
|
||||
guidance=guidance,
|
||||
)
|
||||
|
||||
if self.offload:
|
||||
self.offload_model_to_cpu(self.model)
|
||||
#self.ae.decoder.to(x.device)
|
||||
x = unpack(x.float(), height, width)
|
||||
return (x/0.3611)+0.1159
|
||||
# x = self.ae.decode(x)
|
||||
# self.offload_model_to_cpu(self.ae.decoder)
|
||||
|
||||
# x1 = x.clamp(-1, 1)
|
||||
# x1 = rearrange(x1[-1], "c h w -> h w c")
|
||||
# output_img = Image.fromarray((127.5 * (x1 + 1.0)).cpu().byte().numpy())
|
||||
# return output_img
|
||||
|
||||
def offload_model_to_cpu(self, *models):
|
||||
if not self.offload: return
|
||||
for model in models:
|
||||
model.cpu()
|
||||
torch.cuda.empty_cache()
|
||||
@@ -0,0 +1,304 @@
|
||||
# Copyright (c) 2025 Bytedance Ltd. and/or its affiliates. All rights reserved.
|
||||
# Copyright (c) 2024 Black Forest Labs and The XLabs-AI 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.
|
||||
|
||||
import os
|
||||
from typing import Literal
|
||||
|
||||
import torch
|
||||
from einops import rearrange
|
||||
from PIL import ExifTags, Image
|
||||
import torchvision.transforms.functional as TVF
|
||||
|
||||
from .modules.layers import (
|
||||
DoubleStreamBlockLoraProcessor,
|
||||
DoubleStreamBlockProcessor,
|
||||
SingleStreamBlockLoraProcessor,
|
||||
SingleStreamBlockProcessor,
|
||||
)
|
||||
from .sampling import denoise, get_noise, get_schedule, prepare_multi_ip, unpack
|
||||
from .util import (
|
||||
get_lora_rank,
|
||||
load_ae,
|
||||
load_checkpoint,
|
||||
load_clip,
|
||||
load_flow_model,
|
||||
load_flow_model_only_lora,
|
||||
load_flow_model_quintized,
|
||||
load_t5,
|
||||
)
|
||||
|
||||
|
||||
def find_nearest_scale(image_h, image_w, predefined_scales):
|
||||
"""
|
||||
根据图片的高度和宽度，找到最近的预定义尺度。
|
||||
|
||||
:param image_h: 图片的高度
|
||||
:param image_w: 图片的宽度
|
||||
:param predefined_scales: 预定义尺度列表 [(h1, w1), (h2, w2), ...]
|
||||
:return: 最近的预定义尺度 (h, w)
|
||||
"""
|
||||
# 计算输入图片的长宽比
|
||||
image_ratio = image_h / image_w
|
||||
|
||||
# 初始化变量以存储最小差异和最近的尺度
|
||||
min_diff = float('inf')
|
||||
nearest_scale = None
|
||||
|
||||
# 遍历所有预定义尺度，找到与输入图片长宽比最接近的尺度
|
||||
for scale_h, scale_w in predefined_scales:
|
||||
predefined_ratio = scale_h / scale_w
|
||||
diff = abs(predefined_ratio - image_ratio)
|
||||
|
||||
if diff < min_diff:
|
||||
min_diff = diff
|
||||
nearest_scale = (scale_h, scale_w)
|
||||
|
||||
return nearest_scale
|
||||
|
||||
def preprocess_ref(raw_image: Image.Image, long_size: int = 512):
|
||||
# 获取原始图像的宽度和高度
|
||||
image_w, image_h = raw_image.size
|
||||
|
||||
# 计算长边和短边
|
||||
if image_w >= image_h:
|
||||
new_w = long_size
|
||||
new_h = int((long_size / image_w) * image_h)
|
||||
else:
|
||||
new_h = long_size
|
||||
new_w = int((long_size / image_h) * image_w)
|
||||
|
||||
# 按新的宽高进行等比例缩放
|
||||
raw_image = raw_image.resize((new_w, new_h), resample=Image.LANCZOS)
|
||||
target_w = new_w // 16 * 16
|
||||
target_h = new_h // 16 * 16
|
||||
|
||||
# 计算裁剪的起始坐标以实现中心裁剪
|
||||
left = (new_w - target_w) // 2
|
||||
top = (new_h - target_h) // 2
|
||||
right = left + target_w
|
||||
bottom = top + target_h
|
||||
|
||||
# 进行中心裁剪
|
||||
raw_image = raw_image.crop((left, top, right, bottom))
|
||||
|
||||
# 转换为 RGB 模式
|
||||
raw_image = raw_image.convert("RGB")
|
||||
return raw_image
|
||||
|
||||
class UNOPipeline:
|
||||
def __init__(
|
||||
self,
|
||||
model_type: str,
|
||||
device: torch.device,
|
||||
offload: bool = False,
|
||||
only_lora: bool = False,
|
||||
lora_rank: int = 16
|
||||
):
|
||||
self.device = device
|
||||
self.offload = offload
|
||||
self.model_type = model_type
|
||||
|
||||
self.clip = load_clip(self.device)
|
||||
self.t5 = load_t5(self.device, max_length=512)
|
||||
self.ae = load_ae(model_type, device="cpu" if offload else self.device)
|
||||
if "fp8" in model_type:
|
||||
self.model = load_flow_model_quintized(model_type, device="cpu" if offload else self.device)
|
||||
elif only_lora:
|
||||
self.model = load_flow_model_only_lora(
|
||||
model_type, device="cpu" if offload else self.device, lora_rank=lora_rank
|
||||
)
|
||||
else:
|
||||
self.model = load_flow_model(model_type, device="cpu" if offload else self.device)
|
||||
|
||||
|
||||
def load_ckpt(self, ckpt_path):
|
||||
if ckpt_path is not None:
|
||||
from safetensors.torch import load_file as load_sft
|
||||
print("Loading checkpoint to replace old keys")
|
||||
# load_sft doesn't support torch.device
|
||||
if ckpt_path.endswith('safetensors'):
|
||||
sd = load_sft(ckpt_path, device='cpu')
|
||||
missing, unexpected = self.model.load_state_dict(sd, strict=False, assign=True)
|
||||
else:
|
||||
dit_state = torch.load(ckpt_path, map_location='cpu')
|
||||
sd = {}
|
||||
for k in dit_state.keys():
|
||||
sd[k.replace('module.','')] = dit_state[k]
|
||||
missing, unexpected = self.model.load_state_dict(sd, strict=False, assign=True)
|
||||
self.model.to(str(self.device))
|
||||
print(f"missing keys: {missing}\n\n\n\n\nunexpected keys: {unexpected}")
|
||||
|
||||
def set_lora(self, local_path: str = None, repo_id: str = None,
|
||||
name: str = None, lora_weight: int = 0.7):
|
||||
checkpoint = load_checkpoint(local_path, repo_id, name)
|
||||
self.update_model_with_lora(checkpoint, lora_weight)
|
||||
|
||||
def set_lora_from_collection(self, lora_type: str = "realism", lora_weight: int = 0.7):
|
||||
checkpoint = load_checkpoint(
|
||||
None, self.hf_lora_collection, self.lora_types_to_names[lora_type]
|
||||
)
|
||||
self.update_model_with_lora(checkpoint, lora_weight)
|
||||
|
||||
def update_model_with_lora(self, checkpoint, lora_weight):
|
||||
rank = get_lora_rank(checkpoint)
|
||||
lora_attn_procs = {}
|
||||
|
||||
for name, _ in self.model.attn_processors.items():
|
||||
lora_state_dict = {}
|
||||
for k in checkpoint.keys():
|
||||
if name in k:
|
||||
lora_state_dict[k[len(name) + 1:]] = checkpoint[k] * lora_weight
|
||||
|
||||
if len(lora_state_dict):
|
||||
if name.startswith("single_blocks"):
|
||||
lora_attn_procs[name] = SingleStreamBlockLoraProcessor(dim=3072, rank=rank)
|
||||
else:
|
||||
lora_attn_procs[name] = DoubleStreamBlockLoraProcessor(dim=3072, rank=rank)
|
||||
lora_attn_procs[name].load_state_dict(lora_state_dict)
|
||||
lora_attn_procs[name].to(self.device)
|
||||
else:
|
||||
if name.startswith("single_blocks"):
|
||||
lora_attn_procs[name] = SingleStreamBlockProcessor()
|
||||
else:
|
||||
lora_attn_procs[name] = DoubleStreamBlockProcessor()
|
||||
|
||||
self.model.set_attn_processor(lora_attn_procs)
|
||||
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
prompt: str,
|
||||
width: int = 512,
|
||||
height: int = 512,
|
||||
guidance: float = 4,
|
||||
num_steps: int = 50,
|
||||
seed: int = 123456789,
|
||||
**kwargs
|
||||
):
|
||||
width = 16 * (width // 16)
|
||||
height = 16 * (height // 16)
|
||||
|
||||
return self.forward(
|
||||
prompt,
|
||||
width,
|
||||
height,
|
||||
guidance,
|
||||
num_steps,
|
||||
seed,
|
||||
**kwargs
|
||||
)
|
||||
|
||||
@torch.inference_mode()
|
||||
def gradio_generate(
|
||||
self,
|
||||
prompt: str,
|
||||
width: int,
|
||||
height: int,
|
||||
guidance: float,
|
||||
num_steps: int,
|
||||
seed: int,
|
||||
image_prompt1: Image.Image,
|
||||
image_prompt2: Image.Image,
|
||||
image_prompt3: Image.Image,
|
||||
image_prompt4: Image.Image,
|
||||
):
|
||||
ref_imgs = [image_prompt1, image_prompt2, image_prompt3, image_prompt4]
|
||||
ref_imgs = [img for img in ref_imgs if isinstance(img, Image.Image)]
|
||||
ref_long_side = 512 if len(ref_imgs) <= 1 else 320
|
||||
ref_imgs = [preprocess_ref(img, ref_long_side) for img in ref_imgs]
|
||||
|
||||
seed = seed if seed != -1 else torch.randint(0, 10 ** 8, (1,)).item()
|
||||
|
||||
img = self(prompt=prompt, width=width, height=height, guidance=guidance,
|
||||
num_steps=num_steps, seed=seed, ref_imgs=ref_imgs)
|
||||
|
||||
filename = f"output/gradio/{seed}_{prompt[:20]}.png"
|
||||
os.makedirs(os.path.dirname(filename), exist_ok=True)
|
||||
exif_data = Image.Exif()
|
||||
exif_data[ExifTags.Base.Make] = "UNO"
|
||||
exif_data[ExifTags.Base.Model] = self.model_type
|
||||
info = f"{prompt=}, {seed=}, {width=}, {height=}, {guidance=}, {num_steps=}"
|
||||
exif_data[ExifTags.Base.ImageDescription] = info
|
||||
img.save(filename, format="png", exif=exif_data)
|
||||
return img, filename
|
||||
|
||||
@torch.inference_mode
|
||||
def forward(
|
||||
self,
|
||||
prompt: str,
|
||||
width: int,
|
||||
height: int,
|
||||
guidance: float,
|
||||
num_steps: int,
|
||||
seed: int,
|
||||
ref_imgs: list[Image.Image] | None = None,
|
||||
pe: Literal['d', 'h', 'w', 'o'] = 'd',
|
||||
):
|
||||
x = get_noise(
|
||||
1, height, width, device=self.device,
|
||||
dtype=torch.bfloat16, seed=seed
|
||||
)
|
||||
timesteps = get_schedule(
|
||||
num_steps,
|
||||
(width // 8) * (height // 8) // (16 * 16),
|
||||
shift=True,
|
||||
)
|
||||
if self.offload:
|
||||
self.ae.encoder = self.ae.encoder.to(self.device)
|
||||
x_1_refs = [
|
||||
self.ae.encode(
|
||||
(TVF.to_tensor(ref_img) * 2.0 - 1.0)
|
||||
.unsqueeze(0).to(self.device, torch.float32)
|
||||
).to(torch.bfloat16)
|
||||
for ref_img in ref_imgs
|
||||
]
|
||||
|
||||
if self.offload:
|
||||
self.ae.encoder = self.offload_model_to_cpu(self.ae.encoder)
|
||||
self.t5, self.clip = self.t5.to(self.device), self.clip.to(self.device)
|
||||
inp_cond = prepare_multi_ip(
|
||||
t5=self.t5, clip=self.clip,
|
||||
img=x,
|
||||
prompt=prompt, ref_imgs=x_1_refs, pe=pe
|
||||
)
|
||||
|
||||
if self.offload:
|
||||
self.offload_model_to_cpu(self.t5, self.clip)
|
||||
self.model = self.model.to(self.device)
|
||||
|
||||
x = denoise(
|
||||
self.model,
|
||||
**inp_cond,
|
||||
timesteps=timesteps,
|
||||
guidance=guidance,
|
||||
)
|
||||
|
||||
if self.offload:
|
||||
self.offload_model_to_cpu(self.model)
|
||||
self.ae.decoder.to(x.device)
|
||||
x = unpack(x.float(), height, width)
|
||||
x = self.ae.decode(x)
|
||||
self.offload_model_to_cpu(self.ae.decoder)
|
||||
|
||||
x1 = x.clamp(-1, 1)
|
||||
x1 = rearrange(x1[-1], "c h w -> h w c")
|
||||
output_img = Image.fromarray((127.5 * (x1 + 1.0)).cpu().byte().numpy())
|
||||
return output_img
|
||||
|
||||
def offload_model_to_cpu(self, *models):
|
||||
if not self.offload: return
|
||||
for model in models:
|
||||
model.cpu()
|
||||
torch.cuda.empty_cache()
|
||||
@@ -0,0 +1,388 @@
|
||||
# Copyright (c) 2025 Bytedance Ltd. and/or its affiliates. All rights reserved.
|
||||
# Copyright (c) 2024 Black Forest Labs and The XLabs-AI 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.
|
||||
|
||||
import math
|
||||
from typing import Literal
|
||||
|
||||
import torch
|
||||
from einops import rearrange, repeat
|
||||
from torch import Tensor
|
||||
from tqdm import tqdm
|
||||
|
||||
from .model import Flux
|
||||
from .modules.conditioner import HFEmbedder
|
||||
|
||||
|
||||
def get_noise(
|
||||
num_samples: int,
|
||||
height: int,
|
||||
width: int,
|
||||
device: torch.device,
|
||||
dtype: torch.dtype,
|
||||
seed: int,
|
||||
):
|
||||
return torch.randn(
|
||||
num_samples,
|
||||
16,
|
||||
# allow for packing
|
||||
2 * math.ceil(height / 16),
|
||||
2 * math.ceil(width / 16),
|
||||
device=device,
|
||||
dtype=dtype,
|
||||
generator=torch.Generator(device=device).manual_seed(seed),
|
||||
)
|
||||
|
||||
|
||||
def prepare(
|
||||
t5: HFEmbedder,
|
||||
clip: HFEmbedder,
|
||||
img: Tensor,
|
||||
prompt: str | list[str],
|
||||
ref_img: None | Tensor=None,
|
||||
pe: Literal['d', 'h', 'w', 'o'] ='d'
|
||||
) -> dict[str, Tensor]:
|
||||
assert pe in ['d', 'h', 'w', 'o']
|
||||
bs, c, h, w = img.shape
|
||||
if bs == 1 and not isinstance(prompt, str):
|
||||
bs = len(prompt)
|
||||
|
||||
img = rearrange(img, "b c (h ph) (w pw) -> b (h w) (c ph pw)", ph=2, pw=2)
|
||||
if img.shape[0] == 1 and bs > 1:
|
||||
img = repeat(img, "1 ... -> bs ...", bs=bs)
|
||||
|
||||
img_ids = torch.zeros(h // 2, w // 2, 3)
|
||||
img_ids[..., 1] = img_ids[..., 1] + torch.arange(h // 2)[:, None]
|
||||
img_ids[..., 2] = img_ids[..., 2] + torch.arange(w // 2)[None, :]
|
||||
img_ids = repeat(img_ids, "h w c -> b (h w) c", b=bs)
|
||||
|
||||
if ref_img is not None:
|
||||
_, _, ref_h, ref_w = ref_img.shape
|
||||
ref_img = rearrange(ref_img, "b c (h ph) (w pw) -> b (h w) (c ph pw)", ph=2, pw=2)
|
||||
if ref_img.shape[0] == 1 and bs > 1:
|
||||
ref_img = repeat(ref_img, "1 ... -> bs ...", bs=bs)
|
||||
ref_img_ids = torch.zeros(ref_h // 2, ref_w // 2, 3)
|
||||
# img id分别在宽高偏移各自最大值
|
||||
h_offset = h // 2 if pe in {'d', 'h'} else 0
|
||||
w_offset = w // 2 if pe in {'d', 'w'} else 0
|
||||
ref_img_ids[..., 1] = ref_img_ids[..., 1] + torch.arange(ref_h // 2)[:, None] + h_offset
|
||||
ref_img_ids[..., 2] = ref_img_ids[..., 2] + torch.arange(ref_w // 2)[None, :] + w_offset
|
||||
ref_img_ids = repeat(ref_img_ids, "h w c -> b (h w) c", b=bs)
|
||||
|
||||
if isinstance(prompt, str):
|
||||
prompt = [prompt]
|
||||
txt = t5(prompt)
|
||||
if txt.shape[0] == 1 and bs > 1:
|
||||
txt = repeat(txt, "1 ... -> bs ...", bs=bs)
|
||||
txt_ids = torch.zeros(bs, txt.shape[1], 3)
|
||||
|
||||
vec = clip(prompt)
|
||||
if vec.shape[0] == 1 and bs > 1:
|
||||
vec = repeat(vec, "1 ... -> bs ...", bs=bs)
|
||||
|
||||
if ref_img is not None:
|
||||
return {
|
||||
"img": img,
|
||||
"img_ids": img_ids.to(img.device),
|
||||
"ref_img": ref_img,
|
||||
"ref_img_ids": ref_img_ids.to(img.device),
|
||||
"txt": txt.to(img.device),
|
||||
"txt_ids": txt_ids.to(img.device),
|
||||
"vec": vec.to(img.device),
|
||||
}
|
||||
else:
|
||||
return {
|
||||
"img": img,
|
||||
"img_ids": img_ids.to(img.device),
|
||||
"txt": txt.to(img.device),
|
||||
"txt_ids": txt_ids.to(img.device),
|
||||
"vec": vec.to(img.device),
|
||||
}
|
||||
|
||||
def prepare_wrapper(
|
||||
# t5: HFEmbedder,
|
||||
# clip: HFEmbedder,
|
||||
clip,
|
||||
img: Tensor,
|
||||
prompt: str | list[str],
|
||||
ref_img: None | Tensor=None,
|
||||
pe: Literal['d', 'h', 'w', 'o'] ='d'
|
||||
) -> dict[str, Tensor]:
|
||||
assert pe in ['d', 'h', 'w', 'o']
|
||||
bs, c, h, w = img.shape
|
||||
if bs == 1 and not isinstance(prompt, str):
|
||||
bs = len(prompt)
|
||||
|
||||
img = rearrange(img, "b c (h ph) (w pw) -> b (h w) (c ph pw)", ph=2, pw=2)
|
||||
if img.shape[0] == 1 and bs > 1:
|
||||
img = repeat(img, "1 ... -> bs ...", bs=bs)
|
||||
|
||||
img_ids = torch.zeros(h // 2, w // 2, 3)
|
||||
img_ids[..., 1] = img_ids[..., 1] + torch.arange(h // 2)[:, None]
|
||||
img_ids[..., 2] = img_ids[..., 2] + torch.arange(w // 2)[None, :]
|
||||
img_ids = repeat(img_ids, "h w c -> b (h w) c", b=bs)
|
||||
|
||||
if ref_img is not None:
|
||||
_, _, ref_h, ref_w = ref_img.shape
|
||||
ref_img = rearrange(ref_img, "b c (h ph) (w pw) -> b (h w) (c ph pw)", ph=2, pw=2)
|
||||
if ref_img.shape[0] == 1 and bs > 1:
|
||||
ref_img = repeat(ref_img, "1 ... -> bs ...", bs=bs)
|
||||
ref_img_ids = torch.zeros(ref_h // 2, ref_w // 2, 3)
|
||||
# img id分别在宽高偏移各自最大值
|
||||
h_offset = h // 2 if pe in {'d', 'h'} else 0
|
||||
w_offset = w // 2 if pe in {'d', 'w'} else 0
|
||||
ref_img_ids[..., 1] = ref_img_ids[..., 1] + torch.arange(ref_h // 2)[:, None] + h_offset
|
||||
ref_img_ids[..., 2] = ref_img_ids[..., 2] + torch.arange(ref_w // 2)[None, :] + w_offset
|
||||
ref_img_ids = repeat(ref_img_ids, "h w c -> b (h w) c", b=bs)
|
||||
|
||||
# if isinstance(prompt, str):
|
||||
# prompt = [prompt]
|
||||
tokens = clip.tokenize(prompt)
|
||||
outputs = clip.encode_from_tokens(tokens, return_dict=True)
|
||||
txt = outputs.get("cond")
|
||||
if txt.shape[0] == 1 and bs > 1:
|
||||
txt = repeat(txt, "1 ... -> bs ...", bs=bs)
|
||||
txt_ids = torch.zeros(bs, txt.shape[1], 3)
|
||||
|
||||
vec = outputs.get("pooled_output")
|
||||
if vec.shape[0] == 1 and bs > 1:
|
||||
vec = repeat(vec, "1 ... -> bs ...", bs=bs)
|
||||
|
||||
if ref_img is not None:
|
||||
return {
|
||||
"img": img,
|
||||
"img_ids": img_ids.to(img.device,torch.bfloat16),
|
||||
"ref_img": ref_img.to(img.device,torch.bfloat16),
|
||||
"ref_img_ids": ref_img_ids.to(img.device,torch.bfloat16),
|
||||
"txt": txt.to(img.device,torch.bfloat16),
|
||||
"txt_ids": txt_ids.to(img.device,torch.bfloat16),
|
||||
"vec": vec.to(img.device,torch.bfloat16),
|
||||
}
|
||||
else:
|
||||
return {
|
||||
"img": img,
|
||||
"img_ids": img_ids.to(img.device,torch.bfloat16),
|
||||
"txt": txt.to(img.device,torch.bfloat16),
|
||||
"txt_ids": txt_ids.to(img.device,torch.bfloat16),
|
||||
"vec": vec.to(img.device,torch.bfloat16),
|
||||
}
|
||||
|
||||
def prepare_multi_ip_wrapper(
|
||||
clip,
|
||||
img: Tensor,
|
||||
prompt: str | list[str],
|
||||
ref_imgs: list[Tensor] | None = None,
|
||||
pe: Literal['d', 'h', 'w', 'o'] = 'd'
|
||||
) -> dict[str, Tensor]:
|
||||
assert pe in ['d', 'h', 'w', 'o']
|
||||
bs, c, h, w = img.shape
|
||||
if bs == 1 and not isinstance(prompt, str):
|
||||
bs = len(prompt)
|
||||
|
||||
img = rearrange(img, "b c (h ph) (w pw) -> b (h w) (c ph pw)", ph=2, pw=2)
|
||||
if img.shape[0] == 1 and bs > 1:
|
||||
img = repeat(img, "1 ... -> bs ...", bs=bs)
|
||||
|
||||
img_ids = torch.zeros(h // 2, w // 2, 3)
|
||||
img_ids[..., 1] = img_ids[..., 1] + torch.arange(h // 2)[:, None]
|
||||
img_ids[..., 2] = img_ids[..., 2] + torch.arange(w // 2)[None, :]
|
||||
img_ids = repeat(img_ids, "h w c -> b (h w) c", b=bs)
|
||||
|
||||
ref_img_ids = []
|
||||
ref_imgs_list = []
|
||||
pe_shift_w, pe_shift_h = w // 2, h // 2
|
||||
for ref_img in ref_imgs:
|
||||
_, _, ref_h1, ref_w1 = ref_img.shape
|
||||
ref_img = rearrange(ref_img, "b c (h ph) (w pw) -> b (h w) (c ph pw)", ph=2, pw=2)
|
||||
if ref_img.shape[0] == 1 and bs > 1:
|
||||
ref_img = repeat(ref_img, "1 ... -> bs ...", bs=bs)
|
||||
ref_img_ids1 = torch.zeros(ref_h1 // 2, ref_w1 // 2, 3)
|
||||
# img id分别在宽高偏移各自最大值
|
||||
h_offset = pe_shift_h if pe in {'d', 'h'} else 0
|
||||
w_offset = pe_shift_w if pe in {'d', 'w'} else 0
|
||||
ref_img_ids1[..., 1] = ref_img_ids1[..., 1] + torch.arange(ref_h1 // 2)[:, None] + h_offset
|
||||
ref_img_ids1[..., 2] = ref_img_ids1[..., 2] + torch.arange(ref_w1 // 2)[None, :] + w_offset
|
||||
ref_img_ids1 = repeat(ref_img_ids1, "h w c -> b (h w) c", b=bs)
|
||||
ref_img_ids.append(ref_img_ids1)
|
||||
ref_imgs_list.append(ref_img)
|
||||
|
||||
# 更新pe shift
|
||||
pe_shift_h += ref_h1 // 2
|
||||
pe_shift_w += ref_w1 // 2
|
||||
|
||||
# if isinstance(prompt, str):
|
||||
# prompt = [prompt]
|
||||
|
||||
tokens = clip.tokenize(prompt)
|
||||
outputs = clip.encode_from_tokens(tokens, return_dict=True)
|
||||
txt = outputs.get("cond")
|
||||
if txt.shape[0] == 1 and bs > 1:
|
||||
txt = repeat(txt, "1 ... -> bs ...", bs=bs)
|
||||
txt_ids = torch.zeros(bs, txt.shape[1], 3)
|
||||
|
||||
vec = outputs.get("pooled_output")
|
||||
if vec.shape[0] == 1 and bs > 1:
|
||||
vec = repeat(vec, "1 ... -> bs ...", bs=bs)
|
||||
|
||||
return {
|
||||
"img": img,
|
||||
"img_ids": img_ids.to(img.device,torch.bfloat16),
|
||||
"ref_img": tuple(ref_imgs_list),
|
||||
"ref_img_ids": [ref_img_id.to(img.device,torch.bfloat16) for ref_img_id in ref_img_ids],
|
||||
"txt": txt.to(img.device,torch.bfloat16),
|
||||
"txt_ids": txt_ids.to(img.device,torch.bfloat16),
|
||||
"vec": vec.to(img.device,torch.bfloat16),
|
||||
}
|
||||
|
||||
|
||||
def prepare_multi_ip(
|
||||
t5: HFEmbedder,
|
||||
clip: HFEmbedder,
|
||||
img: Tensor,
|
||||
prompt: str | list[str],
|
||||
ref_imgs: list[Tensor] | None = None,
|
||||
pe: Literal['d', 'h', 'w', 'o'] = 'd'
|
||||
) -> dict[str, Tensor]:
|
||||
assert pe in ['d', 'h', 'w', 'o']
|
||||
bs, c, h, w = img.shape
|
||||
if bs == 1 and not isinstance(prompt, str):
|
||||
bs = len(prompt)
|
||||
|
||||
img = rearrange(img, "b c (h ph) (w pw) -> b (h w) (c ph pw)", ph=2, pw=2)
|
||||
if img.shape[0] == 1 and bs > 1:
|
||||
img = repeat(img, "1 ... -> bs ...", bs=bs)
|
||||
|
||||
img_ids = torch.zeros(h // 2, w // 2, 3)
|
||||
img_ids[..., 1] = img_ids[..., 1] + torch.arange(h // 2)[:, None]
|
||||
img_ids[..., 2] = img_ids[..., 2] + torch.arange(w // 2)[None, :]
|
||||
img_ids = repeat(img_ids, "h w c -> b (h w) c", b=bs)
|
||||
|
||||
ref_img_ids = []
|
||||
ref_imgs_list = []
|
||||
pe_shift_w, pe_shift_h = w // 2, h // 2
|
||||
for ref_img in ref_imgs:
|
||||
_, _, ref_h1, ref_w1 = ref_img.shape
|
||||
ref_img = rearrange(ref_img, "b c (h ph) (w pw) -> b (h w) (c ph pw)", ph=2, pw=2)
|
||||
if ref_img.shape[0] == 1 and bs > 1:
|
||||
ref_img = repeat(ref_img, "1 ... -> bs ...", bs=bs)
|
||||
ref_img_ids1 = torch.zeros(ref_h1 // 2, ref_w1 // 2, 3)
|
||||
# img id分别在宽高偏移各自最大值
|
||||
h_offset = pe_shift_h if pe in {'d', 'h'} else 0
|
||||
w_offset = pe_shift_w if pe in {'d', 'w'} else 0
|
||||
ref_img_ids1[..., 1] = ref_img_ids1[..., 1] + torch.arange(ref_h1 // 2)[:, None] + h_offset
|
||||
ref_img_ids1[..., 2] = ref_img_ids1[..., 2] + torch.arange(ref_w1 // 2)[None, :] + w_offset
|
||||
ref_img_ids1 = repeat(ref_img_ids1, "h w c -> b (h w) c", b=bs)
|
||||
ref_img_ids.append(ref_img_ids1)
|
||||
ref_imgs_list.append(ref_img)
|
||||
|
||||
# 更新pe shift
|
||||
pe_shift_h += ref_h1 // 2
|
||||
pe_shift_w += ref_w1 // 2
|
||||
|
||||
if isinstance(prompt, str):
|
||||
prompt = [prompt]
|
||||
txt = t5(prompt)
|
||||
if txt.shape[0] == 1 and bs > 1:
|
||||
txt = repeat(txt, "1 ... -> bs ...", bs=bs)
|
||||
txt_ids = torch.zeros(bs, txt.shape[1], 3)
|
||||
|
||||
vec = clip(prompt)
|
||||
if vec.shape[0] == 1 and bs > 1:
|
||||
vec = repeat(vec, "1 ... -> bs ...", bs=bs)
|
||||
|
||||
return {
|
||||
"img": img,
|
||||
"img_ids": img_ids.to(img.device),
|
||||
"ref_img": tuple(ref_imgs_list),
|
||||
"ref_img_ids": [ref_img_id.to(img.device) for ref_img_id in ref_img_ids],
|
||||
"txt": txt.to(img.device),
|
||||
"txt_ids": txt_ids.to(img.device),
|
||||
"vec": vec.to(img.device),
|
||||
}
|
||||
def time_shift(mu: float, sigma: float, t: Tensor):
|
||||
return math.exp(mu) / (math.exp(mu) + (1 / t - 1) ** sigma)
|
||||
|
||||
|
||||
def get_lin_function(
|
||||
x1: float = 256, y1: float = 0.5, x2: float = 4096, y2: float = 1.15
|
||||
):
|
||||
m = (y2 - y1) / (x2 - x1)
|
||||
b = y1 - m * x1
|
||||
return lambda x: m * x + b
|
||||
|
||||
|
||||
def get_schedule(
|
||||
num_steps: int,
|
||||
image_seq_len: int,
|
||||
base_shift: float = 0.5,
|
||||
max_shift: float = 1.15,
|
||||
shift: bool = True,
|
||||
) -> list[float]:
|
||||
# extra step for zero
|
||||
timesteps = torch.linspace(1, 0, num_steps + 1)
|
||||
|
||||
# shifting the schedule to favor high timesteps for higher signal images
|
||||
if shift:
|
||||
# eastimate mu based on linear estimation between two points
|
||||
mu = get_lin_function(y1=base_shift, y2=max_shift)(image_seq_len)
|
||||
timesteps = time_shift(mu, 1.0, timesteps)
|
||||
|
||||
return timesteps.tolist()
|
||||
|
||||
|
||||
def denoise(
|
||||
model: Flux,
|
||||
# model input
|
||||
img: Tensor,
|
||||
img_ids: Tensor,
|
||||
txt: Tensor,
|
||||
txt_ids: Tensor,
|
||||
vec: Tensor,
|
||||
# sampling parameters
|
||||
timesteps: list[float],
|
||||
guidance: float = 4.0,
|
||||
ref_img: Tensor=None,
|
||||
ref_img_ids: Tensor=None,
|
||||
):
|
||||
i = 0
|
||||
guidance_vec = torch.full((img.shape[0],), guidance, device=img.device, dtype=img.dtype)
|
||||
#print(img.device)#cuda:0
|
||||
for t_curr, t_prev in tqdm(zip(timesteps[:-1], timesteps[1:]), total=len(timesteps) - 1):
|
||||
t_vec = torch.full((img.shape[0],), t_curr, dtype=img.dtype, device=img.device)
|
||||
#print(t_vec.device,img_ids.device,ref_img_ids,txt_ids.device,vec.device,t_vec.device,guidance_vec.device)
|
||||
pred = model(
|
||||
img=img,
|
||||
img_ids=img_ids,
|
||||
ref_img=ref_img,
|
||||
ref_img_ids=ref_img_ids,
|
||||
txt=txt,
|
||||
txt_ids=txt_ids,
|
||||
y=vec,
|
||||
timesteps=t_vec,
|
||||
guidance=guidance_vec
|
||||
)
|
||||
img = img + (t_prev - t_curr) * pred
|
||||
i += 1
|
||||
return img
|
||||
|
||||
|
||||
def unpack(x: Tensor, height: int, width: int) -> Tensor:
|
||||
return rearrange(
|
||||
x,
|
||||
"b (h w) (c ph pw) -> b c (h ph) (w pw)",
|
||||
h=math.ceil(height / 16),
|
||||
w=math.ceil(width / 16),
|
||||
ph=2,
|
||||
pw=2,
|
||||
)
|
||||
@@ -0,0 +1,495 @@
|
||||
# Copyright (c) 2025 Bytedance Ltd. and/or its affiliates. All rights reserved.
|
||||
# Copyright (c) 2024 Black Forest Labs and The XLabs-AI 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.
|
||||
|
||||
import os
|
||||
from dataclasses import dataclass
|
||||
from datetime import datetime
|
||||
import torch
|
||||
import json
|
||||
import numpy as np
|
||||
from huggingface_hub import hf_hub_download
|
||||
from safetensors import safe_open
|
||||
from safetensors.torch import load_file as load_sft
|
||||
|
||||
from .model import Flux, FluxParams
|
||||
from .modules.autoencoder import AutoEncoder, AutoEncoderParams
|
||||
from .modules.conditioner import HFEmbedder
|
||||
|
||||
import re
|
||||
from .modules.layers import DoubleStreamBlockLoraProcessor, SingleStreamBlockLoraProcessor
|
||||
timestamp_uno = datetime.now().strftime("%Y%m%d_%H%M%S")
|
||||
def load_model(ckpt, device='cpu'):
|
||||
if ckpt.endswith('safetensors'):
|
||||
from safetensors import safe_open
|
||||
pl_sd = {}
|
||||
with safe_open(ckpt, framework="pt", device=device) as f:
|
||||
for k in f.keys():
|
||||
pl_sd[k] = f.get_tensor(k)
|
||||
else:
|
||||
pl_sd = torch.load(ckpt,weights_only=False, map_location=device,)
|
||||
return pl_sd
|
||||
|
||||
def load_safetensors(path):
|
||||
tensors = {}
|
||||
with safe_open(path, framework="pt", device="cpu") as f:
|
||||
for key in f.keys():
|
||||
tensors[key] = f.get_tensor(key)
|
||||
return tensors
|
||||
|
||||
def get_lora_rank(checkpoint):
|
||||
for k in checkpoint.keys():
|
||||
if k.endswith(".down.weight"):
|
||||
return checkpoint[k].shape[0]
|
||||
|
||||
def load_checkpoint(local_path, repo_id, name):
|
||||
if local_path is not None:
|
||||
if '.safetensors' in local_path:
|
||||
print(f"Loading .safetensors checkpoint from {local_path}")
|
||||
checkpoint = load_safetensors(local_path)
|
||||
else:
|
||||
print(f"Loading checkpoint from {local_path}")
|
||||
checkpoint = torch.load(local_path, weights_only=False,map_location='cpu')
|
||||
elif repo_id is not None and name is not None:
|
||||
print(f"Loading checkpoint {name} from repo id {repo_id}")
|
||||
checkpoint = load_from_repo_id(repo_id, name)
|
||||
else:
|
||||
raise ValueError(
|
||||
"LOADING ERROR: you must specify local_path or repo_id with name in HF to download"
|
||||
)
|
||||
return checkpoint
|
||||
|
||||
|
||||
def c_crop(image):
|
||||
width, height = image.size
|
||||
new_size = min(width, height)
|
||||
left = (width - new_size) / 2
|
||||
top = (height - new_size) / 2
|
||||
right = (width + new_size) / 2
|
||||
bottom = (height + new_size) / 2
|
||||
return image.crop((left, top, right, bottom))
|
||||
|
||||
def pad64(x):
|
||||
return int(np.ceil(float(x) / 64.0) * 64 - x)
|
||||
|
||||
def HWC3(x):
|
||||
assert x.dtype == np.uint8
|
||||
if x.ndim == 2:
|
||||
x = x[:, :, None]
|
||||
assert x.ndim == 3
|
||||
H, W, C = x.shape
|
||||
assert C == 1 or C == 3 or C == 4
|
||||
if C == 3:
|
||||
return x
|
||||
if C == 1:
|
||||
return np.concatenate([x, x, x], axis=2)
|
||||
if C == 4:
|
||||
color = x[:, :, 0:3].astype(np.float32)
|
||||
alpha = x[:, :, 3:4].astype(np.float32) / 255.0
|
||||
y = color * alpha + 255.0 * (1.0 - alpha)
|
||||
y = y.clip(0, 255).astype(np.uint8)
|
||||
return y
|
||||
|
||||
@dataclass
|
||||
class ModelSpec:
|
||||
params: FluxParams
|
||||
ae_params: AutoEncoderParams
|
||||
ckpt_path: str | None
|
||||
ae_path: str | None
|
||||
repo_id: str | None
|
||||
repo_flow: str | None
|
||||
repo_ae: str | None
|
||||
repo_id_ae: str | None
|
||||
|
||||
|
||||
configs = {
|
||||
"flux-dev": ModelSpec(
|
||||
repo_id="black-forest-labs/FLUX.1-dev",
|
||||
repo_id_ae="black-forest-labs/FLUX.1-dev",
|
||||
repo_flow="flux1-dev.safetensors",
|
||||
repo_ae="ae.safetensors",
|
||||
ckpt_path=os.getenv("FLUX_DEV"),
|
||||
params=FluxParams(
|
||||
in_channels=64,
|
||||
vec_in_dim=768,
|
||||
context_in_dim=4096,
|
||||
hidden_size=3072,
|
||||
mlp_ratio=4.0,
|
||||
num_heads=24,
|
||||
depth=19,
|
||||
depth_single_blocks=38,
|
||||
axes_dim=[16, 56, 56],
|
||||
theta=10_000,
|
||||
qkv_bias=True,
|
||||
guidance_embed=True,
|
||||
),
|
||||
ae_path=os.getenv("AE"),
|
||||
ae_params=AutoEncoderParams(
|
||||
resolution=256,
|
||||
in_channels=3,
|
||||
ch=128,
|
||||
out_ch=3,
|
||||
ch_mult=[1, 2, 4, 4],
|
||||
num_res_blocks=2,
|
||||
z_channels=16,
|
||||
scale_factor=0.3611,
|
||||
shift_factor=0.1159,
|
||||
),
|
||||
),
|
||||
"flux-dev-fp8": ModelSpec(
|
||||
repo_id="XLabs-AI/flux-dev-fp8",
|
||||
repo_id_ae="black-forest-labs/FLUX.1-dev",
|
||||
repo_flow="flux-dev-fp8.safetensors",
|
||||
repo_ae="ae.safetensors",
|
||||
ckpt_path=os.getenv("FLUX_DEV_FP8"),
|
||||
params=FluxParams(
|
||||
in_channels=64,
|
||||
vec_in_dim=768,
|
||||
context_in_dim=4096,
|
||||
hidden_size=3072,
|
||||
mlp_ratio=4.0,
|
||||
num_heads=24,
|
||||
depth=19,
|
||||
depth_single_blocks=38,
|
||||
axes_dim=[16, 56, 56],
|
||||
theta=10_000,
|
||||
qkv_bias=True,
|
||||
guidance_embed=True,
|
||||
),
|
||||
ae_path=os.getenv("AE"),
|
||||
ae_params=AutoEncoderParams(
|
||||
resolution=256,
|
||||
in_channels=3,
|
||||
ch=128,
|
||||
out_ch=3,
|
||||
ch_mult=[1, 2, 4, 4],
|
||||
num_res_blocks=2,
|
||||
z_channels=16,
|
||||
scale_factor=0.3611,
|
||||
shift_factor=0.1159,
|
||||
),
|
||||
),
|
||||
"flux-schnell": ModelSpec(
|
||||
repo_id="black-forest-labs/FLUX.1-schnell",
|
||||
repo_id_ae="black-forest-labs/FLUX.1-dev",
|
||||
repo_flow="flux1-schnell.safetensors",
|
||||
repo_ae="ae.safetensors",
|
||||
ckpt_path=os.getenv("FLUX_SCHNELL"),
|
||||
params=FluxParams(
|
||||
in_channels=64,
|
||||
vec_in_dim=768,
|
||||
context_in_dim=4096,
|
||||
hidden_size=3072,
|
||||
mlp_ratio=4.0,
|
||||
num_heads=24,
|
||||
depth=19,
|
||||
depth_single_blocks=38,
|
||||
axes_dim=[16, 56, 56],
|
||||
theta=10_000,
|
||||
qkv_bias=True,
|
||||
guidance_embed=False,
|
||||
),
|
||||
ae_path=os.getenv("AE"),
|
||||
ae_params=AutoEncoderParams(
|
||||
resolution=256,
|
||||
in_channels=3,
|
||||
ch=128,
|
||||
out_ch=3,
|
||||
ch_mult=[1, 2, 4, 4],
|
||||
num_res_blocks=2,
|
||||
z_channels=16,
|
||||
scale_factor=0.3611,
|
||||
shift_factor=0.1159,
|
||||
),
|
||||
),
|
||||
}
|
||||
|
||||
|
||||
def print_load_warning(missing: list[str], unexpected: list[str]) -> None:
|
||||
if len(missing) > 0 and len(unexpected) > 0:
|
||||
print(f"Got {len(missing)} missing keys:\n\t" + "\n\t".join(missing))
|
||||
print("\n" + "-" * 79 + "\n")
|
||||
print(f"Got {len(unexpected)} unexpected keys:\n\t" + "\n\t".join(unexpected))
|
||||
elif len(missing) > 0:
|
||||
print(f"Got {len(missing)} missing keys:\n\t" + "\n\t".join(missing))
|
||||
elif len(unexpected) > 0:
|
||||
print(f"Got {len(unexpected)} unexpected keys:\n\t" + "\n\t".join(unexpected))
|
||||
|
||||
def load_from_repo_id(repo_id, checkpoint_name):
|
||||
ckpt_path = hf_hub_download(repo_id, checkpoint_name)
|
||||
sd = load_sft(ckpt_path, device='cpu')
|
||||
return sd
|
||||
|
||||
def load_flow_model(name: str, ckpt_path,use_fp8,device: str | torch.device = "cuda", save_quantezed: bool = False):
|
||||
# Loading Flux
|
||||
print("Init model")
|
||||
# ckpt_path = configs[name].ckpt_path
|
||||
# if (
|
||||
# ckpt_path is None
|
||||
# and configs[name].repo_id is not None
|
||||
# and configs[name].repo_flow is not None
|
||||
# and hf_download
|
||||
# ):
|
||||
# ckpt_path = hf_hub_download(configs[name].repo_id, configs[name].repo_flow)
|
||||
|
||||
with torch.device("meta" if ckpt_path is not None else device):
|
||||
model = Flux(configs[name].params).to(torch.bfloat16)
|
||||
if use_fp8:
|
||||
model=load_flow_model_quintized(name, ckpt_path, device=device,save_quantezed=save_quantezed)
|
||||
else:
|
||||
print("Loading checkpoint")
|
||||
# load_sft doesn't support torch.device
|
||||
sd = load_model(ckpt_path, device=str(device))
|
||||
missing, unexpected = model.load_state_dict(sd, strict=False, assign=True)
|
||||
print_load_warning(missing, unexpected)
|
||||
return model
|
||||
def load_flow_model_(name: str, device: str | torch.device = "cuda", hf_download: bool = True):
|
||||
# Loading Flux
|
||||
print("Init model")
|
||||
ckpt_path = configs[name].ckpt_path
|
||||
if (
|
||||
ckpt_path is None
|
||||
and configs[name].repo_id is not None
|
||||
and configs[name].repo_flow is not None
|
||||
and hf_download
|
||||
):
|
||||
ckpt_path = hf_hub_download(configs[name].repo_id, configs[name].repo_flow)
|
||||
|
||||
with torch.device("meta" if ckpt_path is not None else device):
|
||||
model = Flux(configs[name].params).to(torch.bfloat16)
|
||||
|
||||
if ckpt_path is not None:
|
||||
print("Loading checkpoint")
|
||||
# load_sft doesn't support torch.device
|
||||
sd = load_model(ckpt_path, device=str(device))
|
||||
missing, unexpected = model.load_state_dict(sd, strict=False, assign=True)
|
||||
print_load_warning(missing, unexpected)
|
||||
return model
|
||||
|
||||
|
||||
|
||||
def load_flow_model_only_lora(
|
||||
name: str,
|
||||
ckpt_path,
|
||||
lora_ckpt_path,
|
||||
use_fp8,
|
||||
device: str | torch.device = "cuda",
|
||||
save_quantezed: bool = False,
|
||||
lora_rank: int = 16
|
||||
):
|
||||
# Loading Flux
|
||||
|
||||
with torch.device("meta" if ckpt_path is not None else device):
|
||||
model = Flux(configs[name].params)
|
||||
|
||||
|
||||
model = set_lora(model, lora_rank, device="meta" if lora_ckpt_path is not None else device)
|
||||
|
||||
|
||||
print("Loading lora")
|
||||
lora_sd = load_sft(lora_ckpt_path, device=str(device)) if lora_ckpt_path.endswith("safetensors") else torch.load(lora_ckpt_path,weights_only=False, map_location='cpu')
|
||||
|
||||
#print("Loading main checkpoint")
|
||||
# load_sft doesn't support torch.device
|
||||
|
||||
if ckpt_path.endswith('safetensors'):
|
||||
if not use_fp8:
|
||||
sd = load_sft(ckpt_path, device=str(device))
|
||||
else:
|
||||
sd = load_sft(ckpt_path, device="cpu")
|
||||
sd = {k: v.to(dtype=torch.float8_e4m3fn, device=device) for k, v in sd.items()}
|
||||
if save_quantezed:
|
||||
torch.save(sd, os.path.join(os.path.dirname(ckpt_path), f"flux1-dev-fp8-{timestamp_uno}.safetensors"))
|
||||
sd.update(lora_sd)
|
||||
missing, unexpected = model.load_state_dict(sd, strict=False, assign=True)
|
||||
else:
|
||||
dit_state = torch.load(ckpt_path,weights_only=False, map_location='cpu')
|
||||
sd = {}
|
||||
for k in dit_state.keys():
|
||||
sd[k.replace('module.','')] = dit_state[k]
|
||||
sd.update(lora_sd)
|
||||
missing, unexpected = model.load_state_dict(sd, strict=False, assign=True)
|
||||
model.to(str(device))
|
||||
print_load_warning(missing, unexpected)
|
||||
return model
|
||||
def load_flow_model_only_lora_(
|
||||
name: str,
|
||||
device: str | torch.device = "cuda",
|
||||
hf_download: bool = True,
|
||||
lora_rank: int = 16
|
||||
):
|
||||
# Loading Flux
|
||||
print("Init model")
|
||||
ckpt_path = configs[name].ckpt_path
|
||||
if (
|
||||
ckpt_path is None
|
||||
and configs[name].repo_id is not None
|
||||
and configs[name].repo_flow is not None
|
||||
and hf_download
|
||||
):
|
||||
ckpt_path = hf_hub_download(configs[name].repo_id, configs[name].repo_flow.replace("sft", "safetensors"))
|
||||
|
||||
if hf_download:
|
||||
# lora_ckpt_path = hf_hub_download("bytedance-research/UNO", "dit_lora.safetensors")
|
||||
try:
|
||||
lora_ckpt_path = hf_hub_download("bytedance-research/UNO", "dit_lora.safetensors")
|
||||
except:
|
||||
lora_ckpt_path = os.environ.get("LORA", None)
|
||||
else:
|
||||
lora_ckpt_path = os.environ.get("LORA", None)
|
||||
|
||||
with torch.device("meta" if ckpt_path is not None else device):
|
||||
model = Flux(configs[name].params)
|
||||
|
||||
|
||||
model = set_lora(model, lora_rank, device="meta" if lora_ckpt_path is not None else device)
|
||||
|
||||
if ckpt_path is not None:
|
||||
print("Loading lora")
|
||||
lora_sd = load_sft(lora_ckpt_path, device=str(device)) if lora_ckpt_path.endswith("safetensors")\
|
||||
else torch.load(lora_ckpt_path,weights_only=False, map_location='cpu')
|
||||
|
||||
print("Loading main checkpoint")
|
||||
# load_sft doesn't support torch.device
|
||||
|
||||
if ckpt_path.endswith('safetensors'):
|
||||
sd = load_sft(ckpt_path, device=str(device))
|
||||
sd.update(lora_sd)
|
||||
missing, unexpected = model.load_state_dict(sd, strict=False, assign=True)
|
||||
else:
|
||||
dit_state = torch.load(ckpt_path,weights_only=False, map_location='cpu')
|
||||
sd = {}
|
||||
for k in dit_state.keys():
|
||||
sd[k.replace('module.','')] = dit_state[k]
|
||||
sd.update(lora_sd)
|
||||
missing, unexpected = model.load_state_dict(sd, strict=False, assign=True)
|
||||
model.to(str(device))
|
||||
print_load_warning(missing, unexpected)
|
||||
return model
|
||||
|
||||
def set_lora(
|
||||
model: Flux,
|
||||
lora_rank: int,
|
||||
double_blocks_indices: list[int] | None = None,
|
||||
single_blocks_indices: list[int] | None = None,
|
||||
device: str | torch.device = "cpu",
|
||||
) -> Flux:
|
||||
double_blocks_indices = list(range(model.params.depth)) if double_blocks_indices is None else double_blocks_indices
|
||||
single_blocks_indices = list(range(model.params.depth_single_blocks)) if single_blocks_indices is None \
|
||||
else single_blocks_indices
|
||||
|
||||
lora_attn_procs = {}
|
||||
with torch.device(device):
|
||||
for name, attn_processor in model.attn_processors.items():
|
||||
match = re.search(r'\.(\d+)\.', name)
|
||||
if match:
|
||||
layer_index = int(match.group(1))
|
||||
|
||||
if name.startswith("double_blocks") and layer_index in double_blocks_indices:
|
||||
lora_attn_procs[name] = DoubleStreamBlockLoraProcessor(dim=model.params.hidden_size, rank=lora_rank)
|
||||
elif name.startswith("single_blocks") and layer_index in single_blocks_indices:
|
||||
lora_attn_procs[name] = SingleStreamBlockLoraProcessor(dim=model.params.hidden_size, rank=lora_rank)
|
||||
else:
|
||||
lora_attn_procs[name] = attn_processor
|
||||
model.set_attn_processor(lora_attn_procs)
|
||||
return model
|
||||
|
||||
|
||||
def load_flow_model_quintized(name: str, ckpt_path,device: str | torch.device = "cuda", save_quantezed: bool = True):
|
||||
# Loading Flux
|
||||
#from optimum.quanto import requantize
|
||||
print("Init model")
|
||||
#ckpt_path = configs[name].ckpt_path
|
||||
# if (
|
||||
# ckpt_path is None
|
||||
# and configs[name].repo_id is not None
|
||||
# and configs[name].repo_flow is not None
|
||||
# and hf_download
|
||||
# ):
|
||||
# ckpt_path = hf_hub_download(configs[name].repo_id, configs[name].repo_flow)
|
||||
# json_path = hf_hub_download(configs[name].repo_id, 'flux_dev_quantization_map.json')
|
||||
|
||||
|
||||
model = Flux(configs[name].params).to(torch.bfloat16)
|
||||
|
||||
print("Loading checkpoint")
|
||||
# load_sft doesn't support torch.device
|
||||
sd = load_sft(ckpt_path, device='cpu')
|
||||
sd = {k: v.to(dtype=torch.float8_e4m3fn, device=device) for k, v in sd.items()}
|
||||
if save_quantezed:
|
||||
torch.save(sd, os.path.join(os.path.dirname(ckpt_path), f"flux1-dev-fp8-{timestamp_uno}.safetensors"))
|
||||
model.load_state_dict(sd, assign=True)
|
||||
# with open(json_path, "r") as f:
|
||||
# quantization_map = json.load(f)
|
||||
# print("Start a quantization process...")
|
||||
# requantize(model, sd, quantization_map, device=device)
|
||||
# print("Model is quantized!")
|
||||
return model
|
||||
|
||||
# def load_flow_model_quintized_(name: str, device: str | torch.device = "cuda", hf_download: bool = True):
|
||||
# # Loading Flux
|
||||
# from optimum.quanto import requantize
|
||||
# print("Init model")
|
||||
# ckpt_path = configs[name].ckpt_path
|
||||
# if (
|
||||
# ckpt_path is None
|
||||
# and configs[name].repo_id is not None
|
||||
# and configs[name].repo_flow is not None
|
||||
# and hf_download
|
||||
# ):
|
||||
# ckpt_path = hf_hub_download(configs[name].repo_id, configs[name].repo_flow)
|
||||
# json_path = hf_hub_download(configs[name].repo_id, 'flux_dev_quantization_map.json')
|
||||
|
||||
|
||||
# model = Flux(configs[name].params).to(torch.bfloat16)
|
||||
|
||||
# print("Loading checkpoint")
|
||||
# # load_sft doesn't support torch.device
|
||||
# sd = load_sft(ckpt_path, device='cpu')
|
||||
# with open(json_path, "r") as f:
|
||||
# quantization_map = json.load(f)
|
||||
# print("Start a quantization process...")
|
||||
# requantize(model, sd, quantization_map, device=device)
|
||||
# print("Model is quantized!")
|
||||
# return model
|
||||
def load_t5(device: str | torch.device = "cuda", max_length: int = 512) -> HFEmbedder:
|
||||
# max length 64, 128, 256 and 512 should work (if your sequence is short enough)
|
||||
version = os.environ.get("T5", "xlabs-ai/xflux_text_encoders")
|
||||
return HFEmbedder(version, max_length=max_length, torch_dtype=torch.bfloat16).to(device)
|
||||
|
||||
def load_clip(device: str | torch.device = "cuda") -> HFEmbedder:
|
||||
version = os.environ.get("CLIP", "openai/clip-vit-large-patch14")
|
||||
return HFEmbedder(version, max_length=77, torch_dtype=torch.bfloat16).to(device)
|
||||
|
||||
|
||||
def load_ae(name: str, device: str | torch.device = "cuda", hf_download: bool = True) -> AutoEncoder:
|
||||
ckpt_path = configs[name].ae_path
|
||||
if (
|
||||
ckpt_path is None
|
||||
and configs[name].repo_id is not None
|
||||
and configs[name].repo_ae is not None
|
||||
and hf_download
|
||||
):
|
||||
ckpt_path = hf_hub_download(configs[name].repo_id_ae, configs[name].repo_ae)
|
||||
|
||||
# Loading the autoencoder
|
||||
print("Init AE")
|
||||
with torch.device("meta" if ckpt_path is not None else device):
|
||||
ae = AutoEncoder(configs[name].ae_params)
|
||||
|
||||
if ckpt_path is not None:
|
||||
sd = load_sft(ckpt_path, device=str(device))
|
||||
missing, unexpected = ae.load_state_dict(sd, strict=False, assign=True)
|
||||
print_load_warning(missing, unexpected)
|
||||
return ae
|
||||
@@ -0,0 +1,34 @@
|
||||
# Copyright (c) 2025 Bytedance Ltd. and/or its affiliates. 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.
|
||||
|
||||
import argparse
|
||||
import yaml
|
||||
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--yaml", type=str, required=True)
|
||||
parser.add_argument("--arg", type=str, required=True)
|
||||
args = parser.parse_args()
|
||||
|
||||
|
||||
with open(args.yaml, "r") as f:
|
||||
data = yaml.safe_load(f)
|
||||
|
||||
with open(args.arg, "w") as f:
|
||||
for k, v in data.items():
|
||||
if isinstance(v, list):
|
||||
v = list(map(str, v))
|
||||
v = " ".join(v)
|
||||
if v is None:
|
||||
continue
|
||||
print(f"--{k} {v}", end=" ", file=f)
|
||||
File diff suppressed because one or more lines are too long
Binary file not shown.
|
After Width: | Height: | Size: 467 KiB |
Binary file not shown.
BIN
Binary file not shown.
@@ -157,6 +157,15 @@ def tensortopil_list_upscale(tensor_in,width,height):
|
||||
img_list=[nomarl_upscale(i,width,height) for i in tensor_list]
|
||||
return img_list
|
||||
|
||||
def tensortolist(tensor_in,width,height):
|
||||
d1, _, _, _ = tensor_in.size()
|
||||
if d1 == 1:
|
||||
tensor_list = [nomarl_tensor_upscale(tensor_in,width,height)]
|
||||
else:
|
||||
tensor_list_ = torch.chunk(tensor_in, chunks=d1)
|
||||
tensor_list=[nomarl_tensor_upscale(i,width,height) for i in tensor_list_]
|
||||
return tensor_list
|
||||
|
||||
|
||||
|
||||
def nomarl_tensor_upscale(tensor, width, height):
|
||||
@@ -1718,11 +1727,17 @@ def cf_clip_single(clip,prompt,
|
||||
tokens_p = clip.tokenize(prompt)
|
||||
output_p = clip.encode_from_tokens(tokens_p, return_dict=True) # {"pooled_output":tensor}
|
||||
cond_p = output_p.pop("cond")
|
||||
if cond_p.shape[1] / 77 > 1 :
|
||||
logging.warning("prompt'tokens length is abvoe 77,split it")
|
||||
cond_p = torch.chunk(cond_p, cond_p.shape[1] // 77, dim=1)[0]
|
||||
pool_p=output_p.pop("pooled_output")
|
||||
|
||||
tokens_n = clip.tokenize(negative_prompt)
|
||||
output_n = clip.encode_from_tokens(tokens_n, return_dict=True) # {"pooled_output":tensor}
|
||||
cond_n = output_n.pop("cond")
|
||||
if cond_n.shape[1] / 77 > 1 :
|
||||
logging.warning("prompt'tokens length is abvoe 77,split it")
|
||||
cond_n = torch.chunk(cond_n, cond_n.shape[1] // 77, dim=1)[0]
|
||||
pool_n = output_n.pop("pooled_output")
|
||||
|
||||
return cond_p,cond_n,pool_p,pool_n # TODO: replace the pooled_prompt_embeds with text only prompt
|
||||
@@ -2178,4 +2193,47 @@ def photomaker_clip_v2(clip,model_,prompt_list,negative_prompt,input_id_images,i
|
||||
}
|
||||
return emb_dict
|
||||
|
||||
def Loader_UNO(model,offload,quantize_mode,save_quantezed,lora_rank):
|
||||
from .UNO.uno.flux.pipeline import UNOPipeline, preprocess_ref
|
||||
from accelerate import Accelerator
|
||||
accelerator = Accelerator()
|
||||
device= accelerator.device
|
||||
if isinstance(model, dict):
|
||||
model_path = model["model_path"]
|
||||
lora_ckpt_path= model["lora_ckpt_path"]
|
||||
only_lora=True if lora_ckpt_path is not None else False
|
||||
if lora_ckpt_path:
|
||||
if "dit_lora" not in lora_ckpt_path:
|
||||
raise ValueError("lora_ckpt_path must be a dit_lora checkpoint")
|
||||
|
||||
use_fp8=True if quantize_mode=="fp8" else False
|
||||
|
||||
|
||||
if "schnell" in model_path:
|
||||
model_type="flux-schnell"
|
||||
else:
|
||||
model_type="flux-dev"
|
||||
|
||||
|
||||
if "schnell" in model_path:
|
||||
model_type="flux-schnell"
|
||||
else:
|
||||
model_type="flux-dev"
|
||||
|
||||
if only_lora :
|
||||
from .UNO.uno.flux.util import load_flow_model_only_lora
|
||||
model_ = load_flow_model_only_lora(
|
||||
model_type,model_path,lora_ckpt_path,use_fp8, device="cpu" if offload else device,save_quantezed=save_quantezed, lora_rank=lora_rank)
|
||||
else:
|
||||
from .UNO.uno.flux.util import load_flow_model
|
||||
model_ = load_flow_model(model_type,model_path,use_fp8, device="cpu" if offload else device,save_quantezed=save_quantezed)
|
||||
model = UNOPipeline(
|
||||
model_,
|
||||
device,
|
||||
offload,
|
||||
use_fp8,
|
||||
)
|
||||
else:
|
||||
raise ValueError("must sue lite model node")
|
||||
return model
|
||||
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -560,9 +560,7 @@ class PhotoMakerStableDiffusionXLPipeline(StableDiffusionXLPipeline):
|
||||
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
|
||||
)
|
||||
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)
|
||||
|
||||
Reference in New Issue
Block a user