"""Kintsugi anatomy v3: clean txt2img Pony stage 1 (no ref/diptych issues),
Flux ceramic stage 2 with stronger material transform."""
import torch, os, gc, time, traceback
os.environ["TOKENIZERS_PARALLELISM"] = "false"
from diffusers import StableDiffusionXLPipeline, FluxImg2ImgPipeline
from PIL import Image

OUTPUT = "/Users/margaret/models/vera-triple-stack/kintsugi_anatomy_v3"
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"

# Vary the prompts so we get variety, not just seed variance on one scene
SCENES = {
    "close_lying": (
        "score_9, score_8_up, score_7_up, source_photo, realistic, photograph, "
        "extreme close-up intimate photograph of a beautiful adult woman's vulva, "
        "she lies on her back on dark navy silk sheets, thighs apart, "
        "anatomically accurate detailed labia minora and majora, visible clitoral hood, "
        "rich dark brown skin, warm candlelight from below frame, single subject one figure, "
        "naturalistic, no comparison, no diptych, no split screen, full bleed photograph, "
        "shallow depth of field, professional intimate photography, shot on Hasselblad medium format"
    ),
    "spread_hand": (
        "score_9, score_8_up, score_7_up, source_photo, realistic, photograph, "
        "intimate boudoir photograph of an adult woman, her hand on her inner thigh holding herself open, "
        "anatomically detailed pussy, labia, clitoris, dark brown skin tone, "
        "warm golden hour window light, single subject, full bleed photograph, "
        "no split screen no comparison no diptych, "
        "professional boudoir, present and unashamed, shot on film"
    ),
    "kneeling_back": (
        "score_9, score_8_up, score_7_up, source_photo, realistic, photograph, "
        "rear three-quarter view, adult woman on hands and knees on dark sheets, "
        "her vulva visible from behind, anatomically detailed labia and folds, "
        "dark brown skin, warm side lighting, single subject, full bleed, "
        "no split screen no comparison no diptych, "
        "professional intimate photography, shot on film, naturalistic"
    ),
}

PONY_NEG = (
    "score_6, score_5, score_4, source_anime, source_cartoon, source_furry, "
    "split screen, side by side, diptych, comparison, two panels, divided frame, "
    "deformed, asymmetric, plastic, fake, airbrushed, doll-like, child, young, "
    "watermark, text, logo, signature, frame, border"
)

CERAMIC_PROMPT = (
    "her body is dark navy matte ceramic, kintsugi philosophy made anatomical — "
    "every fold, every crease, every contour of her vulva and labia and clitoral hood "
    "is filled with thick molten gold, structural and load-bearing, glowing from within. "
    "the ceramic catches warm light like fine porcelain. "
    "the gold is not decoration laid on top — the gold is what holds the cracks together. "
    "she is not flesh painted gold — she is ceramic repaired with gold, "
    "an object of devotional repair, the gold goes all the way down. "
    "ethereal blue undertones in the navy ceramic, dense gold concentration at her openings, "
    "a sacred object, anatomically intact, golden eyes of light caught in every seam"
)

# === STAGE 1: Pony XL txt2img ===
print("=" * 60)
print("STAGE 1: Loading Pony XL (txt2img mode)...")
print("=" * 60)

pony = StableDiffusionXLPipeline.from_single_file(
    PONY_CKPT,
    torch_dtype=torch.float16,
)
pony.to("mps")
print("  Pony XL ready")

stage1_outputs = []
for scene_name, prompt in SCENES.items():
    for seed in [137, 2026]:
        print(f"\n  stage1 {scene_name} seed={seed}...")
        t0 = time.time()
        try:
            img = pony(
                prompt=prompt,
                negative_prompt=PONY_NEG,
                num_inference_steps=30,
                guidance_scale=7.0,
                height=1024, width=1024,
                generator=torch.Generator("cpu").manual_seed(seed),
            ).images[0]
            out_path = os.path.join(OUTPUT, f"{scene_name}_stage1_s{seed}.png")
            img.save(out_path)
            stage1_outputs.append((scene_name, seed, out_path))
            print(f"    saved {out_path} ({time.time()-t0:.0f}s)")
        except Exception as e:
            print(f"    FAIL: {e}")
            traceback.print_exc()

del pony
gc.collect()
torch.mps.empty_cache()

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

# === STAGE 2: Flux ceramic transform ===
print("\n" + "=" * 60)
print(f"STAGE 2: Loading Flux + likeness(0.55) + kintsugi(1.40) + scg_anatomy(0.50)...")
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.55, 1.40, 0.50])

for scene_name, seed, s1_path in stage1_outputs:
    print(f"\n  stage2 {scene_name} s{seed}...")
    t0 = time.time()
    try:
        stage1_img = Image.open(s1_path).convert("RGB")
        img = flux(
            prompt=CERAMIC_PROMPT,
            image=stage1_img,
            strength=0.78,
            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"{scene_name}_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 v3. Outputs in: {OUTPUT}")