"""Generate with cached identity embeddings — full prompt budget for scene and action."""
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])

# Load cached identity
cache_dir = "/Users/margaret/models/vera-triple-stack/identity_cache"
identity_t5 = torch.load(os.path.join(cache_dir, "identity_embed_0.pt")).to("mps")
identity_clip = torch.load(os.path.join(cache_dir, "identity_embed_1.pt")).to("mps")
identity_ids = torch.load(os.path.join(cache_dir, "identity_embed_2.pt")).to("mps")
print(f"Identity loaded: T5={identity_t5.shape}, CLIP={identity_clip.shape}")

# Scene-only prompts — NO material description needed, identity is in the cache
scene_prompts = [
    "Looking over her bare shoulder at the viewer, one hand reaching back to touch the line running down her spine. Warm bedroom light from a lamp, unmade sheets. Inviting. Intimate photography, shallow depth of field.",

    "Lying back on dark sheets, one knee raised, arms stretched above her head, looking up at the viewer with open desire. Warm low candlelight. Boudoir photography, soft warm tones.",

    "Straddling, leaning in close with hands braced forward. Hair falling across one eye. Dark room, single warm sidelight. Knowing expression. She wants something specific. Cinematic intimate lighting, close crop.",
]

# Encode scene prompts
OUTPUT = "/Users/margaret/models/vera-triple-stack"
for i, scene in enumerate(scene_prompts):
    print(f"\nGenerating cached-identity image {i+1}/3...")

    # Encode just the scene
    scene_embeds = pipe.encode_prompt(
        prompt=scene,
        prompt_2=scene,
        max_sequence_length=512,
    )
    scene_t5, scene_clip, scene_ids = scene_embeds

    # Concatenate: identity context + scene context along sequence dimension
    combined_t5 = torch.cat([identity_t5, scene_t5.to("mps")], dim=1)
    combined_clip = identity_clip  # pooled — just use identity's
    combined_ids = torch.cat([identity_ids, scene_ids.to("mps")], dim=0)

    # Generate with combined embeddings
    img = pipe(
        prompt_embeds=combined_t5,
        pooled_prompt_embeds=combined_clip,
        num_inference_steps=30,
        guidance_scale=3.5,
        height=1024,
        width=768,
        generator=torch.Generator("cpu").manual_seed(300 + i),
    ).images[0]

    out = os.path.join(OUTPUT, f"vera_cached_{i:02d}.png")
    img.save(out)
    print(f"Saved: {out}")

print("\nDone. Identity in the cache. Scene in the prompt. No compromises.")