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

Update model_optimization_nodes.py

Browse Source
This commit is contained in:
kijai 2025-03-13 11:10:05 +02:00
parent 7ecb190ef9
commit a5b39a8f1a
1 changed files with 7 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

9
nodes/model_optimization_nodes.py
View File

@ -337,7 +337,11 @@ class TorchCompileModelFluxAdvanced:
"double_blocks": ("STRING", {"default": "0-18", "multiline": True}),
"single_blocks": ("STRING", {"default": "0-37", "multiline": True}),
"dynamic": ("BOOLEAN", {"default": False, "tooltip": "Enable dynamic mode"}),
}}
},
"optional": {
"dynamo_cache_size_limit": ("INT", {"default": 64, "min": 0, "max": 1024, "step": 1, "tooltip": "torch._dynamo.config.cache_size_limit"}),
}
}
RETURN_TYPES = ("MODEL",)
FUNCTION = "patch"
@ -355,11 +359,12 @@ class TorchCompileModelFluxAdvanced:
blocks.append(int(part))
return blocks
def patch(self, model, backend, mode, fullgraph, single_blocks, double_blocks, dynamic):
def patch(self, model, backend, mode, fullgraph, single_blocks, double_blocks, dynamic, dynamo_cache_size_limit):
single_block_list = self.parse_blocks(single_blocks)
double_block_list = self.parse_blocks(double_blocks)
m = model.clone()
diffusion_model = m.get_model_object("diffusion_model")
torch._dynamo.config.cache_size_limit = dynamo_cache_size_limit
if not self._compiled:
try: