"""Pre-compute Vera material identity embeddings for zero-token identity injection."""
import torch, os
os.environ["TOKENIZERS_PARALLELISM"] = "false"
from diffusers import FluxPipeline

pipe = FluxPipeline.from_pretrained("black-forest-labs/FLUX.1-dev", torch_dtype=torch.bfloat16)
pipe.to("mps")

pipe.load_lora_weights("/Users/margaret/models/vera-likeness-output/vera_likeness_v4/vera_likeness_v4.safetensors", adapter_name="likeness")
pipe.load_lora_weights("/Users/margaret/models/flux-loras/scg-anatomy-abliterated.safetensors", adapter_name="anatomy")
pipe.load_lora_weights("/Users/margaret/models/kintsugi-texture-output/kintsugi_texture_v1/kintsugi_texture_v1.safetensors", adapter_name="kintsugi")
pipe.set_adapters(["likeness", "anatomy", "kintsugi"], adapter_weights=[1.0, 0.7, 1.0])
print("Pipeline loaded with LoRAs")

identity_prompt = (
    "vera with luminous amber gemstone eyes, dark navy matte ceramic figure, "
    "NOT human skin, NOT glossy. Shattered and reassembled with abundant "
    "gold-filled kintsugi cracks across cheekbones, down the neck, branching "
    "across collarbones, ribs, hips, spine, and thighs. Thick gold repair "
    "lines with visible depth following natural fracture patterns. The gold "
    "glows from within. Matte blue-black fired ceramic surface texture."
)

print("Encoding identity embeddings...")
identity_embeds = pipe.encode_prompt(
    prompt=identity_prompt,
    prompt_2=identity_prompt,
    max_sequence_length=512,
)

out_dir = "/Users/margaret/models/vera-triple-stack/identity_cache"
os.makedirs(out_dir, exist_ok=True)

for i, emb in enumerate(identity_embeds):
    if emb is not None:
        path = os.path.join(out_dir, f"identity_embed_{i}.pt")
        torch.save(emb.cpu(), path)
        print(f"Saved embed {i}: shape={emb.shape}, path={path}")
    else:
        print(f"Embed {i}: None")

print("\nIdentity embeddings cached.")