github/ai-toolkit
1
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

Allow fine tuning pruned versions of chroma. Allow flash attention 2 for chroma if it is installed.

Browse Source
This commit is contained in:
Jaret Burkett
2025-05-21 07:02:50 -06:00
parent 48e11cf843
commit 79499fa795
3 changed files with 43 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files

23
extensions_built_in/diffusion_models/chroma/chroma_model.py
View File

@@ -141,19 +141,38 @@ class ChromaModel(BaseModel):
extras_path = 'ostris/Flex.1-alpha'
self.print_and_status_update("Loading transformer")
chroma_state_dict = load_file(model_path, 'cpu')
# determine number of double and single blocks
double_blocks = 0
single_blocks = 0
for key in chroma_state_dict.keys():
if "double_blocks" in key:
block_num = int(key.split(".")[1]) + 1
if block_num > double_blocks:
double_blocks = block_num
elif "single_blocks" in key:
block_num = int(key.split(".")[1]) + 1
if block_num > single_blocks:
single_blocks = block_num
print(f"Double Blocks: {double_blocks}")
print(f"Single Blocks: {single_blocks}")
chroma_params.depth = double_blocks
chroma_params.depth_single_blocks = single_blocks
transformer = Chroma(chroma_params)
# add dtype, not sure why it doesnt have it
transformer.dtype = dtype
chroma_state_dict = load_file(model_path, 'cpu')
# load the state dict into the model
transformer.load_state_dict(chroma_state_dict)
transformer.to(self.quantize_device, dtype=dtype)
transformer.config = FakeConfig()
transformer.config.num_layers = double_blocks
transformer.config.num_single_layers = single_blocks
if self.model_config.quantize:
# patch the state dict method