Phase 2

extensions_built_in/diffusion_models/flux2/src/model.py

@@ -1,4 +1,6 @@
 import torch
+
+from toolkit.models.v2._mixin import OstrisModelMixin
 from einops import rearrange
 from torch import Tensor, nn
 import torch.utils.checkpoint as ckpt
@@ -54,7 +56,7 @@ class FakeConfig:
         self.patch_size = 1
 
 
-class Flux2(nn.Module):
+class Flux2(nn.Module, OstrisModelMixin):
     def __init__(self, params: Flux2Params):
         super().__init__()
         self.config = FakeConfig()