Spaces:
Running on Zero
Running on Zero
| """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() | |