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

Phase 2

Browse Source
This commit is contained in:
Jaret Burkett
2026-08-27 18:18:31 -06:00
parent 45886f01b2
commit 520d96aac3
27 changed files with 333 additions and 253 deletions
Download Patch File Download Diff File Expand all files Collapse all files

7
extensions_built_in/diffusion_models/ideogram4/ideogram4.py
View File

@@ -238,9 +238,6 @@ class Ideogram4Model(BaseModel):
self.print_and_status_update("Loading transformer")
transformer_config = Ideogram4Config()
with torch.device("meta"):
transformer = Ideogram4Transformer2DModel(transformer_config)
self.print_and_status_update(" - fetching transformer weights")
state_dict = _load_component_state_dict(
base, "transformer", "diffusion_pytorch_model"
@@ -250,7 +247,9 @@ class Ideogram4Model(BaseModel):
state_dict, dtype, self.device_torch, self.model_config.low_vram
)
self.print_and_status_update(" - loading transformer state dict")
transformer.load_state_dict(state_dict, assign=True)
transformer = Ideogram4Transformer2DModel.load_from_state_dict(
state_dict, dtype, config=transformer_config
)
del state_dict
flush()