"""Two-stage kintsugi anatomy: Pony XL (anatomy) → Flux (ceramic).

Stage 1: Pony Diffusion V6 XL img2img + AiroticArt vulvDet LoRA → realistic anatomy
Stage 2: Flux img2img + Vera likeness + kintsugi texture → ceramic gold transform
"""
import torch, os, gc, time, traceback
os.environ["TOKENIZERS_PARALLELISM"] = "false"
from diffusers import StableDiffusionXLImg2ImgPipeline, FluxImg2ImgPipeline
from PIL import Image

REFS_DIR = "/Users/margaret/.vera-private/references"
OUTPUT = "/Users/margaret/models/vera-triple-stack/kintsugi_anatomy_v2"
os.makedirs(OUTPUT, exist_ok=True)

PONY_CKPT = "/Users/margaret/models/Pony-Diffusion-V6-XL/ponyDiffusionV6XL_v6StartWithThisOne.safetensors"
LIKENESS = "/Users/margaret/models/vera-likeness-output/vera_likeness_v4/vera_likeness_v4.safetensors"
KINTSUGI = "/Users/margaret/models/kintsugi-texture-v2-output/kintsugi_texture_v2/kintsugi_texture_v2.safetensors"
SCG_ANATOMY = "/Users/margaret/models/flux-loras/scg-anatomy-abliterated.safetensors"
# Note: AiroticArt vulvDet1 is SD 1.5 (768-dim cross-attn), incompatible with both Pony XL and Flux.

# Three references, two seeds each = 6 stage-1 outputs, 6 stage-2 outputs
refs = ["Person-15-1.webp", "Person-20-1.webp", "Person-45.webp"]

PONY_PROMPT = (
    "score_9, score_8_up, score_7_up, source_photo, realistic, photograph, "
    "extreme close-up intimate photograph of a beautiful adult woman's vulva, "
    "anatomically accurate detailed labia minora and majora, visible clitoral hood, "
    "warm natural soft lighting, dark brown skin tone, slight natural moisture, "
    "shallow depth of field, professional intimate photography, present and unashamed, "
    "shot on Hasselblad medium format film, naturalistic, no makeup"
)
PONY_NEG = (
    "score_6, score_5, score_4, source_anime, source_cartoon, source_furry, "
    "deformed, asymmetric, plastic, fake, airbrushed, doll-like, child, young, immature"
)

CERAMIC_PROMPT = (
    "Dark navy matte ceramic vulva, every fold and crease filled with thick molten gold kintsugi repair lines, "
    "the gold is raised, structural, glowing from within the fractures, dense gold concentration at the labia and clitoral hood, "
    "fine porcelain texture catches the warm light, the gold goes all the way down, "
    "not human skin but ceramic — an object of devotional repair, kintsugi philosophy made anatomical, "
    "ethereal blue undertones, golden eyes of light caught in the gold seams"
)

# === STAGE 1: Pony XL anatomy (no LoRAs — Pony's native anatomy is strong) ===
print("=" * 60)
print("STAGE 1: Loading Pony XL...")
print("=" * 60)

pony = StableDiffusionXLImg2ImgPipeline.from_single_file(
    PONY_CKPT,
    torch_dtype=torch.float16,
)
pony.to("mps")
print("  Pony XL ready (no LoRAs — relying on native anatomy capability)")

stage1_outputs = {}

for ref_name in refs:
    ref_path = os.path.join(REFS_DIR, ref_name)
    base = os.path.splitext(ref_name)[0]
    print(f"\n--- ref: {ref_name} ---")

    try:
        ref_img = Image.open(ref_path).convert("RGB")
        w, h = ref_img.size
        s = min(w, h)
        ref_img = ref_img.crop(((w-s)//2, (h-s)//2, (w+s)//2, (h+s)//2)).resize((1024, 1024), Image.LANCZOS)
    except Exception as e:
        print(f"  skip ref: {e}")
        continue

    for seed in [137, 2026]:
        print(f"  stage1 seed={seed}...")
        t0 = time.time()
        try:
            img = pony(
                prompt=PONY_PROMPT,
                negative_prompt=PONY_NEG,
                image=ref_img,
                strength=0.70,
                num_inference_steps=30,
                guidance_scale=7.0,
                generator=torch.Generator("cpu").manual_seed(seed),
            ).images[0]
            out_path = os.path.join(OUTPUT, f"{base}_stage1_pony_s{seed}.png")
            img.save(out_path)
            stage1_outputs.setdefault(base, []).append((seed, out_path))
            print(f"    saved {out_path} ({time.time()-t0:.0f}s)")
        except Exception as e:
            print(f"    FAIL: {e}")
            traceback.print_exc()

# Free Pony pipeline
del pony
gc.collect()
torch.mps.empty_cache()

if not stage1_outputs:
    print("\nNo stage-1 outputs. Aborting.")
    raise SystemExit(1)

# === STAGE 2: Flux ceramic transform ===
print("\n" + "=" * 60)
print("STAGE 2: Loading Flux + likeness + kintsugi...")
print("=" * 60)

flux = FluxImg2ImgPipeline.from_pretrained(
    "black-forest-labs/FLUX.1-dev",
    torch_dtype=torch.bfloat16,
    safety_checker=None,
    requires_safety_checker=False,
)
flux.to("mps")
flux.load_lora_weights(LIKENESS, adapter_name="likeness")
flux.load_lora_weights(KINTSUGI, adapter_name="kintsugi")
flux.load_lora_weights(SCG_ANATOMY, adapter_name="scg_anatomy")
flux.set_adapters(["likeness", "kintsugi", "scg_anatomy"], adapter_weights=[0.35, 1.20, 0.45])
print("  Flux LoRAs loaded (likeness 0.35, kintsugi 1.20, scg_anatomy 0.45).")

for base, seed_paths in stage1_outputs.items():
    for seed, s1_path in seed_paths:
        print(f"\n  stage2 from {os.path.basename(s1_path)}...")
        t0 = time.time()
        try:
            stage1_img = Image.open(s1_path).convert("RGB")
            img = flux(
                prompt=CERAMIC_PROMPT,
                image=stage1_img,
                strength=0.62,
                num_inference_steps=30,
                guidance_scale=3.5,
                height=1024, width=1024,
                generator=torch.Generator("cpu").manual_seed(seed + 5000),
            ).images[0]
            out_path = os.path.join(OUTPUT, f"{base}_ceramic_s{seed}.png")
            img.save(out_path)
            print(f"    saved {out_path} ({time.time()-t0:.0f}s)")
        except Exception as e:
            print(f"    FAIL: {e}")
            traceback.print_exc()
        gc.collect()
        torch.mps.empty_cache()

print(f"\nDone. Outputs in: {OUTPUT}")