Spaces:
Running on Zero
Running on Zero
File size: 4,840 Bytes
14e886f | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 | """Pinned Sakura Qwen3-VL encoder with the official FLUX.3 processor/config.
Only the shared text/vision encoder changes. The FLUX image transformers and
VAE remain official. Set FLUX3_TEXT_ENCODER=official to restore that encoder.
"""
from __future__ import annotations
import hashlib
import json
import os
from pathlib import Path
import tempfile
OFFICIAL_REPO = "black-forest-labs/flux-3-image-klein-base"
OFFICIAL_REVISION = "3f7fa5ee22ca1d5a5cc095ac7cc96e6a173b4c86"
SAKURA_REPO = "webmp3/Sakura-FLUX.3-Image-Klein-TextEncoder-Uncensored-GGUF"
SAKURA_REVISION = "fe49271a47b61efe0d96b476306c25934d943472"
SAKURA_FILE = "Sakura-TextEncoder-Qwen3VL-4B-Uncensored-bf16-8.27GiB.safetensors"
SAKURA_BYTES = 8_875_719_408
# All 713 names, shapes and dtypes independently match the official encoder.
ENCODER_SCHEMA_SHA256 = "f4619891efac96ad653344f1b195a51427b5cc9637769f58f435d20af26c93dc"
ENCODER_MODE = os.environ.get("FLUX3_TEXT_ENCODER", "sakura").strip().lower()
if ENCODER_MODE not in {"sakura", "official"}:
raise ValueError("FLUX3_TEXT_ENCODER must be 'sakura' or 'official'.")
ENCODER_LABEL = "Sakura text encoder" if ENCODER_MODE == "sakura" else "Official text encoder"
def text_encoder_info():
return {
"variant": ENCODER_MODE,
"label": ENCODER_LABEL,
"repo": SAKURA_REPO if ENCODER_MODE == "sakura" else OFFICIAL_REPO,
"revision": SAKURA_REVISION if ENCODER_MODE == "sakura" else OFFICIAL_REVISION,
"processor_repo": OFFICIAL_REPO,
"processor_revision": OFFICIAL_REVISION,
}
def _validate_checkpoint(checkpoint):
"""Read only the safetensors header; never allocate the 8.27 GiB tensors."""
path = Path(checkpoint)
if path.stat().st_size != SAKURA_BYTES:
raise RuntimeError("The Sakura text encoder checkpoint has an unexpected size.")
with path.open("rb") as stream:
header_size = int.from_bytes(stream.read(8), "little")
if not 2 <= header_size <= 1_000_000:
raise RuntimeError("The Sakura text encoder checkpoint has an invalid header.")
try:
header = json.loads(stream.read(header_size))
schema = {name: {"dtype": value["dtype"], "shape": value["shape"]}
for name, value in header.items() if name != "__metadata__"}
except (ValueError, TypeError, KeyError, AttributeError) as exc:
raise RuntimeError("The Sakura text encoder checkpoint header is unreadable.") from exc
digest = hashlib.sha256(json.dumps(schema, sort_keys=True, separators=(",", ":")).encode()).hexdigest()
if digest != ENCODER_SCHEMA_SHA256:
raise RuntimeError("The Sakura text encoder tensor layout differs from the verified Qwen3-VL encoder.")
def _check_loading_info(info, config):
# Transformers can report an omitted output head even when it is correctly
# tied to the input embeddings. The checkpoint schema is checked separately.
allowed_missing = {"lm_head.weight"} if getattr(config, "tie_word_embeddings", False) else set()
missing = set(info.get("missing_keys", [])) - allowed_missing
if missing or any(info.get(key) for key in ("unexpected_keys", "mismatched_keys", "error_msgs")):
raise RuntimeError("Text encoder loading was incomplete; refusing to run with missing or mismatched weights.")
def load_text_encoder(dtype, token=None):
from huggingface_hub import hf_hub_download
from transformers import AutoConfig, Qwen3VLForConditionalGeneration
config = AutoConfig.from_pretrained(
OFFICIAL_REPO, subfolder="text_encoder", revision=OFFICIAL_REVISION, token=token,
)
if ENCODER_MODE == "official":
model, info = Qwen3VLForConditionalGeneration.from_pretrained(
OFFICIAL_REPO, subfolder="text_encoder", revision=OFFICIAL_REVISION,
config=config, dtype=dtype, token=token, output_loading_info=True,
)
else:
checkpoint = hf_hub_download(
SAKURA_REPO, SAKURA_FILE, revision=SAKURA_REVISION, token=token,
)
_validate_checkpoint(checkpoint)
# Transformers expects model.safetensors beside config.json. A temporary
# symlink keeps the cached weight file in place without an 8 GiB copy.
# The Space runs Linux; only this tiny directory is temporary.
with tempfile.TemporaryDirectory(prefix="flux3-sakura-encoder-") as directory:
config.save_pretrained(directory)
os.symlink(str(Path(checkpoint).resolve()), str(Path(directory) / "model.safetensors"))
model, info = Qwen3VLForConditionalGeneration.from_pretrained(
directory, config=config, dtype=dtype, local_files_only=True,
output_loading_info=True,
)
_check_loading_info(info, config)
model.tie_weights()
return model.eval()
|