Phase 1

extensions_built_in/diffusion_models/f_light/f_light.py

@@ -6,15 +6,15 @@ import yaml
 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.prompt_utils import PromptEmbeds
 from toolkit.samplers.custom_flowmatch_sampler import CustomFlowMatchEulerDiscreteScheduler
 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
 from .src import FLitePipeline, DiT
 
 if TYPE_CHECKING:
@@ -74,11 +74,7 @@ class FLiteModel(BaseModel):
 
         self.print_and_status_update("Loading transformer")
 
-        transformer = DiT.from_pretrained(
-            model_path,
-            subfolder="dit_model",
-            torch_dtype=dtype,
-        )
+        transformer = DiT.load_model(model_path, dtype=dtype)
         
         transformer.to(self.quantize_device, dtype=dtype)
 
@@ -97,31 +93,16 @@ class FLiteModel(BaseModel):
         flush()
 
         self.print_and_status_update("Loading T5")
-        tokenizer = T5TokenizerFast.from_pretrained(
-            extras_path, subfolder="tokenizer", torch_dtype=dtype
+        tokenizer = T5TextEncoder.load_tokenizer(extras_path, subfolder="tokenizer")
+        text_encoder = T5TextEncoder.load_model(
+            extras_path, dtype=dtype, subfolder="text_encoder"
         )
-        text_encoder = T5EncoderModel.from_pretrained(
-            extras_path, subfolder="text_encoder", torch_dtype=dtype
-        )
-        text_encoder.to(self.device_torch, dtype=dtype)
-        flush()
-
-        if self.model_config.quantize_te:
-            self.print_and_status_update("Quantizing T5")
-            quantize(text_encoder, weights=get_qtype(
-                self.model_config.qtype))
-            freeze(text_encoder)
-            flush()
+        self.prepare_text_encoder(text_encoder, dtype=dtype)
 
         self.noise_scheduler = FLiteModel.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/f_light/src/model.py

@@ -7,6 +7,8 @@ import torch.nn.functional as F
 from diffusers.configuration_utils import ConfigMixin, register_to_config
 from diffusers.loaders import FromOriginalModelMixin, PeftAdapterMixin
 from diffusers.models.modeling_utils import ModelMixin
+
+from toolkit.models.v2._mixin import OstrisModelMixin
 from diffusers.utils.accelerate_utils import apply_forward_hook
 from einops import rearrange
 from peft import get_peft_model_state_dict, set_peft_model_state_dict
@@ -302,7 +304,13 @@ def apply_rotary_emb(x, cos, sin):
     return torch.cat([y1, y2], 3).to(dtype=orig_dtype)
 
 
-class DiT(ModelMixin, ConfigMixin, FromOriginalModelMixin, PeftAdapterMixin):  # type: ignore[misc]
+class DiT(ModelMixin, ConfigMixin, FromOriginalModelMixin, PeftAdapterMixin, OstrisModelMixin):  # type: ignore[misc]
+    aitk_subfolder = "dit_model"
+
+    @classmethod
+    def get_transformer_block_names(cls):
+        return ["blocks"]
+
     _supports_gradient_checkpointing = True
     
     @register_to_config