Phase 1

extensions_built_in/diffusion_models/chroma/chroma_model.py

@@ -5,8 +5,9 @@ import torch
 from toolkit.config_modules import GenerateImageConfig, ModelConfig
 from PIL import Image
 from toolkit.models.base_model import BaseModel
+from toolkit.models.v2.text_encoders.t5 import T5TextEncoder
+from toolkit.models.v2.vae.autoencoder_kl import KLVAE
 from toolkit.basic import flush
-from diffusers import AutoencoderKL
 # from toolkit.pixel_shuffle_encoder import AutoencoderPixelMixer
 from toolkit.prompt_utils import PromptEmbeds
 from toolkit.samplers.custom_flowmatch_sampler import CustomFlowMatchEulerDiscreteScheduler
@@ -14,7 +15,6 @@ from toolkit.dequantize import patch_dequantization_on_save
 from toolkit.accelerator import unwrap_model
 from optimum.quanto import freeze, QTensor
 from toolkit.util.quantize import quantize, get_qtype
-from transformers import T5TokenizerFast, T5EncoderModel, CLIPTextModel, CLIPTokenizer
 from .pipeline import ChromaPipeline, prepare_latent_image_ids
 from einops import rearrange, repeat
 import random
@@ -148,37 +148,17 @@ class ChromaModel(BaseModel):
 
         self.print_and_status_update("Loading transformer")
         
-        chroma_state_dict = load_file(model_path, 'cpu')
-        
-        # determine number of double and single blocks
-        double_blocks = 0
-        single_blocks = 0
-        for key in chroma_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}")
-
-        chroma_params.depth = double_blocks
-        chroma_params.depth_single_blocks = single_blocks
-        transformer = Chroma(chroma_params)
-        
+        if model_path.endswith(".safetensors"):
+            transformer = Chroma.load_model(model_path, dtype=dtype)
+        else:
+            transformer = Chroma.load_from_state_dict(load_file(model_path, "cpu"), dtype)
         # add dtype, not sure why it doesnt have it
         transformer.dtype = dtype
-        # load the state dict into the model
-        transformer.load_state_dict(chroma_state_dict)
-        
         transformer.to(self.quantize_device, dtype=dtype)
-        
+
         transformer.config = FakeConfig()
-        transformer.config.num_layers = double_blocks
-        transformer.config.num_single_layers = single_blocks
+        transformer.config.num_layers = transformer.params.depth
+        transformer.config.num_single_layers = transformer.params.depth_single_blocks
 
         if self.model_config.quantize:
             # patch the state dict method
@@ -195,21 +175,9 @@ class ChromaModel(BaseModel):
         flush()
 
         self.print_and_status_update("Loading T5")
-        tokenizer_2 = T5TokenizerFast.from_pretrained(
-            extras_path, subfolder="tokenizer_2", torch_dtype=dtype
-        )
-        text_encoder_2 = T5EncoderModel.from_pretrained(
-            extras_path, subfolder="text_encoder_2", torch_dtype=dtype
-        )
-        text_encoder_2.to(self.device_torch, dtype=dtype)
-        flush()
-
-        if self.model_config.quantize_te:
-            self.print_and_status_update("Quantizing T5")
-            quantize(text_encoder_2, weights=get_qtype(
-                self.model_config.qtype))
-            freeze(text_encoder_2)
-            flush()
+        tokenizer_2 = T5TextEncoder.load_tokenizer(extras_path)
+        text_encoder_2 = T5TextEncoder.load_model(extras_path, dtype=dtype)
+        self.prepare_text_encoder(text_encoder_2, dtype=dtype)
 
         # self.print_and_status_update("Loading CLIP")
         text_encoder = FakeCLIP()
@@ -219,12 +187,7 @@ class ChromaModel(BaseModel):
         self.noise_scheduler = ChromaModel.get_train_scheduler()
         
         self.print_and_status_update("Loading VAE")
