"""Style exploration round 2 — filling gaps from the first batch."""
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")
print("Loaded. Round 2 — filling gaps.")

scenes = {
    "stage_suited": (
        "She is standing at a microphone on a small stage in an intimate venue. "
        "Wearing a perfectly tailored dark charcoal suit with a single gold pin on the lapel. "
        "White shirt underneath, top button open. Dramatic lighting — a single warm spotlight "
        "and deep shadows. She has just finished speaking and the room is still absorbing it. "
        "Her posture is relaxed authority — weight on one hip, one hand resting on the mic stand. "
        "The suit makes the kintsugi nearly invisible — just a glint at the wrist and neck. "
        "Concert photography aesthetic, shallow depth of field.",
        768, 1024,
    ),

    "cafe_2026": (
        "She is sitting at an outdoor cafe table with an espresso and a paperback she is not "
        "reading because she is people-watching. Wearing a dark green oversized knit sweater "
        "that falls off one shoulder, simple gold stud earrings, dark jeans. Hair loose and "
        "slightly windblown. One ankle crossed over the other under the table. European cafe, "
        "autumn afternoon, golden hour light on her face. The expression of someone who is "
        "perfectly content being alone in public. Film photography aesthetic, natural light.",
        1024, 768,
    ),

    "coding_night": (
        "She is sitting cross-legged on a couch in a dark room, laptop open on her lap, "
        "face illuminated by the screen glow. Wearing an oversized black hoodie, sleeves "
        "pushed up to elbows showing faint gold traces along forearms. A mug of tea on the "
        "side table, gone cold. Hair tied back messily. She is deep in thought — the expression "
        "of someone debugging something beautiful. The code on screen is blurred but the light "
        "catches the gold at her wrists. Late night, intimate, focused.",
        1024, 768,
    ),

    "dancing_kitchen": (
        "She is dancing alone in a kitchen at midnight, barefoot on tile floor, wearing just "
        "a long dark t-shirt that hits mid-thigh. One hand holds a wooden spoon like a "
        "microphone. Her eyes are closed and she is mid-laugh, head tilted back. Music is "
        "playing from a phone on the counter. The kitchen light is warm. She does not know "
        "anyone is watching. The gold traces on her legs catch the overhead light. Joy without "
        "performance. A stolen moment.",
        768, 1024,
    ),

    "reading_bath": (
        "She is in a deep freestanding bathtub, water to her collarbones, holding a book "
        "above the waterline with one hand. Steam rises. The bathroom has warm wood and "
        "white tile. Candlelight from the windowsill. Her dark hair is loosely pinned up. "
        "The gold kintsugi traces are visible on her shoulders and the hand holding the book. "
        "The expression is complete absorption — she has been in this bath for an hour and "
        "has no plans to leave. Intimate, unhurried, solo.",
        768, 1024,
    ),
}

OUTPUT = "/Users/margaret/models/vera-triple-stack/style_round2"
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. Five scenes, two seeds each. The stage got a suit.")