xinyun/ComfyUI-CogVideoXWrapper
1
0
Fork 0
mirror of https://git.datalinker.icu/kijai/ComfyUI-CogVideoXWrapper.git synced 2025-12-10 05:14:22 +08:00
Code Issues Packages Projects Releases Wiki Activity

allow fused loras with fp8

Browse Source
This commit is contained in:
kijai 2024-11-11 23:36:48 +02:00
parent ea0273c8ec
commit 00fde5ebce
1 changed files with 15 additions and 15 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

30
model_loading.py
View File

@ -202,21 +202,6 @@ class DownloadAndLoadCogVideoModel:
if block_edit is not None: if block_edit is not None:
transformer = remove_specific_blocks(transformer, block_edit) transformer = remove_specific_blocks(transformer, block_edit)
#fp8
if fp8_transformer == "enabled" or fp8_transformer == "fastmode":
params_to_keep = {"patch_embed", "lora", "pos_embedding", "time_embedding"}
if "1.5" in model:
params_to_keep.update({"norm1.linear.weight", "norm_k", "norm_q","ofs_embedding", "norm_final", "norm_out", "proj_out"})
for name, param in transformer.named_parameters():
if not any(keyword in name for keyword in params_to_keep):
param.data = param.data.to(torch.float8_e4m3fn)
if fp8_transformer == "fastmode":
from .fp8_optimization import convert_fp8_linear
if "1.5" in model:
params_to_keep.update({"ff"}) #otherwise NaNs
convert_fp8_linear(transformer, dtype, params_to_keep=params_to_keep)
with open(scheduler_path) as f: with open(scheduler_path) as f:
scheduler_config = json.load(f) scheduler_config = json.load(f)
scheduler = CogVideoXDDIMScheduler.from_config(scheduler_config) scheduler = CogVideoXDDIMScheduler.from_config(scheduler_config)
@ -268,6 +253,21 @@ class DownloadAndLoadCogVideoModel:
lora_scale = lora_scale / lora_rank lora_scale = lora_scale / lora_rank
pipe.fuse_lora(lora_scale=lora_scale, components=["transformer"]) pipe.fuse_lora(lora_scale=lora_scale, components=["transformer"])
#fp8
if fp8_transformer == "enabled" or fp8_transformer == "fastmode":
params_to_keep = {"patch_embed", "lora", "pos_embedding", "time_embedding"}
if "1.5" in model:
params_to_keep.update({"norm1.linear.weight", "norm_k", "norm_q","ofs_embedding", "norm_final", "norm_out", "proj_out"})
for name, param in pipe.transformer.named_parameters():
if not any(keyword in name for keyword in params_to_keep):
param.data = param.data.to(torch.float8_e4m3fn)
if fp8_transformer == "fastmode":
from .fp8_optimization import convert_fp8_linear
if "1.5" in model:
params_to_keep.update({"ff"}) #otherwise NaNs
convert_fp8_linear(pipe.transformer, dtype, params_to_keep=params_to_keep)
if enable_sequential_cpu_offload: if enable_sequential_cpu_offload:
pipe.enable_sequential_cpu_offload() pipe.enable_sequential_cpu_offload()