Phase 1

extensions_built_in/diffusion_models/chroma/src/model.py

@@ -1,4 +1,6 @@
-from dataclasses import dataclass
+from dataclasses import dataclass, replace
+
+from toolkit.models.v2._mixin import OstrisModelMixin
 
 import torch
 from torch import Tensor, nn
@@ -86,11 +88,40 @@ def modify_mask_to_attend_padding(mask, max_seq_length, num_extra_padding=8):
     return modified_mask
 
 
-class Chroma(nn.Module):
+class Chroma(nn.Module, OstrisModelMixin):
     """
     Transformer model for flow matching on sequences.
     """
 
+    @classmethod
+    def aitk_config_from_state_dict(cls, state_dict):
+        # block counts come from the checkpoint's key indices
+        double_blocks = 0
+        single_blocks = 0
+        for key in state_dict.keys():
+            if "double_blocks" in key:
+                block_num = int(key.split(".")[1]) + 1
+                if block_num > double_blocks:
+                    double_blocks = block_num
+            elif "single_blocks" in key:
+                block_num = int(key.split(".")[1]) + 1
+                if block_num > single_blocks:
+                    single_blocks = block_num
+        print(f"Double Blocks: {double_blocks}")
+        print(f"Single Blocks: {single_blocks}")
+        return replace(
+            chroma_params, depth=double_blocks, depth_single_blocks=single_blocks
+        )
+
+    @classmethod
+    def aitk_from_config(cls, config):
+        with torch.device("meta"):
+            return cls(config)
+
+    @classmethod
+    def get_transformer_block_names(cls):
+        return ["double_blocks", "single_blocks"]
+
     def __init__(self, params: ChromaParams):
         super().__init__()
         self.params = params

extensions_built_in/diffusion_models/chroma/src/radiance.py

@@ -1,4 +1,6 @@
-from dataclasses import dataclass
+from dataclasses import dataclass, replace
+
+from toolkit.models.v2._mixin import OstrisModelMixin
 
 import torch
 from torch import Tensor, nn
@@ -100,11 +102,40 @@ def modify_mask_to_attend_padding(mask, max_seq_length, num_extra_padding=8):
     return modified_mask
 
 
-class Chroma(nn.Module):
+class Chroma(nn.Module, OstrisModelMixin):
     """
     Transformer model for flow matching on sequences.
     """
 
+    @classmethod
+    def aitk_config_from_state_dict(cls, state_dict):
+        # block counts come from the checkpoint's key indices
+        double_blocks = 0
+        single_blocks = 0
+        for key in state_dict.keys():
+            if "double_blocks" in key:
+                block_num = int(key.split(".")[1]) + 1
+                if block_num > double_blocks:
+                    double_blocks = block_num
+            elif "single_blocks" in key:
+                block_num = int(key.split(".")[1]) + 1
+                if block_num > single_blocks:
+                    single_blocks = block_num
+        print(f"Double Blocks: {double_blocks}")
+        print(f"Single Blocks: {single_blocks}")
+        return replace(
+            chroma_params, depth=double_blocks, depth_single_blocks=single_blocks
+        )
+
+    @classmethod
+    def aitk_from_config(cls, config):
+        with torch.device("meta"):
+            return cls(config)
+
+    @classmethod
+    def get_transformer_block_names(cls):
+        return ["double_blocks", "single_blocks"]
+
     def __init__(self, params: ChromaParams):
         super().__init__()
         self.params = params