City
+121
@@ -0,0 +1,121 @@
|
||||
import os
|
||||
import torch
|
||||
import folder_paths
|
||||
from transformers import AutoTokenizer, AutoModelForCausalLM
|
||||
from ..utils.dtype import string_to_dtype
|
||||
from huggingface_hub import snapshot_download
|
||||
|
||||
|
||||
tenc_root = (
|
||||
folder_paths.folder_names_and_paths.get(
|
||||
"text_encoders",
|
||||
folder_paths.folder_names_and_paths.get("clip", [[], set()])
|
||||
)
|
||||
)
|
||||
|
||||
dtypes = [
|
||||
"default",
|
||||
"auto (comfy)",
|
||||
"BF16",
|
||||
"FP32",
|
||||
"FP16",
|
||||
]
|
||||
try: torch.float8_e5m2
|
||||
except AttributeError: print("Torch版本过旧,不支持FP8")
|
||||
else: dtypes += ["FP8 E4M3", "FP8 E5M2"]
|
||||
|
||||
class GemmaLoader:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
devices = ["auto", "cpu", "cuda"]
|
||||
# 支持多GPU
|
||||
for k in range(1, torch.cuda.device_count()):
|
||||
devices.append(f"cuda:{k}")
|
||||
return {
|
||||
"required": {
|
||||
"model_name": (["google/gemma-2-2b-it", "unsloth/gemma-2-2b-it-bnb-4bit"],),
|
||||
"device": (devices, {"default":"cpu"}),
|
||||
"dtype": (dtypes,),
|
||||
}
|
||||
}
|
||||
RETURN_TYPES = ("GEMMA",)
|
||||
FUNCTION = "load_model"
|
||||
CATEGORY = "ExtraModels/Gemma"
|
||||
TITLE = "Gemma Loader"
|
||||
|
||||
def load_model(self, model_name, device, dtype):
|
||||
dtype = string_to_dtype(dtype, "text_encoder")
|
||||
if device == "cpu":
|
||||
assert dtype in [None, torch.float32], f"Can't use dtype '{dtype}' with CPU! Set dtype to 'default'."
|
||||
|
||||
if model_name == 'google/gemma-2-2b-it':
|
||||
text_encoder_dir = os.path.join(folder_paths.models_dir, 'text_encoders', 'models--google--gemma-2-2b-it')
|
||||
if not os.path.exists(os.path.join(text_encoder_dir, 'model.safetensors')):
|
||||
snapshot_download('google/gemma-2-2b-it', local_dir=text_encoder_dir)
|
||||
elif model_name == 'unsloth/gemma-2-2b-it-bnb-4bit':
|
||||
text_encoder_dir = os.path.join(folder_paths.models_dir, 'text_encoders', 'models--unsloth--gemma-2-2b-it-bnb-4bit')
|
||||
if not os.path.exists(os.path.join(text_encoder_dir, 'model.safetensors')):
|
||||
snapshot_download('unsloth/gemma-2-2b-it-bnb-4bit', local_dir=text_encoder_dir)
|
||||
else:
|
||||
raise ValueError('Not implemented!')
|
||||
|
||||
tokenizer = AutoTokenizer.from_pretrained(model_name)
|
||||
text_encoder_model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=dtype)
|
||||
tokenizer.padding_side = "right"
|
||||
text_encoder = text_encoder_model.get_decoder()
|
||||
|
||||
if device != "cpu":
|
||||
text_encoder = text_encoder.to(device)
|
||||
|
||||
return ({
|
||||
"tokenizer": tokenizer,
|
||||
"text_encoder": text_encoder,
|
||||
"text_encoder_model": text_encoder_model
|
||||
},)
|
||||
|
||||
|
||||
class GemmaTextEncode:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"text": ("STRING", {"multiline": True}),
|
||||
"GEMMA": ("GEMMA",),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("CONDITIONING",)
|
||||
FUNCTION = "encode"
|
||||
CATEGORY = "ExtraModels/Gemma"
|
||||
TITLE = "Gemma Text Encode"
|
||||
|
||||
def encode(self, text, GEMMA=None):
|
||||
print(text)
|
||||
tokenizer = GEMMA["tokenizer"]
|
||||
text_encoder = GEMMA["text_encoder"]
|
||||
|
||||
with torch.no_grad():
|
||||
tokens = tokenizer(
|
||||
text,
|
||||
max_length=300,
|
||||
padding="max_length",
|
||||
truncation=True,
|
||||
return_tensors="pt"
|
||||
).to(text_encoder.device)
|
||||
|
||||
cond = text_encoder(tokens.input_ids, tokens.attention_mask)[0]
|
||||
emb_masks = tokens.attention_mask
|
||||
|
||||
cond = cond * emb_masks.unsqueeze(-1)
|
||||
|
||||
return ([[cond, {}]], )
|
||||
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"GemmaLoader": GemmaLoader,
|
||||
"GemmaTextEncode": GemmaTextEncode,
|
||||
}
|
||||
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"GemmaLoader": "Gemma Loader",
|
||||
"GemmaTextEncode": "Gemma Text Encode",
|
||||
}
|
||||
@@ -0,0 +1,98 @@
|
||||
"""
|
||||
List of all Sana model types / settings
|
||||
"""
|
||||
|
||||
sampling_settings = {
|
||||
"shift": 3.0,
|
||||
}
|
||||
|
||||
sana_conf = {
|
||||
"SanaMS_600M_P1_D28": {
|
||||
"target": "SanaMS",
|
||||
"unet_config": {
|
||||
"in_channels": 32,
|
||||
"depth": 28,
|
||||
"hidden_size": 1152,
|
||||
"patch_size": 1,
|
||||
"num_heads": 16,
|
||||
"linear_head_dim": 32,
|
||||
"model_max_length": 300,
|
||||
"y_norm": True,
|
||||
"attn_type": "linear",
|
||||
"ffn_type": "glumbconv",
|
||||
"mlp_ratio": 2.5,
|
||||
"mlp_acts": ["silu", "silu", None],
|
||||
"use_pe": False,
|
||||
"pred_sigma": False,
|
||||
"learn_sigma": False,
|
||||
"fp32_attention": True,
|
||||
},
|
||||
"sampling_settings" : sampling_settings,
|
||||
},
|
||||
"SanaMS_1600M_P1_D20": {
|
||||
"target": "SanaMS",
|
||||
"unet_config": {
|
||||
"in_channels": 32,
|
||||
"depth": 20,
|
||||
"hidden_size": 2240,
|
||||
"patch_size": 1,
|
||||
"num_heads": 20,
|
||||
"linear_head_dim": 32,
|
||||
"model_max_length": 300,
|
||||
"y_norm": True,
|
||||
"attn_type": "linear",
|
||||
"ffn_type": "glumbconv",
|
||||
"mlp_ratio": 2.5,
|
||||
"mlp_acts": ["silu", "silu", None],
|
||||
"use_pe": False,
|
||||
"pred_sigma": False,
|
||||
"learn_sigma": False,
|
||||
"fp32_attention": True,
|
||||
},
|
||||
"sampling_settings" : sampling_settings,
|
||||
},
|
||||
}
|
||||
|
||||
sana_res = {
|
||||
"1024px": { # models/SanaMS 1024x1024
|
||||
'0.25': [512, 2048], '0.26': [512, 1984], '0.27': [512, 1920], '0.28': [512, 1856],
|
||||
'0.32': [576, 1792], '0.33': [576, 1728], '0.35': [576, 1664], '0.40': [640, 1600],
|
||||
'0.42': [640, 1536], '0.48': [704, 1472], '0.50': [704, 1408], '0.52': [704, 1344],
|
||||
'0.57': [768, 1344], '0.60': [768, 1280], '0.68': [832, 1216], '0.72': [832, 1152],
|
||||
'0.78': [896, 1152], '0.82': [896, 1088], '0.88': [960, 1088], '0.94': [960, 1024],
|
||||
'1.00': [1024,1024], '1.07': [1024, 960], '1.13': [1088, 960], '1.21': [1088, 896],
|
||||
'1.29': [1152, 896], '1.38': [1152, 832], '1.46': [1216, 832], '1.67': [1280, 768],
|
||||
'1.75': [1344, 768], '2.00': [1408, 704], '2.09': [1472, 704], '2.40': [1536, 640],
|
||||
'2.50': [1600, 640], '2.89': [1664, 576], '3.00': [1728, 576], '3.11': [1792, 576],
|
||||
'3.62': [1856, 512], '3.75': [1920, 512], '3.88': [1984, 512], '4.00': [2048, 512],
|
||||
},
|
||||
"512px": { # models/SanaMS 512x512
|
||||
'0.25': [256,1024], '0.26': [256, 992], '0.27': [256, 960], '0.28': [256, 928],
|
||||
'0.32': [288, 896], '0.33': [288, 864], '0.35': [288, 832], '0.40': [320, 800],
|
||||
'0.42': [320, 768], '0.48': [352, 736], '0.50': [352, 704], '0.52': [352, 672],
|
||||
'0.57': [384, 672], '0.60': [384, 640], '0.68': [416, 608], '0.72': [416, 576],
|
||||
'0.78': [448, 576], '0.82': [448, 544], '0.88': [480, 544], '0.94': [480, 512],
|
||||
'1.00': [512, 512], '1.07': [512, 480], '1.13': [544, 480], '1.21': [544, 448],
|
||||
'1.29': [576, 448], '1.38': [576, 416], '1.46': [608, 416], '1.67': [640, 384],
|
||||
'1.75': [672, 384], '2.00': [704, 352], '2.09': [736, 352], '2.40': [768, 320],
|
||||
'2.50': [800, 320], '2.89': [832, 288], '3.00': [864, 288], '3.11': [896, 288],
|
||||
'3.62': [928, 256], '3.75': [960, 256], '3.88': [992, 256], '4.00': [1024,256]
|
||||
},
|
||||
"2K": {
|
||||
'0.25': [1024, 4096], '0.26': [1024, 3968], '0.27': [1024, 3840], '0.28': [1024, 3712],
|
||||
'0.32': [1152, 3584], '0.33': [1152, 3456], '0.35': [1152, 3328], '0.40': [1280, 3200],
|
||||
'0.42': [1280, 3072], '0.48': [1408, 2944], '0.50': [1408, 2816], '0.52': [1408, 2688],
|
||||
'0.57': [1536, 2688], '0.60': [1536, 2560], '0.68': [1664, 2432], '0.72': [1664, 2304],
|
||||
'0.78': [1792, 2304], '0.82': [1792, 2176], '0.88': [1920, 2176], '0.94': [1920, 2048],
|
||||
'1.00': [2048, 2048], '1.07': [2048, 1920], '1.13': [2176, 1920], '1.21': [2176, 1792],
|
||||
'1.29': [2304, 1792], '1.38': [2304, 1664], '1.46': [2432, 1664], '1.67': [2560, 1536],
|
||||
'1.75': [2688, 1536], '2.00': [2816, 1408], '2.09': [2944, 1408], '2.40': [3072, 1280],
|
||||
'2.50': [3200, 1280], '2.89': [3328, 1152], '3.00': [3456, 1152], '3.11': [3584, 1152],
|
||||
'3.62': [3712, 1024], '3.75': [3840, 1024], '3.88': [3968, 1024], '4.00': [4096, 1024]
|
||||
}
|
||||
}
|
||||
# These should be the same
|
||||
sana_res.update({
|
||||
"SanaMS_600M_P1_D28": sana_res["1024px"],
|
||||
"SanaMS_1600M_P1_D20": sana_res["1024px"],
|
||||
})
|
||||
@@ -0,0 +1,209 @@
|
||||
# For using the diffusers format weights
|
||||
# Based on the original ComfyUI function +
|
||||
# https://github.com/NVlabs/Sana/blob/main/tools/convert_sana_to_diffusers.py
|
||||
import torch
|
||||
|
||||
|
||||
def get_depth(state_dict):
|
||||
return sum(key.endswith('.attn1.to_k.bias') for key in state_dict.keys())
|
||||
|
||||
def get_lora_depth(state_dict):
|
||||
cnt = max([
|
||||
sum(key.endswith('.attn1.to_k.lora_A.weight') for key in state_dict.keys()),
|
||||
sum(key.endswith('_attn1_to_k.lora_A.weight') for key in state_dict.keys()),
|
||||
sum(key.endswith('.attn1.to_k.lora_up.weight') for key in state_dict.keys()),
|
||||
sum(key.endswith('_attn1_to_k.lora_up.weight') for key in state_dict.keys()),
|
||||
])
|
||||
assert cnt > 0, "Unable to detect model depth!"
|
||||
return cnt
|
||||
|
||||
def get_conversion_map(state_dict):
|
||||
conversion_map = [ # main SD conversion map (Sana reference, HF Diffusers)
|
||||
# Patch embeddings
|
||||
("x_embedder.proj.weight", "pos_embed.proj.weight"),
|
||||
("x_embedder.proj.bias", "pos_embed.proj.bias"),
|
||||
# Caption projection
|
||||
("y_embedder.y_embedding", "caption_projection.y_embedding"),
|
||||
("y_embedder.y_proj.fc1.weight", "caption_projection.linear_1.weight"),
|
||||
("y_embedder.y_proj.fc1.bias", "caption_projection.linear_1.bias"),
|
||||
("y_embedder.y_proj.fc2.weight", "caption_projection.linear_2.weight"),
|
||||
("y_embedder.y_proj.fc2.bias", "caption_projection.linear_2.bias"),
|
||||
# AdaLN-single LN
|
||||
("t_embedder.mlp.0.weight", "adaln_single.emb.timestep_embedder.linear_1.weight"),
|
||||
("t_embedder.mlp.0.bias", "adaln_single.emb.timestep_embedder.linear_1.bias"),
|
||||
("t_embedder.mlp.2.weight", "adaln_single.emb.timestep_embedder.linear_2.weight"),
|
||||
("t_embedder.mlp.2.bias", "adaln_single.emb.timestep_embedder.linear_2.bias"),
|
||||
# Shared norm
|
||||
("t_block.1.weight", "adaln_single.linear.weight"),
|
||||
("t_block.1.bias", "adaln_single.linear.bias"),
|
||||
# Final block
|
||||
("final_layer.linear.weight", "proj_out.weight"),
|
||||
("final_layer.linear.bias", "proj_out.bias"),
|
||||
("final_layer.scale_shift_table", "scale_shift_table"),
|
||||
]
|
||||
|
||||
# Add actual transformer blocks
|
||||
for depth in range(get_depth(state_dict)):
|
||||
# Transformer blocks
|
||||
conversion_map += [
|
||||
(f"blocks.{depth}.scale_shift_table", f"transformer_blocks.{depth}.scale_shift_table"),
|
||||
# Projection
|
||||
(f"blocks.{depth}.attn.proj.weight", f"transformer_blocks.{depth}.attn1.to_out.0.weight"),
|
||||
(f"blocks.{depth}.attn.proj.bias", f"transformer_blocks.{depth}.attn1.to_out.0.bias"),
|
||||
# Feed-forward
|
||||
(f"blocks.{depth}.mlp.fc1.weight", f"transformer_blocks.{depth}.ff.net.0.proj.weight"),
|
||||
(f"blocks.{depth}.mlp.fc1.bias", f"transformer_blocks.{depth}.ff.net.0.proj.bias"),
|
||||
(f"blocks.{depth}.mlp.fc2.weight", f"transformer_blocks.{depth}.ff.net.2.weight"),
|
||||
(f"blocks.{depth}.mlp.fc2.bias", f"transformer_blocks.{depth}.ff.net.2.bias"),
|
||||
# Cross-attention (proj)
|
||||
(f"blocks.{depth}.cross_attn.proj.weight" ,f"transformer_blocks.{depth}.attn2.to_out.0.weight"),
|
||||
(f"blocks.{depth}.cross_attn.proj.bias" ,f"transformer_blocks.{depth}.attn2.to_out.0.bias"),
|
||||
]
|
||||
return conversion_map
|
||||
|
||||
def find_prefix(state_dict, target_key):
|
||||
prefix = ""
|
||||
for k in state_dict.keys():
|
||||
if k.endswith(target_key):
|
||||
prefix = k.split(target_key)[0]
|
||||
break
|
||||
return prefix
|
||||
|
||||
def convert_state_dict(state_dict):
|
||||
cmap = get_conversion_map(state_dict)
|
||||
|
||||
missing = [k for k,v in cmap if v not in state_dict]
|
||||
new_state_dict = {k: state_dict[v] for k,v in cmap if k not in missing}
|
||||
matched = list(v for k,v in cmap if v in state_dict.keys())
|
||||
|
||||
for depth in range(get_depth(state_dict)):
|
||||
for wb in ["weight", "bias"]:
|
||||
# Self Attention
|
||||
key = lambda a: f"transformer_blocks.{depth}.attn1.to_{a}.{wb}"
|
||||
new_state_dict[f"blocks.{depth}.attn.qkv.{wb}"] = torch.cat((
|
||||
state_dict[key('q')], state_dict[key('k')], state_dict[key('v')]
|
||||
), dim=0)
|
||||
matched += [key('q'), key('k'), key('v')]
|
||||
|
||||
# Cross-attention (linear)
|
||||
key = lambda a: f"transformer_blocks.{depth}.attn2.to_{a}.{wb}"
|
||||
new_state_dict[f"blocks.{depth}.cross_attn.q_linear.{wb}"] = state_dict[key('q')]
|
||||
new_state_dict[f"blocks.{depth}.cross_attn.kv_linear.{wb}"] = torch.cat((
|
||||
state_dict[key('k')], state_dict[key('v')]
|
||||
), dim=0)
|
||||
matched += [key('q'), key('k'), key('v')]
|
||||
|
||||
if len(matched) < len(state_dict):
|
||||
print(f"Sana: UNET conversion has leftover keys! ({len(matched)} vs {len(state_dict)})")
|
||||
print(list( set(state_dict.keys()) - set(matched) ))
|
||||
|
||||
if len(missing) > 0:
|
||||
print(f"Sana: UNET conversion has missing keys!")
|
||||
print(missing)
|
||||
|
||||
return new_state_dict
|
||||
|
||||
# Same as above but for LoRA weights:
|
||||
# TODO: Not used yet, need to support LoRA for Sana
|
||||
def convert_lora_state_dict(state_dict, peft=True):
|
||||
# koyha
|
||||
rep_ak = lambda x: x.replace(".weight", ".lora_down.weight")
|
||||
rep_bk = lambda x: x.replace(".weight", ".lora_up.weight")
|
||||
rep_pk = lambda x: x.replace(".weight", ".alpha")
|
||||
if peft: # peft
|
||||
rep_ap = lambda x: x.replace(".weight", ".lora_A.weight")
|
||||
rep_bp = lambda x: x.replace(".weight", ".lora_B.weight")
|
||||
rep_pp = lambda x: x.replace(".weight", ".alpha")
|
||||
|
||||
prefix = find_prefix(state_dict, "adaln_single.linear.lora_A.weight")
|
||||
state_dict = {k[len(prefix):]:v for k,v in state_dict.items()}
|
||||
else: # OneTrainer
|
||||
rep_ap = lambda x: x.replace(".", "_")[:-7] + ".lora_down.weight"
|
||||
rep_bp = lambda x: x.replace(".", "_")[:-7] + ".lora_up.weight"
|
||||
rep_pp = lambda x: x.replace(".", "_")[:-7] + ".alpha"
|
||||
|
||||
prefix = "lora_transformer_"
|
||||
gemma_marker = "lora_te_encoder"
|
||||
gemma_keys = []
|
||||
for key in list(state_dict.keys()):
|
||||
if key.startswith(prefix):
|
||||
state_dict[key[len(prefix):]] = state_dict.pop(key)
|
||||
elif gemma_marker in key:
|
||||
gemma_keys.append(state_dict.pop(key))
|
||||
if len(gemma_keys) > 0:
|
||||
print(f"Text Encoder not supported for Sana LoRA, ignoring {len(gemma_keys)} keys")
|
||||
|
||||
cmap = []
|
||||
cmap_unet = get_conversion_map(state_dict) # todo: 512 model
|
||||
for k, v in cmap_unet:
|
||||
if v.endswith(".weight"):
|
||||
cmap.append((rep_ak(k), rep_ap(v)))
|
||||
cmap.append((rep_bk(k), rep_bp(v)))
|
||||
if not peft:
|
||||
cmap.append((rep_pk(k), rep_pp(v)))
|
||||
|
||||
missing = [k for k,v in cmap if v not in state_dict]
|
||||
new_state_dict = {k: state_dict[v] for k,v in cmap if k not in missing}
|
||||
matched = list(v for k,v in cmap if v in state_dict.keys())
|
||||
|
||||
lora_depth = get_lora_depth(state_dict)
|
||||
for fp, fk in ((rep_ap, rep_ak),(rep_bp, rep_bk)):
|
||||
for depth in range(lora_depth):
|
||||
# Self Attention
|
||||
key = lambda a: fp(f"transformer_blocks.{depth}.attn1.to_{a}.weight")
|
||||
new_state_dict[fk(f"blocks.{depth}.attn.qkv.weight")] = torch.cat((
|
||||
state_dict[key('q')], state_dict[key('k')], state_dict[key('v')]
|
||||
), dim=0)
|
||||
|
||||
matched += [key('q'), key('k'), key('v')]
|
||||
if not peft:
|
||||
akey = lambda a: rep_pp(f"transformer_blocks.{depth}.attn1.to_{a}.weight")
|
||||
new_state_dict[rep_pk((f"blocks.{depth}.attn.qkv.weight"))] = state_dict[akey("q")]
|
||||
matched += [akey('q'), akey('k'), akey('v')]
|
||||
|
||||
# Self Attention projection?
|
||||
key = lambda a: fp(f"transformer_blocks.{depth}.attn1.to_{a}.weight")
|
||||
new_state_dict[fk(f"blocks.{depth}.attn.proj.weight")] = state_dict[key('out.0')]
|
||||
matched += [key('out.0')]
|
||||
|
||||
# Cross-attention (linear)
|
||||
key = lambda a: fp(f"transformer_blocks.{depth}.attn2.to_{a}.weight")
|
||||
new_state_dict[fk(f"blocks.{depth}.cross_attn.q_linear.weight")] = state_dict[key('q')]
|
||||
new_state_dict[fk(f"blocks.{depth}.cross_attn.kv_linear.weight")] = torch.cat((
|
||||
state_dict[key('k')], state_dict[key('v')]
|
||||
), dim=0)
|
||||
matched += [key('q'), key('k'), key('v')]
|
||||
if not peft:
|
||||
akey = lambda a: rep_pp(f"transformer_blocks.{depth}.attn2.to_{a}.weight")
|
||||
new_state_dict[rep_pk((f"blocks.{depth}.cross_attn.q_linear.weight"))] = state_dict[akey("q")]
|
||||
new_state_dict[rep_pk((f"blocks.{depth}.cross_attn.kv_linear.weight"))] = state_dict[akey("k")]
|
||||
matched += [akey('q'), akey('k'), akey('v')]
|
||||
|
||||
# Cross Attention projection?
|
||||
key = lambda a: fp(f"transformer_blocks.{depth}.attn2.to_{a}.weight")
|
||||
new_state_dict[fk(f"blocks.{depth}.cross_attn.proj.weight")] = state_dict[key('out.0')]
|
||||
matched += [key('out.0')]
|
||||
|
||||
try:
|
||||
key = fp(f"transformer_blocks.{depth}.ff.net.0.proj.weight")
|
||||
new_state_dict[fk(f"blocks.{depth}.mlp.fc1.weight")] = state_dict[key]
|
||||
matched += [key]
|
||||
except KeyError:
|
||||
pass
|
||||
|
||||
try:
|
||||
key = fp(f"transformer_blocks.{depth}.ff.net.2.weight")
|
||||
new_state_dict[fk(f"blocks.{depth}.mlp.fc2.weight")] = state_dict[key]
|
||||
matched += [key]
|
||||
except KeyError:
|
||||
pass
|
||||
|
||||
if len(matched) < len(state_dict):
|
||||
print(f"Sana: LoRA conversion has leftover keys! ({len(matched)} vs {len(state_dict)})")
|
||||
print(list( set(state_dict.keys()) - set(matched) ))
|
||||
|
||||
if len(missing) > 0:
|
||||
print(f"Sana: LoRA conversion has missing keys! (probably)")
|
||||
print(missing)
|
||||
|
||||
return new_state_dict
|
||||
+100
@@ -0,0 +1,100 @@
|
||||
import comfy.supported_models_base
|
||||
import comfy.latent_formats
|
||||
import comfy.model_patcher
|
||||
import comfy.model_base
|
||||
import comfy.utils
|
||||
import comfy.conds
|
||||
import torch
|
||||
import math
|
||||
from comfy import model_management
|
||||
from comfy.latent_formats import LatentFormat
|
||||
from .diffusers_convert import convert_state_dict
|
||||
|
||||
|
||||
class SanaLatent(LatentFormat):
|
||||
latent_channels = 32
|
||||
def __init__(self):
|
||||
self.scale_factor = 0.41407
|
||||
|
||||
|
||||
class EXM_Sana(comfy.supported_models_base.BASE):
|
||||
unet_config = {}
|
||||
unet_extra_config = {}
|
||||
latent_format = SanaLatent
|
||||
|
||||
def __init__(self, model_conf):
|
||||
self.model_target = model_conf.get("target")
|
||||
self.unet_config = model_conf.get("unet_config", {})
|
||||
self.sampling_settings = model_conf.get("sampling_settings", {})
|
||||
self.latent_format = self.latent_format()
|
||||
# UNET is handled by extension
|
||||
self.unet_config["disable_unet_model_creation"] = True
|
||||
|
||||
def model_type(self, state_dict, prefix=""):
|
||||
return comfy.model_base.ModelType.FLOW
|
||||
|
||||
|
||||
class EXM_Sana_Model(comfy.model_base.BaseModel):
|
||||
def __init__(self, *args, **kwargs):
|
||||
super().__init__(*args, **kwargs)
|
||||
|
||||
def extra_conds(self, **kwargs):
|
||||
out = super().extra_conds(**kwargs)
|
||||
|
||||
cn_hint = kwargs.get("cn_hint", None)
|
||||
if cn_hint is not None:
|
||||
out["cn_hint"] = comfy.conds.CONDRegular(cn_hint)
|
||||
|
||||
return out
|
||||
|
||||
|
||||
def load_sana(model_path, model_conf):
|
||||
state_dict = comfy.utils.load_torch_file(model_path)
|
||||
state_dict = state_dict.get("model", state_dict)
|
||||
|
||||
# prefix
|
||||
for prefix in ["model.diffusion_model.",]:
|
||||
if any(True for x in state_dict if x.startswith(prefix)):
|
||||
state_dict = {k[len(prefix):]:v for k,v in state_dict.items()}
|
||||
|
||||
# diffusers
|
||||
if "adaln_single.linear.weight" in state_dict:
|
||||
state_dict = convert_state_dict(state_dict) # Diffusers
|
||||
|
||||
parameters = comfy.utils.calculate_parameters(state_dict)
|
||||
unet_dtype = comfy.model_management.unet_dtype()
|
||||
load_device = comfy.model_management.get_torch_device()
|
||||
offload_device = comfy.model_management.unet_offload_device()
|
||||
|
||||
# ignore fp8/etc and use directly for now
|
||||
manual_cast_dtype = model_management.unet_manual_cast(unet_dtype, load_device)
|
||||
if manual_cast_dtype:
|
||||
print(f"Sana: falling back to {manual_cast_dtype}")
|
||||
unet_dtype = manual_cast_dtype
|
||||
|
||||
model_conf = EXM_Sana(model_conf) # convert to object
|
||||
model = EXM_Sana_Model( # same as comfy.model_base.BaseModel
|
||||
model_conf,
|
||||
model_type=comfy.model_base.ModelType.FLOW,
|
||||
device=model_management.get_torch_device()
|
||||
)
|
||||
|
||||
if model_conf.model_target == "SanaMS":
|
||||
from .models.sana_multi_scale import SanaMS
|
||||
model.diffusion_model = SanaMS(**model_conf.unet_config)
|
||||
else:
|
||||
raise NotImplementedError(f"Unknown model target '{model_conf.model_target}'")
|
||||
|
||||
m, u = model.diffusion_model.load_state_dict(state_dict, strict=False)
|
||||
if len(m) > 0: print("Missing UNET keys", m)
|
||||
if len(u) > 0: print("Leftover UNET keys", u)
|
||||
model.diffusion_model.dtype = unet_dtype
|
||||
model.diffusion_model.eval()
|
||||
model.diffusion_model.to(unet_dtype)
|
||||
|
||||
model_patcher = comfy.model_patcher.ModelPatcher(
|
||||
model,
|
||||
load_device = load_device,
|
||||
offload_device = offload_device,
|
||||
)
|
||||
return model_patcher
|
||||
@@ -0,0 +1,59 @@
|
||||
# Copyright 2024 NVIDIA CORPORATION & AFFILIATES
|
||||
#
|
||||
# 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.
|
||||
#
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
import copy
|
||||
|
||||
import torch.nn as nn
|
||||
|
||||
__all__ = ["build_act", "get_act_name"]
|
||||
|
||||
# register activation function here
|
||||
# name: module, kwargs with default values
|
||||
REGISTERED_ACT_DICT: dict[str, tuple[type, dict[str, any]]] = {
|
||||
"relu": (nn.ReLU, {"inplace": True}),
|
||||
"relu6": (nn.ReLU6, {"inplace": True}),
|
||||
"hswish": (nn.Hardswish, {"inplace": True}),
|
||||
"hsigmoid": (nn.Hardsigmoid, {"inplace": True}),
|
||||
"swish": (nn.SiLU, {"inplace": True}),
|
||||
"silu": (nn.SiLU, {"inplace": True}),
|
||||
"tanh": (nn.Tanh, {}),
|
||||
"sigmoid": (nn.Sigmoid, {}),
|
||||
"gelu": (nn.GELU, {"approximate": "tanh"}),
|
||||
"mish": (nn.Mish, {"inplace": True}),
|
||||
"identity": (nn.Identity, {}),
|
||||
}
|
||||
|
||||
|
||||
def build_act(name: str or None, **kwargs) -> nn.Module or None:
|
||||
if name in REGISTERED_ACT_DICT:
|
||||
act_cls, default_args = copy.deepcopy(REGISTERED_ACT_DICT[name])
|
||||
for key in default_args:
|
||||
if key in kwargs:
|
||||
default_args[key] = kwargs[key]
|
||||
return act_cls(**default_args)
|
||||
elif name is None or name.lower() == "none":
|
||||
return None
|
||||
else:
|
||||
raise ValueError(f"do not support: {name}")
|
||||
|
||||
|
||||
def get_act_name(act: nn.Module or None) -> str or None:
|
||||
if act is None:
|
||||
return None
|
||||
module2name = {}
|
||||
for key, config in REGISTERED_ACT_DICT.items():
|
||||
module2name[config[0].__name__] = key
|
||||
return module2name.get(type(act).__name__, "unknown")
|
||||
@@ -0,0 +1,361 @@
|
||||
# Copyright 2024 NVIDIA CORPORATION & AFFILIATES
|
||||
#
|
||||
# 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.
|
||||
#
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
# This file is modified from https://github.com/PixArt-alpha/PixArt-sigma
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from timm.models.vision_transformer import Mlp
|
||||
|
||||
from .act import build_act, get_act_name
|
||||
from .norms import build_norm, get_norm_name
|
||||
from .utils import get_same_padding, val2tuple
|
||||
|
||||
|
||||
class ConvLayer(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
in_dim: int,
|
||||
out_dim: int,
|
||||
kernel_size=3,
|
||||
stride=1,
|
||||
dilation=1,
|
||||
groups=1,
|
||||
padding: int or None = None,
|
||||
use_bias=False,
|
||||
dropout=0.0,
|
||||
norm="bn2d",
|
||||
act="relu",
|
||||
):
|
||||
super().__init__()
|
||||
if padding is None:
|
||||
padding = get_same_padding(kernel_size)
|
||||
padding *= dilation
|
||||
|
||||
self.in_dim = in_dim
|
||||
self.out_dim = out_dim
|
||||
self.kernel_size = kernel_size
|
||||
self.stride = stride
|
||||
self.dilation = dilation
|
||||
self.groups = groups
|
||||
self.padding = padding
|
||||
self.use_bias = use_bias
|
||||
|
||||
self.dropout = nn.Dropout2d(dropout, inplace=False) if dropout > 0 else None
|
||||
self.conv = nn.Conv2d(
|
||||
in_dim,
|
||||
out_dim,
|
||||
kernel_size=(kernel_size, kernel_size),
|
||||
stride=(stride, stride),
|
||||
padding=padding,
|
||||
dilation=(dilation, dilation),
|
||||
groups=groups,
|
||||
bias=use_bias,
|
||||
)
|
||||
self.norm = build_norm(norm, num_features=out_dim)
|
||||
self.act = build_act(act)
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
if self.dropout is not None:
|
||||
x = self.dropout(x)
|
||||
x = self.conv(x)
|
||||
if self.norm:
|
||||
x = self.norm(x)
|
||||
if self.act:
|
||||
x = self.act(x)
|
||||
return x
|
||||
|
||||
|
||||
class GLUMBConv(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
in_features: int,
|
||||
hidden_features: int,
|
||||
out_feature=None,
|
||||
kernel_size=3,
|
||||
stride=1,
|
||||
padding: int or None = None,
|
||||
use_bias=False,
|
||||
norm=(None, None, None),
|
||||
act=("silu", "silu", None),
|
||||
dilation=1,
|
||||
):
|
||||
out_feature = out_feature or in_features
|
||||
super().__init__()
|
||||
use_bias = val2tuple(use_bias, 3)
|
||||
norm = val2tuple(norm, 3)
|
||||
act = val2tuple(act, 3)
|
||||
|
||||
self.glu_act = build_act(act[1], inplace=False)
|
||||
self.inverted_conv = ConvLayer(
|
||||
in_features,
|
||||
hidden_features * 2,
|
||||
1,
|
||||
use_bias=use_bias[0],
|
||||
norm=norm[0],
|
||||
act=act[0],
|
||||
)
|
||||
self.depth_conv = ConvLayer(
|
||||
hidden_features * 2,
|
||||
hidden_features * 2,
|
||||
kernel_size,
|
||||
stride=stride,
|
||||
groups=hidden_features * 2,
|
||||
padding=padding,
|
||||
use_bias=use_bias[1],
|
||||
norm=norm[1],
|
||||
act=None,
|
||||
dilation=dilation,
|
||||
)
|
||||
self.point_conv = ConvLayer(
|
||||
hidden_features,
|
||||
out_feature,
|
||||
1,
|
||||
use_bias=use_bias[2],
|
||||
norm=norm[2],
|
||||
act=act[2],
|
||||
)
|
||||
# from IPython import embed; embed(header='debug dilate conv')
|
||||
|
||||
def forward(self, x: torch.Tensor, HW=None) -> torch.Tensor:
|
||||
B, N, C = x.shape
|
||||
if HW is None:
|
||||
H = W = int(N**0.5)
|
||||
else:
|
||||
H, W = HW
|
||||
|
||||
x = x.reshape(B, H, W, C).permute(0, 3, 1, 2)
|
||||
x = self.inverted_conv(x)
|
||||
x = self.depth_conv(x)
|
||||
|
||||
x, gate = torch.chunk(x, 2, dim=1)
|
||||
gate = self.glu_act(gate)
|
||||
x = x * gate
|
||||
|
||||
x = self.point_conv(x)
|
||||
x = x.reshape(B, C, N).permute(0, 2, 1)
|
||||
|
||||
return x
|
||||
|
||||
|
||||
class SlimGLUMBConv(GLUMBConv):
|
||||
def __init__(self, *args, **kwargs):
|
||||
super().__init__(*args, **kwargs)
|
||||
|
||||
# 移除 self.inverted_conv 层
|
||||
del self.inverted_conv
|
||||
self.out_dim = self.point_conv.out_dim
|
||||
|
||||
def forward(self, x: torch.Tensor, HW=None) -> torch.Tensor:
|
||||
B, N, C = x.shape
|
||||
if HW is None:
|
||||
H = W = int(N**0.5)
|
||||
else:
|
||||
H, W = HW
|
||||
|
||||
# 直接使用 x，跳过 self.inverted_conv 层的调用
|
||||
x = x.reshape(B, H, W, C).permute(0, 3, 1, 2)
|
||||
# x = self.inverted_conv(x)
|
||||
x = self.depth_conv(x)
|
||||
|
||||
x, gate = torch.chunk(x, 2, dim=1)
|
||||
gate = self.glu_act(gate)
|
||||
x = x * gate
|
||||
|
||||
x = self.point_conv(x)
|
||||
x = x.reshape(B, self.out_dim, N).permute(0, 2, 1)
|
||||
|
||||
return x
|
||||
|
||||
|
||||
class MBConvPreGLU(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
in_dim: int,
|
||||
out_dim: int,
|
||||
kernel_size=3,
|
||||
stride=1,
|
||||
mid_dim=None,
|
||||
expand=6,
|
||||
padding: int or None = None,
|
||||
use_bias=False,
|
||||
norm=(None, None, "ln2d"),
|
||||
act=("silu", "silu", None),
|
||||
):
|
||||
super().__init__()
|
||||
use_bias = val2tuple(use_bias, 3)
|
||||
norm = val2tuple(norm, 3)
|
||||
act = val2tuple(act, 3)
|
||||
|
||||
mid_dim = mid_dim or round(in_dim * expand)
|
||||
|
||||
self.inverted_conv = ConvLayer(
|
||||
in_dim,
|
||||
mid_dim * 2,
|
||||
1,
|
||||
use_bias=use_bias[0],
|
||||
norm=norm[0],
|
||||
act=None,
|
||||
)
|
||||
self.glu_act = build_act(act[0], inplace=False)
|
||||
self.depth_conv = ConvLayer(
|
||||
mid_dim,
|
||||
mid_dim,
|
||||
kernel_size,
|
||||
stride=stride,
|
||||
groups=mid_dim,
|
||||
padding=padding,
|
||||
use_bias=use_bias[1],
|
||||
norm=norm[1],
|
||||
act=act[1],
|
||||
)
|
||||
self.point_conv = ConvLayer(
|
||||
mid_dim,
|
||||
out_dim,
|
||||
1,
|
||||
use_bias=use_bias[2],
|
||||
norm=norm[2],
|
||||
act=act[2],
|
||||
)
|
||||
|
||||
def forward(self, x: torch.Tensor, HW=None) -> torch.Tensor:
|
||||
B, N, C = x.shape
|
||||
if HW is None:
|
||||
H = W = int(N**0.5)
|
||||
else:
|
||||
H, W = HW
|
||||
|
||||
x = x.reshape(B, H, W, C).permute(0, 3, 1, 2)
|
||||
|
||||
x = self.inverted_conv(x)
|
||||
x, gate = torch.chunk(x, 2, dim=1)
|
||||
gate = self.glu_act(gate)
|
||||
x = x * gate
|
||||
|
||||
x = self.depth_conv(x)
|
||||
x = self.point_conv(x)
|
||||
|
||||
x = x.reshape(B, C, N).permute(0, 2, 1)
|
||||
return x
|
||||
|
||||
@property
|
||||
def module_str(self) -> str:
|
||||
_str = f"{self.depth_conv.kernel_size}{type(self).__name__}("
|
||||
_str += f"in={self.inverted_conv.in_dim},mid={self.depth_conv.in_dim},out={self.point_conv.out_dim},s={self.depth_conv.stride}"
|
||||
_str += (
|
||||
f",norm={get_norm_name(self.inverted_conv.norm)}"
|
||||
f"+{get_norm_name(self.depth_conv.norm)}"
|
||||
f"+{get_norm_name(self.point_conv.norm)}"
|
||||
)
|
||||
_str += (
|
||||
f",act={get_act_name(self.inverted_conv.act)}"
|
||||
f"+{get_act_name(self.depth_conv.act)}"
|
||||
f"+{get_act_name(self.point_conv.act)}"
|
||||
)
|
||||
_str += f",glu_act={get_act_name(self.glu_act)})"
|
||||
return _str
|
||||
|
||||
|
||||
class DWMlp(Mlp):
|
||||
"""MLP as used in Vision Transformer, MLP-Mixer and related networks"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
in_features,
|
||||
hidden_features=None,
|
||||
out_features=None,
|
||||
act_layer=nn.GELU,
|
||||
bias=True,
|
||||
drop=0.0,
|
||||
kernel_size=3,
|
||||
stride=1,
|
||||
dilation=1,
|
||||
padding=None,
|
||||
):
|
||||
super().__init__(
|
||||
in_features=in_features,
|
||||
hidden_features=hidden_features,
|
||||
out_features=out_features,
|
||||
act_layer=act_layer,
|
||||
bias=bias,
|
||||
drop=drop,
|
||||
)
|
||||
hidden_features = hidden_features or in_features
|
||||
self.hidden_features = hidden_features
|
||||
if padding is None:
|
||||
padding = get_same_padding(kernel_size)
|
||||
padding *= dilation
|
||||
|
||||
self.conv = nn.Conv2d(
|
||||
hidden_features,
|
||||
hidden_features,
|
||||
kernel_size=(kernel_size, kernel_size),
|
||||
stride=(stride, stride),
|
||||
padding=padding,
|
||||
dilation=(dilation, dilation),
|
||||
groups=hidden_features,
|
||||
bias=bias,
|
||||
)
|
||||
|
||||
def forward(self, x, HW=None):
|
||||
B, N, C = x.shape
|
||||
if HW is None:
|
||||
H = W = int(N**0.5)
|
||||
else:
|
||||
H, W = HW
|
||||
x = self.fc1(x)
|
||||
x = self.act(x)
|
||||
x = self.drop1(x)
|
||||
x = x.reshape(B, H, W, self.hidden_features).permute(0, 3, 1, 2)
|
||||
x = self.conv(x)
|
||||
x = x.reshape(B, self.hidden_features, N).permute(0, 2, 1)
|
||||
x = self.fc2(x)
|
||||
x = self.drop2(x)
|
||||
return x
|
||||
|
||||
|
||||
class Mlp(Mlp):
|
||||
"""MLP as used in Vision Transformer, MLP-Mixer and related networks"""
|
||||
|
||||
def __init__(self, in_features, hidden_features=None, out_features=None, act_layer=nn.GELU, bias=True, drop=0.0):
|
||||
super().__init__(
|
||||
in_features=in_features,
|
||||
hidden_features=hidden_features,
|
||||
out_features=out_features,
|
||||
act_layer=act_layer,
|
||||
bias=bias,
|
||||
drop=drop,
|
||||
)
|
||||
|
||||
def forward(self, x, HW=None):
|
||||
x = self.fc1(x)
|
||||
x = self.act(x)
|
||||
x = self.drop1(x)
|
||||
x = self.fc2(x)
|
||||
x = self.drop2(x)
|
||||
return x
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
model = GLUMBConv(
|
||||
1152,
|
||||
1152 * 4,
|
||||
1152,
|
||||
use_bias=(True, True, False),
|
||||
norm=(None, None, None),
|
||||
act=("silu", "silu", None),
|
||||
).cuda()
|
||||
input = torch.randn(4, 256, 1152).cuda()
|
||||
output = model(input)
|
||||
@@ -0,0 +1,225 @@
|
||||
# Copyright 2024 NVIDIA CORPORATION & AFFILIATES
|
||||
#
|
||||
# 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.
|
||||
#
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
import copy
|
||||
import warnings
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from torch.nn.modules.batchnorm import _BatchNorm
|
||||
|
||||
__all__ = ["LayerNorm2d", "build_norm", "get_norm_name", "reset_bn", "remove_bn", "set_norm_eps"]
|
||||
|
||||
|
||||
class LayerNorm2d(nn.LayerNorm):
|
||||
rmsnorm = False
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
out = x if LayerNorm2d.rmsnorm else x - torch.mean(x, dim=1, keepdim=True)
|
||||
out = out / torch.sqrt(torch.square(out).mean(dim=1, keepdim=True) + self.eps)
|
||||
if self.elementwise_affine:
|
||||
out = out * self.weight.view(1, -1, 1, 1) + self.bias.view(1, -1, 1, 1)
|
||||
return out
|
||||
|
||||
def extra_repr(self) -> str:
|
||||
return f"{self.normalized_shape}, eps={self.eps}, elementwise_affine={self.elementwise_affine}, rmsnorm={self.rmsnorm}"
|
||||
|
||||
|
||||
# register normalization function here
|
||||
# name: module, kwargs with default values
|
||||
REGISTERED_NORMALIZATION_DICT: dict[str, tuple[type, dict[str, any]]] = {
|
||||
"bn2d": (nn.BatchNorm2d, {"num_features": None, "eps": 1e-5, "momentum": 0.1, "affine": True}),
|
||||
"syncbn": (nn.SyncBatchNorm, {"num_features": None, "eps": 1e-5, "momentum": 0.1, "affine": True}),
|
||||
"ln": (nn.LayerNorm, {"normalized_shape": None, "eps": 1e-5, "elementwise_affine": True}),
|
||||
"ln2d": (LayerNorm2d, {"normalized_shape": None, "eps": 1e-5, "elementwise_affine": True}),
|
||||
}
|
||||
|
||||
|
||||
def build_norm(name="bn2d", num_features=None, affine=True, **kwargs) -> nn.Module or None:
|
||||
if name in ["ln", "ln2d"]:
|
||||
kwargs["normalized_shape"] = num_features
|
||||
kwargs["elementwise_affine"] = affine
|
||||
else:
|
||||
kwargs["num_features"] = num_features
|
||||
kwargs["affine"] = affine
|
||||
if name in REGISTERED_NORMALIZATION_DICT:
|
||||
norm_cls, default_args = copy.deepcopy(REGISTERED_NORMALIZATION_DICT[name])
|
||||
for key in default_args:
|
||||
if key in kwargs:
|
||||
default_args[key] = kwargs[key]
|
||||
return norm_cls(**default_args)
|
||||
elif name is None or name.lower() == "none":
|
||||
return None
|
||||
else:
|
||||
raise ValueError("do not support: %s" % name)
|
||||
|
||||
|
||||
def get_norm_name(norm: nn.Module or None) -> str or None:
|
||||
if norm is None:
|
||||
return None
|
||||
module2name = {}
|
||||
for key, config in REGISTERED_NORMALIZATION_DICT.items():
|
||||
module2name[config[0].__name__] = key
|
||||
return module2name.get(type(norm).__name__, "unknown")
|
||||
|
||||
|
||||
def reset_bn(
|
||||
model: nn.Module,
|
||||
data_loader: list,
|
||||
sync=True,
|
||||
progress_bar=False,
|
||||
) -> None:
|
||||
import copy
|
||||
|
||||
import torch.nn.functional as F
|
||||
from packages.apps.utils import AverageMeter, is_master, sync_tensor
|
||||
from packages.models.utils import get_device, list_join
|
||||
from tqdm import tqdm
|
||||
|
||||
bn_mean = {}
|
||||
bn_var = {}
|
||||
|
||||
tmp_model = copy.deepcopy(model)
|
||||
for name, m in tmp_model.named_modules():
|
||||
if isinstance(m, _BatchNorm):
|
||||
bn_mean[name] = AverageMeter(is_distributed=False)
|
||||
bn_var[name] = AverageMeter(is_distributed=False)
|
||||
|
||||
def new_forward(bn, mean_est, var_est):
|
||||
def lambda_forward(x):
|
||||
x = x.contiguous()
|
||||
if sync:
|
||||
batch_mean = x.mean(0, keepdim=True).mean(2, keepdim=True).mean(3, keepdim=True) # 1, C, 1, 1
|
||||
batch_mean = sync_tensor(batch_mean, reduce="cat")
|
||||
batch_mean = torch.mean(batch_mean, dim=0, keepdim=True)
|
||||
|
||||
batch_var = (x - batch_mean) * (x - batch_mean)
|
||||
batch_var = batch_var.mean(0, keepdim=True).mean(2, keepdim=True).mean(3, keepdim=True)
|
||||
batch_var = sync_tensor(batch_var, reduce="cat")
|
||||
batch_var = torch.mean(batch_var, dim=0, keepdim=True)
|
||||
else:
|
||||
batch_mean = x.mean(0, keepdim=True).mean(2, keepdim=True).mean(3, keepdim=True) # 1, C, 1, 1
|
||||
batch_var = (x - batch_mean) * (x - batch_mean)
|
||||
batch_var = batch_var.mean(0, keepdim=True).mean(2, keepdim=True).mean(3, keepdim=True)
|
||||
|
||||
batch_mean = torch.squeeze(batch_mean)
|
||||
batch_var = torch.squeeze(batch_var)
|
||||
|
||||
mean_est.update(batch_mean.data, x.size(0))
|
||||
var_est.update(batch_var.data, x.size(0))
|
||||
|
||||
# bn forward using calculated mean & var
|
||||
_feature_dim = batch_mean.shape[0]
|
||||
return F.batch_norm(
|
||||
x,
|
||||
batch_mean,
|
||||
batch_var,
|
||||
bn.weight[:_feature_dim],
|
||||
bn.bias[:_feature_dim],
|
||||
False,
|
||||
0.0,
|
||||
bn.eps,
|
||||
)
|
||||
|
||||
return lambda_forward
|
||||
|
||||
m.forward = new_forward(m, bn_mean[name], bn_var[name])
|
||||
|
||||
# skip if there is no batch normalization layers in the network
|
||||
if len(bn_mean) == 0:
|
||||
return
|
||||
|
||||
tmp_model.eval()
|
||||
with torch.inference_mode():
|
||||
with tqdm(total=len(data_loader), desc="reset bn", disable=not progress_bar or not is_master()) as t:
|
||||
for images in data_loader:
|
||||
images = images.to(get_device(tmp_model))
|
||||
tmp_model(images)
|
||||
t.set_postfix(
|
||||
{
|
||||
"bs": images.size(0),
|
||||
"res": list_join(images.shape[-2:], "x"),
|
||||
}
|
||||
)
|
||||
t.update()
|
||||
|
||||
for name, m in model.named_modules():
|
||||
if name in bn_mean and bn_mean[name].count > 0:
|
||||
feature_dim = bn_mean[name].avg.size(0)
|
||||
assert isinstance(m, _BatchNorm)
|
||||
m.running_mean.data[:feature_dim].copy_(bn_mean[name].avg)
|
||||
m.running_var.data[:feature_dim].copy_(bn_var[name].avg)
|
||||
|
||||
|
||||
def remove_bn(model: nn.Module) -> None:
|
||||
for m in model.modules():
|
||||
if isinstance(m, _BatchNorm):
|
||||
m.weight = m.bias = None
|
||||
m.forward = lambda x: x
|
||||
|
||||
|
||||
def set_norm_eps(model: nn.Module, eps: float or None = None, momentum: float or None = None) -> None:
|
||||
for m in model.modules():
|
||||
if isinstance(m, (nn.GroupNorm, nn.LayerNorm, _BatchNorm)):
|
||||
if eps is not None:
|
||||
m.eps = eps
|
||||
if momentum is not None:
|
||||
m.momentum = momentum
|
||||
|
||||
|
||||
class RMSNorm(torch.nn.Module):
|
||||
def __init__(self, dim: int, scale_factor=1.0, eps: float = 1e-6):
|
||||
"""
|
||||
Initialize the RMSNorm normalization layer.
|
||||
|
||||
Args:
|
||||
dim (int): The dimension of the input tensor.
|
||||
eps (float, optional): A small value added to the denominator for numerical stability. Default is 1e-6.
|
||||
|
||||
Attributes:
|
||||
eps (float): A small value added to the denominator for numerical stability.
|
||||
weight (nn.Parameter): Learnable scaling parameter.
|
||||
|
||||
"""
|
||||
super().__init__()
|
||||
self.eps = eps
|
||||
self.weight = nn.Parameter(torch.ones(dim) * scale_factor)
|
||||
|
||||
def _norm(self, x):
|
||||
"""
|
||||
Apply the RMSNorm normalization to the input tensor.
|
||||
|
||||
Args:
|
||||
x (torch.Tensor): The input tensor.
|
||||
|
||||
Returns:
|
||||
torch.Tensor: The normalized tensor.
|
||||
|
||||
"""
|
||||
return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
|
||||
|
||||
def forward(self, x):
|
||||
"""
|
||||
Forward pass through the RMSNorm layer.
|
||||
|
||||
Args:
|
||||
x (torch.Tensor): The input tensor.
|
||||
|
||||
Returns:
|
||||
torch.Tensor: The output tensor after applying RMSNorm.
|
||||
|
||||
"""
|
||||
return (self.weight * self._norm(x.float())).type_as(x)
|
||||
@@ -0,0 +1,378 @@
|
||||
# Copyright 2024 NVIDIA CORPORATION & AFFILIATES
|
||||
#
|
||||
# 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.
|
||||
#
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
# This file is modified from https://github.com/PixArt-alpha/PixArt-sigma
|
||||
import os
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from timm.models.layers import DropPath
|
||||
|
||||
from .basic_modules import DWMlp, GLUMBConv, MBConvPreGLU, Mlp
|
||||
from .sana_blocks import (
|
||||
Attention,
|
||||
CaptionEmbedder,
|
||||
FlashAttention,
|
||||
LiteLA,
|
||||
MultiHeadCrossAttention,
|
||||
PatchEmbed,
|
||||
T2IFinalLayer,
|
||||
TimestepEmbedder,
|
||||
t2i_modulate,
|
||||
)
|
||||
from .norms import RMSNorm
|
||||
from .utils import auto_grad_checkpoint, to_2tuple
|
||||
|
||||
|
||||
class SanaBlock(nn.Module):
|
||||
"""
|
||||
A Sana block with global shared adaptive layer norm (adaLN-single) conditioning.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
hidden_size,
|
||||
num_heads,
|
||||
mlp_ratio=4.0,
|
||||
drop_path=0,
|
||||
input_size=None,
|
||||
qk_norm=False,
|
||||
attn_type="flash",
|
||||
ffn_type="mlp",
|
||||
mlp_acts=("silu", "silu", None),
|
||||
linear_head_dim=32,
|
||||
**block_kwargs,
|
||||
):
|
||||
super().__init__()
|
||||
self.norm1 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
|
||||
if attn_type == "flash":
|
||||
# flash self attention
|
||||
self.attn = FlashAttention(
|
||||
hidden_size,
|
||||
num_heads=num_heads,
|
||||
qkv_bias=True,
|
||||
qk_norm=qk_norm,
|
||||
**block_kwargs,
|
||||
)
|
||||
elif attn_type == "linear":
|
||||
# linear self attention
|
||||
# TODO: Here the num_heads set to 36 for tmp used
|
||||
self_num_heads = hidden_size // linear_head_dim
|
||||
self.attn = LiteLA(hidden_size, hidden_size, heads=self_num_heads, eps=1e-8, qk_norm=qk_norm)
|
||||
elif attn_type == "vanilla":
|
||||
# vanilla self attention
|
||||
self.attn = Attention(hidden_size, num_heads=num_heads, qkv_bias=True)
|
||||
else:
|
||||
raise ValueError(f"{attn_type} type is not defined.")
|
||||
|
||||
self.cross_attn = MultiHeadCrossAttention(hidden_size, num_heads, **block_kwargs)
|
||||
self.norm2 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
|
||||
# to be compatible with lower version pytorch
|
||||
if ffn_type == "dwmlp":
|
||||
approx_gelu = lambda: nn.GELU(approximate="tanh")
|
||||
self.mlp = DWMlp(
|
||||
in_features=hidden_size, hidden_features=int(hidden_size * mlp_ratio), act_layer=approx_gelu, drop=0
|
||||
)
|
||||
elif ffn_type == "glumbconv":
|
||||
self.mlp = GLUMBConv(
|
||||
in_features=hidden_size,
|
||||
hidden_features=int(hidden_size * mlp_ratio),
|
||||
use_bias=(True, True, False),
|
||||
norm=(None, None, None),
|
||||
act=mlp_acts,
|
||||
)
|
||||
elif ffn_type == "glumbconv_dilate":
|
||||
self.mlp = GLUMBConv(
|
||||
in_features=hidden_size,
|
||||
hidden_features=int(hidden_size * mlp_ratio),
|
||||
use_bias=(True, True, False),
|
||||
norm=(None, None, None),
|
||||
act=mlp_acts,
|
||||
dilation=2,
|
||||
)
|
||||
elif ffn_type == "mbconvpreglu":
|
||||
self.mlp = MBConvPreGLU(
|
||||
in_dim=hidden_size,
|
||||
out_dim=hidden_size,
|
||||
mid_dim=int(hidden_size * mlp_ratio),
|
||||
use_bias=(True, True, False),
|
||||
norm=None,
|
||||
act=("silu", "silu", None),
|
||||
)
|
||||
elif ffn_type == "mlp":
|
||||
approx_gelu = lambda: nn.GELU(approximate="tanh")
|
||||
self.mlp = Mlp(
|
||||
in_features=hidden_size, hidden_features=int(hidden_size * mlp_ratio), act_layer=approx_gelu, drop=0
|
||||
)
|
||||
else:
|
||||
raise ValueError(f"{ffn_type} type is not defined.")
|
||||
self.drop_path = DropPath(drop_path) if drop_path > 0.0 else nn.Identity()
|
||||
self.scale_shift_table = nn.Parameter(torch.randn(6, hidden_size) / hidden_size**0.5)
|
||||
|
||||
def forward(self, x, y, t, mask=None, **kwargs):
|
||||
B, N, C = x.shape
|
||||
|
||||
shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = (
|
||||
self.scale_shift_table[None] + t.reshape(B, 6, -1)
|
||||
).chunk(6, dim=1)
|
||||
x = x + self.drop_path(gate_msa * self.attn(t2i_modulate(self.norm1(x), shift_msa, scale_msa)).reshape(B, N, C))
|
||||
x = x + self.cross_attn(x, y, mask)
|
||||
x = x + self.drop_path(gate_mlp * self.mlp(t2i_modulate(self.norm2(x), shift_mlp, scale_mlp)))
|
||||
|
||||
return x
|
||||
|
||||
|
||||
#############################################################################
|
||||
# Core Sana Model #
|
||||
#################################################################################
|
||||
class Sana(nn.Module):
|
||||
"""
|
||||
Diffusion model with a Transformer backbone.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
input_size=32,
|
||||
patch_size=1,
|
||||
in_channels=32,
|
||||
hidden_size=1152,
|
||||
depth=28,
|
||||
num_heads=36,
|
||||
mlp_ratio=2.5,
|
||||
class_dropout_prob=0.1,
|
||||
pred_sigma=False,
|
||||
drop_path: float = 0.0,
|
||||
caption_channels=2304,
|
||||
pe_interpolation=1.0,
|
||||
config=None,
|
||||
model_max_length=300,
|
||||
qk_norm=False,
|
||||
y_norm=False,
|
||||
norm_eps=1e-5,
|
||||
attn_type="flash",
|
||||
ffn_type="mlp",
|
||||
use_pe=False,
|
||||
y_norm_scale_factor=1.0,
|
||||
patch_embed_kernel=None,
|
||||
mlp_acts=("silu", "silu", None),
|
||||
linear_head_dim=32,
|
||||
**kwargs,
|
||||
):
|
||||
super().__init__()
|
||||
self.pred_sigma = pred_sigma
|
||||
self.in_channels = in_channels
|
||||
self.out_channels = in_channels * 2 if pred_sigma else in_channels
|
||||
self.patch_size = patch_size
|
||||
self.num_heads = num_heads
|
||||
self.pe_interpolation = pe_interpolation
|
||||
self.depth = depth
|
||||
self.use_pe = use_pe
|
||||
self.y_norm = y_norm
|
||||
self.model_max_length = model_max_length
|
||||
self.fp32_attention = kwargs.get("use_fp32_attention", False)
|
||||
|
||||
kernel_size = patch_embed_kernel or patch_size
|
||||
self.x_embedder = PatchEmbed(
|
||||
input_size, patch_size, in_channels, hidden_size, kernel_size=kernel_size, bias=True
|
||||
)
|
||||
self.t_embedder = TimestepEmbedder(hidden_size)
|
||||
num_patches = self.x_embedder.num_patches
|
||||
self.base_size = input_size // self.patch_size
|
||||
# Will use fixed sin-cos embedding:
|
||||
self.register_buffer("pos_embed", torch.zeros(1, num_patches, hidden_size))
|
||||
|
||||
approx_gelu = lambda: nn.GELU(approximate="tanh")
|
||||
self.t_block = nn.Sequential(nn.SiLU(), nn.Linear(hidden_size, 6 * hidden_size, bias=True))
|
||||
self.y_embedder = CaptionEmbedder(
|
||||
in_channels=caption_channels,
|
||||
hidden_size=hidden_size,
|
||||
uncond_prob=class_dropout_prob,
|
||||
act_layer=approx_gelu,
|
||||
token_num=model_max_length,
|
||||
)
|
||||
if self.y_norm:
|
||||
self.attention_y_norm = RMSNorm(hidden_size, scale_factor=y_norm_scale_factor, eps=norm_eps)
|
||||
drop_path = [x.item() for x in torch.linspace(0, drop_path, depth)] # stochastic depth decay rule
|
||||
self.blocks = nn.ModuleList(
|
||||
[
|
||||
SanaBlock(
|
||||
hidden_size,
|
||||
num_heads,
|
||||
mlp_ratio=mlp_ratio,
|
||||
drop_path=drop_path[i],
|
||||
input_size=(input_size // patch_size, input_size // patch_size),
|
||||
qk_norm=qk_norm,
|
||||
attn_type=attn_type,
|
||||
ffn_type=ffn_type,
|
||||
mlp_acts=mlp_acts,
|
||||
linear_head_dim=linear_head_dim,
|
||||
)
|
||||
for i in range(depth)
|
||||
]
|
||||
)
|
||||
self.final_layer = T2IFinalLayer(hidden_size, patch_size, self.out_channels)
|
||||
|
||||
def forward(self, x, timestep, y, mask=None, data_info=None, **kwargs):
|
||||
"""
|
||||
Forward pass of Sana.
|
||||
x: (N, C, H, W) tensor of spatial inputs (images or latent representations of images)
|
||||
t: (N,) tensor of diffusion timesteps
|
||||
y: (N, 1, 120, C) tensor of class labels
|
||||
"""
|
||||
x = x.to(self.dtype)
|
||||
timestep = timestep.to(self.dtype)
|
||||
y = y.to(self.dtype)
|
||||
pos_embed = self.pos_embed.to(self.dtype)
|
||||
self.h, self.w = x.shape[-2] // self.patch_size, x.shape[-1] // self.patch_size
|
||||
if self.use_pe:
|
||||
x = self.x_embedder(x) + pos_embed # (N, T, D), where T = H * W / patch_size ** 2
|
||||
else:
|
||||
x = self.x_embedder(x)
|
||||
t = self.t_embedder(timestep.to(x.dtype)) # (N, D)
|
||||
t0 = self.t_block(t)
|
||||
y = self.y_embedder(y, self.training) # (N, 1, L, D)
|
||||
if self.y_norm:
|
||||
y = self.attention_y_norm(y)
|
||||
if mask is not None:
|
||||
if mask.shape[0] != y.shape[0]:
|
||||
mask = mask.repeat(y.shape[0] // mask.shape[0], 1)
|
||||
mask = mask.squeeze(1).squeeze(1)
|
||||
y = y.squeeze(1).masked_select(mask.unsqueeze(-1) != 0).view(1, -1, x.shape[-1])
|
||||
y_lens = mask.sum(dim=1).tolist()
|
||||
else:
|
||||
y_lens = [y.shape[2]] * y.shape[0]
|
||||
y = y.squeeze(1).view(1, -1, x.shape[-1])
|
||||
for block in self.blocks:
|
||||
x = auto_grad_checkpoint(block, x, y, t0, y_lens) # (N, T, D) #support grad checkpoint
|
||||
x = self.final_layer(x, t) # (N, T, patch_size ** 2 * out_channels)
|
||||
x = self.unpatchify(x) # (N, out_channels, H, W)
|
||||
return x
|
||||
|
||||
def __call__(self, *args, **kwargs):
|
||||
"""
|
||||
This method allows the object to be called like a function.
|
||||
It simply calls the forward method.
|
||||
"""
|
||||
return self.forward(*args, **kwargs)
|
||||
|
||||
def forward_with_dpmsolver(self, x, timestep, y, mask=None, **kwargs):
|
||||
"""
|
||||
dpm solver donnot need variance prediction
|
||||
"""
|
||||
# https://github.com/openai/glide-text2im/blob/main/notebooks/text2im.ipynb
|
||||
model_out = self.forward(x, timestep, y, mask)
|
||||
return model_out.chunk(2, dim=1)[0] if self.pred_sigma else model_out
|
||||
|
||||
def unpatchify(self, x):
|
||||
"""
|
||||
x: (N, T, patch_size**2 * C)
|
||||
imgs: (N, H, W, C)
|
||||
"""
|
||||
c = self.out_channels
|
||||
p = self.x_embedder.patch_size[0]
|
||||
h = w = int(x.shape[1] ** 0.5)
|
||||
assert h * w == x.shape[1]
|
||||
|
||||
x = x.reshape(shape=(x.shape[0], h, w, p, p, c))
|
||||
x = torch.einsum("nhwpqc->nchpwq", x)
|
||||
imgs = x.reshape(shape=(x.shape[0], c, h * p, h * p))
|
||||
return imgs
|
||||
|
||||
def initialize_weights(self):
|
||||
# Initialize transformer layers:
|
||||
def _basic_init(module):
|
||||
if isinstance(module, nn.Linear):
|
||||
torch.nn.init.xavier_uniform_(module.weight)
|
||||
if module.bias is not None:
|
||||
nn.init.constant_(module.bias, 0)
|
||||
|
||||
self.apply(_basic_init)
|
||||
|
||||
if self.use_pe:
|
||||
# Initialize (and freeze) pos_embed by sin-cos embedding:
|
||||
pos_embed = get_2d_sincos_pos_embed(
|
||||
self.pos_embed.shape[-1],
|
||||
int(self.x_embedder.num_patches**0.5),
|
||||
pe_interpolation=self.pe_interpolation,
|
||||
base_size=self.base_size,
|
||||
)
|
||||
self.pos_embed.data.copy_(torch.from_numpy(pos_embed).float().unsqueeze(0))
|
||||
|
||||
# Initialize patch_embed like nn.Linear (instead of nn.Conv2d):
|
||||
w = self.x_embedder.proj.weight.data
|
||||
nn.init.xavier_uniform_(w.view([w.shape[0], -1]))
|
||||
|
||||
# Initialize timestep embedding MLP:
|
||||
nn.init.normal_(self.t_embedder.mlp[0].weight, std=0.02)
|
||||
nn.init.normal_(self.t_embedder.mlp[2].weight, std=0.02)
|
||||
nn.init.normal_(self.t_block[1].weight, std=0.02)
|
||||
|
||||
# Initialize caption embedding MLP:
|
||||
nn.init.normal_(self.y_embedder.y_proj.fc1.weight, std=0.02)
|
||||
nn.init.normal_(self.y_embedder.y_proj.fc2.weight, std=0.02)
|
||||
|
||||
|
||||
def get_2d_sincos_pos_embed(embed_dim, grid_size, cls_token=False, extra_tokens=0, pe_interpolation=1.0, base_size=16):
|
||||
"""
|
||||
grid_size: int of the grid height and width
|
||||
return:
|
||||
pos_embed: [grid_size*grid_size, embed_dim] or [1+grid_size*grid_size, embed_dim] (w/ or w/o cls_token)
|
||||
"""
|
||||
if isinstance(grid_size, int):
|
||||
grid_size = to_2tuple(grid_size)
|
||||
grid_h = np.arange(grid_size[0], dtype=np.float32) / (grid_size[0] / base_size) / pe_interpolation
|
||||
grid_w = np.arange(grid_size[1], dtype=np.float32) / (grid_size[1] / base_size) / pe_interpolation
|
||||
grid = np.meshgrid(grid_w, grid_h) # here w goes first
|
||||
grid = np.stack(grid, axis=0)
|
||||
grid = grid.reshape([2, 1, grid_size[1], grid_size[0]])
|
||||
|
||||
pos_embed = get_2d_sincos_pos_embed_from_grid(embed_dim, grid)
|
||||
if cls_token and extra_tokens > 0:
|
||||
pos_embed = np.concatenate([np.zeros([extra_tokens, embed_dim]), pos_embed], axis=0)
|
||||
return pos_embed
|
||||
|
||||
|
||||
def get_2d_sincos_pos_embed_from_grid(embed_dim, grid):
|
||||
assert embed_dim % 2 == 0
|
||||
|
||||
# use half of dimensions to encode grid_h
|
||||
emb_h = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[0]) # (H*W, D/2)
|
||||
emb_w = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[1]) # (H*W, D/2)
|
||||
|
||||
emb = np.concatenate([emb_h, emb_w], axis=1) # (H*W, D)
|
||||
return emb
|
||||
|
||||
|
||||
def get_1d_sincos_pos_embed_from_grid(embed_dim, pos):
|
||||
"""
|
||||
embed_dim: output dimension for each position
|
||||
pos: a list of positions to be encoded: size (M,)
|
||||
out: (M, D)
|
||||
"""
|
||||
assert embed_dim % 2 == 0
|
||||
omega = np.arange(embed_dim // 2, dtype=np.float64)
|
||||
omega /= embed_dim / 2.0
|
||||
omega = 1.0 / 10000**omega # (D/2,)
|
||||
|
||||
pos = pos.reshape(-1) # (M,)
|
||||
out = np.einsum("m,d->md", pos, omega) # (M, D/2), outer product
|
||||
|
||||
emb_sin = np.sin(out) # (M, D/2)
|
||||
emb_cos = np.cos(out) # (M, D/2)
|
||||
|
||||
emb = np.concatenate([emb_sin, emb_cos], axis=1) # (M, D)
|
||||
return emb
|
||||
@@ -0,0 +1,796 @@
|
||||
# Copyright 2024 NVIDIA CORPORATION & AFFILIATES
|
||||
#
|
||||
# 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.
|
||||
#
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
# This file is modified from https://github.com/PixArt-alpha/PixArt-sigma
|
||||
import math
|
||||
from typing import Optional
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
from einops import rearrange
|
||||
from timm.models.vision_transformer import Attention as Attention_
|
||||
from timm.models.vision_transformer import Mlp
|
||||
from transformers import AutoModelForCausalLM
|
||||
|
||||
from .norms import RMSNorm
|
||||
from .utils import get_same_padding, to_2tuple
|
||||
|
||||
sdpa_32b = None
|
||||
Q_4GB_LIMIT = 32000000
|
||||
"""If q is greater than this, the operation will likely require >4GB VRAM, which will fail on Intel Arc Alchemist GPUs without a workaround."""
|
||||
# 2k = 37 748 736
|
||||
# 1024 = 9 437 184
|
||||
# 2k model goes very slightly over 4GB
|
||||
|
||||
from comfy import model_management
|
||||
if model_management.xformers_enabled():
|
||||
import xformers
|
||||
import xformers.ops
|
||||
else:
|
||||
if model_management.xpu_available:
|
||||
import intel_extension_for_pytorch as ipex
|
||||
import os
|
||||
if not torch.xpu.has_fp64_dtype() and not os.environ.get('IPEX_FORCE_ATTENTION_SLICE', None):
|
||||
from ...utils.IPEX.attention import scaled_dot_product_attention_32_bit
|
||||
sdpa_32b = scaled_dot_product_attention_32_bit
|
||||
print("Using IPEX 4GB SDPA workaround")
|
||||
else:
|
||||
print("No IPEX 4GB workaround")
|
||||
|
||||
|
||||
def modulate(x, shift, scale):
|
||||
return x * (1 + scale.unsqueeze(1)) + shift.unsqueeze(1)
|
||||
|
||||
|
||||
def t2i_modulate(x, shift, scale):
|
||||
return x * (1 + scale) + shift
|
||||
|
||||
|
||||
class MultiHeadCrossAttention(nn.Module):
|
||||
def __init__(self, d_model, num_heads, attn_drop=0.0, proj_drop=0.0, qk_norm=False, **block_kwargs):
|
||||
super().__init__()
|
||||
assert d_model % num_heads == 0, "d_model must be divisible by num_heads"
|
||||
|
||||
self.d_model = d_model
|
||||
self.num_heads = num_heads
|
||||
self.head_dim = d_model // num_heads
|
||||
|
||||
self.q_linear = nn.Linear(d_model, d_model)
|
||||
self.kv_linear = nn.Linear(d_model, d_model * 2)
|
||||
self.attn_drop = nn.Dropout(attn_drop)
|
||||
self.proj = nn.Linear(d_model, d_model)
|
||||
self.proj_drop = nn.Dropout(proj_drop)
|
||||
if qk_norm:
|
||||
# not used for now
|
||||
self.q_norm = RMSNorm(d_model, scale_factor=1.0, eps=1e-6)
|
||||
self.k_norm = RMSNorm(d_model, scale_factor=1.0, eps=1e-6)
|
||||
else:
|
||||
self.q_norm = nn.Identity()
|
||||
self.k_norm = nn.Identity()
|
||||
|
||||
def forward(self, x, cond, mask=None):
|
||||
# query/value: img tokens; key: condition; mask: if padding tokens
|
||||
B, N, C = x.shape
|
||||
|
||||
q = self.q_linear(x).view(1, -1, self.num_heads, self.head_dim)
|
||||
kv = self.kv_linear(cond).view(1, -1, 2, self.num_heads, self.head_dim)
|
||||
k, v = kv.unbind(2)
|
||||
|
||||
if model_management.xformers_enabled():
|
||||
attn_bias = None
|
||||
if mask is not None:
|
||||
attn_bias = xformers.ops.fmha.BlockDiagonalMask.from_seqlens([N] * B, mask)
|
||||
x = xformers.ops.memory_efficient_attention(
|
||||
q, k, v,
|
||||
p=self.attn_drop.p,
|
||||
attn_bias=attn_bias
|
||||
)
|
||||
else:
|
||||
q, k, v = map(lambda t: t.permute(0, 2, 1, 3),(q, k, v),)
|
||||
attn_mask = None
|
||||
if mask is not None and len(mask) > 1:
|
||||
|
||||
# Create equivalent of xformer diagonal block mask, still only correct for square masks
|
||||
# But depth doesn't matter as tensors can expand in that dimension
|
||||
attn_mask_template = torch.ones(
|
||||
[q.shape[2] // B, mask[0]],
|
||||
dtype=torch.bool,
|
||||
device=q.device
|
||||
)
|
||||
attn_mask = torch.block_diag(attn_mask_template)
|
||||
|
||||
# create a mask on the diagonal for each mask in the batch
|
||||
for n in range(B - 1):
|
||||
attn_mask = torch.block_diag(attn_mask, attn_mask_template)
|
||||
|
||||
p = getattr(self.attn_drop, "p", 0) # IPEX.optimize() will turn attn_drop into an Identity()
|
||||
|
||||
if sdpa_32b is not None and (q.element_size() * q.nelement()) > Q_4GB_LIMIT:
|
||||
sdpa = sdpa_32b
|
||||
else:
|
||||
sdpa = torch.nn.functional.scaled_dot_product_attention
|
||||
|
||||
x = sdpa(
|
||||
q, k, v,
|
||||
attn_mask=attn_mask,
|
||||
dropout_p=p
|
||||
).permute(0, 2, 1, 3).contiguous()
|
||||
x = x.view(B, -1, C)
|
||||
x = self.proj(x)
|
||||
x = self.proj_drop(x)
|
||||
return x
|
||||
|
||||
|
||||
class LiteLA(Attention_):
|
||||
r"""Lightweight linear attention"""
|
||||
|
||||
PAD_VAL = 1
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
in_dim: int,
|
||||
out_dim: int,
|
||||
heads: Optional[int] = None,
|
||||
heads_ratio: float = 1.0,
|
||||
dim=32,
|
||||
eps=1e-15,
|
||||
use_bias=False,
|
||||
qk_norm=False,
|
||||
norm_eps=1e-5,
|
||||
):
|
||||
heads = heads or int(out_dim // dim * heads_ratio)
|
||||
super().__init__(in_dim, num_heads=heads, qkv_bias=use_bias)
|
||||
|
||||
self.in_dim = in_dim
|
||||
self.out_dim = out_dim
|
||||
self.heads = heads
|
||||
self.dim = out_dim // heads # TODO: need some change
|
||||
self.eps = eps
|
||||
|
||||
self.kernel_func = nn.ReLU(inplace=False)
|
||||
if qk_norm:
|
||||
self.q_norm = RMSNorm(in_dim, scale_factor=1.0, eps=norm_eps)
|
||||
self.k_norm = RMSNorm(in_dim, scale_factor=1.0, eps=norm_eps)
|
||||
else:
|
||||
self.q_norm = nn.Identity()
|
||||
self.k_norm = nn.Identity()
|
||||
|
||||
def attn_matmul(self, q, k, v: torch.Tensor) -> torch.Tensor:
|
||||
# lightweight linear attention
|
||||
q = self.kernel_func(q) # B, h, h_d, N
|
||||
k = self.kernel_func(k)
|
||||
|
||||
q, k, v = q.float(), k.float(), v.float()
|
||||
|
||||
v = F.pad(v, (0, 0, 0, 1), mode="constant", value=LiteLA.PAD_VAL)
|
||||
vk = torch.matmul(v, k)
|
||||
out = torch.matmul(vk, q)
|
||||
|
||||
if out.dtype in [torch.float16, torch.bfloat16]:
|
||||
out = out.float()
|
||||
out = out[:, :, :-1] / (out[:, :, -1:] + self.eps)
|
||||
|
||||
return out
|
||||
|
||||
def forward(self, x: torch.Tensor, mask=None, HW=None, block_id=None) -> torch.Tensor:
|
||||
B, N, C = x.shape
|
||||
|
||||
qkv = self.qkv(x).reshape(B, N, 3, C)
|
||||
q, k, v = qkv.unbind(2) # B, N, 3, C --> B, N, C
|
||||
dtype = q.dtype
|
||||
|
||||
q = self.q_norm(q).transpose(-1, -2) # (B, N, C) -> (B, C, N)
|
||||
k = self.k_norm(k).transpose(-1, -2) # (B, N, C) -> (B, C, N)
|
||||
v = v.transpose(-1, -2)
|
||||
|
||||
q = q.reshape(B, C // self.dim, self.dim, N) # (B, h, h_d, N)
|
||||
k = k.reshape(B, C // self.dim, self.dim, N).transpose(-1, -2) # (B, h, N, h_d)
|
||||
v = v.reshape(B, C // self.dim, self.dim, N) # (B, h, h_d, N)
|
||||
|
||||
out = self.attn_matmul(q, k, v).to(dtype)
|
||||
|
||||
out = out.view(B, C, N).permute(0, 2, 1) # B, N, C
|
||||
out = self.proj(out)
|
||||
|
||||
if torch.get_autocast_gpu_dtype() == torch.float16:
|
||||
out = out.clip(-65504, 65504)
|
||||
|
||||
return out
|
||||
|
||||
@property
|
||||
def module_str(self) -> str:
|
||||
_str = type(self).__name__ + "("
|
||||
eps = f"{self.eps:.1E}"
|
||||
_str += f"i={self.in_dim},o={self.out_dim},h={self.heads},d={self.dim},eps={eps}"
|
||||
return _str
|
||||
|
||||
def __repr__(self):
|
||||
return f"EPS{self.eps}-" + super().__repr__()
|
||||
|
||||
|
||||
class PAGCFGIdentitySelfAttnProcessorLiteLA:
|
||||
r"""Self Attention with Perturbed Attention & CFG Guidance"""
|
||||
|
||||
def __init__(self, attn):
|
||||
self.attn = attn
|
||||
|
||||
def __call__(self, x: torch.Tensor, mask=None, HW=None, block_id=None) -> torch.Tensor:
|
||||
x_uncond, x_org, x_ptb = x.chunk(3)
|
||||
x_org = torch.cat([x_uncond, x_org])
|
||||
B, N, C = x_org.shape
|
||||
|
||||
qkv = self.attn.qkv(x_org).reshape(B, N, 3, C)
|
||||
# B, N, 3, C --> B, N, C
|
||||
q, k, v = qkv.unbind(2)
|
||||
dtype = q.dtype
|
||||
q = self.attn.q_norm(q).transpose(-1, -2) # (B, N, C) -> (B, C, N)
|
||||
k = self.attn.k_norm(k).transpose(-1, -2) # (B, N, C) -> (B, C, N)
|
||||
v = v.transpose(-1, -2)
|
||||
|
||||
q = q.reshape(B, C // self.attn.dim, self.attn.dim, N) # (B, h, h_d, N)
|
||||
k = k.reshape(B, C // self.attn.dim, self.attn.dim, N).transpose(-1, -2) # (B, h, N, h_d)
|
||||
v = v.reshape(B, C // self.attn.dim, self.attn.dim, N) # (B, h, h_d, N)
|
||||
|
||||
out = self.attn.attn_matmul(q, k, v).to(dtype)
|
||||
|
||||
out = out.view(B, C, N).permute(0, 2, 1) # B, N, C
|
||||
out = self.attn.proj(out)
|
||||
|
||||
# perturbed path (identity attention)
|
||||
v_weight = self.attn.qkv.weight[C * 2 : C * 3, :] # Shape: (dim, dim)
|
||||
if self.attn.qkv.bias:
|
||||
v_bias = self.attn.qkv.bias[C * 2 : C * 3] # Shape: (dim,)
|
||||
x_ptb = (torch.matmul(x_ptb, v_weight.t()) + v_bias).to(dtype)
|
||||
else:
|
||||
x_ptb = torch.matmul(x_ptb, v_weight.t()).to(dtype)
|
||||
x_ptb = self.attn.proj(x_ptb)
|
||||
|
||||
out = torch.cat([out, x_ptb])
|
||||
|
||||
if torch.get_autocast_gpu_dtype() == torch.float16:
|
||||
out = out.clip(-65504, 65504)
|
||||
|
||||
return out
|
||||
|
||||
|
||||
class PAGIdentitySelfAttnProcessorLiteLA:
|
||||
r"""Self Attention with Perturbed Attention Guidance"""
|
||||
|
||||
def __init__(self, attn):
|
||||
self.attn = attn
|
||||
|
||||
def __call__(self, x: torch.Tensor, mask=None, HW=None, block_id=None) -> torch.Tensor:
|
||||
x_org, x_ptb = x.chunk(2)
|
||||
B, N, C = x_org.shape
|
||||
|
||||
qkv = self.attn.qkv(x_org).reshape(B, N, 3, C)
|
||||
# B, N, 3, C --> B, N, C
|
||||
q, k, v = qkv.unbind(2)
|
||||
dtype = q.dtype
|
||||
q = self.attn.q_norm(q).transpose(-1, -2) # (B, N, C) -> (B, C, N)
|
||||
k = self.attn.k_norm(k).transpose(-1, -2) # (B, N, C) -> (B, C, N)
|
||||
v = v.transpose(-1, -2)
|
||||
|
||||
q = q.reshape(B, C // self.attn.dim, self.attn.dim, N) # (B, h, h_d, N)
|
||||
k = k.reshape(B, C // self.attn.dim, self.attn.dim, N).transpose(-1, -2) # (B, h, N, h_d)
|
||||
v = v.reshape(B, C // self.attn.dim, self.attn.dim, N) # (B, h, h_d, N)
|
||||
|
||||
out = self.attn.attn_matmul(q, k, v).to(dtype)
|
||||
|
||||
out = out.view(B, C, N).permute(0, 2, 1) # B, N, C
|
||||
out = self.attn.proj(out)
|
||||
|
||||
# perturbed path (identity attention)
|
||||
v_weight = self.attn.qkv.weight[C * 2 : C * 3, :] # Shape: (dim, dim)
|
||||
if self.attn.qkv.bias:
|
||||
v_bias = self.attn.qkv.bias[C * 2 : C * 3] # Shape: (dim,)
|
||||
x_ptb = (torch.matmul(x_ptb, v_weight.t()) + v_bias).to(dtype)
|
||||
else:
|
||||
x_ptb = torch.matmul(x_ptb, v_weight.t()).to(dtype)
|
||||
x_ptb = self.attn.proj(x_ptb)
|
||||
|
||||
out = torch.cat([out, x_ptb])
|
||||
|
||||
if torch.get_autocast_gpu_dtype() == torch.float16:
|
||||
out = out.clip(-65504, 65504)
|
||||
|
||||
return out
|
||||
|
||||
|
||||
class SelfAttnProcessorLiteLA:
|
||||
r"""Self Attention with Lite Linear Attention"""
|
||||
|
||||
def __init__(self, attn):
|
||||
self.attn = attn
|
||||
|
||||
def __call__(self, x: torch.Tensor, mask=None, HW=None, block_id=None) -> torch.Tensor:
|
||||
B, N, C = x.shape
|
||||
if HW is None:
|
||||
H = W = int(N**0.5)
|
||||
else:
|
||||
H, W = HW
|
||||
qkv = self.attn.qkv(x).reshape(B, N, 3, C)
|
||||
# B, N, 3, C --> B, N, C
|
||||
q, k, v = qkv.unbind(2)
|
||||
dtype = q.dtype
|
||||
q = self.attn.q_norm(q).transpose(-1, -2) # (B, N, C) -> (B, C, N)
|
||||
k = self.attn.k_norm(k).transpose(-1, -2) # (B, N, C) -> (B, C, N)
|
||||
v = v.transpose(-1, -2)
|
||||
|
||||
q = q.reshape(B, C // self.attn.dim, self.attn.dim, N) # (B, h, h_d, N)
|
||||
k = k.reshape(B, C // self.attn.dim, self.attn.dim, N).transpose(-1, -2) # (B, h, N, h_d)
|
||||
v = v.reshape(B, C // self.attn.dim, self.attn.dim, N) # (B, h, h_d, N)
|
||||
|
||||
out = self.attn.attn_matmul(q, k, v).to(dtype)
|
||||
|
||||
out = out.view(B, C, N).permute(0, 2, 1) # B, N, C
|
||||
out = self.attn.proj(out)
|
||||
|
||||
if torch.get_autocast_gpu_dtype() == torch.float16:
|
||||
out = out.clip(-65504, 65504)
|
||||
|
||||
return out
|
||||
|
||||
|
||||
class FlashAttention(Attention_):
|
||||
"""Multi-head Flash Attention block with qk norm."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
dim,
|
||||
num_heads=8,
|
||||
qkv_bias=True,
|
||||
qk_norm=False,
|
||||
**block_kwargs,
|
||||
):
|
||||
"""
|
||||
Args:
|
||||
dim (int): Number of input channels.
|
||||
num_heads (int): Number of attention heads.
|
||||
qkv_bias (bool: If True, add a learnable bias to query, key, value.
|
||||
"""
|
||||
super().__init__(dim, num_heads=num_heads, qkv_bias=qkv_bias, **block_kwargs)
|
||||
|
||||
if qk_norm:
|
||||
self.q_norm = nn.LayerNorm(dim)
|
||||
self.k_norm = nn.LayerNorm(dim)
|
||||
else:
|
||||
self.q_norm = nn.Identity()
|
||||
self.k_norm = nn.Identity()
|
||||
|
||||
def forward(self, x, mask=None, HW=None, block_id=None):
|
||||
B, N, C = x.shape
|
||||
|
||||
qkv = self.qkv(x).reshape(B, N, 3, C)
|
||||
q, k, v = qkv.unbind(2)
|
||||
dtype = q.dtype
|
||||
|
||||
q = self.q_norm(q)
|
||||
k = self.k_norm(k)
|
||||
|
||||
q = q.reshape(B, N, self.num_heads, C // self.num_heads).to(dtype)
|
||||
k = k.reshape(B, N, self.num_heads, C // self.num_heads).to(dtype)
|
||||
v = v.reshape(B, N, self.num_heads, C // self.num_heads).to(dtype)
|
||||
|
||||
use_fp32_attention = getattr(self, "fp32_attention", False) # necessary for NAN loss
|
||||
if use_fp32_attention:
|
||||
q, k, v = q.float(), k.float(), v.float()
|
||||
|
||||
attn_bias = None
|
||||
if mask is not None:
|
||||
attn_bias = torch.zeros([B * self.num_heads, q.shape[1], k.shape[1]], dtype=q.dtype, device=q.device)
|
||||
attn_bias.masked_fill_(mask.squeeze(1).repeat(self.num_heads, 1, 1) == 0, float("-inf"))
|
||||
|
||||
if _xformers_available:
|
||||
x = xformers.ops.memory_efficient_attention(q, k, v, p=self.attn_drop.p, attn_bias=attn_bias)
|
||||
else:
|
||||
q, k, v = q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2)
|
||||
if mask is not None and mask.ndim == 2:
|
||||
mask = (1 - mask.to(x.dtype)) * -10000.0
|
||||
mask = mask[:, None, None].repeat(1, self.num_heads, 1, 1)
|
||||
x = F.scaled_dot_product_attention(q, k, v, attn_mask=mask, dropout_p=0.0, is_causal=False)
|
||||
x = x.transpose(1, 2)
|
||||
|
||||
x = x.view(B, N, C)
|
||||
x = self.proj(x)
|
||||
x = self.proj_drop(x)
|
||||
|
||||
if torch.get_autocast_gpu_dtype() == torch.float16:
|
||||
x = x.clip(-65504, 65504)
|
||||
|
||||
return x
|
||||
|
||||
|
||||
#################################################################################
|
||||
# AMP attention with fp32 softmax to fix loss NaN problem during training #
|
||||
#################################################################################
|
||||
class Attention(Attention_):
|
||||
def forward(self, x, HW=None):
|
||||
B, N, C = x.shape
|
||||
qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)
|
||||
# B,N,3,H,C -> B,H,N,C
|
||||
q, k, v = qkv.unbind(0) # make torchscript happy (cannot use tensor as tuple)
|
||||
use_fp32_attention = getattr(self, "fp32_attention", False)
|
||||
if use_fp32_attention:
|
||||
q, k = q.float(), k.float()
|
||||
|
||||
with torch.cuda.amp.autocast(enabled=not use_fp32_attention):
|
||||
attn = (q @ k.transpose(-2, -1)) * self.scale
|
||||
attn = attn.softmax(dim=-1)
|
||||
|
||||
attn = self.attn_drop(attn)
|
||||
|
||||
x = (attn @ v).transpose(1, 2).reshape(B, N, C)
|
||||
x = self.proj(x)
|
||||
x = self.proj_drop(x)
|
||||
return x
|
||||
|
||||
|
||||
class FinalLayer(nn.Module):
|
||||
"""
|
||||
The final layer of Sana.
|
||||
"""
|
||||
|
||||
def __init__(self, hidden_size, patch_size, out_channels):
|
||||
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, c):
|
||||
shift, scale = self.adaLN_modulation(c).chunk(2, dim=1)
|
||||
x = modulate(self.norm_final(x), shift, scale)
|
||||
x = self.linear(x)
|
||||
return x
|
||||
|
||||
|
||||
class T2IFinalLayer(nn.Module):
|
||||
"""
|
||||
The final layer of Sana.
|
||||
"""
|
||||
|
||||
def __init__(self, hidden_size, patch_size, out_channels):
|
||||
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.scale_shift_table = nn.Parameter(torch.randn(2, hidden_size) / hidden_size**0.5)
|
||||
self.out_channels = out_channels
|
||||
|
||||
def forward(self, x, t):
|
||||
shift, scale = (self.scale_shift_table[None] + t[:, None]).chunk(2, dim=1)
|
||||
x = t2i_modulate(self.norm_final(x), shift, scale)
|
||||
x = self.linear(x)
|
||||
return x
|
||||
|
||||
|
||||
class MaskFinalLayer(nn.Module):
|
||||
"""
|
||||
The final layer of Sana.
|
||||
"""
|
||||
|
||||
def __init__(self, final_hidden_size, c_emb_size, patch_size, out_channels):
|
||||
super().__init__()
|
||||
self.norm_final = nn.LayerNorm(final_hidden_size, elementwise_affine=False, eps=1e-6)
|
||||
self.linear = nn.Linear(final_hidden_size, patch_size * patch_size * out_channels, bias=True)
|
||||
self.adaLN_modulation = nn.Sequential(nn.SiLU(), nn.Linear(c_emb_size, 2 * final_hidden_size, bias=True))
|
||||
|
||||
def forward(self, x, t):
|
||||
shift, scale = self.adaLN_modulation(t).chunk(2, dim=1)
|
||||
x = modulate(self.norm_final(x), shift, scale)
|
||||
x = self.linear(x)
|
||||
return x
|
||||
|
||||
|
||||
class DecoderLayer(nn.Module):
|
||||
"""
|
||||
The final layer of Sana.
|
||||
"""
|
||||
|
||||
def __init__(self, hidden_size, decoder_hidden_size):
|
||||
super().__init__()
|
||||
self.norm_decoder = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
|
||||
self.linear = nn.Linear(hidden_size, decoder_hidden_size, bias=True)
|
||||
self.adaLN_modulation = nn.Sequential(nn.SiLU(), nn.Linear(hidden_size, 2 * hidden_size, bias=True))
|
||||
|
||||
def forward(self, x, t):
|
||||
shift, scale = self.adaLN_modulation(t).chunk(2, dim=1)
|
||||
x = modulate(self.norm_decoder(x), shift, scale)
|
||||
x = self.linear(x)
|
||||
return x
|
||||
|
||||
|
||||
#################################################################################
|
||||
# Embedding Layers for Timesteps and Class Labels #
|
||||
#################################################################################
|
||||
class TimestepEmbedder(nn.Module):
|
||||
"""
|
||||
Embeds scalar timesteps into vector representations.
|
||||
"""
|
||||
|
||||
def __init__(self, hidden_size, frequency_embedding_size=256):
|
||||
super().__init__()
|
||||
self.mlp = nn.Sequential(
|
||||
nn.Linear(frequency_embedding_size, hidden_size, bias=True),
|
||||
nn.SiLU(),
|
||||
nn.Linear(hidden_size, hidden_size, bias=True),
|
||||
)
|
||||
self.frequency_embedding_size = frequency_embedding_size
|
||||
|
||||
@staticmethod
|
||||
def timestep_embedding(t, dim, max_period=10000):
|
||||
"""
|
||||
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.
|
||||
"""
|
||||
# https://github.com/openai/glide-text2im/blob/main/glide_text2im/nn.py
|
||||
half = dim // 2
|
||||
freqs = torch.exp(
|
||||
-math.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32, device=t.device) / half
|
||||
)
|
||||
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)
|
||||
return embedding
|
||||
|
||||
def forward(self, t):
|
||||
t_freq = self.timestep_embedding(t, self.frequency_embedding_size).to(self.dtype)
|
||||
t_emb = self.mlp(t_freq)
|
||||
return t_emb
|
||||
|
||||
@property
|
||||
def dtype(self):
|
||||
try:
|
||||
return next(self.parameters()).dtype
|
||||
except StopIteration:
|
||||
return torch.float32
|
||||
|
||||
|
||||
class SizeEmbedder(TimestepEmbedder):
|
||||
"""
|
||||
Embeds scalar timesteps into vector representations.
|
||||
"""
|
||||
|
||||
def __init__(self, hidden_size, frequency_embedding_size=256):
|
||||
super().__init__(hidden_size=hidden_size, frequency_embedding_size=frequency_embedding_size)
|
||||
self.mlp = nn.Sequential(
|
||||
nn.Linear(frequency_embedding_size, hidden_size, bias=True),
|
||||
nn.SiLU(),
|
||||
nn.Linear(hidden_size, hidden_size, bias=True),
|
||||
)
|
||||
self.frequency_embedding_size = frequency_embedding_size
|
||||
self.outdim = hidden_size
|
||||
|
||||
def forward(self, s, bs):
|
||||
if s.ndim == 1:
|
||||
s = s[:, None]
|
||||
assert s.ndim == 2
|
||||
if s.shape[0] != bs:
|
||||
s = s.repeat(bs // s.shape[0], 1)
|
||||
assert s.shape[0] == bs
|
||||
b, dims = s.shape[0], s.shape[1]
|
||||
s = rearrange(s, "b d -> (b d)")
|
||||
s_freq = self.timestep_embedding(s, self.frequency_embedding_size).to(self.dtype)
|
||||
s_emb = self.mlp(s_freq)
|
||||
s_emb = rearrange(s_emb, "(b d) d2 -> b (d d2)", b=b, d=dims, d2=self.outdim)
|
||||
return s_emb
|
||||
|
||||
@property
|
||||
def dtype(self):
|
||||
try:
|
||||
return next(self.parameters()).dtype
|
||||
except StopIteration:
|
||||
return torch.float32
|
||||
|
||||
|
||||
class LabelEmbedder(nn.Module):
|
||||
"""
|
||||
Embeds class labels into vector representations. Also handles label dropout for classifier-free guidance.
|
||||
"""
|
||||
|
||||
def __init__(self, num_classes, hidden_size, dropout_prob):
|
||||
super().__init__()
|
||||
use_cfg_embedding = dropout_prob > 0
|
||||
self.embedding_table = nn.Embedding(num_classes + use_cfg_embedding, hidden_size)
|
||||
self.num_classes = num_classes
|
||||
self.dropout_prob = dropout_prob
|
||||
|
||||
def token_drop(self, labels, force_drop_ids=None):
|
||||
"""
|
||||
Drops labels to enable classifier-free guidance.
|
||||
"""
|
||||
if force_drop_ids is None:
|
||||
drop_ids = torch.rand(labels.shape[0]).cuda() < self.dropout_prob
|
||||
else:
|
||||
drop_ids = force_drop_ids == 1
|
||||
labels = torch.where(drop_ids, self.num_classes, labels)
|
||||
return labels
|
||||
|
||||
def forward(self, labels, train, force_drop_ids=None):
|
||||
use_dropout = self.dropout_prob > 0
|
||||
if (train and use_dropout) or (force_drop_ids is not None):
|
||||
labels = self.token_drop(labels, force_drop_ids)
|
||||
embeddings = self.embedding_table(labels)
|
||||
return embeddings
|
||||
|
||||
|
||||
class CaptionEmbedder(nn.Module):
|
||||
"""
|
||||
Embeds class labels into vector representations. Also handles label dropout for classifier-free guidance.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
in_channels,
|
||||
hidden_size,
|
||||
uncond_prob,
|
||||
act_layer=nn.GELU(approximate="tanh"),
|
||||
token_num=120,
|
||||
):
|
||||
super().__init__()
|
||||
self.y_proj = Mlp(
|
||||
in_features=in_channels, hidden_features=hidden_size, out_features=hidden_size, act_layer=act_layer, drop=0
|
||||
)
|
||||
self.register_buffer("y_embedding", nn.Parameter(torch.randn(token_num, in_channels) / in_channels**0.5))
|
||||
self.uncond_prob = uncond_prob
|
||||
|
||||
def initialize_gemma_params(self, model_name="google/gemma-2b-it"):
|
||||
num_layers = len(self.custom_gemma_layers)
|
||||
text_encoder = AutoModelForCausalLM.from_pretrained(model_name).get_decoder()
|
||||
pretrained_layers = text_encoder.layers[-num_layers:]
|
||||
for custom_layer, pretrained_layer in zip(self.custom_gemma_layers, pretrained_layers):
|
||||
info = custom_layer.load_state_dict(pretrained_layer.state_dict(), strict=False)
|
||||
print(f"**** {info} ****")
|
||||
print(f"**** Initialized {num_layers} Gemma layers from pretrained model: {model_name} ****")
|
||||
|
||||
def token_drop(self, caption, force_drop_ids=None):
|
||||
"""
|
||||
Drops labels to enable classifier-free guidance.
|
||||
"""
|
||||
if force_drop_ids is None:
|
||||
drop_ids = torch.rand(caption.shape[0]).cuda() < self.uncond_prob
|
||||
else:
|
||||
drop_ids = force_drop_ids == 1
|
||||
caption = torch.where(drop_ids[:, None, None, None], self.y_embedding, caption)
|
||||
return caption
|
||||
|
||||
def forward(self, caption, train, force_drop_ids=None, mask=None):
|
||||
if train:
|
||||
assert caption.shape[2:] == self.y_embedding.shape
|
||||
use_dropout = self.uncond_prob > 0
|
||||
if (train and use_dropout) or (force_drop_ids is not None):
|
||||
caption = self.token_drop(caption, force_drop_ids)
|
||||
|
||||
caption = self.y_proj(caption)
|
||||
|
||||
return caption
|
||||
|
||||
|
||||
class CaptionEmbedderDoubleBr(nn.Module):
|
||||
"""
|
||||
Embeds class labels into vector representations. Also handles label dropout for classifier-free guidance.
|
||||
"""
|
||||
|
||||
def __init__(self, in_channels, hidden_size, uncond_prob, act_layer=nn.GELU(approximate="tanh"), token_num=120):
|
||||
super().__init__()
|
||||
self.proj = Mlp(
|
||||
in_features=in_channels, hidden_features=hidden_size, out_features=hidden_size, act_layer=act_layer, drop=0
|
||||
)
|
||||
self.embedding = nn.Parameter(torch.randn(1, in_channels) / 10**0.5)
|
||||
self.y_embedding = nn.Parameter(torch.randn(token_num, in_channels) / 10**0.5)
|
||||
self.uncond_prob = uncond_prob
|
||||
|
||||
def token_drop(self, global_caption, caption, force_drop_ids=None):
|
||||
"""
|
||||
Drops labels to enable classifier-free guidance.
|
||||
"""
|
||||
if force_drop_ids is None:
|
||||
drop_ids = torch.rand(global_caption.shape[0]).cuda() < self.uncond_prob
|
||||
else:
|
||||
drop_ids = force_drop_ids == 1
|
||||
global_caption = torch.where(drop_ids[:, None], self.embedding, global_caption)
|
||||
caption = torch.where(drop_ids[:, None, None, None], self.y_embedding, caption)
|
||||
return global_caption, caption
|
||||
|
||||
def forward(self, caption, train, force_drop_ids=None):
|
||||
assert caption.shape[2:] == self.y_embedding.shape
|
||||
global_caption = caption.mean(dim=2).squeeze()
|
||||
use_dropout = self.uncond_prob > 0
|
||||
if (train and use_dropout) or (force_drop_ids is not None):
|
||||
global_caption, caption = self.token_drop(global_caption, caption, force_drop_ids)
|
||||
y_embed = self.proj(global_caption)
|
||||
return y_embed, caption
|
||||
|
||||
|
||||
class PatchEmbed(nn.Module):
|
||||
"""2D Image to Patch Embedding"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
img_size=224,
|
||||
patch_size=16,
|
||||
in_chans=3,
|
||||
embed_dim=768,
|
||||
kernel_size=None,
|
||||
padding=0,
|
||||
norm_layer=None,
|
||||
flatten=True,
|
||||
bias=True,
|
||||
):
|
||||
super().__init__()
|
||||
kernel_size = kernel_size or patch_size
|
||||
img_size = to_2tuple(img_size)
|
||||
patch_size = to_2tuple(patch_size)
|
||||
self.img_size = img_size
|
||||
self.patch_size = patch_size
|
||||
self.grid_size = (img_size[0] // patch_size[0], img_size[1] // patch_size[1])
|
||||
self.num_patches = self.grid_size[0] * self.grid_size[1]
|
||||
self.flatten = flatten
|
||||
if not padding and kernel_size % 2 > 0:
|
||||
padding = get_same_padding(kernel_size)
|
||||
self.proj = nn.Conv2d(
|
||||
in_chans, embed_dim, kernel_size=kernel_size, stride=patch_size, padding=padding, bias=bias
|
||||
)
|
||||
self.norm = norm_layer(embed_dim) if norm_layer else nn.Identity()
|
||||
|
||||
def forward(self, x):
|
||||
B, C, H, W = x.shape
|
||||
assert (H == self.img_size[0], f"Input image height ({H}) doesn't match model ({self.img_size[0]}).")
|
||||
assert (W == self.img_size[1], f"Input image width ({W}) doesn't match model ({self.img_size[1]}).")
|
||||
x = self.proj(x)
|
||||
if self.flatten:
|
||||
x = x.flatten(2).transpose(1, 2) # BCHW -> BNC
|
||||
x = self.norm(x)
|
||||
return x
|
||||
|
||||
|
||||
class PatchEmbedMS(nn.Module):
|
||||
"""2D Image to Patch Embedding"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
patch_size=16,
|
||||
in_chans=3,
|
||||
embed_dim=768,
|
||||
kernel_size=None,
|
||||
padding=0,
|
||||
norm_layer=None,
|
||||
flatten=True,
|
||||
bias=True,
|
||||
):
|
||||
super().__init__()
|
||||
kernel_size = kernel_size or patch_size
|
||||
patch_size = to_2tuple(patch_size)
|
||||
self.patch_size = patch_size
|
||||
self.flatten = flatten
|
||||
if not padding and kernel_size % 2 > 0:
|
||||
padding = get_same_padding(kernel_size)
|
||||
self.proj = nn.Conv2d(
|
||||
in_chans, embed_dim, kernel_size=kernel_size, stride=patch_size, padding=padding, bias=bias
|
||||
)
|
||||
self.norm = norm_layer(embed_dim) if norm_layer else nn.Identity()
|
||||
|
||||
def forward(self, x):
|
||||
x = self.proj(x)
|
||||
if self.flatten:
|
||||
x = x.flatten(2).transpose(1, 2) # BCHW -> BNC
|
||||
x = self.norm(x)
|
||||
return x
|
||||
@@ -0,0 +1,373 @@
|
||||
# Copyright 2024 NVIDIA CORPORATION & AFFILIATES
|
||||
#
|
||||
# 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.
|
||||
#
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
# This file is modified from https://github.com/PixArt-alpha/PixArt-sigma
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from timm.models.layers import DropPath
|
||||
|
||||
from .basic_modules import DWMlp, GLUMBConv, MBConvPreGLU, Mlp
|
||||
from .sana import Sana, get_2d_sincos_pos_embed
|
||||
from .sana_blocks import (
|
||||
Attention,
|
||||
CaptionEmbedder,
|
||||
FlashAttention,
|
||||
LiteLA,
|
||||
MultiHeadCrossAttention,
|
||||
PatchEmbedMS,
|
||||
T2IFinalLayer,
|
||||
t2i_modulate,
|
||||
)
|
||||
from .utils import auto_grad_checkpoint
|
||||
|
||||
|
||||
class SanaMSBlock(nn.Module):
|
||||
"""
|
||||
A Sana block with global shared adaptive layer norm zero (adaLN-Zero) conditioning.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
hidden_size,
|
||||
num_heads,
|
||||
mlp_ratio=4.0,
|
||||
drop_path=0.0,
|
||||
input_size=None,
|
||||
qk_norm=False,
|
||||
attn_type="flash",
|
||||
ffn_type="mlp",
|
||||
mlp_acts=("silu", "silu", None),
|
||||
linear_head_dim=32,
|
||||
cross_norm=False,
|
||||
**block_kwargs,
|
||||
):
|
||||
super().__init__()
|
||||
self.hidden_size = hidden_size
|
||||
self.norm1 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
|
||||
if attn_type == "flash":
|
||||
# flash self attention
|
||||
self.attn = FlashAttention(
|
||||
hidden_size,
|
||||
num_heads=num_heads,
|
||||
qkv_bias=True,
|
||||
qk_norm=qk_norm,
|
||||
**block_kwargs,
|
||||
)
|
||||
elif attn_type == "linear":
|
||||
# linear self attention
|
||||
# TODO: Here the num_heads set to 36 for tmp used
|
||||
self_num_heads = hidden_size // linear_head_dim
|
||||
self.attn = LiteLA(hidden_size, hidden_size, heads=self_num_heads, eps=1e-8, qk_norm=qk_norm)
|
||||
elif attn_type == "vanilla":
|
||||
# vanilla self attention
|
||||
self.attn = Attention(hidden_size, num_heads=num_heads, qkv_bias=True)
|
||||
else:
|
||||
raise ValueError(f"{attn_type} type is not defined.")
|
||||
|
||||
self.cross_attn = MultiHeadCrossAttention(hidden_size, num_heads, qk_norm=cross_norm, **block_kwargs)
|
||||
self.norm2 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
|
||||
if ffn_type == "dwmlp":
|
||||
approx_gelu = lambda: nn.GELU(approximate="tanh")
|
||||
self.mlp = DWMlp(
|
||||
in_features=hidden_size, hidden_features=int(hidden_size * mlp_ratio), act_layer=approx_gelu, drop=0
|
||||
)
|
||||
elif ffn_type == "glumbconv":
|
||||
self.mlp = GLUMBConv(
|
||||
in_features=hidden_size,
|
||||
hidden_features=int(hidden_size * mlp_ratio),
|
||||
use_bias=(True, True, False),
|
||||
norm=(None, None, None),
|
||||
act=mlp_acts,
|
||||
)
|
||||
elif ffn_type == "glumbconv_dilate":
|
||||
self.mlp = GLUMBConv(
|
||||
in_features=hidden_size,
|
||||
hidden_features=int(hidden_size * mlp_ratio),
|
||||
use_bias=(True, True, False),
|
||||
norm=(None, None, None),
|
||||
act=mlp_acts,
|
||||
dilation=2,
|
||||
)
|
||||
elif ffn_type == "mlp":
|
||||
approx_gelu = lambda: nn.GELU(approximate="tanh")
|
||||
self.mlp = Mlp(
|
||||
in_features=hidden_size, hidden_features=int(hidden_size * mlp_ratio), act_layer=approx_gelu, drop=0
|
||||
)
|
||||
elif ffn_type == "mbconvpreglu":
|
||||
self.mlp = MBConvPreGLU(
|
||||
in_dim=hidden_size,
|
||||
out_dim=hidden_size,
|
||||
mid_dim=int(hidden_size * mlp_ratio),
|
||||
use_bias=(True, True, False),
|
||||
norm=None,
|
||||
act=mlp_acts,
|
||||
)
|
||||
else:
|
||||
raise ValueError(f"{ffn_type} type is not defined.")
|
||||
self.drop_path = DropPath(drop_path) if drop_path > 0.0 else nn.Identity()
|
||||
self.scale_shift_table = nn.Parameter(torch.randn(6, hidden_size) / hidden_size**0.5)
|
||||
|
||||
def forward(self, x, y, t, mask=None, HW=None, **kwargs):
|
||||
B, N, C = x.shape
|
||||
|
||||
shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = (
|
||||
self.scale_shift_table[None] + t.reshape(B, 6, -1)
|
||||
).chunk(6, dim=1)
|
||||
x = x + self.drop_path(gate_msa * self.attn(t2i_modulate(self.norm1(x), shift_msa, scale_msa), HW=HW))
|
||||
x = x + self.cross_attn(x, y, mask)
|
||||
x = x + self.drop_path(gate_mlp * self.mlp(t2i_modulate(self.norm2(x), shift_mlp, scale_mlp), HW=HW))
|
||||
|
||||
return x
|
||||
|
||||
|
||||
#############################################################################
|
||||
# Core Sana Model #
|
||||
#################################################################################
|
||||
class SanaMS(Sana):
|
||||
"""
|
||||
Diffusion model with a Transformer backbone.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
input_size=32,
|
||||
patch_size=2,
|
||||
in_channels=32,
|
||||
hidden_size=1152,
|
||||
depth=28,
|
||||
num_heads=16,
|
||||
mlp_ratio=4.0,
|
||||
class_dropout_prob=0.1,
|
||||
learn_sigma=False,
|
||||
pred_sigma=False,
|
||||
drop_path: float = 0.0,
|
||||
caption_channels=2304,
|
||||
pe_interpolation=1.0,
|
||||
config=None,
|
||||
model_max_length=300,
|
||||
qk_norm=False,
|
||||
y_norm=False,
|
||||
norm_eps=1e-5,
|
||||
attn_type="linear",
|
||||
ffn_type="glumbconv",
|
||||
use_pe=False,
|
||||
y_norm_scale_factor=1.0,
|
||||
patch_embed_kernel=None,
|
||||
mlp_acts=("silu", "silu", None),
|
||||
linear_head_dim=32,
|
||||
cross_norm=False,
|
||||
**kwargs,
|
||||
):
|
||||
super().__init__(
|
||||
input_size=input_size,
|
||||
patch_size=patch_size,
|
||||
in_channels=in_channels,
|
||||
hidden_size=hidden_size,
|
||||
depth=depth,
|
||||
num_heads=num_heads,
|
||||
mlp_ratio=mlp_ratio,
|
||||
class_dropout_prob=class_dropout_prob,
|
||||
learn_sigma=learn_sigma,
|
||||
pred_sigma=pred_sigma,
|
||||
drop_path=drop_path,
|
||||
caption_channels=caption_channels,
|
||||
pe_interpolation=pe_interpolation,
|
||||
config=config,
|
||||
model_max_length=model_max_length,
|
||||
qk_norm=qk_norm,
|
||||
y_norm=y_norm,
|
||||
norm_eps=norm_eps,
|
||||
attn_type=attn_type,
|
||||
ffn_type=ffn_type,
|
||||
use_pe=use_pe,
|
||||
y_norm_scale_factor=y_norm_scale_factor,
|
||||
patch_embed_kernel=patch_embed_kernel,
|
||||
mlp_acts=mlp_acts,
|
||||
linear_head_dim=linear_head_dim,
|
||||
**kwargs,
|
||||
)
|
||||
self.dtype = torch.get_default_dtype()
|
||||
self.h = self.w = 0
|
||||
approx_gelu = lambda: nn.GELU(approximate="tanh")
|
||||
self.t_block = nn.Sequential(nn.SiLU(), nn.Linear(hidden_size, 6 * hidden_size, bias=True))
|
||||
self.pos_embed_ms = None
|
||||
|
||||
kernel_size = patch_embed_kernel or patch_size
|
||||
self.x_embedder = PatchEmbedMS(patch_size, in_channels, hidden_size, kernel_size=kernel_size, bias=True)
|
||||
self.y_embedder = CaptionEmbedder(
|
||||
in_channels=caption_channels,
|
||||
hidden_size=hidden_size,
|
||||
uncond_prob=class_dropout_prob,
|
||||
act_layer=approx_gelu,
|
||||
token_num=model_max_length,
|
||||
)
|
||||
drop_path = [x.item() for x in torch.linspace(0, drop_path, depth)] # stochastic depth decay rule
|
||||
self.blocks = nn.ModuleList(
|
||||
[
|
||||
SanaMSBlock(
|
||||
hidden_size,
|
||||
num_heads,
|
||||
mlp_ratio=mlp_ratio,
|
||||
drop_path=drop_path[i],
|
||||
input_size=(input_size // patch_size, input_size // patch_size),
|
||||
qk_norm=qk_norm,
|
||||
attn_type=attn_type,
|
||||
ffn_type=ffn_type,
|
||||
mlp_acts=mlp_acts,
|
||||
linear_head_dim=linear_head_dim,
|
||||
cross_norm=cross_norm,
|
||||
)
|
||||
for i in range(depth)
|
||||
]
|
||||
)
|
||||
self.final_layer = T2IFinalLayer(hidden_size, patch_size, self.out_channels)
|
||||
|
||||
self.initialize()
|
||||
|
||||
def forward(self, x, timesteps, context, **kwargs):
|
||||
"""
|
||||
Forward pass that adapts comfy input to original forward function
|
||||
x: (N, C, H, W) tensor of spatial inputs (images or latent representations of images)
|
||||
timesteps: (N,) tensor of diffusion timesteps
|
||||
context: (N, 1, 120, C) conditioning
|
||||
"""
|
||||
## size/ar from cond with fallback based on the latent image shape.
|
||||
bs = x.shape[0]
|
||||
## Still accepts the input w/o that dim but returns garbage
|
||||
if len(context.shape) == 3:
|
||||
context = context.unsqueeze(1)
|
||||
|
||||
## run original forward pass
|
||||
out = self.forward_raw(
|
||||
x = x.to(self.dtype),
|
||||
timestep = timesteps.to(self.dtype),
|
||||
y = context.to(self.dtype),
|
||||
)
|
||||
|
||||
## only return EPS
|
||||
out = out.to(torch.float)
|
||||
|
||||
return out
|
||||
|
||||
def forward_raw(self, x, timestep, y, mask=None, data_info=None, **kwargs):
|
||||
"""
|
||||
Forward pass of Sana.
|
||||
x: (N, C, H, W) tensor of spatial inputs (images or latent representations of images)
|
||||
t: (N,) tensor of diffusion timesteps
|
||||
y: (N, 1, 120, C) tensor of class labels
|
||||
"""
|
||||
bs = x.shape[0]
|
||||
x = x.to(self.dtype)
|
||||
timestep = timestep.to(self.dtype)
|
||||
y = y.to(self.dtype)
|
||||
self.h, self.w = x.shape[-2] // self.patch_size, x.shape[-1] // self.patch_size
|
||||
if self.use_pe:
|
||||
x = self.x_embedder(x)
|
||||
if self.pos_embed_ms is None or self.pos_embed_ms.shape[1:] != x.shape[1:]:
|
||||
self.pos_embed_ms = (
|
||||
torch.from_numpy(
|
||||
get_2d_sincos_pos_embed(
|
||||
self.pos_embed.shape[-1],
|
||||
(self.h, self.w),
|
||||
pe_interpolation=self.pe_interpolation,
|
||||
base_size=self.base_size,
|
||||
)
|
||||
)
|
||||
.unsqueeze(0)
|
||||
.to(x.device)
|
||||
.to(self.dtype)
|
||||
)
|
||||
x += self.pos_embed_ms # (N, T, D), where T = H * W / patch_size ** 2
|
||||
else:
|
||||
x = self.x_embedder(x)
|
||||
|
||||
t = self.t_embedder(timestep) # (N, D)
|
||||
|
||||
y_lens = ((y != 0).sum(dim=3) > 0).sum(dim=2).squeeze().tolist()
|
||||
y_lens = [y_lens[1]] * bs
|
||||
|
||||
mask = torch.zeros((len(y_lens), self.model_max_length), dtype=torch.int).to(x.device)
|
||||
for i, count in enumerate(y_lens):
|
||||
mask[i, :count] = 1
|
||||
|
||||
t0 = self.t_block(t)
|
||||
y = self.y_embedder(y, self.training, mask=mask) # (N, D)
|
||||
if self.y_norm:
|
||||
y = self.attention_y_norm(y)
|
||||
|
||||
y = y.squeeze(1).masked_select(mask.unsqueeze(-1).bool()).view(1, -1, y.shape[-1])
|
||||
|
||||
for block in self.blocks:
|
||||
x = auto_grad_checkpoint(
|
||||
block, x, y, t0, y_lens, (self.h, self.w), **kwargs
|
||||
) # (N, T, D) #support grad checkpoint
|
||||
|
||||
x = self.final_layer(x, t) # (N, T, patch_size ** 2 * out_channels)
|
||||
x = self.unpatchify(x) # (N, out_channels, H, W)
|
||||
|
||||
return x
|
||||
|
||||
def __call__(self, *args, **kwargs):
|
||||
"""
|
||||
This method allows the object to be called like a function.
|
||||
It simply calls the forward method.
|
||||
"""
|
||||
return self.forward(*args, **kwargs)
|
||||
|
||||
def forward_with_dpmsolver(self, x, timestep, y, data_info, **kwargs):
|
||||
"""
|
||||
dpm solver donnot need variance prediction
|
||||
"""
|
||||
# https://github.com/openai/glide-text2im/blob/main/notebooks/text2im.ipynb
|
||||
model_out = self.forward(x, timestep, y, data_info=data_info, **kwargs)
|
||||
return model_out.chunk(2, dim=1)[0] if self.pred_sigma else model_out
|
||||
|
||||
def unpatchify(self, x):
|
||||
"""
|
||||
x: (N, T, patch_size**2 * C)
|
||||
imgs: (N, H, W, C)
|
||||
"""
|
||||
c = self.out_channels
|
||||
p = self.x_embedder.patch_size[0]
|
||||
assert self.h * self.w == x.shape[1]
|
||||
|
||||
x = x.reshape(shape=(x.shape[0], self.h, self.w, p, p, c))
|
||||
x = torch.einsum("nhwpqc->nchpwq", x)
|
||||
imgs = x.reshape(shape=(x.shape[0], c, self.h * p, self.w * p))
|
||||
return imgs
|
||||
|
||||
def initialize(self):
|
||||
# Initialize transformer layers:
|
||||
def _basic_init(module):
|
||||
if isinstance(module, nn.Linear):
|
||||
torch.nn.init.xavier_uniform_(module.weight)
|
||||
if module.bias is not None:
|
||||
nn.init.constant_(module.bias, 0)
|
||||
|
||||
self.apply(_basic_init)
|
||||
|
||||
# Initialize patch_embed like nn.Linear (instead of nn.Conv2d):
|
||||
w = self.x_embedder.proj.weight.data
|
||||
nn.init.xavier_uniform_(w.view([w.shape[0], -1]))
|
||||
|
||||
# Initialize timestep embedding MLP:
|
||||
nn.init.normal_(self.t_embedder.mlp[0].weight, std=0.02)
|
||||
nn.init.normal_(self.t_embedder.mlp[2].weight, std=0.02)
|
||||
nn.init.normal_(self.t_block[1].weight, std=0.02)
|
||||
|
||||
# Initialize caption embedding MLP:
|
||||
nn.init.normal_(self.y_embedder.y_proj.fc1.weight, std=0.02)
|
||||
nn.init.normal_(self.y_embedder.y_proj.fc2.weight, std=0.02)
|
||||
@@ -0,0 +1,96 @@
|
||||
# Copyright 2024 NVIDIA CORPORATION & AFFILIATES
|
||||
#
|
||||
# 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.
|
||||
#
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
from collections.abc import Iterable
|
||||
from itertools import repeat
|
||||
from typing import Union, Tuple
|
||||
|
||||
import torch
|
||||
from torch.utils.checkpoint import checkpoint, checkpoint_sequential
|
||||
|
||||
|
||||
def _ntuple(n):
|
||||
def parse(x):
|
||||
if isinstance(x, Iterable) and not isinstance(x, str):
|
||||
return x
|
||||
return tuple(repeat(x, n))
|
||||
|
||||
return parse
|
||||
|
||||
|
||||
to_1tuple = _ntuple(1)
|
||||
to_2tuple = _ntuple(2)
|
||||
|
||||
|
||||
def auto_grad_checkpoint(module, *args, **kwargs):
|
||||
if getattr(module, "grad_checkpointing", False):
|
||||
if isinstance(module, Iterable):
|
||||
gc_step = module[0].grad_checkpointing_step
|
||||
return checkpoint_sequential(module, gc_step, *args, **kwargs)
|
||||
else:
|
||||
return checkpoint(module, *args, **kwargs)
|
||||
return module(*args, **kwargs)
|
||||
|
||||
|
||||
def checkpoint_sequential(functions, step, input, *args, **kwargs):
|
||||
|
||||
# Hack for keyword-only parameter in a python 2.7-compliant way
|
||||
preserve = kwargs.pop("preserve_rng_state", True)
|
||||
if kwargs:
|
||||
raise ValueError("Unexpected keyword arguments: " + ",".join(arg for arg in kwargs))
|
||||
|
||||
def run_function(start, end, functions):
|
||||
def forward(input):
|
||||
for j in range(start, end + 1):
|
||||
input = functions[j](input, *args)
|
||||
return input
|
||||
|
||||
return forward
|
||||
|
||||
if isinstance(functions, torch.nn.Sequential):
|
||||
functions = list(functions.children())
|
||||
|
||||
# the last chunk has to be non-volatile
|
||||
end = -1
|
||||
segment = len(functions) // step
|
||||
for start in range(0, step * (segment - 1), step):
|
||||
end = start + step - 1
|
||||
input = checkpoint(run_function(start, end, functions), input, preserve_rng_state=preserve)
|
||||
return run_function(end + 1, len(functions) - 1, functions)(input)
|
||||
|
||||
def val2list(x: list or tuple or any, repeat_time=1) -> list: # type: ignore
|
||||
"""Repeat `val` for `repeat_time` times and return the list or val if list/tuple."""
|
||||
if isinstance(x, (list, tuple)):
|
||||
return list(x)
|
||||
return [x for _ in range(repeat_time)]
|
||||
|
||||
def val2tuple(x: list or tuple or any, min_len: int = 1, idx_repeat: int = -1) -> tuple: # type: ignore
|
||||
"""Return tuple with min_len by repeating element at idx_repeat."""
|
||||
# convert to list first
|
||||
x = val2list(x)
|
||||
|
||||
# repeat elements if necessary
|
||||
if len(x) > 0:
|
||||
x[idx_repeat:idx_repeat] = [x[idx_repeat] for _ in range(min_len - len(x))]
|
||||
|
||||
return tuple(x)
|
||||
|
||||
def get_same_padding(kernel_size: Union[int, Tuple[int, ...]]) -> Union[int, Tuple[int, ...]]:
|
||||
if isinstance(kernel_size, tuple):
|
||||
return tuple([get_same_padding(ks) for ks in kernel_size])
|
||||
else:
|
||||
assert kernel_size % 2 > 0, f"kernel size {kernel_size} should be odd number"
|
||||
return kernel_size // 2
|
||||
+148
@@ -0,0 +1,148 @@
|
||||
import torch
|
||||
import folder_paths
|
||||
|
||||
from .conf import sana_conf, sana_res
|
||||
from .loader import load_sana
|
||||
|
||||
dtypes = [
|
||||
"auto",
|
||||
"FP32",
|
||||
"FP16",
|
||||
"BF16"
|
||||
]
|
||||
|
||||
class SanaCheckpointLoader:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"ckpt_name": (folder_paths.get_filename_list("checkpoints"),),
|
||||
"model": (list(sana_conf.keys()),),
|
||||
}
|
||||
}
|
||||
RETURN_TYPES = ("MODEL",)
|
||||
RETURN_NAMES = ("model",)
|
||||
FUNCTION = "load_checkpoint"
|
||||
CATEGORY = "ExtraModels/Sana"
|
||||
TITLE = "Sana Checkpoint Loader"
|
||||
|
||||
def load_checkpoint(self, ckpt_name, model):
|
||||
ckpt_path = folder_paths.get_full_path("checkpoints", ckpt_name)
|
||||
model_conf = sana_conf[model]
|
||||
model = load_sana(
|
||||
model_path = ckpt_path,
|
||||
model_conf = model_conf,
|
||||
)
|
||||
return (model,)
|
||||
|
||||
|
||||
class SanaResolutionSelect():
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"model": (list(sana_res.keys()),),
|
||||
"ratio": (list(sana_res["1024px"].keys()),{"default":"1.00"}),
|
||||
}
|
||||
}
|
||||
RETURN_TYPES = ("INT","INT")
|
||||
RETURN_NAMES = ("width","height")
|
||||
FUNCTION = "get_res"
|
||||
CATEGORY = "ExtraModels/Sana"
|
||||
TITLE = "Sana Resolution Select"
|
||||
|
||||
def get_res(self, model, ratio):
|
||||
width, height = sana_res[model][ratio]
|
||||
return (width,height)
|
||||
|
||||
|
||||
class SanaResolutionCond:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"cond": ("CONDITIONING", ),
|
||||
"width": ("INT", {"default": 1024.0, "min": 0, "max": 8192}),
|
||||
"height": ("INT", {"default": 1024.0, "min": 0, "max": 8192}),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("CONDITIONING",)
|
||||
RETURN_NAMES = ("cond",)
|
||||
FUNCTION = "add_cond"
|
||||
CATEGORY = "ExtraModels/Sana"
|
||||
TITLE = "Sana Resolution Conditioning"
|
||||
|
||||
def add_cond(self, cond, width, height):
|
||||
for c in range(len(cond)):
|
||||
cond[c][1].update({
|
||||
"img_hw": [[height, width]],
|
||||
"aspect_ratio": [[height/width]],
|
||||
})
|
||||
return (cond,)
|
||||
|
||||
|
||||
class SanaTextEncode:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"text": ("STRING", {"multiline": True}),
|
||||
"GEMMA": ("GEMMA",),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("CONDITIONING",)
|
||||
FUNCTION = "encode"
|
||||
CATEGORY = "ExtraModels/Sana"
|
||||
TITLE = "Sana Text Encode"
|
||||
|
||||
def encode(self, text, GEMMA=None):
|
||||
tokenizer = GEMMA["tokenizer"]
|
||||
text_encoder = GEMMA["text_encoder"]
|
||||
|
||||
with torch.no_grad():
|
||||
chi_prompt = "\n".join(preset_te_prompt)
|
||||
full_prompt = chi_prompt + text
|
||||
num_chi_tokens = len(tokenizer.encode(chi_prompt))
|
||||
max_length = num_chi_tokens + 300 - 2
|
||||
|
||||
tokens = tokenizer(
|
||||
[full_prompt],
|
||||
max_length=max_length,
|
||||
padding="max_length",
|
||||
truncation=True,
|
||||
return_tensors="pt"
|
||||
).to(text_encoder.device)
|
||||
|
||||
select_idx = [0] + list(range(-300 + 1, 0))
|
||||
embs = text_encoder(tokens.input_ids, tokens.attention_mask)[0][:, None][:, :, select_idx]
|
||||
emb_masks = tokens.attention_mask[:, select_idx]
|
||||
embs = embs * emb_masks.unsqueeze(-1)
|
||||
|
||||
return ([[embs, {}]], )
|
||||
|
||||
preset_te_prompt = [
|
||||
'Given a user prompt, generate an "Enhanced prompt" that provides detailed visual descriptions suitable for image generation. Evaluate the level of detail in the user prompt:',
|
||||
'- If the prompt is simple, focus on adding specifics about colors, shapes, sizes, textures, and spatial relationships to create vivid and concrete scenes.',
|
||||
'- If the prompt is already detailed, refine and enhance the existing details slightly without overcomplicating.',
|
||||
'Here are examples of how to transform or refine prompts:',
|
||||
'- User Prompt: A cat sleeping -> Enhanced: A small, fluffy white cat curled up in a round shape, sleeping peacefully on a warm sunny windowsill, surrounded by pots of blooming red flowers.',
|
||||
'- User Prompt: A busy city street -> Enhanced: A bustling city street scene at dusk, featuring glowing street lamps, a diverse crowd of people in colorful clothing, and a double-decker bus passing by towering glass skyscrapers.',
|
||||
'Please generate only the enhanced description for the prompt below and avoid including any additional commentary or evaluations:',
|
||||
'User Prompt: '
|
||||
]
|
||||
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"SanaCheckpointLoader" : SanaCheckpointLoader,
|
||||
"SanaResolutionSelect" : SanaResolutionSelect,
|
||||
"SanaTextEncode" : SanaTextEncode,
|
||||
"SanaResolutionCond" : SanaResolutionCond,
|
||||
}
|
||||
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"Sana Checkpoint Loader": "SanaCheckpointLoader",
|
||||
"Sana Resolution Select": "SanaResolutionSelect",
|
||||
"Sana Text Encoder": "SanaTextEncode",
|
||||
"Sana Resolution Cond": "SanaResolutionCond",
|
||||
}
|
||||
+20
@@ -155,5 +155,25 @@ vae_conf = {
|
||||
"ch_mult" : [1, 2, 2, 4],
|
||||
"num_res_blocks" : 2,
|
||||
"attn_resolutions" : [32],
|
||||
},
|
||||
# DCAE configs
|
||||
"dcae-f32c32-sana-1.0": {
|
||||
"type" : "DCAE",
|
||||
"in_channels" : 3,
|
||||
"embed_scale" : 32,
|
||||
"embed_dim" : 32,
|
||||
"encoder_block_type" : ["ResBlock", "ResBlock", "ResBlock", "EViTS5GLU", "EViTS5GLU", "EViTS5GLU"],
|
||||
"encoder_width_list" : [128, 256, 512, 512, 1024, 1024],
|
||||
"encoder_depth_list" : [2, 2, 2, 3, 3, 3],
|
||||
"encoder_norm" : "rms2d",
|
||||
"encoder_act" : "silu",
|
||||
"downsample_block_type" : "Conv",
|
||||
"decoder_block_type" : ["ResBlock", "ResBlock", "ResBlock", "EViTS5GLU", "EViTS5GLU", "EViTS5GLU"],
|
||||
"decoder_width_list" : [128, 256, 512, 512, 1024, 1024],
|
||||
"decoder_depth_list" : [3, 3, 3, 3, 3, 3],
|
||||
"decoder_norm" : "rms2d",
|
||||
"decoder_act" : "silu",
|
||||
"upsample_block_type" : "InterpolateConv",
|
||||
"scaling_factor" : 0.41407
|
||||
}
|
||||
}
|
||||
|
||||
@@ -32,6 +32,9 @@ class EXVAE(comfy.sd.VAE):
|
||||
elif model_conf["type"] == "MoVQ3":
|
||||
from .models.movq3 import MoVQ
|
||||
model = MoVQ(model_conf)
|
||||
elif model_conf["type"] == "DCAE":
|
||||
from .models.dcae import DCAE
|
||||
model = DCAE(**model_conf)
|
||||
else:
|
||||
raise NotImplementedError(f"Unknown VAE type '{model_conf['type']}'")
|
||||
|
||||
|
||||
+1018
File diff suppressed because it is too large
Load Diff
@@ -1,4 +1,6 @@
|
||||
import folder_paths
|
||||
import torch
|
||||
import comfy
|
||||
|
||||
from .conf import vae_conf
|
||||
from .loader import EXVAE
|
||||
@@ -12,6 +14,8 @@ dtypes = [
|
||||
"BF16"
|
||||
]
|
||||
|
||||
MAX_RESOLUTION=16384
|
||||
|
||||
class ExtraVAELoader:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
@@ -33,6 +37,34 @@ class ExtraVAELoader:
|
||||
vae = EXVAE(model_path, model_conf, string_to_dtype(dtype, "vae"))
|
||||
return (vae,)
|
||||
|
||||
|
||||
class EmptyDCAELatentImage:
|
||||
def __init__(self):
|
||||
self.device = comfy.model_management.intermediate_device()
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"width": ("INT", {"default": 512, "min": 16, "max": MAX_RESOLUTION, "step": 8, "tooltip": "The width of the latent images in pixels."}),
|
||||
"height": ("INT", {"default": 512, "min": 16, "max": MAX_RESOLUTION, "step": 8, "tooltip": "The height of the latent images in pixels."}),
|
||||
"batch_size": ("INT", {"default": 1, "min": 1, "max": 4096, "tooltip": "The number of latent images in the batch."})
|
||||
}
|
||||
}
|
||||
RETURN_TYPES = ("LATENT",)
|
||||
OUTPUT_TOOLTIPS = ("The empty latent image batch.",)
|
||||
FUNCTION = "generate"
|
||||
TITLE = "Empty DCAE Latent Image"
|
||||
|
||||
CATEGORY = "latent"
|
||||
DESCRIPTION = "Create a new batch of empty latent images to be denoised via sampling."
|
||||
|
||||
def generate(self, width, height, batch_size=1):
|
||||
latent = torch.zeros([batch_size, 32, height // 32, width // 32], device=self.device)
|
||||
return ({"samples":latent}, )
|
||||
|
||||
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"ExtraVAELoader" : ExtraVAELoader,
|
||||
"EmptyDCAELatentImage" : EmptyDCAELatentImage,
|
||||
}
|
||||
|
||||
@@ -38,5 +38,14 @@ else:
|
||||
from .utils.nodes import NODE_CLASS_MAPPINGS as Extra_Nodes
|
||||
NODE_CLASS_MAPPINGS.update(Extra_Nodes)
|
||||
|
||||
# Sana
|
||||
from .Sana.nodes import NODE_CLASS_MAPPINGS as Sana_Nodes
|
||||
NODE_CLASS_MAPPINGS.update(Sana_Nodes)
|
||||
|
||||
# Gemma
|
||||
from .Gemma.nodes import NODE_CLASS_MAPPINGS as Gemma_Nodes
|
||||
NODE_CLASS_MAPPINGS.update(Gemma_Nodes)
|
||||
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {k:v.TITLE for k,v in NODE_CLASS_MAPPINGS.items()}
|
||||
__all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS']
|
||||
|
||||
|
||||
Reference in New Issue
Block a user