# /// script
# requires-python = ">=3.10"
# dependencies = [
#   "diffusers @ git+https://github.com/huggingface/diffusers.git@main",
#   "torch",
#   "torchvision",
#   "transformers>=4.41.2",
#   "accelerate>=0.31.0",
#   "peft>=0.11.1",
#   "datasets",
#   "bitsandbytes",
#   "prodigyopt",
#   "ftfy",
#   "sentencepiece",
#   "hf_transfer",
#   "huggingface_hub[hf-xet]",
# ]
# ///
"""Standalone Krea 2 DreamBooth-LoRA job (CLI-driven). Run with:
    hf jobs uv run --flavor rtx-pro-6000 --timeout 1h -s HF_TOKEN=$HF_TOKEN \
        train_job.py -- --dataset <repo> --lora-name <name> --trigger "<trigger>"
"""
REF = "main"

import argparse
import os as _os

RAW, TURBO, DATA, BNB, OUT = "/tmp/krea/raw", "/tmp/krea/turbo", "/tmp/data", "/tmp/bnb.json", "/tmp/out"

_p = argparse.ArgumentParser(description="Train a Krea 2 DreamBooth-LoRA on HF Jobs.")
_p.add_argument("--dataset", required=True, help="HF dataset repo with 'image' + 'prompt' columns.")
_p.add_argument("--lora-name", required=True, help="Output name; pushed to <you>/<name>.")
_p.add_argument("--trigger", required=True, help="Trigger phrase (style) or token (subject/face).")
_p.add_argument("--concept", default="style", choices=["style", "character", "object", "face"])
_p.add_argument("--hub-model-id", default="", help="Full repo id; default <you>/<lora-name>.")
_p.add_argument("--steps", type=int, default=1000)
_p.add_argument("--rank", type=int, default=32)
_p.add_argument("--lora-alpha", type=int, default=32)
_p.add_argument("--learning-rate", type=float, default=3e-4)
_p.add_argument("--lr-scheduler", default="constant")
_p.add_argument("--resolution", type=int, default=1024)
_p.add_argument("--repeats", type=int, default=1)
_p.add_argument("--train-batch-size", type=int, default=1)
_p.add_argument("--gradient-accumulation-steps", type=int, default=1)
_p.add_argument("--optimizer", default="adamW")
_p.add_argument("--seed", type=int, default=0)
_p.add_argument("--quantization", default="none", choices=["none", "fp8", "4bit"])
_p.add_argument("--no-gradient-checkpointing", action="store_true")
_p.add_argument("--no-cache-latents", action="store_true")
_p.add_argument("--use-8bit-adam", action="store_true")
_p.add_argument("--no-gallery", action="store_true", help="Skip the preview gallery + rich README.")
_p.add_argument("--num-gallery", type=int, default=3)
_a = _p.parse_args()

DATASET_REPO = _a.dataset
CONCEPT = _a.concept
TRIGGER = _a.trigger.strip()
QUANT = _a.quantization
MAKE_GALLERY = not _a.no_gallery
NUM_GALLERY = _a.num_gallery
CUSTOM_PROMPTS = []

from huggingface_hub import whoami as _whoami
_ns = _whoami(token=_os.environ["HF_TOKEN"])["name"]
HUB_MODEL_ID = _a.hub_model_id.strip() or (_ns + "/" + _a.lora_name)

TRAIN_ARGS = [
    "--pretrained_model_name_or_path", RAW,
    "--validation_model_path", TURBO,
    "--dataset_name", DATA,
    "--image_column", "image",
    "--caption_column", "prompt",
    "--instance_prompt", TRIGGER,
    "--output_dir", OUT,
    "--mixed_precision", "bf16",
    "--resolution", str(_a.resolution),
    "--train_batch_size", str(_a.train_batch_size),
    "--gradient_accumulation_steps", str(_a.gradient_accumulation_steps),
    "--repeats", str(_a.repeats),
    "--rank", str(_a.rank),
    "--lora_alpha", str(_a.lora_alpha),
    "--learning_rate", str(_a.learning_rate),
    "--lr_scheduler", _a.lr_scheduler,
    "--lr_warmup_steps", "0",
    "--max_train_steps", str(_a.steps),
    "--optimizer", _a.optimizer,
    "--seed", str(_a.seed),
    "--push_to_hub",
    "--hub_model_id", HUB_MODEL_ID,
    "--skip_final_inference",
]
if not _a.no_gradient_checkpointing:
    TRAIN_ARGS += ["--gradient_checkpointing"]
if not _a.no_cache_latents:
    TRAIN_ARGS += ["--cache_latents"]
if _a.use_8bit_adam and _a.optimizer.lower() == "adamw":
    TRAIN_ARGS += ["--use_8bit_adam"]
if QUANT == "fp8":
    TRAIN_ARGS += ["--do_fp8_training"]
