#!/usr/bin/env python3
# Copyright (C) 2026 Danilo M. <danix@danix.xyz>
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License version 2 as
# published by the Free Software Foundation.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
# GNU General Public License for more details.

# Local SDXL-base image generation for Intel Arc B580 (XPU backend).
# Bounded scope by design: SDXL base, single 1024x1024 image, no refiner,
# no Flux, no batching. This is what fits comfortably in 12GB VRAM once the
# fp32 VAE upcast is removed via the fp16-fix VAE. See README.md for limits.

import argparse
import sys
import time
from datetime import datetime
from pathlib import Path

import torch
from diffusers import StableDiffusionXLPipeline, AutoencoderKL

MODELS = {
    "sdxl":    "stabilityai/stable-diffusion-xl-base-1.0",  # default, general
    "realvis": "SG161222/RealVisXL_V5.0",                   # photoreal, permissive
    "pony":    "AstraliteHeart/pony-diffusion-v6-xl",       # unrestricted, versatile
}
DEFAULT_MODEL = "sdxl"
# fp16-fix VAE avoids the fp32 upcast that OOMs a 12GB card at decode time.
VAE = "madebyollin/sdxl-vae-fp16-fix"
OUT_DIR = Path("/out")


def resolve_model(key):
    if key not in MODELS:
        valid = ", ".join(sorted(MODELS))
        raise SystemExit(f"unknown model '{key}'. valid: {valid}")
    return MODELS[key]


def build_pipe(model_key):
    if not torch.xpu.is_available():
        # Fail loud and specific. A silent CPU fallback would run but take
        # minutes per image, which is worse than a clear error.
        print(
            "ERROR: torch.xpu is not available. The Level Zero runtime is not "
            "visible inside this container. Check that you passed --device "
            "/dev/dri and that the base image ships the Intel GPU runtime.",
            file=sys.stderr,
        )
        sys.exit(2)

    repo = resolve_model(model_key)
    vae = AutoencoderKL.from_pretrained(VAE, torch_dtype=torch.float16)
    try:
        # SDXL base publishes an fp16 variant (smaller download).
        pipe = StableDiffusionXLPipeline.from_pretrained(
            repo, vae=vae, torch_dtype=torch.float16,
            use_safetensors=True, variant="fp16",
        )
    except (ValueError, OSError, EnvironmentError):
        # Many fine-tunes do not publish a separate fp16 variant; load the
        # default files at fp16 dtype instead.
        print(f"no fp16 variant for {repo}, loading default files", file=sys.stderr)
        pipe = StableDiffusionXLPipeline.from_pretrained(
            repo, vae=vae, torch_dtype=torch.float16, use_safetensors=True,
        )
    return pipe.to("xpu")


def generate(pipe, prompt, negative, steps, guidance, seed, width, height):
    generator = None
    if seed is not None:
        generator = torch.Generator(device="xpu").manual_seed(seed)

    t = time.time()
    image = pipe(
        prompt,
        negative_prompt=negative,
        num_inference_steps=steps,
        guidance_scale=guidance,
        generator=generator,
        width=width,
        height=height,
    ).images[0]
    torch.xpu.synchronize()
    elapsed = time.time() - t
    return image, elapsed


def main():
    ap = argparse.ArgumentParser(
        description="Generate a single SDXL-base image locally on Intel Arc (XPU)."
    )
    ap.add_argument("prompt", help="Text prompt for the image.")
    ap.add_argument(
        "-m",
        "--model",
        default=DEFAULT_MODEL,
        help=f"Model key: {', '.join(sorted(MODELS))} (default: {DEFAULT_MODEL}).",
    )
    ap.add_argument(
        "-n", "--negative", default="", help="Negative prompt (default: none)."
    )
    ap.add_argument(
        "-s", "--steps", type=int, default=25, help="Inference steps (default: 25)."
    )
    ap.add_argument(
        "-g",
        "--guidance",
        type=float,
        default=7.0,
        help="Guidance scale / CFG (default: 7.0).",
    )
    ap.add_argument(
        "-W",
        "--width",
        type=int,
        default=1024,
        help="Image width, multiple of 8 (default: 1024).",
    )
    ap.add_argument(
        "-H",
        "--height",
        type=int,
        default=1024,
        help="Image height, multiple of 8 (default: 1024).",
    )
    ap.add_argument(
        "--seed",
        type=int,
        default=None,
        help="Seed for reproducible output (default: random).",
    )
    ap.add_argument(
        "-o",
        "--out",
        default=None,
        help="Output filename inside /out (default: timestamped).",
    )
    args = ap.parse_args()

    if args.width % 8 or args.height % 8:
        ap.error("--width and --height must be multiples of 8 (SDXL constraint).")

    OUT_DIR.mkdir(parents=True, exist_ok=True)

    if args.out:
        out_path = OUT_DIR / args.out
    else:
        stamp = datetime.now().strftime("%Y%m%d-%H%M%S")
        out_path = OUT_DIR / f"sdxl-{stamp}.png"

    print(f"Loading {args.model} + fp16-fix VAE on XPU ...", file=sys.stderr)
    t0 = time.time()
    pipe = build_pipe(args.model)
    print(f"Loaded in {time.time() - t0:.1f}s", file=sys.stderr)

    image, elapsed = generate(
        pipe, args.prompt, args.negative, args.steps, args.guidance, args.seed,
        args.width, args.height,
    )
    image.save(out_path)
    print(f"Generated in {elapsed:.2f}s")
    print(f"Saved {out_path}")


if __name__ == "__main__":
    main()