example.py

from __future__ import annotations

import argparse
import json
from pathlib import Path

import torch
from diffusers import DiffusionPipeline
from PIL import Image, ImageOps

from face_crop import YuNetFaceCropper


WEIGHT_NAME = "krea2_reid_rank32.safetensors"


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Run Krea 2 ReID reference generation")
    parser.add_argument("--reference", type=Path, required=True)
    parser.add_argument("--output", type=Path, required=True)
    parser.add_argument("--prompt", required=True)
    parser.add_argument("--width", type=int, default=1024)
    parser.add_argument("--height", type=int, default=1024)
    parser.add_argument("--steps", type=int, default=8)
    parser.add_argument("--seed", type=int, default=42)
    parser.add_argument("--lora-scale", type=float, default=1.0)
    parser.add_argument(
        "--no-face-crop",
        action="store_true",
        help="Use the reference exactly as supplied instead of applying the optional YuNet crop",
    )
    parser.add_argument(
        "--save-reference-crop",
        type=Path,
        help="Optionally save the reference actually passed to the generation pipeline",
    )
    parser.add_argument(
        "--full-gpu",
        action="store_true",
        help="Keep the BF16 pipeline on CUDA instead of enabling CPU offload",
    )
    return parser.parse_args()


def main() -> None:
    args = parse_args()
    root = Path(__file__).resolve().parent
    reference = ImageOps.exif_transpose(Image.open(args.reference)).convert("RGB")
    if args.no_face_crop:
        print("Reference preprocessing: automatic face crop disabled")
    else:
        cropper = YuNetFaceCropper(root / "models" / "face_detection_yunet_2023mar_int8.onnx")
        crop_result = cropper.crop(reference)
        reference = crop_result.image
        print("Reference preprocessing: " + json.dumps(crop_result.metadata))

    if args.save_reference_crop is not None:
        args.save_reference_crop.parent.mkdir(parents=True, exist_ok=True)
        reference.save(args.save_reference_crop)

    pipe = DiffusionPipeline.from_pretrained(
        "krea/Krea-2-Turbo",
        custom_pipeline=str(root),
        torch_dtype=torch.bfloat16,
    )
    pipe.load_lora_weights(root, weight_name=WEIGHT_NAME, adapter_name="reid")
    pipe.set_adapters(["reid"], adapter_weights=[args.lora_scale])
    if args.full_gpu:
        pipe.to("cuda")
    else:
        pipe.enable_model_cpu_offload()

    generator = torch.Generator(device="cpu").manual_seed(args.seed)
    result = pipe(
        prompt=args.prompt,
        image=reference,
        width=args.width,
        height=args.height,
        num_inference_steps=args.steps,
        guidance_scale=0.0,
        generator=generator,
        reference_max_pixels=384 * 384,
        vl_image_max_pixels=384 * 384,
        encode_reference_in_prompt=True,
        kv_cache=True,
    ).images[0]

    args.output.parent.mkdir(parents=True, exist_ok=True)
    result.save(args.output)


if __name__ == "__main__":
    main()