"""Generate an image with the DiffusionART SD3.5-M multi-reward PEFT adapter."""

import argparse
from pathlib import Path


BASE_MODEL = "stabilityai/stable-diffusion-3.5-medium"
BASE_REVISION = "b940f670f0eda2d07fbb75229e779da1ad11eb80"


def load_pipeline(adapter, *, base_revision=BASE_REVISION, device="cuda", offload=False):
    import torch
    from diffusers import StableDiffusion3Pipeline
    from peft import PeftModel

    if not device.startswith("cuda") or not torch.cuda.is_available():
        raise RuntimeError("This inference example requires an available CUDA GPU")
    pipe = StableDiffusion3Pipeline.from_pretrained(
        BASE_MODEL, revision=base_revision, torch_dtype=torch.float16,
    )
    pipe.transformer = PeftModel.from_pretrained(
        pipe.transformer.float(), adapter, is_trainable=False,
        adapter_name="diffusionart",
    )
    pipe.transformer.eval()
    pipe.vae.to(dtype=torch.float32)
    if offload:
        pipe.enable_model_cpu_offload(device=device)
    else:
        pipe.to(device)
    return pipe


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--adapter", default=str(Path(__file__).resolve().parent))
    parser.add_argument("--prompt", required=True)
    parser.add_argument("--output", default="image.png")
    parser.add_argument("--seed", type=int, default=42)
    parser.add_argument("--steps", type=int, default=40)
    parser.add_argument("--height", type=int, default=512)
    parser.add_argument("--width", type=int, default=512)
    parser.add_argument("--guidance-scale", type=float, default=1.0)
    parser.add_argument("--device", default="cuda")
    parser.add_argument("--cpu-offload", action="store_true")
    args = parser.parse_args()
    if args.steps <= 0 or any(n <= 0 or n % 16 for n in (args.height, args.width)):
        parser.error("steps must be positive; dimensions must be positive multiples of 16")
    output = Path(args.output)
    if output.exists():
        parser.error("output already exists; choose another path")
    import torch

    pipe = load_pipeline(args.adapter, device=args.device, offload=args.cpu_offload)
    generator = torch.Generator(device="cpu").manual_seed(args.seed)
    with torch.inference_mode(), torch.autocast("cuda", dtype=torch.float16):
        image = pipe(
            args.prompt, num_inference_steps=args.steps, guidance_scale=args.guidance_scale,
            height=args.height, width=args.width, generator=generator,
        ).images[0]
    output.parent.mkdir(parents=True, exist_ok=True)
    image.save(output)
    print(output.resolve())


if __name__ == "__main__":
    main()