Phase 1

extensions_built_in/diffusion_models/prx_pixel_t2i/prx_pixel_t2i.py

@@ -126,9 +126,7 @@ class PRXPixelT2IModel(BaseModel):
         self.print_and_status_update("Loading transformer")
         # from_pretrained reads config.json (bottleneck_size, resolution_embeds,
         # in_channels=3, ...) and the safetensors in one shot.
-        transformer = PRXTransformer2DModel.from_pretrained(
-            model_path, subfolder="transformer", torch_dtype=dtype
-        )
+        transformer = PRXTransformer2DModel.load_model(model_path, dtype=dtype)
         transformer.to(dtype=dtype)
         flush()
 

extensions_built_in/diffusion_models/prx_pixel_t2i/src/transformer_prx.py

@@ -40,6 +40,8 @@ from diffusers.models.attention_dispatch import dispatch_attention_fn
 from diffusers.models.embeddings import get_timestep_embedding
 from diffusers.models.modeling_outputs import Transformer2DModelOutput
 from diffusers.models.modeling_utils import ModelMixin
+
+from toolkit.models.v2._mixin import OstrisModelMixin
 from diffusers.models.normalization import RMSNorm
 
 
@@ -666,7 +668,7 @@ def seq2img(seq: torch.Tensor, patch_size: int, shape: torch.Tensor) -> torch.Te
     return seq
 
 
-class PRXTransformer2DModel(ModelMixin, ConfigMixin, AttentionMixin):
+class PRXTransformer2DModel(ModelMixin, ConfigMixin, AttentionMixin, OstrisModelMixin):
     r"""
     Transformer-based 2D model for text to image generation.
 
@@ -702,6 +704,12 @@ class PRXTransformer2DModel(ModelMixin, ConfigMixin, AttentionMixin):
             `PRXResolutionEmbedder`. Used by the PRX-7B variant.
     """
 
+    aitk_subfolder = "transformer"
+
+    @classmethod
+    def get_transformer_block_names(cls):
+        return ["blocks"]
+
     config_name = "config.json"
     _supports_gradient_checkpointing = True