Update model_loading.py

This commit is contained in:
kijai 2024-11-19 20:27:31 +02:00
parent 1cfe0835f5
commit 41a0f33381

View File

@ -425,7 +425,7 @@ class DownloadAndLoadCogVideoGGUFModel:
},
"optional": {
"block_edit": ("TRANSFORMERBLOCKS", {"default": None}),
"compile_args":("COMPILEARGS", ),
#"compile_args":("COMPILEARGS", ),
"attention_mode": (["sdpa", "sageattn"], {"default": "sdpa"}),
}
}
@ -523,12 +523,6 @@ class DownloadAndLoadCogVideoGGUFModel:
from .fp8_optimization import convert_fp8_linear
convert_fp8_linear(transformer, vae_dtype, params_to_keep=params_to_keep)
if compile_args is not None:
torch._dynamo.config.cache_size_limit = compile_args["dynamo_cache_size_limit"]
for i, block in enumerate(transformer.transformer_blocks):
if "CogVideoXBlock" in str(block):
transformer.transformer_blocks[i] = torch.compile(block, fullgraph=compile_args["fullgraph"], dynamic=compile_args["dynamic"], backend=compile_args["backend"], mode=compile_args["mode"])
with open(scheduler_path) as f:
scheduler_config = json.load(f)