File size: 4,033 Bytes
8a2be73 | 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 | """FLUX.2 Qwen3 instruction cache for Delta-flow tasks that have no previous experiment to copy it from.
Same computation as HumanWAM scripts/flux2/precompute_flux2_qwen3_embeds.py (chat template, enable_thinking False,
padding to max_length 128, hidden states of flux2.text_encoder.OUTPUT_LAYERS_QWEN3 concatenated, stored as bf16 +
bool mask), on CPU in bfloat16 (the GPU original also runs the model in bfloat16). make() first recomputes a reference
prompt whose cache already exists (A3) and refuses to write anything unless the result matches it.
"""
import hashlib
import os
import sys
from pathlib import Path
import torch
_MODEL = {}
def _load(code, qwen_path, flux2_src):
if "m" not in _MODEL:
sys.path.insert(0, str(Path(code) / "src"))
from imagewam.models.backbones.flux2_imports import ensure_flux2_importable
ensure_flux2_importable(str(flux2_src))
from flux2.text_encoder import OUTPUT_LAYERS_QWEN3
from transformers import AutoModelForCausalLM, AutoTokenizer
torch.set_num_threads(min(32, os.cpu_count() or 8))
tok = AutoTokenizer.from_pretrained(str(qwen_path))
model = AutoModelForCausalLM.from_pretrained(str(qwen_path), torch_dtype=torch.bfloat16).eval()
_MODEL["m"] = (tok, model, tuple(OUTPUT_LAYERS_QWEN3))
return _MODEL["m"]
def embed(prompt, code, qwen_path, flux2_src, context_len=128):
tok, model, layers = _load(code, qwen_path, flux2_src)
rendered = tok.apply_chat_template([{"role": "user", "content": prompt}], tokenize=False,
add_generation_prompt=True, enable_thinking=False)
enc = tok([rendered], return_tensors="pt", padding="max_length", truncation=True, max_length=context_len)
with torch.no_grad():
out = model(input_ids=enc.input_ids, attention_mask=enc.attention_mask, output_hidden_states=True,
use_cache=False)
hidden = torch.stack([out.hidden_states[k] for k in layers], dim=1) # [1, c, l, d]
hidden = hidden.permute(0, 2, 1, 3).reshape(1, context_len, -1) # b l (c d)
return dict(text_hidden_states=hidden[0].to(torch.bfloat16).clone(),
text_attention_mask=enc.attention_mask[0].to(torch.bool).clone())
def cache_name(prompt, context_len=128):
return f"{hashlib.sha256(prompt.encode('utf-8')).hexdigest()}.qwen3_flux2_len{context_len}.pt"
def check_reference(ref_prompt, ref_dir, code, qwen_path, flux2_src):
ref = torch.load(Path(ref_dir) / cache_name(ref_prompt), map_location="cpu", weights_only=True)
new = embed(ref_prompt, code, qwen_path, flux2_src)
if not torch.equal(ref["text_attention_mask"].bool(), new["text_attention_mask"]):
raise SystemExit("qwen cache: attention mask differs from the reference cache")
m = ref["text_attention_mask"].bool()
a, b = ref["text_hidden_states"][m].float(), new["text_hidden_states"][m].float()
cos = torch.nn.functional.cosine_similarity(a, b, dim=-1)
rel = ((a - b).norm() / a.norm()).item()
info = dict(min_token_cosine=float(cos.min()), mean_token_cosine=float(cos.mean()), relative_l2=rel,
tokens=int(m.sum()))
if info["min_token_cosine"] < 0.99 or rel > 0.05:
raise SystemExit(f"qwen cache: CPU recomputation does not match the reference cache: {info}")
return info
def make(prompt, out_dir, code, qwen_path, flux2_src, ref_prompt, ref_dir):
out_dir = Path(out_dir)
out_dir.mkdir(parents=True, exist_ok=True)
path = out_dir / cache_name(prompt)
if path.is_file():
return path, None
info = check_reference(ref_prompt, ref_dir, code, qwen_path, flux2_src)
payload = embed(prompt, code, qwen_path, flux2_src)
if not torch.isfinite(payload["text_hidden_states"].float()).all() or not payload["text_attention_mask"].any():
raise SystemExit("qwen cache: invalid embedding")
tmp = out_dir / f".{path.name}.tmp{os.getpid()}"
torch.save(payload, tmp)
os.replace(tmp, path)
return path, info
|