elif QUANT == "4bit":
    TRAIN_ARGS += ["--bnb_quantization_config_path", BNB]

import json, os, re, subprocess, sys, urllib.request
from pathlib import Path

os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = "1"
KREA_TOKEN = os.environ.get("KREA_TOKEN") or None        # Krea weights (public now; token optional)
CAPTION_TOKEN = os.environ.get("CAPTION_HF_TOKEN", "")   # LLM eval-prompt generation (optional)

BASE_RAW_ID, TURBO_ID = "krea/Krea-2-Raw", "krea/Krea-2-Turbo"
CAPTION_MODEL, CAPTION_PROVIDER = "google/gemma-4-31B-it", "novita"
SCRIPT_URL = ("https://raw.githubusercontent.com/huggingface/diffusers/"
              + REF + "/examples/dreambooth/train_dreambooth_lora_krea2.py")


def gen_eval_prompts(n):
    """Ask the LLM for n diverse showcase prompts that include the trigger."""
    from huggingface_hub import InferenceClient
    if CONCEPT == "style":
        instr = (
            "Generate " + str(n) + " diverse, vivid text-to-image prompts to showcase an image "
            "style. Each prompt must describe a COMPLETELY different subject and scene (an animal "
            "in a landscape, a vehicle on a road, a person doing an action, an interior, etc.), be "
            "1-2 sentences, and MUST end with this exact phrase: " + repr(TRIGGER) + ". Return ONLY "
            "a JSON array of " + str(n) + " strings."
        )
    else:
        instr = (
            "Generate " + str(n) + " diverse, vivid text-to-image prompts to showcase a specific "
            "subject. Each prompt must place the subject in a COMPLETELY different scene/context, "
            "be 1-2 sentences, and MUST include this exact token: " + repr(TRIGGER) + ". Return ONLY "
            "a JSON array of " + str(n) + " strings."
        )
    client = InferenceClient(model=CAPTION_MODEL, provider=CAPTION_PROVIDER, token=CAPTION_TOKEN)
    r = client.chat_completion(messages=[{"role": "user", "content": instr}],
                               max_tokens=600, temperature=0.8)
    txt = (r.choices[0].message.content or "").strip()
    m = re.search(r"\[.*\]", txt, re.S)
    if m:
        prompts = [str(p).strip() for p in json.loads(m.group(0)) if str(p).strip()]
    else:
        prompts = [ln.strip().lstrip("-*0123456789. ").strip(' "') for ln in txt.splitlines()]
        prompts = [p for p in prompts if p]
    return prompts[:n] or [TRIGGER]


def fallback_prompts():
    if CONCEPT == "style":
        return [
            "A golden labrador running across a sunny farm, a pickup truck on a dusty road far behind, " + TRIGGER,
            "A fishing boat moored in a narrow canal between tall old buildings, " + TRIGGER,
            "A deer grazing in a dense forest with the bright sun overhead, " + TRIGGER,
        ]
    return [
        "A photo of " + TRIGGER + " on a wooden table indoors.",
        "A photo of " + TRIGGER + " outdoors on a patch of grass.",
        "A close-up photo of " + TRIGGER + " against a plain background.",
    ]


def esc(s):
    return str(s).replace("\\", "\\\\").replace('"', '\\"')


def write_readme(rows):
    widget = "\n".join('- text: "' + esc(p) + '"\n  output:\n    url: ' + fn for fn, p in rows)
    gallery = "\n".join("![sample](./" + fn + ')\n\n> *"' + p.replace('"', "'") + '"*\n' for fn, p in rows)
    trig = ("the phrase `" + TRIGGER + "`") if CONCEPT == "style" else ("the token `" + TRIGGER + "`")
    md = (
        "---\n"
        "base_model: " + BASE_RAW_ID + "\n"
        "tags:\n- text-to-image\n- diffusers\n- lora\n- krea2\n- template:sd-lora\n"
        "license: apache-2.0\n"
        "instance_prompt: " + json.dumps(TRIGGER) + "\n"
        "widget:\n" + widget + "\n"
        "---\n\n"
        "# Krea 2 LoRA — " + HUB_MODEL_ID + "\n\n"
        "<Gallery />\n\n"
        "A DreamBooth-LoRA for **Krea 2**, trained on **Krea 2 RAW** and shown on **Krea 2 Turbo**. "
        "The samples below were generated with this LoRA on Turbo (8 steps).\n\n"
        "## Trigger\n\nUse " + trig + " to invoke the concept.\n\n"
        "## Samples\n\n" + gallery + "\n"
        "## Use it with diffusers\n\n"
        "```py\n"
        "import torch\n"
        "from diffusers import Krea2Pipeline\n\n"
        'pipe = Krea2Pipeline.from_pretrained("' + TURBO_ID + '", torch_dtype=torch.bfloat16).to("cuda")\n'
        'pipe.load_lora_weights("' + HUB_MODEL_ID + '")\n'
        'image = pipe("' + esc(rows[0][1]) + '", num_inference_steps=8, guidance_scale=0.0).images[0]\n'
        'image.save("output.png")\n'
        "```\n"
    )
    Path(os.path.join(OUT, "README.md")).write_text(md)