-        vae = AutoencoderKL.from_pretrained(
-            extras_path,
-            subfolder="vae",
-            torch_dtype=dtype
-        )
-        vae = vae.to(self.device_torch, dtype=dtype)
+        vae = KLVAE.load_model(extras_path, dtype=dtype, device=self.device_torch)
 
         self.print_and_status_update("Making pipe")
 

extensions_built_in/diffusion_models/chroma/chroma_radiance_model.py

@@ -5,8 +5,8 @@ import torch
 from toolkit.config_modules import GenerateImageConfig, ModelConfig
 from PIL import Image
 from toolkit.models.base_model import BaseModel
+from toolkit.models.v2.text_encoders.t5 import T5TextEncoder
 from toolkit.basic import flush
-from diffusers import AutoencoderKL
 # from toolkit.pixel_shuffle_encoder import AutoencoderPixelMixer
 from toolkit.prompt_utils import PromptEmbeds
 from toolkit.samplers.custom_flowmatch_sampler import CustomFlowMatchEulerDiscreteScheduler
@@ -14,7 +14,6 @@ from toolkit.dequantize import patch_dequantization_on_save
 from toolkit.accelerator import unwrap_model
 from optimum.quanto import freeze, QTensor
 from toolkit.util.quantize import quantize, get_qtype
-from transformers import T5TokenizerFast, T5EncoderModel, CLIPTextModel, CLIPTokenizer
 from .pipeline import ChromaPipeline, prepare_latent_image_ids
 from einops import rearrange, repeat
 import random
@@ -152,38 +151,16 @@ class ChromaRadianceModel(BaseModel):
         
         if model_path.endswith('.pth') or model_path.endswith('.pt'):
             chroma_state_dict = torch.load(model_path, map_location='cpu', weights_only=True)
+            transformer = Chroma.load_from_state_dict(chroma_state_dict, dtype)
         else:
-            chroma_state_dict = load_file(model_path, 'cpu')
-        
-        # determine number of double and single blocks
-        double_blocks = 0
-        single_blocks = 0
-        for key in chroma_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}")
-
-        chroma_params.depth = double_blocks
-        chroma_params.depth_single_blocks = single_blocks
-        transformer = Chroma(chroma_params)
-        
+            transformer = Chroma.load_model(model_path, dtype=dtype)
         # add dtype, not sure why it doesnt have it
         transformer.dtype = dtype
-        # load the state dict into the model
-        transformer.load_state_dict(chroma_state_dict)
-        
         transformer.to(self.quantize_device, dtype=dtype)
-        
+
         transformer.config = FakeConfig()
-        transformer.config.num_layers = double_blocks
-        transformer.config.num_single_layers = single_blocks
+        transformer.config.num_layers = transformer.params.depth
+        transformer.config.num_single_layers = transformer.params.depth_single_blocks
 
         if self.model_config.quantize:
             # patch the state dict method
@@ -200,21 +177,9 @@ class ChromaRadianceModel(BaseModel):
         flush()
 
         self.print_and_status_update("Loading T5")
-        tokenizer_2 = T5TokenizerFast.from_pretrained(
-            extras_path, subfolder="tokenizer_2", torch_dtype=dtype
-        )
-        text_encoder_2 = T5EncoderModel.from_pretrained(
-            extras_path, subfolder="text_encoder_2", torch_dtype=dtype
-        )
-        text_encoder_2.to(self.device_torch, dtype=dtype)
-        flush()
-
-        if self.model_config.quantize_te:
-            self.print_and_status_update("Quantizing T5")
-            quantize(text_encoder_2, weights=get_qtype(
-                self.model_config.qtype))
-            freeze(text_encoder_2)
-            flush()
+        tokenizer_2 = T5TextEncoder.load_tokenizer(extras_path)
+        text_encoder_2 = T5TextEncoder.load_model(extras_path, dtype=dtype)
+        self.prepare_text_encoder(text_encoder_2, dtype=dtype)
 
         # self.print_and_status_update("Loading CLIP")
         text_encoder = FakeCLIP()

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