--- license: mit tags: - clip - text-encoder - knowledge-distillation - stable-diffusion - keras - jax pipeline_tag: feature-extraction --- # 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](https://arxiv.org/abs/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: | Source | Weight | Notes | |---|---|---| | DiffusionDB | 0.45 | long, compositional, attribute-rich prompts (the target failure region); lightly cleaned | | LAION-6.5-aesthetics | 0.35 | captions | | ImageNet (new JoyCaption captions) | 0.20 | captions only, no image join | ## 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: ![sample grid](assets/samples_grid.png) ## Usage ```python 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`](https://huggingface.co/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. ```bash pip install torch transformers diffusers h5py huggingface_hub ``` ```python 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.