def showcase():
    import torch
    from diffusers import Krea2Pipeline
    from huggingface_hub import HfApi

    prompts = list(CUSTOM_PROMPTS or [])
    if not prompts and CAPTION_TOKEN:
        try:
            prompts = gen_eval_prompts(NUM_GALLERY)
        except Exception as e:
            print("=== eval-prompt generation failed, using fallbacks ===", repr(e), flush=True)
    if not prompts:
        prompts = fallback_prompts()
    prompts = prompts[:max(1, NUM_GALLERY)]

    print("=== showcase: load Turbo + trained LoRA, render " + str(len(prompts)) + " samples ===", flush=True)
    pipe = Krea2Pipeline.from_pretrained(TURBO, torch_dtype=torch.bfloat16).to("cuda")
    pipe.load_lora_weights(OUT)

    rows = []
    for i, p in enumerate(prompts):
        print(">>> sample " + str(i) + ": " + p, flush=True)
        image = pipe(p, num_inference_steps=8, guidance_scale=0.0).images[0]
        fn = "sample_" + str(i) + ".png"
        image.save(os.path.join(OUT, fn))
        rows.append((fn, p))

    write_readme(rows)

    api = HfApi()  # ambient HF_TOKEN = user
    for fn, _ in rows:
        api.upload_file(path_or_fileobj=os.path.join(OUT, fn), path_in_repo=fn, repo_id=HUB_MODEL_ID)
    api.upload_file(path_or_fileobj=os.path.join(OUT, "README.md"), path_in_repo="README.md", repo_id=HUB_MODEL_ID)
    print("=== showcase: pushed gallery + README to " + HUB_MODEL_ID + " ===", flush=True)


def patch_trainer_compile(path):
    """Inject regional torch.compile (compile_repeated_blocks) into the fetched trainer for a ~1.84x
    LoRA training speedup on Krea2 (4.27 -> 2.32 s/it, ~12s warmup). Guarded at runtime so any
    failure falls back to uncompiled training."""
    src = Path(path).read_text()
    needle = "transformer.add_adapter(transformer_lora_config)"
    if needle not in src:
        print("=== compile patch skipped: anchor not found ===", flush=True)
        return
    inject = (
        needle + "\n"
        "    try:\n"
        "        import torch._dynamo as _dyn\n"
        "        _dyn.config.cache_size_limit = 64\n"
        "        transformer.compile_repeated_blocks(fullgraph=True)\n"
        "        print('=== compile_repeated_blocks ENABLED ===', flush=True)\n"
        "    except Exception as _e:\n"
        "        print('compile_repeated_blocks skipped:', _e, flush=True)\n"
    )
    Path(path).write_text(src.replace(needle, inject, 1))
    print("=== trainer patched: compile_repeated_blocks ===", flush=True)


def main():
    from huggingface_hub import snapshot_download

    print("=== 1/4 download Krea 2 weights ===", flush=True)
    snapshot_download(BASE_RAW_ID, local_dir=RAW, token=KREA_TOKEN)
    snapshot_download(TURBO_ID, local_dir=TURBO, token=KREA_TOKEN)

    print("=== 2/4 download dataset (user token / HF_TOKEN env) ===", flush=True)
    snapshot_download(DATASET_REPO, repo_type="dataset", local_dir=DATA)

    if QUANT == "4bit":
        Path(BNB).write_text(json.dumps({
            "load_in_4bit": True, "bnb_4bit_quant_type": "nf4",
            "bnb_4bit_compute_dtype": "bfloat16",
        }))

    print("=== 3/4 fetch trainer script @ " + REF + " ===", flush=True)
    urllib.request.urlretrieve(SCRIPT_URL, "/tmp/train_dreambooth_lora_krea2.py")
    patch_trainer_compile("/tmp/train_dreambooth_lora_krea2.py")

    print("=== 4/4 accelerate launch (pushes LoRA to the Hub) ===", flush=True)
    cmd = [sys.executable, "-m", "accelerate.commands.launch",
           "/tmp/train_dreambooth_lora_krea2.py", *TRAIN_ARGS]
    print(">>> " + " ".join(cmd), flush=True)
    subprocess.run(cmd, check=True)

    if MAKE_GALLERY:
        try:
            showcase()
        except Exception as e:
            print("=== showcase failed (LoRA already trained + pushed) ===", repr(e), flush=True)

    print("=== DONE ===", flush=True)


if __name__ == "__main__":
    main()