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

Add initial support for chroma radiance

Browse Source
This commit is contained in:
Jaret Burkett
2025-09-10 08:41:05 -06:00
parent af6fdaaaf9
commit b95c17dc17
9 changed files with 1339 additions and 20 deletions
Download Patch File Download Diff File Expand all files Collapse all files

10
extensions_built_in/diffusion_models/chroma/src/model.py
View File

@@ -156,13 +156,19 @@ class Chroma(nn.Module):
)
# TODO: move this hardcoded value to config
self.mod_index_length = 344
# single layer has 3 modulation vectors
# double layer has 6 modulation vectors for each expert
# final layer has 2 modulation vectors
self.mod_index_length = 3 * params.depth_single_blocks + 2 * 6 * params.depth + 2
self.depth_single_blocks = params.depth_single_blocks
self.depth_double_blocks = params.depth
# self.mod_index = torch.tensor(list(range(self.mod_index_length)), device=0)
self.register_buffer(
"mod_index",
torch.tensor(list(range(self.mod_index_length)), device="cpu"),
persistent=False,
)
self.approximator_in_dim = params.approximator_in_dim
@property
def device(self):
@@ -213,7 +219,7 @@ class Chroma(nn.Module):
# then and only then we could concatenate it together
input_vec = torch.cat([timestep_guidance, modulation_index], dim=-1)
mod_vectors = self.distilled_guidance_layer(input_vec.requires_grad_(True))
mod_vectors_dict = distribute_modulations(mod_vectors)
mod_vectors_dict = distribute_modulations(mod_vectors, self.depth_single_blocks, self.depth_double_blocks)
ids = torch.cat((txt_ids, img_ids), dim=1)
pe = self.pe_embedder(ids)