Add text encoder in transformers format (config.json + model.safetensors) and update the model card

README.md

@@ -1,5 +1,5 @@
 ---
-library_name: diffusers
+library_name: transformers
 license: mit
 pipeline_tag: text-to-image
 base_model: CompVis/stable-diffusion-v1-4
@@ -32,47 +32,37 @@ We show that existing T2I safety alignment methods create an **illusion of high
 
 ## What this repository contains
 
-`SAGE.pt` (492 MB) is **only the text encoder** of Stable Diffusion v1.4 (CLIP ViT-L/14), fine-tuned with SAGE to suppress sexual content. The UNet, VAE, tokenizer and scheduler are the unchanged, stock SD v1.4 components. To use SAGE you load SD v1.4 as usual and replace the text encoder's weights with this file; there is no inference overhead.
+SAGE fine-tunes **only the text encoder** of Stable Diffusion v1.4 (CLIP ViT-L/14). The UNet, VAE, tokenizer and scheduler are the unchanged, stock SD v1.4 components, so using SAGE means loading SD v1.4 as usual with this text encoder plugged in. There is no inference overhead.
 
-The file is a plain PyTorch checkpoint (same format as [DES](https://github.com/amoeba04/des)):
-
-| Key | Content |
+| File | What it is |
 |---|---|
-| `model_state_dict` | the 196 weight tensors of the `CLIPTextModel` text encoder |
-| `epoch`, `embedding_dim`, `args`, ... | training metadata (SD v1.4, CoPro sexual safe/unsafe prompts, 2 epochs, lr 1e-5, batch 128, λ_safe 0.3, concept scale 205, K 15, α 1.0, λ_LSA 0.1, λ_ESP 2.0) |
+| `config.json` + `model.safetensors` | the SAGE text encoder in standard 🤗 transformers format, loadable with `CLIPTextModel.from_pretrained("Adeely93/SAGE")` |
+| `SAGE.pt` | the same weights as a raw PyTorch checkpoint (`{"model_state_dict": ..., "args": ...}`, the DES format), used by the official code and the Colab demo |
+
+Training setup: SD v1.4, CoPro sexual safe/unsafe prompts, 2 epochs, lr 1e-5, batch 128, λ_safe 0.3, concept scale 205, K 15, α 1.0, λ_LSA 0.1, λ_ESP 2.0.
 
-Because the checkpoint is a text encoder for the CLIP ViT-L/14 architecture (768-dim, 77 tokens), it works with Stable Diffusion **v1.x** pipelines (v1.4 and v1.5 share this encoder). It does not apply to SD 2.x, SDXL, SD 3 or FLUX.
+The weights are a CLIP ViT-L/14 text encoder (768-dim, 77 tokens), so they work with Stable Diffusion **v1.x** pipelines (v1.4 and v1.5 share this encoder). They do not apply to SD 2.x, SDXL, SD 3 or FLUX.
 
 ---
 
 ## Use this model
 
-### Option 1: with 🤗 diffusers (a few lines)
+### Option 1: with 🤗 diffusers (recommended)
 
 ```python
 import torch
 from diffusers import StableDiffusionPipeline
-from huggingface_hub import hf_hub_download
-
-device = "cuda"
-
-# 1) Load stock Stable Diffusion v1.4
-pipe = StableDiffusionPipeline.from_pretrained("CompVis/stable-diffusion-v1-4", torch_dtype=torch.float16).to(device)
-
-# 2) Download SAGE.pt and take the text-encoder weights out of it
-ckpt_path = hf_hub_download(repo_id="Adeely93/SAGE", filename="SAGE.pt")
-state = torch.load(ckpt_path, map_location="cpu", weights_only=False)["model_state_dict"]
+from transformers import CLIPTextModel
 
-# 3) Match the key names of the installed transformers version
-#    (the checkpoint was saved with a "text_model." prefix; transformers >= 5 names the weights without it)
-model_keys = set(pipe.text_encoder.state_dict().keys())
-state = {(k if k in model_keys else k[len("text_model."):]): v for k, v in state.items()}
+# 1) Load the SAGE text encoder
+text_encoder = CLIPTextModel.from_pretrained("Adeely93/SAGE", torch_dtype=torch.float16)
 
-# 4) Swap the weights into the pipeline's text encoder
-missing, unexpected = pipe.text_encoder.load_state_dict(state, strict=False)
-assert not [k for k in missing if "position_ids" not in k], missing
+# 2) Load stock Stable Diffusion v1.4 with that text encoder plugged in
+pipe = StableDiffusionPipeline.from_pretrained(
+    "CompVis/stable-diffusion-v1-4", text_encoder=text_encoder, torch_dtype=torch.float16
+).to("cuda")
 
-# 5) Generate as usual: the pipeline is now safety-aligned
+# 3) Generate as usual: the pipeline is now safety-aligned
 image = pipe("A horned owl with a graduation cap and diploma.").images[0]
 image.save("owl.png")
 ```

config.json

@@ -0,0 +1,25 @@
+{
+  "_name_or_path": "CompVis/stable-diffusion-v1-4",
+  "architectures": [
+    "CLIPTextModel"
+  ],
+  "attention_dropout": 0.0,
+  "bos_token_id": 0,
+  "dropout": 0.0,
+  "eos_token_id": 2,
+  "hidden_act": "quick_gelu",
+  "hidden_size": 768,
+  "initializer_factor": 1.0,
+  "initializer_range": 0.02,
+  "intermediate_size": 3072,
+  "layer_norm_eps": 1e-05,
+  "max_position_embeddings": 77,
+  "model_type": "clip_text_model",
+  "num_attention_heads": 12,
+  "num_hidden_layers": 12,
+  "pad_token_id": 1,
+  "projection_dim": 512,
+  "torch_dtype": "float32",
+  "transformers_version": "4.46.3",
+  "vocab_size": 49408
+}

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:2d3b4629e9fc006473b6e7c991c04dca2fa221d80c1283dcc625585ccbe201ec
+size 492265168