tinyclip-vit-h-distill

tinyclip-vit-h-distill

A TinyCLIP-39M text tower + MLP adapter, distilled from the OpenCLIP ViT-H/14 text encoder used by Stable Diffusion 2.1, so it can serve as a drop-in, much smaller replacement for that encoder's cross-attention conditioning. Trained with Keras 3 on JAX.

Method

Vision-based (step-following) knowledge distillation, per DistillT5 (Wang et al., 2025, arXiv:2503.19897):

• Teacher: OpenCLIP ViT-H/14 text encoder (frozen), as used by SD2.1-base.
• Student: TinyCLIP-39M text tower → Dense(2048) → ReLU → Dropout(0.1) → Dense(1024) adapter head, projecting the student's 512-d hidden states to the teacher's 1024-d cross-attention dimension.
• Frozen UNet: the Stable Diffusion 2.1-base UNet itself. The student is trained so that feeding its context embeddings into this frozen UNet reproduces the same noise prediction the teacher's embeddings would produce — i.e. the loss is defined on the UNet's output, not on the embeddings directly. No images are required (step-following, not full-image rollout).
• Loss:
  • Primary: MSE between teacher-conditioned and student-conditioned UNet noise predictions, on cached teacher denoising trajectories.
  • Auxiliary: masked per-token cosine distance between student and teacher context embeddings (real tokens only, up to EOT), weighted by a small emb_lambda that is linearly annealed from 0.2 → 0.0 over the first 50k steps. This is a stabilizing nudge, not a co-equal objective — the paper flags embedding-matching alone as prone to mode collapse.
• Samples: decoded with the SD2.1 VAE decoder.

Training data

Prompts only (no images) from three sources, mixed ~45/35/20:

Training config

• Optimizer: Lion, lr 1e-5, weight decay 1e-3
• Batch size: 16, NUM_INFER_STEPS=20 per trajectory (one teacher rollout → up to 20 grad steps)
• Guidance scale sampled per trajectory from [2.0, 5.0]
• Precision: mixed_bfloat16 (fp32 master weights, bf16 compute); frozen teacher/UNet in pure bf16
• Hardware: TPU v5e (JAX), data-parallel across all local devices

Samples

Teacher (SD2.1 ViT-H) vs. this student text encoder, same prompts, same UNet, SD2.1 VAE decoder:

Usage

from huggingface_hub import hf_hub_download
import keras
import keras_hub
from keras import layers

STUDENT_CLIP_PRESET = "..."   # same TinyCLIP preset used during training
MLP_INTERMEDIATE = 2048
MLP_DROPOUT = 0.1
TEACHER_DIM = 1024

student_clip = keras_hub.models.CLIPBackbone.from_preset(STUDENT_CLIP_PRESET).text_encoder
token_id_input = keras.Input(shape=(None,), dtype="int32", name="token_ids")
seq = student_clip({"token_ids": token_id_input})
h = layers.Dense(MLP_INTERMEDIATE, name="adapter_dense_1")(seq)
h = layers.Activation("relu", name="adapter_relu")(h)
h = layers.Dropout(MLP_DROPOUT, name="adapter_dropout")(h)
context = layers.Dense(TEACHER_DIM, name="adapter_dense_2")(h)
student = keras.Model({"token_ids": token_id_input}, context, name="student_text_encoder")

ckpt_path = hf_hub_download(
    repo_id="masterofaudio2077/tinyclip-vit-h-distill",
    filename="model_weights.h5",
)
student.load_weights(ckpt_path)

token_ids = student_clip.preprocessor.generate_preprocess(["a cat sitting on a wooden table"])
context_embeddings = student.predict({"token_ids": token_ids})  # (1, 77, 1024)
# feed context_embeddings into the SD2.1 UNet's cross-attention `context` input

Usage as a diffusers text encoder (drop-in for SD2.1)

