| """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() |
|
|
