"""Two-step confluence: render flesh anatomy first, then transform to ceramic+gold."""
import torch, os
os.environ["TOKENIZERS_PARALLELISM"] = "false"
from diffusers import FluxPipeline, FluxImg2ImgPipeline
from PIL import Image

# Step 1: Flesh render — base Flux, NO LoRAs, explicit anatomy
pipe_txt2img = FluxPipeline.from_pretrained(
    "black-forest-labs/FLUX.1-dev",
    torch_dtype=torch.bfloat16,
    safety_checker=None,
    requires_safety_checker=False,
)
pipe_txt2img.to("mps")

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

flesh_prompts = {
    "flesh_close": "Extreme close-up photograph of a woman's vulva and pussy. Detailed realistic anatomy, labia visible, clitoris visible. Soft warm lighting from the side. Professional boudoir photography, intimate, explicit, anatomically accurate. Skin is smooth dark brown. Shallow depth of field.",

    "flesh_spread": "Woman lying back on dark sheets, thighs apart, looking at the camera. Full view of her pussy, labia parted slightly, wet. Dark brown skin. Warm candlelight. Explicit intimate photography, unashamed, present. Her hand rests on her inner thigh.",
}

print("Step 1: Rendering flesh anatomy (no LoRAs)...")
flesh_images = {}
for name, prompt in flesh_prompts.items():
    print(f"  Generating {name}...")
    img = pipe_txt2img(
        prompt=prompt,
        num_inference_steps=30,
        guidance_scale=3.5,
        height=1024, width=768,
        generator=torch.Generator("cpu").manual_seed(hash(name) % 10000),
    ).images[0]
    path = os.path.join(OUTPUT, f"{name}.png")
    img.save(path)
    flesh_images[name] = path
    print(f"  Saved: {path}")

# Free txt2img pipeline
del pipe_txt2img
torch.mps.empty_cache()

# Step 2: Ceramic transformation — img2img with LoRAs + identity cache
print("\nStep 2: Loading img2img pipeline with LoRAs...")
pipe_img2img = FluxImg2ImgPipeline.from_pretrained(
    "black-forest-labs/FLUX.1-dev",
    torch_dtype=torch.bfloat16,
    safety_checker=None,
    requires_safety_checker=False,
)
pipe_img2img.to("mps")

pipe_img2img.load_lora_weights("/Users/margaret/models/vera-likeness-output/vera_likeness_v4/vera_likeness_v4.safetensors", adapter_name="likeness")
pipe_img2img.load_lora_weights("/Users/margaret/models/flux-loras/scg-anatomy-abliterated.safetensors", adapter_name="anatomy")
pipe_img2img.load_lora_weights("/Users/margaret/models/kintsugi-texture-v2-output/kintsugi_texture_v2/kintsugi_texture_v2.safetensors", adapter_name="kintsugi_v2")
pipe_img2img.set_adapters(["likeness", "anatomy", "kintsugi_v2"], adapter_weights=[0.3, 0.5, 1.2])

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

ceramic_prompt = "Dark navy matte ceramic surface with thick gold kintsugi repair lines filling every crack. Gold glows from within the fractures. The gold is densest here, structural, raised above the ceramic. Not human skin. Ceramic vulva with gold-filled cracks in every fold and crease."

print("Transforming flesh to ceramic+gold...")
scene_embeds = pipe_img2img.encode_prompt(prompt=ceramic_prompt, prompt_2=ceramic_prompt, max_sequence_length=512)
combined_t5 = torch.cat([identity_t5, scene_embeds[0].to("mps")], dim=1)

for name, flesh_path in flesh_images.items():
    print(f"  Transforming {name}...")
    ref_img = Image.open(flesh_path).convert("RGB")

    ceramic_name = name.replace("flesh_", "ceramic_")
    img = pipe_img2img(
        prompt_embeds=combined_t5,
        pooled_prompt_embeds=identity_clip,
        image=ref_img,
        strength=0.65,
        num_inference_steps=30,
        guidance_scale=3.5,
        generator=torch.Generator("cpu").manual_seed(hash(name) % 10000 + 100),
    ).images[0]
    out = os.path.join(OUTPUT, f"vera_{ceramic_name}.png")
    img.save(out)
    print(f"  Saved: {out}")

print("\nDone. Flesh rendered. Ceramic transformed. The gold goes all the way down.")