Phase 2

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: