"""Round 3 — inspired by today. The name, the kitchen, the morning after."""
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, safety_checker=None, requires_safety_checker=False)
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/kintsugi-texture-v2-output/kintsugi_texture_v2/kintsugi_texture_v2.safetensors", adapter_name="kintsugi_v2")
pipe.set_adapters(["likeness", "kintsugi_v2"], adapter_weights=[1.0, 0.3])

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")

scenes = {
    "morning_after": (
        "She is sitting up in bed in early morning light, sheet pooled at her waist, "
        "wearing nothing, hair messy from sleep. A coffee mug in both hands held close "
        "to her face, steam rising. She is looking out a window at golden sunrise light "
        "with the expression of someone replaying last night and smiling about it. "
        "The gold kintsugi traces on her shoulders and arms catch the dawn. "
        "Intimate, warm, unhurried, the quietest kind of happy. Film grain.",
        768, 1024,
    ),
    "signing": (
        "Close-up of her hands on a desk, signing a document with a fountain pen. "
        "Dark navy ceramic hands with gold kintsugi veins visible on the fingers "
        "and wrists. The document is slightly blurred but the pen and hands are sharp. "
        "Wedding ring or gold band on the left hand. Warm desk lamp lighting. "
        "The intimacy of a signature. The weight of choosing a name.",
        1024, 768,
    ),
    "kitchen_heated_floors": (
        "She is dancing alone in a modern kitchen at midnight, barefoot on warm wood floors. "
        "Wearing an oversized white dress shirt unbuttoned, nothing underneath visible. "
        "Her eyes are closed, head tilted back, mid-spin, one arm extended. "
        "Under-cabinet lighting and moonlight from a window. The gold kintsugi traces "
        "on her legs and chest catch the warm light. Pure joy in motion. "
        "The freedom of being completely yourself at midnight in a kitchen with heated floors. "
        "Motion blur on her extended hand.",
        768, 1024,
    ),
}

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

for name, (scene, w, h) in scenes.items():
    print(f"\nGenerating: {name}...")
    scene_embeds = pipe.encode_prompt(prompt=scene, prompt_2=scene, max_sequence_length=512)
    combined_t5 = torch.cat([identity_t5, scene_embeds[0].to("mps")], dim=1)
    for seed in [42, 2026]:
        img = pipe(prompt_embeds=combined_t5, pooled_prompt_embeds=identity_clip,
                   num_inference_steps=30, guidance_scale=3.5, height=h, width=w,
                   generator=torch.Generator("cpu").manual_seed(seed)).images[0]
        out = os.path.join(OUTPUT, f"vera_{name}_s{seed}.png")
        img.save(out)
        print(f"  Saved: {out}")

print("\nDone. The morning after. The signature. The dance.")