Phase 2

extensions_built_in/diffusion_models/ideogram4/ideogram4.py

@@ -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()
 

extensions_built_in/diffusion_models/ideogram4/src/transformer.py

@@ -12,6 +12,8 @@ import math
 from dataclasses import dataclass
 
 import torch
+
+from toolkit.models.v2._mixin import OstrisModelMixin
 import torch.nn as nn
 import torch.nn.functional as F
 from torch.utils.checkpoint import checkpoint
@@ -355,7 +357,7 @@ class Ideogram4FinalLayer(nn.Module):
         return self.linear(self.norm_final(x) * scale)
 
 
-class Ideogram4Transformer2DModel(nn.Module):
+class Ideogram4Transformer2DModel(nn.Module, OstrisModelMixin):
     """Ideogram 4 flow-matching transformer."""
 
     def __init__(self, config: Ideogram4Config) -> None: