Fix a few issues with compile. Changed defaults. Future proofed block layer compile.

This commit is contained in:
Jaret Burkett committed 2026-06-12 11:43:43 -06:00
1 parent a5558ae7d9
commit d83f7dd4d9
4 files changed
+19 -20

No files matched your search

+13 -12
View File
@@ -2099,24 +2099,25 @@ class BaseSDTrainProcess(BaseTrainProcess):
block_compile = getattr(self.model_config, 'block_compile', False) block_compile = getattr(self.model_config, 'block_compile', False)
cache_info = f", cache_size_limit={cache_size_limit}" if cache_size_limit != 8 else "" cache_info = f", cache_size_limit={cache_size_limit}" if cache_size_limit != 8 else ""
# ==================================================== # ====================================================
# BLOCK COMPILE # BLOCK COMPILE
# ==================================================== # ====================================================
if block_compile: if block_compile:
BLOCK_LIST_ATTRS = self.sd.get_transformer_block_names()
if BLOCK_LIST_ATTRS is None or len(BLOCK_LIST_ATTRS) == 0:
BLOCK_LIST_ATTRS = [
'layers',
'transformer_blocks',
'single_transformer_blocks',
'double_stream_blocks',
'single_stream_blocks',
'double_blocks',
'single_blocks',
'blocks',
]
inner_unet = unwrap_model(self.sd.unet) inner_unet = unwrap_model(self.sd.unet)
BLOCK_LIST_ATTRS = [
'layers',
'transformer_blocks',
'single_transformer_blocks',
'double_stream_blocks',
'single_stream_blocks',
'double_blocks',
'single_blocks',
'blocks',
]
compiled_block_count = 0 compiled_block_count = 0
for attr_name in BLOCK_LIST_ATTRS: for attr_name in BLOCK_LIST_ATTRS:
+3 -1
View File
@@ -709,7 +709,9 @@ class ModelConfig:
if self.compile and self.quantize: if self.compile and self.quantize:
print("Quantized model detected - allowing torch.compile (experimental - Does not work with offloaded transformer)") print("Quantized model detected - allowing torch.compile (experimental - Does not work with offloaded transformer)")
# self.compile = False # make it torchao instead of quantio for compatibility with torch compile
if self.qtype == "qfloat8":
self.qtype = "float8"
self.block_compile = kwargs.get("block_compile", False) self.block_compile = kwargs.get("block_compile", False)
self.compile_mode = kwargs.get("compile_mode", "default") self.compile_mode = kwargs.get("compile_mode", "default")
self.compile_fullgraph = kwargs.get("compile_fullgraph", True) self.compile_fullgraph = kwargs.get("compile_fullgraph", True)
+2 -6
View File
@@ -32,11 +32,7 @@ export const defaultSliderConfig: SliderConfig = {
}; };
export const defaultCompileOptions = { export const defaultCompileOptions = {
block_compile: false, block_compile: true,
compile_mode: 'default',
compile_fullgraph: false,
compile_dynamic: false,
cache_size_limit: undefined,
}; };
export const defaultJobConfig: JobConfig = { export const defaultJobConfig: JobConfig = {
@@ -102,7 +98,6 @@ export const defaultJobConfig: JobConfig = {
diff_output_preservation_class: 'person', diff_output_preservation_class: 'person',
switch_boundary_every: 1, switch_boundary_every: 1,
loss_type: 'mse', loss_type: 'mse',
compile: false,
}, },
logging: { logging: {
log_every: 1, log_every: 1,
@@ -117,6 +112,7 @@ export const defaultJobConfig: JobConfig = {
arch: 'flex1', arch: 'flex1',
low_vram: false, low_vram: false,
model_kwargs: {}, model_kwargs: {},
compile: false,
}, },
sample: defaultSampleConfig, sample: defaultSampleConfig,
}, },
+1 -1
View File
@@ -1 +1 @@
VERSION = "0.10.8" VERSION = "0.10.9"