shenglin.net/ComfyUI-MochiWrapper-2
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Allow settin dynamo_cache_size_limit

Browse Source
This commit is contained in:
kijai 2024-11-06 23:38:07 +02:00
parent 4ef7df00c9
commit d3287d61b7
4 changed files with 8 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
mochi_preview/dit/joint_model/asymm_models_joint.py
View File

@ -165,6 +165,7 @@ class AsymmetricAttention(nn.Module):
raise ImportError("Flash RMSNorm not available.")
elif rms_norm_func == "apex":
from apex.normalization import FusedRMSNorm as ApexRMSNorm
@torch.compiler.disable()
class RMSNorm(ApexRMSNorm):
pass
else:
@ -237,7 +238,6 @@ class AsymmetricAttention(nn.Module):
skip_reshape=True
)
return out
def run_attention(
self,
q,

2
mochi_preview/dit/joint_model/layers.py
View File

@ -140,7 +140,7 @@ class PatchEmbed(nn.Module):
x = self.norm(x)
return x
@torch.compiler.disable()
class RMSNorm(torch.nn.Module):
def __init__(self, hidden_size, eps=1e-5, device=None):
super().__init__()

3
mochi_preview/t2v_synth_mochi.py
View File

@ -34,6 +34,7 @@ except:
import torch
import torch.utils.data
import torch._dynamo
from tqdm import tqdm
from comfy.utils import ProgressBar, load_torch_file
@ -161,6 +162,8 @@ class T2VSynthMochiModel:
#torch.compile
if compile_args is not None:
torch._dynamo.config.cache_size_limit = compile_args["dynamo_cache_size_limit"]
log.info(f"Set dynamo cache size limit to {torch._dynamo.config.cache_size_limit}")
if compile_args["compile_dit"]:
for i, block in enumerate(model.blocks):
model.blocks[i] = torch.compile(block, fullgraph=compile_args["fullgraph"], dynamic=compile_args["dynamic"], backend=compile_args["backend"])

4
nodes.py
View File

@ -247,6 +247,7 @@ class MochiTorchCompileSettings:
"compile_dit": ("BOOLEAN", {"default": True, "tooltip": "Compiles all transformer blocks"}),
"compile_final_layer": ("BOOLEAN", {"default": True, "tooltip": "Enable compiling final layer."}),
"dynamic": ("BOOLEAN", {"default": False, "tooltip": "Enable dynamic mode"}),
"dynamo_cache_size_limit": ("INT", {"default": 64, "min": 0, "max": 1024, "step": 1, "tooltip": "torch._dynamo.config.cache_size_limit"}),
},
}
RETURN_TYPES = ("MOCHICOMPILEARGS",)
@ -255,7 +256,7 @@ class MochiTorchCompileSettings:
CATEGORY = "MochiWrapper"
DESCRIPTION = "torch.compile settings, when connected to the model loader, torch.compile of the selected layers is attempted. Requires Triton and torch 2.5.0 is recommended"
def loadmodel(self, backend, fullgraph, mode, compile_dit, compile_final_layer, dynamic):
def loadmodel(self, backend, fullgraph, mode, compile_dit, compile_final_layer, dynamic, dynamo_cache_size_limit):
compile_args = {
"backend": backend,
@ -264,6 +265,7 @@ class MochiTorchCompileSettings:
"compile_dit": compile_dit,
"compile_final_layer": compile_final_layer,
"dynamic": dynamic,
"dynamo_cache_size_limit": dynamo_cache_size_limit,
}
return (compile_args, )