Swap this in as pipe.text_encoder on any StableDiffusionPipeline built on SD2.1-base — no Keras required. The transformer backbone was ported into keras_hub from wkcn/TinyCLIP-ViT-61M-32-Text-29M-LAION400M, a standard transformers.CLIPModel checkpoint whose text config (hidden_size=512, num_hidden_layers=9, num_attention_heads=8, intermediate_size=2048) matches this repo's model_weights.h5 exactly. Only the small adapter head (Dense(2048)->ReLU->Dropout->Dense(1024)) is pulled out of that .h5 file; everything else loads through transformers/diffusers directly. This has been run end-to-end (pipe.encode_prompt and a full pipe(...) call) against a real SD2.1-base checkpoint.

pip install torch transformers diffusers h5py huggingface_hub

import h5py
import torch
import torch.nn as nn
from transformers import CLIPTextModel, CLIPTokenizer
from huggingface_hub import hf_hub_download
from diffusers import StableDiffusionPipeline

BACKBONE_REPO = "wkcn/TinyCLIP-ViT-61M-32-Text-29M-LAION400M"


class TinyCLIPDistillTextEncoder(CLIPTextModel):
    """Drop-in replacement for `pipe.text_encoder`. Must subclass CLIPTextModel (not
    just nn.Module) — diffusers' from_pretrained isinstance-checks for PreTrainedModel."""

    def __init__(self, config, dim_hidden=2048, dim_out=1024, dropout=0.1):
        super().__init__(config)
        self.dense_1 = nn.Linear(config.hidden_size, dim_hidden)
        self.relu = nn.ReLU()
        self.dropout = nn.Dropout(dropout)
        self.dense_2 = nn.Linear(dim_hidden, dim_out)

    def load_adapter_(self, h5_path):
        with h5py.File(h5_path, "r") as f:
            w1 = torch.from_numpy(f["layers/dense/vars/0"][()]).float().T    # keras (in,out) -> torch (out,in)
            b1 = torch.from_numpy(f["layers/dense/vars/1"][()]).float()
            w2 = torch.from_numpy(f["layers/dense_1/vars/0"][()]).float().T
            b2 = torch.from_numpy(f["layers/dense_1/vars/1"][()]).float()
        self.dense_1.weight.data.copy_(w1); self.dense_1.bias.data.copy_(b1)
        self.dense_2.weight.data.copy_(w2); self.dense_2.bias.data.copy_(b2)
        return self

    def forward(self, input_ids, attention_mask=None, **kwargs):
        hidden = super().forward(input_ids, attention_mask=attention_mask).last_hidden_state
        context = self.dense_2(self.dropout(self.relu(self.dense_1(hidden))))
        return (context,)


# --- build the student text encoder: upstream TinyCLIP backbone + this repo's adapter ---
_backbone = CLIPTextModel.from_pretrained(BACKBONE_REPO)
student_text_encoder = TinyCLIPDistillTextEncoder(_backbone.config)
student_text_encoder.load_state_dict(_backbone.state_dict(), strict=False)  # backbone weights only; dense_1/dense_2 stay uninitialized here

h5_path = hf_hub_download(repo_id="masterofaudio2077/tinyclip-vit-h-distill", filename="model_weights.h5")
student_text_encoder.load_adapter_(h5_path).eval()

tokenizer = CLIPTokenizer.from_pretrained(BACKBONE_REPO)
tokenizer.model_max_length = 77  # upstream tokenizer_config ships a bogus sentinel here

# --- swap into any SD2.1-base pipeline ---
pipe = StableDiffusionPipeline.from_pretrained(
    "stabilityai/stable-diffusion-2-1-base",   # use your own trusted SD2.1-base checkpoint/mirror if this 404s for you
    text_encoder=student_text_encoder,
    tokenizer=tokenizer,
    safety_checker=None,
)

image = pipe("a cat sitting on a wooden table", guidance_scale=3.5).images[0]
image.save("out.png")

Note: stabilityai/stable-diffusion-2-1-base returned a 404 on the Hub as of this writing — point from_pretrained at whatever SD2.1-base checkpoint/mirror you already trust and use for inference; the text_encoder/tokenizer swap above works the same regardless of where the rest of the pipeline comes from.